Improve the dimension checks for the FP8 recipes#2894
Improve the dimension checks for the FP8 recipes#2894ptrendx wants to merge 19 commits intoNVIDIA:mainfrom
Conversation
This one-size-fits-all Python guard checked dimensions before the recipe was known, rejecting valid shapes. Dimension validation is handled per-recipe in the C++ quantizer where requirements are known. Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
- MXFP8: only require last dim divisible by 16 (was: both dims by 32). The kernel handles partial blocks via bounds checks and TMA zero-padding. Scale tensors are over-allocated with roundup alignment. - NVFP4: only require last dim divisible by 32 (was: both dims by 16). 4-bit data needs 32 elements for 16-byte alignment. - Float8BlockScaling swizzle: remove data_rows%4 assertion. The kernel already handles non-aligned rows via DIVUP and OOB zero-fill. - Fix integer truncation in MXFP8/NVFP4 get_scale_shape to use ceildiv instead of plain division, ensuring correct scale allocation for non-block-aligned dimensions. Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
- MXFP8: only check last dim divisible by 16 (removed first-dim check) - NVFP4: only check last dim divisible by 32 (removed first-dim check) Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
NVFP4 backward pass uses the Hadamard transform which requires num_rows % 16 == 0. Add a clear error message at the quantizer level when columnwise_usage is enabled, instead of letting the user hit the raw Hadamard kernel assertion. Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
The old name was misleading — it doesn't check whether a tensor can be quantized in general, but whether a local shard's shape supports quantized all-gather without scaling factor blocks spanning across GPU boundaries. The new name reflects the actual purpose. Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Include actual dimension values (m, n, k, lda, ldb) in error messages so users can trace which tensor dimension is misaligned. Reference the cuBLAS documentation for FP8 alignment requirements. Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Remove "Ensure all tensor dimensions..." sentence and parentheses from cuBLAS documentation links for consistent error messages. Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Replace hardcoded alignment values (16, 32) with a helper that computes the minimum element alignment from the data type's bit width and the 16-byte cuBLAS alignment requirement. This ensures alignment checks stay correct if new data types are added. Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Verifies that the relaxed dimension checks allow small M values for recipes that support them on the forward pass. Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Each recipe gets its tightest viable (M, N, K) for both inference and training. The baseline pre-quantizes+dequantizes inputs and weights via the recipe's quantizers so the comparison holds at BF16 tolerance. Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
for more information, see https://pre-commit.ci
Greptile SummaryThis PR refines FP8/FP4 dimension checks to be recipe-aware: removes the over-restrictive Python-level
Confidence Score: 4/5Safe to merge with minor concern about missing alignment padding in the new_zeros dispatch for newly-relaxed MXFP8 shapes. The core dimension-check relaxation is correct and well-reasoned. Previously flagged P0/P1 issues (flat_first_dim NameError, ceildiv inconsistency, fsdp_pre_all_gather floor-div) are resolved. One P2 concern remains: the new_zeros dispatch uses unpadded scale shapes inconsistent with make_empty, which could produce undersized scale buffers for shapes newly enabled by this PR (last_dim = 16). transformer_engine/pytorch/tensor/mxfp8_tensor.py — new_zeros dispatch scale shape computation (lines 576–583) Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[Tensor shape] --> B{last_dim % 16 == 0?}
B -- No --> C[Reject / fallback to BF16]
B -- Yes --> D{Recipe?}
D -- MXFP8 --> E[C++: last_dim % 16 check\nPython make_empty: last_dim % 16\nScale via ceildiv+roundup]
D -- NVFP4 --> F[C++: last_dim % 2, first_dim % 16 if RHT\nPython make_empty: last_dim % 32\nScale via ceildiv+roundup]
D -- Float8 1x128 --> G[swizzle_block_scaling: rows % 4\ncublaslt: lda % 16, m % 8]
E --> H{Distributed?}
F --> H
H -- MXFP8 AllGather --> I[supports_quantized_allgather:\nlast_dim % 16\nfirst_dim % 32 if columnwise]
H -- NVFP4 AllGather --> J[supports_quantized_allgather:\nlast_dim % 32\nfirst_dim % 16 if columnwise]
I -- passes --> K[Quantized AllGather]
J -- passes --> K
I -- fails --> L[Fallback: BF16 AllGather + re-quantize]
J -- fails --> L
E --> M[cuBLAS GEMM:\nlda % 16 for FP8]
F --> N[cuBLAS GEMM:\nlda % 32 for FP4]
Reviews (5): Last reviewed commit: "Scope quantizer.cpp dim checks to quanti..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| Constraints: | ||
| - cuBLAS needs 16B-aligned leading dims: 16 elts for FP8, 32 for FP4. | ||
| - Linear(K, N) fprop has lda=ldb=K, ldc=N (no M alignment needed). | ||
| - Training adds wgrad/dgrad with lda/ldb/ldc spanning M, N, K -> all must be aligned. | ||
| - Float8BlockScaling swizzle requires data_rows (first dim) % 4. | ||
| - NVFP4 RHT requires input first dim % 16 (subsumed by 32-elt alignment for training). | ||
| """ | ||
| if recipe_obj.delayed() or recipe_obj.float8_current_scaling(): |
There was a problem hiding this comment.
K = 16
The test uses K = 32 for MXFP8 inference, so last_dim = 32 is divisible by 32 and the Python //-based scale-shape bug is never triggered. Adding a separate assertion or updating the inference tuple to (1, 16, 16, ...) would catch the ceildiv inconsistency before it reaches production.
The C++ quantizer uses DIVUP for partial trailing blocks, so a last_dim of 16 (now allowed by the relaxed C++ guard) produces one scale. The Python helpers were still floor-dividing, so the corresponding scale tensors collapsed to zero size. Mirrors the C++ behavior across make_empty, get_scale_shape, and the new_zeros torch_dispatch path; the new_zeros path also relaxes its fall-back guard so shapes the quantizer now accepts aren't silently downgraded to a generic tensor. Adds a Python-level test covering last_dim=16 since the end-to-end GEMM test can't exercise it (the MXFP8 GEMM kernel still requires K >= 32). Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
for more information, see https://pre-commit.ci
- nvfp4_tensor.make_empty: flat_first_dim was never defined; use math.prod(shape[:-1]) (which is what the reader would expect from the surrounding code). - quantized_tensor.supports_quantized_allgather: move the unused-argument pragma to disable-next so pylint actually picks it up. Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
…ble by 32. Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
for more information, see https://pre-commit.ci
|
/te-ci |
| def supports_quantized_allgather(self, inp: torch.Tensor) -> bool: | ||
| """Whether tensor shape supports quantized all-gather. | ||
| When False, the distributed all-gather falls back to gathering | ||
| in high precision and quantizing afterward. This is needed when | ||
| the local shard's shape would cause scaling factor blocks to | ||
| span across GPU boundaries. | ||
| """ | ||
| return True |
There was a problem hiding this comment.
- We shouldn't assume that quantized tensors use a block-scale format.
- While we're touching this function, changing the default to
Falseis more general. It will always work, even without a custom all-gather impl.
| def supports_quantized_allgather(self, inp: torch.Tensor) -> bool: | |
| """Whether tensor shape supports quantized all-gather. | |
| When False, the distributed all-gather falls back to gathering | |
| in high precision and quantizing afterward. This is needed when | |
| the local shard's shape would cause scaling factor blocks to | |
| span across GPU boundaries. | |
| """ | |
| return True | |
| def supports_quantized_allgather(self, inp: torch.Tensor) -> bool: | |
| """Whether tensor shape supports quantized all-gather. | |
| When False, the distributed all-gather falls back to gathering | |
| in high precision and quantizing afterward. | |
| """ | |
| return False |
| flat_first_dim = math.prod(shape[:-1]) | ||
| assert flat_first_dim % NVFP4_BLOCK_SCALING_SIZE == 0, ( | ||
| f"Incorrect shape {shape} for NVFP4. Tensor dims must be divisible by" | ||
| f" {NVFP4_BLOCK_SCALING_SIZE}" | ||
| ) |
There was a problem hiding this comment.
It would be safer to keep this check if allocating column-wise data.
There was a problem hiding this comment.
Actually, not sure if this check is needed at all (even the 32 elements on the rowwise side) that was left there. This check comes from the GEMM and so it is checked there. There is also the check in the allgather support function. Here we should just check if the quantization functions support this shape - ultimately we don't know what is the goal of the person invoking this function, maybe they just need this as storage for their custom operations and so they don't care about the GEMM requirements?
| if inp.shape[-1] % 16 != 0: | ||
| return False |
There was a problem hiding this comment.
Wouldn't partial blocks break all-gather? If we wanted to be as permissive as possible, we could reintroduce the check if row-wise usage is requested:
| if inp.shape[-1] % 16 != 0: | |
| return False | |
| if inp.shape[-1] % 16 != 0: | |
| return False | |
| if self.rowwise_usage and inp.shape[-1] % MXFP8_BLOCK_SCALING_SIZE != 0: | |
| return False |
There was a problem hiding this comment.
I don't think so - in the end the MXFP8 tensors are always rowwise in memory so the last dimension does not matter as it is not the dimension that is being allgathered. The dimension that is checked is the first dimension when the columnwise usage is true.
| scale_inv = torch.empty( | ||
| round_up_to_nearest_multiple(math.prod(shape[:-1]), 128), | ||
| round_up_to_nearest_multiple(shape[-1] // MXFP8_BLOCK_SCALING_SIZE, 4), | ||
| round_up_to_nearest_multiple(math.ceil(shape[-1] / MXFP8_BLOCK_SCALING_SIZE), 4), |
There was a problem hiding this comment.
Nit: I'd prefer going through a ceildiv utility function. It's not like FP64's 52 mantissa bits aren't enough, but I have residual superstition about floating point error.
| // Minimum number of elements for 16-byte alignment, given a data type. | ||
| // cuBLAS requires (dim * typeSize) % 16 == 0 for FP8 tensor core usage, | ||
| // i.e. dim % (128 / typeBits) == 0. | ||
| constexpr size_t kAlignmentBytes = 16; | ||
|
|
||
| size_t min_alignment_elements(transformer_engine::DType dtype) { | ||
| return kAlignmentBytes * 8 / transformer_engine::typeToNumBits(dtype); | ||
| } |
There was a problem hiding this comment.
The comment is misplaced. Also, the min_ is kind of redundant.
| // Minimum number of elements for 16-byte alignment, given a data type. | |
| // cuBLAS requires (dim * typeSize) % 16 == 0 for FP8 tensor core usage, | |
| // i.e. dim % (128 / typeBits) == 0. | |
| constexpr size_t kAlignmentBytes = 16; | |
| size_t min_alignment_elements(transformer_engine::DType dtype) { | |
| return kAlignmentBytes * 8 / transformer_engine::typeToNumBits(dtype); | |
| } | |
| constexpr size_t kAlignmentBytes = 16; | |
| // Number of elements for 16-byte alignment, given a data type. | |
| // cuBLAS requires (dim * typeSize) % 16 == 0 for FP8 tensor core usage, | |
| // i.e. dim % (128 / typeBits) == 0. | |
| size_t alignment_elements(transformer_engine::DType dtype) { | |
| return kAlignmentBytes * 8 / transformer_engine::typeToNumBits(dtype); | |
| } |
| def test_linear_tight_dims(recipe, inference, dtype): | ||
| """te.Linear with the tightest M/N/K per recipe, vs a pytorch baseline. | ||
|
|
||
| Previously the Python assert_dim_for_fp8_exec rejected any M not divisible |
There was a problem hiding this comment.
These historical comments are good for checking what Claude did, but they're useless and distracting if they end up in the final code. assert_dim_for_fp8_exec is gone and mentioning it doesn't help you understand the code as it is.
There was a problem hiding this comment.
Missed that during reviewing of the generated code. Will clean.
| torch.backends.cudnn.allow_tf32 = False | ||
| try: | ||
| te_linear = Linear(K, N, bias=False, params_dtype=dtype, device=device) | ||
| torch_linear = torch.nn.Linear(K, N, bias=False, device=device, dtype=dtype) |
There was a problem hiding this comment.
The real baseline isn't PyTorch, it's mathematical ground truth. FP64 compute helps us get as close as possible:
| torch_linear = torch.nn.Linear(K, N, bias=False, device=device, dtype=dtype) | |
| torch_linear = torch.nn.Linear(K, N, bias=False, device="cpu", dtype=torch.float64) |
This also lets us get rid of the TF32 logic.
There was a problem hiding this comment.
I know you like that approach, but I generally disagree - while you are correct that the FP64 would give us the best estimate of the real values that "should" be there, it doesn't necessarily make the best reference. This is because in the case of the quantized execution we know that we deviate from the ground truth (due to various factors: quantization, TF32, etc.) and using the reference that is correct but closer to the actual computation being done helps with tightening the tolerances. And those ultimately provide better confidence in the result being correct.
| # Share weights: TE gets the raw weight (it quantizes internally); the | ||
| # baseline gets dequantize(quantize(W)) so both do the same matmul. | ||
| W = torch.randn(N, K, dtype=dtype, device=device) | ||
| with torch.no_grad(): | ||
| te_linear.weight.copy_(W) | ||
| torch_linear.weight.copy_(weight_quantizer(W).dequantize()) |
There was a problem hiding this comment.
Nit: It's a bit odd that the TE and reference impls don't actually share the same weight values.
| # Share weights: TE gets the raw weight (it quantizes internally); the | |
| # baseline gets dequantize(quantize(W)) so both do the same matmul. | |
| W = torch.randn(N, K, dtype=dtype, device=device) | |
| with torch.no_grad(): | |
| te_linear.weight.copy_(W) | |
| torch_linear.weight.copy_(weight_quantizer(W).dequantize()) | |
| # Share weights | |
| W = torch.randn(N, K, dtype=dtype, device=device) | |
| W = weight_quantizers(W).dequantize() | |
| with torch.no_grad(): | |
| te_linear.weight.copy_(W) | |
| torch_linear.weight.copy_(W) |
| torch.backends.cudnn.allow_tf32 = prev_tf32_cudnn | ||
|
|
||
|
|
||
| def test_mxfp8_scale_shape_partial_block(): |
There was a problem hiding this comment.
This doesn't test numerics, and generally test_quantized_tensor.py is a more natural place for basic quantization tests.
Remove leading-dim alignment checks that duplicate the cuBLAS FP8 alignment requirement already enforced (with clearer error messages) in cublaslt_gemm.cu via min_alignment_elements(): MXFP8 last_dim%16 and NVFP4 last_dim%32 in both create_tensor and get_scale_shape. Keep quantization-specific checks: NVFP4 last_dim%2 for the 4-bit byte-packed storage, and first_dim%16 for NVFP4 with RHT (Hadamard transform). Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
2 participants
Description
Improved the dimension checks for the FP8 recipes. The old one was too restrictive and did not have any recipe-aware logic.
Dimension check changes
assert_dim_for_fp8_exec(utils.py)M % 8 == 0∧last_dim % 16 == 0, run before the recipe is knownMXFP8Quantizer::create_tensor/get_scale_shapefirst_dim % 32∧last_dim % 32last_dim % 16only (TMA 16 B alignment). Scale tensor usesDIVUPfor partial trailing blocksNVFP4Quantizer::create_tensor/get_scale_shapefirst_dim % 32∧last_dim % 32last_dim % 32always;first_dim % 16only when RHT is enabled (Hadamard block size)swizzle_block_scaling.cukerneldata_rows % 4 == 0data_rowscublaslt_gemm.cuMXFP8 pathlda % 16 == 0∧ldb % 16 == 0viamin_alignment_elements(dtype)usingtypeToNumBits(16 B alignment for FP8)cublaslt_gemm.cuNVFP4 pathlda % 32 == 0∧ldb % 32 == 0(16 B alignment for FP4 = 32 elements)cublaslt_gemm.cu1D block-scaled FP8 pathn % 8 == 0for the native 1D block-scaling pathis_quantizable→supports_quantized_allgather(mxfp8_tensor.py)% 32last_dim % 16; alsofirst_dim % 32when columnwise usage (block boundary across ranks)is_quantizable→supports_quantized_allgather(nvfp4_tensor.py)% 32last_dim % 32; alsofirst_dim % 16when columnwise usageFixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: