Skip to content

[Common] Optimized kernels for swizzling and remove obsolete narrow scaled kernels #3682

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
vthumbe1503 wants to merge 2 commits into
base: main
Choose a base branch
from
+561 −718
Open

[Common] Optimized kernels for swizzling and remove obsolete narrow scaled kernels #3682

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8ae90bc
Optimize scale swizzle kernels
vthumbe1503 Oct 23, 2026
d881f0d
Merge branch 'main' into swizzle-per-tile-default
vthumbe1503 Oct 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 0 additions & 6 deletions docs/envvars.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -351,12 +351,6 @@ Kernel Configuration
:Default: ``0``
:Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. Existing transpose operations select their static fallback automatically. NVFP4 4over6 quantization, fused softmax, and normalization paths require their corresponding ``NVTE_BUILD_LEGACY_STATIC_NVFP4``, ``NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX``, or ``NVTE_BUILD_LEGACY_STATIC_NORM`` CMake option to have been enabled when the library was built; otherwise no static fallback is available.

.. envvar:: NVTE_MXFP8_PER_TILE_SWIZZLE

:Type: ``int`` (0 or 1)
:Default: ``0``
:Description: Swizzle MXFP8 scaling factors for GEMM with kernels that give each 128-row tile of scaling factors (row-wise) or each 128-column by 32-row block (column-wise) its own thread block, instead of batching 32 tiles in one block. This lowers the latency of swizzling the scaling factors of tensors with few tiles, such as activations.

.. envvar:: NVTE_USE_CUTLASS_GROUPED_GEMM

:Type: ``int`` (0 or 1)
Expand Down
212 changes: 176 additions & 36 deletions tests/cpp/operator/test_multi_swizzle.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,18 +51,25 @@ void compute_ref_swizzle(const uint8_t *h_input, uint8_t *h_output,
}
}

static void zero_scale_inv_padding(uint8_t *buf,
size_t padded_rows, size_t padded_cols,
size_t orig_rows, size_t orig_cols) {
static void set_scale_inv_padding(uint8_t *buf,
size_t padded_rows, size_t padded_cols,
size_t orig_rows, size_t orig_cols,
uint8_t value) {
for (size_t r = 0; r < padded_rows; ++r) {
for (size_t c = 0; c < padded_cols; ++c) {
if (r >= orig_rows || c >= orig_cols) {
buf[r * padded_cols + c] = 0;
buf[r * padded_cols + c] = value;
}
}
}
}

static void zero_scale_inv_padding(uint8_t *buf,
size_t padded_rows, size_t padded_cols,
size_t orig_rows, size_t orig_cols) {
set_scale_inv_padding(buf, padded_rows, padded_cols, orig_rows, orig_cols, 0);
}

// ===================================================================
// Multi-tensor swizzle test
// ===================================================================
Expand Down Expand Up @@ -141,6 +148,159 @@ void performTestMultiTensorSwizzle(const int num_tensors, const size_t M, const
}
}

void performTestMultiTensorSwizzleNVFP4(const int num_tensors, const size_t M, const size_t K,
bool rowwise) {
using namespace test;
constexpr size_t BLOCK_SIZE = 16;
std::vector<std::unique_ptr<Tensor>> inputs;
std::vector<std::unique_ptr<Tensor>> outputs;
std::vector<NVTETensor> input_handles;
std::vector<NVTETensor> output_handles;

for (int i = 0; i < num_tensors; ++i) {
auto input = std::make_unique<Tensor>("input_" + std::to_string(i),
std::vector<size_t>{M, K}, DType::kFloat4E2M1, rowwise,
!rowwise, NVTE_NVFP4_1D_SCALING);
auto output = std::make_unique<Tensor>("output_" + std::to_string(i),
std::vector<size_t>{M, K}, DType::kFloat4E2M1, rowwise,
!rowwise, NVTE_NVFP4_1D_SCALING);
output->set_with_gemm_swizzled_scales(true);
if (rowwise) {
input->fill_uniform_rowwise_scale_inv();
} else {
input->fill_uniform_columnwise_scale_inv();
}
const NVTEShape scale_shape =
rowwise ? input->rowwise_scale_inv_shape() : input->columnwise_scale_inv_shape();
uint8_t *input_scales = rowwise ? input->rowwise_cpu_scale_inv_ptr<uint8_t>()
: input->columnwise_cpu_scale_inv_ptr<uint8_t>();
const size_t original_rows = rowwise ? M : K;
const size_t original_cols =
rowwise ? (K + BLOCK_SIZE - 1) / BLOCK_SIZE : (M + BLOCK_SIZE - 1) / BLOCK_SIZE;
set_scale_inv_padding(input_scales, scale_shape.data[0], scale_shape.data[1], original_rows,
original_cols, 0xFF);
input->from_cpu();
input_handles.push_back(input->data());
output_handles.push_back(output->data());
inputs.emplace_back(std::move(input));
outputs.emplace_back(std::move(output));
}

nvte_multi_tensor_swizzle_scaling_factors(input_handles.data(), output_handles.data(),
num_tensors, 0);
NVTE_CHECK_CUDA(cudaDeviceSynchronize());
for (int i = 0; i < num_tensors; ++i) {
outputs[i]->to_cpu();
const NVTEShape scale_shape =
rowwise ? inputs[i]->rowwise_scale_inv_shape() : inputs[i]->columnwise_scale_inv_shape();
const uint8_t *input_scales = rowwise ? inputs[i]->rowwise_cpu_scale_inv_ptr<uint8_t>()
: inputs[i]->columnwise_cpu_scale_inv_ptr<uint8_t>();
const uint8_t *actual = rowwise ? outputs[i]->rowwise_cpu_scale_inv_ptr<uint8_t>()
: outputs[i]->columnwise_cpu_scale_inv_ptr<uint8_t>();
std::vector<uint8_t> reference_input(
input_scales, input_scales + scale_shape.data[0] * scale_shape.data[1]);
const size_t original_rows = rowwise ? M : K;
const size_t original_cols =
rowwise ? (K + BLOCK_SIZE - 1) / BLOCK_SIZE : (M + BLOCK_SIZE - 1) / BLOCK_SIZE;
zero_scale_inv_padding(reference_input.data(), scale_shape.data[0], scale_shape.data[1],
original_rows, original_cols);
std::vector<uint8_t> reference(scale_shape.data[0] * scale_shape.data[1]);
compute_ref_swizzle<128, 4, true>(reference_input.data(), reference.data(), scale_shape.data[0],
scale_shape.data[1]);
compareResults("multi_nvfp4_swizzle", actual, reference.data(), reference.size());
}
}

TEST(MultiTensorSwizzleTest, MXFP8ColumnwiseBatchIgnoresUnusedRowwiseBuffer) {
using namespace test;
constexpr int num_tensors = 2;
const std::vector<size_t> shape{256, 128};
std::vector<std::unique_ptr<Tensor>> inputs;
std::vector<std::unique_ptr<Tensor>> outputs;
std::vector<NVTETensor> input_handles;
std::vector<NVTETensor> output_handles;

for (int i = 0; i < num_tensors; ++i) {
const bool has_rowwise_scales = i == 0;
auto input = std::make_unique<Tensor>("input_" + std::to_string(i), shape,
DType::kFloat8E4M3, has_rowwise_scales, true,
NVTE_MXFP8_1D_SCALING);
auto output = std::make_unique<Tensor>("output_" + std::to_string(i), shape,
DType::kFloat8E4M3, has_rowwise_scales, true,
NVTE_MXFP8_1D_SCALING);
if (has_rowwise_scales) input->fill_uniform_rowwise_scale_inv();
input->fill_uniform_columnwise_scale_inv();
output->set_with_gemm_swizzled_scales(true);
input->from_cpu();
input_handles.push_back(input->data());
output_handles.push_back(output->data());
inputs.emplace_back(std::move(input));
outputs.emplace_back(std::move(output));
}

nvte_multi_tensor_swizzle_scaling_factors(input_handles.data(), output_handles.data(),
num_tensors, 0);
NVTE_CHECK_CUDA(cudaDeviceSynchronize());
for (int i = 0; i < num_tensors; ++i) {
outputs[i]->to_cpu();
const NVTEShape scale_shape = inputs[i]->columnwise_scale_inv_shape();
std::vector<uint8_t> reference(scale_shape.data[0] * scale_shape.data[1]);
compute_ref_swizzle<128, 4, false>(
inputs[i]->columnwise_cpu_scale_inv_ptr<uint8_t>(), reference.data(), scale_shape.data[1],
scale_shape.data[0]);
compareResults("multi_mixed_buffers_columnwise_" + std::to_string(i),
outputs[i]->columnwise_cpu_scale_inv_ptr<uint8_t>(), reference.data(),
reference.size());
}
}

TEST(MultiTensorSwizzleTest, NVFP4FullWidthChunkLaunchSplitting) {
for (const bool rowwise : {false, true}) {
performTestMultiTensorSwizzleNVFP4(4, 160, 96, rowwise);
performTestMultiTensorSwizzleNVFP4(4, 128, 128, rowwise);
performTestMultiTensorSwizzleNVFP4(65, 128, 128, rowwise);
}
}

TEST(MultiTensorSwizzleTest, RowwiseMixedSharedMemoryFallback) {
using namespace test;
// The first tensor fits the full-width shared-memory limit, but the second does not.
const std::vector<std::pair<size_t, size_t>> shapes{{128, 7168}, {128, 8064}};
std::vector<std::unique_ptr<Tensor>> inputs;
std::vector<std::unique_ptr<Tensor>> outputs;
std::vector<NVTETensor> input_handles;
std::vector<NVTETensor> output_handles;
for (size_t i = 0; i < shapes.size(); ++i) {
const auto [M, K] = shapes[i];
auto input = std::make_unique<Tensor>("input_" + std::to_string(i),
std::vector<size_t>{M, K}, DType::kFloat8E4M3, true,
false, NVTE_MXFP8_1D_SCALING);
auto output = std::make_unique<Tensor>("output_" + std::to_string(i),
std::vector<size_t>{M, K}, DType::kFloat8E4M3, true,
false, NVTE_MXFP8_1D_SCALING);
output->set_with_gemm_swizzled_scales(true);
input->fill_uniform_rowwise_scale_inv();
input->from_cpu();
input_handles.push_back(input->data());
output_handles.push_back(output->data());
inputs.emplace_back(std::move(input));
outputs.emplace_back(std::move(output));
}

nvte_multi_tensor_swizzle_scaling_factors(
input_handles.data(), output_handles.data(), static_cast<int>(inputs.size()), 0);
NVTE_CHECK_CUDA(cudaDeviceSynchronize());
for (size_t i = 0; i < inputs.size(); ++i) {
outputs[i]->to_cpu();
const NVTEShape scale_shape = inputs[i]->rowwise_scale_inv_shape();
std::vector<uint8_t> reference(scale_shape.data[0] * scale_shape.data[1]);
compute_ref_swizzle<128, 4, true>(inputs[i]->rowwise_cpu_scale_inv_ptr<uint8_t>(),
reference.data(), scale_shape.data[0], scale_shape.data[1]);
compareResults("multi_rowwise_fallback", outputs[i]->rowwise_cpu_scale_inv_ptr<uint8_t>(),
reference.data(), reference.size());
}
}

// ===================================================================
// Multi-tensor unswizzle test (uses single-tensor swizzle to prepare)
// ===================================================================
Expand Down Expand Up @@ -300,22 +460,11 @@ void performTestMultiTensorRoundtrip(const int num_tensors, const size_t M, cons
// Test suites
// ===================================================================

// Sets NVTE_MXFP8_PER_TILE_SWIZZLE while it lives, so that MXFP8 swizzles use the per-tile
// kernels (enable) or the default ones.
class PerTileSwizzle {
public:
explicit PerTileSwizzle(bool enable) {
setenv("NVTE_MXFP8_PER_TILE_SWIZZLE", enable ? "1" : "0", 1);
}
~PerTileSwizzle() { unsetenv("NVTE_MXFP8_PER_TILE_SWIZZLE"); }
};

class MultiTensorSwizzleTestSuite
: public ::testing::TestWithParam<std::tuple<std::tuple<int, size_t, size_t, bool>, bool>> {};
: public ::testing::TestWithParam<std::tuple<int, size_t, size_t, bool>> {};

TEST_P(MultiTensorSwizzleTestSuite, TestMultiTensorSwizzle) {
const auto& [num_tensors, M, K, rowwise] = std::get<0>(GetParam());
const PerTileSwizzle per_tile_swizzle(std::get<1>(GetParam()));
const auto& [num_tensors, M, K, rowwise] = GetParam();
performTestMultiTensorSwizzle(num_tensors, M, K, rowwise);
}

Expand Down Expand Up @@ -343,17 +492,7 @@ TEST_P(MultiTensorRoundtripTestSuite, TestMultiTensorRoundtrip) {

namespace {

// Shapes that exercise the narrow_k kernel (rowwise) / narrow_m kernel (colwise):
// Narrow-K fires when ALL tensors have scale num_tiles_k < TB_DIM (32),
// i.e. padded ceil(K/32) < 128.
// Narrow-M fires analogously for colwise when padded K < 4096
// (since colwise m = K padded to 128, num_tiles_m = m / 128 < 32).
//
// Shapes that bypass narrow and use the regular multi_tensor kernel:
// K >= 4096 makes num_tiles_k >= 32 (rowwise) and num_tiles_m >= 32 (colwise).

std::vector<std::tuple<int, size_t, size_t, bool>> multi_tensor_test_cases = {
// --- Narrow path cases (K small → narrow_k for row, narrow_m for col) ---
// M and K both aligned to 128
{3, 256, 256, true},
{3, 256, 256, false},
Expand All @@ -365,7 +504,7 @@ std::vector<std::tuple<int, size_t, size_t, bool>> multi_tensor_test_cases = {
{3, 192, 256, false},
{2, 64, 256, true},
{2, 64, 256, false},
// Larger narrow K (num_tiles_k = 8, shared mem = 128 KB)
// Larger K
{2, 128, 1024, true},
{2, 128, 1024, false},
// K not divisible by 128
Expand All @@ -380,13 +519,16 @@ std::vector<std::tuple<int, size_t, size_t, bool>> multi_tensor_test_cases = {
{4, 32, 64, true},
{4, 32, 64, false},

// --- Non-narrow path cases (K >= 4096 → regular multi_tensor kernel) ---
{3, 256, 4096, true},
{3, 256, 4096, false},
{2, 128, 8192, true},
{2, 128, 8192, false},
// Rowwise num_tiles_k = 34 selects vec_load_size = 2 (int2 kernel)
{2, 128, 4352, true},
// Wide row-wise scales fall back for an underfilled grid and use the full-width kernel once
// the aggregate grid is large enough.
{2, 2048, 7168, true}, // 32 aggregate row tiles: regular fallback.
{16, 2048, 7168, true}, // 256 aggregate row tiles: rowwise full-width kernel.
// Colwise num_tiles_k = 512 / 32 / 4 = 4 selects vec_load_size = 4 (int4 kernel)
{2, 512, 4096, false},
// MXFP8 swizzle: more tensors than one launch takes, several row-wise tiles per CTA,
Expand All @@ -404,14 +546,12 @@ std::vector<std::tuple<int, size_t, size_t, bool>> multi_tensor_test_cases = {
INSTANTIATE_TEST_SUITE_P(
OperatorTest,
MultiTensorSwizzleTestSuite,
::testing::Combine(::testing::ValuesIn(multi_tensor_test_cases), ::testing::Values(false, true)),
::testing::ValuesIn(multi_tensor_test_cases),
[](const testing::TestParamInfo<MultiTensorSwizzleTestSuite::ParamType>& info) {
const auto& test_case = std::get<0>(info.param);
return "n" + std::to_string(std::get<0>(test_case)) +
"_M" + std::to_string(std::get<1>(test_case)) +
"_K" + std::to_string(std::get<2>(test_case)) +
(std::get<3>(test_case) ? "_row" : "_col") +
(std::get<1>(info.param) ? "_per_tile" : "");
return "n" + std::to_string(std::get<0>(info.param)) +
"_M" + std::to_string(std::get<1>(info.param)) +
"_K" + std::to_string(std::get<2>(info.param)) +
(std::get<3>(info.param) ? "_row" : "_col");
});

INSTANTIATE_TEST_SUITE_P(
Expand Down
Loading
Loading