Function matmul_kernel

src/fused_mm_sampling/tl_matmul.py:35–126

Kernel for computing the matmul C = A x B.T. A has shape (M, K), B is pre-transposed to (K, N), and C has shape (M, N). Uses device-side tensor descriptors (TMA) for efficient memory access. NOTE: tl.dot(a, b.T) does NOT work correctly with TMA-loaded blocks — .T only swaps the logi

(
    # Pointers to matrices
    a_ptr,  # [M, K]
    b_ptr,  # [K, N] (pre-transposed by wrapper)
    c_ptr,
    # Matrix dimensions
    M,  # noqa: N803
    N,  # noqa: N803
    K,  # noqa: N803
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr,  # noqa: N803
    BLOCK_SIZE_N: tl.constexpr,  # noqa: N803
    BLOCK_SIZE_K: tl.constexpr,  # noqa: N803
    GROUP_SIZE_M: tl.constexpr,  # noqa: N803
    ACTIVATION: tl.constexpr,  # noqa: N803
)

)
@triton.jit
def matmul_kernel(
    # Pointers to matrices
    a_ptr,  # [M, K]
    b_ptr,  # [K, N] (pre-transposed by wrapper)
    c_ptr,
    # Matrix dimensions
    M,  # noqa: N803
    N,  # noqa: N803
    K,  # noqa: N803
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr,  # noqa: N803
    BLOCK_SIZE_N: tl.constexpr,  # noqa: N803
    BLOCK_SIZE_K: tl.constexpr,  # noqa: N803
    GROUP_SIZE_M: tl.constexpr,  # noqa: N803
    ACTIVATION: tl.constexpr,  # noqa: N803
):
    """Kernel for computing the matmul C = A x B.T.
    A has shape (M, K), B is pre-transposed to (K, N), and C has shape (M, N).
    Uses device-side tensor descriptors (TMA) for efficient memory access.

    NOTE: tl.dot(a, b.T) does NOT work correctly with TMA-loaded blocks —
    .T only swaps the logical view without rearranging shared memory, but
    tensor core MMA instructions depend on physical layout. And TMA enforces
    strides[-1] == 1 so we can't describe the transpose via strides either.
    The wrapper pre-transposes B to [K, N] contiguous.
    """
    # -----------------------------------------------------------
    # Create device-side tensor descriptors for TMA
    # A: [M, K] row-major contiguous
    # B: [K, N] row-major contiguous (pre-transposed by wrapper)
    # C: [M, N] row-major contiguous
    a_desc = tl.make_tensor_descriptor(
        a_ptr,
        shape=[M, K],
        strides=[K, 1],
        block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K],
    )
    b_desc = tl.make_tensor_descriptor(
        b_ptr,
        shape=[K, N],
        strides=[N, 1],
        block_shape=[BLOCK_SIZE_K, BLOCK_SIZE_N],
    )
    c_desc = tl.make_tensor_descriptor(
        c_ptr,
        shape=[M, N],
        strides=[N, 1],
        block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_N],
    )

    # -----------------------------------------------------------
    # Map program ids `pid` to the block of C it should compute.
    # Process N dimension first, then M (matching fused kernel pattern).
    # This enables processing many N blocks for the same M block, allowing
    # A matrix (small M dimension) to be reused from L2 cache.
    pid_n = tl.program_id(axis=0)  # N dimension first (like vocab in fused kernel)
    pid_m = tl.program_id(axis=1)  # M dimension second (like hidden_states in fused kernel)