From d74fd2cc5e1c7dc28ec03f7bdf16e98679053171 Mon Sep 17 00:00:00 2001 From: DarkSharpness <2040703891@qq.com> Date: Sun, 21 Dec 2025 03:41:02 +0800 Subject: [PATCH] feat: support dpsk v32 flashinfer decode --- .../srt/layers/attention/nsa_backend.py | 53 ++++++++++++++++--- python/sglang/srt/layers/quantization/fp8.py | 1 - python/sglang/srt/server_args.py | 6 ++- 3 files changed, 52 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index a28ad66dcf50..2a54930d77fb 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -234,7 +234,9 @@ def topk_transform( assert False, f"Unsupported {self.topk_transform_method = }" -_NSA_IMPL_T: TypeAlias = Literal["flashmla_sparse", "flashmla_kv", "fa3", "tilelang"] +NSA_IMPL_T: TypeAlias = Literal[ + "flashmla_sparse", "flashmla_kv", "fa3", "tilelang", "flashinfer" +] class NativeSparseAttnBackend( @@ -266,16 +268,18 @@ def __init__( self.num_q_heads = ( model_runner.model_config.num_attention_heads // get_attention_tp_size() ) - self.kv_cache_dim = model_runner.token_to_kv_pool.kv_cache_dim + self.kv_cache_dim: int = model_runner.token_to_kv_pool.kv_cache_dim assert model_runner.req_to_token_pool is not None self.req_to_token = model_runner.req_to_token_pool.req_to_token self.use_mha: bool = False - self.nsa_prefill_impl: _NSA_IMPL_T = ( + self.nsa_prefill_impl: NSA_IMPL_T = ( # type: ignore model_runner.server_args.nsa_prefill_backend ) - self.nsa_decode_impl: _NSA_IMPL_T = model_runner.server_args.nsa_decode_backend + self.nsa_decode_impl: NSA_IMPL_T = ( # type: ignore + model_runner.server_args.nsa_decode_backend + ) self.enable_auto_select_prefill_impl = self.nsa_prefill_impl == "flashmla_auto" self._arange_buf = torch.arange(16384, device=self.device, dtype=torch.int32) @@ -298,8 +302,8 @@ def __init__( self.device_capability = torch.cuda.get_device_capability() self.device_sm_major = self.device_capability[0] - # Allocate global workspace buffer for TRTLLm ragged attention kernel (SM100/B200) - if self.device_sm_major >= 10: + # Allocate global workspace buffer for flashinfer + if self.device_sm_major >= 10 or self.nsa_decode_impl == "flashinfer": global global_workspace_buffer if global_workspace_buffer is None: global_workspace_buffer = torch.empty( @@ -310,6 +314,7 @@ def __init__( self.workspace_buffer = global_workspace_buffer else: self.workspace_buffer = None + print(f"{self.nsa_prefill_impl = } {self.nsa_decode_impl = }") def get_device_int32_arange(self, l: int) -> torch.Tensor: if l > len(self._arange_buf): @@ -1312,6 +1317,16 @@ def forward_decode( logit_cap=layer.logit_cap, page_size=1, ) + elif self.nsa_decode_impl == "flashinfer": + if q_rope is not None: + q_all = torch.cat([q_nope, q_rope], dim=-1) + return self._forward_flashinfer( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + metadata=metadata, + ) elif self.nsa_decode_impl == "aiter": if q_rope is not None: q_all = torch.cat([q_nope, q_rope], dim=-1) @@ -1526,6 +1541,32 @@ def _forward_standard_mha( ver=fa_version, ) + def _forward_flashinfer( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + page_table_1: torch.Tensor, + sm_scale: float, + metadata: NSAMetadata, + ) -> torch.Tensor: + import flashinfer + + assert self.workspace_buffer is not None + return flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( + query=q_all.unsqueeze(1), # TODO(dark): support MTP + kv_cache=kv_cache.view(-1, 1, self.real_page_size, self.kv_cache_dim), + workspace_buffer=self.workspace_buffer, + qk_nope_head_dim=128, + kv_lora_rank=512, + qk_rope_head_dim=64, + block_tables=page_table_1.unsqueeze(1), # NOTE: 1 is MTP length + seq_lens=metadata.nsa_seqlens_expanded, + max_seq_len=metadata.nsa_max_seqlen_q, + sparse_mla_top_k=self.nsa_index_topk, + bmm1_scale=sm_scale, + enable_pdl=True, + ) + def _forward_tilelang( self, q_all: torch.Tensor, diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index e7f2b2a2853f..c47fb8b5c028 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -1389,7 +1389,6 @@ def apply_with_router_logits( if routed_scaling_factor is not None else 1.0 ), - tile_tokens_dim=None, routing_method_type=routing_method_type, use_shuffled_weight=False, tune_max_num_tokens=next_power_of_2(a_q.shape[0]), diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 553cb1c49fce..53c3d478699a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1074,7 +1074,11 @@ def _handle_model_specific_adjustments(self): # set prefill/decode backends for Blackwell. The default settings are for Hopper. if major >= 10: self.nsa_prefill_backend = "flashmla_sparse" - self.nsa_decode_backend = "flashmla_sparse" + self.nsa_decode_backend = ( + "flashinfer" + if self.kv_cache_dtype == "bfloat16" + else "flashmla_sparse" + ) # Logging env vars for NSA from sglang.srt.layers.attention.nsa.utils import (