submission 115249

import torch
from task import input_t, output_t
from itertools import product

import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import make_ptr
import cutlass.utils.blockscaled_layout as blockscaled_utils

# =============================================================================
# AUTOTUNING CONFIGURATION
# =============================================================================
# Enable/disable autotuning (set to False to use heuristic selection)
AUTOTUNE_ENABLED = True

# Benchmark sizes - only autotune for these specific (m, k, l) combinations
BENCHMARK_SIZES = {
    (7168, 16384, 1),
    (4096, 7168, 8),
    (7168, 2048, 4),
}

# Benchmarking parameters
WARMUP_ITERS = 2
BENCHMARK_ITERS = 40

# Search space for autotuning: (M_tile, K_tile, threads_m, threads_k)
# Constraints:
#   - threads_m == M_tile (one thread per M element in tile)
#   - K_tile must divide K dimension evenly (typically powers of 2)
#   - threads_m * threads_k <= 1024 (max threads per block)
AUTOTUNE_SEARCH_SPACE = [
    # M_tile options
    [16, 32, 64, 128],
    # K_tile options
    [32, 64, 128, 256],
    # threads_k options (K-parallelism)
    [4, 8, 16, 32],
]

# Fallback configs for heuristic selection (when autotuning disabled)
FALLBACK_CONFIGS = {
    "large_k_single_batch": (64, 128, 64, 16),
    "large_k_batched": (64, 128, 64, 16),
    "medium": (32, 128, 32, 4),
    "small_k": (64, 64, 64, 8),
}


def select_fallback_config(m: int, k: int, l: int) -> tuple:
    """
    Heuristic configuration selection (used when autotuning is disabled).

    Args:
        m: Number of rows in A matrix
        k: Inner dimension (after x2 adjustment for FP4 packing)
        l: Batch size

    Returns:
        Config tuple (m_tile, k_tile, threads_m, threads_k)
    """
    if l == 1 and m > 4096:
        return FALLBACK_CONFIGS["large_k_single_batch"]
    elif k > 8192:
        return FALLBACK_CONFIGS["large_k_batched"]
    elif k < 4096:
        return FALLBACK_CONFIGS["small_k"]
    else:
        return FALLBACK_CONFIGS["medium"]


# =============================================================================
# DATA TYPES (shared across all configs)
# =============================================================================
ab_dtype = cutlass.Float4E2M1FN  # FP4 data type for A and B
sf_dtype = cutlass.Float8E4M3FN  # FP8 data type for scale factors
c_dtype = cutlass.Float16        # FP16 output type
acc_dtype = cutlass.Float32      # Float32 accumulator
sf_vec_size = 16                 # Scale factor block size (16 elements share one scale)


def ceil_div(a, b):
    return (a + b - 1) // b


# =============================================================================
# KERNEL FACTORY - Creates kernels for each configuration
# =============================================================================

def create_kernel(m_tile: int, k_tile: int, threads_m: int, threads_k: int):
    """
    Create a CuTe kernel with the specified configuration.

    Args:
        m_tile: Tile size for M dimension
        k_tile: Tile size for K dimension
        threads_m: Number of threads for M dimension
        threads_k: Number of threads for K dimension (K-parallelism)

    Returns:
        Tuple of (kernel_func, jit_func) for the given configuration
    """
    # Capture config in closure
    mma_tiler_mnk = (m_tile, 1, k_tile)

    @cute.kernel
    def kernel(
        mA_mkl: cute.Tensor,
        mB_nkl: cute.Tensor,
        mSFA_mkl: cute.Tensor,
        mSFB_nkl: cute.Tensor,
        mC_mnl: cute.Tensor,
    ):
        """
        CuTe kernel for NVFP4 block-scaled GEMV with warp shuffle reduction.

        Uses 2D thread block with coalesced memory access pattern:
        - tidx (x-dimension) maps to K for coalesced loads (K-major layout)
        - tidy (y-dimension) maps to M

        Reduction via warp shuffle instead of shared memory.
        """
        bidx, bidy, bidz = cute.arch.block_idx()
        tidx, tidy, _ = cute.arch.thread_idx()

        # Map threads: x → K dimension (coalesced), y → M dimension
        tk = tidx  # thread over K within tile_K
        tm = tidy  # thread over M within tile_M

        # Extract local tiles using hierarchical layouts
        gA_mkl = cute.local_tile(
            mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
        )
        gSFA_mkl = cute.local_tile(
            mSFA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
        )
        gB_nkl = cute.local_tile(
            mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
        )
        gSFB_nkl = cute.local_tile(
            mSFB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
        )
        gC_mnl = cute.local_tile(
            mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None)
        )

        # Output element for this thread/block
        tCgC = gC_mnl[tm, None, bidx, bidy, bidz]
        tCgC = cute.make_tensor(tCgC.iterator, 1)
        res = cute.zeros_like(tCgC, acc_dtype)

        # K-reduction loop - each thread handles a subset of K-tiles
        k_tile_cnt = gA_mkl.layout[3].shape

        # Hoist RMEM tensor allocations outside the loop
        # Use dummy tile (k=0) to get layout info - actual data loaded inside loop
        tAgA_dummy = gA_mkl[tm, None, bidx, 0, bidz]
        tBgB_dummy = gB_nkl[0, None, bidy, 0, bidz]
        tAgSFA_dummy = gSFA_mkl[tm, None, bidx, 0, bidz]

        tArA = cute.make_rmem_tensor_like(tAgA_dummy, c_dtype)
        tBrB = cute.make_rmem_tensor_like(tBgB_dummy, c_dtype)
        tABrAB = cute.make_rmem_tensor_like(tAgA_dummy, c_dtype)
        tArSFA = cute.make_rmem_tensor_like(tAgSFA_dummy, acc_dtype)
        tBrSFB = cute.make_rmem_tensor_like(tAgSFA_dummy, acc_dtype)
        tSFrSF = cute.make_rmem_tensor_like(tAgSFA_dummy, acc_dtype)

        # Each (tm, tk) thread processes K-tiles: k_tile_idx = tk, tk + threads_k, ...
        for k_tile_idx in cutlass.range(tk, k_tile_cnt, threads_k, unroll_full=True):
            # Load entire K-tile into RMEM (vectorized, handles FP4 packing)
            tAgA = gA_mkl[tm, None, bidx, k_tile_idx, bidz]
            tBgB = gB_nkl[0, None, bidy, k_tile_idx, bidz]
            tAgSFA = gSFA_mkl[tm, None, bidx, k_tile_idx, bidz]
            tBgSFB = gSFB_nkl[0, None, bidy, k_tile_idx, bidz]

            # Load and convert NVFP4/FP8 values into pre-allocated RMEM
            tArA.store(tAgA.load().to(c_dtype))
            tBrB.store(tBgB.load().to(c_dtype))
            tArSFA.store(tAgSFA.load().to(acc_dtype))
            tBrSFB.store(tBgSFB.load().to(acc_dtype))

            # Pre-multiply A*B and SFA*SFB in RMEM
            tABrAB.store(tArA.load() * tBrB.load())
            tSFrSF.store(tArSFA.load() * tBrSFB.load())

            # Accumulate over all K elements within this tile
            for i in cutlass.range_constexpr(k_tile):
                res += tABrAB[i] * tSFrSF[i]

        # =============================
        # Shared memory reduction
        # =============================

        # Allocate shared memory for reduction with stride padding to avoid bank conflicts
        allocator = cutlass.utils.SmemAllocator()
        smem_layout = cute.make_layout(
            (threads_m, threads_k),
            stride=(threads_k + 1, 1)  # +1 padding to avoid bank conflicts
        )
        shared_res = allocator.allocate_tensor(element_type=cutlass.Float32, layout=smem_layout)

        # Store partial result to shared memory
        shared_res[(tm, tk)] = res[0]
        cute.arch.sync_threads()

        # Reduction - only thread with tk==0 writes final result
        if tk == 0:
            out = cute.zeros_like(tCgC, acc_dtype)
            for i in cutlass.range_constexpr(threads_k):
                out += shared_res[(tm, i)]

            # Store final result as FP16
            tCgC.store(out.to(cutlass.Float16))
        return

    @cute.jit
    def jit_func(
        a_ptr: cute.Pointer,
        b_ptr: cute.Pointer,
        sfa_ptr: cute.Pointer,
        sfb_ptr: cute.Pointer,
        c_ptr: cute.Pointer,
        problem_size: tuple,
    ):
        """Host-side JIT function with 2D thread block."""
        m, _, k, l = problem_size

        # Create CuTe Tensors with proper layouts
        a_tensor = cute.make_tensor(
            a_ptr,
            cute.make_layout(
                (m, cute.assume(k, 32), l),
                stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
            ),
        )

        # N dimension padded to 128 for torch._scaled_mm compatibility
        n_padded_128 = 128
        b_tensor = cute.make_tensor(
            b_ptr,
            cute.make_layout(
                (n_padded_128, cute.assume(k, 32), l),
                stride=(cute.assume(k, 32), 1, cute.assume(n_padded_128 * k, 32)),
            ),
        )

        c_tensor = cute.make_tensor(
            c_ptr, cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m))
        )

        # Scale factor layouts matching Tensor Core requirements
        sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a_tensor.shape, sf_vec_size)
        sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)

        sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b_tensor.shape, sf_vec_size)
        sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)

        # Grid: (M_blocks, 1, L)
        grid = (
            cute.ceil_div(c_tensor.shape[0], m_tile),
            1,
            c_tensor.shape[2],
        )

        # Launch kernel with 2D thread block
        # Map x-dimension to K for coalesced loads (K-major layout)
        kernel(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch(
            grid=grid,
            block=[threads_k, threads_m, 1],
            cluster=(1, 1, 1),
        )
        return

    return kernel, jit_func


# =============================================================================
# KERNEL CACHE - Two-level caching for autotuning
# =============================================================================

# Level 1: Cache compiled kernels by configuration tuple
_compiled_kernel_cache = {}

# Level 2: Cache best kernel by input size (m, k, l)
_input_kernel_cache = {}


def get_cache_key(m_tile: int, k_tile: int, threads_m: int, threads_k: int) -> str:
    """Generate cache key for a configuration."""
    return f"{m_tile}x{k_tile}x{threads_m}x{threads_k}"


def get_compiled_kernel(config: tuple):
    """
    Get or compile a kernel for the specified configuration.

    Args:
        config: Tuple (m_tile, k_tile, threads_m, threads_k)

    Returns:
        Compiled kernel function
    """
    global _compiled_kernel_cache

    m_tile, k_tile, threads_m, threads_k = config
    cache_key = get_cache_key(m_tile, k_tile, threads_m, threads_k)

    if cache_key in _compiled_kernel_cache:
        return _compiled_kernel_cache[cache_key]

    # Create kernel and JIT function for this config
    _, jit_func = create_kernel(m_tile, k_tile, threads_m, threads_k)

    # Create CuTe pointers for compilation
    a_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
    b_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
    c_ptr = make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
    sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)
    sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)

    # Compile the kernel
    compiled = cute.compile(
        jit_func, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 0)
    )

    _compiled_kernel_cache[cache_key] = compiled
    return compiled


# =============================================================================
# BENCHMARKING - Measure kernel execution time
# =============================================================================

def benchmark_kernel(kernel_func, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size):
    """
    Benchmark a kernel configuration.

    Args:
        kernel_func: Compiled kernel function
        a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr: CuTe pointers
        problem_size: Tuple (m, n, k, l)

    Returns:
        Average execution time in microseconds
    """
    # Warmup
    for _ in range(WARMUP_ITERS):
        kernel_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size)
    torch.cuda.synchronize()

    # Benchmark
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)

    start_event.record()
    for _ in range(BENCHMARK_ITERS):
        kernel_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size)
    end_event.record()
    torch.cuda.synchronize()

    elapsed_ms = start_event.elapsed_time(end_event)
    return elapsed_ms / BENCHMARK_ITERS * 1000  # Return in microseconds


def is_valid_config(m_tile: int, k_tile: int, threads_k: int, m: int, k: int) -> bool:
    """
    Check if a configuration is valid for the given problem size.

    Args:
        m_tile: Tile size for M dimension
        k_tile: Tile size for K dimension
        threads_k: Number of K-parallel threads
        m: Problem M dimension
        k: Problem K dimension

    Returns:
        True if configuration is valid
    """
    threads_m = m_tile  # threads_m == m_tile constraint

    # Check thread count limit
    if threads_m * threads_k > 1024:
        return False

    # M dimension must be divisible by m_tile
    if m % m_tile != 0:
        return False

    # K dimension must be divisible by k_tile
    if k % k_tile != 0:
        return False

    return True


# =============================================================================
# AUTOTUNING - Find best configuration for input size
# =============================================================================

def autotune_kernel(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size):
    """
    Autotune to find the best kernel configuration for the given problem size.

    Performs exhaustive search over the configuration space, benchmarking
    each valid configuration and returning the best performer.

    Args:
        a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr: CuTe pointers
        problem_size: Tuple (m, n, k, l)

    Returns:
        Best compiled kernel function
    """
    m, n, k, l = problem_size

    best_kernel = None
    best_time = float("inf")
    best_config = None

    # Iterate through search space
    m_tiles, k_tiles, threads_k_options = AUTOTUNE_SEARCH_SPACE

    for m_tile, k_tile, threads_k in product(m_tiles, k_tiles, threads_k_options):
        threads_m = m_tile  # threads_m == m_tile constraint

        # Skip invalid configurations
        if not is_valid_config(m_tile, k_tile, threads_k, m, k):
            continue

        config = (m_tile, k_tile, threads_m, threads_k)

        try:
            # Get or compile kernel for this config
            compiled_kernel = get_compiled_kernel(config)

            # Benchmark
            cur_time = benchmark_kernel(
                compiled_kernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size
            )

            if cur_time < best_time:
                best_time = cur_time
                best_kernel = compiled_kernel
                best_config = config

        except Exception:
            # Skip configurations that fail to compile or run
            continue

    if best_kernel is None:
        # Fallback to heuristic if no valid config found
        fallback_config = select_fallback_config(m, k, l)
        best_kernel = get_compiled_kernel(fallback_config)
        best_config = fallback_config

    return best_kernel, best_config


def get_best_kernel(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size):
    """
    Get the best kernel for the given input size, using cached result if available.

    Args:
        a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr: CuTe pointers
        problem_size: Tuple (m, n, k, l)

    Returns:
        Best compiled kernel function
    """
    global _input_kernel_cache

    m, n, k, l = problem_size
    input_key = (m, k, l)

    if input_key in _input_kernel_cache:
        return _input_kernel_cache[input_key]

    # Only autotune for benchmark sizes, use fallback for tests
    if AUTOTUNE_ENABLED and input_key in BENCHMARK_SIZES:
        best_kernel, _ = autotune_kernel(
            a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size
        )
    else:
        # Use heuristic selection for tests or when autotuning disabled
        fallback_config = select_fallback_config(m, k, l)
        best_kernel = get_compiled_kernel(fallback_config)

    _input_kernel_cache[input_key] = best_kernel
    return best_kernel


def custom_kernel(data: input_t) -> output_t:
    """
    Execute NVFP4 GEMV with autotuning configuration selection.

    When AUTOTUNE_ENABLED=True, performs exhaustive search over the configuration
    space on first call for each unique (m, k, l) combination, then caches the
    best performing kernel for subsequent calls.

    When AUTOTUNE_ENABLED=False, uses heuristic configuration selection.

    Configuration search space can be adjusted in AUTOTUNE_SEARCH_SPACE.
    """
    a, b, _, _, sfa_permuted, sfb_permuted, c = data

    m, k_packed, l = a.shape
    k = k_packed * 2  # Torch uses e2m1_x2, so k is halved
    n = 1

    # Create CuTe pointers from torch tensor data
    a_ptr = make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    b_ptr = make_ptr(ab_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    sfa_ptr = make_ptr(
        sf_dtype, sfa_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
    )
    sfb_ptr = make_ptr(
        sf_dtype, sfb_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
    )

    problem_size = (m, n, k, l)

    # Get best kernel (autotuned or heuristic based on AUTOTUNE_ENABLED)
    compiled_func = get_best_kernel(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size)

    # Execute the kernel
    compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size)

    return c