Function matmul

src/fused_mm_sampling/tl_matmul.py:135–187

(a, b, activation="")

def matmul(a, b, activation=""):
    # Check constraints.
    assert a.shape[1] == b.shape[1], "Incompatible dimensions"
    assert a.is_contiguous(), "Matrix A must be contiguous"
    assert b.is_contiguous(), "Matrix B must be contiguous"
    M, K = a.shape  # noqa: N806
    N, K = b.shape  # noqa: N806

    # TMA requires the innermost (stride-1) dimension to be aligned to 16 bytes.
    # For bfloat16 (2 bytes per element) that means multiples of 8 elements.
    tma_align = 16 // a.element_size()
    if K % tma_align != 0:
        raise ValueError(
            f"K={K} is not a multiple of {tma_align}. "
            f"TMA descriptors require the innermost dimension to be aligned to 16 bytes."
        )
    if N % tma_align != 0:
        raise ValueError(
            f"N={N} is not a multiple of {tma_align}. "
            f"TMA descriptors require the innermost dimension to be aligned to 16 bytes."
        )

    # Pre-transpose B: [N, K] -> [K, N] contiguous.
    # TMA enforces strides[-1]==1, so we can't describe the transpose via strides.
    b_t = b.T.contiguous()
    # Allocates output.
    c = torch.empty((M, N), device=a.device, dtype=torch.bfloat16)

    # TMA descriptors require a global memory allocation
    def alloc_fn(size: int, alignment: int, stream: Optional[int]):
        return torch.empty(size, device="cuda", dtype=torch.int8)

    triton.set_allocator(alloc_fn)

    # 2D launch kernel with N first to match 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.
    def grid(meta):
        return (
            triton.cdiv(N, meta["BLOCK_SIZE_N"]),
            triton.cdiv(M, meta["BLOCK_SIZE_M"]),
        )

    matmul_kernel[grid](
        a,
        b_t,
        c,
        M,  # noqa: N803
        N,  # noqa: N803
        K,  # noqa: N803
        ACTIVATION=activation,  #
    )
    return c


def get_cublas():