vllm-project/vllm · error · NotImplementedError

MoRIIO heterogeneous TP requires replicated KV heads on both

Error message

MoRIIO heterogeneous TP requires replicated KV heads on both prefill and decode. Got total_num_kv_heads={total_num_kv_heads}, local_tp_size={local_tp_size}, remote_tp_size={remote_tp_size}.

What it means

validate_moriio_heterogeneous_tp_kv_heads raises NotImplementedError when min(local_tp_size, remote_tp_size) < total_num_kv_heads for non-MLA models with unequal TP sizes. MoRI-IO heterogeneous TP only supports replicated (not sharded) KV heads on both sides; once TP exceeds the KV-head count, heads would need splitting across ranks, which the transfer path cannot do.

Source

Thrown at vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py:144

            f"local tp_size {local_tp_size} must be a multiple of remote "
            f"tp_size {remote_tp_size} for heterogeneous-TP P/D"
        )
    return local_tp_rank // (local_tp_size // remote_tp_size)


def validate_moriio_heterogeneous_tp_kv_heads(
    local_tp_size: int,
    remote_tp_size: int,
    total_num_kv_heads: int,
    is_mla: bool,
) -> None:
    if is_mla or local_tp_size == remote_tp_size:
        return
    if local_tp_size <= 0 or remote_tp_size <= 0 or total_num_kv_heads <= 0:
        raise ValueError("TP sizes and total_num_kv_heads must be positive")
    if min(local_tp_size, remote_tp_size) >= total_num_kv_heads:
        return
    raise NotImplementedError(
        "MoRIIO heterogeneous TP requires replicated KV heads on both "
        f"prefill and decode. Got total_num_kv_heads={total_num_kv_heads}, "
        f"local_tp_size={local_tp_size}, remote_tp_size={remote_tp_size}."
    )


def get_moriio_expected_ack_count(producer_tp_size: int, consumer_tp_size: int) -> int:
    if producer_tp_size <= 0 or consumer_tp_size <= 0:
        raise ValueError("TP sizes must be positive")
    if consumer_tp_size <= producer_tp_size:
        return 1
    if consumer_tp_size % producer_tp_size != 0:
        raise ValueError(
            f"consumer tp_size {consumer_tp_size} must be a multiple of "
            f"producer tp_size {producer_tp_size} for heterogeneous-TP P/D"
        )
    return consumer_tp_size // producer_tp_size

View on GitHub (pinned to c794754062)

Solutions

1. Use equal TP sizes on prefill and decode (the check is skipped when local_tp_size == remote_tp_size)
2. Or keep both TP sizes <= total_num_kv_heads so heads stay replicated (e.g. TP <= num_key_value_heads)
3. Or use an MLA-architecture model (is_mla=True bypasses the check entirely)
4. As a last resort pick a model variant with more KV heads

Example fix

# before: GQA model with 8 kv heads, prefill TP=4, decode TP=16 -> NotImplementedError
# after: prefill TP=4, decode TP=8 (both <= 8 kv heads, replicated) -> ok

def supports_heterogeneous_tp(local_tp: int, remote_tp: int, total_kv_heads: int, is_mla: bool) -> bool:
    if is_mla or local_tp == remote_tp:
        return True
    return min(local_tp, remote_tp) >= total_kv_heads > 0

assert supports_heterogeneous_tp(p_tp, d_tp, total_kv_heads, is_mla), "KV heads would be sharded; equalize TP or use an MLA model"

def is_kv_head_replicated(local_tp: int, remote_tp: int, total_kv_heads: int) -> bool:
    return min(local_tp, remote_tp) >= total_kv_heads

try:
    validate_moriio_heterogeneous_tp_kv_heads(local_tp, remote_tp, kv_heads, is_mla)
except NotImplementedError:
    # fall back to equal TP deployment; do not retry with same sizes
    redeploy_with_equal_tp()

When it happens

Trigger: GQA model with few KV heads run with unequal P/D TP sizes, e.g. 8 KV heads with prefill TP=4 / decode TP=16 (min=4 < 8 is fine, but TP=16 with 8 heads fails), or MLA=False models like Qwen/Llama GQA variants with kv_heads < larger tp size.

Common situations: Scaling decode TP beyond the model's num_key_value_heads on GQA models; mixing an MHA model (kv_heads == num_heads, usually fine) with a heavily sharded GQA model; switching from an MLA architecture (exempt) to GQA without rebalancing TP.

Related errors

• remote tp_size {remote_tp_size} must be a multiple of local
• local tp_size {local_tp_size} must be a multiple of remote t
• TP sizes and total_num_kv_heads must be positive
• consumer tp_size {consumer_tp_size} must be a multiple of pr
• TP sizes must be positive

AI-assisted analysis of vllm-project/vllm@c794754062 (2026-08-14). Data as JSON: /api/errors/805e537459597a3f. Report an issue: GitHub.