Add ROCm Triton blockwise FP8 grouped GEMM - #716
Conversation
Implement DeepSeek-style blockwise FP8 grouped GEMM for the PyTorch backend,
selected from GroupedLinear under Float8BlockScaling when
NVTE_USE_BLOCKWISE_FP8_GROUPED_GEMM=1.
- Add _forward_blockwise_fp8/_backward_blockwise_fp8 to
transformer_engine/pytorch/module/grouped_linear.py, using 1x128 rowwise
activation quantization, 128x128 weight quantization, and a segment-padded
columnwise activation for wgrad.
- Add new Triton kernels:
- blockwise_fp8_grouped_gemm.py: persistent grouped blockwise FP8 GEMM and
variable-K wgrad kernels.
- blockwise_quantize.py: 1x128 / 128x128 blockwise FP8 quantization
kernels, including segment-padded columnwise quantize for variable M.
- Add is_cdna4() to triton_kernels/common.py for gfx950 detection.
- Add `_is_blockwise_fp8_grouped_gemm_supported()` to centralize the feature-gate checks (HIP, `NVTE_USE_BLOCKWISE_GMM_TRITON=1`, `Float8BlockScaling` layout, and unsupported orchestration options). - Simplify `_forward_blockwise_fp8` by removing inline validation and accepting a pre-uploaded `m_splits_tensor`, avoiding a blocking host-to-device copy of the split sizes. - Switch the environment gate from `NVTE_USE_BLOCKWISE_FP8_GROUPED_GEMM` to `NVTE_USE_BLOCKWISE_GMM_TRITON`.
- Extend the Triton variable-K grouped GEMM kernel to support in-place accumulation and an optional output tensor. - Add fused wgrad handling to GroupedLinear's blockwise FP8 path, including packed main-grad views and first-microbatch accumulation logic. - Remove ROCm test skips for FP8 block scaling in grouped linear tests and switch block-scaling cases to the blockwise Triton backend. - Add unit tests for the blockwise FP8 quantization and grouped GEMM Triton kernels and include them in the PyTorch CI script.
…GEMM - Generalize `_packed_main_grad_view` into `_packed_3d_view` for any sequence of contiguous 2D buffers. - Add `_expert_weights_as_3d` to return a zero-copy `[G, N, K]` view when expert weights are consecutive slices of a single buffer (e.g. `single_grouped_weight`), falling back to `torch.stack` otherwise. - Use the new helper in `_forward_blockwise_fp8` instead of always stacking weights.
- Extend `Fp8Padding.forward` with an optional `m_splits_tensor` argument. - When provided, compute and return padded split sizes as a tensor on the same device, avoiding a blocking host-to-device copy.
- Match Primus-Turbo FP8 quantize tolerances (atol=rtol=0.10) in test_blockwise_fp8.py and the grouped linear blockwise-triton path. - Remove obsolete None-output assertion in test_grouped_linear.py.
- Introduce curated fwd/dgrad autotune configs from Primus-Turbo for the grouped blockwise FP8 persistent GEMM kernel. - Add warm-up tracking so the first autotune call uses balanced group offsets, preventing the cached config from being tied to a single uneven MoE routing. - Refactor the launch helper to use the autotuned kernel and remove the hard-coded block-size heuristic.
…erances - In `grouped_linear.py`, ensure grouped tensors actually share the same underlying storage and that the storage is large enough for all slices before returning a zero-copy 3D view. - Update `test_grouped_linear.py` tolerances for the blockwise Triton path to account for two independent FP8 quantization stacks.
…GEMM - Replace CDNA4-specific checks with CDNA3 detection and remove `is_cdna4`. - Force gfx950 compiler knobs (async_copy, block_pingpong, scalarize) always on. - Gate gfx942 knobs by GEMM layout, disabling them for TN/wgrad to avoid regressions. - Reorder the blockwise FP8 test entry in the PyTorch CI script.
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| rm_s = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M_g | ||
| rn_s = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N | ||
| rn_s = tl.max_contiguous(tl.multiple_of(rn_s, BLOCK_SIZE_N), BLOCK_SIZE_N) | ||
| c_mask = (rm_s[:, None] < M_g) & (rn_s[None, :] < N) | ||
| C_ = C + m_start_g * stride_cm + rm_s[:, None] * stride_cm + rn_s[None, :] * stride_cn | ||
| tl.store(C_, c, c_mask) |
There was a problem hiding this comment.
rn_s is reduced mod N before the mask is built, so rn_s[None, :] < N is always true and the N-side of c_mask is a no-op (same for rm_s < M_g). When N % BLOCK_SIZE_N != 0 the tail tile's out-of-range lanes wrap onto low columns and are stored — but they were scaled by b_s, which is loaded once per tile at pid_n * stride_bs_n (line 369). Those wrapped columns belong to a different N scale block, so the correct values written by pid_n = 0 get overwritten with wrongly-scaled ones.
Concretely with N = 192, BLOCK_SIZE_N = 128: pid_n = 1 covers raw cols 128..255, wraps to [128..191, 0..63], and stores cols 0..63 using scale block 1.
test_grouped_gemm_fp8_blockwise_matches_dequant_ref skips out_n % BLOCK != 0, so this isn't covered by the unit tests, and _is_blockwise_fp8_grouped_gemm_supported has no shape check. Forward uses N = out_features and dgrad uses N = in_features (trans_b=False), so a GroupedLinear whose in_features or out_features isn't a multiple of 128 silently produces wrong numerics instead of falling back to the default path.
Either drop the % N / % M_g and mask on the raw indices, or add in_features % 128 == 0 and out_features % 128 == 0 to the gate.
There was a problem hiding this comment.
Added in_features % 128 == 0 and out_features % 128 == 0 gate since the kernel assumes in and out features are a multiple of 128.
| rm_s = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) | ||
| rn_s = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) | ||
| rn_s = tl.max_contiguous(tl.multiple_of(rn_s % OUT_N, BLOCK_SIZE_N), BLOCK_SIZE_N) | ||
| c_mask = (rm_s[:, None] < OUT_M) & (rn_s[None, :] < OUT_N) | ||
| C_ = C + group_idx.to(tl.int64) * stride_cg + rm_s[:, None] * stride_cm + rn_s[None, :] * stride_cn | ||
| c = acc.to(C.type.element_ty) | ||
| if ACCUMULATE: | ||
| c += tl.load(C_, mask=c_mask, other=0) | ||
| tl.store(C_, c, c_mask) |
There was a problem hiding this comment.
Same vacuous-mask pattern as the forward kernel, but here it breaks ACCUMULATE. rn_s is taken mod OUT_N, so rn_s[None, :] < OUT_N is always true.
Unlike the forward kernel the values are fine (RHS scales are indexed elementwise by rn), but when OUT_N % BLOCK_SIZE_N != 0 the tail tile wraps onto low columns that another tile also owns, and c += tl.load(C_) followed by tl.store adds that contribution twice — on top of a genuine read-modify-write race between the two tiles.
_bwd_autotune_configs() includes BLOCK_SIZE_N = 256, and OUT_N here is in_features, so e.g. in_features = 384 hits it. Only reachable with fuse_wgrad_accumulation, and test_variable_k_wgrad only uses k in {128, 256} — both multiples of every candidate BLOCK_SIZE_N, so it isn't covered.
Masking on the un-wrapped pid_n * BLOCK_SIZE_N + tl.arange(...) (as is already done for rm_s on line 686) fixes it.
There was a problem hiding this comment.
Removed % OUT_N and let the c_mask take care of bounds.
| # Quantize: activation rowwise (1x128 along K), weights 128x128. | ||
| a_row, a_srow = quantize_fp8_blockwise(a, dt, axis=1, block_size=128) | ||
| b_fp8, b_scale = quantize_fp8_blockwise_weight(w, dt, block_size=128) |
There was a problem hiding this comment.
Float8BlockScaling defines fp8_quant_fwd_inp / fwd_weight / bwd_grad = QParams(power_2_scale=not use_f32_scales, amax_epsilon=0.0) — i.e. power-of-2 scales by default, unless NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1.
The Triton launchers support this (pow2=), but neither this call site nor _backward_blockwise_fp8 ever passes it, so the blockwise Triton path always uses fp32 scales. That means the same recipe produces different quantization depending on whether NVTE_USE_BLOCKWISE_GMM_TRITON is set, and it's plausibly a large part of why the test_grouped_linear_accuracy tolerances had to be widened against the sequential reference.
Threading recipe.fp8_quant_fwd_inp.power_2_scale (and the weight/grad equivalents) into pow2= would make the two paths agree.
There was a problem hiding this comment.
Integrated Qparams values to enable or disable power of 2 scales
| def _set_triton_knobs_gfx950() -> None: | ||
| """Force-on AMD compiler knobs for gfx950 (async_copy, block_pingpong, scalarize).""" | ||
| global _KNOBS_SET | ||
| if _KNOBS_SET: | ||
| return | ||
| _KNOBS_SET = True | ||
| os.environ["TRITON_HIP_USE_ASYNC_COPY"] = "1" | ||
| os.environ["AMDGCN_SCALARIZE_PACKED_FOPS"] = "1" | ||
| os.environ["TRITON_HIP_USE_BLOCK_PINGPONG"] = "1" | ||
| if hasattr(triton, "knobs") and hasattr(triton.knobs, "amd"): | ||
| triton.knobs.amd.use_async_copy = True | ||
| triton.knobs.amd.scalarize_packed_fops = True | ||
| triton.knobs.amd.use_block_pingpong = True | ||
|
|
||
|
|
||
|
|
||
| def _set_triton_knobs_gfx942(enable: bool = True): | ||
| """Set AMD Triton knobs on gfx942 (CDNA3). | ||
|
|
||
| ``use_async_copy`` / ``scalarize_packed_fops`` help NT/NN but regress | ||
| TN/wgrad ~5-8% on gfx942, so callers pass ``enable`` from layout. | ||
| """ | ||
| if hasattr(triton, "knobs") and hasattr(triton.knobs, "amd"): | ||
| triton.knobs.amd.use_async_copy = enable | ||
| triton.knobs.amd.scalarize_packed_fops = enable | ||
|
|
||
|
|
||
| def _apply_amd_compiler_knobs(*, is_tn: bool) -> None: | ||
| """gfx942: knobs from GEMM layout. Else (gfx950): always-on gfx950 knobs.""" | ||
| if is_cdna3(): | ||
| _set_triton_knobs_gfx942(enable=not is_tn) | ||
| else: | ||
| _set_triton_knobs_gfx950() |
There was a problem hiding this comment.
These knobs are process-global, and no other Triton kernel in the tree touches them. _set_triton_knobs_gfx942 flips triton.knobs.amd.use_async_copy / scalarize_packed_fops on every call to a public entrypoint, so whichever GEMM ran last silently determines the codegen of the next unrelated Triton kernel that compiles (triton_kernels/cast.py, gmm/, the MXFP8 kernels...). On gfx950 the _KNOBS_SET latch makes it one-way for the lifetime of the process.
The os.environ[...] writes are also likely dead: triton.knobs reads the environment at import time, which is presumably why the triton.knobs.amd.* assignments were added next to them. If so, they're worth dropping rather than leaving as a misleading no-op.
Saving and restoring the previous values around the launch (or a small context manager) would keep the effect scoped to these kernels.
There was a problem hiding this comment.
Created a context helper which enables the knobs before kernel is run and disables it once it gets over.
| loop_k = tl.cdiv(K, BLOCK_SIZE_K) | ||
| if not EVEN_K: | ||
| loop_k -= 1 | ||
| tl.assume(loop_k > 1) |
There was a problem hiding this comment.
tl.assume lowers to llvm.assume, so a false predicate is UB and lets the backend drop or miscompile the loop.
loop_k is cdiv(K, 128), minus 1 when !EVEN_K. It is 1 for K = 128 — which test_grouped_gemm_fp8_blockwise_matches_dequant_ref exercises directly via the ([256], 128, 128) case — and 0 for K < 128 with EVEN_K=False. Both violate loop_k > 1.
| tl.assume(loop_k > 1) | |
| tl.assume(loop_k >= 0) |
| if use_blockwise_triton: | ||
| # Sequential Linear uses TE Float8BlockQuantizer + TE GEMM; this path | ||
| # uses Triton quant + grouped GEMM. Budget two independent FP8 stacks. | ||
| atol, rtol = 0.25, 0.12 | ||
| for o, o_ref in zip(outputs, outputs_ref): | ||
| torch.testing.assert_close(o, o_ref, rtol=rtol, atol=atol) | ||
| if use_blockwise_triton: | ||
| mag = max(float(o.detach().abs().max()), float(o_ref.detach().abs().max())) | ||
| tensor_atol = max(atol, 0.05 * mag) | ||
| torch.testing.assert_close(o, o_ref, rtol=rtol, atol=tensor_atol) | ||
| else: | ||
| torch.testing.assert_close(o, o_ref, rtol=rtol, atol=atol) |
There was a problem hiding this comment.
Two things make this assertion weaker than it looks:
-
atol = max(0.25, 0.05 * max|out|)withrtol = 0.12is close to vacuous — a 5%-of-peak absolute budget would pass even if an entire output tile were mis-scaled (e.g. theN % 128wrap case). Comparing against a blockwise-FP8 reference (quantize → dequantize → matmul, exactly whattest_blockwise_fp8.pyalready builds) instead of the high-precision sequentialLinearwould let the tolerance stay tight and actually exercise the kernel rather than the FP8 format. -
use_blockwise_tritontracks the env var, not whether the Triton path was actually selected._is_blockwise_fp8_grouped_gemm_supportedrejectsuse_bias,save_original_input,unpad_output,actual_m_splits, etc. — sobias=True(half the matrix) runs the default TE path yet still gets the loosened tolerances, silently weakening coverage of a path this PR doesn't touch. Gating on the same predicate as the module would keep the two in sync.
There was a problem hiding this comment.
Reverted back to original tolerances used by use_triton
| packed = _GroupedLinear._packed_3d_view(weights) | ||
| if packed is not None: | ||
| if packed.dtype != dtype: | ||
| packed = packed.to(dtype) | ||
| return packed if packed.is_contiguous() else packed.contiguous() | ||
| return torch.stack([wt.to(dtype).contiguous() for wt in weights], 0).contiguous() |
There was a problem hiding this comment.
Does this hold up when fp8_model_params=True? test_grouped_linear_accuracy parametrizes fp8_model_params × use_triton, so weights here can be Float8BlockwiseQTensor, which QuantizedTensor.__new__ builds via torch.Tensor._make_wrapper_subclass — those have no real storage, so _packed_3d_view's g0.untyped_storage().size() (line 411) is at best returning 0 and at worst raising.
If it does fall through to line 439, wt.to(dtype) dequantizes an already-FP8 parameter which is then re-quantized blockwise by quantize_fp8_blockwise_weight — double quantization, which would show up as exactly the kind of error the widened test tolerances absorb.
Worth either handling QuantizedTensor weights explicitly (use the existing rowwise_data + scales rather than round-tripping) or excluding fp8_model_params in _is_blockwise_fp8_grouped_gemm_supported until it is.
There was a problem hiding this comment.
Excluded fp8 model parameters in _is_blockwise_fp8_grouped_gemm_supported since it's not supported currently. Default path should be good enough.
| # ----------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| def quantize_fp8_blockwise_dual(x: torch.Tensor, dtype: torch.dtype, block_size: int = 128, pow2: bool = False): |
There was a problem hiding this comment.
This file and blockwise_fp8_grouped_gemm.py haven't been run through Black. .pre-commit-config.yaml / qa/format.sh apply black --line-length=100, and there are 10 over-length lines here (61, 99, 100, 234, 245, 247, 258, 277, 302, 320) plus 447, 470 and 637 in grouped_linear.py.
bash qa/format.sh before merge will keep the format job green.
Claude review —
|
…iler-knob scoping - Rename blockwise FP8 helpers to `..._triton` and tighten `_is_blockwise_fp8_triton_grouped_gemm_supported`: require 128-aligned `in_features`/`out_features` and reject `fp8_weights` to avoid double quantization. - Add `pow2` rounding flags to blockwise quantization helpers and plumb them through the grouped-linear forward/backward. - Scope AMD Triton compiler knobs with a context manager so gfx950/gfx942 overrides don't leak into unrelated kernels; wrap the grouped GEMM launch instead of setting globals. - Add unit tests for the 128-alignment gate and wgrad tail-tile coverage; drop the blockwise-triton-specific accuracy tolerance branch and skip FP8-block-scaling cases for CUTLASS/HipKittens/CK ROCm backends.
- Allow already-quantized blockwise FP8 weights (`fp8_model_params`) to be consumed directly in the grouped-linear blockwise Triton path, avoiding dequantize/re-quantize double quantization. - Add `Float8BlockwiseQTensor` wrappers for grouped weight and activation operands, and a zero-copy packed 2D view over per-expert weight tensors. - Replace the separate rowwise/colwise dual-quantize kernel with a unified grouped kernel that can emit both rowwise activation and segment-padded columnwise wgrad operands in one pass. - Rename the raw grouped GEMM kernels to `..._raw` and add public wrappers that extract FP8 data/scales from `Float8BlockwiseQTensor`. - Add `_vk_group_offs` to `Float8BlockwiseQTensorStorage` for variable-K grouped wgrad segment offsets. - Update tests to use the new QTensor APIs and remove the `fp8_weights` rejection gate.
| # Members are views, not the Parameter, so mirror the grad state the | ||
| # autograd Function reads off ``weights[i]``: ``requires_grad`` (gates | ||
| # wgrad), ``main_grad`` (per-expert views into the grouped | ||
| # fuse-accumulation buffer), and ``overwrite_main_grad``. | ||
| want_grad = grouped_weight.requires_grad | ||
| main_grad = getattr(grouped_weight, "main_grad", None) | ||
| per_expert_main_grad = None | ||
| if main_grad is not None: | ||
| per_expert_main_grad = main_grad.view( | ||
| self.num_gemms, self.out_features, self.in_features | ||
| ) | ||
| for i, w in enumerate(weight_tensors): | ||
| if w.requires_grad != want_grad: | ||
| w.requires_grad_(want_grad) | ||
| if per_expert_main_grad is not None: | ||
| w.main_grad = per_expert_main_grad[i] | ||
| if hasattr(grouped_weight, "overwrite_main_grad"): | ||
| w.overwrite_main_grad = grouped_weight.overwrite_main_grad |
There was a problem hiding this comment.
Upstream compatibility / scope: this is an unguarded behavior change to shared, CUDA-reachable code.
_get_weight_tensors() is on the common GroupedLinear path (its result feeds _GroupedLinear.apply(...) at line 2310), and none of this new logic is gated on ROCm or on the blockwise-Triton path. It runs on CUDA too whenever single_grouped_weight is active. Per the fork rules, a change to a code path CUDA also executes is either (a) a generic bug fix that must be called out in the PR description so it can be upstreamed, or (b) something that needs a guard — right now it's neither.
Two concrete concerns beyond the classification:
- The loop mutates the cached
grouped_weight.quantized_tensorsviews in place (requires_grad_,main_grad,overwrite_main_grad). Sincequantized_tensorsis cached on the Parameter, these mutations persist across forwards and across modules that share the storage.requires_grad_()on a non-leaf view raisesRuntimeError, so this depends onsplit_into_quantized_tensors()always returning leaves. main_grad.view(self.num_gemms, self.out_features, self.in_features)assumes the fuse-accumulation buffer is exactly that shape and contiguous. Ifmain_gradis allocated flat with padding (or as a differently-shaped grouped buffer) this raises rather than falling back.
Same question applies to the requires_grad_ mirroring in _get_bias_tensors (lines 2507-2510).
If this is a real bug on the single_grouped_weight path, could you note it in the PR description as a generic fix and add a test that exercises it? I don't see coverage for it in this PR, and NVTE_GROUPED_LINEAR_SINGLE_PARAM is off by default, so a regression here would land silently.
There was a problem hiding this comment.
Added coverage for this issue, this is a bug which is present in ROCm and CUDA for single_grouped_weight. This is fixed in latest TE upstream using different technique. NVIDIA/TransformerEngine@2d80391#diff-e52c6ddc8c0f4d20cb5aa832e92dd3794687bac1f828e367b526f17f524238a5
Will wait for upstream integration, until this this fix is good enough.
| else: | ||
| w = _GroupedLinear._expert_weights_as_3d(weights, activation_dtype) | ||
| qw = quantize_fp8_blockwise_grouped_weight_qtensor(w, dt, pow2=pow2_w) |
There was a problem hiding this comment.
is_first_microbatch is accepted by this forward (line 535) but only ever stored on ctx (line 645) — it never gates weight quantization, and line 659 returns new_workspaces = [None] * num_gemms. The result is that high-precision weights get re-quantized on every microbatch here, whereas the default grouped path (lines 1021-1030) quantizes once when is_first_microbatch is True and reuses the cached workspace via self._fp8_workspaces afterwards.
Correctness is fine, but for gradient accumulation over N microbatches this is N× the weight-quantization work relative to the non-Triton path, which partly works against the point of the fused kernel. The fp8_model_params branch above already sidesteps it; the else branch here is the one that pays.
Not a blocker — but if the caching is deliberately deferred, a short comment saying so would keep the next reader from assuming it's an oversight.
There was a problem hiding this comment.
Added weight workspace caching feature.
Re-reviewFocused on the two commits added since the last review ( Round-1 findings verified fixed — the N-mask/B-scale mismatch (now gated on New findings (3 inline comments): the main one is that Copyright headers: OK — all 8 files in scope carry correct AMD lines, and no NVIDIA year ranges were altered. Verdict: kernel-correctness blockers from round 1 are resolved; one scope/upstream-compat item plus two nits remain. |
…8 weight quantization - Mirror `requires_grad`, `main_grad`, and `overwrite_main_grad` from the grouped Parameter onto per-expert split views in `_get_weight_tensors` / `_get_bias_tensors`, guarding `requires_grad_` to leaves and validating that `main_grad` can alias in place. - Cache the packed blockwise FP8 quantized weight across microbatches in the Triton grouped-GEMM path using `weight_workspaces`. - Add a ROCm autouse `EnvVarCleaner` fixture to snapshot/restore grouped-GEMM Triton backend env vars and prevent leakage between tests, plus a unit test verifying the grad-state mirroring behavior.
| update_ws = is_first_microbatch is None or is_first_microbatch | ||
| cached_qw = weight_workspaces[0] if weight_workspaces else None | ||
| if not update_ws and isinstance(cached_qw, Float8BlockwiseQTensor): | ||
| qw = cached_qw | ||
| else: | ||
| w = _GroupedLinear._expert_weights_as_3d(weights, activation_dtype) | ||
| qw = quantize_fp8_blockwise_grouped_weight_qtensor(w, dt, pow2=pow2_w) | ||
| if cache_weight: | ||
| new_workspaces[0] = qw |
There was a problem hiding this comment.
The new cache shares self._fp8_workspaces keys with the default path, but the two store incompatible layouts under the same key.
forward (line 2350) builds weight_workspaces = [self._fp8_workspaces.get(f"weight{i}") ...] and stores back into self._fp8_workspaces[f"weight{i}"] (line 2391) regardless of which path ran. Here new_workspaces[0] is the packed [G*out_features, in_features] QTensor from quantize_fp8_blockwise_grouped_weight_qtensor, whereas the default path stores expert 0's per-expert [out_features, in_features] workspace under that same "weight0" key.
Neither read validates shape:
- Triton→default:
quantize_weightcalls_is_weight_workspace_valid, whoseisinstancechain coversFloat8TensorStorage/MXFP8TensorStorage/NVFP4TensorStorageonly — aFloat8BlockwiseQTensorfalls through toreturn True, thenworkspace.quantize_(tensor)runs against aG×-too-large buffer. - default→Triton: the
isinstance(cached_qw, Float8BlockwiseQTensor)check on line 617 passes for the per-expert workspace, soqwbecomes a single expert's weight used as the packed operand for allGexperts.
This is reachable because _is_blockwise_fp8_triton_grouped_gemm_supported reads per-call inputs — unpad_output and actual_m_splits are forward() arguments and cpu_offloading is is_cpu_offload_enabled(), a context-manager global — so one module can alternate paths across microbatches while is_first_microbatch keeps the cache live.
A distinct key (e.g. "blockwise_packed_weight") would decouple the two, and a shape check on the cached tensor before line 618 would make a mismatch loud rather than silent.
Worth noting this is untested either way: nothing in tests/pytorch/test_grouped_linear.py passes is_first_microbatch, so cache_weight = is_first_microbatch is not None (line 2349) is always False in CI and the cache-hit branch on lines 617-618 never executes. test_cuda_graphs.py / test_float8_current_scaling_exact.py exercise this for Linear, but not for GroupedLinear under Float8BlockScaling.
| ) from e | ||
| for i, w in enumerate(weight_tensors): | ||
| # ``requires_grad_`` only works on leaves. | ||
| if w.requires_grad != want_grad and w.is_leaf: |
There was a problem hiding this comment.
The and w.is_leaf guard trades a loud failure for a silent one.
Previously a non-leaf view raised RuntimeError from requires_grad_. Now the mirroring is skipped, so a stale requires_grad on the view survives. The case that matters is freezing: grouped_weight.requires_grad_(False) while the views still carry requires_grad=True. _forward_blockwise_fp8_triton reads ctx.weight_requires_grad = weights[0].requires_grad (line 654), and the default path gates wgrad the same way — so a frozen expert would still get a wgrad computed, which is exactly the divergence this loop exists to prevent.
Since the guard is protecting against a state the code can't currently repair, is_leaf is False combined with a requires_grad mismatch is arguably worth raising on rather than passing over silently.
| # requires_grad is mirrored (not merely coincidental): flip the Parameter and refetch. | ||
| grouped_linear.weight.requires_grad_(False) | ||
| for w in grouped_linear._get_weight_tensors(): | ||
| assert w.requires_grad is False |
There was a problem hiding this comment.
The comment says "not merely coincidental", but as written this assertion is coincidental.
GroupedTensorStorage initialises quantized_tensors = None and only populates it on quantize, so for the bf16 grouped weight in this test _get_weight_tensors takes the split_into_quantized_tensors() branch on every call. After requires_grad_(False), the refetch re-derives fresh views from the already-frozen Parameter, which naturally come back requires_grad=False — the mirroring loop on line 2504 never has a mismatch to act on. The same applies to the first loop (line 2035): both sides are True because the views inherit it, not because it was mirrored.
The main_grad aliasing and overwrite_main_grad assertions are genuinely load-bearing; only the requires_grad ones aren't.
Setting grouped_linear.weight.quantized_tensors = weights before the flip would pin the cached-view path and make this exercise the assignment. That also surfaces the is_leaf question from the production comment — if the cached views turn out to be non-leaf, this test would fail rather than pass silently.
Re-reviewScoped to Round-2 findings verified fixed — the unused Verdict: comment. Three new findings, all in this commit — no blockers on the kernels themselves:
Copyright headers: OK — all 8 in-scope files carry correct AMD lines with 2026 end-years, and the NVIDIA lines are unchanged. |
- On `is_first_microbatch=False`, validate the cached blockwise FP8 weight workspace is a `Float8BlockwiseQTensor` of the expected packed shape and raise a clear `RuntimeError` if it is incompatible, instead of silently re-quantizing or using a stale buffer. - Only stage the freshly quantized weight for write-back when a new quantization actually occurs. - Replace the `single_grouped_weight` grad-state mirroring test with a ROCm-only test that verifies the blockwise FP8 packed weight is quantized once, reused across microbatches without re-quantization, remains numerically identical, and fails loudly on an incompatible cached workspace.
- Clamp block amax to 1e-4 in `blockwise_quantize.py` kernels to prevent divide-by-zero/inf scales on all-zero or tiny blocks, matching Primus-Turbo. - Add a comment explaining the hardcoded `NUM_XCDS=8` used for round-robin PID swizzling in ROCm blockwise FP8 grouped GEMM, including SPX/MPX mode caveats. - Update the ROCm grouped-linear test skip message to list only HipKittens/CK as FP8-block-scaling-unsupported backends.
- Update the return tuple in `_GroupedLinear.backward` to include `None` grads for `out` and `dgrad_out`, matching the saved tensor signature used by the forward pass.
- In `get_align_size_for_quantization`, return 128 for HIP blockwise FP8 scaling on CDNA3 (gfx942) when not using the Triton grouped-GEMM path, since the native blockwise FP8 GEMM requires K padded to 128. - Keep the existing 16-byte alignment when the Triton path is enabled, as it handles K-padding internally.
3 participants
Description
Opt-in ROCm Triton path for blockwise FP8 grouped GEMM on PyTorch
GroupedLinear, used for MoE expert GEMMs underFloat8BlockScaling.The default grouped GEMM path does not implement this blockwise layout (activation 1×128 along K, weights 128×128, columnwise segment-padded operand for variable-K wgrad). This PR adds Triton quantization + persistent grouped GEMM kernels and selects them from
GroupedLinearwhenNVTE_USE_BLOCKWISE_GMM_TRITON=1and the call is otherwise compatible. Unsupported configs (bias, fused-pad /unpad_output,save_original_input,backward_override, cpu offloading, debug, non-128-aligned features, non-matching recipe dims, etc.) fall back to the existing path instead of raising.Quantization is always from the original high-precision tensors (no double-quantization). When the caller passes a device
m_splits_tensor, split lengths are taken from that tensor so the path does not issue a blocking H2D of pageable CPUm_splits.Fixes # (issue)
Type of change
single_grouped_weightgrad-state fix, see Changes.Changes
Please list the changes introduced in this PR:
_forward_blockwise_fp8_triton/_backward_blockwise_fp8_tritonon_GroupedLinear, gated by_is_blockwise_fp8_triton_grouped_gemm_supported()andNVTE_USE_BLOCKWISE_GMM_TRITON=1underFloat8BlockScaling(x 1D / w 2D / grad 1D).triton_kernels/blockwise_quantize.py(1×128 activation / 128×128 weight / segment-padded columnwise) andtriton_kernels/blockwise_fp8_grouped_gemm.py(persistent grouped GEMM and variable-K wgrad).quant_fp8_blockwise_grouped_kernelproduces both the rowwise (1×128 along K, forward/dgrad) and the segment-padded columnwise (1×128 along M, variable-K wgrad) operands from a single HBM read when both are needed (training), instead of two separate passes. Aquantize_fp8_blockwise_act_operandsdispatcher selects fused-both / rowwise-only / columnwise-only based on what the caller requests, and returns a singleFloat8BlockwiseQTensorcarrying both operands (columnwise operand + its paddedvk_group_offsride in the QTensor's columnwise slots).@triton.autotuneon the persistent forward/dgrad kernel (keyed onG, N, K) and the variable-K wgrad kernel (keyed onG, OUT_M, OUT_N) over a curated config set. The first call per key warms up on a balancedgroup_offsso the cached config is not locked to one uneven MoE routing.accumulateflag, so wgrad adds directly intomain_gradin-place instead of producing a separate gradient that is summed afterward (removes an extra tensor and add pass).is_first_microbatch/self._fp8_workspacesreuse: the high-precision weight is quantized once (first microbatch) and the packed[G*N, K]Float8BlockwiseQTensoris reused for the rest, avoiding N× re-quantization during gradient accumulation.is_first_microbatch=Nonekeeps the previous always-quantize behavior; thefp8_model_paramscase already skips re-quantization.[G, N, K]weight buffer (single_grouped_weight) to avoid per-forward concatenation of separate per-expert weight tensors. UpstreamGroupedLinearshould pass this through to use it.GroupedLinear._get_weight_tensors()/_get_bias_tensors()split the grouped Parameter into per-expert views (this fork's transitional pre-#3224 design; upstream returns[self.weight]). Those views did not carry the grad state the autograd Function reads offweights[i], sosingle_grouped_weighttraining was broken on every backend:requires_gradwasFalse(wgrad never ran) andweights[i].main_gradraisedAttributeErrorunderfuse_wgrad_accumulation. Now mirrorrequires_grad, per-expert aliasing views ofmain_grad, andoverwrite_main_gradfrom the grouped Parameter. This is transitional and disappears if the fork adopts upstream's[self.weight]design (Improve device-init grouped linear module with single grouped weight support NVIDIA/TransformerEngine#3224). Note:NVTE_GROUPED_LINEAR_SINGLE_PARAMis off by default, so this path was previously untested here; a regression test is added (see below).m_splits_tensorwhen present; otherwise copy CPUm_splitsto device.m_splitslist, forcing a D2H sync. ExtendFp8Padding.forwardwith an optionalm_splits_tensor; when provided, return the padded split sizes as an on-device tensor (rounded up on-device, no D2H/H2D) so callers keep split sizes on the GPU. Default list behavior is unchanged.Checklist: