[Perf][Spec Decode] Enable fused multi-step draft decode for FlashInfer trtllm-gen (#58371)
Signed-off-by: Bowen Wang <bowen@inferact.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Bowen Wang committed
Sep 29, 2026 at 14:06 UTC
35d6fb3187d0a5a70c43ee5b0c3e923a7cb96226
2 files changed
+58
tests/v1/attention/test_trtllm_attention_integration.py
+48
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Integration tests for TRTLLM gen-full attention through FlashInfer."""
4
5
+import dataclasses
6
import unittest.mock
7
from functools import partial
8
@@ -513,3 +514,50 @@ def test_trtllm_gen_nvfp4_kv_integration(batch_spec_name: str):
514
kv_cache_dtype="nvfp4",
515
model_name=MODEL_NVFP4,
516
)
517
+
518
+
519
+@torch.inference_mode()
520
+def test_trtllm_gen_draft_decode_metadata_tracks_in_place_seq_lens(
521
+ monkeypatch: pytest.MonkeyPatch,
522
+):
523
+ """Fused draft steps reuse the step-1 build after advancing seq_lens in
524
+ place, so its decode metadata must match the split path's step-2 rebuild."""
525
+ monkeypatch.setattr(
526
+ "vllm.utils.flashinfer.supports_trtllm_attention",
527
+ unittest.mock.Mock(return_value=True),
528
+ )
529
+ monkeypatch.setattr(
530
+ "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
531
+ _mock_get_per_layer_parameters,
532
+ )
533
+ vllm_config = create_vllm_config(
534
+ model_name=MODEL, block_size=BLOCK_SIZE, num_gpu_blocks=NUM_GPU_BLOCKS
535
+ )
536
+ model_config = vllm_config.model_config
537
+ kv_cache_spec = FullAttentionSpec(
538
+ block_size=BLOCK_SIZE,
539
+ num_kv_heads=model_config.get_num_kv_heads(vllm_config.parallel_config),
540
+ head_size=model_config.get_head_size(),
541
+ dtype=model_config.dtype,
542
+ )
543
+ device = torch.device(f"{DEVICE_TYPE}:0")
544
+ with set_current_vllm_config(vllm_config):
545
+ builder = FlashInferMetadataBuilder(
546
+ kv_cache_spec, ["test_layer_0"], vllm_config, device
547
+ )
548
+ assert builder.supports_draft_decode_metadata_update
549
+
550
+ step1 = create_common_attn_metadata(BATCH_SPECS["decode_only"], BLOCK_SIZE, device)
551
+ reused = builder.build(common_prefix_len=0, common_attn_metadata=step1)
552
+ step1.seq_lens.add_(1)
553
+ step2 = dataclasses.replace(step1, seq_lens_cpu_upper_bound=step1.seq_lens.cpu())
554
+ rebuilt = builder.build(common_prefix_len=0, common_attn_metadata=step2)
555
+
556
+ assert isinstance(reused.decode, FlashInferTrtllmAPIDecode)
557
+ for field in dataclasses.fields(reused.decode):
558
+ value = getattr(reused.decode, field.name)
559
+ expected = getattr(rebuilt.decode, field.name)
560
+ if isinstance(value, torch.Tensor):
561
+ assert torch.equal(value, expected), field.name
562
+ else:
563
+ assert value == expected, field.name
vllm/v1/attention/backends/flashinfer.py
+10
@@ -860,6 +860,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
860
self.use_trtllm_gen_varlen_decode = self._uses_trtllm_gen_varlen_decode(
861
vllm_config, self.flashinfer_trtllm_api_decode_kernel, self.use_dcp
862
)
863
+ # Fused draft steps advance seq_lens in place and keep block tables fixed,
864
+ # so one trtllm-gen decode build serves every step (not DCP-local seq_lens).
865
+ self.supports_draft_decode_metadata_update = (
866
+ self.flashinfer_trtllm_api_decode_kernel
867
+ == FlashInferDecodeKernel.TRTLLM_GEN
868
+ and not self.use_dcp
869
+ )
870
self._init_reorder_batch_threshold(
871
1,
872
supports_spec_as_decode=(
@@ -1837,6 +1844,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1844
attn_metadata.decode = FIDecode(wrapper=decode_wrapper)
1845
return attn_metadata
1846
1847
+ def update_draft_decode_metadata(self, _metadata: FlashInferMetadata) -> None:
1848
+ pass
1849
+
1850
def use_cascade_attention(self, *args, **kwargs) -> bool:
1851
if self.kv_cache_spec.dtype != self.vllm_config.model_config.dtype:
1852
# TODO: The cascade wrapper currently does not support setting