submission 738630

import torch
import triton
import triton.language as tl
from task import input_t, output_t
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import fused_moe

# ---------------------------------------------------------------------------
# MoE Cache-Resident Fusion Strategy
#
# Because the official aiter.fused_moe triggers multiple kernels
# (routing, block quantize on activations, multiple GEMMs), the same
# cache eviction policy we used earlier can be applied. We can manually
# load the input `hidden_states` into a tiny cache-pinning kernel right
# before handing it over to fused_moe.
# ---------------------------------------------------------------------------

@triton.jit
def l2_pin_hidden_states(
    in_ptr, out_ptr,
    N_ELEMENTS,
    BLOCK_SIZE: tl.constexpr
):
    pid = tl.program_id(0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offs < N_ELEMENTS

    # Load from DRAM
    val = tl.load(in_ptr + offs, mask=mask)

    # Store back to DRAM but WITH EVICT_LAST policy!
    # This forces the MI355X's massive L2 Cache to keep this activation tensor warm
    # so the subsequent fused_moe kernel hits it at 5TB/s instead of 1TB/s HBM bandwidth.
    tl.store(out_ptr + offs, val, mask=mask, eviction_policy="evict_last")


def custom_kernel(data: input_t) -> output_t:
    (
        hidden_states,
        gate_up_weight,
        down_weight,
        gate_up_weight_scale,
        down_weight_scale,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        gate_up_weight_scale_shuffled,
        down_weight_scale_shuffled,
        topk_weights,
        topk_ids,
        config,
    ) = data

    hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
    intermediate_pad = config["d_expert_pad"] - config["d_expert"]

    # 1. Force the activations into L2 cache
    M, d_hidden = hidden_states.shape
    total_elements = M * d_hidden

    # Pre-allocate a pinned buffer if needed, or overwrite (here we just use a dummy buffer
    # to trigger the read/write if necessary, but we can actually write back to the SAME tensor!)

    BLOCK_SIZE = 1024
    grid = (triton.cdiv(total_elements, BLOCK_SIZE),)

    # We do a cache-pinning pass over hidden_states.
    l2_pin_hidden_states[grid](
        hidden_states, hidden_states,
        total_elements,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=4
    )

    # 2. Call the highly optimized aiter C++ MoE kernel
    output = fused_moe(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        expert_mask=None,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,
        doweight_stage1=False,
        w1_scale=gate_up_weight_scale_shuffled,
        w2_scale=down_weight_scale_shuffled,
        a1_scale=None,
        a2_scale=None,
        hidden_pad=hidden_pad,
        intermediate_pad=intermediate_pad,
    )

    return output