[perf] prepare_prefill_qkv hook + fp8 quantize jit kernel (#25460)

python/sglang/jit_kernel/fp8_quantize.py

@@ -0,0 +1,157 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _fp8_quantize_kernel(
+    x_ptr,
+    out_ptr,
+    scale_inv,
+    M,
+    x_row_stride,
+    out_row_stride,
+    N: tl.constexpr,
+    FP8_DTYPE: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    ENABLE_PDL: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    m_idx = pid * BLOCK_M + tl.arange(0, BLOCK_M)
+    m_mask = m_idx < M
+    n_idx = tl.arange(0, N)
+
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    x_off = m_idx[:, None] * x_row_stride + n_idx[None, :]
+    x = tl.load(x_ptr + x_off, mask=m_mask[:, None])
+
+    x_fp8 = (x.to(tl.float32) * scale_inv).to(FP8_DTYPE)
+
+    out_off = m_idx[:, None] * out_row_stride + n_idx[None, :]
+    tl.store(out_ptr + out_off, x_fp8, mask=m_mask[:, None])
+
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+def _flatten_to_2d(x: torch.Tensor):
+    """Flatten leading dims onto the row stride; returns (M, N, row_stride).
+
+    Accepts contiguous tensors and last-dim slice views (e.g.
+    ``kv[..., qk_nope:]``) where leading dims still pack onto a uniform row
+    stride.
+    """
+    assert x.stride(-1) == 1, f"expected stride-1 inner dim, got stride={x.stride(-1)}"
+    N = x.shape[-1]
+    if x.ndim == 1:
+        return 1, N, N
+    M = x.numel() // N
+    row_stride = x.stride(-2)
+    for d in range(x.ndim - 2):
+        expected = x.shape[d + 1] * x.stride(d + 1)
+        if x.stride(d) != expected:
+            raise ValueError(
+                f"cannot flatten dim {d}: stride={x.stride(d)} but expected "
+                f"shape[{d+1}]*stride[{d+1}]={expected}. Tensor shape={tuple(x.shape)}, "
+                f"stride={tuple(x.stride())}."
+            )
+    return M, N, row_stride
+
+
+def fp8_quantize(
+    x: torch.Tensor,
+    scale_inv: float = 1.0,
+    out: Optional[torch.Tensor] = None,
+    fp8_dtype: torch.dtype = torch.float8_e4m3fn,
+    enable_pdl: bool = False,
+) -> torch.Tensor:
+    """Cast a BF16/FP16 tensor to FP8 with an optional per-tensor scale.
+
+    Computes ``out = saturate((x * scale_inv) -> fp8)`` element-wise. When
+    ``scale_inv == 1.0`` the multiply is dropped at compile time (pure cast).
+
+    Args:
+        x: BF16 or FP16 tensor. Must have stride(-1) == 1; leading dims must
+           pack uniformly onto the row stride (true for contiguous tensors and
+           for last-dim slice views like ``kv[..., qk_nope:]``).
+        scale_inv: scalar multiplier applied before the cast (i.e. ``1/scale``).
+        out: optional pre-allocated FP8 output. Same shape as ``x``.
+        fp8_dtype: ``torch.float8_e4m3fn`` (default) or ``torch.float8_e5m2``.
+        enable_pdl: opt into Programmatic Dependent Launch (Hopper+).
+
+    Returns:
+        FP8 tensor with the same shape as ``x``.
+    """
+    assert x.dtype in (
+        torch.bfloat16,
+        torch.float16,
+    ), f"fp8_quantize input must be bf16/fp16, got {x.dtype}"
+    assert fp8_dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
+
+    M, N, x_row_stride = _flatten_to_2d(x)
+
+    if out is None:
+        out = torch.empty(x.shape, dtype=fp8_dtype, device=x.device)
+    else:
+        assert out.shape == x.shape and out.dtype == fp8_dtype
+    out_M, _, out_row_stride = _flatten_to_2d(out)
+    assert out_M == M
+
+    fp8_dtype_const = tl.float8e4nv if fp8_dtype is torch.float8_e4m3fn else tl.float8e5
+
+    if M <= 2048:
+        block_m = 4
+    elif M <= 16384:
+        block_m = 16
+    else:
+        block_m = 32
+    num_warps = 4
+    num_stages = 2
+
+    grid = (triton.cdiv(M, block_m),)
+
+    # launch_pdl is NVIDIA-only; the HIP backend rejects unknown kwargs.
+    extra_kwargs = {"launch_pdl": True} if enable_pdl else {}
+
+    _fp8_quantize_kernel[grid](
+        x,
+        out,
+        scale_inv,
+        M,
+        x_row_stride,
+        out_row_stride,
+        N=N,
+        FP8_DTYPE=fp8_dtype_const,
+        BLOCK_M=block_m,
+        ENABLE_PDL=enable_pdl,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        **extra_kwargs,
+    )
+    return out

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

@@ -33,20 +33,26 @@ from typing import TYPE_CHECKING, Optional
 
 import torch
 
+from sglang.jit_kernel.fp8_quantize import fp8_quantize
+from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import mla_kv_pack_quantize_fp8
 from sglang.jit_kernel.utils import is_arch_support_pdl
 from sglang.srt.layers.attention.trtllm_mla_backend import (
     TRTLLMMLABackend,
     TRTLLMMLAMultiStepDraftBackend,
-    _quantize_fp8_qkv,
 )
-from sglang.srt.utils import is_tokenspeed_mla_available
+from sglang.srt.utils import is_flashinfer_available, is_tokenspeed_mla_available
+
+if is_flashinfer_available():
+    import flashinfer.rope as _flashinfer_rope
 
 if is_tokenspeed_mla_available():
     import tokenspeed_mla
 
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
     from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
 
 logger = logging.getLogger(__name__)
 
@@ -77,6 +83,9 @@ def _get_tokenspeed_workspace(
     return _g_tokenspeed_workspace[device]
 
 
+# TODO(Qiaolin-Yu): Merge this attention backend into trtllm_mla_backend.py
+# once the same CuteDSL kernels in flashinfer_trtllm are stable
+# and there is no performance gap compared to this backend.
 class TokenspeedMLABackend(TRTLLMMLABackend):
     """tokenspeed-mla CuTe DSL attention backend (Blackwell SM100, FP8 KV)."""
 
@@ -146,6 +155,121 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
                         enable_ex2_emulation=enable_ex2_emulation,
                     )
 
+    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,
+        positions: torch.Tensor,
+        is_neox: bool,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+    ) -> 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.
+        """
+        num_heads = q_nope.shape[1]
+        seq_len = q_nope.shape[0]
+        q_fp8 = torch.empty(
+            (seq_len, num_heads, qk_nope_head_dim + qk_rope_head_dim),
+            dtype=torch.float8_e4m3fn,
+            device=q_nope.device,
+        )
+        k_fp8 = torch.empty(
+            (seq_len, num_heads, qk_nope_head_dim + qk_rope_head_dim),
+            dtype=torch.float8_e4m3fn,
+            device=k_nope.device,
+        )
+        if seq_len == 0:
+            return q_fp8, k_fp8
+
+        # Broadcast the shared latent k_pe across heads — RoPE is position-only
+        # so per-head outputs are identical, and the cache write below reuses
+        # head 0.
+        if k_pe.dim() == 3 and k_pe.shape[1] == 1:
+            k_pe_expanded = k_pe.expand(-1, num_heads, -1)
+        else:
+            k_pe_expanded = k_pe
+
+        _flashinfer_rope.mla_rope_quantize_fp8(
+            q_rope=q_pe,
+            k_rope=k_pe_expanded,
+            q_nope=q_nope,
+            k_nope=k_nope,
+            cos_sin_cache=cos_sin_cache,
+            pos_ids=positions,
+            is_neox=is_neox,
+            quantize_dtype=torch.float8_e4m3fn,
+            q_rope_out=q_fp8[..., qk_nope_head_dim:],
+            k_rope_out=k_fp8[..., qk_nope_head_dim:],
+            q_nope_out=q_fp8[..., :qk_nope_head_dim],
+            k_nope_out=k_fp8[..., :qk_nope_head_dim],
+            quant_scale_q=1.0,
+            quant_scale_kv=1.0,
+            enable_pdl=is_arch_support_pdl(),
+        )
+        return q_fp8, k_fp8
+
+    def prepare_prefill_qkv(
+        self,
+        *,
+        q: torch.Tensor,
+        q_pe: torch.Tensor,
+        kv_a: torch.Tensor,
+        k_pe: torch.Tensor,
+        positions: torch.Tensor,
+        layer: "DeepseekV2AttentionMLA",
+        forward_batch: "ForwardBatch",
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Build FP8 (Q, K, V) for the FMHA kernel and write FP8 KV cache."""
+        kv = layer.kv_b_proj(kv_a)[0]
+        kv = kv.view(
+            -1, layer.num_local_heads, layer.qk_nope_head_dim + layer.v_head_dim
+        )
+        k_nope = kv[..., : layer.qk_nope_head_dim]
+        v_bf16 = kv[..., layer.qk_nope_head_dim :]
+        q_nope = q[..., : layer.qk_nope_head_dim]
+
+        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,
+            positions=positions,
+            is_neox=getattr(layer.rotary_emb, "is_neox_style", True),
+            qk_nope_head_dim=layer.qk_nope_head_dim,
+            qk_rope_head_dim=layer.qk_rope_head_dim,
+        )
+        v_fp8 = fp8_quantize(v_bf16, enable_pdl=is_arch_support_pdl())
+
+        # k_pe is shared across heads (RoPE is position-only), so head 0
+        # reproduces the original [tokens, 1, qk_rope] latent layout.
+        kv_a_fp8 = fp8_quantize(kv_a, enable_pdl=is_arch_support_pdl())
+        k_pe_fp8 = k_fp8[:, 0:1, layer.qk_nope_head_dim :]
+        forward_batch.token_to_kv_pool.set_mla_kv_buffer(
+            layer.attn_mha,
+            forward_batch.out_cache_loc,
+            kv_a_fp8.unsqueeze(1),
+            k_pe_fp8,
+        )
+        return q_fp8, k_fp8, v_fp8
+
+    def pack_prefix_chunk_kv(
+        self,
+        k_nope: torch.Tensor,
+        k_pe: torch.Tensor,
+        v: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Pack strided ``k_nope``+``k_pe`` into contig FP8 K and quantize
+        strided ``v`` into contig FP8 V in a single kernel.
+        """
+        return mla_kv_pack_quantize_fp8(
+            k_nope, k_pe, v, enable_pdl=is_arch_support_pdl()
+        )
+
     def _run_decode_kernel(
         self,
         query: torch.Tensor,
@@ -194,18 +318,8 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
         return_lse: bool,
         out_buffer: torch.Tensor,
         o_sf_scale: float = 1.0,
-    ):
-        # Quantize to FP8 for the Blackwell FP8 GEMM speedup (mirrors trtllm-gen).
-        # The kernel has no per-tensor scale knob for either K or V, so we
-        # require both ``k_scale_float`` and ``v_scale_float`` to be 1.0.
-        if self.data_type == torch.float8_e4m3fn:
-            q, k, v, k_scale, v_scale = _quantize_fp8_qkv(q, k, v, layer)
-            assert k_scale == 1.0 and v_scale == 1.0, (
-                "tokenspeed_mla prefill kernel has no per-tensor K/V scale "
-                "knob; both k_scale_float and v_scale_float must be 1.0, got "
-                f"k_scale={k_scale}, v_scale={v_scale}."
-            )
-
+    ):  # Q/K/V arrive already in FP8 via the model-side fused path
+        # (prepare_prefill_qkv / pack_prefix_chunk_kv); no quantize here.
         return tokenspeed_mla.tokenspeed_mla_prefill(
             query=q,
             key=k,

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

@@ -1144,7 +1144,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             assert k_rope is None
             chunk_idx = forward_batch.prefix_chunk_idx
 
-            out = torch.zeros(
+            out = torch.empty(
                 q.shape[0],
                 layer.tp_q_head_num,
                 layer.v_head_dim,
@@ -1187,7 +1187,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
 
             return result
         else:
-            out = torch.zeros(
+            out = torch.empty(
                 q.shape[0],
                 q.shape[1],
                 v.shape[2],

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py

@@ -220,6 +220,24 @@ class DeepseekMHAForwardMixin:
             kv_a = self.kv_a_layernorm(kv_a)
 
         k_pe = latent_cache[:, :, self.kv_lora_rank :]
+
+        # Backend prefill hook: the backend owns the BF16->FP8 transition
+        # (fused RoPE + quantize for Q/K, direct FP8 KV-cache write) and
+        # returns FP8 tensors ready for its kernel. Backends without the
+        # hook fall through to the BF16 path below.
+        backend = _resolve_attn_backend(forward_batch)
+        if hasattr(backend, "prepare_prefill_qkv"):
+            q_out, k_out, v_out = backend.prepare_prefill_qkv(
+                q=q,
+                q_pe=q_pe,
+                kv_a=kv_a,
+                k_pe=k_pe,
+                positions=positions,
+                layer=self,
+                forward_batch=forward_batch,
+            )
+            return q_out, k_out, v_out, forward_batch
+
         if self.rotary_emb is not None:
             q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
         q[..., self.qk_nope_head_dim :] = q_pe