feat(kernels): port standalone Kimi K3 kernels (#32890)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: zhangxiaolei <zhangxiaolei.666@bytedance.com>

python/sglang/kernels/jit/csrc/attention/fixup_zero_kv.cuh

@@ -17,37 +17,52 @@ constexpr int kFixupBlockSize = 256;
 
 // -- vectorised zero-fill helpers ------------------------------------------
 
-// Zero-fill `n` elements of type T starting at `ptr`, using float4 stores.
-// `ptr` must be 16-byte aligned (guaranteed by PyTorch allocator).
+// Zero-fill `n` elements of type T starting at `ptr`.
 template <typename T>
 __device__ __forceinline__ void vec_zero_fill(T* ptr, int n) {
-  constexpr int kVec = 16 / sizeof(T);  // elements per float4
-  const int n_vec = n / kVec;           // full vectors
-  float4* dst4 = reinterpret_cast<float4*>(ptr);
-  const float4 z4 = make_float4(0.f, 0.f, 0.f, 0.f);
-  for (int i = threadIdx.x; i < n_vec; i += blockDim.x) {
-    dst4[i] = z4;
-  }
-  // tail elements
-  const int tail_start = n_vec * kVec;
-  for (int i = tail_start + threadIdx.x; i < n; i += blockDim.x) {
-    ptr[i] = static_cast<T>(0);
+  if ((reinterpret_cast<uintptr_t>(ptr) & 0xF) == 0) {
+    // 16-byte aligned -> vectorised float4 stores
+    constexpr int kVec = 16 / sizeof(T);  // elements per float4
+    const int n_vec = n / kVec;           // full vectors
+    float4* dst4 = reinterpret_cast<float4*>(ptr);
+    const float4 z4 = make_float4(0.f, 0.f, 0.f, 0.f);
+    for (int i = threadIdx.x; i < n_vec; i += blockDim.x) {
+      dst4[i] = z4;
+    }
+    // tail elements
+    const int tail_start = n_vec * kVec;
+    for (int i = tail_start + threadIdx.x; i < n; i += blockDim.x) {
+      ptr[i] = static_cast<T>(0);
+    }
+  } else {
+    // misaligned row base -> scalar stores
+    for (int i = threadIdx.x; i < n; i += blockDim.x) {
+      ptr[i] = static_cast<T>(0);
+    }
   }
 }
 
-// Fill `n` float elements with -inf using float4 stores.
+// Fill `n` float elements with -inf.
 __device__ __forceinline__ void vec_neginf_fill(float* ptr, int n) {
-  constexpr int kVec = 4;  // float4 = 4 floats
-  const int n_vec = n / kVec;
-  float4* dst4 = reinterpret_cast<float4*>(ptr);
   const float ninf = -INFINITY;
-  const float4 inf4 = make_float4(ninf, ninf, ninf, ninf);
-  for (int i = threadIdx.x; i < n_vec; i += blockDim.x) {
-    dst4[i] = inf4;
-  }
-  const int tail_start = n_vec * kVec;
-  for (int i = tail_start + threadIdx.x; i < n; i += blockDim.x) {
-    ptr[i] = ninf;
+  if ((reinterpret_cast<uintptr_t>(ptr) & 0xF) == 0) {
+    // 16-byte aligned -> vectorised float4 stores
+    constexpr int kVec = 4;  // float4 = 4 floats
+    const int n_vec = n / kVec;
+    float4* dst4 = reinterpret_cast<float4*>(ptr);
+    const float4 inf4 = make_float4(ninf, ninf, ninf, ninf);
+    for (int i = threadIdx.x; i < n_vec; i += blockDim.x) {
+      dst4[i] = inf4;
+    }
+    const int tail_start = n_vec * kVec;
+    for (int i = tail_start + threadIdx.x; i < n; i += blockDim.x) {
+      ptr[i] = ninf;
+    }
+  } else {
+    // misaligned row base -> scalar stores
+    for (int i = threadIdx.x; i < n; i += blockDim.x) {
+      ptr[i] = ninf;
+    }
   }
 }
 

python/sglang/kernels/jit/csrc/attention/kda_packed_decode.cuh

@@ -0,0 +1,240 @@
+// CUDA port of the triton fused_recurrent_kda_packed_decode_kernel for
+// batched decode. The triton kernel holds a whole [BV, K] fp32 state tile in
+// the registers of a single warp, which caps it at ~5 TB/s of the ~9.6 TB/s
+// this in-place read+write stream can reach (probe: torch inplace mul_).
+// This kernel streams the state row by row instead: one warp per V-row group,
+// each row is a 512B float4 load -> warp-reduced dot -> decayed delta-rule
+// update -> 512B store, so loads pipeline across rows and nothing holds a
+// tile. Setup (l2norm'd q/k, per-K decay, beta) is computed redundantly per
+// warp - the kernel has no __syncthreads at all.
+//
+// Math follows the triton kernel exactly (fp32 throughout, same op order):
+//   h *= exp(g);  t = <h, k>;  v = (v - t) * sigmoid(b);  h += v * k;
+//   o = <h, q>
+// with g = -exp(A_log) * softplus(a + dt_bias) (K3: no lower bound) or
+// lower_bound * sigmoid(exp(A_log) * (a + dt_bias)). Warp-shuffle reduction
+// order differs from tl.sum, so outputs match to ULPs, not bits (validated
+// against the triton kernel with tolerance + GSM8K).
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck
+
+#include <sgl_kernel/type.cuh>   // For bf16_t, fp32_t, device::cast
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+struct KdaPackedDecodeParams {
+  const bf16_t* __restrict__ mixed_qkv;  // [B, 2*H*K + HV*V]
+  const bf16_t* __restrict__ a;          // [B, HV*K]
+  const bf16_t* __restrict__ b;          // [B, HV]
+  const fp32_t* __restrict__ A_log;      // [HV]
+  const fp32_t* __restrict__ dt_bias;    // [HV*K]
+  bf16_t* __restrict__ o;                // [B, HV*V] (contiguous view)
+  fp32_t* __restrict__ state;            // pool, row stride = stride_state
+  const int32_t* __restrict__ indices;   // [B]
+  int64_t stride_mixed;
+  int64_t stride_a;
+  int64_t stride_b;
+  int64_t stride_state;  // elements per pool slot
+  uint32_t H;
+  uint32_t HV;
+  fp32_t scale;
+  fp32_t lower_bound;
+  int32_t use_lower_bound;
+};
+
+__device__ __forceinline__ float warp_allreduce_sum(float v) {
+#if defined(__HIP_PLATFORM_AMD__)
+  constexpr uint64_t kFullMask = 0xffffffffffffffffull;
+#else
+  constexpr uint32_t kFullMask = 0xffffffffu;
+#endif
+#pragma unroll
+  for (int off = 16; off > 0; off >>= 1) {
+    v += __shfl_xor_sync(kFullMask, v, off);
+  }
+  return v;
+}
+
+// K = V = 128 specialization: one lane owns 4 consecutive K-elements (16B).
+template <int kWarps, bool kUsePDL>
+__global__
+__launch_bounds__(kWarps * 32) void kda_packed_decode_kernel(const KdaPackedDecodeParams __grid_constant__ params) {
+  using namespace device;
+  constexpr int K = 128;
+  constexpr int V = 128;
+  constexpr int kElems = 4;  // K / 32 lanes
+
+  const uint32_t i_nh = blockIdx.x;
+  const uint32_t n = i_nh / params.HV;
+  const uint32_t hv = i_nh % params.HV;
+  const uint32_t i_h = hv / (params.HV / params.H);
+  const uint32_t warp = threadIdx.x >> 5;
+  const uint32_t lane = threadIdx.x & 31;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  bf16_t* o_ptr = params.o + (static_cast<int64_t>(n) * params.HV + hv) * V;
+  const int64_t sidx = params.indices[n];
+  if (sidx < 0) {
+    // Padded cuda-graph slot: zero the output, leave the pool untouched.
+    for (uint32_t i = threadIdx.x; i < V; i += kWarps * 32) {
+      o_ptr[i] = cast<bf16_t>(0.0f);
+    }
+    PDLTriggerSecondary<kUsePDL>();
+    return;
+  }
+
+  // --- per-warp redundant setup (no cross-warp synchronization) ---
+  const bf16_t* mixed = params.mixed_qkv + n * params.stride_mixed;
+  const uint32_t e0 = lane * kElems;
+
+  float q[kElems], k[kElems];
+  float q_sq = 0.0f, k_sq = 0.0f;
+#pragma unroll
+  for (int e = 0; e < kElems; ++e) {
+    q[e] = cast<fp32_t>(mixed[i_h * K + e0 + e]);
+    k[e] = cast<fp32_t>(mixed[params.H * K + i_h * K + e0 + e]);
+    q_sq += q[e] * q[e];
+    k_sq += k[e] * k[e];
+  }
+  // tl: q / sqrt(sum(q*q) + 1e-6), then * scale
+  const float q_inv = 1.0f / sqrtf(warp_allreduce_sum(q_sq) + 1e-6f);
+  const float k_inv = 1.0f / sqrtf(warp_allreduce_sum(k_sq) + 1e-6f);
+#pragma unroll
+  for (int e = 0; e < kElems; ++e) {
+    q[e] = q[e] * q_inv * params.scale;
+    k[e] = k[e] * k_inv;
+  }
+
+  const float exp_A = expf(params.A_log[hv]);
+  float decay[kElems];
+#pragma unroll
+  for (int e = 0; e < kElems; ++e) {
+    const float x = cast<fp32_t>(params.a[n * params.stride_a + hv * K + e0 + e]) + params.dt_bias[hv * K + e0 + e];
+    float g;
+    if (params.use_lower_bound) {
+      g = params.lower_bound / (1.0f + expf(-exp_A * x));
+    } else {
+      const float sp = (x <= 20.0f) ? logf(1.0f + expf(x)) : x;
+      g = -exp_A * sp;
+    }
+    decay[e] = expf(g);
+  }
+  const float beta = 1.0f / (1.0f + expf(-cast<fp32_t>(params.b[n * params.stride_b + hv])));
+
+  const bf16_t* v_ptr = mixed + 2 * params.H * K + hv * V;
+  fp32_t* h_base = params.state + sidx * params.stride_state + static_cast<int64_t>(hv) * V * K;
+
+  // --- stream this warp's V-rows: 512B load -> update -> 512B store ---
+  constexpr int kRowsPerWarp = V / kWarps;
+#pragma unroll 4
+  for (int r = warp * kRowsPerWarp; r < (int)(warp + 1) * kRowsPerWarp; ++r) {
+    float4 h4 = *reinterpret_cast<const float4*>(h_base + r * K + e0);
+    float h[kElems] = {h4.x, h4.y, h4.z, h4.w};
+    float t = 0.0f;
+#pragma unroll
+    for (int e = 0; e < kElems; ++e) {
+      h[e] *= decay[e];
+      t += h[e] * k[e];
+    }
+    t = warp_allreduce_sum(t);
+    const float v_new = (cast<fp32_t>(v_ptr[r]) - t) * beta;
+    float o_acc = 0.0f;
+#pragma unroll
+    for (int e = 0; e < kElems; ++e) {
+      h[e] += v_new * k[e];
+      o_acc += h[e] * q[e];
+    }
+    o_acc = warp_allreduce_sum(o_acc);
+    *reinterpret_cast<float4*>(h_base + r * K + e0) = make_float4(h[0], h[1], h[2], h[3]);
+    if (lane == 0) {
+      o_ptr[r] = cast<bf16_t>(o_acc);
+    }
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int kWarps, bool kUsePDL>
+struct KdaPackedDecodeKernel {
+  static constexpr auto kernel = kda_packed_decode_kernel<kWarps, kUsePDL>;
+
+  static void
+  run(const tvm::ffi::TensorView mixed_qkv,
+      const tvm::ffi::TensorView a,
+      const tvm::ffi::TensorView b,
+      const tvm::ffi::TensorView A_log,
+      const tvm::ffi::TensorView dt_bias,
+      const tvm::ffi::TensorView o,
+      const tvm::ffi::TensorView state,
+      const tvm::ffi::TensorView indices,
+      double scale,
+      double lower_bound,
+      bool use_lower_bound,
+      int64_t num_q_heads) {
+    using namespace host;
+
+    auto B_ = SymbolicSize{"batch"};
+    auto MixedDim_ = SymbolicSize{"mixed_dim"};
+    auto ADim_ = SymbolicSize{"a_dim"};
+    auto HV_ = SymbolicSize{"num_v_heads"};
+    auto V_ = SymbolicSize{"head_v_dim"};
+    auto K_ = SymbolicSize{"head_k_dim"};
+    auto Slots_ = SymbolicSize{"pool_slots"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({B_, MixedDim_}).with_dtype<bf16_t>().with_device(device).with_strides({-1, 1}).verify(mixed_qkv);
+    TensorMatcher({B_, ADim_}).with_dtype<bf16_t>().with_device(device).with_strides({-1, 1}).verify(a);
+    TensorMatcher({B_, HV_}).with_dtype<bf16_t>().with_device(device).with_strides({-1, 1}).verify(b);
+    TensorMatcher({HV_}).with_dtype<fp32_t>().with_device(device).verify(A_log);
+    TensorMatcher({ADim_}).with_dtype<fp32_t>().with_device(device).verify(dt_bias);
+    TensorMatcher({B_, HV_, V_}).with_dtype<bf16_t>().with_device(device).verify(o);
+    TensorMatcher({Slots_, HV_, V_, K_})
+        .with_dtype<fp32_t>()
+        .with_device(device)
+        .with_strides({-1, -1, -1, 1})
+        .verify(state);
+    TensorMatcher({B_}).with_dtype<int32_t>().with_device(device).verify(indices);
+
+    const auto B = static_cast<uint32_t>(B_.unwrap());
+    const auto HV = static_cast<uint32_t>(HV_.unwrap());
+    const auto H = static_cast<uint32_t>(num_q_heads);
+    RuntimeCheck(K_.unwrap() == 128 && V_.unwrap() == 128, "kda_packed_decode is specialized for K = V = 128");
+    RuntimeCheck(
+        ADim_.unwrap() == HV * 128 && H > 0 && HV % H == 0, "a/dt_bias must be [*, HV*K] and HV divisible by H");
+    RuntimeCheck(MixedDim_.unwrap() == 2 * H * 128 + HV * 128, "mixed_qkv last dim must be 2*H*K + HV*V");
+    RuntimeCheck(state.stride(1) == 128 * 128 && state.stride(2) == 128, "state inner layout must be dense [HV, V, K]");
+    if (B == 0) return;
+
+    const auto params = KdaPackedDecodeParams{
+        .mixed_qkv = static_cast<const bf16_t*>(mixed_qkv.data_ptr()),
+        .a = static_cast<const bf16_t*>(a.data_ptr()),
+        .b = static_cast<const bf16_t*>(b.data_ptr()),
+        .A_log = static_cast<const fp32_t*>(A_log.data_ptr()),
+        .dt_bias = static_cast<const fp32_t*>(dt_bias.data_ptr()),
+        .o = static_cast<bf16_t*>(o.data_ptr()),
+        .state = static_cast<fp32_t*>(state.data_ptr()),
+        .indices = static_cast<const int32_t*>(indices.data_ptr()),
+        .stride_mixed = mixed_qkv.stride(0),
+        .stride_a = a.stride(0),
+        .stride_b = b.stride(0),
+        .stride_state = state.stride(0),
+        .H = H,
+        .HV = HV,
+        .scale = static_cast<fp32_t>(scale),
+        .lower_bound = static_cast<fp32_t>(lower_bound),
+        .use_lower_bound = use_lower_bound ? 1 : 0,
+    };
+
+    LaunchKernel(B * HV, kWarps * 32, device.unwrap()).enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/kernels/jit/csrc/elementwise/add3.cuh

@@ -0,0 +1,93 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace sglang {
+
+struct Add3Params {
+  const bf16_t* __restrict__ a;  // [N] contiguous
+  const bf16_t* __restrict__ b;  // [N] contiguous
+  const bf16_t* __restrict__ c;  // [N] contiguous
+  bf16_t* __restrict__ out;      // [N] contiguous
+  int64_t n_vecs;                // N / kVecElems
+};
+
+template <bool kUsePDL, bool kPrefetchBC>
+__global__ void add3_kernel(const __grid_constant__ Add3Params params) {
+  constexpr uint32_t kVecPairs = device::kMaxVecBytes / sizeof(bf16x2_t);
+  using vec_t = device::AlignedVector<bf16x2_t, kVecPairs>;
+
+  const int64_t vid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (vid >= params.n_vecs) return;
+
+  // Trigger early, so that the next kernel gets a chance to prefetch.
+  device::PDLTriggerSecondary<kUsePDL>();
+
+  vec_t a, b, c;
+  if constexpr (kPrefetchBC) {
+    b.load(params.b, vid);
+    c.load(params.c, vid);
+    device::PDLWaitPrimary<kUsePDL>();
+    a.load(params.a, vid);
+  } else {
+    device::PDLWaitPrimary<kUsePDL>();
+    a.load(params.a, vid);
+    b.load(params.b, vid);
+    c.load(params.c, vid);
+  }
+
+  vec_t out;
+#pragma unroll
+  for (uint32_t i = 0; i < kVecPairs; ++i) {
+    out[i] = __hadd2(__hadd2(a[i], b[i]), c[i]);
+  }
+  out.store(params.out, vid);
+}
+
+template <bool kUsePDL>
+struct Add3Kernel {
+  static constexpr int64_t kVecElems = device::kMaxVecBytes / sizeof(bf16_t);
+
+  static void launch(
+      const tvm::ffi::TensorView a,
+      const tvm::ffi::TensorView b,
+      const tvm::ffi::TensorView c,
+      const tvm::ffi::TensorView out,
+      const bool prefetch_bc) {
+    using namespace host;
+
+    auto N = SymbolicSize{"numel"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N}).with_dtype<bf16_t>().with_device(device).verify(a).verify(b).verify(c).verify(out);
+    const auto numel = N.unwrap();
+    RuntimeCheck(numel % kVecElems == 0, "numel must be divisible by the vector width");
+    if (numel == 0) return;
+    const auto params = Add3Params{
+        .a = static_cast<const bf16_t*>(a.data_ptr()),
+        .b = static_cast<const bf16_t*>(b.data_ptr()),
+        .c = static_cast<const bf16_t*>(c.data_ptr()),
+        .out = static_cast<bf16_t*>(out.data_ptr()),
+        .n_vecs = numel / kVecElems,
+    };
+    const auto num_threads = [&]() -> int64_t {
+      for (int64_t n : {128, 256, 512}) {
+        if (params.n_vecs <= n * 256) return n;
+      }
+      return 512;
+    }();
+    const auto grid = div_ceil(params.n_vecs, num_threads);
+    const auto kernel = prefetch_bc ? add3_kernel<kUsePDL, true> : add3_kernel<kUsePDL, false>;
+    LaunchKernel(grid, num_threads, device.unwrap()).enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace sglang

python/sglang/kernels/jit/csrc/elementwise/concat_mla.cuh

@@ -186,6 +186,7 @@ constexpr int A_LAST_DIM = 512;
 constexpr int B_LAST_DIM = 64;
 constexpr int OUT_LAST_DIM = A_LAST_DIM + B_LAST_DIM;
 
+template <bool kUsePDL>
 __global__ void concat_mla_absorb_q_kernel(
     bf16_t* a,
     bf16_t* b,
@@ -198,6 +199,8 @@ __global__ void concat_mla_absorb_q_kernel(
     const int b_stride_1,
     const int64_t out_stride_0,
     const int out_stride_1) {
+  device::PDLWaitPrimary<kUsePDL>();
+
   const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
   const int lane_id = get_lane_id();
 
@@ -229,6 +232,8 @@ __global__ void concat_mla_absorb_q_kernel(
     a_buf[i] = *(base_addr + i * 32 + lane_id);
   }
 
+  device::PDLTriggerSecondary<kUsePDL>();
+
   {
     BBufType* base_addr = reinterpret_cast<BBufType*>(out + idx_0 * out_stride_0 + idx_1 * out_stride_1 + A_LAST_DIM);
     *(base_addr + lane_id) = b_buf;
@@ -241,6 +246,7 @@ __global__ void concat_mla_absorb_q_kernel(
   }
 }
 
+template <bool kUsePDL>
 struct ConcatMlaAbsorbQKernel {
   static void run(tvm::ffi::TensorView a, tvm::ffi::TensorView b, tvm::ffi::TensorView out) {
     using namespace host;
@@ -306,19 +312,20 @@ struct ConcatMlaAbsorbQKernel {
     const int grid_size = div_ceil(num_items, num_warps_per_block);
     const int block_size = num_warps_per_block * 32;
 
-    LaunchKernel(grid_size, block_size, device.unwrap())(
-        concat_mla_absorb_q_kernel,
-        static_cast<bf16_t*>(a.data_ptr()),
-        static_cast<bf16_t*>(b.data_ptr()),
-        static_cast<bf16_t*>(out.data_ptr()),
-        num_items,
-        dim_1,
-        S0_a.unwrap(),
-        static_cast<int>(S1_a.unwrap()),
-        S0_b.unwrap(),
-        static_cast<int>(S1_b.unwrap()),
-        S0_out.unwrap(),
-        static_cast<int>(S1_out.unwrap()));
+    LaunchKernel(grid_size, block_size, device.unwrap())
+        .enable_pdl(kUsePDL)(
+            concat_mla_absorb_q_kernel<kUsePDL>,
+            static_cast<bf16_t*>(a.data_ptr()),
+            static_cast<bf16_t*>(b.data_ptr()),
+            static_cast<bf16_t*>(out.data_ptr()),
+            num_items,
+            dim_1,
+            S0_a.unwrap(),
+            static_cast<int>(S1_a.unwrap()),
+            S0_b.unwrap(),
+            static_cast<int>(S1_b.unwrap()),
+            S0_out.unwrap(),
+            static_cast<int>(S1_out.unwrap()));
   }
 };
 

python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_buffer.cuh

@@ -27,6 +27,7 @@
 #include <sgl_kernel/tile.cuh>
 #include <sgl_kernel/utils.cuh>
 #include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
 
 #include <cuda/ptx>
 #include <dlpack/dlpack.h>
@@ -47,40 +48,6 @@ struct SetMlaKVBufferParams {
   uint32_t batch_size;
 };
 
-// Warp-cooperative gmem -> smem copy. Picks the widest vec width that divides
-// both the per-thread share and the byte total. Caller guarantees src is
-// 16-byte aligned (PyTorch tensors are) and dst is the start of a per-warp
-// smem slot (also 16-byte aligned by ``alignas(16)``).
-template <int64_t kBytes>
-SGL_DEVICE void warp_g2s_copy(const void* __restrict__ src, void* __restrict__ dst) {
-  using namespace device;
-  constexpr int64_t kAlignment = (kBytes % (16 * kWarpThreads) == 0)  ? 16
-                                 : (kBytes % (8 * kWarpThreads) == 0) ? 8
-                                 : (kBytes % (4 * kWarpThreads) == 0) ? 4
-                                 : (kBytes % 4 == 0)                  ? 4
-                                                                      : 0;
-  static_assert(kAlignment > 0, "kBytes must be a multiple of 4");
-
-  using vec_t = AlignedStorage<uint32_t, kAlignment / 4>;
-  constexpr auto kLoopBytes = sizeof(vec_t) * kWarpThreads;
-  constexpr auto kLoopCount = kBytes / kLoopBytes;
-  constexpr int64_t kTailVecs = (kBytes - kLoopCount * kLoopBytes) / sizeof(vec_t);
-
-  const auto gmem = tile::Memory<vec_t>::warp();
-
-#pragma unroll
-  for (int64_t i = 0; i < kLoopCount; ++i) {
-    const auto v = gmem.load(src, i);
-    gmem.store(dst, v, i);
-  }
-  if constexpr (kTailVecs > 0) {
-    if (gmem.in_bound(kLoopCount * kWarpThreads + kTailVecs, kLoopCount)) {
-      const auto v = gmem.load(src, kLoopCount);
-      gmem.store(dst, v, kLoopCount);
-    }
-  }
-}
-
 template <int64_t kNopeBytes, int64_t kRopeBytes, int kNumWarps, bool kUsePDL, typename TLoc>
 __global__ void set_mla_kv_buffer_kernel(const __grid_constant__ SetMlaKVBufferParams params) {
   using namespace device;
@@ -104,8 +71,8 @@ __global__ void set_mla_kv_buffer_kernel(const __grid_constant__ SetMlaKVBufferP
   void* const gmem_dst = pointer::offset(params.kv_buffer, loc * params.stride_buffer_bytes);
 
   // Warp-cooperative load (nope, rope) into the per-warp smem slot.
-  warp_g2s_copy<kNopeBytes>(nope_src, &smem[warp_in_cta][0]);
-  warp_g2s_copy<kRopeBytes>(rope_src, &smem[warp_in_cta][kNopeBytes]);
+  warp::copy_bytes<kNopeBytes>(nope_src, &smem[warp_in_cta][0]);
+  warp::copy_bytes<kRopeBytes>(rope_src, &smem[warp_in_cta][kNopeBytes]);
 
   // Fence required: TMA reads smem via the async proxy, normal sts writes
   // through the generic proxy. Without this the TMA engine can observe stale

python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_concat_q.cuh

@@ -0,0 +1,607 @@
+// MLA KV-cache write fused with the Q concat, bf16 and fp8 entry points.
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil
+
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE, PDL helpers
+#include <sgl_kernel/vec.cuh>    // For AlignedVector
+#include <sgl_kernel/warp.cuh>   // For warp::copy_bytes, elect_one_lane, inclusive_sum
+
+#include <cuda/ptx>
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+#include <cuda_fp8.h>
+
+namespace {
+
+struct SetMlaKVConcatQParams {
+  // KV scatter side (byte-typed: dtype-agnostic row copies).
+  const void* __restrict__ k_nope;
+  const void* __restrict__ k_rope;
+  void* __restrict__ kv_buffer;
+  const void* __restrict__ loc;
+  int64_t stride_nope_bytes;
+  int64_t stride_rope_bytes;
+  int64_t stride_buffer_bytes;
+  uint32_t batch_size;
+  // Q concat side (bf16, element strides).
+  const bf16_t* __restrict__ q_nope;
+  const bf16_t* __restrict__ q_rope;
+  bf16_t* __restrict__ q_out;
+  uint32_t num_q_items;  // batch_size * num_heads
+  uint32_t q_dim_1;      // num_heads
+  int64_t qn_stride_0;
+  int32_t qn_stride_1;
+  int64_t qr_stride_0;
+  int32_t qr_stride_1;
+  int64_t qo_stride_0;
+  int32_t qo_stride_1;
+};
+
+template <int64_t kNopeBytes, int64_t kRopeBytes, int kNumWarps, bool kUsePDL, typename TLoc>
+__global__ void set_mla_kv_concat_q_kernel(const __grid_constant__ SetMlaKVConcatQParams params) {
+  using namespace device;
+  static_assert((kNopeBytes + kRopeBytes) % 16 == 0, "TMA bulk store requires total row to be 16-byte aligned");
+
+  constexpr int64_t kRowBytes = kNopeBytes + kRopeBytes;
+  constexpr int kQNopeDim = static_cast<int>(kNopeBytes / sizeof(bf16_t));
+  constexpr int kQRopeDim = static_cast<int>(kRopeBytes / sizeof(bf16_t));
+
+  // Per-warp smem slots for the KV scatter role; concat warps leave theirs idle.
+  __shared__ alignas(16) uint8_t smem[kNumWarps][kRowBytes];
+
+  const uint32_t warp_in_cta = threadIdx.x / kWarpThreads;
+  const uint32_t lane_id = threadIdx.x % kWarpThreads;
+  const uint32_t flat_warp = blockIdx.x * kNumWarps + warp_in_cta;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  if (flat_warp < params.batch_size) {
+    // --- KV scatter role: one warp per token (smem staging + TMA bulk store) ---
+    const uint32_t item_id = flat_warp;
+    const int64_t loc = static_cast<int64_t>(static_cast<const TLoc*>(params.loc)[item_id]);
+
+    const auto nope_src = pointer::offset(params.k_nope, item_id * params.stride_nope_bytes);
+    const auto rope_src = pointer::offset(params.k_rope, item_id * params.stride_rope_bytes);
+    void* const gmem_dst = pointer::offset(params.kv_buffer, loc * params.stride_buffer_bytes);
+
+    warp::copy_bytes<kNopeBytes>(nope_src, &smem[warp_in_cta][0]);
+    warp::copy_bytes<kRopeBytes>(rope_src, &smem[warp_in_cta][kNopeBytes]);
+
+    // TMA reads smem via the async proxy; fence so it can't observe stale sts.
+    __syncwarp();
+    asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
+
+    // elect.sync rather than `lane_id == 0`: the TMA issue must not sit
+    // behind a lane-index predicate (see PR review).
+    if (device::warp::elect_one_lane()) {
+      cuda::ptx::cp_async_bulk(
+          cuda::ptx::space_global,
+          cuda::ptx::space_shared,
+          gmem_dst,
+          &smem[warp_in_cta][0],
+          static_cast<uint32_t>(kRowBytes));
+    }
+
+    // ``wait_group`` (not ``_read``): waits for gmem commit, not just smem reuse.
+    cuda::ptx::cp_async_bulk_commit_group();
+    cuda::ptx::cp_async_bulk_wait_group(cuda::ptx::n32_t<0>{});
+  } else if (flat_warp - params.batch_size < params.num_q_items) {
+    // --- Q concat role: one warp per (token, head) row ---
+    const uint32_t q_item = flat_warp - params.batch_size;
+    const uint32_t idx_0 = q_item / params.q_dim_1;
+    const uint32_t idx_1 = q_item % params.q_dim_1;
+
+    using ABufType = int4;
+    constexpr int kAVecElems = static_cast<int>(sizeof(ABufType) / sizeof(bf16_t));
+    constexpr int kANumUnroll = kQNopeDim / (kAVecElems * kWarpThreads);
+    static_assert(kANumUnroll * kAVecElems * kWarpThreads == kQNopeDim, "nope dim must fill whole int4 warp rounds");
+    using BBufType = int;
+    constexpr int kBVecElems = static_cast<int>(sizeof(BBufType) / sizeof(bf16_t));
+    static_assert(kBVecElems * kWarpThreads == kQRopeDim, "rope dim must be exactly one int warp round");
+
+    const bf16_t* a_row = params.q_nope + idx_0 * params.qn_stride_0 + idx_1 * params.qn_stride_1;
+    const bf16_t* b_row = params.q_rope + idx_0 * params.qr_stride_0 + idx_1 * params.qr_stride_1;
+    bf16_t* o_row = params.q_out + idx_0 * params.qo_stride_0 + idx_1 * params.qo_stride_1;
+
+    ABufType a_buf[kANumUnroll];
+#pragma unroll
+    for (int i = 0; i < kANumUnroll; ++i) {
+      a_buf[i] = reinterpret_cast<const ABufType*>(a_row)[i * kWarpThreads + lane_id];
+    }
+    const BBufType b_buf = reinterpret_cast<const BBufType*>(b_row)[lane_id];
+
+#pragma unroll
+    for (int i = 0; i < kANumUnroll; ++i) {
+      reinterpret_cast<ABufType*>(o_row)[i * kWarpThreads + lane_id] = a_buf[i];
+    }
+    reinterpret_cast<BBufType*>(o_row + kQNopeDim)[lane_id] = b_buf;
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kNopeBytes, int64_t kRopeBytes, bool kUsePDL>
+struct SetMlaKVConcatQKernel {
+  static_assert(kNopeBytes > 0 && kNopeBytes % 4 == 0, "kNopeBytes must be a positive multiple of 4");
+  static_assert(kRopeBytes > 0 && kRopeBytes % 4 == 0, "kRopeBytes must be a positive multiple of 4");
+  static_assert(
+      (kNopeBytes + kRopeBytes) % 16 == 0, "TMA bulk store requires (kNopeBytes + kRopeBytes) to be a multiple of 16");
+
+  static constexpr int64_t kQNopeDim = kNopeBytes / static_cast<int64_t>(sizeof(bf16_t));
+  static constexpr int64_t kQRopeDim = kRopeBytes / static_cast<int64_t>(sizeof(bf16_t));
+
+  template <int kNumWarps, typename TLoc>
+  static constexpr auto kernel = set_mla_kv_concat_q_kernel<kNopeBytes, kRopeBytes, kNumWarps, kUsePDL, TLoc>;
+
+  static void
+  run(tvm::ffi::TensorView kv_buffer,
+      tvm::ffi::TensorView loc,
+      tvm::ffi::TensorView k_nope,
+      tvm::ffi::TensorView k_rope,
+      tvm::ffi::TensorView q_nope,
+      tvm::ffi::TensorView q_rope,
+      tvm::ffi::TensorView q_out,
+      int64_t num_warps_per_block) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto H = SymbolicSize{"num_heads"};
+    auto D_nope = SymbolicSize{"nope_dim"};
+    auto D_rope = SymbolicSize{"rope_dim"};
+    auto D_buf = SymbolicSize{"buffer_last_dim"};
+    auto D_qn = SymbolicSize{"q_nope_dim"};
+    auto D_qr = SymbolicSize{"q_rope_dim"};
+    auto D_qo = SymbolicSize{"q_out_dim"};
+    auto S_nope = SymbolicSize{"nope_stride"};
+    auto S_rope = SymbolicSize{"rope_stride"};
+    auto S_buf = SymbolicSize{"buffer_stride"};
+    auto S_loc = SymbolicSize{"loc_stride"};
+    auto S0_qn = SymbolicSize{"q_nope_stride_0"};
+    auto S1_qn = SymbolicSize{"q_nope_stride_1"};
+    auto S0_qr = SymbolicSize{"q_rope_stride_0"};
+    auto S1_qr = SymbolicSize{"q_rope_stride_1"};
+    auto S0_qo = SymbolicSize{"q_out_stride_0"};
+    auto S1_qo = SymbolicSize{"q_out_stride_1"};
+    auto loc_dtype = SymbolicDType{};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    D_qn.set_value(kQNopeDim);
+    D_qr.set_value(kQRopeDim);
+    D_qo.set_value(kQNopeDim + kQRopeDim);
+
+    TensorMatcher({B, D_nope})  //
+        .with_strides({S_nope, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(k_nope);
+    TensorMatcher({B, D_rope})  //
+        .with_strides({S_rope, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(k_rope);
+    TensorMatcher({-1, D_buf})  //
+        .with_strides({S_buf, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(kv_buffer);
+    TensorMatcher({B})  //
+        .with_strides({S_loc})
+        .with_dtype<int32_t, int64_t>(loc_dtype)
+        .with_device(device)
+        .verify(loc);
+    TensorMatcher({B, H, D_qn})  //
+        .with_strides({S0_qn, S1_qn, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(q_nope);
+    TensorMatcher({B, H, D_qr})  //
+        .with_strides({S0_qr, S1_qr, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(q_rope);
+    TensorMatcher({B, H, D_qo})  //
+        .with_strides({S0_qo, S1_qo, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(q_out);
+
+    constexpr int64_t kDtypeSize = static_cast<int64_t>(sizeof(bf16_t));
+    CHECK_HOST(kNopeBytes == kDtypeSize * D_nope.unwrap())
+        << "kNopeBytes mismatch: expected " << kNopeBytes << ", got " << kDtypeSize * D_nope.unwrap();
+    CHECK_HOST(kRopeBytes == kDtypeSize * D_rope.unwrap())
+        << "kRopeBytes mismatch: expected " << kRopeBytes << ", got " << kDtypeSize * D_rope.unwrap();
+    CHECK_HOST(kDtypeSize * D_buf.unwrap() >= kNopeBytes + kRopeBytes) << "kv_buffer last dim too small";
+    CHECK_HOST(S_loc.unwrap() == 1) << "loc must be contiguous; got stride " << S_loc.unwrap();
+
+    // Alignment tripwires. The device code does 16-byte vector accesses on the
+    // kv row / nope rows / q rows and 4-byte accesses on the rope rows; the
+    // python-side ``covered()`` mirrors these so uncovered layouts fall back
+    // instead of faulting (do NOT assume "PyTorch tensors are aligned" — views
+    // and odd pool pitches break that).
+    const auto aligned = [](const void* ptr, int64_t align) {
+      return reinterpret_cast<uintptr_t>(ptr) % static_cast<uintptr_t>(align) == 0;
+    };
+    CHECK_HOST(aligned(kv_buffer.data_ptr(), 16) && (S_buf.unwrap() * kDtypeSize) % 16 == 0)
+        << "kv_buffer base/row-stride must be 16-byte aligned for TMA bulk store";
+    CHECK_HOST(aligned(k_nope.data_ptr(), 16) && (S_nope.unwrap() * kDtypeSize) % 16 == 0)
+        << "k_nope base/row-stride must be 16-byte aligned";
+    CHECK_HOST(aligned(k_rope.data_ptr(), 4) && (S_rope.unwrap() * kDtypeSize) % 4 == 0)
+        << "k_rope base/row-stride must be 4-byte aligned";
+    CHECK_HOST(
+        aligned(q_nope.data_ptr(), 16) && (S0_qn.unwrap() * kDtypeSize) % 16 == 0 &&
+        (S1_qn.unwrap() * kDtypeSize) % 16 == 0)
+        << "q_nope base/strides must be 16-byte aligned";
+    CHECK_HOST(
+        aligned(q_rope.data_ptr(), 4) && (S0_qr.unwrap() * kDtypeSize) % 4 == 0 &&
+        (S1_qr.unwrap() * kDtypeSize) % 4 == 0)
+        << "q_rope base/strides must be 4-byte aligned";
+    CHECK_HOST(
+        aligned(q_out.data_ptr(), 16) && (S0_qo.unwrap() * kDtypeSize) % 16 == 0 &&
+        (S1_qo.unwrap() * kDtypeSize) % 16 == 0)
+        << "q_out base/strides must be 16-byte aligned";
+
+    const uint32_t batch = static_cast<uint32_t>(B.unwrap());
+    const uint32_t num_heads = static_cast<uint32_t>(H.unwrap());
+    if (batch == 0) return;
+
+    const auto params = SetMlaKVConcatQParams{
+        .k_nope = k_nope.data_ptr(),
+        .k_rope = k_rope.data_ptr(),
+        .kv_buffer = kv_buffer.data_ptr(),
+        .loc = loc.data_ptr(),
+        .stride_nope_bytes = S_nope.unwrap() * kDtypeSize,
+        .stride_rope_bytes = S_rope.unwrap() * kDtypeSize,
+        .stride_buffer_bytes = S_buf.unwrap() * kDtypeSize,
+        .batch_size = batch,
+        .q_nope = static_cast<const bf16_t*>(q_nope.data_ptr()),
+        .q_rope = static_cast<const bf16_t*>(q_rope.data_ptr()),
+        .q_out = static_cast<bf16_t*>(q_out.data_ptr()),
+        .num_q_items = batch * num_heads,
+        .q_dim_1 = num_heads,
+        .qn_stride_0 = S0_qn.unwrap(),
+        .qn_stride_1 = static_cast<int32_t>(S1_qn.unwrap()),
+        .qr_stride_0 = S0_qr.unwrap(),
+        .qr_stride_1 = static_cast<int32_t>(S1_qr.unwrap()),
+        .qo_stride_0 = S0_qo.unwrap(),
+        .qo_stride_1 = static_cast<int32_t>(S1_qo.unwrap()),
+    };
+
+    const auto use_int32 = loc_dtype.is_type<int32_t>();
+    const uint32_t total_warps = params.batch_size + params.num_q_items;
+
+    auto launch = [&]<int kNW>() {
+      const auto kernel_ptr = use_int32 ? kernel<kNW, int32_t> : kernel<kNW, int64_t>;
+      const uint32_t num_blocks = div_ceil(total_warps, static_cast<uint32_t>(kNW));
+      const uint32_t threads_per_block = static_cast<uint32_t>(kNW) * device::kWarpThreads;
+      LaunchKernel(num_blocks, threads_per_block, device.unwrap())  //
+          .enable_pdl(kUsePDL)(kernel_ptr, params);
+    };
+
+    switch (num_warps_per_block) {
+      case 1:
+        launch.template operator()<1>();
+        break;
+      case 2:
+        launch.template operator()<2>();
+        break;
+      case 4:
+        launch.template operator()<4>();
+        break;
+      case 8:
+        launch.template operator()<8>();
+        break;
+      default:
+        Panic("Unsupported num_warps_per_block=", num_warps_per_block);
+    }
+  }
+};
+
+// ---------------------------------------------------------------------------
+// fp8 variant. Shares the translation unit, not the kernel: dims are runtime
+// rather than template parameters, it shards DCP slots (vloc % world != rank),
+// converts per lane instead of bulk-copying, and counts strides in elements.
+// Only the module that instantiates it pays for it.
+// ---------------------------------------------------------------------------
+constexpr int kFp8NopeDim = 512;
+constexpr int kFp8RopeDim = 64;
+constexpr int kFp8RowBytes = kFp8NopeDim + kFp8RopeDim;  // fp8: 1 byte/elem
+
+struct SetMlaKVConcatQFp8Params {
+  // KV quantize + scatter side.
+  const bf16_t* __restrict__ k_nope;
+  const bf16_t* __restrict__ k_rope;
+  uint8_t* __restrict__ kv_buffer;
+  const void* __restrict__ loc;
+  int64_t stride_nope;          // elements
+  int64_t stride_rope;          // elements
+  int64_t stride_buffer_bytes;  // bytes
+  uint32_t batch_size;
+  // DCP cyclic sharding of the KV pool: ``loc`` is VIRTUAL; the physical
+  // row on the owner rank is loc / world, and only the owner
+  // (loc % world == rank) writes. world=1/rank=0 = identity (non-DCP).
+  int32_t dcp_world_size;
+  int32_t dcp_rank;
+  // Q quantize + concat side.
+  const bf16_t* __restrict__ q_nope;
+  const bf16_t* __restrict__ q_rope;
+  uint8_t* __restrict__ q_out;
+  uint32_t num_q_items;  // batch_size * num_heads
+  uint32_t q_dim_1;      // num_heads
+  int64_t qn_stride_0;
+  int32_t qn_stride_1;
+  int64_t qr_stride_0;
+  int32_t qr_stride_1;
+  int64_t qo_stride_0;  // elements (== bytes for fp8)
+  int32_t qo_stride_1;
+};
+
+// 2x bf16 -> 2x fp8 e4m3, float-mediated cvt.rn NOSAT (matches aten: overflow -> NaN).
+SGL_DEVICE uint16_t bf16x2_to_fp8x2(const bf16x2_t v) {
+  const float2 f = __bfloat1622float2(v);
+  return __nv_cvt_float2_to_fp8x2(f, __NV_NOSAT, __NV_E4M3);
+}
+
+// Convert 8 bf16 (one int4 load) to 8 fp8 packed in a uint2.
+SGL_DEVICE uint2 bf16x8_to_fp8x8(const int4 v) {
+  const bf16x2_t* p = reinterpret_cast<const bf16x2_t*>(&v);
+  uint2 out;
+  out.x = static_cast<uint32_t>(bf16x2_to_fp8x2(p[0])) | (static_cast<uint32_t>(bf16x2_to_fp8x2(p[1])) << 16);
+  out.y = static_cast<uint32_t>(bf16x2_to_fp8x2(p[2])) | (static_cast<uint32_t>(bf16x2_to_fp8x2(p[3])) << 16);
+  return out;
+}
+
+template <int kNumWarps, bool kUsePDL, typename TLoc>
+__global__ void set_mla_kv_concat_q_fp8_kernel(const __grid_constant__ SetMlaKVConcatQFp8Params params) {
+  using namespace device;
+
+  // Per-warp smem slots for the KV role (fp8 rows); concat warps leave
+  // theirs idle. 576 % 16 == 0 satisfies the TMA bulk-store requirement.
+  __shared__ alignas(16) uint8_t smem[kNumWarps][kFp8RowBytes];
+
+  const uint32_t warp_in_cta = threadIdx.x / kWarpThreads;
+  const uint32_t lane_id = threadIdx.x % kWarpThreads;
+  const uint32_t flat_warp = blockIdx.x * kNumWarps + warp_in_cta;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  if (flat_warp < params.batch_size) {
+    // --- KV role: quantize one token's row into smem, TMA-scatter it ---
+    const uint32_t item_id = flat_warp;
+    const int64_t vloc = static_cast<int64_t>(static_cast<const TLoc*>(params.loc)[item_id]);
+    // DCP ownership: non-owner ranks write nothing for this token (mirrors
+    // the triton writer's is_valid mask + loc // world translation).
+    if (vloc % params.dcp_world_size != params.dcp_rank) {
+      PDLTriggerSecondary<kUsePDL>();
+      return;
+    }
+    const int64_t loc = vloc / params.dcp_world_size;
+    const bf16_t* nope_src = params.k_nope + item_id * params.stride_nope;
+    const bf16_t* rope_src = params.k_rope + item_id * params.stride_rope;
+
+    // nope: 512 bf16 -> 512 fp8; 16 elems/lane (2 int4 loads -> 1 int4 store).
+    {
+      const int4* src = reinterpret_cast<const int4*>(nope_src);
+      uint2 lo = bf16x8_to_fp8x8(src[lane_id * 2]);
+      uint2 hi = bf16x8_to_fp8x8(src[lane_id * 2 + 1]);
+      reinterpret_cast<int4*>(&smem[warp_in_cta][0])[lane_id] =
+          make_int4(static_cast<int>(lo.x), static_cast<int>(lo.y), static_cast<int>(hi.x), static_cast<int>(hi.y));
+    }
+    // rope: 64 bf16 -> 64 fp8; 2 elems/lane.
+    {
+      const bf16x2_t v = reinterpret_cast<const bf16x2_t*>(rope_src)[lane_id];
+      reinterpret_cast<uint16_t*>(&smem[warp_in_cta][kFp8NopeDim])[lane_id] = bf16x2_to_fp8x2(v);
+    }
+
+    // TMA reads smem via the async proxy; fence so it can't observe stale sts.
+    __syncwarp();
+    asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
+
+    // elect.sync rather than `lane_id == 0`: the TMA issue must not sit
+    // behind a lane-index predicate (same review point as the bf16 variant).
+    if (device::warp::elect_one_lane()) {
+      cuda::ptx::cp_async_bulk(
+          cuda::ptx::space_global,
+          cuda::ptx::space_shared,
+          params.kv_buffer + loc * params.stride_buffer_bytes,
+          &smem[warp_in_cta][0],
+          static_cast<uint32_t>(kFp8RowBytes));
+    }
+    // ``wait_group`` (not ``_read``): waits for gmem commit, not just smem reuse.
+    cuda::ptx::cp_async_bulk_commit_group();
+    cuda::ptx::cp_async_bulk_wait_group(cuda::ptx::n32_t<0>{});
+  } else if (flat_warp - params.batch_size < params.num_q_items) {
+    // --- Q role: quantize one (token, head) row into the fp8 query ---
+    const uint32_t q_item = flat_warp - params.batch_size;
+    const uint32_t idx_0 = q_item / params.q_dim_1;
+    const uint32_t idx_1 = q_item % params.q_dim_1;
+    const bf16_t* a_row = params.q_nope + idx_0 * params.qn_stride_0 + idx_1 * params.qn_stride_1;
+    const bf16_t* b_row = params.q_rope + idx_0 * params.qr_stride_0 + idx_1 * params.qr_stride_1;
+    uint8_t* o_row = params.q_out + idx_0 * params.qo_stride_0 + idx_1 * params.qo_stride_1;
+
+    {
+      const int4* src = reinterpret_cast<const int4*>(a_row);
+      uint2 lo = bf16x8_to_fp8x8(src[lane_id * 2]);
+      uint2 hi = bf16x8_to_fp8x8(src[lane_id * 2 + 1]);
+      reinterpret_cast<int4*>(o_row)[lane_id] =
+          make_int4(static_cast<int>(lo.x), static_cast<int>(lo.y), static_cast<int>(hi.x), static_cast<int>(hi.y));
+    }
+    {
+      const bf16x2_t v = reinterpret_cast<const bf16x2_t*>(b_row)[lane_id];
+      reinterpret_cast<uint16_t*>(o_row + kFp8NopeDim)[lane_id] = bf16x2_to_fp8x2(v);
+    }
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <bool kUsePDL>
+struct SetMlaKVConcatQFp8Kernel {
+  template <int kNumWarps, typename TLoc>
+  static constexpr auto kernel = set_mla_kv_concat_q_fp8_kernel<kNumWarps, kUsePDL, TLoc>;
+
+  static void
+  run(tvm::ffi::TensorView kv_buffer,
+      tvm::ffi::TensorView loc,
+      tvm::ffi::TensorView k_nope,
+      tvm::ffi::TensorView k_rope,
+      tvm::ffi::TensorView q_nope,
+      tvm::ffi::TensorView q_rope,
+      tvm::ffi::TensorView q_out,
+      int64_t num_warps_per_block,
+      int64_t dcp_world_size,
+      int64_t dcp_rank) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto H = SymbolicSize{"num_heads"};
+    auto D_nope = SymbolicSize{"nope_dim"};
+    auto D_rope = SymbolicSize{"rope_dim"};
+    auto D_buf = SymbolicSize{"buffer_last_dim"};
+    auto D_qn = SymbolicSize{"q_nope_dim"};
+    auto D_qr = SymbolicSize{"q_rope_dim"};
+    auto D_qo = SymbolicSize{"q_out_dim"};
+    auto S_nope = SymbolicSize{"nope_stride"};
+    auto S_rope = SymbolicSize{"rope_stride"};
+    auto S_buf = SymbolicSize{"buffer_stride"};
+    auto S_loc = SymbolicSize{"loc_stride"};
+    auto S0_qn = SymbolicSize{"q_nope_stride_0"};
+    auto S1_qn = SymbolicSize{"q_nope_stride_1"};
+    auto S0_qr = SymbolicSize{"q_rope_stride_0"};
+    auto S1_qr = SymbolicSize{"q_rope_stride_1"};
+    auto S0_qo = SymbolicSize{"q_out_stride_0"};
+    auto S1_qo = SymbolicSize{"q_out_stride_1"};
+    auto loc_dtype = SymbolicDType{};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    D_nope.set_value(kFp8NopeDim);
+    D_rope.set_value(kFp8RopeDim);
+    D_qn.set_value(kFp8NopeDim);
+    D_qr.set_value(kFp8RopeDim);
+    D_qo.set_value(kFp8RowBytes);
+
+    TensorMatcher({B, D_nope})  //
+        .with_strides({S_nope, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(k_nope);
+    TensorMatcher({B, D_rope})  //
+        .with_strides({S_rope, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(k_rope);
+    TensorMatcher({-1, D_buf})  //
+        .with_strides({S_buf, 1})
+        .with_dtype<fp8_e4m3_t, uint8_t>()
+        .with_device(device)
+        .verify(kv_buffer);
+    TensorMatcher({B})  //
+        .with_strides({S_loc})
+        .with_dtype<int32_t, int64_t>(loc_dtype)
+        .with_device(device)
+        .verify(loc);
+    TensorMatcher({B, H, D_qn})  //
+        .with_strides({S0_qn, S1_qn, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(q_nope);
+    TensorMatcher({B, H, D_qr})  //
+        .with_strides({S0_qr, S1_qr, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(q_rope);
+    TensorMatcher({B, H, D_qo})  //
+        .with_strides({S0_qo, S1_qo, 1})
+        .with_dtype<fp8_e4m3_t, uint8_t>()
+        .with_device(device)
+        .verify(q_out);
+
+    CHECK_HOST(D_buf.unwrap() >= kFp8RowBytes) << "kv_buffer last dim too small";
+    CHECK_HOST(dcp_world_size >= 1 && dcp_rank >= 0 && dcp_rank < dcp_world_size)
+        << "invalid dcp world/rank: " << dcp_world_size << "/" << dcp_rank;
+    CHECK_HOST(S_loc.unwrap() == 1) << "loc must be contiguous; got stride " << S_loc.unwrap();
+
+    // Alignment tripwires (mirrored by python covered() so uncovered layouts
+    // fall back instead of faulting): 16B vector loads on the bf16 nope/q
+    // rows, 4B on the rope rows, 16B TMA dst rows, 16B int4 stores on q_out.
+    const auto aligned = [](const void* ptr, int64_t align) {
+      return reinterpret_cast<uintptr_t>(ptr) % static_cast<uintptr_t>(align) == 0;
+    };
+    CHECK_HOST(aligned(kv_buffer.data_ptr(), 16) && S_buf.unwrap() % 16 == 0)
+        << "kv_buffer base/row-stride must be 16-byte aligned for TMA bulk store";
+    CHECK_HOST(aligned(k_nope.data_ptr(), 16) && (S_nope.unwrap() * 2) % 16 == 0)
+        << "k_nope base/row-stride must be 16-byte aligned";
+    CHECK_HOST(aligned(k_rope.data_ptr(), 4) && (S_rope.unwrap() * 2) % 4 == 0)
+        << "k_rope base/row-stride must be 4-byte aligned";
+    CHECK_HOST(aligned(q_nope.data_ptr(), 16) && (S0_qn.unwrap() * 2) % 16 == 0 && (S1_qn.unwrap() * 2) % 16 == 0)
+        << "q_nope base/strides must be 16-byte aligned";
+    CHECK_HOST(aligned(q_rope.data_ptr(), 4) && (S0_qr.unwrap() * 2) % 4 == 0 && (S1_qr.unwrap() * 2) % 4 == 0)
+        << "q_rope base/strides must be 4-byte aligned";
+    CHECK_HOST(aligned(q_out.data_ptr(), 16) && S0_qo.unwrap() % 16 == 0 && S1_qo.unwrap() % 16 == 0)
+        << "q_out base/strides must be 16-byte aligned";
+
+    const uint32_t batch = static_cast<uint32_t>(B.unwrap());
+    const uint32_t num_heads = static_cast<uint32_t>(H.unwrap());
+    if (batch == 0) return;
+
+    const auto params = SetMlaKVConcatQFp8Params{
+        .k_nope = static_cast<const bf16_t*>(k_nope.data_ptr()),
+        .k_rope = static_cast<const bf16_t*>(k_rope.data_ptr()),
+        .kv_buffer = static_cast<uint8_t*>(kv_buffer.data_ptr()),
+        .loc = loc.data_ptr(),
+        .stride_nope = S_nope.unwrap(),
+        .stride_rope = S_rope.unwrap(),
+        .stride_buffer_bytes = S_buf.unwrap(),
+        .batch_size = batch,
+        .dcp_world_size = static_cast<int32_t>(dcp_world_size),
+        .dcp_rank = static_cast<int32_t>(dcp_rank),
+        .q_nope = static_cast<const bf16_t*>(q_nope.data_ptr()),
+        .q_rope = static_cast<const bf16_t*>(q_rope.data_ptr()),
+        .q_out = static_cast<uint8_t*>(q_out.data_ptr()),
+        .num_q_items = batch * num_heads,
+        .q_dim_1 = num_heads,
+        .qn_stride_0 = S0_qn.unwrap(),
+        .qn_stride_1 = static_cast<int32_t>(S1_qn.unwrap()),
+        .qr_stride_0 = S0_qr.unwrap(),
+        .qr_stride_1 = static_cast<int32_t>(S1_qr.unwrap()),
+        .qo_stride_0 = S0_qo.unwrap(),
+        .qo_stride_1 = static_cast<int32_t>(S1_qo.unwrap()),
+    };
+
+    const auto use_int32 = loc_dtype.is_type<int32_t>();
+    const uint32_t total_warps = params.batch_size + params.num_q_items;
+
+    auto launch = [&]<int kNW>() {
+      const auto kernel_ptr = use_int32 ? kernel<kNW, int32_t> : kernel<kNW, int64_t>;
+      const uint32_t num_blocks = div_ceil(total_warps, static_cast<uint32_t>(kNW));
+      LaunchKernel(num_blocks, static_cast<uint32_t>(kNW) * device::kWarpThreads, device.unwrap())
+          .enable_pdl(kUsePDL)(kernel_ptr, params);
+    };
+
+    switch (num_warps_per_block) {
+      case 1:
+        launch.template operator()<1>();
+        break;
+      case 2:
+        launch.template operator()<2>();
+        break;
+      case 4:
+        launch.template operator()<4>();
+        break;
+      case 8:
+        launch.template operator()<8>();
+        break;
+      default:
+        Panic("Unsupported num_warps_per_block=", num_warps_per_block);
+    }
+  }
+};
+
+}  // namespace

python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh

@@ -408,11 +408,22 @@ QuantHostContext<Trait> build_quant_context( //
     TensorMatcher({E, N, -1}).with_strides({-1, -1, 1}).with_dtype<T>().with_device(device).verify(input);
     TensorMatcher({E, N, H}).with_strides({-1, -1, 1}).with_dtype<Q>().with_device(device).verify(output_q);
     TensorMatcher({E, N, G}).with_strides({-1, -1, -1}).with_dtype<S>().with_device(device).verify(output_s);
+    CHECK_HOST((input.stride(0) * sizeof(T)) % 32 == 0)
+        << "input expert stride must keep rows 32B-aligned for the vectorized loads";
   } else {
     TensorMatcher({N, -1}).with_strides({-1, 1}).with_dtype<T>().with_device(device).verify(input);
     TensorMatcher({N, H}).with_strides({-1, 1}).with_dtype<Q>().with_device(device).verify(output_q);
     TensorMatcher({N, G}).with_strides({-1, -1}).with_dtype<S>().with_device(device).verify(output_s);
   }
+  // The 32B/lane vectorized loads need every input row to start 32B-aligned
+  // (kMaxVecBytes on Blackwell; over-strict but harmless on Hopper, whose
+  // 16B vectors only need 16). Contiguous allocations always satisfy this; it
+  // only bites hand-made row-strided views, which must keep rows aligned --
+  // rejected loudly here rather than densified silently at the call site.
+  CHECK_HOST(reinterpret_cast<uintptr_t>(input.data_ptr()) % 32 == 0)
+      << "input base pointer must be 32B-aligned for the vectorized loads";
+  CHECK_HOST((input.stride(-2) * sizeof(T)) % 32 == 0)
+      << "input token stride must keep rows 32B-aligned for the vectorized loads";
 
   const uint32_t num_tokens = N.unwrap();
   const uint32_t hidden_size = H.unwrap();

python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh

@@ -0,0 +1,231 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <array>
+#include <bit>
+#include <cstdint>
+#include <utility>
+
+namespace sglang {
+
+using namespace device;
+
+constexpr uint32_t kTinyNGemmVecSize = kMaxVecBytes / sizeof(bf16_t);
+
+template <uint32_t M, uint32_t N, uint32_t K, uint32_t N_SPLIT, typename OutT, bool kUsePDL>
+__global__ __launch_bounds__(K / kTinyNGemmVecSize, 1)  // 1 block per SM
+    void tiny_n_gemm_kernel(OutT* __restrict__ out, const bf16_t* __restrict__ x, const bf16_t* __restrict__ w) {
+  constexpr uint32_t kBlockSize = K / kTinyNGemmVecSize;
+  constexpr uint32_t kNumWarps = kBlockSize / kWarpThreads;
+  static_assert(M * N_SPLIT <= kBlockSize, "output tile must fit one thread each for the final reduce");
+  using vec_t = AlignedVector<bf16_t, kTinyNGemmVecSize>;
+
+  const uint32_t bx = blockIdx.x;
+  const uint32_t tx = threadIdx.x;
+  const bf16_t* w_tile = w + bx * (N_SPLIT * K);
+
+  // Weight prefetch: address is input-independent, load before the PDL wait.
+  vec_t wv[N_SPLIT];
+#pragma unroll
+  for (uint32_t n = 0; n < N_SPLIT; ++n) {
+    wv[n].load(w_tile + n * K, tx);
+  }
+
+  PDLWaitPrimary<kUsePDL>();
+
+  vec_t xv[M];
+#pragma unroll
+  for (uint32_t m = 0; m < M; ++m) {
+    xv[m].load(x + m * K, tx);
+  }
+
+  __shared__ float s_acc[kNumWarps][M * N_SPLIT];
+  const uint32_t warp_id = tx / kWarpThreads;
+
+#pragma unroll
+  for (uint32_t m = 0; m < M; ++m) {
+#pragma unroll
+    for (uint32_t n = 0; n < N_SPLIT; ++n) {
+      float acc = 0.0f;
+#if SGL_ARCH_BLACKWELL_OR_GREATER
+#pragma unroll
+      for (uint32_t i = 0; i < kTinyNGemmVecSize; ++i) {
+        acc = device::math::fma_f32_bf16(xv[m][i], wv[n][i], acc);
+      }
+#else
+      for (uint32_t i = 0; i < kTinyNGemmVecSize / 2; ++i) {
+        const auto [x0, x1] = cast<fp32x2_t>(bf16x2_t{xv[m][2 * i], xv[m][2 * i + 1]});
+        const auto [w0, w1] = cast<fp32x2_t>(bf16x2_t{wv[n][2 * i], wv[n][2 * i + 1]});
+        acc = fmaf(x0, w0, acc);
+        acc = fmaf(x1, w1, acc);
+      }
+#endif
+      // NOTE: broadcast write (all lanes hold the reduced value), safe here.
+      s_acc[warp_id][m * N_SPLIT + n] = warp::reduce_sum(acc);
+    }
+  }
+  PDLTriggerSecondary<kUsePDL>();
+  __syncthreads();
+
+  if (tx < M * N_SPLIT) {
+    float acc[kNumWarps];
+#pragma unroll
+    for (uint32_t i = 0; i < kNumWarps; ++i) {
+      acc[i] = s_acc[i][tx];
+    }
+#pragma unroll
+    for (uint32_t i = 1; i < kNumWarps; ++i) {
+      acc[0] += acc[i];
+    }
+    const uint32_t m = tx / N_SPLIT;
+    const uint32_t n = tx % N_SPLIT;
+    out[m * N + bx * N_SPLIT + n] = cast<OutT>(acc[0]);
+  }
+}
+
+SGL_DEVICE void cp_async_cg_16(void* smem_dst, const void* gmem_src, int32_t vec_offset) {
+  const uint32_t offset = static_cast<uint32_t>(vec_offset * 16);
+#if defined(USE_ROCM)
+  *reinterpret_cast<uint4*>(static_cast<char*>(smem_dst) + offset) =
+      *reinterpret_cast<const uint4*>(static_cast<const char*>(gmem_src) + offset);
+#else
+  const uint32_t smem_addr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst)) + offset;
+  const uint64_t gmem_addr = static_cast<uint64_t>(__cvta_generic_to_global(gmem_src)) + offset;
+  asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" : : "r"(smem_addr), "l"(gmem_addr) : "memory");
+#endif
+}
+
+constexpr uint32_t kTinyKGemmVecSize = 16 / sizeof(bf16_t);  // NOTE: no need to be large
+
+template <uint32_t M, uint32_t N, uint32_t K, uint32_t N_SPLIT, typename OutT, bool kUsePDL>
+__global__ __launch_bounds__(N_SPLIT* K / kTinyKGemmVecSize, 1)  // control the block size
+    void tiny_k_gemm_kernel(
+        OutT* __restrict__ out, const bf16_t* __restrict__ x, const bf16_t* __restrict__ w, const int64_t dx) {
+  using vec_t = AlignedVector<bf16_t, kTinyKGemmVecSize>;
+  constexpr uint32_t kNumKLanes = K / kTinyKGemmVecSize;
+  static_assert(std::has_single_bit(kNumKLanes), "K / vec_size must be a power of 2");
+  static_assert(kNumKLanes <= kWarpThreads, "require in-warp reduction");
+  static_assert((N_SPLIT * K / kTinyKGemmVecSize) % kWarpThreads == 0);
+  const uint32_t bx = blockIdx.x;
+  const uint32_t tx = threadIdx.x;
+  const uint32_t n_idx = bx * N_SPLIT + tx / kNumKLanes;
+  const uint32_t work_id = tx % kNumKLanes;
+  const bf16_t* w_tile = w + n_idx * K;
+
+  // Weight prefetch: address is input-independent, load before the PDL wait.
+  vec_t wv;
+  wv.load(w_tile, work_id);
+
+  PDLWaitPrimary<kUsePDL>();
+  vec_t xv[M];
+#pragma unroll
+  for (uint32_t m = 0; m < M; ++m) {
+    xv[m].load(x + m * dx, work_id);
+  }
+
+#pragma unroll
+  for (uint32_t m = 0; m < M; ++m) {
+    float acc = 0.0f;
+#pragma unroll
+    for (uint32_t i = 0; i < kTinyKGemmVecSize; ++i) {
+      acc = device::math::fma_f32_bf16(xv[m][i], wv[i], acc);
+    }
+    // Broadcast store: every lane of the group holds the reduced sum.
+    out[m * N + n_idx] = cast<OutT>(warp::reduce_sum<kNumKLanes>(acc));
+  }
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+}  // namespace sglang
+
+using namespace sglang;
+
+template <uint32_t N, uint32_t K, uint32_t kMaxM, uint32_t N_SPLIT, typename OutT, bool kUsePDL>
+struct TinyNGemmKernel {
+  static constexpr uint32_t kBlockSize = K / kTinyNGemmVecSize;
+  static constexpr uint32_t kNumBlocks = N / N_SPLIT;
+  static_assert(K % kTinyNGemmVecSize == 0, "K must be divisible by the vector width");
+  static_assert(kBlockSize % kWarpThreads == 0, "K / vec_size must be a multiple of the warp size");
+  static_assert(kBlockSize <= 1024, "K / vec_size exceeds the maximum block size");
+  static_assert(N % N_SPLIT == 0, "N must be divisible by split_n");
+  static_assert(kMaxM * N_SPLIT <= kBlockSize, "max_m * split_n must fit one thread each for the final reduce");
+
+  using KernelFn = void (*)(OutT*, const bf16_t*, const bf16_t*);
+
+  template <std::size_t... I>
+  static constexpr auto make_table(std::index_sequence<I...>) {
+    return std::array<KernelFn, kMaxM + 1>{nullptr, tiny_n_gemm_kernel<I + 1, N, K, N_SPLIT, OutT, kUsePDL>...};
+  }
+  static constexpr auto kTable = make_table(std::make_index_sequence<kMaxM>{});
+
+  static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView w, const tvm::ffi::TensorView out) {
+    using namespace host;
+
+    auto M = SymbolicSize{"num_tokens"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    TensorMatcher({M, K}).with_dtype<bf16_t>().with_device(device).verify(x);
+    TensorMatcher({N, K}).with_dtype<bf16_t>().with_device(device).verify(w);
+    TensorMatcher({M, N}).with_dtype<OutT>().with_device(device).verify(out);
+    const auto num_tokens = static_cast<uint32_t>(M.unwrap());
+    RuntimeCheck(num_tokens >= 1 && num_tokens <= kMaxM);
+    LaunchKernel(kNumBlocks, kBlockSize, device.unwrap())
+        .enable_pdl(kUsePDL)(
+            kTable[num_tokens],
+            static_cast<OutT*>(out.data_ptr()),
+            static_cast<const bf16_t*>(x.data_ptr()),
+            static_cast<const bf16_t*>(w.data_ptr()));
+  }
+};
+
+template <uint32_t N, uint32_t K, uint32_t kMaxM, uint32_t N_SPLIT, typename OutT, bool kUsePDL>
+struct TinyKGemmKernel {
+  static constexpr uint32_t kNumKLanes = K / kTinyKGemmVecSize;
+  static constexpr uint32_t kBlockSize = N_SPLIT * kNumKLanes;
+  static constexpr uint32_t kNumBlocks = N / N_SPLIT;
+  static_assert(K % kTinyKGemmVecSize == 0, "K must be divisible by the vector width");
+  static_assert(N % N_SPLIT == 0, "N must be divisible by split_n");
+  static_assert(kBlockSize % kWarpThreads == 0, "split_n * K-lanes must fill whole warps");
+  static_assert(kBlockSize <= 1024, "split_n * K-lanes exceeds the maximum block size");
+
+  using KernelFn = void (*)(OutT*, const bf16_t*, const bf16_t*, int64_t);
+
+  template <std::size_t... I>
+  static constexpr auto make_table(std::index_sequence<I...>) {
+    return std::array<KernelFn, kMaxM + 1>{nullptr, tiny_k_gemm_kernel<I + 1, N, K, N_SPLIT, OutT, kUsePDL>...};
+  }
+  static constexpr auto kTable = make_table(std::make_index_sequence<kMaxM>{});
+
+  static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView w, const tvm::ffi::TensorView out) {
+    using namespace host;
+
+    auto M = SymbolicSize{"num_tokens"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    // x may be a row-sliced view of a wider fused buffer: allow stride != K.
+    TensorMatcher({M, K}).with_dtype<bf16_t>().with_strides({-1, 1}).with_device(device).verify(x);
+    TensorMatcher({N, K}).with_dtype<bf16_t>().with_device(device).verify(w);
+    TensorMatcher({M, N}).with_dtype<OutT>().with_device(device).verify(out);
+    const auto num_tokens = static_cast<uint32_t>(M.unwrap());
+    const auto x_stride = static_cast<int64_t>(x.stride(0));
+    RuntimeCheck(num_tokens >= 1 && num_tokens <= kMaxM);
+    RuntimeCheck(
+        x_stride * sizeof(bf16_t) % (kTinyKGemmVecSize * sizeof(bf16_t)) == 0,
+        "x rows must stay aligned to the vector width, got stride ",
+        x_stride);
+    LaunchKernel(kNumBlocks, kBlockSize, device.unwrap())
+        .enable_pdl(kUsePDL)(
+            kTable[num_tokens],
+            static_cast<OutT*>(out.data_ptr()),
+            static_cast<const bf16_t*>(x.data_ptr()),
+            static_cast<const bf16_t*>(w.data_ptr()),
+            x_stride);
+  }
+};

python/sglang/kernels/jit/csrc/kimi_k3/attn_res/fused_tma.cuh

@@ -0,0 +1,946 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/mbarrier.cuh>
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include "../../distributed/custom_all_reduce.cuh"
+#include <algorithm>
+#include <array>
+#include <cfloat>
+#include <cstdint>
+#include <utility>
+
+// Local PTX primitives (mbarrier / bulk TMA / tcgen05 / warp-group sync)
+
+namespace ptx {
+
+// ---- bulk 1D TMA (PTX ISA §9.7.9.25) ---------------------------------------
+
+// global -> shared::cluster, completed by an smem mbarrier. Arm `bar` with
+// `mbar_arrive_expect_tx(bar, bytes)` before issuing; `bytes` and both
+// endpoints must be 16-byte aligned.
+static SGL_DEVICE void cp_async_bulk_1d_load(void* smem_dst, const void* gmem_src, uint32_t bytes, uint64_t* bar) {
+  asm volatile(
+      "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
+      " [%0], [%1], %2, [%3];" ::"r"(to_shared(smem_dst)),
+      "l"(gmem_src),
+      "r"(bytes),
+      "r"(to_shared(bar))
+      : "memory");
+}
+
+// Publish mbarrier initialization from the generic proxy before an async engine
+// uses the barrier.
+static SGL_DEVICE void fence_mbarrier_init() {
+  asm volatile("fence.mbarrier_init.release.cluster;");
+}
+
+// ---- warp / warp-group sync (PTX ISA §9.7.4, §9.7.12.6, §9.7.13) -----------
+
+// Partial-CTA rendezvous. `id` must be in [1, 15]; barrier 0 is reserved for
+// the full-CTA barrier behind __syncthreads().
+static SGL_DEVICE void named_barrier_sync(uint32_t id, uint32_t num_threads) {
+  asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(num_threads) : "memory");
+}
+
+// True on exactly one lane of the issuing warp — guards single-issuer sites
+// (mbar init, TMA issue, MMA issue, TMEM alloc) without gating on lane_id.
+static SGL_DEVICE bool elect_one() {
+  uint32_t pred;
+  asm volatile(
+      "{\n\t.reg .pred p;\n\t"
+      "elect.sync _|p, 0xffffffff;\n\t"
+      "selp.b32 %0, 1, 0, p;\n\t}\n"
+      : "=r"(pred));
+  return pred != 0;
+}
+
+// Runtime warp-group register-budget reallocation: widen the epilogue's
+// per-thread budget (so it holds a larger primary array without spilling) by
+// narrowing the mainloop warps, which need few registers.
+//
+// Both forms are `.sync.aligned`: all 128 threads of the issuing warp-group
+// must execute the SAME instruction with the SAME N, from a warp-group
+// boundary (issuing from only one warp of the group hangs). N in [24, 256],
+// multiple of 8, per thread; the CTA total must satisfy
+// sum(warp_group_threads * N) <= 64512 (the safe allocatable cap on B100/B300
+// after ~1024 reserved regs). Caller owns the budgeting — there is no
+// compile-time check, since N per warp-group is orthogonal.
+//
+// This pays only for ASYMMETRIC budgets that ptxas cannot infer from the
+// source. For a symmetric cap, `__launch_bounds__(NUM_THREADS, 1)` is cleaner
+// and measured faster on B100/B300.
+template <int N>
+static SGL_DEVICE void setmaxnreg_dec() {
+  static_assert(N >= 24 && N <= 256, "setmaxnreg N must be in [24, 256]");
+  static_assert((N & 7) == 0, "setmaxnreg N must be a multiple of 8");
+  asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" ::"n"(N));
+}
+
+template <int N>
+static SGL_DEVICE void setmaxnreg_inc() {
+  static_assert(N >= 24 && N <= 256, "setmaxnreg N must be in [24, 256]");
+  static_assert((N & 7) == 0, "setmaxnreg N must be a multiple of 8");
+  asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" ::"n"(N));
+}
+
+// ---- tcgen05 (PTX ISA §9.7.16) ---------------------------------------------
+//
+// Lifecycle (mandatory order, §9.7.16.7.1): alloc (one warp, n_cols a power of
+// 2 in [32, 512], TMEM address written to smem) -> __syncthreads + read taddr
+// -> ld/st -> dealloc -> relinquish before kernel exit.
+//
+// Each warp can only touch its own 32-lane TMEM band (§9.7.16.8.1): warp 0 ->
+// lanes 0-31, warp 1 -> 32-63, and so on. Use `tcgen05_wait_st` /
+// `tcgen05_wait_ld` before consuming the other side of a store / drain.
+static SGL_DEVICE void tcgen05_alloc(uint32_t smem_addr_for_taddr, uint32_t n_cols) {
+  asm volatile(
+      "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(smem_addr_for_taddr), "r"(n_cols));
+}
+
+static SGL_DEVICE void tcgen05_dealloc(uint32_t taddr, uint32_t n_cols) {
+  asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(taddr), "r"(n_cols));
+}
+
+static SGL_DEVICE void tcgen05_relinquish() {
+  asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;");
+}
+
+// .32x32b.x8: 8 b32 per lane = 8 TMEM columns. Per-lane 8 FP32 -> 4 bf16x2
+// packs = one int4, the natural fit for a BF16 epilogue moving a column band
+// with 16-byte smem accesses.
+static SGL_DEVICE void tcgen05_ld_32x32b_x8(
+    uint32_t taddr,
+    uint32_t& r0,
+    uint32_t& r1,
+    uint32_t& r2,
+    uint32_t& r3,
+    uint32_t& r4,
+    uint32_t& r5,
+    uint32_t& r6,
+    uint32_t& r7) {
+  asm volatile(
+      "tcgen05.ld.sync.aligned.32x32b.x8.b32 "
+      " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];"
+      : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3), "=r"(r4), "=r"(r5), "=r"(r6), "=r"(r7)
+      : "r"(taddr));
+}
+
+static SGL_DEVICE void tcgen05_ld_32x32b_x8(uint32_t taddr, uint32_t* dst) {
+  tcgen05_ld_32x32b_x8(taddr, dst[0], dst[1], dst[2], dst[3], dst[4], dst[5], dst[6], dst[7]);
+}
+
+static SGL_DEVICE void tcgen05_st_32x32b_x8(uint32_t taddr, const uint32_t* src) {
+  asm volatile(
+      "tcgen05.st.sync.aligned.32x32b.x8.b32 "
+      " [%8], {%0, %1, %2, %3, %4, %5, %6, %7};"
+      :
+      : "r"(src[0]),
+        "r"(src[1]),
+        "r"(src[2]),
+        "r"(src[3]),
+        "r"(src[4]),
+        "r"(src[5]),
+        "r"(src[6]),
+        "r"(src[7]),
+        "r"(taddr));
+}
+
+static SGL_DEVICE void tcgen05_wait_st() {
+  asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory");
+}
+
+}  // namespace ptx
+
+namespace sglang {
+
+struct AttnResTMAParams {
+  const bf16_t* __restrict__ prefix_sum;  // [T, H]
+  const bf16_t* __restrict__ bank;        // [T, NB_total, H]
+  const bf16_t* __restrict__ cw;          // [H] score norm * proj weight
+  const bf16_t* __restrict__ ow;          // [H] out norm weight
+  bf16_t* __restrict__ out;               // [T, H]
+  // Fused bank write (nullptr = off): per-token destination of the prefix
+  // row snapshot, bank row nvb (strided by stride_bm like the read rows).
+  // The kernel never reads row nvb, so the write races with nothing.
+  bf16_t* __restrict__ prefix_dst;
+  // Optional fused NVLS reduce-scatter source. When input_mc is non-null,
+  // the producer warp materializes this rank's reduced token shard plus the
+  // local residual into prefix_sum/prefix_out before TMA consumes it.
+  const uint8_t* input_mc;
+  const bf16_t* residual;
+  bf16_t* prefix_out;
+  device::distributed::Semaphore* sem_local;
+  uint8_t* sem_mc;
+  uint8_t* output_mc;
+  uint32_t world_size;
+  uint32_t rank;
+  int64_t stride_bm;  // bank stride along T (in elements)
+  float eps;
+  uint32_t num_tokens;
+};
+
+template <int64_t kDim_, uint32_t kNumBankRows_, uint32_t kChunkRows_, uint32_t kConsumerRegs_ = 0>
+struct KimiK3AttnResTrait {
+ public:
+  static constexpr int64_t kDim = kDim_;
+  static constexpr int64_t kTile = 1024;               // one warp-group-wide 16B sweep
+  static constexpr uint32_t kNumRows = kNumBankRows_;  // bank rows; +1 prefix row
+  static constexpr uint32_t kChunkRows = kChunkRows_;  // rows per chunk (one barrier pair per chunk)
+  // Chunk slots in the smem ring. Frozen at 2 (double buffering): 1 stalls
+  // the producer behind the consumers (~10% slower), >2 gains nothing and
+  // costs smem at small T.
+  static constexpr uint32_t kNumStages = 2;
+  static constexpr uint32_t kNumChunks = (kNumRows + 1 + kChunkRows - 1) / kChunkRows;
+  static constexpr uint32_t kNumConsumerWarps = 8;
+  static constexpr uint32_t kConsumerRegs = kConsumerRegs_;
+  static constexpr uint32_t kProducerRegs = 40;
+  static constexpr uint32_t kNumProducerWarps = kConsumerRegs > 0 ? 4 : 1;
+  static constexpr uint32_t kNumWarps = kNumConsumerWarps + kNumProducerWarps;
+  static constexpr uint32_t kNumThreads = kNumWarps * device::kWarpThreads;
+  static constexpr uint32_t kNumConsumerThreads = kNumConsumerWarps * device::kWarpThreads;
+  static_assert(
+      kConsumerRegs == 0 || (kConsumerRegs % 8 == 0 && 24 <= kConsumerRegs && kConsumerRegs <= 256 &&
+                             2 * kConsumerRegs + kProducerRegs <= 512),
+      "consumer register budget exceeds the SM sub-partition file");
+
+  // Consumer tiling (v1 layout): two 128-thread warp groups; group g owns
+  // tiles g, g + 2, ... of the row; each thread owns one 16B vector per tile.
+  static constexpr uint32_t kNumGroups = 2;
+  static constexpr uint32_t kGroupThreads = kNumConsumerThreads / kNumGroups;
+  static constexpr uint32_t kVecElems = 16 / sizeof(bf16_t);  // smem ld/st are 16B max
+  static constexpr uint32_t kNumTiles = kDim / kTile;
+  static constexpr uint32_t kSlicesPerGroup = (kNumTiles + kNumGroups - 1) / kNumGroups;
+  static constexpr uint32_t kAccPerThread = kSlicesPerGroup * kVecElems;
+
+  // TMEM: per group, kTmemColsPerGroup columns of cw then of ow.
+  static constexpr uint32_t kTmemColsPerGroup = 32;
+  static constexpr uint32_t kTmemCols = 2 * kNumGroups * kTmemColsPerGroup;
+  static constexpr uint32_t kConsumerBarId = 1;  // barrier 0 stays __syncthreads'
+
+  static_assert(kDim % kTile == 0, "kDim must be a whole number of tiles");
+  static_assert(kTile == kGroupThreads * kVecElems, "a tile is one group-wide 16B sweep");
+  static_assert(kNumTiles <= kNumGroups * kSlicesPerGroup, "slices must cover all tiles");
+  static_assert(kSlicesPerGroup * kVecElems <= kTmemColsPerGroup, "weight slices must fit their TMEM columns");
+  static_assert(kNumRows >= 1, "need at least one bank row");
+  static_assert(kChunkRows >= 1, "need at least one chunk row");
+
+  struct Smem {
+    uint64_t bar_full[kNumStages];
+    uint64_t bar_free[kNumStages];
+    float warp_rms[kNumConsumerWarps][kChunkRows];
+    float warp_dot[kNumConsumerWarps][kChunkRows];
+    // The out-norm reduction gets its own buffer: it can overlap the next
+    // token's first score reduction.
+    float warp_ssq[kNumConsumerWarps];
+    uint32_t tmem_base;
+    alignas(128) bf16_t buf[kNumStages][kChunkRows][kDim];
+  };
+
+  static SGL_DEVICE void forward(const AttnResTMAParams& params, Smem* smem);
+};
+
+SGL_DEVICE float2 fma_f32x2(float2 a, float2 b, float2 c) {
+  const uint64_t a_bits = reinterpret_cast<const uint64_t&>(a);
+  const uint64_t b_bits = reinterpret_cast<const uint64_t&>(b);
+  const uint64_t c_bits = reinterpret_cast<const uint64_t&>(c);
+  uint64_t result;
+  asm("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(result) : "l"(a_bits), "l"(b_bits), "l"(c_bits));
+  return reinterpret_cast<const float2&>(result);
+}
+
+SGL_DEVICE float2 mul_f32x2(float2 a, float2 b) {
+  const uint64_t a_bits = reinterpret_cast<const uint64_t&>(a);
+  const uint64_t b_bits = reinterpret_cast<const uint64_t&>(b);
+  uint64_t result;
+  asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(result) : "l"(a_bits), "l"(b_bits));
+  return reinterpret_cast<const float2&>(result);
+}
+
+template <int64_t kDim_, uint32_t kNumBankRows_, uint32_t kChunkRows_, uint32_t kConsumerRegs_>
+SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerRegs_>::forward(
+    const AttnResTMAParams& params, Smem* smem) {
+  using namespace device;
+  using row_vec_t = AlignedVector<bf16x2_t, kVecElems / 2>;  // 16 bytes
+  const auto tx = threadIdx.x;
+  const auto warp_id = tx / kWarpThreads;
+  const auto lane_id = tx % kWarpThreads;
+
+  if (warp_id == 0 && lane_id < kNumStages) {
+    ::ptx::mbar_init(&smem->bar_full[lane_id], 1);
+    ::ptx::mbar_init(&smem->bar_free[lane_id], kNumConsumerWarps * kWarpThreads);
+    ::ptx::fence_mbarrier_init();
+  } else if (warp_id == 1) {
+    ::ptx::tcgen05_alloc(::ptx::to_shared(&smem->tmem_base), kTmemCols);
+    ::ptx::tcgen05_relinquish();
+  }
+
+  __syncthreads();
+  if (warp_id >= kNumConsumerWarps) {  // producer warp (group); first warp works
+    if constexpr (kConsumerRegs > 0) ::ptx::setmaxnreg_dec<kProducerRegs>();
+    // TODO: reduce the register usage
+    if (warp_id == kNumConsumerWarps && ::ptx::elect_one()) {
+      uint32_t global_chunks = 0;
+      constexpr uint32_t kRowBytes = kDim * sizeof(bf16_t);
+      for (auto token = blockIdx.x; token < params.num_tokens; token += gridDim.x) {
+#pragma unroll
+        for (uint32_t ci = 0; ci < kNumChunks; ++ci, ++global_chunks) {
+          const uint32_t base_row = ci * kChunkRows;
+          const uint32_t an = (kNumRows + 1 - base_row) < kChunkRows ? (kNumRows + 1 - base_row) : kChunkRows;
+          const auto slot = global_chunks % kNumStages;
+          const auto phase = (global_chunks / kNumStages) & 1;
+          if (global_chunks >= kNumStages) {
+            ::ptx::mbar_wait_parity(&smem->bar_free[slot], phase ^ 1);
+          }
+          // One barrier per chunk; each row still gets its own bulk copy.
+          ::ptx::mbar_arrive_expect_tx(&smem->bar_full[slot], an * kRowBytes);
+#pragma unroll
+          for (uint32_t r = 0; r < an; ++r) {
+            const auto row = base_row + r;
+            const auto src = row == kNumRows ? params.prefix_sum + token * kDim  //
+                                             : params.bank + token * params.stride_bm + row * kDim;
+            // Only prefix_sum is written by the immediately-preceding kernel;
+            // one wait before the first token's prefix load covers the rest.
+            if (token == blockIdx.x && row == kNumRows) PDLWaitPrimary<true>();
+            ::ptx::cp_async_bulk_1d_load(&smem->buf[slot][r], src, kRowBytes, &smem->bar_full[slot]);
+          }
+        }
+      }
+      PDLTriggerSecondary<true>();
+    }
+  } else {  // 2 consumer warp groups; one chunk per rendezvous
+    if constexpr (kConsumerRegs > 0) ::ptx::setmaxnreg_inc<kConsumerRegs>();
+    const auto group = warp_id / (kNumConsumerWarps / kNumGroups);
+    const auto tid_in_group = tx % kGroupThreads;
+    const auto tmem_cw = smem->tmem_base + group * kTmemColsPerGroup;
+    const auto tmem_ow = tmem_cw + kNumGroups * kTmemColsPerGroup;
+
+    // Stage this thread's cw / ow slices into TMEM (read once from gmem).
+    {
+      float staged[kAccPerThread];
+#pragma unroll
+      for (uint32_t si = 0; si < kSlicesPerGroup; ++si) {
+        const auto tile = si * kNumGroups + group;
+        if (tile >= kNumTiles) continue;
+        const auto h_base = tile * kTile + tid_in_group * kVecElems;
+#pragma unroll
+        for (uint32_t j = 0; j < kVecElems; ++j) {
+          staged[si * kVecElems + j] = __bfloat162float(params.cw[h_base + j]);
+        }
+      }
+#pragma unroll
+      for (uint32_t si = 0; si < kSlicesPerGroup; ++si) {
+        ::ptx::tcgen05_st_32x32b_x8(
+            tmem_cw + si * kVecElems, reinterpret_cast<const uint32_t*>(&staged[si * kVecElems]));
+      }
+#pragma unroll
+      for (uint32_t si = 0; si < kSlicesPerGroup; ++si) {
+        const auto tile = si * kNumGroups + group;
+        if (tile >= kNumTiles) continue;
+        const auto h_base = tile * kTile + tid_in_group * kVecElems;
+#pragma unroll
+        for (uint32_t j = 0; j < kVecElems; ++j) {
+          staged[si * kVecElems + j] = __bfloat162float(params.ow[h_base + j]);
+        }
+      }
+#pragma unroll
+      for (uint32_t si = 0; si < kSlicesPerGroup; ++si) {
+        ::ptx::tcgen05_st_32x32b_x8(
+            tmem_ow + si * kVecElems, reinterpret_cast<const uint32_t*>(&staged[si * kVecElems]));
+      }
+      ::ptx::tcgen05_wait_st();
+    }
+
+    uint32_t global_chunks = 0;  // mirrors the producer's chunk counter
+    for (auto token = blockIdx.x; token < params.num_tokens; token += gridDim.x) {
+      float run_max = -FLT_MAX;  // online-softmax state
+      float run_sum = 0.f;
+      float2 acc[kAccPerThread / 2] = {};  // packed fp32x2 accumulator
+
+#pragma unroll
+      for (uint32_t ci = 0; ci < kNumChunks; ++ci, ++global_chunks) {
+        const uint32_t base_row = ci * kChunkRows;
+        // Active rows of this chunk; folds per unrolled iteration.
+        const uint32_t an = (kNumRows + 1 - base_row) < kChunkRows ? (kNumRows + 1 - base_row) : kChunkRows;
+        const auto slot = global_chunks % kNumStages;
+        const auto phase = (global_chunks / kNumStages) & 1;
+        ::ptx::mbar_wait_parity(&smem->bar_full[slot], phase);
+
+        // Score pass: the cw slice is loaded once and reused across the
+        // chunk's rows; each row's 16B slices land in registers. rms/dot
+        // accumulate as packed fp32x2 lanes, folded to scalars just before
+        // the warp reduction.
+        row_vec_t rows[kSlicesPerGroup][kChunkRows];
+        float2 acc_rms2[kChunkRows] = {};
+        float2 acc_dot2[kChunkRows] = {};
+#pragma unroll
+        for (uint32_t si = 0; si < kSlicesPerGroup; ++si) {
+          const auto tile = si * kNumGroups + group;
+          if (tile >= kNumTiles) continue;
+          float q[kVecElems];
+          ::ptx::tcgen05_ld_32x32b_x8(tmem_cw + si * kVecElems, reinterpret_cast<uint32_t*>(q));
+          const auto* q2 = reinterpret_cast<const float2*>(q);
+          const auto offset = tile * kTile + tid_in_group * kVecElems;
+#pragma unroll
+          for (uint32_t r = 0; r < an; ++r) {
+            rows[si][r].load(&smem->buf[slot][r][offset]);
+          }
+#pragma unroll
+          for (uint32_t r = 0; r < an; ++r) {
+#pragma unroll
+            for (uint32_t j = 0; j < kVecElems / 2; ++j) {
+              const auto f = cast<float2>(rows[si][r][j]);
+              acc_rms2[r] = fma_f32x2(f, f, acc_rms2[r]);
+              acc_dot2[r] = fma_f32x2(f, q2[j], acc_dot2[r]);
+            }
+          }
+        }
+        ::ptx::mbar_arrive(&smem->bar_free[slot]);
+
+        // Fused bank write: the prefix row (last row of the last chunk) is
+        // already in registers; snapshot it to bank row nvb with plain
+        // stores — the .write() copy kernel disappears. Placed after the
+        // arrive so the slot handoff is not delayed.
+        if (params.prefix_dst != nullptr && base_row + an == kNumRows + 1) {
+          const uint32_t pr = kNumRows - base_row;
+          auto* dst = params.prefix_dst + static_cast<int64_t>(token) * params.stride_bm;
+#pragma unroll
+          for (uint32_t si = 0; si < kSlicesPerGroup; ++si) {
+            const auto tile = si * kNumGroups + group;
+            if (tile >= kNumTiles) continue;
+            rows[si][pr].store(dst, tile * (kTile / kVecElems) + tid_in_group);
+          }
+        }
+
+        float acc_rms[kChunkRows];
+        float acc_dot[kChunkRows];
+#pragma unroll
+        for (uint32_t r = 0; r < an; ++r) {
+          acc_rms[r] = acc_rms2[r].x + acc_rms2[r].y;
+          acc_dot[r] = acc_dot2[r].x + acc_dot2[r].y;
+        }
+
+#pragma unroll
+        for (int n = 0; n < an; n++) {
+          acc_rms[n] = warp::reduce_sum(acc_rms[n]);
+          acc_dot[n] = warp::reduce_sum(acc_dot[n]);
+        }
+        if (lane_id == 0) {
+#pragma unroll
+          for (uint32_t r = 0; r < an; ++r) {
+            smem->warp_rms[warp_id][r] = acc_rms[r];
+            smem->warp_dot[warp_id][r] = acc_dot[r];
+          }
+        }
+        ::ptx::named_barrier_sync(kConsumerBarId, kNumConsumerThreads);
+        // Lane r totals row r, then broadcasts: an*16 smem loads per warp
+        // instead of per thread.
+        float lane_logit = 0.f;
+        if (lane_id < an) {
+          float total_rms = 0.f;
+          float total_dot = 0.f;
+#pragma unroll
+          for (uint32_t w = 0; w < kNumConsumerWarps; ++w) {
+            total_rms += smem->warp_rms[w][lane_id];
+            total_dot += smem->warp_dot[w][lane_id];
+          }
+          constexpr float kScale = 1.f / static_cast<float>(kDim);
+          lane_logit = total_dot * rsqrtf(total_rms * kScale + params.eps);
+        }
+        float logit[kChunkRows];
+#pragma unroll
+        for (uint32_t r = 0; r < an; ++r) {
+          logit[r] = __shfl_sync(0xffffffffu, lane_logit, r);
+        }
+
+        // Online-softmax fold of the chunk into the running accumulator.
+        float chunk_max = -FLT_MAX;
+#pragma unroll
+        for (uint32_t r = 0; r < an; ++r) {
+          chunk_max = fmaxf(chunk_max, logit[r]);
+        }
+        const float new_max = fmaxf(run_max, chunk_max);
+        const float correction = exp2f((run_max - new_max) * math::log2e);
+        float weight[kChunkRows];
+        float weight_sum = 0.f;
+#pragma unroll
+        for (uint32_t r = 0; r < an; ++r) {
+          weight[r] = exp2f((logit[r] - new_max) * math::log2e);
+          weight_sum += weight[r];
+        }
+        run_sum = run_sum * correction + weight_sum;
+        run_max = new_max;
+
+        // Fold the chunk into the packed accumulator (v1 loop order: scale
+        // once, then rows outer / vector lanes inner, all fp32x2 FMAs).
+        const float2 correction2 = make_float2(correction, correction);
+        float2 weight2[kChunkRows];
+#pragma unroll
+        for (uint32_t r = 0; r < an; ++r) {
+          weight2[r] = make_float2(weight[r], weight[r]);
+        }
+#pragma unroll
+        for (uint32_t si = 0; si < kSlicesPerGroup; ++si) {
+          const auto tile = si * kNumGroups + group;
+          if (tile >= kNumTiles) continue;
+          float2 a[kVecElems / 2];
+#pragma unroll
+          for (uint32_t j = 0; j < kVecElems / 2; ++j) {
+            a[j] = mul_f32x2(acc[si * (kVecElems / 2) + j], correction2);
+          }
+#pragma unroll
+          for (uint32_t r = 0; r < an; ++r) {
+#pragma unroll
+            for (uint32_t j = 0; j < kVecElems / 2; ++j) {
+              a[j] = fma_f32x2(weight2[r], cast<float2>(rows[si][r][j]), a[j]);
+            }
+          }
+#pragma unroll
+          for (uint32_t j = 0; j < kVecElems / 2; ++j) {
+            acc[si * (kVecElems / 2) + j] = a[j];
+          }
+        }
+      }
+
+      // Fused out norm: mixed = acc / run_sum, out = rmsnorm(mixed) * ow.
+      const float inv_sum = 1.f / run_sum;
+      float2 acc_sq2 = make_float2(0.f, 0.f);
+#pragma unroll
+      for (uint32_t j = 0; j < kAccPerThread / 2; ++j) {
+        acc_sq2 = fma_f32x2(acc[j], acc[j], acc_sq2);
+      }
+      float acc_sq = warp::reduce_sum(acc_sq2.x + acc_sq2.y);
+      if (lane_id == 0) smem->warp_ssq[warp_id] = acc_sq;
+      ::ptx::named_barrier_sync(kConsumerBarId, kNumConsumerThreads);
+      float total_sq = 0.f;
+#pragma unroll
+      for (uint32_t w = 0; w < kNumConsumerWarps; ++w) {
+        total_sq += smem->warp_ssq[w];
+      }
+      const float scale = inv_sum * rsqrtf(total_sq * inv_sum * inv_sum / static_cast<float>(kDim) + params.eps);
+      const float2 scale2 = make_float2(scale, scale);
+
+      auto* out_ptr = params.out + static_cast<int64_t>(token) * kDim;
+#pragma unroll
+      for (uint32_t si = 0; si < kSlicesPerGroup; ++si) {
+        const auto tile = si * kNumGroups + group;
+        if (tile >= kNumTiles) continue;
+        float q[kVecElems];
+        ::ptx::tcgen05_ld_32x32b_x8(tmem_ow + si * kVecElems, reinterpret_cast<uint32_t*>(q));
+        const auto* q2 = reinterpret_cast<const float2*>(q);
+        row_vec_t out_vec;
+#pragma unroll
+        for (uint32_t j = 0; j < kVecElems / 2; ++j) {
+          const auto scaled = mul_f32x2(acc[si * (kVecElems / 2) + j], scale2);
+          out_vec[j] = cast<bf16x2_t>(mul_f32x2(scaled, q2[j]));
+        }
+        const auto row_vid = tile * (kTile / kVecElems) + tid_in_group;
+        if (params.output_mc != nullptr) {
+          const auto global_token = static_cast<int64_t>(params.rank) * params.num_tokens + token;
+          const auto global_vid = global_token * (kDim / kVecElems) + row_vid;
+          st_multimem_16B(out_vec, params.output_mc, global_vid);
+        } else {
+          out_vec.store(out_ptr, row_vid);
+        }
+      }
+    }
+    ::ptx::named_barrier_sync(kConsumerBarId, kNumConsumerThreads);
+    if (warp_id == 1) {
+      ::ptx::tcgen05_dealloc(smem->tmem_base, kTmemCols);
+    }
+  }
+}
+
+// kOccupancy > 1 caps the register budget (65536 / (kOccupancy * kNumThreads))
+// so that many CTAs actually co-reside; smem must also fit kOccupancy copies.
+template <typename Trait, uint32_t kOccupancy>
+__global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy)
+    attn_res_fused_tma_kernel(const __grid_constant__ AttnResTMAParams params) {
+  extern __shared__ char smem_raw[];
+  Trait::forward(params, reinterpret_cast<typename Trait::Smem*>(smem_raw));
+}
+
+SGL_DEVICE uint32_t* attn_res_sem_mc_flag(uint8_t* sem_mc, uint32_t block) {
+  static_assert(sizeof(device::distributed::Semaphore) == 128);
+  return reinterpret_cast<uint32_t*>(sem_mc + block * sizeof(device::distributed::Semaphore));
+}
+
+SGL_DEVICE void attn_res_sem_arrive_relaxed(uint32_t* flag) {
+#if SGL_ARCH_HOPPER_OR_GREATER
+  asm volatile("multimem.red.relaxed.sys.global.add.u32 [%0], 1;" ::"l"(flag) : "memory");
+#else
+  assert(false && "multimem red requires Hopper or later");
+#endif
+}
+
+SGL_DEVICE void attn_res_sem_arrive_release(uint32_t* flag) {
+#if SGL_ARCH_HOPPER_OR_GREATER
+  asm volatile("multimem.red.release.sys.global.add.u32 [%0], 1;" ::"l"(flag) : "memory");
+#else
+  assert(false && "multimem red requires Hopper or later");
+#endif
+}
+
+// Fused NVLS pull RS + local residual + attention-residual aggregation.
+// The entry/exit barriers make local o_proj writes visible before the
+// producer's multimem reduction and preserve the shared pull-semaphore
+// protocol used by the neighboring K3 collectives.
+template <typename Trait, uint32_t kOccupancy>
+__global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy)
+    attn_res_fused_pull_rs_kernel(const __grid_constant__ AttnResTMAParams params) {
+  __shared__ uint32_t exit_base;
+  if (threadIdx.x == 0) {
+    auto* semaphore = &params.sem_local[blockIdx.x];
+    const auto reserved = semaphore->counter_ptr()->inc(2 * params.world_size);
+    exit_base = reserved + params.world_size;
+    device::PDLWaitPrimary<true>();
+    attn_res_sem_arrive_relaxed(attn_res_sem_mc_flag(params.sem_mc, blockIdx.x));
+    while (semaphore->get_relaxed() - reserved < params.world_size)
+      ;
+  }
+  __syncthreads();
+
+  // Cooperative NVLS materialization: all TMA producer + consumer threads
+  // participate, so the remote reduction exposes hundreds of outstanding
+  // 16-byte loads per CTA instead of serializing the row through one warp.
+  using pull_vec_t = device::AlignedVector<bf16x2_t, 4>;
+  using SumOp = device::ReductionTrait<device::ReductionOp::SUM, bf16x2_t>;
+  constexpr uint32_t kRowVecs = Trait::kDim * sizeof(bf16_t) / sizeof(pull_vec_t);
+  for (auto token = blockIdx.x; token < params.num_tokens; token += gridDim.x) {
+    auto* prefix = params.prefix_out + static_cast<int64_t>(token) * Trait::kDim;
+    const auto* input_mc = params.input_mc + static_cast<int64_t>(token) * Trait::kDim * sizeof(bf16_t);
+    const auto* residual =
+        params.residual == nullptr ? nullptr : params.residual + static_cast<int64_t>(token) * Trait::kDim;
+    for (uint32_t vid = threadIdx.x; vid < kRowVecs; vid += blockDim.x) {
+      pull_vec_t vec;
+      ld_multimem_16B(vec, input_mc, vid);
+      if (residual != nullptr) {
+        pull_vec_t res;
+        res.load(residual, vid);
+#pragma unroll
+        for (uint32_t j = 0; j < 4; ++j) {
+          vec[j] = SumOp::reduce(vec[j], res[j]);
+        }
+      }
+      vec.store(prefix, vid);
+    }
+  }
+  __threadfence();
+  __syncthreads();
+
+  extern __shared__ char smem_raw[];
+  Trait::forward(params, reinterpret_cast<typename Trait::Smem*>(smem_raw));
+
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    auto* semaphore = &params.sem_local[blockIdx.x];
+    attn_res_sem_arrive_release(attn_res_sem_mc_flag(params.sem_mc, blockIdx.x));
+    while (semaphore->get_acquire() - exit_base < params.world_size)
+      ;
+  }
+}
+
+// Local attention-residual aggregation + direct AG epilogue. The consumer
+// threads already hold each normalized 16B output vector in registers, so
+// Trait::forward multicast-stores those vectors into every peer's symmetric
+// full-token output instead of launching a separate all-gather.
+template <typename Trait, uint32_t kOccupancy>
+__global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy)
+    attn_res_fused_direct_ag_kernel(const __grid_constant__ AttnResTMAParams params) {
+  __shared__ uint32_t exit_base;
+  if (threadIdx.x == 0) {
+    auto* semaphore = &params.sem_local[blockIdx.x];
+    const auto reserved = semaphore->counter_ptr()->inc(2 * params.world_size);
+    exit_base = reserved + params.world_size;
+    attn_res_sem_arrive_relaxed(attn_res_sem_mc_flag(params.sem_mc, blockIdx.x));
+    while (semaphore->get_relaxed() - reserved < params.world_size)
+      ;
+  }
+  __syncthreads();
+
+  extern __shared__ char smem_raw[];
+  Trait::forward(params, reinterpret_cast<typename Trait::Smem*>(smem_raw));
+
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    auto* semaphore = &params.sem_local[blockIdx.x];
+    attn_res_sem_arrive_release(attn_res_sem_mc_flag(params.sem_mc, blockIdx.x));
+    while (semaphore->get_acquire() - exit_base < params.world_size)
+      ;
+  }
+}
+
+}  // namespace sglang
+
+using namespace sglang;
+using host::distributed::CommunicatorRef;
+
+// Host launcher: constexpr kernel table over nvb.
+
+template <int64_t kDim, uint32_t kMaxBankRows, uint32_t kChunkRows, uint32_t kOccupancy, uint32_t kConsumerRegs>
+struct AttnResFusedTmaKernel {
+  using KernelFn = void (*)(const AttnResTMAParams);
+  template <uint32_t kNvb>
+  using Trait = KimiK3AttnResTrait<kDim, kNvb, kChunkRows, kConsumerRegs>;
+  static constexpr uint32_t kNumThreads = Trait<1>::kNumThreads;
+  static constexpr size_t kSmemBytes = sizeof(typename Trait<1>::Smem);
+  // kOccupancy copies of the smem ring must fit one SM (228KB on SM100).
+  static_assert(kOccupancy >= 1 && kOccupancy * kSmemBytes <= 233472 - 1024, "occupancy exceeds the smem budget");
+
+  template <std::size_t... I>
+  static constexpr auto make_table(std::index_sequence<I...>) {
+    return std::array<KernelFn, kMaxBankRows + 1>{nullptr, attn_res_fused_tma_kernel<Trait<I + 1>, kOccupancy>...};
+  }
+  static constexpr auto kTable = make_table(std::make_index_sequence<kMaxBankRows>{});
+  template <std::size_t... I>
+  static constexpr auto make_pull_table(std::index_sequence<I...>) {
+    return std::array<KernelFn, kMaxBankRows + 1>{nullptr, attn_res_fused_pull_rs_kernel<Trait<I + 1>, kOccupancy>...};
+  }
+  static constexpr auto kPullTable = make_pull_table(std::make_index_sequence<kMaxBankRows>{});
+  template <std::size_t... I>
+  static constexpr auto make_ag_table(std::index_sequence<I...>) {
+    return std::array<KernelFn, kMaxBankRows + 1>{
+        nullptr, attn_res_fused_direct_ag_kernel<Trait<I + 1>, kOccupancy>...};
+  }
+  static constexpr auto kAgTable = make_ag_table(std::make_index_sequence<kMaxBankRows>{});
+
+  static void
+  run(const tvm::ffi::TensorView prefix_sum,
+      const tvm::ffi::TensorView bank,
+      const tvm::ffi::TensorView cw,
+      const tvm::ffi::TensorView ow,
+      const tvm::ffi::TensorView out,
+      int64_t nvb,
+      double eps,
+      bool write_prefix) {
+    using namespace host;
+
+    auto T_ = SymbolicSize{"num_tokens"};
+    auto H_ = SymbolicSize{"hidden_size"};
+    auto NB_ = SymbolicSize{"num_bank_slots"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({T_, H_}).with_dtype<bf16_t>().with_device(device).verify(prefix_sum).verify(out);
+    TensorMatcher({T_, NB_, H_}).with_dtype<bf16_t>().with_device(device).verify(bank);
+    TensorMatcher({H_}).with_dtype<bf16_t>().with_device(device).verify(cw).verify(ow);
+
+    const auto num_tokens = static_cast<int64_t>(T_.unwrap());
+    const auto H = static_cast<int64_t>(H_.unwrap());
+    const auto NB = static_cast<int64_t>(NB_.unwrap());
+
+    RuntimeCheck(H == kDim, "attn_res_fused_tma: H must be ", kDim, ", got ", H);
+    RuntimeCheck(
+        1 <= nvb && nvb <= kMaxBankRows && nvb <= NB,
+        "attn_res_fused_tma: nvb must be in [1, ",
+        kMaxBankRows,
+        "] and <= NB, got nvb=",
+        nvb,
+        " NB=",
+        NB);
+    RuntimeCheck(
+        !write_prefix || nvb < NB,
+        "attn_res_fused_tma: write_prefix targets bank row nvb, needs nvb < NB, got nvb=",
+        nvb,
+        " NB=",
+        NB);
+
+    if (num_tokens == 0) return;
+
+    [[maybe_unused]] static const bool attrs_set = [] {
+      for (uint32_t i = 1; i <= kMaxBankRows; ++i) {
+        RuntimeDeviceCheck(cudaFuncSetAttribute(kTable[i], cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes));
+      }
+      return true;
+    }();
+
+    const auto num_sm = runtime::get_sm_count(device.unwrap().device_id);
+    const auto grid = std::min<int64_t>((int64_t)num_sm * kOccupancy, num_tokens);
+    const auto params = AttnResTMAParams{
+        .prefix_sum = static_cast<const bf16_t*>(prefix_sum.data_ptr()),
+        .bank = static_cast<const bf16_t*>(bank.data_ptr()),
+        .cw = static_cast<const bf16_t*>(cw.data_ptr()),
+        .ow = static_cast<const bf16_t*>(ow.data_ptr()),
+        .out = static_cast<bf16_t*>(out.data_ptr()),
+        .prefix_dst = write_prefix ? static_cast<bf16_t*>(bank.data_ptr()) + nvb * H : nullptr,
+        .input_mc = nullptr,
+        .residual = nullptr,
+        .prefix_out = nullptr,
+        .sem_local = nullptr,
+        .sem_mc = nullptr,
+        .output_mc = nullptr,
+        .world_size = 0,
+        .rank = 0,
+        .stride_bm = NB * H,
+        .eps = static_cast<float>(eps),
+        .num_tokens = static_cast<uint32_t>(num_tokens),
+    };
+    LaunchKernel(grid, kNumThreads, device.unwrap(), kSmemBytes).enable_pdl(true)(kTable[nvb], params);
+  }
+
+  static void run_pull_rs(
+      CommunicatorRef ref,
+      const tvm::ffi::TensorView input,
+      std::optional<tvm::ffi::TensorView> residual,
+      const tvm::ffi::TensorView bank,
+      const tvm::ffi::TensorView cw,
+      const tvm::ffi::TensorView ow,
+      const tvm::ffi::TensorView out,
+      const tvm::ffi::TensorView prefix_out,
+      int64_t nvb,
+      double eps,
+      int64_t input_mc_ptr,
+      int64_t sem_mc_ptr,
+      int64_t max_blocks) {
+    using namespace host;
+    const auto& data = *ref.get();
+    auto GT_ = SymbolicSize{"global_tokens"};
+    auto T_ = SymbolicSize{"local_tokens"};
+    auto H_ = SymbolicSize{"hidden_size"};
+    auto NB_ = SymbolicSize{"num_bank_slots"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({GT_, H_}).with_dtype<bf16_t>().with_device(device).verify(input);
+    TensorMatcher({T_, H_}).with_dtype<bf16_t>().with_device(device).verify(out).verify(prefix_out);
+    if (residual.has_value()) {
+      TensorMatcher({T_, H_}).with_dtype<bf16_t>().with_device(device).verify(residual.value());
+    }
+    TensorMatcher({T_, NB_, H_}).with_dtype<bf16_t>().with_device(device).verify(bank);
+    TensorMatcher({H_}).with_dtype<bf16_t>().with_device(device).verify(cw).verify(ow);
+
+    const auto global_tokens = static_cast<int64_t>(GT_.unwrap());
+    const auto num_tokens = static_cast<int64_t>(T_.unwrap());
+    const auto H = static_cast<int64_t>(H_.unwrap());
+    const auto NB = static_cast<int64_t>(NB_.unwrap());
+    RuntimeCheck(data.world_size > 1, "fused pull RS requires world_size > 1");
+    RuntimeCheck(global_tokens == num_tokens * data.world_size, "global tokens must equal local tokens * world size");
+    RuntimeCheck(H == kDim, "fused pull RS: H must be ", kDim, ", got ", H);
+    RuntimeCheck(1 <= nvb && nvb <= kMaxBankRows && nvb <= NB, "fused pull RS: invalid nvb=", nvb, " NB=", NB);
+    RuntimeCheck(input_mc_ptr != 0, "fused pull RS requires multicast input");
+    RuntimeCheck(sem_mc_ptr != 0, "fused pull RS requires multicast semaphores");
+    RuntimeCheck(max_blocks > 0, "fused pull RS requires max_blocks > 0");
+    if (num_tokens == 0) return;
+
+    [[maybe_unused]] static const bool attrs_set = [] {
+      for (uint32_t i = 1; i <= kMaxBankRows; ++i) {
+        RuntimeDeviceCheck(
+            cudaFuncSetAttribute(kPullTable[i], cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes));
+      }
+      return true;
+    }();
+
+    const auto num_sm = runtime::get_sm_count(device.unwrap().device_id);
+    const auto grid = std::min<int64_t>(
+        {static_cast<int64_t>(num_sm) * kOccupancy,
+         num_tokens,
+         max_blocks,
+         static_cast<int64_t>(data.num_pull_blocks)});
+    const auto local_elems = num_tokens * H;
+    const auto params = AttnResTMAParams{
+        .prefix_sum = static_cast<const bf16_t*>(prefix_out.data_ptr()),
+        .bank = static_cast<const bf16_t*>(bank.data_ptr()),
+        .cw = static_cast<const bf16_t*>(cw.data_ptr()),
+        .ow = static_cast<const bf16_t*>(ow.data_ptr()),
+        .out = static_cast<bf16_t*>(out.data_ptr()),
+        .prefix_dst = nullptr,
+        .input_mc = reinterpret_cast<const uint8_t*>(static_cast<uintptr_t>(input_mc_ptr)) +
+                    data.rank * local_elems * sizeof(bf16_t),
+        .residual = residual.has_value() ? static_cast<const bf16_t*>(residual.value().data_ptr()) : nullptr,
+        .prefix_out = static_cast<bf16_t*>(prefix_out.data_ptr()),
+        .sem_local = data.pull_semaphores[data.rank],
+        .sem_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(sem_mc_ptr)),
+        .output_mc = nullptr,
+        .world_size = data.world_size,
+        .rank = data.rank,
+        .stride_bm = NB * H,
+        .eps = static_cast<float>(eps),
+        .num_tokens = static_cast<uint32_t>(num_tokens),
+    };
+    LaunchKernel(grid, kNumThreads, device.unwrap(), kSmemBytes).enable_pdl(true)(kPullTable[nvb], params);
+  }
+
+  static void run_direct_ag(
+      CommunicatorRef ref,
+      const tvm::ffi::TensorView prefix_sum,
+      const tvm::ffi::TensorView bank,
+      const tvm::ffi::TensorView cw,
+      const tvm::ffi::TensorView ow,
+      const tvm::ffi::TensorView out,
+      int64_t nvb,
+      double eps,
+      int64_t output_mc_ptr,
+      int64_t sem_mc_ptr,
+      int64_t max_blocks,
+      bool write_prefix) {
+    using namespace host;
+    const auto& data = *ref.get();
+    auto T_ = SymbolicSize{"local_tokens"};
+    auto GT_ = SymbolicSize{"global_tokens"};
+    auto H_ = SymbolicSize{"hidden_size"};
+    auto NB_ = SymbolicSize{"num_bank_slots"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({T_, H_}).with_dtype<bf16_t>().with_device(device).verify(prefix_sum);
+    TensorMatcher({T_, NB_, H_}).with_dtype<bf16_t>().with_device(device).verify(bank);
+    TensorMatcher({H_}).with_dtype<bf16_t>().with_device(device).verify(cw).verify(ow);
+    TensorMatcher({GT_, H_}).with_dtype<bf16_t>().with_device(device).verify(out);
+
+    const auto num_tokens = static_cast<int64_t>(T_.unwrap());
+    const auto global_tokens = static_cast<int64_t>(GT_.unwrap());
+    const auto H = static_cast<int64_t>(H_.unwrap());
+    const auto NB = static_cast<int64_t>(NB_.unwrap());
+    RuntimeCheck(data.world_size > 1, "fused direct AG requires world_size > 1");
+    RuntimeCheck(global_tokens == num_tokens * data.world_size, "global tokens must equal local tokens * world size");
+    RuntimeCheck(H == kDim, "fused direct AG: H must be ", kDim, ", got ", H);
+    RuntimeCheck(1 <= nvb && nvb <= kMaxBankRows && nvb <= NB, "fused direct AG: invalid nvb=", nvb, " NB=", NB);
+    RuntimeCheck(!write_prefix || nvb < NB, "fused direct AG: write_prefix targets bank row nvb, needs nvb < NB");
+    RuntimeCheck(output_mc_ptr != 0, "fused direct AG requires multicast output");
+    RuntimeCheck(sem_mc_ptr != 0, "fused direct AG requires multicast semaphores");
+    RuntimeCheck(max_blocks > 0, "fused direct AG requires max_blocks > 0");
+    if (num_tokens == 0) return;
+
+    [[maybe_unused]] static const bool attrs_set = [] {
+      for (uint32_t i = 1; i <= kMaxBankRows; ++i) {
+        RuntimeDeviceCheck(cudaFuncSetAttribute(kAgTable[i], cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes));
+      }
+      return true;
+    }();
+
+    const auto num_sm = runtime::get_sm_count(device.unwrap().device_id);
+    const auto grid = std::min<int64_t>(
+        {static_cast<int64_t>(num_sm) * kOccupancy,
+         num_tokens,
+         max_blocks,
+         static_cast<int64_t>(data.num_pull_blocks)});
+    const auto params = AttnResTMAParams{
+        .prefix_sum = static_cast<const bf16_t*>(prefix_sum.data_ptr()),
+        .bank = static_cast<const bf16_t*>(bank.data_ptr()),
+        .cw = static_cast<const bf16_t*>(cw.data_ptr()),
+        .ow = static_cast<const bf16_t*>(ow.data_ptr()),
+        .out = static_cast<bf16_t*>(out.data_ptr()),
+        .prefix_dst = write_prefix ? static_cast<bf16_t*>(bank.data_ptr()) + nvb * H : nullptr,
+        .input_mc = nullptr,
+        .residual = nullptr,
+        .prefix_out = nullptr,
+        .sem_local = data.pull_semaphores[data.rank],
+        .sem_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(sem_mc_ptr)),
+        .output_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(output_mc_ptr)),
+        .world_size = data.world_size,
+        .rank = data.rank,
+        .stride_bm = NB * H,
+        .eps = static_cast<float>(eps),
+        .num_tokens = static_cast<uint32_t>(num_tokens),
+    };
+    LaunchKernel(grid, kNumThreads, device.unwrap(), kSmemBytes).enable_pdl(true)(kAgTable[nvb], params);
+  }
+};

python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh

@@ -0,0 +1,908 @@
+// K3 MNNVL fused all-reduce (bf16-only), two zero-copy algorithm families:
+//
+//   - push (1shot): lamport-style push, but the data staging is a SINGLE
+//     multicast store into slot `rank` of every peer's push workspace
+//     (replacing the 7 unicast stores of the generic kernel), followed by a
+//     local zero-marker polling reduce. Input is read in place and the
+//     result is written back in place — no staging copies. Works for any
+//     input tensor; reuses the CustomAllReduceV2 push workspace + counter.
+//     Best for small messages.
+//
+//   - pull (2shot): low-SM NVLS pull directly ON the input tensor, which
+//     therefore MUST live in (multicast-bound) symmetric memory: each rank
+//     `multimem.ld_reduce`s its shard from the input's multicast address and
+//     `multimem.st`s the result back — in-place, zero copies. Two
+//     NCCL-style ideas keep a handful of blocks (tuned 1~16, 4 at the large
+//     end) at the fabric limit:
+//
+//       * deep pipelining: the copy loop manually keeps `unroll` multimem
+//         loads in flight per thread before the first store, cascading
+//         through halving widths for mid-size tails.
+//       * multicast barriers on the CustomAllReduceV2 pull semaphores: the
+//         generic pull kernels' reservation protocol with each arrival sent
+//         as ONE `multimem.red.add` on the semaphore flag's multicast alias
+//         instead of per-peer unicast reds — identical memory effects, so
+//         both kernel families share the slots freely (single-stream calls
+//         are serialized).
+//
+// Both families fuse either a residual add (`out = allreduce(x) + r`; the
+// residual must be identical on every rank — a fully reduced tensor such as
+// the attn-res prefix sum — or absent) or the RMSNorm epilogue over the
+// latent of the K3 latent|shared MoE buffer (*_norm variants).
+//
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <cooperative_groups.h>
+
+// TODO: remove dependency on the custom_all_reduce, move out common utilities
+#include "../../distributed/custom_all_reduce.cuh"
+#include "ptx_sys.cuh"
+
+namespace sglang {
+
+// Same shape as gemm_ag / gemm_ar: pull the ptx_sys helpers in by name so the
+// call sites below stay unqualified.
+using device::distributed::multimem_red_add_relaxed;
+using device::distributed::multimem_red_add_release;
+
+struct FusionParams {
+  uint8_t* input;              // tensor pointer (in place)
+  const uint8_t* residual;     // may be null (compile-time kHasResidual selects)
+  uint8_t* push_ws_mc;         // multicast VA of the push workspace base
+  uint8_t* push_ws_local;      // local push workspace base (poll side)
+  Counter* push_counter;       // per-block phase counters (local memory)
+  int64_t push_buffer_stride;  // per-buffer bytes (2 * world_size buffers)
+  uint32_t rank;
+  uint32_t num_vecs;  // 16B vectors
+  // *_norm variants only: RMSNorm epilogue over the first num_norm_rows
+  // rows of the [num_vecs / kNormRowVecs, kNormDim] row view
+  const uint8_t* norm_weight;
+  float norm_eps;
+  uint32_t num_norm_rows;
+  uint32_t num_push_counters;  // cluster variant only: full counter array size
+  // finalize_push_norm only: trtllm-gen deferred-finalize inputs (kimi_k3.py
+  // deferred finalize path); `input` is then output-only ([T, kNormDim])
+  const uint8_t* fin_gemm2;    // [P, kNormDim] bf16, permuted rows
+  const uint8_t* fin_idx;      // [T * kFinTopK] int32, -1 = dropped slot
+  const uint8_t* fin_weights;  // [T, kFinTopK] bf16
+};
+
+// The *_norm variants view the input as rows of the K3 latent width (3584
+// bf16 = 448 16B vectors per row) and give the first num_norm_rows an
+// RMSNorm epilogue (K3: the [N latent | 2N shared] MoE buffer with N normed
+// rows, or a latent-only [N, 3584] tensor normed in full).
+constexpr uint32_t kNormDim = 3584;
+constexpr uint32_t kNormRowVecs = kNormDim / 8;                       // 448
+constexpr uint32_t kNormWarps = kNormRowVecs / device::kWarpThreads;  // 14
+
+template <uint32_t kWorldSize, bool kHasResidual, bool kUsePDL>
+__global__ __launch_bounds__(1024, 1) void all_reduce_push_res_kernel(const __grid_constant__ FusionParams params) {
+  using vec_t = device::AlignedVector<bf16x2_t, 4>;
+
+  const auto tx = threadIdx.x;
+  const auto bx = blockIdx.x;
+  const auto global_tid = bx * blockDim.x + tx;
+  const auto num_threads = blockDim.x * gridDim.x;
+  const auto num_vecs = params.num_vecs;
+
+  // prologue: the previous phase flip (counter inc) must be visible
+  device::PDLWaitPrimary<kUsePDL>();
+  const auto phase = params.push_counter[bx].get() & 1;
+  const auto r = params.rank;
+  const auto stride_bytes = params.push_buffer_stride;
+  const auto phase_stride_bytes = (phase * kWorldSize) * stride_bytes;
+  // one multicast store lands this rank's data in slot r of EVERY peer
+  const auto push_ptr = params.push_ws_mc + r * stride_bytes + phase_stride_bytes;
+  const auto poll_ptr = params.push_ws_local + phase_stride_bytes;
+
+  // stage 1: multicast-push local data, remapping all-zero bf16x2 pairs
+  static_assert(fp_trait<bf16_t>::pos_zero == 0, "the empty marker is all-zero bits");
+  constexpr uint32_t kNegZeroPair = 0x8000u;  // {-0.0, +0.0}: sum-neutral, non-zero
+  for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
+    vec_t vec;
+    ld_global_16B(vec, params.input, vid);
+    auto& bits = *reinterpret_cast<uint4*>(&vec);
+    if (bits.x == 0) bits.x = kNegZeroPair;
+    if (bits.y == 0) bits.y = kNegZeroPair;
+    if (bits.z == 0) bits.z = kNegZeroPair;
+    if (bits.w == 0) bits.w = kNegZeroPair;
+    st_multimem_16B(vec, push_ptr, vid);
+  }
+
+  // launch pdl early for low latency case
+  device::PDLTriggerSecondary<kUsePDL>();
+
+  // stage 2: poll all slots, reduce (+ residual), write back in place,
+  // re-establish the empty markers for the next same-phase round
+  vec_t zero_vec;
+  zero_vec.fill(bf16x2_t{get_pos_zero<bf16_t>(), get_pos_zero<bf16_t>()});
+  for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
+    vec_t vec[kWorldSize + kHasResidual];
+    if constexpr (kHasResidual) vec[kWorldSize].load(params.residual, vid);
+    do {
+      bool has_zero = false;
+#pragma unroll
+      for (uint32_t i = 0; i < kWorldSize; ++i) {
+        ld_relaxed_16B(vec[i], poll_ptr + i * stride_bytes, vid);
+        // the producer remapped all-zero pairs, so a written u32 is never
+        // 0: u32 == 0 <=> the 4B atom still holds the empty marker
+        const auto bits = *reinterpret_cast<const uint4*>(&vec[i]);
+        has_zero |= bits.x == 0;
+        has_zero |= bits.y == 0;
+        has_zero |= bits.z == 0;
+        has_zero |= bits.w == 0;
+      }
+      if (!has_zero) break;
+    } while (true);
+    const auto out_vec = reduce(vec);  // fp32 accumulation over 8(+1) inputs
+    st_global_16B(out_vec, params.input, vid);
+#pragma unroll
+    for (uint32_t i = 0; i < kWorldSize; ++i) {
+      st_global_16B(zero_vec, poll_ptr + i * stride_bytes, vid);
+    }
+  }
+
+  // epilogue: flip this block's phase
+  __syncthreads();
+  if (tx == 0) params.push_counter[bx].set(phase ^ 1);
+}
+
+// --- deferred-finalize staging (finalize_push_norm) ------------------------
+// The trtllm-gen MoE with do_finalize=False hands back its finalize inputs
+// (see kernels/ops/moe/trtllm_gen_moe.py); the fused kernel computes the finalize
+// during the push staging pass, so the rank-local latent never materializes.
+
+constexpr uint32_t kFinTopK = 16;
+
+// One 16B vector of the deferred MoE finalize (latent width fixed to kNormDim):
+//   local[t] = sum_k fin_weights[t, k] * fin_gemm2[fin_idx[t*16 + k]]
+// All 16 gathers issue before the FMA chain; threads of the same token
+// broadcast-load the same routing rows.
+SGL_DEVICE device::AlignedVector<bf16x2_t, 4> finalize_vec(const FusionParams& params, uint32_t vid) {
+  using namespace device;
+  constexpr uint32_t kIdxVecSize = kMaxVecBytes / sizeof(int32_t);
+  constexpr uint32_t kWVecSize = kMaxVecBytes / sizeof(bf16_t);
+  constexpr uint32_t kIdxVecs = kFinTopK / kIdxVecSize;  // 2 on SM100+
+  constexpr uint32_t kWVecs = kFinTopK / kWVecSize;      // 1 on SM100+
+
+  const uint32_t token = vid / kNormRowVecs;
+  const uint32_t hvec = vid % kNormRowVecs;
+
+  AlignedVector<int32_t, kIdxVecSize> idx[kIdxVecs];
+#pragma unroll
+  for (uint32_t j = 0; j < kIdxVecs; ++j) {
+    idx[j].load(params.fin_idx, token * kIdxVecs + j);
+  }
+  AlignedVector<bf16_t, kWVecSize> weight[kWVecs];
+#pragma unroll
+  for (uint32_t j = 0; j < kWVecs; ++j) {
+    weight[j].load(params.fin_weights, token * kWVecs + j);
+  }
+
+  const auto* g2 = reinterpret_cast<const bf16_t*>(params.fin_gemm2);
+  AlignedVector<bf16_t, 8> in[kFinTopK];
+#pragma unroll
+  for (uint32_t k = 0; k < kFinTopK; ++k) {
+    const int32_t row = idx[k / kIdxVecSize][k % kIdxVecSize];
+    if (row >= 0) {
+      in[k].load(g2 + static_cast<int64_t>(row) * kNormDim, hvec);
+    }
+  }
+  float acc[8] = {};
+#pragma unroll
+  for (uint32_t k = 0; k < kFinTopK; ++k) {
+    const int32_t row = idx[k / kIdxVecSize][k % kIdxVecSize];
+    if (row < 0) continue;
+    const bf16_t w_k = weight[k / kWVecSize][k % kWVecSize];
+#pragma unroll
+    for (uint32_t i = 0; i < 8; ++i) {
+      acc[i] = device::math::fma_f32_bf16(in[k][i], w_k, acc[i]);
+    }
+  }
+  AlignedVector<bf16x2_t, 4> out;
+#pragma unroll
+  for (uint32_t j = 0; j < 4; ++j) {
+    out[j] = cast<bf16x2_t>(fp32x2_t{acc[2 * j], acc[2 * j + 1]});
+  }
+  return out;
+}
+
+template <typename T2, size_t N, size_t M>
+SGL_DEVICE float reduce_sqr(device::AlignedVector<T2, N>& out_vec, device::AlignedVector<T2, N> (&vec)[M]) {
+  fp32x2_t acc_vec[N];
+#pragma unroll
+  for (size_t i = 0; i < M; ++i) {
+#pragma unroll
+    for (size_t j = 0; j < N; ++j) {
+      const auto [x, y] = device::cast<fp32x2_t>(vec[i][j]);
+      auto& [acc_x, acc_y] = acc_vec[j];
+      acc_x = i == 0 ? x : acc_x + x;
+      acc_y = i == 0 ? y : acc_y + y;
+    }
+  }
+  float sum_eq = 0.0f;
+#pragma unroll
+  for (size_t j = 0; j < N; ++j) {
+    sum_eq += acc_vec[j].x * acc_vec[j].x;
+    sum_eq += acc_vec[j].y * acc_vec[j].y;
+    out_vec[j] = device::cast<T2>(acc_vec[j]);
+  }
+  return sum_eq;
+}
+
+// kFinalize: stage 1 computes the deferred MoE finalize per vector instead of
+// reading a staged input tensor; `input` is then output-only. The host sets
+// num_norm_rows to the full row count (every reduced row is normed).
+template <uint32_t kWorldSize, uint32_t kClusterSize, bool kUsePDL, bool kFinalize = false>
+__global__ __launch_bounds__(kNormRowVecs / kClusterSize) __cluster_dims__(kClusterSize, 1, 1)  //
+    void all_reduce_push_norm_cluster_kernel(const __grid_constant__ FusionParams params) {
+  namespace cg = cooperative_groups;
+  using namespace device;
+  using vec_t = AlignedVector<bf16x2_t, 4>;
+  constexpr uint32_t kBlockSize = kNormRowVecs / kClusterSize;
+  constexpr uint32_t kNumWarps = kBlockSize / kWarpThreads;
+
+  static_assert(kBlockSize % kWarpThreads == 0);
+  static_assert(kNormRowVecs % kClusterSize == 0);
+  static_assert(kNumWarps >= 1);
+
+  const auto tx = threadIdx.x;
+  const auto bx = blockIdx.x;
+  const auto global_tid = bx * kBlockSize + tx;
+  const auto num_vecs = params.num_vecs;
+  const auto num_rows = num_vecs / kNormRowVecs;
+
+  const auto row_idx = bx / kClusterSize;
+  const auto num_row_clusters = gridDim.x / kClusterSize - 1;  // last one is the bumper
+  // stage-1 grid-stride is over the ROW clusters only: the bumper early-returns
+  // and never stages, so it must not be counted or its share of vids is dropped
+  const auto num_threads = kBlockSize * num_row_clusters * kClusterSize;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  // special case: the bumper cluster flips every remaining counter so the
+  // whole array stays globally in phase. Only ONE block does it (threads
+  // grid-stride the counters) — each counter must be inc'd exactly once,
+  // independent of whether kClusterSize is odd or even.
+  if (row_idx == num_row_clusters) {
+    if (bx % kClusterSize == 0) {
+      for (uint32_t r = num_row_clusters + tx; r < params.num_push_counters; r += kBlockSize) {
+        params.push_counter[r].inc(1);
+      }
+    }
+    return PDLTriggerSecondary<kUsePDL>();
+  }
+
+  const auto phase = params.push_counter[row_idx].get() & 1;
+  const auto r = params.rank;
+  const auto stride_bytes = params.push_buffer_stride;
+  const auto phase_stride_bytes = (phase * kWorldSize) * stride_bytes;
+  const auto push_ptr = params.push_ws_mc + r * stride_bytes + phase_stride_bytes;
+  const auto poll_ptr = params.push_ws_local + phase_stride_bytes;
+
+  // stage 1: multicast staging (grid-stride); kFinalize computes each vector
+  // in place of the load
+  static_assert(fp_trait<bf16_t>::pos_zero == 0, "the empty marker is all-zero bits");
+  for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
+    vec_t vec;
+    if constexpr (kFinalize) {
+      vec = finalize_vec(params, vid);
+    } else {
+      ld_global_16B(vec, params.input, vid);
+    }
+    auto& bits = *reinterpret_cast<uint4*>(&vec);
+    if (bits.x == 0) bits.x = fp_trait<bf16_t>::neg_zero;
+    if (bits.y == 0) bits.y = fp_trait<bf16_t>::neg_zero;
+    if (bits.z == 0) bits.z = fp_trait<bf16_t>::neg_zero;
+    if (bits.w == 0) bits.w = fp_trait<bf16_t>::neg_zero;
+    st_multimem_16B(vec, push_ptr, vid);
+  }
+
+  // stage 2: one row per cluster pass (the bumper cluster owns no rows)
+  const auto cluster = cg::this_cluster();
+  const auto cluster_rank = bx % kClusterSize;
+  vec_t w;
+  w.load(params.norm_weight, cluster_rank * kBlockSize + tx);
+  vec_t zero_vec;
+  zero_vec.fill(bf16x2_t{get_pos_zero<bf16_t>(), get_pos_zero<bf16_t>()});
+  __shared__ alignas(8) float smem_raw[2][kClusterSize][kNumWarps];
+  uint32_t parity = 0;
+
+  // NOTE: launch PDL earlier for low latency case
+  PDLTriggerSecondary<kUsePDL>();
+
+  for (auto row = row_idx; row < num_rows; row += num_row_clusters) {
+    const auto vid = row * kNormRowVecs + cluster_rank * kBlockSize + tx;
+    vec_t vec[kWorldSize];
+    do {
+      bool has_zero = false;
+#pragma unroll
+      for (uint32_t i = 0; i < kWorldSize; ++i) {
+        ld_relaxed_16B(vec[i], poll_ptr + i * stride_bytes, vid);
+        const auto bits = *reinterpret_cast<const uint4*>(&vec[i]);
+        has_zero |= bits.x == 0;
+        has_zero |= bits.y == 0;
+        has_zero |= bits.z == 0;
+        has_zero |= bits.w == 0;
+      }
+      if (!has_zero) break;
+    } while (true);
+
+    vec_t out_vec;
+    if (row < params.num_norm_rows) {  // cluster-uniform branch
+      auto& smem = smem_raw[parity];
+      parity ^= 1;
+      // push each warp's partial to EVERY peer's slot for this block: lane p
+      // (p < kClusterSize) targets peer p, warp w selects the [.][w] slot. So
+      // after the barrier every block holds all kClusterSize*kNumWarps
+      // partials in its own smem and reduces them locally — the read side
+      // never touches remote DSMEM, so no post-read barrier is needed to guard
+      // against a peer CTA exiting (parity double-buffers across rows).
+      const auto lane = tx % kWarpThreads;
+      const auto warp = tx / kWarpThreads;
+      const auto warp_sqr = warp::reduce_sum(reduce_sqr(out_vec, vec));
+      if (lane < kClusterSize) {
+        float* dst = cluster.map_shared_rank(&smem[cluster_rank][warp], lane);
+        *dst = warp_sqr;
+      }
+      cluster.sync();
+      // load local
+      float total = 0.0f;
+#pragma unroll
+      for (uint32_t r = 0; r < kClusterSize; ++r) {
+        using vec_t = AlignedVector<float, kNumWarps>;
+        vec_t remote_value;
+        remote_value.load(smem[r]);
+#pragma unroll
+        for (uint32_t w = 0; w < kNumWarps; ++w) {
+          total += remote_value[w];
+        }
+      }
+      const auto norm_factor = math::rsqrt(total / kNormDim + params.norm_eps);
+#pragma unroll
+      for (uint32_t j = 0; j < 4; ++j) {
+        const auto [a, b] = cast<fp32x2_t>(out_vec[j]);
+        const auto [wa, wb] = cast<fp32x2_t>(w[j]);
+        out_vec[j] = cast<bf16x2_t>(fp32x2_t{a * norm_factor * wa, b * norm_factor * wb});
+      }
+    } else {
+      out_vec = reduce(vec);
+    }
+
+    st_global_16B(out_vec, params.input, vid);
+#pragma unroll
+    for (uint32_t i = 0; i < kWorldSize; ++i) {
+      st_global_16B(zero_vec, poll_ptr + i * stride_bytes, vid);
+    }
+  }
+
+  // epilogue: each row cluster flips its own counter; the bumper cluster
+  // flips every remaining one so the whole array stays globally uniform
+  if (cluster_rank == 0 && tx == 0) {
+    params.push_counter[row_idx].set(phase ^ 1);
+  }
+}
+
+// Pull family: low-SM NVLS, reusing the CustomAllReduceV2 pull semaphores
+
+inline constexpr uint32_t kPullBlockSize = 512;
+
+struct PullParams {
+  uint8_t* input_mc;        // multicast VA of the symmetric input
+  const uint8_t* residual;  // may be null (compile-time kHasResidual selects)
+  Semaphore* sem_local;     // this rank's v2 pull semaphores (poll side)
+  uint8_t* sem_mc;          // multicast VA of the pull-semaphore region
+  uint32_t rank;
+  uint32_t world_size;
+  uint32_t num_vecs;  // 16B vectors
+  // pull_norm only: RMSNorm epilogue over the first num_norm_rows rows of
+  // the [num_vecs / kNormRowVecs, kNormDim] row view
+  const uint8_t* norm_weight;
+  float norm_eps;
+  uint32_t num_norm_rows;
+};
+
+// K3 pull barriers reuse the v2 pull-semaphore slots with EXACTLY the
+// generic kernels' reservation protocol — reserve a 2 * world_size flag
+// window on the local m_counter, signal arrival on m_flag, wait for the
+// window to fill — except that each arrival is ONE `multimem.red.add` on
+// the m_flag's multicast alias instead of world_size per-peer unicast reds
+// (same aggregate effect: every rank's flag gains world_size arrivals per
+// phase). Identical memory effects per call, so both kernel families share
+// the slots freely (single-stream calls are serialized).
+//
+// The multicast alias of Semaphore::m_flag (the struct's first member):
+SGL_DEVICE uint32_t* pull_sem_mc_flag(uint8_t* sem_mc, uint32_t block) {
+  static_assert(sizeof(Semaphore) == 128);
+  return reinterpret_cast<uint32_t*>(sem_mc + block * sizeof(Semaphore));
+}
+
+// enter barrier (relaxed): reserve this call's flag window, signal arrival
+// with one multicast red, poll the local flag until all world_size arrivals
+// landed — every rank's producer has finished writing the input. The
+// reservation atomicAdd sits BEFORE the PDL wait: it is safe there (the
+// previous same-slot call's reservation completed at ITS enter, which
+// precedes its launch_dependents and hence this kernel's start, so windows
+// are handed out in stream order) and it keeps the RMW latency off the
+// post-wait critical path. The red must stay AFTER the wait — it asserts
+// the producer grid has flushed. Returns the window base for the exit
+// barrier — meaningful in thread 0 only, the sole barrier poller.
+template <bool kUsePDL>
+SGL_DEVICE uint32_t pull_barrier_enter(const PullParams& params) {
+  uint32_t current = 0;
+  if (threadIdx.x == 0) {
+    const auto semaphore = &params.sem_local[blockIdx.x];
+    const auto reserved = semaphore->counter_ptr()->inc(2 * params.world_size);
+    current = reserved + params.world_size;
+    device::PDLWaitPrimary<kUsePDL>();
+    multimem_red_add_relaxed(pull_sem_mc_flag(params.sem_mc, blockIdx.x));
+    while (semaphore->get_relaxed() - reserved < params.world_size)
+      ;
+  }
+  __syncthreads();
+  return current;
+}
+
+// exit barrier (release/acquire): every peer has finished reading my buffer
+// (and, for 2shot, its broadcast into it is visible) before my next kernel
+// may touch it. Mirrors AllReducePullImpl::sync_exit_pull<true>.
+template <bool kUsePDL>
+SGL_DEVICE void pull_barrier_exit(const PullParams& params, uint32_t current) {
+  device::PDLTriggerSecondary<kUsePDL>();
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    const auto semaphore = &params.sem_local[blockIdx.x];
+    multimem_red_add_release(pull_sem_mc_flag(params.sem_mc, blockIdx.x));
+    while (semaphore->get_acquire() - current < params.world_size)
+      ;
+  }
+}
+
+// One pipelined pass at width kWidth (kWidth multimem loads in flight per
+// thread), then recurse to kWidth/2 for the remainder, down to a plain
+// unroll-1 tail. The cascade matters for mid sizes: a shard smaller than
+// kUnroll * step would otherwise skip the main loop entirely and run the
+// whole range with a single request in flight.
+template <uint32_t kWidth, bool kHasResidual>
+SGL_DEVICE void
+pull_reduce_pass(uint32_t& vid, const uint32_t num_vecs, const uint32_t step, uint8_t* mc_ptr, const uint8_t* res_ptr) {
+  using vec_t = device::AlignedVector<bf16x2_t, 4>;
+  using SumOp = device::ReductionTrait<device::ReductionOp::SUM, bf16x2_t>;
+  for (; vid + (kWidth - 1) * step < num_vecs; vid += kWidth * step) {
+    vec_t vec[kWidth];
+#pragma unroll
+    for (uint32_t u = 0; u < kWidth; ++u) {
+      ld_multimem_16B(vec[u], mc_ptr, vid + u * step);
+    }
+    if constexpr (kHasResidual) {
+#pragma unroll
+      for (uint32_t u = 0; u < kWidth; ++u) {
+        vec_t res_vec;
+        res_vec.load(res_ptr, vid + u * step);
+#pragma unroll
+        for (uint32_t j = 0; j < 4; ++j) {
+          vec[u][j] = SumOp::reduce(vec[u][j], res_vec[j]);
+        }
+      }
+    }
+#pragma unroll
+    for (uint32_t u = 0; u < kWidth; ++u) {
+      st_multimem_16B(vec[u], mc_ptr, vid + u * step);
+    }
+  }
+  if constexpr (kWidth > 1) {
+    pull_reduce_pass<kWidth / 2, kHasResidual>(vid, num_vecs, step, mc_ptr, res_ptr);
+  }
+}
+
+template <uint32_t kUnroll, bool kHasResidual, bool kUsePDL>
+__global__
+__launch_bounds__(kPullBlockSize, 1) void all_reduce_pull_res_kernel(const __grid_constant__ PullParams params) {
+  static_assert(1 <= kUnroll && kUnroll <= 16 && (kUnroll & (kUnroll - 1)) == 0);
+
+  const auto tx = threadIdx.x;
+  const auto bx = blockIdx.x;
+  const auto barrier_window = pull_barrier_enter<kUsePDL>(params);
+
+  // this rank's shard of the 16B-vector range
+  const auto r = params.rank;
+  const auto avg_vecs = params.num_vecs / params.world_size;
+  const auto rem_vecs = params.num_vecs % params.world_size;
+  const auto vec_bias = int64_t(avg_vecs) * r + min(r, rem_vecs);
+  const auto num_vecs = avg_vecs + (r < rem_vecs ? 1 : 0);
+  const auto mc_ptr = params.input_mc + vec_bias * 16;
+  const auto res_ptr = kHasResidual ? params.residual + vec_bias * 16 : nullptr;
+
+  // deep-pipelined body: issue up to kUnroll multimem loads back-to-back
+  // before the first (residual add and) store, keeping kUnroll requests in
+  // flight per thread to hide the NVLink latency with very few blocks; the
+  // remainder cascades through halving widths down to 1.
+  const auto step = kPullBlockSize * gridDim.x;
+  auto vid = bx * kPullBlockSize + tx;
+  pull_reduce_pass<kUnroll, kHasResidual>(vid, num_vecs, step, mc_ptr, res_ptr);
+
+  pull_barrier_exit<kUsePDL>(params, barrier_window);
+}
+
+// Fused RMSNorm over the latent of the K3 latent|shared MoE buffer: the
+// row-structured counterpart of the res kernel with kUnroll ROWS in flight
+// per block. One block pass covers kUnroll consecutive rows: 448 threads
+// each ld_reduce one 16B vector per row back-to-back, warp partials of the
+// normed rows go to smem (one __syncthreads per group), non-norm rows store
+// immediately. smem is parity double-buffered across groups so the next
+// group's partial writes can't race this group's reads (the WAR pair is two
+// groups apart and separated by the intervening group's barrier). Works on
+// this rank's row shard, in place.
+template <uint32_t kUnroll, bool kUsePDL>
+__global__
+__launch_bounds__(kNormRowVecs, 1) void all_reduce_pull_norm_kernel(const __grid_constant__ PullParams params) {
+  using vec_t = device::AlignedVector<bf16x2_t, 4>;
+  using namespace device;
+
+  const auto tx = threadIdx.x;
+  const auto bx = blockIdx.x;
+  const auto barrier_window = pull_barrier_enter<kUsePDL>(params);
+
+  // this rank's shard of the row range
+  const auto num_rows = params.num_vecs / kNormRowVecs;
+  const auto r = params.rank;
+  const auto avg_rows = num_rows / params.world_size;
+  const auto rem_rows = num_rows % params.world_size;
+  const auto row_bias = avg_rows * r + min(r, rem_rows);
+  const auto my_rows = avg_rows + (r < rem_rows ? 1 : 0);
+
+  vec_t wvec;
+  wvec.load(params.norm_weight, tx);
+  __shared__ float smem[2][kUnroll][kNormWarps];
+  const auto warp = tx / kWarpThreads;
+  const auto lane = tx % kWarpThreads;
+  uint32_t parity = 0;
+
+  const auto row_step = gridDim.x * kUnroll;
+  for (auto row0 = bx * kUnroll; row0 < my_rows; row0 += row_step) {
+    const auto cnt = min(kUnroll, my_rows - row0);
+    const auto vid0 = int64_t(row_bias + row0) * kNormRowVecs + tx;
+    vec_t vec[kUnroll];
+#pragma unroll
+    for (uint32_t u = 0; u < kUnroll; ++u) {
+      if (u < cnt) ld_multimem_16B(vec[u], params.input_mc, vid0 + u * kNormRowVecs);
+    }
+    // norm rows: push this warp's partial sum of squares to smem; non-norm
+    // rows (the shared 2/3) don't wait for the barrier — store right away
+    auto& sm = smem[parity];
+    parity ^= 1;
+#pragma unroll
+    for (uint32_t u = 0; u < kUnroll; ++u) {
+      if (u >= cnt) continue;  // block-uniform
+      if (row_bias + row0 + u < params.num_norm_rows) {
+        float sum_of_squares = 0.0f;
+#pragma unroll
+        for (uint32_t j = 0; j < 4; ++j) {
+          const auto [a, b] = cast<fp32x2_t>(vec[u][j]);
+          sum_of_squares += a * a + b * b;
+        }
+        sum_of_squares = warp::reduce_sum(sum_of_squares);
+        if (lane == 0) sm[u][warp] = sum_of_squares;
+      } else {
+        st_multimem_16B(vec[u], params.input_mc, vid0 + u * kNormRowVecs);
+      }
+    }
+    __syncthreads();
+#pragma unroll
+    for (uint32_t u = 0; u < kUnroll; ++u) {
+      if (u >= cnt || row_bias + row0 + u >= params.num_norm_rows) continue;  // block-uniform
+      float total = 0.0f;
+#pragma unroll
+      for (uint32_t w = 0; w < kNormWarps; ++w) {
+        total += sm[u][w];
+      }
+      const auto norm_factor = math::rsqrt(total / kNormDim + params.norm_eps);
+#pragma unroll
+      for (uint32_t j = 0; j < 4; ++j) {
+        const auto [a, b] = cast<fp32x2_t>(vec[u][j]);
+        const auto [wa, wb] = cast<fp32x2_t>(wvec[j]);
+        vec[u][j] = cast<bf16x2_t>(fp32x2_t{a * norm_factor * wa, b * norm_factor * wb});
+      }
+      st_multimem_16B(vec[u], params.input_mc, vid0 + u * kNormRowVecs);
+    }
+  }
+
+  pull_barrier_exit<kUsePDL>(params, barrier_window);
+}
+
+}  // namespace sglang
+
+using namespace sglang;
+
+// Host entry points
+
+template <uint32_t kWorldSize, bool kUsePDL>
+struct AllReduceFusionKernel {
+ private:
+  using TensorView = tvm::ffi::TensorView;
+
+  template <bool kHasResidual>
+  static constexpr auto res_push_kernel = all_reduce_push_res_kernel<kWorldSize, kHasResidual, kUsePDL>;
+
+  static FusionParams
+  make_params(const host::distributed::CommunicatorObj& data, TensorView input, std::optional<TensorView> residual) {
+    using namespace host;
+    SymbolicSize N = {"num_elements"};
+    SymbolicDevice device;
+    device.set_options<kDLCUDA>();
+    if (residual.has_value()) {
+      TensorMatcher({N})  //
+          .with_dtype<bf16_t>()
+          .with_device<kDLCUDA>(device)
+          .verify(input)
+          .verify(residual.value());
+    } else {
+      TensorMatcher({N})  //
+          .with_dtype<bf16_t>()
+          .with_device<kDLCUDA>(device)
+          .verify(input);
+    }
+    const auto num_elems = N.unwrap();
+    CHECK_HOST(data.world_size == kWorldSize);
+    CHECK_HOST(num_elems > 0 && num_elems % 8 == 0);
+    FusionParams params{};
+    params.input = static_cast<uint8_t*>(input.data_ptr());
+    params.residual = residual.has_value() ? static_cast<const uint8_t*>(residual.value().data_ptr()) : nullptr;
+    params.push_ws_mc = nullptr;
+    params.push_ws_local = data.push_workspaces[data.rank];
+    params.push_counter = data.push_counter;
+    params.push_buffer_stride = data.push_bytes;
+    params.rank = data.rank;
+    params.num_vecs = static_cast<uint32_t>(num_elems / 8);
+    return params;
+  }
+
+  /// The input viewed as rows of the norm width (3584 bf16 each); the first
+  /// num_norm_rows get the RMSNorm epilogue, the rest are a plain allreduce
+  /// (K3 uses [N latent rows | 2N shared rows] with num_norm_rows = N, or a
+  /// latent-only [N, 3584] tensor with num_norm_rows = N).
+  static FusionParams make_params_norm(
+      const host::distributed::CommunicatorObj& data,
+      TensorView input,
+      TensorView weight,
+      float eps,
+      int64_t num_norm_rows) {
+    using namespace host;
+    auto params = make_params(data, input, std::nullopt);
+    SymbolicDevice device;
+    device.set_options<kDLCUDA>();
+    TensorMatcher({kNormDim}).with_dtype<bf16_t>().with_device<kDLCUDA>(device).verify(weight);
+    CHECK_HOST(params.num_vecs % kNormRowVecs == 0)
+        << "numel must be a multiple of " << kNormDim << ", got " << int64_t(params.num_vecs) * 8;
+    const auto num_rows = params.num_vecs / kNormRowVecs;
+    CHECK_HOST(0 <= num_norm_rows && num_norm_rows <= num_rows)
+        << "num_norm_rows " << num_norm_rows << " out of range [0, " << num_rows << "]";
+    params.norm_weight = static_cast<const uint8_t*>(weight.data_ptr());
+    params.norm_eps = static_cast<float>(eps);
+    params.num_norm_rows = static_cast<uint32_t>(num_norm_rows);
+    params.num_push_counters = data.num_push_blocks;
+    return params;
+  }
+
+  // Shared pull validation; the reduce is in place on the symmetric input.
+  static PullParams make_pull_params(
+      const host::distributed::CommunicatorObj& data,
+      TensorView input,
+      std::optional<TensorView> residual,
+      int64_t input_mc_ptr,
+      int64_t sem_mc_ptr) {
+    using namespace host;
+    SymbolicSize N = {"num_elements"};
+    SymbolicDevice device;
+    device.set_options<kDLCUDA>();
+    if (residual.has_value()) {
+      TensorMatcher({N})  //
+          .with_dtype<bf16_t>()
+          .with_device<kDLCUDA>(device)
+          .verify(input)
+          .verify(residual.value());
+    } else {
+      TensorMatcher({N})  //
+          .with_dtype<bf16_t>()
+          .with_device<kDLCUDA>(device)
+          .verify(input);
+    }
+    const auto num_elems = N.unwrap();
+    CHECK_HOST(data.world_size == kWorldSize);
+    CHECK_HOST(num_elems > 0 && num_elems % 8 == 0) << "numel must be a positive multiple of 8, got " << num_elems;
+    // headroom below 2^32 so the unrolled loop's `vid + (kUnroll-1)*step`
+    // arithmetic can never wrap around u32
+    CHECK_HOST(num_elems / 8 < (int64_t(1) << 31)) << "numel exceeds the 16B-vector limit";
+    CHECK_HOST(input_mc_ptr != 0) << "pull requires the input's multicast address";
+    CHECK_HOST(sem_mc_ptr != 0) << "pull requires the semaphores' multicast address";
+    PullParams params{};
+    params.input_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(input_mc_ptr));
+    params.residual = residual.has_value() ? static_cast<const uint8_t*>(residual.value().data_ptr()) : nullptr;
+    params.sem_local = data.pull_semaphores[data.rank];
+    params.sem_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(sem_mc_ptr));
+    params.rank = data.rank;
+    params.world_size = data.world_size;
+    params.num_vecs = static_cast<uint32_t>(num_elems / 8);
+    return params;
+  }
+
+  // Runtime unroll dispatch: every supported width is compiled into the
+  // module so the tuned per-size unroll needs no extra JIT builds.
+  template <bool kHasResidual>
+  static void launch_pull_res(const PullParams& params, int64_t num_blocks, int64_t unroll, DLDevice device) {
+    const auto run = [&](auto kernel) {
+      host::LaunchKernel(static_cast<uint32_t>(num_blocks), kPullBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
+    };
+    switch (unroll) {
+      case 2:
+        return run(all_reduce_pull_res_kernel<2, kHasResidual, kUsePDL>);
+      case 4:
+        return run(all_reduce_pull_res_kernel<4, kHasResidual, kUsePDL>);
+      case 8:
+        return run(all_reduce_pull_res_kernel<8, kHasResidual, kUsePDL>);
+      case 16:
+        return run(all_reduce_pull_res_kernel<16, kHasResidual, kUsePDL>);
+      default:
+        CHECK_HOST(false) << "unsupported unroll " << unroll << " (must be 2, 4, 8, or 16)";
+    }
+  }
+
+  static void launch_pull_norm(const PullParams& params, int64_t num_blocks, int64_t unroll, DLDevice device) {
+    const auto run = [&](auto kernel) {
+      host::LaunchKernel(static_cast<uint32_t>(num_blocks), kNormRowVecs, device).enable_pdl(kUsePDL)(kernel, params);
+    };
+    switch (unroll) {
+      case 2:
+        return run(all_reduce_pull_norm_kernel<2, kUsePDL>);
+      case 4:
+        return run(all_reduce_pull_norm_kernel<4, kUsePDL>);
+      case 8:
+        return run(all_reduce_pull_norm_kernel<8, kUsePDL>);
+      case 16:
+        return run(all_reduce_pull_norm_kernel<16, kUsePDL>);
+      default:
+        CHECK_HOST(false) << "unsupported unroll " << unroll << " (must be 2, 4, 8, or 16)";
+    }
+  }
+
+ public:
+  static void push_res(CommunicatorRef ref, TensorView input, std::optional<TensorView> residual, int64_t ws_mc_base) {
+    const auto& data = *ref.get();
+    auto params = make_params(data, input, residual);
+    CHECK_HOST(ws_mc_base != 0) << "push requires a multicast-capable workspace";
+    const int64_t nbytes = int64_t(params.num_vecs) * 16;
+    CHECK_HOST(nbytes <= data.push_bytes)
+        << "input size " << nbytes << " exceeds push workspace size " << data.push_bytes;
+    params.push_ws_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(ws_mc_base));
+    const auto kernel = residual.has_value() ? res_push_kernel<true> : res_push_kernel<false>;
+    host::LaunchKernel(data.num_push_blocks, choose_block_size(params.num_vecs), input.device())
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+
+  static void push_norm(
+      CommunicatorRef ref, TensorView input, TensorView weight, float eps, int64_t num_norm_rows, int64_t ws_mc_base) {
+    constexpr auto kClusterSize = 7;
+    const auto& data = *ref.get();
+    auto params = make_params_norm(data, input, weight, eps, num_norm_rows);
+    CHECK_HOST(ws_mc_base != 0) << "push requires a multicast-capable workspace";
+    const int64_t nbytes = int64_t(params.num_vecs) * 16;
+    CHECK_HOST(nbytes <= data.push_bytes)
+        << "input size " << nbytes << " exceeds push workspace size " << data.push_bytes;
+    params.push_ws_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(ws_mc_base));
+    const auto num_rows = params.num_vecs / kNormRowVecs;
+    constexpr uint32_t kMaxClusters = 96;
+    const auto num_row_clusters = std::max<uint32_t>(std::min(num_rows, kMaxClusters), 1);
+    CHECK_HOST(num_row_clusters < data.num_push_blocks);
+    host::LaunchKernel((num_row_clusters + 1) * kClusterSize, kNormRowVecs / kClusterSize, input.device())
+        .enable_pdl(kUsePDL)(all_reduce_push_norm_cluster_kernel<kWorldSize, kClusterSize, kUsePDL>, params);
+  }
+
+  /// Deferred MoE finalize + 1shot push all-reduce + RMSNorm over EVERY row.
+  /// `out` (flattened [num_tokens * kNormDim] bf16) is output-only: each
+  /// rank's partial latent is computed from the trtllm-gen deferred-finalize
+  /// triple during the staging pass and never materializes in global memory.
+  static void finalize_push_norm(
+      CommunicatorRef ref,
+      TensorView out,
+      TensorView gemm2_out,
+      TensorView permuted_idx,
+      TensorView expert_weights,
+      TensorView weight,
+      float eps,
+      int64_t ws_mc_base) {
+    using namespace host;
+    constexpr auto kClusterSize = 7;
+    const auto& data = *ref.get();
+    // every row of the latent-only output is normed
+    auto params = make_params_norm(data, out, weight, eps, out.size(0) / kNormDim);
+    const auto num_tokens = params.num_vecs / kNormRowVecs;
+
+    auto P = SymbolicSize{"num_permuted_rows"};
+    auto T = SymbolicSize{"num_tokens"};
+    T.set_value(num_tokens);
+    auto K = SymbolicSize{"top_k"};
+    auto TK = SymbolicSize{"num_expanded"};
+    TK.set_value(static_cast<int64_t>(num_tokens) * kFinTopK);
+    SymbolicDevice device;
+    device.set_options<kDLCUDA>();
+    TensorMatcher({P, kNormDim}).with_dtype<bf16_t>().with_device<kDLCUDA>(device).verify(gemm2_out);
+    TensorMatcher({T, K}).with_dtype<bf16_t>().with_device<kDLCUDA>(device).verify(expert_weights);
+    TensorMatcher({TK}).with_dtype<int32_t>().with_device<kDLCUDA>(device).verify(permuted_idx);
+    CHECK_HOST(K.unwrap() == kFinTopK) << "finalize_push_norm is specialized for top_k = " << kFinTopK;
+
+    CHECK_HOST(ws_mc_base != 0) << "push requires a multicast-capable workspace";
+    const int64_t nbytes = int64_t(params.num_vecs) * 16;
+    CHECK_HOST(nbytes <= data.push_bytes)
+        << "output size " << nbytes << " exceeds push workspace size " << data.push_bytes;
+    params.push_ws_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(ws_mc_base));
+    params.fin_gemm2 = static_cast<const uint8_t*>(gemm2_out.data_ptr());
+    params.fin_idx = static_cast<const uint8_t*>(permuted_idx.data_ptr());
+    params.fin_weights = static_cast<const uint8_t*>(expert_weights.data_ptr());
+
+    constexpr uint32_t kMaxClusters = 96;
+    const auto num_row_clusters = std::max<uint32_t>(std::min(num_tokens, kMaxClusters), 1);
+    CHECK_HOST(num_row_clusters < data.num_push_blocks);
+    host::LaunchKernel((num_row_clusters + 1) * kClusterSize, kNormRowVecs / kClusterSize, out.device())
+        .enable_pdl(kUsePDL)(
+            all_reduce_push_norm_cluster_kernel<kWorldSize, kClusterSize, kUsePDL, /*kFinalize=*/true>, params);
+  }
+
+  /// Low-SM NVLS pull (+ optional residual): in-place reduce-scatter +
+  /// broadcast on the symmetric input. `sem_mc_ptr` is the multicast VA of
+  /// the v2 pull-semaphore region; num_blocks — which must be uniform
+  /// across ranks per call — is clamped to the semaphore capacity.
+  static void pull_res(
+      CommunicatorRef ref,
+      TensorView input,
+      std::optional<TensorView> residual,
+      int64_t input_mc_ptr,
+      int64_t sem_mc_ptr,
+      int64_t num_blocks,
+      int64_t unroll) {
+    const auto& data = *ref.get();
+    const auto params = make_pull_params(data, input, residual, input_mc_ptr, sem_mc_ptr);
+    CHECK_HOST(num_blocks >= 1) << "invalid num_blocks: " << num_blocks;
+    num_blocks = std::min<int64_t>(num_blocks, data.num_pull_blocks);
+    if (residual.has_value()) {
+      launch_pull_res<true>(params, num_blocks, unroll, input.device());
+    } else {
+      launch_pull_res<false>(params, num_blocks, unroll, input.device());
+    }
+  }
+
+  /// Low-SM NVLS pull + RMSNorm over the latent of the K3 latent|shared MoE
+  /// buffer ([num_tokens, 3584] latent then [num_tokens, 7168] shared);
+  /// num_tokens and the normed row range are derived from the element count.
+  /// Same semaphore / num_blocks semantics as pull_res.
+  static void pull_norm(
+      CommunicatorRef ref,
+      TensorView input,
+      TensorView weight,
+      double eps,
+      int64_t num_norm_rows,
+      int64_t input_mc_ptr,
+      int64_t sem_mc_ptr,
+      int64_t num_blocks,
+      int64_t unroll) {
+    const auto& data = *ref.get();
+    auto params = make_pull_params(data, input, std::nullopt, input_mc_ptr, sem_mc_ptr);
+    CHECK_HOST(num_blocks >= 1) << "invalid num_blocks: " << num_blocks;
+    num_blocks = std::min<int64_t>(num_blocks, data.num_pull_blocks);
+    using namespace host;
+    SymbolicDevice device;
+    device.set_options<kDLCUDA>();
+    TensorMatcher({kNormDim}).with_dtype<bf16_t>().with_device<kDLCUDA>(device).verify(weight);
+    CHECK_HOST(params.num_vecs % kNormRowVecs == 0)
+        << "numel must be a multiple of " << kNormDim << ", got " << int64_t(params.num_vecs) * 8;
+    const auto num_rows = params.num_vecs / kNormRowVecs;
+    CHECK_HOST(0 <= num_norm_rows && num_norm_rows <= num_rows)
+        << "num_norm_rows " << num_norm_rows << " out of range [0, " << num_rows << "]";
+    params.norm_weight = static_cast<const uint8_t*>(weight.data_ptr());
+    params.norm_eps = static_cast<float>(eps);
+    params.num_norm_rows = static_cast<uint32_t>(num_norm_rows);
+    launch_pull_norm(params, num_blocks, unroll, input.device());
+  }
+};

python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ag.cuh

@@ -0,0 +1,303 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <sgl_kernel/distributed/communicator.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/extra/stl.h>
+
+#include "ptx_sys.cuh"
+#include <array>
+#include <cstdint>
+#include <optional>
+#include <utility>
+
+namespace sglang {
+
+namespace gemm_ag {
+
+using device::distributed::Counter;
+using device::distributed::multimem_store_relaxed;
+
+constexpr uint32_t kWorld = 8;                      // TP world size
+constexpr uint32_t kVecSize = 32 / sizeof(bf16_t);  // 16 bf16 per 32B vector
+constexpr uint32_t kSpinBlock = 128;                // consumer threads per block
+constexpr uint32_t kSpinVec = 16 / sizeof(bf16_t);  // 8 bf16 (16B) per consumer thread
+
+// Producer: per-rank column-slice GEMV, multicast store with Lamport markers.
+
+struct ProducerParams {
+  uint8_t* ws_mc;       // multicast VA of the push workspace base
+  Counter* counter;     // per-block phase counters (READ only here)
+  uint32_t half_bytes;  // bytes per phase half (world_size * push_bytes)
+  uint32_t rank;
+};
+
+template <uint32_t K, uint32_t N, uint32_t M, uint32_t N_SPLIT, bool kUsePDL>
+__global__ __launch_bounds__(K / kVecSize) void gemm_ag_gemv_kernel(
+    const __grid_constant__ ProducerParams params,
+    const bf16_t* __restrict__ x,         // [M, K]
+    const bf16_t* __restrict__ weight) {  // [N, K] FULL replicated weight
+  using namespace device;
+  using vec_t = AlignedVector<bf16_t, kVecSize>;
+  constexpr uint32_t kNLocal = N / kWorld;  // columns computed by this rank
+  constexpr uint32_t kGemvBlock = K / kVecSize;
+  constexpr uint32_t kNumWarps = kGemvBlock / kWarpThreads;
+  static_assert(K % kVecSize == 0, "K must be a multiple of the 32B vector width");
+  static_assert(kGemvBlock % kWarpThreads == 0, "K / vec_size must fill whole warps");
+  static_assert(kGemvBlock <= 1024, "K / vec_size exceeds the maximum block size");
+  static_assert(N % kWorld == 0, "N must split evenly over the TP world");
+  static_assert(kNLocal % N_SPLIT == 0, "the local column slice must split into whole tiles");
+  static_assert(M * N_SPLIT <= kGemvBlock, "output tile must fit one thread each for the final reduce");
+  static_assert(N_SPLIT % 2 == 0, "epilogue stores adjacent column pairs");
+
+  const uint32_t bx = blockIdx.x;
+  const uint32_t tx = threadIdx.x;
+  // this rank's rows of the replicated [N, K] weight, sliced HERE (the
+  // Python side always hands the full weight)
+  const bf16_t* weight_tile = weight + (params.rank * kNLocal + bx * N_SPLIT) * K;
+
+  // weight prefetch before the PDL wait (input-independent addresses)
+  vec_t weight_vec[N_SPLIT];
+#pragma unroll
+  for (uint32_t n = 0; n < N_SPLIT; ++n) {
+    weight_vec[n].load(weight_tile + n * K, tx);
+  }
+
+  PDLWaitPrimary<kUsePDL>();
+  // Every push-workspace consumer flips the WHOLE counter array each round
+  // (each has a tail loop up to num_counters), so all counters hold the same
+  // phase at this point and counter[0] is equivalent to counter[bx]. Reading a
+  // single counter is what frees the producer grid from num_push_blocks.
+  const uint32_t phase = params.counter[0].get() & 1;
+
+  vec_t input_vec[M];
+#pragma unroll
+  for (uint32_t m = 0; m < M; ++m) {
+    input_vec[m].load(x + m * K, tx);
+  }
+
+  __shared__ alignas(16) float s_acc[kNumWarps][M * N_SPLIT];
+  const uint32_t warp_id = tx / kWarpThreads;
+
+#pragma unroll
+  for (uint32_t m = 0; m < M; ++m) {
+#pragma unroll
+    for (uint32_t n = 0; n < N_SPLIT; ++n) {
+      float acc = 0.0f;
+#pragma unroll
+      for (uint32_t i = 0; i < kVecSize; ++i) {
+        acc = device::math::fma_f32_bf16(input_vec[m][i], weight_vec[n][i], acc);
+      }
+      s_acc[warp_id][m * N_SPLIT + n] = warp::reduce_sum(acc);
+    }
+  }
+  __syncthreads();
+  constexpr uint32_t kNumPairs = M * N_SPLIT / 2;
+  if (tx < kNumPairs) {
+    auto packed = load_as<float2>(s_acc[0], tx);
+#pragma unroll
+    for (uint32_t i = 1; i < kNumWarps; ++i) {
+      const auto [lo, hi] = load_as<float2>(s_acc[i], tx);
+      packed.x += lo;
+      packed.y += hi;
+    }
+    const auto pair = cast<bf16x2_t>(packed);
+    auto bits = *reinterpret_cast<const uint32_t*>(&pair);
+    if (bits == 0) bits = 0x8000u;  // -0.0 in the first element: never all-zero
+    const uint32_t m = (2 * tx) / N_SPLIT;
+    const uint32_t n = (2 * tx) % N_SPLIT;  // even column within the tile
+    // bf16 index in the phase half's dense [world][M][N / world] prefix;
+    // one multicast store lands this rank's pair on EVERY peer
+    const uint32_t elem = (params.rank * M + m) * kNLocal + bx * N_SPLIT + n;
+    const auto base = reinterpret_cast<bf16_t*>(params.ws_mc + phase * params.half_bytes);
+    const auto dst = reinterpret_cast<uint32_t*>(base + elem);
+    multimem_store_relaxed(dst, bits);
+  }
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+// Consumer: Lamport spin + add3, one 16B vector (8 bf16) per thread.
+
+struct ConsumerParams {
+  uint8_t* ws_local;      // LOCAL VA of the push workspace base (poll + reset)
+  Counter* counter;       // per-block phase counters (read + flip)
+  uint32_t num_counters;  // full counter array size (num_push_blocks)
+  uint32_t half_bytes;    // bytes per phase half (world_size * push_bytes)
+  const bf16_t* b;        // [M, N]
+  const bf16_t* c;        // may be null
+  bf16_t* out;            // [M, N]
+  uint32_t num_rows;      // M
+};
+
+template <uint32_t N, bool kHasC, bool kUsePDL>
+__global__ void spin_add3_kernel(const __grid_constant__ ConsumerParams params) {
+  using namespace device;
+  using vec_t = AlignedVector<bf16x2_t, kSpinVec / 2>;  // 8 bf16 as 4 pairs
+  constexpr uint32_t kNLocal = N / kWorld;
+  static_assert(N % kSpinVec == 0, "rows must stay 16B aligned");
+  static_assert(kNLocal % kSpinVec == 0, "a vector must never cross a rank block");
+  const auto bx = blockIdx.x;
+  const auto tx = threadIdx.x;
+  const uint32_t tid = bx * kSpinBlock + tx;
+  const uint32_t elem = tid * kSpinVec;  // first bf16 of this thread's vector
+  const uint32_t phase = params.counter[bx].get() & 1;
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  // use the last block to clean up: it flips ITS OWN counter and every one
+  // past the grid (work blocks flip [0, num_blocks - 1) themselves)
+  if (const auto num_blocks = gridDim.x; bx == num_blocks - 1) {
+    [[unlikely]];
+    __syncthreads();  // ensure phase is ready for all threads
+    for (uint32_t i = num_blocks - 1 + tx; i < params.num_counters; i += kSpinBlock) {
+      params.counter[i].set(phase ^ 1);
+    }
+    return void();  // this block is done, no output to write
+  }
+
+  // Deliberately NO PDLWaitPrimary: the dependency is carried through data
+  if (elem < params.num_rows * N) {
+    const auto row = elem / N;
+    const auto col = elem % N;
+    // out[row, col] lives at half[col / kNLocal][row][col % kNLocal] of the
+    // dense [world][M][N / world] prefix of the current phase half
+    const auto base = reinterpret_cast<bf16_t*>(params.ws_local + phase * params.half_bytes);
+    const auto src = base + ((col / kNLocal) * params.num_rows + row) * kNLocal + col % kNLocal;
+    vec_t b_vec, c_vec;
+    b_vec.load(params.b + elem);
+    if constexpr (kHasC) c_vec.load(params.c + elem);
+    // spin until all 4 packed pairs of the vector have landed
+    uint4 raw;
+    do {
+      asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];"
+                   : "=r"(raw.x), "=r"(raw.y), "=r"(raw.z), "=r"(raw.w)
+                   : "l"(src)
+                   : "memory");
+    } while (raw.x == 0 || raw.y == 0 || raw.z == 0 || raw.w == 0);
+    const auto& gathered = *reinterpret_cast<const vec_t*>(&raw);
+    vec_t out_vec;
+#pragma unroll
+    for (uint32_t j = 0; j < kSpinVec / 2; ++j) {
+      using Trait = DTypeTrait<bf16x2_t>;
+      out_vec[j] = Trait::add(gathered[j], b_vec[j]);
+      if constexpr (kHasC) out_vec[j] = Trait::add(out_vec[j], c_vec[j]);
+    }
+    out_vec.store(params.out + elem);
+    AlignedVector<uint32_t, 4> zero;
+    zero.fill(0);
+    zero.store(src);
+  }
+  __syncthreads();
+  if (tx == 0) params.counter[bx].set(phase ^ 1);
+}
+
+}  // namespace gemm_ag
+}  // namespace sglang
+
+using namespace sglang;
+using host::distributed::CommunicatorRef;
+
+// Host entry point (tiny_gemm style: one GEMV instantiation per M in
+// [1, kMaxM] selected through a constexpr function-pointer table, then the
+// spin consumer launched with PDL right behind it). Any (K, N) that passes
+// the kernels' static_asserts works; Kimi-K3 uses (3584, 7168).
+
+template <uint32_t K, uint32_t N, uint32_t kMaxM, bool kUsePDL>
+struct GEMMAGKernel {
+  using TensorView = tvm::ffi::TensorView;
+
+  // Columns of this rank's slice per producer block; sets the grid to
+  // kNLocal / N_SPLIT. Measured on 2x4 GB300 TP8 at bs=1 (three reps each,
+  // mean TPOT): 8 -> grid 112, 8.36 ms; 4 -> 224, 8.29 ms; 2 -> 448, 8.21 ms.
+  // Standalone GEMV at the same shape: 4.15 / 3.20 / 2.56 us, and 16 -> 4.42 us,
+  // so the trend is monotonic and 2 is the floor (the epilogue stores column
+  // pairs, so N_SPLIT must stay even). 8 used to be the largest grid that fit
+  // the old kNumProducerBlocks <= num_push_blocks bound; the producer now reads
+  // a single phase counter, so the grid is free and 112 blocks did not even
+  // fill one per SM.
+  static constexpr uint32_t N_SPLIT = 2;
+  static constexpr uint32_t kNLocal = N / gemm_ag::kWorld;
+  static constexpr uint32_t kGemvBlock = K / gemm_ag::kVecSize;
+  static constexpr uint32_t kNumProducerBlocks = kNLocal / N_SPLIT;
+  static_assert(kNLocal % N_SPLIT == 0);
+
+  using GemvFn = void (*)(gemm_ag::ProducerParams, const bf16_t*, const bf16_t*);
+
+  template <std::size_t... I>
+  static constexpr auto make_table(std::index_sequence<I...>) {
+    return std::array<GemvFn, kMaxM + 1>{nullptr, gemm_ag::gemm_ag_gemv_kernel<K, N, I + 1, N_SPLIT, kUsePDL>...};
+  }
+  static constexpr auto kGemvTable = make_table(std::make_index_sequence<kMaxM>{});
+
+  static void
+  run(CommunicatorRef ref,
+      TensorView x,
+      TensorView weight,
+      TensorView b,
+      std::optional<TensorView> c,
+      TensorView out,
+      intptr_t ws_mc_base) {
+    using namespace host;
+    const auto& data = *ref.get();
+
+    auto M = SymbolicSize{"num_tokens"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    TensorMatcher({M, K}).with_dtype<bf16_t>().with_device(device).verify(x);
+    TensorMatcher({N, K}).with_dtype<bf16_t>().with_device(device).verify(weight);
+    TensorMatcher({M, N}).with_dtype<bf16_t>().with_device(device).verify(b);
+    if (c.has_value()) {
+      TensorMatcher({M, N}).with_dtype<bf16_t>().with_device(device).verify(c.value());
+    }
+    TensorMatcher({M, N}).with_dtype<bf16_t>().with_device(device).verify(out);
+    const auto num_tokens = static_cast<uint32_t>(M.unwrap());
+    CHECK_HOST(num_tokens >= 1 && num_tokens <= kMaxM);
+    CHECK_HOST(data.world_size == gemm_ag::kWorld) << "the kernel is compiled for TP" << gemm_ag::kWorld;
+    CHECK_HOST(ws_mc_base != 0) << "requires a multicast-capable workspace";
+    CHECK_HOST(int64_t(num_tokens) * kNLocal * 2 <= data.push_bytes)
+        << "staging slice exceeds the push slot size " << data.push_bytes;
+    // The producer grid is no longer bound to the counter array: it reads only
+    // counter[0] (see gemm_ag_gemv_kernel). The consumer grid still is.
+    CHECK_HOST(data.num_push_blocks > 0) << "no push blocks available";
+    // producer: GEMV
+    const auto producer_params = gemm_ag::ProducerParams{
+        .ws_mc = reinterpret_cast<uint8_t*>(ws_mc_base),
+        .counter = data.push_counter,
+        .half_bytes = static_cast<uint32_t>(data.push_bytes * data.world_size),
+        .rank = data.rank,
+    };
+    LaunchKernel(kNumProducerBlocks, kGemvBlock, device.unwrap())
+        .enable_pdl(kUsePDL)(
+            kGemvTable[num_tokens],
+            producer_params,
+            static_cast<const bf16_t*>(x.data_ptr()),
+            static_cast<const bf16_t*>(weight.data_ptr()));
+
+    // consumer: spin + add3
+    const auto consumer_params = gemm_ag::ConsumerParams{
+        .ws_local = data.push_workspaces[data.rank],
+        .counter = data.push_counter,
+        .num_counters = data.num_push_blocks,
+        .half_bytes = static_cast<uint32_t>(data.push_bytes * data.world_size),
+        .b = static_cast<const bf16_t*>(b.data_ptr()),
+        .c = c.has_value() ? static_cast<const bf16_t*>(c.value().data_ptr()) : nullptr,
+        .out = static_cast<bf16_t*>(out.data_ptr()),
+        .num_rows = num_tokens,
+    };
+    const auto num_vecs = num_tokens * N / gemm_ag::kSpinVec;
+    const auto num_consumers = host::div_ceil(num_vecs, gemm_ag::kSpinBlock);
+    CHECK_HOST(num_consumers + 1 <= data.num_push_blocks);
+    // use last block to clean up the counter
+    const auto num_consumer_blocks = num_consumers + 1;
+    using gemm_ag::spin_add3_kernel;
+    const auto kernel = c.has_value() ? spin_add3_kernel<N, 1, kUsePDL> : spin_add3_kernel<N, 0, kUsePDL>;
+    host::LaunchKernel(num_consumer_blocks, gemm_ag::kSpinBlock, device.unwrap())
+        .enable_pdl(kUsePDL)(kernel, consumer_params);
+  }
+};

python/sglang/kernels/jit/csrc/kimi_k3/comm/ptx_sys.cuh

@@ -0,0 +1,68 @@
+#pragma once
+
+// System-scope PTX the Kimi K3 collectives need and no shared sglang header
+// wraps. These live beside their only consumers (gemm_ar / gemm_ag) rather than
+// in distributed/communicator.cuh: that header's Semaphore does not use any of
+// them, so putting them there would grow a shared header for one caller's
+// benefit.
+//
+// The `device::distributed` namespace is deliberate -- it is where the rest of
+// the collective vocabulary lives, so call sites and `using` declarations read
+// the same whichever header supplied the symbol.
+
+#include <sgl_kernel/utils.cuh>
+
+#include <cstdint>
+
+namespace device::distributed {
+
+// Peer-visible flag increment. `.sys` scope, relaxed: ordering is established by
+// the surrounding fence, not by this store.
+SGL_DEVICE void red_add_relaxed_sys(uint32_t* ptr, uint32_t val) {
+  asm volatile("red.relaxed.sys.global.add.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");
+}
+
+// Acquire load of a peer-written flag: everything the writer released before
+// its matching store is visible to this thread afterwards.
+SGL_DEVICE uint32_t load_acquire_sys(const uint32_t* ptr) {
+  uint32_t val;
+  asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(ptr) : "memory");
+  return val;
+}
+
+// Publishes every prior write to system scope. Pair with a relaxed flag store so
+// a peer's acquire load of that flag also observes the payload.
+SGL_DEVICE void fence_release_sys() {
+  asm volatile("fence.release.sys;" ::: "memory");
+}
+
+// Device-scope arrival counter. acq_rel so the winner of the count also observes
+// the losers' payload writes.
+SGL_DEVICE uint32_t atomic_add_acq_rel_gpu(uint32_t* ptr, uint32_t val) {
+  uint32_t old;
+  asm volatile("atom.acq_rel.gpu.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(ptr), "r"(val) : "memory");
+  return old;
+}
+
+// One store fanned out to every rank in the multicast team.
+SGL_DEVICE void multimem_store_relaxed(uint32_t* ptr, uint32_t val) {
+  asm volatile("multimem.st.relaxed.sys.global.b32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");
+}
+
+SGL_DEVICE void multimem_red_add_relaxed(uint32_t* mc_flag) {
+#if SGL_ARCH_HOPPER_OR_GREATER
+  asm volatile("multimem.red.relaxed.sys.global.add.u32 [%0], 1;" ::"l"(mc_flag) : "memory");
+#else
+  assert(false && "multimem red is only supported on Hopper or later architecture");
+#endif
+}
+
+SGL_DEVICE void multimem_red_add_release(uint32_t* mc_flag) {
+#if SGL_ARCH_HOPPER_OR_GREATER
+  asm volatile("multimem.red.release.sys.global.add.u32 [%0], 1;" ::"l"(mc_flag) : "memory");
+#else
+  assert(false && "multimem red is only supported on Hopper or later architecture");
+#endif
+}
+
+}  // namespace device::distributed

python/sglang/kernels/jit/csrc/kimi_k3/comm/sp_collective.cuh

@@ -0,0 +1,439 @@
+// K3 SP-MoE row-sharded collectives (bf16):
+//
+//   reduce_scatter_res:
+//     [world * rows, hidden] -> [rows, hidden], optionally adding the
+//     destination rank's residual rows in the reduction epilogue.
+//     Every input vector is written exactly once, to the rank that owns its
+//     row shard; the destination polls and reduces the world producer slots.
+//
+//   all_gather:
+//     [rows, hidden] -> [world * rows, hidden].  Every rank multicast-stores
+//     its local shard once, then every peer polls the rank slots and copies
+//     them into rank-concatenated row order.
+//
+// Both kernels reuse CustomAllReduceV2's double-buffered push workspace and
+// phase counters.  A bumper block advances counters outside the tuned work
+// grid, so calls remain protocol-compatible with all-reduce and gemm_ag.
+
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include "../../distributed/custom_all_reduce.cuh"
+
+namespace sglang::sp_collective {
+
+using device::distributed::Counter;
+using device::distributed::Semaphore;
+
+struct Params {
+  const uint8_t* input;
+  uint8_t* output;
+  const uint8_t* residual;
+  uint8_t* push_workspaces[device::distributed::kMaxWorldSize];
+  uint8_t* push_ws_mc;
+  Counter* counter;
+  Semaphore* sem_local;
+  uint8_t* sem_mc;
+  uint8_t* input_mc;
+  uint8_t* output_mc;
+  int64_t stride_bytes;
+  uint32_t num_counters;
+  uint32_t rank;
+  uint32_t local_vecs;
+  uint32_t residual_is_local;
+};
+
+template <typename Vec>
+SGL_DEVICE void make_nonzero(Vec& vec) {
+  constexpr uint32_t kNegZeroPair = 0x8000u;
+  auto& bits = *reinterpret_cast<uint4*>(&vec);
+  if (bits.x == 0) bits.x = kNegZeroPair;
+  if (bits.y == 0) bits.y = kNegZeroPair;
+  if (bits.z == 0) bits.z = kNegZeroPair;
+  if (bits.w == 0) bits.w = kNegZeroPair;
+}
+
+SGL_DEVICE uint32_t* sem_mc_flag(uint8_t* sem_mc, uint32_t block) {
+  static_assert(sizeof(Semaphore) == 128);
+  return reinterpret_cast<uint32_t*>(sem_mc + block * sizeof(Semaphore));
+}
+
+SGL_DEVICE void sem_arrive_relaxed(uint32_t* flag) {
+#if SGL_ARCH_HOPPER_OR_GREATER
+  asm volatile("multimem.red.relaxed.sys.global.add.u32 [%0], 1;" ::"l"(flag) : "memory");
+#else
+  assert(false && "multimem red requires Hopper or later");
+#endif
+}
+
+SGL_DEVICE void sem_arrive_release(uint32_t* flag) {
+#if SGL_ARCH_HOPPER_OR_GREATER
+  asm volatile("multimem.red.release.sys.global.add.u32 [%0], 1;" ::"l"(flag) : "memory");
+#else
+  assert(false && "multimem red requires Hopper or later");
+#endif
+}
+
+template <typename Vec>
+SGL_DEVICE bool has_empty_marker(const Vec& vec) {
+  const auto bits = *reinterpret_cast<const uint4*>(&vec);
+  return bits.x == 0 || bits.y == 0 || bits.z == 0 || bits.w == 0;
+}
+
+template <typename Vec>
+SGL_DEVICE Vec zero_vec() {
+  Vec zero;
+  zero.fill(bf16x2_t{get_pos_zero<bf16_t>(), get_pos_zero<bf16_t>()});
+  return zero;
+}
+
+template <uint32_t kWorldSize>
+SGL_DEVICE bool bumper_block(const Params& params) {
+  const auto bx = blockIdx.x;
+  if (bx + 1 != gridDim.x) return false;
+  const auto phase = params.counter[bx].get() & 1;
+  __syncthreads();
+  for (uint32_t i = bx + threadIdx.x; i < params.num_counters; i += blockDim.x) {
+    params.counter[i].set(phase ^ 1);
+  }
+  return true;
+}
+
+template <uint32_t kWorldSize, bool kHasResidual, bool kUsePDL>
+__global__ void reduce_scatter_res_kernel(const __grid_constant__ Params params) {
+  using vec_t = device::AlignedVector<bf16x2_t, 4>;  // 16 B
+
+  device::PDLWaitPrimary<kUsePDL>();
+  if (bumper_block<kWorldSize>(params)) {
+    device::PDLTriggerSecondary<kUsePDL>();
+    return;
+  }
+
+  const uint32_t bx = blockIdx.x;
+  const uint32_t tid = bx * blockDim.x + threadIdx.x;
+  const uint32_t num_threads = (gridDim.x - 1) * blockDim.x;
+  const uint32_t phase = params.counter[bx].get() & 1;
+  const auto phase_offset = phase * kWorldSize * params.stride_bytes;
+  const auto producer_offset = phase_offset + params.rank * params.stride_bytes;
+
+  // Each vector goes only to the rank that owns its row shard.
+  for (uint32_t vid = tid; vid < kWorldSize * params.local_vecs; vid += num_threads) {
+    const uint32_t dst_rank = vid / params.local_vecs;
+    const uint32_t local_vid = vid - dst_rank * params.local_vecs;
+    vec_t vec;
+    ld_global_16B(vec, params.input, vid);
+    make_nonzero(vec);
+    st_relaxed_16B(vec, params.push_workspaces[dst_rank] + producer_offset, local_vid);
+  }
+
+  device::PDLTriggerSecondary<kUsePDL>();
+
+  // Poll this rank's shard from all producer slots and reduce locally.
+  const auto poll_base = params.push_workspaces[params.rank] + phase_offset;
+  const auto residual_base = params.residual + (params.residual_is_local ? 0 : params.rank * params.local_vecs * 16);
+  const auto zero = zero_vec<vec_t>();
+  for (uint32_t vid = tid; vid < params.local_vecs; vid += num_threads) {
+    vec_t vec[kWorldSize + kHasResidual];
+    if constexpr (kHasResidual) {
+      ld_global_16B(vec[kWorldSize], residual_base, vid);
+    }
+    do {
+      bool empty = false;
+#pragma unroll
+      for (uint32_t rank = 0; rank < kWorldSize; ++rank) {
+        ld_relaxed_16B(vec[rank], poll_base + rank * params.stride_bytes, vid);
+        empty |= has_empty_marker(vec[rank]);
+      }
+      if (!empty) break;
+    } while (true);
+    const auto out = reduce(vec);
+    st_global_16B(out, params.output, vid);
+#pragma unroll
+    for (uint32_t rank = 0; rank < kWorldSize; ++rank) {
+      st_global_16B(zero, poll_base + rank * params.stride_bytes, vid);
+    }
+  }
+
+  __syncthreads();
+  if (threadIdx.x == 0) params.counter[bx].set(phase ^ 1);
+}
+
+template <uint32_t kWorldSize, bool kUsePDL>
+__global__ void all_gather_kernel(const __grid_constant__ Params params) {
+  using vec_t = device::AlignedVector<bf16x2_t, 4>;  // 16 B
+
+  device::PDLWaitPrimary<kUsePDL>();
+  if (bumper_block<kWorldSize>(params)) {
+    device::PDLTriggerSecondary<kUsePDL>();
+    return;
+  }
+
+  const uint32_t bx = blockIdx.x;
+  const uint32_t tid = bx * blockDim.x + threadIdx.x;
+  const uint32_t num_threads = (gridDim.x - 1) * blockDim.x;
+  const uint32_t phase = params.counter[bx].get() & 1;
+  const auto phase_offset = phase * kWorldSize * params.stride_bytes;
+  const auto producer_offset = phase_offset + params.rank * params.stride_bytes;
+
+  // One multicast store places this rank's shard in the same slot on peers.
+  for (uint32_t vid = tid; vid < params.local_vecs; vid += num_threads) {
+    vec_t vec;
+    ld_global_16B(vec, params.input, vid);
+    make_nonzero(vec);
+    st_multimem_16B(vec, params.push_ws_mc + producer_offset, vid);
+  }
+
+  device::PDLTriggerSecondary<kUsePDL>();
+
+  const auto poll_base = params.push_workspaces[params.rank] + phase_offset;
+  const auto zero = zero_vec<vec_t>();
+  for (uint32_t vid = tid; vid < kWorldSize * params.local_vecs; vid += num_threads) {
+    const uint32_t src_rank = vid / params.local_vecs;
+    const uint32_t local_vid = vid - src_rank * params.local_vecs;
+    const auto src = poll_base + src_rank * params.stride_bytes;
+    vec_t vec;
+    do {
+      ld_relaxed_16B(vec, src, local_vid);
+    } while (has_empty_marker(vec));
+    st_global_16B(vec, params.output, vid);
+    st_global_16B(zero, src, local_vid);
+  }
+
+  __syncthreads();
+  if (threadIdx.x == 0) params.counter[bx].set(phase ^ 1);
+}
+
+// Direct variant: output is multicast-bound symmetric memory.  Each producer
+// writes its rank slice straight into every peer's final output, avoiding the
+// staging read/copy/clear.  The two pull-semaphore barriers preserve protocol
+// compatibility with CustomAllReduceV2 and make the remote writes visible
+// before any rank leaves the kernel.
+template <uint32_t kWorldSize, bool kUsePDL>
+__global__ void all_gather_direct_kernel(const __grid_constant__ Params params) {
+  using vec_t = device::AlignedVector<bf16x2_t, 4>;  // 16 B
+
+  uint32_t exit_base = 0;
+  if (threadIdx.x == 0) {
+    auto* semaphore = &params.sem_local[blockIdx.x];
+    const auto reserved = semaphore->counter_ptr()->inc(2 * kWorldSize);
+    exit_base = reserved + kWorldSize;
+    device::PDLWaitPrimary<kUsePDL>();
+    sem_arrive_relaxed(sem_mc_flag(params.sem_mc, blockIdx.x));
+    while (semaphore->get_relaxed() - reserved < kWorldSize)
+      ;
+  }
+  __syncthreads();
+
+  const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint32_t step = gridDim.x * blockDim.x;
+  const uint32_t dst_bias = params.rank * params.local_vecs;
+  for (uint32_t vid = tid; vid < params.local_vecs; vid += step) {
+    vec_t vec;
+    ld_global_16B(vec, params.input, vid);
+    st_multimem_16B(vec, params.output_mc, dst_bias + vid);
+  }
+
+  device::PDLTriggerSecondary<kUsePDL>();
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    auto* semaphore = &params.sem_local[blockIdx.x];
+    sem_arrive_release(sem_mc_flag(params.sem_mc, blockIdx.x));
+    while (semaphore->get_acquire() - exit_base < kWorldSize)
+      ;
+  }
+}
+
+// NVLS pull variant: o_proj writes its TP-partial result into multicast-bound
+// symmetric memory. Each rank reduces only its owned row shard directly from
+// the multicast alias, so no staging or second broadcast is needed.
+template <uint32_t kWorldSize, bool kHasResidual, bool kUsePDL>
+__global__ void reduce_scatter_pull_kernel(const __grid_constant__ Params params) {
+  using vec_t = device::AlignedVector<bf16x2_t, 4>;  // 16 B
+  using SumOp = device::ReductionTrait<device::ReductionOp::SUM, bf16x2_t>;
+
+  uint32_t exit_base = 0;
+  if (threadIdx.x == 0) {
+    auto* semaphore = &params.sem_local[blockIdx.x];
+    const auto reserved = semaphore->counter_ptr()->inc(2 * kWorldSize);
+    exit_base = reserved + kWorldSize;
+    device::PDLWaitPrimary<kUsePDL>();
+    sem_arrive_relaxed(sem_mc_flag(params.sem_mc, blockIdx.x));
+    while (semaphore->get_relaxed() - reserved < kWorldSize)
+      ;
+  }
+  __syncthreads();
+
+  const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint32_t step = gridDim.x * blockDim.x;
+  const auto* input_mc = params.input_mc + params.rank * params.local_vecs * 16;
+  const auto* residual =
+      kHasResidual ? params.residual + (params.residual_is_local ? 0 : params.rank * params.local_vecs * 16) : nullptr;
+  for (uint32_t vid = tid; vid < params.local_vecs; vid += step) {
+    vec_t vec;
+    ld_multimem_16B(vec, input_mc, vid);
+    if constexpr (kHasResidual) {
+      vec_t res;
+      ld_global_16B(res, residual, vid);
+#pragma unroll
+      for (uint32_t j = 0; j < 4; ++j) {
+        vec[j] = SumOp::reduce(vec[j], res[j]);
+      }
+    }
+    st_global_16B(vec, params.output, vid);
+  }
+
+  device::PDLTriggerSecondary<kUsePDL>();
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    auto* semaphore = &params.sem_local[blockIdx.x];
+    sem_arrive_release(sem_mc_flag(params.sem_mc, blockIdx.x));
+    while (semaphore->get_acquire() - exit_base < kWorldSize)
+      ;
+  }
+}
+
+}  // namespace sglang::sp_collective
+
+using namespace sglang;
+using host::distributed::CommunicatorRef;
+
+template <uint32_t kWorldSize, bool kUsePDL>
+struct SPCollectiveKernel {
+  using TensorView = tvm::ffi::TensorView;
+
+  static sp_collective::Params make_params(
+      const host::distributed::CommunicatorObj& data,
+      TensorView input,
+      TensorView output,
+      std::optional<TensorView> residual,
+      bool residual_is_local,
+      int64_t ws_mc_base) {
+    using namespace host;
+    auto input_elems = SymbolicSize{"input_elems"};
+    auto local_elems = SymbolicSize{"local_elems"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    TensorMatcher({input_elems}).with_dtype<bf16_t>().with_device(device).verify(input);
+    TensorMatcher({local_elems}).with_dtype<bf16_t>().with_device(device).verify(output);
+    if (residual.has_value()) {
+      if (residual_is_local) {
+        TensorMatcher({local_elems}).with_dtype<bf16_t>().with_device(device).verify(residual.value());
+      } else {
+        TensorMatcher({input_elems}).with_dtype<bf16_t>().with_device(device).verify(residual.value());
+      }
+    }
+    CHECK_HOST(data.world_size == kWorldSize);
+    CHECK_HOST(local_elems.unwrap() > 0);
+    CHECK_HOST(input_elems.unwrap() == local_elems.unwrap() * kWorldSize);
+    CHECK_HOST(local_elems.unwrap() % 8 == 0) << "local shard bytes must be 16B aligned";
+    CHECK_HOST(local_elems.unwrap() * sizeof(bf16_t) <= data.push_bytes) << "local shard exceeds a push slot";
+
+    sp_collective::Params params{
+        .input = static_cast<const uint8_t*>(input.data_ptr()),
+        .output = static_cast<uint8_t*>(output.data_ptr()),
+        .residual = residual.has_value() ? static_cast<const uint8_t*>(residual.value().data_ptr()) : nullptr,
+        .push_workspaces = {},
+        .push_ws_mc = reinterpret_cast<uint8_t*>(ws_mc_base),
+        .counter = data.push_counter,
+        .sem_local = data.pull_semaphores[data.rank],
+        .sem_mc = nullptr,
+        .input_mc = nullptr,
+        .output_mc = nullptr,
+        .stride_bytes = data.push_bytes,
+        .num_counters = data.num_push_blocks,
+        .rank = data.rank,
+        .local_vecs = static_cast<uint32_t>(local_elems.unwrap() * sizeof(bf16_t) / 16),
+        .residual_is_local = static_cast<uint32_t>(residual_is_local),
+    };
+    for (uint32_t i = 0; i < kWorldSize; ++i) {
+      params.push_workspaces[i] = data.push_workspaces[i];
+    }
+    return params;
+  }
+
+  static void check_launch(const host::distributed::CommunicatorObj& data, int64_t num_blocks, int64_t block_size) {
+    CHECK_HOST(num_blocks > 0 && num_blocks < data.num_push_blocks);
+    // The RS reduction keeps one 16B vector per producer in registers.
+    // 1024-thread CTAs exceed the GB300 launch resource limit.
+    CHECK_HOST(block_size >= 32 && block_size <= 512 && block_size % 32 == 0);
+  }
+
+  static void reduce_scatter_res(
+      CommunicatorRef ref,
+      TensorView input,
+      TensorView output,
+      std::optional<TensorView> residual,
+      bool residual_is_local,
+      int64_t num_blocks,
+      int64_t block_size) {
+    const auto& data = *ref.get();
+    check_launch(data, num_blocks, block_size);
+    auto params = make_params(data, input, output, residual, residual_is_local, 0);
+    const auto kernel = residual.has_value() ? sp_collective::reduce_scatter_res_kernel<kWorldSize, true, kUsePDL>
+                                             : sp_collective::reduce_scatter_res_kernel<kWorldSize, false, kUsePDL>;
+    host::LaunchKernel(num_blocks + 1, block_size, input.device()).enable_pdl(kUsePDL)(kernel, params);
+  }
+
+  static void all_gather(
+      CommunicatorRef ref,
+      TensorView input,
+      TensorView output,
+      int64_t ws_mc_base,
+      int64_t num_blocks,
+      int64_t block_size) {
+    const auto& data = *ref.get();
+    CHECK_HOST(ws_mc_base != 0) << "all-gather requires multicast workspace";
+    check_launch(data, num_blocks, block_size);
+    // Reuse the RS matcher by swapping input/output roles conceptually.
+    auto params = make_params(data, output, input, std::nullopt, false, ws_mc_base);
+    params.input = static_cast<const uint8_t*>(input.data_ptr());
+    params.output = static_cast<uint8_t*>(output.data_ptr());
+    host::LaunchKernel(num_blocks + 1, block_size, input.device())
+        .enable_pdl(kUsePDL)(sp_collective::all_gather_kernel<kWorldSize, kUsePDL>, params);
+  }
+
+  static void all_gather_direct(
+      CommunicatorRef ref,
+      TensorView input,
+      TensorView output,
+      int64_t output_mc_ptr,
+      int64_t sem_mc_ptr,
+      int64_t num_blocks,
+      int64_t block_size) {
+    const auto& data = *ref.get();
+    CHECK_HOST(output_mc_ptr != 0) << "direct all-gather needs symmetric output";
+    CHECK_HOST(sem_mc_ptr != 0) << "direct all-gather needs multicast semaphores";
+    CHECK_HOST(num_blocks > 0 && num_blocks <= data.num_pull_blocks);
+    CHECK_HOST(block_size >= 32 && block_size <= 1024 && block_size % 32 == 0);
+    auto params = make_params(data, output, input, std::nullopt, false, 0);
+    params.input = static_cast<const uint8_t*>(input.data_ptr());
+    params.output = static_cast<uint8_t*>(output.data_ptr());
+    params.output_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(output_mc_ptr));
+    params.sem_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(sem_mc_ptr));
+    host::LaunchKernel(num_blocks, block_size, input.device())
+        .enable_pdl(kUsePDL)(sp_collective::all_gather_direct_kernel<kWorldSize, kUsePDL>, params);
+  }
+
+  static void reduce_scatter_pull(
+      CommunicatorRef ref,
+      TensorView input,
+      TensorView output,
+      std::optional<TensorView> residual,
+      bool residual_is_local,
+      int64_t input_mc_ptr,
+      int64_t sem_mc_ptr,
+      int64_t num_blocks,
+      int64_t block_size) {
+    const auto& data = *ref.get();
+    CHECK_HOST(input_mc_ptr != 0) << "pull RS needs symmetric input";
+    CHECK_HOST(sem_mc_ptr != 0) << "pull RS needs multicast semaphores";
+    CHECK_HOST(num_blocks > 0 && num_blocks <= data.num_pull_blocks);
+    CHECK_HOST(block_size >= 32 && block_size <= 1024 && block_size % 32 == 0);
+    auto params = make_params(data, input, output, residual, residual_is_local, 0);
+    params.input_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(input_mc_ptr));
+    params.sem_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(sem_mc_ptr));
+    const auto kernel = residual.has_value() ? sp_collective::reduce_scatter_pull_kernel<kWorldSize, true, kUsePDL>
+                                             : sp_collective::reduce_scatter_pull_kernel<kWorldSize, false, kUsePDL>;
+    host::LaunchKernel(num_blocks, block_size, input.device()).enable_pdl(kUsePDL)(kernel, params);
+  }
+};

python/sglang/kernels/jit/csrc/kimi_k3/mla_output_gate.cuh

@@ -0,0 +1,84 @@
+// K3 MLA output gate: out = bf16(bf16(x) * bf16(sigmoid(gate))), replacing
+// the torch.sigmoid + mul elementwise pair (two launches, two memory passes)
+// with one kernel. sigmoid is computed in fp32 and rounded to bf16 before the
+// multiply, reproducing the unfused pair's double rounding bit-for-bit.
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck
+
+#include <sgl_kernel/type.cuh>   // For bf16_t, fp32_t, device::cast
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel
+#include <sgl_kernel/vec.cuh>    // For AlignedVector
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+struct MlaOutputGateParams {
+  const bf16_t* __restrict__ x;     // [N] contiguous (flattened [T, H])
+  const bf16_t* __restrict__ gate;  // [N] contiguous
+  bf16_t* __restrict__ out;         // [N] contiguous
+  uint32_t n_vecs;
+};
+
+template <int kThreads, bool kUsePDL>
+__global__ void mla_output_gate_kernel(const MlaOutputGateParams __grid_constant__ params) {
+  using namespace device;
+
+  constexpr int kVecN = 8;
+  using vec_bf16_t = AlignedVector<bf16_t, kVecN>;
+
+  const uint32_t v = blockIdx.x * kThreads + threadIdx.x;
+  if (v >= params.n_vecs) return;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  vec_bf16_t xv, gv, ov;
+  xv.load(params.x, v);
+  gv.load(params.gate, v);
+#pragma unroll
+  for (int i = 0; i < kVecN; ++i) {
+    // Match torch.sigmoid(bf16): fp32 sigmoid, round to bf16, then the bf16
+    // multiply upcasts both operands to fp32 and rounds once more.
+    const float g = cast<fp32_t>(gv[i]);
+    const bf16_t s = cast<bf16_t>(1.0f / (1.0f + expf(-g)));
+    ov[i] = cast<bf16_t>(cast<fp32_t>(xv[i]) * cast<fp32_t>(s));
+  }
+  ov.store(params.out, v);
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int kThreads, bool kUsePDL>
+struct MlaOutputGateKernel {
+  static constexpr auto kernel = mla_output_gate_kernel<kThreads, kUsePDL>;
+
+  static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView gate, const tvm::ffi::TensorView out) {
+    using namespace host;
+
+    auto N_ = SymbolicSize{"numel"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N_}).with_dtype<bf16_t>().with_device(device).verify(x);
+    TensorMatcher({N_}).with_dtype<bf16_t>().with_device(device).verify(gate);
+    TensorMatcher({N_}).with_dtype<bf16_t>().with_device(device).verify(out);
+
+    const auto N = static_cast<uint32_t>(N_.unwrap());
+    RuntimeCheck(N % 8 == 0, "numel must be divisible by 8");
+    if (N == 0) return;
+
+    const auto params = MlaOutputGateParams{
+        .x = static_cast<const bf16_t*>(x.data_ptr()),
+        .gate = static_cast<const bf16_t*>(gate.data_ptr()),
+        .out = static_cast<bf16_t*>(out.data_ptr()),
+        .n_vecs = N / 8,
+    };
+    const uint32_t n_blocks = (params.n_vecs + kThreads - 1) / kThreads;
+    LaunchKernel(n_blocks, kThreads, device.unwrap()).enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/kernels/jit/csrc/kimi_k3/situ_and_mul.cuh

@@ -0,0 +1,450 @@
+// Kimi K3 SiTU activation kernels: plain elementwise and varlen masked with a
+// grouped-quant epilogue. The shared double-softcap activation is inlined below.
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/type.cuh>   // For dtype_trait, bf16_t, fp32_t, cast
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE, PDL helpers
+#include <sgl_kernel/vec.cuh>    // For AlignedVector
+#include <sgl_kernel/warp.cuh>   // For warp::copy_bytes, elect_one_lane, inclusive_sum
+
+#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+#include <cuda_fp8.h>
+#include <limits>
+#include <type_traits>
+
+namespace sglang {
+
+namespace kimi_k3 {
+
+/// One SiTU element. `sigmoid_fast` is `1/(1+expf(-x))` (math.cuh), i.e. the
+/// same expression both call sites used before they were folded together.
+template <bool kHasLinearBeta>
+SGL_DEVICE float situ_activate(float g, float u, float beta, float inv_beta, float linear_beta, float inv_linear_beta) {
+  const float gate_out = beta * tanhf(g * inv_beta) * device::math::sigmoid_fast(g);
+  float up_out;
+  if constexpr (kHasLinearBeta) {
+    up_out = linear_beta * tanhf(u * inv_linear_beta);
+  } else {
+    up_out = u;
+  }
+  return gate_out * up_out;
+}
+
+}  // namespace kimi_k3
+
+}  // namespace sglang
+
+namespace {
+
+// SiTU (SoftCap-GLU) activation:
+//   gate_out = beta * tanh(gate / beta) * sigmoid(gate)
+//   up_out   = linear_beta * tanh(up / linear_beta)
+//   output   = gate_out * up_out
+//
+// Input: bf16 tensor [N, 2*D] (gate = [:, :D], up = [:, D:])
+// Output: bf16 tensor [N, D]
+
+struct SituAndMulParams {
+  const void* __restrict__ input;
+  void* __restrict__ out;
+  float beta;
+  float inv_beta;
+  float linear_beta;
+  float inv_linear_beta;
+  uint32_t hidden_dim;  // D (output width, half of input last dim)
+  uint32_t num_tokens;
+  uint32_t stride_in_vecs;  // input row stride in vector units (2*D/vec if dense)
+};
+
+template <typename T, bool kHasLinearBeta, bool kUsePDL>
+__global__ void situ_and_mul_kernel(const __grid_constant__ SituAndMulParams params) {
+  using namespace device;
+  constexpr auto kVecSize = kMaxVecBytes / sizeof(T);
+  using vec_t = AlignedVector<T, kMaxVecBytes / sizeof(T)>;
+
+  const auto num_vecs = params.hidden_dim / kVecSize;  // per token
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto token_id = tid / num_vecs;
+
+  if (token_id >= params.num_tokens) return;
+
+  const auto offset = tid % num_vecs;
+  // Input rows may be strided (e.g. a slice of a wider fused-GEMM output);
+  // within a row: gate = [0..D-1], up = [D..2D-1].
+  const auto input_offset = static_cast<uint64_t>(token_id) * params.stride_in_vecs + offset;
+  const auto output_offset = tid;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto gate = load_as<vec_t>(params.input, input_offset);
+  const auto up = load_as<vec_t>(params.input, input_offset + num_vecs);
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  const float beta = params.beta;
+  const float inv_beta = params.inv_beta;
+  const float linear_beta = params.linear_beta;
+  const float inv_linear_beta = params.inv_linear_beta;
+
+  vec_t out;
+#pragma unroll
+  for (int i = 0; i < kVecSize; ++i) {
+    const float g = cast<fp32_t>(gate[i]);
+    const float u = cast<fp32_t>(up[i]);
+
+    out[i] =
+        cast<T>(sglang::kimi_k3::situ_activate<kHasLinearBeta>(g, u, beta, inv_beta, linear_beta, inv_linear_beta));
+  }
+
+  store_as<vec_t>(params.out, out, output_offset);
+}
+
+// Host launcher
+
+template <typename T, bool kUsePDL>
+struct SituAndMulKernel {
+  static constexpr auto kVecSize = device::kMaxVecBytes / sizeof(T);
+  static constexpr auto kBlockSize = 256u;
+
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView out,
+      const double beta,
+      const double linear_beta,
+      const bool has_linear_beta) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto D_in = SymbolicSize{"input_width"};
+    auto D_out = SymbolicSize{"output_width"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D_out})  //
+        .with_dtype<T>()
+        .with_device(device_)
+        .verify(out);
+    TensorMatcher({N, D_in})  //
+        .with_dtype<T>()
+        .with_device(device_)
+        .with_strides({-1, 1})
+        .verify(input);
+
+    const auto hidden_size = static_cast<uint32_t>(D_out.unwrap());
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto device = device_.unwrap();
+
+    if (num_tokens == 0) return;
+    RuntimeCheck(hidden_size * 2 == D_in.unwrap(), "invalid activation dimension: D_out * 2 != D_in");
+    RuntimeCheck(hidden_size % kVecSize == 0, "hidden size must be divisible by vector size");
+    RuntimeCheck(input.stride(0) % kVecSize == 0, "input row stride must be divisible by vector size");
+
+    const auto num_total_items = num_tokens * (hidden_size / kVecSize);
+    RuntimeCheck(num_total_items <= std::numeric_limits<uint32_t>::max(), "too many items for 32-bit indexing");
+
+    const auto num_blocks = div_ceil(static_cast<uint32_t>(num_total_items), kBlockSize);
+    const float beta_f = static_cast<float>(beta);
+    const float linear_beta_f = static_cast<float>(linear_beta);
+
+    const auto params = SituAndMulParams{
+        .input = input.data_ptr(),
+        .out = out.data_ptr(),
+        .beta = beta_f,
+        .inv_beta = 1.0f / beta_f,
+        .linear_beta = linear_beta_f,
+        .inv_linear_beta = linear_beta_f != 0.0f ? 1.0f / linear_beta_f : 0.0f,
+        .hidden_dim = hidden_size,
+        .num_tokens = num_tokens,
+        .stride_in_vecs = static_cast<uint32_t>(input.stride(0) / kVecSize),
+    };
+
+    if (has_linear_beta) {
+      LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(situ_and_mul_kernel<T, true, kUsePDL>, params);
+    } else {
+      LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(situ_and_mul_kernel<T, false, kUsePDL>, params);
+    }
+  }
+};
+
+// ---------------------------------------------------------------------------
+// varlen masked variant with the grouped-quant epilogue. Same activation, a
+// different kernel: __launch_bounds__(1024, 2) plus a per-group scale writeback.
+// ---------------------------------------------------------------------------
+using deepseek_v4::fp8::cast_to_ue8m0;
+using deepseek_v4::fp8::pack_fp8;
+
+struct SituMulQuantVarlenParams {
+  const bf16_t* __restrict__ input;
+  fp8_e4m3_t* __restrict__ output;
+  float* __restrict__ output_scale;
+  const int32_t* __restrict__ masked_m;
+  float beta;         // gate softcap (e.g. 4.0)
+  float linear_beta;  // up softcap (e.g. 25.0)
+  int64_t hidden_dim;
+  uint32_t num_tokens;
+  uint32_t num_experts;
+};
+
+constexpr uint32_t kMaxExperts = 256;
+
+struct alignas(16) CTAWork {
+  uint32_t expert_id;
+  uint32_t expert_token_id;
+  bool valid;
+};
+
+// SiTU (SoftCap-GLU) activation:
+//   gate_out = beta * tanh(gate / beta) * sigmoid(gate)
+//   up_out   = linear_beta * tanh(up / linear_beta)
+//   output   = gate_out * up_out
+// Unlike SiLU, no external swiglu_limit clamp is needed: the tanh softcap
+// inherently bounds the output to |beta * linear_beta| (< FP8_E4M3_MAX).
+template <bool kPrecise = true, typename DType2>
+SGL_DEVICE fp32x2_t
+situ_and_mul(DType2 gate, DType2 up, float beta, float inv_beta, float linear_beta, float inv_linear_beta) {
+  using namespace device;
+  const auto [g0, g1] = cast<fp32x2_t>(gate);
+  const auto [u0, u1] = cast<fp32x2_t>(up);
+  // kHasLinearBeta=true: this path always softcaps the up operand, as before.
+  const float val0 = sglang::kimi_k3::situ_activate<true>(g0, u0, beta, inv_beta, linear_beta, inv_linear_beta);
+  const float val1 = sglang::kimi_k3::situ_activate<true>(g1, u1, beta, inv_beta, linear_beta, inv_linear_beta);
+  if constexpr (kPrecise) {
+    return {val0, val1};
+  } else {
+    return cast<fp32x2_t>(cast<bf16x2_t>(fp32x2_t{val0, val1}));
+  }
+}
+
+[[maybe_unused]]
+SGL_DEVICE CTAWork get_work(const SituMulQuantVarlenParams& params) {
+  // Preconditions:
+  // 1. blockDim.x >= params.num_experts
+  // 2. params.num_experts <= kMaxExperts
+  using namespace device;
+  static_assert(kWarpThreads == 32);
+
+  static __shared__ uint32_t s_warp_sum[32];
+  static __shared__ CTAWork result;
+
+  result.valid = false;
+
+  const uint32_t tx = threadIdx.x;
+  const uint32_t lane_id = tx % kWarpThreads;
+  const uint32_t warp_id = tx / kWarpThreads;
+
+  const uint32_t val = tx < params.num_experts ? params.masked_m[tx] : 0u;
+
+  // Per-warp inclusive scan of masked_m.
+  const uint32_t warp_inclusive = device::warp::inclusive_sum(lane_id, val);
+  const uint32_t warp_exclusive = warp_inclusive - val;
+
+  // Write each warp total.
+  if (lane_id == kWarpThreads - 1) s_warp_sum[warp_id] = warp_inclusive;
+  __syncthreads();
+  const auto tmp_val = lane_id < warp_id ? s_warp_sum[lane_id] : 0u;
+  const auto prefix_exclusive = warp::reduce_sum(tmp_val) + warp_exclusive;
+  const auto bx = blockIdx.x;
+  if (prefix_exclusive <= bx && bx < prefix_exclusive + val) {
+    result = {tx, bx - prefix_exclusive, true};
+  }
+  __syncthreads();
+  return result;
+}
+
+template <bool kScaleUE8M0, bool kTransposed, bool kSwizzle, bool kUsePDL>
+__global__ __launch_bounds__(1024, 2) void  // maximize occupancy
+    situ_mul_quant_varlen_kernel(const SituMulQuantVarlenParams __grid_constant__ params) {
+  using namespace device;
+
+  constexpr uint32_t kGroupSize = 128u;
+  constexpr uint32_t kWorkThreads = 16u;
+  // each thread will handle 8 elements
+  using InputVec = AlignedVector<bf16x2_t, 4>;
+  using OutputVec = AlignedVector<fp8x2_e4m3_t, 4>;
+  static_assert(8 * kWorkThreads == 128, "Invalid tiling");
+  static_assert(!(kTransposed && !kScaleUE8M0), "transposed layout only supports ue8m0");
+
+  const auto [expert_id, token_id, valid] = get_work(params);
+
+  if (!valid) return;
+
+  const auto work_id = threadIdx.x / kWorkThreads;
+
+  const auto offset = expert_id * params.num_tokens + token_id;
+  const auto input = params.input + offset * params.hidden_dim * 2;
+  const auto output = params.output + offset * params.hidden_dim;
+  [[maybe_unused]]
+  const auto output_scale = [&] {
+    const auto num_groups = params.hidden_dim / kGroupSize;
+    if constexpr (kTransposed) {
+      const auto base = reinterpret_cast<uint8_t*>(params.output_scale);
+      // Physical layout is [E, G//4, N] int32.  Each int32 packs 4 consecutive
+      // group scales for the same token, so the byte address is:
+      //   expert_offset + (group/4)*N*4 + token*4 + group%4
+      return base + expert_id * num_groups * params.num_tokens + (work_id / 4u) * (params.num_tokens * 4u) +
+             token_id * 4u + (work_id % 4u);
+    } else {
+      return params.output_scale + offset * num_groups + work_id;
+    }
+  }();
+
+  const float beta = params.beta;
+  const float linear_beta = params.linear_beta;
+  const float inv_beta = 1.0f / beta;
+  const float inv_linear_beta = 1.0f / linear_beta;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  InputVec gate_vec, up_vec;
+  if constexpr (kSwizzle) {
+    // gran=8 interleaved: every 16-element chunk on the N axis is
+    // [gate[0..7], up[0..7]]. Each thread handles 8 consecutive output
+    // elements, so its gate chunk lives at vec index 2*threadIdx.x and its
+    // up chunk at 2*threadIdx.x+1.
+    gate_vec.load(input, threadIdx.x * 2);
+    up_vec.load(input, threadIdx.x * 2 + 1);
+  } else {
+    gate_vec.load(input, threadIdx.x);
+    up_vec.load(input, threadIdx.x + blockDim.x);
+  }
+
+  float local_max = 0.0f;
+  float results[8];
+
+#pragma unroll
+  for (uint32_t i = 0; i < 4; ++i) {
+    const auto [x, y] = situ_and_mul(gate_vec[i], up_vec[i], beta, inv_beta, linear_beta, inv_linear_beta);
+    results[2 * i + 0] = x;
+    results[2 * i + 1] = y;
+    local_max = fmaxf(local_max, fmaxf(fabsf(x), fabsf(y)));
+  }
+
+  local_max = warp::reduce_max<kWorkThreads>(local_max);
+
+  const float absmax = fmaxf(local_max, 1e-10f);
+  float scale;
+  uint32_t ue8m0_exp;
+
+  if constexpr (kScaleUE8M0) {
+    const float raw_scale = absmax / math::FP8_E4M3_MAX;
+    ue8m0_exp = cast_to_ue8m0(raw_scale);
+    scale = __uint_as_float(ue8m0_exp << 23);
+  } else {
+    scale = absmax / math::FP8_E4M3_MAX;
+  }
+  const auto inv_scale = 1.0f / scale;
+
+  OutputVec out_vec;
+#pragma unroll
+  for (uint32_t i = 0; i < 4; ++i) {
+    const float scaled_val0 = results[2 * i + 0] * inv_scale;
+    const float scaled_val1 = results[2 * i + 1] * inv_scale;
+    out_vec[i] = pack_fp8(scaled_val0, scaled_val1);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  out_vec.store(output, threadIdx.x);
+  if constexpr (kTransposed) {
+    *output_scale = ue8m0_exp;
+  } else {
+    *output_scale = scale;
+  }
+}
+
+// ---- Host wrapper
+
+template <int64_t kGroupSize, bool kScaleUE8M0, bool kSwizzle, bool kUsePDL>
+struct SituAndMulMaskedPostQuantKernel {
+  static_assert(kGroupSize == 128);
+  static constexpr auto kernel_normal = situ_mul_quant_varlen_kernel<kScaleUE8M0, false, kSwizzle, kUsePDL>;
+  static constexpr auto kernel_transposed = situ_mul_quant_varlen_kernel<true, true, kSwizzle, kUsePDL>;
+
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView output,
+      const tvm::ffi::TensorView output_scale,
+      const tvm::ffi::TensorView masked_m,
+      const uint32_t topk,
+      const bool transposed,
+      const double beta,
+      const double linear_beta) {
+    using namespace host;
+
+    auto device = SymbolicDevice{};
+    auto E = SymbolicSize{"num_experts"};
+    auto T = SymbolicSize{"num_tokens_padded"};
+    auto D = SymbolicSize{"hidden_dim x 2"};
+    auto N = SymbolicSize{"hidden_dim"};
+    auto G = SymbolicSize{"num_groups"};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({E, T, D})  // input
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(input);
+    TensorMatcher({E, T, N})  // output
+        .with_dtype<fp8_e4m3_t>()
+        .with_device(device)
+        .verify(output);
+    if (!transposed) {
+      TensorMatcher({E, T, G})  //
+          .with_dtype<fp32_t>()
+          .with_device(device)
+          .verify(output_scale);
+    } else {
+      RuntimeCheck(kScaleUE8M0, "transposed layout only supports scale_ue8m0=true");
+      auto G_ = SymbolicSize{"G // 4"};
+      TensorMatcher({E, G_, T})  //
+          .with_dtype<int32_t>()
+          .with_device(device)
+          .verify(output_scale);
+      G.set_value(G_.unwrap() * 4);
+    }
+    TensorMatcher({E})  //
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(masked_m);
+
+    const auto num_experts = static_cast<uint32_t>(E.unwrap());
+    const auto num_tokens = static_cast<uint32_t>(T.unwrap());
+    const auto num_groups = static_cast<uint32_t>(G.unwrap());
+    const auto hidden_dim = N.unwrap();
+
+    RuntimeCheck(D.unwrap() == 2 * hidden_dim, "invalid dimension");
+    RuntimeCheck(hidden_dim % kGroupSize == 0);
+    RuntimeCheck(num_experts <= kMaxExperts, "num_experts exceeds maximum (256)");
+    RuntimeCheck(num_groups * kGroupSize == hidden_dim, "invalid num_groups");
+
+    const auto params = SituMulQuantVarlenParams{
+        .input = static_cast<const bf16_t*>(input.data_ptr()),
+        .output = static_cast<fp8_e4m3_t*>(output.data_ptr()),
+        .output_scale = static_cast<float*>(output_scale.data_ptr()),
+        .masked_m = static_cast<const int32_t*>(masked_m.data_ptr()),
+        .beta = static_cast<float>(beta),
+        .linear_beta = static_cast<float>(linear_beta),
+        .hidden_dim = hidden_dim,
+        .num_tokens = num_tokens,
+        .num_experts = num_experts,
+    };
+
+    const auto num_threads = hidden_dim / 8;
+    RuntimeCheck(num_threads % device::kWarpThreads == 0);
+    RuntimeCheck(num_threads >= num_experts);
+    const auto kernel = transposed ? kernel_transposed : kernel_normal;
+    LaunchKernel(num_tokens * topk, num_threads, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/kernels/jit/csrc/moe/route_quant_fused.cuh

@@ -0,0 +1,142 @@
+// K3 MoE-front prep in one launch: radix routing (+ trtllm packed ids) on the
+// first M CTAs, mxfp8 per-token-group quant of the routed activations on the
+// next M. At decode batch sizes the unfused chain is three tiny back-to-back
+// kernels (route 3.8us + quant 2.6us + pack 1.4us per layer) each leaving the
+// SMs idle; fused, the quant CTAs run concurrently with the routing CTA and
+// the pack is a 16-store epilogue.
+//
+// Both halves are the existing kernels verbatim: route_radix_block is the
+// standalone route_radix body (same TU, same flags — no fast-math, which the
+// routing bit-exactness contract requires and the quant math tolerates: its
+// only transcendentals are explicit intrinsics and exact bit manipulation),
+// and QuantTrait::run is the per_token_group_quant math. Specialized like
+// route_radix itself: 896 experts, top-16, and a 3584-wide bf16 activation row
+// (112 ue8m0 groups of 32 = 224 lanes, exactly the routing block width).
+
+#include "../gemm/per_token_group_quant.cuh"
+#include "route_radix.cuh"
+
+namespace sglang {
+
+struct RouteQuantFusedParams {
+  RouteRadixParams route;
+  QuantKernelParams quant;
+};
+
+// One quant CTA covers one token row: thread pairs (2g, 2g+1) hold group g
+// with lanes (0, 1) — the same subwarp layout the flat quant kernel derives
+// from global_tid, so the group reduction and stores are bit-identical.
+using RouteQuantTrait = QuantTrait<
+    bf16_t,
+    fp8_e4m3_t,
+    /*kGroupSize=*/32,
+    /*kUe8m0=*/true,
+    /*kRowMajor=*/true,
+    /*kAligned=*/true,
+    /*kFuseSiluAndMul=*/false>;
+
+inline constexpr uint32_t kQuantGroupsPerRow_ = LargeRouterRadixTrait::kBlockSize / RouteQuantTrait::kNumLanes;
+inline constexpr uint32_t kQuantHidden_ = kQuantGroupsPerRow_ * RouteQuantTrait::kGroupSize;  // 3584
+
+template <bool kUsePDL, typename TScore>
+__global__ __launch_bounds__(LargeRouterRadixTrait::kBlockSize)  //
+    void route_quant_fused_kernel(const __grid_constant__ RouteQuantFusedParams params) {
+  const auto M = static_cast<uint32_t>(params.route.M);
+  if (blockIdx.x < M) {
+    __shared__ typename LargeRouterRadixTrait::Smem smem;
+    route_radix_block<kUsePDL, TScore>(params.route, smem);
+  } else {
+    // Quant CTAs read the same primary-kernel output (the fused-front GEMM)
+    // as the routing CTAs, so they carry their own PDL wait/trigger.
+    device::PDLWaitPrimary<kUsePDL>();
+    const uint32_t token_idx = blockIdx.x - M;
+    const uint32_t group_idx = threadIdx.x / RouteQuantTrait::kNumLanes;
+    const uint32_t lane_id = threadIdx.x % RouteQuantTrait::kNumLanes;
+    RouteQuantTrait::run(params.quant, /*expert_idx=*/0, token_idx, group_idx, lane_id);
+    device::PDLTriggerSecondary<kUsePDL>();
+  }
+}
+
+}  // namespace sglang
+
+template <bool kUsePDL>
+struct RouteQuantFusedKernel {
+  static void
+  run(const tvm::ffi::TensorView scores,
+      const tvm::ffi::TensorView bias,
+      const tvm::ffi::TensorView out_w,
+      const tvm::ffi::TensorView out_i,
+      const tvm::ffi::TensorView out_packed,
+      const tvm::ffi::TensorView x,
+      const tvm::ffi::TensorView out_q,
+      const tvm::ffi::TensorView out_s,
+      int64_t topk,
+      double routed_scaling_factor,
+      bool renormalize,
+      bool apply_scale) {
+    using namespace host;
+    using Trait = sglang::RouteQuantTrait;
+
+    auto M_ = SymbolicSize{"num_tokens"};
+    auto N_ = SymbolicSize{"num_experts"};
+    auto K_ = SymbolicSize{"topk"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    auto score_dtype = SymbolicDType{};
+    TensorMatcher({M_, N_})
+        .with_dtype<bf16_t, fp32_t>(score_dtype)
+        .with_device(device)
+        .with_strides({-1, 1})
+        .verify(scores);
+    TensorMatcher({N_}).with_dtype<fp32_t>().with_device(device).verify(bias);
+    TensorMatcher({M_, K_}).with_dtype<fp32_t>().with_device(device).verify(out_w);
+    TensorMatcher({M_, K_}).with_dtype<int32_t>().with_device(device).verify(out_i);
+    TensorMatcher({M_, K_}).with_dtype<int32_t>().with_strides({-1, 1}).with_device(device).verify(out_packed);
+
+    RuntimeCheck(
+        N_.unwrap() == sglang::kNumExperts_ && K_.unwrap() == sglang::kTopK_ && topk == sglang::kTopK_,
+        "route_quant_fused is specialized for N=896, K=16");
+    RuntimeCheck(scores.stride(0) % 4 == 0, "route_quant_fused: scores row stride must be a multiple of 4");
+
+    // Quant half: shape/stride/alignment checks + byte-stride munging shared
+    // with the standalone flat kernel.
+    const auto ctx = build_quant_context<Trait, /*kMasked=*/false>(x, out_q, out_s);
+    RuntimeCheck(
+        ctx.params.hidden_size == sglang::kQuantHidden_,
+        "route_quant_fused is specialized for a 3584-wide activation row");
+    RuntimeCheck(
+        ctx.params.num_tokens == static_cast<uint32_t>(M_.unwrap()),
+        "route_quant_fused: scores and activations must have the same token count");
+
+    const auto M = static_cast<uint32_t>(M_.unwrap());
+    if (M == 0) return;
+
+    const auto params = sglang::RouteQuantFusedParams{
+        .route =
+            {scores.data_ptr(),
+             static_cast<const fp32_t*>(bias.data_ptr()),
+             static_cast<fp32_t*>(out_w.data_ptr()),
+             static_cast<int32_t*>(out_i.data_ptr()),
+             static_cast<int32_t*>(out_packed.data_ptr()),
+             static_cast<int>(M),
+             static_cast<long long>(scores.stride(0)),
+             static_cast<long long>(out_w.stride(0)),
+             static_cast<long long>(out_i.stride(0)),
+             static_cast<long long>(out_packed.stride(0)),
+             static_cast<float>(routed_scaling_factor),
+             renormalize ? 1 : 0,
+             apply_scale ? 1 : 0,
+             /*sorted=*/0},
+        .quant = ctx.params,
+    };
+
+    if (score_dtype.is_type<fp32_t>()) {
+      LaunchKernel(2 * M, sglang::LargeRouterRadixTrait::kBlockSize, device.unwrap())
+          .enable_pdl(kUsePDL)(sglang::route_quant_fused_kernel<kUsePDL, fp32_t>, params);
+    } else {
+      LaunchKernel(2 * M, sglang::LargeRouterRadixTrait::kBlockSize, device.unwrap())
+          .enable_pdl(kUsePDL)(sglang::route_quant_fused_kernel<kUsePDL, bf16_t>, params);
+    }
+  }
+};

python/sglang/kernels/jit/csrc/moe/route_radix.cuh

@@ -0,0 +1,753 @@
+// MoE routing by radix select: the standalone route kernel and the fused-gate
+// front end, over one copy of the radix primitives (previously
+// moe/radix_select_common.cuh, folded in once its two consumers became one file).
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil
+
+#include <sgl_kernel/type.cuh>   // For dtype_trait, bf16_t, fp32_t, cast
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE, PDL helpers
+#include <sgl_kernel/vec.cuh>    // For AlignedVector
+#include <sgl_kernel/warp.cuh>   // For warp::copy_bytes, elect_one_lane, inclusive_sum
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace sglang {
+
+namespace moe::radix {
+
+struct RadixSelectBase {
+  static constexpr uint32_t kRadixBits = 8;
+  static constexpr uint32_t kRadixSize = 1 << kRadixBits;
+  static constexpr uint32_t kRadixRounds = 32 / kRadixBits;
+
+  struct alignas(16) MatchBin {
+    uint32_t bin;
+    uint32_t above_count;  // active elements in bins strictly above `bin`
+    uint32_t equal_count;  // active elements in bin `bin`
+  };
+};
+
+inline constexpr float kNanFloor = -1e30f;
+
+// Monotone unsigned key: larger biased -> larger key. Caller must have floored
+// biased-NaN. Canonicalizes -0.0 -> +0.0 so equal values get equal keys.
+SGL_DEVICE uint32_t biased_to_key(float biased) {
+  if (biased == 0.0f) biased = 0.0f;
+  uint32_t u = __float_as_uint(biased);
+  return (u & 0x80000000u) ? ~u : (u | 0x80000000u);
+}
+
+// tl.sigmoid(x) = 1/(1+exp(-x)). Must stay instruction-identical to v1's
+// sigmoid_match so both kernels rank (and weight) identically.
+SGL_DEVICE float sigmoid_match(float x) {
+  return __fdividef(1.0f, 1.0f + __expf(-x));
+}
+
+SGL_DEVICE float nan_floor(float x) {
+  return (x == x) ? x : kNanFloor;
+}
+
+SGL_DEVICE void bar_sync(uint32_t id, uint32_t num_threads) {
+  asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(num_threads) : "memory");
+}
+
+// Exclusive prefix (block-wide, thread-rank order) of `cnt`. Uses
+// smem_warp_sum[kNumWarps]; syncs on entry (so the workspace can be reused
+// across calls) and before the cross-warp read.
+SGL_DEVICE uint32_t block_exclusive_sum(uint32_t cnt, uint32_t lane_id, uint32_t warp_id, uint32_t* smem_warp_sum) {
+  const uint32_t inc = device::warp::inclusive_sum(lane_id, cnt);
+  if (lane_id == 31) smem_warp_sum[warp_id] = inc;
+  __syncthreads();
+  // TODO: replace `__reduce_add_sync` with `warp::reduce_sum`
+  const auto base = __reduce_add_sync(0xFFFFFFFF, lane_id < warp_id ? smem_warp_sum[lane_id] : 0u);
+  return base + inc - cnt;
+}
+
+}  // namespace moe::radix
+
+inline constexpr uint32_t kNumExperts_ = 896;
+inline constexpr uint32_t kTopK_ = 16;
+
+struct LargeRouterRadixTrait : moe::radix::RadixSelectBase {
+  static constexpr uint32_t kNumExperts = kNumExperts_;
+  static constexpr uint32_t kTopK = kTopK_;
+  static constexpr uint32_t kVecSize = 4;
+
+  static constexpr uint32_t kBlockSize = kNumExperts / kVecSize;  // 224 = 7 warps
+  static constexpr uint32_t kNumWarps = kBlockSize / 32;
+  struct Smem {
+    uint32_t warp_sum[3][kNumWarps];  // cross-warp scan workspace
+    MatchBin match[kRadixRounds];
+    uint32_t histogram[kRadixSize];
+    // winner staging (compaction order = expert-id ascending)
+    int32_t wid[kTopK];
+    uint32_t wkey[kTopK];
+    fp32_t wact[kTopK];
+    // sorted staging ((key desc, id asc) order), only used when sorted != 0
+    int32_t sid[kTopK];
+    fp32_t sact[kTopK];
+    fp32_t norm;
+  };
+};
+
+struct RouteRadixParams {
+  const void* __restrict__ scores;  // bf16 or fp32, typed by the kernel template
+  const fp32_t* __restrict__ bias;
+  fp32_t* __restrict__ out_w;
+  int32_t* __restrict__ out_i;
+  // Optional trtllm-gen routed-MoE packing: (id << 16) | bf16(weight) bits,
+  // bit-identical to the standalone triton pack. nullptr skips the store.
+  int32_t* __restrict__ out_packed;
+  int M;
+  long long scores_stride;
+  long long out_w_stride;
+  long long out_i_stride;
+  long long out_packed_stride;
+  float routed_scaling_factor;
+  int renormalize;
+  int apply_scale;
+  int sorted;
+};
+
+// Whole-CTA routing body, callable from other kernels (the fused
+// route+quant launch runs it on its first M CTAs). Routes row `blockIdx.x`;
+// every thread of the 224-wide CTA must enter (block barriers inside).
+template <bool kUsePDL, typename TScore>
+SGL_DEVICE void route_radix_block(const RouteRadixParams& params, typename LargeRouterRadixTrait::Smem& smem) {
+  using namespace device;
+  using T = LargeRouterRadixTrait;
+  constexpr uint32_t kVecSize = T::kVecSize;
+  constexpr uint32_t kRadixLanes = T::kRadixSize / 2;  // 128: 2 bins per thread
+  enum { BAR_RESERVED = 0, BAR_SUM = 1 };
+
+  const auto bx = blockIdx.x;
+  const auto tx = threadIdx.x;
+  const auto warp_id = tx / kWarpThreads;
+  const auto lane_id = tx % kWarpThreads;
+  // grid.x == M exactly; no row guard (an early return would deadlock the
+  // block-wide barriers below).
+
+  // ---- Load + key transform: thread tx owns experts [4*tx, 4*tx+4) ----
+  uint32_t keys[kVecSize];
+  float act[kVecSize];  // raw sigmoid (weight source) — never NaN-sanitized
+  {
+    const auto scores = static_cast<const TScore*>(params.scores) + bx * params.scores_stride;
+    AlignedVector<fp32x2_t, kVecSize / 2> bias_vec;
+    // bf16: 2x bf16x2 (8B row loads); fp32: 2x fp32x2 (16B row loads). The
+    // radix math below is fp32 either way — only the load width differs.
+    AlignedVector<packed_t<TScore>, kVecSize / 2> scores_vec;
+
+    // prefetch bias (frozen weight) before the PDL wait
+    bias_vec.load(params.bias, tx);
+    PDLWaitPrimary<kUsePDL>();
+    scores_vec.load(scores, tx);
+
+#pragma unroll
+    for (uint32_t i = 0; i < kVecSize / 2; ++i) {
+      fp32x2_t xy;
+      if constexpr (std::is_same_v<TScore, fp32_t>) {
+        xy = scores_vec[i];
+      } else {
+        xy = cast<fp32x2_t>(scores_vec[i]);
+      }
+      const auto [x, y] = xy;
+      const auto sx = moe::radix::sigmoid_match(x), sy = moe::radix::sigmoid_match(y);
+      keys[2 * i + 0] = moe::radix::biased_to_key(moe::radix::nan_floor(sx + bias_vec[i].x));
+      keys[2 * i + 1] = moe::radix::biased_to_key(moe::radix::nan_floor(sy + bias_vec[i].y));
+      act[2 * i + 0] = sx;
+      act[2 * i + 1] = sy;
+    }
+  }
+
+  // ---- Radix narrowing, MSB -> LSB ----
+  // Invariants entering round r:
+  //   active[i]      <=> key's top 8r bits == threshold's top 8r bits
+  //   total_active    = size of the active set
+  //   topk            = winners still to take from the active set (1..total_active)
+  bool active[kVecSize];
+#pragma unroll
+  for (uint32_t i = 0; i < kVecSize; ++i) {
+    active[i] = true;
+  }
+
+  uint32_t total_active = T::kNumExperts;
+  uint32_t topk = T::kTopK;
+  uint32_t threshold = 0;      // assembled split-key prefix (unexamined low bits zero)
+  uint32_t examined_mask = 0;  // bits of `threshold` that have been fixed
+  bool take_all_equals = false;
+
+  {
+    AlignedVector<uint32_t, 2> zero;
+    zero.fill(0);
+    if (tx < kRadixLanes) zero.store(smem.histogram, tx);
+
+#pragma unroll
+    for (uint32_t round = 0; round < T::kRadixRounds; ++round) {
+      __syncthreads();  // histogram zeroed & previous match consumed
+      const uint32_t shift = 24 - round * 8;
+      uint32_t bin[kVecSize];
+#pragma unroll
+      for (uint32_t i = 0; i < kVecSize; ++i) {
+        bin[i] = (keys[i] >> shift) & 0xff;
+      }
+#pragma unroll
+      for (uint32_t i = 0; i < kVecSize; ++i) {
+        if (active[i]) atomicAdd(&smem.histogram[bin[i]], 1);
+      }
+      __syncthreads();
+
+      // Split-bin search on 128 threads: thread t owns bins {2t, 2t+1}.
+      // The split bin b is the unique bin with above(b) < topk <= above(b) + hist[b].
+      if (tx < kRadixLanes) {
+        AlignedVector<uint32_t, 2> hist;
+        hist.load(smem.histogram, tx);
+        const auto local_val = hist[0] + hist[1];
+        const auto warp_inc = device::warp::inclusive_sum(lane_id, local_val);
+        if (lane_id == kWarpThreads - 1) smem.warp_sum[0][warp_id] = warp_inc;
+        moe::radix::bar_sync(BAR_SUM, kRadixLanes);
+        const auto inter = __reduce_add_sync(0xFFFFFFFF, lane_id < warp_id ? smem.warp_sum[0][lane_id] : 0u);
+        const auto prefix = inter + warp_inc;        // active elements in bins [0, 2t+1]
+        const auto above_r = total_active - prefix;  // in bins > 2t+1
+        const auto above_m = above_r + hist[1];      // in bins > 2t
+        const auto above_l = above_m + hist[0];      // in bins >= 2t
+        if (above_r < topk && above_m >= topk) {
+          smem.match[round] = {tx * 2 + 1, above_r, hist[1]};
+        } else if (above_m < topk && above_l >= topk) {
+          smem.match[round] = {tx * 2 + 0, above_m, hist[0]};
+        }
+      }
+      __syncthreads();
+
+      const auto [threshold_bin, above_count, equal_count] = smem.match[round];
+      threshold |= threshold_bin << shift;
+      examined_mask |= 0xffu << shift;
+#pragma unroll
+      for (uint32_t i = 0; i < kVecSize; ++i) {
+        active[i] &= (bin[i] == threshold_bin);
+      }
+      total_active = equal_count;
+      topk -= above_count;  // split condition guarantees 1 <= topk <= equal_count
+      if (topk == equal_count) {
+        // The remaining quota exactly covers the equal set: every active
+        // element wins, no deeper narrowing or tie-break needed. At the last
+        // round this is the no-full-key-tie case (the typical one).
+        take_all_equals = true;
+        break;
+      }
+      // Re-zero for the next round (synced by the loop-top barrier). Reaching
+      // round 3 with topk < equal_count means a full-key tie: resolved below
+      // by the smallest-id rank among `active`.
+      if (round + 1 < T::kRadixRounds && tx < kRadixLanes) zero.store(smem.histogram, tx);
+    }
+  }
+
+  // ---- Epilogue: collect the K winners ----
+  // Strict winners: examined bits compare above the split prefix (these were
+  // peeled off `active` in earlier rounds). Equal set (== `active`): take all
+  // (take_all_equals) or the `topk` smallest ids (full-key tie-break).
+  bool selected[kVecSize];
+  if (take_all_equals) {
+#pragma unroll
+    for (uint32_t i = 0; i < kVecSize; ++i) {
+      selected[i] = active[i] || (keys[i] & examined_mask) > threshold;
+    }
+  } else {  // deterministic tie-break
+    uint32_t cnt = 0;
+#pragma unroll
+    for (uint32_t i = 0; i < kVecSize; ++i) {
+      cnt += active[i] ? 1 : 0;
+    }
+    uint32_t rank = moe::radix::block_exclusive_sum(cnt, lane_id, warp_id, smem.warp_sum[1]);
+#pragma unroll
+    for (uint32_t i = 0; i < kVecSize; ++i) {
+      const bool eq_win = active[i] && rank < topk;
+      if (active[i]) ++rank;
+      selected[i] = eq_win || (keys[i] & examined_mask) > threshold;
+    }
+  }
+
+  // Compaction slots in expert-id order (deterministic).
+  uint32_t selected_cnt = 0;
+#pragma unroll
+  for (uint32_t i = 0; i < kVecSize; ++i) {
+    selected_cnt += selected[i] ? 1 : 0;
+  }
+  uint32_t slot = moe::radix::block_exclusive_sum(selected_cnt, lane_id, warp_id, smem.warp_sum[2]);
+#pragma unroll
+  for (uint32_t i = 0; i < kVecSize; ++i) {
+    if (selected[i] && slot < T::kTopK) {
+      smem.wid[slot] = (int32_t)(tx * kVecSize + i);
+      smem.wkey[slot] = keys[i];
+      smem.wact[slot] = act[i];
+      ++slot;
+    }
+  }
+  __syncthreads();
+
+  static_assert(T::kTopK <= kWarpThreads);
+  if (tx < T::kTopK) {
+    uint32_t rank = tx;
+    auto w = smem.wact[tx];
+    const auto id = smem.wid[tx];
+    if (params.sorted) {
+      const uint32_t ka = smem.wkey[tx];
+      const int32_t ia = id;
+      rank = 0;
+#pragma unroll
+      for (uint32_t b = 0; b < T::kTopK; ++b) {
+        if (smem.wkey[b] > ka || (smem.wkey[b] == ka && smem.wid[b] < ia)) ++rank;
+      }
+    }
+    PDLTriggerSecondary<kUsePDL>();
+    float sum = 0.f;
+#pragma unroll
+    for (uint32_t i = 0; i < T::kTopK; ++i) {
+      sum += smem.wact[i];
+    }
+    const auto norm = (sum > 0.0f) ? sum : 1.0f;
+    if (params.renormalize) w = w / norm;
+    if (params.apply_scale) w = w * params.routed_scaling_factor;
+    params.out_w[bx * params.out_w_stride + rank] = w;
+    params.out_i[bx * params.out_i_stride + rank] = id;
+    if (params.out_packed != nullptr) {
+      // (id << 16) | bf16(w) bits — RN float->bf16 matches the triton pack.
+      const auto bits = static_cast<uint32_t>(__bfloat16_as_ushort(__float2bfloat16_rn(w)));
+      params.out_packed[bx * params.out_packed_stride + rank] =
+          static_cast<int32_t>((static_cast<uint32_t>(id) << 16) | bits);
+    }
+  }
+}
+
+template <bool kUsePDL, typename TScore>
+__global__ __launch_bounds__(LargeRouterRadixTrait::kBlockSize)  //
+    void route_radix_kernel(const __grid_constant__ RouteRadixParams params) {
+  __shared__ typename LargeRouterRadixTrait::Smem smem;
+  route_radix_block<kUsePDL, TScore>(params, smem);
+}
+
+// ---------------------------------------------------------------------------
+// fused-gate front end. Same radix primitives above; a separate kernel because it
+// folds the gate and the quant epilogue into one launch. Only the module that
+// instantiates it pays for it -- an uninstantiated template costs parse time, not
+// codegen, which is what lets both share this translation unit.
+// ---------------------------------------------------------------------------
+inline constexpr uint32_t kFGTNumExperts = 896;
+inline constexpr uint32_t kFGTTopK = 16;
+// 7 warps: 896 experts / 224 threads = 4 experts per thread in the epilogue, and
+// one expert per warp per pass in phase 1.  Keeping the block at route_radix's
+// shape lets the epilogue reuse its radix-select verbatim.
+/// Block size is a tunable: it sets how many experts each thread owns in the
+/// radix select (kNumExperts / kBlockSize).  It must be at least kRadixSize/2 =
+/// 128 threads (the split-bin search puts 2 bins per thread) and must divide the
+/// expert count into an even per-thread count (the loads are fp32x2 pairs), so
+/// 224 (4 experts/thread) and 448 (2 experts/thread) are the legal choices for
+/// 896 experts.
+template <uint32_t kBlockSize_>
+struct MoEFrontTrait : moe::radix::RadixSelectBase {
+  static constexpr uint32_t kNumExperts = kFGTNumExperts;
+  static constexpr uint32_t kTopK = kFGTTopK;
+  static constexpr uint32_t kBlockSize = kBlockSize_;
+  static constexpr uint32_t kVecSize = kNumExperts / kBlockSize;  // experts per thread
+  static constexpr uint32_t kNumWarps = kBlockSize / 32;
+
+  static_assert(kNumExperts % kBlockSize == 0, "block size must divide the expert count");
+  static_assert(kVecSize % 2 == 0, "experts per thread must be even (fp32x2 loads)");
+  static_assert(kBlockSize >= kRadixSize / 2, "block must cover the split-bin search lanes");
+
+  struct Smem {
+    uint32_t warp_sum[3][kNumWarps];
+    MatchBin match[kRadixRounds];
+    uint32_t histogram[kRadixSize];
+    int32_t wid[kTopK];
+    uint32_t wkey[kTopK];
+    fp32_t wact[kTopK];
+  };
+};
+
+struct MoEFrontParams {
+  const fp32_t* __restrict__ bias;    // [E] fp32 correction bias
+  const fp32_t* __restrict__ logits;  // [M, logits_stride] fp32, gate slice first
+  fp32_t* __restrict__ out_w;         // [M, topk] fp32
+  int32_t* __restrict__ out_i;        // [M, topk] int32
+  int M;
+  int logits_stride;  // E for the router-only entry, E + latent for the front
+  long long out_w_stride;
+  long long out_i_stride;
+  float routed_scaling_factor;
+  int renormalize;
+  int apply_scale;
+  // Merged-front entry only: the [M, latent] bf16 routed_input to emit.
+  bf16_t* __restrict__ routed_out;
+  int latent;
+  long long routed_stride;
+};
+
+/// Radix-select top-k over one token's fp32 logits.  Lifted from
+/// route_radix.cuh; the only change is the input dtype (fp32 in place of bf16).
+template <bool kUsePDL, typename T>
+SGL_DEVICE void fgt_select_topk(
+    const MoEFrontParams& params, typename T::Smem& smem, int m, uint32_t tx, uint32_t warp_id, uint32_t lane_id) {
+  constexpr uint32_t kVecSize = T::kVecSize;
+  constexpr uint32_t kRadixLanes = T::kRadixSize / 2;
+  enum { BAR_SUM = 1 };
+
+  uint32_t keys[kVecSize];
+  float act[kVecSize];
+  {
+    // 4 experts per thread as two fp32x2 (16B) loads, matching route_radix.
+    device::AlignedVector<fp32x2_t, kVecSize / 2> bias_vec;
+    bias_vec.load(params.bias, tx);
+    device::AlignedVector<fp32x2_t, kVecSize / 2> lv;
+    lv.load(params.logits + (long long)m * params.logits_stride, tx);
+    float logit[kVecSize];
+#pragma unroll
+    for (uint32_t i = 0; i < kVecSize / 2; ++i) {
+      logit[2 * i + 0] = lv[i].x;
+      logit[2 * i + 1] = lv[i].y;
+    }
+#pragma unroll
+    for (uint32_t i = 0; i < kVecSize / 2; ++i) {
+      const float sx = moe::radix::sigmoid_match(logit[2 * i + 0]);
+      const float sy = moe::radix::sigmoid_match(logit[2 * i + 1]);
+      act[2 * i + 0] = sx;
+      act[2 * i + 1] = sy;
+      keys[2 * i + 0] = moe::radix::biased_to_key(moe::radix::nan_floor(sx + bias_vec[i].x));
+      keys[2 * i + 1] = moe::radix::biased_to_key(moe::radix::nan_floor(sy + bias_vec[i].y));
+    }
+  }
+
+  bool active[kVecSize];
+#pragma unroll
+  for (uint32_t i = 0; i < kVecSize; ++i) {
+    active[i] = true;
+  }
+
+  uint32_t total_active = T::kNumExperts;
+  uint32_t topk = T::kTopK;
+  uint32_t threshold = 0;
+  uint32_t examined_mask = 0;
+  bool take_all_equals = false;
+
+  {
+    device::AlignedVector<uint32_t, 2> zero;
+    zero.fill(0);
+    if (tx < kRadixLanes) zero.store(smem.histogram, tx);
+
+#pragma unroll
+    for (uint32_t round = 0; round < T::kRadixRounds; ++round) {
+      __syncthreads();
+      const uint32_t shift = 24 - round * 8;
+      uint32_t bin[kVecSize];
+#pragma unroll
+      for (uint32_t i = 0; i < kVecSize; ++i) {
+        bin[i] = (keys[i] >> shift) & 0xff;
+      }
+#pragma unroll
+      for (uint32_t i = 0; i < kVecSize; ++i) {
+        if (active[i]) atomicAdd(&smem.histogram[bin[i]], 1);
+      }
+      __syncthreads();
+
+      if (tx < kRadixLanes) {
+        device::AlignedVector<uint32_t, 2> hist;
+        hist.load(smem.histogram, tx);
+        const auto local_val = hist[0] + hist[1];
+        const auto warp_inc = device::warp::inclusive_sum(lane_id, local_val);
+        if (lane_id == 31) smem.warp_sum[0][warp_id] = warp_inc;
+        moe::radix::bar_sync(BAR_SUM, kRadixLanes);
+        const auto inter = __reduce_add_sync(0xFFFFFFFF, lane_id < warp_id ? smem.warp_sum[0][lane_id] : 0u);
+        const auto prefix = inter + warp_inc;
+        const auto above_r = total_active - prefix;
+        const auto above_m = above_r + hist[1];
+        const auto above_l = above_m + hist[0];
+        if (above_r < topk && above_m >= topk) {
+          smem.match[round] = {tx * 2 + 1, above_r, hist[1]};
+        } else if (above_m < topk && above_l >= topk) {
+          smem.match[round] = {tx * 2 + 0, above_m, hist[0]};
+        }
+      }
+      __syncthreads();
+
+      const auto [threshold_bin, above_count, equal_count] = smem.match[round];
+      threshold |= threshold_bin << shift;
+      examined_mask |= 0xffu << shift;
+#pragma unroll
+      for (uint32_t i = 0; i < kVecSize; ++i) {
+        active[i] &= (bin[i] == threshold_bin);
+      }
+      total_active = equal_count;
+      topk -= above_count;
+      if (topk == equal_count) {
+        take_all_equals = true;
+        break;
+      }
+      if (round + 1 < T::kRadixRounds && tx < kRadixLanes) zero.store(smem.histogram, tx);
+    }
+  }
+
+  bool selected[kVecSize];
+  if (take_all_equals) {
+#pragma unroll
+    for (uint32_t i = 0; i < kVecSize; ++i) {
+      selected[i] = active[i] || (keys[i] & examined_mask) > threshold;
+    }
+  } else {
+    uint32_t cnt = 0;
+#pragma unroll
+    for (uint32_t i = 0; i < kVecSize; ++i) {
+      cnt += active[i] ? 1 : 0;
+    }
+    uint32_t rank = moe::radix::block_exclusive_sum(cnt, lane_id, warp_id, smem.warp_sum[1]);
+#pragma unroll
+    for (uint32_t i = 0; i < kVecSize; ++i) {
+      const bool eq_win = active[i] && rank < topk;
+      if (active[i]) ++rank;
+      selected[i] = eq_win || (keys[i] & examined_mask) > threshold;
+    }
+  }
+
+  uint32_t selected_cnt = 0;
+#pragma unroll
+  for (uint32_t i = 0; i < kVecSize; ++i) {
+    selected_cnt += selected[i] ? 1 : 0;
+  }
+  uint32_t slot = moe::radix::block_exclusive_sum(selected_cnt, lane_id, warp_id, smem.warp_sum[2]);
+#pragma unroll
+  for (uint32_t i = 0; i < kVecSize; ++i) {
+    if (selected[i] && slot < T::kTopK) {
+      smem.wid[slot] = (int32_t)(tx * kVecSize + i);
+      smem.wact[slot] = act[i];
+      ++slot;
+    }
+  }
+  __syncthreads();
+
+  static_assert(T::kTopK <= 32);
+  if (tx < T::kTopK) {
+    auto wv = smem.wact[tx];
+    const auto id = smem.wid[tx];
+    float sum = 0.f;
+#pragma unroll
+    for (uint32_t i = 0; i < T::kTopK; ++i) {
+      sum += smem.wact[i];
+    }
+    const auto norm = (sum > 0.0f) ? sum : 1.0f;
+    if (params.renormalize) wv = wv / norm;
+    if (params.apply_scale) wv = wv * params.routed_scaling_factor;
+    params.out_w[m * params.out_w_stride + tx] = wv;
+    params.out_i[m * params.out_i_stride + tx] = id;
+  }
+  __syncthreads();  // smem reuse across the token loop
+}
+
+/// Tunables: `kBlockSize` sets the experts-per-thread of the radix select,
+/// `kCastVec` the fp32 elements each thread converts per step (the cast moves
+/// [T, 3584] fp32 in and bf16 out, which dominates the epilogue at large T), and
+/// `kCastFirst` whether the cast is issued before the select (loads in flight
+/// during the radix rounds) or after it.
+template <bool kUsePDL, uint32_t kBlockSize, uint32_t kCastVec, bool kCastFirst>
+__global__ __launch_bounds__(kBlockSize)  //
+    void fused_front_epilogue_kernel(const __grid_constant__ MoEFrontParams params) {
+  using namespace device;
+  using T = MoEFrontTrait<kBlockSize>;
+  __shared__ typename T::Smem smem;
+  const uint32_t tx = threadIdx.x;
+  const int m = (int)blockIdx.x;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  // Cast the latent slice: [E, E + latent) fp32 -> [0, latent) bf16.
+  auto cast_latent = [&]() {
+    const fp32_t* src = params.logits + (long long)m * params.logits_stride + T::kNumExperts;
+    bf16_t* dst = params.routed_out + (long long)m * params.routed_stride;
+    for (int i = (int)tx * kCastVec; i < params.latent; i += (int)kBlockSize * kCastVec) {
+      AlignedVector<fp32x2_t, kCastVec / 2> v;
+      v.load(src, i / kCastVec);
+      AlignedVector<bf16_t, kCastVec> o;
+#pragma unroll
+      for (uint32_t j = 0; j < kCastVec / 2; ++j) {
+        o[2 * j + 0] = cast<bf16_t>(v[j].x);
+        o[2 * j + 1] = cast<bf16_t>(v[j].y);
+      }
+      o.store(dst, i / kCastVec);
+    }
+  };
+
+  if (kCastFirst) cast_latent();
+  fgt_select_topk<kUsePDL, T>(params, smem, m, tx, tx / 32, tx % 32);
+  if (!kCastFirst) cast_latent();
+}
+
+}  // namespace sglang
+
+template <bool kUsePDL>
+struct RouteRadixKernel {
+  static void
+  run(const tvm::ffi::TensorView scores,
+      const tvm::ffi::TensorView bias,
+      const tvm::ffi::TensorView out_w,
+      const tvm::ffi::TensorView out_i,
+      int64_t topk,
+      double routed_scaling_factor,
+      bool renormalize,
+      bool apply_scale,
+      bool sorted) {
+    using namespace host;
+
+    auto M_ = SymbolicSize{"num_tokens"};
+    auto N_ = SymbolicSize{"num_experts"};
+    auto K_ = SymbolicSize{"topk"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    auto score_dtype = SymbolicDType{};
+    TensorMatcher({M_, N_})
+        .with_dtype<bf16_t, fp32_t>(score_dtype)
+        .with_device(device)
+        .with_strides({-1, 1})
+        .verify(scores);
+    TensorMatcher({N_}).with_dtype<fp32_t>().with_device(device).verify(bias);
+    TensorMatcher({M_, K_}).with_dtype<fp32_t>().with_device(device).verify(out_w);
+    TensorMatcher({M_, K_}).with_dtype<int32_t>().with_device(device).verify(out_i);
+
+    RuntimeCheck(
+        N_.unwrap() == sglang::kNumExperts_ && K_.unwrap() == sglang::kTopK_ && topk == sglang::kTopK_,
+        "route_radix is specialized for N=896, K=16");
+    // Vectorized row loads (8B for bf16, 16B for fp32) need aligned row
+    // starts; stride % 4 elements covers both (4 x 2B = 8B / 4 x 4B = 16B).
+    RuntimeCheck(scores.stride(0) % 4 == 0, "route_radix: scores row stride must be a multiple of 4");
+
+    const auto M = static_cast<uint32_t>(M_.unwrap());
+    if (M == 0) return;
+
+    const auto params = sglang::RouteRadixParams{
+        scores.data_ptr(),
+        static_cast<const fp32_t*>(bias.data_ptr()),
+        static_cast<fp32_t*>(out_w.data_ptr()),
+        static_cast<int32_t*>(out_i.data_ptr()),
+        /*out_packed=*/nullptr,
+        static_cast<int>(M),
+        static_cast<long long>(scores.stride(0)),
+        static_cast<long long>(out_w.stride(0)),
+        static_cast<long long>(out_i.stride(0)),
+        /*out_packed_stride=*/0,
+        static_cast<float>(routed_scaling_factor),
+        renormalize ? 1 : 0,
+        apply_scale ? 1 : 0,
+        sorted ? 1 : 0};
+
+    if (score_dtype.is_type<fp32_t>()) {
+      LaunchKernel(M, sglang::LargeRouterRadixTrait::kBlockSize, device.unwrap())
+          .enable_pdl(kUsePDL)(sglang::route_radix_kernel<kUsePDL, fp32_t>, params);
+    } else {
+      LaunchKernel(M, sglang::LargeRouterRadixTrait::kBlockSize, device.unwrap())
+          .enable_pdl(kUsePDL)(sglang::route_radix_kernel<kUsePDL, bf16_t>, params);
+    }
+  }
+};
+
+template <bool kUsePDL>
+struct FusedFrontEpilogueKernel {
+  static void
+  run(const tvm::ffi::TensorView merged,  // [M, E + latent] fp32, row-dense
+      const tvm::ffi::TensorView bias,    // [E] fp32
+      const tvm::ffi::TensorView out_w,   // [M, topk] fp32
+      const tvm::ffi::TensorView out_i,   // [M, topk] int32
+      const tvm::ffi::TensorView routed,  // [M, latent] bf16
+      int64_t topk,
+      double routed_scaling_factor,
+      bool renormalize,
+      bool apply_scale,
+      int64_t block_size,
+      int64_t cast_vec,
+      bool cast_first) {
+    using namespace host;
+
+    auto M_ = SymbolicSize{"num_tokens"};
+    auto W_ = SymbolicSize{"merged_width"};
+    auto L_ = SymbolicSize{"latent"};
+    auto E_ = SymbolicSize{"num_experts"};
+    auto K_ = SymbolicSize{"topk"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({M_, W_}).with_dtype<fp32_t>().with_device(device).with_strides({-1, 1}).verify(merged);
+    TensorMatcher({E_}).with_dtype<fp32_t>().with_device(device).verify(bias);
+    TensorMatcher({M_, K_}).with_dtype<fp32_t>().with_device(device).verify(out_w);
+    TensorMatcher({M_, K_}).with_dtype<int32_t>().with_device(device).verify(out_i);
+    TensorMatcher({M_, L_}).with_dtype<bf16_t>().with_device(device).with_strides({-1, 1}).verify(routed);
+
+    const auto M = static_cast<int>(M_.unwrap());
+    const auto latent = static_cast<int>(L_.unwrap());
+    RuntimeCheck(
+        E_.unwrap() == sglang::kFGTNumExperts && K_.unwrap() == sglang::kFGTTopK && topk == sglang::kFGTTopK,
+        "fused_front_epilogue is specialized for E=896, topk=16");
+    RuntimeCheck(
+        static_cast<int>(W_.unwrap()) == static_cast<int>(sglang::kFGTNumExperts) + latent,
+        "fused_front_epilogue: merged width must be num_experts + latent");
+    // 16B vectorized reads of the fp32 rows and 8B writes of the bf16 rows.
+    RuntimeCheck(latent % 4 == 0, "fused_front_epilogue: latent must be a multiple of 4");
+    RuntimeCheck(
+        merged.stride(0) % 4 == 0 && routed.stride(0) % 4 == 0,
+        "fused_front_epilogue: row strides must be a multiple of 4");
+    if (M == 0) return;
+
+    auto params = sglang::MoEFrontParams{};
+    params.bias = static_cast<const fp32_t*>(bias.data_ptr());
+    params.logits = static_cast<fp32_t*>(merged.data_ptr());
+    params.out_w = static_cast<fp32_t*>(out_w.data_ptr());
+    params.out_i = static_cast<int32_t*>(out_i.data_ptr());
+    params.M = M;
+    params.out_w_stride = static_cast<long long>(out_w.stride(0));
+    params.out_i_stride = static_cast<long long>(out_i.stride(0));
+    params.routed_scaling_factor = static_cast<float>(routed_scaling_factor);
+    params.renormalize = renormalize ? 1 : 0;
+    params.apply_scale = apply_scale ? 1 : 0;
+    params.logits_stride = static_cast<int>(merged.stride(0));
+    params.routed_out = static_cast<bf16_t*>(routed.data_ptr());
+    params.latent = latent;
+    params.routed_stride = static_cast<long long>(routed.stride(0));
+
+    // Tunables come from the JSON config table; see kernels/ops/moe/moe_front.py.
+    // cast_vec * 4 bytes per thread must stay inside the 32B vector-load limit.
+    RuntimeCheck(cast_vec == 2 || cast_vec == 4 || cast_vec == 8, "fused_front_epilogue: cast_vec must be 2, 4 or 8");
+    RuntimeCheck(latent % cast_vec == 0, "fused_front_epilogue: cast_vec must divide latent");
+
+#define SGL_FRONT_LAUNCH(BS, CV, CF)   \
+  LaunchKernel(M, BS, device.unwrap()) \
+      .enable_pdl(kUsePDL)(sglang::fused_front_epilogue_kernel<kUsePDL, BS, CV, CF>, params)
+#define SGL_FRONT_DISPATCH_CV(BS, CF) \
+  do {                                \
+    if (cast_vec == 2) {              \
+      SGL_FRONT_LAUNCH(BS, 2, CF);    \
+    } else if (cast_vec == 4) {       \
+      SGL_FRONT_LAUNCH(BS, 4, CF);    \
+    } else {                          \
+      SGL_FRONT_LAUNCH(BS, 8, CF);    \
+    }                                 \
+  } while (0)
+#define SGL_FRONT_DISPATCH_BS(CF)     \
+  do {                                \
+    if (block_size == 448) {          \
+      SGL_FRONT_DISPATCH_CV(448, CF); \
+    } else {                          \
+      SGL_FRONT_DISPATCH_CV(224, CF); \
+    }                                 \
+  } while (0)
+
+    RuntimeCheck(block_size == 224 || block_size == 448, "fused_front_epilogue: block_size must be 224 or 448");
+    if (cast_first) {
+      SGL_FRONT_DISPATCH_BS(true);
+    } else {
+      SGL_FRONT_DISPATCH_BS(false);
+    }
+#undef SGL_FRONT_DISPATCH_BS
+#undef SGL_FRONT_DISPATCH_CV
+#undef SGL_FRONT_LAUNCH
+  }
+};

python/sglang/kernels/jit/csrc/moe/topk_sum.cuh

@@ -0,0 +1,103 @@
+// Top-k expert-output sum: out[M, K] = sum_j in[M, topk, K].
+//
+// Replaces sgl_kernel's moe_sum_reduce_kernel_general (~5.7us at decode
+// shapes [1, 16, 3584]) with a straightforward vectorized pass (~1.5us):
+// one thread per 8-element vector of K, looping the topk rows in fp32.
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck
+
+#include <sgl_kernel/type.cuh>   // For bf16_t, fp32_t, device::cast
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel
+#include <sgl_kernel/vec.cuh>    // For AlignedVector
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+struct TopkSumParams {
+  const bf16_t* __restrict__ in;  // [M, topk, K] contiguous
+  bf16_t* __restrict__ out;       // [M, K] contiguous
+  uint32_t K;
+  uint32_t topk;
+};
+
+template <int kThreads, bool kUsePDL>
+__global__ void topk_sum_kernel(const TopkSumParams __grid_constant__ params) {
+  using namespace device;
+
+  constexpr int kVecN = 8;  // 8 bf16 = 128 bits
+  using vec_bf16_t = AlignedVector<bf16_t, kVecN>;
+
+  const uint32_t m = blockIdx.y;
+  const uint32_t v = blockIdx.x * kThreads + threadIdx.x;
+  const uint32_t n_vecs = params.K / kVecN;
+  if (v >= n_vecs) return;
+
+  const bf16_t* base = params.in + static_cast<int64_t>(m) * params.topk * params.K;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  float acc[kVecN];
+#pragma unroll
+  for (int i = 0; i < kVecN; ++i) {
+    acc[i] = 0.0f;
+  }
+  for (uint32_t j = 0; j < params.topk; ++j) {
+    vec_bf16_t x;
+    x.load(base + static_cast<int64_t>(j) * params.K, v);
+#pragma unroll
+    for (int i = 0; i < kVecN; ++i) {
+      acc[i] += cast<fp32_t>(x[i]);
+    }
+  }
+
+  vec_bf16_t o;
+#pragma unroll
+  for (int i = 0; i < kVecN; ++i) {
+    o[i] = cast<bf16_t>(acc[i]);
+  }
+  o.store(params.out + static_cast<int64_t>(m) * params.K, v);
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int kThreads, bool kUsePDL>
+struct TopkSumKernel {
+  static constexpr auto kernel = topk_sum_kernel<kThreads, kUsePDL>;
+
+  static void run(const tvm::ffi::TensorView in, const tvm::ffi::TensorView out) {
+    using namespace host;
+
+    auto M_ = SymbolicSize{"num_tokens"};
+    auto T_ = SymbolicSize{"topk"};
+    auto K_ = SymbolicSize{"hidden"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({M_, T_, K_}).with_dtype<bf16_t>().with_device(device).verify(in);
+    TensorMatcher({M_, K_}).with_dtype<bf16_t>().with_device(device).verify(out);
+
+    const auto M = static_cast<uint32_t>(M_.unwrap());
+    const auto topk = static_cast<uint32_t>(T_.unwrap());
+    const auto K = static_cast<uint32_t>(K_.unwrap());
+
+    RuntimeCheck(K % 8 == 0, "K must be divisible by 8 for vectorized loads");
+    if (M == 0) return;
+
+    const auto params = TopkSumParams{
+        .in = static_cast<const bf16_t*>(in.data_ptr()),
+        .out = static_cast<bf16_t*>(out.data_ptr()),
+        .K = K,
+        .topk = topk,
+    };
+
+    const uint32_t n_vecs = K / 8;
+    dim3 grid((n_vecs + kThreads - 1) / kThreads, M);
+    LaunchKernel(grid, kThreads, device.unwrap()).enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/kernels/jit/include/sgl_kernel/math.cuh

@@ -58,6 +58,24 @@ SGL_DEVICE T exp(T a) {
   return DTypeTrait<T>::exp(a);
 }
 
+/// \brief Fast approximate sigmoid for FP32 device code.
+SGL_DEVICE float sigmoid_fast(float x) {
+  return 1.0f / (1.0f + __expf(-x));
+}
+
+/// \brief Fast approximate SiLU for FP32 device code.
+SGL_DEVICE float silu_fast(float x) {
+  return x * sigmoid_fast(x);
+}
+
+/// \brief Fast approximate softplus for FP32 device code.
+///
+/// Values above 20 use the asymptotic result directly, avoiding overflow and
+/// an unnecessary exponential while matching common softplus kernels.
+SGL_DEVICE float softplus_fast(float x) {
+  return x > 20.0f ? x : log1pf(__expf(x));
+}
+
 /// \brief Returns sin(a).
 template <typename T>
 SGL_DEVICE T sin(T a) {
@@ -70,4 +88,20 @@ SGL_DEVICE T cos(T a) {
   return DTypeTrait<T>::cos(a);
 }
 
+// bf16 x bf16 -> fp32 fused multiply-add The mixed-precision PTX
+// instruction saves the explicit converts; the fallback is bit-identical (the
+// bf16 -> f32 conversion is exact, both round once). Shared by tiny_gemm,
+// gemm_ag and ar_fusion.
+SGL_DEVICE float fma_f32_bf16(bf16_t a, bf16_t b, float acc) {
+#if SGL_ARCH_BLACKWELL_OR_GREATER
+  const uint16_t a_bits = __bfloat16_as_ushort(a);
+  const uint16_t b_bits = __bfloat16_as_ushort(b);
+  float result;
+  asm("fma.rn.f32.bf16 %0, %1, %2, %3;" : "=f"(result) : "h"(a_bits), "h"(b_bits), "f"(acc));
+  return result;
+#else
+  return fmaf(cast<fp32_t>(a), cast<fp32_t>(b), acc);
+#endif
+}
+
 }  // namespace device::math

python/sglang/kernels/jit/include/sgl_kernel/mbarrier.cuh

@@ -0,0 +1,68 @@
+#pragma once
+
+// mbarrier PTX wrappers shared by the Kimi K3 kernels that drive TMA by hand:
+// kimi_k3/comm/gemm_ar.cuh and kimi_k3/attn_res/fused_tma.cuh both defined these
+// with identical bodies.
+//
+// The enclosing namespace is the same global `ptx` both files already open, so
+// existing `::ptx::mbar_*` call sites need no change.
+//
+// gemm_ar.cuh keeps mbar_arrive_cluster_release: only it uses that one.
+// attention/kda_prefill.cu duplicates a different set (MMA / ldmatrix) but is
+// built without a sglang include path and cannot consume this header.
+
+#include <sgl_kernel/utils.cuh>
+
+#include <cstdint>
+
+namespace ptx {
+
+// Inline-PTX `.shared` instructions take a 32-bit byte offset in the shared
+// window, not a generic 64-bit pointer.
+template <typename T>
+static SGL_DEVICE uint32_t to_shared(T* ptr) {
+  return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
+}
+
+// ---- mbarrier (PTX ISA §9.7.13.15) -----------------------------------------
+//
+// Only the `try_wait.parity` waiter is wrapped: state-token waits couple
+// arriver and waiter, and mixing state- and parity-tracked codepaths in one
+// kernel is a deadlock risk. The caller owns the phase counter and flips it at
+// the stage wrap (`phase ^= (stage == 0)`).
+//
+// Initial parity, the easy-to-flip part: after `mbar_init` the bar is at
+// parity 0, and each full cycle (count arrivals -> fire -> reset) flips it.
+//   - consumer-first (waits for an external producer's first signal) -> 0.
+//   - producer-first (waits for a consumer to release a slot that no consumer
+//     has touched yet) -> 1, so the first wait is a no-op skip.
+// A consumer-first wait initialized to 1 skips the producer's first signal and
+// blocks forever on the second.
+static SGL_DEVICE void mbar_init(uint64_t* bar, uint32_t count) {
+  asm volatile("mbarrier.init.shared.b64 [%0], %1;" ::"r"(to_shared(bar)), "r"(count));
+}
+
+static SGL_DEVICE uint64_t mbar_arrive(uint64_t* bar) {
+  uint64_t state;
+  asm volatile("mbarrier.arrive.shared.b64 %0, [%1];" : "=l"(state) : "r"(to_shared(bar)));
+  return state;
+}
+
+// Combined arrive + set tx-count, for TMA-load completion.
+static SGL_DEVICE void mbar_arrive_expect_tx(uint64_t* bar, uint32_t bytes) {
+  asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;" ::"r"(to_shared(bar)), "r"(bytes));
+}
+
+// Wait for phase `parity` to complete. Looped because the spec allows spurious
+// early wakeups. Default `.acquire` semantics mean prior `cp.async.bulk` writes
+// tracked by this mbarrier are visible to later generic-proxy reads on this
+// thread with no `fence.proxy.async` (spec §9.7.13.15.16 point 3).
+static SGL_DEVICE void mbar_wait_parity(uint64_t* bar, uint32_t parity) {
+  asm volatile(
+      "{\n\t.reg .pred p;\n\t"
+      "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1;\n\t"
+      "@!p bra WAIT_%=;\n\t}\n" ::"r"(to_shared(bar)),
+      "r"(parity));
+}
+
+}  // namespace ptx

python/sglang/kernels/jit/include/sgl_kernel/warp.cuh

@@ -1,9 +1,11 @@
 /// \file warp.cuh
-/// \brief Warp-level reduction primitives.
+/// \brief Warp-level reduction and cooperative-copy primitives.
 
 #pragma once
 #include <sgl_kernel/math.cuh>
+#include <sgl_kernel/tile.cuh>
 #include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
 
 #include <cstdint>
 #include <type_traits>
@@ -119,4 +121,87 @@ SGL_DEVICE T reduce_min(T value, mask_t active_mask = kFullMask) {
   return reduce<ReductionOp::MIN, kNumThreads, kInner>(value, active_mask);
 }
 
+/// \brief Warp-cooperative gmem -> smem copy of a compile-time byte count.
+///
+/// Picks the widest vector width that divides both the per-thread share and
+/// the byte total. The caller guarantees ``src`` is aligned to the picked
+/// width (16B for kBytes % (16*32) == 0, else 8/4) and ``dst`` is the start
+/// of a 16B-aligned per-warp smem slot.
+// Warp-cooperative byte copy between any two address spaces, vectorised to the
+// widest unit `kBytes` allows. Named for what it does rather than where it is
+// used: the MLA call sites happen to target shared memory, but nothing here is
+// global->shared specific -- no cp.async, no TMA, the payload moves through
+// registers.
+//
+// The strategy was measured against the two async alternatives on B300 (sm_103,
+// 148 SMs), copying one MLA row per warp out of a 512 MB pool so every row
+// streams from HBM (grid 296, 64 rows/warp, 50 launches):
+//
+//   1152 B/warp (bf16, nope 1024 + rope 128)   576 B/warp (fp8, 512 + 64)
+//     this (generic)     47.3 us  3.69 TB/s      40.8 us  2.14 TB/s
+//     cp.async (ldgsts)  73.9 us  2.36 TB/s      56.4 us  1.55 TB/s
+//     cp.async.bulk/TMA  50.0 us  3.50 TB/s      43.2 us  2.02 TB/s
+//
+// The generic path wins at both sizes: a ~1 KB row is too small to amortise
+// cp.async's per-lane 16 B issues or TMA's fixed issue plus mbarrier round trip.
+// Revisit if a call site ever copies substantially more than one row per warp.
+template <int64_t kBytes>
+SGL_DEVICE void copy_bytes(const void* __restrict__ src, void* __restrict__ dst) {
+  constexpr int64_t kAlignment = (kBytes % (16 * kWarpThreads) == 0)  ? 16
+                                 : (kBytes % (8 * kWarpThreads) == 0) ? 8
+                                 : (kBytes % (4 * kWarpThreads) == 0) ? 4
+                                 : (kBytes % 4 == 0)                  ? 4
+                                                                      : 0;
+  static_assert(kAlignment > 0, "kBytes must be a multiple of 4");
+
+  using vec_t = AlignedStorage<uint32_t, kAlignment / 4>;
+  constexpr auto kLoopBytes = sizeof(vec_t) * kWarpThreads;
+  constexpr auto kLoopCount = kBytes / kLoopBytes;
+  constexpr int64_t kTailVecs = (kBytes - kLoopCount * kLoopBytes) / sizeof(vec_t);
+
+  const auto gmem = tile::Memory<vec_t>::warp();
+
+#pragma unroll
+  for (int64_t i = 0; i < kLoopCount; ++i) {
+    const auto v = gmem.load(src, i);
+    gmem.store(dst, v, i);
+  }
+  if constexpr (kTailVecs > 0) {
+    if (gmem.in_bound(kLoopCount * kWarpThreads + kTailVecs, kLoopCount)) {
+      const auto v = gmem.load(src, kLoopCount);
+      gmem.store(dst, v, kLoopCount);
+    }
+  }
+}
+
+/// Inclusive prefix sum across one warp, thread-rank order. Distinct from
+/// reduce_sum above: every lane keeps its own running total rather than the
+/// whole-warp result.
+SGL_DEVICE uint32_t inclusive_sum(uint32_t lane_id, uint32_t val) {
+  static_assert(kWarpThreads == 32);
+#pragma unroll
+  for (uint32_t offset = 1; offset < 32; offset *= 2) {
+    uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset);
+    if (lane_id >= offset) val += n;
+  }
+  return val;
+}
+
+// One elected lane, via elect.sync. Raw PTX rather than cute::elect_one_sync,
+// which would drag the whole CuTe include path into elementwise JIT modules;
+// cuda::ptx has no elect_sync in CUDA 13.0. Use this to gate a single-thread
+// TMA issue instead of a lane-index predicate.
+SGL_DEVICE bool elect_one_lane() {
+  uint32_t pred;
+  asm volatile(
+      "{\n"
+      "  .reg .pred p;\n"
+      "  .reg .b32  r;\n"
+      "  elect.sync r|p, 0xFFFFFFFF;\n"
+      "  selp.b32 %0, 1, 0, p;\n"
+      "}\n"
+      : "=r"(pred));
+  return pred != 0;
+}
+
 }  // namespace device::warp

python/sglang/kernels/jit/utils/compile.py

@@ -173,6 +173,8 @@ def load_jit(
     *args: str,
     cpp_files: List[str] | None = None,
     cuda_files: List[str] | None = None,
+    external_cpp_files: List[str] | None = None,
+    external_cuda_files: List[str] | None = None,
     cpp_wrappers: List[Tuple[str, str]] | None = None,
     cuda_wrappers: List[Tuple[str, str]] | None = None,
     extra_cflags: List[str] | None = None,
@@ -195,6 +197,12 @@ def load_jit(
     :type cpp_files: List[str] | None
     :param cuda_files: A list of CUDA source files.
     :type cuda_files: List[str] | None
+    :param external_cpp_files: A list of caller-resolved C++ source paths outside
+                               the in-tree JIT source directory.
+    :type external_cpp_files: List[str] | None
+    :param external_cuda_files: A list of caller-resolved CUDA source paths outside
+                                the in-tree JIT source directory.
+    :type external_cuda_files: List[str] | None
     :param cpp_wrappers: A list of C++ wrappers, defining the export name and kernel name.
     :type cpp_wrappers: List[Tuple[str, str]] | None
     :param cuda_wrappers: A list of CUDA wrappers, defining the export name and kernel name.
@@ -222,13 +230,26 @@ def load_jit(
 
     cpp_files = cpp_files or []
     cuda_files = cuda_files or []
+    external_cpp_files = external_cpp_files or []
+    external_cuda_files = external_cuda_files or []
     extra_cflags = extra_cflags or []
     extra_cuda_cflags = extra_cuda_cflags or []
     extra_ldflags = extra_ldflags or []
     extra_include_paths = extra_include_paths or []
 
-    cpp_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cpp_files]
-    cuda_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cuda_files]
+    if torch.version.hip is not None:
+        extra_cuda_cflags = [
+            flag
+            for flag in extra_cuda_cflags
+            if flag not in ("--use_fast_math", "-use_fast_math")
+        ]
+
+    cpp_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cpp_files] + [
+        str(pathlib.Path(f).resolve()) for f in external_cpp_files
+    ]
+    cuda_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cuda_files] + [
+        str(pathlib.Path(f).resolve()) for f in external_cuda_files
+    ]
 
     for dep in set(extra_dependencies or []):
         if dep not in REGISTERED_DEPENDENCIES:

python/sglang/kernels/ops/__init__.py

@@ -23,10 +23,12 @@ _GROUPS = (
     "embeddings",
     "gemm",
     "grammar",
+    "kimi_k3",
     "kvcache",
     "layernorm",
     "mamba",
     "memory",
+    "mm",
     "moe",
     "quantization",
     "sampling",

python/sglang/kernels/ops/attention/concat_mla.py

@@ -4,7 +4,12 @@ from typing import TYPE_CHECKING
 
 import torch
 
-from sglang.kernels.jit.utils import cache_once, load_jit
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module
@@ -21,10 +26,12 @@ def _jit_concat_mla_k_module() -> Module:
 
 @cache_once
 def _jit_concat_mla_absorb_q_module() -> Module:
+    args = make_cpp_args(is_arch_support_pdl())
     return load_jit(
         "concat_mla_absorb_q",
+        *args,
         cuda_files=["elementwise/concat_mla.cuh"],
-        cuda_wrappers=[("concat_mla_absorb_q", "ConcatMlaAbsorbQKernel::run")],
+        cuda_wrappers=[("concat_mla_absorb_q", f"ConcatMlaAbsorbQKernel<{args}>::run")],
     )
 
 

python/sglang/kernels/ops/attention/cutedsl_kda.py

@@ -94,6 +94,10 @@ def _define_kernels():
         pool_idx = h0_indices[i_n]
 
         if pool_idx >= 0:
+            # State indexing in int64: envelope-strided pools (unified memory /
+            # page-major) put pool_idx * stride(0) far past 2^31 (the CuTe twin
+            # of the fla/chunk_delta_h.py stride_init_state fix).
+            pool_idx64 = cutlass.Int64(pool_idx)
             k_local = in_warp_tid // V_PER_WARP_SMALL
             v_local = in_warp_tid % V_PER_WARP_SMALL
             v_base = warp_idx * V_PER_WARP_SMALL
@@ -206,7 +210,7 @@ def _define_kernels():
                         h_val = 0.0
                         if v_global_load < v.shape[3]:
                             h_val = cutlass.Float32(
-                                h0_source[(pool_idx, i_hv, v_global_load, k_load)]
+                                h0_source[(pool_idx64, i_hv, v_global_load, k_load)]
                             )
                         sData[(k_load, v_load, stage)] = h_val
 
@@ -263,7 +267,7 @@ def _define_kernels():
                     if k_write < TILE_K:
                         v_global_write = v_tile * TILE_V_SMALL + v_write
                         if v_global_write < v.shape[3]:
-                            h0_source[(pool_idx, i_hv, v_global_write, k_write)] = (
+                            h0_source[(pool_idx64, i_hv, v_global_write, k_write)] = (
                                 sData[(k_write, v_write, stage)]
                             )
 
@@ -311,6 +315,10 @@ def _define_kernels():
         pool_idx = h0_indices[i_n]
 
         if pool_idx >= 0:
+            # State indexing in int64: envelope-strided pools (unified memory /
+            # page-major) put pool_idx * stride(0) far past 2^31 (the CuTe twin
+            # of the fla/chunk_delta_h.py stride_init_state fix).
+            pool_idx64 = cutlass.Int64(pool_idx)
             k_local = in_warp_tid // V_PER_WARP_SMALL
             v_local = in_warp_tid % V_PER_WARP_SMALL
             v_base = warp_idx * V_PER_WARP_SMALL
@@ -423,7 +431,7 @@ def _define_kernels():
                         h_val = 0.0
                         if v_global_load < v.shape[3]:
                             h_val = cutlass.Float32(
-                                h0_source[(pool_idx, i_hv, v_global_load, k_load)]
+                                h0_source[(pool_idx64, i_hv, v_global_load, k_load)]
                             )
                         sData[(k_load, v_load, stage)] = h_val
 
@@ -480,7 +488,7 @@ def _define_kernels():
                     if k_write < TILE_K:
                         v_global_write = v_tile * TILE_V_SMALL + v_write
                         if v_global_write < v.shape[3]:
-                            h0_source[(pool_idx, i_hv, v_global_write, k_write)] = (
+                            h0_source[(pool_idx64, i_hv, v_global_write, k_write)] = (
                                 sData[(k_write, v_write, stage)]
                             )
 
@@ -523,6 +531,10 @@ def _define_kernels():
         pool_idx = h0_indices[i_n]
 
         if pool_idx >= 0:
+            # State indexing in int64: envelope-strided pools (unified memory /
+            # page-major) put pool_idx * stride(0) far past 2^31 (the CuTe twin
+            # of the fla/chunk_delta_h.py stride_init_state fix).
+            pool_idx64 = cutlass.Int64(pool_idx)
             k_local = in_warp_tid // V_PER_WARP
             v_local = in_warp_tid % V_PER_WARP
             v_base = warp_idx * V_PER_WARP
@@ -634,7 +646,7 @@ def _define_kernels():
                         h_val = 0.0
                         if v_global_load < v.shape[3]:
                             h_val = cutlass.Float32(
-                                h0_source[(pool_idx, i_hv, v_global_load, k_load)]
+                                h0_source[(pool_idx64, i_hv, v_global_load, k_load)]
                             )
                         sData[(k_load, v_load, stage)] = h_val
 
@@ -691,7 +703,7 @@ def _define_kernels():
                     if k_write < TILE_K:
                         v_global_write = v_tile * TILE_V + v_write
                         if v_global_write < v.shape[3]:
-                            h0_source[(pool_idx, i_hv, v_global_write, k_write)] = (
+                            h0_source[(pool_idx64, i_hv, v_global_write, k_write)] = (
                                 sData[(k_write, v_write, stage)]
                             )
 
@@ -734,6 +746,10 @@ def _define_kernels():
         pool_idx = h0_indices[i_n]
 
         if pool_idx >= 0:
+            # State indexing in int64: envelope-strided pools (unified memory /
+            # page-major) put pool_idx * stride(0) far past 2^31 (the CuTe twin
+            # of the fla/chunk_delta_h.py stride_init_state fix).
+            pool_idx64 = cutlass.Int64(pool_idx)
             k_local = in_warp_tid // V_PER_WARP
             v_local = in_warp_tid % V_PER_WARP
             v_base = warp_idx * V_PER_WARP
@@ -845,7 +861,7 @@ def _define_kernels():
                         h_val = 0.0
                         if v_global_load < v.shape[3]:
                             h_val = cutlass.Float32(
-                                h0_source[(pool_idx, i_hv, v_global_load, k_load)]
+                                h0_source[(pool_idx64, i_hv, v_global_load, k_load)]
                             )
                         sData[(k_load, v_load, stage)] = h_val
 
@@ -902,7 +918,7 @@ def _define_kernels():
                     if k_write < TILE_K:
                         v_global_write = v_tile * TILE_V + v_write
                         if v_global_write < v.shape[3]:
-                            h0_source[(pool_idx, i_hv, v_global_write, k_write)] = (
+                            h0_source[(pool_idx64, i_hv, v_global_write, k_write)] = (
                                 sData[(k_write, v_write, stage)]
                             )
 
@@ -1223,11 +1239,28 @@ def _get_jit_functions():
     return _jit_functions
 
 
-def _get_compiled_kernel(N, H, HV, K, V, pool_size, use_small_batch, is_varlen_decode):
-    """Get or compile the KDA kernel for given dimensions."""
+def _get_compiled_kernel(N, H, HV, K, V, h0_source, use_small_batch, is_varlen_decode):
+    """Get or compile the KDA kernel for given dimensions.
+
+    ``h0_source`` is the caller's real state pool: ``from_dlpack`` bakes its
+    exact layout into the compiled kernel, so envelope-strided pools (unified
+    memory / page-major, slot stride(0) != HV*V*K) compile against their true
+    slot pitch. The cache key carries the strides alongside the shape.
+    """
     global _compiled_kernels
 
-    key = (N, H, HV, K, V, pool_size, use_small_batch, is_varlen_decode)
+    pool_size = h0_source.shape[0]
+    key = (
+        N,
+        H,
+        HV,
+        K,
+        V,
+        pool_size,
+        tuple(h0_source.stride()),
+        use_small_batch,
+        is_varlen_decode,
+    )
     if key in _compiled_kernels:
         return _compiled_kernels[key]
 
@@ -1250,7 +1283,6 @@ def _get_compiled_kernel(N, H, HV, K, V, pool_size, use_small_batch, is_varlen_d
 
     A_log = torch.zeros(HV, dtype=torch.float32, device="cuda")
     dt_bias = torch.zeros(HV, K, dtype=torch.bfloat16, device="cuda")
-    h0_source = torch.zeros(pool_size, HV, V, K, dtype=torch.float32, device="cuda")
     h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda")
 
     cu_seqlens_tensor = from_dlpack(cu_seqlens, assumed_align=16)
@@ -1261,7 +1293,7 @@ def _get_compiled_kernel(N, H, HV, K, V, pool_size, use_small_batch, is_varlen_d
     b_tensor = from_dlpack(b, assumed_align=16)
     A_log_tensor = from_dlpack(A_log, assumed_align=16)
     dt_bias_tensor = from_dlpack(dt_bias, assumed_align=16)
-    h0_source_tensor = from_dlpack(h0_source, assumed_align=16)
+    h0_source_tensor = from_dlpack(h0_source.detach(), assumed_align=16)
     h0_indices_tensor = from_dlpack(h0_indices, assumed_align=16)
     o_tensor = from_dlpack(o, assumed_align=16)
 
@@ -1304,6 +1336,7 @@ def _get_compiled_kernel(N, H, HV, K, V, pool_size, use_small_batch, is_varlen_d
     logger.info(
         "CuTe DSL KDA kernel compiled: "
         f"N={N}, H={H}, HV={HV}, K={K}, V={V}, pool_size={pool_size}, "
+        f"pool_strides={tuple(h0_source.stride())}, "
         f"small_batch={use_small_batch}, varlen={is_varlen_decode}"
     )
     return compiled_kernel
@@ -1369,6 +1402,9 @@ def cutedsl_fused_sigmoid_gating_kda_update(
 
     State layout contract:
         initial_state_source.shape == (pool_size, HV, V, K)
+        The slot dim may be envelope-strided (stride(0) > HV*V*K under
+        unified-memory / page-major pools); each slot's [HV, V, K] block must
+        be compact. State updates are written back in place through the view.
 
     Dense decode:
         q/k: (N, 1, H, K)
@@ -1453,7 +1489,16 @@ def cutedsl_fused_sigmoid_gating_kda_update(
     A_log = _normalize_A_log(A_log, HV)
     dt_bias = _normalize_dt_bias(dt_bias, HV, K)
 
-    h0_source = h0_source.contiguous()
+    # h0_source may be an envelope-strided pool view (unified memory /
+    # page-major): slot stride(0) is the per-slot envelope pitch, not HV*V*K.
+    # Never .contiguous() it — on a strided view that copies, so the kernel's
+    # in-place state update would land in a dropped temporary. The kernel only
+    # needs each slot's [HV, V, K] block itself to be compact.
+    assert h0_source.stride()[1:] == (V * K, K, 1), (
+        "CuTe DSL KDA decode requires a compact per-slot [HV, V, K] state "
+        f"block; got strides {tuple(h0_source.stride())} for shape "
+        f"{tuple(h0_source.shape)}"
+    )
 
     initial_state_indices = initial_state_indices.contiguous()
     if cu_seqlens is not None:
@@ -1496,7 +1541,7 @@ def cutedsl_fused_sigmoid_gating_kda_update(
     stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
 
     compiled_kernel = _get_compiled_kernel(
-        N, H, HV, K, V, pool_size, use_small_batch, is_varlen_decode
+        N, H, HV, K, V, h0_source, use_small_batch, is_varlen_decode
     )
 
     compiled_kernel(

python/sglang/kernels/ops/attention/fla/fused_recurrent.py

@@ -414,6 +414,7 @@ def fused_recurrent_kda_packed_decode_kernel(
     ht,
     ssm_state_indices,
     scale,
+    lower_bound,
     stride_mixed_qkv_tok: tl.constexpr,
     stride_a_tok: tl.constexpr,
     stride_b_tok: tl.constexpr,
@@ -428,6 +429,7 @@ def fused_recurrent_kda_packed_decode_kernel(
     BV: tl.constexpr,
     SOFTPLUS_THRESHOLD: tl.constexpr,
     USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
+    USE_LOWER_BOUND: tl.constexpr,
 ):
     """KDA packed decode: same shape as the GDN packed decode kernel, but
     with a per-K gate (``a`` is ``[B, HV*K]`` and ``dt_bias`` is ``[HV*K]``),
@@ -475,8 +477,14 @@ def fused_recurrent_kda_packed_decode_kernel(
     A_log_val = tl.load(A_log + i_hv).to(tl.float32)
 
     x = b_a + b_dt
-    softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x)
-    b_g = -tl.exp(A_log_val) * softplus_x  # [BK]
+    if USE_LOWER_BOUND:
+        # KDA safe gate: lower_bound * sigmoid(exp(A_log) * (g + bias)),
+        # matching the chunked prefill kernel (kda.py) and FLA reference.
+        b_g = lower_bound * tl.sigmoid(tl.exp(A_log_val) * x)  # [BK]
+    else:
+        # Standard gate: -exp(A_log) * softplus(g + bias)
+        softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x)
+        b_g = -tl.exp(A_log_val) * softplus_x  # [BK]
 
     b_val = tl.load(b + i_n * stride_b_tok + i_hv).to(tl.float32)
     # Keep beta in fp32 (no bf16 round-trip) to match the generic decode
@@ -508,6 +516,7 @@ def fused_recurrent_kda_packed_decode(
     out: torch.Tensor,
     ssm_state_indices: torch.Tensor,
     use_qk_l2norm_in_kernel: bool = False,
+    lower_bound: Optional[float] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """KDA T=1 decode fast path. Mirrors ``fused_recurrent_gated_delta_rule_packed_decode``
     but the gate ``g`` is a per-K vector instead of a scalar.
@@ -612,6 +621,40 @@ def fused_recurrent_kda_packed_decode(
             f"Invalid head config inferred from mixed_qkv: H={H}, HV={HV}."
         )
 
+    # Batched-decode CUDA fast path:
+    # row-streaming state update reaches the in-place R+W bandwidth of the
+    # part (~9.6 TB/s) where this triton kernel tops out at ~5 TB/s holding a
+    # [BV, K] register tile per warp. ULP-level output differences only
+    # (reduction order); small batches keep triton (launch-bound anyway).
+    if use_qk_l2norm_in_kernel:
+        from sglang.kernels.ops.attention import kda_packed_decode as kda_decode_cuda
+
+        if kda_decode_cuda.covered(
+            mixed_qkv,
+            a,
+            b,
+            A_log,
+            dt_bias,
+            initial_state,
+            out,
+            ssm_state_indices,
+            H,
+        ):
+            kda_decode_cuda.kda_packed_decode(
+                mixed_qkv,
+                a,
+                b,
+                A_log,
+                dt_bias,
+                scale,
+                initial_state,
+                out,
+                ssm_state_indices,
+                H,
+                lower_bound,
+            )
+            return out, initial_state
+
     BK = triton.next_power_of_2(K)
     if triton.cdiv(K, BK) != 1:
         raise ValueError(
@@ -641,6 +684,7 @@ def fused_recurrent_kda_packed_decode(
         ht=initial_state,
         ssm_state_indices=ssm_state_indices,
         scale=scale,
+        lower_bound=lower_bound if lower_bound is not None else 0.0,
         stride_mixed_qkv_tok=stride_mixed_qkv_tok,
         stride_a_tok=stride_a_tok,
         stride_b_tok=stride_b_tok,
@@ -655,6 +699,7 @@ def fused_recurrent_kda_packed_decode(
         BV=BV,
         SOFTPLUS_THRESHOLD=20.0,
         USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
+        USE_LOWER_BOUND=lower_bound is not None,
         num_warps=num_warps,
         num_stages=num_stages,
     )
@@ -830,6 +875,7 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
     o,
     h0_source,
     h0_indices,
+    stride_h0_source,
     cu_seqlens,
     scale,
     intermediate_states_buffer,
@@ -897,22 +943,28 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
 
     b_h = tl.zeros([BV, BK], dtype=tl.float32)
     if USE_INITIAL_STATE:
-        idx = tl.load(h0_indices + i_n)
+        # Slot stride comes from the caller (h0_source.stride(0)): the state pool
+        # may be an envelope-strided view (page-major / unified memory), where the
+        # per-slot pitch spans ALL layers' state, not HV*K*V. int64: envelope
+        # pitches overflow an int32 index product.
+        idx = tl.load(h0_indices + i_n).to(tl.int64)
         # Add bounds checking for idx
         if idx >= 0:  # Assuming negative indices are invalid
             p_h0 = (
                 h0_source
-                + idx * HV * K * V
+                + idx * stride_h0_source
                 + i_hv * K * V
                 + o_v[:, None] * K
                 + o_k[None, :]
             )
             b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
 
-    # Prepare intermediate state cache variables if enabled
+    # Prepare intermediate state cache variables if enabled. int64: the buffer
+    # is contiguous but `cache_idx * cache_steps * HV * K * V` can exceed int32
+    # for large slot counts.
     cache_idx = -1
     if CACHE_INTERMEDIATE_STATES:
-        cache_idx = tl.load(intermediate_state_indices + i_n)
+        cache_idx = tl.load(intermediate_state_indices + i_n).to(tl.int64)
 
     step_idx = 0
     for _ in range(0, T):
@@ -987,11 +1039,11 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
     # Store final state back to h0_source with bounds checking
     # ssm states
     if not DISABLE_STATE_UPDATE:
-        idx = tl.load(h0_indices + i_n)
+        idx = tl.load(h0_indices + i_n).to(tl.int64)
         if idx >= 0:  # Add bounds checking
             p_h0 = (
                 h0_source
-                + idx * HV * K * V
+                + idx * stride_h0_source
                 + i_hv * K * V
                 + o_v[:, None] * K
                 + o_k[None, :]
@@ -1053,6 +1105,11 @@ def fused_recurrent_gated_delta_rule_update_fwd(
         o=o,
         h0_source=initial_state_source,
         h0_indices=initial_state_indices,
+        # Envelope-strided state pools (page-major / unified memory) have a
+        # per-slot pitch != HV*K*V; contiguous pools pass exactly HV*K*V.
+        stride_h0_source=(
+            initial_state_source.stride(0) if initial_state_source is not None else 0
+        ),
         cu_seqlens=cu_seqlens,
         scale=scale,
         intermediate_states_buffer=intermediate_states_buffer,

python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py

@@ -12,6 +12,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
     dt_bias,
     softplus_beta,
     softplus_threshold,
+    lower_bound,
     q,
     k,
     v,
@@ -19,6 +20,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
     o,
     h0_source,
     h0_indices,
+    stride_h0_source,
     cu_seqlens,
     # Parameters for target_verify support (unused for decode)
     intermediate_states_buffer,
@@ -47,10 +49,14 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
     USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
     IS_VARLEN: tl.constexpr,
     IS_KDA: tl.constexpr,
+    USE_LOWER_BOUND: tl.constexpr,
     # Optional flags for target_verify support (default False for decode)
     DISABLE_STATE_UPDATE: tl.constexpr = False,
     CACHE_INTERMEDIATE_STATES: tl.constexpr = False,
     HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: tl.constexpr = False,
+    # ReplaySSM fused ring-write. Pointers stay None and CACHE_RING False for
+    # decode / flag-off -> byte-identical. The gate ring layout follows IS_KDA
+    # (see the store below).
     replayssm_rawv=None,
     replayssm_rawk=None,
     replayssm_g=None,
@@ -104,11 +110,15 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
 
     b_h = tl.zeros([BK, BV], dtype=tl.float32)
     if USE_INITIAL_STATE:
-        idx = tl.load(h0_indices + i_n)
+        # Slot stride comes from the caller (h0_source.stride(0)): the state pool
+        # may be an envelope-strided view (page-major / unified memory), where the
+        # per-slot pitch spans ALL layers' state, not HV*K*V. int64: envelope
+        # pitches overflow an int32 index product.
+        idx = tl.load(h0_indices + i_n).to(tl.int64)
         if idx >= 0:
             p_h0 = (
                 h0_source
-                + idx * HV * K * V
+                + idx * stride_h0_source
                 + i_hv * K * V
                 + o_v[None, :] * K
                 + o_k[:, None]
@@ -128,10 +138,12 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
             retrieve_parent_token_base, mask=mask_retrieve, other=0
         )
 
-    # Prepare intermediate state cache index if enabled
+    # Prepare intermediate state cache index if enabled. int64: the buffer is
+    # contiguous but `cache_idx * cache_steps * HV * K * V` can exceed int32 for
+    # large slot counts.
     cache_idx = -1
     if CACHE_INTERMEDIATE_STATES:
-        cache_idx = tl.load(intermediate_state_indices + i_n)
+        cache_idx = tl.load(intermediate_state_indices + i_n).to(tl.int64)
 
     step_idx = 0
     for _ in range(0, T):
@@ -169,23 +181,32 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
             b_a = tl.load(p_a).to(tl.float32)
             b_dt_bias = tl.load(p_dt_bias).to(tl.float32)
 
-        # Compute g = -exp(A_log) * softplus(a + dt_bias)
         x = b_a + b_dt_bias
-        beta_x = softplus_beta * x
-        # Apply softplus with numerical stability
-        softplus_x = tl.where(
-            beta_x <= softplus_threshold,
-            (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)),
-            x,
-        )
-        b_g = -tl.exp(b_A_log) * softplus_x
+        if USE_LOWER_BOUND:
+            # KDA safe gate: lower_bound * sigmoid(exp(A_log) * (a + dt_bias))
+            b_g = lower_bound * tl.sigmoid(tl.exp(b_A_log) * x)
+        else:
+            # Compute g = -exp(A_log) * softplus(a + dt_bias)
+            beta_x = softplus_beta * x
+            # Apply softplus with numerical stability
+            softplus_x = tl.where(
+                beta_x <= softplus_threshold,
+                (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)),
+                x,
+            )
+            b_g = -tl.exp(b_A_log) * softplus_x
 
         # Compute beta = sigmoid(b)
         b_beta = 1.0 / (1.0 + tl.exp(-b_b))
 
-        # Stored here, pre-l2norm k / pre-delta v, so the commit fold's replay
-        # is bit-identical to the update below; steps >= MAX_CACHE_LEN would
-        # smash the next slot's ring.
+        # fused ring-write: stash this step's raw inputs + in-kernel gate/beta
+        # into the per-slot ring for the commit fold to replay. Must sit here --
+        # b_k is still pre-l2norm, b_v still pre-delta, b_g/b_beta are formed,
+        # so the fold's replay is bit-identical to the update below. rawk uses
+        # the k-head i_h (shared across a GQA group); rawv/g/beta use the v-head
+        # i_hv. step_idx < MAX_CACHE_LEN: absorb-inflated rows can exceed the
+        # ring; the overflow steps are past the committable prefix, so drop them
+        # (writing them would smash the next slot's ring).
         if CACHE_RING:
             ring_slot = tl.load(h0_indices + i_n).to(tl.int64)
             if ring_slot >= 0 and step_idx < MAX_CACHE_LEN:
@@ -208,13 +229,30 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
                         b_k.to(replayssm_rawk.dtype.element_ty),
                         mask=mask_k,
                     )
-                    tl.store(
-                        replayssm_g
-                        + ring_slot * stride_g_slot
-                        + i_hv * MAX_CACHE_LEN
-                        + step_idx,
-                        b_g,
-                    )
+                    # b_g follows IS_KDA: KDA loads a/dt_bias with mask_k, so the
+                    # gate is a per-K vector and the ring row is K wide; GDN's is
+                    # a scalar per (head, step). The two layouts are not
+                    # interchangeable -- storing one into the other's stride is a
+                    # shape error, not a slow path -- and memory_pool.py sizes
+                    # replayssm_g off the same is_kda test.
+                    if IS_KDA:
+                        tl.store(
+                            replayssm_g
+                            + ring_slot * stride_g_slot
+                            + i_hv * MAX_CACHE_LEN * K
+                            + step_idx * K
+                            + o_k,
+                            b_g,
+                            mask=mask_k,
+                        )
+                    else:
+                        tl.store(
+                            replayssm_g
+                            + ring_slot * stride_g_slot
+                            + i_hv * MAX_CACHE_LEN
+                            + step_idx,
+                            b_g,
+                        )
                     if i_k == 0:
                         tl.store(
                             replayssm_beta
@@ -277,11 +315,11 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
     # Store final state back to h0_source with bounds checking
     if not DISABLE_STATE_UPDATE:
         if USE_INITIAL_STATE:
-            idx = tl.load(h0_indices + i_n)
+            idx = tl.load(h0_indices + i_n).to(tl.int64)
             if idx >= 0:
                 p_h0 = (
                     h0_source
-                    + idx * HV * K * V
+                    + idx * stride_h0_source
                     + i_hv * K * V
                     + o_v[None, :] * K
                     + o_k[:, None]
@@ -305,6 +343,7 @@ def fused_sigmoid_gating_delta_rule_update(
     use_qk_l2norm_in_kernel: bool = False,
     cu_seqlens: Optional[torch.Tensor] = None,
     is_kda: bool = False,
+    lower_bound: Optional[float] = None,
     # Optional parameters for target_verify support
     disable_state_update: bool = False,
     intermediate_states_buffer: Optional[torch.Tensor] = None,
@@ -313,6 +352,9 @@ def fused_sigmoid_gating_delta_rule_update(
         int
     ] = None,  # kept for API compat; stride is derived from ``intermediate_states_buffer.shape[1]``
     retrieve_parent_token: Optional[torch.Tensor] = None,
+    # fused ReplaySSM ring-write (spec verify). When cache_ring, each draft step
+    # stores pre-norm k / raw v / gate / beta into these per-slot rings,
+    # replacing the eager ring-write. Off by default -> decode unchanged.
     cache_ring: bool = False,
     replayssm_rawv: Optional[torch.Tensor] = None,
     replayssm_rawk: Optional[torch.Tensor] = None,
@@ -375,14 +417,18 @@ def fused_sigmoid_gating_delta_rule_update(
         else 0
     )
 
+    # ring strides (per-slot rings are contiguous [num_slots, heads, L, dim];
+    # the kernel offsets within a slot with MAX_CACHE_LEN and the dim extents).
     if cache_ring:
-        assert not is_kda, "cache_ring supports GDN only (scalar gate layout)"
         # stride(0) is used as the slot pitch, so a tensor still carrying the
-        # layer dim would scribble outside its slot.
+        # layer dim would scribble outside its slot. The gate ring is the one
+        # whose rank depends on the model: per-K vector for KDA, per-head scalar
+        # for GDN, matching g_shape in memory_pool.py and the IS_KDA branch in
+        # the store above.
         assert (
             replayssm_rawv.dim() == 4
             and replayssm_rawk.dim() == 4
-            and replayssm_g.dim() == 3
+            and replayssm_g.dim() == (4 if is_kda else 3)
             and replayssm_beta.dim() == 3
         ), "cache_ring expects per-layer ring views"
         max_cache_len = replayssm_rawv.shape[-2]
@@ -400,6 +446,7 @@ def fused_sigmoid_gating_delta_rule_update(
         dt_bias=dt_bias,
         softplus_beta=softplus_beta,
         softplus_threshold=softplus_threshold,
+        lower_bound=lower_bound if lower_bound is not None else 0.0,
         q=q,
         k=k,
         v=v,
@@ -407,6 +454,11 @@ def fused_sigmoid_gating_delta_rule_update(
         o=o,
         h0_source=initial_state_source,
         h0_indices=initial_state_indices,
+        # Envelope-strided state pools (page-major / unified memory) have a
+        # per-slot pitch != HV*K*V; contiguous pools pass exactly HV*K*V.
+        stride_h0_source=(
+            initial_state_source.stride(0) if initial_state_source is not None else 0
+        ),
         cu_seqlens=cu_seqlens,
         intermediate_states_buffer=intermediate_states_buffer,
         intermediate_state_indices=intermediate_state_indices,
@@ -433,6 +485,7 @@ def fused_sigmoid_gating_delta_rule_update(
         USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
         IS_VARLEN=cu_seqlens is not None,
         IS_KDA=is_kda,
+        USE_LOWER_BOUND=lower_bound is not None,
         DISABLE_STATE_UPDATE=disable_state_update,
         CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None,
         HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=retrieve_parent_token is not None,

python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py

@@ -38,7 +38,8 @@ Closed-loop exact fold (the state / output error split):
   ``"tf32"`` (~5e-4, tensor-core path; worst case through the (I+A)^{-1}
   amplification 2^(BS-1) still lands at the floor). ``"ieee"`` / ``"tf32x3"``
   remain selectable for ablations. The committed state is untouched by any of
-  these dots (requires the fp32 SSM checkpoint; enforced in server_args).
+  these dots (fp32 SSM checkpoint is the server_args default; a 16-bit
+  checkpoint is allowed with a warning, unvalidated for GDN).
 
 Differences from the vLLM reference:
   * SGLang passes **split** ``q`` / ``k`` / ``v`` tensors (already split + post

python/sglang/kernels/ops/attention/fla/kda.py

@@ -25,8 +25,11 @@ from sglang.kernels.ops.attention.fla.index import (
 from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
 from sglang.kernels.ops.attention.fla.op import exp, exp2, log
 from sglang.kernels.ops.attention.fla.utils import (
+    autotune_cache_kwargs,
     check_shared_mem,
     is_intel,
+    is_nvidia,
+    is_tf32_supported,
 )
 
 if is_intel:
@@ -514,18 +517,8 @@ def chunk_kda_scaled_dot_kkt_fwd(
     return A, Aqk
 
 
-@triton.autotune(
-    configs=[
-        triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
-        for BK in [64, 128]
-        for BV in [64, 128]
-        for num_warps in [2, 4, 8]
-        for num_stages in [2, 3, 4]
-    ],
-    key=["H", "K", "V", "BT", "IS_VARLEN"],
-)
 @triton.jit(do_not_specialize=["T"])
-def recompute_w_u_fwd_kernel(
+def _recompute_w_u_fwd_kernel(
     k,
     kg,
     v,
@@ -645,6 +638,67 @@ def recompute_w_u_fwd_kernel(
         tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
 
 
+_RECOMPUTE_W_U_CONFIGS = [
+    triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
+    for BK in [64, 128]
+    for BV in [64, 128]
+    for num_warps in [2, 4, 8]
+    for num_stages in [2, 3, 4]
+]
+
+recompute_w_u_fwd_kernel = triton.autotune(
+    configs=_RECOMPUTE_W_U_CONFIGS,
+    key=["H", "K", "V", "BT", "IS_VARLEN"],
+    **autotune_cache_kwargs,
+)(_recompute_w_u_fwd_kernel)
+
+_K3_RECOMPUTE_W_U_CONFIGS = {
+    (9, 0): {"BK": 128, "BV": 128, "num_warps": 8, "num_stages": 2},
+    (10, 3): {"BK": 64, "BV": 128, "num_warps": 8, "num_stages": 2},
+}
+
+
+@torch.inference_mode()
+def precompile_k3_recompute_w_u_kernel(
+    *, num_heads: int, dtype: torch.dtype, device: torch.device
+) -> bool:
+    device = torch.device(device)
+    if (
+        not is_nvidia
+        or device.type != "cuda"
+        or torch.cuda.get_device_capability(device) not in _K3_RECOMPUTE_W_U_CONFIGS
+    ):
+        return False
+
+    shape = (1, 1, num_heads, 128)
+    k = torch.zeros(shape, dtype=dtype, device=device)
+    v = torch.zeros_like(k)
+    beta = torch.zeros((1, 1, num_heads), dtype=dtype, device=device)
+    A = torch.zeros((1, 1, num_heads, 64), dtype=dtype, device=device)
+    gk = torch.zeros(shape, dtype=torch.float32, device=device)
+    cu_seqlens = torch.tensor([0, 1], dtype=torch.int64, device=device)
+    recompute_w_u_fwd(k, v, beta, A, gk=gk, cu_seqlens=cu_seqlens)
+    return True
+
+
+def _get_k3_recompute_w_u_config(
+    k: torch.Tensor,
+    gk: torch.Tensor | None,
+    cu_seqlens: torch.LongTensor | None,
+    K: int,
+    V: int,
+    BT: int,
+) -> dict | None:
+    if (
+        not is_nvidia
+        or gk is None
+        or cu_seqlens is None
+        or (K, V, BT) != (128, 128, 64)
+    ):
+        return None
+    return _K3_RECOMPUTE_W_U_CONFIGS.get(torch.cuda.get_device_capability(k.device))
+
+
 def recompute_w_u_fwd(
     k: torch.Tensor,
     v: torch.Tensor,
@@ -664,7 +718,13 @@ def recompute_w_u_fwd(
     w = torch.empty_like(k)
     u = torch.empty_like(v)
     kg = torch.empty_like(k) if gk is not None else None
-    recompute_w_u_fwd_kernel[(NT, B * H)](
+    static_config = _get_k3_recompute_w_u_config(k, gk, cu_seqlens, K, V, BT)
+    kernel = (
+        _recompute_w_u_fwd_kernel
+        if static_config is not None
+        else recompute_w_u_fwd_kernel
+    )
+    kernel[(NT, B * H)](
         k=k,
         kg=kg,
         v=v,
@@ -682,7 +742,8 @@ def recompute_w_u_fwd(
         BT=BT,
         STORE_KG=kg is not None,
         IS_VARLEN=cu_seqlens is not None,
-        DOT_PRECISION="tf32",
+        DOT_PRECISION="tf32" if is_tf32_supported else "ieee",
+        **(static_config or {}),
     )
     return w, u, kg
 
@@ -1126,6 +1187,9 @@ def chunk_kda_fwd(
     del Aqk, v_new
 
     if output_intermediate_states:
+        # h holds the recurrent state at every chunk-size boundary
+        # ([1, NT, H, V, K] packed across cu_seqlens) — the mamba radix
+        # track path snapshots per-chunk states from it during extend.
         return o, h
     del h
     return o

python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py

@@ -0,0 +1,386 @@
+# SPDX-License-Identifier: Apache-2.0
+"""ReplaySSM speculative-decode state commit for KDA (Kimi Delta Attention).
+
+KDA keeps its own recurrent verify kernel for the per-step OUTPUT (unchanged),
+so — unlike the GDN spec kernel (gdn_replayssm_spec_decode.py) — this module does
+NOT reconstruct the verify output. It only replaces the per-draft-token full-SSM
+snapshot cache (``intermediate_ssm``, ``max_running×(γ+1)×[HV,V,K]``, the memory
+hog that collapses dspark concurrency) with a small per-request input window +
+an exact-fold on commit.
+
+Scheme (fold-every-commit, no circular ring / periodic flush):
+  * during verify, the KDA backend stores the draft window's raw inputs
+    (raw v, raw pre-norm k, per-K log-decay gate ``gk`` (fp32), beta (fp32)) into
+    the per-slot ring at positions ``0..spec_len``;
+  * on commit, :func:`commit_kda_replayssm_spec` replays the *accepted* prefix
+    ``0..accept_len`` from the persistent checkpoint ``h0`` into ``h0`` in place —
+    ``h0`` is always the current committed state, so the next verify reads it
+    directly as its initial state (no lag, no chunked reconstruction).
+
+The exact-fold is a BITWISE CLONE of ``fused_recurrent_gated_delta_rule_fwd_kernel``'s
+``IS_KDA`` branch (fused_recurrent.py): same [BK, BV] fp32 tile (K rows, V cols),
+same division-form L2 norm (eps inside sqrt), same per-K gate decay
+``h *= exp(gk)``, same decay→delta→rank-1 op order. Given identical inputs the
+folded checkpoint is bit-identical to the recurrent baseline's committed state
+(the delta-rule recurrence is contractive, so no length-dependent error). Do NOT
+reorder into tl.dot / reciprocal-multiply; keep num_warps=1 so the reduction
+trees match.
+
+Linear chain only (dspark γ chain, topk<=1).
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def kda_replayssm_exact_fold_kernel(
+    h0,  # [num_slots, HV, V, K] fp32 checkpoint (folded in place)
+    rawv_cache,  # [num_slots, HV, L, V]  raw v
+    rawk_cache,  # [num_slots, H,  L, K]  raw pre-norm k
+    gk_cache,  # [num_slots, HV, L, K]  fp32 per-K log-decay gate
+    beta_cache,  # [num_slots, HV, L]     fp32 beta
+    ssm_state_indices,  # [B] int  physical slot per request
+    accept_lens,  # [B] int  committed prefix length per request
+    mamba_track_indices,  # [B] int  extra_buffer track slot (or NULL) per request
+    mamba_steps_to_track,  # [B] int  crossing step (or -1) per request
+    stride_state_slot: tl.constexpr,
+    stride_rawv_slot: tl.constexpr,
+    stride_rawk_slot: tl.constexpr,
+    stride_gk_slot: tl.constexpr,
+    stride_beta_slot: tl.constexpr,
+    stride_state_layer: tl.constexpr,
+    stride_rawv_layer: tl.constexpr,
+    stride_rawk_layer: tl.constexpr,
+    stride_gk_layer: tl.constexpr,
+    stride_beta_layer: tl.constexpr,
+    stride_indices: tl.constexpr,
+    stride_accept: tl.constexpr,
+    stride_track: tl.constexpr,
+    stride_steps: tl.constexpr,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    MAX_CACHE_LEN: tl.constexpr,
+    USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
+    NULL_BLOCK_ID: tl.constexpr,
+    HAS_TRACK: tl.constexpr,
+):
+    i_v = tl.program_id(0)
+    i_n = tl.program_id(1)
+    # program_id(2) packs (layer, v-head): layer-major so a single launch folds
+    # all KDA layers. num_layers=1 launches (per-layer entry) keep i_layer == 0.
+    i_hvl = tl.program_id(2)
+    # int64: layer stride * i_layer overflows int32 at K3 scale (69 layers x
+    # ~34M-element per-layer stride > 2^31).
+    i_layer = (i_hvl // HV).to(tl.int64)
+    i_hv = i_hvl % HV
+    i_h = i_hv // (HV // H)
+    # Shift the layer-indexed bases once; every pointer below is layer-relative.
+    h0 = h0 + i_layer * stride_state_layer
+    rawv_cache = rawv_cache + i_layer * stride_rawv_layer
+    rawk_cache = rawk_cache + i_layer * stride_rawk_layer
+    gk_cache = gk_cache + i_layer * stride_gk_layer
+    beta_cache = beta_cache + i_layer * stride_beta_layer
+
+    state_idx = tl.load(ssm_state_indices + i_n * stride_indices).to(tl.int64)
+    if state_idx <= NULL_BLOCK_ID:
+        return
+    n_commit = tl.load(accept_lens + i_n * stride_accept).to(tl.int32)
+    if n_commit <= 0:
+        return
+
+    # extra_buffer: snapshot the interval-crossing state into the track ping-pong
+    # slot on the step it crosses (mask-gated: -1 track step / NULL slot = skip).
+    if HAS_TRACK:
+        track_idx = tl.load(mamba_track_indices + i_n * stride_track).to(tl.int64)
+        track_step = tl.load(mamba_steps_to_track + i_n * stride_steps).to(tl.int32)
+    else:
+        track_idx = NULL_BLOCK_ID
+        track_step = -1
+
+    o_k = tl.arange(0, BK)
+    o_v = i_v * BV + tl.arange(0, BV)
+    mask_k = o_k < K
+    mask_v = o_v < V
+    mask_h = mask_k[:, None] & mask_v[None, :]
+
+    # [BK, BV] tile: K rows / V cols, matching the recurrent baseline's memory
+    # offset (v * K + k, K contiguous).
+    p_h0 = (
+        h0
+        + state_idx * stride_state_slot
+        + i_hv * V * K
+        + o_v[None, :] * K
+        + o_k[:, None]
+    )
+    b_h = tl.load(p_h0, mask=mask_h, other=0.0).to(tl.float32)
+
+    for t in range(0, n_commit):
+        phys = t.to(tl.int64)
+        b_k = tl.load(
+            rawk_cache
+            + state_idx * stride_rawk_slot
+            + (i_h * MAX_CACHE_LEN + phys) * K
+            + o_k,
+            mask=mask_k,
+            other=0.0,
+        ).to(tl.float32)
+        b_v = tl.load(
+            rawv_cache
+            + state_idx * stride_rawv_slot
+            + (i_hv * MAX_CACHE_LEN + phys) * V
+            + o_v,
+            mask=mask_v,
+            other=0.0,
+        ).to(tl.float32)
+        b_gk = tl.load(
+            gk_cache
+            + state_idx * stride_gk_slot
+            + (i_hv * MAX_CACHE_LEN + phys) * K
+            + o_k,
+            mask=mask_k,
+            other=0.0,
+        ).to(tl.float32)
+        b_beta = tl.load(
+            beta_cache + state_idx * stride_beta_slot + i_hv * MAX_CACHE_LEN + phys
+        ).to(tl.float32)
+
+        # --- verbatim recurrent update, IS_KDA branch (see module docstring) ---
+        if USE_QK_L2NORM_IN_KERNEL:
+            b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6))
+        b_h *= tl.exp(b_gk[:, None])  # per-K gate decay, broadcast over V
+        b_v -= tl.sum(b_h * b_k[:, None], 0)
+        b_v *= b_beta
+        b_h += b_k[:, None] * b_v[None, :]
+
+        # Interval-crossing snapshot -> track slot (state AFTER step `track_step`).
+        if HAS_TRACK:
+            if (t == track_step) and (track_idx > NULL_BLOCK_ID):
+                tl.store(
+                    h0
+                    + track_idx * stride_state_slot
+                    + i_hv * V * K
+                    + o_v[None, :] * K
+                    + o_k[:, None],
+                    b_h.to(h0.dtype.element_ty),
+                    mask=mask_h,
+                )
+
+    tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h)
+
+
+def commit_kda_replayssm_spec(
+    checkpoint_state: torch.Tensor,  # [num_slots, HV, V, K] fp32, folded in place
+    rawv_cache: torch.Tensor,  # [num_slots, HV, L, V]
+    rawk_cache: torch.Tensor,  # [num_slots, H,  L, K]
+    gk_cache: torch.Tensor,  # [num_slots, HV, L, K] fp32
+    beta_cache: torch.Tensor,  # [num_slots, HV, L]    fp32
+    ssm_state_indices: torch.Tensor,  # [B] int
+    accept_lens: torch.Tensor,  # [B] int  (incl. the bonus token)
+    max_cache_len: int,
+    num_k_heads: int,
+    mamba_track_indices: torch.Tensor | None = None,  # [B] extra_buffer track slot
+    mamba_steps_to_track: torch.Tensor | None = None,  # [B] crossing step (or -1)
+    use_qk_l2norm_in_kernel: bool = True,
+    null_block_id: int = 0,
+) -> None:
+    """Replay each request's accepted window into its fp32 checkpoint in place.
+
+    Tiling clones the recurrent kernel (full-K rows, BV = min(np2(V), 32) cols,
+    num_warps=1) so the folded checkpoint is bit-identical to the recurrent
+    baseline's committed state. With extra_buffer (mamba_track_indices given) the
+    same replay snapshots the interval-crossing state into the track slot in one
+    pass, so no separate track scatter / force-flush is needed.
+    """
+    num_slots, HV, V, K = checkpoint_state.shape
+    B = ssm_state_indices.shape[0]
+    BK = triton.next_power_of_2(K)
+    BV = min(triton.next_power_of_2(V), 32)
+    grid = (triton.cdiv(V, BV), B, HV)
+    has_track = mamba_track_indices is not None and mamba_steps_to_track is not None
+    if has_track:
+        track_idx_t = mamba_track_indices
+        steps_t = mamba_steps_to_track
+        stride_track = track_idx_t.stride(0)
+        stride_steps = steps_t.stride(0)
+    else:
+        track_idx_t = ssm_state_indices  # unused (HAS_TRACK False); pass a valid ptr
+        steps_t = accept_lens
+        stride_track = 0
+        stride_steps = 0
+    kda_replayssm_exact_fold_kernel[grid](
+        checkpoint_state,
+        rawv_cache,
+        rawk_cache,
+        gk_cache,
+        beta_cache,
+        ssm_state_indices,
+        accept_lens,
+        track_idx_t,
+        steps_t,
+        checkpoint_state.stride(0),
+        rawv_cache.stride(0),
+        rawk_cache.stride(0),
+        gk_cache.stride(0),
+        beta_cache.stride(0),
+        0,  # stride_*_layer unused: single-layer entry, i_layer == 0
+        0,
+        0,
+        0,
+        0,
+        ssm_state_indices.stride(0),
+        accept_lens.stride(0),
+        stride_track,
+        stride_steps,
+        H=num_k_heads,
+        HV=HV,
+        K=K,
+        V=V,
+        BK=BK,
+        BV=BV,
+        MAX_CACHE_LEN=max_cache_len,
+        USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
+        NULL_BLOCK_ID=null_block_id,
+        HAS_TRACK=has_track,
+        num_warps=1,
+        num_stages=3,
+    )
+
+
+def commit_kda_replayssm_spec_all_layers(
+    checkpoint_state: torch.Tensor,  # [num_layers, num_slots, HV, V, K] fp32, in place
+    rawv_cache: torch.Tensor,  # [num_layers, num_slots, HV, L, V]
+    rawk_cache: torch.Tensor,  # [num_layers, num_slots, H,  L, K]
+    gk_cache: torch.Tensor,  # [num_layers, num_slots, HV, L, K] fp32
+    beta_cache: torch.Tensor,  # [num_layers, num_slots, HV, L]    fp32
+    ssm_state_indices: torch.Tensor,  # [B] int   (shared across layers)
+    accept_lens: torch.Tensor,  # [B] int
+    max_cache_len: int,
+    num_k_heads: int,
+    mamba_track_indices: torch.Tensor | None = None,
+    mamba_steps_to_track: torch.Tensor | None = None,
+    use_qk_l2norm_in_kernel: bool = True,
+    null_block_id: int = 0,
+) -> None:
+    """Fold every layer's accepted window in a single launch.
+
+    Replaces the per-layer Python loop over commit_kda_replayssm_spec (one launch
+    per KDA layer -> ~69 tiny eager launches at bs=1, dispatch-bound). The layer
+    is packed into the head grid axis (program_id(2) = layer * HV + head), so the
+    result is bit-identical to the loop -- each (layer, head, v-tile) block runs
+    the same per-slot recurrent replay. ssm_state_indices / accept_lens / track
+    are per-request and shared across layers.
+    """
+    num_layers, num_slots, HV, V, K = checkpoint_state.shape
+    B = ssm_state_indices.shape[0]
+    BK = triton.next_power_of_2(K)
+    BV = min(triton.next_power_of_2(V), 32)
+    grid = (triton.cdiv(V, BV), B, HV * num_layers)
+    has_track = mamba_track_indices is not None and mamba_steps_to_track is not None
+    if has_track:
+        track_idx_t = mamba_track_indices
+        steps_t = mamba_steps_to_track
+        stride_track = track_idx_t.stride(0)
+        stride_steps = steps_t.stride(0)
+    else:
+        track_idx_t = ssm_state_indices  # unused (HAS_TRACK False); valid ptr
+        steps_t = accept_lens
+        stride_track = 0
+        stride_steps = 0
+    kda_replayssm_exact_fold_kernel[grid](
+        checkpoint_state,
+        rawv_cache,
+        rawk_cache,
+        gk_cache,
+        beta_cache,
+        ssm_state_indices,
+        accept_lens,
+        track_idx_t,
+        steps_t,
+        checkpoint_state.stride(1),
+        rawv_cache.stride(1),
+        rawk_cache.stride(1),
+        gk_cache.stride(1),
+        beta_cache.stride(1),
+        checkpoint_state.stride(0),
+        rawv_cache.stride(0),
+        rawk_cache.stride(0),
+        gk_cache.stride(0),
+        beta_cache.stride(0),
+        ssm_state_indices.stride(0),
+        accept_lens.stride(0),
+        stride_track,
+        stride_steps,
+        H=num_k_heads,
+        HV=HV,
+        K=K,
+        V=V,
+        BK=BK,
+        BV=BV,
+        MAX_CACHE_LEN=max_cache_len,
+        USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
+        NULL_BLOCK_ID=null_block_id,
+        HAS_TRACK=has_track,
+        num_warps=1,
+        num_stages=3,
+    )
+
+
+def commit_kda_replayssm_after_verify(
+    *,
+    spec_state,  # MambaPool.SpeculativeState (all layers)
+    state_batch_indices: torch.Tensor,  # [B] per-req mamba slot
+    accept_lens: torch.Tensor,  # [B] int, incl. the bonus token
+    last_correct_step_indices: torch.Tensor,  # [B] conv rollback target step
+    mamba_track_indices: torch.Tensor | None = None,
+    mamba_steps_to_track: torch.Tensor | None = None,
+    null_block_id: int = -1,
+) -> None:
+    """Fold each layer's accepted window into `temporal` and roll back conv.
+
+    Single commit entry point shared by the generic spec_utils commit and the
+    dspark/dflash direct `update_mamba_state_after_mtp_verify` path. The SSM state
+    lives in the per-slot ring (written during verify); the fold replays the
+    accepted prefix into the fp32 checkpoint, so `temporal` stays current. Conv
+    still needs its usual accept-rollback.
+    """
+    from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
+        fused_conv_window_scatter_with_mask,
+    )
+
+    L = spec_state.replayssm_rawv.shape[-2]
+    num_k_heads = spec_state.replayssm_rawk.shape[2]
+    commit_kda_replayssm_spec_all_layers(
+        checkpoint_state=spec_state.temporal,
+        rawv_cache=spec_state.replayssm_rawv,
+        rawk_cache=spec_state.replayssm_rawk,
+        gk_cache=spec_state.replayssm_g,
+        beta_cache=spec_state.replayssm_beta,
+        ssm_state_indices=state_batch_indices,
+        accept_lens=accept_lens,
+        max_cache_len=L,
+        num_k_heads=num_k_heads,
+        mamba_track_indices=mamba_track_indices,
+        mamba_steps_to_track=mamba_steps_to_track,
+        null_block_id=null_block_id,
+    )
+    # Conv rollback + track-slot conv snapshot, per conv group (fold already did
+    # the ssm side via HAS_TRACK). Loop mirrors the recurrent commit's zip; track
+    # scatter is mask-gated (step -1 => skip).
+    for conv_states, interm_conv in zip(
+        spec_state.conv, spec_state.intermediate_conv_window
+    ):
+        fused_conv_window_scatter_with_mask(
+            conv_states, interm_conv, state_batch_indices, last_correct_step_indices
+        )
+        if mamba_track_indices is not None and mamba_steps_to_track is not None:
+            fused_conv_window_scatter_with_mask(
+                conv_states, interm_conv, mamba_track_indices, mamba_steps_to_track
+            )

python/sglang/kernels/ops/attention/flash_attention_v4.py

@@ -23,6 +23,10 @@ else:
     _flash_attn_import_error = None
 
 
+def is_flash_attention_v4_available() -> bool:
+    return _flash_attn_varlen_func is not None
+
+
 def _maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
     return x.contiguous() if x is not None and x.stride(-1) != 1 else x
 

python/sglang/kernels/ops/attention/kda_fused_decode.py

@@ -0,0 +1,166 @@
+"""Fully fused KDA decode step (Kimi K3 batched decode fast path).
+
+One kernel replaces the three-kernel decode chain
+``causal_conv1d_update -> kda_packed_decode -> rms_norm_gated``: it reads the
+raw (pre-conv) qkv slice straight out of the fused projection GEMM output,
+does the causal conv1d update (conv state shifted in the pool in place), the
+delta-rule recurrence (l2-normed q/k, softplus forget gate, sigmoid beta),
+and the sigmoid-gated output RMSNorm.
+
+Kernel body vendored from the NVIDIA x Moonshot Kimi K3 optimization package
+(see csrc/attention/kda_fused_decode.cuh for provenance and the list of
+integration patches). Specialized for the K3 KDA decode regime:
+K = V = 128, kernel width 4, no lower bound, T = 1 per request.
+The JIT currently instantiates local head counts H = HV in {12, 6, 3}
+(TP8, TP16, and TP32).
+
+The model must hand off the output-norm gate (attempt-and-verify stash on the
+attention layer, see kimi_k3.py), and a covered() check gates supported inputs.
+Everything else falls back to the unfused chain.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+_SUPPORTED_HEADS = {3, 6, 12}
+_CONV_STATE_W = 3  # kernel width 4 -> 3 cached tokens
+
+
+@cache_once
+def _jit_kda_fused_decode_module() -> Module:
+    args = make_cpp_args(is_arch_support_pdl())
+    return load_jit(
+        "kda_fused_decode",
+        *args,
+        cuda_files=["attention/kda_fused_decode.cuh"],
+        cuda_wrappers=[("run", f"KdaFusedDecodeKernel<{args}>::run")],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+    )
+
+
+def covered(
+    mixed_qkv: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    conv_states: torch.Tensor,
+    ssm_states: torch.Tensor,
+    cache_indices: torch.Tensor,
+    onorm_g: torch.Tensor,
+) -> bool:
+    """The kernel is compiled for the K3 KDA decode regime: H heads of 128,
+    packed [T, 3*H*128] qkv rows, transposed [slots, 3, 3*H*128] conv pool, fp32
+    [slots, H, 128, 128] ssm pool (inner-contiguous, any slot pitch — the
+    kernel reads the real slot stride), one token per request."""
+    if ssm_states.ndim < 4:
+        return False
+    H, V, K = ssm_states.shape[-3:]
+    if H not in _SUPPORTED_HEADS:
+        return False
+    seg = H * 128
+    conv_dim = 3 * seg
+    if mixed_qkv.ndim != 2 or mixed_qkv.shape[-1] != conv_dim:
+        return False
+    return (
+        V == 128
+        and K == 128
+        and a.ndim == 2
+        and a.shape[-1] == seg
+        and b.ndim == 2
+        and b.shape[-1] == H
+        and onorm_g.ndim == 2
+        and onorm_g.shape[-1] == seg
+        and conv_states.ndim == 3
+        and conv_states.shape[-2:] == (_CONV_STATE_W, conv_dim)
+        and mixed_qkv.dtype == torch.bfloat16
+        and a.dtype == torch.bfloat16
+        and b.dtype == torch.bfloat16
+        and onorm_g.dtype == torch.bfloat16
+        and conv_states.dtype == torch.bfloat16
+        and ssm_states.dtype == torch.float32
+        and cache_indices.dtype == torch.int32
+        and mixed_qkv.stride(-1) == 1
+        and a.stride(-1) == 1
+        and b.stride(-1) == 1
+        and onorm_g.stride(-1) == 1
+        and conv_states.stride(-1) == 1
+        # Inner [HV, V, K] must be contiguous (the kernel float4-loads V*K
+        # chunks); the slot pitch (stride(-4)) is arbitrary — a locally
+        # allocated pool packs it at HV*V*K, the unified / page-major pools at
+        # the multi-layer envelope. The kernel reads ssm_states.stride(0), so
+        # any slot pitch is fine. (Do NOT use .view(-1, HV, V, K): that fails /
+        # copies on an envelope-strided view.)
+        and ssm_states.stride(-1) == 1
+        and ssm_states.stride(-2) == K
+        and ssm_states.stride(-3) == V * K
+        and cache_indices.is_contiguous()
+    )
+
+
+def kda_fused_decode(
+    mixed_qkv: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    conv_states: torch.Tensor,
+    w_q_t: torch.Tensor,
+    w_k_t: torch.Tensor,
+    w_v_t: torch.Tensor,
+    conv_bias: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    onorm_g: torch.Tensor,
+    onorm_weight: torch.Tensor,
+    ssm_states: torch.Tensor,
+    cache_indices: torch.Tensor,
+    scale: float,
+    onorm_eps: float,
+    lower_bound: Optional[float] = None,
+) -> torch.Tensor:
+    """In-place fused decode step: shifts `conv_states` and updates
+    `ssm_states` rows selected by `cache_indices` (rows < 0 are padded
+    cuda-graph slots and only zero their output), returns the gated-normed
+    attention output [1, B, HV, V] (the packed-decode output layout).
+    Caller must have checked covered()."""
+    B = mixed_qkv.shape[0]
+    H = ssm_states.shape[-3]
+    seg = H * 128
+    out = torch.empty((B, seg), dtype=torch.bfloat16, device=mixed_qkv.device)
+    _jit_kda_fused_decode_module().run(
+        mixed_qkv,
+        a,
+        b,
+        conv_states,
+        w_q_t,
+        w_k_t,
+        w_v_t,
+        conv_bias,
+        A_log,
+        dt_bias,
+        onorm_g,
+        onorm_weight,
+        # Pass the pool view as-is (already [slots, HV, V, K]); the kernel
+        # binding reads its real slot stride via state.stride(0). A
+        # .view(-1, H, 128, 128) here would break on envelope-strided pools
+        # (unified / page-major) — the reshape can't fold a non-dense slot
+        # pitch and would raise / silently copy.
+        ssm_states,
+        cache_indices,
+        out,
+        float(scale),
+        float(onorm_eps),
+        float(lower_bound) if lower_bound is not None else 0.0,
+        lower_bound is not None,
+    )
+    return out.view(1, B, H, 128)

python/sglang/kernels/ops/attention/kda_packed_decode.py

@@ -0,0 +1,120 @@
+"""CUDA KDA packed-decode kernel (batched decode fast path).
+
+Row-streaming port of the triton fused_recurrent_kda_packed_decode_kernel:
+the triton kernel keeps a [BV, K] fp32 state tile in one warp's registers and
+tops out at ~5 TB/s; this kernel streams the state one 512B row at a time and
+reaches the in-place read+write bandwidth of the part (~9.6 TB/s probe).
+Outputs match the triton kernel to ULPs (warp-shuffle reduction order), not
+bits. Unsupported inputs fall back to triton through a covered() check.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+_WARPS: int = 8
+# The row-streaming layout needs enough (batch x head) CTAs to fill the GPU;
+# below this the triton kernel's launch cost is already the floor.
+_MIN_BATCH: int = 8
+
+
+@cache_once
+def _jit_kda_packed_decode_module() -> Module:
+    args = make_cpp_args(_WARPS, is_arch_support_pdl())
+    return load_jit(
+        "kda_packed_decode_" + str(_WARPS),
+        *args,
+        cuda_files=["attention/kda_packed_decode.cuh"],
+        cuda_wrappers=[("run", f"KdaPackedDecodeKernel<{args}>::run")],
+        extra_cuda_cflags=["-O3"],
+    )
+
+
+def covered(
+    mixed_qkv: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    initial_state: torch.Tensor,
+    out: torch.Tensor,
+    ssm_state_indices: torch.Tensor,
+    num_q_heads: int,
+) -> bool:
+    B = mixed_qkv.shape[0]
+    HV, V, K = initial_state.shape[-3:]
+    return (
+        B >= _MIN_BATCH
+        and K == 128
+        and V == 128
+        and HV % max(num_q_heads, 1) == 0
+        # Per-K gate layout. A per-head scalar gate ([B, HV] / [HV], the GDN
+        # shape) is a different kernel, not a slower input for this one.
+        and a.dim() == 2
+        and a.shape[1] == HV * K
+        and dt_bias.numel() == HV * K
+        and mixed_qkv.dtype == torch.bfloat16
+        and a.dtype == torch.bfloat16
+        and b.dtype == torch.bfloat16
+        and A_log.dtype == torch.float32
+        and dt_bias.dtype == torch.float32
+        and initial_state.dtype == torch.float32
+        and out.dtype == torch.bfloat16
+        and ssm_state_indices.dtype == torch.int32
+        and mixed_qkv.stride(-1) == 1
+        and a.stride(-1) == 1
+        and b.stride(-1) == 1
+        and initial_state.stride(-1) == 1
+        and initial_state.stride(-2) == K
+        and initial_state.stride(-3) == V * K
+        and out.is_contiguous()
+        and ssm_state_indices.is_contiguous()
+    )
+
+
+def kda_packed_decode(
+    mixed_qkv: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    out: torch.Tensor,
+    ssm_state_indices: torch.Tensor,
+    num_q_heads: int,
+    lower_bound: Optional[float] = None,
+) -> None:
+    """In-place KDA decode step: updates `initial_state` rows selected by
+    `ssm_state_indices` and writes attention output into `out` ([B, 1, HV, V]).
+    Caller must have checked covered(); q/k l2-norm is always applied
+    (matches the production dispatch)."""
+    B = mixed_qkv.shape[0]
+    HV, V, _ = initial_state.shape[-3:]
+    state = initial_state.view(-1, *initial_state.shape[-3:])
+    _jit_kda_packed_decode_module().run(
+        mixed_qkv,
+        a,
+        b,
+        A_log,
+        dt_bias,
+        out.view(B, HV, V),
+        state,
+        ssm_state_indices,
+        float(scale),
+        float(lower_bound) if lower_bound is not None else 0.0,
+        lower_bound is not None,
+        int(num_q_heads),
+    )

python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_h.py

@@ -698,7 +698,19 @@ def kda_h_cutedsl(
     ``h0``/``ht`` may be the full state pool; ``state_indices`` [N] int32 maps
     each sequence to its row, so state gather/scatter fuses into the kernel's
     TMA load/store (no per-call state intermediates).
+
+    Envelope-strided pools (unified memory / page-major, slot stride(0) !=
+    Hv*V*K) are supported natively: the compile-time fake tensors carry
+    dynamic int64 strides (``make_fake_tensor``), so the TMA descriptors pick
+    the real slot pitch up at launch. The fakes assume 16-element stride
+    divisibility (TMA also needs 16-byte global strides); guard it loudly
+    rather than corrupt state.
     """
+    for name, t in (("h0", h0), ("ht", ht)):
+        assert t.stride(-1) == 1 and all(s % 16 == 0 for s in t.stride()[:-1]), (
+            f"kda_h_cutedsl: {name} strides {tuple(t.stride())} violate the "
+            "16-element divisibility the kernel was compiled with"
+        )
     _, Hv, K_dim = kg.shape
     _, _, V_dim = V.shape
     h_dtype = {torch.bfloat16: BFloat16, torch.float32: Float32}[h0.dtype]

python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# NVIDIA KDA_prefill (Blackwell): optimized chunked KDA forward, an
+# FLA-compatible replacement for chunk_kda_fwd. K1 (gate+cumsum+scale, CuTe)
+# + K2 (intra sub-chunk, CuTe) + K3 (inter-chunk solve, Triton) + K4
+# (W/U/v_new/O/state update, cuTile persistent). 2.3-2.9x vs the FLA
+# reference on B200 in the upstream package tests.
+from .chunk_fwd import chunk_kda_fwd
+
+__all__ = ["chunk_kda_fwd"]

python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py

@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Hand-written PTX/tcgen05 KDA chunked-prefill kernel (GB300 / sm_103a).
+
+Vendored from the upstream ``kda_prefill`` artifact (commit 33583615): the CUDA
+source in ``kernels/jit/csrc/attention/kda_prefill.cu`` plus the FLA-signature
+``chunk_kda_fwd`` wrapper below, which is a drop-in for
+``fla.ops.kda.chunk_fwd.chunk_kda_fwd`` on the inference forward path
+(K = V = 128, chunk_size = 64).
+
+The extension is JIT-compiled with ``torch.utils.cpp_extension`` on first use
+(~1-2 min, cached under ``TORCH_EXTENSIONS_DIR``); concurrent TP ranks
+serialize on torch's build lock and then load the cached .so.
+"""
+
+import os
+
+import torch
+
+K = 128  # head dim (qk == v), fixed by the kernel
+CHUNK = 64  # kernel chunk size, fixed
+
+_EXT_NAME = "kda_prefill_ptx"
+_ext = None
+
+
+def load_ext():
+    """JIT-load the CUDA extension (cached after first call)."""
+    global _ext
+    if _ext is None:
+        from torch.utils import cpp_extension
+
+        src = os.path.join(
+            os.path.dirname(os.path.abspath(__file__)),
+            "../../../../jit/csrc/attention/kda_prefill.cu",
+        )
+        # -lcuda: cuTensorMapEncodeTiled (driver API). The stubs dir covers
+        # boxes whose real libcuda.so lives off the default linker path;
+        # ld.so still binds the driver's libcuda.so.1 at import time.
+        stubs = os.path.join(
+            cpp_extension.CUDA_HOME or "/usr/local/cuda", "lib64", "stubs"
+        )
+        _ext = cpp_extension.load(
+            name=_EXT_NAME,
+            sources=[src],
+            extra_cuda_cflags=[
+                "-O3",
+                "-std=c++20",
+                "-use_fast_math",
+                "-lineinfo",
+                "-gencode",
+                "arch=compute_103a,code=sm_103a",
+            ],
+            extra_cflags=["-O3"],
+            extra_ldflags=[f"-L{stubs}", "-lcuda"],
+        )
+    return _ext
+
+
+def chunk_kda_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor | None,
+    output_final_state: bool,
+    state_v_first: bool = False,
+    cu_seqlens: torch.Tensor | None = None,
+    cu_seqlens_cpu: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_size: int = 64,
+    safe_gate: bool = False,
+    lower_bound: float | None = None,
+    use_gate_in_kernel: bool = False,
+    A_log: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    disable_recompute: bool = False,
+    return_intermediate_states: bool = False,
+    cp_context=None,
+    use_qk_l2norm_in_kernel: bool = False,
+    use_beta_sigmoid_in_kernel: bool = False,
+    allow_neg_eigval: bool = False,
+):
+    """Inference drop-in for fla's chunk_kda_fwd (forward only).
+
+    Argument mapping onto the CUDA kernel:
+      q/k/v [B,T,H,128] bf16 -> flat [T,H,128] (B>1 folds to varlen with a
+        synthetic equal-length cu_seqlens; B==1 squeezes).
+      g [B,T,H,128]: use_gate_in_kernel=False -> pre-transformed glog, fp32
+        narrowed to bf16 here (the kernel's GM=0 contract); =True -> RAW gate
+        input (bf16) with the transform fused in-kernel (GM=1 softplus /
+        GM=2 safe-gate). Following fla, the safe-gate TRANSFORM is selected
+        by `lower_bound is not None`; the `safe_gate` flag alone is fla's
+        intra-path hint and does not change the math here.
+      beta [B,T,H] bf16 (fp32 also accepted; widened to fp32 in the ext).
+      use_qk_l2norm_in_kernel=True accepts raw q/k and applies FLA-compatible
+        L2 normalization (eps=1e-6, bf16 rounding) in the CUDA tile loads.
+      use_beta_sigmoid_in_kernel=True accepts beta logits and fuses sigmoid.
+      cu_seqlens: host values are needed for the kernel's per-sequence piece
+        table -- pass cu_seqlens_cpu to avoid the D2H sync; chunk_indices is
+        accepted and ignored (the kernel derives its own piece table).
+      initial_state [N,H,128,128] fp32 or None (zeros).
+      return_intermediate_states=True returns dense fp32 chunk-boundary states
+        [1, NT, H, 128, 128] at tuple index 10.
+
+    Returns the fla-shaped 12-tuple: (o [B,T,H,128] bf16, final_state
+    [N,H,128,128] fp32 or None, then Nones, ..., h, initial_state).
+    """
+    assert (
+        chunk_size == CHUNK
+    ), f"kda_prefill supports chunk_size={CHUNK} only, got {chunk_size}"
+    if cp_context is not None or disable_recompute:
+        raise NotImplementedError(
+            "kda_prefill is the inference forward path: cp_context, "
+            "and disable_recompute are training-side knobs it does not implement"
+        )
+    if allow_neg_eigval and use_beta_sigmoid_in_kernel:
+        raise NotImplementedError(
+            "allow_neg_eigval=True requires 2*sigmoid(beta), which is not "
+            "implemented by the fused beta path; pass pre-activated beta with "
+            "use_beta_sigmoid_in_kernel=False"
+        )
+    if state_v_first and initial_state is not None:
+        # [V,K]-layout state: pure transpose (K==V==128), exact, ~us/call
+        initial_state = initial_state.transpose(-1, -2).contiguous()
+    assert (
+        q.dim() == 4 and q.shape[-1] == K and v.shape[-1] == K
+    ), f"expected [B,T,H,{K}] q/k/v, got q={tuple(q.shape)} v={tuple(v.shape)}"
+    B, T, H, _ = q.shape
+
+    cu_cpu = None
+    if cu_seqlens is not None or cu_seqlens_cpu is not None:
+        assert B == 1, "cu_seqlens requires B == 1 (flattened varlen batch)"
+        src = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens
+        cu_cpu = torch.as_tensor(src, dtype=torch.int32).cpu()
+    elif B > 1:
+        # eqlen batch == varlen with equal lengths (host-known, no sync)
+        cu_cpu = torch.arange(0, (B + 1) * T, T, dtype=torch.int32)
+
+    Tt = B * T
+    qf = q.reshape(Tt, H, K).contiguous()
+    kf = k.reshape(Tt, H, K).contiguous()
+    vf = v.reshape(Tt, H, K).contiguous()
+    betaf = beta.reshape(Tt, H).contiguous()
+
+    if use_gate_in_kernel:
+        assert (
+            A_log is not None and dt_bias is not None
+        ), "use_gate_in_kernel=True requires A_log and dt_bias"
+        assert g.dtype == torch.bfloat16, f"raw gate input must be bf16, got {g.dtype}"
+        gf = g.reshape(Tt, H, K).contiguous()
+        sg = lower_bound is not None  # fla: lb presence selects safe-gate
+        a_log_flat = A_log.reshape(H).to(torch.float32).contiguous()
+        dtb = dt_bias.reshape(H * K).to(torch.float32).contiguous()
+        lb = float(lower_bound) if lower_bound is not None else 0.0
+    else:
+        gf = g.reshape(Tt, H, K)
+        if gf.dtype != torch.bfloat16:  # pre-transformed glog: bf16 narrow
+            gf = gf.to(torch.bfloat16)
+        gf = gf.contiguous()
+        sg, a_log_flat, dtb, lb = False, None, None, 0.0
+
+    h = None
+    if return_intermediate_states:
+        lens = (cu_cpu[1:] - cu_cpu[:-1]).tolist() if cu_cpu is not None else [Tt]
+        nt = sum((int(length) + CHUNK - 1) // CHUNK for length in lens)
+        h = torch.empty(nt, H, K, K, dtype=torch.float32, device=q.device)
+
+    o, Sf = load_ext().kda_prefill_fwd(
+        qf,
+        kf,
+        vf,
+        gf,
+        betaf,
+        float(scale),
+        initial_state=initial_state,
+        cu_seqlens=cu_cpu,
+        use_gate_in_kernel=use_gate_in_kernel,
+        A_log=a_log_flat,
+        dt_bias=dtb,
+        safe_gate=sg,
+        lower_bound=lb,
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+        use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
+        h_per_chunk=h,
+        h_v_first=state_v_first,
+    )
+
+    o = o.view(B, T, H, K)
+    if state_v_first:
+        Sf = Sf.transpose(-1, -2).contiguous()
+    final = Sf if output_final_state else None
+    return (
+        o,
+        final,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None if h is None else h.unsqueeze(0),
+        initial_state,
+    )

python/sglang/kernels/ops/attention/set_mla_kv_concat_q.py

@@ -0,0 +1,318 @@
+"""Fused MLA decode prepare tail: paged-KV scatter + absorbed-q concat.
+
+One launch replacing the back-to-back ``set_mla_kv_buffer`` +
+``concat_mla_absorb_q`` pair on the trtllm-mla decode graph path. Both
+workloads are launch-bound data movement at decode batch sizes; the fusion
+saves a kernel launch per MLA layer and keeps the PDL chain to the decode
+fmha kernel intact. SM90+ only (TMA bulk store) — gate via ``covered()``.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def set_mla_kv_concat_q_module(
+    nope_bytes: int, rope_bytes: int, use_pdl: bool
+) -> Module:
+    args = make_cpp_args(nope_bytes, rope_bytes, use_pdl)
+    return load_jit(
+        f"set_mla_kv_concat_q_{nope_bytes}_{rope_bytes}",
+        *args,
+        cuda_files=["elementwise/set_mla_kv_concat_q.cuh"],
+        cuda_wrappers=[
+            ("set_mla_kv_concat_q", f"SetMlaKVConcatQKernel<{args}>::run"),
+        ],
+    )
+
+
+@cache_once
+def can_use_set_mla_kv_concat_q(nope_bytes: int, rope_bytes: int) -> bool:
+    """Whether the fused kernel supports these row byte widths on this arch.
+
+    Static gate (per process): SM90+ for the TMA bulk store, plus the
+    compile-time width constraints (total row 16B-aligned for TMA; nope dim
+    filling whole int4 warp rounds and rope dim exactly one int warp round
+    for the concat side — 1024/128 bytes i.e. 512/64 bf16 satisfies both).
+    """
+    if torch.cuda.get_device_capability()[0] < 9:
+        return False
+    if nope_bytes % 4 != 0 or rope_bytes % 4 != 0:
+        return False
+    if (nope_bytes + rope_bytes) % 16 != 0:
+        return False
+    # Concat-side vector layout: int4 (8 bf16) x 32 lanes per round for nope,
+    # one int (2 bf16) x 32 lanes round for rope.
+    if (nope_bytes // 2) % (8 * 32) != 0 or rope_bytes // 2 != 2 * 32:
+        return False
+    try:
+        set_mla_kv_concat_q_module(nope_bytes, rope_bytes, is_arch_support_pdl())
+        return True
+    except Exception:  # pragma: no cover - compile-time only
+        return False
+
+
+def _row_aligned(t: torch.Tensor, align: int) -> bool:
+    """Base pointer and all non-unit strides aligned to ``align`` bytes."""
+    if t.data_ptr() % align != 0:
+        return False
+    esize = t.element_size()
+    return all(s * esize % align == 0 for s in t.stride()[:-1])
+
+
+def covered(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    k_nope: torch.Tensor,
+    k_rope: torch.Tensor,
+    q_nope: torch.Tensor,
+    q_rope: torch.Tensor,
+) -> bool:
+    """Per-call gate mirroring the launcher's alignment/layout tripwires, so
+    uncovered layouts fall back to the two-kernel path instead of faulting.
+
+    Expects the already-flattened views the wrapper launches with:
+    kv_buffer [pages, D], k_nope/k_rope [B, d], q_nope/q_rope [B, H, d],
+    loc [B].
+    """
+    if not (
+        kv_buffer.dtype == torch.bfloat16
+        and k_nope.dtype == torch.bfloat16
+        and k_rope.dtype == torch.bfloat16
+        and q_nope.dtype == torch.bfloat16
+        and q_rope.dtype == torch.bfloat16
+    ):
+        return False
+    if loc.dtype not in (torch.int32, torch.int64):
+        return False
+    if loc.dim() != 1 or not loc.is_contiguous():
+        return False
+    if not (
+        k_nope.shape[0]
+        == k_rope.shape[0]
+        == loc.shape[0]
+        == q_nope.shape[0]
+        == q_rope.shape[0]
+    ):
+        return False
+    if q_nope.shape[1] != q_rope.shape[1]:
+        return False
+    nope_bytes = k_nope.shape[-1] * 2
+    rope_bytes = k_rope.shape[-1] * 2
+    if q_nope.shape[-1] * 2 != nope_bytes or q_rope.shape[-1] * 2 != rope_bytes:
+        return False
+    if not can_use_set_mla_kv_concat_q(nope_bytes, rope_bytes):
+        return False
+    # Last dims must be dense for the vectorised row accesses.
+    if any(t.stride(-1) != 1 for t in (kv_buffer, k_nope, k_rope, q_nope, q_rope)):
+        return False
+    return (
+        _row_aligned(kv_buffer, 16)
+        and _row_aligned(k_nope, 16)
+        and _row_aligned(k_rope, 4)
+        and _row_aligned(q_nope, 16)
+        and _row_aligned(q_rope, 4)
+    )
+
+
+def _pick_num_warps(total_items: int) -> int:
+    # Same GB300 heuristic as set_mla_kv_buffer: small grids favour more CTAs.
+    return 4 if total_items <= 768 else 8
+
+
+def set_mla_kv_concat_q(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: torch.Tensor,
+    q_nope: torch.Tensor,
+    q_rope: torch.Tensor,
+    num_warps: int = 0,
+) -> torch.Tensor:
+    """Scatter [k_nope | k_rope] rows into ``kv_buffer`` at ``loc`` and return
+    the concatenated query [q_nope | q_rope], all in one kernel launch.
+
+    Shapes (leading singleton dims on the k sources are flattened away):
+        kv_buffer:    [num_pages, total_dim] or [num_pages, 1, total_dim]
+        loc:          [n_loc]
+        cache_k_nope: [n_loc, nope_dim] or [n_loc, 1, nope_dim]
+        cache_k_rope: [n_loc, rope_dim] or [n_loc, 1, rope_dim]
+        q_nope:       [n_loc, num_heads, nope_dim]
+        q_rope:       [n_loc, num_heads, rope_dim]
+    Returns:
+        query: [n_loc, num_heads, nope_dim + rope_dim] (new contiguous tensor)
+    """
+    n_loc = loc.shape[0]
+    src_nope = cache_k_nope.view(n_loc, -1) if cache_k_nope.dim() != 2 else cache_k_nope
+    src_rope = cache_k_rope.view(n_loc, -1) if cache_k_rope.dim() != 2 else cache_k_rope
+    buf = kv_buffer.view(kv_buffer.shape[0], -1) if kv_buffer.dim() != 2 else kv_buffer
+
+    q_out = torch.empty(
+        (*q_nope.shape[:-1], q_nope.shape[-1] + q_rope.shape[-1]),
+        dtype=q_nope.dtype,
+        device=q_nope.device,
+    )
+
+    nope_bytes = src_nope.shape[-1] * src_nope.element_size()
+    rope_bytes = src_rope.shape[-1] * src_rope.element_size()
+    if num_warps <= 0:
+        num_warps = _pick_num_warps(n_loc + q_nope.shape[0] * q_nope.shape[1])
+
+    module = set_mla_kv_concat_q_module(nope_bytes, rope_bytes, is_arch_support_pdl())
+    module.set_mla_kv_concat_q(
+        buf, loc, src_nope, src_rope, q_nope, q_rope, q_out, num_warps
+    )
+    return q_out
+
+
+@cache_once
+def set_mla_kv_concat_q_fp8_module(use_pdl: bool) -> Module:
+    args = make_cpp_args(use_pdl)
+    return load_jit(
+        "set_mla_kv_concat_q_fp8",
+        *args,
+        cuda_files=["elementwise/set_mla_kv_concat_q.cuh"],
+        cuda_wrappers=[
+            ("set_mla_kv_concat_q_fp8", f"SetMlaKVConcatQFp8Kernel<{args}>::run"),
+        ],
+    )
+
+
+@cache_once
+def can_use_set_mla_kv_concat_q_fp8() -> bool:
+    """SM90+ (TMA bulk store) and the module compiles. Row widths are fixed
+    at 512/64 (the MLA absorb layout) inside the kernel."""
+    if torch.cuda.get_device_capability()[0] < 9:
+        return False
+    try:
+        set_mla_kv_concat_q_fp8_module(is_arch_support_pdl())
+        return True
+    except Exception:  # pragma: no cover - compile-time only
+        return False
+
+
+def covered_fp8(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    k_nope: torch.Tensor,
+    k_rope: torch.Tensor,
+    q_nope: torch.Tensor,
+    q_rope: torch.Tensor,
+) -> bool:
+    """Per-call gate for the fused fp8 quantize+scatter+concat kernel,
+    mirroring the launcher tripwires. Expects flattened views: kv_buffer
+    [pages, 576] fp8/uint8, k halves [B, 512/64] bf16, q halves
+    [B, H, 512/64] bf16, loc [B]."""
+    if not (
+        kv_buffer.dtype in (torch.float8_e4m3fn, torch.uint8)
+        and k_nope.dtype == torch.bfloat16
+        and k_rope.dtype == torch.bfloat16
+        and q_nope.dtype == torch.bfloat16
+        and q_rope.dtype == torch.bfloat16
+    ):
+        return False
+    if loc.dtype not in (torch.int32, torch.int64):
+        return False
+    if loc.dim() != 1 or not loc.is_contiguous():
+        return False
+    if not (
+        k_nope.shape[0]
+        == k_rope.shape[0]
+        == loc.shape[0]
+        == q_nope.shape[0]
+        == q_rope.shape[0]
+    ):
+        return False
+    if q_nope.shape[1] != q_rope.shape[1]:
+        return False
+    if (
+        k_nope.shape[-1] != 512
+        or k_rope.shape[-1] != 64
+        or q_nope.shape[-1] != 512
+        or q_rope.shape[-1] != 64
+        or kv_buffer.shape[-1] < 576
+    ):
+        return False
+    if not can_use_set_mla_kv_concat_q_fp8():
+        return False
+    if any(t.stride(-1) != 1 for t in (kv_buffer, k_nope, k_rope, q_nope, q_rope)):
+        return False
+    if kv_buffer.data_ptr() % 16 != 0 or kv_buffer.stride(0) % 16 != 0:
+        return False
+    return (
+        _row_aligned(k_nope, 16)
+        and _row_aligned(k_rope, 4)
+        and _row_aligned(q_nope, 16)
+        and _row_aligned(q_rope, 4)
+    )
+
+
+def set_mla_kv_concat_q_fp8(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: torch.Tensor,
+    q_nope: torch.Tensor,
+    q_rope: torch.Tensor,
+    num_warps: int = 0,
+    dcp_world_size: int = 1,
+    dcp_rank: int = 0,
+) -> torch.Tensor:
+    """Quantize bf16 [k_nope | k_rope] rows to fp8-e4m3 and scatter them into
+    ``kv_buffer`` at ``loc``, and return the fp8 concatenated query
+    [q_nope | q_rope], all in one kernel launch (replaces concat + three
+    aten fp8 casts + the KV-row write on the fp8 decode path).
+
+    Under DCP, ``loc`` is VIRTUAL: the physical row is ``loc //
+    dcp_world_size`` and only the owner rank (``loc % dcp_world_size ==
+    dcp_rank``) writes its KV row (query conversion still runs for every
+    token). world=1/rank=0 is the non-DCP identity.
+
+    Shapes (leading singleton dims on the k sources are flattened away):
+        kv_buffer:    [num_pages, 576] fp8_e4m3/uint8 (or [num_pages, 1, 576])
+        loc:          [n_loc]
+        cache_k_nope: [n_loc, 512] bf16   cache_k_rope: [n_loc, 64] bf16
+        q_nope:       [n_loc, H, 512] bf16   q_rope: [n_loc, H, 64] bf16
+    Returns:
+        query: [n_loc, H, 576] float8_e4m3fn (new contiguous tensor)
+    """
+    n_loc = loc.shape[0]
+    src_nope = cache_k_nope.view(n_loc, -1) if cache_k_nope.dim() != 2 else cache_k_nope
+    src_rope = cache_k_rope.view(n_loc, -1) if cache_k_rope.dim() != 2 else cache_k_rope
+    buf = kv_buffer.view(kv_buffer.shape[0], -1) if kv_buffer.dim() != 2 else kv_buffer
+
+    q_out = torch.empty(
+        (q_nope.shape[0], q_nope.shape[1], 576),
+        dtype=torch.float8_e4m3fn,
+        device=q_nope.device,
+    )
+    if num_warps <= 0:
+        num_warps = _pick_num_warps(n_loc + q_nope.shape[0] * q_nope.shape[1])
+
+    module = set_mla_kv_concat_q_fp8_module(is_arch_support_pdl())
+    module.set_mla_kv_concat_q_fp8(
+        buf,
+        loc,
+        src_nope,
+        src_rope,
+        q_nope,
+        q_rope,
+        q_out,
+        num_warps,
+        dcp_world_size,
+        dcp_rank,
+    )
+    return q_out

python/sglang/kernels/ops/attention/verify_splitkv.py

@@ -149,6 +149,8 @@ def _verify_prefix_stage1(
     kv_group_num: tl.constexpr,
     N_SPLITS: tl.constexpr,
     L_EXT: tl.constexpr,  # padded power-of-2 row tile (>= real l_ext)
+    HEAD_DIM: tl.constexpr,
+    V_HEAD_DIM: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_DV: tl.constexpr,
     BLOCK_N: tl.constexpr,
@@ -190,7 +192,11 @@ def _verify_prefix_stage1(
             + cur_head * stride_qh
             + offs_d[None, :]
         )
-        q = tl.load(Q + offs_q, mask=mask_l[:, None], other=0.0)
+        q = tl.load(
+            Q + offs_q,
+            mask=mask_l[:, None] & (offs_d[None, :] < HEAD_DIM),
+            other=0.0,
+        )
         q_k = q.to(K_Buffer.dtype.element_ty)
 
         base_offs_k = cur_kv_head * stride_buf_kh + offs_d[:, None]
@@ -206,7 +212,11 @@ def _verify_prefix_stage1(
             )
             # K block: [D, BLOCK_N]
             offs_buf_k = kv_loc[None, :] * stride_buf_kbs + base_offs_k
-            k = tl.load(K_Buffer + offs_buf_k, mask=n_mask[None, :], other=0.0)
+            k = tl.load(
+                K_Buffer + offs_buf_k,
+                mask=(offs_d[:, None] < HEAD_DIM) & n_mask[None, :],
+                other=0.0,
+            )
             qk = tl.dot(q_k, k)  # [L_EXT, BLOCK_N]
             qk *= sm_scale * k_scale  # fp8 dequant of prefix K (k_scale==1 if bf16)
             # NO causal mask: full prefix is visible to all draft tokens.
@@ -214,7 +224,11 @@ def _verify_prefix_stage1(
 
             # V block: [BLOCK_N, Dv]
             offs_buf_v = kv_loc[:, None] * stride_buf_vbs + base_offs_v
-            v = tl.load(V_Buffer + offs_buf_v, mask=n_mask[:, None], other=0.0)
+            v = tl.load(
+                V_Buffer + offs_buf_v,
+                mask=n_mask[:, None] & (offs_dv[None, :] < V_HEAD_DIM),
+                other=0.0,
+            )
 
             n_e_max = tl.maximum(tl.max(qk, 1), e_max)
             re_scale = tl.exp(e_max - n_e_max)
@@ -234,7 +248,11 @@ def _verify_prefix_stage1(
             + offs_l[:, None] * stride_ol
             + offs_dv[None, :]
         )
-        tl.store(Att_Out + offs_o, acc / e_sum[:, None], mask=mask_l[:, None])
+        tl.store(
+            Att_Out + offs_o,
+            acc / e_sum[:, None],
+            mask=mask_l[:, None] & (offs_dv[None, :] < V_HEAD_DIM),
+        )
 
         offs_lse = (
             cur_batch * stride_lb
@@ -286,6 +304,8 @@ def _verify_combine_stage2(
     kv_group_num: tl.constexpr,
     N_SPLITS: tl.constexpr,
     L_EXT: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    V_HEAD_DIM: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_DV: tl.constexpr,
 ):
@@ -310,7 +330,11 @@ def _verify_combine_stage2(
         + offs_s[:, None] * stride_ls
         + offs_l[None, :]
     )
-    lse = tl.load(offs_lse + Att_Lse)  # [N_SPLITS, L_EXT]
+    lse = tl.load(
+        offs_lse + Att_Lse,
+        mask=mask_l[None, :],
+        other=float("-inf"),
+    )  # [N_SPLITS, L_EXT]
     m_p = tl.max(lse, 0)  # [L_EXT]
     w = tl.exp(lse - m_p[None, :])  # [N_SPLITS, L_EXT]; -inf->0
     denom_p = tl.sum(w, 0)  # [L_EXT]
@@ -324,7 +348,11 @@ def _verify_combine_stage2(
         + offs_l[None, :, None] * stride_ol
         + offs_dv[None, None, :]
     )
-    ao = tl.load(offs_ao + Att_Out)  # [N_SPLITS, L_EXT, Dv]
+    ao = tl.load(
+        offs_ao + Att_Out,
+        mask=mask_l[None, :, None] & (offs_dv[None, None, :] < V_HEAD_DIM),
+        other=0.0,
+    )  # [N_SPLITS, L_EXT, Dv]
     o_prefix = tl.sum(ao * w[:, :, None], 0)  # [L_EXT, Dv]
     o_prefix = o_prefix / denom_p[:, None]
     lse_prefix = m_p + tl.log(denom_p)  # [L_EXT]
@@ -336,20 +364,32 @@ def _verify_combine_stage2(
         + cur_head * stride_qh
         + offs_d[None, :]
     )
-    q = tl.load(Q + offs_q, mask=mask_l[:, None], other=0.0).to(tl.float32)
+    q = tl.load(
+        Q + offs_q,
+        mask=mask_l[:, None] & (offs_d[None, :] < HEAD_DIM),
+        other=0.0,
+    ).to(tl.float32)
 
     offs_ke = (
         (cur_q_start + offs_l)[:, None] * stride_kebs
         + cur_kv_head * stride_keh
         + offs_d[None, :]
     )
-    ke = tl.load(K_Extend + offs_ke, mask=mask_l[:, None], other=0.0).to(tl.float32)
+    ke = tl.load(
+        K_Extend + offs_ke,
+        mask=mask_l[:, None] & (offs_d[None, :] < HEAD_DIM),
+        other=0.0,
+    ).to(tl.float32)
     offs_ve = (
         (cur_q_start + offs_l)[:, None] * stride_vebs
         + cur_kv_head * stride_veh
         + offs_dv[None, :]
     )
-    ve = tl.load(V_Extend + offs_ve, mask=mask_l[:, None], other=0.0).to(tl.float32)
+    ve = tl.load(
+        V_Extend + offs_ve,
+        mask=mask_l[:, None] & (offs_dv[None, :] < V_HEAD_DIM),
+        other=0.0,
+    ).to(tl.float32)
 
     # scores[i,j] = q_i . k_j  (i query, j key)  -> [L_EXT, L_EXT]
     qk = tl.sum(q[:, None, :] * ke[None, :, :], 2) * sm_scale
@@ -374,7 +414,11 @@ def _verify_combine_stage2(
         + cur_head * stride_ooh
         + offs_dv[None, :]
     )
-    tl.store(O_Out + offs_oo, o.to(O_Out.dtype.element_ty), mask=mask_l[:, None])
+    tl.store(
+        O_Out + offs_oo,
+        o.to(O_Out.dtype.element_ty),
+        mask=mask_l[:, None] & (offs_dv[None, :] < V_HEAD_DIM),
+    )
 
 
 class VerifySplitKV:
@@ -471,6 +515,8 @@ class VerifySplitKV:
             kv_group_num=self.group,
             N_SPLITS=self.n_splits,
             L_EXT=self.l_pad,
+            HEAD_DIM=self.head_dim,
+            V_HEAD_DIM=self.v_head_dim,
             BLOCK_DMODEL=triton.next_power_of_2(self.head_dim),
             BLOCK_DV=triton.next_power_of_2(self.v_head_dim),
             BLOCK_N=self.block_n,
@@ -511,6 +557,8 @@ class VerifySplitKV:
             kv_group_num=self.group,
             N_SPLITS=self.n_splits,
             L_EXT=self.l_pad,
+            HEAD_DIM=self.head_dim,
+            V_HEAD_DIM=self.v_head_dim,
             BLOCK_DMODEL=triton.next_power_of_2(self.head_dim),
             BLOCK_DV=triton.next_power_of_2(self.v_head_dim),
             num_warps=1,

python/sglang/kernels/ops/attention/vision_rope.py

@@ -0,0 +1,217 @@
+"""Fused interleaved complex RoPE for vision attention Q/K tensors."""
+
+from __future__ import annotations
+
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+PreparedInplaceComplexRoPE = Tuple[torch.Tensor, torch.Tensor]
+
+
+@triton.jit(do_not_specialize=["n_pairs"])
+def _fused_qk_complex_rope_kernel(
+    q_ptr,
+    k_ptr,
+    freqs_ptr,
+    q_out_ptr,
+    k_out_ptr,
+    n_pairs,
+    n_heads: tl.constexpr,
+    head_dim: tl.constexpr,
+    q_stride_token: tl.constexpr,
+    q_stride_head: tl.constexpr,
+    q_stride_dim: tl.constexpr,
+    k_stride_token: tl.constexpr,
+    k_stride_head: tl.constexpr,
+    k_stride_dim: tl.constexpr,
+    freq_stride_token: tl.constexpr,
+    freq_stride_pair: tl.constexpr,
+    freq_stride_complex: tl.constexpr,
+    BLOCK: tl.constexpr,
+) -> None:
+    pair_offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = pair_offsets < n_pairs
+    pairs_per_row = head_dim // 2
+    row = pair_offsets // pairs_per_row
+    pair = pair_offsets - row * pairs_per_row
+    token = row // n_heads
+    head = row - token * n_heads
+
+    q_base = token * q_stride_token + head * q_stride_head + pair * 2 * q_stride_dim
+    k_base = token * k_stride_token + head * k_stride_head + pair * 2 * k_stride_dim
+    freq_base = token * freq_stride_token + pair * freq_stride_pair
+    cos = tl.load(freqs_ptr + freq_base, mask=mask).to(tl.float32)
+    sin = tl.load(freqs_ptr + freq_base + freq_stride_complex, mask=mask).to(tl.float32)
+    q_real = tl.load(q_ptr + q_base, mask=mask).to(tl.float32)
+    q_imag = tl.load(q_ptr + q_base + q_stride_dim, mask=mask).to(tl.float32)
+    k_real = tl.load(k_ptr + k_base, mask=mask).to(tl.float32)
+    k_imag = tl.load(k_ptr + k_base + k_stride_dim, mask=mask).to(tl.float32)
+
+    out_base = row * head_dim + pair * 2
+    tl.store(
+        q_out_ptr + out_base,
+        tl.fma(-q_imag, sin, q_real * cos),
+        mask=mask,
+    )
+    tl.store(
+        q_out_ptr + out_base + 1,
+        tl.fma(q_real, sin, q_imag * cos),
+        mask=mask,
+    )
+    tl.store(
+        k_out_ptr + out_base,
+        tl.fma(-k_imag, sin, k_real * cos),
+        mask=mask,
+    )
+    tl.store(
+        k_out_ptr + out_base + 1,
+        tl.fma(k_real, sin, k_imag * cos),
+        mask=mask,
+    )
+
+
+def can_use_fused_qk_complex_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> bool:
+    """Whether the NVIDIA fused path supports these vision RoPE tensors."""
+
+    if not (
+        q.is_cuda
+        and k.is_cuda
+        and freqs_cis.is_cuda
+        and q.device == k.device == freqs_cis.device
+    ):
+        return False
+    if q.dtype != k.dtype or q.dtype not in (torch.bfloat16, torch.float16):
+        return False
+    if freqs_cis.dtype != torch.complex64 or q.shape != k.shape or q.ndim < 3:
+        return False
+    if q.shape[-1] % 2 != 0:
+        return False
+    if freqs_cis.shape != q.shape[:-2] + (q.shape[-1] // 2,):
+        return False
+    major, _ = torch.cuda.get_device_capability(q.device)
+    return major >= 9
+
+
+def apply_fused_qk_complex_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Rotate interleaved Q/K pairs with one kernel.
+
+    ``q`` and ``k`` may be strided views of an interleaved QKV projection. The
+    output is contiguous, matching the native complex-multiply implementation.
+    The token count remains a runtime kernel argument so random image sizes do
+    not create new Triton specializations.
+    """
+
+    if not can_use_fused_qk_complex_rope(q, k, freqs_cis):
+        raise ValueError(
+            "Unsupported fused vision RoPE inputs: "
+            f"q={q.shape}/{q.dtype}/{q.device}, "
+            f"k={k.shape}/{k.dtype}/{k.device}, "
+            f"freqs={freqs_cis.shape}/{freqs_cis.dtype}/{freqs_cis.device}"
+        )
+
+    original_shape = q.shape
+    # Preserve the interleaved QKV token stride when the token dimension is 1.
+    # ``view(-1, ...)`` is otherwise free to collapse that singleton stride,
+    # producing a different Triton specialization from real image requests.
+    q_flat = q if q.ndim == 3 else q.view(-1, q.shape[-2], q.shape[-1])
+    k_flat = k if k.ndim == 3 else k.view(-1, k.shape[-2], k.shape[-1])
+    freqs = torch.view_as_real(freqs_cis).view(-1, q.shape[-1] // 2, 2)
+    q_out = torch.empty(q_flat.shape, dtype=q.dtype, device=q.device)
+    k_out = torch.empty(k_flat.shape, dtype=k.dtype, device=k.device)
+
+    block = 128
+    n_pairs = q_flat.numel() // 2
+    _fused_qk_complex_rope_kernel[(triton.cdiv(n_pairs, block),)](
+        q_flat,
+        k_flat,
+        freqs,
+        q_out,
+        k_out,
+        n_pairs,
+        q_flat.shape[1],
+        q_flat.shape[2],
+        q_flat.stride(0),
+        q_flat.stride(1),
+        q_flat.stride(2),
+        k_flat.stride(0),
+        k_flat.stride(1),
+        k_flat.stride(2),
+        freqs.stride(0),
+        freqs.stride(1),
+        freqs.stride(2),
+        BLOCK=block,
+        num_warps=4,
+    )
+    return q_out.view(original_shape), k_out.view(original_shape)
+
+
+def prepare_fused_qk_complex_rope_inplace(
+    freqs_cis: torch.Tensor,
+) -> PreparedInplaceComplexRoPE:
+    """Prepare the cache and positions used by the contiguous in-place kernel."""
+
+    if freqs_cis.dtype != torch.complex64:
+        raise ValueError(
+            "In-place vision RoPE requires complex64 frequencies, got "
+            f"{freqs_cis.dtype}/{freqs_cis.device}"
+        )
+    return (
+        torch.cat((freqs_cis.real, freqs_cis.imag), dim=-1),
+        torch.arange(
+            freqs_cis.size(0),
+            dtype=torch.long,
+            device=freqs_cis.device,
+        ),
+    )
+
+
+def apply_fused_qk_complex_rope_inplace(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    prepared_rope: PreparedInplaceComplexRoPE,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Rotate contiguous Q/K tensors in place with the shared JIT kernel."""
+
+    from sglang.kernels.ops.attention.rope import apply_rope_inplace
+
+    cos_sin_cache, positions = prepared_rope
+    apply_rope_inplace(
+        q,
+        k,
+        cos_sin_cache,
+        positions,
+        is_neox=False,
+        rope_dim=cos_sin_cache.size(-1),
+    )
+    return q, k
+
+
+def precompile_fused_qk_complex_rope(
+    *,
+    num_heads: int,
+    head_dim: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> bool:
+    """Compile the dynamic-token QKV-view specialization before serving."""
+
+    if device.type != "cuda" or dtype not in (torch.bfloat16, torch.float16):
+        return False
+    qkv = torch.empty((1, 3, num_heads, head_dim), dtype=dtype, device=device)
+    q, k, _ = torch.unbind(qkv, dim=1)
+    freqs = torch.ones((1, head_dim // 2), dtype=torch.complex64, device=device)
+    if not can_use_fused_qk_complex_rope(q, k, freqs):
+        return False
+    apply_fused_qk_complex_rope(q, k, freqs)
+    return True

python/sglang/kernels/ops/elementwise/add3.py

@@ -0,0 +1,67 @@
+"""CUDA JIT elementwise 3-way add: out = bf16(bf16(a + b) + c)."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+# The kernel vectorizes by device::kMaxVecBytes (16B pre-Blackwell, 32B on
+# Blackwell+). Requiring divisibility by the widest case keeps covered()
+# arch-independent and never looser than the compiled kernel's check.
+_MAX_VEC_ELEMS: int = 16
+
+
+@cache_once
+def _jit_add3_module() -> Module:
+    args = make_cpp_args(is_arch_support_pdl())
+    return load_jit(
+        "add3_bf16",
+        *args,
+        cuda_files=["elementwise/add3.cuh"],
+        cuda_wrappers=[("run", f"sglang::Add3Kernel<{args}>::launch")],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+    )
+
+
+def covered(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> bool:
+    """Same-shape contiguous CUDA bf16 tensors, numel a multiple of the
+    widest vector (16 elements)."""
+    return (
+        a.dtype == b.dtype == c.dtype == torch.bfloat16
+        and a.shape == b.shape == c.shape
+        and a.is_contiguous()
+        and b.is_contiguous()
+        and c.is_contiguous()
+        and a.is_cuda
+        and a.numel() > 0
+        and a.numel() % _MAX_VEC_ELEMS == 0
+    )
+
+
+def add3(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    c: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+    prefetch_bc: bool = False,
+) -> torch.Tensor:
+    """out = bf16(bf16(a + b) + c); double rounding matches the unfused add
+    pair bit-for-bit. With prefetch_bc, b/c are loaded before the PDL wait —
+    only safe when their producers are at least two kernels back."""
+    if out is None:
+        out = torch.empty_like(a)
+    module = _jit_add3_module()
+    module.run(a.view(-1), b.view(-1), c.view(-1), out.view(-1), prefetch_bc)
+    return out

python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py

@@ -1294,12 +1294,39 @@ def _pick_tactic(m: int, n: int, k: int) -> int:
     return best
 
 
+# Kimi-K3 per-rank dense-GEMM shapes (TP8). They sit in this heuristic's
+# unmeasured region (hidden=7168 inputs fail k > 6144; o_proj-style k=1536
+# fails k < 2048), but TGV wins 1.04-2.43x on every one of them on GB300
+# (L2-defeating weight rotation + CUDA-graph timing; serving A/B confirmed
+# e2e with GSM8K parity). kv_a (n=576) loses (0.84x) and stays out. Gated to
+# small decode batches; larger m stays on the measured heuristic below.
+_K3_TGV_WIN_SHAPES = frozenset(
+    {
+        (6144, 7168),  # KDA fused qkvg
+        (6016, 7168),  # merged MoE front (gate_up | router | latent down)
+        (7168, 1536),  # KDA / MLA o_proj
+        (1536, 7168),  # MLA q_a / shared gate_up
+        (2304, 1536),  # MLA q_b
+        (3584, 7168),  # MoE latent down (unfused fallback)
+        (7168, 3584),  # MoE latent up
+        (7168, 768),  # shared down
+    }
+)
+
+
 def use_cutedsl_bf16_gemm(m: int, n: int, k: int) -> bool:
     """TGV-vs-cuBLAS (``F.linear``) decision, CUPTI-measured on B300 under CUDA
     graph capture (cold L2). Conservative: ties and unmeasured regions fall
     back to cuBLAS."""
+    if m <= 0:
+        # Empty batch: DP-attention idle groups run a 0-token dummy forward to
+        # keep the mlp-sync lockstep. A 0-CTA TGV grid is a driver-level
+        # CUDA_ERROR_INVALID_VALUE, so leave empty inputs to cuBLAS.
+        return False
     if k % 8 != 0:  # TMA requires 16B-aligned rows
         return False
+    if m <= 8 and (n, k) in _K3_TGV_WIN_SHAPES:
+        return True
     if n < 1024 or k < 2048 or k > 6144:
         return False
     ragged = m % 16 != 0
@@ -1333,6 +1360,10 @@ def _tgv_bf16_gemm_run(
     out = torch.empty(
         (x.shape[0], weight.shape[0]), dtype=torch.bfloat16, device=x.device
     )
+    if x.shape[0] == 0:
+        # Match cuBLAS/F.linear semantics for empty batches; a 0-CTA launch
+        # would fail with CUDA_ERROR_INVALID_VALUE.
+        return out
     return _run_tgv(
         x,
         weight.t(),
@@ -1343,6 +1374,33 @@ def _tgv_bf16_gemm_run(
     )
 
 
+def _tgv_bf16_gemm_out_run(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    out: torch.Tensor,
+    bias: Optional[torch.Tensor],
+) -> None:
+    if get_device_sm() not in (100, 103):
+        raise RuntimeError("cutedsl_bf16_gemm requires SM100/SM103 (Blackwell)")
+    assert x.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16
+    assert out.dtype == torch.bfloat16 and out.device == x.device
+    assert x.ndim == 2 and weight.ndim == 2 and out.ndim == 2
+    assert x.stride(-1) == 1, "x must be K-major [M, K]"
+    assert weight.stride(-1) == 1, "weight must be K-major [N, K]"
+    assert out.is_contiguous() and out.shape == (x.shape[0], weight.shape[0])
+    if x.shape[0] == 0:
+        return None
+    _run_tgv(
+        x,
+        weight.t(),
+        bias,
+        out,
+        pdl=True,
+        tactic=_pick_tactic(x.shape[0], weight.shape[0], weight.shape[1]),
+    )
+    return None
+
+
 def _tgv_bf16_gemm_fake(
     x: torch.Tensor,
     weight: torch.Tensor,
@@ -1351,6 +1409,15 @@ def _tgv_bf16_gemm_fake(
     return x.new_empty((x.shape[0], weight.shape[0]))
 
 
+def _tgv_bf16_gemm_out_fake(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    out: torch.Tensor,
+    bias: Optional[torch.Tensor],
+) -> None:
+    return None
+
+
 direct_register_custom_op(
     op_name="cutedsl_tgv_bf16_gemm",
     op_func=_tgv_bf16_gemm_run,
@@ -1358,6 +1425,13 @@ direct_register_custom_op(
     fake_impl=_tgv_bf16_gemm_fake,
 )
 
+direct_register_custom_op(
+    op_name="cutedsl_tgv_bf16_gemm_out",
+    op_func=_tgv_bf16_gemm_out_run,
+    mutates_args=["out"],
+    fake_impl=_tgv_bf16_gemm_out_fake,
+)
+
 
 @debug_kernel_api
 def cutedsl_bf16_gemm(
@@ -1367,3 +1441,15 @@ def cutedsl_bf16_gemm(
 ) -> torch.Tensor:
     """out[M, N] = x[M, K] @ weight[N, K].T (+ bias[N]), all bf16, fp32 accum."""
     return torch.ops.sglang.cutedsl_tgv_bf16_gemm(x, weight, bias)
+
+
+@debug_kernel_api
+def cutedsl_bf16_gemm_out(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    out: torch.Tensor,
+    bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Write the BF16 GEMM directly into a caller-owned contiguous tensor."""
+    torch.ops.sglang.cutedsl_tgv_bf16_gemm_out(x, weight, out, bias)
+    return out

python/sglang/kernels/ops/gemm/tiny_gemm.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+_MAX_M_DEFAULT: int = 16
+
+
+@cache_once
+def _jit_tiny_gemm_module(
+    n: int, k: int, max_m: int, split_n: int, out_dtype: torch.dtype
+) -> Module:
+    args = make_cpp_args(n, k, max_m, split_n, out_dtype, is_arch_support_pdl())
+    return load_jit(
+        "tiny_gemm",
+        *args,
+        cuda_files=["gemm/tiny_gemm.cuh"],
+        cuda_wrappers=[("run", f"TinyNGemmKernel<{args}>::run")],
+        extra_cuda_cflags=["-O3"],
+    )
+
+
+@cache_once
+def _jit_tiny_k_gemm_module(
+    n: int, k: int, max_m: int, n_unroll: int, out_dtype: torch.dtype
+) -> Module:
+    args = make_cpp_args(n, k, max_m, n_unroll, out_dtype, is_arch_support_pdl())
+    return load_jit(
+        "tiny_k_gemm",
+        *args,
+        cuda_files=["gemm/tiny_gemm.cuh"],
+        cuda_wrappers=[("run", f"TinyKGemmKernel<{args}>::run")],
+        extra_cuda_cflags=["-O3"],
+    )
+
+
+def _vec_elems() -> int:
+    """bf16 elements per vectorized load; mirrors kMaxVecBytes in utils.cuh."""
+    from sglang.kernels.jit.utils import get_jit_cuda_arch
+
+    cuda = tuple(int(v) for v in (torch.version.cuda or "0.0").split(".")[:2])
+    return 16 if get_jit_cuda_arch().major >= 10 and cuda >= (12, 9) else 8
+
+
+def _default_split_n(n: int, k: int, max_m: int, device: torch.device) -> int:
+    """Smallest divisor of n whose n / split_n blocks fit in one wave, subject
+    to the max_m * split_n <= K / vec_elems block-size constraint; falls back
+    to the largest split_n satisfying the constraint (multi-wave grid)."""
+    sm_count = torch.cuda.get_device_properties(device).multi_processor_count
+    split_cap = (k // _vec_elems()) // max_m
+    divisors = [d for d in range(1, min(n, split_cap) + 1) if n % d == 0]
+    if not divisors:
+        raise RuntimeError(
+            f"tiny_gemm: no valid split_n for N={n}, K={k}, max_m={max_m};"
+            " lower max_m"
+        )
+    for split in divisors:
+        if n // split <= sm_count:
+            return split
+    return divisors[-1]
+
+
+def tiny_n_gemm_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    *,
+    out_dtype: Optional[torch.dtype] = None,
+    split_n: Optional[int] = None,
+    max_m: int = _MAX_M_DEFAULT,
+) -> torch.Tensor:
+    n = w.shape[0]
+    k = x.shape[1]
+    if out is None:
+        out_dtype = out_dtype or torch.bfloat16
+        out = torch.empty((x.shape[0], n), dtype=out_dtype, device=x.device)
+    else:
+        assert out_dtype is None or out_dtype == out.dtype
+    if split_n is None:
+        split_n = _default_split_n(n, k, max_m, x.device)
+    module = _jit_tiny_gemm_module(n, k, max_m, split_n, out.dtype)
+    module.run(x, w, out)
+    return out
+
+
+def _default_k_split_n(n: int, k: int) -> int:
+    """Smallest divisor of n whose n / split_n blocks fit one wave, with
+    split_n * K-lanes whole-warp aligned and within the block-size limit."""
+    lanes = k // 8  # fixed 16-byte vectors in the K variant
+    candidates = [
+        d
+        for d in range(1, n + 1)
+        if n % d == 0 and d * lanes % 32 == 0 and d * lanes <= 1024
+    ]
+    if not candidates:
+        raise RuntimeError(f"tiny_k_gemm: no valid split_n for N={n}, K={k}")
+    sm_count = torch.cuda.get_device_properties(0).multi_processor_count
+    for d in candidates:
+        if n // d <= sm_count:
+            return d
+    return candidates[-1]
+
+
+def tiny_k_gemm_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    *,
+    out_dtype: Optional[torch.dtype] = None,
+    split_n: Optional[int] = None,
+    max_m: int = _MAX_M_DEFAULT,
+) -> torch.Tensor:
+    """Small-K / large-N variant: K / 8 lanes of one warp reduce the K
+    dimension for one output column; each block covers split_n columns and the
+    exact N / split_n grid fills the SMs (no tail). Requires K / 8 to be a
+    power of 2 and <= 32 (e.g. K = 128/256). x may be a row-sliced view as
+    long as rows stay 16-byte aligned.
+
+    split_n trades block count for block size; the default picks the smallest
+    divisor of N that fits one wave (12 for [1536, 128] on B200: 128 blocks
+    of 6 warps)."""
+    n = w.shape[0]
+    k = x.shape[1]
+    if out is None:
+        out_dtype = out_dtype or torch.bfloat16
+        out = torch.empty((x.shape[0], n), dtype=out_dtype, device=x.device)
+    else:
+        assert out_dtype is None or out_dtype == out.dtype
+    if split_n is None:
+        split_n = _default_k_split_n(n, k)
+    module = _jit_tiny_k_gemm_module(n, k, max_m, split_n, out.dtype)
+    module.run(x, w, out)
+    return out

python/sglang/kernels/ops/kimi_k3/__init__.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    import torch
+
+_K3_N_GEMM_DISPATCH_MAP = {
+    (144, 7168): 16,
+    (896, 7168): 8,
+}
+_K3_K_GEMM_DISPATCH_MAP = {
+    (1536, 128): 12,
+}
+
+
+def situ_and_mul(
+    input: torch.Tensor,
+    out: Optional[torch.Tensor],
+    beta: float,
+    linear_beta: Optional[float],
+) -> torch.Tensor:
+    from .activation import situ_and_mul as impl
+
+    return impl(input, out, beta, linear_beta)
+
+
+def situ_and_mul_masked_post_quant(
+    input: torch.Tensor,
+    output: torch.Tensor,
+    output_scale: torch.Tensor,
+    quant_group_size: int,
+    masked_m: torch.Tensor,
+    beta: float,
+    linear_beta: float,
+    scale_ue8m0: bool = False,
+    topk: int = 8,
+    transposed: bool = False,
+    swizzle: bool = False,
+) -> None:
+    from .moe import situ_and_mul_masked_post_quant as impl
+
+    return impl(
+        input,
+        output,
+        output_scale,
+        quant_group_size,
+        masked_m,
+        beta,
+        linear_beta,
+        scale_ue8m0,
+        topk,
+        transposed,
+        swizzle,
+    )
+
+
+def kimi_k3_tiny_gemm(
+    x: torch.Tensor,
+    w: torch.Tensor,
+) -> torch.Tensor:
+    import torch
+
+    from ..gemm.tiny_gemm import tiny_k_gemm_bf16, tiny_n_gemm_bf16
+
+    m, k = x.shape
+    n, _ = w.shape
+    if max_num_tokens := _K3_N_GEMM_DISPATCH_MAP.get((n, k)):
+        if 0 < m <= max_num_tokens:
+            return tiny_n_gemm_bf16(x, w)
+    if max_num_tokens := _K3_K_GEMM_DISPATCH_MAP.get((n, k)):
+        if 0 < m <= max_num_tokens:
+            return tiny_k_gemm_bf16(x, w)
+    return torch.nn.functional.linear(x, w)
+
+
+__all__ = [
+    "situ_and_mul",
+    "situ_and_mul_masked_post_quant",
+    "kimi_k3_tiny_gemm",
+]

python/sglang/kernels/ops/kimi_k3/activation.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    get_jit_cuda_arch,
+    is_arch_support_pdl,
+    is_hip_runtime,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+def _make_name(*args):
+    return "kimi_k3_" + "_".join(str(a) for a in args)
+
+
+def _fast_math_flags() -> list[str]:
+    # Mirrors sgl-kernel's CMake policy: fast-math on SM90, precise on
+    # SM100+ (Blackwell needs bit-exact expf), off on HIP (clang rejects).
+    if is_hip_runtime():
+        return []
+    if get_jit_cuda_arch().major >= 10:
+        return []
+    return ["--use_fast_math"]
+
+
+@cache_once
+def _jit_situ_and_mul_module(dtype: torch.dtype) -> Module:
+    """Compile and cache the JIT SiTU-and-mul module for a given dtype."""
+    args = make_cpp_args(dtype, is_arch_support_pdl())
+    return load_jit(
+        _make_name("situ_and_mul"),
+        *args,
+        cuda_files=["kimi_k3/situ_and_mul.cuh"],
+        cuda_wrappers=[("run", f"SituAndMulKernel<{args}>::run")],
+        extra_cuda_cflags=_fast_math_flags(),
+    )
+
+
+def situ_and_mul(
+    input: torch.Tensor,
+    out: Optional[torch.Tensor],
+    beta: float,
+    linear_beta: Optional[float],
+) -> torch.Tensor:
+    """Fused SiTU (SoftCap-GLU) activation: bf16 -> bf16.
+
+    gate_out = beta * tanh(gate / beta) * sigmoid(gate)
+    up_out   = linear_beta * tanh(up / linear_beta)  [if linear_beta is not None]
+    output   = gate_out * up_out
+
+    Parameters
+    ----------
+    input       : bf16 CUDA tensor [*, 2*D]
+    out         : optional pre-allocated bf16 CUDA tensor [*, D]
+    beta        : gate softcap scalar (e.g. 4.0)
+    linear_beta : up softcap scalar (e.g. 25.0), or None to skip
+    """
+    hidden_size = input.shape[-1] // 2
+    if out is None:
+        out = input.new_empty(*input.shape[:-1], hidden_size)
+
+    # 2D inputs may be row-strided (e.g. a slice of a fused-GEMM output);
+    # higher-rank inputs keep the dense-view path.
+    if input.dim() == 2 and input.stride(1) == 1:
+        input_2d = input
+    else:
+        input_2d = input.contiguous().view(-1, hidden_size * 2)
+    out_2d = out.view(-1, hidden_size)
+
+    has_linear_beta = linear_beta is not None
+    module = _jit_situ_and_mul_module(input.dtype)
+    module.run(
+        input_2d,
+        out_2d,
+        float(beta),
+        float(linear_beta) if has_linear_beta else 0.0,
+        has_linear_beta,
+    )
+    return out

python/sglang/kernels/ops/kimi_k3/all_reduce.py

@@ -0,0 +1,380 @@
+"""K3 MNNVL fused all-reduce (bf16): zero-copy AR and AR+RMSNorm.
+
+Four entry points over ``csrc/kimi_k3/comm/ar_fusion.cuh``, spanning two
+algorithm families x two epilogues:
+
+============  =========================  ==================================
+              res (+ optional residual)  norm (fused RMSNorm on the latent)
+============  =========================  ==================================
+push (1shot)  :func:`all_reduce_push_res`   :func:`all_reduce_push_norm`
+pull (2shot)  :func:`all_reduce_pull_res`   :func:`all_reduce_pull_norm`
+============  =========================  ==================================
+
+* **push** — 1shot multicast-push. Works on ANY contiguous bf16 tensor
+  (input is read and written in place); reuses the CustomAllReduceV2 push
+  workspace, so the caller passes the workspace slab's multicast base.
+  Best for small messages. Needs :func:`register_comm`.
+* **pull** — low-SM NVLS 2shot ON the input, which must be allocated from
+  multicast-bound symmetric memory (the caller passes its multicast VA):
+  reduce-scatter + broadcast in place.
+  Launch geometry defaults to :data:`RES_TUNING` /
+  :data:`NORM_TUNING` and can be overridden per call via ``num_blocks`` /
+  ``unroll`` (``num_blocks`` must be uniform across ranks per call).
+  Barriers reuse the CustomAllReduceV2 pull semaphores (same reservation
+  protocol as the generic pull kernels, signaled via one multicast red), so
+  :func:`register_comm` additionally needs the semaphore region's multicast
+  VA (``CustomAllReduceV2.pull_sem_mc_ptr``).
+
+Epilogue contracts: the ``res`` residual must be identical on every rank (a
+fully reduced tensor such as the attn-res prefix sum) or absent; the
+``norm`` input is the K3 latent|shared MoE buffer ([N, 3584] latent then
+[N, 7168] shared, contiguous — the row layout is derived and hardcoded
+C++-side).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, NamedTuple, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+    from sglang.kernels.ops.communication.all_reduce import Communicator
+
+
+@cache_once
+def _jit_module(world_size: int) -> Module:
+    args = make_cpp_args(world_size, is_arch_support_pdl())
+    cls = f"AllReduceFusionKernel<{args}>"
+    return load_jit(
+        "kimi_k3_all_reduce",
+        *args,
+        cuda_files=["kimi_k3/comm/ar_fusion.cuh"],
+        cuda_wrappers=[
+            ("push_res", f"{cls}::push_res"),
+            ("push_norm", f"{cls}::push_norm"),
+            ("pull_res", f"{cls}::pull_res"),
+            ("pull_norm", f"{cls}::pull_norm"),
+            ("finalize_push_norm", f"{cls}::finalize_push_norm"),
+        ],
+        extra_cuda_cflags=["-O3"],
+    )
+
+
+# Storage plane: the CustomAllReduceV2 Communicator
+
+
+class _CommEntry(NamedTuple):
+    obj: Communicator  # sgl.Communicator
+    pull_sem_mc_ptr: int
+
+
+_COMM_MAP: dict[int, _CommEntry] = {}
+
+
+def register_comm(comm: Communicator, *, pull_sem_mc_ptr: int = 0) -> None:
+    """Register the CustomAllReduceV2 storage plane.
+
+    The push kernels only need ``comm``; the pull kernels additionally need
+    ``pull_sem_mc_ptr`` (``CustomAllReduceV2.pull_sem_mc_ptr``), the
+    multicast VA of the pull-semaphore region their barriers reuse.
+    """
+    # world_size is the whole key, so at most one communicator per size can be
+    # registered in a process. That matches how these ops are called -- the custom
+    # ops below take world_size and nothing else, so a second group of the same
+    # size would silently inherit the first one's peer pointers and semaphores,
+    # and the symptom would be a hang or corruption rather than an error. Assert
+    # it instead of letting the overwrite happen; widening to per-group handles
+    # means changing the custom-op signatures, which is a separate change.
+    prev = _COMM_MAP.get(comm.world_size)
+    assert prev is None or prev.obj is comm, (
+        f"a different communicator is already registered for world_size="
+        f"{comm.world_size}; these ops key only on world_size, so two groups of "
+        f"the same size cannot coexist in one process"
+    )
+    _COMM_MAP[comm.world_size] = _CommEntry(obj=comm, pull_sem_mc_ptr=pull_sem_mc_ptr)
+
+
+class PullTuning(NamedTuple):
+    num_blocks: int
+    unroll: int  # 2, 4, 8, or 16 (every width is compiled into the module)
+
+
+class PullTuningTable(NamedTuple):
+    bands: tuple[tuple[int, PullTuning], ...]
+    fallback: PullTuning
+
+    def lookup(self, nbytes: int) -> PullTuning:
+        for max_bytes, tuning in self.bands:
+            if nbytes <= max_bytes:
+                return tuning
+        return self.fallback
+
+
+_KB, _MB = 1024, 1024 * 1024
+
+RES_TUNING = PullTuningTable(
+    bands=(
+        (128 * _KB, PullTuning(num_blocks=1, unroll=8)),
+        (4 * _MB, PullTuning(num_blocks=16, unroll=8)),
+    ),
+    fallback=PullTuning(num_blocks=4, unroll=8),
+)
+
+NORM_TUNING = PullTuningTable(
+    bands=(
+        (512 * _KB, PullTuning(num_blocks=2, unroll=4)),
+        (2 * _MB, PullTuning(num_blocks=12, unroll=2)),
+    ),
+    fallback=PullTuning(num_blocks=24, unroll=2),
+)
+
+
+def _resolve_tuning(
+    table: PullTuningTable,
+    *,
+    nbytes: int,
+    num_blocks: Optional[int],
+    unroll: Optional[int],
+) -> PullTuning:
+    """Per-size tuned config, with explicit overrides taking precedence
+    (C++-side, the block count is clamped to the semaphore capacity)."""
+    tuned = table.lookup(nbytes)
+    return PullTuning(
+        num_blocks=num_blocks or tuned.num_blocks,
+        unroll=unroll or tuned.unroll,
+    )
+
+
+# Custom ops, one per C++ entry point
+
+
+@register_custom_op(mutates_args=["x"])
+def _push_res_op(
+    world_size: int,
+    x: torch.Tensor,
+    residual: Optional[torch.Tensor],
+    ws_mc_base: int,
+) -> None:
+    comm = _COMM_MAP[world_size].obj
+    _jit_module(world_size).push_res(comm, x.view(-1), residual, ws_mc_base)
+
+
+@register_custom_op(mutates_args=["x"])
+def _push_norm_op(
+    world_size: int,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    num_norm_rows: int,
+    ws_mc_base: int,
+) -> None:
+    comm = _COMM_MAP[world_size].obj
+    _jit_module(world_size).push_norm(
+        comm, x.view(-1), weight, eps, num_norm_rows, ws_mc_base
+    )
+
+
+@register_custom_op(mutates_args=["out"])
+def _finalize_push_norm_op(
+    world_size: int,
+    out: torch.Tensor,
+    gemm2_out: torch.Tensor,
+    expanded_idx_to_permuted_idx: torch.Tensor,
+    expert_weights: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    ws_mc_base: int,
+) -> None:
+    comm = _COMM_MAP[world_size].obj
+    _jit_module(world_size).finalize_push_norm(
+        comm,
+        out.view(-1),
+        gemm2_out,
+        expanded_idx_to_permuted_idx,
+        expert_weights,
+        weight,
+        eps,
+        ws_mc_base,
+    )
+
+
+@register_custom_op(mutates_args=["x"])
+def _pull_res_op(
+    world_size: int,
+    x: torch.Tensor,
+    residual: Optional[torch.Tensor],
+    input_mc_ptr: int,
+    num_blocks: int,
+    unroll: int,
+) -> None:
+    entry = _COMM_MAP[world_size]
+    _jit_module(world_size).pull_res(
+        entry.obj,
+        x.view(-1),
+        residual,
+        input_mc_ptr,
+        entry.pull_sem_mc_ptr,
+        num_blocks,
+        unroll,
+    )
+
+
+@register_custom_op(mutates_args=["x"])
+def _pull_norm_op(
+    world_size: int,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    num_norm_rows: int,
+    input_mc_ptr: int,
+    num_blocks: int,
+    unroll: int,
+) -> None:
+    entry = _COMM_MAP[world_size]
+    _jit_module(world_size).pull_norm(
+        entry.obj,
+        x.view(-1),
+        weight,
+        eps,
+        num_norm_rows,
+        input_mc_ptr,
+        entry.pull_sem_mc_ptr,
+        num_blocks,
+        unroll,
+    )
+
+
+def all_reduce_push_res(
+    world_size: int,
+    x: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    *,
+    ws_mc_base: int,
+) -> torch.Tensor:
+    """In-place ``x = allreduce(x) [+ residual]`` via 1shot multicast push.
+
+    ``x`` may be any contiguous bf16 CUDA tensor whose byte size fits the
+    registered push workspace. ``ws_mc_base`` is the multicast VA of the v2
+    workspace slab base. Call :func:`register_comm` once beforehand.
+    """
+    residual_ = residual.view(-1) if residual is not None else None
+    _push_res_op(world_size, x, residual_, ws_mc_base)
+    return x
+
+
+def all_reduce_push_norm(
+    world_size: int,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    *,
+    num_norm_rows: int,
+    ws_mc_base: int,
+) -> torch.Tensor:
+    """In-place allreduce via 1shot multicast push + RMSNorm over the first
+    ``num_norm_rows`` rows of ``x`` viewed as [numel / 3584, 3584]."""
+    _push_norm_op(world_size, x, weight, eps, num_norm_rows, ws_mc_base)
+    return x
+
+
+def finalize_all_reduce_push_norm(
+    world_size: int,
+    out: torch.Tensor,
+    gemm2_out: torch.Tensor,
+    expanded_idx_to_permuted_idx: torch.Tensor,
+    expert_weights: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    *,
+    ws_mc_base: int,
+) -> torch.Tensor:
+    """Deferred MoE finalize + 1shot push all-reduce + RMSNorm on EVERY row.
+
+    ``out`` ([T, 3584] bf16) is output-only; each rank's partial latent
+    (``sum_k expert_weights[t, k] * gemm2_out[idx[t*16 + k]]``, -1 slots
+    skipped) is computed during the multicast staging pass from the
+    trtllm-gen deferred-finalize triple (``do_finalize=False``) and never
+    materializes in global memory. top_k is fixed to 16 (K3)."""
+    _finalize_push_norm_op(
+        world_size,
+        out,
+        gemm2_out,
+        expanded_idx_to_permuted_idx,
+        expert_weights,
+        weight,
+        eps,
+        ws_mc_base,
+    )
+    return out
+
+
+def all_reduce_pull_res(
+    world_size: int,
+    x: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    *,
+    input_mc_ptr: int,
+    num_blocks: Optional[int] = None,
+    unroll: Optional[int] = None,
+) -> torch.Tensor:
+    """In-place ``x = allreduce(x) [+ residual]`` via low-SM NVLS 2shot.
+
+    ``x`` MUST be allocated from multicast-bound symmetric memory and
+    ``input_mc_ptr`` must be its multicast VA. Call :func:`register_comm`
+    (with ``pull_sem_mc_ptr``) once beforehand.
+    """
+    tuning = _resolve_tuning(
+        RES_TUNING,
+        nbytes=x.numel() * x.element_size(),
+        num_blocks=num_blocks,
+        unroll=unroll,
+    )
+    residual_ = residual.view(-1) if residual is not None else None
+    _pull_res_op(
+        world_size, x, residual_, input_mc_ptr, tuning.num_blocks, tuning.unroll
+    )
+    return x
+
+
+def all_reduce_pull_norm(
+    world_size: int,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    num_norm_rows: int,
+    input_mc_ptr: int,
+    num_blocks: Optional[int] = None,
+    unroll: Optional[int] = None,
+) -> torch.Tensor:
+    """In-place allreduce via low-SM NVLS 2shot + RMSNorm over the first
+    ``num_norm_rows`` rows of ``x`` viewed as [numel / 3584, 3584]; ``x``
+    must live in multicast-bound symmetric memory."""
+    tuning = _resolve_tuning(
+        NORM_TUNING,
+        nbytes=x.nbytes,
+        num_blocks=num_blocks,
+        unroll=unroll,
+    )
+    _pull_norm_op(
+        world_size,
+        x,
+        weight,
+        eps,
+        num_norm_rows,
+        input_mc_ptr,
+        tuning.num_blocks,
+        tuning.unroll,
+    )
+    return x

python/sglang/kernels/ops/kimi_k3/attn_res.py

@@ -0,0 +1,301 @@
+"""CUDA JIT wrapper for the Kimi-K3 SM100 attention-residual kernel."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    load_jit,
+    make_cpp_args,
+    override_jit_cuda_arch,
+)
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+    from sglang.kernels.ops.communication.all_reduce import Communicator
+
+_DIM: int = 7168  # K3 hidden size, template parameter of the TMA kernel
+_MAX_BANK_ROWS: int = 8  # K3 has <= 8 snapshots, upper bound of the nvb dispatch tables
+
+
+def _make_name(*args):
+    return "kimi_k3_attn_res_" + "_".join(str(a) for a in args)
+
+
+@cache_once
+def _jit_fused_tma_module(
+    chunk_rows: int, occupancy: int, consumer_regs: int
+) -> Module:
+    """Compile and cache the warp-specialized TMA aggregation kernel (per-row
+    bulk copies into chunk slots; chunk_rows / occupancy / consumer_regs are
+    tuning knobs). The smem ring is frozen at 2 chunk slots and PDL is always
+    on: the kernel targets SM100+, where both are unconditional wins."""
+    major, minor = torch.cuda.get_device_capability()
+    if major < 10:
+        raise RuntimeError(
+            "attn_res_fused_tma requires SM100+ (tcgen05, cp.async.bulk)"
+        )
+    args = make_cpp_args(
+        _DIM,
+        _MAX_BANK_ROWS,
+        chunk_rows,
+        occupancy,
+        consumer_regs,
+    )
+    with override_jit_cuda_arch(major, minor, suffix="a"):
+        return load_jit(
+            _make_name("fused_tma"),
+            *args,
+            cuda_files=["kimi_k3/attn_res/fused_tma.cuh"],
+            cuda_wrappers=[
+                ("run", f"AttnResFusedTmaKernel<{args}>::run"),
+                ("run_pull_rs", f"AttnResFusedTmaKernel<{args}>::run_pull_rs"),
+                ("run_direct_ag", f"AttnResFusedTmaKernel<{args}>::run_direct_ag"),
+            ],
+            extra_cuda_cflags=["-O3", "--use_fast_math"],
+        )
+
+
+# Benchmarked-best (chunk_rows, occupancy, consumer_regs) per nvb
+# (GB200/GB300-class, H=7168). nvb=1 is latency-bound per token, so 2 CTAs/SM
+# (occupancy=2, which excludes setmaxnreg: 4*Nc + 2*Np > 512) wins large T by
+# ~17%; nvb=4/8 fill 5-row chunks exactly ((nvb+1) % 5 == 0 or covers it in
+# 1-2 chunks); everything else is fastest on the balanced 4-row chunk. All
+# entries use the setmaxnreg producer/consumer split. consumer_regs sits on
+# the 200..232 performance plateau (below 200 the consumer loop starves); we
+# take 200, not the 232 budget cap, so each SMSP keeps 2K registers free for
+# PDL-overlapped neighbor kernels instead of allocating them idle.
+_TMA_BEST_CONFIG: dict[int, tuple[int, int, int]] = {
+    1: (2, 2, 0),
+    2: (4, 1, 200),
+    3: (4, 1, 200),
+    4: (5, 1, 200),
+    5: (3, 1, 200),
+    6: (4, 1, 200),
+    7: (4, 1, 200),
+    8: (5, 1, 200),
+}
+
+
+def _tuning(nvb: int, num_tokens: int) -> tuple[int, int, int]:
+    """(chunk_rows, occupancy, consumer_regs) for this aggregation point.
+
+    Shared by all three entry points below: they run the same kernel template
+    and differ only in the collective fused onto it."""
+    if not 1 <= nvb <= _MAX_BANK_ROWS:
+        raise ValueError(f"attn_res: nvb must be in [1, {_MAX_BANK_ROWS}], got {nvb}")
+    best = _TMA_BEST_CONFIG[nvb]
+    if best[1] > 1 and num_tokens < 128:
+        # occupancy=2 only pays off once there are enough tokens to fill both
+        # CTAs per SM; below that its tighter register budget just costs ~10%.
+        best = (4, 1, 200)
+    return best
+
+
+_COMM_MAP: dict[int, Communicator] = {}
+_PULL_SEM_MC_MAP: dict[int, int] = {}
+
+
+def register_comm(comm: Communicator, *, pull_sem_mc_ptr: int) -> None:
+    # One communicator per world_size per process -- see the note in
+    # kimi_k3/all_reduce.py::register_comm. The ops key only on world_size, so an
+    # overwrite here would hand the old group's callers the new group's peer
+    # pointers.
+    prev = _COMM_MAP.get(comm.world_size)
+    assert prev is None or prev is comm, (
+        f"a different communicator is already registered for world_size="
+        f"{comm.world_size}"
+    )
+    _COMM_MAP[comm.world_size] = comm
+    _PULL_SEM_MC_MAP[comm.world_size] = pull_sem_mc_ptr
+
+
+@register_custom_op(mutates_args=["out", "prefix_out"])
+def _attn_res_fused_pull_rs_op(
+    world_size: int,
+    input: torch.Tensor,
+    residual: torch.Tensor | None,
+    bank: torch.Tensor,
+    cw: torch.Tensor,
+    ow: torch.Tensor,
+    out: torch.Tensor,
+    prefix_out: torch.Tensor,
+    nvb: int,
+    eps: float,
+    input_mc_ptr: int,
+    max_blocks: int,
+    chunk_rows: int,
+    occupancy: int,
+    consumer_regs: int,
+) -> None:
+    _jit_fused_tma_module(chunk_rows, occupancy, consumer_regs).run_pull_rs(
+        _COMM_MAP[world_size],
+        input,
+        residual,
+        bank,
+        cw,
+        ow,
+        out,
+        prefix_out,
+        nvb,
+        eps,
+        input_mc_ptr,
+        _PULL_SEM_MC_MAP[world_size],
+        max_blocks,
+    )
+
+
+def attn_res_fused_tma(
+    prefix_sum: torch.Tensor,
+    bank: torch.Tensor,
+    cw: torch.Tensor,
+    ow: torch.Tensor,
+    out: torch.Tensor,
+    nvb: int,
+    eps: float,
+    *,
+    write_prefix: bool = False,
+) -> None:
+    """Warp-specialized TMA aggregation (score -> online softmax -> weighted
+    combine -> fused output RMSNorm), one persistent CTA per SM: a producer
+    warp (group) fetches rows with one bulk copy each into chunk slots of
+    `chunk_rows` rows (one barrier pair per chunk, double-buffered ring of 2
+    slots); the 8 consumer warps score and fold one chunk per rendezvous,
+    with cw / ow staged in TMEM.
+
+    Restrictions: H == 7168, nvb in [1, 8], SM100a+. The launch config comes
+    from _TMA_BEST_CONFIG via _tuning() (each combination compiles its own
+    module).
+
+    Parameters
+    ----------
+    prefix_sum : [T, H] bf16
+    bank       : [T, NB, H] bf16 (rows 0..nvb-1 are aggregated)
+    cw         : [H] bf16 — precomputed score_norm_weight * proj_weight
+    ow         : [H] bf16 — output RMSNorm weight
+    out        : [T, H] bf16 output buffer
+    nvb        : number of valid bank rows (1..8)
+    eps        : RMSNorm epsilon (shared by score and output norms)
+    write_prefix : also snapshot the prefix row into bank[:, nvb, :]
+                 (bit-exact copy, fused into the score pass which already
+                 has the row in registers); requires NB > nvb
+    """
+    _jit_fused_tma_module(*_tuning(nvb, prefix_sum.shape[0])).run(
+        prefix_sum, bank, cw, ow, out, nvb, eps, write_prefix
+    )
+
+
+@register_custom_op(mutates_args=["bank", "out"])
+def _attn_res_fused_direct_ag_op(
+    world_size: int,
+    prefix_sum: torch.Tensor,
+    bank: torch.Tensor,
+    cw: torch.Tensor,
+    ow: torch.Tensor,
+    out: torch.Tensor,
+    nvb: int,
+    eps: float,
+    output_mc_ptr: int,
+    max_blocks: int,
+    write_prefix: bool,
+    chunk_rows: int,
+    occupancy: int,
+    consumer_regs: int,
+) -> None:
+    _jit_fused_tma_module(chunk_rows, occupancy, consumer_regs).run_direct_ag(
+        _COMM_MAP[world_size],
+        prefix_sum,
+        bank,
+        cw,
+        ow,
+        out,
+        nvb,
+        eps,
+        output_mc_ptr,
+        _PULL_SEM_MC_MAP[world_size],
+        max_blocks,
+        write_prefix,
+    )
+
+
+def attn_res_fused_direct_ag(
+    world_size: int,
+    prefix_sum: torch.Tensor,
+    bank: torch.Tensor,
+    cw: torch.Tensor,
+    ow: torch.Tensor,
+    out: torch.Tensor,
+    nvb: int,
+    eps: float,
+    *,
+    output_mc_ptr: int,
+    max_blocks: int = 128,
+    write_prefix: bool = False,
+) -> torch.Tensor:
+    """Fuse local attention-residual aggregation with direct multicast AG.
+
+    `out` is the full symmetric [world * local_tokens, H] output. Each rank
+    aggregates its local token shard and multicast-stores normalized vectors
+    directly from consumer registers into that rank's slice on every peer.
+    """
+    _attn_res_fused_direct_ag_op(
+        world_size,
+        prefix_sum,
+        bank,
+        cw,
+        ow,
+        out,
+        nvb,
+        eps,
+        output_mc_ptr,
+        max_blocks,
+        write_prefix,
+        *_tuning(nvb, prefix_sum.shape[0]),
+    )
+    return out
+
+
+def attn_res_fused_pull_rs(
+    world_size: int,
+    input: torch.Tensor,
+    residual: torch.Tensor | None,
+    bank: torch.Tensor,
+    cw: torch.Tensor,
+    ow: torch.Tensor,
+    out: torch.Tensor,
+    prefix_out: torch.Tensor,
+    nvb: int,
+    eps: float,
+    *,
+    input_mc_ptr: int,
+    max_blocks: int = 128,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fused NVLS pull RS + local residual + TMA attention aggregation.
+
+    `input` is the full TP-partial o_proj tensor in multicast symmetric
+    memory. Every rank reduces only its contiguous token shard, optionally
+    adds its already-local residual, materializes that prefix in `prefix_out`,
+    and feeds it directly to the fused attention-residual aggregation/norm.
+    """
+    _attn_res_fused_pull_rs_op(
+        world_size,
+        input,
+        residual,
+        bank,
+        cw,
+        ow,
+        out,
+        prefix_out,
+        nvb,
+        eps,
+        input_mc_ptr,
+        max_blocks,
+        # prefix_out, not input: the shard this rank actually aggregates.
+        *_tuning(nvb, prefix_out.shape[0]),
+    )
+    return out, prefix_out

python/sglang/kernels/ops/kimi_k3/configs/sp_collective/world=4,H=7168,device_name=NVIDIA_GB300.json

@@ -0,0 +1,271 @@
+{
+  "source": {
+    "code_commits": [
+      "c4ad2e84b",
+      "5bf69bd32"
+    ],
+    "date": "2026-07-25",
+    "device": "NVIDIA GB300",
+    "nodes": 1,
+    "gpus_per_node": 4,
+    "push_slot_bytes": 33554432,
+    "raw_results": [
+      "standalone-c4ad2e84b-run3.log",
+      "standalone-large-c4ad2e84b.json",
+      "standalone-8192-tiebreak-c4ad2e84b.json",
+      "fused-small-5bf69bd32.json",
+      "fused-medium-5bf69bd32.json",
+      "fused-8192-5bf69bd32.json"
+    ]
+  },
+  "selection": {
+    "rule": "nearest global-token bucket at or below the workload",
+    "fallback": "nccl",
+    "note": "Standalone RS push won through T=2048 and NVLS pull won at T=4096/8192. Fused pull RS+residual+attention-residual wins through T=256; the separate path wins from T=512. Fused attention-residual+direct AG wins every measured bucket T=4..8192. The T=8192 standalone and fused AG custom wins were repeated with higher iteration counts. T=9364 is the first world=4 shape whose per-rank shard exceeds the 32 MiB symmetric slot, so it introduces explicit NCCL/separate safety boundaries."
+  },
+  "configs": {
+    "reduce_scatter": {
+      "4": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 512
+      },
+      "8": {
+        "strategy": "push",
+        "num_blocks": 96,
+        "block_size": 256
+      },
+      "16": {
+        "strategy": "push",
+        "num_blocks": 96,
+        "block_size": 512
+      },
+      "32": {
+        "strategy": "push",
+        "num_blocks": 96,
+        "block_size": 512
+      },
+      "64": {
+        "strategy": "push",
+        "num_blocks": 64,
+        "block_size": 256
+      },
+      "128": {
+        "strategy": "push",
+        "num_blocks": 64,
+        "block_size": 512
+      },
+      "256": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 256
+      },
+      "512": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 256
+      },
+      "1024": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 512
+      },
+      "2048": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 512
+      },
+      "4096": {
+        "strategy": "pull",
+        "num_blocks": 96,
+        "block_size": 1024
+      },
+      "8192": {
+        "strategy": "pull",
+        "num_blocks": 96,
+        "block_size": 1024
+      },
+      "9364": {
+        "strategy": "nccl"
+      },
+      "16384": {
+        "strategy": "nccl"
+      }
+    },
+    "all_gather": {
+      "4": {
+        "strategy": "direct",
+        "num_blocks": 8,
+        "block_size": 1024
+      },
+      "8": {
+        "strategy": "push",
+        "num_blocks": 32,
+        "block_size": 128
+      },
+      "16": {
+        "strategy": "push",
+        "num_blocks": 64,
+        "block_size": 256
+      },
+      "32": {
+        "strategy": "direct",
+        "num_blocks": 64,
+        "block_size": 128
+      },
+      "64": {
+        "strategy": "push",
+        "num_blocks": 32,
+        "block_size": 128
+      },
+      "128": {
+        "strategy": "push",
+        "num_blocks": 96,
+        "block_size": 256
+      },
+      "256": {
+        "strategy": "direct",
+        "num_blocks": 64,
+        "block_size": 256
+      },
+      "512": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 512
+      },
+      "1024": {
+        "strategy": "direct",
+        "num_blocks": 96,
+        "block_size": 1024
+      },
+      "2048": {
+        "strategy": "direct",
+        "num_blocks": 64,
+        "block_size": 512
+      },
+      "4096": {
+        "strategy": "direct",
+        "num_blocks": 64,
+        "block_size": 256
+      },
+      "8192": {
+        "strategy": "direct",
+        "num_blocks": 64,
+        "block_size": 1024
+      },
+      "9364": {
+        "strategy": "nccl"
+      },
+      "16384": {
+        "strategy": "nccl"
+      }
+    },
+    "reduce_scatter_attn_res": {
+      "4": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "8": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "16": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "32": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "64": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "128": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "256": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "512": {
+        "strategy": "separate"
+      },
+      "1024": {
+        "strategy": "separate"
+      },
+      "2048": {
+        "strategy": "separate"
+      },
+      "4096": {
+        "strategy": "separate"
+      },
+      "8192": {
+        "strategy": "separate"
+      },
+      "9364": {
+        "strategy": "separate"
+      },
+      "16384": {
+        "strategy": "separate"
+      }
+    },
+    "attn_res_all_gather": {
+      "4": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "8": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "16": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "32": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "64": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "128": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "256": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "512": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "1024": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "2048": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "4096": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "8192": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "9364": {
+        "strategy": "separate"
+      },
+      "16384": {
+        "strategy": "separate"
+      }
+    }
+  }
+}

python/sglang/kernels/ops/kimi_k3/configs/sp_collective/world=8,H=7168,device_name=NVIDIA_GB300.json

@@ -0,0 +1,235 @@
+{
+  "source": {
+    "code_commits": [
+      "8acabdbb7",
+      "e6bbf14d7",
+      "a74213e8c",
+      "7b1d23a0b",
+      "a132f1025",
+      "f574cdc8a"
+    ],
+    "date": "2026-07-25",
+    "device": "NVIDIA GB300",
+    "nodes": 2,
+    "gpus_per_node": 4,
+    "push_slot_bytes": 33554432,
+    "raw_results": [
+      "tune-8acabdbb7.json",
+      "ag-strategies-e6bbf14d7.json",
+      "rs-strategies-a74213e8c.json",
+      "sp-attn-res-composite-full-7b1d23a0b.json",
+      "fused-pull-attn-res-full-a132f1025.json",
+      "fused-direct-ag-full-f574cdc8a.json"
+    ]
+  },
+  "selection": {
+    "rule": "nearest global-token bucket at or below the workload",
+    "fallback": "nccl",
+    "note": "Standalone RS push won through T=1024 and NVLS pull won at T=2048/4096. Fused pull RS+residual+attention-residual wins through T=512; the separate path wins from T=1024. Fused attention-residual+direct AG wins every measured bucket T=8..4096. NCCL won standalone collectives at T=8192 and T=16384; fusion falls back outside its measured envelope."
+  },
+  "configs": {
+    "reduce_scatter": {
+      "8": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 512
+      },
+      "16": {
+        "strategy": "push",
+        "num_blocks": 4,
+        "block_size": 512
+      },
+      "32": {
+        "strategy": "push",
+        "num_blocks": 96,
+        "block_size": 512
+      },
+      "64": {
+        "strategy": "push",
+        "num_blocks": 16,
+        "block_size": 512
+      },
+      "128": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 512
+      },
+      "256": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 512
+      },
+      "512": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 512
+      },
+      "1024": {
+        "strategy": "push",
+        "num_blocks": 96,
+        "block_size": 512
+      },
+      "2048": {
+        "strategy": "pull",
+        "num_blocks": 96,
+        "block_size": 1024
+      },
+      "4096": {
+        "strategy": "pull",
+        "num_blocks": 96,
+        "block_size": 1024
+      },
+      "8192": {
+        "strategy": "nccl"
+      },
+      "16384": {
+        "strategy": "nccl"
+      }
+    },
+    "all_gather": {
+      "8": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 128
+      },
+      "16": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 512
+      },
+      "32": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 512
+      },
+      "64": {
+        "strategy": "push",
+        "num_blocks": 32,
+        "block_size": 256
+      },
+      "128": {
+        "strategy": "push",
+        "num_blocks": 64,
+        "block_size": 512
+      },
+      "256": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 512
+      },
+      "512": {
+        "strategy": "push",
+        "num_blocks": 128,
+        "block_size": 512
+      },
+      "1024": {
+        "strategy": "direct",
+        "num_blocks": 2,
+        "block_size": 1024
+      },
+      "2048": {
+        "strategy": "direct",
+        "num_blocks": 2,
+        "block_size": 1024
+      },
+      "4096": {
+        "strategy": "direct",
+        "num_blocks": 2,
+        "block_size": 1024
+      },
+      "8192": {
+        "strategy": "nccl"
+      },
+      "16384": {
+        "strategy": "nccl"
+      }
+    },
+    "reduce_scatter_attn_res": {
+      "8": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "16": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "32": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "64": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "128": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "256": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "512": {
+        "strategy": "fused_pull",
+        "max_blocks": 64
+      },
+      "1024": {
+        "strategy": "separate"
+      },
+      "2048": {
+        "strategy": "separate"
+      },
+      "4096": {
+        "strategy": "separate"
+      },
+      "8192": {
+        "strategy": "separate"
+      }
+    },
+    "attn_res_all_gather": {
+      "8": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "16": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "32": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "64": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "128": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "256": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "512": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "1024": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "2048": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "4096": {
+        "strategy": "fused_direct",
+        "max_blocks": 64
+      },
+      "8192": {
+        "strategy": "separate"
+      }
+    }
+  }
+}

python/sglang/kernels/ops/kimi_k3/gemm_ag.py

@@ -0,0 +1,88 @@
+"""K3 column-parallel up_proj + multicast all-gather + add3 (bf16, TP8).
+
+One entry point over ``csrc/kimi_k3/comm/gemm_ag.cuh``: for the latent MoE
+up_proj ([M, 3584] x [3584, 7168]) at small decode M, every rank computes
+only its 896-column slice of the replicated GEMM (the C++ side slices the
+full weight itself), multicast-stores it into the CustomAllReduceV2 push
+workspace (one more user of its double-buffer phase protocol), and a
+Lamport-spin consumer assembles ``out = up_proj(x) + b (+ c)`` — reading
+1/8 of the weight bytes per rank instead of all of them. Needs
+:func:`sglang.kernels.ops.kimi_k3.all_reduce.register_comm` once beforehand
+(the same registration the push all-reduce uses).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.kernels.ops.kimi_k3.all_reduce import _COMM_MAP
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+# Kimi-K3 up_proj dims (the kernel template takes any K/N passing its
+# static_asserts; this module instantiates the K3 shape).
+K = 3584
+N = 7168
+
+# Largest decode batch the kernel wins at (crossover vs the replicated
+# cublas GEMM + add3 tail is ~13-14 tokens on B200x8); also the GEMV
+# function-table size.
+MAX_TOKENS = 12
+
+
+@cache_once
+def _jit_module() -> Module:
+    args = make_cpp_args(K, N, MAX_TOKENS, is_arch_support_pdl())
+    cls = f"GEMMAGKernel<{args}>"
+    return load_jit(
+        "kimi_k3_gemm_ag",
+        *args,
+        cuda_files=["kimi_k3/comm/gemm_ag.cuh"],
+        cuda_wrappers=[("run", f"{cls}::run")],
+        extra_cuda_cflags=["-O3"],
+    )
+
+
+@register_custom_op(mutates_args=["out"])
+def _gemm_ag_op(
+    world_size: int,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    b: torch.Tensor,
+    c: Optional[torch.Tensor],
+    out: torch.Tensor,
+    ws_mc_base: int,
+) -> None:
+    comm = _COMM_MAP[world_size].obj
+    _jit_module().run(comm, x, weight, b, c, out, ws_mc_base)
+
+
+def gemm_ag_up_proj(
+    world_size: int,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    b: torch.Tensor,
+    c: Optional[torch.Tensor],
+    out: torch.Tensor,
+    *,
+    ws_mc_base: int,
+) -> torch.Tensor:
+    """``out = x @ weight.T (allgathered) + b (+ c)``, all bf16.
+
+    ``x`` is [M, 3584] with M in [1, MAX_TOKENS]; ``weight`` is the FULL
+    replicated [7168, 3584] up_proj weight (each rank reads only its own
+    row block); ``b`` / ``c`` / ``out`` are [M, 7168] (``out`` is
+    output-only). ``ws_mc_base`` is the multicast VA of the v2 workspace
+    slab base (``comm.mc_base_ptr``)."""
+    _gemm_ag_op(world_size, x, weight, b, c, out, ws_mc_base)
+    return out

python/sglang/kernels/ops/kimi_k3/gemm_ar.py

@@ -0,0 +1,216 @@
+"""K3 fused o_proj GEMM + all-reduce for decode (bf16, TP row-parallel).
+
+One entry point over ``csrc/kimi_k3/comm/gemm_ar.cuh``: a single kernel per
+rank computes the local ``x_r [M, K] @ W_r [7168, K]^T`` partial AND the
+cross-rank sum — the epilogue pushes finished tiles straight into a
+peer-mapped P2P comm region, one flag boundary, then a tile-local reduce
+writes the fully reduced ``out [M, 7168]`` on every rank. Replaces the
+o_proj GEMM + NCCL all-reduce pair with one launch (see GEMM_AR_README.md).
+
+Contracts:
+
+* bf16 only; ``out = sum_r bf16(x_r @ W_r^T)`` (partials round to bf16
+  pre-sum — same numerics as the unfused bf16 GEMM + ring AR).
+* M in [1, 512]; internally rounded up to a tuned cell {8, 16, 32, 64,
+  128, 256, 512}. ``out`` is allocated with ``cell`` rows and sliced.
+* SM100+ with full NVLink P2P (fabric/MNNVL across nodes); perf-tuned on
+  GB300 (sm_103a).
+* CUDA-graph compatible: the per-cell launch epoch lives in device memory
+  (read at kernel entry, bumped by a trailing kernel), so replays advance
+  it naturally. Each dispatch cell owns its own flag-ring family — no
+  host-side ring reset, ever.
+* All TP ranks must call :func:`o_proj_gemm_ar` with the same M in
+  lockstep (same stream order of cells on every rank).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, NamedTuple, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+N = 7168  # K3 hidden size (OPROJ_N compile-time default in gemm_ar.cuh)
+MAX_TOKENS = 512  # kMMax
+
+
+@cache_once
+def _jit_module(k: int, world_size: int) -> Module:
+    args = make_cpp_args(
+        k,
+        world_size,
+        is_arch_support_pdl(),
+    )
+    cls = f"GemmArKernel<{args}>"
+    return load_jit(
+        "kimi_k3_gemm_ar",
+        *args,
+        cuda_files=["kimi_k3/comm/gemm_ar.cuh"],
+        cuda_wrappers=[
+            ("run", f"{cls}::run"),
+            ("set_bases", f"{cls}::set_bases"),
+            ("region_nbytes", f"{cls}::region_nbytes"),
+            ("gather_words", f"{cls}::gather_words"),
+            ("num_fams", f"{cls}::num_fams"),
+        ],
+        extra_cuda_cflags=["-O3"],
+        extra_dependencies=["cutlass"],
+    )
+
+
+class _State(NamedTuple):
+    world_size: int
+    rank: int
+    region: tuple  # (slab tensor, peer buffer views) — keeps mappings alive
+    uc_bases: torch.Tensor  # [R] int64 CPU: per-rank UC VAs of the region
+    gather: torch.Tensor  # [kFams * 2 * kRing] int32 CUDA, device-local
+    epochs: torch.Tensor  # [kFams] int32 CUDA: device-resident CTA ticket counters
+
+
+_STATE: Optional[_State] = None
+
+
+def init(
+    *,
+    world_size: int,
+    rank: int,
+    group: torch.distributed.ProcessGroup,
+    k: int,
+) -> None:
+    """Allocate + rendezvous the P2P comm region (collective; call once from
+    every TP rank, BEFORE any CUDA-graph capture). ``group`` is the TP CPU
+    (gloo) group used for the symm-mem rendezvous."""
+    global _STATE
+    if _STATE is not None:
+        return
+    # the empty_strided_p2p + get_buffer path is the one that exchanges
+    # fabric handles and maps every peer (incl. cross-node MNNVL) into this
+    # process — the mem-pool rendezvous(tensor, group_name) API leaves
+    # remote-node (and sometimes even local) peers unmapped.
+    from torch._C._distributed_c10d import _SymmetricMemory
+
+    mod = _jit_module(k, world_size)
+    nbytes = int(mod.region_nbytes())
+    device = torch.device("cuda", torch.cuda.current_device())
+    if torch.__version__ < "2.11.0":
+        import torch.distributed._symmetric_memory as torch_symm_mem
+
+        torch_symm_mem.enable_symm_mem_for_group(group.group_name)
+    region = _SymmetricMemory.empty_strided_p2p(
+        (nbytes,), [1], torch.uint8, device, group.group_name
+    )
+    symm = _SymmetricMemory.rendezvous(region)
+    region.zero_()
+    torch.cuda.synchronize()
+    torch.distributed.barrier(group=group)
+    # keep the peer buffer tensors alive alongside the region
+    peer_bufs = [symm.get_buffer(r, [nbytes], torch.uint8) for r in range(world_size)]
+    ptrs = [t.data_ptr() for t in peer_bufs]
+    import logging
+
+    logging.getLogger(__name__).info(
+        "gemm_ar comm region: rank=%d nbytes=%d uc_bases=%s",
+        rank,
+        nbytes,
+        [hex(p) for p in ptrs],
+    )
+    assert all(p != 0 for p in ptrs), f"gemm_ar: null peer pointers {ptrs}"
+    # explicit cpu: model build may run under a cuda default-device context,
+    # and a silently-cuda tensor here means the host-side deref in set_bases
+    # reads a device pointer (segfault)
+    uc_bases = torch.tensor(ptrs, dtype=torch.int64, device="cpu")
+    gather = torch.zeros(int(mod.gather_words()), dtype=torch.int32, device=device)
+    epochs = torch.zeros(int(mod.num_fams()), dtype=torch.int32, device=device)
+    torch.cuda.synchronize()
+    _STATE = _State(
+        world_size=world_size,
+        rank=rank,
+        region=(region, peer_bufs),
+        uc_bases=uc_bases,
+        gather=gather,
+        epochs=epochs,
+    )
+
+
+def initialized() -> bool:
+    return _STATE is not None
+
+
+@cache_once
+def _module_with_bases(k: int, world_size: int) -> Module:
+    """The per-K JIT module with the comm-region base addresses stashed
+    host-side (per-call CPU-tensor derefs from inside the custom op segfault
+    under the sglang scheduler, so the module holds them in a static)."""
+    state = _STATE
+    assert state is not None
+    mod = _jit_module(k, world_size)
+    mod.set_bases(state.uc_bases)
+    return mod
+
+
+def fits(x: torch.Tensor) -> bool:
+    """Whether this o_proj input can take the fused GEMM+AR path."""
+    return (
+        _STATE is not None
+        and x.dim() == 2
+        and x.dtype == torch.bfloat16
+        and 0 < x.shape[0] <= MAX_TOKENS
+        and x.stride(1) == 1
+        and x.stride(0) == x.shape[1]
+    )
+
+
+@register_custom_op(mutates_args=["out", "epochs"])
+def _gemm_ar_op(
+    k: int,
+    world_size: int,
+    out: torch.Tensor,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    gather: torch.Tensor,
+    epochs: torch.Tensor,
+    my_rank: int,
+) -> None:
+    _module_with_bases(k, world_size).run(out, x, weight, gather, epochs, my_rank)
+
+
+def _cell_of(m: int) -> int:
+    for c in (8, 16, 32, 64, 128, 256, 512):
+        if m <= c:
+            return c
+    raise ValueError(f"gemm_ar: M={m} outside [1, {MAX_TOKENS}]")
+
+
+def o_proj_gemm_ar(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    """Fully reduced ``sum_r x_r @ weight_r^T`` on every rank, one kernel.
+
+    ``x`` is the TP-local [M, K] o_proj input, ``weight`` the TP-local
+    [7168, K] o_proj weight shard. Caller checked :func:`fits`; all ranks
+    call in lockstep with the same M.
+    """
+    state = _STATE
+    assert state is not None
+    m = x.shape[0]
+    cell = _cell_of(m)
+    out = torch.empty((cell, N), dtype=torch.bfloat16, device=x.device)
+    _gemm_ar_op(
+        weight.shape[1],
+        state.world_size,
+        out,
+        x,
+        weight,
+        state.gather,
+        state.epochs,
+        state.rank,
+    )
+    return out[:m]

python/sglang/kernels/ops/kimi_k3/mla_output_gate.py

@@ -0,0 +1,51 @@
+"""CUDA JIT K3 MLA output gate: out = x * sigmoid(gate) in one kernel."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+_THREADS: int = 256
+
+
+@cache_once
+def _jit_mla_output_gate_module() -> Module:
+    args = make_cpp_args(_THREADS, is_arch_support_pdl())
+    return load_jit(
+        "kimi_k3_mla_output_gate_" + str(_THREADS),
+        *args,
+        cuda_files=["kimi_k3/mla_output_gate.cuh"],
+        cuda_wrappers=[("run", f"MlaOutputGateKernel<{args}>::run")],
+        extra_cuda_cflags=["-O3"],
+    )
+
+
+def covered(x: torch.Tensor, gate: torch.Tensor) -> bool:
+    return (
+        x.dtype == torch.bfloat16
+        and gate.dtype == torch.bfloat16
+        and x.shape == gate.shape
+        and x.is_contiguous()
+        and gate.is_contiguous()
+        and x.numel() % 8 == 0
+        and x.numel() > 0
+    )
+
+
+def kimi_k3_mla_output_gate(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+    """out = bf16(x * bf16(sigmoid(gate))); double rounding matches the
+    unfused torch.sigmoid + mul pair bit-for-bit. Caller checks covered()."""
+    out = torch.empty_like(x)
+    _jit_mla_output_gate_module().run(x.view(-1), gate.view(-1), out.view(-1))
+    return out

python/sglang/kernels/ops/kimi_k3/moe.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+
+def _make_name(*args):
+    return "kimi_k3_" + "_".join(str(a) for a in args)
+
+
+@cache_once
+def _jit_situ_mul_quant_varlen_module(
+    quant_group_size: int,
+    scale_ue8m0: bool,
+    swizzle: bool,
+):
+    args = make_cpp_args(
+        quant_group_size,
+        scale_ue8m0,
+        swizzle,
+        is_arch_support_pdl(),
+    )
+    return load_jit(
+        _make_name("situ_mul_quant_varlen"),
+        *args,
+        cuda_files=["kimi_k3/situ_and_mul.cuh"],
+        cuda_wrappers=[("run", f"SituAndMulMaskedPostQuantKernel<{args}>::run")],
+        extra_cuda_cflags=["-use_fast_math"],
+    )
+
+
+def situ_and_mul_masked_post_quant(
+    input: torch.Tensor,
+    output: torch.Tensor,
+    output_scale: torch.Tensor,
+    quant_group_size: int,
+    masked_m: torch.Tensor,
+    beta: float,
+    linear_beta: float,
+    scale_ue8m0: bool = False,
+    topk: int = 8,
+    transposed: bool = False,
+    swizzle: bool = False,
+) -> None:
+    module = _jit_situ_mul_quant_varlen_module(quant_group_size, scale_ue8m0, swizzle)
+    module.run(
+        input,
+        output,
+        output_scale,
+        masked_m,
+        topk,
+        transposed,
+        float(beta),
+        float(linear_beta),
+    )

python/sglang/kernels/ops/kimi_k3/sp_collective.py

@@ -0,0 +1,338 @@
+"""K3 SP-MoE bf16 reduce-scatter and all-gather over MNNVL push memory."""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import TYPE_CHECKING, NamedTuple, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+    from sglang.kernels.ops.communication.all_reduce import Communicator
+
+
+class Tuning(NamedTuple):
+    num_blocks: int
+    block_size: int
+
+
+class Dispatch(NamedTuple):
+    strategy: str
+    tuning: Tuning
+
+
+class FusionDispatch(NamedTuple):
+    strategy: str
+    max_blocks: int
+
+
+# Safe seed values. GB300 production values are loaded by the layer glue from
+# the checked-in JSON after the sweep settles.
+DEFAULT_TUNING = Tuning(num_blocks=16, block_size=256)
+_CONFIG_DIR = os.path.join(os.path.dirname(__file__), "configs", "sp_collective")
+_TABLES: dict[str, Optional[dict]] = {}
+
+
+def _device_name(device: torch.device) -> str:
+    return torch.cuda.get_device_name(device).replace(" ", "_").replace("/", "_")
+
+
+def _table(world_size: int, hidden_size: int, device: torch.device) -> Optional[dict]:
+    path = os.path.join(
+        _CONFIG_DIR,
+        (
+            f"world={world_size},H={hidden_size},"
+            f"device_name={_device_name(device)}.json"
+        ),
+    )
+    if path not in _TABLES:
+        if os.path.exists(path):
+            with open(path) as f:
+                _TABLES[path] = json.load(f)
+        else:
+            _TABLES[path] = None
+    return _TABLES[path]
+
+
+def get_dispatch(
+    kind: str,
+    world_size: int,
+    hidden_size: int,
+    num_tokens: int,
+    device: torch.device,
+) -> Optional[Dispatch]:
+    """Return the tuned strategy, or None when the table selects NCCL."""
+    table = _table(world_size, hidden_size, device)
+    if table is None:
+        return None
+    raw_configs = table["configs"].get(kind)
+    if raw_configs is None:
+        return None
+    configs = {int(k): v for k, v in raw_configs.items()}
+    bucket = min(configs)
+    for candidate in sorted(configs):
+        if candidate <= num_tokens:
+            bucket = candidate
+        else:
+            break
+    config = configs[bucket]
+    if config["strategy"] == "nccl":
+        return None
+    return Dispatch(
+        config["strategy"],
+        Tuning(config["num_blocks"], config["block_size"]),
+    )
+
+
+def get_tuning(
+    kind: str,
+    world_size: int,
+    hidden_size: int,
+    num_tokens: int,
+    device: torch.device,
+) -> Optional[Tuning]:
+    """Compatibility helper for callers that only support staging push."""
+    dispatch = get_dispatch(kind, world_size, hidden_size, num_tokens, device)
+    if dispatch is None or dispatch.strategy != "push":
+        return None
+    return dispatch.tuning
+
+
+def get_fusion_dispatch(
+    kind: str,
+    world_size: int,
+    hidden_size: int,
+    num_tokens: int,
+    device: torch.device,
+) -> Optional[FusionDispatch]:
+    """Return a measured fused strategy, or None for the separate path."""
+    table = _table(world_size, hidden_size, device)
+    if table is None:
+        return None
+    raw_configs = table["configs"].get(kind)
+    if raw_configs is None:
+        return None
+    configs = {int(k): v for k, v in raw_configs.items()}
+    bucket = min(configs)
+    for candidate in sorted(configs):
+        if candidate <= num_tokens:
+            bucket = candidate
+        else:
+            break
+    config = configs[bucket]
+    if config["strategy"] == "separate":
+        return None
+    return FusionDispatch(config["strategy"], config["max_blocks"])
+
+
+@cache_once
+def _jit_module(world_size: int) -> Module:
+    args = make_cpp_args(world_size, is_arch_support_pdl())
+    cls = f"SPCollectiveKernel<{args}>"
+    return load_jit(
+        "kimi_k3_sp_collective",
+        *args,
+        cuda_files=["kimi_k3/comm/sp_collective.cuh"],
+        cuda_wrappers=[
+            ("reduce_scatter_res", f"{cls}::reduce_scatter_res"),
+            ("reduce_scatter_pull", f"{cls}::reduce_scatter_pull"),
+            ("all_gather", f"{cls}::all_gather"),
+            ("all_gather_direct", f"{cls}::all_gather_direct"),
+        ],
+        extra_cuda_cflags=["-O3"],
+    )
+
+
+_COMM_MAP: dict[int, Communicator] = {}
+_PULL_SEM_MC_MAP: dict[int, int] = {}
+
+
+def register_comm(comm: Communicator, *, pull_sem_mc_ptr: int = 0) -> None:
+    # One communicator per world_size per process -- see the note in
+    # kimi_k3/all_reduce.py::register_comm. The ops key only on world_size, so an
+    # overwrite here would hand the old group's callers the new group's peer
+    # pointers.
+    prev = _COMM_MAP.get(comm.world_size)
+    assert prev is None or prev is comm, (
+        f"a different communicator is already registered for world_size="
+        f"{comm.world_size}"
+    )
+    _COMM_MAP[comm.world_size] = comm
+    _PULL_SEM_MC_MAP[comm.world_size] = pull_sem_mc_ptr
+
+
+@register_custom_op(mutates_args=["output"])
+def _reduce_scatter_res_op(
+    world_size: int,
+    input: torch.Tensor,
+    output: torch.Tensor,
+    residual: Optional[torch.Tensor],
+    residual_is_local: bool,
+    num_blocks: int,
+    block_size: int,
+) -> None:
+    _jit_module(world_size).reduce_scatter_res(
+        _COMM_MAP[world_size],
+        input.view(-1),
+        output.view(-1),
+        None if residual is None else residual.view(-1),
+        residual_is_local,
+        num_blocks,
+        block_size,
+    )
+
+
+@register_custom_op(mutates_args=["output"])
+def _reduce_scatter_pull_op(
+    world_size: int,
+    input: torch.Tensor,
+    output: torch.Tensor,
+    residual: Optional[torch.Tensor],
+    residual_is_local: bool,
+    input_mc_ptr: int,
+    num_blocks: int,
+    block_size: int,
+) -> None:
+    _jit_module(world_size).reduce_scatter_pull(
+        _COMM_MAP[world_size],
+        input.view(-1),
+        output.view(-1),
+        None if residual is None else residual.view(-1),
+        residual_is_local,
+        input_mc_ptr,
+        _PULL_SEM_MC_MAP[world_size],
+        num_blocks,
+        block_size,
+    )
+
+
+@register_custom_op(mutates_args=["output"])
+def _all_gather_op(
+    world_size: int,
+    input: torch.Tensor,
+    output: torch.Tensor,
+    ws_mc_base: int,
+    num_blocks: int,
+    block_size: int,
+) -> None:
+    _jit_module(world_size).all_gather(
+        _COMM_MAP[world_size],
+        input.view(-1),
+        output.view(-1),
+        ws_mc_base,
+        num_blocks,
+        block_size,
+    )
+
+
+@register_custom_op(mutates_args=["output"])
+def _all_gather_direct_op(
+    world_size: int,
+    input: torch.Tensor,
+    output: torch.Tensor,
+    output_mc_ptr: int,
+    num_blocks: int,
+    block_size: int,
+) -> None:
+    _jit_module(world_size).all_gather_direct(
+        _COMM_MAP[world_size],
+        input.view(-1),
+        output.view(-1),
+        output_mc_ptr,
+        _PULL_SEM_MC_MAP[world_size],
+        num_blocks,
+        block_size,
+    )
+
+
+def reduce_scatter_res(
+    world_size: int,
+    input: torch.Tensor,
+    output: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    *,
+    tuning: Tuning = DEFAULT_TUNING,
+) -> torch.Tensor:
+    residual_is_local = residual is not None and residual.numel() == output.numel()
+    _reduce_scatter_res_op(
+        world_size,
+        input,
+        output,
+        residual,
+        residual_is_local,
+        tuning.num_blocks,
+        tuning.block_size,
+    )
+    return output
+
+
+def reduce_scatter_pull(
+    world_size: int,
+    input: torch.Tensor,
+    output: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    *,
+    input_mc_ptr: int,
+    tuning: Tuning = DEFAULT_TUNING,
+) -> torch.Tensor:
+    residual_is_local = residual is not None and residual.numel() == output.numel()
+    _reduce_scatter_pull_op(
+        world_size,
+        input,
+        output,
+        residual,
+        residual_is_local,
+        input_mc_ptr,
+        tuning.num_blocks,
+        tuning.block_size,
+    )
+    return output
+
+
+def all_gather(
+    world_size: int,
+    input: torch.Tensor,
+    output: torch.Tensor,
+    *,
+    ws_mc_base: int,
+    tuning: Tuning = DEFAULT_TUNING,
+) -> torch.Tensor:
+    _all_gather_op(
+        world_size,
+        input,
+        output,
+        ws_mc_base,
+        tuning.num_blocks,
+        tuning.block_size,
+    )
+    return output
+
+
+def all_gather_direct(
+    world_size: int,
+    input: torch.Tensor,
+    output: torch.Tensor,
+    *,
+    output_mc_ptr: int,
+    tuning: Tuning = DEFAULT_TUNING,
+) -> torch.Tensor:
+    _all_gather_direct_op(
+        world_size,
+        input,
+        output,
+        output_mc_ptr,
+        tuning.num_blocks,
+        tuning.block_size,
+    )
+    return output

python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py

@@ -501,3 +501,101 @@ def scatter_mamba_states_after_mtp_verify(
                 mamba_track_indices,
                 mamba_steps_to_track,
             )
+
+
+@triton.jit
+def track_mamba_states_all_layers_kernel(
+    conv_states_ptr,  # [num_layers, pool_size, ...] full conv pool
+    ssm_states_ptr,  # [num_layers, pool_size, ...] full ssm pool
+    cache_indices_ptr,
+    mamba_track_mask_ptr,
+    mamba_track_indices_ptr,
+    conv_layer_stride,
+    conv_row_stride,
+    ssm_layer_stride,
+    ssm_row_stride,
+    batch_size,
+    conv_state_numel_per_row: tl.constexpr,
+    ssm_state_numel_per_row: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    check_freed_slots: tl.constexpr,
+):
+    """All-layers variant of track_mamba_state_if_needed_kernel: one launch
+    covers every mamba layer (grid = num_layers * batch_size) instead of one
+    launch per layer. The track mask / source / destination indices are shared
+    across layers, so a single launch at the end of the step is equivalent to
+    the per-layer launches (each layer's state is final by then)."""
+    pid = tl.program_id(0)
+    layer_idx = (pid // batch_size).to(tl.int64)
+    batch_idx = pid % batch_size
+
+    track_mask = tl.load(mamba_track_mask_ptr + batch_idx)
+    if not track_mask:
+        return
+
+    src_idx = tl.load(cache_indices_ptr + batch_idx).to(tl.int64)
+    dst_idx = tl.load(mamba_track_indices_ptr + batch_idx).to(tl.int64)
+    if check_freed_slots:
+        if src_idx < 0 or dst_idx < 0:
+            return
+
+    conv_base = conv_states_ptr + layer_idx * conv_layer_stride
+    for offset in range(0, conv_state_numel_per_row, BLOCK_SIZE):
+        element_indices = offset + tl.arange(0, BLOCK_SIZE)
+        mask = element_indices < conv_state_numel_per_row
+        data = tl.load(
+            conv_base + src_idx * conv_row_stride + element_indices,
+            mask=mask,
+            other=0.0,
+        )
+        tl.store(
+            conv_base + dst_idx * conv_row_stride + element_indices, data, mask=mask
+        )
+
+    ssm_base = ssm_states_ptr + layer_idx * ssm_layer_stride
+    for offset in range(0, ssm_state_numel_per_row, BLOCK_SIZE):
+        element_indices = offset + tl.arange(0, BLOCK_SIZE)
+        mask = element_indices < ssm_state_numel_per_row
+        data = tl.load(
+            ssm_base + src_idx * ssm_row_stride + element_indices, mask=mask, other=0.0
+        )
+        tl.store(ssm_base + dst_idx * ssm_row_stride + element_indices, data, mask=mask)
+
+
+def track_mamba_states_all_layers(
+    conv_states_pool: torch.Tensor,
+    ssm_states_pool: torch.Tensor,
+    cache_indices: torch.Tensor,
+    mamba_track_mask: torch.Tensor,
+    mamba_track_indices: torch.Tensor,
+    batch_size: int,
+    check_freed_slots: bool = False,
+):
+    """Track conv/ssm states for ALL mamba layers in one launch.
+
+    conv_states_pool / ssm_states_pool are the full [num_layers, pool_size,
+    ...] pools; per-row copy semantics are identical to
+    track_mamba_states_if_needed applied per layer.
+    """
+    num_layers = conv_states_pool.shape[0]
+    conv_state_numel_per_row = conv_states_pool[0, 0].numel()
+    ssm_state_numel_per_row = ssm_states_pool[0, 0].numel()
+
+    BLOCK_SIZE = 1024
+    grid = (num_layers * batch_size,)
+    track_mamba_states_all_layers_kernel[grid](
+        conv_states_pool,
+        ssm_states_pool,
+        cache_indices,
+        mamba_track_mask,
+        mamba_track_indices,
+        conv_states_pool.stride(0),
+        conv_states_pool.stride(1),
+        ssm_states_pool.stride(0),
+        ssm_states_pool.stride(1),
+        batch_size,
+        conv_state_numel_per_row,
+        ssm_state_numel_per_row,
+        BLOCK_SIZE,
+        check_freed_slots,
+    )

python/sglang/kernels/ops/memory/gpu_tensor_hash.py

@@ -40,7 +40,7 @@ def _fmix32(x, C1: tl.constexpr, C2: tl.constexpr):
     return x
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["n_u32", "seed1", "seed2"])
 def hash_tiles32_kernel_blocked(
     in_ptr,
     out_ptr,
@@ -100,7 +100,7 @@ def hash_tiles32_kernel_blocked(
     tl.store(out_ptr + pid, out)
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["n_elems"])
 def add_tree_reduce_u64_kernel(in_ptr, out_ptr, n_elems, CHUNK: tl.constexpr):
     pid = tl.program_id(axis=0)
     start = pid * CHUNK

python/sglang/kernels/ops/mm/__init__.py

@@ -0,0 +1,3 @@
+"""Multimodal kernels."""
+
+__all__ = ["process"]

python/sglang/kernels/ops/mm/process/__init__.py

@@ -0,0 +1,5 @@
+"""Multimodal input-processing kernels."""
+
+from sglang.kernels.ops.mm.process.image import normalize_and_patchify
+
+__all__ = ["normalize_and_patchify"]

python/sglang/kernels/ops/mm/process/image.py

@@ -0,0 +1,124 @@
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+_MAX_TRITON_ELEMENTS = 2**31 - 1
+
+
+@triton.jit
+def _normalize_and_patchify_kernel(
+    input_ptr,
+    scale_ptr,
+    bias_ptr,
+    output_ptr,
+    channels: tl.constexpr,
+    input_height,
+    input_width,
+    grid_height,
+    grid_width,
+    patch_size: tl.constexpr,
+    element_count,
+    block_size: tl.constexpr,
+):
+    offsets = tl.program_id(0) * block_size + tl.arange(0, block_size)
+    mask = offsets < element_count
+
+    patch_area = patch_size * patch_size
+    patch_offset = offsets % patch_area
+    patch_x = patch_offset % patch_size
+    patch_y = patch_offset // patch_size
+    remaining = offsets // patch_area
+    channel = remaining % channels
+    remaining = remaining // channels
+    grid_x = remaining % grid_width
+    remaining = remaining // grid_width
+    grid_y = remaining % grid_height
+    batch = remaining // grid_height
+
+    input_y = grid_y * patch_size + patch_y
+    input_x = grid_x * patch_size + patch_x
+    input_mask = mask & (input_y < input_height) & (input_x < input_width)
+    input_offset = (
+        (batch * channels + channel) * input_height + input_y
+    ) * input_width + input_x
+    value = tl.load(input_ptr + input_offset, mask=input_mask, other=0.0)
+    scale = tl.load(scale_ptr + channel, mask=mask)
+    bias = tl.load(bias_ptr + channel, mask=mask)
+    tl.store(output_ptr + offsets, value * scale + bias, mask=mask)
+
+
+def _normalize_and_patchify_torch(
+    image: torch.Tensor,
+    image_scale: torch.Tensor,
+    image_bias: torch.Tensor,
+    patch_size: int,
+    padded_height: int,
+    padded_width: int,
+) -> torch.Tensor:
+    pad_height = padded_height - image.shape[-2]
+    pad_width = padded_width - image.shape[-1]
+    if pad_height > 0 or pad_width > 0:
+        image = F.pad(image, (0, pad_width, 0, pad_height), value=0.0)
+    image = torch.addcmul(image_bias, image, image_scale)
+    batch, channels, height, width = image.shape
+    grid_height = height // patch_size
+    grid_width = width // patch_size
+    image = image.view(batch, channels, grid_height, patch_size, grid_width, patch_size)
+    return image.permute(0, 2, 4, 1, 3, 5).reshape(
+        batch, -1, channels, patch_size, patch_size
+    )
+
+
+def normalize_and_patchify(
+    image: torch.Tensor,
+    image_scale: torch.Tensor,
+    image_bias: torch.Tensor,
+    patch_size: int,
+    padded_height: int,
+    padded_width: int,
+) -> torch.Tensor:
+    batch, channels, _, _ = image.shape
+    grid_height = padded_height // patch_size
+    grid_width = padded_width // patch_size
+    element_count = (
+        batch * grid_height * grid_width * channels * patch_size * patch_size
+    )
+    if (
+        not image.is_cuda
+        or torch.version.hip is not None
+        or not image.is_contiguous()
+        or not image_scale.is_contiguous()
+        or not image_bias.is_contiguous()
+        or element_count > _MAX_TRITON_ELEMENTS
+    ):
+        return _normalize_and_patchify_torch(
+            image,
+            image_scale,
+            image_bias,
+            patch_size,
+            padded_height,
+            padded_width,
+        )
+
+    output = torch.empty(
+        (batch, grid_height * grid_width, channels, patch_size, patch_size),
+        dtype=image.dtype,
+        device=image.device,
+    )
+    block_size = 256
+    _normalize_and_patchify_kernel[(triton.cdiv(element_count, block_size),)](
+        image,
+        image_scale,
+        image_bias,
+        output,
+        channels,
+        image.shape[-2],
+        image.shape[-1],
+        grid_height,
+        grid_width,
+        patch_size,
+        element_count,
+        block_size,
+    )
+    return output

python/sglang/kernels/ops/moe/configs/moe_front/epilogue,E=896,topk=16,device_name=NVIDIA_GB300.json

@@ -0,0 +1,154 @@
+{
+  "meta": {
+    "device": "NVIDIA GB300",
+    "device_name": "NVIDIA_GB300",
+    "torch": "2.11.0+cu130",
+    "experts": 896,
+    "latent": 3584,
+    "topk": 16,
+    "inner": 50,
+    "reps": 20
+  },
+  "configs": {
+    "1": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "2": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "3": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "4": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "6": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": true
+    },
+    "8": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "12": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "16": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "24": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "32": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "48": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "64": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "96": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "128": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "192": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "256": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "384": {
+      "block_size": 448,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "512": {
+      "block_size": 224,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "768": {
+      "block_size": 224,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "1024": {
+      "block_size": 224,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "1536": {
+      "block_size": 224,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "2048": {
+      "block_size": 224,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "3072": {
+      "block_size": 224,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "4096": {
+      "block_size": 224,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "6144": {
+      "block_size": 224,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "8192": {
+      "block_size": 224,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "12288": {
+      "block_size": 224,
+      "cast_vec": 8,
+      "cast_first": false
+    },
+    "16384": {
+      "block_size": 224,
+      "cast_vec": 8,
+      "cast_first": false
+    }
+  }
+}

python/sglang/kernels/ops/moe/configs/moe_front/strategy,E=896,topk=16,device_name=NVIDIA_GB300.json

@@ -0,0 +1,37 @@
+{
+  "meta": {
+    "device": "NVIDIA GB300",
+    "hidden": 7168,
+    "experts": 896,
+    "latent": 3584,
+    "topk": 16,
+    "routing_dtype": "fp32",
+    "selection": "exact token count; overlap must win all 3 repeats and median speedup must exceed 1%"
+  },
+  "configs": {
+    "736": "overlap",
+    "768": "overlap",
+    "896": "overlap",
+    "960": "overlap",
+    "1280": "overlap",
+    "1536": "overlap",
+    "1792": "overlap",
+    "2048": "overlap",
+    "2304": "overlap",
+    "2560": "overlap",
+    "2816": "overlap",
+    "3328": "overlap",
+    "3584": "overlap",
+    "3840": "overlap",
+    "4352": "overlap",
+    "4608": "overlap",
+    "4864": "overlap",
+    "5632": "overlap",
+    "5888": "overlap",
+    "6144": "overlap",
+    "7424": "overlap",
+    "8448": "overlap",
+    "8704": "overlap",
+    "8960": "overlap"
+  }
+}

python/sglang/kernels/ops/moe/moe_front.py

@@ -0,0 +1,235 @@
+"""K3 MoE front: merged gate + routed_expert_down_proj GEMM, and the fp32 router.
+
+The unfused MoE front -- the path every EP-a2a / WideEP deployment takes -- runs
+three ops over the same `hidden_states [T, 7168]`:
+
+    router_logits = gate(hidden_states)                     # [896, 7168]  12.85 MB
+    topk_output   = topk(hidden_states, router_logits)
+    routed_input  = routed_expert_down_proj(hidden_states)   # [3584, 7168] 51.4 MB
+
+Plain [M, 896] fp32 logits go to route_radix. This module covers what that
+cannot: the merged front.
+
+**fused_front** -- the two GEMMs share their input, so their weights are merged
+and one cuBLAS GEMM emits `[T, 896 + 3584]` fp32; a single epilogue kernel then
+runs the top-k on the gate slice and casts the latent slice to bf16.  Routing
+stays bit-identical to the fp32 path and routed_input comes out dense.
+
+**overlap** -- the original fp32 gate + top-k run on the model's side stream
+while the latent down-projection runs on the main stream. The streams join
+before expert dispatch, and the side stream is then reused by the existing
+shared-expert overlap.
+
+Measured in-graph on a GB300, us per MoE layer:
+
+    T            512    768   1024   1280   2048   2560   4096   8192  16384
+    baseline    31.1   37.6   50.9   54.3   91.5   99.3  152.7  304.7  629.3
+    merged      26.1   37.3   46.7      -      -      -      -      -      -
+    overlap     29.0   33.7   46.5   52.4   90.5   96.4  151.2  307.1  637.8
+
+The fastest strategy is non-monotonic because cuBLAS changes GEMM algorithms at
+specific row counts. A GB300 JSON table therefore opts exact, repeatedly
+measured token counts into overlap; unmeasured shapes keep the conservative
+defaults (merged through 1024, unfused above it). This captures the clear 768
+and 2560 wins without regressing 512, 8192, or 16384.
+
+A bf16-output merged GEMM was measured too (fastest at T=1, 11.8 us).  It is not
+used: bf16 rounds the router logits and moves the selected expert set on 2-25% of
+rows depending on T, for ~1.2 us -- and from T>=8 it loses to the fp32 variant
+anyway, because its routed_input is a strided slice a dense-input runner must copy.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+import json
+import os
+
+NUM_EXPERTS = 896
+TOPK = 16
+
+# Above this token count the merged GEMM stops paying by default; see the table
+# above. A device strategy table may override individual, measured token counts
+# with the dual-stream overlap.
+MERGED_FRONT_MAX_TOKENS = 1024
+
+_CONFIG_DIR = os.path.join(os.path.dirname(__file__), "configs", "moe_front")
+
+# Kernel tunables, per token bucket, from the JSON table.
+#   block_size  threads per CTA; sets experts-per-thread in the radix select
+#               (896 / block_size). 224 -> 4, 448 -> 2.
+#   cast_vec    fp32 elements each thread converts per step in the latent cast.
+#   cast_first  issue the cast before the select (loads in flight during the
+#               radix rounds) or after it.
+# Fallbacks when no tuned table matches the device. Both were the sweep's most
+# common winners: cast_vec 8 is 32 B/thread, the Blackwell vector-load limit, and
+# it won at every one of the 28 token counts measured; issuing the cast after the
+# select beat issuing it before almost everywhere.
+DEFAULT_EPILOGUE_CONFIG = {"block_size": 224, "cast_vec": 8, "cast_first": False}
+
+_tables = {}
+
+
+def _table(kind: str, device_name: str):
+    path = os.path.join(
+        _CONFIG_DIR,
+        f"{kind},E={NUM_EXPERTS},topk={TOPK},device_name={device_name}.json",
+    )
+    if path not in _tables:
+        table = None
+        if os.path.exists(path):
+            with open(path) as f:
+                table = {int(k): v for k, v in json.load(f)["configs"].items()}
+        _tables[path] = table
+    return _tables[path]
+
+
+def _device_name(device) -> str:
+    return torch.cuda.get_device_name(device).replace(" ", "_").replace("/", "_")
+
+
+def get_config(kind: str, num_tokens: int, device, default: dict) -> dict:
+    """Tuned config for the nearest token bucket at or below `num_tokens`."""
+    table = _table(kind, _device_name(device))
+    if not table:
+        return dict(default)
+    pick = min(table)
+    for k in sorted(table):
+        if k <= num_tokens:
+            pick = k
+        else:
+            break
+    return dict(table[pick])
+
+
+def get_front_strategy(num_tokens: int, device) -> str:
+    """Return the measured front strategy for this exact workload.
+
+    GEMM algorithm changes make the merged/overlap crossover non-monotonic in
+    M. Therefore strategy tables are exact-match only. Unmeasured shapes keep
+    the robust defaults: merged fp32 through 1024 tokens, then unfused.
+    """
+    table = _table("strategy", _device_name(device))
+    if table is not None and num_tokens in table:
+        return str(table[num_tokens])
+    return "merged_fp32" if num_tokens <= MERGED_FRONT_MAX_TOKENS else "unfused"
+
+
+@cache_once
+def _jit_module() -> Module:
+    args = make_cpp_args(is_arch_support_pdl())
+    return load_jit(
+        "moe_front",
+        *args,
+        cuda_files=["moe/route_radix.cuh"],
+        cuda_wrappers=[
+            ("front_epilogue", f"FusedFrontEpilogueKernel<{args}>::run"),
+        ],
+        # No fast-math: scoring and expert-id selection must stay comparable to
+        # route_radix / the Triton router under ties and NaN.
+        extra_cuda_cflags=["-O3"],
+    )
+
+
+@cache_once
+def available() -> bool:
+    import logging
+
+    try:
+        _jit_module()
+        return True
+    except Exception as e:  # pragma: no cover - toolchain dependent
+        logging.getLogger(__name__).warning(
+            f"Failed to load the JIT MoE front kernels: {e}"
+        )
+        return False
+
+
+# merged front: [gate | down] GEMM -> top-k + routed_input
+
+
+def fused_front_covered(
+    hidden_states: torch.Tensor,
+    merged_weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    topk: int,
+    latent: int,
+) -> bool:
+    """[T<=MERGED_FRONT_MAX_TOKENS, 7168] bf16 x [896 + latent, 7168] bf16, fp32
+    bias, top-16, latent a multiple of 4."""
+    return (
+        hidden_states.dim() == 2
+        and merged_weight.dim() == 2
+        and hidden_states.dtype == torch.bfloat16
+        and merged_weight.dtype == torch.bfloat16
+        and bias is not None
+        and bias.dtype == torch.float32
+        and bias.numel() == NUM_EXPERTS
+        and int(topk) == TOPK
+        and merged_weight.shape[0] == NUM_EXPERTS + latent
+        and merged_weight.shape[1] == hidden_states.shape[1]
+        and latent % 4 == 0
+        and 0 < hidden_states.shape[0] <= MERGED_FRONT_MAX_TOKENS
+        and hidden_states.stride(1) == 1
+        and merged_weight.stride(1) == 1
+    )
+
+
+def fused_front(
+    hidden_states: torch.Tensor,
+    merged_weight: torch.Tensor,
+    correction_bias: torch.Tensor,
+    latent: int,
+    topk: int = TOPK,
+    renormalize: bool = True,
+    routed_scaling_factor: float = 1.0,
+    apply_routed_scaling_factor_on_output: bool = False,
+    config: Optional[dict] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Merged front GEMM + fused top-k/cast epilogue.
+
+    Returns ``(topk_weights [M, topk] fp32, topk_ids [M, topk] int32,
+    routed_input [M, latent] bf16)``.
+    """
+    M = hidden_states.shape[0]
+    device = hidden_states.device
+    if config is None:
+        config = get_config("epilogue", M, device, DEFAULT_EPILOGUE_CONFIG)
+
+    # fp32 out keeps routing exact; the extra output traffic versus bf16 is
+    # M x (896 + latent) x 2 bytes, negligible against the 64 MB weight read at
+    # the sizes this path serves.
+    merged = torch.mm(hidden_states, merged_weight.t(), out_dtype=torch.float32)
+
+    weights = torch.empty((M, topk), dtype=torch.float32, device=device)
+    ids = torch.empty((M, topk), dtype=torch.int32, device=device)
+    routed = torch.empty((M, latent), dtype=torch.bfloat16, device=device)
+
+    _jit_module().front_epilogue(
+        merged,
+        correction_bias,
+        weights,
+        ids,
+        routed,
+        topk,
+        float(routed_scaling_factor if routed_scaling_factor is not None else 1.0),
+        bool(renormalize),
+        bool(apply_routed_scaling_factor_on_output),
+        int(config["block_size"]),
+        int(config["cast_vec"]),
+        bool(config["cast_first"]),
+    )
+    return weights, ids, routed

python/sglang/kernels/ops/moe/moe_fused_gate.py

@@ -9,11 +9,11 @@ import triton.language as tl
 
 from sglang.kernel_api_logging import debug_kernel_api
 from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl, load_jit
+from sglang.kernels.ops.moe import moe_route_radix
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module
 
-
 _SCORING_FUNC_MAP = {
     "sigmoid": 0,
     "sqrtsoftplus": 1,
@@ -290,6 +290,30 @@ def moe_fused_gate(
     if routed_scaling_factor is None:
         routed_scaling_factor = 1.0
 
+    # K3 radix-select fast path: native-CUDA radix-select replaces the 16
+    # dependent argmax rounds (single CTA per token; ids bit-identical to this
+    # triton kernel incl. ties).
+    # The radix kernel keeps keys register-resident and returns winners in
+    # expert-id order (skipping the biased-descending sort; downstream MoE
+    # kernels are order-insensitive). It is 3.1-3.5x faster than the Triton
+    # kernel at [1..8192, 896] top-16 on B200.
+    if (
+        scoring_func.lower() == "sigmoid"
+        and num_fused_shared_experts == 0
+        and num_expert_group <= 1
+        and moe_softcapping == 0.0
+    ):
+        radix_args = (
+            scores,
+            bias,
+            topk,
+            renormalize,
+            routed_scaling_factor,
+            apply_routed_scaling_factor_on_output,
+        )
+        if moe_route_radix.covered(scores, bias, topk):
+            return moe_route_radix.route_radix(*radix_args, sorted=False)
+
     M, N = scores.shape
     K = topk
     K_routed = topk - num_fused_shared_experts
@@ -309,7 +333,10 @@ def moe_fused_gate(
     # stay occupancy-bound. Swept on H100/B200: this beats the AOT kernels across
     # shapes, whereas larger tiles / more warps regress (register pressure).
     BLOCK_M = max(1, min(4, 256 // BLOCK_N))
-    num_warps = 1
+    # For wide rows (e.g. Kimi K3: 896 experts, BLOCK_N 1024) the K sequential
+    # argmax passes dominate and benefit from more warps despite the
+    # cross-warp reduction cost.
+    num_warps = 1 if BLOCK_N <= 512 else 4
     grid = (triton.cdiv(M, BLOCK_M),)
     use_pdl = is_arch_support_pdl()
     extra = {"launch_pdl": True} if use_pdl else {}

python/sglang/kernels/ops/moe/moe_route_quant_fused.py

@@ -0,0 +1,125 @@
+"""Fused K3 MoE-front prep: radix routing + trtllm id pack + mxfp8 quant.
+
+One launch replaces the three tiny kernels between the K3 fused-front GEMM and
+the trtllm-gen routed-MoE op at decode batch sizes (route_radix -> triton
+(id<<16|bf16(w)) pack -> per_token_group_quant, ~7.5us busy + 2 extra launches
+per MoE layer with the SMs near idle). CTAs [0, M) run the route_radix body
+with the pack folded into its epilogue; CTAs [M, 2M) run the
+per_token_group_quant math, one CTA per token row. Both halves reuse the
+standalone kernels' device code, so ids/weights, packed ids, and quantized
+activations are bit-identical to the unfused chain.
+
+Specialized like route_radix itself: 896 experts, top-16, bf16/fp32 scores,
+and a 3584-wide bf16 activation row quantized to fp8 with row-major packed
+UE8M0 group-32 scales (the trtllm-gen SiTU MoE input format). Wired into
+serving through sglang.srt.layers.moe.route_quant_handoff.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Tuple
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.kernels.ops.moe import moe_route_radix
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+_HIDDEN = 3584
+_GROUP_SIZE = 32
+_NUM_GROUPS = _HIDDEN // _GROUP_SIZE
+# Fusion trades the flat quant grid for one 224-thread CTA per token; that (and
+# the win itself, which is launch overhead) only makes sense at small decode
+# batches. Above the cap the callers run the unfused chain.
+_MAX_TOKENS = 64
+
+
+@cache_once
+def _jit_module() -> Module:
+    args = make_cpp_args(is_arch_support_pdl())
+    return load_jit(
+        "moe_route_quant_fused",
+        *args,
+        cuda_files=["moe/route_quant_fused.cuh"],
+        cuda_wrappers=[("run", f"RouteQuantFusedKernel<{args}>::run")],
+        # No fast-math: the routing half must stay bit-identical to
+        # route_radix (see its module comment); the quant half's math is
+        # fast-math-independent (explicit intrinsics + bit manipulation).
+        extra_cuda_cflags=["-O3"],
+    )
+
+
+@cache_once
+def available() -> bool:
+    import logging
+
+    try:
+        _jit_module()
+        return True
+    except Exception as e:  # pragma: no cover - toolchain dependent
+        logging.getLogger(__name__).warning(
+            f"Failed to load the JIT fused route+quant kernel: {e}"
+        )
+        return False
+
+
+def covered(
+    scores: torch.Tensor, bias: torch.Tensor, topk: int, x: torch.Tensor
+) -> bool:
+    """route_radix coverage plus the quant half: [M<=64, 3584] bf16 rows with
+    32B-aligned starts (base and stride), same token count as the scores."""
+    return (
+        moe_route_radix.covered(scores, bias, topk)
+        and x.dim() == 2
+        and x.shape[0] == scores.shape[0]
+        and 0 < x.shape[0] <= _MAX_TOKENS
+        and x.shape[1] == _HIDDEN
+        and x.dtype == torch.bfloat16
+        and x.stride(1) == 1
+        and x.data_ptr() % 32 == 0
+        and (x.stride(0) * x.element_size()) % 32 == 0
+    )
+
+
+def route_quant_fused(
+    scores: torch.Tensor,
+    bias: torch.Tensor,
+    x: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    routed_scaling_factor: float,
+    apply_scale: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Returns ``(weights [M, topk] fp32, ids [M, topk] int32, packed [M, topk]
+    int32, x_q [M, 3584] fp8_e4m3, x_s [M, 28] int32 row-major packed UE8M0)``.
+    Caller must have checked covered(). Winners come out in expert-id-ascending
+    order (the standalone production dispatch's sorted=False)."""
+    M = scores.shape[0]
+    device = scores.device
+    out_w = torch.empty((M, topk), dtype=torch.float32, device=device)
+    out_i = torch.empty((M, topk), dtype=torch.int32, device=device)
+    out_packed = torch.empty((M, topk), dtype=torch.int32, device=device)
+    out_q = torch.empty((M, _HIDDEN), dtype=torch.float8_e4m3fn, device=device)
+    out_s = torch.empty((M, _NUM_GROUPS // 4), dtype=torch.int32, device=device)
+    _jit_module().run(
+        scores,
+        bias,
+        out_w,
+        out_i,
+        out_packed,
+        x,
+        out_q,
+        out_s,
+        topk,
+        float(routed_scaling_factor),
+        bool(renormalize),
+        bool(apply_scale),
+    )
+    return out_w, out_i, out_packed, out_q, out_s

python/sglang/kernels/ops/moe/moe_route_radix.py

@@ -0,0 +1,98 @@
+"""Native-CUDA radix-select router for K3 routing (all batch sizes).
+
+Keys and activations stay in registers (224 threads, 4 experts each), the
+split-bin search runs on warp scans instead of cub, rounds exit early when the
+top-k separates on a byte boundary, and the (biased desc, id asc) output sort
+is optional. Consumers that only gather by expert id can pass sorted=False and
+skip the epilogue rank-sort entirely.
+
+Dispatched automatically from moe_fused_gate for covered inputs; the production
+dispatch uses sorted=False. It is 3.1-3.5x faster than the Triton router at
+[1..8192, 896] top-16 on B200. Correctness coverage lives in
+test_kimi_k3_prerequisite_ops.py, against a pure-torch fp32 oracle rather than the
+Triton router: moe_fused_gate dispatches back here for every input this kernel
+covers, so using it as the reference compares the kernel with itself.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Tuple
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+_NUM_EXPERTS = 896
+_TOPK = 16
+
+
+@cache_once
+def _jit_route_radix_module() -> Module:
+    args = make_cpp_args(is_arch_support_pdl())
+    return load_jit(
+        "moe_route_radix",
+        *args,
+        cuda_files=["moe/route_radix.cuh"],
+        cuda_wrappers=[("run", f"RouteRadixKernel<{args}>::run")],
+        # No fast-math: expert-id selection must stay bit-identical to the
+        # Triton router under ties/NaN.
+        extra_cuda_cflags=["-O3"],
+    )
+
+
+def covered(scores: torch.Tensor, bias: torch.Tensor, topk: int) -> bool:
+    """Specialized for K3 decode routing: [M, 896] bf16 or fp32
+    row-contiguous scores (8B/16B-aligned rows), fp32 bias, top-16."""
+    return (
+        scores.dim() == 2
+        and scores.size(1) == _NUM_EXPERTS
+        and int(topk) == _TOPK
+        and scores.dtype in (torch.bfloat16, torch.float32)
+        and bias.dtype == torch.float32
+        and scores.stride(1) == 1
+        and scores.stride(0) % 4 == 0
+        and bias.is_contiguous()
+    )
+
+
+def route_radix(
+    scores: torch.Tensor,
+    bias: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    routed_scaling_factor: float,
+    apply_scale: bool,
+    sorted: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Returns (weights [M, topk] fp32, ids [M, topk] int32). Caller must have
+    checked covered().
+
+    Default sorted=False: winners come out in expert-id-ascending order
+    (downstream MoE kernels are order-insensitive) and the epilogue rank-sort
+    is skipped. sorted=True restores the Triton router's (biased desc, id asc)
+    output order. Either way the winner set matches Triton exactly; the renorm
+    sum is taken in the respective output order, so weights may differ by
+    <= ~1 ulp."""
+    M = scores.shape[0]
+    out_w = torch.empty((M, topk), dtype=torch.float32, device=scores.device)
+    out_i = torch.empty((M, topk), dtype=torch.int32, device=scores.device)
+    _jit_route_radix_module().run(
+        scores,
+        bias,
+        out_w,
+        out_i,
+        topk,
+        float(routed_scaling_factor),
+        bool(renormalize),
+        bool(apply_scale),
+        bool(sorted),
+    )
+    return out_w, out_i

python/sglang/kernels/ops/moe/moe_topk_sum.py

@@ -0,0 +1,37 @@
+"""CUDA JIT top-k expert-output sum: out[M, K] = in[M, topk, K].sum(dim=1)."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+_THREADS: int = 256
+
+
+@cache_once
+def _jit_topk_sum_module() -> Module:
+    args = make_cpp_args(_THREADS, is_arch_support_pdl())
+    return load_jit(
+        "moe_topk_sum_" + str(_THREADS),
+        *args,
+        cuda_files=["moe/topk_sum.cuh"],
+        cuda_wrappers=[("run", f"TopkSumKernel<{args}>::run")],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+    )
+
+
+def moe_topk_sum(x: torch.Tensor, out: torch.Tensor) -> torch.Tensor:
+    """out[M, K] = x[M, topk, K].sum(dim=1) for contiguous bf16 tensors."""
+    _jit_topk_sum_module().run(x, out)
+    return out

python/sglang/kernels/ops/moe/pack_topk_ids.py

@@ -9,6 +9,8 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.kernels.jit.utils import is_arch_support_pdl
+
 
 class PackTopkIds:
 
@@ -51,12 +53,16 @@ class PackTopkIds:
 
         BLOCK_SIZE = 1024
         grid = (triton.cdiv(numel, BLOCK_SIZE),)
+        pdl_kwargs = (
+            {"USE_PDL": True, "launch_pdl": True} if is_arch_support_pdl() else {}
+        )
         _pack_topk_ids_triton_kernel[grid](
             topk_ids,
             topk_weights,
             out,
             numel,
             BLOCK_SIZE=BLOCK_SIZE,
+            **pdl_kwargs,
         )
         return out
 
@@ -68,14 +74,21 @@ def _pack_topk_ids_triton_kernel(
     out_ptr,
     numel,
     BLOCK_SIZE: tl.constexpr,
+    USE_PDL: tl.constexpr = False,
 ):
     pid = tl.program_id(0)
     offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
     mask = offsets < numel
 
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+
     ids = tl.load(topk_ids_ptr + offsets, mask=mask, other=0)
     w = tl.load(topk_weights_ptr + offsets, mask=mask, other=0.0)
 
+    if USE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
     w_bf16 = w.to(tl.bfloat16)
     w_i16 = w_bf16.to(tl.int16, bitcast=True)
     w_i32 = w_i16.to(tl.int32) & 0xFFFF

python/sglang/kernels/ops/moe/trtllm_gen_moe.py

@@ -0,0 +1,528 @@
+"""TRT-LLM-gen fused MoE (SiTU) compiled through the sglang JIT system.
+
+Builds the trtllm-gen fused-MoE host/runner sources with sglang's own
+tvm-ffi ``load_jit`` from a **self-contained cubin pool**
+(``SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL``): a downloadable directory holding
+
+  * the prebuilt SiTU cubins (``local/``) + ``config.json`` +
+    ``flashinferMetaInfo.h``,
+  * the flat batched-gemm ABI headers (staged into a
+    ``trtllmGen_bmm_export/``-shaped include tree at build time),
+  * an ``overlay/`` with only the sources/headers that differ from the
+    public ``flashinfer`` pip package.
+
+Every unmodified source and the CUTLASS headers come from the installed
+``flashinfer`` package's ``data/`` tree (the wheel ships it for its own
+JIT), so running this backend needs exactly one download and one env var —
+no extra source checkout.
+
+This module vendors only glue:
+
+  * header staging: the pool ships the batched-gemm ABI headers flat; they
+    are copied into a content-addressed include tree shaped like
+    ``flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/``;
+  * JIT build of the 12 launcher/runner/routing sources with the private
+    ABI defines (``TLLM_GEN_LOCAL_CUBINS_ABI`` etc.);
+  * the ctypes cubin-loader callback (the .so asks for cubins by absolute
+    path + sha256; we read them from the pool);
+  * a thin ``trtllm_fp4_block_scale_moe`` wrapper (FromLogits routing,
+    ``do_finalize=True``); kernel tile config ("tactic") defaults to the
+    runner's built-in heuristic — pass an explicit one for tuned setups.
+
+Validated for the Kimi K3 decode/prefill MoE regime: MxFP4 weights with
+bf16 (w4a16) or MxFP8 (w4a8) activations, ``ActivationType.Situ`` (SiTuGlu:
+``a*tanh(g/a)*sigmoid(g) * b*tanh(u/b)``), DeepSeekV3/noaux_tc routing.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import hashlib
+import logging
+import os
+import pathlib
+import shutil
+from typing import TYPE_CHECKING, Optional, Sequence
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    get_jit_cuda_arch,
+    load_jit,
+    override_jit_cuda_arch,
+)
+from sglang.srt.environ import envs
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+# ActivationType / RoutingMethodType values from trtllm-gen's tllm_enums
+# (kept as plain ints here to avoid importing anything for them).
+ACTIVATION_SITU = 9
+ROUTING_DEEPSEEK_V3 = 2
+_ROUTING_TOPK = 5
+_ROUTING_INPUT_FROM_LOGITS = 0
+# NOTE: the enum VALUES start at 0; the "Mode 1/2/3" wording in upstream
+# comments is documentation numbering, not the enum value.
+_ROUTING_INPUT_PACKED = 1
+
+# Batched-gemm ABI headers shipped flat in the cubin pool; the launcher
+# includes them as flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/<h>.
+_BMM_EXPORT_HEADERS = [
+    "BatchedGemmEnums.h",
+    "BatchedGemmInterface.h",
+    "BatchedGemmOptions.h",
+    "Enums.h",
+    "GemmGatedActOptions.h",
+    "GemmOptions.h",
+    "KernelParams.h",
+    "KernelParamsDecl.h",
+    "KernelTraits.h",
+    "TmaDescriptor.h",
+    "trtllm/gen/CommonUtils.h",
+    "trtllm/gen/CudaArchDecl.h",
+    "trtllm/gen/CudaKernelLauncher.h",
+    "trtllm/gen/DtypeDecl.h",
+    "trtllm/gen/MmaDecl.h",
+    "trtllm/gen/SfLayoutDecl.h",
+    "trtllm/gen/SparsityDecl.h",
+]
+
+_SOURCES = [
+    "csrc/nv_internal/cpp/kernels/quantization.cu",
+    "csrc/nv_internal/cpp/common/envUtils.cpp",
+    "csrc/nv_internal/cpp/common/logger.cpp",
+    "csrc/nv_internal/cpp/common/stringUtils.cpp",
+    "csrc/nv_internal/cpp/common/tllmException.cpp",
+    "csrc/nv_internal/cpp/common/memoryUtils.cu",
+    "csrc/trtllm_fused_moe_kernel_launcher.cu",
+    "csrc/trtllm_fused_moe_runner.cu",
+    "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_deepseek.cu",
+    "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_llama4.cu",
+    "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cu",
+    "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_common.cu",
+    "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu",
+    "csrc/trtllm_batched_gemm_runner.cu",
+]
+
+
+logger = logging.getLogger(__name__)
+
+
+def cubin_pool_dir() -> Optional[pathlib.Path]:
+    p = envs.SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL.get()
+    if not p:
+        return None
+    pool = pathlib.Path(p)
+    return pool if pool.is_dir() else None
+
+
+def _flashinfer_data_dir() -> Optional[pathlib.Path]:
+    """The installed public flashinfer package's JIT source tree (ships
+    csrc/, include/ and its pinned cutlass), used as the base layer under
+    the pool's overlay."""
+    try:
+        import flashinfer  # noqa: PLC0415
+    except ImportError:
+        return None
+    data = pathlib.Path(flashinfer.__file__).parent / "data"
+    return data if (data / "csrc").is_dir() else None
+
+
+def available() -> bool:
+    pool = cubin_pool_dir()
+    return (
+        pool is not None
+        and (pool / "flashinferMetaInfo.h").is_file()
+        and (pool / "local").is_dir()
+        # Modified sources ship in the pool's overlay/, everything else
+        # compiles from the installed flashinfer package.
+        and (pool / "overlay" / "csrc").is_dir()
+        and _flashinfer_data_dir() is not None
+    )
+
+
+def _stage_headers(pool: pathlib.Path) -> pathlib.Path:
+    """Copy the pool's ABI headers into a content-addressed include tree."""
+    meta = (pool / "flashinferMetaInfo.h").read_bytes()
+    tag = hashlib.sha256(meta).hexdigest()[:12]
+    cache = pathlib.Path(
+        os.environ.get("TVM_FFI_CACHE_DIR", "~/.cache/tvm-ffi")
+    ).expanduser()
+    root = cache / "trtllm_gen_moe_headers" / tag
+    dest = root / "flashinfer" / "trtllm" / "batched_gemm" / "trtllmGen_bmm_export"
+    stamp = root / ".staged"
+    if not stamp.is_file():
+        for name in _BMM_EXPORT_HEADERS:
+            target = dest / name
+            target.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copyfile(pool / name, target)
+        shutil.copyfile(pool / "flashinferMetaInfo.h", dest / "flashinferMetaInfo.h")
+        stamp.touch()
+    return root
+
+
+def _cuda_home() -> pathlib.Path:
+    home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
+    if not home:
+        nvcc = shutil.which("nvcc")
+        home = str(pathlib.Path(nvcc).parent.parent) if nvcc else "/usr/local/cuda"
+    return pathlib.Path(home)
+
+
+def _cuda_include_dir() -> str:
+    return str(_cuda_home() / "include")
+
+
+def _cuda_stub_ldflags() -> list[str]:
+    """-L flags for the libcuda driver stub, so -lcuda links in bare build
+    environments (containers without the driver lib on the default linker
+    path); the real driver is dlopened at runtime as usual."""
+    home = _cuda_home()
+    stubs = [
+        home / "lib64" / "stubs",
+        *home.glob("targets/*/lib/stubs"),
+    ]
+    return [f"-L{s}" for s in stubs if s.is_dir()]
+
+
+_CUBIN_CB_KEEPALIVE = {}
+
+
+def _setup_cubin_loader(so_path: str, pool_local: pathlib.Path) -> None:
+    """Register the ctypes callback the .so uses to fetch cubins by name.
+
+    The runner requests ``<TLLM_GEN_GEMM_CUBIN_PATH>/<kernel>`` (absolute,
+    because the pool path is baked in at compile time); we read the bytes
+    and hand them back via FlashInferSetCurrentCubin.
+    """
+    if so_path in _CUBIN_CB_KEEPALIVE:
+        return
+    lib = ctypes.CDLL(so_path)
+    cb_type = ctypes.CFUNCTYPE(None, ctypes.c_char_p, ctypes.c_char_p)
+
+    def _get_cubin(name: bytes, sha256: bytes) -> None:
+        rel = name.decode()
+        path = pathlib.Path(rel)
+        if not path.is_absolute():
+            path = pool_local / rel
+        if path.suffix != ".cubin":
+            path = path.with_name(path.name + ".cubin")
+        data = path.read_bytes()
+        want = sha256.decode()
+        if want:
+            got = hashlib.sha256(data).hexdigest()
+            if got != want:
+                raise RuntimeError(
+                    f"cubin sha mismatch for {path}: want {want} got {got}"
+                )
+        lib.FlashInferSetCurrentCubin(
+            ctypes.cast(ctypes.create_string_buffer(data, len(data)), ctypes.c_char_p),
+            ctypes.c_int(len(data)),
+        )
+
+    cb = cb_type(_get_cubin)
+    _CUBIN_CB_KEEPALIVE[so_path] = (lib, cb)
+    lib.FlashInferSetCubinCallback(cb)
+
+
+@cache_once
+def _jit_trtllm_gen_moe_module() -> Module:
+    pool = cubin_pool_dir()
+    fi_data = _flashinfer_data_dir()
+    if pool is None or not (pool / "overlay" / "csrc").is_dir() or fi_data is None:
+        raise RuntimeError(
+            "trtllm-gen MoE sources not found: point "
+            "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL at an unpacked cubin pool "
+            "(cubins + flat ABI headers + overlay/) and install the public "
+            "flashinfer package."
+        )
+    # Overlay first (its modified sources/headers shadow the public copies),
+    # installed flashinfer data as the base.
+    src_roots = [pool / "overlay", fi_data]
+    include_roots = [pool / "overlay", fi_data]
+
+    def _resolve_source(rel: str) -> str:
+        for root in src_roots:
+            cand = root / rel
+            if cand.is_file():
+                return str(cand)
+        raise RuntimeError(f"trtllm-gen MoE source not found in any root: {rel}")
+
+    staged = _stage_headers(pool)
+    meta_tag = staged.name
+    cubin_path = str((pool / "local").resolve())
+
+    cache = pathlib.Path(
+        os.environ.get("TVM_FFI_CACHE_DIR", "~/.cache/tvm-ffi")
+    ).expanduser()
+    # Flags are not part of load_jit's source hash: fold the pool identity
+    # (meta hash + path) into the module marker so a pool change rebuilds.
+    path_tag = hashlib.sha256(cubin_path.encode()).hexdigest()[:8]
+    build_dir = cache / f"sgl_trtllm_gen_moe_{meta_tag}_{path_tag}"
+
+    cpp_files = [_resolve_source(s) for s in _SOURCES if s.endswith(".cpp")]
+    cuda_files = [_resolve_source(s) for s in _SOURCES if s.endswith(".cu")]
+    # quantization.cu emits fp4 cvt instructions (.e2m1x2) that need the
+    # arch-specific feature set: compile for sm_XXXa, not plain sm_XXX.
+    # The trtllm-gen cubins themselves are prebuilt (sm100f) and loaded at
+    # runtime, unaffected by this flag.
+    arch = get_jit_cuda_arch()
+    with override_jit_cuda_arch(arch.major, arch.minor, "a"):
+        module = load_jit(
+            "trtllm_gen_moe",
+            meta_tag,
+            path_tag,
+            external_cpp_files=cpp_files,
+            external_cuda_files=cuda_files,
+            header_only=False,  # the launcher exports its own tvm-ffi functions
+            extra_cflags=["-fvisibility=hidden"],
+            extra_cuda_cflags=[
+                "-DTLLM_GEN_EXPORT_INTERFACE",
+                "-DTLLM_GEN_EXPORT_FLASHINFER",
+                "-DTLLM_ENABLE_CUDA",
+                "-DENABLE_BF16",
+                "-DENABLE_FP8",
+                "-DENABLE_FP4",
+                "-DCUTLASS_ENABLE_GDC_FOR_SM100=1",
+                "-DTLLM_GEN_LOCAL_CUBINS_ABI",
+                "-DFLASHINFER_PRIVATE_MOE_FFI_NAMES",
+                "-DFLASHINFER_PRIVATE_MOE_LEAN_ROUTING",
+                f'-DTLLM_GEN_GEMM_CUBIN_PATH=\\"{cubin_path}\\"',
+                "-Xcompiler=-fvisibility=hidden",
+            ],
+            extra_ldflags=[*_cuda_stub_ldflags(), "-lcuda", "-lnvrtc"],
+            extra_include_paths=[
+                str(staged),
+                str(
+                    staged
+                    / "flashinfer"
+                    / "trtllm"
+                    / "batched_gemm"
+                    / "trtllmGen_bmm_export"
+                ),
+                # Per-root include layout: include/, csrc/, csrc/nv_internal/,
+                # csrc/nv_internal/include/, plus the flashinfer package's
+                # pinned CUTLASS (data/cutlass/). The overlay root comes first
+                # so modified headers shadow the public copies.
+                *[
+                    str(root / sub)
+                    for root in include_roots
+                    for sub in (
+                        "include",
+                        "csrc",
+                        "csrc/nv_internal",
+                        "csrc/nv_internal/include",
+                    )
+                ],
+                *[
+                    str(root / "cutlass" / "include")
+                    for root in include_roots
+                    if (root / "cutlass" / "include").is_dir()
+                ],
+                # Host .cpp files (g++) need the CUDA headers explicitly; nvcc
+                # adds them implicitly for .cu. CUDA 13's bundled CCCL is
+                # used as-is (mixing another pinned CCCL with the toolkit's
+                # explodes).
+                _cuda_include_dir(),
+            ],
+            build_directory=str(build_dir),
+        )
+    so_files = sorted(build_dir.glob("*.so"))
+    if not so_files:
+        raise RuntimeError(f"no built .so under {build_dir}")
+    _setup_cubin_loader(str(so_files[-1]), pool / "local")
+    return module
+
+
+def trtllm_fp4_block_scale_moe(
+    routing_logits: torch.Tensor,
+    routing_bias: Optional[torch.Tensor],
+    hidden_states: torch.Tensor,
+    hidden_states_scale: Optional[torch.Tensor],
+    gemm1_weights: torch.Tensor,
+    gemm1_weights_scale: torch.Tensor,
+    gemm1_alpha: Optional[torch.Tensor],
+    gemm1_beta: Optional[torch.Tensor],
+    gemm2_weights: torch.Tensor,
+    gemm2_weights_scale: torch.Tensor,
+    output1_scale_scalar: Optional[torch.Tensor],
+    output1_scale_gate_scalar: Optional[torch.Tensor],
+    output2_scale_scalar: Optional[torch.Tensor],
+    num_experts: int,
+    top_k: int,
+    n_group: Optional[int],
+    topk_group: Optional[int],
+    intermediate_size: int,
+    routed_scaling_factor: Optional[float],
+    routing_method_type: int = ROUTING_DEEPSEEK_V3,
+    activation_type: int = ACTIVATION_SITU,
+    norm_topk_prob: bool = True,
+    local_expert_offset: int = 0,
+    local_num_experts: Optional[int] = None,
+    tactic: Sequence[int] = (-1, -1),
+    output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP4 block-scale MoE with routing from logits and finalize fused.
+
+    ``hidden_states``: bf16 ``[T, hidden]`` (w4a16) or MxFP8-packed uint8
+    with ``hidden_states_scale`` (w4a8). Weights are trtllm-gen shuffled
+    MxFP4 (uint8 packed, fp8 block scales, MajorK). ``tactic`` is the
+    (gemm1, gemm2) config index pair; ``(-1, -1)`` = runner heuristic.
+    """
+    module = _jit_trtllm_gen_moe_module()
+    # The FFI launcher reads these as dense row-major; a strided slice
+    # (e.g. a fused-GEMM split) would silently mis-route.
+    routing_logits = routing_logits.contiguous()
+    hidden_states = hidden_states.contiguous()
+    num_tokens = routing_logits.shape[0]
+    hidden_size = hidden_states.shape[-1]
+    if hidden_states.dtype == torch.uint8:
+        hidden_size *= 2
+    device = hidden_states.device
+    topk_ids = torch.empty(num_tokens, top_k, dtype=torch.int32, device=device)
+    topk_weights = torch.empty(
+        num_tokens, top_k, dtype=routing_logits.dtype, device=device
+    )
+    if output is None:
+        output = torch.empty(
+            num_tokens, hidden_size, dtype=torch.bfloat16, device=device
+        )
+    module.trtllm_fp4_block_scale_moe_private(
+        _ROUTING_INPUT_FROM_LOGITS,
+        routing_logits,
+        topk_ids,
+        topk_weights,
+        routing_bias,
+        hidden_states,
+        hidden_states_scale,
+        gemm1_weights,
+        gemm1_weights_scale,
+        None,  # gemm1_bias
+        gemm1_alpha,
+        gemm1_beta,
+        None,  # gemm1_clamp_limit
+        gemm2_weights,
+        gemm2_weights_scale,
+        None,  # gemm2_bias
+        output1_scale_scalar,
+        output1_scale_gate_scalar,
+        output2_scale_scalar,
+        None,  # per_token_scale
+        num_experts,
+        top_k,
+        n_group,
+        topk_group,
+        intermediate_size,
+        local_expert_offset,
+        num_experts if local_num_experts is None else local_num_experts,
+        routed_scaling_factor,
+        routing_method_type,
+        True,  # do_finalize
+        True,  # enable_pdl
+        activation_type,
+        output,
+        list(tactic),
+        norm_topk_prob,
+        None,  # routing_replay_out
+    )
+    return output
+
+
+def trtllm_fp4_block_scale_routed_moe(
+    packed_topk_ids: torch.Tensor,
+    hidden_states: torch.Tensor,
+    hidden_states_scale: Optional[torch.Tensor],
+    gemm1_weights: torch.Tensor,
+    gemm1_weights_scale: torch.Tensor,
+    gemm1_alpha: Optional[torch.Tensor],
+    gemm1_beta: Optional[torch.Tensor],
+    gemm2_weights: torch.Tensor,
+    gemm2_weights_scale: torch.Tensor,
+    output1_scale_scalar: Optional[torch.Tensor],
+    output1_scale_gate_scalar: Optional[torch.Tensor],
+    output2_scale_scalar: Optional[torch.Tensor],
+    num_experts: int,
+    top_k: int,
+    intermediate_size: int,
+    activation_type: int = ACTIVATION_SITU,
+    local_expert_offset: int = 0,
+    local_num_experts: Optional[int] = None,
+    tactic: Sequence[int] = (-1, -1),
+    output: Optional[torch.Tensor] = None,
+    do_finalize: bool = True,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """FP4 block-scale MoE with PRECOMPUTED routing (PackedPrecomputed).
+
+    ``packed_topk_ids``: int32 ``[T, top_k]`` with ``(expert_id << 16) |
+    bf16-weight-bits`` (PackTopkIds layout) — selection and weights come
+    from the caller's router, the in-op routing kernels are skipped. This
+    is the fast path at small T, where the in-op single-CTA routing kernel
+    (~22 µs at 896 experts) costs more than an external radix router.
+
+    ``do_finalize=False`` skips the in-op finalize (top-k weighted
+    unpermute) and returns its inputs instead:
+    ``(gemm2_output [padded_rows, hidden] bf16 in permuted layout,
+    topk_weights [T, top_k] bf16 unpacked from packed_topk_ids,
+    expanded_idx_to_permuted_idx [T*top_k] int32 with -1 = dropped slot)``.
+    ``output`` is left unwritten in that mode.
+    """
+    module = _jit_trtllm_gen_moe_module()
+    hidden_states = hidden_states.contiguous()
+    num_tokens = packed_topk_ids.shape[0]
+    hidden_size = hidden_states.shape[-1]
+    if hidden_states.dtype == torch.uint8:
+        hidden_size *= 2
+    device = hidden_states.device
+    # Mode 2 unpacks the weights in-kernel; this is its output buffer.
+    topk_weights = torch.empty(num_tokens, top_k, dtype=torch.bfloat16, device=device)
+    if output is None:
+        output = torch.empty(
+            num_tokens, hidden_size, dtype=torch.bfloat16, device=device
+        )
+    result = module.trtllm_fp4_block_scale_moe_private(
+        _ROUTING_INPUT_PACKED,
+        None,  # routing_logits
+        packed_topk_ids.contiguous(),
+        topk_weights,
+        None,  # routing_bias (already applied by the external router)
+        hidden_states,
+        hidden_states_scale,
+        gemm1_weights,
+        gemm1_weights_scale,
+        None,  # gemm1_bias
+        gemm1_alpha,
+        gemm1_beta,
+        None,  # gemm1_clamp_limit
+        gemm2_weights,
+        gemm2_weights_scale,
+        None,  # gemm2_bias
+        output1_scale_scalar,
+        output1_scale_gate_scalar,
+        output2_scale_scalar,
+        None,  # per_token_scale
+        num_experts,
+        top_k,
+        None,  # n_group
+        None,  # topk_group
+        intermediate_size,
+        local_expert_offset,
+        num_experts if local_num_experts is None else local_num_experts,
+        1.0,  # routed_scaling_factor (already applied by the router)
+        _ROUTING_TOPK,  # routing_method_type (unused for precomputed)
+        do_finalize,
+        True,  # enable_pdl
+        activation_type,
+        output,
+        list(tactic),
+        True,  # norm_topk_prob (unused for precomputed)
+        None,  # routing_replay_out
+    )
+    if do_finalize:
+        return output
+    # Deferred: [gemm2_output, expert_weights (None in packed mode — the
+    # weights live in the topk_weights buffer mode 2 unpacked into),
+    # expanded_idx_to_permuted_idx]. Index access — iterating the tvm-ffi
+    # Array yields one-shot dlpack capsules.
+    return result[0], topk_weights, result[2]

python/sglang/kernels/ops/sampling/top_p_renorm_triton.py

@@ -0,0 +1,127 @@
+"""ROCm-compatible top-p probability renormalization fallback."""
+
+from __future__ import annotations
+
+from typing import Union
+
+import torch
+import triton
+import triton.language as tl
+
+_BLOCK_SIZE = 1024
+
+
+@triton.jit
+def _mask_and_partial_sum_kernel(
+    probs_ptr,
+    pivots_ptr,
+    out_ptr,
+    partial_sums_ptr,
+    vocab_size: tl.constexpr,
+    num_chunks: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    chunk = tl.program_id(1)
+    offsets = chunk * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < vocab_size
+    row_offsets = row * vocab_size + offsets
+
+    probs = tl.load(probs_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32)
+    pivot = tl.load(pivots_ptr + row)
+    kept = tl.where(mask & (probs >= pivot), probs, 0.0)
+
+    tl.store(out_ptr + row_offsets, kept, mask=mask)
+    tl.store(partial_sums_ptr + row * num_chunks + chunk, tl.sum(kept, axis=0))
+
+
+@triton.jit
+def _normalize_kernel(
+    out_ptr,
+    row_sums_ptr,
+    numel,
+    vocab_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < numel
+    row = offsets // vocab_size
+    values = tl.load(out_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+    denominator = tl.load(row_sums_ptr + row, mask=mask, other=1.0)
+    tl.store(out_ptr + offsets, values / denominator, mask=mask)
+
+
+def top_p_renorm_probs_triton(
+    probs: torch.Tensor, top_p: Union[torch.Tensor, float]
+) -> torch.Tensor:
+    """Apply exact top-p thresholding and renormalize each probability row.
+
+    Sorting and prefix sums use PyTorch's device kernels because a vocabulary-sized
+    in-register Triton sort does not scale to 100K+ vocabularies. Triton performs
+    the bandwidth-heavy masking, partial reduction, and normalization.
+    """
+    if probs.ndim != 2:
+        raise ValueError(f"probs must be 2D, got shape={tuple(probs.shape)}")
+    if not probs.is_cuda:
+        raise ValueError("top_p_renorm_probs_triton requires a CUDA/HIP tensor")
+
+    probs_fp32 = probs.float().contiguous()
+    batch_size, vocab_size = probs_fp32.shape
+    if batch_size == 0 or vocab_size == 0:
+        return probs_fp32
+
+    if isinstance(top_p, torch.Tensor):
+        top_ps = top_p.to(device=probs.device, dtype=torch.float32).reshape(-1)
+        if top_ps.numel() == 1:
+            top_ps = top_ps.expand(batch_size)
+        elif top_ps.numel() != batch_size:
+            raise ValueError(
+                f"top_p must be scalar or have one value per row, got "
+                f"{top_ps.numel()} values for {batch_size} rows"
+            )
+    else:
+        if not 0.0 < float(top_p) <= 1.0:
+            raise ValueError("top_p values must be in (0, 1]")
+        top_ps = torch.full(
+            (batch_size,), float(top_p), device=probs.device, dtype=torch.float32
+        )
+
+    # Match FlashInfer's threshold semantics: sort ascending, discard the prefix
+    # whose cumulative mass is below 1 - p, and retain all ties at the pivot.
+    sorted_probs = torch.sort(probs_fp32, dim=-1).values
+    cdf = torch.cumsum(sorted_probs, dim=-1)
+    cutoff = torch.searchsorted(cdf, (1.0 - top_ps).unsqueeze(1), right=False).squeeze(
+        1
+    )
+    cutoff.clamp_(max=vocab_size - 1)
+    pivots = sorted_probs.gather(1, cutoff.unsqueeze(1)).squeeze(1).contiguous()
+
+    num_chunks = triton.cdiv(vocab_size, _BLOCK_SIZE)
+    out = torch.empty_like(probs_fp32)
+    partial_sums = torch.empty(
+        (batch_size, num_chunks), device=probs.device, dtype=torch.float32
+    )
+    _mask_and_partial_sum_kernel[(batch_size, num_chunks)](
+        probs_fp32,
+        pivots,
+        out,
+        partial_sums,
+        vocab_size=vocab_size,
+        num_chunks=num_chunks,
+        BLOCK_SIZE=_BLOCK_SIZE,
+        num_warps=8,
+    )
+
+    row_sums = partial_sums.sum(dim=1)
+    _normalize_kernel[(triton.cdiv(out.numel(), _BLOCK_SIZE),)](
+        out,
+        row_sums,
+        out.numel(),
+        vocab_size=vocab_size,
+        BLOCK_SIZE=_BLOCK_SIZE,
+        num_warps=8,
+    )
+    return out
+
+
+__all__ = ["top_p_renorm_probs_triton"]

python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py

@@ -211,6 +211,12 @@ class CustomAllReduceV2:
 
         multicast_ptr = int(symm_mem.multicast_ptr)
         can_multicast = multicast_ptr != 0
+        # multicast VA of the slab base (== the push workspace, at offset 0);
+        # consumed by the K3 all_reduce push kernel
+        self.mc_base_ptr = multicast_ptr if can_multicast else 0
+        # multicast VA of the pull-semaphore region; the K3 pull kernels reuse
+        # these semaphores (same reservation protocol, multicast-signaled)
+        self.pull_sem_mc_ptr = multicast_ptr + pull_sem_offset if can_multicast else 0
         pull_mc_workspace = multicast_ptr + pull_ws_offset if can_multicast else None
         if not can_multicast or cfg.num_mc_blocks is None:
             self.config = self.config._replace(num_mc_blocks=None)

python/sglang/srt/environ.py

@@ -661,6 +661,9 @@ class Envs:
     # Launch the TRT-LLM MoE grouped GEMMs with PDL only at or below this
     # token count.
     SGLANG_TRTLLM_MOE_PDL_MAX_TOKENS = EnvInt(8192)
+    # Unpacked cubin pool for the JIT-built trtllm-gen fused MoE (cubins + flat
+    # ABI headers + overlay/). Unset means the path is unavailable, not empty.
+    SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL = EnvStr(None)
     # SGLang needs to know FlashInfer NVFP4 4over6 config to compute the global scale factor.
     FLASHINFER_NVFP4_4OVER6 = EnvBool(False)
     FLASHINFER_NVFP4_4OVER6_E4M3_USE_256 = EnvBool(False)