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, 05:45 UTC
[Bugfix] JAIS: Only apply ALiBi when position_embedding_type='alibi' (#37820)
Co-authored-by: r266-tech <r266-tech@users.noreply.github.com>
r266-tech committed
Mar 23, 2026 at 15:36 UTC
02e6efe56d545d50d342356630430b2ad9583ca9
1 file changed
+8
-5
vllm/model_executor/models/jais.py
+8
-5
@@ -117,11 +117,14 @@ class JAISAttention(nn.Module):
117
prefix=f"{prefix}.c_proj",
118
)
119
120
- tp_rank = get_tensor_model_parallel_rank()
121
- head_start = tp_rank * self.num_heads
122
- head_end = (tp_rank + 1) * self.num_heads
123
- alibi_slopes = _get_alibi_slopes(total_num_heads)
124
- alibi_slopes = alibi_slopes[head_start:head_end]
120
+ self.use_alibi = config.position_embedding_type == "alibi"
121
+ alibi_slopes = None
122
+ if self.use_alibi:
123
+ tp_rank = get_tensor_model_parallel_rank()
124
+ head_start = tp_rank * self.num_heads
125
+ head_end = (tp_rank + 1) * self.num_heads
126
+ alibi_slopes = _get_alibi_slopes(total_num_heads)
127
+ alibi_slopes = alibi_slopes[head_start:head_end]
128
self.attn = Attention(
129
self.num_heads,
130
self.head_dim,