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.
| run_default_fa 1 triton_kernels/test_cast_mxfp8.py | ||
| run_default_fa 1 triton_kernels/test_cast_mxfp4.py | ||
| run_default_fa 1 triton_kernels/test_grouped_gemm.py | ||
| run_default_fa 1 triton_kernels/test_blockwise_fp8.py |
There was a problem hiding this comment.
Nit: keel alphabetical order, put it before triton_kernels/test_cast
| triton.knobs.amd.scalarize_packed_fops = True | ||
| triton.knobs.amd.use_block_pingpong = True | ||
| else: | ||
| os.environ.setdefault("TRITON_HIP_USE_ASYNC_COPY", "1") |
There was a problem hiding this comment.
Why setdefault not set? With current implementation if branch forces settings while else branch only sets defaults
There was a problem hiding this comment.
Changed to force setting.
| if is_cdna4(): | ||
| set_triton_knobs_gfx950() | ||
| else: | ||
| _set_amd_knobs(enable=True) |
There was a problem hiding this comment.
From the comments, enable should be calculated form layout. Also, since knobs win/loss is gfx942 specific, it should be rather if is_cdna3() ... else
| triton.knobs.amd.scalarize_packed_fops = enable | ||
|
|
||
|
|
||
| NUM_XCDS = 8 |
There was a problem hiding this comment.
In MPX mode it will be different number. Why is it needed?
There was a problem hiding this comment.
This is used for swizzling of PIDs across dies for better L2 cache utilization.
Seems like we can't get this info programatically, so it is hardcoded.
| @triton.jit | ||
| def compute_scale_and_quant(x_tile, x_tile_abs, axis, FP8_MAX, ROUND_POW2: tl.constexpr): | ||
| x_tile_max = tl.max(x_tile_abs, axis=axis, keep_dims=True) | ||
| x_tile_max = tl.maximum(x_tile_max, 1e-4) |
There was a problem hiding this comment.
Scale is FP8_MAX / amax, this is to ensure that the minimum value of amax isn't 0 so we don't accidentally divide by 0. The value was chosen by primus turbo's implementation.
…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
| pytest.skip("FP8 parameters are not supported in debug mode.") | ||
| if IS_HIP_EXTENSION and recipe is not None and recipe.float8_block_scaling(): | ||
| pytest.skip("ROCm grouped GEMM does not yet support FP8 block scaling.") | ||
| skip_unsupported_backward_override( |
There was a problem hiding this comment.
This test — and test_padding_grouped_linear_accuracy_save_original_input below — has no use_triton parameter, so it never sets NVTE_USE_BLOCKWISE_GMM_TRITON. Removing the "ROCm grouped GEMM does not yet support FP8 block scaling." skip therefore enables Float8BlockScaling on the existing ROCm grouped GEMM path, which this PR doesn't change. The same applies to test_grouped_linear_accuracy with use_triton=False.
Was that skip already stale (i.e. the default path gained block-scaling support separately), or are these newly-enabled combinations expected to pass by way of something in this PR? If the default path still doesn't support it, the skip should stay and only be lifted for the Triton path.
There was a problem hiding this comment.
the skip was stale. Float8BlockScaling passes on test_padding_grouped_linear_accuracy, test_padding_grouped_linear_accuracy_save_original_input, and test_grouped_linear_accuracy for both use_triton=True/False. So lifting it there is correct.
Added skip for blockwise scaling in test_grouped_linear_accuracy_rocm_backends.
| 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.
| delay_wgrad_compute, | ||
| ) | ||
|
|
||
| use_blockwise_triton = os.getenv("NVTE_USE_BLOCKWISE_GMM_TRITON", "0") == "1" |
There was a problem hiding this comment.
use_blockwise_triton is computed here but never read anywhere in the function — leftover from the earlier tolerance-branching approach that was reverted. Worth dropping so it doesn't look like a live gate.
| 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. |
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: