submission 750036

#!POPCORN leaderboard amd-mxfp4-mm
#!POPCORN gpu MI355X
"""
v28: FULL SHAPE SPECIALIZATION. Template NTW + K128 + K32.

v27 WIN (LB 8.16): template-NTW -> `bid/NTW` magic-mul instead of s_div_u32.
  Same VGPR, -0.2 to -0.5us. m=64 ranked variance 0.09->0.05.

v28 EXTENDS: K128 is ALSO template. Cascading effects:
  1. K-loop bound `for(ks<K128;ks+=KU)` compile-time -> FULLY UNROLLS.
     m=64: K128=16 KU=4 -> 4 straight-line iterations. Zero induction/branch.
     m=256: K128=12 KU=4 -> 3 iterations.
  2. Ald_stride=K128*1024, Asd_stride=K128*64 -> constexpr -> every LDS
     offset `r*stride + ksi*1024 + L*16` is base+literal, no SGPR math.
  3. Prologue loop `for(g<MR*16*K32;g+=nthr)` — at m=64 W=8 MR=1 K32=64:
     1024 groups / 512 threads = 2 iters EXACTLY. With template K32:
     compile-time known -> UNROLLS. Two inlined hw_quant32, no loop.
  4. Bsc preload: KP = K128/2 already template (good). `for(p<KP)` already
     unrolls. No change needed.

`fgemm_alds3f<W,MR,KU,KP,BSC,NTW,K128T>`: K passed for ptr arithmetic
(Bsh_L offset uses K*8), but K128/K32 derive from template K128T.

SHAPE TABLE (template NTW,K128 pairs):
  m=64  k=2048: NTW=56  K128=16  (K32=64)
  m=256 k=1536: NTW=24  K128=12  (K32=48)
  SECRET m=64/m=16 k=1536: NTW=24  K128=12  (shared!)
  SECRET m=256 k=512:      NTW=45  K128=4   (W=4)
  ~4 unique kernels total. Compile budget trivial.

If K-loop unroll composes with v27's NTW win: maybe -0.3 to -0.8us more.
If prologue unroll matters: bonus on top.
"""
import os, sys, time
os.environ.setdefault("PYTORCH_ROCM_ARCH", "gfx950")
import warnings; warnings.filterwarnings("ignore")
import torch
import triton
import triton.language as tl

_L = lambda *a: print(*a, file=sys.stderr, flush=True)


_HIP_SRC = r"""
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#include <cstdint>

typedef int   i32x4 __attribute__((ext_vector_type(4)));
typedef int   i32x8 __attribute__((ext_vector_type(8)));
typedef float f32x4 __attribute__((ext_vector_type(4)));
typedef __hip_bfloat16 bf16;

__device__ __forceinline__ uint32_t f2u(float x){
  union{float f;uint32_t u;}c;c.f=x;return c.u;}
__device__ __forceinline__ float e8f(uint8_t e){
  union{uint32_t u;float f;}c;c.u=(uint32_t)e<<23;return c.f;}
#define QCV(o,a,b,s,bs) __builtin_amdgcn_cvt_scalef32_pk_fp4_f32((o),(a),(b),(s),(bs))

__device__ __forceinline__ i32x8 w8(i32x4 x){
  i32x8 r={0,0,0,0,0,0,0,0};r[0]=x[0];r[1]=x[1];r[2]=x[2];r[3]=x[3];return r;}

__device__ __forceinline__
void hw_quant32(const int4* __restrict__ a4,i32x4& o,int& e8){
  bf16 ab[32] __attribute__((aligned(16)));
  *reinterpret_cast<int4*>(&ab[ 0])=a4[0];
  *reinterpret_cast<int4*>(&ab[ 8])=a4[1];
  *reinterpret_cast<int4*>(&ab[16])=a4[2];
  *reinterpret_cast<int4*>(&ab[24])=a4[3];
  float v[32];float amax=0.f;
  #pragma unroll
  for(int i=0;i<32;++i){v[i]=(float)ab[i];
    float t=__builtin_fabsf(v[i]);amax=t>amax?t:amax;}
  uint32_t au=(f2u(amax)+0x200000u)&0xFF800000u;
  int su=au?(int)((au>>23)&0xFFu)-129:-127;
  su=su<-127?-127:(su>127?127:su);e8=su+127;
  float bsc=e8f((uint8_t)e8);
  int w0=0,w1=0,w2=0,w3=0;
  w0=QCV(w0,v[ 0],v[ 1],bsc,0);w0=QCV(w0,v[ 2],v[ 3],bsc,1);
  w0=QCV(w0,v[ 4],v[ 5],bsc,2);w0=QCV(w0,v[ 6],v[ 7],bsc,3);
  w1=QCV(w1,v[ 8],v[ 9],bsc,0);w1=QCV(w1,v[10],v[11],bsc,1);
  w1=QCV(w1,v[12],v[13],bsc,2);w1=QCV(w1,v[14],v[15],bsc,3);
  w2=QCV(w2,v[16],v[17],bsc,0);w2=QCV(w2,v[18],v[19],bsc,1);
  w2=QCV(w2,v[20],v[21],bsc,2);w2=QCV(w2,v[22],v[23],bsc,3);
  w3=QCV(w3,v[24],v[25],bsc,0);w3=QCV(w3,v[26],v[27],bsc,1);
  w3=QCV(w3,v[28],v[29],bsc,2);w3=QCV(w3,v[30],v[31],bsc,3);
  o=(i32x4){w0,w1,w2,w3};
}

#define ALDS_PROLOGUE(M_REP) \
  const int K32=K>>5, K128=K>>7; \
  const long Ald_stride=(long)K128*1024; \
  const long Asd_stride=(long)K128*64; \
  const long Asd_base=(long)(M_REP)*Ald_stride; \
  extern __shared__ uint8_t _sh[]; \
  uint8_t* Ald=_sh; uint8_t* Asd=_sh+Asd_base; \
  { const int nthr=WAVES*64; \
    const int ngrp=(M_REP)*16*K32; \
    const int m_base=m_tile*(M_REP)*16; \
    for(int g=tid; g<ngrp; g+=nthr){ \
      const int r=g/K32; const int kb=g%K32; \
      const int r_tile=r>>4; const int r16=r&15; \
      const int k128s=kb>>2; const int kg4=kb&3; \
      const int Lw=kg4*16+r16; \
      const int m=m_base+r; \
      i32x4 o={0,0,0,0}; int e8=0; \
      if(m<M){ \
        const bf16* Ap=A+(long)m*K+(long)kb*32; \
        int4 ai[4]; \
        ai[0]=*reinterpret_cast<const int4*>(Ap); \
        ai[1]=*reinterpret_cast<const int4*>(Ap+8); \
        ai[2]=*reinterpret_cast<const int4*>(Ap+16); \
        ai[3]=*reinterpret_cast<const int4*>(Ap+24); \
        hw_quant32(ai,o,e8); \
      } \
      *reinterpret_cast<i32x4*>(Ald+(long)r_tile*Ald_stride+(long)k128s*1024+Lw*16)=o; \
      Asd[(long)r_tile*Asd_stride+(long)k128s*64+Lw]=(uint8_t)e8; \
    } } \
  __syncthreads()

#define ALDS_SETUP(M_REP) \
  if(!vn) return; \
  const long n_col=(long)n_tile*16+m16; \
  const long bsc_row=(n_col>>5)*(sn8*256)+(n_col&15)*4+((n_col>>4)&1)+(long)kg*64; \
  const uint8_t* Bsc_r=Bsc+bsc_row; \
  const uint8_t* Bsh_L=Bsh+(long)n_tile*(long)K*8+L*16; \
  const uint8_t* Ald_L=Ald+L*16; \
  const uint8_t* Asd_L=Asd+L; \
  f32x4 acc[M_REP]; \
  _Pragma("unroll") for(int r=0;r<(M_REP);++r) acc[r]=(f32x4){0,0,0,0}; \
  int mmsk[M_REP]; \
  _Pragma("unroll") for(int r=0;r<(M_REP);++r){ \
    int m=(m_tile*(M_REP)+r)*16+m16; mmsk[r]=(m<M)?0xFF:0; }

#define ALDS_STORE(M_REP) \
  _Pragma("unroll") for(int r=0;r<(M_REP);++r) \
    _Pragma("unroll") for(int i=0;i<4;++i){ \
      int mo=(m_tile*(M_REP)+r)*16+kg*4+i; \
      if(mo<M) C[(long)mo*N+n_col]=(bf16)acc[r][i]; }

#define LDA(abuf,sbuf,r,ksi) do{ \
    abuf=*reinterpret_cast<const i32x4*>(Ald_L+(long)(r)*Ald_stride+(long)(ksi)*1024); \
    sbuf=(int)Asd_L[(long)(r)*Asd_stride+(long)(ksi)*64]&mmsk[r]; }while(0)

// ───── alds3b: Bsc preload. REQUIRES K128 even, KP = K128/2 (template). ─────
template<int WAVES,int M_REP,int KU,int KP>
__global__ __launch_bounds__(WAVES*64)
void fgemm_alds3b(
    const bf16* __restrict__ A,const uint8_t* __restrict__ Bsh,
    const uint8_t* __restrict__ Bsc,bf16* __restrict__ C,
    int M,int N,int K,long sn8,int NT)
{
  const int tid=threadIdx.x,L=tid&63,w=tid>>6;
  const int m16=L&15,kg=L>>4;const int bid=blockIdx.x;
  const int NTW=(NT+WAVES-1)/WAVES;
  const int m_tile=bid/NTW, ntg=bid%NTW;
  const int n_tile=ntg*WAVES+w;const bool vn=(n_tile<NT);
  ALDS_PROLOGUE(M_REP);
  ALDS_SETUP(M_REP);
  // Bsc preload: KP int32s, each contains bytes for ks=2p and ks=2p+1.
  int bsc_pk[KP];
  #pragma unroll
  for(int p=0;p<KP;++p)
    bsc_pk[p]=*reinterpret_cast<const int*>(Bsc_r+(long)p*256);
  for(int ks=0;ks<K128;ks+=KU){
    i32x4 bb[KU]; int bsv[KU];
    i32x4 ab[KU][M_REP]; int asv[KU][M_REP];
    #pragma unroll
    for(int u=0;u<KU;++u){
      const int ksi=ks+u;
      bb[u]=*reinterpret_cast<const i32x4*>(Bsh_L+(long)ksi*1024);
      bsv[u]=(bsc_pk[ksi>>1] >> ((ksi&1)*16)) & 0xFF;
      #pragma unroll
      for(int r=0;r<M_REP;++r) LDA(ab[u][r],asv[u][r],r,ksi);
    }
    __builtin_amdgcn_sched_barrier(0);
    #pragma unroll
    for(int u=0;u<KU;++u){
      i32x8 b8=w8(bb[u]);
      #pragma unroll
      for(int r=0;r<M_REP;++r)
        acc[r]=__builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(
            w8(ab[u][r]),b8,acc[r],4,4,0,asv[u][r],0,bsv[u]);
    }
  }
  ALDS_STORE(M_REP);
}

// ═══════ fgemm_alds3f: FULL TEMPLATE (NTW + K128). v28. ═══════
// ALL shape-derived ints are compile-time:
//   NTW     -> bid/NTW magic mul              (v27 win)
//   K128T   -> K-loop bound constexpr -> LLVM FULLY UNROLLS K-loop
//   K32=K128T*4 -> prologue loop bound constexpr -> LLVM may unroll prologue
//   Ald_stride=K128T*1024 -> constexpr -> all LDS offsets base+literal
//   KP=K128T/2 (if BSC) -> Bsc preload fully unrolls (already did in v27)
// K runtime ONLY for: A/Bsh base pointer strides (M*K, n_tile*K*8).
// CONSTRAINT: runtime K must satisfy K>>7 == K128T. Python-enforced.
template<int WAVES,int M_REP,int KU,int KP,int BSC,int NTW,int K128T>
__global__ __launch_bounds__(WAVES*64)
void fgemm_alds3f(
    const bf16* __restrict__ A,const uint8_t* __restrict__ Bsh,
    const uint8_t* __restrict__ Bsc,bf16* __restrict__ C,
    int M,int N,int K,long sn8,int NT)
{
  constexpr int K128 = K128T;                    // K-loop bound: constexpr
  constexpr int K32  = K128T * 4;                // Prologue bound: constexpr
  constexpr long Ald_stride = (long)K128T * 1024;// LDS stride: constexpr
  constexpr long Asd_stride = (long)K128T * 64;
  constexpr long Asd_base   = (long)M_REP * Ald_stride;
  constexpr int  nthr       = WAVES * 64;
  constexpr int  ngrp       = M_REP * 16 * K32;  // Prologue groups: constexpr

  const int tid=threadIdx.x,L=tid&63,w=tid>>6;
  const int m16=L&15,kg=L>>4;const int bid=blockIdx.x;
  const int m_tile=bid/NTW, ntg=bid%NTW;         // NTW template (v27)
  const int n_tile=ntg*WAVES+w;const bool vn=(n_tile<NT);

  extern __shared__ uint8_t _sh[];
  uint8_t* Ald=_sh; uint8_t* Asd=_sh+Asd_base;
  { // PROLOGUE: ngrp/nthr iters, constexpr -> unroll candidate.
    const int m_base=m_tile*M_REP*16;
    #pragma unroll
    for(int g=tid; g<ngrp; g+=nthr){
      const int r=g/K32; const int kb=g%K32;     // K32 template -> magic div
      const int r_tile=r>>4; const int r16=r&15;
      const int k128s=kb>>2; const int kg4=kb&3;
      const int Lw=kg4*16+r16;
      const int m=m_base+r;
      i32x4 o={0,0,0,0}; int e8=0;
      if(m<M){
        const bf16* Ap=A+(long)m*K+(long)kb*32;  // K runtime (ptr stride)
        int4 ai[4];
        ai[0]=*reinterpret_cast<const int4*>(Ap);
        ai[1]=*reinterpret_cast<const int4*>(Ap+8);
        ai[2]=*reinterpret_cast<const int4*>(Ap+16);
        ai[3]=*reinterpret_cast<const int4*>(Ap+24);
        hw_quant32(ai,o,e8);
      }
      *reinterpret_cast<i32x4*>(Ald+(long)r_tile*Ald_stride+(long)k128s*1024+Lw*16)=o;
      Asd[(long)r_tile*Asd_stride+(long)k128s*64+Lw]=(uint8_t)e8;
    }
  }
  __syncthreads();
  if(!vn) return;

  const long n_col=(long)n_tile*16+m16;
  const long bsc_row=(n_col>>5)*(sn8*256)+(n_col&15)*4+((n_col>>4)&1)+(long)kg*64;
  const uint8_t* Bsc_r=Bsc+bsc_row;
  const uint8_t* Bsh_L=Bsh+(long)n_tile*(long)K*8+L*16;  // K runtime
  const uint8_t* Ald_L=Ald+L*16;
  const uint8_t* Asd_L=Asd+L;
  f32x4 acc[M_REP];
  #pragma unroll
  for(int r=0;r<M_REP;++r) acc[r]=(f32x4){0,0,0,0};
  int mmsk[M_REP];
  #pragma unroll
  for(int r=0;r<M_REP;++r){
    int m=(m_tile*M_REP+r)*16+m16;mmsk[r]=(m<M)?0xFF:0;}

  int bsc_pk[KP>0?KP:1];
  if constexpr(BSC){
    #pragma unroll
    for(int p=0;p<KP;++p)bsc_pk[p]=*reinterpret_cast<const int*>(Bsc_r+(long)p*256);
  }

  // K-LOOP: K128 constexpr -> LLVM unrolls this COMPLETELY (K128/KU iters).
  // sched_barrier kept (v19 proved it's the KU-batch mechanism).
  #pragma unroll
  for(int ks=0;ks<K128;ks+=KU){
    i32x4 bb[KU];int bsv[KU];i32x4 ab[KU][M_REP];int asv[KU][M_REP];
    #pragma unroll
    for(int u=0;u<KU;++u){const int ksi=ks+u;
      bb[u]=*reinterpret_cast<const i32x4*>(Bsh_L+(long)ksi*1024);
      if constexpr(BSC){
        bsv[u]=(bsc_pk[ksi>>1]>>((ksi&1)*16))&0xFF;
      } else {
        bsv[u]=(int)Bsc_r[(long)(ksi>>1)*256+(ksi&1)*2];
      }
      #pragma unroll
      for(int r=0;r<M_REP;++r){
        ab[u][r]=*reinterpret_cast<const i32x4*>(Ald_L+(long)r*Ald_stride+(long)ksi*1024);
        asv[u][r]=(int)Asd_L[(long)r*Asd_stride+(long)ksi*64]&mmsk[r];
      }}
    __builtin_amdgcn_sched_barrier(0);
    #pragma unroll
    for(int u=0;u<KU;++u){i32x8 b8=w8(bb[u]);
      #pragma unroll
      for(int r=0;r<M_REP;++r)
        acc[r]=__builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(
            w8(ab[u][r]),b8,acc[r],4,4,0,asv[u][r],0,bsv[u]);}}

  #pragma unroll
  for(int r=0;r<M_REP;++r)
    #pragma unroll
    for(int i=0;i<4;++i){int mo=(m_tile*M_REP+r)*16+kg*4+i;
      if(mo<M)C[(long)mo*N+n_col]=(bf16)acc[r][i];}
}

// ───── alds3bt: alds3b + TEMPLATE NTW (v26 discovery) ─────
// ONLY delta vs alds3b: NTW is template. bid/NTW -> magic mul (not s_div_u32).
// v26 head-to-head: m=256 9.74 vs alds3b 10.23 SAME RUN SAME VGPR. -0.49us FREE.
template<int WAVES,int M_REP,int KU,int KP,int NTW>
__global__ __launch_bounds__(WAVES*64)
void fgemm_alds3bt(
    const bf16* __restrict__ A,const uint8_t* __restrict__ Bsh,
    const uint8_t* __restrict__ Bsc,bf16* __restrict__ C,
    int M,int N,int K,long sn8,int NT)
{
  const int tid=threadIdx.x,L=tid&63,w=tid>>6;
  const int m16=L&15,kg=L>>4;const int bid=blockIdx.x;
  // NTW is TEMPLATE -> bid/NTW = magic mul+shift at compile time.
  const int m_tile=bid/NTW, ntg=bid%NTW;
  const int n_tile=ntg*WAVES+w;const bool vn=(n_tile<NT);
  ALDS_PROLOGUE(M_REP);
  if(!vn) return;
  const long n_col=(long)n_tile*16+m16;
  const long bsc_row=(n_col>>5)*(sn8*256)+(n_col&15)*4+((n_col>>4)&1)+(long)kg*64;
  const uint8_t* Bsc_r=Bsc+bsc_row;
  const uint8_t* Bsh_L=Bsh+(long)n_tile*(long)K*8+L*16;
  const uint8_t* Ald_L=Ald+L*16;const uint8_t* Asd_L=Asd+L;
  f32x4 acc[M_REP];
  #pragma unroll
  for(int r=0;r<M_REP;++r) acc[r]=(f32x4){0,0,0,0};
  int mmsk[M_REP];
  #pragma unroll
  for(int r=0;r<M_REP;++r){
    int m=(m_tile*M_REP+r)*16+m16;mmsk[r]=(m<M)?0xFF:0;}
  int bsc_pk[KP];
  #pragma unroll
  for(int p=0;p<KP;++p)bsc_pk[p]=*reinterpret_cast<const int*>(Bsc_r+(long)p*256);
  for(int ks=0;ks<K128;ks+=KU){
    i32x4 bb[KU];int bsv[KU];i32x4 ab[KU][M_REP];int asv[KU][M_REP];
    #pragma unroll
    for(int u=0;u<KU;++u){const int ksi=ks+u;
      bb[u]=*reinterpret_cast<const i32x4*>(Bsh_L+(long)ksi*1024);
      bsv[u]=(bsc_pk[ksi>>1]>>((ksi&1)*16))&0xFF;
      #pragma unroll
      for(int r=0;r<M_REP;++r)LDA(ab[u][r],asv[u][r],r,ksi);}
    __builtin_amdgcn_sched_barrier(0);
    #pragma unroll
    for(int u=0;u<KU;++u){i32x8 b8=w8(bb[u]);
      #pragma unroll
      for(int r=0;r<M_REP;++r)
        acc[r]=__builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(
            w8(ab[u][r]),b8,acc[r],4,4,0,asv[u][r],0,bsv[u]);}}
  #pragma unroll
  for(int r=0;r<M_REP;++r)
    #pragma unroll
    for(int i=0;i<4;++i){int mo=(m_tile*M_REP+r)*16+kg*4+i;
      if(mo<M)C[(long)mo*N+n_col]=(bf16)acc[r][i];}
}

// ───── alds3t: alds3 + TEMPLATE NTW (no bsc) ─────
template<int WAVES,int M_REP,int KU,int NTW>
__global__ __launch_bounds__(WAVES*64)
void fgemm_alds3t(
    const bf16* __restrict__ A,const uint8_t* __restrict__ Bsh,
    const uint8_t* __restrict__ Bsc,bf16* __restrict__ C,
    int M,int N,int K,long sn8,int NT)
{
  const int tid=threadIdx.x,L=tid&63,w=tid>>6;
  const int m16=L&15,kg=L>>4;const int bid=blockIdx.x;
  const int m_tile=bid/NTW, ntg=bid%NTW;
  const int n_tile=ntg*WAVES+w;const bool vn=(n_tile<NT);
  ALDS_PROLOGUE(M_REP);
  if(!vn) return;
  const long n_col=(long)n_tile*16+m16;
  const long bsc_row=(n_col>>5)*(sn8*256)+(n_col&15)*4+((n_col>>4)&1)+(long)kg*64;
  const uint8_t* Bsc_r=Bsc+bsc_row;
  const uint8_t* Bsh_L=Bsh+(long)n_tile*(long)K*8+L*16;
  const uint8_t* Ald_L=Ald+L*16;const uint8_t* Asd_L=Asd+L;
  f32x4 acc[M_REP];
  #pragma unroll
  for(int r=0;r<M_REP;++r) acc[r]=(f32x4){0,0,0,0};
  int mmsk[M_REP];
  #pragma unroll
  for(int r=0;r<M_REP;++r){
    int m=(m_tile*M_REP+r)*16+m16;mmsk[r]=(m<M)?0xFF:0;}
  for(int ks=0;ks<K128;ks+=KU){
    i32x4 bb[KU];int bsv[KU];i32x4 ab[KU][M_REP];int asv[KU][M_REP];
    #pragma unroll
    for(int u=0;u<KU;++u){const int ksi=ks+u;
      bb[u]=*reinterpret_cast<const i32x4*>(Bsh_L+(long)ksi*1024);
      bsv[u]=(int)Bsc_r[(long)(ksi>>1)*256+(ksi&1)*2];
      #pragma unroll
      for(int r=0;r<M_REP;++r)LDA(ab[u][r],asv[u][r],r,ksi);}
    __builtin_amdgcn_sched_barrier(0);
    #pragma unroll
    for(int u=0;u<KU;++u){i32x8 b8=w8(bb[u]);
      #pragma unroll
      for(int r=0;r<M_REP;++r)
        acc[r]=__builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(
            w8(ab[u][r]),b8,acc[r],4,4,0,asv[u][r],0,bsv[u]);}}
  #pragma unroll
  for(int r=0;r<M_REP;++r)
    #pragma unroll
    for(int i=0;i<4;++i){int mo=(m_tile*M_REP+r)*16+kg*4+i;
      if(mo<M)C[(long)mo*N+n_col]=(bf16)acc[r][i];}
}

// ───── alds3: v18 fallback (for K128 odd or KP unavailable) ─────
template<int WAVES,int M_REP,int KU>
__global__ __launch_bounds__(WAVES*64)
void fgemm_alds3(
    const bf16* __restrict__ A,const uint8_t* __restrict__ Bsh,
    const uint8_t* __restrict__ Bsc,bf16* __restrict__ C,
    int M,int N,int K,long sn8,int NT)
{
  const int tid=threadIdx.x,L=tid&63,w=tid>>6;
  const int m16=L&15,kg=L>>4;const int bid=blockIdx.x;
  const int NTW=(NT+WAVES-1)/WAVES;
  const int m_tile=bid/NTW, ntg=bid%NTW;
  const int n_tile=ntg*WAVES+w;const bool vn=(n_tile<NT);
  ALDS_PROLOGUE(M_REP);
  ALDS_SETUP(M_REP);
  for(int ks=0;ks<K128;ks+=KU){
    i32x4 bb[KU]; int bsv[KU];
    i32x4 ab[KU][M_REP]; int asv[KU][M_REP];
    #pragma unroll
    for(int u=0;u<KU;++u){
      const int ksi=ks+u;
      bb[u]=*reinterpret_cast<const i32x4*>(Bsh_L+(long)ksi*1024);
      bsv[u]=(int)Bsc_r[(long)(ksi>>1)*256+(ksi&1)*2];
      #pragma unroll
      for(int r=0;r<M_REP;++r) LDA(ab[u][r],asv[u][r],r,ksi);
    }
    __builtin_amdgcn_sched_barrier(0);
    #pragma unroll
    for(int u=0;u<KU;++u){
      i32x8 b8=w8(bb[u]);
      #pragma unroll
      for(int r=0;r<M_REP;++r)
        acc[r]=__builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(
            w8(ab[u][r]),b8,acc[r],4,4,0,asv[u][r],0,bsv[u]);
    }
  }
  ALDS_STORE(M_REP);
}

// ═══ alds_sk (v18 — unchanged, proven) ═══
template<int WAVES,int KU>
__global__ __launch_bounds__(WAVES*64)
void fgemm_alds_sk(
    const bf16* __restrict__ A,const uint8_t* __restrict__ Bsh,
    const uint8_t* __restrict__ Bsc,float* __restrict__ Cf,
    int M,int N,int K,long sn8,int NT,int SK)
{
  const int tid=threadIdx.x,L=tid&63,w=tid>>6;
  const int m16=L&15,kg=L>>4;const int bid=blockIdx.x;
  const int NTW=(NT+WAVES-1)/WAVES;
  const int pk=bid/NTW, ntg=bid%NTW;
  const int n_tile=ntg*WAVES+w;const bool vn=(n_tile<NT);
  const int K32=K>>5, K128=K>>7;
  const int Kps128=(K128+SK-1)/SK;
  const int ks_lo=pk*Kps128;const int ks_hi=min(ks_lo+Kps128,K128);
  const int nsl=ks_hi-ks_lo;if(nsl<=0)return;
  extern __shared__ uint8_t _sh[];
  uint8_t* Ald=_sh;uint8_t* Asd=_sh+(long)nsl*1024;
  { const int nthr=WAVES*64;const int nkg=nsl*4;
    const int ngrp=16*nkg;const int kb_lo=ks_lo*4;
    for(int g=tid;g<ngrp;g+=nthr){
      const int r=g/nkg;const int kbs=g%nkg;const int kb=kb_lo+kbs;
      const int k128s=kbs>>2;const int kg4=kbs&3;const int Lw=kg4*16+r;
      i32x4 o={0,0,0,0};int e8=0;
      if(r<M){const bf16* Ap=A+(long)r*K+(long)kb*32;int4 ai[4];
        ai[0]=*reinterpret_cast<const int4*>(Ap);
        ai[1]=*reinterpret_cast<const int4*>(Ap+8);
        ai[2]=*reinterpret_cast<const int4*>(Ap+16);
        ai[3]=*reinterpret_cast<const int4*>(Ap+24);
        hw_quant32(ai,o,e8);}
      *reinterpret_cast<i32x4*>(Ald+(long)k128s*1024+Lw*16)=o;
      Asd[(long)k128s*64+Lw]=(uint8_t)e8;}}
  __syncthreads();if(!vn)return;
  const long n_col=(long)n_tile*16+m16;
  const long bsc_row=(n_col>>5)*(sn8*256)+(n_col&15)*4+((n_col>>4)&1)+(long)kg*64;
  const uint8_t* Bsc_r=Bsc+bsc_row;
  const uint8_t* Bsh_L=Bsh+(long)n_tile*(long)K*8+L*16;
  const uint8_t* Ald_L=Ald+L*16;const uint8_t* Asd_L=Asd+L;
  f32x4 acc={0,0,0,0};const int mmsk=(m16<M)?0xFF:0;
  for(int ks_l=0;ks_l<nsl;ks_l+=KU){
    i32x4 bb[KU];int bsv[KU];i32x4 ab[KU];int asv[KU];
    const int lim=min(KU,nsl-ks_l);
    #pragma unroll
    for(int u=0;u<KU;++u){const int ksi_l=ks_l+u,ksi_g=ks_lo+ksi_l;
      if(u<lim){bb[u]=*reinterpret_cast<const i32x4*>(Bsh_L+(long)ksi_g*1024);
        bsv[u]=(int)Bsc_r[(long)(ksi_g>>1)*256+(ksi_g&1)*2];
        ab[u]=*reinterpret_cast<const i32x4*>(Ald_L+(long)ksi_l*1024);
        asv[u]=(int)Asd_L[(long)ksi_l*64]&mmsk;
      }else{bb[u]=(i32x4){0,0,0,0};bsv[u]=0;ab[u]=(i32x4){0,0,0,0};asv[u]=0;}}
    __builtin_amdgcn_sched_barrier(0);
    #pragma unroll
    for(int u=0;u<KU;++u)
      acc=__builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(
          w8(ab[u]),w8(bb[u]),acc,4,4,0,asv[u],0,bsv[u]);}
  #pragma unroll
  for(int i=0;i<4;++i){int mo=kg*4+i;
    if(mo<M)atomicAdd(&Cf[(long)mo*N+n_col],acc[i]);}
}

__global__ __launch_bounds__(256)
void cast_f32_bf16_z(const float* __restrict__ S,bf16* __restrict__ C,
                     float* __restrict__ Sz,long N){
  long g=(long)blockIdx.x*256+threadIdx.x;
  if(g<N){C[g]=(bf16)S[g];Sz[g]=0.0f;}
}

// ═══ fqn / fq (small-M / safety) ═══
template<int WAVES,int N_REP>
__global__ __launch_bounds__(WAVES*64)
void fgemm_fqn(
    const bf16* __restrict__ A,const uint8_t* __restrict__ Bsh,
    const uint8_t* __restrict__ Bsc,bf16* __restrict__ C,
    int M,int N,int K,long sn8,int NT)
{
  const int tid=threadIdx.x,L=tid&63,w=tid>>6;
  const int m16=L&15,kg=L>>4;const int bid=blockIdx.x;
  const int NTG=(NT+N_REP-1)/N_REP;
  const int m_tile=bid/NTG,ntg=bid%NTG;
  long ksz=((K/128+WAVES-1)/WAVES)*128;
  long k_lo=(long)w*ksz,k_hi=min(k_lo+ksz,(long)K);
  long n_col[N_REP];int vnm[N_REP];const uint8_t* Bsh_t[N_REP];
  #pragma unroll
  for(int nr=0;nr<N_REP;++nr){int nt=ntg*N_REP+nr;int v=(nt<NT);
    vnm[nr]=v?0xFF:0;long ntr=v?nt:0;
    n_col[nr]=ntr*16+m16;Bsh_t[nr]=Bsh+ntr*(long)K*8;}
  const int m_row=m_tile*16+m16;const bool vm=m_row<M;
  const long mrow=vm?m_row:0;
  f32x4 acc[N_REP];
  #pragma unroll
  for(int nr=0;nr<N_REP;++nr)acc[nr]=(f32x4){0,0,0,0};
  for(long k=k_lo;k<k_hi;k+=128){
    long kb_=(k>>5)*256+L*16,ks=(k>>5)+kg;
    i32x4 bb[N_REP];int bsv[N_REP];
    #pragma unroll
    for(int nr=0;nr<N_REP;++nr){
      bb[nr]=*reinterpret_cast<const i32x4*>(Bsh_t[nr]+kb_);
      long c=ks;
      bsv[nr]=(int)Bsc[(n_col[nr]>>5)*(sn8*256)+(n_col[nr]&15)*4+((n_col[nr]>>4)&1)
                      +(c>>3)*256+(c&3)*64+((c>>2)&1)*2]&vnm[nr];}
    const long kba=k+(long)kg*32;const bf16* Ap=A+mrow*K+kba;int4 ai[4];
    ai[0]=*reinterpret_cast<const int4*>(Ap);
    ai[1]=*reinterpret_cast<const int4*>(Ap+8);
    ai[2]=*reinterpret_cast<const int4*>(Ap+16);
    ai[3]=*reinterpret_cast<const int4*>(Ap+24);
    i32x4 a4;int a_sc;hw_quant32(ai,a4,a_sc);if(!vm)a_sc=0;
    i32x8 a8=w8(a4);
    #pragma unroll
    for(int nr=0;nr<N_REP;++nr)
      acc[nr]=__builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(
          a8,w8(bb[nr]),acc[nr],4,4,0,a_sc,0,bsv[nr]);}
  extern __shared__ float red[];
  #pragma unroll
  for(int nr=0;nr<N_REP;++nr)
    #pragma unroll
    for(int i=0;i<4;++i)red[((long)w*N_REP+nr)*256+L*4+i]=acc[nr][i];
  __syncthreads();if(w!=0)return;
  #pragma unroll
  for(int nr=0;nr<N_REP;++nr)
    #pragma unroll
    for(int i=0;i<4;++i){float s=0;
      #pragma unroll
      for(int ww=0;ww<WAVES;++ww)s+=red[((long)ww*N_REP+nr)*256+L*4+i];
      acc[nr][i]=s;}
  #pragma unroll
  for(int i=0;i<4;++i){int mo=m_tile*16+kg*4+i;if(mo>=M)continue;
    #pragma unroll
    for(int nr=0;nr<N_REP;++nr)
      if(vnm[nr])C[(long)mo*N+n_col[nr]]=(bf16)acc[nr][i];}
}

template<int WAVES>
__global__ __launch_bounds__(WAVES*64)
void fgemm_fq(
    const bf16* __restrict__ A,const uint8_t* __restrict__ Bsh,
    const uint8_t* __restrict__ Bsc,bf16* __restrict__ C,
    int M,int N,int K,long sn8,int NT)
{
  const int tid=threadIdx.x,L=tid&63,w=tid>>6;
  const int m16=L&15,kg=L>>4;const int bid=blockIdx.x;
  int m_tile=bid/NT,n_tile=bid%NT;
  long ksz=((K/128+WAVES-1)/WAVES)*128;
  long k_lo=(long)w*ksz,k_hi=min(k_lo+ksz,(long)K);
  const long n_col=(long)n_tile*16+m16;
  const uint8_t* Bsh_t=Bsh+(long)n_tile*(long)K*8;
  f32x4 acc={0,0,0,0};const int m_row=m_tile*16+m16;const bool vm=m_row<M;
  for(long k=k_lo;k<k_hi;k+=128){
    i32x4 b4=*reinterpret_cast<const i32x4*>(Bsh_t+(k>>5)*256+L*16);
    long c=(k>>5)+kg;
    int b_sc=(int)Bsc[(n_col>>5)*(sn8*256)+(n_col&15)*4+((n_col>>4)&1)
                     +(c>>3)*256+(c&3)*64+((c>>2)&1)*2];
    const long kb=k+(long)kg*32;
    const bf16* Ap=A+(long)(vm?m_row:0)*K+kb;int4 ai[4];
    ai[0]=*reinterpret_cast<const int4*>(Ap);
    ai[1]=*reinterpret_cast<const int4*>(Ap+8);
    ai[2]=*reinterpret_cast<const int4*>(Ap+16);
    ai[3]=*reinterpret_cast<const int4*>(Ap+24);
    i32x4 a4;int a_sc;hw_quant32(ai,a4,a_sc);if(!vm)a_sc=0;
    acc=__builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(
        w8(a4),w8(b4),acc,4,4,0,a_sc,0,b_sc);}
  extern __shared__ float red[];
  #pragma unroll
  for(int i=0;i<4;++i)red[(long)w*256+L*4+i]=acc[i];
  __syncthreads();if(w!=0)return;
  #pragma unroll
  for(int i=0;i<4;++i){float s=0;
    #pragma unroll
    for(int ww=0;ww<WAVES;++ww)s+=red[(long)ww*256+L*4+i];acc[i]=s;}
  #pragma unroll
  for(int i=0;i<4;++i){int mo=m_tile*16+kg*4+i;
    if(mo<M)C[(long)mo*N+n_col]=(bf16)acc[i];}
}

#include <torch/extension.h>

template<int W,int MR,int KU,int KP>
static void _ga3b(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,int64_t NT){
  const int64_t K128=K>>7;const int64_t MT=(M+16*MR-1)/(16*MR);
  const int64_t NTW=(NT+W-1)/W;const int64_t gx=MT*NTW;
  const int64_t lds=(int64_t)MR*K128*1088;
  static bool _s=false;if(!_s){(void)hipFuncSetAttribute(
    (const void*)fgemm_alds3b<W,MR,KU,KP>,hipFuncAttributeMaxDynamicSharedMemorySize,160*1024);_s=true;}
  fgemm_alds3b<W,MR,KU,KP><<<dim3(gx),dim3(W*64),lds,0>>>(
    reinterpret_cast<const bf16*>(A.data_ptr()),
    Bsh.data_ptr<uint8_t>(),Bsc.data_ptr<uint8_t>(),
    reinterpret_cast<bf16*>(C.data_ptr()),(int)M,(int)N,(int)K,sn8,(int)NT);
}
template<int W,int MR,int KU>
static void _ga3(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,int64_t NT){
  const int64_t K128=K>>7;const int64_t MT=(M+16*MR-1)/(16*MR);
  const int64_t NTW=(NT+W-1)/W;const int64_t gx=MT*NTW;
  const int64_t lds=(int64_t)MR*K128*1088;
  static bool _s=false;if(!_s){(void)hipFuncSetAttribute(
    (const void*)fgemm_alds3<W,MR,KU>,hipFuncAttributeMaxDynamicSharedMemorySize,160*1024);_s=true;}
  fgemm_alds3<W,MR,KU><<<dim3(gx),dim3(W*64),lds,0>>>(
    reinterpret_cast<const bf16*>(A.data_ptr()),
    Bsh.data_ptr<uint8_t>(),Bsc.data_ptr<uint8_t>(),
    reinterpret_cast<bf16*>(C.data_ptr()),(int)M,(int)N,(int)K,sn8,(int)NT);
}

// Full-template launcher. NTW + K128T both baked in. K runtime for ptr stride.
template<int W,int MR,int KU,int KP,int BSC,int NTW,int K128T>
static void _ga3f(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,int64_t NT){
  const int64_t MT=(M+16*MR-1)/(16*MR);
  const int64_t gx=MT*(int64_t)NTW;
  constexpr int64_t lds=(int64_t)MR*K128T*1088;
  static bool _s=false;if(!_s){(void)hipFuncSetAttribute(
    (const void*)fgemm_alds3f<W,MR,KU,KP,BSC,NTW,K128T>,
    hipFuncAttributeMaxDynamicSharedMemorySize,160*1024);_s=true;}
  fgemm_alds3f<W,MR,KU,KP,BSC,NTW,K128T><<<dim3(gx),dim3(W*64),lds,0>>>(
    reinterpret_cast<const bf16*>(A.data_ptr()),
    Bsh.data_ptr<uint8_t>(),Bsc.data_ptr<uint8_t>(),
    reinterpret_cast<bf16*>(C.data_ptr()),(int)M,(int)N,(int)K,sn8,(int)NT);
}

// Full-template dispatch. (NTW, K128T) pair must match compiled instance.
// BENCH: (56,16) m=64 | (24,12) m=256 & secret m=64/m=16
// SECRET: (45,4) m=256/k=512 W=4 | (23,4) m=32/k=512 W=8 | (17,56) m=8 no-alds
int64_t launch_alds3f(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,int64_t NT,
    int64_t W,int64_t MR,int64_t KU,int64_t BSC){
  int64_t K128=K>>7;
  int64_t NTW=(NT+W-1)/W;
  if((int64_t)MR*K128*1088>160*1024)return -2;
  if(K128%KU!=0)return -3;
  if(BSC&&(K128&1))return -4;
  // KP = K128/2 when BSC=1, dummy=1 when BSC=0. Folded into dispatch macros.
  #define D0(Ww,Rr,Uu,Nw,Kt) if(W==Ww&&MR==Rr&&KU==Uu&&BSC==0&&NTW==Nw&&K128==Kt){ \
      _ga3f<Ww,Rr,Uu,1,0,Nw,Kt>(A,Bsh,Bsc,C,M,N,K,sn8,NT);return 0;}
  #define D1(Ww,Rr,Uu,Nw,Kt) if(W==Ww&&MR==Rr&&KU==Uu&&BSC==1&&NTW==Nw&&K128==Kt){ \
      _ga3f<Ww,Rr,Uu,(Kt)/2,1,Nw,Kt>(A,Bsh,Bsc,C,M,N,K,sn8,NT);return 0;}
  // ── (NTW=56, K128=16): m=64 n=7168 k=2048 W=8. KU in {4,8,16} ──
  D0(8,1,4, 56,16);D0(8,1,8, 56,16);D0(8,1,16, 56,16);
  D1(8,1,4, 56,16);D1(8,1,8, 56,16);
  // W=4 variant: NTW=112
  D0(4,1,4, 112,16);D0(4,1,8, 112,16);D1(4,1,4, 112,16);D1(4,1,8, 112,16);
  // ── (NTW=24, K128=12): m=256 & secret-m=64/m=16 n=3072 k=1536 W=8. KU in {4,6,12} ──
  D1(8,1,4, 24,12);D1(8,1,6, 24,12);D1(8,1,12, 24,12);
  D0(8,1,4, 24,12);D0(8,1,6, 24,12);D0(8,1,12, 24,12);
  // W=4: NTW=48
  D1(4,1,4, 48,12);D1(4,1,6, 48,12);D1(4,1,12, 48,12);
  D0(4,1,4, 48,12);D0(4,1,6, 48,12);D0(4,1,12, 48,12);
  // ── (NTW=45/23, K128=4): m=256/m=32 n=2880 k=512. KU=4 only ──
  D1(4,1,4, 45,4);D0(4,1,4, 45,4);D1(8,1,4, 23,4);D0(8,1,4, 23,4);
  // ── (NTW=32/64, K128=4): m=32 n=4096 k=512 ──
  D1(8,1,4, 32,4);D0(8,1,4, 32,4);D1(4,1,4, 64,4);D0(4,1,4, 64,4);
  // ── (NTW=24/48, K128=8): secret coverage ──
  D1(8,1,4, 24,8);D1(8,1,8, 24,8);D0(8,1,4, 24,8);D0(8,1,8, 24,8);
  D1(4,1,4, 48,8);D1(4,1,8, 48,8);
  // ── (NTW=56/112, K128=8): more secret coverage ──
  D1(8,1,4, 56,8);D1(8,1,8, 56,8);D0(8,1,4, 56,8);
  #undef D0
  #undef D1
  return -1;
}

// Template-NTW launchers. grid = MT*NTW (same as runtime variants, but NTW
// is baked into the kernel binary). Dispatch requires exact NTW match.
template<int W,int MR,int KU,int KP,int NTW>
static void _ga3bt(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,int64_t NT){
  const int64_t K128=K>>7;const int64_t MT=(M+16*MR-1)/(16*MR);
  const int64_t gx=MT*(int64_t)NTW;
  const int64_t lds=(int64_t)MR*K128*1088;
  static bool _s=false;if(!_s){(void)hipFuncSetAttribute(
    (const void*)fgemm_alds3bt<W,MR,KU,KP,NTW>,hipFuncAttributeMaxDynamicSharedMemorySize,160*1024);_s=true;}
  fgemm_alds3bt<W,MR,KU,KP,NTW><<<dim3(gx),dim3(W*64),lds,0>>>(
    reinterpret_cast<const bf16*>(A.data_ptr()),
    Bsh.data_ptr<uint8_t>(),Bsc.data_ptr<uint8_t>(),
    reinterpret_cast<bf16*>(C.data_ptr()),(int)M,(int)N,(int)K,sn8,(int)NT);
}
template<int W,int MR,int KU,int NTW>
static void _ga3t(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,int64_t NT){
  const int64_t K128=K>>7;const int64_t MT=(M+16*MR-1)/(16*MR);
  const int64_t gx=MT*(int64_t)NTW;
  const int64_t lds=(int64_t)MR*K128*1088;
  static bool _s=false;if(!_s){(void)hipFuncSetAttribute(
    (const void*)fgemm_alds3t<W,MR,KU,NTW>,hipFuncAttributeMaxDynamicSharedMemorySize,160*1024);_s=true;}
  fgemm_alds3t<W,MR,KU,NTW><<<dim3(gx),dim3(W*64),lds,0>>>(
    reinterpret_cast<const bf16*>(A.data_ptr()),
    Bsh.data_ptr<uint8_t>(),Bsc.data_ptr<uint8_t>(),
    reinterpret_cast<bf16*>(C.data_ptr()),(int)M,(int)N,(int)K,sn8,(int)NT);
}

// Unified template dispatch. BSC: 0=no-preload(alds3t), 1=preload(alds3bt).
// NTW_runtime must MATCH a compiled template instance (bench+secret shapes).
// BENCH:  m=64/n=7168  W=8 -> NT=448 NTW=56. W=4 -> NTW=112.
//         m=256/n=3072 W=8 -> NT=192 NTW=24. W=4 -> NTW=48.
// SECRET: m=8/n=2112   W=8 -> NT=132 NTW=17. (K128=56 -> no alds, skip)
//         m=16/n=3072  W=8 -> NT=192 NTW=24. (K128=12)
//         m=64/n=3072  W=8 -> NT=192 NTW=24. (K128=12)
//         m=256/n=2880 W=8 -> NT=180 NTW=23(prime). W=4 -> NTW=45.
int64_t launch_alds3t(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,int64_t NT,
    int64_t W,int64_t MR,int64_t KU,int64_t BSC){
  int64_t K128=K>>7;int64_t KP=K128/2;
  int64_t NTW=(NT+W-1)/W;
  if((int64_t)MR*K128*1088>160*1024)return -2;
  if(K128%KU!=0)return -3;
  if(BSC&&(K128&1))return -4;
  #define D0(Ww,Rr,Uu,Nw) if(W==Ww&&MR==Rr&&KU==Uu&&BSC==0&&NTW==Nw){ \
      _ga3t<Ww,Rr,Uu,Nw>(A,Bsh,Bsc,C,M,N,K,sn8,NT);return 0;}
  #define D1(Ww,Rr,Uu,Kp,Nw) if(W==Ww&&MR==Rr&&KU==Uu&&BSC==1&&KP==Kp&&NTW==Nw){ \
      _ga3bt<Ww,Rr,Uu,Kp,Nw>(A,Bsh,Bsc,C,M,N,K,sn8,NT);return 0;}
  // ── m=64 n=7168 k=2048: NTW=56(W8)/112(W4) K128=16 KP=8 ──
  D0(8,1,4, 56);D0(8,1,8, 56);D1(8,1,4,8, 56);D1(8,1,8,8, 56);
  D0(4,1,4, 112);D0(4,1,8, 112);D1(4,1,4,8, 112);D1(4,1,8,8, 112);
  // ── m=256 n=3072 k=1536: NTW=24(W8)/48(W4) K128=12 KP=6 ──
  // ── SECRET m=64/m=16 n=3072 k=1536: SAME NTW=24/48 K128=12 KP=6 ──
  D1(8,1,4,6, 24);D1(8,1,6,6, 24);D0(8,1,4, 24);D0(8,1,6, 24);D0(8,1,12, 24);
  D1(4,1,4,6, 48);D1(4,1,6,6, 48);D1(4,1,12,6, 48);D0(4,1,4, 48);D0(4,1,6, 48);D0(4,1,12, 48);
  // ── SECRET m=256 n=2880 k=512: NTW=23(W8)/45(W4) K128=4 KP=2 ──
  //   NTW=23 prime -> W=4 only (NTW=45 = 9*5)
  D1(4,1,4,2, 45);D0(4,1,4, 45);D1(8,1,4,2, 23);D0(8,1,4, 23);
  // ── m=32 n=2880 k=512: NTW=23(W8)/45(W4). m=32 n=4096: NTW=32/64 K128=4 KP=2 ──
  D1(4,1,4,2, 64);D0(4,1,4, 64);D1(8,1,4,2, 32);D0(8,1,4, 32);
  // ── Generic K128=8 KP=4 (secret shape space) ──
  D1(8,1,4,4, 24);D1(8,1,8,4, 24);D1(4,1,4,4, 48);D1(4,1,8,4, 48);
  D1(8,1,4,4, 56);D1(8,1,8,4, 56);
  #undef D0
  #undef D1
  return -1;
}

int64_t launch_alds3b(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,
    int64_t NT,int64_t W,int64_t MR,int64_t KU){
  int64_t K128=K>>7;int64_t KP=K128/2;
  if((int64_t)MR*K128*1088>160*1024)return -2;
  if(K128%KU!=0)return -3;
  if(K128&1)return -4;
  #define D(Ww,Rr,Uu,Kp) if(W==Ww&&MR==Rr&&KU==Uu&&KP==Kp){ \
      _ga3b<Ww,Rr,Uu,Kp>(A,Bsh,Bsc,C,M,N,K,sn8,NT);return 0;}
  // Bench winners
  D(8,1,8,8);D(8,1,4,8);   // m=64 K128=16 -> KP=8
  D(8,1,4,6);D(8,1,6,6);D(8,2,4,6);D(8,2,6,6);  // m=256 K128=12 -> KP=6
  D(4,1,4,2);D(8,1,4,2);   // m=32 K128=4 -> KP=2
  // Secret shapes: broad KP coverage
  D(4,1,4,6);D(4,1,6,6);D(4,1,12,6);D(4,2,4,6);D(4,2,6,6);
  D(4,1,4,8);D(4,1,8,8);D(4,2,4,8);D(4,2,8,8);D(8,2,4,8);D(8,2,8,8);
  D(8,1,4,4);D(8,1,8,4);D(4,1,4,4);D(4,1,8,4);   // K128=8
  #undef D
  return -1;
}

int64_t launch_alds3(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,
    int64_t NT,int64_t W,int64_t MR,int64_t KU){
  int64_t K128=K>>7;
  if((int64_t)MR*K128*1088>160*1024)return -2;
  if(K128%KU!=0)return -3;
  #define D(Ww,Rr,Uu) if(W==Ww&&MR==Rr&&KU==Uu){ \
      _ga3<Ww,Rr,Uu>(A,Bsh,Bsc,C,M,N,K,sn8,NT);return 0;}
  D(8,1,4);D(8,1,6);D(8,1,8);D(8,1,12);
  D(4,1,4);D(4,1,6);D(4,1,8);D(4,1,12);
  D(8,2,4);D(8,2,6);D(4,2,4);D(4,2,6);
  #undef D
  return -1;
}

template<int W,int KU>
static void _galdsk(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor Cf,int64_t M,int64_t N,int64_t K,int64_t sn8,int64_t NT,int64_t SK){
  const int64_t K128=K>>7;const int64_t Kps128=(K128+SK-1)/SK;
  const int64_t NTW=(NT+W-1)/W;const int64_t gx=NTW*SK;
  const int64_t lds=Kps128*1088;
  static bool _s=false;if(!_s){(void)hipFuncSetAttribute(
    (const void*)fgemm_alds_sk<W,KU>,hipFuncAttributeMaxDynamicSharedMemorySize,160*1024);_s=true;}
  fgemm_alds_sk<W,KU><<<dim3(gx),dim3(W*64),lds,0>>>(
    reinterpret_cast<const bf16*>(A.data_ptr()),
    Bsh.data_ptr<uint8_t>(),Bsc.data_ptr<uint8_t>(),
    Cf.data_ptr<float>(),(int)M,(int)N,(int)K,sn8,(int)NT,(int)SK);
}
void go_cast(torch::Tensor Cf,torch::Tensor C,torch::Tensor Cfz,int64_t Ne){
  int64_t g=(Ne+255)/256;
  cast_f32_bf16_z<<<dim3(g),dim3(256),0,0>>>(
    Cf.data_ptr<float>(),reinterpret_cast<bf16*>(C.data_ptr()),
    Cfz.data_ptr<float>(),Ne);
}
int64_t launch_alds_sk(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor Cf,int64_t M,int64_t N,int64_t K,int64_t sn8,
    int64_t NT,int64_t SK,int64_t W,int64_t KU){
  int64_t K128=K>>7;int64_t Kps128=(K128+SK-1)/SK;
  if(Kps128*1088>160*1024)return -2;
  #define D(Ww,Uu) if(W==Ww&&KU==Uu){_galdsk<Ww,Uu>(A,Bsh,Bsc,Cf,M,N,K,sn8,NT,SK);return 0;}
  D(4,4);D(4,7);D(8,4);D(8,7);
  #undef D
  return -1;
}

template<int W,int NR>
static void _gfqn(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,int64_t MT,int64_t NT){
  int64_t NTG=(NT+NR-1)/NR;int64_t gx=MT*NTG,lds=(int64_t)W*NR*256*4;
  fgemm_fqn<W,NR><<<dim3(gx),dim3(W*64),lds,0>>>(
    reinterpret_cast<const bf16*>(A.data_ptr()),
    Bsh.data_ptr<uint8_t>(),Bsc.data_ptr<uint8_t>(),
    reinterpret_cast<bf16*>(C.data_ptr()),(int)M,(int)N,(int)K,sn8,(int)NT);
}
template<int W>
static void _gfq(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,int64_t MT,int64_t NT){
  int64_t gx=MT*NT,lds=(int64_t)W*256*4;
  fgemm_fq<W><<<dim3(gx),dim3(W*64),lds,0>>>(
    reinterpret_cast<const bf16*>(A.data_ptr()),
    Bsh.data_ptr<uint8_t>(),Bsc.data_ptr<uint8_t>(),
    reinterpret_cast<bf16*>(C.data_ptr()),(int)M,(int)N,(int)K,sn8,(int)NT);
}
int64_t launch_fqn(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,
    int64_t MT,int64_t NT,int64_t W,int64_t NR){
  #define D(Ww,Nn) if(W==Ww&&NR==Nn){_gfqn<Ww,Nn>(A,Bsh,Bsc,C,M,N,K,sn8,MT,NT);return 0;}
  D(4,1);D(4,2);D(2,2);D(8,1);D(8,2);
  #undef D
  return -1;
}
int64_t launch_fq(torch::Tensor A,torch::Tensor Bsh,torch::Tensor Bsc,
    torch::Tensor C,int64_t M,int64_t N,int64_t K,int64_t sn8,
    int64_t MT,int64_t NT,int64_t W){
  #define D(Ww) if(W==Ww){_gfq<Ww>(A,Bsh,Bsc,C,M,N,K,sn8,MT,NT);return 0;}
  D(4);D(8);
  #undef D
  return -1;
}

void probe(){
  hipFuncAttributes a;
  #define P(k,s) (void)hipFuncGetAttributes(&a,(const void*)k); \
    printf("[v28] %-44s VGPR=%3d spill=%zu\n",s,a.numRegs,a.localSizeBytes);
  // CRITICAL: full-template VGPR vs NTW-only template (v27).
  //   v27: a3t<8,1,4,N56>=55.  a3f must be <=55 (K-loop unroll doesn't add regs
  //   if liveranges stay per-KU-block; it SHOULD since acc[] is the only cross).
  //   RISK: #pragma unroll on K-loop could fuse bb[KU] across iters -> bloat.
  P((fgemm_alds3t<8,1,4,56>),           "a3t<8,1,4,N56>         v27 ref=55");
  P((fgemm_alds3f<8,1,4,1,0,56,16>),    "a3f<8,1,4,nb,N56,K16>  <- m=64 v28");
  P((fgemm_alds3f<8,1,8,1,0,56,16>),    "a3f<8,1,8,nb,N56,K16>  KU=8");
  P((fgemm_alds3f<8,1,16,1,0,56,16>),   "a3f<8,1,16,nb,N56,K16> KU=K128 1iter!");
  P((fgemm_alds3bt<8,1,4,6,24>),        "a3bt<8,1,4,6,N24>      v27 ref=59");
  P((fgemm_alds3f<8,1,4,6,1,24,12>),    "a3f<8,1,4,b6,N24,K12>  <- m=256 v28");
  P((fgemm_alds3f<8,1,6,6,1,24,12>),    "a3f<8,1,6,b6,N24,K12>  KU=6");
  P((fgemm_alds3f<8,1,12,6,1,24,12>),   "a3f<8,1,12,b6,N24,K12> KU=K128 1iter!");
  P((fgemm_alds3f<4,1,4,2,1,45,4>),     "a3f<4,1,4,b2,N45,K4>   secret m256");
  P((fgemm_alds3f<4,1,12,6,1,48,12>),   "a3f<4,1,12,b6,N48,K12> secret m64 1iter");
  P((fgemm_alds_sk<8,7>),               "alds_sk<8,7>");
  P((fgemm_fqn<4,2>),                   "fqn<4,2>");
  #undef P
}
"""

_CPP = r"""
#include <torch/extension.h>
int64_t launch_alds3f(torch::Tensor,torch::Tensor,torch::Tensor,torch::Tensor,
    int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t);
int64_t launch_alds3t(torch::Tensor,torch::Tensor,torch::Tensor,torch::Tensor,
    int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t);
int64_t launch_alds3b(torch::Tensor,torch::Tensor,torch::Tensor,torch::Tensor,
    int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t);
int64_t launch_alds3(torch::Tensor,torch::Tensor,torch::Tensor,torch::Tensor,
    int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t);
int64_t launch_alds_sk(torch::Tensor,torch::Tensor,torch::Tensor,torch::Tensor,
    int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t);
void go_cast(torch::Tensor,torch::Tensor,torch::Tensor,int64_t);
int64_t launch_fqn(torch::Tensor,torch::Tensor,torch::Tensor,torch::Tensor,
    int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t);
int64_t launch_fq(torch::Tensor,torch::Tensor,torch::Tensor,torch::Tensor,
    int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t);
void probe();
"""

_hip = None
try:
    from torch.utils.cpp_extension import load_inline
    _t0 = time.time()
    _hip = load_inline(name="v28_full", cpp_sources=_CPP,
        cuda_sources=_HIP_SRC,
        functions=["launch_alds3f","launch_alds3t","launch_alds3b","launch_alds3",
                   "launch_alds_sk","go_cast","launch_fqn","launch_fq","probe"],
        with_cuda=True,
        extra_cuda_cflags=["-O3","--offload-arch=gfx950","-ffast-math",
                           "-mllvm","-amdgpu-early-inline-all=true",
                           "-mllvm","-amdgpu-function-calls=false",
                           "-munsafe-fp-atomics"],
        verbose=False)
    _L(f"[v28] HIP compiled {time.time()-_t0:.1f}s"); _hip.probe()
except Exception as ex:
    import traceback
    _L(f"[v28] HIP FAIL: {type(ex).__name__}: {str(ex)[:2000]}")
    for ln in traceback.format_exc().splitlines()[-25:]:
        _L(f"   {ln[:200]}")


@triton.jit
def _sh_row(r,sn8):return (r//32)*(sn8*256)+(r%16)*4+(r//16)%2
@triton.jit
def _sh_col(c):return (c//8)*256+(c%4)*64+(c//4)%2*2
def _make_gemm_ref():
    from aiter.ops.triton._triton_kernels.quant.quant import _mxfp4_quant_op
    @triton.jit
    def _k(A,Bq,Bsc,C,M,N,K,sAm,sBn,sn8,
           BM:tl.constexpr,BN:tl.constexpr,BK:tl.constexpr,EN:tl.constexpr):
        pid=tl.program_id(0);nn=tl.cdiv(N,BN);pm=pid//nn;pn=pid%nn
        om=pm*BM+tl.arange(0,BM);on=pn*BN+tl.arange(0,BN)
        o64=on.to(tl.int64);mm=om<M;mn=on<N
        rk=tl.arange(0,BK);r2=tl.arange(0,BK//2);r32=tl.arange(0,BK//32)
        bp=Bq+o64[:,None]*sBn+r2[None,:]
        br=_sh_row(o64,sn8);acc=tl.zeros((BM,BN),dtype=tl.float32)
        ap=A+om[:,None].to(tl.int64)*sAm+rk[None,:]
        for kk in tl.range(0,K,BK):
            ab=tl.load(ap,mask=mm[:,None],other=0.)
            af,asc=_mxfp4_quant_op(ab.to(tl.float32),BK,BM,32);ap+=BK
            if EN:
                bf=tl.load(bp);bs=tl.load(Bsc+br[:,None]+_sh_col(kk//32+r32)[None,:])
            else:
                bf=tl.load(bp,mask=mn[:,None],other=0)
                bs=tl.load(Bsc+br[:,None]+_sh_col(kk//32+r32)[None,:],mask=mn[:,None],other=0)
            acc=tl.dot_scaled(af,asc,"e2m1",tl.trans(bf),bs,"e2m1",acc);bp+=BK//2
        cm=mm[:,None]&mn[None,:]
        tl.store(C+om[:,None].to(tl.int64)*N+on[None,:],acc.to(tl.bfloat16),mask=cm)
    return _k
_gemm_ref=_make_gemm_ref()


# ═══════ HARDCODED TABLE (v28: full-template NTW+K128) ═══════
# "a3f" = full-template alds3f. NTW & K128 both baked in. BSC via if constexpr.
# K-loop AND prologue loop compile-time-unrolled.
# Fallback chain: a3f -> a3t (NTW-only) -> a3b/a3 (runtime).
_HARD = {
    (4,   2880, 512 ): ("fqn", 4, 2),
    (16,  2112, 7168): ("aldsk", 8, 8, 7),
    (32,  4096, 512 ): ("fqn", 4, 2),
    (32,  2880, 512 ): ("fqn", 4, 2),
    (64,  7168, 2048): ("a3f", 8, 1, 4, 0),  # full tpl NTW=56 K128=16. m=64 no-bsc (VGPR safety)
    (256, 3072, 1536): ("a3f", 8, 1, 4, 1),  # full tpl NTW=24 K128=12 bsc KP=6
}

def _pick_cands(m,n,k):
    NT=-(-n//16);MT16=-(-m//16);K128=k//128
    h=_HARD.get((m,n,k))
    if h: return [h]
    out=[]
    divs=[d for d in(4,6,8,12,16) if K128%d==0 and K128*1088<=160*1024]
    bsc_ok=(K128%2==0)and((K128//2)in(2,4,6,8))
    # a3f (full tpl) > a3t (NTW-only) > a3b/a3 (runtime). Cascading fallback.
    if MT16>=2 and divs:
        for W in(8,4):
            for KU in divs:
                if bsc_ok:out.append(("a3f",W,1,KU,1))   # full-tpl bsc
                out.append(("a3f",W,1,KU,0))             # full-tpl no-bsc
                if bsc_ok:out.append(("a3t",W,1,KU,1))   # NTW-tpl bsc (fallback)
                out.append(("a3t",W,1,KU,0))
                if bsc_ok:out.append(("a3b",W,1,KU))     # runtime fallback
                out.append(("a3",W,1,KU))
        if MT16>=8:
            for KU in[d for d in(4,6)if K128%d==0 and 2*K128*1088<=160*1024]:
                if bsc_ok:out.append(("a3b",8,2,KU))
                out.append(("a3",8,2,KU))
    # aldsk for small-M large-K
    if MT16<=2 and K128>=16:
        for SK in(8,14,7):
            if SK>K128:continue
            out.append(("aldsk",8,SK,7))
            out.append(("aldsk",4,SK,4))
    # safety
    out.append(("fqn",4,2))
    out.append(("fq",min(8,K128)))
    return out


_L2=torch.empty(512*1024*1024,dtype=torch.int8,device="cuda")
def _tcold(fn,n=7):
    for _ in range(2):fn()
    torch.cuda.synchronize()
    evs=[(torch.cuda.Event(True),torch.cuda.Event(True))for _ in range(n)]
    for e0,e1 in evs:_L2.zero_();e0.record();fn();e1.record()
    torch.cuda.synchronize()
    ts=sorted(e0.elapsed_time(e1)for e0,e1 in evs)
    return sum(ts[1:-1])*1000/(n-2)


_ST={}
def _build(data):
    A,B,Bq_,Bsh_,Bsc_=data
    m,k=A.shape;n=B.shape[0]
    sn=Bsc_.shape[1];sn8=sn//8;dev=A.device
    NT=-(-n//16);MT16=-(-m//16);K128=k//128;mn=m*n
    Bq=Bq_.view(torch.uint8);Bsh=Bsh_.view(torch.uint8);Bsc=Bsc_.view(torch.uint8)
    C=torch.empty((m,n),dtype=torch.bfloat16,device=dev)
    Cf=torch.zeros((m,n),dtype=torch.float32,device=dev)
    Cf2=torch.zeros((m,n),dtype=torch.float32,device=dev)
    _pp=[Cf,Cf2]

    def _run(cfg,_A,_Bq,_Bsh,_Bsc):
        kind=cfg[0]
        if kind=="a3f":
            _,Wv,MR,KU,BSC=cfg
            rc=_hip.launch_alds3f(_A,_Bsh,_Bsc,C,m,n,k,sn8,NT,Wv,MR,KU,BSC)
            if rc!=0:raise RuntimeError(f"a3f rc={rc}")
            return C
        if kind=="a3t":
            _,Wv,MR,KU,BSC=cfg
            rc=_hip.launch_alds3t(_A,_Bsh,_Bsc,C,m,n,k,sn8,NT,Wv,MR,KU,BSC)
            if rc!=0:raise RuntimeError(f"a3t rc={rc}")
            return C
        if kind=="a3b":
            _,Wv,MR,KU=cfg
            rc=_hip.launch_alds3b(_A,_Bsh,_Bsc,C,m,n,k,sn8,NT,Wv,MR,KU)
            if rc!=0:raise RuntimeError(f"a3b rc={rc}")
            return C
        if kind=="a3":
            _,Wv,MR,KU=cfg
            rc=_hip.launch_alds3(_A,_Bsh,_Bsc,C,m,n,k,sn8,NT,Wv,MR,KU)
            if rc!=0:raise RuntimeError(f"a3 rc={rc}")
            return C
        if kind=="aldsk":
            _,Wv,SK,KU=cfg
            rc=_hip.launch_alds_sk(_A,_Bsh,_Bsc,_pp[0],m,n,k,sn8,NT,SK,Wv,KU)
            if rc!=0:raise RuntimeError(f"aldsk rc={rc}")
            _hip.go_cast(_pp[0],C,_pp[1],mn)
            _pp[0],_pp[1]=_pp[1],_pp[0]
            return C
        if kind=="fqn":
            _,Wv,NR=cfg
            rc=_hip.launch_fqn(_A,_Bsh,_Bsc,C,m,n,k,sn8,MT16,NT,Wv,NR)
            if rc!=0:raise RuntimeError(f"fqn rc={rc}")
            return C
        if kind=="fq":
            _,Wv=cfg
            rc=_hip.launch_fq(_A,_Bsh,_Bsc,C,m,n,k,sn8,MT16,NT,Wv)
            if rc!=0:raise RuntimeError(f"fq rc={rc}")
            return C
        raise RuntimeError(f"?{cfg}")

    def _ref_tri(_A,_Bq,_Bsc):
        Cref=torch.empty_like(C)
        BK=min(512,k);gx=MT16*triton.cdiv(n,32)
        _gemm_ref[(gx,)](_A,_Bq,_Bsc,Cref,m,n,k,k,k//2,sn8,
            BM=16,BN=32,BK=BK,EN=(n%32==0),
            num_warps=8,num_stages=2,matrix_instr_nonkdim=16)
        return Cref

    rf=_ref_tri(A,Bq,Bsc).float();mag=rf.abs().mean().item()+1e-9

    cands=_pick_cands(m,n,k)
    is_hard=(m,n,k) in _HARD
    _L(f"\n[v28 m={m} n={n} k={k}] {'HARD' if is_hard else 'PICK'}: {len(cands)}c NTW8={-(-NT//8)} K128={K128}")

    if is_hard:
        cfg=cands[0]
        try:
            C.fill_(float('nan'))
            o=_run(cfg,A,Bq,Bsh,Bsc);torch.cuda.synchronize()
            err=((o.float()-rf).abs().mean()/mag).item()
            if err<5e-3:
                A2=torch.randn_like(A)
                rf2=_ref_tri(A2,Bq,Bsc).float()
                C.fill_(float('nan'))
                o2=_run(cfg,A2,Bq,Bsh,Bsc);torch.cuda.synchronize()
                e2=((o2.float()-rf2).abs().mean()/(rf2.abs().mean()+1e-9)).item()
                if e2<5e-3:
                    _L(f"  {cfg} chk={err:.3%},{e2:.3%} OK")
                    return {"cfg":cfg,"run":_run,"C":C}
                _L(f"  RECHECK FAIL {cfg} e2={e2:.2%}")
            else:
                _L(f"  ERR {cfg} {err:.2%}")
        except Exception as e:
            _L(f"  EXC {cfg} {type(e).__name__}:{e}")

    best=None;bt=1e18;log=[];t0=time.time();nmiss=0
    for cfg in cands:
        if time.time()-t0>40:break
        try:
            C.fill_(float('nan'))
            o=_run(cfg,A,Bq,Bsh,Bsc);torch.cuda.synchronize()
            err=((o.float()-rf).abs().mean()/mag).item()
            if not(err<5e-3):_L(f"  {cfg}:ERR{err:.2%}");continue
            t=_tcold(lambda c=cfg:_run(c,A,Bq,Bsh,Bsc))
            log.append((cfg,t))
            if t<bt:bt,best=t,cfg;_L(f"  {cfg}:{t:.2f}us*")
        except RuntimeError as e:
            # a3t rc=-1 = NTW template not compiled (secret shapes). Silent skip.
            if "rc=-1" in str(e):nmiss+=1
            else:_L(f"  {cfg}:EXC{str(e)[:100]}")
            torch.cuda.synchronize()
        except Exception as e:
            _L(f"  {cfg}:EXC{type(e).__name__}:{str(e)[:100]}")
            torch.cuda.synchronize()
    if nmiss:_L(f"  (tpl miss={nmiss})")
    if best is None:
        _L("  ->fq");best=("fq",min(8,K128))
    try:
        A2=torch.randn_like(A)
        rf2=_ref_tri(A2,Bq,Bsc).float()
        C.fill_(float('nan'))
        o2=_run(best,A2,Bq,Bsh,Bsc);torch.cuda.synchronize()
        e2=((o2.float()-rf2).abs().mean()/(rf2.abs().mean()+1e-9)).item()
        if not(e2<5e-3):_L(f"  RECHECK FAIL {best} {e2:.2%}");best=("fq",min(8,K128))
    except Exception as e:_L(f"  recheck exc {e}")
    log.sort(key=lambda x:x[1])
    for c,t in log[:10]:_L(f"  top{c}:{t:.2f}")
    _L(f"  ->best={best}@{bt:.2f}us")
    return {"cfg":best,"run":_run,"C":C}


def custom_kernel(data):
    A=data[0];m,k=A.shape;n=data[2].shape[0]
    S=_ST.get((m,n,k))
    if S is None:
        S=_build(data);_ST[(m,n,k)]=S
    return S["run"](S["cfg"],A,
        data[2].view(torch.uint8),
        data[3].view(torch.uint8),data[4].view(torch.uint8))