From 148f681a24db93d563673c51385def39c1460053 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Wed, 7 Oct 2026 23:57:34 -0700 Subject: [PATCH 01/24] [Common][PyTorch] Add NVFP4 row-scaled grouped cast with graph-safe path Single-launch grouped row-scaled NVFP4 cast + fused amax, plus graph-safe amax/cast variants driven by device first_dims/tensor_offsets for capture-once/replay MoE. Wired into PyTorch (ep.nvfp4_group_row_scaled_cast) with C++ and PyTorch capture/replay tests. Signed-off-by: Cael Ling --- tests/cpp/operator/CMakeLists.txt | 1 + .../test_row_scaled_nvfp4_grouped_cast.cu | 759 +++++++++++++++++ tests/cpp/test_common.h | 12 +- ..._nvfp4_group_row_scaled_cast_graph_safe.py | 154 ++++ .../common/cast/cast_grouped.cu | 93 +++ .../nvfp4/group_quantize_transpose_nvfp4.cuh | 760 +++++++++++++++++- .../common/include/transformer_engine/cast.h | 68 ++ transformer_engine/pytorch/csrc/extensions.h | 5 + .../pytorch/csrc/extensions/cast.cpp | 135 +++- .../pytorch/csrc/extensions/pybind.cpp | 5 + transformer_engine/pytorch/ep.py | 43 + 11 files changed, 2012 insertions(+), 23 deletions(-) create mode 100644 tests/cpp/operator/test_row_scaled_nvfp4_grouped_cast.cu create mode 100644 tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index c5acd043363..3732c1c5fed 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -16,6 +16,7 @@ add_executable(test_operator test_cast_mxfp8_grouped_scaled_swiglu.cu test_cast_nvfp4_transpose.cu test_row_scaled_nvfp4_grouped_amax.cu + test_row_scaled_nvfp4_grouped_cast.cu test_cast_float8blockwise.cu test_cast_float8blockwise_grouped.cu test_dequantize_mxfp8.cu diff --git a/tests/cpp/operator/test_row_scaled_nvfp4_grouped_cast.cu b/tests/cpp/operator/test_row_scaled_nvfp4_grouped_cast.cu new file mode 100644 index 00000000000..9e3af23fdcf --- /dev/null +++ b/tests/cpp/operator/test_row_scaled_nvfp4_grouped_cast.cu @@ -0,0 +1,759 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include +#include "../test_common.h" +#include "transformer_engine/transformer_engine.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +// (per-expert row counts, shared column count, request transpose); all values +// are multiples of 128 so every 128-row chunk stays inside a single expert. +struct GroupedCastConfig { + std::vector Ms; + size_t K; + bool transpose; +}; + +std::vector> make_outputs(const std::string& tag, + const std::vector& Ms, size_t K, + bool transpose) { + std::vector> outs; + for (size_t i = 0; i < Ms.size(); ++i) { + const std::vector shape = {Ms[i], K}; + auto t = std::make_unique(tag + std::to_string(i), shape, DType::kFloat4E2M1, + /*rowwise=*/true, /*columnwise=*/transpose, + NVTE_NVFP4_1D_SCALING); + t->set_row_scaled_nvfp4(true); + outs.push_back(std::move(t)); + } + return outs; +} + +// The grouped row-scaled cast reuses the single-tensor encode verbatim, so for +// every expert its FP4 data, E4M3 block scales and amaxes must be bit-identical +// to quantizing that expert on its own via nvte_quantize_v2. +void performGroupedRowScaledCastTest(const GroupedCastConfig& cfg) { + const std::vector& Ms = cfg.Ms; + const size_t K = cfg.K; + const bool transpose = cfg.transpose; + const size_t num_tensors = Ms.size(); + size_t sum_M = 0; + for (size_t m : Ms) sum_M += m; + + // Per-expert BF16 inputs, packed into one buffer with identical bytes so the + // grouped kernel and the per-expert oracle quantize the same values. + std::vector> ins; + for (size_t i = 0; i < num_tensors; ++i) { + const std::vector shape = {Ms[i], K}; + ins.push_back(std::make_unique("in_" + std::to_string(i), shape, DType::kBFloat16)); + fillCase(ins[i].get(), InputsFillCase::uniform); + ins[i]->to_cpu(); + } + + Tensor packed("packed_input", std::vector{sum_M, K}, DType::kBFloat16); + bf16* pdst = packed.rowwise_cpu_dptr(); + size_t row_off = 0; + for (size_t i = 0; i < num_tensors; ++i) { + std::copy_n(ins[i]->rowwise_cpu_dptr(), Ms[i] * K, pdst + row_off * K); + row_off += Ms[i]; + } + packed.from_cpu(); + + const std::vector splits(Ms); + + // System under test: grouped amax (Stage 1) + grouped row-scaled cast (Stage 2). + auto grouped = make_outputs("grp_", Ms, K, transpose); + std::vector grouped_handles; + for (auto& t : grouped) grouped_handles.push_back(t->data()); + nvte_group_nvfp4_compute_amax(packed.data(), grouped_handles.data(), splits.data(), num_tensors, + 0); + nvte_group_nvfp4_row_scaled_cast_with_amax(packed.data(), grouped_handles.data(), splits.data(), + num_tensors, nullptr, 0); + + // Oracle: quantize each expert independently (compute_amaxes + tuned_1D) via + // the generic nvte_quantize_v2 dispatch. + auto oracle = make_outputs("ref_", Ms, K, transpose); + QuantizationConfigWrapper quant_config; + quant_config.set_stochastic_rounding(false); + for (size_t i = 0; i < num_tensors; ++i) { + nvte_quantize_v2(ins[i]->data(), oracle[i]->data(), quant_config, 0); + } + + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + + for (size_t i = 0; i < num_tensors; ++i) { + grouped[i]->to_cpu(); + oracle[i]->to_cpu(); + const size_t M = Ms[i]; + + // Rowwise FP4 data: M x K/2 packed bytes, fully written (no padding since + // K % 128 == 0). Read as fp4e2m1 (the tensor dtype) then compare raw bytes. + { + const uint8_t* g = reinterpret_cast(grouped[i]->rowwise_cpu_dptr()); + const uint8_t* r = reinterpret_cast(oracle[i]->rowwise_cpu_dptr()); + size_t mismatches = 0; + for (size_t b = 0; b < M * K / 2; ++b) { + if (g[b] != r[b]) ++mismatches; + } + ASSERT_EQ(mismatches, 0u) << "rowwise FP4 data mismatch: expert " << i << " (" << mismatches + << "/" << (M * K / 2) << " bytes)"; + } + + // Rowwise E4M3 block scales. + { + const std::array sd = get_scale_tensor_dims(M, K, 1, 16); + size_t mismatches = 0; + compare_scaling_factors("rowwise_scales_" + std::to_string(i), + grouped[i]->rowwise_cpu_scale_inv_ptr(), + oracle[i]->rowwise_cpu_scale_inv_ptr(), sd[0], sd[1], + sd[3], mismatches); + ASSERT_EQ(mismatches, 0u) << "rowwise scale mismatch: expert " << i; + } + + // Rowwise per-row amax. + { + const float* g = grouped[i]->cpu_rowwise_amax_ptr(); + const float* r = oracle[i]->cpu_rowwise_amax_ptr(); + for (size_t row = 0; row < M; ++row) { + ASSERT_EQ(g[row], r[row]) << "rowwise amax mismatch: expert " << i << " row " << row; + } + } + + if (!transpose) continue; + + // Columnwise (transpose) FP4 data: K x M/2 packed bytes. + { + const uint8_t* g = + reinterpret_cast(grouped[i]->columnwise_cpu_dptr()); + const uint8_t* r = + reinterpret_cast(oracle[i]->columnwise_cpu_dptr()); + size_t mismatches = 0; + for (size_t b = 0; b < K * M / 2; ++b) { + if (g[b] != r[b]) ++mismatches; + } + ASSERT_EQ(mismatches, 0u) << "columnwise FP4 data mismatch: expert " << i << " (" << mismatches + << "/" << (K * M / 2) << " bytes)"; + } + + // Columnwise E4M3 block scales. + { + const std::array sd_t = get_scale_tensor_dims(K, M, 1, 16); + size_t mismatches = 0; + compare_scaling_factors("columnwise_scales_" + std::to_string(i), + grouped[i]->columnwise_cpu_scale_inv_ptr(), + oracle[i]->columnwise_cpu_scale_inv_ptr(), sd_t[0], + sd_t[1], sd_t[3], mismatches); + ASSERT_EQ(mismatches, 0u) << "columnwise scale mismatch: expert " << i; + } + + // Columnwise per-col amax. + { + const float* g = grouped[i]->cpu_columnwise_amax_ptr(); + const float* r = oracle[i]->cpu_columnwise_amax_ptr(); + for (size_t col = 0; col < K; ++col) { + ASSERT_EQ(g[col], r[col]) << "columnwise amax mismatch: expert " << i << " col " << col; + } + } + } +} + +// Host-split reference for one routing, gathered into the contiguous per-direction +// layout the graph-safe kernel writes. +struct GroupedReference { + size_t sum_M = 0; + size_t K = 0; + size_t row_scale_cols = 0; // rowwise E4M3 scale stride (shared across experts) + std::vector row_data; + std::vector row_scale; + std::vector row_amax; + std::vector col_data; + std::vector col_scale; + std::vector col_amax; + bool transpose = false; +}; + +std::vector pack_inputs(const std::vector>& ins, + const std::vector& Ms, size_t K) { + size_t sum_M = 0; + for (size_t m : Ms) sum_M += m; + std::vector packed(sum_M * K); + size_t row_off = 0; + for (size_t i = 0; i < ins.size(); ++i) { + std::copy_n(ins[i]->rowwise_cpu_dptr(), Ms[i] * K, packed.data() + row_off * K); + row_off += Ms[i]; + } + return packed; +} + +GroupedReference build_reference(const std::vector>& ins, + const std::vector& Ms, size_t K, bool transpose) { + const size_t num_tensors = Ms.size(); + size_t sum_M = 0; + for (size_t m : Ms) sum_M += m; + + const std::vector packed_host = pack_inputs(ins, Ms, K); + Tensor packed("ref_packed", std::vector{sum_M, K}, DType::kBFloat16); + std::copy_n(packed_host.data(), sum_M * K, packed.rowwise_cpu_dptr()); + packed.from_cpu(); + + const std::vector splits(Ms); + auto ref = make_outputs("gsref_", Ms, K, transpose); + std::vector handles; + for (auto& t : ref) handles.push_back(t->data()); + nvte_group_nvfp4_compute_amax(packed.data(), handles.data(), splits.data(), num_tensors, 0); + nvte_group_nvfp4_row_scaled_cast_with_amax(packed.data(), handles.data(), splits.data(), + num_tensors, nullptr, 0); + cudaDeviceSynchronize(); + + GroupedReference r; + r.sum_M = sum_M; + r.K = K; + r.transpose = transpose; + r.row_scale_cols = K / 16; // contiguous (unpadded) rowwise E4M3 scale stride + + r.row_data.resize(sum_M * K / 2); + r.row_scale.resize(sum_M * r.row_scale_cols); + r.row_amax.resize(sum_M); + if (transpose) { + r.col_data.resize((K / 2) * sum_M); + r.col_scale.resize((K / 16) * sum_M); + r.col_amax.resize(num_tensors * K); + } + + // Gather per-expert scales into the contiguous (unpadded) layout the graph-safe + // kernel writes, stripping any allocation padding in the source stride. + auto gather_scale = [](const fp8e4m3* src, fp8e4m3* dst, size_t rows, size_t cols, + size_t src_stride) { + for (size_t row = 0; row < rows; ++row) { + std::copy_n(src + row * src_stride, cols, dst + row * cols); + } + }; + + size_t row_off = 0; + size_t col_data_off = 0; + size_t col_scale_off = 0; + for (size_t i = 0; i < num_tensors; ++i) { + ref[i]->to_cpu(); + const size_t M = Ms[i]; + std::copy_n(reinterpret_cast(ref[i]->rowwise_cpu_dptr()), M * K / 2, + r.row_data.data() + row_off * (K / 2)); + const std::array sd = get_scale_tensor_dims(M, K, 1, 16); + gather_scale(ref[i]->rowwise_cpu_scale_inv_ptr(), + reinterpret_cast(r.row_scale.data()) + row_off * r.row_scale_cols, sd[0], + sd[1], sd[3]); + std::copy_n(ref[i]->cpu_rowwise_amax_ptr(), M, r.row_amax.data() + row_off); + if (transpose) { + std::copy_n(reinterpret_cast(ref[i]->columnwise_cpu_dptr()), + K * M / 2, r.col_data.data() + col_data_off); + const std::array sd_t = get_scale_tensor_dims(K, M, 1, 16); + gather_scale(ref[i]->columnwise_cpu_scale_inv_ptr(), + reinterpret_cast(r.col_scale.data()) + col_scale_off, sd_t[0], sd_t[1], + sd_t[3]); + std::copy_n(ref[i]->cpu_columnwise_amax_ptr(), K, r.col_amax.data() + i * K); + col_data_off += K * M / 2; + col_scale_off += K * (M / 16); + } + row_off += M; + } + return r; +} + +// Contiguous device buffers for one input/output grouped tensor pair, sized to a +// fixed capacity so a captured graph can be replayed with a redistributed routing. +struct GraphSafeBuffers { + size_t capacity_rows = 0; + size_t K = 0; + size_t num_tensors = 0; + size_t row_scale_cols = 0; + bool transpose = false; + + CudaPtr in_data; + CudaPtr out_data; + CudaPtr out_scale; + CudaPtr out_amax; + CudaPtr out_col_data; + CudaPtr out_col_scale; + CudaPtr out_col_amax; + CudaPtr offsets; + CudaPtr first_dims; + + std::unique_ptr input; + std::unique_ptr output; + + void allocate(size_t capacity_rows_, size_t K_, size_t num_tensors_, size_t row_scale_cols_, + bool transpose_) { + capacity_rows = capacity_rows_; + K = K_; + num_tensors = num_tensors_; + row_scale_cols = row_scale_cols_; + transpose = transpose_; + + in_data = cuda_alloc(capacity_rows * K * sizeof(bf16)); + out_data = cuda_alloc(capacity_rows * (K / 2)); + out_scale = cuda_alloc(capacity_rows * row_scale_cols); + out_amax = cuda_alloc(capacity_rows * sizeof(float)); + offsets = cuda_alloc((num_tensors + 1) * sizeof(int64_t)); + first_dims = cuda_alloc(num_tensors * sizeof(int64_t)); + if (transpose) { + out_col_data = cuda_alloc((K / 2) * capacity_rows); + out_col_scale = cuda_alloc((K / 16) * capacity_rows); + out_col_amax = cuda_alloc(num_tensors * K * sizeof(float)); + } + + input = std::make_unique(num_tensors, + std::vector{capacity_rows, K}); + output = std::make_unique(num_tensors, + std::vector{capacity_rows, K}); + + input->set_rowwise_data(in_data.get(), DType::kBFloat16, + std::vector{capacity_rows, K}); + input->set_tensor_offsets(offsets.get(), DType::kInt64, + std::vector{num_tensors + 1}); + input->set_first_dims(first_dims.get(), DType::kInt64, std::vector{num_tensors}); + + output->set_rowwise_data(out_data.get(), DType::kFloat4E2M1, + std::vector{capacity_rows, K}); + output->set_rowwise_scale_inv(out_scale.get(), DType::kFloat8E4M3, + std::vector{capacity_rows, row_scale_cols}); + output->set_amax(out_amax.get(), DType::kFloat32, std::vector{capacity_rows}); + output->set_tensor_offsets(offsets.get(), DType::kInt64, + std::vector{num_tensors + 1}); + output->set_first_dims(first_dims.get(), DType::kInt64, std::vector{num_tensors}); + if (transpose) { + output->set_columnwise_data(out_col_data.get(), DType::kFloat4E2M1, + std::vector{K, capacity_rows}); + output->set_columnwise_scale_inv(out_col_scale.get(), DType::kFloat8E4M3, + std::vector{K, capacity_rows / 16}); + output->set_columnwise_amax(out_col_amax.get(), DType::kFloat32, + std::vector{num_tensors * K}); + } + } + + // Program device metadata + input for one routing. seed_amax: fill amax from the + // reference (cast-only tests) or poison it (so a missing amax write is caught). + void set_routing(const std::vector& Ms, const GroupedReference& ref, + const std::vector& packed_in, bool seed_amax = true) { + std::vector off(num_tensors + 1, 0); + std::vector fd(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + fd[i] = static_cast(Ms[i]); + off[i + 1] = off[i] + static_cast(Ms[i] * K); + } + NVTE_CHECK_CUDA(cudaMemcpy(offsets.get(), off.data(), (num_tensors + 1) * sizeof(int64_t), + cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(first_dims.get(), fd.data(), num_tensors * sizeof(int64_t), + cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(in_data.get(), packed_in.data(), ref.sum_M * K * sizeof(bf16), + cudaMemcpyHostToDevice)); + if (seed_amax) { + NVTE_CHECK_CUDA(cudaMemcpy(out_amax.get(), ref.row_amax.data(), ref.sum_M * sizeof(float), + cudaMemcpyHostToDevice)); + if (transpose) { + NVTE_CHECK_CUDA(cudaMemcpy(out_col_amax.get(), ref.col_amax.data(), + num_tensors * K * sizeof(float), cudaMemcpyHostToDevice)); + } + } else { + std::vector poison_row(capacity_rows, -1.0f); + NVTE_CHECK_CUDA(cudaMemcpy(out_amax.get(), poison_row.data(), capacity_rows * sizeof(float), + cudaMemcpyHostToDevice)); + if (transpose) { + std::vector poison_col(num_tensors * K, -1.0f); + NVTE_CHECK_CUDA(cudaMemcpy(out_col_amax.get(), poison_col.data(), + num_tensors * K * sizeof(float), cudaMemcpyHostToDevice)); + } + } + } +}; + +// The graph-safe amax must reproduce the host-split amax exactly: rowwise amax +// packed by global row [sum_M], columnwise amax packed per expert [num_tensors*K]. +void expect_amax_matches_reference(const GraphSafeBuffers& b, const GroupedReference& ref) { + const size_t sum_M = ref.sum_M; + std::vector h_amax(sum_M); + NVTE_CHECK_CUDA( + cudaMemcpy(h_amax.data(), b.out_amax.get(), sum_M * sizeof(float), cudaMemcpyDeviceToHost)); + for (size_t i = 0; i < sum_M; ++i) + ASSERT_EQ(h_amax[i], ref.row_amax[i]) << "rowwise amax mismatch at row " << i; + + if (!ref.transpose) return; + std::vector h_col_amax(b.num_tensors * ref.K); + NVTE_CHECK_CUDA(cudaMemcpy(h_col_amax.data(), b.out_col_amax.get(), + b.num_tensors * ref.K * sizeof(float), cudaMemcpyDeviceToHost)); + for (size_t i = 0; i < b.num_tensors * ref.K; ++i) + ASSERT_EQ(h_col_amax[i], ref.col_amax[i]) << "columnwise amax mismatch at index " << i; +} + +void expect_matches_reference(const GraphSafeBuffers& b, const GroupedReference& ref) { + const size_t sum_M = ref.sum_M; + const size_t K = ref.K; + + std::vector h_data(sum_M * K / 2); + std::vector h_scale(sum_M * ref.row_scale_cols); + std::vector h_amax(sum_M); + NVTE_CHECK_CUDA(cudaMemcpy(h_data.data(), b.out_data.get(), h_data.size(), cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA( + cudaMemcpy(h_scale.data(), b.out_scale.get(), h_scale.size(), cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA( + cudaMemcpy(h_amax.data(), b.out_amax.get(), sum_M * sizeof(float), cudaMemcpyDeviceToHost)); + + size_t data_mismatch = 0; + for (size_t i = 0; i < h_data.size(); ++i) + if (h_data[i] != ref.row_data[i]) ++data_mismatch; + ASSERT_EQ(data_mismatch, 0u) << "rowwise FP4 data mismatch (" << data_mismatch << "/" + << h_data.size() << ")"; + size_t scale_mismatch = 0; + for (size_t i = 0; i < h_scale.size(); ++i) + if (h_scale[i] != ref.row_scale[i]) ++scale_mismatch; + ASSERT_EQ(scale_mismatch, 0u) << "rowwise scale mismatch (" << scale_mismatch << "/" + << h_scale.size() << ")"; + for (size_t i = 0; i < sum_M; ++i) + ASSERT_EQ(h_amax[i], ref.row_amax[i]) << "rowwise amax mismatch at row " << i; + + if (!ref.transpose) return; + + std::vector h_col_data((K / 2) * sum_M); + std::vector h_col_scale((K / 16) * sum_M); + NVTE_CHECK_CUDA(cudaMemcpy(h_col_data.data(), b.out_col_data.get(), h_col_data.size(), + cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(h_col_scale.data(), b.out_col_scale.get(), h_col_scale.size(), + cudaMemcpyDeviceToHost)); + size_t col_data_mismatch = 0; + for (size_t i = 0; i < h_col_data.size(); ++i) + if (h_col_data[i] != ref.col_data[i]) ++col_data_mismatch; + ASSERT_EQ(col_data_mismatch, 0u) << "columnwise FP4 data mismatch (" << col_data_mismatch << "/" + << h_col_data.size() << ")"; + size_t col_scale_mismatch = 0; + for (size_t i = 0; i < h_col_scale.size(); ++i) + if (h_col_scale[i] != ref.col_scale[i]) ++col_scale_mismatch; + ASSERT_EQ(col_scale_mismatch, 0u) << "columnwise scale mismatch (" << col_scale_mismatch << "/" + << h_col_scale.size() << ")"; +} + +std::vector> make_random_inputs(const std::vector& Ms, size_t K) { + std::vector> ins; + for (size_t i = 0; i < Ms.size(); ++i) { + ins.push_back( + std::make_unique("gsin_" + std::to_string(i), std::vector{Ms[i], K}, + DType::kBFloat16)); + fillCase(ins[i].get(), InputsFillCase::uniform); + ins[i]->to_cpu(); + } + return ins; +} + +// Graph-safe cast must be byte-identical to the host-split cast. +void performGraphSafeEqualityTest(const GroupedCastConfig& cfg) { + const std::vector& Ms = cfg.Ms; + const size_t K = cfg.K; + const bool transpose = cfg.transpose; + const size_t num_tensors = Ms.size(); + size_t sum_M = 0; + for (size_t m : Ms) sum_M += m; + + auto ins = make_random_inputs(Ms, K); + const GroupedReference ref = build_reference(ins, Ms, K, transpose); + const std::vector packed_in = pack_inputs(ins, Ms, K); + + GraphSafeBuffers b; + b.allocate(sum_M, K, num_tensors, ref.row_scale_cols, transpose); + b.set_routing(Ms, ref, packed_in); + + nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(b.input->data(), b.output->data(), 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + + expect_matches_reference(b, ref); +} + +// Capture the graph-safe cast once, then replay after redistributing the routing +// on device; both replays must match their host-split reference. +void performGraphSafeReplayTest(size_t K, bool transpose) { + const std::vector Ms_a = {256, 128, 128}; + const std::vector Ms_b = {128, 256, 128}; + const size_t num_tensors = Ms_a.size(); + const size_t capacity_rows = 768; // > sum_M (512): exercises the capacity noop + + auto ins_a = make_random_inputs(Ms_a, K); + auto ins_b = make_random_inputs(Ms_b, K); + const GroupedReference ref_a = build_reference(ins_a, Ms_a, K, transpose); + const GroupedReference ref_b = build_reference(ins_b, Ms_b, K, transpose); + const std::vector packed_a = pack_inputs(ins_a, Ms_a, K); + const std::vector packed_b = pack_inputs(ins_b, Ms_b, K); + + GraphSafeBuffers b; + b.allocate(capacity_rows, K, num_tensors, ref_a.row_scale_cols, transpose); + b.set_routing(Ms_a, ref_a, packed_a); + + cudaStream_t stream; + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); + + // Warm up outside capture so one-time setup (e.g. cudaFuncSetAttribute) does + // not land in the captured graph. + nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(b.input->data(), b.output->data(), stream); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + cudaGraph_t graph; + cudaGraphExec_t exec; + NVTE_CHECK_CUDA(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed)); + nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(b.input->data(), b.output->data(), stream); + NVTE_CHECK_CUDA(cudaStreamEndCapture(stream, &graph)); + NVTE_CHECK_CUDA(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0)); + + NVTE_CHECK_CUDA(cudaGraphLaunch(exec, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + expect_matches_reference(b, ref_a); + + // Redistribute the routing on device and replay the same captured graph. + b.set_routing(Ms_b, ref_b, packed_b); + NVTE_CHECK_CUDA(cudaGraphLaunch(exec, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + expect_matches_reference(b, ref_b); + + NVTE_CHECK_CUDA(cudaGraphExecDestroy(exec)); + NVTE_CHECK_CUDA(cudaGraphDestroy(graph)); + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// Graph-safe amax must be byte-identical to the host-split amax. +void performGraphSafeAmaxEqualityTest(const GroupedCastConfig& cfg) { + const std::vector& Ms = cfg.Ms; + const size_t K = cfg.K; + const bool transpose = cfg.transpose; + const size_t num_tensors = Ms.size(); + size_t sum_M = 0; + for (size_t m : Ms) sum_M += m; + + auto ins = make_random_inputs(Ms, K); + const GroupedReference ref = build_reference(ins, Ms, K, transpose); + const std::vector packed_in = pack_inputs(ins, Ms, K); + + GraphSafeBuffers b; + b.allocate(sum_M, K, num_tensors, ref.row_scale_cols, transpose); + b.set_routing(Ms, ref, packed_in, /*seed_amax=*/false); + + nvte_group_nvfp4_compute_amax_graph_safe(b.input->data(), b.output->data(), 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + + expect_amax_matches_reference(b, ref); +} + +// Full graph-safe path (amax then cast, no host-seeded amax); must be byte-identical +// to the host-split reference. +void performGraphSafeFullEqualityTest(const GroupedCastConfig& cfg) { + const std::vector& Ms = cfg.Ms; + const size_t K = cfg.K; + const bool transpose = cfg.transpose; + const size_t num_tensors = Ms.size(); + size_t sum_M = 0; + for (size_t m : Ms) sum_M += m; + + auto ins = make_random_inputs(Ms, K); + const GroupedReference ref = build_reference(ins, Ms, K, transpose); + const std::vector packed_in = pack_inputs(ins, Ms, K); + + GraphSafeBuffers b; + b.allocate(sum_M, K, num_tensors, ref.row_scale_cols, transpose); + b.set_routing(Ms, ref, packed_in, /*seed_amax=*/false); + + nvte_group_nvfp4_compute_amax_graph_safe(b.input->data(), b.output->data(), 0); + nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(b.input->data(), b.output->data(), 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + + expect_amax_matches_reference(b, ref); + expect_matches_reference(b, ref); +} + +// Capture the full graph-safe path (amax + cast) once, then replay after +// redistributing the routing on device; amax is recomputed on every replay. +void performGraphSafeFullReplayTest(size_t K, bool transpose) { + const std::vector Ms_a = {256, 128, 128}; + const std::vector Ms_b = {128, 256, 128}; + const size_t num_tensors = Ms_a.size(); + const size_t capacity_rows = 768; // > sum_M (512): exercises the capacity noop + + auto ins_a = make_random_inputs(Ms_a, K); + auto ins_b = make_random_inputs(Ms_b, K); + const GroupedReference ref_a = build_reference(ins_a, Ms_a, K, transpose); + const GroupedReference ref_b = build_reference(ins_b, Ms_b, K, transpose); + const std::vector packed_a = pack_inputs(ins_a, Ms_a, K); + const std::vector packed_b = pack_inputs(ins_b, Ms_b, K); + + GraphSafeBuffers b; + b.allocate(capacity_rows, K, num_tensors, ref_a.row_scale_cols, transpose); + b.set_routing(Ms_a, ref_a, packed_a, /*seed_amax=*/false); + + cudaStream_t stream; + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); + + // Warm up outside capture so one-time setup (e.g. cudaFuncSetAttribute) does + // not land in the captured graph. + nvte_group_nvfp4_compute_amax_graph_safe(b.input->data(), b.output->data(), stream); + nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(b.input->data(), b.output->data(), stream); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + cudaGraph_t graph; + cudaGraphExec_t exec; + NVTE_CHECK_CUDA(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed)); + nvte_group_nvfp4_compute_amax_graph_safe(b.input->data(), b.output->data(), stream); + nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(b.input->data(), b.output->data(), stream); + NVTE_CHECK_CUDA(cudaStreamEndCapture(stream, &graph)); + NVTE_CHECK_CUDA(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0)); + + NVTE_CHECK_CUDA(cudaGraphLaunch(exec, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + expect_amax_matches_reference(b, ref_a); + expect_matches_reference(b, ref_a); + + // Redistribute the routing on device and replay the same captured graph. + b.set_routing(Ms_b, ref_b, packed_b, /*seed_amax=*/false); + NVTE_CHECK_CUDA(cudaGraphLaunch(exec, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + expect_amax_matches_reference(b, ref_b); + expect_matches_reference(b, ref_b); + + NVTE_CHECK_CUDA(cudaGraphExecDestroy(exec)); + NVTE_CHECK_CUDA(cudaGraphDestroy(graph)); + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); +} + +} // namespace + +class NVFP4GroupedRowScaledCastTestSuite : public ::testing::TestWithParam {}; + +TEST_P(NVFP4GroupedRowScaledCastTestSuite, MatchesPerExpertOracle) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + performGroupedRowScaledCastTest(GetParam()); +} + +INSTANTIATE_TEST_SUITE_P( + NVFP4GroupedRowScaledCastRowwise, NVFP4GroupedRowScaledCastTestSuite, + ::testing::Values(GroupedCastConfig{{128}, 128, false}, + GroupedCastConfig{{128, 128}, 128, false}, + GroupedCastConfig{{128, 256, 128}, 256, false}, + GroupedCastConfig{{256, 128}, 512, false}, + GroupedCastConfig{{384, 128, 256}, 128, false}, + GroupedCastConfig{{512}, 1024, false})); + +INSTANTIATE_TEST_SUITE_P( + NVFP4GroupedRowScaledCastTranspose, NVFP4GroupedRowScaledCastTestSuite, + ::testing::Values(GroupedCastConfig{{128}, 128, true}, + GroupedCastConfig{{128, 128}, 128, true}, + GroupedCastConfig{{128, 256, 128}, 256, true}, + GroupedCastConfig{{256, 128}, 512, true}, + GroupedCastConfig{{384, 128, 256}, 128, true}, + GroupedCastConfig{{512}, 1024, true})); + +class NVFP4GraphSafeGroupedRowScaledCastTestSuite + : public ::testing::TestWithParam {}; + +TEST_P(NVFP4GraphSafeGroupedRowScaledCastTestSuite, MatchesHostSplit) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + performGraphSafeEqualityTest(GetParam()); +} + +INSTANTIATE_TEST_SUITE_P( + NVFP4GraphSafeGroupedRowScaledCast, NVFP4GraphSafeGroupedRowScaledCastTestSuite, + ::testing::Values(GroupedCastConfig{{128}, 128, false}, + GroupedCastConfig{{128, 256, 128}, 256, false}, + GroupedCastConfig{{384, 128, 256}, 128, false}, + GroupedCastConfig{{128}, 128, true}, + GroupedCastConfig{{128, 256, 128}, 256, true}, + GroupedCastConfig{{384, 128, 256}, 128, true})); + +class NVFP4GraphSafeGroupedRowScaledCastReplaySuite + : public ::testing::TestWithParam {}; + +TEST_P(NVFP4GraphSafeGroupedRowScaledCastReplaySuite, ReplayWithChangedRouting) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + performGraphSafeReplayTest(GetParam().K, GetParam().transpose); +} + +INSTANTIATE_TEST_SUITE_P(NVFP4GraphSafeGroupedRowScaledCastReplay, + NVFP4GraphSafeGroupedRowScaledCastReplaySuite, + ::testing::Values(GroupedCastConfig{{}, 256, false}, + GroupedCastConfig{{}, 256, true})); + +class NVFP4GraphSafeGroupedAmaxTestSuite : public ::testing::TestWithParam {}; + +TEST_P(NVFP4GraphSafeGroupedAmaxTestSuite, MatchesHostSplit) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + performGraphSafeAmaxEqualityTest(GetParam()); +} + +INSTANTIATE_TEST_SUITE_P( + NVFP4GraphSafeGroupedAmax, NVFP4GraphSafeGroupedAmaxTestSuite, + ::testing::Values(GroupedCastConfig{{128}, 128, false}, + GroupedCastConfig{{128, 256, 128}, 256, false}, + GroupedCastConfig{{384, 128, 256}, 128, false}, + GroupedCastConfig{{128}, 128, true}, + GroupedCastConfig{{128, 256, 128}, 256, true}, + GroupedCastConfig{{384, 128, 256}, 128, true})); + +class NVFP4GraphSafeGroupedFullTestSuite : public ::testing::TestWithParam {}; + +TEST_P(NVFP4GraphSafeGroupedFullTestSuite, AmaxThenCastMatchesHostSplit) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + performGraphSafeFullEqualityTest(GetParam()); +} + +INSTANTIATE_TEST_SUITE_P( + NVFP4GraphSafeGroupedFull, NVFP4GraphSafeGroupedFullTestSuite, + ::testing::Values(GroupedCastConfig{{128}, 128, false}, + GroupedCastConfig{{128, 256, 128}, 256, false}, + GroupedCastConfig{{384, 128, 256}, 128, false}, + GroupedCastConfig{{128}, 128, true}, + GroupedCastConfig{{128, 256, 128}, 256, true}, + GroupedCastConfig{{384, 128, 256}, 128, true})); + +class NVFP4GraphSafeGroupedFullReplaySuite : public ::testing::TestWithParam {}; + +TEST_P(NVFP4GraphSafeGroupedFullReplaySuite, ReplayWithChangedRouting) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + performGraphSafeFullReplayTest(GetParam().K, GetParam().transpose); +} + +INSTANTIATE_TEST_SUITE_P(NVFP4GraphSafeGroupedFullReplay, NVFP4GraphSafeGroupedFullReplaySuite, + ::testing::Values(GroupedCastConfig{{}, 256, false}, + GroupedCastConfig{{}, 256, true})); diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 1b5b95d7314..cd42847afb9 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -108,11 +108,15 @@ struct TypeInfo { struct Helper { constexpr static DType getType() { constexpr int i = static_cast(current); - if (std::is_same::type>::value) { - return current; - } else { - return Helper(i + 1)>::getType(); + // The DType enum carries values (e.g. kFloat8UE5M3) whose type is only + // added to ``types`` on CUDA >= 13.4, so guard the tuple access to keep + // the enum walk from indexing past the tuple on older toolkits. + if constexpr (i < static_cast(std::tuple_size::value)) { + if (std::is_same::type>::value) { + return current; + } } + return Helper(i + 1)>::getType(); } }; diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py new file mode 100644 index 00000000000..3d683132c39 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py @@ -0,0 +1,154 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Graph-safe grouped row-scaled NVFP4 cast (amax + cast) from device-side routing. + +Exercises ``transformer_engine.pytorch.ep.nvfp4_group_row_scaled_cast`` against the +host-split row-scaled cast (``tex.split_quantize`` with row-scaled NVFP4 quantizers). +""" + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.pytorch import NVFP4Quantizer +from transformer_engine.pytorch.ep import nvfp4_group_row_scaled_cast + +recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) + + +def _row_scaled_quantizer() -> NVFP4Quantizer: + return NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=False, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=True, + ) + + +def _host_split_reference(x_valid: torch.Tensor, split_sections: list[int]): + quantizers = [_row_scaled_quantizer() for _ in split_sections] + return tex.split_quantize(x_valid, split_sections, quantizers) + + +def _gather_reference(reference, split_sections: list[int], N: int): + """Concatenate the per-expert host-split outputs into contiguous [sum_M, *] buffers.""" + live = [(r, m) for r, m in zip(reference, split_sections) if m > 0] + data = torch.cat([r._rowwise_data.view(torch.uint8).reshape(m, N // 2) for r, m in live], dim=0) + scale = torch.cat([r._rowwise_scale_inv.reshape(m, -1) for r, m in live], dim=0) + amax = torch.cat([r._amax_rowwise.reshape(m) for r, m in live], dim=0) + return data, scale, amax + + +def _assert_matches_reference(grouped, reference, split_sections: list[int], N: int) -> None: + """Compare the grouped output's raw device buffers (valid region) to the reference. + + Reads raw buffers rather than ``split_into_quantized_tensors`` because the wrapper's + per-expert shapes are frozen at capture time while the device buffers are recomputed. + """ + sum_m = sum(split_sections) + ref_data, ref_scale, ref_amax = _gather_reference(reference, split_sections, N) + + capacity = int(grouped.logical_shape[0]) + scale_stride = grouped.scale_inv.numel() // capacity + g_data = grouped.rowwise_data.view(torch.uint8).reshape(capacity, N // 2)[:sum_m] + g_scale = grouped.scale_inv.view(torch.uint8).reshape(capacity, scale_stride)[:sum_m] + g_amax = grouped.amax.reshape(-1)[:sum_m] + + torch.testing.assert_close(g_data, ref_data, atol=0, rtol=0) + torch.testing.assert_close(g_scale[:, : ref_scale.shape[1]], ref_scale, atol=0, rtol=0) + torch.testing.assert_close(g_amax, ref_amax, atol=0, rtol=0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "split_sections, N", + [ + ([128], 128), + ([128, 256, 128], 256), + ([384, 128, 256], 128), + ([256, 128], 512), + ], +) +def test_matches_host_split(split_sections: list[int], N: int) -> None: + """Graph-safe amax + cast is byte-identical to the host-split row-scaled cast.""" + sum_m = sum(split_sections) + torch.manual_seed(0) + x = torch.randn((sum_m, N), dtype=torch.bfloat16, device="cuda") + tokens_per_expert = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + grouped = nvfp4_group_row_scaled_cast(x, tokens_per_expert) + reference = _host_split_reference(x, split_sections) + + _assert_matches_reference(grouped, reference, split_sections, N) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "split_sections, N, capacity", + [ + ([128, 128], 256, 512), + ([256, 128, 128], 128, 768), + ], +) +def test_paged_stashing_capacity_noop(split_sections: list[int], N: int, capacity: int) -> None: + """Capacity rows beyond the live token sum are a device-side no-op.""" + sum_m = sum(split_sections) + assert capacity > sum_m + torch.manual_seed(1) + x = torch.randn((capacity, N), dtype=torch.bfloat16, device="cuda") + x_valid = x[:sum_m, :].clone() + tokens_per_expert = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + grouped = nvfp4_group_row_scaled_cast(x, tokens_per_expert) + reference = _host_split_reference(x_valid, split_sections) + + _assert_matches_reference(grouped, reference, split_sections, N) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_capture_once_replay_changed_routing() -> None: + """Capture the amax + cast once, then replay it with redistributed routing.""" + N = 256 + num_experts = 3 + capacity = 768 + splits_a = [256, 128, 128] + splits_b = [128, 256, 128] + + tokens_per_expert = torch.empty(num_experts, dtype=torch.int64, device="cuda") + x = torch.empty((capacity, N), dtype=torch.bfloat16, device="cuda") + + torch.manual_seed(2) + x_a = torch.randn((sum(splits_a), N), dtype=torch.bfloat16, device="cuda") + x_b = torch.randn((sum(splits_b), N), dtype=torch.bfloat16, device="cuda") + + def program_routing(splits: list[int], x_src: torch.Tensor) -> None: + tokens_per_expert.copy_(torch.tensor(splits, dtype=torch.int64, device="cuda")) + x.zero_() + x[: x_src.shape[0], :].copy_(x_src) + + # Warm up outside capture so one-time setup does not land in the captured graph. + program_routing(splits_a, x_a) + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + _ = nvfp4_group_row_scaled_cast(x, tokens_per_expert) + torch.cuda.current_stream().wait_stream(warmup_stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + grouped = nvfp4_group_row_scaled_cast(x, tokens_per_expert) + + program_routing(splits_a, x_a) + graph.replay() + torch.cuda.synchronize() + _assert_matches_reference(grouped, _host_split_reference(x_a, splits_a), splits_a, N) + + program_routing(splits_b, x_b) + graph.replay() + torch.cuda.synchronize() + _assert_matches_reference(grouped, _host_split_reference(x_b, splits_b), splits_b, N) diff --git a/transformer_engine/common/cast/cast_grouped.cu b/transformer_engine/common/cast/cast_grouped.cu index 49ab39f7e79..c34fc1acb1b 100644 --- a/transformer_engine/common/cast/cast_grouped.cu +++ b/transformer_engine/common/cast/cast_grouped.cu @@ -114,3 +114,96 @@ void nvte_group_nvfp4_compute_amax(const NVTETensor input, NVTETensor *outputs, dispatch::nvfp4::group_compute_fused_amax(*input_tensor, /*noop=*/nullptr, output_tensors, split_sections, num_tensors, stream); } + +void nvte_group_nvfp4_row_scaled_cast_with_amax(const NVTETensor input, NVTETensor *outputs, + const size_t *split_sections, size_t num_tensors, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_nvfp4_row_scaled_cast_with_amax); + using namespace transformer_engine; + + const Tensor *input_tensor = convertNVTETensorCheck(input); + std::vector output_tensors; + output_tensors.reserve(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + output_tensors.push_back(convertNVTETensorCheck(outputs[i])); + } + + QuantizationConfig quant_config_cpp; + if (quant_config != nullptr) { + quant_config_cpp = *reinterpret_cast(quant_config); + } + + Tensor dummy_tensor; + Tensor *noop_tensor = &dummy_tensor; + if (quant_config_cpp.noop_tensor != nullptr) { + noop_tensor = convertNVTETensorCheck(quant_config_cpp.noop_tensor); + } + + // The scale dtype is shared across experts; pick it from the first non-empty output. + DType scale_dtype = DType::kFloat8E4M3; + for (size_t i = 0; i < num_tensors; ++i) { + if (split_sections[i] == 0) continue; + const Tensor *o = output_tensors[i]; + if (o->scale_inv.has_data()) { + scale_dtype = o->scale_inv.dtype; + } else if (o->columnwise_scale_inv.has_data()) { + scale_dtype = o->columnwise_scale_inv.dtype; + } + break; + } + + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + dispatch::nvfp4::group_row_scaled_cast(*input_tensor, noop_tensor, output_tensors, + split_sections, num_tensors, + &quant_config_cpp, stream);); +} + +void nvte_group_nvfp4_compute_amax_graph_safe(const NVTEGroupedTensor input, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_nvfp4_compute_amax_graph_safe); + using namespace transformer_engine; + + GroupedTensor *input_tensor = convertNVTEGroupedTensorCheck(input); + GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output); + + if (input_tensor->num_tensors == 0) { + return; + } + + dispatch::nvfp4::group_compute_fused_amax_graph_safe(input_tensor, /*noop=*/nullptr, + output_tensor, stream); +} + +void nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(const NVTEGroupedTensor input, + NVTEGroupedTensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe); + using namespace transformer_engine; + + GroupedTensor *input_tensor = convertNVTEGroupedTensorCheck(input); + GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output); + + if (input_tensor->num_tensors == 0) { + return; + } + + QuantizationConfig quant_config_cpp; + + Tensor dummy_noop; + Tensor *noop_tensor = &dummy_noop; + + // The scale dtype is shared across experts; pick it from the output buffers. + DType scale_dtype = DType::kFloat8E4M3; + if (output_tensor->scale_inv.has_data()) { + scale_dtype = output_tensor->scale_inv.dtype; + } else if (output_tensor->columnwise_scale_inv.has_data()) { + scale_dtype = output_tensor->columnwise_scale_inv.dtype; + } + + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + dispatch::nvfp4::group_row_scaled_cast_graph_safe( + input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream);); +} 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 f506b764947..1850dd49a0b 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -23,7 +23,9 @@ #include "../../util/math.h" #include "../../util/ptx_arch_spec.cuh" #include "../../utils.cuh" +#include "../core/grouped_layout.cuh" #include "core_nvfp4.cuh" +#include "specialized/quantize_transpose_nvfp4_tuned_1D.cuh" namespace transformer_engine { namespace dispatch { @@ -99,16 +101,13 @@ static __global__ void group_fused_amax_zero_kernel(GroupedFusedAmaxArgs args, c } } +// Shared amax body; host-split and graph-safe kernels differ only in the prologue. template -__global__ void __launch_bounds__(FA_THREADS_NUM) - group_compute_fused_amax_kernel(const __grid_constant__ CUtensorMap tensor_map_input, - GroupedFusedAmaxArgs args, const float *noop, const size_t rows, - const size_t cols) { +__device__ __forceinline__ void group_compute_fused_amax_body(const CUtensorMap &tensor_map_input, + float *const row_amax_out, + float *const col_amax_out, + const int split_start) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - if (noop != nullptr && noop[0] == 1.0f) { - return; - } - using IType = __nv_bfloat16; using IType2 = ptx::FPx2; using IType3D = IType[FA_BUFFS_NUM][FA_TILE_DIM_Y][FA_TILE_DIM_X]; @@ -117,15 +116,6 @@ __global__ void __launch_bounds__(FA_THREADS_NUM) const int tid = threadIdx.x; const int block_offset_Y = blockIdx.y * FA_CHUNK_DIM_Y; const int block_offset_X = blockIdx.x * FA_CHUNK_DIM_X; - - // The 128-row chunk is inside a single expert. Select its per-expert outputs - // and remap the global row to a tensor-local row. - const int tensor_id = GetTensorId(args, block_offset_Y); - const int split_start = args.split_sections_range[tensor_id]; - float *const row_amax_out = - DO_ROW ? reinterpret_cast(args.row_amax_list[tensor_id]) : nullptr; - float *const col_amax_out = - DO_COL ? reinterpret_cast(args.col_amax_list[tensor_id]) : nullptr; const int local_row = block_offset_Y - split_start + tid; const int global_col = block_offset_X + tid; @@ -262,6 +252,84 @@ __global__ void __launch_bounds__(FA_THREADS_NUM) #endif // __CUDA_ARCH__ >= 1000 } +template +__global__ void __launch_bounds__(FA_THREADS_NUM) + group_compute_fused_amax_kernel(const __grid_constant__ CUtensorMap tensor_map_input, + GroupedFusedAmaxArgs args, const float *noop, const size_t rows, + const size_t cols) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + const int block_offset_Y = blockIdx.y * FA_CHUNK_DIM_Y; + // The 128-row chunk is inside a single expert; select its per-expert outputs. + const int tensor_id = GetTensorId(args, block_offset_Y); + const int split_start = args.split_sections_range[tensor_id]; + float *const row_amax_out = + DO_ROW ? reinterpret_cast(args.row_amax_list[tensor_id]) : nullptr; + float *const col_amax_out = + DO_COL ? reinterpret_cast(args.col_amax_list[tensor_id]) : nullptr; + + group_compute_fused_amax_body(tensor_map_input, row_amax_out, col_amax_out, + split_start); +#else + NVTE_DEVICE_ERROR("Grouped fused amax kernel requires SM 10.0+ (Blackwell)."); +#endif // __CUDA_ARCH__ >= 1000 +} + +// Zero the capacity-sized grouped amax buffers before atomicMax. +static __global__ void group_fused_amax_zero_graph_safe_kernel(float *row_amax_all, + float *col_amax_all, + const size_t row_count, + const size_t col_count, + const float *noop) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + const size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t stride = static_cast(gridDim.x) * blockDim.x; + if (row_amax_all != nullptr) { + for (size_t i = idx; i < row_count; i += stride) row_amax_all[i] = 0.0f; + } + if (col_amax_all != nullptr) { + for (size_t i = idx; i < col_count; i += stride) col_amax_all[i] = 0.0f; + } +} + +// Graph-safe: per-expert outputs and row start derived on device from offsets; +// CTAs past the live row sum noop. +template +__global__ void __launch_bounds__(FA_THREADS_NUM) + group_compute_fused_amax_graph_safe_kernel(const __grid_constant__ CUtensorMap tensor_map_input, + float *const row_amax_all, float *const col_amax_all, + const int64_t *const __restrict__ offsets_ptr, + const size_t num_tensors, const float *noop, + const size_t cols) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + const size_t block_offset_Y = static_cast(blockIdx.y) * FA_CHUNK_DIM_Y; + const size_t global_offset = block_offset_Y * cols; + if (global_offset >= static_cast(offsets_ptr[num_tensors])) { + return; + } + + const size_t tensor_id = + dispatch::common::find_tensor_from_offsets(offsets_ptr, num_tensors, global_offset); + const int split_start = static_cast(static_cast(offsets_ptr[tensor_id]) / cols); + float *const row_amax_out = DO_ROW ? row_amax_all + split_start : nullptr; + float *const col_amax_out = DO_COL ? col_amax_all + tensor_id * cols : nullptr; + + group_compute_fused_amax_body(tensor_map_input, row_amax_out, col_amax_out, + split_start); +#else + NVTE_DEVICE_ERROR("Grouped fused amax kernel requires SM 10.0+ (Blackwell)."); +#endif // __CUDA_ARCH__ >= 1000 +} + #endif // FP4_TYPE_SUPPORTED } // namespace group_row_scaled_amax_kernel @@ -349,6 +417,81 @@ inline void group_compute_fused_amax(const Tensor &input, const Tensor *noop, #endif // FP4_TYPE_SUPPORTED } +// Graph-safe counterpart of group_compute_fused_amax: routing comes from device +// output->tensor_offsets, amax goes to the contiguous rowwise [sum_M] / columnwise +// [num_tensors * K] buffers. BF16-only, VARYING_FIRST_DIM. +inline void group_compute_fused_amax_graph_safe(const GroupedTensor *input, const Tensor *noop, + GroupedTensor *output, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace group_row_scaled_amax_kernel; + + NVTE_CHECK(input->dtype() == DType::kBFloat16, "Grouped fused amax requires BF16 input."); + NVTE_CHECK(output->all_same_last_dim() || output->all_same_shape(), + "Graph-safe grouped fused amax only supports a common last dimension " + "(VARYING_FIRST_DIM)."); + + const size_t num_tensors = input->num_tensors; + const size_t capacity_rows = input->logical_shape.data[0]; + const size_t cols = input->logical_shape.data[1]; + NVTE_CHECK(cols % FA_CHUNK_DIM_X == 0, "Grouped fused amax requires cols to be a multiple of ", + FA_CHUNK_DIM_X, " (got ", cols, ")."); + NVTE_CHECK(capacity_rows % FA_CHUNK_DIM_Y == 0, + "Graph-safe grouped fused amax requires the capacity row dimension to be a multiple " + "of ", + FA_CHUNK_DIM_Y, " (got ", capacity_rows, ")."); + + float *const row_amax_all = reinterpret_cast(output->amax.dptr); + float *const col_amax_all = reinterpret_cast(output->columnwise_amax.dptr); + const bool do_row = row_amax_all != nullptr; + const bool do_col = col_amax_all != nullptr; + if (!do_row && !do_col) return; + + const int64_t *offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); + NVTE_CHECK(offsets_ptr != nullptr, + "Graph-safe grouped fused amax requires device tensor_offsets."); + + checkCuDriverContext(stream); + const float *noop_ptr = (noop != nullptr && noop->data.dptr != nullptr) + ? reinterpret_cast(noop->data.dptr) + : nullptr; + + // Zero the capacity-sized amax buffers before atomicMax. + { + const size_t row_count = do_row ? capacity_rows : 0; + const size_t col_count = do_col ? num_tensors * cols : 0; + const size_t total = std::max(row_count, col_count); + const unsigned zero_blocks = static_cast(DIVUP(total, static_cast(256))) + 1u; + group_fused_amax_zero_graph_safe_kernel<<>>( + row_amax_all, col_amax_all, row_count, col_count, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + + alignas(64) CUtensorMap tensor_map_input{}; + create_2D_tensor_map(tensor_map_input, input->data, capacity_rows, cols, FA_TILE_DIM_Y, + FA_TILE_DIM_X, cols, 0, sizeof(__nv_bfloat16) * 8); + + constexpr size_t buff_elems_total = FA_BUFFS_NUM * FA_BUFF_IN_SIZE; + constexpr size_t buff_size_aligned_in = + DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(__nv_bfloat16), TMA_SHMEM_ALIGNMENT); + constexpr size_t dshmem_size = buff_size_aligned_in + TMA_SHMEM_ALIGNMENT; + + const dim3 grid(static_cast(cols / FA_CHUNK_DIM_X), + static_cast(capacity_rows / FA_CHUNK_DIM_Y), 1); + const dim3 block(FA_THREADS_NUM, 1, 1); + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + do_row, DO_ROW, TRANSFORMER_ENGINE_SWITCH_CONDITION(do_col, DO_COL, { + auto kernel = group_compute_fused_amax_graph_safe_kernel; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); + kernel<<>>(tensor_map_input, row_amax_all, col_amax_all, + offsets_ptr, num_tensors, noop_ptr, cols); + })); + NVTE_CHECK_CUDA(cudaGetLastError()); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + namespace group_quantize_transpose_kernel { using namespace core; @@ -1227,6 +1370,589 @@ void group_quantize_transpose(const Tensor &input, const Tensor *noop, #endif // FP4_TYPE_SUPPORTED } +// Single-launch grouped (MoE) row-scaled NVFP4 cast. Reuses the single-tensor +// tuned_1D encode helpers verbatim (so the output is byte-identical to the +// per-expert loop) and replaces tuned_1D's TMA output stores with a per-expert +// global scatter, since each expert's transpose output has its own M_i stride. +namespace group_row_scaled_cast_kernel { + +using namespace quantize_transpose_tuned_kernel; +using namespace ptx; + +#if FP4_TYPE_SUPPORTED + +constexpr int kMaxCastTensorsPerKernel = 64; + +// Per-expert output / amax pointers. The packed (sum_M, K) input is covered by a +// single TMA map; data_stride_row = cols/2 and data_stride_col = M_i/2 are +// derived in-kernel from split_sections_range. +struct GroupedRowScaledCastArgs { + void *q_row_list[kMaxCastTensorsPerKernel]; // uint8* [M_i, cols/2] + void *s_dec_row_list[kMaxCastTensorsPerKernel]; // e4m3* [M_i, scale_stride_row] + void *q_col_list[kMaxCastTensorsPerKernel]; // uint8* [cols, M_i/2] + void *s_dec_col_list[kMaxCastTensorsPerKernel]; // e4m3* [cols, scale_stride_col] + void *row_amax_list[kMaxCastTensorsPerKernel]; // float* [M_i] + void *col_amax_list[kMaxCastTensorsPerKernel]; // float* [cols] + int scale_stride_col[kMaxCastTensorsPerKernel]; // per-expert colwise SF row stride + int split_sections_range[kMaxCastTensorsPerKernel + 1]; // prefix sum of M_i + int num_tensors; +}; + +__device__ __forceinline__ int GetTensorId(const GroupedRowScaledCastArgs &args, int global_row) { + int tensor_id = 0; + while (args.split_sections_range[tensor_id + 1] <= global_row) { + ++tensor_id; + } + return tensor_id; +} + +// Per-expert encode + scatter body shared by the host-split and graph-safe +// grouped row-scaled cast kernels. The two kernels differ only in how they +// resolve the per-expert bases (host-baked args vs device metadata); everything +// from the SMEM layout onward is identical. +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, + const float *const col_amax_base, const size_t scale_stride_col, const int M_t, + const int local_row_base, const int ctaid_X, const int ctaid_Y, const size_t cols, + const size_t scale_stride_row) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr bool USE_STOCHASTIC_ROUNDING = false; // row-scaled rejects SR + constexpr bool ROW_SCALED_NVFP4 = true; + + // SR off: the RNG stays inert (random_uint4 is never consumed); it exists only + // to satisfy the shared encode-helper signatures. + RNG_t rng; + rng.init(0, 0, 0); + uint4 random_uint4 = uint4{0, 0, 0, 0}; + int rnd_idx = 0; + + const bool leading_thread = (threadIdx.x == 0); + + const int block_offset_Y = ctaid_Y * TunableConfig::CHUNK_DIM_Y; + const int block_offset_X = ctaid_X * TunableConfig::CHUNK_DIM_X; + + const size_t data_stride_row = cols / 2; + const size_t data_stride_col = static_cast(M_t) / 2; + + // SMEM layout: identical to the single-tensor tuned_1D kernel. + constexpr int buff_elems = BUFF_DIM_Y * BUFF_IN_DIM_X; + constexpr int buff_elems_total_in = BUFFS_NUM_IN * buff_elems; + constexpr int buff_size_aligned_in = + DIVUP_TO_MULTIPLE(buff_elems_total_in * sizeof(IType), TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT * BUFF_OUT_SIZE, TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out_t = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT_TR * BUFF_OUT_TR_SIZE, TMA_SHMEM_ALIGNMENT); + constexpr int in_mem = buff_size_aligned_in; + constexpr int out_mem_rowwise_data = buff_size_aligned_out; + constexpr int out_mem_colwise_data = RETURN_TRANSPOSE ? buff_size_aligned_out_t : 0; + constexpr int out_mem_rowwise_scales = DIVUP_TO_MULTIPLE( + TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); + + extern __shared__ char dynamic_shmem[]; + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); + + IType *sIn_ptr = reinterpret_cast(dshmem); + fp4e2m1x2 *sOut_ptr = reinterpret_cast(dshmem + in_mem); + fp4e2m1x2 *sOut_tr_ptr = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); + ScaleType *sSFrowwise_ptr = + reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); + ScaleType *sSFcolwise_ptr = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); + + auto &sIn = *reinterpret_cast(sIn_ptr); + + constexpr int shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; + + __shared__ uint64_t IN_buff_readable_mbar[BUFFS_NUM]; + + if (leading_thread) { +#pragma unroll + for (int buff = 0; buff < BUFFS_NUM; ++buff) { + ptx::mbarrier_init(&IN_buff_readable_mbar[buff], 1); + } + ptx::fence_proxy_async_shared_cta(); + } + __syncthreads(); + + // Prefetch the first input tile(s). +#pragma unroll + for (int stage = 0; stage < TunableConfig::PREFETCH_STAGES; ++stage) { + const int buff_in_p = stage; + const int stage_Y = stage / STAGES_X; + const int stage_X = stage % STAGES_X; + const int global_offset_Y = block_offset_Y + stage_Y * TILE_DIM_Y; + const int global_offset_X = block_offset_X + stage_X * TILE_DIM_X; + if (leading_thread) { + ptx::mbarrier_arrive_expect_tx(&IN_buff_readable_mbar[buff_in_p], shmem_buff_size); + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&sIn[buff_in_p]), + reinterpret_cast(&tensor_map_input), global_offset_X, global_offset_Y, + &IN_buff_readable_mbar[buff_in_p]); + } + } + + int buff_in = 0; + int buff_out = 0; + int buff_out_tr = 0; + int IN_buff_readable_parity[BUFFS_NUM] = {0, 0}; + +#pragma unroll + for (int stage = 0; stage < STAGES; ++stage) { + const int stage_Y = stage / STAGES_X; + const int stage_X = stage % STAGES_X; + + // Prefetch the next input tile. + if (stage < STAGES - TunableConfig::PREFETCH_STAGES) { + const int next_prefetch_buff = (buff_in + TunableConfig::PREFETCH_STAGES) % BUFFS_NUM; + const int next_prefetch_stage = stage + TunableConfig::PREFETCH_STAGES; + const int next_stage_Y = next_prefetch_stage / STAGES_X; + const int next_stage_X = next_prefetch_stage % STAGES_X; + const int next_global_offset_Y = block_offset_Y + next_stage_Y * TILE_DIM_Y; + const int next_global_offset_X = block_offset_X + next_stage_X * TILE_DIM_X; + if (leading_thread) { + ptx::mbarrier_arrive_expect_tx(&IN_buff_readable_mbar[next_prefetch_buff], shmem_buff_size); + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&sIn[next_prefetch_buff]), + reinterpret_cast(&tensor_map_input), next_global_offset_X, + next_global_offset_Y, &IN_buff_readable_mbar[next_prefetch_buff]); + } + ptx::fence_proxy_async_shared_cta(); + } + + // Wait for the current input tile to land. + ptx::mbarrier_wait_parity_acquire_cta_shared_cta(&IN_buff_readable_mbar[buff_in], + IN_buff_readable_parity[buff_in]); + IN_buff_readable_parity[buff_in] ^= 1; + + // Per-row encode (byte-identical to the single-tensor path): pass the + // per-expert amax base and a local row offset so amax[row_idx] indexes the + // expert-local row. + rowwise_scaling( + sIn_ptr, sOut_ptr, sSFrowwise_ptr, /*S_enc_rowwise=*/1.0f, stage_Y, stage_X, buff_in, + buff_out, row_amax_base, static_cast(local_row_base), static_cast(M_t), rng, + random_uint4, rnd_idx); + + if constexpr (RETURN_TRANSPOSE) { + colwise_scaling( + sIn_ptr, sOut_tr_ptr, sSFcolwise_ptr, /*S_enc_colwise=*/1.0f, stage_Y, stage_X, buff_in, + buff_out_tr, col_amax_base, static_cast(block_offset_X), cols, rng, random_uint4, + rnd_idx); + } + + // Make the encode-helper SMEM writes visible before the scatter. + __syncthreads(); + + // Scatter the rowwise FP4 tile to the per-expert buffer: a verbatim copy of + // the 64-row x 32-byte tile (2 threads x 16 B per row). + { + auto &sOut = *reinterpret_cast(sOut_ptr); + const int row_in_tile = static_cast(threadIdx.x) >> 1; + const int half = static_cast(threadIdx.x) & 1; + const int local_row = local_row_base + stage_Y * TILE_DIM_Y + row_in_tile; + const size_t byte_off_X = + static_cast(block_offset_X / 2) + stage_X * (TILE_DIM_X / 2) + half * 16; + const uint4 *src = reinterpret_cast(&sOut[buff_out][row_in_tile][half * 16]); + uint4 *dst = reinterpret_cast( + q_row_base + static_cast(local_row) * data_stride_row + byte_off_X); + *dst = *src; + } + if constexpr (RETURN_TRANSPOSE) { + auto &sOut_tr = *reinterpret_cast(sOut_tr_ptr); + const int col_in_tile = static_cast(threadIdx.x) >> 1; + const int half = static_cast(threadIdx.x) & 1; + const int global_col = block_offset_X + stage_X * TILE_DIM_X + col_in_tile; + const size_t byte_off_M = + static_cast(local_row_base / 2) + stage_Y * (TILE_DIM_Y / 2) + half * 16; + const uint4 *src = + reinterpret_cast(&sOut_tr[buff_out_tr][col_in_tile][half * 16]); + uint4 *dst = reinterpret_cast( + q_col_base + static_cast(global_col) * data_stride_col + byte_off_M); + *dst = *src; + } + + // Sync so the scatter completes before the next stage overwrites the buffer. + __syncthreads(); + + buff_in = (buff_in + 1) % BUFFS_NUM_IN; + buff_out = (buff_out + 1) % BUFFS_NUM_OUT; + buff_out_tr = (buff_out_tr + 1) % BUFFS_NUM_OUT_TR; + } + + // SF epilogue: store the per-chunk SMEM scale tiles to the per-expert buffers. + { + auto &sSFrowwise = *reinterpret_cast *>(sSFrowwise_ptr); + using ScalesVec = Vec; + const size_t scales_block_offset_X_rowwise = static_cast(ctaid_X) * SCALES_PER_CHUNK_X; + for (int row = static_cast(threadIdx.x); row < TunableConfig::CHUNK_DIM_Y; + 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 (RETURN_TRANSPOSE) { + auto &sSFcolwise = *reinterpret_cast *>(sSFcolwise_ptr); + using ScalesVec = Vec; + const size_t local_block_offset_M = static_cast(local_row_base) / SCALE_DIM; + for (int row_tr = static_cast(threadIdx.x); row_tr < TunableConfig::CHUNK_DIM_X; + 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 (leading_thread) { +#pragma unroll + for (int buff = 0; buff < BUFFS_NUM; ++buff) { + ptx::mbarrier_invalid(&IN_buff_readable_mbar[buff]); + } + } +#else + NVTE_DEVICE_ERROR("Grouped row-scaled NVFP4 cast kernel requires SM 10.0+ (Blackwell)."); +#endif // __CUDA_ARCH__ >= 1000 +} + +// Host-split kernel: per-expert bases come from the host-baked GroupedRowScaledCastArgs. +template +__global__ void __launch_bounds__(THREADS_NUM) + group_row_scaled_cast_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, + const __grid_constant__ GroupedRowScaledCastArgs args, + const float *noop, const size_t rows, const size_t cols, + const size_t scale_stride_row) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + const int ctaid_X = blockIdx.x; + const int ctaid_Y = blockIdx.y; + const int block_offset_Y = ctaid_Y * TunableConfig::CHUNK_DIM_Y; + + // A 128-row chunk lies fully inside one expert (every split is 128-aligned). + const int tensor_id = GetTensorId(args, block_offset_Y); + const int split_start = args.split_sections_range[tensor_id]; + const int M_t = args.split_sections_range[tensor_id + 1] - split_start; + const int local_row_base = block_offset_Y - split_start; + + uint8_t *const q_row_base = reinterpret_cast(args.q_row_list[tensor_id]); + ScaleType *const s_dec_row_base = reinterpret_cast(args.s_dec_row_list[tensor_id]); + const float *const row_amax_base = reinterpret_cast(args.row_amax_list[tensor_id]); + + uint8_t *const q_col_base = + RETURN_TRANSPOSE ? reinterpret_cast(args.q_col_list[tensor_id]) : nullptr; + ScaleType *const s_dec_col_base = + RETURN_TRANSPOSE ? reinterpret_cast(args.s_dec_col_list[tensor_id]) : nullptr; + const float *const col_amax_base = + RETURN_TRANSPOSE ? reinterpret_cast(args.col_amax_list[tensor_id]) : nullptr; + const size_t scale_stride_col = + RETURN_TRANSPOSE ? static_cast(args.scale_stride_col[tensor_id]) : 0; + + 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); +#else + NVTE_DEVICE_ERROR("Grouped row-scaled NVFP4 cast kernel requires SM 10.0+ (Blackwell)."); +#endif // __CUDA_ARCH__ >= 1000 +} + +// Graph-safe kernel: per-expert bases derived on device from offsets / first_dims; +// grid covers the input capacity, CTAs past the live row sum noop. +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, + ScaleType *const s_dec_col_all, const float *const col_amax_all, + const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, + const size_t num_tensors, const float *noop, const size_t cols, const size_t scale_stride_row) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + const int ctaid_X = blockIdx.x; + const int ctaid_Y = blockIdx.y; + const size_t block_offset_Y = static_cast(ctaid_Y) * TunableConfig::CHUNK_DIM_Y; + + const size_t global_offset = block_offset_Y * cols; + if (global_offset >= static_cast(offsets_ptr[num_tensors])) { + return; + } + + const size_t tensor_id = + dispatch::common::find_tensor_from_offsets(offsets_ptr, num_tensors, global_offset); + const size_t row_start = static_cast(offsets_ptr[tensor_id]) / cols; + const int M_t = static_cast(first_dims_ptr[tensor_id]); + const int local_row_base = static_cast(block_offset_Y - row_start); + + // Rowwise buffers are contiguous by global row; transpose buffers are packed + // per expert with column stride M_t. + const size_t scale_stride_col = RETURN_TRANSPOSE ? static_cast(M_t) / SCALE_DIM : 0; + + uint8_t *const q_row_base = q_row_all + row_start * (cols / 2); + ScaleType *const s_dec_row_base = s_dec_row_all + row_start * scale_stride_row; + const float *const row_amax_base = row_amax_all + row_start; + + uint8_t *const q_col_base = RETURN_TRANSPOSE ? q_col_all + (cols / 2) * row_start : nullptr; + ScaleType *const s_dec_col_base = + 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( + 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); +#else + NVTE_DEVICE_ERROR("Grouped row-scaled NVFP4 cast kernel requires SM 10.0+ (Blackwell)."); +#endif // __CUDA_ARCH__ >= 1000 +} + +#endif // FP4_TYPE_SUPPORTED + +} // namespace group_row_scaled_cast_kernel + +// Host launcher for the grouped row-scaled NVFP4 cast. BF16-only; consumes the +// per-expert amax computed by nvte_group_nvfp4_compute_amax. +template +inline void group_row_scaled_cast(const Tensor &input, const Tensor *noop, + std::vector &output_list, const size_t *split_sections, + size_t num_tensors, const QuantizationConfig *quant_config, + cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace group_row_scaled_cast_kernel; + using namespace quantize_transpose_tuned_kernel; + + NVTE_CHECK(num_tensors <= static_cast(kMaxCastTensorsPerKernel), + "Number of tensors should be <= ", kMaxCastTensorsPerKernel); + NVTE_CHECK(input.dtype() == DType::kBFloat16, + "Grouped row-scaled NVFP4 cast requires BF16 input."); + + const bool use_fast_math = quant_config ? quant_config->use_fast_math : false; + + checkCuDriverContext(stream); + CheckNoopTensor(*noop, "cast_noop"); + CheckInputTensor(input, "input"); + NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); + + const auto [rows, cols] = input.flat_2d_dims(); + NVTE_CHECK(cols % TunableConfig::CHUNK_DIM_X == 0, + "Grouped row-scaled NVFP4 cast requires cols to be a multiple of ", + TunableConfig::CHUNK_DIM_X, " (got ", cols, ")."); + + // First non-empty output defines the shared layout (transpose + rowwise SF stride). + Tensor *first_output = nullptr; + for (size_t i = 0; i < num_tensors; ++i) { + if (split_sections[i] != 0 && output_list[i]->has_data()) { + first_output = output_list[i]; + break; + } + } + NVTE_CHECK(first_output != nullptr, "No non-empty output tensor found."); + const bool return_transpose = first_output->has_columnwise_data(); + const size_t scale_stride_row = first_output->scale_inv.shape[1]; + + GroupedRowScaledCastArgs args; + args.num_tensors = 0; + args.split_sections_range[0] = 0; + for (size_t i = 0; i < num_tensors; ++i) { + if (split_sections[i] == 0) continue; + NVTE_CHECK(split_sections[i] % TunableConfig::CHUNK_DIM_Y == 0, + "Grouped row-scaled NVFP4 cast requires each split to be a multiple of ", + TunableConfig::CHUNK_DIM_Y, " (split ", i, " = ", split_sections[i], ")."); + Tensor *o = output_list[i]; + NVTE_CHECK(o->amax.dptr != nullptr, "Grouped row-scaled NVFP4 cast requires rowwise amax."); + const int idx = args.num_tensors; + args.q_row_list[idx] = o->data.dptr; + args.s_dec_row_list[idx] = o->scale_inv.dptr; + args.row_amax_list[idx] = o->amax.dptr; + if (return_transpose) { + NVTE_CHECK(o->has_columnwise_data() && o->columnwise_scale_inv.dptr != nullptr && + o->columnwise_amax.dptr != nullptr, + "Grouped row-scaled NVFP4 transpose cast requires columnwise data/scale/amax for " + "every expert."); + args.q_col_list[idx] = o->columnwise_data.dptr; + args.s_dec_col_list[idx] = o->columnwise_scale_inv.dptr; + args.col_amax_list[idx] = o->columnwise_amax.dptr; + args.scale_stride_col[idx] = static_cast(o->columnwise_scale_inv.shape[1]); + } else { + args.q_col_list[idx] = nullptr; + args.s_dec_col_list[idx] = nullptr; + args.col_amax_list[idx] = nullptr; + args.scale_stride_col[idx] = 0; + } + args.split_sections_range[idx + 1] = + args.split_sections_range[idx] + static_cast(split_sections[i]); + args.num_tensors++; + } + NVTE_CHECK(args.split_sections_range[args.num_tensors] == static_cast(rows), + "Sum of split sections must equal the number of input rows."); + + const float *noop_ptr = reinterpret_cast(noop->data.dptr); + + alignas(64) CUtensorMap tensor_map_input{}; + create_2D_tensor_map(tensor_map_input, input.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, cols, 0, + sizeof(IType) * 8); + + constexpr int buff_elems = BUFF_DIM_Y * BUFF_DIM_X; + constexpr int buff_elems_total_in = BUFFS_NUM_IN * buff_elems; + constexpr int buff_size_aligned_in = + DIVUP_TO_MULTIPLE(buff_elems_total_in * sizeof(IType), TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT * BUFF_OUT_SIZE, TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out_t = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT_TR * BUFF_OUT_TR_SIZE, TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_scales = DIVUP_TO_MULTIPLE( + TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_scales_transpose = DIVUP_TO_MULTIPLE( + TunableConfig::CHUNK_DIM_X * SCALES_PER_CHUNK_Y * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); + + const dim3 grid(static_cast(cols / TunableConfig::CHUNK_DIM_X), + static_cast(rows / TunableConfig::CHUNK_DIM_Y), 1); + const dim3 block(THREADS_NUM, 1, 1); + + 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_kernel; + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + kernel<<>>(tensor_map_input, args, noop_ptr, rows, 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); +#endif // FP4_TYPE_SUPPORTED +} + +// Graph-safe counterpart of group_row_scaled_cast: routing comes from device +// output->first_dims / tensor_offsets. Grid covers the input capacity; CTAs past +// the live row sum noop. BF16-only; consumes the output's per-expert amax. +template +inline void group_row_scaled_cast_graph_safe(const GroupedTensor *input, const Tensor *noop, + GroupedTensor *output, + const QuantizationConfig *quant_config, + cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace group_row_scaled_cast_kernel; + using namespace quantize_transpose_tuned_kernel; + + NVTE_CHECK(input->dtype() == DType::kBFloat16, + "Grouped row-scaled NVFP4 cast requires BF16 input."); + NVTE_CHECK(output->all_same_last_dim() || output->all_same_shape(), + "Graph-safe grouped row-scaled NVFP4 cast only supports a common last dimension " + "(VARYING_FIRST_DIM)."); + + const bool use_fast_math = quant_config ? quant_config->use_fast_math : false; + + checkCuDriverContext(stream); + CheckNoopTensor(*noop, "cast_noop"); + + const size_t num_tensors = input->num_tensors; + const size_t capacity_rows = input->logical_shape.data[0]; + const size_t cols = input->logical_shape.data[1]; + NVTE_CHECK(cols % TunableConfig::CHUNK_DIM_X == 0, + "Grouped row-scaled NVFP4 cast requires cols to be a multiple of ", + TunableConfig::CHUNK_DIM_X, " (got ", cols, ")."); + NVTE_CHECK(capacity_rows % TunableConfig::CHUNK_DIM_Y == 0, + "Graph-safe grouped row-scaled NVFP4 cast requires the capacity row dimension to be a " + "multiple of ", + TunableConfig::CHUNK_DIM_Y, " (got ", capacity_rows, ")."); + + const bool return_transpose = output->has_columnwise_data(); + 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. + size_t scale_stride_row = 0; + { + const auto &scale_shape = output->scale_inv.shape; + const size_t scale_ndim = scale_shape.size(); + if (scale_ndim >= 2) { + scale_stride_row = scale_shape[scale_ndim - 1]; + } else { + size_t scale_numel = 1; + for (size_t i = 0; i < scale_ndim; ++i) scale_numel *= scale_shape[i]; + scale_stride_row = scale_numel / capacity_rows; + } + } + + const int64_t *offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); + const int64_t *first_dims_ptr = reinterpret_cast(output->first_dims.dptr); + NVTE_CHECK(offsets_ptr != nullptr && first_dims_ptr != nullptr, + "Graph-safe cast requires device tensor_offsets and first_dims."); + + uint8_t *q_row_all = reinterpret_cast(output->data.dptr); + ScaleType *s_dec_row_all = reinterpret_cast(output->scale_inv.dptr); + const float *row_amax_all = reinterpret_cast(output->amax.dptr); + uint8_t *q_col_all = + return_transpose ? reinterpret_cast(output->columnwise_data.dptr) : nullptr; + ScaleType *s_dec_col_all = + return_transpose ? reinterpret_cast(output->columnwise_scale_inv.dptr) : nullptr; + const float *col_amax_all = + return_transpose ? reinterpret_cast(output->columnwise_amax.dptr) : nullptr; + if (return_transpose) { + NVTE_CHECK(s_dec_col_all != nullptr && col_amax_all != nullptr, + "Graph-safe transpose cast requires columnwise scale and amax buffers."); + } + + const float *noop_ptr = reinterpret_cast(noop->data.dptr); + + alignas(64) CUtensorMap tensor_map_input{}; + create_2D_tensor_map(tensor_map_input, input->data, capacity_rows, cols, BUFF_DIM_Y, BUFF_DIM_X, + cols, 0, sizeof(IType) * 8); + + constexpr int buff_elems = BUFF_DIM_Y * BUFF_DIM_X; + constexpr int buff_elems_total_in = BUFFS_NUM_IN * buff_elems; + constexpr int buff_size_aligned_in = + DIVUP_TO_MULTIPLE(buff_elems_total_in * sizeof(IType), TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT * BUFF_OUT_SIZE, TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out_t = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT_TR * BUFF_OUT_TR_SIZE, TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_scales = DIVUP_TO_MULTIPLE( + TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_scales_transpose = DIVUP_TO_MULTIPLE( + TunableConfig::CHUNK_DIM_X * SCALES_PER_CHUNK_Y * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); + + const dim3 grid(static_cast(cols / TunableConfig::CHUNK_DIM_X), + static_cast(capacity_rows / TunableConfig::CHUNK_DIM_Y), 1); + const dim3 block(THREADS_NUM, 1, 1); + + 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); + })); + NVTE_CHECK_CUDA(cudaGetLastError()); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + } // namespace nvfp4 } // namespace dispatch } // namespace transformer_engine diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index ca3eb13ed24..17b6b7fff37 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -514,6 +514,74 @@ void nvte_group_nvfp4_compute_amax(const NVTETensor input, NVTETensor *outputs, const size_t *split_sections, size_t num_tensors, cudaStream_t stream); +/*! \brief Casts a grouped row-scaled NVFP4 input to FP4 in a single kernel launch. + * + * Single-launch counterpart of the per-expert row-scaled NVFP4 cast. The input + * is one packed (sum_M, K) BF16 buffer; `split_sections` gives the per-expert + * row counts (each a multiple of 128) and K must be a multiple of 128. The + * per-expert rowwise amax [M_i] and (when the columnwise/transpose output is + * allocated) columnwise amax [K] must already be populated, e.g. by + * nvte_group_nvfp4_compute_amax. For expert i the rowwise FP4 data and E4M3 + * scales are written to `outputs[i]`, and the transpose counterpart to + * `outputs[i]`'s columnwise buffers when present. The output is byte-identical + * to quantizing each expert independently. + * + * \param[in] input Packed grouped input tensor (BF16). + * \param[in,out] outputs Per-expert output tensors (amax pre-populated). + * \param[in] split_sections Per-expert row counts (each a multiple of 128). + * \param[in] num_tensors Number of output tensors. + * \param[in] quant_config (Optional) Quantization configurations. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_nvfp4_row_scaled_cast_with_amax(const NVTETensor input, NVTETensor *outputs, + const size_t *split_sections, size_t num_tensors, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream); + +/*! \brief CUDA-graph-safe grouped row-scaled NVFP4 fused amax. + * + * Device-metadata counterpart of nvte_group_nvfp4_compute_amax. Instead of host + * `split_sections`, the per-expert offsets are read from the output grouped + * tensor's device `tensor_offsets`, so a dynamic-routing MoE amax can be + * captured once and replayed with different routing. + * + * The input is one packed (capacity, K) BF16 buffer. The rowwise amax [sum_M] + * is written to the contiguous `output->amax` and the columnwise amax + * [num_tensors * K] to `output->columnwise_amax`, whichever are allocated. Only + * the const-last-dimension layout (VARYING_FIRST_DIM: common K) is supported; + * K must be a multiple of 128 and each per-expert row count a multiple of 128. + * + * \param[in] input Packed grouped input tensor (BF16) with device metadata. + * \param[in,out] output Grouped output tensor receiving the amax vectors. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_nvfp4_compute_amax_graph_safe(const NVTEGroupedTensor input, + NVTEGroupedTensor output, cudaStream_t stream); + +/*! \brief CUDA-graph-safe grouped row-scaled NVFP4 cast. + * + * Device-metadata counterpart of nvte_group_nvfp4_row_scaled_cast_with_amax. + * Instead of host `split_sections`, the per-expert row counts and offsets are + * read from the input/output grouped tensors' device `first_dims` and + * `tensor_offsets`, so a dynamic-routing MoE cast can be captured once and + * replayed with different routing without rebaking kernel parameters. + * + * Both tensors use the single-buffer grouped layout: the input is one packed + * (capacity, K) BF16 buffer, and the output carries contiguous rowwise data / + * E4M3 scales (and, when allocated, columnwise data / E4M3 scales) plus the + * pre-populated rowwise amax [sum_M] and columnwise amax [num_tensors * K]. + * Only the const-last-dimension layout (VARYING_FIRST_DIM: common K, per-expert + * row count) is supported; each per-expert row count must be a multiple of 128 + * and K a multiple of 128. The output is byte-identical to the host-split cast. + * + * \param[in] input Packed grouped input tensor (BF16) with device metadata. + * \param[in,out] output Grouped output tensor (amax pre-populated) with device metadata. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(const NVTEGroupedTensor input, + NVTEGroupedTensor output, + cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 955763f43b1..6ec981bf116 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -420,6 +420,11 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q const at::Tensor &columnwise_amax, std::optional tensor_offsets); +py::object nvfp4_group_row_scaled_cast_graph_safe(const at::Tensor &tensor, py::handle quantizer, + const size_t num_tensors, + const at::Tensor &first_dims, + const at::Tensor &tensor_offsets); + py::object group_dequantize(const py::handle &input, DType otype); py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 06d03cd6e1c..741de9726af 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -635,6 +635,73 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q return py::reinterpret_borrow(grouped_output_py); } +// Graph-safe grouped row-scaled NVFP4 cast: input and routing (first_dims, +// tensor_offsets) live on device, so it can be captured and replayed. +py::object nvfp4_group_row_scaled_cast_graph_safe(const at::Tensor &tensor, py::handle quantizer, + const size_t num_tensors, + const at::Tensor &first_dims, + const at::Tensor &tensor_offsets) { + using namespace transformer_engine::pytorch::detail; + init_extension(); + + NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D"); + NVTE_CHECK(tensor.scalar_type() == at::kBFloat16, + "Graph-safe grouped row-scaled NVFP4 cast requires a BF16 input."); + NVTE_CHECK(first_dims.is_cuda() && tensor_offsets.is_cuda(), + "first_dims and tensor_offsets must be CUDA tensors."); + NVTE_CHECK(first_dims.scalar_type() == at::kLong && tensor_offsets.scalar_type() == at::kLong, + "first_dims and tensor_offsets must be int64."); + NVTE_CHECK(first_dims.numel() == static_cast(num_tensors), + "first_dims must contain one row count per expert."); + NVTE_CHECK(tensor_offsets.numel() == static_cast(num_tensors) + 1, + "tensor_offsets must contain num_tensors + 1 element offsets."); + NVTE_CHECK(IsNVFP4Quantizers(quantizer.ptr()), + "Graph-safe grouped row-scaled NVFP4 cast only supports NVFP4 quantizers."); + + auto quantizer_cpp = convert_quantizer(quantizer); + NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); + NVTE_CHECK(nvfp4_quantizer_cpp->row_scaled_nvfp4, + "Graph-safe grouped cast requires a row-scaled NVFP4 quantizer."); + + std::vector logical_shape; + for (const auto &d : tensor.sizes()) { + logical_shape.push_back(d); + } + const auto logical_first_dim = logical_shape[0]; // capacity (rows) + const auto logical_last_dim = logical_shape[1]; // K (common last dim) + + std::optional first_dims_opt = first_dims; + std::optional last_dims_opt = std::nullopt; + std::optional tensor_offsets_opt = tensor_offsets; + + auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); + grouped_input_tensor.set_rowwise_data( + tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), getTensorShape(tensor)); + grouped_input_tensor.set_first_dims(first_dims.data_ptr(), DType::kInt64, + getTensorShape(first_dims)); + grouped_input_tensor.set_tensor_offsets(tensor_offsets.data_ptr(), DType::kInt64, + getTensorShape(tensor_offsets)); + + auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( + num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), + py::reinterpret_borrow(quantizer), first_dims_opt, last_dims_opt, + tensor_offsets_opt, logical_first_dim, logical_last_dim); + + if (logical_first_dim == 0 || logical_last_dim == 0) { + return py::reinterpret_borrow(grouped_output_py); + } + + auto stream = at::cuda::getCurrentCUDAStream(); + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_nvfp4_compute_amax_graph_safe(grouped_input_tensor.data(), + grouped_output_tensor_cpp.data(), stream); + nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(grouped_input_tensor.data(), + grouped_output_tensor_cpp.data(), stream); + }); + + return py::reinterpret_borrow(grouped_output_py); +} + py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, @@ -1958,6 +2025,43 @@ void split_quantize_nvfp4_impl_helper(const TensorWrapper &input, } } +// Single-launch grouped row-scaled NVFP4 quantize: per-expert amax followed by a +// single grouped cast, replacing the per-expert quantize loop. Byte-identical to +// quantizing each expert on its own. +void split_quantize_nvfp4_row_scaled_impl(const TensorWrapper &input, + const std::vector &input_list, + std::vector &output_list, + const std::vector &split_sections, + const std::vector &quantizers) { + const size_t num_tensors = input_list.size(); + const auto &quantizer = *quantizers.front(); + NVTE_CHECK(!quantizer.stochastic_rounding, + "Row-scaled NVFP4 quantization does not support stochastic rounding."); + + std::vector nvte_tensor_output_list; + nvte_tensor_output_list.reserve(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + nvte_tensor_output_list.push_back(output_list[i].data()); + } + + QuantizationConfigWrapper quant_config; + const auto use_fast_math = transformer_engine::getenv("NVTE_USE_FAST_MATH"); + if (use_fast_math) { + quant_config.set_use_fast_math(true); + } + + auto stream = at::cuda::getCurrentCUDAStream(); + + // Per-expert rowwise (and columnwise, when transpose is requested) amax. + nvte_group_nvfp4_compute_amax(input.data(), nvte_tensor_output_list.data(), split_sections.data(), + num_tensors, stream); + + // Single grouped cast consuming the amax computed above. + nvte_group_nvfp4_row_scaled_cast_with_amax(input.data(), nvte_tensor_output_list.data(), + split_sections.data(), num_tensors, quant_config, + stream); +} + void split_quantize_nvfp4_impl(const TensorWrapper &input, const std::vector &input_list, std::vector &output_list, @@ -2097,7 +2201,7 @@ std::vector split_quantize(const at::Tensor &tensor, // Choose implementation for allocating and populating tensors enum class AllocationMethod { UNFUSED, BULK_FP8_BLOCKWISE, BULK_MXFP8, BULK_NVFP4 }; - enum class QuantizationMethod { UNFUSED, FUSED_NVFP4 }; + enum class QuantizationMethod { UNFUSED, FUSED_NVFP4, FUSED_NVFP4_ROW_SCALED }; AllocationMethod allocation_method = AllocationMethod::UNFUSED; QuantizationMethod quantization_method = QuantizationMethod::UNFUSED; if (!disable_bulk_allocation) { @@ -2122,7 +2226,23 @@ std::vector split_quantize(const at::Tensor &tensor, return static_cast(quantizer.get())->row_scaled_nvfp4; }); if (has_row_scaled_nvfp4) { - quantization_method = QuantizationMethod::UNFUSED; + // The grouped row-scaled cast requires every split to be 128-aligned (a + // 128-row chunk must stay inside one expert) and does not cover the + // 4over6 path; otherwise fall back to the per-expert path. + const bool all_splits_128 = std::all_of(split_sections.begin(), split_sections.end(), + [](size_t s) { return s % 128 == 0; }); + const bool any_4over6 = std::any_of( + quantizer_cpp_list.begin(), quantizer_cpp_list.end(), + [](const std::unique_ptr &quantizer) { + return static_cast(quantizer.get())->nvfp4_4over6_mode != + kNVTENVFP44Over6Disabled; + }); + const bool grouped_row_scaled_cast = + transformer_engine::getenv("NVTE_NVFP4_GROUPED_ROW_SCALED_CAST", true); + quantization_method = (grouped_row_scaled_cast && all_splits_128 && !any_4over6 && + input_dtype == DType::kBFloat16) + ? QuantizationMethod::FUSED_NVFP4_ROW_SCALED + : QuantizationMethod::UNFUSED; } else { quantization_method = QuantizationMethod::FUSED_NVFP4; } @@ -2203,6 +2323,17 @@ std::vector split_quantize(const at::Tensor &tensor, nvfp4_quantizers); break; } + case QuantizationMethod::FUSED_NVFP4_ROW_SCALED: { + // Single-launch grouped row-scaled NVFP4 quantize. + auto input_nvte = makeTransformerEngineTensor(input_dptr, input_shape, input_dtype); + std::vector nvfp4_quantizers; + for (auto &quantizer : quantizer_cpp_list) { + nvfp4_quantizers.push_back(static_cast(quantizer.get())); + } + split_quantize_nvfp4_row_scaled_impl(input_nvte, input_list, output_cpp_list, split_sections, + nvfp4_quantizers); + break; + } default: // General multi-tensor quantization multi_tensor_quantize_impl(input_list, quantizer_list, quantizer_cpp_list, output_cpp_list); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 2046408eba8..29c12754e4a 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -184,6 +184,11 @@ void bind_quantize_with_amax_extensions(py::module_ &m) { py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), py::arg("last_dims") = py::none(), py::arg("rowwise_amax"), py::arg("columnwise_amax"), py::arg("tensor_offsets") = py::none()); + m.def("nvfp4_group_row_scaled_cast_graph_safe", + transformer_engine::pytorch::nvfp4_group_row_scaled_cast_graph_safe, py::arg("tensor"), + py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("tensor_offsets"), + "CUDA-graph-safe grouped row-scaled NVFP4 cast (amax + cast) from device routing"); } } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 997f15561e5..2a3ce5feb55 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -38,6 +38,7 @@ "ep_finalize", "ep_dispatch", "ep_combine", + "nvfp4_group_row_scaled_cast", "symm_mem_alloc", "release_symm_mem_pool", "is_symm_backed", @@ -1127,6 +1128,48 @@ def _make_grouped_mxfp8(data, scale_inv, token_counts, fp8_dtype, fake_dtype): ) +def nvfp4_group_row_scaled_cast( + tokens: torch.Tensor, + tokens_per_expert: torch.Tensor, + *, + fp4_dtype: DType = DType.kFloat4E2M1, + scale_dtype: DType = DType.kFloat8E4M3, +): + """Graph-safe grouped row-scaled NVFP4 cast (amax + cast) from device routing. + + ``tokens`` is a packed (capacity, hidden) bfloat16 buffer; ``tokens_per_expert`` + (int64 [num_experts]) is the 128-aligned per-expert row count. Routing stays on + device, so the cast can be captured and replayed. Returns a row-scaled NVFP4 + ``GroupedTensor``. + """ + from .tensor.nvfp4_tensor import NVFP4Quantizer + + if tokens.dim() != 2: + raise ValueError("tokens must be 2D [capacity_rows, hidden].") + if tokens.dtype is not torch.bfloat16: + raise TypeError(f"tokens must be bfloat16; got dtype={tokens.dtype}.") + if tokens_per_expert.dtype is not torch.int64: + raise TypeError(f"tokens_per_expert must be int64; got dtype={tokens_per_expert.dtype}.") + + _, hidden = tokens.shape + num_tensors = tokens_per_expert.numel() + quantizer = NVFP4Quantizer( + fp4_dtype=fp4_dtype, + scale_dtype=scale_dtype, + rowwise=True, + columnwise=False, + row_scaled_nvfp4=True, + ) + tensor_offsets = tex.splits_to_offsets(tokens_per_expert, hidden) + return tex.nvfp4_group_row_scaled_cast_graph_safe( + tokens, + quantizer, + num_tensors, + tokens_per_expert, + tensor_offsets, + ) + + def ep_dispatch( buffer: EpBuffer, tokens: torch.Tensor, From b58bbe1fcb49e812ab5786a4eb2c2b2acc6a5a7a Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Thu, 8 Oct 2026 02:06:15 -0700 Subject: [PATCH 02/24] Address review: guard grouped row-scaled cast dispatch and inputs Restrict the FUSED_NVFP4_ROW_SCALED selection to shapes the grouped kernels accept (128-aligned last dim, <= 64 experts) so unsupported requests fall back to the per-expert path. Return early for empty routing, make the graph-safe binding's input contiguous, and reject with_rht/with_2d_quantization/with_amax_reduction on both row-scaled entries. Add regression tests for the four guarded paths. Signed-off-by: Cael Ling --- ..._nvfp4_group_row_scaled_cast_graph_safe.py | 83 +++++++++++++++++++ .../pytorch/csrc/extensions/cast.cpp | 40 +++++++-- 2 files changed, 114 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py index 3d683132c39..8d3e69e2faf 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py @@ -152,3 +152,86 @@ def program_routing(splits: list[int], x_src: torch.Tensor) -> None: graph.replay() torch.cuda.synchronize() _assert_matches_reference(grouped, _host_split_reference(x_b, splits_b), splits_b, N) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "split_sections, N", + [ + ([128, 128], 64), + ([128] * 65, 128), + ], +) +def test_split_quantize_falls_back_to_per_expert(split_sections: list[int], N: int) -> None: + """Shapes the grouped kernel rejects (last dim not 128-aligned, >64 experts) cast per-expert.""" + sum_m = sum(split_sections) + torch.manual_seed(3) + x = torch.randn((sum_m, N), dtype=torch.bfloat16, device="cuda") + + outputs = tex.split_quantize( + x, split_sections, [_row_scaled_quantizer() for _ in split_sections] + ) + + assert len(outputs) == len(split_sections) + for out, m in zip(outputs, split_sections): + assert tuple(out.shape) == (m, N) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_split_quantize_empty_routing() -> None: + """Empty row-scaled routing returns without launching grouped work.""" + split_sections = [0, 0] + x = torch.empty((0, 128), dtype=torch.bfloat16, device="cuda") + + outputs = tex.split_quantize( + x, split_sections, [_row_scaled_quantizer() for _ in split_sections] + ) + + assert len(outputs) == len(split_sections) + for out in outputs: + assert tuple(out.shape) == (0, 128) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_non_contiguous_input_matches_contiguous() -> None: + """A non-contiguous input is made contiguous before the graph-safe cast.""" + split_sections = [128, 128] + N = 256 + sum_m = sum(split_sections) + torch.manual_seed(4) + x_t = torch.randn((N, sum_m), dtype=torch.bfloat16, device="cuda").t() + assert not x_t.is_contiguous() + tokens_per_expert = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + grouped = nvfp4_group_row_scaled_cast(x_t, tokens_per_expert) + reference = _host_split_reference(x_t.contiguous(), split_sections) + + _assert_matches_reference(grouped, reference, split_sections, N) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_rejects_unsupported_quantizer_settings() -> None: + """with_2d_quantization is rejected on both the host-split and graph-safe row-scaled paths.""" + + def _quantizer_with_2d() -> NVFP4Quantizer: + return NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=False, + row_scaled_nvfp4=True, + with_2d_quantization=True, + ) + + split_sections = [128, 128] + N = 128 + x = torch.randn((sum(split_sections), N), dtype=torch.bfloat16, device="cuda") + + with pytest.raises(RuntimeError): + tex.split_quantize(x, split_sections, [_quantizer_with_2d() for _ in split_sections]) + + tokens_per_expert = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + tensor_offsets = tex.splits_to_offsets(tokens_per_expert, N) + with pytest.raises(RuntimeError): + tex.nvfp4_group_row_scaled_cast_graph_safe( + x, _quantizer_with_2d(), len(split_sections), tokens_per_expert, tensor_offsets + ) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 741de9726af..13e038ea6ed 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -662,9 +662,21 @@ py::object nvfp4_group_row_scaled_cast_graph_safe(const at::Tensor &tensor, py:: NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); NVTE_CHECK(nvfp4_quantizer_cpp->row_scaled_nvfp4, "Graph-safe grouped cast requires a row-scaled NVFP4 quantizer."); + NVTE_CHECK(!nvfp4_quantizer_cpp->with_rht, + "Graph-safe grouped row-scaled NVFP4 cast does not support RHT."); + NVTE_CHECK(!nvfp4_quantizer_cpp->with_2d_quantization, + "Graph-safe grouped row-scaled NVFP4 cast does not support 2D quantization."); + NVTE_CHECK(!nvfp4_quantizer_cpp->stochastic_rounding, + "Graph-safe grouped row-scaled NVFP4 cast does not support stochastic rounding."); + NVTE_CHECK(!nvfp4_quantizer_cpp->with_amax_reduction, + "Graph-safe grouped row-scaled NVFP4 cast does not support amax reduction."); + NVTE_CHECK(nvfp4_quantizer_cpp->nvfp4_4over6_mode == kNVTENVFP44Over6Disabled, + "Graph-safe grouped row-scaled NVFP4 cast does not support 4over6."); + + auto input_contiguous = tensor.contiguous(); std::vector logical_shape; - for (const auto &d : tensor.sizes()) { + for (const auto &d : input_contiguous.sizes()) { logical_shape.push_back(d); } const auto logical_first_dim = logical_shape[0]; // capacity (rows) @@ -675,15 +687,16 @@ py::object nvfp4_group_row_scaled_cast_graph_safe(const at::Tensor &tensor, py:: std::optional tensor_offsets_opt = tensor_offsets; auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); - grouped_input_tensor.set_rowwise_data( - tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), getTensorShape(tensor)); + grouped_input_tensor.set_rowwise_data(input_contiguous.data_ptr(), + GetTransformerEngineDType(input_contiguous.scalar_type()), + getTensorShape(input_contiguous)); grouped_input_tensor.set_first_dims(first_dims.data_ptr(), DType::kInt64, getTensorShape(first_dims)); grouped_input_tensor.set_tensor_offsets(tensor_offsets.data_ptr(), DType::kInt64, getTensorShape(tensor_offsets)); auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( - num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), + num_tensors, logical_shape, GetTransformerEngineDType(input_contiguous.scalar_type()), py::reinterpret_borrow(quantizer), first_dims_opt, last_dims_opt, tensor_offsets_opt, logical_first_dim, logical_last_dim); @@ -2034,9 +2047,17 @@ void split_quantize_nvfp4_row_scaled_impl(const TensorWrapper &input, const std::vector &split_sections, const std::vector &quantizers) { const size_t num_tensors = input_list.size(); + if (num_tensors == 0 || input.numel() == 0) { + return; + } const auto &quantizer = *quantizers.front(); NVTE_CHECK(!quantizer.stochastic_rounding, "Row-scaled NVFP4 quantization does not support stochastic rounding."); + NVTE_CHECK(!quantizer.with_rht, "Row-scaled NVFP4 quantization does not support RHT."); + NVTE_CHECK(!quantizer.with_2d_quantization, + "Row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK(!quantizer.with_amax_reduction, + "Row-scaled NVFP4 quantization does not support amax reduction."); std::vector nvte_tensor_output_list; nvte_tensor_output_list.reserve(num_tensors); @@ -2226,11 +2247,12 @@ std::vector split_quantize(const at::Tensor &tensor, return static_cast(quantizer.get())->row_scaled_nvfp4; }); if (has_row_scaled_nvfp4) { - // The grouped row-scaled cast requires every split to be 128-aligned (a - // 128-row chunk must stay inside one expert) and does not cover the - // 4over6 path; otherwise fall back to the per-expert path. + // Grouped kernels need 128-aligned splits and last dim, <= 64 experts, and + // no 4over6; otherwise fall back to the per-expert path. const bool all_splits_128 = std::all_of(split_sections.begin(), split_sections.end(), [](size_t s) { return s % 128 == 0; }); + const bool last_dim_128 = input_shape.back() % 128 == 0; + const bool within_expert_cap = num_splits <= 64; const bool any_4over6 = std::any_of( quantizer_cpp_list.begin(), quantizer_cpp_list.end(), [](const std::unique_ptr &quantizer) { @@ -2239,8 +2261,8 @@ std::vector split_quantize(const at::Tensor &tensor, }); const bool grouped_row_scaled_cast = transformer_engine::getenv("NVTE_NVFP4_GROUPED_ROW_SCALED_CAST", true); - quantization_method = (grouped_row_scaled_cast && all_splits_128 && !any_4over6 && - input_dtype == DType::kBFloat16) + quantization_method = (grouped_row_scaled_cast && all_splits_128 && last_dim_128 && + within_expert_cap && !any_4over6 && input_dtype == DType::kBFloat16) ? QuantizationMethod::FUSED_NVFP4_ROW_SCALED : QuantizationMethod::UNFUSED; } else { From 0af4dd939e90fae39aa9e2a8afcf98d64c822294 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Fri, 9 Oct 2026 05:02:51 -0700 Subject: [PATCH 03/24] Guard row-scaled grouped cast against non-2D, partial, and strided routing The grouped row-scaled cast treats its input as 2D [rows, last_dim] indexed by split counts, so route higher-rank inputs and partial splits (split sum below the row count) to the per-expert path instead of the grouped kernel. Make first_dims and tensor_offsets contiguous in the graph-safe cast so a strided routing view is not misread. Add fallback and strided-routing tests. Signed-off-by: Cael Ling --- ..._nvfp4_group_row_scaled_cast_graph_safe.py | 53 +++++++++++++++++++ .../pytorch/csrc/extensions/cast.cpp | 26 ++++++--- 2 files changed, 72 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py index 8d3e69e2faf..be35ce80d7e 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py @@ -177,6 +177,39 @@ def test_split_quantize_falls_back_to_per_expert(split_sections: list[int], N: i assert tuple(out.shape) == (m, N) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "input_shape, split_sections", + [ + ((256, 2, 128), [128, 128]), # higher-rank: not a 2D [rows, last_dim] input + ((512, 128), [128, 128]), # partial: splits cover 256 of 512 rows + ], +) +def test_split_quantize_non_grouped_shapes_fall_back( + input_shape, split_sections: list[int] +) -> None: + """Higher-rank inputs and partial splits bypass the 2D grouped kernel and cast per-expert, + matching an individual quantize of each split region.""" + torch.manual_seed(6) + x = torch.randn(input_shape, dtype=torch.bfloat16, device="cuda") + + outputs = tex.split_quantize( + x, split_sections, [_row_scaled_quantizer() for _ in split_sections] + ) + + assert len(outputs) == len(split_sections) + offset = 0 + for out, m in zip(outputs, split_sections): + ref = _row_scaled_quantizer()(x[offset : offset + m]) + torch.testing.assert_close( + out._rowwise_data.reshape(-1).view(torch.uint8), + ref._rowwise_data.reshape(-1).view(torch.uint8), + atol=0, + rtol=0, + ) + offset += m + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) def test_split_quantize_empty_routing() -> None: """Empty row-scaled routing returns without launching grouped work.""" @@ -209,6 +242,26 @@ def test_non_contiguous_input_matches_contiguous() -> None: _assert_matches_reference(grouped, reference, split_sections, N) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_strided_first_dims_matches_contiguous() -> None: + """A strided first_dims view is made contiguous before the cast, so it is not misread.""" + split_sections = [256, 128] + N = 128 + sum_m = sum(split_sections) + torch.manual_seed(5) + x = torch.randn((sum_m, N), dtype=torch.bfloat16, device="cuda") + + # [256, 0, 128, 0][::2] == [256, 128] but strided; a raw data_ptr read would see [256, 0]. + strided = torch.tensor([256, 0, 128, 0], dtype=torch.int64, device="cuda")[::2] + assert not strided.is_contiguous() + assert strided.tolist() == split_sections + + grouped = nvfp4_group_row_scaled_cast(x, strided) + reference = _host_split_reference(x, split_sections) + + _assert_matches_reference(grouped, reference, split_sections, N) + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) def test_rejects_unsupported_quantizer_settings() -> None: """with_2d_quantization is rejected on both the host-split and graph-safe row-scaled paths.""" diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 13e038ea6ed..ca9bcf82bab 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -674,6 +674,10 @@ py::object nvfp4_group_row_scaled_cast_graph_safe(const at::Tensor &tensor, py:: "Graph-safe grouped row-scaled NVFP4 cast does not support 4over6."); auto input_contiguous = tensor.contiguous(); + // The kernels read num_tensors consecutive int64 entries from these buffers, so a + // strided view (e.g. a slice) would be misread. Make them contiguous before use. + auto first_dims_contiguous = first_dims.contiguous(); + auto tensor_offsets_contiguous = tensor_offsets.contiguous(); std::vector logical_shape; for (const auto &d : input_contiguous.sizes()) { @@ -682,18 +686,18 @@ py::object nvfp4_group_row_scaled_cast_graph_safe(const at::Tensor &tensor, py:: const auto logical_first_dim = logical_shape[0]; // capacity (rows) const auto logical_last_dim = logical_shape[1]; // K (common last dim) - std::optional first_dims_opt = first_dims; + std::optional first_dims_opt = first_dims_contiguous; std::optional last_dims_opt = std::nullopt; - std::optional tensor_offsets_opt = tensor_offsets; + std::optional tensor_offsets_opt = tensor_offsets_contiguous; auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); grouped_input_tensor.set_rowwise_data(input_contiguous.data_ptr(), GetTransformerEngineDType(input_contiguous.scalar_type()), getTensorShape(input_contiguous)); - grouped_input_tensor.set_first_dims(first_dims.data_ptr(), DType::kInt64, - getTensorShape(first_dims)); - grouped_input_tensor.set_tensor_offsets(tensor_offsets.data_ptr(), DType::kInt64, - getTensorShape(tensor_offsets)); + grouped_input_tensor.set_first_dims(first_dims_contiguous.data_ptr(), DType::kInt64, + getTensorShape(first_dims_contiguous)); + grouped_input_tensor.set_tensor_offsets(tensor_offsets_contiguous.data_ptr(), DType::kInt64, + getTensorShape(tensor_offsets_contiguous)); auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( num_tensors, logical_shape, GetTransformerEngineDType(input_contiguous.scalar_type()), @@ -2253,6 +2257,13 @@ std::vector split_quantize(const at::Tensor &tensor, [](size_t s) { return s % 128 == 0; }); const bool last_dim_128 = input_shape.back() % 128 == 0; const bool within_expert_cap = num_splits <= 64; + // The grouped kernels treat the input as 2D [rows, last_dim] and index rows by + // split_sections, so they require a 2D input whose splits cover every row. For + // higher-rank or partially-split inputs, fall back to the per-expert path. + const bool input_2d = input_shape.size() == 2; + size_t split_sum = 0; + for (size_t s : split_sections) split_sum += s; + const bool splits_cover_input = input_2d && split_sum == input_shape[0]; const bool any_4over6 = std::any_of( quantizer_cpp_list.begin(), quantizer_cpp_list.end(), [](const std::unique_ptr &quantizer) { @@ -2262,7 +2273,8 @@ std::vector split_quantize(const at::Tensor &tensor, const bool grouped_row_scaled_cast = transformer_engine::getenv("NVTE_NVFP4_GROUPED_ROW_SCALED_CAST", true); quantization_method = (grouped_row_scaled_cast && all_splits_128 && last_dim_128 && - within_expert_cap && !any_4over6 && input_dtype == DType::kBFloat16) + within_expert_cap && input_2d && splits_cover_input && !any_4over6 && + input_dtype == DType::kBFloat16) ? QuantizationMethod::FUSED_NVFP4_ROW_SCALED : QuantizationMethod::UNFUSED; } else { From 07de76969febb07f1a87ac27de07328927c7d593 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 12:06:57 +0000 Subject: [PATCH 04/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/csrc/extensions/cast.cpp | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index ca9bcf82bab..b4354daca54 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -2272,11 +2272,11 @@ std::vector split_quantize(const at::Tensor &tensor, }); const bool grouped_row_scaled_cast = transformer_engine::getenv("NVTE_NVFP4_GROUPED_ROW_SCALED_CAST", true); - quantization_method = (grouped_row_scaled_cast && all_splits_128 && last_dim_128 && - within_expert_cap && input_2d && splits_cover_input && !any_4over6 && - input_dtype == DType::kBFloat16) - ? QuantizationMethod::FUSED_NVFP4_ROW_SCALED - : QuantizationMethod::UNFUSED; + quantization_method = + (grouped_row_scaled_cast && all_splits_128 && last_dim_128 && within_expert_cap && + input_2d && splits_cover_input && !any_4over6 && input_dtype == DType::kBFloat16) + ? QuantizationMethod::FUSED_NVFP4_ROW_SCALED + : QuantizationMethod::UNFUSED; } else { quantization_method = QuantizationMethod::FUSED_NVFP4; } From 5536e0c19550c241be4d3e625382b72ebcb3020b Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Fri, 9 Oct 2026 05:49:44 -0700 Subject: [PATCH 05/24] Fall back to per-expert cast for mixed NVFP4 row/tensor-scaled splits The bulk NVFP4 split_quantize allocation sizes its amax buffer from the first quantizer's row-scale setting, so a group mixing row-scaled and tensor-scaled quantizers writes per-row amax into a per-tensor buffer and corrupts neighbors. Detect the mixed case and fall back to the unfused per-split path. Signed-off-by: Cael Ling --- ..._nvfp4_group_row_scaled_cast_graph_safe.py | 39 +++++++++++++++++++ .../pytorch/csrc/extensions/cast.cpp | 21 +++++++--- 2 files changed, 54 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py index be35ce80d7e..293a802a706 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py @@ -210,6 +210,45 @@ def test_split_quantize_non_grouped_shapes_fall_back( offset += m +def _tensor_scaled_quantizer() -> NVFP4Quantizer: + return NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=False, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=False, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_split_quantize_mixed_row_scaled_falls_back() -> None: + """A split_quantize that mixes tensor-scaled and row-scaled NVFP4 quantizers casts each split + per-expert rather than taking the grouped row-scaled path (whose bulk amax, sized from the + first quantizer, would be overrun by the row-scaled kernel).""" + split_sections = [128, 128] + N = 128 + torch.manual_seed(7) + x = torch.randn((sum(split_sections), N), dtype=torch.bfloat16, device="cuda") + + quantizers = [_tensor_scaled_quantizer(), _row_scaled_quantizer()] + outputs = tex.split_quantize(x, split_sections, quantizers) + + assert len(outputs) == len(split_sections) + offset = 0 + for out, m, make_q in zip( + outputs, split_sections, [_tensor_scaled_quantizer, _row_scaled_quantizer] + ): + ref = make_q()(x[offset : offset + m]) + torch.testing.assert_close( + out._rowwise_data.reshape(-1).view(torch.uint8), + ref._rowwise_data.reshape(-1).view(torch.uint8), + atol=0, + rtol=0, + ) + offset += m + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) def test_split_quantize_empty_routing() -> None: """Empty row-scaled routing returns without launching grouped work.""" diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index b4354daca54..43b1c8fb367 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -2245,12 +2245,21 @@ std::vector split_quantize(const at::Tensor &tensor, return detail::IsNVFP4Quantizers(quantizer.ptr()); })) { allocation_method = AllocationMethod::BULK_NVFP4; - const bool has_row_scaled_nvfp4 = - std::any_of(quantizer_cpp_list.begin(), quantizer_cpp_list.end(), - [](const std::unique_ptr &quantizer) { - return static_cast(quantizer.get())->row_scaled_nvfp4; - }); - if (has_row_scaled_nvfp4) { + auto is_row_scaled = [](const std::unique_ptr &quantizer) { + return static_cast(quantizer.get())->row_scaled_nvfp4; + }; + const bool all_row_scaled = + std::all_of(quantizer_cpp_list.begin(), quantizer_cpp_list.end(), is_row_scaled); + const bool any_row_scaled = + std::any_of(quantizer_cpp_list.begin(), quantizer_cpp_list.end(), is_row_scaled); + if (any_row_scaled && !all_row_scaled) { + // Mixed row-scaled and tensor-scaled NVFP4 quantizers can't share the bulk NVFP4 + // allocation: its amax is sized from the first quantizer's row-scale setting, so a + // grouped cast would write per-row amax into a per-tensor buffer. Cast each split + // on its own instead. + allocation_method = AllocationMethod::UNFUSED; + quantization_method = QuantizationMethod::UNFUSED; + } else if (all_row_scaled) { // Grouped kernels need 128-aligned splits and last dim, <= 64 experts, and // no 4over6; otherwise fall back to the per-expert path. const bool all_splits_128 = std::all_of(split_sections.begin(), split_sections.end(), From 91e5585e786d43fce8e53e44cd5eb1cd40feb3f9 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Thu, 8 Oct 2026 22:47:30 -0700 Subject: [PATCH 06/24] Add row-scaled NVFP4 support to grouped-tensor GEMM Bake the cuBLAS-swizzled scale layout directly into the row-scaled grouped NVFP4 cast when optimize_for_gemm is set, mirroring the per-tensor and RHT cast-fusion paths. The grouped-tensor NVFP4 GEMM requires swizzled scales, and the standalone grouped swizzle supports only MXFP8, so the non-RHT row-scaled path had no swizzle producer. Consume row-scaled operands in general_grouped_gemm_for_grouped_tensor by neutralizing the per-row/per-column outer scale, running the per-tensor grouped GEMM into an FP32 scratch, and applying the outer scale afterward. This covers both the rowwise direction (fprop) and the columnwise direction (wgrad). Allow columnwise usage for row-scaled grouped NVFP4 tensors, lifting the prior restriction, and size the columnwise amax per column so wgrad can consume the transpose. The cast kernel already produced the columnwise transpose, amax, and swizzled scales. Add GEMM-level fprop, CUDA-graph, and wgrad parity tests against a per-expert dense row-scaled GEMM over byte-identical operands. Signed-off-by: Cael Ling --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 257 +++++++++++++++++- .../nvfp4/group_quantize_transpose_nvfp4.cuh | 110 +++++--- .../pytorch/cpp_extensions/gemm.py | 103 ++++++- transformer_engine/pytorch/csrc/quantizer.cpp | 9 +- .../tensor/storage/grouped_tensor_storage.py | 24 +- 5 files changed, 446 insertions(+), 57 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 5062cf6b42f..cd5299a8fd3 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -8,14 +8,25 @@ 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 recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) +# The grouped-tensor GEMM tests below need the graph-safe grouped row-scaled cast (PR #3649). +row_scaled_grouped_cast_available = hasattr(tex, "nvfp4_group_row_scaled_cast_graph_safe") +reason_for_no_row_scaled_grouped_cast = ( + "Grouped row-scaled NVFP4 cast (nvfp4_group_row_scaled_cast_graph_safe) is unavailable." +) + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( @@ -838,3 +849,247 @@ 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" + + +def _row_scaled_grouped_tensor_gemm_supported() -> bool: + if not torch.cuda.is_available(): + return False + if torch.cuda.get_device_capability() < (10, 0): + return False + return tex.get_cublasLt_version() >= 130300 + + +_reason_for_no_grouped_tensor_gemm = ( + "Row-scaled grouped-tensor NVFP4 GEMM requires Blackwell (SM100) and cuBLAS 13.3+." +) + + +def _make_row_scaled_grouped_activation( + m_splits: list[int], + k: int, + *, + columnwise: bool, + tokens: torch.Tensor | None = None, +) -> GroupedTensor: + """Pack per-expert bf16 activations into one row-scaled NVFP4 GroupedTensor. + + The graph-safe grouped row-scaled cast (PR #3649) requires each expert row count and + ``k`` to 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, + ) + # Bake cuBLAS-swizzled scales in the cast so the GEMM needs no separate swizzle pass. + quantizer.optimize_for_gemm = True + offsets = tex.splits_to_offsets(tokens_per_expert, k) + return tex.nvfp4_group_row_scaled_cast_graph_safe( + tokens, quantizer, len(m_splits), tokens_per_expert, offsets + ) + + +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.skipif( + not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast +) +@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: + """Forward (TN) grouped-tensor GEMM with a row-scaled activation matches a per-expert + dense GEMM. Only the activation is row-scaled (rowwise RHS), so this covers post_scale_rows. + """ + if not _row_scaled_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() + + # Identical FP4 operands and post-scale; only grouped-vs-single cuBLAS accumulation differs. + 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.skipif( + not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast +) +@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: + """The row-scaled grouped-tensor GEMM is CUDA-graph capturable: a captured GEMM replayed + after refreshing the activation buffers in place matches an eager recompute. + """ + if not _row_scaled_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") + + # Warmup to initialize kernels/allocator before capture. + torch.cuda.synchronize() + _run() + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _run() + + # Refresh the activation in place with freshly cast values, then replay. + 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.skipif( + not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast +) +@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: + """Weight-gradient (NT) grouped-tensor GEMM with two row-scaled operands matches a per-expert + dense GEMM. Both X and dY are consumed columnwise, covering post_scale_cols, post_scale_rows, + and the columnwise swizzled-scale epilogue. + """ + if not _row_scaled_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() + + # Identical FP4 operands and post-scales; only grouped-vs-single cuBLAS accumulation differs. + 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) 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 1850dd49a0b..9d499ffc141 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 { @@ -1410,7 +1411,8 @@ __device__ __forceinline__ int GetTensorId(const GroupedRowScaledCastArgs &args, // grouped row-scaled cast kernels. The two kernels differ only in how they // resolve the per-expert bases (host-baked args vs device metadata); everything // from the SMEM layout onward is identical. -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, @@ -1590,8 +1592,20 @@ __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) { + // gemm_swizzled_scale_idx takes the compact scale column count + // (scale_stride_row = K / SCALE_DIM); local_row is expert-local. + const ScaleType *scales = reinterpret_cast(&scales_vec); +#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[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) { @@ -1602,8 +1616,19 @@ __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) { + // Columnwise counterpart: column count is scale_stride_col = M_t / SCALE_DIM, + // and global_col indexes the K rows. + const ScaleType *scales = reinterpret_cast(&scales_vec); +#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[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); + } } } @@ -1619,7 +1644,8 @@ __device__ __forceinline__ void group_row_scaled_cast_body( } // Host-split kernel: per-expert bases come from the host-baked GroupedRowScaledCastArgs. -template +template __global__ void __launch_bounds__(THREADS_NUM) group_row_scaled_cast_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ GroupedRowScaledCastArgs args, @@ -1653,7 +1679,7 @@ __global__ void __launch_bounds__(THREADS_NUM) const size_t scale_stride_col = RETURN_TRANSPOSE ? static_cast(args.scale_stride_col[tensor_id]) : 0; - 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); @@ -1664,7 +1690,8 @@ __global__ void __launch_bounds__(THREADS_NUM) // Graph-safe kernel: per-expert bases derived on device from offsets / first_dims; // 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, @@ -1704,7 +1731,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); @@ -1756,6 +1783,7 @@ inline void group_row_scaled_cast(const Tensor &input, const Tensor *noop, NVTE_CHECK(first_output != nullptr, "No non-empty output tensor found."); const bool return_transpose = first_output->has_columnwise_data(); const size_t scale_stride_row = first_output->scale_inv.shape[1]; + const bool with_gemm_swizzled = first_output->with_gemm_swizzled_scales; GroupedRowScaledCastArgs args; args.num_tensors = 0; @@ -1818,19 +1846,22 @@ inline void group_row_scaled_cast(const Tensor &input, const Tensor *noop, 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_kernel; - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - kernel<<>>(tensor_map_input, args, noop_ptr, rows, 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_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + kernel<<>>(tensor_map_input, args, noop_ptr, rows, + 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); @@ -1872,6 +1903,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. @@ -1932,21 +1964,25 @@ 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; + 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 6939847c8dc..1f61f17b584 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1139,12 +1139,42 @@ 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.") + # Row-scaled operands carry an outer FP32 scale the per-tensor grouped GEMM cannot + # apply. Neutralize it, GEMM into an FP32 scratch, then post-scale the output + # (mirrors the dense general_gemm path). + def _is_row_scaled(operand) -> bool: + return isinstance(operand, GroupedTensorStorage) and operand.row_scaled_nvfp4 + + row_scaled = _is_row_scaled(A) or _is_row_scaled(B) + row_scaled_accumulate = accumulate + post_scale_rows = None # per-row outer scale -> output rows (from B) + post_scale_cols = None # per-group per-col outer scale -> output cols (from A) + if row_scaled: + if bias is not None: + raise NotImplementedError("Row-scaled NVFP4 grouped GEMM does not support fused bias.") + # A's scale lands on output columns, B's on output rows; rowwise amax for an + # untransposed operand, columnwise for a transposed one (see + # _nvfp4_row_scaled_gemm_inputs). + if _is_row_scaled(A): + A = A.copy() + if transa: + post_scale_cols = A.amax + A.amax = post_scale_cols.new_ones(A.num_tensors) + else: + post_scale_cols = A.columnwise_amax + A.columnwise_amax = post_scale_cols.new_ones(A.num_tensors) + A.row_scaled_nvfp4 = False + if _is_row_scaled(B): + B = B.copy() + if transb: + post_scale_rows = B.columnwise_amax + B.columnwise_amax = post_scale_rows.new_ones(B.num_tensors) + else: + post_scale_rows = B.amax + B.amax = post_scale_rows.new_ones(B.num_tensors) + B.row_scaled_nvfp4 = False + # Post-scale happens after the GEMM, so cuBLAS must not accumulate itself. + accumulate = False def _is_fp8_blockwise(operand) -> bool: if isinstance(operand, (list, tuple)): @@ -1209,14 +1239,37 @@ def _is_fp8_blockwise(operand) -> bool: sm_count = get_sm_count() sm_count = sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))) - return grouped_gemm_impl( + if not row_scaled: + return grouped_gemm_impl( + A, + transa, + B, + transb, + out, + bias, + bias_scale, + alpha, + beta, + workspace_setup, + workspace_cublas, + use_split_accumulator, + sm_count, + ) + + 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 + grouped_gemm_impl( A, transa, B, transb, - out, - bias, - bias_scale, + scratch, + None, + None, alpha, beta, workspace_setup, @@ -1224,3 +1277,33 @@ def _is_fp8_blockwise(operand) -> bool: use_split_accumulator, sm_count, ) + + def _store(dst: torch.Tensor, src: torch.Tensor) -> None: + if row_scaled_accumulate: + dst.add_(src.to(dst.dtype)) + else: + dst.copy_(src.to(dst.dtype)) + + if is_discrete_out: + for i in range(num_tensors): + scaled = scratch[i] + if post_scale_rows is not None: + rows = scaled.shape[-2] + row_scale_i = post_scale_rows.view(num_tensors, rows)[i] + scaled = scaled.mul_(row_scale_i.reshape((1,) * (scaled.dim() - 2) + (rows, 1))) + if post_scale_cols is not None: + cols = scaled.shape[-1] + col_scale_i = post_scale_cols.view(num_tensors, cols)[i] + scaled = scaled.mul_(col_scale_i.reshape((1,) * (scaled.dim() - 1) + (cols,))) + _store(out[i], scaled) + else: + last_dim = out.logical_shape[-1] + flat = scratch.rowwise_data.view(-1, last_dim) + if post_scale_rows is not None: + flat.mul_(post_scale_rows.view(-1, 1)) + if post_scale_cols is not None: + flat.view(num_tensors, -1, last_dim).mul_( + post_scale_cols.view(num_tensors, 1, last_dim) + ) + _store(out.rowwise_data.view(-1, last_dim), flat) + return out diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index fba268d2723..2492f2df25b 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2183,8 +2183,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; @@ -2207,7 +2205,12 @@ 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); + // Per-column for row-scaled, one per group otherwise (matches rowwise amax). + 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 6a48f85c6c7..9eafacb2701 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -860,13 +860,13 @@ 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." - ) + # Per-row/per-column for row-scaled, one per group otherwise. 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 +899,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 @@ -1116,6 +1118,7 @@ def split_into_quantized_tensors( columnwise_scale_inv_offsets.append(cum) self.columnwise_scale_inv_offsets = columnwise_scale_inv_offsets nvfp4_rowwise_amax_offsets = None + nvfp4_columnwise_amax_offsets = None row_scaled_nvfp4 = self.row_scaled_nvfp4 nvfp4_use_4over6 = self.nvfp4_use_4over6 nvfp4_e4m3_max = self.nvfp4_e4m3_max @@ -1123,9 +1126,13 @@ def split_into_quantized_tensors( if recipe.nvfp4() and row_scaled_nvfp4: cum = 0 nvfp4_rowwise_amax_offsets = [0] + ccum = 0 + nvfp4_columnwise_amax_offsets = [0] for i in range(self.num_tensors): cum += math.prod(self.tensor_shapes[i][:-1]) nvfp4_rowwise_amax_offsets.append(cum) + ccum += self.tensor_shapes[i][-1] + nvfp4_columnwise_amax_offsets.append(ccum) for i in range(self.num_tensors): quantizer = self.quantizer @@ -1329,7 +1336,12 @@ def split_into_quantized_tensors( amax_rowwise = self.amax[i : i + 1] if self.columnwise_amax is not None: - amax_columnwise = self.columnwise_amax[i : i + 1] + if nvfp4_columnwise_amax_offsets is not None: + camax_start = nvfp4_columnwise_amax_offsets[i] + camax_end = nvfp4_columnwise_amax_offsets[i + 1] + amax_columnwise = self.columnwise_amax[camax_start:camax_end] + else: + amax_columnwise = self.columnwise_amax[i : i + 1] if quantizer.internal: nvfp4_tensor_class = NVFP4TensorStorage From 26ab746660468418d7ab51e48a3b549b8da29cfa Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 09:22:15 +0000 Subject: [PATCH 07/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../nvfp4/group_quantize_transpose_nvfp4.cuh | 24 +++++++++---------- 1 file changed, 11 insertions(+), 13 deletions(-) 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 9d499ffc141..5ceeca6d151 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -1851,12 +1851,12 @@ inline void group_row_scaled_cast(const Tensor &input, const Tensor *noop, 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_kernel; + 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_kernel; NVTE_CHECK_CUDA(cudaFuncSetAttribute( kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); kernel<<>>(tensor_map_input, args, noop_ptr, rows, @@ -1969,13 +1969,11 @@ inline void group_row_scaled_cast_graph_safe(const GroupedTensor *input, const T 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; + 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<<>>( From 2ce6e701ecec8a4ad3da6810152a47a724dfd6df Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Fri, 9 Oct 2026 04:10:41 -0700 Subject: [PATCH 08/24] Harden row-scaled grouped GEMM post-scale and operand checks Reject explicit beta and listed row-scaled operands for the row-scaled NVFP4 grouped path instead of silently mishandling them. Apply the discrete-out post-scale with per-expert row/column offsets so uneven group sizes are handled, and guard the grouped-out column post-scale on a uniform per-expert output row count. Add a discrete-out wgrad test. Signed-off-by: Cael Ling --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 36 +++++++++++++++++++ .../pytorch/cpp_extensions/gemm.py | 36 ++++++++++++++++--- 2 files changed, 68 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index cd5299a8fd3..ff174363273 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -1093,3 +1093,39 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad( 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.skipif( + not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast +) +@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: + """Weight-gradient (NT) grouped-tensor GEMM into a list output exercises the discrete-out + post-scale path (per-expert row/column slices) against a per-expert dense GEMM. + """ + if not _row_scaled_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) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 1f61f17b584..313587afd77 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1145,6 +1145,16 @@ def general_grouped_gemm_for_grouped_tensor( def _is_row_scaled(operand) -> bool: return isinstance(operand, GroupedTensorStorage) and operand.row_scaled_nvfp4 + def _list_has_row_scaled(operand) -> bool: + return isinstance(operand, (list, tuple)) and any( + getattr(t, "_row_scaled_nvfp4", False) for t in operand + ) + + if _list_has_row_scaled(A) or _list_has_row_scaled(B): + raise NotImplementedError( + "Row-scaled NVFP4 operands passed as a list are not supported; pass a GroupedTensor." + ) + row_scaled = _is_row_scaled(A) or _is_row_scaled(B) row_scaled_accumulate = accumulate post_scale_rows = None # per-row outer scale -> output rows (from B) @@ -1152,6 +1162,10 @@ def _is_row_scaled(operand) -> bool: if row_scaled: 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." + ) # A's scale lands on output columns, B's on output rows; rowwise amax for an # untransposed operand, columnwise for a transposed one (see # _nvfp4_row_scaled_gemm_inputs). @@ -1285,23 +1299,37 @@ def _store(dst: torch.Tensor, src: torch.Tensor) -> None: dst.copy_(src.to(dst.dtype)) if is_discrete_out: + # Slice the per-row/per-column scales by each output's actual shape so uneven + # per-expert row/column counts are handled correctly. + row_offset = 0 + col_offset = 0 for i in range(num_tensors): scaled = scratch[i] + rows = scaled.shape[-2] + cols = scaled.shape[-1] if post_scale_rows is not None: - rows = scaled.shape[-2] - row_scale_i = post_scale_rows.view(num_tensors, rows)[i] + row_scale_i = post_scale_rows[row_offset : row_offset + rows] scaled = scaled.mul_(row_scale_i.reshape((1,) * (scaled.dim() - 2) + (rows, 1))) if post_scale_cols is not None: - cols = scaled.shape[-1] - col_scale_i = post_scale_cols.view(num_tensors, cols)[i] + col_scale_i = post_scale_cols[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: last_dim = out.logical_shape[-1] flat = scratch.rowwise_data.view(-1, last_dim) if post_scale_rows is not None: + # Per output row; handles uneven per-expert row counts directly. flat.mul_(post_scale_rows.view(-1, 1)) if post_scale_cols is not None: + # Per-expert, per-column. Only set for wgrad, where every expert has the + # same output row count, so an equal split over num_tensors is exact. + total_rows = flat.shape[0] + if total_rows % num_tensors != 0: + raise NotImplementedError( + "Row-scaled column post-scale requires a uniform per-expert output row count." + ) flat.view(num_tensors, -1, last_dim).mul_( post_scale_cols.view(num_tensors, 1, last_dim) ) From 32598473939b25e49180af16fe147fe6aada7b9b Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Fri, 9 Oct 2026 04:45:40 -0700 Subject: [PATCH 09/24] Apply row-scaled grouped GEMM column scales over actual group rows The grouped-out column post-scale split the output into equal-sized groups, which is wrong for uneven per-expert row counts and fails outright when the total is not divisible. Apply each group's column scale over its actual rows using host-side shapes, falling back to device first_dims only for varying- dimension tensors (never built under graph capture). Add an uneven TN test. Signed-off-by: Cael Ling --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 38 +++++++++++++++++++ .../pytorch/cpp_extensions/gemm.py | 25 +++++++----- 2 files changed, 53 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index ff174363273..1f180a1f55c 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -1129,3 +1129,41 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad_discrete_out( 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.skipif( + not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast +) +@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: + """TN grouped-tensor GEMM with both operands row-scaled and uneven output rows exercises the + grouped-out column post-scale over each group's actual rows (not an equal split). + """ + if not _row_scaled_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 + + # A (weights): uniform [n, k] per expert; B (activation): uneven [m_i, k]. + 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) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 313587afd77..100c4d86aab 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1323,15 +1323,20 @@ def _store(dst: torch.Tensor, src: torch.Tensor) -> None: # Per output row; handles uneven per-expert row counts directly. flat.mul_(post_scale_rows.view(-1, 1)) if post_scale_cols is not None: - # Per-expert, per-column. Only set for wgrad, where every expert has the - # same output row count, so an equal split over num_tensors is exact. - total_rows = flat.shape[0] - if total_rows % num_tensors != 0: - raise NotImplementedError( - "Row-scaled column post-scale requires a uniform per-expert output row count." - ) - flat.view(num_tensors, -1, last_dim).mul_( - post_scale_cols.view(num_tensors, 1, last_dim) - ) + # Per-expert, per-column: apply each group's scale over its actual rows. + # Uniform-construction exposes host tensor_shapes; varying-dim tensors carry + # first_dims and are never built under graph capture, so reading them is safe. + if scratch.tensor_shapes: + row_counts = [s[0] for s in scratch.tensor_shapes] + elif out.first_dims is None: + row_counts = [flat.shape[0] // num_tensors] * num_tensors + else: + row_counts = out.first_dims.tolist() + col_scales = post_scale_cols.view(num_tensors, last_dim) + row_offset = 0 + for i in range(num_tensors): + rows = row_counts[i] + flat[row_offset : row_offset + rows].mul_(col_scales[i].unsqueeze(0)) + row_offset += rows _store(out.rowwise_data.view(-1, last_dim), flat) return out From 07ceda5ac9b8c87744acd1508a2bc6a3fcb090a9 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Fri, 9 Oct 2026 06:01:04 -0700 Subject: [PATCH 10/24] Make grouped row-scaled GEMM column post-scale graph-safe The grouped single-buffer output path read per-expert row counts via first_dims.tolist(), which forces a device-to-host sync and breaks CUDA graph capture (the MXFP8 EP dispatch builds first_dims-backed outputs under capture). Expand the per-expert column scales with a device-side repeat_interleave whose output_size avoids the sync. Also reject a varying per-expert output width, which the shared-width expansion cannot represent. Signed-off-by: Cael Ling --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 59 +++++++++++++++++++ .../pytorch/cpp_extensions/gemm.py | 32 +++++----- 2 files changed, 76 insertions(+), 15 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 1f180a1f55c..77ad382396b 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -1167,3 +1167,62 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_uneven_colscale( 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.skipif( + not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast +) +@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: + """The grouped-out column post-scale (two row-scaled operands, uneven rows) is CUDA-graph + capturable: the device per-expert row counts drive the column-scale expansion with no host + sync, so a captured GEMM replayed after refreshing both operands matches an eager recompute. + """ + if not _row_scaled_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() + + # Refresh both operands in place with freshly cast values, then replay. + 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) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 100c4d86aab..a8084a1c515 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1317,26 +1317,28 @@ def _store(dst: torch.Tensor, src: torch.Tensor) -> None: row_offset += rows col_offset += cols else: + # Single packed output. A varying per-expert output width would need per-expert + # slicing over mismatched column counts; the per-column scale expansion below + # assumes one shared width. + 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] flat = scratch.rowwise_data.view(-1, last_dim) if post_scale_rows is not None: - # Per output row; handles uneven per-expert row counts directly. + # One scale per output row. flat.mul_(post_scale_rows.view(-1, 1)) if post_scale_cols is not None: - # Per-expert, per-column: apply each group's scale over its actual rows. - # Uniform-construction exposes host tensor_shapes; varying-dim tensors carry - # first_dims and are never built under graph capture, so reading them is safe. - if scratch.tensor_shapes: - row_counts = [s[0] for s in scratch.tensor_shapes] - elif out.first_dims is None: - row_counts = [flat.shape[0] // num_tensors] * num_tensors - else: - row_counts = out.first_dims.tolist() + # Per-expert, per-column. Expand each group's column scale over its rows from + # the device per-expert row counts. Passing output_size keeps repeat_interleave + # from syncing the counts to host, so this stays valid under CUDA-graph capture. col_scales = post_scale_cols.view(num_tensors, last_dim) - row_offset = 0 - for i in range(num_tensors): - rows = row_counts[i] - flat[row_offset : row_offset + rows].mul_(col_scales[i].unsqueeze(0)) - row_offset += rows + if out.first_dims is None: + rows = flat.shape[0] // 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=flat.shape[0])) _store(out.rowwise_data.view(-1, last_dim), flat) return out From e2517b9602e1e70251e9eaa6a149fdfdbfac736b Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Fri, 9 Oct 2026 06:44:56 -0700 Subject: [PATCH 11/24] Fix grouped row-scaled GEMM post-scale gating and cuBLAS test skip Restrict the grouped single-buffer post-scale to the live row count from logical_shape so a capacity-padded output does not misalign the per-row and per-group scales, and keep rejecting a varying per-expert output width (grouped GEMM outputs share the non-contracted dimension, so no TE caller produces one). Raise the test skip threshold to cuBLAS 13.4, matching CUBLAS_NVFP4_GROUPED_GEMM_VERSION, so the tests skip instead of failing the GEMM version check on 13.3. Signed-off-by: Cael Ling --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 70 ++++++++++++++++++- .../pytorch/cpp_extensions/gemm.py | 19 +++-- 2 files changed, 80 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 77ad382396b..9e044f5dddf 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -856,11 +856,13 @@ def _row_scaled_grouped_tensor_gemm_supported() -> bool: return False if torch.cuda.get_device_capability() < (10, 0): return False - return tex.get_cublasLt_version() >= 130300 + # cuBLAS enables grouped NVFP4 GEMM at 13.4 (CUBLAS_NVFP4_GROUPED_GEMM_VERSION 130400); + # below that the GEMM rejects the call, so skip rather than fail the version check. + return tex.get_cublasLt_version() >= 130400 _reason_for_no_grouped_tensor_gemm = ( - "Row-scaled grouped-tensor NVFP4 GEMM requires Blackwell (SM100) and cuBLAS 13.3+." + "Row-scaled grouped-tensor NVFP4 GEMM requires Blackwell (SM100) and cuBLAS 13.4+." ) @@ -1226,3 +1228,67 @@ def _run(): 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.skipif( + not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast +) +@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: + """A grouped output whose backing buffer carries extra capacity rows beyond the live region + is scaled only over the live rows. The post-scale restricts itself to logical_shape[0] + instead of the full buffer, so the extra rows neither break the per-row multiply nor change + the result, which still matches a per-expert dense GEMM. + """ + if not _row_scaled_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) + # Enlarge the backing buffer past the live region (logical_shape keeps the live size), as a + # capacity-padded output would. Groups are still written at their live offsets; without the + # live-row slice the per-row multiply would hit the extra rows and misalign. + 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/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index a8084a1c515..98481c0aab0 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1317,15 +1317,20 @@ def _store(dst: torch.Tensor, src: torch.Tensor) -> None: row_offset += rows col_offset += cols else: - # Single packed output. A varying per-expert output width would need per-expert - # slicing over mismatched column counts; the per-column scale expansion below - # assumes one shared width. + # Single packed output. Groups are stacked along rows (uniform or varying first + # dim) and share one output width, so the scales below index the buffer by row. + # A varying per-expert output width packs groups along columns instead, which this + # row-wise expansion cannot represent; no TE caller produces such an output, since + # grouped GEMM outputs share the non-contracted free dimension. 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] - flat = scratch.rowwise_data.view(-1, last_dim) + # logical_shape[0] is the live row count (host int). Restrict to it so a buffer + # carrying extra capacity rows does not misalign with the per-row/per-group scales. + live_rows = out.logical_shape[0] + flat = scratch.rowwise_data.view(-1, last_dim)[:live_rows] if post_scale_rows is not None: # One scale per output row. flat.mul_(post_scale_rows.view(-1, 1)) @@ -1335,10 +1340,10 @@ def _store(dst: torch.Tensor, src: torch.Tensor) -> None: # from syncing the counts to host, so this stays valid under CUDA-graph capture. col_scales = post_scale_cols.view(num_tensors, last_dim) if out.first_dims is None: - rows = flat.shape[0] // num_tensors + 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=flat.shape[0])) - _store(out.rowwise_data.view(-1, last_dim), flat) + 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 From 913fbf9bc02460636cc5008203b4aeef1066729a Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Wed, 21 Oct 2026 22:07:53 -0700 Subject: [PATCH 12/24] Route row-scaled NVFP4 grouped quantize through nvte_group_quantize Add the kNVTEGroupedRowScaledNVFP4 grouped tensor param and dispatch row-scaled NVFP4 outputs in nvte_group_quantize to the graph-safe amax and cast kernels. Handle row-scaled quantizers in tex.group_quantize. Remove the row-scaled-specific C APIs, the host-split row-scaled cast kernel, tex.nvfp4_group_row_scaled_cast_graph_safe, the ep.py helper, and the fused row-scaled split_quantize path. Add a top-level static_assert for unregistered types in TypeInfo. Signed-off-by: Cael Ling --- .../test_row_scaled_nvfp4_grouped_cast.cu | 474 +++--------------- tests/cpp/test_common.h | 6 +- ..._nvfp4_group_row_scaled_cast_graph_safe.py | 329 ------------ .../test_nvfp4_group_row_scaled_quantize.py | 272 ++++++++++ .../common/cast/cast_grouped.cu | 93 ---- .../common/cast/dispatch/quantize.cuh | 24 + .../nvfp4/group_quantize_transpose_nvfp4.cuh | 199 +------- transformer_engine/common/common.h | 10 +- .../common/include/transformer_engine/cast.h | 69 +-- .../transformer_engine/transformer_engine.h | 13 + .../common/transformer_engine.cpp | 6 + transformer_engine/pytorch/csrc/extensions.h | 5 - .../pytorch/csrc/extensions/cast.cpp | 240 +++------ .../pytorch/csrc/extensions/pybind.cpp | 5 - transformer_engine/pytorch/csrc/quantizer.cpp | 1 + .../pytorch/csrc/type_converters.cpp | 4 + transformer_engine/pytorch/ep.py | 43 -- 17 files changed, 464 insertions(+), 1329 deletions(-) delete mode 100644 tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py create mode 100644 tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py diff --git a/tests/cpp/operator/test_row_scaled_nvfp4_grouped_cast.cu b/tests/cpp/operator/test_row_scaled_nvfp4_grouped_cast.cu index 9e3af23fdcf..50b57304b0f 100644 --- a/tests/cpp/operator/test_row_scaled_nvfp4_grouped_cast.cu +++ b/tests/cpp/operator/test_row_scaled_nvfp4_grouped_cast.cu @@ -6,10 +6,7 @@ #include #include -#include #include -#include -#include #include #include #include @@ -29,8 +26,7 @@ using namespace test; namespace { -// (per-expert row counts, shared column count, request transpose); all values -// are multiples of 128 so every 128-row chunk stays inside a single expert. +// All row counts and K must be multiples of 128. struct GroupedCastConfig { std::vector Ms; size_t K; @@ -52,137 +48,31 @@ std::vector> make_outputs(const std::string& tag, return outs; } -// The grouped row-scaled cast reuses the single-tensor encode verbatim, so for -// every expert its FP4 data, E4M3 block scales and amaxes must be bit-identical -// to quantizing that expert on its own via nvte_quantize_v2. -void performGroupedRowScaledCastTest(const GroupedCastConfig& cfg) { - const std::vector& Ms = cfg.Ms; - const size_t K = cfg.K; - const bool transpose = cfg.transpose; - const size_t num_tensors = Ms.size(); - size_t sum_M = 0; - for (size_t m : Ms) sum_M += m; - - // Per-expert BF16 inputs, packed into one buffer with identical bytes so the - // grouped kernel and the per-expert oracle quantize the same values. +std::vector> make_random_inputs(const std::vector& Ms, size_t K) { std::vector> ins; - for (size_t i = 0; i < num_tensors; ++i) { - const std::vector shape = {Ms[i], K}; - ins.push_back(std::make_unique("in_" + std::to_string(i), shape, DType::kBFloat16)); + for (size_t i = 0; i < Ms.size(); ++i) { + ins.push_back( + std::make_unique("in_" + std::to_string(i), std::vector{Ms[i], K}, + DType::kBFloat16)); fillCase(ins[i].get(), InputsFillCase::uniform); ins[i]->to_cpu(); } + return ins; +} - Tensor packed("packed_input", std::vector{sum_M, K}, DType::kBFloat16); - bf16* pdst = packed.rowwise_cpu_dptr(); +std::vector pack_inputs(const std::vector>& ins, + const std::vector& Ms, size_t K) { + size_t sum_M = 0; + for (size_t m : Ms) sum_M += m; + std::vector packed(sum_M * K); size_t row_off = 0; - for (size_t i = 0; i < num_tensors; ++i) { - std::copy_n(ins[i]->rowwise_cpu_dptr(), Ms[i] * K, pdst + row_off * K); + for (size_t i = 0; i < ins.size(); ++i) { + std::copy_n(ins[i]->rowwise_cpu_dptr(), Ms[i] * K, packed.data() + row_off * K); row_off += Ms[i]; } - packed.from_cpu(); - - const std::vector splits(Ms); - - // System under test: grouped amax (Stage 1) + grouped row-scaled cast (Stage 2). - auto grouped = make_outputs("grp_", Ms, K, transpose); - std::vector grouped_handles; - for (auto& t : grouped) grouped_handles.push_back(t->data()); - nvte_group_nvfp4_compute_amax(packed.data(), grouped_handles.data(), splits.data(), num_tensors, - 0); - nvte_group_nvfp4_row_scaled_cast_with_amax(packed.data(), grouped_handles.data(), splits.data(), - num_tensors, nullptr, 0); - - // Oracle: quantize each expert independently (compute_amaxes + tuned_1D) via - // the generic nvte_quantize_v2 dispatch. - auto oracle = make_outputs("ref_", Ms, K, transpose); - QuantizationConfigWrapper quant_config; - quant_config.set_stochastic_rounding(false); - for (size_t i = 0; i < num_tensors; ++i) { - nvte_quantize_v2(ins[i]->data(), oracle[i]->data(), quant_config, 0); - } - - cudaDeviceSynchronize(); - ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); - - for (size_t i = 0; i < num_tensors; ++i) { - grouped[i]->to_cpu(); - oracle[i]->to_cpu(); - const size_t M = Ms[i]; - - // Rowwise FP4 data: M x K/2 packed bytes, fully written (no padding since - // K % 128 == 0). Read as fp4e2m1 (the tensor dtype) then compare raw bytes. - { - const uint8_t* g = reinterpret_cast(grouped[i]->rowwise_cpu_dptr()); - const uint8_t* r = reinterpret_cast(oracle[i]->rowwise_cpu_dptr()); - size_t mismatches = 0; - for (size_t b = 0; b < M * K / 2; ++b) { - if (g[b] != r[b]) ++mismatches; - } - ASSERT_EQ(mismatches, 0u) << "rowwise FP4 data mismatch: expert " << i << " (" << mismatches - << "/" << (M * K / 2) << " bytes)"; - } - - // Rowwise E4M3 block scales. - { - const std::array sd = get_scale_tensor_dims(M, K, 1, 16); - size_t mismatches = 0; - compare_scaling_factors("rowwise_scales_" + std::to_string(i), - grouped[i]->rowwise_cpu_scale_inv_ptr(), - oracle[i]->rowwise_cpu_scale_inv_ptr(), sd[0], sd[1], - sd[3], mismatches); - ASSERT_EQ(mismatches, 0u) << "rowwise scale mismatch: expert " << i; - } - - // Rowwise per-row amax. - { - const float* g = grouped[i]->cpu_rowwise_amax_ptr(); - const float* r = oracle[i]->cpu_rowwise_amax_ptr(); - for (size_t row = 0; row < M; ++row) { - ASSERT_EQ(g[row], r[row]) << "rowwise amax mismatch: expert " << i << " row " << row; - } - } - - if (!transpose) continue; - - // Columnwise (transpose) FP4 data: K x M/2 packed bytes. - { - const uint8_t* g = - reinterpret_cast(grouped[i]->columnwise_cpu_dptr()); - const uint8_t* r = - reinterpret_cast(oracle[i]->columnwise_cpu_dptr()); - size_t mismatches = 0; - for (size_t b = 0; b < K * M / 2; ++b) { - if (g[b] != r[b]) ++mismatches; - } - ASSERT_EQ(mismatches, 0u) << "columnwise FP4 data mismatch: expert " << i << " (" << mismatches - << "/" << (K * M / 2) << " bytes)"; - } - - // Columnwise E4M3 block scales. - { - const std::array sd_t = get_scale_tensor_dims(K, M, 1, 16); - size_t mismatches = 0; - compare_scaling_factors("columnwise_scales_" + std::to_string(i), - grouped[i]->columnwise_cpu_scale_inv_ptr(), - oracle[i]->columnwise_cpu_scale_inv_ptr(), sd_t[0], - sd_t[1], sd_t[3], mismatches); - ASSERT_EQ(mismatches, 0u) << "columnwise scale mismatch: expert " << i; - } - - // Columnwise per-col amax. - { - const float* g = grouped[i]->cpu_columnwise_amax_ptr(); - const float* r = oracle[i]->cpu_columnwise_amax_ptr(); - for (size_t col = 0; col < K; ++col) { - ASSERT_EQ(g[col], r[col]) << "columnwise amax mismatch: expert " << i << " col " << col; - } - } - } + return packed; } -// Host-split reference for one routing, gathered into the contiguous per-direction -// layout the graph-safe kernel writes. struct GroupedReference { size_t sum_M = 0; size_t K = 0; @@ -196,37 +86,18 @@ struct GroupedReference { bool transpose = false; }; -std::vector pack_inputs(const std::vector>& ins, - const std::vector& Ms, size_t K) { - size_t sum_M = 0; - for (size_t m : Ms) sum_M += m; - std::vector packed(sum_M * K); - size_t row_off = 0; - for (size_t i = 0; i < ins.size(); ++i) { - std::copy_n(ins[i]->rowwise_cpu_dptr(), Ms[i] * K, packed.data() + row_off * K); - row_off += Ms[i]; - } - return packed; -} - GroupedReference build_reference(const std::vector>& ins, const std::vector& Ms, size_t K, bool transpose) { const size_t num_tensors = Ms.size(); size_t sum_M = 0; for (size_t m : Ms) sum_M += m; - const std::vector packed_host = pack_inputs(ins, Ms, K); - Tensor packed("ref_packed", std::vector{sum_M, K}, DType::kBFloat16); - std::copy_n(packed_host.data(), sum_M * K, packed.rowwise_cpu_dptr()); - packed.from_cpu(); - - const std::vector splits(Ms); - auto ref = make_outputs("gsref_", Ms, K, transpose); - std::vector handles; - for (auto& t : ref) handles.push_back(t->data()); - nvte_group_nvfp4_compute_amax(packed.data(), handles.data(), splits.data(), num_tensors, 0); - nvte_group_nvfp4_row_scaled_cast_with_amax(packed.data(), handles.data(), splits.data(), - num_tensors, nullptr, 0); + auto ref = make_outputs("ref_", Ms, K, transpose); + QuantizationConfigWrapper quant_config; + quant_config.set_stochastic_rounding(false); + for (size_t i = 0; i < num_tensors; ++i) { + nvte_quantize_v2(ins[i]->data(), ref[i]->data(), quant_config, 0); + } cudaDeviceSynchronize(); GroupedReference r; @@ -244,8 +115,6 @@ GroupedReference build_reference(const std::vector>& ins r.col_amax.resize(num_tensors * K); } - // Gather per-expert scales into the contiguous (unpadded) layout the graph-safe - // kernel writes, stripping any allocation padding in the source stride. auto gather_scale = [](const fp8e4m3* src, fp8e4m3* dst, size_t rows, size_t cols, size_t src_stride) { for (size_t row = 0; row < rows; ++row) { @@ -282,9 +151,7 @@ GroupedReference build_reference(const std::vector>& ins return r; } -// Contiguous device buffers for one input/output grouped tensor pair, sized to a -// fixed capacity so a captured graph can be replayed with a redistributed routing. -struct GraphSafeBuffers { +struct GroupedBuffersNVFP4 { size_t capacity_rows = 0; size_t K = 0; size_t num_tensors = 0; @@ -326,8 +193,8 @@ struct GraphSafeBuffers { input = std::make_unique(num_tensors, std::vector{capacity_rows, K}); - output = std::make_unique(num_tensors, - std::vector{capacity_rows, K}); + output = std::make_unique( + num_tensors, std::vector{capacity_rows, K}, NVTE_NVFP4_1D_SCALING); input->set_rowwise_data(in_data.get(), DType::kBFloat16, std::vector{capacity_rows, K}); @@ -335,6 +202,7 @@ struct GraphSafeBuffers { std::vector{num_tensors + 1}); input->set_first_dims(first_dims.get(), DType::kInt64, std::vector{num_tensors}); + output->set_row_scaled_nvfp4(true); output->set_rowwise_data(out_data.get(), DType::kFloat4E2M1, std::vector{capacity_rows, K}); output->set_rowwise_scale_inv(out_scale.get(), DType::kFloat8E4M3, @@ -353,10 +221,8 @@ struct GraphSafeBuffers { } } - // Program device metadata + input for one routing. seed_amax: fill amax from the - // reference (cast-only tests) or poison it (so a missing amax write is caught). void set_routing(const std::vector& Ms, const GroupedReference& ref, - const std::vector& packed_in, bool seed_amax = true) { + const std::vector& packed_in) { std::vector off(num_tensors + 1, 0); std::vector fd(num_tensors); for (size_t i = 0; i < num_tensors; ++i) { @@ -369,45 +235,23 @@ struct GraphSafeBuffers { cudaMemcpyHostToDevice)); NVTE_CHECK_CUDA(cudaMemcpy(in_data.get(), packed_in.data(), ref.sum_M * K * sizeof(bf16), cudaMemcpyHostToDevice)); - if (seed_amax) { - NVTE_CHECK_CUDA(cudaMemcpy(out_amax.get(), ref.row_amax.data(), ref.sum_M * sizeof(float), - cudaMemcpyHostToDevice)); - if (transpose) { - NVTE_CHECK_CUDA(cudaMemcpy(out_col_amax.get(), ref.col_amax.data(), - num_tensors * K * sizeof(float), cudaMemcpyHostToDevice)); - } - } else { - std::vector poison_row(capacity_rows, -1.0f); - NVTE_CHECK_CUDA(cudaMemcpy(out_amax.get(), poison_row.data(), capacity_rows * sizeof(float), - cudaMemcpyHostToDevice)); - if (transpose) { - std::vector poison_col(num_tensors * K, -1.0f); - NVTE_CHECK_CUDA(cudaMemcpy(out_col_amax.get(), poison_col.data(), - num_tensors * K * sizeof(float), cudaMemcpyHostToDevice)); - } + std::vector poison_row(capacity_rows, -1.0f); + NVTE_CHECK_CUDA(cudaMemcpy(out_amax.get(), poison_row.data(), capacity_rows * sizeof(float), + cudaMemcpyHostToDevice)); + if (transpose) { + std::vector poison_col(num_tensors * K, -1.0f); + NVTE_CHECK_CUDA(cudaMemcpy(out_col_amax.get(), poison_col.data(), + num_tensors * K * sizeof(float), cudaMemcpyHostToDevice)); } } -}; - -// The graph-safe amax must reproduce the host-split amax exactly: rowwise amax -// packed by global row [sum_M], columnwise amax packed per expert [num_tensors*K]. -void expect_amax_matches_reference(const GraphSafeBuffers& b, const GroupedReference& ref) { - const size_t sum_M = ref.sum_M; - std::vector h_amax(sum_M); - NVTE_CHECK_CUDA( - cudaMemcpy(h_amax.data(), b.out_amax.get(), sum_M * sizeof(float), cudaMemcpyDeviceToHost)); - for (size_t i = 0; i < sum_M; ++i) - ASSERT_EQ(h_amax[i], ref.row_amax[i]) << "rowwise amax mismatch at row " << i; - if (!ref.transpose) return; - std::vector h_col_amax(b.num_tensors * ref.K); - NVTE_CHECK_CUDA(cudaMemcpy(h_col_amax.data(), b.out_col_amax.get(), - b.num_tensors * ref.K * sizeof(float), cudaMemcpyDeviceToHost)); - for (size_t i = 0; i < b.num_tensors * ref.K; ++i) - ASSERT_EQ(h_col_amax[i], ref.col_amax[i]) << "columnwise amax mismatch at index " << i; -} + void quantize(cudaStream_t stream) { + QuantizationConfigWrapper quant_config; + nvte_group_quantize(input->data(), output->data(), quant_config, stream); + } +}; -void expect_matches_reference(const GraphSafeBuffers& b, const GroupedReference& ref) { +void expect_matches_reference(const GroupedBuffersNVFP4& b, const GroupedReference& ref) { const size_t sum_M = ref.sum_M; const size_t K = ref.K; @@ -420,6 +264,8 @@ void expect_matches_reference(const GraphSafeBuffers& b, const GroupedReference& NVTE_CHECK_CUDA( cudaMemcpy(h_amax.data(), b.out_amax.get(), sum_M * sizeof(float), cudaMemcpyDeviceToHost)); + for (size_t i = 0; i < sum_M; ++i) + ASSERT_EQ(h_amax[i], ref.row_amax[i]) << "rowwise amax mismatch at row " << i; size_t data_mismatch = 0; for (size_t i = 0; i < h_data.size(); ++i) if (h_data[i] != ref.row_data[i]) ++data_mismatch; @@ -430,17 +276,20 @@ void expect_matches_reference(const GraphSafeBuffers& b, const GroupedReference& if (h_scale[i] != ref.row_scale[i]) ++scale_mismatch; ASSERT_EQ(scale_mismatch, 0u) << "rowwise scale mismatch (" << scale_mismatch << "/" << h_scale.size() << ")"; - for (size_t i = 0; i < sum_M; ++i) - ASSERT_EQ(h_amax[i], ref.row_amax[i]) << "rowwise amax mismatch at row " << i; if (!ref.transpose) return; + std::vector h_col_amax(b.num_tensors * K); std::vector h_col_data((K / 2) * sum_M); std::vector h_col_scale((K / 16) * sum_M); + NVTE_CHECK_CUDA(cudaMemcpy(h_col_amax.data(), b.out_col_amax.get(), + h_col_amax.size() * sizeof(float), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(h_col_data.data(), b.out_col_data.get(), h_col_data.size(), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(h_col_scale.data(), b.out_col_scale.get(), h_col_scale.size(), cudaMemcpyDeviceToHost)); + for (size_t i = 0; i < h_col_amax.size(); ++i) + ASSERT_EQ(h_col_amax[i], ref.col_amax[i]) << "columnwise amax mismatch at index " << i; size_t col_data_mismatch = 0; for (size_t i = 0; i < h_col_data.size(); ++i) if (h_col_data[i] != ref.col_data[i]) ++col_data_mismatch; @@ -453,24 +302,10 @@ void expect_matches_reference(const GraphSafeBuffers& b, const GroupedReference& << h_col_scale.size() << ")"; } -std::vector> make_random_inputs(const std::vector& Ms, size_t K) { - std::vector> ins; - for (size_t i = 0; i < Ms.size(); ++i) { - ins.push_back( - std::make_unique("gsin_" + std::to_string(i), std::vector{Ms[i], K}, - DType::kBFloat16)); - fillCase(ins[i].get(), InputsFillCase::uniform); - ins[i]->to_cpu(); - } - return ins; -} - -// Graph-safe cast must be byte-identical to the host-split cast. -void performGraphSafeEqualityTest(const GroupedCastConfig& cfg) { +void performGroupedRowScaledQuantizeTest(const GroupedCastConfig& cfg) { const std::vector& Ms = cfg.Ms; const size_t K = cfg.K; const bool transpose = cfg.transpose; - const size_t num_tensors = Ms.size(); size_t sum_M = 0; for (size_t m : Ms) sum_M += m; @@ -478,20 +313,18 @@ void performGraphSafeEqualityTest(const GroupedCastConfig& cfg) { const GroupedReference ref = build_reference(ins, Ms, K, transpose); const std::vector packed_in = pack_inputs(ins, Ms, K); - GraphSafeBuffers b; - b.allocate(sum_M, K, num_tensors, ref.row_scale_cols, transpose); + GroupedBuffersNVFP4 b; + b.allocate(sum_M, K, Ms.size(), ref.row_scale_cols, transpose); b.set_routing(Ms, ref, packed_in); - nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(b.input->data(), b.output->data(), 0); + b.quantize(0); cudaDeviceSynchronize(); ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); expect_matches_reference(b, ref); } -// Capture the graph-safe cast once, then replay after redistributing the routing -// on device; both replays must match their host-split reference. -void performGraphSafeReplayTest(size_t K, bool transpose) { +void performGroupedRowScaledQuantizeReplayTest(size_t K, bool transpose) { const std::vector Ms_a = {256, 128, 128}; const std::vector Ms_b = {128, 256, 128}; const size_t num_tensors = Ms_a.size(); @@ -504,7 +337,7 @@ void performGraphSafeReplayTest(size_t K, bool transpose) { const std::vector packed_a = pack_inputs(ins_a, Ms_a, K); const std::vector packed_b = pack_inputs(ins_b, Ms_b, K); - GraphSafeBuffers b; + GroupedBuffersNVFP4 b; b.allocate(capacity_rows, K, num_tensors, ref_a.row_scale_cols, transpose); b.set_routing(Ms_a, ref_a, packed_a); @@ -513,13 +346,14 @@ void performGraphSafeReplayTest(size_t K, bool transpose) { // Warm up outside capture so one-time setup (e.g. cudaFuncSetAttribute) does // not land in the captured graph. - nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(b.input->data(), b.output->data(), stream); + b.quantize(stream); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + b.set_routing(Ms_a, ref_a, packed_a); cudaGraph_t graph; cudaGraphExec_t exec; NVTE_CHECK_CUDA(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed)); - nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(b.input->data(), b.output->data(), stream); + b.quantize(stream); NVTE_CHECK_CUDA(cudaStreamEndCapture(stream, &graph)); NVTE_CHECK_CUDA(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0)); @@ -528,7 +362,6 @@ void performGraphSafeReplayTest(size_t K, bool transpose) { ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); expect_matches_reference(b, ref_a); - // Redistribute the routing on device and replay the same captured graph. b.set_routing(Ms_b, ref_b, packed_b); NVTE_CHECK_CUDA(cudaGraphLaunch(exec, stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); @@ -540,125 +373,20 @@ void performGraphSafeReplayTest(size_t K, bool transpose) { NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } -// Graph-safe amax must be byte-identical to the host-split amax. -void performGraphSafeAmaxEqualityTest(const GroupedCastConfig& cfg) { - const std::vector& Ms = cfg.Ms; - const size_t K = cfg.K; - const bool transpose = cfg.transpose; - const size_t num_tensors = Ms.size(); - size_t sum_M = 0; - for (size_t m : Ms) sum_M += m; - - auto ins = make_random_inputs(Ms, K); - const GroupedReference ref = build_reference(ins, Ms, K, transpose); - const std::vector packed_in = pack_inputs(ins, Ms, K); - - GraphSafeBuffers b; - b.allocate(sum_M, K, num_tensors, ref.row_scale_cols, transpose); - b.set_routing(Ms, ref, packed_in, /*seed_amax=*/false); - - nvte_group_nvfp4_compute_amax_graph_safe(b.input->data(), b.output->data(), 0); - cudaDeviceSynchronize(); - ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); - - expect_amax_matches_reference(b, ref); -} - -// Full graph-safe path (amax then cast, no host-seeded amax); must be byte-identical -// to the host-split reference. -void performGraphSafeFullEqualityTest(const GroupedCastConfig& cfg) { - const std::vector& Ms = cfg.Ms; - const size_t K = cfg.K; - const bool transpose = cfg.transpose; - const size_t num_tensors = Ms.size(); - size_t sum_M = 0; - for (size_t m : Ms) sum_M += m; - - auto ins = make_random_inputs(Ms, K); - const GroupedReference ref = build_reference(ins, Ms, K, transpose); - const std::vector packed_in = pack_inputs(ins, Ms, K); - - GraphSafeBuffers b; - b.allocate(sum_M, K, num_tensors, ref.row_scale_cols, transpose); - b.set_routing(Ms, ref, packed_in, /*seed_amax=*/false); - - nvte_group_nvfp4_compute_amax_graph_safe(b.input->data(), b.output->data(), 0); - nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(b.input->data(), b.output->data(), 0); - cudaDeviceSynchronize(); - ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); - - expect_amax_matches_reference(b, ref); - expect_matches_reference(b, ref); -} - -// Capture the full graph-safe path (amax + cast) once, then replay after -// redistributing the routing on device; amax is recomputed on every replay. -void performGraphSafeFullReplayTest(size_t K, bool transpose) { - const std::vector Ms_a = {256, 128, 128}; - const std::vector Ms_b = {128, 256, 128}; - const size_t num_tensors = Ms_a.size(); - const size_t capacity_rows = 768; // > sum_M (512): exercises the capacity noop - - auto ins_a = make_random_inputs(Ms_a, K); - auto ins_b = make_random_inputs(Ms_b, K); - const GroupedReference ref_a = build_reference(ins_a, Ms_a, K, transpose); - const GroupedReference ref_b = build_reference(ins_b, Ms_b, K, transpose); - const std::vector packed_a = pack_inputs(ins_a, Ms_a, K); - const std::vector packed_b = pack_inputs(ins_b, Ms_b, K); - - GraphSafeBuffers b; - b.allocate(capacity_rows, K, num_tensors, ref_a.row_scale_cols, transpose); - b.set_routing(Ms_a, ref_a, packed_a, /*seed_amax=*/false); - - cudaStream_t stream; - NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - - // Warm up outside capture so one-time setup (e.g. cudaFuncSetAttribute) does - // not land in the captured graph. - nvte_group_nvfp4_compute_amax_graph_safe(b.input->data(), b.output->data(), stream); - nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(b.input->data(), b.output->data(), stream); - NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - - cudaGraph_t graph; - cudaGraphExec_t exec; - NVTE_CHECK_CUDA(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed)); - nvte_group_nvfp4_compute_amax_graph_safe(b.input->data(), b.output->data(), stream); - nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(b.input->data(), b.output->data(), stream); - NVTE_CHECK_CUDA(cudaStreamEndCapture(stream, &graph)); - NVTE_CHECK_CUDA(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0)); - - NVTE_CHECK_CUDA(cudaGraphLaunch(exec, stream)); - NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); - expect_amax_matches_reference(b, ref_a); - expect_matches_reference(b, ref_a); - - // Redistribute the routing on device and replay the same captured graph. - b.set_routing(Ms_b, ref_b, packed_b, /*seed_amax=*/false); - NVTE_CHECK_CUDA(cudaGraphLaunch(exec, stream)); - NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); - expect_amax_matches_reference(b, ref_b); - expect_matches_reference(b, ref_b); - - NVTE_CHECK_CUDA(cudaGraphExecDestroy(exec)); - NVTE_CHECK_CUDA(cudaGraphDestroy(graph)); - NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); -} - } // namespace -class NVFP4GroupedRowScaledCastTestSuite : public ::testing::TestWithParam {}; +class NVFP4GroupedRowScaledQuantizeTestSuite + : public ::testing::TestWithParam {}; -TEST_P(NVFP4GroupedRowScaledCastTestSuite, MatchesPerExpertOracle) { +TEST_P(NVFP4GroupedRowScaledQuantizeTestSuite, MatchesPerExpertQuantize) { if (getDeviceComputeCapability() < blackwellComputeCapability) { GTEST_SKIP(); } - performGroupedRowScaledCastTest(GetParam()); + performGroupedRowScaledQuantizeTest(GetParam()); } INSTANTIATE_TEST_SUITE_P( - NVFP4GroupedRowScaledCastRowwise, NVFP4GroupedRowScaledCastTestSuite, + NVFP4GroupedRowScaledQuantizeRowwise, NVFP4GroupedRowScaledQuantizeTestSuite, ::testing::Values(GroupedCastConfig{{128}, 128, false}, GroupedCastConfig{{128, 128}, 128, false}, GroupedCastConfig{{128, 256, 128}, 256, false}, @@ -667,7 +395,7 @@ INSTANTIATE_TEST_SUITE_P( GroupedCastConfig{{512}, 1024, false})); INSTANTIATE_TEST_SUITE_P( - NVFP4GroupedRowScaledCastTranspose, NVFP4GroupedRowScaledCastTestSuite, + NVFP4GroupedRowScaledQuantizeTranspose, NVFP4GroupedRowScaledQuantizeTestSuite, ::testing::Values(GroupedCastConfig{{128}, 128, true}, GroupedCastConfig{{128, 128}, 128, true}, GroupedCastConfig{{128, 256, 128}, 256, true}, @@ -675,85 +403,17 @@ INSTANTIATE_TEST_SUITE_P( GroupedCastConfig{{384, 128, 256}, 128, true}, GroupedCastConfig{{512}, 1024, true})); -class NVFP4GraphSafeGroupedRowScaledCastTestSuite - : public ::testing::TestWithParam {}; - -TEST_P(NVFP4GraphSafeGroupedRowScaledCastTestSuite, MatchesHostSplit) { - if (getDeviceComputeCapability() < blackwellComputeCapability) { - GTEST_SKIP(); - } - performGraphSafeEqualityTest(GetParam()); -} - -INSTANTIATE_TEST_SUITE_P( - NVFP4GraphSafeGroupedRowScaledCast, NVFP4GraphSafeGroupedRowScaledCastTestSuite, - ::testing::Values(GroupedCastConfig{{128}, 128, false}, - GroupedCastConfig{{128, 256, 128}, 256, false}, - GroupedCastConfig{{384, 128, 256}, 128, false}, - GroupedCastConfig{{128}, 128, true}, - GroupedCastConfig{{128, 256, 128}, 256, true}, - GroupedCastConfig{{384, 128, 256}, 128, true})); - -class NVFP4GraphSafeGroupedRowScaledCastReplaySuite +class NVFP4GroupedRowScaledQuantizeReplaySuite : public ::testing::TestWithParam {}; -TEST_P(NVFP4GraphSafeGroupedRowScaledCastReplaySuite, ReplayWithChangedRouting) { - if (getDeviceComputeCapability() < blackwellComputeCapability) { - GTEST_SKIP(); - } - performGraphSafeReplayTest(GetParam().K, GetParam().transpose); -} - -INSTANTIATE_TEST_SUITE_P(NVFP4GraphSafeGroupedRowScaledCastReplay, - NVFP4GraphSafeGroupedRowScaledCastReplaySuite, - ::testing::Values(GroupedCastConfig{{}, 256, false}, - GroupedCastConfig{{}, 256, true})); - -class NVFP4GraphSafeGroupedAmaxTestSuite : public ::testing::TestWithParam {}; - -TEST_P(NVFP4GraphSafeGroupedAmaxTestSuite, MatchesHostSplit) { - if (getDeviceComputeCapability() < blackwellComputeCapability) { - GTEST_SKIP(); - } - performGraphSafeAmaxEqualityTest(GetParam()); -} - -INSTANTIATE_TEST_SUITE_P( - NVFP4GraphSafeGroupedAmax, NVFP4GraphSafeGroupedAmaxTestSuite, - ::testing::Values(GroupedCastConfig{{128}, 128, false}, - GroupedCastConfig{{128, 256, 128}, 256, false}, - GroupedCastConfig{{384, 128, 256}, 128, false}, - GroupedCastConfig{{128}, 128, true}, - GroupedCastConfig{{128, 256, 128}, 256, true}, - GroupedCastConfig{{384, 128, 256}, 128, true})); - -class NVFP4GraphSafeGroupedFullTestSuite : public ::testing::TestWithParam {}; - -TEST_P(NVFP4GraphSafeGroupedFullTestSuite, AmaxThenCastMatchesHostSplit) { - if (getDeviceComputeCapability() < blackwellComputeCapability) { - GTEST_SKIP(); - } - performGraphSafeFullEqualityTest(GetParam()); -} - -INSTANTIATE_TEST_SUITE_P( - NVFP4GraphSafeGroupedFull, NVFP4GraphSafeGroupedFullTestSuite, - ::testing::Values(GroupedCastConfig{{128}, 128, false}, - GroupedCastConfig{{128, 256, 128}, 256, false}, - GroupedCastConfig{{384, 128, 256}, 128, false}, - GroupedCastConfig{{128}, 128, true}, - GroupedCastConfig{{128, 256, 128}, 256, true}, - GroupedCastConfig{{384, 128, 256}, 128, true})); - -class NVFP4GraphSafeGroupedFullReplaySuite : public ::testing::TestWithParam {}; - -TEST_P(NVFP4GraphSafeGroupedFullReplaySuite, ReplayWithChangedRouting) { +TEST_P(NVFP4GroupedRowScaledQuantizeReplaySuite, ReplayWithChangedRouting) { if (getDeviceComputeCapability() < blackwellComputeCapability) { GTEST_SKIP(); } - performGraphSafeFullReplayTest(GetParam().K, GetParam().transpose); + performGroupedRowScaledQuantizeReplayTest(GetParam().K, GetParam().transpose); } -INSTANTIATE_TEST_SUITE_P(NVFP4GraphSafeGroupedFullReplay, NVFP4GraphSafeGroupedFullReplaySuite, +INSTANTIATE_TEST_SUITE_P(NVFP4GroupedRowScaledQuantizeReplay, + NVFP4GroupedRowScaledQuantizeReplaySuite, ::testing::Values(GroupedCastConfig{{}, 256, false}, GroupedCastConfig{{}, 256, true})); diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index cd42847afb9..5beed3c47c7 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -108,9 +108,7 @@ struct TypeInfo { struct Helper { constexpr static DType getType() { constexpr int i = static_cast(current); - // The DType enum carries values (e.g. kFloat8UE5M3) whose type is only - // added to ``types`` on CUDA >= 13.4, so guard the tuple access to keep - // the enum walk from indexing past the tuple on older toolkits. + // fp8ue5m3 is only in types on CUDA >= 13.4. if constexpr (i < static_cast(std::tuple_size::value)) { if (std::is_same::type>::value) { return current; @@ -129,6 +127,8 @@ struct TypeInfo { template constexpr static DType getType() { + static_assert(Helper::getType() != DType::kNumTypes, + "Type is not registered"); return Helper::getType(); } diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py deleted file mode 100644 index 293a802a706..00000000000 --- a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_cast_graph_safe.py +++ /dev/null @@ -1,329 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Graph-safe grouped row-scaled NVFP4 cast (amax + cast) from device-side routing. - -Exercises ``transformer_engine.pytorch.ep.nvfp4_group_row_scaled_cast`` against the -host-split row-scaled cast (``tex.split_quantize`` with row-scaled NVFP4 quantizers). -""" - -import pytest -import torch - -import transformer_engine.pytorch as te -import transformer_engine_torch as tex -from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.ep import nvfp4_group_row_scaled_cast - -recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) - - -def _row_scaled_quantizer() -> NVFP4Quantizer: - return NVFP4Quantizer( - fp4_dtype=te.DType.kFloat4E2M1, - rowwise=True, - columnwise=False, - with_rht=False, - with_post_rht_amax=False, - row_scaled_nvfp4=True, - ) - - -def _host_split_reference(x_valid: torch.Tensor, split_sections: list[int]): - quantizers = [_row_scaled_quantizer() for _ in split_sections] - return tex.split_quantize(x_valid, split_sections, quantizers) - - -def _gather_reference(reference, split_sections: list[int], N: int): - """Concatenate the per-expert host-split outputs into contiguous [sum_M, *] buffers.""" - live = [(r, m) for r, m in zip(reference, split_sections) if m > 0] - data = torch.cat([r._rowwise_data.view(torch.uint8).reshape(m, N // 2) for r, m in live], dim=0) - scale = torch.cat([r._rowwise_scale_inv.reshape(m, -1) for r, m in live], dim=0) - amax = torch.cat([r._amax_rowwise.reshape(m) for r, m in live], dim=0) - return data, scale, amax - - -def _assert_matches_reference(grouped, reference, split_sections: list[int], N: int) -> None: - """Compare the grouped output's raw device buffers (valid region) to the reference. - - Reads raw buffers rather than ``split_into_quantized_tensors`` because the wrapper's - per-expert shapes are frozen at capture time while the device buffers are recomputed. - """ - sum_m = sum(split_sections) - ref_data, ref_scale, ref_amax = _gather_reference(reference, split_sections, N) - - capacity = int(grouped.logical_shape[0]) - scale_stride = grouped.scale_inv.numel() // capacity - g_data = grouped.rowwise_data.view(torch.uint8).reshape(capacity, N // 2)[:sum_m] - g_scale = grouped.scale_inv.view(torch.uint8).reshape(capacity, scale_stride)[:sum_m] - g_amax = grouped.amax.reshape(-1)[:sum_m] - - torch.testing.assert_close(g_data, ref_data, atol=0, rtol=0) - torch.testing.assert_close(g_scale[:, : ref_scale.shape[1]], ref_scale, atol=0, rtol=0) - torch.testing.assert_close(g_amax, ref_amax, atol=0, rtol=0) - - -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.parametrize( - "split_sections, N", - [ - ([128], 128), - ([128, 256, 128], 256), - ([384, 128, 256], 128), - ([256, 128], 512), - ], -) -def test_matches_host_split(split_sections: list[int], N: int) -> None: - """Graph-safe amax + cast is byte-identical to the host-split row-scaled cast.""" - sum_m = sum(split_sections) - torch.manual_seed(0) - x = torch.randn((sum_m, N), dtype=torch.bfloat16, device="cuda") - tokens_per_expert = torch.tensor(split_sections, dtype=torch.int64, device="cuda") - - grouped = nvfp4_group_row_scaled_cast(x, tokens_per_expert) - reference = _host_split_reference(x, split_sections) - - _assert_matches_reference(grouped, reference, split_sections, N) - - -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.parametrize( - "split_sections, N, capacity", - [ - ([128, 128], 256, 512), - ([256, 128, 128], 128, 768), - ], -) -def test_paged_stashing_capacity_noop(split_sections: list[int], N: int, capacity: int) -> None: - """Capacity rows beyond the live token sum are a device-side no-op.""" - sum_m = sum(split_sections) - assert capacity > sum_m - torch.manual_seed(1) - x = torch.randn((capacity, N), dtype=torch.bfloat16, device="cuda") - x_valid = x[:sum_m, :].clone() - tokens_per_expert = torch.tensor(split_sections, dtype=torch.int64, device="cuda") - - grouped = nvfp4_group_row_scaled_cast(x, tokens_per_expert) - reference = _host_split_reference(x_valid, split_sections) - - _assert_matches_reference(grouped, reference, split_sections, N) - - -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -def test_capture_once_replay_changed_routing() -> None: - """Capture the amax + cast once, then replay it with redistributed routing.""" - N = 256 - num_experts = 3 - capacity = 768 - splits_a = [256, 128, 128] - splits_b = [128, 256, 128] - - tokens_per_expert = torch.empty(num_experts, dtype=torch.int64, device="cuda") - x = torch.empty((capacity, N), dtype=torch.bfloat16, device="cuda") - - torch.manual_seed(2) - x_a = torch.randn((sum(splits_a), N), dtype=torch.bfloat16, device="cuda") - x_b = torch.randn((sum(splits_b), N), dtype=torch.bfloat16, device="cuda") - - def program_routing(splits: list[int], x_src: torch.Tensor) -> None: - tokens_per_expert.copy_(torch.tensor(splits, dtype=torch.int64, device="cuda")) - x.zero_() - x[: x_src.shape[0], :].copy_(x_src) - - # Warm up outside capture so one-time setup does not land in the captured graph. - program_routing(splits_a, x_a) - warmup_stream = torch.cuda.Stream() - warmup_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup_stream): - _ = nvfp4_group_row_scaled_cast(x, tokens_per_expert) - torch.cuda.current_stream().wait_stream(warmup_stream) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - grouped = nvfp4_group_row_scaled_cast(x, tokens_per_expert) - - program_routing(splits_a, x_a) - graph.replay() - torch.cuda.synchronize() - _assert_matches_reference(grouped, _host_split_reference(x_a, splits_a), splits_a, N) - - program_routing(splits_b, x_b) - graph.replay() - torch.cuda.synchronize() - _assert_matches_reference(grouped, _host_split_reference(x_b, splits_b), splits_b, N) - - -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.parametrize( - "split_sections, N", - [ - ([128, 128], 64), - ([128] * 65, 128), - ], -) -def test_split_quantize_falls_back_to_per_expert(split_sections: list[int], N: int) -> None: - """Shapes the grouped kernel rejects (last dim not 128-aligned, >64 experts) cast per-expert.""" - sum_m = sum(split_sections) - torch.manual_seed(3) - x = torch.randn((sum_m, N), dtype=torch.bfloat16, device="cuda") - - outputs = tex.split_quantize( - x, split_sections, [_row_scaled_quantizer() for _ in split_sections] - ) - - assert len(outputs) == len(split_sections) - for out, m in zip(outputs, split_sections): - assert tuple(out.shape) == (m, N) - - -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.parametrize( - "input_shape, split_sections", - [ - ((256, 2, 128), [128, 128]), # higher-rank: not a 2D [rows, last_dim] input - ((512, 128), [128, 128]), # partial: splits cover 256 of 512 rows - ], -) -def test_split_quantize_non_grouped_shapes_fall_back( - input_shape, split_sections: list[int] -) -> None: - """Higher-rank inputs and partial splits bypass the 2D grouped kernel and cast per-expert, - matching an individual quantize of each split region.""" - torch.manual_seed(6) - x = torch.randn(input_shape, dtype=torch.bfloat16, device="cuda") - - outputs = tex.split_quantize( - x, split_sections, [_row_scaled_quantizer() for _ in split_sections] - ) - - assert len(outputs) == len(split_sections) - offset = 0 - for out, m in zip(outputs, split_sections): - ref = _row_scaled_quantizer()(x[offset : offset + m]) - torch.testing.assert_close( - out._rowwise_data.reshape(-1).view(torch.uint8), - ref._rowwise_data.reshape(-1).view(torch.uint8), - atol=0, - rtol=0, - ) - offset += m - - -def _tensor_scaled_quantizer() -> NVFP4Quantizer: - return NVFP4Quantizer( - fp4_dtype=te.DType.kFloat4E2M1, - rowwise=True, - columnwise=False, - with_rht=False, - with_post_rht_amax=False, - row_scaled_nvfp4=False, - ) - - -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -def test_split_quantize_mixed_row_scaled_falls_back() -> None: - """A split_quantize that mixes tensor-scaled and row-scaled NVFP4 quantizers casts each split - per-expert rather than taking the grouped row-scaled path (whose bulk amax, sized from the - first quantizer, would be overrun by the row-scaled kernel).""" - split_sections = [128, 128] - N = 128 - torch.manual_seed(7) - x = torch.randn((sum(split_sections), N), dtype=torch.bfloat16, device="cuda") - - quantizers = [_tensor_scaled_quantizer(), _row_scaled_quantizer()] - outputs = tex.split_quantize(x, split_sections, quantizers) - - assert len(outputs) == len(split_sections) - offset = 0 - for out, m, make_q in zip( - outputs, split_sections, [_tensor_scaled_quantizer, _row_scaled_quantizer] - ): - ref = make_q()(x[offset : offset + m]) - torch.testing.assert_close( - out._rowwise_data.reshape(-1).view(torch.uint8), - ref._rowwise_data.reshape(-1).view(torch.uint8), - atol=0, - rtol=0, - ) - offset += m - - -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -def test_split_quantize_empty_routing() -> None: - """Empty row-scaled routing returns without launching grouped work.""" - split_sections = [0, 0] - x = torch.empty((0, 128), dtype=torch.bfloat16, device="cuda") - - outputs = tex.split_quantize( - x, split_sections, [_row_scaled_quantizer() for _ in split_sections] - ) - - assert len(outputs) == len(split_sections) - for out in outputs: - assert tuple(out.shape) == (0, 128) - - -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -def test_non_contiguous_input_matches_contiguous() -> None: - """A non-contiguous input is made contiguous before the graph-safe cast.""" - split_sections = [128, 128] - N = 256 - sum_m = sum(split_sections) - torch.manual_seed(4) - x_t = torch.randn((N, sum_m), dtype=torch.bfloat16, device="cuda").t() - assert not x_t.is_contiguous() - tokens_per_expert = torch.tensor(split_sections, dtype=torch.int64, device="cuda") - - grouped = nvfp4_group_row_scaled_cast(x_t, tokens_per_expert) - reference = _host_split_reference(x_t.contiguous(), split_sections) - - _assert_matches_reference(grouped, reference, split_sections, N) - - -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -def test_strided_first_dims_matches_contiguous() -> None: - """A strided first_dims view is made contiguous before the cast, so it is not misread.""" - split_sections = [256, 128] - N = 128 - sum_m = sum(split_sections) - torch.manual_seed(5) - x = torch.randn((sum_m, N), dtype=torch.bfloat16, device="cuda") - - # [256, 0, 128, 0][::2] == [256, 128] but strided; a raw data_ptr read would see [256, 0]. - strided = torch.tensor([256, 0, 128, 0], dtype=torch.int64, device="cuda")[::2] - assert not strided.is_contiguous() - assert strided.tolist() == split_sections - - grouped = nvfp4_group_row_scaled_cast(x, strided) - reference = _host_split_reference(x, split_sections) - - _assert_matches_reference(grouped, reference, split_sections, N) - - -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -def test_rejects_unsupported_quantizer_settings() -> None: - """with_2d_quantization is rejected on both the host-split and graph-safe row-scaled paths.""" - - def _quantizer_with_2d() -> NVFP4Quantizer: - return NVFP4Quantizer( - fp4_dtype=te.DType.kFloat4E2M1, - rowwise=True, - columnwise=False, - row_scaled_nvfp4=True, - with_2d_quantization=True, - ) - - split_sections = [128, 128] - N = 128 - x = torch.randn((sum(split_sections), N), dtype=torch.bfloat16, device="cuda") - - with pytest.raises(RuntimeError): - tex.split_quantize(x, split_sections, [_quantizer_with_2d() for _ in split_sections]) - - tokens_per_expert = torch.tensor(split_sections, dtype=torch.int64, device="cuda") - tensor_offsets = tex.splits_to_offsets(tokens_per_expert, N) - with pytest.raises(RuntimeError): - tex.nvfp4_group_row_scaled_cast_graph_safe( - x, _quantizer_with_2d(), len(split_sections), tokens_per_expert, tensor_offsets - ) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py new file mode 100644 index 00000000000..567065a9a73 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py @@ -0,0 +1,272 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Grouped row-scaled NVFP4 quantization through tex.group_quantize.""" + +from typing import Optional + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.pytorch import NVFP4Quantizer + +recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) + + +def _row_scaled_quantizer(**kwargs) -> NVFP4Quantizer: + return NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=False, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=True, + **kwargs, + ) + + +def _tensor_scaled_quantizer() -> NVFP4Quantizer: + return NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=False, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=False, + ) + + +def _group_quantize(x: torch.Tensor, num_tensors: int, first_dims: Optional[torch.Tensor]): + return tex.group_quantize(x, _row_scaled_quantizer(), num_tensors, first_dims) + + +def _per_expert_reference(x_valid: torch.Tensor, split_sections: list[int]): + N = x_valid.shape[1] + data, scale, amax = [], [], [] + offset = 0 + for m in split_sections: + if m == 0: + continue + ref = _row_scaled_quantizer()(x_valid[offset : offset + m]) + data.append(ref._rowwise_data.view(torch.uint8).reshape(m, N // 2)) + scale.append(ref._rowwise_scale_inv.view(torch.uint8).reshape(m, -1)[:, : N // 16]) + amax.append(ref._amax_rowwise.reshape(m)) + offset += m + return torch.cat(data, dim=0), torch.cat(scale, dim=0), torch.cat(amax, dim=0) + + +def _assert_matches_reference(grouped, x_valid: torch.Tensor, split_sections: list[int]) -> None: + N = x_valid.shape[1] + sum_m = sum(split_sections) + ref_data, ref_scale, ref_amax = _per_expert_reference(x_valid, split_sections) + + capacity = int(grouped.logical_shape[0]) + scale_stride = grouped.scale_inv.numel() // capacity + g_data = grouped.rowwise_data.view(torch.uint8).reshape(capacity, N // 2)[:sum_m] + g_scale = grouped.scale_inv.view(torch.uint8).reshape(capacity, scale_stride)[:sum_m] + g_amax = grouped.amax.reshape(-1)[:sum_m] + + torch.testing.assert_close(g_amax, ref_amax, atol=0, rtol=0) + torch.testing.assert_close(g_data, ref_data, atol=0, rtol=0) + torch.testing.assert_close(g_scale[:, : N // 16], ref_scale, atol=0, rtol=0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "split_sections, N", + [ + ([128], 128), + ([128, 256, 128], 256), + ([384, 128, 256], 128), + ([256, 128], 512), + ], +) +def test_matches_per_expert_quantize(split_sections: list[int], N: int) -> None: + torch.manual_seed(0) + x = torch.randn((sum(split_sections), N), dtype=torch.bfloat16, device="cuda") + first_dims = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + grouped = _group_quantize(x, len(split_sections), first_dims) + + assert grouped.row_scaled_nvfp4 + _assert_matches_reference(grouped, x, split_sections) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("num_tensors, m, N", [(1, 128, 128), (4, 256, 256)]) +def test_uniform_shapes(num_tensors: int, m: int, N: int) -> None: + torch.manual_seed(1) + x = torch.randn((num_tensors * m, N), dtype=torch.bfloat16, device="cuda") + + grouped = _group_quantize(x, num_tensors, None) + + _assert_matches_reference(grouped, x, [m] * num_tensors) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "split_sections, N, capacity", + [ + ([128, 128], 256, 512), + ([256, 128, 128], 128, 768), + ], +) +def test_paged_stashing_capacity_noop(split_sections: list[int], N: int, capacity: int) -> None: + sum_m = sum(split_sections) + assert capacity > sum_m + torch.manual_seed(2) + x = torch.randn((capacity, N), dtype=torch.bfloat16, device="cuda") + first_dims = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + grouped = _group_quantize(x, len(split_sections), first_dims) + + _assert_matches_reference(grouped, x[:sum_m].clone(), split_sections) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_capture_once_replay_changed_routing() -> None: + N = 256 + num_experts = 3 + capacity = 768 + splits_a = [256, 128, 128] + splits_b = [128, 256, 128] + + first_dims = torch.empty(num_experts, dtype=torch.int64, device="cuda") + x = torch.empty((capacity, N), dtype=torch.bfloat16, device="cuda") + + torch.manual_seed(3) + x_a = torch.randn((sum(splits_a), N), dtype=torch.bfloat16, device="cuda") + x_b = torch.randn((sum(splits_b), N), dtype=torch.bfloat16, device="cuda") + + def program_routing(splits: list[int], x_src: torch.Tensor) -> None: + first_dims.copy_(torch.tensor(splits, dtype=torch.int64, device="cuda")) + x.zero_() + x[: x_src.shape[0], :].copy_(x_src) + + # Warm up outside capture so one-time setup does not land in the captured graph. + program_routing(splits_a, x_a) + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + _ = _group_quantize(x, num_experts, first_dims) + torch.cuda.current_stream().wait_stream(warmup_stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + grouped = _group_quantize(x, num_experts, first_dims) + + program_routing(splits_a, x_a) + graph.replay() + torch.cuda.synchronize() + _assert_matches_reference(grouped, x_a, splits_a) + + program_routing(splits_b, x_b) + graph.replay() + torch.cuda.synchronize() + _assert_matches_reference(grouped, x_b, splits_b) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_non_contiguous_input_matches_contiguous() -> None: + split_sections = [128, 128] + N = 256 + torch.manual_seed(4) + x_t = torch.randn((N, sum(split_sections)), dtype=torch.bfloat16, device="cuda").t() + assert not x_t.is_contiguous() + first_dims = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + grouped = _group_quantize(x_t, len(split_sections), first_dims) + + _assert_matches_reference(grouped, x_t.contiguous(), split_sections) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_rejects_strided_first_dims() -> None: + x = torch.randn((384, 128), dtype=torch.bfloat16, device="cuda") + strided = torch.tensor([256, 0, 128, 0], dtype=torch.int64, device="cuda")[::2] + assert not strided.is_contiguous() + + with pytest.raises(RuntimeError): + _group_quantize(x, 2, strided) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_rejects_unsupported_quantizer_settings() -> None: + x = torch.randn((256, 128), dtype=torch.bfloat16, device="cuda") + first_dims = torch.tensor([128, 128], dtype=torch.int64, device="cuda") + + with pytest.raises(RuntimeError): + tex.group_quantize(x, _row_scaled_quantizer(with_2d_quantization=True), 2, first_dims) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "input_shape, split_sections", + [ + ((384, 128), [256, 128]), + ((256, 64), [128, 128]), + ((256, 2, 128), [128, 128]), + ((512, 128), [128, 128]), + ], +) +def test_split_quantize_matches_per_expert(input_shape, split_sections: list[int]) -> None: + torch.manual_seed(5) + x = torch.randn(input_shape, dtype=torch.bfloat16, device="cuda") + + outputs = tex.split_quantize( + x, split_sections, [_row_scaled_quantizer() for _ in split_sections] + ) + + assert len(outputs) == len(split_sections) + offset = 0 + for out, m in zip(outputs, split_sections): + ref = _row_scaled_quantizer()(x[offset : offset + m]) + torch.testing.assert_close( + out._rowwise_data.reshape(-1).view(torch.uint8), + ref._rowwise_data.reshape(-1).view(torch.uint8), + atol=0, + rtol=0, + ) + offset += m + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_split_quantize_mixed_row_scaled_falls_back() -> None: + split_sections = [128, 128] + N = 128 + torch.manual_seed(6) + x = torch.randn((sum(split_sections), N), dtype=torch.bfloat16, device="cuda") + + quantizers = [_tensor_scaled_quantizer(), _row_scaled_quantizer()] + outputs = tex.split_quantize(x, split_sections, quantizers) + + assert len(outputs) == len(split_sections) + offset = 0 + for out, m, make_q in zip( + outputs, split_sections, [_tensor_scaled_quantizer, _row_scaled_quantizer] + ): + ref = make_q()(x[offset : offset + m]) + torch.testing.assert_close( + out._rowwise_data.reshape(-1).view(torch.uint8), + ref._rowwise_data.reshape(-1).view(torch.uint8), + atol=0, + rtol=0, + ) + offset += m + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_split_quantize_empty_routing() -> None: + split_sections = [0, 0] + x = torch.empty((0, 128), dtype=torch.bfloat16, device="cuda") + + outputs = tex.split_quantize( + x, split_sections, [_row_scaled_quantizer() for _ in split_sections] + ) + + assert len(outputs) == len(split_sections) + for out in outputs: + assert tuple(out.shape) == (0, 128) diff --git a/transformer_engine/common/cast/cast_grouped.cu b/transformer_engine/common/cast/cast_grouped.cu index c34fc1acb1b..49ab39f7e79 100644 --- a/transformer_engine/common/cast/cast_grouped.cu +++ b/transformer_engine/common/cast/cast_grouped.cu @@ -114,96 +114,3 @@ void nvte_group_nvfp4_compute_amax(const NVTETensor input, NVTETensor *outputs, dispatch::nvfp4::group_compute_fused_amax(*input_tensor, /*noop=*/nullptr, output_tensors, split_sections, num_tensors, stream); } - -void nvte_group_nvfp4_row_scaled_cast_with_amax(const NVTETensor input, NVTETensor *outputs, - const size_t *split_sections, size_t num_tensors, - const NVTEQuantizationConfig quant_config, - cudaStream_t stream) { - NVTE_API_CALL(nvte_group_nvfp4_row_scaled_cast_with_amax); - using namespace transformer_engine; - - const Tensor *input_tensor = convertNVTETensorCheck(input); - std::vector output_tensors; - output_tensors.reserve(num_tensors); - for (size_t i = 0; i < num_tensors; ++i) { - output_tensors.push_back(convertNVTETensorCheck(outputs[i])); - } - - QuantizationConfig quant_config_cpp; - if (quant_config != nullptr) { - quant_config_cpp = *reinterpret_cast(quant_config); - } - - Tensor dummy_tensor; - Tensor *noop_tensor = &dummy_tensor; - if (quant_config_cpp.noop_tensor != nullptr) { - noop_tensor = convertNVTETensorCheck(quant_config_cpp.noop_tensor); - } - - // The scale dtype is shared across experts; pick it from the first non-empty output. - DType scale_dtype = DType::kFloat8E4M3; - for (size_t i = 0; i < num_tensors; ++i) { - if (split_sections[i] == 0) continue; - const Tensor *o = output_tensors[i]; - if (o->scale_inv.has_data()) { - scale_dtype = o->scale_inv.dtype; - } else if (o->columnwise_scale_inv.has_data()) { - scale_dtype = o->columnwise_scale_inv.dtype; - } - break; - } - - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( - scale_dtype, ScaleType, - dispatch::nvfp4::group_row_scaled_cast(*input_tensor, noop_tensor, output_tensors, - split_sections, num_tensors, - &quant_config_cpp, stream);); -} - -void nvte_group_nvfp4_compute_amax_graph_safe(const NVTEGroupedTensor input, - NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_nvfp4_compute_amax_graph_safe); - using namespace transformer_engine; - - GroupedTensor *input_tensor = convertNVTEGroupedTensorCheck(input); - GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output); - - if (input_tensor->num_tensors == 0) { - return; - } - - dispatch::nvfp4::group_compute_fused_amax_graph_safe(input_tensor, /*noop=*/nullptr, - output_tensor, stream); -} - -void nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(const NVTEGroupedTensor input, - NVTEGroupedTensor output, - cudaStream_t stream) { - NVTE_API_CALL(nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe); - using namespace transformer_engine; - - GroupedTensor *input_tensor = convertNVTEGroupedTensorCheck(input); - GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output); - - if (input_tensor->num_tensors == 0) { - return; - } - - QuantizationConfig quant_config_cpp; - - Tensor dummy_noop; - Tensor *noop_tensor = &dummy_noop; - - // The scale dtype is shared across experts; pick it from the output buffers. - DType scale_dtype = DType::kFloat8E4M3; - if (output_tensor->scale_inv.has_data()) { - scale_dtype = output_tensor->scale_inv.dtype; - } else if (output_tensor->columnwise_scale_inv.has_data()) { - scale_dtype = output_tensor->columnwise_scale_inv.dtype; - } - - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( - scale_dtype, ScaleType, - dispatch::nvfp4::group_row_scaled_cast_graph_safe( - input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream);); -} diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index aa1d16fd894..105cd2b8d1c 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -560,6 +560,30 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor quant_config_cpp.force_pow_2_scales, stream); break; } + case NVTE_NVFP4_1D_SCALING: { + NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_NVFP4_1D_SCALING."); + NVTE_CHECK(output_tensor->row_scaled_nvfp4, + "Grouped NVFP4 quantization is only implemented for row-scaled NVFP4."); + NVTE_CHECK(!quant_config_cpp.stochastic_rounding, + "Grouped row-scaled NVFP4 quantization does not support stochastic rounding."); + NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, + "Grouped row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK(quant_config_cpp.nvfp4_4over6_mode == kNVTENVFP44Over6Disabled, + "Grouped row-scaled NVFP4 quantization does not support 4over6."); + if (input_tensor->num_tensors == 0) { + break; + } + nvfp4::group_compute_fused_amax_graph_safe(input_tensor, noop_tensor, output_tensor, stream); + const DType scale_dtype = output_tensor->scale_inv.has_data() + ? output_tensor->scale_inv.dtype + : output_tensor->columnwise_scale_inv.dtype; + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + nvfp4::group_row_scaled_cast_graph_safe(input_tensor, noop_tensor, + output_tensor, &quant_config_cpp, + stream);); + break; + } default: NVTE_ERROR("Not implemented scaling mode: " + to_string(scaling_mode) + "."); } 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 1850dd49a0b..c7e3e0f7d4e 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -1381,35 +1381,6 @@ using namespace ptx; #if FP4_TYPE_SUPPORTED -constexpr int kMaxCastTensorsPerKernel = 64; - -// Per-expert output / amax pointers. The packed (sum_M, K) input is covered by a -// single TMA map; data_stride_row = cols/2 and data_stride_col = M_i/2 are -// derived in-kernel from split_sections_range. -struct GroupedRowScaledCastArgs { - void *q_row_list[kMaxCastTensorsPerKernel]; // uint8* [M_i, cols/2] - void *s_dec_row_list[kMaxCastTensorsPerKernel]; // e4m3* [M_i, scale_stride_row] - void *q_col_list[kMaxCastTensorsPerKernel]; // uint8* [cols, M_i/2] - void *s_dec_col_list[kMaxCastTensorsPerKernel]; // e4m3* [cols, scale_stride_col] - void *row_amax_list[kMaxCastTensorsPerKernel]; // float* [M_i] - void *col_amax_list[kMaxCastTensorsPerKernel]; // float* [cols] - int scale_stride_col[kMaxCastTensorsPerKernel]; // per-expert colwise SF row stride - int split_sections_range[kMaxCastTensorsPerKernel + 1]; // prefix sum of M_i - int num_tensors; -}; - -__device__ __forceinline__ int GetTensorId(const GroupedRowScaledCastArgs &args, int global_row) { - int tensor_id = 0; - while (args.split_sections_range[tensor_id + 1] <= global_row) { - ++tensor_id; - } - return tensor_id; -} - -// Per-expert encode + scatter body shared by the host-split and graph-safe -// grouped row-scaled cast kernels. The two kernels differ only in how they -// resolve the per-expert bases (host-baked args vs device metadata); everything -// from the SMEM layout onward is identical. 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, @@ -1618,52 +1589,7 @@ __device__ __forceinline__ void group_row_scaled_cast_body( #endif // __CUDA_ARCH__ >= 1000 } -// Host-split kernel: per-expert bases come from the host-baked GroupedRowScaledCastArgs. -template -__global__ void __launch_bounds__(THREADS_NUM) - group_row_scaled_cast_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, - const __grid_constant__ GroupedRowScaledCastArgs args, - const float *noop, const size_t rows, const size_t cols, - const size_t scale_stride_row) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - if (noop != nullptr && noop[0] == 1.0f) { - return; - } - - const int ctaid_X = blockIdx.x; - const int ctaid_Y = blockIdx.y; - const int block_offset_Y = ctaid_Y * TunableConfig::CHUNK_DIM_Y; - - // A 128-row chunk lies fully inside one expert (every split is 128-aligned). - const int tensor_id = GetTensorId(args, block_offset_Y); - const int split_start = args.split_sections_range[tensor_id]; - const int M_t = args.split_sections_range[tensor_id + 1] - split_start; - const int local_row_base = block_offset_Y - split_start; - - uint8_t *const q_row_base = reinterpret_cast(args.q_row_list[tensor_id]); - ScaleType *const s_dec_row_base = reinterpret_cast(args.s_dec_row_list[tensor_id]); - const float *const row_amax_base = reinterpret_cast(args.row_amax_list[tensor_id]); - - uint8_t *const q_col_base = - RETURN_TRANSPOSE ? reinterpret_cast(args.q_col_list[tensor_id]) : nullptr; - ScaleType *const s_dec_col_base = - RETURN_TRANSPOSE ? reinterpret_cast(args.s_dec_col_list[tensor_id]) : nullptr; - const float *const col_amax_base = - RETURN_TRANSPOSE ? reinterpret_cast(args.col_amax_list[tensor_id]) : nullptr; - const size_t scale_stride_col = - RETURN_TRANSPOSE ? static_cast(args.scale_stride_col[tensor_id]) : 0; - - 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); -#else - NVTE_DEVICE_ERROR("Grouped row-scaled NVFP4 cast kernel requires SM 10.0+ (Blackwell)."); -#endif // __CUDA_ARCH__ >= 1000 -} - -// Graph-safe kernel: per-expert bases derived on device from offsets / first_dims; -// grid covers the input capacity, CTAs past the live row sum noop. +// Grid covers the input capacity; CTAs past the live row sum noop. 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, @@ -1717,129 +1643,6 @@ __global__ void __launch_bounds__(THREADS_NUM) group_row_scaled_cast_nvfp4_graph } // namespace group_row_scaled_cast_kernel -// Host launcher for the grouped row-scaled NVFP4 cast. BF16-only; consumes the -// per-expert amax computed by nvte_group_nvfp4_compute_amax. -template -inline void group_row_scaled_cast(const Tensor &input, const Tensor *noop, - std::vector &output_list, const size_t *split_sections, - size_t num_tensors, const QuantizationConfig *quant_config, - cudaStream_t stream) { -#if FP4_TYPE_SUPPORTED - using namespace group_row_scaled_cast_kernel; - using namespace quantize_transpose_tuned_kernel; - - NVTE_CHECK(num_tensors <= static_cast(kMaxCastTensorsPerKernel), - "Number of tensors should be <= ", kMaxCastTensorsPerKernel); - NVTE_CHECK(input.dtype() == DType::kBFloat16, - "Grouped row-scaled NVFP4 cast requires BF16 input."); - - const bool use_fast_math = quant_config ? quant_config->use_fast_math : false; - - checkCuDriverContext(stream); - CheckNoopTensor(*noop, "cast_noop"); - CheckInputTensor(input, "input"); - NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); - - const auto [rows, cols] = input.flat_2d_dims(); - NVTE_CHECK(cols % TunableConfig::CHUNK_DIM_X == 0, - "Grouped row-scaled NVFP4 cast requires cols to be a multiple of ", - TunableConfig::CHUNK_DIM_X, " (got ", cols, ")."); - - // First non-empty output defines the shared layout (transpose + rowwise SF stride). - Tensor *first_output = nullptr; - for (size_t i = 0; i < num_tensors; ++i) { - if (split_sections[i] != 0 && output_list[i]->has_data()) { - first_output = output_list[i]; - break; - } - } - NVTE_CHECK(first_output != nullptr, "No non-empty output tensor found."); - const bool return_transpose = first_output->has_columnwise_data(); - const size_t scale_stride_row = first_output->scale_inv.shape[1]; - - GroupedRowScaledCastArgs args; - args.num_tensors = 0; - args.split_sections_range[0] = 0; - for (size_t i = 0; i < num_tensors; ++i) { - if (split_sections[i] == 0) continue; - NVTE_CHECK(split_sections[i] % TunableConfig::CHUNK_DIM_Y == 0, - "Grouped row-scaled NVFP4 cast requires each split to be a multiple of ", - TunableConfig::CHUNK_DIM_Y, " (split ", i, " = ", split_sections[i], ")."); - Tensor *o = output_list[i]; - NVTE_CHECK(o->amax.dptr != nullptr, "Grouped row-scaled NVFP4 cast requires rowwise amax."); - const int idx = args.num_tensors; - args.q_row_list[idx] = o->data.dptr; - args.s_dec_row_list[idx] = o->scale_inv.dptr; - args.row_amax_list[idx] = o->amax.dptr; - if (return_transpose) { - NVTE_CHECK(o->has_columnwise_data() && o->columnwise_scale_inv.dptr != nullptr && - o->columnwise_amax.dptr != nullptr, - "Grouped row-scaled NVFP4 transpose cast requires columnwise data/scale/amax for " - "every expert."); - args.q_col_list[idx] = o->columnwise_data.dptr; - args.s_dec_col_list[idx] = o->columnwise_scale_inv.dptr; - args.col_amax_list[idx] = o->columnwise_amax.dptr; - args.scale_stride_col[idx] = static_cast(o->columnwise_scale_inv.shape[1]); - } else { - args.q_col_list[idx] = nullptr; - args.s_dec_col_list[idx] = nullptr; - args.col_amax_list[idx] = nullptr; - args.scale_stride_col[idx] = 0; - } - args.split_sections_range[idx + 1] = - args.split_sections_range[idx] + static_cast(split_sections[i]); - args.num_tensors++; - } - NVTE_CHECK(args.split_sections_range[args.num_tensors] == static_cast(rows), - "Sum of split sections must equal the number of input rows."); - - const float *noop_ptr = reinterpret_cast(noop->data.dptr); - - alignas(64) CUtensorMap tensor_map_input{}; - create_2D_tensor_map(tensor_map_input, input.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, cols, 0, - sizeof(IType) * 8); - - constexpr int buff_elems = BUFF_DIM_Y * BUFF_DIM_X; - constexpr int buff_elems_total_in = BUFFS_NUM_IN * buff_elems; - constexpr int buff_size_aligned_in = - DIVUP_TO_MULTIPLE(buff_elems_total_in * sizeof(IType), TMA_SHMEM_ALIGNMENT); - constexpr int buff_size_aligned_out = - DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT * BUFF_OUT_SIZE, TMA_SHMEM_ALIGNMENT); - constexpr int buff_size_aligned_out_t = - DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT_TR * BUFF_OUT_TR_SIZE, TMA_SHMEM_ALIGNMENT); - constexpr int buff_size_scales = DIVUP_TO_MULTIPLE( - TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); - constexpr int buff_size_scales_transpose = DIVUP_TO_MULTIPLE( - TunableConfig::CHUNK_DIM_X * SCALES_PER_CHUNK_Y * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); - - const dim3 grid(static_cast(cols / TunableConfig::CHUNK_DIM_X), - static_cast(rows / TunableConfig::CHUNK_DIM_Y), 1); - const dim3 block(THREADS_NUM, 1, 1); - - 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_kernel; - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - kernel<<>>(tensor_map_input, args, noop_ptr, rows, 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); -#endif // FP4_TYPE_SUPPORTED -} - -// Graph-safe counterpart of group_row_scaled_cast: routing comes from device -// output->first_dims / tensor_offsets. Grid covers the input capacity; CTAs past -// the live row sum noop. BF16-only; consumes the output's per-expert amax. template inline void group_row_scaled_cast_graph_safe(const GroupedTensor *input, const Tensor *noop, GroupedTensor *output, diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 4b64d91b932..f95b6740b09 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -545,6 +545,9 @@ struct GroupedTensor { */ bool with_gemm_swizzled_scales = false; + /*! \brief Whether NVFP4 data uses row scaling instead of tensor scaling */ + bool row_scaled_nvfp4 = false; + /*! Map from NVTEGroupedTensorParam to parameter sizes */ static constexpr size_t attr_sizes[] = { sizeof(NVTEBasicTensor), // kNVTEGroupedRowwiseData @@ -557,7 +560,8 @@ struct GroupedTensor { sizeof(NVTEBasicTensor), // kNVTEGroupedFirstDims sizeof(NVTEBasicTensor), // kNVTEGroupedLastDims sizeof(NVTEBasicTensor), // kNVTEGroupedTensorOffsets - sizeof(uint8_t) // kNVTEGroupedWithGEMMSwizzledScales + sizeof(uint8_t), // kNVTEGroupedWithGEMMSwizzledScales + sizeof(uint8_t) // kNVTEGroupedRowScaledNVFP4 }; GroupedTensor(NVTEScalingMode scaling_mode, size_t num_tensors) @@ -575,7 +579,8 @@ struct GroupedTensor { tensor_offsets(nullptr, std::vector{0}, DType::kInt64), logical_shape(nvte_make_shape(nullptr, 1)), nvte_tensor(0), - with_gemm_swizzled_scales(false) {} + with_gemm_swizzled_scales(false), + row_scaled_nvfp4(false) {} explicit operator NVTEGroupedTensor() const noexcept { return nvte_tensor; } @@ -628,6 +633,7 @@ struct GroupedTensor { scaling_mode = NVTE_DELAYED_TENSOR_SCALING; nvte_tensor = 0; with_gemm_swizzled_scales = false; + row_scaled_nvfp4 = false; } }; diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 17b6b7fff37..774146dea6b 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -93,6 +93,7 @@ void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t strea * The type of quantized tensor in the output depends on the scaling mode of the output * tensor. See file level comments. * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. + * NVFP4 output requires row scaling, BF16 input, and dimensions that are multiples of 128. * * \param[in] input Input grouped tensor to be cast. * \param[in,out] output Output grouped tensor. @@ -514,74 +515,6 @@ void nvte_group_nvfp4_compute_amax(const NVTETensor input, NVTETensor *outputs, const size_t *split_sections, size_t num_tensors, cudaStream_t stream); -/*! \brief Casts a grouped row-scaled NVFP4 input to FP4 in a single kernel launch. - * - * Single-launch counterpart of the per-expert row-scaled NVFP4 cast. The input - * is one packed (sum_M, K) BF16 buffer; `split_sections` gives the per-expert - * row counts (each a multiple of 128) and K must be a multiple of 128. The - * per-expert rowwise amax [M_i] and (when the columnwise/transpose output is - * allocated) columnwise amax [K] must already be populated, e.g. by - * nvte_group_nvfp4_compute_amax. For expert i the rowwise FP4 data and E4M3 - * scales are written to `outputs[i]`, and the transpose counterpart to - * `outputs[i]`'s columnwise buffers when present. The output is byte-identical - * to quantizing each expert independently. - * - * \param[in] input Packed grouped input tensor (BF16). - * \param[in,out] outputs Per-expert output tensors (amax pre-populated). - * \param[in] split_sections Per-expert row counts (each a multiple of 128). - * \param[in] num_tensors Number of output tensors. - * \param[in] quant_config (Optional) Quantization configurations. - * \param[in] stream CUDA stream used for the operation. - */ -void nvte_group_nvfp4_row_scaled_cast_with_amax(const NVTETensor input, NVTETensor *outputs, - const size_t *split_sections, size_t num_tensors, - const NVTEQuantizationConfig quant_config, - cudaStream_t stream); - -/*! \brief CUDA-graph-safe grouped row-scaled NVFP4 fused amax. - * - * Device-metadata counterpart of nvte_group_nvfp4_compute_amax. Instead of host - * `split_sections`, the per-expert offsets are read from the output grouped - * tensor's device `tensor_offsets`, so a dynamic-routing MoE amax can be - * captured once and replayed with different routing. - * - * The input is one packed (capacity, K) BF16 buffer. The rowwise amax [sum_M] - * is written to the contiguous `output->amax` and the columnwise amax - * [num_tensors * K] to `output->columnwise_amax`, whichever are allocated. Only - * the const-last-dimension layout (VARYING_FIRST_DIM: common K) is supported; - * K must be a multiple of 128 and each per-expert row count a multiple of 128. - * - * \param[in] input Packed grouped input tensor (BF16) with device metadata. - * \param[in,out] output Grouped output tensor receiving the amax vectors. - * \param[in] stream CUDA stream used for the operation. - */ -void nvte_group_nvfp4_compute_amax_graph_safe(const NVTEGroupedTensor input, - NVTEGroupedTensor output, cudaStream_t stream); - -/*! \brief CUDA-graph-safe grouped row-scaled NVFP4 cast. - * - * Device-metadata counterpart of nvte_group_nvfp4_row_scaled_cast_with_amax. - * Instead of host `split_sections`, the per-expert row counts and offsets are - * read from the input/output grouped tensors' device `first_dims` and - * `tensor_offsets`, so a dynamic-routing MoE cast can be captured once and - * replayed with different routing without rebaking kernel parameters. - * - * Both tensors use the single-buffer grouped layout: the input is one packed - * (capacity, K) BF16 buffer, and the output carries contiguous rowwise data / - * E4M3 scales (and, when allocated, columnwise data / E4M3 scales) plus the - * pre-populated rowwise amax [sum_M] and columnwise amax [num_tensors * K]. - * Only the const-last-dimension layout (VARYING_FIRST_DIM: common K, per-expert - * row count) is supported; each per-expert row count must be a multiple of 128 - * and K a multiple of 128. The output is byte-identical to the host-split cast. - * - * \param[in] input Packed grouped input tensor (BF16) with device metadata. - * \param[in,out] output Grouped output tensor (amax pre-populated) with device metadata. - * \param[in] stream CUDA stream used for the operation. - */ -void nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(const NVTEGroupedTensor input, - NVTEGroupedTensor output, - cudaStream_t stream); - #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 4c0e844c8c4..1adb47773ef 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -570,6 +570,8 @@ enum NVTEGroupedTensorParam { 9, /*!< Tensor offsets for contiguous layout (device pointer to int64_t array) */ kNVTEGroupedWithGEMMSwizzledScales = 10, /*!< Whether scaling factors are in format expected by GEMM */ + kNVTEGroupedRowScaledNVFP4 = + 11, /*!< Whether NVFP4 data uses row scaling instead of tensor scaling */ kNVTENumGroupedTensorParams }; @@ -1302,6 +1304,11 @@ class GroupedTensorWrapper { nvte_set_grouped_tensor_param(tensor_, kNVTEGroupedWithGEMMSwizzledScales, &val, sizeof(val)); } + void set_row_scaled_nvfp4(bool row_scaled_nvfp4) { + const auto val = static_cast(row_scaled_nvfp4); + nvte_set_grouped_tensor_param(tensor_, kNVTEGroupedRowScaledNVFP4, &val, sizeof(val)); + } + // Parameter getters NVTEBasicTensor get_parameter(const NVTEGroupedTensorParam param) const noexcept { NVTEBasicTensor ret; @@ -1348,6 +1355,12 @@ class GroupedTensorWrapper { return static_cast(val); } + bool get_row_scaled_nvfp4() const { + uint8_t val = 0; + nvte_get_grouped_tensor_param(tensor_, kNVTEGroupedRowScaledNVFP4, &val, sizeof(val), nullptr); + return static_cast(val); + } + /*! \brief Get an underlying NVTEGroupedTensor. * * \return NVTEGroupedTensor held by this GroupedTensorWrapper. diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index ad0d47ac99e..66693dca3a1 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -1330,6 +1330,9 @@ void nvte_set_grouped_tensor_param(NVTEGroupedTensor tensor, NVTEGroupedTensorPa case kNVTEGroupedWithGEMMSwizzledScales: t.with_gemm_swizzled_scales = static_cast(*reinterpret_cast(buf)); break; + case kNVTEGroupedRowScaledNVFP4: + t.row_scaled_nvfp4 = static_cast(*reinterpret_cast(buf)); + break; default: NVTE_ERROR("Unsupported grouped tensor parameter (", static_cast(param), ")"); } @@ -1425,6 +1428,9 @@ void nvte_get_grouped_tensor_param(const NVTEGroupedTensor tensor, NVTEGroupedTe case kNVTEGroupedWithGEMMSwizzledScales: *reinterpret_cast(buf) = static_cast(t->with_gemm_swizzled_scales); break; + case kNVTEGroupedRowScaledNVFP4: + *reinterpret_cast(buf) = static_cast(t->row_scaled_nvfp4); + break; default: NVTE_ERROR("Unsupported grouped tensor parameter (", static_cast(param), ")"); } diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 6ec981bf116..955763f43b1 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -420,11 +420,6 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q const at::Tensor &columnwise_amax, std::optional tensor_offsets); -py::object nvfp4_group_row_scaled_cast_graph_safe(const at::Tensor &tensor, py::handle quantizer, - const size_t num_tensors, - const at::Tensor &first_dims, - const at::Tensor &tensor_offsets); - py::object group_dequantize(const py::handle &input, DType otype); py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 43b1c8fb367..9b6e754c429 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -211,6 +211,35 @@ void group_quantize_nvfp4_impl(const GroupedTensorWrapper &grouped_input_tensor, }); } +void group_quantize_nvfp4_row_scaled_impl(const GroupedTensorWrapper &grouped_input_tensor, + GroupedTensorWrapper &grouped_output_tensor, + const NVFP4Quantizer *nvfp4_quantizer_cpp, + const std::optional &noop_flag, + cudaStream_t stream) { + NVTE_CHECK(!nvfp4_quantizer_cpp->with_rht, + "Grouped row-scaled NVFP4 quantization does not support RHT."); + NVTE_CHECK(!nvfp4_quantizer_cpp->with_2d_quantization, + "Grouped row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK(!nvfp4_quantizer_cpp->stochastic_rounding, + "Grouped row-scaled NVFP4 quantization does not support stochastic rounding."); + NVTE_CHECK(!nvfp4_quantizer_cpp->with_amax_reduction, + "Grouped row-scaled NVFP4 quantization does not support amax reduction."); + NVTE_CHECK(nvfp4_quantizer_cpp->nvfp4_4over6_mode == kNVTENVFP44Over6Disabled, + "Grouped row-scaled NVFP4 quantization does not support 4over6."); + + QuantizationConfigWrapper quant_config_cpp; + if (noop_flag.has_value()) { + quant_config_cpp.set_noop_tensor(noop_flag->data()); + } + if (transformer_engine::getenv("NVTE_USE_FAST_MATH")) { + quant_config_cpp.set_use_fast_math(true); + } + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor.data(), + quant_config_cpp, stream); + }); +} + float fp8_max_for_dtype(const DType dtype) { if (!is_fp8_dtype(dtype)) { NVTE_ERROR("Expected FP8 dtype for grouped current-scaling quantization, got ", @@ -276,14 +305,15 @@ void compute_grouped_fp8_current_scaling_amax_and_scale( } // namespace -py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, +py::object group_quantize(const at::Tensor &input, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, std::optional tensor_offsets, std::optional noop_flag, const py::object &output) { using namespace transformer_engine::pytorch::detail; init_extension(); - NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D"); + NVTE_CHECK(input.dim() == 2, "Tensor must be 2D"); + const at::Tensor tensor = input.contiguous(); // No-op flag for CUDA graph weight caching (skip_fp8_weight_update). std::optional noop_flag_cpp; @@ -302,6 +332,26 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const auto quantizer_cpp = convert_quantizer(quantizer); + const bool row_scaled_nvfp4 = + detail::IsNVFP4Quantizers(quantizer.ptr()) && + static_cast(quantizer_cpp.get())->row_scaled_nvfp4; + if (row_scaled_nvfp4) { + 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."); + } else { + NVTE_CHECK(num_tensors > 0 && logical_first_dim % num_tensors == 0, + "group_quantize: row-scaled NVFP4 with uniform shapes requires the first " + "dimension (", + logical_first_dim, ") to be divisible by num_tensors (", num_tensors, ")."); + first_dims = at::full({static_cast(num_tensors)}, + static_cast(logical_first_dim / num_tensors), + tensor.options().dtype(at::kLong)); + } + } + // Create input GroupedTensor. auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); grouped_input_tensor.set_rowwise_data(tensor.data_ptr(), @@ -370,9 +420,15 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const NVTE_CHECK(!last_dims.has_value(), "group_quantize: varying last dim is not supported with NVFP4."); NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); - group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp, - nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream(), - !nvfp4_quantizer_cpp->disable_second_level_scale); + if (row_scaled_nvfp4) { + group_quantize_nvfp4_row_scaled_impl(grouped_input_tensor, grouped_output_tensor_cpp, + nvfp4_quantizer_cpp, noop_flag_cpp, + at::cuda::getCurrentCUDAStream()); + } else { + group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp, + nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream(), + !nvfp4_quantizer_cpp->disable_second_level_scale); + } break; } case GroupedQuantizationMode::FP8_CURRENT_SCALING_GROUPED_QUANTIZE: { @@ -635,90 +691,6 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q return py::reinterpret_borrow(grouped_output_py); } -// Graph-safe grouped row-scaled NVFP4 cast: input and routing (first_dims, -// tensor_offsets) live on device, so it can be captured and replayed. -py::object nvfp4_group_row_scaled_cast_graph_safe(const at::Tensor &tensor, py::handle quantizer, - const size_t num_tensors, - const at::Tensor &first_dims, - const at::Tensor &tensor_offsets) { - using namespace transformer_engine::pytorch::detail; - init_extension(); - - NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D"); - NVTE_CHECK(tensor.scalar_type() == at::kBFloat16, - "Graph-safe grouped row-scaled NVFP4 cast requires a BF16 input."); - NVTE_CHECK(first_dims.is_cuda() && tensor_offsets.is_cuda(), - "first_dims and tensor_offsets must be CUDA tensors."); - NVTE_CHECK(first_dims.scalar_type() == at::kLong && tensor_offsets.scalar_type() == at::kLong, - "first_dims and tensor_offsets must be int64."); - NVTE_CHECK(first_dims.numel() == static_cast(num_tensors), - "first_dims must contain one row count per expert."); - NVTE_CHECK(tensor_offsets.numel() == static_cast(num_tensors) + 1, - "tensor_offsets must contain num_tensors + 1 element offsets."); - NVTE_CHECK(IsNVFP4Quantizers(quantizer.ptr()), - "Graph-safe grouped row-scaled NVFP4 cast only supports NVFP4 quantizers."); - - auto quantizer_cpp = convert_quantizer(quantizer); - NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); - NVTE_CHECK(nvfp4_quantizer_cpp->row_scaled_nvfp4, - "Graph-safe grouped cast requires a row-scaled NVFP4 quantizer."); - NVTE_CHECK(!nvfp4_quantizer_cpp->with_rht, - "Graph-safe grouped row-scaled NVFP4 cast does not support RHT."); - NVTE_CHECK(!nvfp4_quantizer_cpp->with_2d_quantization, - "Graph-safe grouped row-scaled NVFP4 cast does not support 2D quantization."); - NVTE_CHECK(!nvfp4_quantizer_cpp->stochastic_rounding, - "Graph-safe grouped row-scaled NVFP4 cast does not support stochastic rounding."); - NVTE_CHECK(!nvfp4_quantizer_cpp->with_amax_reduction, - "Graph-safe grouped row-scaled NVFP4 cast does not support amax reduction."); - NVTE_CHECK(nvfp4_quantizer_cpp->nvfp4_4over6_mode == kNVTENVFP44Over6Disabled, - "Graph-safe grouped row-scaled NVFP4 cast does not support 4over6."); - - auto input_contiguous = tensor.contiguous(); - // The kernels read num_tensors consecutive int64 entries from these buffers, so a - // strided view (e.g. a slice) would be misread. Make them contiguous before use. - auto first_dims_contiguous = first_dims.contiguous(); - auto tensor_offsets_contiguous = tensor_offsets.contiguous(); - - std::vector logical_shape; - for (const auto &d : input_contiguous.sizes()) { - logical_shape.push_back(d); - } - const auto logical_first_dim = logical_shape[0]; // capacity (rows) - const auto logical_last_dim = logical_shape[1]; // K (common last dim) - - std::optional first_dims_opt = first_dims_contiguous; - std::optional last_dims_opt = std::nullopt; - std::optional tensor_offsets_opt = tensor_offsets_contiguous; - - auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); - grouped_input_tensor.set_rowwise_data(input_contiguous.data_ptr(), - GetTransformerEngineDType(input_contiguous.scalar_type()), - getTensorShape(input_contiguous)); - grouped_input_tensor.set_first_dims(first_dims_contiguous.data_ptr(), DType::kInt64, - getTensorShape(first_dims_contiguous)); - grouped_input_tensor.set_tensor_offsets(tensor_offsets_contiguous.data_ptr(), DType::kInt64, - getTensorShape(tensor_offsets_contiguous)); - - auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( - num_tensors, logical_shape, GetTransformerEngineDType(input_contiguous.scalar_type()), - py::reinterpret_borrow(quantizer), first_dims_opt, last_dims_opt, - tensor_offsets_opt, logical_first_dim, logical_last_dim); - - if (logical_first_dim == 0 || logical_last_dim == 0) { - return py::reinterpret_borrow(grouped_output_py); - } - - auto stream = at::cuda::getCurrentCUDAStream(); - NVTE_SCOPED_GIL_RELEASE({ - nvte_group_nvfp4_compute_amax_graph_safe(grouped_input_tensor.data(), - grouped_output_tensor_cpp.data(), stream); - nvte_group_nvfp4_row_scaled_cast_with_amax_graph_safe(grouped_input_tensor.data(), - grouped_output_tensor_cpp.data(), stream); - }); - - return py::reinterpret_borrow(grouped_output_py); -} - py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, @@ -2042,51 +2014,6 @@ void split_quantize_nvfp4_impl_helper(const TensorWrapper &input, } } -// Single-launch grouped row-scaled NVFP4 quantize: per-expert amax followed by a -// single grouped cast, replacing the per-expert quantize loop. Byte-identical to -// quantizing each expert on its own. -void split_quantize_nvfp4_row_scaled_impl(const TensorWrapper &input, - const std::vector &input_list, - std::vector &output_list, - const std::vector &split_sections, - const std::vector &quantizers) { - const size_t num_tensors = input_list.size(); - if (num_tensors == 0 || input.numel() == 0) { - return; - } - const auto &quantizer = *quantizers.front(); - NVTE_CHECK(!quantizer.stochastic_rounding, - "Row-scaled NVFP4 quantization does not support stochastic rounding."); - NVTE_CHECK(!quantizer.with_rht, "Row-scaled NVFP4 quantization does not support RHT."); - NVTE_CHECK(!quantizer.with_2d_quantization, - "Row-scaled NVFP4 quantization does not support 2D quantization."); - NVTE_CHECK(!quantizer.with_amax_reduction, - "Row-scaled NVFP4 quantization does not support amax reduction."); - - std::vector nvte_tensor_output_list; - nvte_tensor_output_list.reserve(num_tensors); - for (size_t i = 0; i < num_tensors; ++i) { - nvte_tensor_output_list.push_back(output_list[i].data()); - } - - QuantizationConfigWrapper quant_config; - const auto use_fast_math = transformer_engine::getenv("NVTE_USE_FAST_MATH"); - if (use_fast_math) { - quant_config.set_use_fast_math(true); - } - - auto stream = at::cuda::getCurrentCUDAStream(); - - // Per-expert rowwise (and columnwise, when transpose is requested) amax. - nvte_group_nvfp4_compute_amax(input.data(), nvte_tensor_output_list.data(), split_sections.data(), - num_tensors, stream); - - // Single grouped cast consuming the amax computed above. - nvte_group_nvfp4_row_scaled_cast_with_amax(input.data(), nvte_tensor_output_list.data(), - split_sections.data(), num_tensors, quant_config, - stream); -} - void split_quantize_nvfp4_impl(const TensorWrapper &input, const std::vector &input_list, std::vector &output_list, @@ -2226,7 +2153,7 @@ std::vector split_quantize(const at::Tensor &tensor, // Choose implementation for allocating and populating tensors enum class AllocationMethod { UNFUSED, BULK_FP8_BLOCKWISE, BULK_MXFP8, BULK_NVFP4 }; - enum class QuantizationMethod { UNFUSED, FUSED_NVFP4, FUSED_NVFP4_ROW_SCALED }; + enum class QuantizationMethod { UNFUSED, FUSED_NVFP4 }; AllocationMethod allocation_method = AllocationMethod::UNFUSED; QuantizationMethod quantization_method = QuantizationMethod::UNFUSED; if (!disable_bulk_allocation) { @@ -2253,39 +2180,11 @@ std::vector split_quantize(const at::Tensor &tensor, const bool any_row_scaled = std::any_of(quantizer_cpp_list.begin(), quantizer_cpp_list.end(), is_row_scaled); if (any_row_scaled && !all_row_scaled) { - // Mixed row-scaled and tensor-scaled NVFP4 quantizers can't share the bulk NVFP4 - // allocation: its amax is sized from the first quantizer's row-scale setting, so a - // grouped cast would write per-row amax into a per-tensor buffer. Cast each split - // on its own instead. + // Bulk NVFP4 allocation sizes amax from the first quantizer. allocation_method = AllocationMethod::UNFUSED; quantization_method = QuantizationMethod::UNFUSED; } else if (all_row_scaled) { - // Grouped kernels need 128-aligned splits and last dim, <= 64 experts, and - // no 4over6; otherwise fall back to the per-expert path. - const bool all_splits_128 = std::all_of(split_sections.begin(), split_sections.end(), - [](size_t s) { return s % 128 == 0; }); - const bool last_dim_128 = input_shape.back() % 128 == 0; - const bool within_expert_cap = num_splits <= 64; - // The grouped kernels treat the input as 2D [rows, last_dim] and index rows by - // split_sections, so they require a 2D input whose splits cover every row. For - // higher-rank or partially-split inputs, fall back to the per-expert path. - const bool input_2d = input_shape.size() == 2; - size_t split_sum = 0; - for (size_t s : split_sections) split_sum += s; - const bool splits_cover_input = input_2d && split_sum == input_shape[0]; - const bool any_4over6 = std::any_of( - quantizer_cpp_list.begin(), quantizer_cpp_list.end(), - [](const std::unique_ptr &quantizer) { - return static_cast(quantizer.get())->nvfp4_4over6_mode != - kNVTENVFP44Over6Disabled; - }); - const bool grouped_row_scaled_cast = - transformer_engine::getenv("NVTE_NVFP4_GROUPED_ROW_SCALED_CAST", true); - quantization_method = - (grouped_row_scaled_cast && all_splits_128 && last_dim_128 && within_expert_cap && - input_2d && splits_cover_input && !any_4over6 && input_dtype == DType::kBFloat16) - ? QuantizationMethod::FUSED_NVFP4_ROW_SCALED - : QuantizationMethod::UNFUSED; + quantization_method = QuantizationMethod::UNFUSED; } else { quantization_method = QuantizationMethod::FUSED_NVFP4; } @@ -2366,17 +2265,6 @@ std::vector split_quantize(const at::Tensor &tensor, nvfp4_quantizers); break; } - case QuantizationMethod::FUSED_NVFP4_ROW_SCALED: { - // Single-launch grouped row-scaled NVFP4 quantize. - auto input_nvte = makeTransformerEngineTensor(input_dptr, input_shape, input_dtype); - std::vector nvfp4_quantizers; - for (auto &quantizer : quantizer_cpp_list) { - nvfp4_quantizers.push_back(static_cast(quantizer.get())); - } - split_quantize_nvfp4_row_scaled_impl(input_nvte, input_list, output_cpp_list, split_sections, - nvfp4_quantizers); - break; - } default: // General multi-tensor quantization multi_tensor_quantize_impl(input_list, quantizer_list, quantizer_cpp_list, output_cpp_list); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 29c12754e4a..2046408eba8 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -184,11 +184,6 @@ void bind_quantize_with_amax_extensions(py::module_ &m) { py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), py::arg("last_dims") = py::none(), py::arg("rowwise_amax"), py::arg("columnwise_amax"), py::arg("tensor_offsets") = py::none()); - m.def("nvfp4_group_row_scaled_cast_graph_safe", - transformer_engine::pytorch::nvfp4_group_row_scaled_cast_graph_safe, py::arg("tensor"), - py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), - py::arg("tensor_offsets"), - "CUDA-graph-safe grouped row-scaled NVFP4 cast (amax + cast) from device routing"); } } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index fba268d2723..5d8c37bd108 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2242,6 +2242,7 @@ std::pair NVFP4Quantizer::create_grouped_tenso } out_cpp.set_with_gemm_swizzled_scales(this->optimize_for_gemm); + out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); py::dict kwargs; diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index 7f377583d18..f9e4f487549 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -319,6 +319,10 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { } ret.set_with_gemm_swizzled_scales(with_gemm_swizzled); + if (py::hasattr(tensor, "_row_scaled_nvfp4")) { + ret.set_row_scaled_nvfp4(tensor.attr("_row_scaled_nvfp4").cast()); + } + return ret; } diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 2a3ce5feb55..997f15561e5 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -38,7 +38,6 @@ "ep_finalize", "ep_dispatch", "ep_combine", - "nvfp4_group_row_scaled_cast", "symm_mem_alloc", "release_symm_mem_pool", "is_symm_backed", @@ -1128,48 +1127,6 @@ def _make_grouped_mxfp8(data, scale_inv, token_counts, fp8_dtype, fake_dtype): ) -def nvfp4_group_row_scaled_cast( - tokens: torch.Tensor, - tokens_per_expert: torch.Tensor, - *, - fp4_dtype: DType = DType.kFloat4E2M1, - scale_dtype: DType = DType.kFloat8E4M3, -): - """Graph-safe grouped row-scaled NVFP4 cast (amax + cast) from device routing. - - ``tokens`` is a packed (capacity, hidden) bfloat16 buffer; ``tokens_per_expert`` - (int64 [num_experts]) is the 128-aligned per-expert row count. Routing stays on - device, so the cast can be captured and replayed. Returns a row-scaled NVFP4 - ``GroupedTensor``. - """ - from .tensor.nvfp4_tensor import NVFP4Quantizer - - if tokens.dim() != 2: - raise ValueError("tokens must be 2D [capacity_rows, hidden].") - if tokens.dtype is not torch.bfloat16: - raise TypeError(f"tokens must be bfloat16; got dtype={tokens.dtype}.") - if tokens_per_expert.dtype is not torch.int64: - raise TypeError(f"tokens_per_expert must be int64; got dtype={tokens_per_expert.dtype}.") - - _, hidden = tokens.shape - num_tensors = tokens_per_expert.numel() - quantizer = NVFP4Quantizer( - fp4_dtype=fp4_dtype, - scale_dtype=scale_dtype, - rowwise=True, - columnwise=False, - row_scaled_nvfp4=True, - ) - tensor_offsets = tex.splits_to_offsets(tokens_per_expert, hidden) - return tex.nvfp4_group_row_scaled_cast_graph_safe( - tokens, - quantizer, - num_tensors, - tokens_per_expert, - tensor_offsets, - ) - - def ep_dispatch( buffer: EpBuffer, tokens: torch.Tensor, From 3076ff704d5417f385d0153b4d65f51d87dc31e6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Oct 2026 06:56:48 +0000 Subject: [PATCH 13/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/cast/dispatch/quantize.cuh | 5 ++--- transformer_engine/pytorch/csrc/extensions/cast.cpp | 4 ++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 105cd2b8d1c..3be5bcf6cfe 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -579,9 +579,8 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor : output_tensor->columnwise_scale_inv.dtype; TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( scale_dtype, ScaleType, - nvfp4::group_row_scaled_cast_graph_safe(input_tensor, noop_tensor, - output_tensor, &quant_config_cpp, - stream);); + nvfp4::group_row_scaled_cast_graph_safe( + input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream);); break; } default: diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 9b6e754c429..2e0d7ccb67d 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -235,8 +235,8 @@ void group_quantize_nvfp4_row_scaled_impl(const GroupedTensorWrapper &grouped_in quant_config_cpp.set_use_fast_math(true); } NVTE_SCOPED_GIL_RELEASE({ - nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor.data(), - quant_config_cpp, stream); + nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor.data(), quant_config_cpp, + stream); }); } From de54786896a45b68032f17f070acf943680407b1 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Thu, 22 Oct 2026 01:26:41 -0700 Subject: [PATCH 14/24] Require 128-row experts for uniform row-scaled NVFP4 group_quantize Signed-off-by: Cael Ling --- .../pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py | 8 ++++++++ transformer_engine/pytorch/csrc/extensions/cast.cpp | 5 +++-- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py index 567065a9a73..a740e0ad2cc 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py @@ -183,6 +183,14 @@ def test_non_contiguous_input_matches_contiguous() -> None: _assert_matches_reference(grouped, x_t.contiguous(), split_sections) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_rejects_unaligned_uniform_shapes() -> None: + x = torch.randn((128, 128), dtype=torch.bfloat16, device="cuda") + + with pytest.raises(RuntimeError): + _group_quantize(x, 2, None) + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) def test_rejects_strided_first_dims() -> None: x = torch.randn((384, 128), dtype=torch.bfloat16, device="cuda") diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 2e0d7ccb67d..e89c4daf090 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -342,10 +342,11 @@ py::object group_quantize(const at::Tensor &input, py::handle quantizer, const s NVTE_CHECK(first_dims->is_contiguous() && first_dims->scalar_type() == at::kLong, "group_quantize: first_dims must be a contiguous int64 tensor."); } else { - NVTE_CHECK(num_tensors > 0 && logical_first_dim % num_tensors == 0, + NVTE_CHECK(num_tensors > 0 && logical_first_dim % (num_tensors * 128) == 0, "group_quantize: row-scaled NVFP4 with uniform shapes requires the first " "dimension (", - logical_first_dim, ") to be divisible by num_tensors (", num_tensors, ")."); + logical_first_dim, ") to be a multiple of 128 * num_tensors (", num_tensors, + ")."); first_dims = at::full({static_cast(num_tensors)}, static_cast(logical_first_dim / num_tensors), tensor.options().dtype(at::kLong)); From 3d656d21b24c5ccf4420f59fb34788e52e52373c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Oct 2026 08:49:53 +0000 Subject: [PATCH 15/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/csrc/extensions/cast.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index e89c4daf090..142481493d1 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -345,8 +345,7 @@ py::object group_quantize(const at::Tensor &input, py::handle quantizer, const s NVTE_CHECK(num_tensors > 0 && logical_first_dim % (num_tensors * 128) == 0, "group_quantize: row-scaled NVFP4 with uniform shapes requires the first " "dimension (", - logical_first_dim, ") to be a multiple of 128 * num_tensors (", num_tensors, - ")."); + logical_first_dim, ") to be a multiple of 128 * num_tensors (", num_tensors, ")."); first_dims = at::full({static_cast(num_tensors)}, static_cast(logical_first_dim / num_tensors), tensor.options().dtype(at::kLong)); From 1d9a31879c4fba88611f9ea97006a71df9c2e0b4 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Thu, 22 Oct 2026 02:48:23 -0700 Subject: [PATCH 16/24] Use tex.group_quantize for row-scaled grouped GEMM tests Signed-off-by: Cael Ling --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 35 ++------------------ 1 file changed, 2 insertions(+), 33 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 9e044f5dddf..5cf8ae45f73 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -21,12 +21,6 @@ recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) -# The grouped-tensor GEMM tests below need the graph-safe grouped row-scaled cast (PR #3649). -row_scaled_grouped_cast_available = hasattr(tex, "nvfp4_group_row_scaled_cast_graph_safe") -reason_for_no_row_scaled_grouped_cast = ( - "Grouped row-scaled NVFP4 cast (nvfp4_group_row_scaled_cast_graph_safe) is unavailable." -) - @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( @@ -875,8 +869,7 @@ def _make_row_scaled_grouped_activation( ) -> GroupedTensor: """Pack per-expert bf16 activations into one row-scaled NVFP4 GroupedTensor. - The graph-safe grouped row-scaled cast (PR #3649) requires each expert row count and - ``k`` to be multiples of 128. + Each expert row count and ``k`` must be multiples of 128. """ device = "cuda" if tokens is None: @@ -894,10 +887,7 @@ def _make_row_scaled_grouped_activation( ) # Bake cuBLAS-swizzled scales in the cast so the GEMM needs no separate swizzle pass. quantizer.optimize_for_gemm = True - offsets = tex.splits_to_offsets(tokens_per_expert, k) - return tex.nvfp4_group_row_scaled_cast_graph_safe( - tokens, quantizer, len(m_splits), tokens_per_expert, offsets - ) + return tex.group_quantize(tokens, quantizer, len(m_splits), tokens_per_expert) def _make_output_grouped_tensor( @@ -916,9 +906,6 @@ def _make_output_grouped_tensor( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.skipif( - not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast -) @pytest.mark.parametrize( "m_splits, k, n", [ @@ -979,9 +966,6 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_fprop( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.skipif( - not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast -) @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 @@ -1047,9 +1031,6 @@ def _run(): @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.skipif( - not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast -) @pytest.mark.parametrize( "m_splits, k, n", [ @@ -1098,9 +1079,6 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.skipif( - not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast -) @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 @@ -1134,9 +1112,6 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad_discrete_out( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.skipif( - not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast -) @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 @@ -1172,9 +1147,6 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_uneven_colscale( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.skipif( - not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast -) @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 @@ -1231,9 +1203,6 @@ def _run(): @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.skipif( - not row_scaled_grouped_cast_available, reason=reason_for_no_row_scaled_grouped_cast -) @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 From 035a322fb476a77ba65bedb87409f745159c454b Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Thu, 22 Oct 2026 04:02:42 -0700 Subject: [PATCH 17/24] Trim comments in row-scaled grouped GEMM Signed-off-by: Cael Ling --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 49 +++---------------- .../nvfp4/group_quantize_transpose_nvfp4.cuh | 4 -- .../pytorch/cpp_extensions/gemm.py | 21 +------- transformer_engine/pytorch/csrc/quantizer.cpp | 1 - .../tensor/storage/grouped_tensor_storage.py | 1 - 5 files changed, 9 insertions(+), 67 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 5cf8ae45f73..9f744352756 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -850,8 +850,6 @@ def _row_scaled_grouped_tensor_gemm_supported() -> bool: return False if torch.cuda.get_device_capability() < (10, 0): return False - # cuBLAS enables grouped NVFP4 GEMM at 13.4 (CUBLAS_NVFP4_GROUPED_GEMM_VERSION 130400); - # below that the GEMM rejects the call, so skip rather than fail the version check. return tex.get_cublasLt_version() >= 130400 @@ -867,10 +865,7 @@ def _make_row_scaled_grouped_activation( columnwise: bool, tokens: torch.Tensor | None = None, ) -> GroupedTensor: - """Pack per-expert bf16 activations into one row-scaled NVFP4 GroupedTensor. - - Each expert row count and ``k`` must be multiples of 128. - """ + """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) @@ -885,7 +880,6 @@ def _make_row_scaled_grouped_activation( with_post_rht_amax=False, row_scaled_nvfp4=True, ) - # Bake cuBLAS-swizzled scales in the cast so the GEMM needs no separate swizzle pass. quantizer.optimize_for_gemm = True return tex.group_quantize(tokens, quantizer, len(m_splits), tokens_per_expert) @@ -919,9 +913,7 @@ def _make_output_grouped_tensor( def test_nvfp4_row_scaled_grouped_tensor_gemm_fprop( m_splits: list[int], k: int, n: int, out_dtype: torch.dtype ) -> None: - """Forward (TN) grouped-tensor GEMM with a row-scaled activation matches a per-expert - dense GEMM. Only the activation is row-scaled (rowwise RHS), so this covers post_scale_rows. - """ + """Row-scaled grouped-tensor TN GEMM matches a per-expert dense GEMM.""" if not _row_scaled_grouped_tensor_gemm_supported(): pytest.skip(_reason_for_no_grouped_tensor_gemm) @@ -956,7 +948,6 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_fprop( general_grouped_gemm_for_grouped_tensor(weights, activation, grouped_out, layout="TN") actual = grouped_out.split_into_quantized_tensors() - # Identical FP4 operands and post-scale; only grouped-vs-single cuBLAS accumulation differs. if out_dtype == torch.bfloat16: tols = dict(atol=0.0, rtol=7.8e-3) else: @@ -970,9 +961,7 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_fprop( def test_nvfp4_row_scaled_grouped_tensor_gemm_cudagraph( m_splits: list[int], k: int, n: int ) -> None: - """The row-scaled grouped-tensor GEMM is CUDA-graph capturable: a captured GEMM replayed - after refreshing the activation buffers in place matches an eager recompute. - """ + """A captured row-scaled grouped-tensor GEMM replays correctly.""" if not _row_scaled_grouped_tensor_gemm_supported(): pytest.skip(_reason_for_no_grouped_tensor_gemm) @@ -1005,7 +994,6 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_cudagraph( def _run(): general_grouped_gemm_for_grouped_tensor(weights, activation, grouped_out, layout="TN") - # Warmup to initialize kernels/allocator before capture. torch.cuda.synchronize() _run() torch.cuda.synchronize() @@ -1014,7 +1002,6 @@ def _run(): with torch.cuda.graph(graph): _run() - # Refresh the activation in place with freshly cast values, then replay. 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) @@ -1044,10 +1031,7 @@ def _run(): def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad( m_splits: list[int], k: int, n: int, out_dtype: torch.dtype ) -> None: - """Weight-gradient (NT) grouped-tensor GEMM with two row-scaled operands matches a per-expert - dense GEMM. Both X and dY are consumed columnwise, covering post_scale_cols, post_scale_rows, - and the columnwise swizzled-scale epilogue. - """ + """Row-scaled grouped-tensor NT GEMM matches a per-expert dense GEMM.""" if not _row_scaled_grouped_tensor_gemm_supported(): pytest.skip(_reason_for_no_grouped_tensor_gemm) @@ -1069,7 +1053,6 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad( general_grouped_gemm_for_grouped_tensor(x, dy, grouped_out, layout="NT") actual = grouped_out.split_into_quantized_tensors() - # Identical FP4 operands and post-scales; only grouped-vs-single cuBLAS accumulation differs. if out_dtype == torch.bfloat16: tols = dict(atol=0.0, rtol=7.8e-3) else: @@ -1083,9 +1066,7 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad( def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad_discrete_out( m_splits: list[int], k: int, n: int ) -> None: - """Weight-gradient (NT) grouped-tensor GEMM into a list output exercises the discrete-out - post-scale path (per-expert row/column slices) against a per-expert dense GEMM. - """ + """Row-scaled NT GEMM into a list output matches a per-expert dense GEMM.""" if not _row_scaled_grouped_tensor_gemm_supported(): pytest.skip(_reason_for_no_grouped_tensor_gemm) @@ -1116,9 +1097,7 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad_discrete_out( def test_nvfp4_row_scaled_grouped_tensor_gemm_uneven_colscale( m_splits: list[int], k: int, n: int ) -> None: - """TN grouped-tensor GEMM with both operands row-scaled and uneven output rows exercises the - grouped-out column post-scale over each group's actual rows (not an equal split). - """ + """Column post-scale with uneven groups matches a per-expert dense GEMM.""" if not _row_scaled_grouped_tensor_gemm_supported(): pytest.skip(_reason_for_no_grouped_tensor_gemm) @@ -1127,7 +1106,6 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_uneven_colscale( num_gemms = len(m_splits) out_dtype = torch.float32 - # A (weights): uniform [n, k] per expert; B (activation): uneven [m_i, k]. 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() @@ -1151,10 +1129,7 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_uneven_colscale( def test_nvfp4_row_scaled_grouped_tensor_gemm_colscale_cudagraph( m_splits: list[int], k: int, n: int ) -> None: - """The grouped-out column post-scale (two row-scaled operands, uneven rows) is CUDA-graph - capturable: the device per-expert row counts drive the column-scale expansion with no host - sync, so a captured GEMM replayed after refreshing both operands matches an eager recompute. - """ + """A captured column post-scale with uneven groups replays correctly.""" if not _row_scaled_grouped_tensor_gemm_supported(): pytest.skip(_reason_for_no_grouped_tensor_gemm) @@ -1184,7 +1159,6 @@ def _run(): with torch.cuda.graph(graph): _run() - # Refresh both operands in place with freshly cast values, then replay. 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)): @@ -1207,11 +1181,7 @@ def _run(): def test_nvfp4_row_scaled_grouped_tensor_gemm_output_capacity( m_splits: list[int], k: int, n: int ) -> None: - """A grouped output whose backing buffer carries extra capacity rows beyond the live region - is scaled only over the live rows. The post-scale restricts itself to logical_shape[0] - instead of the full buffer, so the extra rows neither break the per-row multiply nor change - the result, which still matches a per-expert dense GEMM. - """ + """Output capacity rows beyond the live region do not affect the result.""" if not _row_scaled_grouped_tensor_gemm_supported(): pytest.skip(_reason_for_no_grouped_tensor_gemm) @@ -1243,9 +1213,6 @@ def test_nvfp4_row_scaled_grouped_tensor_gemm_output_capacity( ) grouped_out = _make_output_grouped_tensor(m_splits, n, device, out_dtype) - # Enlarge the backing buffer past the live region (logical_shape keeps the live size), as a - # capacity-padded output would. Groups are still written at their live offsets; without the - # live-row slice the per-row multiply would hit the extra rows and misalign. live_numel = grouped_out.rowwise_data.numel() padded = torch.full( (live_numel + 128 * n,), 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 7c9d8a30ef6..85284942e2c 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -1564,8 +1564,6 @@ __device__ __forceinline__ void group_row_scaled_cast_body( ScalesVec &scales_vec = *reinterpret_cast(sSFrowwise[row]); const size_t local_row = static_cast(local_row_base) + row; if constexpr (WITH_GEMM_SWIZZLED_SCALES) { - // gemm_swizzled_scale_idx takes the compact scale column count - // (scale_stride_row = K / SCALE_DIM); local_row is expert-local. const ScaleType *scales = reinterpret_cast(&scales_vec); #pragma unroll for (int k = 0; k < SCALES_PER_CHUNK_X; ++k) { @@ -1588,8 +1586,6 @@ __device__ __forceinline__ void group_row_scaled_cast_body( ScalesVec &scales_vec = *reinterpret_cast(sSFcolwise[row_tr]); const size_t global_col = static_cast(block_offset_X) + row_tr; if constexpr (WITH_GEMM_SWIZZLED_SCALES) { - // Columnwise counterpart: column count is scale_stride_col = M_t / SCALE_DIM, - // and global_col indexes the K rows. const ScaleType *scales = reinterpret_cast(&scales_vec); #pragma unroll for (int k = 0; k < SCALES_PER_CHUNK_Y; ++k) { diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 98481c0aab0..b6010a72bad 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1139,9 +1139,6 @@ 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.") - # Row-scaled operands carry an outer FP32 scale the per-tensor grouped GEMM cannot - # apply. Neutralize it, GEMM into an FP32 scratch, then post-scale the output - # (mirrors the dense general_gemm path). def _is_row_scaled(operand) -> bool: return isinstance(operand, GroupedTensorStorage) and operand.row_scaled_nvfp4 @@ -1166,9 +1163,6 @@ def _list_has_row_scaled(operand) -> bool: raise NotImplementedError( "Row-scaled NVFP4 grouped GEMM does not support explicit beta; use accumulate." ) - # A's scale lands on output columns, B's on output rows; rowwise amax for an - # untransposed operand, columnwise for a transposed one (see - # _nvfp4_row_scaled_gemm_inputs). if _is_row_scaled(A): A = A.copy() if transa: @@ -1187,7 +1181,6 @@ def _list_has_row_scaled(operand) -> bool: post_scale_rows = B.amax B.amax = post_scale_rows.new_ones(B.num_tensors) B.row_scaled_nvfp4 = False - # Post-scale happens after the GEMM, so cuBLAS must not accumulate itself. accumulate = False def _is_fp8_blockwise(operand) -> bool: @@ -1299,8 +1292,6 @@ def _store(dst: torch.Tensor, src: torch.Tensor) -> None: dst.copy_(src.to(dst.dtype)) if is_discrete_out: - # Slice the per-row/per-column scales by each output's actual shape so uneven - # per-expert row/column counts are handled correctly. row_offset = 0 col_offset = 0 for i in range(num_tensors): @@ -1317,27 +1308,17 @@ def _store(dst: torch.Tensor, src: torch.Tensor) -> None: row_offset += rows col_offset += cols else: - # Single packed output. Groups are stacked along rows (uniform or varying first - # dim) and share one output width, so the scales below index the buffer by row. - # A varying per-expert output width packs groups along columns instead, which this - # row-wise expansion cannot represent; no TE caller produces such an output, since - # grouped GEMM outputs share the non-contracted free dimension. 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] - # logical_shape[0] is the live row count (host int). Restrict to it so a buffer - # carrying extra capacity rows does not misalign with the per-row/per-group scales. live_rows = out.logical_shape[0] flat = scratch.rowwise_data.view(-1, last_dim)[:live_rows] if post_scale_rows is not None: - # One scale per output row. flat.mul_(post_scale_rows.view(-1, 1)) if post_scale_cols is not None: - # Per-expert, per-column. Expand each group's column scale over its rows from - # the device per-expert row counts. Passing output_size keeps repeat_interleave - # from syncing the counts to host, so this stays valid under CUDA-graph capture. + # output_size avoids a host sync under CUDA-graph capture. col_scales = post_scale_cols.view(num_tensors, last_dim) if out.first_dims is None: rows = live_rows // num_tensors diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 669c9eec5be..8f36532453f 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2205,7 +2205,6 @@ 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) { - // Per-column for row-scaled, one per group otherwise (matches rowwise amax). const int64_t col_amax_elements = row_scaled_nvfp4 ? static_cast(num_tensors) * static_cast(logical_last_dim) diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 9eafacb2701..334f158cd8b 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -860,7 +860,6 @@ def make_grouped_tensor( raise ValueError( "Row-scaled NVFP4 grouped quantization requires rowwise usage." ) - # Per-row/per-column for row-scaled, one per group otherwise. total_amax_elements = ( sum(math.prod(s[:-1]) for s in shape) if row_scaled_nvfp4 else num_tensors ) From 055277ed66ccb9aad707aa392f7f565bc7d1ab5b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Oct 2026 11:09:31 +0000 Subject: [PATCH 18/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/csrc/extensions/cast.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index e89c4daf090..142481493d1 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -345,8 +345,7 @@ py::object group_quantize(const at::Tensor &input, py::handle quantizer, const s NVTE_CHECK(num_tensors > 0 && logical_first_dim % (num_tensors * 128) == 0, "group_quantize: row-scaled NVFP4 with uniform shapes requires the first " "dimension (", - logical_first_dim, ") to be a multiple of 128 * num_tensors (", num_tensors, - ")."); + logical_first_dim, ") to be a multiple of 128 * num_tensors (", num_tensors, ")."); first_dims = at::full({static_cast(num_tensors)}, static_cast(logical_first_dim / num_tensors), tensor.options().dtype(at::kLong)); From 0432c8ec5e079e764ab3abd5cafa46c4cdcade57 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Mon, 26 Oct 2026 15:22:00 -0700 Subject: [PATCH 19/24] Fold row-scaled NVFP4 amax into nvte_group_amax Dispatch nvte_group_amax and nvte_group_amax_graph_safe on the output tensor format so that row-scaled NVFP4 outputs get per-row amaxes, and deprecate nvte_group_nvfp4_compute_amax in favor of nvte_group_amax. Rename the row-scaled amax helpers to group_row_amax(_graph_safe). In split_quantize, fall back to unfused allocation and quantization whenever any NVFP4 quantizer is row-scaled. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Tim Moon --- .../test_row_scaled_nvfp4_grouped_amax.cu | 10 +++---- tests/cpp/test_common.h | 1 - .../common/cast/cast_grouped.cu | 13 +-------- .../common/cast/dispatch/quantize.cuh | 2 +- .../nvfp4/group_quantize_transpose_nvfp4.cuh | 14 ++++----- .../graph_safe_group_hadamard_transform.cu | 11 ++++++- .../group_hadamard_transform.cu | 29 ++++++++++++++++++- .../common/include/transformer_engine/cast.h | 17 +---------- .../include/transformer_engine/multi_tensor.h | 12 ++++++-- .../pytorch/csrc/extensions/cast.cpp | 19 +++++------- 10 files changed, 69 insertions(+), 59 deletions(-) diff --git a/tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu b/tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu index 8b883f25e25..c4b05ff5f5a 100644 --- a/tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu +++ b/tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu @@ -80,12 +80,12 @@ void performGroupedFusedAmaxTest(const GroupedAmaxConfig& cfg) { auto outs = make_row_scaled_outputs("gout_", Ms, K); std::vector out_handles; for (auto& t : outs) out_handles.push_back(t->data()); - nvte_group_nvfp4_compute_amax(packed.data(), out_handles.data(), splits.data(), num_tensors, 0); + nvte_group_amax(packed.data(), out_handles.data(), splits.data(), num_tensors, 0); auto outs2 = make_row_scaled_outputs("gout2_", Ms, K); std::vector out_handles2; for (auto& t : outs2) out_handles2.push_back(t->data()); - nvte_group_nvfp4_compute_amax(packed.data(), out_handles2.data(), splits.data(), num_tensors, 0); + nvte_group_amax(packed.data(), out_handles2.data(), splits.data(), num_tensors, 0); cudaDeviceSynchronize(); ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); @@ -187,7 +187,7 @@ TEST(NVFP4GroupedRowScaledAmaxTestSuite, HeterogeneousColumnwise) { std::vector handles; for (auto& t : outs) handles.push_back(t->data()); const std::vector splits(Ms); - nvte_group_nvfp4_compute_amax(packed.data(), handles.data(), splits.data(), num_tensors, 0); + nvte_group_amax(packed.data(), handles.data(), splits.data(), num_tensors, 0); cudaDeviceSynchronize(); ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); @@ -265,7 +265,7 @@ void runLoopVsGroupedBench(size_t num_tensors, size_t M, size_t K, int iters) { double grouped = median_ms( [&]() { - nvte_group_nvfp4_compute_amax(packed.data(), h_g.data(), splits.data(), num_tensors, 0); + nvte_group_amax(packed.data(), h_g.data(), splits.data(), num_tensors, 0); }, iters); @@ -274,7 +274,7 @@ void runLoopVsGroupedBench(size_t num_tensors, size_t M, size_t K, int iters) { for (size_t i = 0; i < num_tensors; ++i) { NVTETensor one = outs_l[i]->data(); size_t s1 = M; - nvte_group_nvfp4_compute_amax(ins[i]->data(), &one, &s1, 1, 0); + nvte_group_amax(ins[i]->data(), &one, &s1, 1, 0); } }, iters); diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 5beed3c47c7..58ea430d8a6 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -108,7 +108,6 @@ struct TypeInfo { struct Helper { constexpr static DType getType() { constexpr int i = static_cast(current); - // fp8ue5m3 is only in types on CUDA >= 13.4. if constexpr (i < static_cast(std::tuple_size::value)) { if (std::is_same::type>::value) { return current; diff --git a/transformer_engine/common/cast/cast_grouped.cu b/transformer_engine/common/cast/cast_grouped.cu index 49ab39f7e79..11314eaa6ca 100644 --- a/transformer_engine/common/cast/cast_grouped.cu +++ b/transformer_engine/common/cast/cast_grouped.cu @@ -101,16 +101,5 @@ void nvte_group_nvfp4_quantize_with_amax(const NVTETensor input, NVTETensor *out void nvte_group_nvfp4_compute_amax(const NVTETensor input, NVTETensor *outputs, const size_t *split_sections, const size_t num_tensors, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_nvfp4_compute_amax); - using namespace transformer_engine; - - const Tensor *input_tensor = convertNVTETensorCheck(input); - std::vector output_tensors; - output_tensors.reserve(num_tensors); - for (size_t i = 0; i < num_tensors; ++i) { - output_tensors.push_back(convertNVTETensorCheck(outputs[i])); - } - - dispatch::nvfp4::group_compute_fused_amax(*input_tensor, /*noop=*/nullptr, output_tensors, - split_sections, num_tensors, stream); + nvte_group_amax(input, outputs, split_sections, num_tensors, stream); } diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 3be5bcf6cfe..13c3694d040 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -573,7 +573,7 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor if (input_tensor->num_tensors == 0) { break; } - nvfp4::group_compute_fused_amax_graph_safe(input_tensor, noop_tensor, output_tensor, stream); + nvfp4::group_row_amax_graph_safe(input_tensor, noop_tensor, output_tensor, stream); const DType scale_dtype = output_tensor->scale_inv.has_data() ? output_tensor->scale_inv.dtype : output_tensor->columnwise_scale_inv.dtype; 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 c7e3e0f7d4e..dd406e51370 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -334,10 +334,10 @@ __global__ void __launch_bounds__(FA_THREADS_NUM) } // namespace group_row_scaled_amax_kernel -inline void group_compute_fused_amax(const Tensor &input, const Tensor *noop, - std::vector &output_list, - const size_t *split_sections, size_t num_tensors, - cudaStream_t stream) { +inline void group_row_amax(const Tensor &input, const Tensor *noop, + std::vector &output_list, + const size_t *split_sections, size_t num_tensors, + cudaStream_t stream) { #if FP4_TYPE_SUPPORTED using namespace group_row_scaled_amax_kernel; @@ -417,11 +417,11 @@ inline void group_compute_fused_amax(const Tensor &input, const Tensor *noop, #endif // FP4_TYPE_SUPPORTED } -// Graph-safe counterpart of group_compute_fused_amax: routing comes from device +// Graph-safe calculation of amaxes for row-scaled NVFP4: routing comes from device // output->tensor_offsets, amax goes to the contiguous rowwise [sum_M] / columnwise // [num_tensors * K] buffers. BF16-only, VARYING_FIRST_DIM. -inline void group_compute_fused_amax_graph_safe(const GroupedTensor *input, const Tensor *noop, - GroupedTensor *output, cudaStream_t stream) { +inline void group_row_amax_graph_safe(const GroupedTensor *input, const Tensor *noop, + GroupedTensor *output, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED using namespace group_row_scaled_amax_kernel; diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu index 2316d9697ae..0528ac0927a 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu @@ -14,6 +14,7 @@ #include +#include "common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh" #include "common/common.h" #include "common/util/ptx.cuh" #include "common/utils.cuh" @@ -563,7 +564,7 @@ void nvte_group_hadamard_transform_amax_graph_safe(const NVTEGroupedTensor input static_cast(random_sign_mask_t), false, stream); } -// Grouped-tensor amax without doing hadamard transform +// Grouped-tensor amax void nvte_group_amax_graph_safe(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_group_amax_graph_safe); @@ -576,5 +577,13 @@ void nvte_group_amax_graph_safe(const NVTEGroupedTensor input, NVTEGroupedTensor return; } + // Per-row amax + if (output_tensor->scaling_mode == NVTE_NVFP4_1D_SCALING && output_tensor->row_scaled_nvfp4) { + dispatch::nvfp4::group_row_amax_graph_safe(input_tensor, /*noop=*/nullptr, output_tensor, + stream); + return; + } + + // Per-tensor amax group_hadamard_transform_amax_graph_safe(input_tensor, output_tensor, 0, 0, true, stream); } diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu index 24d06e5d25e..5949dcaf4ae 100644 --- a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu @@ -4,6 +4,7 @@ * See LICENSE for license information. ************************************************************************/ +#include #include #include #include @@ -14,6 +15,7 @@ #include +#include "common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh" #include "common/common.h" #include "common/util/ptx.cuh" #include "common/utils.cuh" @@ -582,7 +584,7 @@ void nvte_group_hadamard_transform_amax(const NVTETensor input, NVTETensor* outp static_cast(random_sign_mask_t), false, stream); } -// Grouped-tensor amax without doing hadamard transform +// Grouped-tensor amax void nvte_group_amax(const NVTETensor input, NVTETensor* outputs, const size_t* split_sections, size_t num_tensors, cudaStream_t stream) { NVTE_API_CALL(nvte_group_amax); @@ -597,6 +599,31 @@ void nvte_group_amax(const NVTETensor input, NVTETensor* outputs, const size_t* output_list[i] = convertNVTETensorCheck(outputs[i]); } + // Check that tensors have uniform scaling modes + const auto scaling_mode = output_list[0]->scaling_mode; + for (const auto *t : output_list) { + NVTE_CHECK(t->scaling_mode == scaling_mode, + "Tensors have inconsistent scaling modes (found ", + to_string(scaling_mode), " and ", to_string(t->scaling_mode), ")."); + } + + // Compute per-row amax if outputs are row-scaled NVFP4 tensors + if (scaling_mode == NVTE_NVFP4_1D_SCALING) { + auto is_row_scaled_nvfp4 = [](const Tensor *t) -> bool { return t->row_scaled_nvfp4; }; + const bool all_row_scaled_nvfp4 = std::all_of(output_list.begin(), output_list.end(), + is_row_scaled_nvfp4); + const bool any_row_scaled_nvfp4 = std::any_of(output_list.begin(), output_list.end(), + is_row_scaled_nvfp4); + NVTE_CHECK(all_row_scaled_nvfp4 == any_row_scaled_nvfp4, + "Output tensors are not consistent between row-scaling and tensor-scaling."); + if (all_row_scaled_nvfp4) { + dispatch::nvfp4::group_row_amax(*input_tensor, /*noop=*/nullptr, output_list, + split_sections, num_tensors, stream); + return; + } + } + + // Per-tensor amax group_hadamard_transform_amax(*input_tensor, output_list, split_sections, num_tensors, 0, 0, true, stream); } diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 774146dea6b..0c0c0d89d4d 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -495,22 +495,7 @@ void nvte_group_nvfp4_quantize_with_amax(const NVTETensor input, NVTETensor *out const NVTEQuantizationConfig quant_config, cudaStream_t stream); -/*! \brief Computes per-expert rowwise and/or columnwise amax for a grouped - * row-scaled NVFP4 input in a single kernel launch. - * - * The input is one packed (sum_M, K) BF16 buffer; `split_sections` gives the - * per-expert row counts (each a multiple of 128) and K must also be a multiple - * of 128. For expert i the rowwise amax [M_i] is written to `outputs[i]->amax` - * and the columnwise amax [K] to `outputs[i]->columnwise_amax`, whichever are - * allocated; experts may opt into either direction independently (a null buffer - * is skipped). - * - * \param[in] input Packed grouped input tensor (BF16). - * \param[in,out] outputs Per-expert output tensors receiving the amax vectors. - * \param[in] split_sections Per-expert row counts. - * \param[in] num_tensors Number of output tensors. - * \param[in] stream CUDA stream used for the operation. - */ +/*! \deprecated Use nvte_group_amax instead. */ void nvte_group_nvfp4_compute_amax(const NVTETensor input, NVTETensor *outputs, const size_t *split_sections, size_t num_tensors, cudaStream_t stream); diff --git a/transformer_engine/common/include/transformer_engine/multi_tensor.h b/transformer_engine/common/include/transformer_engine/multi_tensor.h index 09ab260f15b..30f7669392e 100644 --- a/transformer_engine/common/include/transformer_engine/multi_tensor.h +++ b/transformer_engine/common/include/transformer_engine/multi_tensor.h @@ -302,10 +302,12 @@ void nvte_multi_tensor_compute_scale_inv_e8m0_cuda(int chunk_size, NVTETensor ** * This function is experimental and the API is not stable. * * For each tensor split, compute the maximum absolute value (amax) - * and populate the amax of the corresponding output tensor. + * and populate the amax buffers of the corresponding output tensor. + * All output tensors must have the same tensor format, and the + * tensor format determines which amaxes are computed. * * \param[in] input Input tensor. - * \param[in,out] outputs Array of output tensors. Only the amax is updated. + * \param[in,out] outputs Array of output tensors. Only amax buffers are updated. * \param[in] split_sections Size of each tensor split along dimension 0. * \param[in] num_tensors Number of tensor splits. * \param[in] stream CUDA stream used for the operation. @@ -313,10 +315,14 @@ void nvte_multi_tensor_compute_scale_inv_e8m0_cuda(int chunk_size, NVTETensor ** void nvte_group_amax(const NVTETensor input, NVTETensor *outputs, const size_t *split_sections, size_t num_tensors, cudaStream_t stream); -/*! \brief Grouped-tensor amax without doing hadamard transform. +/*! \brief Graph-safe variant of nvte_group_amax for grouped tensors. * * This function is experimental and the API is not stable. * + * Compute the amax for each tensor in the group and populate the amax + * buffers of the output tensor. The output tensor format determines + * which amaxes are computed. + * * \param[in] input NVTEGroupedTensor Input tensor. * \param[in,out] output NVTEGroupedTensor Output tensor. * \param[in] stream CUDA stream used for the operation. diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 142481493d1..3d61dc525f8 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -2171,21 +2171,16 @@ std::vector split_quantize(const at::Tensor &tensor, [](const py::handle &quantizer) -> bool { return detail::IsNVFP4Quantizers(quantizer.ptr()); })) { - allocation_method = AllocationMethod::BULK_NVFP4; - auto is_row_scaled = [](const std::unique_ptr &quantizer) { - return static_cast(quantizer.get())->row_scaled_nvfp4; - }; - const bool all_row_scaled = - std::all_of(quantizer_cpp_list.begin(), quantizer_cpp_list.end(), is_row_scaled); - const bool any_row_scaled = - std::any_of(quantizer_cpp_list.begin(), quantizer_cpp_list.end(), is_row_scaled); - if (any_row_scaled && !all_row_scaled) { - // Bulk NVFP4 allocation sizes amax from the first quantizer. + const bool has_row_scaled_nvfp4 = + std::any_of(quantizer_cpp_list.begin(), quantizer_cpp_list.end(), + [](const std::unique_ptr &quantizer) { + return static_cast(quantizer.get())->row_scaled_nvfp4; + }); + if (has_row_scaled_nvfp4) { allocation_method = AllocationMethod::UNFUSED; quantization_method = QuantizationMethod::UNFUSED; - } else if (all_row_scaled) { - quantization_method = QuantizationMethod::UNFUSED; } else { + allocation_method = AllocationMethod::BULK_NVFP4; quantization_method = QuantizationMethod::FUSED_NVFP4; } } From 55f69ca27de2e3c24b4ef5543cbd70b55ae0be59 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:28:36 +0000 Subject: [PATCH 20/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../nvfp4/group_quantize_transpose_nvfp4.cuh | 5 ++--- .../group_hadamard_transform.cu | 21 +++++++++---------- 2 files changed, 12 insertions(+), 14 deletions(-) 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 dd406e51370..167c7c173e4 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -335,9 +335,8 @@ __global__ void __launch_bounds__(FA_THREADS_NUM) } // namespace group_row_scaled_amax_kernel inline void group_row_amax(const Tensor &input, const Tensor *noop, - std::vector &output_list, - const size_t *split_sections, size_t num_tensors, - cudaStream_t stream) { + std::vector &output_list, const size_t *split_sections, + size_t num_tensors, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED using namespace group_row_scaled_amax_kernel; diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu index 5949dcaf4ae..f8c04bea0cb 100644 --- a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu @@ -4,7 +4,6 @@ * See LICENSE for license information. ************************************************************************/ -#include #include #include #include @@ -13,6 +12,7 @@ #include #include +#include #include #include "common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh" @@ -601,24 +601,23 @@ void nvte_group_amax(const NVTETensor input, NVTETensor* outputs, const size_t* // Check that tensors have uniform scaling modes const auto scaling_mode = output_list[0]->scaling_mode; - for (const auto *t : output_list) { - NVTE_CHECK(t->scaling_mode == scaling_mode, - "Tensors have inconsistent scaling modes (found ", + for (const auto* t : output_list) { + NVTE_CHECK(t->scaling_mode == scaling_mode, "Tensors have inconsistent scaling modes (found ", to_string(scaling_mode), " and ", to_string(t->scaling_mode), ")."); } // Compute per-row amax if outputs are row-scaled NVFP4 tensors if (scaling_mode == NVTE_NVFP4_1D_SCALING) { - auto is_row_scaled_nvfp4 = [](const Tensor *t) -> bool { return t->row_scaled_nvfp4; }; - const bool all_row_scaled_nvfp4 = std::all_of(output_list.begin(), output_list.end(), - is_row_scaled_nvfp4); - const bool any_row_scaled_nvfp4 = std::any_of(output_list.begin(), output_list.end(), - is_row_scaled_nvfp4); + auto is_row_scaled_nvfp4 = [](const Tensor* t) -> bool { return t->row_scaled_nvfp4; }; + const bool all_row_scaled_nvfp4 = + std::all_of(output_list.begin(), output_list.end(), is_row_scaled_nvfp4); + const bool any_row_scaled_nvfp4 = + std::any_of(output_list.begin(), output_list.end(), is_row_scaled_nvfp4); NVTE_CHECK(all_row_scaled_nvfp4 == any_row_scaled_nvfp4, "Output tensors are not consistent between row-scaling and tensor-scaling."); if (all_row_scaled_nvfp4) { - dispatch::nvfp4::group_row_amax(*input_tensor, /*noop=*/nullptr, output_list, - split_sections, num_tensors, stream); + dispatch::nvfp4::group_row_amax(*input_tensor, /*noop=*/nullptr, output_list, split_sections, + num_tensors, stream); return; } } From 32c57f2f14c95da9157dbc8a5dc0ecaeefd96661 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Mon, 26 Oct 2026 15:38:43 -0700 Subject: [PATCH 21/24] Fix missing include Review suggestion from @greptile-apps Signed-off-by: Tim Moon --- tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu b/tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu index c4b05ff5f5a..226e5c4c2ed 100644 --- a/tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu +++ b/tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu @@ -17,7 +17,7 @@ #include #include -#include +#include #include "../test_common.h" #include "transformer_engine/transformer_engine.h" From 07f45a161245cfd7637f0ca73745556b4080edd8 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Mon, 26 Oct 2026 17:02:13 -0700 Subject: [PATCH 22/24] Add missing header include Signed-off-by: Tim Moon --- transformer_engine/common/cast/cast_grouped.cu | 1 + 1 file changed, 1 insertion(+) diff --git a/transformer_engine/common/cast/cast_grouped.cu b/transformer_engine/common/cast/cast_grouped.cu index 11314eaa6ca..0a33f62d5fb 100644 --- a/transformer_engine/common/cast/cast_grouped.cu +++ b/transformer_engine/common/cast/cast_grouped.cu @@ -8,6 +8,7 @@ #include #include #include +#include #include #include From 67e7ea669710ce5f6b970ca0c4071147c0d98f9e Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Fri, 6 Nov 2026 05:36:47 -0800 Subject: [PATCH 23/24] Address review feedback on row-scaled NVFP4 grouped GEMM - Move row-scaled handling out of general_grouped_gemm_for_grouped_tensor into a dedicated _grouped_gemm_row_scaled_nvfp4 helper. - Compute amax offsets for all NVFP4 grouped tensors and slice amax through a single code path. - Index Vec elements directly instead of reinterpret_cast. - Rename the grouped-tensor GEMM test gate and split skip reasons. - Check that first_dims is a CUDA tensor with num_tensors entries. Signed-off-by: Cael Ling --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 33 ++-- .../nvfp4/group_quantize_transpose_nvfp4.cuh | 7 +- .../pytorch/cpp_extensions/gemm.py | 182 ++++++++++-------- .../pytorch/csrc/extensions/cast.cpp | 8 +- .../tensor/storage/grouped_tensor_storage.py | 47 +++-- 5 files changed, 152 insertions(+), 125 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 9f744352756..5a6d93f4ee5 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -845,17 +845,14 @@ def test_nvfp4_ue5m3_gemm_leaves_operands_unswizzled(M: int, K: int, N: int): assert rel_err < 5e-3, f"relative error {rel_err:.2e} is too large" -def _row_scaled_grouped_tensor_gemm_supported() -> bool: - if not torch.cuda.is_available(): - return False - if torch.cuda.get_device_capability() < (10, 0): - return False - return tex.get_cublasLt_version() >= 130400 - - -_reason_for_no_grouped_tensor_gemm = ( - "Row-scaled grouped-tensor NVFP4 GEMM requires Blackwell (SM100) and cuBLAS 13.4+." -) +_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( @@ -914,7 +911,7 @@ 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 _row_scaled_grouped_tensor_gemm_supported(): + if not _grouped_tensor_gemm_supported: pytest.skip(_reason_for_no_grouped_tensor_gemm) torch.manual_seed(7) @@ -962,7 +959,7 @@ 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 _row_scaled_grouped_tensor_gemm_supported(): + if not _grouped_tensor_gemm_supported: pytest.skip(_reason_for_no_grouped_tensor_gemm) torch.manual_seed(11) @@ -1032,7 +1029,7 @@ 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 _row_scaled_grouped_tensor_gemm_supported(): + if not _grouped_tensor_gemm_supported: pytest.skip(_reason_for_no_grouped_tensor_gemm) torch.manual_seed(13) @@ -1067,7 +1064,7 @@ 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 _row_scaled_grouped_tensor_gemm_supported(): + if not _grouped_tensor_gemm_supported: pytest.skip(_reason_for_no_grouped_tensor_gemm) torch.manual_seed(17) @@ -1098,7 +1095,7 @@ 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 _row_scaled_grouped_tensor_gemm_supported(): + if not _grouped_tensor_gemm_supported: pytest.skip(_reason_for_no_grouped_tensor_gemm) torch.manual_seed(23) @@ -1130,7 +1127,7 @@ 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 _row_scaled_grouped_tensor_gemm_supported(): + if not _grouped_tensor_gemm_supported: pytest.skip(_reason_for_no_grouped_tensor_gemm) torch.manual_seed(29) @@ -1182,7 +1179,7 @@ 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 _row_scaled_grouped_tensor_gemm_supported(): + if not _grouped_tensor_gemm_supported: pytest.skip(_reason_for_no_grouped_tensor_gemm) torch.manual_seed(31) 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 ce91ab33a16..190b7b3e21a 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -1573,12 +1573,11 @@ __device__ __forceinline__ void group_row_scaled_cast_body( ScalesVec &scales_vec = *reinterpret_cast(sSFrowwise[row]); const size_t local_row = static_cast(local_row_base) + row; if constexpr (WITH_GEMM_SWIZZLED_SCALES) { - const ScaleType *scales = reinterpret_cast(&scales_vec); #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[k]; + scales_vec.data.elt[k]; } } else { const size_t scale_idx = local_row * scale_stride_row + scales_block_offset_X_rowwise; @@ -1595,11 +1594,11 @@ __device__ __forceinline__ void group_row_scaled_cast_body( ScalesVec &scales_vec = *reinterpret_cast(sSFcolwise[row_tr]); const size_t global_col = static_cast(block_offset_X) + row_tr; if constexpr (WITH_GEMM_SWIZZLED_SCALES) { - const ScaleType *scales = reinterpret_cast(&scales_vec); #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[k]; + scale_stride_col)] = + scales_vec.data.elt[k]; } } else { const size_t scale_idx = global_col * scale_stride_col + local_block_offset_M; diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 0b18c7779c0..feafa851804 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1139,50 +1139,19 @@ 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.") - def _is_row_scaled(operand) -> bool: - return isinstance(operand, GroupedTensorStorage) and operand.row_scaled_nvfp4 - - def _list_has_row_scaled(operand) -> bool: - return isinstance(operand, (list, tuple)) and any( - getattr(t, "_row_scaled_nvfp4", False) for t in operand - ) - - if _list_has_row_scaled(A) or _list_has_row_scaled(B): - raise NotImplementedError( - "Row-scaled NVFP4 operands passed as a list are not supported; pass a GroupedTensor." + 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, ) - row_scaled = _is_row_scaled(A) or _is_row_scaled(B) - row_scaled_accumulate = accumulate - post_scale_rows = None # per-row outer scale -> output rows (from B) - post_scale_cols = None # per-group per-col outer scale -> output cols (from A) - if row_scaled: - 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." - ) - if _is_row_scaled(A): - A = A.copy() - if transa: - post_scale_cols = A.amax - A.amax = post_scale_cols.new_ones(A.num_tensors) - else: - post_scale_cols = A.columnwise_amax - A.columnwise_amax = post_scale_cols.new_ones(A.num_tensors) - A.row_scaled_nvfp4 = False - if _is_row_scaled(B): - B = B.copy() - if transb: - post_scale_rows = B.columnwise_amax - B.columnwise_amax = post_scale_rows.new_ones(B.num_tensors) - else: - post_scale_rows = B.amax - B.amax = post_scale_rows.new_ones(B.num_tensors) - B.row_scaled_nvfp4 = False - accumulate = False - def _is_fp8_blockwise(operand) -> bool: if isinstance(operand, (list, tuple)): return any(isinstance(t, Float8BlockwiseQTensorStorage) for t in operand) @@ -1248,51 +1217,110 @@ def _is_fp8_blockwise(operand) -> bool: sm_count = get_sm_count() sm_count = sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))) - if not row_scaled: - return grouped_gemm_impl( - A, - transa, - B, - transb, - out, - bias, - bias_scale, - alpha, - beta, - workspace_setup, - workspace_cublas, - use_split_accumulator, - sm_count, + return grouped_gemm_impl( + A, + transa, + B, + transb, + out, + bias, + bias_scale, + alpha, + beta, + workspace_setup, + workspace_cublas, + 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 - grouped_gemm_impl( + general_grouped_gemm_for_grouped_tensor( A, - transa, B, - transb, scratch, - None, - None, - alpha, - beta, - workspace_setup, - workspace_cublas, - use_split_accumulator, - sm_count, + layout=layout, + accumulate=False, + use_split_accumulator=use_split_accumulator, + alpha=alpha, ) def _store(dst: torch.Tensor, src: torch.Tensor) -> None: - if row_scaled_accumulate: + 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 @@ -1300,11 +1328,11 @@ def _store(dst: torch.Tensor, src: torch.Tensor) -> None: scaled = scratch[i] rows = scaled.shape[-2] cols = scaled.shape[-1] - if post_scale_rows is not None: - row_scale_i = post_scale_rows[row_offset : row_offset + rows] + 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 post_scale_cols is not None: - col_scale_i = post_scale_cols[col_offset : col_offset + cols] + 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 @@ -1317,11 +1345,11 @@ def _store(dst: torch.Tensor, src: torch.Tensor) -> None: last_dim = out.logical_shape[-1] live_rows = out.logical_shape[0] flat = scratch.rowwise_data.view(-1, last_dim)[:live_rows] - if post_scale_rows is not None: - flat.mul_(post_scale_rows.view(-1, 1)) - if post_scale_cols is not None: + 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 = post_scale_cols.view(num_tensors, last_dim) + 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) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 3d61dc525f8..7010e615bdd 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/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 334f158cd8b..ea28b4c6bed 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -1116,22 +1116,29 @@ def split_into_quantized_tensors( 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 - nvfp4_columnwise_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] - ccum = 0 - nvfp4_columnwise_amax_offsets = [0] - for i in range(self.num_tensors): - cum += math.prod(self.tensor_shapes[i][:-1]) - nvfp4_rowwise_amax_offsets.append(cum) - ccum += self.tensor_shapes[i][-1] - nvfp4_columnwise_amax_offsets.append(ccum) + + # 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 @@ -1327,20 +1334,12 @@ def split_into_quantized_tensors( ) 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 = self.amax[amax_start:amax_end] - else: - amax_rowwise = self.amax[i : i + 1] + amax_start, amax_end = nvfp4_rowwise_amax_offsets[i : i + 2] + amax_rowwise = self.amax[amax_start:amax_end] if self.columnwise_amax is not None: - if nvfp4_columnwise_amax_offsets is not None: - camax_start = nvfp4_columnwise_amax_offsets[i] - camax_end = nvfp4_columnwise_amax_offsets[i + 1] - amax_columnwise = self.columnwise_amax[camax_start:camax_end] - else: - amax_columnwise = self.columnwise_amax[i : i + 1] + amax_start, amax_end = nvfp4_columnwise_amax_offsets[i : i + 2] + amax_columnwise = self.columnwise_amax[amax_start:amax_end] if quantizer.internal: nvfp4_tensor_class = NVFP4TensorStorage From f3aedc9f1355f1fcd69e418f0c5480987d1ed6f5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 6 Nov 2026 13:57:12 +0000 Subject: [PATCH 24/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh | 5 ++--- transformer_engine/pytorch/cpp_extensions/gemm.py | 3 ++- 2 files changed, 4 insertions(+), 4 deletions(-) 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 190b7b3e21a..10b5bddc325 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -1596,9 +1596,8 @@ __device__ __forceinline__ void group_row_scaled_cast_body( 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]; + 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; diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index feafa851804..8a8bb81e6ea 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1270,7 +1270,8 @@ def _grouped_gemm_row_scaled_nvfp4( 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." + "Row-scaled NVFP4 operands passed as a list are not supported; pass a" + " GroupedTensor." ) transa = layout[0] == "T" transb = layout[1] == "T"