diff --git a/python/sglang/kernels/ops/attention/dsv4/topk.py b/python/sglang/kernels/ops/attention/dsv4/topk.py index 24c3bbd96394..4b9593ad6e21 100644 --- a/python/sglang/kernels/ops/attention/dsv4/topk.py +++ b/python/sglang/kernels/ops/attention/dsv4/topk.py @@ -126,30 +126,6 @@ def topk_transform_paged( ) -def topk_transform_paged_torch( - 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: - """The torch ``topk_transform_paged``: top-``k`` (``k = out_page_indices.shape[1]``) - of each row within its first ``seq_lens[b]`` columns, ascending, as pool slots - through ``page_tables`` and as positions when given; ``-1`` padded.""" - topk = out_page_indices.shape[1] - columns = torch.arange(scores.shape[1], device=seq_lens.device) - lens_c = seq_lens.unsqueeze(-1) - # Columns past a row's length hold garbage. - s = scores.masked_fill(columns[None, :] >= lens_c, -torch.inf) - idx = s.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values - reach = idx < lens_c - slots = page_tables.gather(-1, idx // page_size) * page_size + (idx % page_size) - out_page_indices.copy_(torch.where(reach, slots, -1).to(torch.int32)) - if out_raw_indices is not None: - out_raw_indices.copy_(torch.where(reach, idx, -1).to(torch.int32)) - - # metadata is (batch+1, 2) int32: row 0 = {cluster_threshold, num_cluster_items}; # rows 1..N = {batch_id, seq_len} of items routed to the persistent cluster pool. _PLAN_METADATA_INTS_PER_BATCH = 2 diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e88ab4262fc4..4dbce1c0f18e 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1741,9 +1741,6 @@ class Envs: SGLANG_OPT_USE_ONLINE_COMPRESS = EnvBool(False) SGLANG_EXPERIMENTAL_ONLINE_C128_MTP = EnvBool(False) SGLANG_DSV4_COMPRESS_STATE_DTYPE = EnvStr("float32") - # Run the DeepSeek-V4.1 ratio-1/2 prefill indexer on the torch path instead - # of the DeepGEMM dense fp4 logits kernel (test oracle / fallback). - SGLANG_DSV41_TORCH_PREFILL_INDEXER = EnvBool(False) # DeepSeek-V4.1 ratio-1/2 index layers: the DeepGEMM fp4 logits kernels store bf16 instead # of fp32 (the reference ranks bf16 scores); every consumer of those logits must take bf16. SGLANG_OPT_DSV41_BF16_INDEX_LOGITS = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index ab7b80facd7f..e7602f5b287f 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -83,7 +83,6 @@ CapturedPrefillInputs, DecodeInputs, PrefillInputs, - has_dense_fp4_indexer, is_sm100_or_newer, make_candidate_indexer, make_full_topk_indexer, @@ -1012,13 +1011,8 @@ def _pad_rows(t: torch.Tensor, *, num_rows: int) -> torch.Tensor: return torch.nn.functional.pad(t, (0, num_rows - t.shape[0])) -# Rows per logits chunk of the ratio-1/2 indexer metadata under the prefill CUDA -# graph; its width is _prefill_graph_context_width. -_PREFILL_GRAPH_INDEXER_ROW_CHUNK = 2048 - - -# The candidate window (2,048 blocks of 8 tokens): inside it the captured full -# top-k equals the two-level selection. +# The replayed context limit when the prefill graph's metadata is built at a fixed +# width (no two-level candidates); the candidate window, 2,048 blocks of 8 tokens. _PREFILL_GRAPH_DEFAULT_MAX_SEQ_LEN = 16 * 1024 @@ -1085,30 +1079,11 @@ def copy_(self, other: DSV4Metadata): ) self.sparse_prefill_cache = None - def refresh_for_breakable_cuda_graph_replay_( - self, - static_metadata: DSV4Metadata, - *, - two_level: bool, - ): + def refresh_for_breakable_cuda_graph_replay_(self, static_metadata: DSV4Metadata): self.core_attn_metadata.refresh_for_breakable_cuda_graph_replay_( static_metadata.core_attn_metadata ) maybe_copy_inplace(self.indexer_metadata, src=static_metadata.indexer_metadata) - if not two_level: - # The captured full top-k fallback reads the ratio-1/2 lengths and - # page tables by address: refresh them in place. - maybe_copy_inplace( - self.c1_indexer_metadata, src=static_metadata.c1_indexer_metadata - ) - maybe_copy_inplace( - self.c2_indexer_metadata, src=static_metadata.c2_indexer_metadata - ) - else: - # Two-level candidates: only the eager full top-k breaks read these, so - # they take the live objects; captured consumers read the bridge tensors. - self.c1_indexer_metadata = static_metadata.c1_indexer_metadata - self.c2_indexer_metadata = static_metadata.c2_indexer_metadata maybe_copy_inplace( self.low_ratio_req_indices, src=static_metadata.low_ratio_req_indices ) @@ -1607,59 +1582,20 @@ def create(compress_ratio: Literal[4, 128]): ), ) if use_prefill_cuda_graph and self.low_ratio_prefill_graph: - low = core_attn_metadata.low_ratios - metadata.c1_indexer_metadata = ( - self._low_ratio_prefill_indexer_metadata(core_attn_metadata, 1) - if 1 in low - else None - ) - metadata.c2_indexer_metadata = ( - self._low_ratio_prefill_indexer_metadata(core_attn_metadata, 2) - if 2 in low - else None - ) metadata.low_ratio_req_indices = req_pool_indices_repeated.to(torch.int64) metadata.low_ratio_pos_i64 = core_attn_metadata.positions_casual.to( torch.int64 ) return metadata - def _low_ratio_prefill_indexer_metadata( - self, core: DSV4AttnMetadata, compress_ratio: int - ) -> PagedIndexerMetadata: - num_pages = core.page_table.shape[1] - max_seq_len = self._prefill_graph_seq_limit() - if max_seq_len is not None: - num_pages = min(max_seq_len // self.page_size, num_pages) - index_page_size = self.token_to_kv_pool.get_index_k_page_size(compress_ratio) - page_table = expand_index_page_table( - core.page_table[:, :num_pages], - full_page_size=self.page_size, - compress_ratio=compress_ratio, - index_page_size=index_page_size, - ) - # Unclamped: a token with no completed group scores nothing, as in eager. - c_seq_lens = (core.seq_lens_casual // compress_ratio).to(torch.int32) - row_chunk = _PREFILL_GRAPH_INDEXER_ROW_CHUNK - return PagedIndexerMetadata( - page_size=self.page_size, - compressed_page_size=index_page_size, - page_table=page_table, - compressed_seq_lens=c_seq_lens, - use_topk_v2=False, - use_prefill_cuda_graph=True, - compress_ratio=compress_ratio, - row_chunk=row_chunk if row_chunk < c_seq_lens.shape[0] else 0, - ) - @property def low_ratio_prefill_graph(self) -> bool: - """Whether ratio-1/2 sources use captured projections and indexer metadata.""" - return ( - bool(self.low_ratios) - and has_dense_fp4_indexer() - and is_sm100_or_newer() - and get_parallel().attn_cp_size == 1 + """Whether ratio-1/2 sources run in the prefill graph with two-level candidates: + the full top-k layers score at eager breaks, and the consumers score the + published candidate blocks in the graph. Otherwise the sources are one eager + break.""" + return bool(self.low_ratios) and isinstance( + self.prefill_candidates, SparseTableBackend ) def can_run_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool: @@ -3128,10 +3064,7 @@ def prepare_forward_metadata_for_breakable_cuda_graph_replay( use_prefill_cuda_graph=True, ) assert isinstance(capture_metadata, DSV4Metadata) - capture_metadata.refresh_for_breakable_cuda_graph_replay_( - static_metadata, - two_level=self._prefill_graph_two_level(), - ) + capture_metadata.refresh_for_breakable_cuda_graph_replay_(static_metadata) self.prefill_graph_replay_batch = forward_batch self.forward_metadata = capture_metadata self._activate_request_window() @@ -3156,10 +3089,7 @@ def _prefill_graph_tail_metadata( if held is None: self._bcg_tail_metadata[rows] = fresh return fresh - held.refresh_for_breakable_cuda_graph_replay_( - fresh, - two_level=self._prefill_graph_two_level(), - ) + held.refresh_for_breakable_cuda_graph_replay_(fresh) held.core_attn_metadata.swa_out_cache_loc = ( fresh.core_attn_metadata.swa_out_cache_loc ) @@ -3302,11 +3232,10 @@ def forward_low_ratio_sources( and self._low_ratio_in_prefill_graph() ): bufs = self._source_projection_buffers(x.shape[0], layer.compress_ratio) - two_level = self._prefill_graph_two_level() captured_index = ( run_indexer and layer.indexer is not None - and (layer.indexer.uses_candidates or not two_level) + and layer.indexer.uses_candidates ) _bcg_low_ratio_source_projections( layer, x, q_lora, pos, bufs, run_compressor, captured_index @@ -3761,35 +3690,19 @@ def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> Non ratio = layer.compress_ratio indexer = layer.indexer core = self.forward_metadata.core_metadata - metadata = ( - self.forward_metadata.c1_indexer_metadata - if ratio == 1 - else self.forward_metadata.c2_indexer_metadata - ) - assert metadata is not None, f"no prefill graph indexer metadata for {ratio = }" inputs = CapturedPrefillInputs( indexer=indexer, layer_id=layer.layer_id, q=projected_q, weights=projected_w, - paged_metadata=metadata, + index_page_size=self.token_to_kv_pool.get_index_k_page_size(ratio), out_page_indices=core.sparse_page_indices(ratio), out_raw_indices=core.sparse_raw_indices(ratio), ) - # At any width: inside the candidate window the published table is the - # identity, and past it the full top-k would be wrong. - if indexer.uses_candidates and self._prefill_graph_two_level(): - self.prefill_candidates.consume_prefill_captured( - inputs, self.forward_metadata.candidate_metadata - ) - return - if indexer.uses_candidates or indexer.is_candidate_source: - # Inside the window every block is a candidate: the full top-k applies. - width = metadata.max_compressed_seq_len - assert ( - width <= indexer.candidate_topk_blocks * indexer.candidate_block_size - ), f"prefill graph indexer width {width} exceeds the candidate window" - self.full_topk_indexer.topk_prefill_captured(inputs) + # At any width: the consumer scores only the blocks the source published. + self.prefill_candidates.consume_prefill_captured( + inputs, self.forward_metadata.candidate_metadata + ) def _make_low_ratio_decode_indexer_inputs(self, layer, x, q_lora, req, pos, mode): ratio = layer.compress_ratio @@ -3888,25 +3801,15 @@ def _publish_candidate_metadata(self, published: Optional[CandidateMetadata]): return self.forward_metadata.candidate_metadata = published - def _prefill_graph_two_level(self) -> bool: - """Whether the full top-k layers score at eager breaks and the consumers score - the published candidate blocks inside the prefill graph.""" - return isinstance(self.prefill_candidates, SparseTableBackend) - def _prefill_graph_replays_any_context(self) -> bool: """Whether a prefill graph replay admits any context length: with two-level candidates only c4/c128 compress metadata would be captured at a width.""" - return ( - self.low_ratio_prefill_graph - and self._prefill_graph_two_level() - and not self.has_c4 - and not self.has_c128 - ) + return self.low_ratio_prefill_graph and not self.has_c4 and not self.has_c128 @property def prefill_cuda_graph_default_max_seq_len(self) -> Optional[int]: - """DeepSeek-V4.1's captured indexer width; prefill context parallelism builds - its own metadata and needs none.""" + """DeepSeek-V4.1's fixed metadata width without two-level candidates; prefill + context parallelism builds its own metadata and needs none.""" if ( self.low_ratios != (1, 2) or get_parallel().attn_cp_size > 1 diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 3c1e85e7d800..f6ce39ae12f9 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -889,7 +889,8 @@ class DeepseekV4HipRadixBackend( forward_low_ratio_sources = DeepseekV4AttnBackend.forward_low_ratio_sources _forward_low_ratio_sources_cp = DeepseekV4AttnBackend._forward_low_ratio_sources_cp write_tail_source_kv = DeepseekV4AttnBackend.write_tail_source_kv - low_ratio_prefill_graph = DeepseekV4AttnBackend.low_ratio_prefill_graph + # No two-level candidates on HIP: the ratio-1/2 sources stay one eager break. + low_ratio_prefill_graph = False # This attribute is used in python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py. # PrefillCudaGraphRunner looks this up via getattr; contexts over prefill.max_seq_len run eager. can_run_prefill_cuda_graph = DeepseekV4AttnBackend.can_run_prefill_cuda_graph @@ -3491,11 +3492,6 @@ def _low_ratio_index_topk( self, layer, x, q_lora, pos, forward_batch=forward_batch ) elif forward_batch.forward_mode.is_extend(): - if envs.SGLANG_DSV41_TORCH_PREFILL_INDEXER.get(): - raise NotImplementedError( - "SGLANG_DSV41_TORCH_PREFILL_INDEXER (the torch prefill indexer " - "oracle) is not supported on HIP" - ) # HIP rejects V4.1 prefill CP, the only caller passing rows_per_request assert rows_per_request is None, "V4.1 prefill CP is not supported on HIP" assert ( diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index af39d75f4a98..77d2dcfc232b 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -185,9 +185,6 @@ class PagedIndexerMetadata: use_prefill_cuda_graph: bool = False # Indexer source compression ratio: 4 for c4, 1 or 2 for the dsv41 sources. compress_ratio: int = 4 - # Rows per logits chunk for the prefill CUDA graph low-ratio indexer; 0 plans - # all rows at once. - row_chunk: int = 0 rows_per_request: int = 1 # A list when the dynamic-budget forward is row-chunked: one schedule per chunk. deep_gemm_metadata: Any = field(init=False, repr=False) @@ -254,18 +251,7 @@ def __post_init__(self): self.rows_per_chunk, self.mqa_logits_budget_bytes, ) - if self.row_chunk > 0: - self.deep_gemm_metadata = torch.stack( - [ - get_paged_mqa_logits_metadata( - compressed_seq_lens[_s : _s + self.row_chunk], - self.compressed_page_size, - deep_gemm.get_num_sms(), - ) - for _s in range(0, compressed_seq_lens.shape[0], self.row_chunk) - ] - ) - elif self.rows_per_chunk is not None: + if self.rows_per_chunk is not None: # Chunk metadata is shared by all indexer layers in this forward. self.deep_gemm_metadata = [ get_paged_mqa_logits_metadata( @@ -339,9 +325,7 @@ def max_compressed_seq_len(self) -> int: def row_chunks(self): num_rows = self.compressed_seq_lens.shape[0] - if self.row_chunk > 0: - rows_per_chunk = self.row_chunk - elif isinstance(self.deep_gemm_metadata, list): + if isinstance(self.deep_gemm_metadata, list): assert self.rows_per_chunk is not None, ( "chunked DeepGEMM metadata requires rows_per_chunk" ) @@ -385,7 +369,6 @@ def copy_(self, other: PagedIndexerMetadata): "page_size", "compressed_page_size", "compress_ratio", - "row_chunk", "force_deep_gemm_metadata", "use_prefill_cuda_graph", "use_topk_v2", diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py index 59e9e54ac681..c2ecd658cd50 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py @@ -59,18 +59,7 @@ def has_dense_fp4_indexer() -> bool: @functools.cache def _use_deep_gemm_prefill() -> bool: - # A CP rank scores rank-local rows, which only DeepGEMM serves, so CP ignores the env. - from sglang.srt.environ import envs - from sglang.srt.runtime_context import get_parallel - - return ( - is_sm100_or_newer() - and has_dense_fp4_indexer() - and ( - get_parallel().attn_cp_size > 1 - or not envs.SGLANG_DSV41_TORCH_PREFILL_INDEXER.get() - ) - ) + return is_sm100_or_newer() and has_dense_fp4_indexer() def make_full_topk_indexer( @@ -127,8 +116,8 @@ def make_candidate_indexer( candidate_topk_blocks=candidate_topk_blocks, candidate_block_size=candidate_block_size, ) - # A CP rank's rows are an interleaved subset of the batch; the torch prefill - # env keeps the sparse table on decode only. + # A CP rank's rows are an interleaved subset of the batch, and without the dense + # fp4 indexer prefill scores on torch: both keep the sparse table on decode only. if not use_deep_gemm_prefill or get_parallel().attn_cp_size > 1: return dense_blocks, sparse_table return sparse_table, sparse_table diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py index ec0a973928e2..9414866007af 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py @@ -14,7 +14,6 @@ from sglang.kernels.ops.attention.dsv4.index_logits import ( deep_gemm_fp4_paged_mqa_logits, ) -from sglang.kernels.ops.attention.dsv4.topk import topk_transform_paged_torch from sglang.srt.layers.attention.dsv4.indexer import topk_transform_paged_from_metadata from .scoring import ( @@ -22,14 +21,11 @@ dense_prefill_topk, get_deep_gemm_decode_data, get_deep_gemm_prefill_data, - get_index_k_cache, prefill_requests, - quantize_index_q, select_decode, write_prefill, ) from .types import ( - CapturedPrefillInputs, DecodeInputs, PrefillInputs, ) @@ -61,9 +57,6 @@ def topk_prefill(self, inputs: PrefillInputs) -> None: else: self._torch_prefill(inputs) - def topk_prefill_captured(self, inputs: CapturedPrefillInputs) -> None: - self._deep_gemm_prefill_captured(inputs) - def topk_decode(self, inputs: DecodeInputs) -> None: if self.use_deep_gemm_decode: self._deep_gemm_decode(inputs) @@ -156,42 +149,3 @@ def _torch_decode(self, inputs: DecodeInputs) -> None: ) return select_decode(inputs, d, inputs.indexer.index_topk) - - def _deep_gemm_prefill_captured(self, inputs: CapturedPrefillInputs) -> None: - indexer = inputs.indexer - metadata = inputs.paged_metadata - assert indexer.n_local_heads == indexer.n_heads - q_fp4, q_sf = quantize_index_q(inputs.q) - num_tokens, num_heads = q_fp4.shape[0], q_fp4.shape[1] - q_fp4 = q_fp4.view(num_tokens, 1, num_heads, 64) - q_sf = q_sf.view(num_tokens, 1, num_heads) - weights = inputs.weights.float() - page_size = metadata.compressed_page_size - k_cache = get_index_k_cache( - token_to_kv_pool=self.token_to_kv_pool, - layer_id=inputs.layer_id, - page_size=page_size, - ) - - width = metadata.max_compressed_seq_len - lens = metadata.compressed_seq_lens - page_table = metadata.page_table - topk = min(indexer.index_topk, width) - for rows, plan in metadata.row_chunks(): - logits = deep_gemm_fp4_paged_mqa_logits( - (q_fp4[rows], q_sf[rows]), - k_cache, - weights[rows], - lens[rows], - page_table[rows], - plan, - width, - ) - topk_transform_paged_torch( - logits, - lens[rows], - page_table[rows], - inputs.out_page_indices[rows, :topk], - page_size, - inputs.out_raw_indices[rows, :topk], - ) diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py index 34a48c314c27..07a1e9ed3ba1 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py @@ -200,7 +200,7 @@ def consume_prefill_captured( k_cache = get_index_k_cache( token_to_kv_pool=self.token_to_kv_pool, layer_id=inputs.layer_id, - page_size=inputs.paged_metadata.compressed_page_size, + page_size=inputs.index_page_size, ) logits = _prefill_table_logits(published, q_fp4, q_sf, inputs.weights, k_cache) topk = min(indexer.index_topk, published.blocks.shape[1] * CANDIDATE_BLOCK_SIZE) diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py index a36df1b43261..3dd03d51b98e 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py @@ -104,7 +104,7 @@ class CapturedPrefillInputs(msgspec.Struct, frozen=True, kw_only=True): layer_id: int q: torch.Tensor # [rows, heads, 128] bf16, roped weights: torch.Tensor # [rows, heads] bf16 head weights - paged_metadata: PagedIndexerMetadata + index_page_size: int # index-K pool page size at this ratio # The selection as pool slots and as request-local compressed positions. out_page_indices: torch.Tensor out_raw_indices: torch.Tensor