A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 21:15 UTC
[Bugfix] Fix spec decode test failures on Blackwell (SM100+) (#39546)
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com> Signed-off-by: Rishi Puri <puririshi98@berkeley.edu> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Stefano Castagnetta <scastagnetta@nvidia.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
Rishi Puri committed
Apr 21, 2026 at 13:21 UTC
9f39b380d070d2f60d28a152b9cbd05fea91a821
2 files changed
+52
-5
.buildkite/test_areas/spec_decode.yaml
+34
@@ -12,6 +12,17 @@ steps:
12
commands:
13
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
14
15
+- label: Spec Decode Eagle Nightly B200
16
+ timeout_in_minutes: 30
17
+ device: b200
18
+ optional: true
19
+ source_file_dependencies:
20
+ - vllm/v1/spec_decode/
21
+ - vllm/v1/worker/gpu/spec_decode/
22
+ - tests/v1/e2e/spec_decode/
23
+ commands:
24
+ - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
25
+
26
- label: Spec Decode Speculators + MTP
27
timeout_in_minutes: 30
28
device: h200_18gb
@@ -23,6 +34,18 @@ steps:
34
commands:
35
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
36
37
+- label: Spec Decode Speculators + MTP Nightly B200
38
+ timeout_in_minutes: 30
39
+ device: b200
40
+ optional: true
41
+ source_file_dependencies:
42
+ - vllm/v1/spec_decode/
43
+ - vllm/v1/worker/gpu/spec_decode/
44
+ - vllm/transformers_utils/configs/speculators/
45
+ - tests/v1/e2e/spec_decode/
46
+ commands:
47
+ - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
48
+
49
- label: Spec Decode Ngram + Suffix
50
timeout_in_minutes: 30
51
device: h200_18gb
@@ -43,6 +66,17 @@ steps:
66
commands:
67
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
68
69
+- label: Spec Decode Draft Model Nightly B200
70
+ timeout_in_minutes: 30
71
+ device: b200
72
+ optional: true
73
+ source_file_dependencies:
74
+ - vllm/v1/spec_decode/
75
+ - vllm/v1/worker/gpu/spec_decode/
76
+ - tests/v1/e2e/spec_decode/
77
+ commands:
78
+ - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
79
+
80
- label: DFlash Speculators Correctness
81
timeout_in_minutes: 30
82
device: h100
vllm/v1/attention/backends/flashinfer.py
+18
-5
@@ -920,9 +920,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
920
all_uses_trtllm = (num_prefills == 0 or prefill_use_trtllm) and (
921
num_decodes == 0 or decode_use_trtllm
922
)
923
- is_only_trtllm_decode = num_prefills == 0 and (
924
- num_decodes > 0 and decode_use_trtllm
925
- )
923
924
if not all_uses_trtllm:
925
if self.has_sinks:
@@ -968,7 +965,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
965
966
# Guard access to seq_lens_cpu, which may not always be needed
967
# and can be expensive to retrieve in async mode.
971
- needs_seq_lens_cpu = self.use_dcp or use_cascade or not is_only_trtllm_decode
968
+ # When all attention (both prefill and decode) uses TRTLLM,
969
+ # seq_lens_cpu is not needed since TRTLLM paths use GPU tensors
970
+ # (block_tables, seq_lens) directly.
971
+ needs_seq_lens_cpu = self.use_dcp or use_cascade or not all_uses_trtllm
972
seq_lens_cpu = common_attn_metadata.seq_lens_cpu if needs_seq_lens_cpu else None
973
seq_lens_np = seq_lens_cpu.numpy() if seq_lens_cpu is not None else None
974
num_blocks_np = (
@@ -1006,7 +1006,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1006
num_blocks_np -= num_common_kv_blocks
1007
1008
# Compute paged_kv_indices if necessary
1009
- needs_paged_kv_indices = use_cascade or not is_only_trtllm_decode
1009
+ # paged_kv_indices is only needed for FlashInfer native paths;
1010
+ # TRTLLM paths use block_tables directly on GPU.
1011
+ needs_paged_kv_indices = use_cascade or not all_uses_trtllm
1012
if needs_paged_kv_indices:
1013
assert num_blocks_np is not None
1014
assert seq_lens_np is not None
@@ -1083,9 +1085,20 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1085
qo_indptr_prefill_gpu = (
1086
qo_indptr[prefill_start:] - qo_indptr[prefill_start]
1087
)
1088
+ # Compute cum_seq_lens_kv on GPU to avoid CPU sync.
1089
+ # This is the cumulative sum of the number of KV cache
1090
+ # blocks per prefill request.
1091
+ prefill_seq_lens = seq_lens[prefill_start:]
1092
+ num_blocks_per_req = (prefill_seq_lens + page_size - 1) // page_size
1093
paged_kv_indptr_prefill_gpu = self.paged_kv_indptr.gpu[
1094
prefill_start : num_reqs + 1
1095
]
1096
+ paged_kv_indptr_prefill_gpu[0] = 0
1097
+ torch.cumsum(
1098
+ num_blocks_per_req,
1099
+ dim=0,
1100
+ out=paged_kv_indptr_prefill_gpu[1:],
1101
+ )
1102
# Compute max_q_len for prefill requests
1103
query_lens_prefill_cpu = (
1104
qo_indptr_prefill_cpu[1:] - qo_indptr_prefill_cpu[:-1]