From 3247c50d08bcd956d08aa360a06ce0758f6bb86d Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Tue, 18 Aug 2026 16:48:54 +0000 Subject: [PATCH 01/13] Add ROCm Triton blockwise FP8 grouped GEMM path Implement DeepSeek-style blockwise FP8 grouped GEMM for the PyTorch backend, selected from GroupedLinear under Float8BlockScaling when NVTE_USE_BLOCKWISE_FP8_GROUPED_GEMM=1. - Add _forward_blockwise_fp8/_backward_blockwise_fp8 to transformer_engine/pytorch/module/grouped_linear.py, using 1x128 rowwise activation quantization, 128x128 weight quantization, and a segment-padded columnwise activation for wgrad. - Add new Triton kernels: - blockwise_fp8_grouped_gemm.py: persistent grouped blockwise FP8 GEMM and variable-K wgrad kernels. - blockwise_quantize.py: 1x128 / 128x128 blockwise FP8 quantization kernels, including segment-padded columnwise quantize for variable M. - Add is_cdna4() to triton_kernels/common.py for gfx950 detection. --- .../pytorch/module/grouped_linear.py | 234 +++++++ .../blockwise_fp8_grouped_gemm.py | 572 ++++++++++++++++++ .../triton_kernels/blockwise_quantize.py | 328 ++++++++++ .../pytorch/triton_kernels/common.py | 3 + 4 files changed, 1137 insertions(+) create mode 100644 transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py create mode 100644 transformer_engine/pytorch/triton_kernels/blockwise_quantize.py diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index f534da5c3b..e9b7d84eec 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -399,6 +399,205 @@ def _forward_grouped_tensor( return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces + @staticmethod + def _forward_blockwise_fp8( + ctx, + *, + inp, + m_splits, + weights, + biases, + use_bias, + fp8, + recipe, + wgrad_store, + weight_quantizers, + fuse_wgrad_accumulation, + cpu_offloading, + activation_dtype, + is_grad_enabled, + save_original_input, + debug, + actual_m_splits, + unpad_output, + backward_override, + ): + """DeepSeek-style blockwise FP8 grouped GEMM forward (ROCm Triton). + + Selected from :meth:`forward` under the ``Float8BlockScaling`` recipe when + ``NVTE_USE_BLOCKWISE_FP8_GROUPED_GEMM=1``. Activations are quantized rowwise + (1x128 along K), weights 128x128; the backward wgrad (see + :meth:`_backward_blockwise_fp8`) uses a segment-padded columnwise operand + quantized from the original high-precision input (no double-quantization). + Recipe configurations and orchestration features it does not yet support are + rejected instead of silently ignored. + """ + from ..triton_kernels.common import te_dtype_to_torch_dtype + from ..triton_kernels.blockwise_quantize import ( + quantize_fp8_blockwise, + quantize_fp8_blockwise_weight, + quantize_fp8_blockwise_segment_m, + ) + from ..triton_kernels.blockwise_fp8_grouped_gemm import ( + grouped_gemm_fp8_blockwise_triton_kernel, + ) + + num_gemms = len(m_splits) + + # fp8 + block-scaling recipe are internal invariants (asserts, guaranteed by the + # caller's gate); the rest are user-reachable (raise). + assert fp8, "blockwise grouped FP8 path requires fp8=True" + assert recipe.float8_block_scaling(), "blockwise grouped FP8 path requires Float8BlockScaling" + + # The Triton kernels hardcode the DeepSeek layout (x rowwise 1x128, w 128x128, + # grad rowwise 1x128). Reject recipe block-scaling dims that don't match rather + # than silently producing wrong numerics. + if ( + recipe.x_block_scaling_dim != 1 + or recipe.w_block_scaling_dim != 2 + or recipe.grad_block_scaling_dim != 1 + ): + raise NotImplementedError( + "the blockwise grouped FP8 path only supports x_block_scaling_dim=1, " + "w_block_scaling_dim=2, grad_block_scaling_dim=1 (got " + f"{recipe.x_block_scaling_dim}, {recipe.w_block_scaling_dim}, " + f"{recipe.grad_block_scaling_dim})" + ) + if use_bias: + raise NotImplementedError("bias is not supported in the blockwise grouped FP8 path yet") + if backward_override is not None: + raise NotImplementedError( + "backward_override is not supported in the blockwise grouped FP8 path yet" + ) + if fuse_wgrad_accumulation: + raise NotImplementedError( + "fuse_wgrad_accumulation (gradient_accumulation_fusion) is not yet supported in " + "the ROCm blockwise grouped FP8 path. Pass --no-gradient-accumulation-fusion to the " + "training script: wgrad is then returned as a plain gradient and Megatron's DDP " + "post-hook accumulates it into the fp32 main_grad (numerically equivalent for bring-up)." + ) + if cpu_offloading: + raise NotImplementedError( + "cpu_offloading is not supported in the blockwise grouped FP8 path yet" + ) + if save_original_input: + raise NotImplementedError( + "save_original_input is not supported in the blockwise grouped FP8 path yet" + ) + if wgrad_store is not None and wgrad_store.delay_wgrad_compute(): + raise NotImplementedError( + "delayed wgrad is not supported in the blockwise grouped FP8 path yet" + ) + if debug: + raise NotImplementedError( + "debug quantization is not supported in the blockwise grouped FP8 path yet" + ) + if unpad_output or ( + actual_m_splits is not None and list(actual_m_splits) != list(m_splits) + ): + raise NotImplementedError( + "the ROCm fused-pad / unpad_output path is not supported in the blockwise grouped " + "FP8 path yet" + ) + + # Resolve the FP8 torch dtype via the quantizer, normalizing through ``tex.DType`` + # so ``te_dtype_to_torch_dtype`` picks the arch-correct variant (fnuz on CDNA3). + if weight_quantizers and weight_quantizers[0] is not None: + dt = te_dtype_to_torch_dtype(tex.DType(int(weight_quantizers[0].dtype))) + else: + dt = torch.float8_e4m3fn + in_features = weights[0].size(-1) + out_features = weights[0].size(0) + device = inp.device + + inp_shape = inp.shape + a = inp.reshape(-1, in_features).to(activation_dtype).contiguous() + + # Group offsets built on-device from the split tensor (graph-capture safe, + # no device->host sync). + group_lens = m_splits.to(device=device, dtype=torch.int64) + group_offs = torch.zeros(num_gemms + 1, dtype=torch.int64, device=device) + group_offs[1:] = torch.cumsum(group_lens, 0) + + w = torch.stack([wt.to(activation_dtype).contiguous() for wt in weights], 0).contiguous() + + # Quantize: activation rowwise (1x128 along K), weights 128x128. + a_row, a_srow = quantize_fp8_blockwise(a, dt, axis=1, block_size=128) + b_fp8, b_scale = quantize_fp8_blockwise_weight(w, dt, block_size=128) + + # Forward GEMM: out[seg] = A[seg] @ W[g]^T (trans_b=True). + out = grouped_gemm_fp8_blockwise_triton_kernel( + a_row, b_fp8, a_srow, b_scale, group_offs, trans_b=True, out_dtype=activation_dtype + ) + + if is_grad_enabled: + ctx.use_blockwise_fp8 = True + # Segment-padded columnwise activation for the variable-K wgrad. + a_col, a_scol, _vk_lens, vk_offs = quantize_fp8_blockwise_segment_m( + a, dt, 128, group_lens, group_offs + ) + ctx.save_for_backward(a_col, a_scol, b_fp8, b_scale, group_lens, group_offs, vk_offs) + ctx.dt = dt + ctx.activation_dtype = activation_dtype + ctx.num_gemms = num_gemms + ctx.inp_shape = inp_shape + ctx.out_features = out_features + ctx.requires_dgrad = inp.requires_grad + ctx.weight_requires_grad = weights[0].requires_grad + ctx.reduce_and_update_bwd_fp8_tensors = False + + new_workspaces = [None] * num_gemms + return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces + + @staticmethod + def _backward_blockwise_fp8(ctx, grad_output): + """Backward path paired with :meth:`_forward_blockwise_fp8`.""" + from ..triton_kernels.blockwise_quantize import ( + quantize_fp8_blockwise, + quantize_fp8_blockwise_segment_m, + ) + from ..triton_kernels.blockwise_fp8_grouped_gemm import ( + grouped_gemm_fp8_blockwise_triton_kernel, + grouped_gemm_fp8_blockwise_variable_k_triton_kernel, + ) + + a_col, a_scol, b_fp8, b_scale, group_lens, group_offs, vk_offs = ctx.saved_tensors + dt = ctx.dt + g_out = grad_output.reshape(-1, ctx.out_features).contiguous() + + # dgrad: dX[seg] = dY[seg] @ W[g] (trans_b=False) -> [M_total, K]. + dgrad = None + if ctx.requires_dgrad: + go_row, go_srow = quantize_fp8_blockwise(g_out, dt, axis=1, block_size=128) + dgrad = grouped_gemm_fp8_blockwise_triton_kernel( + go_row, b_fp8, go_srow, b_scale, group_offs, trans_b=False, + out_dtype=ctx.activation_dtype, + ) + + # wgrad: dW[g] = dY[g]^T @ X[g] (variable-K) -> [G, N, K]. + wgrad_list = [None] * ctx.num_gemms + if ctx.weight_requires_grad: + go_col, go_scol, _l, _o = quantize_fp8_blockwise_segment_m( + g_out, dt, 128, group_lens, group_offs + ) + dW = grouped_gemm_fp8_blockwise_variable_k_triton_kernel( + go_col, a_col, go_scol, a_scol, vk_offs, out_dtype=ctx.activation_dtype + ) + wgrad_list = [dW[g].contiguous() for g in range(ctx.num_gemms)] + + grad_biases = [None] * ctx.num_gemms # bias rejected in forward + + # Grads match forward inputs: (inp, m_splits, dispatched_probs, non_tensor_args, + # *weights, *biases). + return ( + dgrad.view(ctx.inp_shape) if dgrad is not None else None, + None, # m_splits + None, # dispatched_probs + None, # non_tensor_args + *wgrad_list, + *grad_biases, + ) + # pylint: disable=keyword-arg-before-vararg @staticmethod def forward( @@ -454,6 +653,39 @@ def forward( device = inp.device weight_requires_grad = weights[0].requires_grad + # DeepSeek-style blockwise FP8 grouped GEMM (ROCm Triton) opt-in. Only engages under + # the Float8BlockScaling recipe (whose quantizers already carry blockwise semantics); + # it runs its own quantization + grouped GEMM and returns early, bypassing the default + # quantizer setup below. + use_blockwise_fp8 = ( + IS_HIP_EXTENSION + and fp8 + and os.getenv("NVTE_USE_BLOCKWISE_FP8_GROUPED_GEMM", "0") == "1" + and FP8GlobalStateManager.get_fp8_recipe().float8_block_scaling() + ) + if use_blockwise_fp8: + return _GroupedLinear._forward_blockwise_fp8( + ctx, + inp=inp, + m_splits=m_splits, + weights=weights, + biases=biases, + use_bias=use_bias, + fp8=fp8, + recipe=FP8GlobalStateManager.get_fp8_recipe(), + wgrad_store=wgrad_store, + weight_quantizers=weight_quantizers, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + cpu_offloading=cpu_offloading, + activation_dtype=activation_dtype, + is_grad_enabled=is_grad_enabled, + save_original_input=save_original_input, + debug=debug, + actual_m_splits=actual_m_splits, + unpad_output=unpad_output, + backward_override=backward_override, + ) + # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): if FP8GlobalStateManager.get_fp8_recipe().custom(): @@ -951,6 +1183,8 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring with get_nvtx_range_context("_GroupedLinear_backward"): + if getattr(ctx, "use_blockwise_fp8", False): + return _GroupedLinear._backward_blockwise_fp8(ctx, grad_output) if ctx.use_grouped_tensor_path: return _GroupedLinear._backward_grouped_tensor(ctx, grad_output) diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py new file mode 100644 index 0000000000..f31de8b0ff --- /dev/null +++ b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py @@ -0,0 +1,572 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +# +# Blockwise FP8 grouped GEMM Triton kernels (MoE), adapted from AMD Primus-Turbo +# (primus_turbo/triton/grouped_gemm/grouped_gemm_fp8_kernel.py). + +import os + +import torch +import triton +import triton.language as tl + +from .common import is_cdna4 +from .gmm.pid_preprocessing import remap_xcd_chunked + + +def get_num_cus() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +# AMD gfx950 compiler knobs. +_KNOBS_SET = False + + +def set_triton_knobs_gfx950() -> None: + """Enable AMD compiler knobs for gfx950 (async_copy, block_pingpong, scalarize).""" + global _KNOBS_SET + if _KNOBS_SET: + return + _KNOBS_SET = True + if hasattr(triton, "knobs") and hasattr(triton.knobs, "amd"): + triton.knobs.amd.use_async_copy = True + triton.knobs.amd.scalarize_packed_fops = True + triton.knobs.amd.use_block_pingpong = True + else: + os.environ.setdefault("TRITON_HIP_USE_ASYNC_COPY", "1") + os.environ.setdefault("AMDGCN_SCALARIZE_PACKED_FOPS", "1") + os.environ.setdefault("TRITON_HIP_USE_BLOCK_PINGPONG", "1") + + +def _set_amd_knobs(enable: bool = True): + """Set AMD-specific Triton knobs (non-gfx950 fallback). + NOTE: use_async_copy and scalarize_packed_fops help NT/NN but regress + TN/wgrad ~5-8% on gfx942, so callers gate ``enable`` per layout. + """ + if hasattr(triton, "knobs") and hasattr(triton.knobs, "amd"): + triton.knobs.amd.use_async_copy = enable + triton.knobs.amd.scalarize_packed_fops = enable + + +NUM_XCDS = 8 + + +# Blockwise grouped FP8 kernel and public entrypoint + +@triton.jit() +def _grouped_blockwise_fp8_persistent_gemm_kernel( + # Pointers + A, # [M_total, K] FP8 + B, # [G, ?, ?] FP8 + C, # [M_total, N] + A_scales_ptr, # [K//128, M_total] float32 (pre-transposed for coalesced access) + B_scales_ptr, # [G, ?, ?] float32 (block-wise, layout depends on trans_b) + group_offs_ptr, # [G+1] int64 + # Dimensions + G, # number of groups (runtime) + N, + K, + # Strides + stride_am, # A row stride + stride_bg, # B group stride: b.stride(0) + stride_bn, # B N-stride (within a group) + stride_cm, # C row stride + stride_cn, # C col stride + # A_scales strides (pre-transposed: [K//128, M_total]) + stride_as_k, # A_scales_t.stride(0) + stride_as_m, # A_scales_t.stride(1) + # B_scales strides + stride_bs_g, # B_scales.stride(0) — group stride + stride_bs_n, # stride along N-block dimension + stride_bs_k, # stride along K-block dimension + # Constexpr strides + stride_ak: tl.constexpr, + stride_bk: tl.constexpr, + # Tile config + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_SMS: tl.constexpr, + NUM_XCDS: tl.constexpr, + CHUNK_SIZE: tl.constexpr, + EVEN_K: tl.constexpr, + CACHE_MODIFIER: tl.constexpr, +): + """Persistent grouped block-wise FP8 GEMM kernel (CPU-sync-free).""" + pid = tl.program_id(0) + if NUM_XCDS != 1: + pid = remap_xcd_chunked(pid, NUM_SMS, NUM_XCDS, CHUNK_SIZE) + + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + # ── Compute total tiles across all groups ── + total_tiles: tl.int32 = 0 + for _g in range(G): + m_g = (tl.load(group_offs_ptr + _g + 1) - tl.load(group_offs_ptr + _g)).to(tl.int32) + total_tiles += tl.cdiv(m_g, BLOCK_SIZE_M) * num_pid_n + + tl.assume(stride_am > 0) + tl.assume(stride_ak > 0) + tl.assume(stride_bn > 0) + tl.assume(stride_bk > 0) + tl.assume(stride_cm > 0) + tl.assume(stride_cn > 0) + + acc_dtype = tl.float32 + + for global_tile_id in range(pid, total_tiles, NUM_SMS): + # ── Find group via linear scan (O(G)) ── + group_idx: tl.int32 = 0 + tile_start: tl.int32 = 0 + cumsum: tl.int32 = 0 + for _g in range(G): + m_g_i = (tl.load(group_offs_ptr + _g + 1) - tl.load(group_offs_ptr + _g)).to(tl.int32) + tiles_g = tl.cdiv(m_g_i, BLOCK_SIZE_M) * num_pid_n + new_cumsum = cumsum + tiles_g + if global_tile_id >= new_cumsum: + group_idx = _g + 1 + tile_start = new_cumsum + cumsum = new_cumsum + + # ── Group-local tile → (pid_m, pid_n) ── + local_tile = global_tile_id - tile_start + m_start_g = tl.load(group_offs_ptr + group_idx) # int64 + M_g = (tl.load(group_offs_ptr + group_idx + 1) - m_start_g).to(tl.int32) + tiles_m_g = tl.cdiv(M_g, BLOCK_SIZE_M) + + num_pid_in_group = GROUP_SIZE_M * num_pid_n + swizzle_group = local_tile // num_pid_in_group + first_pid_m = swizzle_group * GROUP_SIZE_M + group_size_m = min(tiles_m_g - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((local_tile % num_pid_in_group) % group_size_m) + pid_n = (local_tile % num_pid_in_group) // group_size_m + tl.assume(pid_m >= 0) + tl.assume(pid_n >= 0) + + # ── Address computation ── + rm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M_g + rn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + rk = tl.arange(0, BLOCK_SIZE_K) + rn = tl.max_contiguous(tl.multiple_of(rn, BLOCK_SIZE_N), BLOCK_SIZE_N) + + group_offset_b = group_idx.to(tl.int64) * stride_bg + A_BASE = A + m_start_g * stride_am + rm[:, None] * stride_am + rk[None, :] * stride_ak + B_BASE = B + group_offset_b + rk[:, None] * stride_bk + rn[None, :] * stride_bn + + # A_scales pointer: pre-transposed [K//128, M_total] + as_ptrs_base = A_scales_ptr + (m_start_g + rm.to(tl.int64)) * stride_as_m + + # B_scales pointer: B_scales[g, pn, ki] (2D block scaling) + bs_ptr_base = B_scales_ptr + group_idx.to(tl.int64) * stride_bs_g + pid_n * stride_bs_n + + # ── K-loop with block-wise scaling (EVEN_K pattern) ── + loop_k = tl.cdiv(K, BLOCK_SIZE_K) + if not EVEN_K: + loop_k -= 1 + tl.assume(loop_k > 1) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype) + + for ki in range(0, loop_k): + if stride_ak == 1: + a = tl.load(tl.multiple_of(A_BASE, (1, 16)), cache_modifier=CACHE_MODIFIER) + else: + a = tl.load(tl.multiple_of(A_BASE, (16, 1)), cache_modifier=CACHE_MODIFIER) + + if stride_bk == 1: + b = tl.load(tl.multiple_of(B_BASE, (16, 1)), cache_modifier=CACHE_MODIFIER) + else: + b = tl.load(tl.multiple_of(B_BASE, (1, 16)), cache_modifier=CACHE_MODIFIER) + + partial = tl.dot(a, b) + + # Block-wise scales: a_s is [BLOCK_M] vector, b_s is scalar + a_s = tl.load(as_ptrs_base + ki * stride_as_k) + b_s = tl.load(bs_ptr_base + ki * stride_bs_k) + acc += partial * (a_s * b_s)[:, None] + + A_BASE += BLOCK_SIZE_K * stride_ak + B_BASE += BLOCK_SIZE_K * stride_bk + + if not EVEN_K: + # ── Last partial K-block (masked) ── + rk_last = loop_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + A_LAST = A + m_start_g * stride_am + rm[:, None] * stride_am + rk_last[None, :] * stride_ak + B_LAST = B + group_offset_b + rk_last[:, None] * stride_bk + rn[None, :] * stride_bn + if stride_ak == 1: + A_LAST = tl.multiple_of(A_LAST, (1, 16)) + else: + A_LAST = tl.multiple_of(A_LAST, (16, 1)) + if stride_bk == 1: + B_LAST = tl.multiple_of(B_LAST, (16, 1)) + else: + B_LAST = tl.multiple_of(B_LAST, (1, 16)) + a = tl.load(A_LAST, mask=rk_last[None, :] < K, other=0.0, cache_modifier=CACHE_MODIFIER) + b = tl.load(B_LAST, mask=rk_last[:, None] < K, other=0.0, cache_modifier=CACHE_MODIFIER) + partial = tl.dot(a, b) + a_s = tl.load(as_ptrs_base + loop_k * stride_as_k) + b_s = tl.load(bs_ptr_base + loop_k * stride_bs_k) + acc += partial * (a_s * b_s)[:, None] + + # ── Store output ── + c = acc.to(C.type.element_ty) + rm_s = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M_g + rn_s = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + rn_s = tl.max_contiguous(tl.multiple_of(rn_s, BLOCK_SIZE_N), BLOCK_SIZE_N) + c_mask = (rm_s[:, None] < M_g) & (rn_s[None, :] < N) + C_ = C + m_start_g * stride_cm + rm_s[:, None] * stride_cm + rn_s[None, :] * stride_cn + tl.store(C_, c, c_mask) + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Blockwise FP8 Variable-K Backward Kernel (persistent, CPU-sync-free) +# +# Computes: C[g] = LHS[g]^T @ RHS[g] with 1D+1D block-wise scales +# ═══════════════════════════════════════════════════════════════════════════════ + + + +def grouped_gemm_fp8_blockwise_triton_kernel( + a: torch.Tensor, + b: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + group_offs: torch.Tensor, + trans_b: bool = True, + out_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Persistent grouped block-wise FP8 GEMM (CPU-sync-free) using Triton. + + Computes: out[offs[g]:offs[g+1], :] = A[offs[g]:offs[g+1], :] @ B_view[g] + with block-wise scaling for each K-block. + + Args: + a: [M_total, K] FP8 input (trans_a=False always). + b: [G, N, K] (if trans_b=True) or [G, K, N] FP8 weights. + a_scales: [M_total, K//128] float32, block-wise scale for A. + b_scales: [G, ceil(N/128), ceil(K/128)] or [G, ceil(K/128), ceil(N/128)] float32. + group_offs: [G+1] int64 prefix sum of group lengths. + trans_b: If True, b[g] is [N, K] (transposed). + out_dtype: Output dtype (default bfloat16). + + Returns: + [M_total, N] output in out_dtype. + """ + if is_cdna4(): + set_triton_knobs_gfx950() + else: + _set_amd_knobs(enable=True) + + assert a.ndim == 2, f"a must be 2D, got {a.shape}" + assert b.ndim == 3, f"b must be 3D, got {b.shape}" + assert b_scales.ndim == 3, f"b_scales must be 3D, got {b_scales.shape}" + + M_total, K = a.shape + G = b.shape[0] + + if trans_b: + N = b.shape[1] + stride_bk = b.stride(2) + stride_bn = b.stride(1) + stride_bs_n = b_scales.stride(1) + stride_bs_k = b_scales.stride(2) + else: + N = b.shape[2] + stride_bk = b.stride(1) + stride_bn = b.stride(2) + stride_bs_n = b_scales.stride(2) + stride_bs_k = b_scales.stride(1) + + stride_bg = b.stride(0) + stride_ak = a.stride(1) + + out = torch.empty((M_total, N), device=a.device, dtype=out_dtype) + A_scales_t = a_scales.T.contiguous() + num_sms = get_num_cus() + + blk_m = 256 + blk_n = 128 # Keep 128 to match B_scale block alignment + blk_k = 128 + even_k = K % blk_k == 0 + + # GROUP_SIZE_M heuristic (match tensorwise) + tiles_m_per_group = (M_total + G * blk_m - 1) // (G * blk_m) + tiles_n = (N + blk_n - 1) // blk_n + group_m = 8 if min(tiles_m_per_group, tiles_n) < 16 else 4 + + _grouped_blockwise_fp8_persistent_gemm_kernel[(num_sms,)]( + a, + b, + out, + A_scales_t, + b_scales, + group_offs, + G, + N, + K, + a.stride(0), + stride_bg, + stride_bn, + out.stride(0), + out.stride(1), + A_scales_t.stride(0), + A_scales_t.stride(1), + b_scales.stride(0), + stride_bs_n, + stride_bs_k, + stride_ak=stride_ak, + stride_bk=stride_bk, + BLOCK_SIZE_M=blk_m, + BLOCK_SIZE_N=blk_n, + BLOCK_SIZE_K=blk_k, + GROUP_SIZE_M=group_m, + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + CHUNK_SIZE=32, + EVEN_K=even_k, + CACHE_MODIFIER=".ca", + num_warps=8, + num_stages=1, # 256×128×128 needs 48KB/stage; 2 stages=96KB > 64KB LDS + waves_per_eu=0, + matrix_instr_nonkdim=16, + kpack=1, + ) + return out + + +# ── Blockwise FP8 Variable-K Backward Public API ── + + + + +@triton.jit() +def _grouped_blockwise_fp8_variable_k_gemm_kernel( + # C[g] = LHS_g^T @ RHS_g * block_scales + LHS, # [M_padded_total, OUT_M] FP8 + RHS, # [M_padded_total, OUT_N] FP8 + C, # [G, OUT_M, OUT_N] + LHS_scales_ptr, # [ceil(M_padded/128), OUT_M] float32 + RHS_scales_ptr, # [ceil(M_padded/128), OUT_N] float32 + group_offs_ptr, # [G+1] int64 (padded segment offsets, each aligned to 128) + G, # number of groups + OUT_M, + OUT_N, + # Strides + stride_lhs_m, + stride_rhs_m, + stride_cg, + stride_cm, + stride_cn, + # LHS_scales strides + stride_ls_0, + stride_ls_1, + # RHS_scales strides + stride_rs_0, + stride_rs_1, + # Constexpr strides + stride_lhs_n: tl.constexpr, + stride_rhs_n: tl.constexpr, + # Tile config + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_SMS: tl.constexpr, + NUM_XCDS: tl.constexpr, + CHUNK_SIZE: tl.constexpr, + CACHE_MODIFIER: tl.constexpr, +): + """Persistent grouped block-wise FP8 variable-K GEMM kernel (backward, CPU-sync-free). + + All groups share the same output dims (OUT_M × OUT_N), only the inner product + dimension M_g varies per group. 1D+1D scale pattern for TN/CRR layout. + + NOTE: Data is segment-padded to BLOCK_SIZE_K (128) boundaries by + quant_fp8_blockwise_segment_m_impl, so M_g is always a multiple of + BLOCK_SIZE_K. No masking is needed in the K-loop. + """ + pid = tl.program_id(0) + if NUM_XCDS != 1: + pid = remap_xcd_chunked(pid, NUM_SMS, NUM_XCDS, CHUNK_SIZE) + + tiles_m = tl.cdiv(OUT_M, BLOCK_SIZE_M) + tiles_n = tl.cdiv(OUT_N, BLOCK_SIZE_N) + tiles_per_group = tiles_m * tiles_n + total_tiles = G * tiles_per_group + + tl.assume(stride_lhs_m > 0) + tl.assume(stride_lhs_n > 0) + tl.assume(stride_rhs_m > 0) + tl.assume(stride_rhs_n > 0) + tl.assume(stride_cm > 0) + tl.assume(stride_cn > 0) + + acc_dtype = tl.float32 + + for global_tile in range(pid, total_tiles, NUM_SMS): + # ── Map to (group, local_tile) ── + group_idx = global_tile // tiles_per_group + local_tile = global_tile - group_idx * tiles_per_group + + # ── Swizzle local tile → (pid_m, pid_n) ── + num_pid_in_group = GROUP_SIZE_M * tiles_n + swizzle_group = local_tile // num_pid_in_group + first_pid_m = swizzle_group * GROUP_SIZE_M + group_size_m = min(tiles_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((local_tile % num_pid_in_group) % group_size_m) + pid_n = (local_tile % num_pid_in_group) // group_size_m + tl.assume(pid_m >= 0) + tl.assume(pid_n >= 0) + + # ── Group boundaries ── + m_start = tl.load(group_offs_ptr + group_idx) # int64 + M_g = (tl.load(group_offs_ptr + group_idx + 1) - m_start).to(tl.int32) + + # ── Output indices ── + rm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % OUT_M + rn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % OUT_N + rk = tl.arange(0, BLOCK_SIZE_K) + rn = tl.max_contiguous(tl.multiple_of(rn, BLOCK_SIZE_N), BLOCK_SIZE_N) + + # ── Base pointers ── + LHS_BASE = LHS + m_start * stride_lhs_m + rm[:, None] * stride_lhs_n + rk[None, :] * stride_lhs_m + RHS_BASE = RHS + m_start * stride_rhs_m + rk[:, None] * stride_rhs_m + rn[None, :] * stride_rhs_n + + scale_row_start = m_start // BLOCK_SIZE_K + + # ── K-loop over M_g with block-wise 1D+1D scaling ── + # M_g is always a multiple of BLOCK_SIZE_K (data padded), so no masking needed. + loop_k = M_g // BLOCK_SIZE_K + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype) + + for k in range(loop_k): + if stride_lhs_n == 1: + a = tl.load( + tl.multiple_of(LHS_BASE, (16, 1)), + cache_modifier=CACHE_MODIFIER, + ) + else: + a = tl.load( + tl.multiple_of(LHS_BASE, (1, 16)), + cache_modifier=CACHE_MODIFIER, + ) + + if stride_rhs_n == 1: + b = tl.load( + tl.multiple_of(RHS_BASE, (1, 16)), + cache_modifier=CACHE_MODIFIER, + ) + else: + b = tl.load( + tl.multiple_of(RHS_BASE, (16, 1)), + cache_modifier=CACHE_MODIFIER, + ) + + partial = tl.dot(a, b) + + # 1D+1D block-wise scales + scale_row = scale_row_start + k + a_s = tl.load(LHS_scales_ptr + scale_row * stride_ls_0 + rm * stride_ls_1) + b_s = tl.load(RHS_scales_ptr + scale_row * stride_rs_0 + rn * stride_rs_1) + acc += partial * a_s[:, None] * b_s[None, :] + + LHS_BASE += BLOCK_SIZE_K * stride_lhs_m + RHS_BASE += BLOCK_SIZE_K * stride_rhs_m + + # ── Store output ── + c = acc.to(C.type.element_ty) + rm_s = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + rn_s = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + rn_s = tl.max_contiguous(tl.multiple_of(rn_s % OUT_N, BLOCK_SIZE_N), BLOCK_SIZE_N) + c_mask = (rm_s[:, None] < OUT_M) & (rn_s[None, :] < OUT_N) + C_ = C + group_idx.to(tl.int64) * stride_cg + rm_s[:, None] * stride_cm + rn_s[None, :] * stride_cn + tl.store(C_, c, c_mask) + + +# ── Blockwise FP8 Forward Public API ── + + + +def grouped_gemm_fp8_blockwise_variable_k_triton_kernel( + lhs: torch.Tensor, + rhs: torch.Tensor, + lhs_scales: torch.Tensor, + rhs_scales: torch.Tensor, + group_offs: torch.Tensor, + out_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Variable-K grouped block-wise FP8 GEMM (backward, 1D+1D scaling) using Triton. + + Computes: C[g] = lhs[offs[g]:offs[g+1]]^T @ rhs[offs[g]:offs[g+1]] + with 1D+1D block-wise scaling applied in the K-loop. + + Output: [G, OUT_M, OUT_N]. + + Args: + lhs: [M_padded_total, OUT_M] FP8 (segment-padded, each segment aligned to 128). + rhs: [M_padded_total, OUT_N] FP8. + lhs_scales: [ceil(M_padded/128), OUT_M] float32. + rhs_scales: [ceil(M_padded/128), OUT_N] float32. + group_offs: [G+1] int64 padded segment offsets. + out_dtype: Output dtype (default bfloat16). + + Returns: + [G, OUT_M, OUT_N] output. + """ + if is_cdna4(): + set_triton_knobs_gfx950() + else: + _set_amd_knobs(enable=False) + + assert lhs.ndim == 2 and rhs.ndim == 2 + assert lhs.shape[0] == rhs.shape[0] + OUT_M = lhs.shape[1] + OUT_N = rhs.shape[1] + G = group_offs.shape[0] - 1 + + out = torch.empty((G, OUT_M, OUT_N), device=lhs.device, dtype=out_dtype) + num_sms = get_num_cus() + + # Use 128x128 tiles to reduce register pressure from double-accumulator + # (partial + acc both need full tile VGPRs for blockwise scale application). + # With 256x256 tiles + 8 warps, 2 accumulator sets need ~290 VGPRs/wave + # which exceeds the 256 limit at 2 waves/SIMD, causing spilling. + # 128x128 with 4 warps keeps VGPRs at ~170/wave, fitting 2 waves/SIMD. + _grouped_blockwise_fp8_variable_k_gemm_kernel[(num_sms,)]( + lhs, + rhs, + out, + lhs_scales, + rhs_scales, + group_offs, + G, + OUT_M, + OUT_N, + lhs.stride(0), + rhs.stride(0), + out.stride(0), + out.stride(1), + out.stride(2), + lhs_scales.stride(0), + lhs_scales.stride(1), + rhs_scales.stride(0), + rhs_scales.stride(1), + stride_lhs_n=lhs.stride(1), + stride_rhs_n=rhs.stride(1), + BLOCK_SIZE_M=128, + BLOCK_SIZE_N=128, + BLOCK_SIZE_K=128, + GROUP_SIZE_M=4, + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + CHUNK_SIZE=32, + CACHE_MODIFIER=".ca", + num_warps=4, + num_stages=2, + waves_per_eu=0, + matrix_instr_nonkdim=16, + kpack=1, + ) + return out diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py b/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py new file mode 100644 index 0000000000..73fa02e264 --- /dev/null +++ b/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py @@ -0,0 +1,328 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +# +# Blockwise FP8 quantization Triton kernels (1x128 activation, 128x128 weight), +# adapted from AMD Primus-Turbo (primus_turbo/triton/quantization/quant_blockwise.py +# and primus_turbo/pytorch/kernels/quantization/quantization_impl.py). + +import torch +import triton +import triton.language as tl + +__all__ = [ + "quantize_fp8_blockwise", + "quantize_fp8_blockwise_dual", + "quantize_fp8_blockwise_weight", + "quantize_fp8_blockwise_segment_m", +] + + +@triton.jit +def _floor_to_pow2(scale): + scale_bits = scale.to(tl.uint32, bitcast=True) & 0xFF800000 + return scale_bits.to(tl.float32, bitcast=True) + + +@triton.jit +def compute_scale_and_quant(x_tile, x_tile_abs, axis, FP8_MAX, ROUND_POW2: tl.constexpr): + x_tile_max = tl.max(x_tile_abs, axis=axis, keep_dims=True) + x_tile_max = tl.maximum(x_tile_max, 1e-4) + x_scales_tile = FP8_MAX / x_tile_max + if ROUND_POW2: + x_scales_tile = _floor_to_pow2(x_scales_tile) + x_fp8_tile = x_tile * x_scales_tile + x_fp8_tile = tl.clamp(x_fp8_tile, min=-FP8_MAX, max=FP8_MAX) + return x_fp8_tile, x_scales_tile + + +@triton.jit +def quant_fp8_blockwise_kernel( + x_ptr, + x_fp8_ptr, + x_scales_ptr, + M, + N, + BLOCK_SIZE: tl.constexpr, + FP8_MAX: tl.constexpr, + AXIS: tl.constexpr, + ROUND_POW2: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + offs_m = tl.cast(pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) + offs_n = tl.cast(pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) + mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + + x_ptrs = x_ptr + offs_m[:, None] * N + offs_n[None, :] + x_tile = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) + x_tile_abs = tl.abs(x_tile) + + x_fp8_tile, x_scales_tile = compute_scale_and_quant(x_tile, x_tile_abs, AXIS, FP8_MAX, ROUND_POW2) + + x_fp8_ptrs = x_fp8_ptr + offs_m[:, None] * N + offs_n[None, :] + tl.store(x_fp8_ptrs, x_fp8_tile.to(x_fp8_ptr.dtype.element_ty), mask=mask) + + if AXIS == 1: + scale_offs = offs_m * tl.cdiv(N, BLOCK_SIZE) + pid_n + scale_mask = offs_m < M + else: + scale_offs = pid_m * N + offs_n + scale_mask = offs_n < N + x_scales_tile_inv = tl.reshape(1.0 / x_scales_tile, BLOCK_SIZE) + tl.store(x_scales_ptr + scale_offs, x_scales_tile_inv, mask=scale_mask) + + +@triton.jit +def quant_fp8_blockwise_dual_kernel( + x_ptr, + x_fp8_row_ptr, + x_scales_row_ptr, + x_fp8_col_ptr, + x_scales_col_ptr, + M, + N, + BLOCK_SIZE: tl.constexpr, + FP8_MAX: tl.constexpr, + ROUND_POW2: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + offs_m = tl.cast(pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) + offs_n = tl.cast(pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) + mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + + x_ptrs = x_ptr + offs_m[:, None] * N + offs_n[None, :] + x_tile = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) + x_tile_abs = tl.abs(x_tile) + + x_fp8_row_tile, x_scales_row_tile = compute_scale_and_quant(x_tile, x_tile_abs, 1, FP8_MAX, ROUND_POW2) + x_fp8_col_tile, x_scales_col_tile = compute_scale_and_quant(x_tile, x_tile_abs, 0, FP8_MAX, ROUND_POW2) + + x_fp8_row_ptrs = x_fp8_row_ptr + offs_m[:, None] * N + offs_n[None, :] + tl.store(x_fp8_row_ptrs, x_fp8_row_tile.to(x_fp8_row_ptr.dtype.element_ty), mask=mask) + + x_fp8_col_ptrs = x_fp8_col_ptr + offs_m[:, None] * N + offs_n[None, :] + tl.store(x_fp8_col_ptrs, x_fp8_col_tile.to(x_fp8_col_ptr.dtype.element_ty), mask=mask) + + row_scale_offs = offs_m * tl.cdiv(N, BLOCK_SIZE) + pid_n + row_scale_mask = offs_m < M + x_scales_row_tile_inv = tl.reshape(1.0 / x_scales_row_tile, BLOCK_SIZE) + tl.store(x_scales_row_ptr + row_scale_offs, x_scales_row_tile_inv, mask=row_scale_mask) + + col_scale_offs = pid_m * N + offs_n + col_scale_mask = offs_n < N + x_scales_col_tile_inv = tl.reshape(1.0 / x_scales_col_tile, BLOCK_SIZE) + tl.store(x_scales_col_ptr + col_scale_offs, x_scales_col_tile_inv, mask=col_scale_mask) + + +@triton.jit +def quant_fp8_blockwise_for_weight_kernel( + w_ptr, + w_fp8_ptr, + w_scales_ptr, + M, + N, + BLOCK_SIZE: tl.constexpr, + FP8_MAX: tl.constexpr, + ROUND_POW2: tl.constexpr, +): + bid = tl.program_id(axis=0) + pid_m = tl.program_id(axis=1) + pid_n = tl.program_id(axis=2) + + batch_offset_w = bid * M * N + batch_offset_scales = bid * tl.cdiv(M, BLOCK_SIZE) * tl.cdiv(N, BLOCK_SIZE) + + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + + w_ptrs = w_ptr + batch_offset_w + offs_m[:, None] * N + offs_n[None, :] + w_tile = tl.load(w_ptrs, mask=mask, other=0.0).to(tl.float32) + + w_tile_abs = tl.abs(w_tile) + w_tile_max = tl.max(w_tile_abs) + w_tile_max = tl.maximum(w_tile_max, 1e-4) + w_scales = FP8_MAX / w_tile_max + if ROUND_POW2: + w_scales = _floor_to_pow2(w_scales) + w_fp8_tile = w_tile * w_scales + w_fp8_tile = tl.clamp(w_fp8_tile, min=-FP8_MAX, max=FP8_MAX) + + w_fp8_ptrs = w_fp8_ptr + batch_offset_w + offs_m[:, None] * N + offs_n[None, :] + tl.store(w_fp8_ptrs, w_fp8_tile.to(w_fp8_ptr.dtype.element_ty), mask=mask) + scale_offs = batch_offset_scales + pid_m * tl.cdiv(N, BLOCK_SIZE) + pid_n + w_scales_inv = 1.0 / w_scales + tl.store(w_scales_ptr + scale_offs, w_scales_inv) + + +@triton.jit +def quant_fp8_blockwise_segment_m_kernel( + x_ptr, # Input tensor [M_in, N] + x_fp8_ptr, # Output tensor [M_out, N] (padded) + x_scales_ptr, # Output scales [M_out // BLOCK_SIZE, N] + group_offs_ptr, # Original group offsets [B+1] + padded_group_offs_ptr, # Padded group offsets [B+1] + N, + num_groups, + BLOCK_SIZE: tl.constexpr, + FP8_MAX: tl.constexpr, + ROUND_POW2: tl.constexpr, +): + """Colwise (axis=0) blockwise quantize with per-segment M padding. + + Reads from the original (unpadded) tensor and writes to a segment-aligned + padded output, so each MoE group starts on a BLOCK_SIZE boundary. Used to + produce the column-wise FP8 operand for the grouped backward (wgrad), always + quantizing from the original high-precision input (no double-quantization). + """ + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + + M_padded = tl.load(padded_group_offs_ptr + num_groups) + block_start = pid_m * BLOCK_SIZE + if block_start >= M_padded: + return + + group_id = 0 + for g in range(num_groups): + padded_start = tl.load(padded_group_offs_ptr + g) + padded_end = tl.load(padded_group_offs_ptr + g + 1) + if block_start >= padded_start and block_start < padded_end: + group_id = g + + orig_group_start = tl.load(group_offs_ptr + group_id) + orig_group_end = tl.load(group_offs_ptr + group_id + 1) + padded_group_start = tl.load(padded_group_offs_ptr + group_id) + + offs_m_out = tl.cast(pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) + offs_n = tl.cast(pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) + offs_m_in = orig_group_start + (offs_m_out - padded_group_start) + + mask = ( + (offs_m_in[:, None] >= orig_group_start) + & (offs_m_in[:, None] < orig_group_end) + & (offs_n[None, :] < N) + ) + + x_ptrs = x_ptr + offs_m_in[:, None] * N + offs_n[None, :] + x_tile = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) + x_tile_abs = tl.abs(x_tile) + + x_fp8_tile, x_scales_tile = compute_scale_and_quant(x_tile, x_tile_abs, 0, FP8_MAX, ROUND_POW2) + + x_fp8_ptrs = x_fp8_ptr + offs_m_out[:, None] * N + offs_n[None, :] + out_mask = (offs_m_out[:, None] < M_padded) & (offs_n[None, :] < N) + tl.store(x_fp8_ptrs, x_fp8_tile.to(x_fp8_ptr.dtype.element_ty), mask=out_mask) + + scale_offs = pid_m * N + offs_n + scale_mask = (pid_m < tl.cdiv(M_padded, BLOCK_SIZE)) & (offs_n < N) + x_scales_tile_inv = tl.reshape(1.0 / x_scales_tile, BLOCK_SIZE) + tl.store(x_scales_ptr + scale_offs, x_scales_tile_inv, mask=scale_mask) + + +# ----------------------------------------------------------------------------- +# Host launchers +# +# Plain functions, not ``torch.library.custom_op``. A custom_op wrapper may be +# needed if torch.compile hits the inductor ``identify_mutated_tensors`` bug +# seen on gfx942 + Triton 3.7. +# ----------------------------------------------------------------------------- + + +def quantize_fp8_blockwise_dual(x: torch.Tensor, dtype: torch.dtype, block_size: int = 128, pow2: bool = False): + """Blockwise-quantize a 2D tensor in BOTH row (1xB) and column (Bx1) modes in one pass. + + Returns (x_fp8_row, x_scales_row, x_fp8_col, x_scales_col); scales hold the + dequant scale (amax/FP8_MAX) as fp32. + """ + assert x.is_contiguous() and x.dim() == 2, "Input must be 2D and contiguous" + M, N = x.shape + fp8_max = torch.finfo(dtype).max + + x_fp8_row = torch.empty((M, N), dtype=dtype, device=x.device) + x_scales_row = torch.empty((M, triton.cdiv(N, block_size)), dtype=torch.float32, device=x.device) + x_fp8_col = torch.empty((M, N), dtype=dtype, device=x.device) + x_scales_col = torch.empty((triton.cdiv(M, block_size), N), dtype=torch.float32, device=x.device) + + grid = (triton.cdiv(M, block_size), triton.cdiv(N, block_size)) + quant_fp8_blockwise_dual_kernel[grid]( + x, x_fp8_row, x_scales_row, x_fp8_col, x_scales_col, M, N, + BLOCK_SIZE=block_size, FP8_MAX=fp8_max, + ROUND_POW2=pow2, + ) + return x_fp8_row, x_scales_row, x_fp8_col, x_scales_col + + +def quantize_fp8_blockwise(x: torch.Tensor, dtype: torch.dtype, axis: int, block_size: int = 128, pow2: bool = False): + """Single-direction blockwise quantize. axis=1 -> rowwise (1xB), axis=0 -> colwise (Bx1).""" + assert x.is_contiguous() and x.dim() == 2, "Input must be 2D and contiguous" + M, N = x.shape + fp8_max = torch.finfo(dtype).max + x_fp8 = torch.empty((M, N), dtype=dtype, device=x.device) + if axis == 1: + scales = torch.empty((M, triton.cdiv(N, block_size)), dtype=torch.float32, device=x.device) + else: + scales = torch.empty((triton.cdiv(M, block_size), N), dtype=torch.float32, device=x.device) + + grid = (triton.cdiv(M, block_size), triton.cdiv(N, block_size)) + quant_fp8_blockwise_kernel[grid]( + x, x_fp8, scales, M, N, BLOCK_SIZE=block_size, FP8_MAX=fp8_max, AXIS=axis, + ROUND_POW2=pow2, + ) + return x_fp8, scales + + +def quantize_fp8_blockwise_weight(w: torch.Tensor, dtype: torch.dtype, block_size: int = 128, pow2: bool = False): + """128x128 weight blockwise quantize. w is [B, M, N] (or [M, N], promoted to B=1).""" + squeeze = False + if w.dim() == 2: + w = w.unsqueeze(0) + squeeze = True + assert w.is_contiguous() and w.dim() == 3, "Weight must be 3D [B,M,N] and contiguous" + B, M, N = w.shape + fp8_max = torch.finfo(dtype).max + + w_fp8 = torch.empty((B, M, N), dtype=dtype, device=w.device) + w_scales = torch.empty( + (B, triton.cdiv(M, block_size), triton.cdiv(N, block_size)), + dtype=torch.float32, device=w.device, + ) + grid = (B, triton.cdiv(M, block_size), triton.cdiv(N, block_size)) + quant_fp8_blockwise_for_weight_kernel[grid]( + w, w_fp8, w_scales, M, N, BLOCK_SIZE=block_size, FP8_MAX=fp8_max, + ROUND_POW2=pow2, + ) + if squeeze: + return w_fp8.squeeze(0), w_scales.squeeze(0) + return w_fp8, w_scales + + +def quantize_fp8_blockwise_segment_m(x, dtype, block_size, group_lens, group_offs, pow2: bool = False): + """Colwise blockwise quantize with per-segment (MoE group) M padding. + + Returns (x_fp8 [M_pad, N], x_scales, var_k_group_lens [B], var_k_group_offs [B+1]). + Each group is padded up to a multiple of block_size. Allocates with an upper + bound (M + B*block_size) to avoid a device->host sync (graph-capture safe). + """ + assert x.is_contiguous() and x.dim() == 2, "Input must be 2D and contiguous" + M, N = x.shape + num_groups = group_lens.size(0) + fp8_max = torch.finfo(dtype).max + + var_k_group_lens = ((group_lens + block_size - 1) // block_size) * block_size + var_k_group_offs = torch.zeros(num_groups + 1, dtype=torch.int64, device=x.device) + var_k_group_offs[1:] = torch.cumsum(var_k_group_lens, dim=0) + + m_padded_max = M + num_groups * block_size + x_fp8 = torch.zeros((m_padded_max, N), dtype=dtype, device=x.device) + x_scales = torch.zeros((triton.cdiv(m_padded_max, block_size), N), dtype=torch.float32, device=x.device) + + grid = (triton.cdiv(m_padded_max, block_size), triton.cdiv(N, block_size)) + quant_fp8_blockwise_segment_m_kernel[grid]( + x, x_fp8, x_scales, group_offs, var_k_group_offs, N, num_groups, + BLOCK_SIZE=block_size, FP8_MAX=fp8_max, + ROUND_POW2=pow2, + ) + return x_fp8, x_scales, var_k_group_lens, var_k_group_offs diff --git a/transformer_engine/pytorch/triton_kernels/common.py b/transformer_engine/pytorch/triton_kernels/common.py index 212c8cd695..2487ee3fe1 100644 --- a/transformer_engine/pytorch/triton_kernels/common.py +++ b/transformer_engine/pytorch/triton_kernels/common.py @@ -14,6 +14,9 @@ def get_arch(): def is_cdna3(): return get_arch() == "gfx942" +def is_cdna4(): + return get_arch() == "gfx950" + get_torch_e4m3_type = lambda: torch.float8_e4m3fn if not is_cdna3() else torch.float8_e4m3fnuz get_torch_e5m2_type = lambda: torch.float8_e5m2 if not is_cdna3() else torch.float8_e5m2fnuz From 3a4951db5dc63f6517de9583b7989ff574ca1ccd Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Tue, 25 Aug 2026 18:00:20 +0000 Subject: [PATCH 02/13] Refactor ROCm blockwise FP8 grouped GEMM gating in GroupedLinear - Add `_is_blockwise_fp8_grouped_gemm_supported()` to centralize the feature-gate checks (HIP, `NVTE_USE_BLOCKWISE_GMM_TRITON=1`, `Float8BlockScaling` layout, and unsupported orchestration options). - Simplify `_forward_blockwise_fp8` by removing inline validation and accepting a pre-uploaded `m_splits_tensor`, avoiding a blocking host-to-device copy of the split sizes. - Switch the environment gate from `NVTE_USE_BLOCKWISE_FP8_GROUPED_GEMM` to `NVTE_USE_BLOCKWISE_GMM_TRITON`. --- .../pytorch/module/grouped_linear.py | 184 ++++++++---------- 1 file changed, 81 insertions(+), 103 deletions(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index e9b7d84eec..bdaceb9f31 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -399,38 +399,79 @@ def _forward_grouped_tensor( return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces + @staticmethod + def _is_blockwise_fp8_grouped_gemm_supported( + *, + fp8, + recipe, + use_bias, + backward_override, + fuse_wgrad_accumulation, + cpu_offloading, + save_original_input, + wgrad_store, + debug, + unpad_output, + actual_m_splits, + m_splits, + ) -> bool: + """Return whether the ROCm Triton blockwise FP8 grouped GEMM path can run this call. + + Requires ``NVTE_USE_BLOCKWISE_GMM_TRITON=1``, ``Float8BlockScaling``, and the + hardcoded blockwise layout (activation 1x128 along K, weights 128x128, grad 1x128). + Unsupported orchestration features fall back to the default path. + """ + if not IS_HIP_EXTENSION: + return False + if os.getenv("NVTE_USE_BLOCKWISE_GMM_TRITON", "0") != "1": + return False + if not fp8 or recipe is None or not recipe.float8_block_scaling(): + return False + if ( + recipe.x_block_scaling_dim != 1 + or recipe.w_block_scaling_dim != 2 + or recipe.grad_block_scaling_dim != 1 + ): + return False + if use_bias: + return False + if backward_override is not None: + return False + if fuse_wgrad_accumulation: + return False + if cpu_offloading: + return False + if save_original_input: + return False + if wgrad_store is not None and wgrad_store.delay_wgrad_compute(): + return False + if debug: + return False + if unpad_output or ( + actual_m_splits is not None and list(actual_m_splits) != list(m_splits) + ): + return False + return True + @staticmethod def _forward_blockwise_fp8( ctx, *, inp, m_splits, + m_splits_tensor, weights, - biases, - use_bias, - fp8, - recipe, - wgrad_store, weight_quantizers, - fuse_wgrad_accumulation, - cpu_offloading, activation_dtype, is_grad_enabled, - save_original_input, - debug, - actual_m_splits, - unpad_output, - backward_override, ): - """DeepSeek-style blockwise FP8 grouped GEMM forward (ROCm Triton). + """Blockwise FP8 grouped GEMM forward (ROCm Triton). Selected from :meth:`forward` under the ``Float8BlockScaling`` recipe when - ``NVTE_USE_BLOCKWISE_FP8_GROUPED_GEMM=1``. Activations are quantized rowwise - (1x128 along K), weights 128x128; the backward wgrad (see + :meth:`_is_blockwise_fp8_grouped_gemm_supported` is true. Activations are + quantized rowwise (1x128 along K), weights 128x128; the backward wgrad (see :meth:`_backward_blockwise_fp8`) uses a segment-padded columnwise operand quantized from the original high-precision input (no double-quantization). - Recipe configurations and orchestration features it does not yet support are - rejected instead of silently ignored. """ from ..triton_kernels.common import te_dtype_to_torch_dtype from ..triton_kernels.blockwise_quantize import ( @@ -444,62 +485,6 @@ def _forward_blockwise_fp8( num_gemms = len(m_splits) - # fp8 + block-scaling recipe are internal invariants (asserts, guaranteed by the - # caller's gate); the rest are user-reachable (raise). - assert fp8, "blockwise grouped FP8 path requires fp8=True" - assert recipe.float8_block_scaling(), "blockwise grouped FP8 path requires Float8BlockScaling" - - # The Triton kernels hardcode the DeepSeek layout (x rowwise 1x128, w 128x128, - # grad rowwise 1x128). Reject recipe block-scaling dims that don't match rather - # than silently producing wrong numerics. - if ( - recipe.x_block_scaling_dim != 1 - or recipe.w_block_scaling_dim != 2 - or recipe.grad_block_scaling_dim != 1 - ): - raise NotImplementedError( - "the blockwise grouped FP8 path only supports x_block_scaling_dim=1, " - "w_block_scaling_dim=2, grad_block_scaling_dim=1 (got " - f"{recipe.x_block_scaling_dim}, {recipe.w_block_scaling_dim}, " - f"{recipe.grad_block_scaling_dim})" - ) - if use_bias: - raise NotImplementedError("bias is not supported in the blockwise grouped FP8 path yet") - if backward_override is not None: - raise NotImplementedError( - "backward_override is not supported in the blockwise grouped FP8 path yet" - ) - if fuse_wgrad_accumulation: - raise NotImplementedError( - "fuse_wgrad_accumulation (gradient_accumulation_fusion) is not yet supported in " - "the ROCm blockwise grouped FP8 path. Pass --no-gradient-accumulation-fusion to the " - "training script: wgrad is then returned as a plain gradient and Megatron's DDP " - "post-hook accumulates it into the fp32 main_grad (numerically equivalent for bring-up)." - ) - if cpu_offloading: - raise NotImplementedError( - "cpu_offloading is not supported in the blockwise grouped FP8 path yet" - ) - if save_original_input: - raise NotImplementedError( - "save_original_input is not supported in the blockwise grouped FP8 path yet" - ) - if wgrad_store is not None and wgrad_store.delay_wgrad_compute(): - raise NotImplementedError( - "delayed wgrad is not supported in the blockwise grouped FP8 path yet" - ) - if debug: - raise NotImplementedError( - "debug quantization is not supported in the blockwise grouped FP8 path yet" - ) - if unpad_output or ( - actual_m_splits is not None and list(actual_m_splits) != list(m_splits) - ): - raise NotImplementedError( - "the ROCm fused-pad / unpad_output path is not supported in the blockwise grouped " - "FP8 path yet" - ) - # Resolve the FP8 torch dtype via the quantizer, normalizing through ``tex.DType`` # so ``te_dtype_to_torch_dtype`` picks the arch-correct variant (fnuz on CDNA3). if weight_quantizers and weight_quantizers[0] is not None: @@ -513,9 +498,10 @@ def _forward_blockwise_fp8( inp_shape = inp.shape a = inp.reshape(-1, in_features).to(activation_dtype).contiguous() - # Group offsets built on-device from the split tensor (graph-capture safe, - # no device->host sync). - group_lens = m_splits.to(device=device, dtype=torch.int64) + # Prefer the caller-provided GPU splits tensor (avoids a blocking H2D of + # pageable m_splits). Fall back to m_splits when it is not passed. + split_src = m_splits_tensor if m_splits_tensor is not None else m_splits + group_lens = split_src.to(device=device, dtype=torch.int64) group_offs = torch.zeros(num_gemms + 1, dtype=torch.int64, device=device) group_offs[1:] = torch.cumsum(group_lens, 0) @@ -585,14 +571,12 @@ def _backward_blockwise_fp8(ctx, grad_output): ) wgrad_list = [dW[g].contiguous() for g in range(ctx.num_gemms)] - grad_biases = [None] * ctx.num_gemms # bias rejected in forward + grad_biases = [None] * ctx.num_gemms # bias not supported on this path - # Grads match forward inputs: (inp, m_splits, dispatched_probs, non_tensor_args, - # *weights, *biases). + # Grads match forward inputs: (inp, m_splits, non_tensor_args, *weights, *biases). return ( dgrad.view(ctx.inp_shape) if dgrad is not None else None, None, # m_splits - None, # dispatched_probs None, # non_tensor_args *wgrad_list, *grad_biases, @@ -653,37 +637,31 @@ def forward( device = inp.device weight_requires_grad = weights[0].requires_grad - # DeepSeek-style blockwise FP8 grouped GEMM (ROCm Triton) opt-in. Only engages under - # the Float8BlockScaling recipe (whose quantizers already carry blockwise semantics); - # it runs its own quantization + grouped GEMM and returns early, bypassing the default - # quantizer setup below. - use_blockwise_fp8 = ( - IS_HIP_EXTENSION - and fp8 - and os.getenv("NVTE_USE_BLOCKWISE_FP8_GROUPED_GEMM", "0") == "1" - and FP8GlobalStateManager.get_fp8_recipe().float8_block_scaling() - ) - if use_blockwise_fp8: + # Blockwise FP8 grouped GEMM (ROCm Triton) opt-in. Runs its own quantization + + # grouped GEMM and returns early, bypassing the default quantizer setup below. + if _GroupedLinear._is_blockwise_fp8_grouped_gemm_supported( + fp8=fp8, + recipe=FP8GlobalStateManager.get_fp8_recipe() if fp8 else None, + use_bias=use_bias, + backward_override=backward_override, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + cpu_offloading=cpu_offloading, + save_original_input=save_original_input, + wgrad_store=wgrad_store, + debug=debug, + unpad_output=unpad_output, + actual_m_splits=actual_m_splits, + m_splits=m_splits, + ): return _GroupedLinear._forward_blockwise_fp8( ctx, inp=inp, m_splits=m_splits, + m_splits_tensor=m_splits_tensor, weights=weights, - biases=biases, - use_bias=use_bias, - fp8=fp8, - recipe=FP8GlobalStateManager.get_fp8_recipe(), - wgrad_store=wgrad_store, weight_quantizers=weight_quantizers, - fuse_wgrad_accumulation=fuse_wgrad_accumulation, - cpu_offloading=cpu_offloading, activation_dtype=activation_dtype, is_grad_enabled=is_grad_enabled, - save_original_input=save_original_input, - debug=debug, - actual_m_splits=actual_m_splits, - unpad_output=unpad_output, - backward_override=backward_override, ) # Configure quantizers From bfea740a643bb61a3cd44ee476800d6c627f2a14 Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Tue, 25 Aug 2026 19:04:43 +0000 Subject: [PATCH 03/13] Enable fused wgrad accumulation for ROCm blockwise FP8 grouped GEMM - Extend the Triton variable-K grouped GEMM kernel to support in-place accumulation and an optional output tensor. - Add fused wgrad handling to GroupedLinear's blockwise FP8 path, including packed main-grad views and first-microbatch accumulation logic. - Remove ROCm test skips for FP8 block scaling in grouped linear tests and switch block-scaling cases to the blockwise Triton backend. - Add unit tests for the blockwise FP8 quantization and grouped GEMM Triton kernels and include them in the PyTorch CI script. --- ci/pytorch.sh | 1 + tests/pytorch/test_grouped_linear.py | 26 +- .../triton_kernels/test_blockwise_fp8.py | 297 ++++++++++++++++++ .../pytorch/module/grouped_linear.py | 184 ++++++++--- .../blockwise_fp8_grouped_gemm.py | 30 +- 5 files changed, 469 insertions(+), 69 deletions(-) create mode 100644 tests/pytorch/triton_kernels/test_blockwise_fp8.py diff --git a/ci/pytorch.sh b/ci/pytorch.sh index b5241f3d75..bae141110c 100755 --- a/ci/pytorch.sh +++ b/ci/pytorch.sh @@ -98,6 +98,7 @@ run_test_config(){ run_default_fa 1 triton_kernels/test_cast_mxfp8.py run_default_fa 1 triton_kernels/test_cast_mxfp4.py run_default_fa 1 triton_kernels/test_grouped_gemm.py + run_default_fa 1 triton_kernels/test_blockwise_fp8.py run_default_fa 1 triton_kernels/test_utils.py NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa 1 triton_kernels/test_norms.py NVTE_ROCM_ENABLE_MXFP8=1 NVTE_TEST_TRITON_AUTOTUNE=1 run_default_fa_lbl "autotune" 3 triton_kernels/test_norms.py diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 51c16769a3..c440ea7dfc 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -300,8 +300,6 @@ def test_grouped_linear_accuracy( pytest.skip("Triton grouped gemm is only supported on HIP.") if IS_HIP_EXTENSION and dtype not in (torch.float32,) and fuse_wgrad_accumulation and not fp8: pytest.skip(f"ROCm does not support fused wgrad accumulation for {dtype}.") - if IS_HIP_EXTENSION and recipe is not None and recipe.float8_block_scaling(): - pytest.skip("ROCm grouped GEMM does not yet support FP8 block scaling.") if fp8 and fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: pytest.skip("FP8 parameters are not supported in debug mode.") if NVTE_TEST_NVINSPECT_ENABLED and delay_wgrad_compute: @@ -321,7 +319,10 @@ def test_grouped_linear_accuracy( ) if use_triton: - os.environ["NVTE_USE_GROUPED_GEMM_TRITON"] = "1" + if recipe is not None and recipe.float8_block_scaling(): + os.environ["NVTE_USE_BLOCKWISE_GMM_TRITON"] = "1" + else: + os.environ["NVTE_USE_GROUPED_GEMM_TRITON"] = "1" with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): grouped_linear = GroupedLinear( @@ -385,8 +386,10 @@ def test_grouped_linear_accuracy( delay_wgrad_compute, ) + use_blockwise_triton = os.getenv("NVTE_USE_BLOCKWISE_GMM_TRITON", "0") == "1" if use_triton: os.environ.pop("NVTE_USE_GROUPED_GEMM_TRITON", None) + os.environ.pop("NVTE_USE_BLOCKWISE_GMM_TRITON", None) atol, rtol = 0, 0 if use_cutlass: @@ -398,8 +401,19 @@ def test_grouped_linear_accuracy( if dtype == torch.float32: atol = 2.6e-6 rtol = 5e-2 + if use_blockwise_triton: + # tests/pytorch/triton_kernels/test_blockwise_fp8.py dequant checks. + atol, rtol = 0.25, 0.12 for o, o_ref in zip(outputs, outputs_ref): - torch.testing.assert_close(o, o_ref, rtol=rtol, atol=atol) + if use_blockwise_triton: + if o is None: + assert o_ref is None + continue + mag = max(float(o.detach().abs().max()), float(o_ref.detach().abs().max())) + tensor_atol = max(atol, 0.05 * mag) + torch.testing.assert_close(o, o_ref, rtol=rtol, atol=tensor_atol) + else: + torch.testing.assert_close(o, o_ref, rtol=rtol, atol=atol) @pytest.mark.skipif( @@ -754,8 +768,6 @@ def test_padding_grouped_linear_accuracy( ): if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: pytest.skip("FP8 parameters are not supported in debug mode.") - if IS_HIP_EXTENSION and recipe is not None and recipe.float8_block_scaling(): - pytest.skip("ROCm grouped GEMM does not yet support FP8 block scaling.") skip_unsupported_backward_override( "grouped_linear", recipe, getattr(recipe, "backward_override", None) ) @@ -836,8 +848,6 @@ def test_padding_grouped_linear_accuracy_save_original_input( pytest.skip("FP8 parameters are not supported in debug mode.") if fp8 and recipe.delayed(): pytest.skip("DelayedScaling recipe is not supported with save_original_input") - if IS_HIP_EXTENSION and recipe is not None and recipe.float8_block_scaling(): - pytest.skip("ROCm grouped GEMM does not yet support FP8 block scaling.") skip_unsupported_backward_override( "grouped_linear", recipe, getattr(recipe, "backward_override", None) ) diff --git a/tests/pytorch/triton_kernels/test_blockwise_fp8.py b/tests/pytorch/triton_kernels/test_blockwise_fp8.py new file mode 100644 index 0000000000..3791696197 --- /dev/null +++ b/tests/pytorch/triton_kernels/test_blockwise_fp8.py @@ -0,0 +1,297 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Unit tests for ROCm Triton blockwise FP8 quantization and grouped GEMM.""" + +import pytest +import torch +from torch.utils.cpp_extension import IS_HIP_EXTENSION + +from transformer_engine.pytorch.triton_kernels.common import get_torch_e4m3_type +from transformer_engine.pytorch.triton_kernels.blockwise_quantize import ( + quantize_fp8_blockwise, + quantize_fp8_blockwise_dual, + quantize_fp8_blockwise_weight, + quantize_fp8_blockwise_segment_m, +) +from transformer_engine.pytorch.triton_kernels.blockwise_fp8_grouped_gemm import ( + grouped_gemm_fp8_blockwise_triton_kernel, + grouped_gemm_fp8_blockwise_variable_k_triton_kernel, +) + +pytestmark = pytest.mark.skipif(not IS_HIP_EXTENSION, reason="ROCm Triton blockwise kernels only") + +BLOCK = 128 +FP8_DTYPE = get_torch_e4m3_type() +IN_DTYPES = [torch.bfloat16, torch.float16] + + +def _cdiv(n, d): + return (n + d - 1) // d + + +def _floor_to_pow2(scale: torch.Tensor) -> torch.Tensor: + # 0xFF800000 as signed int32: keep the exponent, zero the mantissa. + bits = scale.to(torch.float32).contiguous().view(torch.int32) + return (bits & -8388608).view(torch.float32) + + +def _group_offs(splits, device="cuda"): + offs = torch.zeros(len(splits) + 1, dtype=torch.int64, device=device) + offs[1:] = torch.cumsum(torch.tensor(splits, dtype=torch.int64, device=device), 0) + return offs + + +def _ref_rowwise_quantize(x: torch.Tensor, dtype: torch.dtype, block: int = BLOCK, pow2: bool = False): + m, n = x.shape + fp8_max = torch.finfo(dtype).max + nb = _cdiv(n, block) + x_pad = torch.nn.functional.pad(x.float(), (0, nb * block - n)) + tiles = x_pad.reshape(m, nb, block) + amax = tiles.abs().amax(dim=-1).clamp_min(1e-4) + scale = fp8_max / amax + if pow2: + scale = _floor_to_pow2(scale) + q = (tiles * scale.unsqueeze(-1)).clamp(-fp8_max, fp8_max).to(dtype) + q = q.reshape(m, nb * block)[:, :n].contiguous() + return q, (1.0 / scale).contiguous() + + +def _ref_colwise_quantize(x: torch.Tensor, dtype: torch.dtype, block: int = BLOCK, pow2: bool = False): + m, n = x.shape + fp8_max = torch.finfo(dtype).max + mb = _cdiv(m, block) + x_pad = torch.nn.functional.pad(x.float(), (0, 0, 0, mb * block - m)) + tiles = x_pad.reshape(mb, block, n) + amax = tiles.abs().amax(dim=1).clamp_min(1e-4) + scale = fp8_max / amax + if pow2: + scale = _floor_to_pow2(scale) + q = (tiles * scale.unsqueeze(1)).clamp(-fp8_max, fp8_max).to(dtype) + q = q.reshape(mb * block, n)[:m].contiguous() + return q, (1.0 / scale).contiguous() + + +def _ref_weight_quantize(w: torch.Tensor, dtype: torch.dtype, block: int = BLOCK, pow2: bool = False): + g, m, n = w.shape + fp8_max = torch.finfo(dtype).max + mb, nb = _cdiv(m, block), _cdiv(n, block) + w_pad = torch.nn.functional.pad(w.float(), (0, nb * block - n, 0, mb * block - m)) + tiles = w_pad.reshape(g, mb, block, nb, block) + amax = tiles.abs().amax(dim=(2, 4)).clamp_min(1e-4) + scale = fp8_max / amax + if pow2: + scale = _floor_to_pow2(scale) + q = (tiles * scale[:, :, None, :, None]).clamp(-fp8_max, fp8_max).to(dtype) + q = q.reshape(g, mb * block, nb * block)[:, :m, :n].contiguous() + return q, (1.0 / scale).contiguous() + + +def _dequant_rowwise(q, s, n, block=BLOCK): + return q.float() * s.repeat_interleave(block, dim=1)[:, :n] + + +def _dequant_colwise(q, s, m, block=BLOCK): + return q.float() * s.repeat_interleave(block, dim=0)[:m] + + +@pytest.mark.parametrize("shape", [(128, 256), (200, 128), (256, 384), (64, 192)]) +@pytest.mark.parametrize("dtype", IN_DTYPES, ids=str) +@pytest.mark.parametrize("axis", [0, 1]) +@pytest.mark.parametrize("pow2", [False, True]) +def test_quantize_fp8_blockwise(shape, dtype, axis, pow2): + x = torch.randn(*shape, dtype=dtype, device="cuda") + q, s = quantize_fp8_blockwise(x, FP8_DTYPE, axis=axis, block_size=BLOCK, pow2=pow2) + if axis == 1: + q_ref, s_ref = _ref_rowwise_quantize(x, FP8_DTYPE, pow2=pow2) + dq = _dequant_rowwise(q, s, x.shape[1]) + dq_ref = _dequant_rowwise(q_ref, s_ref, x.shape[1]) + else: + q_ref, s_ref = _ref_colwise_quantize(x, FP8_DTYPE, pow2=pow2) + dq = _dequant_colwise(q, s, x.shape[0]) + dq_ref = _dequant_colwise(q_ref, s_ref, x.shape[0]) + torch.testing.assert_close(s, s_ref, atol=1e-5, rtol=1e-4) + torch.testing.assert_close(dq, dq_ref, atol=0.25, rtol=0.12) + + +@pytest.mark.parametrize("shape", [(128, 256), (200, 128)]) +@pytest.mark.parametrize("dtype", IN_DTYPES, ids=str) +@pytest.mark.parametrize("pow2", [False, True]) +def test_quantize_fp8_blockwise_dual(shape, dtype, pow2): + x = torch.randn(*shape, dtype=dtype, device="cuda") + q_row, s_row, q_col, s_col = quantize_fp8_blockwise_dual(x, FP8_DTYPE, block_size=BLOCK, pow2=pow2) + q_row_ref, s_row_ref = _ref_rowwise_quantize(x, FP8_DTYPE, pow2=pow2) + q_col_ref, s_col_ref = _ref_colwise_quantize(x, FP8_DTYPE, pow2=pow2) + torch.testing.assert_close(s_row, s_row_ref, atol=1e-5, rtol=1e-4) + torch.testing.assert_close(s_col, s_col_ref, atol=1e-5, rtol=1e-4) + torch.testing.assert_close( + _dequant_rowwise(q_row, s_row, x.shape[1]), + _dequant_rowwise(q_row_ref, s_row_ref, x.shape[1]), + atol=0.25, + rtol=0.12, + ) + torch.testing.assert_close( + _dequant_colwise(q_col, s_col, x.shape[0]), + _dequant_colwise(q_col_ref, s_col_ref, x.shape[0]), + atol=0.25, + rtol=0.12, + ) + + +@pytest.mark.parametrize("shape", [(2, 256, 256), (3, 128, 256), (1, 200, 192), (256, 256)]) +@pytest.mark.parametrize("dtype", IN_DTYPES, ids=str) +@pytest.mark.parametrize("pow2", [False, True]) +def test_quantize_fp8_blockwise_weight(shape, dtype, pow2): + w = torch.randn(*shape, dtype=dtype, device="cuda") + q, s = quantize_fp8_blockwise_weight(w, FP8_DTYPE, block_size=BLOCK, pow2=pow2) + w3 = w if w.dim() == 3 else w.unsqueeze(0) + q_ref, s_ref = _ref_weight_quantize(w3, FP8_DTYPE, pow2=pow2) + if w.dim() == 2: + q_ref, s_ref = q_ref.squeeze(0), s_ref.squeeze(0) + torch.testing.assert_close(s, s_ref, atol=1e-5, rtol=1e-4) + n = w.shape[-1] + m = w.shape[-2] + dq = q.float() * ( + s.repeat_interleave(BLOCK, dim=-2).repeat_interleave(BLOCK, dim=-1)[..., :m, :n] + ) + dq_ref = q_ref.float() * ( + s_ref.repeat_interleave(BLOCK, dim=-2).repeat_interleave(BLOCK, dim=-1)[..., :m, :n] + ) + torch.testing.assert_close(dq, dq_ref, atol=0.25, rtol=0.12) + + +@pytest.mark.parametrize( + "m_splits", + [ + [80, 176], + [128, 128], + [0, 256], + [64, 0, 192], + [200, 56], + ], +) +@pytest.mark.parametrize("n", [128, 256, 192]) +@pytest.mark.parametrize("dtype", IN_DTYPES, ids=str) +def test_quantize_fp8_blockwise_segment_m(m_splits, n, dtype): + m = sum(m_splits) + x = torch.randn(m, n, dtype=dtype, device="cuda") + group_lens = torch.tensor(m_splits, dtype=torch.int64, device="cuda") + group_offs = _group_offs(m_splits) + x_fp8, scales, vk_lens, vk_offs = quantize_fp8_blockwise_segment_m( + x, FP8_DTYPE, BLOCK, group_lens, group_offs + ) + expected_lens = [((s + BLOCK - 1) // BLOCK) * BLOCK for s in m_splits] + expected_offs = [0] + for length in expected_lens: + expected_offs.append(expected_offs[-1] + length) + assert vk_lens.tolist() == expected_lens + assert vk_offs.tolist() == expected_offs + cursor = 0 + for valid, padded in zip(m_splits, expected_lens): + if padded > valid: + assert torch.all(x_fp8[cursor + valid : cursor + padded].float() == 0) + cursor += padded + assert x_fp8.shape[0] >= expected_offs[-1] + assert scales.dtype == torch.float32 + + +@pytest.mark.parametrize( + "splits,n,k", + [ + ([64, 96, 96], 256, 256), + ([128, 128], 256, 192), + ([80, 176], 192, 256), + ([0, 128, 128], 256, 256), + ([256], 128, 128), + ], +) +@pytest.mark.parametrize("trans_b", [True, False]) +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=str) +@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16], ids=str) +def test_grouped_gemm_fp8_blockwise_matches_dequant_ref(splits, n, k, trans_b, in_dtype, out_dtype): + g = len(splits) + m = sum(splits) + out_n = n if trans_b else k + if out_n % BLOCK != 0: + pytest.skip("blockwise grouped GEMM uses 128-wide N tiles") + a_k = k if trans_b else n + a = torch.randn(m, a_k, dtype=in_dtype, device="cuda") + b = torch.randn(g, n, k, dtype=in_dtype, device="cuda") + a_fp8, a_s = quantize_fp8_blockwise(a, FP8_DTYPE, axis=1, block_size=BLOCK) + b_fp8, b_s = quantize_fp8_blockwise_weight(b, FP8_DTYPE, block_size=BLOCK) + offs = _group_offs(splits) + + out = grouped_gemm_fp8_blockwise_triton_kernel( + a_fp8, b_fp8, a_s, b_s, offs, trans_b=trans_b, out_dtype=out_dtype + ) + assert out.dtype == out_dtype + + a_dq = _dequant_rowwise(a_fp8, a_s, a_k) + b_dq = b_fp8.float() * ( + b_s.repeat_interleave(BLOCK, dim=1).repeat_interleave(BLOCK, dim=2)[:, :n, :k] + ) + ref = torch.zeros(m, out_n, dtype=torch.float32, device="cuda") + for i in range(g): + sl = slice(int(offs[i]), int(offs[i + 1])) + if sl.start == sl.stop: + continue + ref[sl] = a_dq[sl] @ (b_dq[i].T if trans_b else b_dq[i]) + torch.testing.assert_close(out.float(), ref, atol=0.15, rtol=0.05) + + +@pytest.mark.parametrize( + "splits,n,k", + [ + ([128, 128], 128, 128), + ([80, 176], 256, 128), + ([0, 256], 256, 128), + ([64, 96, 96], 128, 256), + ], +) +@pytest.mark.parametrize("accumulate", [False, True]) +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=str) +@pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16], ids=str) +def test_variable_k_wgrad(splits, n, k, accumulate, in_dtype, out_dtype): + g = len(splits) + tokens = sum(splits) + dy = torch.randn(tokens, n, dtype=in_dtype, device="cuda") + x = torch.randn(tokens, k, dtype=in_dtype, device="cuda") + group_lens = torch.tensor(splits, dtype=torch.int64, device="cuda") + group_offs = _group_offs(splits) + go_col, go_s, _, vk_offs = quantize_fp8_blockwise_segment_m( + dy, FP8_DTYPE, BLOCK, group_lens, group_offs + ) + x_col, x_s, _, _ = quantize_fp8_blockwise_segment_m( + x, FP8_DTYPE, BLOCK, group_lens, group_offs + ) + + fresh = grouped_gemm_fp8_blockwise_variable_k_triton_kernel( + go_col, x_col, go_s, x_s, vk_offs, out_dtype=out_dtype, accumulate=False + ) + assert fresh.dtype == out_dtype + assert fresh.shape == (g, n, k) + + go_dq = _dequant_colwise(go_col, go_s, go_col.shape[0]) + x_dq = _dequant_colwise(x_col, x_s, x_col.shape[0]) + ref = torch.zeros(g, n, k, dtype=torch.float32, device="cuda") + for i in range(g): + sl = slice(int(vk_offs[i]), int(vk_offs[i + 1])) + if sl.start == sl.stop: + continue + ref[i] = go_dq[sl].T @ x_dq[sl] + torch.testing.assert_close(fresh.float(), ref, atol=0.15, rtol=0.05) + + if accumulate: + main_grad = torch.randn(g, n, k, dtype=out_dtype, device="cuda") + expected = main_grad.clone() + grouped_gemm_fp8_blockwise_variable_k_triton_kernel( + go_col, + x_col, + go_s, + x_s, + vk_offs, + out=main_grad, + accumulate=True, + ) + torch.testing.assert_close(main_grad.float(), (expected + fresh).float(), atol=1e-3, rtol=1e-4) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index bdaceb9f31..39adb7da5d 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -399,6 +399,35 @@ def _forward_grouped_tensor( return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces + @staticmethod + def _packed_main_grad_view(main_grads): + """[G, N, K] view of consecutive contiguous ``main_grad`` buffers, or None.""" + g0 = main_grads[0] + if g0 is None or g0.ndim != 2 or not g0.is_contiguous(): + return None + n, k = g0.shape + step = g0.numel() * g0.element_size() + for i, g in enumerate(main_grads): + if ( + g is None + or g.dtype != g0.dtype + or g.device != g0.device + or tuple(g.shape) != (n, k) + or not g.is_contiguous() + or g.data_ptr() != g0.data_ptr() + i * step + ): + return None + return g0.as_strided((len(main_grads), n, k), (n * k, k, 1)) + + @staticmethod + def _handle_fused_wgrad(weight, main_grad): + """Megatron DDP hook: mark wgrad as consumed and return a dummy (or None).""" + if hasattr(weight, "grad_added_to_main_grad"): + weight.grad_added_to_main_grad = True + shape = list(main_grad.shape) if main_grad is not None else list(weight.shape) + return get_dummy_wgrad(shape, weight.dtype, zero=getattr(weight, "zero_out_wgrad", False)) + return None + @staticmethod def _is_blockwise_fp8_grouped_gemm_supported( *, @@ -406,52 +435,31 @@ def _is_blockwise_fp8_grouped_gemm_supported( recipe, use_bias, backward_override, - fuse_wgrad_accumulation, cpu_offloading, save_original_input, - wgrad_store, debug, unpad_output, actual_m_splits, - m_splits, ) -> bool: - """Return whether the ROCm Triton blockwise FP8 grouped GEMM path can run this call. - - Requires ``NVTE_USE_BLOCKWISE_GMM_TRITON=1``, ``Float8BlockScaling``, and the - hardcoded blockwise layout (activation 1x128 along K, weights 128x128, grad 1x128). - Unsupported orchestration features fall back to the default path. - """ - if not IS_HIP_EXTENSION: - return False - if os.getenv("NVTE_USE_BLOCKWISE_GMM_TRITON", "0") != "1": - return False - if not fp8 or recipe is None or not recipe.float8_block_scaling(): - return False - if ( - recipe.x_block_scaling_dim != 1 - or recipe.w_block_scaling_dim != 2 - or recipe.grad_block_scaling_dim != 1 - ): - return False - if use_bias: - return False - if backward_override is not None: - return False - if fuse_wgrad_accumulation: - return False - if cpu_offloading: - return False - if save_original_input: - return False - if wgrad_store is not None and wgrad_store.delay_wgrad_compute(): - return False - if debug: - return False - if unpad_output or ( - actual_m_splits is not None and list(actual_m_splits) != list(m_splits) - ): - return False - return True + """ROCm Triton blockwise FP8 grouped GEMM: HIP, env, Float8BlockScaling 1x128/128x128.""" + return ( + IS_HIP_EXTENSION + and os.getenv("NVTE_USE_BLOCKWISE_GMM_TRITON", "0") == "1" + and fp8 + and recipe is not None + and recipe.float8_block_scaling() + and (recipe.x_block_scaling_dim, recipe.w_block_scaling_dim, recipe.grad_block_scaling_dim) + == (1, 2, 1) + and not ( + use_bias + or backward_override + or cpu_offloading + or save_original_input + or debug + or unpad_output + or actual_m_splits is not None + ) + ) @staticmethod def _forward_blockwise_fp8( @@ -464,6 +472,9 @@ def _forward_blockwise_fp8( weight_quantizers, activation_dtype, is_grad_enabled, + fuse_wgrad_accumulation, + is_first_microbatch, + wgrad_store, ): """Blockwise FP8 grouped GEMM forward (ROCm Triton). @@ -531,6 +542,20 @@ def _forward_blockwise_fp8( ctx.requires_dgrad = inp.requires_grad ctx.weight_requires_grad = weights[0].requires_grad ctx.reduce_and_update_bwd_fp8_tensors = False + ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation + ctx.is_first_microbatch = is_first_microbatch + ctx.wgrad_store = wgrad_store + if fuse_wgrad_accumulation and ctx.weight_requires_grad: + ctx.origin_weight_refs = [weakref.ref(w) for w in weights] + ctx.origin_weights_overwrite_main_grad = getattr( + weights[0], "overwrite_main_grad", False + ) + if hasattr(weights[0], "__fsdp_param__"): + ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] + else: + ctx.main_grad_funcs = [ + lambda j=i: weights[j].main_grad for i in range(num_gemms) + ] new_workspaces = [None] * num_gemms return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @@ -566,10 +591,75 @@ def _backward_blockwise_fp8(ctx, grad_output): go_col, go_scol, _l, _o = quantize_fp8_blockwise_segment_m( g_out, dt, 128, group_lens, group_offs ) - dW = grouped_gemm_fp8_blockwise_variable_k_triton_kernel( - go_col, a_col, go_scol, a_scol, vk_offs, out_dtype=ctx.activation_dtype - ) - wgrad_list = [dW[g].contiguous() for g in range(ctx.num_gemms)] + fuse = getattr(ctx, "fuse_wgrad_accumulation", False) + origin_weights = [None] * ctx.num_gemms + main_grads = [None] * ctx.num_gemms + accumulate = False + packed_out = None + if fuse: + origin_weight_refs = ctx.origin_weight_refs + ctx.origin_weight_refs = None + origin_weights = [ref() if ref is not None else None for ref in origin_weight_refs] + assert all( + w is not None for w in origin_weights + ), "weight was removed while fuse_wgrad_accumulation=True" + main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] + for origin_weight, main_grad in zip(origin_weights, main_grads): + if main_grad is not None: + origin_weight.main_grad = main_grad + if ctx.is_first_microbatch is not None: + accumulate = not ctx.is_first_microbatch + else: + accumulate = True + if getattr(ctx, "origin_weights_overwrite_main_grad", False): + accumulate = False + packed_out = _GroupedLinear._packed_main_grad_view(main_grads) + wgrad_list = main_grads + out_dtype = main_grads[0].dtype if main_grads[0] is not None else ctx.activation_dtype + else: + out_dtype = ctx.activation_dtype + wgrad_3d = torch.empty( + (ctx.num_gemms, ctx.out_features, a_col.shape[1]), + dtype=out_dtype, + device=a_col.device, + ) + packed_out = wgrad_3d + wgrad_list = [wgrad_3d[g] for g in range(ctx.num_gemms)] + + def grouped_gemm_wgrad(*_unused): + dW = grouped_gemm_fp8_blockwise_variable_k_triton_kernel( + go_col, + a_col, + go_scol, + a_scol, + vk_offs, + out_dtype=out_dtype, + out=packed_out, + accumulate=accumulate and packed_out is not None and fuse, + ) + if fuse and packed_out is None: + for g, mg in enumerate(main_grads): + if mg is None: + continue + if accumulate: + mg.add_(dW[g]) + else: + mg.copy_(dW[g]) + # Signature matches WeightGradStore / GroupedLinear.backward_dw. + return None, [None] * ctx.num_gemms, None + + wgrad_store = getattr(ctx, "wgrad_store", None) + if wgrad_store is not None and wgrad_store.delay_wgrad_compute(): + # tensor_list[2] is the wgrad buffers backward_dw assigns to .grad. + wgrad_store.put([go_col, a_col, wgrad_list], grouped_gemm_wgrad) + else: + grouped_gemm_wgrad() + + if fuse: + wgrad_list = [ + _GroupedLinear._handle_fused_wgrad(weight, main_grad) + for weight, main_grad in zip(origin_weights, main_grads) + ] grad_biases = [None] * ctx.num_gemms # bias not supported on this path @@ -644,14 +734,11 @@ def forward( recipe=FP8GlobalStateManager.get_fp8_recipe() if fp8 else None, use_bias=use_bias, backward_override=backward_override, - fuse_wgrad_accumulation=fuse_wgrad_accumulation, cpu_offloading=cpu_offloading, save_original_input=save_original_input, - wgrad_store=wgrad_store, debug=debug, unpad_output=unpad_output, actual_m_splits=actual_m_splits, - m_splits=m_splits, ): return _GroupedLinear._forward_blockwise_fp8( ctx, @@ -662,6 +749,9 @@ def forward( weight_quantizers=weight_quantizers, activation_dtype=activation_dtype, is_grad_enabled=is_grad_enabled, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + is_first_microbatch=is_first_microbatch, + wgrad_store=wgrad_store, ) # Configure quantizers diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py index f31de8b0ff..b8d266f1fd 100644 --- a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py +++ b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py @@ -377,6 +377,7 @@ def _grouped_blockwise_fp8_variable_k_gemm_kernel( NUM_XCDS: tl.constexpr, CHUNK_SIZE: tl.constexpr, CACHE_MODIFIER: tl.constexpr, + ACCUMULATE: tl.constexpr, ): """Persistent grouped block-wise FP8 variable-K GEMM kernel (backward, CPU-sync-free). @@ -476,12 +477,14 @@ def _grouped_blockwise_fp8_variable_k_gemm_kernel( RHS_BASE += BLOCK_SIZE_K * stride_rhs_m # ── Store output ── - c = acc.to(C.type.element_ty) rm_s = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) rn_s = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) rn_s = tl.max_contiguous(tl.multiple_of(rn_s % OUT_N, BLOCK_SIZE_N), BLOCK_SIZE_N) c_mask = (rm_s[:, None] < OUT_M) & (rn_s[None, :] < OUT_N) C_ = C + group_idx.to(tl.int64) * stride_cg + rm_s[:, None] * stride_cm + rn_s[None, :] * stride_cn + c = acc.to(C.type.element_ty) + if ACCUMULATE: + c += tl.load(C_, mask=c_mask, other=0) tl.store(C_, c, c_mask) @@ -496,24 +499,16 @@ def grouped_gemm_fp8_blockwise_variable_k_triton_kernel( rhs_scales: torch.Tensor, group_offs: torch.Tensor, out_dtype: torch.dtype = torch.bfloat16, + out: torch.Tensor = None, + accumulate: bool = False, ) -> torch.Tensor: """Variable-K grouped block-wise FP8 GEMM (backward, 1D+1D scaling) using Triton. Computes: C[g] = lhs[offs[g]:offs[g+1]]^T @ rhs[offs[g]:offs[g+1]] with 1D+1D block-wise scaling applied in the K-loop. - Output: [G, OUT_M, OUT_N]. - - Args: - lhs: [M_padded_total, OUT_M] FP8 (segment-padded, each segment aligned to 128). - rhs: [M_padded_total, OUT_N] FP8. - lhs_scales: [ceil(M_padded/128), OUT_M] float32. - rhs_scales: [ceil(M_padded/128), OUT_N] float32. - group_offs: [G+1] int64 padded segment offsets. - out_dtype: Output dtype (default bfloat16). - - Returns: - [G, OUT_M, OUT_N] output. + Output: [G, OUT_M, OUT_N]. When ``out`` is provided the kernel writes in-place + (and adds into it if ``accumulate``). ``out_dtype`` is ignored in that case. """ if is_cdna4(): set_triton_knobs_gfx950() @@ -526,7 +521,13 @@ def grouped_gemm_fp8_blockwise_variable_k_triton_kernel( OUT_N = rhs.shape[1] G = group_offs.shape[0] - 1 - out = torch.empty((G, OUT_M, OUT_N), device=lhs.device, dtype=out_dtype) + if out is None: + assert not accumulate, "accumulate=True requires an existing out tensor" + out = torch.empty((G, OUT_M, OUT_N), device=lhs.device, dtype=out_dtype) + else: + assert out.shape == (G, OUT_M, OUT_N), ( + f"out must be {(G, OUT_M, OUT_N)}, got {tuple(out.shape)}" + ) num_sms = get_num_cus() # Use 128x128 tiles to reduce register pressure from double-accumulator @@ -563,6 +564,7 @@ def grouped_gemm_fp8_blockwise_variable_k_triton_kernel( NUM_XCDS=NUM_XCDS, CHUNK_SIZE=32, CACHE_MODIFIER=".ca", + ACCUMULATE=accumulate, num_warps=4, num_stages=2, waves_per_eu=0, From 2cc840fbef63018a2c075a83b6fce923676d9644 Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Tue, 25 Aug 2026 20:57:13 +0000 Subject: [PATCH 04/13] Avoid copying already-packed expert weights in blockwise FP8 grouped GEMM - Generalize `_packed_main_grad_view` into `_packed_3d_view` for any sequence of contiguous 2D buffers. - Add `_expert_weights_as_3d` to return a zero-copy `[G, N, K]` view when expert weights are consecutive slices of a single buffer (e.g. `single_grouped_weight`), falling back to `torch.stack` otherwise. - Use the new helper in `_forward_blockwise_fp8` instead of always stacking weights. --- .../pytorch/module/grouped_linear.py | 29 ++++++++++++++----- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 39adb7da5d..8b998a14d4 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -400,14 +400,14 @@ def _forward_grouped_tensor( return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @staticmethod - def _packed_main_grad_view(main_grads): - """[G, N, K] view of consecutive contiguous ``main_grad`` buffers, or None.""" - g0 = main_grads[0] + def _packed_3d_view(tensors): + """[G, N, K] view of consecutive contiguous 2D buffers, or None.""" + g0 = tensors[0] if g0 is None or g0.ndim != 2 or not g0.is_contiguous(): return None n, k = g0.shape step = g0.numel() * g0.element_size() - for i, g in enumerate(main_grads): + for i, g in enumerate(tensors): if ( g is None or g.dtype != g0.dtype @@ -417,7 +417,22 @@ def _packed_main_grad_view(main_grads): or g.data_ptr() != g0.data_ptr() + i * step ): return None - return g0.as_strided((len(main_grads), n, k), (n * k, k, 1)) + return g0.as_strided((len(tensors), n, k), (n * k, k, 1)) + + @staticmethod + def _expert_weights_as_3d(weights, dtype): + """[G, N, K] expert weights without stacking when storage is already packed. + + ``single_grouped_weight`` keeps experts as consecutive slices of one + ``GroupedTensor.rowwise_data`` buffer; those slices are a zero-copy view. + Discrete ``weight0..weightN`` still need a stack. + """ + packed = _GroupedLinear._packed_3d_view(weights) + if packed is not None: + if packed.dtype != dtype: + packed = packed.to(dtype) + return packed if packed.is_contiguous() else packed.contiguous() + return torch.stack([wt.to(dtype).contiguous() for wt in weights], 0).contiguous() @staticmethod def _handle_fused_wgrad(weight, main_grad): @@ -516,7 +531,7 @@ def _forward_blockwise_fp8( group_offs = torch.zeros(num_gemms + 1, dtype=torch.int64, device=device) group_offs[1:] = torch.cumsum(group_lens, 0) - w = torch.stack([wt.to(activation_dtype).contiguous() for wt in weights], 0).contiguous() + w = _GroupedLinear._expert_weights_as_3d(weights, activation_dtype) # Quantize: activation rowwise (1x128 along K), weights 128x128. a_row, a_srow = quantize_fp8_blockwise(a, dt, axis=1, block_size=128) @@ -613,7 +628,7 @@ def _backward_blockwise_fp8(ctx, grad_output): accumulate = True if getattr(ctx, "origin_weights_overwrite_main_grad", False): accumulate = False - packed_out = _GroupedLinear._packed_main_grad_view(main_grads) + packed_out = _GroupedLinear._packed_3d_view(main_grads) wgrad_list = main_grads out_dtype = main_grads[0].dtype if main_grads[0] is not None else ctx.activation_dtype else: From 09e72bcfbad783d68272a5bccd84dc0fdb7c0b41 Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Tue, 25 Aug 2026 21:47:22 +0000 Subject: [PATCH 05/13] Return GPU tensor split sizes from FP8 padding when given device copy - Extend `Fp8Padding.forward` with an optional `m_splits_tensor` argument. - When provided, compute and return padded split sizes as a tensor on the same device, avoiding a blocking host-to-device copy. --- transformer_engine/pytorch/module/fp8_padding.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/module/fp8_padding.py b/transformer_engine/pytorch/module/fp8_padding.py index 3a0073a492..41ff2607ba 100644 --- a/transformer_engine/pytorch/module/fp8_padding.py +++ b/transformer_engine/pytorch/module/fp8_padding.py @@ -6,7 +6,7 @@ """FP8 Padding API""" -from typing import List, Optional, Tuple +from typing import List, Optional, Tuple, Union import torch @@ -113,7 +113,8 @@ def forward( self, inp: torch.Tensor, m_splits: List[int], - ) -> Tuple[torch.Tensor, List[int]]: + m_splits_tensor: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, Union[List[int], torch.Tensor]]: """ Apply the padding to the input. @@ -123,6 +124,10 @@ def forward( Input tensor. m_splits : List[int] List of integers representing the split of the input tensor. + m_splits_tensor : torch.Tensor, optional + Device copy of ``m_splits``. When provided, the padded split sizes are + returned as a tensor on the same device (rounded up on-device), instead + of a Python list, so the caller can keep the splits on the GPU. """ assert len(m_splits) == self.num_gemms, "Number of splits should match number of GEMMs." @@ -136,6 +141,8 @@ def forward( ] # no padding needed if m_splits == padded_m_splits: + if m_splits_tensor is not None: + return inp, m_splits_tensor return inp, m_splits is_grad_enabled = torch.is_grad_enabled() @@ -154,4 +161,7 @@ def forward( ) out = fn(*autograd_ctx, inp, non_tensor_args) + if m_splits_tensor is not None: + align = self.align_size + return out, ((m_splits_tensor + (align - 1)) // align) * align return out, padded_m_splits From 1006109b4dff5a591bad9d3a0dc4b23b949e6424 Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Tue, 25 Aug 2026 23:38:00 +0000 Subject: [PATCH 06/13] Tighten FP8 blockwise quantization test tolerances - Match Primus-Turbo FP8 quantize tolerances (atol=rtol=0.10) in test_blockwise_fp8.py and the grouped linear blockwise-triton path. - Remove obsolete None-output assertion in test_grouped_linear.py. --- tests/pytorch/test_grouped_linear.py | 7 ++----- tests/pytorch/triton_kernels/test_blockwise_fp8.py | 12 ++++++------ 2 files changed, 8 insertions(+), 11 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index c440ea7dfc..6fb4ce2a8b 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -402,13 +402,10 @@ def test_grouped_linear_accuracy( atol = 2.6e-6 rtol = 5e-2 if use_blockwise_triton: - # tests/pytorch/triton_kernels/test_blockwise_fp8.py dequant checks. - atol, rtol = 0.25, 0.12 + # Match Primus-Turbo FP8 quantize tols (test_blockwise_fp8.py). + atol, rtol = 0.10, 0.10 for o, o_ref in zip(outputs, outputs_ref): if use_blockwise_triton: - if o is None: - assert o_ref is None - continue mag = max(float(o.detach().abs().max()), float(o_ref.detach().abs().max())) tensor_atol = max(atol, 0.05 * mag) torch.testing.assert_close(o, o_ref, rtol=rtol, atol=tensor_atol) diff --git a/tests/pytorch/triton_kernels/test_blockwise_fp8.py b/tests/pytorch/triton_kernels/test_blockwise_fp8.py index 3791696197..4b18c8d6ab 100644 --- a/tests/pytorch/triton_kernels/test_blockwise_fp8.py +++ b/tests/pytorch/triton_kernels/test_blockwise_fp8.py @@ -112,7 +112,7 @@ def test_quantize_fp8_blockwise(shape, dtype, axis, pow2): dq = _dequant_colwise(q, s, x.shape[0]) dq_ref = _dequant_colwise(q_ref, s_ref, x.shape[0]) torch.testing.assert_close(s, s_ref, atol=1e-5, rtol=1e-4) - torch.testing.assert_close(dq, dq_ref, atol=0.25, rtol=0.12) + torch.testing.assert_close(dq, dq_ref, atol=0.10, rtol=0.10) @pytest.mark.parametrize("shape", [(128, 256), (200, 128)]) @@ -128,14 +128,14 @@ def test_quantize_fp8_blockwise_dual(shape, dtype, pow2): torch.testing.assert_close( _dequant_rowwise(q_row, s_row, x.shape[1]), _dequant_rowwise(q_row_ref, s_row_ref, x.shape[1]), - atol=0.25, - rtol=0.12, + atol=0.10, + rtol=0.10, ) torch.testing.assert_close( _dequant_colwise(q_col, s_col, x.shape[0]), _dequant_colwise(q_col_ref, s_col_ref, x.shape[0]), - atol=0.25, - rtol=0.12, + atol=0.10, + rtol=0.10, ) @@ -158,7 +158,7 @@ def test_quantize_fp8_blockwise_weight(shape, dtype, pow2): dq_ref = q_ref.float() * ( s_ref.repeat_interleave(BLOCK, dim=-2).repeat_interleave(BLOCK, dim=-1)[..., :m, :n] ) - torch.testing.assert_close(dq, dq_ref, atol=0.25, rtol=0.12) + torch.testing.assert_close(dq, dq_ref, atol=0.10, rtol=0.10) @pytest.mark.parametrize( From 4dcedb32023cfbb077484be9f5ffc58d0f1eb47b Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Tue, 25 Aug 2026 23:46:04 +0000 Subject: [PATCH 07/13] Add Triton autotuning for ROCm blockwise FP8 grouped GEMM - Introduce curated fwd/dgrad autotune configs from Primus-Turbo for the grouped blockwise FP8 persistent GEMM kernel. - Add warm-up tracking so the first autotune call uses balanced group offsets, preventing the cached config from being tied to a single uneven MoE routing. - Refactor the launch helper to use the autotuned kernel and remove the hard-coded block-size heuristic. --- .../blockwise_fp8_grouped_gemm.py | 377 ++++++++++++++---- 1 file changed, 289 insertions(+), 88 deletions(-) diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py index b8d266f1fd..ccd45f3914 100644 --- a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py +++ b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py @@ -51,9 +51,207 @@ def _set_amd_knobs(enable: bool = True): NUM_XCDS = 8 +# First call per autotune key uses balanced group_offs so the cached config +# is not locked to one uneven MoE routing (Triton keys on G/N/K, not splits). +_grouped_blockwise_warmed = set() +_grouped_blockwise_vk_warmed = set() + + +def _get_grouped_blockwise_autotune_configs(): + """Curated fwd/dgrad tiles from Primus-Turbo (32-config sweep → these 8). + + BLOCK_SIZE_N is pinned to 128: the kernel loads one B scale per N-tile + (BN == SCALE_BLOCK_N). BN=256 applies the wrong scale to half the columns. + """ + return [ + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 32, + }, + num_warps=4, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=4, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 64, + }, + num_warps=4, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=1, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 64, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 64, + }, + num_warps=8, + num_stages=1, + ), + ] + + +def _bwd_autotune_configs(): + """Curated variable-K wgrad tiles from Primus-Turbo. BLOCK_SIZE_K=128 always.""" + return [ + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=1, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=1, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 32, + }, + num_warps=4, + num_stages=2, + ), + ] + # Blockwise grouped FP8 kernel and public entrypoint +@triton.autotune(configs=_get_grouped_blockwise_autotune_configs(), key=["G", "N", "K"]) @triton.jit() def _grouped_blockwise_fp8_persistent_gemm_kernel( # Pointers @@ -285,54 +483,51 @@ def grouped_gemm_fp8_blockwise_triton_kernel( out = torch.empty((M_total, N), device=a.device, dtype=out_dtype) A_scales_t = a_scales.T.contiguous() num_sms = get_num_cus() + even_k = K % 128 == 0 + + def _launch(c_out, offs): + _grouped_blockwise_fp8_persistent_gemm_kernel[(num_sms,)]( + a, + b, + c_out, + A_scales_t, + b_scales, + offs, + G, + N, + K, + a.stride(0), + stride_bg, + stride_bn, + c_out.stride(0), + c_out.stride(1), + A_scales_t.stride(0), + A_scales_t.stride(1), + b_scales.stride(0), + stride_bs_n, + stride_bs_k, + stride_ak=stride_ak, + stride_bk=stride_bk, + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + EVEN_K=even_k, + CACHE_MODIFIER=".ca", + waves_per_eu=0, + matrix_instr_nonkdim=16, + kpack=1, + ) - blk_m = 256 - blk_n = 128 # Keep 128 to match B_scale block alignment - blk_k = 128 - even_k = K % blk_k == 0 - - # GROUP_SIZE_M heuristic (match tensorwise) - tiles_m_per_group = (M_total + G * blk_m - 1) // (G * blk_m) - tiles_n = (N + blk_n - 1) // blk_n - group_m = 8 if min(tiles_m_per_group, tiles_n) < 16 else 4 - - _grouped_blockwise_fp8_persistent_gemm_kernel[(num_sms,)]( - a, - b, - out, - A_scales_t, - b_scales, - group_offs, - G, - N, - K, - a.stride(0), - stride_bg, - stride_bn, - out.stride(0), - out.stride(1), - A_scales_t.stride(0), - A_scales_t.stride(1), - b_scales.stride(0), - stride_bs_n, - stride_bs_k, - stride_ak=stride_ak, - stride_bk=stride_bk, - BLOCK_SIZE_M=blk_m, - BLOCK_SIZE_N=blk_n, - BLOCK_SIZE_K=blk_k, - GROUP_SIZE_M=group_m, - NUM_SMS=num_sms, - NUM_XCDS=NUM_XCDS, - CHUNK_SIZE=32, - EVEN_K=even_k, - CACHE_MODIFIER=".ca", - num_warps=8, - num_stages=1, # 256×128×128 needs 48KB/stage; 2 stages=96KB > 64KB LDS - waves_per_eu=0, - matrix_instr_nonkdim=16, - kpack=1, - ) + # Triton's autotune key is (G, N, K) only. Prime once with balanced offs + # so the cached config is not locked to the first uneven MoE routing. + warm_key = (G, N, K, even_k) + if warm_key not in _grouped_blockwise_warmed: + _grouped_blockwise_warmed.add(warm_key) + per = M_total // max(G, 1) + bal_offs = torch.arange(G + 1, device=group_offs.device, dtype=group_offs.dtype) * per + bal_offs[-1] = M_total + _launch(torch.empty_like(out), bal_offs) + + _launch(out, group_offs) return out @@ -341,6 +536,10 @@ def grouped_gemm_fp8_blockwise_triton_kernel( +@triton.autotune( + configs=_bwd_autotune_configs(), + key=["G", "OUT_M", "OUT_N"], +) @triton.jit() def _grouped_blockwise_fp8_variable_k_gemm_kernel( # C[g] = LHS_g^T @ RHS_g * block_scales @@ -530,45 +729,47 @@ def grouped_gemm_fp8_blockwise_variable_k_triton_kernel( ) num_sms = get_num_cus() - # Use 128x128 tiles to reduce register pressure from double-accumulator - # (partial + acc both need full tile VGPRs for blockwise scale application). - # With 256x256 tiles + 8 warps, 2 accumulator sets need ~290 VGPRs/wave - # which exceeds the 256 limit at 2 waves/SIMD, causing spilling. - # 128x128 with 4 warps keeps VGPRs at ~170/wave, fitting 2 waves/SIMD. - _grouped_blockwise_fp8_variable_k_gemm_kernel[(num_sms,)]( - lhs, - rhs, - out, - lhs_scales, - rhs_scales, - group_offs, - G, - OUT_M, - OUT_N, - lhs.stride(0), - rhs.stride(0), - out.stride(0), - out.stride(1), - out.stride(2), - lhs_scales.stride(0), - lhs_scales.stride(1), - rhs_scales.stride(0), - rhs_scales.stride(1), - stride_lhs_n=lhs.stride(1), - stride_rhs_n=rhs.stride(1), - BLOCK_SIZE_M=128, - BLOCK_SIZE_N=128, - BLOCK_SIZE_K=128, - GROUP_SIZE_M=4, - NUM_SMS=num_sms, - NUM_XCDS=NUM_XCDS, - CHUNK_SIZE=32, - CACHE_MODIFIER=".ca", - ACCUMULATE=accumulate, - num_warps=4, - num_stages=2, - waves_per_eu=0, - matrix_instr_nonkdim=16, - kpack=1, - ) + def _launch(c_out, offs, accumulate_flag): + _grouped_blockwise_fp8_variable_k_gemm_kernel[(num_sms,)]( + lhs, + rhs, + c_out, + lhs_scales, + rhs_scales, + offs, + G, + OUT_M, + OUT_N, + lhs.stride(0), + rhs.stride(0), + c_out.stride(0), + c_out.stride(1), + c_out.stride(2), + lhs_scales.stride(0), + lhs_scales.stride(1), + rhs_scales.stride(0), + rhs_scales.stride(1), + stride_lhs_n=lhs.stride(1), + stride_rhs_n=rhs.stride(1), + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + CACHE_MODIFIER=".ca", + ACCUMULATE=accumulate_flag, + waves_per_eu=0, + matrix_instr_nonkdim=16, + kpack=1, + ) + + # Autotune key is (G, OUT_M, OUT_N). Prime on a scratch buffer with + # balanced padded offs so accumulate=True never benchmarks into ``out``. + warm_key = (G, OUT_M, OUT_N, out.dtype, accumulate) + if warm_key not in _grouped_blockwise_vk_warmed: + _grouped_blockwise_vk_warmed.add(warm_key) + m_padded = lhs.shape[0] + per = max((m_padded // max(G, 1)) // 128 * 128, 128) + bal_offs = torch.arange(G + 1, device=group_offs.device, dtype=group_offs.dtype) * per + bal_offs[-1] = m_padded + _launch(torch.zeros_like(out), bal_offs, False) + + _launch(out, group_offs, accumulate) return out From 71edbb8e79ffb8d1e73d17de3dcc35e196b9b609 Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Tue, 25 Aug 2026 23:54:54 +0000 Subject: [PATCH 08/13] Fix grouped weight buffer validation and relax blockwise FP8 test tolerances - In `grouped_linear.py`, ensure grouped tensors actually share the same underlying storage and that the storage is large enough for all slices before returning a zero-copy 3D view. - Update `test_grouped_linear.py` tolerances for the blockwise Triton path to account for two independent FP8 quantization stacks. --- tests/pytorch/test_grouped_linear.py | 5 +++-- transformer_engine/pytorch/module/grouped_linear.py | 4 ++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 6fb4ce2a8b..c5765e21aa 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -402,8 +402,9 @@ def test_grouped_linear_accuracy( atol = 2.6e-6 rtol = 5e-2 if use_blockwise_triton: - # Match Primus-Turbo FP8 quantize tols (test_blockwise_fp8.py). - atol, rtol = 0.10, 0.10 + # Sequential Linear uses TE Float8BlockQuantizer + TE GEMM; this path + # uses Triton quant + grouped GEMM. Budget two independent FP8 stacks. + atol, rtol = 0.25, 0.12 for o, o_ref in zip(outputs, outputs_ref): if use_blockwise_triton: mag = max(float(o.detach().abs().max()), float(o_ref.detach().abs().max())) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 8b998a14d4..8bd8e0fac1 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -407,6 +407,9 @@ def _packed_3d_view(tensors): return None n, k = g0.shape step = g0.numel() * g0.element_size() + nbytes_needed = g0.storage_offset() * g0.element_size() + len(tensors) * step + if g0.untyped_storage().size() < nbytes_needed: + return None for i, g in enumerate(tensors): if ( g is None @@ -414,6 +417,7 @@ def _packed_3d_view(tensors): or g.device != g0.device or tuple(g.shape) != (n, k) or not g.is_contiguous() + or g.untyped_storage().data_ptr() != g0.untyped_storage().data_ptr() or g.data_ptr() != g0.data_ptr() + i * step ): return None From 8dd644127151b9db6602df3cb4462690c118e4f3 Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Wed, 26 Aug 2026 05:04:59 +0000 Subject: [PATCH 09/13] Refactor AMD Triton compiler knob handling for blockwise FP8 grouped GEMM - Replace CDNA4-specific checks with CDNA3 detection and remove `is_cdna4`. - Force gfx950 compiler knobs (async_copy, block_pingpong, scalarize) always on. - Gate gfx942 knobs by GEMM layout, disabling them for TN/wgrad to avoid regressions. - Reorder the blockwise FP8 test entry in the PyTorch CI script. --- ci/pytorch.sh | 2 +- .../blockwise_fp8_grouped_gemm.py | 43 +++++++++++-------- .../pytorch/triton_kernels/common.py | 3 -- 3 files changed, 25 insertions(+), 23 deletions(-) diff --git a/ci/pytorch.sh b/ci/pytorch.sh index bae141110c..80eada7b6e 100755 --- a/ci/pytorch.sh +++ b/ci/pytorch.sh @@ -94,11 +94,11 @@ run_test_config(){ NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 run_default_fa_lbl "deterministic" 3 attention/test_attention.py -k "test_deterministic_bwd_ck" run_default_fa 1 attention/test_cp_utils.py run_default_fa 1 attention/test_kv_cache.py + run_default_fa 1 triton_kernels/test_blockwise_fp8.py run_default_fa 1 triton_kernels/test_cast.py run_default_fa 1 triton_kernels/test_cast_mxfp8.py run_default_fa 1 triton_kernels/test_cast_mxfp4.py run_default_fa 1 triton_kernels/test_grouped_gemm.py - run_default_fa 1 triton_kernels/test_blockwise_fp8.py run_default_fa 1 triton_kernels/test_utils.py NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa 1 triton_kernels/test_norms.py NVTE_ROCM_ENABLE_MXFP8=1 NVTE_TEST_TRITON_AUTOTUNE=1 run_default_fa_lbl "autotune" 3 triton_kernels/test_norms.py diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py index ccd45f3914..e9b48a5a0d 100644 --- a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py +++ b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py @@ -11,7 +11,7 @@ import triton import triton.language as tl -from .common import is_cdna4 +from .common import is_cdna3 from .gmm.pid_preprocessing import remap_xcd_chunked @@ -23,32 +23,41 @@ def get_num_cus() -> int: _KNOBS_SET = False -def set_triton_knobs_gfx950() -> None: - """Enable AMD compiler knobs for gfx950 (async_copy, block_pingpong, scalarize).""" +def _set_triton_knobs_gfx950() -> None: + """Force-on AMD compiler knobs for gfx950 (async_copy, block_pingpong, scalarize).""" global _KNOBS_SET if _KNOBS_SET: return _KNOBS_SET = True + os.environ["TRITON_HIP_USE_ASYNC_COPY"] = "1" + os.environ["AMDGCN_SCALARIZE_PACKED_FOPS"] = "1" + os.environ["TRITON_HIP_USE_BLOCK_PINGPONG"] = "1" if hasattr(triton, "knobs") and hasattr(triton.knobs, "amd"): triton.knobs.amd.use_async_copy = True triton.knobs.amd.scalarize_packed_fops = True triton.knobs.amd.use_block_pingpong = True - else: - os.environ.setdefault("TRITON_HIP_USE_ASYNC_COPY", "1") - os.environ.setdefault("AMDGCN_SCALARIZE_PACKED_FOPS", "1") - os.environ.setdefault("TRITON_HIP_USE_BLOCK_PINGPONG", "1") -def _set_amd_knobs(enable: bool = True): - """Set AMD-specific Triton knobs (non-gfx950 fallback). - NOTE: use_async_copy and scalarize_packed_fops help NT/NN but regress - TN/wgrad ~5-8% on gfx942, so callers gate ``enable`` per layout. + +def _set_triton_knobs_gfx942(enable: bool = True): + """Set AMD Triton knobs on gfx942 (CDNA3). + + ``use_async_copy`` / ``scalarize_packed_fops`` help NT/NN but regress + TN/wgrad ~5-8% on gfx942, so callers pass ``enable`` from layout. """ if hasattr(triton, "knobs") and hasattr(triton.knobs, "amd"): triton.knobs.amd.use_async_copy = enable triton.knobs.amd.scalarize_packed_fops = enable +def _apply_amd_compiler_knobs(*, is_tn: bool) -> None: + """gfx942: knobs from GEMM layout. Else (gfx950): always-on gfx950 knobs.""" + if is_cdna3(): + _set_triton_knobs_gfx942(enable=not is_tn) + else: + _set_triton_knobs_gfx950() + + NUM_XCDS = 8 # First call per autotune key uses balanced group_offs so the cached config @@ -452,10 +461,8 @@ def grouped_gemm_fp8_blockwise_triton_kernel( Returns: [M_total, N] output in out_dtype. """ - if is_cdna4(): - set_triton_knobs_gfx950() - else: - _set_amd_knobs(enable=True) + # trans_a is always False here: NT if trans_b else NN (never TN). + _apply_amd_compiler_knobs(is_tn=False) assert a.ndim == 2, f"a must be 2D, got {a.shape}" assert b.ndim == 3, f"b must be 3D, got {b.shape}" @@ -709,10 +716,8 @@ def grouped_gemm_fp8_blockwise_variable_k_triton_kernel( Output: [G, OUT_M, OUT_N]. When ``out`` is provided the kernel writes in-place (and adds into it if ``accumulate``). ``out_dtype`` is ignored in that case. """ - if is_cdna4(): - set_triton_knobs_gfx950() - else: - _set_amd_knobs(enable=False) + # Variable-K wgrad is TN (lhs^T @ rhs). + _apply_amd_compiler_knobs(is_tn=True) assert lhs.ndim == 2 and rhs.ndim == 2 assert lhs.shape[0] == rhs.shape[0] diff --git a/transformer_engine/pytorch/triton_kernels/common.py b/transformer_engine/pytorch/triton_kernels/common.py index 2487ee3fe1..212c8cd695 100644 --- a/transformer_engine/pytorch/triton_kernels/common.py +++ b/transformer_engine/pytorch/triton_kernels/common.py @@ -14,9 +14,6 @@ def get_arch(): def is_cdna3(): return get_arch() == "gfx942" -def is_cdna4(): - return get_arch() == "gfx950" - get_torch_e4m3_type = lambda: torch.float8_e4m3fn if not is_cdna3() else torch.float8_e4m3fnuz get_torch_e5m2_type = lambda: torch.float8_e5m2 if not is_cdna3() else torch.float8_e5m2fnuz From 04efa5b177d6557cf00befe4e7a132c192285abf Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Wed, 26 Aug 2026 18:50:23 +0000 Subject: [PATCH 10/13] Refine ROCm blockwise FP8 grouped GEMM gating, quantization, and compiler-knob scoping - Rename blockwise FP8 helpers to `..._triton` and tighten `_is_blockwise_fp8_triton_grouped_gemm_supported`: require 128-aligned `in_features`/`out_features` and reject `fp8_weights` to avoid double quantization. - Add `pow2` rounding flags to blockwise quantization helpers and plumb them through the grouped-linear forward/backward. - Scope AMD Triton compiler knobs with a context manager so gfx950/gfx942 overrides don't leak into unrelated kernels; wrap the grouped GEMM launch instead of setting globals. - Add unit tests for the 128-alignment gate and wgrad tail-tile coverage; drop the blockwise-triton-specific accuracy tolerance branch and skip FP8-block-scaling cases for CUTLASS/HipKittens/CK ROCm backends. --- tests/pytorch/test_grouped_linear.py | 71 ++++++- .../triton_kernels/test_blockwise_fp8.py | 8 + .../pytorch/module/grouped_linear.py | 167 ++++++++++++----- .../blockwise_fp8_grouped_gemm.py | 174 +++++++++++------- .../triton_kernels/blockwise_quantize.py | 90 ++++++--- 5 files changed, 367 insertions(+), 143 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index c5765e21aa..aedf97f5a9 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -48,6 +48,7 @@ skip_unsupported_backward_override, ) from triton_kernels.test_common import get_tolerances +from transformer_engine.pytorch.triton_kernels.common import get_torch_e4m3_type # Only run FP8 tests on supported devices. fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) @@ -401,17 +402,61 @@ def test_grouped_linear_accuracy( if dtype == torch.float32: atol = 2.6e-6 rtol = 5e-2 - if use_blockwise_triton: - # Sequential Linear uses TE Float8BlockQuantizer + TE GEMM; this path - # uses Triton quant + grouped GEMM. Budget two independent FP8 stacks. - atol, rtol = 0.25, 0.12 for o, o_ref in zip(outputs, outputs_ref): - if use_blockwise_triton: - mag = max(float(o.detach().abs().max()), float(o_ref.detach().abs().max())) - tensor_atol = max(atol, 0.05 * mag) - torch.testing.assert_close(o, o_ref, rtol=rtol, atol=tensor_atol) - else: - torch.testing.assert_close(o, o_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(o, o_ref, rtol=rtol, atol=atol) + + +@pytest.mark.skipif(not IS_HIP_EXTENSION, reason="Blockwise FP8 grouped GEMM is ROCm-only.") +@pytest.mark.parametrize( + "in_features,out_features,expected", + [ + (256, 512, True), + (384, 512, True), # 384 % 128 == 0 (not a multiple of 256, still supported) + (192, 512, False), # in_features not a multiple of 128 + (256, 320, False), # out_features not a multiple of 128 + ], +) +def test_blockwise_fp8_gate_requires_128_aligned_features( + in_features, out_features, expected, monkeypatch +): + """The forward/dgrad kernel applies one scalar B scale per 128-wide N-tile, + so non-128-aligned in/out features must fall back to the default path.""" + from transformer_engine.pytorch.module.grouped_linear import _GroupedLinear + + monkeypatch.setenv("NVTE_USE_BLOCKWISE_GMM_TRITON", "1") + supported = _GroupedLinear._is_blockwise_fp8_triton_grouped_gemm_supported( + fp8=True, + recipe=recipe.Float8BlockScaling(), + use_bias=False, + backward_override=None, + cpu_offloading=False, + save_original_input=False, + debug=False, + unpad_output=False, + actual_m_splits=None, + in_features=in_features, + out_features=out_features, + fp8_weights=False, + ) + assert supported is expected + + # Quantized weight params (fp8_model_params) must always fall back to avoid + # dequantize -> re-quantize double quantization. + supported_fp8_weights = _GroupedLinear._is_blockwise_fp8_triton_grouped_gemm_supported( + fp8=True, + recipe=recipe.Float8BlockScaling(), + use_bias=False, + backward_override=None, + cpu_offloading=False, + save_original_input=False, + debug=False, + unpad_output=False, + actual_m_splits=None, + in_features=in_features, + out_features=out_features, + fp8_weights=True, + ) + assert supported_fp8_weights is False @pytest.mark.skipif( @@ -484,6 +529,12 @@ def test_grouped_linear_accuracy_rocm_backends( ): pytest.skip("CK MXFP8 grouped GEMM only supported on gfx1250.") + if recipe is not None and recipe.float8_block_scaling(): + # The CUTLASS/HipKittens/CK grouped GEMM backends do not support FP8 block + # scaling (they abort at runtime). The blockwise path is covered by the + # default and Triton backends in test_grouped_linear_accuracy instead. + pytest.skip("CUTLASS/HipKittens/CK grouped GEMM backends do not support FP8 block scaling.") + monkeypatch.setenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "1") monkeypatch.delenv("NVTE_USE_HIPKITTENS_GROUPED_GEMM", raising=False) monkeypatch.delenv("NVTE_USE_CK_GROUPED_GEMM", raising=False) diff --git a/tests/pytorch/triton_kernels/test_blockwise_fp8.py b/tests/pytorch/triton_kernels/test_blockwise_fp8.py index 4b18c8d6ab..21cde0ca7c 100644 --- a/tests/pytorch/triton_kernels/test_blockwise_fp8.py +++ b/tests/pytorch/triton_kernels/test_blockwise_fp8.py @@ -247,6 +247,14 @@ def test_grouped_gemm_fp8_blockwise_matches_dequant_ref(splits, n, k, trans_b, i ([80, 176], 256, 128), ([0, 256], 256, 128), ([64, 96, 96], 128, 256), + # OUT_N == k == 192 is not a multiple of BLOCK_SIZE_N (128 or 256), so the + # wgrad kernel always hits an N tail tile regardless of the autotuned + # config. Covers tail-tile addressing / the un-wrapped store mask for both + # the fresh and ACCUMULATE paths. (Note: the wrapped-store race the mask + # removes converges to the correct value here, so this guards addressing, + # not the race itself.) + ([128, 128], 256, 192), + ([64, 96, 96], 128, 192), ], ) @pytest.mark.parametrize("accumulate", [False, True]) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 8bd8e0fac1..d49be027dd 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -444,11 +444,13 @@ def _handle_fused_wgrad(weight, main_grad): if hasattr(weight, "grad_added_to_main_grad"): weight.grad_added_to_main_grad = True shape = list(main_grad.shape) if main_grad is not None else list(weight.shape) - return get_dummy_wgrad(shape, weight.dtype, zero=getattr(weight, "zero_out_wgrad", False)) + return get_dummy_wgrad( + shape, weight.dtype, zero=getattr(weight, "zero_out_wgrad", False) + ) return None @staticmethod - def _is_blockwise_fp8_grouped_gemm_supported( + def _is_blockwise_fp8_triton_grouped_gemm_supported( *, fp8, recipe, @@ -459,6 +461,9 @@ def _is_blockwise_fp8_grouped_gemm_supported( debug, unpad_output, actual_m_splits, + in_features, + out_features, + fp8_weights, ) -> bool: """ROCm Triton blockwise FP8 grouped GEMM: HIP, env, Float8BlockScaling 1x128/128x128.""" return ( @@ -467,8 +472,23 @@ def _is_blockwise_fp8_grouped_gemm_supported( and fp8 and recipe is not None and recipe.float8_block_scaling() - and (recipe.x_block_scaling_dim, recipe.w_block_scaling_dim, recipe.grad_block_scaling_dim) + and ( + recipe.x_block_scaling_dim, + recipe.w_block_scaling_dim, + recipe.grad_block_scaling_dim, + ) == (1, 2, 1) + # The path quantizes high-precision weights itself. Quantized weight + # params (fp8_model_params) would be dequantized here and re-quantized + # by the kernel (double quantization), so fall back to the default path + # until the stored rowwise_data + scales are consumed directly. + and not fp8_weights + # The forward/dgrad kernel applies one scalar B scale per N-tile + # (BLOCK_SIZE_N == 128 == scale-block-N), so an N that is not a + # multiple of 128 would apply the wrong scale to the tail tile. + # Forward's N is out_features, dgrad's N is in_features. + and in_features % 128 == 0 + and out_features % 128 == 0 and not ( use_bias or backward_override @@ -481,7 +501,7 @@ def _is_blockwise_fp8_grouped_gemm_supported( ) @staticmethod - def _forward_blockwise_fp8( + def _forward_blockwise_fp8_triton( ctx, *, inp, @@ -494,13 +514,16 @@ def _forward_blockwise_fp8( fuse_wgrad_accumulation, is_first_microbatch, wgrad_store, + pow2_x=True, + pow2_w=True, + pow2_grad=True, ): """Blockwise FP8 grouped GEMM forward (ROCm Triton). Selected from :meth:`forward` under the ``Float8BlockScaling`` recipe when - :meth:`_is_blockwise_fp8_grouped_gemm_supported` is true. Activations are + :meth:`_is_blockwise_fp8_triton_grouped_gemm_supported` is true. Activations are quantized rowwise (1x128 along K), weights 128x128; the backward wgrad (see - :meth:`_backward_blockwise_fp8`) uses a segment-padded columnwise operand + :meth:`_backward_blockwise_fp8_triton`) uses a segment-padded columnwise operand quantized from the original high-precision input (no double-quantization). """ from ..triton_kernels.common import te_dtype_to_torch_dtype @@ -538,8 +561,8 @@ def _forward_blockwise_fp8( w = _GroupedLinear._expert_weights_as_3d(weights, activation_dtype) # Quantize: activation rowwise (1x128 along K), weights 128x128. - a_row, a_srow = quantize_fp8_blockwise(a, dt, axis=1, block_size=128) - b_fp8, b_scale = quantize_fp8_blockwise_weight(w, dt, block_size=128) + a_row, a_srow = quantize_fp8_blockwise(a, dt, axis=1, block_size=128, pow2=pow2_x) + b_fp8, b_scale = quantize_fp8_blockwise_weight(w, dt, block_size=128, pow2=pow2_w) # Forward GEMM: out[seg] = A[seg] @ W[g]^T (trans_b=True). out = grouped_gemm_fp8_blockwise_triton_kernel( @@ -547,13 +570,14 @@ def _forward_blockwise_fp8( ) if is_grad_enabled: - ctx.use_blockwise_fp8 = True + ctx.use_blockwise_fp8_triton = True # Segment-padded columnwise activation for the variable-K wgrad. a_col, a_scol, _vk_lens, vk_offs = quantize_fp8_blockwise_segment_m( - a, dt, 128, group_lens, group_offs + a, dt, 128, group_lens, group_offs, pow2=pow2_x ) ctx.save_for_backward(a_col, a_scol, b_fp8, b_scale, group_lens, group_offs, vk_offs) ctx.dt = dt + ctx.pow2_grad = pow2_grad ctx.activation_dtype = activation_dtype ctx.num_gemms = num_gemms ctx.inp_shape = inp_shape @@ -580,8 +604,8 @@ def _forward_blockwise_fp8( return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @staticmethod - def _backward_blockwise_fp8(ctx, grad_output): - """Backward path paired with :meth:`_forward_blockwise_fp8`.""" + def _backward_blockwise_fp8_triton(ctx, grad_output): + """Backward path paired with :meth:`_forward_blockwise_fp8_triton`.""" from ..triton_kernels.blockwise_quantize import ( quantize_fp8_blockwise, quantize_fp8_blockwise_segment_m, @@ -598,9 +622,16 @@ def _backward_blockwise_fp8(ctx, grad_output): # dgrad: dX[seg] = dY[seg] @ W[g] (trans_b=False) -> [M_total, K]. dgrad = None if ctx.requires_dgrad: - go_row, go_srow = quantize_fp8_blockwise(g_out, dt, axis=1, block_size=128) + go_row, go_srow = quantize_fp8_blockwise( + g_out, dt, axis=1, block_size=128, pow2=ctx.pow2_grad + ) dgrad = grouped_gemm_fp8_blockwise_triton_kernel( - go_row, b_fp8, go_srow, b_scale, group_offs, trans_b=False, + go_row, + b_fp8, + go_srow, + b_scale, + group_offs, + trans_b=False, out_dtype=ctx.activation_dtype, ) @@ -608,7 +639,7 @@ def _backward_blockwise_fp8(ctx, grad_output): wgrad_list = [None] * ctx.num_gemms if ctx.weight_requires_grad: go_col, go_scol, _l, _o = quantize_fp8_blockwise_segment_m( - g_out, dt, 128, group_lens, group_offs + g_out, dt, 128, group_lens, group_offs, pow2=ctx.pow2_grad ) fuse = getattr(ctx, "fuse_wgrad_accumulation", False) origin_weights = [None] * ctx.num_gemms @@ -634,7 +665,9 @@ def _backward_blockwise_fp8(ctx, grad_output): accumulate = False packed_out = _GroupedLinear._packed_3d_view(main_grads) wgrad_list = main_grads - out_dtype = main_grads[0].dtype if main_grads[0] is not None else ctx.activation_dtype + out_dtype = ( + main_grads[0].dtype if main_grads[0] is not None else ctx.activation_dtype + ) else: out_dtype = ctx.activation_dtype wgrad_3d = torch.empty( @@ -738,7 +771,12 @@ def forward( save_original_input = True # Check if Triton kernel should be used - use_grouped_gemm_triton = IS_HIP_EXTENSION and os.getenv("NVTE_USE_GROUPED_GEMM_TRITON", "0") == "1" and not fp8 and not fuse_wgrad_accumulation + use_grouped_gemm_triton = ( + IS_HIP_EXTENSION + and os.getenv("NVTE_USE_GROUPED_GEMM_TRITON", "0") == "1" + and not fp8 + and not fuse_wgrad_accumulation + ) num_gemms = len(m_splits) weights = weights_and_biases[:num_gemms] @@ -748,9 +786,10 @@ def forward( # Blockwise FP8 grouped GEMM (ROCm Triton) opt-in. Runs its own quantization + # grouped GEMM and returns early, bypassing the default quantizer setup below. - if _GroupedLinear._is_blockwise_fp8_grouped_gemm_supported( + blockwise_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + if _GroupedLinear._is_blockwise_fp8_triton_grouped_gemm_supported( fp8=fp8, - recipe=FP8GlobalStateManager.get_fp8_recipe() if fp8 else None, + recipe=blockwise_recipe, use_bias=use_bias, backward_override=backward_override, cpu_offloading=cpu_offloading, @@ -758,8 +797,11 @@ def forward( debug=debug, unpad_output=unpad_output, actual_m_splits=actual_m_splits, + in_features=weights[0].size(-1), + out_features=weights[0].size(0), + fp8_weights=isinstance(weights[0], QuantizedTensorStorage), ): - return _GroupedLinear._forward_blockwise_fp8( + return _GroupedLinear._forward_blockwise_fp8_triton( ctx, inp=inp, m_splits=m_splits, @@ -771,6 +813,11 @@ def forward( fuse_wgrad_accumulation=fuse_wgrad_accumulation, is_first_microbatch=is_first_microbatch, wgrad_store=wgrad_store, + # Match the recipe's scale rounding (pow2 by default) so the Triton + # path quantizes identically to the default Float8BlockScaling path. + pow2_x=blockwise_recipe.fp8_quant_fwd_inp.power_2_scale, + pow2_w=blockwise_recipe.fp8_quant_fwd_weight.power_2_scale, + pow2_grad=blockwise_recipe.fp8_quant_bwd_grad.power_2_scale, ) # Configure quantizers @@ -872,12 +919,19 @@ def forward( # tensors (like scales), but bulk allocation shares storage across all tensors, # so if scales can't be offloaded, nothing in the group can be offloaded. fused_padding_kwargs = {} - if actual_m_splits is not None and IS_HIP_EXTENSION \ - and inp_view.shape[0] == sum(actual_m_splits): + if ( + actual_m_splits is not None + and IS_HIP_EXTENSION + and inp_view.shape[0] == sum(actual_m_splits) + ): fused_padding_kwargs["valid_split_sections"] = actual_m_splits inputmats = tex.split_quantize( - inp_view, m_splits, input_quantizers, - disable_bulk_allocation=cpu_offloading, **fused_padding_kwargs) + inp_view, + m_splits, + input_quantizers, + disable_bulk_allocation=cpu_offloading, + **fused_padding_kwargs, + ) elif debug: inputmats = DebugQuantizer.multi_tensor_quantize( inp_view, input_quantizers, m_splits, activation_dtype @@ -951,12 +1005,17 @@ def forward( **kwargs, ) - output_unpadded = False - if unpad_output and actual_m_splits is not None and IS_HIP_EXTENSION and actual_m_splits != m_splits: + if ( + unpad_output + and actual_m_splits is not None + and IS_HIP_EXTENSION + and actual_m_splits != m_splits + ): out_unpadded = torch.empty( [sum(actual_m_splits), out.shape[-1]], - dtype=out.dtype, device=out.device, + dtype=out.dtype, + device=out.device, ) tex.fused_multi_row_unpadding(out, out_unpadded, m_splits, actual_m_splits) out = out_unpadded @@ -1270,8 +1329,8 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring with get_nvtx_range_context("_GroupedLinear_backward"): - if getattr(ctx, "use_blockwise_fp8", False): - return _GroupedLinear._backward_blockwise_fp8(ctx, grad_output) + if getattr(ctx, "use_blockwise_fp8_triton", False): + return _GroupedLinear._backward_blockwise_fp8_triton(ctx, grad_output) if ctx.use_grouped_tensor_path: return _GroupedLinear._backward_grouped_tensor(ctx, grad_output) @@ -1279,7 +1338,7 @@ def backward( N = ctx.num_gemms num_inputs = ctx.num_input_tensors inputmats = saved_tensors[:num_inputs] - weights = saved_tensors[num_inputs: num_inputs + N] + weights = saved_tensors[num_inputs : num_inputs + N] saved_weights = saved_tensors[num_inputs + N : num_inputs + 2 * N] biases = saved_tensors[num_inputs + 2 * N : num_inputs + 3 * N] @@ -1331,13 +1390,19 @@ def backward( for i in range(ctx.num_gemms): grad_biases[i] = grad_output_mats[i].sum(dim=0) grad_output = tex.split_quantize( - grad_output_view, ctx.m_splits, - ctx.grad_output_quantizers, **bwd_fused_kwargs) + grad_output_view, + ctx.m_splits, + ctx.grad_output_quantizers, + **bwd_fused_kwargs, + ) else: # Multi-tensor quantize grad_output = tex.split_quantize( - grad_output_view, ctx.m_splits, - ctx.grad_output_quantizers, **bwd_fused_kwargs) + grad_output_view, + ctx.m_splits, + ctx.grad_output_quantizers, + **bwd_fused_kwargs, + ) elif ctx.debug: grad_output_mats = torch.split(grad_output_view, ctx.m_splits) for i in range(ctx.num_gemms): @@ -1353,9 +1418,9 @@ def backward( # wgrad GEMM. if not ctx.use_grouped_gemm_triton: grad_output = torch.split( - cast_if_needed(grad_output_view, ctx.activation_dtype), - ctx.m_splits, - ) + cast_if_needed(grad_output_view, ctx.activation_dtype), + ctx.m_splits, + ) else: grad_output = [cast_if_needed(grad_output_view, ctx.activation_dtype)] @@ -1427,14 +1492,21 @@ def backward( **kwargs, ) - if ctx.actual_m_splits is not None and ctx.actual_m_splits != ctx.m_splits \ - and not ctx.output_unpadded: + if ( + ctx.actual_m_splits is not None + and ctx.actual_m_splits != ctx.m_splits + and not ctx.output_unpadded + ): dgrad_unpadded = torch.empty( (sum(ctx.actual_m_splits), dgrad.shape[-1]), - dtype=dgrad.dtype, device=dgrad.device, + dtype=dgrad.dtype, + device=dgrad.device, ) tex.fused_multi_row_unpadding( - dgrad, dgrad_unpadded, ctx.m_splits, ctx.actual_m_splits, + dgrad, + dgrad_unpadded, + ctx.m_splits, + ctx.actual_m_splits, ) dgrad = dgrad_unpadded @@ -1461,7 +1533,7 @@ def backward( wgrad_list = torch.empty( (ctx.num_gemms, weights[0].size(0), weights[0].size(1)), dtype=ctx.activation_dtype, - device=ctx.device + device=ctx.device, ) if ctx.save_original_input: @@ -1480,12 +1552,15 @@ def backward( inputmats: list if ctx.fp8 and not ctx.debug: save_fused_kwargs = {} - if ctx.actual_m_splits is not None and IS_HIP_EXTENSION \ - and inp_view.shape[0] == sum(ctx.actual_m_splits): + if ( + ctx.actual_m_splits is not None + and IS_HIP_EXTENSION + and inp_view.shape[0] == sum(ctx.actual_m_splits) + ): save_fused_kwargs["valid_split_sections"] = ctx.actual_m_splits inputmats = tex.split_quantize( - inp_view, ctx.m_splits, ctx.input_quantizers, - **save_fused_kwargs) + inp_view, ctx.m_splits, ctx.input_quantizers, **save_fused_kwargs + ) elif ctx.debug: inputmats = DebugQuantizer.multi_tensor_quantize( inp_view, diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py index e9b48a5a0d..0f255763a1 100644 --- a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py +++ b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py @@ -5,6 +5,7 @@ # Blockwise FP8 grouped GEMM Triton kernels (MoE), adapted from AMD Primus-Turbo # (primus_turbo/triton/grouped_gemm/grouped_gemm_fp8_kernel.py). +import contextlib import os import torch @@ -19,43 +20,69 @@ def get_num_cus() -> int: return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count -# AMD gfx950 compiler knobs. -_KNOBS_SET = False +# ── AMD compiler knobs (scoped) ── +# The AMD codegen toggles live on ``triton.knobs.amd`` and are process-global: they +# are read when a kernel is compiled, so leaving them set would silently change the +# codegen of every unrelated Triton kernel compiled afterwards (cast, gmm, MXFP8, +# ...). We therefore apply them only around our own launches and restore the prior +# state on exit. Triton's setter keeps the backing env var in sync with the +# attribute, so we snapshot/restore both. +_AMD_KNOB_ENV = { + "use_async_copy": "TRITON_HIP_USE_ASYNC_COPY", + "scalarize_packed_fops": "AMDGCN_SCALARIZE_PACKED_FOPS", + "use_block_pingpong": "TRITON_HIP_USE_BLOCK_PINGPONG", +} -def _set_triton_knobs_gfx950() -> None: - """Force-on AMD compiler knobs for gfx950 (async_copy, block_pingpong, scalarize).""" - global _KNOBS_SET - if _KNOBS_SET: - return - _KNOBS_SET = True - os.environ["TRITON_HIP_USE_ASYNC_COPY"] = "1" - os.environ["AMDGCN_SCALARIZE_PACKED_FOPS"] = "1" - os.environ["TRITON_HIP_USE_BLOCK_PINGPONG"] = "1" - if hasattr(triton, "knobs") and hasattr(triton.knobs, "amd"): - triton.knobs.amd.use_async_copy = True - triton.knobs.amd.scalarize_packed_fops = True - triton.knobs.amd.use_block_pingpong = True +def _amd_knob_overrides(*, is_tn: bool) -> dict: + """Knob values to apply for the enclosed GEMM launch. + gfx950 (CDNA4): force async_copy / scalarize / block_pingpong on. + gfx942 (CDNA3): async_copy / scalarize help NT/NN but regress TN/wgrad + ~5-8%, so gate them on the GEMM layout (block_pingpong left at its default). + """ + if is_cdna3(): + enable = not is_tn + return {"use_async_copy": enable, "scalarize_packed_fops": enable} + return { + "use_async_copy": True, + "scalarize_packed_fops": True, + "use_block_pingpong": True, + } -def _set_triton_knobs_gfx942(enable: bool = True): - """Set AMD Triton knobs on gfx942 (CDNA3). +@contextlib.contextmanager +def _amd_compiler_knobs(*, is_tn: bool): + """Temporarily apply AMD Triton compiler knobs, restoring prior state on exit. - ``use_async_copy`` / ``scalarize_packed_fops`` help NT/NN but regress - TN/wgrad ~5-8% on gfx942, so callers pass ``enable`` from layout. + Scoped rather than set once because the knobs are process-global and read at + kernel compile time; leaving them set would leak into other kernels' codegen. """ - if hasattr(triton, "knobs") and hasattr(triton.knobs, "amd"): - triton.knobs.amd.use_async_copy = enable - triton.knobs.amd.scalarize_packed_fops = enable - - -def _apply_amd_compiler_knobs(*, is_tn: bool) -> None: - """gfx942: knobs from GEMM layout. Else (gfx950): always-on gfx950 knobs.""" - if is_cdna3(): - _set_triton_knobs_gfx942(enable=not is_tn) - else: - _set_triton_knobs_gfx950() + amd = getattr(getattr(triton, "knobs", None), "amd", None) + if amd is None: + yield + return + overrides = _amd_knob_overrides(is_tn=is_tn) + saved = { + name: (name in amd.__dict__, amd.__dict__.get(name), os.environ.get(env_key)) + for name, env_key in _AMD_KNOB_ENV.items() + if name in overrides + } + try: + for name, value in overrides.items(): + setattr(amd, name, value) + yield + finally: + for name, (had_override, override_val, env_val) in saved.items(): + if had_override: + amd.__dict__[name] = override_val + else: + amd.__dict__.pop(name, None) + env_key = _AMD_KNOB_ENV[name] + if env_val is None: + os.environ.pop(env_key, None) + else: + os.environ[env_key] = env_val NUM_XCDS = 8 @@ -260,6 +287,7 @@ def _bwd_autotune_configs(): # Blockwise grouped FP8 kernel and public entrypoint + @triton.autotune(configs=_get_grouped_blockwise_autotune_configs(), key=["G", "N", "K"]) @triton.jit() def _grouped_blockwise_fp8_persistent_gemm_kernel( @@ -371,8 +399,13 @@ def _grouped_blockwise_fp8_persistent_gemm_kernel( # ── K-loop with block-wise scaling (EVEN_K pattern) ── loop_k = tl.cdiv(K, BLOCK_SIZE_K) if not EVEN_K: + # Full K-blocks only; the ragged tail is handled below. This can be 0 + # (K < BLOCK_SIZE_K), so only assume non-negativity. loop_k -= 1 - tl.assume(loop_k > 1) + tl.assume(loop_k >= 0) + else: + # K is a positive multiple of BLOCK_SIZE_K, so there is >= 1 full block. + tl.assume(loop_k >= 1) acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype) @@ -400,7 +433,9 @@ def _grouped_blockwise_fp8_persistent_gemm_kernel( if not EVEN_K: # ── Last partial K-block (masked) ── rk_last = loop_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) - A_LAST = A + m_start_g * stride_am + rm[:, None] * stride_am + rk_last[None, :] * stride_ak + A_LAST = ( + A + m_start_g * stride_am + rm[:, None] * stride_am + rk_last[None, :] * stride_ak + ) B_LAST = B + group_offset_b + rk_last[:, None] * stride_bk + rn[None, :] * stride_bn if stride_ak == 1: A_LAST = tl.multiple_of(A_LAST, (1, 16)) @@ -434,7 +469,6 @@ def _grouped_blockwise_fp8_persistent_gemm_kernel( # ═══════════════════════════════════════════════════════════════════════════════ - def grouped_gemm_fp8_blockwise_triton_kernel( a: torch.Tensor, b: torch.Tensor, @@ -461,9 +495,6 @@ def grouped_gemm_fp8_blockwise_triton_kernel( Returns: [M_total, N] output in out_dtype. """ - # trans_a is always False here: NT if trans_b else NN (never TN). - _apply_amd_compiler_knobs(is_tn=False) - assert a.ndim == 2, f"a must be 2D, got {a.shape}" assert b.ndim == 3, f"b must be 3D, got {b.shape}" assert b_scales.ndim == 3, f"b_scales must be 3D, got {b_scales.shape}" @@ -527,22 +558,23 @@ def _launch(c_out, offs): # Triton's autotune key is (G, N, K) only. Prime once with balanced offs # so the cached config is not locked to the first uneven MoE routing. warm_key = (G, N, K, even_k) - if warm_key not in _grouped_blockwise_warmed: - _grouped_blockwise_warmed.add(warm_key) - per = M_total // max(G, 1) - bal_offs = torch.arange(G + 1, device=group_offs.device, dtype=group_offs.dtype) * per - bal_offs[-1] = M_total - _launch(torch.empty_like(out), bal_offs) - - _launch(out, group_offs) + # trans_a is always False here: forward/dgrad are NT/NN (never TN). Scope the + # AMD knobs to these launches so they don't leak into other Triton kernels. + with _amd_compiler_knobs(is_tn=False): + if warm_key not in _grouped_blockwise_warmed: + _grouped_blockwise_warmed.add(warm_key) + per = M_total // max(G, 1) + bal_offs = torch.arange(G + 1, device=group_offs.device, dtype=group_offs.dtype) * per + bal_offs[-1] = M_total + _launch(torch.empty_like(out), bal_offs) + + _launch(out, group_offs) return out # ── Blockwise FP8 Variable-K Backward Public API ── - - @triton.autotune( configs=_bwd_autotune_configs(), key=["G", "OUT_M", "OUT_N"], @@ -638,8 +670,12 @@ def _grouped_blockwise_fp8_variable_k_gemm_kernel( rn = tl.max_contiguous(tl.multiple_of(rn, BLOCK_SIZE_N), BLOCK_SIZE_N) # ── Base pointers ── - LHS_BASE = LHS + m_start * stride_lhs_m + rm[:, None] * stride_lhs_n + rk[None, :] * stride_lhs_m - RHS_BASE = RHS + m_start * stride_rhs_m + rk[:, None] * stride_rhs_m + rn[None, :] * stride_rhs_n + LHS_BASE = ( + LHS + m_start * stride_lhs_m + rm[:, None] * stride_lhs_n + rk[None, :] * stride_lhs_m + ) + RHS_BASE = ( + RHS + m_start * stride_rhs_m + rk[:, None] * stride_rhs_m + rn[None, :] * stride_rhs_n + ) scale_row_start = m_start // BLOCK_SIZE_K @@ -685,9 +721,14 @@ def _grouped_blockwise_fp8_variable_k_gemm_kernel( # ── Store output ── rm_s = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) rn_s = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - rn_s = tl.max_contiguous(tl.multiple_of(rn_s % OUT_N, BLOCK_SIZE_N), BLOCK_SIZE_N) + rn_s = tl.max_contiguous(tl.multiple_of(rn_s, BLOCK_SIZE_N), BLOCK_SIZE_N) c_mask = (rm_s[:, None] < OUT_M) & (rn_s[None, :] < OUT_N) - C_ = C + group_idx.to(tl.int64) * stride_cg + rm_s[:, None] * stride_cm + rn_s[None, :] * stride_cn + C_ = ( + C + + group_idx.to(tl.int64) * stride_cg + + rm_s[:, None] * stride_cm + + rn_s[None, :] * stride_cn + ) c = acc.to(C.type.element_ty) if ACCUMULATE: c += tl.load(C_, mask=c_mask, other=0) @@ -697,7 +738,6 @@ def _grouped_blockwise_fp8_variable_k_gemm_kernel( # ── Blockwise FP8 Forward Public API ── - def grouped_gemm_fp8_blockwise_variable_k_triton_kernel( lhs: torch.Tensor, rhs: torch.Tensor, @@ -716,9 +756,6 @@ def grouped_gemm_fp8_blockwise_variable_k_triton_kernel( Output: [G, OUT_M, OUT_N]. When ``out`` is provided the kernel writes in-place (and adds into it if ``accumulate``). ``out_dtype`` is ignored in that case. """ - # Variable-K wgrad is TN (lhs^T @ rhs). - _apply_amd_compiler_knobs(is_tn=True) - assert lhs.ndim == 2 and rhs.ndim == 2 assert lhs.shape[0] == rhs.shape[0] OUT_M = lhs.shape[1] @@ -729,9 +766,11 @@ def grouped_gemm_fp8_blockwise_variable_k_triton_kernel( assert not accumulate, "accumulate=True requires an existing out tensor" out = torch.empty((G, OUT_M, OUT_N), device=lhs.device, dtype=out_dtype) else: - assert out.shape == (G, OUT_M, OUT_N), ( - f"out must be {(G, OUT_M, OUT_N)}, got {tuple(out.shape)}" - ) + assert out.shape == ( + G, + OUT_M, + OUT_N, + ), f"out must be {(G, OUT_M, OUT_N)}, got {tuple(out.shape)}" num_sms = get_num_cus() def _launch(c_out, offs, accumulate_flag): @@ -768,13 +807,16 @@ def _launch(c_out, offs, accumulate_flag): # Autotune key is (G, OUT_M, OUT_N). Prime on a scratch buffer with # balanced padded offs so accumulate=True never benchmarks into ``out``. warm_key = (G, OUT_M, OUT_N, out.dtype, accumulate) - if warm_key not in _grouped_blockwise_vk_warmed: - _grouped_blockwise_vk_warmed.add(warm_key) - m_padded = lhs.shape[0] - per = max((m_padded // max(G, 1)) // 128 * 128, 128) - bal_offs = torch.arange(G + 1, device=group_offs.device, dtype=group_offs.dtype) * per - bal_offs[-1] = m_padded - _launch(torch.zeros_like(out), bal_offs, False) - - _launch(out, group_offs, accumulate) + # Variable-K wgrad is TN (lhs^T @ rhs). Scope the AMD knobs to these launches + # so they don't leak into other Triton kernels. + with _amd_compiler_knobs(is_tn=True): + if warm_key not in _grouped_blockwise_vk_warmed: + _grouped_blockwise_vk_warmed.add(warm_key) + m_padded = lhs.shape[0] + per = max((m_padded // max(G, 1)) // 128 * 128, 128) + bal_offs = torch.arange(G + 1, device=group_offs.device, dtype=group_offs.dtype) * per + bal_offs[-1] = m_padded + _launch(torch.zeros_like(out), bal_offs, False) + + _launch(out, group_offs, accumulate) return out diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py b/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py index 73fa02e264..21eb2d5d9d 100644 --- a/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py +++ b/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py @@ -58,7 +58,9 @@ def quant_fp8_blockwise_kernel( x_tile = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) x_tile_abs = tl.abs(x_tile) - x_fp8_tile, x_scales_tile = compute_scale_and_quant(x_tile, x_tile_abs, AXIS, FP8_MAX, ROUND_POW2) + x_fp8_tile, x_scales_tile = compute_scale_and_quant( + x_tile, x_tile_abs, AXIS, FP8_MAX, ROUND_POW2 + ) x_fp8_ptrs = x_fp8_ptr + offs_m[:, None] * N + offs_n[None, :] tl.store(x_fp8_ptrs, x_fp8_tile.to(x_fp8_ptr.dtype.element_ty), mask=mask) @@ -96,8 +98,12 @@ def quant_fp8_blockwise_dual_kernel( x_tile = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) x_tile_abs = tl.abs(x_tile) - x_fp8_row_tile, x_scales_row_tile = compute_scale_and_quant(x_tile, x_tile_abs, 1, FP8_MAX, ROUND_POW2) - x_fp8_col_tile, x_scales_col_tile = compute_scale_and_quant(x_tile, x_tile_abs, 0, FP8_MAX, ROUND_POW2) + x_fp8_row_tile, x_scales_row_tile = compute_scale_and_quant( + x_tile, x_tile_abs, 1, FP8_MAX, ROUND_POW2 + ) + x_fp8_col_tile, x_scales_col_tile = compute_scale_and_quant( + x_tile, x_tile_abs, 0, FP8_MAX, ROUND_POW2 + ) x_fp8_row_ptrs = x_fp8_row_ptr + offs_m[:, None] * N + offs_n[None, :] tl.store(x_fp8_row_ptrs, x_fp8_row_tile.to(x_fp8_row_ptr.dtype.element_ty), mask=mask) @@ -159,10 +165,10 @@ def quant_fp8_blockwise_for_weight_kernel( @triton.jit def quant_fp8_blockwise_segment_m_kernel( - x_ptr, # Input tensor [M_in, N] - x_fp8_ptr, # Output tensor [M_out, N] (padded) - x_scales_ptr, # Output scales [M_out // BLOCK_SIZE, N] - group_offs_ptr, # Original group offsets [B+1] + x_ptr, # Input tensor [M_in, N] + x_fp8_ptr, # Output tensor [M_out, N] (padded) + x_scales_ptr, # Output scales [M_out // BLOCK_SIZE, N] + group_offs_ptr, # Original group offsets [B+1] padded_group_offs_ptr, # Padded group offsets [B+1] N, num_groups, @@ -231,7 +237,9 @@ def quant_fp8_blockwise_segment_m_kernel( # ----------------------------------------------------------------------------- -def quantize_fp8_blockwise_dual(x: torch.Tensor, dtype: torch.dtype, block_size: int = 128, pow2: bool = False): +def quantize_fp8_blockwise_dual( + x: torch.Tensor, dtype: torch.dtype, block_size: int = 128, pow2: bool = False +): """Blockwise-quantize a 2D tensor in BOTH row (1xB) and column (Bx1) modes in one pass. Returns (x_fp8_row, x_scales_row, x_fp8_col, x_scales_col); scales hold the @@ -242,20 +250,33 @@ def quantize_fp8_blockwise_dual(x: torch.Tensor, dtype: torch.dtype, block_size: fp8_max = torch.finfo(dtype).max x_fp8_row = torch.empty((M, N), dtype=dtype, device=x.device) - x_scales_row = torch.empty((M, triton.cdiv(N, block_size)), dtype=torch.float32, device=x.device) + x_scales_row = torch.empty( + (M, triton.cdiv(N, block_size)), dtype=torch.float32, device=x.device + ) x_fp8_col = torch.empty((M, N), dtype=dtype, device=x.device) - x_scales_col = torch.empty((triton.cdiv(M, block_size), N), dtype=torch.float32, device=x.device) + x_scales_col = torch.empty( + (triton.cdiv(M, block_size), N), dtype=torch.float32, device=x.device + ) grid = (triton.cdiv(M, block_size), triton.cdiv(N, block_size)) quant_fp8_blockwise_dual_kernel[grid]( - x, x_fp8_row, x_scales_row, x_fp8_col, x_scales_col, M, N, - BLOCK_SIZE=block_size, FP8_MAX=fp8_max, + x, + x_fp8_row, + x_scales_row, + x_fp8_col, + x_scales_col, + M, + N, + BLOCK_SIZE=block_size, + FP8_MAX=fp8_max, ROUND_POW2=pow2, ) return x_fp8_row, x_scales_row, x_fp8_col, x_scales_col -def quantize_fp8_blockwise(x: torch.Tensor, dtype: torch.dtype, axis: int, block_size: int = 128, pow2: bool = False): +def quantize_fp8_blockwise( + x: torch.Tensor, dtype: torch.dtype, axis: int, block_size: int = 128, pow2: bool = False +): """Single-direction blockwise quantize. axis=1 -> rowwise (1xB), axis=0 -> colwise (Bx1).""" assert x.is_contiguous() and x.dim() == 2, "Input must be 2D and contiguous" M, N = x.shape @@ -268,13 +289,22 @@ def quantize_fp8_blockwise(x: torch.Tensor, dtype: torch.dtype, axis: int, block grid = (triton.cdiv(M, block_size), triton.cdiv(N, block_size)) quant_fp8_blockwise_kernel[grid]( - x, x_fp8, scales, M, N, BLOCK_SIZE=block_size, FP8_MAX=fp8_max, AXIS=axis, + x, + x_fp8, + scales, + M, + N, + BLOCK_SIZE=block_size, + FP8_MAX=fp8_max, + AXIS=axis, ROUND_POW2=pow2, ) return x_fp8, scales -def quantize_fp8_blockwise_weight(w: torch.Tensor, dtype: torch.dtype, block_size: int = 128, pow2: bool = False): +def quantize_fp8_blockwise_weight( + w: torch.Tensor, dtype: torch.dtype, block_size: int = 128, pow2: bool = False +): """128x128 weight blockwise quantize. w is [B, M, N] (or [M, N], promoted to B=1).""" squeeze = False if w.dim() == 2: @@ -287,11 +317,18 @@ def quantize_fp8_blockwise_weight(w: torch.Tensor, dtype: torch.dtype, block_siz w_fp8 = torch.empty((B, M, N), dtype=dtype, device=w.device) w_scales = torch.empty( (B, triton.cdiv(M, block_size), triton.cdiv(N, block_size)), - dtype=torch.float32, device=w.device, + dtype=torch.float32, + device=w.device, ) grid = (B, triton.cdiv(M, block_size), triton.cdiv(N, block_size)) quant_fp8_blockwise_for_weight_kernel[grid]( - w, w_fp8, w_scales, M, N, BLOCK_SIZE=block_size, FP8_MAX=fp8_max, + w, + w_fp8, + w_scales, + M, + N, + BLOCK_SIZE=block_size, + FP8_MAX=fp8_max, ROUND_POW2=pow2, ) if squeeze: @@ -299,7 +336,9 @@ def quantize_fp8_blockwise_weight(w: torch.Tensor, dtype: torch.dtype, block_siz return w_fp8, w_scales -def quantize_fp8_blockwise_segment_m(x, dtype, block_size, group_lens, group_offs, pow2: bool = False): +def quantize_fp8_blockwise_segment_m( + x, dtype, block_size, group_lens, group_offs, pow2: bool = False +): """Colwise blockwise quantize with per-segment (MoE group) M padding. Returns (x_fp8 [M_pad, N], x_scales, var_k_group_lens [B], var_k_group_offs [B+1]). @@ -317,12 +356,21 @@ def quantize_fp8_blockwise_segment_m(x, dtype, block_size, group_lens, group_off m_padded_max = M + num_groups * block_size x_fp8 = torch.zeros((m_padded_max, N), dtype=dtype, device=x.device) - x_scales = torch.zeros((triton.cdiv(m_padded_max, block_size), N), dtype=torch.float32, device=x.device) + x_scales = torch.zeros( + (triton.cdiv(m_padded_max, block_size), N), dtype=torch.float32, device=x.device + ) grid = (triton.cdiv(m_padded_max, block_size), triton.cdiv(N, block_size)) quant_fp8_blockwise_segment_m_kernel[grid]( - x, x_fp8, x_scales, group_offs, var_k_group_offs, N, num_groups, - BLOCK_SIZE=block_size, FP8_MAX=fp8_max, + x, + x_fp8, + x_scales, + group_offs, + var_k_group_offs, + N, + num_groups, + BLOCK_SIZE=block_size, + FP8_MAX=fp8_max, ROUND_POW2=pow2, ) return x_fp8, x_scales, var_k_group_lens, var_k_group_offs From 01f95ae54edee70c49d25bab9d52aafb5133f62e Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Wed, 26 Aug 2026 23:55:35 +0000 Subject: [PATCH 11/13] Support fp8_model_params in ROCm blockwise FP8 grouped GEMM - Allow already-quantized blockwise FP8 weights (`fp8_model_params`) to be consumed directly in the grouped-linear blockwise Triton path, avoiding dequantize/re-quantize double quantization. - Add `Float8BlockwiseQTensor` wrappers for grouped weight and activation operands, and a zero-copy packed 2D view over per-expert weight tensors. - Replace the separate rowwise/colwise dual-quantize kernel with a unified grouped kernel that can emit both rowwise activation and segment-padded columnwise wgrad operands in one pass. - Rename the raw grouped GEMM kernels to `..._raw` and add public wrappers that extract FP8 data/scales from `Float8BlockwiseQTensor`. - Add `_vk_group_offs` to `Float8BlockwiseQTensorStorage` for variable-K grouped wgrad segment offsets. - Update tests to use the new QTensor APIs and remove the `fp8_weights` rejection gate. --- tests/pytorch/test_grouped_linear.py | 20 - .../triton_kernels/test_blockwise_fp8.py | 55 +-- .../pytorch/module/grouped_linear.py | 178 +++++-- .../float8_blockwise_tensor_storage.py | 8 + .../blockwise_fp8_grouped_gemm.py | 111 ++++- .../triton_kernels/blockwise_quantize.py | 454 +++++++++++++----- 6 files changed, 614 insertions(+), 212 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index aedf97f5a9..1be232e9a3 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -48,7 +48,6 @@ skip_unsupported_backward_override, ) from triton_kernels.test_common import get_tolerances -from transformer_engine.pytorch.triton_kernels.common import get_torch_e4m3_type # Only run FP8 tests on supported devices. fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) @@ -436,28 +435,9 @@ def test_blockwise_fp8_gate_requires_128_aligned_features( actual_m_splits=None, in_features=in_features, out_features=out_features, - fp8_weights=False, ) assert supported is expected - # Quantized weight params (fp8_model_params) must always fall back to avoid - # dequantize -> re-quantize double quantization. - supported_fp8_weights = _GroupedLinear._is_blockwise_fp8_triton_grouped_gemm_supported( - fp8=True, - recipe=recipe.Float8BlockScaling(), - use_bias=False, - backward_override=None, - cpu_offloading=False, - save_original_input=False, - debug=False, - unpad_output=False, - actual_m_splits=None, - in_features=in_features, - out_features=out_features, - fp8_weights=True, - ) - assert supported_fp8_weights is False - @pytest.mark.skipif( torch.cuda.get_device_capability() != (9, 0), diff --git a/tests/pytorch/triton_kernels/test_blockwise_fp8.py b/tests/pytorch/triton_kernels/test_blockwise_fp8.py index 21cde0ca7c..5a0d5d3176 100644 --- a/tests/pytorch/triton_kernels/test_blockwise_fp8.py +++ b/tests/pytorch/triton_kernels/test_blockwise_fp8.py @@ -11,13 +11,12 @@ from transformer_engine.pytorch.triton_kernels.common import get_torch_e4m3_type from transformer_engine.pytorch.triton_kernels.blockwise_quantize import ( quantize_fp8_blockwise, - quantize_fp8_blockwise_dual, quantize_fp8_blockwise_weight, quantize_fp8_blockwise_segment_m, ) from transformer_engine.pytorch.triton_kernels.blockwise_fp8_grouped_gemm import ( - grouped_gemm_fp8_blockwise_triton_kernel, - grouped_gemm_fp8_blockwise_variable_k_triton_kernel, + _grouped_gemm_fp8_blockwise_raw, + _grouped_gemm_fp8_blockwise_variable_k_raw, ) pytestmark = pytest.mark.skipif(not IS_HIP_EXTENSION, reason="ROCm Triton blockwise kernels only") @@ -43,7 +42,9 @@ def _group_offs(splits, device="cuda"): return offs -def _ref_rowwise_quantize(x: torch.Tensor, dtype: torch.dtype, block: int = BLOCK, pow2: bool = False): +def _ref_rowwise_quantize( + x: torch.Tensor, dtype: torch.dtype, block: int = BLOCK, pow2: bool = False +): m, n = x.shape fp8_max = torch.finfo(dtype).max nb = _cdiv(n, block) @@ -58,7 +59,9 @@ def _ref_rowwise_quantize(x: torch.Tensor, dtype: torch.dtype, block: int = BLOC return q, (1.0 / scale).contiguous() -def _ref_colwise_quantize(x: torch.Tensor, dtype: torch.dtype, block: int = BLOCK, pow2: bool = False): +def _ref_colwise_quantize( + x: torch.Tensor, dtype: torch.dtype, block: int = BLOCK, pow2: bool = False +): m, n = x.shape fp8_max = torch.finfo(dtype).max mb = _cdiv(m, block) @@ -73,7 +76,9 @@ def _ref_colwise_quantize(x: torch.Tensor, dtype: torch.dtype, block: int = BLOC return q, (1.0 / scale).contiguous() -def _ref_weight_quantize(w: torch.Tensor, dtype: torch.dtype, block: int = BLOCK, pow2: bool = False): +def _ref_weight_quantize( + w: torch.Tensor, dtype: torch.dtype, block: int = BLOCK, pow2: bool = False +): g, m, n = w.shape fp8_max = torch.finfo(dtype).max mb, nb = _cdiv(m, block), _cdiv(n, block) @@ -115,30 +120,6 @@ def test_quantize_fp8_blockwise(shape, dtype, axis, pow2): torch.testing.assert_close(dq, dq_ref, atol=0.10, rtol=0.10) -@pytest.mark.parametrize("shape", [(128, 256), (200, 128)]) -@pytest.mark.parametrize("dtype", IN_DTYPES, ids=str) -@pytest.mark.parametrize("pow2", [False, True]) -def test_quantize_fp8_blockwise_dual(shape, dtype, pow2): - x = torch.randn(*shape, dtype=dtype, device="cuda") - q_row, s_row, q_col, s_col = quantize_fp8_blockwise_dual(x, FP8_DTYPE, block_size=BLOCK, pow2=pow2) - q_row_ref, s_row_ref = _ref_rowwise_quantize(x, FP8_DTYPE, pow2=pow2) - q_col_ref, s_col_ref = _ref_colwise_quantize(x, FP8_DTYPE, pow2=pow2) - torch.testing.assert_close(s_row, s_row_ref, atol=1e-5, rtol=1e-4) - torch.testing.assert_close(s_col, s_col_ref, atol=1e-5, rtol=1e-4) - torch.testing.assert_close( - _dequant_rowwise(q_row, s_row, x.shape[1]), - _dequant_rowwise(q_row_ref, s_row_ref, x.shape[1]), - atol=0.10, - rtol=0.10, - ) - torch.testing.assert_close( - _dequant_colwise(q_col, s_col, x.shape[0]), - _dequant_colwise(q_col_ref, s_col_ref, x.shape[0]), - atol=0.10, - rtol=0.10, - ) - - @pytest.mark.parametrize("shape", [(2, 256, 256), (3, 128, 256), (1, 200, 192), (256, 256)]) @pytest.mark.parametrize("dtype", IN_DTYPES, ids=str) @pytest.mark.parametrize("pow2", [False, True]) @@ -222,7 +203,7 @@ def test_grouped_gemm_fp8_blockwise_matches_dequant_ref(splits, n, k, trans_b, i b_fp8, b_s = quantize_fp8_blockwise_weight(b, FP8_DTYPE, block_size=BLOCK) offs = _group_offs(splits) - out = grouped_gemm_fp8_blockwise_triton_kernel( + out = _grouped_gemm_fp8_blockwise_raw( a_fp8, b_fp8, a_s, b_s, offs, trans_b=trans_b, out_dtype=out_dtype ) assert out.dtype == out_dtype @@ -270,11 +251,9 @@ def test_variable_k_wgrad(splits, n, k, accumulate, in_dtype, out_dtype): go_col, go_s, _, vk_offs = quantize_fp8_blockwise_segment_m( dy, FP8_DTYPE, BLOCK, group_lens, group_offs ) - x_col, x_s, _, _ = quantize_fp8_blockwise_segment_m( - x, FP8_DTYPE, BLOCK, group_lens, group_offs - ) + x_col, x_s, _, _ = quantize_fp8_blockwise_segment_m(x, FP8_DTYPE, BLOCK, group_lens, group_offs) - fresh = grouped_gemm_fp8_blockwise_variable_k_triton_kernel( + fresh = _grouped_gemm_fp8_blockwise_variable_k_raw( go_col, x_col, go_s, x_s, vk_offs, out_dtype=out_dtype, accumulate=False ) assert fresh.dtype == out_dtype @@ -293,7 +272,7 @@ def test_variable_k_wgrad(splits, n, k, accumulate, in_dtype, out_dtype): if accumulate: main_grad = torch.randn(g, n, k, dtype=out_dtype, device="cuda") expected = main_grad.clone() - grouped_gemm_fp8_blockwise_variable_k_triton_kernel( + _grouped_gemm_fp8_blockwise_variable_k_raw( go_col, x_col, go_s, @@ -302,4 +281,6 @@ def test_variable_k_wgrad(splits, n, k, accumulate, in_dtype, out_dtype): out=main_grad, accumulate=True, ) - torch.testing.assert_close(main_grad.float(), (expected + fresh).float(), atol=1e-3, rtol=1e-4) + torch.testing.assert_close( + main_grad.float(), (expected + fresh).float(), atol=1e-3, rtol=1e-4 + ) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index d49be027dd..b0f353519e 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -438,6 +438,32 @@ def _expert_weights_as_3d(weights, dtype): return packed if packed.is_contiguous() else packed.contiguous() return torch.stack([wt.to(dtype).contiguous() for wt in weights], 0).contiguous() + @staticmethod + def _grouped_weight_blockwise_qtensor(weights, activation_dtype, dt, pow2_w): + """Packed 2D ``[G*N, K]`` grouped-weight ``Float8BlockwiseQTensor`` from a list + of per-expert blockwise ``Float8BlockwiseQTensor`` (``fp8_model_params``). + + When the per-expert ``_rowwise_data`` / ``_rowwise_scale_inv`` are consecutive + slices of one contiguous buffer (``single_grouped_weight``), the packed + ``[G*N, K]`` data and ``[G*(N/128), K/128]`` scales are reconstructed as a + zero-copy view. Discrete params fall back to a ``torch.stack`` copy. + """ + from ..triton_kernels.blockwise_fp8_grouped_gemm import _stack_weight_qtensors + from ..triton_kernels.blockwise_quantize import ( + wrap_fp8_blockwise_grouped_weight_qtensor, + ) + from ..triton_kernels.common import te_dtype_to_torch_dtype + + # Zero-copy [G, N, K] / [G, bn, bk] views when the per-expert payloads are + # consecutive slices of one buffer; the wrap builder flattens them to 2D. + b_fp8 = _GroupedLinear._packed_3d_view([w._rowwise_data for w in weights]) + b_scale = _GroupedLinear._packed_3d_view([w._rowwise_scale_inv for w in weights]) + if b_fp8 is None or b_scale is None: + b_fp8, b_scale = _stack_weight_qtensors(weights, te_dtype_to_torch_dtype(dt)) + return wrap_fp8_blockwise_grouped_weight_qtensor( + b_fp8, b_scale, activation_dtype, dt, pow2=pow2_w + ) + @staticmethod def _handle_fused_wgrad(weight, main_grad): """Megatron DDP hook: mark wgrad as consumed and return a dummy (or None).""" @@ -463,7 +489,6 @@ def _is_blockwise_fp8_triton_grouped_gemm_supported( actual_m_splits, in_features, out_features, - fp8_weights, ) -> bool: """ROCm Triton blockwise FP8 grouped GEMM: HIP, env, Float8BlockScaling 1x128/128x128.""" return ( @@ -478,11 +503,6 @@ def _is_blockwise_fp8_triton_grouped_gemm_supported( recipe.grad_block_scaling_dim, ) == (1, 2, 1) - # The path quantizes high-precision weights itself. Quantized weight - # params (fp8_model_params) would be dequantized here and re-quantized - # by the kernel (double quantization), so fall back to the default path - # until the stored rowwise_data + scales are consumed directly. - and not fp8_weights # The forward/dgrad kernel applies one scalar B scale per N-tile # (BLOCK_SIZE_N == 128 == scale-block-N), so an N that is not a # multiple of 128 would apply the wrong scale to the tail tile. @@ -522,28 +542,31 @@ def _forward_blockwise_fp8_triton( Selected from :meth:`forward` under the ``Float8BlockScaling`` recipe when :meth:`_is_blockwise_fp8_triton_grouped_gemm_supported` is true. Activations are - quantized rowwise (1x128 along K), weights 128x128; the backward wgrad (see + quantized rowwise (1x128 along K) and weights 128x128, both wrapped as + :class:`Float8BlockwiseQTensor`. When the weights are already blockwise-quantized + params (``fp8_model_params``) their stored rowwise data + scales are reused + directly instead of re-quantizing. The backward wgrad (see :meth:`_backward_blockwise_fp8_triton`) uses a segment-padded columnwise operand - quantized from the original high-precision input (no double-quantization). + quantized from the original high-precision input. """ - from ..triton_kernels.common import te_dtype_to_torch_dtype from ..triton_kernels.blockwise_quantize import ( - quantize_fp8_blockwise, - quantize_fp8_blockwise_weight, - quantize_fp8_blockwise_segment_m, + quantize_fp8_blockwise_act_operands, + quantize_fp8_blockwise_grouped_weight_qtensor, ) from ..triton_kernels.blockwise_fp8_grouped_gemm import ( grouped_gemm_fp8_blockwise_triton_kernel, ) + from ..tensor.float8_blockwise_tensor import Float8BlockwiseQTensor num_gemms = len(m_splits) - # Resolve the FP8 torch dtype via the quantizer, normalizing through ``tex.DType`` - # so ``te_dtype_to_torch_dtype`` picks the arch-correct variant (fnuz on CDNA3). + # Resolve the FP8 dtype from the weight quantizer as a TE ``DType`` (the + # QTensor's native dtype). The Triton builders map it to the arch-correct + # torch FP8 dtype only at the kernel/``.view`` boundary. if weight_quantizers and weight_quantizers[0] is not None: - dt = te_dtype_to_torch_dtype(tex.DType(int(weight_quantizers[0].dtype))) + dt = tex.DType(int(weight_quantizers[0].dtype)) else: - dt = torch.float8_e4m3fn + dt = tex.DType.kFloat8E4M3 in_features = weights[0].size(-1) out_features = weights[0].size(0) device = inp.device @@ -558,25 +581,58 @@ def _forward_blockwise_fp8_triton( group_offs = torch.zeros(num_gemms + 1, dtype=torch.int64, device=device) group_offs[1:] = torch.cumsum(group_lens, 0) - w = _GroupedLinear._expert_weights_as_3d(weights, activation_dtype) + # Quantize the activation rowwise (1x128 along K) as a Float8BlockwiseQTensor. + # When gradients are needed the same pass also emits the segment-padded + # columnwise operand for the variable-K wgrad into ``qa``'s columnwise slots. + qa = quantize_fp8_blockwise_act_operands( + a, + dt, + 128, + group_lens, + group_offs, + rowwise=True, + columnwise=is_grad_enabled, + pow2=pow2_x, + ) - # Quantize: activation rowwise (1x128 along K), weights 128x128. - a_row, a_srow = quantize_fp8_blockwise(a, dt, axis=1, block_size=128, pow2=pow2_x) - b_fp8, b_scale = quantize_fp8_blockwise_weight(w, dt, block_size=128, pow2=pow2_w) + # Weight operand: a packed 2D [G*N, K] Float8BlockwiseQTensor. When the + # params are already blockwise-quantized (fp8_model_params) reuse their + # stored rowwise data + scales; otherwise quantize the high-precision + # weights once over the packed 3D view (zero-copy when + # ``single_grouped_weight``), no stack. + if isinstance(weights[0], Float8BlockwiseQTensor): + # fp8_model_params: reuse the stored rowwise data + scales directly (no + # re-quantization). + qw = _GroupedLinear._grouped_weight_blockwise_qtensor( + weights, activation_dtype, dt, pow2_w + ) + else: + w = _GroupedLinear._expert_weights_as_3d(weights, activation_dtype) + qw = quantize_fp8_blockwise_grouped_weight_qtensor(w, dt, pow2=pow2_w) # Forward GEMM: out[seg] = A[seg] @ W[g]^T (trans_b=True). out = grouped_gemm_fp8_blockwise_triton_kernel( - a_row, b_fp8, a_srow, b_scale, group_offs, trans_b=True, out_dtype=activation_dtype + qa, qw, group_offs, trans_b=True, out_dtype=activation_dtype ) if is_grad_enabled: ctx.use_blockwise_fp8_triton = True - # Segment-padded columnwise activation for the variable-K wgrad. - a_col, a_scol, _vk_lens, vk_offs = quantize_fp8_blockwise_segment_m( - a, dt, 128, group_lens, group_offs, pow2=pow2_x + # The segment-padded columnwise activation for the variable-K wgrad rides + # in ``qa``'s columnwise slots (produced by the fused pass above). Save + # those raw fields plus the packed weight QTensor's 2D fields; both + # rewrapped into Float8BlockwiseQTensors in backward. + ctx.save_for_backward( + qa._columnwise_data, + qa._columnwise_scale_inv, + qw._rowwise_data, + qw._rowwise_scale_inv, + group_lens, + group_offs, + qa._vk_group_offs, ) - ctx.save_for_backward(a_col, a_scol, b_fp8, b_scale, group_lens, group_offs, vk_offs) ctx.dt = dt + ctx.pow2_x = pow2_x + ctx.pow2_w = pow2_w ctx.pow2_grad = pow2_grad ctx.activation_dtype = activation_dtype ctx.num_gemms = num_gemms @@ -607,8 +663,9 @@ def _forward_blockwise_fp8_triton( def _backward_blockwise_fp8_triton(ctx, grad_output): """Backward path paired with :meth:`_forward_blockwise_fp8_triton`.""" from ..triton_kernels.blockwise_quantize import ( - quantize_fp8_blockwise, - quantize_fp8_blockwise_segment_m, + quantize_fp8_blockwise_act_operands, + wrap_fp8_blockwise_grouped_weight_qtensor, + wrap_fp8_blockwise_segment_m_qtensor, ) from ..triton_kernels.blockwise_fp8_grouped_gemm import ( grouped_gemm_fp8_blockwise_triton_kernel, @@ -619,17 +676,31 @@ def _backward_blockwise_fp8_triton(ctx, grad_output): dt = ctx.dt g_out = grad_output.reshape(-1, ctx.out_features).contiguous() + # Grad-output operands in one QTensor (``qdgrad``): rowwise for dgrad and/or + # segment-padded columnwise (variable-K) for wgrad. When both grads are + # needed they share a single fused pass over g_out; otherwise only the + # requested operand is produced. + qdgrad = quantize_fp8_blockwise_act_operands( + g_out, + dt, + 128, + group_lens, + group_offs, + rowwise=ctx.requires_dgrad, + columnwise=ctx.weight_requires_grad, + pow2=ctx.pow2_grad, + ) + # dgrad: dX[seg] = dY[seg] @ W[g] (trans_b=False) -> [M_total, K]. dgrad = None if ctx.requires_dgrad: - go_row, go_srow = quantize_fp8_blockwise( - g_out, dt, axis=1, block_size=128, pow2=ctx.pow2_grad + # Rewrap the saved packed weight fields as a 2D [G*N, K] QTensor. + qw = wrap_fp8_blockwise_grouped_weight_qtensor( + b_fp8, b_scale, ctx.activation_dtype, dt, pow2=ctx.pow2_w ) dgrad = grouped_gemm_fp8_blockwise_triton_kernel( - go_row, - b_fp8, - go_srow, - b_scale, + qdgrad, + qw, group_offs, trans_b=False, out_dtype=ctx.activation_dtype, @@ -638,8 +709,10 @@ def _backward_blockwise_fp8_triton(ctx, grad_output): # wgrad: dW[g] = dY[g]^T @ X[g] (variable-K) -> [G, N, K]. wgrad_list = [None] * ctx.num_gemms if ctx.weight_requires_grad: - go_col, go_scol, _l, _o = quantize_fp8_blockwise_segment_m( - g_out, dt, 128, group_lens, group_offs, pow2=ctx.pow2_grad + # qdgrad carries the segment-padded columnwise grad output (wgrad operand). + # Rewrap the saved activation columnwise fields as a columnwise-only QTensor. + a_operand = wrap_fp8_blockwise_segment_m_qtensor( + a_col, a_scol, vk_offs, ctx.activation_dtype, dt, pow2=ctx.pow2_x ) fuse = getattr(ctx, "fuse_wgrad_accumulation", False) origin_weights = [None] * ctx.num_gemms @@ -680,11 +753,8 @@ def _backward_blockwise_fp8_triton(ctx, grad_output): def grouped_gemm_wgrad(*_unused): dW = grouped_gemm_fp8_blockwise_variable_k_triton_kernel( - go_col, - a_col, - go_scol, - a_scol, - vk_offs, + qdgrad, + a_operand, out_dtype=out_dtype, out=packed_out, accumulate=accumulate and packed_out is not None and fuse, @@ -703,7 +773,10 @@ def grouped_gemm_wgrad(*_unused): wgrad_store = getattr(ctx, "wgrad_store", None) if wgrad_store is not None and wgrad_store.delay_wgrad_compute(): # tensor_list[2] is the wgrad buffers backward_dw assigns to .grad. - wgrad_store.put([go_col, a_col, wgrad_list], grouped_gemm_wgrad) + wgrad_store.put( + [qdgrad._columnwise_data, a_operand._columnwise_data, wgrad_list], + grouped_gemm_wgrad, + ) else: grouped_gemm_wgrad() @@ -799,7 +872,6 @@ def forward( actual_m_splits=actual_m_splits, in_features=weights[0].size(-1), out_features=weights[0].size(0), - fp8_weights=isinstance(weights[0], QuantizedTensorStorage), ): return _GroupedLinear._forward_blockwise_fp8_triton( ctx, @@ -2394,6 +2466,24 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage if weight_tensors is None: # TODO(ksivaman): Remove this after GEMM integration. weight_tensors = grouped_weight.split_into_quantized_tensors() + # Members are views, not the Parameter, so mirror the grad state the + # autograd Function reads off ``weights[i]``: ``requires_grad`` (gates + # wgrad), ``main_grad`` (per-expert views into the grouped + # fuse-accumulation buffer), and ``overwrite_main_grad``. + want_grad = grouped_weight.requires_grad + main_grad = getattr(grouped_weight, "main_grad", None) + per_expert_main_grad = None + if main_grad is not None: + per_expert_main_grad = main_grad.view( + self.num_gemms, self.out_features, self.in_features + ) + for i, w in enumerate(weight_tensors): + if w.requires_grad != want_grad: + w.requires_grad_(want_grad) + if per_expert_main_grad is not None: + w.main_grad = per_expert_main_grad[i] + if hasattr(grouped_weight, "overwrite_main_grad"): + w.overwrite_main_grad = grouped_weight.overwrite_main_grad else: weight_tensors = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] if not self.fp8 and any(isinstance(w, QuantizedTensorStorage) for w in weight_tensors): @@ -2414,6 +2504,10 @@ def _get_bias_tensors(self) -> List[torch.Tensor]: parts = grouped_bias.quantized_tensors if parts is None: parts = grouped_bias.split_into_quantized_tensors() + want_grad = grouped_bias.requires_grad + for p in parts: + if p.requires_grad != want_grad: + p.requires_grad_(want_grad) return [p.reshape(-1) for p in parts] return [getattr(self, f"bias{i}") for i in range(self.num_gemms)] diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index f7a3dae70b..3d4bab6d4a 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -1,3 +1,5 @@ +# This file was modified for portability to AMDGPU +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. @@ -34,6 +36,10 @@ class Float8BlockwiseQTensorStorage(QuantizedTensorStorage): _rowwise_scale_inv: Optional[torch.Tensor] _columnwise_scale_inv: Optional[torch.Tensor] _is_2D_scaled: bool + # ROCm-only: padded per-segment (MoE group) offsets for the variable-K grouped + # wgrad. Set only when the columnwise operand is segment-padded (columnwise data + # is [M_pad, N] rather than the standard transposed [N, M]); otherwise None. + _vk_group_offs: Optional[torch.Tensor] def __new__( cls, @@ -46,6 +52,7 @@ def __new__( is_2D_scaled: bool, *args, fake_dtype: Optional[torch.dtype] = None, + vk_group_offs: Optional[torch.Tensor] = None, **kwargs, ): if cls is Float8BlockwiseQTensorStorage: @@ -60,6 +67,7 @@ def __new__( instance._rowwise_scale_inv = rowwise_scale_inv instance._columnwise_scale_inv = columnwise_scale_inv instance._is_2D_scaled = is_2D_scaled + instance._vk_group_offs = vk_group_offs return instance diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py index 0f255763a1..948e637e3b 100644 --- a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py +++ b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py @@ -469,7 +469,7 @@ def _grouped_blockwise_fp8_persistent_gemm_kernel( # ═══════════════════════════════════════════════════════════════════════════════ -def grouped_gemm_fp8_blockwise_triton_kernel( +def _grouped_gemm_fp8_blockwise_raw( a: torch.Tensor, b: torch.Tensor, a_scales: torch.Tensor, @@ -477,6 +477,7 @@ def grouped_gemm_fp8_blockwise_triton_kernel( group_offs: torch.Tensor, trans_b: bool = True, out_dtype: torch.dtype = torch.bfloat16, + a_scales_pretransposed: bool = False, ) -> torch.Tensor: """Persistent grouped block-wise FP8 GEMM (CPU-sync-free) using Triton. @@ -486,7 +487,9 @@ def grouped_gemm_fp8_blockwise_triton_kernel( Args: a: [M_total, K] FP8 input (trans_a=False always). b: [G, N, K] (if trans_b=True) or [G, K, N] FP8 weights. - a_scales: [M_total, K//128] float32, block-wise scale for A. + a_scales: [M_total, K//128] float32, block-wise scale for A. When + ``a_scales_pretransposed`` is True this is already ``[K//128, M_total]`` + (the GEMM-ready layout) and the internal transpose is skipped. b_scales: [G, ceil(N/128), ceil(K/128)] or [G, ceil(K/128), ceil(N/128)] float32. group_offs: [G+1] int64 prefix sum of group lengths. trans_b: If True, b[g] is [N, K] (transposed). @@ -519,7 +522,7 @@ def grouped_gemm_fp8_blockwise_triton_kernel( stride_ak = a.stride(1) out = torch.empty((M_total, N), device=a.device, dtype=out_dtype) - A_scales_t = a_scales.T.contiguous() + A_scales_t = a_scales if a_scales_pretransposed else a_scales.T.contiguous() num_sms = get_num_cus() even_k = K % 128 == 0 @@ -572,6 +575,75 @@ def _launch(c_out, offs): return out +def _stack_weight_qtensors(weights, torch_fp8_dtype): + """Stack per-expert weight ``Float8BlockwiseQTensor`` into packed grouped + buffers: FP8 data ``[G, N, K]`` and scales ``[G, ceil(N/128), ceil(K/128)]``.""" + b_fp8 = torch.stack([w._rowwise_data.view(torch_fp8_dtype) for w in weights], 0) + b_scales = torch.stack([w._rowwise_scale_inv for w in weights], 0) + return b_fp8, b_scales + + +def grouped_gemm_fp8_blockwise_triton_kernel( + a, + b, + group_offs: torch.Tensor, + trans_b: bool = True, + out_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Grouped block-wise FP8 GEMM over ``Float8BlockwiseQTensor`` operands. + + ``a`` is a 1x128 rowwise-scaled ``Float8BlockwiseQTensor`` (``[M_total, K]``); + its stored ``rowwise_scale_inv`` is already in the ``[K//128, M_total]`` GEMM + layout. ``b`` is the weight operand in one of three forms: + + * a packed weight ``Float8BlockwiseQTensor`` -- either 2D ``[G*N, K]`` (scale + ``[G*(N/128), K/128]``, reshaped to ``[G, N, K]`` here using ``G`` from + ``group_offs``) or already 3D ``[G, N, K]``, + * a list of per-expert 128x128 weight ``Float8BlockwiseQTensor`` (stacked here), or + * a pre-stacked ``(b_fp8, b_scales)`` tuple of raw tensors (used by the backward + pass to reuse the buffers built during the forward pass). + + Raw FP8 fields are extracted here and handed to + :func:`_grouped_gemm_fp8_blockwise_raw`. + """ + import transformer_engine_torch as tex + from ..tensor.float8_blockwise_tensor import Float8BlockwiseQTensor + from .common import te_dtype_to_torch_dtype + + a_dt = te_dtype_to_torch_dtype(tex.DType(int(a._fp8_dtype))) + a_fp8 = a._rowwise_data.view(a_dt) + a_scales = a._rowwise_scale_inv # already [K//128, M_total] + + if isinstance(b, Float8BlockwiseQTensor): + b_dt = te_dtype_to_torch_dtype(tex.DType(int(b._fp8_dtype))) + b_fp8 = b._rowwise_data.view(b_dt) + b_scales = b._rowwise_scale_inv + if b_fp8.dim() == 2: + # Packed 2D [G*N, K] grouped weight: recover [G, N, K] using G from offs. + groups = group_offs.shape[0] - 1 + n = b_fp8.shape[0] // groups + b_fp8 = b_fp8.view(groups, n, b_fp8.shape[1]) + bn = b_scales.shape[0] // groups + b_scales = b_scales.view(groups, bn, b_scales.shape[1]) + elif isinstance(b, (list, tuple)) and isinstance(b[0], Float8BlockwiseQTensor): + b_dt = te_dtype_to_torch_dtype(tex.DType(int(b[0]._fp8_dtype))) + b_fp8, b_scales = _stack_weight_qtensors(b, b_dt) + else: + # Pre-stacked raw (b_fp8, b_scales); b_fp8 already an FP8-typed tensor. + b_fp8, b_scales = b + + return _grouped_gemm_fp8_blockwise_raw( + a_fp8, + b_fp8, + a_scales, + b_scales, + group_offs, + trans_b=trans_b, + out_dtype=out_dtype, + a_scales_pretransposed=True, + ) + + # ── Blockwise FP8 Variable-K Backward Public API ── @@ -739,6 +811,39 @@ def _grouped_blockwise_fp8_variable_k_gemm_kernel( def grouped_gemm_fp8_blockwise_variable_k_triton_kernel( + lhs, + rhs, + out_dtype: torch.dtype = torch.bfloat16, + out: torch.Tensor = None, + accumulate: bool = False, +) -> torch.Tensor: + """Variable-K grouped block-wise FP8 GEMM (backward, 1D+1D scaling) using Triton. + + ``lhs`` and ``rhs`` are ``Float8BlockwiseQTensor`` operands carrying the + segment-padded columnwise-quantized grad and activation in their columnwise + slots (``_columnwise_data`` ``[M_pad, N]``, ``_columnwise_scale_inv`` + ``[ceil(M_pad/128), N]``) plus the shared padded segment offsets in + ``_vk_group_offs``. Raw FP8 fields are extracted here (uint8 data viewed back to + the FP8 dtype) and handed to :func:`_grouped_gemm_fp8_blockwise_variable_k_raw`. + """ + import transformer_engine_torch as tex + from .common import te_dtype_to_torch_dtype + + lhs_dt = te_dtype_to_torch_dtype(tex.DType(int(lhs._fp8_dtype))) + rhs_dt = te_dtype_to_torch_dtype(tex.DType(int(rhs._fp8_dtype))) + return _grouped_gemm_fp8_blockwise_variable_k_raw( + lhs._columnwise_data.view(lhs_dt), + rhs._columnwise_data.view(rhs_dt), + lhs._columnwise_scale_inv, + rhs._columnwise_scale_inv, + lhs._vk_group_offs, + out_dtype=out_dtype, + out=out, + accumulate=accumulate, + ) + + +def _grouped_gemm_fp8_blockwise_variable_k_raw( lhs: torch.Tensor, rhs: torch.Tensor, lhs_scales: torch.Tensor, diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py b/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py index 21eb2d5d9d..9bc965c707 100644 --- a/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py +++ b/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py @@ -9,12 +9,21 @@ import torch import triton import triton.language as tl +import transformer_engine_torch as tex + +from .common import te_dtype_to_torch_dtype __all__ = [ "quantize_fp8_blockwise", - "quantize_fp8_blockwise_dual", "quantize_fp8_blockwise_weight", "quantize_fp8_blockwise_segment_m", + "quantize_fp8_blockwise_act_qtensor", + "quantize_fp8_blockwise_grouped_weight_qtensor", + "wrap_fp8_blockwise_grouped_weight_qtensor", + "quantize_fp8_blockwise_segment_m_operand", + "wrap_fp8_blockwise_segment_m_qtensor", + "quantize_fp8_blockwise_act_and_segment_m", + "quantize_fp8_blockwise_act_operands", ] @@ -75,53 +84,6 @@ def quant_fp8_blockwise_kernel( tl.store(x_scales_ptr + scale_offs, x_scales_tile_inv, mask=scale_mask) -@triton.jit -def quant_fp8_blockwise_dual_kernel( - x_ptr, - x_fp8_row_ptr, - x_scales_row_ptr, - x_fp8_col_ptr, - x_scales_col_ptr, - M, - N, - BLOCK_SIZE: tl.constexpr, - FP8_MAX: tl.constexpr, - ROUND_POW2: tl.constexpr, -): - pid_m = tl.program_id(axis=0) - pid_n = tl.program_id(axis=1) - offs_m = tl.cast(pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) - offs_n = tl.cast(pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) - mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) - - x_ptrs = x_ptr + offs_m[:, None] * N + offs_n[None, :] - x_tile = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) - x_tile_abs = tl.abs(x_tile) - - x_fp8_row_tile, x_scales_row_tile = compute_scale_and_quant( - x_tile, x_tile_abs, 1, FP8_MAX, ROUND_POW2 - ) - x_fp8_col_tile, x_scales_col_tile = compute_scale_and_quant( - x_tile, x_tile_abs, 0, FP8_MAX, ROUND_POW2 - ) - - x_fp8_row_ptrs = x_fp8_row_ptr + offs_m[:, None] * N + offs_n[None, :] - tl.store(x_fp8_row_ptrs, x_fp8_row_tile.to(x_fp8_row_ptr.dtype.element_ty), mask=mask) - - x_fp8_col_ptrs = x_fp8_col_ptr + offs_m[:, None] * N + offs_n[None, :] - tl.store(x_fp8_col_ptrs, x_fp8_col_tile.to(x_fp8_col_ptr.dtype.element_ty), mask=mask) - - row_scale_offs = offs_m * tl.cdiv(N, BLOCK_SIZE) + pid_n - row_scale_mask = offs_m < M - x_scales_row_tile_inv = tl.reshape(1.0 / x_scales_row_tile, BLOCK_SIZE) - tl.store(x_scales_row_ptr + row_scale_offs, x_scales_row_tile_inv, mask=row_scale_mask) - - col_scale_offs = pid_m * N + offs_n - col_scale_mask = offs_n < N - x_scales_col_tile_inv = tl.reshape(1.0 / x_scales_col_tile, BLOCK_SIZE) - tl.store(x_scales_col_ptr + col_scale_offs, x_scales_col_tile_inv, mask=col_scale_mask) - - @triton.jit def quant_fp8_blockwise_for_weight_kernel( w_ptr, @@ -164,10 +126,12 @@ def quant_fp8_blockwise_for_weight_kernel( @triton.jit -def quant_fp8_blockwise_segment_m_kernel( +def quant_fp8_blockwise_grouped_kernel( x_ptr, # Input tensor [M_in, N] - x_fp8_ptr, # Output tensor [M_out, N] (padded) - x_scales_ptr, # Output scales [M_out // BLOCK_SIZE, N] + x_fp8_row_ptr, # Rowwise output [M_in, N] (unused when ROWWISE is False) + x_scales_row_ptr, # Rowwise scales [M_in, ceil(N/BLOCK_SIZE)] (unused when ROWWISE is False) + x_fp8_col_ptr, # Colwise padded output [M_pad, N] (unused when COLUMNWISE is False) + x_scales_col_ptr, # Colwise scales [ceil(M_pad/BLOCK_SIZE), N] (unused when COLUMNWISE is False) group_offs_ptr, # Original group offsets [B+1] padded_group_offs_ptr, # Padded group offsets [B+1] N, @@ -175,14 +139,28 @@ def quant_fp8_blockwise_segment_m_kernel( BLOCK_SIZE: tl.constexpr, FP8_MAX: tl.constexpr, ROUND_POW2: tl.constexpr, + ROWWISE: tl.constexpr, + COLUMNWISE: tl.constexpr, ): - """Colwise (axis=0) blockwise quantize with per-segment M padding. - - Reads from the original (unpadded) tensor and writes to a segment-aligned - padded output, so each MoE group starts on a BLOCK_SIZE boundary. Used to - produce the column-wise FP8 operand for the grouped backward (wgrad), always - quantizing from the original high-precision input (no double-quantization). + """Grouped blockwise quantize: rowwise and/or segment-padded columnwise. + + The grid iterates the *segment-padded* output tiles (each MoE group padded up + to a ``BLOCK_SIZE`` boundary along M). Each program reads one ``[BLOCK_SIZE, + BLOCK_SIZE]`` block of the original (unpadded) input once and, from that single + tile, emits the operand(s) selected by the ``constexpr`` flags: + + * ``COLUMNWISE`` (axis=0, BLOCKx1 along M) -- writes FP8 data + 1D scales to the + *segment-padded* positions; this is the per-segment M padding used by the + variable-K wgrad (each group starts on a BLOCK_SIZE boundary). + * ``ROWWISE`` (axis=1, 1xBLOCK along N) -- writes FP8 data + 1D scales to the + *original* (unpadded) positions (the forward/dgrad activation operand); + padding rows are skipped. + + With both flags set this is the fused activation + wgrad quantize (one HBM read + of the input); with only ``COLUMNWISE`` it is the plain segment-padded colwise + quantize. """ + tl.static_assert(ROWWISE or COLUMNWISE, "at least one of ROWWISE/COLUMNWISE must be set") pid_m = tl.program_id(axis=0) pid_n = tl.program_id(axis=1) @@ -206,26 +184,38 @@ def quant_fp8_blockwise_segment_m_kernel( offs_n = tl.cast(pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) offs_m_in = orig_group_start + (offs_m_out - padded_group_start) - mask = ( - (offs_m_in[:, None] >= orig_group_start) - & (offs_m_in[:, None] < orig_group_end) - & (offs_n[None, :] < N) - ) + row_valid = (offs_m_in >= orig_group_start) & (offs_m_in < orig_group_end) + mask = row_valid[:, None] & (offs_n[None, :] < N) x_ptrs = x_ptr + offs_m_in[:, None] * N + offs_n[None, :] x_tile = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) x_tile_abs = tl.abs(x_tile) - x_fp8_tile, x_scales_tile = compute_scale_and_quant(x_tile, x_tile_abs, 0, FP8_MAX, ROUND_POW2) - - x_fp8_ptrs = x_fp8_ptr + offs_m_out[:, None] * N + offs_n[None, :] - out_mask = (offs_m_out[:, None] < M_padded) & (offs_n[None, :] < N) - tl.store(x_fp8_ptrs, x_fp8_tile.to(x_fp8_ptr.dtype.element_ty), mask=out_mask) - - scale_offs = pid_m * N + offs_n - scale_mask = (pid_m < tl.cdiv(M_padded, BLOCK_SIZE)) & (offs_n < N) - x_scales_tile_inv = tl.reshape(1.0 / x_scales_tile, BLOCK_SIZE) - tl.store(x_scales_ptr + scale_offs, x_scales_tile_inv, mask=scale_mask) + if COLUMNWISE: + # axis=0 -> segment-padded output positions. + x_fp8_col_tile, x_scales_col_tile = compute_scale_and_quant( + x_tile, x_tile_abs, 0, FP8_MAX, ROUND_POW2 + ) + x_fp8_col_ptrs = x_fp8_col_ptr + offs_m_out[:, None] * N + offs_n[None, :] + col_out_mask = (offs_m_out[:, None] < M_padded) & (offs_n[None, :] < N) + tl.store( + x_fp8_col_ptrs, x_fp8_col_tile.to(x_fp8_col_ptr.dtype.element_ty), mask=col_out_mask + ) + col_scale_offs = pid_m * N + offs_n + col_scale_mask = (pid_m < tl.cdiv(M_padded, BLOCK_SIZE)) & (offs_n < N) + x_scales_col_tile_inv = tl.reshape(1.0 / x_scales_col_tile, BLOCK_SIZE) + tl.store(x_scales_col_ptr + col_scale_offs, x_scales_col_tile_inv, mask=col_scale_mask) + + if ROWWISE: + # axis=1 -> original (unpadded) positions; padding rows are skipped. + x_fp8_row_tile, x_scales_row_tile = compute_scale_and_quant( + x_tile, x_tile_abs, 1, FP8_MAX, ROUND_POW2 + ) + x_fp8_row_ptrs = x_fp8_row_ptr + offs_m_in[:, None] * N + offs_n[None, :] + tl.store(x_fp8_row_ptrs, x_fp8_row_tile.to(x_fp8_row_ptr.dtype.element_ty), mask=mask) + row_scale_offs = offs_m_in * tl.cdiv(N, BLOCK_SIZE) + pid_n + x_scales_row_tile_inv = tl.reshape(1.0 / x_scales_row_tile, BLOCK_SIZE) + tl.store(x_scales_row_ptr + row_scale_offs, x_scales_row_tile_inv, mask=row_valid) # ----------------------------------------------------------------------------- @@ -237,43 +227,6 @@ def quant_fp8_blockwise_segment_m_kernel( # ----------------------------------------------------------------------------- -def quantize_fp8_blockwise_dual( - x: torch.Tensor, dtype: torch.dtype, block_size: int = 128, pow2: bool = False -): - """Blockwise-quantize a 2D tensor in BOTH row (1xB) and column (Bx1) modes in one pass. - - Returns (x_fp8_row, x_scales_row, x_fp8_col, x_scales_col); scales hold the - dequant scale (amax/FP8_MAX) as fp32. - """ - assert x.is_contiguous() and x.dim() == 2, "Input must be 2D and contiguous" - M, N = x.shape - fp8_max = torch.finfo(dtype).max - - x_fp8_row = torch.empty((M, N), dtype=dtype, device=x.device) - x_scales_row = torch.empty( - (M, triton.cdiv(N, block_size)), dtype=torch.float32, device=x.device - ) - x_fp8_col = torch.empty((M, N), dtype=dtype, device=x.device) - x_scales_col = torch.empty( - (triton.cdiv(M, block_size), N), dtype=torch.float32, device=x.device - ) - - grid = (triton.cdiv(M, block_size), triton.cdiv(N, block_size)) - quant_fp8_blockwise_dual_kernel[grid]( - x, - x_fp8_row, - x_scales_row, - x_fp8_col, - x_scales_col, - M, - N, - BLOCK_SIZE=block_size, - FP8_MAX=fp8_max, - ROUND_POW2=pow2, - ) - return x_fp8_row, x_scales_row, x_fp8_col, x_scales_col - - def quantize_fp8_blockwise( x: torch.Tensor, dtype: torch.dtype, axis: int, block_size: int = 128, pow2: bool = False ): @@ -361,10 +314,14 @@ def quantize_fp8_blockwise_segment_m( ) grid = (triton.cdiv(m_padded_max, block_size), triton.cdiv(N, block_size)) - quant_fp8_blockwise_segment_m_kernel[grid]( + # Columnwise-only: the rowwise pointer args are unused (compiled out) and are + # given the colwise buffers as harmless placeholders. + quant_fp8_blockwise_grouped_kernel[grid]( x, x_fp8, x_scales, + x_fp8, + x_scales, group_offs, var_k_group_offs, N, @@ -372,5 +329,282 @@ def quantize_fp8_blockwise_segment_m( BLOCK_SIZE=block_size, FP8_MAX=fp8_max, ROUND_POW2=pow2, + ROWWISE=False, + COLUMNWISE=True, ) return x_fp8, x_scales, var_k_group_lens, var_k_group_offs + + +def _torch_fp8_dtype(te_fp8_dtype): + """Arch-correct torch FP8 dtype for a TE ``DType`` (kernel launch / ``.view``).""" + return te_dtype_to_torch_dtype(tex.DType(int(te_fp8_dtype))) + + +def _make_blockwise_qtensor( + x_fp8, + scale_inv, + *, + orig_shape, + orig_dtype, + te_fp8_dtype, + is_2D_scaled, + pow2, + columnwise_data=None, + columnwise_scale_inv=None, + vk_group_offs=None, +): + """Wrap raw Triton blockwise-quant outputs in a ``Float8BlockwiseQTensor``. + + ``scale_inv`` (rowwise) must already be in the TE (ROCm) GEMM layout: + ``[ceil(M/128), ceil(N/128)]`` for 2D (weight) scaling and ``[ceil(K/128), M]`` + for 1D (rowwise activation) scaling. + + The optional columnwise operand carries the **segment-padded** wgrad data + ``[M_pad, N]`` and its 1D block scales ``[ceil(M_pad/128), N]`` -- note this is + *not* the standard transposed ``[N, M]`` columnwise layout; the grouped + variable-K wgrad kernel consumes it together with ``vk_group_offs`` (the padded + per-segment offsets). Either operand may be omitted (``None``) when only the + other is needed. + """ + from ..tensor.float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer + + rowwise = x_fp8 is not None + columnwise = columnwise_data is not None + quantizer = Float8BlockQuantizer( + fp8_dtype=te_fp8_dtype, + rowwise=rowwise, + columnwise=columnwise, + block_scaling_dim=2 if is_2D_scaled else 1, + force_pow_2_scales=pow2, + ) + ref = x_fp8 if rowwise else columnwise_data + return Float8BlockwiseQTensor( + shape=tuple(orig_shape), + dtype=orig_dtype, + fp8_dtype=te_fp8_dtype, + rowwise_data=x_fp8.view(torch.uint8) if rowwise else None, + rowwise_scale_inv=scale_inv, + columnwise_data=columnwise_data.view(torch.uint8) if columnwise else None, + columnwise_scale_inv=columnwise_scale_inv, + quantizer=quantizer, + is_2D_scaled=is_2D_scaled, + device=ref.device, + vk_group_offs=vk_group_offs, + ) + + +def quantize_fp8_blockwise_act_qtensor(a, fp8_dtype, pow2: bool = False): + """1x128 rowwise activation quantize returning a ``Float8BlockwiseQTensor``. + + ``fp8_dtype`` is a TE ``DType`` (the QTensor's native dtype); it is mapped to + the arch-correct torch FP8 dtype only for the Triton launch. + """ + x_fp8, scales = quantize_fp8_blockwise( + a, _torch_fp8_dtype(fp8_dtype), axis=1, block_size=128, pow2=pow2 + ) + # Triton emits scales as [M, ceil(K/128)]; TE stores the 1D rowwise scale + # transposed as [ceil(K/128), M] (the GEMM-ready layout). + return _make_blockwise_qtensor( + x_fp8, + scales.t().contiguous(), + orig_shape=a.shape, + orig_dtype=a.dtype, + te_fp8_dtype=fp8_dtype, + is_2D_scaled=False, + pow2=pow2, + ) + + +def wrap_fp8_blockwise_segment_m_qtensor( + data, scale_inv, vk_group_offs, orig_dtype, fp8_dtype, pow2: bool = False +): + """Wrap segment-padded columnwise FP8 fields as a columnwise-only QTensor. + + The segment-padded wgrad operand lives in the ``columnwise_data`` / + ``columnwise_scale_inv`` slots (``[M_pad, N]`` data, ``[ceil(M_pad/128), N]`` + scales) with the padded per-segment offsets in ``_vk_group_offs``. The rowwise + operand is left ``None``; only the variable-K wgrad kernel reads this tensor. + """ + return _make_blockwise_qtensor( + None, + None, + orig_shape=tuple(data.shape), + orig_dtype=orig_dtype, + te_fp8_dtype=fp8_dtype, + is_2D_scaled=False, + pow2=pow2, + columnwise_data=data, + columnwise_scale_inv=scale_inv, + vk_group_offs=vk_group_offs, + ) + + +def quantize_fp8_blockwise_segment_m_operand( + x, fp8_dtype, block_size, group_lens, group_offs, pow2: bool = False +): + """Segment-padded columnwise quantize returning a columnwise-only QTensor. + + ``fp8_dtype`` is a TE ``DType``; the torch FP8 dtype is derived only for the + Triton launch. + """ + x_fp8, x_scales, _vk_lens, vk_offs = quantize_fp8_blockwise_segment_m( + x, _torch_fp8_dtype(fp8_dtype), block_size, group_lens, group_offs, pow2=pow2 + ) + return wrap_fp8_blockwise_segment_m_qtensor( + x_fp8, x_scales, vk_offs, x.dtype, fp8_dtype, pow2=pow2 + ) + + +def quantize_fp8_blockwise_act_and_segment_m( + a, fp8_dtype, block_size, group_lens, group_offs, pow2: bool = False +): + """Fused rowwise activation + segment-padded columnwise quantize in one pass. + + Reads ``a`` ``[M, K]`` from HBM once and returns a single + ``Float8BlockwiseQTensor`` carrying **both** operands the grouped + forward/backward consume: + + * rowwise (1x128 along K) in ``rowwise_data`` with the scale stored transposed + ``[ceil(K/128), M]`` (the forward/dgrad GEMM layout), identical to + :func:`quantize_fp8_blockwise_act_qtensor`. + * segment-padded columnwise (1x128 along M, each MoE group padded to + ``block_size``) in ``columnwise_data`` + ``_vk_group_offs`` (the variable-K + wgrad operand), identical to :func:`quantize_fp8_blockwise_segment_m_operand`. + + Both share the same ``pow2`` rounding and source ``a``; this replaces the two + separate passes (rowwise + segment-padded colwise) with a single launch. Only + worthwhile when the wgrad operand is needed (training). + """ + assert a.is_contiguous() and a.dim() == 2, "Input must be 2D and contiguous" + M, N = a.shape + num_groups = group_lens.size(0) + torch_fp8_dtype = _torch_fp8_dtype(fp8_dtype) + fp8_max = torch.finfo(torch_fp8_dtype).max + + var_k_group_lens = ((group_lens + block_size - 1) // block_size) * block_size + var_k_group_offs = torch.zeros(num_groups + 1, dtype=torch.int64, device=a.device) + var_k_group_offs[1:] = torch.cumsum(var_k_group_lens, dim=0) + + x_fp8_row = torch.empty((M, N), dtype=torch_fp8_dtype, device=a.device) + x_scales_row = torch.empty( + (M, triton.cdiv(N, block_size)), dtype=torch.float32, device=a.device + ) + + m_padded_max = M + num_groups * block_size + x_fp8_col = torch.zeros((m_padded_max, N), dtype=torch_fp8_dtype, device=a.device) + x_scales_col = torch.zeros( + (triton.cdiv(m_padded_max, block_size), N), dtype=torch.float32, device=a.device + ) + + grid = (triton.cdiv(m_padded_max, block_size), triton.cdiv(N, block_size)) + quant_fp8_blockwise_grouped_kernel[grid]( + a, + x_fp8_row, + x_scales_row, + x_fp8_col, + x_scales_col, + group_offs, + var_k_group_offs, + N, + num_groups, + BLOCK_SIZE=block_size, + FP8_MAX=fp8_max, + ROUND_POW2=pow2, + ROWWISE=True, + COLUMNWISE=True, + ) + + # Triton emits rowwise scales as [M, ceil(K/128)]; TE stores the 1D rowwise + # scale transposed as [ceil(K/128), M] (the GEMM-ready layout). The + # segment-padded columnwise operand rides along in the columnwise slots. + return _make_blockwise_qtensor( + x_fp8_row, + x_scales_row.t().contiguous(), + orig_shape=a.shape, + orig_dtype=a.dtype, + te_fp8_dtype=fp8_dtype, + is_2D_scaled=False, + pow2=pow2, + columnwise_data=x_fp8_col, + columnwise_scale_inv=x_scales_col, + vk_group_offs=var_k_group_offs, + ) + + +def quantize_fp8_blockwise_act_operands( + a, + fp8_dtype, + block_size, + group_lens, + group_offs, + *, + rowwise, + columnwise, + pow2: bool = False, +): + """Quantize a grouped activation into the operand(s) the caller needs. + + Dispatches over the (``rowwise``, ``columnwise``) request: + + * both -> :func:`quantize_fp8_blockwise_act_and_segment_m` (one fused HBM read). + * rowwise only -> :func:`quantize_fp8_blockwise_act_qtensor` (forward/dgrad). + * columnwise only -> :func:`quantize_fp8_blockwise_segment_m_operand` (wgrad). + * neither -> ``None``. + + ``block_size`` / ``group_lens`` / ``group_offs`` are only consumed by the + segment-padded columnwise path. + """ + if rowwise and columnwise: + return quantize_fp8_blockwise_act_and_segment_m( + a, fp8_dtype, block_size, group_lens, group_offs, pow2=pow2 + ) + if rowwise: + return quantize_fp8_blockwise_act_qtensor(a, fp8_dtype, pow2=pow2) + if columnwise: + return quantize_fp8_blockwise_segment_m_operand( + a, fp8_dtype, block_size, group_lens, group_offs, pow2=pow2 + ) + return None + + +def wrap_fp8_blockwise_grouped_weight_qtensor( + rowwise_data, rowwise_scale_inv, orig_dtype, fp8_dtype, pow2: bool = True +): + """Wrap packed grouped-weight FP8 data + scales as a 2D ``[G*N, K]`` QTensor. + + Because ``N`` (out_features) is a multiple of 128, expert boundaries align + with the 128x128 blocks, so a packed ``[G, N, K]`` weight is bit-identical to + a single 2D ``[G*N, K]`` blockwise tensor with the scale flattened to + ``[G*(N/128), K/128]`` -- and it dequantizes correctly. ``rowwise_data`` may + be ``[G, N, K]`` or already 2D ``[G*N, K]`` (scales correspondingly 3D or 2D). + """ + if rowwise_data.dim() == 3: + gg, nn, kk = rowwise_data.shape + rowwise_data = rowwise_data.reshape(gg * nn, kk) + rowwise_scale_inv = rowwise_scale_inv.reshape( + gg * rowwise_scale_inv.shape[1], rowwise_scale_inv.shape[2] + ) + return _make_blockwise_qtensor( + rowwise_data, + rowwise_scale_inv, + orig_shape=tuple(rowwise_data.shape), + orig_dtype=orig_dtype, + te_fp8_dtype=fp8_dtype, + is_2D_scaled=True, + pow2=pow2, + ) + + +def quantize_fp8_blockwise_grouped_weight_qtensor(w, fp8_dtype, pow2: bool = False): + """128x128 grouped-weight quantize of ``[G, N, K]`` -> 2D ``[G*N, K]`` QTensor. + + ``N`` must be a multiple of 128 so expert boundaries align with 128-blocks + (see :func:`wrap_fp8_blockwise_grouped_weight_qtensor`). One packed quantize + launch over ``[G, N, K]``; no per-expert launches, no stack. + """ + assert w.dim() == 3, "quantize_fp8_blockwise_grouped_weight_qtensor expects [G, N, K]" + assert w.shape[1] % 128 == 0, "grouped weight N (out_features) must be a multiple of 128" + x_fp8, scales = quantize_fp8_blockwise_weight( + w, _torch_fp8_dtype(fp8_dtype), block_size=128, pow2=pow2 + ) + return wrap_fp8_blockwise_grouped_weight_qtensor(x_fp8, scales, w.dtype, fp8_dtype, pow2=pow2) From d0f81c408b072e1cb5494a84a93d705c31036e2d Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Thu, 27 Aug 2026 23:11:18 +0000 Subject: [PATCH 12/13] Fix single_grouped_weight grad-state mirroring and cache blockwise FP8 weight quantization - Mirror `requires_grad`, `main_grad`, and `overwrite_main_grad` from the grouped Parameter onto per-expert split views in `_get_weight_tensors` / `_get_bias_tensors`, guarding `requires_grad_` to leaves and validating that `main_grad` can alias in place. - Cache the packed blockwise FP8 quantized weight across microbatches in the Triton grouped-GEMM path using `weight_workspaces`. - Add a ROCm autouse `EnvVarCleaner` fixture to snapshot/restore grouped-GEMM Triton backend env vars and prevent leakage between tests, plus a unit test verifying the grad-state mirroring behavior. --- tests/pytorch/test_grouped_linear.py | 79 +++++++++++++++++-- .../pytorch/module/grouped_linear.py | 50 +++++++++--- 2 files changed, 112 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 1be232e9a3..fabb2c39c1 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -58,6 +58,18 @@ seed = 1234 reset_rng_states() +if IS_HIP_EXTENSION: + from utils import EnvVarCleaner + + @pytest.fixture(autouse=True) + def reset_grouped_gemm_backend(): + # Snapshot/restore the process-global Triton grouped-GEMM env vars so a test + # that sets them (and may raise before its own cleanup) cannot leak the + # backend into unrelated tests. + env = EnvVarCleaner(["NVTE_USE_BLOCKWISE_GMM_TRITON", "NVTE_USE_GROUPED_GEMM_TRITON"]) + yield + + NVTE_TEST_NVINSPECT_ENABLED = int(os.environ.get("NVTE_TEST_NVINSPECT_ENABLED", "0")) if NVTE_TEST_NVINSPECT_ENABLED: @@ -386,11 +398,6 @@ def test_grouped_linear_accuracy( delay_wgrad_compute, ) - use_blockwise_triton = os.getenv("NVTE_USE_BLOCKWISE_GMM_TRITON", "0") == "1" - if use_triton: - os.environ.pop("NVTE_USE_GROUPED_GEMM_TRITON", None) - os.environ.pop("NVTE_USE_BLOCKWISE_GMM_TRITON", None) - atol, rtol = 0, 0 if use_cutlass: atol, rtol = 1e-3, 1e-3 @@ -1992,6 +1999,68 @@ def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monk grouped_linear.backward_dw() +def test_single_grouped_weight_mirrors_grad_state(monkeypatch): + """Per-expert views must carry the grad state the autograd Function reads off + ``weights[i]`` / ``biases[i]``. + + Regression for the generic ``single_grouped_weight`` fix: ``_get_weight_tensors`` / + ``_get_bias_tensors`` return split views of the grouped Parameter, which do not + inherit ``requires_grad`` (gates wgrad), ``main_grad`` (fuse-accumulation), or + ``overwrite_main_grad``. ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` is off by default, so a + regression here would otherwise land silently on both CUDA and ROCm. + """ + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + num_gemms, in_features, out_features = 3, 128, 128 + grouped_linear = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=True, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + single_grouped_bias=True, + ) + assert grouped_linear.single_grouped_weight and grouped_linear.single_grouped_bias + + # Megatron-style fuse-accumulation buffer on the grouped weight Parameter. + grouped_linear.weight.main_grad = torch.zeros( + num_gemms * out_features, in_features, dtype=torch.float32, device="cuda" + ) + grouped_linear.weight.overwrite_main_grad = False + + weights = grouped_linear._get_weight_tensors() + assert len(weights) == num_gemms + for i, w in enumerate(weights): + assert w.requires_grad == grouped_linear.weight.requires_grad + assert w.overwrite_main_grad is False + # main_grad must be an aliasing view so per-expert wgrad accumulates in place. + w.main_grad.add_(i + 1.0) + grouped_main_grad = grouped_linear.weight.main_grad.view(num_gemms, out_features, in_features) + for i in range(num_gemms): + assert torch.all(grouped_main_grad[i] == (i + 1.0)) + + # requires_grad is mirrored (not merely coincidental): flip the Parameter and refetch. + grouped_linear.weight.requires_grad_(False) + for w in grouped_linear._get_weight_tensors(): + assert w.requires_grad is False + + # Bias views mirror requires_grad (bias never uses main_grad). + biases = grouped_linear._get_bias_tensors() + assert len(biases) == num_gemms + for b in biases: + assert b.requires_grad == grouped_linear.bias.requires_grad + + # A main_grad that cannot alias as [G*out, in] must raise a clear error (an + # aliasing view is required; a silent copy would drop accumulation). + grouped_linear.weight.requires_grad_(True) + grouped_linear.weight.main_grad = torch.zeros( + num_gemms * out_features + 1, in_features, dtype=torch.float32, device="cuda" + ) + with pytest.raises(RuntimeError, match="accumulate in place"): + grouped_linear._get_weight_tensors() + + @pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): """Non-RHT NVFP4 falls back to the legacy path; check it stays numerically correct. diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index b0f353519e..64de640e09 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -534,6 +534,8 @@ def _forward_blockwise_fp8_triton( fuse_wgrad_accumulation, is_first_microbatch, wgrad_store, + weight_workspaces=None, + cache_weight=False, pow2_x=True, pow2_w=True, pow2_grad=True, @@ -600,6 +602,7 @@ def _forward_blockwise_fp8_triton( # stored rowwise data + scales; otherwise quantize the high-precision # weights once over the packed 3D view (zero-copy when # ``single_grouped_weight``), no stack. + new_workspaces = [None] * num_gemms if isinstance(weights[0], Float8BlockwiseQTensor): # fp8_model_params: reuse the stored rowwise data + scales directly (no # re-quantization). @@ -607,8 +610,17 @@ def _forward_blockwise_fp8_triton( weights, activation_dtype, dt, pow2_w ) else: - w = _GroupedLinear._expert_weights_as_3d(weights, activation_dtype) - qw = quantize_fp8_blockwise_grouped_weight_qtensor(w, dt, pow2=pow2_w) + # High-precision weights: quantize the whole packed weight into a single + # Float8BlockwiseQTensor and cache it across microbatches. + update_ws = is_first_microbatch is None or is_first_microbatch + cached_qw = weight_workspaces[0] if weight_workspaces else None + if not update_ws and isinstance(cached_qw, Float8BlockwiseQTensor): + qw = cached_qw + else: + w = _GroupedLinear._expert_weights_as_3d(weights, activation_dtype) + qw = quantize_fp8_blockwise_grouped_weight_qtensor(w, dt, pow2=pow2_w) + if cache_weight: + new_workspaces[0] = qw # Forward GEMM: out[seg] = A[seg] @ W[g]^T (trans_b=True). out = grouped_gemm_fp8_blockwise_triton_kernel( @@ -656,7 +668,6 @@ def _forward_blockwise_fp8_triton( lambda j=i: weights[j].main_grad for i in range(num_gemms) ] - new_workspaces = [None] * num_gemms return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @staticmethod @@ -885,6 +896,8 @@ def forward( fuse_wgrad_accumulation=fuse_wgrad_accumulation, is_first_microbatch=is_first_microbatch, wgrad_store=wgrad_store, + weight_workspaces=weight_workspaces, + cache_weight=cache_weight, # Match the recipe's scale rounding (pow2 by default) so the Triton # path quantizes identically to the default Float8BlockScaling path. pow2_x=blockwise_recipe.fp8_quant_fwd_inp.power_2_scale, @@ -2466,19 +2479,29 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage if weight_tensors is None: # TODO(ksivaman): Remove this after GEMM integration. weight_tensors = grouped_weight.split_into_quantized_tensors() - # Members are views, not the Parameter, so mirror the grad state the - # autograd Function reads off ``weights[i]``: ``requires_grad`` (gates - # wgrad), ``main_grad`` (per-expert views into the grouped - # fuse-accumulation buffer), and ``overwrite_main_grad``. + # ROCm + CUDA fix: Split views don't carry the grad state that the autograd Function reads + # off ``weights[i]``, so mirror it from the grouped Parameter (needed on + # every backend; transitional until ``_get_weight_tensors`` returns the + # Parameter directly). want_grad = grouped_weight.requires_grad main_grad = getattr(grouped_weight, "main_grad", None) per_expert_main_grad = None if main_grad is not None: - per_expert_main_grad = main_grad.view( - self.num_gemms, self.out_features, self.in_features - ) + # Must be an aliasing view (not a copy) so per-expert wgrad accumulates + # in place; requires a contiguous [num_gemms, out_features, in_features]. + try: + per_expert_main_grad = main_grad.view( + self.num_gemms, self.out_features, self.in_features + ) + except RuntimeError as e: + raise RuntimeError( + "single_grouped_weight with fuse_wgrad_accumulation requires main_grad to " + "be a contiguous [num_gemms * out_features, in_features] buffer so per-expert" + f" wgrad can accumulate in place (got shape {tuple(main_grad.shape)})." + ) from e for i, w in enumerate(weight_tensors): - if w.requires_grad != want_grad: + # ``requires_grad_`` only works on leaves. + if w.requires_grad != want_grad and w.is_leaf: w.requires_grad_(want_grad) if per_expert_main_grad is not None: w.main_grad = per_expert_main_grad[i] @@ -2504,9 +2527,12 @@ def _get_bias_tensors(self) -> List[torch.Tensor]: parts = grouped_bias.quantized_tensors if parts is None: parts = grouped_bias.split_into_quantized_tensors() + # Mirror ``requires_grad`` onto the per-expert views (bias never uses + # ``main_grad``: fuse-accumulation is weights-only). ``requires_grad_`` + # only works on leaves, so guard to avoid a RuntimeError on a non-leaf view. want_grad = grouped_bias.requires_grad for p in parts: - if p.requires_grad != want_grad: + if p.requires_grad != want_grad and p.is_leaf: p.requires_grad_(want_grad) return [p.reshape(-1) for p in parts] return [getattr(self, f"bias{i}") for i in range(self.num_gemms)] From 835742dc5934e25a33d7e6027c90001f597852db Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Fri, 28 Aug 2026 16:40:57 +0000 Subject: [PATCH 13/13] Fix blockwise FP8 grouped-linear weight cache reuse and add unit test - On `is_first_microbatch=False`, validate the cached blockwise FP8 weight workspace is a `Float8BlockwiseQTensor` of the expected packed shape and raise a clear `RuntimeError` if it is incompatible, instead of silently re-quantizing or using a stale buffer. - Only stage the freshly quantized weight for write-back when a new quantization actually occurs. - Replace the `single_grouped_weight` grad-state mirroring test with a ROCm-only test that verifies the blockwise FP8 packed weight is quantized once, reused across microbatches without re-quantization, remains numerically identical, and fails loudly on an incompatible cached workspace. --- tests/pytorch/test_grouped_linear.py | 107 +++++++++--------- .../pytorch/module/grouped_linear.py | 20 +++- 2 files changed, 70 insertions(+), 57 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index fabb2c39c1..ead7966a57 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -1999,66 +1999,67 @@ def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monk grouped_linear.backward_dw() -def test_single_grouped_weight_mirrors_grad_state(monkeypatch): - """Per-expert views must carry the grad state the autograd Function reads off - ``weights[i]`` / ``biases[i]``. - - Regression for the generic ``single_grouped_weight`` fix: ``_get_weight_tensors`` / - ``_get_bias_tensors`` return split views of the grouped Parameter, which do not - inherit ``requires_grad`` (gates wgrad), ``main_grad`` (fuse-accumulation), or - ``overwrite_main_grad``. ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` is off by default, so a - regression here would otherwise land silently on both CUDA and ROCm. +@pytest.mark.skipif(not IS_HIP_EXTENSION, reason="Blockwise FP8 grouped GEMM is ROCm-only.") +@pytest.mark.skipif(not fp8_block_scaling_available, reason="FP8 block scaling is unsupported.") +def test_blockwise_fp8_weight_cache_reuse(monkeypatch): + """``is_first_microbatch`` caches the packed blockwise weight across microbatches. + + CI never passes ``is_first_microbatch``, so the blockwise cache-hit branch is + otherwise unexercised. Verify that (1) the first microbatch quantizes once and + stores the packed weight in the single workspace slot (``weight0``), (2) later + microbatches reuse it with no re-quantization, (3) the reuse is numerically + identical to quantizing fresh, and (4) an incompatible cached workspace fails loud + instead of running a wrong weight. """ - monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") - num_gemms, in_features, out_features = 3, 128, 128 + monkeypatch.setenv("NVTE_USE_BLOCKWISE_GMM_TRITON", "1") + + from transformer_engine.pytorch.triton_kernels import blockwise_quantize as _bwq + + # Count quantization launches; the forward re-imports this symbol per call, so + # patching the source module is observed. + n_quant = {"count": 0} + _orig_quant = _bwq.quantize_fp8_blockwise_grouped_weight_qtensor + + def _counting_quant(*args, **kwargs): + n_quant["count"] += 1 + return _orig_quant(*args, **kwargs) + + monkeypatch.setattr( + _bwq, "quantize_fp8_blockwise_grouped_weight_qtensor", _counting_quant + ) + + num_gemms, in_features, out_features = 2, 256, 512 + fp8_recipe = recipe.Float8BlockScaling() grouped_linear = GroupedLinear( num_gemms, in_features, out_features, - bias=True, + bias=False, params_dtype=torch.bfloat16, device="cuda", - single_grouped_weight=True, - single_grouped_bias=True, - ) - assert grouped_linear.single_grouped_weight and grouped_linear.single_grouped_bias - - # Megatron-style fuse-accumulation buffer on the grouped weight Parameter. - grouped_linear.weight.main_grad = torch.zeros( - num_gemms * out_features, in_features, dtype=torch.float32, device="cuda" - ) - grouped_linear.weight.overwrite_main_grad = False - - weights = grouped_linear._get_weight_tensors() - assert len(weights) == num_gemms - for i, w in enumerate(weights): - assert w.requires_grad == grouped_linear.weight.requires_grad - assert w.overwrite_main_grad is False - # main_grad must be an aliasing view so per-expert wgrad accumulates in place. - w.main_grad.add_(i + 1.0) - grouped_main_grad = grouped_linear.weight.main_grad.view(num_gemms, out_features, in_features) - for i in range(num_gemms): - assert torch.all(grouped_main_grad[i] == (i + 1.0)) - - # requires_grad is mirrored (not merely coincidental): flip the Parameter and refetch. - grouped_linear.weight.requires_grad_(False) - for w in grouped_linear._get_weight_tensors(): - assert w.requires_grad is False - - # Bias views mirror requires_grad (bias never uses main_grad). - biases = grouped_linear._get_bias_tensors() - assert len(biases) == num_gemms - for b in biases: - assert b.requires_grad == grouped_linear.bias.requires_grad - - # A main_grad that cannot alias as [G*out, in] must raise a clear error (an - # aliasing view is required; a silent copy would drop accumulation). - grouped_linear.weight.requires_grad_(True) - grouped_linear.weight.main_grad = torch.zeros( - num_gemms * out_features + 1, in_features, dtype=torch.float32, device="cuda" - ) - with pytest.raises(RuntimeError, match="accumulate in place"): - grouped_linear._get_weight_tensors() + ).eval() + + m_splits = [32, 32] + x = torch.randn(sum(m_splits), in_features, dtype=torch.bfloat16, device="cuda") + + with torch.no_grad(), autocast(enabled=True, recipe=fp8_recipe): + out_first = grouped_linear(x, m_splits, is_first_microbatch=True) + assert n_quant["count"] == 1, "first microbatch should quantize the weight once" + assert "weight0" in grouped_linear._fp8_workspaces + cached = grouped_linear._fp8_workspaces["weight0"] + + with torch.no_grad(), autocast(enabled=True, recipe=fp8_recipe): + out_second = grouped_linear(x, m_splits, is_first_microbatch=False) + assert n_quant["count"] == 1, "later microbatch should reuse the cache, not re-quantize" + assert grouped_linear._fp8_workspaces["weight0"] is cached + torch.testing.assert_close(out_second, out_first, rtol=0, atol=0) + + # A live is_first_microbatch=False cache whose layout does not match the current + # weight must raise, not silently GEMM a wrong weight. + grouped_linear._fp8_workspaces["weight0"] = torch.empty(1, dtype=torch.uint8, device="cuda") + with torch.no_grad(), autocast(enabled=True, recipe=fp8_recipe): + with pytest.raises(RuntimeError, match="Cached blockwise weight workspace"): + grouped_linear(x, m_splits, is_first_microbatch=False) @pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 64de640e09..56c1837660 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -611,16 +611,28 @@ def _forward_blockwise_fp8_triton( ) else: # High-precision weights: quantize the whole packed weight into a single - # Float8BlockwiseQTensor and cache it across microbatches. + # Float8BlockwiseQTensor and cache it across microbatches. On a cache hit + # (``is_first_microbatch=False``) reuse the stored packed weight untouched so + # its identity is preserved; only a fresh quantization is staged for write-back. update_ws = is_first_microbatch is None or is_first_microbatch cached_qw = weight_workspaces[0] if weight_workspaces else None - if not update_ws and isinstance(cached_qw, Float8BlockwiseQTensor): + if not update_ws: + expected_wshape = (num_gemms * out_features, in_features) + if not ( + isinstance(cached_qw, Float8BlockwiseQTensor) + and tuple(cached_qw.shape) == expected_wshape + ): + raise RuntimeError( + "Cached blockwise weight workspace is incompatible with the " + "current weights; expected a Float8BlockwiseQTensor of shape " + f"{expected_wshape}." + ) qw = cached_qw else: w = _GroupedLinear._expert_weights_as_3d(weights, activation_dtype) qw = quantize_fp8_blockwise_grouped_weight_qtensor(w, dt, pow2=pow2_w) - if cache_weight: - new_workspaces[0] = qw + if cache_weight: + new_workspaces[0] = qw # Forward GEMM: out[seg] = A[seg] @ W[g]^T (trans_b=True). out = grouped_gemm_fp8_blockwise_triton_kernel(