From 8bdd8d847dcbdafd80e666ad7fb797a18da62975 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Fri, 2 Oct 2026 20:05:27 -0400 Subject: [PATCH] [GLM 5.3 Perf] Enable fused multi-step decode, 13.3% E2E throughput improvement for concurrency 1 (#57443) Signed-off-by: yewentao256 --- .../test_indexer_deepseek_v4_slot_mapping.py | 68 +++++++++++++++++++ .../attention/test_kpool_tail_slot_mapping.py | 10 +++ .../test_eagle_draft_attn_metadata.py | 2 + .../test_gpu_autoregressive_speculator.py | 1 + .../backends/mla/flashattn_mla_sparse.py | 4 ++ vllm/v1/attention/backends/mla/indexer.py | 63 +++++++++++++++++ vllm/v1/worker/gpu/spec_decode/speculator.py | 1 + 7 files changed, 149 insertions(+) diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 5bb6eb4c32..8acfc0a334 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -23,6 +23,8 @@ from vllm.v1.attention.backends.mla.compressor_utils import ( from vllm.v1.attention.backends.mla.indexer import ( BuildPrefillChunkMetadataKernel, DeepseekV4IndexerBackend, + DeepSeekV32IndexerDecodeMetadata, + DeepseekV32IndexerMetadata, DeepseekV32IndexerMetadataBuilder, DeepseekV41IndexerBackend, get_max_prefill_buffer_size, @@ -99,6 +101,72 @@ def test_fused_indexer_decode_metadata(query_lens, padding): torch.testing.assert_close(per_req[:reqs], lengths) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_indexer_draft_decode_metadata_update_is_capture_safe(): + if not indexer_module.has_deep_gemm(): + pytest.skip("requires DeepGEMM") + device = torch.device("cuda") + num_reqs = 3 + seq_lens = torch.tensor([17, 33, 65], dtype=torch.int32, device=device) + decode_seq_lens = torch.zeros((num_reqs, 1), dtype=torch.int32, device=device) + block_table = torch.arange(num_reqs * 4, dtype=torch.int32, device=device).view( + num_reqs, 4 + ) + num_sms = torch.cuda.get_device_properties(device).multi_processor_count + schedule_metadata = indexer_module.get_paged_mqa_logits_metadata( + torch.div(seq_lens, 16, rounding_mode="floor").unsqueeze(1), + 64, + num_sms, + ).clone() + + builder = object.__new__(DeepseekV32IndexerMetadataBuilder) + builder.dcp_world_size = 1 + builder.compress_ratio = 16 + builder.num_sms = num_sms + builder.arange_buffer = torch.arange(num_reqs + 1, dtype=torch.int32, device=device) + builder.kv_cache_spec = SimpleNamespace(num_states=64) + metadata = DeepseekV32IndexerMetadata( + seq_lens=seq_lens, + max_seq_len=128, + slot_mapping=torch.zeros(num_reqs, dtype=torch.int64, device=device), + num_decodes=num_reqs, + num_decode_tokens=num_reqs, + num_prefills=0, + num_prefill_tokens=0, + decode=DeepSeekV32IndexerDecodeMetadata( + block_table=block_table, + seq_lens=decode_seq_lens, + decode_lens=torch.zeros(num_reqs, dtype=torch.int32, device=device), + requires_padding=False, + schedule_metadata=schedule_metadata, + ), + ) + + builder.update_draft_decode_metadata(metadata) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + builder.update_draft_decode_metadata(metadata) + + seq_lens.copy_(torch.tensor([32, 48, 80], dtype=torch.int32, device=device)) + graph.replay() + torch.accelerator.synchronize() + + expected_decode_seq_lens = torch.div(seq_lens, 16, rounding_mode="floor") + torch.testing.assert_close(decode_seq_lens.flatten(), expected_decode_seq_lens) + torch.testing.assert_close( + schedule_metadata, + indexer_module.get_paged_mqa_logits_metadata( + expected_decode_seq_lens.unsqueeze(1), 64, num_sms + ), + ) + assert metadata.decode is not None + assert torch.all(metadata.decode.decode_lens == 1) + expected_slot_mapping = block_table[:, 0].to(torch.int64) * 64 + torch.tensor( + [1, 2, 4], dtype=torch.int64, device=device + ) + torch.testing.assert_close(metadata.slot_mapping, expected_slot_mapping) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize( "query_lens", [[1], [0, 257, 1, 0, 3], [6] * 64, [2, 6, 4, 0], [0, 0, 0]] diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index 157366e49c..da25eaefb2 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -300,6 +300,16 @@ def test_builder_build_falls_back_without_positions(): assert meta.slot_mapping is cam.slot_mapping +def test_builder_updates_draft_mapping(): + builder = make_tail_builder() + cam = make_common_metadata([[15], [16]], [5, 9]) + meta = KpoolTailMetadataBuilder.build(builder, 0, cam) + assert cam.positions is not None + cam.positions.add_(1) + builder.update_draft_decode_metadata(meta) + assert meta.slot_mapping[:2].tolist() == [5 * KPOOL, 9 * KPOOL + 1] + + def test_builder_reuses_slot_mapping_storage(): builder = make_tail_builder() first = make_common_metadata([list(range(10))], [5]) diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py index aa6a65661d..57831fe4d3 100644 --- a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py +++ b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py @@ -39,6 +39,7 @@ def _make_fake_speculator( and a draft model. """ fake_input_buffers = SimpleNamespace( + positions=torch.arange(max_num_tokens, dtype=torch.int64), query_start_loc=torch.zeros(max_num_reqs + 1, dtype=torch.int32), seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32), dcp_local_seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32), @@ -113,6 +114,7 @@ def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound(): captured = _run_build(fake, num_reqs=3, num_reqs_padded=4, base=base, step=2) + assert torch.equal(captured["positions"], fake.input_buffers.positions[:4]) bound = captured["seq_lens_cpu_upper_bound"] assert isinstance(bound, torch.Tensor), ( "seq_lens_cpu_upper_bound must be a tensor, not None" diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index fde3605a20..a87c647256 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -85,6 +85,7 @@ def test_pcp_draft_metadata_keeps_graph_padding_in_decode(cg_mode): speculator.max_model_len = speculator.draft_max_seq_len = 32 speculator.draft_is_prefilling = torch.zeros(4, dtype=torch.bool) speculator.input_buffers = SimpleNamespace( + positions=torch.arange(4, dtype=torch.int64), query_start_loc=torch.tensor([0, 1, 2, 2, 2], dtype=torch.int32), seq_lens=torch.tensor([11, 21, 0, 0], dtype=torch.int32), ) diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index f067915c87..3ca88992aa 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -130,6 +130,10 @@ class FlashAttnMLASparseMetadataBuilder( ) threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256) self._init_reorder_batch_threshold(threshold, supports_spec_as_decode=True) + self.supports_draft_decode_metadata_update = self.dcp_world_size == 1 + + def update_draft_decode_metadata(self, _metadata: SparseMLACommonMetadata) -> None: + pass class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 20a52af82b..5cfaf6ca5f 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -669,6 +669,7 @@ class DeepseekV32IndexerMetadata: decode: DeepSeekV32IndexerDecodeMetadata | None = None prefill: DeepseekV32IndexerPrefillMetadata | None = None + positions: torch.Tensor | None = None @triton.jit(do_not_specialize=["num_reqs", "num_actual_tokens", "num_tokens"]) @@ -758,6 +759,7 @@ class KpoolTailMetadataBuilder(AttentionMetadataBuilder): _cudagraph_support = AttentionCGSupport.ALWAYS supports_update_block_table = False + supports_draft_decode_metadata_update = True reorder_batch_threshold = None def __init__( @@ -807,6 +809,21 @@ class KpoolTailMetadataBuilder(AttentionMetadataBuilder): num_decode_tokens=num_decode_tokens, num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, + positions=positions, + ) + + def update_draft_decode_metadata( + self, + metadata: DeepseekV32IndexerMetadata, + ) -> None: + assert metadata.positions is not None + slot_mapping = metadata.slot_mapping[: metadata.num_decode_tokens] + slot_mapping.div_(self.kv_cache_spec.block_size, rounding_mode="floor") + slot_mapping.mul_(self.kv_cache_spec.block_size) + slot_mapping.add_( + metadata.positions[: metadata.num_decode_tokens].remainder( + self.kv_cache_spec.block_size + ) ) @@ -1016,6 +1033,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): ) self.indexer_decode_block_table_buffer: torch.Tensor | None = None self._max_num_batched_tokens = scheduler_config.max_num_batched_tokens + # DCP not supported yet + self.supports_draft_decode_metadata_update = self.dcp_world_size == 1 def _dcp_localize_decode_seq_lens( self, @@ -1668,6 +1687,50 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): return attn_metadata + def update_draft_decode_metadata( + self, + metadata: DeepseekV32IndexerMetadata, + ) -> None: + decode = metadata.decode + if decode is None or metadata.num_decode_tokens == 0: + return + + assert metadata.num_prefills == 0 + assert metadata.num_decodes == metadata.num_decode_tokens + assert decode.seq_lens.numel() == metadata.num_decode_tokens + assert self.dcp_world_size == 1 + + if self.compress_ratio > 1: + get_compressed_slot_mapping( + metadata.num_decode_tokens, + self.arange_buffer[: metadata.num_decode_tokens], + self.arange_buffer[: metadata.num_decode_tokens + 1], + metadata.seq_lens, + decode.block_table, + self.kv_cache_spec.num_states, + self.compress_ratio, + out=metadata.slot_mapping, + ) + torch.div( + metadata.seq_lens, + self.compress_ratio, + rounding_mode="floor", + out=decode.seq_lens.view(-1), + ) + else: + decode.seq_lens.view(-1).copy_(metadata.seq_lens) + decode.decode_lens.fill_(1) + + if current_platform.is_cuda() and has_deep_gemm(): + schedule_metadata = get_paged_mqa_logits_metadata( + decode.seq_lens, + self.kv_cache_spec.num_states, + self.num_sms, + indices=decode.indices, + ) + assert schedule_metadata.shape == decode.schedule_metadata.shape + decode.schedule_metadata.copy_(schedule_metadata) + def build_prefill_chunk_metadata( start_idx: int, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 0728ddb87c..77a9ece360 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -373,6 +373,7 @@ class DraftModelSpeculator(BaseSpeculator): kv_cache_config=self.kv_cache_config, causal=causal, seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound, + positions=self.input_buffers.positions[:num_tokens], is_prefilling=self.draft_is_prefilling[:num_reqs_padded], ) return attn_metadata