Fix(es): align expert_shared golden with kernel int32 matmul order - #847
Fix(es): align expert_shared golden with kernel int32 matmul order#847lwDavid wants to merge 1 commit into
Conversation
golden_expert_shared computed gate/up/W2 as fp32 matmuls of pre-dequantized operands, but the kernel does exact INT8xINT8->INT32 accumulate then dequant. The golden's fp32 accumulation rounding flipped bf16 last-bit ties on the final cast, causing flaky precision FAILs (~3/5 runs on a5). Mirror the kernel's 'accumulate-int32, then one fp32 dequant mul' order in the golden. Integer accumulation is exact and order-independent, so the golden's int32 matmul is bit-identical to the kernel's tiled int32 accumulate regardless of K-tiling. Verified: 5/5 PASS on a5 (was flaky).
📝 WalkthroughWalkthrough
ChangesExpert reference numerics
Estimated code review effort: 2 (Simple) | ~10 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@models/deepseek/v4-pro/expert_shared.py`:
- Around line 322-324: Split the semicolon-separated assignments in the shared
weight initialization block into separate statements, keeping each tensor and
scale assignment unchanged and preserving the existing symbols w1_i8, w1_scale,
w3_i8, w3_scale, w2_i8, and w2_scale.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 2b626c61-1a7a-4a84-8f1e-325c2cdafa9d
📒 Files selected for processing (1)
models/deepseek/v4-pro/expert_shared.py
| w1_i8 = tensors["shared_w1"]; w1_scale = tensors["shared_w1_scale"].float() # [MOE_INTER, D], [MOE_INTER] | ||
| w3_i8 = tensors["shared_w3"]; w3_scale = tensors["shared_w3_scale"].float() | ||
| w2_i8 = tensors["shared_w2"]; w2_scale = tensors["shared_w2_scale"].float() # [D, MOE_INTER], [D] |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Split the semicolon-separated assignments.
Lines 322-324 violate Ruff E702 and can fail linting.
Proposed fix
- w1_i8 = tensors["shared_w1"]; w1_scale = tensors["shared_w1_scale"].float() # [MOE_INTER, D], [MOE_INTER]
- w3_i8 = tensors["shared_w3"]; w3_scale = tensors["shared_w3_scale"].float()
- w2_i8 = tensors["shared_w2"]; w2_scale = tensors["shared_w2_scale"].float() # [D, MOE_INTER], [D]
+ w1_i8 = tensors["shared_w1"]
+ w1_scale = tensors["shared_w1_scale"].float()
+ w3_i8 = tensors["shared_w3"]
+ w3_scale = tensors["shared_w3_scale"].float()
+ w2_i8 = tensors["shared_w2"]
+ w2_scale = tensors["shared_w2_scale"].float()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| w1_i8 = tensors["shared_w1"]; w1_scale = tensors["shared_w1_scale"].float() # [MOE_INTER, D], [MOE_INTER] | |
| w3_i8 = tensors["shared_w3"]; w3_scale = tensors["shared_w3_scale"].float() | |
| w2_i8 = tensors["shared_w2"]; w2_scale = tensors["shared_w2_scale"].float() # [D, MOE_INTER], [D] | |
| w1_i8 = tensors["shared_w1"] | |
| w1_scale = tensors["shared_w1_scale"].float() | |
| w3_i8 = tensors["shared_w3"] | |
| w3_scale = tensors["shared_w3_scale"].float() | |
| w2_i8 = tensors["shared_w2"] | |
| w2_scale = tensors["shared_w2_scale"].float() |
🧰 Tools
🪛 Ruff (0.15.21)
[error] 322-322: Multiple statements on one line (semicolon)
(E702)
[error] 323-323: Multiple statements on one line (semicolon)
(E702)
[error] 324-324: Multiple statements on one line (semicolon)
(E702)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@models/deepseek/v4-pro/expert_shared.py` around lines 322 - 324, Split the
semicolon-separated assignments in the shared weight initialization block into
separate statements, keeping each tensor and scale assignment unchanged and
preserving the existing symbols w1_i8, w1_scale, w3_i8, w3_scale, w2_i8, and
w2_scale.
Source: Linters/SAST tools
1 participant
What & why
expert_sharedwas a flaky precision FAIL (~3/5 runs on a5). Root cause:golden_expert_sharedcomputed all three matmuls (gate/up/W2) as fp32 matmuls of pre-dequantized operands, while the kernel does exact INT8×INT8→INT32 accumulate then dequant. The golden's fp32 accumulation rounding flipped bf16 last-bit ties on the finalfp32→bf16cast.Fix
Mirror the kernel's "accumulate-int32, then one fp32 dequant mul" order in
golden_expert_shared. Integer accumulation is exact and order-independent, so the golden's int32 matmul is bit-identical to the kernel's tiled int32 accumulate regardless of K-tiling.This is a golden-side fix: the kernel was already numerically correct (int32 matmul is the exact dot product), and the golden was an imperfect fp32 emulator of the device's int8 path. Per
docs/precision-tuning.md§2, for an algebraically-equivalent reorder the golden mirrors the kernel's order. Aligning the golden also made the test stricter (a real kernel bug would still fail against the exact-int32 reference).Verification
5/5 PASS on a5 (was flaky ~3/5).
expert_sharedwas the only v4-pro single-card case hitting this; the int32 alignment makes it deterministic.