diff --git a/python/sglang/jit_kernel/fp8_quantize.py b/python/sglang/jit_kernel/fp8_quantize.py new file mode 100644 index 000000000..fcb77b26b --- /dev/null +++ b/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 diff --git a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py index 8e4fcad80..9ab17576a 100644 --- a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 5ccac1171..87f9c281d 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/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], diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index fdea0f8f3..75019ba11 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/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