From 0ffd0a3995e537c2ec6f450cbb7d4b4db44550aa Mon Sep 17 00:00:00 2001 From: Rain Jiang <96632942+rainj-me@users.noreply.github.com> Date: Sun, 15 Feb 2026 17:29:54 -0800 Subject: [PATCH] Nsa trtllm mla sparse fp8 support with Deepseek v3.2 NVFP4 (#18389) --- docs/advanced_features/server_arguments.md | 4 +- docs/basic_usage/deepseek_v32.md | 4 + .../srt/layers/attention/nsa_backend.py | 238 +++++++++++++----- .../layers/attention/trtllm_mla_backend.py | 110 +------- python/sglang/srt/layers/attention/utils.py | 99 ++++++++ python/sglang/srt/mem_cache/memory_pool.py | 32 ++- .../model_runner_kv_cache_mixin.py | 40 +++ python/sglang/srt/models/deepseek_v2.py | 6 + .../hicache/test_nsa_pool_host_unit.py | 1 + test/registered/kernels/test_nsa_indexer.py | 1 + 10 files changed, 352 insertions(+), 183 deletions(-) diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index e6930c841..48fb777c0 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -266,8 +266,8 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--sampling-backend` | Choose the kernels for sampling layers. | `None` | `flashinfer`, `pytorch`, `ascend` | | `--grammar-backend` | Choose the backend for grammar-guided decoding. | `None` | `xgrammar`, `outlines`, `llguidance`, `none` | | `--mm-attention-backend` | Set multimodal attention backend. | `None` | `sdpa`, `fa3`, `fa4`, `triton_attn`, `ascend_attn`, `aiter_attn` | -| `--nsa-prefill-backend` | Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention). | `flashmla_sparse` | `flashmla_sparse`, `flashmla_kv`, `flashmla_auto`, `fa3`, `tilelang`, `aiter` | -| `--nsa-decode-backend` | Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding. | `fa3` | `flashmla_sparse`, `flashmla_kv`, `fa3`, `tilelang`, `aiter` | +| `--nsa-prefill-backend` | Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention). | `flashmla_sparse` | `flashmla_sparse`, `flashmla_kv`, `flashmla_auto`, `fa3`, `tilelang`, `aiter`, `trtllm` | +| `--nsa-decode-backend` | Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding. | `fa3` | `flashmla_sparse`, `flashmla_kv`, `fa3`, `tilelang`, `aiter`, `trtllm` | | `--fp8-gemm-backend` | Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (optimal for Blackwell and low-latency), 'cutlass' (optimal for Hopper/Blackwell GPUs and high-throughput), 'triton' (fallback, widely compatible), 'aiter' (ROCm only). **NOTE**: This replaces the deprecated environment variables SGLANG_ENABLE_FLASHINFER_FP8_GEMM and SGLANG_SUPPORT_CUTLASS_BLOCK_FP8. | `auto` | `auto`, `deep_gemm`, `flashinfer_trtllm`, `cutlass`, `triton`, `aiter` | | `--fp4-gemm-backend` | Choose the runner backend for NVFP4 GEMM operations. Options: 'flashinfer_cutlass' (default), 'auto' (auto-selects between flashinfer_cudnn/flashinfer_cutlass based on CUDA/cuDNN version), 'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), 'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling). All backends are from FlashInfer; when FlashInfer is unavailable, sgl-kernel CUTLASS is used as an automatic fallback. **NOTE**: This replaces the deprecated environment variable SGLANG_FLASHINFER_FP4_GEMM_BACKEND. | `flashinfer_cutlass` | `auto`, `flashinfer_cudnn`, `flashinfer_cutlass`, `flashinfer_trtllm` | | `--disable-flashinfer-autotune` | Flashinfer autotune is enabled by default. Set this flag to disable the autotune. | `False` | bool flag (set to enable) | diff --git a/docs/basic_usage/deepseek_v32.md b/docs/basic_usage/deepseek_v32.md index 4389c2836..56f37f29c 100644 --- a/docs/basic_usage/deepseek_v32.md +++ b/docs/basic_usage/deepseek_v32.md @@ -66,9 +66,13 @@ python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --n - `fa3`: `flash_attn_with_kvcache` kernel from `flash_attn` library. Can only run on Hopper GPUs. It requires bf16 q, kv inputs. - `tilelang`: `tilelang` implementation that can run on GPU, HPU and NPU. - `aiter`: Aiter kernel on AMD HPUs. Can only be used as decode kernel. + - `trtllm`: `trtllm-mla` sparse kernel from flashinfer library. Only run on blackwell GPUs. It requires QKV bf16 or QKV fp8. - On the basis of performance benchmarks, the default configuration on H200 and B200 are set as follows : - H200: `flashmla_sparse` prefill attention (short-seq prefill uses MHA via FlashAttention varlen), `fa3` decode attention, `bf16` kv cache dtype. - B200: `flashmla_auto` prefill attention (short-seq prefill uses MHA via TRT-LLM ragged), `flashmla_kv` decode attention, `fp8_e4m3` kv cache dtype. `flashmla_auto` enables automatic selection of either `flashmla_sparse` or `flashmla_kv` kernel for prefill based on KV cache dtype, hardware, and heuristics. When FP8 KV cache is enabled and `total_kv_tokens < total_q_tokens * 512`, it uses the `flashmla_sparse` kernel; otherwise, it falls back to the `flashmla_kv` kernel. The heuristics may need to be tuned if the performance of either the `flashmla_sparse` or `flashmla_kv` kernel changes significantly. +- On Blackwell platform, with slightly accuracy drop, the performance can boost up to 3x-5x + - B200: by choosing `trtllm` for both `--nsa-prefill-backend` and `--nsa-decode-backend`, the prefill attention use MHA via TRT-LLM ragged for both short and long sequence (**accuracy impact**). Combine the `trtllm` with `fp8_e4m3` kv cache, the kv cache dim is `576` (kv_lora_rank + qk_rope_head_dim) (**accuracy impact**), compare to the combination of `flashmla_auto` and `fp8_e4m` kv cache dim is `656` (kv_lora_rank + scale storage (kv_lora_rank // quant_block_size * 4 bytes) + rope dimension storage). + ## Multi-token Prediction SGLang implements Multi-Token Prediction (MTP) for DeepSeek V3.2 based on [EAGLE speculative decoding](https://docs.sglang.io/advanced_features/speculative_decoding.html#EAGLE-Decoding). With this optimization, the decoding speed can be improved significantly on small batch sizes. Please look at [this PR](https://github.com/sgl-project/sglang/pull/11652) for more information. diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 57193d49f..1d2480027 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -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, diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index bebe41037..e5489ce7a 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -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) diff --git a/python/sglang/srt/layers/attention/utils.py b/python/sglang/srt/layers/attention/utils.py index 6cc08f72f..44d5edaaf 100644 --- a/python/sglang/srt/layers/attention/utils.py +++ b/python/sglang/srt/layers/attention/utils.py @@ -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) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index a5ad78b89..1d917137c 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -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__( diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 2f30f6a55..9852aaff3 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -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, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 35fe54e28..1583dd788 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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 ( diff --git a/test/registered/hicache/test_nsa_pool_host_unit.py b/test/registered/hicache/test_nsa_pool_host_unit.py index 7c1f53f93..c945b3684 100644 --- a/test/registered/hicache/test_nsa_pool_host_unit.py +++ b/test/registered/hicache/test_nsa_pool_host_unit.py @@ -50,6 +50,7 @@ class TestNSAHiCacheTransfer(unittest.TestCase): layer_num=layer_num, device="cuda", enable_memory_saver=False, + kv_cache_dim=576, index_head_dim=128, ) pin_memory = io_backend == "kernel" diff --git a/test/registered/kernels/test_nsa_indexer.py b/test/registered/kernels/test_nsa_indexer.py index 0a1e59654..834b1039c 100644 --- a/test/registered/kernels/test_nsa_indexer.py +++ b/test/registered/kernels/test_nsa_indexer.py @@ -232,6 +232,7 @@ class MockModelRunner: device=self.device, index_head_dim=self.config["index_head_dim"], enable_memory_saver=False, + kv_cache_dim=self.config["kv_lora_rank"] + self.config["qk_rope_head_dim"], ) # Required by backend with NSA-specific attributes