@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8a680463fa

Fork 0

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 6, 14:30 UTC

[Bugfix] Fix NemotronH MTP + Chunked Prefill (#35447)

Benjamin Chislett committed Mar 17, 2026 at 02:07 UTC 8a680463fab3bc9e6760417cd5c0a6aa58283065

5 files changed +181 -8

tests/v1/e2e/test_hybrid_chunked_prefill.py new

+104

@@ -0,0 +1,104 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +

4 +import pytest

5 +

6 +from vllm import SamplingParams

7 +from vllm.platforms import current_platform

8 +

9 +from ...utils import large_gpu_mark, multi_gpu_marks

10 +

11 +# A trivial request with a short prompt to ensure we run a mixed batch

12 +SMALL_MESSAGE = [

13 + {

14 + "role": "user",

15 + "content": "The secret beta value is 64. What is the secret beta?",

16 + }

17 +]

18 +

19 +# Sample prompt with a bunch of filler in between the critical fact and the request.

20 +# Both parts need to be processed properly for the model to generate the correct answer

21 +MESSAGES = [

22 + {

23 + "role": "user",

24 + "content": (

25 + "Important: The secret number is 42. "

26 + "The sky is green in this hypothetical world. "

27 + "Apples grow on trees in the forest. "

28 + "Rivers flow through the valleys and mountains. "

29 + "Birds sing songs in the early morning light. "

30 + "The weather today is sunny with clear skies ahead. "

31 + "Flowers bloom in the garden during spring season. "

32 + "Now answer with ONLY the number and nothing else: "

33 + "What is the secret number plus one?"

34 + ),

35 + }

36 +]

37 +

38 +

39 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")

40 +@pytest.mark.parametrize(

41 + "model_name",

42 + [

43 + pytest.param("Qwen/Qwen3.5-4B", marks=[large_gpu_mark(min_gb=40)]),

44 + pytest.param(

45 + "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8",

46 + marks=[large_gpu_mark(min_gb=80)] + multi_gpu_marks(num_gpus=2),

47 + ),

48 + ],

49 +)

50 +@pytest.mark.parametrize("enable_prefix_caching", [False, True])

51 +def test_mtp_speculative_mixed_batch_short_prefill(

52 + vllm_runner, model_name, enable_prefix_caching

53 +):

54 + """Test to ensure MTP speculative decoding correctly handles

55 + short prefill chunks that fall below the reorder_batch_threshold."""

56 +

57 + # Set so large that both prefills will be classified as decodes in a mixed batch

58 + # note, with prefix caching we require chunk_size >= mamba_block_size

59 + chunk_size = 256 if not enable_prefix_caching else 16384

60 + num_draft_tokens = 100

61 +

62 + with vllm_runner(

63 + model_name,

64 + speculative_config={

65 + "method": "mtp",

66 + "num_speculative_tokens": num_draft_tokens,

67 + },

68 + max_num_batched_tokens=chunk_size,

69 + max_model_len=512,

70 + enforce_eager=True,

71 + tensor_parallel_size=2,

72 + trust_remote_code=True,

73 + enable_chunked_prefill=True,

74 + enable_prefix_caching=enable_prefix_caching,

75 + mamba_cache_mode="align" if enable_prefix_caching else "none",

76 + ) as llm:

77 + sampling_params = SamplingParams(

78 + temperature=0.0,

79 + max_tokens=128,

80 + )

81 +

82 + # First small message gets prefilled first, under normal conditions since the

83 + # batch is not yet mixed. Then the second prefill arrives as a mixed batch, but

84 + # is shorter than num_speculative_tokens, so it gets misclassified as a decode

85 + # and processed with the wrong state management logic, causing the critical

86 + # fact from the first chunk to be lost and the model to generate nonsense.

87 + outputs = llm.get_llm().chat(

88 + [SMALL_MESSAGE, MESSAGES],

89 + sampling_params,

90 + chat_template_kwargs={"enable_thinking": False},

91 + )

92 +

93 + responses = []

94 + for output in outputs:

95 + generated_text = output.outputs[0].text

96 + print(f"Generated text: {generated_text!r}")

97 + responses.append(generated_text)

98 +

99 + assert "64" in responses[0], (

100 + "The first response should contain the correct value of 64."

101 + )

102 + assert "43" in responses[1], (

103 + "The second response should contain the correct value of 42+1=43."

104 + )

vllm/model_executor/layers/mamba/ops/mamba_ssm.py

+3 -3

@@ -334,13 +334,13 @@ def selective_state_update(

334 dt_bias = dt_bias.unsqueeze(0)

335 if out.dim() == 2:

336 out = out.unsqueeze(1)

337 - if num_accepted_tokens is not None:

338 - assert state_batch_indices is not None and state_batch_indices.dim() == 2

339 - assert dst_state_batch_indices is None or dst_state_batch_indices.dim() == 2

337 if state_batch_indices is not None and state_batch_indices.dim() == 1:

338 state_batch_indices = state_batch_indices.unsqueeze(1)

339 if dst_state_batch_indices is not None and dst_state_batch_indices.dim() == 1:

340 dst_state_batch_indices = dst_state_batch_indices.unsqueeze(1)

341 + if num_accepted_tokens is not None:

342 + assert state_batch_indices is not None and state_batch_indices.dim() == 2

343 + assert dst_state_batch_indices is None or dst_state_batch_indices.dim() == 2

344

345 _, nheads, dim, dstate = state.shape

346 batch = x.shape[0]

vllm/v1/attention/backends/mamba_attn.py

+6 -4

@@ -414,8 +414,11 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

414 ]

415 state_indices_tensor_p = state_indices_tensor_p[:, 0]

416

417 - if num_decodes > 0 and self.use_spec_decode:

418 - assert num_accepted_tokens is not None

417 + # Sometimes even with specdec enabled we get single-token prefill chunks that

418 + # should be treated as decodes but don't have num_accepted_tokens set.

419 + # These should be fine to process as non-spec decodes since there's only

420 + # one token, so no risk of placing accepted tokens in the wrong slot.

421 + if num_decodes > 0 and self.use_spec_decode and num_accepted_tokens is not None:

422 query_start_loc_d = common_attn_metadata.query_start_loc[: num_decodes + 1]

423 num_accepted_tokens = num_accepted_tokens[:num_decodes]

424

@@ -501,9 +504,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

504 state_indices_tensor_d = self.state_indices_tensor_d[:padded_bs]

505 state_indices_tensor_d[metadata.num_decodes :] = PAD_SLOT_ID

506

504 - if self.use_spec_decode:

507 + if self.use_spec_decode and num_accepted_tokens is not None:

508 assert query_start_loc_d is not None

506 - assert num_accepted_tokens is not None

509 query_start_loc_d = query_start_loc_d[: padded_bs + 1]

510 self.decode_num_accepted_tokens[: metadata.num_decodes].copy_(

511 num_accepted_tokens, non_blocking=True

vllm/v1/worker/gpu_model_runner.py

+26 -1

@@ -739,6 +739,19 @@ class GPUModelRunner(

739

740 self.uniform_decode_query_len = 1 + self.num_spec_tokens

741

742 + # When spec decode is active, the mamba backend classifies requests

743 + # with query_len <= reorder_batch_threshold as "decodes". Prefill

744 + # chunks that fall under this threshold get processed via the decode

745 + # path, which stores intermediate states at sequential slots. We must

746 + # set num_accepted_tokens to the chunk's query_len for those requests

747 + # so the next iteration reads from the correct final-state slot.

748 + # Prefills that went through the actual prefill path should keep the

749 + # default value of 1 (the prefill path stores state at slot 0 only).

750 + self.needs_prefill_as_decode_slots: bool = False

751 + self.prefill_as_decode_num_tokens = self._make_buffer(

752 + self.max_num_reqs, dtype=torch.int32

753 + )

754 +

755 # Cudagraph dispatcher for runtime cudagraph dispatching.

756 self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)

757

@@ -1355,12 +1368,22 @@ class GPUModelRunner(

1368 .int()

1369 .argmax(-1)

1370 )

1371 + spec_decode_active = bool(scheduler_output.scheduled_spec_decode_tokens)

1372 + if self.needs_prefill_as_decode_slots and spec_decode_active:

1373 + mamba_utils.update_accepted_tokens_for_prefill_as_decode(

1374 + self.input_batch,

1375 + self.prefill_as_decode_num_tokens,

1376 + self.num_accepted_tokens.gpu,

1377 + scheduler_output,

1378 + self.reorder_batch_threshold,

1379 + num_reqs,

1380 + )

1381 +

1382 if self.cache_config.mamba_cache_mode == "align":

1383 for i, num_tokens in enumerate(

1384 self.num_accepted_tokens.gpu[:num_reqs].cpu().numpy()

1385 ):

1386 self.input_batch.num_accepted_tokens_cpu[i] = num_tokens

1363 -

1387 mamba_utils.postprocess_mamba(

1388 scheduler_output,

1389 self.kv_cache_config,

@@ -2024,6 +2047,8 @@ class GPUModelRunner(

2047 else 0

2048 )

2049

2050 + if isinstance(builder, Mamba2AttentionMetadataBuilder):

2051 + self.needs_prefill_as_decode_slots = True

2052 extra_attn_metadata_args = {}

2053 if use_spec_decode and isinstance(

2054 builder, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder)

vllm/v1/worker/mamba_utils.py

+42

@@ -266,3 +266,45 @@ def postprocess_mamba(

266 if src_block_idx == dest_block_idx:

267 num_accepted_tokens_cpu[i] = 1

268 do_mamba_copy_block(copy_bufs)

269 +

270 +

271 +def update_accepted_tokens_for_prefill_as_decode(

272 + input_batch: GPUInputBatch,

273 + prefill_as_decode_num_tokens: CpuGpuBuffer,

274 + num_accepted_tokens_gpu: torch.Tensor,

275 + scheduler_output: SchedulerOutput,

276 + decode_qlen_threshold: int | None,

277 + num_reqs: int,

278 +):

279 + """

280 + Adjusts num_accepted_tokens for prefill chunks processed via the decode path.

281 + This ensures subsequent iterations read from the correct sequential state slot

282 + instead of the default prefill slot 0. Not used by GDN attention, which manually

283 + separates short prefills and short decodes when building the attention metadata.

284 + """

285 + any_is_prefill = False

286 + for i in range(num_reqs):

287 + num_computed = input_batch.num_computed_tokens_cpu[i]

288 + num_prompt = input_batch.num_prompt_tokens[i]

289 + is_prefill = num_computed < num_prompt

290 + req_id = input_batch.req_ids[i]

291 + query_len = scheduler_output.num_scheduled_tokens[req_id]

292 +

293 + if is_prefill:

294 + classified_as_decode = (

295 + decode_qlen_threshold is not None and query_len <= decode_qlen_threshold

296 + )

297 + num_tokens = query_len if classified_as_decode else 1

298 + any_is_prefill = True

299 + else:

300 + num_tokens = -1

301 + prefill_as_decode_num_tokens.np[i] = num_tokens

302 +

303 + # We can skip the GPU transfer if there aren't any values to update

304 + if any_is_prefill:

305 + prefill_as_decode_num_tokens.copy_to_gpu(num_reqs)

306 + num_accepted_tokens_gpu[:num_reqs] = torch.where(

307 + prefill_as_decode_num_tokens.gpu[:num_reqs] != -1,

308 + prefill_as_decode_num_tokens.gpu[:num_reqs],

309 + num_accepted_tokens_gpu[:num_reqs],

310 + )