sgl-project/sglang · error · ValueError

num_key_value_heads must be divisible by attention TP

Error message

num_key_value_heads must be divisible by attention TP

What it means

When total_num_kv_heads >= attn_tp_size, KV heads must shard evenly: num_key_value_heads % attn_tp_size == 0. Otherwise replication/sharding of GQA KV heads is undefined.

Source

Thrown at python/sglang/srt/models/interns2_mobius.py:571

        config: InternS2MobiusTextConfig,
        layer_id: int,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        alt_stream: torch.cuda.Stream | None = None,
    ) -> None:
        nn.Module.__init__(self)
        self.config = config
        self.hidden_size = config.hidden_size
        self.attn_tp_rank = get_parallel().attn_tp_rank
        self.attn_tp_size = get_parallel().attn_tp_size
        self.total_num_heads = config.num_attention_heads
        if self.total_num_heads % self.attn_tp_size != 0:
            raise ValueError("num_attention_heads must be divisible by attention TP")
        self.num_heads = self.total_num_heads // self.attn_tp_size
        self.total_num_kv_heads = config.num_key_value_heads
        if self.total_num_kv_heads >= self.attn_tp_size:
            if self.total_num_kv_heads % self.attn_tp_size != 0:
                raise ValueError(
                    "num_key_value_heads must be divisible by attention TP"
                )
        elif self.attn_tp_size % self.total_num_kv_heads != 0:
            raise ValueError("attention TP must be divisible by num_key_value_heads")
        self.num_kv_heads = max(1, self.total_num_kv_heads // self.attn_tp_size)
        self.head_dim = config.head_dim or (self.hidden_size // self.num_heads)
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5
        self.max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
        self.rope_theta, rope_scaling = get_rope_config(config)
        self.partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
        self.layer_id = layer_id
        if rope_scaling and not ("rope_type" in rope_scaling or "type" in rope_scaling):
            rope_scaling = None
        self.attn_output_gate = getattr(config, "attn_output_gate", True)
        self.rotary_emb = get_rope(
            head_size=self.head_dim,

View on GitHub (pinned to 0132848349)

Solutions

  1. Choose TP size that divides num_key_value_heads (e.g., tp=6 or 4 for 12 KV heads)
  2. Enable or adjust --dp-size / dp-attention so attn_tp_size divides KV heads
  3. Fall back to tp=1 for odd head counts

Example fix

# before
--tp 8   # 12 kv heads, 12 % 8 != 0

# after
--tp 4    # 12 % 4 == 0
Defensive patterns

Strategy: validation

Validate before calling

kv, tp = config.num_key_value_heads, get_parallel().attn_tp_size
assert kv < tp or kv % tp == 0

Type guard

def kv_heads_shard_ok(kv: int, tp: int) -> bool:
    return kv % tp == 0 if kv >= tp else tp % kv == 0

Prevention

When it happens

Trigger: Building the attention with attn_tp_size exceeding but not dividing num_key_value_heads, e.g., 6 KV heads with attn_tp_size 8 (6 < 8 falls to the elif, but 12 KV heads with tp=8 hits this branch: 12%8!=0).

Common situations: GQA models with few KV heads launched under high TP, or DP-attention setups that change the effective attention TP group size.

Related errors


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