submission 99712

#!POPCORN leaderboard nvfp4_gemv
import os
import torch
from pathlib import Path
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

# B200 / sm90 target
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "9.0+PTX")


cpp_src = r"""
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <ATen/core/dispatch/Dispatcher.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_runtime.h>
#include <vector>

using torch::indexing::Slice;

static inline at::Tensor call_scaled_mm(
    const at::Tensor& A,
    const at::Tensor& BT,
    const at::Tensor& SA,
    const at::Tensor& SB
) {
  static auto op = c10::Dispatcher::singleton()
      .findSchemaOrThrow("aten::_scaled_mm", "")
      .typed<at::Tensor (const at::Tensor&, const at::Tensor&, const at::Tensor&, const at::Tensor&,
                         const c10::optional<at::Tensor>&, const c10::optional<at::Tensor>&,
                         c10::optional<c10::ScalarType>, bool)>();
  return op.call(A, BT, SA, SB, c10::nullopt, c10::nullopt,
                 c10::optional<c10::ScalarType>(at::kHalf), false);
}

void gemv_nvfp4_batched(
    at::Tensor a,
    at::Tensor b,
    at::Tensor sfa_perm,
    at::Tensor sfb_perm,
    at::Tensor c
) {
  const int64_t L = a.size(2);
  const int64_t N_b = b.size(0);
  auto device = a.device();
  at::cuda::CUDAGuard guard(device);
  auto parent_stream = at::cuda::getCurrentCUDAStream(device.index());

  // b_l: (L, N, K/2) view
  at::Tensor b_l = b.permute({2, 0, 1});
  // Flatten scales inside C++ to avoid Python-side overhead.
  at::Tensor sa = sfa_perm.permute({5, 2, 4, 0, 1, 3}).contiguous().view({L, -1});
  at::Tensor sb = sfb_perm.permute({5, 2, 4, 0, 1, 3}).contiguous().view({L, -1});

  if (L <= 8) {
      // Use worker streams with parent-stream waits to overlap small L.
      std::vector<cudaEvent_t> events;
      events.reserve(L);
      for (int64_t l = 0; l < L; ++l) {
          auto stream = at::cuda::getStreamFromPool(false, device.index());
          at::cuda::CUDAStreamGuard stream_guard(stream);

          at::Tensor A_ = a.index({Slice(), Slice(), l});
          at::Tensor BT_ = b_l.index({l}).transpose(0, 1);
          at::Tensor SA_ = sa.index({l});
          at::Tensor SB_ = sb.index({l});

          at::Tensor out = call_scaled_mm(A_, BT_, SA_, SB_);

          if (N_b == 1) {
              c.index_put_({Slice(), Slice(), l}, out);
          } else {
              c.index_put_({Slice(), 0, l}, out.index({Slice(), 0}));
          }

          cudaEvent_t ev;
          cudaEventCreateWithFlags(&ev, cudaEventDisableTiming);
          cudaEventRecord(ev, stream.stream());
          cudaStreamWaitEvent(parent_stream.stream(), ev, 0);
          events.push_back(ev);
      }
      for (auto ev : events) cudaEventDestroy(ev);
  } else {
      // Fallback: serial loop for other L on parent stream.
      for (int64_t l = 0; l < L; ++l) {
          at::Tensor A_ = a.index({Slice(), Slice(), l});
          at::Tensor BT_ = b_l.index({l}).transpose(0, 1);
          at::Tensor SA_ = sa.index({l});
          at::Tensor SB_ = sb.index({l});
          at::Tensor out = call_scaled_mm(A_, BT_, SA_, SB_);
          if (N_b == 1) c.index_put_({Slice(), Slice(), l}, out);
          else c.index_put_({Slice(), 0, l}, out.index({Slice(), 0}));
      }
  }
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("gemv_nvfp4_batched", &gemv_nvfp4_batched, "Batched NVFP4 GEMV");
}
"""

# Cache build in local dir to avoid repeated compile.
build_dir = Path(__file__).resolve().parent / ".torch_extensions"
build_dir.mkdir(parents=True, exist_ok=True)

mod = load_inline(
    name="nvfp4_gemv_cpp_stream",
    cpp_sources=cpp_src,
    with_cuda=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "--relocatable-device-code=false"],
    build_directory=str(build_dir),
    verbose=False,
)


@torch.no_grad()
def custom_kernel(data: input_t) -> output_t:
    a, b, _, _, sfa_perm, sfb_perm, c = data

    if not torch.cuda.is_available():
        return c

    mod.gemv_nvfp4_batched(a, b, sfa_perm, sfb_perm, c)
    torch.cuda.synchronize()
    return c