submission 188824

"""
Single-file submission that builds a CUTLASS NVFP4 GEMM kernel via
torch.utils.cpp_extension.load_inline.

Variant: NVF4 schedule (1Sm), ClusterShape 1x2x1, Tile 128x64x256.
Note: The v2 kernel name uses C/D=bf16, but the nvfp4_gemm task expects fp16 output,
so this submission keeps C/D as fp16.
"""

import hashlib
import os
from torch.utils.cpp_extension import load_inline
import torch
from task import input_t, output_t

# Ensure we compile for SM100a (tcgen05.mma.blockscaled).
# Many top submissions also pass an explicit -gencode, but this helps torch's build logic too.
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "10.0a")

_mod = None


def _cutlass_paths() -> list[str]:
    # Try multiple common locations for CUTLASS
    candidates = [
        # Repo checkout (this workspace)
        os.path.abspath(
            os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..", "cutlass")
        ),
        os.path.join(os.getcwd(), "cutlass"),
        os.getenv("CUTLASS_PATH"),  # 1. Explicit env var
        "/mnt/cutlass",  # 2. Standard location
        "/mnt/SFS-jpKCDTAn/cutlass",  # 3. Known SFS location
    ]

    cutlass_root = None
    for path in candidates:
        if (
            path
            and os.path.isdir(path)
            and os.path.exists(os.path.join(path, "include", "cutlass", "cutlass.h"))
        ):
            cutlass_root = path
            break

    if cutlass_root is None:
        # Fallback to /mnt/cutlass and hop for the best, or print warning
        cutlass_root = "/mnt/cutlass"
        print(
            f"Warning: Could not find CUTLASS in candidates: {candidates}. Defaulting to {cutlass_root}"
        )

    return [
        os.path.join(cutlass_root, "include"),
        os.path.join(cutlass_root, "tools", "util", "include"),
    ]


def _pack_scale_factors_blocked(sf: torch.Tensor) -> torch.Tensor:
    """
    Pack scale factors from [MN, K/16, L] into the 1D blocked layout expected by SM100 block-scaled kernels.

    Matches reference.py's `to_blocked()` ordering, but runs on GPU and supports batching in L.
    Returns a flat 1D tensor with batches concatenated (L-major).
    """
    sf = sf.contiguous()
    rows, cols, batch = sf.shape

    if rows % 128 != 0 or cols % 4 != 0:
        raise RuntimeError(
            f"Scale factors must be padded: rows%128==0 and cols%4==0 (got rows={rows}, cols={cols})"
        )

    n_row_blocks = rows // 128
    n_col_blocks = cols // 4

    # [rows, cols, L] -> [L, rows, cols]
    sf_lrc = sf.permute(2, 0, 1).contiguous()

    # [L, n_row_blocks, 128, n_col_blocks, 4] -> [L, n_row_blocks, n_col_blocks, 128, 4]
    blocks = sf_lrc.view(batch, n_row_blocks, 128, n_col_blocks, 4).permute(
        0, 1, 3, 2, 4
    )

    # blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16).flatten()
    rearranged = (
        blocks.reshape(batch, -1, 4, 32, 4).transpose(2, 3).reshape(batch, -1, 32, 16)
    )

    # Concatenate batches (L-major) into a flat buffer
    return rearranged.reshape(batch, -1).contiguous().view(-1)


def _pack_scale_factors_from_provided(
    sf_permuted: torch.Tensor, batch: int
) -> torch.Tensor:
    """
    Use the generator-provided blocked layout tensor, materialized to a packed 1D buffer.

    For this competition problem, `l` is always 1 in tests/benchmarks, so we can avoid any
    additional permutation to make L-major. If `batch > 1`, we conservatively permute L to front.
    """
    if sf_permuted.is_contiguous():
        t = sf_permuted
    else:
        t = sf_permuted.contiguous()

    if batch == 1:
        return t.view(-1)

    # General case (not hit by current nvfp4_gemm tests/benchmarks): make L-major.
    dims = list(range(t.dim()))
    perm = [dims[-1]] + dims[:-1]
    return t.permute(*perm).contiguous().view(-1)


def _sf_permuted_to_packed_1d_no_copy(
    sf_permuted: torch.Tensor, mn: int, k: int, l: int
) -> torch.Tensor:
    """
    Zero-copy view of the generator's scale-factor backing storage as a packed 1D tensor.

    reference.py creates the permuted tensor via:
      base = rand_int_tensor.to(fp8).permute(3,4,1,5,2,0)

    The underlying storage is contiguous in the *pre-permute* base shape:
      (L, ceil_div(MN,128), ceil_div(sf_k,4), 32, 4, 4)

    This function reconstructs that contiguous base view via as_strided(), then flattens it.
    """
    # sf_k = ceil_div(k, 16) ; k is divisible by 256 in this problem so sf_k == k//16
    sf_k = (k + 15) // 16
    rest_mn = (mn + 127) // 128
    rest_k = (sf_k + 3) // 4

    # Base contiguous layout (element strides) for shape (L, rest_mn, rest_k, 32, 4, 4)
    s5 = 1
    s4 = 4
    s3 = 16
    s2 = 32 * 4 * 4  # 512
    s1 = rest_k * s2
    s0 = rest_mn * s1

    base = sf_permuted.as_strided(
        size=(l, rest_mn, rest_k, 32, 4, 4),
        stride=(s0, s1, s2, s3, s4, s5),
        storage_offset=sf_permuted.storage_offset(),
    )
    return base.view(-1)


def compile_kernel():
    """
    Compile and load the CUDA extension (cached).

    Called by the eval harness in leaderboard mode to exclude compilation time.
    """
    global _mod
    if _mod is not None:
        return _mod

    cpp_src = r"""
    #include <torch/extension.h>
    #include <tuple>
    #include <cstdint>

    int nvfp4_gemm_run(
        uint64_t a_ptr,
        uint64_t b_ptr,
        uint64_t sfa_ptr,
        uint64_t sfb_ptr,
        uint64_t c_ptr,
        int m, int n, int k, int l);

    std::tuple<int, int> nvfp4_gemm_last_error();

    long long nvfp4_gemm_sfa_numel(int m, int n, int k, int l);
    long long nvfp4_gemm_sfb_numel(int m, int n, int k, int l);
    """

    cuda_src = r"""
    #include <cuda.h>
    #include <cuda_runtime.h>
    #include <cstdlib>
    #include <cstdint>
    #include <tuple>
    #include <torch/extension.h>
    #include <ATen/ATen.h>

    #include "cutlass/cutlass.h"
    #include "cute/tensor.hpp"
    #include "cutlass/tensor_ref.h"
    #include "cutlass/gemm/dispatch_policy.hpp"
    #include "cutlass/gemm/collective/collective_builder.hpp"
    #include "cutlass/epilogue/collective/collective_builder.hpp"
    #include "cutlass/detail/sm100_blockscaled_layout.hpp"
    #include "cutlass/gemm/device/gemm_universal_adapter.h"
    #include "cutlass/gemm/kernel/gemm_universal.hpp"
    #include "cutlass/util/packed_stride.hpp"

    using namespace cute;

    using ElementA     = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
    using LayoutATag   = cutlass::layout::RowMajor;
    constexpr int AlignmentA = 32;

    using ElementB     = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
    using LayoutBTag   = cutlass::layout::ColumnMajor;
    constexpr int AlignmentB = 32;

    using ElementD     = cutlass::half_t;
    using ElementC     = cutlass::half_t;
    using LayoutCTag   = cutlass::layout::RowMajor;
    using LayoutDTag   = cutlass::layout::RowMajor;
    constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
    constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;

    using ElementSF    = typename ElementA::ScaleFactorType;  // float_ue4m3_t
    using ElementAccumulator = float;
    using ArchTag            = cutlass::arch::Sm100;
    using OperatorClass      = cutlass::arch::OpClassBlockScaledTensorOp;

    // Variant C: NVF4 schedule + N-multicast cluster (1x2x1)
    // Closest Top20 family name:
    //   cutlass3x_sm100_bstensorop_gemm_ue4m3xe2m1_ue4m3xe2m1_f32_bf16_bf16_128x64x256_*_tnt_align32_o_vs16_1sm_epi_nosmem
    using MmaTileShape = Shape<_128, _64, _256>;
    using ClusterShape = Shape<_1, _2, _1>;

    using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
        ArchTag, OperatorClass,
        MmaTileShape, ClusterShape,
        cutlass::epilogue::collective::EpilogueTileAuto,
        ElementAccumulator, ElementAccumulator,
        ElementC, LayoutCTag, AlignmentC,
        ElementD, LayoutDTag, AlignmentD,
        cutlass::epilogue::NoSmemWarpSpecialized1Sm
    >::CollectiveOp;

    using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
        ArchTag, OperatorClass,
        ElementA, LayoutATag, AlignmentA,
        ElementB, LayoutBTag, AlignmentB,
        ElementAccumulator,
        MmaTileShape, ClusterShape,
        cutlass::gemm::collective::StageCountAutoCarveout<
            static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
        cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100
    >::CollectiveOp;

    using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
        Shape<int,int,int,int>,
        CollectiveMainloop,
        CollectiveEpilogue,
        void>;

    using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;

    using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
    using StrideA   = typename Gemm::GemmKernel::StrideA;
    using StrideB   = typename Gemm::GemmKernel::StrideB;
    using StrideC   = typename Gemm::GemmKernel::StrideC;
    using StrideD   = typename Gemm::GemmKernel::StrideD;

    namespace {
    static int g_last_cutlass_status = int(cutlass::Status::kSuccess);
    static int g_last_cuda_error = int(cudaSuccess);

    inline bool debug_sync_enabled() {
      const char* env = std::getenv("NVFP4_CUTLASS_DEBUG_SYNC");
      return env && env[0] == '1';
    }

    int nvfp4_gemm_impl(
        void* a_ptr, void* b_ptr,
        void* sfa_ptr, void* sfb_ptr,
        void* c_ptr,
        int m, int n, int k, int l,
        int* out_cutlass_status,
        int* out_cuda_error
    ) {
      if (out_cutlass_status) *out_cutlass_status = int(cutlass::Status::kSuccess);
      if (out_cuda_error) *out_cuda_error = int(cudaSuccess);

      auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, l});
      auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, l});
      auto stride_C = cutlass::make_cute_packed_stride(StrideC{}, {m, n, l});
      auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, l});

      auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(m, n, k, l));
      auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(m, n, k, l));

      auto* ptr_A = reinterpret_cast<typename ElementA::DataType*>(a_ptr);
      auto* ptr_B = reinterpret_cast<typename ElementB::DataType*>(b_ptr);
      auto* ptr_SFA = reinterpret_cast<ElementSF*>(sfa_ptr);
      auto* ptr_SFB = reinterpret_cast<ElementSF*>(sfb_ptr);
      auto* ptr_C = reinterpret_cast<ElementC*>(c_ptr);
      auto* ptr_D = reinterpret_cast<ElementD*>(c_ptr);

      typename Gemm::Arguments arguments{
          cutlass::gemm::GemmUniversalMode::kGemm,
          {m, n, k, l},
          {
              ptr_A, stride_A,
              ptr_B, stride_B,
              ptr_SFA, layout_SFA,
              ptr_SFB, layout_SFB
          },
          {
              {1.0f, 0.0f},
              ptr_C, stride_C,
              ptr_D, stride_D
          }
      };

      Gemm gemm;

      size_t workspace_size = Gemm::get_workspace_size(arguments);
      torch::Tensor workspace_tensor;
      void* workspace_ptr = nullptr;
      if (workspace_size > 0) {
        workspace_tensor = torch::empty(
          {static_cast<long long>(workspace_size)},
          torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA)
        );
        workspace_ptr = workspace_tensor.data_ptr();
        if (workspace_ptr == nullptr) {
          if (out_cuda_error) *out_cuda_error = int(cudaErrorMemoryAllocation);
          return -4;
        }
      }

      auto status = gemm.can_implement(arguments);
      if (out_cutlass_status) *out_cutlass_status = int(status);
      if (status != cutlass::Status::kSuccess) return -1;

      status = gemm.initialize(arguments, workspace_ptr);
      if (out_cutlass_status) *out_cutlass_status = int(status);
      if (status != cutlass::Status::kSuccess) {
        return -2;
      }

      status = gemm.run();
      if (out_cutlass_status) *out_cutlass_status = int(status);
      if (status != cutlass::Status::kSuccess) {
        return -3;
      }

      cudaError_t err = cudaGetLastError();
      if (out_cuda_error) *out_cuda_error = int(err);
      if (err != cudaSuccess) return -5;

      if (debug_sync_enabled()) {
        err = cudaDeviceSynchronize();
        if (out_cuda_error) *out_cuda_error = int(err);
        if (err != cudaSuccess) return -6;
      }

      return 0;
    }
    } // namespace

    int nvfp4_gemm_run(
        uint64_t a_ptr,
        uint64_t b_ptr,
        uint64_t sfa_ptr,
        uint64_t sfb_ptr,
        uint64_t c_ptr,
        int m, int n, int k, int l
    ) {
      int cutlass_status = int(cutlass::Status::kSuccess);
      int cuda_error = int(cudaSuccess);
      int rc = nvfp4_gemm_impl(
          reinterpret_cast<void*>(a_ptr),
          reinterpret_cast<void*>(b_ptr),
          reinterpret_cast<void*>(sfa_ptr),
          reinterpret_cast<void*>(sfb_ptr),
          reinterpret_cast<void*>(c_ptr),
          m, n, k, l,
          &cutlass_status,
          &cuda_error
      );
      g_last_cutlass_status = cutlass_status;
      g_last_cuda_error = cuda_error;
      return rc;
    }

    std::tuple<int, int> nvfp4_gemm_last_error() {
      return {g_last_cutlass_status, g_last_cuda_error};
    }

    long long nvfp4_gemm_sfa_numel(int m, int n, int k, int l) {
      auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(m, n, k, l));
      auto filtered = cute::filter_zeros(layout_SFA);
      return static_cast<long long>(cute::size(filtered));
    }

    long long nvfp4_gemm_sfb_numel(int m, int n, int k, int l) {
      auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(m, n, k, l));
      auto filtered = cute::filter_zeros(layout_SFB);
      return static_cast<long long>(cute::size(filtered));
    }
    """

    # Make the module name stable for caching but unique to this source.
    digest = hashlib.md5((cpp_src + cuda_src).encode("utf-8")).hexdigest()[:10]
    name = f"nvfp4_gemm_inline_{digest}"

    _mod = load_inline(
        name=name,
        cpp_sources=cpp_src,
        cuda_sources=cuda_src,
        functions=[
            "nvfp4_gemm_run",
            "nvfp4_gemm_last_error",
            "nvfp4_gemm_sfa_numel",
            "nvfp4_gemm_sfb_numel",
        ],
        extra_include_paths=_cutlass_paths(),
        extra_cflags=["-O3"],
        extra_cuda_cflags=[
            "-O3",
            "--use_fast_math",
            "--expt-relaxed-constexpr",
            "--expt-extended-lambda",
            "-std=c++17",
            "-w",
            "-gencode=arch=compute_100a,code=sm_100a",
            "--ptxas-options=--gpu-name=sm_100a",
            "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1",
            "-DCUTLASS_ARCH_MMA_SM100_SUPPORTED=1",
            "-DCUTLASS_ARCH_MMA_SM100A_SUPPORTED=1",
        ],
        extra_ldflags=["-lcuda", "-lcudart"],
        verbose=bool(int(os.getenv("NVFP4_EXT_VERBOSE", "0"))),
    )

    return _mod


def custom_kernel(data: input_t) -> output_t:
    a, b, sfa, sfb, _, _, c = data

    m, k_half, batch = a.shape
    k = k_half * 2
    n, _, _ = b.shape

    # Packing strategy:
    # - "provided_nocopy" (default): use the generator's sfa_permuted/sfb_permuted WITHOUT any materialization/caching.
    # - "provided": use the generator's sfa_permuted/sfb_permuted and materialize contiguously (copy).
    # - "blocked": compute the blocked packing from sfa/sfb (copy; mostly for validation).
    sf_packing = os.getenv("NVFP4_SF_PACKING", "provided_nocopy")

    if sf_packing == "provided_nocopy":
        sfa_permuted = data[4]
        sfb_permuted = data[5]
        sfa_buf = _sf_permuted_to_packed_1d_no_copy(sfa_permuted, m, k, batch)
        sfb_buf = _sf_permuted_to_packed_1d_no_copy(sfb_permuted, n, k, batch)
    elif sf_packing == "provided":
        sfa_permuted = data[4]
        sfb_permuted = data[5]
        sfa_buf = _pack_scale_factors_from_provided(sfa_permuted, batch)
        sfb_buf = _pack_scale_factors_from_provided(sfb_permuted, batch)
    elif sf_packing == "blocked":
        sfa_buf = _pack_scale_factors_blocked(sfa)
        sfb_buf = _pack_scale_factors_blocked(sfb)
    else:
        raise RuntimeError(
            f"Unknown NVFP4_SF_PACKING={sf_packing!r} "
            f"(expected 'provided_nocopy', 'provided', or 'blocked')"
        )

    mod = compile_kernel()

    if os.getenv("NVFP4_DEBUG_LAYOUT") == "1":
        expected_sfa = int(mod.nvfp4_gemm_sfa_numel(m, n, k, batch))
        expected_sfb = int(mod.nvfp4_gemm_sfb_numel(m, n, k, batch))
        if sfa_buf.numel() != expected_sfa:
            raise RuntimeError(
                f"SFA numel mismatch: got {sfa_buf.numel()} expected {expected_sfa}"
            )
        if sfb_buf.numel() != expected_sfb:
            raise RuntimeError(
                f"SFB numel mismatch: got {sfb_buf.numel()} expected {expected_sfb}"
            )

    rc = int(
        mod.nvfp4_gemm_run(
            a.data_ptr(),
            b.data_ptr(),
            sfa_buf.data_ptr(),
            sfb_buf.data_ptr(),
            c.data_ptr(),
            m,
            n,
            k,
            batch,
        )
    )
    if rc != 0:
        cutlass_status, cuda_error = mod.nvfp4_gemm_last_error()
        raise RuntimeError(
            f"nvfp4_gemm_run failed: rc={rc} cutlass_status={int(cutlass_status)} cuda_error={int(cuda_error)}"
        )

    return c

# Leaderboard (2025-12-21) - modal_harness.cli nvfp4_gemm
# Variant: KernelTmaWarpSpecialized1SmNvf4Sm100, MmaTile=128x64x256, Cluster=1x2x1, epilogue=NoSmemWarpSpecialized1Sm
# Verification: PASSED
# Geometric Mean: 12070.83 ns (12.071 us)
# ```
# [0] m: 128; n: 7168; k: 16384; l: 1; seed: 1111
#     mean: 19825.11 ns (19.825 us)
# [1] m: 128; n: 4096; k: 7168; l: 1; seed: 1111
#     mean: 10669.64 ns (10.670 us)
# [2] m: 128; n: 7168; k: 2048; l: 1; seed: 1111
#     mean: 8314.70 ns (8.315 us)
# ```
# Leaderboard (2025-12-21) - modal_harness.cli nvfp4_gemm
# Target kernel (v2/top20 #1): cutlass3x_sm100_bstensorop_gemm_ue4m3xe2m1_ue4m3xe2m1_f32_bf16_bf16_128x64x256_0x0x1_0_tnt_align32_o_vs16_1sm_epi_nosmem
# (Adapted for task output: C/D=fp16)
# Verification: PASSED
# Geometric Mean: 20237.25 ns (20.237 us)
# ```
# [0] m: 128; n: 7168; k: 16384; l: 1; seed: 1111
#     mean: 46085.12 ns (46.085 us)
# [1] m: 128; n: 4096; k: 7168; l: 1; seed: 1111
#     mean: 14983.02 ns (14.983 us)
# [2] m: 128; n: 7168; k: 2048; l: 1; seed: 1111
#     mean: 12003.13 ns (12.003 us)
# ```