refactor: fp8 gemm policy layering with swap-NN and narrow-N ctas

Kernel restructured CUTLASS-style: Fp8GemmPolicy as the kernel's single template parameter (traits + operand layouts + scheduling knobs), the body split into Fp8GemmTileScheduler / Fp8CollectiveMainloop / Fp8CollectiveEpilogue collectives, and the entry split into canonicalize_gemm -> plan_gemm -> launch_plan behind fp8::gemm.

- NN (dual-N-contiguous) problems run as their transpose: the swap in canonicalize_gemm plus an out-transposed epilogue removes one kernel instantiation per (format, tile config)
- new 128x64 narrow CTA (8 warps of 32x32) serves the sub-wave band once its grid passes ~3/8 of a wave: +7..77% there (128x4096x4096 116->131T, 1024^3 131->174T, 4096x384x4096 147->242T, 8192x128x4096 131->233T); decode, the padding band and multi-wave shapes unchanged
- launch_with_smem no longer swallows cudaFuncSetAttribute failures
- fp8_test: GPU-side fp32 reference (O(m*n) compare instead of O(m*n*k) host loop), production-dispatch cases for the NN swap and the plan selection; dead transpose_layout trait removed

Device: NVIDIA RTX 6000D (sm_120, 156 SMs), CUDA 13.1, torch 2.11.0+cu130. Kernel-only bench vs CUTLASS 4.8.0 sm120 dense fp8: ahead up to 1.68x below one wave (512^3 44 vs 26T, 64x4096x4096 95 vs 62T), within ~7% in the DRAM-streaming regime (8192^3 248 vs 266T).

csrc/kernels/fp8/common.h

@@ -29,23 +29,6 @@ enum class FP8Format : int {
 struct RowMajor {};
 struct ColMajor {};
 
-// Transpose of a layout tag: the same buffer with the rows and contract dims
-// swapped. B's tag is relative to the canonical [K][N] GEMM matrix, so the
-// stage-load (which views any operand as [rows][contract]) sees the transposed
-// tag — this trait makes that inversion explicit.
-template <typename Layout>
-struct transpose_layout;
-template <>
-struct transpose_layout<RowMajor> {
-    using type = ColMajor;
-};
-template <>
-struct transpose_layout<ColMajor> {
-    using type = RowMajor;
-};
-template <typename Layout>
-using transpose_layout_t = typename transpose_layout<Layout>::type;
-
 // Compile-time tile configuration, mirroring KernelTraits<HEAD_DIM, BC,
 // WARPS, STAGES> in the attention kernels. `Fmt` selects the FP8 conversion
 // and the MMA PTX mnemonic; the remaining parameters shape the CTA tile, the
@@ -112,6 +95,13 @@ struct FP8Params {
     void* __restrict__ out_ptr = nullptr;
 
     const float* __restrict__ scale = nullptr;
+    // Transposed-output mode (set by dispatch_fp8_gemm's swap for NN
+    // problems): the kernel computes E[N'][M'] over swapped operands and the
+    // epilogue scatters into the caller's [M][N] row-major buffer, so
+    // D[row][col] lives at out[col * p.m + row] — p.m/p.n are the swapped
+    // problem's dims and the D row stride is p.m. Zero in the plain
+    // orientation.
+    int out_transposed = 0;
     // Shapes. `int` covers every realistic LLM shape; the kernels promote
     // to int64 for all pointer arithmetic.
     int m, n, k;

csrc/kernels/fp8/ops.cu

@@ -70,25 +70,9 @@ void pack_gemm(FP8Params& p, const void* a, const void* b, void* output,
     p.b_ld = static_cast<int>(b_ld);
 }
 
-template <FP8Format Fmt, int Variant>
+// Layout dispatch (the NN swap in canonicalize_gemm) and launch planning
-void launch_variant(const FP8Params& p, cudaStream_t stream) {
+// (plan_gemm/launch_plan) live in gemm.cuh behind fp8::gemm — pure CUDA,
-    using LayoutA = std::conditional_t<(Variant & 2) != 0, ColMajor, RowMajor>;
+// shared with the C test suite.
-    using LayoutB = std::conditional_t<(Variant & 1) != 0, ColMajor, RowMajor>;
-    launch_fp8_gemm<Fmt, LayoutA, LayoutB>(p, stream);
-}
-
-template <FP8Format Fmt>
-void dispatch_gemm(const FP8Params& p, cudaStream_t stream, bool trans_a,
-                   bool trans_b) {
-    const int variant = (static_cast<int>(trans_a) << 1) |
-                        static_cast<int>(trans_b);
-    switch (variant) {
-        case 0: launch_variant<Fmt, 0>(p, stream); break;
-        case 1: launch_variant<Fmt, 1>(p, stream); break;
-        case 2: launch_variant<Fmt, 2>(p, stream); break;
-        case 3: launch_variant<Fmt, 3>(p, stream); break;
-    }
-}
 
 // Inner-layout resolution for one GEMM operand. The user flag names the
 // math (0 = tensor's last two dims are [rows][contract], 1 = transposed);
@@ -223,9 +207,9 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
     p.b_batch_stride = (batch_b == 1 && batch > 1) ? 0 : b_bstride;
     p.out_batch_stride = m * n;
     if (a.scalar_type() == torch::kFloat8_e4m3fn)
-        dispatch_gemm<FP8Format::E4M3>(p, stream.stream(), tag_a, tag_b);
+        gemm<FP8Format::E4M3>(p, stream.stream(), tag_a, tag_b);
     else
-        dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), tag_a, tag_b);
+        gemm<FP8Format::E5M2>(p, stream.stream(), tag_a, tag_b);
     C10_CUDA_CHECK(cudaGetLastError());
     return output;
 }

csrc/tests/fp8_test.cu

@@ -170,9 +170,37 @@ static bool test_single_mma() {
 // Part 2: GEMM correctness — layouts x K-tiles vs fp32 CPU reference
 // ---------------------------------------------------------------------------
 
+// Naive fp32 reference on the GPU: same layout interpretation as the CPU
+// loop it replaces (O(m*n) to check instead of O(m*n*k) to compute).
+__global__ static void
+naive_gemm_ref(const __nv_fp8_e4m3* a, const __nv_fp8_e4m3* b, float* out,
+               int m, int n, int k, int a_ld, int b_ld, int a_rm, int b_rm) {
+    const int i = blockIdx.y * 32 + threadIdx.y;
+    const int j = blockIdx.x * 32 + threadIdx.x;
+    if (i >= m || j >= n) return;
+    float acc = 0.f;
+    for (int kk = 0; kk < k; ++kk) {
+        float av = a_rm ? (float)a[i * a_ld + kk] : (float)a[kk * a_ld + i];
+        float bv = b_rm ? (float)b[kk * b_ld + j] : (float)b[j * b_ld + kk];
+        acc += av * bv;
+    }
+    out[i * n + j] = acc;
+}
+
+// Big-CTA policies for the direct-layout cases: kK/Stages vary per case;
+// the fast interior loop follows the dual-congruous rule, grouped raster 8
+// matches the production dispatch.
+template <typename LA, typename LB>
+constexpr bool kCaseFast =
+    !std::is_same_v<LA, ColMajor> && !std::is_same_v<LB, RowMajor>;
+template <typename LA, typename LB, int kK, int Stages>
+using CasePolicy =
+    Fp8GemmPolicy<FP8Format::E4M3, 128, 128, LA, LB, 64, 32, kK, Stages, 8,
+                  false, false, kCaseFast<LA, LB>>;
+
 template <typename LA, typename LB, int kK, int Stages>
 static bool run_gemm_case(const float* ha, const float* hb, int m, int n,
-                          int k, int a_ld, int b_ld) {
+                          int k, int a_ld, int b_ld, int dispatch = 0) {
     __nv_fp8_e4m3 *da, *db;
     __nv_bfloat16* dout;
     float* dscale;
@@ -205,7 +233,27 @@ static bool run_gemm_case(const float* ha, const float* hb, int m, int n,
     p.k = k;
     p.a_ld = a_ld;
     p.b_ld = b_ld;
-    launch_fp8_gemm<FP8Format::E4M3, LA, LB, kK, Stages>(p, 0);
+    float* d_ref;
+    cudaMalloc(&d_ref, (size_t)m * n * 4);
+    naive_gemm_ref<<<dim3((n + 31) / 32, (m + 31) / 32), dim3(32, 32)>>>(
+        da, db, d_ref, m, n, k, a_ld, b_ld,
+        !std::is_same_v<LA, ColMajor>, !std::is_same_v<LB, ColMajor>);
+    std::vector<float> href((size_t)m * n);
+    cudaMemcpy(href.data(), d_ref, href.size() * 4, cudaMemcpyDeviceToHost);
+    cudaFree(d_ref);
+
+    if (dispatch == 1)
+        // Production route, NN: the dual-N-contiguous problem has no
+        // dedicated instantiation — canonicalize_gemm swaps to the
+        // transposed <ColMajor, ColMajor> kernel with its out-transposed
+        // epilogue (see gemm.cuh).
+        gemm<FP8Format::E4M3>(p, 0, false, false);
+    else if (dispatch == 2)
+        // Production route, NT: exercises plan_gemm's small/narrow/big
+        // selection for this shape.
+        gemm<FP8Format::E4M3>(p, 0, false, true);
+    else
+        launch_policy<CasePolicy<LA, LB, kK, Stages>>(p, 0);
     cudaError_t e = cudaDeviceSynchronize();
     if (e != cudaSuccess) {
         printf("  CUDA err: %s\n", cudaGetErrorString(e));
@@ -218,20 +266,7 @@ static bool run_gemm_case(const float* ha, const float* hb, int m, int n,
     bool ok = true;
     for (int i = 0; i < m && ok; ++i) {
         for (int j = 0; j < n && ok; ++j) {
-            float ref = 0;
+            const float ref = href[(size_t)i * n + j];
-            for (int kk = 0; kk < k; ++kk) {
-                // A reference reads the actual uploaded buffer: LA ColMajor
-                // means the buffer is [K][M] (ha_t), else [M][K].
-                float av = std::is_same_v<LA, ColMajor>
-                               ? (float)__nv_fp8_e4m3(ha[kk * m + i])
-                               : (float)__nv_fp8_e4m3(ha[i * k + kk]);
-                float bv;
-                if (std::is_same_v<LB, ColMajor>)
-                    bv = (float)__nv_fp8_e4m3(hb[j * k + kk]);
-                else
-                    bv = (float)__nv_fp8_e4m3(hb[kk * n + j]);
-                ref += av * bv;
-            }
             float got =
                 __bfloat162float(__ushort_as_bfloat16(hb16[i * n + j]));
             float err = fabsf(got - ref);
@@ -254,6 +289,7 @@ static bool test_gemm() {
     } cfgs[] = {
         {128, 128, 128}, {256, 128, 256}, {128, 256, 64},
         {100, 130, 96},  {64, 64, 160},   {300, 200, 320},
+        {2048, 256, 512}, {1024, 1024, 512},
     };
     bool all = true;
     for (auto& c : cfgs) {
@@ -274,12 +310,12 @@ static bool test_gemm() {
         printf(" NT K64:");
         all &= run_gemm_case<RowMajor, ColMajor, 64, 2>(ha, hb_colmajor, c.m,
                                                         c.n, c.k, c.k, c.k);
-        printf(" NN K32:");
+        printf(" NN swap:");
-        all &= run_gemm_case<RowMajor, RowMajor, 32, 3>(ha, hb_rowmajor, c.m,
+        all &= run_gemm_case<RowMajor, RowMajor, 64, 2>(
-                                                        c.n, c.k, c.k, c.n);
+            ha, hb_rowmajor, c.m, c.n, c.k, c.k, c.n, /*dispatch=*/1);
-        printf(" NN K64:");
+        printf(" NT disp:");
-        all &= run_gemm_case<RowMajor, RowMajor, 64, 2>(ha, hb_rowmajor, c.m,
+        all &= run_gemm_case<RowMajor, ColMajor, 64, 2>(
-                                                        c.n, c.k, c.k, c.n);
+            ha, hb_colmajor, c.m, c.n, c.k, c.k, c.k, /*dispatch=*/2);
         printf(" TN K32:");
         all &= run_gemm_case<ColMajor, ColMajor, 32, 3>(ha_t, hb_colmajor, c.m,
                                                         c.n, c.k, c.m, c.k);

docs/developer/cuda_kernels.md

@@ -46,9 +46,9 @@ style as attention, but split into **three** files:
 
 | File | Role |
 |------|------|
-| `fp8/common.h` | `FP8Format` enum (E4M3/E5M2), `Fp8GemmTraits<Fmt, BlockM, BlockN, K, Stages>`, `FP8Params` POD — no torch |
+| `fp8/common.h` | `FP8Format` enum (E4M3/E5M2), `Fp8GemmTraits<Fmt, BlockM, BlockN, K, Stages>`, `Fp8GemmPolicy` (traits + layouts + scheduling knobs — the kernel's single template parameter), `FP8Params` POD — no torch |
 | `fp8/quantize.cuh` | pure-CUDA device code: `fp8_quantize_kernel<Fmt, InT>` (bf16/fp16/fp32 → FP8 + amax, `quant_in_traits<InT>` vectorized unpack) — no torch |
-| `fp8/gemm.cuh` | pure-CUDA device code: `fp8_gemm_kernel` (pre-quantized GEMM; 64×64 / 128×128 CTA picked at runtime by `prefer_small_cta`, 64×32 warp tiles, multi-stage cp.async, transposed-operand layouts) — no torch |
+| `fp8/gemm.cuh` | pure-CUDA device code: CUTLASS-style collectives (`Fp8GemmTileScheduler` / `Fp8CollectiveMainloop` / `Fp8CollectiveEpilogue`) around `fp8_gemm_kernel<Policy>` (pre-quantized GEMM; 64×64 / 128×64 / 128×128 CTA picked by `plan_gemm`, multi-stage cp.async, transposed-operand layouts, NN routed through a swap + out-transposed epilogue) — no torch. Entry: `gemm<Fmt>(params, stream, trans_a, trans_b)` = `canonicalize_gemm` → `plan_gemm` → `launch_plan` |
 | `fp8/ops.cu` | binding only: `check_fp8_device` (sm_89+), param packing, launch dispatch, pybind → module `fp8_ops` |
 
 Scale semantics: `quantize` takes the quantization *multiplier*; the