PR #37411: fix(jais): only apply ALiBi when position_embedding_type is 'alibi'

• Repository: vllm-project/vllm
• Author: jigangz
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/37411

Description (problem / solution / changelog)

Summary

Fixes #37400

In JAISAttention.__init__, ALiBi slopes were computed and passed to Attention unconditionally, regardless of position_embedding_type. This caused models with position_embedding_type="learned" to apply both learned positional embeddings (wpe) and ALiBi simultaneously.

Changes

Added a conditional check in JAISAttention.__init__:

if config.position_embedding_type == "alibi":
    alibi_slopes = _get_alibi_slopes(total_num_heads)
    alibi_slopes = alibi_slopes[head_start:head_end]
else:
    alibi_slopes = None

This is consistent with how JAISModel already handles wpe — only creating learned position embeddings when position_embedding_type != "alibi".

Testing

• Verified the fix is consistent with the existing wpe guard in JAISModel
• No functional tests added as this requires a JAIS model checkpoint to test

Changed files

• vllm/model_executor/models/jais.py (modified, +5/-2)

PR #37474: jais: only enable ALiBi when position_embedding_type == "alibi"

• Repository: vllm-project/vllm
• Author: ahmedabbas104
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/37474

Description (problem / solution / changelog)

Summary This change fixes a correctness issue in the JAIS model implementation where the ALiBi positional bias was unconditionally constructed and applied even when the model configuration used learned positional embeddings. Applying both learned positional embeddings and an ALiBi bias results in double positional encoding, which can lead to incorrect attention behavior for JAIS variants configured with learned positional embeddings.

What I changed File: vllm/model_executor/models/jais.py Change: ALiBi slopes are now only computed and passed to the Attention layer when the HF config explicitly sets position_embedding_type == "alibi". When position_embedding_type is other values (e.g., "learned"), alibi_slopes is set to None so no ALiBi bias is applied.

Code excerpt (conceptual)

Before: alibi_slopes = _get_alibi_slopes(total_num_heads) alibi_slopes = alibi_slopes[head_start:head_end] self.attn = Attention(..., alibi_slopes=alibi_slopes, ...)

After: use_alibi = getattr(config, "position_embedding_type", "") == "alibi" alibi_slopes = _get_alibi_slopes(...) if use_alibi else None self.attn = Attention(..., alibi_slopes=alibi_slopes, ...)

Why this fixes the issue The JAISModel already conditionally creates the learned positional embedding (wpe) when position_embedding_type != "alibi", but the attention code was still enabling ALiBi unconditionally. This patch makes the attention implementation consistent with the model’s positional embedding config and prevents applying two forms of positional encoding simultaneously.

Testing & validation I created a working branch (agent/37400-fix) with the change and pushed it to a fork (ahmedabbas104/vllm). Recommended tests to run upstream or CI: Unit tests covering attention behavior (existing attention tests exercise alibi paths; ensure none fail). Model initialization checks for JAIS configs with position_embedding_type == "learned" and "alibi". Minimal forward pass smoke test for a JAIS config using learned embeddings to confirm outputs run without ALiBi.

Backward compatibility No behavior change for models configured with position_embedding_type == "alibi". Models that rely on learned embeddings will no longer see redundant ALiBi bias.

Related issue Fixes / addresses: vllm-project/vllm#37400

Changed files

• vllm/model_executor/models/jais.py (modified, +10/-2)