submission 84465

"""
Optimization 39: PyTorch Multi-Batch Optimization

Hypothesis: Maybe we can use PyTorch for multi-batch too if we optimize it properly
- Pre-convert all scales at once
- Use batched operations
- Minimize Python loop overhead

opt36 results: 48.2/89.1/30.7 µs
Goal: Improve Bench 1 and 2 using smarter PyTorch batching

Submit using:
popcorn-cli submit --gpu NVIDIA --leaderboard nvfp4_gemv --mode leaderboard --no-tui opt39_pytorch_multi_batch.py 2>&1 | tee opt39_submission.log
"""

import torch
from task import input_t, output_t

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


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


# Compiled scale conversion
@torch.compile
def to_blocked_gpu_compiled(input_matrix):
    """Compiled GPU-optimized blocked format conversion"""
    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()


@torch.compile(fullgraph=False)
def pytorch_multi_batch_optimized(a, b, scale_a_list, scale_b_list, c):
    """Optimized multi-batch using PyTorch with pre-converted scales"""
    _, _, l = c.shape

    for batch_idx in range(l):
        res = torch._scaled_mm(
            a[:, :, batch_idx],
            b[:, :, batch_idx].transpose(0, 1),
            scale_a_list[batch_idx],
            scale_b_list[batch_idx],
            bias=None,
            out_dtype=torch.float16,
        )
        c[:, 0, batch_idx] = res[:, 0]

    return c


def pytorch_path(a, b, sfa_cpu, sfb_cpu, c):
    """PyTorch path with batch optimization"""
    _, _, l = c.shape

    # Move to GPU once
    sfa_gpu = sfa_cpu.cuda()
    sfb_gpu = sfb_cpu.cuda()

    # Pre-convert all scales
    scale_a_list = []
    scale_b_list = []
    for batch_idx in range(l):
        scale_a = to_blocked_gpu_compiled(sfa_gpu[:, :, batch_idx])
        scale_b = to_blocked_gpu_compiled(sfb_gpu[:, :, batch_idx])
        scale_a_list.append(scale_a)
        scale_b_list.append(scale_b)

    # Run batched computation
    return pytorch_multi_batch_optimized(a, b, scale_a_list, scale_b_list, c)


# CuTeDSL fallback (proven K=256 configuration)
_cutedsl_compiled = None


def get_cutedsl_kernel():
    """Get compiled CuTeDSL kernel (K=256 - proven optimal)"""
    global _cutedsl_compiled

    if _cutedsl_compiled is not None:
        return _cutedsl_compiled

    mma_tiler_mnk = (128, 1, 256)
    ab_dtype = cutlass.Float4E2M1FN
    sf_dtype = cutlass.Float8E4M3FN
    c_dtype = cutlass.Float16
    sf_vec_size = 16
    threads_per_cta = 128

    @cute.kernel
    def kernel(
        mA_mkl: cute.Tensor,
        mB_nkl: cute.Tensor,
        mSFA_mkl: cute.Tensor,
        mSFB_nkl: cute.Tensor,
        mC_mnl: cute.Tensor,
    ):
        bidx, bidy, bidz = cute.arch.block_idx()
        tidx, _, _ = cute.arch.thread_idx()

        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)
        )

        tCgC = gC_mnl[tidx, None, bidx, bidy, bidz]
        tCgC = cute.make_tensor(tCgC.iterator, 1)
        res = cute.zeros_like(tCgC, cutlass.Float32)

        k_tile_cnt = gA_mkl.layout[3].shape
        for k_tile_idx in range(k_tile_cnt):
            tAgA = gA_mkl[tidx, None, bidx, k_tile_idx, bidz]
            tBgB = gB_nkl[0, None, bidy, k_tile_idx, bidz]
            tAgSFA = gSFA_mkl[tidx, None, bidx, k_tile_idx, bidz]
            tBgSFB = gSFB_nkl[0, None, bidy, k_tile_idx, bidz]

            tArA = cute.make_rmem_tensor_like(tAgA, cutlass.Float32)
            tBrB = cute.make_rmem_tensor_like(tBgB, cutlass.Float32)
            tArSFA = cute.make_rmem_tensor_like(tAgSFA, cutlass.Float32)
            tBrSFB = cute.make_rmem_tensor_like(tBgSFB, cutlass.Float32)

            tArA.store(tAgA.load().to(cutlass.Float32))
            tBrB.store(tBgB.load().to(cutlass.Float32))
            tArSFA.store(tAgSFA.load().to(cutlass.Float32))
            tBrSFB.store(tBgSFB.load().to(cutlass.Float32))

            for i in cutlass.range_constexpr(256):
                res += tArA[i] * tArSFA[i] * tBrB[i] * tBrSFB[i]

        tCgC.store(res.to(cutlass.Float16))
        return

    @cute.jit
    def jit_kernel(
        a_ptr: cute.Pointer,
        b_ptr: cute.Pointer,
        sfa_ptr: cute.Pointer,
        sfb_ptr: cute.Pointer,
        c_ptr: cute.Pointer,
        problem_size: tuple,
    ):
        m, _, k, l = problem_size
        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_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))
        )

        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 = (
            cute.ceil_div(c_tensor.shape[0], 128),
            1,
            c_tensor.shape[2],
        )

        kernel(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch(
            grid=grid,
            block=[threads_per_cta, 1, 1],
            cluster=(1, 1, 1),
        )
        return

    # Compile
    ab_dtype = cutlass.Float4E2M1FN
    sf_dtype = cutlass.Float8E4M3FN
    c_dtype = cutlass.Float16

    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)

    _cutedsl_compiled = cute.compile(jit_kernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 0))
    return _cutedsl_compiled


def cutedsl_path(a, b, sfa_permuted, sfb_permuted, c):
    """CuTeDSL path with K=256"""
    m, k_packed, l = a.shape
    k = k_packed * 2
    n = 1

    compiled_kernel = get_cutedsl_kernel()

    ab_dtype = cutlass.Float4E2M1FN
    sf_dtype = cutlass.Float8E4M3FN
    c_dtype = cutlass.Float16

    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)

    compiled_kernel(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l))
    return c


def custom_kernel(data: input_t) -> output_t:
    """
    Hybrid with adaptive strategy:
    - Try PyTorch for small batch counts
    - Fall back to CuTeDSL for larger batches
    """
    a, b, sfa_cpu, sfb_cpu, sfa_permuted, sfb_permuted, c = data

    _, _, l = c.shape

    # Adaptive threshold: PyTorch is good up to l=2, CuTeDSL better for l>2
    if l <= 2:
        return pytorch_path(a, b, sfa_cpu, sfb_cpu, c)
    else:
        return cutedsl_path(a, b, sfa_permuted, sfb_permuted, c)