Skip to content

[Bugfix][Attention] Forward per-head FP8 descales through FA4 - #51363

Merged
yewentao256 merged 11 commits into
vllm-project:mainfrom
yiliu30:fix-fa4
Aug 11, 2026
Merged

yewentao256 merged 11 commits into
vllm-project:mainfrom
yiliu30:fix-fa4

Conversation

@yiliu30

@yiliu30 yiliu30 commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Summary

Forward per-head FP8 Q/K/V descales from vLLM's FlashAttention interface into the FA4 CuTe forward kernel.

Root cause

The FA3 path forwarded q_descale, k_descale, and v_descale, but the FA4 path dropped them. Per-head FP8 Q/K/V tensors were therefore consumed without their calibration scales, producing garbage output for llm-compressor attention/KV-quantized checkpoints.

Changes

  • Pass Q/K/V descales to the FA4 _flash_attn_fwd call.
  • Add a regression test using non-unity per-KV-head Q/K/V descales.

Validation

--------------------------------------------------
Prompt: 'Hello, my name is'
Generated text: '罨 Snapshot cyl罨mseanga乐趣热点 Streaming乐趣乐趣属于自己呼应乐趣乐趣乐趣'
--------------------------------------------------
Prompt: 'The president of the United States is'
Generated text: '魔法师和持平继续保持 sustained继续保持疗程 healed healed healed丝丝/detail久久unsubscribe Congratulations.ObjectModel pinpoint'
  • after fix
--------------------------------------------------
Prompt: 'Hello, my name is'
Generated text: ' Yizhou, and I am a student at the University of Toronto. I'
--------------------------------------------------
Prompt: 'The president of the United States is'
Generated text: ' the head of state and head of government of the United States, and the leader'
--------------------------------------------------

Duplicate-work check

Related issue #40444 and open PR #40448 were checked. PR #40448 clarifies the FlashInfer limitation and does not fix FA4 descale forwarding, so this is materially different kernel work.

AI assistance

AI assistance was used to prepare this PR. The human submitter must review every changed line, understand the fix, and confirm the evaluation results before requesting review.

cc @brian-dellabetta

Signed-off-by: yiliu30 <yi4.liu@intel.com>
@mergify mergify Bot added the bug Something isn't working label Aug 7, 2026
Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: yiliu30 <yi4.liu@intel.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@brian-dellabetta brian-dellabetta left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @yiliu30 !

@yewentao256 yewentao256 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the work!

Comment thread tests/kernels/attention/test_flash_attn.py Outdated
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Signed-off-by: Yi Liu <yi4.liu@intel.com>
@yiliu30
yiliu30 requested a review from yewentao256 August 8, 2026 10:13

@yewentao256 yewentao256 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks for the work!

@yewentao256 yewentao256 added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 8, 2026
@yewentao256

Copy link
Copy Markdown
Member

/ci run

@github-actions

github-actions Bot commented Aug 8, 2026

Copy link
Copy Markdown

@yiliu30, CI is now available for this PR.

  • /ci run starts a CI build.
  • /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.

@github-actions

github-actions Bot commented Aug 8, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #83008 for commit b25b73f31e91.

@yiliu30

yiliu30 commented Aug 11, 2026

Copy link
Copy Markdown
Contributor Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #83244 for commit 4fbb14d850d4.

@yiliu30

yiliu30 commented Aug 11, 2026

Copy link
Copy Markdown
Contributor Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #83254 for commit 235e526fb7ef.

Signed-off-by: yiliu30 <yi4.liu@intel.com>
@yiliu30

yiliu30 commented Aug 11, 2026

Copy link
Copy Markdown
Contributor Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #83318 for commit 3c645977c571.

@yiliu30

yiliu30 commented Aug 11, 2026

Copy link
Copy Markdown
Contributor Author

Hi @yewentao256 the CI was passed. Could you help get it merge? Thanks!

@yewentao256
yewentao256 merged commit 4f2f31b into vllm-project:main Aug 11, 2026
59 checks passed
zyp2014 pushed a commit to zyp2014/vllm that referenced this pull request Aug 21, 2026
…roject#51363)

Signed-off-by: yiliu30 <yi4.liu@intel.com>
Signed-off-by: Yi Liu <yi4.liu@intel.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working ready ONLY add when PR is ready to merge/full CI is needed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants