Skip to content

feat: add quantized tensor support for blockwise grouped gemm #422

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
xiaobochen-amd merged 1 commit into from
Jul 21, 2026
+163 −36
Merged

feat: add quantized tensor support for blockwise grouped gemm #422

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
103 changes: 67 additions & 36 deletions primus_turbo/pytorch/ops/grouped_gemm_fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,10 +33,12 @@
group_offs_from_lens,
)
from primus_turbo.pytorch.kernels.quantization.quantization_impl import (
grouped_quantize_mxfp8_flydsl_impl,
quant_fp8_blockwise_for_weight_impl,
quant_fp8_blockwise_segment_m_row_col_impl,
quantize_mxfp8_flydsl_impl,
)
from primus_turbo.pytorch.ops.quantization import (
grouped_quantize_fp8_with_trans,
quantize_fp8_with_trans,
)

__all__ = [
Expand Down Expand Up @@ -71,11 +73,15 @@ def _deter_use_nt_layout_gemm_in_bwd(trans_a: bool, trans_b: bool):


class FP8GroupedGemmBlockFunc(torch.autograd.Function):
"""BLOCKWISE grouped GEMM autograd"""

@staticmethod
def forward(
ctx,
a: torch.Tensor,
b: torch.Tensor,
a: Union[torch.Tensor, QuantizedTensor],
b: Union[torch.Tensor, QuantizedTensor],
a_t: Optional[QuantizedTensor], # not used
b_t: Optional[QuantizedTensor], # not used
group_lens: torch.Tensor, # [B,] int64
group_offs: torch.Tensor, # [B + 1,] int64
trans_b: bool,
Expand All @@ -85,6 +91,11 @@ def forward(
):
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."
raise NotImplementedError(
"FP8GroupedGemmBlockFunc does not support a pre-quantized activation `a`"
)
assert a.ndim == 2, "Input tensor must be 2-dimensional."
assert b.ndim == 3, "Weight tensor must be 3-dimensional."
assert group_lens.size(0) == b.size(0), "group_lens size must match b size(0)."
Expand All @@ -93,23 +104,26 @@ def forward(
a_dtype = _get_fp8_dtype(config.format, True)
b_dtype = _get_fp8_dtype(config.format, True)

# One bf16 read of `a` → row-wise (fwd) + segment-padded col-wise (bwd wgrad).
# Row scales are pre-shuffled to the persistent GEMM's scale order. gemm_other_dim
# = fwd-GEMM N lets the quant pick the HIP fast path on small GEMMs.
# --- 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 +107 to 111

b_fp8, b_scale_inv = quant_fp8_blockwise_for_weight_impl(b, b_dtype, block_size=config.block_size)
# --- B side: 2D-block weight, reused unchanged in fwd + bwd. If the caller
# pre-quantized it as a QuantizedTensor, reuse its buffers directly. ---
b_scaling_recipe = ScalingRecipe(use_2d_block=True)
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)
Comment on lines +116 to +120

out = grouped_gemm_fp8_impl(
a_fp8_row,
b_fp8,
a_scale_inv_row,
b_scale_inv,
a_scale_row,
b_scale,
group_lens,
group_offs,
trans_a=False,
Expand All @@ -122,9 +136,9 @@ def forward(

ctx.save_for_backward(
a_fp8_col,
a_scale_inv_col,
a_scale_col,
b_fp8,
b_scale_inv,
b_scale,
group_lens,
group_offs,
)
Expand All @@ -142,23 +156,22 @@ def backward(ctx, grad_out):

(
a_fp8_col,
a_scale_inv_col,
a_scale_col,
b_fp8,
b_scale_inv,
b_scale,
group_lens,
group_offs,
) = ctx.saved_tensors
block_size = ctx.config.block_size
grad_out_dtype = _get_fp8_dtype(ctx.config.format, False)

# One bf16 read of grad_out → row-wise (dgrad) + segment-padded col-wise (wgrad).
# gemm_other_dim = bwd-GEMM K lets the quant pick the HIP fast path on small GEMMs.
# --- grad_out: fused row + segment-padded col grouped quant in one bf16 read.
gemm_k = b_fp8.size(-1) if ctx.trans_b else b_fp8.size(-2)
(
grad_out_fp8_row,
grad_out_fp8_col,
grad_out_scale_inv_row,
grad_out_scale_inv_col,
grad_out_scale_row,
grad_out_scale_col,
var_k_group_lens,
var_k_group_offs,
) = quant_fp8_blockwise_segment_m_row_col_impl(
Expand All @@ -169,8 +182,8 @@ def backward(ctx, grad_out):
grad_a = grouped_gemm_fp8_impl(
grad_out_fp8_row,
b_fp8,
grad_out_scale_inv_row,
b_scale_inv,
grad_out_scale_row,
b_scale,
group_lens,
group_offs,
trans_a=False,
Expand All @@ -184,8 +197,8 @@ def backward(ctx, grad_out):
grad_b = grouped_gemm_fp8_variable_k_impl(
a_fp8_col,
grad_out_fp8_col,
a_scale_inv_col,
grad_out_scale_inv_col,
a_scale_col,
grad_out_scale_col,
var_k_group_lens,
var_k_group_offs,
trans_a=not ctx.trans_a,
Expand All @@ -200,6 +213,8 @@ def backward(ctx, grad_out):
return (
grad_a, # a
grad_b, # b
None, # a_t
None, # b_t
None, # group_lens
None, # group_offs
None, # trans_b
Expand Down Expand Up @@ -604,7 +619,7 @@ def forward(

a_scaling_recipe = ScalingRecipe()
if not isinstance(a, QuantizedTensor):
# NOTE: If a is not a QuantizedTensor use grouped_quantize_mxfp8_flydsl_impl to avoid call dequantize.
# NOTE: If a is not a QuantizedTensor use grouped_quantize_fp8_with_trans to avoid call dequantize.
(
a_fp8_row,
a_scale_row,
Expand All @@ -614,9 +629,10 @@ def forward(
group_offs_padded_rowwise,
_,
_,
) = grouped_quantize_mxfp8_flydsl_impl(
) = grouped_quantize_fp8_with_trans(
a,
a_dtype,
config.granularity,
group_lens,
group_offs,
block_size=config.block_size,
Expand Down Expand Up @@ -650,13 +666,15 @@ def forward(

b_scaling_recipe = ScalingRecipe(use_2d_block=True)
if not isinstance(b, QuantizedTensor):
# NOTE: If b is not a QuantizedTensor use quantize_mxfp8_flydsl_impl to avoid call dequantize.
b_fp8_row, b_scale_row, b_fp8_col, b_scale_col = quantize_mxfp8_flydsl_impl(
# NOTE: If b is not a QuantizedTensor use quantize_fp8_with_trans to avoid call dequantize.

b_fp8_row, b_scale_row, b_fp8_col, b_scale_col = quantize_fp8_with_trans(
b,
b_dtype,
config.granularity,
block_size=config.block_size,
scaling_recipe=b_scaling_recipe,
scaling_recipe_for_trans=b_scaling_recipe,
scaling_recipe=ScalingRecipe(use_2d_block=True),
scaling_recipe_for_trans=ScalingRecipe(use_2d_block=True),
)
else:
quantized_b = b
Expand Down Expand Up @@ -730,9 +748,10 @@ def backward(ctx, grad_out: torch.Tensor):
group_offs_padded_rowwise,
group_lens_padded_colwise,
group_offs_padded_colwise,
) = grouped_quantize_mxfp8_flydsl_impl(
) = grouped_quantize_fp8_with_trans(
grad_out,
grad_out_dtype,
ctx.config.granularity,
group_lens,
group_offs,
block_size=ctx.config.block_size,
Expand Down Expand Up @@ -878,9 +897,21 @@ def grouped_gemm_fp8(
num_cu,
)
elif config.granularity == ScalingGranularity.BLOCKWISE:
# BLOCKWISE only accepts raw tensors today; preserve existing assertion
# behaviour in ``FP8GroupedGemmBlockFunc.forward``.
return FP8GroupedGemmBlockFunc.apply(a, b, group_lens, group_offs, trans_b, out_dtype, config, num_cu)
# BLOCKWISE accepts a pre-quantized 2D-block weight (``b``); the activation
# ``a`` must stay a raw tensor (fused pre-shuffled quant). ``a_data_t`` /
# ``b_data_t`` are unused by ``FP8GroupedGemmBlockFunc``.
return FP8GroupedGemmBlockFunc.apply(
a_data,
b_data,
a_data_t,
b_data_t,
group_lens,
group_offs,
trans_b,
out_dtype,
config,
num_cu,
)
elif config.granularity == ScalingGranularity.MX_BLOCKWISE:
return FP8GroupedGemmMXFunc.apply(
a_data,
Expand Down
96 changes: 96 additions & 0 deletions tests/pytorch/ops/test_grouped_gemm_fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -799,6 +799,102 @@ def test_grouped_gemm_fp8_mx_blockwise_quantized_tensor(
)


@pytest.mark.parametrize("B", B_VALUES)
@pytest.mark.parametrize("M", M_VALUES)
@pytest.mark.parametrize("NK", NK_VALUES)
@pytest.mark.parametrize("ori_dtype", ORI_DTYPE_VALUES)
@pytest.mark.parametrize("format", FORMAT_VALUES)
@pytest.mark.parametrize("block_size", [128])
@pytest.mark.parametrize("trans_b", TRANS_B_VALUES)
@pytest.mark.parametrize("balance", BALANCE_VALUES)
@pytest.mark.parametrize("backend", [None, BackendType.TRITON])
@pytest.mark.parametrize("auto_tune", [False, True])
def test_grouped_gemm_fp8_blockwise_weight_quantized_tensor(
B, M, NK, ori_dtype, format, block_size, trans_b, balance, backend, auto_tune
):
"""BLOCKWISE grouped_gemm with a pre-quantized 2D-block weight (``b``); ``a`` raw.

NOTE: Grouped BLOCKWISE only supports the weight side as a QuantizedTensor: the
activation ``a`` must stay a raw tensor. So ``a`` is kept high-precision and
only ``b`` is externally quantized.
"""
N, K = NK
seed = 42
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

# Skip redundant test: auto_tune is ignored when backend is explicitly specified
if backend is not None and auto_tune:
pytest.skip("auto_tune is ignored when backend is explicitly specified")

if _check_hit_int32_limit(B, M, N, K):
pytest.skip("Shape hits int32 indexing limit (numel >= 2**31).")

GlobalBackendManager.set_grouped_gemm_backend(backend)
GlobalBackendManager.set_auto_tune(auto_tune)

device = "cuda:0"
group_lens = generate_grouped_gemm_group_lens(B, M, balance=balance).to(device)
print(
f"\n[QT-BLOCKWISE-weight] B={B}, M={M}, N={N}, K={K}, ori_dtype={ori_dtype}, "
f"format={format}, block_size={block_size}, trans_b={trans_b}, balance={balance}, "
f"backend={backend}, auto_tune={auto_tune}"
)

b_shape = (B, N, K) if trans_b else (B, K, N)
a = torch.randn((B * M, K), dtype=ori_dtype, device=device, requires_grad=True)
b = torch.randn(b_shape, dtype=ori_dtype, device=device, requires_grad=True)
a_ref = a.detach().clone().requires_grad_(True)
b_ref = b.detach().clone().requires_grad_(True)
torch.cuda.synchronize()

# Externally quantize only the weight as a 2D-block QuantizedTensor.
fwd_dtype = _get_fp8_dtype(format, is_fwd=True)
b_scaling_recipe = ScalingRecipe(use_2d_block=True)
qt_b = QuantizedTensor.quantize(
b,
fwd_dtype,
ScalingGranularity.BLOCKWISE,
axis=-1 if trans_b else -2,
block_size=block_size,
scaling_recipe=b_scaling_recipe,
)

# Reference
out_ref = grouped_gemm_ref(a_ref, b_ref, group_lens, trans_b)
grad_out = torch.randn_like(out_ref)
out_ref.backward(grad_out)
torch.cuda.synchronize()

config = Float8QuantConfig(format=format, granularity=ScalingGranularity.BLOCKWISE, block_size=block_size)
out = grouped_gemm_fp8(a, qt_b, group_lens, trans_b=trans_b, config=config)
out.backward(grad_out)
torch.cuda.synchronize()

# Check Shape
assert out.shape == out_ref.shape
assert a.grad is not None and a.grad.shape == a_ref.grad.shape
assert qt_b.grad is not None and qt_b.grad.shape == b.shape

# Check Results
snr_threshold = 25 if format == Format.E4M3 else 20

out_snr = compute_snr(out_ref, out)
a_grad_snr = compute_snr(a_ref.grad, a.grad)
b_grad_snr = compute_snr(b_ref.grad, qt_b.grad)
print(
f"[QT-BLOCKWISE-weight] Out-SNR={out_snr:.2f} dB, "
f"AGrad-SNR={a_grad_snr:.2f} dB, BGrad-SNR={b_grad_snr:.2f} dB"
)
assert out_snr > snr_threshold, f"out_snr={out_snr:.2f} too low"
assert a_grad_snr > snr_threshold, f"a_grad_snr={a_grad_snr:.2f} too low"
assert b_grad_snr > snr_threshold, f"b_grad_snr={b_grad_snr:.2f} too low"

# Reset config and caches
GlobalBackendManager.reset()


def _test_grouped_gemm_fp8_hipgraph_test(
B: int,
M: int,
Expand Down
Loading