submission 415875

from __future__ import annotations

import os
from typing import List

import torch
from torch.utils.cpp_extension import load_inline









































_FORCE_NO_GROUPED = False
_FORCE_BN64 = False
_FORCE_RAW_SF = False

_EXT_READY = False
_OPS_READY = False

_GEMM = None
_GEMM_GROUPED = None




_DIAG = os.environ.get("NVFP4_GGEMM_DIAG", "0") == "1"
_DIAG_ONCE = True


def _diag_write(msg: str) -> None:
    if not _DIAG:
        return
    try:
        os.write(2, (msg + "\n").encode("utf-8"))
    except Exception:
        pass


def _load_ext() -> None:
    global _EXT_READY, _OPS_READY, _GEMM, _GEMM_GROUPED
    if _EXT_READY:
        return

    cuda_src = r"""
#include <cuda.h>
#include <cudaTypedefs.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>

#include <torch/extension.h>
#include <torch/library.h>
#include <ATen/ATen.h>

#include <cstdint>
#include <vector>
#include <array>

// 默认关闭检查，极致压缩 host 热路径分支
#ifndef NVFP4_GGEMM_CHECK
#define NVFP4_GGEMM_CHECK 0
#endif

#if NVFP4_GGEMM_CHECK
#define GG_CHECK(x, msg) TORCH_CHECK((x), msg)
#else
#define GG_CHECK(x, msg) ((void)0)
#endif

constexpr int WARP_SIZE = 32;
constexpr int MMA_K = 64;

// cache-policy 仅为性能 hint：恢复常量立即数，避免 createpolicy 指令的 CTA 固定开销
constexpr uint64_t EVICT_FIRST = 0x12F0000000000000ULL;
constexpr uint64_t EVICT_LAST  = 0x14F0000000000000ULL;

__device__ __forceinline__ constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3FFFFULL) >> 4ULL; }

__device__ __forceinline__ uint32_t elect_sync() {
  uint32_t pred = 0;
  asm volatile(
    "{\n\t"
    ".reg .pred %%px;\n\t"
    "elect.sync _|%%px, %1;\n\t"
    "@%%px mov.s32 %0, 1;\n\t"
    "}"
    : "+r"(pred)
    : "r"(0xFFFFFFFF)
  );
  return pred;
}

__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) {
  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));
}

__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) {
  uint32_t ticks = 0x989680;
  asm volatile(
    "{\n\t"
    ".reg .pred P1;\n\t"
    "LAB_WAIT:\n\t"
    "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"
    "@P1 bra.uni DONE;\n\t"
    "bra.uni LAB_WAIT;\n\t"
    "DONE:\n\t"
    "}"
    :: "r"(mbar_addr), "r"(phase), "r"(ticks)
  );
}

__device__ __forceinline__ void tma_gmem2smem(int dst, const void *src, int size, int mbar_addr, uint64_t cache_policy) {
  asm volatile(
    "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint "
    "[%0], [%1], %2, [%3], %4;"
    :: "r"(dst), "l"(src), "r"(size), "r"(mbar_addr), "l"(cache_policy)
  );
}

__device__ __forceinline__ void tma_3d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr, uint64_t cache_policy) {
  asm volatile(
    "cp.async.bulk.tensor.3d.shared::cta.global.mbarrier::complete_tx::bytes.cta_group::1.L2::cache_hint "
    "[%0], [%1, {%2, %3, %4}], [%5], %6;"
    :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "l"(cache_policy)
    : "memory"
  );
}

__device__ __forceinline__ void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
  asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));
}

__device__ __forceinline__ void tcgen05_mma_nvfp4(
  uint64_t a_desc,
  uint64_t b_desc,
  uint32_t i_desc,
  int scale_A_tmem,
  int scale_B_tmem,
  int enable_input_d
) {
  const int d_tmem = 0;
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %6, 0;\n\t"
    "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\n\t"
    "}"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d)
  );
}

struct SHAPE {
  static constexpr char _16x256b[] = ".16x256b";
};
struct NUM {
  static constexpr char x8[]  = ".x8";
  static constexpr char x16[] = ".x16";
};

template <const char *SHAPE_V, const char *NUM_V>
__device__ __forceinline__ void tcgen05_ld_32regs(float *tmp, int row, int col) {
  asm volatile(
    "tcgen05.ld.sync.aligned%33%34.b32 "
    "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
    "  %8,  %9, %10, %11, %12, %13, %14, %15, "
    " %16, %17, %18, %19, %20, %21, %22, %23, "
    " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];"
    : "=f"(tmp[ 0]), "=f"(tmp[ 1]), "=f"(tmp[ 2]), "=f"(tmp[ 3]), "=f"(tmp[ 4]), "=f"(tmp[ 5]), "=f"(tmp[ 6]), "=f"(tmp[ 7]),
      "=f"(tmp[ 8]), "=f"(tmp[ 9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15]),
      "=f"(tmp[16]), "=f"(tmp[17]), "=f"(tmp[18]), "=f"(tmp[19]), "=f"(tmp[20]), "=f"(tmp[21]), "=f"(tmp[22]), "=f"(tmp[23]),
      "=f"(tmp[24]), "=f"(tmp[25]), "=f"(tmp[26]), "=f"(tmp[27]), "=f"(tmp[28]), "=f"(tmp[29]), "=f"(tmp[30]), "=f"(tmp[31])
    : "r"((row << 16) | col), "C"(SHAPE_V), "C"(NUM_V));
}

template <const char *SHAPE_V, const char *NUM_V>
__device__ __forceinline__ void tcgen05_ld_64regs(float *tmp, int row, int col) {
  asm volatile(
    "tcgen05.ld.sync.aligned%65%66.b32 "
    "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
    "  %8,  %9, %10, %11, %12, %13, %14, %15, "
    " %16, %17, %18, %19, %20, %21, %22, %23, "
    " %24, %25, %26, %27, %28, %29, %30, %31, "
    " %32, %33, %34, %35, %36, %37, %38, %39, "
    " %40, %41, %42, %43, %44, %45, %46, %47, "
    " %48, %49, %50, %51, %52, %53, %54, %55, "
    " %56, %57, %58, %59, %60, %61, %62, %63}, [%64];"
    : "=f"(tmp[ 0]), "=f"(tmp[ 1]), "=f"(tmp[ 2]), "=f"(tmp[ 3]), "=f"(tmp[ 4]), "=f"(tmp[ 5]), "=f"(tmp[ 6]), "=f"(tmp[ 7]),
      "=f"(tmp[ 8]), "=f"(tmp[ 9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15]),
      "=f"(tmp[16]), "=f"(tmp[17]), "=f"(tmp[18]), "=f"(tmp[19]), "=f"(tmp[20]), "=f"(tmp[21]), "=f"(tmp[22]), "=f"(tmp[23]),
      "=f"(tmp[24]), "=f"(tmp[25]), "=f"(tmp[26]), "=f"(tmp[27]), "=f"(tmp[28]), "=f"(tmp[29]), "=f"(tmp[30]), "=f"(tmp[31]),
      "=f"(tmp[32]), "=f"(tmp[33]), "=f"(tmp[34]), "=f"(tmp[35]), "=f"(tmp[36]), "=f"(tmp[37]), "=f"(tmp[38]), "=f"(tmp[39]),
      "=f"(tmp[40]), "=f"(tmp[41]), "=f"(tmp[42]), "=f"(tmp[43]), "=f"(tmp[44]), "=f"(tmp[45]), "=f"(tmp[46]), "=f"(tmp[47]),
      "=f"(tmp[48]), "=f"(tmp[49]), "=f"(tmp[50]), "=f"(tmp[51]), "=f"(tmp[52]), "=f"(tmp[53]), "=f"(tmp[54]), "=f"(tmp[55]),
      "=f"(tmp[56]), "=f"(tmp[57]), "=f"(tmp[58]), "=f"(tmp[59]), "=f"(tmp[60]), "=f"(tmp[61]), "=f"(tmp[62]), "=f"(tmp[63])
    : "r"((row << 16) | col), "C"(SHAPE_V), "C"(NUM_V));
}

__device__ __forceinline__ void tcgen05_ld_16x256bx8(float *tmp, int row, int col) {
  tcgen05_ld_32regs<SHAPE::_16x256b, NUM::x8>(tmp, row, col);
}
__device__ __forceinline__ void tcgen05_ld_16x256bx16(float *tmp, int row, int col) {
  tcgen05_ld_64regs<SHAPE::_16x256b, NUM::x16>(tmp, row, col);
}

static __forceinline__ void check_cu(CUresult err) {
  if (err == CUDA_SUCCESS) return;
  const char *msg = "unknown";
  cuGetErrorString(err, &msg);
  TORCH_CHECK(false, msg);
}

struct TmapKey {
  uint64_t ptr;
  uint64_t global_height;
  uint64_t global_width;
  uint32_t shared_height;
  uint32_t shared_width;
  uint32_t swizzle;
  uint32_t l2_promotion;
  int32_t dev;
};

static __forceinline__ bool tmap_key_eq(const TmapKey &a, const TmapKey &b) {
  return a.ptr == b.ptr
      && a.global_height == b.global_height
      && a.global_width == b.global_width
      && a.shared_height == b.shared_height
      && a.shared_width == b.shared_width
      && a.swizzle == b.swizzle
      && a.l2_promotion == b.l2_promotion
      && a.dev == b.dev;
}

template <int CAP>
struct TmapCache {
  std::array<TmapKey, CAP> keys;
  std::array<CUtensorMap, CAP> vals;
  std::array<uint8_t, CAP> used;
  int head;

  TmapCache() : used{}, head(0) {}

  bool lookup(const TmapKey &k, CUtensorMap *out) {
    #pragma unroll
    for (int i = 0; i < CAP; i++) {
      if (used[(size_t)i] && tmap_key_eq(keys[(size_t)i], k)) {
        *out = vals[(size_t)i];
        return true;
      }
    }
    return false;
  }

  void insert(const TmapKey &k, const CUtensorMap &v) {
    keys[(size_t)head] = k;
    vals[(size_t)head] = v;
    used[(size_t)head] = 1;
    head++;
    if (head >= CAP) head = 0;
  }
};

static TmapCache<64> g_tmap_cache;

static __forceinline__ void init_AB_tmap(
  CUtensorMap *tmap,
  const char *ptr,
  uint64_t global_height,
  uint64_t global_width,
  uint32_t shared_height,
  uint32_t shared_width,
  CUtensorMapSwizzle swizzle,
  CUtensorMapL2promotion l2_promotion
) {
  int dev = 0;
  cudaGetDevice(&dev);
  TmapKey key;
  key.ptr = (uint64_t)ptr;
  key.global_height = global_height;
  key.global_width = global_width;
  key.shared_height = shared_height;
  key.shared_width = shared_width;
  key.swizzle = (uint32_t)swizzle;
  key.l2_promotion = (uint32_t)l2_promotion;
  key.dev = (int32_t)dev;

  if (g_tmap_cache.lookup(key, tmap)) return;

  constexpr uint32_t rank = 3;
  uint64_t globalDim[rank]       = {256, global_height, global_width / 256};
  uint64_t globalStrides[rank-1] = {global_width / 2, 128};
  uint32_t boxDim[rank]          = {256, shared_height, shared_width / 256};
  uint32_t elementStrides[rank]  = {1, 1, 1};

  auto err = cuTensorMapEncodeTiled(
    tmap,
    CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B,
    rank,
    (void *)ptr,
    globalDim,
    globalStrides,
    boxDim,
    elementStrides,
    CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
    swizzle,
    l2_promotion,
    // 重要说明：
    // - 这里不依赖“越界区域数值被 0 填充”的语义（不同 datatype/驱动下可能不保证）。
    // - 正确性只依赖：越界对应的行最终不会写回到 C（epilogue 有 row<M 的写保护）。
    CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
  );
  check_cu(err);
  g_tmap_cache.insert(key, *tmap);
}

// grouped 元数据：一次性搬到 device，kernel 只读 descs[gid]
struct __align__(16) GroupDesc {
  CUtensorMap A_tmap;
  CUtensorMap B_tmap;
  uint64_t SFA_ptr;
  uint64_t SFB_ptr;
  uint64_t C_ptr;
  int M;
  int N;
  int K;
  uint8_t swap;
  uint8_t _pad0;
  uint16_t _pad1;
};

// 1D 紧凑 block 列表：避免 3D grid 的无效 block 调度
struct __align__(16) BlockDesc {
  int gid;
  int bid_m;
  int bid_n;
  int _pad;
};

__device__ __forceinline__ int active_threads_128(int M, int off_m) {
  int rem = M - off_m;
  if (rem <= 0) return 0;
  int thr = (rem + 31) & ~31;
  if (thr > 128) thr = 128;
  return thr;
}

// non-grouped 内核已不再使用：gemm() 复用 grouped 实现
// 直接禁用可减少 NVCC 编译与加载压力，降低评测卡 processing 概率
#if 0
template <int BLOCK_N, int NUM_STAGES>
__global__ __launch_bounds__(128 + 2 * WARP_SIZE, (BLOCK_N == 128 && (NUM_STAGES == 2 || NUM_STAGES == 3)) ? 2 : 1)
void kernel(
  const __grid_constant__ CUtensorMap A_tmap,
  const __grid_constant__ CUtensorMap B_tmap,
  const char *SFA_ptr,
  const char *SFB_ptr,
  half *C_ptr,
  int M, int N, int K
) {
  constexpr int BLOCK_M = 128;
  constexpr int BLOCK_K = 256;

  const int tid = (int)threadIdx.x;
  const int bid_n = (int)blockIdx.x;
  const int bid_m = (int)blockIdx.y;

  const int lane_id = tid & (WARP_SIZE - 1);
  const int warp_id = tid >> 5;

  const int off_m = bid_m * BLOCK_M;
  const int off_n = bid_n * BLOCK_N;

  constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE + 2;

  extern __shared__ __align__(1024) char smem_ptr[];
  const int smem = (int)__cvta_generic_to_shared(smem_ptr);
  constexpr int A_size   = BLOCK_M * BLOCK_K / 2;
  constexpr int B_size   = BLOCK_N * BLOCK_K / 2;
  constexpr int SFA_size = 128 * BLOCK_K / 16;
  constexpr int SFB_size = 128 * BLOCK_K / 16;
  constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;

  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ int64_t mbars[NUM_STAGES * 2 + 1];
  const int tma_mbar_addr = (int)__cvta_generic_to_shared(mbars);
  const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
  const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;

  constexpr int SFA_tmem = BLOCK_N;
  constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);

  if (warp_id == 0 && elect_sync()) {
    #pragma unroll
    for (int i = 0; i < NUM_STAGES * 2 + 1; i++) mbarrier_init(tma_mbar_addr + i * 8, 1);
    asm volatile("fence.mbarrier_init.release.cluster;");
  } else if (warp_id == 1) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 2));
  }
  __syncthreads();

  const int num_iters = K / BLOCK_K;

  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    const int grid_m = (M + (BLOCK_M - 1)) >> 7;
    const int grid_n = (N + (BLOCK_N - 1)) / BLOCK_N;
    // 回滚为可解释的阈值：当 N 的网格不小于 M 时，优先保留 A，B 走更激进的驱逐
    const bool swap = (grid_n >= grid_m);
    const uint64_t cache_A = swap ? EVICT_LAST : EVICT_FIRST;
    const uint64_t cache_B = swap ? EVICT_FIRST : EVICT_LAST;

    auto issue_tma = [&](int iter_k, int stage_id) {
      const int mbar_addr = tma_mbar_addr + stage_id * 8;
      const int A_smem = smem + stage_id * STAGE_SIZE;
      const int B_smem = A_smem + A_size;
      const int SFA_smem = B_smem + B_size;
      const int SFB_smem = SFA_smem + SFA_size;

      const int off_k = iter_k * BLOCK_K;
      tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A);
      tma_3d_gmem2smem(B_smem, &B_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);

      const int rest_k = K / 16 / 4;
      const char *SFA_src = SFA_ptr + ((off_m / 128) * rest_k + off_k / (16 * 4)) * 512;
      const char *SFB_src = SFB_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
      tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, cache_A);
      tma_gmem2smem(SFB_smem, SFB_src, SFB_size, mbar_addr, cache_B);

      asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
                   :: "r"(mbar_addr), "r"(STAGE_SIZE) : "memory");
    };

    const int init_stage = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;
    for (int iter_k = 0; iter_k < init_stage; iter_k++) issue_tma(iter_k, iter_k);

    for (int iter_k = NUM_STAGES; iter_k < num_iters; iter_k++) {
      const int stage_id = iter_k % NUM_STAGES;
      const int mma_phase = (iter_k / NUM_STAGES - 1) & 1;
      mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);
      issue_tma(iter_k, stage_id);
    }
  } else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
    constexpr int MMA_N = BLOCK_N;
    constexpr int MMA_M = 128;
    constexpr uint32_t i_desc = (1U << 7U)
                              | (1U << 10U)
                              | ((uint32_t)MMA_N >> 3U << 17U)
                              | ((uint32_t)MMA_M >> 7U << 27U);

    auto make_desc_AB = [] __device__ (int addr) -> uint64_t {
      const int SBO = 8 * 128;
      return desc_encode((uint64_t)addr) | (desc_encode((uint64_t)SBO) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
    };
    auto make_desc_SF = [] __device__ (int addr) -> uint64_t {
      const int SBO = 8 * 16;
      return desc_encode((uint64_t)addr) | (desc_encode((uint64_t)SBO) << 32ULL) | (1ULL << 46ULL);
    };

    for (int iter_k = 0; iter_k < num_iters; iter_k++) {
      const int stage_id = iter_k % NUM_STAGES;
      const int tma_phase = (iter_k / NUM_STAGES) & 1;
      mbarrier_wait(tma_mbar_addr + stage_id * 8, tma_phase);

      const int A_smem = smem + stage_id * STAGE_SIZE;
      const int B_smem = A_smem + A_size;
      const int SFA_smem = B_smem + B_size;
      const int SFB_smem = SFA_smem + SFA_size;

      const uint64_t SFA_desc = make_desc_SF(0) + ((uint64_t)SFA_smem >> 4ULL);
      const uint64_t SFB_desc = make_desc_SF(0) + ((uint64_t)SFB_smem >> 4ULL);

      #pragma unroll
      for (int k = 0; k < BLOCK_K / MMA_K; k++) {
        uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
        uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
        tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
        tcgen05_cp_nvfp4(SFB_tmem + k * 4, sfb_desc);
      }

      #pragma unroll
      for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {
        #pragma unroll
        for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
          uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
          uint64_t b_desc = make_desc_AB(B_smem + k1 * BLOCK_N * 128 + k2 * 32);

          const int k_sf = k1 * 4 + k2;
          const int scale_A_tmem = SFA_tmem + k_sf * 4;
          const int scale_B_tmem = SFB_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
          const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
          tcgen05_mma_nvfp4(a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d);
        }
      }

      asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                   :: "r"(mma_mbar_addr + stage_id * 8) : "memory");
    }

    asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                 :: "r"(mainloop_mbar_addr) : "memory");
  } else if (tid < BLOCK_M) {
    const int thr = active_threads_128(M, off_m);
    if (tid >= thr) return;

    mbarrier_wait(mainloop_mbar_addr, 0);
    asm volatile("tcgen05.fence::after_thread_sync;");
    // 题面约束 N 整除 BLOCK_N；且 B 的 TMA 读阶段本就依赖该条件
    // 因此移除 full_n 慢路径以减少分支与指令数

    #pragma unroll
    for (int m = 0; m < 2; m++) {
      const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);
      const int row2 = row + 8;

      if constexpr (BLOCK_N == 128) {
        float tmp[32];
        #pragma unroll
        for (int seg = 0; seg < 2; seg++) {
          const int col_base = off_n + seg * 64 + ((lane_id & 3) << 1);
          tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, seg * 64);
          asm volatile("tcgen05.wait::ld.sync.aligned;");
          if (row < M) {
            half *out0 = C_ptr + row * N + col_base;
            #pragma unroll
            for (int i = 0; i < 8; i++) {
              reinterpret_cast<half2 *>(out0 + i * 8)[0] =
                  __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});
            }
          }
          if (row2 < M) {
            half *out1 = C_ptr + row2 * N + col_base;
            #pragma unroll
            for (int i = 0; i < 8; i++) {
              reinterpret_cast<half2 *>(out1 + i * 8)[0] =
                  __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});
            }
          }
        }
      } else {
        float tmp[32];
        const int col_base = off_n + ((lane_id & 3) << 1);
        tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, 0);
        asm volatile("tcgen05.wait::ld.sync.aligned;");
        if (row < M) {
          half *out0 = C_ptr + row * N + col_base;
          #pragma unroll
          for (int i = 0; i < 8; i++) {
            reinterpret_cast<half2 *>(out0 + i * 8)[0] =
                __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});
          }
        }
        if (row2 < M) {
          half *out1 = C_ptr + row2 * N + col_base;
          #pragma unroll
          for (int i = 0; i < 8; i++) {
            reinterpret_cast<half2 *>(out1 + i * 8)[0] =
                __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});
          }
        }
      }
    }

    asm volatile("bar.sync 1, %0;" :: "r"(thr) : "memory");
    if (warp_id == 0) asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(BLOCK_N * 2));
  }
}

#endif
template <int BLOCK_N, int BLOCK_K, int NUM_STAGES>
__global__ __launch_bounds__(
  128 + 2 * WARP_SIZE,
  (BLOCK_K == 512) ? 1 : ((BLOCK_N == 128 && (NUM_STAGES == 2 || NUM_STAGES == 3)) ? 2 : 1)
)
void kernel_grouped(const GroupDesc *descs, const BlockDesc *blocks) {
  constexpr int BLOCK_M = 128;
  static_assert(BLOCK_K == 256 || BLOCK_K == 512, "bad BLOCK_K");
  constexpr int Z_STEP = BLOCK_K / 256;
  constexpr int BK_LOG2 = (BLOCK_K == 256) ? 8 : 9;

  const int block_linear = (int)blockIdx.x;
  const int tid = (int)threadIdx.x;
  const int lane_id = tid & (WARP_SIZE - 1);
  const int warp_id = tid >> 5;

  // BlockDesc 只做一次 16B global load，然后 CTA 广播（短 K 固定开销更敏感）
  __shared__ int4 bd_smem;
  if (warp_id == 0 && elect_sync()) {
    bd_smem = reinterpret_cast<const int4 *>(blocks)[block_linear];
  }

  constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE + 2;

  extern __shared__ __align__(1024) char smem_ptr[];
  const int smem = (int)__cvta_generic_to_shared(smem_ptr);
  constexpr int A_size   = BLOCK_M * BLOCK_K / 2;
  constexpr int B_size   = BLOCK_N * BLOCK_K / 2;
  constexpr int SFA_size = 128 * BLOCK_K / 16;
  constexpr int SFB_size_128 = 128 * BLOCK_K / 16;
  constexpr int SFB_size = (BLOCK_N == 256) ? (SFB_size_128 * 2) : SFB_size_128;
  constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;

  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ int64_t mbars[NUM_STAGES * 2 + 1];
  const int tma_mbar_addr = (int)__cvta_generic_to_shared(mbars);
  const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
  const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;

  constexpr int SFA_tmem = BLOCK_N;
  constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);

  if (warp_id == 0 && elect_sync()) {
    #pragma unroll
    for (int i = 0; i < NUM_STAGES * 2 + 1; i++) mbarrier_init(tma_mbar_addr + i * 8, 1);
    asm volatile("fence.mbarrier_init.release.cluster;");
  } else if (warp_id == 1) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 2));
  }
  __syncthreads();

  const int gid = bd_smem.x;
  const int bid_m = bd_smem.y;
  const int bid_n = bd_smem.z;
  const int off_m = bid_m * BLOCK_M;
  const int off_n = bid_n * BLOCK_N;

  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    const GroupDesc *desc = descs + gid;
    const int K = desc->K;
    const int num_iters = K >> BK_LOG2;  // K / BLOCK_K

    const CUtensorMap *A_tmap = &desc->A_tmap;
    const CUtensorMap *B_tmap = &desc->B_tmap;
    const char *SFA_ptr = (const char *)desc->SFA_ptr;
    const char *SFB_ptr = (const char *)desc->SFB_ptr;

    const bool swap = (desc->swap != 0);
    const uint64_t cache_A = swap ? EVICT_LAST  : EVICT_FIRST;
    const uint64_t cache_B = swap ? EVICT_FIRST : EVICT_LAST;

    const int rest_k = K >> 6;  // K / 64
    const int n_tile_128 = (BLOCK_N == 64) ? (bid_n >> 1) : (bid_n * (BLOCK_N / 128));
    const int sfa_tile_base = bid_m * rest_k;
    const int sfb_tile_base = n_tile_128 * rest_k;
    const int sf_step = (BLOCK_K >> 6) * 512;
    const char *sfa_iter = SFA_ptr + (size_t)sfa_tile_base * 512;
    const char *sfb_iter0 = SFB_ptr + (size_t)sfb_tile_base * 512;
    const char *sfb_iter1 = (BLOCK_N == 256) ? (sfb_iter0 + (size_t)rest_k * 512) : sfb_iter0;

    auto issue_tma = [&](int iter_k, int stage_id, const char *sfa_src, const char *sfb_src0, const char *sfb_src1) {
      const int mbar_addr = tma_mbar_addr + stage_id * 8;
      const int A_smem = smem + stage_id * STAGE_SIZE;
      const int B_smem = A_smem + A_size;
      const int SFA_smem = B_smem + B_size;
      const int SFB_smem = SFA_smem + SFA_size;

      const int z = iter_k * Z_STEP;
      tma_3d_gmem2smem(A_smem, A_tmap, 0, off_m, z, mbar_addr, cache_A);
      tma_3d_gmem2smem(B_smem, B_tmap, 0, off_n, z, mbar_addr, cache_B);

      tma_gmem2smem(SFA_smem, sfa_src, SFA_size, mbar_addr, cache_A);
      if constexpr (BLOCK_N == 256) {
        tma_gmem2smem(SFB_smem, sfb_src0, SFB_size_128, mbar_addr, cache_B);
        tma_gmem2smem(SFB_smem + SFB_size_128, sfb_src1, SFB_size_128, mbar_addr, cache_B);
      } else {
        tma_gmem2smem(SFB_smem, sfb_src0, SFB_size_128, mbar_addr, cache_B);
      }

      asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
                   :: "r"(mbar_addr), "r"(STAGE_SIZE) : "memory");
    };

    const int init_stage = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;
    for (int iter_k = 0; iter_k < init_stage; iter_k++) {
      issue_tma(iter_k, iter_k, sfa_iter, sfb_iter0, sfb_iter1);
      sfa_iter += sf_step;
      sfb_iter0 += sf_step;
      if constexpr (BLOCK_N == 256) sfb_iter1 += sf_step;
    }

    for (int iter_k = NUM_STAGES; iter_k < num_iters; iter_k++) {
      const int stage_id = iter_k % NUM_STAGES;
      const int mma_phase = (iter_k / NUM_STAGES - 1) & 1;
      mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);
      issue_tma(iter_k, stage_id, sfa_iter, sfb_iter0, sfb_iter1);
      sfa_iter += sf_step;
      sfb_iter0 += sf_step;
      if constexpr (BLOCK_N == 256) sfb_iter1 += sf_step;
    }
  } else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
    const GroupDesc *desc = descs + gid;
    const int K = desc->K;
    const int num_iters = K >> BK_LOG2;  // K / BLOCK_K

    constexpr int MMA_N = BLOCK_N;
    constexpr int MMA_M = 128;
    constexpr uint32_t i_desc = (1U << 7U)
                              | (1U << 10U)
                              | ((uint32_t)MMA_N >> 3U << 17U)
                              | ((uint32_t)MMA_M >> 7U << 27U);

    auto make_desc_AB = [] __device__ (int addr) -> uint64_t {
      const int SBO = 8 * 128;
      return desc_encode((uint64_t)addr) | (desc_encode((uint64_t)SBO) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
    };
    auto make_desc_SF = [] __device__ (int addr) -> uint64_t {
      const int SBO = 8 * 16;
      return desc_encode((uint64_t)addr) | (desc_encode((uint64_t)SBO) << 32ULL) | (1ULL << 46ULL);
    };

    for (int iter_k = 0; iter_k < num_iters; iter_k++) {
      const int stage_id = iter_k % NUM_STAGES;
      const int tma_phase = (iter_k / NUM_STAGES) & 1;
      mbarrier_wait(tma_mbar_addr + stage_id * 8, tma_phase);

      const int A_smem = smem + stage_id * STAGE_SIZE;
      const int B_smem = A_smem + A_size;
      const int SFA_smem = B_smem + B_size;
      const int SFB_smem = SFA_smem + SFA_size;

      const uint64_t SFA_desc = make_desc_SF(0) + ((uint64_t)SFA_smem >> 4ULL);
      const uint64_t SFB_desc = make_desc_SF(0) + ((uint64_t)SFB_smem >> 4ULL);

      #pragma unroll
      for (int k = 0; k < BLOCK_K / MMA_K; k++) {
        uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
        tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);

        if constexpr (BLOCK_N == 256) {
          uint64_t sfb0 = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
          uint64_t sfb1 = SFB_desc + ((uint64_t)SFB_size_128 >> 4ULL) + (uint64_t)k * (512ULL >> 4ULL);
          tcgen05_cp_nvfp4(SFB_tmem + k * 8 + 0, sfb0);
          tcgen05_cp_nvfp4(SFB_tmem + k * 8 + 4, sfb1);
        } else {
          uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
          tcgen05_cp_nvfp4(SFB_tmem + k * 4, sfb_desc);
        }
      }

      #pragma unroll
      for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {
        #pragma unroll
        for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
          uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
          uint64_t b_desc = make_desc_AB(B_smem + k1 * BLOCK_N * 128 + k2 * 32);

          const int k_sf = k1 * 4 + k2;
          const int scale_A_tmem = SFA_tmem + k_sf * 4;
          const int scale_B_tmem = (BLOCK_N == 256)
              ? (SFB_tmem + k_sf * 8)
              : (SFB_tmem + k_sf * 4 + ((BLOCK_N == 64) ? ((bid_n & 1) * 2) : 0));
          const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
          tcgen05_mma_nvfp4(a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d);
        }
      }

      asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                   :: "r"(mma_mbar_addr + stage_id * 8) : "memory");
    }

    asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                 :: "r"(mainloop_mbar_addr) : "memory");
  } else if (tid < BLOCK_M) {
    const GroupDesc *desc = descs + gid;
    int M = 0;
    int N = 0;
    uint64_t C_ptr_u64 = 0;
    if (lane_id == 0) {
      M = desc->M;
      N = desc->N;
      C_ptr_u64 = desc->C_ptr;
    }
    constexpr uint32_t FULL_MASK = 0xFFFFFFFFu;
    M = __shfl_sync(FULL_MASK, M, 0);
    N = __shfl_sync(FULL_MASK, N, 0);
    uint32_t c_lo = (uint32_t)C_ptr_u64;
    uint32_t c_hi = (uint32_t)(C_ptr_u64 >> 32);
    c_lo = __shfl_sync(FULL_MASK, c_lo, 0);
    c_hi = __shfl_sync(FULL_MASK, c_hi, 0);
    half *C_ptr = (half *)((((uint64_t)c_hi) << 32ULL) | (uint64_t)c_lo);

    const int thr = active_threads_128(M, off_m);
    if (tid >= thr) return;

    mbarrier_wait(mainloop_mbar_addr, 0);
    asm volatile("tcgen05.fence::after_thread_sync;");

    // 题面约束 N 整除 BLOCK_N；且 B 的 TMA 读阶段本就依赖该条件
    // 因此移除 full_n 慢路径以减少分支与指令数

    #pragma unroll
    for (int m = 0; m < 2; m++) {
      const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);
      const int row2 = row + 8;

      float tmp[32];
      constexpr int SEG_N = BLOCK_N / 64;
      #pragma unroll
      for (int seg = 0; seg < SEG_N; seg++) {
        const int col_base = off_n + seg * 64 + ((lane_id & 3) << 1);
        tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, seg * 64);
        asm volatile("tcgen05.wait::ld.sync.aligned;");

        if (row < M) {
          half *out0 = C_ptr + row * N + col_base;
          #pragma unroll
          for (int i = 0; i < 8; i++) {
            reinterpret_cast<half2 *>(out0 + i * 8)[0] =
                __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});
          }
        }
        if (row2 < M) {
          half *out1 = C_ptr + row2 * N + col_base;
          #pragma unroll
          for (int i = 0; i < 8; i++) {
            reinterpret_cast<half2 *>(out1 + i * 8)[0] =
                __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});
          }
        }
      }
    }

    asm volatile("bar.sync 1, %0;" :: "r"(thr) : "memory");
    if (warp_id == 0) asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(BLOCK_N * 2));
  }
}

#if 0
template <int BLOCK_N, int NUM_STAGES>
static __forceinline__ void gemm_launch(
  const at::Tensor& A,
  const at::Tensor& B,
  const at::Tensor& SFA,
  const at::Tensor& SFB,
        at::Tensor& C,
  int M, int N, int K
) {
  const int Apad = (int)A.size(0);
  const char *A_ptr = (const char *)A.data_ptr();
  const char *B_ptr = (const char *)B.data_ptr();
  const char *SFA_ptr = (const char *)SFA.data_ptr();
  const char *SFB_ptr = (const char *)SFB.data_ptr();
  half *C_ptr = (half *)C.data_ptr<at::Half>();

  CUtensorMap A_tmap, B_tmap;
  init_AB_tmap(&A_tmap, A_ptr, (uint64_t)Apad, (uint64_t)K, 128, 256,
               CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
               CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE);
  init_AB_tmap(&B_tmap, B_ptr, (uint64_t)N,    (uint64_t)K, (uint32_t)BLOCK_N, 256,
               CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
               CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE);

  const int grid_m = (M + 127) / 128;
  const int grid_n = (N + BLOCK_N - 1) / BLOCK_N;

  const int tb_size = 128 + 2 * WARP_SIZE;
  const int A_size = 128 * 256 / 2;
  const int B_size = BLOCK_N * 256 / 2;
  const int SF_size = 128 * 256 / 16;
  const int smem_size = (A_size + B_size + SF_size * 2) * NUM_STAGES;

  auto k = kernel<BLOCK_N, NUM_STAGES>;
  // 只在首次使用该设备时设置一次，避免每次调用都走一次 runtime API
  static int last_dev = -1;
  int dev = -1;
  cudaGetDevice(&dev);
  if (dev != last_dev) {
    auto err = cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
    GG_CHECK(err == cudaSuccess, "cudaFuncSetAttribute failed");
    last_dev = dev;
  }

  dim3 grid((unsigned)grid_n, (unsigned)grid_m, 1);
  k<<<grid, tb_size, smem_size>>>(A_tmap, B_tmap, SFA_ptr, SFB_ptr, C_ptr, M, N, K);
}
#endif

static __forceinline__ int pick_stage_from_K(int K) {
  // K 一定是 256 的倍数
  const int iters = K >> 8;
  constexpr int STAGE1_MAX_ITERS = 8;
  constexpr int STAGE3_MAX_ITERS = 16;
  constexpr int STAGE4_ITERS16_ENABLE = 1;
  if (iters <= STAGE1_MAX_ITERS) return 1;
  if (iters < STAGE3_MAX_ITERS) return 3;
  if (iters == STAGE3_MAX_ITERS) return STAGE4_ITERS16_ENABLE ? 4 : 3;
  return 4;
}

// BN128 的 stage=4 控制：用“单变量阈值”便于做单变量对照
// - 设为 1<<30 等超大值可等价禁用 stage=4（强制回退到 stage=3）。
constexpr int BN128_STAGE4_ITERS = 24;
constexpr int BN128_STAGE4_ITERS16_ENABLE = 1;

void gemm_grouped(
  at::TensorList A_list,
  at::TensorList B_list,
  at::TensorList SFA_list,
  at::TensorList SFB_list,
  at::TensorList C_list,
  bool force_bn64
);

at::Tensor gemm(
  const at::Tensor& A,
  const at::Tensor& B,
  const at::Tensor& SFA,
  const at::Tensor& SFB,
        at::Tensor& C,
  int64_t M,
  int64_t N,
  int64_t K
) {
  (void)M;
  (void)N;
  (void)K;
  // 非 grouped 路径复用 grouped kernel（减少重复模板实例与编译/加载压力）
  std::array<at::Tensor, 1> A_list{A};
  std::array<at::Tensor, 1> B_list{B};
  std::array<at::Tensor, 1> SFA_list{SFA};
  std::array<at::Tensor, 1> SFB_list{SFB};
  std::array<at::Tensor, 1> C_list{C};
  gemm_grouped(A_list, B_list, SFA_list, SFB_list, C_list, false);
  return C;
}

struct GroupWorkspace {
  at::Tensor descs_d;
  at::Tensor blocks_d;
  int64_t cap_G;
  int64_t cap_blocks;
  int64_t dev;
  uint64_t last_hash1;
  uint64_t last_hash2;
  int64_t last_G;
  int last_block_n;
  int last_block_k;
  int last_num_blocks;
  int last_max_K;
  int last_stage;
  // 小 G（<=8）优先走全量 pack 对比：消除 hash 碰撞并减少混合乘法
  std::array<std::array<uint64_t, 8>, 8> last_pack_small;
  std::array<uint64_t, 8> last_a0;
  std::array<uint64_t, 8> last_a1;
  std::array<uint64_t, 8> last_a2;
  std::array<uint64_t, 8> last_a3;
  GroupWorkspace()
      : cap_G(0),
        cap_blocks(0),
        dev(-1),
        last_hash1(0),
        last_hash2(0),
        last_G(0),
        last_block_n(0),
        last_block_k(0),
        last_num_blocks(0),
        last_max_K(0),
        last_stage(0),
        last_pack_small{},
        last_a0{0},
        last_a1{0},
        last_a2{0},
        last_a3{0} {}
};

static GroupWorkspace g_ws;

static __forceinline__ void ensure_ws(int64_t dev, int64_t G) {
  if (g_ws.dev != dev || g_ws.cap_G < G || !g_ws.descs_d.defined()) {
    g_ws.dev = dev;
    g_ws.cap_G = G;
    g_ws.cap_blocks = 0;
    at::TensorOptions opt_u8 = at::TensorOptions().device(at::kCUDA, (int)dev).dtype(at::kByte);
    g_ws.descs_d = at::empty({G, (int64_t)sizeof(GroupDesc)}, opt_u8);
    g_ws.blocks_d = at::Tensor();
    g_ws.last_hash1 = 0;
    g_ws.last_hash2 = 0;
    g_ws.last_G = 0;
    g_ws.last_block_n = 0;
    g_ws.last_block_k = 0;
    g_ws.last_num_blocks = 0;
    g_ws.last_max_K = 0;
    g_ws.last_stage = 0;
    g_ws.last_pack_small = {};
    g_ws.last_a0 = {};
    g_ws.last_a1 = {};
    g_ws.last_a2 = {};
    g_ws.last_a3 = {};
  }
}

static __forceinline__ uint64_t fnv1a_mix_u64(uint64_t h, uint64_t x) {
  h ^= x;
  h *= 1099511628211ULL;
  return h;
}

static __forceinline__ uint64_t fmix_u64(uint64_t x) {
  x ^= x >> 33;
  x *= 0xff51afd7ed558ccdULL;
  x ^= x >> 33;
  x *= 0xc4ceb9fe1a85ec53ULL;
  x ^= x >> 33;
  return x;
}

template <int BLOCK_N, int BLOCK_K, int NUM_STAGES>
static __forceinline__ void grouped_launch(
  int64_t dev,
  int num_blocks
) {
  const int tb_size = 128 + 2 * WARP_SIZE;
  const int A_size = 128 * BLOCK_K / 2;
  const int B_size = BLOCK_N * BLOCK_K / 2;
  const int SFA_size = 128 * BLOCK_K / 16;
  const int SFB_size = (BLOCK_N == 256) ? (2 * SFA_size) : SFA_size;
  const int smem_size = (A_size + B_size + SFA_size + SFB_size) * NUM_STAGES;

  auto k = kernel_grouped<BLOCK_N, BLOCK_K, NUM_STAGES>;
  static int last_dev = -1;
  if ((int)dev != last_dev) {
    auto err = cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
    GG_CHECK(err == cudaSuccess, "cudaFuncSetAttribute failed");
    last_dev = (int)dev;
  }

  dim3 grid((unsigned)num_blocks, 1U, 1U);
  k<<<grid, tb_size, smem_size>>>(
    (const GroupDesc *)g_ws.descs_d.data_ptr(),
    (const BlockDesc *)g_ws.blocks_d.data_ptr()
  );
}

void gemm_grouped(
  at::TensorList A_list,
  at::TensorList B_list,
  at::TensorList SFA_list,
  at::TensorList SFB_list,
  at::TensorList C_list,
  bool force_bn64
) {
  const int64_t G = (int64_t)A_list.size();
  if (G <= 0) return;

  const auto &A0 = A_list[0];
  const int64_t dev = (int64_t)A0.get_device();
  if (g_ws.dev != dev) cudaSetDevice((int)dev);
  ensure_ws(dev, G);

  bool use_bn128 = !force_bn64;
  bool use_bn256 = !force_bn64;
  for (int i = 0; i < (int)G; i++) {
    const auto &C = C_list[i];
    const int N = (int)C.size(1);
    if (N < 128 || ((N & 127) != 0)) { use_bn128 = false; use_bn256 = false; break; }
    if (N < 256 || ((N & 255) != 0)) use_bn256 = false;
  }

  // 先构建轻量锚点 + 双 hash：命中则直接 launch（避免全量 sig 构建与比较）
  std::array<uint64_t, 8> sig_a0{};
  std::array<uint64_t, 8> sig_a1{};
  std::array<uint64_t, 8> sig_a2{};
  std::array<uint64_t, 8> sig_a3{};
  std::array<std::array<uint64_t, 8>, 8> pack_small{};
  const bool use_small_pack = (G <= 8);
  const int g_int = (int)G;
  const int a0 = 0;
  const int a1 = (g_int > 1) ? 1 : 0;
  const int a2 = (g_int > 2) ? (g_int - 2) : (g_int - 1);
  const int a3 = g_int - 1;

  uint64_t h1 = 0;
  uint64_t h2 = 0;
  if (!use_small_pack) {
    h1 = 1469598103934665603ULL;
    h1 = fnv1a_mix_u64(h1, (uint64_t)G);

    h2 = 0x243f6a8885a308d3ULL;
    h2 ^= (uint64_t)G;
  }

  int num_blocks_64 = 0;
  int num_blocks_128 = 0;
  int num_blocks_256 = 0;
  int max_K = 0;
  bool all_k_divisible_512 = true;
  bool all_k_equal = true;
  int first_K = -1;

  for (int i = 0; i < (int)G; i++) {
    const auto &A = A_list[i];
    const auto &B = B_list[i];
    const auto &SFA = SFA_list[i];
    const auto &SFB = SFB_list[i];
    const auto &C = C_list[i];

    const int M = (int)C.size(0);
    const int N = (int)C.size(1);
    const int K = (int)A.size(1) * 2;
    const int Apad = (int)A.size(0);
    all_k_divisible_512 = all_k_divisible_512 && ((K & 511) == 0);
    if (first_K < 0) first_K = K;
    else if (K != first_K) all_k_equal = false;

    const int grid_m = (M + 127) / 128;
    num_blocks_64 += grid_m * ((N + 63) / 64);
    if (use_bn128) num_blocks_128 += grid_m * ((N + 127) / 128);
    if (use_bn256) num_blocks_256 += grid_m * ((N + 255) / 256);
    if (K > max_K) max_K = K;

    const uint64_t A_ptr = (uint64_t)A.data_ptr();
    const uint64_t B_ptr = (uint64_t)B.data_ptr();
    const uint64_t SFA_ptr = (uint64_t)SFA.data_ptr();
    const uint64_t SFB_ptr = (uint64_t)SFB.data_ptr();
    const uint64_t C_ptr = (uint64_t)C.data_ptr<at::Half>();

    const std::array<uint64_t, 8> pack = {
      A_ptr,
      B_ptr,
      SFA_ptr,
      SFB_ptr,
      C_ptr,
      (uint64_t)M,
      (uint64_t)N,
      ((uint64_t)K << 32) | (uint64_t)(uint32_t)Apad,
    };

    if (use_small_pack) {
      pack_small[(size_t)i] = pack;
    }

    if (i == a0) sig_a0 = pack;
    if (i == a1) sig_a1 = pack;
    if (i == a2) sig_a2 = pack;
    if (i == a3) sig_a3 = pack;

    if (!use_small_pack) {
      h1 = fnv1a_mix_u64(h1, pack[0]); h2 = fmix_u64(h2 ^ pack[0]);
      h1 = fnv1a_mix_u64(h1, pack[1]); h2 = fmix_u64(h2 ^ pack[1]);
      h1 = fnv1a_mix_u64(h1, pack[2]); h2 = fmix_u64(h2 ^ pack[2]);
      h1 = fnv1a_mix_u64(h1, pack[3]); h2 = fmix_u64(h2 ^ pack[3]);
      h1 = fnv1a_mix_u64(h1, pack[4]); h2 = fmix_u64(h2 ^ pack[4]);
      h1 = fnv1a_mix_u64(h1, pack[5]); h2 = fmix_u64(h2 ^ pack[5]);
      h1 = fnv1a_mix_u64(h1, pack[6]); h2 = fmix_u64(h2 ^ pack[6]);
      h1 = fnv1a_mix_u64(h1, pack[7]); h2 = fmix_u64(h2 ^ pack[7]);
    }
  }

  // BK512：仅当本次 grouped 的所有 K 都可被 512 整除时启用，避免 num_iters 向下取整导致 silent wrong。
  // 额外门控：当 max_K 过长时优先回到 BK256+更深流水，避免 BK512(stage<=2) overlap 不足导致退化。
  // - all_k_equal 仅用于后续策略扩展，不作为硬性约束
  int block_k = 256;
  if (use_bn128 && all_k_divisible_512 && max_K <= 4096) block_k = 512;
  if (block_k == 512) GG_CHECK(all_k_divisible_512, "BK512 requires all K % 512 == 0");

  // stage 决策：从“块数阈值”升级为“占用/波次模型”（依赖设备属性与 smem_size 上界，且按 device 缓存）
  static int cached_prop_dev = -1;
  static int cached_sm_count = 0;
  static int cached_smem_per_sm = 0;
  static int cached_threads_per_sm = 0;
  static int cached_blocks_per_sm = 0;
  static int cached_smem_per_block = 0;
  if ((int)dev != cached_prop_dev) {
    int v = 0;
    cudaDeviceGetAttribute(&v, cudaDevAttrMultiProcessorCount, (int)dev);
    cached_sm_count = v;
    cudaDeviceGetAttribute(&v, cudaDevAttrMaxSharedMemoryPerMultiprocessor, (int)dev);
    cached_smem_per_sm = v;
    cudaDeviceGetAttribute(&v, cudaDevAttrMaxThreadsPerMultiProcessor, (int)dev);
    cached_threads_per_sm = v;
    cudaDeviceGetAttribute(&v, cudaDevAttrMaxBlocksPerMultiprocessor, (int)dev);
    cached_blocks_per_sm = v;
    cudaDeviceGetAttribute(&v, cudaDevAttrMaxSharedMemoryPerBlockOptin, (int)dev);
    cached_smem_per_block = v;
    cached_prop_dev = (int)dev;
  }
  const int sm_count = cached_sm_count;

  // BLOCK_N 选择：BN256 仅在并行度足够（至少 1 wave）时启用，避免小 G case 并行不足导致 SM 空转
  int block_n = 64;
  int num_blocks = num_blocks_64;
  if (use_bn128) {
    const bool bn256_ok = use_bn256 && (sm_count <= 0 || num_blocks_256 >= sm_count);
    block_n = bn256_ok ? 256 : 128;
    num_blocks = bn256_ok ? num_blocks_256 : num_blocks_128;
  }

  // 将 dispatch 关键信息纳入 hash（仅在大 G 路径使用 hash）
  if (!use_small_pack) {
    h1 = fnv1a_mix_u64(h1, (uint64_t)block_n);
    h1 = fnv1a_mix_u64(h1, (uint64_t)block_k);
    h2 = fmix_u64(h2 ^ (uint64_t)block_n);
    h2 = fmix_u64(h2 ^ (uint64_t)block_k);
  }

  auto smem_bytes = [&](int bn, int bk, int stg) -> int {
    const int A_size = 128 * bk / 2;
    const int B_size = bn * bk / 2;
    const int SFA_size = 128 * bk / 16;
    const int SFB_size = (bn == 256) ? (2 * SFA_size) : SFA_size;
    return (A_size + B_size + SFA_size + SFB_size) * stg;
  };
  auto stage_allowed = [&](int bn, int bk, int stg) -> bool {
    if (cached_smem_per_block <= 0) return true;
    return smem_bytes(bn, bk, stg) <= cached_smem_per_block;
  };
  // 将寄存器限制纳入 waves 估算：用 occupancy API 获取真实 active blocks/SM，并按 dev+variant 缓存。
  static int cached_occ_dev = -1;
  static int occ_cache[3][2][5];
  if ((int)dev != cached_occ_dev) {
    for (int bi = 0; bi < 3; bi++) {
      for (int ki = 0; ki < 2; ki++) {
        for (int si = 0; si < 5; si++) occ_cache[bi][ki][si] = 0;
      }
    }
    cached_occ_dev = (int)dev;
  }

  auto active_blocks_per_sm = [&](int bn, int bk, int stg) -> int {
    const int bi = (bn == 64) ? 0 : (bn == 128 ? 1 : 2);
    const int ki = (bk == 256) ? 0 : 1;
    int &slot = occ_cache[bi][ki][stg];
    if (slot != 0) return slot;

    int blocks = 1;
    const int tb = 128 + 2 * WARP_SIZE;
    const int smem = smem_bytes(bn, bk, stg);

    auto occ_query = [&](auto k) {
      cudaError_t aerr = cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
      GG_CHECK(aerr == cudaSuccess, "cudaFuncSetAttribute failed");
      int b = 0;
      cudaError_t err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&b, k, tb, smem);
      if (err == cudaSuccess && b > 0) blocks = b;
    };

    if (bn == 64) {
      if (bk == 256) {
        if (stg == 1) occ_query(kernel_grouped<64, 256, 1>);
        else if (stg == 3) occ_query(kernel_grouped<64, 256, 3>);
        else occ_query(kernel_grouped<64, 256, 4>);
      }
    } else if (bn == 128) {
      if (bk == 512) {
        if (stg <= 1) occ_query(kernel_grouped<128, 512, 1>);
        else if (stg == 2) occ_query(kernel_grouped<128, 512, 2>);
        else occ_query(kernel_grouped<128, 512, 3>);
      } else {
        if (stg <= 1) occ_query(kernel_grouped<128, 256, 1>);
        else if (stg == 2) occ_query(kernel_grouped<128, 256, 2>);
        else if (stg == 3) occ_query(kernel_grouped<128, 256, 3>);
        else occ_query(kernel_grouped<128, 256, 4>);
      }
    } else {
      if (bk == 512) {
        if (stg <= 1) occ_query(kernel_grouped<256, 512, 1>);
        else occ_query(kernel_grouped<256, 512, 2>);
      } else {
        if (stg <= 1) occ_query(kernel_grouped<256, 256, 1>);
        else if (stg == 2) occ_query(kernel_grouped<256, 256, 2>);
        else if (stg == 3) occ_query(kernel_grouped<256, 256, 3>);
        else occ_query(kernel_grouped<256, 256, 4>);
      }
    }

    if (blocks < 1) blocks = 1;
    slot = blocks;
    return blocks;
  };
  auto cap_blocks = [&](int bn, int bk, int stg) -> int64_t {
    if (!stage_allowed(bn, bk, stg)) return 0;
    const int a = active_blocks_per_sm(bn, bk, stg);
    return (int64_t)sm_count * (int64_t)a;
  };

  const int iters = (block_k == 512) ? (max_K >> 9) : (max_K >> 8);
  int64_t cap4 = 0;
  bool stage4_16_ok = true;
  bool stage4_8_ok = true;
  if (block_k == 256) {
    cap4 = cap_blocks(block_n, block_k, 4);
    stage4_16_ok = (cap4 <= 0) ? true : ((int64_t)num_blocks * 4 >= cap4 * 3);
    stage4_8_ok  = (cap4 <= 0) ? true : ((int64_t)num_blocks >= cap4 * 2);
  }

  int stage = 0;
  if (use_bn128) {
    if (block_k == 512) {
      // BK512 下允许 {1,2,3}：用更深的 stage 改善 TMA↔MMA overlap（受 SMEM 上限约束）
      if (iters <= 3) stage = 1;
      else if (iters <= 6) stage = 2;
      else stage = 3;
      // BN256+BK512 的 SMEM 压力极高，stage=3 通常不成立（也没必要编译更多变体）
      if (block_n == 256 && stage > 2) stage = 2;
    } else {
      if (iters <= 7) stage = 1;
      else if (iters == 8) stage = stage4_8_ok ? 4 : 2;
      else if (iters >= BN128_STAGE4_ITERS || (BN128_STAGE4_ITERS16_ENABLE && iters == 16 && stage4_16_ok)) stage = 4;
      else stage = 3;
    }
    while (stage > 1 && (!stage_allowed(block_n, block_k, stage))) stage--;
    if (stage < 1) stage = 1;
  } else {
    stage = pick_stage_from_K(max_K);
    if (((max_K >> 8) == 16) && stage == 4 && (!stage4_16_ok)) stage = 3;
    if (stage == 2) stage = 1;
    while (stage > 1 && (!stage_allowed(block_n, block_k, stage))) {
      if (stage == 4) stage = 3;
      else stage = 1;
    }
  }
  bool hit = (g_ws.last_G == G)
          && (g_ws.last_block_n == block_n)
          && (g_ws.last_block_k == block_k)
          && (g_ws.last_num_blocks == num_blocks)
          && (g_ws.last_max_K == max_K)
          && (g_ws.last_stage == stage)
          && g_ws.descs_d.defined()
          && g_ws.blocks_d.defined();
  if (hit) {
    if (use_small_pack) {
      #pragma unroll
      for (int i = 0; i < 8; i++) {
        if (i >= (int)G) break;
        if (g_ws.last_pack_small[(size_t)i] != pack_small[(size_t)i]) { hit = false; break; }
      }
    } else {
      hit = hit
         && (g_ws.last_hash1 == h1)
         && (g_ws.last_hash2 == h2)
         && (g_ws.last_a0 == sig_a0)
         && (g_ws.last_a1 == sig_a1)
         && (g_ws.last_a2 == sig_a2)
         && (g_ws.last_a3 == sig_a3);
    }
  }

  if (!hit) {
    std::vector<GroupDesc> descs((size_t)G);
    std::vector<BlockDesc> blocks((size_t)num_blocks);
    int block_cursor = 0;
    for (int i = 0; i < (int)G; i++) {
      const auto &A = A_list[i];
      const auto &B = B_list[i];
      const auto &SFA = SFA_list[i];
      const auto &SFB = SFB_list[i];
      const auto &C = C_list[i];

      const int M = (int)C.size(0);
      const int N = (int)C.size(1);
      const int K = (int)A.size(1) * 2;
      GroupDesc d;
      const CUtensorMapSwizzle sw = CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B;
      // A 的 l2Promotion 按 K 分段：短 K 避免过取与缓存压力；长 K 才启用更大粒度
      const CUtensorMapL2promotion a_l2 = (K >= 4096)
          ? CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B
          : CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE;
      // 关键优化：A 使用真实 M 的 TensorMap，以消除 Python 侧的 A padding copy（避免额外 device 拷贝与 launch）。
      // 注意：不依赖 OOB 的“数值 0-fill”；正确性只依赖 row>=M 的行不会写回到 C（见 epilogue 写回保护）。
      init_AB_tmap(&d.A_tmap, (const char *)A.data_ptr(), (uint64_t)M, (uint64_t)K, 128, (uint32_t)block_k,
                   sw,
                   a_l2);
      init_AB_tmap(&d.B_tmap, (const char *)B.data_ptr(), (uint64_t)N,    (uint64_t)K, (uint32_t)block_n, (uint32_t)block_k,
                   sw,
                   CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE);
      d.SFA_ptr = (uint64_t)SFA.data_ptr();
      d.SFB_ptr = (uint64_t)SFB.data_ptr();
      d.C_ptr = (uint64_t)C.data_ptr<at::Half>();
      d.M = M;
      d.N = N;
      d.K = K;
      const int grid_m = (M + 127) / 128;
      const int grid_n = (N + block_n - 1) / block_n;
      // cache hint：短 K 时更偏向保留 A/SFA，以提升同一 bid_m 下跨 bid_n 的复用命中
      const bool swap = (grid_n >= grid_m) || (K == 2048);
      d.swap = (uint8_t)(swap ? 1 : 0);
      d._pad0 = 0;
      d._pad1 = 0;
      descs[(size_t)i] = d;

      // 构建紧凑 block 列表：按 gid -> bid_m -> bid_n 排序，利于同一 A tile 的邻近调度复用
      for (int bm = 0; bm < grid_m; bm++) {
        for (int bn = 0; bn < grid_n; bn++) {
          BlockDesc bd;
          bd.gid = i;
          bd.bid_m = bm;
          bd.bid_n = bn;
          bd._pad = 0;
          blocks[(size_t)block_cursor] = bd;
          block_cursor++;
        }
      }
    }
    GG_CHECK(block_cursor == num_blocks, "bad block list");

    // blocks 缓冲按需扩容（只在 !hit 分支发生）
    if (!g_ws.blocks_d.defined() || g_ws.cap_blocks < (int64_t)num_blocks) {
      g_ws.cap_blocks = (int64_t)num_blocks;
      at::TensorOptions opt_u8 = at::TensorOptions().device(at::kCUDA, (int)dev).dtype(at::kByte);
      g_ws.blocks_d = at::empty({(int64_t)num_blocks, (int64_t)sizeof(BlockDesc)}, opt_u8);
    }

    // 单次异步拷贝：避免 host 阻塞（与后续 kernel launch 同一默认队列保证顺序）
    cudaError_t cperr = cudaMemcpyAsync(
      g_ws.descs_d.data_ptr(),
      descs.data(),
      (size_t)G * sizeof(GroupDesc),
      cudaMemcpyHostToDevice,
      0
    );
    GG_CHECK(cperr == cudaSuccess, "memcpy desc fail");
    cudaError_t bperr = cudaMemcpyAsync(
      g_ws.blocks_d.data_ptr(),
      blocks.data(),
      (size_t)num_blocks * sizeof(BlockDesc),
      cudaMemcpyHostToDevice,
      0
    );
    GG_CHECK(bperr == cudaSuccess, "memcpy block fail");

    g_ws.last_hash1 = h1;
    g_ws.last_hash2 = h2;
    g_ws.last_G = G;
    g_ws.last_block_n = block_n;
    g_ws.last_block_k = block_k;
    g_ws.last_num_blocks = num_blocks;
    g_ws.last_max_K = max_K;
    g_ws.last_stage = stage;
    if (use_small_pack) g_ws.last_pack_small = pack_small;
    g_ws.last_a0 = sig_a0;
    g_ws.last_a1 = sig_a1;
    g_ws.last_a2 = sig_a2;
    g_ws.last_a3 = sig_a3;
  } else {
    num_blocks = g_ws.last_num_blocks;
  }

  if (use_bn128) {
    if (block_k == 512) {
      if (block_n == 256) {
        if (stage == 2) grouped_launch<256, 512, 2>(dev, num_blocks);
        else grouped_launch<256, 512, 1>(dev, num_blocks);
      } else {
        if (stage == 3) grouped_launch<128, 512, 3>(dev, num_blocks);
        else if (stage == 2) grouped_launch<128, 512, 2>(dev, num_blocks);
        else grouped_launch<128, 512, 1>(dev, num_blocks);
      }
    } else if (block_n == 256) {
      if (stage == 1) grouped_launch<256, 256, 1>(dev, num_blocks);
      else if (stage == 2) grouped_launch<256, 256, 2>(dev, num_blocks);
      else if (stage == 4) grouped_launch<256, 256, 4>(dev, num_blocks);
      else grouped_launch<256, 256, 3>(dev, num_blocks);
    } else {
      if (stage == 1) grouped_launch<128, 256, 1>(dev, num_blocks);
      else if (stage == 2) grouped_launch<128, 256, 2>(dev, num_blocks);
      else if (stage == 4) grouped_launch<128, 256, 4>(dev, num_blocks);
      else grouped_launch<128, 256, 3>(dev, num_blocks);
    }
  } else {
    if (stage == 1) grouped_launch<64, 256, 1>(dev, num_blocks);
    else if (stage == 3) grouped_launch<64, 256, 3>(dev, num_blocks);
    else grouped_launch<64, 256, 4>(dev, num_blocks);
  }

  GG_CHECK(cudaGetLastError() == cudaSuccess, "kernel launch failed");
}

TORCH_LIBRARY(nvfp4_group_gemm_opt, m) {
  m.def("gemm(Tensor A, Tensor B, Tensor SFA, Tensor SFB, Tensor(a!) C, int M, int N, int K) -> Tensor");
  m.impl("gemm", &gemm);
  m.def("gemm_grouped(Tensor[] A, Tensor[] B, Tensor[] SFA, Tensor[] SFB, Tensor[] C, bool force_bn64) -> ()");
  m.impl("gemm_grouped", &gemm_grouped);
}
"""

    build_dir = os.path.join(os.path.dirname(__file__), ".build_nvfp4_group_gemm_opt")
    os.makedirs(build_dir, exist_ok=True)

    _diag_write("[diag] load_ext begin")
    load_inline(
        name="nvfp4_group_gemm_opt_ext",
        cpp_sources="",
        cuda_sources=cuda_src,
        functions=None,
        extra_cflags=["-O3"],
        extra_cuda_cflags=[
            "-O3",
            "-gencode=arch=compute_100a,code=sm_100a",
            "--use_fast_math",
            "--expt-extended-lambda",
            "--expt-relaxed-constexpr",
            "--relocatable-device-code=false",
            "-std=c++17",
        ],
        extra_ldflags=["-lcuda"],
        with_cuda=True,
        is_python_module=False,
        no_implicit_headers=True,
        build_directory=build_dir,
        verbose=False,
    )
    _diag_write("[diag] load_ext end")

    _EXT_READY = True
    _OPS_READY = False
    _GEMM = None
    _GEMM_GROUPED = None


def _get_ops():
    global _OPS_READY, _GEMM, _GEMM_GROUPED
    if not _EXT_READY:
        _load_ext()
    if not _OPS_READY:
        _GEMM = torch.ops.nvfp4_group_gemm_opt.gemm
        _GEMM_GROUPED = torch.ops.nvfp4_group_gemm_opt.gemm_grouped
        _OPS_READY = True
    return _GEMM, _GEMM_GROUPED


def _as_u8(x: torch.Tensor) -> torch.Tensor:
    if x.dtype == torch.uint8:
        return x
    if x.element_size() != 1:
        raise RuntimeError("packed tensor must have 1-byte elements")
    return x.view(torch.uint8)


def _reorder_scale_from_raw(scale_u8_2d: torch.Tensor, rows_pad: int) -> torch.Tensor:
    if scale_u8_2d.dim() != 2:
        raise RuntimeError("scale must be 2D")
    rows = int(scale_u8_2d.size(0))
    k16 = int(scale_u8_2d.size(1))
    if (k16 & 3) != 0:
        raise RuntimeError("K//16 must be multiple of 4")
    if (rows_pad & 127) != 0:
        raise RuntimeError("rows_pad must be multiple of 128")
    blk_m = rows_pad // 128
    blk_k = k16 // 4
    buf = torch.empty((rows_pad, k16), device=scale_u8_2d.device, dtype=torch.uint8)
    buf[:rows].copy_(scale_u8_2d)
    v = buf.view(blk_m, 32, 4, blk_k, 4).permute(0, 3, 1, 2, 4).contiguous()
    return v


def _must_sfp_layout(x: torch.Tensor) -> None:
    
    if (not x.is_cuda) or (x.dim() != 6) or (x.element_size() != 1) or (int(x.storage_offset()) != 0):
        raise RuntimeError("bad sfx_p layout")
    if int(x.size(0)) != 32 or int(x.size(1)) != 4 or int(x.size(3)) != 4:
        raise RuntimeError("bad sfx_p shape")


def custom_kernel(data):
    abc_tensors, sfasfb_tensors, sfasfb_reordered_tensors, problem_sizes = data

    gemm, gemm_grouped = _get_ops()
    global _DIAG_ONCE
    if _DIAG and _DIAG_ONCE:
        _diag_write("[diag] first custom_kernel call")
        _DIAG_ONCE = False

    g = len(problem_sizes)
    if g == 0:
        return []

    dev0 = int(abc_tensors[0][0].device.index)
    all_l1 = True
    for i in range(g):
        if int(problem_sizes[i][3]) != 1:
            all_l1 = False
            break
        if int(abc_tensors[i][0].device.index) != dev0:
            all_l1 = False
            break

    as_u8 = _as_u8
    must_sfp = _must_sfp_layout
    reorder_scale = _reorder_scale_from_raw

    outs: List[torch.Tensor] = []

    if all_l1 and (not _FORCE_NO_GROUPED):
        
        all_c_contig = True
        for i in range(g):
            if not abc_tensors[i][2].is_contiguous():
                all_c_contig = False
                break

        a_list = [None] * g
        b_list = [None] * g
        sfa_list = [None] * g
        sfb_list = [None] * g
        c_list = [None] * g
        need_copy: List[tuple[torch.Tensor, torch.Tensor]] = []
        if not all_c_contig:
            outs = [None] * g

        for i in range(g):
            a, b, c = abc_tensors[i]
            sfa, sfb = sfasfb_tensors[i]
            sfa_p, sfb_p = sfasfb_reordered_tensors[i]
            m, n, _k, _l = problem_sizes[i]

            m_int = int(m)
            n_int = int(n)

            c_out = c
            if all_c_contig:
                c_tmp = c_out
            else:
                outs[i] = c_out
                if not c_out.is_contiguous():
                    c_tmp = torch.empty_like(c_out, memory_format=torch.contiguous_format)
                    need_copy.append((c_out, c_tmp))
                else:
                    c_tmp = c_out

            a_u8 = as_u8(a)
            b_u8 = as_u8(b)
            if not a_u8.is_contiguous():
                raise RuntimeError("grouped 快路要求 A contiguous")
            if not b_u8.is_contiguous():
                raise RuntimeError("grouped 快路要求 B contiguous")

            if _FORCE_RAW_SF:
                m_pad = ((m_int + 127) // 128) * 128
                n_pad = ((n_int + 127) // 128) * 128
                sfa2 = as_u8(sfa[..., 0]).contiguous()
                sfb2 = as_u8(sfb[..., 0]).contiguous()
                sfa_arg = reorder_scale(sfa2, m_pad)
                sfb_arg = reorder_scale(sfb2, n_pad)
            else:
                must_sfp(sfa_p)
                must_sfp(sfb_p)
                sfa_arg = sfa_p
                sfb_arg = sfb_p

            a_list[i] = a_u8
            b_list[i] = b_u8
            sfa_list[i] = sfa_arg
            sfb_list[i] = sfb_arg
            c_list[i] = c_tmp

        gemm_grouped(a_list, b_list, sfa_list, sfb_list, c_list, _FORCE_BN64)

        if all_c_contig:
            return c_list

        for c_out, c_tmp in need_copy:
            c_out.copy_(c_tmp)
        return outs

    for i in range(g):
        a, b, c = abc_tensors[i]
        sfa, sfb = sfasfb_tensors[i]
        sfa_p, sfb_p = sfasfb_reordered_tensors[i]
        m, n, k, l = problem_sizes[i]

        m_int = int(m)
        n_int = int(n)
        k_int = int(k)
        l_int = int(l)

        c_out = c
        if not c_out.is_contiguous():
            c_tmp = torch.empty_like(c_out, memory_format=torch.contiguous_format)
        else:
            c_tmp = c_out

        if l_int == 1:
            a_u8 = as_u8(a).contiguous()
            b_u8 = as_u8(b).contiguous()

            if _FORCE_RAW_SF:
                m_pad = ((m_int + 127) // 128) * 128
                n_pad = ((n_int + 127) // 128) * 128
                sfa2 = as_u8(sfa[..., 0]).contiguous()
                sfb2 = as_u8(sfb[..., 0]).contiguous()
                sfa_arg = reorder_scale(sfa2, m_pad)
                sfb_arg = reorder_scale(sfb2, n_pad)
            else:
                must_sfp(sfa_p)
                must_sfp(sfb_p)
                sfa_arg = sfa_p
                sfb_arg = sfb_p

            gemm(a_u8, b_u8, sfa_arg, sfb_arg, c_tmp, m_int, n_int, k_int)
        else:
            for li in range(l_int):
                a2 = as_u8(a[..., li]).contiguous().unsqueeze(-1)
                b2 = as_u8(b[..., li]).contiguous().unsqueeze(-1)

                m_pad = ((m_int + 127) // 128) * 128
                n_pad = ((n_int + 127) // 128) * 128

                sfa2 = as_u8(sfa[..., li]).contiguous()
                sfb2 = as_u8(sfb[..., li]).contiguous()
                sfa_r = reorder_scale(sfa2, m_pad)
                sfb_r = reorder_scale(sfb2, n_pad)

                c2 = torch.empty((m_int, n_int, 1), device=c_tmp.device, dtype=torch.float16)
                gemm(a2, b2, sfa_r, sfb_r, c2, m_int, n_int, k_int)
                c_tmp[..., li].copy_(c2[..., 0])

        if c_tmp is not c_out:
            c_out.copy_(c_tmp)
        outs.append(c_out)

    return outs


__all__ = ["custom_kernel"]