FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

[Common] Let MXFP8 quantization always reduce dbias in colwise by kainzhong · Pull Request #3439 · NVIDIA/TransformerEngine · GitHub

/ TransformerEngine Public

[Common] Let MXFP8 quantization always reduce dbias in colwise - #3439

Merged
ptrendx merged 2 commits into
NVIDIA:mainfrom
kainzhong:mxfp8_improve_dbias
Sep 2, 2026
ConversationCommits (2)ChecksFiles changed
Merged

[Common] Let MXFP8 quantization always reduce dbias in colwise#3439
ptrendx merged 2 commits into
NVIDIA:mainfrom
kainzhong:mxfp8_improve_dbias

Conversation

kainzhong commented Aug 28, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Update: I really can't find out why in the CUDA kernel this optimization regresses on FP32+DBIAS+DACT, but for my CuTeDSL kernel it still improves on this case. I don't find anything that could explain this

Description

This PR improve MXFP8 kernel's performance for DBIAS reduction by disabling row-wise DBIAS reduction and letting column-wise do it instead. By doing this we can eliminate the row-wise reduction register array and skip the thread data exchange where each thread owns a column instead of the row it had after the row-wise quantization, and reduce this column fragment and write the result in the workspace buffer.

The only exception is when we have FP32+DBIAS+DACT, where registers are no longer bottleneck, so this optimization's cost (extra SMEM traffic when we cache activations in the column-wise reduction only pass) is more than its gain (register save doesn't alleviate the bottleneck) so we disable it for this case only.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • DBIAS reduction always happens in the column-wise direction even we don't quantize column-wisely, in which case the column-wise pass only does DBIAS reduction and will not quantize
  • If activation is involved, column-wise will also cache the activation result into the SMEM buffer for row-wise pass to use
  • Column-wise pass now also use the half precision path with DBIAS (independent optimization)

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

[Common] Let MXFP8 quantization always reduce dbias in colwise
f90b9f3 
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

kainzhong commented Aug 28, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author 
  dbias — bf16

  ┌────────────┬─────┬───────────────────────┬──────────────────────┬───────────────────────┐
  │   shape    │ sw  │          row          │         col          │         both          │
  ├────────────┼─────┼───────────────────────┼──────────────────────┼───────────────────────┤
  │ 4096x4096  │ on  │ 18.68→16.04 (−14.13%) │ 16.21→15.97 (−1.48%) │ 22.62→20.47 (−9.50%)  │
  ├────────────┼─────┼───────────────────────┼──────────────────────┼───────────────────────┤
  │ 4096x4096  │ off │ 19.19→16.34 (−14.85%) │ 16.17→15.95 (−1.36%) │ 22.42→20.32 (−9.37%)  │
  ├────────────┼─────┼───────────────────────┼──────────────────────┼───────────────────────┤
  │ 4096x8192  │ on  │ 28.11→24.60 (−12.49%) │ 24.79→24.44 (−1.41%) │ 35.54→31.97 (−10.05%) │
  ├────────────┼─────┼───────────────────────┼──────────────────────┼───────────────────────┤
  │ 4096x8192  │ off │ 30.42→24.94 (−18.01%) │ 24.66→24.19 (−1.91%) │ 35.26→32.17 (−8.76%)  │
  ├────────────┼─────┼───────────────────────┼──────────────────────┼───────────────────────┤
  │ 4096x14336 │ on  │ 42.18→36.77 (−12.83%) │ 36.94→36.59 (−0.95%) │ 53.63→46.77 (−12.79%) │
  ├────────────┼─────┼───────────────────────┼──────────────────────┼───────────────────────┤
  │ 4096x14336 │ off │ 45.49→37.69 (−17.15%) │ 36.52→36.29 (−0.63%) │ 53.08→47.70 (−10.14%) │
  ├────────────┼─────┼───────────────────────┼──────────────────────┼───────────────────────┤
  │ 8192x8192  │ on  │ 52.70→47.00 (−10.82%) │ 46.92→46.76 (−0.34%) │ 65.19→57.97 (−11.08%) │
  ├────────────┼─────┼───────────────────────┼──────────────────────┼───────────────────────┤
  │ 8192x8192  │ off │ 56.19→47.83 (−14.88%) │ 46.78→46.45 (−0.71%) │ 64.60→58.28 (−9.78%)  │
  └────────────┴─────┴───────────────────────┴──────────────────────┴───────────────────────┘

  dbias — fp32

  ┌────────────┬─────┬──────────────────────┬──────────────────────┬──────────────────────┐
  │   shape    │ sw  │         row          │         col          │         both         │
  ├────────────┼─────┼──────────────────────┼──────────────────────┼──────────────────────┤
  │ 4096x4096  │ on  │ 22.99→23.24 (+1.09%) │ 20.41→20.40 (−0.05%) │ 25.48→25.40 (−0.31%) │
  ├────────────┼─────┼──────────────────────┼──────────────────────┼──────────────────────┤
  │ 4096x4096  │ off │ 23.28→23.10 (−0.77%) │ 20.36→20.34 (−0.10%) │ 25.70→25.74 (+0.16%) │
  ├────────────┼─────┼──────────────────────┼──────────────────────┼──────────────────────┤
  │ 4096x8192  │ on  │ 38.30→35.48 (−7.36%) │ 32.46→32.44 (−0.06%) │ 41.98→42.02 (+0.10%) │
  ├────────────┼─────┼──────────────────────┼──────────────────────┼──────────────────────┤
  │ 4096x8192  │ off │ 38.94→35.84 (−7.96%) │ 32.34→32.31 (−0.09%) │ 42.40→42.43 (+0.07%) │
  ├────────────┼─────┼──────────────────────┼──────────────────────┼──────────────────────┤
  │ 4096x14336 │ on  │ 57.55→54.65 (−5.04%) │ 51.17→51.22 (+0.10%) │ 64.15→64.11 (−0.06%) │
  ├────────────┼─────┼──────────────────────┼──────────────────────┼──────────────────────┤
  │ 4096x14336 │ off │ 58.23→55.04 (−5.48%) │ 51.03→51.00 (−0.06%) │ 64.32→64.43 (+0.17%) │
  ├────────────┼─────┼──────────────────────┼──────────────────────┼──────────────────────┤
  │ 8192x8192  │ on  │ 66.68→63.46 (−4.83%) │ 60.05→60.07 (+0.03%) │ 74.08→74.14 (+0.08%) │
  ├────────────┼─────┼──────────────────────┼──────────────────────┼──────────────────────┤
  │ 8192x8192  │ off │ 67.40→63.94 (−5.13%) │ 59.90→59.85 (−0.08%) │ 74.46→74.48 (+0.03%) │
  └────────────┴─────┴──────────────────────┴──────────────────────┴──────────────────────┘

  dbias_dgelu — bf16

  ┌────────────┬─────┬────────────────────────┬────────────────────────┬────────────────────────┐
  │   shape    │ sw  │          row           │          col           │          both          │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 4096x4096  │ on  │ 41.30→39.68 (−3.92%)   │ 39.76→40.28 (+1.31%)   │ 52.98→53.15 (+0.32%)   │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 4096x4096  │ off │ 40.79→39.60 (−2.92%)   │ 40.60→39.93 (−1.65%)   │ 51.95→51.88 (−0.13%)   │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 4096x8192  │ on  │ 69.87→68.83 (−1.49%)   │ 67.09→67.05 (−0.06%)   │ 94.39→94.15 (−0.25%)   │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 4096x8192  │ off │ 69.51→68.44 (−1.54%)   │ 67.50→67.54 (+0.06%)   │ 91.86→91.94 (+0.09%)   │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 4096x14336 │ on  │ 111.42→109.08 (−2.10%) │ 106.53→106.73 (+0.19%) │ 152.74→152.81 (+0.05%) │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 4096x14336 │ off │ 110.74→107.72 (−2.73%) │ 106.11→106.44 (+0.31%) │ 148.74→148.79 (+0.03%) │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 8192x8192  │ on  │ 136.68→134.01 (−1.95%) │ 131.71→132.03 (+0.24%) │ 184.72→185.07 (+0.19%) │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 8192x8192  │ off │ 136.41→132.64 (−2.76%) │ 130.95→130.91 (−0.03%) │ 180.21→180.61 (+0.22%) │
  └────────────┴─────┴────────────────────────┴────────────────────────┴────────────────────────┘

  dbias_dgelu — fp32 (not affected due to the exception gate added)

  ┌────────────┬─────┬────────────────────────┬────────────────────────┬────────────────────────┐
  │   shape    │ sw  │          row           │          col           │          both          │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 4096x4096  │ on  │ 43.43→43.62 (+0.44%)   │ 42.65→42.61 (−0.09%)   │ 55.18→54.98 (−0.36%)   │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 4096x4096  │ off │ 41.33→41.35 (+0.05%)   │ 42.76→42.52 (−0.56%)   │ 54.30→54.34 (+0.07%)   │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 4096x8192  │ on  │ 71.63→71.63 (+0.00%)   │ 69.77→69.75 (−0.03%)   │ 94.57→94.59 (+0.02%)   │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 4096x8192  │ off │ 68.67→68.62 (−0.07%)   │ 69.09→69.12 (+0.04%)   │ 93.34→93.37 (+0.03%)   │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 4096x14336 │ on  │ 112.95→112.98 (+0.03%) │ 108.58→108.55 (−0.03%) │ 151.92→151.92 (+0.00%) │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 4096x14336 │ off │ 107.32→107.43 (+0.10%) │ 107.55→107.55 (+0.00%) │ 149.76→149.77 (+0.01%) │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 8192x8192  │ on  │ 132.03→131.90 (−0.10%) │ 126.91→127.05 (+0.11%) │ 176.97→176.96 (−0.01%) │
  ├────────────┼─────┼────────────────────────┼────────────────────────┼────────────────────────┤
  │ 8192x8192  │ off │ 126.50→126.33 (−0.13%) │ 125.65→125.78 (+0.10%) │ 174.76→174.74 (−0.01%) │
  └────────────┴─────┴────────────────────────┴────────────────────────┴────────────────────────┘

  Every cell within ±0.6% — the gate makes this config compile to main's code path exactly.

kainzhong requested a review from tdophung August 28, 2026 20:33

greptile-apps Bot commented Aug 28, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR moves most fused MXFP8 dbias reductions to the columnwise traversal while retaining the rowwise path for the FP32 activation case.

  • Adds a columnwise reduction-only pass when columnwise quantization is disabled.
  • Reuses cached activation values between columnwise and rowwise passes.
  • Enables the half-precision quantization path when dbias is fused.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename Overview
transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh Refactors fused MXFP8 dbias reduction direction, activation caching, and half-precision dispatch without an eligible blocking failure established in this follow-up.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[MXFP8 quantization with fused dbias] --> B{Columnwise scaling enabled?}
    B -->|Yes| C[Columnwise quantization and dbias reduction]
    B -->|No| D{FP32 input with activation?}
    D -->|Yes| E[Rowwise quantization and dbias reduction]
    D -->|No| F[Columnwise dbias-only pass]
    F --> G[Cache fused activation when needed]
    G --> H[Rowwise quantization consumes cache]
    C --> I[Write partial dbias workspace]
    E --> I
    F --> I
Loading

Reviews (2): Last reviewed commit: "fix" | Re-trigger Greptile

kainzhong marked this pull request as draft August 28, 2026 21:17
fix
eb3f010 
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
kainzhong force-pushed the mxfp8_improve_dbias branch from 8fcac77 to eb3f010 Compare August 28, 2026 23:48
kainzhong marked this pull request as ready for review August 28, 2026 23:51

ptrendx commented Sep 1, 2026

Copy link
Copy Markdown
Member

/te-ci

ptrendx self-assigned this Sep 1, 2026
ptrendx merged commit 44158bf into NVIDIA:main Sep 2, 2026
39 of 47 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Show hidden characters
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

ptrendx ptrendx approved these changes
Oleg-Goncharov Awaiting requested review from Oleg-Goncharov Oleg-Goncharov is a code owner
tdophung Awaiting requested review from tdophung

Assignees

ptrendx

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

Back | FazBrowse Home | New Git URL