diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 5062cf6b42f..5a6d93f4ee5 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -8,9 +8,14 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.constants import TE_DType from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm +from transformer_engine.pytorch.cpp_extensions import ( + general_gemm, + general_grouped_gemm, + general_grouped_gemm_for_grouped_tensor, +) from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import reference_utils +from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor from transformer_engine.pytorch.tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage @@ -838,3 +843,385 @@ def test_nvfp4_ue5m3_gemm_leaves_operands_unswizzled(M: int, K: int, N: int): ref = x_q.dequantize(dtype=torch.float32) @ w_q.dequantize(dtype=torch.float32).t() rel_err = (y1.float() - ref).norm() / ref.norm() assert rel_err < 5e-3, f"relative error {rel_err:.2e} is too large" + + +_grouped_tensor_gemm_supported = True +_reason_for_no_grouped_tensor_gemm = "" +if torch.cuda.get_device_capability() < (10, 0): + _grouped_tensor_gemm_supported = False + _reason_for_no_grouped_tensor_gemm = "NVFP4 grouped-tensor GEMM requires Blackwell+." +elif tex.get_cublasLt_version() < 130400: + _grouped_tensor_gemm_supported = False + _reason_for_no_grouped_tensor_gemm = "NVFP4 grouped-tensor GEMM requires cuBLAS 13.4+." + + +def _make_row_scaled_grouped_activation( + m_splits: list[int], + k: int, + *, + columnwise: bool, + tokens: torch.Tensor | None = None, +) -> GroupedTensor: + """Cast tokens to a row-scaled NVFP4 GroupedTensor; splits and k must be multiples of 128.""" + device = "cuda" + if tokens is None: + tokens = torch.randn(sum(m_splits), k, dtype=torch.bfloat16, device=device) + tokens_per_expert = torch.tensor(m_splits, dtype=torch.int64, device=device) + quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=columnwise, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=True, + ) + quantizer.optimize_for_gemm = True + return tex.group_quantize(tokens, quantizer, len(m_splits), tokens_per_expert) + + +def _make_output_grouped_tensor( + m_splits: list[int], last_dim: int, device: torch.device, dtype: torch.dtype +) -> GroupedTensor: + return GroupedTensor.make_grouped_tensor( + num_tensors=len(m_splits), + first_dims=torch.tensor(m_splits, device=device, dtype=torch.int64), + last_dims=None, + logical_first_dim=sum(m_splits), + logical_last_dim=last_dim, + quantizer=None, + device=device, + dtype=dtype, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "m_splits, k, n", + [ + ([128, 128], 128, 128), + ([128, 256, 128], 256, 256), + ([256, 128, 384], 512, 256), + ([128, 256, 384, 512], 256, 512), + ], +) +@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str) +def test_nvfp4_row_scaled_grouped_tensor_gemm_fprop( + m_splits: list[int], k: int, n: int, out_dtype: torch.dtype +) -> None: + """Row-scaled grouped-tensor TN GEMM matches a per-expert dense GEMM.""" + if not _grouped_tensor_gemm_supported: + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(7) + device = torch.device("cuda") + num_gemms = len(m_splits) + + activation = _make_row_scaled_grouped_activation(m_splits, k, columnwise=False) + activation_splits = activation.split_into_quantized_tensors() + + w_quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + ) + w_quantizer.optimize_for_gemm = True + weights = [] + expected = [] + for i in range(num_gemms): + w = torch.randn(n, k, dtype=torch.bfloat16, device=device) + w_nvfp4 = w_quantizer(w) + weights.append(w_nvfp4) + expected.append( + general_gemm(w_nvfp4, activation_splits[i], out_dtype=out_dtype, layout="TN")[0] + ) + + grouped_out = _make_output_grouped_tensor(m_splits, n, device, out_dtype) + general_grouped_gemm_for_grouped_tensor(weights, activation, grouped_out, layout="TN") + actual = grouped_out.split_into_quantized_tensors() + + if out_dtype == torch.bfloat16: + tols = dict(atol=0.0, rtol=7.8e-3) + else: + tols = dict(atol=1e-4, rtol=1e-4) + for got, ref in zip(actual, expected): + torch.testing.assert_close(got, ref, **tols) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("m_splits, k, n", [([128, 256, 128], 256, 256)]) +def test_nvfp4_row_scaled_grouped_tensor_gemm_cudagraph( + m_splits: list[int], k: int, n: int +) -> None: + """A captured row-scaled grouped-tensor GEMM replays correctly.""" + if not _grouped_tensor_gemm_supported: + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(11) + device = torch.device("cuda") + num_gemms = len(m_splits) + out_dtype = torch.float32 + + static_tokens = torch.randn(sum(m_splits), k, dtype=torch.bfloat16, device=device) + activation = _make_row_scaled_grouped_activation( + m_splits, k, columnwise=False, tokens=static_tokens + ) + + w_quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + ) + w_quantizer.optimize_for_gemm = True + weights = [ + w_quantizer(torch.randn(n, k, dtype=torch.bfloat16, device=device)) + for _ in range(num_gemms) + ] + grouped_out = _make_output_grouped_tensor(m_splits, n, device, out_dtype) + + def _run(): + general_grouped_gemm_for_grouped_tensor(weights, activation, grouped_out, layout="TN") + + torch.cuda.synchronize() + _run() + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _run() + + fresh = _make_row_scaled_grouped_activation(m_splits, k, columnwise=False) + activation.rowwise_data.copy_(fresh.rowwise_data) + activation.scale_inv.copy_(fresh.scale_inv) + activation.amax.copy_(fresh.amax) + graph.replay() + torch.cuda.synchronize() + replayed = [t.clone() for t in grouped_out.split_into_quantized_tensors()] + + _run() + torch.cuda.synchronize() + expected = grouped_out.split_into_quantized_tensors() + for got, ref in zip(replayed, expected): + torch.testing.assert_close(got, ref, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "m_splits, k, n", + [ + ([128, 128], 128, 128), + ([128, 256, 128], 256, 256), + ([256, 128, 384], 512, 256), + ([128, 256, 384, 512], 256, 512), + ], +) +@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str) +def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad( + m_splits: list[int], k: int, n: int, out_dtype: torch.dtype +) -> None: + """Row-scaled grouped-tensor NT GEMM matches a per-expert dense GEMM.""" + if not _grouped_tensor_gemm_supported: + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(13) + device = torch.device("cuda") + num_gemms = len(m_splits) + + x = _make_row_scaled_grouped_activation(m_splits, k, columnwise=True) + dy = _make_row_scaled_grouped_activation(m_splits, n, columnwise=True) + x_splits = x.split_into_quantized_tensors() + dy_splits = dy.split_into_quantized_tensors() + + expected = [ + general_gemm(x_splits[i], dy_splits[i], out_dtype=out_dtype, layout="NT")[0] + for i in range(num_gemms) + ] + + grouped_out = _make_output_grouped_tensor([n] * num_gemms, k, device, out_dtype) + general_grouped_gemm_for_grouped_tensor(x, dy, grouped_out, layout="NT") + actual = grouped_out.split_into_quantized_tensors() + + if out_dtype == torch.bfloat16: + tols = dict(atol=0.0, rtol=7.8e-3) + else: + tols = dict(atol=1e-4, rtol=1e-4) + for got, ref in zip(actual, expected): + torch.testing.assert_close(got, ref, **tols) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("m_splits, k, n", [([256, 128, 384], 512, 256)]) +def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad_discrete_out( + m_splits: list[int], k: int, n: int +) -> None: + """Row-scaled NT GEMM into a list output matches a per-expert dense GEMM.""" + if not _grouped_tensor_gemm_supported: + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(17) + device = torch.device("cuda") + num_gemms = len(m_splits) + out_dtype = torch.float32 + + x = _make_row_scaled_grouped_activation(m_splits, k, columnwise=True) + dy = _make_row_scaled_grouped_activation(m_splits, n, columnwise=True) + x_splits = x.split_into_quantized_tensors() + dy_splits = dy.split_into_quantized_tensors() + + expected = [ + general_gemm(x_splits[i], dy_splits[i], out_dtype=out_dtype, layout="NT")[0] + for i in range(num_gemms) + ] + + out_list = [torch.empty(n, k, dtype=out_dtype, device=device) for _ in range(num_gemms)] + general_grouped_gemm_for_grouped_tensor(x, dy, out_list, layout="NT") + + for got, ref in zip(out_list, expected): + torch.testing.assert_close(got, ref, atol=1e-4, rtol=1e-4) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("m_splits, k, n", [([128, 256, 128], 256, 256)]) +def test_nvfp4_row_scaled_grouped_tensor_gemm_uneven_colscale( + m_splits: list[int], k: int, n: int +) -> None: + """Column post-scale with uneven groups matches a per-expert dense GEMM.""" + if not _grouped_tensor_gemm_supported: + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(23) + device = torch.device("cuda") + num_gemms = len(m_splits) + out_dtype = torch.float32 + + weights = _make_row_scaled_grouped_activation([n] * num_gemms, k, columnwise=False) + activation = _make_row_scaled_grouped_activation(m_splits, k, columnwise=False) + w_splits = weights.split_into_quantized_tensors() + act_splits = activation.split_into_quantized_tensors() + + expected = [ + general_gemm(w_splits[i], act_splits[i], out_dtype=out_dtype, layout="TN")[0] + for i in range(num_gemms) + ] + + grouped_out = _make_output_grouped_tensor(m_splits, n, device, out_dtype) + general_grouped_gemm_for_grouped_tensor(weights, activation, grouped_out, layout="TN") + actual = grouped_out.split_into_quantized_tensors() + + for got, ref in zip(actual, expected): + torch.testing.assert_close(got, ref, atol=1e-4, rtol=1e-4) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("m_splits, k, n", [([128, 256, 128], 256, 256)]) +def test_nvfp4_row_scaled_grouped_tensor_gemm_colscale_cudagraph( + m_splits: list[int], k: int, n: int +) -> None: + """A captured column post-scale with uneven groups replays correctly.""" + if not _grouped_tensor_gemm_supported: + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(29) + device = torch.device("cuda") + num_gemms = len(m_splits) + out_dtype = torch.float32 + + static_w = torch.randn(n * num_gemms, k, dtype=torch.bfloat16, device=device) + static_act = torch.randn(sum(m_splits), k, dtype=torch.bfloat16, device=device) + weights = _make_row_scaled_grouped_activation( + [n] * num_gemms, k, columnwise=False, tokens=static_w + ) + activation = _make_row_scaled_grouped_activation( + m_splits, k, columnwise=False, tokens=static_act + ) + grouped_out = _make_output_grouped_tensor(m_splits, n, device, out_dtype) + + def _run(): + general_grouped_gemm_for_grouped_tensor(weights, activation, grouped_out, layout="TN") + + torch.cuda.synchronize() + _run() + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _run() + + fresh_w = _make_row_scaled_grouped_activation([n] * num_gemms, k, columnwise=False) + fresh_act = _make_row_scaled_grouped_activation(m_splits, k, columnwise=False) + for dst, src in ((weights, fresh_w), (activation, fresh_act)): + dst.rowwise_data.copy_(src.rowwise_data) + dst.scale_inv.copy_(src.scale_inv) + dst.amax.copy_(src.amax) + graph.replay() + torch.cuda.synchronize() + replayed = [t.clone() for t in grouped_out.split_into_quantized_tensors()] + + _run() + torch.cuda.synchronize() + expected = grouped_out.split_into_quantized_tensors() + for got, ref in zip(replayed, expected): + torch.testing.assert_close(got, ref, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("m_splits, k, n", [([128, 256, 128], 256, 256)]) +def test_nvfp4_row_scaled_grouped_tensor_gemm_output_capacity( + m_splits: list[int], k: int, n: int +) -> None: + """Output capacity rows beyond the live region do not affect the result.""" + if not _grouped_tensor_gemm_supported: + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(31) + device = torch.device("cuda") + num_gemms = len(m_splits) + out_dtype = torch.float32 + + activation = _make_row_scaled_grouped_activation(m_splits, k, columnwise=False) + activation_splits = activation.split_into_quantized_tensors() + + w_quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + ) + w_quantizer.optimize_for_gemm = True + weights = [] + expected = [] + for i in range(num_gemms): + w_nvfp4 = w_quantizer(torch.randn(n, k, dtype=torch.bfloat16, device=device)) + weights.append(w_nvfp4) + expected.append( + general_gemm(w_nvfp4, activation_splits[i], out_dtype=out_dtype, layout="TN")[0] + ) + + grouped_out = _make_output_grouped_tensor(m_splits, n, device, out_dtype) + live_numel = grouped_out.rowwise_data.numel() + padded = torch.full( + (live_numel + 128 * n,), + float("nan"), + dtype=grouped_out.rowwise_data.dtype, + device=device, + ) + padded[:live_numel].copy_(grouped_out.rowwise_data) + grouped_out.rowwise_data = padded + + general_grouped_gemm_for_grouped_tensor(weights, activation, grouped_out, layout="TN") + actual = grouped_out.split_into_quantized_tensors() + + for got, ref in zip(actual, expected): + torch.testing.assert_close(got, ref, atol=1e-4, rtol=1e-4) diff --git a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh index cfc5dd42c81..10b5bddc325 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -26,6 +26,7 @@ #include "../core/grouped_layout.cuh" #include "core_nvfp4.cuh" #include "specialized/quantize_transpose_nvfp4_tuned_1D.cuh" +#include "swizzle.cuh" namespace transformer_engine { namespace dispatch { @@ -1390,7 +1391,8 @@ using namespace ptx; #if FP4_TYPE_SUPPORTED -template +template __device__ __forceinline__ void group_row_scaled_cast_body( const CUtensorMap &tensor_map_input, uint8_t *const q_row_base, ScaleType *const s_dec_row_base, const float *const row_amax_base, uint8_t *const q_col_base, ScaleType *const s_dec_col_base, @@ -1570,8 +1572,17 @@ __device__ __forceinline__ void group_row_scaled_cast_body( row += THREADS_NUM) { ScalesVec &scales_vec = *reinterpret_cast(sSFrowwise[row]); const size_t local_row = static_cast(local_row_base) + row; - const size_t scale_idx = local_row * scale_stride_row + scales_block_offset_X_rowwise; - scales_vec.store_to_elts(&s_dec_row_base[scale_idx], 0, SCALES_PER_CHUNK_X); + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { +#pragma unroll + for (int k = 0; k < SCALES_PER_CHUNK_X; ++k) { + const size_t col = scales_block_offset_X_rowwise + k; + s_dec_row_base[swizzle::gemm_swizzled_scale_idx(local_row, col, scale_stride_row)] = + scales_vec.data.elt[k]; + } + } else { + const size_t scale_idx = local_row * scale_stride_row + scales_block_offset_X_rowwise; + scales_vec.store_to_elts(&s_dec_row_base[scale_idx], 0, SCALES_PER_CHUNK_X); + } } } if constexpr (RETURN_TRANSPOSE) { @@ -1582,8 +1593,16 @@ __device__ __forceinline__ void group_row_scaled_cast_body( row_tr += THREADS_NUM) { ScalesVec &scales_vec = *reinterpret_cast(sSFcolwise[row_tr]); const size_t global_col = static_cast(block_offset_X) + row_tr; - const size_t scale_idx = global_col * scale_stride_col + local_block_offset_M; - scales_vec.store_to_elts(&s_dec_col_base[scale_idx], 0, SCALES_PER_CHUNK_Y); + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { +#pragma unroll + for (int k = 0; k < SCALES_PER_CHUNK_Y; ++k) { + s_dec_col_base[swizzle::gemm_swizzled_scale_idx( + global_col, local_block_offset_M + k, scale_stride_col)] = scales_vec.data.elt[k]; + } + } else { + const size_t scale_idx = global_col * scale_stride_col + local_block_offset_M; + scales_vec.store_to_elts(&s_dec_col_base[scale_idx], 0, SCALES_PER_CHUNK_Y); + } } } @@ -1599,7 +1618,8 @@ __device__ __forceinline__ void group_row_scaled_cast_body( } // Grid covers the input capacity; CTAs past the live row sum noop. -template +template __global__ void __launch_bounds__(THREADS_NUM) group_row_scaled_cast_nvfp4_graph_safe_kernel( const __grid_constant__ CUtensorMap tensor_map_input, uint8_t *const q_row_all, ScaleType *const s_dec_row_all, const float *const row_amax_all, uint8_t *const q_col_all, @@ -1639,7 +1659,7 @@ __global__ void __launch_bounds__(THREADS_NUM) group_row_scaled_cast_nvfp4_graph RETURN_TRANSPOSE ? s_dec_col_all + (cols / SCALE_DIM) * row_start : nullptr; const float *const col_amax_base = RETURN_TRANSPOSE ? col_amax_all + tensor_id * cols : nullptr; - group_row_scaled_cast_body( + group_row_scaled_cast_body( tensor_map_input, q_row_base, s_dec_row_base, row_amax_base, q_col_base, s_dec_col_base, col_amax_base, scale_stride_col, M_t, local_row_base, ctaid_X, ctaid_Y, cols, scale_stride_row); @@ -1684,6 +1704,7 @@ inline void group_row_scaled_cast_graph_safe(const GroupedTensor *input, const T TunableConfig::CHUNK_DIM_Y, " (got ", capacity_rows, ")."); const bool return_transpose = output->has_columnwise_data(); + const bool with_gemm_swizzled = output->with_gemm_swizzled_scales; NVTE_CHECK(output->scale_inv.dptr != nullptr, "Graph-safe cast requires rowwise scales."); NVTE_CHECK(output->amax.dptr != nullptr, "Graph-safe cast requires rowwise amax."); // Per-row scale stride, from a 2D [rows, stride] scale or a flat 1D buffer. @@ -1744,21 +1765,23 @@ inline void group_row_scaled_cast_graph_safe(const GroupedTensor *input, const T TRANSFORMER_ENGINE_SWITCH_CONDITION( use_fast_math, USE_FAST_MATH, - TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { - const int out_mem_colwise_data = RETURN_TRANSPOSE ? buff_size_aligned_out_t : 0; - const int out_scales_transpose_mem = RETURN_TRANSPOSE ? buff_size_scales_transpose : 0; - const int dshmem_size = buff_size_aligned_in + buff_size_aligned_out + - out_mem_colwise_data + buff_size_scales + out_scales_transpose_mem + - TMA_SHMEM_ALIGNMENT; - auto kernel = group_row_scaled_cast_nvfp4_graph_safe_kernel; - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - kernel<<>>(tensor_map_input, q_row_all, s_dec_row_all, - row_amax_all, q_col_all, s_dec_col_all, - col_amax_all, offsets_ptr, first_dims_ptr, - num_tensors, noop_ptr, cols, scale_stride_row); - })); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + return_transpose, RETURN_TRANSPOSE, + TRANSFORMER_ENGINE_SWITCH_CONDITION(with_gemm_swizzled, WITH_GEMM_SWIZZLED_SCALES, { + const int out_mem_colwise_data = RETURN_TRANSPOSE ? buff_size_aligned_out_t : 0; + const int out_scales_transpose_mem = RETURN_TRANSPOSE ? buff_size_scales_transpose : 0; + const int dshmem_size = buff_size_aligned_in + buff_size_aligned_out + + out_mem_colwise_data + buff_size_scales + + out_scales_transpose_mem + TMA_SHMEM_ALIGNMENT; + auto kernel = group_row_scaled_cast_nvfp4_graph_safe_kernel< + ScaleType, USE_FAST_MATH, RETURN_TRANSPOSE, WITH_GEMM_SWIZZLED_SCALES>; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + kernel<<>>( + tensor_map_input, q_row_all, s_dec_row_all, row_amax_all, q_col_all, s_dec_col_all, + col_amax_all, offsets_ptr, first_dims_ptr, num_tensors, noop_ptr, cols, + scale_stride_row); + }))); NVTE_CHECK_CUDA(cudaGetLastError()); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 404a6e65613..328d57d6aab 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1237,12 +1237,18 @@ def general_grouped_gemm_for_grouped_tensor( if is_discrete_in and is_discrete_out: raise ValueError("Both A and out are discrete. This is not supported yet.") - if isinstance(A, GroupedTensorStorage) and A.row_scaled_nvfp4: - raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") - if isinstance(B, GroupedTensorStorage) and B.row_scaled_nvfp4: - raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") - if isinstance(out, GroupedTensorStorage) and out.row_scaled_nvfp4: - raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") + if _is_row_scaled_nvfp4(A) or _is_row_scaled_nvfp4(B): + return _grouped_gemm_row_scaled_nvfp4( + A, + B, + out, + layout=layout, + accumulate=accumulate, + use_split_accumulator=use_split_accumulator, + bias=bias, + beta=beta, + alpha=alpha, + ) # Discrete expert tensors share a scaling mode, so inspect only the first one. is_fp8_blockwise_A = ( @@ -1323,3 +1329,130 @@ def general_grouped_gemm_for_grouped_tensor( use_split_accumulator, sm_count, ) + + +def _is_row_scaled_nvfp4(tensors) -> bool: + """Whether a grouped GEMM operand carries row-scaled NVFP4 data.""" + if isinstance(tensors, GroupedTensorStorage): + return tensors.row_scaled_nvfp4 + if isinstance(tensors, (list, tuple)): + return any(_is_nvfp4_row_scaled_tensor(t) for t in tensors) + return False + + +def _grouped_gemm_row_scaled_nvfp4( + A, + B, + out, + *, + layout: str, + accumulate: bool, + use_split_accumulator: bool, + bias, + beta: Optional[torch.Tensor], + alpha: Optional[torch.Tensor], +) -> Union[torch.Tensor, List[torch.Tensor]]: + """Grouped GEMM with row-scaled NVFP4 operands. + + Runs the grouped GEMM with unit global scales into an FP32 buffer, then applies + the per-row scales of B to the output rows and the per-column scales of A to the + output columns. + """ + if bias is not None: + raise NotImplementedError("Row-scaled NVFP4 grouped GEMM does not support fused bias.") + if beta is not None: + raise NotImplementedError( + "Row-scaled NVFP4 grouped GEMM does not support explicit beta; use accumulate." + ) + for operand in (A, B): + if isinstance(operand, (list, tuple)) and _is_row_scaled_nvfp4(operand): + raise NotImplementedError( + "Row-scaled NVFP4 operands passed as a list are not supported; pass a" + " GroupedTensor." + ) + transa = layout[0] == "T" + transb = layout[1] == "T" + + col_scales = None + if _is_row_scaled_nvfp4(A): + A = A.copy() + if transa: + col_scales = A.amax + A.amax = col_scales.new_ones(A.num_tensors) + else: + col_scales = A.columnwise_amax + A.columnwise_amax = col_scales.new_ones(A.num_tensors) + A.row_scaled_nvfp4 = False + + row_scales = None + if _is_row_scaled_nvfp4(B): + B = B.copy() + if transb: + row_scales = B.columnwise_amax + B.columnwise_amax = row_scales.new_ones(B.num_tensors) + else: + row_scales = B.amax + B.amax = row_scales.new_ones(B.num_tensors) + B.row_scaled_nvfp4 = False + + is_discrete_out = isinstance(out, list) + if is_discrete_out: + scratch = [torch.empty_like(o, dtype=torch.float32) for o in out] + else: + scratch = out.copy() + scratch.rowwise_data = torch.empty_like(out.rowwise_data, dtype=torch.float32) + scratch.fake_dtype = torch.float32 + general_grouped_gemm_for_grouped_tensor( + A, + B, + scratch, + layout=layout, + accumulate=False, + use_split_accumulator=use_split_accumulator, + alpha=alpha, + ) + + def _store(dst: torch.Tensor, src: torch.Tensor) -> None: + if accumulate: + dst.add_(src.to(dst.dtype)) + else: + dst.copy_(src.to(dst.dtype)) + + num_tensors = B.num_tensors + if is_discrete_out: + row_offset = 0 + col_offset = 0 + for i in range(num_tensors): + scaled = scratch[i] + rows = scaled.shape[-2] + cols = scaled.shape[-1] + if row_scales is not None: + row_scale_i = row_scales[row_offset : row_offset + rows] + scaled = scaled.mul_(row_scale_i.reshape((1,) * (scaled.dim() - 2) + (rows, 1))) + if col_scales is not None: + col_scale_i = col_scales[col_offset : col_offset + cols] + scaled = scaled.mul_(col_scale_i.reshape((1,) * (scaled.dim() - 1) + (cols,))) + _store(out[i], scaled) + row_offset += rows + col_offset += cols + else: + if out.last_dims is not None: + raise NotImplementedError( + "Row-scaled grouped GEMM does not support a varying per-expert output width." + ) + last_dim = out.logical_shape[-1] + live_rows = out.logical_shape[0] + flat = scratch.rowwise_data.view(-1, last_dim)[:live_rows] + if row_scales is not None: + flat.mul_(row_scales.view(-1, 1)) + if col_scales is not None: + # output_size avoids a host sync under CUDA-graph capture. + col_scales = col_scales.view(num_tensors, last_dim) + if out.first_dims is None: + rows = live_rows // num_tensors + repeats = torch.full((num_tensors,), rows, dtype=torch.int64, device=flat.device) + else: + repeats = out.first_dims + flat.mul_(col_scales.repeat_interleave(repeats, dim=0, output_size=live_rows)) + _store(out.rowwise_data.view(-1, last_dim)[:live_rows], flat) + return out diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8c9b6b1079d..09b82351885 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -339,8 +339,12 @@ py::object group_quantize(const at::Tensor &input, py::handle quantizer, const s NVTE_CHECK(output.is_none(), "group_quantize: output reuse is not supported with row-scaled NVFP4."); if (first_dims.has_value()) { - NVTE_CHECK(first_dims->is_contiguous() && first_dims->scalar_type() == at::kLong, - "group_quantize: first_dims must be a contiguous int64 tensor."); + NVTE_CHECK(first_dims->is_cuda() && first_dims->is_contiguous() && + first_dims->scalar_type() == at::kLong, + "group_quantize: first_dims must be a contiguous int64 CUDA tensor."); + NVTE_CHECK(static_cast(first_dims->numel()) == num_tensors, + "group_quantize: first_dims has ", first_dims->numel(), + " entries, but num_tensors is ", num_tensors, "."); } else { NVTE_CHECK(num_tensors > 0 && logical_first_dim % (num_tensors * 128) == 0, "group_quantize: row-scaled NVFP4 with uniform shapes requires the first " diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 74f71713e51..180e4580207 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2206,8 +2206,6 @@ std::pair NVFP4Quantizer::create_grouped_tenso const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 grouped quantization requires rowwise usage."); - NVTE_CHECK(!columnwise_usage, - "Row-scaled NVFP4 grouped quantization does not support columnwise usage."); } const int64_t total_data_elements = total_elements / 2; @@ -2230,7 +2228,11 @@ std::pair NVFP4Quantizer::create_grouped_tenso const int64_t total_scale_elements = static_cast(product(scale_shape)); columnwise_scale_inv = at::empty({total_scale_elements}, uint8_opts); if (!disable_second_level_scale) { - columnwise_amax = at::empty({static_cast(num_tensors)}, float_opts); + const int64_t col_amax_elements = + row_scaled_nvfp4 + ? static_cast(num_tensors) * static_cast(logical_last_dim) + : static_cast(num_tensors); + columnwise_amax = at::empty({col_amax_elements}, float_opts); } } diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index d85146351a8..6f70fc3e9d3 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -860,13 +860,12 @@ def make_grouped_tensor( raise ValueError( "Row-scaled NVFP4 grouped quantization requires rowwise usage." ) - if columnwise_usage: - raise ValueError( - "Row-scaled NVFP4 grouped quantization does not support columnwise usage." - ) total_amax_elements = ( sum(math.prod(s[:-1]) for s in shape) if row_scaled_nvfp4 else num_tensors ) + total_columnwise_amax_elements = ( + sum(s[-1] for s in shape) if row_scaled_nvfp4 else num_tensors + ) if rowwise_usage: # Allocate rowwise data buffer (1D flattened, uint8, but FP4 packs 2 values per byte) @@ -899,7 +898,9 @@ def make_grouped_tensor( total_columnwise_scale_elements, dtype=torch.uint8, device=device ) if not disable_second_level_scale: - columnwise_amax = torch.empty(num_tensors, dtype=torch.float32, device=device) + columnwise_amax = torch.empty( + total_columnwise_amax_elements, dtype=torch.float32, device=device + ) elif compatible_recipe.float8_block_scaling(): scale_inv_dtype = DType.kFloat32 @@ -1136,17 +1137,29 @@ def flattened_buffer(buffer: Optional[torch.Tensor]) -> Optional[torch.Tensor]: cum += math.prod(scale_shape) columnwise_scale_inv_offsets.append(cum) self.columnwise_scale_inv_offsets = columnwise_scale_inv_offsets - nvfp4_rowwise_amax_offsets = None + row_scaled_nvfp4 = self.row_scaled_nvfp4 nvfp4_use_4over6 = self.nvfp4_use_4over6 nvfp4_e4m3_max = self.nvfp4_e4m3_max scale_inv_dtype = self.scale_inv_dtype - if recipe.nvfp4() and row_scaled_nvfp4: - cum = 0 - nvfp4_rowwise_amax_offsets = [0] - for i in range(self.num_tensors): - cum += math.prod(self.tensor_shapes[i][:-1]) - nvfp4_rowwise_amax_offsets.append(cum) + + # Determine amax offsets for NVFP4 tensors + nvfp4_rowwise_amax_offsets = None + nvfp4_columnwise_amax_offsets = None + if recipe.nvfp4(): + if row_scaled_nvfp4: + off_row = 0 + off_col = 0 + nvfp4_rowwise_amax_offsets = [0] + nvfp4_columnwise_amax_offsets = [0] + for shape in self.tensor_shapes: + off_row += math.prod(shape[:-1]) + off_col += shape[-1] + nvfp4_rowwise_amax_offsets.append(off_row) + nvfp4_columnwise_amax_offsets.append(off_col) + else: + nvfp4_rowwise_amax_offsets = list(range(self.num_tensors + 1)) + nvfp4_columnwise_amax_offsets = list(range(self.num_tensors + 1)) for i in range(self.num_tensors): quantizer = self.quantizer @@ -1342,15 +1355,12 @@ def flattened_buffer(buffer: Optional[torch.Tensor]) -> Optional[torch.Tensor]: ].view(cscale_shape) if self.amax is not None: - if nvfp4_rowwise_amax_offsets is not None: - amax_start = nvfp4_rowwise_amax_offsets[i] - amax_end = nvfp4_rowwise_amax_offsets[i + 1] - amax_rowwise = amax_buffer[amax_start:amax_end] - else: - amax_rowwise = amax_buffer[i : i + 1] + amax_start, amax_end = nvfp4_rowwise_amax_offsets[i : i + 2] + amax_rowwise = amax_buffer[amax_start:amax_end] if self.columnwise_amax is not None: - amax_columnwise = columnwise_amax_buffer[i : i + 1] + amax_start, amax_end = nvfp4_columnwise_amax_offsets[i : i + 2] + amax_columnwise = columnwise_amax_buffer[amax_start:amax_end] if quantizer.internal: nvfp4_tensor_class = NVFP4TensorStorage