"""FP8 e4m3 GEMM for B200 (SM100): hand-written tcgen05 2-SM-MMA CUDA kernel
with Triton fallbacks, dispatched per shape, replayed via CUDA graphs.

y = (x @ w.T) * weight_scale, x fp8e4m3 (M,K), w fp8e4m3 (N,K), out bf16.

Primary path (compute-bound shapes): a self-contained CUDA/PTX kernel doing a
real fp8 x fp8 tensor-core MMA (tcgen05.mma.cta_group::2, fp32 accumulate in
TMEM, M=256 tiles across a 2-CTA cluster, B split N-wise across the pair).
TMA loads use 3D atom-major tensor maps (the TMA engine is instruction-rate
limited at ~5.5M instr/s/SM, so bigger boxes matter), a 6-deep mbarrier
pipeline, double-buffered TMEM accumulators, and a swizzled-SMEM TMA-store
epilogue with the per-output-channel scale fused in.

Skinny shapes (decode-style M<=64) use a Triton split-K kernel with a
partials+reduce finish; these are DRAM-bound.

K not a multiple of 16 (TMA byte-stride rule) is handled by zero-padding to a
multiple of 128: x is re-padded every call through a persistent buffer via a
funnel-shift copy kernel (aligned u32 loads + byte recombination -- ~1.7x
faster than cudaMemcpy2D for odd strides); the padded weight is cached and
invalidated with the tensor _version counter, because the numeric-stress
harness mutates `weight` in place between forwards.

Forward is replayed through CUDA graphs keyed on (x.data_ptr, shape): a graph
re-reads live bytes at recorded addresses, so replay stays correct when tensor
contents change; new address or shape triggers a fresh capture. This removes
~10-20us of Python/launch overhead per call.
"""
import os

import torch
import torch.nn as nn
import triton
import triton.language as tl

E4M3_MAX = 448.0
_NUM_SMS = None


def _alloc(size, alignment, stream):
    return torch.empty(size, device="cuda", dtype=torch.int8)


triton.set_allocator(_alloc)


def _num_sms():
    global _NUM_SMS
    if _NUM_SMS is None:
        _NUM_SMS = torch.cuda.get_device_properties(0).multi_processor_count
    return _NUM_SMS


# ==========================================================================
# Hand-written SM100 tcgen05 kernel (primary path)
# ==========================================================================
_CUDA_SRC = r"""
// SM100 (B200) fp8e4m3 x fp8e4m3 -> bf16 GEMM with per-output-channel scale.
// Hand-rolled tcgen05 kernel: 2-CTA cluster MMA (cta_group::2, M=256), TMA
// double-sided pipeline, TMEM double-buffered accumulators, TMA-store epilogue.
//
// y[m, n] = (sum_k x[m, k] * w[n, k]) * scale[n]
//   x: (M, K) fp8e4m3 row-major, K % 16 == 0 (byte stride rule)
//   w: (N, K) fp8e4m3 row-major
//   scale: (N,) fp32; y: (M, N) bf16
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cstdint>

// ------------------------------ device PTX helpers ------------------------
#define DEVICE __device__ __forceinline__

DEVICE uint32_t smem_u32(const void* p) {
  return static_cast<uint32_t>(__cvta_generic_to_shared(p));
}

DEVICE bool elect_one_sync() {
  uint32_t pred = 0;
  asm volatile(
      "{\n .reg .pred %%px;\n elect.sync _|%%px, 0xffffffff;\n selp.b32 %0, 1, 0, %%px;\n}\n"
      : "+r"(pred)::"memory");
  return pred != 0;
}

DEVICE void mbar_init(uint64_t* bar, uint32_t count) {
  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(smem_u32(bar)), "r"(count));
}

DEVICE void fence_barrier_init() {
  asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
}

DEVICE void mbar_wait(uint64_t* bar, uint32_t phase) {
  asm volatile(
      "{\n"
      ".reg .pred P;\n"
      "WAIT:\n"
      "mbarrier.test_wait.parity.shared::cta.b64 P, [%0], %1;\n"
      "@P bra DONE;\n"
      "bra WAIT;\n"
      "DONE:\n"
      "}\n" ::"r"(smem_u32(bar)),
      "r"(phase));
}

DEVICE void mbar_wait_park(uint64_t* bar, uint32_t phase) {
  asm volatile(
      "{\n"
      ".reg .pred P;\n"
      "WAIT:\n"
      "mbarrier.try_wait.parity.shared::cta.b64 P, [%0], %1, 10000000;\n"
      "@P bra DONE;\n"
      "bra WAIT;\n"
      "DONE:\n"
      "}\n" ::"r"(smem_u32(bar)),
      "r"(phase));
}

DEVICE void mbar_arrive_expect_tx(uint64_t* bar, uint32_t tx) {
  asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smem_u32(bar)),
               "r"(tx)
               : "memory");
}

DEVICE void mbar_arrive_local(uint64_t* bar) {
  asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(smem_u32(bar)) : "memory");
}

// Arrive at the same barrier in the cluster-rank-0 CTA (remote for rank 1).
DEVICE void mbar_arrive_cluster0(uint64_t* bar) {
  uint32_t addr = smem_u32(bar);
  uint32_t remote;
  asm volatile("mapa.shared::cluster.u32 %0, %1, 0;" : "=r"(remote) : "r"(addr));
  asm volatile("mbarrier.arrive.release.cluster.shared::cluster.b64 _, [%0];" ::"r"(remote)
               : "memory");
}

DEVICE void cluster_sync() {
  asm volatile("barrier.cluster.arrive.aligned;" ::: "memory");
  asm volatile("barrier.cluster.wait.aligned;" ::: "memory");
}

DEVICE uint32_t cluster_ctarank() {
  uint32_t r;
  asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r));
  return r;
}

DEVICE void prefetch_tmap(const void* p) {
  asm volatile("prefetch.tensormap [%0];" ::"l"(p) : "memory");
}

// TMA load (cta_group::2): both CTAs issue; tx bytes credited to rank-0 barrier.
DEVICE void tma_load_2sm(const void* tmap, uint64_t* bar, void* smem, int32_t c0, int32_t c1) {
  uint32_t mbar = smem_u32(bar) & 0xFEFFFFFF;  // clear peer bit -> leader CTA
  constexpr uint64_t kEvictNormal = 0x1000000000000000ull;
  asm volatile(
      "cp.async.bulk.tensor.2d.cta_group::2.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
      " [%0], [%1, {%3, %4}], [%2], %5;" ::"r"(smem_u32(smem)),
      "l"(reinterpret_cast<uint64_t>(tmap)), "r"(mbar), "r"(c0), "r"(c1), "l"(kEvictNormal)
      : "memory");
}

// 3D TMA load, cta_group::2: tx credited to leader CTA's barrier (peer bit cleared)
DEVICE void tma_load3d(const void* tmap, uint64_t* bar, void* smem, int32_t c1, int32_t c2) {
  uint32_t mbar = smem_u32(bar) & 0xFEFFFFFF;
  constexpr uint64_t kEvictNormal = 0x1000000000000000ull;
  asm volatile(
      "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
      " [%0], [%1, {0, %3, %4}], [%2], %5;" ::"r"(smem_u32(smem)),
      "l"(reinterpret_cast<uint64_t>(tmap)), "r"(mbar), "r"(c1), "r"(c2), "l"(kEvictNormal)
      : "memory");
}

// Plain 1-CTA TMA load (for the 1SM kernel variant)
DEVICE void tma_load_1sm(const void* tmap, uint64_t* bar, void* smem, int32_t c0, int32_t c1) {
  constexpr uint64_t kEvictNormal = 0x1000000000000000ull;
  asm volatile(
      "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
      " [%0], [%1, {%3, %4}], [%2], %5;" ::"r"(smem_u32(smem)),
      "l"(reinterpret_cast<uint64_t>(tmap)), "r"(smem_u32(bar)), "r"(c0), "r"(c1),
      "l"(kEvictNormal)
      : "memory");
}

DEVICE void tma_store_2d(const void* tmap, const void* smem, int32_t c0, int32_t c1) {
  asm volatile(
      "cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group [%0, {%2, %3}], [%1];" ::"l"(
          reinterpret_cast<uint64_t>(tmap)),
      "r"(smem_u32(smem)), "r"(c0), "r"(c1)
      : "memory");
}

DEVICE void tma_store_commit() { asm volatile("cp.async.bulk.commit_group;" ::: "memory"); }

template <int N>
DEVICE void tma_store_wait() {
  asm volatile("cp.async.bulk.wait_group.read %0;" ::"n"(N) : "memory");
}

DEVICE void fence_async_proxy() {
  asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
}

// tcgen05
DEVICE void tmem_alloc_2sm(uint32_t* dst_smem, uint32_t ncols) {
  asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(
                   smem_u32(dst_smem)),
               "r"(ncols));
}
DEVICE void tmem_relinquish_2sm() {
  asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;" ::);
}
DEVICE void tmem_dealloc_2sm(uint32_t addr, uint32_t ncols) {
  asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" ::"r"(addr), "r"(ncols));
}
DEVICE void tmem_alloc_1sm(uint32_t* dst_smem, uint32_t ncols) {
  asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(
                   smem_u32(dst_smem)),
               "r"(ncols));
}
DEVICE void tmem_relinquish_1sm() {
  asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;" ::);
}
DEVICE void tmem_dealloc_1sm(uint32_t addr, uint32_t ncols) {
  asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(addr), "r"(ncols));
}

DEVICE void tcgen05_fence_after_thread_sync() {
  asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
}
DEVICE void tcgen05_fence_before_thread_sync() {
  asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory");
}
DEVICE void tcgen05_wait_ld() { asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); }

// MMA: kind::f8f6f4, SS. scale_c: 0 -> overwrite acc, 1 -> accumulate.
DEVICE void umma_f8_2sm(uint32_t tmem_c, uint64_t desc_a, uint64_t desc_b, uint32_t idesc,
                        uint32_t scale_c) {
  asm volatile(
      "{\n.reg .pred p;\n setp.ne.b32 p, %4, 0;\n"
      "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, {%5, %6, %7, %8, %9, %10, %11, %12}, p;\n}\n"
      :
      : "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(idesc), "r"(scale_c), "r"(0), "r"(0), "r"(0),
        "r"(0), "r"(0), "r"(0), "r"(0), "r"(0));
}
DEVICE void umma_f8_1sm(uint32_t tmem_c, uint64_t desc_a, uint64_t desc_b, uint32_t idesc,
                        uint32_t scale_c) {
  asm volatile(
      "{\n.reg .pred p;\n setp.ne.b32 p, %4, 0;\n"
      "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, {%5, %6, %7, %8}, p;\n}\n"
      :
      : "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(idesc), "r"(scale_c), "r"(0), "r"(0), "r"(0),
        "r"(0));
}

// tcgen05.commit -> mbarrier arrive, multicast to both cluster CTAs
DEVICE void umma_commit_multicast2(uint64_t* bar) {
  asm volatile(
      "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;" ::
          "r"(smem_u32(bar)),
      "h"((uint16_t)0b11)
      : "memory");
}
DEVICE void umma_commit_1sm(uint64_t* bar) {
  asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" ::"r"(
                   smem_u32(bar))
               : "memory");
}

// TMEM load: 32 lanes x 8 consecutive 32-bit columns per lane
DEVICE void tmem_ld_32x32b_x8(uint32_t addr, uint32_t v[8]) {
  asm volatile(
      "tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];"
      : "=r"(v[0]), "=r"(v[1]), "=r"(v[2]), "=r"(v[3]), "=r"(v[4]), "=r"(v[5]), "=r"(v[6]),
        "=r"(v[7])
      : "r"(addr));
}

DEVICE uint32_t pack_bf16(float lo, float hi) {
  __nv_bfloat162 h = __float22bfloat162_rn({lo, hi});
  return *reinterpret_cast<uint32_t*>(&h);
}

DEVICE void st_shared_v4(void* p, uint32_t a, uint32_t b, uint32_t c, uint32_t d) {
  asm volatile("st.shared.v4.b32 [%0], {%1, %2, %3, %4};" ::"r"(smem_u32(p)), "r"(a), "r"(b),
               "r"(c), "r"(d));
}

DEVICE void named_barrier_sync(uint32_t id, uint32_t nthreads) {
  asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(nthreads) : "memory");
}

// SMEM descriptor for K-major SW128 tiles: SBO=1024B, LBO=0, version=1, layout=SWIZZLE_128B(2)
DEVICE uint64_t make_desc_k_major_sw128(const void* smem) {
  const uint32_t lo = smem_u32(smem) >> 4;
  const uint32_t hi = 64u /*SBO 1024>>4*/ | (1u << 14) /*version*/ | (2u << 29) /*SW128*/;
  return (uint64_t)lo | ((uint64_t)hi << 32);
}

// ------------------------------ kernel ------------------------------------
// Per-CTA: A tile 128 x BK, B tile (BN/2) x BK (2SM: pair covers 256 x BN).
template <uint32_t BN, uint32_t BK, uint32_t STAGES, uint32_t GROUP_M, uint32_t MODE = 0>
__global__ void __launch_bounds__(256, 1) fp8_gemm_2sm_kernel(
    const __grid_constant__ CUtensorMap tmap_a, const __grid_constant__ CUtensorMap tmap_b,
    const __grid_constant__ CUtensorMap tmap_d, const float* __restrict__ scale, uint32_t M,
    uint32_t N, uint32_t K, uint32_t num_k_blocks) {
#if __CUDA_ARCH__ >= 1000
  constexpr uint32_t BM = 256;               // pair-wide M tile
  constexpr uint32_t LOAD_M = 128;           // per-CTA A rows
  constexpr uint32_t LOAD_N = BN / 2;        // per-CTA B rows (N-split across pair)
  constexpr uint32_t ACC_STAGES = 2;
  constexpr uint32_t ATOMS = BK / 128;       // 128B swizzle atoms per stage
  constexpr uint32_t SMEM_A_STAGE = LOAD_M * BK;      // bytes (fp8)
  constexpr uint32_t SMEM_B_STAGE = LOAD_N * BK;      // bytes
  constexpr uint32_t STORE_N = 64;                    // bf16 cols per CD stage (128B swizzle)
  constexpr uint32_t CD_STAGES = 2;
  constexpr uint32_t SMEM_CD_STAGE = 128 * STORE_N * 2;  // 16KB

  extern __shared__ __align__(1024) uint8_t smem[];
  uint8_t* smem_cd = smem;                                   // CD_STAGES * 16KB
  uint8_t* smem_a = smem + CD_STAGES * SMEM_CD_STAGE;        // STAGES * SMEM_A_STAGE
  uint8_t* smem_b = smem_a + STAGES * SMEM_A_STAGE;          // STAGES * SMEM_B_STAGE
  uint64_t* bars = reinterpret_cast<uint64_t*>(smem_b + STAGES * SMEM_B_STAGE);
  uint64_t* full_bar = bars;                    // [STAGES]
  uint64_t* empty_bar = bars + STAGES;          // [STAGES]
  uint64_t* tmem_full = bars + 2 * STAGES;      // [ACC_STAGES]
  uint64_t* tmem_empty = bars + 2 * STAGES + ACC_STAGES;  // [ACC_STAGES]
  uint32_t* tmem_ptr = reinterpret_cast<uint32_t*>(bars + 2 * STAGES + 2 * ACC_STAGES);

  const uint32_t warp_idx = threadIdx.x / 32;
  const uint32_t lane_idx = threadIdx.x % 32;
  const uint32_t rank = cluster_ctarank();  // 0 = leader
  const bool is_leader = rank == 0;

  if (warp_idx == 0 && elect_one_sync()) {
    prefetch_tmap(&tmap_a);
    prefetch_tmap(&tmap_b);
    prefetch_tmap(&tmap_d);
  }

  cluster_sync();  // rendezvous before 2-CTA TMEM alloc

  if (warp_idx == 1 && elect_one_sync()) {
    for (uint32_t i = 0; i < STAGES; ++i) {
      mbar_init(&full_bar[i], 1);   // leader's arrive_expect_tx only; peer tx auto-credits
      mbar_init(&empty_bar[i], 1);  // one multicast commit per CTA
    }
    for (uint32_t i = 0; i < ACC_STAGES; ++i) {
      mbar_init(&tmem_full[i], 1);           // one multicast commit
      mbar_init(&tmem_empty[i], 2);          // one elected arrive per CTA
    }
    fence_barrier_init();
  } else if (warp_idx == 2) {
    tmem_alloc_2sm(tmem_ptr, 512);
    tmem_relinquish_2sm();
  }
  cluster_sync();

  // ---- persistent tile schedule (grouped along M) ----
  const uint32_t num_pid_m = (M + BM - 1) / BM;
  const uint32_t num_pid_n = (N + BN - 1) / BN;
  const uint32_t num_tiles = num_pid_m * num_pid_n;
  const uint32_t cluster_id = blockIdx.x / 2;
  const uint32_t num_clusters = gridDim.x / 2;
  const uint32_t pids_per_group = GROUP_M * num_pid_n;

  auto tile_coords = [&](uint32_t tile, uint32_t& pm, uint32_t& pn) {
    uint32_t group = tile / pids_per_group;
    uint32_t first_m = group * GROUP_M;
    uint32_t gsz = min(num_pid_m - first_m, GROUP_M);
    pm = first_m + (tile % pids_per_group) % gsz;
    pn = (tile % pids_per_group) / gsz;
  };

  if (warp_idx == 0) {
    // ---- TMA producer ----
    if (elect_one_sync()) {
      uint32_t stage = 0, phase = 0;
      for (uint32_t tile = cluster_id; tile < num_tiles; tile += num_clusters) {
        uint32_t pm, pn;
        tile_coords(tile, pm, pn);
        const int32_t m_idx = pm * BM + rank * LOAD_M;
        const int32_t n_idx = pn * BN + rank * LOAD_N;
        for (uint32_t kb = 0; kb < num_k_blocks; ++kb) {
          mbar_wait(&empty_bar[stage], phase ^ 1);
          if (is_leader)
            mbar_arrive_expect_tx(&full_bar[stage], 2 * (SMEM_A_STAGE + SMEM_B_STAGE));
          tma_load3d(&tmap_a, &full_bar[stage], smem_a + stage * SMEM_A_STAGE, m_idx, kb * ATOMS);
          tma_load3d(&tmap_b, &full_bar[stage], smem_b + stage * SMEM_B_STAGE, n_idx, kb * ATOMS);
          stage = stage + 1 == STAGES ? 0 : stage + 1;
          phase ^= (stage == 0);
        }
      }
    }
  } else if (warp_idx == 1 && is_leader) {
    // ---- MMA issuer (leader CTA only) ----
    // idesc: c_format=F32, e4m3 x e4m3, M=256, N=BN, K-major both
    const uint32_t idesc = (1u << 4) | ((BN >> 3) << 17) | ((256u >> 4) << 24);
    uint32_t stage = 0, phase = 0;
    uint32_t iter = 0;
    for (uint32_t tile = cluster_id; tile < num_tiles; tile += num_clusters, ++iter) {
      const uint32_t acc_stage = iter % ACC_STAGES;
      const uint32_t acc_phase = (iter / ACC_STAGES) & 1;
      mbar_wait(&tmem_empty[acc_stage], acc_phase ^ 1);
      tcgen05_fence_after_thread_sync();

      for (uint32_t kb = 0; kb < num_k_blocks; ++kb) {
        mbar_wait(&full_bar[stage], phase);
        tcgen05_fence_after_thread_sync();
        if (elect_one_sync()) {
          uint64_t desc_a = make_desc_k_major_sw128(smem_a + stage * SMEM_A_STAGE);
          uint64_t desc_b = make_desc_k_major_sw128(smem_b + stage * SMEM_B_STAGE);
          const uint32_t tmem_c = acc_stage * BN;
          uint32_t accum = (kb > 0) ? 1u : 0u;
#pragma unroll
          for (uint32_t u = 0; u < BK / 32; ++u) {
            // per-atom SMEM step = LOAD_M rows * 128B; within-atom step = 32B
            const uint32_t lo_off = (u / 4) * (LOAD_M * 128 / 16) + (u % 4) * 2;
            umma_f8_2sm(tmem_c, desc_a + lo_off, desc_b + (u / 4) * (LOAD_N * 128 / 16) + (u % 4) * 2,
                        idesc, accum | (u > 0 ? 1u : 0u));
          }
          umma_commit_multicast2(&empty_bar[stage]);
          if (kb == num_k_blocks - 1) umma_commit_multicast2(&tmem_full[acc_stage]);
        }
        __syncwarp();
        stage = stage + 1 == STAGES ? 0 : stage + 1;
        phase ^= (stage == 0);
      }
    }
    // allow safe teardown: wait for last accumulator to be drained
    if (iter > 0) {
      const uint32_t last = iter - 1;
      mbar_wait(&tmem_empty[last % ACC_STAGES], (last / ACC_STAGES) & 1);
    }
  } else if (warp_idx >= 4) {
    // ---- epilogue: 4 warps, 128 threads ----
    const uint32_t ep_warp = warp_idx - 4;
    const uint32_t row = ep_warp * 32 + lane_idx;  // row in this CTA's 128-row half
    uint32_t cd_stage = 0;
    uint32_t iter = 0;
    for (uint32_t tile = cluster_id; tile < num_tiles; tile += num_clusters, ++iter) {
      uint32_t pm, pn;
      tile_coords(tile, pm, pn);
      const uint32_t acc_stage = iter % ACC_STAGES;
      const uint32_t acc_phase = (iter / ACC_STAGES) & 1;
      mbar_wait(&tmem_full[acc_stage], acc_phase);
      tcgen05_fence_after_thread_sync();

      const int32_t m_idx = pm * BM + rank * LOAD_M;  // this CTA's 128 rows
#pragma unroll
      for (uint32_t s = 0; s < BN / STORE_N; ++s) {
        // wait for this CD buffer's previous TMA store to drain
        if (ep_warp == 0) tma_store_wait<CD_STAGES - 1>();
        named_barrier_sync(1, 128);

        uint8_t* cd = smem_cd + cd_stage * SMEM_CD_STAGE;
        const bool n_aligned = (N % BN) == 0;
#pragma unroll
        for (uint32_t i = 0; i < STORE_N / 8; i += 4) {
          uint32_t v[32];
          tmem_ld_32x32b_x8((acc_stage * BN) + s * STORE_N + i * 8, v);
          tmem_ld_32x32b_x8((acc_stage * BN) + s * STORE_N + i * 8 + 8, v + 8);
          tmem_ld_32x32b_x8((acc_stage * BN) + s * STORE_N + i * 8 + 16, v + 16);
          tmem_ld_32x32b_x8((acc_stage * BN) + s * STORE_N + i * 8 + 24, v + 24);
          tcgen05_wait_ld();
          const uint32_t n0 = pn * BN + s * STORE_N + i * 8;
          float f[32];
          if (n_aligned) {
#pragma unroll
            for (uint32_t j = 0; j < 32; ++j)
              f[j] = reinterpret_cast<float&>(v[j]) * __ldg(scale + n0 + j);
          } else {
#pragma unroll
            for (uint32_t j = 0; j < 32; ++j)
              f[j] = reinterpret_cast<float&>(v[j]) * __ldg(scale + min(n0 + j, N - 1));
          }
#pragma unroll
          for (uint32_t q = 0; q < 4; ++q) {
            uint8_t* p = cd + row * 128 + (((i + q) ^ (row % 8)) * 16);
            st_shared_v4(p, pack_bf16(f[q * 8 + 0], f[q * 8 + 1]), pack_bf16(f[q * 8 + 2], f[q * 8 + 3]),
                         pack_bf16(f[q * 8 + 4], f[q * 8 + 5]), pack_bf16(f[q * 8 + 6], f[q * 8 + 7]));
          }
        }
        if (s == BN / STORE_N - 1) tcgen05_fence_before_thread_sync();
        fence_async_proxy();
        named_barrier_sync(1, 128);
        if (ep_warp == 0 && elect_one_sync()) {
          if (s == BN / STORE_N - 1)
            mbar_arrive_cluster0(&tmem_empty[acc_stage]);  // rank0-local, rank1-remote
          tma_store_2d(&tmap_d, cd, pn * BN + s * STORE_N, m_idx);
          tma_store_commit();
        }
        cd_stage ^= 1;
      }
    }
    if (ep_warp == 0) tma_store_wait<0>();
  }

  cluster_sync();
  if (warp_idx == 2) tmem_dealloc_2sm(0, 512);
#endif
}

// ------------------------------ host side ---------------------------------
// {128B, rows, K/128 atoms} view with atom-major boxes {128, box_rows, 2}
static CUtensorMap make_tmap_3d_atoms(const void* ptr, uint64_t rows, uint64_t K,
                                      uint32_t box_rows, uint32_t box_atoms) {
  CUtensorMap tmap;
  uint64_t dims[3] = {128, rows, K / 128};
  uint64_t strides[2] = {K, 128};
  uint32_t box[3] = {128, box_rows, box_atoms};
  uint32_t es[3] = {1, 1, 1};
  CUresult res = cuTensorMapEncodeTiled(
      &tmap, CU_TENSOR_MAP_DATA_TYPE_UINT8, 3, const_cast<void*>(ptr), dims, strides, box, es,
      CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
      CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
  TORCH_CHECK(res == CUDA_SUCCESS, "cuTensorMapEncodeTiled 3d failed: ", (int)res);
  return tmap;
}

static CUtensorMap make_tmap_2d(const void* ptr, uint64_t inner, uint64_t outer,
                                uint64_t stride_bytes, uint32_t box_inner, uint32_t box_outer,
                                CUtensorMapDataType dtype, CUtensorMapSwizzle swizzle) {
  CUtensorMap tmap;
  uint64_t dims[2] = {inner, outer};
  uint64_t strides[1] = {stride_bytes};
  uint32_t box[2] = {box_inner, box_outer};
  uint32_t elem_strides[2] = {1, 1};
  CUresult res = cuTensorMapEncodeTiled(
      &tmap, dtype, 2, const_cast<void*>(ptr), dims, strides, box, elem_strides,
      CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
      CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
  TORCH_CHECK(res == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed: ", (int)res);
  return tmap;
}

template <uint32_t BN, uint32_t BK, uint32_t STAGES, uint32_t GROUP_M, uint32_t MODE = 0>
void launch_2sm(const torch::Tensor& x, const torch::Tensor& w, const torch::Tensor& scale,
                torch::Tensor& y, int sms) {
  const uint32_t M = x.size(0), K = x.size(1), N = w.size(0);
  constexpr uint32_t SMEM_A = 128 * BK, SMEM_B = (BN / 2) * BK;
  constexpr uint32_t CD_BYTES = 2 * 128 * 64 * 2;
  constexpr uint32_t BAR_BYTES = (2 * STAGES + 4) * 8 + 16;
  const uint32_t smem_size = CD_BYTES + STAGES * (SMEM_A + SMEM_B) + BAR_BYTES;

  TORCH_CHECK(K % 128 == 0, "hand kernel requires K % 128 == 0");
  auto tmap_a = make_tmap_3d_atoms(x.data_ptr(), M, K, 128, BK / 128);
  auto tmap_b = make_tmap_3d_atoms(w.data_ptr(), N, K, BN / 2, BK / 128);
  auto tmap_d = make_tmap_2d(y.data_ptr(), N, M, (uint64_t)N * 2, 64, 128,
                             CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, CU_TENSOR_MAP_SWIZZLE_128B);

  auto kernel = fp8_gemm_2sm_kernel<BN, BK, STAGES, GROUP_M, MODE>;
  static bool attr_set = false;
  if (!attr_set) {
    cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 232448);
    attr_set = true;
  }

  const uint32_t num_k_blocks = (K + BK - 1) / BK;
  cudaLaunchConfig_t cfg = {};
  cudaLaunchAttribute attrs[1];
  attrs[0].id = cudaLaunchAttributeClusterDimension;
  attrs[0].val.clusterDim = {2, 1, 1};
  cfg.gridDim = dim3(sms, 1, 1);
  cfg.blockDim = dim3(256, 1, 1);
  cfg.dynamicSmemBytes = smem_size;
  cfg.stream = at::cuda::getCurrentCUDAStream();
  cfg.attrs = attrs;
  cfg.numAttrs = 1;
  cudaError_t err = cudaLaunchKernelEx(&cfg, kernel, tmap_a, tmap_b, tmap_d,
                                       scale.data_ptr<float>(), M, N, K, num_k_blocks);
  TORCH_CHECK(err == cudaSuccess, "launch failed: ", cudaGetErrorString(err));
}

void gemm_fp8_2sm(torch::Tensor x, torch::Tensor w, torch::Tensor scale, torch::Tensor y,
                  int64_t sms) {
  TORCH_CHECK(x.size(1) == w.size(1), "K mismatch");
  TORCH_CHECK(x.size(1) % 128 == 0, "K must be a multiple of 128");
  TORCH_CHECK(x.is_contiguous() && w.is_contiguous() && y.is_contiguous());
  launch_2sm<256, 128, 6, 8>(x, w, scale, y, sms);
}


DEVICE void tma_load3d_1sm_hint(const void* tmap, uint64_t* bar, void* smem, int32_t c1, int32_t c2,
                                uint64_t hint) {
  asm volatile(
      "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
      " [%0], [%1, {0, %3, %4}], [%2], %5;" ::"r"(smem_u32(smem)),
      "l"((uint64_t)tmap), "r"(smem_u32(bar)), "r"(c1), "r"(c2), "l"(hint) : "memory");
}

// ------------------- skinny-M kernel -------------------

template <uint32_t BN, uint32_t BK, uint32_t STAGES, bool EVICT_FIRST>
__global__ void __launch_bounds__(256, 1) fp8_gemm_skinny_kernel(
    const __grid_constant__ CUtensorMap tmap_a, const __grid_constant__ CUtensorMap tmap_b,
    float* __restrict__ partials, uint32_t M, uint32_t N, uint32_t K, uint32_t k_per_split) {
#if __CUDA_ARCH__ >= 1000
  constexpr uint32_t ROWS = 128;                       // MMA M (rows beyond M are TMA-zero)
  constexpr uint32_t ATOMS = BK / 128;
  constexpr uint32_t SMEM_A_STAGE = ROWS * BK;         // 32KB at BK=256
  constexpr uint32_t SMEM_B_STAGE = BN * BK;

  extern __shared__ __align__(1024) uint8_t smem[];
  uint8_t* smem_a = smem;
  uint8_t* smem_b = smem + STAGES * SMEM_A_STAGE;
  uint64_t* bars = reinterpret_cast<uint64_t*>(smem_b + STAGES * SMEM_B_STAGE);
  uint64_t* full_bar = bars;
  uint64_t* empty_bar = bars + STAGES;
  uint64_t* tmem_full = bars + 2 * STAGES;
  uint32_t* tmem_ptr = reinterpret_cast<uint32_t*>(bars + 2 * STAGES + 1);

  const uint32_t warp_idx = threadIdx.x / 32;
  const uint32_t lane_idx = threadIdx.x % 32;

  if (warp_idx == 0 && elect_one_sync()) {
    asm volatile("prefetch.tensormap [%0];" ::"l"((const void*)&tmap_a) : "memory");
    asm volatile("prefetch.tensormap [%0];" ::"l"((const void*)&tmap_b) : "memory");
  }
  if (warp_idx == 1 && elect_one_sync()) {
    for (uint32_t i = 0; i < STAGES; ++i) {
      mbar_init(&full_bar[i], 1);
      mbar_init(&empty_bar[i], 1);
    }
    mbar_init(tmem_full, 1);
    asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
  } else if (warp_idx == 2) {
    tmem_alloc_1sm(tmem_ptr, BN >= 128 ? BN : 128);
    tmem_relinquish_1sm();
  }
  __syncthreads();

  const uint32_t pid_n = blockIdx.x;
  const uint32_t pid_s = blockIdx.y;
  const int32_t n_idx = pid_n * BN;
  const uint32_t k_lo_atom = pid_s * (k_per_split / 128);
  const uint32_t num_kb = k_per_split / BK;

  constexpr uint64_t kHintNormal = 0x1000000000000000ull;
  constexpr uint64_t kHintEvictFirst = 0x12F0000000000000ull;
  constexpr uint64_t kHint = EVICT_FIRST ? kHintEvictFirst : kHintNormal;

  if (warp_idx == 0) {
    if (elect_one_sync()) {
      uint32_t stage = 0, phase = 0;
      for (uint32_t kb = 0; kb < num_kb; ++kb) {
        mbar_wait(&empty_bar[stage], phase ^ 1);
        mbar_arrive_expect_tx(&full_bar[stage], SMEM_A_STAGE + SMEM_B_STAGE);
        tma_load3d_1sm_hint(&tmap_a, &full_bar[stage], smem_a + stage * SMEM_A_STAGE, 0,
                       k_lo_atom + kb * ATOMS, kHintNormal);
        tma_load3d_1sm_hint(&tmap_b, &full_bar[stage], smem_b + stage * SMEM_B_STAGE, n_idx,
                       k_lo_atom + kb * ATOMS, kHint);
        stage = stage + 1 == STAGES ? 0 : stage + 1;
        phase ^= (stage == 0);
      }
    }
  } else if (warp_idx == 1) {
    const uint32_t idesc = (1u << 4) | ((BN >> 3) << 17) | ((ROWS >> 4) << 24);
    uint32_t stage = 0, phase = 0;
    for (uint32_t kb = 0; kb < num_kb; ++kb) {
      mbar_wait(&full_bar[stage], phase);
      tcgen05_fence_after_thread_sync();
      if (elect_one_sync()) {
        uint64_t desc_a = make_desc_k_major_sw128(smem_a + stage * SMEM_A_STAGE);
        uint64_t desc_b = make_desc_k_major_sw128(smem_b + stage * SMEM_B_STAGE);
#pragma unroll
        for (uint32_t u = 0; u < BK / 32; ++u) {
          const uint32_t a_off = (u / 4) * (ROWS * 128 / 16) + (u % 4) * 2;
          const uint32_t b_off = (u / 4) * (BN * 128 / 16) + (u % 4) * 2;
          umma_f8_1sm(0, desc_a + a_off, desc_b + b_off, idesc, (kb > 0 || u > 0) ? 1u : 0u);
        }
        umma_commit_1sm(&empty_bar[stage]);
        if (kb == num_kb - 1) umma_commit_1sm(tmem_full);
      }
      __syncwarp();
      stage = stage + 1 == STAGES ? 0 : stage + 1;
      phase ^= (stage == 0);
    }
  } else if (warp_idx >= 4) {
    const uint32_t ep_warp = warp_idx - 4;
    const uint32_t row = ep_warp * 32 + lane_idx;
    mbar_wait(tmem_full, 0);
    tcgen05_fence_after_thread_sync();
    // tcgen05.ld is warp-scoped: the whole warp must execute it uniformly, so
    // gate per warp and predicate only the stores.
    if (ep_warp * 32 < M) {
      const bool row_ok = row < M;
      float* out = partials + (uint64_t)pid_s * M * N + (uint64_t)row * N + n_idx;
#pragma unroll
      for (uint32_t i = 0; i < BN / 8; ++i) {
        uint32_t v[8];
        tmem_ld_32x32b_x8(i * 8, v);
        tcgen05_wait_ld();
        const uint32_t n0 = n_idx + i * 8;
        if (row_ok) {
          if (n0 + 8 <= N) {
            float4 lo = {reinterpret_cast<float&>(v[0]), reinterpret_cast<float&>(v[1]),
                         reinterpret_cast<float&>(v[2]), reinterpret_cast<float&>(v[3])};
            float4 hi = {reinterpret_cast<float&>(v[4]), reinterpret_cast<float&>(v[5]),
                         reinterpret_cast<float&>(v[6]), reinterpret_cast<float&>(v[7])};
            *reinterpret_cast<float4*>(out + i * 8) = lo;
            *reinterpret_cast<float4*>(out + i * 8 + 4) = hi;
          } else {
#pragma unroll
            for (uint32_t j = 0; j < 8; ++j)
              if (n0 + j < N) out[i * 8 + j] = reinterpret_cast<float&>(v[j]);
          }
        }
      }
    }
  }

  __syncthreads();
  if (warp_idx == 2) tmem_dealloc_1sm(0, BN >= 128 ? BN : 128);
#endif
}

void gemm_fp8_skinny(torch::Tensor x, torch::Tensor w, torch::Tensor partials, int64_t split,
                     int64_t evict_first) {
  const uint32_t M = x.size(0), K = x.size(1), N = w.size(0);
  TORCH_CHECK(K % 128 == 0 && M <= 128);
  constexpr uint32_t BN = 128, BK = 256, STAGES = 3;
  const uint32_t k_per_split = (K / split + BK - 1) / BK * BK;

  auto tmap_a = make_tmap_3d_atoms(x.data_ptr(), M, K, 128, BK / 128);
  auto tmap_b = make_tmap_3d_atoms(w.data_ptr(), N, K, BN, BK / 128);

  const uint32_t smem = STAGES * (128 * BK + BN * BK) + (2 * STAGES + 1) * 8 + 16;
  dim3 grid((N + BN - 1) / BN, split);
  auto stream = at::cuda::getCurrentCUDAStream();
  if (evict_first) {
    auto k = fp8_gemm_skinny_kernel<BN, BK, STAGES, true>;
    static bool set1 = false;
    if (!set1) { cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, 232448); set1 = true; }
    k<<<grid, 256, smem, stream>>>(tmap_a, tmap_b, partials.data_ptr<float>(), M, N, K, k_per_split);
  } else {
    auto k = fp8_gemm_skinny_kernel<BN, BK, STAGES, false>;
    static bool set2 = false;
    if (!set2) { cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, 232448); set2 = true; }
    k<<<grid, 256, smem, stream>>>(tmap_a, tmap_b, partials.data_ptr<float>(), M, N, K, k_per_split);
  }
  cudaError_t err = cudaGetLastError();
  TORCH_CHECK(err == cudaSuccess, "skinny launch failed: ", cudaGetErrorString(err));
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("gemm_fp8_2sm", &gemm_fp8_2sm, "fp8 GEMM 2SM tcgen05");
  m.def("gemm_fp8_skinny", &gemm_fp8_skinny, "fp8 skinny GEMM 1SM tcgen05");
}

"""

_HAND = None
if os.environ.get("KBH_DISABLE_HAND", "0") != "1":
    try:
        if torch.cuda.get_device_capability(0)[0] == 10:
            from torch.utils.cpp_extension import load_inline

            # tcgen05 needs the arch-specific target; the container pre-sets a
            # multi-arch TORCH_CUDA_ARCH_LIST that would break the build.
            _saved_arch = os.environ.get("TORCH_CUDA_ARCH_LIST")
            os.environ["TORCH_CUDA_ARCH_LIST"] = "10.0a"
            try:
                _HAND = load_inline(
                    name="fp8_gemm_sm100_hand",
                    cpp_sources=[],
                    cuda_sources=[_CUDA_SRC],
                    extra_cuda_cflags=["-O3", "--use_fast_math", "-std=c++17"],
                    extra_ldflags=["-lcuda"],
                    verbose=False,
                )
            finally:
                if _saved_arch is None:
                    os.environ.pop("TORCH_CUDA_ARCH_LIST", None)
                else:
                    os.environ["TORCH_CUDA_ARCH_LIST"] = _saved_arch
    except Exception:
        _HAND = None


def _hand_ok(M, N, K_pad):
    return (
        _HAND is not None
        and M >= 128
        and N % 16 == 0
        and K_pad % 128 == 0
        and torch.cuda.get_device_capability(0)[0] == 10
    )


# ==========================================================================
# Triton kernels: persistent fallback GEMM, skinny split-K, pad copy
# ==========================================================================
@triton.jit
def _gemm_persistent(
    x_ptr, w_ptr, y_ptr, scale_ptr,
    M, N, K_pad,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr, NUM_SMS: tl.constexpr, SUBTILE: tl.constexpr,
):
    start_pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    k_tiles = tl.cdiv(K_pad, BLOCK_K)
    num_tiles = num_pid_m * num_pid_n

    x_desc = tl.make_tensor_descriptor(x_ptr, [M, K_pad], [K_pad, 1], [BLOCK_M, BLOCK_K])
    w_desc = tl.make_tensor_descriptor(w_ptr, [N, K_pad], [K_pad, 1], [BLOCK_N, BLOCK_K])
    y_desc = tl.make_tensor_descriptor(
        y_ptr, [M, N], [N, 1],
        [BLOCK_M, BLOCK_N // 2 if SUBTILE else BLOCK_N],
    )

    num_pid_in_group = GROUP_M * num_pid_n

    for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True, warp_specialize=True):
        group_id = tile_id // 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 + (tile_id % num_pid_in_group) % group_size_m
        pid_n = (tile_id % num_pid_in_group) // group_size_m

        off_m = pid_m * BLOCK_M
        off_n = pid_n * BLOCK_N

        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
        for k in tl.range(0, k_tiles):
            x = x_desc.load([off_m, k * BLOCK_K])
            w = w_desc.load([off_n, k * BLOCK_K])
            acc = tl.dot(x, w.T, acc)

        offs_n = off_n + tl.arange(0, BLOCK_N)
        scale = tl.load(scale_ptr + offs_n, mask=offs_n < N, other=0.0).to(tl.float32)
        acc = acc * scale[None, :]
        if SUBTILE:
            acc0, acc1 = acc.reshape(BLOCK_M, 2, BLOCK_N // 2).permute(0, 2, 1).split()
            y_desc.store([off_m, off_n], acc0.to(tl.bfloat16))
            y_desc.store([off_m, off_n + BLOCK_N // 2], acc1.to(tl.bfloat16))
        else:
            y_desc.store([off_m, off_n], acc.to(tl.bfloat16))


@triton.jit
def _gemm_skinny(
    x_ptr, w_ptr, p_ptr,
    M, N, K,
    K_PER: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
    pid_n = tl.program_id(0)
    pid_s = tl.program_id(1)
    x_desc = tl.make_tensor_descriptor(x_ptr, [M, K], [K, 1], [BLOCK_M, BLOCK_K])
    w_desc = tl.make_tensor_descriptor(w_ptr, [N, K], [K, 1], [BLOCK_N, BLOCK_K])
    off_n = pid_n * BLOCK_N
    k_lo = pid_s * K_PER
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in tl.range(0, K_PER // BLOCK_K, flatten=True):
        x = x_desc.load([0, k_lo + k * BLOCK_K])
        w = w_desc.load([off_n, k_lo + k * BLOCK_K])
        acc = tl.dot(x, w.T, acc)
    offs_m = tl.arange(0, BLOCK_M)
    offs_n = off_n + tl.arange(0, BLOCK_N)
    ptrs = p_ptr + pid_s.to(tl.int64) * M * N + offs_m[:, None] * N + offs_n[None, :]
    tl.store(ptrs, acc, mask=(offs_m[:, None] < M))


@triton.jit
def _skinny_reduce(p_ptr, scale_ptr, y_ptr, MN, N, S: tl.constexpr, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < MN
    acc = tl.zeros((BLOCK,), dtype=tl.float32)
    for s_i in tl.static_range(S):
        acc += tl.load(p_ptr + s_i * MN + offs, mask=mask, other=0.0)
    sc = tl.load(scale_ptr + offs % N, mask=mask, other=0.0)
    tl.store(y_ptr + offs, (acc * sc).to(tl.bfloat16), mask=mask)


@triton.jit
def _pad_rows_u32(src_ptr, dst_ptr, M, K, K_pad,
                  N_U32: tl.constexpr, BLOCK: tl.constexpr):
    """Copy (M,K) bytes with odd row stride K into (M,K_pad) zero-padded.

    Vectorized despite misalignment: aligned u32 pair loads + funnel shift,
    aligned u32 stores. Dst pad columns beyond the partial u32 stay untouched
    (pre-zeroed buffer)."""
    pid = tl.program_id(0)
    row = pid // tl.cdiv(N_U32, BLOCK)
    chunk = pid % tl.cdiv(N_U32, BLOCK)
    jc = chunk * BLOCK + tl.arange(0, BLOCK)
    in_row = jc < N_U32

    row_start = row * K
    sh = (row_start % 4).to(tl.uint32)
    idx = row_start // 4 + jc
    total_u32 = (M * K + 3) // 4

    lo = tl.load(src_ptr + idx, mask=in_row & (idx < total_u32), other=0,
                 eviction_policy="evict_first").to(tl.uint32, bitcast=True)
    hi = tl.load(src_ptr + idx + 1, mask=in_row & (idx + 1 < total_u32), other=0,
                 eviction_policy="evict_first").to(tl.uint32, bitcast=True)
    val = tl.where(sh == 0, lo, (lo >> (8 * sh)) | (hi << (32 - 8 * sh)))

    rem = K - jc * 4
    val = tl.where(rem >= 4, val,
                   val & ((tl.full((BLOCK,), 1, tl.uint32) << (8 * tl.maximum(rem, 0).to(tl.uint32))) - 1))

    tl.store(dst_ptr + row * (K_pad // 4) + jc, val.to(tl.int32, bitcast=True), mask=in_row)


def _pad_copy(x: torch.Tensor, xpad: torch.Tensor):
    M, K = x.shape
    K_pad = xpad.shape[1]
    n_u32 = (K + 3) // 4
    src32 = x.view(-1)[: M * K // 4 * 4].view(torch.int32)
    grid = (M * triton.cdiv(n_u32, 512),)
    _pad_rows_u32[grid](src32, xpad.view(-1).view(torch.int32), M, K, K_pad,
                        n_u32, 512, num_warps=4)


# ==========================================================================
# config selection
# ==========================================================================
# triton persistent fallback: (BM, BN, BK, GROUP_M, stages, warps, subtile)
_PERSISTENT_CFGS = {
    (4096, 4096): (128, 128, 128, 8, 6, 8, False),
    (4096, 14336): (256, 128, 128, 8, 4, 8, True),
}
_PERSISTENT_DEFAULT = (128, 128, 128, 8, 6, 8, False)
# skinny: (BM, BN, BK, SPLIT, stages, warps)
_SKINNY_CFG = (64, 128, 256, 2, 4, 4)


def _pad_k(K):
    """Padded K when K breaks TMA's 16B-stride rule: round to 128B so every
    TMA box row stays DRAM-sector aligned (16B-only alignment measurably
    loses ~25% mainloop throughput), and so the hand kernel's K%128 rule
    holds."""
    return K if K % 16 == 0 else (K + 127) // 128 * 128


class _Plan:
    __slots__ = ("launch", "graph", "y")

    def __init__(self, launch, y):
        self.launch = launch
        self.graph = None
        self.y = y


class Model(nn.Module):
    """y = ((x @ w.T) * weight_scale).to(bf16) via fp8 tensor-core MMA."""

    def __init__(self, M: int, N: int, K: int):
        super().__init__()
        self.M, self.N, self.K = M, N, K
        w = torch.empty(N, K, dtype=torch.bfloat16)
        nn.init.normal_(w, std=0.02)
        s = (w.float().abs().amax(dim=1, keepdim=True) / E4M3_MAX).clamp(min=1e-12)
        w_fp8 = (w.float() / s).to(torch.float8_e4m3fn)
        self.register_buffer("weight", w_fp8)
        self.register_buffer("weight_scale", s.squeeze(1).to(torch.float32))
        # runtime state (not in state_dict)
        self._wpad = None
        self._wpad_key = None
        self._xpad = None
        self._partials = None
        self._plans = {}
        self._graphs_ok = True

    # -- padding ----------------------------------------------------------
    def _padded_weight(self) -> torch.Tensor:
        w = self.weight
        K = w.shape[1]
        Kp = _pad_k(K)
        key = (w.data_ptr(), w._version)
        if self._wpad_key != key:
            if self._wpad is None or self._wpad.shape != (w.shape[0], Kp):
                self._wpad = torch.zeros(w.shape[0], Kp, device=w.device, dtype=w.dtype)
            _pad_copy(w, self._wpad)
            self._wpad_key = key
        return self._wpad

    # -- launchers ----------------------------------------------------------
    def _make_launch(self, x, w, y, M, N, K_pad):
        scale = self.weight_scale
        K_orig = x.shape[1]
        needs_pad = K_pad != K_orig

        if M <= 64:
            if _HAND is not None and K_pad % 128 == 0 and M <= 128:
                SPLIT = 2
                if self._partials is None or self._partials.shape != (SPLIT, M, N):
                    self._partials = torch.empty(SPLIT, M, N, device=x.device, dtype=torch.float32)
                P = self._partials
                grid2 = (triton.cdiv(M * N, 2048),)
                xp = self._xpad if needs_pad else x

                def launch():
                    if needs_pad:
                        _pad_copy(x, self._xpad)
                    _HAND.gemm_fp8_skinny(xp, w, P, SPLIT, 1)
                    _skinny_reduce[grid2](P, scale, y, M * N, N, SPLIT, 2048, num_warps=4)
                return launch
            BM, BN, BK, SPLIT, stages, warps = _SKINNY_CFG
            k_per = triton.cdiv(triton.cdiv(K_pad, SPLIT), BK) * BK
            n_split = max(1, triton.cdiv(K_pad, k_per))
            if self._partials is None or self._partials.shape != (n_split, M, N):
                self._partials = torch.empty(n_split, M, N, device=x.device, dtype=torch.float32)
            P = self._partials
            grid = (triton.cdiv(N, BN), n_split)
            grid2 = (triton.cdiv(M * N, 2048),)
            xp = self._xpad if needs_pad else x

            def launch():
                if needs_pad:
                    _pad_copy(x, self._xpad)
                _gemm_skinny[grid](xp, w, P, M, N, K_pad, k_per, BM, BN, BK,
                                   num_stages=stages, num_warps=warps)
                _skinny_reduce[grid2](P, scale, y, M * N, N, n_split, 2048, num_warps=4)
        elif _hand_ok(M, N, K_pad):
            # Balance persistent waves: with 256x256 tiles, pick the smallest
            # cluster count that keeps the same wave count as the full GPU
            # (e.g. 256 tiles -> 64 clusters x 4 exact waves beats 74 ragged).
            max_clusters = _num_sms() // 2
            tiles = ((M + 255) // 256) * ((N + 255) // 256)
            waves = -(-tiles // max_clusters)
            clusters = min(max_clusters, -(-tiles // waves))
            sms = 2 * clusters
            xp = self._xpad if needs_pad else x

            def launch():
                if needs_pad:
                    _pad_copy(x, self._xpad)
                _HAND.gemm_fp8_2sm(xp, w, scale, y, sms)
        else:
            BM, BN, BK, GM, stages, warps, subtile = _PERSISTENT_CFGS.get(
                (M, N), _PERSISTENT_DEFAULT)
            NUM_SMS = _num_sms()
            grid = (min(NUM_SMS, triton.cdiv(M, BM) * triton.cdiv(N, BN)),)
            xp = self._xpad if needs_pad else x

            def launch():
                if needs_pad:
                    _pad_copy(x, self._xpad)
                _gemm_persistent[grid](xp, w, y, scale, M, N, K_pad,
                                       BM, BN, BK, GM, NUM_SMS, subtile,
                                       num_stages=stages, num_warps=warps)
        return launch

    # -- forward ------------------------------------------------------------
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_contiguous():
            x = x.contiguous()
        M, K = x.shape
        N = self.weight.shape[0]
        Kp = _pad_k(K)
        needs_pad = Kp != K

        if needs_pad:
            self._padded_weight()  # refresh cache if weight mutated (checked every call)
            w = self._wpad
            if self._xpad is None or self._xpad.shape != (M, Kp):
                self._xpad = torch.zeros(M, Kp, device=x.device, dtype=x.dtype)
        else:
            w = self.weight

        key = (x.data_ptr(), M, K, x.stride(0))
        plan = self._plans.get(key)
        if plan is None:
            y = torch.empty(M, N, device=x.device, dtype=torch.bfloat16)
            launch = self._make_launch(x, w, y, M, N, Kp)
            plan = _Plan(launch, y)
            launch()  # compile + warm
            if self._graphs_ok:
                try:
                    torch.cuda.synchronize()
                    g = torch.cuda.CUDAGraph()
                    with torch.cuda.graph(g):
                        launch()
                    plan.graph = g
                except Exception:
                    self._graphs_ok = False
                    plan.graph = None
            if len(self._plans) > 64:
                self._plans.clear()
            self._plans[key] = plan
            if plan.graph is not None:
                plan.graph.replay()
            return plan.y

        if plan.graph is not None:
            plan.graph.replay()
        else:
            plan.launch()
        return plan.y


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]