diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index d5f564982..07e960938 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -3486,17 +3486,44 @@ class DeepseekV4AttnBackend( self.candidate_indexer.publish_decode(inputs, page_indices, raw_indices) ) return - logits = deep_gemm_fp4_paged_mqa_logits( - (q_fp4, q_sf), - k_cache, - weights, - metadata.compressed_seq_lens, - metadata.page_table, - metadata.deep_gemm_metadata, - metadata.max_compressed_seq_len, - ) - # TODO(dark): add bf16 topk - topk_transform_paged_from_metadata(logits, metadata, page_indices, raw_indices) + if isinstance(metadata.deep_gemm_metadata, list): + topk_plans = metadata.topk_metadata_chunks + assert not metadata.use_topk_v2 or topk_plans is not None + for chunk_idx, (rows, plan) in enumerate(metadata.row_chunks()): + logits = deep_gemm_fp4_paged_mqa_logits( + (q_fp4[rows], q_sf[rows]), + k_cache, + weights[rows], + metadata.compressed_seq_lens[rows], + metadata.page_table[rows], + plan, + metadata.max_compressed_seq_len, + ) + # TODO(dark): add bf16 topk + topk_transform_paged_from_metadata( + logits, + metadata, + page_indices, + raw_indices, + rows=rows, + topk_metadata=( + topk_plans[chunk_idx] if topk_plans is not None else None + ), + ) + else: + logits = deep_gemm_fp4_paged_mqa_logits( + (q_fp4, q_sf), + k_cache, + weights, + metadata.compressed_seq_lens, + metadata.page_table, + metadata.deep_gemm_metadata, + metadata.max_compressed_seq_len, + ) + # TODO(dark): add bf16 topk + topk_transform_paged_from_metadata( + logits, metadata, page_indices, raw_indices + ) # TODO(candidate): Hopper decode still publishes / consumes masks inline (torch # top-k); move into the candidate indexer with the prefill paths. diff --git a/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py b/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py index 41d4dd591..37f7cb591 100644 --- a/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py +++ b/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py @@ -21,6 +21,7 @@ from sglang.srt.layers.attention.dsv4.candidate_indexer import ( ) from sglang.srt.layers.attention.dsv4.indexer import ( deep_gemm_fp4_paged_mqa_logits, + topk_transform_paged_from_metadata, ) CANDIDATE_BLOCK_SIZE = 8 # positions per block; DeepGEMM accepts 8 or 16 @@ -176,6 +177,10 @@ class DeepGemmCandidateIndexer: metadata.""" metadata = inputs.metadata seq_lens = metadata.compressed_seq_lens.reshape(-1) + if isinstance(metadata.deep_gemm_metadata, list): + return self._publish_decode_chunked( + inputs, page_indices, raw_indices, seq_lens + ) logits = deep_gemm_fp4_paged_mqa_logits( (inputs.q_fp4, inputs.q_sf), inputs.k_cache, @@ -231,6 +236,83 @@ class DeepGemmCandidateIndexer: ready=ready, ) + def _publish_decode_chunked( + self, + inputs: IndexerInputs, + page_indices: torch.Tensor, + raw_indices: Optional[torch.Tensor], + seq_lens: torch.Tensor, + ) -> SparseBlockTable: + """Publish an eager forward whose dense logits are bounded by row chunks. + + CUDA-graph metadata always carries one tensor schedule and keeps using the + asynchronous fast path above. The exceptional eager path stays on the + current stream so each chunk's full logits can be released before the next. + """ + metadata = inputs.metadata + block_chunks = [] + phys_block_chunks = [] + valid_len_chunks = [] + topk_plans = metadata.topk_metadata_chunks + assert not metadata.use_topk_v2 or topk_plans is not None + + for chunk_idx, (rows, plan) in enumerate(metadata.row_chunks()): + logits = deep_gemm_fp4_paged_mqa_logits( + (inputs.q_fp4[rows], inputs.q_sf[rows]), + inputs.k_cache, + inputs.weights[rows], + metadata.compressed_seq_lens[rows], + metadata.page_table[rows], + plan, + metadata.max_compressed_seq_len, + ) + topk_transform_paged_from_metadata( + logits, + metadata, + page_indices, + raw_indices, + rows=rows, + topk_metadata=( + topk_plans[chunk_idx] if topk_plans is not None else None + ), + ) + + chunk_seq_lens = seq_lens[rows] + nblocks, row_valid_lens = candidate_row_lens( + chunk_seq_lens, self.topk_blocks + ) + blocks = amax_topk_blocks(logits, chunk_seq_lens, nblocks, self.topk_blocks) + phys_blocks = sort_candidate_blocks( + blocks, + chunk_seq_lens, + metadata.page_table[rows], + metadata.compressed_page_size, + ) + block_chunks.append(blocks) + phys_block_chunks.append(phys_blocks) + valid_len_chunks.append(row_valid_lens) + + blocks = torch.cat(block_chunks) + phys_blocks = torch.cat(phys_block_chunks) + row_valid_lens = torch.cat(valid_len_chunks) + schedule = build_sparse_indexer_schedule( + blocks, + seq_lens, + metadata.page_table, + metadata.compressed_page_size, + inputs.q_fp4.dtype, + self._request_ids(inputs.request_ids, inputs.num_rows, blocks.device), + ) + ready = torch.cuda.Event() + ready.record(torch.cuda.current_stream()) + return SparseBlockTable( + blocks=blocks, + schedule=schedule, + phys_blocks=phys_blocks, + valid_lens=row_valid_lens, + ready=ready, + ) + def _scores(self, table: SparseBlockTable, inputs: IndexerInputs) -> torch.Tensor: return sparse_logits( inputs.q_fp4, diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 28d521854..31b123070 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -475,27 +475,40 @@ def topk_transform_paged_from_metadata( metadata, page_indices: torch.Tensor, raw_indices: Optional[torch.Tensor] = None, + *, + rows: Optional[slice] = None, + topk_metadata: Optional[torch.Tensor] = None, ) -> None: """Pool slots into ``page_indices`` (``-1`` past the valid count) and, when given, positions into ``raw_indices``; ``metadata`` is a ``PagedIndexerMetadata``.""" + if rows is None: + seq_lens = metadata.compressed_seq_lens + page_table = metadata.page_table + out_page_indices = page_indices + out_raw_indices = raw_indices + else: + seq_lens = metadata.compressed_seq_lens[rows] + page_table = metadata.page_table[rows] + out_page_indices = page_indices[rows] + out_raw_indices = raw_indices[rows] if raw_indices is not None else None if metadata.use_topk_v2: topk_transform_paged_v2( logits, - metadata.compressed_seq_lens, - metadata.page_table, - page_indices, + seq_lens, + page_table, + out_page_indices, metadata.compressed_page_size, - metadata.topk_metadata, - raw_indices, + metadata.topk_metadata if topk_metadata is None else topk_metadata, + out_raw_indices, ) else: topk_transform_paged( logits, - metadata.compressed_seq_lens, - metadata.page_table, - page_indices, + seq_lens, + page_table, + out_page_indices, metadata.compressed_page_size, - raw_indices, + out_raw_indices, ) diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index da0bf0d84..019850d05 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -21,6 +21,7 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, ) +from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode from sglang.srt.utils import is_hip, is_sm120_supported, is_xpu logger = logging.getLogger(__name__) @@ -306,7 +307,8 @@ class PagedIndexerMetadata: ): return None if ( - torch.cuda.is_current_stream_capturing() + get_is_capture_mode() + or torch.cuda.is_current_stream_capturing() or is_in_breakable_cuda_graph() or is_in_tc_piecewise_cuda_graph() ): @@ -325,14 +327,27 @@ class PagedIndexerMetadata: def row_chunks(self): num_rows = self.compressed_seq_lens.shape[0] - if self.row_chunk <= 0: + if self.row_chunk > 0: + rows_per_chunk = self.row_chunk + elif isinstance(self.deep_gemm_metadata, list): + assert self.rows_per_chunk is not None, ( + "chunked DeepGEMM metadata requires rows_per_chunk" + ) + rows_per_chunk = self.rows_per_chunk + else: return [(slice(0, num_rows), self.deep_gemm_metadata)] - return [ - (slice(start, min(start + self.row_chunk, num_rows)), plan) + + chunks = [ + (slice(start, min(start + rows_per_chunk, num_rows)), plan) for start, plan in zip( - range(0, num_rows, self.row_chunk), self.deep_gemm_metadata + range(0, num_rows, rows_per_chunk), self.deep_gemm_metadata ) ] + assert chunks and chunks[-1][0].stop == num_rows, ( + f"chunk schedules do not cover all rows: {num_rows=} {rows_per_chunk=} " + f"{len(chunks)=}" + ) + return chunks def copy_(self, other: PagedIndexerMetadata): # A chunked schedule list has no in-place copy; rebind it instead. diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py index ad06b05e8..3dbb65fab 100644 --- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py +++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py @@ -752,6 +752,33 @@ class TestPagedIndexerMetadataChunking(CustomTestCase): row chunks the indexer loops over; a mismatch would silently score rows with another chunk's schedule.""" + def test_capture_warmup_skips_dynamic_budget_but_eager_forward_uses_it(self): + metadata = SimpleNamespace( + use_prefill_cuda_graph=False, + compressed_seq_lens=SimpleNamespace( + is_cuda=True, device=SimpleNamespace(index=0) + ), + max_compressed_seq_len=65536, + ) + for capture_mode in (True, False): + with ( + self.subTest(capture_mode=capture_mode), + patch(f"{_METADATA}.get_is_capture_mode", return_value=capture_mode), + patch("torch.cuda.is_current_stream_capturing", return_value=False), + patch(f"{_METADATA}.is_in_breakable_cuda_graph", return_value=False), + patch(f"{_METADATA}.is_in_tc_piecewise_cuda_graph", return_value=False), + patch( + f"{_METADATA}.mqa_logits_budget_bytes", return_value=4096 + ) as budget, + ): + result = PagedIndexerMetadata._mqa_logits_budget(metadata, num_rows=256) + if capture_mode: + self.assertIsNone(result) + budget.assert_not_called() + else: + self.assertEqual(result, 4096) + budget.assert_called_once_with(device_index=0, allow_sync=True) + def _build(self, *, num_rows: int, budget, use_topk_v2: bool): deep_gemm = SimpleNamespace( get_num_sms=MagicMock(return_value=1), @@ -806,6 +833,12 @@ class TestPagedIndexerMetadataChunking(CustomTestCase): self.assertIsInstance(metadata.deep_gemm_metadata, list) self.assertEqual(len(metadata.deep_gemm_metadata), len(chunks)) + metadata_chunks = metadata.row_chunks() + self.assertEqual([rows for rows, _ in metadata_chunks], chunks) + for (_, actual_plan), expected_plan in zip( + metadata_chunks, metadata.deep_gemm_metadata + ): + self.assertIs(actual_plan, expected_plan) schedule_rows = [ call.args[0] for call in deep_gemm.get_paged_mqa_logits_metadata.call_args_list @@ -875,6 +908,92 @@ class TestChunkedTopKMatchesUnchunked(CustomTestCase): self.assertTrue(torch.equal(run(rows_per_chunk), expected)) +class TestChunkedCandidatePublisher(CustomTestCase): + def test_each_deep_gemm_call_receives_one_tensor_schedule(self): + from sglang.srt.layers.attention.dsv4 import candidate_indexer_deep_gemm as mod + + num_rows, width = 5, 16 + chunks = [slice(0, 2), slice(2, 4), slice(4, 5)] + plans = [torch.tensor([i], dtype=torch.uint8) for i in range(len(chunks))] + topk_plans = [torch.tensor([i], dtype=torch.int32) for i in range(len(chunks))] + metadata = SimpleNamespace( + compressed_seq_lens=torch.full((num_rows, 1), width, dtype=torch.int32), + page_table=torch.zeros((num_rows, 1), dtype=torch.int32), + deep_gemm_metadata=plans, + max_compressed_seq_len=width, + compressed_page_size=64, + topk_metadata_chunks=topk_plans, + use_topk_v2=True, + row_chunks=lambda: list(zip(chunks, plans)), + ) + inputs = SimpleNamespace( + q_fp4=torch.zeros((num_rows, 1, 2, 64), dtype=torch.int8), + q_sf=torch.zeros((num_rows, 1, 2), dtype=torch.int32), + k_cache=torch.zeros((1, 64, 1, 68), dtype=torch.uint8), + weights=torch.zeros((num_rows, 2), dtype=torch.float32), + metadata=metadata, + request_ids=torch.arange(num_rows), + num_rows=num_rows, + ) + page_indices = torch.full((num_rows, 4), -1, dtype=torch.int32) + raw_indices = torch.full_like(page_indices, -1) + indexer = object.__new__(mod.DeepGemmCandidateIndexer) + indexer.topk_blocks = 2 + + deep_gemm = MagicMock( + side_effect=lambda q, *_args: torch.zeros( + (q[0].shape[0], width), dtype=torch.float32 + ) + ) + topk = MagicMock() + event = MagicMock() + stream = MagicMock() + with ( + patch.object(mod, "deep_gemm_fp4_paged_mqa_logits", deep_gemm), + patch.object(mod, "topk_transform_paged_from_metadata", topk), + patch.object( + mod, + "candidate_row_lens", + side_effect=lambda lens, _topk: ( + torch.ones_like(lens), + lens.clone(), + ), + ), + patch.object( + mod, + "amax_topk_blocks", + side_effect=lambda _logits, lens, _nblocks, topk_blocks: torch.zeros( + (lens.shape[0], topk_blocks), dtype=torch.int32 + ), + ), + patch.object( + mod, + "sort_candidate_blocks", + side_effect=lambda blocks, *_args: blocks + 1, + ), + patch.object( + mod, + "build_sparse_indexer_schedule", + return_value=torch.tensor([7], dtype=torch.uint8), + ), + patch.object(mod.torch.cuda, "Event", return_value=event), + patch.object(mod.torch.cuda, "current_stream", return_value=stream), + ): + table = indexer.publish_decode(inputs, page_indices, raw_indices) + + self.assertEqual(deep_gemm.call_count, len(chunks)) + for call, plan in zip(deep_gemm.call_args_list, plans): + self.assertIs(call.args[5], plan) + self.assertIsInstance(call.args[5], torch.Tensor) + self.assertEqual([call.kwargs["rows"] for call in topk.call_args_list], chunks) + for call, plan in zip(topk.call_args_list, topk_plans): + self.assertIs(call.kwargs["topk_metadata"], plan) + self.assertEqual(table.blocks.shape, (num_rows, indexer.topk_blocks)) + self.assertEqual(table.phys_blocks.shape, table.blocks.shape) + self.assertEqual(table.valid_lens.shape, (num_rows,)) + event.record.assert_called_once_with(stream) + + class TestCandidateIndexerGating(CustomTestCase): def test_candidate_indexer_gating(self): from sglang.srt.layers.attention.dsv4 import candidate_indexer