iframe-proxy

Ankit-Jaiswal-AMD · 2026-09-10T10:20:12Z

Purpose

This PR adds support for fp32 in attention sinks, to support running the gpt-oss-20B-Bf16 model on fp16. It converts all the tensors other than bf16 to fp32 to compute attention sinks.
It also updates the unit tests and adds test coverage for fp16.

Test Plan

Run the unit tests for the added support:
python -m pytest tests/kernels/attention/test_cpu_attn.py -q
Install the public fork in the validation env and run a real CPU accuracy benchmark:
python -m pip install . --no-build-isolation --no-deps --force-reinstall

lm_eval --model vllm \
    --model_args pretrained=unsloth/gpt-oss-20b-BF16,dtype=half \
    --tasks gsm8k \
    --batch_size auto \
    --trust_remote_code \
    --apply_chat_template \
    --num_fewshot 5 \
    --gen_kwargs "max_gen_toks=2048" \
    --log_samples \
    --output_path <output_path>

Test Result

The Unit tests passed.
1242 passed, 1478 skipped, 14 warnings in 818.95s (0:13:38)
The accuracy benchmark passes, with comparable accuracy.

vllm ({'pretrained': 'unsloth/gpt-oss-20b-BF16', 'dtype': 'half'}), gen_kwargs: ({'max_gen_toks': 2048}), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.8779|±  |0.0090|
|     |       |strict-match    |     5|exact_match|↑  |0.0205|±  |0.0039|

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

github-actions · 2026-09-10T10:20:30Z

amukho

Can you please update the tag to use [CPU] instead of [ZENTORCH] in this PR since the fix is generic for all CPU backends?
The changes look good to me...

amukho

Looks good to me...

dllehr-amd · 2026-09-17T15:56:51Z

+            cached = getattr(self, "_sinks_fp32", None)
+            if cached is None:
+                self._sinks_fp32 = s_aux.to(torch.float32)
+            s_aux = self._sinks_fp32


Just a quick question would it make sense to convert self.sinks = self._sinks_fp32? and leave our code just passing self.sinks?

Yes that would be a another way to go ahead as well. The main reason for creating another variable was to not modify the user set sinks during the forward pass. I have modified the code as per you comment.

dllehr-amd · 2026-09-17T15:58:04Z

+        # are loaded, rather than at __init__ time).
+        s_aux = self.sinks
+        if s_aux is not None and s_aux.dtype != torch.bfloat16:
+            cached = getattr(self, "_sinks_fp32", None)


similar to my comment below if we converted self.sinks to fp32 when needed could we just check the sinks dtype and be done with it?

Modified the code as per the comment.

github-actions · 2026-09-17T16:00:54Z

✅ @Ankit-Jaiswal-AMD, CI is now available for this PR.

/ci run starts upstream CI; /amd-ci run starts AMD CI only.
Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
/ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
/amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
/ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

dllehr-amd · 2026-09-17T16:01:36Z

/ci run

github-actions · 2026-09-17T16:02:01Z

❌ This PR is 499 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

mergify · 2026-09-18T05:52:08Z

Hi @Ankit-Jaiswal-AMD, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Ankit-Jaiswal-AMD · 2026-09-18T06:12:57Z

/ci run

github-actions · 2026-09-18T06:13:16Z

❌ This PR is 1 commit behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

Ankit-Jaiswal-AMD · 2026-09-18T06:44:27Z

/ci run

github-actions · 2026-09-18T06:44:45Z

✅ Triggered Buildkite CI #89793 for commit 2acc9698daa9.

Signed-off-by: Ankit Jaiswal <ankit.jaiswal@amd.com>

Ankit-Jaiswal-AMD · 2026-09-22T06:18:30Z

/ci run

github-actions · 2026-09-22T06:18:56Z

✅ Triggered Buildkite CI #90380 for commit e2d380218f76.

Ankit-Jaiswal-AMD · 2026-09-22T08:22:06Z

/ci retry

github-actions · 2026-09-22T08:22:25Z

✅ Queued 6 failed job(s) for retry in Buildkite CI #90380.

Ankit-Jaiswal-AMD · 2026-09-22T10:39:10Z

/ci retry

github-actions · 2026-09-22T10:39:27Z

✅ Queued 1 failed job(s) for retry in Buildkite CI #90380.

dllehr-amd · 2026-09-22T23:29:43Z

/ci retry

github-actions · 2026-09-22T23:30:03Z

dllehr-amd

Thanks @Ankit-Jaiswal-AMD for addressing my concerns. LGTM!

Signed-off-by: Ankit Jaiswal <ankit.jaiswal@amd.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>

Ankit-Jaiswal-AMD requested review from AndreasKaratzas, LucasWilkinson, MatthewBonanni, WoosukKwon, bigPYJ1151, mgoin, tlrmchlsmth, yewentao256 and zyongye as code owners September 10, 2026 10:20

claude Bot reviewed Sep 10, 2026

View reviewed changes

mergify Bot added the cpu Related to CPU backends label Sep 10, 2026

amukho suggested changes Sep 10, 2026

View reviewed changes

Ankit-Jaiswal-AMD changed the title ~~[ZENTORCH] Adds support for fp32 attention sinks~~ [CPU] Adds support for fp32 attention sinks Sep 11, 2026

amukho approved these changes Sep 11, 2026

View reviewed changes

Ankit-Jaiswal-AMD force-pushed the feature/fp32_attention_sinks branch from e342fb3 to d39b59e Compare September 11, 2026 06:35

dllehr-amd reviewed Sep 17, 2026

View reviewed changes

dllehr-amd added the ready ONLY add when PR is ready to merge/full CI is needed label Sep 17, 2026

Ankit-Jaiswal-AMD force-pushed the feature/fp32_attention_sinks branch from d39b59e to 67ec8b4 Compare September 18, 2026 05:44

Ankit-Jaiswal-AMD force-pushed the feature/fp32_attention_sinks branch 2 times, most recently from 8a57822 to 2acc969 Compare September 18, 2026 06:43

[CPU] Adds support for fp32 attention sinks

e2d3802

Signed-off-by: Ankit Jaiswal <ankit.jaiswal@amd.com>

Ankit-Jaiswal-AMD force-pushed the feature/fp32_attention_sinks branch from 2acc969 to e2d3802 Compare September 22, 2026 06:18

mganczarenko mentioned this pull request Sep 22, 2026

[CI Failure]: (MI355) Core Operation Kernels Shard 1: kernels/test_concat_mla_q.py::test_concat_mla_q_fp8_nope_and_rope #58134

Closed

3 tasks

Ankit-Jaiswal-AMD requested a review from dllehr-amd September 22, 2026 11:12

Merge branch 'main' into feature/fp32_attention_sinks

9dbecf3

dllehr-amd approved these changes Sep 23, 2026

View reviewed changes

dllehr-amd merged commit 88afb9d into vllm-project:main Sep 23, 2026
18 checks passed

Sunbelt Computer Software

PL/B Language Development and Support

Conversation

Ankit-Jaiswal-AMD commented Sep 10, 2026

Purpose

Test Plan

Test Result

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

github-actions Bot commented Sep 10, 2026

Uh oh!

amukho left a comment

Choose a reason for hiding this comment

Uh oh!

amukho left a comment

Choose a reason for hiding this comment

Uh oh!

dllehr-amd Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

Ankit-Jaiswal-AMD Sep 18, 2026

Choose a reason for hiding this comment

Uh oh!

dllehr-amd Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

Ankit-Jaiswal-AMD Sep 18, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Sep 17, 2026

Uh oh!

dllehr-amd commented Sep 17, 2026

Uh oh!

github-actions Bot commented Sep 17, 2026

Uh oh!

mergify Bot commented Sep 18, 2026

Uh oh!

Ankit-Jaiswal-AMD commented Sep 18, 2026

Uh oh!

github-actions Bot commented Sep 18, 2026

Uh oh!

Ankit-Jaiswal-AMD commented Sep 18, 2026

Uh oh!

github-actions Bot commented Sep 18, 2026

Uh oh!

Ankit-Jaiswal-AMD commented Sep 22, 2026

Uh oh!

github-actions Bot commented Sep 22, 2026

Uh oh!

Ankit-Jaiswal-AMD commented Sep 22, 2026

Uh oh!

github-actions Bot commented Sep 22, 2026

Uh oh!

Ankit-Jaiswal-AMD commented Sep 22, 2026

Uh oh!

github-actions Bot commented Sep 22, 2026

Uh oh!

dllehr-amd commented Sep 22, 2026

Uh oh!

github-actions Bot commented Sep 22, 2026

Uh oh!

dllehr-amd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants