Repository navigation
[Common/PyTorch] Enable FP8BS for Grouped GEMM on Blackwell via MXFP8 Broadcast - #3614
Merged
denera merged 8 commits intoOct 19, 2026
Merged
Conversation
denera
requested review from
Oleg-Goncharov,
ksivaman,
ptrendx and
timmoon10
as code owners
October 2, 2026 13:00
denera
force-pushed
the
pytorch/grouped-linear-fp8bs-blackwell
branch
from
October 2, 2026 13:01
0778269 to
55c2ca4
Compare
Collaborator
Author
|
/te-ci pytorch |
Contributor
 [Medium impact] The PR appears safe to merge, with the previous binding failure fixed and no new actionable issues found. SummaryThis PR enables grouped FP8 block scaling on Blackwell by converting its scales to the MXFP8 format.
Review was source-only; no runtime tests were run. Diagram%%{init: {'theme': 'neutral'}}%%
flowchart LR
A["Grouped FP8 block-scaled operands"] --> B{"GPU architecture"}
B -->|Hopper| C["Native grouped FP8 block-scaling GEMM"]
B -->|Blackwell| D["Broadcast and swizzle scales to MXFP8"]
D --> E["MXFP8 grouped GEMM using TN views"]
C --> F["Output"]
E --> F
Reviews (7) · Last reviewed commit: "Added compile-time static check for grou..." · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
denera
force-pushed
the
pytorch/grouped-linear-fp8bs-blackwell
branch
from
October 2, 2026 13:31
55c2ca4 to
1229f6d
Compare
Collaborator
Author
|
/te-ci pytorch |
vthumbe1503
reviewed
Oct 5, 2026
Collaborator
There was a problem hiding this comment.
Thanks @denera! Mostly LGTM with some minor comments on CPU overheads. Could we also add some perf numbers for this kernel in the description?
Comment on lines
+600
to
+602
| NVTE_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, | ||
| static_cast<int>(smem_bytes))); | ||
| } |
Collaborator
There was a problem hiding this comment.
We could cache this as well.
Comment on lines
+594
to
+599
| NVTE_CHECK_CUDA(cudaDeviceGetAttribute( | ||
| &max_smem_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin, cuda::current_device())); | ||
| NVTE_CHECK(smem_bytes <= static_cast<size_t>(max_smem_bytes), "Too many tensors (", | ||
| input->num_tensors, ") to convert FP8 block scaling to MXFP8 scales: needs ", | ||
| smem_bytes, " bytes of shared memory, but the device allows ", max_smem_bytes, | ||
| "."); |
Collaborator
There was a problem hiding this comment.
We can cache this value per device in cuda_runtime.cpp to avoid cpu overheads.
|
|
||
| // Emulate the FP8 block scaling recipe with MXFP8 on Blackwell and newer as it is not | ||
| // natively supported by cublasLt grouped GEMM (mirrors generic_gemm). | ||
| std::vector<at::Tensor> mxfp8_emulation_scales; |
Collaborator
There was a problem hiding this comment.
Is there a need to define this std::vector and emplace the scales in it? It doesnt seem to be used anywhere
| mxfp8_emulation_scales.emplace_back(convert_block_scaling_to_mxfp8_tensor(A_tensor, transa)); | ||
| } | ||
| mxfp8_emulation_scales.emplace_back( | ||
| convert_grouped_block_scaling_to_mxfp8_tensor(grouped_B, !transb)); |
Collaborator
There was a problem hiding this comment.
Isnt grouped_swizzle function the right place for this implementation(I dont exactly remember the API).
| std::vector<at::Tensor> mxfp8_emulation_scales; | ||
| if (needs_fp8_block_scaling_emulation(operand_modes)) { | ||
| for (auto& A_tensor : te_A_wrappers) { | ||
| mxfp8_emulation_scales.emplace_back(convert_block_scaling_to_mxfp8_tensor(A_tensor, transa)); |
Collaborator
There was a problem hiding this comment.
For weights we would be calling multi_tensor_swizzle somewhere. Maybe that is the right place for this implementation?
Comment on lines
+280
to
+282
| __host__ __device__ __forceinline__ size_t round_up(const size_t a, const size_t b) { | ||
| return ceil_div(a, b) * b; | ||
| } |
Collaborator
There was a problem hiding this comment.
Could we reuse DIVUP_TO_MULTIPLE() from common/common.h for the corresponding calculations below?
Comment on lines
+277
to
+279
| __host__ __device__ __forceinline__ size_t ceil_div(const size_t a, const size_t b) { | ||
| return (a + b - 1) / b; | ||
| } |
Collaborator
There was a problem hiding this comment.
Could we reuse DIVUP() from common/common.h for the corresponding calculations below?
denera
force-pushed
the
pytorch/grouped-linear-fp8bs-blackwell
branch
from
October 6, 2026 19:49
1229f6d to
0d8f8d5
Compare
denera
force-pushed
the
pytorch/grouped-linear-fp8bs-blackwell
branch
from
October 8, 2026 13:59
0d8f8d5 to
d8c2ac1
Compare
…ewer Drop the Hopper-only arch gates from the grouped FP8 block-scaling quantize and dequantize kernels. On SM100 and newer the recipe is emulated with MXFP8 GEMM, so grouped quantize now rejects non-power-of-2 scales there, matching the non-grouped quantize kernels. Signed-off-by: Alp Dener <adener@nvidia.com>
…e conversion Add nvte_swizzle_grouped_block_scaling_to_mxfp8_scaling_factors, the grouped variant of nvte_swizzle_block_scaling_to_mxfp8_scaling_factors used to emulate FP8 block scaling with MXFP8 grouped GEMM. A persistent kernel reads per-tensor dims on the device, so captured launches stay correct across CUDA graph replays, and tabulates each tensor's first tile and scale offsets in shared memory. It then walks 128x128 tiles across all tensors, finds each tile's tensor by binary search, and writes the swizzled E8M0 scales at the offsets the grouped GEMM expects. nvte_get_grouped_block_scaling_to_mxfp8_scale_inv_size returns the output buffer size, which depends only on the logical shape. Add nvte_multi_tensor_swizzle_block_scaling_to_mxfp8_scaling_factors, which converts a list of tensors, such as per-expert weights, in one launch per up to 128 tensors instead of one launch per tensor. All conversion kernels share the per-tile code. All three APIs are covered by a host-reference unit test. Signed-off-by: Alp Dener <adener@nvidia.com>
cuBLASLt grouped GEMM has no native FP8 block scaling on Blackwell and newer. As in non-grouped GEMM, the grouped-tensor GEMM bindings now convert FP8 block-scaling operands to MXFP8 by broadcasting their scales, viewing transposed columnwise data as rowwise data of the transposed tensors, and force the TN layout. The grouped GEMM itself only ever sees an MXFP8 GEMM. Lists of discrete operands, such as per-expert weights, are converted in one multi-tensor launch. The multi-stream grouped GEMM now does the same instead of converting one tensor at a time. Signed-off-by: Alp Dener <adener@nvidia.com>
…tests The exact operands used by the grouped FP8 block-scaling GEMM tests gave nearly every 128x128 block the same power-of-2 scale, so a wrong scale offset, orientation or broadcast in the MXFP8 emulation would not change the numerics. Multiply each operand tile by its own random power-of-2 factor, which every quantization block absorbs exactly into its scale, so quantization stays lossless while per-block scales now differ. Signed-off-by: Alp Dener <adener@nvidia.com>
Admit FP8 block scaling on the GroupedLinear module and fusible-op grouped-tensor paths wherever MXFP8 is supported, now that the grouped GEMM bindings emulate it with MXFP8 on Blackwell and newer, and drop the module's Blackwell error. Reject non-power-of-2 FP8 block-scaling recipes on Blackwell and newer when entering autocast, since MXFP8 emulation needs power-of-2 scales; previously non-grouped GEMM only failed later inside quantization. Signed-off-by: Alp Dener <adener@nvidia.com>
Capture grouped quantize and grouped-tensor GEMM in one CUDA graph for MXFP8 and FP8 block scaling (TN, NN and NT layouts), then replay it after updating the device-side split sizes in place, including zero-token experts and fewer tokens than the buffer capacity. Each run is compared against an exact high-precision reference over each expert's real region. On Blackwell and newer this covers the FP8 block-scaling MXFP8 emulation under graph replay. Also tighten the mixed-operand rejection test to match the binding's error message. Signed-off-by: Alp Dener <adener@nvidia.com>
…S scale conversion Document in both GroupedLinear docstrings (module and fusible op) and the module's grouped-tensor path predicate that, with MXFP8 or FP8 block scaling, the grouped-tensor path requires every split size to be a multiple of 128, which is checked only on the device. For the fusible op, this is the alignment that split_alignment declares. With FP8 block scaling now on this path on Blackwell, fusible-op users no longer fall back to split quantization there. Check that the per-tensor dims passed to nvte_swizzle_grouped_block_scaling_to_mxfp8_scaling_factors are int64, as the kernel reads them. Say to unset NVTE_FP8_BLOCK_SCALING_FP32_SCALES before importing Transformer Engine in the Blackwell non-power-of-2 rejection message, because the recipe defaults read it at import time. Signed-off-by: Alp Dener <adener@nvidia.com>
denera
force-pushed
the
pytorch/grouped-linear-fp8bs-blackwell
branch
from
October 15, 2026 19:02
d8c2ac1 to
667b675
Compare
Collaborator
Author
|
/te-ci pytorch |
0z5a
reviewed
Oct 19, 2026
0z5a
left a comment
0z5a
left a comment
Contributor
There was a problem hiding this comment.
I tested 667b6751 on B200 (sm_100), built with NVTE_CUDA_ARCHS=100a and NVTE_WITH_NCCL_EP=0.
One setup note first. With the cuBLAS that ships with the CUDA 13.0 toolkit (cuBLASLt 13.1), none of this PR's grouped-tensor path runs: is_module_grouped_tensor_path_supported returns False, and 936 cases in the subset below are skipped with "Grouped GEMM requires cuBLAS 13.3+". To exercise the PR I rebuilt against cuBLAS 13.8.1 headers and libraries (runtime tex.get_cublasLt_version() = 130801).
Results with that build:
| Tests | Result |
|---|---|
test_grouped_linear.py -k "cuda_graph or grouped_tensor or single_grouped or grouped_parameter or fused_path" |
872 passed, 77 skipped, 0 failed |
FP8 block-scaling CUDA-graph cases: test_grouped_gemm_grouped_tensor_cuda_graph_changed_splits[fp8_block_scaling-{TN,NN,NT}], test_grouped_linear_fused_path_cuda_graph_safe[{True,False}-fp8_block_scaling] |
5/5 passed |
test_grouped_linear.py -k "Float8BlockScaling or fp8_block_scaling" |
610 passed, 72 skipped |
test_grouped_tensor.py + test_recipe.py |
393 passed, 22 skipped |
The skips are the declared ones: "NT (wgrad) has no weight operand" (72), "fp8_model_params requires FP8" (4), and NVFP4 single_grouped_weight (1). I did not run the C++ test_swizzle_block_scaling. Hopper and sm_103 were not covered.
One observation outside this PR's diff:
is_module_grouped_tensor_path_supportedgates on the runtime cuBLASLt version, but the grouped-GEMM entry points are compiled out when TE is built against cuBLAS < 13.3.- With an older build running against a newer cuBLASLt (for example after upgrading
nvidia-cublas), the predicate says yes andGroupedLinear(use_grouped_tensor=True)raisesnvte_get_grouped_gemm_setup_workspace_size requires cuBLAS 13.3+, but compile-time cuBLAS version is ...instead of falling back. - That is how 856 cases of the same subset failed before I rebuilt. Checking a compile-time capability in the predicate would turn this into a clean fallback.
Collaborator
Author
This is a very fair point and we should definitely guard this against the compile-time static cuBLAS version. @vthumbe1503 I added a fix for this but it invalidated the earlier review. Could you take another look and approve again so we can merge? |
…NOT built with grouped GEMM support even if cuBLAS runtime version supports it. Signed-off-by: Alp Dener <adener@nvidia.com>
denera
force-pushed
the
pytorch/grouped-linear-fp8bs-blackwell
branch
from
October 19, 2026 22:29
eb14ae6 to
de53764
Compare
vthumbe1503
approved these changes
Oct 19, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Extends
Float8BlockScalingon the grouped-tensor path from Hopper-only to Blackwell and newer. This covers grouped quantization, cuBLASLt grouped GEMM, andte.pytorch.GroupedLinear(module and fusible op). The new path applies wherever MXFP8 grouped GEMM is supported (SM100/SM103/SM107/SM110) by broadcasting FP32 power-of-2 scales into swizzled E8M0 MXFP8 scales.Fixes #3520
Performance
Broadcast + Swizzle
Grouped tensors w/
nvte_swizzle_grouped_block_scaling_to_mxfp8_scaling_factorsList of Tensors w/
nvte_multi_tensor_swizzle_block_scaling_to_mxfp8_scaling_factorsConversion Cost per GroupedGemm in an MoE Layer
Based on 16 local experts × 2560 tokens, with 4096 hidden size and 1024 expert FFN.
GroupedLinear
Based on 4096 hidden size and 1024 FC2 FFN.
End-to-End Timing (ms)
Breakdown per FWD+BWD (µs)
Notes
FP8BS GroupedLinear slowdown vs. MXFP8 is due to the quantization, not the conversion.
FP8BS forces TN layout in both non-grouped and grouped cases:
Optimizing FP8BS quantization kernels (both grouped and non-grouped) on Blackwell deferred to another PR.
Type of change
Checklist: