cublaslt_nvfp4_scaled_mm_n2048_k2048

import torch
from torchao.prototype.mx_formats.utils import to_blocked


def _unswizzle_sf(sf, row, col, vec=16):
    f = vec * 4
    nmt = (row + 127) // 128
    nkt = (col + f - 1) // f
    r = sf.view(nmt, nkt, 32, 4, 4).transpose(1, 3).reshape(nmt * 32 * 4, nkt * 4)
    return r[:row, : col // vec].contiguous()


def run(A_fp4, A_scale, B_fp4, B_scale, alpha):
    # cuBLASLt NVFP4 GEMM via torch._scaled_mm. Stored UE4M3 block scales are in the
    # flashinfer 128x4 swizzle; un-swizzle to logical [*, K/16] then re-block for cuBLAS.
    m = A_fp4.shape[0]
    k = A_fp4.shape[1] * 2
    n = B_fp4.shape[0]
    a_sc = _unswizzle_sf(A_scale.view(torch.uint8), m, k, 16).view(torch.float8_e4m3fn).reshape(m, k // 16)
    b_sc = _unswizzle_sf(B_scale.view(torch.uint8), n, k, 16).view(torch.float8_e4m3fn).reshape(n, k // 16)
    a_blk = to_blocked(a_sc)
    b_blk = to_blocked(b_sc)
    out = torch._scaled_mm(
        A_fp4.view(torch.float4_e2m1fn_x2),
        B_fp4.view(torch.float4_e2m1fn_x2).T,
        a_blk, b_blk, out_dtype=torch.bfloat16,
    )
    return (out.float() * alpha).to(torch.bfloat16)