[PyTorch] Enable fused FP8 block-scaling path in GroupedLinear module by denera · Pull Request #3171 · NVIDIA/TransformerEngine

denera · 2026-07-02T20:35:20Z

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

DGRAD and WGRAD grouped GEMMs executing serially on the same stream appear to violate the workspace serialization under some specific graph replay conditions, leading to a deadlock in WGRAD. PR resolves the deadlock by allocating and caching TWO 32MiB cuBLAS workspaces in order to isolate DGRAD and WGRAD workspaces.
PR enables pow_2 constrained scales for grouped FP8 block-scaling at the kernel level for numerical parity with the split-quantize path on SM90, but the FP8BS->MXFP8 scale-broadcast mechanism remains disabled on SM100+.

Type of change

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

Checklist:

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

denera · 2026-07-02T20:42:44Z

/te-ci core pytorch

greptile-apps · 2026-07-02T20:45:54Z

Greptile Summary

This PR integrates the grouped FP8 block-scaling quantization kernels into the fused GroupedTensor path in GroupedLinear, including support for power-of-2 constrained scales (force_pow_2_scales) and a deadlock fix for CUDA-graph replay.

pow_2_scales threading: The force_pow_2_scales parameter is removed from the "unsupported" gate and instead forwarded all the way down through group_quantize_blockwise_{1d,2d} to the individual CUDA kernels, enabling numerical parity with the split-quantize path on SM90.
cuBLAS workspace deadlock fix: Replaces per-call torch.empty(...) allocations with lru_cache-backed persistent tensors per (device, layout) triplet, preventing PyTorch's caching allocator from aliasing the DGRAD (NN) and WGRAD (NT) cuBLAS workspaces within a CUDA graph, which caused a stale-descriptor deadlock on the 2nd graph replay.
Blackwell hard error: On SM100/SM110 with NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1, the new Float8BlockQuantizer branch raises RuntimeError rather than silently falling back; SM120+ still returns False from the fused-path gate and falls back cleanly.

Confidence Score: 5/5

The change is safe to merge. Core logic — pow_2_scales threading, cuBLAS workspace isolation, and the Blackwell hard-error guard — is correct, well-tested, and has no observable correctness regressions on the supported Hopper path.

The pow_2_scales parameter is correctly forwarded through every kernel call site and its previous guard-rejection has been cleanly removed. The cuBLAS workspace deadlock fix correctly identifies the root cause (caching allocator aliasing within a CUDA graph) and isolates each GEMM layout to its own persistent workspace. The fuse_bgrad condition for FP8 block scaling handles both the rowwise-present and rowwise-absent cases correctly and consistently with the MXFP8 precedent. Tests cover the new SM90 execution path and the Blackwell error path. The two observations flagged are design/maintenance notes rather than correctness issues.

transformer_engine/pytorch/cpp_extensions/gemm.py — the setup workspace cache key does not include layout, so concurrent multi-stream grouped GEMMs of different layouts on the same device would share one buffer (safe for the current single-stream backward, but worth revisiting if TE adds multi-stream overlap).

Important Files Changed

Filename	Overview
transformer_engine/pytorch/cpp_extensions/gemm.py	Adds lru_cache-backed persistent workspace tensors per (device, layout); introduces _is_fp8_blockwise guard to force use_split_accumulator=True for FP8 block-scaling GEMMs; setup workspace cached per (device, num_tensors) shared across layouts.
transformer_engine/pytorch/module/grouped_linear.py	Adds Float8BlockQuantizer branch to _use_grouped_gemm_path (SM90 returns True, SM100+ raises); threads fuse_bgrad condition for FP8BlockScaling requiring rowwise output; backward logic is correct.
transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh	Adds pow_2_scales parameter to all four grouped kernels (1D colwise, 1D TMA, 2D TMA colwise/rowwise) and forwards it to compute_scale_from_types; inline dispatch functions updated accordingly.
transformer_engine/common/cast/dispatch/quantize.cuh	Removes the NVTE_CHECK guards that blocked force_pow_2_scales=True in grouped FP8 block-scaling; now forwards the flag to group_quantize_blockwise_{1d,2d} for both forward and backward paths.
transformer_engine/common/include/transformer_engine/cast.h	nvte_group_quantize_dbias signature gains a new NVTEQuantizationConfig parameter; all in-repo callers updated to pass nullptr/quant_config_cpp, but this is a source-breaking change for any out-of-tree users of the C API.
transformer_engine/pytorch/csrc/extensions/cast.cpp	bgrad_group_quantize now reads force_pow_2_scales and amax_epsilon off Float8BlockQuantizer and passes them as a QuantizationConfigWrapper to nvte_group_quantize_dbias.
transformer_engine/pytorch/csrc/quantizer.cpp	Removes the NVTE_CHECK guard that rejected force_pow_2_scales=True in Float8BlockQuantizer::create_grouped_tensor since the kernel now handles it.
tests/pytorch/test_grouped_linear.py	Adds Float8BlockScaling to two existing test parametrizations with correct SM90-only skip guards, and adds a new test for the Blackwell hard-error path.
transformer_engine/common/cast/cast_grouped_dbias.cu	Forwards the new quant_config parameter from the nvte_group_quantize_dbias public API to group_quantize_bwd_helper; previously hardcoded nullptr.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    GL["GroupedLinear backward"]
    CHECK_FP8{"ctx.fp8?"}
    GET_Q["grad_output_quantizer\nset_usage(rowwise=requires_dgrad,\ncolumnwise=weights_requires_grad)"]
    FUSE_BGRAD{"fuse_bgrad?\n(MXFP8 OR\nFP8Block+requires_dgrad)"}
    BGRAD_Q["tex.bgrad_group_quantize\n→ grouped_dy + dbias_packed"]
    Q_ONLY["tex.group_quantize\n→ grouped_dy (col-only)"]
    BF16_PATH["_make_grouped_tensor\n(BF16/FP16)"]
    DBIAS_FALLBACK["compute_grouped_dbias\n(float dy)"]
    DGRAD_GEMM["DGRAD GEMM (NN)\n_get_grouped_cublas_workspace(dev, 'NN')"]
    WGRAD_GEMM["WGRAD GEMM (NT)\n_get_grouped_cublas_workspace(dev, 'NT')"]
    SETUP_WS["_get_grouped_gemm_setup_workspace\n(dev, num_tensors) shared across layouts"]

    GL --> CHECK_FP8
    CHECK_FP8 -->|yes| GET_Q
    CHECK_FP8 -->|no| BF16_PATH
    GET_Q --> FUSE_BGRAD
    FUSE_BGRAD -->|yes| BGRAD_Q
    FUSE_BGRAD -->|no| Q_ONLY
    BF16_PATH --> DGRAD_GEMM
    BGRAD_Q --> DGRAD_GEMM
    Q_ONLY -->|use_bias=True| DBIAS_FALLBACK
    Q_ONLY --> DGRAD_GEMM
    DGRAD_GEMM --> SETUP_WS
    DGRAD_GEMM --> WGRAD_GEMM
    WGRAD_GEMM --> SETUP_WS

_{Reviews (8): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

timmoon10 · 2026-07-06T20:43:34Z

+        * 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


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.

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.

My complaint is more about commment style, not functionality. Power-of-2 is a minor advanced config, and it should be documented where it is relevant. However, it seems the agent struggled with it at some point and now it brings it up as a top-level consideration everywhere even if it would be distracting to future users.

There's only one place I see any mention of the power-of-2 in comments: https://github.com/NVIDIA/TransformerEngine/pull/3171/changes#diff-e52c6ddc8c0f4d20cb5aa832e92dd3794687bac1f828e367b526f17f524238a5R112-R113

I think there might have been more before because the PR originally did not support power-of-2 scales, and there had to be a lot of exceptions/branching to reconcile that with the fact that non-grouped/unfused FP8BS does (in fact it's the default, I believe).

Those comments/branching are removed now that grouped FP8BS has parity with non-grouped.

vthumbe1503 · 2026-07-07T19:02:21Z

+    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)


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

denera · 2026-07-09T04:29:36Z

/te-ci pytorch

vthumbe1503 · 2026-07-17T17:27:28Z

+    # out-discreteness. A wgrad into a single grouped weight-grad (single_grouped_weight)
+    # has a GroupedTensor out, not a list, so an is_discrete_out proxy would collide it
+    # with dgrad on slot 0. fprop (TN) and dgrad (NN) can share slot 0; only NT is isolated.
+    workspace_cublas = _get_grouped_cublas_workspace(device.index, 1 if transb else 0)


Lets just include layout(NN, TN, NT) as part of the function signature, so that we get different cached workspace for each layout. Although we face this problem only for wgrad now in TE. There are also some other reports of even fwd and dgrad conflicting with each other.

vthumbe1503 · 2026-07-17T17:40:00Z

+    # Spy on the grouped quantize entry point so a predicate that silently
+    # declines the fused path fails the test instead of vacuously matching the
+    # legacy reference.
+    group_quantize_calls = 0


This seems too much to be honest. We already have cuda graph test that verifies the graph safeness. If it is graph safe, it has to go through the group quantize route. So I dont think, counting the group_quantize calls and verifying is necessary

denera · 2026-07-21T21:42:54Z

/te-ci pytorch

timmoon10

This is functionally ready, although I would update the comments so future developers can understand the cuBLAS GGEMM bug with stale TMA descriptors. As it is, the comments make it seem the bug is in the memory allocator, which is wrong.

timmoon10 · 2026-07-22T17:31:18Z

+    This must not be allocated per call: under CUDA-graph capture a per-call
+    allocation's block returns to the shared capture pool as soon as the Python
+    reference dies, so the forward and backward graphs can alias the same block
+    and the captured GEMM's pointer/dimension arrays get overwritten at replay.
+    Consecutive GEMMs reusing one workspace are ordered by the stream, matching
+    how the non-grouped path shares its cached cuBLAS workspace.


Functionally I'm ok with caching the GGEMM workspace (it's a bit wasteful, but I don't expect the workspace to be that large anyways). This explanation is misleading though. It claims there's a bug when the forward and backward graphs alias the same workspace buffer, and the fix is to force aliasing the same buffer? The real bug is not in the workspace or the memory allocator, but in the cuBLAS kernel TMA descriptors. Caching the workspace is a hacky workaround until the cuBLAS fix is available.

Yeah, this comment is no longer relevant. It's left over from debugging when it wasn't yet clear whether it was the setup workspace or cuBLAS workspace causing the issue, but since then we've narrowed it down to stale descriptors in the cuBLAS workspace and the setup workspace does not need to be cached per GEMM layout (just number of tensors). I removed this comment entirely now.

timmoon10 · 2026-07-22T17:38:31Z

+    Two grouped-tensor GEMMs that share a single cuBLAS workspace can deadlock on the
+    second CUDA-graph replay: the grouped kernels interact through the shared workspace
+    and the second matmul hangs. It is deterministic per graph geometry and reproduces
+    on cuBLAS 13.5/13.6/13.7 (i.e. not fixed upstream as of 13.7); giving each GEMM its
+    own persistent workspace avoids it. The backward wgrad (NT) is the case seen in TE,
+    but fprop (TN) and dgrad (NN) have also been reported to conflict, so we isolate by
+    layout: TN, NN, and NT each get a distinct persistent workspace. Each is a single
+    persistent allocation, so CUDA-graph capture safety is preserved.


Now that we've gotten to the bottom of the bug (stale TMA descriptors), we should include that rather than this exploratory bug report. We should also note that the bug has been fixed in cuBLAS and that this hacky fix can be reverted in the future.

Yes, this was leftover from before cuBLAS devs diagnosed the issue. I updated the comment to reflect what the actual cause is.

timmoon10 · 2026-07-22T17:47:18Z

+        * 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


My complaint is more about commment style, not functionality. Power-of-2 is a minor advanced config, and it should be documented where it is relevant. However, it seems the agent struggled with it at some point and now it brings it up as a top-level consideration everywhere even if it would be distracting to future users.

…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

…; drop redundant test spy - _get_grouped_cublas_workspace now keys the persistent workspace on the grouped GEMM layout, so fprop (TN), dgrad (NN), and wgrad (NT) each get a distinct workspace. The previous NT-vs-rest scheme left fprop and dgrad sharing one workspace; those have also been reported to conflict under CUDA-graph replay. Documents that the deadlock is deterministic and present through cuBLAS 13.7. - Drop the group_quantize call-counting spy in test_grouped_linear_grouped_tensor_path_matches_legacy; fused-path engagement is covered by the graph-safe test. Signed-off-by: Alp Dener <adener@nvidia.com>

…d deadlocks Signed-off-by: Alp Dener <adener@nvidia.com>

for more information, see https://pre-commit.ci

…#3171) * [Common/PyTorch] Support power-of-2 scales in grouped FP8 block-scaling 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> * [PyTorch] Enable fused grouped FP8 block-scaling path in GroupedLinear 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> * [PyTorch] Enable FP8 block-scaling in GroupedLinear fusible op 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> * [PyTorch] Use persistent workspaces in grouped-tensor GEMM 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> * [PyTorch] Fix grouped FP8 block-scaling CUDA-graph deadlock via per-role 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> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Address review: document split-accumulator override, fix stale 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> * [PyTorch] Revert fusible-ops FP8 block-scaling; scope PR to GroupedLinear 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> * [PyTorch] Isolate grouped wgrad cuBLAS workspace by NT layout, not out-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> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Address review: isolate grouped cuBLAS workspace per layout; drop redundant test spy - _get_grouped_cublas_workspace now keys the persistent workspace on the grouped GEMM layout, so fprop (TN), dgrad (NN), and wgrad (NT) each get a distinct workspace. The previous NT-vs-rest scheme left fprop and dgrad sharing one workspace; those have also been reported to conflict under CUDA-graph replay. Documents that the deadlock is deterministic and present through cuBLAS 13.7. - Drop the group_quantize call-counting spy in test_grouped_linear_grouped_tensor_path_matches_legacy; fused-path engagement is covered by the graph-safe test. Signed-off-by: Alp Dener <adener@nvidia.com> * updated grouped GEMM workspace comment on stale TMA descriptor related deadlocks Signed-off-by: Alp Dener <adener@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Alp Dener <adener@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

denera requested review from ptrendx and vthumbe1503 July 2, 2026 20:35

denera self-assigned this Jul 2, 2026

denera requested review from Oleg-Goncharov, ksivaman and timmoon10 as code owners July 2, 2026 20:35

denera added FP8 MoE 2.18 labels Jul 2, 2026

denera force-pushed the pytorch/grouped-linear-fused-fp8bs branch from bbfa1da to b35924c Compare July 2, 2026 20:35

denera linked an issue Jul 2, 2026 that may be closed by this pull request

Blockwise (1x128 and 128x128) FP8 grouped quantization #2525

Closed

greptile-apps Bot reviewed Jul 2, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/cpp_extensions/gemm.py

timmoon10 requested changes Jul 6, 2026

View reviewed changes

vthumbe1503 reviewed Jul 7, 2026

View reviewed changes

denera changed the title ~~[PyTorch] Enable fused FP8 block-scaling path in GroupedLinear module and fusible ops~~ [PyTorch] Enable fused FP8 block-scaling path in GroupedLinear module Jul 7, 2026

denera requested review from timmoon10 and vthumbe1503 July 7, 2026 22:39

denera force-pushed the pytorch/grouped-linear-fused-fp8bs branch 3 times, most recently from 05ef7a7 to e365a88 Compare July 9, 2026 04:29

vthumbe1503 reviewed Jul 17, 2026

View reviewed changes

denera force-pushed the pytorch/grouped-linear-fused-fp8bs branch 2 times, most recently from 8745513 to 5dd2248 Compare July 21, 2026 21:29

vthumbe1503 previously approved these changes Jul 21, 2026

View reviewed changes

timmoon10 previously approved these changes Jul 22, 2026

View reviewed changes

denera dismissed stale reviews from timmoon10 and vthumbe1503 via 196bb76 July 22, 2026 20:13

denera and others added 13 commits July 22, 2026 20:22

[pre-commit.ci] auto fixes from pre-commit.com hooks

cd8d42b

for more information, see https://pre-commit.ci

[pre-commit.ci] auto fixes from pre-commit.com hooks

85baa24

for more information, see https://pre-commit.ci

updated grouped GEMM workspace comment on stale TMA descriptor relate…

3dc9a8f

…d deadlocks Signed-off-by: Alp Dener <adener@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

1393778

for more information, see https://pre-commit.ci

denera force-pushed the pytorch/grouped-linear-fused-fp8bs branch from 906654b to 1393778 Compare July 22, 2026 20:23

denera requested review from timmoon10 and vthumbe1503 July 22, 2026 20:30

vthumbe1503 approved these changes Jul 22, 2026

View reviewed changes

denera merged commit 098b496 into NVIDIA:main Jul 22, 2026
10 of 14 checks passed

		* 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

Uh oh!

Conversation

denera commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Notes

Type of change

Checklist:

Uh oh!

denera commented Jul 2, 2026

Uh oh!

greptile-apps Bot commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Flowchart

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

denera commented Jul 9, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

denera commented Jul 21, 2026

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

denera commented Jul 2, 2026 •

edited

Loading

greptile-apps Bot commented Jul 2, 2026 •

edited

Loading