Skip to content

[Feature] support dpsk v32 flashinfer decode #15546

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
DarkSharpness wants to merge 2 commits into from
Closed

[Feature] support dpsk v32 flashinfer decode #15546

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d74fd2c
feat: support dpsk v32 flashinfer decode
DarkSharpness Dec 20, 2025
87cb459
Merge branch 'main' into dpsk_v32_fi
Fridge003 Jan 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 47 additions & 6 deletions python/sglang/srt/layers/attention/nsa_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -254,7 +254,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(
Expand Down Expand Up @@ -286,16 +288,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)
Expand All @@ -318,8 +322,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(
Expand All @@ -330,6 +334,7 @@ def __init__(
self.workspace_buffer = global_workspace_buffer
else:
self.workspace_buffer = None
print(f"{self.nsa_prefill_impl = } {self.nsa_decode_impl = }")

@gemini-code-assist gemini-code-assist Bot Dec 20, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This print statement appears to be for debugging. Please remove it before merging to keep the production logs clean.

@DarkSharpness DarkSharpness Dec 20, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

true

@Fridge003 Fridge003 Jan 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
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):
Expand Down Expand Up @@ -1433,6 +1438,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)

@Fridge003 Fridge003 Jan 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we apply _concat_mla_absorb_q_general(q_nope, q_rope), which might be faster

return self._forward_flashinfer(
q_all=q_all,
kv_cache=kv_cache,
page_table_1=page_table_1,
sm_scale=layer.scaling,
metadata=metadata,
)
Comment on lines +1444 to +1450

@gemini-code-assist gemini-code-assist Bot Dec 20, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To avoid hardcoding model-specific parameters in _forward_flashinfer, please pass the layer object to it. This will allow accessing the model configuration dynamically in the next step.

Suggested change
return self._forward_flashinfer(
q_all=q_all,
kv_cache=kv_cache,
page_table_1=page_table_1,
sm_scale=layer.scaling,
metadata=metadata,
)
return self._forward_flashinfer(
q_all=q_all,
kv_cache=kv_cache,
page_table_1=page_table_1,
sm_scale=layer.scaling,
metadata=metadata,
layer=layer,
)

elif self.nsa_decode_impl == "aiter":
if q_rope is not None:
q_all = torch.cat([q_nope, q_rope], dim=-1)
Expand Down Expand Up @@ -1647,6 +1662,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,
)
Comment on lines +1665 to +1689

@gemini-code-assist gemini-code-assist Bot Dec 20, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Instead of hardcoding values for qk_nope_head_dim, kv_lora_rank, and qk_rope_head_dim, please retrieve them from the layer.model_config. This makes the implementation more generic and avoids magic numbers. This change depends on passing the layer object from forward_decode.

    def _forward_flashinfer(
        self,
        q_all: torch.Tensor,
        kv_cache: torch.Tensor,
        page_table_1: torch.Tensor,
        sm_scale: float,
        metadata: NSAMetadata,
        layer: "RadixAttention",
    ) -> 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=layer.model_config.qk_nope_head_dim,
            kv_lora_rank=layer.model_config.kv_lora_rank,
            qk_rope_head_dim=layer.model_config.qk_rope_head_dim,
            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,
Expand Down
6 changes: 5 additions & 1 deletion python/sglang/srt/server_args.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1171,7 +1171,11 @@ def _handle_model_specific_adjustments(self):
# The default settings (P=flashmla_sparse, D=fa3) 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"
)

if self.enable_nsa_prefill_context_parallel:
assert (
Expand Down
Loading