chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Nsa trtllm mla sparse fp8 support with Deepseek v3.2 NVFP4 (#18389)

Browse Source
This commit is contained in:
Rain Jiang
2026-02-15 17:29:54 -08:00
committed by GitHub
parent 8290171f52
commit 0ffd0a3995
10 changed files with 352 additions and 183 deletions
Download Patch File Download Diff File Expand all files Collapse all files

238
python/sglang/srt/layers/attention/nsa_backend.py
View File

@@ -33,7 +33,10 @@ from sglang.srt.layers.attention.nsa.utils import (
nsa_cp_round_robin_split_q_seqs,
pad_nsa_cache_seqlens,
)
from sglang.srt.layers.attention.trtllm_mla_backend import _concat_mla_absorb_q_general
from sglang.srt.layers.attention.utils import (
concat_mla_absorb_q_general,
mla_quantize_and_rope_for_fp8,
)
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.utils import is_cuda, is_hip
@@ -340,6 +343,7 @@ class NativeSparseAttnBackend(
self.device_capability = torch.cuda.get_device_capability()
self.device_sm_major = self.device_capability[0]
self.kv_cache_dtype = model_runner.kv_cache_dtype
# Allocate global workspace buffer for TRT-LLM kernels (ragged attention on SM100/B200, or trtllm decode)
if self.device_sm_major >= 10 or self.nsa_decode_impl == "trtllm":
@@ -1299,8 +1303,41 @@ class NativeSparseAttnBackend(
q_rope: Optional[torch.Tensor] = None,
k_rope: Optional[torch.Tensor] = None,
topk_indices: Optional[torch.Tensor] = None,
cos_sin_cache: Optional[torch.Tensor] = None,
is_neox: Optional[bool] = False,
llama_4_scaling: Optional[torch.Tensor] = None,
) -> torch.Tensor:
causal = not layer.is_cross_attention
metadata = self.forward_metadata
assert causal, "NSA is causal only"
nsa_impl = (
self.nsa_decode_impl
if (
forward_batch.forward_mode.is_target_verify()
or forward_batch.forward_mode.is_draft_extend(include_v2=True)
)
else self.nsa_prefill_impl
)
if nsa_impl == "trtllm" and not self.use_mha:
return self._forward_trtllm(
q,
k,
v,
layer,
forward_batch,
metadata.nsa_cache_seqlens_int32,
save_kv_cache,
q_rope,
k_rope,
topk_indices,
cos_sin_cache,
is_neox,
llama_4_scaling,
)
if k is not None:
assert v is not None
if save_kv_cache:
@@ -1316,10 +1353,6 @@ class NativeSparseAttnBackend(
k_rope,
)
metadata = self.forward_metadata
causal = not layer.is_cross_attention
assert causal, "NSA is causal only"
# Use MHA kernel if in MHA_ONE_SHOT mode
if self.use_mha:
assert k is not None and v is not None
@@ -1381,18 +1414,9 @@ class NativeSparseAttnBackend(
page_size=1,
)
nsa_impl = (
self.nsa_decode_impl
if (
forward_batch.forward_mode.is_target_verify()
or forward_batch.forward_mode.is_draft_extend(include_v2=True)
)
else self.nsa_prefill_impl
)
if nsa_impl == "tilelang":
if q_rope is not None:
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_tilelang(
q_all=q_all,
kv_cache=kv_cache,
@@ -1402,7 +1426,7 @@ class NativeSparseAttnBackend(
)
elif nsa_impl == "flashmla_sparse":
if q_rope is not None:
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
if topk_transform_method == TopkTransformMethod.RAGGED:
if any(forward_batch.extend_prefix_lens_cpu):
@@ -1426,7 +1450,7 @@ class NativeSparseAttnBackend(
)
elif nsa_impl == "flashmla_kv":
if q_rope is not None:
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_flashmla_kv(
q_all=q_all,
kv_cache=kv_cache,
@@ -1452,21 +1476,6 @@ class NativeSparseAttnBackend(
logit_cap=layer.logit_cap,
page_size=1,
)
elif nsa_impl == "trtllm":
assert forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend(
include_v2=True
), "TRT-LLM NSA only supports target_verify/draft_extend; normal extend untested."
if q_rope is not None:
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
# Use expanded seq_lens for per-token decode in target_verify/draft_extend.
return self._forward_trtllm(
q_all=q_all,
kv_cache=kv_cache,
page_table_1=page_table_1,
metadata=metadata,
sm_scale=layer.scaling,
seq_lens=metadata.nsa_cache_seqlens_int32,
)
else:
raise ValueError(f"Unsupported {nsa_impl = }")
@@ -1482,7 +1491,32 @@ class NativeSparseAttnBackend(
q_rope: Optional[torch.Tensor] = None,
k_rope: Optional[torch.Tensor] = None,
topk_indices: Optional[torch.Tensor] = None,
cos_sin_cache: Optional[torch.Tensor] = None,
is_neox: Optional[bool] = False,
llama_4_scaling: Optional[torch.Tensor] = None,
) -> torch.Tensor:
causal = not layer.is_cross_attention
metadata = self.forward_metadata
assert causal, "NSA is causal only"
if self.nsa_decode_impl == "trtllm":
return self._forward_trtllm(
q,
k,
v,
layer,
forward_batch,
metadata.cache_seqlens_int32,
save_kv_cache,
q_rope,
k_rope,
topk_indices,
cos_sin_cache,
is_neox,
llama_4_scaling,
)
if k is not None:
assert v is not None
if save_kv_cache:
@@ -1498,10 +1532,6 @@ class NativeSparseAttnBackend(
k_rope,
)
metadata = self.forward_metadata
causal = not layer.is_cross_attention
assert causal, "NSA is causal only"
# Do absorbed multi-latent attention
kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
if q_rope is not None:
@@ -1529,7 +1559,7 @@ class NativeSparseAttnBackend(
if self.nsa_decode_impl == "flashmla_sparse":
if q_rope is not None:
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_flashmla_sparse(
q_all=q_all,
kv_cache=kv_cache,
@@ -1539,7 +1569,7 @@ class NativeSparseAttnBackend(
)
elif self.nsa_decode_impl == "flashmla_kv":
if q_rope is not None:
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_flashmla_kv(
q_all=q_all,
kv_cache=kv_cache,
@@ -1552,7 +1582,7 @@ class NativeSparseAttnBackend(
)
elif self.nsa_decode_impl == "tilelang":
if q_rope is not None:
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_tilelang(
q_all=q_all,
kv_cache=kv_cache,
@@ -1587,18 +1617,6 @@ class NativeSparseAttnBackend(
bs=forward_batch.batch_size,
)
elif self.nsa_decode_impl == "trtllm":
if q_rope is not None:
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_trtllm(
q_all=q_all,
kv_cache=kv_cache,
page_table_1=page_table_1,
metadata=metadata,
sm_scale=layer.scaling,
seq_lens=metadata.cache_seqlens_int32,
)
else:
assert False, f"Unsupported {self.nsa_decode_impl = }"
@@ -1860,22 +1878,103 @@ class NativeSparseAttnBackend(
def _forward_trtllm(
self,
q_all: torch.Tensor,
kv_cache: torch.Tensor,
page_table_1: torch.Tensor,
metadata: NSAMetadata,
sm_scale: float,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
layer: RadixAttention,
forward_batch: ForwardBatch,
seq_lens: torch.Tensor,
save_kv_cache=True,
# For multi-head latent attention
q_rope: Optional[torch.Tensor] = None,
k_rope: Optional[torch.Tensor] = None,
topk_indices: Optional[torch.Tensor] = None,
cos_sin_cache: Optional[torch.Tensor] = None,
is_neox: Optional[bool] = False,
llama_4_scaling: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Forward using TRT-LLM sparse MLA kernel."""
import flashinfer.decode
metadata = self.forward_metadata
merge_query = q_rope is not None
if self.kv_cache_dtype == torch.float8_e4m3fn:
# For FP8 path, we quantize the query and rope parts and merge them into a single tensor
# Note: rope application in deepseek_v2.py:forward_absorb_prepare is skipped for FP8 decode path of this trtllm_mla backend
assert q_rope is not None, "For FP8 path q_rope should not be None."
assert k_rope is not None, "For FP8 path k_rope should not be None."
assert (
cos_sin_cache is not None
), "For FP8 path cos_sin_cache should not be None."
q, k, k_rope = mla_quantize_and_rope_for_fp8(
q,
q_rope,
k.squeeze(1),
k_rope.squeeze(1),
forward_batch.positions,
cos_sin_cache,
is_neox,
self.kv_lora_rank,
self.qk_rope_head_dim,
)
merge_query = False
# Save KV cache if requested
if save_kv_cache:
assert (
k is not None and k_rope is not None
), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None."
cache_loc = (
forward_batch.out_cache_loc
if not layer.is_cross_attention
else forward_batch.encoder_out_cache_loc
)
forward_batch.token_to_kv_pool.set_mla_kv_buffer(
layer, cache_loc, k, k_rope
)
k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
kv_cache = k_cache.view(-1, self.real_page_size, self.kv_cache_dim).unsqueeze(1)
if merge_query:
q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
q_rope_reshaped = q_rope.view(
-1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
)
q_all = concat_mla_absorb_q_general(q_nope, q_rope_reshaped)
else:
q_all = q.view(-1, layer.tp_q_head_num, layer.head_dim)
# Align topk_indices with q dimensions
if topk_indices is not None:
topk_indices = self._pad_topk_indices(topk_indices, q.shape[0])
if envs.SGLANG_NSA_FUSE_TOPK.get():
page_table_1 = topk_indices
else:
page_table_1 = transform_index_page_table_decode(
page_table=metadata.page_table_1,
topk_indices=topk_indices,
page_size=1,
)
q_scale = 1.0
k_scale = (
layer.k_scale_float
if getattr(layer, "k_scale_float", None) is not None
else 1.0
)
bmm1_scale = q_scale * k_scale * layer.scaling
batch_size = page_table_1.shape[0]
_, num_heads, head_dim = q_all.shape
q = q_all.view(batch_size, 1, num_heads, head_dim)
kv = kv_cache.view(-1, 1, self.real_page_size, self.kv_cache_dim)
block_tables = page_table_1.unsqueeze(1)
seq_lens = metadata.cache_seqlens_int32 if seq_lens is None else seq_lens
out = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
query=q,
@@ -1888,7 +1987,7 @@ class NativeSparseAttnBackend(
seq_lens=seq_lens,
max_seq_len=metadata.max_seq_len_k,
sparse_mla_top_k=self.nsa_index_topk,
bmm1_scale=sm_scale,
bmm1_scale=bmm1_scale,
backend="trtllm-gen",
)
# Output: [batch, q_len=1, heads, v_dim] -> [batch, heads, v_dim]
@@ -1933,12 +2032,19 @@ class NativeSparseAttnBackend(
sum_seq_lens = sum(forward_batch.seq_lens_cpu)
device_sm = get_device_sm()
# when nsa prefill impl is trtllm, use its max chunk capacity as mha max kv len
mha_max_kv_len = (
forward_batch.get_max_chunk_capacity()
if self.nsa_prefill_impl == "trtllm"
else self.nsa_index_topk
)
# Requirements: H200/B200, short sequences, supported dtype, fits in chunk
self.use_mha = (
(
device_sm == 90 or (device_sm >= 100 and device_sm < 110)
) # SM90/SM100 only
and max_kv_len <= self.nsa_index_topk # Short enough for MHA
and max_kv_len <= mha_max_kv_len # Short enough for MHA
and forward_batch.token_to_kv_pool.dtype
in [torch.bfloat16, torch.float8_e4m3fn]
and sum_seq_lens
@@ -2022,7 +2128,7 @@ class NativeSparseAttnMultiStepBackend:
self.topk = topk
self.speculative_num_steps = speculative_num_steps
self.attn_backends = []
for i in range(self.speculative_num_steps):
for i in range(self.speculative_num_steps - 1):
self.attn_backends.append(
NativeSparseAttnBackend(
model_runner,
@@ -2037,11 +2143,11 @@ class NativeSparseAttnMultiStepBackend:
self.attn_backends[i].init_forward_metadata(forward_batch)
def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
for i in range(self.speculative_num_steps):
for i in range(self.speculative_num_steps - 1):
self.attn_backends[i].init_cuda_graph_state(max_bs, max_num_tokens)
def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch):
for i in range(self.speculative_num_steps):
for i in range(self.speculative_num_steps - 1):
self.attn_backends[i].init_forward_metadata_capture_cuda_graph(
forward_batch.batch_size,
forward_batch.batch_size * self.topk,
@@ -2068,7 +2174,7 @@ class NativeSparseAttnMultiStepBackend:
# Use multi-backend fused copy when we have 3 or more backends
# This is 3x faster than calling the single-backend copy 3 times
if self.speculative_num_steps >= 3:
if self.speculative_num_steps > 3:
try:
from sglang.jit_kernel.fused_metadata_copy import (
fused_metadata_copy_multi_cuda,
@@ -2187,7 +2293,7 @@ class NativeSparseAttnMultiStepBackend:
)
# Copy remaining backends one by one (if > 3 backends)
for i in range(3, self.speculative_num_steps):
for i in range(3, self.speculative_num_steps - 1):
self.attn_backends[
i
].init_forward_metadata_replay_cuda_graph_from_precomputed(
@@ -2205,7 +2311,7 @@ class NativeSparseAttnMultiStepBackend:
print(
f"Warning: Multi-backend fused metadata copy kernel failed with error: {e}, falling back to loop."
)
for i in range(self.speculative_num_steps):
for i in range(self.speculative_num_steps - 1):
self.attn_backends[
i
].init_forward_metadata_replay_cuda_graph_from_precomputed(
@@ -2215,7 +2321,7 @@ class NativeSparseAttnMultiStepBackend:
)
else:
# Less than 3 backends: copy to each backend individually
for i in range(self.speculative_num_steps):
for i in range(self.speculative_num_steps - 1):
self.attn_backends[
i
].init_forward_metadata_replay_cuda_graph_from_precomputed(
@@ -2225,7 +2331,7 @@ class NativeSparseAttnMultiStepBackend:
)
else:
# Fallback: compute metadata separately for each backend
for i in range(self.speculative_num_steps):
for i in range(self.speculative_num_steps - 1):
self.attn_backends[i].init_forward_metadata_replay_cuda_graph(
bs=bs,
req_pool_indices=forward_batch.req_pool_indices,

110
python/sglang/srt/layers/attention/trtllm_mla_backend.py
View File

@@ -19,14 +19,16 @@ from sglang.srt.layers.attention.flashinfer_mla_backend import (
FlashInferMLAMultiStepDraftBackend,
)
from sglang.srt.layers.attention.utils import (
concat_mla_absorb_q_general,
create_flashmla_kv_indices_triton,
get_num_page_per_block_flashmla,
mla_quantize_and_rope_for_fp8,
)
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import is_cuda, is_flashinfer_available, is_float4_e2m1fn_x2
from sglang.srt.utils import is_flashinfer_available, is_float4_e2m1fn_x2
if is_flashinfer_available():
import flashinfer
@@ -38,11 +40,6 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
_is_cuda = is_cuda()
if _is_cuda:
from sgl_kernel import concat_mla_absorb_q
# Constants
DEFAULT_WORKSPACE_SIZE_MB = 150 # Memory workspace size in MB
@@ -669,84 +666,6 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
def init_mha_chunk_metadata(self, forward_batch: ForwardBatch):
super().init_mha_chunk_metadata(forward_batch, disable_flashinfer_ragged=True)
def quantize_and_rope_for_fp8(
self,
q_nope: torch.Tensor,
q_rope: torch.Tensor,
k_nope: torch.Tensor,
k_rope: torch.Tensor,
forward_batch: ForwardBatch,
cos_sin_cache: torch.Tensor,
is_neox: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Quantize and apply RoPE for FP8 attention path.
This function handles the FP8 quantization and RoPE application for MLA attention.
It takes separate query/key nope and rope components, applies RoPE to the rope parts,
quantizes all components to FP8, and merges the query components into a single tensor.
Args:
q_nope: Query no-position-encoding component [seq_len, num_heads, kv_lora_rank]
- expected dtype: torch.bfloat16
q_rope: Query RoPE component [seq_len, num_heads, qk_rope_head_dim]
- expected dtype: torch.bfloat16
k_nope: Key no-position-encoding component [seq_len, num_heads, kv_lora_rank]
- expected dtype: torch.bfloat16
k_rope: Key RoPE component [seq_len, num_heads, qk_rope_head_dim]
- expected dtype: torch.bfloat16
forward_batch: Forward batch containing position information
cos_sin_cache: Precomputed cosine/sine cache for RoPE
- expected dtype: matches q_/k_ input dtype (torch.bfloat16)
is_neox: Whether to use NeoX-style RoPE (interleaved) or GPT-style (half rotation)
Returns:
tuple: (merged_q_out, k_nope_out, k_rope_out) quantized to FP8
- merged_q_out: [seq_len, num_heads, kv_lora_rank + qk_rope_head_dim], dtype=torch.float8_e4m3fn
- k_nope_out: [seq_len, num_heads, kv_lora_rank], dtype=torch.float8_e4m3fn
- k_rope_out: [seq_len, num_heads, qk_rope_head_dim], dtype=torch.float8_e4m3fn
"""
attn_dtype = torch.float8_e4m3fn
q_len, num_heads = q_rope.shape[0], q_rope.shape[1]
# Allocate output tensors with FP8 dtype
# Query output will contain merged nope + rope components
q_out = q_rope.new_empty(
q_len,
num_heads,
self.kv_lora_rank + self.qk_rope_head_dim,
dtype=attn_dtype,
)
# Key outputs maintain original shapes but with FP8 dtype
k_rope_out = k_rope.new_empty(k_rope.shape, dtype=attn_dtype)
k_nope_out = k_nope.new_empty(k_nope.shape, dtype=attn_dtype)
# Apply RoPE and quantize all components in a single fused kernel call
# This kernel handles:
# 1. RoPE application to q_rope and k_rope using cos_sin_cache and positions
# 2. Quantization of all components to FP8 format
# 3. Output placement into pre-allocated tensors
flashinfer.rope.mla_rope_quantize_fp8(
q_rope=q_rope,
k_rope=k_rope,
q_nope=q_nope,
k_nope=k_nope,
cos_sin_cache=cos_sin_cache,
pos_ids=forward_batch.positions,
is_neox=is_neox,
quantize_dtype=attn_dtype,
# Output tensor slicing: q_out contains [nope_part, rope_part]
q_rope_out=q_out[..., self.kv_lora_rank :], # RoPE part goes to end
k_rope_out=k_rope_out,
q_nope_out=q_out[..., : self.kv_lora_rank], # Nope part goes to beginning
k_nope_out=k_nope_out,
# Quantization scales (set to 1.0 for no additional scaling)
quant_scale_q=1.0,
quant_scale_kv=1.0,
)
return q_out, k_nope_out, k_rope_out
def pad_draft_extend_query(
self,
q: torch.Tensor,
@@ -849,14 +768,16 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
assert all(
x is not None for x in [q_rope, k_rope, cos_sin_cache]
), "For FP8 path and using flashinfer.rope.mla_rope_quantize we need all of q_rope, k_rope and cos_sin_cache to be not None."
q, k, k_rope = self.quantize_and_rope_for_fp8(
q, k, k_rope = mla_quantize_and_rope_for_fp8(
q,
q_rope,
k.squeeze(1),
k_rope.squeeze(1),
forward_batch,
forward_batch.positions,
cos_sin_cache,
is_neox,
self.kv_lora_rank,
self.qk_rope_head_dim,
)
merge_query = False
@@ -876,7 +797,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
q_rope_reshaped = q_rope.view(
-1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
)
query = _concat_mla_absorb_q_general(q_nope, q_rope_reshaped)
query = concat_mla_absorb_q_general(q_nope, q_rope_reshaped)
else:
# For FP8 path, we already have the query and rope parts merged because of the quantize_and_rope_for_fp8 function
query = q.view(-1, layer.tp_q_head_num, layer.head_dim)
@@ -986,14 +907,16 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
assert all(
x is not None for x in [q_rope, k_rope, cos_sin_cache]
), "For FP8 path and using flashinfer.rope.mla_rope_quantize we need all of q_rope, k_rope and cos_sin_cache to be not None."
q, k, k_rope = self.quantize_and_rope_for_fp8(
q, k, k_rope = mla_quantize_and_rope_for_fp8(
q,
q_rope,
k.squeeze(1),
k_rope.squeeze(1),
forward_batch,
forward_batch.positions,
cos_sin_cache,
is_neox,
self.kv_lora_rank,
self.qk_rope_head_dim,
)
merge_query = False
@@ -1014,7 +937,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
q_rope_reshaped = q_rope.view(
-1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
)
q = _concat_mla_absorb_q_general(q_nope, q_rope_reshaped)
q = concat_mla_absorb_q_general(q_nope, q_rope_reshaped)
q = q.view(-1, layer.tp_q_head_num, layer.head_dim)
@@ -1232,10 +1155,3 @@ class TRTLLMMLAMultiStepDraftBackend(FlashInferMLAMultiStepDraftBackend):
kv_indptr_buf=self.kv_indptr[i],
q_indptr_decode_buf=self.q_indptr_decode,
)
def _concat_mla_absorb_q_general(q_nope, q_rope):
if _is_cuda and q_nope.shape[-1] == 512 and q_rope.shape[-1] == 64:
return concat_mla_absorb_q(q_nope, q_rope)
else:
return torch.cat([q_nope, q_rope], dim=-1)

99
python/sglang/srt/layers/attention/utils.py
View File

@@ -2,9 +2,16 @@ import torch
import triton
import triton.language as tl
from sglang.srt.utils import is_cuda
_FLASHMLA_CREATE_KV_BLOCK_SIZE = 4096
FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON = tl.constexpr(_FLASHMLA_CREATE_KV_BLOCK_SIZE)
_is_cuda = is_cuda()
if _is_cuda:
from sgl_kernel import concat_mla_absorb_q
@triton.jit
def create_flashinfer_kv_indices_triton(
@@ -312,3 +319,95 @@ def canonicalize_stride(tensor: torch.Tensor) -> torch.Tensor:
new_strides[i] = new_strides[i + 1] * sizes[i + 1]
return tensor.as_strided(sizes, new_strides)
def mla_quantize_and_rope_for_fp8(
q_nope: torch.Tensor,
q_rope: torch.Tensor,
k_nope: torch.Tensor,
k_rope: torch.Tensor,
pos_ids: torch.Tensor,
cos_sin_cache: torch.Tensor,
is_neox: bool,
kv_lora_rank: int,
qk_rope_head_dim: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
import flashinfer.rope
"""Quantize and apply RoPE for FP8 attention path.
This function handles the FP8 quantization and RoPE application for MLA attention.
It takes separate query/key nope and rope components, applies RoPE to the rope parts,
quantizes all components to FP8, and merges the query components into a single tensor.
Args:
q_nope: Query no-position-encoding component [seq_len, num_heads, kv_lora_rank]
- expected dtype: torch.bfloat16
q_rope: Query RoPE component [seq_len, num_heads, qk_rope_head_dim]
- expected dtype: torch.bfloat16
k_nope: Key no-position-encoding component [seq_len, num_heads, kv_lora_rank]
- expected dtype: torch.bfloat16
k_rope: Key RoPE component [seq_len, num_heads, qk_rope_head_dim]
- expected dtype: torch.bfloat16
pos_ids: Position indices for each token
- expected dtype: torch.int64 or torch.int32
cos_sin_cache: Precomputed cosine/sine cache for RoPE
- expected dtype: matches q_/k_ input dtype (torch.bfloat16)
is_neox: Whether to use NeoX-style RoPE (interleaved) or GPT-style (half rotation)
kv_lora_rank: Dimension of the no-position-encoding component
qk_rope_head_dim: Dimension of the RoPE component
Returns:
tuple: (merged_q_out, k_nope_out, k_rope_out) quantized to FP8
- merged_q_out: [seq_len, num_heads, kv_lora_rank + qk_rope_head_dim], dtype=torch.float8_e4m3fn
- k_nope_out: [seq_len, num_heads, kv_lora_rank], dtype=torch.float8_e4m3fn
- k_rope_out: [seq_len, num_heads, qk_rope_head_dim], dtype=torch.float8_e4m3fn
"""
attn_dtype = torch.float8_e4m3fn
q_len, num_heads = q_rope.shape[0], q_rope.shape[1]
# Allocate output tensors with FP8 dtype
# Query output will contain merged nope + rope components
q_out = q_rope.new_empty(
q_len,
num_heads,
kv_lora_rank + qk_rope_head_dim,
dtype=attn_dtype,
)
# Key outputs maintain original shapes but with FP8 dtype
k_rope_out = k_rope.new_empty(k_rope.shape, dtype=attn_dtype)
k_nope_out = k_nope.new_empty(k_nope.shape, dtype=attn_dtype)
# Apply RoPE and quantize all components in a single fused kernel call
# This kernel handles:
# 1. RoPE application to q_rope and k_rope using cos_sin_cache and positions
# 2. Quantization of all components to FP8 format
# 3. Output placement into pre-allocated tensors
flashinfer.rope.mla_rope_quantize_fp8(
q_rope=q_rope,
k_rope=k_rope,
q_nope=q_nope,
k_nope=k_nope,
cos_sin_cache=cos_sin_cache,
pos_ids=pos_ids,
is_neox=is_neox,
quantize_dtype=attn_dtype,
# Output tensor slicing: q_out contains [nope_part, rope_part]
q_rope_out=q_out[..., kv_lora_rank:], # RoPE part goes to end
k_rope_out=k_rope_out,
q_nope_out=q_out[..., :kv_lora_rank], # Nope part goes to beginning
k_nope_out=k_nope_out,
# Quantization scales (set to 1.0 for no additional scaling)
quant_scale_q=1.0,
quant_scale_kv=1.0,
)
return q_out, k_nope_out, k_rope_out
def concat_mla_absorb_q_general(q_nope, q_rope):
if _is_cuda and q_nope.shape[-1] == 512 and q_rope.shape[-1] == 64:
return concat_mla_absorb_q(q_nope, q_rope)
else:
return torch.cat([q_nope, q_rope], dim=-1)

32
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -1404,13 +1404,16 @@ class MLATokenToKVPool(KVCache):
self.kv_lora_rank = kv_lora_rank
self.qk_rope_head_dim = qk_rope_head_dim
self.use_nsa = use_nsa
self.nsa_kv_cache_store_fp8 = use_nsa and dtype == torch.float8_e4m3fn
assert not (
self.nsa_kv_cache_store_fp8 and override_kv_cache_dim is None
), "override_kv_cache_dim must be provided when using NSA with FP8 kv cache storage"
self.nsa_kv_cache_store_fp8 = (
use_nsa
and dtype == torch.float8_e4m3fn
and override_kv_cache_dim is not None
)
# When override_kv_cache_dim is provided with nsa model, we assume the
# override kv cache dim is correct and use it directly.
self.kv_cache_dim = (
override_kv_cache_dim
if self.use_nsa and self.nsa_kv_cache_store_fp8
if self.nsa_kv_cache_store_fp8
else (kv_lora_rank + qk_rope_head_dim)
)
@@ -1492,7 +1495,7 @@ class MLATokenToKVPool(KVCache):
cache_v: torch.Tensor,
):
layer_id = layer.layer_id
assert not (self.use_nsa and self.nsa_kv_cache_store_fp8)
assert not self.nsa_kv_cache_store_fp8
if cache_k.dtype != self.dtype:
cache_k = cache_k.to(self.dtype)
@@ -1512,7 +1515,7 @@ class MLATokenToKVPool(KVCache):
):
layer_id = layer.layer_id
if self.use_nsa and self.nsa_kv_cache_store_fp8:
if self.nsa_kv_cache_store_fp8:
# OPTIMIZATION: Quantize k_nope and k_rope separately to avoid concat overhead
# This also enables reuse of set_mla_kv_buffer_triton two-tensor write path
# quantize_k_cache_separate returns (nope_part, rope_part) as uint8 bytes
@@ -1661,7 +1664,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
cache_v: torch.Tensor,
):
layer_id = layer.layer_id
assert not (self.use_nsa and self.nsa_kv_cache_store_fp8)
assert not self.nsa_kv_cache_store_fp8
if cache_k.dtype != self.dtype:
from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil
@@ -1686,7 +1689,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
):
layer_id = layer.layer_id
if self.use_nsa and self.nsa_kv_cache_store_fp8:
if self.nsa_kv_cache_store_fp8:
# original cache_k: (num_tokens, num_heads 1, hidden 576); we unsqueeze the page_size=1 dim here
# TODO no need to cat
cache_k = torch.cat([cache_k_nope, cache_k_rope], dim=-1)
@@ -1740,20 +1743,13 @@ class NSATokenToKVPool(MLATokenToKVPool):
device: str,
index_head_dim: int,
enable_memory_saver: bool,
kv_cache_dim: int,
start_layer: Optional[int] = None,
end_layer: Optional[int] = None,
):
assert (
kv_lora_rank % self.quant_block_size == 0
), f"kv_lora_rank {kv_lora_rank} must be multiple of quant_block_size {self.quant_block_size}"
# Calculate override_kv_cache_dim for FP8 storage:
# kv_lora_rank + scale storage (kv_lora_rank // quant_block_size * 4 bytes) + rope dimension storage
# Note: rope dimension is stored in original dtype (bf16), not quantized to fp8
override_dim = (
kv_lora_rank
+ kv_lora_rank // self.quant_block_size * 4
+ qk_rope_head_dim * self.rope_storage_dtype.itemsize
kv_cache_dim if kv_cache_dim != kv_lora_rank + qk_rope_head_dim else None
)
super().__init__(

40
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
View File

@@ -209,6 +209,45 @@ class ModelRunnerKVCacheMixin:
)
return total_rest_memory - mamba_state_memory
def calculate_mla_kv_cache_dim(self: ModelRunner) -> int:
is_nsa_model = is_deepseek_nsa(self.model_config.hf_config)
kv_cache_dtype = self.kv_cache_dtype
kv_lora_rank = self.model_config.kv_lora_rank
qk_rope_head_dim = self.model_config.qk_rope_head_dim
kv_cache_dim = kv_lora_rank + qk_rope_head_dim # default mla kv cache dim
# For non-NSA models, MLA kv cache dim is simply kv_lora_rank + qk_rope_head_dim
if not is_nsa_model:
return kv_cache_dim
# TRTLLM backend does not override kv_cache_dim for MLA kv cache
# Assuming nsa prefill and decode backends are the same when using trtllm MLA backend,
# since it is not compatible for trtllm and other mla attn backend due to the different
# kv cache layout.
if (
self.server_args.nsa_prefill_backend == "trtllm"
or self.server_args.nsa_decode_backend == "trtllm"
):
return kv_cache_dim
quant_block_size = NSATokenToKVPool.quant_block_size
rope_storage_dtype = NSATokenToKVPool.rope_storage_dtype
# Calculate override_kv_cache_dim for FP8 storage for non-trtllm attention backends:
# kv_lora_rank + scale storage (kv_lora_rank // quant_block_size * 4 bytes) + rope dimension storage
# Note: rope dimension is stored in original dtype (bf16), not quantized to fp8
if kv_cache_dtype == torch.float8_e4m3fn:
assert (
kv_lora_rank % quant_block_size == 0
), f"kv_lora_rank {kv_lora_rank} must be multiple of quant_block_size {quant_block_size}"
return (
kv_lora_rank
+ kv_lora_rank // quant_block_size * 4
+ qk_rope_head_dim * rope_storage_dtype.itemsize
)
return kv_cache_dim
def set_num_tokens_hybrid_swa(self: ModelRunner):
page_size = self.server_args.page_size
@@ -492,6 +531,7 @@ class ModelRunnerKVCacheMixin:
qk_rope_head_dim=self.model_config.qk_rope_head_dim,
layer_num=self.num_effective_layers,
device=self.device,
kv_cache_dim=self.calculate_mla_kv_cache_dim(),
enable_memory_saver=self.server_args.enable_memory_saver,
start_layer=self.start_layer,
end_layer=self.end_layer,

6
python/sglang/srt/models/deepseek_v2.py
View File

@@ -1493,6 +1493,12 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin):
"""
Check if we should skip rope and do fused rope+quantize for TRTLLM MLA decode in fp8_e4m3 path.
"""
if self.current_attention_backend == "nsa":
return (
get_global_server_args().nsa_decode_backend == "trtllm"
or get_global_server_args().nsa_prefill_backend == "trtllm"
) and forward_batch.attn_backend.kv_cache_dtype == torch.float8_e4m3fn
return (
self.current_attention_backend == "trtllm_mla"
and (