[Common, PyTorch] Fuse NVFP4 scale swizzling on SM120 and SM121

tests/cpp/operator/test_cast_nvfp4_transpose.cu

@@ -1681,24 +1681,122 @@ TEST(NVFP4UE5M3ReferenceTest, Grouped) {
 }
 #endif
 
-TEST(CastNVFP4RHTTest, RejectsSwizzledScalesOnSM12x) {
+TEST(CastNVFP4RHTTest, RejectsUnalignedSwizzledScalesOnSM12x) {
     const int sm_arch = getDeviceComputeCapability();
     if (sm_arch != 120 && sm_arch != 121) {
         GTEST_SKIP();
     }
 
-    Tensor input("input", std::vector<size_t>{128, 128}, DType::kBFloat16);
-    Tensor output("output", std::vector<size_t>{128, 128}, DType::kFloat4E2M1,
+    // In-kernel GEMM-swizzled scales require 128-aligned rows and cols.
+    Tensor input("input", std::vector<size_t>{64, 128}, DType::kBFloat16);
+    Tensor output("output", std::vector<size_t>{64, 128}, DType::kFloat4E2M1,
                   false, true, NVTE_NVFP4_1D_SCALING);
     Tensor hadamard("hadamard", std::vector<size_t>{16, 16}, DType::kBFloat16);
     output.set_with_gemm_swizzled_scales(true);
 
     try {
         nvte_quantize_with_hadamard_transform(input.data(), output.data(), hadamard.data(),
                                               nullptr, nullptr);
-        FAIL() << "Expected the SM12x RHT path to reject swizzled scales";
+        FAIL() << "Expected the SM12x RHT path to reject swizzled scales for 64 rows";
     } catch (const std::runtime_error& e) {
-        EXPECT_NE(std::string(e.what()).find("does not support GEMM-swizzled scales"),
-                  std::string::npos);
+        EXPECT_NE(std::string(e.what()).find("require rows and cols to be multiples of 128"),
+                  std::string::npos) << e.what();
+    }
+}
+
+TEST(CastNVFP4RHTTest, SwizzledScalesMatchCompactOnSM12x) {
+    const int sm_arch = getDeviceComputeCapability();
+    if (sm_arch != 120 && sm_arch != 121) {
+        GTEST_SKIP();
+    }
+
+    constexpr size_t rows = 256;
+    constexpr size_t cols = 384;
+    const std::vector<size_t> shape{rows, cols};
+
+    Tensor input("input", shape, DType::kBFloat16);
+    fillUniform(&input);
+    const bf16 *input_data = input.rowwise_cpu_dptr<bf16>();
+    float amax = 0.0f;
+    for (size_t i = 0; i < rows * cols; ++i) {
+        amax = fmaxf(amax, fabsf(static_cast<float>(input_data[i])));
+    }
+
+    // Orthonormal 16-point Sylvester Hadamard matrix: H[i][j] = (-1)^popcount(i & j) / 4.
+    Tensor hadamard("hadamard", std::vector<size_t>{16, 16}, DType::kBFloat16);
+    bf16 *hadamard_data = hadamard.rowwise_cpu_dptr<bf16>();
+    for (unsigned i = 0; i < 16; ++i) {
+        for (unsigned j = 0; j < 16; ++j) {
+            hadamard_data[i * 16 + j] =
+                static_cast<bf16>(__builtin_popcount(i & j) % 2 ? -0.25f : 0.25f);
+        }
+    }
+    hadamard.from_cpu();
+
+    // Rowwise scales form a [rows, cols/16] matrix; columnwise scales belong to the transposed
+    // data and form a [cols, rows/16] matrix. With 128-aligned dims neither buffer is padded.
+    const auto dims = get_scale_tensor_dims(rows, cols, 1, 16);
+    const auto dims_t = get_scale_tensor_dims(cols, rows, 1, 16);
+    for (const auto &d : {dims, dims_t}) {
+        ASSERT_EQ(d[0], d[2]);
+        ASSERT_EQ(d[1], d[3]);
+    }
+
+    // Permutes a compact [M, K] scale matrix into cuBLAS's block-scaling layout (128x4 tiles
+    // stored as 32 rows of 16 bytes), as compute_ref_swizzle<128, 4, true> in test_swizzle.cu.
+    const auto swizzle_ref = [](const uint8_t *scales, const size_t M, const size_t K) {
+        std::vector<uint8_t> out(M * K);
+        for (size_t m = 0; m < M; ++m) {
+            for (size_t k = 0; k < K; ++k) {
+                const size_t tile_offset = (m / 128) * 128 * K + (k / 4) * 512;
+                out[tile_offset + (m % 32) * 16 + (m % 128) / 32 * 4 + k % 4] = scales[m * K + k];
+            }
+        }
+        return out;
+    };
+    const auto check_direction = [&](const std::string &name, const fp4e2m1 *data,
+                                     const fp4e2m1 *swizzled_data, const uint8_t *scales,
+                                     const uint8_t *swizzled_scales, const size_t M,
+                                     const size_t K) {
+        compareResults(name + "_data", reinterpret_cast<const uint8_t *>(swizzled_data),
+                       reinterpret_cast<const uint8_t *>(data), rows * cols / 2);
+        // Buffers start zeroed, so a block skipped by both kernels would otherwise go unnoticed.
+        ASSERT_TRUE(std::none_of(scales, scales + M * K, [](uint8_t s) { return s == 0; }))
+            << name << " compact scales contain unwritten (zero) entries";
+        compareResults(name + "_scales", swizzled_scales, swizzle_ref(scales, M, K).data(),
+                       M * K);
+    };
+
+    for (const bool rowwise : {false, true}) {
+        SCOPED_TRACE(rowwise ? "rowwise+columnwise" : "columnwise-only");
+        Tensor compact("compact", shape, DType::kFloat4E2M1, rowwise, true,
+                       NVTE_NVFP4_1D_SCALING);
+        Tensor swizzled("swizzled", shape, DType::kFloat4E2M1, rowwise, true,
+                        NVTE_NVFP4_1D_SCALING);
+        swizzled.set_with_gemm_swizzled_scales(true);
+        for (Tensor *output : {&compact, &swizzled}) {
+            if (rowwise) {
+                output->set_amax(amax);
+            }
+            // Each RHT output sums 16 terms of +-x/4, so 4 * amax bounds the columnwise amax.
+            output->set_tensor_amax_columnwise(4.0f * amax);
+            nvte_quantize_with_hadamard_transform(input.data(), output->data(), hadamard.data(),
+                                                  nullptr, nullptr);
+        }
+        ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+        ASSERT_EQ(cudaGetLastError(), cudaSuccess);
+        compact.to_cpu();
+        swizzled.to_cpu();
+
+        check_direction("columnwise", compact.columnwise_cpu_dptr<fp4e2m1>(),
+                        swizzled.columnwise_cpu_dptr<fp4e2m1>(),
+                        compact.columnwise_cpu_scale_inv_ptr<uint8_t>(),
+                        swizzled.columnwise_cpu_scale_inv_ptr<uint8_t>(), dims_t[0], dims_t[1]);
+        if (rowwise) {
+            check_direction("rowwise", compact.rowwise_cpu_dptr<fp4e2m1>(),
+                            swizzled.rowwise_cpu_dptr<fp4e2m1>(),
+                            compact.rowwise_cpu_scale_inv_ptr<uint8_t>(),
+                            swizzled.rowwise_cpu_scale_inv_ptr<uint8_t>(), dims[0], dims[1]);
+        }
     }
 }

tests/pytorch/nvfp4/test_nvfp4_rht_quantize_swizzle_fusion.py

@@ -208,23 +208,25 @@ def _make_swizzle_fusion_quantizer() -> NVFP4Quantizer:
 
 @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
 @pytest.mark.parametrize(
-    "M, N, expected_swizzled",
+    "M, N, expected_swizzled, expected_swizzled_sm12x",
     [
         # Eligible: rows%64==0 AND cols%128==0 -> the fused RHT cast-fusion
         # kernel can bake the GEMM-swizzled SF in directly.
-        (64, 128, True),
-        (128, 256, True),
+        (64, 128, True, False),
+        (128, 256, True, True),
         # Ineligible (cols%128 != 0) -> the fused kernel is gated off, so
         # create_tensor reports swizzled=False (the fused kernel did not bake
         # it in). This is the case hit in production at irregular shapes like
         # (8192, 11328) where 11328%128==64.
-        (64, 144, False),
-        (128, 144, False),
+        (64, 144, False, False),
+        (128, 144, False, False),
         # Ineligible (rows%64 != 0) -> same gating.
-        (48, 128, False),
+        (48, 128, False, False),
     ],
 )
-def test_nvfp4_rht_swizzle_fusion_shape_gate(M: int, N: int, expected_swizzled: bool) -> None:
+def test_nvfp4_rht_swizzle_fusion_shape_gate(
+    M: int, N: int, expected_swizzled: bool, expected_swizzled_sm12x: bool
+) -> None:
     """The fused-kernel shape gate must reflect cast-fusion eligibility.
 
     Only ``row_cast_col_hadamard_transform_cast_fusion.cu`` can bake
@@ -243,6 +245,10 @@ def test_nvfp4_rht_swizzle_fusion_shape_gate(M: int, N: int, expected_swizzled:
     ineligible shapes too and flips the flag back to True; see
     ``test_nvfp4_rht_swizzle_fusion_end_to_end_swizzled``.)
     """
+    # The SM12x 1D path needs full 128-row scale tiles; TMEM architectures
+    # retain their existing 64-row eligibility.
+    if torch.cuda.get_device_capability() in ((12, 0), (12, 1)):
+        expected_swizzled = expected_swizzled_sm12x
     quantizer = _make_swizzle_fusion_quantizer()
     out = quantizer.make_empty([M, N], dtype=torch.bfloat16)
     assert out._with_gemm_swizzled_scales is expected_swizzled, (

tests/pytorch/nvfp4/test_nvfp4_sm12x_swizzle_fusion.py

@@ -0,0 +1,110 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""SM12x scale-layout fidelity and allocation gates for fused NVFP4 quantization."""
+
+import pytest
+import torch
+
+import transformer_engine.pytorch as te
+from transformer_engine.pytorch import NVFP4Quantizer
+
+from nvfp4_utils import swizzle_nvfp4_scale
+
+
+@pytest.fixture(autouse=True)
+def require_sm12x():
+    """These kernels require an SM120 or SM121 device."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    if torch.cuda.get_device_capability() not in ((12, 0), (12, 1)):
+        pytest.skip("SM120/SM121 is required")
+    available, reason = te.is_nvfp4_available(return_reason=True)
+    if not available:
+        pytest.skip(reason)
+
+
+@pytest.mark.parametrize("shape", [(128, 128), (256, 512), (384, 256)])
+@pytest.mark.parametrize("direction", ["row", "column", "both"])
+@pytest.mark.parametrize("kind", ["1d", "2d", "rht"])
+@pytest.mark.parametrize("stochastic_rounding", [False, True])
+@pytest.mark.parametrize("scale_dtype", [te.DType.kFloat8E4M3, te.DType.kFloat8UE5M3])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+def test_sm12x_swizzled_outputs_match_compact_reference(
+    shape, direction, kind, stochastic_rounding, scale_dtype, dtype
+):
+    """Check FP4 bytes, both amaxes and an independent scale permutation."""
+    if kind == "rht" and dtype != torch.bfloat16:
+        pytest.skip("NVFP4 RHT only supports BF16 input")
+    torch.manual_seed(1234)
+    x = torch.randn(shape, device="cuda", dtype=dtype)
+    compact = NVFP4Quantizer(
+        rowwise=direction != "column",
+        columnwise=direction != "row",
+        with_2d_quantization=kind == "2d",
+        with_rht=kind == "rht",
+        with_post_rht_amax=kind == "rht",
+        with_random_sign_mask=False,
+        stochastic_rounding=stochastic_rounding,
+        scale_dtype=scale_dtype,
+    )
+    compact.optimize_for_gemm = False
+    fused = compact.copy()
+    fused.optimize_for_gemm = True
+    # Test allocation separately: the fallback also sets this flag after quantization.
+    allocated = fused.make_empty(shape, dtype=x.dtype)
+    assert allocated._with_gemm_swizzled_scales
+
+    torch.manual_seed(5678)
+    expected = compact(x)
+    torch.manual_seed(5678)
+    actual = fused(x)
+    for columnwise, prefix in [(False, "rowwise"), (True, "columnwise")]:
+        ref_data = getattr(expected, f"_{prefix}_data")
+        if ref_data is None:
+            assert getattr(actual, f"_{prefix}_data") is None
+            continue
+        assert torch.equal(getattr(actual, f"_{prefix}_data"), ref_data)
+        suffix = "columnwise" if columnwise else "rowwise"
+        assert torch.equal(getattr(actual, f"_amax_{suffix}"), getattr(expected, f"_amax_{suffix}"))
+        reference_scales = swizzle_nvfp4_scale(
+            *shape, getattr(expected, f"_{prefix}_scale_inv"), columnwise=columnwise
+        )
+        assert torch.equal(
+            getattr(actual, f"_{prefix}_scale_inv").view(torch.uint8),
+            reference_scales.view(torch.uint8),
+        )
+
+
+@pytest.mark.parametrize("kind", ["1d", "2d", "rht"])
+@pytest.mark.parametrize(
+    "shape,eligible", [((128, 128), True), ((64, 128), False), ((128, 192), False)]
+)
+def test_sm12x_in_kernel_swizzle_shape_gate(kind, shape, eligible):
+    """Partial scale tiles must retain compact allocation and the existing fallback."""
+    q = NVFP4Quantizer(
+        with_2d_quantization=kind == "2d",
+        with_rht=kind == "rht",
+        with_post_rht_amax=kind == "rht",
+    )
+    q.optimize_for_gemm = True
+    out = q.make_empty(shape, dtype=torch.bfloat16)
+    assert out._with_gemm_swizzled_scales is eligible
+
+
+@pytest.mark.parametrize("with_rht", [False, True])
+@pytest.mark.parametrize("special_recipe", ["row_scaled", "4over6"])
+def test_sm12x_special_recipes_keep_compact_allocation(with_rht, special_recipe):
+    """Unsupported recipes must be rejected before selecting the RHT scale layout."""
+    q = NVFP4Quantizer(
+        rowwise=True,
+        columnwise=False,
+        with_rht=with_rht,
+        with_post_rht_amax=with_rht,
+        row_scaled_nvfp4=special_recipe == "row_scaled",
+        nvfp4_use_4over6=special_recipe == "4over6",
+    )
+    q.optimize_for_gemm = True
+    out = q.make_empty((128, 128), dtype=torch.bfloat16)
+    assert not out._with_gemm_swizzled_scales