[ROCm][Kimi-K3] Support MoRIIO consumers with specialized decode

docs/features/moriio_connector_usage.md

@@ -144,6 +144,12 @@ For DSpark, configure the same draft model on both instances through `--speculat
 
 Hybrid READ keeps its existing single-transferable-attention-group requirement. Plain-attention WRITE retains its existing notification format.
 
+### Kimi-K3 specialized decode
+
+Kimi-K3 on MI355X supports specialized decode on built-in MoRIIO READ and WRITE consumers. Set `--additional-config '{"decode_backend":"latency_specialized"}'` on the consumer and keep the producer on the default `decode_backend=auto`.
+
+Use matching transfer modes and cache layouts. See [supported deployments](https://github.com/vllm-project/vllm/blob/main/vllm/models/kimi_k3/amd/specialized_decode/README.md#supported-deployments) for the hardware, state dtype, and execution requirements.
+
 ### Different prefill and decode TP sizes
 
 In WRITE mode, the prefill and decode instances can use different tensor-parallel sizes, as long as one is a multiple of the other. Each prefill rank writes its shard into the decode ranks that hold it:

tests/models/kimi_k3/amd_specialized_decode/test_runtime.py

@@ -8,7 +8,12 @@
 import pytest
 import torch
 
-from vllm.config import CompilationMode, CUDAGraphMode
+from vllm.config import (
+    CompilationMode,
+    CUDAGraphMode,
+    KVTransferConfig,
+    SpeculativeConfig,
+)
 from vllm.models.kimi_k3.amd.specialized_decode import runtime
 from vllm.models.kimi_k3.amd.specialized_decode.runtime import (
     FFN,
@@ -141,6 +146,71 @@ def test_supported_config_is_accepted(on_mi355x):
     assert specialized_decode_config(_vllm_config()) == DecodeConfig(q=8)
 
 
+@pytest.fixture
+def supported_vllm_config(tmp_path, monkeypatch, on_mi355x):
+    from tests.v1.attention.utils import create_vllm_config
+    from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    KimiLinearConfig(
+        architectures=["KimiLinearForCausalLM"],
+        hidden_size=7168,
+        num_experts=896,
+        attn_res_block_size=12,
+        linear_attn_config={
+            "kda_layers": [1],
+            "full_attn_layers": [2],
+            "num_heads": 96,
+            "head_dim": 128,
+            "short_conv_kernel_size": 4,
+        },
+    ).save_pretrained(tmp_path)
+    config = create_vllm_config(
+        model_name=str(tmp_path), tensor_parallel_size=8, dtype=torch.bfloat16
+    )
+    config.additional_config = {"decode_backend": "latency_specialized"}
+    config.speculative_config = SpeculativeConfig(
+        method="ngram", num_speculative_tokens=7
+    )
+    return config
+
+
+@pytest.mark.parametrize("read_mode", [False, True])
+@pytest.mark.parametrize("backend", ["rdma", "xgmi"])
+def test_moriio_consumer_is_supported(supported_vllm_config, read_mode, backend):
+    supported_vllm_config.kv_transfer_config = KVTransferConfig(
+        kv_connector="MoRIIOConnector",
+        kv_role="kv_consumer",
+        kv_connector_extra_config={"read_mode": read_mode, "backend": backend},
+    )
+    assert specialized_decode_config(supported_vllm_config) == DecodeConfig(q=8)
+
+
+@pytest.mark.parametrize(
+    "overrides",
+    [
+        {"kv_connector": "NixlConnector"},
+        {"kv_connector": "MultiConnector"},
+        {"kv_connector_module_path": "custom.connector"},
+        {"kv_role": "kv_producer"},
+        {"kv_role": "kv_both"},
+    ],
+)
+def test_unsupported_kv_transfer_is_refused(supported_vllm_config, overrides):
+    supported_vllm_config.kv_transfer_config = KVTransferConfig(
+        **(
+            {
+                "kv_connector": "MoRIIOConnector",
+                "kv_role": "kv_consumer",
+                "kv_connector_extra_config": {"read_mode": True},
+            }
+            | overrides
+        )
+    )
+    with pytest.raises(ValueError, match="KV transfer connectors"):
+        specialized_decode_config(supported_vllm_config)
+
+
 @pytest.mark.parametrize(
     "over, reason",
     [
@@ -153,7 +223,14 @@ def test_supported_config_is_accepted(on_mi355x):
         ),
         ({"parallel_config.pipeline_parallel_size": 2}, "pipeline parallelism"),
         ({"parallel_config.decode_context_parallel_size": 2}, "decode context"),
-        ({"kv_transfer_config": object()}, "KV transfer connectors"),
+        (
+            {
+                "kv_transfer_config": KVTransferConfig(
+                    kv_connector="NixlConnector", kv_role="kv_consumer"
+                )
+            },
+            "KV transfer connectors",
+        ),
         ({"speculative_config": None}, "without speculative decoding"),
         ({"speculative_config.num_speculative_tokens": 8}, "more than 7"),
         ({"speculative_config.enable_adaptive_verification": True}, "adaptive"),

vllm/models/kimi_k3/amd/specialized_decode/README.md

@@ -31,12 +31,16 @@ TP8 on MI355X with all 256 CUs, decode context parallel 1 or 8, speculative
 decoding with at most 7 speculative tokens, Model Runner V2, BF16 conv / FP32
 recurrent KDA states.
 
+KV transfer support is limited to built-in MoRIIO READ and WRITE consumers.
+Keep the producer on `decode_backend=auto` and use matching cache layouts.
+
 If you request the specialized decode for anything else, startup fails with one
 `ValueError` that lists every unsupported setting:
 
 - **Config, in the model's `__init__`:** pipeline, data, expert or prefill
-  context parallelism, EPLB, DBO, LoRA, KV connectors, routed-expert capture,
-  sleep mode, CPU weight offload, weight updates, no speculative decoding,
+  context parallelism, EPLB, DBO, LoRA, unsupported KV connectors,
+  routed-expert capture, sleep mode, CPU weight offload, weight updates,
+  no speculative decoding,
   adaptive verification, dynamic speculative decoding, ReplaySSM / RecoverSSM,
   the V1 model runner, `STOCK_TORCH_COMPILE`, another model's shapes, another GPU.
 - **Loaded weights, when the model state builds the executor:** each kernel

vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py

@@ -13,6 +13,7 @@
 
 import torch
 
+from vllm.config import KVTransferConfig
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
@@ -103,7 +104,10 @@ def unsupported(vllm_config) -> list[str]:
         # the kernels read one batch's attention metadata per step
         (pc.use_ubatching, "DBO (dual-batch overlap)"),
         (vllm_config.lora_config is not None, "LoRA"),
-        (vllm_config.kv_transfer_config is not None, "KV transfer connectors"),
+        (
+            not supports_kv_transfer(vllm_config.kv_transfer_config),
+            "KV transfer connectors other than a MoRIIO consumer",
+        ),
         (
             vllm_config.aux_output_config.enable_return_routed_experts,
             "routed-expert capture",
@@ -162,6 +166,15 @@ def unsupported(vllm_config) -> list[str]:
     return reasons + platform_reasons()
 
 
+def supports_kv_transfer(config: KVTransferConfig | None) -> bool:
+    # MoRIIO consumers finish loads before execution, without per-layer hooks.
+    return config is None or (
+        config.kv_connector == "MoRIIOConnector"
+        and config.kv_connector_module_path is None
+        and config.kv_role == "kv_consumer"
+    )
+
+
 def text_shapes(text) -> tuple:
     kda = getattr(text, "linear_attn_config", None) or {}
     return (