{"record":{"id":"805e537459597a3f","repo":"vllm-project/vllm","slug":"moriio-heterogeneous-tp-requires-replicated-kv-hea","errorCode":null,"errorMessage":"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}.","messagePattern":"MoRIIO heterogeneous TP requires replicated KV heads on both prefill and decode\\. Got total_num_kv_heads=(.+?), local_tp_size=(.+?), remote_tp_size=(.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py","lineNumber":144,"sourceCode":"            f\"local tp_size {local_tp_size} must be a multiple of remote \"\n            f\"tp_size {remote_tp_size} for heterogeneous-TP P/D\"\n        )\n    return local_tp_rank // (local_tp_size // remote_tp_size)\n\n\ndef validate_moriio_heterogeneous_tp_kv_heads(\n    local_tp_size: int,\n    remote_tp_size: int,\n    total_num_kv_heads: int,\n    is_mla: bool,\n) -> None:\n    if is_mla or local_tp_size == remote_tp_size:\n        return\n    if local_tp_size <= 0 or remote_tp_size <= 0 or total_num_kv_heads <= 0:\n        raise ValueError(\"TP sizes and total_num_kv_heads must be positive\")\n    if min(local_tp_size, remote_tp_size) >= total_num_kv_heads:\n        return\n    raise NotImplementedError(\n        \"MoRIIO heterogeneous TP requires replicated KV heads on both \"\n        f\"prefill and decode. Got total_num_kv_heads={total_num_kv_heads}, \"\n        f\"local_tp_size={local_tp_size}, remote_tp_size={remote_tp_size}.\"\n    )\n\n\ndef get_moriio_expected_ack_count(producer_tp_size: int, consumer_tp_size: int) -> int:\n    if producer_tp_size <= 0 or consumer_tp_size <= 0:\n        raise ValueError(\"TP sizes must be positive\")\n    if consumer_tp_size <= producer_tp_size:\n        return 1\n    if consumer_tp_size % producer_tp_size != 0:\n        raise ValueError(\n            f\"consumer tp_size {consumer_tp_size} must be a multiple of \"\n            f\"producer tp_size {producer_tp_size} for heterogeneous-TP P/D\"\n        )\n    return consumer_tp_size // producer_tp_size\n","sourceCodeStart":126,"sourceCodeEnd":162,"githubUrl":"https://github.com/vllm-project/vllm/blob/c794754062d49a8fdb63ab3c5215b488b865030c/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py#L126-L162","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Use equal TP sizes on prefill and decode (the check is skipped when local_tp_size == remote_tp_size)","Or keep both TP sizes <= total_num_kv_heads so heads stay replicated (e.g. TP <= num_key_value_heads)","Or use an MLA-architecture model (is_mla=True bypasses the check entirely)","As a last resort pick a model variant with more KV heads"],"exampleFix":"# before: GQA model with 8 kv heads, prefill TP=4, decode TP=16 -> NotImplementedError\n# after: prefill TP=4, decode TP=8 (both <= 8 kv heads, replicated) -> ok","handlingStrategy":"validation","validationCode":"def supports_heterogeneous_tp(local_tp: int, remote_tp: int, total_kv_heads: int, is_mla: bool) -> bool:\n    if is_mla or local_tp == remote_tp:\n        return True\n    return min(local_tp, remote_tp) >= total_kv_heads > 0\n\nassert supports_heterogeneous_tp(p_tp, d_tp, total_kv_heads, is_mla), \"KV heads would be sharded; equalize TP or use an MLA model\"","typeGuard":"def is_kv_head_replicated(local_tp: int, remote_tp: int, total_kv_heads: int) -> bool:\n    return min(local_tp, remote_tp) >= total_kv_heads","tryCatchPattern":"try:\n    validate_moriio_heterogeneous_tp_kv_heads(local_tp, remote_tp, kv_heads, is_mla)\nexcept NotImplementedError:\n    # fall back to equal TP deployment; do not retry with same sizes\n    redeploy_with_equal_tp()","preventionTips":["For GQA models keep both TP sizes <= num_key_value_heads, or equal","Check num_key_value_heads against planned TP before choosing an asymmetric topology","Remember MLA models bypass this constraint"],"tags":["tensor-parallel","kv-transfer","gqa","heterogeneous-tp","model-config"],"backgroundTag":null,"analyzedSha":"c794754062d49a8fdb63ab3c5215b488b865030c","analyzedAt":"2026-08-14T21:17:39.825Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}