Add 2d quant for mxfp8 - #2634
Conversation
Greptile SummaryThe PR adds opt-in 2D block quantization for MXFP8 and propagates the setting from recipes through the PyTorch and common C++ quantization paths.
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains in the current code, and all previously reported concerns are fixed or invalid. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart LR
Recipe["MXFP8BlockScaling<br/>enable_2d_quantization"] --> State["MXFP8 recipe state"]
State --> Quantizer["MXFP8Quantizer<br/>with_2d_quantization"]
Quantizer --> Binding["PyTorch C++ binding"]
Binding --> Config["NVTE QuantizationConfig"]
Config --> Dispatch["MXFP8 cast dispatch"]
Dispatch --> Kernel["32x32 block-amax<br/>and E8M0 scale kernel"]
Kernel --> Tensor["MXFP8 data and<br/>row/column scale metadata"]
Reviews (7): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| e8m0_t scale_from_shmem; | ||
| if (thread_lane < THREADS_X) { | ||
| scale_from_shmem = block_scales_2d[thread_lane]; | ||
| } | ||
| // Broadcast: each thread gets scale from lane matching its tid_X_rowwise | ||
| biased_exponent = __shfl_sync(0xffffffff, scale_from_shmem, tid_X_rowwise); |
There was a problem hiding this comment.
scale_from_shmem is potentially uninitialized for threads where thread_lane >= THREADS_X. While __shfl_sync only reads from lanes specified by tid_X_rowwise (which should be < THREADS_X), it's safer to initialize this variable.
| e8m0_t scale_from_shmem; | |
| if (thread_lane < THREADS_X) { | |
| scale_from_shmem = block_scales_2d[thread_lane]; | |
| } | |
| // Broadcast: each thread gets scale from lane matching its tid_X_rowwise | |
| biased_exponent = __shfl_sync(0xffffffff, scale_from_shmem, tid_X_rowwise); | |
| e8m0_t scale_from_shmem = 0; | |
| if (thread_lane < THREADS_X) { | |
| scale_from_shmem = block_scales_2d[thread_lane]; | |
| } |
| @@ -420,7 +421,8 @@ struct QuantizationConfig { | |||
| sizeof(NVTETensor), // rng_seed and offset | |||
| sizeof(uint8_t), // nvfp4_2d_quantization | |||
| sizeof(uint8_t), // stochastic_rounding | |||
There was a problem hiding this comment.
QuantizationConfig layout mismatch
QuantizationConfig::attr_sizes[] was extended with mxfp8_2d_quantization (common.h:415-423), but QuantizationConfig itself uses bool fields. If nvte_set/get_quantization_config_attribute assumes all attributes are packed according to attr_sizes, adding an extra byte-sized attribute without updating any corresponding packing/unpacking logic can break attribute indexing for subsequent attributes (or any code that iterates kNVTEQuantizationConfigNumAttributes). Please double-check the code that uses attr_sizes to ensure the new attribute is reflected everywhere it’s consumed (and that kNVTEQuantizationConfigNumAttributes matches the size of attr_sizes).
| ) # (num_block_rows, num_block_cols, 32, 32) | ||
|
|
||
| # Compute amax for each 32x32 block | ||
| block_amax = torch.amax( | ||
| torch.abs(x_blocks.to(torch.float32)), dim=(-1, -2) | ||
| ) # (num_block_rows, num_block_cols) | ||
|
|
||
| # Convert to E8M0 scale inverse | ||
| block_scale_e8m0 = float_to_e8m0(block_amax) # (num_block_rows, num_block_cols) |
There was a problem hiding this comment.
Incorrect float bit-cast
float_to_e8m0 does val_u32 = val.view(torch.int32) (test_mxfp8_2d_quantize.py:104-106). On PyTorch, .view(dtype) is a numeric cast, not a bit reinterpretation. This makes the reference implementation compute wrong exponents/mantissas and can cause false failures/passes.
Use val.view(torch.int32) only if you’ve explicitly reinterpreted bytes (e.g., via val.view(torch.uint8) + view(torch.int32) on the same storage) or use val.to(torch.int32) with torch.frexp/torch.bitwise_* alternatives. As written, the reference is not modeling the GPU’s IEEE754 bit extraction.
Additional Comments (1)
In If 2D block scaling is only valid when both rowwise+colwise outputs are requested, it should be validated (error) instead of overriding |
Signed-off-by: kunlunl <kunlunl@nvidia.com>
for more information, see https://pre-commit.ci
| False, | ||
| ] | ||
|
|
||
|
|
There was a problem hiding this comment.
MXFP8BlockScaling() uses os.getenv("NVTE_MXFP8_ENABLE_2D_QUANTIZATION", "0") == "1" as its default for enable_2d_quantization. If a developer or CI job has this env var set to "1", the assertion mxfp8_recipe.enable_2d_quantization is False will always fail — even though the recipe is behaving exactly as designed. The test should either monkeypatch the env var to "0" before constructing the recipe, or assert on the attribute when the recipe is explicitly constructed with enable_2d_quantization=False.
Signed-off-by: kunlunl <kunlunl@nvidia.com>
|
Hi @kunlunl, could you please share performance benchmarks for this kernel? |
|
Benchmarked the direct MXFP8 quantization kernel path on a B200 node with preallocated output tensors, timing only
Raw data:
|
| """Reference MXFP8 2D quantization using one scale per 32x32 block.""" | ||
| rows, cols = x.shape | ||
| assert rows % MXFP8_BLOCK_SIZE == 0 | ||
| assert cols % MXFP8_BLOCK_SIZE == 0 |
There was a problem hiding this comment.
Considering this is the only test of this functionality (and it would be better if we also had C++ tests covering it, since it is a change in common), we should test also the cases where the shapes are less nice, and so should not impose those requirements. The MXFP8 quantization kernel does not actually need them - the only real restriction if the 16B alignment from TMA, so the cols % 16 should be 0.
There was a problem hiding this comment.
Partially addressed in Python and fully covered at the common C++ level. I added a C++ operator test in tests/cpp/operator/test_cast_mxfp8.cu with an independent 32x32-block CPU reference, exact E8M0 scale comparison, and FP8 data comparison. It covers rowwise-only, colwise-only, and bidirectional modes over a broader shape matrix including less-nice shapes such as {1,16}, {16,48}, {65,80}, {127,400}, and {993,512}.
For the Python tests, I kept the current shape list 32-aligned. The Python MXFP8Quantizer public allocation path still enforces the existing 32-aligned MXFP8 storage contract, and relaxing it to cols % 16 == 0 would broaden Python allocator behavior beyond the intended scope of this PR. The less-nice shape coverage is therefore in the common C++ test, which exercises the kernel/API path directly.
| If set to `True`, 2D block scaling is used for weight tensors. | ||
| """ | ||
|
|
||
| enable_2d_quantization: bool = os.getenv("NVTE_MXFP8_ENABLE_2D_QUANTIZATION", "0") == "1" |
There was a problem hiding this comment.
This should be the last argument, otherwise it is a breaking API change.
There was a problem hiding this comment.
Checked and addressed.
|
|
||
| def _use_2d_quantization(idx: int) -> bool: | ||
| role = self._slot_role(idx) | ||
| return role.module_type in ("linear", "grouped_linear") and role.tensor_type == "weight" |
There was a problem hiding this comment.
If we want to use it in grouped linear then we should also enable the grouped quantization kernel with this feature.
There was a problem hiding this comment.
Addressed. I restricted the recipe policy to regular linear weight roles only. grouped_linear and other unsupported or unknown roles stay on the existing 1D MXFP8 path. There is also a test that verifies grouped-linear weight roles do not enable 2D quantization.
| (64, 64), | ||
| (128, 128), | ||
| (256, 256), | ||
| (256, 1024), | ||
| (1024, 256), | ||
| (256, 288), | ||
| (320, 320), | ||
| (352, 256), | ||
| (2048, 2048), | ||
| (1024, 2048), | ||
| (2048, 1024), |
There was a problem hiding this comment.
I don't think we need all these shapes to be honest and that increases the CI workload. Could we be more thoughtful here?
There was a problem hiding this comment.
Addressed with a smaller Python shape set. The Python test matrix now keeps representative 32-aligned shapes rather than the full original list. I did keep (320, 320) and (352, 256) because they were "less nice".
| #pragma unroll | ||
| for (int i = 16; i > 0; i /= 2) { | ||
| thread_amax = fmaxf(thread_amax, __shfl_xor_sync(0xffffffff, thread_amax, i)); | ||
| } | ||
| } |
There was a problem hiding this comment.
I'm pretty sure we already have the warp-level reduction somewhere (I think in utils.cuh). If that one is good, please use that instead. If not, please improve it rather than adding another reduction code here.
There was a problem hiding this comment.
Addressed. The colwise 2D reduction now uses the existing warp-level max reduction helper.
| float amax_2d = 0.0f; | ||
| #pragma unroll | ||
| for (int i = 0; i < THREADS_Y; ++i) { | ||
| amax_2d = fmaxf(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); |
There was a problem hiding this comment.
Do we need to use FP32 here? Maybe we could use BF16 version for the BF16 input data.
There was a problem hiding this comment.
Addressed. The rowwise 2D amax staging type now uses IType for the cast-only BF16/FP16 path, reducing shared-memory pressure and avoiding unnecessary FP32 staging there. It still uses float for FP32 input and for activation/dbias paths where keeping FP32 intermediate precision is safer.
Signed-off-by: kunlunl <kunlunl@nvidia.com>
for more information, see https://pre-commit.ci
github-actions
Bot
added
the
community-contribution
| fp8_dpa: bool = False | ||
| fp8_mha: bool = False | ||
| backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) | ||
| enable_2d_quantization: bool = os.getenv("NVTE_MXFP8_ENABLE_2D_QUANTIZATION", "0") == "1" |
There was a problem hiding this comment.
If the env variable is not strictly needed I would prefer to not introduce one, but rather rely on the user just specifying the option in their MXFP8 recipe.
|
|
||
| def _use_2d_quantization(idx: int) -> bool: | ||
| role = self._slot_role(idx) | ||
| return role.module_type == "linear" and role.tensor_type == "weight" |
There was a problem hiding this comment.
There is an issue here with the GroupedLinear - it uses the grouped quantization, which was not updated to be 2D-quant aware, so will break. We should add the 2D awareness to the grouped quantize kernel too.
Signed-off-by: kunlunl <kunlunl@nvidia.com>
for more information, see https://pre-commit.ci
|
/te-ci |
4 participants
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: