perf: transpose-quantize backward operands to route all gemms nt

- quantize gains out_layout (0 row-major / 1 transposed / 2 single-read dual-write); modes 1/2 run a new 32x32 smem-tile transpose kernel
- backward feeds g8/w8T and g8T/x8T to trans_b=True gemms, dropping the NN-swap and TT crosswise kernels from training; fp8 weights keep the swap fallback
- a 64x64 tile variant tied on the real step mix and was reverted; noted in the kernel header

Benchmark: NVIDIA L20, 1.2B model, full train step fwd+bwd+CE
- M=8192: fp8 551.8 -> 532.2 ms, 1.21x -> 1.26x vs bf16; M=2048 0.90x -> 0.95x
- kernel-level grad_x +3.7..12.4%, grad_w +13.8..20.8%; layouts byte-exact, fp8 tests 36/36

astrai/extension/fp8.py

@@ -407,11 +407,21 @@ class _LinearFp8(torch.autograd.Function):
                 meta.g.seed(g2, fmt)
             sg = meta.g.scale.clone()
             sw, sx = _sw_fwd, _sx_fwd
-        g8, amax_g = quantize(g2, sg.reciprocal(), fmt)
-        x8, _ = quantize(x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt)
-        w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0]
-        grad_x = mm_fp8(g8, w8, sg * sw).reshape(x.shape)  # g8[m,n] @ w8[n,k]
-        grad_w = mm_fp8(g8, x8, sg * sx, trans_a=True)  # g8.T @ x8
+        # Backward GEMMs route through the NT fast path via transposed
+        # quantize outputs: g8 [m,n] with w8T [k,n] (trans_b=True) gives
+        # grad_x, g8T [n,m] with x8T [k,m] gives grad_w — no NN-swap or TT
+        # crosswise kernel in the training path. g is consumed in both
+        # orientations, so one dual-layout pass feeds both.
+        g8, g8T, amax_g = quantize(g2, sg.reciprocal(), fmt, layout=2)
+        x8T, _ = quantize(x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt, layout=1)
+        if _is_fp8(w.dtype):
+            # Pre-quantized weight has no transposed copy: keep the swap
+            # path for grad_x (grad_w is unaffected).
+            grad_x = mm_fp8(g8, w, sg * sw).reshape(x.shape)
+        else:
+            w8T, _ = quantize(w, sw.reciprocal(), fmt, layout=1)
+            grad_x = mm_fp8(g8, w8T, sg * sw, trans_b=True).reshape(x.shape)
+        grad_w = mm_fp8(g8T, x8T, sg * sx, trans_b=True)  # g8.T @ x8
         # bias-free linears must not pay the column-sum
         # reduce: g2.sum(0) is another full read of the gradient.
         grad_b = g2.sum(0).to(torch.bfloat16) if ctx.needs_input_grad[2] else None

astrai/extension/ops/fp8.py

@@ -63,6 +63,71 @@ def _fp8_quantize_fake(x, scale, fmt):
 _QUANT_INPUT_DTYPES = (torch.bfloat16, torch.float16, torch.float32)
 
 
+@custom_op("custom::fp8_quantize_t", mutates_args=())
+def fp8_quantize_t(
+    x: torch.Tensor, scale: torch.Tensor, fmt: int
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Transposed-output variant of fp8_quantize: returns ``(x8T, amax)``
+    where ``x8T`` is the [cols][rows] row-major transpose of the quantized
+    input (the K-contiguous operand orientation for NT GEMMs)."""
+
+
+@fp8_quantize_t.register_fake
+def _fp8_quantize_t_fake(x, scale, fmt):
+    dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn
+    rows, cols = x.shape[-2], x.shape[-1]
+    return (
+        torch.empty((*x.shape[:-2], cols, rows), device=x.device, dtype=dtype),
+        torch.empty(1, device=x.device, dtype=torch.float32),
+    )
+
+
+@fp8_quantize_t.register_kernel("cuda")
+def _fp8_quantize_t_cuda(x, scale, fmt):
+    if x.dtype not in _QUANT_INPUT_DTYPES:
+        raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
+    return get_module("fp8_ops").quantize(x, scale, int(fmt), 1)
+
+
+@fp8_quantize_t.register_kernel("cpu")
+def _fp8_quantize_t_cpu(x, scale, fmt):
+    x8, amax = _fp8_quantize_cpu(x, scale, fmt)
+    return x8.transpose(-2, -1).contiguous(), amax
+
+
+@custom_op("custom::fp8_quantize_dual", mutates_args=())
+def fp8_quantize_dual(
+    x: torch.Tensor, scale: torch.Tensor, fmt: int
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Dual-orientation quantize: one read of ``x`` produces both the
+    row-major ``x8`` and its transposed ``x8T`` (plus ``amax``), for tensors
+    consumed by GEMMs on both orientations (backward ``g``)."""
+
+
+@fp8_quantize_dual.register_fake
+def _fp8_quantize_dual_fake(x, scale, fmt):
+    dtype = torch.float8_e5m2 if fmt == 1 else torch.float8_e4m3fn
+    rows, cols = x.shape[-2], x.shape[-1]
+    return (
+        torch.empty(x.shape, device=x.device, dtype=dtype),
+        torch.empty((*x.shape[:-2], cols, rows), device=x.device, dtype=dtype),
+        torch.empty(1, device=x.device, dtype=torch.float32),
+    )
+
+
+@fp8_quantize_dual.register_kernel("cuda")
+def _fp8_quantize_dual_cuda(x, scale, fmt):
+    if x.dtype not in _QUANT_INPUT_DTYPES:
+        raise TypeError(f"fp8 quantize requires bf16/fp16/fp32 input, got {x.dtype}")
+    return get_module("fp8_ops").quantize(x, scale, int(fmt), 2)
+
+
+@fp8_quantize_dual.register_kernel("cpu")
+def _fp8_quantize_dual_cpu(x, scale, fmt):
+    x8, amax = _fp8_quantize_cpu(x, scale, fmt)
+    return x8, x8.transpose(-2, -1).contiguous(), amax
+
+
 @fp8_quantize.register_kernel("cuda")
 def _fp8_quantize_cuda(x, scale, fmt):
     if x.dtype not in _QUANT_INPUT_DTYPES:
@@ -125,13 +190,16 @@ def _fp8_gemm_cpu(a, b, scale, trans_a=0, trans_b=0, bias=None):
 
 
 def quantize(
-    x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3"
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Float (bf16/fp16/fp32) -> FP8 quantize with fused amax; returns
-    ``(x8, amax)``.
+    x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3", layout: int = 0
+) -> tuple:
+    """Float (bf16/fp16/fp32) -> FP8 quantize with fused amax.
 
     ``scale`` is the quantization multiplier (device scalar); ``fmt`` selects
-    E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor.
+    E4M3 or E5M2. ``amax`` is a fresh 1-element float32 tensor. ``layout``
+    picks the output orientation: 0 = row-major ``(x8, amax)``; 1 =
+    transposed ``[cols][rows]`` ``(x8T, amax)`` — the K-contiguous operand
+    orientation NT GEMMs want; 2 = both from one read ``(x8, x8T, amax)``
+    (for tensors consumed in both orientations, e.g. backward ``g``).
     """
     # Hot-path bypass of the torch.library dispatch (~5us/call, ~40% of a
     # 512-wide GEMM): real CUDA tensors of a supported dtype go straight to
@@ -144,8 +212,12 @@ def quantize(
         and x.dtype in _QUANT_INPUT_DTYPES
         and fmt in _FMT_TO_INT
     ):
-        return get_module("fp8_ops").quantize(x, scale, _FMT_TO_INT[fmt])
-    return fp8_quantize(x, scale, _fmt_int(fmt))
+        return get_module("fp8_ops").quantize(x, scale, _FMT_TO_INT[fmt], layout)
+    if layout == 0:
+        return fp8_quantize(x, scale, _fmt_int(fmt))
+    if layout == 1:
+        return fp8_quantize_t(x, scale, _fmt_int(fmt))
+    return fp8_quantize_dual(x, scale, _fmt_int(fmt))
 
 
 def mm_fp8(

csrc/kernels/fp8/common.h

@@ -71,12 +71,24 @@ struct FP8QuantizeParams {
     // the binding and receives the raw-domain absolute maximum.
     const void* __restrict__ input_ptr = nullptr;
     void* __restrict__ output_ptr = nullptr;
+    void* __restrict__ output_transposed_ptr = nullptr;
+    // Transposed-output destination ([cols][rows]); the output-layout modes:
+    //   0 = row-major only (output_ptr; the vectorized elementwise kernel)
+    //   1 = transposed only (output_transposed_ptr; the tiled kernel)
+    //   2 = both destinations in one read of the input (the tiled kernel)
+    // Modes 1/2 exist so crosswise-layout GEMM operands (NN grad_x, TT
+    // grad_w) can be produced K-contiguous instead, routing every training
+    // GEMM through the dual-congruous NT fast path.
+    int out_layout = 0;
 
     const float* __restrict__ scale = nullptr;
     float* __restrict__ amax = nullptr;
 
-    // Element count (only the elementwise quantize kernel uses it).
+    // Element count (only the elementwise quantize kernel uses it); the
+    // tiled kernel views the same buffer as [rows][cols] row-major.
     int total = 0;
+    int rows = 0;
+    int cols = 0;
 };
 
 // Unified GEMM parameter POD, mirroring AttentionParams: one struct flows

csrc/kernels/fp8/ops.cu

@@ -46,16 +46,6 @@ void check_scale(const torch::Tensor& scale, const torch::Tensor& input) {
                 "scale must be a CUDA float32 scalar on the input device");
 }
 
-void pack_quantize(FP8QuantizeParams& p, const void* input, void* output,
-                   const torch::Tensor& scale, torch::Tensor& amax,
-                   int64_t total) {
-    p.input_ptr = input;
-    p.output_ptr = output;
-    p.scale = scale.data_ptr<float>();
-    p.amax = amax.data_ptr<float>();
-    p.total = static_cast<int>(total);
-}
-
 void pack_gemm(FP8Params& p, const void* a, const void* b, void* output,
                const torch::Tensor& scale, int64_t m, int64_t n, int64_t k,
                int64_t a_ld, int64_t b_ld) {
@@ -99,11 +89,45 @@ bool resolve_operand(const torch::Tensor& t_in, bool flag, int64_t& ld,
     return flag ^ col_major;
 }
 
+// Dtype x format switch shared by both quantize kernels; Tiled selects the
+// transpose kernel (out_layout 1/2) over the vectorized elementwise one.
+template <bool Tiled, FP8Format Fmt, typename InT>
+void launch_one(const FP8QuantizeParams& p, cudaStream_t stream) {
+    if constexpr (Tiled)
+        launch_fp8_quantize_tiled<Fmt, InT>(p, stream);
+    else
+        launch_fp8_quantize<Fmt, InT>(p, stream);
+}
+
+template <bool Tiled>
+void launch_quantize_for(const torch::Tensor& x, const FP8QuantizeParams& p,
+                         bool e5m2, cudaStream_t stream) {
+    if (x.scalar_type() == torch::kHalf) {
+        if (e5m2)
+            launch_one<Tiled, FP8Format::E5M2, __half>(p, stream);
+        else
+            launch_one<Tiled, FP8Format::E4M3, __half>(p, stream);
+    } else if (x.scalar_type() == torch::kFloat32) {
+        if (e5m2)
+            launch_one<Tiled, FP8Format::E5M2, float>(p, stream);
+        else
+            launch_one<Tiled, FP8Format::E4M3, float>(p, stream);
+    } else {
+        if (e5m2)
+            launch_one<Tiled, FP8Format::E5M2, __nv_bfloat16>(p, stream);
+        else
+            launch_one<Tiled, FP8Format::E4M3, __nv_bfloat16>(p, stream);
+    }
+}
+
 }  // namespace
 
-std::tuple<torch::Tensor, torch::Tensor> quantize(torch::Tensor x,
-                                                  torch::Tensor scale,
-                                                  int64_t fmt) {
+// Output-layout dispatch: 0 = [rows][cols] row-major (the historic 2-tuple
+// return), 1 = transposed [cols][rows] only (2-tuple), 2 = both orientations
+// from a single read of the input (3-tuple). Layouts 1/2 feed the NT GEMM
+// fast path from crosswise consumers (backward grad_x / grad_w).
+py::object quantize(torch::Tensor x, torch::Tensor scale, int64_t fmt,
+                    int64_t layout) {
     TORCH_CHECK(x.is_cuda(), "CUDA tensors required");
     TORCH_CHECK(x.scalar_type() == torch::kBFloat16 ||
                     x.scalar_type() == torch::kHalf ||
@@ -112,39 +136,44 @@ std::tuple<torch::Tensor, torch::Tensor> quantize(torch::Tensor x,
     TORCH_CHECK(fmt == static_cast<int64_t>(FP8Format::E4M3) ||
                     fmt == static_cast<int64_t>(FP8Format::E5M2),
                 "unsupported quantization type: expected E4M3 (0) or E5M2 (1)");
+    TORCH_CHECK(layout >= 0 && layout <= 2,
+                "layout must be 0 (row-major), 1 (transposed) or 2 (both)");
+    TORCH_CHECK(layout == 0 || x.dim() >= 2,
+                "transposed quantize layouts need a 2D+ tensor");
     check_scale(scale, x);
     check_fp8_device(x);
     const at::cuda::OptionalCUDAGuard guard(x.device());
     auto stream = at::cuda::getCurrentCUDAStream();
     auto input = x.contiguous();
-    auto output = torch::empty_like(
-        input, input.options().dtype(fmt ? torch::kFloat8_e5m2
-                                          : torch::kFloat8_e4m3fn));
+    auto out_opts = input.options().dtype(
+        fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn);
     auto amax = torch::zeros({1}, input.options().dtype(torch::kFloat32));
+
     FP8QuantizeParams p;
-    pack_quantize(p, input.data_ptr(), output.data_ptr(), scale, amax,
-                  input.numel());
-    const bool e5m2 = fmt == static_cast<int64_t>(FP8Format::E5M2);
-    if (x.scalar_type() == torch::kHalf) {
-        if (e5m2)
-            launch_fp8_quantize<FP8Format::E5M2, __half>(p, stream.stream());
-        else
-            launch_fp8_quantize<FP8Format::E4M3, __half>(p, stream.stream());
-    } else if (x.scalar_type() == torch::kFloat32) {
-        if (e5m2)
-            launch_fp8_quantize<FP8Format::E5M2, float>(p, stream.stream());
-        else
-            launch_fp8_quantize<FP8Format::E4M3, float>(p, stream.stream());
-    } else {
-        if (e5m2)
-            launch_fp8_quantize<FP8Format::E5M2, __nv_bfloat16>(
-                p, stream.stream());
-        else
-            launch_fp8_quantize<FP8Format::E4M3, __nv_bfloat16>(
-                p, stream.stream());
+    p.input_ptr = input.data_ptr();
+    p.scale = scale.data_ptr<float>();
+    p.amax = amax.data_ptr<float>();
+    p.total = static_cast<int>(input.numel());
+    p.out_layout = static_cast<int>(layout);
+    p.rows = static_cast<int>(input.size(-2));
+    p.cols = static_cast<int>(input.size(-1));
+    torch::Tensor output, output_t;
+    if (layout == 0 || layout == 2) {
+        output = torch::empty_like(input, out_opts);
+        p.output_ptr = output.data_ptr();
     }
+    if (layout >= 1) {
+        output_t = torch::empty({input.size(-1), input.size(-2)}, out_opts);
+        p.output_transposed_ptr = output_t.data_ptr();
+    }
+    const bool e5m2 = fmt == static_cast<int64_t>(FP8Format::E5M2);
+    if (layout != 0)
+        launch_quantize_for<true>(input, p, e5m2, stream.stream());
+    else
+        launch_quantize_for<false>(input, p, e5m2, stream.stream());
     C10_CUDA_CHECK(cudaGetLastError());
-    return {output, amax};
+    if (layout == 2) return py::make_tuple(output, output_t, amax);
+    return py::make_tuple(layout == 1 ? output_t : output, amax);
 }
 
 torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
@@ -219,7 +248,7 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
 // bias argument through untouched instead of normalizing it host-side).
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("quantize", &quantize, py::arg("x"), py::arg("scale"),
-          py::arg("fmt"));
+          py::arg("fmt"), py::arg("layout") = 0);
     m.def(
         "mm_fp8",
         [](torch::Tensor a, torch::Tensor b, torch::Tensor scale,

csrc/kernels/fp8/quantize.cuh

@@ -169,5 +169,89 @@ void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
     fp8_quantize_kernel<Fmt, InT><<<blocks, kThreads, 0, stream>>>(p);
 }
 
+// Tiled transpose quantize (out_layout 1/2): reads the [rows][cols]
+// row-major input once and writes the fp8 bytes transposed ([cols][rows],
+// so the contract dim lands K-contiguous for NT GEMM operands) and, in
+// mode 2, the plain row-major copy too. A 32x32 tile stages through shared
+// memory: input-row-major loads and output writes both stay coalesced, and
+// the byte-wide staging is conflict-free — the +4 pad makes the store
+// stride 9 (words) coprime with the 32 banks and the load is a 32-byte
+// broadcast segment. A 64x64 split-half variant (16 elems/thread, paired
+// 2-byte scatter stores) measured +21% on L2-resident shapes but -3..5%
+// on the DRAM-bound ones that carry the training traffic (occupancy and
+// memory-level parallelism, not instruction count, gate the DRAM regime);
+// weighted by the real step's mix the two tie, so the simpler tile stays.
+template <FP8Format Fmt, typename InT>
+__global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
+    constexpr int kTile = 32;
+    __shared__ uint8_t tile[kTile][kTile + 4];
+    const float mult = *p.scale;
+    const auto* x = static_cast<const InT*>(p.input_ptr);
+    const int r0 = blockIdx.y * kTile;
+    const int c0 = blockIdx.x * kTile;
+    const int r = r0 + threadIdx.y * 4;
+    const int c = c0 + threadIdx.x;
+
+    uint8_t q[4];
+    float local_amax = 0.0f;
+#pragma unroll
+    for (int j = 0; j < 4; ++j) {
+        q[j] = 0;
+        if (r + j < p.rows && c < p.cols) {
+            const float v =
+                quant_in_traits<InT>::to_float(x[(int64_t)(r + j) * p.cols + c]);
+            local_amax = fmaxf(local_amax, fabsf(v));
+            if constexpr (Fmt == FP8Format::E5M2)
+                q[j] = __nv_fp8_e5m2(v * mult).__x;
+            else
+                q[j] = __nv_fp8_e4m3(v * mult).__x;
+        }
+    }
+    if (p.out_layout == 2) {
+        uint8_t* out = static_cast<uint8_t*>(p.output_ptr);
+#pragma unroll
+        for (int j = 0; j < 4; ++j)
+            if (r + j < p.rows && c < p.cols)
+                out[(int64_t)(r + j) * p.cols + c] = q[j];
+    }
+#pragma unroll
+    for (int j = 0; j < 4; ++j) tile[threadIdx.x][threadIdx.y * 4 + j] = q[j];
+    __syncthreads();
+    // Transposed scatter: output element (c, r) lives at c * rows + r; r
+    // tracks threadIdx.x so each warp writes one contiguous run. The read
+    // swaps the staging indices — tile[col][row] was written, so the value
+    // for input (r0+tx, c0+ty*4+j) sits at tile[ty*4+j][tx].
+    uint8_t* out_t = static_cast<uint8_t*>(p.output_transposed_ptr);
+#pragma unroll
+    for (int j = 0; j < 4; ++j) {
+        const int oc = c0 + threadIdx.y * 4 + j;
+        if (oc < p.cols && r0 + threadIdx.x < p.rows)
+            out_t[(int64_t)oc * p.rows + r0 + threadIdx.x] =
+                tile[threadIdx.y * 4 + j][threadIdx.x];
+    }
+    if (p.amax) {
+        local_amax = warp_reduce_max(local_amax);
+        __shared__ float slots[8];
+        // blockDim.x is 32, so warp id == threadIdx.y; only complete warps
+        // exist (blockDim.y == 8).
+        if (threadIdx.x == 0) slots[threadIdx.y] = local_amax;
+        __syncthreads();
+        if (threadIdx.x == 0 && threadIdx.y == 0) {
+            float v = 0.0f;
+            for (int w = 0; w < (int)blockDim.y; ++w) v = fmaxf(v, slots[w]);
+            atomic_max_float(p.amax, v);
+        }
+    }
+}
+
+template <FP8Format Fmt, typename InT>
+void launch_fp8_quantize_tiled(const FP8QuantizeParams& p,
+                               cudaStream_t stream) {
+    const dim3 grid((p.cols + 31) / 32, (p.rows + 31) / 32);
+    if (grid.x == 0 || grid.y == 0) return;
+    fp8_quantize_tiled_kernel<Fmt, InT>
+        <<<grid, dim3(32, 8), 0, stream>>>(p);
+}
+
 }  // namespace fp8
 }  // namespace astrai