Support NoPE layers in the tokenspeed_mla FP8 prefill hook (#38152)

Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>

python/sglang/srt/layers/attention/tokenspeed_mla_backend.py

@@ -180,13 +180,13 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
                         enable_ex2_emulation=enable_ex2_emulation,
                     )
 
+    @staticmethod
     def _fused_rope_fp8_quantize(
-        self,
         q_nope: torch.Tensor,
         q_pe: torch.Tensor,
         k_nope: torch.Tensor,
         k_pe: torch.Tensor,
-        cos_sin_cache: torch.Tensor,
+        cos_sin_cache: Optional[torch.Tensor],
         positions: torch.Tensor,
         is_neox: bool,
         qk_nope_head_dim: int,
@@ -194,6 +194,9 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Fused RoPE + FP8 quantize that also packs nope+pe along the last
         dim, so FMHA consumes contig FP8 Q/K without an extra concat or cast.
+
+        ``cos_sin_cache`` is None for NoPE layers (``skip_rope``); they keep the
+        FP8 quantize and the packed layout, only the rotation is dropped.
         """
         num_heads = q_nope.shape[1]
         seq_len = q_nope.shape[0]
@@ -218,6 +221,14 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
         else:
             k_pe_expanded = k_pe
 
+        if cos_sin_cache is None:
+            # NoPE layer: quantize straight into the packed buffers.
+            q_fp8[..., :qk_nope_head_dim].copy_(q_nope)
+            q_fp8[..., qk_nope_head_dim:].copy_(q_pe)
+            k_fp8[..., :qk_nope_head_dim].copy_(k_nope)
+            k_fp8[..., qk_nope_head_dim:].copy_(k_pe_expanded)
+            return q_fp8, k_fp8
+
         _flashinfer_rope.mla_rope_quantize_fp8(
             q_rope=q_pe,
             k_rope=k_pe_expanded,
@@ -257,14 +268,15 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
         v_bf16 = kv[..., layer.qk_nope_head_dim :]
         q_nope = q[..., : layer.qk_nope_head_dim]
 
+        rotary_emb = layer.rotary_emb
         q_fp8, k_fp8 = self._fused_rope_fp8_quantize(
             q_nope=q_nope,
             q_pe=q_pe,
             k_nope=k_nope,
             k_pe=k_pe,
-            cos_sin_cache=layer.rotary_emb.cos_sin_cache,
+            cos_sin_cache=None if rotary_emb is None else rotary_emb.cos_sin_cache,
             positions=positions,
-            is_neox=getattr(layer.rotary_emb, "is_neox_style", True),
+            is_neox=True if rotary_emb is None else rotary_emb.is_neox_style,
             qk_nope_head_dim=layer.qk_nope_head_dim,
             qk_rope_head_dim=layer.qk_rope_head_dim,
         )