[DSA] Use cos_sin_cache for DSA indexer fusion (#29613)

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,13 +33,27 @@ constexpr uint32_t kFusedKIndexerNumWarps = kFusedKIndexerBlockSize / device::kW
 
 #define K_INDEXER_KERNEL __global__ __launch_bounds__(kFusedKIndexerBlockSize, 16)
 
+template <int64_t kRopeDim>
+SGL_DEVICE device::AlignedVector<float, 4>
+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<float, 4> 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 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;
@@ -71,7 +85,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope(const __grid_constant__ FusedKIn
 
   const auto input_ptr = static_cast<const DType*>(params.k_input) + work_id * params.k_input_stride_batch;
   const auto position = static_cast<int32_t>(static_cast<const PosT*>(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<kUsePDL>();
   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<kRopeDim>(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<float>()
         .with_device(device_)
-        .verify(freqs_cis);
+        .verify(cos_sin_cache);
     auto pos_dtype = SymbolicDType{};
     TensorMatcher({B})  //
         .with_dtype<int32_t, int64_t>(pos_dtype)
@@ -191,7 +205,7 @@ struct FusedKIndexerNormRopeKernel {
         .k_out = k_out.data_ptr(),
         .weight = static_cast<const float*>(weight.data_ptr()),
         .bias = static_cast<const float*>(bias.data_ptr()),
-        .freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
+        .cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
         .positions = positions.data_ptr(),
         .k_input_stride_batch = k_input.stride(0),
         .batch_size = batch_size,
@@ -215,7 +229,7 @@ struct FusedKIndexerNormRopeStoreParams {
   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 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;
@@ -249,7 +263,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope_store(const __grid_constant__ Fu
 
   const auto input_ptr = static_cast<const DType*>(params.k_input) + work_id * params.k_input_stride_batch;
   const auto position = static_cast<int32_t>(static_cast<const PosT*>(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<kUsePDL>();
   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<kRopeDim>(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<float>()
         .with_device(device_)
-        .verify(freqs_cis);
+        .verify(cos_sin_cache);
     auto pos_dtype = SymbolicDType{};
     TensorMatcher({B})  //
         .with_dtype<int32_t, int64_t>(pos_dtype)
@@ -397,7 +411,7 @@ struct FusedKIndexerNormRopeStoreKernel {
         .indices = indices.data_ptr(),
         .weight = static_cast<const float*>(weight.data_ptr()),
         .bias = static_cast<const float*>(bias.data_ptr()),
-        .freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
+        .cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
         .positions = positions.data_ptr(),
         .k_input_stride_batch = k_input.stride(0),
         .batch_size = batch_size,

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 <int64_t kRopeDim>
+SGL_DEVICE device::AlignedVector<float, 4>
+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<float, 4> 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.
 // ============================================================================
@@ -423,7 +437,10 @@ struct FusedQIndexerRopeHadamardQuantParams {
   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
+  // 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;
@@ -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<const DType*>(params.q_input) + work_id * kHeadDim;
   const auto position = static_cast<int32_t>(static_cast<const PosT*>(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<kRopeDim>(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<float>(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<float>()
         .with_device(device_)
-        .verify(freqs_cis);
+        .verify(rope_cache);
     auto pos_dtype = SymbolicDType{};
     TensorMatcher({B})  //
         .with_dtype<int32_t, int64_t>(pos_dtype)
@@ -632,7 +655,7 @@ struct FusedQIndexerRopeHadamardQuantKernel {
         .weight = weight.data_ptr(),
         .weights_out = static_cast<float*>(weights_out.data_ptr()),
         .weight_scale = static_cast<float>(weight_scale),
-        .freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
+        .rope_cache = static_cast<const float*>(rope_cache.data_ptr()),
         .positions = positions.data_ptr(),
         .weight_stride_batch = weight.stride(0),
         .batch_size = batch_size,

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),
     )

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

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,
         )
 

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