sgl-project/sglang · error · NotImplementedError

topk kernels only support k <= 32: {k=}

Error message

topk kernels only support k <= 32: {k=}

What it means

NotImplementedError from gate_topk: the streaming top-k Triton kernel uses power-of-two block logic sized for at most 32 selections per row, so k > 32 cannot be handled and the function refuses loudly instead of producing wrong results.

Source

Thrown at python/sglang/kernels/ops/moe/gate_topk.py:142

    Stable implementation of torch.topk(..., dim=-1) that is most efficient
    for small values of k.
    """
    assert x.is_contiguous(), f"{x.shape=} {x.stride()=}"
    assert x.ndim == 2, f"{x.shape=}"
    assert x.numel() <= 2**31, f"assumes int32 indexing: {x.shape=}"
    n_rows, n_cols = x.shape
    if return_values:
        values = torch.empty((n_rows, k), dtype=x.dtype, device=x.device)
    else:
        values = None
    # int32 indices: column ids fit int32 (numel <= 2**31, asserted above) and the
    # sole caller (Inkling gate) feeds the SRT MoeRunner topk-packing, which requires
    # int32. The kernel store casts to the buffer dtype, so this emits int32
    # directly — no separate .to(int32) downstream.
    indices = torch.empty((n_rows, k), dtype=torch.int32, device=x.device)
    if k > 32:
        # For larger topk, we need to reevaluate the kernel strategy
        raise NotImplementedError(f"topk kernels only support k <= 32: {k=}")

    if _impl == "streaming":
        BLOCK_SIZE_N = 32
        BLOCK_SIZE_M = 32
        grid = (triton.cdiv(n_rows, BLOCK_SIZE_M),)
        _streaming_topk_kernel[grid](
            x_ptr=x,
            stride_xm=x.stride(0),
            values_ptr=values,
            indices_ptr=indices,
            M=n_rows,
            N=n_cols,
            N_PAD=triton.cdiv(n_cols, BLOCK_SIZE_N) * BLOCK_SIZE_N,
            K=k,
            K_POW2=triton.next_power_of_2(k),
            BLOCK_SIZE_M=BLOCK_SIZE_M,
            BLOCK_SIZE_N=BLOCK_SIZE_N,
            RETURN_VALUES=return_values,

View on GitHub (pinned to 0132848349)

Solutions

  1. Reduce k to <= 32 if the model tolerates it (check the model's routing config)
  2. Fall back to torch.topk-based routing instead of the fused Triton kernel for k > 32
  3. If the workload genuinely needs k > 32, extend the kernel strategy as the error message suggests (new kernel path)

Example fix

// before
ids, w = gate_topk(logits, k=64)
// after
ids, w = gate_topk(logits, k=32)  # or torch.topk(logits, 64, dim=-1)
Defensive patterns

Strategy: fallback

Validate before calling

if k > 32:
    vals, ids = torch.topk(logits, k, dim=-1)
else:
    ids = gate_topk(logits, k)

Prevention

When it happens

Trigger: Calling gate_topk (from the routing forward) with topk > 32, e.g. a model config with num_selected_experts=64 or a very large top-k routing setting.

Common situations: Enabling a new model whose router selects more than 32 experts; experimenting with high top-k for routing quality; overriding topk via CLI/config without checking kernel limits.

Related errors


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