@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9dcf6bf344

[BugFix] Fix DP token padding in dflash attention metadata (#56181)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>

Giancarlo Delfin committed Sep 11, 2026 at 14:33 UTC 9dcf6bf344caa7793bae0b45a7896d3f8e03a01a

5 files changed +97 -97

tests/v1/spec_decode/test_eagle_draft_attn_metadata.py

+41 -23

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Unit tests for the EAGLE speculator's draft attention metadata builder.

5 -These tests guard the regression where ``_build_draft_attn_metadata`` did

5 +These tests guard the regression where ``_build_uniform_attn_metadata`` did

6 not populate ``seq_lens_cpu_upper_bound`` on the per-step

7 ``CommonAttentionMetadata``. Several downstream attention backends and

8 helpers (``split_decodes_prefills_and_extends``, the MLA indexer,

@@ -13,11 +13,14 @@ backends (e.g. ``ROCM_AITER_FA`` with eagle/eagle3 spec decode):

13 AssertionError: assert common_attn_metadata.seq_lens_cpu_upper_bound is not None

14 """

16 -from types import SimpleNamespace

16 +from types import MethodType, SimpleNamespace

17 from unittest.mock import patch

19 +import numpy as np

20 import torch

22 +from vllm.config.compilation import CUDAGraphMode

23 +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor

24 from vllm.v1.worker.gpu.spec_decode import speculator as base_speculator

25 from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator

@@ -30,7 +33,7 @@ def _make_fake_speculator(

33 draft_max_seq_len: int = 1024,

34 ) -> SimpleNamespace:

35 """Build a fake EagleSpeculator with just the attributes used by

33 - ``_build_draft_attn_metadata``. We deliberately avoid constructing a

36 + ``_build_uniform_attn_metadata``. We deliberately avoid constructing a

37 real ``EagleSpeculator`` because that requires a full ``VllmConfig``

38 and a draft model.

39 """

@@ -46,8 +49,8 @@ def _make_fake_speculator(

49 cp_rank=0,

50 cp_interleave=1,

51 )

49 - return SimpleNamespace(

50 - arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"),

52 + fake = SimpleNamespace(

53 + arange_np=np.arange(max_num_reqs + 1, dtype=np.int32),

54 draft_is_prefilling=torch.zeros(max_num_reqs, dtype=torch.bool),

55 block_tables=fake_block_tables,

56 input_buffers=fake_input_buffers,

@@ -56,21 +59,42 @@ def _make_fake_speculator(

59 max_model_len=max_model_len,

60 draft_max_seq_len=draft_max_seq_len,

61 )

62 + # The uniform wrapper delegates through self; bind the real implementation.

63 + fake._build_attn_metadata = MethodType(EagleSpeculator._build_attn_metadata, fake)

64 + return fake

61 -def _run_build(fake, *, num_reqs, num_reqs_padded, num_tokens_padded, base, step):

67 +def _run_build(

68 + fake,

69 + *,

70 + num_reqs,

71 + num_reqs_padded,

72 + base,

73 + step,

74 + num_query_per_req=1,

75 + cg_mode=CUDAGraphMode.FULL,

76 +):

77 captured: dict[str, object] = {}

79 def fake_build_attn_metadata(**kwargs):

80 captured.update(kwargs)

81 return {}

83 + # Request padding only occurs under FULL graphs, where the captured batch

84 + # is a uniform decode of num_reqs_padded * num_query_per_req tokens.

85 + batch_desc = BatchExecutionDescriptor(

86 + cg_mode=cg_mode,

87 + num_tokens=num_reqs_padded * num_query_per_req,

88 + num_reqs=num_reqs_padded,

89 + uniform_token_count=num_query_per_req,

90 + )

91 +

92 with patch.object(base_speculator, "build_attn_metadata", fake_build_attn_metadata):

69 - EagleSpeculator._build_draft_attn_metadata(

93 + EagleSpeculator._build_uniform_attn_metadata(

94 fake, # type: ignore[arg-type]

95 + batch_desc=batch_desc,

96 num_reqs=num_reqs,

72 - num_reqs_padded=num_reqs_padded,

73 - num_tokens_padded=num_tokens_padded,

97 + num_query_per_req=num_query_per_req,

98 seq_lens_cpu_upper_bound=base,

99 step=step,

100 )

@@ -85,9 +109,7 @@ def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound():

109 fake = _make_fake_speculator()

110 base = torch.tensor([100, 200, 300, 0], dtype=torch.int32)

111

88 - captured = _run_build(

89 - fake, num_reqs=3, num_reqs_padded=4, num_tokens_padded=4, base=base, step=2

90 - )

112 + captured = _run_build(fake, num_reqs=3, num_reqs_padded=4, base=base, step=2)

113

114 bound = captured["seq_lens_cpu_upper_bound"]

115 assert isinstance(bound, torch.Tensor), (

@@ -102,20 +124,19 @@ def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound():

124 assert torch.equal(bound, torch.tensor([102, 202, 302, 0], dtype=torch.int32))

125

126

105 -def test_build_draft_attn_metadata_handles_zero_unpadded_reqs():

106 - """Edge case: when ``num_reqs == 0`` the upper-bound tensor must

107 - still be a valid all-zero tensor of length ``num_reqs_padded``."""

127 +def test_build_draft_attn_metadata_zeroes_padded_upper_bound_tail():

128 + """The padded tail of the upper-bound tensor is zeroed, so it stays a

129 + valid tensor of length ``num_reqs_padded`` regardless of padding."""

130 fake = _make_fake_speculator()

131 base = torch.zeros(2, dtype=torch.int32)

132

111 - captured = _run_build(

112 - fake, num_reqs=0, num_reqs_padded=2, num_tokens_padded=2, base=base, step=1

113 - )

133 + captured = _run_build(fake, num_reqs=1, num_reqs_padded=2, base=base, step=1)

134

135 bound = captured["seq_lens_cpu_upper_bound"]

136 assert isinstance(bound, torch.Tensor)

137 assert bound.shape == (2,)

118 - assert torch.equal(bound, torch.zeros(2, dtype=torch.int32))

138 + # base is all zeros, so the real entry is 0 + step and the pad entry is 0.

139 + assert torch.equal(bound, torch.tensor([1, 0], dtype=torch.int32))

140

141

142 def test_build_draft_attn_metadata_clamps_to_max_model_len():

@@ -124,9 +145,7 @@ def test_build_draft_attn_metadata_clamps_to_max_model_len():

145 fake = _make_fake_speculator(max_model_len=1024)

146 base = torch.tensor([1023, 500], dtype=torch.int32)

147

127 - captured = _run_build(

128 - fake, num_reqs=2, num_reqs_padded=2, num_tokens_padded=2, base=base, step=3

129 - )

148 + captured = _run_build(fake, num_reqs=2, num_reqs_padded=2, base=base, step=3)

149

150 bound = captured["seq_lens_cpu_upper_bound"]

151 # 1023 + 3 = 1026 -> clamped to 1024; 500 + 3 = 503 unaffected.

@@ -154,7 +173,6 @@ def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens():

173 fake,

174 num_reqs=3,

175 num_reqs_padded=4,

157 - num_tokens_padded=4,

176 base=torch.tensor([5, 9, 16]),

177 step=0,

178 )

vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py

+6 -8

@@ -532,11 +532,10 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

532 slot_mappings_by_layer = build_slot_mappings_by_layer(

533 slot_mappings, self.kv_cache_config

534 )

535 - attn_metadata = self._build_draft_attn_metadata(

535 + attn_metadata = self._build_uniform_attn_metadata(

536 num_reqs=num_reqs,

537 - num_reqs_padded=batch_desc.num_reqs or num_reqs,

538 - # One query per request; exclude DP-only model padding.

539 - num_tokens_padded=batch_desc.num_reqs or num_reqs,

537 + batch_desc=batch_desc,

538 + num_query_per_req=1,

539 seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

540 step=step,

541 )

@@ -581,11 +580,10 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

580 slot_mappings_by_layer = build_slot_mappings_by_layer(

581 slot_mappings, self.kv_cache_config

582 )

584 - attn_metadata = self._build_draft_attn_metadata(

583 + attn_metadata = self._build_uniform_attn_metadata(

584 num_reqs=num_reqs,

586 - num_reqs_padded=batch_desc.num_reqs or num_reqs,

587 - # One query per request; exclude DP-only model padding.

588 - num_tokens_padded=batch_desc.num_reqs or num_reqs,

585 + batch_desc=batch_desc,

586 + num_query_per_req=1,

587 seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

588 step=1,

589 )

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

+3 -33

@@ -1,10 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import copy

4 -from collections.abc import Mapping

4 from typing import Any

7 -import numpy as np

6 import torch

7 import torch.nn as nn

@@ -299,33 +297,6 @@ class DFlashSpeculator(DraftModelSpeculator):

297 num_reqs, self.num_speculative_steps

298 )

299

302 - def _build_draft_attn_metadata(

303 - self,

304 - num_reqs: int,

305 - num_reqs_padded: int,

306 - num_tokens_padded: int,

307 - seq_lens_cpu_upper_bound: torch.Tensor,

308 - step: int,

309 - num_query_per_req: int | None = None,

310 - causal: bool | Mapping[int, bool] = False,

311 - query_start_loc_np: np.ndarray | None = None,

312 - dcp_local_seq_lens: torch.Tensor | None = None,

313 - ) -> dict[str, Any] | None:

314 - if not self.draft_attn_layer_names:

315 - return None

316 - assert num_query_per_req is None # Omitted for DFlash, read from self instead

317 - return super()._build_draft_attn_metadata(

318 - num_reqs,

319 - num_reqs_padded,

320 - num_tokens_padded,

321 - seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

322 - step=step,

323 - num_query_per_req=self.num_query_per_req,

324 - causal=causal,

325 - query_start_loc_np=query_start_loc_np,

326 - dcp_local_seq_lens=dcp_local_seq_lens,

327 - )

328 -

300 @torch.inference_mode()

301 def propose(

302 self,

@@ -471,7 +442,6 @@ class DFlashSpeculator(DraftModelSpeculator):

442 need_eager=is_profile,

443 dp_sync=batch_sync,

444 )

474 - num_reqs_padded = batch_desc.num_reqs or num_reqs

445 num_tokens_padded = batch_desc.num_tokens

446 num_tokens_across_dp = (

447 batch_sync.num_tokens_across_dp if batch_sync is not None else None

@@ -479,10 +449,10 @@ class DFlashSpeculator(DraftModelSpeculator):

449

450 # Rebuild the draft attention metadata even when replaying the FULL

451 # graph so that any attention metadata builder state is updated.

482 - draft_attn_metadata = self._build_draft_attn_metadata(

452 + draft_attn_metadata = self._build_uniform_attn_metadata(

453 num_reqs=num_reqs,

484 - num_reqs_padded=num_reqs_padded,

485 - num_tokens_padded=num_tokens_padded,

454 + batch_desc=batch_desc,

455 + num_query_per_req=self.num_query_per_req,

456 seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound,

457 step=self.num_query_per_req,

458 causal=self._group_causal,

vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py

+2 -3

@@ -226,10 +226,9 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):

226 slot_mappings = build_slot_mappings_by_layer(

227 slot_mappings_tensor, self.kv_cache_config

228 )

229 - draft_attn_metadata = self._build_draft_attn_metadata(

229 + draft_attn_metadata = self._build_attn_metadata(

230 num_reqs=num_reqs,

231 - num_reqs_padded=batch_desc.num_reqs or num_reqs,

232 - num_tokens_padded=batch_desc.num_tokens,

231 + batch_desc=batch_desc,

232 query_start_loc_np=input_batch.query_start_loc_np,

233 seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

234 step=0,

vllm/v1/worker/gpu/spec_decode/speculator.py

+45 -30

@@ -23,6 +23,7 @@ from vllm.v1.worker.gpu.attn_utils import (

23 )

24 from vllm.v1.worker.gpu.block_table import BlockTables

25 from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens

26 +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor

27 from vllm.v1.worker.gpu.dp_utils import DPSyncState

28 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

29 from vllm.v1.worker.gpu.model_states.interface import ModelState

@@ -149,9 +150,7 @@ class DraftModelSpeculator(BaseSpeculator):

150 dtype=torch.int64,

151 device=device,

152 )

152 - self.arange = torch.arange(

153 - self.max_num_reqs + 1, dtype=torch.int32, device="cpu"

154 - )

153 + self.arange_np = np.arange(self.max_num_reqs + 1, dtype=np.int32)

154 self.draft_is_prefilling = torch.zeros(self.max_num_reqs, dtype=torch.bool)

155

156 self.draft_logits: torch.Tensor | None = None

@@ -254,42 +253,37 @@ class DraftModelSpeculator(BaseSpeculator):

253 self.target_input_buffers = target_input_buffers

254 self.target_attn_groups = target_attn_groups

255

257 - def _build_draft_attn_metadata(

256 + def _build_attn_metadata(

257 self,

258 num_reqs: int,

260 - num_reqs_padded: int,

261 - num_tokens_padded: int,

259 + batch_desc: BatchExecutionDescriptor,

260 + query_start_loc_np: np.ndarray,

261 seq_lens_cpu_upper_bound: torch.Tensor,

262 step: int,

264 - num_query_per_req: int = 1,

263 causal: bool | Mapping[int, bool] = True,

266 - query_start_loc_np: np.ndarray | None = None,

264 dcp_local_seq_lens: torch.Tensor | None = None,

265 ) -> dict[str, Any] | None:

269 - if query_start_loc_np is not None:

270 - # Non-uniform query layout (e.g. multi-module MTP's mixed

271 - # prefill/decode queries); num_query_per_req is ignored.

272 - query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32)

273 - query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy(

274 - query_start_loc_np[: num_reqs + 1]

275 - )

276 - query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs]

277 - max_query_len = int(

278 - (query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max()

279 - )

280 - else:

281 - # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req.

282 - # Clamp keeps the series non-decreasing past num_reqs, which some

283 - # attention backends require.

284 - query_start_loc_cpu = (

285 - torch.clamp(self.arange[: num_reqs_padded + 1], max=num_reqs)

286 - * num_query_per_req

287 - )

288 - max_query_len = num_query_per_req

266 + num_reqs_padded = batch_desc.num_reqs or num_reqs

267 + # A FULL graph replays a captured shape whose padded requests each hold

268 + # a full query width, so attention must see the padded token count.

269 + # PIECEWISE/eager needs the actual token count, because batch_desc may

270 + # carry graph or DP padding that no request owns, which would desync it

271 + # from query_start_loc.

272 + num_tokens = (

273 + batch_desc.num_tokens

274 + if batch_desc.cg_mode == CUDAGraphMode.FULL

275 + else int(query_start_loc_np[-1])

276 + )

277 + query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32)

278 + query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy(

279 + query_start_loc_np[: num_reqs + 1]

280 + )

281 + query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs]

282 + max_query_len = int((query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max())

283 block_tables = [

284 x[:num_reqs_padded] for x in self.block_tables.input_block_tables

285 ]

292 - slot_mappings = self.block_tables.slot_mappings[:, :num_tokens_padded]

286 + slot_mappings = self.block_tables.slot_mappings[:, :num_tokens]

287 draft_seq_lens_cpu_upper_bound = torch.zeros(

288 num_reqs_padded, dtype=torch.int32, device="cpu"

289 )

@@ -313,7 +307,7 @@ class DraftModelSpeculator(BaseSpeculator):

307 attn_metadata = build_attn_metadata(

308 attn_groups=self.attn_groups,

309 num_reqs=num_reqs_padded,

316 - num_tokens=num_tokens_padded,

310 + num_tokens=num_tokens,

311 query_start_loc_gpu=self.input_buffers.query_start_loc[

312 : num_reqs_padded + 1

313 ],

@@ -434,3 +428,24 @@ class DraftModelSpeculator(BaseSpeculator):

428 uniform_token_count=num_query_per_req,

429 eager=False,

430 ), num_batch_tokens

431 +

432 + def _build_uniform_attn_metadata(

433 + self,

434 + batch_desc: BatchExecutionDescriptor,

435 + num_reqs: int,

436 + num_query_per_req: int,

437 + seq_lens_cpu_upper_bound: torch.Tensor,

438 + step: int,

439 + causal: bool | Mapping[int, bool] = True,

440 + dcp_local_seq_lens: torch.Tensor | None = None,

441 + ) -> dict[str, Any] | None:

442 + query_start_loc_np = self.arange_np[: num_reqs + 1] * num_query_per_req

443 + return self._build_attn_metadata(

444 + num_reqs=num_reqs,

445 + batch_desc=batch_desc,

446 + query_start_loc_np=query_start_loc_np,

447 + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

448 + step=step,

449 + causal=causal,

450 + dcp_local_seq_lens=dcp_local_seq_lens,

451 + )