Resolved ALIBI bias regression due to porting flat PA

[tannervoas742]

Requires associated changes on vllm-hpu-extension PR

Changes:

• Added back alibi biases to decode stage.
• Optimized ALiBI memory usage.
  • Added environment variable "VLLM_PROMPT_ALIBI_MAX_SEQ_LEN" to allow
    large models to run with restricted prompt lengths.
  • Prompt biases instantiated once in init rather than each
    forward.
  • Prompt and decode biases are shared across encoder/decoder layers.
• Added environment variable "VLLM_ALIBI_USE_FLOAT32_BIASES" to resolve
  accuracy issue on long sequences.
• Updated jais, mpt, falcon, baichuan, and bloom to work with ALiBI.
  • Due to bloom's 176B parameter size I was unable to test this model.
    Its changes are the simplest though.
• Works in lazy and eager mode.
• ALiBI is restricted to "VLLM_PROMPT_USE_FUSEDSDPA=false", and
  "VLLM_CONTIGUOUS_PA=true".
• Add position offsets to improve quality on BS > 1 with sequences of
  varying length.
• BS > 1 may have accuracy issues if on FW < 1.19.0. This is due to
  limitation in softmax. Resolved on FW >= 1.19.0.
• NTT patch for GQA

Co-authored-by: Tanner Voas [email protected]
Co-authored-by: Haihao Xiang [email protected]
Signed-off-by: Tanner Voas [email protected]

[tannervoas742]

@itaraban @madamczykhabana @kzawora-intel has anyone gotten a chance to review this PR and the associated one on vllm-hpu-extension. I just pushed out a significant update that minimizes changes to non-alibi code sections. It also has significant accuracy and memory optimization changes.

With the current changes ALiBi is now fully functional as long as FW >= 1.19.0 is being used.

Please help review. Any feedback would be appreciated.

[tannervoas742]

@michalkuligowski @kwisniewski98 conflicts have been resolved. Yapf and ruff issues should also be resolved now.

[PatrykWo]

@kwisniewski98 please resolve conflicts before merge, plus some tests failed.

[tannervoas742]

@kwisniewski98 please resolve conflicts before merge, plus some tests failed.

Resolved merge conflicts. Waiting for tests to re-run.

[kwisniewski98]

@tannervoas742 Could you please switch for now vllm-hpu-extension version to branch that you've added to check if tests will go through?

[tannervoas742]

@tannervoas742 Could you please switch for now vllm-hpu-extension version to branch that you've added to check if tests will go through?

Updated the reference. The cpu-test passed when I ran it locally.

[PatrykWo]

@tannervoas742 please resolve issue in yapf

yapf.....................................................................Failed

• hook id: yapf
• exit code: 1

[tannervoas742]

@tannervoas742 please resolve issue in yapf

yapf.....................................................................Failed

• hook id: yapf
• exit code: 1

Fixed. I setup my pre-commit hooks locally and they passed now including yapf.

Problems resolved

[PatrykWo]

@tannervoas742 please rebase

[tannervoas742]

@tannervoas742 please rebase

Rebased

[tannervoas742]

Freshly rebased. ALiBi fixes still work. Pre-commit hooks all passed.

[tannervoas742]

@kwisniewski98 @PatrykWo @michalkuligowski any progress towards getting this merged? There are customers in PRC who care about ALiBi support in vLLM.