fix: attention_unfused alpha missing norm_factor squaring by delock · Pull Request #8666 · deepspeedai/DeepSpeed

delock · 2026-09-25T14:59:20Z

Summary

attention_unfused in pt_binding.cpp uses norm_factor directly as the QK^T scale, but ds_attention.py passes 1/sqrt(sqrt(head_dim)) which is head_dim^(-1/4). The correct softmax scale is head_dim^(-1/2), requiring the value to be squared.

The bug

// attention_unfused (line ~197):
float alpha = norm_factor;  // head_dim^(-1/4) — WRONG

// Fallback binding in the same file (line ~378):
float alpha = norm_factor * norm_factor;  // head_dim^(-1/2) — CORRECT

Same file, two bindings — one squares, one doesn't. The one that doesn't is a bug.

How norm_factor is computed (ds_attention.py)

self.norm_factor = math.sqrt(head_dim)        # sqrt(64) = 8
self.norm_factor = math.sqrt(self.norm_factor) # sqrt(8) = 2.83 = head_dim^(1/4)
# Passed to binding as: 1 / norm_factor = head_dim^(-1/4)

The binding is expected to square it back to head_dim^(-1/2).

Impact

For head_dim=64: correct scale is 1/8, buggy scale is 1/2.83 — a 2.8× difference in QK^T score magnitude, causing significantly different softmax distributions.

Test plan

Verify attention output matches F.scaled_dot_product_attention reference
Test with multiple head_dim values (64, 128, 256)

ds_attention passes 1/sqrt(sqrt(head_dim)); attention_unfused used it directly as QK^T scale (head_dim^(-1/4)). The correct softmax scale is head_dim^(-1/2), requiring alpha = norm_factor * norm_factor. The fallback binding in the same file (line 378) already does this — this fix aligns attention_unfused with it. Signed-off-by: Guokai Ma <guokai.ma@intel.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ceef30ff39

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2026-09-25T15:02:24Z

-    float alpha = norm_factor;
+    // ds_attention passes 1/sqrt(sqrt(head_dim)); square it back to the
+    // 1/sqrt(head_dim) softmax scale, matching the fallback binding below.
+    float alpha = norm_factor * norm_factor;


Add an integration test for the changed attention scale

This changes an inference-level attention result but adds no test coverage (the commit's stated test plan is still unchecked). Add and execute a real-device integration test that compares this path's output against an independent attention reference, including the relevant head dimensions, before merging.

AGENTS.md reference: AGENTS.md:L36-L37

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-09-25T15:02:24Z

-    float alpha = norm_factor;
+    // ds_attention passes 1/sqrt(sqrt(head_dim)); square it back to the
+    // 1/sqrt(head_dim) softmax scale, matching the fallback binding below.
+    float alpha = norm_factor * norm_factor;


Add the required Signed-off-by trailer

Commit 10ac5597212375b3a0c04a9bad296eb378c65b98 is a non-merge commit, and git log -1 --format='%b' confirms that its message contains no Signed-off-by: trailer. Add the required DCO sign-off before merging.

AGENTS.md reference: AGENTS.md:L8-L9

Useful? React with 👍 / 👎.

Copilot

Copilot review overview

🟡 Changes recommended

The corrected C++ path lacks regression coverage, and its new comment inaccurately describes non-default configurations.

Review effort: Balanced
Findings: 1 · 1

Open (2)

Add regression coverage for C++ softmax_context_int8 scaling · New
Clarify norm_factor as the caller-provided pre-square scale · New

What changed in this PR

Corrects attention scaling by squaring norm_factor before the unfused QKᵀ operation.

Changes:

Applies the expected softmax scaling.
Documents alignment with the other attention binding.

File	Description
`csrc/transformer/inference/csrc/pt_binding.cpp`	Corrects unfused attention scaling.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

-    float alpha = norm_factor;
+    // ds_attention passes 1/sqrt(sqrt(head_dim)); square it back to the
+    // 1/sqrt(head_dim) softmax scale, matching the fallback binding below.
+    float alpha = norm_factor * norm_factor;


+    // ds_attention passes 1/sqrt(sqrt(head_dim)); square it back to the
+    // 1/sqrt(head_dim) softmax scale, matching the fallback binding below.


Add test_softmax_context_int8 comparing the legacy int8 attention binding against an independently scaled reference softmax(q@k^T * norm_factor^2) @ v for head_dim 64/128/256 and seq_len 16/128, and clarify that norm_factor is the caller-provided pre-square scale so the comment stays accurate under non-default configurations (use_mup, inverse-layer scaling). Verified on RTX 4080 with torch 2.12.1+cu130: the test fails on the pre-fix binding and passes after the alpha squaring fix. Signed-off-by: Guokai Ma <guokai.ma@intel.com>

delock · 2026-10-06T04:50:29Z

close as this does not block #8381 works

delock requested a review from tjruwase as a code owner September 25, 2026 14:59

chatgpt-codex-connector Bot reviewed Sep 25, 2026

View reviewed changes

pengdurice added the will-review label Sep 25, 2026

delock mentioned this pull request Sep 26, 2026

feat(rollout): segment-based kernel injection for hybrid-engine greedy decode #8602

Draft

3 of 4 tasks

delock added merge-train PRs on merge-train (without additional commits) would be cherrypicked in an aggregate commit and removed merge-train PRs on merge-train (without additional commits) would be cherrypicked in an aggregate commit labels Sep 29, 2026

hwchen2017 requested a balanced review from Copilot September 29, 2026 04:14

Copilot started reviewing on behalf of hwchen2017 September 29, 2026 04:14

View session

Copilot AI reviewed Sep 29, 2026

View reviewed changes

delock requested review from loadams and tohtana as code owners October 3, 2026 16:22

delock force-pushed the fix/pt-binding-alpha-squaring branch from d0dc79f to 82ecffe Compare October 4, 2026 03:24

delock closed this Oct 6, 2026

Navigation Menu

Repository navigation

fix: attention_unfused alpha missing norm_factor squaring - #8666

fix: attention_unfused alpha missing norm_factor squaring#8666
delock wants to merge 2 commits into
masterfrom
fix/pt-binding-alpha-squaring

delock commented Sep 25, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Sep 25, 2026

Uh oh!

chatgpt-codex-connector Bot Sep 25, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

delock commented Oct 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		// ds_attention passes 1/sqrt(sqrt(head_dim)); square it back to the
		// 1/sqrt(head_dim) softmax scale, matching the fallback binding below.

Conversation

delock commented Sep 25, 2026

Summary

The bug

How norm_factor is computed (ds_attention.py)

Impact

Test plan

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

delock commented Oct 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants