From f7986c8603f7f665eb48e280c3b1fb0228f71e45 Mon Sep 17 00:00:00 2001 From: Zhihao Wang <101526713+xiuhu17@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:45:44 -0700 Subject: [PATCH] [RL] DSV4: dispatch indexer topk_transform_512 through DSATopKBackend (#31087) Signed-off-by: zhihaow6 --- .../layers/attention/deepseek_v4_backend.py | 4 + .../srt/layers/attention/dsv4/indexer.py | 101 ++++++++++++++++-- 2 files changed, 96 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 51b7a091d..fadc8c1a4 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -40,6 +40,7 @@ from sglang.kernels.ops.speculative.dspark.dspark_attn_metadata import ( ) from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend from sglang.srt.layers.attention.dsv4.compressor_v2 import ( CompressorBackendMixin, FusedCompressMetadata, @@ -528,6 +529,9 @@ class DeepseekV4AttnBackend( self.enable_deepseek_v4_fp4_indexer: bool = ( model_runner.server_args.enable_deepseek_v4_fp4_indexer ) + self.dsa_topk_backend: DSATopKBackend = DSATopKBackend( + model_runner.server_args.dsa_topk_backend + ) self.topk = model_runner.server_args.speculative_eagle_topk or 0 assert self.topk in [0, 1], "MTP Topk > 1 not supported for DeepSeek V4" self.mtp_enabled = self.topk > 0 diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index c5bc40368..76f5954e0 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -1,6 +1,16 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, List, Optional, Tuple, TypeAlias, Union +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Optional, + Tuple, + TypeAlias, + Union, +) import torch import torch.nn as nn @@ -15,6 +25,7 @@ from sglang.kernels.ops.attention.dsv4 import ( from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.environ import envs +from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend from sglang.srt.layers.attention.dsv4.compressor import Compressor from sglang.srt.layers.attention.dsv4.metadata import ( NonPagedIndexerPlan, @@ -245,18 +256,17 @@ def fp8_paged_mqa_logits_torch_sm120( return logits -def topk_transform_512_pytorch_vectorized( +def _topk_transform_512_vectorized( scores: torch.Tensor, seq_lens: torch.Tensor, page_tables: torch.Tensor, out_page_indices: torch.Tensor, page_size: int, out_raw_indices: Optional[torch.Tensor] = None, + topk_op: Callable[..., Tuple[torch.Tensor, torch.Tensor]] = torch.topk, + topk_op_kwargs: Optional[Dict[str, object]] = None, + contiguous_topk_input: bool = False, ) -> None: - """Vectorized PyTorch fallback for topk_transform_512. - All helper tensors (arange, zeros) are cached to avoid device-tensor - creation during HIP/CUDA graph capture.""" - TOPK = out_page_indices.shape[1] batch_size = scores.shape[0] max_seq_len = scores.shape[1] @@ -283,9 +293,13 @@ def topk_transform_512_pytorch_vectorized( masked_scores.masked_fill_(~valid_mask, float("-inf")) actual_k = min(TOPK, max_seq_len) - _, raw_indices = torch.topk( - masked_scores, k=actual_k, dim=1, largest=True, sorted=False + topk_kwargs = ( + {"dim": 1, "largest": True, "sorted": False} + if topk_op_kwargs is None + else topk_op_kwargs ) + topk_input = masked_scores.contiguous() if contiguous_topk_input else masked_scores + _, raw_indices = topk_op(topk_input, actual_k, **topk_kwargs) raw_indices = raw_indices.to(torch.int32) if actual_k < TOPK: @@ -330,10 +344,67 @@ def topk_transform_512_pytorch_vectorized( out_raw_indices.copy_(raw_indices) +def topk_transform_512_pytorch_vectorized( + scores: torch.Tensor, + seq_lens: torch.Tensor, + page_tables: torch.Tensor, + out_page_indices: torch.Tensor, + page_size: int, + out_raw_indices: Optional[torch.Tensor] = None, +) -> None: + """Vectorized PyTorch fallback for topk_transform_512. + All helper tensors (arange, zeros) are cached to avoid device-tensor + creation during HIP/CUDA graph capture.""" + + _topk_transform_512_vectorized( + scores, + seq_lens, + page_tables, + out_page_indices, + page_size, + out_raw_indices, + topk_op=torch.topk, + topk_op_kwargs={"dim": 1, "largest": True, "sorted": False}, + ) + + +def topk_transform_512_flashinfer_unfused( + scores: torch.Tensor, + seq_lens: torch.Tensor, + page_tables: torch.Tensor, + out_page_indices: torch.Tensor, + page_size: int, + out_raw_indices: Optional[torch.Tensor] = None, +) -> None: + import flashinfer + + from sglang.srt.layers.attention.dsa.dsa_topk_backend import ( + _flashinfer_tie_break_value, + ) + + _topk_transform_512_vectorized( + scores, + seq_lens, + page_tables, + out_page_indices, + page_size, + out_raw_indices, + topk_op=flashinfer.top_k, + topk_op_kwargs={ + "sorted": False, + "deterministic": envs.SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC.get(), + "tie_break": _flashinfer_tie_break_value(), + "dsa_graph_safe": True, + }, + contiguous_topk_input=True, + ) + + class C4IndexerBackendMixin: def __init__(self): super().__init__() self.debug_use_external_c4_sparse_indices: bool = False + self.dsa_topk_backend: DSATopKBackend = DSATopKBackend.SGL_KERNEL def _forward_prepare_multi_stream( self, @@ -732,7 +803,10 @@ class C4IndexerBackendMixin: elif core_metadata.c4_sparse_raw_indices is not None: raw_indices = core_metadata.c4_sparse_raw_indices - if envs.SGLANG_TOPK_TRANSFORM_512_TORCH.get(): + if ( + envs.SGLANG_TOPK_TRANSFORM_512_TORCH.get() + or self.dsa_topk_backend.is_torch() + ): topk_transform_512_pytorch_vectorized( logits, c4_seq_lens, @@ -741,6 +815,15 @@ class C4IndexerBackendMixin: indexer_metadata.c4_page_size, raw_indices, ) + elif self.dsa_topk_backend.is_flashinfer(): + topk_transform_512_flashinfer_unfused( + logits, + c4_seq_lens, + page_table, + c4_sparse_page_indices, + indexer_metadata.c4_page_size, + raw_indices, + ) elif envs.SGLANG_OPT_USE_TOPK_V2.get() and raw_indices is None: topk_transform_512_v2( logits,