feat: add quantized tensor support for blockwise grouped gemm#422
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends the FP8 BLOCKWISE grouped GEMM path to accept a pre-quantized 2D-block QuantizedTensor for the weight (b), avoiding an internal dequantize/re-quantize cycle, and migrates the MX_BLOCKWISE path to unified granularity-aware quantization entry points.
Changes:
- Allow
FP8GroupedGemmBlockFuncto takebas aQuantizedTensorand reuseqdata/scale_invin forward and backward; explicitly reject pre-quantized activations for BLOCKWISE. - Align BLOCKWISE autograd function signatures/dispatch with other grouped GEMM paths by adding placeholder
a_t/b_targuments and updating backward returns. - Add a new test covering BLOCKWISE grouped GEMM with an externally pre-quantized 2D-block weight.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| tests/pytorch/ops/test_grouped_gemm_fp8.py | Adds coverage for BLOCKWISE grouped GEMM when b is provided as a pre-quantized 2D-block QuantizedTensor. |
| primus_turbo/pytorch/ops/grouped_gemm_fp8.py | Enables BLOCKWISE grouped GEMM to reuse pre-quantized weight buffers and migrates MX_BLOCKWISE to unified quantization helpers. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| assert config.granularity == ScalingGranularity.BLOCKWISE | ||
| assert config.block_size in [128], "Only block_size 128 is supported currently." | ||
| if isinstance(a, QuantizedTensor): | ||
| # TODO(ruibin): grouped BLOCKWISE emits fused pre-shuffled row / segment-padded col scales, from the quant kernel, it is not compatible with the QuantizedTensor." |
Comment on lines
+107
to
111
| # --- A side: fused row + segment-padded col grouped quant in one bf16 read. | ||
| gemm_n = b.size(-2) if trans_b else b.size(-1) | ||
| a_fp8_row, a_fp8_col, a_scale_inv_row, a_scale_inv_col, _, _ = ( | ||
| quant_fp8_blockwise_segment_m_row_col_impl( | ||
| a, a_dtype, config.block_size, group_lens, group_offs, gemm_other_dim=gemm_n | ||
| ) | ||
| a_fp8_row, a_fp8_col, a_scale_row, a_scale_col, _, _ = quant_fp8_blockwise_segment_m_row_col_impl( | ||
| a, a_dtype, config.block_size, group_lens, group_offs, gemm_other_dim=gemm_n | ||
| ) |
Comment on lines
+116
to
+120
| if isinstance(b, QuantizedTensor): | ||
| check_quantized_tensor(b, config, scaling_recipe=b_scaling_recipe) | ||
| b_fp8, b_scale = b.qdata, b.scale_inv | ||
| else: | ||
| b_fp8, b_scale = quant_fp8_blockwise_for_weight_impl(b, b_dtype, block_size=config.block_size) |
xiaobochen-amd
approved these changes
Jul 21, 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
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
This PR adds
QuantizedTensorsupport to the BLOCKWISE FP8 grouped GEMM path, so a pre-quantized 2D-block weight can be passed directly intogrouped_gemm_fp8without going through an internal dequantize/re-quantize round trip.Previously the BLOCKWISE grouped GEMM only accepted raw high-precision tensors and asserted on anything else.
It now accepts a
QuantizedTensorweight (b) and reuses itsqdata/scale_invbuffers directly in both forward and backward.The activation (
a) must remain a raw tensor, because the grouped BLOCKWISE quant kernel emits fused pre-shuffled row scales plus segment-padded column scales that are not compatible with the standardQuantizedTensorlayout.As part of this change the MX_BLOCKWISE path is migrated onto the unified, granularity-aware quant entry points (
quantize_fp8_with_trans/grouped_quantize_fp8_with_trans), replacing the MX-specific*_mxfp8_flydsl_implhelpers so both granularities share one quant interface.Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
FP8GroupedGemmBlockFuncto accept a pre-quantized 2D-block weightbas aQuantizedTensor, reusing itsqdata/scale_invin forward and backward instead of re-quantizing.a_t/b_tplaceholder arguments (currently unused) to theFP8GroupedGemmBlockFunc.forwardsignature and update the backward return tuple and thegrouped_gemm_fp8dispatch accordingly, aligning the API with the other grouped GEMM functions.NotImplementedErrorwhen a pre-quantized activationais supplied to the BLOCKWISE path, since the fused pre-shuffled row / segment-padded col scales from the grouped quant kernel are not compatible withQuantizedTensor.QuantizedTensorviacheck_quantized_tensoragainst theScalingRecipe(use_2d_block=True)recipe.*_scale_inv_*to*_scale_*for consistency across the forward/backward code paths.quantize_fp8_with_trans/grouped_quantize_fp8_with_trans(passingconfig.granularity) in place of the MX-specific*_mxfp8_flydsl_implfunctions.test_grouped_gemm_fp8_blockwise_weight_quantized_tensorcovering a pre-quantized 2D-block weight with a raw activation, checking output and gradient SNR across formats, backends, auto-tune, and transpose configurations.Checklist: