Function fused_gate_up_kernel

kernels/fused_mlp.py:26–99

Fused kernel: computes activation(X @ W_gate^T) * (X @ W_up^T). W_gate and W_up are [intermediate_size, hidden_size] (transposed access). X is [M, K], output is [M, N] where N = intermediate_size, K = hidden_size.

(
    X_ptr,
    W_gate_ptr,
    W_up_ptr,
    Out_ptr,
    M, N, K,
    stride_xm, stride_xk,
    stride_wgk, stride_wgn,
    stride_wuk, stride_wun,
    stride_om, stride_on,
    USE_SILU: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
)

@triton.jit
def fused_gate_up_kernel(
    X_ptr,
    W_gate_ptr,
    W_up_ptr,
    Out_ptr,
    M, N, K,
    stride_xm, stride_xk,
    stride_wgk, stride_wgn,
    stride_wuk, stride_wun,
    stride_om, stride_on,
    USE_SILU: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
):
    """
    Fused kernel: computes activation(X @ W_gate^T) * (X @ W_up^T).
    W_gate and W_up are [intermediate_size, hidden_size] (transposed access).
    X is [M, K], output is [M, N] where N = intermediate_size, K = hidden_size.
    """
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)

    # Pointers for X
    x_ptrs = X_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk

    # Pointers for W_gate (shape [K, N] after transpose -- stored as [N, K])
    wg_ptrs = W_gate_ptr + offs_k[:, None] * stride_wgk + offs_n[None, :] * stride_wgn
    # Pointers for W_up
    wu_ptrs = W_up_ptr + offs_k[:, None] * stride_wuk + offs_n[None, :] * stride_wun

    # Accumulators
    acc_gate = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    acc_up = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    for k_start in range(0, K, BLOCK_SIZE_K):
        k_offs = k_start + offs_k
        # Load X tile
        x_mask = (offs_m[:, None] < M) & (k_offs[None, :] < K)
        x = tl.load(x_ptrs, mask=x_mask, other=0.0)

        # Load W_gate tile
        wg_mask = (k_offs[:, None] < K) & (offs_n[None, :] < N)
        wg = tl.load(wg_ptrs, mask=wg_mask, other=0.0)

        # Load W_up tile
        wu_mask = (k_offs[:, None] < K) & (offs_n[None, :] < N)
        wu = tl.load(wu_ptrs, mask=wu_mask, other=0.0)

        acc_gate += tl.dot(x, wg)
        acc_up += tl.dot(x, wu)

        x_ptrs += BLOCK_SIZE_K * stride_xk
        wg_ptrs += BLOCK_SIZE_K * stride_wgk