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(