"""FP8 e4m3 x e4m3 GEMM with per-output-channel dequant scale.

y = (x @ weight.T) * weight_scale in bf16.

Engines:
  1. CUTLASS SM90 warp-specialized FP8 GEMM (TMA + wgmma, FP32 accumulate,
     custom EVT epilogue applying the per-N-channel scale) for compute-bound
     shapes. Vendored headers live in `cutlass_3.9/` next to this file; the
     extension is JIT-built on first use and silently falls back if it can't
     be built (the Triton path below handles every shape on its own).
  2. Triton tl.dot fp8 kernel (tensor-core MMA, fp32 accumulate) for
     skinny-M shapes and as a universal fallback.
  3. K not a multiple of 16: operands are zero-padded along K (the pad adds
     zeros, which is mathematically exact) and the weight pad is cached
     (input-derived, invalidated by tensor version).

Engine choice per (M, N, K) is made by a tiny one-shot micro-tuner on first
encounter of a shape, so any shape gets the best available backend.
"""
import os
import time
import torch
import torch.nn as nn
import triton
import triton.language as tl

E4M3_MAX = 448.0

# ---------------------------------------------------------------------------
# Triton kernels
# ---------------------------------------------------------------------------


def _gen_configs():
    cfgs = []
    for bm, bn, bk, g, w, s in [
        (128, 128, 64, 8, 8, 4),
        (128, 128, 128, 8, 8, 3),
        (128, 256, 64, 8, 8, 3),
        (128, 256, 128, 8, 8, 3),
        (64, 128, 128, 4, 4, 4),
        (64, 256, 128, 4, 4, 4),
        (128, 64, 128, 4, 4, 4),
    ]:
        cfgs.append(
            triton.Config(
                {"BM": bm, "BN": bn, "BK": bk, "GROUP_M": g}, num_warps=w, num_stages=s
            )
        )
    for bm, bn, bk, g, w, s in [
        (16, 128, 128, 1, 4, 4),
        (16, 128, 128, 1, 4, 5),
        (16, 128, 256, 1, 4, 3),
        (32, 128, 128, 1, 4, 5),
    ]:
        cfgs.append(
            triton.Config(
                {"BM": bm, "BN": bn, "BK": bk, "GROUP_M": g}, num_warps=w, num_stages=s
            )
        )
    return cfgs


@triton.autotune(configs=_gen_configs(), key=["M", "N", "K"])
@triton.jit
def _fp8_gemm_kernel(
    a_ptr, b_ptr, c_ptr, s_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bn, stride_bk,
    stride_cm, stride_cn,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, GROUP_M: tl.constexpr,
):
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BM)
    num_pid_n = tl.cdiv(N, BN)
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_m = pid_m * BM + tl.arange(0, BM)
    offs_n = pid_n * BN + tl.arange(0, BN)
    offs_k = tl.arange(0, BK)
    m_mask = offs_m < M
    n_mask = offs_n < N

    a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
    b_ptrs = b_ptr + offs_n[None, :] * stride_bn + offs_k[:, None] * stride_bk

    acc = tl.zeros((BM, BN), dtype=tl.float32)
    k_full = K // BK
    for _ in range(0, k_full):
        a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
        b = tl.load(b_ptrs, mask=n_mask[None, :], other=0.0)
        # max_num_imprecise_acc=BK: promote from imprecise wgmma accumulation
        # into fp32 FFMA every block (correctness for long K; ~5-10% cost).
        acc = tl.dot(a, b, acc, max_num_imprecise_acc=BK)
        a_ptrs += BK * stride_ak
        b_ptrs += BK * stride_bk
    if K % BK != 0:
        k_rem = K - k_full * BK
        a = tl.load(
            a_ptrs, mask=m_mask[:, None] & (offs_k[None, :] < k_rem), other=0.0
        )
        b = tl.load(
            b_ptrs, mask=n_mask[None, :] & (offs_k[:, None] < k_rem), other=0.0
        )
        acc = tl.dot(a, b, acc, max_num_imprecise_acc=BK)

    s = tl.load(s_ptr + offs_n, mask=n_mask, other=0.0)
    c = acc * s[None, :]
    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(c_ptrs, c.to(tl.bfloat16), mask=m_mask[:, None] & n_mask[None, :])


def _triton_gemm(x, w, w_scale, y=None):
    M, K = x.shape
    N = w.shape[0]
    if y is None:
        y = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
    grid = lambda meta: (  # noqa: E731
        triton.cdiv(M, meta["BM"]) * triton.cdiv(N, meta["BN"]),
    )
    _fp8_gemm_kernel[grid](
        x, w, y, w_scale,
        M, N, K,
        x.stride(0), x.stride(1),
        w.stride(0), w.stride(1),
        y.stride(0), y.stride(1),
    )
    return y


# --- K-padding kernels ------------------------------------------------------


@triton.jit
def _pad_k_kernel(src, dst, R, K, K_PADDED, BK: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_k = tl.program_id(1)
    offs_k = pid_k * BK + tl.arange(0, BK)
    mask = offs_k < K
    val = tl.load(src + pid_m * K + offs_k, mask=mask, other=0.0)
    tl.store(dst + pid_m * K_PADDED + offs_k, val, mask=offs_k < K_PADDED)


@triton.jit
def _pad_k_shl(src_i32, dst_i32, R, K, K_PAD, TOTAL_SRC_DW, BLOCK: tl.constexpr):
    # dst-side aligned int32 ops; src bytes gathered via funnel shift of two
    # aligned dwords (little-endian). Zeros beyond K.
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    kpad_dw = K_PAD // 4
    total_dst_dw = R * kpad_dw
    valid = offs < total_dst_dw
    r = offs // kpad_dw
    c = (offs % kpad_dw) * 4
    p = r * K + c
    inside = c < K
    d0 = p // 4
    sh8 = (p % 4) * 8
    v0 = tl.load(src_i32 + d0, mask=valid & inside, other=0)
    v1 = tl.load(
        src_i32 + d0 + 1, mask=valid & inside & (sh8 > 0) & (d0 + 1 < TOTAL_SRC_DW),
        other=0,
    )
    sh32 = tl.where(sh8 > 0, 32 - sh8, 0)
    u0 = v0.to(tl.uint32, bitcast=True)
    u1 = v1.to(tl.uint32, bitcast=True)
    val = (u0 >> sh8) | tl.where(sh8 > 0, u1 << sh32, 0)
    rem = K - c
    keep = tl.where(rem >= 4, -1, (1 << (rem * 8)) - 1).to(tl.uint32, bitcast=True)
    val = tl.where(inside & valid, val & keep, 0)
    tl.store(dst_i32 + offs, val.to(tl.int32, bitcast=True), mask=valid)


def _pad_k(t, k_pad):
    """(R, K) opaque 1-byte-dtype -> (R, k_pad), zero-filled tail."""
    R, K = t.shape
    out = torch.empty((R, k_pad), dtype=t.dtype, device=t.device)
    if K >= 4 and (R * K) % 4 == 0 and k_pad % 4 == 0:
        src = t.view(-1).view(torch.int32)
        dst = out.view(-1).view(torch.int32)
        BLOCK = 512
        grid = (triton.cdiv(R * (k_pad // 4), BLOCK),)
        _pad_k_shl[grid](src, dst, R, K, k_pad, (R * K + 3) // 4, BLOCK=BLOCK,
                         num_warps=4)
        return out
    BK = 512
    grid = (R, triton.cdiv(k_pad, BK))
    _pad_k_kernel[grid](t, out, R, K, k_pad, BK=BK)
    return out


class _WPadCache:
    """Cached zero-padded-K copy of a weight tensor, invalidated by version."""

    def __init__(self):
        self._entry = None

    def get(self, w, k_pad):
        key = (w.data_ptr(), w._version, tuple(w.shape), k_pad, w.device.index)
        if self._entry is not None and self._entry[0] == key:
            return self._entry[1]
        t = _pad_k(w, k_pad)
        self._entry = (key, t)
        return t


_w_pad_cache = _WPadCache()

# ---------------------------------------------------------------------------
# CUTLASS extension (optional, JIT-built from vendored headers)
# ---------------------------------------------------------------------------

_HERE = os.path.dirname(os.path.abspath(__file__))
_CUDA_SRC = r"""
// FP8 e4m3 x e4m3 GEMM: D = (A @ B^T) * scale_n, bf16 out.
// A: (M,K) fp8 row-major (K-contig). B: (N,K) fp8 row-major (K-contig)
//   == CUTLASS LayoutB ColumnMajor (K,N). scale: (N,) fp32. D: (M,N) bf16 row-major.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cute/tensor.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/epilogue/dispatch_policy.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp"

using namespace cute;

namespace fp8gemm {

using ElementA = cutlass::float_e4m3_t;
using LayoutA = cutlass::layout::RowMajor;
constexpr int AlignmentA = 16;
using ElementB = cutlass::float_e4m3_t;
using LayoutB = cutlass::layout::ColumnMajor;
constexpr int AlignmentB = 16;

using ElementC = cutlass::bfloat16_t;
using LayoutC = cutlass::layout::RowMajor;
constexpr int AlignmentC = 8;
using ElementD = cutlass::bfloat16_t;
using LayoutD = cutlass::layout::RowMajor;
constexpr int AlignmentD = 8;

using ElementAccumulator = float;
using ElementCompute = float;
using ArchTag = cutlass::arch::Sm90;
using OperatorClass = cutlass::arch::OpClassTensorOp;

template <class TileShapeMNK, class ClusterShapeMNK, class KernelSched,
          class EpiSched = cutlass::epilogue::TmaWarpSpecializedCooperative>
struct GemmTraits {
  using TileShape = TileShapeMNK;
  using ClusterShape = ClusterShapeMNK;
  using KernelSchedule = KernelSched;
  using EpilogueSchedule = EpiSched;
  static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;

  // D = acc * per-N scale (row broadcast of an (N,) vector, in fp32, cast bf16)
  using CustomEVT = cutlass::epilogue::fusion::Sm90EVT<
      cutlass::epilogue::fusion::Sm90Compute<cutlass::multiplies, ElementD, ElementCompute, RoundStyle>,
      cutlass::epilogue::fusion::Sm90RowBroadcast<
          0, TileShapeMNK, float, ElementCompute, Stride<_0, _1, int64_t>>,
      cutlass::epilogue::fusion::Sm90AccFetch>;

  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
      ArchTag, OperatorClass,
      TileShape, ClusterShape,
      cutlass::epilogue::collective::EpilogueTileAuto,
      ElementAccumulator, ElementCompute,
      ElementC, LayoutC, AlignmentC,
      ElementD, LayoutD, AlignmentD,
      EpilogueSchedule,
      CustomEVT>::CollectiveOp;

  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
      ArchTag, OperatorClass,
      ElementA, LayoutA, AlignmentA,
      ElementB, LayoutB, AlignmentB,
      ElementAccumulator,
      TileShape, ClusterShape,
      cutlass::gemm::collective::StageCountAutoCarveout<
          static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
      KernelSchedule>::CollectiveOp;

  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
      cutlass::gemm::PersistentScheduler>;

  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};

template <class Traits>
inline cutlass::Status gemm_can_implement(const typename Traits::Gemm::Arguments& args) {
  return Traits::Gemm::can_implement(args);
}

template <class Traits>
void run_gemm(const torch::Tensor& a, const torch::Tensor& b,
              const torch::Tensor& scale, torch::Tensor& d,
              int max_swizzle = 1, int raster_dir = 1, int promo_interval = 4) {
  using Gemm = typename Traits::Gemm;
  int M = a.size(0), K = a.size(1), N = b.size(0);

  typename Gemm::Arguments args{
      cutlass::gemm::GemmUniversalMode::kGemm,
      {M, N, K, 1},
      {(const ElementA*)a.data_ptr(), cute::make_stride(int64_t(K), Int<1>{}, int64_t(0)),
       (const ElementB*)b.data_ptr(), cute::make_stride(int64_t(K), Int<1>{}, int64_t(0))},
      {// epilogue
       {  // fusion args (EVT): {row_broadcast args, acc args(empty), compute args(empty)}
           {(const float*)scale.data_ptr(), 0.0f, {_0{}, _1{}, int64_t(0)}},
           {},
           {},
       },
       (const ElementC*)d.data_ptr(),  // C unused; alias D (beta path never taken)
       cute::make_stride(int64_t(N), Int<1>{}, int64_t(0)),
       (ElementD*)d.data_ptr(),
       cute::make_stride(int64_t(N), Int<1>{}, int64_t(0))}};

  args.scheduler.max_swizzle_size = max_swizzle;
  args.scheduler.raster_order =
      raster_dir == 0
          ? cutlass::gemm::kernel::detail::RasterOrderOptions::AlongM
          : cutlass::gemm::kernel::detail::RasterOrderOptions::AlongN;

  // Fast-accum FFMA promotion interval (multiple of TileK/GMMA_K = 4)
  args.mainloop.mma_promotion_interval = (uint32_t)promo_interval;

  // can_implement + workspace size are pointer-independent; cache per shape/schedule.
  struct ShapeCache {
    cutlass::Status status = cutlass::Status::kInvalid;
    size_t ws_size = 0;
    int max_swizzle = -1;
    int raster_dir = -1;
    int M = -1, N = -1, K = -1;
  };
  static thread_local ShapeCache cache;
  if (cache.M != M || cache.N != N || cache.K != K ||
      cache.max_swizzle != max_swizzle || cache.raster_dir != raster_dir) {
    cache.status = gemm_can_implement<Traits>(args);
    cache.ws_size = cache.status == cutlass::Status::kSuccess
                        ? Gemm::get_workspace_size(args)
                        : 0;
    cache.max_swizzle = max_swizzle;
    cache.raster_dir = raster_dir;
    cache.M = M; cache.N = N; cache.K = K;
  }
  TORCH_CHECK(cache.status == cutlass::Status::kSuccess,
              "CUTLASS can_implement failed: ", (int)cache.status);
  torch::Tensor ws;
  void* ws_ptr = nullptr;
  if (cache.ws_size > 0) {
    ws = torch::empty({(int64_t)cache.ws_size},
                      torch::TensorOptions().dtype(torch::kUInt8).device(a.device()));
    ws_ptr = ws.data_ptr();
  }
  Gemm gemm;
  auto status = gemm.initialize(args, ws_ptr);
  TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS initialize failed: ",
              (int)status);
  status = gemm.run(at::cuda::getCurrentCUDAStream());
  TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS run failed: ", (int)status);
}

}  // namespace fp8gemm

torch::Tensor fp8_gemm_scaled(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
                              int64_t variant, int64_t max_swizzle, int64_t raster_dir,
                              int64_t promo_interval) {
  TORCH_CHECK(a.is_cuda() && b.is_cuda() && scale.is_cuda());
  TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn && b.dtype() == torch::kFloat8_e4m3fn);
  TORCH_CHECK(scale.dtype() == torch::kFloat32);
  TORCH_CHECK(a.is_contiguous() && b.is_contiguous() && scale.is_contiguous());
  const auto& a_c = a;
  const auto& b_c = b;
  const auto& s_c = scale;
  int M = a_c.size(0), N = b_c.size(0);
  auto d = torch::empty({M, N}, torch::TensorOptions()
                                       .dtype(torch::kBFloat16)
                                       .device(a_c.device()));
  using COOP_FA = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
  using T1 = fp8gemm::GemmTraits<Shape<_128, _128, _128>, Shape<_2, _1, _1>, COOP_FA>;
  using T2 = fp8gemm::GemmTraits<Shape<_128, _256, _128>, Shape<_2, _2, _1>, COOP_FA>;
  using T3 = T1;  // was: COOP non-fast-accum (slower)
  using T4 = T1;  // was: pingpong (broken with custom EVT)
  using T5 = fp8gemm::GemmTraits<Shape<_128, _128, _128>, Shape<_1, _1, _1>, COOP_FA>;
  using T6 = fp8gemm::GemmTraits<Shape<_256, _128, _128>, Shape<_2, _1, _1>, COOP_FA>;
  using T7 = fp8gemm::GemmTraits<Shape<_128, _128, _128>, Shape<_4, _1, _1>, COOP_FA>;
  using T8 = T1;
  using T9 = T1;
  using T10 = T1;
  using T11 = fp8gemm::GemmTraits<Shape<_256, _128, _128>, Shape<_2, _2, _1>, COOP_FA>;
  using T12 = fp8gemm::GemmTraits<Shape<_256, _128, _128>, Shape<_4, _1, _1>, COOP_FA>;
  using T13 = fp8gemm::GemmTraits<Shape<_128, _128, _256>, Shape<_2, _1, _1>, COOP_FA>;
  using T14 = fp8gemm::GemmTraits<Shape<_128, _256, _128>, Shape<_2, _1, _1>, COOP_FA>;
  using T15 = T1;
  using T16 = T1;
  using T17 = T1;
  using T18 = T1;
  using T19 = T1;
  switch (variant) {
    case 0: fp8gemm::run_gemm<T1>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 1: fp8gemm::run_gemm<T2>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 2: fp8gemm::run_gemm<T3>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 3: fp8gemm::run_gemm<T4>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 4: fp8gemm::run_gemm<T5>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 5: fp8gemm::run_gemm<T6>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 6: fp8gemm::run_gemm<T7>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 7: fp8gemm::run_gemm<T8>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 8: fp8gemm::run_gemm<T9>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 9: fp8gemm::run_gemm<T10>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 10: fp8gemm::run_gemm<T11>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 11: fp8gemm::run_gemm<T12>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 12: fp8gemm::run_gemm<T13>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 13: fp8gemm::run_gemm<T14>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 14: fp8gemm::run_gemm<T15>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 15: fp8gemm::run_gemm<T16>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 16: fp8gemm::run_gemm<T17>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 17: fp8gemm::run_gemm<T18>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 18: fp8gemm::run_gemm<T19>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    default: fp8gemm::run_gemm<T1>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval);
  }
  return d;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("fp8_gemm_scaled", &fp8_gemm_scaled, "fp8 gemm with per-n scale");
}

"""

_CUTLASS_DIR = os.path.join(_HERE, "cutlass_3.9")
_EXT = None
_EXT_FAILED = False


def _load_cutlass_ext():
    global _EXT, _EXT_FAILED
    if _EXT is not None or _EXT_FAILED:
        return _EXT
    try:
        from torch.utils.cpp_extension import load

        build_dir = os.path.join(_HERE, "build")
        os.makedirs(build_dir, exist_ok=True)
        # torch's jit-compile lock persists after killed processes; a lock
        # whose mtime is >15 min old is certainly abandoned — remove it.
        lock_path = os.path.join(build_dir, "lock")
        try:
            if os.path.exists(lock_path):
                age = time.time() - os.path.getmtime(lock_path)
                if age > 900:
                    os.remove(lock_path)
        except OSError:
            pass
        src_path = os.path.join(_HERE, "fp8_gemm_cutlass.cu")
        try:
            if not os.path.exists(src_path) or \
                    open(src_path).read().strip() != _CUDA_SRC.strip():
                with open(src_path, "w") as f:
                    f.write(_CUDA_SRC.strip() + "\n")
        except OSError:
            pass
        build_dir = os.path.join(_HERE, "build")
        os.makedirs(build_dir, exist_ok=True)
        _EXT = load(
            name="fp8_gemm_cutlass",
            sources=[os.path.join(_HERE, "fp8_gemm_cutlass.cu")],
            extra_include_paths=[os.path.join(_CUTLASS_DIR, "include")],
            extra_cuda_cflags=[
                "-O3", "-std=c++17", "--expt-relaxed-constexpr",
                "--expt-extended-lambda", "-gencode=arch=compute_90a,code=sm_90a",
                "--use_fast_math", "-DNDEBUG",
            ],
            extra_cflags=["-O3", "-std=c++17"],
            build_directory=build_dir,
            verbose=False,
        )
    except Exception:
        _EXT = None
        _EXT_FAILED = True
    return _EXT


# CUTLASS (variant, max_swizzle, raster_order) candidates per problem size.
def _cutlass_candidates(M, N, K):
    return [(0, 1, 1), (4, 1, 1), (5, 4, 1), (5, 8, 1), (1, 1, 1),
            (10, 1, 1), (11, 1, 1), (7, 1, 1), (13, 1, 1)]


# ---------------------------------------------------------------------------
# Dispatch with one-shot micro-autotune per shape
# ---------------------------------------------------------------------------

_plan_cache = {}
_l2_scratch = None


def _flush_l2():
    global _l2_scratch
    if _l2_scratch is None:
        _l2_scratch = torch.empty(128 * 1024 * 1024 // 4, dtype=torch.float32, device="cuda")
    _l2_scratch.zero_()


def _bench_once(fn, iters=9):
    fn()  # warm
    torch.cuda.synchronize()
    times = []
    for _ in range(iters):
        _flush_l2()
        torch.cuda.synchronize()
        s = torch.cuda.Event(enable_timing=True)
        e = torch.cuda.Event(enable_timing=True)
        s.record()
        fn()
        e.record()
        torch.cuda.synchronize()
        times.append(s.elapsed_time(e))
    times.sort()
    return times[len(times) // 2]


def _make_plan(x, w, w_scale):
    """Choose the fastest engine for this (M, N, K). Returns a callable."""
    M, K = x.shape
    N = w.shape[0]
    candidates = []

    ext = None
    if K % 16 == 0 and (2 * N) % 16 == 0:
        ext = _load_cutlass_ext()
    if ext is not None:
        for v, sw, ro in _cutlass_candidates(M, N, K):
            def cutlass_cand(a, b, c, v=v, sw=sw, ro=ro):
                return ext.fp8_gemm_scaled(a, b, c, v, sw, ro, 4)
            candidates.append(("cutlass", cutlass_cand))

    if K % 16 != 0:
        # zero-pad operands to a 16B-aligned K (exact: pad contributes 0s)
        k_pad = (K + 127) // 128 * 128
        ext2 = _load_cutlass_ext()

        def triton_padded(a, b, c, k_pad=k_pad):
            wp = _w_pad_cache.get(b, k_pad)  # re-validated against b version
            xp = _pad_k(a, k_pad)
            return _triton_gemm(xp, wp, c)

        candidates.append(("triton_pad", triton_padded))

        if ext2 is not None and M * N >= 1 << 22:
            for v, sw, ro in [(0, 1, 1), (10, 1, 1), (5, 4, 1)]:
                def cutlass_pad(a, b, c, k_pad=k_pad, v=v, sw=sw, ro=ro):
                    wp = _w_pad_cache.get(b, k_pad)
                    xp = _pad_k(a, k_pad)
                    return ext2.fp8_gemm_scaled(xp, wp, c, v, sw, ro, 4)
                candidates.append(("cutlass_pad", cutlass_pad))

    # Plain triton always works.
    candidates.append(("triton", _triton_gemm))

    if len(candidates) == 1:
        return candidates[0][1]

    best, best_t, best_name = candidates[-1][1], float("inf"), None
    for name, cand in candidates:
        try:
            cand(x, w, w_scale)  # correctness smoke (throws on failure)
            t = _bench_once(lambda: cand(x, w, w_scale))
        except Exception:
            continue
        if t < best_t:
            best, best_t, best_name = cand, t, name
    return best


def fp8_gemm(x, w, w_scale):
    if not x.is_contiguous():
        x = x.contiguous()
    if not w.is_contiguous():
        w = w.contiguous()
    M, K = x.shape
    N = w.shape[0]
    key = (M, N, K, x.device.index)
    plan = _plan_cache.get(key)
    if plan is None:
        plan = _make_plan(x, w, w_scale)
        _plan_cache[key] = plan
    return plan(x, w, w_scale)


class Model(nn.Module):
    """Same interface as reference.Model: buffers `weight` (N,K) fp8 and
    `weight_scale` (N,) fp32; forward returns y = (x @ w.T) * scale in bf16."""

    def __init__(self, M: int, N: int, K: int):
        super().__init__()
        self.M, self.N, self.K = M, N, K
        self.register_buffer("weight", torch.zeros(N, K, dtype=torch.float8_e4m3fn))
        self.register_buffer("weight_scale", torch.ones(N, dtype=torch.float32))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return fp8_gemm(x, self.weight, self.weight_scale)


M = 4096
N = 4096
K = 4096


def get_inputs():
    x = (torch.rand(M, K) * 8 - 4).to(torch.float8_e4m3fn)
    return [x]


def get_init_inputs():
    return [M, N, K]


# ==================================================================
# ===== sidecar: fp8_gemm_cutlass.cu (11027 bytes, loaded by solution.py) =====
# ==================================================================

// FP8 e4m3 x e4m3 GEMM: D = (A @ B^T) * scale_n, bf16 out.
// A: (M,K) fp8 row-major (K-contig). B: (N,K) fp8 row-major (K-contig)
//   == CUTLASS LayoutB ColumnMajor (K,N). scale: (N,) fp32. D: (M,N) bf16 row-major.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cute/tensor.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/epilogue/dispatch_policy.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp"

using namespace cute;

namespace fp8gemm {

using ElementA = cutlass::float_e4m3_t;
using LayoutA = cutlass::layout::RowMajor;
constexpr int AlignmentA = 16;
using ElementB = cutlass::float_e4m3_t;
using LayoutB = cutlass::layout::ColumnMajor;
constexpr int AlignmentB = 16;

using ElementC = cutlass::bfloat16_t;
using LayoutC = cutlass::layout::RowMajor;
constexpr int AlignmentC = 8;
using ElementD = cutlass::bfloat16_t;
using LayoutD = cutlass::layout::RowMajor;
constexpr int AlignmentD = 8;

using ElementAccumulator = float;
using ElementCompute = float;
using ArchTag = cutlass::arch::Sm90;
using OperatorClass = cutlass::arch::OpClassTensorOp;

template <class TileShapeMNK, class ClusterShapeMNK, class KernelSched,
          class EpiSched = cutlass::epilogue::TmaWarpSpecializedCooperative>
struct GemmTraits {
  using TileShape = TileShapeMNK;
  using ClusterShape = ClusterShapeMNK;
  using KernelSchedule = KernelSched;
  using EpilogueSchedule = EpiSched;
  static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;

  // D = acc * per-N scale (row broadcast of an (N,) vector, in fp32, cast bf16)
  using CustomEVT = cutlass::epilogue::fusion::Sm90EVT<
      cutlass::epilogue::fusion::Sm90Compute<cutlass::multiplies, ElementD, ElementCompute, RoundStyle>,
      cutlass::epilogue::fusion::Sm90RowBroadcast<
          0, TileShapeMNK, float, ElementCompute, Stride<_0, _1, int64_t>>,
      cutlass::epilogue::fusion::Sm90AccFetch>;

  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
      ArchTag, OperatorClass,
      TileShape, ClusterShape,
      cutlass::epilogue::collective::EpilogueTileAuto,
      ElementAccumulator, ElementCompute,
      ElementC, LayoutC, AlignmentC,
      ElementD, LayoutD, AlignmentD,
      EpilogueSchedule,
      CustomEVT>::CollectiveOp;

  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
      ArchTag, OperatorClass,
      ElementA, LayoutA, AlignmentA,
      ElementB, LayoutB, AlignmentB,
      ElementAccumulator,
      TileShape, ClusterShape,
      cutlass::gemm::collective::StageCountAutoCarveout<
          static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
      KernelSchedule>::CollectiveOp;

  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
      cutlass::gemm::PersistentScheduler>;

  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};

template <class Traits>
inline cutlass::Status gemm_can_implement(const typename Traits::Gemm::Arguments& args) {
  return Traits::Gemm::can_implement(args);
}

template <class Traits>
void run_gemm(const torch::Tensor& a, const torch::Tensor& b,
              const torch::Tensor& scale, torch::Tensor& d,
              int max_swizzle = 1, int raster_dir = 1, int promo_interval = 4) {
  using Gemm = typename Traits::Gemm;
  int M = a.size(0), K = a.size(1), N = b.size(0);

  typename Gemm::Arguments args{
      cutlass::gemm::GemmUniversalMode::kGemm,
      {M, N, K, 1},
      {(const ElementA*)a.data_ptr(), cute::make_stride(int64_t(K), Int<1>{}, int64_t(0)),
       (const ElementB*)b.data_ptr(), cute::make_stride(int64_t(K), Int<1>{}, int64_t(0))},
      {// epilogue
       {  // fusion args (EVT): {row_broadcast args, acc args(empty), compute args(empty)}
           {(const float*)scale.data_ptr(), 0.0f, {_0{}, _1{}, int64_t(0)}},
           {},
           {},
       },
       (const ElementC*)d.data_ptr(),  // C unused; alias D (beta path never taken)
       cute::make_stride(int64_t(N), Int<1>{}, int64_t(0)),
       (ElementD*)d.data_ptr(),
       cute::make_stride(int64_t(N), Int<1>{}, int64_t(0))}};

  args.scheduler.max_swizzle_size = max_swizzle;
  args.scheduler.raster_order =
      raster_dir == 0
          ? cutlass::gemm::kernel::detail::RasterOrderOptions::AlongM
          : cutlass::gemm::kernel::detail::RasterOrderOptions::AlongN;

  // Fast-accum FFMA promotion interval (multiple of TileK/GMMA_K = 4)
  args.mainloop.mma_promotion_interval = (uint32_t)promo_interval;

  // can_implement + workspace size are pointer-independent; cache per shape/schedule.
  struct ShapeCache {
    cutlass::Status status = cutlass::Status::kInvalid;
    size_t ws_size = 0;
    int max_swizzle = -1;
    int raster_dir = -1;
    int M = -1, N = -1, K = -1;
  };
  static thread_local ShapeCache cache;
  if (cache.M != M || cache.N != N || cache.K != K ||
      cache.max_swizzle != max_swizzle || cache.raster_dir != raster_dir) {
    cache.status = gemm_can_implement<Traits>(args);
    cache.ws_size = cache.status == cutlass::Status::kSuccess
                        ? Gemm::get_workspace_size(args)
                        : 0;
    cache.max_swizzle = max_swizzle;
    cache.raster_dir = raster_dir;
    cache.M = M; cache.N = N; cache.K = K;
  }
  TORCH_CHECK(cache.status == cutlass::Status::kSuccess,
              "CUTLASS can_implement failed: ", (int)cache.status);
  torch::Tensor ws;
  void* ws_ptr = nullptr;
  if (cache.ws_size > 0) {
    ws = torch::empty({(int64_t)cache.ws_size},
                      torch::TensorOptions().dtype(torch::kUInt8).device(a.device()));
    ws_ptr = ws.data_ptr();
  }
  Gemm gemm;
  auto status = gemm.initialize(args, ws_ptr);
  TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS initialize failed: ",
              (int)status);
  status = gemm.run(at::cuda::getCurrentCUDAStream());
  TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS run failed: ", (int)status);
}

}  // namespace fp8gemm

torch::Tensor fp8_gemm_scaled(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
                              int64_t variant, int64_t max_swizzle, int64_t raster_dir,
                              int64_t promo_interval) {
  TORCH_CHECK(a.is_cuda() && b.is_cuda() && scale.is_cuda());
  TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn && b.dtype() == torch::kFloat8_e4m3fn);
  TORCH_CHECK(scale.dtype() == torch::kFloat32);
  TORCH_CHECK(a.is_contiguous() && b.is_contiguous() && scale.is_contiguous());
  const auto& a_c = a;
  const auto& b_c = b;
  const auto& s_c = scale;
  int M = a_c.size(0), N = b_c.size(0);
  auto d = torch::empty({M, N}, torch::TensorOptions()
                                       .dtype(torch::kBFloat16)
                                       .device(a_c.device()));
  using COOP_FA = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
  using T1 = fp8gemm::GemmTraits<Shape<_128, _128, _128>, Shape<_2, _1, _1>, COOP_FA>;
  using T2 = fp8gemm::GemmTraits<Shape<_128, _256, _128>, Shape<_2, _2, _1>, COOP_FA>;
  using T3 = T1;  // was: COOP non-fast-accum (slower)
  using T4 = T1;  // was: pingpong (broken with custom EVT)
  using T5 = fp8gemm::GemmTraits<Shape<_128, _128, _128>, Shape<_1, _1, _1>, COOP_FA>;
  using T6 = fp8gemm::GemmTraits<Shape<_256, _128, _128>, Shape<_2, _1, _1>, COOP_FA>;
  using T7 = fp8gemm::GemmTraits<Shape<_128, _128, _128>, Shape<_4, _1, _1>, COOP_FA>;
  using T8 = T1;
  using T9 = T1;
  using T10 = T1;
  using T11 = fp8gemm::GemmTraits<Shape<_256, _128, _128>, Shape<_2, _2, _1>, COOP_FA>;
  using T12 = fp8gemm::GemmTraits<Shape<_256, _128, _128>, Shape<_4, _1, _1>, COOP_FA>;
  using T13 = fp8gemm::GemmTraits<Shape<_128, _128, _256>, Shape<_2, _1, _1>, COOP_FA>;
  using T14 = fp8gemm::GemmTraits<Shape<_128, _256, _128>, Shape<_2, _1, _1>, COOP_FA>;
  using T15 = T1;
  using T16 = T1;
  using T17 = T1;
  using T18 = T1;
  using T19 = T1;
  switch (variant) {
    case 0: fp8gemm::run_gemm<T1>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 1: fp8gemm::run_gemm<T2>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 2: fp8gemm::run_gemm<T3>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 3: fp8gemm::run_gemm<T4>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 4: fp8gemm::run_gemm<T5>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 5: fp8gemm::run_gemm<T6>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 6: fp8gemm::run_gemm<T7>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 7: fp8gemm::run_gemm<T8>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 8: fp8gemm::run_gemm<T9>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 9: fp8gemm::run_gemm<T10>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 10: fp8gemm::run_gemm<T11>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 11: fp8gemm::run_gemm<T12>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 12: fp8gemm::run_gemm<T13>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 13: fp8gemm::run_gemm<T14>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 14: fp8gemm::run_gemm<T15>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 15: fp8gemm::run_gemm<T16>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 16: fp8gemm::run_gemm<T17>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 17: fp8gemm::run_gemm<T18>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    case 18: fp8gemm::run_gemm<T19>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval); break;
    default: fp8gemm::run_gemm<T1>(a_c, b_c, s_c, d, (int)max_swizzle, (int)raster_dir, (int)promo_interval);
  }
  return d;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("fp8_gemm_scaled", &fp8_gemm_scaled, "fp8 gemm with per-n scale");
}