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 14, 05:30 UTC
[bug fix] Fix llama4 spec decoding (#22691)
Signed-off-by: qizixi <qizixi@meta.com> Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
qizixi committed
Aug 19, 2025 at 17:53 UTC
5bfe0dea7a3422f027ebb8fe36ce50ed5aa68d86
1 file changed
+4
-2
vllm/model_executor/models/llama4.py
+4
-2
@@ -195,7 +195,9 @@ class Llama4Attention(nn.Module):
195
is_neox_style=is_neox_style,
196
) if not self.nope else None
197
198
- attn_cls = Attention if self.nope else ChunkedLocalAttention
198
+ use_chunked_local_attn = not self.nope and config.attention_chunk_size
199
+ attn_cls = (ChunkedLocalAttention
200
+ if use_chunked_local_attn else Attention)
201
self.attn = attn_cls(
202
self.num_heads,
203
self.head_dim,
@@ -206,7 +208,7 @@ class Llama4Attention(nn.Module):
208
prefix=f"{prefix}.attn",
209
**({
210
"attention_chunk_size": config.attention_chunk_size
209
- } if not self.nope else {}))
211
+ } if use_chunked_local_attn else {}))
212
213
def _get_attn_scale(self, positions: torch.Tensor) -> torch.Tensor:
214
floor = torch.floor((positions + 1.0) / self.floor_scale)