Function kernel_fn

kernels/fused_mlp.py:102–170

Entry point called by bench.py. Must match reference.fused_mlp_ref signature. SwiGLU MLP: hidden = activation(x @ w_gate.T) * (x @ w_up.T) out = hidden @ w_down.T Args: x: [batch, hidden_size] or [batch, seq_len, hidden_size] w_gate: [intermediate_size, hid

(
    x: torch.Tensor,
    w_gate: torch.Tensor,
    w_up: torch.Tensor,
    w_down: torch.Tensor,
    activation: str = "silu",
)

def kernel_fn(
    x: torch.Tensor,
    w_gate: torch.Tensor,
    w_up: torch.Tensor,
    w_down: torch.Tensor,
    activation: str = "silu",
) -> torch.Tensor:
    """
    Entry point called by bench.py. Must match reference.fused_mlp_ref signature.

    SwiGLU MLP:
      hidden = activation(x @ w_gate.T) * (x @ w_up.T)
      out = hidden @ w_down.T

    Args:
        x: [batch, hidden_size] or [batch, seq_len, hidden_size]
        w_gate: [intermediate_size, hidden_size]
        w_up: [intermediate_size, hidden_size]
        w_down: [hidden_size, intermediate_size]
        activation: "silu" or "gelu"
    """
    assert x.is_cuda

    # Handle multi-dim input
    orig_shape = x.shape
    if x.ndim > 2:
        x = x.view(-1, x.shape[-1])

    M, K = x.shape
    N, K2 = w_gate.shape
    assert K == K2, f"Hidden dim mismatch: x has {K}, w_gate has {K2}"
    assert w_up.shape == (N, K), f"w_up shape mismatch"

    hidden = torch.empty((M, N), device=x.device, dtype=x.dtype)

    BLOCK_SIZE_M = 64
    BLOCK_SIZE_N = 64
    BLOCK_SIZE_K = 32

    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))

    # W_gate and W_up are [N, K]. We access them as transposed: X[M,K] @ W^T[K,N]
    # So stride_wgk corresponds to stride along the K dimension (stride(1) of [N,K])
    # and stride_wgn corresponds to stride along N dimension (stride(0) of [N,K])
    fused_gate_up_kernel[grid](
        x,
        w_gate,
        w_up,
        hidden,
        M, N, K,
        x.stride(0), x.stride(1),
        w_gate.stride(1), w_gate.stride(0),  # transposed: K-stride, N-stride
        w_up.stride(1), w_up.stride(0),      # transposed: K-stride, N-stride
        hidden.stride(0), hidden.stride(1),
        USE_SILU=(activation == "silu"),
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K,