[PyTorch] Enable fused FP8 block-scaling path in GroupedLinear module#3171
[PyTorch] Enable fused FP8 block-scaling path in GroupedLinear module#3171denera wants to merge 10 commits into
Conversation
bbfa1da to
b35924c
Compare
|
/te-ci core pytorch |
Greptile SummaryThis PR integrates the grouped FP8 block-scaling quantization kernels into the fused
Confidence Score: 5/5Safe to merge — the changes are well-scoped, gate-guarded by an env flag, and backed by new tests for both the happy path (SM90) and the explicit error path (Blackwell). All three mutually-dependent pieces — kernel-level pow-2 scale support, quant_config threading through the C++ API, and the two-slot persistent workspace fix — are coherent and correctly wired together. The Blackwell guard prevents silent fall-through when the fused path is explicitly requested. The test suite adds monkeypatching to confirm the fused code path is actually exercised, not vacuously matched by the reference path. No files require special attention; the most load-bearing logic (workspace caching and the Blackwell raise) is straightforward and well-commented. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Caller
participant GroupedLinear
participant group_quantize
participant bgrad_group_quantize
participant general_grouped_gemm
Caller->>GroupedLinear: forward(x, m_splits)
GroupedLinear->>GroupedLinear: _can_use_grouped_tensor_path check
GroupedLinear->>group_quantize: quantize input FP8BS rowwise+columnwise
GroupedLinear->>general_grouped_gemm: fprop TN GEMM using slot 0 workspace
GroupedLinear-->>Caller: output
Caller->>GroupedLinear: backward(dy)
GroupedLinear->>GroupedLinear: "set_usage rowwise=requires_dgrad columnwise=weights_req_grad"
alt use_bias AND requires_dgrad
GroupedLinear->>bgrad_group_quantize: quantize dy with fused dbias reduction
Note over bgrad_group_quantize: force_pow_2_scales threaded via quant_config
else fuse_bgrad is False
GroupedLinear->>group_quantize: quantize dy columnwise only
GroupedLinear->>GroupedLinear: compute_grouped_dbias Triton fallback
end
GroupedLinear->>general_grouped_gemm: dgrad NN GEMM using slot 0 workspace
GroupedLinear->>general_grouped_gemm: wgrad NT GEMM using slot 1 workspace isolated
GroupedLinear-->>Caller: dx dw dbias
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant Caller
participant GroupedLinear
participant group_quantize
participant bgrad_group_quantize
participant general_grouped_gemm
Caller->>GroupedLinear: forward(x, m_splits)
GroupedLinear->>GroupedLinear: _can_use_grouped_tensor_path check
GroupedLinear->>group_quantize: quantize input FP8BS rowwise+columnwise
GroupedLinear->>general_grouped_gemm: fprop TN GEMM using slot 0 workspace
GroupedLinear-->>Caller: output
Caller->>GroupedLinear: backward(dy)
GroupedLinear->>GroupedLinear: "set_usage rowwise=requires_dgrad columnwise=weights_req_grad"
alt use_bias AND requires_dgrad
GroupedLinear->>bgrad_group_quantize: quantize dy with fused dbias reduction
Note over bgrad_group_quantize: force_pow_2_scales threaded via quant_config
else fuse_bgrad is False
GroupedLinear->>group_quantize: quantize dy columnwise only
GroupedLinear->>GroupedLinear: compute_grouped_dbias Triton fallback
end
GroupedLinear->>general_grouped_gemm: dgrad NN GEMM using slot 0 workspace
GroupedLinear->>general_grouped_gemm: wgrad NT GEMM using slot 1 workspace isolated
GroupedLinear-->>Caller: dx dw dbias
Reviews (5): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
| # Whether this operation supports the FP8 block-scaling recipe. Most basic | ||
| # ops have not been audited for it, so they opt in explicitly. | ||
| supports_float8_block_scaling: bool = False |
There was a problem hiding this comment.
Putting something so specific in the base class is quite hacky. As far as I can tell, this is only used to get around a recipe check in reset_recipe_state, when it would be better to just remove that check.
| # Whether this operation supports the FP8 block-scaling recipe. Most basic | |
| # ops have not been audited for it, so they opt in explicitly. | |
| supports_float8_block_scaling: bool = False |
There was a problem hiding this comment.
Right, this becomes unnecessary if simply adding FP8BS into the _quantization_list works out of the box.
| * FP8 block scaling runs the graph-safe flow on Hopper (CC 9.0) only, including | ||
| power-of-2 scales. On other architectures it falls back to the split-quantize |
There was a problem hiding this comment.
Nit: If power-of-2 doesn't matter, why bring it up? It seems like Claude struggled with this, and now it brings it up everywhere even when it won't be helpful to downstream users.
There was a problem hiding this comment.
Split-quantize path supports power-of-2 constrained scales and disallowing that in the fused-quantize path means we either silently ignore the config flag for this in the quantizer or we introduce a bunch of ugly code branching. I wasn't a fan of either of those options, so I enabled support power-of-2 scales in the kernel because: 1) it's trivial, 2) makes the split- and fused-quantize paths numerically equivalent for every config, and 3) lays down the foundation for enabling grouped FP8BS on Blackwell+ using the same MXFP8 scale-broadcast trick we use for non-grouped FP8BS in a separate/future PR.
| dgrad (``slot`` 0) and wgrad (``slot`` 1) must NOT share a single cuBLAS | ||
| workspace. The grouped GEMM keeps a grid-synchronization flag in the first | ||
| bytes of the workspace (cuBLAS emits a memset that zeros it before each | ||
| matmul). When two grouped matmuls in a replayed CUDA graph share one | ||
| workspace, that flag is aliased, and the second matmul's cooperative kernel | ||
| deadlocks on its second replay with cuBLAS 13.6 (and corrupts the last | ||
| expert's wgrad on cuBLAS < 13.6) -- the two matmuls are strictly stream- | ||
| ordered, so this is shared-workspace reuse, not concurrent co-scheduling. | ||
| Giving each its own persistent buffer removes the hazard; each slot is a | ||
| single persistent allocation, so CUDA-graph capture safety is preserved. | ||
| These are dedicated to the grouped path and are intentionally not shared with | ||
| the non-grouped GEMM workspace. |
There was a problem hiding this comment.
Wait, is this true? This is saying that the kernel will fail if the workspace has certain values. This means that it's not safe to allocate with torch.emtpy and we need to allocate with torch.zero.
There was a problem hiding this comment.
This is true and I have a deterministic reproducer for it, but it's not really about the torch.empty vs torch.zero allocation issue.
It is very specifically affecting only the 1Dx1D grouped FP8BS GEMM kernel (only used in WGRAD, while FWD and DGRAD use 1Dx2D and 2Dx1D), and only under graph replay. DGRAD GEMM is leaving behind stuff in the cuBLAS workspace that's corrupting/deadlocking the WGRAD GEMM, but it manifests only in the 2nd graph replay (first execution after capture is fine).
Zeroing out the workspace in between DGRAD and WGRAD GEMMs also resolves the issue, but I opted to utilize workspace isolation instead to avoid introducing extra work of zeroing out a shared workspace into the backward pass.
The catch is that the deterministic reproducer for this is still PyTorch-dependent. I cannot reproduce this in pure C++ using just the TE/common API. So at least by process of elimination, there seems to be an interaction I don't quite understand between graph capture/replay and PyTorch's allocation pool.
| _quantization_list.append("nvfp4_4over6") | ||
|
|
||
| # GroupedLinear is currently the only basic op that supports FP8 block scaling. | ||
| _grouped_linear_quantization_list: list[Optional[str]] = list(_quantization_list) |
There was a problem hiding this comment.
Why not just add FP8 block scaling to _quantization_list?
| _quantization_list.append("nvfp4_4over6") | ||
|
|
||
| # GroupedLinear is currently the only basic op that supports FP8 block scaling. | ||
| _grouped_linear_quantization_list: list[Optional[str]] = list(_quantization_list) |
There was a problem hiding this comment.
Have we tried just adding FP8 block scaling to _quantization_list? The quantization infrastructure should be general enough for it to work out of the box. If it works automatically without non-trivial changes, then I don't see a reason not to include it in the tests.
There was a problem hiding this comment.
No, I haven't attempted this but you're right, if it's that simple, there's no reason to not try and include it if it works. I'll give it a try and see what happens.
There was a problem hiding this comment.
Float8BlockwiseQTensor isn't propagating dtype recasts (e.g.op.half()) and does not support .view() reshape, so it looks like there's meaningful work that needs to be done before we can broadly enable FP8BS in fusible ops.
I'd suggest maybe scaling this PR back to strip out fusible ops support entirely, implement integration strictly in the GroupedLinear module, and shift the fusible ops FP8BS (both non-grouped and grouped) support to a separate PR. Is that acceptable?
CC: @vthumbe1503 @ptrendx
| workspace_setup = _get_grouped_gemm_setup_workspace(device.index, num_tensors) | ||
| # dgrad and wgrad must use distinct persistent cuBLAS workspaces: sharing one | ||
| # deadlocks under CUDA-graph replay on cuBLAS 13.6 (see _get_grouped_cublas_workspace). | ||
| workspace_cublas = _get_grouped_cublas_workspace(device.index, 1 if is_discrete_out else 0) |
There was a problem hiding this comment.
For single_grouped_weight case, we would need single grouped tensor output instead discrete output even for wgrad.
So this logic wont give different workspaces for wgrad and dgrad in that case
9b4518d to
5392125
Compare
…ng quantize The default Float8BlockScaling recipe constrains scales to powers of 2, so the fused grouped path must honor the flag to stay numerically consistent with the unfused path. Thread a runtime pow_2_scales argument through the grouped quantize kernels (the shared scale helper already implements the rounding) and drop the force_pow_2_scales rejections. Also add a quantization-config parameter to nvte_group_quantize_dbias, which previously had no way to receive force_pow_2_scales or amax_epsilon on the bgrad path. Signed-off-by: Alp Dener <adener@nvidia.com>
…r module Admit Float8BlockQuantizer in the fused GroupedTensor path on Hopper. The existing usage flags already match the Hopper TN-only mapping and the grouped GEMM selects transposed columnwise storage for NN/NT layouts, so only the path predicate changes. The fused path is an explicit opt-in via NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM, so raise on Blackwell (SM100/SM110) instead of silently falling back; the fused path has no MXFP8-broadcast emulation. Extend the fused dbias path (tex.bgrad_group_quantize) to FP8 block scaling when dgrad is required (dbias is computed in the rowwise pass). Add fp8_block_scaling to the fused-path tests with a Hopper-only gate, assert the fused path engages via a group_quantize spy, and add a Blackwell error-path test. Signed-off-by: Alp Dener <adener@nvidia.com>
Replace the blanket FP8 block-scaling rejection in BasicOperation.reset_recipe_state with a per-op supports_float8_block_scaling flag and opt in the GroupedLinear op. Mirror the module-path predicate and fused-bgrad changes; since the graph-safe flow is default-on here (no env-var opt-in), other architectures fall back to the split-quantize flow instead of raising. Force use_split_accumulator=True for FP8 block-scaling operands in general_grouped_gemm_for_grouped_tensor, matching non-grouped general_gemm: cuBLAS has no fast-accum FP8 block-scaling algorithm, so the ops-layer forward failed algo selection without it. Add fp8_block_scaling coverage to the ops GroupedLinear tests. The CUDA-graph-safe test skips it for now: the replayed wgrad for the last expert diverges between replays depending on process allocation history; under investigation. Graph capture remains covered by the module-path test. Signed-off-by: Alp Dener <adener@nvidia.com>
general_grouped_gemm_for_grouped_tensor allocated its setup workspace (the cuBLAS per-group pointer/dimension arrays) and its cuBLAS workspace with per-call torch.empty. Under make_graphed_callables the forward and backward graphs share one capture memory pool, and a per-call allocation's block returns to that pool as soon as the Python reference dies, so blocks alias across the two graphs and captured kernels from one graph overwrite the GEMM metadata the other graph reads at replay. Observed as allocation-history-dependent failures in the ops-layer GroupedLinear cuda-graph test: capture-time cublasLtMatmulAlgoGetHeuristic NOT_SUPPORTED errors and corrupted wgrad outputs. This is also the likely mechanism behind the FP8 block-scaling wgrad corruption under CUDA graphs previously observed on Hopper and attributed to cuBLAS. Cache the setup workspace per (device, group size) and reuse the cached per-device cuBLAS workspace from the non-grouped path; consecutive GEMMs reusing one workspace are ordered by the stream. Signed-off-by: Alp Dener <adener@nvidia.com>
…ole cuBLAS workspaces The grouped-tensor GEMM path shared one persistent cuBLAS workspace across all grouped matmuls. cuBLAS's grouped GEMM keeps a grid-synchronization flag in the first bytes of that workspace and zeros it (via a captured memset) before each matmul. When the dgrad and wgrad grouped matmuls of a GroupedLinear backward share one workspace inside a replayed CUDA graph, that flag is aliased between the two matmuls; on the second graph replay the second matmul's cooperative kernel deadlocks with cuBLAS 13.6 (and corrupts the last expert's wgrad on cuBLAS < 13.6). The two matmuls are strictly stream-ordered (single stream, all-DEFAULT graph edges, no programmatic dependent launch), so this is shared-workspace reuse, not concurrent co-scheduling. Give dgrad/forward (slot 0) and wgrad (slot 1) distinct persistent cuBLAS workspaces, dedicated to the grouped path. Each slot remains a single persistent allocation, so CUDA-graph capture safety is preserved. Also drop the cuBLAS-version gate that skipped the FP8 block-scaling GroupedLinear CUDA-graph test, so it now exercises the fix on all supported cuBLAS versions. Signed-off-by: Alp Dener <adener@nvidia.com>
for more information, see https://pre-commit.ci
…ale dbias comment - general_grouped_gemm_for_grouped_tensor: expand the comment to state that the fused grouped FP8 block-scaling GEMM forces use_split_accumulator=True and intentionally overrides the caller-supplied value, consistent with the Float8BlockScaling recipe (which fixes it True for fprop/dgrad/wgrad). - Float8BlockScaling recipe docstring: document that FP8 block scaling always uses split accumulation and that the fused grouped GEMM path ignores any caller- or recipe-supplied use_split_accumulator value. - GroupedLinear ops backward: correct the stale "BF16/FP16 path" comment; that branch also handles quantized paths where bgrad fusion did not apply (e.g. FP8 block scaling without a dgrad pass). Signed-off-by: Alp Dener <adener@nvidia.com>
…near module Restrict this PR to the GroupedLinear module fused-quantize path. Revert the fusible-ops FP8 block-scaling enablement -- the BasicOperation opt-in gate, the GroupedLinear op support, and the fusible-ops test coverage -- back to main. Enabling fusible-ops FP8 block-scaling for both grouped and non-grouped paths is deferred to a separate PR. The blanket FP8 block-scaling rejection in BasicOperation.reset_recipe_state is restored. The split-accumulator guard in general_grouped_gemm_for_grouped_tensor is retained: it is correct for the module's FP8 block-scaling grouped GEMM. Signed-off-by: Alp Dener <adener@nvidia.com>
…t-discreteness _get_grouped_cublas_workspace slots were keyed on is_discrete_out as a proxy for "this is the wgrad GEMM", which only holds when wgrad writes a list of per-expert grads. With single_grouped_weight=True, wgrad writes a single grouped weight-grad (GroupedTensor out, not a list), so is_discrete_out is False and it collided with dgrad on slot 0 -- reintroducing the FP8 block-scaling grid-sync-flag aliasing deadlock/corruption under CUDA-graph replay. Key the slot on the wgrad layout (NT / transb) instead: fprop (TN) and dgrad (NN) share slot 0, wgrad (NT) is always isolated on slot 1. Signed-off-by: Alp Dener <adener@nvidia.com>
for more information, see https://pre-commit.ci
05ef7a7 to
e365a88
Compare
|
/te-ci pytorch |
Description
This PR integrates the new grouped FP8 block-scaling quantization kernels into the fused GroupedTensor path in GroupedLinear.
Integration into fusible ops is deferred to a follow-up PR as it requires additional effort to also enable non-grouped FP8 block-scaling at the same time, considered out-of-scope for this PR.
Notes
Type of change
Checklist: