{"record":{"id":"92b22a5a2b77689d","repo":"sgl-project/sglang","slug":"gemm-ar-m-m-outside-1-max-tokens","errorCode":null,"errorMessage":"gemm_ar: M={m} outside [1, {MAX_TOKENS}]","messagePattern":"gemm_ar: M=(.+?) outside \\[1, (.+?)\\]","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/kimi_k3/gemm_ar.py","lineNumber":191,"sourceCode":"@register_custom_op(mutates_args=[\"out\", \"epochs\"])\ndef _gemm_ar_op(\n    k: int,\n    world_size: int,\n    out: torch.Tensor,\n    x: torch.Tensor,\n    weight: torch.Tensor,\n    gather: torch.Tensor,\n    epochs: torch.Tensor,\n    my_rank: int,\n) -> None:\n    _module_with_bases(k, world_size).run(out, x, weight, gather, epochs, my_rank)\n\n\ndef _cell_of(m: int) -> int:\n    for c in (8, 16, 32, 64, 128, 256, 512):\n        if m <= c:\n            return c\n    raise ValueError(f\"gemm_ar: M={m} outside [1, {MAX_TOKENS}]\")\n\n\ndef o_proj_gemm_ar(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:\n    \"\"\"Fully reduced ``sum_r x_r @ weight_r^T`` on every rank, one kernel.\n\n    ``x`` is the TP-local [M, K] o_proj input, ``weight`` the TP-local\n    [7168, K] o_proj weight shard. Caller checked :func:`fits`; all ranks\n    call in lockstep with the same M.\n    \"\"\"\n    state = _STATE\n    assert state is not None\n    m = x.shape[0]\n    cell = _cell_of(m)\n    out = torch.empty((cell, N), dtype=torch.bfloat16, device=x.device)\n    _gemm_ar_op(\n        weight.shape[1],\n        state.world_size,\n        out,","sourceCodeStart":173,"sourceCodeEnd":209,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/kimi_k3/gemm_ar.py#L173-L209","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Route batches with M > MAX_TOKENS to the regular GEMM + all-reduce path instead of o_proj_gemm_ar","Chunk the input into <=MAX_TOKENS row blocks and call the kernel per chunk","If the kernel is intended for larger M, raise MAX_TOKENS and extend the cell tuple accordingly (requires revalidating perf)"],"exampleFix":"# before\ny = o_proj_gemm_ar(x, w)  # x.shape[0] == 640 > 512\n# after\nif x.shape[0] <= MAX_TOKENS:\n    y = o_proj_gemm_ar(x, w)\nelse:\n    y = tensor_model_parallel_all_reduce(x @ w.t())","handlingStrategy":"validation","validationCode":"M = x.shape[0]\nif M > MAX_TOKENS:\n    y = tensor_model_parallel_all_reduce(x @ w.t())\nelse:\n    y = o_proj_gemm_ar(x, w)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Bound decode/speculative batch sizes to MAX_TOKENS","Gate dispatch on token count in the o_proj backend selector","Add a regression test that M == MAX_TOKENS + 1 routes to fallback"],"tags":["gemm","all-reduce","token-limit","shape-validation"],"backgroundTag":"input-exceeds-limit","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}