Quantize NVFP4 1D from registers instead of through TMA staging

tests/cpp/operator/CMakeLists.txt

@@ -15,6 +15,7 @@ add_executable(test_operator
                test_cast_mxfp8_grouped.cu
                test_cast_mxfp8_grouped_scaled_swiglu.cu
                test_cast_nvfp4_transpose.cu
+               test_cast_nvfp4_sr_streams.cu
                test_row_scaled_nvfp4_grouped_amax.cu
                test_row_scaled_nvfp4_grouped_cast.cu
                test_cast_float8blockwise.cu

tests/cpp/operator/test_cast_nvfp4_transpose.cu

@@ -845,7 +845,8 @@ void performTest(float (*OP)(const float),
                  const NVFP4ScalingMode scaling_mode = NVFP4ScalingMode::Block1D,
                  const NVTENVFP44Over6Mode mode = kNVTENVFP44Over6Disabled,
                  const int scale_max = 0,
-                 const bool use_4over6_err_use_fast_math = false) {
+                 const bool use_4over6_err_use_fast_math = false,
+                 const bool rowwise_only = false) {
     using namespace test;
     const bool use_4over6 = mode != kNVTENVFP44Over6Disabled;
 
@@ -863,7 +864,7 @@ void performTest(float (*OP)(const float),
     const bool is_2d_quantization = use_2d_quantization(scaling_mode);
     const bool row_scaled_nvfp4 = scaling_mode == NVFP4ScalingMode::RowScaled1D;
     const bool rowwise = true;
-    const bool columnwise = !row_scaled_nvfp4;
+    const bool columnwise = !row_scaled_nvfp4 && !rowwise_only;
 
     const size_t rows = first_dimension(shape);
     const size_t cols = last_dimension(shape);
@@ -1083,7 +1084,7 @@ void performTest(float (*OP)(const float),
                                         cols,
                                         unpadded_blocks_X,
                                         scales_stride);
-        if (!row_scaled_nvfp4) {
+        if (columnwise) {
             compare_nvfp4_4over6_candidates<ScaleType>("output_t",
                                             output.columnwise_cpu_dptr<fp4e2m1>(),
                                             output.columnwise_cpu_scale_inv_ptr<ScaleType>(),
@@ -1099,7 +1100,7 @@ void performTest(float (*OP)(const float),
     } else {
         // Set dump_data=true to enable dumping tensor data to files for analysis
         compareResults_nvfp4(output, ref_output.get(), ref_output_t.get(), rows, cols, atol, rtol,
-                             true, false, !row_scaled_nvfp4);
+                             true, false, columnwise);
 
         size_t scale_mismatches_num = 0;
         if constexpr (std::is_same_v<ScaleType, fp8e4m3>) {
@@ -1112,7 +1113,7 @@ void performTest(float (*OP)(const float),
                 unpadded_blocks_Y, unpadded_blocks_X, scales_stride);
         }
 
-        if (!row_scaled_nvfp4) {
+        if (columnwise) {
             if constexpr (std::is_same_v<ScaleType, fp8e4m3>) {
                 compare_scaling_factors<ScaleType>(
                     "scales_t", output.columnwise_cpu_scale_inv_ptr<ScaleType>(),
@@ -1365,7 +1366,8 @@ class FusedCastTransposeNVFP4TestSuite : public ::testing::TestWithParam
                 bool,
                 NVFP4ScalingMode,
                 NVFP4FourOverSixTestConfig,
-                transformer_engine::DType>> {};
+                transformer_engine::DType,
+                bool>> {};
 
 TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) {
     // Skip tests for pre-Blackwell architectures
@@ -1383,6 +1385,7 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) {
     const NVFP4ScalingMode scaling_mode = std::get<4>(GetParam());
     const NVFP4FourOverSixTestConfig config = std::get<5>(GetParam());
     const DType scale_type = std::get<6>(GetParam());
+    const bool rowwise_only = std::get<7>(GetParam());
 
     // Skip tests if the input tensor is 1D
     if (tensor_dims.size() < 2) {
@@ -1403,12 +1406,12 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) {
         if (scale_type == DType::kFloat8E4M3) {
             performTest<InputType, fp8e4m3>(
                 OP, tensor_dims, use_fast_math, scaling_mode, config.mode, config.scale_max,
-                config.err_use_fast_math);
+                config.err_use_fast_math, rowwise_only);
 #if CUDA_VERSION >= 13040
         } else if (scale_type == DType::kFloat8UE5M3) {
             performTest<InputType, fp8ue5m3>(
                 OP, tensor_dims, use_fast_math, scaling_mode, config.mode, config.scale_max,
-                config.err_use_fast_math);
+                config.err_use_fast_math, rowwise_only);
 #endif
         } else {
             FAIL() << "Unsupported NVFP4 scale dtype " << static_cast<int>(scale_type);
@@ -1467,6 +1470,9 @@ std::string test_name(const FusedCastTransposeNVFP4TestSuite::ParamType& param)
     if (std::get<6>(param) != DType::kFloat8E4M3) {
         name += "X" + test::typeName(std::get<6>(param));
     }
+    if (std::get<7>(param)) {
+        name += "XROWWISE_ONLY";
+    }
     return name;
 }
 
@@ -1477,10 +1483,11 @@ INSTANTIATE_TEST_SUITE_P(
         ::testing::ValuesIn(Activation_types),           // activation_type
         ::testing::ValuesIn(tensor_dims),                // tensor_dims
         ::testing::Values(DType::kBFloat16),             // input_type
-        ::testing::Values(false),                       // use_fast_math
+        ::testing::Values(false, true),                 // use_fast_math
         ::testing::Values(NVFP4ScalingMode::Block1D),   // scaling_mode
         ::testing::Values(NVFP4FourOverSixTestConfig{}), // four_over_six_config
-        ::testing::Values(DType::kFloat8E4M3)),          // scale_type
+        ::testing::Values(DType::kFloat8E4M3),           // scale_type
+        ::testing::Values(false, true)),                 // rowwise_only
     [](const testing::TestParamInfo<FusedCastTransposeNVFP4TestSuite::ParamType>& info) {
         return test_name(info.param);
     });
@@ -1495,7 +1502,8 @@ INSTANTIATE_TEST_SUITE_P(
         ::testing::Values(false),                           // use_fast_math
         ::testing::Values(NVFP4ScalingMode::RowScaled1D),   // scaling_mode
         ::testing::Values(NVFP4FourOverSixTestConfig{}),   // four_over_six_config
-        ::testing::Values(DType::kFloat8E4M3)),            // scale_type
+        ::testing::Values(DType::kFloat8E4M3),             // scale_type
+        ::testing::Values(false)),                         // rowwise_only
     [](const testing::TestParamInfo<FusedCastTransposeNVFP4TestSuite::ParamType>& info) {
         return test_name(info.param);
     });
@@ -1520,7 +1528,8 @@ INSTANTIATE_TEST_SUITE_P(
             NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMAE, 256, true},
             NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 256, false},
             NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 256, true}), // four_over_six_config
-        ::testing::Values(DType::kFloat8E4M3)),                             // scale_type
+        ::testing::Values(DType::kFloat8E4M3),                              // scale_type
+        ::testing::Values(false)),                                          // rowwise_only
     [](const testing::TestParamInfo<FusedCastTransposeNVFP4TestSuite::ParamType>& info) {
         return test_name(info.param);
     });
@@ -1586,7 +1595,8 @@ INSTANTIATE_TEST_SUITE_P(
         ::testing::Values(NVFP4ScalingMode::Block1D,
                           NVFP4ScalingMode::Block2D),   // scaling_mode
         ::testing::Values(NVFP4FourOverSixTestConfig{}),  // four_over_six_config
-        ::testing::Values(DType::kFloat8UE5M3)),
+        ::testing::Values(DType::kFloat8UE5M3),  // scale_type
+        ::testing::Values(false, true)),  // rowwise_only
     [](const testing::TestParamInfo<FusedCastTransposeNVFP4TestSuite::ParamType>& info) {
         return test_name(info.param);
     });

transformer_engine/common/cast/nvfp4/nvfp4_block_philox.cuh

@@ -0,0 +1,176 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+/*! \file nvfp4_block_philox.cuh
+ *  \brief Philox4x32 draws for NVFP4 stochastic rounding, one per output block, several
+ *         streams at a time.
+ *
+ *  Each draw is exactly what curanddx's philox4x32_native_state<Rounds> returns from
+ *  init(seed, subsequence, offset) followed by one generate4(): the counter is
+ *  {offset_lo, offset_hi, subsequence_lo, subsequence_hi}.  Only the cost differs.  Callers
+ *  give every output block its own subsequence and take one counter step from rng_state[1],
+ *  so the counter never moves past the offset the framework reserved, whatever the grid.
+ *
+ *  Two savings over a plain per-thread generator:
+ *   - Every round-0 term that reads only the launch-wide offset is evaluated once
+ *     (make_philox).
+ *   - When a thread draws for many blocks that share one subsequence word (its block column,
+ *     say), placing that word in the subsequence's LOW half makes both of round 0's products
+ *     thread constants, and the constant parts of rounds 0..2 fold away (make_fold).  The
+ *     other half carries the varying coordinate XORed with the launch constant p.b, a
+ *     bijection that cancels round 0's own XOR.
+ *  Several streams advance together so the key schedule is paid once per batch.
+ *
+ *  The round count follows NVTE_BUILD_NUM_PHILOX_ROUNDS, as TE's other Philox users do.
+ */
+
+#ifndef TRANSFORMER_ENGINE_NVFP4_BLOCK_PHILOX_CUH_
+#define TRANSFORMER_ENGINE_NVFP4_BLOCK_PHILOX_CUH_
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+namespace transformer_engine {
+namespace dispatch {
+namespace nvfp4 {
+namespace block_philox {
+
+#ifdef NVTE_BUILD_NUM_PHILOX_ROUNDS
+constexpr int kRounds = NVTE_BUILD_NUM_PHILOX_ROUNDS;
+#else
+constexpr int kRounds = 10;
+#endif
+// make_fold consumes rounds 0..2, so it only applies to builds with at least 3 rounds; callers
+// must fall back to a plain generator otherwise.
+constexpr bool kFoldSupported = kRounds >= 3;
+
+// Subsequence of each output block, as (low word, logical high word).  The low word is the
+// coordinate a thread keeps fixed while it walks the other, so make_fold can fold it.  Its
+// parity tells the directions apart, so the rowwise and columnwise draws of one element never
+// share a subsequence; within a direction the pair is the block's coordinates.
+//   rowwise block (row, block_col):                  (2 * block_col,       row)
+//   columnwise block (row_block, col), col = 4q + i: (2 * (4q) + 1, 4 * row_block + i)
+// The column's offset inside its group of four moves to the high word because a thread owns
+// four adjacent columns and folds against their base.
+__host__ __device__ constexpr uint32_t rowwise_sub_lo(uint32_t block_col) { return block_col << 1; }
+__host__ __device__ constexpr uint32_t rowwise_sub_hi(uint32_t row) { return row; }
+__host__ __device__ constexpr uint32_t columnwise_sub_lo(uint32_t quad_base_col) {
+  return (quad_base_col << 1) | 1u;
+}
+__host__ __device__ constexpr uint32_t columnwise_sub_hi(uint32_t row_block, uint32_t col_in_quad) {
+  return (row_block << 2) | col_in_quad;
+}
+
+constexpr uint32_t kM0 = 0xD2511F53u;
+constexpr uint32_t kM1 = 0xCD9E8D57u;
+constexpr uint32_t kW0 = 0x9E3779B9u;
+constexpr uint32_t kW1 = 0xBB67AE85u;
+
+// Both halves of a 32 x 32 -> 64-bit product in one IMAD.WIDE.U32, written as PTX like
+// curanddx's mulhilo32.  As inline asm the product chain keeps the place the kernels give it,
+// right after their input loads, where it fills the wait for them; as plain C++ the compiler
+// sank it past the loads' consumers and the kernels stalled on the loads instead.
+__device__ __forceinline__ void mulhilo(uint32_t a, uint32_t b, uint32_t &lo, uint32_t &hi) {
+  uint64_t p;
+  asm("mul.wide.u32 %0, %1, %2;" : "=l"(p) : "r"(a), "r"(b));
+  lo = static_cast<uint32_t>(p);
+  hi = static_cast<uint32_t>(p >> 32);
+}
+
+// Three-way XOR in one LOP3.
+__device__ __forceinline__ uint32_t xor3(uint32_t a, uint32_t b, uint32_t c) {
+  uint32_t d;
+  asm("lop3.b32 %0, %1, %2, %3, 0x96;" : "=r"(d) : "r"(a), "r"(b), "r"(c));
+  return d;
+}
+
+// Round-0 terms that read only the key and the launch-wide offset.
+struct PhiloxState {
+  uint2 key;   // the key round 1 uses
+  uint32_t a;  // offset_hi ^ k0         -> round-0 c.x = hi(sub_lo * M1) ^ a
+  uint32_t b;  // hi(offset_lo * M0) ^ k1 -> round-0 c.z = b ^ sub_hi
+  uint32_t w;  // lo(offset_lo * M0)      -> round-0 c.w
+};
+
+__device__ __forceinline__ PhiloxState make_philox(uint64_t seed, uint64_t offset) {
+  const uint32_t k0 = static_cast<uint32_t>(seed);
+  const uint32_t k1 = static_cast<uint32_t>(seed >> 32);
+  uint32_t lo0, hi0;
+  mulhilo(static_cast<uint32_t>(offset), kM0, lo0, hi0);
+  PhiloxState p;
+  p.a = static_cast<uint32_t>(offset >> 32) ^ k0;
+  p.b = hi0 ^ k1;
+  p.w = lo0;
+  p.key = make_uint2(k0 + kW0, k1 + kW1);
+  return p;
+}
+
+// Rounds 0..2 for a fixed subsequence low word.
+struct PhiloxFold {
+  uint32_t A;     // c1_0 ^ k1.x   -> c0_1 = hi1 ^ A
+  uint32_t B;     // hi1_2 ^ k2.x  -> c0_2 = B ^ c1_1
+  uint32_t D;     // c3_1 ^ k2.y   -> c2_2 = hi0 ^ D
+  uint32_t c1_2;  // round-2 c1, constant for a fixed low word
+  uint2 k3;       // key as of round 3
+};
+
+__device__ __forceinline__ PhiloxFold make_fold(const PhiloxState &p, uint32_t sub_lo) {
+  uint32_t lo0, hi0, lo1, hi1;
+  mulhilo(sub_lo, kM1, lo1, hi1);  // round 0: both products are constant
+  const uint32_t c0_0 = hi1 ^ p.a, c1_0 = lo1;
+  uint2 k = p.key;               // round-1 key
+  mulhilo(c0_0, kM0, lo0, hi0);  // round 1: the .x product is constant
+  PhiloxFold f;
+  f.A = c1_0 ^ k.x;
+  const uint32_t c2_1 = xor3(hi0, p.w, k.y), c3_1 = lo0;
+  k.x += kW0;
+  k.y += kW1;                    // round-2 key
+  mulhilo(c2_1, kM1, lo1, hi1);  // round 2: the .z product is constant
+  f.B = hi1 ^ k.x;
+  f.c1_2 = lo1;
+  f.D = c3_1 ^ k.y;
+  k.x += kW0;
+  k.y += kW1;  // round-3 key
+  f.k3 = k;
+  return f;
+}
+
+// N draws that share the fold's low word.  logical_hi[i] is the varying coordinate; the
+// draw's subsequence high word is logical_hi[i] ^ p.b.
+template <int N>
+__device__ __forceinline__ void philox_fold_n(const uint32_t (&logical_hi)[N], const PhiloxFold &f,
+                                              uint4 (&c)[N]) {
+#pragma unroll
+  for (int i = 0; i < N; ++i) {  // the varying remnant of rounds 0..2
+    uint32_t lo0, hi0, lo1, hi1;
+    mulhilo(logical_hi[i], kM1, lo1, hi1);
+    mulhilo(hi1 ^ f.A, kM0, lo0, hi0);
+    c[i] = make_uint4(f.B ^ lo1, f.c1_2, hi0 ^ f.D, lo0);
+  }
+  uint2 k = f.k3;
+#pragma unroll
+  for (int r = 3; r < kRounds; ++r) {
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      uint32_t lo0, hi0, lo1, hi1;
+      mulhilo(c[i].x, kM0, lo0, hi0);
+      mulhilo(c[i].z, kM1, lo1, hi1);
+      c[i] = make_uint4(xor3(hi1, c[i].y, k.x), lo1, xor3(hi0, c[i].w, k.y), lo0);
+    }
+    if (r < kRounds - 1) {
+      k.x += kW0;
+      k.y += kW1;
+    }
+  }
+}
+
+}  // namespace block_philox
+}  // namespace nvfp4
+}  // namespace dispatch
+}  // namespace transformer_engine
+
+#endif  // TRANSFORMER_ENGINE_NVFP4_BLOCK_PHILOX_CUH_

transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh

@@ -24,6 +24,7 @@
 #include "../../util/ptx_arch_spec.cuh"
 #include "../../utils.cuh"
 #include "core_nvfp4.cuh"
+#include "specialized/quantize_transpose_nvfp4_register_1D.cuh"
 #include "specialized/quantize_transpose_nvfp4_tuned_1D.cuh"
 #include "swizzle.cuh"
 
@@ -1860,6 +1861,12 @@ void quantize_transpose_impl(const Tensor &input, const Tensor *noop, Tensor *ou
   }
 
   if (!use_2d_quantization && (input.dtype() == DType::kBFloat16)) {
+    // The register-resident kernel serves every 1D request it supports; row-scaled NVFP4,
+    // fast math, swizzled scales and RHT stay on the TMA kernel.
+    if (!apply_columnwise_rht && register_1D_supported(input, *output, quant_config)) {
+      quantize_transpose_register_1D<ScaleType>(input, noop, output, quant_config, stream);
+      return;
+    }
     quantize_transpose_tuned_1D<ScaleType>(input, noop, output, quant_config, stream);
     return;
   }