Add Inkling model support (#31681)

Co-authored-by: Chunan Zeng <zcnrex@gmail.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
Co-authored-by: Yanbin Jiang <jybsuper@gmail.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: Qiaolin Yu <qiaolin.yu@radixark.ai>
Co-authored-by: Zhichen Zeng <zczeng@uw.edu>
Co-authored-by: Aurick Qiao <aurick@thinkingmachines.ai>
Co-authored-by: Joseph <jk@thinkingmachines.ai>

python/pyproject.toml

@@ -33,6 +33,7 @@ dependencies = [
   "flash-attn-4==4.0.0b15",
   "flashinfer_python[cu13]==0.6.14", # keep it aligned with jit-cache version in Dockerfile
   "gguf",
+  "helion==0.2.6",
   "humming-kernels[cu13]==0.1.10",
   "interegular",
   "IPython",
@@ -42,6 +43,7 @@ dependencies = [
   "modelscope",
   "msgspec",
   "ninja",
+  "numba==0.65.1",
   "numpy",
   "nvidia-cutlass-dsl[cu13]==4.5.2",
   "nvidia-mathdx==25.6.0",
@@ -228,5 +230,11 @@ target = "sglang.srt.grpc._core"
 path = "../rust/sglang-grpc/Cargo.toml"
 binding = "PyO3"
 
+[[tool.setuptools-rust.ext-modules]]
+target = "sglang.srt.multimodal._core"
+path = "../rust/sglang-mm/Cargo.toml"
+binding = "PyO3"
+debug = false
+
 [tool.kernels.dependencies]
 "kernels-community/sgl-flash-attn3" = 1

python/pyproject_other.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=61.0", "setuptools-scm>=8.0", "wheel"]
+requires = ["setuptools>=61.0", "setuptools-rust>=1.10", "setuptools-scm>=8.0", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -30,6 +30,7 @@ runtime_common = [
   "einops",
   "fastapi",
   "gguf",
+  "helion==0.2.6",
   "interegular",
   "IPython",
   "llguidance>=1.7.6,<2.0.0",
@@ -191,6 +192,13 @@ dev_mps = ["sglang[all_mps]", "sglang[test]"]
 [project.scripts]
 sglang = "sglang.cli.main:main"
 
+# Rust-accelerated multimodal preprocessing (sglang.srt.multimodal._core).
+# grpc is intentionally omitted here (it needs proto/tonic); ROCm only builds mm.
+[[tool.setuptools-rust.ext-modules]]
+target = "sglang.srt.multimodal._core"
+path = "../rust/sglang-mm/Cargo.toml"
+binding = "PyO3"
+
 [tool.setuptools.package-data]
 "sglang" = [
   "srt/**/*",

python/sglang/jit_kernel/csrc/inkling/causal_conv1d.cuh

@@ -0,0 +1,217 @@
+// Depthwise causal conv1d (extend/prefill) with the W-1 prefix taps gathered
+// directly from sconv_cache.
+//
+// Semantics:
+// For packed token t in sequence s (bos = cu_seqlens[s], slot = safe_idx[s]) and
+// tap iw in 0..W-1, shifted = t - (W-1) + iw:
+//   shifted >= bos (in-seq history)          -> tap = x[shifted, d]
+//   shifted <  bos, pp=shifted-bos+(W-1)>=0  -> tap = cache[slot, pp, d]
+//                                               (* cache_mask[s] when !IS_DECODE)
+//   else                                     -> tap = 0
+//   out[t,d] = act(sum_iw tap*weight[d,iw]) (+ x[t,d] if residual), fp32 accum.
+// in_x / in_prefix are mutually exclusive, so the fp32 tap sum is bit-identical to
+// the Triton bf16 add (one operand is always 0).
+//
+// Channel-independent control is shared by two channels packed as bf16x2.
+// Each thread keeps a token strip and its prefix window in registers across taps.
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil
+
+#include <sgl_kernel/type.cuh>   // For fp32_t / bf16_t aliases
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+#include <cuda_bf16.h>
+
+namespace {
+
+struct CausalConv1dParams {
+  const void* __restrict__ x;           // [T, D]
+  const void* __restrict__ cache;       // [max_slots, W-1, D]
+  const void* __restrict__ safe_idx;    // int64 [nseq]  cache slot per sequence
+  const void* __restrict__ cache_mask;  // bool  [nseq,1,1] raw metadata
+  const void* __restrict__ weight;      // [D, W]
+  const void* __restrict__ cu;          // int64 [nseq+1] packed sequence starts
+  const void* __restrict__ seq_idx;     // int32 [T]     sequence id per token
+  void* __restrict__ y;                 // [T, D] contiguous output
+  int64_t x_stride_t;
+  int64_t cache_stride_slot;
+  int64_t cache_stride_w;
+  int64_t cache_mask_stride;
+  int64_t weight_stride_d;
+  int64_t y_stride_t;
+  uint32_t T;
+  uint32_t D;
+};
+
+constexpr int kConvBlockT = 4;          // tokens per thread strip
+constexpr uint32_t kConvThreads = 256;  // threads per block (each owns 2 channels)
+
+// blockIdx.x = token strip (BLOCK_T tokens); each thread owns channel pair (c0, c0+1).
+// Requires bf16, D even, and unit channel/row-inner stride (host-checked).
+template <int W, bool USE_SILU, bool USE_RESIDUAL, bool IS_DECODE, typename DType>
+__global__ void causal_conv1d_kernel(const __grid_constant__ CausalConv1dParams p) {
+  constexpr int BT = kConvBlockT;
+  constexpr int WIN = BT + (W - 1);
+
+  __shared__ int s_bos[BT];
+  __shared__ int s_slot[BT];
+  __shared__ float s_m[BT];
+
+  const int T = static_cast<int>(p.T);
+  const int t0 = static_cast<int>(blockIdx.x) * BT;
+
+  if (threadIdx.x < static_cast<uint32_t>(BT)) {
+    const int j = static_cast<int>(threadIdx.x);
+    const int t = t0 + j;
+    if (t < T) {
+      const int seq = static_cast<const int32_t*>(p.seq_idx)[t];
+      s_bos[j] = static_cast<int>(static_cast<const int64_t*>(p.cu)[seq]);
+      s_slot[j] = static_cast<int>(static_cast<const int64_t*>(p.safe_idx)[seq]);
+      if constexpr (!IS_DECODE) {
+        s_m[j] = static_cast<const bool*>(p.cache_mask)[static_cast<int64_t>(seq) * p.cache_mask_stride] ? 1.0f : 0.0f;
+      }
+    }
+  }
+  __syncthreads();
+
+  const int c0 = (blockIdx.y * kConvThreads + threadIdx.x) * 2;  // this thread's channel pair
+  if (c0 >= static_cast<int>(p.D)) return;
+
+  const int sxt = static_cast<int>(p.x_stride_t);
+  const int syt = static_cast<int>(p.y_stride_t);
+  const int swd = static_cast<int>(p.weight_stride_d);
+
+  const auto* xp = static_cast<const __nv_bfloat16*>(p.x);
+  const auto* cp = static_cast<const __nv_bfloat16*>(p.cache);
+  const auto* wp = static_cast<const __nv_bfloat16*>(p.weight);
+  auto* yp = static_cast<__nv_bfloat16*>(p.y);
+
+  // Window (bf16x2 per row), read once into registers.
+  __nv_bfloat162 xr[WIN];
+#pragma unroll
+  for (int i = 0; i < WIN; ++i) {
+    const int row = t0 - (W - 1) + i;
+    xr[i] = (row >= 0 && row < T) ? *reinterpret_cast<const __nv_bfloat162*>(&xp[row * sxt + c0])
+                                  : __float2bfloat162_rn(0.0f);
+  }
+  // Weight taps for the two channels (weight[c0, iw], weight[c0+1, iw]).
+  float2 wv[W];
+#pragma unroll
+  for (int iw = 0; iw < W; ++iw) {
+    wv[iw] = make_float2(__bfloat162float(wp[c0 * swd + iw]), __bfloat162float(wp[(c0 + 1) * swd + iw]));
+  }
+
+#pragma unroll
+  for (int j = 0; j < BT; ++j) {
+    const int t = t0 + j;
+    if (t >= T) break;
+    const int bos = s_bos[j];
+    const float2 x_cur = __bfloat1622float2(xr[j + (W - 1)]);  // tap iw == W-1
+
+    float acc0 = 0.0f, acc1 = 0.0f;
+#pragma unroll
+    for (int iw = 0; iw < W; ++iw) {
+      float2 tap;
+      if (iw == W - 1) {
+        tap = x_cur;
+      } else {
+        const int shifted = t - (W - 1) + iw;  // < T always
+        tap = (shifted >= bos) ? __bfloat1622float2(xr[j + iw]) : make_float2(0.0f, 0.0f);
+        const int prefix_pos = shifted - bos + (W - 1);
+        if (shifted < bos && prefix_pos >= 0 && prefix_pos < (W - 1)) {  // rare: seq start
+          const int64_t coff = static_cast<int64_t>(s_slot[j]) * p.cache_stride_slot +
+                               static_cast<int64_t>(prefix_pos) * p.cache_stride_w + static_cast<int64_t>(c0);
+          float2 pv = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&cp[coff]));
+          if constexpr (!IS_DECODE) {
+            pv.x *= s_m[j];
+            pv.y *= s_m[j];
+          }
+          tap.x += pv.x;
+          tap.y += pv.y;
+        }
+      }
+      acc0 += tap.x * wv[iw].x;
+      acc1 += tap.y * wv[iw].y;
+    }
+
+    if constexpr (USE_SILU) {
+      acc0 = __fdividef(acc0, 1.0f + __expf(-acc0));  // silu = x*sigmoid(x)
+      acc1 = __fdividef(acc1, 1.0f + __expf(-acc1));
+    }
+    if constexpr (USE_RESIDUAL) {
+      acc0 += x_cur.x;
+      acc1 += x_cur.y;
+    }
+    *reinterpret_cast<__nv_bfloat162*>(&yp[t * syt + c0]) = __floats2bfloat162_rn(acc0, acc1);
+  }
+}
+
+template <int W, bool USE_SILU, bool USE_RESIDUAL, bool IS_DECODE, typename DType>
+struct CausalConv1dKernel {
+  static void
+  run(tvm::ffi::TensorView x,
+      tvm::ffi::TensorView cache,
+      tvm::ffi::TensorView safe_idx,
+      tvm::ffi::TensorView cache_mask,
+      tvm::ffi::TensorView weight,
+      tvm::ffi::TensorView cu,
+      tvm::ffi::TensorView seq_idx,
+      tvm::ffi::TensorView y) {
+    using namespace host;
+    auto T = SymbolicSize{"T"};
+    auto D = SymbolicSize{"D"};
+    auto Wd = SymbolicSize{"W"};
+    auto Km1 = SymbolicSize{"W_minus_1"};
+    auto NS = SymbolicSize{"nseq"};
+    auto dev = SymbolicDevice{};
+    dev.set_options<kDLCUDA>();
+    Wd.set_value(W);
+    Km1.set_value(W - 1);
+
+    // x may be a non-contiguous row view (stride_t arbitrary) but must be
+    // channel-contiguous. cache_mask is torch-bool (verify shape/device only).
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(x);
+    TensorMatcher({-1, Km1, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({NS}).with_dtype<int64_t>().with_device(dev).verify(safe_idx);
+    TensorMatcher({NS, 1, 1}).with_device(dev).verify(cache_mask);
+    TensorMatcher({D, Wd}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(weight);
+    TensorMatcher({-1}).with_dtype<int64_t>().with_device(dev).verify(cu);
+    TensorMatcher({T}).with_dtype<int32_t>().with_device(dev).verify(seq_idx);
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(y);
+    RuntimeCheck(cu.size(0) == NS.unwrap() + 1, "cu must have length nseq+1");
+    RuntimeCheck(sizeof(DType) == 2, "causal_conv1d: bf16x2 kernel requires a 16-bit dtype");
+    RuntimeCheck(D.unwrap() % 2 == 0, "causal_conv1d: D must be even for the bf16x2 kernel");
+    RuntimeCheck(cache.stride(2) == 1, "causal_conv1d: sconv_cache must be channel-contiguous");
+
+    const auto params = CausalConv1dParams{
+        .x = x.data_ptr(),
+        .cache = cache.data_ptr(),
+        .safe_idx = safe_idx.data_ptr(),
+        .cache_mask = cache_mask.data_ptr(),
+        .weight = weight.data_ptr(),
+        .cu = cu.data_ptr(),
+        .seq_idx = seq_idx.data_ptr(),
+        .y = y.data_ptr(),
+        .x_stride_t = x.stride(0),
+        .cache_stride_slot = cache.stride(0),
+        .cache_stride_w = cache.stride(1),
+        .cache_mask_stride = cache_mask.stride(0),
+        .weight_stride_d = weight.stride(0),
+        .y_stride_t = y.stride(0),
+        .T = static_cast<uint32_t>(T.unwrap()),
+        .D = static_cast<uint32_t>(D.unwrap()),
+    };
+
+    const uint32_t d_pairs = params.D / 2;
+    const dim3 grid{div_ceil(params.T, static_cast<uint32_t>(kConvBlockT)), div_ceil(d_pairs, kConvThreads)};
+    const dim3 block{kConvThreads};
+    constexpr auto kernel = causal_conv1d_kernel<W, USE_SILU, USE_RESIDUAL, IS_DECODE, DType>;
+    LaunchKernel(grid, block, dev.unwrap())(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/inkling/draft_extend_sconv.cuh

@@ -0,0 +1,147 @@
+// Fused draft-extend convolution-cache update.
+//
+// Speculative draft-extend: for each sequence b (slot ci = cache_indices[b]) the new
+// conv state is the length-W1 window of the "virtual padded" stream
+//     virtual = [ sconv_cache[ci] (W1 rows) ++ hidden[b, 0:T] (T rows) ]
+// starting at num_accepted_tokens[b]:  new[w] = virtual[n_acc + w], w in 0..W1-1
+//     n_acc + w <  W1  -> sconv_cache[ci, n_acc + w]   (initial state)
+//     n_acc + w >= W1  -> hidden[b*T + (n_acc + w - W1)]  (a draft token)
+// written back to sconv_cache[ci]. With tracking, the window at track_step[b] is also
+// written to sconv_cache[mamba_track_indices[b]] wherever crossed[b].
+// Pure copy/select (BIT-EXACT). Init state loaded to registers before writes (RAW-safe);
+// 2 channels/thread as bf16x2. Requires bf16 + even D + channel-contiguous.
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+#include <cuda_bf16.h>
+
+namespace {
+
+struct DraftExtendParams {
+  const void* __restrict__ hidden;         // [B*T, D], channel-contiguous
+  void* __restrict__ cache;                // [pool, W1, D], in-place
+  const void* __restrict__ cache_indices;  // int32 [B]
+  const void* __restrict__ num_accepted;   // int32 [B]
+  const void* __restrict__ crossed;        // bool  [B]   (DO_TRACK only)
+  const void* __restrict__ track_step;     // int32 [B]   (DO_TRACK only)
+  const void* __restrict__ track_indices;  // int64 [B]   (DO_TRACK only)
+  int64_t hs_stride_t;
+  int64_t cache_stride_slot;
+  int64_t cache_stride_w;
+  uint32_t D;
+  uint32_t T;  // draft_token_num
+};
+
+constexpr uint32_t kDEThreads = 256;
+
+template <int W1, bool DO_TRACK, typename DType>
+__global__ void draft_extend_kernel(const __grid_constant__ DraftExtendParams p) {
+  const int b = blockIdx.y;
+  const int c0 = (blockIdx.x * kDEThreads + threadIdx.x) * 2;
+  if (c0 >= static_cast<int>(p.D)) return;
+
+  const int ci = static_cast<const int32_t*>(p.cache_indices)[b];
+  const auto* hp = static_cast<const __nv_bfloat16*>(p.hidden);
+  auto* cp = static_cast<__nv_bfloat16*>(p.cache);
+  const int cw = static_cast<int>(p.cache_stride_w);
+  const int T = static_cast<int>(p.T);
+  const int b_off = b * T;  // hidden row base for this sequence
+  const int64_t src_slot_base = static_cast<int64_t>(ci) * p.cache_stride_slot + c0;
+
+  // Initial state -> registers (RAW-safe against the cache[ci] writes below).
+  __nv_bfloat162 init_reg[W1];
+#pragma unroll
+  for (int w = 0; w < W1; ++w) {
+    init_reg[w] = *reinterpret_cast<const __nv_bfloat162*>(&cp[src_slot_base + static_cast<int64_t>(w) * cw]);
+  }
+
+  // Select the window at `at` from the virtual stream and write it to cache[dst_base].
+  auto emit = [&](int at, int64_t dst_base) {
+#pragma unroll
+    for (int w = 0; w < W1; ++w) {
+      const int pos = at + w;
+      __nv_bfloat162 v;
+      if (pos < W1) {
+        v = init_reg[0];
+#pragma unroll
+        for (int src = 0; src < W1; ++src) {
+          if (src == pos) v = init_reg[src];
+        }
+      } else {
+        const int row = b_off + (pos - W1);
+        v = *reinterpret_cast<const __nv_bfloat162*>(&hp[static_cast<int64_t>(row) * p.hs_stride_t + c0]);
+      }
+      *reinterpret_cast<__nv_bfloat162*>(&cp[dst_base + static_cast<int64_t>(w) * cw]) = v;
+    }
+  };
+
+  const int n_acc = static_cast<const int32_t*>(p.num_accepted)[b];
+  if constexpr (DO_TRACK) {
+    // Track window first (reads init_reg, distinct dst slot) then the main window.
+    if (static_cast<const bool*>(p.crossed)[b]) {
+      const int tstep = static_cast<const int32_t*>(p.track_step)[b];
+      const int64_t tslot = static_cast<const int64_t*>(p.track_indices)[b];
+      emit(tstep, tslot * p.cache_stride_slot + c0);
+    }
+  }
+  emit(n_acc, src_slot_base);
+}
+
+template <int W1, bool DO_TRACK, typename DType>
+struct DraftExtendSconvKernel {
+  static void
+  run(tvm::ffi::TensorView hidden,
+      tvm::ffi::TensorView cache,
+      tvm::ffi::TensorView cache_indices,
+      tvm::ffi::TensorView num_accepted,
+      int64_t draft_token_num,
+      tvm::ffi::TensorView crossed,
+      tvm::ffi::TensorView track_step,
+      tvm::ffi::TensorView track_indices) {
+    using namespace host;
+    auto BT = SymbolicSize{"B_times_T"};
+    auto D = SymbolicSize{"D"};
+    auto W1s = SymbolicSize{"W_minus_1"};
+    auto B = SymbolicSize{"B"};
+    auto dev = SymbolicDevice{};
+    dev.set_options<kDLCUDA>();
+    W1s.set_value(W1);
+
+    TensorMatcher({BT, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(hidden);
+    TensorMatcher({-1, W1s, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({B}).with_dtype<int32_t>().with_device(dev).verify(cache_indices);
+    TensorMatcher({B}).with_dtype<int32_t>().with_device(dev).verify(num_accepted);
+    RuntimeCheck(sizeof(DType) == 2, "draft_extend: bf16x2 kernel requires 16-bit dtype");
+    RuntimeCheck(D.unwrap() % 2 == 0, "draft_extend: D must be even for the bf16x2 kernel");
+    RuntimeCheck(cache.stride(2) == 1, "draft_extend: cache must be channel-contiguous");
+
+    const auto params = DraftExtendParams{
+        .hidden = hidden.data_ptr(),
+        .cache = cache.data_ptr(),
+        .cache_indices = cache_indices.data_ptr(),
+        .num_accepted = num_accepted.data_ptr(),
+        .crossed = DO_TRACK ? crossed.data_ptr() : nullptr,
+        .track_step = DO_TRACK ? track_step.data_ptr() : nullptr,
+        .track_indices = DO_TRACK ? track_indices.data_ptr() : nullptr,
+        .hs_stride_t = hidden.stride(0),
+        .cache_stride_slot = cache.stride(0),
+        .cache_stride_w = cache.stride(1),
+        .D = static_cast<uint32_t>(D.unwrap()),
+        .T = static_cast<uint32_t>(draft_token_num),
+    };
+
+    const uint32_t d_pairs = params.D / 2;
+    const dim3 grid{div_ceil(d_pairs, kDEThreads), static_cast<uint32_t>(B.unwrap())};
+    const dim3 block{kDEThreads};
+    constexpr auto kernel = draft_extend_kernel<W1, DO_TRACK, DType>;
+    LaunchKernel(grid, block, dev.unwrap())(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/inkling/fused_decode_update.cuh

@@ -0,0 +1,187 @@
+// Fused decode causal_conv1d, cache shift-update, and optional track copy.
+//
+// Decode: each token t is its own sequence (bos=t). Per token:
+//   conv: acc = sum_{iw<W-1} cache[slot, iw, d]*cache_mask[t] * weight[d, iw]
+//              + x[t, d] * weight[d, W-1]
+//         y[t,d] = act(acc) (+ x[t,d] if residual).
+//   update (valid lanes, ci != PAD): shift the state left, append current token --
+//     new[iw] = cache[slot, iw+1]*cache_mask[t] (iw < W-2);  new[W-2] = x[t].
+//   track (DO_TRACK): the same post-update window is also written to
+//     cache[track_indices[t]] wherever track_mask[t] (prefix-cache ping-pong slot).
+// Working slots and ping-pong track slots are pairwise-distinct, so writes never race.
+// Cache history is loaded to registers BEFORE any write (RAW-safe). 2 channels/thread
+// as bf16x2; conv accumulates in fp32 (matches the fp32 reference), update is a
+// bit-exact bf16 move. Requires bf16 + even D + channel-contiguous cache/x/y.
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+#include <cuda_bf16.h>
+
+namespace {
+
+struct DecodeUpdateParams {
+  const void* __restrict__ x;              // [T, D], channel-contiguous
+  void* __restrict__ cache;                // [pool, W-1, D], in-place update
+  const void* __restrict__ cache_indices;  // int32 [T]  (PAD == -1)
+  const void* __restrict__ cache_mask;     // bool  [T]
+  const void* __restrict__ weight;         // [D, W]
+  void* __restrict__ y;                    // [T, D] contiguous output
+  const void* __restrict__ track_mask;     // bool  [T]   (DO_TRACK only)
+  const void* __restrict__ track_indices;  // int64 [T]   (DO_TRACK only)
+  int64_t x_stride_t;
+  int64_t cache_stride_slot;
+  int64_t cache_stride_w;
+  int64_t weight_stride_d;
+  int64_t y_stride_t;
+  int64_t track_idx_stride;
+  uint32_t D;
+};
+
+constexpr uint32_t kDecThreads = 256;
+constexpr int kPadSlot = -1;
+
+template <int W, bool USE_SILU, bool USE_RESIDUAL, bool DO_TRACK, typename DType>
+__global__ void fused_decode_update_kernel(const __grid_constant__ DecodeUpdateParams p) {
+  constexpr int W1 = W - 1;  // number of cached history taps / conv-state rows
+  const int t = blockIdx.y;
+  const int ci = static_cast<const int32_t*>(p.cache_indices)[t];
+  const bool valid = ci != kPadSlot;
+  const int slot = valid ? ci : 0;  // clamp: PAD lanes still emit y (discarded), no cache write
+
+  const int c0 = (blockIdx.x * kDecThreads + threadIdx.x) * 2;
+  if (c0 >= static_cast<int>(p.D)) return;
+
+  const float cm = static_cast<const bool*>(p.cache_mask)[t] ? 1.0f : 0.0f;
+  const auto* xp = static_cast<const __nv_bfloat16*>(p.x);
+  const auto* wp = static_cast<const __nv_bfloat16*>(p.weight);
+  auto* cp = static_cast<__nv_bfloat16*>(p.cache);
+  auto* yp = static_cast<__nv_bfloat16*>(p.y);
+  const int cw = static_cast<int>(p.cache_stride_w);
+  const int swd = static_cast<int>(p.weight_stride_d);
+  const int64_t cache_base = static_cast<int64_t>(slot) * p.cache_stride_slot + c0;
+
+  // History taps -> registers (RAW-safe against the update writes below).
+  __nv_bfloat162 hist[W1];
+#pragma unroll
+  for (int w = 0; w < W1; ++w) {
+    hist[w] = *reinterpret_cast<const __nv_bfloat162*>(&cp[cache_base + static_cast<int64_t>(w) * cw]);
+  }
+  const __nv_bfloat162 xv = *reinterpret_cast<const __nv_bfloat162*>(&xp[static_cast<int64_t>(t) * p.x_stride_t + c0]);
+  const float2 xf = __bfloat1622float2(xv);
+
+  float2 wv[W];
+#pragma unroll
+  for (int iw = 0; iw < W; ++iw) {
+    wv[iw] = make_float2(__bfloat162float(wp[c0 * swd + iw]), __bfloat162float(wp[(c0 + 1) * swd + iw]));
+  }
+
+  // ---- conv (fp32 accum): W-1 cached taps (gated by cache_mask) + current token ----
+  float acc0 = 0.0f, acc1 = 0.0f;
+#pragma unroll
+  for (int iw = 0; iw < W1; ++iw) {
+    const float2 h = __bfloat1622float2(hist[iw]);
+    acc0 += h.x * cm * wv[iw].x;
+    acc1 += h.y * cm * wv[iw].y;
+  }
+  acc0 += xf.x * wv[W1].x;
+  acc1 += xf.y * wv[W1].y;
+  if constexpr (USE_SILU) {
+    acc0 = __fdividef(acc0, 1.0f + __expf(-acc0));
+    acc1 = __fdividef(acc1, 1.0f + __expf(-acc1));
+  }
+  if constexpr (USE_RESIDUAL) {
+    acc0 += xf.x;
+    acc1 += xf.y;
+  }
+  *reinterpret_cast<__nv_bfloat162*>(&yp[static_cast<int64_t>(t) * p.y_stride_t + c0]) =
+      __floats2bfloat162_rn(acc0, acc1);
+
+  if (!valid) return;
+
+  // ---- update: shift state left (gated by cache_mask), append current token ----
+  const __nv_bfloat162 zero = __float2bfloat162_rn(0.0f);
+  int64_t track_base = 0;
+  bool do_tr = false;
+  if constexpr (DO_TRACK) {
+    do_tr = static_cast<const bool*>(p.track_mask)[t];
+    if (do_tr) {
+      const int64_t tslot = static_cast<const int64_t*>(p.track_indices)[static_cast<int64_t>(t) * p.track_idx_stride];
+      track_base = tslot * p.cache_stride_slot + c0;
+    }
+  }
+#pragma unroll
+  for (int iw = 0; iw < W1; ++iw) {
+    const __nv_bfloat162 nv = (iw < W1 - 1) ? ((cm != 0.0f) ? hist[iw + 1] : zero) : xv;
+    *reinterpret_cast<__nv_bfloat162*>(&cp[cache_base + static_cast<int64_t>(iw) * cw]) = nv;
+    if constexpr (DO_TRACK) {
+      if (do_tr) {
+        *reinterpret_cast<__nv_bfloat162*>(&cp[track_base + static_cast<int64_t>(iw) * cw]) = nv;
+      }
+    }
+  }
+}
+
+template <int W, bool USE_SILU, bool USE_RESIDUAL, bool DO_TRACK, typename DType>
+struct FusedDecodeUpdateKernel {
+  static void
+  run(tvm::ffi::TensorView x,
+      tvm::ffi::TensorView cache,
+      tvm::ffi::TensorView cache_indices,
+      tvm::ffi::TensorView cache_mask,
+      tvm::ffi::TensorView weight,
+      tvm::ffi::TensorView y,
+      tvm::ffi::TensorView track_mask,
+      tvm::ffi::TensorView track_indices) {
+    using namespace host;
+    auto T = SymbolicSize{"T"};
+    auto D = SymbolicSize{"D"};
+    auto Wd = SymbolicSize{"W"};
+    auto W1s = SymbolicSize{"W_minus_1"};
+    auto dev = SymbolicDevice{};
+    dev.set_options<kDLCUDA>();
+    Wd.set_value(W);
+    W1s.set_value(W - 1);
+
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(x);
+    TensorMatcher({-1, W1s, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({T}).with_dtype<int32_t>().with_device(dev).verify(cache_indices);
+    TensorMatcher({T}).with_device(dev).verify(cache_mask);
+    TensorMatcher({D, Wd}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(weight);
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(y);
+    RuntimeCheck(sizeof(DType) == 2, "fused_decode: bf16x2 kernel requires 16-bit dtype");
+    RuntimeCheck(D.unwrap() % 2 == 0, "fused_decode: D must be even for the bf16x2 kernel");
+    RuntimeCheck(cache.stride(2) == 1, "fused_decode: cache must be channel-contiguous");
+
+    const auto params = DecodeUpdateParams{
+        .x = x.data_ptr(),
+        .cache = cache.data_ptr(),
+        .cache_indices = cache_indices.data_ptr(),
+        .cache_mask = cache_mask.data_ptr(),
+        .weight = weight.data_ptr(),
+        .y = y.data_ptr(),
+        .track_mask = DO_TRACK ? track_mask.data_ptr() : nullptr,
+        .track_indices = DO_TRACK ? track_indices.data_ptr() : nullptr,
+        .x_stride_t = x.stride(0),
+        .cache_stride_slot = cache.stride(0),
+        .cache_stride_w = cache.stride(1),
+        .weight_stride_d = weight.stride(0),
+        .y_stride_t = y.stride(0),
+        .track_idx_stride = DO_TRACK ? track_indices.stride(0) : 0,
+        .D = static_cast<uint32_t>(D.unwrap()),
+    };
+
+    const uint32_t d_pairs = params.D / 2;
+    const dim3 grid{div_ceil(d_pairs, kDecThreads), static_cast<uint32_t>(T.unwrap())};
+    const dim3 block{kDecThreads};
+    constexpr auto kernel = fused_decode_update_kernel<W, USE_SILU, USE_RESIDUAL, DO_TRACK, DType>;
+    LaunchKernel(grid, block, dev.unwrap())(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/inkling/gather_scatter_sconv.cuh

@@ -0,0 +1,109 @@
+// Fused gather and scatter into sconv_cache.
+//
+// For each batch element b where mask[b] is true, copy the W1 = W-1 token rows
+// hidden_states[track_idx[b, w]]  ->  sconv_cache[dst[b], w]  (w = 0..W1-1).
+// Masked-out lanes are left untouched. Pure copy (no arithmetic) => BIT-EXACT.
+// 2 channels/thread packed as bf16x2. Requires bf16 + even D + channel-contiguous.
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+#include <cuda_bf16.h>
+
+namespace {
+
+struct GatherScatterParams {
+  const void* __restrict__ hidden;     // [T, D], channel-contiguous
+  void* __restrict__ cache;            // [pool, W1, D], in-place scatter target
+  const void* __restrict__ track_idx;  // int32 [B, W1]
+  const void* __restrict__ mask;       // bool  [B]
+  const void* __restrict__ dst;        // int64 [B]
+  int64_t hs_stride_t;
+  int64_t cache_stride_slot;
+  int64_t cache_stride_w;
+  int64_t track_stride_b;
+  int64_t track_stride_w;
+  int64_t dst_stride_b;
+  uint32_t D;
+};
+
+constexpr uint32_t kGSThreads = 256;
+
+template <int W1, typename DType>
+__global__ void gather_scatter_kernel(const __grid_constant__ GatherScatterParams p) {
+  const int b = blockIdx.y;
+  if (!static_cast<const bool*>(p.mask)[b]) return;  // masked-out lane: untouched
+
+  const int c0 = (blockIdx.x * kGSThreads + threadIdx.x) * 2;
+  if (c0 >= static_cast<int>(p.D)) return;
+
+  const auto* hp = static_cast<const __nv_bfloat16*>(p.hidden);
+  auto* cp = static_cast<__nv_bfloat16*>(p.cache);
+  const int64_t dst_slot = static_cast<const int64_t*>(p.dst)[static_cast<int64_t>(b) * p.dst_stride_b];
+  const int64_t cache_base = dst_slot * p.cache_stride_slot + c0;
+  const int64_t track_base = static_cast<int64_t>(b) * p.track_stride_b;
+
+#pragma unroll
+  for (int w = 0; w < W1; ++w) {
+    const int64_t src_t =
+        static_cast<const int32_t*>(p.track_idx)[track_base + static_cast<int64_t>(w) * p.track_stride_w];
+    const __nv_bfloat162 v = *reinterpret_cast<const __nv_bfloat162*>(&hp[src_t * p.hs_stride_t + c0]);
+    *reinterpret_cast<__nv_bfloat162*>(&cp[cache_base + static_cast<int64_t>(w) * p.cache_stride_w]) = v;
+  }
+}
+
+template <int W1, typename DType>
+struct GatherScatterSconvKernel {
+  static void
+  run(tvm::ffi::TensorView hidden,
+      tvm::ffi::TensorView cache,
+      tvm::ffi::TensorView track_idx,
+      tvm::ffi::TensorView mask,
+      tvm::ffi::TensorView dst) {
+    using namespace host;
+    auto T = SymbolicSize{"T"};
+    auto D = SymbolicSize{"D"};
+    auto W1s = SymbolicSize{"W_minus_1"};
+    auto B = SymbolicSize{"B"};
+    auto dev = SymbolicDevice{};
+    dev.set_options<kDLCUDA>();
+    W1s.set_value(W1);
+
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(hidden);
+    TensorMatcher({-1, W1s, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({B, W1s}).with_dtype<int32_t>().with_device(dev).verify(track_idx);
+    TensorMatcher({B}).with_device(dev).verify(mask);
+    TensorMatcher({B}).with_dtype<int64_t>().with_device(dev).verify(dst);
+    RuntimeCheck(sizeof(DType) == 2, "gather_scatter: bf16x2 kernel requires 16-bit dtype");
+    RuntimeCheck(D.unwrap() % 2 == 0, "gather_scatter: D must be even for the bf16x2 kernel");
+    RuntimeCheck(cache.stride(2) == 1, "gather_scatter: cache must be channel-contiguous");
+
+    const auto params = GatherScatterParams{
+        .hidden = hidden.data_ptr(),
+        .cache = cache.data_ptr(),
+        .track_idx = track_idx.data_ptr(),
+        .mask = mask.data_ptr(),
+        .dst = dst.data_ptr(),
+        .hs_stride_t = hidden.stride(0),
+        .cache_stride_slot = cache.stride(0),
+        .cache_stride_w = cache.stride(1),
+        .track_stride_b = track_idx.stride(0),
+        .track_stride_w = track_idx.stride(1),
+        .dst_stride_b = dst.stride(0),
+        .D = static_cast<uint32_t>(D.unwrap()),
+    };
+
+    const uint32_t d_pairs = params.D / 2;
+    const dim3 grid{div_ceil(d_pairs, kGSThreads), static_cast<uint32_t>(B.unwrap())};
+    const dim3 block{kGSThreads};
+    constexpr auto kernel = gather_scatter_kernel<W1, DType>;
+    LaunchKernel(grid, block, dev.unwrap())(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/inkling/inkling_all_reduce.cuh

@@ -0,0 +1,672 @@
+// Two-shot (reduce-scatter + all-gather) all-reduce over a torch
+// symmetric-memory buffer.
+//
+// It operates IN PLACE on the peer symm buffers: the producer (e.g. the wo_ud /
+// MoE-combine GEMM) writes its local shard straight into THIS rank's symm buffer
+// (via get_ar_buffer), so there is no stage-in copy; the reduced result is left
+// in the buffer and handed back to Python as a view, so there is no copy-out.
+//
+// Correctness (two-shot is race-safe in place): rank r owns the disjoint vec
+// slice [local_vec_start, local_vec_finish); it reads every peer's slice, sums,
+// and broadcasts the sum back to every peer's slice. Only rank r ever writes
+// slice S_r (in any buffer), so there is no write-write conflict, and each
+// per-element load completes before its store (data dependency).
+//
+// Two variants:
+//   * ..._kernel  (v1): no in-kernel sync; the caller fences with the symm-mem
+//     handle's barrier() on each side (3 launches total).
+//   * ..._fused_kernel (v2): an in-kernel per-block system barrier (entry:
+//     producers done + visible; exit: broadcasts done + visible), so the whole
+//     all-reduce is a single launch. The barrier uses a DEDICATED symmetric
+//     flags buffer (independent of torch's signal pad, so no interference with
+//     multimem) and a device-resident monotonic epoch counter per block, which
+//     keeps advancing across launches -- including CUDA-graph replays -- so
+//     flags never go stale (spin is `flag < epoch`, epoch strictly increasing).
+//
+// Fusion seam: the reduced `result` Storage below is where an epilogue (RMSNorm
+// / short-conv / bias) plugs in -- applied in registers before the broadcast
+// store, so the normed/conv'd result never makes an extra HBM round trip.
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <dlpack/dlpack.h>
+
+#include "inkling_ar_barrier.cuh"
+#include <bit>
+#include <cstdint>
+#include <mutex>
+#include <type_traits>
+#include <unordered_map>
+
+namespace {
+
+template <typename DType, uint32_t kNumGPU>
+struct InklingAllReduceTrait {
+  static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  static constexpr uint32_t kElemsPerVec = kVecSize * 2;
+  using DType2 = packed_t<DType>;
+  using Storage = device::AlignedVector<DType2, kVecSize>;
+  static_assert(sizeof(Storage) == 16 && alignof(Storage) == 16, "Storage must be 16B");
+  static_assert(std::has_single_bit(kNumGPU), "kNumGPU must be a power of two");
+};
+
+// Register-level fused add of two vecs (fp32 math, ONE round to DType) -- the
+// exact numerics of torch.add on two bf16 tensors, so fusing the shared-expert
+// partials stays bit-identical to the unfused {torch.add -> AR} chain.
+template <typename DType>
+__device__ __forceinline__ typename InklingAllReduceTrait<DType, 2>::Storage add_vec_rn(
+    const typename InklingAllReduceTrait<DType, 2>::Storage& a,
+    const typename InklingAllReduceTrait<DType, 2>::Storage& b) {
+  using namespace device;
+  using Trait = InklingAllReduceTrait<DType, 2>;  // kNumGPU-independent
+  using DType2 = typename Trait::DType2;
+  typename Trait::Storage out;
+#pragma unroll
+  for (uint32_t j = 0; j < Trait::kVecSize; ++j) {
+    const fp32x2_t x = cast<fp32x2_t>(a[j]);
+    const fp32x2_t y = cast<fp32x2_t>(b[j]);
+    fp32x2_t s;
+    s.x = x.x + y.x;
+    s.y = x.y + y.y;
+    out[j] = cast<DType2>(s);
+  }
+  return out;
+}
+
+// Fused-shared PROLOGUE for the pull-based kernels (v2/v3/v3b/v4): fold this
+// rank's LOCAL shared-expert partials into its own symm input region before
+// the ENTRY barrier, so every peer's ld_reduce / peer-read sums
+// (routed_r + shared_r) across ranks. The entry barrier must then run in
+// publish mode (grid_system_barrier, publish_writes=true): these are in-kernel
+// stores by ALL CTAs, not prior-kernel stores, so each CTA has to
+// system-publish them before the leader's release. (The per-block barrier
+// cannot order this: block b's fold range is not the range peer block b
+// reads.) The push-based kernels (v5 & the fused decode family) instead fold
+// in registers at the push -- see the shared branch in the push loop.
+template <typename DType, uint32_t kNumGPU>
+__device__ __forceinline__ void
+fold_shared_local(DType* __restrict__ buf, const DType* __restrict__ shared, uint32_t num_items) {
+  using Trait = InklingAllReduceTrait<DType, kNumGPU>;
+  using Storage = typename Trait::Storage;
+  const uint32_t total_vec = num_items / Trait::kElemsPerVec;
+  const uint32_t stride = gridDim.x * blockDim.x;
+  for (uint32_t v = blockIdx.x * blockDim.x + threadIdx.x; v < total_vec; v += stride) {
+    Storage a, b;
+    a.load(buf, v);
+    b.load(shared, v);
+    add_vec_rn<DType>(a, b).store(buf, v);
+  }
+}
+
+// Two-shot partition: contiguous, warp-aligned vec slice per rank. Returns
+// {start, count} in vec units (empty for trailing ranks when the range is small).
+template <typename DType, uint32_t kNumGPU>
+__device__ __forceinline__ uint2 rank_vec_slice(uint32_t rank, uint32_t num_items) {
+  using namespace device;
+  using Trait = InklingAllReduceTrait<DType, kNumGPU>;
+  const uint32_t total_vec = num_items / Trait::kElemsPerVec;
+  const uint32_t vec_per_rank = div_ceil(div_ceil(total_vec, kNumGPU), kWarpThreads) * kWarpThreads;
+  const uint32_t start = min(rank * vec_per_rank, total_vec);
+  const uint32_t finish = min(start + vec_per_rank, total_vec);
+  return {start, finish - start};
+}
+
+// Offset each peer pointer to this rank's slice, and return the slice's local
+// vec count.
+template <typename DType, uint32_t kNumGPU>
+__device__ __forceinline__ uint32_t
+slice_setup(DType* (&input)[kNumGPU], void* const* peer_ptrs, uint32_t rank, uint32_t num_items) {
+  using Trait = InklingAllReduceTrait<DType, kNumGPU>;
+  const uint2 slice = rank_vec_slice<DType, kNumGPU>(rank, num_items);
+  const uint32_t base = slice.x * Trait::kElemsPerVec;
+#pragma unroll
+  for (uint32_t i = 0; i < kNumGPU; ++i)
+    input[i] = static_cast<DType*>(peer_ptrs[i]) + base;
+  return slice.y;  // local vec count
+}
+
+template <typename DType, uint32_t kNumGPU>
+__device__ __forceinline__ void two_shot_reduce_local(DType* (&input)[kNumGPU], uint32_t local_vecs) {
+  using namespace device;
+  using Trait = InklingAllReduceTrait<DType, kNumGPU>;
+  using Storage = typename Trait::Storage;
+  using DType2 = typename Trait::DType2;
+  constexpr uint32_t kVecSize = Trait::kVecSize;
+  const uint32_t stride = gridDim.x * blockDim.x;
+  for (uint32_t v = blockIdx.x * blockDim.x + threadIdx.x; v < local_vecs; v += stride) {
+    Storage s[kNumGPU];
+#pragma unroll
+    for (uint32_t i = 0; i < kNumGPU; ++i)
+      s[i].load(input[i], v);
+    Storage result;
+#pragma unroll
+    for (uint32_t j = 0; j < kVecSize; ++j) {
+      fp32x2_t acc = cast<fp32x2_t>(s[0][j]);
+#pragma unroll
+      for (uint32_t i = 1; i < kNumGPU; ++i) {
+        const fp32x2_t x = cast<fp32x2_t>(s[i][j]);
+        acc.x += x.x;
+        acc.y += x.y;
+      }
+      result[j] = cast<DType2>(acc);  // <-- EPILOGUE SEAM
+    }
+#pragma unroll
+    for (uint32_t i = 0; i < kNumGPU; ++i)
+      result.store(input[i], v);
+  }
+}
+
+// v1: no in-kernel barrier (caller fences via hdl.barrier()).
+template <typename DType, uint32_t kNumGPU>
+__global__ __launch_bounds__(1024, 1) void inkling_two_shot_all_reduce_kernel(
+    void* const* __restrict__ peer_ptrs, const uint32_t rank, const uint32_t num_items) {
+  DType* input[kNumGPU];
+  const uint32_t local_vecs = slice_setup<DType, kNumGPU>(input, peer_ptrs, rank, num_items);
+  two_shot_reduce_local<DType, kNumGPU>(input, local_vecs);
+}
+
+// v2: single-launch, fused entry + exit system barrier. `shared` (optional):
+// this rank's LOCAL shared-expert partials, folded into its own buffer before
+// the entry barrier (which then must publish -- see fold_shared_local).
+template <typename DType, uint32_t kNumGPU>
+__global__ __launch_bounds__(1024, 1) void inkling_two_shot_all_reduce_fused_kernel(
+    void* const* __restrict__ peer_ptrs,
+    void* const* __restrict__ flag_ptrs,
+    uint32_t* __restrict__ state,
+    const DType* __restrict__ shared,
+    const uint32_t rank,
+    const uint32_t num_items) {
+  DType* input[kNumGPU];
+  const uint32_t local_vecs = slice_setup<DType, kNumGPU>(input, peer_ptrs, rank, num_items);
+  if (shared != nullptr) {
+    fold_shared_local<DType, kNumGPU>(static_cast<DType*>(peer_ptrs[rank]), shared, num_items);
+  }
+  // ENTRY: producers done + visible (publish the fold's in-kernel stores too).
+  inkling_ar::grid_system_barrier<kNumGPU>(state, flag_ptrs, rank, 0, /*publish_writes=*/shared != nullptr);
+  two_shot_reduce_local<DType, kNumGPU>(input, local_vecs);
+  inkling_ar::grid_system_barrier<kNumGPU>(
+      state, flag_ptrs, rank, 1, /*publish_writes=*/true);  // EXIT: broadcasts done + visible
+}
+
+// Multimem one-shot all-reduce: uses the NVLink multicast ld_reduce/st hardware
+// instructions on the symm buffer's multicast pointer -- the same in-switch
+// reduce torch's multimem_all_reduce_ uses -- so it matches multimem for the
+// tiny, latency-bound decode messages where two-shot's N peer reads lose. Reduce
+// is one transaction (hardware sums all GPUs); scatter partition keeps the store
+// traffic minimal. bf16-only (multimem.add supports .bf16x2 on sm90/sm100).
+// kPerBlockBarrier swaps both barriers for block_system_barrier (per-block
+// peer handshake, no grid funnel). Correct for the two-shot too: any peer
+// block's ENTRY signal proves that peer's producer kernel completed (kernel
+// serialization on its stream), and kernel end is a grid-wide join, so my
+// per-block EXIT waits compose into "every peer block's broadcasts done"
+// before my consumer can run. The two calls share the per-block epoch slot
+// (it just advances twice per launch).
+template <typename DType, uint32_t kNumGPU, bool kPerBlockBarrier>
+__global__ __launch_bounds__(1024, 1) void inkling_multimem_one_shot_fused_kernel(
+    DType* __restrict__ mc_ptr,     // multicast base pointer (covers all peers)
+    DType* __restrict__ local_ptr,  // this rank's LOCAL base of the same buffer
+    void* const* __restrict__ flag_ptrs,
+    uint32_t* __restrict__ state,
+    const DType* __restrict__ shared,  // optional LOCAL shared-expert partials
+    const uint32_t rank,
+    const uint32_t num_items) {
+  using namespace device;
+  using Trait = InklingAllReduceTrait<DType, kNumGPU>;
+  static_assert(std::is_same_v<DType, bf16_t>, "multimem.add path is bf16-only");
+  constexpr uint32_t kElemsPerVec = Trait::kElemsPerVec;  // 8 bf16 = 16 B
+
+  const uint2 slice = rank_vec_slice<DType, kNumGPU>(rank, num_items);
+  const uint32_t local_vecs = slice.y;
+  DType* mc = mc_ptr + slice.x * kElemsPerVec;
+
+  if (shared != nullptr) {
+    // Fold covers the FULL range while each peer ld_reduces only its slice, so
+    // the per-block handshake cannot order it -- use the publishing grid
+    // barrier for entry even in v3b (exit stays per-block).
+    fold_shared_local<DType, kNumGPU>(local_ptr, shared, num_items);
+    inkling_ar::grid_system_barrier<kNumGPU>(state, flag_ptrs, rank, 0, /*publish_writes=*/true);
+  } else if constexpr (kPerBlockBarrier) {
+    inkling_ar::block_system_barrier<kNumGPU>(state, flag_ptrs, rank);  // ENTRY
+  } else {
+    inkling_ar::grid_system_barrier<kNumGPU>(
+        state, flag_ptrs, rank, 0, /*publish_writes=*/false);  // ENTRY: producers done + visible
+  }
+  const uint32_t stride = gridDim.x * blockDim.x;
+  for (uint32_t v = blockIdx.x * blockDim.x + threadIdx.x; v < local_vecs; v += stride) {
+    DType* addr = mc + v * kElemsPerVec;  // 16 B, 16-B aligned
+    uint32_t r0, r1, r2, r3;
+    // hardware reduce across all GPUs mapped to the multicast region.
+    asm volatile("multimem.ld_reduce.relaxed.sys.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];"
+                 : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
+                 : "l"(addr));
+    // <-- EPILOGUE SEAM (norm / sconv / bias on {r0..r3} before broadcast)
+    // broadcast the reduced slice to every GPU.
+    asm volatile(
+        "multimem.st.relaxed.sys.global.v4.bf16x2 [%0], {%1,%2,%3,%4};" ::"l"(addr), "r"(r0), "r"(r1), "r"(r2), "r"(r3)
+        : "memory");
+  }
+  if constexpr (kPerBlockBarrier) {
+    inkling_ar::block_system_barrier<kNumGPU>(state, flag_ptrs, rank);  // EXIT (release signal publishes)
+  } else {
+    inkling_ar::grid_system_barrier<kNumGPU>(
+        state, flag_ptrs, rank, 1, /*publish_writes=*/true);  // EXIT: broadcasts done + visible
+  }
+}
+
+// One-shot PUSH all-reduce (v5): each rank multicast-STORES its full input into
+// its per-rank slot of a symmetric staging area (the NVSwitch replicates the
+// slot to every GPU), ONE grid barrier waits for all pushes to land, then each
+// rank reduces the N staged shards LOCALLY (fp32 accum) into a LOCAL output.
+// Single barrier total: the push needs no entry barrier (it publishes only this
+// rank's own producer data, stream-ordered locally) and the local output needs
+// no exit barrier. Staging reuse is caller-managed (A/B rotation, like v4's
+// input; the next AR's barrier proves peers consumed the old buffer).
+//
+// vs v3/mm (two-shot): drops one full cross-GPU barrier round trip -- wins the
+// latency-bound band. vs v4 (full one-shot ld_reduce): switch REPLICATION is
+// cheap where the switch's reduce engine serializes N redundant full-range
+// reduces, so this scales past v4's 2-row ceiling. Fabric cost: n egress,
+// (N-1)*n ingress per GPU; local HBM/L2: N*n read + n write. Like v4, each rank
+// holds the FULL row at the epilogue seam (natural RMSNorm-fusion base).
+// bf16-only (multimem.st .bf16x2).
+//
+// kPerBlockBarrier selects block_system_barrier (per-block peer handshake, no
+// grid funnel -- the multi-block latency winner) over the single-leader grid
+// barrier. Safe here because the reduce loop reads exactly the vec ranges the
+// blockIdx-matched pushes wrote.
+template <typename DType, uint32_t kNumGPU, bool kPerBlockBarrier>
+__global__ __launch_bounds__(1024, 1) void inkling_multimem_push_oneshot_kernel(
+    const DType* __restrict__ in_ptr,     // local input (producer's partial sums)
+    DType* __restrict__ mc_stage_ptr,     // multicast staging base (slot r at r*num_items)
+    const DType* __restrict__ stage_ptr,  // this GPU's LOCAL view of the staging base
+    DType* __restrict__ out_ptr,          // local output
+    void* const* __restrict__ flag_ptrs,
+    uint32_t* __restrict__ state,
+    const DType* __restrict__ shared,  // optional LOCAL shared-expert partials
+    const uint32_t rank,
+    const uint32_t num_items) {
+  using namespace device;
+  using Trait = InklingAllReduceTrait<DType, kNumGPU>;
+  static_assert(std::is_same_v<DType, bf16_t>, "multimem path is bf16-only");
+  constexpr uint32_t kElemsPerVec = Trait::kElemsPerVec;  // 8 bf16 = 16 B
+  const uint32_t total_vec = num_items / kElemsPerVec;
+  const uint32_t stride = gridDim.x * blockDim.x;
+
+  // Phase 1: push. One multicast store per vec; the switch fans it out to every
+  // GPU's replica of slot `rank` (including our own). With `shared`, the
+  // shared-expert partials fold into the pushed value in registers (fp32 add,
+  // one bf16 round -- torch.add numerics) at ZERO extra fabric or HBM traffic;
+  // both barrier flavors stay valid because push/reduce mappings are unchanged.
+  DType* slot = mc_stage_ptr + rank * num_items;
+  if (shared != nullptr) {
+    using Storage = typename Trait::Storage;
+    for (uint32_t v = blockIdx.x * blockDim.x + threadIdx.x; v < total_vec; v += stride) {
+      Storage a, b;
+      a.load(in_ptr, v);
+      b.load(shared, v);
+      const Storage s = add_vec_rn<DType>(a, b);
+      const uint4 d = *reinterpret_cast<const uint4*>(&s);
+      asm volatile("multimem.st.relaxed.sys.global.v4.bf16x2 [%0], {%1,%2,%3,%4};" ::"l"(slot + v * kElemsPerVec),
+                   "r"(d.x),
+                   "r"(d.y),
+                   "r"(d.z),
+                   "r"(d.w)
+                   : "memory");
+    }
+  } else {
+    for (uint32_t v = blockIdx.x * blockDim.x + threadIdx.x; v < total_vec; v += stride) {
+      const uint4 d = *reinterpret_cast<const uint4*>(in_ptr + v * kElemsPerVec);
+      asm volatile("multimem.st.relaxed.sys.global.v4.bf16x2 [%0], {%1,%2,%3,%4};" ::"l"(slot + v * kElemsPerVec),
+                   "r"(d.x),
+                   "r"(d.y),
+                   "r"(d.z),
+                   "r"(d.w)
+                   : "memory");
+    }
+  }
+
+  // Single barrier: publish our pushes and wait until every rank's pushes for
+  // OUR ranges have landed in this GPU's local staging copy.
+  if constexpr (kPerBlockBarrier) {
+    inkling_ar::block_system_barrier<kNumGPU>(state, flag_ptrs, rank);
+  } else {
+    inkling_ar::grid_system_barrier<kNumGPU>(state, flag_ptrs, rank, 0, /*publish_writes=*/true);
+  }
+
+  // Phase 2: local reduce of the N staged shards -- all-local reads (the pushes
+  // just landed in L2), fp32 accumulation.
+  using Storage = typename Trait::Storage;
+  using DType2 = typename Trait::DType2;
+  constexpr uint32_t kVecSize = Trait::kVecSize;
+  for (uint32_t v = blockIdx.x * blockDim.x + threadIdx.x; v < total_vec; v += stride) {
+    Storage s[kNumGPU];
+#pragma unroll
+    for (uint32_t i = 0; i < kNumGPU; ++i)
+      s[i].load(stage_ptr + i * num_items, v);
+    Storage result;
+#pragma unroll
+    for (uint32_t j = 0; j < kVecSize; ++j) {
+      fp32x2_t acc = cast<fp32x2_t>(s[0][j]);
+#pragma unroll
+      for (uint32_t i = 1; i < kNumGPU; ++i) {
+        const fp32x2_t x = cast<fp32x2_t>(s[i][j]);
+        acc.x += x.x;
+        acc.y += x.y;
+      }
+      result[j] = cast<DType2>(acc);  // <-- EPILOGUE SEAM (full row on-rank)
+    }
+    result.store(out_ptr, v);
+  }
+}
+
+// Full one-shot: every rank ld_reduces the ENTIRE range (multicast hardware sum
+// -> full result), writing it to a LOCAL output buffer. No broadcast and NO exit
+// barrier -- the result is complete on this rank, and input-buffer reuse is the
+// caller's responsibility (double-buffer the input). Halving the barrier count
+// wins for tiny, latency-bound (decode) messages. bf16-only.
+template <typename DType, uint32_t kNumGPU>
+__global__ __launch_bounds__(1024, 1) void inkling_multimem_full_oneshot_kernel(
+    DType* __restrict__ mc_ptr,        // multicast input base (covers all peers)
+    DType* __restrict__ local_in_ptr,  // this rank's LOCAL base of the input
+    DType* __restrict__ out_ptr,       // local output base
+    void* const* __restrict__ flag_ptrs,
+    uint32_t* __restrict__ state,
+    const DType* __restrict__ shared,  // optional LOCAL shared-expert partials
+    const uint32_t rank,
+    const uint32_t num_items) {
+  using namespace device;
+  using Trait = InklingAllReduceTrait<DType, kNumGPU>;
+  static_assert(std::is_same_v<DType, bf16_t>, "multimem.add path is bf16-only");
+  constexpr uint32_t kElemsPerVec = Trait::kElemsPerVec;  // 8 bf16 = 16 B
+  const uint32_t total_vec = num_items / kElemsPerVec;
+
+  if (shared != nullptr) {
+    // Fold into this rank's (double-buffered) input region; the publishing
+    // entry barrier then orders it for every peer's ld_reduce. v4 fires only
+    // for 1-2 rows, so the extra local pass is negligible next to the
+    // torch.add launch it replaces.
+    fold_shared_local<DType, kNumGPU>(local_in_ptr, shared, num_items);
+  }
+  inkling_ar::grid_system_barrier<kNumGPU>(
+      state, flag_ptrs, rank, 0, /*publish_writes=*/shared != nullptr);  // ENTRY only (single barrier)
+  const uint32_t stride = gridDim.x * blockDim.x;
+  for (uint32_t v = blockIdx.x * blockDim.x + threadIdx.x; v < total_vec; v += stride) {
+    DType* in = mc_ptr + v * kElemsPerVec;
+    uint32_t r0, r1, r2, r3;
+    asm volatile("multimem.ld_reduce.relaxed.sys.global.add.v4.bf16x2 {%0,%1,%2,%3}, [%4];"
+                 : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
+                 : "l"(in));
+    // <-- EPILOGUE SEAM (norm / sconv / bias on {r0..r3} before the local store)
+    *reinterpret_cast<uint4*>(out_ptr + v * kElemsPerVec) = make_uint4(r0, r1, r2, r3);
+  }
+  // NO exit barrier: result is local & complete; input reuse is caller-managed.
+}
+
+// Blocks needed to cover this rank's two-shot slice (v1/v2/v3 partition).
+template <typename DType, uint32_t kNumGPU>
+uint32_t work_num_blocks(uint32_t n, uint32_t block_size) {
+  using Trait = InklingAllReduceTrait<DType, kNumGPU>;
+  const uint32_t total_vec = n / Trait::kElemsPerVec;
+  const uint32_t vec_per_rank =
+      host::div_ceil(host::div_ceil(total_vec, kNumGPU), device::kWarpThreads) * device::kWarpThreads;
+  return max(1u, host::div_ceil(vec_per_rank, block_size));
+}
+
+// Blocks needed to cover the FULL vec range (the full one-shot kernel reads
+// the entire range on every rank, not a per-rank slice).
+template <typename DType>
+uint32_t full_range_num_blocks(uint32_t n, uint32_t block_size) {
+  constexpr uint32_t kElemsPerVec = InklingAllReduceTrait<DType, 2>::kElemsPerVec;  // kNumGPU-independent
+  return max(1u, host::div_ceil(n / kElemsPerVec, block_size));
+}
+
+// Max blocks that are simultaneously resident for `kernel` at `block_size`.
+// The grid-level barrier REQUIRES all launched blocks to be co-resident (the
+// leader waits for every block to arrive); launching more would deadlock, so
+// the fused kernels cap their grid at this. Small messages need far fewer.
+// Cached per (kernel, block_size, device): the occupancy query costs ~a few us
+// on every eager launch of a latency-bound AR otherwise.
+template <typename Kernel>
+uint32_t max_resident_blocks(Kernel kernel, uint32_t block_size, DLDevice device) {
+  using namespace host;
+  static std::mutex mu;
+  static std::unordered_map<uint64_t, uint32_t> cache;
+  const uint64_t key = (std::bit_cast<uint64_t>(reinterpret_cast<void*>(kernel)) << 12) ^
+                       (static_cast<uint64_t>(block_size) << 8) ^ static_cast<uint64_t>(device.device_id);
+  {
+    std::lock_guard<std::mutex> lk(mu);
+    if (auto it = cache.find(key); it != cache.end()) return it->second;
+  }
+  int sm_count = 0;
+  cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device.device_id);
+  RuntimeCheck(sm_count > 0, "failed to query multiProcessorCount");
+  const uint32_t bps = runtime::get_blocks_per_sm(kernel, block_size);
+  RuntimeCheck(bps > 0, "kernel has zero occupancy at block_size ", block_size);
+  const uint32_t result = static_cast<uint32_t>(sm_count) * bps;
+  std::lock_guard<std::mutex> lk(mu);
+  cache.emplace(key, result);
+  return result;
+}
+
+// Optional shared-expert partials: numel == 0 -> disabled (nullptr); else a
+// LOCAL contiguous tensor covering num_items, folded in-kernel.
+template <typename DType>
+const DType* shared_ptr_or_null(tvm::ffi::TensorView shared, int64_t num_items) {
+  using namespace host;
+  if (shared.numel() == 0) return nullptr;
+  RuntimeCheck(shared.IsContiguous(), "shared must be contiguous");
+  RuntimeCheck(is_type<DType>(shared.dtype()), "shared dtype mismatch");
+  RuntimeCheck(shared.numel() >= num_items, "shared smaller than num_items");
+  RuntimeCheck(std::bit_cast<intptr_t>(shared.data_ptr()) % 16 == 0, "shared not 16B aligned");
+  return reinterpret_cast<const DType*>(shared.data_ptr());
+}
+
+template <typename DType, uint32_t kNumGPU>
+void validate(tvm::ffi::TensorView buf, int64_t peer_ptrs_dev, int64_t rank, int64_t num_items, uint32_t& n) {
+  using namespace host;
+  using Trait = InklingAllReduceTrait<DType, kNumGPU>;
+  n = static_cast<uint32_t>(num_items);
+  RuntimeCheck(buf.IsContiguous(), "buffer must be contiguous");
+  RuntimeCheck(buf.device().device_type == kDLCUDA, "buffer must be on a CUDA device");
+  RuntimeCheck(is_type<DType>(buf.dtype()), "buffer dtype mismatch");
+  RuntimeCheck(static_cast<int64_t>(n) == num_items, "num_items exceeds 4G");
+  RuntimeCheck(buf.numel() >= num_items, "buffer smaller than num_items");
+  RuntimeCheck(n % Trait::kElemsPerVec == 0, "num_items must be a multiple of ", Trait::kElemsPerVec);
+  RuntimeCheck(std::bit_cast<intptr_t>(buf.data_ptr()) % 16 == 0, "buffer not 16B aligned");
+  RuntimeCheck(peer_ptrs_dev != 0, "peer_ptrs_dev is null");
+  RuntimeCheck(rank >= 0 && rank < kNumGPU, "rank out of range");
+}
+
+template <typename DType, uint32_t kNumGPU>
+void inkling_two_shot_all_reduce(
+    tvm::ffi::TensorView local_buffer, int64_t peer_ptrs_dev, int64_t rank, int64_t num_items) {
+  using namespace host;
+  uint32_t n;
+  validate<DType, kNumGPU>(local_buffer, peer_ptrs_dev, rank, num_items, n);
+  const auto device = local_buffer.device();
+  const uint32_t num_blocks = work_num_blocks<DType, kNumGPU>(n, 1024u);  // no in-kernel barrier -> uncapped
+  const auto stream = LaunchKernel::resolve_device(device);
+  LaunchKernel(num_blocks, 1024u, stream)(
+      inkling_two_shot_all_reduce_kernel<DType, kNumGPU>,
+      reinterpret_cast<void* const*>(peer_ptrs_dev),
+      static_cast<uint32_t>(rank),
+      n);
+}
+
+template <typename DType, uint32_t kNumGPU>
+void inkling_two_shot_all_reduce_fused(
+    tvm::ffi::TensorView local_buffer,
+    int64_t data_ptrs_dev,
+    int64_t flag_ptrs_dev,
+    int64_t state_ptr,
+    int64_t rank,
+    int64_t num_items,
+    int64_t nb_override,
+    int64_t bs_override,
+    tvm::ffi::TensorView shared) {
+  using namespace host;
+  uint32_t n;
+  validate<DType, kNumGPU>(local_buffer, data_ptrs_dev, rank, num_items, n);
+  RuntimeCheck(flag_ptrs_dev != 0, "flag_ptrs_dev is null");
+  RuntimeCheck(state_ptr != 0, "state_ptr is null");
+  const DType* shared_ptr = shared_ptr_or_null<DType>(shared, num_items);
+  const auto device = local_buffer.device();
+  const auto kernel = inkling_two_shot_all_reduce_fused_kernel<DType, kNumGPU>;
+  const uint32_t block_size = bs_override > 0 ? static_cast<uint32_t>(bs_override) : 1024u;
+  const uint32_t cap = max_resident_blocks(kernel, block_size, device);
+  const uint32_t num_blocks = nb_override > 0 ? min(static_cast<uint32_t>(nb_override), cap)
+                                              : min(work_num_blocks<DType, kNumGPU>(n, block_size), cap);
+  const auto stream = LaunchKernel::resolve_device(device);
+  LaunchKernel(num_blocks, block_size, stream)(
+      kernel,
+      reinterpret_cast<void* const*>(data_ptrs_dev),
+      reinterpret_cast<void* const*>(flag_ptrs_dev),
+      reinterpret_cast<uint32_t*>(state_ptr),
+      shared_ptr,
+      static_cast<uint32_t>(rank),
+      n);
+}
+
+template <typename DType, uint32_t kNumGPU>
+void inkling_multimem_one_shot_fused(
+    tvm::ffi::TensorView local_buffer,
+    int64_t multicast_ptr,
+    int64_t flag_ptrs_dev,
+    int64_t state_ptr,
+    int64_t rank,
+    int64_t num_items,
+    int64_t nb_override,
+    int64_t bs_override,
+    int64_t per_block_barrier,
+    tvm::ffi::TensorView shared) {
+  using namespace host;
+  uint32_t n;
+  // validate uses the local buffer view only for device/dtype/shape; the kernel
+  // operates on the multicast pointer (plus the local view for the shared fold).
+  validate<DType, kNumGPU>(local_buffer, multicast_ptr, rank, num_items, n);
+  RuntimeCheck(flag_ptrs_dev != 0, "flag_ptrs_dev is null");
+  RuntimeCheck(state_ptr != 0, "state_ptr is null");
+  RuntimeCheck(multicast_ptr % 16 == 0, "multicast_ptr not 16B aligned");
+  const DType* shared_ptr = shared_ptr_or_null<DType>(shared, num_items);
+  const auto device = local_buffer.device();
+  const auto kernel = per_block_barrier ? inkling_multimem_one_shot_fused_kernel<DType, kNumGPU, true>
+                                        : inkling_multimem_one_shot_fused_kernel<DType, kNumGPU, false>;
+  const uint32_t block_size = bs_override > 0 ? static_cast<uint32_t>(bs_override) : 1024u;
+  uint32_t cap = max_resident_blocks(kernel, block_size, device);
+  if (per_block_barrier) cap = min(cap, inkling_ar::kMaxBarrierBlocks);
+  const uint32_t num_blocks = nb_override > 0 ? min(static_cast<uint32_t>(nb_override), cap)
+                                              : min(work_num_blocks<DType, kNumGPU>(n, block_size), cap);
+  const auto stream = LaunchKernel::resolve_device(device);
+  LaunchKernel(num_blocks, block_size, stream)(
+      kernel,
+      reinterpret_cast<DType*>(multicast_ptr),
+      reinterpret_cast<DType*>(local_buffer.data_ptr()),
+      reinterpret_cast<void* const*>(flag_ptrs_dev),
+      reinterpret_cast<uint32_t*>(state_ptr),
+      shared_ptr,
+      static_cast<uint32_t>(rank),
+      n);
+}
+
+template <typename DType, uint32_t kNumGPU>
+void inkling_multimem_push_oneshot(
+    tvm::ffi::TensorView in_buffer,
+    tvm::ffi::TensorView out_buffer,
+    int64_t mc_stage_ptr,
+    int64_t local_stage_ptr,
+    int64_t flag_ptrs_dev,
+    int64_t state_ptr,
+    int64_t rank,
+    int64_t num_items,
+    int64_t nb_override,
+    int64_t bs_override,
+    int64_t per_block_barrier,
+    tvm::ffi::TensorView shared) {
+  using namespace host;
+  uint32_t n;
+  // in_buffer is any LOCAL contiguous bf16 tensor (need not be a symm buffer);
+  // validate() covers contiguity/dtype/alignment; mc_stage stands in for the
+  // pointer null check.
+  validate<DType, kNumGPU>(in_buffer, mc_stage_ptr, rank, num_items, n);
+  const DType* shared_ptr = shared_ptr_or_null<DType>(shared, num_items);
+  RuntimeCheck(out_buffer.IsContiguous(), "out must be contiguous");
+  RuntimeCheck(is_type<DType>(out_buffer.dtype()), "out dtype mismatch");
+  RuntimeCheck(out_buffer.numel() >= num_items, "out smaller than num_items");
+  RuntimeCheck(std::bit_cast<intptr_t>(out_buffer.data_ptr()) % 16 == 0, "out not 16B aligned");
+  RuntimeCheck(flag_ptrs_dev != 0, "flag_ptrs_dev is null");
+  RuntimeCheck(state_ptr != 0, "state_ptr is null");
+  RuntimeCheck(local_stage_ptr != 0, "local_stage_ptr is null");
+  RuntimeCheck(mc_stage_ptr % 16 == 0, "mc_stage_ptr not 16B aligned");
+  RuntimeCheck(local_stage_ptr % 16 == 0, "local_stage_ptr not 16B aligned");
+  const auto device = in_buffer.device();
+  const auto kernel = per_block_barrier ? inkling_multimem_push_oneshot_kernel<DType, kNumGPU, true>
+                                        : inkling_multimem_push_oneshot_kernel<DType, kNumGPU, false>;
+  const uint32_t block_size = bs_override > 0 ? static_cast<uint32_t>(bs_override) : 1024u;
+  uint32_t cap = max_resident_blocks(kernel, block_size, device);
+  // The per-block barrier has kMaxBarrierBlocks flag/epoch slots per rank.
+  if (per_block_barrier) cap = min(cap, inkling_ar::kMaxBarrierBlocks);
+  const uint32_t num_blocks = nb_override > 0 ? min(static_cast<uint32_t>(nb_override), cap)
+                                              : min(full_range_num_blocks<DType>(n, block_size), cap);
+  const auto stream = LaunchKernel::resolve_device(device);
+  LaunchKernel(num_blocks, block_size, stream)(
+      kernel,
+      reinterpret_cast<const DType*>(in_buffer.data_ptr()),
+      reinterpret_cast<DType*>(mc_stage_ptr),
+      reinterpret_cast<const DType*>(local_stage_ptr),
+      reinterpret_cast<DType*>(out_buffer.data_ptr()),
+      reinterpret_cast<void* const*>(flag_ptrs_dev),
+      reinterpret_cast<uint32_t*>(state_ptr),
+      shared_ptr,
+      static_cast<uint32_t>(rank),
+      n);
+}
+
+template <typename DType, uint32_t kNumGPU>
+void inkling_multimem_full_oneshot(
+    tvm::ffi::TensorView in_buffer,
+    tvm::ffi::TensorView out_buffer,
+    int64_t multicast_ptr,
+    int64_t flag_ptrs_dev,
+    int64_t state_ptr,
+    int64_t rank,
+    int64_t num_items,
+    int64_t nb_override,
+    int64_t bs_override,
+    tvm::ffi::TensorView shared) {
+  using namespace host;
+  uint32_t n;
+  validate<DType, kNumGPU>(in_buffer, multicast_ptr, rank, num_items, n);
+  RuntimeCheck(out_buffer.IsContiguous(), "out must be contiguous");
+  RuntimeCheck(is_type<DType>(out_buffer.dtype()), "out dtype mismatch");
+  RuntimeCheck(out_buffer.numel() >= num_items, "out smaller than num_items");
+  RuntimeCheck(std::bit_cast<intptr_t>(out_buffer.data_ptr()) % 16 == 0, "out not 16B aligned");
+  RuntimeCheck(flag_ptrs_dev != 0, "flag_ptrs_dev is null");
+  RuntimeCheck(state_ptr != 0, "state_ptr is null");
+  RuntimeCheck(multicast_ptr % 16 == 0, "multicast_ptr not 16B aligned");
+  const DType* shared_ptr = shared_ptr_or_null<DType>(shared, num_items);
+  const auto device = in_buffer.device();
+  const auto kernel = inkling_multimem_full_oneshot_kernel<DType, kNumGPU>;
+  const uint32_t block_size = bs_override > 0 ? static_cast<uint32_t>(bs_override) : 1024u;
+  const uint32_t cap = max_resident_blocks(kernel, block_size, device);
+  const uint32_t num_blocks = nb_override > 0 ? min(static_cast<uint32_t>(nb_override), cap)
+                                              : min(full_range_num_blocks<DType>(n, block_size), cap);
+  const auto stream = LaunchKernel::resolve_device(device);
+  LaunchKernel(num_blocks, block_size, stream)(
+      kernel,
+      reinterpret_cast<DType*>(multicast_ptr),
+      reinterpret_cast<DType*>(in_buffer.data_ptr()),
+      reinterpret_cast<DType*>(out_buffer.data_ptr()),
+      reinterpret_cast<void* const*>(flag_ptrs_dev),
+      reinterpret_cast<uint32_t*>(state_ptr),
+      shared_ptr,
+      static_cast<uint32_t>(rank),
+      n);
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/inkling/inkling_ar_barrier.cuh

@@ -0,0 +1,193 @@
+// Cross-GPU barrier primitives shared by the Inkling custom all-reduce kernels
+// (inkling_all_reduce.cuh) and the fused AR+sconv+norm decode kernel
+// (inkling_ar_fused_decode.cuh). Two designs are provided: a single-leader
+// grid barrier and a per-block variant.
+//
+// Resources (see inkling_all_reduce.py):
+//   * flags: DEDICATED symmetric uint32 buffer, zero-initialized at setup:
+//     kNumGPU single-leader slots (one per peer), then
+//     kNumGPU * kMaxBarrierBlocks per-(writer, block) slots.
+//   * state: device-LOCAL uint32 buffer: [arrival0, arrival1, release0,
+//     release1, xepoch] padded to kLeaderStateWords, then kMaxBarrierBlocks
+//     per-block epochs. All epochs are monotonic (mod 2^32, wrap-safe compares)
+//     and advance under CUDA-graph replay, so flags never go stale.
+
+#pragma once
+
+#include <cstdint>
+
+namespace inkling_ar {
+
+constexpr uint32_t kLeaderStateWords = 8;
+constexpr uint32_t kMaxBarrierBlocks = 256;
+
+// Grid-level system barrier across all ranks. Two levels:
+//   1. Grid: every block arrives at a self-resetting device counter
+//      (atomicInc wraps at gridDim.x-1); the last arriver is the leader.
+//   2. Cross-GPU: ONLY the leader block does the peer release/acquire
+//      signal/wait, so that O(1) cost is independent of gridDim.x (the reason
+//      the old per-block barrier was slow for many-block launches).
+// The leader then bumps a release counter; followers spin on it (device scope).
+//
+// `st` is a device-local uint32 state buffer: [arrival0, arrival1, release0,
+// release1, xepoch]. idx 0/1 selects the entry/exit instances (distinct grid
+// counters so the two barriers in one kernel don't collide). xepoch is a single
+// monotonic cross-GPU epoch (entry uses e, exit uses e+1) -- consistent across
+// ranks (SPMD) and advancing under CUDA-graph replay, so flags never go stale.
+// s_prev is read BEFORE arriving, and the leader (last arriver) bumps release
+// only after all blocks arrived, so no follower can miss the bump (no deadlock).
+template <uint32_t kNumGPU>
+__device__ __forceinline__ void grid_system_barrier(
+    uint32_t* __restrict__ st, void* const* __restrict__ flag_ptrs, uint32_t rank, uint32_t idx, bool publish_writes) {
+  // publish_writes=true (EXIT barriers): every CTA flushes its just-written
+  // reduced/broadcast slices to SYSTEM scope BEFORE it signals arrival, so the
+  // single leader's `st.release.sys` publishes ALL blocks' stores rather than
+  // only the leader thread's own. Without this, a multi-block launch (the tuned
+  // v2/v3 configs) lets a peer leave the exit barrier and read a slice a
+  // non-leader CTA wrote but never system-published. ONE fence per CTA suffices:
+  // the __syncthreads below orders every thread's stores before thread 0's
+  // fence (CTA-scope happens-before), and `fence.sys + relaxed arrival` is a
+  // release pattern, so the arrival publishes the whole CTA's stores. ENTRY
+  // barriers pass false: the data they gate on was written by a prior kernel
+  // and is already uniformly visible, which the leader's release then promotes
+  // for free. (The solo path needs no fence either way: its st.release.sys
+  // signals below are themselves release ops ordered after the __syncthreads.)
+  uint32_t* xepoch = st + 4;
+  __shared__ uint32_t s_e;
+  __shared__ uint32_t s_prev;
+  __shared__ int s_leader;
+  const bool solo = (gridDim.x == 1u);  // token=1 etc.: the sole block IS the grid
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    if (solo) {
+      s_leader = 1;  // skip the grid arrival/release bookkeeping entirely
+    } else {
+      if (publish_writes) __threadfence_system();               // release pattern with the arrive below
+      s_prev = *static_cast<volatile uint32_t*>(st + 2 + idx);  // pre-barrier release
+      // Self-resetting arrive (atomicInc semantics: wrap at gridDim.x-1).
+      // acq_rel: the release side pairs with the fence above (publishing this
+      // CTA's stores); the acquire side lets the last arriver (leader) inherit
+      // every earlier CTA's release pattern, so its st.release.sys to the peers
+      // covers the whole grid's writes.
+      uint32_t old;
+      asm volatile("atom.acq_rel.gpu.global.inc.u32 %0, [%1], %2;"
+                   : "=r"(old)
+                   : "l"(st + idx), "r"(gridDim.x - 1u)
+                   : "memory");
+      s_leader = (old == gridDim.x - 1u) ? 1 : 0;
+    }
+  }
+  __syncthreads();
+  if (s_leader) {
+    if (threadIdx.x == 0) {
+      const uint32_t e = *xepoch + 1u;
+      *xepoch = e;
+      s_e = e;
+    }
+    __syncthreads();
+    const uint32_t e = s_e;
+    // Cross-GPU arrive+wait with release/acquire at system scope. The release
+    // store publishes THIS (leader) thread's system-visible writes and the
+    // acquire spin makes the peer's visible -- far cheaper than a full
+    // threadfence_system here. Data written by OTHER (non-leader) CTAs is made
+    // system-visible by the publish_writes=true fence they each ran before
+    // arriving (see top), so the leader's single release covers the whole grid.
+    if (threadIdx.x < kNumGPU) {
+      const uint32_t peer = threadIdx.x;
+      uint32_t* remote = static_cast<uint32_t*>(flag_ptrs[peer]) + rank;
+      asm volatile("st.release.sys.global.u32 [%0], %1;" ::"l"(remote), "r"(e) : "memory");
+      uint32_t* mine = static_cast<uint32_t*>(flag_ptrs[rank]) + peer;
+      uint32_t got;
+      do {
+        asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(got) : "l"(mine) : "memory");
+      } while (static_cast<int32_t>(got - e) < 0);  // wrap-safe: epoch is mod-2^32
+    }
+    __syncthreads();
+    if (!solo && threadIdx.x == 0) {
+      // Release-ordered bump: pairs with the followers' ld.acquire.gpu so the
+      // leader's acquired peer state (and its xepoch store above) is visible to
+      // them -- a relaxed atomicAdd would leave that handoff formally unordered.
+      asm volatile("red.release.gpu.global.add.u32 [%0], %1;" ::"l"(st + 2 + idx), "r"(1u) : "memory");
+    }
+  } else {
+    if (threadIdx.x == 0) {
+      // `release` is this rank's LOCAL counter -> device-scope acquire suffices.
+      uint32_t* rel = st + 2 + idx;
+      uint32_t got;
+      do {
+        asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(got) : "l"(rel) : "memory");
+      } while (static_cast<int32_t>(got - s_prev) <= 0);  // wrap-safe
+    }
+    __syncthreads();
+  }
+}
+
+// Device-LOCAL grid sync (no cross-GPU traffic): all blocks arrive at a
+// self-resetting counter (state word 5), the last arriver bumps a release
+// counter (word 6), followers spin on it -- the grid level of
+// grid_system_barrier without the peer handshake. Words 5/6 are spare in the
+// kLeaderStateWords block. Requires all blocks co-resident (the launch cap the
+// fused kernels already apply). Used by the two-phase {AR + scattered sconv}
+// kernel to publish its local scratch between the reduce and conv phases.
+__device__ __forceinline__ void grid_local_sync(uint32_t* __restrict__ st) {
+  __syncthreads();
+  if (gridDim.x > 1u) {
+    if (threadIdx.x == 0) {
+      uint32_t* arrive = st + 5;
+      uint32_t* release = st + 6;
+      const uint32_t prev = *static_cast<volatile uint32_t*>(release);
+      uint32_t old;
+      asm volatile("atom.acq_rel.gpu.global.inc.u32 %0, [%1], %2;"
+                   : "=r"(old)
+                   : "l"(arrive), "r"(gridDim.x - 1u)
+                   : "memory");
+      if (old == gridDim.x - 1u) {
+        asm volatile("red.release.gpu.global.add.u32 [%0], %1;" ::"l"(release), "r"(1u) : "memory");
+      } else {
+        uint32_t got;
+        do {
+          asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(got) : "l"(release) : "memory");
+        } while (static_cast<int32_t>(got - prev) <= 0);  // wrap-safe
+      }
+    }
+    __syncthreads();
+  }
+}
+
+// Per-block cross-GPU barrier (no grid funnel): block b handshakes ONLY with
+// block b on each peer -- one NVLink round trip per block, all blocks in
+// parallel, no arrival/release atomics and no leader serialization. Valid
+// whenever the consumer phase reads exactly the ranges its blockIdx-matched
+// producers wrote (true for the push one-shot: its push and reduce loops use
+// the same grid-stride mapping, and every rank launches the same grid). The
+// signal is a release store, which covers the CTA's prior (multicast) stores
+// via the preceding __syncthreads -- no explicit fence needed. Epochs live in
+// per-block device-local slots (monotonic across launches and CUDA-graph
+// replays, like xepoch).
+template <uint32_t kNumGPU>
+__device__ __forceinline__ void
+block_system_barrier(uint32_t* __restrict__ st, void* const* __restrict__ flag_ptrs, uint32_t rank) {
+  __shared__ uint32_t s_e;
+  __syncthreads();  // CTA stores done before the release signals below
+  if (threadIdx.x == 0) {
+    uint32_t* epoch = st + kLeaderStateWords + blockIdx.x;
+    const uint32_t e = *epoch + 1u;
+    *epoch = e;
+    s_e = e;
+  }
+  __syncthreads();
+  const uint32_t e = s_e;
+  if (threadIdx.x < kNumGPU) {
+    const uint32_t peer = threadIdx.x;
+    uint32_t* remote = static_cast<uint32_t*>(flag_ptrs[peer]) + kNumGPU + rank * kMaxBarrierBlocks + blockIdx.x;
+    asm volatile("st.release.sys.global.u32 [%0], %1;" ::"l"(remote), "r"(e) : "memory");
+    uint32_t* mine = static_cast<uint32_t*>(flag_ptrs[rank]) + kNumGPU + peer * kMaxBarrierBlocks + blockIdx.x;
+    uint32_t got;
+    do {
+      asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(got) : "l"(mine) : "memory");
+    } while (static_cast<int32_t>(got - e) < 0);  // wrap-safe
+  }
+  __syncthreads();
+}
+
+}  // namespace inkling_ar

python/sglang/jit_kernel/csrc/inkling/inkling_ar_fused_decode.cuh

@@ -0,0 +1,830 @@
+// Fused decode {all-reduce -> mlp/attn sconv -> residual-add + RMSNorm} for the
+// Inkling (Moonrise) small-batch decode path -- the v5 push one-shot all-reduce
+// (inkling_all_reduce.cuh) with the EPILOGUE SEAM filled in by the decode short-conv
+// (fused_decode_update.cuh semantics) and the fused-add RMSNorm.
+//
+// Replaces THREE kernels (AR + fused_decode_update + fused_add_rmsnorm) and
+// their intermediate HBM round trips with ONE launch per (AR, sconv, norm)
+// chain. Layout: ONE BLOCK PER TOKEN (decode rows are few and the RMSNorm needs
+// a per-row cross-hidden reduction), VPT 16B vecs (8 channels each) per thread
+// -- a TUNED knob: fewer/fatter threads buy load ILP and a cheaper block
+// reduction; more threads buy parallelism. Phases:
+//
+//   0. prefetch: sconv metadata, conv history and conv weights load FIRST --
+//      none depend on the producer kernel's output, and their HBM latency
+//      hides under the cross-GPU barrier below. (This -- not PDL -- is where
+//      the fused kernel's latency win comes from: the producer GEMMs never
+//      trigger programmatic launch early, so the PDL wait is a no-op in
+//      practice and the launch attribute only pipelines the launch tail.)
+//   1. push:    griddepcontrol.wait, then multicast-store this rank's partial
+//               row into staging slot (rank*T + t)*D; issue the residual load.
+//   2. barrier: per-block peer handshake (block t <-> peers' block t).
+//   3. reduce:  fp32 sum of the kNumGPU staged shards; round to bf16 `xb`
+//               (bit-identical to what the unfused AR would have stored).
+//   4. sconv:   decode causal_conv1d on xb (W-1 cached taps gated by
+//               cache_mask + current token), optional SiLU, optional +xb
+//               residual; cache shift-update (+ optional track-copy) --
+//               identical semantics to fused_decode_update_kernel.
+//   5. norm:    r = residual_in + y (fp32); block-reduce sum(r^2); write
+//               residual_out = bf16(r) and hs_out = bf16(r * rsqrt(mean+eps)
+//               * gamma)  (fused_add_rmsnorm semantics).
+//
+// Staging reuse is caller-managed (A/B rotation shared with v5 -- this kernel
+// IS a v5 AR occupying one rotation slot). PAD rows (cache_indices == -1)
+// still compute y/hs but never write the cache, matching the unfused kernel.
+// bf16-only.
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include "inkling_ar_barrier.cuh"
+#include <bit>
+#include <cstdint>
+#include <cuda_bf16.h>
+#include <type_traits>
+
+namespace {
+
+constexpr int kPadSlot = -1;
+constexpr uint32_t kVecElems = 8;  // bf16x8 = 16 B
+
+// Register-level fused add of two bf16x8 vecs (fp32 math, ONE round to bf16)
+// -- torch.add numerics, so folding the shared-expert partials into the push
+// stays bit-identical to the unfused {torch.add -> AR} chain.
+__device__ __forceinline__ uint4 add_bf16x8_rn(const uint4 a, const uint4 b) {
+  const auto* a2 = reinterpret_cast<const __nv_bfloat162*>(&a);
+  const auto* b2 = reinterpret_cast<const __nv_bfloat162*>(&b);
+  uint4 out;
+  auto* o2 = reinterpret_cast<__nv_bfloat162*>(&out);
+#pragma unroll
+  for (int j = 0; j < 4; ++j) {
+    const float2 x = __bfloat1622float2(a2[j]);
+    const float2 y = __bfloat1622float2(b2[j]);
+    o2[j] = __floats2bfloat162_rn(x.x + y.x, x.y + y.y);
+  }
+  return out;
+}
+
+struct ArSconvNormParams {
+  // AR
+  const void* __restrict__ in;      // [T, D] partial sums (LOCAL tensor)
+  const void* __restrict__ shared;  // optional [T, D] shared-expert partials (LOCAL)
+  void* __restrict__ mc_stage;      // multicast staging base (>= kNumGPU*T*D elems)
+  const void* __restrict__ stage;   // this GPU's local view of the staging base
+  void* const* __restrict__ flag_ptrs;
+  uint32_t* __restrict__ state;
+  // sconv (fused_decode_update semantics)
+  void* __restrict__ cache;                // [pool, W-1, D], in-place update
+  const void* __restrict__ cache_indices;  // int32 [T] (PAD == -1)
+  const void* __restrict__ cache_mask;     // bool  [T]
+  const void* __restrict__ conv_weight;    // [D, W]
+  const void* __restrict__ track_mask;     // bool  [T]  (DO_TRACK only)
+  const void* __restrict__ track_indices;  // int64 [T]  (DO_TRACK only)
+  // norm
+  const void* __restrict__ residual_in;  // [T, D]
+  void* __restrict__ residual_out;       // [T, D]
+  void* __restrict__ hs_out;             // [T, D]
+  const void* __restrict__ norm_weight;  // [D]
+  float eps;
+  // strides (elements)
+  int64_t in_stride_t;
+  int64_t shared_stride_t;
+  int64_t res_in_stride_t;
+  int64_t res_out_stride_t;
+  int64_t hs_stride_t;
+  int64_t cache_stride_slot;
+  int64_t cache_stride_w;
+  int64_t conv_weight_stride_d;
+  int64_t track_idx_stride;
+  uint32_t rank;
+  uint32_t T;
+  uint32_t D;
+};
+
+// VPT = 16B vecs handled per thread (tuning knob; see the header comment).
+// Vec i of a thread is at index threadIdx.x + i*blockDim.x (warp-coalesced).
+template <typename DType, uint32_t kNumGPU, int W, bool USE_SILU, bool USE_RESIDUAL, bool DO_TRACK, int VPT>
+__global__ __launch_bounds__(1024, 1) void inkling_ar_sconv_norm_kernel(const __grid_constant__ ArSconvNormParams p) {
+  static_assert(std::is_same_v<DType, __nv_bfloat16>, "multimem push path is bf16-only");
+  constexpr int W1 = W - 1;
+  const uint32_t t = blockIdx.x;
+  const uint32_t vecs = p.D / kVecElems;
+
+  uint32_t c0[VPT];
+  bool act[VPT];
+#pragma unroll
+  for (int i = 0; i < VPT; ++i) {
+    const uint32_t v = threadIdx.x + i * blockDim.x;
+    act[i] = v < vecs;
+    c0[i] = (act[i] ? v : 0) * kVecElems;  // clamp: inactive lanes never store
+  }
+
+  // ---- 0. prefetch (independent of the producer's output) ----
+  const int ci = static_cast<const int32_t*>(p.cache_indices)[t];
+  const bool valid = ci != kPadSlot;
+  const int slot_id = valid ? ci : 0;  // PAD lanes still emit y, never write cache
+  const float cm = static_cast<const bool*>(p.cache_mask)[t] ? 1.0f : 0.0f;
+  auto* cp = static_cast<__nv_bfloat16*>(p.cache);
+  const auto* wp = static_cast<const __nv_bfloat16*>(p.conv_weight);
+  uint4 hist_raw[VPT][W1];
+  __nv_bfloat16 wtaps[VPT][kVecElems][W];
+#pragma unroll
+  for (int i = 0; i < VPT; ++i) {
+    if (!act[i]) continue;
+    const int64_t cache_base = static_cast<int64_t>(slot_id) * p.cache_stride_slot + c0[i];
+#pragma unroll
+    for (int w = 0; w < W1; ++w) {
+      hist_raw[i][w] = *reinterpret_cast<const uint4*>(&cp[cache_base + w * p.cache_stride_w]);
+    }
+#pragma unroll
+    for (int j = 0; j < static_cast<int>(kVecElems); ++j) {
+      const int64_t wrow = static_cast<int64_t>(c0[i] + j) * p.conv_weight_stride_d;
+      if constexpr (W == 4) {
+        // One 8B load per channel row (bf16 x4, 8B-aligned for contiguous [D, W]).
+        if (p.conv_weight_stride_d == W) {
+          *reinterpret_cast<uint2*>(wtaps[i][j]) = *reinterpret_cast<const uint2*>(wp + wrow);
+          continue;
+        }
+      }
+#pragma unroll
+      for (int w = 0; w < W; ++w)
+        wtaps[i][j][w] = wp[wrow + w];
+    }
+  }
+
+  // ---- 1. push: wait for the producer's output (PDL; no-op without a PDL
+  // launch or an early-triggering producer), multicast-store this rank's
+  // partial row, and issue the residual load (it lands under the barrier). ----
+  asm volatile("griddepcontrol.wait;" ::: "memory");
+  const auto* in_row = static_cast<const __nv_bfloat16*>(p.in) + t * p.in_stride_t;
+  const auto* sh_row =
+      p.shared == nullptr ? nullptr : static_cast<const __nv_bfloat16*>(p.shared) + t * p.shared_stride_t;
+  auto* slot = static_cast<__nv_bfloat16*>(p.mc_stage) + (static_cast<uint64_t>(p.rank) * p.T + t) * p.D;
+  uint4 res_raw[VPT];
+#pragma unroll
+  for (int i = 0; i < VPT; ++i) {
+    if (!act[i]) continue;
+    uint4 d = *reinterpret_cast<const uint4*>(in_row + c0[i]);
+    if (sh_row != nullptr) {
+      d = add_bf16x8_rn(d, *reinterpret_cast<const uint4*>(sh_row + c0[i]));
+    }
+    asm volatile("multimem.st.relaxed.sys.global.v4.bf16x2 [%0], {%1,%2,%3,%4};" ::"l"(slot + c0[i]),
+                 "r"(d.x),
+                 "r"(d.y),
+                 "r"(d.z),
+                 "r"(d.w)
+                 : "memory");
+    res_raw[i] = *reinterpret_cast<const uint4*>(
+        static_cast<const __nv_bfloat16*>(p.residual_in) + t * p.res_in_stride_t + c0[i]);
+  }
+
+  // ---- 2. per-block barrier: all ranks' row-t pushes have landed locally ----
+  inkling_ar::block_system_barrier<kNumGPU>(p.state, p.flag_ptrs, p.rank);
+  // Inactive lanes must NOT exit: they participate in the norm's __syncthreads
+  // and full-mask warp shuffles below (sumsq contribution 0).
+
+  float r[VPT][kVecElems];
+  float sumsq = 0.0f;
+  const auto* stage = static_cast<const __nv_bfloat16*>(p.stage);
+#pragma unroll
+  for (int i = 0; i < VPT; ++i) {
+    if (!act[i]) continue;
+    // ---- 3. reduce: fp32 sum of the kNumGPU staged shards; round to bf16 ----
+    float xf[kVecElems];
+#pragma unroll
+    for (int j = 0; j < static_cast<int>(kVecElems); ++j)
+      xf[j] = 0.0f;
+#pragma unroll
+    for (uint32_t rr = 0; rr < kNumGPU; ++rr) {
+      const uint4 d = *reinterpret_cast<const uint4*>(stage + (static_cast<uint64_t>(rr) * p.T + t) * p.D + c0[i]);
+      const auto* h2 = reinterpret_cast<const __nv_bfloat162*>(&d);
+#pragma unroll
+      for (int j = 0; j < 4; ++j) {
+        const float2 f = __bfloat1622float2(h2[j]);
+        xf[2 * j] += f.x;
+        xf[2 * j + 1] += f.y;
+      }
+    }
+    // Round to bf16 exactly as the unfused AR's store would (the sconv below
+    // and the cache append must see the same bits the unfused path sees).
+    __nv_bfloat162 xb2[4];
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+      xb2[j] = __floats2bfloat162_rn(xf[2 * j], xf[2 * j + 1]);
+
+    // ---- 4. sconv: conv over W-1 cached taps (prefetched) + current token ----
+    float y[kVecElems];
+#pragma unroll
+    for (int j = 0; j < static_cast<int>(kVecElems); ++j) {
+      const float xj = __bfloat162float(reinterpret_cast<const __nv_bfloat16*>(xb2)[j]);
+      float acc = 0.0f;
+#pragma unroll
+      for (int w = 0; w < W1; ++w) {
+        const float h = __bfloat162float(reinterpret_cast<const __nv_bfloat16*>(&hist_raw[i][w])[j]);
+        acc += h * cm * __bfloat162float(wtaps[i][j][w]);
+      }
+      acc += xj * __bfloat162float(wtaps[i][j][W1]);
+      if constexpr (USE_SILU) acc = __fdividef(acc, 1.0f + __expf(-acc));
+      if constexpr (USE_RESIDUAL) acc += xj;
+      y[j] = acc;
+    }
+
+    if (valid) {
+      // Shift state left (gated by cache_mask), append current token (xb).
+      const int64_t cache_base = static_cast<int64_t>(slot_id) * p.cache_stride_slot + c0[i];
+      int64_t track_base = 0;
+      bool do_tr = false;
+      if constexpr (DO_TRACK) {
+        do_tr = static_cast<const bool*>(p.track_mask)[t];
+        if (do_tr) {
+          const int64_t tslot =
+              static_cast<const int64_t*>(p.track_indices)[static_cast<int64_t>(t) * p.track_idx_stride];
+          track_base = tslot * p.cache_stride_slot + c0[i];
+        }
+      }
+      const uint4 zero = make_uint4(0, 0, 0, 0);
+#pragma unroll
+      for (int w = 0; w < W1; ++w) {
+        const uint4 nv =
+            (w < W1 - 1) ? ((cm != 0.0f) ? hist_raw[i][w + 1] : zero) : *reinterpret_cast<const uint4*>(xb2);
+        *reinterpret_cast<uint4*>(&cp[cache_base + w * p.cache_stride_w]) = nv;
+        if constexpr (DO_TRACK) {
+          if (do_tr) {
+            *reinterpret_cast<uint4*>(&cp[track_base + w * p.cache_stride_w]) = nv;
+          }
+        }
+      }
+    }
+
+    // ---- 5a. residual add (fused_add_rmsnorm semantics) ----
+    // yb: round the sconv output to bf16 first -- the unfused path writes y to
+    // HBM as bf16 before the norm kernel reads it back.
+#pragma unroll
+    for (int j = 0; j < static_cast<int>(kVecElems); ++j) {
+      const float yb = __bfloat162float(__float2bfloat16_rn(y[j]));
+      r[i][j] = yb + __bfloat162float(reinterpret_cast<const __nv_bfloat16*>(&res_raw[i])[j]);
+      sumsq += r[i][j] * r[i][j];
+    }
+  }
+
+  // ---- 5b. block reduction of sumsq (warp shuffle + one smem slot/warp) ----
+  __shared__ float s_warp[32];
+  __shared__ float s_inv;
+  const uint32_t lane = threadIdx.x & 31u;
+  const uint32_t warp = threadIdx.x >> 5;
+#pragma unroll
+  for (int off = 16; off > 0; off >>= 1)
+    sumsq += __shfl_down_sync(~0u, sumsq, off);
+  if (lane == 0) s_warp[warp] = sumsq;
+  __syncthreads();
+  if (warp == 0) {
+    const uint32_t nwarps = (blockDim.x + 31u) >> 5;
+    float total = (lane < nwarps && lane < 32u) ? s_warp[lane] : 0.0f;
+#pragma unroll
+    for (int off = 16; off > 0; off >>= 1)
+      total += __shfl_down_sync(~0u, total, off);
+    if (lane == 0) s_inv = rsqrtf(total / static_cast<float>(p.D) + p.eps);
+  }
+  __syncthreads();
+  const float inv = s_inv;
+
+  const auto* gw = static_cast<const __nv_bfloat16*>(p.norm_weight);
+  auto* res_out = static_cast<__nv_bfloat16*>(p.residual_out) + t * p.res_out_stride_t;
+  auto* hs_out = static_cast<__nv_bfloat16*>(p.hs_out) + t * p.hs_stride_t;
+#pragma unroll
+  for (int i = 0; i < VPT; ++i) {
+    if (!act[i]) continue;
+    __nv_bfloat162 ro[4], ho[4];
+#pragma unroll
+    for (int j = 0; j < 4; ++j) {
+      const float g0 = __bfloat162float(gw[c0[i] + 2 * j]);
+      const float g1 = __bfloat162float(gw[c0[i] + 2 * j + 1]);
+      ro[j] = __floats2bfloat162_rn(r[i][2 * j], r[i][2 * j + 1]);
+      ho[j] = __floats2bfloat162_rn(r[i][2 * j] * inv * g0, r[i][2 * j + 1] * inv * g1);
+    }
+    *reinterpret_cast<uint4*>(res_out + c0[i]) = *reinterpret_cast<const uint4*>(ro);
+    *reinterpret_cast<uint4*>(hs_out + c0[i]) = *reinterpret_cast<const uint4*>(ho);
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Target-verify variant: {AR -> extend-style causal_conv1d ->
+// save_intermediate_conv_windows -> add+RMSNorm} in one launch. Every sequence
+// has exactly `q` (draft_token_num) consecutive tokens; token t belongs to
+// seq = t/q with bos = seq*q. The conv's cross-token taps are RE-REDUCED from
+// the v5 staging buffer (any block can rebuild any token's reduced row by
+// summing the staged shards -- ~kNumGPU x 16B extra local L2 reads per tap, no
+// cross-block dependency). The conv does NOT update the working cache at
+// verify; instead the per-position windows are written to intermediate_out
+// (consumed by update_conv_state_after_mtp_verify), whose values are exactly
+// the cache prefix rows and the re-reduced x this kernel already holds.
+struct ArSconvNormVerifyParams {
+  const void* __restrict__ in;      // [T, D] partial sums (LOCAL tensor)
+  const void* __restrict__ shared;  // optional [T, D] shared-expert partials (LOCAL)
+  void* __restrict__ mc_stage;      // multicast staging base
+  const void* __restrict__ stage;   // this GPU's local view of the staging base
+  void* const* __restrict__ flag_ptrs;
+  uint32_t* __restrict__ state;
+  const void* __restrict__ cache;          // [pool, W-1, D] (read-only here)
+  const void* __restrict__ cache_indices;  // int32 [B] per-SEQ slot (PAD == -1)
+  const void* __restrict__ cache_mask;     // bool  [B] per-SEQ prefix gate
+  const void* __restrict__ conv_weight;    // [D, W]
+  void* __restrict__ inter_out;            // [max_bs, q, W-1, D]
+  const void* __restrict__ residual_in;    // [T, D]
+  void* __restrict__ residual_out;         // [T, D]
+  void* __restrict__ hs_out;               // [T, D]
+  const void* __restrict__ norm_weight;    // [D]
+  float eps;
+  int64_t in_stride_t;
+  int64_t shared_stride_t;
+  int64_t res_in_stride_t;
+  int64_t res_out_stride_t;
+  int64_t hs_stride_t;
+  int64_t cache_stride_slot;
+  int64_t cache_stride_w;
+  int64_t conv_weight_stride_d;
+  int64_t inter_stride_b;
+  int64_t inter_stride_t;
+  int64_t inter_stride_w;
+  uint32_t rank;
+  uint32_t T;
+  uint32_t D;
+  uint32_t q;  // draft_token_num
+};
+
+template <typename DType, uint32_t kNumGPU, int W, bool USE_SILU, bool USE_RESIDUAL>
+__global__
+__launch_bounds__(1024, 1) void inkling_ar_sconv_norm_verify_kernel(const __grid_constant__ ArSconvNormVerifyParams p) {
+  static_assert(std::is_same_v<DType, __nv_bfloat16>, "multimem push path is bf16-only");
+  constexpr int W1 = W - 1;
+  const uint32_t vecs = p.D / kVecElems;
+  const uint32_t v = threadIdx.x;  // one 16B vec (8 channels) per thread
+  const bool active = v < vecs;
+  const uint32_t c0 = (active ? v : 0) * kVecElems;
+  const uint32_t stride_t = gridDim.x;  // grid-stride over tokens
+
+  // Conv weights are token-independent (per channel) -- load once.
+  const auto* wp = static_cast<const __nv_bfloat16*>(p.conv_weight);
+  __nv_bfloat16 wtaps[kVecElems][W];
+  if (active) {
+#pragma unroll
+    for (int j = 0; j < static_cast<int>(kVecElems); ++j) {
+      const int64_t wrow = static_cast<int64_t>(c0 + j) * p.conv_weight_stride_d;
+      if constexpr (W == 4) {
+        if (p.conv_weight_stride_d == W) {
+          *reinterpret_cast<uint2*>(wtaps[j]) = *reinterpret_cast<const uint2*>(wp + wrow);
+          continue;
+        }
+      }
+#pragma unroll
+      for (int w = 0; w < W; ++w)
+        wtaps[j][w] = wp[wrow + w];
+    }
+  }
+
+  // ---- Phase 1: push every assigned row into staging (PDL-gated input). ----
+  // A single grid barrier (below) then makes ALL rows' pushes visible on this
+  // GPU, so Phase 2's cross-token (neighbor) staging reads are race-free -- the
+  // per-block barrier only synchronized the same blockIdx across ranks and did
+  // NOT order block t-j's push before block t's read.
+  asm volatile("griddepcontrol.wait;" ::: "memory");
+  auto* mc = static_cast<__nv_bfloat16*>(p.mc_stage);
+  const auto* in = static_cast<const __nv_bfloat16*>(p.in);
+  const auto* sh = static_cast<const __nv_bfloat16*>(p.shared);
+  if (active) {
+    for (uint32_t t = blockIdx.x; t < p.T; t += stride_t) {
+      uint4 d = *reinterpret_cast<const uint4*>(in + t * p.in_stride_t + c0);
+      if (sh != nullptr) {
+        // Fold the shared-expert partials in registers (torch.add numerics);
+        // the staged value then matches the unfused pre-added input, so the
+        // cross-token re-reduces below stay bit-identical too.
+        d = add_bf16x8_rn(d, *reinterpret_cast<const uint4*>(sh + t * p.shared_stride_t + c0));
+      }
+      auto* slot = mc + (static_cast<uint64_t>(p.rank) * p.T + t) * p.D + c0;
+      asm volatile("multimem.st.relaxed.sys.global.v4.bf16x2 [%0], {%1,%2,%3,%4};" ::"l"(slot),
+                   "r"(d.x),
+                   "r"(d.y),
+                   "r"(d.z),
+                   "r"(d.w)
+                   : "memory");
+    }
+  }
+
+  // ---- Grid barrier: all pushes done + system-visible across all ranks. ----
+  inkling_ar::grid_system_barrier<kNumGPU>(
+      p.state,
+      p.flag_ptrs,
+      p.rank,
+      0,
+      /*publish_writes=*/true);
+
+  // ---- Phase 2: reduce + conv + save_windows + add-RMSNorm per row. ----
+  const auto* stage = static_cast<const __nv_bfloat16*>(p.stage);
+  const auto* cp = static_cast<const __nv_bfloat16*>(p.cache);
+  const auto* gw = static_cast<const __nv_bfloat16*>(p.norm_weight);
+  __shared__ float s_warp[32];
+  __shared__ float s_inv;
+  const uint32_t lane = threadIdx.x & 31u;
+  const uint32_t warp = threadIdx.x >> 5;
+
+  auto reduce_row = [&](uint32_t row, __nv_bfloat162* out2) {
+    float xf[kVecElems];
+#pragma unroll
+    for (int j = 0; j < static_cast<int>(kVecElems); ++j)
+      xf[j] = 0.0f;
+#pragma unroll
+    for (uint32_t rr = 0; rr < kNumGPU; ++rr) {
+      const uint4 d = *reinterpret_cast<const uint4*>(stage + (static_cast<uint64_t>(rr) * p.T + row) * p.D + c0);
+      const auto* h2 = reinterpret_cast<const __nv_bfloat162*>(&d);
+#pragma unroll
+      for (int j = 0; j < 4; ++j) {
+        const float2 f = __bfloat1622float2(h2[j]);
+        xf[2 * j] += f.x;
+        xf[2 * j + 1] += f.y;
+      }
+    }
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+      out2[j] = __floats2bfloat162_rn(xf[2 * j], xf[2 * j + 1]);
+  };
+
+  for (uint32_t t = blockIdx.x; t < p.T; t += stride_t) {
+    const uint32_t seq = t / p.q;
+    const uint32_t tq = t - seq * p.q;
+    const int bos = static_cast<int>(seq * p.q);
+    const int ci = static_cast<const int32_t*>(p.cache_indices)[seq];
+    const bool valid = ci != kPadSlot;
+    const int slot_id = valid ? ci : 0;
+    const float cm = (valid && static_cast<const bool*>(p.cache_mask)[seq]) ? 1.0f : 0.0f;
+    const int64_t cache_base = static_cast<int64_t>(slot_id) * p.cache_stride_slot + c0;
+
+    float r[kVecElems];
+    float sumsq = 0.0f;
+    if (active) {
+      uint4 pref_raw[W1];
+#pragma unroll
+      for (int w = 0; w < W1; ++w) {
+        pref_raw[w] = *reinterpret_cast<const uint4*>(&cp[cache_base + w * p.cache_stride_w]);
+      }
+      const uint4 res_raw = *reinterpret_cast<const uint4*>(
+          static_cast<const __nv_bfloat16*>(p.residual_in) + t * p.res_in_stride_t + c0);
+
+      __nv_bfloat162 xb2[4];      // own row
+      __nv_bfloat162 xn2[W1][4];  // neighbors t-1 .. t-(W-1), where in-seq
+      reduce_row(t, xb2);
+#pragma unroll
+      for (int j = 1; j <= W1; ++j) {
+        const int n = static_cast<int>(t) - j;
+        if (n >= bos) reduce_row(static_cast<uint32_t>(n), xn2[j - 1]);
+      }
+
+      // conv (jit causal_conv1d semantics, fp32 accum, ascending tap order).
+      float y[kVecElems];
+#pragma unroll
+      for (int j = 0; j < static_cast<int>(kVecElems); ++j) {
+        const float xj = __bfloat162float(reinterpret_cast<const __nv_bfloat16*>(xb2)[j]);
+        float acc = 0.0f;
+#pragma unroll
+        for (int iw = 0; iw < W1; ++iw) {
+          const int shifted = static_cast<int>(t) - W1 + iw;
+          float tap = 0.0f;
+          if (shifted >= bos) {
+            tap = __bfloat162float(reinterpret_cast<const __nv_bfloat16*>(xn2[W1 - 1 - iw])[j]);
+          } else {
+            const int prefix_pos = shifted - bos + W1;
+            if (prefix_pos >= 0) {
+              tap = cm * __bfloat162float(reinterpret_cast<const __nv_bfloat16*>(&pref_raw[prefix_pos])[j]);
+            }
+          }
+          acc += tap * __bfloat162float(wtaps[j][iw]);
+        }
+        acc += xj * __bfloat162float(wtaps[j][W1]);
+        if constexpr (USE_SILU) acc = __fdividef(acc, 1.0f + __expf(-acc));
+        if constexpr (USE_RESIDUAL) acc += xj;
+        y[j] = acc;
+      }
+
+      // save_intermediate_conv_windows: window after draft position tq is raw
+      // copies of {cache prefix rows | reduced x rows} (no cm gating).
+      if (valid) {
+        auto* op = static_cast<__nv_bfloat16*>(p.inter_out) + static_cast<int64_t>(seq) * p.inter_stride_b +
+                   static_cast<int64_t>(tq) * p.inter_stride_t + c0;
+#pragma unroll
+        for (int w = 0; w < W1; ++w) {
+          const int position = static_cast<int>(tq) + 1 + w;
+          uint4 val;
+          if (position < W1) {
+            val = pref_raw[position];
+          } else {
+            const int g = bos + position - W1;
+            val = (g == static_cast<int>(t)) ? *reinterpret_cast<const uint4*>(xb2)
+                                             : *reinterpret_cast<const uint4*>(xn2[t - g - 1]);
+          }
+          *reinterpret_cast<uint4*>(op + w * p.inter_stride_w) = val;
+        }
+      }
+
+      // residual add.
+#pragma unroll
+      for (int j = 0; j < static_cast<int>(kVecElems); ++j) {
+        const float yb = __bfloat162float(__float2bfloat16_rn(y[j]));
+        r[j] = yb + __bfloat162float(reinterpret_cast<const __nv_bfloat16*>(&res_raw)[j]);
+        sumsq += r[j] * r[j];
+      }
+    }
+
+    // block reduction of sumsq (all threads participate; inactive contribute 0).
+    __syncthreads();  // protect s_warp/s_inv reuse across the token loop
+    float ss = sumsq;
+#pragma unroll
+    for (int off = 16; off > 0; off >>= 1)
+      ss += __shfl_down_sync(~0u, ss, off);
+    if (lane == 0) s_warp[warp] = ss;
+    __syncthreads();
+    if (warp == 0) {
+      const uint32_t nwarps = (blockDim.x + 31u) >> 5;
+      float total = (lane < nwarps && lane < 32u) ? s_warp[lane] : 0.0f;
+#pragma unroll
+      for (int off = 16; off > 0; off >>= 1)
+        total += __shfl_down_sync(~0u, total, off);
+      if (lane == 0) s_inv = rsqrtf(total / static_cast<float>(p.D) + p.eps);
+    }
+    __syncthreads();
+    const float inv = s_inv;
+
+    if (active) {
+      auto* res_out = static_cast<__nv_bfloat16*>(p.residual_out) + t * p.res_out_stride_t;
+      auto* hs_out = static_cast<__nv_bfloat16*>(p.hs_out) + t * p.hs_stride_t;
+      __nv_bfloat162 ro[4], ho[4];
+#pragma unroll
+      for (int j = 0; j < 4; ++j) {
+        const float g0 = __bfloat162float(gw[c0 + 2 * j]);
+        const float g1 = __bfloat162float(gw[c0 + 2 * j + 1]);
+        ro[j] = __floats2bfloat162_rn(r[2 * j], r[2 * j + 1]);
+        ho[j] = __floats2bfloat162_rn(r[2 * j] * inv * g0, r[2 * j + 1] * inv * g1);
+      }
+      *reinterpret_cast<uint4*>(res_out + c0) = *reinterpret_cast<const uint4*>(ro);
+      *reinterpret_cast<uint4*>(hs_out + c0) = *reinterpret_cast<const uint4*>(ho);
+    }
+  }
+}
+
+template <typename DType, uint32_t kNumGPU, int W, bool USE_SILU, bool USE_RESIDUAL, bool DO_TRACK>
+struct ArSconvNormKernel {
+  template <int VPT>
+  static void launch(const ArSconvNormParams& params, uint32_t t_num, uint32_t vecs, DLDevice dev, bool pdl) {
+    using namespace host;
+    const uint32_t block = min(1024u, div_ceil(div_ceil(vecs, VPT), 32u) * 32u);
+    constexpr auto kernel = inkling_ar_sconv_norm_kernel<DType, kNumGPU, W, USE_SILU, USE_RESIDUAL, DO_TRACK, VPT>;
+    LaunchKernel(dim3{t_num}, dim3{block}, dev).enable_pdl(pdl)(kernel, params);
+  }
+
+  static void
+  run(tvm::ffi::TensorView in,
+      tvm::ffi::TensorView residual_in,
+      tvm::ffi::TensorView residual_out,
+      tvm::ffi::TensorView hs_out,
+      tvm::ffi::TensorView norm_weight,
+      double eps,
+      tvm::ffi::TensorView cache,
+      tvm::ffi::TensorView cache_indices,
+      tvm::ffi::TensorView cache_mask,
+      tvm::ffi::TensorView conv_weight,
+      tvm::ffi::TensorView track_mask,
+      tvm::ffi::TensorView track_indices,
+      int64_t mc_stage_ptr,
+      int64_t local_stage_ptr,
+      int64_t flag_ptrs_dev,
+      int64_t state_ptr,
+      int64_t rank,
+      int64_t enable_pdl,
+      int64_t vecs_per_thread,
+      tvm::ffi::TensorView shared) {
+    using namespace host;
+    auto T = SymbolicSize{"T"};
+    auto D = SymbolicSize{"D"};
+    auto Wd = SymbolicSize{"W"};
+    auto W1s = SymbolicSize{"W_minus_1"};
+    auto dev = SymbolicDevice{};
+    dev.set_options<kDLCUDA>();
+    Wd.set_value(W);
+    W1s.set_value(W - 1);
+
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(in);
+    const bool do_shared = shared.numel() > 0;
+    if (do_shared) {
+      TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(shared);
+      RuntimeCheck(shared.stride(0) % kVecElems == 0, "shared row stride must keep 16B alignment");
+      RuntimeCheck(std::bit_cast<intptr_t>(shared.data_ptr()) % 16 == 0, "shared not 16B aligned");
+    }
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(residual_in);
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(residual_out);
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(hs_out);
+    TensorMatcher({D}).with_dtype<DType>().with_device(dev).verify(norm_weight);
+    TensorMatcher({-1, W1s, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({T}).with_dtype<int32_t>().with_device(dev).verify(cache_indices);
+    TensorMatcher({T}).with_device(dev).verify(cache_mask);
+    TensorMatcher({D, Wd}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(conv_weight);
+    const uint32_t t_num = static_cast<uint32_t>(T.unwrap());
+    const uint32_t d_num = static_cast<uint32_t>(D.unwrap());
+    const uint32_t vecs = d_num / kVecElems;
+    RuntimeCheck(
+        t_num >= 1 && t_num <= inkling_ar::kMaxBarrierBlocks,
+        "T must be in [1, kMaxBarrierBlocks] (one barrier slot per token)");
+    RuntimeCheck(d_num % kVecElems == 0, "D must be a multiple of 8");
+    RuntimeCheck(cache.stride(2) == 1, "cache must be channel-contiguous");
+    RuntimeCheck(mc_stage_ptr % 16 == 0, "mc_stage_ptr not 16B aligned");
+    RuntimeCheck(local_stage_ptr != 0 && local_stage_ptr % 16 == 0, "bad local_stage_ptr");
+    RuntimeCheck(flag_ptrs_dev != 0, "flag_ptrs_dev is null");
+    RuntimeCheck(state_ptr != 0, "state_ptr is null");
+    RuntimeCheck(rank >= 0 && rank < kNumGPU, "rank out of range");
+    RuntimeCheck(in.stride(0) % kVecElems == 0, "in row stride must keep 16B alignment");
+    RuntimeCheck(std::bit_cast<intptr_t>(in.data_ptr()) % 16 == 0, "in not 16B aligned");
+
+    const auto params = ArSconvNormParams{
+        .in = in.data_ptr(),
+        .shared = do_shared ? shared.data_ptr() : nullptr,
+        .mc_stage = reinterpret_cast<void*>(mc_stage_ptr),
+        .stage = reinterpret_cast<const void*>(local_stage_ptr),
+        .flag_ptrs = reinterpret_cast<void* const*>(flag_ptrs_dev),
+        .state = reinterpret_cast<uint32_t*>(state_ptr),
+        .cache = cache.data_ptr(),
+        .cache_indices = cache_indices.data_ptr(),
+        .cache_mask = cache_mask.data_ptr(),
+        .conv_weight = conv_weight.data_ptr(),
+        .track_mask = DO_TRACK ? track_mask.data_ptr() : nullptr,
+        .track_indices = DO_TRACK ? track_indices.data_ptr() : nullptr,
+        .residual_in = residual_in.data_ptr(),
+        .residual_out = residual_out.data_ptr(),
+        .hs_out = hs_out.data_ptr(),
+        .norm_weight = norm_weight.data_ptr(),
+        .eps = static_cast<float>(eps),
+        .in_stride_t = in.stride(0),
+        .shared_stride_t = do_shared ? shared.stride(0) : 0,
+        .res_in_stride_t = residual_in.stride(0),
+        .res_out_stride_t = residual_out.stride(0),
+        .hs_stride_t = hs_out.stride(0),
+        .cache_stride_slot = cache.stride(0),
+        .cache_stride_w = cache.stride(1),
+        .conv_weight_stride_d = conv_weight.stride(0),
+        .track_idx_stride = DO_TRACK ? track_indices.stride(0) : 0,
+        .rank = static_cast<uint32_t>(rank),
+        .T = t_num,
+        .D = d_num,
+    };
+
+    // vecs_per_thread (VPT) is the tuned knob; 0 -> 1. Each VPT must still fit
+    // one block (div_ceil(vecs, VPT) <= 1024).
+    const int vpt = vecs_per_thread > 0 ? static_cast<int>(vecs_per_thread) : 1;
+    const bool pdl = enable_pdl != 0;
+    switch (vpt) {
+      case 1:
+        RuntimeCheck(vecs <= 1024, "D/8 must fit one block at VPT=1");
+        launch<1>(params, t_num, vecs, dev.unwrap(), pdl);
+        break;
+      case 2:
+        launch<2>(params, t_num, vecs, dev.unwrap(), pdl);
+        break;
+      case 3:
+        launch<3>(params, t_num, vecs, dev.unwrap(), pdl);
+        break;
+      case 4:
+        launch<4>(params, t_num, vecs, dev.unwrap(), pdl);
+        break;
+      case 6:
+        launch<6>(params, t_num, vecs, dev.unwrap(), pdl);
+        break;
+      default:
+        RuntimeCheck(false, "unsupported vecs_per_thread (use 1/2/3/4/6)");
+    }
+  }
+};
+
+// Host wrapper for the target-verify variant. DO_TRACK is accepted (to share
+// the module's template-arg string) but unused -- verify never tracks.
+template <typename DType, uint32_t kNumGPU, int W, bool USE_SILU, bool USE_RESIDUAL, bool DO_TRACK>
+struct ArSconvNormVerifyKernel {
+  static void
+  run(tvm::ffi::TensorView in,
+      tvm::ffi::TensorView residual_in,
+      tvm::ffi::TensorView residual_out,
+      tvm::ffi::TensorView hs_out,
+      tvm::ffi::TensorView norm_weight,
+      double eps,
+      tvm::ffi::TensorView cache,
+      tvm::ffi::TensorView cache_indices,
+      tvm::ffi::TensorView cache_mask,
+      tvm::ffi::TensorView conv_weight,
+      tvm::ffi::TensorView inter_out,
+      int64_t q,
+      int64_t mc_stage_ptr,
+      int64_t local_stage_ptr,
+      int64_t flag_ptrs_dev,
+      int64_t state_ptr,
+      int64_t rank,
+      int64_t enable_pdl,
+      tvm::ffi::TensorView shared) {
+    using namespace host;
+    auto T = SymbolicSize{"T"};
+    auto B = SymbolicSize{"B"};
+    auto D = SymbolicSize{"D"};
+    auto Wd = SymbolicSize{"W"};
+    auto W1s = SymbolicSize{"W_minus_1"};
+    auto dev = SymbolicDevice{};
+    dev.set_options<kDLCUDA>();
+    Wd.set_value(W);
+    W1s.set_value(W - 1);
+
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(in);
+    const bool do_shared = shared.numel() > 0;
+    if (do_shared) {
+      TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(shared);
+      RuntimeCheck(shared.stride(0) % kVecElems == 0, "shared row stride must keep 16B alignment");
+      RuntimeCheck(std::bit_cast<intptr_t>(shared.data_ptr()) % 16 == 0, "shared not 16B aligned");
+    }
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(residual_in);
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(residual_out);
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(hs_out);
+    TensorMatcher({D}).with_dtype<DType>().with_device(dev).verify(norm_weight);
+    TensorMatcher({-1, W1s, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({B}).with_dtype<int32_t>().with_device(dev).verify(cache_indices);
+    TensorMatcher({B}).with_device(dev).verify(cache_mask);
+    TensorMatcher({D, Wd}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(conv_weight);
+    const uint32_t t_num = static_cast<uint32_t>(T.unwrap());
+    const uint32_t b_num = static_cast<uint32_t>(B.unwrap());
+    const uint32_t d_num = static_cast<uint32_t>(D.unwrap());
+    RuntimeCheck(q > 0 && t_num == b_num * static_cast<uint32_t>(q), "T must equal B * draft_token_num");
+    RuntimeCheck(
+        t_num >= 1 && t_num <= inkling_ar::kMaxBarrierBlocks,
+        "T must be in [1, kMaxBarrierBlocks] (one barrier slot per token)");
+    RuntimeCheck(d_num % kVecElems == 0, "D must be a multiple of 8");
+    RuntimeCheck(d_num / kVecElems <= 1024, "D/8 must fit one block");
+    RuntimeCheck(cache.stride(2) == 1, "cache must be channel-contiguous");
+    // inter_out: [max_bs, q, W-1, D], channel-contiguous, batch B rows used.
+    auto MB = SymbolicSize{"max_bs"};
+    auto Qs = SymbolicSize{"q"};
+    Qs.set_value(q);
+    TensorMatcher({MB, Qs, W1s, D}).with_dtype<DType>().with_device(dev).verify(inter_out);
+    RuntimeCheck(MB.unwrap() >= b_num, "inter_out batch dim too small");
+    RuntimeCheck(inter_out.stride(3) == 1, "inter_out must be channel-contiguous");
+    RuntimeCheck(mc_stage_ptr % 16 == 0, "mc_stage_ptr not 16B aligned");
+    RuntimeCheck(local_stage_ptr != 0 && local_stage_ptr % 16 == 0, "bad local_stage_ptr");
+    RuntimeCheck(flag_ptrs_dev != 0 && state_ptr != 0, "null barrier resources");
+    RuntimeCheck(rank >= 0 && rank < kNumGPU, "rank out of range");
+    RuntimeCheck(in.stride(0) % kVecElems == 0, "in row stride must keep 16B alignment");
+    RuntimeCheck(std::bit_cast<intptr_t>(in.data_ptr()) % 16 == 0, "in not 16B aligned");
+
+    const auto params = ArSconvNormVerifyParams{
+        .in = in.data_ptr(),
+        .shared = do_shared ? shared.data_ptr() : nullptr,
+        .mc_stage = reinterpret_cast<void*>(mc_stage_ptr),
+        .stage = reinterpret_cast<const void*>(local_stage_ptr),
+        .flag_ptrs = reinterpret_cast<void* const*>(flag_ptrs_dev),
+        .state = reinterpret_cast<uint32_t*>(state_ptr),
+        .cache = cache.data_ptr(),
+        .cache_indices = cache_indices.data_ptr(),
+        .cache_mask = cache_mask.data_ptr(),
+        .conv_weight = conv_weight.data_ptr(),
+        .inter_out = inter_out.data_ptr(),
+        .residual_in = residual_in.data_ptr(),
+        .residual_out = residual_out.data_ptr(),
+        .hs_out = hs_out.data_ptr(),
+        .norm_weight = norm_weight.data_ptr(),
+        .eps = static_cast<float>(eps),
+        .in_stride_t = in.stride(0),
+        .shared_stride_t = do_shared ? shared.stride(0) : 0,
+        .res_in_stride_t = residual_in.stride(0),
+        .res_out_stride_t = residual_out.stride(0),
+        .hs_stride_t = hs_out.stride(0),
+        .cache_stride_slot = cache.stride(0),
+        .cache_stride_w = cache.stride(1),
+        .conv_weight_stride_d = conv_weight.stride(0),
+        .inter_stride_b = inter_out.stride(0),
+        .inter_stride_t = inter_out.stride(1),
+        .inter_stride_w = inter_out.stride(2),
+        .rank = static_cast<uint32_t>(rank),
+        .T = t_num,
+        .D = d_num,
+        .q = static_cast<uint32_t>(q),
+    };
+
+    const uint32_t block = min(1024u, div_ceil(d_num / kVecElems, 32u) * 32u);
+    constexpr auto kernel = inkling_ar_sconv_norm_verify_kernel<DType, kNumGPU, W, USE_SILU, USE_RESIDUAL>;
+    // The kernel grid-strides over tokens with ONE grid_system_barrier between
+    // the push and the neighbor-reading reduce, so all blocks must be
+    // co-resident (else the leader waits forever). Cap the grid at the
+    // occupancy limit; the token loop covers any remaining rows.
+    const uint32_t bps = host::runtime::get_blocks_per_sm(kernel, block);
+    const uint32_t cap = host::runtime::get_sm_count(dev.unwrap().device_id) * max(1u, bps);
+    const uint32_t grid = min(t_num, cap);
+    LaunchKernel(dim3{grid}, dim3{block}, dev.unwrap()).enable_pdl(enable_pdl != 0)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/inkling/inkling_rel_proj.cuh

@@ -0,0 +1,145 @@
+// Latency-lean rel_logits projection for SMALL token counts:
+// out[t, h, :] = bf16(sum_d fp32(r[t, h, d]) * fp32(proj[d, :])) with an
+// optional per-token tau prescale folded in registers (the shipped prescale
+// semantics: r*tau rounds to bf16 BEFORE the dot, matching
+// {row_scale -> einsum} exactly).
+//
+// At t=1 the cuBLAS GEMM ([16,16]@[16,1024]) is pure launch + entry overhead
+// (~1.6 us for ~64 KB of traffic); this kernel is a no-smem no-sync grid of
+// independent 8-wide dots reading proj straight from L2 (32 KB, hot across
+// decode steps), so its floor is the launch itself. An earlier smem-staged
+// bandwidth-oriented kernel lost to cuBLAS at EVERY size -- this one is only
+// dispatched inside its measured small-t band; large t stays on cuBLAS.
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil
+
+#include <sgl_kernel/type.cuh>   // For bf16_t/fp32_t aliases
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, PDL helpers
+#include <sgl_kernel/vec.cuh>    // For AlignedVector (16B loads)
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+constexpr uint32_t kRpVec = 8;  // bf16x8 = 16 B
+constexpr uint32_t kRpBlock = 256;
+
+template <int kDRel, bool kUsePDL, bool kHasTau>
+__global__ __launch_bounds__(kRpBlock, 1) void rel_proj_small_t_kernel(
+    const bf16_t* __restrict__ r,     // [t, h, kDRel], token rows strided
+    const fp32_t* __restrict__ tau,   // [t]; unread when !kHasTau
+    const bf16_t* __restrict__ proj,  // [kDRel, e] contiguous
+    bf16_t* __restrict__ out,         // [t, h, e] contiguous
+    const int64_t r_stride_t,         // elems between token rows
+    const uint32_t h,
+    const uint32_t e,
+    const uint32_t t) {
+  using namespace device;
+  PDLWaitPrimary<kUsePDL>();
+  const uint32_t evecs = e / kRpVec;
+  const uint32_t total = t * h * evecs;
+  for (uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += gridDim.x * blockDim.x) {
+    const uint32_t ev = idx % evecs;
+    const uint32_t th = idx / evecs;
+    const uint32_t ti = th / h;
+    const uint32_t hi = th % h;
+
+    // r[ti, hi, :] once into registers (2x 16B for kDRel=16), tau folded
+    // with the prescale rounding (bf16 round before the dot).
+    const bf16_t* rrow = r + static_cast<int64_t>(ti) * r_stride_t + static_cast<int64_t>(hi) * kDRel;
+    float rv[kDRel];
+#pragma unroll
+    for (int d = 0; d < kDRel; d += static_cast<int>(kRpVec)) {
+      AlignedVector<bf16_t, kRpVec> a;
+      a.load(rrow, d / static_cast<int>(kRpVec));
+#pragma unroll
+      for (int k = 0; k < static_cast<int>(kRpVec); ++k) {
+        if constexpr (kHasTau) {
+          rv[d + k] = static_cast<float>(static_cast<bf16_t>(static_cast<float>(a[k]) * tau[ti]));
+        } else {
+          rv[d + k] = static_cast<float>(a[k]);
+        }
+      }
+    }
+
+    float acc[kRpVec] = {};
+#pragma unroll
+    for (int d = 0; d < kDRel; ++d) {
+      AlignedVector<bf16_t, kRpVec> p;
+      p.load(proj + static_cast<int64_t>(d) * e, ev);
+#pragma unroll
+      for (int k = 0; k < static_cast<int>(kRpVec); ++k) {
+        acc[k] += rv[d] * static_cast<float>(p[k]);
+      }
+    }
+
+    AlignedVector<bf16_t, kRpVec> o;
+#pragma unroll
+    for (int k = 0; k < static_cast<int>(kRpVec); ++k) {
+      o[k] = static_cast<bf16_t>(acc[k]);
+    }
+    o.store(out, idx);
+  }
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int kDRel, bool kUsePDL>
+void rel_proj_small_t(
+    tvm::ffi::TensorView r,
+    tvm::ffi::TensorView tau,  // numel-0 sentinel = no prescale
+    tvm::ffi::TensorView proj,
+    tvm::ffi::TensorView out) {
+  using namespace host;
+  auto T = SymbolicSize{"t"};
+  auto H = SymbolicSize{"h"};
+  auto D = SymbolicSize{"d_rel"};
+  auto E = SymbolicSize{"e"};
+  auto dev = SymbolicDevice{};
+  dev.set_options<kDLCUDA>();
+
+  TensorMatcher({T, H, D}).with_dtype<bf16_t>().with_device(dev).with_strides({-1, D, 1}).verify(r);
+  TensorMatcher({D, E}).with_dtype<bf16_t>().with_device(dev).verify(proj);
+  TensorMatcher({T, H, E}).with_dtype<bf16_t>().with_device(dev).verify(out);
+
+  const uint32_t t = static_cast<uint32_t>(T.unwrap());
+  const uint32_t h = static_cast<uint32_t>(H.unwrap());
+  const uint32_t e = static_cast<uint32_t>(E.unwrap());
+  RuntimeCheck(D.unwrap() == kDRel, "d_rel must be ", kDRel);
+  static_assert(kDRel % static_cast<int>(kRpVec) == 0, "d_rel must be a vector multiple (r loads are 16B)");
+  RuntimeCheck(e % kRpVec == 0, "e must be a multiple of ", kRpVec);
+  RuntimeCheck((r.stride(0) * 2) % 16 == 0, "r token stride must keep 16B alignment");
+  RuntimeCheck(std::bit_cast<intptr_t>(r.data_ptr()) % 16 == 0, "r not 16B aligned");
+  RuntimeCheck(std::bit_cast<intptr_t>(proj.data_ptr()) % 16 == 0, "proj not 16B aligned");
+
+  const bool has_tau = tau.numel() > 0;
+  if (has_tau) {
+    TensorMatcher({T}).with_dtype<fp32_t>().with_device(dev).verify(tau);
+  }
+
+  const uint32_t total = t * h * (e / kRpVec);
+  const uint32_t grid = div_ceil(total, kRpBlock);
+  auto launch = [&](auto kernel) {
+    LaunchKernel(grid, kRpBlock, dev.unwrap())
+        .enable_pdl(kUsePDL)(
+            kernel,
+            static_cast<const bf16_t*>(r.data_ptr()),
+            has_tau ? static_cast<const fp32_t*>(tau.data_ptr()) : nullptr,
+            static_cast<const bf16_t*>(proj.data_ptr()),
+            static_cast<bf16_t*>(out.data_ptr()),
+            r.stride(0),
+            h,
+            e,
+            t);
+  };
+  if (has_tau) {
+    launch(rel_proj_small_t_kernel<kDRel, kUsePDL, true>);
+  } else {
+    launch(rel_proj_small_t_kernel<kDRel, kUsePDL, false>);
+  }
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/inkling/inkling_row_scale.cuh

@@ -0,0 +1,116 @@
+// Vectorized per-row scale for the Inkling log-scaling tau paths:
+// out[row, :] = bf16(fp32(x[row, :]) * tau[row]) -- the apply_log_scaling_tau
+// contract (fp32 multiply, one bf16 round), replacing the scalar triton
+// kernel (per-ELEMENT int64 div/mod + tau load; ~1.7 us at 512 B in-graph,
+// ~2.5x off the copy floor at 16k rows) with 16 B vector loads/stores and one
+// row divide per vector. x may be row-strided (a slice of the packed qkvr
+// projection); out is contiguous.
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil
+
+#include <sgl_kernel/runtime.cuh>  // For get_blocks_per_sm / get_sm_count
+#include <sgl_kernel/type.cuh>     // For bf16_t/fp32_t aliases
+#include <sgl_kernel/utils.cuh>    // For LaunchKernel, PDL helpers
+#include <sgl_kernel/vec.cuh>      // For AlignedVector (16B loads)
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+constexpr uint32_t kRsVec = 8;  // bf16x8 = 16 B
+constexpr uint32_t kRsBlock = 256;
+
+// kHasTau=false is the pure row-compaction flavor (tau may be nullptr): same
+// vectorized strided-rows -> contiguous copy, no multiply. It replaces the
+// TensorIterator copy hidden inside einsum's reshape of the strided r operand
+// (measured ~2.3 us slower per call at decode sizes).
+template <bool kUsePDL, bool kHasTau>
+__global__ __launch_bounds__(kRsBlock, 1) void row_scale_kernel(
+    const bf16_t* __restrict__ x,    // [rows, inner], row-strided
+    const fp32_t* __restrict__ tau,  // [rows]; unread when !kHasTau
+    bf16_t* __restrict__ out,        // [rows, inner] contiguous
+    const int64_t x_stride_row,      // elems
+    const uint32_t inner,
+    const uint32_t rows) {
+  using namespace device;
+  PDLWaitPrimary<kUsePDL>();
+  const uint32_t vrow = inner / kRsVec;
+  const uint32_t total = rows * vrow;
+  for (uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += gridDim.x * blockDim.x) {
+    const uint32_t row = idx / vrow;
+    const uint32_t v = idx % vrow;
+    AlignedVector<bf16_t, kRsVec> a;
+    a.load(x + static_cast<int64_t>(row) * x_stride_row, v);
+    if constexpr (kHasTau) {
+      const float tv = tau[row];
+#pragma unroll
+      for (int k = 0; k < static_cast<int>(kRsVec); ++k) {
+        a[k] = static_cast<bf16_t>(static_cast<float>(a[k]) * tv);
+      }
+    }
+    a.store(out, idx);
+  }
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <bool kUsePDL, bool kHasTau>
+void row_scale_launch(
+    tvm::ffi::TensorView x,
+    const fp32_t* tau_ptr,
+    tvm::ffi::TensorView out,
+    host::SymbolicSize& R,
+    host::SymbolicSize& N,
+    host::SymbolicDevice& dev) {
+  using namespace host;
+  TensorMatcher({R, N}).with_dtype<bf16_t>().with_device(dev).with_strides({-1, 1}).verify(x);
+  TensorMatcher({R, N}).with_dtype<bf16_t>().with_device(dev).verify(out);
+
+  const uint32_t rows = static_cast<uint32_t>(R.unwrap());
+  const uint32_t inner = static_cast<uint32_t>(N.unwrap());
+  RuntimeCheck(inner % kRsVec == 0, "inner must be a multiple of ", kRsVec);
+  RuntimeCheck((x.stride(0) * 2) % 16 == 0, "x row stride must keep 16B alignment");
+  RuntimeCheck(std::bit_cast<intptr_t>(x.data_ptr()) % 16 == 0, "x not 16B aligned");
+
+  const auto kernel = row_scale_kernel<kUsePDL, kHasTau>;
+  const uint32_t sm = runtime::get_sm_count(dev.unwrap().device_id);
+  const uint32_t bps = runtime::get_blocks_per_sm(kernel, kRsBlock);
+  const uint32_t want = div_ceil(rows * (inner / kRsVec), kRsBlock);
+  const uint32_t grid = std::min(sm * std::max(1u, bps), std::max(1u, want));
+  LaunchKernel(grid, kRsBlock, dev.unwrap())
+      .enable_pdl(kUsePDL)(
+          kernel,
+          static_cast<const bf16_t*>(x.data_ptr()),
+          tau_ptr,
+          static_cast<bf16_t*>(out.data_ptr()),
+          x.stride(0),
+          inner,
+          rows);
+}
+
+template <bool kUsePDL>
+void row_scale(tvm::ffi::TensorView x, tvm::ffi::TensorView tau, tvm::ffi::TensorView out) {
+  using namespace host;
+  auto R = SymbolicSize{"rows"};
+  auto N = SymbolicSize{"inner"};
+  auto dev = SymbolicDevice{};
+  dev.set_options<kDLCUDA>();
+  TensorMatcher({R}).with_dtype<fp32_t>().with_device(dev).verify(tau);
+  row_scale_launch<kUsePDL, true>(x, static_cast<const fp32_t*>(tau.data_ptr()), out, R, N, dev);
+}
+
+// Pure compaction: out = contiguous copy of the row-strided x (no tau).
+template <bool kUsePDL>
+void row_compact(tvm::ffi::TensorView x, tvm::ffi::TensorView out) {
+  using namespace host;
+  auto R = SymbolicSize{"rows"};
+  auto N = SymbolicSize{"inner"};
+  auto dev = SymbolicDevice{};
+  dev.set_options<kDLCUDA>();
+  row_scale_launch<kUsePDL, false>(x, nullptr, out, R, N, dev);
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/inkling/update_sconv_cache.cuh

@@ -0,0 +1,138 @@
+// Update the convolution cache from an extend/prefill token stream.
+//
+// For each sequence b with slot ci = cache_indices[b] and query range
+// [start, end) (query_start_loc), the new conv state is the last W1 = W-1 entries of
+// the virtual stream  [ old_state (W1 rows, gated by has_initial_state) ++ x[start:end] ]:
+//     new_state[w] = virtual[qlen + w]   for w in 0..W1-1  (qlen = end - start)
+//       qlen + w >= W1  -> x[end - W1 + w, d]                          (a "current" token)
+//       qlen + w <  W1  -> old_cache[slot, w + qlen, d] * has_state    (shifted state)
+// PAD (ci == -1) or empty (qlen <= 0) lanes are left untouched. This is a pure
+// select/copy (no arithmetic) => must be BIT-EXACT (bf16 values moved verbatim).
+//
+// RAW-safe: each thread loads all W1 old_cache rows into registers BEFORE writing any,
+// so the in-place writes never clobber a not-yet-read shift source. 2 channels/thread
+// are packed as bf16x2 (32-bit) to halve the moves. Requires bf16 + even D.
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+#include <cuda_bf16.h>
+
+namespace {
+
+struct UpdateSconvParams {
+  const void* __restrict__ x;              // [T, D], channel-contiguous
+  void* __restrict__ cache;                // [max_slots, W1, D], in-place update
+  const void* __restrict__ cache_indices;  // int32 [B]
+  const void* __restrict__ has_state;      // bool  [B]
+  const void* __restrict__ qsl;            // int32 [B+1] query_start_loc
+  int64_t x_stride_t;
+  int64_t cache_stride_slot;
+  int64_t cache_stride_w;
+  uint32_t D;
+};
+
+constexpr uint32_t kUpdThreads = 256;  // threads/block, each owns a channel pair
+constexpr int kPadSlot = -1;
+
+template <int W1, typename DType>
+__global__ void update_sconv_cache_kernel(const __grid_constant__ UpdateSconvParams p) {
+  const int b = blockIdx.y;
+  const int ci = static_cast<const int32_t*>(p.cache_indices)[b];
+  const int start = static_cast<const int32_t*>(p.qsl)[b];
+  const int end = static_cast<const int32_t*>(p.qsl)[b + 1];
+  const int qlen = end - start;
+  if (ci == kPadSlot || qlen <= 0) return;  // PAD / empty lane: untouched
+
+  const int c0 = (blockIdx.x * kUpdThreads + threadIdx.x) * 2;
+  if (c0 >= static_cast<int>(p.D)) return;
+
+  const bool hs = static_cast<const bool*>(p.has_state)[b];
+  const auto* xp = static_cast<const __nv_bfloat16*>(p.x);
+  auto* cp = static_cast<__nv_bfloat16*>(p.cache);
+  const int cw = static_cast<int>(p.cache_stride_w);
+  const int64_t slot_base = static_cast<int64_t>(ci) * p.cache_stride_slot + c0;
+
+  // Load all old-state rows into registers first (RAW-safe against the writes below).
+  __nv_bfloat162 old_reg[W1];
+#pragma unroll
+  for (int w = 0; w < W1; ++w) {
+    old_reg[w] = *reinterpret_cast<const __nv_bfloat162*>(&cp[slot_base + static_cast<int64_t>(w) * cw]);
+  }
+  const __nv_bfloat162 zero = __float2bfloat162_rn(0.0f);
+
+#pragma unroll
+  for (int w = 0; w < W1; ++w) {
+    __nv_bfloat162 nv;
+    if (qlen >= (W1 - w)) {
+      // current token from x: index end - W1 + w >= start >= 0
+      const int x_idx = end - W1 + w;
+      nv = *reinterpret_cast<const __nv_bfloat162*>(&xp[static_cast<int64_t>(x_idx) * p.x_stride_t + c0]);
+    } else {
+      // shifted state old_cache[w + qlen] (w+qlen in [0, W1)), gated by has_state
+      __nv_bfloat162 shift = zero;
+#pragma unroll
+      for (int src = 0; src < W1; ++src) {
+        if (src == w + qlen) shift = old_reg[src];
+      }
+      nv = hs ? shift : zero;
+    }
+    *reinterpret_cast<__nv_bfloat162*>(&cp[slot_base + static_cast<int64_t>(w) * cw]) = nv;
+  }
+}
+
+template <int W1, typename DType>
+struct UpdateSconvCacheKernel {
+  static void
+  run(tvm::ffi::TensorView x,
+      tvm::ffi::TensorView cache,
+      tvm::ffi::TensorView cache_indices,
+      tvm::ffi::TensorView has_state,
+      tvm::ffi::TensorView qsl) {
+    using namespace host;
+    auto T = SymbolicSize{"T"};
+    auto D = SymbolicSize{"D"};
+    auto W1s = SymbolicSize{"W_minus_1"};
+    auto B = SymbolicSize{"B"};
+    auto dev = SymbolicDevice{};
+    dev.set_options<kDLCUDA>();
+    W1s.set_value(W1);
+
+    // x channel-contiguous (may be a non-contiguous row view); cache contiguous
+    // [slots, W1, D]. cache_indices/qsl int32, has_state torch-bool (shape/device only).
+    TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(x);
+    TensorMatcher({-1, W1s, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({B}).with_dtype<int32_t>().with_device(dev).verify(cache_indices);
+    TensorMatcher({B}).with_device(dev).verify(has_state);
+    TensorMatcher({-1}).with_dtype<int32_t>().with_device(dev).verify(qsl);
+    RuntimeCheck(qsl.size(0) == B.unwrap() + 1, "qsl must have length B+1");
+    RuntimeCheck(sizeof(DType) == 2, "update_sconv_cache: bf16x2 kernel requires 16-bit dtype");
+    RuntimeCheck(D.unwrap() % 2 == 0, "update_sconv_cache: D must be even for the bf16x2 kernel");
+    RuntimeCheck(cache.stride(2) == 1, "update_sconv_cache: cache must be channel-contiguous");
+
+    const auto params = UpdateSconvParams{
+        .x = x.data_ptr(),
+        .cache = cache.data_ptr(),
+        .cache_indices = cache_indices.data_ptr(),
+        .has_state = has_state.data_ptr(),
+        .qsl = qsl.data_ptr(),
+        .x_stride_t = x.stride(0),
+        .cache_stride_slot = cache.stride(0),
+        .cache_stride_w = cache.stride(1),
+        .D = static_cast<uint32_t>(D.unwrap()),
+    };
+
+    const uint32_t d_pairs = params.D / 2;
+    const dim3 grid{div_ceil(d_pairs, kUpdThreads), static_cast<uint32_t>(B.unwrap())};
+    const dim3 block{kUpdThreads};
+    constexpr auto kernel = update_sconv_cache_kernel<W1, DType>;
+    LaunchKernel(grid, block, dev.unwrap())(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/trtllm_lora_temp/moe_lora_merged_align_kernel.cu

@@ -20,11 +20,9 @@ limitations under the License.
 // and compute the merged virtual id inline (mirrors _fused_virtual_topk_ids),
 // so virtual_topk_ids is never materialized to global memory.
 //
-// Commit 1 scope: pure fusion (inline virtual id), NO EP skip. Output is
-// bucket-for-bucket equivalent to the old path (dropped/-1 tokens still land in
-// the sentinel bucket 0), so it can be asserted equal to the old kernels.
-// Only the `64 < num_buckets <= 1024` branch is implemented here; other expert
-// counts keep the old path (handled by the Python dispatcher).
+// Shared-outer and compact EP routing support up to 1024 effective buckets,
+// using fused scatter for eligible shapes and two kernels otherwise. Larger
+// domains keep the old path through the Python dispatcher.
 
 #include <sgl_kernel/tensor.h>
 #include <sgl_kernel/utils.h>

python/sglang/jit_kernel/flash_attention.py

@@ -41,6 +41,11 @@ def flash_attn_with_kvcache(
     sinks=None,
     score_mod=None,
     aux_tensors=None,
+    sfq=None,
+    sfk=None,
+    sfv=None,
+    rel_bias=None,
+    rel_bias_prep_cache=None,
     ver=3,
     out=None,
 ):
@@ -202,6 +207,11 @@ def flash_attn_with_kvcache(
             sinks=sinks,
             score_mod=score_mod,
             aux_tensors=aux_tensors,
+            sfq=sfq,
+            sfk=sfk,
+            sfv=sfv,
+            rel_bias=rel_bias,
+            rel_bias_prep_cache=rel_bias_prep_cache,
             return_softmax_lse=return_softmax_lse,
         )
     else:
@@ -236,6 +246,11 @@ def flash_attn_varlen_func(
     sinks=None,
     score_mod=None,
     aux_tensors=None,
+    sfq=None,
+    sfk=None,
+    sfv=None,
+    rel_bias=None,
+    rel_bias_prep_cache=None,
     ver=3,
     out=None,
 ):
@@ -294,6 +309,14 @@ def flash_attn_varlen_func(
             pack_gqa=pack_gqa,
             score_mod=score_mod,
             aux_tensors=aux_tensors,
+            q_descale=q_descale,
+            k_descale=k_descale,
+            v_descale=v_descale,
+            sfq=sfq,
+            sfk=sfk,
+            sfv=sfv,
+            rel_bias=rel_bias,
+            rel_bias_prep_cache=rel_bias_prep_cache,
             return_softmax_lse=return_softmax_lse,
         )
     else:

python/sglang/jit_kernel/flash_attention_v4.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from typing import Callable, Optional, Tuple, Union
 
 import torch
@@ -7,7 +8,14 @@ import torch
 from sglang.kernel_api_logging import debug_kernel_api
 
 try:
-    from flash_attn.cute import flash_attn_varlen_func as _flash_attn_varlen_func
+    if os.environ.get("SGLANG_INKLING_FA4_USE_PIP") == "1":
+        # A/B debug escape hatch: route through the pip flash-attn-4 package
+        # (dev's stack). rel_bias is vendored-only, so SHEARED must be 0.
+        from flash_attn.cute import flash_attn_varlen_func as _flash_attn_varlen_func
+    else:
+        from sglang.jit_kernel.flash_attn.cute import (
+            flash_attn_varlen_func as _flash_attn_varlen_func,
+        )
 except Exception as _e:  # pragma: no cover
     _flash_attn_varlen_func = None
     _flash_attn_import_error = _e
@@ -41,12 +49,30 @@ def flash_attn_varlen_func(
     pack_gqa: Optional[bool] = None,
     score_mod: Optional[Callable] = None,
     aux_tensors: Optional[list] = None,
+    q_descale: Optional[
+        torch.Tensor
+    ] = None,  # legacy per-tensor FP8 descale scalar (fp8_e4m3/e5m2 KV)
+    k_descale: Optional[torch.Tensor] = None,  # legacy per-tensor FP8 descale scalar
+    v_descale: Optional[torch.Tensor] = None,  # legacy per-tensor FP8 descale scalar
+    sfq: Optional[
+        torch.Tensor
+    ] = None,  # MXFP8 UE8M0 per-32-elem block scales (block-scaled QK^T)
+    sfk: Optional[
+        torch.Tensor
+    ] = None,  # MXFP8 UE8M0 per-32-elem block scales (block-scaled QK^T)
+    sfv: Optional[
+        torch.Tensor
+    ] = None,  # MXFP8 UE8M0 per-32-elem block scales (in-kernel V dequant)
+    rel_bias: Optional[torch.Tensor] = None,
+    rel_bias_prep_cache: Optional[dict] = None,
     return_softmax_lse: bool = False,
+    **_: object,
 ):
     if _flash_attn_varlen_func is None:  # pragma: no cover
         raise ImportError(
-            "Vendored FlashAttention CUTE is not available (cannot import "
-            "flash_attn.cute). Please check your source tree."
+            "FlashAttention-4 CUTE is not available. Install flash-attn-4 with "
+            "its CUDA/CUTE dependencies, or run from a source tree where the "
+            "vendored FA4 package is importable."
         ) from _flash_attn_import_error
 
     q, k, v = [_maybe_contiguous(t) for t in (q, k, v)]
@@ -62,6 +88,32 @@ def flash_attn_varlen_func(
     if window_size == (-1, -1):
         window_size = (None, None)
 
+    # sf* = MXFP8 UE8M0 block scale factors (per-32-element), for the
+    # block-scaled QK^T / V-dequant path. *_descale = the legacy per-tensor
+    # FP8 descale scalars (kv_cache_dtype fp8_e4m3/fp8_e5m2). Only one group is
+    # ever populated for a given call. Non-None kwargs only, so bf16/other calls
+    # don't hand these to the kernel.
+    sf_kwargs = {}
+    if sfq is not None:
+        sf_kwargs["sfq"] = sfq
+    if sfk is not None:
+        sf_kwargs["sfk"] = sfk
+    if sfv is not None:
+        sf_kwargs["sfv"] = sfv
+
+    descale_kwargs = {}
+    if q_descale is not None:
+        descale_kwargs["q_descale"] = q_descale
+    if k_descale is not None:
+        descale_kwargs["k_descale"] = k_descale
+    if v_descale is not None:
+        descale_kwargs["v_descale"] = v_descale
+
+    rel_bias_kwargs = {}
+    if rel_bias is not None:
+        rel_bias_kwargs["rel_bias"] = rel_bias
+    if rel_bias_prep_cache is not None:
+        rel_bias_kwargs["rel_bias_prep_cache"] = rel_bias_prep_cache
     result = _flash_attn_varlen_func(
         q=q,
         k=k,
@@ -83,6 +135,9 @@ def flash_attn_varlen_func(
         score_mod=score_mod,
         aux_tensors=aux_tensors,
         return_lse=return_softmax_lse,
+        **sf_kwargs,
+        **descale_kwargs,
+        **rel_bias_kwargs,
     )
 
     if return_softmax_lse:
@@ -126,6 +181,11 @@ def flash_attn_with_kvcache(
     sinks: Optional[torch.Tensor] = None,
     score_mod: Optional[Callable] = None,
     aux_tensors: Optional[list] = None,
+    sfq: Optional[torch.Tensor] = None,
+    sfk: Optional[torch.Tensor] = None,
+    sfv: Optional[torch.Tensor] = None,
+    rel_bias: Optional[torch.Tensor] = None,
+    rel_bias_prep_cache: Optional[dict] = None,
     return_softmax_lse: bool = False,
     **_: object,
 ):
@@ -137,9 +197,6 @@ def flash_attn_with_kvcache(
         raise NotImplementedError(
             "FA4 path does not support non-consecutive batch indices or left padding."
         )
-    if q_descale is not None or k_descale is not None or v_descale is not None:
-        raise NotImplementedError("FA4 path does not support descale.")
-
     if isinstance(cache_seqlens, int):
         cache_seqlens = torch.full(
             (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device
@@ -157,11 +214,19 @@ def flash_attn_with_kvcache(
         causal=causal,
         softcap=softcap if softcap != 0.0 else None,
         window_size=window_size,
-        num_splits=num_splits if num_splits != 0 else 1,
+        num_splits=num_splits,
         pack_gqa=pack_gqa,
         learnable_sink=sinks,
         score_mod=score_mod,
         aux_tensors=aux_tensors,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
+        sfq=sfq,
+        sfk=sfk,
+        sfv=sfv,
+        rel_bias=rel_bias,
+        rel_bias_prep_cache=rel_bias_prep_cache,
         return_softmax_lse=True,
     )
 

python/sglang/jit_kernel/flash_attn/cute/.flake8

@@ -0,0 +1,4 @@
+[flake8]
+max-line-length = 100
+# W503: line break before binary operator
+ignore = E731, E741, F841, W503

python/sglang/jit_kernel/flash_attn/cute/AUTHORS

@@ -0,0 +1,8 @@
+Tri Dao
+Jay Shah
+Ted Zadouri
+Markus Hoehnerbach
+Vijay Thakkar
+Timmy Liu
+Driss Guessous
+Reuben Stern

python/sglang/jit_kernel/flash_attn/cute/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

python/sglang/jit_kernel/flash_attn/cute/MANIFEST.in

@@ -0,0 +1,5 @@
+global-exclude *.egg-info/*
+prune flash_attn_4.egg-info
+prune flash_attn.egg-info
+prune build
+prune dist

python/sglang/jit_kernel/flash_attn/cute/README.md

@@ -0,0 +1,33 @@
+# FlashAttention-4 (CuTeDSL)
+
+FlashAttention-4 is a CuTeDSL-based implementation of FlashAttention for Hopper and Blackwell GPUs.
+
+## Installation
+
+```sh
+pip install flash-attn-4
+```
+
+If you're on CUDA 13, install with the `cu13` extra for best performance:
+
+```sh
+pip install "flash-attn-4[cu13]"
+```
+
+## Usage
+
+```python
+from flash_attn.cute import flash_attn_func, flash_attn_varlen_func
+
+out = flash_attn_func(q, k, v, causal=True)
+```
+
+## Development
+
+```sh
+git clone https://github.com/Dao-AILab/flash-attention.git
+cd flash-attention
+pip install -e "flash_attn/cute[dev]"       # CUDA 12.x
+pip install -e "flash_attn/cute[dev,cu13]"  # CUDA 13.x (e.g. B200)
+pytest tests/cute/
+```

python/sglang/jit_kernel/flash_attn/cute/__init__.py

@@ -0,0 +1,18 @@
+"""Flash Attention CUTE (CUDA Template Engine) implementation."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("fa4")
+except PackageNotFoundError:
+    __version__ = "0.0.0"
+
+from .interface import (
+    flash_attn_func,
+    flash_attn_varlen_func,
+)
+
+__all__ = [
+    "flash_attn_func",
+    "flash_attn_varlen_func",
+]

python/sglang/jit_kernel/flash_attn/cute/ampere_helpers.py

@@ -0,0 +1,122 @@
+# Copyright (c) 2025, Tri Dao.
+from typing import Callable, Optional, Type
+
+import cutlass
+import cutlass.cute as cute
+
+
+def get_smem_layout_atom(
+    dtype: Type[cutlass.Numeric], k_dim: int
+) -> cute.ComposedLayout:
+    dtype_byte = cutlass.const_expr(dtype.width // 8)
+    bytes_per_row = cutlass.const_expr(k_dim * dtype_byte)
+    smem_k_block_size = (
+        cutlass.const_expr(
+            128
+            if bytes_per_row % 128 == 0
+            else (
+                64
+                if bytes_per_row % 64 == 0
+                else (32 if bytes_per_row % 32 == 0 else 16)
+            )
+        )
+        // dtype_byte
+    )
+    swizzle_bits = (
+        4
+        if smem_k_block_size == 128
+        else (3 if smem_k_block_size == 64 else (2 if smem_k_block_size == 32 else 1))
+    )
+    swizzle_base = 2 if dtype_byte == 4 else (3 if dtype_byte == 2 else 4)
+    return cute.make_composed_layout(
+        cute.make_swizzle(swizzle_bits, swizzle_base, swizzle_base),
+        0,
+        cute.make_ordered_layout(
+            (8 if cutlass.const_expr(k_dim % 32 == 0) else 16, smem_k_block_size),
+            order=(1, 0),
+        ),
+    )
+
+
+@cute.jit
+def gemm(
+    tiled_mma: cute.TiledMma,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    tCsA: cute.Tensor,
+    tCsB: cute.Tensor,
+    smem_thr_copy_A: cute.TiledCopy,
+    smem_thr_copy_B: cute.TiledCopy,
+    hook_fn: Optional[Callable] = None,
+    A_in_regs: cutlass.Constexpr[bool] = False,
+    B_in_regs: cutlass.Constexpr[bool] = False,
+    swap_AB: cutlass.Constexpr[bool] = False,
+) -> None:
+    if cutlass.const_expr(swap_AB):
+        gemm(
+            tiled_mma,
+            acc,
+            tCrB,
+            tCrA,
+            tCsB,
+            tCsA,
+            smem_thr_copy_B,
+            smem_thr_copy_A,
+            hook_fn,
+            A_in_regs=B_in_regs,
+            B_in_regs=A_in_regs,
+            swap_AB=False,
+        )
+    else:
+        tCrA_copy_view = smem_thr_copy_A.retile(tCrA)
+        tCrB_copy_view = smem_thr_copy_B.retile(tCrB)
+        if cutlass.const_expr(not A_in_regs):
+            cute.copy(
+                smem_thr_copy_A, tCsA[None, None, 0], tCrA_copy_view[None, None, 0]
+            )
+        if cutlass.const_expr(not B_in_regs):
+            cute.copy(
+                smem_thr_copy_B, tCsB[None, None, 0], tCrB_copy_view[None, None, 0]
+            )
+        for k in cutlass.range_constexpr(cute.size(tCsA.shape[2])):
+            if k < cute.size(tCsA.shape[2]) - 1:
+                if cutlass.const_expr(not A_in_regs):
+                    cute.copy(
+                        smem_thr_copy_A,
+                        tCsA[None, None, k + 1],
+                        tCrA_copy_view[None, None, k + 1],
+                    )
+                if cutlass.const_expr(not B_in_regs):
+                    cute.copy(
+                        smem_thr_copy_B,
+                        tCsB[None, None, k + 1],
+                        tCrB_copy_view[None, None, k + 1],
+                    )
+            cute.gemm(tiled_mma, acc, tCrA[None, None, k], tCrB[None, None, k], acc)
+            if cutlass.const_expr(k == 0 and hook_fn is not None):
+                hook_fn()
+
+
+@cute.jit
+def gemm_rs(
+    tiled_mma: cute.TiledMma,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    tCsB: cute.Tensor,
+    smem_thr_copy_B: cute.TiledCopy,
+    hook_fn: Optional[Callable] = None,
+) -> None:
+    tCrB_copy_view = smem_thr_copy_B.retile(tCrB)
+    cute.copy(smem_thr_copy_B, tCsB[None, None, 0], tCrB_copy_view[None, None, 0])
+    for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])):
+        if cutlass.const_expr(k < cute.size(tCrA.shape[2]) - 1):
+            cute.copy(
+                smem_thr_copy_B,
+                tCsB[None, None, k + 1],
+                tCrB_copy_view[None, None, k + 1],
+            )
+        cute.gemm(tiled_mma, acc, tCrA[None, None, k], tCrB[None, None, k], acc)
+        if cutlass.const_expr(k == 0 and hook_fn is not None):
+            hook_fn()

python/sglang/jit_kernel/flash_attn/cute/barrier.py

@@ -0,0 +1,76 @@
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+
+@dsl_user_op
+def ld_acquire(lock_ptr: cute.Pointer, *, loc=None, ip=None) -> cutlass.Int32:
+    lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value()
+    state = llvm.inline_asm(
+        T.i32(),
+        [lock_ptr_i64],
+        "ld.global.acquire.gpu.b32 $0, [$1];",
+        "=r,l",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+    return cutlass.Int32(state)
+
+
+@dsl_user_op
+def red_relaxed(
+    lock_ptr: cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=None, ip=None
+) -> None:
+    lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value()
+    llvm.inline_asm(
+        None,
+        [lock_ptr_i64, Int32(val).ir_value(loc=loc, ip=ip)],
+        "red.relaxed.gpu.global.add.s32 [$0], $1;",
+        "l,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def red_release(
+    lock_ptr: cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=None, ip=None
+) -> None:
+    lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value()
+    llvm.inline_asm(
+        None,
+        [lock_ptr_i64, Int32(val).ir_value(loc=loc, ip=ip)],
+        "red.release.gpu.global.add.s32 [$0], $1;",
+        "l,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@cute.jit
+def wait_eq(
+    lock_ptr: cute.Pointer, thread_idx: int | Int32, flag_offset: int, val: Int32
+) -> None:
+    flag_ptr = lock_ptr + flag_offset
+    if thread_idx == 0:
+        read_val = Int32(0)
+        while read_val != val:
+            read_val = ld_acquire(flag_ptr)
+
+
+@cute.jit
+def arrive_inc(
+    lock_ptr: cute.Pointer,
+    thread_idx: int | Int32,
+    flag_offset: int,
+    val: cutlass.Constexpr[Int32],
+) -> None:
+    flag_ptr = lock_ptr + flag_offset
+    if thread_idx == 0:
+        red_release(flag_ptr, val)
+        # red_relaxed(flag_ptr, val)

python/sglang/jit_kernel/flash_attn/cute/bench_utils.py

@@ -0,0 +1,261 @@
+"""Shared benchmark utilities: attention_ref, cuDNN helpers, flops calculation."""
+
+import math
+
+import torch
+
+try:
+    import cudnn
+except ImportError:
+    cudnn = None
+
+
+# ── FLOPS calculation ────────────────────────────────────────────────────────
+
+
+def flops(
+    batch,
+    nheads,
+    seqlen_q,
+    seqlen_k,
+    headdim,
+    headdim_v,
+    causal=False,
+    window_size=(None, None),
+    has_qv=False,
+):
+    if causal:
+        avg_seqlen = (max(0, seqlen_k - seqlen_q) + seqlen_k) / 2
+    else:
+        if window_size == (None, None):
+            avg_seqlen = seqlen_k
+        else:
+            row_idx = torch.arange(seqlen_q, device="cuda")
+            col_left = (
+                torch.maximum(
+                    row_idx + seqlen_k - seqlen_q - window_size[0], torch.tensor(0)
+                )
+                if window_size[0] is not None
+                else torch.zeros_like(row_idx)
+            )
+            col_right = (
+                torch.minimum(
+                    row_idx + seqlen_k - seqlen_q + window_size[1],
+                    torch.tensor(seqlen_k - 1),
+                )
+                if window_size[1] is not None
+                else torch.full_like(row_idx, seqlen_k - 1)
+            )
+            avg_seqlen = (col_right - col_left + 1).float().mean().item()
+    eff_headdim = headdim + headdim_v if has_qv else headdim
+    return batch * nheads * 2 * seqlen_q * avg_seqlen * (eff_headdim + headdim_v)
+
+
+# ── Bandwidth calculation ────────────────────────────────────────────────────
+
+
+def bandwidth_fwd_bytes(
+    batch,
+    nheads,
+    nheads_kv,
+    seqlen_q,
+    seqlen_k,
+    headdim,
+    headdim_v,
+    dtype_bytes=2,
+    has_qv=False,
+):
+    """HBM traffic for one attention pass: read Q,K,V + write O."""
+    q = batch * nheads * seqlen_q * headdim
+    qv = batch * nheads * seqlen_q * headdim_v if has_qv else 0
+    k = batch * nheads_kv * seqlen_k * headdim
+    v = batch * nheads_kv * seqlen_k * headdim_v
+    o = batch * nheads * seqlen_q * headdim_v
+    return (q + qv + k + v + o) * dtype_bytes
+
+
+def bandwidth_bwd_bytes(
+    batch, nheads, nheads_kv, seqlen_q, seqlen_k, headdim, headdim_v, dtype_bytes=2
+):
+    """HBM traffic for one attention pass: read Q,K,V,dO + write dQ,dK,dV."""
+    q = batch * nheads * seqlen_q * headdim
+    k = batch * nheads_kv * seqlen_k * headdim
+    v = batch * nheads_kv * seqlen_k * headdim_v
+    do = batch * nheads * seqlen_q * headdim_v
+    dq = q
+    dk = k
+    dv = v
+    return (q + k + v + do + dq + dk + dv) * dtype_bytes
+
+
+# ── Reference attention ─────────────────────────────────────────────────────
+
+_attention_ref_mask_cache = {}
+
+
+def attention_ref(q, k, v, causal=False):
+    """Standard attention reference implementation.
+
+    Args:
+        q, k, v: (batch, seqlen, nheads, headdim) tensors.
+        causal: whether to apply causal mask.
+    """
+    softmax_scale = 1.0 / math.sqrt(q.shape[-1])
+    scores = torch.einsum("bthd,bshd->bhts", q * softmax_scale, k)
+    if causal:
+        if scores.shape[-2] not in _attention_ref_mask_cache:
+            mask = torch.tril(
+                torch.ones(scores.shape[-2:], device=scores.device, dtype=torch.bool),
+                diagonal=0,
+            )
+            _attention_ref_mask_cache[scores.shape[-2]] = mask
+        else:
+            mask = _attention_ref_mask_cache[scores.shape[-2]]
+        scores = scores.masked_fill(mask, float("-inf"))
+    attn = torch.softmax(scores, dim=-1)
+    return torch.einsum("bhts,bshd->bthd", attn, v)
+
+
+# ── cuDNN graph helpers ─────────────────────────────────────────────────────
+
+_TORCH_TO_CUDNN_DTYPE = {
+    torch.float16: "HALF",
+    torch.bfloat16: "BFLOAT16",
+    torch.float32: "FLOAT",
+    torch.int32: "INT32",
+    torch.int64: "INT64",
+}
+
+
+def _build_cudnn_graph(io_dtype, tensors, build_fn):
+    """Build a cuDNN graph.  Returns (graph, variant_pack, workspace)."""
+    assert cudnn is not None, "cuDNN is not available"
+    cudnn_dtype = getattr(cudnn.data_type, _TORCH_TO_CUDNN_DTYPE[io_dtype])
+    graph = cudnn.pygraph(
+        io_data_type=cudnn_dtype,
+        intermediate_data_type=cudnn.data_type.FLOAT,
+        compute_data_type=cudnn.data_type.FLOAT,
+    )
+    graph_tensors = {name: graph.tensor_like(t.detach()) for name, t in tensors.items()}
+    variant_pack = build_fn(graph, graph_tensors)
+    graph.validate()
+    graph.build_operation_graph()
+    graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
+    graph.check_support()
+    graph.build_plans()
+    workspace = torch.empty(
+        graph.get_workspace_size(), device="cuda", dtype=torch.uint8
+    )
+    return graph, variant_pack, workspace
+
+
+def cudnn_fwd_setup(q, k, v, causal=False, window_size_left=None):
+    """Build a cuDNN forward SDPA graph.
+
+    Args:
+        q, k, v: (batch, nheads, seqlen, headdim) tensors (cuDNN layout).
+        causal: whether to apply causal mask.
+        window_size_left: sliding window size (None for no window).
+
+    Returns:
+        (fwd_fn, o_gpu, stats_gpu) where fwd_fn is a zero-arg callable.
+    """
+    b, nheads, seqlen_q, headdim = q.shape
+    headdim_v = v.shape[-1]
+    o_gpu = torch.empty(b, nheads, seqlen_q, headdim_v, dtype=q.dtype, device=q.device)
+    stats_gpu = torch.empty(
+        b, nheads, seqlen_q, 1, dtype=torch.float32, device=q.device
+    )
+
+    def build(graph, gt):
+        o, stats = graph.sdpa(
+            name="sdpa",
+            q=gt["q"],
+            k=gt["k"],
+            v=gt["v"],
+            is_inference=False,
+            attn_scale=1.0 / math.sqrt(headdim),
+            use_causal_mask=causal or window_size_left is not None,
+            sliding_window_length=(
+                window_size_left
+                if window_size_left is not None and not causal
+                else None
+            ),
+        )
+        o.set_output(True).set_dim(o_gpu.shape).set_stride(o_gpu.stride())
+        stats.set_output(True).set_data_type(cudnn.data_type.FLOAT)
+        return {gt["q"]: q, gt["k"]: k, gt["v"]: v, o: o_gpu, stats: stats_gpu}
+
+    graph, variant_pack, workspace = _build_cudnn_graph(
+        q.dtype, {"q": q, "k": k, "v": v}, build
+    )
+
+    def fwd_fn():
+        graph.execute(variant_pack, workspace)
+        return o_gpu
+
+    return fwd_fn, o_gpu, stats_gpu
+
+
+def cudnn_bwd_setup(q, k, v, o, g, lse, causal=False, window_size_left=None):
+    """Build a cuDNN backward SDPA graph.
+
+    Args:
+        q, k, v, o, g, lse: (batch, nheads, seqlen, dim) tensors (cuDNN layout).
+        causal: whether to apply causal mask.
+        window_size_left: sliding window size (None for no window).
+
+    Returns:
+        bwd_fn: zero-arg callable that returns (dq, dk, dv).
+    """
+    headdim = q.shape[-1]
+    dq_gpu, dk_gpu, dv_gpu = (
+        torch.empty_like(q),
+        torch.empty_like(k),
+        torch.empty_like(v),
+    )
+
+    def build(graph, gt):
+        dq, dk, dv = graph.sdpa_backward(
+            name="sdpa_backward",
+            q=gt["q"],
+            k=gt["k"],
+            v=gt["v"],
+            o=gt["o"],
+            dO=gt["g"],
+            stats=gt["lse"],
+            attn_scale=1.0 / math.sqrt(headdim),
+            use_causal_mask=causal or window_size_left is not None,
+            sliding_window_length=(
+                window_size_left
+                if window_size_left is not None and not causal
+                else None
+            ),
+            use_deterministic_algorithm=False,
+        )
+        dq.set_output(True).set_dim(dq_gpu.shape).set_stride(dq_gpu.stride())
+        dk.set_output(True).set_dim(dk_gpu.shape).set_stride(dk_gpu.stride())
+        dv.set_output(True).set_dim(dv_gpu.shape).set_stride(dv_gpu.stride())
+        return {
+            gt["q"]: q,
+            gt["k"]: k,
+            gt["v"]: v,
+            gt["o"]: o,
+            gt["g"]: g,
+            gt["lse"]: lse,
+            dq: dq_gpu,
+            dk: dk_gpu,
+            dv: dv_gpu,
+        }
+
+    graph, variant_pack, workspace = _build_cudnn_graph(
+        q.dtype,
+        {"q": q, "k": k, "v": v, "o": o, "g": g, "lse": lse},
+        build,
+    )
+
+    def bwd_fn():
+        graph.execute(variant_pack, workspace)
+        return dq_gpu, dk_gpu, dv_gpu
+
+    return bwd_fn

python/sglang/jit_kernel/flash_attn/cute/benchmark.py

@@ -0,0 +1,281 @@
+# Copyright (c) 2023, Tri Dao.
+"""Useful functions for writing test code."""
+
+import torch
+import torch.utils.benchmark as benchmark
+
+
+def benchmark_forward(
+    fn,
+    *inputs,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward pass of an arbitrary function."""
+    if verbose:
+        print(desc, "- Forward pass")
+
+    def amp_wrapper(*inputs, **kwinputs):
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            fn(*inputs, **kwinputs)
+
+    t = benchmark.Timer(
+        stmt="fn_amp(*inputs, **kwinputs)",
+        globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs},
+        num_threads=torch.get_num_threads(),
+    )
+    m = t.timeit(repeats)
+    if verbose:
+        print(m)
+    return t, m
+
+
+def benchmark_backward(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the backward pass of an arbitrary function."""
+    if verbose:
+        print(desc, "- Backward pass")
+    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+        y = fn(*inputs, **kwinputs)
+        if type(y) is tuple:
+            y = y[0]
+    if grad is None:
+        grad = torch.randn_like(y)
+    else:
+        if grad.shape != y.shape:
+            raise RuntimeError("Grad shape does not match output shape")
+
+    def f(*inputs, y, grad):
+        # Set .grad to None to avoid extra operation of gradient accumulation
+        for x in inputs:
+            if isinstance(x, torch.Tensor):
+                x.grad = None
+        y.backward(grad, retain_graph=True)
+
+    t = benchmark.Timer(
+        stmt="f(*inputs, y=y, grad=grad)",
+        globals={"f": f, "inputs": inputs, "y": y, "grad": grad},
+        num_threads=torch.get_num_threads(),
+    )
+    m = t.timeit(repeats)
+    if verbose:
+        print(m)
+    return t, m
+
+
+def benchmark_combined(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
+    if verbose:
+        print(desc, "- Forward + Backward pass")
+    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+        y = fn(*inputs, **kwinputs)
+        if type(y) is tuple:
+            y = y[0]
+    if grad is None:
+        grad = torch.randn_like(y)
+    else:
+        if grad.shape != y.shape:
+            raise RuntimeError("Grad shape does not match output shape")
+
+    def f(grad, *inputs, **kwinputs):
+        for x in inputs:
+            if isinstance(x, torch.Tensor):
+                x.grad = None
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            y = fn(*inputs, **kwinputs)
+            if type(y) is tuple:
+                y = y[0]
+        y.backward(grad, retain_graph=True)
+
+    t = benchmark.Timer(
+        stmt="f(grad, *inputs, **kwinputs)",
+        globals={
+            "f": f,
+            "fn": fn,
+            "inputs": inputs,
+            "grad": grad,
+            "kwinputs": kwinputs,
+        },
+        num_threads=torch.get_num_threads(),
+    )
+    m = t.timeit(repeats)
+    if verbose:
+        print(m)
+    return t, m
+
+
+def benchmark_fwd_bwd(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
+    return (
+        benchmark_forward(
+            fn,
+            *inputs,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+        benchmark_backward(
+            fn,
+            *inputs,
+            grad=grad,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+    )
+
+
+def benchmark_all(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
+    return (
+        benchmark_forward(
+            fn,
+            *inputs,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+        benchmark_backward(
+            fn,
+            *inputs,
+            grad=grad,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+        benchmark_combined(
+            fn,
+            *inputs,
+            grad=grad,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+    )
+
+
+def pytorch_profiler(
+    fn,
+    *inputs,
+    trace_filename=None,
+    backward=False,
+    amp=False,
+    amp_dtype=torch.float16,
+    cpu=False,
+    verbose=True,
+    **kwinputs,
+):
+    """Wrap benchmark functions in Pytorch profiler to see CUDA information."""
+    if backward:
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            out = fn(*inputs, **kwinputs)
+            if type(out) is tuple:
+                out = out[0]
+            g = torch.randn_like(out)
+    for _ in range(30):  # Warm up
+        if backward:
+            for x in inputs:
+                if isinstance(x, torch.Tensor):
+                    x.grad = None
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            out = fn(*inputs, **kwinputs)
+            if type(out) is tuple:
+                out = out[0]
+        # Backward should be done outside autocast
+        if backward:
+            out.backward(g, retain_graph=True)
+    activities = ([torch.profiler.ProfilerActivity.CPU] if cpu else []) + [
+        torch.profiler.ProfilerActivity.CUDA
+    ]
+    with torch.profiler.profile(
+        activities=activities,
+        record_shapes=True,
+        # profile_memory=True,
+        with_stack=True,
+    ) as prof:
+        if backward:
+            for x in inputs:
+                if isinstance(x, torch.Tensor):
+                    x.grad = None
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            out = fn(*inputs, **kwinputs)
+            if type(out) is tuple:
+                out = out[0]
+        if backward:
+            out.backward(g, retain_graph=True)
+    if verbose:
+        # print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=50))
+        print(prof.key_averages().table(row_limit=50))
+    if trace_filename is not None:
+        prof.export_chrome_trace(trace_filename)
+
+
+def benchmark_memory(fn, *inputs, desc="", verbose=True, **kwinputs):
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    torch.cuda.synchronize()
+    fn(*inputs, **kwinputs)
+    torch.cuda.synchronize()
+    mem = torch.cuda.max_memory_allocated() / ((2**20) * 1000)
+    if verbose:
+        print(f"{desc} max memory: {mem}GB")
+    torch.cuda.empty_cache()
+    return mem

python/sglang/jit_kernel/flash_attn/cute/benchmark_flash_attention_fp8.py

@@ -0,0 +1,487 @@
+# Benchmark FP8 attention for FA4 (CuTe-DSL) on SM100.
+#
+# Run (recommended):
+#   python -m flash_attn.cute.benchmark_flash_attention_fp8
+#
+# Notes:
+# - This is intended to be used while bringing up FP8 support for SM100.
+# - FP8 correctness depends on descales + max-offset scaling being implemented in the SM100 kernel.
+#   This script optionally checks output vs a BF16 PyTorch baseline on dequantized FP8 inputs.
+#
+# Adapted from: `hopper/benchmark_flash_attention_fp8.py`
+
+from __future__ import annotations
+
+import argparse
+import inspect
+import math
+import time
+from typing import Iterable
+
+import torch
+from einops import rearrange
+
+from sglang.jit_kernel.flash_attn.cute.benchmark import benchmark_forward
+from sglang.jit_kernel.flash_attn.cute.interface import (
+    _flash_attn_fwd as flash_attn_cute_fwd,
+)
+
+try:
+    import cudnn
+except ImportError:
+    cudnn = None
+
+
+def _torch_float8_dtype(name: str) -> torch.dtype:
+    if name in ("fp8", "fp8_e4m3", "fp8_e4m3fn"):
+        return torch.float8_e4m3fn
+    if name in ("fp8_e5m2", "fp8_e5m2fn"):
+        return torch.float8_e5m2
+    raise ValueError(f"Unsupported fp8 dtype name: {name}")
+
+
+def _parse_int_list(csv: str) -> list[int]:
+    out: list[int] = []
+    for part in csv.split(","):
+        part = part.strip()
+        if not part:
+            continue
+        out.append(int(part))
+    return out
+
+
+def attention_pytorch(qkv: torch.Tensor, causal: bool) -> torch.Tensor:
+    """
+    qkv: (batch, seqlen, 3, nheads, headdim)
+    out: (batch, seqlen, nheads, headdim)
+    """
+    batch_size, seqlen, _, nheads, d = qkv.shape
+    q, k, v = qkv.unbind(dim=2)
+    q = rearrange(q, "b t h d -> (b h) t d")
+    k = rearrange(k, "b s h d -> (b h) d s")
+    softmax_scale = 1.0 / math.sqrt(d)
+    scores = torch.empty(
+        batch_size * nheads, seqlen, seqlen, dtype=qkv.dtype, device=qkv.device
+    )
+    scores = rearrange(
+        torch.baddbmm(scores, q, k, beta=0, alpha=softmax_scale),
+        "(b h) t s -> b h t s",
+        h=nheads,
+    )
+    if causal:
+        causal_mask = torch.triu(
+            torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1
+        )
+        scores = scores + causal_mask.to(dtype=scores.dtype)
+    attention = torch.softmax(scores, dim=-1)
+    output = torch.einsum("bhts,bshd->bthd", attention, v)
+    return output.to(dtype=qkv.dtype)
+
+
+def flops(batch: int, seqlen: int, headdim: int, nheads: int, causal: bool) -> int:
+    # Matches the hopper benchmark’s convention.
+    return 4 * batch * seqlen**2 * nheads * headdim // (2 if causal else 1)
+
+
+def efficiency(flop: int, seconds: float) -> float:
+    return (flop / seconds / 1e12) if not math.isnan(seconds) else 0.0
+
+
+def time_fwd(fn, *args, repeats: int, **kwargs) -> float:
+    time.sleep(1)  # reduce residual throttling effects between benchmarks
+    _, m = benchmark_forward(fn, *args, repeats=repeats, verbose=False, **kwargs)
+    return float(m.mean)
+
+
+def convert_to_cudnn_type(torch_type):
+    if torch_type == torch.float16:
+        return cudnn.data_type.HALF
+    if torch_type == torch.bfloat16:
+        return cudnn.data_type.BFLOAT16
+    if torch_type == torch.float32:
+        return cudnn.data_type.FLOAT
+    if torch_type == torch.int32:
+        return cudnn.data_type.INT32
+    if torch_type == torch.int64:
+        return cudnn.data_type.INT64
+    if torch_type == torch.float8_e4m3fn:
+        return cudnn.data_type.FP8_E4M3
+    if torch_type == torch.float8_e5m2:
+        return cudnn.data_type.FP8_E5M2
+    raise ValueError("Unsupported tensor data type.")
+
+
+def cudnn_sdpa_fp8_setup(qkv: torch.Tensor, seqlen_q: int, seqlen_k: int, causal: bool):
+    """Minimal cudnn.fp8 sdpa runner (optional)."""
+    assert cudnn is not None, "cudnn python bindings not available"
+    b, _, _, nheads, headdim = qkv.shape
+    o_gpu = torch.zeros(
+        b, seqlen_q, nheads, headdim, dtype=qkv.dtype, device=qkv.device
+    )
+    o_gpu_transposed = torch.as_strided(
+        o_gpu,
+        [b, nheads, seqlen_q, headdim],
+        [nheads * seqlen_q * headdim, headdim, nheads * headdim, 1],
+    )
+    amax_s_gpu = torch.empty(1, 1, 1, 1, dtype=torch.float32, device=qkv.device)
+    amax_o_gpu = torch.empty(1, 1, 1, 1, dtype=torch.float32, device=qkv.device)
+
+    graph = cudnn.pygraph(
+        io_data_type=convert_to_cudnn_type(qkv.dtype),
+        intermediate_data_type=cudnn.data_type.FLOAT,
+        compute_data_type=cudnn.data_type.FLOAT,
+    )
+
+    new_q = torch.as_strided(
+        qkv,
+        [b, nheads, seqlen_q, headdim],
+        [seqlen_q * nheads * headdim * 3, headdim, headdim * nheads * 3, 1],
+        storage_offset=0,
+    )
+    q = graph.tensor(
+        name="Q",
+        dim=list(new_q.shape),
+        stride=list(new_q.stride()),
+        data_type=convert_to_cudnn_type(qkv.dtype),
+    )
+
+    new_k = torch.as_strided(
+        qkv,
+        [b, nheads, seqlen_k, headdim],
+        [seqlen_k * nheads * headdim * 3, headdim, headdim * nheads * 3, 1],
+        storage_offset=nheads * headdim,
+    )
+    k = graph.tensor(
+        name="K",
+        dim=list(new_k.shape),
+        stride=list(new_k.stride()),
+        data_type=convert_to_cudnn_type(qkv.dtype),
+    )
+
+    new_v = torch.as_strided(
+        qkv,
+        [b, nheads, seqlen_k, headdim],
+        [seqlen_k * nheads * headdim * 3, headdim, headdim * nheads * 3, 1],
+        storage_offset=nheads * headdim * 2,
+    )
+    v = graph.tensor(
+        name="V",
+        dim=list(new_v.shape),
+        stride=list(new_v.stride()),
+        data_type=convert_to_cudnn_type(qkv.dtype),
+    )
+
+    def _scale_tensor():
+        return graph.tensor(
+            dim=[1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT
+        )
+
+    default_scale_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float32, device="cuda")
+    descale_q = _scale_tensor()
+    descale_k = _scale_tensor()
+    descale_v = _scale_tensor()
+    descale_s = _scale_tensor()
+    scale_s = _scale_tensor()
+    scale_o = _scale_tensor()
+
+    o, _, amax_s, amax_o = graph.sdpa_fp8(
+        q=q,
+        k=k,
+        v=v,
+        descale_q=descale_q,
+        descale_k=descale_k,
+        descale_v=descale_v,
+        descale_s=descale_s,
+        scale_s=scale_s,
+        scale_o=scale_o,
+        is_inference=True,
+        attn_scale=1.0 / math.sqrt(headdim),
+        use_causal_mask=causal,
+        name="sdpa",
+    )
+    o.set_output(True).set_dim(o_gpu_transposed.shape).set_stride(
+        o_gpu_transposed.stride()
+    )
+    amax_s.set_output(False).set_dim(amax_s_gpu.shape).set_stride(amax_s_gpu.stride())
+    amax_o.set_output(False).set_dim(amax_o_gpu.shape).set_stride(amax_o_gpu.stride())
+
+    graph.validate()
+    graph.build_operation_graph()
+    graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
+    graph.check_support()
+    graph.build_plans()
+
+    variant_pack = {
+        q: new_q,
+        k: new_k,
+        v: new_v,
+        descale_q: default_scale_gpu,
+        descale_k: default_scale_gpu,
+        descale_v: default_scale_gpu,
+        descale_s: default_scale_gpu,
+        scale_s: default_scale_gpu,
+        scale_o: default_scale_gpu,
+        o: o_gpu_transposed,
+        amax_s: amax_s_gpu,
+        amax_o: amax_o_gpu,
+    }
+    workspace = torch.empty(
+        graph.get_workspace_size(), device="cuda", dtype=torch.uint8
+    )
+
+    def run():
+        graph.execute(variant_pack, workspace)
+        return o_gpu
+
+    return run
+
+
+def _maybe_pass_descales(callable_, **kwargs):
+    sig = inspect.signature(callable_)
+    return {k: v for k, v in kwargs.items() if k in sig.parameters}
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repeats", type=int, default=30)
+    parser.add_argument("--dim", type=int, default=2048)
+    parser.add_argument("--headdims", default="64,128")
+    parser.add_argument("--dtype", default="fp8_e4m3fn")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument(
+        "--check",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Enable correctness checks vs BF16 PyTorch baseline.",
+    )
+    parser.add_argument(
+        "--check-quantization-only",
+        action="store_true",
+        help="Check FP8 kernel vs dequantized-FP8 baseline (quantization error only).",
+    )
+    parser.add_argument("--atol-bf16", type=float, default=0.10)
+    parser.add_argument("--rtol-bf16", type=float, default=0.10)
+    parser.add_argument("--atol-fp8", type=float, default=0.50)
+    parser.add_argument("--rtol-fp8", type=float, default=0.50)
+    parser.add_argument("--run-cudnn", action="store_true")
+    args = parser.parse_args(list(argv) if argv is not None else None)
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    major, minor = torch.cuda.get_device_capability()
+    if major != 10:
+        raise RuntimeError(
+            f"This benchmark is for SM100 (compute capability 10.x). Got {major}.{minor}."
+        )
+
+    torch.manual_seed(args.seed)
+    device = "cuda"
+    fp8_dtype = _torch_float8_dtype(args.dtype)
+    headdim_vals = _parse_int_list(args.headdims)
+    bs_seqlen_vals = [
+        (32, 512),
+        (16, 1024),
+        (8, 2048),
+        (4, 4096),
+        (2, 8192),
+        (1, 16384),
+    ]
+
+    methods = ["Pytorch", "FA4-CuTe-BF16", "FA4-CuTe-FP8"] + (
+        ["cuDNN-FP8"] if args.run_cudnn and cudnn is not None else []
+    )
+
+    fp8_failures = []
+
+    for headdim in headdim_vals:
+        for causal in (False, True):
+            for batch, seqlen in bs_seqlen_vals:
+                torch.cuda.empty_cache()
+                nheads = args.dim // headdim
+                if args.dim % headdim != 0:
+                    raise ValueError(
+                        f"--dim must be divisible by headdim ({args.dim=} {headdim=})"
+                    )
+
+                q_bf16 = torch.randn(
+                    batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16
+                )
+                k_bf16 = torch.randn(
+                    batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16
+                )
+                v_bf16 = torch.randn(
+                    batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16
+                )
+                qkv_bf16 = torch.stack([q_bf16, k_bf16, v_bf16], dim=2)
+
+                times = {}
+                speeds = {}
+
+                out_ref_bf16 = None
+                try:
+                    out_ref_bf16 = attention_pytorch(
+                        qkv_bf16, causal=causal
+                    )  # warmup / reference
+                    t = time_fwd(
+                        attention_pytorch, qkv_bf16, causal=causal, repeats=args.repeats
+                    )
+                    times["Pytorch"] = t
+                except RuntimeError as e:
+                    if "out of memory" in str(e).lower():
+                        times["Pytorch"] = float("nan")
+                        out_ref_bf16 = None
+                    else:
+                        raise
+
+                # FA4 / CuTe BF16 baseline
+                try:
+                    softmax_scale = headdim**-0.5
+                    out_fa4_bf16, _ = flash_attn_cute_fwd(
+                        q_bf16,
+                        k_bf16,
+                        v_bf16,
+                        softmax_scale=softmax_scale,
+                        causal=causal,
+                    )  # warmup / compile
+                    t = time_fwd(
+                        flash_attn_cute_fwd,
+                        q_bf16,
+                        k_bf16,
+                        v_bf16,
+                        softmax_scale=softmax_scale,
+                        causal=causal,
+                        repeats=args.repeats,
+                    )
+                    times["FA4-CuTe-BF16"] = t
+                    if args.check and out_ref_bf16 is not None:
+                        torch.testing.assert_close(
+                            out_fa4_bf16,
+                            out_ref_bf16,
+                            atol=args.atol_bf16,
+                            rtol=args.rtol_bf16,
+                        )
+                except Exception as e:
+                    # Treat as fatal: BF16 kernel should be usable for basic sanity checking.
+                    raise RuntimeError("FA4-CuTe BF16 baseline failed") from e
+
+                # FA4 / CuTe FP8
+                q_fp8 = q_bf16.to(fp8_dtype)
+                k_fp8 = k_bf16.to(fp8_dtype)
+                v_fp8 = v_bf16.to(fp8_dtype)
+
+                # Placeholder descales (FA3-style: per-(batch, kv_head)).
+                q_descale = torch.ones(
+                    batch, nheads, device=device, dtype=torch.float32
+                )
+                k_descale = torch.ones(
+                    batch, nheads, device=device, dtype=torch.float32
+                )
+                v_descale = torch.ones(
+                    batch, nheads, device=device, dtype=torch.float32
+                )
+
+                # Optional: FP8 reference baseline (dequantized FP8 -> PyTorch) for quantization-error-only checks
+                out_ref_fp8 = None
+                if args.check and args.check_quantization_only:
+                    try:
+                        # Dequantize FP8 inputs back to BF16 (applying descales)
+                        q_ref_fp8 = (
+                            q_fp8.to(torch.bfloat16) * q_descale[:, None, :, None]
+                        ).to(torch.bfloat16)
+                        k_ref_fp8 = (
+                            k_fp8.to(torch.bfloat16) * k_descale[:, None, :, None]
+                        ).to(torch.bfloat16)
+                        v_ref_fp8 = (
+                            v_fp8.to(torch.bfloat16) * v_descale[:, None, :, None]
+                        ).to(torch.bfloat16)
+                        qkv_ref_fp8 = torch.stack(
+                            [q_ref_fp8, k_ref_fp8, v_ref_fp8], dim=2
+                        )
+                        out_ref_fp8 = attention_pytorch(qkv_ref_fp8, causal=causal)
+                    except RuntimeError as e:
+                        if "out of memory" in str(e).lower():
+                            out_ref_fp8 = None
+                        else:
+                            raise
+
+                fa4_kwargs = dict(softmax_scale=softmax_scale, causal=causal)
+                fa4_kwargs.update(
+                    _maybe_pass_descales(
+                        flash_attn_cute_fwd,
+                        q_descale=q_descale,
+                        k_descale=k_descale,
+                        v_descale=v_descale,
+                    )
+                )
+
+                try:
+                    # Warmup/compile (will raise until FP8 is implemented)
+                    out_fa4_fp8, _ = flash_attn_cute_fwd(
+                        q_fp8, k_fp8, v_fp8, **fa4_kwargs
+                    )
+                    t = time_fwd(
+                        flash_attn_cute_fwd,
+                        q_fp8,
+                        k_fp8,
+                        v_fp8,
+                        repeats=args.repeats,
+                        **fa4_kwargs,
+                    )
+                    times["FA4-CuTe-FP8"] = t
+                    if args.check:
+                        # Choose baseline: quantization-only (dequantized FP8) or full (BF16)
+                        if args.check_quantization_only:
+                            ref_baseline = out_ref_fp8
+                        else:
+                            ref_baseline = out_ref_bf16
+
+                        if ref_baseline is not None:
+                            torch.testing.assert_close(
+                                out_fa4_fp8,
+                                ref_baseline,
+                                atol=args.atol_fp8,
+                                rtol=args.rtol_fp8,
+                            )
+                except Exception as e:
+                    fp8_failures.append((causal, headdim, batch, seqlen, repr(e)))
+                    times["FA4-CuTe-FP8"] = float("nan")
+
+                if args.run_cudnn and cudnn is not None:
+                    qkv_fp8 = qkv_bf16.to(fp8_dtype)
+                    runner = cudnn_sdpa_fp8_setup(
+                        qkv_fp8, seqlen, seqlen, causal=causal
+                    )
+                    _ = runner()  # warmup
+                    t = time_fwd(lambda: runner(), repeats=args.repeats)
+                    times["cuDNN-FP8"] = t
+
+                print(
+                    f"### causal={causal}, headdim={headdim}, batch={batch}, seqlen={seqlen} ###"
+                )
+                for method in methods:
+                    t = times.get(method, float("nan"))
+                    speeds[method] = efficiency(
+                        flops(batch, seqlen, headdim, nheads, causal), t
+                    )
+                    if math.isnan(t):
+                        print(f"{method} fwd: (skipped)")
+                    else:
+                        print(
+                            f"{method} fwd: {speeds[method]:.2f} TFLOPs/s, {t * 1e3:.3f} ms"
+                        )
+                if math.isnan(times.get("FA4-CuTe-FP8", float("nan"))):
+                    print("FA4-CuTe-FP8 status: FAILED")
+
+    if fp8_failures:
+        print(f"\nFP8 failures: {len(fp8_failures)} (showing first 5)")
+        for causal, headdim, batch, seqlen, err in fp8_failures[:5]:
+            print(
+                f"- causal={causal} headdim={headdim} batch={batch} seqlen={seqlen}: {err}"
+            )
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

python/sglang/jit_kernel/flash_attn/cute/block_info.py

@@ -0,0 +1,203 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, const_expr
+
+from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK, SeqlenInfoQKNewK
+
+
+@dataclass(frozen=True)
+class BlockInfo:
+    tile_m: cutlass.Constexpr[int]
+    tile_n: cutlass.Constexpr[int]
+    is_causal: cutlass.Constexpr[bool]
+    is_local: cutlass.Constexpr[bool] = False
+    is_split_kv: cutlass.Constexpr[bool] = False
+    window_size_left: Optional[Int32] = None
+    window_size_right: Optional[Int32] = None
+    qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1
+
+    @cute.jit
+    def get_n_idx_left_right(
+        self,
+        seqlen_info: SeqlenInfoQK,
+        m_idx: Int32,
+    ) -> Tuple[Int32, Int32]:
+        m_idx_actual = m_idx // self.qhead_per_kvhead_packgqa
+        if const_expr(
+            self.is_causal or (self.is_local and self.window_size_right is not None)
+        ):
+            n_idx_right = m_idx_actual + 1 + seqlen_info.seqlen_k - seqlen_info.seqlen_q
+            if const_expr(self.window_size_right is not None):
+                n_idx_right += self.window_size_right
+        else:
+            n_idx_right = seqlen_info.seqlen_k
+        if const_expr(self.is_local and self.window_size_left is not None):
+            n_idx_left = (
+                m_idx_actual
+                + seqlen_info.seqlen_k
+                - seqlen_info.seqlen_q
+                - self.window_size_left
+            )
+            n_idx_left = cutlass.max(n_idx_left, 0)
+        else:
+            n_idx_left = 0
+        # inclusive n_idx_left, exclusive n_idx_right
+        # e.g. for causal, return (0, m_idx + 1)
+        return n_idx_left, n_idx_right
+
+    @cute.jit
+    def get_n_block_min_max(
+        self,
+        seqlen_info: SeqlenInfoQK,
+        m_block: Int32,
+        split_idx: Int32 = 0,
+        num_splits: Int32 = 1,
+        half_tile_m: bool = False,
+        absolute: bool = False,
+        half_tile_n: bool = False,
+    ) -> Tuple[Int32, Int32]:
+        tile_m = self.tile_m // 2 if const_expr(half_tile_m) else self.tile_m
+        tile_n = self.tile_n // 2 if const_expr(half_tile_n) else self.tile_n
+        n_block_max = cute.ceil_div(seqlen_info.seqlen_k, tile_n)
+        if const_expr(
+            self.is_causal or (self.is_local and self.window_size_right is not None)
+        ):
+            m_idx_max = (m_block + 1) * tile_m
+            if const_expr(self.qhead_per_kvhead_packgqa > 1):
+                m_idx_max = cute.ceil_div(m_idx_max, self.qhead_per_kvhead_packgqa)
+            n_idx = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q
+            n_idx_right = (
+                n_idx if const_expr(self.is_causal) else n_idx + self.window_size_right
+            )
+            n_block_max = min(n_block_max, cute.ceil_div(n_idx_right, tile_n))
+        n_block_min = 0
+        if const_expr(self.is_local and self.window_size_left is not None):
+            m_idx_min = m_block * tile_m
+            if const_expr(self.qhead_per_kvhead_packgqa > 1):
+                m_idx_min = m_idx_min // self.qhead_per_kvhead_packgqa
+            n_idx = m_idx_min + seqlen_info.seqlen_k - seqlen_info.seqlen_q
+            n_idx_left = n_idx - self.window_size_left
+            n_block_min = cutlass.max(n_idx_left // tile_n, 0)
+        if cutlass.const_expr(self.is_split_kv and not absolute):
+            num_n_blocks_per_split = (
+                Int32(0)
+                if n_block_max <= n_block_min
+                else (n_block_max - n_block_min + num_splits - 1) // num_splits
+            )
+            n_block_min = n_block_min + split_idx * num_n_blocks_per_split
+            n_block_max = cutlass.min(n_block_min + num_n_blocks_per_split, n_block_max)
+        return n_block_min, n_block_max
+
+    @cute.jit
+    def get_m_block_min_max(
+        self, seqlen_info: SeqlenInfoQK, n_block: Int32
+    ) -> Tuple[Int32, Int32]:
+        m_block_max = cute.ceil_div(seqlen_info.seqlen_q, self.tile_m)
+        m_block_min = 0
+        if const_expr(
+            self.is_causal or (self.is_local and self.window_size_right is not None)
+        ):
+            n_idx_min = n_block * self.tile_n
+            m_idx = n_idx_min + seqlen_info.seqlen_q - seqlen_info.seqlen_k
+            m_idx_right = (
+                m_idx if const_expr(self.is_causal) else m_idx - self.window_size_right
+            )
+            m_block_min = max(m_block_min, m_idx_right // self.tile_m)
+        if const_expr(self.is_local and self.window_size_left is not None):
+            n_idx_max = (n_block + 1) * self.tile_n
+            m_idx = n_idx_max + seqlen_info.seqlen_q - seqlen_info.seqlen_k
+            m_idx_left = m_idx + self.window_size_left
+            m_block_max = min(m_block_max, cute.ceil_div(m_idx_left, self.tile_m))
+        return m_block_min, m_block_max
+
+    @cute.jit
+    def get_n_block_k_new_min_max(
+        self,
+        seqlen_info: SeqlenInfoQKNewK,
+        m_block: Int32,
+        split_idx: Int32 = 0,
+        num_splits: Int32 = 1,
+    ) -> Tuple[Int32, Int32]:
+        """Get the block range for new K tokens (append KV).
+
+        First computes the full n_block range via get_n_block_min_max, then maps
+        those blocks into the new-K index space by subtracting seqlen_k_og.
+        """
+        n_block_min, n_block_max = self.get_n_block_min_max(
+            seqlen_info,
+            m_block,
+            split_idx,
+            num_splits,
+        )
+        idx_k_new_min = cutlass.max(
+            n_block_min * self.tile_n - seqlen_info.seqlen_k_og, 0
+        )
+        idx_k_new_max = cutlass.min(
+            n_block_max * self.tile_n - seqlen_info.seqlen_k_og,
+            seqlen_info.seqlen_k_new,
+        )
+        n_block_new_min = idx_k_new_min // self.tile_n
+        n_block_new_max = (
+            cute.ceil_div(idx_k_new_max, self.tile_n)
+            if idx_k_new_max > idx_k_new_min
+            else n_block_new_min
+        )
+        return n_block_new_min, n_block_new_max
+
+    @cute.jit
+    def get_n_block_min_causal_local_mask(
+        self,
+        seqlen_info: SeqlenInfoQK,
+        m_block: Int32,
+        n_block_min: Int32,
+    ) -> Int32:
+        """If we have separate iterations with causal or local masking at the start, where do we stop"""
+        m_idx_min = m_block * self.tile_m
+        if const_expr(self.qhead_per_kvhead_packgqa > 1):
+            m_idx_min = m_idx_min // self.qhead_per_kvhead_packgqa
+        n_idx = m_idx_min + seqlen_info.seqlen_k - seqlen_info.seqlen_q
+        n_idx_right = (
+            n_idx
+            if const_expr(not self.is_local or self.window_size_right is None)
+            else n_idx + self.window_size_right
+        )
+        return cutlass.max(n_block_min, n_idx_right // self.tile_n)
+
+    @cute.jit
+    def get_n_block_min_before_local_mask(
+        self,
+        seqlen_info: SeqlenInfoQK,
+        m_block: Int32,
+        n_block_min: Int32,
+    ) -> Int32:
+        """If we have separate iterations with local masking at the end, where do we stop the non-masked iterations"""
+        if const_expr(not self.is_local or self.window_size_left is None):
+            return n_block_min
+        else:
+            m_idx_max = (m_block + 1) * self.tile_m
+            if const_expr(self.qhead_per_kvhead_packgqa > 1):
+                m_idx_max = cute.ceil_div(m_idx_max, self.qhead_per_kvhead_packgqa)
+            n_idx = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q
+            n_idx_left = n_idx - self.window_size_left
+            return cutlass.max(n_block_min, cute.ceil_div(n_idx_left, self.tile_n))
+
+    @cute.jit
+    def get_n_block_max_for_m_block(
+        self,
+        seqlen_info: SeqlenInfoQK,
+        m_block: Int32,
+    ) -> Int32:
+        n_block_max = cute.ceil_div(seqlen_info.seqlen_k, self.tile_n)
+        if const_expr(self.is_causal or self.window_size_right is not None):
+            m_idx_max = (m_block + 1) * self.tile_m
+            if const_expr(self.qhead_per_kvhead_packgqa > 1):
+                m_idx_max = cute.ceil_div(m_idx_max, self.qhead_per_kvhead_packgqa)
+            n_idx_right = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q
+            if const_expr(self.window_size_right is not None):
+                n_idx_right += self.window_size_right
+            n_block_max = min(n_block_max, cute.ceil_div(n_idx_right, self.tile_n))
+        return n_block_max

python/sglang/jit_kernel/flash_attn/cute/block_sparsity.py

@@ -0,0 +1,722 @@
+"""
+Block-sparsity utilities for FlexAttention
+"""
+
+from typing import Callable, NamedTuple, Tuple
+
+import cutlass.cute as cute
+import torch
+
+from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import (
+    get_broadcast_dims,
+    to_cute_tensor,
+)
+
+
+def ceildiv(a: int, b: int) -> int:
+    return (a + b - 1) // b
+
+
+class BlockSparseTensors(NamedTuple):
+    mask_block_cnt: cute.Tensor
+    mask_block_idx: cute.Tensor
+    full_block_cnt: cute.Tensor | None = None
+    full_block_idx: cute.Tensor | None = None
+    cu_total_m_blocks: cute.Tensor | None = None
+    cu_block_idx_offsets: cute.Tensor | None = None
+    dq_write_order: cute.Tensor | None = None
+    dq_write_order_full: cute.Tensor | None = None
+
+    def __new_from_mlir_values__(self, values):
+        new_fields = []
+        idx = 0
+        for original in self:
+            if original is None:
+                new_fields.append(None)
+            else:
+                new_fields.append(values[idx])
+                idx += 1
+        return BlockSparseTensors(*new_fields)
+
+
+class BlockSparseTensorsTorch(NamedTuple):
+    mask_block_cnt: torch.Tensor
+    mask_block_idx: torch.Tensor
+    full_block_cnt: torch.Tensor | None = None
+    full_block_idx: torch.Tensor | None = None
+    cu_total_m_blocks: torch.Tensor | None = None
+    cu_block_idx_offsets: torch.Tensor | None = None
+    block_size: tuple[int, int] | None = None
+    dq_write_order: torch.Tensor | None = None
+    dq_write_order_full: torch.Tensor | None = None
+    spt: bool | None = None
+
+
+def _ordered_to_dense_simple(
+    num_blocks: torch.Tensor,
+    indices: torch.Tensor,
+    num_cols: int,
+) -> torch.Tensor:
+    """Convert ordered sparse representation to dense binary matrix.
+
+    Args:
+        num_blocks: [B, H, num_rows] count of valid entries per row
+        indices: [B, H, num_rows, max_entries] column indices (valid entries packed left)
+        num_cols: total number of columns
+
+    Returns:
+        dense: [B, H, num_rows, num_cols] binary int32 matrix
+    """
+    B, H, num_rows, max_entries = indices.shape
+    device = indices.device
+    dense = torch.zeros(B, H, num_rows, num_cols + 1, dtype=torch.int32, device=device)
+    col_range = torch.arange(max_entries, device=device)
+    valid = col_range[None, None, None, :] < num_blocks[:, :, :, None]
+    safe_indices = torch.where(valid, indices.long(), num_cols)
+    row_idx = torch.arange(num_rows, device=device)[None, None, :, None].expand_as(
+        indices
+    )
+    b_idx = torch.arange(B, device=device)[:, None, None, None].expand_as(indices)
+    h_idx = torch.arange(H, device=device)[None, :, None, None].expand_as(indices)
+    dense[b_idx, h_idx, row_idx, safe_indices] = 1
+    return dense[:, :, :, :num_cols]
+
+
+def compute_dq_write_order(
+    fwd_mask_cnt: torch.Tensor,
+    fwd_mask_idx: torch.Tensor,
+    fwd_full_cnt: torch.Tensor | None,
+    fwd_full_idx: torch.Tensor | None,
+    bwd_mask_cnt: torch.Tensor,
+    bwd_mask_idx: torch.Tensor,
+    bwd_full_cnt: torch.Tensor | None,
+    bwd_full_idx: torch.Tensor | None,
+    spt: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Compute dQ write-order metadata for deterministic block-sparse backward.
+
+    For each (n_block, i) in the backward iteration, computes the semaphore
+    lock value: the rank of n_block in the combined (partial + full) sorted
+    contributor list for the target m_block.
+
+    Lock values are assigned in ascending n_block order (or descending if spt=True)
+    to guarantee deadlock-freedom with the CTA scheduling order.
+
+    Args:
+        fwd_mask_cnt: [B, H, num_m_blocks] partial contributor counts per m_block
+        fwd_mask_idx: [B, H, num_m_blocks, max_kv] partial contributor n_block indices (ascending)
+        fwd_full_cnt: [B, H, num_m_blocks] full contributor counts per m_block (optional)
+        fwd_full_idx: [B, H, num_m_blocks, max_kv] full contributor n_block indices (optional)
+        bwd_mask_cnt: [B, H, num_n_blocks] partial iteration counts per n_block
+        bwd_mask_idx: [B, H, num_n_blocks, max_q] partial iteration m_block indices
+        bwd_full_cnt: [B, H, num_n_blocks] full iteration counts per n_block (optional)
+        bwd_full_idx: [B, H, num_n_blocks, max_q] full iteration m_block indices (optional)
+        spt: if True, reverse ordering (highest n_block gets lock_value=0)
+
+    Returns:
+        (dq_write_order, dq_write_order_full): tensors parallel to bwd_mask_idx
+        and bwd_full_idx respectively, containing lock values.
+    """
+    device = fwd_mask_idx.device
+    B, H, num_m, max_kv_partial = fwd_mask_idx.shape
+    _, _, num_n, max_q_partial = bwd_mask_idx.shape
+
+    has_full = fwd_full_cnt is not None and fwd_full_idx is not None
+
+    dense_partial = _ordered_to_dense_simple(fwd_mask_cnt, fwd_mask_idx, num_n)
+    if has_full:
+        dense_full = _ordered_to_dense_simple(fwd_full_cnt, fwd_full_idx, num_n)
+        dense = (dense_partial + dense_full).clamp(max=1)
+    else:
+        dense = dense_partial
+
+    cumsum = dense.cumsum(dim=-1)
+    rank_table = (cumsum - dense).to(torch.int32)
+
+    if spt:
+        total_per_m = cumsum[:, :, :, -1:]
+        rank_table = (total_per_m - 1 - rank_table).to(torch.int32)
+
+    def _gather_write_order(bwd_idx, bwd_cnt):
+        b_i = torch.arange(B, device=device)[:, None, None, None].expand_as(bwd_idx)
+        h_i = torch.arange(H, device=device)[None, :, None, None].expand_as(bwd_idx)
+        n_i = torch.arange(bwd_idx.shape[2], device=device)[
+            None, None, :, None
+        ].expand_as(bwd_idx)
+        m_vals = bwd_idx.long().clamp(0, num_m - 1)
+        return rank_table[b_i, h_i, m_vals, n_i].to(torch.int32)
+
+    dq_write_order = _gather_write_order(bwd_mask_idx, bwd_mask_cnt)
+
+    dq_write_order_full = None
+    if has_full and bwd_full_cnt is not None and bwd_full_idx is not None:
+        dq_write_order_full = _gather_write_order(bwd_full_idx, bwd_full_cnt)
+
+    return dq_write_order, dq_write_order_full
+
+
+def compute_dq_write_order_from_block_mask(
+    block_mask,
+    spt: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    (
+        _seq_q,
+        _seq_k,
+        kv_mask_cnt,
+        kv_mask_idx,
+        full_kv_cnt,
+        full_kv_idx,
+        q_mask_cnt,
+        q_mask_idx,
+        full_q_cnt,
+        full_q_idx,
+        *_,
+    ) = block_mask.as_tuple()
+    return compute_dq_write_order(
+        kv_mask_cnt,
+        kv_mask_idx,
+        full_kv_cnt,
+        full_kv_idx,
+        q_mask_cnt,
+        q_mask_idx,
+        full_q_cnt,
+        full_q_idx,
+        spt=spt,
+    )
+
+
+def get_sparse_q_block_size(
+    tensors: BlockSparseTensorsTorch | None,
+    seqlen_q: int,
+) -> int | None:
+    """Return the Q sparse block size, or None when sparsity is unset or ambiguous."""
+    if tensors is None:
+        return None
+    if tensors.block_size is not None:
+        return tensors.block_size[0]
+    num_m_blocks = tensors.mask_block_idx.shape[2]
+    min_block_size = ceildiv(seqlen_q, num_m_blocks)
+    max_block_size = (
+        seqlen_q if num_m_blocks == 1 else (seqlen_q - 1) // (num_m_blocks - 1)
+    )
+    if min_block_size != max_block_size:
+        return None
+    return min_block_size
+
+
+def _expand_sparsity_tensor(
+    tensor: torch.Tensor,
+    expected_shape: Tuple[int, ...],
+    tensor_name: str,
+    context: str | None,
+    hint: str | Callable[[], str] | None,
+) -> torch.Tensor:
+    """Check if we need to expand the tensor to expected shape, and do so if possible."""
+    needs_expand = tensor.shape != expected_shape
+    if not needs_expand:
+        return tensor
+    can_expand = all(
+        map(lambda cur, tgt: cur == tgt or cur == 1, tensor.shape, expected_shape)
+    )
+    if not can_expand:
+        context_clause = f" ({context})" if context else ""
+        resolved_hint = hint() if callable(hint) else hint
+        hint_clause = f" Hint: {resolved_hint}" if resolved_hint else ""
+        raise ValueError(
+            f"{tensor_name}{context_clause} with shape {tensor.shape} cannot be expanded to expected shape {expected_shape}."
+            f"{hint_clause}"
+        )
+    return tensor.expand(*expected_shape)
+
+
+def _check_and_expand_block(
+    name: str,
+    cnt: torch.Tensor | None,
+    idx: torch.Tensor | None,
+    expected_count_shape: Tuple[int, ...],
+    expected_index_shape: Tuple[int, ...],
+    context: str | None,
+    hint: str | Callable[[], str] | None,
+) -> Tuple[torch.Tensor | None, torch.Tensor | None]:
+    if (cnt is None) != (idx is None):
+        raise ValueError(
+            f"{name}_block_cnt and {name}_block_idx must both be provided or both be None"
+        )
+    if cnt is None or idx is None:
+        return None, None
+    if cnt.dtype != torch.int32 or idx.dtype != torch.int32:
+        raise ValueError(f"{name}_block tensors must have dtype torch.int32")
+    if cnt.device != idx.device:
+        raise ValueError(
+            f"{name}_block_cnt and {name}_block_idx must be on the same device"
+        )
+    if not cnt.is_cuda or not idx.is_cuda:
+        raise ValueError(f"{name}_block tensors must live on CUDA")
+    expanded_cnt = _expand_sparsity_tensor(
+        cnt, expected_count_shape, f"{name}_block_cnt", context, hint
+    )
+    # [Note] Allow Compact block sparse indices
+    # Allow the last dimension (n_blocks) of idx to be <= expected, since
+    # FA4 only accesses indices 0..cnt-1 per query tile. This enables compact
+    # index tensors that avoid O(N^2) memory at long sequence lengths.
+    if idx.ndim == 4 and idx.shape[3] <= expected_index_shape[3]:
+        expected_index_shape = (*expected_index_shape[:3], idx.shape[3])
+    expanded_idx = _expand_sparsity_tensor(
+        idx, expected_index_shape, f"{name}_block_idx", context, hint
+    )
+    return expanded_cnt, expanded_idx
+
+
+def _check_and_expand_metadata_tensor(
+    name: str,
+    tensor: torch.Tensor | None,
+    expected_shape: Tuple[int, ...],
+    context: str | None,
+    hint: str | Callable[[], str] | None,
+    device: torch.device,
+) -> torch.Tensor | None:
+    if tensor is None:
+        return None
+    if tensor.dtype != torch.int32:
+        raise ValueError(f"{name} must have dtype torch.int32")
+    if tensor.device != device:
+        raise ValueError(f"{name} must be on the same device as block sparse tensors")
+    if not tensor.is_cuda:
+        raise ValueError(f"{name} must live on CUDA")
+    return _expand_sparsity_tensor(tensor, expected_shape, name, context, hint)
+
+
+def get_block_sparse_expected_shapes(
+    batch_size: int,
+    num_head: int,
+    seqlen_q: int,
+    seqlen_k: int,
+    m_block_size: int,
+    n_block_size: int,
+    q_stage: int,
+) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]:
+    """Return (expected_count_shape, expected_index_shape) for block sparse normalization."""
+    m_block_size_effective = q_stage * m_block_size
+    expected_m_blocks = ceildiv(seqlen_q, m_block_size_effective)
+    expected_n_blocks = ceildiv(seqlen_k, n_block_size)
+    expected_count_shape = (batch_size, num_head, expected_m_blocks)
+    expected_index_shape = (batch_size, num_head, expected_m_blocks, expected_n_blocks)
+    return expected_count_shape, expected_index_shape
+
+
+def infer_block_sparse_expected_shapes(
+    tensors: BlockSparseTensorsTorch,
+    *,
+    batch_size: int,
+    num_head: int,
+    seqlen_q: int,
+    seqlen_k: int,
+    m_block_size: int,
+    n_block_size: int,
+    q_stage: int,
+    context: str,
+    sparse_block_size_q: int | None = None,
+    sparse_block_size_kv: int | None = None,
+) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int], int]:
+    """Infer shapes and scaling for block-sparse tensors.
+
+    Expectations:
+    - mask_block_cnt is (B, H, M) and mask_block_idx is (B, H, M, N).
+    - Batch/head dims may be 1 for broadcast, or match the requested sizes.
+    - sparse_block_size_kv must match tile_n.
+    - sparse_block_size_q must be a multiple of q_stage * tile_m.
+    - If sparse_block_size_q is omitted and seqlen_q/num_m_blocks is ambiguous,
+      the caller must provide block_size to disambiguate. TODO will make this required in a future PR.
+    """
+    base_m_block = q_stage * m_block_size
+    base_n_block = n_block_size
+    if sparse_block_size_kv is None:
+        sparse_block_size_kv = base_n_block
+    if sparse_block_size_kv != base_n_block:
+        raise ValueError(
+            f"Block sparse tensors{context} require BLOCK_SIZE_KV={base_n_block}."
+        )
+    if tensors.mask_block_idx is None:
+        raise ValueError(
+            "mask_block_cnt and mask_block_idx must be provided for block sparsity."
+        )
+    num_m_blocks = tensors.mask_block_idx.shape[2]
+
+    if sparse_block_size_q is None:
+        sparse_block_size_q = get_sparse_q_block_size(tensors, seqlen_q)
+        if sparse_block_size_q is None and base_m_block != 1:
+            raise ValueError(
+                f"Block sparse tensors{context} require explicit sparse_block_size[0] "
+                f"to disambiguate block size for seqlen_q={seqlen_q} and num_m_blocks={num_m_blocks}."
+            )
+        if sparse_block_size_q is None:
+            sparse_block_size_q = ceildiv(seqlen_q, num_m_blocks)
+
+    if sparse_block_size_q % base_m_block != 0:
+        raise ValueError(
+            f"Block sparse tensors{context} have block size {sparse_block_size_q}, "
+            f"which must be a multiple of {base_m_block}."
+        )
+
+    expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q)
+    expected_n_blocks = ceildiv(seqlen_k, sparse_block_size_kv)
+    q_subtile_factor = sparse_block_size_q // base_m_block
+    expected_count_shape = (batch_size, num_head, expected_m_blocks)
+    expected_index_shape = (batch_size, num_head, expected_m_blocks, expected_n_blocks)
+
+    mask_block_cnt = tensors.mask_block_cnt
+    mask_block_idx = tensors.mask_block_idx
+    if mask_block_cnt is None or mask_block_idx is None:
+        raise ValueError(
+            "mask_block_cnt and mask_block_idx must be provided for block sparsity."
+        )
+    if mask_block_cnt.ndim != 3 or mask_block_idx.ndim != 4:
+        raise ValueError(
+            f"Block sparse tensors{context} must have shapes (B, H, M) and (B, H, M, N)."
+        )
+    for dim_name, cur, tgt in (
+        ("batch", mask_block_cnt.shape[0], expected_count_shape[0]),
+        ("head", mask_block_cnt.shape[1], expected_count_shape[1]),
+    ):
+        if cur != tgt and cur != 1:
+            raise ValueError(
+                f"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1."
+            )
+    for dim_name, cur, tgt in (
+        ("batch", mask_block_idx.shape[0], expected_index_shape[0]),
+        ("head", mask_block_idx.shape[1], expected_index_shape[1]),
+    ):
+        if cur != tgt and cur != 1:
+            raise ValueError(
+                f"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1."
+            )
+    if mask_block_cnt.shape[2] != mask_block_idx.shape[2]:
+        raise ValueError(
+            f"Block sparse tensors{context} must share the same m-block dimension."
+        )
+    # [Note] Allow Compact block sparse indices: FA4 only accesses indices 0..cnt-1
+    # per query tile, so idx.shape[3] can be <= expected_n_blocks.
+    if mask_block_idx.shape[3] > expected_n_blocks:
+        raise ValueError(
+            f"Block sparse tensors{context} n-block dimension must be <= {expected_n_blocks}."
+        )
+    if expected_m_blocks != num_m_blocks:
+        raise ValueError(
+            f"Block sparse tensors{context} m-block dimension {num_m_blocks} does not match "
+            f"sparse_block_size_q={sparse_block_size_q}. "
+            f"Set BlockSparseTensorsTorch.block_size to match the BlockMask BLOCK_SIZE."
+        )
+    return expected_count_shape, expected_index_shape, q_subtile_factor
+
+
+def get_block_sparse_expected_shapes_bwd(
+    batch_size: int,
+    num_head: int,
+    seqlen_q: int,
+    seqlen_k: int,
+    m_block_size: int,
+    n_block_size: int,
+    subtile_factor: int,
+) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]:
+    """Return (expected_count_shape, expected_index_shape) for backward block sparse normalization.
+
+    Backward uses Q-direction indexing (transposed from forward), where shapes are
+    indexed by N-blocks first, then M-blocks. The sparse_block_size_q is determined
+    by subtile_factor * m_block_size.
+    """
+    sparse_block_size_q = subtile_factor * m_block_size
+    expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q)
+    expected_n_blocks = ceildiv(seqlen_k, n_block_size)
+    expected_count_shape = (batch_size, num_head, expected_n_blocks)
+    expected_index_shape = (batch_size, num_head, expected_n_blocks, expected_m_blocks)
+    return expected_count_shape, expected_index_shape
+
+
+def normalize_block_sparse_tensors(
+    tensors: BlockSparseTensorsTorch,
+    *,
+    expected_count_shape: Tuple[int, ...],
+    expected_index_shape: Tuple[int, ...],
+    context: str | None = None,
+    hint: str | Callable[[], str] | None = None,
+) -> BlockSparseTensorsTorch:
+    if tensors.mask_block_cnt is None or tensors.mask_block_idx is None:
+        raise ValueError(
+            "mask_block_cnt and mask_block_idx must be provided for block sparsity."
+        )
+
+    mask_cnt, mask_idx = _check_and_expand_block(
+        "mask",
+        tensors.mask_block_cnt,
+        tensors.mask_block_idx,
+        expected_count_shape,
+        expected_index_shape,
+        context,
+        hint,
+    )
+    if mask_cnt is None or mask_idx is None:
+        raise ValueError(
+            "mask_block_cnt and mask_block_idx must be provided for block sparsity."
+        )
+
+    full_cnt, full_idx = _check_and_expand_block(
+        "full",
+        tensors.full_block_cnt,
+        tensors.full_block_idx,
+        expected_count_shape,
+        expected_index_shape,
+        context,
+        hint,
+    )
+    if full_cnt is not None and mask_cnt.device != full_cnt.device:
+        raise ValueError("All block sparse tensors must be on the same device")
+
+    dq_write_order = _check_and_expand_metadata_tensor(
+        "dq_write_order",
+        tensors.dq_write_order,
+        tuple(mask_idx.shape),
+        context,
+        hint,
+        mask_cnt.device,
+    )
+    dq_write_order_full = _check_and_expand_metadata_tensor(
+        "dq_write_order_full",
+        tensors.dq_write_order_full,
+        tuple(full_idx.shape) if full_idx is not None else expected_index_shape,
+        context,
+        hint,
+        mask_cnt.device,
+    )
+    spt = tensors.spt
+    if spt is not None and not isinstance(spt, bool):
+        raise ValueError("spt must be a bool when provided")
+    if spt is not None and dq_write_order is None:
+        raise ValueError("spt requires dq_write_order to be provided")
+
+    return BlockSparseTensorsTorch(
+        mask_block_cnt=mask_cnt,
+        mask_block_idx=mask_idx,
+        full_block_cnt=full_cnt,
+        full_block_idx=full_idx,
+        cu_total_m_blocks=tensors.cu_total_m_blocks,
+        cu_block_idx_offsets=tensors.cu_block_idx_offsets,
+        block_size=tensors.block_size,
+        dq_write_order=dq_write_order,
+        dq_write_order_full=dq_write_order_full,
+        spt=spt,
+    )
+
+
+def is_block_sparsity_enabled(tensors: BlockSparseTensorsTorch) -> bool:
+    return any(t is not None for t in (tensors.full_block_cnt, tensors.mask_block_cnt))
+
+
+def get_block_sparse_broadcast_pattern(
+    tensors: BlockSparseTensorsTorch,
+) -> Tuple[Tuple[bool, ...], ...] | None:
+    """Return broadcast pattern for block sparse tensors by checking actual strides.
+
+    Returns a tuple of broadcast patterns (one per tensor) where each pattern
+    is a tuple of bools indicating which dims have stride=0.
+    This is used in compile keys to ensure kernels are recompiled when
+    broadcast patterns change, since CuTe's mark_layout_dynamic() keeps
+    stride=0 as static.
+
+    The tensors should already be expanded/normalized before calling this function.
+
+    Returns None if block sparsity is not enabled.
+    """
+    if not is_block_sparsity_enabled(tensors):
+        return None
+
+    patterns = []
+    for tensor in (
+        tensors.mask_block_cnt,
+        tensors.mask_block_idx,
+        tensors.full_block_cnt,
+        tensors.full_block_idx,
+        tensors.dq_write_order,
+        tensors.dq_write_order_full,
+    ):
+        if tensor is not None:
+            patterns.append(get_broadcast_dims(tensor))
+        else:
+            patterns.append(None)
+    return tuple(patterns)
+
+
+def normalize_block_sparse_config(
+    tensors: BlockSparseTensorsTorch,
+    *,
+    batch_size: int,
+    num_head: int,
+    seqlen_q: int,
+    seqlen_k: int,
+    block_size: tuple[int, int],
+    q_stage: int,
+) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None, int]:
+    """Validate the block-sparse config, infer expected shapes, and normalize.
+
+    Handles both fixed-length (3D `[B, H, M]` / 4D `[B, H, M, N]`) and varlen
+    (2D `[H, total_m_blocks]` / `[H, total_n_blocks]`) layouts. Varlen is
+    detected by `tensors.cu_total_m_blocks is not None` and forces
+    `q_subtile_factor == 1` (TODO: potentially remove this restriction).
+    """
+    m_block_size, n_block_size = block_size
+    if tensors.block_size is None:
+        sparse_block_size_q, sparse_block_size_kv = None, n_block_size
+    else:
+        sparse_block_size_q, sparse_block_size_kv = tensors.block_size
+    if sparse_block_size_kv != n_block_size:
+        raise ValueError(
+            f"Block sparsity requires sparse_block_size[1]={n_block_size} to match tile_n."
+        )
+    if tensors.cu_total_m_blocks is not None:
+        base_m_block = q_stage * m_block_size
+        if sparse_block_size_q is not None and sparse_block_size_q != base_m_block:
+            raise ValueError(
+                f"Varlen block sparsity requires sparse_block_size[0]={base_m_block} "
+                f"(= q_stage * tile_m); got {sparse_block_size_q}."
+            )
+        total_m_blocks = tensors.mask_block_cnt.shape[-1]
+        total_n_blocks = tensors.mask_block_idx.shape[-1]
+        expected_count_shape = (num_head, total_m_blocks)
+        expected_index_shape = (num_head, total_n_blocks)
+        q_subtile_factor = 1
+    else:
+        expected_count_shape, expected_index_shape, q_subtile_factor = (
+            infer_block_sparse_expected_shapes(
+                tensors,
+                batch_size=batch_size,
+                num_head=num_head,
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                m_block_size=m_block_size,
+                n_block_size=n_block_size,
+                q_stage=q_stage,
+                context="forward",
+                sparse_block_size_q=sparse_block_size_q,
+                sparse_block_size_kv=sparse_block_size_kv,
+            )
+        )
+    normalized_tensors = normalize_block_sparse_tensors(
+        tensors,
+        expected_count_shape=expected_count_shape,
+        expected_index_shape=expected_index_shape,
+    )
+    return (
+        normalized_tensors,
+        get_block_sparse_broadcast_pattern(normalized_tensors),
+        q_subtile_factor,
+    )
+
+
+def normalize_block_sparse_config_bwd(
+    tensors: BlockSparseTensorsTorch,
+    *,
+    batch_size: int,
+    num_head: int,
+    seqlen_q: int,
+    seqlen_k: int,
+    block_size: tuple[int, int],
+    subtile_factor: int,
+) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None]:
+    m_block_size, n_block_size = block_size
+    if tensors.block_size is None:
+        sparse_block_size_q, sparse_block_size_kv = (
+            subtile_factor * m_block_size,
+            n_block_size,
+        )
+    else:
+        sparse_block_size_q, sparse_block_size_kv = tensors.block_size
+    if sparse_block_size_q != subtile_factor * m_block_size:
+        raise ValueError(
+            f"Block sparsity expects sparse_block_size_q={subtile_factor * m_block_size} "
+            f"for subtile_factor={subtile_factor}."
+        )
+    if sparse_block_size_kv != n_block_size:
+        raise ValueError(
+            f"Block sparsity expects sparse_block_size[1]={n_block_size} to match tile_n."
+        )
+    expected_count_shape, expected_index_shape = get_block_sparse_expected_shapes_bwd(
+        batch_size,
+        num_head,
+        seqlen_q,
+        seqlen_k,
+        m_block_size,
+        n_block_size,
+        subtile_factor,
+    )
+    normalized_tensors = normalize_block_sparse_tensors(
+        tensors,
+        expected_count_shape=expected_count_shape,
+        expected_index_shape=expected_index_shape,
+        context="_flash_attn_bwd",
+        hint=lambda: (
+            f"Backward expects Q-direction block-sparse tensors (q_mask_cnt/q_mask_idx, "
+            f"and optionally full_q_cnt/full_q_idx). Regenerate the backward BlockMask with "
+            f"BLOCK_SIZE=({subtile_factor * m_block_size}, {n_block_size})."
+        ),
+    )
+    return normalized_tensors, get_block_sparse_broadcast_pattern(normalized_tensors)
+
+
+def to_cute_block_sparse_tensors(
+    tensors: BlockSparseTensorsTorch, enable_tvm_ffi: bool = True
+) -> BlockSparseTensors | None:
+    """Convert torch block sparsity tensors to CuTe tensors, optionally for tvm ffi"""
+    if not is_block_sparsity_enabled(tensors):
+        return None
+    mask_block_cnt_tensor, mask_block_idx_tensor = [
+        to_cute_tensor(
+            t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi
+        )
+        for t in (tensors.mask_block_cnt, tensors.mask_block_idx)
+    ]
+    full_block_cnt_tensor, full_block_idx_tensor = [
+        (
+            to_cute_tensor(
+                t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi
+            )
+            if t is not None
+            else None
+        )
+        for t in (tensors.full_block_cnt, tensors.full_block_idx)
+    ]
+    cu_total_m_blocks_tensor, cu_block_idx_offsets_tensor = [
+        (
+            to_cute_tensor(
+                t, assumed_align=4, leading_dim=0, enable_tvm_ffi=enable_tvm_ffi
+            )
+            if t is not None
+            else None
+        )
+        for t in (tensors.cu_total_m_blocks, tensors.cu_block_idx_offsets)
+    ]
+    dq_write_order_tensor, dq_write_order_full_tensor = [
+        (
+            to_cute_tensor(
+                t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi
+            )
+            if t is not None
+            else None
+        )
+        for t in (tensors.dq_write_order, tensors.dq_write_order_full)
+    ]
+
+    return BlockSparseTensors(
+        mask_block_cnt_tensor,
+        mask_block_idx_tensor,
+        full_block_cnt_tensor,
+        full_block_idx_tensor,
+        cu_total_m_blocks_tensor,
+        cu_block_idx_offsets_tensor,
+        dq_write_order_tensor,
+        dq_write_order_full_tensor,
+    )
+
+
+def fast_sampling(mask_mod):
+    """Convenience decorator to mark mask_mod as safe for 5-point fast sampling"""
+    mask_mod.use_fast_sampling = True
+    return mask_mod

python/sglang/jit_kernel/flash_attn/cute/cache_utils.py

@@ -0,0 +1,289 @@
+# Manage Ahead-of-Time (AOT) compiled kernels
+import ctypes
+import fcntl
+import hashlib
+import os
+import pickle
+import sys
+import tempfile
+import time
+from functools import lru_cache
+from getpass import getuser
+from pathlib import Path
+from typing import Hashable, TypeAlias
+
+import cutlass
+import cutlass.cute as cute
+import tvm_ffi
+from cutlass.cutlass_dsl import JitCompiledFunction
+
+from sglang.jit_kernel.flash_attn.cute.fa_logging import fa_log
+
+# Pre-load cute DSL runtime libraries with RTLD_GLOBAL so that their symbols
+# (e.g. _cudaLibraryLoadData) are visible to .so modules loaded later via dlopen.
+# Upstream cute.runtime.load_module loads these without RTLD_GLOBAL, which causes
+# "undefined symbol" errors when loading cached kernels from disk.
+for _lib_path in cute.runtime.find_runtime_libraries(enable_tvm_ffi=False):
+    if Path(_lib_path).exists():
+        ctypes.CDLL(_lib_path, mode=ctypes.RTLD_GLOBAL)
+
+CompileKeyType: TypeAlias = tuple[Hashable, ...]
+CallableFunction: TypeAlias = JitCompiledFunction | tvm_ffi.Function
+
+# Enable cache via `FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1`
+CUTE_DSL_CACHE_ENABLED: bool = (
+    os.getenv("FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED", "0") == "1"
+)
+
+
+# Customize cache dir via `FLASH_ATTENTION_CUTE_DSL_CACHE_DIR`, default is
+# `/tmp/${USER}/flash_attention_cute_dsl_cache``
+CUTE_DSL_CACHE_DIR: str | None = os.getenv("FLASH_ATTENTION_CUTE_DSL_CACHE_DIR", None)
+
+
+def get_cache_path() -> Path:
+    if CUTE_DSL_CACHE_DIR is not None:
+        cache_dir = Path(CUTE_DSL_CACHE_DIR)
+    else:
+        cache_dir = (
+            Path(tempfile.gettempdir()) / getuser() / "flash_attention_cute_dsl_cache"
+        )
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
+@lru_cache(maxsize=1)
+def _compute_source_fingerprint() -> str:
+    """
+    Hash all CuTe Python sources plus runtime ABI stamps into a short fingerprint.
+
+    The fingerprint changes whenever:
+    - Any .py file under flash_attn/cute is added, removed, renamed, or modified.
+    - The Python minor version changes (e.g. 3.13 -> 3.14).
+    - The cutlass or tvm_ffi package version changes.
+
+    Computed once per process and cached.
+    """
+    cute_root = Path(__file__).resolve().parent
+    h = hashlib.sha256()
+
+    h.update(f"py{sys.version_info.major}.{sys.version_info.minor}".encode())
+    h.update(f"cutlass={cutlass.__version__}".encode())
+    h.update(f"tvm_ffi={tvm_ffi.__version__}".encode())
+
+    for src in sorted(cute_root.rglob("*.py")):
+        if not src.is_file():
+            continue
+        h.update(src.relative_to(cute_root).as_posix().encode())
+        content = src.read_bytes()
+        h.update(len(content).to_bytes(8, "little"))
+        h.update(content)
+
+    return h.hexdigest()
+
+
+class FileLock:
+    """Context manager for advisory file locks using fcntl.flock.
+
+    Supports exclusive (write) and shared (read) locks.
+    Always blocks with polling until the lock is acquired or timeout is reached.
+
+    Usage:
+        with FileLock(lock_path, exclusive=True, timeout=15, label="abc"):
+            # do work under lock
+    """
+
+    def __init__(
+        self,
+        lock_path: Path,
+        exclusive: bool,
+        timeout: float = 15,
+        label: str = "",
+    ):
+        """
+        Args:
+            lock_path: Path to the lock file on disk.
+            exclusive: True for exclusive (write) lock, False for shared (read) lock.
+            timeout: Max seconds to wait for lock acquisition before raising RuntimeError.
+            label: Optional human-readable label for error messages.
+        """
+        self.lock_path: Path = lock_path
+        self.exclusive: bool = exclusive
+        self.timeout: float = timeout
+        self.label: str = label
+        self._fd: int = -1
+
+    @property
+    def _lock_label(self) -> str:
+        kind = "exclusive" if self.exclusive else "shared"
+        return f"{kind} {self.label}" if self.label else kind
+
+    def __enter__(self) -> "FileLock":
+        open_flags = (
+            os.O_WRONLY | os.O_CREAT if self.exclusive else os.O_RDONLY | os.O_CREAT
+        )
+        lock_type = fcntl.LOCK_EX if self.exclusive else fcntl.LOCK_SH
+
+        self._fd = os.open(str(self.lock_path), open_flags)
+
+        deadline = time.monotonic() + self.timeout
+        acquired = False
+        while time.monotonic() < deadline:
+            try:
+                fcntl.flock(self._fd, lock_type | fcntl.LOCK_NB)
+                acquired = True
+                break
+            except OSError:
+                time.sleep(0.1)
+        if not acquired:
+            os.close(self._fd)
+            self._fd = None
+            raise RuntimeError(
+                f"Timed out after {self.timeout}s waiting for "
+                f"{self._lock_label} lock: {self.lock_path}"
+            )
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        if self._fd is not None:
+            fcntl.flock(self._fd, fcntl.LOCK_UN)
+            os.close(self._fd)
+            self._fd = None
+
+
+class JITCache:
+    """
+    In-memory cache for compiled functions.
+    """
+
+    def __init__(self):
+        self.cache: dict[CompileKeyType, CallableFunction] = {}
+
+    def __setitem__(self, key: CompileKeyType, fn: JitCompiledFunction) -> None:
+        self.cache[key] = fn
+
+    def __getitem__(self, key: CompileKeyType) -> CallableFunction:
+        return self.cache[key]
+
+    def __contains__(self, key: CompileKeyType) -> bool:
+        return key in self.cache
+
+    def clear(self) -> None:
+        """
+        Clear in-memory cache of compiled functions
+        """
+        self.cache.clear()
+
+
+class JITPersistentCache(JITCache):
+    """
+    In-memory cache for compiled functions, which is also backed by persistent storage.
+    Use cutedsl ahead-of-time (AOT) compilation, only supporting enable_tvm_ffi=True
+    """
+
+    EXPORT_FUNCTION_PREFIX = "func"
+    LOCK_TIMEOUT_SECONDS = 15
+
+    def __init__(self, cache_path: Path):
+        super().__init__()
+        cache_path.mkdir(parents=True, exist_ok=True)
+        self.cache_path: Path = cache_path
+
+    def __setitem__(self, key: CompileKeyType, fn: JitCompiledFunction) -> None:
+        JITCache.__setitem__(self, key, fn)
+        self._try_export_to_storage(key, fn)
+
+    def __getitem__(self, key: CompileKeyType) -> CallableFunction:
+        # Use __contains__ to try populating in-memory cache with persistent storage
+        self.__contains__(key)
+        return JITCache.__getitem__(self, key)
+
+    def __contains__(self, key: CompileKeyType) -> bool:
+        # Checks in-memory cache first, then tries loading from storage.
+        # When returning True, guarantees the in-memory cache is populated.
+        if JITCache.__contains__(self, key):
+            return True
+        return self._try_load_from_storage(key)
+
+    def _try_load_from_storage(self, key: CompileKeyType) -> bool:
+        """
+        Try to load a function from persistent storage into in-memory cache.
+        Returns True if loaded successfully, False if not found on disk.
+        Holds a shared lock during loading to prevent concurrent writes.
+        """
+        sha256_hex = self._key_to_hash(key)
+        obj_path = self.cache_path / f"{sha256_hex}.o"
+        with FileLock(
+            self._lock_path(sha256_hex),
+            exclusive=False,
+            timeout=self.LOCK_TIMEOUT_SECONDS,
+            label=sha256_hex,
+        ):
+            if obj_path.exists():
+                fa_log(1, f"Loading compiled function from disk: {obj_path}")
+                m = cute.runtime.load_module(str(obj_path), enable_tvm_ffi=True)
+                fn = getattr(m, self.EXPORT_FUNCTION_PREFIX)
+                JITCache.__setitem__(self, key, fn)
+                return True
+            else:
+                fa_log(1, f"Cache miss on disk for key hash {sha256_hex}")
+        return False
+
+    def _try_export_to_storage(
+        self, key: CompileKeyType, fn: JitCompiledFunction
+    ) -> None:
+        """Export a compiled function to persistent storage under exclusive lock."""
+        sha256_hex = self._key_to_hash(key)
+        with FileLock(
+            self._lock_path(sha256_hex),
+            exclusive=True,
+            timeout=self.LOCK_TIMEOUT_SECONDS,
+            label=sha256_hex,
+        ):
+            obj_path = self.cache_path / f"{sha256_hex}.o"
+            if obj_path.exists():
+                # Another process already exported.
+                fa_log(1, f"Skipping export, already on disk: {obj_path}")
+                return
+            fa_log(1, f"Exporting compiled function to disk: {obj_path}")
+            fn.export_to_c(
+                object_file_path=str(obj_path),
+                function_name=self.EXPORT_FUNCTION_PREFIX,
+            )
+            fa_log(1, f"Successfully exported compiled function to disk: {obj_path}")
+
+    def _key_to_hash(self, key: CompileKeyType) -> str:
+        return hashlib.sha256(pickle.dumps(key)).hexdigest()
+
+    def _lock_path(self, sha256_hex: str) -> Path:
+        return self.cache_path / f"{sha256_hex}.lock"
+
+    def clear(self) -> None:
+        """
+        Not only clear the in-memory cache. Also purge persistent compilation cache.
+        """
+        fa_log(1, f"Clearing persistent cache at {self.cache_path}")
+        super().clear()
+        for child in self.cache_path.iterdir():
+            child.unlink()
+
+
+def get_jit_cache(name: str | None = None) -> JITCache:
+    """
+    JIT cache factory.
+    `name` is an optional identifier to create subdirectories to manage cache.
+
+    When persistent caching is enabled, artifacts are namespaced under a
+    source fingerprint directory so that code or dependency changes
+    automatically invalidate stale entries.
+    """
+    if CUTE_DSL_CACHE_ENABLED:
+        path = get_cache_path() / _compute_source_fingerprint()
+        if name:
+            path = path / name
+        fa_log(1, f"Creating persistent JIT cache at {path}")
+        return JITPersistentCache(path)
+    else:
+        fa_log(1, "Persistent cache disabled, using in-memory JIT cache")
+        return JITCache()

python/sglang/jit_kernel/flash_attn/cute/compute_block_sparsity.py

@@ -0,0 +1,591 @@
+from functools import partial
+from typing import Callable, Optional, Tuple
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import Boolean, Int8, Int32, const_expr
+
+from sglang.jit_kernel.flash_attn.cute.block_sparse_utils import (
+    get_curr_blocksparse_tensors,
+)
+from sglang.jit_kernel.flash_attn.cute.block_sparsity import (
+    BlockSparseTensors,
+    BlockSparseTensorsTorch,
+    to_cute_block_sparse_tensors,
+)
+from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import (
+    get_aux_tensor_metadata,
+    to_cute_aux_tensor,
+    to_cute_tensor,
+)
+from sglang.jit_kernel.flash_attn.cute.mask import call_mask_mod
+from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK
+from sglang.jit_kernel.flash_attn.cute.testing import is_fake_mode
+from sglang.jit_kernel.flash_attn.cute.utils import (
+    AuxData,
+    get_batch_from_cu_tensor,
+    hash_callable,
+    scalar_to_ssa,
+    ssa_to_scalar,
+)
+
+
+class BlockSparsityKernel:
+    """Block sparsity kernel for FlexAttention.
+
+    This kernel computes `mask_mod` for every token of each block
+    to determine if an n block is full, masked, or neither.
+
+    Writes block counts and indices to a BlockSparseTensors object.
+
+    When use_fast_sampling=True, uses 5-point sampling (4 corners + center)
+    which is much faster but only suitable for masks where this is sufficient.
+
+    TODO:
+        - optimize mask_mod evaluation
+        - transposed tensors for bwd pass
+    """
+
+    def __init__(
+        self,
+        mask_mod: Callable,
+        tile_mn: Tuple[int, int],
+        compute_full_blocks: bool = True,
+        use_aux_tensors: bool = False,
+        use_fast_sampling: bool = False,
+    ):
+        self.mask_mod = mask_mod
+        self.tile_mn = tile_mn
+        self.compute_full_blocks = compute_full_blocks
+        self.use_aux_tensors = use_aux_tensors
+        self.use_fast_sampling = use_fast_sampling
+
+    @cute.jit
+    def __call__(
+        self,
+        blocksparse_tensors: BlockSparseTensors,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        aux_data: AuxData = AuxData(),
+    ):
+        (
+            mask_cnt,
+            mask_idx,
+            full_cnt,
+            full_idx,
+            mCuTotalMBlocks,
+            mCuBlockIdxOffsets,
+            *_,
+        ) = blocksparse_tensors
+
+        self.is_varlen_q = const_expr(mCuSeqlensQ is not None)
+
+        if const_expr(self.compute_full_blocks):
+            assert (
+                full_cnt is not None and full_idx is not None
+            ), "full block tensors must be provided when computing full blocks"
+        if const_expr(not self.is_varlen_q):
+            batch_size, num_heads, num_m_blocks, _ = mask_idx.shape
+            total_m_blocks = batch_size * num_m_blocks
+        else:
+            assert const_expr(
+                mCuTotalMBlocks is not None
+            ), "mCuTotalMBlocks must be provided when varlen q"
+            num_heads, total_m_blocks = mask_cnt.shape  # num_m_blocks is total_m_blocks
+            batch_size = mCuSeqlensQ.shape[0] - 1
+
+        if const_expr(self.use_fast_sampling):
+            num_threads = 5
+            self.num_warps = 1
+        else:
+            num_threads = self.tile_mn[0]
+            self.num_warps = (num_threads + 32 - 1) // 32
+
+        if const_expr(not self.is_varlen_q):
+            grid = [num_m_blocks, num_heads, batch_size]
+        else:
+            grid = [total_m_blocks, num_heads, 1]
+
+        self.kernel(
+            blocksparse_tensors,
+            seqlen_q,
+            seqlen_k,
+            batch_size,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mSeqUsedQ,
+            mSeqUsedK,
+            mCuTotalMBlocks,
+            mCuBlockIdxOffsets,
+            aux_data,
+        ).launch(grid=grid, block=[num_threads, 1, 1])
+
+    @cute.kernel
+    def kernel(
+        self,
+        blocksparse_tensors: BlockSparseTensors,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        batch_size: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        mCuBlockIdxOffsets: Optional[cute.Tensor] = None,
+        aux_data: AuxData = AuxData(),
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        warp_idx = cute.arch.warp_idx()
+        lane_id = cute.arch.lane_idx()
+
+        ssa = partial(scalar_to_ssa, dtype=Int32)
+
+        @cute.struct
+        class SharedStorage:
+            reduction_buffer_smem: cute.struct.Align[
+                cute.struct.MemRange[cutlass.Int8, 2 * self.num_warps], 1024
+            ]
+
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage, 16)
+
+        reduction_buffer = storage.reduction_buffer_smem.get_tensor(
+            cute.make_layout((self.num_warps, 2))
+        )
+        SeqlenInfoCls = partial(
+            SeqlenInfoQK.create,
+            seqlen_q_static=seqlen_q,
+            seqlen_k_static=seqlen_k,
+            mCuSeqlensQ=mCuSeqlensQ,
+            mCuSeqlensK=mCuSeqlensK,
+            mSeqUsedQ=mSeqUsedQ,
+            mSeqUsedK=mSeqUsedK,
+            mCuTotalMBlocks=mCuTotalMBlocks,
+            mCuBlockIdxOffsets=mCuBlockIdxOffsets,
+            tile_m=self.tile_mn[0],
+            tile_n=self.tile_mn[1],
+        )
+
+        if const_expr(not self.is_varlen_q):
+            m_block, head_idx, batch_idx = cute.arch.block_idx()
+        else:
+            global_m_block, head_idx, _ = cute.arch.block_idx()
+            batch_idx = get_batch_from_cu_tensor(global_m_block, mCuTotalMBlocks)
+            m_block = global_m_block - mCuTotalMBlocks[batch_idx]
+
+        seqlen = SeqlenInfoCls(batch_idx)
+        seqlen_q = seqlen.seqlen_q
+        seqlen_k = seqlen.seqlen_k
+        global_m_block = seqlen.m_block_offset + m_block
+
+        num_n_blocks = (seqlen_k + self.tile_mn[1] - 1) // self.tile_mn[1]
+
+        _, curr_mask_idx, _, curr_full_idx = get_curr_blocksparse_tensors(
+            batch_idx, head_idx, m_block, blocksparse_tensors, seqlen
+        )
+
+        num_mask_blocks = Int32(0)
+        num_full_blocks = Int32(0)
+
+        m_base = m_block * self.tile_mn[0]
+        if const_expr(self.use_fast_sampling):
+            # Loop-invariant per-thread q_idx for the 5 sample points
+            # (tidx 0, 1: top corners; 2, 3: bottom corners; 4: center).
+            q_idx_sample = m_base
+            if tidx == 2 or tidx == 3:
+                q_idx_sample = cutlass.min(m_base + self.tile_mn[0] - 1, seqlen_q - 1)
+            elif tidx == 4:
+                q_idx_sample = (
+                    m_base + cutlass.min(seqlen_q - m_base, self.tile_mn[0]) // 2
+                )
+        else:
+            q_idx_thread = m_base + tidx
+            thread_in_bounds = Boolean(
+                tidx < self.tile_mn[0] and q_idx_thread < seqlen_q
+            )
+
+        for n_block in cutlass.range(num_n_blocks):
+            n_base = n_block * self.tile_mn[1]
+
+            if const_expr(self.use_fast_sampling):
+                # 5-point sampling (4 corners + center). Interior n_blocks
+                # (n_base + tile_n <= seqlen_k) skip the OOB clamp on the right /
+                # center samples.
+                is_interior = (n_base + self.tile_mn[1]) <= seqlen_k
+                n_right = Int32(0)
+                n_mid = Int32(0)
+                if is_interior:
+                    n_right = n_base + self.tile_mn[1] - 1
+                    n_mid = n_base + self.tile_mn[1] // 2
+                else:
+                    n_right = cutlass.min(n_base + self.tile_mn[1] - 1, seqlen_k - 1)
+                    n_mid = (
+                        n_base + cutlass.min(seqlen_k - n_base, self.tile_mn[1]) // 2
+                    )
+
+                kv_idx = n_base
+                if tidx == 1 or tidx == 3:
+                    kv_idx = n_right
+                elif tidx == 4:
+                    kv_idx = n_mid
+
+                thread_result = Boolean(False)
+                thread_is_valid = Boolean(False)
+                if tidx < 5:
+                    thread_is_valid = Boolean(True)
+                    thread_result = ssa_to_scalar(
+                        call_mask_mod(
+                            self.mask_mod,
+                            ssa(batch_idx),
+                            ssa(head_idx),
+                            ssa(q_idx_sample),
+                            ssa(kv_idx),
+                            seqlen,
+                            aux_data,
+                        )
+                    )
+
+                has_unmasked = cute.arch.vote_any_sync(thread_result & thread_is_valid)
+                has_masked = cute.arch.vote_any_sync(
+                    Boolean(not thread_result) & thread_is_valid
+                )
+
+            else:
+                # Full path. Interior blocks (n_base + tile_n <= seqlen_k) drop the
+                # per-element bound check; the boundary block (at most one) keeps it.
+                thread_has_unmasked = Boolean(False)
+                thread_has_masked = Boolean(False)
+                kv_idx = Int32(0)
+                is_interior = (n_base + self.tile_mn[1]) <= seqlen_k
+
+                if is_interior:
+                    if thread_in_bounds:
+                        for c in cutlass.range(self.tile_mn[1], unroll_full=True):
+                            mask_val = ssa_to_scalar(
+                                call_mask_mod(
+                                    self.mask_mod,
+                                    ssa(batch_idx),
+                                    ssa(head_idx),
+                                    ssa(q_idx_thread),
+                                    ssa(n_base + c),
+                                    seqlen,
+                                    aux_data,
+                                )
+                            )
+                            thread_has_unmasked |= Boolean(mask_val)
+                            thread_has_masked |= Boolean(not mask_val)
+                else:
+                    if thread_in_bounds:
+                        for c in cutlass.range(self.tile_mn[1], unroll_full=True):
+                            kv_idx = n_base + c
+                            if kv_idx < seqlen_k:
+                                mask_val = ssa_to_scalar(
+                                    call_mask_mod(
+                                        self.mask_mod,
+                                        ssa(batch_idx),
+                                        ssa(head_idx),
+                                        ssa(q_idx_thread),
+                                        ssa(kv_idx),
+                                        seqlen,
+                                        aux_data,
+                                    )
+                                )
+                                thread_has_unmasked |= Boolean(mask_val)
+                                thread_has_masked |= Boolean(not mask_val)
+
+                warp_unmasked = cute.arch.vote_any_sync(
+                    thread_has_unmasked & thread_in_bounds
+                )
+                warp_masked = cute.arch.vote_any_sync(
+                    thread_has_masked & thread_in_bounds
+                )
+                if lane_id == 0:
+                    reduction_buffer[warp_idx, 0] = (
+                        Int8(1) if warp_unmasked else Int8(0)
+                    )
+                    reduction_buffer[warp_idx, 1] = Int8(1) if warp_masked else Int8(0)
+                cute.arch.sync_threads()
+
+                # Cross-warp OR via warp 0; thread 0 (lane 0 of warp 0) holds the result.
+                has_unmasked = Boolean(False)
+                has_masked = Boolean(False)
+                if warp_idx == 0:
+                    lane_unmasked = Boolean(False)
+                    lane_masked = Boolean(False)
+                    if lane_id < self.num_warps:
+                        lane_unmasked = reduction_buffer[lane_id, 0] != Int8(0)
+                        lane_masked = reduction_buffer[lane_id, 1] != Int8(0)
+                    has_unmasked = cute.arch.vote_any_sync(lane_unmasked)
+                    has_masked = cute.arch.vote_any_sync(lane_masked)
+
+            # Only thread 0 updates the output arrays (common to both paths)
+            if tidx == 0:
+                # Block classification based on what we found:
+                # - If has_masked and has_unmasked: partial block (needs masking)
+                # - If only has_unmasked: full block (no masking needed)
+                # - If only has_masked: skip this block entirely
+                is_partial = Boolean(has_masked and has_unmasked)
+                is_full = Boolean(has_unmasked and (not has_masked))
+
+                if is_partial:
+                    curr_mask_idx[num_mask_blocks] = n_block
+                    num_mask_blocks += 1
+                elif is_full and const_expr(self.compute_full_blocks):
+                    curr_full_idx[num_full_blocks] = n_block
+                    num_full_blocks += 1
+
+        # Only thread 0 writes back the counts
+        if tidx == 0:
+            mask_cnt, _, full_cnt, *_ = blocksparse_tensors
+            if const_expr(self.is_varlen_q):
+                mask_cnt[head_idx, global_m_block] = num_mask_blocks
+                if const_expr(self.compute_full_blocks):
+                    full_cnt[head_idx, global_m_block] = num_full_blocks
+            else:
+                mask_cnt[batch_idx, head_idx, m_block] = num_mask_blocks
+                if const_expr(self.compute_full_blocks):
+                    full_cnt[batch_idx, head_idx, m_block] = num_full_blocks
+
+
+def compute_block_sparsity(
+    tile_m,
+    tile_n,
+    batch_size,
+    num_heads,
+    seqlen_q,
+    seqlen_k,
+    mask_mod: Callable,
+    aux_tensors: Optional[list],
+    device,
+    aux_scalars: Optional[tuple] = None,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    cu_total_m_blocks: Optional[torch.Tensor] = None,
+    cu_block_idx_offsets: Optional[torch.Tensor] = None,
+    compute_full_blocks: bool = True,
+    use_fast_sampling: bool = False,
+) -> BlockSparseTensorsTorch:
+    """
+    Computes block sparsity for a given `mask_mod`.
+
+    Args:
+        tile_m: The tile size for the m dimension.
+        tile_n: The tile size for the n dimension.
+        batch_size: The batch size.
+        num_heads: The number of heads.
+        seqlen_q: The sequence length for the query.
+        seqlen_k: The sequence length for the key.
+        mask_mod: The `mask_mod` callable to use.
+        aux_tensors: A list of auxiliary tensors.
+        device: The device to use.
+        cu_seqlens_q: Cumulative q sequence lengths for varlen
+        cu_seqlens_k: Cumulative k sequence lengths for varlen
+        seqused_q: Per-batch effective q sequence lengths
+        seqused_k: Per-batch effective k sequence lengths
+        cu_total_m_blocks: Cumulative total m blocks tensor for varlen q
+        cu_block_idx_offsets: Cumulative offsets into the packed mask_block_idx /
+            full_block_idx tensors per batch (== cumsum of M_b * N_b).
+        compute_full_blocks: Whether to compute full blocks. If False, only partially-masked blocks are computed.
+        use_fast_sampling: Whether to use 5-point sampling (4 corners + center). This is much faster, but only suitable for masks where this check is sufficient.
+
+    Returns:
+        BlockSparseTensorsTorch
+    """
+    aux_scalars = tuple(aux_scalars) if aux_scalars else None
+
+    # Check if mask_mod is marked as suitable for 5-point sampling
+    use_fast_sampling = getattr(mask_mod, "use_fast_sampling", use_fast_sampling)
+
+    num_m_blocks = (seqlen_q + tile_m - 1) // tile_m
+    num_n_blocks = (seqlen_k + tile_n - 1) // tile_n
+
+    if cu_seqlens_q is not None:
+        assert (
+            cu_total_m_blocks is not None
+        ), "total m blocks must be provided when varlen q"
+        total_m_blocks = cu_total_m_blocks[-1].item()
+        if cu_block_idx_offsets is None and (
+            cu_seqlens_k is not None or seqused_k is not None
+        ):
+            # Derive cu_block_idx_offsets from per-batch K seqlens.
+            cu_block_idx_offsets_list = [0]
+            for batch_idx in range(batch_size):
+                batch_seqlen_q = (
+                    cu_seqlens_q[batch_idx + 1].item() - cu_seqlens_q[batch_idx].item()
+                )
+                if cu_seqlens_k is not None:
+                    batch_seqlen_k = (
+                        cu_seqlens_k[batch_idx + 1].item()
+                        - cu_seqlens_k[batch_idx].item()
+                    )
+                else:
+                    batch_seqlen_k = seqused_k[batch_idx].item()
+                num_m_blocks_batch = (batch_seqlen_q + tile_m - 1) // tile_m
+                num_n_blocks_batch = (batch_seqlen_k + tile_n - 1) // tile_n
+                cu_block_idx_offsets_list.append(
+                    cu_block_idx_offsets_list[-1]
+                    + num_m_blocks_batch * num_n_blocks_batch
+                )
+            cu_block_idx_offsets = torch.tensor(
+                cu_block_idx_offsets_list, dtype=torch.int32, device=device
+            )
+        if cu_block_idx_offsets is not None:
+            total_n_blocks = cu_block_idx_offsets[-1].item()
+        else:
+            # Uniform-K varlen-Q: every batch has the same K seqlen.
+            total_n_blocks = total_m_blocks * num_n_blocks
+
+        mask_block_cnt = torch.zeros(
+            (num_heads, total_m_blocks), device=device, dtype=torch.int32
+        )
+        mask_block_idx = torch.zeros(
+            (num_heads, total_n_blocks), device=device, dtype=torch.int32
+        )
+        full_block_cnt = (
+            torch.zeros((num_heads, total_m_blocks), device=device, dtype=torch.int32)
+            if compute_full_blocks
+            else None
+        )
+        full_block_idx = (
+            torch.zeros((num_heads, total_n_blocks), device=device, dtype=torch.int32)
+            if compute_full_blocks
+            else None
+        )
+    else:
+        total_m_blocks = batch_size * num_m_blocks
+        total_n_blocks = batch_size * num_m_blocks * num_n_blocks
+
+        mask_block_cnt = torch.zeros(
+            (batch_size, num_heads, num_m_blocks), device=device, dtype=torch.int32
+        )
+        mask_block_idx = torch.zeros(
+            (batch_size, num_heads, num_m_blocks, num_n_blocks),
+            device=device,
+            dtype=torch.int32,
+        )
+        full_block_cnt = (
+            torch.zeros(
+                (batch_size, num_heads, num_m_blocks), device=device, dtype=torch.int32
+            )
+            if compute_full_blocks
+            else None
+        )
+        full_block_idx = (
+            torch.zeros(
+                (batch_size, num_heads, num_m_blocks, num_n_blocks),
+                device=device,
+                dtype=torch.int32,
+            )
+            if compute_full_blocks
+            else None
+        )
+
+    blocksparse_tensors_torch = BlockSparseTensorsTorch(
+        mask_block_cnt=mask_block_cnt,
+        mask_block_idx=mask_block_idx,
+        full_block_cnt=full_block_cnt,
+        full_block_idx=full_block_idx,
+        cu_total_m_blocks=cu_total_m_blocks,
+        cu_block_idx_offsets=cu_block_idx_offsets,
+        block_size=(tile_m, tile_n),
+    )
+
+    mask_mod_hash = hash_callable(mask_mod)
+    if aux_tensors is not None:
+        aux_tensor_metadata = get_aux_tensor_metadata(aux_tensors)
+    else:
+        aux_tensor_metadata = None
+    aux_scalar_metadata = (
+        tuple(type(s) for s in aux_scalars) if aux_scalars is not None else None
+    )
+
+    compile_key = (
+        tile_m,
+        tile_n,
+        mask_mod_hash,
+        aux_tensor_metadata,
+        aux_scalar_metadata,
+        compute_full_blocks,
+        cu_seqlens_q is None,
+        cu_seqlens_k is None,
+        seqused_q is None,
+        seqused_k is None,
+        aux_tensors is not None,
+        use_fast_sampling,
+    )
+    if compile_key not in compute_block_sparsity.compile_cache:
+        (
+            cu_seqlens_q_tensor,
+            cu_seqlens_k_tensor,
+            seqused_q_tensor,
+            seqused_k_tensor,
+        ) = [
+            to_cute_tensor(t, assumed_align=4, leading_dim=0) if t is not None else None
+            for t in (
+                cu_seqlens_q,
+                cu_seqlens_k,
+                seqused_q,
+                seqused_k,
+            )
+        ]
+        blocksparse_tensors = to_cute_block_sparse_tensors(
+            blocksparse_tensors_torch, enable_tvm_ffi=True
+        )
+        if aux_tensors is not None:
+            cute_aux_tensors = [to_cute_aux_tensor(buf) for buf in aux_tensors]
+        else:
+            cute_aux_tensors = None
+        kernel = BlockSparsityKernel(
+            mask_mod,
+            tile_mn=(tile_m, tile_n),
+            compute_full_blocks=compute_full_blocks,
+            use_aux_tensors=aux_tensors is not None,
+            use_fast_sampling=use_fast_sampling,
+        )
+
+        compute_block_sparsity.compile_cache[compile_key] = cute.compile(
+            kernel,
+            blocksparse_tensors,
+            seqlen_q,
+            seqlen_k,
+            cu_seqlens_q_tensor,
+            cu_seqlens_k_tensor,
+            seqused_q_tensor,
+            seqused_k_tensor,
+            AuxData(cute_aux_tensors, aux_scalars),
+            options="--enable-tvm-ffi",
+        )
+
+    if not is_fake_mode():
+        compute_block_sparsity.compile_cache[compile_key](
+            (
+                blocksparse_tensors_torch.mask_block_cnt,
+                blocksparse_tensors_torch.mask_block_idx,
+                blocksparse_tensors_torch.full_block_cnt,
+                blocksparse_tensors_torch.full_block_idx,
+                blocksparse_tensors_torch.cu_total_m_blocks,
+                blocksparse_tensors_torch.cu_block_idx_offsets,
+                blocksparse_tensors_torch.dq_write_order,
+                blocksparse_tensors_torch.dq_write_order_full,
+            ),
+            seqlen_q,
+            seqlen_k,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+            AuxData(aux_tensors, aux_scalars),
+        )
+
+    return blocksparse_tensors_torch
+
+
+compute_block_sparsity.compile_cache = {}

python/sglang/jit_kernel/flash_attn/cute/copy_utils.py

@@ -0,0 +1,402 @@
+# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
+
+import math
+from typing import Callable, Optional, Type
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass import Float32, Int32, const_expr
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.nvgpu import cpasync
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+
+@dsl_user_op
+def cvt_copy(
+    atom: cute.CopyAtom,
+    src: cute.Tensor,
+    dst: cute.Tensor,
+    *,
+    pred: Optional[cute.Tensor] = None,
+    loc=None,
+    ip=None,
+    **kwargs,
+) -> None:
+    assert (
+        isinstance(src.iterator, cute.Pointer)
+        and src.memspace == cute.AddressSpace.rmem
+    )
+    if const_expr(src.element_type != dst.element_type):
+        src_cvt = cute.make_fragment_like(src, dst.element_type, loc=loc, ip=ip)
+        src_cvt.store(src.load().to(dst.element_type))
+        src = src_cvt
+    cute.copy(atom, src, dst, pred=pred, loc=loc, ip=ip, **kwargs)
+
+
+@dsl_user_op
+def load_s2r(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
+    dst = cute.make_fragment_like(src, src.element_type, loc=loc, ip=ip)
+    cute.autovec_copy(src, dst, loc=loc, ip=ip)
+    return dst
+
+
+@dsl_user_op
+def get_copy_atom(
+    dtype: Type[cutlass.Numeric],
+    num_copy_elems: int,
+    is_async: bool = False,
+    *,
+    loc=None,
+    ip=None,
+) -> cute.CopyAtom:
+    num_copy_bits = const_expr(min(128, num_copy_elems * dtype.width))
+    copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
+    return cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
+
+
+@dsl_user_op
+def make_tmem_copy(
+    tmem_copy_atom: cute.CopyAtom, num_wg: int = 1, *, loc=None, ip=None
+) -> cute.CopyAtom:
+    num_dp, num_bits, num_rep, _ = sm100_utils.get_tmem_copy_properties(tmem_copy_atom)
+    assert num_dp == 32
+    assert num_bits == 32
+    tiler_mn = (cute.make_layout((128 * num_rep * num_wg // 32, 32), stride=(32, 1)),)
+    layout_tv = cute.make_layout(
+        ((32, 4, num_wg), (num_rep, 32)),
+        stride=((0, 1, 4 * num_rep), (4, 4 * num_rep * num_wg)),
+    )
+    return cute.make_tiled_copy(tmem_copy_atom, layout_tv, tiler_mn)
+
+
+@dsl_user_op
+def copy(
+    src: cute.Tensor,
+    dst: cute.Tensor,
+    *,
+    pred: Optional[cute.Tensor] = None,
+    num_copy_elems: int = 1,
+    is_async: bool = False,
+    loc=None,
+    ip=None,
+    **kwargs,
+) -> None:
+    copy_atom = get_copy_atom(src.element_type, num_copy_elems, is_async)
+    cute.copy(copy_atom, src, dst, pred=pred, loc=loc, ip=ip, **kwargs)
+
+
+def tiled_copy_1d(
+    dtype: Type[cutlass.Numeric],
+    num_threads: int,
+    num_copy_elems: int = 1,
+    is_async: bool = False,
+) -> cute.TiledCopy:
+    num_copy_bits = num_copy_elems * dtype.width
+    copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
+    copy_atom = cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
+    thr_layout = cute.make_layout(num_threads)
+    val_layout = cute.make_layout(num_copy_elems)
+    return cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
+
+
+def tiled_copy_2d(
+    dtype: Type[cutlass.Numeric],
+    major_mode_size: int,
+    num_threads: int,
+    is_async: bool = False,
+) -> cute.TiledCopy:
+    num_copy_bits = math.gcd(major_mode_size, 128 // dtype.width) * dtype.width
+    copy_elems = num_copy_bits // dtype.width
+    copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
+    copy_atom = cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
+    gmem_threads_per_row = major_mode_size // copy_elems
+    assert num_threads % gmem_threads_per_row == 0
+    thr_layout = cute.make_ordered_layout(
+        (num_threads // gmem_threads_per_row, gmem_threads_per_row),
+        order=(1, 0),
+    )
+    val_layout = cute.make_layout((1, copy_elems))
+    return cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
+
+
+@dsl_user_op
+def atomic_add_fp32x4(
+    a: Float32,
+    b: Float32,
+    c: Float32,
+    d: Float32,
+    gmem_ptr: cute.Pointer,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value()
+    # cache_hint = cutlass.Int64(0x12F0000000000000)
+    llvm.inline_asm(
+        None,
+        [
+            gmem_ptr_i64,
+            Float32(a).ir_value(loc=loc, ip=ip),
+            Float32(b).ir_value(loc=loc, ip=ip),
+            Float32(c).ir_value(loc=loc, ip=ip),
+            Float32(d).ir_value(loc=loc, ip=ip),
+        ],
+        # [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip), cache_hint.ir_value()],
+        "{\n\t"
+        # ".reg .b128 abcd;\n\t"
+        # "mov.b128 abcd, {$1, $2, $3, $4};\n\t"
+        ".reg .v4 .f32 abcd;\n\t"
+        # "mov.b128 abcd, {$1, $2, $3, $4};\n\t"
+        "mov.f32 abcd.x, $1;\n\t"
+        "mov.f32 abcd.y, $2;\n\t"
+        "mov.f32 abcd.z, $3;\n\t"
+        "mov.f32 abcd.w, $4;\n\t"
+        "red.global.add.v4.f32 [$0], abcd;\n\t"
+        # "red.global.add.L2::cache_hint.v4.f32 [$0], abcd, 0x14F0000000000000;\n\t"
+        "}\n",
+        # "red.global.add.L2::cache_hint.f32 [$0], $1, 0x12F0000000000000;",
+        # "red.global.add.L2::cache_hint.f32 [$0], $1, $2;",
+        "l,f,f,f,f",
+        # "l,f,l",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def set_block_rank(
+    smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: Int32, *, loc=None, ip=None
+) -> Int32:
+    """Map the given smem pointer to the address at another CTA rank in the cluster."""
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    return Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()],
+            "mapa.shared::cluster.u32 $0, $1, $2;",
+            "=r,r,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@dsl_user_op
+def store_shared_remote_fp32x4(
+    a: Float32,
+    b: Float32,
+    c: Float32,
+    d: Float32,
+    smem_ptr: cute.Pointer,
+    mbar_ptr: cute.Pointer,
+    peer_cta_rank_in_cluster: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    remote_smem_ptr_i32 = set_block_rank(
+        smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
+    ).ir_value()
+    remote_mbar_ptr_i32 = set_block_rank(
+        mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
+    ).ir_value()
+    llvm.inline_asm(
+        None,
+        [
+            remote_smem_ptr_i32,
+            remote_mbar_ptr_i32,
+            Float32(a).ir_value(loc=loc, ip=ip),
+            Float32(b).ir_value(loc=loc, ip=ip),
+            Float32(c).ir_value(loc=loc, ip=ip),
+            Float32(d).ir_value(loc=loc, ip=ip),
+        ],
+        "{\n\t"
+        ".reg .v4 .f32 abcd;\n\t"
+        "mov.f32 abcd.x, $2;\n\t"
+        "mov.f32 abcd.y, $3;\n\t"
+        "mov.f32 abcd.z, $4;\n\t"
+        "mov.f32 abcd.w, $5;\n\t"
+        "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.f32 [$0], abcd, [$1];\n\t"
+        "}\n",
+        "r,r,f,f,f,f",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def cpasync_bulk_s2cluster(
+    smem_src_ptr: cute.Pointer,
+    smem_dst_ptr: cute.Pointer,
+    mbar_ptr: cute.Pointer,
+    size: int | Int32,
+    peer_cta_rank_in_cluster: Int32,
+    *,
+    loc=None,
+    ip=None,
+):
+    smem_src_ptr_i32 = smem_src_ptr.toint(loc=loc, ip=ip).ir_value()
+    smem_dst_ptr_i32 = set_block_rank(
+        smem_dst_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
+    ).ir_value()
+    mbar_ptr_i32 = set_block_rank(
+        mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
+    ).ir_value()
+    llvm.inline_asm(
+        None,
+        [
+            smem_dst_ptr_i32,
+            smem_src_ptr_i32,
+            mbar_ptr_i32,
+            Int32(size).ir_value(loc=loc, ip=ip),
+        ],
+        "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes [$0], [$1], $3, [$2];",
+        "r,r,r,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def cpasync_bulk_g2s(
+    gmem_ptr: cute.Pointer,
+    smem_ptr: cute.Pointer,
+    tma_bar_ptr: cute.Pointer,
+    size: int | Int32,
+    *,
+    loc=None,
+    ip=None,
+):
+    gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value()
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    mbar_ptr_i32 = tma_bar_ptr.toint(loc=loc, ip=ip).ir_value()
+    llvm.inline_asm(
+        None,
+        [gmem_ptr_i64, smem_ptr_i32, mbar_ptr_i32, Int32(size).ir_value()],
+        "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [$1], [$0], $3, [$2];",
+        "l,r,r,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def cpasync_reduce_bulk_add_f32(
+    smem_ptr: cute.Pointer,
+    gmem_ptr: cute.Pointer,
+    store_bytes: int | Int32,
+    *,
+    loc=None,
+    ip=None,
+):
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    # cache_hint = cutlass.Int64(0x14F0000000000000)  # EVICT_LAST
+    llvm.inline_asm(
+        None,
+        [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value()],
+        "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [$0], [$1], $2;",
+        "l,r,r",
+        # [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value(), cache_hint.ir_value()],
+        # "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.f32 [$0], [$1], $2, $3;",
+        # "l,r,r,l",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+def cpasync_bulk_get_copy_fn(
+    src_tensor: cute.Tensor,
+    dst_tensor: cute.Tensor,
+    single_stage: bool = False,
+    **kwargs,
+) -> Callable:
+    # src_is_smem = const_expr(
+    #     isinstance(src_tensor.iterator, cute.Pointer)
+    #     and src_tensor.memspace == cute.AddressSpace.smem
+    # )
+    group_rank_src = const_expr(cute.rank(src_tensor) - (1 if not single_stage else 0))
+    group_rank_dst = const_expr(cute.rank(dst_tensor) - (1 if not single_stage else 0))
+    # ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK)
+    src = cute.group_modes(src_tensor, 0, group_rank_src)
+    dst = cute.group_modes(dst_tensor, 0, group_rank_dst)
+
+    def copy_bulk(src_idx, dst_idx, **new_kwargs):
+        size = const_expr(cute.size(src.shape[:-1]) * src.element_type.width // 8)
+        cpasync_bulk_g2s(
+            src[None, src_idx].iterator,
+            dst[None, dst_idx].iterator,
+            size=size,
+            **new_kwargs,
+            **kwargs,
+        )
+
+    def copy_bulk_single_stage(**new_kwargs):
+        size = const_expr(cute.size(src.shape) * src.element_type.width // 8)
+        cpasync_bulk_g2s(src.iterator, dst.iterator, size=size, **new_kwargs, **kwargs)
+
+    return copy_bulk if const_expr(not single_stage) else copy_bulk_single_stage
+
+
+def tma_get_copy_fn(
+    atom: cute.CopyAtom,
+    cta_coord: cute.Coord,
+    cta_layout: cute.Layout,
+    src_tensor: cute.Tensor,
+    dst_tensor: cute.Tensor,
+    filter_zeros: bool = False,
+    single_stage: bool = False,
+    **kwargs,
+) -> Callable:
+    src_is_smem = const_expr(
+        isinstance(src_tensor.iterator, cute.Pointer)
+        and src_tensor.memspace == cute.AddressSpace.smem
+    )
+    smem_tensor, gmem_tensor = (
+        (src_tensor, dst_tensor) if src_is_smem else (dst_tensor, src_tensor)
+    )
+    group_rank_smem = const_expr(
+        cute.rank(smem_tensor) - (1 if not single_stage else 0)
+    )
+    group_rank_gmem = const_expr(
+        cute.rank(gmem_tensor) - (1 if not single_stage else 0)
+    )
+    # ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK)
+    s, g = cpasync.tma_partition(
+        atom,
+        cta_coord,
+        cta_layout,
+        cute.group_modes(smem_tensor, 0, group_rank_smem),
+        cute.group_modes(gmem_tensor, 0, group_rank_gmem),
+    )
+    if const_expr(filter_zeros):
+        s = cute.filter_zeros(s)
+        g = cute.filter_zeros(g)
+    src, dst = (s, g) if src_is_smem else (g, s)
+
+    def copy_tma(src_idx, dst_idx, **new_kwargs):
+        cute.copy(atom, src[None, src_idx], dst[None, dst_idx], **new_kwargs, **kwargs)
+
+    def copy_tma_single_stage(**new_kwargs):
+        cute.copy(atom, src, dst, **new_kwargs, **kwargs)
+
+    return (copy_tma if const_expr(not single_stage) else copy_tma_single_stage), s, g
+
+
+def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsync):
+    def copy_fn(src_idx, producer_state: cutlass.pipeline.PipelineState, **new_kwargs):
+        copy(
+            src_idx=src_idx,
+            dst_idx=producer_state.index,
+            tma_bar_ptr=pipeline.producer_get_barrier(producer_state),
+            **new_kwargs,
+        )
+
+    return copy_fn

python/sglang/jit_kernel/flash_attn/cute/cu_blocks_kernels.py

@@ -0,0 +1,123 @@
+from typing import Callable
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, const_expr
+
+
+class CuSeqlensToBlocksKernel:
+    """Single-CTA prep for block-packed shear scheduling: computes the cumulative
+    per-batch group-block counts and the block -> batch index map in one launch."""
+
+    def __init__(
+        self,
+        tile: int = 128,
+        num_threads: int = 1024,
+        seqlen_multiple: int = 1,
+        use_pdl: bool = False,
+    ):
+        self.tile = tile
+        self.num_threads = num_threads
+        assert num_threads % 32 == 0
+        self.num_warps = num_threads // cute.arch.WARP_SIZE
+        self.seqlen_multiple = seqlen_multiple
+        self.use_pdl = use_pdl
+
+    @cute.jit
+    def __call__(
+        self,
+        mCuBlocks: cute.Tensor,
+        mCuSeqlens: cute.Tensor,
+        mBlocksToBatchIdx: cute.Tensor,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
+    ):
+        @cute.struct
+        class SharedStorage:
+            warp_block_count: cute.struct.MemRange[Int32, self.num_warps]
+            cu_blocks: cute.struct.MemRange[Int32, self.num_threads + 1]
+
+        self.kernel(
+            mCuBlocks,
+            mCuSeqlens,
+            mBlocksToBatchIdx,
+            SharedStorage,
+        ).launch(
+            grid=[1, 1, 1],
+            block=[self.num_threads, 1, 1],
+            stream=stream,
+            use_pdl=self.use_pdl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mCuBlocks: cute.Tensor,
+        mCuSeqlens: cute.Tensor,
+        mBlocksToBatchIdx: cute.Tensor,
+        SharedStorage: cutlass.Constexpr[Callable],
+    ):
+        if const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+            cute.arch.griddepcontrol_launch_dependents()
+
+        batch_size = mCuBlocks.shape[0] - 1
+        batch_idx = cute.arch.thread_idx()[0]
+        lane_idx = cute.arch.lane_idx()
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+        warp_block_count = storage.warp_block_count.get_tensor(
+            cute.make_layout(self.num_warps)
+        )
+        sCuBlocks = storage.cu_blocks.get_tensor(cute.make_layout(self.num_threads + 1))
+
+        if batch_idx == 0:
+            mCuBlocks[0] = 0
+            sCuBlocks[0] = 0
+
+        seqlen = Int32(0)
+        if batch_idx < batch_size:
+            seqlen = mCuSeqlens[batch_idx + 1] - mCuSeqlens[batch_idx]
+        seqlen *= self.seqlen_multiple
+        num_blocks = (seqlen + self.tile - 1) // self.tile
+
+        total_blocks_for_batch = num_blocks
+        for delta in (1, 2, 4, 8, 16):
+            other = cute.arch.shuffle_sync_up(
+                total_blocks_for_batch, delta, mask_and_clamp=0
+            )
+            if lane_idx >= delta:
+                total_blocks_for_batch += other
+
+        if lane_idx == 31:
+            warp_block_count[warp_idx] = total_blocks_for_batch
+
+        cute.arch.sync_threads()
+
+        if warp_idx * 32 < batch_size:
+            for idx in cutlass.range(warp_idx):
+                total_blocks_for_batch += warp_block_count[idx]
+
+            if batch_idx < batch_size:
+                mCuBlocks[batch_idx + 1] = total_blocks_for_batch
+                sCuBlocks[batch_idx + 1] = total_blocks_for_batch
+
+        cute.arch.sync_threads()
+
+        total_blocks = sCuBlocks[batch_size]
+        num_iters = (total_blocks + self.num_threads - 1) // self.num_threads
+        for it in cutlass.range(num_iters, unroll=1):
+            block = it * self.num_threads + batch_idx
+            if block < total_blocks:
+                lo = Int32(0)
+                hi = Int32(batch_size)
+                while lo < hi:
+                    mid = (lo + hi) // 2
+                    if sCuBlocks[mid + 1] <= block:
+                        lo = mid + 1
+                    else:
+                        hi = mid
+                mBlocksToBatchIdx[block] = lo

python/sglang/jit_kernel/flash_attn/cute/cute_dsl_ptxas.py

@@ -0,0 +1,159 @@
+"""
+System ptxas replacement for CUTLASS DSL.
+Environment variables:
+    CUTE_DSL_PTXAS_PATH    - Path to ptxas (e.g., /usr/local/cuda/bin/ptxas)
+    CUTE_DSL_PTXAS_VERBOSE - Set to 1 for verbose output
+"""
+
+import ctypes
+import os
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+import cutlass
+
+CUTE_DSL_PTXAS_PATH = os.environ.get("CUTE_DSL_PTXAS_PATH", None)
+VERBOSE = os.environ.get("CUTE_DSL_PTXAS_VERBOSE", "0") == "1"
+
+_original_load_cuda_library = None
+_user_wanted_ptx = False  # True if user originally set CUTE_DSL_KEEP_PTX=1
+
+
+def _log(msg):
+    if VERBOSE:
+        print(f"[ptxas] {msg}", file=sys.stderr)
+
+
+def _get_ptx(compiled_func) -> tuple[str, Path] | None:
+    """Find and read PTX file, stripping null bytes."""
+    func_name = getattr(compiled_func, "function_name", None)
+    if not func_name:
+        return None
+
+    dump_dir = os.environ.get("CUTE_DSL_DUMP_DIR", Path.cwd())
+    for ptx_path in Path(dump_dir).glob(f"*{func_name}*.ptx"):
+        content = ptx_path.read_text().rstrip("\x00")
+        if ".entry " in content and content.rstrip().endswith("}"):
+            _log(f"Found PTX: {ptx_path}")
+            return content, ptx_path
+    return None
+
+
+def _compile_ptx(ptx_path: Path, ptx_content: str) -> bytes:
+    """Compile PTX to cubin using system ptxas."""
+    # Extract arch from PTX
+    match = re.search(r"\.target\s+(sm_\d+[a-z]?)", ptx_content)
+    arch = match.group(1) if match else "sm_90a"
+
+    # Write stripped content back if needed
+    if ptx_path.read_text() != ptx_content:
+        ptx_path.write_text(ptx_content)
+
+    # Compile
+    cubin_tmp = ptx_path.with_suffix(".cubin.tmp")
+    try:
+        assert CUTE_DSL_PTXAS_PATH is not None
+        result = subprocess.run(
+            [
+                CUTE_DSL_PTXAS_PATH,
+                f"-arch={arch}",
+                "-O3",
+                "-o",
+                str(cubin_tmp),
+                str(ptx_path),
+            ],
+            capture_output=True,
+            text=True,
+        )
+        if result.returncode != 0:
+            raise RuntimeError(f"ptxas failed: {result.stderr}")
+
+        cubin_data = cubin_tmp.read_bytes()
+        _log(f"Compiled {ptx_path.name} -> {len(cubin_data)} bytes ({arch})")
+
+        # Save cubin if CUTE_DSL_KEEP_CUBIN is set
+        if os.environ.get("CUTE_DSL_KEEP_CUBIN", "0") == "1":
+            cubin_out = ptx_path.with_suffix(".cubin")
+            cubin_out.write_bytes(cubin_data)
+            _log(f"Saved: {cubin_out}")
+
+        return cubin_data
+    finally:
+        cubin_tmp.unlink(missing_ok=True)
+
+
+def _patched_load_cuda_library(self):
+    """Replacement for _load_cuda_library that uses system ptxas."""
+
+    result = _get_ptx(self)
+    if not result:
+        _log("PTX not found, falling back to embedded ptxas")
+        return _original_load_cuda_library(self)
+
+    ptx_content, ptx_path = result
+
+    try:
+        cubin = _compile_ptx(ptx_path, ptx_content)
+    except Exception as e:
+        _log(f"Compilation failed ({e}), falling back to embedded ptxas")
+        return _original_load_cuda_library(self)
+
+    # Load cubin
+    import cuda.bindings.runtime as cuda_runtime
+
+    err, library = cuda_runtime.cudaLibraryLoadData(cubin, None, None, 0, None, None, 0)
+    if err != cuda_runtime.cudaError_t.cudaSuccess:
+        _log(f"cudaLibraryLoadData failed ({err}), falling back to embedded ptxas")
+        return _original_load_cuda_library(self)
+
+    # Register kernels on all devices
+    _, cuda_load_to_device = self._get_cuda_init_and_load()
+    lib_ptr = ctypes.c_void_p(int(library))
+    dev_id = ctypes.c_int32(0)
+    err_val = ctypes.c_int32(0)
+    args = (ctypes.c_void_p * 3)(
+        ctypes.cast(ctypes.pointer(lib_ptr), ctypes.c_void_p),
+        ctypes.cast(ctypes.pointer(dev_id), ctypes.c_void_p),
+        ctypes.cast(ctypes.pointer(err_val), ctypes.c_void_p),
+    )
+
+    for dev in range(self.num_devices):
+        dev_id.value = dev
+        cuda_load_to_device(args)
+        if err_val.value != 0:
+            _log("cuda_load_to_device failed, falling back to embedded ptxas")
+            return _original_load_cuda_library(self)
+
+    _log(f"Loaded kernel from {ptx_path.name}")
+
+    # Delete PTX if user didn't originally want it kept
+    if not _user_wanted_ptx:
+        ptx_path.unlink(missing_ok=True)
+
+    return [cuda_runtime.cudaLibrary_t(lib_ptr.value)]
+
+
+def patch():
+    """Install system ptxas hook. Call before importing cutlass."""
+    global _original_load_cuda_library, _user_wanted_ptx
+
+    assert CUTE_DSL_PTXAS_PATH is not None
+    if not os.path.isfile(CUTE_DSL_PTXAS_PATH) or not os.access(
+        CUTE_DSL_PTXAS_PATH, os.X_OK
+    ):
+        raise RuntimeError(f"ptxas not found: {CUTE_DSL_PTXAS_PATH}")
+
+    # Track if user originally wanted PTX kept
+    _user_wanted_ptx = os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1"
+    # os.environ['CUTE_DSL_KEEP_PTX'] = '1'
+    assert (
+        os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1"
+    ), "Require CUTE_DSL_KEEP_PTX=1 to use system's ptxas"
+
+    cls = cutlass.cutlass_dsl.cuda_jit_executor.CudaDialectJitCompiledFunction
+    _original_load_cuda_library = cls._load_cuda_library
+    cls._load_cuda_library = _patched_load_cuda_library
+    _log("Patch applied")
+    return

python/sglang/jit_kernel/flash_attn/cute/cute_dsl_utils.py

@@ -0,0 +1,176 @@
+# Copyright (c) 2025, Tri Dao.
+
+from functools import lru_cache
+from typing import Tuple
+
+import torch
+
+try:
+    from triton.tools.disasm import extract
+except ImportError:
+    extract = None
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import NumericMeta
+
+StaticTypes = (cutlass.Constexpr, NumericMeta, int, bool, str, float, type(None))
+
+
+load_cubin_module_data_og = cutlass.base_dsl.runtime.cuda.load_cubin_module_data
+cute_compile_og = cute.compile
+
+
+torch2cute_dtype_map = {
+    torch.float16: cutlass.Float16,
+    torch.bfloat16: cutlass.BFloat16,
+    torch.float32: cutlass.Float32,
+    torch.float8_e4m3fn: cutlass.Float8E4M3FN,
+    torch.float8_e5m2: cutlass.Float8E5M2,
+}
+
+
+@lru_cache
+def get_max_active_clusters(cluster_size):
+    return cutlass.utils.HardwareInfo().get_max_active_clusters(
+        cluster_size=cluster_size
+    )
+
+
+@lru_cache
+def get_device_capacity(device: torch.device = None) -> Tuple[int, int]:
+    return torch.cuda.get_device_capability(device)
+
+
+def assume_strides_aligned(t, align=16):
+    """Assume all strides except the last are divisible by `align` bytes (128
+    bits by default; 4 bytes for the packed UE8M0 scale-factor tensors).
+
+    Python int strides (e.g., stride=0 from GQA expand) are kept as-is
+    since they're static and don't need alignment assumptions.
+    """
+    divby = (align * 8) // t.element_type.width
+    strides = tuple(
+        s if isinstance(s, int) else cute.assume(s, divby=divby) for s in t.stride[:-1]
+    )
+    return (*strides, t.stride[-1])
+
+
+def assume_tensor_aligned(t, align=16):
+    """Rebuild a tensor with aligned stride assumptions. Passes through None."""
+    if t is None:
+        return None
+    return cute.make_tensor(
+        t.iterator,
+        cute.make_layout(t.shape, stride=assume_strides_aligned(t, align=align)),
+    )
+
+
+def to_cute_tensor(
+    t, assumed_align=16, leading_dim=-1, fully_dynamic=False, enable_tvm_ffi=True
+):
+    """Convert torch tensor to cute tensor for TVM FFI. leading_dim=-1 defaults to t.ndim-1."""
+    if t is None:
+        return None
+    # NOTE: torch 2.9.1 doesn't support fp8 via DLPack but 2.11.0 nightly does
+    # currently export raw bytes as uint8 and tell cutlass correct type
+    # can directly export as fp8 when torch supports it
+    if t.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+        tensor = from_dlpack(
+            t.view(torch.uint8).detach(),
+            assumed_align=assumed_align,
+            enable_tvm_ffi=enable_tvm_ffi,
+        )
+        tensor.element_type = (
+            cutlass.Float8E4M3FN
+            if t.dtype == torch.float8_e4m3fn
+            else cutlass.Float8E5M2
+        )
+    else:
+        tensor = from_dlpack(
+            t.detach(), assumed_align=assumed_align, enable_tvm_ffi=enable_tvm_ffi
+        )
+    if fully_dynamic:
+        return tensor.mark_layout_dynamic()
+    if leading_dim == -1:
+        leading_dim = t.ndim - 1
+    return tensor.mark_layout_dynamic(leading_dim=leading_dim)
+
+
+def to_cute_aux_tensor(t, enable_tvm_ffi=True):
+    """Convert torch tensor to cute tensor for TVM FFI, tailored to FlexAttention aux tensors.
+    This allows the user to specify alignment and leading dimension for aux tensors used in
+    custom score_mod callables.
+    """
+    assumed_align: int = getattr(t, "__assumed_align__", None)
+    leading_dim: int = getattr(t, "__leading_dim__", None)
+    fully_dynamic: bool = leading_dim is None
+
+    return to_cute_tensor(
+        t,
+        assumed_align=assumed_align,
+        leading_dim=leading_dim,
+        fully_dynamic=fully_dynamic,
+        enable_tvm_ffi=enable_tvm_ffi,
+    )
+
+
+def get_aux_tensor_metadata(aux_tensors):
+    return tuple(
+        (
+            getattr(t, "__assumed_align__", 0),
+            getattr(t, "__leading_dim__", -1),
+            hasattr(t, "__leading_dim__"),
+        )
+        for t in aux_tensors
+    )
+
+
+def get_broadcast_dims(tensor: torch.Tensor) -> Tuple[bool, ...]:
+    """Return tuple of bools indicating which dims have stride=0 (broadcast).
+
+    This is useful for compile keys since CuTe's mark_layout_dynamic() keeps
+    stride=0 as static, meaning kernels compiled with different broadcast
+    patterns are not interchangeable.
+    """
+    return tuple(s == 0 for s in tensor.stride())
+
+
+# credit: monellz (https://github.com/NVIDIA/cutlass/issues/2658#issuecomment-3630564264)
+def dump_kernel_attributes(compiled_kernel):
+    import torch
+    from cuda.bindings import driver
+    from cutlass.utils import HardwareInfo
+
+    device_id = torch.cuda.current_device()
+    hardware_info = HardwareInfo(device_id=device_id)
+    cubin_data = compiled_kernel.artifacts.CUBIN
+    assert (
+        cubin_data is not None
+    ), "cubin_data is None, need '--keep-cubin' option when compiling"
+    cuda_library = hardware_info._checkCudaErrors(
+        driver.cuLibraryLoadData(cubin_data, None, None, 0, None, None, 0)
+    )
+    kernels = hardware_info._checkCudaErrors(
+        driver.cuLibraryEnumerateKernels(1, cuda_library)
+    )
+    kernel = hardware_info._checkCudaErrors(driver.cuKernelGetFunction(kernels[0]))
+    # more metrics: https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__EXEC.html#group__CUDA__EXEC_1g5e92a1b0d8d1b82cb00dcfb2de15961b
+    local_size_bytes = hardware_info._checkCudaErrors(
+        driver.cuFuncGetAttribute(
+            driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES,
+            kernel,
+        )
+    )
+    num_regs = hardware_info._checkCudaErrors(
+        driver.cuFuncGetAttribute(
+            driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS,
+            kernel,
+        )
+    )
+
+    print("--- Kernel Info ---")
+    print(f"local_size_bytes: {local_size_bytes}")
+    print(f"num_regs: {num_regs}")
+    print("--- End Kernel Info ---")

python/sglang/jit_kernel/flash_attn/cute/fa_logging.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2025, Tri Dao.
+
+"""Unified FlashAttention logging controlled by a single ``FA_LOG_LEVEL`` env var.
+
+Host-side messages go through Python ``logging`` (logger name ``flash_attn``).
+A default ``StreamHandler`` is attached automatically when ``FA_LOG_LEVEL >= 1``
+so that standalone scripts get output without extra setup; applications that
+configure their own logging can remove or replace it via the standard API.
+
+FA_LOG_LEVEL mapping::
+
+    0  off       nothing logged
+    1  host      host-side summaries only (no kernel printf)
+    2  kernel    host + curated kernel traces
+    3  max       host + all kernel traces (noisy, perf hit)
+
+Set via environment variable::
+
+    FA_LOG_LEVEL=1 python train.py
+
+Device-side ``cute.printf`` calls are compile-time eliminated via
+``cutlass.const_expr`` when the log level is below the callsite threshold,
+so there is zero performance cost when device logging is off.
+Changing the log level after kernel compilation requires a recompile
+(the level participates in the forward compile key).
+"""
+
+import logging
+import os
+import sys
+
+import cutlass.cute as cute
+from cutlass import const_expr
+
+_LOG_LEVEL_NAMES = {"off": 0, "host": 1, "kernel": 2, "max": 3}
+
+
+def _parse_log_level(raw: str) -> int:
+    if raw in _LOG_LEVEL_NAMES:
+        return _LOG_LEVEL_NAMES[raw]
+    try:
+        level = int(raw)
+    except ValueError:
+        return 0
+    return max(0, min(level, 3))
+
+
+_fa_log_level: int = _parse_log_level(os.environ.get("FA_LOG_LEVEL", "0"))
+
+_logger = logging.getLogger("flash_attn")
+_logger.addHandler(logging.NullHandler())
+_default_handler: logging.Handler | None = None
+
+
+def _configure_default_handler() -> None:
+    global _default_handler
+    if _fa_log_level >= 1:
+        if _default_handler is None:
+            _default_handler = logging.StreamHandler(sys.stdout)
+            _default_handler.setFormatter(logging.Formatter("[FA] %(message)s"))
+            _logger.addHandler(_default_handler)
+        _logger.setLevel(logging.DEBUG)
+    else:
+        if _default_handler is not None:
+            _logger.removeHandler(_default_handler)
+            _default_handler = None
+        _logger.setLevel(logging.WARNING)
+
+
+_configure_default_handler()
+
+
+def get_fa_log_level() -> int:
+    return _fa_log_level
+
+
+def set_fa_log_level(level: int | str) -> None:
+    """Set the FA log level programmatically.
+
+    Host logging takes effect immediately.  Device logging changes only
+    affect kernels compiled after this call (new compile-key selection).
+    """
+    global _fa_log_level
+    if isinstance(level, str):
+        level = _parse_log_level(level)
+    _fa_log_level = max(0, min(int(level), 3))
+    _configure_default_handler()
+
+
+def fa_log(level: int, msg: str):
+    if _fa_log_level >= level:
+        _logger.info(msg)
+
+
+def fa_printf(level: int, fmt, *args):
+    if const_expr(_fa_log_level >= level):
+        cute.printf(fmt, *args)

python/sglang/jit_kernel/flash_attn/cute/fast_math.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2025, Tri Dao.
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32
+
+
+@cute.jit
+def clz(x: Int32) -> Int32:
+    # for i in cutlass.range_constexpr(32):
+    #     if (1 << (31 - i)) & x:
+    #         return Int32(i)
+    # return Int32(32)
+    # Early exit is not supported yet
+    res = Int32(32)
+    done = False
+    for i in cutlass.range(32):
+        if ((1 << (31 - i)) & x) and not done:
+            res = Int32(i)
+            done = True
+    return res

python/sglang/jit_kernel/flash_attn/cute/flash_fwd_combine.py

@@ -0,0 +1,767 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+# A reimplementation of https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_fwd_combine_kernel.h
+# from Cutlass C++ to Cute-DSL.
+import math
+from functools import partial
+from typing import Optional, Type
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+from cutlass import Boolean, Float32, Int32, const_expr
+from cutlass.cute import FastDivmodDivisor
+from cutlass.cute.nvgpu import cpasync
+
+from sglang.jit_kernel.flash_attn.cute import utils
+from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
+from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfo
+
+
+class FlashAttentionForwardCombine:
+    def __init__(
+        self,
+        dtype: Type[cutlass.Numeric],
+        dtype_partial: Type[cutlass.Numeric],
+        head_dim: int,
+        tile_m: int = 8,
+        k_block_size: int = 64,
+        log_max_splits: int = 4,
+        num_threads: int = 256,
+        stages: int = 4,
+        use_pdl: bool = False,
+    ):
+        """
+        Forward combine kernel for split attention computation.
+
+        :param dtype: output data type
+        :param dtype_partial: partial accumulation data type
+        :param head_dim: head dimension
+        :param tile_m: m block size
+        :param k_block_size: k block size
+        :param log_max_splits: log2 of maximum splits
+        :param num_threads: number of threads
+        :param varlen: whether using variable length sequences
+        :param stages: number of pipeline stages
+        """
+        self.dtype = dtype
+        self.dtype_partial = dtype_partial
+        self.head_dim = head_dim
+        self.tile_m = tile_m
+        self.k_block_size = k_block_size
+        self.max_splits = 1 << log_max_splits
+        self.num_threads = num_threads
+        self.is_even_k = head_dim % k_block_size == 0
+        self.stages = stages
+        self.use_pdl = use_pdl
+
+    @staticmethod
+    def can_implement(
+        dtype,
+        dtype_partial,
+        head_dim,
+        tile_m,
+        k_block_size,
+        log_max_splits,
+        num_threads,
+    ) -> bool:
+        """Check if the kernel can be implemented with the given parameters."""
+        if dtype not in [cutlass.Float16, cutlass.BFloat16, cutlass.Float32]:
+            return False
+        if dtype_partial not in [cutlass.Float16, cutlass.BFloat16, Float32]:
+            return False
+        if head_dim % 8 != 0:
+            return False
+        if num_threads % 32 != 0:
+            return False
+        if tile_m % 8 != 0:
+            return False
+        max_splits = 1 << log_max_splits
+        if max_splits > 256:
+            return False
+        if (tile_m * max_splits) % num_threads != 0:
+            return False
+        return True
+
+    def _setup_attributes(self):
+        # GMEM copy setup for O partial
+        universal_copy_bits = 128
+        async_copy_elems = universal_copy_bits // self.dtype_partial.width
+        assert self.k_block_size % async_copy_elems == 0
+
+        k_block_gmem = (
+            128
+            if self.k_block_size % 128 == 0
+            else (64 if self.k_block_size % 64 == 0 else 32)
+        )
+        gmem_threads_per_row = k_block_gmem // async_copy_elems
+        assert self.num_threads % gmem_threads_per_row == 0
+
+        # Async copy atom for O partial load
+        atom_async_copy_partial = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+            self.dtype_partial,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        tOpartial_layout = cute.make_ordered_layout(
+            (self.num_threads // gmem_threads_per_row, gmem_threads_per_row),
+            order=(1, 0),
+        )
+        vOpartial_layout = cute.make_layout((1, async_copy_elems))  # 4 vals per load
+        self.gmem_tiled_copy_O_partial = cute.make_tiled_copy_tv(
+            atom_async_copy_partial, tOpartial_layout, vOpartial_layout
+        )
+
+        # GMEM copy setup for final O (use universal copy for store)
+        atom_universal_copy = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            self.dtype,
+            num_bits_per_copy=async_copy_elems * self.dtype.width,
+        )
+        self.gmem_tiled_copy_O = cute.make_tiled_copy_tv(
+            atom_universal_copy,
+            tOpartial_layout,
+            vOpartial_layout,  # 4 vals per store
+        )
+
+        # LSE copy setup with async copy (alignment = 1)
+        lse_copy_bits = Float32.width  # 1 element per copy, width is in bits
+        m_block_smem = (
+            128
+            if self.tile_m % 128 == 0
+            else (
+                64
+                if self.tile_m % 64 == 0
+                else (
+                    32
+                    if self.tile_m % 32 == 0
+                    else (16 if self.tile_m % 16 == 0 else 8)
+                )
+            )
+        )
+        gmem_threads_per_row_lse = m_block_smem
+        assert self.num_threads % gmem_threads_per_row_lse == 0
+
+        # Async copy atom for LSE load
+        atom_async_copy_lse = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS),
+            Float32,
+            num_bits_per_copy=lse_copy_bits,
+        )
+        tLSE_layout = cute.make_ordered_layout(
+            (self.num_threads // gmem_threads_per_row_lse, gmem_threads_per_row_lse),
+            order=(1, 0),
+        )
+        vLSE_layout = cute.make_layout(1)
+        self.gmem_tiled_copy_LSE = cute.make_tiled_copy_tv(
+            atom_async_copy_lse, tLSE_layout, vLSE_layout
+        )
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Shared memory
+        # ///////////////////////////////////////////////////////////////////////////////
+
+        # Shared memory to register copy for LSE
+        self.smem_threads_per_col_lse = self.num_threads // m_block_smem
+        assert 32 % self.smem_threads_per_col_lse == 0  # Must divide warp size
+
+        s2r_layout_atom_lse = cute.make_ordered_layout(
+            (
+                self.smem_threads_per_col_lse,
+                self.num_threads // self.smem_threads_per_col_lse,
+            ),
+            order=(0, 1),
+        )
+        self.s2r_tiled_copy_LSE = cute.make_tiled_copy_tv(
+            cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32),
+            s2r_layout_atom_lse,
+            cute.make_layout(1),
+        )
+
+        # LSE shared memory layout with swizzling to avoid bank conflicts
+        # This works for kBlockMSmem = 8, 16, 32, 64, 128, no bank conflicts
+        if const_expr(m_block_smem == 8):
+            smem_lse_swizzle = cute.make_swizzle(5, 0, 5)
+        elif const_expr(m_block_smem == 16):
+            smem_lse_swizzle = cute.make_swizzle(4, 0, 4)
+        else:
+            smem_lse_swizzle = cute.make_swizzle(3, 2, 3)
+        smem_layout_atom_lse = cute.make_composed_layout(
+            smem_lse_swizzle,
+            0,
+            cute.make_ordered_layout((8, m_block_smem), order=(1, 0)),
+        )
+        self.smem_layout_lse = cute.tile_to_shape(
+            smem_layout_atom_lse, (self.max_splits, self.tile_m), (0, 1)
+        )
+
+        # O partial shared memory layout (simple layout for pipeline stages)
+        self.smem_layout_o = cute.make_ordered_layout(
+            (self.tile_m, self.k_block_size, self.stages), order=(1, 0, 2)
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        mO_partial: cute.Tensor,
+        mLSE_partial: cute.Tensor,
+        mO: cute.Tensor,
+        mLSE: Optional[cute.Tensor] = None,
+        cu_seqlens: Optional[cute.Tensor] = None,
+        seqused: Optional[cute.Tensor] = None,
+        num_splits_dynamic_ptr: Optional[cute.Tensor] = None,
+        varlen_batch_idx: Optional[cute.Tensor] = None,
+        semaphore_to_reset: Optional[cute.Tensor] = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
+    ):
+        # Type checking
+        if const_expr(not (mO_partial.element_type == self.dtype_partial)):
+            raise TypeError("O partial tensor must match dtype_partial")
+        if const_expr(not (mO.element_type == self.dtype)):
+            raise TypeError("O tensor must match dtype")
+        if const_expr(mLSE_partial.element_type not in [Float32]):
+            raise TypeError("LSE partial tensor must be Float32")
+        if const_expr(mLSE is not None and mLSE.element_type not in [Float32]):
+            raise TypeError("LSE tensor must be Float32")
+
+        # Shape validation - input tensors are in user format, need to be converted to kernel format
+        if const_expr(len(mO_partial.shape) not in [4, 5]):
+            raise ValueError(
+                "O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)"
+            )
+        if const_expr(len(mLSE_partial.shape) not in [3, 4]):
+            raise ValueError(
+                "LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)"
+            )
+        if const_expr(len(mO.shape) not in [3, 4]):
+            raise ValueError(
+                "O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim)"
+            )
+        if const_expr(mLSE is not None and len(mLSE.shape) not in [2, 3]):
+            raise ValueError(
+                "LSE tensor must have 2 or 3 dimensions: (batch, seqlen, nheads) or (total_q, nheads)"
+            )
+
+        mO_partial, mO = [assume_tensor_aligned(t) for t in (mO_partial, mO)]
+        # (num_splits, b, seqlen, h, d) -> (seqlen, d, num_splits, h, b)
+        # or (num_splits, total_q, h, d) -> (total_q, d, num_splits, h)
+        O_partial_layout_transpose = (
+            [2, 4, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 3, 0, 2]
+        )
+        # (b, seqlen, h, d) -> (seqlen, d, h, b) or (total_q, h, d) -> (total_q, d, h)
+        mO_partial = cute.make_tensor(
+            mO_partial.iterator,
+            cute.select(mO_partial.layout, mode=O_partial_layout_transpose),
+        )
+        O_layout_transpose = (
+            [1, 3, 2, 0] if const_expr(cu_seqlens is None) else [0, 2, 1]
+        )
+        mO = cute.make_tensor(
+            mO.iterator, cute.select(mO.layout, mode=O_layout_transpose)
+        )
+        # (num_splits, b, seqlen, h) -> (seqlen, num_splits, h, b)
+        # or (num_splits, total_q, h) -> (total_q, num_splits, h)
+        LSE_partial_layout_transpose = (
+            [2, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 0, 2]
+        )
+        mLSE_partial = cute.make_tensor(
+            mLSE_partial.iterator,
+            cute.select(mLSE_partial.layout, mode=LSE_partial_layout_transpose),
+        )
+        # (b, seqlen, h) -> (seqlen, h, b) or (total_q, h) -> (total_q, h)
+        LSE_layout_transpose = [1, 2, 0] if const_expr(cu_seqlens is None) else [0, 1]
+        mLSE = (
+            cute.make_tensor(
+                mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose)
+            )
+            if mLSE is not None
+            else None
+        )
+
+        # Determine if we have variable length sequences
+        varlen = const_expr(cu_seqlens is not None or seqused is not None)
+
+        self._setup_attributes()
+
+        @cute.struct
+        class SharedStorage:
+            sLSE: cute.struct.Align[
+                cute.struct.MemRange[Float32, cute.cosize(self.smem_layout_lse)], 128
+            ]
+            sMaxValidSplit: cute.struct.Align[
+                cute.struct.MemRange[Int32, self.tile_m], 128
+            ]
+            sO: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.dtype_partial, cute.cosize(self.smem_layout_o)
+                ],
+                128,
+            ]
+
+        smem_size = SharedStorage.size_in_bytes()
+
+        # Grid dimensions: (ceil_div(seqlen, m_block), ceil_div(head_dim, k_block), num_head * batch)
+        seqlen = mO_partial.shape[0]
+        num_head = mO_partial.shape[3]
+        batch_size = (
+            mO_partial.shape[4]
+            if const_expr(cu_seqlens is None)
+            else Int32(cu_seqlens.shape[0] - 1)
+        )
+
+        # Create FastDivmodDivisor objects for efficient division
+        seqlen_divmod = FastDivmodDivisor(seqlen)
+        head_divmod = FastDivmodDivisor(num_head)
+
+        grid_dim = (
+            cute.ceil_div(seqlen * num_head, self.tile_m),
+            cute.ceil_div(self.head_dim, self.k_block_size),
+            batch_size,
+        )
+
+        self.kernel(
+            mO_partial,
+            mLSE_partial,
+            mO,
+            mLSE,
+            cu_seqlens,
+            seqused,
+            num_splits_dynamic_ptr,
+            varlen_batch_idx,
+            semaphore_to_reset,
+            SharedStorage,
+            self.smem_layout_lse,
+            self.smem_layout_o,
+            self.gmem_tiled_copy_O_partial,
+            self.gmem_tiled_copy_O,
+            self.gmem_tiled_copy_LSE,
+            self.s2r_tiled_copy_LSE,
+            seqlen_divmod,
+            head_divmod,
+            varlen,
+        ).launch(
+            grid=grid_dim,
+            block=[self.num_threads, 1, 1],
+            smem=smem_size,
+            stream=stream,
+            use_pdl=self.use_pdl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mO_partial: cute.Tensor,
+        mLSE_partial: cute.Tensor,
+        mO: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        cu_seqlens: Optional[cute.Tensor],
+        seqused: Optional[cute.Tensor],
+        num_splits_dynamic_ptr: Optional[cute.Tensor],
+        varlen_batch_idx: Optional[cute.Tensor],
+        semaphore_to_reset: Optional[cute.Tensor],
+        SharedStorage: cutlass.Constexpr,
+        smem_layout_lse: cute.Layout | cute.ComposedLayout,
+        smem_layout_o: cute.Layout,
+        gmem_tiled_copy_O_partial: cute.TiledCopy,
+        gmem_tiled_copy_O: cute.TiledCopy,
+        gmem_tiled_copy_LSE: cute.TiledCopy,
+        s2r_tiled_copy_LSE: cute.TiledCopy,
+        seqlen_divmod: FastDivmodDivisor,
+        head_divmod: FastDivmodDivisor,
+        varlen: cutlass.Constexpr[bool],
+    ):
+        # Thread and block indices
+        tidx, _, _ = cute.arch.thread_idx()
+        m_block, k_block, maybe_virtual_batch = cute.arch.block_idx()
+
+        # Map virtual batch index to real batch index (for persistent tile schedulers)
+        batch_idx = (
+            varlen_batch_idx[maybe_virtual_batch]
+            if const_expr(varlen_batch_idx is not None)
+            else maybe_virtual_batch
+        )
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Get shared memory buffer
+        # ///////////////////////////////////////////////////////////////////////////////
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+        sLSE = storage.sLSE.get_tensor(smem_layout_lse)
+        sMaxValidSplit = storage.sMaxValidSplit.get_tensor((self.tile_m,))
+        sO = storage.sO.get_tensor(smem_layout_o)
+
+        # Handle semaphore reset — wait for dependent grids first
+        if const_expr(semaphore_to_reset is not None):
+            if (
+                tidx == 0
+                and m_block == cute.arch.grid_dim()[0] - 1
+                and k_block == cute.arch.grid_dim()[1] - 1
+                and maybe_virtual_batch == cute.arch.grid_dim()[2] - 1
+            ):
+                if const_expr(self.use_pdl):
+                    cute.arch.griddepcontrol_wait()
+                semaphore_to_reset[0] = 0
+
+        # Get number of splits (use maybe_virtual_batch for per-batch-slot splits)
+        num_splits = (
+            num_splits_dynamic_ptr[maybe_virtual_batch]
+            if const_expr(num_splits_dynamic_ptr is not None)
+            else mLSE_partial.shape[1]
+        )
+        # Handle variable length sequences using SeqlenInfo
+        seqlen_info = SeqlenInfo.create(
+            batch_idx=batch_idx,
+            seqlen_static=mO_partial.shape[0],
+            cu_seqlens=cu_seqlens,
+            seqused=seqused,
+            # Don't need to pass in tile size since we won't use offset_padded
+        )
+        seqlen, offset = seqlen_info.seqlen, seqlen_info.offset
+
+        # Extract number of heads (head index will be determined dynamically)
+        num_head = mO_partial.shape[3]
+        max_idx = seqlen * num_head
+
+        # Early exit for single split if dynamic
+        if (const_expr(num_splits_dynamic_ptr is None) or num_splits > 1) and (
+            const_expr(not varlen) or m_block * self.tile_m < max_idx
+        ):
+            # Wait for dependent grids (e.g., the main attention kernel that produces O_partial/LSE_partial)
+            if const_expr(self.use_pdl):
+                cute.arch.griddepcontrol_wait()
+
+            # ===============================
+            # Step 1: Load LSE_partial from gmem to shared memory
+            # ===============================
+
+            mLSE_partial_cur = seqlen_info.offset_batch(mLSE_partial, batch_idx, dim=3)
+            mLSE_partial_copy = cute.tiled_divide(mLSE_partial_cur, (1,))
+            gmem_thr_copy_LSE = gmem_tiled_copy_LSE.get_slice(tidx)
+            tLSEsLSE = gmem_thr_copy_LSE.partition_D(sLSE)
+            # Create identity tensor for coordinate tracking
+            cLSE = cute.make_identity_tensor((self.max_splits, self.tile_m))
+            tLSEcLSE = gmem_thr_copy_LSE.partition_S(cLSE)
+
+            # Load LSE partial values
+            for m in cutlass.range(cute.size(tLSEcLSE, mode=[2]), unroll_full=True):
+                mi = tLSEcLSE[0, 0, m][1]  # Get m coordinate
+                idx = m_block * self.tile_m + mi
+                if idx < max_idx:
+                    # Calculate actual sequence position and head using FastDivmodDivisor
+                    if const_expr(not varlen):
+                        head_idx, m_idx = divmod(idx, seqlen_divmod)
+                    else:
+                        head_idx = idx // seqlen
+                        m_idx = idx - head_idx * seqlen
+                    mLSE_partial_cur_copy = mLSE_partial_copy[
+                        None, m_idx, None, head_idx
+                    ]
+                    for s in cutlass.range(
+                        cute.size(tLSEcLSE, mode=[1]), unroll_full=True
+                    ):
+                        si = tLSEcLSE[0, s, 0][0]  # Get split coordinate
+                        if si < num_splits:
+                            cute.copy(
+                                gmem_thr_copy_LSE,
+                                mLSE_partial_cur_copy[None, si],
+                                tLSEsLSE[None, s, m],
+                            )
+                        else:
+                            tLSEsLSE[None, s, m].fill(-Float32.inf)
+                # Don't need to zero out the rest of the LSEs, as we will not write the output to gmem
+            cute.arch.cp_async_commit_group()
+
+            # ===============================
+            # Step 2: Load O_partial for pipeline stages
+            # ===============================
+
+            gmem_thr_copy_O_partial = gmem_tiled_copy_O_partial.get_slice(tidx)
+            cO = cute.make_identity_tensor((self.tile_m, self.k_block_size))
+            tOcO = gmem_thr_copy_O_partial.partition_D(cO)
+            tOsO_partial = gmem_thr_copy_O_partial.partition_D(sO)
+            mO_partial_cur = seqlen_info.offset_batch(mO_partial, batch_idx, dim=4)
+
+            # Precompute these values to avoid recomputing them in the loop
+            num_rows = const_expr(cute.size(tOcO, mode=[1]))
+            tOmidx = cute.make_rmem_tensor(num_rows, cutlass.Int32)
+            tOhidx = cute.make_rmem_tensor(num_rows, cutlass.Int32)
+            tOrOptr = cute.make_rmem_tensor(num_rows, cutlass.Int64)
+            for m in cutlass.range(num_rows, unroll_full=True):
+                mi = tOcO[0, m, 0][0]  # m coordinate
+                idx = m_block * self.tile_m + mi
+                if const_expr(not varlen):
+                    tOhidx[m], tOmidx[m] = divmod(idx, seqlen_divmod)
+                else:
+                    tOhidx[m] = idx // seqlen
+                    tOmidx[m] = idx - tOhidx[m] * seqlen
+                tOrOptr[m] = utils.elem_pointer(
+                    mO_partial_cur,
+                    (tOmidx[m], k_block * self.k_block_size, 0, tOhidx[m]),
+                ).toint()
+                if idx >= max_idx:
+                    tOhidx[m] = -1
+
+            tOpO = None
+            if const_expr(not self.is_even_k):
+                tOpO = cute.make_rmem_tensor(cute.size(tOcO, mode=[2]), Boolean)
+                for k in cutlass.range(cute.size(tOpO), unroll_full=True):
+                    tOpO[k] = (
+                        tOcO[0, 0, k][1]
+                        < mO_partial.shape[1] - k_block * self.k_block_size
+                    )
+                # if cute.arch.thread_idx()[0] == 0 and k_block == 1: cute.print_tensor(tOpO)
+
+            load_O_partial = partial(
+                self.load_O_partial,
+                gmem_tiled_copy_O_partial,
+                tOrOptr,
+                tOsO_partial,
+                tOhidx,
+                tOpO,
+                tOcO,
+                mO_partial_cur.layout,
+            )
+
+            # Load first few stages of O_partial
+            for stage in cutlass.range(self.stages - 1, unroll_full=True):
+                if stage < num_splits:
+                    load_O_partial(stage, stage)
+                cute.arch.cp_async_commit_group()
+
+            # ===============================
+            # Step 3: Load and transpose LSE from smem to registers
+            # ===============================
+
+            # Wait for LSE and initial O partial stages to complete
+            cute.arch.cp_async_wait_group(self.stages - 1)
+            cute.arch.sync_threads()
+            # if cute.arch.thread_idx()[0] == 0:
+            #     # cute.print_tensor(sLSE)
+            #     for i in range(64):
+            #         cute.printf("sLSE[%d, 0] = %f", i, sLSE[i, 0])
+            # cute.arch.sync_threads()
+
+            s2r_thr_copy_LSE = s2r_tiled_copy_LSE.get_slice(tidx)
+            ts2rsLSE = s2r_thr_copy_LSE.partition_S(sLSE)
+            ts2rrLSE = cute.make_rmem_tensor_like(ts2rsLSE)
+            cute.copy(s2r_tiled_copy_LSE, ts2rsLSE, ts2rrLSE)
+
+            # ===============================
+            # Step 4: Compute final LSE along split dimension
+            # ===============================
+
+            lse_sum = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Float32)
+            ts2rcLSE = s2r_thr_copy_LSE.partition_D(cLSE)
+            # We compute the max valid split for each row to short-circuit the computation later
+            max_valid_split = cute.make_rmem_tensor(
+                cute.size(ts2rrLSE, mode=[2]), Int32
+            )
+            assert cute.size(ts2rrLSE, mode=[0]) == 1
+            # Compute max, scales, and final LSE for each row
+            for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
+                # Find max LSE value across splits
+                threads_per_col = const_expr(self.smem_threads_per_col_lse)
+                lse_max = cute.arch.warp_reduction_max(
+                    ts2rrLSE[None, None, m]
+                    .load()
+                    .reduce(
+                        cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0
+                    ),
+                    threads_in_group=threads_per_col,
+                )
+                # if cute.arch.thread_idx()[0] == 0: cute.printf(lse_max)
+                # Find max valid split index
+                max_valid_idx = -1
+                for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True):
+                    if ts2rrLSE[0, s, m] != -Float32.inf:
+                        max_valid_idx = ts2rcLSE[0, s, 0][0]  # Get split coordinate
+                # if cute.arch.thread_idx()[0] < 32: cute.printf(max_valid_idx)
+                max_valid_split[m] = cute.arch.warp_reduction_max(
+                    max_valid_idx, threads_in_group=threads_per_col
+                )
+                # Compute exp scales and sum
+                lse_max_cur = (
+                    0.0 if lse_max == -Float32.inf else lse_max
+                )  # In case all local LSEs are -inf
+                LOG2_E = math.log2(math.e)
+                lse_sum_cur = 0.0
+                for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True):
+                    scale = cute.math.exp2(
+                        ts2rrLSE[0, s, m] * LOG2_E - (lse_max_cur * LOG2_E),
+                        fastmath=True,
+                    )
+                    lse_sum_cur += scale
+                    ts2rrLSE[0, s, m] = scale  # Store scale for later use
+                lse_sum_cur = cute.arch.warp_reduction_sum(
+                    lse_sum_cur, threads_in_group=threads_per_col
+                )
+                lse_sum[m] = cute.math.log(lse_sum_cur, fastmath=True) + lse_max
+                # Normalize scales
+                inv_sum = (
+                    0.0
+                    if (lse_sum_cur == 0.0 or lse_sum_cur != lse_sum_cur)
+                    else 1.0 / lse_sum_cur
+                )
+                ts2rrLSE[None, None, m].store(ts2rrLSE[None, None, m].load() * inv_sum)
+            # Store the scales exp(lse - lse_logsum) back to smem
+            cute.copy(s2r_tiled_copy_LSE, ts2rrLSE, ts2rsLSE)
+
+            # Store max valid split to smem
+            for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
+                if ts2rcLSE[0, 0, m][0] == 0:  # Only thread responsible for s=0 writes
+                    mi = ts2rcLSE[0, 0, m][1]
+                    if mi < self.tile_m:
+                        sMaxValidSplit[mi] = max_valid_split[m]
+
+            # ===============================
+            # Step 5: Store final LSE to gmem
+            # ===============================
+
+            if const_expr(mLSE is not None):
+                if const_expr(cu_seqlens is None):
+                    mLSE_cur = mLSE[None, None, batch_idx]
+                else:
+                    mLSE_cur = cute.domain_offset((offset, 0), mLSE)
+                if k_block == 0:  # Only first k_block writes LSE when mLSE is provided
+                    for m in cutlass.range(
+                        cute.size(ts2rrLSE, mode=[2]), unroll_full=True
+                    ):
+                        if (
+                            ts2rcLSE[0, 0, m][0] == 0
+                        ):  # Only thread responsible for s=0 writes
+                            mi = ts2rcLSE[0, 0, m][1]
+                            idx = m_block * self.tile_m + mi
+                            if idx < max_idx:
+                                if const_expr(not varlen):
+                                    head_idx, m_idx = divmod(idx, seqlen_divmod)
+                                else:
+                                    head_idx = idx // seqlen
+                                    m_idx = idx - head_idx * seqlen
+                                mLSE_cur[m_idx, head_idx] = lse_sum[m]
+
+            # ===============================
+            # Step 6: Read O_partial and accumulate final O
+            # ===============================
+
+            cute.arch.sync_threads()
+
+            # Get max valid split for this thread
+            thr_max_valid_split = sMaxValidSplit[tOcO[0, 0, 0][0]]
+            for m in cutlass.range(1, cute.size(tOcO, mode=[1]), unroll_full=True):
+                thr_max_valid_split = max(
+                    thr_max_valid_split, sMaxValidSplit[tOcO[0, m, 0][0]]
+                )
+
+            tOrO_partial = cute.make_rmem_tensor_like(tOsO_partial[None, None, None, 0])
+            tOrO = cute.make_rmem_tensor_like(tOrO_partial, Float32)
+            tOrO.fill(0.0)
+
+            stage_load = self.stages - 1
+            stage_compute = 0
+
+            # Main accumulation loop
+            for s in cutlass.range(thr_max_valid_split + 1, unroll=4):
+                # Get scales for this split
+                scale = cute.make_rmem_tensor(num_rows, Float32)
+                for m in cutlass.range(num_rows, unroll_full=True):
+                    scale[m] = sLSE[s, tOcO[0, m, 0][0]]  # Get scale from smem
+
+                # Load next stage if needed
+                split_to_load = s + self.stages - 1
+                if split_to_load <= thr_max_valid_split:
+                    load_O_partial(split_to_load, stage_load)
+                cute.arch.cp_async_commit_group()
+                stage_load = 0 if stage_load == self.stages - 1 else stage_load + 1
+
+                # Wait for the current stage to be ready
+                cute.arch.cp_async_wait_group(self.stages - 1)
+                # We don't need __syncthreads() because each thread is just reading its own data from smem
+                # Copy from smem to registers
+                cute.autovec_copy(
+                    tOsO_partial[None, None, None, stage_compute], tOrO_partial
+                )
+                stage_compute = (
+                    0 if stage_compute == self.stages - 1 else stage_compute + 1
+                )
+
+                # Accumulate scaled partial results
+                for m in cutlass.range(num_rows, unroll_full=True):
+                    if tOhidx[m] >= 0 and scale[m] > 0.0:
+                        tOrO[None, m, None].store(
+                            tOrO[None, m, None].load()
+                            + scale[m] * tOrO_partial[None, m, None].load().to(Float32)
+                        )
+
+            # ===============================
+            # Step 7: Write final O to gmem
+            # ===============================
+
+            rO = cute.make_rmem_tensor_like(tOrO, self.dtype)
+            rO.store(tOrO.load().to(self.dtype))
+            mO_cur = seqlen_info.offset_batch(mO, batch_idx, dim=3)
+            if const_expr(cu_seqlens is None):
+                mO_cur = mO[None, None, None, batch_idx]
+            else:
+                mO_cur = cute.domain_offset((offset, 0, 0), mO)
+            mO_cur = utils.domain_offset_aligned(
+                (0, k_block * self.k_block_size, 0), mO_cur
+            )
+            elems_per_store = const_expr(
+                cute.size(gmem_tiled_copy_O.layout_tv_tiled[1])
+            )
+            # mO_cur_copy = cute.tiled_divide(mO_cur, (1, elems_per_store,))
+            gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx)
+            # Write final results
+            for m in cutlass.range(num_rows, unroll_full=True):
+                if tOhidx[m] >= 0:
+                    mO_cur_copy = cute.tiled_divide(
+                        mO_cur[tOmidx[m], None, tOhidx[m]], (elems_per_store,)
+                    )
+                    for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True):
+                        k_idx = tOcO[0, 0, k][1] // elems_per_store
+                        if const_expr(self.is_even_k) or tOpO[k]:
+                            cute.copy(
+                                gmem_thr_copy_O,
+                                rO[None, m, k],
+                                mO_cur_copy[None, k_idx],
+                            )
+
+    @cute.jit
+    def load_O_partial(
+        self,
+        gmem_tiled_copy_O_partial: cute.TiledCopy,
+        tOrOptr: cute.Tensor,
+        tOsO_partial: cute.Tensor,
+        tOhidx: cute.Tensor,
+        tOpO: Optional[cute.Tensor],
+        tOcO: cute.Tensor,
+        mO_cur_partial_layout: cute.Layout,
+        split: Int32,
+        stage: Int32,
+    ) -> None:
+        elems_per_load = const_expr(
+            cute.size(gmem_tiled_copy_O_partial.layout_tv_tiled[1])
+        )
+        tOsO_partial_cur = tOsO_partial[None, None, None, stage]
+        for m in cutlass.range(cute.size(tOcO, [1]), unroll_full=True):
+            if tOhidx[m] >= 0:
+                o_gmem_ptr = cute.make_ptr(
+                    tOsO_partial.element_type,
+                    tOrOptr[m],
+                    cute.AddressSpace.gmem,
+                    assumed_align=16,
+                )
+                mO_partial_cur = cute.make_tensor(
+                    o_gmem_ptr, cute.slice_(mO_cur_partial_layout, (0, None, None, 0))
+                )
+                mO_partial_cur_copy = cute.tiled_divide(
+                    mO_partial_cur, (elems_per_load,)
+                )
+                for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True):
+                    k_idx = tOcO[0, 0, k][1] // elems_per_load
+                    if const_expr(tOpO is None) or tOpO[k]:
+                        cute.copy(
+                            gmem_tiled_copy_O_partial,
+                            mO_partial_cur_copy[None, k_idx, split],
+                            tOsO_partial_cur[None, m, k],
+                        )

python/sglang/jit_kernel/flash_attn/cute/flash_fwd_sm120.py

@@ -0,0 +1,61 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+# SM120 (Blackwell GeForce / DGX Spark) forward pass.
+#
+# SM120 uses the same SM80-era MMA instructions (mma.sync.aligned.m16n8k16) but has
+# a smaller shared memory capacity (99 KB vs 163 KB on SM80). This module subclasses
+# FlashAttentionForwardSm80 and overrides the SMEM capacity check accordingly.
+
+import cutlass
+import cutlass.utils as utils_basic
+
+from sglang.jit_kernel.flash_attn.cute.flash_fwd import FlashAttentionForwardSm80
+
+
+class FlashAttentionForwardSm120(FlashAttentionForwardSm80):
+    # Keep arch = 80 to use CpAsync code paths (no TMA for output).
+    # The compilation target is determined by the GPU at compile time, not this field.
+    arch = 80
+
+    @staticmethod
+    def can_implement(
+        dtype,
+        head_dim,
+        head_dim_v,
+        tile_m,
+        tile_n,
+        num_stages,
+        num_threads,
+        is_causal,
+        Q_in_regs=False,
+    ) -> bool:
+        """Check if the kernel can be implemented on SM120.
+
+        Same logic as SM80 but uses SM120's shared memory capacity (99 KB).
+        """
+        if dtype not in [cutlass.Float16, cutlass.BFloat16]:
+            return False
+        if head_dim % 8 != 0:
+            return False
+        if head_dim_v % 8 != 0:
+            return False
+        if tile_n % 16 != 0:
+            return False
+        if num_threads % 32 != 0:
+            return False
+        # Shared memory usage: Q tile + (K tile + V tile)
+        smem_usage_Q = tile_m * head_dim * 2
+        smem_usage_K = tile_n * head_dim * num_stages * 2
+        smem_usage_V = tile_n * head_dim_v * num_stages * 2
+        smem_usage_QV = (
+            (smem_usage_Q + smem_usage_V)
+            if not Q_in_regs
+            else max(smem_usage_Q, smem_usage_V)
+        )
+        smem_usage = smem_usage_QV + smem_usage_K
+        # SM120 has 99 KB shared memory (vs 163 KB on SM80)
+        smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_120")
+        if smem_usage > smem_capacity:
+            return False
+        if (tile_m * 2) % num_threads != 0:
+            return False
+        return True

python/sglang/jit_kernel/flash_attn/cute/mma_sm100_desc.py

@@ -0,0 +1,319 @@
+# Copyright (c) 2025, Tri Dao.
+# Ported Cutlass code from C++ to Python:
+# https://github.com/NVIDIA/cutlass/blob/main/include/cute/arch/mma_sm100_desc.hpp
+# https://github.com/NVIDIA/cutlass/blob/main/include/cute/atom/mma_traits_sm100.hpp
+
+from enum import IntEnum
+
+import cutlass
+import cutlass.cute as cute
+
+# ---------------------------------------------------------------------------
+# Enumerations that match the HW encodings (values MUST stay identical)
+# ---------------------------------------------------------------------------
+
+
+class Major(IntEnum):  # matrix “layout” in the ISA docs
+    K = 0
+    MN = 1
+
+
+class ScaleIn(IntEnum):  # negate flags
+    One = 0
+    Neg = 1
+
+
+class Saturate(IntEnum):
+    False_ = 0
+    True_ = 1
+
+
+class CFormat(IntEnum):  # 2-bit field (bits 4-5)
+    F16 = 0
+    F32 = 1
+    S32 = 2
+
+
+class F16F32Format(IntEnum):  # 3-bit field (A/B element type)
+    F16 = 0
+    BF16 = 1
+    TF32 = 2
+
+
+class S8Format(IntEnum):
+    UINT8 = 0
+    INT8 = 1
+
+
+class MXF8F6F4Format(IntEnum):
+    E4M3 = 0
+    E5M2 = 1
+    E2M3 = 3
+    E3M2 = 4
+    E2M1 = 5
+
+
+class MaxShift(IntEnum):
+    NoShift = 0
+    MaxShift8 = 1
+    MaxShift16 = 2
+    MaxShift32 = 3
+
+
+# ---------------------------------------------------------------------------
+# CUTLASS-type → encoding helpers
+# ---------------------------------------------------------------------------
+
+
+def to_UMMA_format(cutlass_type) -> int:
+    """
+    Map a CUTLASS scalar class to the 3-bit encoding for Matrix A/B.
+    """
+    if cutlass_type is cutlass.Int8:
+        return S8Format.INT8
+    # Unsigned 8-bit (if available in your CUTLASS build)
+    if cutlass_type is cutlass.Uint8:
+        return S8Format.UINT8
+    # FP-16 / BF-16
+    if cutlass_type is cutlass.Float16:
+        return F16F32Format.F16
+    if cutlass_type is cutlass.BFloat16:
+        return F16F32Format.BF16
+    # TensorFloat-32 (8-bit exponent, 10-bit mantissa packed in 19 bits)
+    if cutlass_type is cutlass.TFloat32:
+        return F16F32Format.TF32
+    # Float-8 / Float-6 / Float-4 – add whenever CUTLASS exposes them
+    if cutlass_type is cutlass.Float8E4M3FN:
+        return MXF8F6F4Format.E4M3
+    if cutlass_type is cutlass.Float8E5M2:
+        return MXF8F6F4Format.E5M2
+    raise TypeError(f"Unsupported CUTLASS scalar type for A/B: {cutlass_type!r}")
+
+
+def to_C_format(cutlass_type) -> int:
+    """
+    Map a CUTLASS scalar class to the 2-bit accumulator encoding.
+    """
+    if cutlass_type is cutlass.Float16:
+        return CFormat.F16
+    if cutlass_type is cutlass.Float32:
+        return CFormat.F32
+    if cutlass_type is cutlass.Int32:
+        return CFormat.S32
+    raise TypeError(
+        f"Unsupported CUTLASS scalar type for accumulator: {cutlass_type!r}"
+    )
+
+
+# ---------------------------------------------------------------------------
+# The constructor – accepts only CUTLASS scalar classes
+# ---------------------------------------------------------------------------
+
+
+def make_instr_desc(
+    a_type,  # CUTLASS scalar class, e.g. cutlass.Int8
+    b_type,
+    c_type,
+    M: int,  # 64, 128 or 256
+    N: int,  # 8 … 256 (multiple of 8)
+    a_major: Major,
+    b_major: Major,
+    a_neg: ScaleIn = ScaleIn.One,
+    b_neg: ScaleIn = ScaleIn.One,
+    c_sat: Saturate = Saturate.False_,
+    is_sparse: bool = False,
+    max_shift: MaxShift = MaxShift.NoShift,
+) -> int:
+    """
+    Build the 32-bit instruction descriptor for Blackwell MMA.
+    All matrix/accumulator **types must be CUTLASS scalar classes** –
+    passing integers is forbidden.
+    """
+    # --- encode element formats -------------------------------------------------
+    a_fmt = int(to_UMMA_format(a_type))
+    b_fmt = int(to_UMMA_format(b_type))
+    c_fmt = int(to_C_format(c_type))
+
+    # --- range checks on M/N -----------------------------------------------------
+    if M not in (64, 128, 256):
+        raise ValueError("M must be 64, 128 or 256")
+    if N < 8 or N > 256 or (N & 7):
+        raise ValueError("N must be a multiple of 8 in the range 8…256")
+
+    m_dim = M >> 4  # 5-bit field
+    n_dim = N >> 3  # 6-bit field
+
+    # fmt: off
+    # --- pack the bit-fields -----------------------------------------------------
+    desc = 0
+    desc |= (0                 & 0x3) << 0        # sparse_id2 (always 0 here)
+    desc |= (int(is_sparse)    & 0x1) << 2        # sparse_flag
+    desc |= (int(c_sat)        & 0x1) << 3        # saturate
+    desc |= (c_fmt             & 0x3) << 4        # c_format
+    desc |= (a_fmt             & 0x7) << 7        # a_format
+    desc |= (b_fmt             & 0x7) << 10       # b_format
+    desc |= (int(a_neg)        & 0x1) << 13       # a_negate
+    desc |= (int(b_neg)        & 0x1) << 14       # b_negate
+    desc |= (int(a_major)      & 0x1) << 15       # a_major
+    desc |= (int(b_major)      & 0x1) << 16       # b_major
+    desc |= (n_dim             & 0x3F) << 17      # n_dim (6 bits)
+    desc |= (m_dim             & 0x1F) << 24      # m_dim (5 bits)
+    desc |= (int(max_shift)    & 0x3) << 30       # max_shift (2 bits)
+    # fmt: on
+
+    return desc & 0xFFFF_FFFF  # ensure 32-bit result
+
+
+def mma_op_to_idesc(op: cute.nvgpu.tcgen05.mma.MmaOp):
+    return make_instr_desc(
+        op.a_dtype,
+        op.b_dtype,
+        op.acc_dtype,
+        op.shape_mnk[0],
+        op.shape_mnk[1],
+        (
+            Major.K
+            if op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K
+            else Major.MN
+        ),
+        (
+            Major.K
+            if op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K
+            else Major.MN
+        ),
+    )
+
+
+class LayoutType(IntEnum):  # occupies the top-3 bits [61:64)
+    SWIZZLE_NONE = 0  # (a.k.a. “INTERLEAVE” in older docs)
+    SWIZZLE_128B_BASE32B = 1
+    SWIZZLE_128B = 2
+    SWIZZLE_64B = 4
+    SWIZZLE_32B = 6
+    # values 3,5,7 are reserved / illegal for UMMA
+
+
+# ---------------------------------------------------------------------------
+#  Helpers – figure out the SWIZZLE_* family from the tensor layout
+# ---------------------------------------------------------------------------
+
+
+def _layout_type(swizzle: cute.Swizzle) -> LayoutType:
+    B, M, S = swizzle.num_bits, swizzle.num_base, swizzle.num_shift
+
+    if M == 4:  # Swizzle<*,4,3>
+        if S != 3:
+            raise ValueError("Unexpected swizzle shift – want S==3 for M==4")
+        return {
+            0: LayoutType.SWIZZLE_NONE,
+            1: LayoutType.SWIZZLE_32B,
+            2: LayoutType.SWIZZLE_64B,
+            3: LayoutType.SWIZZLE_128B,
+        }[
+            B
+        ]  # KeyError ⇒ invalid B→ raise
+    if M == 5:  # Swizzle<2,5,2> (the only legal triple for M==5)
+        if (B, S) != (2, 2):
+            raise ValueError("Only Swizzle<2,5,2> supported for 128B_BASE32B")
+        return LayoutType.SWIZZLE_128B_BASE32B
+
+    # Any other (M,B,S) triple is not a UMMA-legal shared-memory layout
+    raise ValueError("Unsupported swizzle triple for UMMA smem descriptor")
+
+
+def make_smem_desc_base(
+    layout: cute.Layout, swizzle: cute.Swizzle, major: Major
+) -> int:
+    """
+    Convert a 2-D *shared-memory* Cute layout into the Blackwell 64-bit
+    smem-descriptor, without the smem start address.
+    layout must correspond to layout of an uint128 tensor.
+    """
+    # ------------------------------------------------------------------ meta
+    layout_type = _layout_type(swizzle)  # resolve SWIZZLE_* family
+
+    VERSION = 1  # bits 46–47
+    LBO_MODE = 0  # bit  52
+    BASE_OFFSET = 0  # bits 49–51   (CUTLASS always 0)
+
+    # ---------------------------------------------------------- strides  (units: uint128_t = 16 B)
+    swizzle_atom_mn_size = {
+        LayoutType.SWIZZLE_NONE: 1,
+        LayoutType.SWIZZLE_32B: 2,
+        LayoutType.SWIZZLE_64B: 4,
+        LayoutType.SWIZZLE_128B: 8,
+        LayoutType.SWIZZLE_128B_BASE32B: 8,
+    }[layout_type]
+
+    if major is Major.MN:
+        swizzle_atom_k_size = 4 if layout_type is LayoutType.SWIZZLE_128B_BASE32B else 8
+        canonical_layout = cute.logical_divide(
+            layout, (swizzle_atom_mn_size, swizzle_atom_k_size)
+        )
+        if not cute.is_congruent(canonical_layout, ((1, 1), (1, 1))):
+            raise ValueError(
+                "Not a canonical UMMA_MN Layout: Expected profile failure."
+            )
+        stride_00 = canonical_layout.stride[0][0]
+        if layout_type is not LayoutType.SWIZZLE_NONE and stride_00 != 1:
+            raise ValueError("Not a canonical UMMA_MN Layout: Expected stride failure.")
+        stride_10 = canonical_layout.stride[1][0]
+        if stride_10 != swizzle_atom_mn_size:
+            raise ValueError("Not a canonical UMMA_MN Layout: Expected stride failure.")
+        stride_01, stride_11 = (
+            canonical_layout.stride[0][1],
+            canonical_layout.stride[1][1],
+        )
+        if layout_type is LayoutType.SWIZZLE_NONE:
+            stride_byte_offset, leading_byte_offset = stride_01, stride_11
+        else:
+            stride_byte_offset, leading_byte_offset = stride_11, stride_01
+    else:
+        if layout_type == LayoutType.SWIZZLE_128B_BASE32B:
+            raise ValueError("SWIZZLE_128B_BASE32B is invalid for Major-K")
+        if not cute.size(layout.shape[0]) % 8 == 0:
+            raise ValueError(
+                "Not a canonical UMMA_K Layout: Expected MN-size multiple of 8."
+            )
+        canonical_layout = cute.logical_divide(layout, (8, 2))
+        if not cute.is_congruent(canonical_layout, ((1, 1), (1, 1))):
+            raise ValueError("Not a canonical UMMA_K Layout: Expected profile failure.")
+        stride_00 = canonical_layout.stride[0][0]
+        if stride_00 != swizzle_atom_mn_size:
+            raise ValueError("Not a canonical UMMA_K Layout: Expected stride failure.")
+        stride_10 = canonical_layout.stride[1][0]
+        if layout_type is not LayoutType.SWIZZLE_NONE and stride_10 != 1:
+            raise ValueError("Not a canonical UMMA_K Layout: Expected stride failure.")
+        stride_01 = canonical_layout.stride[0][1]
+        stride_byte_offset, leading_byte_offset = stride_01, stride_10
+
+    # ------------------------------------------------------------------ pack
+    desc = 0
+    # leading_byte_offset_  [16:30)
+    desc |= (leading_byte_offset & 0x3FFF) << 16
+    # stride_byte_offset_   [32:46)
+    desc |= (stride_byte_offset & 0x3FFF) << 32
+    # version_             [46:48)
+    desc |= (VERSION & 0x3) << 46
+    # base_offset_         [49:52)
+    desc |= (BASE_OFFSET & 0x7) << 49
+    # lbo_mode_            [52:53)
+    desc |= (LBO_MODE & 0x1) << 52
+    # layout_type_         [61:64)
+    desc |= (int(layout_type) & 0x7) << 61
+
+    return desc & 0xFFFF_FFFF_FFFF_FFFF  # force 64-bit width
+
+
+def make_smem_desc_start_addr(start_addr: cute.Pointer) -> cutlass.Int32:
+    # 14 bits, remove 4 LSB (bits 0-13 in desc)
+    return (start_addr.toint() & 0x3FFFF) >> 4
+
+
+def smem_desc_base_from_tensor(sA: cute.Tensor, major: Major) -> int:
+    sA_swizzle = sA.iterator.type.swizzle_type
+    return make_smem_desc_base(
+        cute.recast_layout(128, sA.element_type.width, sA.layout[0]),
+        sA_swizzle,
+        major,
+    )

python/sglang/jit_kernel/flash_attn/cute/named_barrier.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+
+import enum
+
+
+class NamedBarrierFwd(enum.IntEnum):
+    Epilogue = enum.auto()  # starts from 1 as barrier 0 is reserved for sync_threads()
+    WarpSchedulerWG1 = enum.auto()
+    WarpSchedulerWG2 = enum.auto()
+    WarpSchedulerWG3 = enum.auto()
+    PFull = enum.auto()
+    PEmpty = enum.auto()
+
+
+class NamedBarrierFwdSm100(enum.IntEnum):
+    Epilogue = enum.auto()  # starts from 1 as barrier 0 is reserved for sync_threads()
+    TmemPtr = enum.auto()
+    SoftmaxStatsW0 = enum.auto()
+    SoftmaxStatsW1 = enum.auto()
+    SoftmaxStatsW2 = enum.auto()
+    SoftmaxStatsW3 = enum.auto()
+    SoftmaxStatsW4 = enum.auto()
+    SoftmaxStatsW5 = enum.auto()
+    SoftmaxStatsW6 = enum.auto()
+    SoftmaxStatsW7 = enum.auto()
+    Softmax = enum.auto()
+    Correction = enum.auto()
+
+
+class NamedBarrierBwd(enum.IntEnum):
+    Epilogue = enum.auto()
+    WarpSchedulerWG1 = enum.auto()
+    WarpSchedulerWG2 = enum.auto()
+    WarpSchedulerWG3 = enum.auto()
+    PdS = enum.auto()
+    dQFullWG0 = enum.auto()
+    dQFullWG1 = enum.auto()
+    dQFullWG2 = enum.auto()
+    dQEmptyWG0 = enum.auto()
+    dQEmptyWG1 = enum.auto()
+    dQEmptyWG2 = enum.auto()
+
+
+class NamedBarrierBwdSm100(enum.IntEnum):
+    EpilogueWG1 = enum.auto()
+    EpilogueWG2 = enum.auto()
+    Compute = enum.auto()
+    dQaccReduce = enum.auto()
+    TmemPtr = enum.auto()
+
+
+class NamedBarrierFwdSm100_MLA2CTA(enum.IntEnum):
+    Epilogue = enum.auto()
+    TmemPtr = enum.auto()
+    Cpasync = enum.auto()
+    Softmax = enum.auto()
+    SoftmaxStatsFull = enum.auto()
+    SoftmaxStatsEmpty = enum.auto()

python/sglang/jit_kernel/flash_attn/cute/pack_gqa.py

@@ -0,0 +1,300 @@
+# Copyright (c) 2025, Tri Dao.
+
+from dataclasses import dataclass
+from typing import Tuple, Union
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.nvgpu import cpasync
+from quack import layout_utils
+
+import sglang.jit_kernel.flash_attn.cute.utils as utils
+
+
+def pack_gqa_layout(T, qhead_per_kvhead, nheads_kv, head_idx):
+    """Reshape a tensor to fold qhead_per_kvhead into the seqlen dimension (mode 0).
+
+    The head dimension is at mode ``head_idx``.  Modes before it (1..head_idx-1)
+    are kept as-is (e.g. headdim for Q/O tensors), and modes after it are kept
+    as-is (e.g. batch).
+
+    For Q/O tensors (head_idx=2):
+        (seqlen_q, headdim, nheads, batch, ...) -> ((qhead_per_kvhead, seqlen_q), headdim, nheads_kv, batch, ...)
+    For LSE tensors (head_idx=1):
+        (seqlen_q, nheads, batch, ...) -> ((qhead_per_kvhead, seqlen_q), nheads_kv, batch, ...)
+    """
+    head_stride = T.stride[head_idx]
+    shape_packed = (
+        (qhead_per_kvhead, T.shape[0]),
+        *[T.shape[i] for i in range(1, head_idx)],
+        nheads_kv,
+        *[T.shape[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    stride_packed = (
+        (head_stride, T.stride[0]),
+        *[T.stride[i] for i in range(1, head_idx)],
+        head_stride * qhead_per_kvhead,
+        *[T.stride[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    return cute.make_tensor(
+        T.iterator, cute.make_layout(shape_packed, stride=stride_packed)
+    )
+
+
+def make_packgqa_tiled_tma_atom(
+    op: cute.atom.CopyOp,
+    gmem_tensor: cute.Tensor,
+    smem_layout: Union[cute.Layout, cute.ComposedLayout],
+    cta_tiler: Tuple[int, int],
+    qhead_per_kvhead: int,
+    head_idx: int,
+):
+    # This packing and unpacking of the layout is so that we keep the same TMA dimension as usual.
+    # e.g. for (seqlen, d, nheads, b) layout, we still have 4D TMA after packing to
+    # ((nheads, seqlen), d, b).
+    # If we instead pack directly to ((qhead_per_kvhead, seqlen), d, nheads_kv, b) we'd have 5D TMA.
+    # Pack headdim and seqlen dim into 1: (seqlen, d, nheads, b) -> ((nheads, seqlen), d, b)
+    gmem_tensor = layout_utils.select(
+        gmem_tensor,
+        [head_idx, *range(head_idx), *range(head_idx + 1, cute.rank(gmem_tensor))],
+    )
+    gmem_tensor = cute.group_modes(gmem_tensor, 0, 2)
+    assert (
+        cta_tiler[0] % qhead_per_kvhead == 0
+    ), "CTA tile size in the seqlen dimension must be divisible by qhead_per_kvhead"
+    tma_atom, tma_tensor = cpasync.make_tiled_tma_atom(
+        op,
+        gmem_tensor,
+        smem_layout,
+        (
+            (qhead_per_kvhead, cta_tiler[0] // qhead_per_kvhead),
+            cta_tiler[1],
+        ),  # No mcast
+    )
+    # Unpack from ((nheads, seqlen), d, b) -> ((qhead_per_kvhead, seqlen), d, nheads_kv, b)
+    T = tma_tensor
+    shape_packed = (
+        (qhead_per_kvhead, T.shape[0][1]),
+        *[T.shape[i] for i in range(1, head_idx)],
+        T.shape[0][0] // qhead_per_kvhead,
+        *[T.shape[i] for i in range(head_idx, len(T.shape))],
+    )
+    stride_packed = (
+        *[T.stride[i] for i in range(head_idx)],
+        T.stride[0][0] * qhead_per_kvhead,
+        *[T.stride[i] for i in range(head_idx, len(T.shape))],
+    )
+    tma_tensor = cute.make_tensor(
+        T.iterator, cute.make_layout(shape_packed, stride=stride_packed)
+    )
+    return tma_atom, tma_tensor
+
+
+def unpack_gqa_layout(T, qhead_per_kvhead, head_idx):
+    """Reverse of pack_gqa_layout: unfold qhead_per_kvhead from the seqlen dimension (mode 0).
+
+    The head dimension is at mode ``head_idx``.  Modes before it (1..head_idx-1)
+    are kept as-is (e.g. headdim for Q/O tensors), and modes after it are kept
+    as-is (e.g. batch).
+
+    For Q/O tensors (head_idx=2):
+        ((qhead_per_kvhead, seqlen_q), headdim, nheads_kv, batch, ...) -> (seqlen_q, headdim, nheads, batch, ...)
+    For LSE tensors (head_idx=1):
+        ((qhead_per_kvhead, seqlen_q), nheads_kv, batch, ...) -> (seqlen_q, nheads, batch, ...)
+    """
+    seqlen_stride = T.stride[0][1]
+    head_stride = T.stride[0][0]
+    shape_unpacked = (
+        T.shape[0][1],
+        *[T.shape[i] for i in range(1, head_idx)],
+        T.shape[head_idx] * qhead_per_kvhead,
+        *[T.shape[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    stride_unpacked = (
+        seqlen_stride,
+        *[T.stride[i] for i in range(1, head_idx)],
+        head_stride,
+        *[T.stride[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    return cute.make_tensor(
+        T.iterator, cute.make_layout(shape_unpacked, stride=stride_unpacked)
+    )
+
+
+@dataclass
+class PackGQA:
+    m_block_size: cutlass.Constexpr[int]
+    head_dim_padded: cutlass.Constexpr[int]
+    check_hdim_oob: cutlass.Constexpr[bool]
+    qhead_per_kvhead: cutlass.Constexpr[bool]
+
+    @cute.jit
+    def compute_ptr(
+        self,
+        tensor: cute.Tensor,
+        cRows: cute.Tensor,
+        tidx: cutlass.Int32,
+        block: cutlass.Int32,
+        threads_per_row: cutlass.Constexpr[int],
+        num_threads: cutlass.Constexpr[int],
+    ):
+        num_ptr_per_thread = cute.ceil_div(cute.size(cRows), threads_per_row)
+        tPrPtr = cute.make_rmem_tensor(num_ptr_per_thread, cutlass.Int64)
+        for i in cutlass.range_constexpr(num_ptr_per_thread):
+            row = i * num_threads + cRows[tidx % threads_per_row][0]
+            idx = block * self.m_block_size + row
+            m_idx = idx // self.qhead_per_kvhead
+            h_idx = idx - m_idx * self.qhead_per_kvhead
+            tPrPtr[i] = utils.elem_pointer(tensor, ((h_idx, m_idx),)).toint()
+        return tPrPtr
+
+    @cute.jit
+    def load_Q(
+        self,
+        mQ: cute.Tensor,  # ((qhead_per_kvhead, seqlen_q), headdim)
+        sQ: cute.Tensor,  # (m_block_size, head_dim_padded)
+        gmem_tiled_copy: cute.TiledCopy,
+        tidx: cutlass.Int32,
+        block: cutlass.Int32,
+        seqlen: cutlass.Int32,
+    ):
+        gmem_thr_copy = gmem_tiled_copy.get_slice(tidx)
+        cQ = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded))
+        tQsQ = gmem_thr_copy.partition_D(sQ)
+        tQcQ = gmem_thr_copy.partition_S(cQ)
+        t0QcQ = gmem_thr_copy.get_slice(0).partition_S(cQ)
+        tQpQ = utils.predicate_k(tQcQ, limit=mQ.shape[1])
+        tQcQ_row = tQcQ[0, None, 0]
+        threads_per_row = gmem_tiled_copy.layout_tv_tiled.shape[0][0]
+        assert (
+            cute.arch.WARP_SIZE % threads_per_row == 0
+        ), "threads_per_row must divide WARP_SIZE"
+        num_threads = gmem_tiled_copy.size
+        tPrQPtr = self.compute_ptr(
+            mQ[None, 0], tQcQ_row, tidx, block, threads_per_row, num_threads
+        )
+        for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])):
+            q_ptr_i64 = utils.shuffle_sync(
+                tPrQPtr[m // threads_per_row],
+                m % threads_per_row,
+                width=threads_per_row,
+            )
+            q_gmem_ptr = cute.make_ptr(
+                mQ.element_type, q_ptr_i64, cute.AddressSpace.gmem, assumed_align=16
+            )
+            if (
+                t0QcQ[0, m, 0][0]
+                < seqlen * self.qhead_per_kvhead
+                - block * self.m_block_size
+                - tQcQ_row[0][0]
+            ):
+                mQ_cur = cute.make_tensor(q_gmem_ptr, (self.head_dim_padded,))
+                elems_per_load = cute.size(tQsQ.shape[0][0])
+                mQ_cur_copy = cute.tiled_divide(mQ_cur, (elems_per_load,))
+                for k in cutlass.range_constexpr(cute.size(tQsQ.shape[2])):
+                    ki = tQcQ[0, 0, k][1] // elems_per_load
+                    cute.copy(
+                        gmem_thr_copy,
+                        mQ_cur_copy[None, ki],
+                        tQsQ[None, m, k],
+                        pred=(
+                            tQpQ[None, m, k]
+                            if cutlass.const_expr(self.check_hdim_oob)
+                            else None
+                        ),
+                    )
+            # We don't need to clear the sQ smem tiles since we'll only write out the valid outputs
+
+    @cute.jit
+    def store_LSE(
+        self,
+        mLSE: cute.Tensor,  # (qhead_per_kvhead, seqlen_q)
+        tLSErLSE: cute.Tensor,  # (m_block_size, head_dim_padded)
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        block: cutlass.Int32,
+        seqlen: cutlass.Int32,
+    ):
+        thr_mma = tiled_mma.get_slice(tidx)
+        caccO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded))
+        taccOcO = thr_mma.partition_C(caccO)
+        taccOcO_row = layout_utils.reshape_acc_to_mn(taccOcO)[None, 0]
+        assert cute.size(tLSErLSE) == cute.size(taccOcO_row)
+        threads_per_row = tiled_mma.tv_layout_C.shape[0][0]
+        assert (
+            cute.arch.WARP_SIZE % threads_per_row == 0
+        ), "threads_per_row must divide WARP_SIZE"
+        assert cute.size(tLSErLSE) <= threads_per_row
+        num_threads = tiled_mma.size
+        tPrLSEPtr = self.compute_ptr(
+            mLSE, taccOcO_row, tidx, block, threads_per_row, num_threads
+        )
+        for m in cutlass.range_constexpr(cute.size(tLSErLSE)):
+            lse_ptr_i64 = utils.shuffle_sync(
+                tPrLSEPtr[m // threads_per_row],
+                m % threads_per_row,
+                width=threads_per_row,
+            )
+            lse_gmem_ptr = cute.make_ptr(
+                mLSE.element_type, lse_ptr_i64, cute.AddressSpace.gmem, assumed_align=4
+            )
+            row = block * self.m_block_size + taccOcO_row[m][0]
+            # Only the thread corresponding to column 0 writes out the lse to gmem
+            if taccOcO[0][1] == 0 and row < seqlen * self.qhead_per_kvhead:
+                mLSE_copy = cute.make_tensor(lse_gmem_ptr, (1,))
+                mLSE_copy[0] = tLSErLSE[m]
+
+    @cute.jit
+    def store_O(
+        self,
+        mO: cute.Tensor,  # ((qhead_per_kvhead, seqlen_q), headdim)
+        tOrO: cute.Tensor,  # (m_block_size, head_dim_padded) split across threads according to gmem_tiled_copy
+        gmem_tiled_copy: cute.TiledCopy,
+        tidx: cutlass.Int32,
+        block: cutlass.Int32,
+        seqlen: cutlass.Int32,
+    ):
+        gmem_thr_copy = gmem_tiled_copy.get_slice(tidx)
+        cO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded))
+        tOcO = gmem_thr_copy.partition_S(cO)
+        t0OcO = gmem_thr_copy.get_slice(0).partition_S(cO)
+        tOpO = utils.predicate_k(tOcO, limit=mO.shape[1])
+        tOcO_row = tOcO[0, None, 0]
+        threads_per_row = gmem_tiled_copy.layout_tv_tiled.shape[0][0]
+        assert (
+            cute.arch.WARP_SIZE % threads_per_row == 0
+        ), "threads_per_row must divide WARP_SIZE"
+        num_threads = gmem_tiled_copy.size
+        tPrOPtr = self.compute_ptr(
+            mO[None, 0], tOcO_row, tidx, block, threads_per_row, num_threads
+        )
+        for m in cutlass.range_constexpr(cute.size(tOrO.shape[1])):
+            o_ptr_i64 = utils.shuffle_sync(
+                tPrOPtr[m // threads_per_row],
+                m % threads_per_row,
+                width=threads_per_row,
+            )
+            o_gmem_ptr = cute.make_ptr(
+                mO.element_type, o_ptr_i64, cute.AddressSpace.gmem, assumed_align=16
+            )
+            if (
+                t0OcO[0, m, 0][0]
+                < seqlen * self.qhead_per_kvhead
+                - block * self.m_block_size
+                - tOcO_row[0][0]
+            ):
+                mO_cur = cute.make_tensor(o_gmem_ptr, (self.head_dim_padded,))
+                elems_per_load = cute.size(tOrO.shape[0][0])
+                mO_cur_copy = cute.tiled_divide(mO_cur, (elems_per_load,))
+                for k in cutlass.range_constexpr(cute.size(tOrO.shape[2])):
+                    ki = tOcO[0, 0, k][1] // elems_per_load
+                    cute.copy(
+                        gmem_thr_copy,
+                        tOrO[None, m, k],
+                        mO_cur_copy[None, ki],
+                        pred=(
+                            tOpO[None, m, k]
+                            if cutlass.const_expr(self.check_hdim_oob)
+                            else None
+                        ),
+                    )

python/sglang/jit_kernel/flash_attn/cute/paged_kv.py

@@ -0,0 +1,393 @@
+import math
+from dataclasses import dataclass
+from typing import Optional, Type
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, const_expr
+from cutlass.cute import FastDivmodDivisor
+from cutlass.cute.nvgpu import cpasync
+from quack.cute_dsl_utils import ParamsBase
+
+from sglang.jit_kernel.flash_attn.cute import utils
+
+
+@dataclass
+class PagedKVManager(ParamsBase):
+    mPageTable: cute.Tensor
+    mK_paged: cute.Tensor
+    mV_paged: cute.Tensor
+    mSFK_paged: Optional[cute.Tensor]
+    mSFV_paged: Optional[cute.Tensor]
+    thread_idx: Int32
+
+    page_size_divmod: FastDivmodDivisor
+    seqlen_k: Int32
+    leftpad_k: Int32
+    n_block_size: Int32
+    num_threads: cutlass.Constexpr[Int32]
+    head_dim_padded: cutlass.Constexpr[Int32]
+    head_dim_v_padded: cutlass.Constexpr[Int32]
+
+    arch: cutlass.Constexpr[Int32]
+    v_gmem_transposed: cutlass.Constexpr[bool]
+
+    gmem_threads_per_row: cutlass.Constexpr[Int32]
+    page_entry_per_thread: Int32
+    async_copy_elems: Int32
+
+    gmem_tiled_copy_KV: cute.TiledCopy
+    gmem_thr_copy_KV: cute.TiledCopy
+    gmem_tiled_copy_sf_KV: Optional[cute.TiledCopy]
+    gmem_thr_copy_sf_KV: Optional[cute.TiledCopy]
+    tPrPage: cute.Tensor
+    tPrPageOffset: cute.Tensor
+    tKpK: cute.Tensor
+    tVpV: cute.Tensor
+
+    @staticmethod
+    def create(
+        mPageTable: cute.Tensor,
+        mK_paged: cute.Tensor,
+        mV_paged: cute.Tensor,
+        page_size_divmod: FastDivmodDivisor,
+        bidb: Int32,
+        bidh: Int32,
+        thread_idx: Int32,
+        seqlen_k: Int32,
+        leftpad_k: Int32,
+        n_block_size: cutlass.Constexpr[Int32],
+        head_dim_padded: cutlass.Constexpr[Int32],
+        head_dim_v_padded: cutlass.Constexpr[Int32],
+        num_threads: cutlass.Constexpr[Int32],
+        dtype: Type[cutlass.Numeric],
+        mSFK_paged: Optional[cute.Tensor] = None,
+        mSFV_paged: Optional[cute.Tensor] = None,
+        arch: cutlass.Constexpr[int] = 100,
+    ):
+        # SM100 transposes V in gmem to (dv, page_size, num_pages);
+        # SM90 keeps V as (page_size, dv, num_pages), same layout as K.
+        v_gmem_transposed = arch != 90
+        universal_copy_bits = 128
+        async_copy_elems = universal_copy_bits // dtype.width
+        dtype_bytes = dtype.width // 8
+        gmem_k_block_size = math.gcd(
+            head_dim_padded,
+            head_dim_v_padded,
+            128 // dtype_bytes,
+        )
+        assert gmem_k_block_size % async_copy_elems == 0
+        gmem_threads_per_row = gmem_k_block_size // async_copy_elems
+        assert cute.arch.WARP_SIZE % gmem_threads_per_row == 0
+        atom_async_copy = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+            dtype,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        thr_layout = cute.make_ordered_layout(
+            (num_threads // gmem_threads_per_row, gmem_threads_per_row),
+            order=(1, 0),
+        )
+        val_layout = cute.make_layout((1, async_copy_elems))
+        gmem_tiled_copy_KV = cute.make_tiled_copy_tv(
+            atom_async_copy, thr_layout, val_layout
+        )
+        gmem_thr_copy_KV = gmem_tiled_copy_KV.get_slice(thread_idx)
+        page_entry_per_thread = n_block_size // num_threads
+
+        if const_expr(mSFK_paged is not None or mSFV_paged is not None):
+            atom_async_copy_sf = cute.make_copy_atom(
+                cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS),
+                dtype,
+                num_bits_per_copy=32,
+            )
+            thr_layout_sf = cute.make_ordered_layout(
+                ((num_threads // gmem_threads_per_row, gmem_threads_per_row), 1),
+                order=((1, 0), 2),
+            )
+            val_layout_sf = cute.make_layout((1, 4))
+            gmem_tiled_copy_sf_KV = cute.make_tiled_copy_tv(
+                atom_async_copy_sf,
+                thr_layout_sf,
+                val_layout_sf,
+            )
+            gmem_thr_copy_sf_KV = gmem_tiled_copy_sf_KV.get_slice(thread_idx)
+        else:
+            gmem_tiled_copy_sf_KV = None
+            gmem_thr_copy_sf_KV = None
+
+        tPrPage = cute.make_rmem_tensor((page_entry_per_thread,), Int32)
+        tPrPageOffset = cute.make_rmem_tensor((page_entry_per_thread,), Int32)
+
+        mPageTable = mPageTable[bidb, None]
+        mK_paged = mK_paged[None, None, bidh, None]
+        mV_paged = mV_paged[None, None, bidh, None]
+
+        if const_expr(mSFK_paged is not None):
+            mSFK_paged = mSFK_paged[None, None, bidh, None]
+        if const_expr(mSFV_paged is not None):
+            mSFV_paged = mSFV_paged[None, None, bidh, None]
+
+        cK = cute.make_identity_tensor((n_block_size, head_dim_padded))
+        tKcK = gmem_thr_copy_KV.partition_S(cK)
+        tKpK = utils.predicate_k(tKcK, limit=mK_paged.shape[1])
+
+        if const_expr(head_dim_padded == head_dim_v_padded):
+            tVpV = tKpK
+        else:
+            cV = cute.make_identity_tensor((n_block_size, head_dim_v_padded))
+            tVcV = gmem_thr_copy_KV.partition_S(cV)
+            # When V is transposed in gmem, dv is shape[0]; otherwise dv is shape[1] (same as K)
+            V_limit = cute.size(mV_paged.shape[0 if v_gmem_transposed else 1])
+            tVpV = utils.predicate_k(tVcV, limit=V_limit)
+
+        return PagedKVManager(
+            mPageTable,
+            mK_paged,
+            mV_paged,
+            mSFK_paged,
+            mSFV_paged,
+            thread_idx,
+            page_size_divmod,
+            seqlen_k,
+            leftpad_k,
+            n_block_size,
+            num_threads,
+            head_dim_padded,
+            head_dim_v_padded,
+            arch,
+            v_gmem_transposed,
+            gmem_threads_per_row,
+            page_entry_per_thread,
+            async_copy_elems,
+            gmem_tiled_copy_KV,
+            gmem_thr_copy_KV,
+            gmem_tiled_copy_sf_KV,
+            gmem_thr_copy_sf_KV,
+            tPrPage,
+            tPrPageOffset,
+            tKpK,
+            tVpV,
+        )
+
+    @cute.jit
+    def load_page_table(self, n_block: Int32):
+        for i in cutlass.range(self.page_entry_per_thread, unroll=1):
+            row = (
+                i * self.num_threads
+                + (self.thread_idx % self.gmem_threads_per_row)
+                * (self.num_threads // self.gmem_threads_per_row)
+                + (self.thread_idx // self.gmem_threads_per_row)
+            )
+            row_idx = n_block * self.n_block_size + row
+
+            page_idx, page_offset = divmod(
+                row_idx + self.leftpad_k, self.page_size_divmod
+            )
+
+            is_valid = (
+                (i + 1) * self.num_threads <= self.n_block_size
+                or row < self.n_block_size
+            ) and row_idx < self.seqlen_k
+            page = self.mPageTable[page_idx] if is_valid else 0
+
+            self.tPrPage[i] = page
+            self.tPrPageOffset[i] = page_offset
+
+    @cute.jit
+    def compute_X_ptr(self, K_or_V: str, d_offset: int = 0):
+        tPrXPtr = cute.make_rmem_tensor((self.page_entry_per_thread,), cutlass.Int64)
+        mX = self.mK_paged if const_expr(K_or_V == "K") else self.mV_paged
+        # K is always (page_size, d, num_pages). V matches K when not transposed,
+        # but is (dv, page_size, num_pages) when transposed (SM100).
+        transposed = const_expr(K_or_V == "V" and self.v_gmem_transposed)
+        for i in cutlass.range(self.page_entry_per_thread, unroll=1):
+            page = self.tPrPage[i]
+            page_offset = self.tPrPageOffset[i]
+            if const_expr(transposed):
+                tPrXPtr[i] = utils.elem_pointer(
+                    mX, (d_offset, page_offset, page)
+                ).toint()
+            else:
+                tPrXPtr[i] = utils.elem_pointer(
+                    mX, (page_offset, d_offset, page)
+                ).toint()
+        return tPrXPtr
+
+    @cute.jit
+    def _flatten_smem_sm100(self, sX: cute.Tensor, K_or_V: str):
+        """Flatten SM100 smem ((a,b), cta_split, k) to (a,(b,k)); transpose V to (d,page_size)."""
+        sX_pi = cute.make_tensor(
+            sX.iterator,
+            cute.make_layout(
+                (sX.shape[0][0], (sX.shape[0][1], sX.shape[2])),
+                stride=(sX.stride[0][0], (sX.stride[0][1], sX.stride[2])),
+            ),
+        )
+        if const_expr(K_or_V == "V"):
+            sX_pi = cute.make_tensor(
+                sX_pi.iterator, cute.select(sX_pi.layout, mode=[1, 0])
+            )
+        return sX_pi
+
+    @cute.jit
+    def _copy_row_async(
+        self,
+        tXsX: cute.Tensor,
+        tXcX: cute.Tensor,
+        mX_paged_cur_copy: cute.Tensor,
+        m: Int32,
+        should_load: cute.Tensor,
+    ):
+        """Issue cp.async copies for one row across all k-tiles."""
+        for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])):
+            ki = tXcX[0, 0, k][1] // self.async_copy_elems
+            mX_paged_cur_copy_ki = mX_paged_cur_copy[None, ki]
+            tXsX_k = tXsX[None, m, k]
+            mX_paged_cur_copy_ki = cute.make_tensor(
+                mX_paged_cur_copy_ki.iterator, tXsX_k.layout
+            )
+            cute.copy(
+                self.gmem_tiled_copy_KV,
+                mX_paged_cur_copy_ki,
+                tXsX_k,
+                pred=should_load,
+            )
+
+    @cute.jit
+    def compute_sf_X_ptr(self, K_or_V: str):
+        tPrXPtr = cute.make_rmem_tensor((self.page_entry_per_thread,), cutlass.Int64)
+        for i in cutlass.range(self.page_entry_per_thread, unroll=1):
+            page = self.tPrPage[i]
+            page_offset = self.tPrPageOffset[i]
+            if const_expr(K_or_V == "K"):
+                tPrXPtr[i] = utils.elem_pointer(
+                    self.mSFK_paged, (page_offset, 0, page)
+                ).toint()
+            else:
+                tPrXPtr[i] = utils.elem_pointer(
+                    self.mSFV_paged, (0, page_offset, page)
+                ).toint()
+        return tPrXPtr
+
+    @cute.jit
+    def load_KV(self, n_block: Int32, sX: cute.Tensor, K_or_V: str):
+        assert K_or_V in ("K", "V")
+
+        tPrXPtr = self.compute_X_ptr(K_or_V)
+
+        if const_expr(self.arch == 90):
+            # SM90: sX is already stage-sliced by caller (sK[None, None, stage]).
+            # Flatten hierarchical modes to get (n_block_size, head_dim).
+            sX_pi = cute.group_modes(sX, 0, 1)
+            # SM90 does NOT transpose V here (it's transposed via utils.transpose_view before MMA)
+        else:
+            sX_pi = self._flatten_smem_sm100(sX, K_or_V)
+
+        head_dim = (
+            self.head_dim_v_padded
+            if const_expr(K_or_V == "V")
+            else self.head_dim_padded
+        )
+        cX = cute.make_identity_tensor((self.n_block_size, head_dim))
+        tXsX = self.gmem_thr_copy_KV.partition_D(sX_pi)
+        tXcX = self.gmem_thr_copy_KV.partition_S(cX)
+        tXc0X = self.gmem_thr_copy_KV.get_slice(0).partition_S(cX)
+
+        seqlenk_row_limit = (
+            self.seqlen_k - n_block * self.n_block_size - tXcX[0][0]
+            if n_block >= 0
+            else 0
+        )
+        for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])):
+            row_valid = tXc0X[0, m, 0][0] < seqlenk_row_limit
+            should_load = cute.make_fragment_like(tXsX[(0, None), m, 0], cute.Boolean)
+            should_load.fill(row_valid)
+
+            x_ptr_i64 = utils.shuffle_sync(
+                tPrXPtr[m // self.gmem_threads_per_row],
+                m % self.gmem_threads_per_row,
+                width=self.gmem_threads_per_row,
+            )
+            x_gmem_ptr = cute.make_ptr(
+                self.mK_paged.element_type,
+                x_ptr_i64,
+                cute.AddressSpace.gmem,
+                assumed_align=16,
+            )
+            mX_paged_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,)))
+            mX_paged_cur_copy = cute.tiled_divide(
+                mX_paged_cur, (self.async_copy_elems,)
+            )
+            self._copy_row_async(tXsX, tXcX, mX_paged_cur_copy, m, should_load)
+
+    @cute.jit
+    def load_sf_KV(self, n_block: Int32, sSFX: cute.Tensor, K_or_V: str):
+        # sSFX expected as SFK or SFV
+        assert (
+            cute.rank(sSFX) == 3
+        ), f"mismatched rank for sSFX, expected 3 but got {cute.rank(sSFX)}"
+        assert self.gmem_thr_copy_sf_KV is not None
+        # sSFK:  tensor<ptr<f8E8M0FNU, smem, align<1024>> o ((((32,4),1),(32,1)),1,4,2):((((16,4),0),(0,0)),0,1,512)>
+        # sSFV:  tensor<ptr<f8E8M0FNU, smem, align<1024>> o ((((32,4),1),(32,1)),1,4,2):((((16,4),0),(0,0)),0,1,512)>
+
+        head_dim = (
+            self.head_dim_v_padded
+            if const_expr(K_or_V == "V")
+            else self.head_dim_padded
+        )
+
+        sSFX_cpt = cute.filter_zeros(sSFX)
+        sSFX_cpt_shape_nd = (self.n_block_size, head_dim // 32)
+        sSFX_cpt_layout_nd = cute.make_ordered_layout(
+            sSFX_cpt_shape_nd,
+            order=(0, 1),
+        )
+        # (tile_n, 4)
+        sSFX_cpt_nd = cute.composition(sSFX_cpt, sSFX_cpt_layout_nd)
+
+        cX = cute.make_identity_tensor(sSFX_cpt_shape_nd)
+        # ((V, 1), M, 1)
+        tXsX = self.gmem_thr_copy_sf_KV.partition_D(sSFX_cpt_nd)
+        tXcX = self.gmem_thr_copy_sf_KV.partition_S(cX)
+        tXc0X = self.gmem_thr_copy_sf_KV.get_slice(0).partition_S(cX)
+
+        seqlenk_row_limit = (
+            self.seqlen_k - n_block * self.n_block_size - tXcX[0][0]
+            if n_block >= 0
+            else 0
+        )
+
+        tPrSFXPtr = self.compute_sf_X_ptr(K_or_V)
+        assert cute.size(tPrSFXPtr) == cute.size(
+            tXsX, mode=[1]
+        ), "SFX pointer size mismatch"
+
+        # loop over rows
+        for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])):
+            row_valid = tXc0X[0, m, 0][0] < seqlenk_row_limit
+            should_load = cute.make_fragment_like(
+                tXsX[(0, None), m, None], cute.Boolean
+            )
+            should_load.fill(row_valid)
+
+            # Make gmem tensor of size (4,) using tPrSFXPtr
+            # Simplified version of load_KV, no shuffle, 4 elements to copy (hdim = 128)
+            sfx_ptr_i64 = tPrSFXPtr[m]
+            sfx_gmem_ptr = cute.make_ptr(
+                self.mSFK_paged.element_type,
+                sfx_ptr_i64,
+                cute.AddressSpace.gmem,
+                assumed_align=4,
+            )
+            sf_frg_layout = cute.make_layout(((head_dim // 32, 1), 1))
+            mSFX_paged_cur = cute.make_tensor(sfx_gmem_ptr, sf_frg_layout)
+            assert cute.size(mSFX_paged_cur) == cute.size(
+                tXsX[None, 0, None]
+            ), "SFX gmem-smem tensor size mismatch"
+            cute.copy(
+                self.gmem_tiled_copy_sf_KV,
+                mSFX_paged_cur,
+                tXsX[None, m, None],
+                pred=should_load,
+            )

python/sglang/jit_kernel/flash_attn/cute/pipeline.py

@@ -0,0 +1,412 @@
+# Copyright (c) 2025, Tri Dao.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import cutlass.cute as cute
+from cutlass import Boolean, Int32, const_expr
+from cutlass.cutlass_dsl import dsl_user_op, if_generate
+from cutlass.pipeline import NamedBarrier as NamedBarrierOg
+from cutlass.pipeline import PipelineAsync as PipelineAsyncOg
+from cutlass.pipeline import PipelineAsyncUmma as PipelineAsyncUmmaOg
+from cutlass.pipeline import PipelineCpAsync as PipelineCpAsyncOg
+from cutlass.pipeline import PipelineState
+from cutlass.pipeline import PipelineTmaAsync as PipelineTmaAsyncOg
+from cutlass.pipeline import PipelineTmaUmma as PipelineTmaUmmaOg
+from cutlass.pipeline import PipelineUmmaAsync as PipelineUmmaAsyncOg
+from cutlass.pipeline import PipelineUserType
+
+
+def _override_create(parent_cls, child_cls):
+    """Create a static factory that constructs parent_cls then re-classes to child_cls."""
+
+    @staticmethod
+    def create(*args, **kwargs):
+        obj = parent_cls.create(*args, **kwargs)
+        # Can't assign to __class__ directly since the dataclass is frozen
+        object.__setattr__(obj, "__class__", child_cls)
+        return obj
+
+    return create
+
+
+def _make_state(index: Int32, phase: Int32) -> PipelineState:
+    """Construct a PipelineState from index and phase (count/stages unused by callers)."""
+    return PipelineState(stages=0, count=Int32(0), index=index, phase=phase)
+
+
+class PipelineStateSimple:
+    """
+    Pipeline state contains an index and phase bit corresponding to the current position in the circular buffer.
+    Use a single Int32 to store both the index and phase bit, then we use divmod to get the
+    index and phase. If stages is a power of 2, divmod turns into bit twiddling.
+    """
+
+    def __init__(self, stages: int, phase_index: Int32):
+        self._stages = stages
+        self._phase_index = phase_index
+
+    def clone(self) -> "PipelineStateSimple":
+        return PipelineStateSimple(self.stages, self._phase_index)
+
+    @property
+    def stages(self) -> int:
+        return self._stages
+
+    @property
+    def index(self) -> Int32:
+        if const_expr(self._stages == 1):
+            return Int32(0)
+        else:
+            return self._phase_index % self._stages
+
+    @property
+    def phase(self) -> Int32:
+        # PTX docs say that the phase parity needs to be 0 or 1, so by right we need to
+        # take modulo 2. But in practice just passing the phase in without modulo works fine.
+        if const_expr(self._stages == 1):
+            return self._phase_index
+        else:
+            return self._phase_index // self._stages
+
+    def advance(self):
+        if const_expr(self._stages == 1):
+            self._phase_index ^= 1
+        else:
+            self._phase_index += 1
+
+    def __extract_mlir_values__(self):
+        phase_index = self._phase_index
+        return [phase_index.ir_value()]
+
+    def __new_from_mlir_values__(self, values):
+        return PipelineStateSimple(self.stages, Int32(values[0]))
+
+
+def make_pipeline_state(type: PipelineUserType, stages: int):
+    """
+    Creates a pipeline state. Producers are assumed to start with an empty buffer and have a flipped phase bit of 1.
+    """
+    if type is PipelineUserType.Producer:
+        return PipelineStateSimple(stages, Int32(stages))
+    elif type is PipelineUserType.Consumer:
+        return PipelineStateSimple(stages, Int32(0))
+    else:
+        assert (
+            False
+        ), "Error: invalid PipelineUserType specified for make_pipeline_state."
+
+
+# ── Shared helpers ───────────────────────────────────────────────────────────
+
+
+def _call_with_elect_one(parent_method, self, state, elect_one, syncwarp, loc, ip):
+    """Optionally wrap a parent pipeline method call in sync_warp + elect_one."""
+    if const_expr(elect_one):
+        if const_expr(syncwarp):
+            cute.arch.sync_warp()
+        with cute.arch.elect_one():
+            parent_method(self, state, loc=loc, ip=ip)
+    else:
+        parent_method(self, state, loc=loc, ip=ip)
+
+
+# ── Mixin: _w_index / _w_index_phase variants that delegate to parent ───────
+# Each parent class has PipelineState-based methods (producer_acquire, producer_commit,
+# consumer_wait, consumer_release). The _w_index_phase variants just construct a
+# PipelineState from (index, phase) and delegate.
+
+
+class _PipelineIndexPhaseMixin:
+    """Mixin providing _w_index_phase / _w_index methods that delegate to PipelineState-based parents."""
+
+    @dsl_user_op
+    def producer_acquire_w_index_phase(
+        self,
+        index: Int32,
+        phase: Int32,
+        try_acquire_token: Optional[Boolean] = None,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        state = _make_state(index, phase)
+        # Call the parent's producer_acquire (which takes PipelineState)
+        self.producer_acquire(state, try_acquire_token, loc=loc, ip=ip)
+
+    @dsl_user_op
+    def producer_commit_w_index(self, index: Int32, *, loc=None, ip=None):
+        state = _make_state(index, Int32(0))
+        self.producer_commit(state, loc=loc, ip=ip)
+
+    @dsl_user_op
+    def consumer_wait_w_index_phase(
+        self,
+        index: Int32,
+        phase: Int32,
+        try_wait_token: Optional[Boolean] = None,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        state = _make_state(index, phase)
+        self.consumer_wait(state, try_wait_token, loc=loc, ip=ip)
+
+    @dsl_user_op
+    def consumer_release_w_index(self, index: Int32, *, loc=None, ip=None):
+        state = _make_state(index, Int32(0))
+        self.consumer_release(state, loc=loc, ip=ip)
+
+
+# ── NamedBarrier ─────────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class NamedBarrier(NamedBarrierOg):
+    create = _override_create(NamedBarrierOg, None)  # patched below
+
+    @dsl_user_op
+    def arrive_w_index(self, index: Int32, *, loc=None, ip=None) -> None:
+        """
+        The aligned flavor of arrive is used when all threads in the CTA will execute the
+        same instruction. See PTX documentation.
+        """
+        cute.arch.barrier_arrive(
+            barrier_id=self.barrier_id + index,
+            number_of_threads=self.num_threads,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def arrive_and_wait_w_index(self, index: Int32, *, loc=None, ip=None) -> None:
+        cute.arch.barrier(
+            barrier_id=self.barrier_id + index,
+            number_of_threads=self.num_threads,
+            loc=loc,
+            ip=ip,
+        )
+
+
+NamedBarrier.create = _override_create(NamedBarrierOg, NamedBarrier)
+
+
+# ── PipelineAsync ────────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PipelineAsync(_PipelineIndexPhaseMixin, PipelineAsyncOg):
+    """
+    PipelineAsync with optional elect_one for producer_commit and consumer_release.
+
+    When elect_one_*=True (set at create time), only one elected thread per warp
+    signals the barrier arrive. This is useful when the mask count is set to 1 per warp.
+
+    Args (to create):
+        elect_one_commit: If True, only elected thread signals producer_commit.
+        syncwarp_before_commit: If True (default), issue syncwarp before elect_one.
+        elect_one_release: If True, only elected thread signals consumer_release.
+        syncwarp_before_release: If True (default), issue syncwarp before elect_one.
+            Set syncwarp to False when threads are already converged (e.g. after wgmma wait_group).
+    """
+
+    _elect_one_commit: bool = False
+    _syncwarp_before_commit: bool = True
+    _elect_one_release: bool = False
+    _syncwarp_before_release: bool = True
+
+    @staticmethod
+    def create(
+        *args,
+        elect_one_commit: bool = False,
+        syncwarp_before_commit: bool = True,
+        elect_one_release: bool = False,
+        syncwarp_before_release: bool = True,
+        **kwargs,
+    ):
+        obj = PipelineAsyncOg.create(*args, **kwargs)
+        object.__setattr__(obj, "__class__", PipelineAsync)
+        object.__setattr__(obj, "_elect_one_commit", elect_one_commit)
+        object.__setattr__(obj, "_syncwarp_before_commit", syncwarp_before_commit)
+        object.__setattr__(obj, "_elect_one_release", elect_one_release)
+        object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release)
+        return obj
+
+    @dsl_user_op
+    def producer_commit(self, state: PipelineState, *, loc=None, ip=None):
+        _call_with_elect_one(
+            PipelineAsyncOg.producer_commit,
+            self,
+            state,
+            self._elect_one_commit,
+            self._syncwarp_before_commit,
+            loc,
+            ip,
+        )
+
+    @dsl_user_op
+    def consumer_release(self, state: PipelineState, *, loc=None, ip=None):
+        _call_with_elect_one(
+            PipelineAsyncOg.consumer_release,
+            self,
+            state,
+            self._elect_one_release,
+            self._syncwarp_before_release,
+            loc,
+            ip,
+        )
+
+    # _w_index variants inherited from _PipelineIndexPhaseMixin, which delegate
+    # to producer_commit / consumer_release above.
+
+
+# ── PipelineCpAsync ──────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PipelineCpAsync(_PipelineIndexPhaseMixin, PipelineCpAsyncOg):
+    _elect_one_release: bool = False
+    _syncwarp_before_release: bool = True
+
+    @staticmethod
+    def create(
+        *args,
+        elect_one_release: bool = False,
+        syncwarp_before_release: bool = True,
+        **kwargs,
+    ):
+        obj = PipelineCpAsyncOg.create(*args, **kwargs)
+        object.__setattr__(obj, "__class__", PipelineCpAsync)
+        object.__setattr__(obj, "_elect_one_release", elect_one_release)
+        object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release)
+        return obj
+
+    @dsl_user_op
+    def consumer_release(self, state: PipelineState, *, loc=None, ip=None):
+        _call_with_elect_one(
+            PipelineCpAsyncOg.consumer_release,
+            self,
+            state,
+            self._elect_one_release,
+            self._syncwarp_before_release,
+            loc,
+            ip,
+        )
+
+    # _w_index variants inherited from _PipelineIndexPhaseMixin.
+
+
+# ── PipelineTmaAsync ────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PipelineTmaAsync(_PipelineIndexPhaseMixin, PipelineTmaAsyncOg):
+    """Override producer_acquire to take in extra_tx_count parameter."""
+
+    @dsl_user_op
+    def producer_acquire(
+        self,
+        state: PipelineState,
+        try_acquire_token: Optional[Boolean] = None,
+        extra_tx_count: int = 0,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        """
+        TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks.
+        """
+        if_generate(
+            try_acquire_token is None or try_acquire_token == 0,
+            lambda: self.sync_object_empty.wait(
+                state.index, state.phase, loc=loc, ip=ip
+            ),
+            loc=loc,
+            ip=ip,
+        )
+        if const_expr(extra_tx_count == 0):
+            self.sync_object_full.arrive(
+                state.index, self.producer_mask, loc=loc, ip=ip
+            )
+        else:
+            tx_count = self.sync_object_full.tx_count + extra_tx_count
+            self.sync_object_full.arrive_and_expect_tx(
+                state.index, tx_count, loc=loc, ip=ip
+            )
+
+
+PipelineTmaAsync.create = _override_create(PipelineTmaAsyncOg, PipelineTmaAsync)
+
+
+# ── PipelineTmaUmma ─────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PipelineTmaUmma(_PipelineIndexPhaseMixin, PipelineTmaUmmaOg):
+    """Override producer_acquire to take in extra_tx_count parameter."""
+
+    @dsl_user_op
+    def producer_acquire(
+        self,
+        state: PipelineState,
+        try_acquire_token: Optional[Boolean] = None,
+        extra_tx_count: int = 0,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        """
+        TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks.
+        """
+        if_generate(
+            try_acquire_token is None or try_acquire_token == 0,
+            lambda: self.sync_object_empty.wait(
+                state.index, state.phase, loc=loc, ip=ip
+            ),
+            loc=loc,
+            ip=ip,
+        )
+        if const_expr(extra_tx_count == 0):
+            if_generate(
+                self.is_leader_cta,
+                lambda: self.sync_object_full.arrive(
+                    state.index, self.producer_mask, loc=loc, ip=ip
+                ),
+                loc=loc,
+                ip=ip,
+            )
+        else:
+            tx_count = self.sync_object_full.tx_count + extra_tx_count
+            if_generate(
+                self.is_leader_cta,
+                lambda: self.sync_object_full.arrive_and_expect_tx(
+                    state.index, tx_count, loc=loc, ip=ip
+                ),
+                loc=loc,
+                ip=ip,
+            )
+
+
+PipelineTmaUmma.create = _override_create(PipelineTmaUmmaOg, PipelineTmaUmma)
+
+
+# ── PipelineUmmaAsync ───────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PipelineUmmaAsync(_PipelineIndexPhaseMixin, PipelineUmmaAsyncOg):
+    pass
+
+
+PipelineUmmaAsync.create = _override_create(PipelineUmmaAsyncOg, PipelineUmmaAsync)
+
+
+# ── PipelineAsyncUmma ───────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PipelineAsyncUmma(_PipelineIndexPhaseMixin, PipelineAsyncUmmaOg):
+    pass
+
+
+PipelineAsyncUmma.create = _override_create(PipelineAsyncUmmaOg, PipelineAsyncUmma)

python/sglang/jit_kernel/flash_attn/cute/pyproject.toml

@@ -0,0 +1,75 @@
+[build-system]
+requires = ["setuptools>=75", "setuptools-scm>=8"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "flash-attn-4"
+dynamic = ["version"]
+description = "Flash Attention CUTE (CUDA Template Engine) implementation"
+readme = "README.md"
+requires-python = ">=3.10"
+license = {text = "BSD 3-Clause License"}
+authors = [
+    {name = "Tri Dao"},
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "License :: OSI Approved :: BSD License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+]
+
+dependencies = [
+    "nvidia-cutlass-dsl>=4.5.2",
+    "torch",
+    "einops",
+    "typing_extensions",
+    "apache-tvm-ffi>=0.1.5,<0.2",
+    "torch-c-dlpack-ext",
+    "quack-kernels>=0.5.0",
+]
+
+[project.optional-dependencies]
+cu13 = ["nvidia-cutlass-dsl[cu13]>=4.5.2"]
+dev = [
+    "pytest",
+    "pytest-xdist",
+    "ruff",
+]
+
+[project.urls]
+Homepage = "https://github.com/Dao-AILab/flash-attention"
+Repository = "https://github.com/Dao-AILab/flash-attention"
+
+[tool.setuptools]
+packages = ["flash_attn.cute"]
+package-dir = {"flash_attn.cute" = "."}
+
+[tool.setuptools_scm]
+root = "../.."
+tag_regex = "^fa4-v(?P<version>.+)$"
+git_describe_command = "git describe --dirty --tags --long --match 'fa4-v*'"
+fallback_version = "0.0.0"
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+    { index = "pytorch-cu130", marker = "extra == 'cu13'" },
+]
+
+[tool.ruff]
+line-length = 100
+
+[tool.ruff.lint]
+ignore = [
+    "E731",  # do not assign a lambda expression, use a def
+    "E741",  # Do not use variables named 'I', 'O', or 'l'
+    "F841",  # local variable is assigned to but never used
+    "D102",  # Missing docstring in public methods
+]

python/sglang/jit_kernel/flash_attn/cute/seqlen_info.py

@@ -0,0 +1,331 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, const_expr
+from quack import copy_utils
+
+"""
+This consolidates all the info related to sequence length. This is so that we can do all
+the gmem reads once at the beginning of each tile, rather than having to repeat these reads
+to compute various things like n_block_min, n_block_max, etc.
+"""
+
+
+@dataclass(frozen=True)
+class SeqlenInfo:
+    offset: Int32
+    offset_padded: Int32
+    seqlen: Int32
+    has_cu_seqlens: cutlass.Constexpr[bool] = False
+
+    @staticmethod
+    def create(
+        batch_idx: Int32,
+        seqlen_static: Int32,
+        cu_seqlens: Optional[cute.Tensor] = None,
+        seqused: Optional[cute.Tensor] = None,
+        tile: cutlass.Constexpr[int] = 128,
+    ):
+        offset = 0 if const_expr(cu_seqlens is None) else cu_seqlens[batch_idx]
+        offset_padded = (
+            0
+            if const_expr(cu_seqlens is None)
+            # Add divby so that the compiler knows the alignment when moving by offset_padded
+            else cute.assume((offset + batch_idx * tile) // tile * tile, divby=tile)
+        )
+        if const_expr(seqused is not None):
+            seqlen = seqused[batch_idx]
+        elif const_expr(cu_seqlens is not None):
+            seqlen = cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx]
+        else:
+            seqlen = seqlen_static
+        return SeqlenInfo(
+            offset, offset_padded, seqlen, has_cu_seqlens=cu_seqlens is not None
+        )
+
+    def offset_batch(
+        self,
+        mT: cute.Tensor,
+        batch_idx: Int32,
+        dim: int,
+        padded: cutlass.Constexpr[bool] = False,
+        multiple: int = 1,
+    ) -> cute.Tensor:
+        """Offset a tensor by batch index. batch dim is at position `dim`, seqlen is at dim=0."""
+        if const_expr(not self.has_cu_seqlens):
+            idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mT) - 1 - dim)
+            return mT[idx]
+        else:
+            off = multiple * (
+                self.offset if const_expr(not padded) else self.offset_padded
+            )
+            offset = off if const_expr(cute.rank(mT.shape[0]) == 1) else (0, off)
+            idx = (offset,) + (None,) * (cute.rank(mT) - 1)
+            return cute.domain_offset(idx, mT)
+
+
+@dataclass(frozen=True)
+class SeqlenInfoQK:
+    offset_q: Int32
+    offset_k: Int32
+    padded_offset_q: Int32
+    padded_offset_k: Int32
+    seqlen_q: Int32
+    seqlen_k: Int32
+    m_block_offset: Int32
+    block_idx_offset: Int32
+    num_n_blocks: Int32
+    has_cu_seqlens_q: cutlass.Constexpr[bool]
+    has_cu_seqlens_k: cutlass.Constexpr[bool]
+    has_seqused_q: cutlass.Constexpr[bool]
+    has_seqused_k: cutlass.Constexpr[bool]
+
+    @staticmethod
+    def create(
+        batch_idx: Int32,
+        seqlen_q_static: Int32,
+        seqlen_k_static: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        mCuBlockIdxOffsets: Optional[cute.Tensor] = None,
+        tile_m: cutlass.Constexpr[Int32] = 128,
+        tile_n: cutlass.Constexpr[Int32] = 128,
+    ):
+        offset_q = 0 if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ[batch_idx]
+        offset_k = 0 if const_expr(mCuSeqlensK is None) else mCuSeqlensK[batch_idx]
+        padded_offset_q = (
+            0
+            if const_expr(mCuSeqlensQ is None)
+            else cute.assume(
+                (offset_q + batch_idx * tile_m) // tile_m * tile_m, divby=tile_m
+            )
+        )
+        padded_offset_k = (
+            0
+            if const_expr(mCuSeqlensK is None)
+            else cute.assume(
+                (offset_k + batch_idx * tile_n) // tile_n * tile_n, divby=tile_n
+            )
+        )
+        if const_expr(mSeqUsedQ is not None):
+            seqlen_q = mSeqUsedQ[batch_idx]
+        else:
+            seqlen_q = (
+                seqlen_q_static
+                if const_expr(mCuSeqlensQ is None)
+                else mCuSeqlensQ[batch_idx + 1] - offset_q
+            )
+        if const_expr(mSeqUsedK is not None):
+            seqlen_k = mSeqUsedK[batch_idx]
+        else:
+            seqlen_k = (
+                seqlen_k_static
+                if const_expr(mCuSeqlensK is None)
+                else mCuSeqlensK[batch_idx + 1] - offset_k
+            )
+        m_block_offset = (
+            0 if const_expr(mCuTotalMBlocks is None) else mCuTotalMBlocks[batch_idx]
+        )
+        num_n_blocks = (seqlen_k + tile_n - 1) // tile_n
+        block_idx_offset = (
+            mCuBlockIdxOffsets[batch_idx]
+            if const_expr(mCuBlockIdxOffsets is not None)
+            else m_block_offset * num_n_blocks
+        )
+        return SeqlenInfoQK(
+            offset_q,
+            offset_k,
+            padded_offset_q,
+            padded_offset_k,
+            seqlen_q,
+            seqlen_k,
+            m_block_offset,
+            block_idx_offset,
+            num_n_blocks,
+            has_cu_seqlens_q=mCuSeqlensQ is not None,
+            has_cu_seqlens_k=mCuSeqlensK is not None,
+            has_seqused_q=mSeqUsedQ is not None,
+            has_seqused_k=mSeqUsedK is not None,
+        )
+
+    def offset_batch_Q(
+        self,
+        mQ: cute.Tensor,
+        batch_idx: Int32,
+        dim: int,
+        padded: cutlass.Constexpr[bool] = False,
+        ragged: cutlass.Constexpr[bool] = False,
+    ) -> cute.Tensor:
+        """Seqlen must be the first dimension of mQ"""
+        if const_expr(not ragged):
+            if const_expr(not self.has_cu_seqlens_q):
+                idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim)
+                return mQ[idx]
+            else:
+                offset_q = (
+                    self.offset_q if const_expr(not padded) else self.padded_offset_q
+                )
+                offset_q = (
+                    offset_q
+                    if const_expr(cute.rank(mQ.shape[0]) == 1)
+                    else (None, offset_q)
+                )
+                idx = (offset_q,) + (None,) * (cute.rank(mQ) - 1)
+                return cute.domain_offset(idx, mQ)
+        else:
+            if const_expr(not self.has_cu_seqlens_q):
+                offset_q = 0
+                idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim)
+                mQ = mQ[idx]
+            else:
+                offset_q = (
+                    self.offset_q if const_expr(not padded) else self.padded_offset_q
+                )
+            if const_expr(cute.rank(mQ.shape[0]) == 1):
+                return copy_utils.offset_ragged_tensor(
+                    mQ, offset_q, self.seqlen_q, ragged_dim=0, ptr_shift=True
+                )
+            else:  # PackGQA
+                assert cute.rank(mQ.shape[0]) == 2
+                # Unpack before calling offset_ragged_tensor, then pack
+                idx = ((None, None),) + (None,) * (cute.rank(mQ) - 1)
+                mQ = mQ[idx]
+                mQ = copy_utils.offset_ragged_tensor(
+                    mQ, offset_q, self.seqlen_q, ragged_dim=1, ptr_shift=True
+                )
+                return cute.group_modes(mQ, 0, 2)
+
+    def offset_batch_K(
+        self,
+        mK: cute.Tensor,
+        batch_idx: Int32,
+        dim: int,
+        padded: cutlass.Constexpr[bool] = False,
+        ragged: cutlass.Constexpr[bool] = False,
+        multiple: int = 1,
+    ) -> cute.Tensor:
+        """Seqlen must be the first dimension of mK"""
+        if const_expr(not ragged):
+            if const_expr(not self.has_cu_seqlens_k):
+                idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim)
+                return mK[idx]
+            else:
+                offset_k = (
+                    self.offset_k if const_expr(not padded) else self.padded_offset_k
+                )
+                offset_k *= multiple
+                idx = (offset_k,) + (None,) * (cute.rank(mK) - 1)
+                return cute.domain_offset(idx, mK)
+        else:
+            if const_expr(not self.has_cu_seqlens_k):
+                offset_k = 0
+                idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim)
+                mK = mK[idx]
+            else:
+                offset_k = (
+                    self.offset_k if const_expr(not padded) else self.padded_offset_k
+                )
+                offset_k *= multiple
+            return copy_utils.offset_ragged_tensor(
+                mK, offset_k, self.seqlen_k, ragged_dim=0, ptr_shift=True
+            )
+
+
+@dataclass(frozen=True)
+class SeqlenInfoQKNewK:
+    """Sequence length info for append-KV with left-padding and new K support.
+
+    Extends SeqlenInfoQK with:
+    - leftpad_k: left padding for K (tokens to skip at the start of the KV cache)
+    - offset_k_new: offset into the new K tensor
+    - seqlen_k_og: original K length (before appending new K), excluding leftpad
+    - seqlen_k_new: length of new K to append
+    - seqlen_k: total K length (seqlen_k_og + seqlen_k_new)
+    - seqlen_rotary: position for rotary embedding computation
+    """
+
+    leftpad_k: Int32
+    offset_q: Int32
+    offset_k: Int32
+    offset_k_new: Int32
+    seqlen_q: Int32
+    seqlen_k_og: Int32
+    seqlen_k_new: Int32
+    seqlen_k: Int32
+    seqlen_rotary: Int32
+
+    @staticmethod
+    def create(
+        batch_idx: Int32,
+        seqlen_q_static: Int32,
+        seqlen_k_static: Int32,
+        shape_K_new_0: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mCuSeqlensKNew: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        mLeftpadK: Optional[cute.Tensor] = None,
+        mSeqlensRotary: Optional[cute.Tensor] = None,
+    ):
+        leftpad_k = 0 if const_expr(mLeftpadK is None) else mLeftpadK[batch_idx]
+        offset_q = 0 if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ[batch_idx]
+        if const_expr(mCuSeqlensK is not None):
+            offset_k = mCuSeqlensK[batch_idx] + leftpad_k
+        else:
+            offset_k = leftpad_k if const_expr(mCuSeqlensQ is not None) else 0
+        offset_k_new = (
+            0 if const_expr(mCuSeqlensKNew is None) else mCuSeqlensKNew[batch_idx]
+        )
+        # seqlen_q
+        if const_expr(mSeqUsedQ is not None):
+            seqlen_q = mSeqUsedQ[batch_idx]
+        elif const_expr(mCuSeqlensQ is not None):
+            seqlen_q = mCuSeqlensQ[batch_idx + 1] - mCuSeqlensQ[batch_idx]
+        else:
+            seqlen_q = seqlen_q_static
+        # seqlen_k_og: original K length (excluding leftpad)
+        if const_expr(mSeqUsedK is not None):
+            seqlen_k_og = mSeqUsedK[batch_idx] - leftpad_k
+        elif const_expr(mCuSeqlensK is not None):
+            seqlen_k_og = (
+                mCuSeqlensK[batch_idx + 1] - mCuSeqlensK[batch_idx] - leftpad_k
+            )
+        else:
+            seqlen_k_og = (
+                seqlen_k_static - leftpad_k
+                if const_expr(mCuSeqlensQ is not None)
+                else seqlen_k_static
+            )
+        # seqlen_k_new
+        if const_expr(mCuSeqlensKNew is None):
+            seqlen_k_new = 0 if const_expr(mCuSeqlensQ is None) else shape_K_new_0
+        else:
+            seqlen_k_new = mCuSeqlensKNew[batch_idx + 1] - mCuSeqlensKNew[batch_idx]
+        seqlen_k = (
+            seqlen_k_og
+            if const_expr(mCuSeqlensQ is None)
+            else seqlen_k_og + seqlen_k_new
+        )
+
+        # seqlen_rotary: defaults to seqlen_k_og + leftpad_k unless explicitly provided
+        if const_expr(mSeqlensRotary is not None):
+            seqlen_rotary = mSeqlensRotary[batch_idx]
+        else:
+            seqlen_rotary = seqlen_k_og + leftpad_k
+        return SeqlenInfoQKNewK(
+            leftpad_k,
+            offset_q,
+            offset_k,
+            offset_k_new,
+            seqlen_q,
+            seqlen_k_og,
+            seqlen_k_new,
+            seqlen_k,
+            seqlen_rotary,
+        )

python/sglang/jit_kernel/flash_attn/cute/shearing_bias.py

@@ -0,0 +1,573 @@
+# Copyright (c) 2026, Colfax International.
+
+import math
+from functools import partial
+from typing import Callable, Optional
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, const_expr
+
+from sglang.jit_kernel.flash_attn.cute.block_info import BlockInfo
+from sglang.jit_kernel.flash_attn.cute.copy_utils import tiled_copy_2d
+from sglang.jit_kernel.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
+from sglang.jit_kernel.flash_attn.cute.pack_gqa import pack_gqa_layout
+from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK
+from sglang.jit_kernel.flash_attn.cute.tile_scheduler import (
+    ParamsBase,
+    SingleTileScheduler,
+    SingleTileVarlenScheduler,
+    TileSchedulerArguments,
+)
+from sglang.jit_kernel.flash_attn.cute.utils import get_batch_from_cu_tensor
+
+
+class ShearingBias:
+    def __init__(
+        self,
+        rel_extent: int = 512,
+        is_causal: bool = True,
+        is_local: bool = False,
+        pack_gqa: bool = False,
+        qhead_per_kvhead: cutlass.Constexpr[int] = 1,
+        rows_per_cta: int = 4,
+        tile_m: int = 128,
+        max_m_blocks_leq_one: bool = False,
+        use_pdl: bool = False,
+        clamp_subtiles: bool = True,
+    ):
+        self.is_causal = is_causal
+        self.is_local = is_local
+        assert is_causal or is_local, "Doesn't make sense otherwise"
+        self.pack_gqa = pack_gqa
+        self.qhead_per_kvhead = qhead_per_kvhead
+        if self.pack_gqa:
+            assert (
+                128 % self.qhead_per_kvhead == 0
+            ), "pack_gqa only supported when qhead_per_kvhead divides 128"
+        self.qhead_per_kvhead_packgqa = qhead_per_kvhead if self.pack_gqa else 1
+        self.rel_extent = rel_extent
+        assert rel_extent % 128 == 0
+        self.rel_extent_padded = rel_extent + 256
+        self.num_bias_blocks_padded = (self.rel_extent_padded) // 128
+        # tuneable parameters
+        assert rows_per_cta % 4 == 0
+        self.rows_per_cta = rows_per_cta
+        self.num_threads = self.rows_per_cta * 32
+        self.cta_tiler = (self.rows_per_cta, self.rel_extent)
+        self.cta_out_tiler = (self.rows_per_cta, self.rel_extent_padded)
+
+        self.buffer_align_bytes = 1024
+
+        self.max_m_blocks_leq_one = max_m_blocks_leq_one
+        self.use_pdl = use_pdl
+
+        # only used with block packed scheduling
+        self.tile_m = tile_m
+        # Shrink the subtile grid dim to the rows a block can actually hold
+        # (decode blocks hold qhead_per_kvhead*seqlen_q rows, not tile_m).
+        self.clamp_subtiles = clamp_subtiles
+
+    @cute.jit
+    def __call__(
+        self,
+        mPreBias: cute.Tensor,  # (b, s_q, h, rel_extent) or (total_q, h, rel_extent)
+        mBias: cute.Tensor,  # (b, s_q, h, rel_extent_padded) or (total_q, h, rel_extent_padded)
+        max_seqlen_q: Int32 | int,
+        max_seqlen_k: Int32 | int,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        mBlocksToBatchIdx: Optional[cute.Tensor] = None,
+        window_size_left: Int32 | int | None = None,
+        window_size_right: Int32 | int | None = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
+    ):
+        assert mPreBias.element_type == mBias.element_type
+        self.bias_dtype = mBias.element_type
+
+        right_pad_value = -Float32.inf
+        left_pad_value = (
+            -Float32.inf if const_expr(window_size_left is not None) else 0.0
+        )
+
+        self.vec_size = 32 // self.bias_dtype.width
+        self.cols_per_iter = 32 * self.vec_size
+        assert self.vec_size <= 2
+        assert 128 % self.cols_per_iter == 0
+
+        max_seqlen_k = Int32(max_seqlen_k)
+        if const_expr(window_size_left is not None):
+            window_size_left = Int32(window_size_left)
+        if const_expr(window_size_right is not None):
+            window_size_right = Int32(window_size_right)
+
+        mPreBias, mBias = [assume_tensor_aligned(t) for t in (mPreBias, mBias)]
+        # (s_q, rel_extent, h, b) or (total_q, rel_extent, h)
+        Q_layout_transpose = (
+            [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1]
+        )
+        mPreBias, mBias = [
+            cute.make_tensor(t.iterator, cute.select(t.layout, mode=Q_layout_transpose))
+            for t in (mPreBias, mBias)
+        ]
+
+        if const_expr(self.pack_gqa):
+            nheads_kv = mBias.shape[2] // self.qhead_per_kvhead
+            mPreBias, mBias = [
+                pack_gqa_layout(t, self.qhead_per_kvhead, nheads_kv, head_idx=2)
+                for t in (mPreBias, mBias)
+            ]
+
+        # SMEM layouts
+        prebias_tile_shape = (self.rows_per_cta, self.rel_extent)
+        bias_tile_shape = (self.rows_per_cta, self.rel_extent_padded)
+
+        sPreBias_layout = cute.make_ordered_layout(prebias_tile_shape, order=(1, 0))
+        sBias_layout = cute.make_ordered_layout(bias_tile_shape, order=(1, 0))
+        sPreBias_size = cute.cosize(sPreBias_layout)
+        sBias_size = cute.cosize(sBias_layout)
+
+        in_major_size = math.gcd(256, self.rel_extent)
+        assert in_major_size % 128 == 0
+        self.num_g2s_threads = self.num_threads if in_major_size == 256 else 128
+        g2s_tiled_copy = tiled_copy_2d(
+            self.bias_dtype,
+            math.gcd(256, self.rel_extent),
+            self.num_g2s_threads,
+            is_async=True,
+        )
+
+        out_major_size = math.gcd(256, self.rel_extent_padded)
+        assert out_major_size % 128 == 0
+        self.num_s2g_threads = self.num_threads if out_major_size == 256 else 128
+        s2g_tiled_copy = tiled_copy_2d(
+            self.bias_dtype,
+            math.gcd(256, self.rel_extent_padded),
+            self.num_s2g_threads,
+        )
+
+        @cute.struct
+        class SharedStorage:
+            sPreBias: cute.struct.Align[
+                cute.struct.MemRange[self.bias_dtype, sPreBias_size],
+                self.buffer_align_bytes,
+            ]
+            sBias: cute.struct.Align[
+                cute.struct.MemRange[self.bias_dtype, sBias_size],
+                self.buffer_align_bytes,
+            ]
+
+        self.shared_storage = SharedStorage
+
+        varlen_q = mCuSeqlensQ is not None or mSeqUsedQ is not None
+        self.use_block_packed_scheduling = (
+            mCuTotalMBlocks is not None
+            and mCuSeqlensQ is not None
+            and not self.max_m_blocks_leq_one
+            # and False
+        )
+
+        if const_expr(varlen_q and not self.max_m_blocks_leq_one):
+            if const_expr(self.use_block_packed_scheduling):
+                TileScheduler = SingleTileScheduler
+            else:
+                TileScheduler = SingleTileVarlenScheduler
+        else:
+            TileScheduler = SingleTileScheduler
+
+        batch_size = (
+            cute.size(mPreBias.shape[3])
+            if const_expr(mCuSeqlensQ is None)
+            else cute.size(mCuSeqlensQ.shape[0] - 1)
+        )
+        eff_seqlen_q = (
+            max_seqlen_q
+            if const_expr(not self.pack_gqa)
+            else max_seqlen_q * self.qhead_per_kvhead
+        )
+        total_q = (
+            cute.size(mPreBias.shape[0])
+            if const_expr(mCuSeqlensQ is not None)
+            else cute.size(mPreBias.shape[0]) * cute.size(mPreBias.shape[3])
+        )
+        # same formula as in varlen scheduler -- only used with block packed scheduling
+        total_blocks_max = (total_q + batch_size * (self.tile_m - 1)) // self.tile_m
+
+        num_blocks_for_sched = (
+            cute.ceil_div(eff_seqlen_q, self.rows_per_cta)
+            if const_expr(not self.use_block_packed_scheduling)
+            else total_blocks_max
+        )
+        if const_expr(not self.use_block_packed_scheduling):
+            batch_size_for_sched = batch_size
+        elif const_expr(self.clamp_subtiles):
+            # A block covers at most min(tile_m, eff_seqlen_q) valid rows; subtiles
+            # past that would fail the per-row seqlen guards and exit immediately.
+            batch_size_for_sched = cute.ceil_div(
+                min(self.tile_m, eff_seqlen_q), self.rows_per_cta
+            )
+        else:
+            batch_size_for_sched = self.tile_m // self.rows_per_cta
+
+        tile_sched_args = TileSchedulerArguments(
+            num_blocks_for_sched,
+            cute.size(mPreBias.shape[2]),
+            batch_size_for_sched,
+            1,
+            1,
+            1,
+            1,
+            total_q=total_q,
+            tile_shape_mn=self.cta_tiler,
+            mCuSeqlensQ=mCuSeqlensQ,
+            mSeqUsedQ=mSeqUsedQ,
+            qhead_per_kvhead_packgqa=self.qhead_per_kvhead_packgqa,
+            element_size=self.bias_dtype.width // 8,
+        )
+        tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
+        self.tile_scheduler_cls = TileScheduler
+        grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
+
+        self.kernel(
+            mPreBias,
+            mBias,
+            left_pad_value,
+            right_pad_value,
+            max_seqlen_k,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mSeqUsedQ,
+            mSeqUsedK,
+            mCuTotalMBlocks,
+            mBlocksToBatchIdx,
+            sPreBias_layout,
+            sBias_layout,
+            window_size_left,
+            window_size_right,
+            g2s_tiled_copy,
+            s2g_tiled_copy,
+            SharedStorage,
+            tile_sched_params,
+        ).launch(
+            grid=grid_dim,
+            block=(self.num_threads, 1, 1),
+            stream=stream,
+            use_pdl=self.use_pdl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mPreBias: cute.Tensor,
+        mBias: cute.Tensor,
+        left_pad_value: cutlass.Float32,
+        right_pad_value: cutlass.Float32,
+        max_seqlen_k: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mCuSeqlensK: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        mSeqUsedK: Optional[cute.Tensor],
+        mCuTotalMBlocks: Optional[cute.Tensor],
+        mBlocksToBatchIdx: Optional[cute.Tensor],
+        sPreBias_layout: cute.ComposedLayout | cute.Layout,
+        sBias_layout: cute.ComposedLayout | cute.Layout,
+        window_size_left: Optional[Int32],
+        window_size_right: Optional[Int32],
+        g2s_tiled_copy: cute.TiledCopy,
+        s2g_tiled_copy: cute.TiledCopy,
+        SharedStorage: cutlass.Constexpr[Callable],
+        tile_sched_params: ParamsBase,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        lane_idx = cute.arch.lane_idx()
+
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+
+        sPreBias = storage.sPreBias.get_tensor(sPreBias_layout)
+        sBias = storage.sBias.get_tensor(sBias_layout)
+
+        TileSchedulerCls = partial(self.tile_scheduler_cls.create, tile_sched_params)
+        tile_scheduler = TileSchedulerCls()
+        work_tile = tile_scheduler.initial_work_tile_info()
+        # if pack_gqa, head_idx means head_idx_kv
+        m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+        subtile_idx = batch_idx if const_expr(self.use_block_packed_scheduling) else 0
+
+        if const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+            cute.arch.griddepcontrol_launch_dependents()
+
+        is_valid_tile = work_tile.is_valid_tile
+        if const_expr(self.use_block_packed_scheduling):
+            batch_size = mCuTotalMBlocks.shape[0] - 1
+            is_valid_tile = m_block < mCuTotalMBlocks[batch_size]
+
+        if is_valid_tile:
+            if const_expr(self.use_block_packed_scheduling):
+                if const_expr(mBlocksToBatchIdx is not None):
+                    batch_idx = mBlocksToBatchIdx[m_block]
+                else:
+                    batch_idx = get_batch_from_cu_tensor(m_block, mCuTotalMBlocks)
+                # get local m_block for batch
+                m_block -= mCuTotalMBlocks[batch_idx]
+                m_block = m_block * (self.tile_m // self.rows_per_cta) + subtile_idx
+            seqlen_info = SeqlenInfoQK.create(
+                batch_idx=batch_idx,
+                seqlen_q_static=(
+                    mPreBias.shape[0]
+                    if const_expr(not self.pack_gqa)
+                    else mPreBias.shape[0][1]
+                ),
+                seqlen_k_static=max_seqlen_k,
+                mCuSeqlensQ=mCuSeqlensQ,
+                mCuSeqlensK=mCuSeqlensK,
+                mSeqUsedQ=mSeqUsedQ,
+                mSeqUsedK=mSeqUsedK,
+            )
+
+            block_info = BlockInfo(
+                128,
+                128,
+                self.is_causal,
+                self.is_local,
+                window_size_left=window_size_left,
+                window_size_right=window_size_right,
+                qhead_per_kvhead_packgqa=self.qhead_per_kvhead_packgqa,
+            )
+
+            # (seqlen, rel_extent) or ((seqlen, qhead_per_kvhead), rel_extent)
+            mPreBias_cur = seqlen_info.offset_batch_Q(mPreBias, batch_idx, dim=3)[
+                None, None, head_idx
+            ]
+            # (rows_per_cta, rel_extent)
+            gPreBias = cute.local_tile(mPreBias_cur, self.cta_tiler, (m_block, 0))
+            cPreBias = cute.make_identity_tensor(self.cta_tiler)
+
+            g2s_thr_copy = g2s_tiled_copy.get_slice(tidx)
+
+            # (V, M, N)
+            tBgPreBias = g2s_thr_copy.partition_S(gPreBias)
+            tBsPreBias = g2s_thr_copy.partition_D(sPreBias)
+            tBcPreBias = g2s_thr_copy.partition_S(cPreBias)
+
+            if (
+                const_expr(self.num_g2s_threads == self.num_threads)
+                or warp_idx < self.num_g2s_threads // 32
+            ):
+                num_rows_per_load = tBgPreBias.shape[1]
+                for m in cutlass.range_constexpr(num_rows_per_load):
+                    local_m_idx = tBcPreBias[0, m, 0][0]
+                    load_m_idx = local_m_idx + m_block * self.rows_per_cta
+                    local_m_idx_in_bounds = (
+                        const_expr(self.rows_per_cta % 8 == 0)
+                        or local_m_idx < self.rows_per_cta
+                    )
+                    load_m_idx_in_bounds = (
+                        load_m_idx // self.qhead_per_kvhead_packgqa
+                        < seqlen_info.seqlen_q
+                    )
+                    if local_m_idx_in_bounds and load_m_idx_in_bounds:
+                        cute.copy(
+                            g2s_tiled_copy,
+                            tBgPreBias[None, m, None],
+                            tBsPreBias[None, m, None],
+                        )
+
+            cute.arch.cp_async_commit_group()
+
+            # Convention: inclusive min, exclusive max
+            m_idx = m_block * self.rows_per_cta + warp_idx
+            attn_m_block = m_idx // 128
+
+            _, attn_n_block_max = block_info.get_n_block_min_max(
+                seqlen_info,
+                attn_m_block,
+            )
+
+            n_idx_left, n_idx_right = block_info.get_n_idx_left_right(
+                seqlen_info, m_idx
+            )
+            num_bias_vals = n_idx_right - max(n_idx_left, n_idx_right - self.rel_extent)
+            is_even = n_idx_right % 2 == 0
+
+            # get bias block and idx bounds for row
+            n_block_for_rel0 = (n_idx_right - 1) // 128  # inclusive
+            bias_block_idx_right = 1 + max(
+                self.rel_extent_padded // 128 - (attn_n_block_max - n_block_for_rel0), 0
+            )
+            bias_idx_right = (
+                (bias_block_idx_right - 1) * 128 + ((n_idx_right - 1) % 128) + 1
+            )
+            bias_idx_left = max(0, bias_idx_right - num_bias_vals)
+            bias_block_idx_left = bias_idx_left // 128
+            # num_bias_blocks = self.num_bias_blocks_padded - bias_block_idx_left
+            # num_right_padding_blocks = 0
+            num_bias_blocks = (
+                bias_block_idx_right - bias_block_idx_left if num_bias_vals > 0 else 0
+            )
+            num_right_padding_blocks = (
+                self.num_bias_blocks_padded - bias_block_idx_right
+                if num_bias_vals > 0
+                else self.num_bias_blocks_padded
+            )
+            # might help compiler unroll loops
+            num_bias_blocks = min(num_bias_blocks, self.num_bias_blocks_padded)
+            num_right_padding_blocks = min(
+                num_right_padding_blocks, self.num_bias_blocks_padded
+            )
+
+            sPreBias_row = cute.flat_divide(
+                sPreBias[(warp_idx, None)], (self.vec_size,)
+            )
+            sBias_row = cute.flat_divide(sBias[(warp_idx, None)], (self.vec_size,))
+            sBias_row_vec4 = cute.flat_divide(sBias[(warp_idx, None)], (4,))
+
+            bias_idx = (
+                self.rel_extent_padded + lane_idx * self.vec_size - self.cols_per_iter
+            )
+
+            cute.arch.cp_async_wait_group(0)
+            cute.arch.sync_threads()
+
+            if m_idx // self.qhead_per_kvhead_packgqa < seqlen_info.seqlen_q:
+                # We can try handling right padding separately
+                for i in cutlass.range(num_right_padding_blocks, unroll_full=True):
+                    bias_frg = cute.make_rmem_tensor((4,), dtype=self.bias_dtype)
+                    bias_frg.fill(self.bias_dtype(right_pad_value))
+                    bias_right_pad_idx = (
+                        self.num_bias_blocks_padded - 1 - i
+                    ) * 32 + lane_idx
+                    cute.autovec_copy(
+                        bias_frg, sBias_row_vec4[None, bias_right_pad_idx]
+                    )
+                    bias_idx -= 128
+
+                for _ in cutlass.range(num_bias_blocks, unroll_full=True):
+                    # 2 subblocks for half bias dtype
+                    for _ in cutlass.range_constexpr(128 // self.cols_per_iter):
+                        prebias_idx = bias_idx_right - 1 - bias_idx
+
+                        # (vec_size, lower/upper)
+                        prebias_frg = cute.make_rmem_tensor(
+                            (self.vec_size, self.vec_size), dtype=self.bias_dtype
+                        )
+
+                        in_bounds = (
+                            prebias_idx >= 0
+                            and prebias_idx - self.vec_size + 1 < num_bias_vals
+                        )
+                        prebias_idx_lower = (
+                            prebias_idx - 1 if is_even else max(prebias_idx - 2, 0)
+                        )
+                        prebias_idx_upper = (
+                            prebias_idx - 1
+                            if is_even
+                            else min(prebias_idx, self.rel_extent - 2)
+                        )
+
+                        if in_bounds:
+                            cute.autovec_copy(
+                                sPreBias_row[None, prebias_idx_lower // 2],
+                                prebias_frg[None, 0],
+                            )
+                            if const_expr(self.vec_size == 2) and not is_even:
+                                cute.autovec_copy(
+                                    sPreBias_row[None, prebias_idx_upper // 2],
+                                    prebias_frg[None, 1],
+                                )
+
+                        bias_frg = cute.make_rmem_tensor(
+                            (self.vec_size,), dtype=self.bias_dtype
+                        )
+                        bias_frg.fill(self.bias_dtype(left_pad_value))
+
+                        if const_expr(self.vec_size == 1):
+                            if in_bounds:
+                                bias_frg[0] = prebias_frg[0, 0]
+                            elif prebias_idx < 0:
+                                bias_frg.fill(self.bias_dtype(right_pad_value))
+                        else:
+                            if in_bounds:
+                                if is_even:
+                                    # reverse: [prebias_idx, prebias_idx-1] = bias frg
+                                    bias_frg[0] = prebias_frg[1, 0]
+                                    bias_frg[1] = prebias_frg[0, 0]
+                                else:
+                                    # lower = [2x-2, 2x-1], upper = [2x, 2x+1], 2x = prebias_idx
+                                    # want bias = [2x, 2x-1]
+                                    bias_frg[0] = prebias_frg[0, 1]
+                                    bias_frg[1] = prebias_frg[1, 0]
+                            elif prebias_idx < 0:
+                                bias_frg.fill(self.bias_dtype(right_pad_value))
+
+                        cute.autovec_copy(bias_frg, sBias_row[None, bias_idx // 2])
+                        bias_idx -= self.cols_per_iter
+                        cute.arch.sync_warp()
+
+                # Handle edge cases. For N = rel_extent:
+                # [0, -1], -1 at bias_idx_right and [N, N-1], N-1 at bias_idx_left
+                if not is_even and num_bias_vals > 0:
+                    sBias[(warp_idx, bias_idx_right)] = self.bias_dtype(right_pad_value)
+                    if bias_idx_left - 1 >= 0:
+                        sBias[(warp_idx, bias_idx_left - 1)] = self.bias_dtype(
+                            left_pad_value
+                        )
+
+                num_left_padding_blocks = min(
+                    self.num_bias_blocks_padded
+                    - num_bias_blocks
+                    - num_right_padding_blocks,
+                    self.num_bias_blocks_padded,
+                )
+                for i in cutlass.range(num_left_padding_blocks, unroll_full=True):
+                    bias_left_pad_idx = i * 32 + lane_idx
+                    bias_frg = cute.make_rmem_tensor((4,), dtype=self.bias_dtype)
+                    bias_frg.fill(self.bias_dtype(left_pad_value))
+                    cute.autovec_copy(bias_frg, sBias_row_vec4[None, bias_left_pad_idx])
+
+            cute.arch.sync_threads()
+
+            s2g_thr_copy = s2g_tiled_copy.get_slice(tidx)
+
+            # (seqlen, rel_extent_padded)
+            mBias_cur = seqlen_info.offset_batch_Q(mBias, batch_idx, dim=3)[
+                None, None, head_idx
+            ]
+            # (rows_per_cta, rel_extent_padded)
+            gBias = cute.local_tile(mBias_cur, self.cta_out_tiler, (m_block, 0))
+            cBias = cute.make_identity_tensor(self.cta_out_tiler)
+
+            # (V, M, N)
+            tBsBias = s2g_thr_copy.partition_S(sBias)
+            tBgBias = s2g_thr_copy.partition_D(gBias)
+            tBcBias = s2g_thr_copy.partition_D(cBias)
+
+            if (
+                const_expr(self.num_s2g_threads == self.num_threads)
+                or warp_idx < self.num_s2g_threads // 32
+            ):
+                num_rows_per_store = tBgBias.shape[1]
+                for m in cutlass.range_constexpr(num_rows_per_store):
+                    local_m_idx = tBcBias[0, m, 0][0]
+                    store_m_idx = local_m_idx + m_block * self.rows_per_cta
+                    local_m_idx_in_bounds = (
+                        const_expr(self.rows_per_cta % 8 == 0)
+                        or local_m_idx < self.rows_per_cta
+                    )
+                    store_m_idx_in_bounds = (
+                        store_m_idx // self.qhead_per_kvhead_packgqa
+                        < seqlen_info.seqlen_q
+                    )
+                    if local_m_idx_in_bounds and store_m_idx_in_bounds:
+                        cute.copy(
+                            s2g_tiled_copy,
+                            tBsBias[None, m, None],
+                            tBgBias[None, m, None],
+                        )

python/sglang/jit_kernel/flash_attn/cute/sm90_config_search.py

@@ -0,0 +1,416 @@
+"""Search feasible SM90 fwd/bwd attention configs for given (head_dim, head_dim_v).
+
+Enumerates tile sizes, swap modes, atom layouts, and staging options.
+Checks GMMA divisibility, register budget, and shared memory budget.
+
+Usage:
+    python flash_attn/cute/sm90_config_search.py --headdim 128
+    python flash_attn/cute/sm90_config_search.py --mode fwd --headdim 192-128
+    python flash_attn/cute/sm90_config_search.py --mode bwd --headdim 192 --tile-n 64,96
+"""
+
+import math
+
+# H100 hardware limits
+SMEM_LIMIT = 224 * 1024  # 228 KB minus ~3 KB for LSE, dPsum, mbarriers
+REG_LIMITS = {2: 216, 3: 128}  # per-WG budget: 2WG=240-24, 3WG=160-32
+THREADS_PER_WG = 128
+
+
+def _divisors(n):
+    return [d for d in range(1, n + 1) if n % d == 0]
+
+
+def _acc_regs(M, N, num_wg):
+    """Accumulator registers per thread per WG."""
+    return M * N // (num_wg * THREADS_PER_WG)
+
+
+def _check_mma(M, N, num_wg, atom_layout_m, swap_AB):
+    """Check MMA feasibility. Returns regs per WG, or None if infeasible.
+
+    GMMA atom M=64. Swap exchanges (M, N) and atom layout.
+    Requires: M divisible by (atom_layout_m * 64), N by (atom_layout_n * 8).
+    """
+    if swap_AB:
+        M, N = N, M
+        atom_layout_m = num_wg // atom_layout_m
+    atom_layout_n = num_wg // atom_layout_m
+    if M % (atom_layout_m * 64) != 0 or N % (atom_layout_n * 8) != 0:
+        return None
+    return _acc_regs(M, N, num_wg)
+
+
+def _mma_traffic(M_eff, N_eff, K_red, num_wg, wg_n, is_rs=False):
+    """Total SMEM read traffic for one MMA (all WGs combined).
+
+    num_instr = (M_eff / 64) * wg_n instructions total.
+    Each reads A(64, K_red) and B(N_eff/wg_n, K_red) from smem (bf16).
+    """
+    num_instr = (M_eff // 64) * wg_n
+    A_per = 64 * K_red * 2 if not is_rs else 0
+    B_per = (N_eff // wg_n) * K_red * 2
+    return num_instr * (A_per + B_per)
+
+
+# ============================================================================
+# Backward
+# ============================================================================
+
+
+def _check_bwd_config(
+    hdim,
+    hdimv,
+    tile_m,
+    tile_n,
+    num_wg,
+    SdP_swapAB,
+    dKV_swapAB,
+    dQ_swapAB,
+    AtomLayoutMSdP,
+    AtomLayoutNdKV,
+    AtomLayoutMdQ,
+):
+    reg_limit = REG_LIMITS[num_wg]
+
+    # MMA feasibility
+    regs_SdP = _check_mma(tile_m, tile_n, num_wg, AtomLayoutMSdP, SdP_swapAB)
+    regs_dK = _check_mma(tile_n, hdim, num_wg, AtomLayoutNdKV, dKV_swapAB)
+    regs_dV = _check_mma(tile_n, hdimv, num_wg, AtomLayoutNdKV, dKV_swapAB)
+    regs_dQ = _check_mma(tile_m, hdim, num_wg, AtomLayoutMdQ, dQ_swapAB)
+    if any(r is None for r in (regs_SdP, regs_dK, regs_dV, regs_dQ)):
+        return None
+
+    # Peak regs: max(S+dP, dQ) + dK + dV
+    total_regs = max(2 * regs_SdP, regs_dQ) + regs_dK + regs_dV
+    if total_regs > reg_limit:
+        return None
+
+    # SMEM
+    mma_dkv_is_rs = (
+        AtomLayoutMSdP == 1
+        and AtomLayoutNdKV == num_wg
+        and SdP_swapAB
+        and not dKV_swapAB
+    )
+    Q_stage, PdS_stage = 2, 1
+
+    for dO_stage in (2, 1):
+        sQ = tile_m * hdim * 2 * Q_stage
+        sK = tile_n * hdim * 2
+        sV = tile_n * hdimv * 2
+        sdO = tile_m * hdimv * 2 * dO_stage
+        sPdS = tile_m * tile_n * 2 * PdS_stage
+        sP = sPdS if not mma_dkv_is_rs else 0
+        sdQaccum = tile_m * hdim * 4
+        smem = sQ + sK + sV + sdO + sP + sPdS + sdQaccum
+        if smem <= SMEM_LIMIT:
+            break
+    else:
+        return None
+
+    # SMEM traffic
+    def _swap(a, b, s):
+        return (b, a) if s else (a, b)
+
+    def _wg_n(al_m, s):
+        return al_m if s else num_wg // al_m
+
+    M_s, N_s = _swap(tile_m, tile_n, SdP_swapAB)
+    wn_SdP = _wg_n(AtomLayoutMSdP, SdP_swapAB)
+    traffic_S = _mma_traffic(M_s, N_s, hdim, num_wg, wn_SdP)
+    traffic_dP = _mma_traffic(M_s, N_s, hdimv, num_wg, wn_SdP)
+
+    wn_dKV = _wg_n(AtomLayoutNdKV, dKV_swapAB)
+    M_dv, N_dv = _swap(tile_n, hdimv, dKV_swapAB)
+    traffic_dV = _mma_traffic(M_dv, N_dv, tile_m, num_wg, wn_dKV, is_rs=mma_dkv_is_rs)
+    M_dk, N_dk = _swap(tile_n, hdim, dKV_swapAB)
+    traffic_dK = _mma_traffic(M_dk, N_dk, tile_m, num_wg, wn_dKV, is_rs=mma_dkv_is_rs)
+
+    M_dq, N_dq = _swap(tile_m, hdim, dQ_swapAB)
+    wn_dQ = _wg_n(AtomLayoutMdQ, dQ_swapAB)
+    traffic_dQ = _mma_traffic(M_dq, N_dq, tile_n, num_wg, wn_dQ)
+
+    traffic_P_store = tile_m * tile_n * 2 if not mma_dkv_is_rs else 0
+    traffic_dS_store = tile_m * tile_n * 2
+    traffic_dQ_smem = tile_m * hdim * 4 * 2  # store + TMA load
+
+    smem_traffic = (
+        traffic_S
+        + traffic_dP
+        + traffic_dV
+        + traffic_dK
+        + traffic_dQ
+        + traffic_P_store
+        + traffic_dS_store
+        + traffic_dQ_smem
+    )
+
+    return dict(
+        tile_m=tile_m,
+        tile_n=tile_n,
+        num_wg=num_wg,
+        Q_stage=Q_stage,
+        dO_stage=dO_stage,
+        PdS_stage=PdS_stage,
+        SdP_swapAB=SdP_swapAB,
+        dKV_swapAB=dKV_swapAB,
+        dQ_swapAB=dQ_swapAB,
+        AtomLayoutMSdP=AtomLayoutMSdP,
+        AtomLayoutNdKV=AtomLayoutNdKV,
+        AtomLayoutMdQ=AtomLayoutMdQ,
+        mma_dkv_is_rs=mma_dkv_is_rs,
+        regs_SdP=regs_SdP,
+        regs_dK=regs_dK,
+        regs_dV=regs_dV,
+        regs_dQ=regs_dQ,
+        total_regs=total_regs,
+        reg_limit=reg_limit,
+        smem_bytes=smem,
+        smem_kb=smem / 1024,
+        smem_traffic=smem_traffic,
+        smem_traffic_kb=smem_traffic / 1024,
+        smem_traffic_per_block=smem_traffic / (tile_m * tile_n),
+    )
+
+
+def find_feasible_bwd_configs(
+    head_dim,
+    head_dim_v=None,
+    tile_m_choices=(64, 80, 96, 112, 128),
+    tile_n_choices=(64, 80, 96, 112, 128),
+):
+    if head_dim_v is None:
+        head_dim_v = head_dim
+    hdim = int(math.ceil(head_dim / 32) * 32)
+    hdimv = int(math.ceil(head_dim_v / 32) * 32)
+
+    results = []
+    for num_wg in (2, 3):
+        divs = _divisors(num_wg)
+        for tile_m in tile_m_choices:
+            for tile_n in tile_n_choices:
+                for SdP_swap in (False, True):
+                    if (tile_n if SdP_swap else tile_m) % 64 != 0:
+                        continue
+                    for dKV_swap in (False, True):
+                        if not dKV_swap and tile_n % 64 != 0:
+                            continue
+                        if dKV_swap and (hdim % 64 != 0 or hdimv % 64 != 0):
+                            continue
+                        for dQ_swap in (False, True):
+                            if (hdim if dQ_swap else tile_m) % 64 != 0:
+                                continue
+                            for a1 in divs:
+                                for a2 in divs:
+                                    for a3 in divs:
+                                        cfg = _check_bwd_config(
+                                            hdim,
+                                            hdimv,
+                                            tile_m,
+                                            tile_n,
+                                            num_wg,
+                                            SdP_swap,
+                                            dKV_swap,
+                                            dQ_swap,
+                                            a1,
+                                            a2,
+                                            a3,
+                                        )
+                                        if cfg is not None:
+                                            results.append(cfg)
+
+    results.sort(
+        key=lambda c: (-c["tile_n"], -c["tile_m"], c["smem_traffic_per_block"])
+    )
+    return results
+
+
+def print_bwd_configs(configs, max_results=20):
+    if not configs:
+        print("No feasible configs found!")
+        return
+    n = min(len(configs), max_results)
+    print(f"Found {len(configs)} feasible configs (showing top {n}):\n")
+    hdr = (
+        f"{'wg':>2} {'tm':>3} {'tn':>3}  "
+        f"{'SdP':>3} {'dKV':>3} {'dQ':>3}  "
+        f"{'aSdP':>4} {'adKV':>4} {'adQ':>4}  "
+        f"{'Qs':>2} {'dOs':>3}  "
+        f"{'rS':>3} {'rdK':>3} {'rdV':>3} {'rdQ':>3} {'tot':>4}/{'':<3}  "
+        f"{'smem':>5}  {'traffic':>7}  {'tr/blk':>6}"
+    )
+    print(hdr)
+    print("-" * len(hdr))
+    B = lambda b: "T" if b else "F"
+    for c in configs[:max_results]:
+        print(
+            f"{c['num_wg']:>2} {c['tile_m']:>3} {c['tile_n']:>3}  "
+            f"{B(c['SdP_swapAB']):>3} {B(c['dKV_swapAB']):>3} {B(c['dQ_swapAB']):>3}  "
+            f"{c['AtomLayoutMSdP']:>4} {c['AtomLayoutNdKV']:>4} {c['AtomLayoutMdQ']:>4}  "
+            f"{c['Q_stage']:>2} {c['dO_stage']:>3}  "
+            f"{c['regs_SdP']:>3} {c['regs_dK']:>3} {c['regs_dV']:>3} {c['regs_dQ']:>3} "
+            f"{c['total_regs']:>4}/{c['reg_limit']:<3}  "
+            f"{c['smem_kb']:>4.0f}K  "
+            f"{c['smem_traffic_kb']:>6.0f}K  "
+            f"{c['smem_traffic_per_block']:>6.1f}"
+        )
+
+
+# ============================================================================
+# Forward
+# ============================================================================
+
+
+def _check_fwd_config(hdim, hdimv, tile_n, num_wg, pv_is_rs, overlap_wg):
+    reg_limit = REG_LIMITS[num_wg]
+    tile_m = num_wg * 64
+
+    if tile_n % 8 != 0:
+        return None
+
+    regs_S = _acc_regs(tile_m, tile_n, num_wg)
+    regs_O = _acc_regs(tile_m, hdimv, num_wg)
+    regs_P = regs_S // 2  # bf16 = half of f32
+
+    if overlap_wg:
+        total_regs = regs_S + regs_P + regs_O
+    else:
+        total_regs = regs_S + regs_O
+
+    if total_regs > reg_limit:
+        return None
+
+    # SMEM: 1 stage Q, 2 stages K/V, O overlaps Q, sP if not RS
+    sQ = tile_m * hdim * 2
+    sK = tile_n * hdim * 2 * 2
+    sV = tile_n * hdimv * 2 * 2
+    sO = tile_m * hdimv * 2
+    sP = tile_m * tile_n * 2 if not pv_is_rs else 0
+    smem = max(sQ, sO) + sK + sV + sP
+    if smem > SMEM_LIMIT:
+        return None
+
+    # SMEM traffic: num_instr = num_wg (all WGs in M, wg_n=1)
+    traffic_S = num_wg * (64 * hdim * 2 + tile_n * hdim * 2)
+    A_pv = 64 * tile_n * 2 if not pv_is_rs else 0
+    traffic_O = num_wg * (A_pv + hdimv * tile_n * 2)
+    traffic_P_store = tile_m * tile_n * 2 if not pv_is_rs else 0
+    smem_traffic = traffic_S + traffic_O + traffic_P_store
+
+    return dict(
+        tile_m=tile_m,
+        tile_n=tile_n,
+        num_wg=num_wg,
+        pv_is_rs=pv_is_rs,
+        overlap_wg=overlap_wg,
+        regs_S=regs_S,
+        regs_O=regs_O,
+        regs_P=regs_P,
+        total_regs=total_regs,
+        reg_limit=reg_limit,
+        smem_bytes=smem,
+        smem_kb=smem / 1024,
+        smem_traffic=smem_traffic,
+        smem_traffic_kb=smem_traffic / 1024,
+        smem_traffic_per_block=smem_traffic / (tile_m * tile_n),
+    )
+
+
+def find_feasible_fwd_configs(
+    head_dim, head_dim_v=None, tile_n_choices=(64, 80, 96, 112, 128, 144, 160, 176, 192)
+):
+    if head_dim_v is None:
+        head_dim_v = head_dim
+    hdim = int(math.ceil(head_dim / 32) * 32)
+    hdimv = int(math.ceil(head_dim_v / 32) * 32)
+
+    results = []
+    for num_wg in (2, 3):
+        for tile_n in tile_n_choices:
+            for pv_is_rs in (True, False):
+                for overlap_wg in (True, False):
+                    cfg = _check_fwd_config(
+                        hdim, hdimv, tile_n, num_wg, pv_is_rs, overlap_wg
+                    )
+                    if cfg is not None:
+                        results.append(cfg)
+
+    results.sort(key=lambda c: (-c["tile_n"], c["smem_traffic_per_block"]))
+    return results
+
+
+def print_fwd_configs(configs, max_results=20):
+    if not configs:
+        print("No feasible configs found!")
+        return
+    n = min(len(configs), max_results)
+    print(f"Found {len(configs)} feasible configs (showing top {n}):\n")
+    hdr = (
+        f"{'wg':>2} {'tm':>3} {'tn':>3}  "
+        f"{'RS':>2} {'olap':>4}  "
+        f"{'rS':>3} {'rP':>3} {'rO':>3} {'tot':>4}/{'':<3}  "
+        f"{'smem':>5}  {'traffic':>7}  {'tr/blk':>6}"
+    )
+    print(hdr)
+    print("-" * len(hdr))
+    B = lambda b: "T" if b else "F"
+    for c in configs[:max_results]:
+        print(
+            f"{c['num_wg']:>2} {c['tile_m']:>3} {c['tile_n']:>3}  "
+            f"{B(c['pv_is_rs']):>2} {B(c['overlap_wg']):>4}  "
+            f"{c['regs_S']:>3} {c['regs_P']:>3} {c['regs_O']:>3} "
+            f"{c['total_regs']:>4}/{c['reg_limit']:<3}  "
+            f"{c['smem_kb']:>4.0f}K  "
+            f"{c['smem_traffic_kb']:>6.0f}K  "
+            f"{c['smem_traffic_per_block']:>6.1f}"
+        )
+
+
+# ============================================================================
+# CLI
+# ============================================================================
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Search feasible SM90 MMA configs")
+    parser.add_argument("--mode", choices=["fwd", "bwd", "both"], default="both")
+    parser.add_argument(
+        "--headdim",
+        type=str,
+        default="128",
+        help="Head dim, or hdim-hdimv (e.g. 192-128)",
+    )
+    parser.add_argument(
+        "--tile-m", type=str, default="64,80,96,112,128", help="Bwd tile_m choices"
+    )
+    parser.add_argument(
+        "--tile-n",
+        type=str,
+        default=None,
+        help="tile_n choices (default: fwd up to 192, bwd up to 128)",
+    )
+    parser.add_argument("-n", "--num-results", type=int, default=30)
+    args = parser.parse_args()
+
+    parts = args.headdim.split("-")
+    hdim = int(parts[0])
+    hdimv = int(parts[1]) if len(parts) > 1 else hdim
+
+    TN_FWD = "64,80,96,112,128,144,160,176,192"
+    TN_BWD = "64,80,96,112,128"
+
+    if args.mode in ("fwd", "both"):
+        tn = tuple(int(x) for x in (args.tile_n or TN_FWD).split(","))
+        print(f"=== FWD configs: hdim={hdim}, hdimv={hdimv} ===\n")
+        print_fwd_configs(find_feasible_fwd_configs(hdim, hdimv, tn), args.num_results)
+        print()
+
+    if args.mode in ("bwd", "both"):
+        tm = tuple(int(x) for x in args.tile_m.split(","))
+        tn = tuple(int(x) for x in (args.tile_n or TN_BWD).split(","))
+        print(f"=== BWD configs: hdim={hdim}, hdimv={hdimv} ===\n")
+        print_bwd_configs(
+            find_feasible_bwd_configs(hdim, hdimv, tm, tn), args.num_results
+        )

python/sglang/jit_kernel/flash_attn/cute/softmax.py

@@ -0,0 +1,759 @@
+# Copyright (c) 2025, Tri Dao.
+
+import math
+import operator
+from dataclasses import dataclass
+from typing import Tuple
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Boolean, Float32
+from quack import layout_utils
+from quack.cute_dsl_utils import ParamsBase
+
+import sglang.jit_kernel.flash_attn.cute.utils as utils
+from sglang.jit_kernel.flash_attn.cute.seqlen_info import SeqlenInfoQK
+from sglang.jit_kernel.flash_attn.cute.utils import AuxData
+
+
+@cute.jit
+def call_score_mod(
+    score_mod: cutlass.Constexpr,
+    score,
+    batch_idx,
+    head_idx,
+    q_idx,
+    kv_idx,
+    seqlen_info,
+    aux_data: AuxData,
+):
+    aux_tensors = aux_data.tensors if aux_data.tensors is not None else ()
+    # Compatibility shim for pre-aux_scalars score_mod callables.
+    if cutlass.const_expr(aux_data.scalars is not None):
+        return score_mod(
+            score,
+            batch_idx,
+            head_idx,
+            q_idx=q_idx,
+            kv_idx=kv_idx,
+            seqlen_info=seqlen_info,
+            aux_tensors=aux_tensors,
+            aux_scalars=aux_data.scalars,
+        )
+    return score_mod(
+        score,
+        batch_idx,
+        head_idx,
+        q_idx=q_idx,
+        kv_idx=kv_idx,
+        seqlen_info=seqlen_info,
+        aux_tensors=aux_tensors,
+    )
+
+
+@cute.jit
+def call_score_mod_bwd(
+    score_mod_bwd: cutlass.Constexpr,
+    grad,
+    score,
+    batch_idx,
+    head_idx,
+    q_idx,
+    kv_idx,
+    seqlen_info,
+    aux_data: AuxData,
+):
+    aux_tensors = aux_data.tensors if aux_data.tensors is not None else ()
+    # Compatibility shim for pre-aux_scalars score_mod_bwd callables.
+    if cutlass.const_expr(aux_data.scalars is not None):
+        return score_mod_bwd(
+            grad,
+            score,
+            batch_idx,
+            head_idx,
+            q_idx=q_idx,
+            kv_idx=kv_idx,
+            seqlen_info=seqlen_info,
+            aux_tensors=aux_tensors,
+            aux_scalars=aux_data.scalars,
+        )
+    return score_mod_bwd(
+        grad,
+        score,
+        batch_idx,
+        head_idx,
+        q_idx=q_idx,
+        kv_idx=kv_idx,
+        seqlen_info=seqlen_info,
+        aux_tensors=aux_tensors,
+    )
+
+
+@dataclass
+class Softmax(ParamsBase):
+    scale_log2: Float32
+    num_rows: cutlass.Constexpr[int]
+    row_max: cute.Tensor
+    row_sum: cute.Tensor
+    arch: cutlass.Constexpr[int] = 80
+    softmax_scale: Float32 | None = None
+
+    @staticmethod
+    def create(
+        scale_log2: Float32,
+        num_rows: cutlass.Constexpr[int],
+        arch: cutlass.Constexpr[int] = 80,
+        softmax_scale: Float32 | None = None,
+    ):
+        row_max = cute.make_rmem_tensor(num_rows, Float32)
+        row_sum = cute.make_rmem_tensor(num_rows, Float32)
+        return Softmax(scale_log2, num_rows, row_max, row_sum, arch, softmax_scale)
+
+    def reset(self) -> None:
+        self.row_max.fill(-Float32.inf)
+        self.row_sum.fill(0.0)
+
+    def _compute_row_max(
+        self, acc_S_row: cute.TensorSSA, init_val: float | Float32 | None = None
+    ) -> Float32:
+        return utils.fmax_reduce(acc_S_row, init_val, arch=self.arch)
+
+    def _compute_row_sum(
+        self, acc_S_row_exp: cute.TensorSSA, init_val: float | Float32 | None = None
+    ) -> Float32:
+        return utils.fadd_reduce(acc_S_row_exp, init_val, arch=self.arch)
+
+    @cute.jit
+    def online_softmax(
+        self,
+        acc_S: cute.Tensor,
+        is_first: cutlass.Constexpr[bool] = False,
+        check_inf: cutlass.Constexpr[bool] = True,
+    ) -> cute.Tensor:
+        """Apply online softmax and return the row_scale to rescale O.
+
+        :param acc_S: acc_S tensor
+        :type acc_S: cute.Tensor
+        :param is_first: is first n_block
+        :type is_first: cutlass.Constexpr
+        """
+        # Change acc_S to M,N layout view.
+        acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S)
+        row_scale = cute.make_fragment_like(self.row_max, Float32)
+
+        row_max = self.row_max
+        row_sum = self.row_sum
+        scale_log2 = self.scale_log2
+        arch = self.arch
+
+        # Each iteration processes one row of acc_S
+        for r in cutlass.range(cute.size(row_max), unroll_full=True):
+            acc_S_row = acc_S_mn[r, None].load()  # (n_block_size)
+
+            row_max_cur = utils.fmax_reduce(
+                acc_S_row,
+                init_val=row_max[r] if cutlass.const_expr(not is_first) else None,
+                arch=arch,
+            )
+
+            row_max_cur = cute.arch.warp_reduction_max(row_max_cur, threads_in_group=4)
+            # Update row_max before changing row_max_cur to safe value for -inf
+            row_max_prev = row_max[r]
+            row_max[r] = row_max_cur
+
+            if cutlass.const_expr(check_inf):
+                row_max_cur = 0.0 if row_max_cur == -Float32.inf else row_max_cur
+
+            if cutlass.const_expr(is_first):
+                row_max_cur_scaled = row_max_cur * scale_log2
+                acc_S_row_exp = cute.math.exp2(
+                    acc_S_row * scale_log2 - row_max_cur_scaled, fastmath=True
+                )
+                acc_S_row_sum = utils.fadd_reduce(
+                    acc_S_row_exp, init_val=None, arch=arch
+                )
+                row_scale[r] = 1.0
+            else:
+                row_max_cur_scaled = row_max_cur * scale_log2
+                acc_S_row_exp = cute.math.exp2(
+                    acc_S_row * scale_log2 - row_max_cur_scaled, fastmath=True
+                )
+                # row_scale[r] = cute.math.exp2(row_max_prev * self.scale_log2 - row_max_cur_scaled)
+                row_scale[r] = cute.math.exp2(
+                    (row_max_prev - row_max_cur) * scale_log2, fastmath=True
+                )
+                acc_S_row_sum = utils.fadd_reduce(
+                    acc_S_row_exp, init_val=row_sum[r] * row_scale[r], arch=arch
+                )
+
+            row_sum[r] = acc_S_row_sum
+            acc_S_mn[r, None].store(acc_S_row_exp)
+
+        return row_scale
+
+    @cute.jit
+    def finalize(
+        self, final_scale: Float32 = 1.0, sink_val: Float32 | cute.Tensor | None = None
+    ) -> cute.Tensor:
+        """Finalize the online softmax by computing the scale and logsumexp."""
+        if cutlass.const_expr(
+            sink_val is not None and isinstance(sink_val, cute.Tensor)
+        ):
+            assert cute.size(sink_val) == cute.size(self.row_sum)
+        row_sum = self.row_sum
+        row_max = self.row_max
+        scale_log2 = self.scale_log2
+
+        # quad reduction for row_sum as we didn't do it during each iteration of online softmax
+        row_sum.store(utils.warp_reduce(row_sum.load(), operator.add, width=4))
+        row_scale = cute.make_fragment_like(row_max, Float32)
+
+        for r in cutlass.range(cute.size(row_sum), unroll_full=True):
+            if cutlass.const_expr(sink_val is not None):
+                sink_val_cur = (
+                    sink_val if not isinstance(sink_val, cute.Tensor) else sink_val[r]
+                )
+                LOG2_E = math.log2(math.e)
+                if row_max[r] == -Float32.inf:
+                    # Fully-masked / empty row (can happen with SplitKV when a split's
+                    # blocks are all outside the local window)
+                    row_max[r] = sink_val_cur * (LOG2_E / scale_log2)
+                    row_sum[r] = 1.0
+                else:
+                    row_sum[r] += cute.math.exp2(
+                        sink_val_cur * LOG2_E - row_max[r] * scale_log2, fastmath=True
+                    )
+
+            # if row_sum is zero or nan, set acc_O_mn_row to 1.0
+            acc_O_mn_row_is_zero_or_nan = row_sum[r] == 0.0 or row_sum[r] != row_sum[r]
+            row_scale[r] = (
+                cute.arch.rcp_approx(
+                    row_sum[r] if not acc_O_mn_row_is_zero_or_nan else 1.0
+                )
+            ) * final_scale
+            row_sum_cur = row_sum[r]
+            LN2 = math.log(2.0)
+            row_sum[r] = (
+                (row_max[r] * scale_log2 + cute.math.log2(row_sum_cur, fastmath=True))
+                * LN2
+                if not acc_O_mn_row_is_zero_or_nan
+                else -Float32.inf
+            )
+        return row_scale
+
+    @cute.jit
+    def rescale_O(self, acc_O: cute.Tensor, row_scale: cute.Tensor) -> None:
+        """Scale each row of acc_O by the given scale tensor.
+        :param acc_O: input tensor
+        :type acc_O: cute.Tensor
+        :param row_scale: row_scale tensor
+        :type row_scale: cute.Tensor
+        """
+        acc_O_mn = layout_utils.reshape_acc_to_mn(acc_O)
+        assert cute.size(row_scale) == cute.size(acc_O_mn, mode=[0])
+        for r in cutlass.range(cute.size(row_scale), unroll_full=True):
+            acc_O_mn[r, None].store(acc_O_mn[r, None].load() * row_scale[r])
+
+
+@dataclass
+class SoftmaxSm100(Softmax):
+    rescale_threshold: cutlass.Constexpr[float] = 0.0
+    max_offset: cutlass.Constexpr[int] = 0
+
+    @staticmethod
+    def create(
+        scale_log2: Float32,
+        rescale_threshold: cutlass.Constexpr[float] = 0.0,
+        softmax_scale: Float32 | None = None,
+        max_offset: cutlass.Constexpr[int] = 0,
+    ):
+        num_rows = 1
+        arch = 100
+        row_max = cute.make_rmem_tensor(num_rows, Float32)
+        row_sum = cute.make_rmem_tensor(num_rows, Float32)
+        return SoftmaxSm100(
+            scale_log2,
+            num_rows,
+            row_max,
+            row_sum,
+            arch,
+            softmax_scale,
+            rescale_threshold=rescale_threshold,
+            max_offset=max_offset,
+        )
+
+    @cute.jit
+    def compute_row_max_local(
+        self, acc_S_row: cute.TensorSSA, is_first: Boolean
+    ) -> Float32:
+        if cutlass.const_expr(is_first):
+            row_max_new = self._compute_row_max(acc_S_row)
+        else:
+            row_max_old = self.row_max[0]
+            row_max_new = self._compute_row_max(acc_S_row, init_val=row_max_old)
+        return row_max_new
+
+    @cute.jit
+    def update_row_max_from_local(
+        self,
+        row_max_new: Float32,
+        is_first: Boolean,
+    ) -> Tuple[Float32, Float32]:
+        if cutlass.const_expr(is_first):
+            row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0
+            acc_scale = 0.0
+        else:
+            row_max_old = self.row_max[0]
+            row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0
+            acc_scale_ = (row_max_old - row_max_safe) * self.scale_log2
+            acc_scale = cute.math.exp2(acc_scale_)
+            if cutlass.const_expr(self.rescale_threshold > 0.0):
+                if acc_scale_ >= -self.rescale_threshold:
+                    row_max_new = row_max_old
+                    row_max_safe = row_max_old
+                    acc_scale = 1.0
+        self.row_max[0] = row_max_new
+        return row_max_safe, acc_scale
+
+    @cute.jit
+    def update_row_max(
+        self, acc_S_row: cute.TensorSSA, is_first: int
+    ) -> Tuple[Float32, Float32]:
+        if cutlass.const_expr(is_first):
+            row_max_new = self._compute_row_max(acc_S_row)
+            row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0
+            acc_scale = 0.0
+        else:
+            row_max_old = self.row_max[0]
+            row_max_new = self._compute_row_max(acc_S_row, init_val=row_max_old)
+            row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0
+            acc_scale_ = (row_max_old - row_max_safe) * self.scale_log2
+            acc_scale = cute.math.exp2(acc_scale_, fastmath=True)
+            if cutlass.const_expr(self.rescale_threshold > 0.0):
+                if acc_scale_ >= -self.rescale_threshold:
+                    row_max_new = row_max_old
+                    row_max_safe = row_max_old
+                    acc_scale = 1.0
+        self.row_max[0] = row_max_new
+        return row_max_safe, acc_scale
+
+    def update_row_sum(
+        self, acc_S_row_exp: cute.TensorSSA, row_scale: Float32, is_first: int = False
+    ) -> None:
+        init_val = (
+            self.row_sum[0] * row_scale if cutlass.const_expr(not is_first) else None
+        )
+        # self.row_sum[0] = self._compute_row_sum(acc_S_row_exp, init_val=self.row_sum[0] * row_scale)
+        self.row_sum[0] = self._compute_row_sum(acc_S_row_exp, init_val=init_val)
+        # tmp = self._compute_row_sum(acc_S_row_exp)
+        # self.row_sum[0] = self.row_sum[0] * row_scale + tmp
+
+    @cute.jit
+    def scale_subtract_rowmax(
+        self,
+        acc_S_row: cute.Tensor,
+        row_max: Float32,
+    ):
+        assert (
+            cute.size(acc_S_row.shape) % 2 == 0
+        ), "acc_S_row must have an even number of elements"
+        row_max_scaled = row_max * self.scale_log2
+        max_offset = Float32(self.max_offset)
+        bias = max_offset - row_max_scaled
+        for i in cutlass.range(0, cute.size(acc_S_row.shape), 2, unroll_full=True):
+            acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2(
+                (acc_S_row[i], acc_S_row[i + 1]),
+                (self.scale_log2, self.scale_log2),
+                (bias, bias),
+            )
+
+    @cute.jit
+    def apply_exp2_convert(
+        self,
+        acc_S_row: cute.Tensor,
+        acc_S_row_converted: cute.Tensor,
+        ex2_emu_freq: cutlass.Constexpr[int] = 0,
+        ex2_emu_res: cutlass.Constexpr[int] = 4,
+        ex2_emu_start_frg: cutlass.Constexpr[int] = 0,
+    ):
+        assert (
+            cute.size(acc_S_row.shape) % 2 == 0
+        ), "acc_S_row must have an even number of elements"
+        frg_tile = 32
+        assert frg_tile % 2 == 0
+        frg_cnt = cute.size(acc_S_row) // frg_tile
+        assert cute.size(acc_S_row) % frg_tile == 0
+        acc_S_row_frg = cute.logical_divide(acc_S_row, cute.make_layout(frg_tile))
+        acc_S_row_converted_frg = cute.logical_divide(
+            acc_S_row_converted, cute.make_layout(frg_tile)
+        )
+        for j in cutlass.range_constexpr(frg_cnt):
+            for k in cutlass.range_constexpr(0, cute.size(acc_S_row_frg, mode=[0]), 2):
+                # acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True)
+                # acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True)
+                if cutlass.const_expr(ex2_emu_freq == 0):
+                    acc_S_row_frg[k, j] = cute.math.exp2(
+                        acc_S_row_frg[k, j], fastmath=True
+                    )
+                    acc_S_row_frg[k + 1, j] = cute.math.exp2(
+                        acc_S_row_frg[k + 1, j], fastmath=True
+                    )
+                else:
+                    if cutlass.const_expr(
+                        k % ex2_emu_freq < ex2_emu_freq - ex2_emu_res
+                        or j >= frg_cnt - 1
+                        or j < ex2_emu_start_frg
+                    ):
+                        acc_S_row_frg[k, j] = cute.math.exp2(
+                            acc_S_row_frg[k, j], fastmath=True
+                        )
+                        acc_S_row_frg[k + 1, j] = cute.math.exp2(
+                            acc_S_row_frg[k + 1, j], fastmath=True
+                        )
+                    else:
+                        # acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = utils.e2e_asm2(acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j])
+                        acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = (
+                            utils.ex2_emulation_2(
+                                acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j]
+                            )
+                        )
+            acc_S_row_converted_frg[None, j].store(
+                acc_S_row_frg[None, j].load().to(acc_S_row_converted.element_type)
+            )
+
+    @cute.jit
+    def scale_apply_exp2_convert(
+        self,
+        acc_S_row: cute.Tensor,
+        row_max: Float32,
+        acc_S_row_converted: cute.Tensor,
+    ):
+        assert (
+            cute.size(acc_S_row.shape) % 2 == 0
+        ), "acc_S_row must have an even number of elements"
+        minus_row_max_scaled = -row_max * self.scale_log2
+        for i in cutlass.range_constexpr(0, cute.size(acc_S_row.shape), 2):
+            acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2(
+                (acc_S_row[i], acc_S_row[i + 1]),
+                (self.scale_log2, self.scale_log2),
+                (minus_row_max_scaled, minus_row_max_scaled),
+            )
+
+        # for i in cutlass.range_constexpr(0, cute.size(acc_S_row.shape), 2):
+        #     acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2(
+        #         (acc_S_row[i], acc_S_row[i + 1]),
+        #         (self.scale_log2, self.scale_log2),
+        #         (minus_row_max_scaled, minus_row_max_scaled),
+        #     )
+        #     acc_S_row[i] = cute.math.exp2(acc_S_row[i], fastmath=True)
+        #     acc_S_row[i + 1] = cute.math.exp2(acc_S_row[i + 1], fastmath=True)
+
+        frg_tile = 32
+        assert frg_tile % 2 == 0
+        frg_cnt = cute.size(acc_S_row) // frg_tile
+        assert cute.size(acc_S_row) % frg_tile == 0
+        acc_S_row_frg = cute.logical_divide(acc_S_row, cute.make_layout(frg_tile))
+        acc_S_row_converted_frg = cute.logical_divide(
+            acc_S_row_converted, cute.make_layout(frg_tile)
+        )
+        for j in cutlass.range_constexpr(frg_cnt):
+            for k in cutlass.range_constexpr(0, cute.size(acc_S_row_frg, mode=[0]), 2):
+                # acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = (
+                #     cute.arch.fma_packed_f32x2(
+                #         (acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j]),
+                #         (self.scale_log2, self.scale_log2),
+                #         (minus_row_max_scaled, minus_row_max_scaled),
+                #     )
+                # )
+                # acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True)
+                # acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True)
+                acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True)
+                acc_S_row_frg[k + 1, j] = cute.math.exp2(
+                    acc_S_row_frg[k + 1, j], fastmath=True
+                )
+            acc_S_row_converted_frg[None, j].store(
+                acc_S_row_frg[None, j].load().to(acc_S_row_converted.element_type)
+            )
+
+
+@cute.jit
+def floor_if_packed(
+    q_idx,
+    qhead_per_kvhead: cutlass.Constexpr[int],
+) -> cute.Tensor:
+    """Convert q_idx to packed format for Pack-GQA."""
+    if cutlass.const_expr(qhead_per_kvhead == 1):
+        return q_idx
+    return q_idx // qhead_per_kvhead
+
+
+@cute.jit
+def apply_score_mod_inner(
+    score_tensor,
+    index_tensor,
+    score_mod: cutlass.Constexpr,
+    batch_idx,
+    head_idx,
+    softmax_scale,
+    vec_size: cutlass.Constexpr,
+    qk_acc_dtype: cutlass.Constexpr,
+    aux_data: AuxData,
+    fastdiv_mods,
+    seqlen_info: SeqlenInfoQK,
+    constant_q_idx: cutlass.Constexpr,
+    qhead_per_kvhead: cutlass.Constexpr[int] = 1,
+    transpose_indices: cutlass.Constexpr[bool] = False,
+):
+    """Shared implementation for applying score modification.
+
+    Args:
+        score_tensor: The scores to modify (acc_S for flash_fwd, tSrS_t2r for sm100)
+        index_tensor: Index positions (tScS for flash_fwd, tScS_t2r for sm100)
+        score_mod: The score modification function to apply
+        batch_idx: Batch index
+        head_idx: Head index
+        softmax_scale: Scale to apply
+        vec_size: Vector size for processing elements
+        qk_acc_dtype: Data type for accumulator
+        aux_tensors: Optional aux_tensors for FlexAttention
+        aux_scalars: Optional runtime scalar captures for FlexAttention
+        fastdiv_mods: Tuple of (seqlen_q_divmod, seqlen_k_divmod) for wrapping
+        seqlen_info: Sequence length info
+        constant_q_idx: If provided, use this constant for all q_idx values
+                        If None, compute q_idx per-element
+        qhead_per_kvhead_packgqa: Pack-GQA replication factor. Divide q_idx by this
+                                  when greater than 1 so score mods see logical heads.
+        transpose_indices: If True, swap q_idx/kv_idx in index_tensor (for bwd kernel where S is transposed)
+    """
+    # Index positions in the index_tensor tuple
+    # Forward: index_tensor[...][0] = q_idx, index_tensor[...][1] = kv_idx
+    # Backward (transposed): index_tensor[...][0] = kv_idx, index_tensor[...][1] = q_idx
+    if cutlass.const_expr(transpose_indices):
+        q_idx_pos = cutlass.const_expr(1)
+        kv_idx_pos = cutlass.const_expr(0)
+    else:
+        q_idx_pos = cutlass.const_expr(0)
+        kv_idx_pos = cutlass.const_expr(1)
+
+    n_vals = cutlass.const_expr(cute.size(score_tensor.shape))
+    score_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype)
+    kv_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+
+    # SSA values for batch (constant across all elements)
+    batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32).broadcast_to(
+        (vec_size,)
+    )
+
+    # Handle q_idx based on whether it's constant
+    q_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+
+    # For Pack-GQA with non-constant q_idx, we need per-element head indices
+    # since a thread may process multiple query head indices
+    if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
+        head_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+
+    for i in cutlass.range(0, n_vals, vec_size, unroll_full=True):
+        for j in cutlass.range(vec_size, unroll_full=True):
+            score_vec[j] = score_tensor[i + j] * softmax_scale
+
+            # Extract head offset from packed q_idx for Pack-GQA
+            if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
+                q_idx_packed = index_tensor[i + j][q_idx_pos]
+                # Building up the logical q_head idx: final_q_head = kv_head * qhead_per_kvhead + (q_physical % qhead_per_kvhead)
+                q_idx_logical = q_idx_packed // qhead_per_kvhead
+                head_offset = q_idx_packed - q_idx_logical * qhead_per_kvhead
+                head_idx_vec[j] = head_idx * qhead_per_kvhead + head_offset
+
+            # If we will do loads we mod, in order to not read OOB
+            if cutlass.const_expr(
+                aux_data.tensors is not None and fastdiv_mods is not None
+            ):
+                if cutlass.const_expr(constant_q_idx is None):
+                    seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods
+                    q_idx_floored = floor_if_packed(
+                        index_tensor[i + j][q_idx_pos], qhead_per_kvhead
+                    )
+                    _, q_idx_wrapped = divmod(q_idx_floored, seqlen_q_divmod)
+                    q_idx_vec[j] = q_idx_wrapped
+                else:
+                    _, seqlen_k_divmod = fastdiv_mods
+
+                _, kv_idx_wrapped = divmod(
+                    index_tensor[i + j][kv_idx_pos], seqlen_k_divmod
+                )
+                kv_idx_vec[j] = kv_idx_wrapped
+            else:
+                # No bounds checking - direct indexing
+                if constant_q_idx is None:
+                    q_idx_vec[j] = floor_if_packed(
+                        index_tensor[i + j][q_idx_pos], qhead_per_kvhead
+                    )
+                kv_idx_vec[j] = index_tensor[i + j][kv_idx_pos]
+
+        # Convert to SSA for score_mod call
+        score_ssa = score_vec.load()
+        kv_idx_ssa = kv_idx_vec.load()
+        if cutlass.const_expr(constant_q_idx is None):
+            q_idx_ssa = q_idx_vec.load()
+        else:
+            # NB we do not apply Pack-GQA division here, as constant_q_idx is assumed to already be logical
+            q_idx_const = constant_q_idx
+            q_idx_ssa = utils.scalar_to_ssa(q_idx_const, cutlass.Int32).broadcast_to(
+                (vec_size,)
+            )
+
+        # Compute head_idx_ssa: per-element for Pack-GQA with non-constant q_idx, constant otherwise
+        if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
+            head_idx_ssa = head_idx_vec.load()
+        else:
+            head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32).broadcast_to(
+                (vec_size,)
+            )
+
+        post_mod_scores = call_score_mod(
+            score_mod,
+            score_ssa,
+            batch_idx_ssa,
+            head_idx_ssa,
+            q_idx_ssa,
+            kv_idx_ssa,
+            seqlen_info,
+            aux_data,
+        )
+
+        # Write back modified scores
+        score_vec.store(post_mod_scores)
+        for j in cutlass.range(vec_size, unroll_full=True):
+            score_tensor[i + j] = score_vec[j]
+
+
+@cute.jit
+def apply_score_mod_bwd_inner(
+    grad_tensor,
+    score_tensor,
+    index_tensor,
+    score_mod_bwd: cutlass.Constexpr,
+    batch_idx,
+    head_idx,
+    softmax_scale,
+    vec_size: cutlass.Constexpr,
+    qk_acc_dtype: cutlass.Constexpr,
+    aux_data: AuxData,
+    fastdiv_mods,
+    seqlen_info,
+    constant_q_idx: cutlass.Constexpr,
+    qhead_per_kvhead: cutlass.Constexpr[int] = 1,
+    transpose_indices: cutlass.Constexpr[bool] = False,
+):
+    """Apply backward score modification (joint graph).
+
+    Args:
+        grad_tensor: in/out: dlogits rewritten in-place with d(scaled_scores)
+        score_tensor: pre-mod scores (unscaled QK tile), scaled by softmax_scale internally
+        index_tensor: Index positions (same as forward)
+        score_mod_bwd: The backward score modification function (joint graph)
+        batch_idx: Batch index
+        head_idx: Head index
+        softmax_scale: Scale to apply to score_tensor
+        vec_size: Vector size for processing elements
+        qk_acc_dtype: Data type for accumulator
+        aux_tensors: Optional aux_tensors for FlexAttention
+        aux_scalars: Optional runtime scalar captures for FlexAttention
+        fastdiv_mods: Tuple of (seqlen_q_divmod, seqlen_k_divmod) for wrapping
+        seqlen_info: Sequence length info
+        constant_q_idx: If provided, use this constant for all q_idx values
+        qhead_per_kvhead: Pack-GQA replication factor
+        transpose_indices: If True, swap q_idx/kv_idx in index_tensor
+    """
+    # Index positions in the index_tensor tuple
+    # Forward: index_tensor[...][0] = q_idx, index_tensor[...][1] = kv_idx
+    # Backward (transposed): index_tensor[...][0] = kv_idx, index_tensor[...][1] = q_idx
+    if cutlass.const_expr(transpose_indices):
+        q_idx_pos = cutlass.const_expr(1)
+        kv_idx_pos = cutlass.const_expr(0)
+    else:
+        q_idx_pos = cutlass.const_expr(0)
+        kv_idx_pos = cutlass.const_expr(1)
+    n_vals = cutlass.const_expr(cute.size(grad_tensor.shape))
+    grad_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype)
+    score_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype)
+    kv_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+    batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32).broadcast_to(
+        (vec_size,)
+    )
+    q_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+
+    # For Pack-GQA with non-constant q_idx, we need per-element head indices
+    if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
+        head_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+
+    for i in cutlass.range(0, n_vals, vec_size, unroll_full=True):
+        for j in cutlass.range(vec_size, unroll_full=True):
+            grad_vec[j] = grad_tensor[i + j]
+            # Scale score so joint graph sees same value as forward score_mod
+            score_vec[j] = score_tensor[i + j] * softmax_scale
+
+            if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
+                q_idx_packed = index_tensor[i + j][q_idx_pos]
+                q_idx_logical = q_idx_packed // qhead_per_kvhead
+                head_offset = q_idx_packed - q_idx_logical * qhead_per_kvhead
+                head_idx_vec[j] = head_idx * qhead_per_kvhead + head_offset
+
+            if cutlass.const_expr(
+                aux_data.tensors is not None and fastdiv_mods is not None
+            ):
+                if cutlass.const_expr(constant_q_idx is None):
+                    seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods
+                    q_idx_floored = floor_if_packed(
+                        index_tensor[i + j][q_idx_pos], qhead_per_kvhead
+                    )
+                    _, q_idx_wrapped = divmod(q_idx_floored, seqlen_q_divmod)
+                    q_idx_vec[j] = q_idx_wrapped
+                else:
+                    _, seqlen_k_divmod = fastdiv_mods
+
+                _, kv_idx_wrapped = divmod(
+                    index_tensor[i + j][kv_idx_pos], seqlen_k_divmod
+                )
+                kv_idx_vec[j] = kv_idx_wrapped
+            else:
+                # No bounds checking - direct indexing
+                if constant_q_idx is None:
+                    q_idx_vec[j] = floor_if_packed(
+                        index_tensor[i + j][q_idx_pos], qhead_per_kvhead
+                    )
+                kv_idx_vec[j] = index_tensor[i + j][kv_idx_pos]
+
+        grad_ssa = grad_vec.load()
+        score_ssa = score_vec.load()
+        kv_idx_ssa = kv_idx_vec.load()
+
+        if cutlass.const_expr(constant_q_idx is None):
+            q_idx_ssa = q_idx_vec.load()
+        else:
+            q_idx_ssa = utils.scalar_to_ssa(constant_q_idx, cutlass.Int32).broadcast_to(
+                (vec_size,)
+            )
+
+        if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
+            head_idx_ssa = head_idx_vec.load()
+        else:
+            head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32).broadcast_to(
+                (vec_size,)
+            )
+
+        grad_out_ssa = call_score_mod_bwd(
+            score_mod_bwd,
+            grad_ssa,
+            score_ssa,
+            batch_idx_ssa,
+            head_idx_ssa,
+            q_idx_ssa,
+            kv_idx_ssa,
+            seqlen_info,
+            aux_data,
+        )
+
+        grad_vec.store(grad_out_ssa)
+        for j in cutlass.range(vec_size, unroll_full=True):
+            grad_tensor[i + j] = grad_vec[j]

python/sglang/jit_kernel/flash_attn/cute/testing.py

@@ -0,0 +1,580 @@
+import math
+from contextlib import nullcontext
+from functools import wraps
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch._guards import active_fake_mode
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+
+class IndexFirstAxis(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, indices):
+        ctx.save_for_backward(indices)
+        assert input.ndim >= 2
+        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
+        second_dim = other_shape.numel()
+        return torch.gather(
+            rearrange(input, "b ... -> b (...)"),
+            0,
+            repeat(indices, "z -> z d", d=second_dim),
+        ).reshape(-1, *other_shape)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (indices,) = ctx.saved_tensors
+        assert grad_output.ndim >= 2
+        other_shape = grad_output.shape[1:]
+        grad_output = rearrange(grad_output, "b ... -> b (...)")
+        grad_input = torch.zeros(
+            [ctx.first_axis_dim, grad_output.shape[1]],
+            device=grad_output.device,
+            dtype=grad_output.dtype,
+        )
+        grad_input.scatter_(
+            0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output
+        )
+        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
+
+
+index_first_axis = IndexFirstAxis.apply
+
+
+class IndexPutFirstAxis(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, values, indices, first_axis_dim):
+        ctx.save_for_backward(indices)
+        assert indices.ndim == 1
+        assert values.ndim >= 2
+        output = torch.zeros(
+            first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype
+        )
+        output[indices] = values
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (indices,) = ctx.saved_tensors
+        grad_values = grad_output[indices]
+        return grad_values, None, None
+
+
+index_put_first_axis = IndexPutFirstAxis.apply
+
+
+def unpad_input(hidden_states, attention_mask, unused_mask=None):
+    all_masks = (
+        (attention_mask + unused_mask) if unused_mask is not None else attention_mask
+    )
+    seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32)
+    used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    in_fake_mode = active_fake_mode() is not None
+    if not in_fake_mode:
+        indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten()
+        max_seqlen_in_batch = seqlens_in_batch.max().item()
+    else:
+        # torch.nonzero and .item() are not supported in FakeTensorMode
+        batch_size, seqlen = attention_mask.shape
+        indices = torch.arange(batch_size * seqlen, device=hidden_states.device)
+        max_seqlen_in_batch = seqlen
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    return (
+        index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices),
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+        used_seqlens_in_batch,
+    )
+
+
+def pad_input(hidden_states, indices, batch, seqlen):
+    output = index_put_first_axis(hidden_states, indices, batch * seqlen)
+    return rearrange(output, "(b s) ... -> b s ...", b=batch)
+
+
+def generate_random_padding_mask(
+    max_seqlen, batch_size, device, mode="random", zero_lengths=False, min_seqlen=None
+):
+    assert mode in ["full", "random", "third"]
+    min_seqlen = min_seqlen if min_seqlen is not None else 0 if zero_lengths else 1
+    if mode == "full":
+        lengths = torch.full(
+            (batch_size, 1), max_seqlen, device=device, dtype=torch.int32
+        )
+    elif mode == "random":
+        lengths = torch.randint(
+            max(min_seqlen, max_seqlen - 20),
+            max_seqlen + 1,
+            (batch_size, 1),
+            device=device,
+        )
+    else:
+        lengths = torch.randint(
+            max(min_seqlen, max_seqlen // 3),
+            max_seqlen + 1,
+            (batch_size, 1),
+            device=device,
+        )
+
+    if zero_lengths:
+        for i in range(batch_size):
+            if i % 5 == 0:
+                lengths[i] = 0
+        lengths[-1] = 0
+    padding_mask = (
+        repeat(torch.arange(max_seqlen, device=device), "s -> b s", b=batch_size)
+        < lengths
+    )
+    return padding_mask
+
+
+def generate_qkv(
+    q,
+    k,
+    v,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    qv=None,
+    kvpacked=False,
+    qkvpacked=False,
+    query_unused_mask=None,
+    key_unused_mask=None,
+):
+    assert not (kvpacked and qkvpacked)
+    batch_size, seqlen_q, nheads, d = q.shape
+    d_v = v.shape[-1]
+    _, seqlen_k, nheads_k, _ = k.shape
+    assert k.shape == (batch_size, seqlen_k, nheads_k, d)
+    assert v.shape == (batch_size, seqlen_k, nheads_k, d_v)
+    if query_unused_mask is not None or key_unused_mask is not None:
+        assert not kvpacked
+        assert not qkvpacked
+
+    if query_padding_mask is not None:
+        q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, seqused_q = unpad_input(
+            q, query_padding_mask, query_unused_mask
+        )
+        output_pad_fn = lambda output_unpad: pad_input(
+            output_unpad, indices_q, batch_size, seqlen_q
+        )
+        qv_unpad = (
+            rearrange(qv, "b s ... -> (b s) ...")[indices_q] if qv is not None else None
+        )
+    else:
+        q_unpad = rearrange(q, "b s h d -> (b s) h d")
+        cu_seqlens_q = torch.arange(
+            0,
+            (batch_size + 1) * seqlen_q,
+            step=seqlen_q,
+            dtype=torch.int32,
+            device=q_unpad.device,
+        )
+        seqused_q = None
+        max_seqlen_q = seqlen_q
+        output_pad_fn = lambda output_unpad: rearrange(
+            output_unpad, "(b s) h d -> b s h d", b=batch_size
+        )
+        qv_unpad = rearrange(qv, "b s ... -> (b s) ...") if qv is not None else None
+
+    if key_padding_mask is not None:
+        k_unpad, indices_k, cu_seqlens_k, max_seqlen_k, seqused_k = unpad_input(
+            k, key_padding_mask, key_unused_mask
+        )
+        v_unpad, *_ = unpad_input(v, key_padding_mask, key_unused_mask)
+    else:
+        k_unpad = rearrange(k, "b s h d -> (b s) h d")
+        v_unpad = rearrange(v, "b s h d -> (b s) h d")
+        cu_seqlens_k = torch.arange(
+            0,
+            (batch_size + 1) * seqlen_k,
+            step=seqlen_k,
+            dtype=torch.int32,
+            device=k_unpad.device,
+        )
+        seqused_k = None
+        max_seqlen_k = seqlen_k
+
+    if qkvpacked:
+        assert (query_padding_mask == key_padding_mask).all()
+        assert nheads == nheads_k
+        qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1)
+        qkv = torch.stack([q, k, v], dim=2)
+        if query_padding_mask is not None:
+            dqkv_pad_fn = lambda dqkv_unpad: pad_input(
+                dqkv_unpad, indices_q, batch_size, seqlen_q
+            )
+        else:
+            dqkv_pad_fn = lambda dqkv_unpad: rearrange(
+                dqkv_unpad, "(b s) t h d -> b s t h d", b=batch_size
+            )
+        return (
+            qkv_unpad.detach().requires_grad_(),
+            cu_seqlens_q,
+            max_seqlen_q,
+            qkv.detach().requires_grad_(),
+            output_pad_fn,
+            dqkv_pad_fn,
+        )
+    elif kvpacked:
+        kv_unpad = torch.stack([k_unpad, v_unpad], dim=1)
+        kv = torch.stack([k, v], dim=2)
+        dq_pad_fn = output_pad_fn
+        if key_padding_mask is not None:
+            dkv_pad_fn = lambda dkv_unpad: pad_input(
+                dkv_unpad, indices_k, batch_size, seqlen_k
+            )
+        else:
+            dkv_pad_fn = lambda dkv_unpad: rearrange(
+                dkv_unpad, "(b s) t h d -> b s t h d", b=batch_size
+            )
+        return (
+            q_unpad.detach().requires_grad_(),
+            kv_unpad.detach().requires_grad_(),
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            q.detach().requires_grad_(),
+            kv.detach().requires_grad_(),
+            output_pad_fn,
+            dq_pad_fn,
+            dkv_pad_fn,
+        )
+    else:
+        dq_pad_fn = output_pad_fn
+        if key_padding_mask is not None:
+            dk_pad_fn = lambda dk_unpad: pad_input(
+                dk_unpad, indices_k, batch_size, seqlen_k
+            )
+        else:
+            dk_pad_fn = lambda dk_unpad: rearrange(
+                dk_unpad, "(b s) h d -> b s h d", b=batch_size
+            )
+        return (
+            q_unpad.detach().requires_grad_(),
+            k_unpad.detach().requires_grad_(),
+            v_unpad.detach().requires_grad_(),
+            qv_unpad.detach() if qv is not None else None,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            q.detach().requires_grad_(),
+            k.detach().requires_grad_(),
+            v.detach().requires_grad_(),
+            qv.detach() if qv is not None else None,
+            output_pad_fn,
+            dq_pad_fn,
+            dk_pad_fn,
+        )
+
+
+def construct_local_mask(
+    seqlen_q,
+    seqlen_k,
+    window_size=(None, None),
+    sink_token_length=0,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    key_leftpad=None,
+    device=None,
+):
+    row_idx = rearrange(
+        torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1"
+    )
+    col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long)
+    if key_leftpad is not None:
+        key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1")
+        col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0])
+        col_idx = torch.where(col_idx >= key_leftpad, col_idx - key_leftpad, 2**32)
+    sk = (
+        seqlen_k
+        if key_padding_mask is None
+        else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1")
+    )
+    sq = (
+        seqlen_q
+        if query_padding_mask is None
+        else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1")
+    )
+    if window_size[0] is None:
+        return col_idx > row_idx + sk - sq + window_size[1]
+    else:
+        sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk
+        if window_size[1] is None:
+            local_mask_left = col_idx > sk
+        else:
+            local_mask_left = col_idx > torch.minimum(
+                row_idx + sk - sq + window_size[1], sk
+            )
+        return torch.logical_or(
+            local_mask_left,
+            torch.logical_and(
+                col_idx < row_idx + sk - sq - window_size[0],
+                col_idx >= sink_token_length,
+            ),
+        )
+
+
+def construct_chunk_mask(
+    seqlen_q,
+    seqlen_k,
+    attention_chunk,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    key_leftpad=None,
+    device=None,
+):
+    row_idx = rearrange(
+        torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1"
+    )
+    col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long)
+    if key_leftpad is not None:
+        key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1")
+        col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0])
+        col_idx = torch.where(col_idx >= key_leftpad, col_idx - key_leftpad, 2**32)
+    sk = (
+        seqlen_k
+        if key_padding_mask is None
+        else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1")
+    )
+    sq = (
+        seqlen_q
+        if query_padding_mask is None
+        else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1")
+    )
+    sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk
+    col_limit_left_chunk = row_idx + sk - sq - (row_idx + sk - sq) % attention_chunk
+    return torch.logical_or(
+        col_idx < col_limit_left_chunk,
+        col_idx >= col_limit_left_chunk + attention_chunk,
+    )
+
+
+def attention_ref(
+    q,
+    k,
+    v,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    key_leftpad=None,
+    attn_bias=None,
+    dropout_p=0.0,
+    dropout_mask=None,
+    causal=False,
+    qv=None,
+    q_descale=None,
+    k_descale=None,
+    v_descale=None,
+    window_size=(None, None),
+    attention_chunk=0,
+    sink_token_length=0,
+    learnable_sink: Optional[torch.Tensor] = None,
+    softcap=0.0,
+    upcast=True,
+    reorder_ops=False,
+    intermediate_dtype=None,
+    return_lse=False,
+    gather_kv_indices=None,
+    rel_bias: Optional[torch.Tensor] = None,  # [b, seqlen_q, h, rel_extent]
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+):
+    assert v is not None
+    has_qk = q is not None and k is not None
+    assert has_qk or qv is not None
+    if causal:
+        window_size = (window_size[0], 0)
+    dtype_og = v.dtype
+    q_shape = q.shape if q is not None else qv.shape
+    if upcast:
+        q, k, v, qv = [t.float() if t is not None else None for t in (q, k, v, qv)]
+    if q_descale is not None:
+        q_descale = repeat(q_descale, "b h -> b 1 (h g) 1", g=q_shape[2] // v.shape[2])
+        q, qv = [
+            (t.float() * q_descale).to(t.dtype) if t is not None else None
+            for t in (q, qv)
+        ]
+    if k_descale is not None:
+        k = (k.float() * rearrange(k_descale, "b h -> b 1 h 1")).to(dtype=k.dtype)
+    if v_descale is not None:
+        v = (v.float() * rearrange(v_descale, "b h -> b 1 h 1")).to(dtype=v.dtype)
+    seqlen_q, seqlen_k = q_shape[1], v.shape[1]
+    k, v = [
+        (
+            repeat(t, "b s h d -> b s (h g) d", g=q_shape[2] // t.shape[2])
+            if t is not None
+            else None
+        )
+        for t in (k, v)
+    ]
+    d = q_shape[-1]  # == dv for qv
+    dv = v.shape[-1]
+    softmax_scale = 1.0 / math.sqrt(d if qv is None or q is None else d + dv)
+    if has_qk:
+        scores = torch.einsum(
+            "bthd,bshd->bhts",
+            q if reorder_ops else q * softmax_scale,
+            k * softmax_scale if reorder_ops else k,
+        )
+    if qv is not None:
+        qv_scores = torch.einsum(
+            "bthd,bshd->bhts",
+            qv if reorder_ops else qv * softmax_scale,
+            v * softmax_scale if reorder_ops else v,
+        )
+        scores = qv_scores if not has_qk else scores + qv_scores
+    if softcap > 0:
+        scores = torch.tanh(scores / softcap) * softcap
+    if key_padding_mask is not None:
+        scores.masked_fill_(
+            rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf")
+        )
+    local_mask = None
+    if window_size[0] is not None or window_size[1] is not None:
+        local_mask = construct_local_mask(
+            seqlen_q,
+            seqlen_k,
+            window_size,
+            sink_token_length,
+            query_padding_mask,
+            key_padding_mask,
+            key_leftpad=key_leftpad,
+            device=v.device,
+        )
+    if attention_chunk > 0:
+        chunk_mask = construct_chunk_mask(
+            seqlen_q,
+            seqlen_k,
+            attention_chunk,
+            query_padding_mask,
+            key_padding_mask,
+            key_leftpad=key_leftpad,
+            device=v.device,
+        )
+        local_mask = (
+            torch.logical_or(local_mask, chunk_mask)
+            if local_mask is not None
+            else chunk_mask
+        )
+    if gather_kv_indices is not None:
+        batch = q_shape[0]
+        topk_len = gather_kv_indices.shape[2]
+        if topk_len < seqlen_k:
+            topk_index_mask = torch.full(
+                (batch, seqlen_q, seqlen_k), False, device="cuda"
+            ).scatter_(-1, gather_kv_indices, True)
+            scores.masked_fill_(
+                rearrange(~topk_index_mask, "b t s -> b 1 t s"), float("-inf")
+            )
+    if local_mask is not None:
+        scores.masked_fill_(local_mask, float("-inf"))
+    if attn_bias is not None:
+        scores = scores + attn_bias
+    if rel_bias is not None:
+        # Reference for Inkling sheared bias: gather rel_bias[i, h, i - j] into the [t, s] score grid.
+        rel_extent = rel_bias.shape[-1]
+        if cu_seqlens_q is not None:
+            seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1]
+        elif seqused_q is not None:
+            seqlens_q = seqused_q
+        else:
+            seqlens_q = torch.full(
+                (q.shape[0],), seqlen_q, device=q.device, dtype=torch.long
+            )
+        if cu_seqlens_k is not None:
+            seqlens_k = cu_seqlens_k[1:] - cu_seqlens_k[:-1]
+        elif seqused_k is not None:
+            seqlens_k = seqused_k
+        else:
+            seqlens_k = torch.full(
+                (q.shape[0],), seqlen_k, device=q.device, dtype=torch.long
+            )
+        seqlen_offset = (seqlens_k - seqlens_q).to(torch.long)  # [b]
+        q_idx = torch.arange(seqlen_q, device=q.device, dtype=torch.long)
+        kv_idx = torch.arange(seqlen_k, device=q.device, dtype=torch.long)
+        rel_dist = (
+            q_idx.unsqueeze(1) - kv_idx.unsqueeze(0) + seqlen_offset.view(-1, 1, 1)
+        )  # [b, seqlen_q, seqlen_k]
+        safe_dist = rel_dist.clamp(0, rel_extent - 1)
+        is_within_window = (rel_dist >= 0) & (rel_dist < rel_extent)
+        idx = safe_dist.unsqueeze(2).expand(-1, -1, rel_bias.shape[2], -1)
+        abs_bias = rel_bias.gather(dim=-1, index=idx)  # [b, seqlen_q, h, seqlen_k]
+        abs_bias = rearrange(abs_bias, "b t h s -> b h t s")
+        abs_bias = abs_bias.masked_fill(
+            rearrange(~is_within_window, "b t s -> b 1 t s"), 0.0
+        )
+        scores = scores + abs_bias
+    # After all masks are applied, before softmax:
+    # scores shape: [b, h, t, s]
+    lse = torch.logsumexp(scores, dim=-1)  # [b, h, t]
+    if learnable_sink is None:
+        attention = torch.softmax(scores, dim=-1).to(v.dtype)
+    else:
+        scores_fp32 = scores.to(torch.float32)
+        logits_max = torch.amax(scores_fp32, dim=-1, keepdim=True)
+        learnable_sink = rearrange(learnable_sink, "h -> h 1 1")
+        logits_or_sinks_max = torch.maximum(learnable_sink, logits_max)
+        unnormalized_scores = torch.exp(scores_fp32 - logits_or_sinks_max)
+        normalizer = unnormalized_scores.sum(dim=-1, keepdim=True) + torch.exp(
+            learnable_sink - logits_or_sinks_max
+        )
+        # LSE with sink: log(Z) = log(normalizer) + max
+        lse = (torch.log(normalizer.squeeze(-1)) + logits_or_sinks_max.squeeze(-1)).to(
+            dtype_og
+        )
+        attention = (unnormalized_scores / normalizer).to(v.dtype)
+    if query_padding_mask is not None:
+        attention = attention.masked_fill(
+            rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0
+        )
+    if key_padding_mask is not None:
+        attention = attention.masked_fill(
+            rearrange(~key_padding_mask, "b s -> b 1 1 s"), 0.0
+        )
+    if local_mask is not None:
+        attention = attention.masked_fill(
+            torch.all(local_mask, dim=-1, keepdim=True), 0.0
+        )
+    dropout_scaling = 1.0 / (1 - dropout_p)
+    if dropout_mask is not None:
+        attention_drop = attention.masked_fill(~dropout_mask, 0.0)
+    else:
+        attention_drop = attention
+    if intermediate_dtype is not None:
+        attention_drop = attention_drop.to(intermediate_dtype).to(attention_drop.dtype)
+    output = torch.einsum("bhts,bshd->bthd", attention_drop, v * dropout_scaling)
+    if query_padding_mask is not None:
+        output.masked_fill_(rearrange(~query_padding_mask, "b s -> b s 1 1"), 0.0)
+    if return_lse:
+        return output.to(dtype_og), attention.to(dtype_og), lse.to(dtype_og)
+    return output.to(dtype=dtype_og), attention.to(dtype=dtype_og)
+
+
+def maybe_fake_tensor_mode(fake: bool = True):
+    """
+    One way to populate/pre-compile cache is to use torch fake tensor mode,
+    which does not allocate actual GPU tensors but retains tensor shape/dtype
+    metadata for cute.compile.
+    """
+
+    def decorator(fn):
+        @wraps(fn)
+        def wrapper(*args, **kwargs):
+            with FakeTensorMode() if fake else nullcontext():
+                return fn(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def is_fake_mode() -> bool:
+    return active_fake_mode() is not None

python/sglang/jit_kernel/flash_attn/cute/topk_gather_kv.py

@@ -0,0 +1,289 @@
+import math
+import operator
+from dataclasses import dataclass
+from typing import Optional, Type
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+from cutlass import Boolean, Int32, Uint32, const_expr
+from cutlass.cute.nvgpu import cpasync
+from quack.cute_dsl_utils import ParamsBase
+
+from sglang.jit_kernel.flash_attn.cute import utils
+from sglang.jit_kernel.flash_attn.cute.utils import warp_reduce
+
+
+@dataclass
+class CpasyncGatherKVManager(ParamsBase):
+    mIndexTopk: cute.Tensor
+    sBitmask: Optional[cute.Tensor]
+
+    cta_rank_in_cluster: Int32
+    thread_idx: Int32
+    warp_idx: Int32
+
+    topk_length: Int32
+    seqlen_k_limit: Int32
+    tile_n: Int32
+    num_threads: cutlass.Constexpr[Int32]
+    hdim: cutlass.Constexpr[Int32]
+    hdim_v: cutlass.Constexpr[Int32]
+    num_hdimv_splits: cutlass.Constexpr[Int32]
+    cta_group_size: cutlass.Constexpr[Int32]
+
+    gmem_threads_per_row: cutlass.Constexpr[Int32]
+    topk_indices_per_thread: Int32
+    async_copy_elems: Int32
+
+    gmem_tiled_copy_KV: cute.TiledCopy
+    gmem_thr_copy_KV: cute.TiledCopy
+
+    rTopk: cute.Tensor
+    rTopkHalf: cute.Tensor
+    # for bitmask
+    rTopk_NonInterleaved: cute.Tensor
+
+    pipeline_bitmask: Optional[pipeline.PipelineAsync]
+    cpasync_barrier: Optional[pipeline.NamedBarrier]
+
+    disable_bitmask: cutlass.Constexpr[Boolean]
+
+    @staticmethod
+    def create(
+        mIndexTopk: cute.Tensor,
+        cta_rank_in_cluster: Int32,
+        thread_idx: Int32,
+        warp_idx: Int32,
+        topk_length: Int32,
+        seqlen_k_limit: Int32,
+        tile_n: cutlass.Constexpr[Int32],
+        hdim: cutlass.Constexpr[Int32],
+        hdim_v: cutlass.Constexpr[Int32],
+        num_hdimv_splits: cutlass.Constexpr[Int32],
+        num_threads: cutlass.Constexpr[Int32],
+        dtype: Type[cutlass.Numeric],
+        cta_group_size: cutlass.Constexpr[Int32],
+        cpasync_barrier: Optional[pipeline.NamedBarrier] = None,
+        disable_bitmask: cutlass.Constexpr[Boolean] = True,
+        sBitmask: Optional[cute.Tensor] = None,
+        pipeline_bitmask: Optional[pipeline.PipelineAsync] = None,
+    ):
+        assert tile_n % num_threads == 0
+        assert num_threads == 128
+        assert hdim % 64 == 0
+        assert (hdim_v // num_hdimv_splits // cta_group_size) % 64 == 0
+        assert num_threads % cute.arch.WARP_SIZE == 0
+        universal_copy_bits = 128
+        async_copy_elems = universal_copy_bits // dtype.width
+        dtype_bytes = dtype.width // 8
+        # assumes hdim is never part of transposed operand
+        gmem_k_block_size = math.gcd(
+            hdim,
+            hdim_v // num_hdimv_splits // cta_group_size,
+            128 // dtype_bytes,
+        )
+        assert gmem_k_block_size % async_copy_elems == 0
+        gmem_threads_per_row = gmem_k_block_size // async_copy_elems
+        assert cute.arch.WARP_SIZE % gmem_threads_per_row == 0
+        atom_async_copy = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+            dtype,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        thr_layout = cute.make_ordered_layout(
+            (num_threads // gmem_threads_per_row, gmem_threads_per_row),
+            order=(1, 0),
+        )
+        val_layout = cute.make_layout((1, async_copy_elems))
+        gmem_tiled_copy_KV = cute.make_tiled_copy_tv(
+            atom_async_copy, thr_layout, val_layout
+        )
+        gmem_thr_copy_KV = gmem_tiled_copy_KV.get_slice(thread_idx)
+        topk_indices_per_thread = tile_n // num_threads
+
+        rTopk = cute.make_rmem_tensor((topk_indices_per_thread,), Int32)
+        rTopkHalf = cute.make_rmem_tensor((topk_indices_per_thread,), Int32)
+        rTopk_NonInterleaved = cute.make_rmem_tensor((topk_indices_per_thread,), Int32)
+
+        return CpasyncGatherKVManager(
+            mIndexTopk,
+            sBitmask,
+            cta_rank_in_cluster,
+            thread_idx,
+            warp_idx,
+            topk_length,
+            seqlen_k_limit,
+            tile_n,
+            num_threads,
+            hdim,
+            hdim_v,
+            num_hdimv_splits,
+            cta_group_size,
+            gmem_threads_per_row,
+            topk_indices_per_thread,
+            async_copy_elems,
+            gmem_tiled_copy_KV,
+            gmem_thr_copy_KV,
+            rTopk,
+            rTopkHalf,
+            rTopk_NonInterleaved,
+            pipeline_bitmask,
+            cpasync_barrier,
+            disable_bitmask,
+        )
+
+    @cute.jit
+    def load_index_topk(
+        self,
+        n_block: Int32,
+        transpose: bool,
+    ):
+        entries_per_thread = self.topk_indices_per_thread
+        rTopk = self.rTopk if const_expr(transpose) else self.rTopkHalf
+
+        for i in cutlass.range_constexpr(entries_per_thread):
+            row = (
+                i * self.num_threads
+                + (self.thread_idx % self.gmem_threads_per_row)
+                * (self.num_threads // self.gmem_threads_per_row)
+                + (self.thread_idx // self.gmem_threads_per_row)
+            )
+            # need this if not offset in load_X
+            # if const_expr(not transpose):
+            #     row += self.cta_rank_in_cluster * (self.tile_n//self.cta_group_size)
+            #     row = row % self.tile_n
+            row_idx = n_block * self.tile_n + row
+            rTopk[i] = self.mIndexTopk[row_idx]
+
+            if const_expr(not transpose and not self.disable_bitmask):
+                row_non_interleaved = i * self.num_threads + self.thread_idx
+                row_idx_non_interleaved = n_block * self.tile_n + row_non_interleaved
+                self.rTopk_NonInterleaved[0] = self.mIndexTopk[row_idx_non_interleaved]
+
+    @cute.jit
+    def compute_bitmask(
+        self,
+        producer_state_bitmask,
+    ):
+        assert self.pipeline_bitmask is not None, "pipeline_bitmask not provided"
+        assert self.cpasync_barrier is not None, "cpasync barrier not provided"
+
+        lane_idx = cute.arch.lane_idx()
+        assert cute.size(self.rTopk_NonInterleaved) == 1
+        bitmask = Uint32(0)
+
+        # Step 1. Construct per-thread bitmask
+        topk_idx = self.rTopk_NonInterleaved[0]
+        is_valid = topk_idx >= 0 and topk_idx < self.seqlen_k_limit
+        if is_valid:
+            bitmask = Uint32(1 << lane_idx)
+
+        # Step 2. Warp shuffle bitwise OR = add since indices are exclusive.
+        bitmask = warp_reduce(bitmask, operator.add)
+
+        self.pipeline_bitmask.producer_acquire(producer_state_bitmask)
+        # store to smem and sync threads
+        if lane_idx == 0:
+            self.sBitmask[self.warp_idx, producer_state_bitmask.index] = bitmask
+        self.cpasync_barrier.arrive_and_wait()
+
+        self.pipeline_bitmask.producer_commit(producer_state_bitmask)
+        producer_state_bitmask.advance()
+        return producer_state_bitmask
+
+    @cute.jit
+    def compute_X_ptr(
+        self,
+        mX: cute.Tensor,
+        transpose: bool,
+        d_offset: int = 0,
+    ):
+        entries_per_thread = self.topk_indices_per_thread
+        tPrXPtr = cute.make_rmem_tensor((entries_per_thread,), cutlass.Int64)
+        tPrRowValid = cute.make_rmem_tensor((entries_per_thread,), cutlass.Int32)
+        rTopk = self.rTopk if const_expr(transpose) else self.rTopkHalf
+
+        for i in cutlass.range_constexpr(entries_per_thread):
+            topk_idx = rTopk[i]
+            if const_expr(not self.disable_bitmask):
+                row_valid = topk_idx >= 0 and topk_idx < self.seqlen_k_limit
+                tPrRowValid[i] = row_valid
+            if const_expr(not transpose):
+                tPrXPtr[i] = utils.elem_pointer(mX, (topk_idx, d_offset)).toint()
+            else:
+                tPrXPtr[i] = utils.elem_pointer(mX, (d_offset, topk_idx)).toint()
+
+        return tPrXPtr, tPrRowValid
+
+    @cute.jit
+    def load_X(
+        self,
+        mX: cute.Tensor,
+        sX: cute.Tensor,
+        transpose: bool,
+        K_or_V: str,
+        d_offset: int = 0,
+    ):
+        assert K_or_V in ("K", "V")
+        cta_tile_n = (
+            self.tile_n if const_expr(transpose) else self.tile_n // self.cta_group_size
+        )
+        head_dim = (
+            self.hdim
+            if const_expr(K_or_V == "K")
+            else self.hdim_v // self.num_hdimv_splits
+        )
+        if const_expr(transpose):
+            head_dim = head_dim // self.cta_group_size
+        order = (1, 0) if const_expr(transpose) else (0, 1)
+
+        sX_nd_layout = cute.make_ordered_layout((cta_tile_n, head_dim), order=order)
+        sX_nd = cute.composition(sX, sX_nd_layout)
+
+        cX = cute.make_identity_tensor((cta_tile_n, head_dim))
+        tXsX = self.gmem_thr_copy_KV.partition_D(sX_nd)
+        tXcX = self.gmem_thr_copy_KV.partition_S(cX)
+
+        tPrXPtr, tPrRowValid = self.compute_X_ptr(mX, transpose, d_offset)
+
+        if const_expr(not transpose):
+            offset = self.cta_rank_in_cluster * (
+                self.gmem_threads_per_row // self.cta_group_size
+            )
+        else:
+            offset = 0
+
+        for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])):
+            if const_expr(not self.disable_bitmask):
+                row_valid = utils.shuffle_sync(
+                    tPrRowValid[m // self.gmem_threads_per_row],
+                    (m + offset) % self.gmem_threads_per_row,
+                    width=self.gmem_threads_per_row,
+                )
+                should_load = cute.make_fragment_like(tXsX[(0, None), m, 0], Boolean)
+                should_load.fill(Boolean(row_valid))
+            x_ptr_i64 = utils.shuffle_sync(
+                tPrXPtr[m // self.gmem_threads_per_row],
+                (m + offset) % self.gmem_threads_per_row,
+                width=self.gmem_threads_per_row,
+            )
+            x_gmem_ptr = cute.make_ptr(
+                mX.element_type, x_ptr_i64, cute.AddressSpace.gmem, assumed_align=16
+            )
+            mX_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,)))
+            mX_cur_copy = cute.tiled_divide(mX_cur, (self.async_copy_elems,))
+
+            for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])):
+                ki = tXcX[0, 0, k][1] // self.async_copy_elems
+                mX_cur_copy_ki = mX_cur_copy[None, ki]
+                tXsX_k = tXsX[None, m, k]
+                mX_cur_copy_ki = cute.make_tensor(
+                    mX_cur_copy_ki.iterator, tXsX_k.layout
+                )
+                cute.copy(
+                    self.gmem_tiled_copy_KV,
+                    mX_cur_copy_ki,
+                    tXsX_k,
+                    pred=should_load if const_expr(not self.disable_bitmask) else None,
+                )

python/sglang/jit_kernel/inkling_all_reduce.py

@@ -0,0 +1,357 @@
+"""CUDA-JIT all-reduce kernels for Inkling symmetric-memory buffers.
+
+The producer writes its local shard into the symmetric buffer, and the reduced
+result remains there so callers do not need staging or copy-out kernels.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, empty_sentinel, load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_inkling_all_reduce_module(dtype: torch.dtype, world_size: int) -> Module:
+    args = make_cpp_args(dtype, world_size)
+    return load_jit(
+        "inkling_all_reduce",
+        *args,
+        cuda_files=["inkling/inkling_all_reduce.cuh"],
+        cuda_wrappers=[
+            ("two_shot_all_reduce", f"inkling_two_shot_all_reduce<{args}>"),
+            ("two_shot_all_reduce_fused", f"inkling_two_shot_all_reduce_fused<{args}>"),
+            ("multimem_one_shot_fused", f"inkling_multimem_one_shot_fused<{args}>"),
+            ("multimem_full_oneshot", f"inkling_multimem_full_oneshot<{args}>"),
+            ("multimem_push_oneshot", f"inkling_multimem_push_oneshot<{args}>"),
+        ],
+    )
+
+
+# Barrier resources for the fused kernels:
+#   * flags: a DEDICATED symmetric uint32 buffer, zero-initialized once at
+#     setup: `world_size` single-leader slots (one per peer), then
+#     world_size * MAX_BARRIER_BLOCKS per-(writer, block) slots for the
+#     per-block barrier (v5's multi-block flavor).
+#   * state: a device-LOCAL uint32 buffer: the 5 words
+#     [arrival0, arrival1, release0, release1, xepoch] padded to 8, then
+#     MAX_BARRIER_BLOCKS per-block epochs; persists across calls and advances
+#     under CUDA-graph replay.
+# Keep these sizes aligned with the CUDA barrier implementation.
+MAX_BARRIER_BLOCKS = 256
+STATE_SIZE = 8 + MAX_BARRIER_BLOCKS
+
+
+def flags_numel(world_size: int) -> int:
+    return world_size * (1 + MAX_BARRIER_BLOCKS)
+
+
+# Tuned (kernel, num_blocks, block_size) per reduction row count. Kernels:
+# "v5"=push one-shot with per-block barriers (single barrier, out-of-place),
+# "mm"=torch multimem, "v2"=two-shot explicit,
+# "v3"=two-shot multimem (single-leader barriers), "v3b"=v3 with per-block
+# barriers, and "v4"=full one-shot. nb/bs are 0 for "mm". Tables are keyed
+# by world size; TP4 is the fallback.
+_AR_TUNED_TP4 = {
+    1: ("v5", 1, 1024),
+    2: ("v5", 1, 1024),
+    3: ("v5", 8, 512),
+    4: ("v5", 8, 512),
+    6: ("v5", 8, 1024),
+    8: ("v5", 8, 1024),
+    12: ("v5", 8, 512),
+    16: ("v5", 8, 512),
+    24: ("v5", 8, 1024),
+    32: ("v5", 8, 1024),
+    48: ("v5", 48, 1024),
+    64: ("v5", 48, 1024),
+    96: ("v5", 64, 1024),
+    128: ("mm", 0, 0),
+    192: ("mm", 0, 0),
+    256: ("v3b", 64, 1024),
+    384: ("v3b", 32, 1024),
+    512: ("v3b", 32, 1024),
+    768: ("v3b", 48, 512),
+    1024: ("v3b", 32, 1024),
+    1536: ("v3", 64, 512),
+    2048: ("v3", 64, 512),
+    3072: ("v3", 96, 512),
+    4096: ("v3", 96, 512),
+    6144: ("v3", 64, 512),
+    8192: ("v3", 32, 1024),
+    12288: ("v3", 96, 512),
+    16384: ("v3", 96, 512),
+}
+# TP8 uses full one-shot for the smallest shapes, multimem through the
+# medium-sized range, and two-shot multimem for larger reductions.
+_AR_TUNED_TP8 = {
+    1: ("v4", 1, 1024),
+    2: ("v4", 1, 1024),
+    3: ("mm", 0, 0),
+    4: ("mm", 0, 0),
+    6: ("mm", 0, 0),
+    8: ("mm", 0, 0),
+    12: ("mm", 0, 0),
+    16: ("mm", 0, 0),
+    24: ("mm", 0, 0),
+    32: ("mm", 0, 0),
+    48: ("mm", 0, 0),
+    64: ("mm", 0, 0),
+    96: ("mm", 0, 0),
+    128: ("mm", 0, 0),
+    192: ("mm", 0, 0),
+    256: ("mm", 0, 0),
+    384: ("mm", 0, 0),
+    512: ("mm", 0, 0),
+    768: ("mm", 0, 0),
+    1024: ("v3", 32, 512),
+    1536: ("v3", 16, 1024),
+    2048: ("v3", 32, 512),
+    3072: ("v3", 48, 512),
+    4096: ("v3", 48, 512),
+    6144: ("v3", 64, 512),
+    8192: ("v3", 96, 256),
+    12288: ("v3", 64, 512),
+    16384: ("v3", 64, 512),
+}
+_AR_TUNED = {4: _AR_TUNED_TP4, 8: _AR_TUNED_TP8}
+_AR_TUNED_TOKENS = sorted(_AR_TUNED_TP4)  # same token grid for every table
+assert all(
+    set(t) == set(_AR_TUNED_TP4) for t in _AR_TUNED.values()
+), "all tuned tables must share the same token grid"
+
+
+def select_ar_config(num_tokens: int, world_size: int = 4):
+    """Return (kernel, num_blocks, block_size) for a ``[num_tokens, hidden]``
+    reduction, from the autotuned table for ``world_size`` (round up to the
+    nearest tested shape). Untuned world sizes fall back to the TP4 table.
+    ``kernel`` is one of "v5"/"v4"/"mm"/"v2"/"v3"/"v3b"."""
+    table = _AR_TUNED.get(world_size, _AR_TUNED_TP4)
+    for t in _AR_TUNED_TOKENS:
+        if num_tokens <= t:
+            return table[t]
+    return table[_AR_TUNED_TOKENS[-1]]
+
+
+def compile_inkling_all_reduce(dtype: torch.dtype, world_size: int) -> None:
+    """Warm the JIT module for (dtype, world_size) so the first call is cheap."""
+    _jit_inkling_all_reduce_module(dtype, world_size)
+
+
+def inkling_two_shot_all_reduce(
+    buffer: torch.Tensor,
+    peer_ptrs_dev: int,
+    rank: int,
+    world_size: int,
+    num_items: int,
+) -> None:
+    """Two-shot all-reduce in place over ``num_items`` elements of the symm buffer.
+
+    Args:
+        buffer: this rank's symm buffer (1D, contiguous, bf16), sliced to
+            ``num_items``; used for device/dtype validation. The producer must
+            have already written this rank's shard into it.
+        peer_ptrs_dev: ``hdl.buffer_ptrs_dev`` -- device address of the array of
+            ``world_size`` peer buffer base pointers.
+        rank: this rank within the TP group.
+        world_size: TP world size (compile-time template arg).
+        num_items: number of elements to reduce (multiple of 8 for bf16).
+
+    The caller is responsible for ``hdl.barrier()`` before (producers done) and
+    after (result visible) this call.
+    """
+    module = _jit_inkling_all_reduce_module(buffer.dtype, world_size)
+    module.two_shot_all_reduce(buffer, peer_ptrs_dev, rank, num_items)
+
+
+def inkling_two_shot_all_reduce_fused(
+    buffer: torch.Tensor,
+    data_ptrs_dev: int,
+    flag_ptrs_dev: int,
+    state_ptr: int,
+    rank: int,
+    world_size: int,
+    num_items: int,
+    num_blocks: int = 0,
+    block_size: int = 0,
+    shared: torch.Tensor | None = None,
+) -> None:
+    """Single-launch two-shot all-reduce with an in-kernel grid-level barrier.
+
+    Args:
+        buffer: this rank's symm data buffer (bf16), sliced to ``num_items``.
+        data_ptrs_dev: ``hdl.buffer_ptrs_dev`` for the data buffer.
+        flag_ptrs_dev: ``buffer_ptrs_dev`` of a DEDICATED symm ``uint32[world_size]``
+            flags buffer, zero-initialized once at setup.
+        state_ptr: ``data_ptr()`` of a device-local ``uint32[STATE_SIZE]`` barrier
+            state buffer (persists across calls; advances under graph replay).
+        rank, world_size: TP coordinates (world_size is a template arg).
+        num_items: elements to reduce (multiple of 8 for bf16).
+
+    No external barrier needed -- the kernel fences both sides itself.
+    """
+    module = _jit_inkling_all_reduce_module(buffer.dtype, world_size)
+    module.two_shot_all_reduce_fused(
+        buffer,
+        data_ptrs_dev,
+        flag_ptrs_dev,
+        state_ptr,
+        rank,
+        num_items,
+        num_blocks,
+        block_size,
+        shared if shared is not None else empty_sentinel(buffer.device, buffer.dtype),
+    )
+
+
+def inkling_multimem_one_shot_fused(
+    buffer: torch.Tensor,
+    multicast_ptr: int,
+    flag_ptrs_dev: int,
+    state_ptr: int,
+    rank: int,
+    world_size: int,
+    num_items: int,
+    num_blocks: int = 0,
+    block_size: int = 0,
+    per_block_barrier: bool = False,
+    shared: torch.Tensor | None = None,
+) -> None:
+    """Single-launch multimem one-shot all-reduce (NVLink multicast ld_reduce/st).
+
+    Matches torch multimem for tiny, latency-bound (decode) messages, in a kernel
+    we own so norm/sconv can fuse at the epilogue seam.
+
+    Args:
+        buffer: this rank's symm data buffer (bf16), sliced to ``num_items``.
+        multicast_ptr: ``hdl.multicast_ptr`` for the data buffer (must be != 0).
+        flag_ptrs_dev, state_ptr: dedicated barrier flags + local state buffer
+            (same as the fused two-shot).
+        rank, world_size, num_items: as above.
+        per_block_barrier: use per-block peer handshakes for both barriers (no
+            grid funnel; capped at MAX_BARRIER_BLOCKS blocks).
+    """
+    module = _jit_inkling_all_reduce_module(buffer.dtype, world_size)
+    module.multimem_one_shot_fused(
+        buffer,
+        multicast_ptr,
+        flag_ptrs_dev,
+        state_ptr,
+        rank,
+        num_items,
+        num_blocks,
+        block_size,
+        int(per_block_barrier),
+        shared if shared is not None else empty_sentinel(buffer.device, buffer.dtype),
+    )
+
+
+def inkling_multimem_full_oneshot(
+    in_buffer: torch.Tensor,
+    out_buffer: torch.Tensor,
+    multicast_ptr: int,
+    flag_ptrs_dev: int,
+    state_ptr: int,
+    rank: int,
+    world_size: int,
+    num_items: int,
+    num_blocks: int = 0,
+    block_size: int = 0,
+    shared: torch.Tensor | None = None,
+) -> None:
+    """Full one-shot all-reduce with a SINGLE (entry-only) barrier.
+
+    Every rank ld_reduces the entire range (multicast hardware sum) into its
+    local ``out_buffer`` -- no broadcast, no exit barrier. Fastest for tiny
+    latency-bound messages, but the caller MUST double-buffer ``in_buffer`` (its
+    reuse is not fenced by this kernel; the next AR's entry barrier orders it).
+
+    Args:
+        in_buffer: this rank's symm data buffer (bf16), sliced to ``num_items``.
+        out_buffer: local output buffer (bf16, >= num_items); receives the sum.
+        multicast_ptr: ``hdl.multicast_ptr`` of the in_buffer.
+        flag_ptrs_dev, state_ptr: barrier flags + local state (as above).
+        rank, world_size, num_items: as above.
+    """
+    module = _jit_inkling_all_reduce_module(in_buffer.dtype, world_size)
+    module.multimem_full_oneshot(
+        in_buffer,
+        out_buffer,
+        multicast_ptr,
+        flag_ptrs_dev,
+        state_ptr,
+        rank,
+        num_items,
+        num_blocks,
+        block_size,
+        (
+            shared
+            if shared is not None
+            else empty_sentinel(in_buffer.device, in_buffer.dtype)
+        ),
+    )
+
+
+def inkling_multimem_push_oneshot(
+    in_buffer: torch.Tensor,
+    out_buffer: torch.Tensor,
+    mc_stage_ptr: int,
+    local_stage_ptr: int,
+    flag_ptrs_dev: int,
+    state_ptr: int,
+    rank: int,
+    world_size: int,
+    num_items: int,
+    num_blocks: int = 0,
+    block_size: int = 0,
+    per_block_barrier: bool = False,
+    shared: torch.Tensor | None = None,
+) -> None:
+    """One-shot PUSH all-reduce (v5) with a SINGLE mid barrier.
+
+    Each rank multicast-stores its full input into its per-rank slot of the
+    symmetric staging area (slot ``r`` at elem offset ``r * num_items``), the
+    barrier waits for all pushes to land, then each rank reduces the
+    ``world_size`` staged shards locally (fp32 accum) into ``out_buffer``.
+    Drops one barrier round trip vs the two-shot kernels, and each rank holds
+    the full row at the epilogue seam (norm-fusion base, like v4 but scaling
+    past 2 rows).
+
+    Args:
+        in_buffer: this rank's LOCAL input (any contiguous 16B-aligned bf16
+            tensor -- need not be a symm buffer; it is only read locally).
+        out_buffer: local output buffer (bf16, >= num_items); receives the sum.
+        mc_stage_ptr: multicast address of the staging area (>= world_size *
+            num_items elems). The caller MUST double-buffer the staging area
+            (A/B rotation; the next AR's barrier orders the reuse, like v4).
+        local_stage_ptr: this GPU's local address of the same staging area.
+        flag_ptrs_dev, state_ptr: barrier flags + local state (as above).
+        rank, world_size, num_items: as above.
+        per_block_barrier: use the per-block peer handshake (no grid funnel;
+            capped at MAX_BARRIER_BLOCKS blocks) instead of the single-leader
+            grid barrier -- the multi-block latency winner.
+    """
+    module = _jit_inkling_all_reduce_module(in_buffer.dtype, world_size)
+    module.multimem_push_oneshot(
+        in_buffer,
+        out_buffer,
+        mc_stage_ptr,
+        local_stage_ptr,
+        flag_ptrs_dev,
+        state_ptr,
+        rank,
+        num_items,
+        num_blocks,
+        block_size,
+        int(per_block_barrier),
+        (
+            shared
+            if shared is not None
+            else empty_sentinel(in_buffer.device, in_buffer.dtype)
+        ),
+    )

python/sglang/jit_kernel/inkling_ar_fused.py

@@ -0,0 +1,222 @@
+"""Fused all-reduce, decode short-convolution, and RMSNorm for Inkling.
+
+The small-batch decode kernel processes one token per block.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, empty_sentinel, load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_ar_fused_module(
+    dtype: torch.dtype,
+    world_size: int,
+    w: int,
+    use_silu: bool,
+    use_residual: bool,
+    do_track: bool,
+) -> Module:
+    args = make_cpp_args(dtype, world_size, w, use_silu, use_residual, do_track)
+    return load_jit(
+        "inkling_ar_fused_decode",
+        *args,
+        cuda_files=["inkling/inkling_ar_fused_decode.cuh"],
+        cuda_wrappers=[
+            ("ar_sconv_norm", f"ArSconvNormKernel<{args}>::run"),
+            ("ar_sconv_norm_verify", f"ArSconvNormVerifyKernel<{args}>::run"),
+        ],
+    )
+
+
+# Tuned vectors per thread by decode row count; round up to the next entry.
+_FUSED_VPT_TUNED = {1: 1, 2: 1, 4: 1, 8: 1, 16: 1, 32: 1, 64: 1, 96: 1}
+_FUSED_VPT_TOKENS = sorted(_FUSED_VPT_TUNED)
+
+
+def select_fused_vpt(num_tokens: int) -> int:
+    for t in _FUSED_VPT_TOKENS:
+        if num_tokens <= t:
+            return _FUSED_VPT_TUNED[t]
+    return _FUSED_VPT_TUNED[_FUSED_VPT_TOKENS[-1]]
+
+
+def compile_inkling_ar_sconv_norm(
+    dtype: torch.dtype,
+    world_size: int,
+    w: int,
+    use_silu: bool,
+    use_residual: bool,
+    do_track: bool,
+) -> None:
+    """Warm the JIT module so the first fused call is cheap."""
+    _jit_ar_fused_module(dtype, world_size, w, use_silu, use_residual, do_track)
+
+
+def inkling_ar_sconv_norm(
+    in_partial: torch.Tensor,
+    residual_in: torch.Tensor,
+    residual_out: torch.Tensor,
+    hs_out: torch.Tensor,
+    norm_weight: torch.Tensor,
+    eps: float,
+    sconv_cache: torch.Tensor,
+    cache_indices: torch.Tensor,
+    cache_mask: torch.Tensor,
+    conv_weight: torch.Tensor,
+    mc_stage_ptr: int,
+    local_stage_ptr: int,
+    flag_ptrs_dev: int,
+    state_ptr: int,
+    rank: int,
+    world_size: int,
+    activation: str | None = None,
+    use_residual: bool = True,
+    track_mask: torch.Tensor | None = None,
+    track_indices: torch.Tensor | None = None,
+    enable_pdl: bool = True,
+    vecs_per_thread: int = 0,
+    shared: torch.Tensor | None = None,
+) -> None:
+    """Fused AR + decode sconv + add-RMSNorm over ``[T, D]`` decode rows.
+
+    Args:
+        in_partial: this rank's LOCAL partial sums (``[T, D]`` bf16, contiguous
+            rows, 16B-aligned) -- e.g. the MoE combine output with
+            ``reduce=False``. Read locally only (no stage-in copy).
+        shared: optional LOCAL ``[T, D]`` shared-expert partials, folded into
+            the pushed value in registers (fp32 add, one bf16 round --
+            torch.add numerics), replacing the separate ``routed + shared``
+            add kernel at zero extra traffic. All ranks must agree on passing it.
+        residual_in / residual_out: the residual stream before/after the fused
+            add (may alias); ``hs_out``: the normed output.
+        norm_weight, eps: RMSNorm gamma (``[D]`` bf16) and epsilon.
+        sconv_cache..conv_weight, track_*: exactly the tensors
+            ``fused_causal_conv1d_update_decode`` takes; the conv state is
+            shift-updated in place, identically to the unfused kernel.
+        mc_stage_ptr / local_stage_ptr: multicast + local address of the v5
+            staging rotation slot (>= world_size*T*D elems; caller rotates A/B,
+            same reuse-distance rule as v5).
+        flag_ptrs_dev / state_ptr / rank / world_size: barrier resources
+            (shared with the other fused AR kernels).
+    """
+    if activation == "swish":
+        activation = "silu"
+    use_silu = activation in ("silu", "swish")
+    do_track = track_mask is not None
+    w = conv_weight.shape[1]
+    if do_track:
+        tm = track_mask.reshape(-1)
+        ti = track_indices
+    else:  # dummies; DO_TRACK=false never reads them
+        tm = torch.empty(0, dtype=torch.bool, device=in_partial.device)
+        ti = torch.empty(0, dtype=torch.int64, device=in_partial.device)
+    module = _jit_ar_fused_module(
+        in_partial.dtype, world_size, w, use_silu, use_residual, do_track
+    )
+    if vecs_per_thread <= 0:
+        vecs_per_thread = select_fused_vpt(in_partial.shape[0])
+    sh = (
+        shared
+        if shared is not None
+        else empty_sentinel(in_partial.device, in_partial.dtype)
+    )
+    module.ar_sconv_norm(
+        in_partial,
+        residual_in,
+        residual_out,
+        hs_out,
+        norm_weight,
+        float(eps),
+        sconv_cache,
+        cache_indices,
+        cache_mask.reshape(-1),
+        conv_weight,
+        tm,
+        ti,
+        mc_stage_ptr,
+        local_stage_ptr,
+        flag_ptrs_dev,
+        state_ptr,
+        rank,
+        int(enable_pdl),
+        int(vecs_per_thread),
+        sh,
+    )
+
+
+def inkling_ar_sconv_norm_verify(
+    in_partial: torch.Tensor,
+    residual_in: torch.Tensor,
+    residual_out: torch.Tensor,
+    hs_out: torch.Tensor,
+    norm_weight: torch.Tensor,
+    eps: float,
+    sconv_cache: torch.Tensor,
+    cache_indices: torch.Tensor,
+    cache_mask: torch.Tensor,
+    conv_weight: torch.Tensor,
+    inter_out: torch.Tensor,
+    draft_token_num: int,
+    mc_stage_ptr: int,
+    local_stage_ptr: int,
+    flag_ptrs_dev: int,
+    state_ptr: int,
+    rank: int,
+    world_size: int,
+    activation: str | None = None,
+    use_residual: bool = True,
+    enable_pdl: bool = True,
+    shared: torch.Tensor | None = None,
+) -> None:
+    """Target-verify fused {AR -> causal_conv1d -> save_intermediate_conv_windows
+    -> add+RMSNorm} over ``[B*draft_token_num, D]`` rows.
+
+    ``cache_indices``/``cache_mask`` are per-SEQUENCE (``[B]``); the working
+    conv cache is read-only (the per-position windows go to ``inter_out``,
+    exactly like ``save_intermediate_conv_windows``). Cross-token conv taps are
+    re-reduced from the v5 staging slot, so the same rotation rules as
+    ``inkling_ar_sconv_norm`` apply. ``shared``: optional LOCAL ``[T, D]``
+    shared-expert partials folded into the push (torch.add numerics).
+    """
+    if activation == "swish":
+        activation = "silu"
+    use_silu = activation in ("silu", "swish")
+    w = conv_weight.shape[1]
+    # do_track slot in the module key is unused by the verify kernel.
+    module = _jit_ar_fused_module(
+        in_partial.dtype, world_size, w, use_silu, use_residual, False
+    )
+    sh = (
+        shared
+        if shared is not None
+        else empty_sentinel(in_partial.device, in_partial.dtype)
+    )
+    module.ar_sconv_norm_verify(
+        in_partial,
+        residual_in,
+        residual_out,
+        hs_out,
+        norm_weight,
+        float(eps),
+        sconv_cache,
+        cache_indices.to(torch.int32),
+        cache_mask,
+        conv_weight,
+        inter_out,
+        int(draft_token_num),
+        mc_stage_ptr,
+        local_stage_ptr,
+        flag_ptrs_dev,
+        state_ptr,
+        rank,
+        int(enable_pdl),
+        sh,
+    )

python/sglang/jit_kernel/inkling_ar_scattered_sconv.py

@@ -0,0 +1,349 @@
+"""Fused all-reduce and scattered short-convolution for Inkling.
+
+The kernel reduces a per-rank hidden-channel slice, applies causal convolution,
+and updates the convolution and prefix caches in one launch.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_ar_scattered_sconv_module(
+    dtype: torch.dtype,
+    world_size: int,
+    w: int,
+    use_silu: bool,
+    use_residual: bool,
+) -> Module:
+    args = make_cpp_args(dtype, world_size, w, use_silu, use_residual)
+    return load_jit(
+        "inkling_ar_scattered_sconv",
+        *args,
+        cuda_files=["inkling/inkling_ar_scattered_sconv.cuh"],
+        cuda_wrappers=[
+            ("ar_scattered_sconv", f"ArScatteredSconvKernel<{args}>::run"),
+            ("ar_banded_sconv", f"ArBandedSconvKernel<{args}>::run"),
+            ("ar_ssconv_norm_decode", f"SsconvNormDecodeKernel<{args}>::run"),
+            ("ar_col_decode", f"ColDecodeKernel<{args}>::run"),
+        ],
+    )
+
+
+def compile_inkling_ar_scattered_sconv(
+    dtype: torch.dtype,
+    world_size: int,
+    w: int,
+    use_silu: bool,
+    use_residual: bool,
+) -> None:
+    """Warm the JIT module so the first fused call is cheap."""
+    _jit_ar_scattered_sconv_module(dtype, world_size, w, use_silu, use_residual)
+
+
+def inkling_ar_scattered_sconv(
+    in_buffer: torch.Tensor,
+    x_scratch: torch.Tensor,
+    sconv_cache: torch.Tensor,
+    safe_idx: torch.Tensor,
+    cache_mask: torch.Tensor,
+    cache_indices: torch.Tensor,
+    has_initial_state: torch.Tensor,
+    cu: torch.Tensor,
+    si: torch.Tensor,
+    weight: torch.Tensor,
+    track_rows: torch.Tensor,
+    track_mask: torch.Tensor,
+    track_dst: torch.Tensor,
+    mc_in: int,
+    mc_out: int,
+    flag_ptrs_dev: int,
+    state_ptr: int,
+    rank: int,
+    world_size: int,
+    *,
+    activation: str | None,
+    use_residual: bool,
+    num_blocks: int = 0,
+    block_size: int = 0,
+    per_block_barrier: bool = False,
+    track_from_cache: bool = False,
+    out_local: torch.Tensor | None = None,
+    norm_gamma: torch.Tensor | None = None,
+    norm_residual: torch.Tensor | None = None,
+    norm_out: torch.Tensor | None = None,
+    norm_eps: float = 0.0,
+    need_scratch: bool = True,
+    use_stream: bool = False,
+    stream_walk: int = 0,
+    full_update: bool = False,
+    cache_col0: int = 0,
+) -> None:
+    """Run the fused kernel. ``in_buffer`` is this rank's [T, H] view of the
+    input symm region (partial sums already written by the producer);
+    ``mc_in`` / ``mc_out`` are the multicast pointers of the input and OUT
+    regions. On return the OUT region holds the gathered post-conv [T, H] on
+    every rank and ``x_scratch`` holds the reduced pre-conv [T, Hc] shard.
+
+    Tracking: empty ``track_mask`` disables it. ``track_from_cache`` (decode)
+    snapshots the post-update conv window to ``track_dst`` (``track_rows`` may
+    be empty); otherwise ``track_rows`` gathers pre-conv rows (extend).
+
+    Fused add+RMSNorm tail (decode/verify): pass ``out_local`` (this rank's
+    [T, H] OUT view), ``norm_gamma``/``norm_residual``/``norm_out``/``norm_eps``.
+    Works under either barrier mode. The residual is updated in place;
+    ``norm_out`` receives the normed hidden.
+
+    FULL-WIDTH mode (non-scattered sconv): ``full_update=True`` with
+    ``sconv_cache`` the replicated [slots, W-1, H] tensor, ``weight`` this
+    rank's contiguous [Hc, W] row slice and ``cache_col0 = rank * Hc``. Conv
+    still runs column-sharded; phase 3 updates/tracks ALL H cache columns on
+    every rank (window rows re-ld_reduced full-width) so the replicated cache
+    stays coherent. Verify (``need_scratch``) is unsupported full-width."""
+    w = weight.shape[1]
+    use_silu = activation in ("silu", "swish")
+    module = _jit_ar_scattered_sconv_module(
+        in_buffer.dtype, world_size, w, use_silu, use_residual
+    )
+    do_norm = norm_gamma is not None
+    if do_norm:
+        assert (
+            out_local is not None and norm_residual is not None and norm_out is not None
+        )
+    else:
+        empty = in_buffer.new_empty((0,))
+        out_local = norm_gamma = norm_residual = norm_out = empty
+    module.ar_scattered_sconv(
+        in_buffer,
+        x_scratch,
+        sconv_cache,
+        safe_idx,
+        cache_mask,
+        cache_indices,
+        has_initial_state,
+        cu,
+        si,
+        weight,
+        track_rows,
+        track_mask,
+        track_dst,
+        out_local,
+        norm_gamma,
+        norm_residual,
+        norm_out,
+        mc_in,
+        mc_out,
+        flag_ptrs_dev,
+        state_ptr,
+        rank,
+        num_blocks,
+        block_size,
+        per_block_barrier,
+        track_from_cache,
+        norm_eps,
+        need_scratch,
+        use_stream,
+        stream_walk,
+        full_update,
+        cache_col0,
+    )
+
+
+def inkling_ar_ssconv_norm_decode(
+    in_partials: torch.Tensor,
+    residual_in: torch.Tensor,
+    residual_out: torch.Tensor,
+    hs_out: torch.Tensor,
+    norm_weight: torch.Tensor,
+    norm_eps: float,
+    sconv_cache: torch.Tensor,
+    cache_indices: torch.Tensor,
+    cache_mask: torch.Tensor,
+    conv_weight_full: torch.Tensor,
+    track_mask: torch.Tensor,
+    track_indices: torch.Tensor,
+    mc_stage: int,
+    local_stage: int,
+    mc_wstage: int,
+    local_wstage: int,
+    flag_ptrs_dev: int,
+    state_ptr: int,
+    rank: int,
+    world_size: int,
+    *,
+    activation: str | None,
+    use_residual: bool,
+    vecs_per_thread: int = 0,
+) -> None:
+    """ONE-SHOT decode {AR + scattered sconv + add-RMSNorm}: v5 push pattern
+    with the cache-window shard co-pushed so every rank convs full width from
+    ONE barrier. ``sconv_cache`` is the SHARDED [pool, W-1, Hc] cache (only
+    this rank's columns are updated/tracked); ``conv_weight_full`` must be the
+    UNSHARDED [D, W] taps. ``mc_stage``/``local_stage`` = one v5 rotation slot
+    ([world, T, D]); ``mc_wstage``/``local_wstage`` = a rotating [T, W-1, D]
+    window-staging half. Pass empty ``track_mask`` to disable tracking
+    (post-update-window snapshot semantics otherwise)."""
+    w = conv_weight_full.shape[1]
+    use_silu = activation in ("silu", "swish")
+    module = _jit_ar_scattered_sconv_module(
+        in_partials.dtype, world_size, w, use_silu, use_residual
+    )
+    module.ar_ssconv_norm_decode(
+        in_partials,
+        residual_in,
+        residual_out,
+        hs_out,
+        norm_weight,
+        norm_eps,
+        sconv_cache,
+        cache_indices,
+        cache_mask,
+        conv_weight_full,
+        track_mask,
+        track_indices,
+        mc_stage,
+        local_stage,
+        mc_wstage,
+        local_wstage,
+        flag_ptrs_dev,
+        state_ptr,
+        rank,
+        vecs_per_thread,
+    )
+
+
+def inkling_ar_col_decode(
+    in_buffer: torch.Tensor,
+    out_local: torch.Tensor,
+    residual_in: torch.Tensor,
+    residual_out: torch.Tensor,
+    hs_out: torch.Tensor,
+    norm_weight: torch.Tensor,
+    norm_eps: float,
+    sconv_cache: torch.Tensor,
+    cache_indices: torch.Tensor,
+    cache_mask: torch.Tensor,
+    weight_shard: torch.Tensor,
+    track_mask: torch.Tensor,
+    track_dst: torch.Tensor,
+    mc_in: int,
+    mc_out: int,
+    flag_ptrs_dev: int,
+    state_ptr: int,
+    rank: int,
+    world_size: int,
+    *,
+    activation: str | None,
+    use_residual: bool,
+    vecs_per_thread: int = 0,
+) -> None:
+    """Dedicated small-batch column decode: one block per token row, block-scoped two-round barriers, prefetch under the entry spin, conv from registers on the owner shard, inline cache update (+ decode track), and the full-row add+RMSNorm after the exit round. Decode-only (single-token sequences; every tap is cache prefix). in_buffer/out_local are this rank's views of the input/OUT symm regions."""
+    w = weight_shard.shape[1]
+    use_silu = activation in ("silu", "swish")
+    module = _jit_ar_scattered_sconv_module(
+        in_buffer.dtype, world_size, w, use_silu, use_residual
+    )
+    module.ar_col_decode(
+        in_buffer,
+        out_local,
+        residual_in,
+        residual_out,
+        hs_out,
+        norm_weight,
+        norm_eps,
+        sconv_cache,
+        cache_indices,
+        cache_mask,
+        weight_shard,
+        track_mask,
+        track_dst,
+        mc_in,
+        mc_out,
+        flag_ptrs_dev,
+        state_ptr,
+        rank,
+        vecs_per_thread,
+    )
+
+
+def inkling_ar_banded_sconv(
+    in_buffer: torch.Tensor,
+    scratch: torch.Tensor,
+    sconv_cache: torch.Tensor,
+    safe_idx: torch.Tensor,
+    cache_mask: torch.Tensor,
+    cache_indices: torch.Tensor,
+    has_initial_state: torch.Tensor,
+    cu: torch.Tensor,
+    si: torch.Tensor,
+    weight: torch.Tensor,
+    track_rows: torch.Tensor,
+    track_mask: torch.Tensor,
+    track_dst: torch.Tensor,
+    mc_in: int,
+    mc_out: int,
+    flag_ptrs_dev: int,
+    state_ptr: int,
+    rank: int,
+    world_size: int,
+    *,
+    activation: str | None,
+    use_residual: bool,
+    num_blocks: int = 0,
+    block_size: int = 0,
+    per_block_barrier: bool = False,
+    debug_phase: int = 0,
+    mc_wstage: int = 0,
+    local_wstage: int = 0,
+) -> None:
+    """Token-banded fused {v3 AR + sconv}: contiguous band slices (v3-class
+    switch-transaction efficiency), in-kernel conv-state update + track.
+    ``scratch`` must be [ceil(T/world) + W-1, H]. Pass empty (numel-0)
+    ``track_rows`` to disable the track path.
+
+    Full-width mode (``sconv_cache`` [pool, W-1, H], default): the production
+    {v3 AR + sconv} fusion; every rank keeps the complete cache.
+    SCATTERED mode (``sconv_cache`` [pool, W-1, H/world] + ``mc_wstage``/
+    ``local_wstage`` pointing at a [B, W-1, H] staging region): each rank
+    pushes its cache-window shard pre-barrier (full-width taps come from the
+    staging), convs its contiguous token band full-width, and updates/tracks
+    only its own cache columns. ``weight`` must be the FULL [H, W] taps."""
+    w = weight.shape[1]
+    use_silu = activation in ("silu", "swish")
+    module = _jit_ar_scattered_sconv_module(
+        in_buffer.dtype, world_size, w, use_silu, use_residual
+    )
+    module.ar_banded_sconv(
+        in_buffer,
+        scratch,
+        sconv_cache,
+        safe_idx,
+        cache_mask,
+        cache_indices,
+        has_initial_state,
+        cu,
+        si,
+        weight,
+        track_rows,
+        track_mask,
+        track_dst,
+        mc_in,
+        mc_out,
+        flag_ptrs_dev,
+        state_ptr,
+        rank,
+        num_blocks,
+        block_size,
+        per_block_barrier,
+        debug_phase,
+        mc_wstage,
+        local_wstage,
+    )

python/sglang/jit_kernel/inkling_attn_prologue.py

@@ -0,0 +1,400 @@
+"""Fused target-verify attention prologue: {k/v sconv + save_windows + qk-norm
++ KV-cache store} in one kernel (csrc/tml/inkling_attn_prologue_fused.cuh)."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    empty_sentinel,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_attn_prologue_module(
+    dtype: torch.dtype,
+    w: int,
+    use_silu: bool,
+    use_residual: bool,
+    use_mxfp8: bool,
+) -> Module:
+    args = make_cpp_args(
+        dtype, w, use_silu, use_residual, use_mxfp8, is_arch_support_pdl()
+    )
+    return load_jit(
+        "inkling_attn_prologue_fused",
+        *args,
+        cuda_files=["inkling/inkling_attn_prologue_fused.cuh"],
+        cuda_wrappers=[
+            ("attn_prologue", f"AttnPrologueKernel<{args}>::run"),
+            ("attn_prologue_decode", f"AttnPrologueDecodeKernel<{args}>::run"),
+            ("attn_prologue_extend", f"AttnPrologueExtendKernel<{args}>::run"),
+        ],
+    )
+
+
+def compile_inkling_attn_prologue(
+    dtype: torch.dtype,
+    w: int,
+    use_silu: bool,
+    use_residual: bool,
+    use_mxfp8: bool = False,
+) -> None:
+    _jit_attn_prologue_module(dtype, w, use_silu, use_residual, use_mxfp8)
+
+
+def inkling_attn_prologue_verify(
+    qkvr: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    cache_indices: torch.Tensor,
+    cache_mask: torch.Tensor,
+    k_weight: torch.Tensor,
+    v_weight: torch.Tensor,
+    k_inter: torch.Tensor,
+    v_inter: torch.Tensor,
+    q_gamma: torch.Tensor,
+    k_gamma: torch.Tensor,
+    eps: float,
+    loc: torch.Tensor,
+    k_buf: torch.Tensor,
+    v_buf: torch.Tensor,
+    q_off: int,
+    k_off: int,
+    v_off: int,
+    dq: int,
+    dkv: int,
+    draft_token_num: int,
+    activation: str | None = None,
+    use_residual: bool = True,
+    do_store: bool = True,
+    mxfp8_quant: bool = False,
+    sfk: torch.Tensor | None = None,
+    sfv: torch.Tensor | None = None,
+    page_size: int = 128,
+    log_scaling_tau: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Returns fresh contiguous (q_normed, k_normed, v_conv) [T, dq/dkv];
+    KV rows are also scattered into k_buf/v_buf at ``loc`` (the attention call
+    should pass save_kv_cache=False)."""
+    t = qkvr.shape[0]
+    if mxfp8_quant:
+        if dq % 128 != 0 or dkv % 128 != 0:
+            raise ValueError("MXFP8 fused prologue requires head_dim-aligned Q/K/V.")
+        if sfk is None or sfv is None:
+            raise ValueError("MXFP8 fused prologue requires K/V scale buffers.")
+        sf_shape = (k_buf.shape[0] // page_size, dkv // 128, 32, page_size // 32, 4)
+        if sfk.shape != sf_shape or sfv.shape != sf_shape:
+            raise ValueError(
+                "MXFP8 fused prologue requires interleaved K/V scale buffers "
+                f"with shape {sf_shape}, got {tuple(sfk.shape)} and {tuple(sfv.shape)}."
+            )
+        if not sfk.is_contiguous() or not sfv.is_contiguous():
+            raise ValueError(
+                "MXFP8 fused prologue requires contiguous interleaved SFK/SFV."
+            )
+        q_out = torch.empty(t, dq, dtype=torch.float8_e4m3fn, device=qkvr.device)
+        sfq_u8 = torch.empty(
+            (t, dq // 128, 128 // 32), dtype=torch.uint8, device=qkvr.device
+        )
+        sfk_u8 = sfk.view(torch.uint8)
+        sfv_u8 = sfv.view(torch.uint8)
+    else:
+        q_out = torch.empty(t, dq, dtype=qkvr.dtype, device=qkvr.device)
+        sfq_u8 = torch.empty(0, dtype=torch.uint8, device=qkvr.device)
+        sfk_u8 = torch.empty(0, dtype=torch.uint8, device=qkvr.device)
+        sfv_u8 = torch.empty(0, dtype=torch.uint8, device=qkvr.device)
+    k_out = torch.empty(t, dkv, dtype=qkvr.dtype, device=qkvr.device)
+    v_out = torch.empty(t, dkv, dtype=qkvr.dtype, device=qkvr.device)
+    if activation == "swish":
+        activation = "silu"
+    use_silu = activation in ("silu", "swish")
+    w = k_weight.shape[1]
+    module = _jit_attn_prologue_module(
+        qkvr.dtype, w, use_silu, use_residual, mxfp8_quant
+    )
+    hkv = dkv // 128
+    module.attn_prologue(
+        qkvr,
+        k_cache,
+        v_cache,
+        cache_indices.to(torch.int32),
+        cache_mask,
+        k_weight,
+        v_weight,
+        k_inter,
+        v_inter,
+        q_gamma,
+        k_gamma,
+        float(eps),
+        q_out,
+        k_out,
+        v_out,
+        loc,
+        k_buf.view(-1, hkv * 128),
+        v_buf.view(-1, hkv * 128),
+        sfq_u8,
+        sfk_u8,
+        sfv_u8,
+        int(q_off),
+        int(k_off),
+        int(v_off),
+        int(draft_token_num),
+        int(do_store),
+        int(page_size),
+        (
+            log_scaling_tau.reshape(-1).float()
+            if log_scaling_tau is not None
+            else empty_sentinel(qkvr.device, torch.float32)
+        ),
+    )
+    q_scale = sfq_u8.view(torch.float8_e8m0fnu) if mxfp8_quant else None
+    return q_out, k_out, v_out, q_scale
+
+
+def inkling_attn_prologue_extend(
+    qkvr: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    cache_indices: torch.Tensor,
+    cache_mask: torch.Tensor,
+    has_initial_state: torch.Tensor,
+    cu: torch.Tensor,
+    si: torch.Tensor,
+    k_weight: torch.Tensor,
+    v_weight: torch.Tensor,
+    track_rows: torch.Tensor,
+    track_mask: torch.Tensor,
+    track_dst: torch.Tensor,
+    q_gamma: torch.Tensor,
+    k_gamma: torch.Tensor,
+    eps: float,
+    loc: torch.Tensor,
+    k_buf: torch.Tensor,
+    v_buf: torch.Tensor,
+    q_off: int,
+    k_off: int,
+    v_off: int,
+    dq: int,
+    dkv: int,
+    activation: str | None = None,
+    use_residual: bool = True,
+    do_store: bool = True,
+    mxfp8_quant: bool = False,
+    sfk: torch.Tensor | None = None,
+    sfv: torch.Tensor | None = None,
+    page_size: int = 128,
+    do_cache_update: bool = True,
+    log_scaling_tau: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Extend (prefill) analog of ``inkling_attn_prologue_verify``: varlen
+    sequences via ``cu``/``si``, no window save; instead a tiny trailing
+    kernel does the k/v conv-cache update at sequence ends (+ the extend
+    prefix-cache track when ``track_mask`` is non-empty -- pass empty tensors
+    to disable). Returns fresh contiguous (q_normed, k_normed, v_conv) and
+    scatters KV rows into k_buf/v_buf at ``loc`` when ``do_store`` (the
+    attention call should then pass save_kv_cache=False)."""
+    t = qkvr.shape[0]
+    if mxfp8_quant:
+        if dq % 128 != 0 or dkv % 128 != 0:
+            raise ValueError("MXFP8 fused prologue requires head_dim-aligned Q/K/V.")
+        if sfk is None or sfv is None:
+            raise ValueError("MXFP8 fused prologue requires K/V scale buffers.")
+        sf_shape = (k_buf.shape[0] // page_size, dkv // 128, 32, page_size // 32, 4)
+        if sfk.shape != sf_shape or sfv.shape != sf_shape:
+            raise ValueError(
+                "MXFP8 fused prologue requires interleaved K/V scale buffers "
+                f"with shape {sf_shape}, got {tuple(sfk.shape)} and {tuple(sfv.shape)}."
+            )
+        if not sfk.is_contiguous() or not sfv.is_contiguous():
+            raise ValueError(
+                "MXFP8 fused prologue requires contiguous interleaved SFK/SFV."
+            )
+        q_out = torch.empty(t, dq, dtype=torch.float8_e4m3fn, device=qkvr.device)
+        sfq_u8 = torch.empty(
+            (t, dq // 128, 128 // 32), dtype=torch.uint8, device=qkvr.device
+        )
+        sfk_u8 = sfk.view(torch.uint8)
+        sfv_u8 = sfv.view(torch.uint8)
+    else:
+        q_out = torch.empty(t, dq, dtype=qkvr.dtype, device=qkvr.device)
+        sfq_u8 = torch.empty(0, dtype=torch.uint8, device=qkvr.device)
+        sfk_u8 = torch.empty(0, dtype=torch.uint8, device=qkvr.device)
+        sfv_u8 = torch.empty(0, dtype=torch.uint8, device=qkvr.device)
+    k_out = torch.empty(t, dkv, dtype=qkvr.dtype, device=qkvr.device)
+    v_out = torch.empty(t, dkv, dtype=qkvr.dtype, device=qkvr.device)
+    if activation == "swish":
+        activation = "silu"
+    use_silu = activation in ("silu", "swish")
+    w = k_weight.shape[1]
+    module = _jit_attn_prologue_module(
+        qkvr.dtype, w, use_silu, use_residual, mxfp8_quant
+    )
+    hkv = dkv // 128
+    module.attn_prologue_extend(
+        qkvr,
+        k_cache,
+        v_cache,
+        cache_indices.to(torch.int32),
+        cache_mask,
+        has_initial_state,
+        cu,
+        si,
+        k_weight,
+        v_weight,
+        track_rows,
+        track_mask,
+        track_dst,
+        q_gamma,
+        k_gamma,
+        float(eps),
+        q_out,
+        k_out,
+        v_out,
+        loc,
+        k_buf.view(-1, hkv * 128),
+        v_buf.view(-1, hkv * 128),
+        sfq_u8,
+        sfk_u8,
+        sfv_u8,
+        int(q_off),
+        int(k_off),
+        int(v_off),
+        int(do_store),
+        int(page_size),
+        int(do_cache_update),
+        (
+            log_scaling_tau.reshape(-1).float()
+            if log_scaling_tau is not None
+            else empty_sentinel(qkvr.device, torch.float32)
+        ),
+    )
+    q_scale = sfq_u8.view(torch.float8_e8m0fnu) if mxfp8_quant else None
+    return q_out, k_out, v_out, q_scale
+
+
+def inkling_attn_prologue_decode(
+    qkvr: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    cache_indices: torch.Tensor,
+    cache_mask: torch.Tensor,
+    k_weight: torch.Tensor,
+    v_weight: torch.Tensor,
+    q_gamma: torch.Tensor,
+    k_gamma: torch.Tensor,
+    eps: float,
+    loc: torch.Tensor,
+    k_buf: torch.Tensor,
+    v_buf: torch.Tensor,
+    q_off: int,
+    k_off: int,
+    v_off: int,
+    dq: int,
+    dkv: int,
+    activation: str | None = None,
+    use_residual: bool = True,
+    track_mask: torch.Tensor | None = None,
+    track_indices: torch.Tensor | None = None,
+    do_store: bool = True,
+    mxfp8_quant: bool = False,
+    sfk: torch.Tensor | None = None,
+    sfv: torch.Tensor | None = None,
+    page_size: int = 128,
+    log_scaling_tau: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Decode {k/v decode-conv + conv-cache shift-update (+track) + qk-norm
+    (+ KV store)} in one kernel. Returns fresh (q_normed, k_normed, v_conv).
+    The k/v conv caches are shift-updated in place (fused_decode_update
+    semantics). With ``do_store`` the KV rows are scattered into k_buf/v_buf at
+    ``loc``; MXFP8 mode also quantizes Q and writes interleaved K/V scales."""
+    t = qkvr.shape[0]
+    if mxfp8_quant:
+        if dq % 128 != 0 or dkv % 128 != 0:
+            raise ValueError(
+                "MXFP8 fused decode prologue requires head_dim-aligned Q/K/V."
+            )
+        if sfk is None or sfv is None:
+            raise ValueError("MXFP8 fused decode prologue requires K/V scale buffers.")
+        sf_shape = (k_buf.shape[0] // page_size, dkv // 128, 32, page_size // 32, 4)
+        if sfk.shape != sf_shape or sfv.shape != sf_shape:
+            raise ValueError(
+                "MXFP8 fused decode prologue requires interleaved K/V scale buffers "
+                f"with shape {sf_shape}, got {tuple(sfk.shape)} and {tuple(sfv.shape)}."
+            )
+        if not sfk.is_contiguous() or not sfv.is_contiguous():
+            raise ValueError(
+                "MXFP8 fused decode prologue requires contiguous interleaved SFK/SFV."
+            )
+        q_out = torch.empty(t, dq, dtype=torch.float8_e4m3fn, device=qkvr.device)
+        sfq_u8 = torch.empty(
+            (t, dq // 128, 128 // 32), dtype=torch.uint8, device=qkvr.device
+        )
+        sfk_u8 = sfk.view(torch.uint8)
+        sfv_u8 = sfv.view(torch.uint8)
+    else:
+        q_out = torch.empty(t, dq, dtype=qkvr.dtype, device=qkvr.device)
+        sfq_u8 = torch.empty(0, dtype=torch.uint8, device=qkvr.device)
+        sfk_u8 = torch.empty(0, dtype=torch.uint8, device=qkvr.device)
+        sfv_u8 = torch.empty(0, dtype=torch.uint8, device=qkvr.device)
+    k_out = torch.empty(t, dkv, dtype=qkvr.dtype, device=qkvr.device)
+    v_out = torch.empty(t, dkv, dtype=qkvr.dtype, device=qkvr.device)
+    if activation == "swish":
+        activation = "silu"
+    use_silu = activation in ("silu", "swish")
+    w = k_weight.shape[1]
+    do_track = track_mask is not None
+    if do_track:
+        tm, ti = track_mask.reshape(-1), track_indices
+    else:
+        tm = torch.empty(0, dtype=torch.bool, device=qkvr.device)
+        ti = torch.empty(0, dtype=torch.int64, device=qkvr.device)
+    hkv = dkv // 128
+    module = _jit_attn_prologue_module(
+        qkvr.dtype, w, use_silu, use_residual, mxfp8_quant
+    )
+    module.attn_prologue_decode(
+        qkvr,
+        k_cache,
+        v_cache,
+        cache_indices.to(torch.int32),
+        cache_mask,
+        k_weight,
+        v_weight,
+        tm,
+        ti,
+        q_gamma,
+        k_gamma,
+        float(eps),
+        q_out,
+        k_out,
+        v_out,
+        loc,
+        k_buf.view(-1, hkv * 128),
+        v_buf.view(-1, hkv * 128),
+        sfq_u8,
+        sfk_u8,
+        sfv_u8,
+        int(q_off),
+        int(k_off),
+        int(v_off),
+        int(do_track),
+        int(do_store),
+        int(page_size),
+        (
+            log_scaling_tau.reshape(-1).float()
+            if log_scaling_tau is not None
+            else empty_sentinel(qkvr.device, torch.float32)
+        ),
+    )
+    q_scale = sfq_u8.view(torch.float8_e8m0fnu) if mxfp8_quant else None
+    return q_out, k_out, v_out, q_scale

python/sglang/jit_kernel/inkling_gate_topk_renorm.py

@@ -0,0 +1,325 @@
+"""Shape-specialized Inkling MoE gate top-k + renorm JIT kernels.
+
+Three families, all specialized for the Inkling gate layout (logits
+``[tokens, 258]`` fp32 = 256 routed + 2 shared experts, top-6 selection by
+``sigmoid(logit) + bias``, logsigmoid renorm over selected ++ shared):
+
+- ``inkling_gate_topk_renorm``      -- v1 warp-per-row gate (int64 indices).
+- ``inkling_gate_topk_renorm_v2``   -- v2 gate: wide vector loads, int32
+  indices, optional PDL, in-register raw-logit carry (no re-gather).
+- ``inkling_gate_gemv`` / ``inkling_gate_gemv_fused`` -- expert-per-block GEMV
+  of the gate linear (x [tokens, 6144] bf16 @ W [264, 6144] bf16 -> fp32
+  logits), standalone or with the gate epilogue fused into the same launch
+  (last finishing block runs it; ticket+workspace are cached per device).
+
+NOTE: the fused/gemv wrappers cache CUDA buffers and JIT-compile on first use;
+run them eagerly once before CUDA-graph capture.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+_LOGITS_PAD = 264  # fp32 logits row pitch shared with the padded gate GEMM
+_HIDDEN = 6144
+_TOPK = 6
+_N_SHARED = 2
+_FUSED_MAX_TOKENS = 64
+
+
+@cache_once
+def _jit_module() -> Module:
+    return load_jit(
+        "inkling_gate_topk_renorm",
+        "fast_math",
+        cuda_files=["moe/inkling_gate_topk_renorm.cuh"],
+        cuda_wrappers=[
+            ("inkling_gate_topk_renorm", "inkling_gate_topk_renorm"),
+            ("inkling_gate_topk_renorm_packed", "inkling_gate_topk_renorm_packed"),
+            ("inkling_gate_topk_renorm_v2", "inkling_gate_topk_renorm_v2"),
+            (
+                "inkling_gate_topk_renorm_v2_packed",
+                "inkling_gate_topk_renorm_v2_packed",
+            ),
+            ("inkling_gate_gemv", "inkling_gate_gemv"),
+            ("inkling_gate_gemv_fused", "inkling_gate_gemv_fused"),
+            ("inkling_gate_gemv_fused_packed", "inkling_gate_gemv_fused_packed"),
+        ],
+        extra_cuda_cflags=["-use_fast_math"],
+    )
+
+
+def _launch_inkling_gate_topk_renorm(
+    logits: torch.Tensor,
+    bias: torch.Tensor,
+    global_scale: torch.Tensor,
+    routed_w: torch.Tensor,
+    shared_w: torch.Tensor,
+    indices: torch.Tensor,
+    route_scale: float,
+) -> None:
+    module = _jit_module()
+    module.inkling_gate_topk_renorm(
+        logits, bias, global_scale, routed_w, shared_w, indices, float(route_scale)
+    )
+
+
+def _check_gate_inputs(
+    logits: torch.Tensor, bias: torch.Tensor, global_scale: torch.Tensor
+) -> None:
+    assert logits.is_cuda and logits.dtype == torch.float32 and logits.dim() == 2
+    assert logits.shape[1] == 258 and logits.stride(1) == 1
+    assert bias.is_cuda and bias.dtype == torch.float32 and bias.shape == (256,)
+    assert global_scale.is_cuda and global_scale.dtype == torch.float32
+    assert global_scale.numel() == 1
+
+
+def inkling_gate_topk_renorm(
+    logits: torch.Tensor,
+    bias: torch.Tensor,
+    global_scale: torch.Tensor,
+    route_scale: float,
+    *,
+    return_packed: bool = False,
+) -> (
+    tuple[torch.Tensor, torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor]
+):
+    """Select top-6 routed experts from 256 and renorm with 2 shared experts.
+
+    This is specialized for the Inkling fused gate layout:
+    ``logits`` is ``[tokens, 258]`` fp32, where columns ``0:256`` are routed
+    experts and columns ``256:258`` are shared experts. The top-k selection key is
+    ``sigmoid(logits[:, :256]) + bias``; renorm is over sigmoid(raw logits) for
+    the selected routed experts plus both shared experts.
+
+    ``return_packed=True`` emits the FlashInfer routed-MoE pack instead of the
+    routed_w + indices pair: ``packed[t,6]`` int32 = ``(expert_id << 16) | bf16
+    weight bits``. Returns ``(packed, shared_w)``.
+    """
+    _check_gate_inputs(logits, bias, global_scale)
+
+    tokens = logits.shape[0]
+    shared_w = torch.empty(
+        (tokens, _N_SHARED), dtype=torch.float32, device=logits.device
+    )
+    if return_packed:
+        packed = torch.empty((tokens, _TOPK), dtype=torch.int32, device=logits.device)
+        if tokens == 0:
+            return packed, shared_w
+        _jit_module().inkling_gate_topk_renorm_packed(
+            logits,
+            bias.contiguous(),
+            global_scale.contiguous(),
+            packed,
+            shared_w,
+            float(route_scale),
+        )
+        return packed, shared_w
+
+    routed_w = torch.empty((tokens, _TOPK), dtype=torch.float32, device=logits.device)
+    indices = torch.empty((tokens, _TOPK), dtype=torch.int64, device=logits.device)
+    if tokens == 0:
+        return routed_w, shared_w, indices
+
+    _launch_inkling_gate_topk_renorm(
+        logits,
+        bias.contiguous(),
+        global_scale.contiguous(),
+        routed_w,
+        shared_w,
+        indices,
+        route_scale,
+    )
+    return routed_w, shared_w, indices
+
+
+def inkling_gate_topk_renorm_v2(
+    logits: torch.Tensor,
+    bias: torch.Tensor,
+    global_scale: torch.Tensor,
+    route_scale: float,
+    *,
+    return_packed: bool = False,
+    enable_pdl: bool = False,
+    warps_per_block: int = 0,
+) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor, torch.Tensor | None]:
+    """v2 gate kernel; same math as v1 but int32 indices and optional PDL.
+
+    Returns ``(routed_w, indices, shared_w, packed)`` where the unused half is
+    ``None`` depending on ``return_packed`` -- mirroring the triton
+    ``sigmoid_gate_topk_renorm`` contract. ``warps_per_block`` in
+    ``{0 (auto), 1, 2, 4, 8}`` selects the launch shape.
+
+    Requires 32B-aligned logits rows: the production ``[tokens, 264]``-padded
+    GEMM output sliced to ``[:, :258]`` qualifies.
+    """
+    _check_gate_inputs(logits, bias, global_scale)
+    assert logits.stride(0) % 8 == 0, f"rows must be 32B-aligned: {logits.stride()=}"
+
+    tokens = logits.shape[0]
+    shared_w = torch.empty(
+        (tokens, _N_SHARED), dtype=torch.float32, device=logits.device
+    )
+    if return_packed:
+        packed = torch.empty((tokens, _TOPK), dtype=torch.int32, device=logits.device)
+        if tokens > 0:
+            _jit_module().inkling_gate_topk_renorm_v2_packed(
+                logits,
+                bias.contiguous(),
+                global_scale.contiguous(),
+                packed,
+                shared_w,
+                float(route_scale),
+                bool(enable_pdl),
+                int(warps_per_block),
+            )
+        return None, None, shared_w, packed
+
+    routed_w = torch.empty((tokens, _TOPK), dtype=torch.float32, device=logits.device)
+    indices = torch.empty((tokens, _TOPK), dtype=torch.int32, device=logits.device)
+    if tokens > 0:
+        _jit_module().inkling_gate_topk_renorm_v2(
+            logits,
+            bias.contiguous(),
+            global_scale.contiguous(),
+            routed_w,
+            shared_w,
+            indices,
+            float(route_scale),
+            bool(enable_pdl),
+            int(warps_per_block),
+        )
+    return routed_w, indices, shared_w, None
+
+
+def _check_gemv_inputs(x: torch.Tensor, weight: torch.Tensor) -> None:
+    assert x.is_cuda and x.dtype == torch.bfloat16 and x.dim() == 2
+    assert x.shape[1] == _HIDDEN and x.stride(1) == 1 and x.stride(0) == _HIDDEN
+    assert weight.is_cuda and weight.dtype == torch.bfloat16 and weight.dim() == 2
+    assert weight.shape[0] >= 258 and weight.shape[1] == _HIDDEN
+    assert weight.stride(1) == 1 and weight.stride(0) == _HIDDEN
+
+
+def inkling_gate_gemv(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    *,
+    enable_pdl: bool = False,
+    experts_per_block: int = 0,
+) -> torch.Tensor:
+    """Gate linear as an expert-per-block GEMV: returns fp32 logits [tokens, 258].
+
+    Drop-in for ``inkling_fused_gate_linear_with_fp32_out`` (the returned view
+    shares the same padded [tokens, 264] layout). Meant for small token counts
+    where the PDL split pair (this + v2 gate) beats cublas + gate.
+    """
+    _check_gemv_inputs(x, weight)
+    tokens = x.shape[0]
+    logits = torch.empty((tokens, _LOGITS_PAD), dtype=torch.float32, device=x.device)
+    if tokens > 0:
+        _jit_module().inkling_gate_gemv(
+            x, weight, logits, bool(enable_pdl), int(experts_per_block)
+        )
+    return logits[:, :258]
+
+
+# Per-device (workspace [64, 264] fp32, ticket int32[1]) reused by every fused
+# call. The kernel resets the ticket to zero on completion, so the buffers are
+# CUDA-graph replay-safe; allocate them eagerly (warmup) before graph capture.
+_fused_scratch: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
+
+
+def _get_fused_scratch(device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+    key = device.index if device.index is not None else torch.cuda.current_device()
+    scratch = _fused_scratch.get(key)
+    if scratch is None:
+        # Allocating inside CUDA graph capture would place the persistent
+        # buffers in the capture pool, where other graphs' replays can reuse
+        # (clobber) them. Call ensure_gate_gemv_fused_scratch() eagerly first
+        # (InklingGate.__init__ does).
+        assert (
+            not torch.cuda.is_current_stream_capturing()
+        ), "fused gate scratch must be allocated before CUDA graph capture"
+        workspace = torch.empty(
+            (_FUSED_MAX_TOKENS, _LOGITS_PAD), dtype=torch.float32, device=device
+        )
+        ticket = torch.zeros((1,), dtype=torch.int32, device=device)
+        scratch = (workspace, ticket)
+        _fused_scratch[key] = scratch
+    return scratch
+
+
+def ensure_gate_gemv_fused_scratch(device: torch.device) -> None:
+    """Eagerly allocate the fused-gate workspace/ticket (call at model init,
+    before any CUDA graph capture)."""
+    _get_fused_scratch(device)
+
+
+def inkling_gate_gemv_fused(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    global_scale: torch.Tensor,
+    route_scale: float,
+    *,
+    return_packed: bool = False,
+    enable_pdl: bool = False,
+    experts_per_block: int = 0,
+) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor, torch.Tensor | None]:
+    """Fully fused Inkling gate: GEMV + sigmoid+bias top-6 + renorm, one launch.
+
+    ``x`` is ``[tokens, 6144]`` bf16 (tokens <= 64), ``weight`` the padded
+    ``[264, 6144]`` bf16 gate weight. Output contract matches
+    ``sigmoid_gate_topk_renorm``: ``(routed_w, indices, shared_w, packed)``.
+    """
+    _check_gemv_inputs(x, weight)
+    tokens = x.shape[0]
+    assert tokens <= _FUSED_MAX_TOKENS, f"fused gate supports <= 64 tokens: {tokens=}"
+    assert bias.is_cuda and bias.dtype == torch.float32 and bias.shape == (256,)
+    assert global_scale.is_cuda and global_scale.dtype == torch.float32
+
+    workspace, ticket = _get_fused_scratch(x.device)
+    shared_w = torch.empty((tokens, _N_SHARED), dtype=torch.float32, device=x.device)
+    if return_packed:
+        packed = torch.empty((tokens, _TOPK), dtype=torch.int32, device=x.device)
+        if tokens > 0:
+            _jit_module().inkling_gate_gemv_fused_packed(
+                x,
+                weight,
+                bias.contiguous(),
+                global_scale.contiguous(),
+                workspace,
+                ticket,
+                packed,
+                shared_w,
+                float(route_scale),
+                bool(enable_pdl),
+                int(experts_per_block),
+            )
+        return None, None, shared_w, packed
+
+    routed_w = torch.empty((tokens, _TOPK), dtype=torch.float32, device=x.device)
+    indices = torch.empty((tokens, _TOPK), dtype=torch.int32, device=x.device)
+    if tokens > 0:
+        _jit_module().inkling_gate_gemv_fused(
+            x,
+            weight,
+            bias.contiguous(),
+            global_scale.contiguous(),
+            workspace,
+            ticket,
+            routed_w,
+            shared_w,
+            indices,
+            float(route_scale),
+            bool(enable_pdl),
+            int(experts_per_block),
+        )
+    return routed_w, indices, shared_w, None

python/sglang/jit_kernel/inkling_rel_proj.py

@@ -0,0 +1,53 @@
+"""CUDA-JIT latency-lean rel_logits projection for SMALL token counts, with
+the optional log-scaling tau prescale folded in registers. See
+csrc/tml/inkling_rel_proj.cuh; cuBLAS keeps everything above the measured
+small-t band (an earlier bandwidth-oriented custom kernel lost to it at every
+size)."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    empty_sentinel,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_rel_proj_module(d_rel: int, use_pdl: bool) -> Module:
+    args = make_cpp_args(d_rel, use_pdl)
+    return load_jit(
+        "inkling_rel_proj",
+        *args,
+        cuda_files=["inkling/inkling_rel_proj.cuh"],
+        cuda_wrappers=[("run", f"rel_proj_small_t<{args}>")],
+    )
+
+
+def rel_proj_small_t(
+    r: torch.Tensor,
+    proj: torch.Tensor,
+    tau: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """``r``: [t, h, d_rel] bf16, token rows possibly strided ((h*d_rel)-
+    contiguous inner, 16B-aligned); ``proj``: [d_rel, e] bf16 contiguous;
+    ``tau``: optional fp32 [t] prescale (rounds r*tau to bf16 before the dot,
+    the shipped prescale semantics). Returns contiguous [t, h, e]."""
+    if out is None:
+        out = torch.empty(
+            (r.shape[0], r.shape[1], proj.shape[1]), dtype=r.dtype, device=r.device
+        )
+    module = _jit_rel_proj_module(r.shape[2], is_arch_support_pdl())
+    sh = tau if tau is not None else empty_sentinel(r.device, torch.float32)
+    module.run(r, sh, proj, out)
+    return out

python/sglang/jit_kernel/inkling_row_scale.py

@@ -0,0 +1,59 @@
+"""CUDA-JIT vectorized per-row scale (the apply_log_scaling_tau contract):
+``out[row, :] = bf16(fp32(x[row, :]) * tau[row])``. See
+csrc/tml/inkling_row_scale.cuh; the scalar triton kernel remains the fallback
+for non-bf16 / unaligned inputs."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.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 _jit_row_scale_module(use_pdl: bool) -> Module:
+    args = make_cpp_args(use_pdl)
+    return load_jit(
+        "inkling_row_scale",
+        *args,
+        cuda_files=["inkling/inkling_row_scale.cuh"],
+        cuda_wrappers=[
+            ("run", f"row_scale<{args}>"),
+            ("run_compact", f"row_compact<{args}>"),
+        ],
+    )
+
+
+def row_scale_bf16(
+    x: torch.Tensor, tau: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """``x``: [rows, inner] bf16, possibly row-strided (inner contiguous,
+    inner % 8 == 0, 16B-aligned rows); ``tau``: fp32 [rows]. Returns a fresh
+    contiguous scaled tensor (bit-identical to the triton kernel's output)."""
+    if out is None:
+        out = torch.empty(x.shape, dtype=x.dtype, device=x.device)
+    module = _jit_row_scale_module(is_arch_support_pdl())
+    module.run(x, tau, out)
+    return out
+
+
+def row_compact_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Contiguous copy of row-strided ``x`` ([rows, inner] bf16, inner
+    contiguous, inner % 8 == 0, 16B-aligned rows) -- the tau-less flavor of
+    ``row_scale_bf16``. Beats the TensorIterator strided copy that einsum's
+    reshape would otherwise run on such inputs."""
+    if out is None:
+        out = torch.empty(x.shape, dtype=x.dtype, device=x.device)
+    module = _jit_row_scale_module(is_arch_support_pdl())
+    module.run_compact(x, out)
+    return out

python/sglang/jit_kernel/inkling_sconv.py

@@ -0,0 +1,239 @@
+"""CUDA-JIT implementations of the Inkling short-convolution kernels.
+
+Their signatures match the Triton entrypoints so model layers can select either
+backend without adapting arguments.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_causal_conv1d_module(
+    w: int,
+    use_silu: bool,
+    use_residual: bool,
+    is_decode: bool,
+    dtype: torch.dtype,
+) -> Module:
+    args = make_cpp_args(w, use_silu, use_residual, is_decode, dtype)
+    return load_jit(
+        "inkling_causal_conv1d",
+        *args,
+        cuda_files=["inkling/causal_conv1d.cuh"],
+        cuda_wrappers=[("causal_conv1d", f"CausalConv1dKernel<{args}>::run")],
+    )
+
+
+def causal_conv1d(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    sconv_cache: torch.Tensor,
+    cache_mask: torch.Tensor,
+    safe_idx: torch.Tensor,
+    cu: torch.Tensor,
+    si: torch.Tensor,
+    activation: str | None = None,
+    use_residual: bool = True,
+    is_decode: bool = False,
+) -> torch.Tensor:
+    """Apply depthwise causal convolution to a packed token stream.
+
+    Depthwise causal conv1d over a packed ``[T, D]`` token stream, with the W-1
+    prefix taps gathered directly from ``sconv_cache`` (no intermediate prefix
+    tensor). Metadata args (cache_mask, safe_idx, cu, si) are precomputed once
+    per forward pass and reused across layers.
+    """
+    if activation == "swish":
+        activation = "silu"
+
+    T = x.shape[0]
+    if T == 0:
+        return torch.empty_like(x)
+
+    D = x.shape[1]
+    W = weight.shape[1]
+    use_silu = activation in ("silu", "swish")
+
+    # Contiguous [T, D] output (strides (D, 1)) regardless of x's layout.
+    y = torch.empty(T, D, dtype=x.dtype, device=x.device)
+
+    module = _jit_causal_conv1d_module(W, use_silu, use_residual, is_decode, x.dtype)
+    module.causal_conv1d(x, sconv_cache, safe_idx, cache_mask, weight, cu, si, y)
+    return y
+
+
+@cache_once
+def _jit_update_sconv_cache_module(w1: int, dtype: torch.dtype) -> Module:
+    args = make_cpp_args(w1, dtype)
+    return load_jit(
+        "inkling_update_sconv_cache",
+        *args,
+        cuda_files=["inkling/update_sconv_cache.cuh"],
+        cuda_wrappers=[("update_sconv_cache", f"UpdateSconvCacheKernel<{args}>::run")],
+    )
+
+
+def update_sconv_cache(
+    x: torch.Tensor,
+    sconv_cache: torch.Tensor,
+    cache_indices: torch.Tensor,
+    has_initial_state: torch.Tensor,
+    query_start_loc: torch.Tensor,
+) -> None:
+    """Update each sequence's convolution cache in place.
+
+    Shift-updates each sequence's conv state to the last W-1 entries of
+    ``[old_state(gated) ++ x[start:end]]``; PAD / empty lanes are untouched. Pure
+    bit-exact select/copy.
+    """
+    W1 = sconv_cache.shape[1]
+    module = _jit_update_sconv_cache_module(W1, x.dtype)
+    module.update_sconv_cache(
+        x, sconv_cache, cache_indices, has_initial_state, query_start_loc
+    )
+
+
+@cache_once
+def _jit_gather_scatter_sconv_module(w1: int, dtype: torch.dtype) -> Module:
+    args = make_cpp_args(w1, dtype)
+    return load_jit(
+        "inkling_gather_scatter_sconv",
+        *args,
+        cuda_files=["inkling/gather_scatter_sconv.cuh"],
+        cuda_wrappers=[("gather_scatter", f"GatherScatterSconvKernel<{args}>::run")],
+    )
+
+
+def fused_gather_scatter_to_sconv_cache(
+    hidden_states: torch.Tensor,
+    sconv_cache: torch.Tensor,
+    track_conv_indices: torch.Tensor,
+    mask: torch.Tensor,
+    dst_indices: torch.Tensor,
+) -> None:
+    """Gather selected hidden-state rows into the convolution cache.
+
+    Scatters masked rows ``hidden_states[track_conv_indices[b, w]]`` into
+    ``sconv_cache[dst_indices[b], w]`` in-place; masked-out lanes untouched.
+    Bit-exact copy. (track int32, dst int64, per the model contract.)
+    """
+    W1 = sconv_cache.shape[1]
+    module = _jit_gather_scatter_sconv_module(W1, hidden_states.dtype)
+    module.gather_scatter(
+        hidden_states, sconv_cache, track_conv_indices, mask, dst_indices
+    )
+
+
+@cache_once
+def _jit_fused_decode_update_module(
+    w: int, use_silu: bool, use_residual: bool, do_track: bool, dtype: torch.dtype
+) -> Module:
+    args = make_cpp_args(w, use_silu, use_residual, do_track, dtype)
+    return load_jit(
+        "inkling_fused_decode_update",
+        *args,
+        cuda_files=["inkling/fused_decode_update.cuh"],
+        cuda_wrappers=[
+            ("fused_decode_update", f"FusedDecodeUpdateKernel<{args}>::run")
+        ],
+    )
+
+
+def fused_causal_conv1d_update_decode(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    sconv_cache: torch.Tensor,
+    cache_indices: torch.Tensor,
+    cache_mask: torch.Tensor,
+    activation: str | None = None,
+    use_residual: bool = True,
+    track_mask: torch.Tensor | None = None,
+    track_indices: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Apply decode convolution and update its cache in one kernel.
+
+    Decode conv (W-1 cached taps + current token) fused with the cache shift-update
+    (+ optional prefix-cache track-copy). Returns a contiguous ``[T, D]`` output.
+    """
+    if activation == "swish":
+        activation = "silu"
+    T, D = x.shape
+    W = weight.shape[1]
+    use_silu = activation in ("silu", "swish")
+    do_track = track_mask is not None
+
+    cm = cache_mask.reshape(-1)
+    y = torch.empty(T, D, dtype=x.dtype, device=x.device)
+    if do_track:
+        tm = track_mask.reshape(-1)
+        ti = track_indices
+    else:  # dummy tensors satisfy the signature; DO_TRACK=false never reads them
+        tm = torch.empty(0, dtype=torch.bool, device=x.device)
+        ti = torch.empty(0, dtype=torch.int64, device=x.device)
+
+    module = _jit_fused_decode_update_module(
+        W, use_silu, use_residual, do_track, x.dtype
+    )
+    module.fused_decode_update(x, sconv_cache, cache_indices, cm, weight, y, tm, ti)
+    return y
+
+
+@cache_once
+def _jit_draft_extend_sconv_module(
+    w1: int, do_track: bool, dtype: torch.dtype
+) -> Module:
+    args = make_cpp_args(w1, do_track, dtype)
+    return load_jit(
+        "inkling_draft_extend_sconv",
+        *args,
+        cuda_files=["inkling/draft_extend_sconv.cuh"],
+        cuda_wrappers=[("draft_extend", f"DraftExtendSconvKernel<{args}>::run")],
+    )
+
+
+def fused_draft_extend_sconv_cache(
+    hidden_states: torch.Tensor,
+    sconv_cache: torch.Tensor,
+    cache_indices: torch.Tensor,
+    num_accepted_tokens: torch.Tensor,
+    draft_token_num: int,
+    do_tracking: bool = False,
+    crossed: torch.Tensor | None = None,
+    track_step: torch.Tensor | None = None,
+    mamba_track_indices: torch.Tensor | None = None,
+) -> None:
+    """Update draft-extend convolution state in place.
+
+    Selects each sequence's length-(W-1) conv-state window from the virtual
+    ``[sconv_cache[ci] ++ hidden[b]]`` stream at ``num_accepted_tokens[b]`` (and, if
+    tracking, at ``track_step[b]`` into ``mamba_track_indices[b]`` where crossed).
+    Bit-exact copy.
+    """
+    W1 = sconv_cache.shape[1]
+    module = _jit_draft_extend_sconv_module(W1, do_tracking, hidden_states.dtype)
+    dev = hidden_states.device
+    if do_tracking:
+        cr, ts, mti = crossed, track_step, mamba_track_indices
+    else:  # dummies; DO_TRACK=false never reads them
+        cr = torch.empty(0, dtype=torch.bool, device=dev)
+        ts = torch.empty(0, dtype=torch.int32, device=dev)
+        mti = torch.empty(0, dtype=torch.int64, device=dev)
+    module.draft_extend(
+        hidden_states,
+        sconv_cache,
+        cache_indices,
+        num_accepted_tokens,
+        int(draft_token_num),
+        cr,
+        ts,
+        mti,
+    )

python/sglang/jit_kernel/tests/test_moe_preprocess.py

@@ -0,0 +1,93 @@
+"""fused_moe_preprocess must be bit-identical to the torch.sort-based path,
+and the grouped GEMM must produce identical results under both block_size_m
+configs (the block schedule and kernel config are chosen together).
+"""
+
+import pytest
+import torch
+
+from sglang.srt.layers.moe.moe_runner.triton_utils.inkling_moe import (
+    SMALL_M_BLOCK_SIZE_M,
+    compute_grouped_gemm_metadata,
+    fused_moe_preprocess,
+    get_src2dst,
+    grouped_gemm_triton,
+)
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
+
+E = 256
+TOPK = 6
+
+
+def _reference(topk_ids_flat: torch.Tensor):
+    reorder_topk_ids, reorder_ids = torch.sort(
+        topk_ids_flat.to(torch.int16), stable=True
+    )
+    src2dst = get_src2dst(reorder_ids)
+    meta = compute_grouped_gemm_metadata(
+        reorder_topk_ids, E, block_size_m=SMALL_M_BLOCK_SIZE_M
+    )
+    return (src2dst, *meta, reorder_topk_ids)
+
+
+def _ids(tokens: int, seed: int, skew: bool = False) -> torch.Tensor:
+    torch.manual_seed(seed)
+    if skew:  # all tokens on few experts (stresses multi-block experts)
+        return torch.randint(0, 3, (tokens * TOPK,), dtype=torch.int32, device="cuda")
+    return (
+        torch.stack([torch.randperm(E, device="cuda")[:TOPK] for _ in range(tokens)])
+        .view(-1)
+        .to(torch.int32)
+    )
+
+
+@requires_cuda
+@pytest.mark.parametrize("tokens", [1, 2, 7, 32, 64, 170, 341])  # n = 6*T <= 2048
+@pytest.mark.parametrize("skew", [False, True])
+def test_matches_sort_path(tokens: int, skew: bool):
+    ids = _ids(tokens, seed=tokens, skew=skew)
+    ref = _reference(ids)
+    got = fused_moe_preprocess(ids, E)
+    names = [
+        "src2dst",
+        "num_tokens_per_expert",
+        "expert_token_offs",
+        "expert_block_offs",
+        "expert_block_schedule",
+        "reorder_topk_ids",
+    ]
+    for tag, g, r in zip(names, got, ref):
+        assert g.shape == r.shape, (tag, g.shape, r.shape)
+        assert torch.equal(g.long(), r.long()), (
+            tag,
+            g[: min(16, g.numel())],
+            r[: min(16, r.numel())],
+        )
+
+
+@requires_cuda
+@pytest.mark.parametrize("tokens", [1, 16, 64])
+def test_grouped_gemm_small_config_matches(tokens: int):
+    """GEMM output must be identical whichever (block_size_m, config) runs."""
+    torch.manual_seed(tokens)
+    ids = _ids(tokens, seed=tokens)
+    m, k, n = tokens * TOPK, 768, 1024
+    a = (torch.randn(m, k, device="cuda") * 0.05).to(torch.bfloat16)
+    b = (torch.randn(E, n, k, device="cuda") * 0.02).to(torch.bfloat16)
+
+    sorted_ids, _ = torch.sort(ids.to(torch.int16), stable=True)
+    meta128 = compute_grouped_gemm_metadata(sorted_ids, E)
+    out128 = grouped_gemm_triton(a, b, E, *meta128)
+
+    pre = fused_moe_preprocess(ids, E)
+    out16 = grouped_gemm_triton(a, b, E, *pre[1:5], block_size_m=SMALL_M_BLOCK_SIZE_M)
+    # both are fp32-accumulated bf16 tensor-core dots; BLOCK_K differs so
+    # accumulation grouping may differ by a few ulp
+    torch.testing.assert_close(out16.float(), out128.float(), atol=1e-3, rtol=1e-3)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-x"]))

python/sglang/jit_kernel/tests/test_sconv_decode_metadata.py

@@ -0,0 +1,70 @@
+"""fused_decode_sconv_metadata must be bit-identical to the unfused prep.
+
+The unfused reference is the exact op sequence `_prepare_decode_sconv_metadata`
+used to launch: two arange calls + ones + precompute_helion_decode_metadata
+(!= PAD, &, clamp, long, arange x2).
+"""
+
+import pytest
+import torch
+
+from sglang.srt.models.inkling_common.kernels.sconv import (
+    PAD_SLOT_ID,
+    fused_decode_sconv_metadata,
+    precompute_helion_decode_metadata,
+)
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
+
+# cross the BLOCK=1024 grid boundary and hit odd sizes
+BATCH_SIZES = [1, 2, 3, 17, 64, 160, 257, 1023, 1024, 1025]
+
+
+def _reference(B: int, cache_indices: torch.Tensor):
+    device = cache_indices.device
+    query_start_loc = torch.arange(B + 1, dtype=torch.int32, device=device)
+    has_initial_state = torch.ones(B, dtype=torch.bool, device=device)
+    precomputed = precompute_helion_decode_metadata(
+        B=B, W=4, cache_indices=cache_indices, has_initial_state=has_initial_state
+    )
+    return query_start_loc, has_initial_state, precomputed
+
+
+@requires_cuda
+@pytest.mark.parametrize("b", BATCH_SIZES)
+@pytest.mark.parametrize("idx_dtype", [torch.int32, torch.int64])
+def test_matches_unfused(b: int, idx_dtype: torch.dtype):
+    torch.manual_seed(b)
+    cache_indices = torch.randint(0, 4096, (b,), dtype=idx_dtype, device="cuda")
+    # sprinkle PAD slots (cudagraph padding lanes)
+    pad = torch.rand(b, device="cuda") < 0.25
+    cache_indices[pad] = PAD_SLOT_ID
+
+    ref_qsl, ref_his, ref_meta = _reference(b, cache_indices)
+    qsl, his, meta = fused_decode_sconv_metadata(B=b, cache_indices=cache_indices)
+
+    for tag, got, ref in (
+        ("query_start_loc", qsl, ref_qsl),
+        ("has_initial_state", his, ref_his),
+        ("cache_mask", meta["cache_mask"], ref_meta["cache_mask"]),
+        ("safe_idx", meta["safe_idx"], ref_meta["safe_idx"]),
+        ("cu", meta["cu"], ref_meta["cu"]),
+        ("si", meta["si"], ref_meta["si"]),
+    ):
+        assert got.dtype == ref.dtype, (tag, got.dtype, ref.dtype)
+        assert got.shape == ref.shape, (tag, got.shape, ref.shape)
+        assert torch.equal(got, ref), tag
+
+
+@requires_cuda
+def test_all_pad():
+    cache_indices = torch.full((8,), PAD_SLOT_ID, dtype=torch.int32, device="cuda")
+    _, _, meta = fused_decode_sconv_metadata(B=8, cache_indices=cache_indices)
+    assert not meta["cache_mask"].any()
+    assert (meta["safe_idx"] == 0).all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-x"]))

python/sglang/jit_kernel/tests/test_sconv_extend_metadata.py

@@ -0,0 +1,174 @@
+"""fused_extend_sconv_metadata must be bit-identical to the unfused prep.
+
+The unfused reference is the exact op sequence _prepare_extend_common_metadata
++ precompute_helion_extend_metadata used to launch: zeros + cumsum + slice-copy
+(or arange + ones for verify) + the has_initial_state compare, then != PAD, &,
+clamp, long, to(int64), arange, searchsorted, clamp, to(int32).
+"""
+
+import pytest
+import torch
+
+from sglang.srt.models.inkling_common.kernels.sconv import (
+    HIS_ONES,
+    HIS_PREFIX,
+    HIS_SEQ_MINUS_EXT,
+    HIS_ZEROS,
+    PAD_SLOT_ID,
+    fused_extend_sconv_metadata,
+    precompute_helion_extend_metadata,
+)
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
+
+# cross si tiles (BLOCK_T=256) and the single-tile B bound
+BATCH_SIZES = [1, 2, 7, 64, 257, 1023]
+
+
+def _ref_extend(B, extend_seq_lens, his_mode, his_src, cache_indices, T):
+    device = cache_indices.device
+    query_start_loc = torch.zeros(B + 1, dtype=torch.int32, device=device)
+    query_start_loc[1:] = extend_seq_lens.cumsum(dim=0)
+    if his_mode == HIS_ZEROS:
+        has_initial_state = torch.zeros(B, dtype=torch.bool, device=device)
+    elif his_mode == HIS_PREFIX:
+        has_initial_state = his_src > 0
+    else:  # HIS_SEQ_MINUS_EXT
+        has_initial_state = (his_src[:B] - extend_seq_lens) > 0
+    meta = precompute_helion_extend_metadata(
+        B=B,
+        T=T,
+        W=4,
+        cache_indices=cache_indices,
+        has_initial_state=has_initial_state,
+        query_start_loc=query_start_loc,
+    )
+    return query_start_loc, has_initial_state, meta
+
+
+def _ref_verify(B, draft_token_num, cache_indices):
+    device = cache_indices.device
+    query_start_loc = torch.arange(
+        0, (B + 1) * draft_token_num, draft_token_num, dtype=torch.int32, device=device
+    )
+    has_initial_state = torch.ones(B, dtype=torch.bool, device=device)
+    meta = precompute_helion_extend_metadata(
+        B=B,
+        T=B * draft_token_num,
+        W=4,
+        cache_indices=cache_indices,
+        has_initial_state=has_initial_state,
+        query_start_loc=query_start_loc,
+    )
+    return query_start_loc, has_initial_state, meta
+
+
+def _assert_equal(got, ref):
+    for tag, g, r in (
+        ("query_start_loc", got[0], ref[0]),
+        ("has_initial_state", got[1], ref[1]),
+        ("cache_mask", got[2]["cache_mask"], ref[2]["cache_mask"]),
+        ("safe_idx", got[2]["safe_idx"], ref[2]["safe_idx"]),
+        ("cu", got[2]["cu"], ref[2]["cu"]),
+        ("si", got[2]["si"], ref[2]["si"]),
+    ):
+        assert g.dtype == r.dtype, (tag, g.dtype, r.dtype)
+        assert g.shape == r.shape, (tag, g.shape, r.shape)
+        assert torch.equal(g, r), tag
+
+
+def _cache_indices(b, idx_dtype):
+    ci = torch.randint(0, 4096, (b,), dtype=idx_dtype, device="cuda")
+    pad = torch.rand(b, device="cuda") < 0.25
+    ci[pad] = PAD_SLOT_ID
+    return ci
+
+
+@requires_cuda
+@pytest.mark.parametrize("b", BATCH_SIZES)
+@pytest.mark.parametrize("his_mode", [HIS_ZEROS, HIS_PREFIX, HIS_SEQ_MINUS_EXT])
+@pytest.mark.parametrize("lens_dtype", [torch.int32, torch.int64])
+def test_extend_matches_unfused(b, his_mode, lens_dtype):
+    torch.manual_seed(b * 10 + his_mode)
+    lens = torch.randint(0, 33, (b,), dtype=lens_dtype, device="cuda")
+    lens[torch.rand(b, device="cuda") < 0.2] = 0  # zero-length sequences
+    T = int(lens.sum().item())
+    cache_indices = _cache_indices(b, torch.int32)
+    if his_mode == HIS_PREFIX:
+        his_src = torch.randint(0, 3, (b,), dtype=lens_dtype, device="cuda")
+    elif his_mode == HIS_SEQ_MINUS_EXT:
+        his_src = lens + torch.randint(0, 2, (b,), dtype=lens_dtype, device="cuda")
+    else:
+        his_src = None
+
+    ref = _ref_extend(b, lens, his_mode, his_src, cache_indices, T)
+    got = fused_extend_sconv_metadata(
+        B=b,
+        T=T,
+        cache_indices=cache_indices,
+        his_mode=his_mode,
+        extend_seq_lens=lens,
+        his_src=his_src,
+    )
+    assert got is not None
+    _assert_equal(got, ref)
+
+
+@requires_cuda
+@pytest.mark.parametrize("b", BATCH_SIZES)
+@pytest.mark.parametrize("draft_token_num", [1, 9])
+def test_verify_matches_unfused(b, draft_token_num):
+    torch.manual_seed(b)
+    cache_indices = _cache_indices(b, torch.int64)
+    ref = _ref_verify(b, draft_token_num, cache_indices)
+    got = fused_extend_sconv_metadata(
+        B=b,
+        T=b * draft_token_num,
+        cache_indices=cache_indices,
+        his_mode=HIS_ONES,
+        draft_token_num=draft_token_num,
+    )
+    assert got is not None
+    _assert_equal(got, ref)
+
+
+@requires_cuda
+def test_cu_not_spanning_T():
+    """Dummy capture sequences: cu stops short of T; trailing si rows clamp to
+    B-1 exactly like the reference's searchsorted + clamp."""
+    b = 5
+    lens = torch.tensor([3, 0, 4, 0, 2], dtype=torch.int64, device="cuda")
+    T = int(lens.sum().item()) + 17
+    cache_indices = _cache_indices(b, torch.int32)
+    seq_lens = lens + 1
+    ref = _ref_extend(b, lens, HIS_SEQ_MINUS_EXT, seq_lens, cache_indices, T)
+    got = fused_extend_sconv_metadata(
+        B=b,
+        T=T,
+        cache_indices=cache_indices,
+        his_mode=HIS_SEQ_MINUS_EXT,
+        extend_seq_lens=lens,
+        his_src=seq_lens,
+    )
+    assert got is not None
+    _assert_equal(got, ref)
+
+
+@requires_cuda
+def test_fallback_past_batch_bound():
+    b = 1024  # > _FUSED_EXTEND_MAX_B
+    lens = torch.ones(b, dtype=torch.int64, device="cuda")
+    got = fused_extend_sconv_metadata(
+        B=b,
+        T=b,
+        cache_indices=_cache_indices(b, torch.int32),
+        his_mode=HIS_ZEROS,
+        extend_seq_lens=lens,
+    )
+    assert got is None
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-x"]))

python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh

@@ -187,42 +187,58 @@ inline void launchFusedActivationQuant(
     tensorrt_llm::QuantizationSFLayout sfLayout,
     bool disableFp4FastMath,
     cudaStream_t stream) {
-  constexpr uint32_t BLOCK_SIZE = 128;  // == innerHalf/16 for inter=2048 (one SF block per thread)
-  dim3 const grid(m), block(BLOCK_SIZE);
-
-  auto launch = [&](auto layoutTag, auto fastMathTag) {
-    fusedActivationQuantKernel<BLOCK_SIZE, decltype(layoutTag)::value, decltype(fastMathTag)::value>
-        <<<grid, block, 0, stream>>>(
-            m,
-            innerHalf,
-            innerDim,
-            gateUp,
-            loraDelta,
-            loraInputOut,
-            expandedIdxToPermutedIdx,
-            globalScaleInv,
-            weightOutput,
-            scaleOutput,
-            perTokenScaleOutput);
-  };
-  auto withFastMath = [&](auto layoutTag) {
-    if (disableFp4FastMath) {
-      launch(layoutTag, std::integral_constant<bool, true>{});
+  // One SF block per thread, no stride loop: BLOCK_SIZE must cover innerHalf/16 (a fixed
+  // 128 left cols [2048,inter) unwritten at Inkling EP8's inter=3072 -> NaN from the down GEMM).
+  uint32_t const numVecs = static_cast<uint32_t>(innerHalf) / 16;
+  auto dispatchBlock = [&](auto blockTag) {
+    constexpr uint32_t BLOCK_SIZE = decltype(blockTag)::value;
+    dim3 const grid(m), block(BLOCK_SIZE);
+    auto launch = [&](auto layoutTag, auto fastMathTag) {
+      fusedActivationQuantKernel<BLOCK_SIZE, decltype(layoutTag)::value, decltype(fastMathTag)::value>
+          <<<grid, block, 0, stream>>>(
+              m,
+              innerHalf,
+              innerDim,
+              gateUp,
+              loraDelta,
+              loraInputOut,
+              expandedIdxToPermutedIdx,
+              globalScaleInv,
+              weightOutput,
+              scaleOutput,
+              perTokenScaleOutput);
+    };
+    auto withFastMath = [&](auto layoutTag) {
+      if (disableFp4FastMath) {
+        launch(layoutTag, std::integral_constant<bool, true>{});
+      } else {
+        launch(layoutTag, std::integral_constant<bool, false>{});
+      }
+    };
+    if (sfLayout == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
+      withFastMath(
+          std::integral_constant<
+              tensorrt_llm::QuantizationSFLayout,
+              tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4>{});
+    } else if (sfLayout == tensorrt_llm::QuantizationSFLayout::LINEAR) {
+      withFastMath(
+          std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::LINEAR>{});
     } else {
-      launch(layoutTag, std::integral_constant<bool, false>{});
+      withFastMath(
+          std::integral_constant<
+              tensorrt_llm::QuantizationSFLayout,
+              tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4>{});
     }
   };
-  if (sfLayout == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
-    withFastMath(
-        std::integral_constant<
-            tensorrt_llm::QuantizationSFLayout,
-            tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4>{});
-  } else if (sfLayout == tensorrt_llm::QuantizationSFLayout::LINEAR) {
-    withFastMath(
-        std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::LINEAR>{});
+  if (numVecs <= 128) {
+    dispatchBlock(std::integral_constant<uint32_t, 128>{});
+  } else if (numVecs <= 256) {
+    dispatchBlock(std::integral_constant<uint32_t, 256>{});
+  } else if (numVecs <= 512) {
+    dispatchBlock(std::integral_constant<uint32_t, 512>{});
   } else {
-    withFastMath(
-        std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4>{});
+    // Callers guard on numVecs <= 512 and fall back to the unfused chain.
+    dispatchBlock(std::integral_constant<uint32_t, 1024>{});
   }
 }
 

python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu

@@ -3116,9 +3116,9 @@ class FP4BlockScaleLoraLauncher {
     static int const fuseActQuant = envFlag("SGLANG_OPT_FUSED_MOE_ACTIVATION_QUANT_FUSE") ? 1 : 0;
     static int const actOptMode = envFlag("SGLANG_OPT_FUSED_MOE_ACTIVATION_VEC") ? 1 : 0;
 
-    if (fuseActQuant) {
+    if (fuseActQuant && inter / 16 <= 512) {
       // Fused: gate_up (interleaved) + lora_delta -> act_fp4/sf/per_token + activation_lora_input,
-      // without materializing activated_bf16. inter must be a multiple of 16 (always true here).
+      // without materializing activated_bf16. >512 SF vecs/row falls to the unfused chain below.
       flashinfer::sgl_fused_act_quant::launchFusedActivationQuant(
           num_tokens * top_k,
           inter,

python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_runner.cu

@@ -533,6 +533,9 @@ void Runner::run(
       ptrCtaIdxXyToBatchIdx,
       ptrCtaIdxXyToMnLimit,
       ptrNumNonExitingCtas,
+#if SGLANG_FLASHINFER_HAS_PERMUTED_BIAS_ROW_IDX
+      /* permutedIdxToBiasRowIdx */ nullptr,
+#endif
       bmm1Workspace,
       stream,
       device,
@@ -712,6 +715,9 @@ void Runner::run(
       ptrCtaIdxXyToBatchIdx,
       ptrCtaIdxXyToMnLimit,
       ptrNumNonExitingCtas,
+#if SGLANG_FLASHINFER_HAS_PERMUTED_BIAS_ROW_IDX
+      /* permutedIdxToBiasRowIdx */ nullptr,
+#endif
       bmm2Workspace,
       stream,
       device,

python/sglang/jit_kernel/trtllm_lora_temp/jit.py

@@ -24,6 +24,12 @@ def gen_sgl_trtllm_gen_fused_moe_sm100_module():
     flashinfer_data_dir = Path(flashinfer.__file__).resolve().parent / "data"
     flashinfer_csrc_dir = flashinfer_data_dir / "csrc"
     flashinfer_include_dir = flashinfer_data_dir / "include"
+    kernel_runner_header = (
+        flashinfer_include_dir / "flashinfer/trtllm/batched_gemm/KernelRunner.h"
+    )
+    has_permuted_bias_row_idx = (
+        "permutedIdxToBiasRowIdx" in kernel_runner_header.read_text()
+    )
 
     include_path = f"{ArtifactPath.TRTLLM_GEN_BMM}/include"
     header_name = "flashinferMetaInfo"
@@ -85,6 +91,7 @@ def gen_sgl_trtllm_gen_fused_moe_sm100_module():
             "-DENABLE_FP8",
             "-DENABLE_FP4",
             "-DCUTLASS_ENABLE_GDC_FOR_SM100=1",
+            f"-DSGLANG_FLASHINFER_HAS_PERMUTED_BIAS_ROW_IDX={int(has_permuted_bias_row_idx)}",
             f'-DTLLM_GEN_GEMM_CUBIN_PATH=\\"{ArtifactPath.TRTLLM_GEN_BMM}\\"',
         ]
         + nvcc_flags,

python/sglang/jit_kernel/trtllm_lora_temp/moe_lora_merged_align.py

@@ -23,15 +23,6 @@ def _jit_module(dtype: torch.dtype) -> Module:
     )
 
 
-def supports_merged_align(virtual_num_experts: int) -> bool:
-    """Commit-1 kernel only implements the (64, 1024] bucket-count branch.
-
-    The bucket count is virtual_num_experts + 1 (the +1 sentinel bucket). Other
-    regimes (small-batch <=64, v2 >1024) keep the old path."""
-    num_buckets = virtual_num_experts + 1
-    return 64 < num_buckets <= 1024
-
-
 def moe_lora_merged_align(
     topk_ids: torch.Tensor,
     token_lora_mapping: torch.Tensor,

python/sglang/jit_kernel/utils/__init__.py

@@ -7,6 +7,7 @@ from sglang.jit_kernel.utils.arch import (
 )
 from sglang.jit_kernel.utils.common import (
     cache_once,
+    empty_sentinel,
     get_ci_test_range,
     is_hip_runtime,
     is_musa_runtime,
@@ -16,6 +17,7 @@ from sglang.jit_kernel.utils.common import (
 from sglang.jit_kernel.utils.compile import KERNEL_PATH, load_jit, make_cpp_args
 
 __all__ = [
+    "empty_sentinel",
     "should_run_full_tests",
     "get_ci_test_range",
     "cache_once",

python/sglang/jit_kernel/utils/common.py

@@ -41,6 +41,15 @@ def cache_once(fn: F) -> F:
     return wrapper  # type: ignore
 
 
+@functools.lru_cache(maxsize=None)
+def empty_sentinel(device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    """Cached 0-element tensor for optional-tensor FFI slots (the numel-0
+    "not present" convention). Allocating a fresh empty per call costs
+    ~1.2us CPU on eager paths; the sentinel is never read, so one cached
+    instance per (device, dtype) is safe to share."""
+    return torch.empty(0, dtype=dtype, device=device)
+
+
 @cache_once
 def is_hip_runtime() -> bool:
     return bool(torch.version.hip)

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

@@ -25,6 +25,7 @@ import logging
 import triton
 import triton.language as tl
 
+from sglang.kernels.ops.attention.score_mod import unpack_aux_tensors
 from sglang.srt.utils import is_hip
 
 _is_hip = is_hip()
@@ -129,6 +130,11 @@ def _fwd_kernel_stage1(
     Lv: tl.constexpr,
     xai_temperature_len: tl.constexpr,
     PAGE_SIZE: tl.constexpr,
+    SCORE_MOD: tl.constexpr = None,
+    Aux0=None,
+    aux0_stride_t=0,
+    aux0_stride_h=0,
+    aux0_len=0,
 ):
     # int64 to avoid overflow of flat offsets into Mid_O when
     # batch * num_head * max_kv_splits * head_dim exceeds 2**31.
@@ -206,6 +212,20 @@ def _fwd_kernel_stage1(
             if xai_temperature_len > 0:
                 qk *= xai_temperature_reg
 
+            if SCORE_MOD is not None:
+                qk = SCORE_MOD(
+                    qk,
+                    cur_batch_seq_len - 1,
+                    offs_n,
+                    cur_batch,
+                    cur_head,
+                    offs_n < split_kv_end,
+                    Aux0,
+                    aux0_stride_t,
+                    aux0_stride_h,
+                    aux0_len,
+                )
+
             qk = tl.where(offs_n < split_kv_end, qk, float("-inf"))
 
             if PAGE_SIZE == 1:
@@ -275,6 +295,8 @@ def _decode_att_m_fwd(
     logit_cap,
     xai_temperature_len=-1,
     page_size: int = 1,
+    score_mod=None,
+    aux_tensors=None,
 ):
     BLOCK = 64
     # [TODO] work around SGPR limit on MI3xx
@@ -311,6 +333,10 @@ def _decode_att_m_fwd(
         v_buffer, page_size
     )
 
+    aux0, aux0_stride_t, aux0_stride_h, aux0_len = unpack_aux_tensors(
+        score_mod, aux_tensors
+    )
+
     _fwd_kernel_stage1[grid](
         q,
         k_buffer,
@@ -346,6 +372,11 @@ def _decode_att_m_fwd(
         Lk=Lk,
         Lv=Lv,
         PAGE_SIZE=page_size,
+        SCORE_MOD=score_mod,
+        Aux0=aux0,
+        aux0_stride_t=aux0_stride_t,
+        aux0_stride_h=aux0_stride_h,
+        aux0_len=aux0_len,
     )
 
 
@@ -389,6 +420,11 @@ def _fwd_grouped_kernel_stage1(
     HAS_MLA: tl.constexpr = False,
     USE_PDL: tl.constexpr = False,
     PAGE_SIZE: tl.constexpr = 1,
+    SCORE_MOD: tl.constexpr = None,
+    Aux0=None,
+    aux0_stride_t=0,
+    aux0_stride_h=0,
+    aux0_len=0,
 ):
     # int64 to avoid overflow of flat offsets into Mid_O when
     # batch * num_head * max_kv_splits * head_dim exceeds 2**31.
@@ -500,6 +536,20 @@ def _fwd_grouped_kernel_stage1(
             if xai_temperature_len > 0:
                 qk *= xai_temperature_reg[:, None]
 
+            if SCORE_MOD is not None:
+                qk = SCORE_MOD(
+                    qk,
+                    cur_batch_seq_len - 1,
+                    offs_n[None, :],
+                    cur_batch,
+                    cur_head[:, None],
+                    mask_h[:, None] & (offs_n[None, :] < split_kv_end),
+                    Aux0,
+                    aux0_stride_t,
+                    aux0_stride_h,
+                    aux0_len,
+                )
+
             qk = tl.where(
                 mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf")
             )
@@ -574,6 +624,8 @@ def _decode_grouped_att_m_fwd(
     has_mla=False,
     use_pdl=False,
     page_size: int = 1,
+    score_mod=None,
+    aux_tensors=None,
 ):
     BLOCK = 32
     Lk = k_buffer.shape[-1]
@@ -623,6 +675,10 @@ def _decode_grouped_att_m_fwd(
         v_buffer, page_size
     )
 
+    aux0, aux0_stride_t, aux0_stride_h, aux0_len = unpack_aux_tensors(
+        score_mod, aux_tensors
+    )
+
     _fwd_grouped_kernel_stage1[grid](
         q,
         k_buffer,
@@ -663,6 +719,11 @@ def _decode_grouped_att_m_fwd(
         HAS_MLA=has_mla,
         USE_PDL=use_pdl,
         PAGE_SIZE=page_size,
+        SCORE_MOD=score_mod,
+        Aux0=aux0,
+        aux0_stride_t=aux0_stride_t,
+        aux0_stride_h=aux0_stride_h,
+        aux0_len=aux0_len,
         **extra_kargs,
     )
 
@@ -814,6 +875,8 @@ def decode_attention_fwd_normal(
     sinks=None,
     xai_temperature_len=-1,
     page_size: int = 1,
+    score_mod=None,
+    aux_tensors=None,
 ):
     _decode_att_m_fwd(
         q,
@@ -829,6 +892,8 @@ def decode_attention_fwd_normal(
         logit_cap,
         xai_temperature_len,
         page_size=page_size,
+        score_mod=score_mod,
+        aux_tensors=aux_tensors,
     )
     _decode_softmax_reducev_fwd(
         attn_logits,
@@ -863,6 +928,8 @@ def decode_attention_fwd_grouped(
     has_mla=False,
     use_pdl=False,
     page_size: int = 1,
+    score_mod=None,
+    aux_tensors=None,
 ):
     _decode_grouped_att_m_fwd(
         q,
@@ -880,6 +947,8 @@ def decode_attention_fwd_grouped(
         has_mla=has_mla,
         use_pdl=use_pdl,
         page_size=page_size,
+        score_mod=score_mod,
+        aux_tensors=aux_tensors,
     )
     _decode_softmax_reducev_fwd(
         attn_logits,
@@ -916,6 +985,8 @@ def decode_attention_fwd(
     has_mla=False,
     use_pdl=False,
     page_size: int = 1,
+    score_mod=None,
+    aux_tensors=None,
 ):
     assert max_kv_splits == attn_logits.shape[2]
     assert q.shape[0] <= kv_indptr.shape[0] - 1
@@ -944,6 +1015,8 @@ def decode_attention_fwd(
             sinks=sinks,
             xai_temperature_len=xai_temperature_len,
             page_size=page_size,
+            score_mod=score_mod,
+            aux_tensors=aux_tensors,
         )
     else:
         # GQA/MQA/MLA
@@ -966,4 +1039,6 @@ def decode_attention_fwd(
             has_mla=has_mla,
             use_pdl=use_pdl,
             page_size=page_size,
+            score_mod=score_mod,
+            aux_tensors=aux_tensors,
         )

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

@@ -24,6 +24,7 @@ from sglang.kernels.ops.attention.decode_attention import _extract_kv_strides
 from sglang.kernels.ops.attention.prefill_attention import (
     context_attention_fwd,
 )
+from sglang.kernels.ops.attention.score_mod import unpack_aux_tensors
 from sglang.srt.utils import is_cuda, is_gfx95_supported, is_hip
 
 _is_cuda = is_cuda()
@@ -295,6 +296,11 @@ def _fwd_kernel(
     STORE_TRANSPOSE: tl.constexpr,
     HAS_SINK: tl.constexpr,
     PAGE_SIZE: tl.constexpr = 1,
+    SCORE_MOD: tl.constexpr = None,
+    Aux0=None,
+    aux0_stride_t=0,
+    aux0_stride_h=0,
+    aux0_len=0,
 ):
     cur_seq = tl.program_id(0)
     cur_head = tl.program_id(1)
@@ -450,6 +456,22 @@ def _fwd_kernel(
             if xai_temperature_len > 0:
                 qk *= xai_temperature_reg[:, None]
 
+            if SCORE_MOD is not None:
+                qk = SCORE_MOD(
+                    qk,
+                    (cur_seq_len_prefix + cur_block_m * BLOCK_M + offs_m)[:, None],
+                    start_n + offs_n[None, :],
+                    (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m)[
+                        :, None
+                    ],
+                    cur_head,
+                    final_mask,
+                    Aux0,
+                    aux0_stride_t,
+                    aux0_stride_h,
+                    aux0_len,
+                )
+
             qk = tl.where(final_mask, qk, float("-inf"))
 
             row_max = tl.max(qk, 1)
@@ -565,6 +587,22 @@ def _fwd_kernel(
             if xai_temperature_len > 0:
                 qk *= xai_temperature_reg[:, None]
 
+            if SCORE_MOD is not None:
+                qk = SCORE_MOD(
+                    qk,
+                    (cur_seq_len_prefix + cur_block_m * BLOCK_M + offs_m)[:, None],
+                    cur_seq_len_prefix + start_n + offs_n[None, :],
+                    (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m)[
+                        :, None
+                    ],
+                    cur_head,
+                    final_mask,
+                    Aux0,
+                    aux0_stride_t,
+                    aux0_stride_h,
+                    aux0_len,
+                )
+
             qk = tl.where(final_mask, qk, float("-inf"))
 
             row_max = tl.max(qk, 1)
@@ -646,6 +684,8 @@ def extend_attention_fwd(
     skip_prefix=False,
     skip_extend=False,
     page_size: int = 1,
+    score_mod=None,
+    aux_tensors=None,
 ):
     """
     q_extend, k_extend, v_extend, o_extend: contiguous tensors
@@ -656,6 +696,8 @@ def extend_attention_fwd(
     written to it (used by DCP to merge partial attention across ranks).
     ``skip_prefix`` / ``skip_extend`` skip the prefix-KV / current-chunk stage
     respectively so DCP can compute those two parts separately.
+    ``score_mod`` / ``aux_tensors`` add a custom term to the attention logits;
+    see triton_ops/score_mod.py for the contract.
     """
     Lq, Lk, Lv = (
         q_extend.shape[-1],
@@ -695,6 +737,10 @@ def extend_attention_fwd(
         v_buffer, page_size
     )
 
+    aux0, aux0_stride_t, aux0_stride_h, aux0_len = unpack_aux_tensors(
+        score_mod, aux_tensors
+    )
+
     _fwd_kernel[grid](
         q_extend,
         k_extend,
@@ -751,6 +797,11 @@ def extend_attention_fwd(
         HAS_SINK=HAS_SINK,
         STORE_TRANSPOSE=_is_hip,
         PAGE_SIZE=page_size,
+        SCORE_MOD=score_mod,
+        Aux0=aux0,
+        aux0_stride_t=aux0_stride_t,
+        aux0_stride_h=aux0_stride_h,
+        aux0_len=aux0_len,
         num_warps=num_warps,
         num_stages=num_stages,
         **extra_kargs,
@@ -838,6 +889,11 @@ def _fwd_kernel_unified(
     USE_CUSTOM_MASK: tl.constexpr,
     HAS_SINK: tl.constexpr,
     PAGE_SIZE: tl.constexpr = 1,
+    SCORE_MOD: tl.constexpr = None,
+    Aux0=None,
+    aux0_stride_t=0,
+    aux0_stride_h=0,
+    aux0_len=0,
 ):
     """
     Unified 1-stage kernel for deterministic extend attention.
@@ -1026,6 +1082,20 @@ def _fwd_kernel_unified(
             if xai_temperature_len > 0:
                 qk *= xai_temperature_reg[:, None]
 
+            if SCORE_MOD is not None:
+                qk = SCORE_MOD(
+                    qk,
+                    (cur_seq_prefix_len + cur_block_m * BLOCK_M + offs_m)[:, None],
+                    start_n + offs_n[None, :],
+                    (cur_seq_q_start_idx + cur_block_m * BLOCK_M + offs_m)[:, None],
+                    cur_head,
+                    final_mask,
+                    Aux0,
+                    aux0_stride_t,
+                    aux0_stride_h,
+                    aux0_len,
+                )
+
             qk = tl.where(final_mask, qk, float("-inf"))
 
             # Online softmax
@@ -1101,6 +1171,8 @@ def extend_attention_fwd_unified(
     window_start_pos=None,
     xai_temperature_len=-1,
     page_size: int = 1,
+    score_mod=None,
+    aux_tensors=None,
 ):
     """
     Unified 1-stage extend attention for deterministic inference.
@@ -1162,6 +1234,10 @@ def extend_attention_fwd_unified(
         v_buffer, page_size
     )
 
+    aux0, aux0_stride_t, aux0_stride_h, aux0_len = unpack_aux_tensors(
+        score_mod, aux_tensors
+    )
+
     _fwd_kernel_unified[grid](
         q,
         o,
@@ -1204,6 +1280,11 @@ def extend_attention_fwd_unified(
         USE_CUSTOM_MASK=USE_CUSTOM_MASK,
         HAS_SINK=HAS_SINK,
         PAGE_SIZE=page_size,
+        SCORE_MOD=score_mod,
+        Aux0=aux0,
+        aux0_stride_t=aux0_stride_t,
+        aux0_stride_h=aux0_stride_h,
+        aux0_len=aux0_len,
         num_warps=num_warps,
         num_stages=num_stages,
         **extra_kargs,

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

@@ -0,0 +1,70 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _apply_log_scaling_tau_kernel(
+    x_ptr,
+    tau_ptr,  # [rows] fp32 (flattened per-row scale)
+    out_ptr,  # [rows, inner] contiguous, same dtype as x
+    x_row_stride,
+    inner,
+    total,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid.to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < total
+    row = offs // inner
+    col = offs % inner
+    x = tl.load(x_ptr + row * x_row_stride + col, mask=mask).to(tl.float32)
+    tau = tl.load(tau_ptr + row, mask=mask)
+    y = x * tau
+    tl.store(out_ptr + offs, y.to(out_ptr.dtype.element_ty), mask=mask)
+
+
+def apply_log_scaling_tau(x: torch.Tensor, tau: torch.Tensor) -> torch.Tensor:
+    """out = (x.float() * tau).to(x.dtype) with tau broadcast per leading row,
+    fused into one launch. x may carry a leading-dim stride (the q slice of the
+    fused qkvr output); its trailing dims must be contiguous. No dynamo: the
+    torch.compile'd predecessor's call sites spanned enough rank /
+    dispatch-key / 0-1 specialization variants (target + de-tied MTP heads) to
+    exceed the recompile limit, which crashed (fullgraph) or wedged capture
+    (raised limit)."""
+    rows = x.shape[0]
+    inner = x.numel() // rows if rows else 0
+    inner_contiguous = x.stride(-1) == 1 and (
+        x.dim() == 2 or x.stride(-2) == x.shape[-1] * x.stride(-1)
+    )
+    if rows == 0 or inner == 0 or not inner_contiguous:
+        return (x.float() * tau).to(x.dtype)
+
+    if (
+        x.is_cuda
+        and x.dtype == torch.bfloat16
+        and inner % 8 == 0
+        and x.data_ptr() % 16 == 0
+        and (x.stride(0) * 2) % 16 == 0
+    ):
+        # Vectorized JIT kernel (16B loads, one row divide per vector) --
+        # bit-identical output (same fp32-mul + bf16-round), ~2-3x the
+        # scalar triton kernel below at every size.
+        from sglang.jit_kernel.inkling_row_scale import row_scale_bf16
+
+        x2d = torch.as_strided(x, (rows, inner), (x.stride(0), 1))
+        return row_scale_bf16(x2d, tau.reshape(rows).float()).view(x.shape)
+
+    out = torch.empty(x.shape, dtype=x.dtype, device=x.device)
+    total = rows * inner
+    BLOCK = 1024
+    _apply_log_scaling_tau_kernel[(triton.cdiv(total, BLOCK),)](
+        x,
+        tau.reshape(rows).to(torch.float32),
+        out,
+        x.stride(0),
+        inner,
+        total,
+        BLOCK=BLOCK,
+    )
+    return out

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

@@ -322,6 +322,216 @@ def _fused_metadata_kernel_ps1_no_swa(
     tl.store(page_table + pt_offsets, page_index, mask=mask, cache_modifier=".cg")
 
 
+@triton.jit
+def _draft_extend_metadata_kernel(
+    # Input tensors
+    seq_lens,
+    seq_lens_stride_0,
+    extend_seq_lens,
+    extend_seq_lens_stride_0,
+    req_to_token,
+    req_to_token_stride_0,
+    req_to_token_stride_1,
+    req_pool_indices,
+    req_pool_indices_stride_0,
+    # Output buffers
+    cache_seqlens_int32,
+    cache_seqlens_int32_stride_0,
+    cu_seqlens_k,
+    cu_seqlens_k_stride_0,
+    cu_seqlens_q,
+    cu_seqlens_q_stride_0,
+    page_table,
+    page_table_stride_0,
+    page_table_stride_1,
+    full_to_swa_index_mapping,
+    swa_page_table,
+    out_cache_loc,
+    swa_out_cache_loc,
+    # Scalar parameters
+    B,
+    max_seq_pages,
+    tokens_per_req,
+    PAGE_SIZE_ONE: tl.constexpr,
+    SHIFT: tl.constexpr,
+    BLOCK_COLS: tl.constexpr,
+    HAS_SWA: tl.constexpr,
+    OUT_BLOCK: tl.constexpr,
+):
+    pid_b = tl.program_id(0)  # batch index
+    pid_c = tl.program_id(1)  # column chunk index
+
+    # 1. Prefix sums (only one block does them): cache_seqlens + cu_seqlens_k
+    #    from seq_lens, cu_seqlens_q from extend_seq_lens.
+    if pid_b == 0 and pid_c == 0:
+        acc_k = 0
+        acc_q = 0
+        for idx in range(B):
+            seq = tl.load(seq_lens + idx * seq_lens_stride_0).to(tl.int32)
+            tl.store(cache_seqlens_int32 + idx * cache_seqlens_int32_stride_0, seq)
+            tl.store(cu_seqlens_k + idx * cu_seqlens_k_stride_0, acc_k)
+            acc_k += seq
+            ext = tl.load(extend_seq_lens + idx * extend_seq_lens_stride_0).to(tl.int32)
+            tl.store(cu_seqlens_q + idx * cu_seqlens_q_stride_0, acc_q)
+            acc_q += ext
+        tl.store(cu_seqlens_k + B * cu_seqlens_k_stride_0, acc_k)
+        tl.store(cu_seqlens_q + B * cu_seqlens_q_stride_0, acc_q)
+
+    # 2. SWA write-loc translation for this request's extend tokens. Runs
+    #    before the seq_len early-return so padded rows (seq_len 0) keep
+    #    swa_out_cache_loc consistent with out_cache_loc.
+    if HAS_SWA:
+        if pid_c == 0:
+            tok_idx = tl.arange(0, OUT_BLOCK)
+            tok_mask = tok_idx < tokens_per_req
+            tok_offsets = pid_b * tokens_per_req + tok_idx
+            full_locs = tl.load(out_cache_loc + tok_offsets, mask=tok_mask, other=0)
+            swa_locs = tl.load(
+                full_to_swa_index_mapping + full_locs, mask=tok_mask, other=0
+            )
+            tl.store(swa_out_cache_loc + tok_offsets, swa_locs, mask=tok_mask)
+
+    # 3. Page-table gather for this batch row and column chunk, self-guarded
+    #    on the device-side seq_len (no host max; tails keep stale values the
+    #    attention kernels never read past cache_seqlens).
+    if max_seq_pages == 0:
+        return
+
+    seq_len = tl.load(seq_lens + pid_b * seq_lens_stride_0).to(tl.int32)
+    if PAGE_SIZE_ONE:
+        num_live_pages = seq_len
+    else:
+        num_live_pages = (seq_len + (1 << SHIFT) - 1) >> SHIFT
+    num_live_pages = tl.minimum(num_live_pages, max_seq_pages)
+    if pid_c * BLOCK_COLS >= num_live_pages:
+        return
+
+    row_idx = tl.load(req_pool_indices + pid_b * req_pool_indices_stride_0)
+    row_offset = row_idx * req_to_token_stride_0
+
+    col_offsets = pid_c * BLOCK_COLS + tl.arange(0, BLOCK_COLS)
+    mask = col_offsets < num_live_pages
+
+    if PAGE_SIZE_ONE:
+        col_idx = col_offsets
+    else:
+        col_idx = col_offsets << SHIFT
+
+    rt_offsets = row_offset + col_idx * req_to_token_stride_1
+    page_index = tl.load(
+        req_to_token + rt_offsets, mask=mask, other=0, cache_modifier=".cg"
+    )
+
+    if PAGE_SIZE_ONE:
+        page_table_val = page_index
+    else:
+        page_table_val = page_index >> SHIFT
+
+    pt_offsets = pid_b * page_table_stride_0 + col_offsets * page_table_stride_1
+    tl.store(page_table + pt_offsets, page_table_val, mask=mask, cache_modifier=".cg")
+
+    if HAS_SWA:
+        swa_loc = tl.load(full_to_swa_index_mapping + page_index, mask=mask, other=0)
+        if PAGE_SIZE_ONE:
+            swa_page_table_val = swa_loc
+        else:
+            swa_page_table_val = swa_loc >> SHIFT
+        tl.store(
+            swa_page_table + pt_offsets,
+            swa_page_table_val.to(tl.int32),
+            mask=mask,
+            cache_modifier=".cg",
+        )
+
+
+def draft_extend_set_metadata(
+    cache_seqlens_int32: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    page_table: torch.Tensor,
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    seq_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    page_size: int,
+    full_to_swa_index_mapping: Optional[torch.Tensor] = None,
+    swa_page_table: Optional[torch.Tensor] = None,
+    out_cache_loc: Optional[torch.Tensor] = None,
+    swa_out_cache_loc: Optional[torch.Tensor] = None,
+):
+    """Fused, graph-recordable DRAFT_EXTEND_V2 metadata update (one launch):
+      1. cache_seqlens = seq_lens (int32 cast)
+      2. cu_seqlens_k = pad(cumsum(cache_seqlens))
+      3. cu_seqlens_q = pad(cumsum(extend_seq_lens))
+      4. page_table[:, :pages(seq_len)] = req_to_token[pool_idx, ::page_size] // page_size
+      5. (SWA pools) swa_page_table likewise via the full->swa lookup, and
+         swa_out_cache_loc = full_to_swa_index_mapping[out_cache_loc]
+
+    The page gathers self-guard on the device-side seq_lens (no host max);
+    row tails keep stale values that attention kernels never read past
+    cache_seqlens, matching the eager replay path's bounded writes.
+    """
+    assert (
+        page_size > 0 and (page_size & (page_size - 1)) == 0
+    ), f"page_size must be a power of two, got {page_size}"
+
+    batch_size = cache_seqlens_int32.shape[0]
+    max_seq_pages = page_table.shape[1]
+
+    has_swa = full_to_swa_index_mapping is not None
+    if has_swa:
+        assert swa_page_table is not None
+        assert swa_page_table.shape == page_table.shape
+        assert swa_page_table.stride() == page_table.stride()
+        assert out_cache_loc is not None and swa_out_cache_loc is not None
+        num_out_tokens = out_cache_loc.shape[0]
+        assert swa_out_cache_loc.shape[0] == num_out_tokens
+        assert num_out_tokens > 0 and num_out_tokens % batch_size == 0
+        tokens_per_req = num_out_tokens // batch_size
+        out_block = triton.next_power_of_2(tokens_per_req)
+    else:
+        tokens_per_req = 0
+        out_block = 1
+
+    BLOCK_COLS = 256
+    grid = (batch_size, max(1, triton.cdiv(max_seq_pages, BLOCK_COLS)))
+
+    _draft_extend_metadata_kernel[grid](
+        seq_lens,
+        seq_lens.stride(0),
+        extend_seq_lens,
+        extend_seq_lens.stride(0),
+        req_to_token,
+        req_to_token.stride(0),
+        req_to_token.stride(1),
+        req_pool_indices,
+        req_pool_indices.stride(0),
+        cache_seqlens_int32,
+        cache_seqlens_int32.stride(0),
+        cu_seqlens_k,
+        cu_seqlens_k.stride(0),
+        cu_seqlens_q,
+        cu_seqlens_q.stride(0),
+        page_table,
+        page_table.stride(0),
+        page_table.stride(1),
+        full_to_swa_index_mapping,
+        swa_page_table,
+        out_cache_loc,
+        swa_out_cache_loc,
+        batch_size,
+        max_seq_pages,
+        tokens_per_req,
+        PAGE_SIZE_ONE=page_size == 1,
+        SHIFT=(page_size).bit_length() - 1 if page_size > 1 else 0,
+        BLOCK_COLS=BLOCK_COLS,
+        num_warps=8,
+        num_stages=3,
+        HAS_SWA=has_swa,
+        OUT_BLOCK=out_block,
+    )
+
+
 def normal_decode_set_metadata(
     cache_seqlens_int32: torch.Tensor,
     cu_seqlens_k: torch.Tensor,

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

@@ -0,0 +1,56 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Generic score_mod for the Triton attention kernels, mirroring FA4's
+``score_mod``/``aux_tensors``. A Triton score_mod is a ``@triton.jit`` function
+inlined into the kernels as a constexpr argument:
+
+    score_mod(qk, q_pos, kv_pos, q_idx, head, mask,
+              Aux0, aux0_stride_t, aux0_stride_h, aux0_len) -> qk
+
+The kernels pre-broadcast q_pos/kv_pos/q_idx/head to ``qk``'s shape, so an
+elementwise score_mod works at every call site. ``aux_tensors`` supports one
+3D tensor ``[num_q_tokens, num_q_heads, D]`` with a contiguous last dim.
+"""
+
+import triton
+import triton.language as tl
+
+
+def unpack_aux_tensors(score_mod, aux_tensors):
+    if score_mod is None:
+        return None, 0, 0, 0
+    assert (
+        aux_tensors is not None and len(aux_tensors) == 1
+    ), "Triton score_mod currently requires exactly one aux tensor"
+    aux0 = aux_tensors[0]
+    assert aux0.dim() == 3 and aux0.stride(2) == 1, (
+        f"aux_tensors[0] must be 3D with a contiguous last dim, "
+        f"got shape={tuple(aux0.shape)} stride={aux0.stride()}"
+    )
+    return aux0, aux0.stride(0), aux0.stride(1), aux0.shape[2]
+
+
+@triton.jit
+def relative_bias_score_mod(
+    qk, q_pos, kv_pos, q_idx, head, mask, Aux0, aux0_stride_t, aux0_stride_h, aux0_len
+):
+    """Add ``Aux0[q_idx, head, q_pos - kv_pos]`` when 0 <= q_pos - kv_pos < aux0_len."""
+    rel_dist = q_pos - kv_pos
+    rel_idx = tl.minimum(tl.maximum(rel_dist, 0), aux0_len - 1)
+    bias = tl.load(
+        Aux0 + q_idx * aux0_stride_t + head * aux0_stride_h + rel_idx,
+        mask=mask & (rel_dist >= 0) & (rel_dist < aux0_len),
+        other=0.0,
+    )
+    return qk + bias

python/sglang/kernels/ops/gemm/trtllm_lora_temp/gate_up_lora_b.py

@@ -147,7 +147,7 @@ def _gate_up_lora_b_kernel(
 
     # Store result to output matrix
     partial_sum *= scaling
-    partial_sum = partial_sum.to(x.dtype.element_ty)
+    partial_sum = partial_sum.to(output.dtype.element_ty)
     output_ptr = (
         output
         + n_start * output_stride_1
@@ -214,6 +214,7 @@ def gate_up_lora_b_fwd(
         )
         and s * r >= _CUBLAS_MIN_S_RANK
         and gate_up_lora_b.shape[0] == 1
+        and x.dtype == gate_up_lora_b.dtype
     ):  # single-adapter fast path: only valid with one resident slot
         return _gate_up_lora_b_cublas(
             x, gate_up_lora_b, batch_info, output_dim, base_output

python/sglang/kernels/ops/gemm/trtllm_lora_temp/qkv_lora_b.py

@@ -236,18 +236,17 @@ def qkv_lora_b_fwd(
         and batch_info.max_len >= _CUBLAS_MIN_MAX_LEN
         and qkv_lora_b.shape[0]
         == 1  # single-adapter fast path: only valid with one resident slot
+        and x.dtype == qkv_lora_b.dtype
     ):
         return _qkv_lora_b_cublas(
             x, qkv_lora_b, batch_info, output_offset_cpu, base_output, n_slices
         )
 
     BLOCK_S = 16
-    BLOCK_R = triton.next_power_of_2(r)
-    # BLOCK_OUT stays 64: with the 1-adapter cuBLAS dispatch the Triton path
-    # only runs for decode-sized batches, where 128 halves the grid (96->48
-    # programs on Kimi r16 bs64) and slows the kernel ~60% (11.4->18.5us, B200).
-    # Re-swept for the store path on GB200: 32 vs 64 is within noise (one preset
-    # marginally each way), so the single value is kept for both writebacks.
+    # Pad to >=16 for Triton MMA K>=16 (rank<16 adapters); k_offset < K=r masks the
+    # padded contraction rows to 0, so the result is unchanged.
+    BLOCK_R = max(16, triton.next_power_of_2(r))
+    # Keep one output tile size for both writeback paths.
     BLOCK_OUT = 64
 
     grid_b = (

python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_a.py

@@ -43,6 +43,7 @@ def _sgemm_lora_a_kernel(
     BLOCK_K: tl.constexpr,
     SPLIT_K: tl.constexpr = 1,
     ENABLE_PDL: tl.constexpr = False,
+    PADDED_RANK: tl.constexpr = False,
 ):
     """
     Computes a segmented batched matrix multiplication for the LoRA A matrix.
@@ -81,7 +82,8 @@ def _sgemm_lora_a_kernel(
         return
 
     # Adjust N (stack_num * max_rank) to this adapter's actual rank.
-    N = tl.minimum(N, rank * stack_num)
+    if not PADDED_RANK:
+        N = tl.minimum(N, rank * stack_num)
 
     # The tile in output matrix will have (pid_s, pid_n) as id
     num_pid_n = tl.cdiv(N, BLOCK_N)
@@ -212,11 +214,12 @@ def sgemm_lora_a_fwd(
 
     launch_kwargs = {}
     if split_k > 1:
-        # out_alloc_stream (SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC) is intentionally NOT honored here:
-        # torch.zeros launches its memset on the alloc stream, which would race the side-stream
-        # shrink without extra ordering. No current config exercises split-K together with the
-        # two-stream main-alloc overlap (qwen3.5 leaves split-K off; kimi is single-stream-coherent).
-        output = torch.zeros((S, R), device=x.device, dtype=torch.float32)
+        if out_alloc_stream is not None:
+            with torch.cuda.stream(out_alloc_stream):
+                output = torch.empty((S, R), device=x.device, dtype=torch.float32)
+            output.zero_()
+        else:
+            output = torch.zeros((S, R), device=x.device, dtype=torch.float32)
         launch_kwargs = {
             "num_warps": 2 if split_k <= 4 else 4,
             "num_stages": 3,
@@ -267,3 +270,67 @@ def sgemm_lora_a_fwd(
     # split_k>1 returns the fp32 accumulator directly; the LoRA-B expand casts x to the weight dtype
     # on-load (fused), dropping the standalone fp32->bf16 copy kernel. split_k==1 already returns x.dtype.
     return output
+
+
+def shared_sink_sgemm_lora_a_fwd(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    *,
+    stack_num: int,
+    padded_rank: bool,
+    out_alloc_stream=None,
+) -> torch.Tensor:
+    """Shared-sink shrink with the measured fixed-width schedule."""
+    assert x.is_contiguous()
+    assert weights.is_contiguous()
+    assert x.ndim == 2
+    assert weights.ndim == 3
+
+    num_tokens = x.shape[0]
+    rank_width = weights.shape[-2]
+    input_width = weights.shape[-1]
+    assert x.shape[-1] == input_width
+
+    block_s = 16
+    block_k = 256
+    block_rank = 16
+    grid = (
+        triton.cdiv(batch_info.max_len, block_s) * triton.cdiv(rank_width, block_rank),
+        batch_info.bs,
+    )
+
+    if out_alloc_stream is None:
+        output = torch.empty((num_tokens, rank_width), device=x.device, dtype=x.dtype)
+    else:
+        with torch.cuda.stream(out_alloc_stream):
+            output = torch.empty(
+                (num_tokens, rank_width), device=x.device, dtype=x.dtype
+            )
+
+    _sgemm_lora_a_kernel[grid](
+        x,
+        weights,
+        output,
+        rank_width,
+        input_width,
+        stack_num,
+        x.stride(0),
+        x.stride(1),
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        batch_info.permutation is not None,
+        block_s,
+        block_rank,
+        block_k,
+        PADDED_RANK=padded_rank,
+    )
+    return output

python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_b.py

@@ -66,6 +66,10 @@ def _sgemm_lora_b_kernel(
     # For fused output scaling
     scalings,
     ENABLE_PDL: tl.constexpr = False,
+    APPLY_SCALING: tl.constexpr = True,
+    PADDED_RANK: tl.constexpr = True,
+    FLAT_GRID: tl.constexpr = False,
+    ATOMIC_ADD: tl.constexpr = True,
 ):
     """
     Computes a segmented batched matrix multiplication for the LoRA B matrix
@@ -84,9 +88,16 @@ def _sgemm_lora_b_kernel(
             the base model's output for a fused add operation.
     """
 
-    pid_s = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    batch_id = tl.program_id(axis=2)
+    if FLAT_GRID:
+        pid = tl.program_id(axis=0)
+        batch_id = tl.program_id(axis=1)
+        num_pid_n = tl.cdiv(N, BLOCK_N)
+        pid_s = pid // num_pid_n
+        pid_n = pid % num_pid_n
+    else:
+        pid_s = tl.program_id(axis=0)
+        pid_n = tl.program_id(axis=1)
+        batch_id = tl.program_id(axis=2)
     w_index = tl.load(weight_indices + batch_id)
     rank = tl.load(lora_ranks + w_index)
 
@@ -98,7 +109,9 @@ def _sgemm_lora_b_kernel(
     if pid_s * BLOCK_S >= seg_len:  # also covers seg_len == 0
         return
     seg_start = tl.load(seg_indptr + batch_id)
-    scaling = tl.load(scalings + w_index)
+    scaling = tl.load(scalings + w_index) if APPLY_SCALING else 1.0
+    if not PADDED_RANK:
+        K = tl.minimum(K, rank)
 
     s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
     n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
@@ -122,19 +135,24 @@ def _sgemm_lora_b_kernel(
     )
     output_mask = (s_offset[:, None] < seg_len) & n_mask
 
-    x_tile = tl.load(
-        x_ptrs,
-        mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K),
-        other=0.0,
-    )
-    w_tile = tl.load(
-        w_ptrs,
-        mask=(k_offset[:, None] < K) & n_mask,
-        other=0.0,
-    )
-
-    # cast fused: the split-K shrink returns fp32, plain path bf16 (no-op)
-    partial_sum = tl.dot(x_tile.to(w_tile.dtype), w_tile) * scaling
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        k_remaining = K - k * BLOCK_K
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < k_remaining),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < k_remaining) & n_mask,
+            other=0.0,
+        )
+        # The split-K shrink returns fp32; cast it on-load to the weight dtype.
+        partial_sum += tl.dot(x_tile.to(w_tile.dtype), w_tile)
+        x_ptrs += BLOCK_K * x_stride_1
+        w_ptrs += BLOCK_K * w_stride_2
+    partial_sum *= scaling
 
     # All input reads are done; hint the runtime to launch the dependent kernel.
     if ENABLE_PDL:
@@ -143,7 +161,11 @@ def _sgemm_lora_b_kernel(
     # Store result to output matrix (cast to the OUTPUT dtype: x may be the fp32
     # split-K shrink accumulator while base_output is bf16)
     partial_sum = partial_sum.to(output.dtype.element_ty)
-    tl.atomic_add(output_ptr, partial_sum, mask=output_mask, sem="relaxed")
+    if ATOMIC_ADD:
+        tl.atomic_add(output_ptr, partial_sum, mask=output_mask, sem="relaxed")
+    else:
+        partial_sum += tl.load(output_ptr, mask=output_mask, other=0.0)
+        tl.store(output_ptr, partial_sum, mask=output_mask)
 
 
 def sgemm_lora_b_fwd(
@@ -174,11 +196,14 @@ def sgemm_lora_b_fwd(
         )
         and S * R >= _CUBLAS_MIN_S_RANK
         and weights.shape[0] == 1
+        and x.dtype == weights.dtype
     ):  # single-adapter fast path: only valid with one resident slot
         return _sgemm_lora_b_cublas(x, weights, batch_info, base_output)
     # Block shapes
     BLOCK_S = 16
-    BLOCK_R = triton.next_power_of_2(R)
+    # Pad to >=16 for Triton MMA K>=16 (rank<16 adapters); k_offset < K=R masks the
+    # padded contraction rows to 0, so the result is unchanged.
+    BLOCK_R = max(16, triton.next_power_of_2(R))
     BLOCK_N = 256
 
     grid = (
@@ -221,3 +246,66 @@ def sgemm_lora_b_fwd(
         **pdl_kwargs,
     )
     return output
+
+
+def shared_sink_sgemm_lora_b_fwd(
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    base_output: torch.Tensor = None,
+    *,
+    apply_scaling: bool,
+    padded_rank: bool,
+) -> torch.Tensor:
+    """Shared-sink expand with the measured fixed-width schedule."""
+    assert x.is_contiguous()
+    assert weights.is_contiguous()
+    assert x.ndim == 2
+    assert weights.ndim == 3
+
+    num_tokens = x.shape[0]
+    output_width = weights.shape[-2]
+    rank_width = weights.shape[-1]
+    assert x.shape[-1] == rank_width
+
+    block_s = 16
+    block_rank = 16
+    block_n = 256
+    grid = (
+        triton.cdiv(batch_info.max_len, block_s) * triton.cdiv(output_width, block_n),
+        batch_info.bs,
+    )
+    output = (
+        torch.zeros((num_tokens, output_width), device=x.device, dtype=x.dtype)
+        if base_output is None
+        else base_output
+    )
+    _sgemm_lora_b_kernel[grid](
+        x,
+        weights,
+        output,
+        output_width,
+        rank_width,
+        x.stride(0),
+        x.stride(1),
+        weights.stride(0),
+        weights.stride(1),
+        weights.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        batch_info.lora_ranks,
+        batch_info.permutation,
+        batch_info.permutation is not None,
+        block_s,
+        block_n,
+        block_rank,
+        batch_info.scalings,
+        APPLY_SCALING=apply_scaling,
+        PADDED_RANK=padded_rank,
+        FLAT_GRID=True,
+        ATOMIC_ADD=False,
+    )
+    return output

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

@@ -452,3 +452,52 @@ def fused_conv_window_scatter_with_mask(
         dst_req_size,
         BLOCK_SIZE=BLOCK_SIZE,
     )
+
+
+def scatter_mamba_states_after_mtp_verify(
+    mamba_caches,
+    state_indices_tensor: torch.Tensor,
+    last_correct_step_indices: torch.Tensor,
+    mamba_track_indices: torch.Tensor | None,
+    mamba_steps_to_track: torch.Tensor | None,
+) -> None:
+    """Scatter per-step verify states (ssm + all conv types) into the
+    persistent caches, plus the interval-crossing track slots."""
+    ssm_states = mamba_caches.temporal
+    intermediate_state_cache = mamba_caches.intermediate_ssm
+
+    if ssm_states.numel() > 0:
+        fused_mamba_state_scatter_with_mask(
+            ssm_states,
+            intermediate_state_cache,
+            state_indices_tensor,
+            last_correct_step_indices,
+        )
+    for conv_states, intermediate_conv_window_cache in zip(
+        mamba_caches.conv, mamba_caches.intermediate_conv_window
+    ):
+        fused_conv_window_scatter_with_mask(
+            conv_states,
+            intermediate_conv_window_cache,
+            state_indices_tensor,
+            last_correct_step_indices,
+        )
+
+    if mamba_track_indices is not None:
+        assert mamba_steps_to_track is not None
+        if ssm_states.numel() > 0:
+            fused_mamba_state_scatter_with_mask(
+                ssm_states,
+                intermediate_state_cache,
+                mamba_track_indices,
+                mamba_steps_to_track,
+            )
+        for conv_states, intermediate_conv_window_cache in zip(
+            mamba_caches.conv, mamba_caches.intermediate_conv_window
+        ):
+            fused_conv_window_scatter_with_mask(
+                conv_states,
+                intermediate_conv_window_cache,
+                mamba_track_indices,
+                mamba_steps_to_track,
+            )

python/sglang/kernels/ops/moe/trtllm_lora_temp/virtual_experts.py

@@ -326,14 +326,12 @@ def _invoke_moe_lora_shrink_splitk(
     N = weight.shape[1]
     K = weight.shape[2]
     BLOCK_SIZE_M = config["BLOCK_SIZE_M"]
-    BLOCK_SIZE_N = triton.next_power_of_2(N)
+    BLOCK_SIZE_N = min(128, triton.next_power_of_2(N))
     BLOCK_SIZE_K = 256
     GROUP_SIZE_M = config.get("GROUP_SIZE_M", 1)
 
     num_m_blocks = triton.cdiv(sorted_token_ids.shape[0], BLOCK_SIZE_M)
-    num_n_blocks = triton.cdiv(
-        N, BLOCK_SIZE_N
-    )  # == 1, BLOCK_SIZE_N == next_pow2(N) >= N
+    num_n_blocks = triton.cdiv(N, BLOCK_SIZE_N)
     base_grid = num_m_blocks * num_n_blocks
     # Single source of truth shared with the caller's zero-intermediate decision:
     # split-K accumulation REQUIRES a pre-zeroed output, so the predicted and
@@ -378,26 +376,19 @@ def _get_moe_lora_shrink_split_k(
     sorted_token_ids: torch.Tensor,
     config: dict[str, Any],
 ) -> int:
-    """Rank-tiered split-K occupancy fill (PR #26899).
+    """Choose split-K from rank and available occupancy.
 
-    The K reduction (e.g. 7168 / 256 = 28 iters) dominates this skinny-N grouped
-    GEMV, so splitting K stays useful well past full SM occupancy -- a plain
-    `1 if base_grid >= num_sm else ...` rule collapses SPLIT_K too early and
-    costs up to ~2x at the decode/prefill border. Skinnier ranks want more
-    splits (their output tile carries less work). The target / tiers were picked
-    from an offline per-M B200 sweep over E in {48,96,384}, N in {16,32,64};
-    this heuristic lands within ~5% of the per-shape tuned optimum across the
-    decode regime.
+    Skinny output ranks benefit from more K splits because each output tile
+    carries less work. Block sizes must mirror _invoke_moe_lora_shrink_splitk.
 
-    Block sizes must mirror _invoke_moe_lora_shrink_splitk (BLOCK_SIZE_N =
-    next_pow2(N) -> one N block; BLOCK_SIZE_K = 256).
     """
     N = weight.shape[1]
     K = weight.shape[2]
     block_size_m = config["BLOCK_SIZE_M"]
+    block_size_n = min(128, triton.next_power_of_2(N))
     block_size_k = 256
     num_m_blocks = triton.cdiv(sorted_token_ids.shape[0], block_size_m)
-    base_grid = num_m_blocks  # num_n_blocks == 1: BLOCK_SIZE_N == next_pow2(N) >= N
+    base_grid = num_m_blocks * triton.cdiv(N, block_size_n)
     target = 512 if N <= 16 else 384 if N <= 32 else 256
     max_split_k = max(1, K // block_size_k)
     return max(1, min(triton.cdiv(target, base_grid), max_split_k, 8))
@@ -416,17 +407,10 @@ def _align_block_size_jit(
     block_size: int,
     num_experts: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """CUDA JIT align_block_size for num_experts > 1024 (up to 8191).
+    """CUDA JIT alignment for up to 8191 experts.
 
-    Uses the v2 kernel from moe_align_kernel.cu which supports large expert
-    counts via per-thread multi-expert processing and a two-level warp scan,
-    replacing the previous pure-PyTorch fallback that had excessive CPU overhead
-    from 15+ individual kernel launches and torch.argsort.
-
-    The JIT kernel uses a +1 offset convention: topk_ids are shifted by +1 so
-    that the EP sentinel value (-1) maps to bucket 0. The kernel internally
-    handles histogram, padded prefix-sum, expert_ids assignment, and token
-    scattering in just 2–3 CUDA kernel launches.
+    Expert IDs are shifted by one so ``-1`` maps to a sentinel bucket. The
+    fused allocation stays int4-aligned for the kernel's vectorized clear.
     """
     assert num_experts <= 8191, (
         f"_align_block_size_jit supports at most 8191 experts "
@@ -642,6 +626,10 @@ def _merged_experts_fused_moe_lora_add_impl(
     stage: str = "all",
     intermediate_buffer: torch.Tensor | None = None,
     expand_wait_event: "torch.cuda.Event | None" = None,
+    broadcast_intermediate: bool = False,
+    prewarm_a_routing: bool = True,
+    prewarm_b_routing: bool = True,
+    zero_intermediate: bool = False,
 ) -> "torch.Tensor | None":
     """
     1. Prepare virtual expert routing metadata from topk_ids + token_lora_mapping * num_experts.
@@ -657,12 +645,14 @@ def _merged_experts_fused_moe_lora_add_impl(
     - ``"expand"``: routing-B + LoRA-B expand/add only; requires ``intermediate_buffer`` =
       the tensor produced by the ``"shrink"`` stage.
 
-    EP: when `local_num_experts` (< global) is given, this rank only computes the
-    delta for the experts it owns. We keep the GLOBAL expert ids + global contiguous
-    weights (so the merged-weight reshape stays a free view) and mask non-owned
-    [token, k] slots to the -1 sentinel inside `_fused_virtual_topk_ids_kernel`; the
-    grid shrinks via the per-rank trim in `_get_routing`. Slicing the weight's expert
-    dim instead would force the reshape to copy every step (non-contiguous fold).
+    ``prewarm_a_routing`` and ``prewarm_b_routing`` let a staged caller skip
+    routing for a weight that it replaces with a dense operation. The B flag
+    also controls the automatic expand-route prewarm performed by ``"shrink"``.
+    ``broadcast_intermediate`` is an expand-only mode where one rank vector per
+    token is reused for every routed expert.
+
+    EP accepts either global weights/IDs with a local range or already-localized
+    weights/IDs from the standard dispatcher.
     """
     max_loras, _, max_lora_rank, _ = lora_a.shape
     # Global per-expert dim of the LoRA weights. lora_a may be shared-outer (expert
@@ -759,24 +749,18 @@ def _merged_experts_fused_moe_lora_add_impl(
             if cached is not None:
                 return cached
 
-        # Fused LoRA-local align: one kernel does inline virtual id + EP skip +
-        # compact (local experts) + single-block scatter, replacing the 3-kernel
-        # (_fused_virtual_topk_ids + moe_align + count_and_sort) pipeline. Two
-        # single-adapter (max_loras==1) regimes are fused; everything else falls back:
-        #   - per-expert EP path (ep_local): compact local-expert histogram.
-        #   - shared-outer path (shared_outer): lora-id routing (compute_virtual_id
-        #     uses base=0; the kernel + launcher already size num_experts_for_weight=1
-        #     and have no bucket-count blocker, so it just needs compact=False — compact
-        #     + shared_outer would mis-map the id as base-offset). This is the opt1
-        #     align/sort fusion: shared-outer used to fall through to the unfused
-        #     _fused_virtual_topk_ids + moe_align_block_size_small_batch pair (~10.2us/
-        #     layer at decode bs16); now it takes the single fused launch.
-        # Decode-only: the fused kernel's single-block scatter targets the small
-        # decode batch; prefill (>= 512 tokens) keeps the multi-block old path.
+        # Shared-outer routing has one bucket per adapter, so the same merged
+        # align remains valid for multi-LoRA. Compact EP routing stays single-slot.
+        compact_merged = ep_local and not shared_outer and max_loras == 1
+        bucket_experts = (
+            local_num_experts
+            if compact_merged
+            else (1 if shared_outer else num_experts) * max_loras
+        )
         if (
             lora_envs.SGLANG_OPT_LORA_FUSED_MERGED_ALIGN.get()
-            and max_loras == 1
-            and (shared_outer or ep_local)
+            and (shared_outer or compact_merged)
+            and bucket_experts + 1 <= 1024
             and topk_ids.shape[0] < 512
         ):
             from sglang.jit_kernel.trtllm_lora_temp.moe_lora_merged_align import (
@@ -799,9 +783,7 @@ def _merged_experts_fused_moe_lora_add_impl(
                 local_expert_offset,
                 local_num_experts,
                 do_skip=True,
-                # compact local-expert histogram is only valid for the per-expert EP
-                # path; shared_outer routes by lora id (base=0) so it must stay global.
-                compact=not shared_outer,
+                compact=compact_merged,
             )
             result = (
                 sorted_token_ids,
@@ -874,6 +856,12 @@ def _merged_experts_fused_moe_lora_add_impl(
         "expand",
         "routing",
     ), f"invalid stage {stage!r}"
+    if broadcast_intermediate:
+        assert stage == "expand"
+        assert use_direct_expand_add
+        assert intermediate_buffer is not None
+        assert intermediate_buffer.ndim == 2
+        assert intermediate_buffer.shape[0] == token_lora_mapping.shape[0]
     lora_a_virtual = _merge_lora_expert_weight(lora_a)
     lora_b_virtual = _merge_lora_expert_weight(lora_b)
     num_experts_a = lora_a.shape[1]
@@ -891,20 +879,32 @@ def _merged_experts_fused_moe_lora_add_impl(
         a_cfg = _get_shrink_stage_config(lora_a_virtual, token_lora_mapping.shape[0])
         if lora_envs.SGLANG_OPT_LORA_SHRINK_TUNE.get():
             a_cfg = {**a_cfg, "BLOCK_SIZE_M": 16}
-        _get_routing(
-            topk_ids,
-            token_lora_mapping,
-            num_experts_a,
-            experts_shared_outer_loras_a,
-            a_cfg["BLOCK_SIZE_M"],
-        )
-        _get_routing(
-            topk_ids,
-            token_lora_mapping,
-            num_experts_b,
-            experts_shared_outer_loras_b,
-            b_stage_config["BLOCK_SIZE_M"],
-        )
+        # Match the actual shrink-stage override below. Without this, callers
+        # that admit prefill-shaped batches into a side stream prewarm block-32
+        # routing here, then miss the cache when shrink switches to the B-stage
+        # block size. The miss allocates routing buffers on the side stream
+        # during capture, violating the allocation guarantee of stage='routing'.
+        if (
+            lora_envs.SGLANG_OPT_LORA_PREFILL_ROUTING_REUSE.get()
+            and token_lora_mapping.shape[0] >= 512
+        ):
+            a_cfg["BLOCK_SIZE_M"] = b_stage_config["BLOCK_SIZE_M"]
+        if prewarm_a_routing:
+            _get_routing(
+                topk_ids,
+                token_lora_mapping,
+                num_experts_a,
+                experts_shared_outer_loras_a,
+                a_cfg["BLOCK_SIZE_M"],
+            )
+        if prewarm_b_routing:
+            _get_routing(
+                topk_ids,
+                token_lora_mapping,
+                num_experts_b,
+                experts_shared_outer_loras_b,
+                b_stage_config["BLOCK_SIZE_M"],
+            )
         return None
 
     intermediate = intermediate_buffer
@@ -913,21 +913,14 @@ def _merged_experts_fused_moe_lora_add_impl(
             lora_a_virtual, token_lora_mapping.shape[0]
         )
         if lora_envs.SGLANG_OPT_LORA_SHRINK_TUNE.get():
-            # GB200 hand-tune knob (test-only) on top of PR #26899's heuristic config. The launcher
-            # pins BLOCK_SIZE_N (next_pow2(rank)) and BLOCK_SIZE_K (256), so only M/warps/stages apply.
+            # Test-only override; the launcher fixes N and K block sizes.
             a_stage_config = {
                 **a_stage_config,
                 "BLOCK_SIZE_M": 16,
                 "num_warps": 4,
                 "num_stages": 4,
             }
-        # F1-① prefill routing reuse: the A stage routes with BLOCK_SIZE_M 32 at prefill
-        # but the B stage with the tuned fused-moe config (typically 64), so the
-        # (num_experts, shared_outer, block_size) routing_cache key never matches across
-        # stages and the align/sort pipeline reruns per stage (4x/layer at prefill).
-        # Matching the A stage's routing block to the B stage's collapses them to one
-        # align/sort per layer-forward. Decode (<512 tokens) keeps the opt1 fused
-        # merged-align path and its tuned shrink block untouched.
+        # Match routing block sizes so prefill stages can share cached alignment.
         if (
             lora_envs.SGLANG_OPT_LORA_PREFILL_ROUTING_REUSE.get()
             and token_lora_mapping.shape[0] >= 512
@@ -957,8 +950,10 @@ def _merged_experts_fused_moe_lora_add_impl(
         # non-owned blocks (never reads them), but a shared-outer expand routes by lora id
         # and would read them into the real (all-reduced) output -> must zero. split_k > 1
         # also needs a zeroed buffer for its accumulation.
-        zero_intermediate = intermediate_split_k > 1 or (
-            ep_local and experts_shared_outer_loras_b
+        must_zero_intermediate = (
+            zero_intermediate
+            or intermediate_split_k > 1
+            or (ep_local and experts_shared_outer_loras_b)
         )
         if intermediate is None:
             intermediate = (
@@ -967,14 +962,14 @@ def _merged_experts_fused_moe_lora_add_impl(
                     dtype=hidden_states.dtype,
                     device=hidden_states.device,
                 )
-                if zero_intermediate
+                if must_zero_intermediate
                 else torch.empty(
                     intermediate_shape,
                     dtype=hidden_states.dtype,
                     device=hidden_states.device,
                 )
             )
-        elif zero_intermediate:
+        elif must_zero_intermediate:
             # Caller-provided buffer (allocated on the consumer stream): zero it in-stream.
             intermediate.zero_()
 
@@ -993,7 +988,7 @@ def _merged_experts_fused_moe_lora_add_impl(
         if stage == "shrink":
             # Pre-warm the routing-B cache on this (side) stream so the later "expand" stage
             # launches no routing kernels — they overlap finalize together with the shrink.
-            if routing_cache is not None:
+            if routing_cache is not None and prewarm_b_routing:
                 _get_routing(
                     topk_ids,
                     token_lora_mapping,
@@ -1041,8 +1036,13 @@ def _merged_experts_fused_moe_lora_add_impl(
             b_stage_config,
             mul_routed_weight,
             fuse_sum_all_reduce,
+            broadcast_intermediate=broadcast_intermediate,
         )
     else:
+        assert not broadcast_intermediate, (
+            "broadcasted LoRA-A intermediates require the rank-specialized "
+            "direct expand kernel"
+        )
         invoke_fused_moe_kernel(
             intermediate_flat,
             lora_b_virtual,
@@ -1131,6 +1131,10 @@ def merged_experts_fused_moe_lora_add(
     stage: str = "all",
     intermediate_buffer: torch.Tensor | None = None,
     expand_wait_event: "torch.cuda.Event | None" = None,
+    broadcast_intermediate: bool = False,
+    prewarm_a_routing: bool = True,
+    prewarm_b_routing: bool = True,
+    zero_intermediate: bool = False,
 ) -> "torch.Tensor | None":
     """Public API: wraps the registered op with routing_cache support."""
     return _merged_experts_fused_moe_lora_add_impl(
@@ -1153,4 +1157,8 @@ def merged_experts_fused_moe_lora_add(
         stage=stage,
         intermediate_buffer=intermediate_buffer,
         expand_wait_event=expand_wait_event,
+        broadcast_intermediate=broadcast_intermediate,
+        prewarm_a_routing=prewarm_a_routing,
+        prewarm_b_routing=prewarm_b_routing,
+        zero_intermediate=zero_intermediate,
     )

python/sglang/kernels/ops/speculative/multi_layer_eagle.py

@@ -12,6 +12,7 @@
 # limitations under the License.
 # ==============================================================================
 
+import torch
 import triton
 import triton.language as tl
 
@@ -94,3 +95,636 @@ def rotate_input_ids(
         BLOCK_SIZE=BLOCK_SIZE,
     )
     return input_ids
+
+
+@triton.jit
+def stash_append_boundary_state_kernel(
+    # flat sources (decode: predict + verify FULL hiddens; prefill: rotated
+    # input_ids + target FULL hiddens)
+    src_tokens_ptr,
+    src_hiddens_ptr,  # [num_src_rows, hidden]
+    src_row_ends_ptr,  # [bs] exclusive end row of each request's source segment
+    num_available_ptr,  # [bs] committed rows at the segment tail (accept_lens / extend len)
+    req_pool_indices_ptr,  # [bs]
+    # stash (per request, rolling last `front` committed (token, base-hidden)
+    # pairs; slot j of a request at boundary B holds position B - front + j)
+    stash_tokens_ptr,  # [req_pool_size, front] int64
+    stash_hiddens_ptr,  # [req_pool_size, front, hidden]
+    stash_valid_lens_ptr,  # [req_pool_size] int32, count of valid tail slots
+    front: tl.constexpr,
+    hidden_dim: tl.constexpr,
+    SET_VALID: tl.constexpr,  # prefill: valid = m; decode: valid = min(valid + m, front)
+    BLOCK_H: tl.constexpr,
+):
+    """Roll the per-request boundary stash forward by m = min(available, front)
+    newly committed (token, base-hidden) pairs taken from the source tail
+    rows [end - m, end). Kept old pairs shift down (reads stay ahead of
+    writes, ascending order)."""
+    pid = tl.program_id(0)
+    rpi = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
+    end = tl.load(src_row_ends_ptr + pid).to(tl.int64)
+    avail = tl.load(num_available_ptr + pid).to(tl.int64)
+    m = tl.minimum(avail, front)
+    keep = front - m
+
+    h_off = tl.arange(0, BLOCK_H)
+
+    # 1) Shift the kept tail of the old stash to the front: new[i] = old[i + m]
+    for i in range(0, keep):
+        src_t = tl.load(stash_tokens_ptr + rpi * front + i + m)
+        tl.store(stash_tokens_ptr + rpi * front + i, src_t)
+        for hb in range(0, hidden_dim, BLOCK_H):
+            hmask = (hb + h_off) < hidden_dim
+            src_h = tl.load(
+                stash_hiddens_ptr + (rpi * front + i + m) * hidden_dim + hb + h_off,
+                mask=hmask,
+            )
+            tl.store(
+                stash_hiddens_ptr + (rpi * front + i) * hidden_dim + hb + h_off,
+                src_h,
+                mask=hmask,
+            )
+
+    # 2) Append the m newest committed pairs from the source tail.
+    for i in range(0, m):
+        row = end - m + i
+        dst = rpi * front + keep + i
+        tok = tl.load(src_tokens_ptr + row)
+        tl.store(stash_tokens_ptr + dst, tok)
+        for hb in range(0, hidden_dim, BLOCK_H):
+            hmask = (hb + h_off) < hidden_dim
+            src_h = tl.load(src_hiddens_ptr + row * hidden_dim + hb + h_off, mask=hmask)
+            tl.store(
+                stash_hiddens_ptr + dst * hidden_dim + hb + h_off, src_h, mask=hmask
+            )
+
+    if SET_VALID:
+        valid = m
+    else:
+        valid = tl.minimum(tl.load(stash_valid_lens_ptr + rpi).to(tl.int64) + m, front)
+    tl.store(stash_valid_lens_ptr + rpi, valid.to(tl.int32))
+
+
+def stash_append_boundary_state_triton(
+    src_tokens,
+    src_hiddens,
+    src_row_ends,
+    num_available,
+    req_pool_indices,
+    stash_tokens,
+    stash_hiddens,
+    stash_valid_lens,
+    set_valid: bool,
+):
+    """Append newly committed (token, base-hidden) pairs to the rolling
+    boundary stash (see kernel docstring). Decode: sources are (predict,
+    verify FULL hiddens) with ends = i*W + accept_lens. Prefill: sources are
+    (post-rotation input_ids, target FULL hiddens) with ends = start + len."""
+    bs = req_pool_indices.shape[0]
+    if bs == 0:
+        return
+    stash_append_boundary_state_kernel[(bs,)](
+        src_tokens,
+        src_hiddens,
+        src_row_ends,
+        num_available,
+        req_pool_indices,
+        stash_tokens,
+        stash_hiddens,
+        stash_valid_lens,
+        front=stash_tokens.shape[1],
+        hidden_dim=stash_hiddens.shape[2],
+        SET_VALID=set_valid,
+        BLOCK_H=1024,
+    )
+
+
+@triton.jit
+def fill_widened_draft_extend_inputs_kernel(
+    # outputs: the widened per-request window buffers, width = W + front
+    input_ids_ptr,  # [bs * width]
+    hidden_ptr,  # [bs * width, hidden]
+    # sources
+    predict_ptr,  # [bs * W] verify-sampled successor per verify row
+    verify_hidden_ptr,  # [bs * W, hidden] target verify hiddens (FULL capture)
+    stash_tokens_ptr,  # [req_pool_size, front]
+    stash_hiddens_ptr,  # [req_pool_size, front, hidden]
+    stash_valid_lens_ptr,  # [req_pool_size]
+    seq_lens_ptr,  # [bs] PRE-verify seq_lens (window base = seq_lens - front)
+    req_pool_indices_ptr,  # [bs]
+    draft_token_num: tl.constexpr,  # W
+    front: tl.constexpr,  # F_total
+    hidden_dim: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    """Materialize the widened depth-0 window's input tokens and hiddens: front
+    rows (j < front) source from stash slot j, original rows (j >= front) from
+    predict/verify hiddens; data-invalid front rows are zeroed. Locs/positions
+    are computed separately by compute_widened_draft_extend_locs_positions."""
+    pid = tl.program_id(0)
+    rpi = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
+    seq_len = tl.load(seq_lens_ptr + pid).to(tl.int64)
+    valid_len = tl.load(stash_valid_lens_ptr + rpi).to(tl.int64)
+
+    # Rows below this hold no usable stash data (unseeded slot or p < 0).
+    first_valid = tl.maximum(tl.maximum(front - valid_len, front - seq_len), 0)
+
+    h_off = tl.arange(0, BLOCK_H)
+    width = draft_token_num + front
+
+    for j in range(0, width):
+        row = pid * width + j
+        if j >= front:
+            src = pid * draft_token_num + j - front
+            tok = tl.load(predict_ptr + src)
+            tl.store(input_ids_ptr + row, tok)
+            for hb in range(0, hidden_dim, BLOCK_H):
+                hmask = (hb + h_off) < hidden_dim
+                src_h = tl.load(
+                    verify_hidden_ptr + src * hidden_dim + hb + h_off, mask=hmask
+                )
+                tl.store(hidden_ptr + row * hidden_dim + hb + h_off, src_h, mask=hmask)
+        else:
+            if j >= first_valid:
+                tok = tl.load(stash_tokens_ptr + rpi * front + j)
+                tl.store(input_ids_ptr + row, tok)
+                for hb in range(0, hidden_dim, BLOCK_H):
+                    hmask = (hb + h_off) < hidden_dim
+                    src_h = tl.load(
+                        stash_hiddens_ptr + (rpi * front + j) * hidden_dim + hb + h_off,
+                        mask=hmask,
+                    )
+                    tl.store(
+                        hidden_ptr + row * hidden_dim + hb + h_off, src_h, mask=hmask
+                    )
+            else:
+                tl.store(input_ids_ptr + row, 0)
+                for hb in range(0, hidden_dim, BLOCK_H):
+                    hmask = (hb + h_off) < hidden_dim
+                    tl.store(
+                        hidden_ptr + row * hidden_dim + hb + h_off, 0.0, mask=hmask
+                    )
+
+
+def fill_widened_draft_extend_inputs_triton(
+    input_ids,
+    hidden_states,
+    predict,
+    verify_hiddens,
+    stash_tokens,
+    stash_hiddens,
+    stash_valid_lens,
+    seq_lens,
+    req_pool_indices,
+    draft_token_num: int,
+):
+    """Fill the widened window's input tokens and hiddens in place (see kernel
+    docstring). Must run AFTER verify sampling (reads predict / hiddens) and
+    BEFORE the stash update for this iteration (the stash is still based at
+    the pre-verify boundary)."""
+    bs = req_pool_indices.shape[0]
+    if bs == 0:
+        return
+    fill_widened_draft_extend_inputs_kernel[(bs,)](
+        input_ids,
+        hidden_states,
+        predict,
+        verify_hiddens,
+        stash_tokens,
+        stash_hiddens,
+        stash_valid_lens,
+        seq_lens,
+        req_pool_indices,
+        draft_token_num=draft_token_num,
+        front=stash_tokens.shape[1],
+        hidden_dim=stash_hiddens.shape[2],
+        BLOCK_H=1024,
+    )
+
+
+@triton.jit
+def _wide_row_softmax_partials_kernel(
+    logits_ptr,  # [bs, vocab] fp32
+    temperatures_ptr,  # [bs, 1] fp32 (dummy when HAS_TEMPS is False)
+    partial_max_ptr,  # [bs, nblocks] fp32
+    partial_sum_ptr,  # [bs, nblocks] fp32
+    vocab,
+    nblocks,
+    HAS_TEMPS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    blk = tl.program_id(1)
+    offs = blk * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < vocab
+    z = tl.load(
+        logits_ptr + row.to(tl.int64) * vocab + offs, mask=mask, other=-float("inf")
+    )
+    if HAS_TEMPS:
+        z = z / tl.load(temperatures_ptr + row)
+    m = tl.max(z, axis=0)
+    s = tl.sum(tl.exp(z - m), axis=0)
+    tl.store(partial_max_ptr + row * nblocks + blk, m)
+    tl.store(partial_sum_ptr + row * nblocks + blk, s)
+
+
+@triton.jit
+def _wide_row_softmax_finalize_kernel(
+    partial_max_ptr,
+    partial_sum_ptr,
+    row_max_ptr,  # [bs] fp32
+    row_sum_ptr,  # [bs] fp32
+    nblocks,
+    NBLOCK_POW2: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offs = tl.arange(0, NBLOCK_POW2)
+    mask = offs < nblocks
+    m = tl.load(partial_max_ptr + row * nblocks + offs, mask=mask, other=-float("inf"))
+    s = tl.load(partial_sum_ptr + row * nblocks + offs, mask=mask, other=0.0)
+    gm = tl.max(m, axis=0)
+    gs = tl.sum(s * tl.exp(m - gm), axis=0)
+    tl.store(row_max_ptr + row, gm)
+    tl.store(row_sum_ptr + row, gs)
+
+
+@triton.jit
+def _wide_row_softmax_write_kernel(
+    logits_ptr,
+    temperatures_ptr,
+    row_max_ptr,
+    row_sum_ptr,
+    out_ptr,  # [bs, out_row_stride] fp32; row i written at i * out_row_stride
+    vocab,
+    out_row_stride,
+    HAS_TEMPS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    blk = tl.program_id(1)
+    offs = blk * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < vocab
+    z = tl.load(
+        logits_ptr + row.to(tl.int64) * vocab + offs, mask=mask, other=-float("inf")
+    )
+    if HAS_TEMPS:
+        z = z / tl.load(temperatures_ptr + row)
+    gm = tl.load(row_max_ptr + row)
+    gs = tl.load(row_sum_ptr + row)
+    q = tl.exp(z - gm) / gs
+    tl.store(out_ptr + row.to(tl.int64) * out_row_stride + offs, q, mask=mask)
+
+
+def wide_row_softmax_triton(
+    logits: torch.Tensor,
+    temperatures,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    """Column-parallel softmax over very wide fp32 rows, optionally with
+    per-row temperature (q = softmax(logits / T)), written into ``out``
+    (any row stride >= vocab). torch.softmax gives one block per row, which
+    serializes a single wide draft-vocab row onto one SM."""
+    bs, vocab = logits.shape
+    BLOCK = 4096
+    nblocks = triton.cdiv(vocab, BLOCK)
+    partial_max = torch.empty((bs, nblocks), dtype=torch.float32, device=logits.device)
+    partial_sum = torch.empty((bs, nblocks), dtype=torch.float32, device=logits.device)
+    row_max = torch.empty((bs,), dtype=torch.float32, device=logits.device)
+    row_sum = torch.empty((bs,), dtype=torch.float32, device=logits.device)
+    has_temps = temperatures is not None
+    dummy = row_max
+    _wide_row_softmax_partials_kernel[(bs, nblocks)](
+        logits,
+        temperatures if has_temps else dummy,
+        partial_max,
+        partial_sum,
+        vocab,
+        nblocks,
+        HAS_TEMPS=has_temps,
+        BLOCK=BLOCK,
+    )
+    _wide_row_softmax_finalize_kernel[(bs,)](
+        partial_max,
+        partial_sum,
+        row_max,
+        row_sum,
+        nblocks,
+        NBLOCK_POW2=triton.next_power_of_2(nblocks),
+    )
+    _wide_row_softmax_write_kernel[(bs, nblocks)](
+        logits,
+        temperatures if has_temps else dummy,
+        row_max,
+        row_sum,
+        out,
+        vocab,
+        out.stride(0),
+        HAS_TEMPS=has_temps,
+        BLOCK=BLOCK,
+    )
+    return out
+
+
+@triton.jit
+def compute_widened_draft_extend_locs_positions_kernel(
+    seq_lens_ptr,
+    req_pool_indices_ptr,
+    req_to_token_ptr,
+    stash_valid_lens_ptr,
+    locs_ptr,  # [bs * width] int64
+    positions_ptr,  # [bs * width] int64
+    req_to_token_stride,
+    front,
+    num_warmup_tokens,
+    width,
+    WIDTH_BLOCK: tl.constexpr,
+):
+    """Per-request widened-window locs + positions: pos = seq_len - front + j;
+    rows below first_valid = max(front - stash_valid, front - seq_len, 0) hold
+    no stash data (positions zeroed), and the first num_warmup_tokens valid
+    front rows write to sacrificial loc 0."""
+    pid = tl.program_id(0)
+    offs = tl.arange(0, WIDTH_BLOCK)
+    wmask = offs < width
+    offs64 = offs.to(tl.int64)
+
+    seq_len = tl.load(seq_lens_ptr + pid).to(tl.int64)
+    rpi = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
+    valid_len = tl.load(stash_valid_lens_ptr + rpi).to(tl.int64)
+
+    pos = seq_len - front + offs64
+    first_valid = tl.maximum(tl.maximum(front - valid_len, front - seq_len), 0)
+    data_valid = offs64 >= first_valid
+    write_real = offs64 >= tl.minimum(first_valid + num_warmup_tokens, front)
+
+    tok = tl.load(
+        req_to_token_ptr + rpi * req_to_token_stride + tl.maximum(pos, 0),
+        mask=wmask,
+        other=0,
+    ).to(tl.int64)
+    locs = tl.where(write_real, tok, 0)
+    positions = tl.where(data_valid, pos, 0)
+
+    base = pid.to(tl.int64) * width
+    tl.store(locs_ptr + base + offs, locs, mask=wmask)
+    tl.store(positions_ptr + base + offs, positions, mask=wmask)
+
+
+def compute_widened_draft_extend_locs_positions_triton(
+    seq_lens,
+    req_pool_indices,
+    req_to_token,
+    stash_valid_lens,
+    draft_token_num: int,
+    num_front_tokens: int,
+    num_warmup_tokens: int,
+):
+    width = draft_token_num + num_front_tokens
+    bs = seq_lens.shape[0]
+    locs = torch.empty((bs * width,), dtype=torch.int64, device=seq_lens.device)
+    positions = torch.empty((bs * width,), dtype=torch.int64, device=seq_lens.device)
+    if bs > 0:
+        compute_widened_draft_extend_locs_positions_kernel[(bs,)](
+            seq_lens,
+            req_pool_indices,
+            req_to_token,
+            stash_valid_lens,
+            locs,
+            positions,
+            req_to_token.stride(0),
+            num_front_tokens,
+            num_warmup_tokens,
+            width,
+            WIDTH_BLOCK=triton.next_power_of_2(width),
+        )
+    return locs, positions
+
+
+@triton.jit
+def fill_draft_extend_prepare_buffers_kernel(
+    # persistent per-token buffers (length max_num_token, int64)
+    input_ids_ptr,
+    positions_ptr,
+    out_cache_loc_ptr,
+    # per-token sources (length num_tokens)
+    src_input_ids_ptr,
+    src_positions_ptr,
+    src_out_cache_loc_ptr,
+    # persistent per-request buffers (length max_bs)
+    seq_lens_ptr,  # int32
+    req_pool_indices_ptr,  # int64
+    num_correct_drafts_ptr,  # int32
+    num_accept_tokens_ptr,  # int32
+    select_index_ptr,  # int64
+    temperatures_ptr,  # float32 [max_bs, 1] (dummy when HAS_TEMPS is False)
+    # per-request sources (length raw_bs)
+    src_seq_lens_ptr,
+    src_req_pool_indices_ptr,
+    src_num_correct_drafts_ptr,
+    src_num_accept_tokens_ptr,
+    src_temperatures_ptr,  # dummy when HAS_TEMPS is False
+    # chain hidden window, flat [num_tokens * hidden] (dummies when HAS_HIDDEN
+    # is False)
+    hidden_states_ptr,
+    src_hidden_states_ptr,
+    # gathered-buffer mirrors (dummies when HAS_GATHERED is False)
+    global_num_tokens_ptr,
+    global_num_tokens_for_logprob_ptr,
+    # scalars
+    num_tokens,
+    max_num_token,
+    raw_bs,
+    bs,
+    max_bs,
+    num_tokens_per_bs,
+    num_front_tokens,
+    seq_len_fill_value,
+    hidden_numel,
+    num_global,
+    num_token_programs,
+    HAS_TEMPS: tl.constexpr,
+    HAS_HIDDEN: tl.constexpr,
+    HAS_GATHERED: tl.constexpr,
+    BLOCK_TOK: tl.constexpr,
+    BLOCK_HIDDEN: tl.constexpr,
+    GLOBAL_BLOCK: tl.constexpr,
+):
+    """The whole draft-extend prepare() buffer population in one launch;
+    program roles split by flat program id:
+
+    - [0, num_token_programs): input_ids / positions / out_cache_loc; rows
+      < num_tokens take the source, the tail up to max_num_token is zeroed.
+    - [num_token_programs, +max_bs): one program per request row. Real rows
+      [0, raw_bs) take source values; padded rows [raw_bs, bs) take the pad
+      sentinels the graphs rely on (seq_len fill value, num_accept_tokens = -1,
+      temperatures = 1.0); rows >= bs are untouched except seq_lens, which is
+      fully reset. select_index = i*window + front + num_correct_drafts
+      (padded rows keep their stale num_correct_drafts, whose gather result
+      is discarded).
+    - num_token_programs + max_bs: the DP gathered-buffer fills.
+    - the rest: flat copy of the chain hidden window's real rows (the padded
+      tail is never read).
+    """
+    pid = tl.program_id(0)
+
+    if pid < num_token_programs:
+        tok_offs = pid * BLOCK_TOK + tl.arange(0, BLOCK_TOK)
+        store_mask = tok_offs < max_num_token
+        copy_mask = tok_offs < num_tokens
+
+        tok = tl.load(src_input_ids_ptr + tok_offs, mask=copy_mask, other=0).to(
+            tl.int64
+        )
+        tl.store(input_ids_ptr + tok_offs, tok, mask=store_mask)
+        tok = tl.load(src_positions_ptr + tok_offs, mask=copy_mask, other=0).to(
+            tl.int64
+        )
+        tl.store(positions_ptr + tok_offs, tok, mask=store_mask)
+        tok = tl.load(src_out_cache_loc_ptr + tok_offs, mask=copy_mask, other=0).to(
+            tl.int64
+        )
+        tl.store(out_cache_loc_ptr + tok_offs, tok, mask=store_mask)
+    elif pid < num_token_programs + max_bs:
+        i = pid - num_token_programs
+        is_real = i < raw_bs
+        is_pad = (i >= raw_bs) & (i < bs)
+        in_bs = i < bs
+
+        sl = tl.load(src_seq_lens_ptr + i, mask=is_real, other=seq_len_fill_value)
+        tl.store(seq_lens_ptr + i, sl.to(tl.int32))
+
+        rpi = tl.load(src_req_pool_indices_ptr + i, mask=is_real, other=0).to(tl.int64)
+        tl.store(req_pool_indices_ptr + i, rpi, mask=is_real)
+
+        # The stale count must be read BEFORE the real-row store below.
+        ncd_stale = tl.load(num_correct_drafts_ptr + i, mask=is_pad, other=0).to(
+            tl.int64
+        )
+        ncd_src = tl.load(src_num_correct_drafts_ptr + i, mask=is_real, other=0).to(
+            tl.int64
+        )
+        tl.store(num_correct_drafts_ptr + i, ncd_src.to(tl.int32), mask=is_real)
+        ncd = tl.where(is_real, ncd_src, ncd_stale)
+
+        # Padded rows get -1 so the sconv commit skips their live mamba slots.
+        nat = tl.load(src_num_accept_tokens_ptr + i, mask=is_real, other=-1).to(
+            tl.int32
+        )
+        tl.store(num_accept_tokens_ptr + i, nat, mask=in_bs)
+
+        si = i.to(tl.int64) * num_tokens_per_bs + num_front_tokens + ncd
+        tl.store(select_index_ptr + i, si, mask=in_bs)
+
+        if HAS_TEMPS:
+            t = tl.load(src_temperatures_ptr + i, mask=is_real, other=1.0)
+            tl.store(temperatures_ptr + i, t, mask=in_bs)
+    elif pid == num_token_programs + max_bs:
+        if HAS_GATHERED:
+            g_offs = tl.arange(0, GLOBAL_BLOCK)
+            g_mask = g_offs < num_global
+            g_vals = tl.zeros((GLOBAL_BLOCK,), dtype=tl.int32) + bs * num_tokens_per_bs
+            tl.store(global_num_tokens_ptr + g_offs, g_vals, mask=g_mask)
+            tl.store(global_num_tokens_for_logprob_ptr + g_offs, g_vals, mask=g_mask)
+    else:
+        if HAS_HIDDEN:
+            h_base = pid - num_token_programs - max_bs - 1
+            h_offs = h_base.to(tl.int64) * BLOCK_HIDDEN + tl.arange(0, BLOCK_HIDDEN)
+            h_mask = h_offs < hidden_numel
+            h_vals = tl.load(src_hidden_states_ptr + h_offs, mask=h_mask)
+            tl.store(hidden_states_ptr + h_offs, h_vals, mask=h_mask)
+
+
+def fill_draft_extend_prepare_buffers_triton(
+    input_ids,
+    positions,
+    out_cache_loc,
+    src_input_ids,
+    src_positions,
+    src_out_cache_loc,
+    seq_lens,
+    req_pool_indices,
+    num_correct_drafts,
+    num_accept_tokens,
+    select_index,
+    temperatures,
+    src_seq_lens,
+    src_req_pool_indices,
+    src_num_correct_drafts,
+    src_num_accept_tokens,
+    src_temperatures,
+    hidden_states,
+    src_hidden_states,
+    global_num_tokens,
+    global_num_tokens_for_logprob,
+    raw_bs,
+    bs,
+    num_tokens_per_bs,
+    num_front_tokens,
+    seq_len_fill_value,
+):
+    max_num_token = input_ids.shape[0]
+    num_tokens = src_input_ids.shape[0]
+    max_bs = seq_lens.shape[0]
+    has_temps = temperatures is not None
+    has_hidden = src_hidden_states is not None
+    has_gathered = global_num_tokens is not None
+
+    BLOCK_TOK = 1024
+    BLOCK_HIDDEN = 2048
+    num_token_programs = triton.cdiv(max_num_token, BLOCK_TOK)
+
+    if has_hidden:
+        hidden_numel = num_tokens * hidden_states.shape[1]
+        num_hidden_programs = triton.cdiv(hidden_numel, BLOCK_HIDDEN)
+    else:
+        hidden_numel = 0
+        num_hidden_programs = 0
+
+    if has_gathered:
+        num_global = global_num_tokens.shape[0]
+        global_block = triton.next_power_of_2(num_global)
+    else:
+        num_global = 0
+        global_block = 1
+
+    grid = (num_token_programs + max_bs + 1 + num_hidden_programs,)
+    fill_draft_extend_prepare_buffers_kernel[grid](
+        input_ids,
+        positions,
+        out_cache_loc,
+        src_input_ids,
+        src_positions,
+        src_out_cache_loc,
+        seq_lens,
+        req_pool_indices,
+        num_correct_drafts,
+        num_accept_tokens,
+        select_index,
+        temperatures if has_temps else seq_lens,
+        src_seq_lens,
+        src_req_pool_indices,
+        src_num_correct_drafts,
+        src_num_accept_tokens,
+        src_temperatures if has_temps else seq_lens,
+        hidden_states if has_hidden else seq_lens,
+        src_hidden_states if has_hidden else seq_lens,
+        global_num_tokens if has_gathered else seq_lens,
+        global_num_tokens_for_logprob if has_gathered else seq_lens,
+        num_tokens,
+        max_num_token,
+        raw_bs,
+        bs,
+        max_bs,
+        num_tokens_per_bs,
+        num_front_tokens,
+        seq_len_fill_value,
+        hidden_numel,
+        num_global,
+        num_token_programs,
+        HAS_TEMPS=has_temps,
+        HAS_HIDDEN=has_hidden,
+        HAS_GATHERED=has_gathered,
+        BLOCK_TOK=BLOCK_TOK,
+        BLOCK_HIDDEN=BLOCK_HIDDEN,
+        GLOBAL_BLOCK=global_block,
+    )

python/sglang/srt/arg_groups/overrides.py

@@ -824,6 +824,57 @@ def _deepseek_v4_overrides(server_args: Any, hf_config: Any) -> dict:
     return overrides
 
 
+@_register_for(
+    "InklingForConditionalGeneration",
+    "InklingForConditionalGenerationMTP",
+)
+def _inkling_overrides(server_args: Any, hf_config: Any) -> dict:
+    """Inkling architecture defaults: SWA / mamba KV-pool ratios tuned for the
+    hybrid-SWA layout, the extra-buffer mamba strategy, and the unified radix
+    tree (which Inkling requires — models/inkling.py asserts it). The full-graph
+    prefill default is set separately (inline, before cuda-graph resolution) —
+    see ServerArgs.__post_init__ / _apply_inkling_prefill_cuda_graph_default. The
+    server-arg defaults each yield to an explicit user value (compared against
+    the ServerArgs class default); the prefill declaration is materialized
+    before _parse_cuda_graph_config folds cuda_graph_backend_prefill into
+    prefill.backend, and an explicit --cuda-graph-backend-prefill /
+    --disable-prefill-cuda-graph still wins. The unified-radix env write follows
+    the MiniMax-M3 handler precedent (env is not a resolvable server-arg)."""
+    from sglang.srt.server_args import ServerArgs
+
+    overrides: Dict[str, Any] = {}
+    # NOTE: the full-graph prefill default is NOT set here. cuda-graph config is
+    # resolved in __post_init__ before declarations are materialized, so a
+    # cuda_graph_backend_prefill declared here lands too late (the breakable
+    # default would already have been auto-disabled for this multimodal arch).
+    # It is set inline before _handle_cuda_graph_config instead.
+    if server_args.swa_full_tokens_ratio == ServerArgs.swa_full_tokens_ratio:
+        overrides["swa_full_tokens_ratio"] = 0.1
+    if server_args.mamba_full_memory_ratio == ServerArgs.mamba_full_memory_ratio:
+        overrides["mamba_full_memory_ratio"] = 0.1
+    # Inkling requires the extra-buffer mamba strategy (inkling.py asserts
+    # enable_mamba_extra_buffer()); the generic "auto" resolution does not cover
+    # Inkling, so pin it here. Yields to an explicit --mamba-scheduler-strategy.
+    if server_args.mamba_radix_cache_strategy == ServerArgs.mamba_radix_cache_strategy:
+        overrides["mamba_radix_cache_strategy"] = "extra_buffer"
+    # Inkling attention runs only on the fa4 (Blackwell) or triton backends --
+    # models/inkling_common/attn.py asserts attention_backend in {fa4, triton}.
+    # The generic resolver would otherwise pick trtllm_mha (SM100) / fa3
+    # (Hopper), so a bare launch fails on the first attention forward. Pin a
+    # supported default when the user left every attention-backend flag unset
+    # (mirrors the MiniMax-M3 SM100 fa4-default above); an explicit
+    # --attention-backend / --prefill/decode-attention-backend still wins.
+    if server_args.is_attention_backend_not_set():
+        inkling_attn_backend = "fa4" if is_sm100_supported() else "triton"
+        overrides["attention_backend"] = inkling_attn_backend
+        logger.info(
+            f"Use {inkling_attn_backend} as the attention backend for Inkling "
+            "(requires fa4 or triton)."
+        )
+    envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.set(True)
+    return overrides
+
+
 @_register_for("NemotronHForCausalLM", "NemotronHPuzzleForCausalLM")
 def _nemotron_h_overrides(server_args: Any, hf_config: Any) -> dict:
     """NemotronH quantization / MoE runner / attention backend defaults

python/sglang/srt/configs/__init__.py

@@ -9,6 +9,12 @@ from sglang.srt.configs.dots_vlm import DotsVLMConfig
 from sglang.srt.configs.exaone import ExaoneConfig
 from sglang.srt.configs.falcon_h1 import FalconH1Config
 from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig
+from sglang.srt.configs.inkling import (
+    InklingAudioConfig,
+    InklingMMConfig,
+    InklingModelConfig,
+    InklingVisionConfig,
+)
 from sglang.srt.configs.interns2preview import InternS2PreviewConfig
 from sglang.srt.configs.janus_pro import MultiModalityConfig
 from sglang.srt.configs.jet_nemotron import JetNemotronConfig
@@ -86,6 +92,10 @@ __all__ = [
     "MiniMaxM3VLConfig",
     "Step3p7Config",
     "Qwen3ASRConfig",
+    "InklingAudioConfig",
+    "InklingMMConfig",
+    "InklingModelConfig",
+    "InklingVisionConfig",
     "UnlimitedVLConfig",
     "ZayaConfig",
 ]

python/sglang/srt/configs/hybrid_arch.py

@@ -6,6 +6,8 @@ from sglang.srt.configs import (
     BailingHybridConfig,
     FalconH1Config,
     GraniteMoeHybridConfig,
+    InklingMMConfig,
+    InklingModelConfig,
     InternS2PreviewConfig,
     JetNemotronConfig,
     JetVLMConfig,
@@ -76,6 +78,11 @@ def mamba2_config(model_config: ModelConfig):
         | ZayaConfig,
     ):
         return config
+    if isinstance(config, InklingModelConfig):
+        return config if config.mamba2_cache_params is not None else None
+    if isinstance(config, InklingMMConfig):
+        text_config = config.text_config
+        return text_config if text_config.mamba2_cache_params is not None else None
     if isinstance(config, NemotronH_Nano_VL_V2_Config):
         return config.llm_config
 

python/sglang/srt/configs/inkling.py

@@ -0,0 +1,431 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal, Optional
+
+import torch
+from transformers import CONFIG_MAPPING
+from transformers.configuration_utils import PretrainedConfig
+
+from sglang.srt.configs.mamba_utils import BaseLinearStateParams
+
+
+class InklingModelConfig(PretrainedConfig):
+    model_type = "inkling_model"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        *,
+        vocab_size: int = 201024,
+        hidden_size: int = 1536,
+        intermediate_size: int = 768,
+        dense_intermediate_size: Optional[int] = None,
+        num_hidden_layers: int = 16,
+        num_attention_heads: int = 12,
+        num_key_value_heads: int = 4,
+        head_dim: Optional[int] = None,
+        v_head_dim: Optional[int] = None,
+        d_rel: int = 16,
+        rel_extent: int = 1024,
+        local_layer_ids: Optional[list[int]] = None,
+        sliding_window_size: int = 512,
+        swa_num_attention_heads: Optional[int] = None,
+        swa_num_key_value_heads: Optional[int] = None,
+        swa_head_dim: Optional[int] = None,
+        swa_v_head_dim: Optional[int] = None,
+        mtp_local_layer_ids: Optional[list[int]] = None,
+        mtp_local_extent: Optional[int] = None,
+        mtp_swa_num_attention_heads: Optional[int] = None,
+        mtp_swa_num_key_value_heads: Optional[int] = None,
+        mtp_swa_head_dim: Optional[int] = None,
+        rms_norm_eps: float = 1e-6,
+        hidden_act: str = "silu",
+        q_bias: bool = False,
+        o_bias: bool = False,
+        use_embed_norm: bool = False,
+        use_sconv: bool = False,
+        sconv_kernel_size: int = 4,
+        chain_hidden_post_norm: bool = False,
+        dense_mlp_idx: int = 0,
+        n_routed_experts: int = 0,
+        n_shared_experts: int = 0,
+        num_experts_per_tok: int = 1,
+        route_scale: float = 1.0,
+        use_gate_bias: bool = False,
+        use_global_scale: bool = False,
+        norm_after_topk: bool = True,
+        gate_activation: Literal["sigmoid", "softmax"] = "sigmoid",
+        shared_expert_sink: bool = False,
+        shared_experts_size: int = 1,
+        inference_moe_w13_interleaved: bool = True,
+        log_scaling_n_floor: int | None = None,
+        log_scaling_alpha: float = 0.1,
+        unpadded_vocab_size: Optional[int] = None,
+        padded_vocab_size: Optional[int] = None,
+        logits_mup_width_multiplier: Optional[float] = None,
+        final_logit_softcapping: Optional[float] = None,
+        num_nextn_predict_layers: int = 8,
+        tie_word_embeddings: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        if head_dim is None:
+            head_dim = hidden_size // num_attention_heads
+        if v_head_dim is None:
+            v_head_dim = head_dim
+        if swa_num_attention_heads is None:
+            swa_num_attention_heads = num_attention_heads
+        if swa_num_key_value_heads is None:
+            swa_num_key_value_heads = num_key_value_heads
+        if swa_head_dim is None:
+            swa_head_dim = head_dim
+        if swa_v_head_dim is None:
+            swa_v_head_dim = swa_head_dim
+        if dense_intermediate_size is None:
+            dense_intermediate_size = intermediate_size
+        if local_layer_ids is None:
+            local_layer_ids = []
+        # Per-depth banded MTP attention: a depth listed in mtp_local_layer_ids
+        # is a sliding-window block with its own window (mtp_local_extent);
+        # other depths stay full-attention.
+        if mtp_local_layer_ids is None:
+            mtp_local_layer_ids = []
+        if mtp_local_extent is None:
+            mtp_local_extent = sliding_window_size
+        if mtp_swa_num_attention_heads is None:
+            mtp_swa_num_attention_heads = swa_num_attention_heads
+        if mtp_swa_num_key_value_heads is None:
+            mtp_swa_num_key_value_heads = swa_num_key_value_heads
+        if mtp_swa_head_dim is None:
+            mtp_swa_head_dim = swa_head_dim
+        if mtp_local_layer_ids:
+            local_id_set = set(mtp_local_layer_ids)
+            assert len(local_id_set) == len(
+                mtp_local_layer_ids
+            ), f"mtp_local_layer_ids must be unique: {mtp_local_layer_ids}"
+            assert all(0 <= i < num_nextn_predict_layers for i in local_id_set), (
+                f"mtp_local_layer_ids must be in [0, {num_nextn_predict_layers}): "
+                f"{mtp_local_layer_ids}"
+            )
+            # The draft KV pool and the sconv conv-state cache are still sized
+            # from the trunk's swa geometry; a head geometry that differs is not
+            # wired through yet.
+            assert (
+                mtp_swa_num_key_value_heads == swa_num_key_value_heads
+                and mtp_swa_head_dim == swa_head_dim
+            ), (
+                "banded MTP head geometry must match the trunk swa geometry: "
+                f"kv_heads {mtp_swa_num_key_value_heads} vs "
+                f"{swa_num_key_value_heads}, head_dim {mtp_swa_head_dim} vs "
+                f"{swa_head_dim}"
+            )
+
+        if padded_vocab_size is None:
+            padded_vocab_size = vocab_size
+            vocab_size = (
+                unpadded_vocab_size
+                if (
+                    unpadded_vocab_size is not None
+                    and unpadded_vocab_size < padded_vocab_size
+                )
+                else vocab_size
+            )
+
+        self.vocab_size = vocab_size
+        self.padded_vocab_size = padded_vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.dense_intermediate_size = dense_intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.v_head_dim = v_head_dim
+        self.d_rel = d_rel
+        self.rel_extent = rel_extent
+        self.local_layer_ids = local_layer_ids
+        self.sliding_window_size = sliding_window_size
+        self.swa_num_attention_heads = swa_num_attention_heads
+        self.swa_num_key_value_heads = swa_num_key_value_heads
+        self.swa_head_dim = swa_head_dim
+        self.swa_v_head_dim = swa_v_head_dim
+        self.mtp_local_layer_ids = mtp_local_layer_ids
+        self.mtp_local_extent = mtp_local_extent
+        self.mtp_swa_num_attention_heads = mtp_swa_num_attention_heads
+        self.mtp_swa_num_key_value_heads = mtp_swa_num_key_value_heads
+        self.mtp_swa_head_dim = mtp_swa_head_dim
+        self.rms_norm_eps = rms_norm_eps
+        self.hidden_act = hidden_act
+        self.q_bias = q_bias
+        self.o_bias = o_bias
+        self.use_embed_norm = use_embed_norm
+        self.use_sconv = use_sconv
+        self.sconv_kernel_size = sconv_kernel_size
+        self.chain_hidden_post_norm = chain_hidden_post_norm
+        self.dense_mlp_idx = dense_mlp_idx
+        self.n_routed_experts = n_routed_experts
+        self.num_experts = n_routed_experts
+        self.n_shared_experts = n_shared_experts
+        self.num_shared_experts = n_shared_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.route_scale = route_scale
+        self.use_gate_bias = use_gate_bias
+        self.use_global_scale = use_global_scale
+        self.norm_after_topk = norm_after_topk
+        self.gate_activation = gate_activation
+        self.shared_expert_sink = shared_expert_sink
+        self.shared_experts_size = shared_experts_size
+        self.inference_moe_w13_interleaved = inference_moe_w13_interleaved
+        self.log_scaling_n_floor = log_scaling_n_floor
+        self.log_scaling_alpha = log_scaling_alpha
+        self.unpadded_vocab_size = self.vocab_size
+        self.logits_mup_width_multiplier = logits_mup_width_multiplier
+        self.final_logit_softcapping = final_logit_softcapping
+        self.num_nextn_predict_layers = num_nextn_predict_layers
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    @property
+    def conv_layer_ids(self) -> list[int]:
+        return list(range(self.num_hidden_layers))
+
+    @property
+    def linear_layer_ids(self) -> list[int]:
+        return self.conv_layer_ids
+
+    @property
+    def full_attention_layer_ids(self) -> list[int]:
+        return list(range(self.num_hidden_layers))
+
+    @property
+    def mamba_chunk_size(self) -> int:
+        # Floor at 64: mamba_cache_chunk_size = max(mamba_chunk_size, page_size),
+        # and a floor of 1 lets the radix tree adopt another request's KV at
+        # tiny shared prefixes, whose different kernel-rounding perturbs decode logits.
+        return 64
+
+    @property
+    def mamba2_cache_params(self) -> Optional[InklingConvCacheParams]:
+        from sglang.srt.runtime_context import get_parallel
+
+        try:
+            tp_size = get_parallel().attn_tp_size
+        except (AssertionError, RuntimeError):
+            tp_size = 1
+
+        def tp_local_kv_conv_dim(num_kv_heads: int, head_dim: int) -> int:
+            return max(1, num_kv_heads // tp_size) * head_dim
+
+        full_kv_conv_dim = tp_local_kv_conv_dim(self.num_key_value_heads, self.head_dim)
+        local_kv_conv_dim = tp_local_kv_conv_dim(
+            self.swa_num_key_value_heads, self.swa_head_dim
+        )
+        stream_dim = self.hidden_size
+        from sglang.srt.runtime_context import get_server_args
+
+        if get_server_args().enable_scattered_sconv:
+            # Scattered sconv: the attn/mlp output sconvs run on the [T, H/P]
+            # hidden shard, so their conv-state caches shard with them.
+            assert (
+                self.hidden_size % tp_size == 0
+            ), f"hidden_size {self.hidden_size} not divisible by attn tp {tp_size}"
+            stream_dim = self.hidden_size // tp_size
+        conv_len = self.sconv_kernel_size - 1
+        shape = InklingConvStateShape(
+            conv=[
+                (conv_len, full_kv_conv_dim),
+                (conv_len, full_kv_conv_dim),
+                (conv_len, local_kv_conv_dim),
+                (conv_len, local_kv_conv_dim),
+                (conv_len, stream_dim),
+                (conv_len, stream_dim),
+            ],
+            temporal=(0, 0, 0),
+        )
+        dtype = InklingStateDType(conv=torch.bfloat16, temporal=torch.bfloat16)
+        return InklingConvCacheParams(
+            shape=shape, layers=self.conv_layer_ids, dtype=dtype
+        )
+
+
+class InklingAudioConfig(PretrainedConfig):
+    model_type = "inkling_audio_model"
+
+    def __init__(
+        self,
+        *,
+        decoder_dmodel: Optional[int] = None,
+        n_mel_bins: int = 80,
+        mel_vocab_size: int = 16,
+        dmel_min_value: float = -1.5,
+        dmel_max_value: float = 2.0,
+        use_audio_norm: bool = False,
+        audio_mode: Literal["dmel", "flow"] = "dmel",
+        **kwargs: Any,
+    ) -> None:
+        self.decoder_dmodel = decoder_dmodel
+        self.n_mel_bins = n_mel_bins
+        self.mel_vocab_size = mel_vocab_size
+        self.dmel_min_value = dmel_min_value
+        self.dmel_max_value = dmel_max_value
+        self.use_audio_norm = use_audio_norm
+        self.audio_mode = audio_mode
+        super().__init__(**kwargs)
+
+
+class InklingVisionConfig(PretrainedConfig):
+    model_type = "inkling_vision_model"
+
+    def __init__(
+        self,
+        *,
+        vision_encoder_type: Literal["linear", "hmlp"] = "hmlp",
+        decoder_dmodel: Optional[int] = None,
+        patch_size: int = 16,
+        temporal_patch_size: int = 1,
+        n_channels: int = 3,
+        n_layers: int = 1,
+        use_vision_norm: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        self.vision_encoder_type = vision_encoder_type
+        self.decoder_dmodel = decoder_dmodel
+        self.patch_size = patch_size
+        self.temporal_patch_size = temporal_patch_size
+        self.n_channels = n_channels
+        self.n_layers = n_layers
+        self.use_vision_norm = use_vision_norm
+        super().__init__(**kwargs)
+
+
+class InklingMMConfig(PretrainedConfig):
+    model_type = "inkling_mm_model"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    sub_configs = {
+        "text_config": InklingModelConfig,
+        "audio_config": InklingAudioConfig,
+        "vision_config": InklingVisionConfig,
+    }
+
+    def __init__(
+        self,
+        *,
+        text_config: Optional[dict[str, Any] | InklingModelConfig] = None,
+        audio_config: Optional[dict[str, Any] | InklingAudioConfig] = None,
+        vision_config: Optional[dict[str, Any] | InklingVisionConfig] = None,
+        mtp_config: Optional[dict[str, Any]] = None,
+        tie_word_embeddings: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        self.mtp_config = mtp_config
+        self.text_config = (
+            text_config
+            if isinstance(text_config, InklingModelConfig)
+            else InklingModelConfig(**(text_config or {}))
+        )
+        if isinstance(mtp_config, dict) and mtp_config.get("local_layer_ids"):
+            # Banded MTP head: the checkpoint declares its sliding-window draft
+            # depths on mtp_config. Canonicalize onto text_config so every
+            # consumer (hybrid layer-id split, draft pool routing, the MTP block
+            # construction) reads one source of truth.
+            self.text_config.mtp_local_layer_ids = list(mtp_config["local_layer_ids"])
+            if mtp_config.get("local_extent") is not None:
+                self.text_config.mtp_local_extent = mtp_config["local_extent"]
+        self.audio_config = (
+            audio_config
+            if isinstance(audio_config, InklingAudioConfig)
+            else InklingAudioConfig(**(audio_config or {}))
+        )
+        self.vision_config = (
+            vision_config
+            if isinstance(vision_config, InklingVisionConfig)
+            else InklingVisionConfig(**(vision_config or {}))
+        )
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    def get_text_config(self, *args: Any, **kwargs: Any) -> InklingModelConfig:
+        return self.text_config
+
+    @property
+    def vocab_size(self) -> int:
+        return self.text_config.vocab_size
+
+    @property
+    def hidden_size(self) -> int:
+        return self.text_config.hidden_size
+
+    @property
+    def num_hidden_layers(self) -> int:
+        return self.text_config.num_hidden_layers
+
+    @property
+    def num_attention_heads(self) -> int:
+        return self.text_config.num_attention_heads
+
+    @property
+    def num_key_value_heads(self) -> int:
+        return self.text_config.num_key_value_heads
+
+    @property
+    def head_dim(self) -> int:
+        return self.text_config.head_dim
+
+    @property
+    def full_attention_layer_ids(self) -> list[int]:
+        return self.text_config.full_attention_layer_ids
+
+    @property
+    def linear_layer_ids(self) -> list[int]:
+        return self.text_config.linear_layer_ids
+
+    @property
+    def conv_layer_ids(self) -> list[int]:
+        return self.text_config.conv_layer_ids
+
+    @property
+    def mamba_chunk_size(self) -> int:
+        return self.text_config.mamba_chunk_size
+
+    @property
+    def mamba2_cache_params(self) -> Optional[InklingConvCacheParams]:
+        return self.text_config.mamba2_cache_params
+
+
+@dataclass(kw_only=True, frozen=True)
+class InklingConvStateShape:
+    conv: list[tuple[int, int]]
+    temporal: tuple[int, int, int]
+
+    # Conv tuples read (K-1, dim) — the overlapping dedup view would alias
+    # along the dim axis, so the dedup conv-intermediate layout must stay off.
+    disable_conv_window_dedup: bool = True
+
+
+@dataclass(kw_only=True, frozen=True)
+class InklingStateDType:
+    conv: torch.dtype = torch.bfloat16
+    temporal: torch.dtype = torch.bfloat16
+
+
+@dataclass(kw_only=True, frozen=True)
+class InklingConvCacheParams(BaseLinearStateParams):
+    dtype: InklingStateDType = field(default_factory=InklingStateDType)
+    shape: InklingConvStateShape
+
+
+for _model_type, _config_cls in {
+    "inkling_model": InklingModelConfig,
+    "inkling_audio_model": InklingAudioConfig,
+    "inkling_vision_model": InklingVisionConfig,
+    "inkling_mm_model": InklingMMConfig,
+}.items():
+    try:
+        CONFIG_MAPPING.register(_model_type, _config_cls)
+    except Exception:
+        CONFIG_MAPPING._extra_content[_model_type] = _config_cls

python/sglang/srt/configs/mamba_utils.py

@@ -137,6 +137,10 @@ class Mamba2StateShape:
     conv: list[tuple[int, int]]
     temporal: tuple[int, int, int]
 
+    # Conv tuples read (dim, K-1) — the window axis is last, which the
+    # deduplicated conv-intermediate layout requires.
+    disable_conv_window_dedup: bool = False
+
     intermediate_size: int
     conv_dim: int
     ssm_state_size: int
@@ -217,6 +221,10 @@ class KimiLinearStateShape:
     conv: List[tuple[int, int]]
     temporal: tuple[int, int, int]
 
+    # Conv tuples read (K-1, dim) — the overlapping dedup view would alias
+    # along the dim axis, so the dedup conv-intermediate layout must stay off.
+    disable_conv_window_dedup: bool = True
+
     num_heads: int
     head_dim: int
     num_k_heads: int