sgl-project/sglang · error · ValueError

KDA num_heads ({num_heads}) must be divisible by shard tp_si

Error message

KDA num_heads ({num_heads}) must be divisible by shard tp_size ({tp_size})

What it means

Kimi-Linear's KDA (Kimi Delta Attention) state shards its num_heads across tensor-parallel ranks; _get_kda_local_num_heads (kimi_linear.py:58) requires num_heads % tp_size == 0 and otherwise raises at model init. Head counts cannot be split unevenly, so TP degrees like 3, 5, 6, 7 are invalid for a given head count.

Source

Thrown at python/sglang/srt/models/kimi_linear.py:58

)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import (
    default_weight_loader,
    maybe_remap_kv_scale_name,
    sharded_weight_loader,
)
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA as KimiMLAAttention
from sglang.srt.models.llama import LlamaMLP as KimiMLP
from sglang.srt.models.transformers import maybe_prefix
from sglang.srt.runtime_context import get_parallel, get_stream
from sglang.srt.utils import is_xpu, make_layers
from sglang.srt.utils.common import BumpAllocator, add_prefix, set_weight_attrs


def _get_kda_local_num_heads(num_heads: int, tp_size: int) -> int:
    if num_heads % tp_size != 0:
        raise ValueError(
            f"KDA num_heads ({num_heads}) must be divisible by shard tp_size ({tp_size})"
        )
    return num_heads // tp_size


def _materialize_residual_stream(
    hidden_states: torch.Tensor, residual: Optional[torch.Tensor]
) -> torch.Tensor:
    return hidden_states if residual is None else hidden_states + residual


class KimiMoE(nn.Module):
    def __init__(
        self,
        config: KimiLinearConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        layer_idx: int = 0,

View on GitHub (pinned to 0132848349)

Solutions

  1. Pick a tp_size that divides the KDA num_heads from config (commonly powers of 2: 1/2/4/8)
  2. Check config (e.g. num_kda_heads / linear attention heads) before choosing --tp-size
  3. Use expert/data parallelism instead of an odd TP degree if GPU count forces it

Example fix

# before
python -m sglang.launch_server --tp-size 6 ...
# after
python -m sglang.launch_server --tp-size 4 ...
Defensive patterns

Strategy: validation

Validate before calling

assert cfg.num_kda_heads % args.tp_size == 0, f"{cfg.num_kda_heads} % {args.tp_size}"

Prevention

When it happens

Trigger: Launching with --tp-size N where N does not divide the config's KDA num_heads (e.g. tp=6 with 32 heads: 32 % 6 != 0).

Common situations: Choosing TP size based on GPU count without checking head divisibility; mixing KDA (linear attention) models with TP values valid only for standard transformers.

Related errors


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