huggingface/transformers · critical · ValueError

Number of key value heads {} must be divisible by tensor par

Error message

Number of key value heads {} must be divisible by tensor parallel size {}.

What it means

ValueError raised when tensor parallelism shards the K/V projections (tp_plan contains layers.*.self_attn.k_proj / v_proj) but config.num_key_value_heads is not divisible by the TP size. Each GPU must receive a whole number of KV heads for GQA sharding; otherwise attention shapes become inconsistent across ranks.

Source

Thrown at src/transformers/generation/continuous_batching/cache.py:207

            sliding_window = config.sliding_window if group_types[i] == "sliding_attention" else 1
            for j, layer in enumerate(group):
                self.layer_index_to_group_indices[layer] = (i, j)
                self.sliding_windows[layer] = sliding_window

        # Check if the KV heads are part of the TP plan. If they are not, the cache does not need plan for TP.
        # TODO: this is fragile. If your model fails to TP properly because of this, please open an issue.
        kv_is_tp = True
        for key in ["layers.*.self_attn.k_proj", "layers.*.self_attn.v_proj"]:
            if not (key in tp_plan or "model." + key in tp_plan):
                kv_is_tp = False
                break

        # If the KV heads are TP'ed, each KV head is dispatched to a different GPU, so the effective number of KV heads
        # per GPU is simply divided by the TP size
        tp_size = distributed_helper.tp_size
        if tp_size > 1 and kv_is_tp:
            if self.num_key_value_heads % tp_size != 0:
                raise ValueError(
                    f"Number of key value heads {self.num_key_value_heads} must be divisible by tensor parallel size {tp_size}."
                )
            self.num_key_value_heads //= tp_size

        # If somehow the max memory percent is not yet resolved, resolve it conservatively
        if continuous_batching_config.max_memory_percent is None:
            resolve_max_memory_percent(cb_config=continuous_batching_config, has_logit_processors=True)

        max_batch_tokens, num_blocks = PagedAttentionMemoryHandler(
            config=config,
            continuous_batching_config=continuous_batching_config,
            dtype=self.dtype,
            group_types=group_types,
            group_size=group_size,
        ).infer_max_batch_tokens_and_num_blocks()

        # For TP, align max_batch_tokens and num_blocks to the minimal value across the TP group
        if tp_size > 1:

View on GitHub (pinned to a597f97485)

Solutions

  1. Choose a tensor_parallel_size that divides num_key_value_heads (e.g. for 8 KV heads use TP in {1,2,4,8})
  2. Or use data/pipeline parallelism instead of TP for models with few KV heads
  3. Verify config.num_key_value_heads and the tp_plan; if KV projections should be replicated, remove them from the TP plan

Example fix

# before
devices = ['cuda:0','cuda:1','cuda:2','cuda:3']  # TP=4, model has num_key_value_heads=8 -> ok; with 6 heads -> error
# after: pick TP that divides KV heads
devices = ['cuda:0','cuda:1']  # TP=2 divides 6 and 8
Defensive patterns

Strategy: validation

Validate before calling

kv = getattr(config, 'num_key_value_heads', None) or config.num_attention_heads
for tp in candidate_tp_sizes:
    if kv % tp == 0:
        return tp
raise ValueError(f'no valid TP size divides {kv} KV heads')

Type guard

def tp_compatible(config, tp_size: int) -> bool:
    kv = getattr(config, 'num_key_value_heads', None) or config.num_attention_heads
    return tp_size <= 1 or kv % tp_size == 0

Prevention

When it happens

Trigger: Running multi-GPU continuous batching with tensor_parallel_size=N where N does not divide num_key_value_heads (e.g. 7 KV heads with TP=4); models with prime/small KV-head counts (e.g. num_key_value_heads=8 with TP=6); custom tp_plan strings that accidentally match k_proj/v_proj.

Common situations: Scaling a GQA model (Llama-style with few KV heads) across more GPUs than the KV head count divides; mixing tp_plan styles ('model.layers...' vs 'layers...') so KV sharding is detected unexpectedly.

Related errors


AI-assisted analysis of huggingface/transformers@a597f97485 (2026-08-14). Data as JSON: /api/errors/2c595fa3e19176e3. Report an issue: GitHub.