Align Hopper FP8 blockwise quant to power-of-2 scales; fix tl.trans

[haok1402]

TLDR: Brings the Hopper FP8 path onto the same power-of-2 block scaling that NVIDIA TransformerEngine uses by default. In TE, fp32-mantissa scales are a Hopper-only opt-out (NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1) and Blackwell is pow2-only. See FP8 Blockwise Scaling.

What changed

• Fixed a correctness bug: Triton 3.7.1 tl.trans corrupts fp8 tiles, producing a wrong transposed FP8 weight (used in dgrad). Fix: store to transposed addresses instead of tl.trans on fp8.
• Renamed the functions fused_* -> fp8cast_*: every kernel in this module is fused, so fused_ said nothing; fp8cast_ uses DeepGEMM's own "cast" verb.
• Renamed the scale flag SCALING_MODE -> NATIVE_E8M0 (it just means native-SM100 vs emulated-SM90 E8M0; both are pow2) and corrected the docstrings.
• Tests: moved tests/test_fp8_quantize_kernels.py -> tests/operators/test_deepgemm_quantize.py; fixed the reference to always use E8M0 (was previously device-gated and wrong on Hopper -> 39 failures).

Validation

• pytest tests/operators/test_deepgemm_quantize.py -> 48 passed (was 39 failing on H100).
• Minimal standalone repro of the Triton bug (workspace/triton_fp8_trans_repro.py, included below): kernel1 (tl.trans) mismatches 4090/16384; kernel2 (transposed-address store, the fix) passes with 0.
• dsv2-lite FP8 training run to confirm the loss curve is unchanged vs. reference (fix only touches dgrad).

Files

• pithtrain/operators/deepgemm_quantize.py
• tests/operators/test_deepgemm_quantize.py

Reproduction of the Triton bug

"""
On an NVIDIA H100 80GB HBM3 (Python 3.14.3, Triton 3.7.1, torch 2.12.1+cu130),
the following outputs were observed:

kernel1 (tl.trans):      mismatch 4090 / 16384
kernel2 (address store): mismatch 0 / 16384
"""

import torch
import triton
import triton.language as tl

BLOCK = 128


@triton.jit
def kernel1(x_ptr, o_ptr, t_ptr, BLOCK: tl.constexpr = BLOCK):
    r = tl.arange(0, BLOCK)
    off = r[:, None] * BLOCK + r[None, :]
    fp8 = tl.load(x_ptr + off).to(tl.float32).to(tl.float8e4nv)
    tl.store(o_ptr + off, fp8)
    tl.store(t_ptr + off, tl.trans(fp8))


@triton.jit
def kernel2(x_ptr, o_ptr, t_ptr, BLOCK: tl.constexpr = BLOCK):
    r = tl.arange(0, BLOCK)
    off = r[:, None] * BLOCK + r[None, :]
    tsp = r[None, :] * BLOCK + r[:, None]
    fp8 = tl.load(x_ptr + off).to(tl.float32).to(tl.float8e4nv)
    tl.store(o_ptr + off, fp8)
    tl.store(t_ptr + tsp, fp8)


def run(kernel):
    x = torch.randn(BLOCK, BLOCK, device="cuda", dtype=torch.bfloat16)
    o = torch.empty(BLOCK, BLOCK, device="cuda", dtype=torch.float8_e4m3fn)
    t = torch.empty(BLOCK, BLOCK, device="cuda", dtype=torch.float8_e4m3fn)
    kernel[(1,)](x, o, t)
    return (t.float() != o.float().T).sum().item()


print("kernel1 (tl.trans):      mismatch", run(kernel1), "/", BLOCK * BLOCK)
print("kernel2 (address store): mismatch", run(kernel2), "/", BLOCK * BLOCK)

[haok1402]

@claude review

[gemini-code-assist]

Code Review

This pull request refactors the Triton quantization kernels in deepgemm_quantize.py to use E8M0 power-of-2 scaling factors across both Hopper (SM90) and Blackwell (SM100+) architectures, while renaming the public functions to use a cleaner fp8cast_ prefix. Additionally, it addresses a Triton bug with tl.trans on FP8 tiles by writing to swapped-index addresses. The review feedback highlights that this swapped-index approach results in uncoalesced global memory stores, which can degrade performance. It is recommended to instead perform the transposition in float32 before casting to fp8 to maintain coalesced memory writes while still avoiding the Triton bug.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[haok1402]

Performance Review

Benchmarked the three transposed-store strategies on an idle H100 (do_bench). No measurable gap.

         shape |       trans (base) |      address (cur) |   bf16 (coalesced)
     2048x2048 |   11.5us  1460GB/s |   11.4us  1468GB/s |   11.4us  1476GB/s |
     4096x4096 |   30.1us  2232GB/s |   29.7us  2259GB/s |   29.8us  2254GB/s |
     7168x2048 |   27.0us  2172GB/s |   26.8us  2191GB/s |   26.3us  2232GB/s |
     2048x7168 |   27.0us  2171GB/s |   26.9us  2185GB/s |   26.8us  2192GB/s |
     8192x8192 |  101.1us  2654GB/s |  101.2us  2652GB/s |  100.7us  2666GB/s |

All three within ~1%. No regression.

Also compared end-to-end training throughput (dsv2-lite, 64 steps, steady-state):

              | before (tl.trans) | after (address-store)
tokens/sec    |       54,914      |        55,356
step-time (s) |        38.19      |         37.89

+0.81% after. Within noise. No regression at the training level either.

[haok1402]

The committed evidence is the unit suite ("48 passed"), which does validate the transposed FP8 output bit-exactly against the reference — strong for the kernel. But the claimed "dsv2-lite FP8 training run to confirm the loss curve is unchanged" is asserted only; no loss table, log, or wandb link is committed, and the standalone repro referenced as workspace/triton_fp8_trans_repro.py is not in the branch (only inlined in the PR body).

Fixed. See https://wandb.ai/pithtrain/hopper-deepgemm/workspace?nw=28pacu7obnf

[haok1402]

@claude review

[claude]

Performance Review

No performance issues found. The only perf-bearing change is the blockwise-transpose kernels dropping tl.trans(fp8_tile) for a swapped-index store (everything else is renames, docs, and a benchmark config). The theoretical concern — uncoalesced transposed FP8 stores — is well settled by the committed evidence: a do_bench microbench across 5 shapes puts address within ~1% of both tl.trans and the coalesced-bf16 reference (~2666 GB/s at 8192², i.e. near-peak HBM, so it's coalescing fine in practice), and e2e dsv2-lite (64 steps, steady state) holds throughput (54,914 → 55,356 tok/s, +0.81%, within noise) with an unchanged loss curve. This is measured on the exact Hopper emulated-E8M0 path the change touches. No device syncs, graph breaks, or comm-overlap regressions are introduced. The batched kernel isn't microbenchmarked in isolation, but it's the dominant quant workload in the dsv2-lite e2e run, so it's covered.

[claude]

Compactness Review

No compactness issues. This PR is net-subtractive on surface: it replaces the SCALING_MODE string dispatch ("e8m0"/"fp32" + static_assert) with a plain NATIVE_E8M0 boolean, drops the dead _BLOCK_K/_FP8_MAX constants and the redundant fused_/_cast_to_fp8 name fragments, and swaps tl.trans for a direct swapped-index store. No new wrappers, factories, config knobs, or defensive guards. The one added surface (__all__) matches existing convention (pithtrain/modules/checkpoint.py) and the internals are already underscore-hidden, so it does not warrant a change.

[claude]

Consistency Review

The rename (fused_*_cast_to_fp8 → fp8cast_*), the SCALING_MODE/_USE_E8M0_SCALES → NATIVE_E8M0 recast, and the tests/test_fp8_quantize_kernels.py → tests/operators/test_deepgemm_quantize.py move are propagated cleanly: every importer (linear.py, grouped_linear.py, gpt_oss.py, ring_attention.py), the test file, AGENTS.md, and docs/architecture.md are consistent with the new names and the "E8M0 power-of-2 on both archs" story. No stray old symbol, old test path, or removed constant (_BLOCK_K, _FP8_MAX) survives anywhere in the repo or skills.

One stale reference, outside this PR's diff so it has no inline thread: pithtrain/operators/deepgemm_quantize.py:658 — the _fp8cast_blockwise_transpose_kernel docstring still says the tile is stored "transposed via tl.trans, to out_t (K, M)", but this PR removed tl.trans from exactly this kernel; the transpose is now done with swapped-index stores. It should read: the quantized tile is stored to both out (M, K) and, transposed to out_t (K, M) via swapped-index stores.

[claude]

Correctness: A few lines below (line 558), this kernel still transposes an fp8 tile with tl.trans(blk_fp8) — the exact pattern the PR's own minimal repro shows corrupts (4090/16384 mismatches on H100 / Triton 3.7.1) and that the two blockwise kernels were rewritten to avoid. The same hazard also remains in _fp8cast_rowwise_transpose_kernel at line 414 (tl.trans(t_fp8)).

Both functions feed the transposed FP8 (input_t_fp8 / grad_t_fp8 / x_t_fp8) into the dgrad/wgrad GEMMs (linear.py:54,76, ring_attention.py:645), so a miscompile here silently corrupts gradients while forward/loss stay clean.

The unit suite passes for these two, so the trigger is narrower than "tl.trans corrupts fp8 tiles" as stated — plausibly it fires only when the same fp8 SSA value is both stored directly and transposed (as in the fixed kernels, which reuse fp8_tile), whereas here a distinct value feeds tl.trans. That distinction is codegen-dependent: a Triton/register-allocation change could flip these to corrupt with no test signal.

Please either apply the same swapped-index store here and at line 414, or document the precise trigger and extend the repro to this two-distinct-values shape to prove immunity. (Couldn't verify empirically — no CUDA in this review env.)

[claude]

Correctness Review

Scope: rename fused_* → fp8cast_*, recast SCALING_MODE → NATIVE_E8M0 (both branches were already pow2, so no numeric change), the tl.trans→swapped-index store fix in the two blockwise-transpose kernels, and an example-script/wandb config update. I traced the two rewritten stores (2D _fp8cast_blockwise_transpose_kernel and batched) element-by-element — index math, masks, and scale_t symmetry are all correct, and the new bit-exact tests validate the transposed FP8 outputs against the reference. All call-site renames are consistent with no stragglers.

One correctness concern (inline): the PR's stated root cause — "Triton 3.7.1 tl.trans corrupts fp8 tiles," backed by a minimal 128×128 repro — is fixed in only 2 of the 4 kernels that transpose fp8. _fp8cast_rowwise_transpose_kernel:414 and _fp8cast_rowwise_blockwise_transpose_kernel:558 still call tl.trans on 128×128 fp8 tiles, and their outputs feed the dgrad/wgrad GEMMs. The unit suite passes for them (so the real trigger is narrower than stated — likely same-value-reuse), but the discrepancy is unexplained and codegen-fragile: a Triton/reg-alloc shift could silently corrupt gradients with no test signal. Ask is to fix uniformly or document/repro the precise trigger.

Evidence is otherwise solid: 48 bit-exact unit tests plus the committed wandb dsv2-lite FP8 loss curve cover the behavior-affecting path.