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 10, 20:15 UTC
[Bugfix] Fix block size used in EAGLE slot mapping (#31540)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Jan 1, 2026 at 22:32 UTC
ea53ca5e8583b2154ffd8041dfc8735929a58105
1 file changed
+8
-9
vllm/v1/spec_decode/eagle.py
+8
-9
@@ -71,7 +71,6 @@ class EagleProposer:
71
self.device = device
72
self.dtype = vllm_config.model_config.dtype
73
self.max_model_len = vllm_config.model_config.max_model_len
74
- self.block_size = vllm_config.cache_config.block_size
74
self.dp_rank = vllm_config.parallel_config.data_parallel_rank
75
self.num_speculative_tokens = self.speculative_config.num_speculative_tokens
76
self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
@@ -470,22 +469,23 @@ class EagleProposer:
469
common_attn_metadata._num_computed_tokens_cpu += 1
470
471
# Compute the slot mapping.
472
+ block_size = attn_metadata_builder.kv_cache_spec.block_size
473
if self.uses_mrope:
474
# all dimensions of positions are the same
475
- block_numbers = clamped_positions[0] // self.block_size
475
+ block_numbers = clamped_positions[0] // block_size
476
else:
477
- block_numbers = clamped_positions // self.block_size
477
+ block_numbers = clamped_positions // block_size
478
block_ids = common_attn_metadata.block_table_tensor.gather(
479
dim=1, index=block_numbers.view(-1, 1)
480
)
481
block_ids = block_ids.view(-1)
482
if self.uses_mrope:
483
common_attn_metadata.slot_mapping = (
484
- block_ids * self.block_size + clamped_positions[0] % self.block_size
484
+ block_ids * block_size + clamped_positions[0] % block_size
485
)
486
else:
487
common_attn_metadata.slot_mapping = (
488
- block_ids * self.block_size + clamped_positions % self.block_size
488
+ block_ids * block_size + clamped_positions % block_size
489
)
490
# Mask out the slot mappings that exceed the max model length.
491
# Otherwise, the KV cache will be inadvertently updated with the
@@ -800,12 +800,11 @@ class EagleProposer:
800
attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1)
801
802
# Compute the slot mapping.
803
+ block_size = tree_attn_metadata_builder.kv_cache_spec.block_size
804
query_positions = flattened_draft_positions[:, level : level + query_len]
804
- block_numbers = query_positions // self.block_size
805
+ block_numbers = query_positions // block_size
806
block_ids = attn_metadata.block_table.gather(dim=1, index=block_numbers)
806
- slot_mapping = (
807
- block_ids * self.block_size + query_positions % self.block_size
808
- )
807
+ slot_mapping = block_ids * block_size + query_positions % block_size
808
# Mask out the slot mappings that exceed the max model length.
809
# Otherwise, the KV cache will be inadvertently updated with the
810
# padding tokens.