@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 57c629e9c1

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 7, 05:15 UTC

[Bugfix] Fix block_size for hybrid model MTP (#36036)

Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>

Benjamin Chislett committed Mar 5, 2026 at 01:10 UTC 57c629e9c1ce10ae649c5cb7411770ac31240bb0

2 files changed +25 -17

tests/v1/spec_decode/test_eagle.py

+14 -10

@@ -37,6 +37,8 @@ eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"

37 eagle3_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"

38 ar_draft_model_dir = "amd/PARD-Llama-3.2-1B" # Compatible with parallel and AR drafting

40 +BLOCK_SIZE = 16

41 +

43 def _create_proposer(

44 method: str,

@@ -91,9 +93,11 @@ def _create_proposer(

93 )

95 if "eagle" in method:

94 - return EagleProposer(vllm_config=vllm_config, device=device)

96 + proposer = EagleProposer(vllm_config=vllm_config, device=device)

97 else:

96 - return DraftModelProposer(vllm_config=vllm_config, device=device)

98 + proposer = DraftModelProposer(vllm_config=vllm_config, device=device)

99 + proposer.block_size = BLOCK_SIZE

100 + return proposer

101

102

103 def test_prepare_next_token_ids():

@@ -163,7 +167,7 @@ def test_prepare_next_token_ids():

167

168 common_attn_metadata = create_common_attn_metadata(

169 batch_spec,

166 - block_size=16,

170 + block_size=BLOCK_SIZE,

171 device=device,

172 )

173

@@ -207,7 +211,7 @@ def test_prepare_inputs():

211

212 common_attn_metadata = create_common_attn_metadata(

213 batch_spec,

210 - block_size=16,

214 + block_size=BLOCK_SIZE,

215 device=device,

216 )

217

@@ -302,7 +306,7 @@ def test_prepare_inputs_padded():

306

307 common_attn_metadata = create_common_attn_metadata(

308 batch_spec,

305 - block_size=16,

309 + block_size=BLOCK_SIZE,

310 device=device,

311 )

312

@@ -371,7 +375,7 @@ def test_set_inputs_first_pass_default_eagle():

375

376 common_attn_metadata = create_common_attn_metadata(

377 batch_spec,

374 - block_size=16,

378 + block_size=BLOCK_SIZE,

379 device=device,

380 )

381

@@ -462,7 +466,7 @@ def test_set_inputs_first_pass_draft_model():

466 device = torch.device(current_platform.device_type)

467

468 num_speculative_tokens = 2

465 - block_size = 16

469 + block_size = BLOCK_SIZE

470

471 # Create a proposer configured as a draft model (pass_hidden_states=False)

472 # We need to mock this since _create_proposer defaults to EAGLE

@@ -600,7 +604,7 @@ def test_set_inputs_first_pass_parallel_drafting():

604 device = torch.device(current_platform.device_type)

605

606 num_speculative_tokens = 3

603 - block_size = 16

607 + block_size = BLOCK_SIZE

608

609 proposer = _create_proposer("eagle", num_speculative_tokens, parallel_drafting=True)

610

@@ -926,7 +930,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):

930

931 common_attn_metadata = create_common_attn_metadata(

932 batch_spec,

929 - block_size=16,

933 + block_size=BLOCK_SIZE,

934 device=device,

935 )

936

@@ -1123,7 +1127,7 @@ def test_propose_tree(spec_token_tree):

1127 )

1128 common_attn_metadata = create_common_attn_metadata(

1129 batch_spec,

1126 - block_size=16,

1130 + block_size=BLOCK_SIZE,

1131 device=device,

1132 )

1133 sampling_metadata = mock.MagicMock()

vllm/v1/spec_decode/eagle.py

+11 -7

@@ -162,6 +162,9 @@ class SpecDecodeBaseProposer:

162 (self.max_num_tokens, self.hidden_size), dtype=self.dtype, device=device

163 )

164

165 + # Will be set when we initialize the attention backend

166 + self.block_size: int = -1

167 +

168 # We need +1 here because the arange is used to set query_start_loc,

169 # which has one more element than batch_size.

170 max_num_slots_for_arange = max(max_batch_size + 1, self.max_num_tokens)

@@ -583,8 +586,8 @@ class SpecDecodeBaseProposer:

586 common_attn_metadata._num_computed_tokens_cpu += 1

587

588 # Compute the slot mapping.

586 - # Use the first draft attention group's kv_cache_spec for block_size

587 - block_size = self.draft_attn_groups[0].kv_cache_spec.block_size

589 + block_size = self.block_size

590 + assert block_size > 0, "block_size has not been initialized."

591 if self.uses_mrope:

592 # all dimensions of positions are the same

593 block_numbers = clamped_positions[0] // block_size

@@ -778,17 +781,14 @@ class SpecDecodeBaseProposer:

781 # 2.

782 # Recompute the slot mapping based on the new positions and

783 # rejection mask.

781 - # Use the first draft attention group's kv_cache_spec for block_size

782 - # (all draft layers share the same kv-cache group)

783 - assert len(self.draft_attn_groups) > 0

784 - block_size = self.draft_attn_groups[0].kv_cache_spec.block_size

784 + assert self.block_size > 0, "block_size has not been initialized."

785 new_slot_mapping = compute_new_slot_mapping(

786 cad=cad,

787 new_positions=self.positions[:total_num_output_tokens],

788 is_rejected_token_mask=self.is_rejected_token_mask[

789 :total_num_output_tokens

790 ],

791 - block_size=block_size,

791 + block_size=self.block_size,

792 num_new_tokens=self.net_num_new_slots_per_request,

793 max_model_len=self.max_model_len,

794 )

@@ -1635,6 +1635,10 @@ class SpecDecodeBaseProposer:

1635 attention_groups[backend_key].layer_names.append(layer_name)

1636

1637 self.draft_attn_groups = list(attention_groups.values())

1638 + self.block_size = (

1639 + self.draft_attn_groups[0].get_metadata_builder().kv_cache_spec.block_size

1640 + )

1641 + logger.debug("Using block size %d for drafting layers", self.block_size)

1642

1643 def _determine_batch_execution_and_padding(

1644 self,