chore: import upstream snapshot with attribution

python/sglang/jit_kernel/csrc/lplb/dispatch_probability.cuh

@@ -0,0 +1,122 @@
+// LP probability dispatch kernel: collapse the ~7 torch ops in
+// _topk_ids_logical_to_physical_probability into one launch.
+//
+// Python equivalent:
+//
+//     topk_probs = log2phy_prob[topk_ids]                           # gather
+//     row_sums   = topk_probs.sum(dim=-1)                            # reduce
+//     fallback   = (log2phy_map[topk_ids] >= 0).float()              # gather + cast
+//     topk_probs = where(row_sums > 0, topk_probs, fallback)         # cmp + select
+//     chosen     = multinomial(topk_probs, 1).flatten()              # sample
+//     out        = log2phy_map[topk_ids, chosen]                     # gather
+//
+// Each thread handles one (token, slot) in the flattened topk_ids.
+// `random_vals` is pre-generated by the caller via torch.rand (one kernel
+// launch, ~5 µs); we sample the multinomial via prefix-sum + comparison
+// against `random_vals[i] * row_sum`.
+//
+// Templated on (MAX_COPIES, BLOCK_DIM). MAX_COPIES is small (typically 2-3)
+// so the per-row prefix sum unrolls into a few instructions per thread.
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+template <int MAX_COPIES, int BLOCK_DIM>
+__global__ void dispatch_probability_kernel(
+    int32_t* __restrict__ out_topk_ids,       // (N,)
+    const int32_t* __restrict__ in_topk_ids,  // (N,)
+    const float* __restrict__ log2phy_prob,   // (NUM_LOGICAL, MAX_COPIES)
+    const int32_t* __restrict__ log2phy_map,  // (NUM_LOGICAL, MAX_COPIES)
+    const float* __restrict__ random_vals,    // (N,)
+    int N) {
+  const int idx = blockIdx.x * BLOCK_DIM + threadIdx.x;
+  if (idx >= N) return;
+
+  const int32_t logical_id = in_topk_ids[idx];
+  const int32_t* row_map = log2phy_map + logical_id * MAX_COPIES;
+  const float* row_prob = log2phy_prob + logical_id * MAX_COPIES;
+
+  float probs[MAX_COPIES];
+  int32_t maps[MAX_COPIES];
+  float row_sum = 0.f;
+
+#pragma unroll
+  for (int c = 0; c < MAX_COPIES; c++) {
+    maps[c] = row_map[c];
+    probs[c] = row_prob[c];
+    row_sum += probs[c];
+  }
+
+  // Fallback: if all LP probs for this row are 0, sample uniformly from the
+  // valid physical copies (map != -1).
+  if (row_sum <= 0.f) {
+    row_sum = 0.f;
+#pragma unroll
+    for (int c = 0; c < MAX_COPIES; c++) {
+      probs[c] = (maps[c] >= 0) ? 1.0f : 0.0f;
+      row_sum += probs[c];
+    }
+  }
+
+  // Multinomial sample: smallest c such that cumsum[0..c] > u * row_sum.
+  // Implemented branch-free for unroll friendliness: chosen accumulates the
+  // largest index where cumsum[..c] is still <= u, then we add 1 (clamped to
+  // MAX_COPIES-1 for the all-cumsum-<=-u edge case from float rounding).
+  const float u = random_vals[idx] * row_sum;
+  float cum = 0.f;
+  int chosen = 0;
+#pragma unroll
+  for (int c = 0; c < MAX_COPIES; c++) {
+    cum += probs[c];
+    if (u >= cum) chosen = c + 1;
+  }
+  if (chosen >= MAX_COPIES) chosen = MAX_COPIES - 1;
+
+  out_topk_ids[idx] = maps[chosen];
+}
+
+template <int MAX_COPIES, int BLOCK_DIM>
+void dispatch_probability(
+    tvm::ffi::TensorView out_topk_ids,   // (N,) int32
+    tvm::ffi::TensorView in_topk_ids,    // (N,) int32
+    tvm::ffi::TensorView log2phy_prob,   // (NUM_LOGICAL, MAX_COPIES) float32
+    tvm::ffi::TensorView log2phy_map,    // (NUM_LOGICAL, MAX_COPIES) int32
+    tvm::ffi::TensorView random_vals) {  // (N,) float32
+  using namespace host;
+
+  SymbolicSize N{"num_topk_entries"};
+  SymbolicSize NUM_LOGICAL{"num_logical"};
+  SymbolicDevice device_;
+
+  TensorMatcher({N}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(out_topk_ids).verify(in_topk_ids);
+  TensorMatcher({NUM_LOGICAL, MAX_COPIES}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(log2phy_prob);
+  TensorMatcher({NUM_LOGICAL, MAX_COPIES}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(log2phy_map);
+  TensorMatcher({N}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(random_vals);
+
+  const int n = static_cast<int>(N.unwrap());
+  const int grid = (n + BLOCK_DIM - 1) / BLOCK_DIM;
+  const DLDevice device = device_.unwrap();
+
+  using KernelT = void (*)(int32_t*, const int32_t*, const float*, const int32_t*, const float*, int);
+  KernelT kernel = dispatch_probability_kernel<MAX_COPIES, BLOCK_DIM>;
+
+  LaunchKernel(grid, BLOCK_DIM, device)(
+      kernel,
+      static_cast<int32_t*>(out_topk_ids.data_ptr()),
+      static_cast<const int32_t*>(in_topk_ids.data_ptr()),
+      static_cast<const float*>(log2phy_prob.data_ptr()),
+      static_cast<const int32_t*>(log2phy_map.data_ptr()),
+      static_cast<const float*>(random_vals.data_ptr()),
+      n);
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/lplb/ipm.cuh

@@ -0,0 +1,275 @@
+// Single-SM Interior Point Method (IPM) LP solver.
+//
+// Solves min c^T x subject to A x = b, x >= 0 with a barrier method:
+//   for step in 0..NUM_ITERS:
+//     ax2  = A * x^2
+//     ax2a = ax2 @ A^T          (cuBLASDx GEMM)
+//     ax2c = ax2 @ c            (cuBLASDx GEMM)
+//     d    = solve(ax2a, ax2c)  (cuSolverDx Cholesky/POSV)
+//     r    = d^T @ A
+//     d    = x * (c - r)
+//     x   *= 1 - 0.999 * d / max(d)
+//
+// Convergence is checked at the end and the kernel writes 0.5 to every
+// element on non-convergence (matches the historical Numba behavior).
+//
+// Adapted from DeepSeek-AI/LPLB's `minilp.cu`. Templated on (NC, NV,
+// BLOCK_DIM, SM_VER, NUM_ITERS) so each unique shape is compiled once via
+// sglang's tvm-ffi load_jit cache.
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+#include <cublasdx.hpp>
+
+namespace {
+
+template <int NC, int NV>
+struct ipm_smem {
+  float b[NC];
+  float a[NC][NV];
+  float c[NV];
+  float ax2[NC][NV];
+  float ax2a[NC][NC];
+  float x[NV];
+  float ax2c[NC];
+  float r[NV];
+  float d[NV];
+  float alpha;
+  bool avail_flag;
+};
+
+// In-place Cholesky factorization a = L L^T (lower triangle), no external
+// linkage. Replaces cuSolverDx::posv to keep the kernel self-contained
+// under sglang's tvm-ffi load_jit (which uses plain c++ for the final
+// link step and so cannot satisfy cuSolverDx's device-link requirement).
+//
+// For the typical LPLB shape (N <= 32) the algorithm is dwarfed by the
+// cuBLASDx GEMMs.
+template <int N, int BLOCK_DIM>
+__device__ __forceinline__ void cholesky_factor(float a[N][N]) {
+  const int tid = threadIdx.x;
+  for (int k = 0; k < N; k++) {
+    if (tid == 0) {
+      // Clamp the pivot away from zero before sqrtf. Numerical drift in
+      // the IPM iterations can push a[k][k] slightly negative on
+      // otherwise-PSD matrices, which would produce NaN and propagate
+      // through the rest of the solve. The convergence check at the end
+      // of the kernel writes 0.5 on non-convergence regardless.
+      a[k][k] = sqrtf(fmaxf(a[k][k], 1e-12f));
+    }
+    __syncthreads();
+    const float pivot = a[k][k];
+    for (int i = k + 1 + tid; i < N; i += BLOCK_DIM) {
+      a[i][k] /= pivot;
+    }
+    __syncthreads();
+    // Schur complement: a[i][j] -= a[i][k] * a[j][k] for j>k, i>=j
+    for (int idx = tid; idx < N * N; idx += BLOCK_DIM) {
+      const int i = idx / N, j = idx % N;
+      if (j > k && i >= j && i < N) {
+        a[i][j] -= a[i][k] * a[j][k];
+      }
+    }
+    __syncthreads();
+  }
+}
+
+// Solve L L^T x = b in-place on b, where L is the lower triangle of a
+// (filled by `cholesky_factor`). Forward then back substitution; both
+// run on a single thread because N is small and the inner loops have
+// loop-carried dependencies that don't parallelize cheaply.
+template <int N, int BLOCK_DIM>
+__device__ __forceinline__ void cholesky_apply(const float a[N][N], float b[N]) {
+  const int tid = threadIdx.x;
+  if (tid == 0) {
+    for (int i = 0; i < N; i++) {
+      float s = b[i];
+      for (int j = 0; j < i; j++) {
+        s -= a[i][j] * b[j];
+      }
+      b[i] = s / a[i][i];
+    }
+    for (int i = N - 1; i >= 0; i--) {
+      float s = b[i];
+      for (int j = i + 1; j < N; j++) {
+        s -= a[j][i] * b[j];
+      }
+      b[i] = s / a[i][i];
+    }
+  }
+  __syncthreads();
+}
+
+template <int N, int SM_VER, int BLOCK_DIM>
+__device__ __forceinline__ void cholesky_solve(float a[N][N], float b[N]) {
+  cholesky_factor<N, BLOCK_DIM>(a);
+  cholesky_apply<N, BLOCK_DIM>(a, b);
+}
+
+template <int M, int N, int K, int SM_VER, int BLOCK_DIM>
+__device__ __forceinline__ void matmul_NT(float* a, float* b, float* c) {
+  decltype(cublasdx::Size<M, N, K>() + cublasdx::Function<cublasdx::function::MM>() + cublasdx::Arrangement<cublasdx::row_major, cublasdx::col_major>() + cublasdx::SM<SM_VER>() + cublasdx::Block() + cublasdx::BlockDim<BLOCK_DIM>())()
+      .execute(1.f, a, b, 0.f, c);
+}
+
+template <int M, int N, int K, int SM_VER, int BLOCK_DIM>
+__device__ __forceinline__ void matmul_NN(float* a, float* b, float* c) {
+  decltype(cublasdx::Size<M, N, K>() + cublasdx::Function<cublasdx::function::MM>() + cublasdx::Arrangement<cublasdx::row_major, cublasdx::row_major>() + cublasdx::SM<SM_VER>() + cublasdx::Block() + cublasdx::BlockDim<BLOCK_DIM>())()
+      .execute(1.f, a, b, 0.f, c);
+}
+
+template <int NC, int NV, int BLOCK_DIM, int SM_VER, int NUM_ITERS>
+__global__ void ipm_solve_kernel(
+    float* __restrict__ result,
+    const float* __restrict__ input_a,
+    const float* __restrict__ input_b,
+    const float* __restrict__ input_c) {
+  using SMem = ipm_smem<NC, NV>;
+  extern __shared__ unsigned char raw_smem[];
+  SMem* smem = reinterpret_cast<SMem*>(raw_smem);
+
+  const int tid = threadIdx.x;
+  const int dim = blockDim.x;
+
+  auto& a = smem->a;
+  auto& b = smem->b;
+  auto& c = smem->c;
+
+  // Load A, b, c into shared memory (single block, no grid index).
+  for (int i = tid; i < NC * NV; i += dim) {
+    int ic = i / NV, iv = i % NV;
+    a[ic][iv] = input_a[i];
+  }
+  __syncthreads();
+  for (int i = tid; i < NC; i += dim) {
+    b[i] = input_b[i];
+  }
+  __syncthreads();
+  for (int i = tid; i < NV; i += dim) {
+    c[i] = input_c[i];
+  }
+  __syncthreads();
+
+  auto& ax2 = smem->ax2;
+  auto& ax2a = smem->ax2a;
+  auto& x = smem->x;
+  auto& ax2c = smem->ax2c;
+  auto& r = smem->r;
+  auto& d = smem->d;
+  auto& alpha = smem->alpha;
+  // d_max and max_residual are warp-scope reductions; only valid for tid<32.
+  float d_max = 0.f;
+  float max_residual = 0.f;
+
+  for (int j = tid; j < NV; j += dim) {
+    x[j] = 1.f;
+  }
+  __syncthreads();
+
+  for (int step = 0; step < NUM_ITERS; step++) {
+    for (int ij = tid; ij < NC * NV; ij += dim) {
+      int i = ij / NV, j = ij % NV;
+      ax2[i][j] = a[i][j] * x[j] * x[j];
+    }
+    __syncthreads();
+
+    matmul_NT<NC, NC, NV, SM_VER, BLOCK_DIM>(ax2[0], a[0], ax2a[0]);
+    matmul_NT<NC, 1, NV, SM_VER, BLOCK_DIM>(ax2[0], c, ax2c);
+    cholesky_solve<NC, SM_VER, BLOCK_DIM>(ax2a, ax2c);
+    matmul_NN<1, NV, NC, SM_VER, BLOCK_DIM>(ax2c, a[0], r);
+
+    if (tid < 32) {
+      d_max = 0.f;
+      for (int j = tid; j < NV; j += 32) {
+        float val = x[j] * (c[j] - r[j]);
+        d[j] = val;
+        d_max = fmaxf(d_max, val);
+      }
+      for (int offset = 16; offset > 0; offset >>= 1) {
+        d_max = fmaxf(d_max, __shfl_xor_sync(0xffffffff, d_max, offset));
+      }
+      if (tid == 0) {
+        // Guard against d_max <= 0 from a degenerate / numerically-stuck
+        // iteration. A non-positive d_max would yield inf/NaN from the
+        // division and corrupt x on the next update. The 1.0 fallback
+        // produces a no-op step (x *= 1 - 1*0 = x) so the solver simply
+        // stalls rather than diverges, and the convergence check at the
+        // end of the kernel writes 0.5 if d_max stays small.
+        alpha = (d_max > 1e-9f) ? (0.999f / d_max) : 1.0f;
+      }
+    }
+    __syncthreads();
+
+    for (int j = tid; j < NV; j += dim) {
+      x[j] *= 1.f - alpha * d[j];
+    }
+    __syncthreads();
+  }
+
+  // Compute residual ‖A x - b‖_inf for the convergence check.
+  matmul_NT<NC, 1, NV, SM_VER, BLOCK_DIM>(a[0], x, ax2c);
+  if (tid < 32) {
+    max_residual = 0.f;
+    for (int i = tid; i < NC; i += 32) {
+      max_residual = fmaxf(max_residual, fabsf(ax2c[i] - b[i]));
+    }
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      max_residual = fmaxf(max_residual, __shfl_down_sync(0xffffffff, max_residual, offset));
+    }
+  }
+
+  auto& avail_flag = smem->avail_flag;
+  if (tid == 0) {
+    avail_flag = (d_max < 0.1f && x[NV - 1] >= 0.f && x[NV - 1] < 1e-4f && max_residual < 0.05f);
+  }
+  __syncthreads();
+
+  if (!avail_flag) {
+    for (int i = tid; i < NV; i += dim) {
+      result[i] = 0.5f;
+    }
+  } else {
+    for (int i = tid; i < NV; i += dim) {
+      result[i] = x[i];
+    }
+  }
+}
+
+template <int NC, int NV, int BLOCK_DIM, int SM_VER, int NUM_ITERS>
+void ipm_solve(tvm::ffi::TensorView A, tvm::ffi::TensorView b, tvm::ffi::TensorView c, tvm::ffi::TensorView result) {
+  using namespace host;
+
+  SymbolicDevice device_;
+  TensorMatcher({NC, NV}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(A);
+  TensorMatcher({NC}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(b);
+  TensorMatcher({NV}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(c);
+  TensorMatcher({NV}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(result);
+
+  const DLDevice device = device_.unwrap();
+  const size_t smem_bytes = sizeof(ipm_smem<NC, NV>);
+
+  using KernelT = void (*)(float*, const float*, const float*, const float*);
+  KernelT kernel = ipm_solve_kernel<NC, NV, BLOCK_DIM, SM_VER, NUM_ITERS>;
+
+  // Opt in to >48 KB dynamic shared memory if needed (Hopper supports up to
+  // 228 KB per block).
+  if (smem_bytes > 48 * 1024) {
+    cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(smem_bytes));
+  }
+
+  LaunchKernel(/*grid_dim=*/1, /*block_dim=*/BLOCK_DIM, device, smem_bytes)(
+      kernel,
+      static_cast<float*>(result.data_ptr()),
+      static_cast<const float*>(A.data_ptr()),
+      static_cast<const float*>(b.data_ptr()),
+      static_cast<const float*>(c.data_ptr()));
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/lplb/lp_post.cuh

@@ -0,0 +1,117 @@
+// LP post kernel: build the final log2phy_prob tensor from the IPM output
+// `x` and the prep's `t1`.
+//
+// Python equivalent (~5 torch ops; this kernel is one launch):
+//
+//     x_ratios   = clamp(x[:NUM_RED_PHY], min=0)
+//     phy_prob   = zeros(NUM_SINGLE + NUM_RED_PHY + 1)         # +1 = sink slot
+//     phy_prob[phy_replicated] = x_ratios
+//     phy_prob[phy_single]      = t1
+//     log2phy_prob = take(phy_prob, log2phy)                    # (-1 wraps to sink)
+//
+// `log2phy` may contain -1 for unused replicas (DP-attention padding); we
+// emulate torch.take's wrap-around by adding `phy_prob_size` to negative
+// indices, which lands at the always-zero sink slot.
+//
+// Single-block launch. `phy_prob` lives in shared memory.
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+template <
+    int NUM_LOGICAL,
+    int MAX_COPIES,
+    int NUM_SINGLE,
+    int NUM_RED_PHY,
+    int BLOCK_DIM>
+__global__ void lp_post_kernel(
+    float* __restrict__ log2phy_prob,            // (NUM_LOGICAL, MAX_COPIES) — written
+    const float* __restrict__ x,                 // (NV,) — IPM output
+    const float* __restrict__ t1,                // (NUM_SINGLE,) — from prep
+    const int64_t* __restrict__ phy_single,      // (NUM_SINGLE,)
+    const int64_t* __restrict__ phy_replicated,  // (NUM_RED_PHY,)
+    const int64_t* __restrict__ log2phy) {       // (NUM_LOGICAL, MAX_COPIES)
+  constexpr int PHY_PROB_SIZE = NUM_SINGLE + NUM_RED_PHY + 1;
+
+  extern __shared__ unsigned char raw_smem[];
+  float* phy_prob = reinterpret_cast<float*>(raw_smem);
+
+  const int tid = threadIdx.x;
+
+  // Stage 1: zero-init phy_prob (covers the sink slot at index PHY_PROB_SIZE-1).
+  for (int i = tid; i < PHY_PROB_SIZE; i += BLOCK_DIM) {
+    phy_prob[i] = 0.f;
+  }
+  __syncthreads();
+
+  // Stage 2: scatter x_ratios = clamp(x[:NUM_RED_PHY], min=0) at phy_replicated.
+  for (int i = tid; i < NUM_RED_PHY; i += BLOCK_DIM) {
+    int64_t idx = phy_replicated[i];
+    phy_prob[idx] = fmaxf(x[i], 0.f);
+  }
+  // Stage 3: scatter t1 at phy_single.
+  for (int i = tid; i < NUM_SINGLE; i += BLOCK_DIM) {
+    int64_t idx = phy_single[i];
+    phy_prob[idx] = t1[i];
+  }
+  __syncthreads();
+
+  // Stage 4: gather log2phy_prob[i,j] = phy_prob[log2phy[i,j]].
+  // -1 entries wrap to the sink slot (PHY_PROB_SIZE - 1), which is 0.
+  const int total = NUM_LOGICAL * MAX_COPIES;
+  for (int idx = tid; idx < total; idx += BLOCK_DIM) {
+    int64_t k = log2phy[idx];
+    if (k < 0) k += PHY_PROB_SIZE;
+    log2phy_prob[idx] = phy_prob[k];
+  }
+}
+
+template <int NUM_LOGICAL, int MAX_COPIES, int NUM_SINGLE, int NUM_RED_PHY, int BLOCK_DIM>
+void lp_post(
+    tvm::ffi::TensorView log2phy_prob,
+    tvm::ffi::TensorView x,
+    tvm::ffi::TensorView t1,
+    tvm::ffi::TensorView phy_single,
+    tvm::ffi::TensorView phy_replicated,
+    tvm::ffi::TensorView log2phy) {
+  using namespace host;
+
+  SymbolicDevice device_;
+  TensorMatcher({NUM_LOGICAL, MAX_COPIES}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(log2phy_prob);
+  TensorMatcher({NUM_SINGLE}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(t1);
+  TensorMatcher({NUM_SINGLE}).with_dtype<int64_t>().with_device<kDLCUDA>(device_).verify(phy_single);
+  TensorMatcher({NUM_RED_PHY}).with_dtype<int64_t>().with_device<kDLCUDA>(device_).verify(phy_replicated);
+  TensorMatcher({NUM_LOGICAL, MAX_COPIES}).with_dtype<int64_t>().with_device<kDLCUDA>(device_).verify(log2phy);
+  // x has shape (NV,) which we don't constrain at this layer.
+
+  constexpr int PHY_PROB_SIZE = NUM_SINGLE + NUM_RED_PHY + 1;
+  const size_t smem_bytes = PHY_PROB_SIZE * sizeof(float);
+
+  using KernelT = void (*)(float*, const float*, const float*, const int64_t*, const int64_t*, const int64_t*);
+  KernelT kernel = lp_post_kernel<NUM_LOGICAL, MAX_COPIES, NUM_SINGLE, NUM_RED_PHY, BLOCK_DIM>;
+
+  if (smem_bytes > 48 * 1024) {
+    cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(smem_bytes));
+  }
+
+  const DLDevice device = device_.unwrap();
+  LaunchKernel(/*grid=*/1, /*block=*/BLOCK_DIM, device, smem_bytes)(
+      kernel,
+      static_cast<float*>(log2phy_prob.data_ptr()),
+      static_cast<const float*>(x.data_ptr()),
+      static_cast<const float*>(t1.data_ptr()),
+      static_cast<const int64_t*>(phy_single.data_ptr()),
+      static_cast<const int64_t*>(phy_replicated.data_ptr()),
+      static_cast<const int64_t*>(log2phy.data_ptr()));
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/lplb/lp_prep.cuh

@@ -0,0 +1,182 @@
+// LP prep kernel: build IPM inputs (A_full last column, b, t1) from
+// global_counts and the constant per-rank metadata (log_single,
+// log_replicated, B1, A_base_row_sum).
+//
+// Python equivalent (~8 torch ops; this kernel is one launch):
+//
+//     total       = global_counts.sum()
+//     counts_norm = global_counts / total.clamp(min=1.0)
+//     t1          = counts_norm[log_single]                      # (NUM_SINGLE,)
+//     b1          = counts_norm[log_replicated]                  # (NUM_RED_LOG,)
+//     b2          = -(B1 @ t1).flatten()                          # (NUM_GPUS,)
+//     b           = cat(b1, b2)                                   # (NC,)
+//     A_full[:, -1] = b - A_base_row_sum                          # last column only
+//
+// `A_full` is pre-allocated by the caller with shape (NC, NV); its first
+// NV-1 columns are pre-filled with A_base.copy_() at solver init and not
+// touched by this kernel.
+//
+// Single-block launch. All intermediate state lives in shared memory.
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+template <int NC, int NV, int NUM_SINGLE, int NUM_RED_LOG, int NUM_GPUS, int BLOCK_DIM>
+__global__ void lp_prep_kernel(
+    float* __restrict__ A_full,                  // (NC, NV) — last column is written
+    float* __restrict__ b,                       // (NC,)    — written
+    float* __restrict__ t1,                      // (NUM_SINGLE,) — written
+    const float* __restrict__ global_counts,     // (num_logical,)
+    const int64_t* __restrict__ log_single,      // (NUM_SINGLE,)
+    const int64_t* __restrict__ log_replicated,  // (NUM_RED_LOG,)
+    const float* __restrict__ B1,                // (NUM_GPUS, NUM_SINGLE)
+    const float* __restrict__ A_base_row_sum) {  // (NC,)
+  static_assert(NC == NUM_RED_LOG + NUM_GPUS, "NC must equal NUM_RED_LOG + NUM_GPUS");
+  constexpr int WARPS_PER_BLOCK = BLOCK_DIM / 32;
+
+  extern __shared__ unsigned char raw_smem[];
+  // Layout: shared_t1 [NUM_SINGLE] | shared_b1 [NUM_RED_LOG] | reduce_buf [WARPS_PER_BLOCK] | total_pad [1]
+  float* shared_t1 = reinterpret_cast<float*>(raw_smem);
+  float* shared_b1 = shared_t1 + NUM_SINGLE;
+  float* reduce_buf = shared_b1 + NUM_RED_LOG;
+  float* shared_total = reduce_buf + WARPS_PER_BLOCK;
+
+  const int tid = threadIdx.x;
+  const int lane = tid & 31;
+  const int warp_id = tid >> 5;
+
+  // ---- Stage 1: gather raw t1 / b1 + partial sum for total ----
+  float local_sum = 0.f;
+  for (int i = tid; i < NUM_SINGLE; i += BLOCK_DIM) {
+    float v = global_counts[log_single[i]];
+    shared_t1[i] = v;  // raw, scaled below
+    local_sum += v;
+  }
+  for (int i = tid; i < NUM_RED_LOG; i += BLOCK_DIM) {
+    float v = global_counts[log_replicated[i]];
+    shared_b1[i] = v;
+    local_sum += v;
+  }
+
+  // Block-level reduction: warp shuffle -> shared mem -> warp 0 final reduce.
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    local_sum += __shfl_xor_sync(0xffffffff, local_sum, offset);
+  }
+  if (lane == 0) {
+    reduce_buf[warp_id] = local_sum;
+  }
+  __syncthreads();
+  if (warp_id == 0) {
+    float v = (tid < WARPS_PER_BLOCK) ? reduce_buf[tid] : 0.f;
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      v += __shfl_xor_sync(0xffffffff, v, offset);
+    }
+    if (tid == 0) {
+      shared_total[0] = fmaxf(v, 1.0f);  // clamp(min=1.0)
+    }
+  }
+  __syncthreads();
+  const float inv_total = 1.0f / shared_total[0];
+
+  // ---- Stage 2: scale t1 (keep in shmem for matmul, also write to global)
+  //                  scale b[0..NUM_RED_LOG] = scaled b1
+  for (int i = tid; i < NUM_SINGLE; i += BLOCK_DIM) {
+    float v = shared_t1[i] * inv_total;
+    shared_t1[i] = v;
+    t1[i] = v;
+  }
+  for (int i = tid; i < NUM_RED_LOG; i += BLOCK_DIM) {
+    b[i] = shared_b1[i] * inv_total;
+  }
+  __syncthreads();
+
+  // ---- Stage 3: b[NUM_RED_LOG + j] = -(B1[j] · t1) for j in [0, NUM_GPUS).
+  // Sequential GEMV across NUM_GPUS=16 outputs; each is a 240-wide dot
+  // product reduced across the block. Cheap at this size.
+  for (int j = 0; j < NUM_GPUS; j++) {
+    float dot = 0.f;
+    for (int k = tid; k < NUM_SINGLE; k += BLOCK_DIM) {
+      dot += B1[j * NUM_SINGLE + k] * shared_t1[k];
+    }
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      dot += __shfl_xor_sync(0xffffffff, dot, offset);
+    }
+    if (lane == 0) {
+      reduce_buf[warp_id] = dot;
+    }
+    __syncthreads();
+    if (warp_id == 0) {
+      float v = (tid < WARPS_PER_BLOCK) ? reduce_buf[tid] : 0.f;
+      for (int offset = 16; offset > 0; offset >>= 1) {
+        v += __shfl_xor_sync(0xffffffff, v, offset);
+      }
+      if (tid == 0) {
+        b[NUM_RED_LOG + j] = -v;
+      }
+    }
+    __syncthreads();
+  }
+
+  // ---- Stage 4: A_full[i][NV-1] = b[i] - A_base_row_sum[i].
+  // First NV-1 columns of A_full are pre-filled with A_base at solver init,
+  // so we only write the last column here.
+  for (int i = tid; i < NC; i += BLOCK_DIM) {
+    A_full[i * NV + (NV - 1)] = b[i] - A_base_row_sum[i];
+  }
+}
+
+template <int NC, int NV, int NUM_SINGLE, int NUM_RED_LOG, int NUM_GPUS, int BLOCK_DIM>
+void lp_prep(
+    tvm::ffi::TensorView A_full,
+    tvm::ffi::TensorView b,
+    tvm::ffi::TensorView t1,
+    tvm::ffi::TensorView global_counts,
+    tvm::ffi::TensorView log_single,
+    tvm::ffi::TensorView log_replicated,
+    tvm::ffi::TensorView B1,
+    tvm::ffi::TensorView A_base_row_sum) {
+  using namespace host;
+
+  SymbolicDevice device_;
+  TensorMatcher({NC, NV}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(A_full);
+  TensorMatcher({NC}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(b);
+  TensorMatcher({NUM_SINGLE}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(t1);
+  TensorMatcher({NUM_SINGLE}).with_dtype<int64_t>().with_device<kDLCUDA>(device_).verify(log_single);
+  TensorMatcher({NUM_RED_LOG}).with_dtype<int64_t>().with_device<kDLCUDA>(device_).verify(log_replicated);
+  TensorMatcher({NUM_GPUS, NUM_SINGLE}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(B1);
+  TensorMatcher({NC}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(A_base_row_sum);
+
+  constexpr int WARPS_PER_BLOCK = BLOCK_DIM / 32;
+  const size_t smem_bytes = (NUM_SINGLE + NUM_RED_LOG + WARPS_PER_BLOCK + 1) * sizeof(float);
+
+  using KernelT =
+      void (*)(float*, float*, float*, const float*, const int64_t*, const int64_t*, const float*, const float*);
+  KernelT kernel = lp_prep_kernel<NC, NV, NUM_SINGLE, NUM_RED_LOG, NUM_GPUS, BLOCK_DIM>;
+
+  if (smem_bytes > 48 * 1024) {
+    cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(smem_bytes));
+  }
+
+  const DLDevice device = device_.unwrap();
+  LaunchKernel(/*grid=*/1, /*block=*/BLOCK_DIM, device, smem_bytes)(
+      kernel,
+      static_cast<float*>(A_full.data_ptr()),
+      static_cast<float*>(b.data_ptr()),
+      static_cast<float*>(t1.data_ptr()),
+      static_cast<const float*>(global_counts.data_ptr()),
+      static_cast<const int64_t*>(log_single.data_ptr()),
+      static_cast<const int64_t*>(log_replicated.data_ptr()),
+      static_cast<const float*>(B1.data_ptr()),
+      static_cast<const float*>(A_base_row_sum.data_ptr()));
+}
+
+}  // namespace