[Bugfix][Spec Decode] Qualify Transformers backend attention layer names with the model prefix (#59293)
Signed-off-by: Farhad Al-Amin <99006654+Legend-2727@users.noreply.github.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Farhad Al-Amin committed
Oct 2, 2026 at 16:21 UTC
af2a582d5c8e3ec1edccc304a3d2527f1ea3e3e3
2 files changed
+24
-8
tests/v1/e2e/spec_decode/draft_model/test_draft_model.py
+15
@@ -150,6 +150,21 @@ def test_draft_model_quantization(
150
assert_draft_model_correctness(sd_case, vllm_runner)
151
152
153
+@single_gpu_only
154
+def test_draft_model_transformers_backend(vllm_runner):
155
+ """Ensure that Attention prefixes don't collide between the target and
156
+ draft models when both models use the Transformers modelling backend."""
157
+ sd_case = ArgsTest(
158
+ target_model="allenai/OLMo-2-0425-1B-Instruct",
159
+ draft_model="allenai/OLMo-2-0425-1B-Instruct",
160
+ sampling_config=greedy_sampling(),
161
+ num_speculative_tokens=3,
162
+ expected_acceptance_len=0.98 * (3 + 1), # epsilon discount of K + 1
163
+ expected_acceptance_rate=0.98, # slight epsilon
164
+ )
165
+ assert_draft_model_correctness(sd_case, vllm_runner)
166
+
167
+
168
@multi_gpu_only(num_gpus=2)
169
def test_draft_model_tensor_parallelism(vllm_runner):
170
"""Ensure spec decode works when running with TP > 1."""
vllm/model_executor/models/transformers/base.py
+9
-8
@@ -184,7 +184,7 @@ class Base(
184
# Substitute remaining layers with vLLM's layers as needed
185
self.recursive_replace()
186
# Create attention instances for KV cache allocation
187
- self._create_attention_instances()
187
+ self._create_attention_instances(prefix)
188
189
# Initialize any parameters that have not had their modules replaced
190
self.init_parameters(self.model)
@@ -579,7 +579,7 @@ class Base(
579
580
self.hf_to_vllm_mapper |= WeightsMapper(orig_to_new_stacked=orig_to_new_stacked)
581
582
- def _create_attention_instances(self):
582
+ def _create_attention_instances(self, prefix: str = ""):
583
"""Create `Attention` instances to inform KV cache allocation."""
584
text_config = self.text_config
585
attn_cls = self._get_attn_cls()
@@ -613,8 +613,8 @@ class Base(
613
f"Layer {layer} does not dispatch through the Transformers "
614
"attention interface and vLLM has no other way to handle it."
615
)
616
- prefix, attn_fuser = self.attention_fusers[i]
617
- attn_module = self.get_submodule(prefix)
616
+ attn_prefix, attn_fuser = self.attention_fusers[i]
617
+ attn_module = self.get_submodule(attn_prefix)
618
619
# `[i]` is the whole-model config unless the checkpoint is
620
# heterogeneous, in which case it is this layer's own geometry.
@@ -632,17 +632,18 @@ class Base(
632
scale=scale,
633
cache_config=self.cache_config,
634
quant_config=self.quant_config,
635
- prefix=f"{i}.attn",
635
+ prefix=maybe_prefix(prefix, f"{attn_prefix}.{VLLM_ATTN_ATTR}"),
636
)
637
638
if attn_cls is MLAAttention:
639
mla_fuser = next(
640
- (f for f in self.fusers[prefix] if isinstance(f, MLAFuser)), None
640
+ (f for f in self.fusers[attn_prefix] if isinstance(f, MLAFuser)),
641
+ None,
642
)
643
if mla_fuser is None:
644
raise ValueError(
644
- f"Layer {i} ({prefix}) did not fuse as MLA, but the rest of "
645
- "the model did, so its latent KV cache cannot be allocated."
645
+ f"Layer {i} ({attn_prefix}) did not fuse as MLA, but the rest "
646
+ "of the model did, so its latent KV cache cannot be allocated."
647
)
648
dims = get_mla_dims(self.model_config)
649
kwargs.update(