sgl-project/sglang · error · RuntimeError

tiny_k_gemm: no valid split_n for N={n}, K={k}

Error message

tiny_k_gemm: no valid split_n for N={n}, K={k}

What it means

tiny_k_gemm_bf16 picks a split_n divisor d of N such that d*lanes is a multiple of 32 and d*lanes <= 1024. If no divisor of N satisfies those constraints, the candidate list is empty and the kernel refuses to launch because no valid work decomposition exists.

Source

Thrown at python/sglang/kernels/ops/gemm/tiny_gemm.py:107

        assert out_dtype is None or out_dtype == out.dtype
    if split_n is None:
        split_n = _default_split_n(n, k, max_m, x.device)
    module = _jit_tiny_gemm_module(n, k, max_m, split_n, out.dtype)
    module.run(x, w, out)
    return out


def _default_k_split_n(n: int, k: int) -> int:
    """Smallest divisor of n whose n / split_n blocks fit one wave, with
    split_n * K-lanes whole-warp aligned and within the block-size limit."""
    lanes = k // 8  # fixed 16-byte vectors in the K variant
    candidates = [
        d
        for d in range(1, n + 1)
        if n % d == 0 and d * lanes % 32 == 0 and d * lanes <= 1024
    ]
    if not candidates:
        raise RuntimeError(f"tiny_k_gemm: no valid split_n for N={n}, K={k}")
    sm_count = torch.cuda.get_device_properties(0).multi_processor_count
    for d in candidates:
        if n // d <= sm_count:
            return d
    return candidates[-1]


def tiny_k_gemm_bf16(
    x: torch.Tensor,
    w: torch.Tensor,
    out: Optional[torch.Tensor] = None,
    *,
    out_dtype: Optional[torch.dtype] = None,
    split_n: Optional[int] = None,
    max_m: int = _MAX_M_DEFAULT,
) -> torch.Tensor:
    """Small-K / large-N variant: K / 8 lanes of one warp reduce the K
    dimension for one output column; each block covers split_n columns and the

View on GitHub (pinned to 0132848349)

Solutions

  1. Check N and lanes; if N is odd or prime relative to 32/lanes divisibility, pad the output dimension N up to the next multiple that has a valid divisor (e.g. next multiple of 32/lanes)
  2. Adjust the lanes argument so d*lanes can be a multiple of 32 with d dividing N
  3. Fall back to a standard torch.matmul path for this shape if padding is not acceptable

Example fix

# before
y = tiny_k_gemm_bf16(x, w, n=1150, lanes=1)  # 1150 has no valid split_n
# after
pad = (-n) % 32
y = tiny_k_gemm_bf16(x, w, n=n+pad, lanes=1)[:n]
Defensive patterns

Strategy: validation

Validate before calling

lanes = 1  # given
valid = any(n % d == 0 and d*lanes % 32 == 0 and d*lanes <= 1024 for d in range(1, n+1))
if not valid:
    n_padded = n + (-n) % 32  # pad then slice output

Type guard

def tiny_k_gemm_shape_ok(n: int, lanes: int) -> bool:
    return any(n % d == 0 and d*lanes % 32 == 0 and d*lanes <= 1024 for d in range(1, n+1))

Prevention

When it happens

Trigger: Calling tiny_k_gemm_bf16 with an N that has no divisor d meeting (n % d == 0 and d*lanes % 32 == 0 and d*lanes <= 1024) — typically odd/prime N combined with a lanes value whose multiples never hit a multiple of 32 within the 1024 cap (e.g. lanes=3 with odd N).

Common situations: Non-power-of-two output dimensions in small GEMMs, unusual head sizes or intermediate dims, or a lanes parameter changed during tuning that breaks divisibility.

Related errors


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