[GLM 5.3 Perf] Enable fused multi-step decode, 13.3% E2E throughput improvement for concurrency 1 (#57443)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

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]]

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])

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"

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),
     )

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]):

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,

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