submission 755021

import torch
import triton
from triton import language as tl
from typing import TypeVar

# ----------------------------------------------------------------------
# Type aliases
# ----------------------------------------------------------------------
input_t = TypeVar("input_t", bound=torch.Tensor)
output_t = TypeVar("output_t", bound=tuple[torch.Tensor, torch.Tensor])

fp4x2 = torch.float4_e2m1fn_x2
fp8_e8m0 = torch.float8_e8m0fnu
bf16 = torch.bfloat16

# ----------------------------------------------------------------------
# Buffer cache – avoids re‑allocation on repeated calls
# ----------------------------------------------------------------------
_QUANT_BUFFER_CACHE: dict[tuple[int, int, int, int], tuple[torch.Tensor, torch.Tensor]] = {}

def _get_quant_buffers(x: torch.Tensor, m: int, n: int, scale_n: int) -> tuple[torch.Tensor, torch.Tensor]:
    """Allocate (or reuse) the output buffers for the quantized tensor and its block scales."""
    device_idx = x.device.index if x.device.index is not None else 0
    key = (device_idx, m, n, scale_n)
    buffers = _QUANT_BUFFER_CACHE.get(key)
    if buffers is None:
        # FP4 packed tensor (2 FP4 values per uint8)
        x_fp4 = torch.empty((m, n // 2), dtype=torch.uint8, device=x.device)
        # Block‑scale buffer is padded to a multiple of 256 rows and a multiple of 8 columns.
        # Initialise to 127 (zero exponent) so padded entries are already correct.
        blockscale_e8m0 = torch.full((triton.cdiv(m, 256) * 256, scale_n), 127, dtype=torch.uint8, device=x.device)
        buffers = (x_fp4, blockscale_e8m0)
        _QUANT_BUFFER_CACHE[key] = buffers
    return buffers

# ----------------------------------------------------------------------
# Triton kernel – MXFP4 quantization + shuffled block‑scale storage
# ----------------------------------------------------------------------
@triton.jit
def _dynamic_mxfp4_quant_kernel(
    x_ptr,               # bfloat16 input
    x_fp4_ptr,           # uint8 output (packed FP4x2)
    bs_ptr,              # uint8 block‑scale output (FP8_E8M0)
    stride_x_m,
    stride_x_n,
    stride_x_fp4_m,
    stride_x_fp4_n,
    M: tl.constexpr,
    N: tl.constexpr,
    scaleN_valid: tl.constexpr,   # number of valid 32‑col blocks per row
    scaleN_pad: tl.constexpr,     # column padding (multiple of 8) required by shuffle layout
    BLOCK_SIZE: tl.constexpr,
    MXFP4_QUANT_BLOCK_SIZE: tl.constexpr,
):
    pid_m = tl.program_id(0)   # block row
    pid_n = tl.program_id(1)   # block column (scale block)

    # Cast strides to 64‑bit for safety
    stride_x_m = tl.cast(stride_x_m, tl.int64)
    stride_x_n = tl.cast(stride_x_n, tl.int64)
    stride_x_fp4_m = tl.cast(stride_x_fp4_m, tl.int64)
    stride_x_fp4_n = tl.cast(stride_x_fp4_n, tl.int64)

    # -----------------------------------------------------------------
    # Load a tile of the input matrix (bfloat16 → float32)
    # -----------------------------------------------------------------
    x_offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    x_offs_n = pid_n * MXFP4_QUANT_BLOCK_SIZE + tl.arange(0, MXFP4_QUANT_BLOCK_SIZE)
    x_offs = x_offs_m[:, None] * stride_x_m + x_offs_n[None, :] * stride_x_n
    x_mask = (x_offs_m < M)[:, None] & (x_offs_n < N)[None, :]
    x = tl.load(x_ptr + x_offs, mask=x_mask).to(tl.float32)

    # -----------------------------------------------------------------
    # Per‑row scaling (FP8‑E8M0 unbiased exponent)
    # -----------------------------------------------------------------
    amax = tl.max(tl.abs(x), axis=1, keep_dims=True)
    amax = amax.to(tl.int32, bitcast=True)
    amax = (amax + 0x200000).to(tl.uint32, bitcast=True) & 0xFF800000
    amax = amax.to(tl.float32, bitcast=True)
    scale_e8m0_unbiased = tl.log2(amax).floor() - 2
    scale_e8m0_unbiased = tl.clamp(scale_e8m0_unbiased, min=-127, max=127)
    quant_scale = tl.exp2(-scale_e8m0_unbiased)
    qx = x * quant_scale
    bs_e8m0 = scale_e8m0_unbiased.to(tl.uint8) + 127

    # -----------------------------------------------------------------
    # FP4 conversion constants
    # -----------------------------------------------------------------
    EXP_BIAS_FP32: tl.constexpr = 127
    EXP_BIAS_FP4: tl.constexpr = 1
    EBITS_F32: tl.constexpr = 8
    EBITS_FP4: tl.constexpr = 2
    MBITS_F32: tl.constexpr = 23
    MBITS_FP4: tl.constexpr = 1

    max_normal: tl.constexpr = 6   # 2^6 is the largest normal FP4 value
    min_normal: tl.constexpr = 1   # 2^1 is the smallest normal FP4 value

    # -----------------------------------------------------------------
    # Bit‑level FP4 conversion (normal, denormal, saturate)
    # -----------------------------------------------------------------
    qx = qx.to(tl.uint32, bitcast=True)
    s = qx & 0x80000000               # sign bit mask
    qx = qx ^ s                       # clear sign for magnitude processing

    qx_fp32 = qx.to(tl.float32, bitcast=True)

    saturate_mask = qx_fp32 >= max_normal
    denormal_mask = (~saturate_mask) & (qx_fp32 < min_normal)
    normal_mask   = ~ (saturate_mask | denormal_mask)

    # ----- denormal path -----
    denorm_exp: tl.constexpr = ((EXP_BIAS_FP32 - EXP_BIAS_FP4) + (MBITS_F32 - MBITS_FP4) + 1)
    denorm_mask_int: tl.constexpr = denorm_exp << MBITS_F32
    denorm_mask_float: tl.constexpr = tl.cast(denorm_mask_int, tl.float32, bitcast=True)

    denormal_x = qx_fp32 + denorm_mask_float
    denormal_x = denormal_x.to(tl.uint32, bitcast=True)
    denormal_x = denormal_x - denorm_mask_int
    denormal_x = denormal_x.to(tl.uint8)

    # ----- normal path -----
    normal_x = qx
    mant_odd = (normal_x >> (MBITS_F32 - MBITS_FP4)) & 1
    val_to_add = ((EXP_BIAS_FP4 - EXP_BIAS_FP32) << MBITS_F32) + (1 << 21) - 1
    normal_x = normal_x + val_to_add
    normal_x = normal_x + mant_odd
    normal_x = normal_x >> (MBITS_F32 - MBITS_FP4)
    normal_x = normal_x.to(tl.uint8)

    # ----- combine paths (default saturates to 0x7) -----
    e2m1_value = tl.full((BLOCK_SIZE, MXFP4_QUANT_BLOCK_SIZE), 0x7, dtype=tl.uint8)
    e2m1_value = tl.where(normal_mask, normal_x, e2m1_value)
    e2m1_value = tl.where(denormal_mask, denormal_x, e2m1_value)

    # ----- sign insertion -----
    sign_lp = s >> (MBITS_F32 + EBITS_F32 - MBITS_FP4 - EBITS_FP4)
    sign_lp = sign_lp.to(tl.uint8)
    e2m1_value = e2m1_value | sign_lp

    # -----------------------------------------------------------------
    # Pack two FP4 values into a single uint8 (FP4x2 layout)
    # -----------------------------------------------------------------
    e2m1_value = tl.reshape(e2m1_value, [BLOCK_SIZE, MXFP4_QUANT_BLOCK_SIZE // 2, 2])
    evens, odds = tl.split(e2m1_value)
    out_tensor = evens | (odds << 4)

    # -----------------------------------------------------------------
    # Store the packed FP4x2 tensor
    # -----------------------------------------------------------------
    out_offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    out_offs_n = pid_n * (MXFP4_QUANT_BLOCK_SIZE // 2) + tl.arange(0, MXFP4_QUANT_BLOCK_SIZE // 2)
    out_offs = out_offs_m[:, None] * stride_x_fp4_m + out_offs_n[None, :] * stride_x_fp4_n
    out_mask = (out_offs_m < M)[:, None] & (out_offs_n < (N // 2))[None, :]
    tl.store(x_fp4_ptr + out_offs, out_tensor, mask=out_mask)

    # -----------------------------------------------------------------
    # Store block‑scale (shuffled layout)
    # -----------------------------------------------------------------
    bs_offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

    # Shuffle index computation using bit‑wise ops
    bs_offs_0 = bs_offs_m[:, None] >> 5          # //32
    bs_offs_1 = bs_offs_m[:, None] & 31          # %32
    bs_offs_2 = bs_offs_1 & 15                   # %16
    bs_offs_1 = bs_offs_1 >> 4                    # //16
    bs_offs_3 = pid_n >> 3                       # //8
    bs_offs_4 = pid_n & 7                        # %8
    bs_offs_5 = bs_offs_4 & 3                     # %4
    bs_offs_4 = bs_offs_4 >> 2                    # //4

    bs_offs = (
        bs_offs_1
        + (bs_offs_4 << 1)
        + (bs_offs_2 << 2)
        + (bs_offs_5 << 6)
        + (bs_offs_3 << 8)
        + (bs_offs_0 * (2 * 16 * scaleN_pad))
    )

    row_valid = bs_offs_m < M
    col_valid = pid_n < scaleN_valid
    bs_mask = row_valid[:, None] & col_valid

    tl.store(bs_ptr + bs_offs, bs_e8m0, mask=bs_mask)

def dynamic_mxfp4_quant(x: torch.Tensor, scaling_mode: str = "even") -> tuple[torch.Tensor, torch.Tensor]:
    """Quantize a bfloat16 matrix to packed FP4x2 + block‑scale FP8_E8M0."""
    M, N = x.shape

    # Choose a row block size that balances launch overhead and occupancy
    BLOCK_SIZE = 32 if M <= 64 else 64
    MXFP4_QUANT_BLOCK_SIZE = 32

    # Number of 32‑col blocks (valid) and padded column count for shuffle layout
    scaleN_valid = triton.cdiv(N, MXFP4_QUANT_BLOCK_SIZE)
    scaleN_pad = triton.cdiv(scaleN_valid, 8) * 8

    # Allocate (or reuse) output buffers
    x_fp4, blockscale_e8m0 = _get_quant_buffers(x, M, N, scaleN_pad)

    # Grid: one program per BLOCK_SIZE rows, one per valid scale block column
    grid = (triton.cdiv(M, BLOCK_SIZE), scaleN_valid)

    _dynamic_mxfp4_quant_kernel[grid](
        x,
        x_fp4,
        blockscale_e8m0,
        *x.stride(),
        *x_fp4.stride(),
        M,
        N,
        scaleN_valid,
        scaleN_pad,
        BLOCK_SIZE,
        MXFP4_QUANT_BLOCK_SIZE,
        num_warps=4,
        num_stages=1,
    )

    return (x_fp4.view(fp4x2), blockscale_e8m0.view(fp8_e8m0))

def quant_mxfp4(A: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    A_q, A_scale_sh = dynamic_mxfp4_quant(A)
    return A_q, A_scale_sh

def custom_kernel(data: input_t) -> output_t:
    A, _, _, B_shuffle, B_scale_sh = data

    A_q, A_scale_sh = quant_mxfp4(A)
    import aiter
    out_gemm = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=bf16,
        bpreshuffle=True,
    )
    return out_gemm