From 4e843f12165705b036718ecce8cdfe9d4dce1691 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Thu, 26 Feb 2026 10:23:10 +0800 Subject: [PATCH] [DeepSeek-V3.2][JIT-kernel] Support nsa fuse store indexer k cache (#19148) Co-authored-by: luoyuan.luo Co-authored-by: DarkSharpness <76582120+darksharpness@users.noreply.github.com> --- .../csrc/nsa/fused_store_index_cache.cuh | 124 ++++++++++++++++++ .../jit_kernel/fused_store_index_cache.py | 103 +++++++++++++++ python/sglang/jit_kernel/utils.py | 1 + .../srt/layers/attention/nsa/nsa_indexer.py | 100 +++++++++++--- 4 files changed, 307 insertions(+), 21 deletions(-) create mode 100644 python/sglang/jit_kernel/csrc/nsa/fused_store_index_cache.cuh create mode 100644 python/sglang/jit_kernel/fused_store_index_cache.py diff --git a/python/sglang/jit_kernel/csrc/nsa/fused_store_index_cache.cuh b/python/sglang/jit_kernel/csrc/nsa/fused_store_index_cache.cuh new file mode 100644 index 000000000..e649fda57 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/nsa/fused_store_index_cache.cuh @@ -0,0 +1,124 @@ +#include +#include + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include + +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 +__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(); // wait for primary kernel + + // prefetch the index + const auto index = static_cast(indices)[global_wid]; + // always load the value from input (don't store if invalid) + using KeyT2 = packed_t; + using InStorage = AlignedVector; + using OutStorage = AlignedVector; + const auto elems = static_cast(input)[global_tid]; + const auto [x0, x1] = cast(elems[0]); + const auto [y0, y1] = cast(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(value_ptr)[lane_id] = result; + static_cast(scale_ptr)[0] = scale; + + PDLTriggerSecondary(); // launch secondary kernel +} + +template +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; + + 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(); + TensorMatcher({N, 128}) // input + .with_dtype() + .with_device(device_) + .verify(input); + TensorMatcher({-1, -1}) // cache + .with_strides({kPageBytes, 1}) + .with_dtype() + .with_device(device_) + .verify(cache); + TensorMatcher({N}) // indices + .with_dtype() + .with_device(device_) + .verify(indices); + const auto num_tokens = static_cast(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 diff --git a/python/sglang/jit_kernel/fused_store_index_cache.py b/python/sglang/jit_kernel/fused_store_index_cache.py new file mode 100644 index 000000000..dcfbf4585 --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py index d1f87821f..e5de7df01 100644 --- a/python/sglang/jit_kernel/utils.py +++ b/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", } diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 73c1e865e..323f77c67 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/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, - index_k=k_fp8, - index_k_scale=k_scale, + key=key, + act_quant=act_quant, ) # 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: