From b6fceaa7891e3840062f04894968654391503201 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Tue, 30 Jun 2026 14:49:00 +0800 Subject: [PATCH] [DSA] Use cos_sin_cache for DSA indexer fusion (#29613) --- .../csrc/deepseek_v32/indexer_k.cuh | 62 ++++++++----- .../csrc/deepseek_v4/main_norm_rope.cuh | 43 ++++++--- python/sglang/jit_kernel/dsv32/elementwise.py | 10 +-- python/sglang/jit_kernel/dsv4/elementwise.py | 5 +- .../srt/layers/attention/dsa/dsa_indexer.py | 26 ++---- .../jit/test_dsv32_indexer_fusion.py | 87 ++++++++++++++++--- 6 files changed, 160 insertions(+), 73 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh b/python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh index 0cf4ea46d..056f393d0 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh @@ -1,7 +1,7 @@ // DeepSeek-V3.2 only. // // DSA indexer K kernels: single-head LayerNorm (not RMS), ropes the leading -// kRopeDim dims, and fp8-quantizes the un-rotated activations. V3.2 drops the +// kRopeDim dims, and fp8-quantizes the rotated activations. V3.2 drops the // Hadamard incoherence rotation; it is logit-preserving (see main_norm_rope.cuh). // // Independent of the wk + weights_proj GEMM fusion (dsa_indexer.py): `k_input` @@ -33,14 +33,28 @@ constexpr uint32_t kFusedKIndexerNumWarps = kFusedKIndexerBlockSize / device::kW #define K_INDEXER_KERNEL __global__ __launch_bounds__(kFusedKIndexerBlockSize, 16) +template +SGL_DEVICE device::AlignedVector +load_rope_first_cos_sin(const float* __restrict__ cos_sin_cache, int32_t lane_id) { + constexpr int64_t kHalfRopeDim = kRopeDim / 2; + const int32_t pair0 = lane_id * 2; + const int32_t pair1 = pair0 + 1; + device::AlignedVector freq; + freq[0] = cos_sin_cache[pair0]; + freq[1] = cos_sin_cache[kHalfRopeDim + pair0]; + freq[2] = cos_sin_cache[pair1]; + freq[3] = cos_sin_cache[kHalfRopeDim + pair1]; + return freq; +} + // Indexer K: LayerNorm + RoPE -> bf16. struct FusedKIndexerNormRopeParams { - const void* __restrict__ k_input; // (B, 128) DType - void* __restrict__ k_out; // (B, 128) DType - const float* __restrict__ weight; // (128,) fp32 -- LayerNorm gamma - const float* __restrict__ bias; // (128,) fp32 -- LayerNorm beta - const float* __restrict__ freqs_cis; // (max_pos, 64) fp32 - const void* __restrict__ positions; // (B,) PosT + const void* __restrict__ k_input; // (B, 128) DType + void* __restrict__ k_out; // (B, 128) DType + const float* __restrict__ weight; // (128,) fp32 -- LayerNorm gamma + const float* __restrict__ bias; // (128,) fp32 -- LayerNorm beta + const float* __restrict__ cos_sin_cache; // (max_pos, 64) fp32 [cos..., sin...] + const void* __restrict__ positions; // (B,) PosT // Row stride for `k_input` in elements (caller passes the wk slice directly). int64_t k_input_stride_batch; uint32_t batch_size; @@ -71,7 +85,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope(const __grid_constant__ FusedKIn const auto input_ptr = static_cast(params.k_input) + work_id * params.k_input_stride_batch; const auto position = static_cast(static_cast(params.positions)[work_id]); - const auto freqs_cis = params.freqs_cis + position * kRopeDim; + const auto cos_sin_cache = params.cos_sin_cache + position * kRopeDim; PDLWaitPrimary(); Float4 data, freq, gamma, beta; @@ -82,7 +96,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope(const __grid_constant__ FusedKIn input_vec.load(input_ptr, lane_id); gamma.load(params.weight, lane_id); beta.load(params.bias, lane_id); - if (is_rope_lane) freq.load(freqs_cis, lane_id); + if (is_rope_lane) freq = load_rope_first_cos_sin(cos_sin_cache, lane_id); float sum = 0.0f; #pragma unroll @@ -144,7 +158,7 @@ struct FusedKIndexerNormRopeKernel { const tvm::ffi::TensorView k_out, const tvm::ffi::TensorView weight, const tvm::ffi::TensorView bias, - const tvm::ffi::TensorView freqs_cis, + const tvm::ffi::TensorView cos_sin_cache, const tvm::ffi::TensorView positions, double eps) { using namespace host; @@ -176,7 +190,7 @@ struct FusedKIndexerNormRopeKernel { TensorMatcher({-1, kRopeDim}) // .with_dtype() .with_device(device_) - .verify(freqs_cis); + .verify(cos_sin_cache); auto pos_dtype = SymbolicDType{}; TensorMatcher({B}) // .with_dtype(pos_dtype) @@ -191,7 +205,7 @@ struct FusedKIndexerNormRopeKernel { .k_out = k_out.data_ptr(), .weight = static_cast(weight.data_ptr()), .bias = static_cast(bias.data_ptr()), - .freqs_cis = static_cast(freqs_cis.data_ptr()), + .cos_sin_cache = static_cast(cos_sin_cache.data_ptr()), .positions = positions.data_ptr(), .k_input_stride_batch = k_input.stride(0), .batch_size = batch_size, @@ -210,13 +224,13 @@ struct FusedKIndexerNormRopeKernel { // launch. Page layout matches fused_store_index_cache.cuh: each page is // 132*page_size bytes (128*page_size fp8 keys, then 4*page_size fp32 scales). struct FusedKIndexerNormRopeStoreParams { - const void* __restrict__ k_input; // (B, 128) DType - void* __restrict__ cache; // (num_pages, 132*page_size) uint8 - const void* __restrict__ indices; // (B,) int64 -- out_cache_loc - const float* __restrict__ weight; // (128,) fp32 -- LayerNorm gamma - const float* __restrict__ bias; // (128,) fp32 -- LayerNorm beta - const float* __restrict__ freqs_cis; // (max_pos, 64) fp32 - const void* __restrict__ positions; // (B,) PosT + const void* __restrict__ k_input; // (B, 128) DType + void* __restrict__ cache; // (num_pages, 132*page_size) uint8 + const void* __restrict__ indices; // (B,) int64 -- out_cache_loc + const float* __restrict__ weight; // (128,) fp32 -- LayerNorm gamma + const float* __restrict__ bias; // (128,) fp32 -- LayerNorm beta + const float* __restrict__ cos_sin_cache; // (max_pos, 64) fp32 [cos..., sin...] + const void* __restrict__ positions; // (B,) PosT // Row stride for `k_input` (caller passes the non-contiguous wk slice directly). int64_t k_input_stride_batch; uint32_t batch_size; @@ -249,7 +263,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope_store(const __grid_constant__ Fu const auto input_ptr = static_cast(params.k_input) + work_id * params.k_input_stride_batch; const auto position = static_cast(static_cast(params.positions)[work_id]); - const auto freqs_cis = params.freqs_cis + position * kRopeDim; + const auto cos_sin_cache = params.cos_sin_cache + position * kRopeDim; PDLWaitPrimary(); Float4 data, freq, gamma, beta; @@ -260,7 +274,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope_store(const __grid_constant__ Fu input_vec.load(input_ptr, lane_id); gamma.load(params.weight, lane_id); beta.load(params.bias, lane_id); - if (is_rope_lane) freq.load(freqs_cis, lane_id); + if (is_rope_lane) freq = load_rope_first_cos_sin(cos_sin_cache, lane_id); float sum = 0.0f; #pragma unroll @@ -345,7 +359,7 @@ struct FusedKIndexerNormRopeStoreKernel { const tvm::ffi::TensorView indices, const tvm::ffi::TensorView weight, const tvm::ffi::TensorView bias, - const tvm::ffi::TensorView freqs_cis, + const tvm::ffi::TensorView cos_sin_cache, const tvm::ffi::TensorView positions, double eps) { using namespace host; @@ -381,7 +395,7 @@ struct FusedKIndexerNormRopeStoreKernel { TensorMatcher({-1, kRopeDim}) // .with_dtype() .with_device(device_) - .verify(freqs_cis); + .verify(cos_sin_cache); auto pos_dtype = SymbolicDType{}; TensorMatcher({B}) // .with_dtype(pos_dtype) @@ -397,7 +411,7 @@ struct FusedKIndexerNormRopeStoreKernel { .indices = indices.data_ptr(), .weight = static_cast(weight.data_ptr()), .bias = static_cast(bias.data_ptr()), - .freqs_cis = static_cast(freqs_cis.data_ptr()), + .cos_sin_cache = static_cast(cos_sin_cache.data_ptr()), .positions = positions.data_ptr(), .k_input_stride_batch = k_input.stride(0), .batch_size = batch_size, diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh index 430914606..edcb6b0a9 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh @@ -46,6 +46,20 @@ constexpr uint32_t kFusedKNumWarps = kFusedKBlockSize / device::kWarpThreads; #define Q_KERNEL __global__ __launch_bounds__(kFusedQBlockSize, 16) #define K_KERNEL __global__ __launch_bounds__(kFusedKBlockSize, 8) +template +SGL_DEVICE device::AlignedVector +load_rope_first_cos_sin(const float* __restrict__ cos_sin_cache, int32_t lane_id) { + constexpr int64_t kHalfRopeDim = kRopeDim / 2; + const int32_t pair0 = lane_id * 2; + const int32_t pair1 = pair0 + 1; + device::AlignedVector freq; + freq[0] = cos_sin_cache[pair0]; + freq[1] = cos_sin_cache[kHalfRopeDim + pair0]; + freq[2] = cos_sin_cache[pair1]; + freq[3] = cos_sin_cache[kHalfRopeDim + pair1]; + return freq; +} + // ============================================================================ // Q kernel: warp-per-(token, head) rmsnorm-self + RoPE + write to q_out. // ============================================================================ @@ -420,11 +434,14 @@ struct FusedQIndexerRopeHadamardQuantParams { // weights_out[b, h] = weight[b, h] * weight_scale * q_scale[b, h]. // q_scale is computed internally and not exposed -- the only consumer of // it is `weights_out`. - const void* __restrict__ weight; // (B, num_heads) DType - float* __restrict__ weights_out; // (B, num_heads) fp32 (== (B, H, 1) flat) - float weight_scale; // scalar c4_indexer.weight_scale - const float* __restrict__ freqs_cis; // (max_pos, 64) fp32 - const void* __restrict__ positions; // (B,) PosT + const void* __restrict__ weight; // (B, num_heads) DType + float* __restrict__ weights_out; // (B, num_heads) fp32 (== (B, H, 1) flat) + float weight_scale; // scalar c4_indexer.weight_scale + // Template-dependent layout: + // kRopeFirst=false: (max_pos, 64) fp32 interleaved [cos0, sin0, ...] + // kRopeFirst=true : (max_pos, 64) fp32 halves [cos..., sin...] + const float* __restrict__ rope_cache; + const void* __restrict__ positions; // (B,) PosT // Row stride for `weight` (caller passes the non-contiguous wk slice directly). int64_t weight_stride_batch; uint32_t batch_size; @@ -460,7 +477,7 @@ Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQ const uint32_t batch_id = work_id / params.num_heads; const auto input_ptr = static_cast(params.q_input) + work_id * kHeadDim; const auto position = static_cast(static_cast(params.positions)[batch_id]); - const auto freqs_cis = params.freqs_cis + position * kRopeDim; + const auto rope_cache = params.rope_cache + position * kRopeDim; // Lane 0 prefetches the weight scalar for this (token, head) work item. // Weight is (B, num_heads) DType; we need one scalar per warp -- offload @@ -477,7 +494,13 @@ Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQ { Storage input_vec; input_vec.load(input_ptr, lane_id); - if (is_rope_lane) freq.load(freqs_cis, kRopeFirst ? lane_id : (lane_id - (kWarpThreads - kRopeSize))); + if (is_rope_lane) { + if constexpr (kRopeFirst) { + freq = load_rope_first_cos_sin(rope_cache, lane_id); + } else { + freq.load(rope_cache, lane_id - (kWarpThreads - kRopeSize)); + } + } #pragma unroll for (int i = 0; i < kVecSize; ++i) { data[i] = cast(input_vec[i]); @@ -567,7 +590,7 @@ struct FusedQIndexerRopeHadamardQuantKernel { const tvm::ffi::TensorView weight, const tvm::ffi::TensorView weights_out, double weight_scale, - const tvm::ffi::TensorView freqs_cis, + const tvm::ffi::TensorView rope_cache, const tvm::ffi::TensorView positions) { using namespace host; constexpr int64_t kHeadDim = 128; @@ -603,7 +626,7 @@ struct FusedQIndexerRopeHadamardQuantKernel { TensorMatcher({-1, kRopeDim}) // .with_dtype() .with_device(device_) - .verify(freqs_cis); + .verify(rope_cache); auto pos_dtype = SymbolicDType{}; TensorMatcher({B}) // .with_dtype(pos_dtype) @@ -632,7 +655,7 @@ struct FusedQIndexerRopeHadamardQuantKernel { .weight = weight.data_ptr(), .weights_out = static_cast(weights_out.data_ptr()), .weight_scale = static_cast(weight_scale), - .freqs_cis = static_cast(freqs_cis.data_ptr()), + .rope_cache = static_cast(rope_cache.data_ptr()), .positions = positions.data_ptr(), .weight_stride_batch = weight.stride(0), .batch_size = batch_size, diff --git a/python/sglang/jit_kernel/dsv32/elementwise.py b/python/sglang/jit_kernel/dsv32/elementwise.py index 39f48b357..4c07afe4b 100644 --- a/python/sglang/jit_kernel/dsv32/elementwise.py +++ b/python/sglang/jit_kernel/dsv32/elementwise.py @@ -43,11 +43,10 @@ def fused_k_indexer_norm_rope( weight: torch.Tensor, bias: torch.Tensor, eps: float, - freqs_cis: torch.Tensor, + cos_sin_cache: torch.Tensor, positions: torch.Tensor, ) -> torch.Tensor: """V3.2 indexer K: LayerNorm + RoPE on leading dims -> bf16. CUDA only.""" - freqs_real = torch.view_as_real(freqs_cis).flatten(-2) # k_input may be a non-contiguous wk slice; output is always contiguous. k_out = torch.empty(k_input.shape, dtype=k_input.dtype, device=k_input.device) module = _jit_k_indexer_norm_rope_module(k_input.dtype) @@ -56,7 +55,7 @@ def fused_k_indexer_norm_rope( k_out, weight, bias, - freqs_real, + cos_sin_cache, positions, float(eps), ) @@ -70,13 +69,12 @@ def fused_k_indexer_norm_rope_store( weight: torch.Tensor, bias: torch.Tensor, eps: float, - freqs_cis: torch.Tensor, + cos_sin_cache: torch.Tensor, positions: torch.Tensor, page_size: int, ) -> None: """V3.2 indexer K + fused store: LayerNorm + RoPE on leading dims + fp8 act-quant + paged index-k cache write, in one launch. CUDA only.""" - freqs_real = torch.view_as_real(freqs_cis).flatten(-2) if not out_cache_loc.is_contiguous(): out_cache_loc = out_cache_loc.contiguous() module = _jit_k_indexer_norm_rope_store_module(k_input.dtype, page_size) @@ -86,7 +84,7 @@ def fused_k_indexer_norm_rope_store( out_cache_loc, weight, bias, - freqs_real, + cos_sin_cache, positions, float(eps), ) diff --git a/python/sglang/jit_kernel/dsv4/elementwise.py b/python/sglang/jit_kernel/dsv4/elementwise.py index 96bc8c1a2..e8c743595 100644 --- a/python/sglang/jit_kernel/dsv4/elementwise.py +++ b/python/sglang/jit_kernel/dsv4/elementwise.py @@ -188,11 +188,10 @@ def fused_q_indexer_rope_first_quant( q_input: torch.Tensor, weight: torch.Tensor, weight_scale: float, - freqs_cis: torch.Tensor, + cos_sin_cache: torch.Tensor, positions: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """DeepSeek-V3.2 only. Indexer Q: RoPE on the leading dims + fp8 act-quant. CUDA only.""" - freqs_real = torch.view_as_real(freqs_cis).flatten(-2) q_fp8 = torch.empty(q_input.shape, dtype=torch.float8_e4m3fn, device=q_input.device) weights_out = torch.empty( (*q_input.shape[:-1], 1), dtype=torch.float32, device=q_input.device @@ -204,7 +203,7 @@ def fused_q_indexer_rope_first_quant( weight, weights_out, float(weight_scale), - freqs_real, + cos_sin_cache, positions, ) return q_fp8, weights_out diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index b03deeb65..1fb37764c 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -327,16 +327,6 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor: return hadamard_transform(x, scale=hidden_size**-0.5) -def _shared_indexer_freqs_cis(rotary_emb: torch.nn.Module) -> torch.Tensor: - cached = getattr(rotary_emb, "_dsa_indexer_freqs_cis", None) - if cached is None: - c = rotary_emb.cos_sin_cache.to(torch.float32) - half = c.shape[-1] // 2 - cached = torch.complex(c[:, :half].contiguous(), c[:, half:].contiguous()) - rotary_emb._dsa_indexer_freqs_cis = cached - return cached - - class Indexer(MultiPlatformOp): _MQA_LOGITS_BYTES_PER_ELEM = 4 _MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000 @@ -437,10 +427,6 @@ class Indexer(MultiPlatformOp): self.scale_fmt = scale_fmt self.softmax_scale = self.head_dim**-0.5 - self._indexer_freqs_cis: Optional[torch.Tensor] = None - if _use_dsa_indexer_fusion: - self._indexer_freqs_cis = _shared_indexer_freqs_cis(self.rotary_emb) - @contextlib.contextmanager def _with_real_sm_count(self): # When pipeline parallelism is enabled, each PP rank initiates a recv operation after the _pp_launch_batch @@ -456,6 +442,10 @@ class Indexer(MultiPlatformOp): else: yield + @property + def _indexer_cos_sin_cache(self) -> torch.Tensor: + return self.rotary_emb.cos_sin_cache + def _weights_proj_bf16_in_fp32_out( self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]] ) -> torch.Tensor: @@ -662,7 +652,7 @@ class Indexer(MultiPlatformOp): self.k_norm.weight, self.k_norm.bias, self.k_norm.variance_epsilon, - self._indexer_freqs_cis, + self._indexer_cos_sin_cache, positions, page_size, ) @@ -674,7 +664,7 @@ class Indexer(MultiPlatformOp): self.k_norm.weight, self.k_norm.bias, self.k_norm.variance_epsilon, - self._indexer_freqs_cis, + self._indexer_cos_sin_cache, positions, ) self._store_index_k_cache( @@ -726,7 +716,7 @@ class Indexer(MultiPlatformOp): q.contiguous(), weights_raw, q_scale_gate, - self._indexer_freqs_cis, + self._indexer_cos_sin_cache, positions, ) @@ -763,7 +753,7 @@ class Indexer(MultiPlatformOp): q.contiguous(), weights_raw, q_scale_gate, - self._indexer_freqs_cis, + self._indexer_cos_sin_cache, positions, ) diff --git a/test/registered/jit/test_dsv32_indexer_fusion.py b/test/registered/jit/test_dsv32_indexer_fusion.py index 6c068b9ba..b01dce1a2 100644 --- a/test/registered/jit/test_dsv32_indexer_fusion.py +++ b/test/registered/jit/test_dsv32_indexer_fusion.py @@ -56,9 +56,9 @@ def _make_inputs(B, seed=0, pos_dtype=torch.int32): dev = "cuda" cos = torch.randn(MAX_POS, HALF, device=dev, generator=g) sin = torch.randn(MAX_POS, HALF, device=dev, generator=g) - freqs_cis = torch.complex(cos, sin) + cos_sin_cache = torch.cat((cos, sin), dim=-1) positions = torch.randint(0, 4096, (B,), device=dev, dtype=pos_dtype, generator=g) - return cos, sin, freqs_cis, positions + return cos, sin, cos_sin_cache, positions def _rope_first(x, cos_p, sin_p): @@ -78,12 +78,12 @@ def test_k_norm_rope_matches_reference(): _skip_if_unavailable() dev = "cuda" B = 37 - cos, sin, freqs_cis, positions = _make_inputs(B) + cos, sin, cos_sin_cache, positions = _make_inputs(B) key = torch.randn(B, HEAD_DIM, dtype=torch.bfloat16, device=dev) weight = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev) bias = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev) - out = fused_k_indexer_norm_rope(key, weight, bias, EPS, freqs_cis, positions) + out = fused_k_indexer_norm_rope(key, weight, bias, EPS, cos_sin_cache, positions) torch.cuda.synchronize() normed = torch.nn.functional.layer_norm( @@ -106,7 +106,7 @@ def test_k_store_matches_unfused(strided): pytest.skip("fused store JIT unavailable") dev = "cuda" B, n_heads = 41, 64 - cos, sin, freqs_cis, positions = _make_inputs(B) + cos, sin, cos_sin_cache, positions = _make_inputs(B) weight = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev) bias = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev) @@ -125,11 +125,13 @@ def test_k_store_matches_unfused(strided): ) buf_fused = torch.zeros_like(buf_ref) - key_bf16 = fused_k_indexer_norm_rope(key, weight, bias, EPS, freqs_cis, positions) + key_bf16 = fused_k_indexer_norm_rope( + key, weight, bias, EPS, cos_sin_cache, positions + ) fused_store_index_k_cache(key_bf16, buf_ref, loc, PAGE_SIZE) fused_k_indexer_norm_rope_store( - key, buf_fused, loc, weight, bias, EPS, freqs_cis, positions, PAGE_SIZE + key, buf_fused, loc, weight, bias, EPS, cos_sin_cache, positions, PAGE_SIZE ) torch.cuda.synchronize() @@ -144,13 +146,13 @@ def test_q_rope_quant_matches_reference(pos_dtype): _skip_if_unavailable() dev = "cuda" B, n_heads = 37, 64 - cos, sin, freqs_cis, positions = _make_inputs(B, pos_dtype=pos_dtype) + cos, sin, cos_sin_cache, positions = _make_inputs(B, pos_dtype=pos_dtype) q = torch.randn(B, n_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev) weight = torch.randn(B, n_heads, dtype=torch.bfloat16, device=dev) weight_scale = 0.137 q_fp8, weights_out = fused_q_indexer_rope_first_quant( - q, weight, weight_scale, freqs_cis, positions + q, weight, weight_scale, cos_sin_cache, positions ) torch.cuda.synchronize() @@ -181,7 +183,7 @@ def test_q_strided_weight_matches_contiguous(): _skip_if_unavailable() dev = "cuda" B, n_heads = 29, 64 - cos, sin, freqs_cis, positions = _make_inputs(B) + cos, sin, cos_sin_cache, positions = _make_inputs(B) q = torch.randn(B, n_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev) # weights_raw = kw[:, head_dim:] is a non-contiguous slice. kw = torch.randn(B, HEAD_DIM + n_heads, dtype=torch.bfloat16, device=dev) @@ -190,16 +192,77 @@ def test_q_strided_weight_matches_contiguous(): assert not w_strided.is_contiguous() a_fp8, a_w = fused_q_indexer_rope_first_quant( - q, w_strided, 0.137, freqs_cis, positions + q, w_strided, 0.137, cos_sin_cache, positions ) b_fp8, b_w = fused_q_indexer_rope_first_quant( - q, w_contig, 0.137, freqs_cis, positions + q, w_contig, 0.137, cos_sin_cache, positions ) torch.cuda.synchronize() assert torch.equal(a_fp8, b_fp8) assert torch.equal(a_w, b_w) +def test_indexer_uses_replaced_rope_cache_for_fused_kernels(): + _skip_if_unavailable() + from sglang.srt.layers.attention.dsa.dsa_indexer import Indexer + + dev = "cuda" + B, n_heads = 7, 64 + old_len = 16 + new_len = 128 + g = torch.Generator(device=dev).manual_seed(123) + old_cache = torch.randn(old_len, ROPE_DIM, dtype=torch.float32, device=dev) + cos = torch.randn(new_len, HALF, dtype=torch.float32, device=dev, generator=g) + sin = torch.randn(new_len, HALF, dtype=torch.float32, device=dev, generator=g) + grown_cache = torch.cat((cos, sin), dim=-1) + positions = torch.arange(old_len, old_len + B, device=dev, dtype=torch.int32) + + class DummyRotary: + pass + + rotary_emb = DummyRotary() + rotary_emb.cos_sin_cache = old_cache + indexer = Indexer.__new__(Indexer) + indexer.rotary_emb = rotary_emb + assert indexer._indexer_cos_sin_cache.data_ptr() == old_cache.data_ptr() + + rotary_emb.cos_sin_cache = grown_cache + assert indexer._indexer_cos_sin_cache.data_ptr() == grown_cache.data_ptr() + + key = torch.randn(B, HEAD_DIM, dtype=torch.bfloat16, device=dev, generator=g) + k_weight = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev, generator=g) + k_bias = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev, generator=g) + k_out = fused_k_indexer_norm_rope( + key, k_weight, k_bias, EPS, indexer._indexer_cos_sin_cache, positions + ) + + normed = torch.nn.functional.layer_norm( + key.float(), (HEAD_DIM,), weight=k_weight, bias=k_bias, eps=EPS + ) + k_ref = _rope_first(normed, cos[positions.long()], sin[positions.long()]) + torch.testing.assert_close(k_out.float(), k_ref, atol=0.06, rtol=0.0) + + q = torch.randn(B, n_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev, generator=g) + q_weight = torch.randn(B, n_heads, dtype=torch.bfloat16, device=dev, generator=g) + weight_scale = 0.137 + q_fp8, weights_out = fused_q_indexer_rope_first_quant( + q, q_weight, weight_scale, indexer._indexer_cos_sin_cache, positions + ) + torch.cuda.synchronize() + + q_ref = _rope_first( + q.float(), cos[positions.long()][:, None, :], sin[positions.long()][:, None, :] + ) + scale = torch.clamp(q_ref.abs().amax(dim=-1, keepdim=True), min=1e-4) / FP8_MAX + torch.testing.assert_close( + weights_out.squeeze(-1), + q_weight.float() * weight_scale * scale.squeeze(-1), + atol=1e-3, + rtol=1e-3, + ) + assert ((q_fp8.float() * scale - q_ref).abs() <= 0.0625 * q_ref.abs() + scale).all() + + if __name__ == "__main__": import sys