submission 69507

import torch
from task import input_t, output_t

# --- your existing helpers, unchanged ---
sf_vec_size = 16
def ceil_div(a, b): return (a + b - 1) // b
def to_blocked(input_matrix):
    rows, cols = input_matrix.shape
    n_row_blocks = ceil_div(rows, 128)
    n_col_blocks = ceil_div(cols, 4)
    blocks = input_matrix.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
    rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
    return rearranged.flatten()

# ---- Inline cuBLASLt bridge (uses CUTLASS/CuTe-backed kernels under the hood) ----
_ext = None
try:
    from torch.utils.cpp_extension import load_inline
    cuda_src = r"""
#include <cuda.h>
#include <cuda_runtime.h>
#include <cublasLt.h>
#include <stdexcept>
#include <string>

#define CHECK_CUDA(x) do{cudaError_t e=(x); if(e!=cudaSuccess){throw std::runtime_error(std::string("CUDA error: ")+cudaGetErrorString(e));}}while(0)
#define CHECK_Lt(x) do{cublasStatus_t s=(x); if(s!=CUBLAS_STATUS_SUCCESS){throw std::runtime_error(std::string("cuBLASLt error code ")+std::to_string((int)s));}}while(0)

extern "C" {

struct Ctx {
  cublasLtHandle_t lt;
  void* ws; size_t ws_bytes;
  Ctx(): lt(nullptr), ws(nullptr), ws_bytes(32<<20) {
    CHECK_Lt(cublasLtCreate(&lt));
    CHECK_CUDA(cudaMalloc(&ws, ws_bytes));
  }
  ~Ctx(){ if(ws) cudaFree(ws); if(lt) cublasLtDestroy(lt); }
};
static Ctx* g = nullptr;

void nvfp4_n1_init(){ if(!g) g = new Ctx(); }
void nvfp4_n1_destroy(){ if(g){ delete g; g=nullptr; } }

// NOTE: This assumes CUDA 12.4+ with FP4/FP8 support in cuBLASLt.
// We compute C[m,1] = A[m,k2] (fp4x2) * B^T[k2,1] (fp4x2), with external NVFP4(1x16) scales.
// scaleA length = m*sf_k, scaleB length = 1*sf_k  (N==1!)
int nvfp4_gemv_n1(
  const void* A,        // [m, k2]  fp4x2
  const void* BT_col0,  // [k2, 1]  fp4x2   (first column of B^T)
  const void* scaleA,   // [m*sf_k] fp8 e4m3
  const void* scaleB,   // [sf_k]   fp8 e4m3  (only the first column)
  void*       C,        // [m, 1]   fp16
  int m, int k2, int sf_k,
  cudaStream_t stream)
{
  if(!g) return -100;

  // Datatypes (enum values are toolkit-dependent; we rely on cuBLASLt to accept them)
  const cudaDataType_t compute = CUDA_R_16F;

  cublasLtMatmulDesc_t op;
  CHECK_Lt(cublasLtMatmulDescCreate(&op, compute));

  // A layout: rows=m, cols=k2, ld=k2, type=FP4 (packed x2)
  cublasLtMatrixLayout_t Adesc;
  CHECK_Lt(cublasLtMatrixLayoutCreate(&Adesc, (cudaDataType_t)0x20 /*FP4 e2m1 x2*/, m, k2, k2));

  // B^T layout: rows=k2, cols=1, ld=1, type=FP4 (packed x2)
  cublasLtMatrixLayout_t Bdesc;
  CHECK_Lt(cublasLtMatrixLayoutCreate(&Bdesc, (cudaDataType_t)0x20 /*FP4 e2m1 x2*/, k2, 1, 1));

  // C layout: rows=m, cols=1, ld=1, type=FP16
  cublasLtMatrixLayout_t Cdesc;
  CHECK_Lt(cublasLtMatrixLayoutCreate(&Cdesc, CUDA_R_16F, m, 1, 1));

  // Attach external per-block scales (NVFP4 blockwise 1x16)
  // Attribute IDs are forward-looking; if not supported, the call will fail and Python will fall back.
  CHECK_Lt(cublasLtMatrixLayoutSetAttribute(
      Adesc, (cublasLtMatrixLayoutAttribute_t)0x380 /*SCALE_VECTOR*/, &scaleA, sizeof(scaleA)));
  CHECK_Lt(cublasLtMatrixLayoutSetAttribute(
      Bdesc, (cublasLtMatrixLayoutAttribute_t)0x380 /*SCALE_VECTOR*/, &scaleB, sizeof(scaleB)));
  CHECK_Lt(cublasLtMatrixLayoutSetAttribute(
      Adesc, (cublasLtMatrixLayoutAttribute_t)0x381 /*SCALE_BLOCK_K*/, &sf_k, sizeof(int)));
  CHECK_Lt(cublasLtMatrixLayoutSetAttribute(
      Bdesc, (cublasLtMatrixLayoutAttribute_t)0x381 /*SCALE_BLOCK_K*/, &sf_k, sizeof(int)));

  // alpha,beta
  __half alpha = __float2half(1.0f), beta = __float2half(0.0f);

  // Preference & heuristic
  cublasLtMatmulPreference_t pref; CHECK_Lt(cublasLtMatmulPreferenceCreate(&pref));
  CHECK_Lt(cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &g->ws_bytes, sizeof(g->ws_bytes)));
  cublasLtMatmulHeuristicResult_t heur; int nres=0;
  CHECK_Lt(cublasLtMatmulAlgoGetHeuristic(g->lt, op, Adesc, Bdesc, Cdesc, Cdesc, pref, 1, &heur, &nres));
  if(nres==0){ cublasLtMatmulPreferenceDestroy(pref);
               cublasLtMatrixLayoutDestroy(Adesc);
               cublasLtMatrixLayoutDestroy(Bdesc);
               cublasLtMatrixLayoutDestroy(Cdesc);
               cublasLtMatmulDescDestroy(op);
               return -2; }

  CHECK_Lt(cublasLtMatmul(g->lt, op,
                          &alpha,
                          A, Adesc,
                          BT_col0, Bdesc,
                          &beta,
                          C, Cdesc,
                          C, Cdesc,
                          &heur.algo,
                          g->ws, g->ws_bytes,
                          stream));

  cublasLtMatmulPreferenceDestroy(pref);
  cublasLtMatrixLayoutDestroy(Adesc);
  cublasLtMatrixLayoutDestroy(Bdesc);
  cublasLtMatrixLayoutDestroy(Cdesc);
  cublasLtMatmulDescDestroy(op);
  return 0;
}

} // extern "C"
"""
    _ext = load_inline(
        name="nvfp4_n1_bridge",
        cpp_sources="",
        cuda_sources=cuda_src,
        functions=["nvfp4_n1_init","nvfp4_n1_destroy","nvfp4_gemv_n1"],
        extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17"],
        extra_ldflags=["-lcublasLt","-lcublas"],
        verbose=False,
    )
    _ext.nvfp4_n1_init()
except Exception:
    _ext = None

def _have_ext(): return _ext is not None

# ---------------- Aggressive kernel: N=1 GEMV via cuBLASLt (with safe fallback) ----------------
def custom_kernel(
    data: input_t,
) -> output_t:
    a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data

    device = a_ref.device
    m, k2, l = a_ref.shape
    n_padded = b_ref.shape[0]
    assert n_padded == 128, "B must be padded to 128 in N for NVFP4"
    k = k2 * 2
    sf_k = ceil_div(k, 16)

    # Pre-transposed view (no .contiguous() on FP4)
    bT_all = b_ref.transpose(0, 1)  # [k/2, 128, l]

    # Utility: pack 2D scales and return a contiguous 1D fp8 vector (on CUDA)
    def pack_scales_1d(sf2d: torch.Tensor) -> torch.Tensor:
        return to_blocked(sf2d.to(device=device, non_blocking=True)).contiguous()

    if _have_ext():
        # Fast cuBLASLt path: compute only column 0 — N=1 kernel
        for li in range(l):
            # Pack A scales (length m*sf_k)
            scale_a = pack_scales_1d(sfa_ref_cpu[:, :, li])
            if scale_a.numel() != m * sf_k:
                raise RuntimeError(f"scale_a numel {scale_a.numel()} != {m*sf_k}")

            # Pack B scales (length 128*sf_k); take only column 0 → first sf_k entries
            scale_b_full = pack_scales_1d(sfb_ref_cpu[:, :, li])
            if scale_b_full.numel() != 128 * sf_k:
                raise RuntimeError(f"scale_b numel {scale_b_full.numel()} != {128*sf_k}")
            scale_b_col0 = scale_b_full[:sf_k].contiguous()

            # B^T first column: [k/2, 1] view (no FP4 copies!)
            BT_col0 = bT_all[:, 0:1, li]

            # Output column buffer [m,1] in fp16; we write direct into c_ref column
            Cout_col = torch.empty((m, 1), dtype=torch.float16, device=device)

            # Call cuBLASLt GEMV N=1
            err = _ext.nvfp4_gemv_n1(
                a_ref[:, :, li].data_ptr(),
                BT_col0.data_ptr(),
                scale_a.data_ptr(),
                scale_b_col0.data_ptr(),
                Cout_col.data_ptr(),
                int(m), int(k2), int(sf_k),
                torch.cuda.current_stream(device=device).cuda_stream
            )
            if err != 0:
                # If extension or attrs unsupported here, fall back below.
                _ext = None
                break

            # Store result
            c_ref[:, 0, li].copy_(Cout_col[:, 0], non_blocking=True)

        if _have_ext():
            torch.cuda.synchronize(device)
            return c_ref

    # ---------------- Safe PyTorch fallback (computes 128 columns, then slice) ----------------
    # NOTE: This is slower because it does full m×128 work; only used if the extension path is unavailable.
    for li in range(l):
        scale_a = pack_scales_1d(sfa_ref_cpu[:, :, li])
        scale_b = pack_scales_1d(sfb_ref_cpu[:, :, li])
        res = torch._scaled_mm(
            a_ref[:, :, li],
            bT_all[:, :, li],
            scale_a,
            scale_b,
            bias=None,
            out_dtype=torch.float16,
        )
        c_ref[:, 0, li].copy_(res[:, 0], non_blocking=True)
    torch.cuda.synchronize(device)
    return c_ref