perf: speed up fp8 gemm across small and large shapes

- parameterize warp tile (WarpM/WarpN) in Fp8GemmTraits; MMA loops, fragment arrays and epilogue scale with kMt/kNt instead of the fixed 64x32/4x4, enabling cuBLAS-style 64x64 CTAs of 32x32 warps
- dispatch by output tiling (grid-searched via csrc/tests/fp8_sweep.cu): fewer than 48 output tiles take 64x64/32x32 with a lean ring (4 CTAs/SM fill the wave-quantization gap: 512^3 goes 16 -> 64 CTAs); larger shapes keep 128x128 with the kStages+1 ring
- kStages+1 canonic ring rotation drops the post-compute barrier on the congruous path (one __syncthreads per k-tile); LeanRing keeps the kStages ring for the small CTA; direct-crosswise operands always rotate kStages+1 (their prefetch issues right after barrier 1 and would race a lean ring - caught by the pure C layout suite)
- stage the bf16 epilogue through the reclaimed operand smem: swizzled scatter + barrier + coalesced 16B copy-out replaces 8 disjoint 16B per-warp segments (~50% write efficiency before)
- hoist per-lane ldmatrix swizzle offsets out of the mainloop (stage-relative table + ring-base add) so the innermost loop stops recomputing IMAD/LOP3 address chains
- bypass the torch.library dispatch for real CUDA tensors in quantize/mm_fp8 wrappers (~5us/call, ~40% of a 512-wide call's wall time); fake/subclass tensors keep the custom_op route

vs the previous kernel + python path, wall clock on NT squares: 512^3 52 -> 13us (4.0x, 5.2 -> 20.5 TF, now 1.36x cuBLAS _scaled_mm), 1024^3 1.05x, 2048^3 1.02x (46.9 -> 48.2 TF kernel-only); correctness: 4 layouts x 6 shapes pure C suite PASS, 588 pytest PASS

csrc/kernels/fp8/common.h

@@ -48,25 +48,36 @@ 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 and
-// the cp.async pipeline depth.
-template <FP8Format Fmt, int BlockM, int BlockN, int K, int Stages>
+// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile, the
+// warp tile (WarpM x WarpN — e.g. 64x32 on the 128x128 CTA, or 32x32 on the
+// cuBLAS-style 64x64 small CTA that lifts small-shape occupancy) and the
+// cp.async pipeline depth.
+template <FP8Format Fmt, int BlockM, int BlockN, int K, int Stages,
+          int WarpM = 64, int WarpN = 32>
 struct Fp8GemmTraits {
     static constexpr FP8Format kFormat = Fmt;
     static constexpr int kBlockM = BlockM;
     static constexpr int kBlockN = BlockN;
     static constexpr int kK = K;
     static constexpr int kStages = Stages;
+    static constexpr int kWarpM = WarpM;
+    static constexpr int kWarpN = WarpN;
     static constexpr bool kIsE5M2 = (Fmt == FP8Format::E5M2);
     static constexpr __nv_fp8_interpretation_t kNvFormat =
         kIsE5M2 ? __NV_E5M2 : __NV_E4M3;
     static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f;
 
-    // Derived launch geometry: 64x32 warp tiles give the CTA thread count.
-    // The shared-memory budget is layout-aware (crosswise operands add K-
-    // major staging + a canonical buffer), so it lives in Fp8GemmSmem in
-    // gemm.cuh together with the resident-CTA hint for __launch_bounds__.
-    static constexpr int kCtaThreads = (BlockM / 64) * (BlockN / 32) * 32;
+    // Derived launch geometry: WarpM x WarpN warp tiles tile the CTA. The
+    // shared-memory budget is layout-aware (crosswise operands add K-major
+    // staging + a canonical buffer), so it lives in Fp8GemmSmem in gemm.cuh
+    // together with the resident-CTA hint for __launch_bounds__.
+    static constexpr int kWarpsM = BlockM / WarpM;
+    static constexpr int kWarpsN = BlockN / WarpN;
+    static constexpr int kCtaThreads = kWarpsM * kWarpsN * 32;
+    static_assert(kWarpsM * WarpM == BlockM && kWarpsN * WarpN == BlockN,
+                  "warp tiles must exactly tile the CTA");
+    static_assert(WarpM % 16 == 0 && WarpN % 8 == 0,
+                  "warp tile must be a multiple of the m16n8 MMA shape");
 };
 
 // Quantize-kernel parameter POD: float input (bf16 / fp16 / fp32) -> FP8