submission 243406

import torch
import torch.nn.functional as F

def custom_kernel(data):
    """
    Optimized B200 Dual GEMM (Eager Mode).
    - Removes CUDA Graphs (causes 500 error).
    - Uses efficient vectorized permutation for scales to ensure zero-copy views in the loop.
    - Corrects scale layout to match torch._scaled_mm expectation.
    """
    # Unpack inputs
    a = data[0]
    b1 = data[1]
    b2 = data[2]
    
    # CuTe-layout scales: (32, 4, M_blk, 4, K_blk, L)
    # or similar spatial dims. We assume the spatial structure matches the 5 dims
    # that need permuting (2, 4, 0, 1, 3).
    sfa_perm = data[6]
    sfb1_perm = data[7]
    sfb2_perm = data[8]
    
    c = data[9]
    M, N, L = c.shape
    
    # Pre-allocate output buffers to avoid dynamic allocation overhead in loop
    out1 = torch.empty((M, N), dtype=torch.float32, device="cuda")
    out2 = torch.empty((M, N), dtype=torch.float32, device="cuda")
    
    # --- Vectorized Scale Preparation ---
    # Goal: Transform scales to (L, Flattened_Blocked_Size) so that scales[i] is contiguous.
    # Current assumption: sfa_perm has L as the LAST dimension (dim 5).
    # Target permutation for spatial dims: (2, 4, 0, 1, 3).
    # Combined permutation: (5, 2, 4, 0, 1, 3) -> moves L to front, fixes layout.
    
    # Use reshape to collapse spatial dims after permute.
    # We perform the permute+reshape once. This is ONE kernel call per scale tensor.
    
    sfa_prepared = sfa_perm.permute(5, 2, 4, 0, 1, 3).reshape(L, -1).contiguous()
    sfb1_prepared = sfb1_perm.permute(5, 2, 4, 0, 1, 3).reshape(L, -1).contiguous()
    sfb2_prepared = sfb2_perm.permute(5, 2, 4, 0, 1, 3).reshape(L, -1).contiguous()
    
    for i in range(L):
        # Slice: scales[i] is now (FlatSize,). 
        # Since we made it (L, ...) and contiguous, this slice is a pointer offset.
        # Zero copy.
        scale_a = sfa_prepared[i]
        scale_b1 = sfb1_prepared[i]
        scale_b2 = sfb2_prepared[i]
        
        # Hardware GEMM (A @ B1)
        torch._scaled_mm(
            a[:, :, i],
            b1[:, :, i].transpose(0, 1),
            scale_a,
            scale_b1,
            bias=None,
            out_dtype=torch.float32,
            out=out1
        )
        
        # Hardware GEMM (A @ B2)
        torch._scaled_mm(
            a[:, :, i],
            b2[:, :, i].transpose(0, 1),
            scale_a,
            scale_b2,
            bias=None,
            out_dtype=torch.float32,
            out=out2
        )
        
        # Fused Epilogue
        F.silu(out1, inplace=True)
        out1.mul_(out2)
        c[:, :, i].copy_(out1)

    return c