refactor: dedupe fp8 kernel helpers and trim comments

- merge the quantize launchers into one Tiled template; extract shared cvt_fp8/publish_amax helpers and replace the dtype x format ladder with two-level template dispatch
- fold the gemm interior/generic operand loads into one kInterior template and the fast/generic async loads into load_async<kFast>; Policy carries the smem budget
- compress kernel comments to the load-bearing invariants, dropping measured-number essays; Policy signature and kernel code unchanged

Benchmark: NVIDIA L20, 1.2B model train step fwd+bwd+CE
- M=8192: fp8 532.2 -> 530.4 ms (1.26x, noise); tests/extension 65 passed, quantize layouts byte-exact, NT routing diff 0.0

csrc/kernels/fp8/common.h

@@ -11,30 +11,23 @@
 namespace astrai {
 namespace fp8 {
 
-// Compile-time FP8 format: E4M3 (forward / high precision, max 448) or
+// Compile-time FP8 format: E4M3 (forward, max 448) or E5M2 (gradients,
-// E5M2 (gradient / large dynamic range, max 57344).
+// max 57344).
 enum class FP8Format : int {
     E4M3 = 0,
     E5M2 = 1,
 };
 
-// Operand memory layouts as types (CUTLASS-style tags). The tag names the
+// Operand storage tags (CUTLASS-style) relative to the canonical matrices
-// storage order of the raw buffer relative to the operand's canonical GEMM
+// A [M][K] / B [K][N]: A RowMajor = [M][K] (default), A ColMajor = [K][M],
-// matrix — A is [M][K], B is [K][N]:
+// B RowMajor = [K][N], B ColMajor = [N][K] (the nn.Linear weight). Selection
-//   A RowMajor = [M][K] storage (K-contiguous rows; the default)
+// is by type at compile time (see gemm.cuh's stage loads).
-//   A ColMajor = [K][M] storage (M-contiguous; A^T)
-//   B RowMajor = [K][N] storage (N-contiguous; the plain a @ b operand)
-//   B ColMajor = [N][K] storage (K-contiguous; the nn.Linear weight layout)
-// Empty tags: selection happens by type at compile time (see load_operand_tile).
 struct RowMajor {};
 struct ColMajor {};
 
-// Compile-time tile configuration, mirroring KernelTraits<HEAD_DIM, BC,
+// Compile-time tile configuration, mirroring KernelTraits in the attention
-// WARPS, STAGES> in the attention kernels. `Fmt` selects the FP8 conversion
+// kernels: CTA tile, warp tile (WarpM x WarpN — e.g. 64x32 on the 128x128
-// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile, the
+// CTA, 32x32 on the 64x64 small CTA) and cp.async pipeline depth.
-// 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 {
@@ -50,10 +43,8 @@ struct Fp8GemmTraits {
         kIsE5M2 ? __NV_E5M2 : __NV_E4M3;
     static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f;
 
-    // Derived launch geometry: WarpM x WarpN warp tiles tile the CTA. The
+    // Derived geometry: warp tiles tile the CTA. The smem budget is
-    // shared-memory budget is layout-aware (crosswise operands add K-major
+    // layout-aware, so it lives in Fp8GemmSmem (gemm.cuh).
-    // 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;
@@ -63,73 +54,54 @@ struct Fp8GemmTraits {
                   "warp tile must be a multiple of the m16n8 MMA shape");
 };
 
-// Quantize-kernel parameter POD: float input (bf16 / fp16 / fp32) -> FP8
+// Quantize-kernel parameter POD: float input -> FP8 with fused amax.
-// with fused amax.
 struct FP8QuantizeParams {
-    // Float input and FP8 output buffers; scale is the quantization
-    // multiplier (device scalar). amax (may be null) is zero-initialized by
-    // 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;
+    void* __restrict__ output_transposed_ptr = nullptr;  // [cols][rows]
-    // Transposed-output destination ([cols][rows]); the output-layout modes:
+    // Output layout: 0 = row-major only, 1 = transposed only, 2 = both from
-    //   0 = row-major only (output_ptr; the vectorized elementwise kernel)
+    // a single read. Modes 1/2 produce K-contiguous operands so crosswise
-    //   1 = transposed only (output_transposed_ptr; the tiled kernel)
+    // consumers (backward grad_x / grad_w) route through the NT fast path.
-    //   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;
+    const float* __restrict__ scale = nullptr;  // device multiplier
-    float* __restrict__ amax = nullptr;
+    float* __restrict__ amax = nullptr;         // raw-domain max out
 
-    // Element count (only the elementwise quantize kernel uses it); the
+    // Element count (elementwise kernel); the tiled kernel views the same
-    // tiled kernel views the same buffer as [rows][cols] row-major.
+    // buffer as [rows][cols] row-major.
     int total = 0;
     int rows = 0;
     int cols = 0;
 };
 
 // Unified GEMM parameter POD, mirroring AttentionParams: one struct flows
-// through the pre-quantized GEMM kernels. Each kernel touches only the
+// through the kernels; each kernel touches only the fields it needs.
-// fields it needs; buffers are raw pointers packed by the torch binding.
-// Pointer members default to null so optional paths cannot hold garbage.
 struct FP8Params {
-    // Inputs: a/b are FP8 for the pre-quantized path. Scales are
+    // FP8 operands + output; scales are quantization steps (device
-    // quantization steps (device scalars).
+    // scalars). Optional bf16 bias fuses into the epilogue (fp32 add before
-    // Optional bf16 bias broadcast over output rows (fused into the epilogue
+    // the single bf16 rounding); null disables.
-    // before the bf16 rounding, so it adds in fp32 — one rounding fewer than
-    // the separate out + bias elementwise kernel it replaces). Null disables.
     const void* __restrict__ a_ptr = nullptr;
     const void* __restrict__ b_ptr = nullptr;
     const void* __restrict__ bias_ptr = nullptr;
     void* __restrict__ out_ptr = nullptr;
 
     const float* __restrict__ scale = nullptr;
-    // Transposed-output mode (set by dispatch_fp8_gemm's swap for NN
+    // NN-swap mode (canonicalize_gemm): the kernel computes the transposed
-    // problems): the kernel computes E[N'][M'] over swapped operands and the
+    // problem and the epilogue scatters D[row][col] to out[col * p.m + row]
-    // epilogue scatters into the caller's [M][N] row-major buffer, so
+    // in the caller's [M][N] buffer. Zero in the plain orientation.
-    // 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
+    int m, n, k;  // int covers LLM shapes; kernels promote to int64
-    // to int64 for all pointer arithmetic.
-    int m, n, k;
 
-    // Batched (bmm) geometry: grid.z slices step the operand/output pointers
+    // Batched (bmm) geometry: grid.z steps these element strides (0
-    // by these element strides (0 broadcasts the operand across batches).
+    // broadcasts the operand across batches).
     int batch = 1;
     int64_t a_batch_stride = 0;
     int64_t b_batch_stride = 0;
     int64_t out_batch_stride = 0;
 
-    // Physical leading dimensions (column count, i.e. row stride) of A and
+    // Physical leading dims (row strides) of A and B; the binding packs
-    // B. For a non-transposed operand the stride equals the contract dim;
+    // them so the kernel reads each buffer naturally or transposed per the
-    // for a transposed operand it is the operand's own column count. The
+    // LayoutA/LayoutB tags.
-    // binding packs these so the kernel reads both buffers either naturally
-    // or transposed depending on the LayoutA/LayoutB tags (see gemm.cuh).
     int a_ld, b_ld;
 };
 

csrc/kernels/fp8/ops.cu

@@ -6,7 +6,6 @@
 
 #include <cstdint>
 #include <mutex>
-#include <tuple>
 #include <unordered_map>
 
 #include "../common/device.cuh"
@@ -46,32 +45,13 @@ void check_scale(const torch::Tensor& scale, const torch::Tensor& input) {
                 "scale must be a CUDA float32 scalar on the input device");
 }
 
-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) {
-    p.a_ptr = a;
-    p.b_ptr = b;
-    p.out_ptr = output;
-    p.scale = scale.data_ptr<float>();
-    p.m = static_cast<int>(m);
-    p.n = static_cast<int>(n);
-    p.k = static_cast<int>(k);
-    p.a_ld = static_cast<int>(a_ld);
-    p.b_ld = static_cast<int>(b_ld);
-}
-
-// Layout dispatch (the NN swap in canonicalize_gemm) and launch planning
-// (plan_gemm/launch_plan) live in gemm.cuh behind fp8::gemm — pure CUDA,
-// shared with the C test suite.
-
 // Inner-layout resolution for one GEMM operand. The user flag names the
-// math (0 = tensor's last two dims are [rows][contract], 1 = transposed);
+// math (0 = last two dims are [rows][contract], 1 = transposed); the
-// the storage may independently be a col-major view (.t() of a contiguous
+// storage may independently be a col-major view (.t() of a contiguous
 // buffer), which folds into the returned dispatch flag at zero copy — the
 // kernel's LayoutA/LayoutB tags cover both storages. m/n/k derive from the
-// user flag only; the fold never swaps them (see the layout table in
+// user flag only. Tensors whose inner dims are neither natural layout fall
-// gemm.cuh). Tensors whose inner dims are neither natural layout fall back
+// back to .contiguous().
-// to .contiguous().
 bool resolve_operand(const torch::Tensor& t_in, bool flag, int64_t& ld,
                      int64_t& batch_stride, torch::Tensor& storage) {
     torch::Tensor t = t_in;
@@ -89,43 +69,36 @@ 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
+// Dtype dispatch over the unified quantize launcher.
-// transpose kernel (out_layout 1/2) over the vectorized elementwise one.
+template <bool Tiled, FP8Format Fmt>
-template <bool Tiled, FP8Format Fmt, typename InT>
+void launch_for_dtype(const torch::Tensor& x, const FP8QuantizeParams& p,
-void launch_one(const FP8QuantizeParams& p, cudaStream_t stream) {
+                      cudaStream_t stream) {
-    if constexpr (Tiled)
+    switch (x.scalar_type()) {
-        launch_fp8_quantize_tiled<Fmt, InT>(p, stream);
+    case torch::kHalf:
-    else
+        launch_fp8_quantize<Fmt, __half, Tiled>(p, stream);
-        launch_fp8_quantize<Fmt, InT>(p, stream);
+        break;
+    case torch::kFloat32:
+        launch_fp8_quantize<Fmt, float, Tiled>(p, stream);
+        break;
+    default:
+        launch_fp8_quantize<Fmt, __nv_bfloat16, Tiled>(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)
-        if (e5m2)
+        launch_for_dtype<Tiled, FP8Format::E5M2>(x, p, stream);
-            launch_one<Tiled, FP8Format::E5M2, __half>(p, stream);
+    else
-        else
+        launch_for_dtype<Tiled, FP8Format::E4M3>(x, p, stream);
-            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
 
-// Output-layout dispatch: 0 = [rows][cols] row-major (the historic 2-tuple
+// Output-layout dispatch: 0 = [rows][cols] row-major (2-tuple return),
-// return), 1 = transposed [cols][rows] only (2-tuple), 2 = both orientations
+// 1 = transposed [cols][rows] only (2-tuple), 2 = both orientations from a
-// from a single read of the input (3-tuple). Layouts 1/2 feed the NT GEMM
+// single read (3-tuple). Layouts 1/2 feed the NT GEMM fast path.
-// 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");
@@ -219,8 +192,15 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
             ? torch::empty({batch, m, n}, a.options().dtype(torch::kBFloat16))
             : torch::empty({m, n}, a.options().dtype(torch::kBFloat16));
     FP8Params p;
-    pack_gemm(p, a_st.data_ptr(), b_st.data_ptr(), output.data_ptr(), scale,
+    p.a_ptr = a_st.data_ptr();
-              m, n, k, a_ld, b_ld);
+    p.b_ptr = b_st.data_ptr();
+    p.out_ptr = output.data_ptr();
+    p.scale = scale.data_ptr<float>();
+    p.m = static_cast<int>(m);
+    p.n = static_cast<int>(n);
+    p.k = static_cast<int>(k);
+    p.a_ld = static_cast<int>(a_ld);
+    p.b_ld = static_cast<int>(b_ld);
     // Fused epilogue bias (bf16, broadcast over rows and batches). An
     // undefined or 0-element tensor keeps the plain scaled output.
     if (bias.defined() && bias.numel() > 0) {
@@ -243,9 +223,8 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
     return output;
 }
 
-// mm_fp8 binding: Python None and an omitted argument both mean "no bias"
+// mm_fp8 binding: Python None and an omitted argument both mean "no bias",
-// (resolved to an undefined tensor here, so every Python layer can pass its
+// so every Python layer can pass its bias argument through untouched.
-// 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("layout") = 0);

csrc/kernels/fp8/quantize.cuh

@@ -1,10 +1,7 @@
 #pragma once
-// FP8 quantize device code — pure CUDA, no torch. Any float input element
+// FP8 quantize device code — pure CUDA, no torch: kernels take the
-// type (bf16 / fp16 / fp32) converts to E4M3 or E5M2 with a fused amax over
+// FP8QuantizeParams POD, format and input type ride on template parameters,
-// the raw (unscaled) values. Mirrors the GEMM file's split: kernels take the
+// and the launcher is shared by the torch binding and the C tests.
-// FP8QuantizeParams POD, formats and input types ride on template parameters,
-// and the launcher is a plain function usable from both the torch binding and
-// pure C tests.
 
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
@@ -18,8 +15,8 @@
 namespace astrai {
 namespace fp8 {
 
-// Input element type traits: one element -> float, and the vectorized
+// Input element type traits: one element -> float, and the unpack of one
-// unpack of one 16-byte load into kVecElems floats.
+// 16-byte load into kVecElems floats.
 template <typename InT>
 struct quant_in_traits;
 
@@ -31,12 +28,13 @@ struct quant_in_traits<__nv_bfloat16> {
     }
     static __device__ __forceinline__ void load_vec(const uint4& raw,
                                                     float* f) {
-        const unsigned w[4] = {raw.x, raw.y, raw.z, raw.w};
+        const __nv_bfloat162* b2 =
+            reinterpret_cast<const __nv_bfloat162*>(&raw);
 #pragma unroll
         for (int j = 0; j < 4; ++j) {
-            f[2 * j] =
+            const float2 p = __bfloat1622float2(b2[j]);
-                __bfloat162float(__ushort_as_bfloat16(w[j] & 0xffffu));
+            f[2 * j] = p.x;
-            f[2 * j + 1] = __bfloat162float(__ushort_as_bfloat16(w[j] >> 16));
+            f[2 * j + 1] = p.y;
         }
     }
 };
@@ -65,15 +63,22 @@ struct quant_in_traits<float> {
     static __device__ __forceinline__ float to_float(float v) { return v; }
     static __device__ __forceinline__ void load_vec(const uint4& raw,
                                                     float* f) {
-        f[0] = __uint_as_float(raw.x);
+        const unsigned* w = reinterpret_cast<const unsigned*>(&raw);
-        f[1] = __uint_as_float(raw.y);
+#pragma unroll
-        f[2] = __uint_as_float(raw.z);
+        for (int j = 0; j < 4; ++j) f[j] = __uint_as_float(w[j]);
-        f[3] = __uint_as_float(raw.w);
     }
 };
 
-// Convert one float pair to one packed fp8 pair. The stored bytes see
+// One float -> one fp8 byte (round-nearest-even + satfinite).
-// value * mult (round-nearest-even + satfinite).
+template <FP8Format Fmt>
+__device__ __forceinline__ uint8_t cvt_fp8(float v) {
+    if constexpr (Fmt == FP8Format::E5M2)
+        return __nv_fp8_e5m2(v).__x;
+    else
+        return __nv_fp8_e4m3(v).__x;
+}
+
+// One float pair -> one packed fp8x2 word (round-nearest-even + satfinite).
 template <FP8Format Fmt>
 __device__ __forceinline__ unsigned cvt_fp8x2(float a, float b) {
     constexpr __nv_fp8_interpretation_t kFmt =
@@ -82,23 +87,36 @@ __device__ __forceinline__ unsigned cvt_fp8x2(float a, float b) {
         make_float2(a, b), __NV_SATFINITE, kFmt));
 }
 
-// Quantize kernel: float input -> FP8 (E4M3 or E5M2), fused amax over raw
+// Block-wide amax reduce -> one atomic per block: warp-reduce, park one
-// values.
+// value per warp, thread 0 folds. kWarps must cover the block's warp count.
+template <int kWarps>
+__device__ __forceinline__ void publish_amax(float* amax, float v) {
+    v = warp_reduce_max(v);
+    __shared__ float slots[kWarps];
+    const int tid = threadIdx.y * blockDim.x + threadIdx.x;
+    if ((tid & 31) == 0) slots[tid >> 5] = v;
+    __syncthreads();
+    if (tid == 0) {
+#pragma unroll
+        for (int w = 1; w < kWarps; ++w) v = fmaxf(v, slots[w]);
+        atomic_max_float(amax, v);
+    }
+}
+
+// Elementwise quantize kernel (out_layout 0): vectorized 16B loads -> fp8
+// stores, fused amax over raw values.
 template <FP8Format Fmt, typename InT>
 __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
     const float mult = *p.scale;
     const auto* x = static_cast<const InT*>(p.input_ptr);
-    void* x8 = p.output_ptr;
+    uint8_t* x8 = static_cast<uint8_t*>(p.output_ptr);
-    float* amax = p.amax;
     float local_amax = 0.0f;
     const int64_t stride = (int64_t)blockDim.x * gridDim.x;
 
-    // Vectorized body: one 16B load -> kVecElems fp8 bytes per step (8
+    // One 16B load -> kVecElems bytes per step. Torch allocations are >=16B
-    // elements for 16-bit inputs, 4 for fp32). Torch allocations are >=16B
+    // aligned, so element 0 keeps the uint4 access natural; a misaligned
-    // aligned and the binding passes freshly allocated contiguous buffers,
+    // base (odd storage offset view) falls to the scalar tail via
-    // so element 0 keeps the uint4 access natural; a misaligned base
+    // total_vec = 0.
-    // (contiguous view with an odd storage offset) falls back to the scalar
-    // loop below via total_vec = 0.
     constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
     const bool aligned =
         ((reinterpret_cast<uintptr_t>(x) |
@@ -125,8 +143,7 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
             packed[j] = (lo & 0xffffu) | (hi << 16);
         }
         if constexpr (kVecElems == 8)
-            reinterpret_cast<uint2*>(x8)[i] =
+            reinterpret_cast<uint2*>(x8)[i] = make_uint2(packed[0], packed[1]);
-                make_uint2(packed[0], packed[1]);
         else
             reinterpret_cast<unsigned*>(x8)[i] = packed[0];
     }
@@ -136,51 +153,20 @@ __global__ void fp8_quantize_kernel(FP8QuantizeParams p) {
          i < p.total; i += stride) {
         const float v = quant_in_traits<InT>::to_float(x[i]);
         local_amax = fmaxf(local_amax, fabsf(v));
-        if constexpr (Fmt == FP8Format::E5M2) {
+        x8[i] = cvt_fp8<Fmt>(v * mult);
-            reinterpret_cast<__nv_fp8_e5m2*>(x8)[i] =
-                __nv_fp8_e5m2(v * mult);
-        } else {
-            reinterpret_cast<__nv_fp8_e4m3*>(x8)[i] =
-                __nv_fp8_e4m3(v * mult);
-        }
-    }
-    if (amax) {
-        local_amax = warp_reduce_max(local_amax);
-        __shared__ float slots[32];
-        if ((threadIdx.x & 31) == 0) slots[threadIdx.x >> 5] = local_amax;
-        __syncthreads();
-        if (threadIdx.x == 0) {
-            float v = 0.0f;
-            for (int w = 0; w < (blockDim.x >> 5); ++w)
-                v = fmaxf(v, slots[w]);
-            atomic_max_float(amax, v);
-        }
     }
+    if (p.amax) publish_amax<8>(p.amax, local_amax);
 }
 
-template <FP8Format Fmt, typename InT>
+// Tiled transpose quantize (out_layout 1/2): reads the [rows][cols] input
-void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
+// once and writes the fp8 bytes transposed ([cols][rows], so the contract
-    constexpr int kThreads = 256;
+// dim lands K-contiguous for NT GEMM operands) and, in mode 2, the row-major
-    // One block per 256 vectors; at least one block so the scalar tail of a
+// copy too. A 32x32 tile stages through shared memory: loads and writes
-    // tiny / misaligned tensor is still covered.
+// both stay coalesced, and the byte-wide staging is conflict-free — the +4
-    constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
+// pad makes the store stride 9 words (coprime with the 32 banks) and the
-    int64_t blocks = (p.total / kVecElems + kThreads - 1) / kThreads;
+// read is a 32-byte broadcast segment. (A 64x64 split-half variant measured
-    if (blocks < 1) blocks = 1;
+// +21% L2-resident but -3..5% DRAM-bound; the real step mix ties, so the
-    fp8_quantize_kernel<Fmt, InT><<<blocks, kThreads, 0, stream>>>(p);
+// simpler tile stays.)
-}
-
-// 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;
@@ -201,10 +187,7 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
             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] = cvt_fp8<Fmt>(v * mult);
-                q[j] = __nv_fp8_e5m2(v * mult).__x;
-            else
-                q[j] = __nv_fp8_e4m3(v * mult).__x;
         }
     }
     if (p.out_layout == 2) {
@@ -218,9 +201,9 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
     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
+    // tracks threadIdx.x so each warp writes one contiguous run. tile was
-    // swaps the staging indices — tile[col][row] was written, so the value
+    // written as tile[col][row], so input (r0+tx, c0+ty*4+j) reads back
-    // for input (r0+tx, c0+ty*4+j) sits at tile[ty*4+j][tx].
+    // from 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) {
@@ -229,28 +212,27 @@ __global__ void fp8_quantize_tiled_kernel(FP8QuantizeParams p) {
             out_t[(int64_t)oc * p.rows + r0 + threadIdx.x] =
                 tile[threadIdx.y * 4 + j][threadIdx.x];
     }
-    if (p.amax) {
+    if (p.amax) publish_amax<8>(p.amax, local_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>
+// Unified quantize launcher: Tiled selects the transpose kernel (out_layout
-void launch_fp8_quantize_tiled(const FP8QuantizeParams& p,
+// 1/2) over the vectorized elementwise one.
-                               cudaStream_t stream) {
+template <FP8Format Fmt, typename InT, bool Tiled = false>
-    const dim3 grid((p.cols + 31) / 32, (p.rows + 31) / 32);
+void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) {
-    if (grid.x == 0 || grid.y == 0) return;
+    if constexpr (Tiled) {
-    fp8_quantize_tiled_kernel<Fmt, InT>
+        const dim3 grid((p.cols + 31) / 32, (p.rows + 31) / 32);
-        <<<grid, dim3(32, 8), 0, stream>>>(p);
+        if (grid.x == 0 || grid.y == 0) return;
+        fp8_quantize_tiled_kernel<Fmt, InT>
+            <<<grid, dim3(32, 8), 0, stream>>>(p);
+    } else {
+        constexpr int kThreads = 256;
+        constexpr int kVecElems = quant_in_traits<InT>::kVecElems;
+        // One block per 256 vectors; at least one block so a tiny or
+        // misaligned tensor's scalar tail is still covered.
+        int64_t blocks = (p.total / kVecElems + kThreads - 1) / kThreads;
+        if (blocks < 1) blocks = 1;
+        fp8_quantize_kernel<Fmt, InT><<<blocks, kThreads, 0, stream>>>(p);
+    }
 }
 
 }  // namespace fp8