diff --git a/primus_turbo/pytorch/ops/grouped_gemm_fp8.py b/primus_turbo/pytorch/ops/grouped_gemm_fp8.py index 95db3011b..7876d0508 100644 --- a/primus_turbo/pytorch/ops/grouped_gemm_fp8.py +++ b/primus_turbo/pytorch/ops/grouped_gemm_fp8.py @@ -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__ = [ @@ -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, @@ -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)." @@ -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 ) - 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) 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, @@ -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, ) @@ -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( @@ -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, @@ -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, @@ -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 @@ -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, @@ -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, @@ -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 @@ -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, @@ -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, diff --git a/tests/pytorch/ops/test_grouped_gemm_fp8.py b/tests/pytorch/ops/test_grouped_gemm_fp8.py index 7ff5bbc45..ca846e65e 100644 --- a/tests/pytorch/ops/test_grouped_gemm_fp8.py +++ b/tests/pytorch/ops/test_grouped_gemm_fp8.py @@ -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,