[DeepSeek-V3.2][JIT-kernel] Support nsa fuse store indexer k cache (#19148)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Co-authored-by: DarkSharpness <76582120+darksharpness@users.noreply.github.com>

python/sglang/jit_kernel/csrc/nsa/fused_store_index_cache.cuh

@@ -0,0 +1,124 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <bit>
+#include <cstdint>
+#include <cuda_fp8.h>
+
+namespace {
+
+struct FusedStoreCacheParam {
+  const void* __restrict__ input;
+  void* __restrict__ cache;
+  const void* __restrict__ indices;
+  uint32_t num_tokens;
+};
+
+[[maybe_unused]]
+SGL_DEVICE float fp8_e4m3_clip(float val) {
+  namespace math = device::math;
+  return math::max(math::min(val, math::FP8_E4M3_MAX), -math::FP8_E4M3_MAX);
+}
+
+[[maybe_unused]]
+SGL_DEVICE fp8x2_e4m3_t pack_fp8(float x, float y) {
+  return fp8x2_e4m3_t{fp32x2_t{fp8_e4m3_clip(x), fp8_e4m3_clip(y)}};
+}
+
+template <typename KeyT, typename IndicesT, uint32_t kPageBits, bool kUsePDL>
+__global__ void fused_store_indexer_cache(const __grid_constant__ FusedStoreCacheParam param) {
+  using namespace device;
+
+  /// NOTE: 132 = 128 + 4
+  constexpr int64_t kPageBytes = 132 << kPageBits;
+
+  // each warp handles 128 elements, each block handles multiple rows
+  const auto& [input, cache, indices, num_tokens] = param;
+  const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto global_wid = global_tid / 32;
+  const auto lane_id = threadIdx.x % 32;
+
+  if (global_wid >= num_tokens) return;
+
+  PDLWaitPrimary<kUsePDL>();  // wait for primary kernel
+
+  // prefetch the index
+  const auto index = static_cast<const IndicesT*>(indices)[global_wid];
+  // always load the value from input (don't store if invalid)
+  using KeyT2 = packed_t<KeyT>;
+  using InStorage = AlignedVector<KeyT2, 2>;
+  using OutStorage = AlignedVector<fp8x2_e4m3_t, 2>;
+  const auto elems = static_cast<const InStorage*>(input)[global_tid];
+  const auto [x0, x1] = cast<fp32x2_t>(elems[0]);
+  const auto [y0, y1] = cast<fp32x2_t>(elems[1]);
+  const auto local_max = fmaxf(fmaxf(fabs(x0), fabs(x1)), fmaxf(fabs(y0), fabs(y1)));
+  const auto abs_max = warp::reduce_max(local_max);
+  // use normal fp32 scale
+  const auto scale = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
+  const auto inv_scale = 1.0f / scale;
+  const int32_t page = index >> kPageBits;
+  const int32_t offset = index & ((1 << kPageBits) - 1);
+  const auto page_ptr = pointer::offset(cache, page * kPageBytes);
+  const auto value_ptr = pointer::offset(page_ptr, offset * 128);
+  const auto scale_ptr = pointer::offset(page_ptr, 128 << kPageBits, offset * 4);
+  OutStorage result;
+  result[0] = pack_fp8(x0 * inv_scale, x1 * inv_scale);
+  result[1] = pack_fp8(y0 * inv_scale, y1 * inv_scale);
+  static_cast<OutStorage*>(value_ptr)[lane_id] = result;
+  static_cast<float*>(scale_ptr)[0] = scale;
+
+  PDLTriggerSecondary<kUsePDL>();  // launch secondary kernel
+}
+
+template <typename KeyT, typename IndicesT, uint32_t kPageSize, bool kUsePDL>
+struct FusedStoreCacheIndexerKernel {
+  static constexpr int32_t kLogSize = std::countr_zero(kPageSize);
+  /// NOTE: 132 = 128 + 4 (128 represent K and 4 represent scale)
+  static constexpr int64_t kPageBytes = 132 * kPageSize;
+  static constexpr auto kernel = fused_store_indexer_cache<KeyT, IndicesT, kLogSize, kUsePDL>;
+
+  static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
+  static_assert(1 << kLogSize == kPageSize);
+
+  static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView cache, tvm::ffi::TensorView indices) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+    TensorMatcher({N, 128})  // input
+        .with_dtype<KeyT>()
+        .with_device(device_)
+        .verify(input);
+    TensorMatcher({-1, -1})  // cache
+        .with_strides({kPageBytes, 1})
+        .with_dtype<uint8_t>()
+        .with_device(device_)
+        .verify(cache);
+    TensorMatcher({N})  // indices
+        .with_dtype<IndicesT>()
+        .with_device(device_)
+        .verify(indices);
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto params = FusedStoreCacheParam{
+        .input = input.data_ptr(),
+        .cache = cache.data_ptr(),
+        .indices = indices.data_ptr(),
+        .num_tokens = num_tokens,
+    };
+    const auto kBlockSize = 128;
+    const auto num_blocks = div_ceil(num_tokens * 32, kBlockSize);
+    LaunchKernel(num_blocks, kBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/fused_store_index_cache.py

@@ -0,0 +1,103 @@
+"""
+This module provides JIT-compiled CUDA kernels for fusing multiple tensor
+copy operations into single kernel launches, reducing kernel launch overhead
+and improving CUDA graph replay performance.
+
+The kernels are compiled on-demand using TVM FFI and cached for subsequent use.
+"""
+
+from __future__ import annotations
+
+import logging
+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
+
+logger = logging.getLogger(__name__)
+
+
+@cache_once
+def _jit_nsa_fused_store_module(
+    key_dtype: torch.dtype, indices_dtype: torch.dtype, page_size: int
+) -> Module:
+    """
+    Build a JIT module that exposes:
+      module.fused_store_index_k_cache(input_bf16, index_k_with_scale_u8, loc_i64)
+    """
+    args = make_cpp_args(key_dtype, indices_dtype, page_size, is_arch_support_pdl())
+    return load_jit(
+        "fused_store_index_k_cache",
+        *args,
+        cuda_files=["nsa/fused_store_index_cache.cuh"],
+        cuda_wrappers=[
+            (
+                "fused_store_index_k_cache",
+                # - Float  = bf16_t (sgl_kernel/type.cuh)
+                # - IndicesT = int64_t (out_cache_loc is int64 in SGLang SetKAndS)
+                # - kPageSize = 64 (CUDA NSA)
+                f"FusedStoreCacheIndexerKernel<{args}>::run",
+            )
+        ],
+    )
+
+
+@cache_once
+def can_use_nsa_fused_store(
+    key_dtype: torch.dtype, indices_dtype: torch.dtype, page_size: int
+) -> bool:
+    logger = logging.getLogger(__name__)
+    try:
+        _jit_nsa_fused_store_module(key_dtype, indices_dtype, page_size)
+        return True
+    except Exception as e:
+        logger.warning(f"Failed to load nsa fused store JIT kernel: {e}")
+        return False
+
+
+def fused_store_index_k_cache(
+    key: torch.Tensor,
+    index_k_with_scale: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    page_size: int = 64,
+) -> None:
+    """
+    Fused: quantize bf16 key (N,128) -> fp8 + fp32 scale and write into NSATokenToKVPool.index_k_with_scale_buffer.
+
+    key:            (num_tokens, 128) bf16 (or reshapeable to it)
+    index_k_with_scale:  (num_pages, 64*(128+4)) uint8
+    out_cache_loc:       (num_tokens,) int64 token indices in TokenToKVPool
+    """
+    assert key.is_cuda
+    assert index_k_with_scale.is_cuda
+    assert out_cache_loc.is_cuda
+
+    # 1) normalize shapes
+    if key.dim() != 2:
+        key = key.view(-1, key.shape[-1])
+    assert key.shape[1] == 128, f"expected key last-dim=128, got {key.shape}"
+
+    # 2) dtypes
+    assert key.dtype == torch.bfloat16, f"{key.dtype=}"
+    assert index_k_with_scale.dtype == torch.uint8, f"{index_k_with_scale.dtype=}"
+    assert out_cache_loc.dtype == torch.int64, f"{out_cache_loc.dtype=}"
+
+    # 3) contiguity
+    if not key.is_contiguous():
+        key = key.contiguous()
+    if not out_cache_loc.is_contiguous():
+        out_cache_loc = out_cache_loc.contiguous()
+    if not index_k_with_scale.is_contiguous():
+        index_k_with_scale = index_k_with_scale.contiguous()
+
+    module = _jit_nsa_fused_store_module(key.dtype, out_cache_loc.dtype, page_size)
+    module.fused_store_index_k_cache(key, index_k_with_scale, out_cache_loc)

python/sglang/jit_kernel/utils.py

@@ -78,6 +78,7 @@ CPP_DTYPE_MAP = {
     torch.float8_e4m3fn: "fp8_e4m3_t",
     torch.bfloat16: "bf16_t",
     torch.int8: "int8_t",
+    torch.int64: "int64_t",
 }
 
 

python/sglang/srt/layers/attention/nsa/nsa_indexer.py

@@ -7,6 +7,10 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 import torch
 from einops import rearrange
 
+from sglang.jit_kernel.fused_store_index_cache import (
+    can_use_nsa_fused_store,
+    fused_store_index_k_cache,
+)
 from sglang.srt.environ import envs
 from sglang.srt.layers.layernorm import LayerNorm
 from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
@@ -670,15 +674,15 @@ class Indexer(MultiPlatformOp):
 
         # Fast path: only compute and store k cache, skip all q and weights ops
         key = self._get_k_bf16(x, positions, enable_dual_stream)
-        k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt)
 
         if not forward_batch.out_cache_loc.is_contiguous():
             forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous()
-        forward_batch.token_to_kv_pool.set_index_k_scale_buffer(
+
+        self._store_index_k_cache(
+            forward_batch=forward_batch,
             layer_id=layer_id,
-            loc=forward_batch.out_cache_loc,
+            key=key,
-            index_k=k_fp8,
+            act_quant=act_quant,
-            index_k_scale=k_scale,
         )
 
         # MHA doesn't need topk_indices
@@ -928,6 +932,58 @@ class Indexer(MultiPlatformOp):
         topk_indices = torch.cat(topk_indices_list, dim=0)
         return topk_indices
 
+    def _store_index_k_cache(
+        self,
+        forward_batch: ForwardBatch,
+        layer_id: int,
+        key: torch.Tensor,
+        *,
+        act_quant=None,  # fallback only
+    ) -> None:
+        """
+        Store NSA indexer K cache for current step.
+
+        Preferred: fused_store_index_k_cache(key, cache, out_cache_loc, page_size)
+        Fallback : act_quant(key) + token_to_kv_pool.set_index_k_scale_buffer(...)
+        """
+
+        # Fast path: JIT fused store (CUDA, page_size=64, non-fnuz)
+        if (
+            _is_cuda
+            and (not _is_fp8_fnuz)
+            and can_use_nsa_fused_store(
+                key.dtype,
+                forward_batch.out_cache_loc.dtype,
+                forward_batch.token_to_kv_pool.page_size,
+            )
+        ):
+            # NOTE: wrapper already normalizes shape/contiguity and asserts dtypes.
+            buf = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer(
+                layer_id=layer_id
+            )
+            fused_store_index_k_cache(
+                key,
+                buf,
+                forward_batch.out_cache_loc,
+                forward_batch.token_to_kv_pool.page_size,
+            )
+            return
+
+        # Fallback: original path
+        assert act_quant is not None
+        k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt)
+
+        out_loc = forward_batch.out_cache_loc
+        if not out_loc.is_contiguous():
+            out_loc = out_loc.contiguous()
+
+        forward_batch.token_to_kv_pool.set_index_k_scale_buffer(
+            layer_id=layer_id,
+            loc=out_loc,
+            index_k=k_fp8,
+            index_k_scale=k_scale,
+        )
+
     def forward_cuda(
         self,
         x: torch.Tensor,
@@ -994,7 +1050,12 @@ class Indexer(MultiPlatformOp):
             )
             q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
             with torch.cuda.stream(self.alt_stream):
-                k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt)
+                self._store_index_k_cache(
+                    forward_batch=forward_batch,
+                    layer_id=layer_id,
+                    key=key,
+                    act_quant=act_quant,
+                )
             current_stream.wait_stream(self.alt_stream)
             weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
         else:
@@ -1008,11 +1069,21 @@ class Indexer(MultiPlatformOp):
 
                 q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
                 with torch.cuda.stream(self.alt_stream):
-                    k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt)
+                    self._store_index_k_cache(
+                        forward_batch=forward_batch,
+                        layer_id=layer_id,
+                        key=key,
+                        act_quant=act_quant,
+                    )
                 current_stream.wait_stream(self.alt_stream)
             else:
                 q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
-                k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt)
+                self._store_index_k_cache(
+                    forward_batch=forward_batch,
+                    layer_id=layer_id,
+                    key=key,
+                    act_quant=act_quant,
+                )
 
             # `_get_logits_head_gate` expects a Tensor. For tuple activations, dequantize
             # to a float tensor here (callsite), keeping `_get_logits_head_gate` backend-agnostic.
@@ -1048,19 +1119,6 @@ class Indexer(MultiPlatformOp):
 
             weights = self._get_logits_head_gate(x_for_gate, q_scale)
 
-        # k_fp8: (seq_len, head_dim) fp8_e4m3fn
-        # k_buffer: (num_total_tokens + page_size, head_dim) fp8_e4m3fn
-        # k_scale: (seq_len, head_dim // block_size = 1) fp8_e4m3fn
-        # k_scale_cache: (num_total_tokens + page_size, head_dim // block_size = 1) fp8_e4m3fn
-        if not forward_batch.out_cache_loc.is_contiguous():
-            forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous()
-        forward_batch.token_to_kv_pool.set_index_k_scale_buffer(
-            layer_id=layer_id,
-            loc=forward_batch.out_cache_loc,
-            index_k=k_fp8,
-            index_k_scale=k_scale,
-        )
-
         if _is_cuda or _is_hip:
             assert forward_batch.seq_lens_cpu is not None
             if len(forward_batch.seq_lens_cpu) == 0: