sgl-project/sglang · critical · AssertionError

mega MoE: num_tokens={num_tokens} exceeds SGLANG_OPT_DEEPGEM

Error message

mega MoE: num_tokens={num_tokens} exceeds SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK={num_max_tokens_per_rank}; K3 has no non-mega fallback — raise the env var to cover the per-rank rows

What it means

Kimi K3's mega-MoE DeepGEMM path allocates symmetric buffers sized by SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK and asserts routed_input rows fit. Unlike other models, K3 has no non-mega fallback, so overflow is a hard failure rather than a graceful path switch.

Source

Thrown at python/sglang/srt/models/kimi_k3.py:785

        backend (combine returns fully-summed rows; `_reduce_latent` then only
        applies the norm)."""
        import deep_gemm

        from sglang.kernels.ops.attention.dsv4 import mega_moe_pre_dispatch
        from sglang.srt.distributed.parallel_state import get_moe_ep_group
        from sglang.srt.environ import envs
        from sglang.srt.layers.moe.mega_moe import _get_mega_moe_symm_buffer

        # In SP-MoE mode (KimiK3DecoderLayer reduce-scatters the o_proj
        # output) the incoming rows are already this rank's token shard, so
        # the fused a2a below dispatches each token exactly once. On the
        # non-scattered fallback path the rows are the full batch (redundant
        # across ranks but correct).
        num_tokens = routed_input.shape[0]
        num_max_tokens_per_rank = (
            envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get()
        )
        assert num_tokens <= num_max_tokens_per_rank, (
            f"mega MoE: num_tokens={num_tokens} exceeds "
            f"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK="
            f"{num_max_tokens_per_rank}; K3 has no non-mega fallback — raise "
            f"the env var to cover the per-rank rows"
        )
        buf = _get_mega_moe_symm_buffer(
            get_moe_ep_group().device_group,
            num_experts=self.experts.num_experts,
            num_max_tokens_per_rank=num_max_tokens_per_rank,
            num_topk=self._mega_top_k,
            hidden=self.moe_hidden_size,
            intermediate_hidden=self._mega_intermediate_size,
        )

        if num_tokens > 0:
            topk_ids_in = topk_output.topk_ids.to(torch.int32)
            topk_weights_in = topk_output.topk_weights.to(torch.float32)
        else:

View on GitHub (pinned to 0132848349)

Solutions

Raise SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK above the per-rank row count (prefill chunk or full batch, whichever path runs)
Lower --chunked-prefill-size / batch size to fit the current cap
If memory-bound, reduce parallelism per node or free memory before raising the cap since buffers scale with it

Example fix

# before
python -m sglang.launch_server --model moonshot-ai/Kimi-K3 --chunked-prefill-size 32768
# after
SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=65536 \
  python -m sglang.launch_server --model moonshot-ai/Kimi-K3 --chunked-prefill-size 32768

Defensive patterns

Strategy: validation

Validate before calling

cap = int(os.environ.get("SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK", "0") or 0)
needed = max(chunked_prefill_size, batch_rows)
assert cap >= needed, f"set SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK>={needed}"

Prevention

For Kimi K3, always set the mega-MoE cap to at least the largest per-rank batch (there is no fallback)
Budget GPU memory: symmetric buffers scale with the cap
Re-verify after changing chunked-prefill or DP settings

When it happens

Trigger: Running Kimi K3 with the mega MoE (OPT DeepGEMM) path when routed_input.shape[0] (per-rank scattered rows, or full batch rows in the non-scattered path) exceeds the env-var limit — e.g. big prefill chunks or a whole-batch fallback batch.

Common situations: Serving Kimi K3 with large --chunked-prefill-size or high concurrency; the default mega-MoE token cap being smaller than the actual per-rank workload.

Related errors

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