sgl-project/sglang · error · ValueError

gemm_ar: M={m} outside [1, {MAX_TOKENS}]

Error message

gemm_ar: M={m} outside [1, {MAX_TOKENS}]

What it means

o_proj_gemm_ar buckets the token count M into fixed cell sizes (8,16,32,...,512) for its all-reduce-fused GEMM; MAX_TOKENS caps the largest cell at 512. M larger than every cell (or M<=0) falls through the loop and raises.

Source

Thrown at python/sglang/kernels/ops/kimi_k3/gemm_ar.py:191

@register_custom_op(mutates_args=["out", "epochs"])
def _gemm_ar_op(
    k: int,
    world_size: int,
    out: torch.Tensor,
    x: torch.Tensor,
    weight: torch.Tensor,
    gather: torch.Tensor,
    epochs: torch.Tensor,
    my_rank: int,
) -> None:
    _module_with_bases(k, world_size).run(out, x, weight, gather, epochs, my_rank)


def _cell_of(m: int) -> int:
    for c in (8, 16, 32, 64, 128, 256, 512):
        if m <= c:
            return c
    raise ValueError(f"gemm_ar: M={m} outside [1, {MAX_TOKENS}]")


def o_proj_gemm_ar(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
    """Fully reduced ``sum_r x_r @ weight_r^T`` on every rank, one kernel.

    ``x`` is the TP-local [M, K] o_proj input, ``weight`` the TP-local
    [7168, K] o_proj weight shard. Caller checked :func:`fits`; all ranks
    call in lockstep with the same M.
    """
    state = _STATE
    assert state is not None
    m = x.shape[0]
    cell = _cell_of(m)
    out = torch.empty((cell, N), dtype=torch.bfloat16, device=x.device)
    _gemm_ar_op(
        weight.shape[1],
        state.world_size,
        out,

View on GitHub (pinned to 0132848349)

Solutions

1. Route batches with M > MAX_TOKENS to the regular GEMM + all-reduce path instead of o_proj_gemm_ar
2. Chunk the input into <=MAX_TOKENS row blocks and call the kernel per chunk
3. If the kernel is intended for larger M, raise MAX_TOKENS and extend the cell tuple accordingly (requires revalidating perf)

Example fix

# before
y = o_proj_gemm_ar(x, w)  # x.shape[0] == 640 > 512
# after
if x.shape[0] <= MAX_TOKENS:
    y = o_proj_gemm_ar(x, w)
else:
    y = tensor_model_parallel_all_reduce(x @ w.t())

M = x.shape[0]
if M > MAX_TOKENS:
    y = tensor_model_parallel_all_reduce(x @ w.t())
else:
    y = o_proj_gemm_ar(x, w)

When it happens

Trigger: Calling o_proj_gemm_ar with an M (rows of the TP-local o_proj input) greater than 512 (MAX_TOKENS), e.g. a prefill chunk or large target-verify batch routed into this decode-oriented fused kernel.

Common situations: Using the decode/MTP fused GEMM path during prefill or large batched speculative verification; growing max_num_tokens or speculative num_spec beyond what the kernel was compiled for.

Related errors

• tiny_k_gemm: no valid split_n for N={n}, K={k}
• The pointers must be multiple of 16 bytes.
• The last dimension ({input.shape[-1]}) x itemsize ({input.dt
• rope_pool_fused expects q/k/v to be 3-D
• rope_pool_fused expects positions/slots to be 1-D

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/92b22a5a2b77689d. Report an issue: GitHub.