@triton.autotune(configs=autotune_configs, key=['col_sparse_dim',
    'inner_dim', 'sparse_dim'])
@triton.jit
def output_sparse_matmul_kernel(lhs_ptr, rhs_ptr, out_ptr, expert_ends_ptr,
    row_dim: tl.constexpr, col_sparse_dim: tl.constexpr, inner_dim: tl.
    constexpr, sparse_dim: tl.constexpr, padded_sparse_dim: tl.constexpr,
    lhs_stride_row: tl.constexpr, lhs_stride_inner: tl.constexpr,
    rhs_stride_inner: tl.constexpr, rhs_stride_col: tl.constexpr,
    out_stride_row: tl.constexpr, out_stride_col: tl.constexpr, accumulate:
    tl.constexpr, block_size_row: tl.constexpr, block_size_col: tl.
    constexpr, block_size_inner: tl.constexpr, group_size_row: tl.constexpr):
    tl.static_assert(row_dim % block_size_row == 0)
    tl.static_assert(col_sparse_dim % block_size_col == 0)
    tl.static_assert(inner_dim % block_size_inner == 0)
    tl.static_assert(sparse_dim <= padded_sparse_dim)
    pid_row, pid_col = tl.swizzle2d(tl.program_id(axis=0), tl.program_id(
        axis=1), row_dim // block_size_row, col_sparse_dim //
        block_size_col, group_size_row)
    row_offset = pid_row * block_size_row
    col_sparse_offset = pid_col * block_size_col
    sparse_range = tl.arange(0, padded_sparse_dim)
    expert_ends = tl.load(expert_ends_ptr + sparse_range, mask=sparse_range <
        sparse_dim, other=row_dim)
    sparse_index = tl.sum((expert_ends <= row_offset).to(tl.int64))
    if sparse_index == sparse_dim:
        return
    col_dense_offset = col_sparse_offset + sparse_index * col_sparse_dim
    row_range = tl.arange(0, block_size_row)[:, None]
    col_range = tl.arange(0, block_size_col)[None, :]
    inner_range = tl.arange(0, block_size_inner)
    lhs_ptr += (row_offset + row_range) * lhs_stride_row + inner_range[None, :
        ] * lhs_stride_inner
    rhs_ptr += inner_range[:, None] * rhs_stride_inner + (col_dense_offset +
        col_range) * rhs_stride_col
    out_ptr += (row_offset + row_range) * out_stride_row + (col_sparse_offset +
        col_range) * out_stride_col
    out = tl.dot(tl.load(lhs_ptr), tl.load(rhs_ptr), out_dtype=tl.float32)
    for k in range(1, inner_dim // block_size_inner):
        lhs_ptr += block_size_inner * lhs_stride_inner
        rhs_ptr += block_size_inner * rhs_stride_inner
        out += tl.dot(tl.load(lhs_ptr), tl.load(rhs_ptr))
    if accumulate:
        out += tl.load(out_ptr)
    tl.store(out_ptr, out)
