[Feature][MoRIIO] Support hybrid WRITE transfers with DSpark (#60544)
Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
YukioZzz
co-authored by
OpenAI Codex
tjtanaa
Claude Opus 5.5
OpenAI Codex
tjtanaa
Claude Opus 5.5
parent
a0e7137c4b
commit
41b8c15809
@@ -164,6 +164,59 @@ steps:
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
# The shared KV Connectors steps above exclude vllm/distributed/kv_transfer/ from
|
||||
# their triggers, and CUDA skips MoRIIO's mori/aiter-gated tests. Run the whole
|
||||
# MoRIIO suite on ROCm whenever the connector or its tests change.
|
||||
- label: ":amd: (MI355 DPX) MoRIIO Connector Unit Tests"
|
||||
key: amd-moriio-connector-unit-mi355
|
||||
dind: false
|
||||
device: mi355_dpx
|
||||
num_devices: 1
|
||||
timeout_in_minutes: 30
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/moriio/
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/base.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
|
||||
- examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
|
||||
- tests/v1/kv_connector/unit/test_moriio_
|
||||
- tests/v1/kv_connector/unit/utils.py
|
||||
- requirements/kv_connectors_rocm.txt
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s -rs v1/kv_connector/unit/test_moriio_*.py
|
||||
|
||||
# Single-node MoRIIO WRITE P/D on dummy weights: prefill and decode on separate
|
||||
# GPUs over xGMI, compared exactly against a direct serve.
|
||||
- label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N"
|
||||
key: amd-moriio-write-dummy-e2e-mi355
|
||||
dind: false
|
||||
device: mi355_2
|
||||
num_devices: 2
|
||||
# ~3.5-5 min per case; pytest-shard hashes test ids, so with 20 cases the
|
||||
# largest of 6 shards holds 5 (~22 min on MI355). Every case picks free ports
|
||||
# for its servers and proxy and its own compile-cache dir.
|
||||
parallelism: 6
|
||||
timeout_in_minutes: 45
|
||||
working_dir: /vllm-workspace/tests
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/moriio/
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/base.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
|
||||
- vllm/model_executor/model_loader/dummy_loader.py
|
||||
- examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
|
||||
- tests/v1/kv_connector/moriio_integration/
|
||||
- requirements/kv_connectors_rocm.txt
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
|
||||
- label: ":nvidia: (H200 MIG 35GB) V1 Metrics + LM Eval"
|
||||
device: h200_35gb
|
||||
key: v1-metrics-lmeval
|
||||
|
||||
@@ -136,6 +136,14 @@ WRITE mode is used by default. READ mode can be configured by setting `--kv-tran
|
||||
!!! note
|
||||
`notify_port` is used as a *base* port: each (DP rank, TP rank) pair within an instance uses `notify_port + offset` where the offset is based on the rank. Make sure the range starting at `notify_port` is free on the host.
|
||||
|
||||
### Hybrid attention and recurrent state
|
||||
|
||||
WRITE mode supports hybrid attention plus Mamba/KDA models, including DSpark, with equal producer and consumer TP sizes. Use matching cache layouts on both instances, `--mamba-cache-mode align`, and `read_mode: false` in each instance's `kv_connector_extra_config`. Both peers must support hybrid WRITE. Sliding-window WRITE and hybrid speculative methods other than DSpark remain unsupported.
|
||||
|
||||
For DSpark, configure the same draft model on both instances through `--speculative-config`. Target and draft attention groups retain their own block IDs and token spans; only the running recurrent state is transferred, not speculative scratch slots. The producer transfers the state after the first N-1 prompt tokens, and the consumer recomputes the last prompt token. Convolution and SSM transfers must both finish before the layer is complete.
|
||||
|
||||
Hybrid READ keeps its existing single-transferable-attention-group requirement. Plain-attention WRITE retains its existing notification format.
|
||||
|
||||
### Transport configuration
|
||||
|
||||
MoRI has two transport backends: RDMA and xGMI. You can select backend using `--kv-transfer-config.kv_connector_extra_config.backend $BACKEND`, with `$BACKEND` being `rdma` or `xgmi`. RDMA is the default backend and should be used in multi-node deployments.
|
||||
|
||||
@@ -473,9 +473,11 @@ async def stop_profile():
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--port", type=int, default=10001)
|
||||
# Port the prefill/decode servers register on (their proxy_ping_port).
|
||||
parser.add_argument("--service-discovery-port", type=int, default=36367)
|
||||
args = parser.parse_args()
|
||||
|
||||
t = start_service_discovery("0.0.0.0", 36367)
|
||||
t = start_service_discovery("0.0.0.0", args.service_discovery_port)
|
||||
# High-concurrency hardening. Quart's app.run() uses a shallow listen
|
||||
# backlog (100) and, with app.debug=True, adds per-request overhead that
|
||||
# starves the single accept loop. Under a burst of ~512 simultaneous client
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -0,0 +1,55 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Widen the dummy weight init for the MoRIIO dummy-weight P/D test.
|
||||
|
||||
The dummy loader fills float weights in +-1e-3, which keeps every activation so
|
||||
small that outputs barely depend on the input, so a wrong KV prefix would go
|
||||
unnoticed. The P/D test names ``WideDummyWeights`` as each server's
|
||||
``--worker-extension-cls``. Workers import it before they load the model, which
|
||||
re-registers the ``dummy`` load format in every worker process with a loader
|
||||
that fills 16/32-bit float weights in +-``MORIIO_DUMMY_WEIGHT_RANGE``; 8-bit
|
||||
weights and scales keep the default init. The format keeps the name ``dummy``
|
||||
so every config path that special-cases dummy weights still applies.
|
||||
|
||||
The test process itself never imports this module.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm.config import ModelConfig
|
||||
from vllm.model_executor.model_loader import register_model_loader
|
||||
from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader
|
||||
from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info
|
||||
from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
|
||||
from vllm.model_executor.model_loader.weight_utils import (
|
||||
initialize_single_dummy_weight,
|
||||
)
|
||||
|
||||
_RANGE = float(os.environ["MORIIO_DUMMY_WEIGHT_RANGE"])
|
||||
|
||||
|
||||
@register_model_loader("dummy")
|
||||
class WideDummyModelLoader(DummyModelLoader):
|
||||
def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
|
||||
# DummyModelLoader.load_weights with a wider init for 16/32-bit floats.
|
||||
for layer in model.modules():
|
||||
info = get_layerwise_info(layer)
|
||||
if info.can_load():
|
||||
self._process_online_quant_layer(layer, info)
|
||||
continue
|
||||
for name, tensor in get_layer_tensors(layer).items():
|
||||
if name in layer._non_persistent_buffers_set:
|
||||
continue
|
||||
if tensor.is_floating_point() and torch.finfo(tensor.dtype).bits >= 16:
|
||||
initialize_single_dummy_weight(
|
||||
tensor, low=-_RANGE, high=_RANGE, name=name
|
||||
)
|
||||
else:
|
||||
initialize_single_dummy_weight(tensor, name=name)
|
||||
|
||||
|
||||
class WideDummyWeights:
|
||||
"""Worker extension that only exists to import this module in workers."""
|
||||
@@ -0,0 +1,736 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Single-node MoRIIO WRITE P/D checks on dummy weights.
|
||||
|
||||
Each case serves a real model architecture with ``--load-format dummy``, widened
|
||||
by ``moriio_wide_dummy`` (see ``DUMMY_WEIGHT_RANGE``): the
|
||||
config keeps the real hidden size, head geometry, MLA ranks and recurrent-state
|
||||
shapes, so the connector registers and transfers exactly the production cache
|
||||
layout, but the layer stack is cut to a few layers that still contain every
|
||||
cache-group kind and the MoE is cut to a handful of experts. Dummy weights are
|
||||
seeded, so the prefill, decode and reference servers hold identical weights.
|
||||
Cases marked ``HMA_off`` run every server with
|
||||
``--disable-hybrid-kv-cache-manager``, so all layers, sliding-window ones
|
||||
included, share a single cache group.
|
||||
|
||||
Prefill runs on one GPU, decode on another, and MoRI moves the KV over xGMI.
|
||||
Either way decode computes the last prompt token itself. For hybrid models
|
||||
prefill computes only the first ``n - 1`` tokens and transfers the KV and the
|
||||
recurrent state after them. For attention-only models prefill computes and
|
||||
transfers all ``n``, and decode's scheduler, seeing the whole prompt cached,
|
||||
recomputes the last token. A direct serve splits the same prompt the same way
|
||||
when the first ``n - 1`` tokens end on a chunk boundary that both servers share,
|
||||
so the prompts are chosen to make that happen: one is a single chunk plus one
|
||||
token, the others are a whole number of hybrid cache blocks plus one token
|
||||
(prefix caching ends chunks on block boundaries) or, for attention-only models,
|
||||
a whole number of chunks plus one token. With the split matched, a correct
|
||||
transfer reproduces the direct serve's first-token log-probabilities exactly,
|
||||
so the comparison needs no noise allowance. The
|
||||
servers compile with Inductor's deterministic mode, since otherwise each one
|
||||
benchmarks its own reduction launch configs and separately started servers
|
||||
can drift apart by a few bf16 logit steps. Guards:
|
||||
|
||||
* decode must report an external prefix-cache hit for every transferred token
|
||||
of every prompt, so the match is not a local recompute;
|
||||
* changing the first half of a prompt must move the direct serve's log-probs by
|
||||
at least ``MIN_CONTEXT_EFFECT``, so a wrong prefix would be visible;
|
||||
* ``test_corrupted_prefill_state_is_detected`` makes prefill compute its KV with
|
||||
a different norm epsilon (same shapes, different values) and requires every
|
||||
prompt to fail the comparison.
|
||||
"""
|
||||
|
||||
import atexit
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from collections.abc import Callable, Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
|
||||
_ROOT = Path(__file__).resolve().parents[4]
|
||||
_PROXY = (
|
||||
_ROOT / "examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py"
|
||||
)
|
||||
|
||||
# Max allowed difference between centered first-token log-probabilities of the
|
||||
# P/D and direct outputs, over the direct serve's top-k tokens. The split is
|
||||
# matched, so correct transfers produce identical logits; this only absorbs the
|
||||
# fp32 rounding of the log-softmax (~1e-7), far below one bf16 logit step.
|
||||
TOLERANCE = float(os.getenv("MORIIO_DUMMY_E2E_TOL", "1e-5"))
|
||||
# Several bf16 steps of the logits at DUMMY_WEIGHT_RANGE.
|
||||
MIN_CONTEXT_EFFECT = 0.05
|
||||
# Half-width of the uniform dummy init for 16/32-bit float weights, applied in
|
||||
# every server worker by moriio_wide_dummy.py. The loader default (1e-3) shrinks
|
||||
# every activation so far that a context change moves the first-token log-probs
|
||||
# by ~1e-4; at 0.1 it moves them by >0.1 and a corrupted prefill changes the
|
||||
# top-k set entirely.
|
||||
DUMMY_WEIGHT_RANGE = 0.1
|
||||
_WIDE_DUMMY_DIR = Path(__file__).resolve().parent
|
||||
TOP_K = 20
|
||||
MAX_TOKENS = 4
|
||||
MAX_MODEL_LEN = 4096
|
||||
|
||||
# Prefill token budget of every server. Prompts longer than it are chunked, so
|
||||
# the final-chunk save path runs; the short prompt fills exactly one chunk.
|
||||
PREFILL_CHUNK = 299
|
||||
SHORT_PROMPT_LEN = PREFILL_CHUNK + 1
|
||||
# Hybrid blocks per aligned prompt, before the trailing token.
|
||||
ALIGNED_BLOCKS = (1, 2)
|
||||
# Prefill chunks per long attention-only prompt, before the trailing token.
|
||||
ALIGNED_CHUNKS = (2, 4)
|
||||
|
||||
|
||||
def _mori_available() -> bool:
|
||||
try:
|
||||
import mori.io # noqa: F401
|
||||
except ImportError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.skipif(
|
||||
not current_platform.is_rocm() or not _mori_available(),
|
||||
reason="MoRIIO P/D needs ROCm and MoRI",
|
||||
),
|
||||
pytest.mark.skipif(
|
||||
current_platform.is_rocm() and current_platform.device_count() < 2,
|
||||
reason="needs one GPU for prefill and one for decode",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model architectures
|
||||
|
||||
|
||||
def _kimi_linear_overrides(cfg) -> dict[str, Any]:
|
||||
# 3 KDA layers then 1 MLA layer; layer 0 stays dense (first_k_dense_replace).
|
||||
linear = dict(cfg.linear_attn_config)
|
||||
linear["kda_layers"] = [1, 2, 3]
|
||||
linear["full_attn_layers"] = [4]
|
||||
groups = max(1, cfg.num_expert_group)
|
||||
return {
|
||||
"num_hidden_layers": 4,
|
||||
"linear_attn_config": linear,
|
||||
"num_experts": 2 * groups,
|
||||
"num_experts_per_token": 2,
|
||||
"topk_group": min(cfg.topk_group, groups),
|
||||
}
|
||||
|
||||
|
||||
def _qwen3_5_moe_overrides(cfg) -> dict[str, Any]:
|
||||
# 3 Gated DeltaNet layers then 1 full-attention layer.
|
||||
return {
|
||||
"text_config": {
|
||||
"num_hidden_layers": 4,
|
||||
"layer_types": ["linear_attention"] * 3 + ["full_attention"],
|
||||
"num_experts": 4,
|
||||
"num_experts_per_tok": 2,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _qwen3_dense_overrides(cfg) -> dict[str, Any]:
|
||||
# Attention-only: decode gets the KV for the whole prompt.
|
||||
return {"num_hidden_layers": 4}
|
||||
|
||||
|
||||
def _deepseek_v3_overrides(cfg) -> dict[str, Any]:
|
||||
# MLA attention-only: layer 0 dense, layer 1 MoE, no MTP layer.
|
||||
return {
|
||||
"num_hidden_layers": 2,
|
||||
"first_k_dense_replace": 1,
|
||||
"n_routed_experts": 8,
|
||||
"num_experts_per_tok": 2,
|
||||
"n_group": 1,
|
||||
"topk_group": 1,
|
||||
"num_nextn_predict_layers": 0,
|
||||
}
|
||||
|
||||
|
||||
def _gemma3_overrides(cfg) -> dict[str, Any]:
|
||||
# Five sliding-window layers then one global layer, as in the real
|
||||
# 5:1 pattern. Without the hybrid KV cache manager vLLM gives the sliding
|
||||
# layers full-attention cache specs, so all six share one group.
|
||||
return {
|
||||
"num_hidden_layers": 6,
|
||||
"layer_types": ["sliding_attention"] * 5 + ["full_attention"],
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ArchCase:
|
||||
model: str
|
||||
overrides: Callable[[Any], dict[str, Any]]
|
||||
# Norm-epsilon override that perturbs every computed KV/state value
|
||||
# without changing a single shape. Used by the corruption control.
|
||||
corrupt_overrides: dict[str, Any]
|
||||
env: dict[str, str] = field(default_factory=dict)
|
||||
args: tuple[str, ...] = ()
|
||||
# Has recurrent-state (mamba-style) layers, so prefill stops at n - 1.
|
||||
hybrid: bool = True
|
||||
|
||||
|
||||
ARCHS: dict[str, ArchCase] = {
|
||||
"kimi_linear": ArchCase(
|
||||
model="moonshotai/Kimi-Linear-48B-A3B-Instruct",
|
||||
overrides=_kimi_linear_overrides,
|
||||
corrupt_overrides={"rms_norm_eps": 1.0},
|
||||
env={"VLLM_ROCM_USE_AITER": "1"},
|
||||
args=("--trust-remote-code", "--mamba-ssm-cache-dtype", "float32"),
|
||||
),
|
||||
"qwen3_5_moe": ArchCase(
|
||||
model="Qwen/Qwen3.5-35B-A3B",
|
||||
overrides=_qwen3_5_moe_overrides,
|
||||
corrupt_overrides={"text_config": {"rms_norm_eps": 1.0}},
|
||||
# AITER's GDN conv-update kernel rejects the DS conv-state layout that
|
||||
# MoRIIO's recurrent-state transfer requires.
|
||||
env={"VLLM_ROCM_USE_AITER": "0"},
|
||||
args=("--language-model-only",),
|
||||
),
|
||||
"qwen3_dense": ArchCase(
|
||||
model="Qwen/Qwen3-32B",
|
||||
overrides=_qwen3_dense_overrides,
|
||||
corrupt_overrides={"rms_norm_eps": 1.0},
|
||||
env={"VLLM_ROCM_USE_AITER": "0"},
|
||||
hybrid=False,
|
||||
),
|
||||
"deepseek_v3": ArchCase(
|
||||
model="deepseek-ai/DeepSeek-V3",
|
||||
overrides=_deepseek_v3_overrides,
|
||||
corrupt_overrides={"rms_norm_eps": 1.0},
|
||||
env={"VLLM_ROCM_USE_AITER": "1"},
|
||||
hybrid=False,
|
||||
),
|
||||
# Text-only Gemma 3. The multimodal Gemma3ForConditionalGeneration takes
|
||||
# the inputs_embeds path on KV consumers (enable_mm_embeds is forced on
|
||||
# there) and computes differently from a direct serve without any transfer.
|
||||
"gemma3": ArchCase(
|
||||
model="google/gemma-3-1b-it",
|
||||
overrides=_gemma3_overrides,
|
||||
corrupt_overrides={"rms_norm_eps": 1.0},
|
||||
env={"VLLM_ROCM_USE_AITER": "0"},
|
||||
hybrid=False,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _arch_overrides(arch: ArchCase) -> dict[str, Any]:
|
||||
from vllm.transformers_utils.config import get_config
|
||||
|
||||
return arch.overrides(get_config(arch.model, trust_remote_code=True))
|
||||
|
||||
|
||||
def _merge(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]:
|
||||
merged = dict(base)
|
||||
for key, value in extra.items():
|
||||
if isinstance(value, dict) and isinstance(merged.get(key), dict):
|
||||
merged[key] = _merge(merged[key], value)
|
||||
else:
|
||||
merged[key] = value
|
||||
return merged
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Serving
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Case:
|
||||
arch: str
|
||||
attention_backend: str
|
||||
prefill_graph: str
|
||||
decode_graph: str
|
||||
# Run every server of the case without the hybrid KV cache manager.
|
||||
hma_off: bool = False
|
||||
|
||||
@property
|
||||
def extra_args(self) -> tuple[str, ...]:
|
||||
return ("--disable-hybrid-kv-cache-manager",) if self.hma_off else ()
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
return (
|
||||
f"{self.arch}-{self.attention_backend}-"
|
||||
f"P_{self.prefill_graph}-D_{self.decode_graph}"
|
||||
+ ("-HMA_off" if self.hma_off else "")
|
||||
)
|
||||
|
||||
|
||||
_DEFAULT = "FULL_AND_PIECEWISE"
|
||||
CASES = [
|
||||
# Kimi-Linear: KDA + MLA groups.
|
||||
Case("kimi_linear", "ROCM_AITER_MLA", _DEFAULT, _DEFAULT),
|
||||
Case("kimi_linear", "ROCM_AITER_MLA", "NONE", "NONE"),
|
||||
Case("kimi_linear", "ROCM_AITER_MLA", "PIECEWISE", "PIECEWISE"),
|
||||
Case("kimi_linear", "ROCM_AITER_MLA", _DEFAULT, "FULL_DECODE_ONLY"),
|
||||
Case("kimi_linear", "TRITON_MLA", _DEFAULT, _DEFAULT),
|
||||
# Qwen3.5: Gated DeltaNet + full-attention groups.
|
||||
Case("qwen3_5_moe", "TRITON_ATTN", _DEFAULT, _DEFAULT),
|
||||
Case("qwen3_5_moe", "TRITON_ATTN", "NONE", "NONE"),
|
||||
Case("qwen3_5_moe", "TRITON_ATTN", "PIECEWISE", "PIECEWISE"),
|
||||
Case("qwen3_5_moe", "TRITON_ATTN", _DEFAULT, "FULL_DECODE_ONLY"),
|
||||
# ROCM_ATTN keeps a (2, num_blocks, ...) KV layout and opts out of KV
|
||||
# connectors, so the second ROCm attention backend is the unified one.
|
||||
Case("qwen3_5_moe", "ROCM_AITER_UNIFIED_ATTN", _DEFAULT, _DEFAULT),
|
||||
# Attention-only: prefill transfers the whole prompt.
|
||||
Case("qwen3_dense", "TRITON_ATTN", _DEFAULT, _DEFAULT),
|
||||
Case("qwen3_dense", "ROCM_AITER_UNIFIED_ATTN", _DEFAULT, _DEFAULT),
|
||||
# Without the hybrid KV cache manager: one cache group per model.
|
||||
Case("qwen3_dense", "TRITON_ATTN", _DEFAULT, _DEFAULT, hma_off=True),
|
||||
Case("deepseek_v3", "ROCM_AITER_MLA", _DEFAULT, _DEFAULT, hma_off=True),
|
||||
Case("gemma3", "TRITON_ATTN", _DEFAULT, _DEFAULT, hma_off=True),
|
||||
]
|
||||
|
||||
|
||||
def _graph_args(mode: str) -> list[str]:
|
||||
if mode == "NONE":
|
||||
return ["--enforce-eager"]
|
||||
# Inductor otherwise picks reduction launch configs by timing them, so each
|
||||
# separately compiled server may sum in a different order and drift by a
|
||||
# few bf16 logit steps; deterministic mode fixes the choice.
|
||||
config = {
|
||||
"cudagraph_mode": mode,
|
||||
"inductor_compile_config": {"deterministic": True},
|
||||
}
|
||||
return ["--compilation-config", json.dumps(config)]
|
||||
|
||||
|
||||
def _visible_devices() -> tuple[str, list[str]]:
|
||||
selectors = ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES")
|
||||
selector = next((k for k in selectors if k in os.environ), selectors[0])
|
||||
visible = os.getenv(selector)
|
||||
devices = (
|
||||
visible.split(",")
|
||||
if visible
|
||||
else [str(i) for i in range(current_platform.device_count())]
|
||||
)
|
||||
return selector, devices
|
||||
|
||||
|
||||
def _device_env(selector: str, device: str) -> dict[str, str]:
|
||||
# vLLM mirrors HIP_VISIBLE_DEVICES into CUDA_VISIBLE_DEVICES at import, so
|
||||
# the test process carries both; pin every selector or the server refuses
|
||||
# the mismatch. ROCR narrows the set first, so HIP/CUDA then index into it.
|
||||
env = {selector: device}
|
||||
for key in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
|
||||
if key != selector:
|
||||
env[key] = "0" if selector == "ROCR_VISIBLE_DEVICES" else device
|
||||
return env
|
||||
|
||||
|
||||
# KV consumers force enable_mm_embeds on, which moves multimodal-capable models
|
||||
# (Qwen3.5) from input_ids to inputs_embeds, but the AOT compile-cache key does
|
||||
# not cover it, so a consumer can load a producer's graph and crash in
|
||||
# profile_run. Keep each role's compile cache apart until
|
||||
# vllm-project/vllm#54642 lands (vllm-project/vllm#50891).
|
||||
_CACHE_ROOT = Path(tempfile.mkdtemp(prefix="moriio_dummy_e2e_cache_"))
|
||||
atexit.register(shutil.rmtree, _CACHE_ROOT, ignore_errors=True)
|
||||
|
||||
|
||||
def _server(
|
||||
arch: ArchCase,
|
||||
overrides: dict[str, Any],
|
||||
backend: str,
|
||||
graph: str,
|
||||
device: str,
|
||||
selector: str,
|
||||
*,
|
||||
port: int | None = None,
|
||||
extra_args: tuple[str, ...] = (),
|
||||
kv_transfer_config: dict[str, Any] | None = None,
|
||||
) -> RemoteOpenAIServer:
|
||||
port = port or get_open_port()
|
||||
args = [
|
||||
*arch.args,
|
||||
*extra_args,
|
||||
"--port",
|
||||
str(port),
|
||||
"--load-format",
|
||||
"dummy",
|
||||
"--worker-extension-cls",
|
||||
"moriio_wide_dummy.WideDummyWeights",
|
||||
"--hf-overrides",
|
||||
json.dumps(overrides),
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--attention-backend",
|
||||
backend,
|
||||
"--max-model-len",
|
||||
str(MAX_MODEL_LEN),
|
||||
"--max-num-batched-tokens",
|
||||
str(PREFILL_CHUNK),
|
||||
"--max-num-seqs",
|
||||
"16",
|
||||
"--kv-cache-memory-bytes",
|
||||
str(2 * 1024**3),
|
||||
"--gpu-memory-utilization",
|
||||
"0.5",
|
||||
"--enable-prefix-caching",
|
||||
*_graph_args(graph),
|
||||
]
|
||||
role = "direct"
|
||||
if kv_transfer_config is not None:
|
||||
args += ["--kv-transfer-config", json.dumps(kv_transfer_config)]
|
||||
role = kv_transfer_config["kv_role"]
|
||||
env = {
|
||||
"PATH": str(Path(sys.executable).parent) + os.pathsep + os.environ["PATH"],
|
||||
**_device_env(selector, device),
|
||||
"VLLM_HOST_IP": "127.0.0.1",
|
||||
"VLLM_PORT": str(get_open_port()),
|
||||
"VLLM_WORKER_MULTIPROC_METHOD": "spawn",
|
||||
"VLLM_SSM_CONV_STATE_LAYOUT": "DS",
|
||||
"VLLM_CACHE_ROOT": str(_CACHE_ROOT / role),
|
||||
# Only this directory, so the source tree cannot shadow installed vllm.
|
||||
"PYTHONPATH": os.pathsep.join(
|
||||
filter(None, [str(_WIDE_DUMMY_DIR), os.environ.get("PYTHONPATH")])
|
||||
),
|
||||
"MORIIO_DUMMY_WEIGHT_RANGE": str(DUMMY_WEIGHT_RANGE),
|
||||
**arch.env,
|
||||
}
|
||||
return RemoteOpenAIServer(
|
||||
arch.model, args, env_dict=env, auto_port=False, max_wait_seconds=900
|
||||
)
|
||||
|
||||
|
||||
def _moriio_config(role: str, http_port: int, ping_port: int) -> dict[str, Any]:
|
||||
return {
|
||||
"kv_connector": "MoRIIOConnector",
|
||||
"kv_role": role,
|
||||
"kv_connector_extra_config": {
|
||||
"backend": "xgmi",
|
||||
"read_mode": False,
|
||||
"host_ip": "127.0.0.1",
|
||||
"proxy_ip": "127.0.0.1",
|
||||
"proxy_ping_port": str(ping_port),
|
||||
"http_port": str(http_port),
|
||||
"handshake_port": str(get_open_port()),
|
||||
"notify_port": str(get_open_port()),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _proxy(ping_port: int) -> Iterator[str]:
|
||||
# Free ports for both sockets, so cases from parallel shards on one host
|
||||
# do not collide.
|
||||
port = get_open_port()
|
||||
process = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
str(_PROXY),
|
||||
"--port",
|
||||
str(port),
|
||||
"--service-discovery-port",
|
||||
str(ping_port),
|
||||
],
|
||||
start_new_session=True,
|
||||
)
|
||||
try:
|
||||
deadline = time.monotonic() + 60
|
||||
while time.monotonic() < deadline:
|
||||
assert process.poll() is None, "MoRIIO proxy exited during startup"
|
||||
try:
|
||||
requests.get(f"http://127.0.0.1:{port}/", timeout=1)
|
||||
break
|
||||
except requests.ConnectionError:
|
||||
time.sleep(0.2)
|
||||
else:
|
||||
raise TimeoutError("MoRIIO proxy did not start")
|
||||
yield f"http://127.0.0.1:{port}"
|
||||
finally:
|
||||
if process.poll() is None:
|
||||
os.killpg(process.pid, signal.SIGTERM)
|
||||
try:
|
||||
process.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
process.wait(timeout=10)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _pd_pair(
|
||||
arch: ArchCase,
|
||||
prefill_overrides: dict[str, Any],
|
||||
decode_overrides: dict[str, Any],
|
||||
case: Case,
|
||||
) -> Iterator[tuple[str, RemoteOpenAIServer]]:
|
||||
selector, devices = _visible_devices()
|
||||
assert len(devices) >= 2, "needs one GPU for prefill and one for decode"
|
||||
prefill_port, decode_port, ping_port = (get_open_port() for _ in range(3))
|
||||
with _proxy(ping_port) as proxy_url, contextlib.ExitStack() as stack:
|
||||
servers: list[RemoteOpenAIServer] = []
|
||||
stack.callback(RemoteOpenAIServer.shutdown_many, servers)
|
||||
with ThreadPoolExecutor(2) as pool:
|
||||
futures = [
|
||||
pool.submit(
|
||||
_server,
|
||||
arch,
|
||||
prefill_overrides,
|
||||
case.attention_backend,
|
||||
case.prefill_graph,
|
||||
devices[0],
|
||||
selector,
|
||||
port=prefill_port,
|
||||
extra_args=case.extra_args,
|
||||
kv_transfer_config=_moriio_config(
|
||||
"kv_producer", prefill_port, ping_port
|
||||
),
|
||||
),
|
||||
pool.submit(
|
||||
_server,
|
||||
arch,
|
||||
decode_overrides,
|
||||
case.attention_backend,
|
||||
case.decode_graph,
|
||||
devices[1],
|
||||
selector,
|
||||
port=decode_port,
|
||||
extra_args=case.extra_args,
|
||||
kv_transfer_config=_moriio_config(
|
||||
"kv_consumer", decode_port, ping_port
|
||||
),
|
||||
),
|
||||
]
|
||||
for future in futures:
|
||||
servers.append(future.result())
|
||||
_wait_for_registration(proxy_url, arch.model)
|
||||
yield proxy_url, servers[1]
|
||||
|
||||
|
||||
def _wait_for_registration(proxy_url: str, model: str) -> None:
|
||||
deadline = time.monotonic() + 120
|
||||
while time.monotonic() < deadline:
|
||||
# The proxy hands decode max_tokens - 1 (prefill emits the first
|
||||
# token), so the probe needs at least 2 for decode to accept it.
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{proxy_url}/v1/completions",
|
||||
json={"model": model, "prompt": [1000, 1001], "max_tokens": 2},
|
||||
timeout=300,
|
||||
)
|
||||
# The proxy streams the decode body, so a decode-side error can
|
||||
# surface as a truncated response rather than a status code.
|
||||
if response.status_code == 200 and response.json()["choices"]:
|
||||
return
|
||||
except requests.RequestException:
|
||||
pass
|
||||
time.sleep(1)
|
||||
raise TimeoutError("prefill/decode never registered with the MoRIIO proxy")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Measurement
|
||||
|
||||
|
||||
def _metric(server: RemoteOpenAIServer, name: str) -> list[Any]:
|
||||
from prometheus_client.parser import text_string_to_metric_families
|
||||
|
||||
response = requests.get(server.url_for("metrics"), timeout=10)
|
||||
response.raise_for_status()
|
||||
return [
|
||||
sample
|
||||
for family in text_string_to_metric_families(response.text)
|
||||
for sample in family.samples
|
||||
if sample.name == name
|
||||
]
|
||||
|
||||
|
||||
def _block_size(server: RemoteOpenAIServer) -> int:
|
||||
# Hybrid models pad the attention block to the recurrent-state page, so the
|
||||
# effective size is only known once the server has started.
|
||||
sizes = {
|
||||
int(sample.labels["block_size"])
|
||||
for sample in _metric(server, "vllm:cache_config_info")
|
||||
}
|
||||
assert len(sizes) == 1, f"expected one cache block size, got {sizes}"
|
||||
return sizes.pop()
|
||||
|
||||
|
||||
def _prompt_lengths(arch: ArchCase, server: RemoteOpenAIServer) -> list[int]:
|
||||
# Hybrid chunks end on cache-block boundaries; attention-only chunks end on
|
||||
# the prefill budget, so the last token is computed on its own either way.
|
||||
if arch.hybrid:
|
||||
unit, counts = _block_size(server), ALIGNED_BLOCKS
|
||||
else:
|
||||
unit, counts = PREFILL_CHUNK, ALIGNED_CHUNKS
|
||||
aligned = [
|
||||
k * unit + 1 for k in counts if k * unit + 1 + MAX_TOKENS <= MAX_MODEL_LEN
|
||||
]
|
||||
assert aligned, f"alignment unit {unit} leaves no aligned prompt"
|
||||
return [SHORT_PROMPT_LEN, *aligned]
|
||||
|
||||
|
||||
def _prompts(lengths: list[int]) -> list[list[int]]:
|
||||
rng = random.Random(0)
|
||||
return [[rng.randrange(1000, 30000) for _ in range(n)] for n in lengths]
|
||||
|
||||
|
||||
def _first_step(url: str, model: str, prompt: list[int]) -> dict[str, float]:
|
||||
response = requests.post(
|
||||
f"{url}/v1/completions",
|
||||
json={
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"max_tokens": MAX_TOKENS,
|
||||
"temperature": 0,
|
||||
"logprobs": TOP_K,
|
||||
"return_tokens_as_token_ids": True,
|
||||
},
|
||||
timeout=600,
|
||||
)
|
||||
response.raise_for_status()
|
||||
top = response.json()["choices"][0]["logprobs"]["top_logprobs"]
|
||||
assert top, "server returned no logprobs"
|
||||
return top[0]
|
||||
|
||||
|
||||
def _distance(ref: dict[str, float], got: dict[str, float]) -> float:
|
||||
"""Max difference of centered log-probs over the shared top-k tokens.
|
||||
|
||||
Centering removes the log-sum-exp offset, leaving the logit differences the
|
||||
KV actually determines.
|
||||
"""
|
||||
common = sorted(ref.keys() & got.keys())
|
||||
if len(common) < len(ref) // 2:
|
||||
return float("inf")
|
||||
ref_mean = sum(ref[t] for t in common) / len(common)
|
||||
got_mean = sum(got[t] for t in common) / len(common)
|
||||
return max(abs((ref[t] - ref_mean) - (got[t] - got_mean)) for t in common)
|
||||
|
||||
|
||||
def _external_hits(server: RemoteOpenAIServer) -> float:
|
||||
samples = _metric(server, "vllm:external_prefix_cache_hits_total")
|
||||
return sum(sample.value for sample in samples)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Reference:
|
||||
prompts: list[list[int]]
|
||||
logprobs: list[dict[str, float]]
|
||||
hybrid: bool
|
||||
|
||||
@property
|
||||
def expected_external_hits(self) -> int:
|
||||
# Hybrid prefill stops one token short; attention-only prefill hands
|
||||
# decode the whole prompt.
|
||||
return sum(len(p) - int(self.hybrid) for p in self.prompts)
|
||||
|
||||
|
||||
def _reference(arch: ArchCase, overrides: dict[str, Any], case: Case) -> Reference:
|
||||
"""Direct serve; also checks that the comparison can see a wrong prefix."""
|
||||
selector, devices = _visible_devices()
|
||||
# Prefill graphs run the chunks and the decode-only last step runs the same
|
||||
# captured graph in every decode mode used here, so the prefill mode
|
||||
# reproduces both halves of the P/D computation.
|
||||
with _server(
|
||||
arch,
|
||||
overrides,
|
||||
case.attention_backend,
|
||||
case.prefill_graph,
|
||||
devices[0],
|
||||
selector,
|
||||
extra_args=case.extra_args,
|
||||
) as server:
|
||||
url = server.url_root
|
||||
lengths = _prompt_lengths(arch, server)
|
||||
prompts = _prompts(lengths)
|
||||
refs = [_first_step(url, arch.model, p) for p in prompts]
|
||||
# Same tail, different early context: the comparison has to see it.
|
||||
for prompt, ref in zip(prompts, refs):
|
||||
altered = [t + 1 for t in prompt[: len(prompt) // 2]]
|
||||
altered += prompt[len(prompt) // 2 :]
|
||||
moved = _distance(ref, _first_step(url, arch.model, altered))
|
||||
print(f"[sensitivity] len={len(prompt)} context-change distance={moved}")
|
||||
assert moved >= MIN_CONTEXT_EFFECT, (
|
||||
f"first-token log-probs barely depend on the context (moved "
|
||||
f"{moved:.4g} for len {len(prompt)}); a wrong prefix would go "
|
||||
"unnoticed"
|
||||
)
|
||||
# Determinism of the reference itself (served from its prefix cache).
|
||||
again = _first_step(url, arch.model, prompts[-1])
|
||||
print(f"[noise] repeat distance={_distance(refs[-1], again)}")
|
||||
return Reference(prompts, refs, arch.hybrid)
|
||||
|
||||
|
||||
def _pd_distances(
|
||||
arch: ArchCase,
|
||||
prefill_overrides: dict[str, Any],
|
||||
decode_overrides: dict[str, Any],
|
||||
case: Case,
|
||||
ref: Reference,
|
||||
) -> tuple[list[float], float]:
|
||||
with _pd_pair(arch, prefill_overrides, decode_overrides, case) as (
|
||||
proxy_url,
|
||||
decode,
|
||||
):
|
||||
hits_before = _external_hits(decode)
|
||||
distances = [
|
||||
_distance(expected, _first_step(proxy_url, arch.model, prompt))
|
||||
for prompt, expected in zip(ref.prompts, ref.logprobs)
|
||||
]
|
||||
hits = _external_hits(decode) - hits_before
|
||||
lengths = [len(p) for p in ref.prompts]
|
||||
print(f"[pd] {case.id} lens={lengths} distances={distances} external_hits={hits}")
|
||||
return distances, hits
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", CASES, ids=lambda c: c.id)
|
||||
def test_write_matches_direct_serve(case: Case):
|
||||
arch = ARCHS[case.arch]
|
||||
overrides = _arch_overrides(arch)
|
||||
ref = _reference(arch, overrides, case)
|
||||
distances, hits = _pd_distances(arch, overrides, overrides, case, ref)
|
||||
|
||||
# Every transferred token of every prompt must come from prefill.
|
||||
assert hits == ref.expected_external_hits, (
|
||||
f"decode reported {hits} external prefix-cache hits, expected "
|
||||
f"{ref.expected_external_hits}; it did not use all the transferred KV"
|
||||
)
|
||||
lengths = [len(p) for p in ref.prompts]
|
||||
bad = {n: d for n, d in zip(lengths, distances) if not d <= TOLERANCE}
|
||||
assert not bad, f"P/D first-token log-probs differ from direct serve: {bad}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("arch_name", sorted(ARCHS))
|
||||
def test_corrupted_prefill_state_is_detected(arch_name: str):
|
||||
"""Negative control: KV computed with different values must fail."""
|
||||
arch = ARCHS[arch_name]
|
||||
default = next(c for c in CASES if c.arch == arch_name)
|
||||
case = Case(
|
||||
arch_name, default.attention_backend, _DEFAULT, _DEFAULT, default.hma_off
|
||||
)
|
||||
overrides = _arch_overrides(arch)
|
||||
ref = _reference(arch, overrides, case)
|
||||
corrupted = _merge(overrides, arch.corrupt_overrides)
|
||||
distances, hits = _pd_distances(arch, corrupted, overrides, case, ref)
|
||||
|
||||
assert hits == ref.expected_external_hits, (
|
||||
"decode did not use the transferred KV; the control proves nothing"
|
||||
)
|
||||
assert all(d > TOLERANCE for d in distances), (
|
||||
f"P/D outputs with wrongly computed prefill KV still matched the "
|
||||
f"reference (distances {distances}); the comparison cannot detect "
|
||||
"transfer corruption"
|
||||
)
|
||||
@@ -146,6 +146,7 @@ def _write_consumer_scheduler_for_finished_request(tp_size: int = 2):
|
||||
scheduler.is_producer = False
|
||||
scheduler.mode = MoRIIOMode.WRITE
|
||||
scheduler.tp_size = tp_size
|
||||
scheduler._has_mamba = False
|
||||
scheduler._reqs_need_recv = {}
|
||||
scheduler.unmap_request_id = MagicMock()
|
||||
return scheduler
|
||||
@@ -158,6 +159,10 @@ def _write_producer_scheduler(block_size: int = 1) -> Any:
|
||||
scheduler.mode = MoRIIOMode.WRITE
|
||||
scheduler.is_producer = True
|
||||
scheduler.block_size = block_size
|
||||
scheduler._has_mamba = False
|
||||
scheduler.kv_cache_config = SimpleNamespace(
|
||||
kv_cache_groups=[SimpleNamespace(kv_cache_spec=object())]
|
||||
)
|
||||
scheduler.blocks_per_sw = [0]
|
||||
scheduler._group_block_sizes = [block_size]
|
||||
scheduler.transfer_id_to_request_id = {}
|
||||
@@ -944,6 +949,7 @@ def test_clamp_to_prompt_blocks_clamps_each_group_by_its_own_block_size():
|
||||
# their own prompt-block count; a sliding-window group is left untouched.
|
||||
scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler)
|
||||
scheduler.blocks_per_sw = [0, 0, 3]
|
||||
scheduler._has_mamba = False
|
||||
scheduler._group_block_sizes = [16, 32, 16]
|
||||
|
||||
req = SimpleNamespace(num_prompt_tokens=40)
|
||||
|
||||
@@ -12,12 +12,13 @@ wiring.
|
||||
import importlib
|
||||
from dataclasses import asdict
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.v1.core.block_pool import BlockPool
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.core.sched.output import KVConnectorBlockState, SchedulerOutput
|
||||
from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec, UniformTypeKVCacheSpecs
|
||||
|
||||
@@ -52,6 +53,7 @@ class _FakeScheduler(moriio_connector.MoRIIOConnectorScheduler): # type: ignore
|
||||
attributes read by the method under test are set."""
|
||||
|
||||
def __init__(self, **attrs):
|
||||
self.mode = MoRIIOMode.READ
|
||||
self._mamba_group_ids: list[int] = []
|
||||
self._attn_group_ids: list[int] = [0]
|
||||
self._num_ssm_scratch_blocks = 0
|
||||
@@ -287,7 +289,8 @@ def test_scheduler_rejects_mamba_group_without_two_states():
|
||||
moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config)
|
||||
|
||||
|
||||
def test_scheduler_rejects_non_dspark_speculative_hybrid_read():
|
||||
@pytest.mark.parametrize("read_mode", [True, False])
|
||||
def test_scheduler_rejects_non_dspark_speculative_hybrid(read_mode):
|
||||
config = SimpleNamespace(
|
||||
kv_cache_groups=[
|
||||
SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
|
||||
@@ -298,7 +301,9 @@ def test_scheduler_rejects_non_dspark_speculative_hybrid_read():
|
||||
|
||||
with pytest.raises(moriio_common.MoRIIOError, match="speculative decoding"):
|
||||
moriio_connector.MoRIIOConnectorScheduler(
|
||||
_gate_vllm_config(speculative_config=_spec_config("ngram")),
|
||||
_gate_vllm_config(
|
||||
read_mode=read_mode, speculative_config=_spec_config("ngram")
|
||||
),
|
||||
"engine",
|
||||
config,
|
||||
)
|
||||
@@ -336,14 +341,16 @@ def _dspark_wrapped_attention_config(*, dcp_sharded: bool = True):
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("restore_moriio_role")
|
||||
@pytest.mark.parametrize("read_mode", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
"dcp_size,dcp_sharded,expected_tail", [(1, True, 4), (4, True, 1), (4, False, 4)]
|
||||
)
|
||||
def test_scheduler_accepts_dspark_with_wrapped_full_attention(
|
||||
dcp_size, dcp_sharded, expected_tail
|
||||
dcp_size, dcp_sharded, expected_tail, read_mode
|
||||
):
|
||||
config = _dspark_wrapped_attention_config(dcp_sharded=dcp_sharded)
|
||||
vllm_config = _gate_vllm_config(
|
||||
read_mode=read_mode,
|
||||
speculative_config=_spec_config("dspark"),
|
||||
block_size=4,
|
||||
num_lookahead_tokens=7,
|
||||
@@ -384,19 +391,234 @@ def test_scheduler_accepts_dspark_with_wrapped_full_attention(
|
||||
)
|
||||
|
||||
|
||||
def test_scheduler_rejects_hybrid_write():
|
||||
config = SimpleNamespace(
|
||||
kv_cache_groups=[
|
||||
SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
|
||||
SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()),
|
||||
],
|
||||
@pytest.mark.usefixtures("restore_moriio_role")
|
||||
def test_hybrid_write_clamps_attention_and_excludes_scratch_slots():
|
||||
config = _dspark_wrapped_attention_config()
|
||||
config.select_transfer_block_ids = lambda blocks: tuple(blocks)
|
||||
sched = moriio_connector.MoRIIOConnectorScheduler(
|
||||
_gate_vllm_config(read_mode=False, dcp_size=2), "engine", config
|
||||
)
|
||||
config.transfer_groups = tuple(config.kv_cache_groups)
|
||||
assert sched._clamp_to_prompt_blocks(
|
||||
SimpleNamespace(num_prompt_tokens=9), [[10, 11, 12, 13], [1, 2, 3, 4]]
|
||||
) == [[11], [1, 2]]
|
||||
|
||||
with pytest.raises(moriio_common.MoRIIOError, match="READ mode only"):
|
||||
moriio_connector.MoRIIOConnectorScheduler(
|
||||
_gate_vllm_config(read_mode=False), "engine", config
|
||||
|
||||
@pytest.mark.usefixtures("restore_moriio_role")
|
||||
def test_hybrid_write_resolves_current_state_instead_of_cached_block_ids():
|
||||
config = _dspark_wrapped_attention_config()
|
||||
config.select_transfer_block_ids = lambda blocks: tuple(blocks)
|
||||
sched = moriio_connector.MoRIIOConnectorScheduler(
|
||||
_gate_vllm_config(read_mode=False, dcp_size=2), "engine", config
|
||||
)
|
||||
state = KVConnectorBlockState(
|
||||
req_ids={"req"},
|
||||
resolve_block_ids=lambda _: ([0, 51, 52, 53], [31, 32, 33]),
|
||||
boundary_state_offloads={},
|
||||
)
|
||||
assert sched._get_write_blocks(
|
||||
SimpleNamespace(request_id="req", num_prompt_tokens=9),
|
||||
[[10, 11, 12], [1]],
|
||||
SimpleNamespace(kv_connector_block_state=state),
|
||||
) == [[51], [31, 32]]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("restore_moriio_role")
|
||||
@pytest.mark.parametrize("external_tokens", [0, 9])
|
||||
def test_hybrid_write_notifies_grouped_slots_or_releases_a_complete_local_hit(
|
||||
external_tokens,
|
||||
):
|
||||
config = _dspark_wrapped_attention_config()
|
||||
config.select_transfer_block_ids = lambda blocks: tuple(blocks)
|
||||
sched = moriio_connector.MoRIIOConnectorScheduler(
|
||||
_gate_vllm_config(read_mode=False), "engine", config
|
||||
)
|
||||
sched.send_notify_block = Mock()
|
||||
sched._send_transfer_release = Mock()
|
||||
req = _mk_request(
|
||||
list(range(10)),
|
||||
params={
|
||||
"do_remote_prefill": True,
|
||||
"transfer_id": "tx",
|
||||
"remote_host": "127.0.0.2",
|
||||
"remote_notify_port": 7000,
|
||||
},
|
||||
)
|
||||
req.request_id = "req"
|
||||
sched.update_state_after_alloc(
|
||||
req, _FakeBlocks(([10, 11, 12], [1, 2, 3])), external_tokens
|
||||
)
|
||||
if external_tokens:
|
||||
sched.send_notify_block.assert_called_once_with(
|
||||
req_id="req",
|
||||
transfer_id="tx",
|
||||
block_notify_list=[[10], [1, 2, 3]],
|
||||
host="127.0.0.2",
|
||||
port=7000,
|
||||
)
|
||||
sched._send_transfer_release.assert_not_called()
|
||||
else:
|
||||
sched.send_notify_block.assert_not_called()
|
||||
sched._send_transfer_release.assert_called_once_with("tx", "127.0.0.2", 7000)
|
||||
assert not req.kv_transfer_params["do_remote_prefill"]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("restore_moriio_role")
|
||||
def test_hybrid_write_keeps_separate_draft_groups_and_their_block_spans():
|
||||
config = _dspark_wrapped_attention_config()
|
||||
# A disabled cache is absent from the wire. Target and draft may have
|
||||
# different token spans, and neither can borrow the other's block IDs.
|
||||
draft = SimpleNamespace(
|
||||
enable_kv_transfer=True,
|
||||
kv_cache_spec=FullAttentionSpec(
|
||||
block_size=4,
|
||||
num_kv_heads=1,
|
||||
head_size=2,
|
||||
dtype=torch.float32,
|
||||
dcp_sharded=False,
|
||||
),
|
||||
)
|
||||
disabled = SimpleNamespace(enable_kv_transfer=False, kv_cache_spec=object())
|
||||
config.kv_cache_groups.insert(1, disabled)
|
||||
config.kv_cache_groups.append(draft)
|
||||
config.transfer_groups = tuple(
|
||||
g for g in config.kv_cache_groups if g.enable_kv_transfer
|
||||
)
|
||||
config.select_transfer_block_ids = lambda blocks: (blocks[0], blocks[2], blocks[3])
|
||||
sched = moriio_connector.MoRIIOConnectorScheduler(
|
||||
_gate_vllm_config(
|
||||
read_mode=False, dcp_size=2, speculative_config=_spec_config("dspark")
|
||||
),
|
||||
"engine",
|
||||
config,
|
||||
)
|
||||
req = _mk_request(
|
||||
list(range(9)),
|
||||
params={
|
||||
"do_remote_decode": True,
|
||||
"transfer_id": "tx",
|
||||
},
|
||||
)
|
||||
req.request_id = "req"
|
||||
req.status = moriio_connector.RequestStatus.FINISHED_LENGTH_CAPPED
|
||||
blocks = ([10, 11, 12, 13], [99], [20, 21, 22], [30, 31, 32, 33])
|
||||
expected = [[11], [20, 21], [30, 31, 32]]
|
||||
assert sched._clamp_to_prompt_blocks(req, blocks) == expected
|
||||
# Exercise the connector entry point too: completion must not flatten
|
||||
# a multi-group WRITE payload through READ's single-attention helper.
|
||||
delay_free, params = _FakeConnector(sched).request_finished_all_groups(req, blocks)
|
||||
assert delay_free
|
||||
assert params["remote_block_ids"] == expected
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("restore_moriio_role")
|
||||
@pytest.mark.parametrize(
|
||||
"blocks,delay_free",
|
||||
[(([0, 51, 52, 53], [31, 32, 33]), True), (([], []), False)],
|
||||
)
|
||||
def test_hybrid_write_request_finished_defers_free_only_with_blocks_to_send(
|
||||
blocks, delay_free
|
||||
):
|
||||
config = _dspark_wrapped_attention_config()
|
||||
config.select_transfer_block_ids = lambda blocks: tuple(blocks)
|
||||
sched = moriio_connector.MoRIIOConnectorScheduler(
|
||||
_gate_vllm_config(read_mode=False, dcp_size=2), "engine", config
|
||||
)
|
||||
sched.is_producer = True
|
||||
req = _mk_request(
|
||||
list(range(9)), params={"do_remote_decode": True, "transfer_id": "tx"}
|
||||
)
|
||||
req.request_id = "req"
|
||||
req.status = moriio_connector.RequestStatus.FINISHED_LENGTH_CAPPED
|
||||
|
||||
assert sched.request_finished(req, blocks)[0] is delay_free
|
||||
# The WRITE ACK, not the READ watchdog, releases the deferred blocks.
|
||||
if delay_free:
|
||||
assert sched._deferred_send_deadlines["req"][1] == "tx"
|
||||
else:
|
||||
assert "req" not in sched._deferred_send_deadlines
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("restore_moriio_role")
|
||||
def test_hybrid_write_build_meta_saves_current_state_on_the_final_chunk():
|
||||
"""Chunked prefill must not ship block ids cached at admission: align-mode
|
||||
state slots move between chunks without appending a block."""
|
||||
config = _dspark_wrapped_attention_config()
|
||||
config.select_transfer_block_ids = lambda blocks: tuple(blocks)
|
||||
sched = moriio_connector.MoRIIOConnectorScheduler(
|
||||
_gate_vllm_config(read_mode=False, dcp_size=2), "engine", config
|
||||
)
|
||||
sched.is_producer = True
|
||||
set_role(moriio_common.ROLE.PRODUCER)
|
||||
params = {
|
||||
"do_remote_decode": True,
|
||||
"transfer_id": "tx",
|
||||
"remote_engine_id": "decode",
|
||||
"remote_block_ids": [],
|
||||
"remote_host": "127.0.0.1",
|
||||
"remote_handshake_port": 6000,
|
||||
"remote_notify_port": 6001,
|
||||
}
|
||||
req = _mk_request(list(range(9)), params=params)
|
||||
req.request_id = "req"
|
||||
req.num_computed_tokens = 0
|
||||
sched._reqs_need_save["req"] = (req, [[10, 11, 12], [1]])
|
||||
sched._req_kv_params["req"] = dict(params)
|
||||
|
||||
def _output(cached_req_ids, new_block_ids, num_scheduled):
|
||||
return SimpleNamespace(
|
||||
scheduled_cached_reqs=SimpleNamespace(
|
||||
req_ids=cached_req_ids, new_block_ids=new_block_ids
|
||||
),
|
||||
num_scheduled_tokens={"req": num_scheduled},
|
||||
new_block_ids_to_zero=None,
|
||||
kv_connector_block_state=KVConnectorBlockState(
|
||||
req_ids={"req"},
|
||||
resolve_block_ids=lambda _: ([0, 51, 52, 53], [31, 32, 33]),
|
||||
boundary_state_offloads={},
|
||||
),
|
||||
)
|
||||
|
||||
# First chunk: not final, so nothing is saved yet.
|
||||
meta = sched.build_connector_meta(_output([], [], 4))
|
||||
assert not meta.reqs_to_save
|
||||
assert "req" in sched._reqs_need_pending_save
|
||||
|
||||
# Final chunk: stale cached ids plus new ids are ignored in favour of
|
||||
# the scheduler's current table, clamped to the prompt.
|
||||
req.num_computed_tokens = 4
|
||||
meta = sched.build_connector_meta(_output(["req"], [([77], [78])], 5))
|
||||
assert meta.reqs_to_save["req"].local_block_ids == [[51], [31, 32]]
|
||||
assert "req" not in sched._reqs_need_pending_save
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mode,deadlines,expected",
|
||||
[
|
||||
(MoRIIOMode.WRITE, {"req": (0.0, "tx")}, True),
|
||||
(MoRIIOMode.WRITE, {}, False),
|
||||
(MoRIIOMode.READ, {"req": (0.0, "tx")}, False),
|
||||
],
|
||||
)
|
||||
def test_has_pending_push_work_tracks_write_deferred_frees(mode, deadlines, expected):
|
||||
connector = _FakeConnector(_FakeScheduler(_deferred_send_deadlines=deadlines))
|
||||
connector.mode = mode
|
||||
assert connector.has_pending_push_work() is expected
|
||||
|
||||
|
||||
def test_has_pending_push_work_is_false_on_worker_role():
|
||||
connector = _FakeConnector(None)
|
||||
connector.mode = MoRIIOMode.WRITE
|
||||
assert connector.has_pending_push_work() is False
|
||||
|
||||
|
||||
def test_loaded_kv_cache_groups_are_every_transfer_group():
|
||||
# The connector restores the request's own blocks, not a hash-addressed
|
||||
# prefix, so a transfer-enabled Mamba group counts as loaded too.
|
||||
connector = _FakeConnector(None)
|
||||
connector._kv_cache_config = SimpleNamespace(
|
||||
transfer_group_ids=(0, 2), prefix_cacheable_group_ids=(0,)
|
||||
)
|
||||
assert connector.get_loaded_kv_cache_group_ids(SimpleNamespace()) == (0, 2)
|
||||
|
||||
|
||||
def test_split_block_groups_rejects_non_hybrid_use():
|
||||
@@ -728,11 +950,15 @@ def test_session_build_rejects_per_layer_region_count_mismatch():
|
||||
worker._get_built_session("prefill")
|
||||
|
||||
|
||||
def test_register_kv_caches_rejects_hybrid_write_before_registration():
|
||||
worker = _FakeWorker(mode=MoRIIOMode.WRITE, _transfer_layer_names={"kda.0"})
|
||||
def test_register_kv_caches_rejects_invalid_hybrid_write_state_layout():
|
||||
worker = _FakeWorker(
|
||||
mode=MoRIIOMode.WRITE,
|
||||
_transfer_layer_names={"kda.0"},
|
||||
layer_to_spec={"kda.0": _mamba_spec(1)},
|
||||
)
|
||||
worker._is_mamba_layer = lambda _layer_name: True
|
||||
|
||||
with pytest.raises(moriio_common.MoRIIOError, match="READ mode only"):
|
||||
with pytest.raises(moriio_common.MoRIIOError, match="exactly two Mamba states"):
|
||||
worker.register_kv_caches({"kda.0": object()})
|
||||
|
||||
|
||||
@@ -933,6 +1159,40 @@ def test_read_matches_only_requests_with_pending_remote_prefill(params, expected
|
||||
assert scheduler.get_num_new_matched_tokens(request, 0) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"has_mamba,expected", [(False, (32, True)), (True, (31, True))]
|
||||
)
|
||||
def test_write_consumer_matches_full_prompt_unless_hybrid(has_mamba, expected):
|
||||
"""Without a mamba group (attention-only, or HMA disabled) prefill pushed
|
||||
KV for all N tokens; with one it pushed h(N-1)."""
|
||||
scheduler = _FakeScheduler(
|
||||
is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=has_mamba
|
||||
)
|
||||
request = SimpleNamespace(
|
||||
num_prompt_tokens=32, kv_transfer_params={"do_remote_prefill": True}
|
||||
)
|
||||
|
||||
assert scheduler.get_num_new_matched_tokens(request, 0) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"is_producer,has_mamba,expected_prompt_len",
|
||||
[(True, True, 4), (True, False, 5), (False, True, 5)],
|
||||
)
|
||||
def test_on_new_request_truncates_only_hybrid_producers(
|
||||
is_producer, has_mamba, expected_prompt_len
|
||||
):
|
||||
"""Attention-only prefill (including models run with HMA disabled)
|
||||
computes all N prompt tokens; only a hybrid producer drops the last."""
|
||||
scheduler = _FakeScheduler(is_producer=is_producer, _has_mamba=has_mamba)
|
||||
request = _mk_request([10, 11, 12, 13, 14], params={"do_remote_decode": True})
|
||||
|
||||
scheduler.on_new_request(request)
|
||||
|
||||
assert request.num_prompt_tokens == expected_prompt_len
|
||||
assert len(request.prompt_token_ids) == expected_prompt_len
|
||||
|
||||
|
||||
def test_read_rematch_after_alloc_keeps_queued_transfer():
|
||||
"""Clearing the flag blocks another external hit, not the queued READ."""
|
||||
sched = _make_read_scheduler()
|
||||
@@ -981,12 +1241,37 @@ def test_get_num_new_matched_tokens_write_plain_keeps_all_tokens():
|
||||
assert is_async is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("params", [None, {}, {"do_remote_prefill": False}])
|
||||
def test_write_does_not_wait_again_after_receive_or_for_local_requests(params):
|
||||
sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=True)
|
||||
req = _mk_request(list(range(10)))
|
||||
req.kv_transfer_params = params
|
||||
assert sched.get_num_new_matched_tokens(req, 0) == (0, False)
|
||||
|
||||
|
||||
def test_write_never_declares_an_empty_async_load():
|
||||
sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=False)
|
||||
req = _mk_request(list(range(10)), params={"do_remote_prefill": True})
|
||||
assert sched.get_num_new_matched_tokens(req, 10) == (0, False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prompt_len,computed,expected", [(16, 0, 15), (17, 16, 0), (1, 0, 0)]
|
||||
)
|
||||
def test_hybrid_write_advertises_only_the_state_computed_by_prefill(
|
||||
prompt_len, computed, expected
|
||||
):
|
||||
sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=True)
|
||||
req = _mk_request(list(range(prompt_len)), params={"do_remote_prefill": True})
|
||||
assert sched.get_num_new_matched_tokens(req, computed) == (expected, expected > 0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "num_computed_tokens", "expected", "is_async"),
|
||||
[
|
||||
(MoRIIOMode.READ, 0, 9, False),
|
||||
(MoRIIOMode.READ, 10, 0, False),
|
||||
(MoRIIOMode.WRITE, 2, 8, True),
|
||||
(MoRIIOMode.WRITE, 2, 7, True),
|
||||
],
|
||||
)
|
||||
def test_get_num_new_matched_tokens_supports_embeds_only_prompts(
|
||||
|
||||
@@ -7,6 +7,7 @@ import threading
|
||||
from collections import OrderedDict
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -106,7 +107,7 @@ def _write_task(layer_name: str, transfer_id: str = "xfer") -> Any:
|
||||
request_id="req",
|
||||
transfer_id=transfer_id,
|
||||
dst_engine_id="remote-engine",
|
||||
local_block_ids=[1, 3],
|
||||
local_block_ids=[[1, 3]],
|
||||
remote_block_ids_hint=None,
|
||||
layer_name=layer_name,
|
||||
event=None,
|
||||
@@ -345,6 +346,8 @@ def test_write_transfer_plan_caches_offsets_per_layer_geometry():
|
||||
calls: list[str] = []
|
||||
|
||||
class FakeWorker:
|
||||
_has_mamba: bool
|
||||
layer_to_group: dict[str, int]
|
||||
kv_caches: dict[str, torch.Tensor]
|
||||
layer_name_to_local_kv_cache_metadata: dict[str, list[Any]]
|
||||
|
||||
@@ -364,11 +367,13 @@ def test_write_transfer_plan_caches_offsets_per_layer_geometry():
|
||||
return ([call_id], [call_id + 10], [call_id + 20])
|
||||
|
||||
fake_worker = FakeWorker()
|
||||
fake_worker._has_mamba = False
|
||||
fake_worker.layer_to_group = dict.fromkeys(kv_caches, 0)
|
||||
fake_worker.kv_caches = kv_caches
|
||||
fake_worker.layer_name_to_local_kv_cache_metadata = {name: [] for name in kv_caches}
|
||||
writer = MoRIIOWriter.__new__(MoRIIOWriter)
|
||||
writer._worker_ref = lambda: fake_worker
|
||||
request_info = RemoteAllocInfo(block_ids=[4, 5])
|
||||
request_info = RemoteAllocInfo(block_ids=[[4, 5]])
|
||||
remote_meta = _remote_meta()
|
||||
|
||||
dense0_plan = writer._prepare_transfer_plan(
|
||||
@@ -395,7 +400,7 @@ def test_write_transfer_plan_caches_offsets_per_layer_geometry():
|
||||
|
||||
|
||||
def test_write_scheduler_deduplicates_layers_and_seals_expected_count():
|
||||
request_info = RemoteAllocInfo(block_ids=[4, 5])
|
||||
request_info = RemoteAllocInfo(block_ids=[[4, 5]])
|
||||
wrapper = _wrapper_for_messages()
|
||||
wrapper.done_remote_allocate_req_dict["xfer"] = request_info
|
||||
writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper))
|
||||
@@ -411,6 +416,80 @@ def test_write_scheduler_deduplicates_layers_and_seals_expected_count():
|
||||
assert writer._sealed_writes["xfer"] == 2
|
||||
|
||||
|
||||
def test_hybrid_write_attention_uses_its_own_group_not_the_first_group():
|
||||
worker = SimpleNamespace(
|
||||
_has_mamba=True,
|
||||
kv_caches={"draft": torch.empty((8, 1, 4, 3))},
|
||||
layer_to_group={"draft": 2},
|
||||
_region_session_indices=lambda _: [5],
|
||||
_compute_block_transfer_offsets=Mock(return_value=([1], [2], [3])),
|
||||
)
|
||||
writer = make_moriio_writer(worker)
|
||||
task = _write_task("draft")
|
||||
task.local_block_ids = [[10], [20], [30, 31]]
|
||||
remote = RemoteAllocInfo(block_ids=[[40], [50], [60, 61, 62]])
|
||||
meta = _remote_meta()
|
||||
plan = writer._prepare_transfer_plan(task, remote, meta)
|
||||
worker._compute_block_transfer_offsets.assert_called_once_with(
|
||||
"draft", [30, 31], [60, 61], meta, remote_engine_id=task.dst_engine_id
|
||||
)
|
||||
assert plan.sess_idx == 5
|
||||
|
||||
|
||||
def test_hybrid_write_posts_both_regions_before_layer_completion():
|
||||
from unittest.mock import Mock
|
||||
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
|
||||
MoRIIOConnectorWorker,
|
||||
)
|
||||
|
||||
wrapper = _wrapper_for_messages()
|
||||
info = RemoteAllocInfo(block_ids=[[], [7], [9]])
|
||||
wrapper.done_remote_allocate_req_dict["xfer"] = info
|
||||
wrapper.write_remote_data = Mock(side_effect=["conv-status", "ssm-status"])
|
||||
worker = SimpleNamespace(
|
||||
moriio_wrapper=wrapper,
|
||||
tp_rank=0,
|
||||
get_engine_name_with_dp=lambda name, rank: name,
|
||||
_get_built_session=lambda name: (["conv-session", "ssm-session"], None),
|
||||
_is_mamba_layer=lambda name: True,
|
||||
_mamba_payload_index_by_layer={"kda": 2},
|
||||
_region_session_indices=lambda name: [0, 1],
|
||||
_compute_mamba_transfer_offsets=Mock(
|
||||
return_value=([10, 20, 30, 40], [50, 60, 70, 80], [2, 4, 8, 16], 3)
|
||||
),
|
||||
)
|
||||
worker._mamba_blocks_for_layer = lambda name, blocks: (
|
||||
MoRIIOConnectorWorker._mamba_blocks_for_layer(worker, name, blocks)
|
||||
)
|
||||
writer = make_moriio_writer(worker)
|
||||
task = _write_task("kda")
|
||||
task.local_block_ids = [[], [1], [3]]
|
||||
task.remote_tp_size = 2
|
||||
task.event = Mock()
|
||||
writer._execute_write_task(task)
|
||||
|
||||
task.event.synchronize.assert_called_once()
|
||||
worker._compute_mamba_transfer_offsets.assert_called_once_with(
|
||||
"kda", [3], [9], remote_tp_size=2
|
||||
)
|
||||
assert wrapper.write_remote_data.call_args_list[0].args == (
|
||||
[2, 4, 8],
|
||||
[10, 20, 30],
|
||||
[50, 60, 70],
|
||||
"conv-session",
|
||||
)
|
||||
assert wrapper.write_remote_data.call_args_list[1].args == (
|
||||
[16],
|
||||
[40],
|
||||
[80],
|
||||
"ssm-session",
|
||||
)
|
||||
assert info.transfer_statuses == ["conv-status", "ssm-status"]
|
||||
assert info.writes_done == 1
|
||||
assert not info.completion_notified
|
||||
|
||||
|
||||
def test_write_completion_notifies_once_after_all_sealed_writes_finish():
|
||||
class FakeWrapper:
|
||||
def __init__(self):
|
||||
@@ -445,7 +524,7 @@ def test_write_completion_notifies_once_after_all_sealed_writes_finish():
|
||||
)
|
||||
|
||||
wrapper = FakeWrapper()
|
||||
request_info = RemoteAllocInfo(block_ids=[4, 5], writes_expected=2)
|
||||
request_info = RemoteAllocInfo(block_ids=[[4, 5]], writes_expected=2)
|
||||
request_info.transfer_statuses.extend(["status-a", "status-b"])
|
||||
request_info.completion_request_id = "req"
|
||||
# Final notify port is resolved in _execute_write_task; finalize reads it
|
||||
@@ -487,6 +566,7 @@ def test_moriio_wrapper_waits_scoped_statuses_without_global_drain():
|
||||
wrapper = MoRIIOWrapper.__new__(MoRIIOWrapper)
|
||||
wrapper.lock = threading.Lock()
|
||||
wrapper._transfer_timeout = 1
|
||||
wrapper.moriio_engine = object()
|
||||
global_status = FakeStatus()
|
||||
scoped_status = FakeStatus()
|
||||
wrapper.transfer_status = [global_status]
|
||||
@@ -500,7 +580,7 @@ def test_moriio_wrapper_waits_scoped_statuses_without_global_drain():
|
||||
|
||||
def test_write_failure_marks_terminal_and_clears_scheduled_state():
|
||||
wrapper = _wrapper_for_messages()
|
||||
wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[4, 5])
|
||||
wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[[4, 5]])
|
||||
writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper))
|
||||
writer._scheduled_writes["xfer"] = 2
|
||||
writer._scheduled_layers["xfer"] = {"dense0", "indexer"}
|
||||
@@ -518,7 +598,7 @@ def test_write_failure_marks_terminal_and_clears_scheduled_state():
|
||||
|
||||
def test_schedule_write_rejects_terminal_transfer_without_recreating_state():
|
||||
wrapper = _wrapper_for_messages()
|
||||
wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[4, 5])
|
||||
wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[[4, 5]])
|
||||
writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper))
|
||||
writer._scheduled_writes["xfer"] = 1
|
||||
writer._scheduled_layers["xfer"] = {"dense0"}
|
||||
@@ -610,7 +690,7 @@ def test_moriio_wrapper_routes_valid_messages(role, payload, expected):
|
||||
|
||||
if expected == "remote_blocks":
|
||||
request_info = wrapper.done_remote_allocate_req_dict["xfer"]
|
||||
assert request_info.block_ids == [4, 5]
|
||||
assert request_info.block_ids == [[4, 5]]
|
||||
assert request_info.decode_dp_rank == 3
|
||||
elif expected == "write_done":
|
||||
assert wrapper.done_write_cache_req_ids == ["xfer"]
|
||||
|
||||
@@ -11,6 +11,8 @@ import importlib
|
||||
import importlib.util
|
||||
import threading
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -127,3 +129,62 @@ def test_wait_blocks_until_terminal():
|
||||
assert time.monotonic() - start >= 0.15
|
||||
finally:
|
||||
timer.cancel()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("timeout,expected_ms", [(0.05, 50), (0.0001, 1)])
|
||||
def test_native_wait_receives_all_statuses_and_a_positive_deadline(
|
||||
timeout, expected_ms
|
||||
):
|
||||
statuses = [_pending(), _ok()]
|
||||
engine = SimpleNamespace(
|
||||
wait_all=Mock(return_value=moriio_engine.StatusCode.SUCCESS)
|
||||
)
|
||||
wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=timeout)
|
||||
wrapper.waiting_for_transfer_complete(statuses)
|
||||
engine.wait_all.assert_called_once_with(statuses, timeout_ms=expected_ms)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"result,statuses,message",
|
||||
[
|
||||
("ERR_RDMA_OP", [_pending(), _bad()], "SQ full"),
|
||||
("IN_PROGRESS", [_pending(), _ok()], "timed out"),
|
||||
("ERR_INVALID_ARGS", [_pending()], "ERR_INVALID_ARGS"),
|
||||
],
|
||||
)
|
||||
def test_native_wait_reports_failure_or_timeout(result, statuses, message):
|
||||
engine = SimpleNamespace(
|
||||
wait_all=Mock(return_value=getattr(moriio_engine.StatusCode, result))
|
||||
)
|
||||
wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05)
|
||||
with pytest.raises(TransferError, match=message):
|
||||
wrapper.waiting_for_transfer_complete(statuses)
|
||||
|
||||
|
||||
def test_wait_falls_back_without_native_status_enum(monkeypatch):
|
||||
monkeypatch.setattr(moriio_engine, "StatusCode", None)
|
||||
engine = SimpleNamespace(wait_all=Mock())
|
||||
wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05)
|
||||
with pytest.raises(TransferError, match="SQ full"):
|
||||
wrapper.waiting_for_transfer_complete([_bad(), _ok()])
|
||||
engine.wait_all.assert_not_called()
|
||||
|
||||
|
||||
def test_nonblocking_poll_does_not_call_native_wait():
|
||||
engine = SimpleNamespace(wait_all=Mock())
|
||||
wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05)
|
||||
assert wrapper.poll_transfer_batch([_pending()]) is TransferBatchState.PENDING
|
||||
wrapper.waiting_for_transfer_complete([])
|
||||
engine.wait_all.assert_not_called()
|
||||
|
||||
|
||||
def test_native_wait_drains_only_the_unscoped_status_snapshot():
|
||||
statuses = [_pending(), _ok()]
|
||||
engine = SimpleNamespace(
|
||||
wait_all=Mock(return_value=moriio_engine.StatusCode.SUCCESS)
|
||||
)
|
||||
wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05)
|
||||
wrapper.transfer_status.extend(statuses)
|
||||
wrapper.waiting_for_transfer_complete()
|
||||
engine.wait_all.assert_called_once_with(statuses, timeout_ms=50)
|
||||
assert wrapper.transfer_status == []
|
||||
|
||||
@@ -14,6 +14,8 @@ import pytest
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.moriio import moriio_engine
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
|
||||
ROLE,
|
||||
LayerTransferPlan,
|
||||
MoRIIOError,
|
||||
MoRIIOMode,
|
||||
RemoteAllocInfo,
|
||||
ReqMeta,
|
||||
@@ -96,7 +98,7 @@ def _write_task(
|
||||
request_id=request_id,
|
||||
transfer_id=transfer_id,
|
||||
dst_engine_id="remote-engine",
|
||||
local_block_ids=[1, 3],
|
||||
local_block_ids=[[1, 3]],
|
||||
remote_block_ids_hint=None,
|
||||
layer_name="layer0",
|
||||
event=None,
|
||||
@@ -249,6 +251,7 @@ def test_finalize_single_pod_port_uses_global_rank():
|
||||
|
||||
def test_write_blocks_for_req_forwards_routing_to_schedule_write():
|
||||
worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
|
||||
worker._has_mamba = False
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
def _capture_schedule_write_blocks(**kwargs: Any) -> None:
|
||||
@@ -266,13 +269,14 @@ def test_write_blocks_for_req_forwards_routing_to_schedule_write():
|
||||
assert captured["multi_pod_hosts"] == ["10.0.0.1", "10.0.0.2"]
|
||||
assert captured["remote_dp_size_local"] == 8
|
||||
assert captured["remote_ip"] == "10.0.0.1"
|
||||
assert captured["local_block_ids"] == [1]
|
||||
assert captured["remote_block_ids"] == [2]
|
||||
assert captured["local_block_ids"] == [[1]]
|
||||
assert captured["remote_block_ids"] == [[2]]
|
||||
assert not hasattr(worker, "multi_pod_hosts")
|
||||
|
||||
|
||||
def test_write_blocks_for_req_falls_back_when_multi_pod_hosts_empty():
|
||||
worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
|
||||
worker._has_mamba = False
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
def _capture_schedule_write_blocks(**kwargs: Any) -> None:
|
||||
@@ -290,8 +294,8 @@ def test_write_blocks_for_req_falls_back_when_multi_pod_hosts_empty():
|
||||
|
||||
assert captured["multi_pod_hosts"] == ["10.0.0.9"]
|
||||
assert captured["remote_dp_size_local"] == 4
|
||||
assert captured["local_block_ids"] == [1]
|
||||
assert captured["remote_block_ids"] == [2]
|
||||
assert captured["local_block_ids"] == [[1]]
|
||||
assert captured["remote_block_ids"] == [[2]]
|
||||
|
||||
|
||||
def test_deferred_write_timeout_fails_transfer_on_every_decode_rank():
|
||||
@@ -373,3 +377,73 @@ def test_write_failed_is_reported_as_failed_recving_on_owning_rank():
|
||||
assert results.failed_recving == {"req-1"}
|
||||
assert wrapper.pop_finished_write_req_ids() == set()
|
||||
assert list(worker._unmatched_write_failures) == ["tOtherRank"]
|
||||
|
||||
|
||||
def test_hybrid_write_rejects_unequal_tp_before_scheduling_any_layer():
|
||||
from unittest.mock import Mock
|
||||
|
||||
worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
|
||||
worker._has_mamba = True
|
||||
worker.world_size = 2
|
||||
worker.schedule_write_blocks = Mock()
|
||||
|
||||
with pytest.raises(NotImplementedError, match="heterogeneous-TP"):
|
||||
worker._write_blocks_for_req("req", _req_meta(), "attention", kv_layer=None)
|
||||
|
||||
worker.schedule_write_blocks.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"notify,expected",
|
||||
[
|
||||
# Plain WRITE peers still send one flat attention list.
|
||||
([7, 8], [[7, 8]]),
|
||||
# Hybrid WRITE sends one list per transfer group, kept as-is even
|
||||
# when the first group is empty (full attention hit, state only).
|
||||
([[7], [9]], [[7], [9]]),
|
||||
([[], [9]], [[], [9]]),
|
||||
],
|
||||
)
|
||||
def test_remote_blocks_message_normalizes_to_grouped_block_ids(notify, expected):
|
||||
wrapper = _real_wrapper()
|
||||
with patch.object(moriio_engine, "get_role", return_value=ROLE.PRODUCER):
|
||||
wrapper._handle_remote_blocks_message(
|
||||
{"transfer_id": "tA", "block_notify_list": notify, "decode_rank": 2}
|
||||
)
|
||||
info = wrapper.done_remote_allocate_req_dict["tA"]
|
||||
assert info.block_ids == expected
|
||||
assert info.decode_dp_rank == 2
|
||||
|
||||
|
||||
def test_empty_layer_plan_posts_no_rdma_write():
|
||||
from unittest.mock import Mock
|
||||
|
||||
wrapper = SimpleNamespace(write_remote_data=Mock(), write_remote_data_single=Mock())
|
||||
writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper))
|
||||
plan = LayerTransferPlan(
|
||||
request_id="req",
|
||||
transfer_id="tA",
|
||||
layer_name="layer0",
|
||||
sess_idx=0,
|
||||
transfer_local_offsets=[],
|
||||
transfer_remote_offsets=[],
|
||||
transfer_sizes=[],
|
||||
)
|
||||
|
||||
assert writer._do_layer_write(plan, ["session"]) == []
|
||||
wrapper.write_remote_data.assert_not_called()
|
||||
wrapper.write_remote_data_single.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("local_state", [[], [1, 2]])
|
||||
def test_hybrid_write_requires_exactly_one_running_state_slot(local_state):
|
||||
worker = SimpleNamespace(
|
||||
_mamba_blocks_for_layer=lambda layer_name, block_ids: block_ids[1],
|
||||
)
|
||||
writer = make_moriio_writer(worker)
|
||||
task = _write_task("tA", "10.0.0.1")
|
||||
task.layer_name = "kda"
|
||||
task.local_block_ids = [[], local_state]
|
||||
|
||||
with pytest.raises(MoRIIOError, match="one running-state slot"):
|
||||
writer._prepare_mamba_transfer_plans(task, RemoteAllocInfo(block_ids=[[], [9]]))
|
||||
|
||||
@@ -63,12 +63,13 @@ class WriteTask:
|
||||
request_id: ReqId
|
||||
transfer_id: TransferId
|
||||
dst_engine_id: str
|
||||
local_block_ids: list[int]
|
||||
remote_block_ids_hint: list[int] | None
|
||||
local_block_ids: BlockIds
|
||||
remote_block_ids_hint: BlockIds | None
|
||||
layer_name: str
|
||||
event: torch.cuda.Event
|
||||
remote_notify_port: int
|
||||
remote_ip: str
|
||||
remote_tp_size: int | None = None
|
||||
multi_pod_hosts: list[str] = field(default_factory=list)
|
||||
remote_dp_size_local: int = 0
|
||||
remote_dp_size: int = 1
|
||||
@@ -94,7 +95,7 @@ class LayerTransferPlan:
|
||||
class RemoteAllocInfo:
|
||||
"""Information about remote block allocation."""
|
||||
|
||||
block_ids: list[int]
|
||||
block_ids: BlockIds | None
|
||||
writes_done: int = 0
|
||||
writes_expected: int | None = None
|
||||
decode_dp_rank: int = 0
|
||||
|
||||
@@ -338,7 +338,10 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA):
|
||||
block_ids: BlockIds,
|
||||
) -> tuple[bool, dict[str, Any] | None]:
|
||||
assert self.connector_scheduler is not None
|
||||
if self.connector_scheduler._has_mamba:
|
||||
if (
|
||||
self.connector_scheduler._has_mamba
|
||||
and self.connector_scheduler.mode == MoRIIOMode.READ
|
||||
):
|
||||
attn_block_ids, mamba_block_groups = (
|
||||
self.connector_scheduler.split_block_groups(block_ids)
|
||||
)
|
||||
@@ -420,6 +423,16 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA):
|
||||
)
|
||||
self.connector_worker.wait_for_save(self._connector_metadata)
|
||||
|
||||
def has_pending_push_work(self) -> bool:
|
||||
return (
|
||||
self.mode == MoRIIOMode.WRITE
|
||||
and self.connector_scheduler is not None
|
||||
and bool(self.connector_scheduler._deferred_send_deadlines)
|
||||
)
|
||||
|
||||
def get_loaded_kv_cache_group_ids(self, request: "Request") -> tuple[int, ...]:
|
||||
return self._kv_cache_config.transfer_group_ids
|
||||
|
||||
def shutdown(self):
|
||||
if self.connector_worker is not None:
|
||||
self.connector_worker.shutdown()
|
||||
@@ -463,11 +476,10 @@ def _validate_hybrid_speculation(vllm_config: VllmConfig) -> None:
|
||||
speculative_config = vllm_config.speculative_config
|
||||
if speculative_config is None:
|
||||
return
|
||||
# DSpark is allowlisted because it is the only hybrid speculative method
|
||||
# validated end to end with MoRIIO READ.
|
||||
# Other hybrid speculative methods need their own transfer validation.
|
||||
if not speculative_config.use_dspark():
|
||||
raise MoRIIOError(
|
||||
"MoRIIO hybrid READ supports DSpark speculative decoding only, got "
|
||||
"MoRIIO hybrid transfer supports DSpark speculative decoding only, got "
|
||||
f"method={speculative_config.method!r}"
|
||||
)
|
||||
|
||||
@@ -480,14 +492,14 @@ def _validate_mamba_specs(specs: Collection[MambaSpec]) -> MambaSpec | None:
|
||||
for spec in specs:
|
||||
if len(spec.shapes) != 2 or len(spec.dtypes) != 2:
|
||||
raise MoRIIOError(
|
||||
"MoRIIO hybrid READ supports exactly two Mamba states "
|
||||
"MoRIIO hybrid transfer supports exactly two Mamba states "
|
||||
"(conv and SSM), got "
|
||||
f"{len(spec.shapes)} shapes and {len(spec.dtypes)} dtypes"
|
||||
)
|
||||
reference = specs[0]
|
||||
if any(spec != reference for spec in specs[1:]):
|
||||
raise MoRIIOError(
|
||||
"MoRIIO hybrid READ requires all transferable Mamba groups "
|
||||
"MoRIIO hybrid transfer requires all transferable Mamba groups "
|
||||
"to use the same cache spec"
|
||||
)
|
||||
return reference
|
||||
@@ -515,7 +527,9 @@ class MoRIIOConnectorScheduler:
|
||||
self._has_mamba = bool(self._mamba_group_ids)
|
||||
mamba_spec = None
|
||||
if self._has_mamba:
|
||||
if len(self._attn_group_ids) != 1:
|
||||
if not self._attn_group_ids:
|
||||
raise MoRIIOError("MoRIIO hybrid transfer requires an attention group")
|
||||
if self.mode == MoRIIOMode.READ and len(self._attn_group_ids) != 1:
|
||||
raise MoRIIOError(
|
||||
"MoRIIO hybrid READ requires exactly one transferable "
|
||||
"attention group, got "
|
||||
@@ -529,11 +543,6 @@ class MoRIIOConnectorScheduler:
|
||||
]
|
||||
mamba_spec = _validate_mamba_specs(mamba_specs)
|
||||
_validate_hybrid_speculation(vllm_config)
|
||||
if self.mode != MoRIIOMode.READ:
|
||||
raise MoRIIOError(
|
||||
"MoRIIO hybrid (mamba/KDA) transfer is implemented for READ "
|
||||
"mode only; set kv_connector_extra_config.read_mode=true"
|
||||
)
|
||||
self._num_ssm_scratch_blocks = (
|
||||
mamba_spec.num_speculative_blocks if mamba_spec is not None else 0
|
||||
)
|
||||
@@ -581,11 +590,13 @@ class MoRIIOConnectorScheduler:
|
||||
f"cache group spec(s): {sorted(set(unsupported))}. Pass "
|
||||
"--disable-hybrid-kv-cache-manager to run without HMA."
|
||||
)
|
||||
if self.mode == MoRIIOMode.WRITE:
|
||||
# TODO(simondanielsson): support HMA in WRITE mode
|
||||
if self.mode == MoRIIOMode.WRITE and any(
|
||||
isinstance(g.kv_cache_spec, SlidingWindowSpec)
|
||||
for g in kv_cache_config.transfer_groups
|
||||
):
|
||||
raise NotImplementedError(
|
||||
"MoRIIO WRITE mode does not support hybrid KV cache groups "
|
||||
"(sliding-window attention). Use READ mode, or pass "
|
||||
"MoRIIO WRITE mode does not support sliding-window "
|
||||
"cache groups. Use READ mode, or pass "
|
||||
"--disable-hybrid-kv-cache-manager."
|
||||
)
|
||||
|
||||
@@ -784,10 +795,11 @@ class MoRIIOConnectorScheduler:
|
||||
if not params.get("do_remote_prefill"):
|
||||
return 0, False
|
||||
if self.mode == MoRIIOMode.WRITE:
|
||||
# MoriiO in write mode, no remote prefill. Hybrid models never get
|
||||
# here: register_kv_caches refuses WRITE mode for them, so there is
|
||||
# no recurrent-state accounting to do.
|
||||
return num_external_tokens, True
|
||||
if self._has_mamba:
|
||||
# P pushes h(N-1). Do not publish it under the N-token hash
|
||||
# when the scheduler caches a completed asynchronous load.
|
||||
num_external_tokens = max(num_external_tokens - 1, 0)
|
||||
return num_external_tokens, num_external_tokens > 0
|
||||
|
||||
# READ mode always recomputes the last token locally on the decoder.
|
||||
#
|
||||
@@ -1178,6 +1190,16 @@ class MoRIIOConnectorScheduler:
|
||||
f"remote_notify_port={remote_notify_port!r})"
|
||||
)
|
||||
|
||||
block_notify_list: list[int] | list[list[int]]
|
||||
if self._has_mamba:
|
||||
block_notify_list = list(
|
||||
self._clamp_to_prompt_blocks(
|
||||
request, blocks.get_block_ids()
|
||||
)
|
||||
)
|
||||
else:
|
||||
block_notify_list = blocks.get_block_ids()[0]
|
||||
|
||||
# Wide-EP multi-pod: a pod binds notify sockets only for
|
||||
# its LOCAL ranks, so the port offset must use the per-pod
|
||||
# local rank (% dp_local), not the global rank. Single-pod
|
||||
@@ -1204,7 +1226,7 @@ class MoRIIOConnectorScheduler:
|
||||
self.send_notify_block(
|
||||
req_id=request.request_id,
|
||||
transfer_id=request.kv_transfer_params["transfer_id"],
|
||||
block_notify_list=blocks.get_block_ids()[0],
|
||||
block_notify_list=block_notify_list,
|
||||
host=_notify_host,
|
||||
port=target_port,
|
||||
)
|
||||
@@ -1227,6 +1249,26 @@ class MoRIIOConnectorScheduler:
|
||||
sliding-window groups are left unchanged.
|
||||
|
||||
"""
|
||||
if self._has_mamba:
|
||||
# WRITE preserves transfer-group order, including a separate drafter
|
||||
# attention group. READ retains its attention-first wire format.
|
||||
transfer_blocks = self.kv_cache_config.select_transfer_block_ids(block_ids)
|
||||
dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size
|
||||
clamped_hybrid: list[list[int]] = []
|
||||
for group, blocks in zip(
|
||||
self.kv_cache_config.transfer_groups, transfer_blocks, strict=True
|
||||
):
|
||||
spec = group.kv_cache_spec
|
||||
if isinstance(spec, MambaSpec):
|
||||
clamped_hybrid.append(
|
||||
clip_ssm_state_blocks(list(blocks), spec.num_speculative_blocks)
|
||||
)
|
||||
else:
|
||||
block_size = spec.block_size * (dcp_size if spec.dcp_sharded else 1)
|
||||
clamped_hybrid.append(
|
||||
list(blocks[: cdiv(req.num_prompt_tokens, block_size)])
|
||||
)
|
||||
return clamped_hybrid
|
||||
clamped = list(block_ids)
|
||||
changed = False
|
||||
for g, group_blocks in enumerate(block_ids):
|
||||
@@ -1252,6 +1294,19 @@ class MoRIIOConnectorScheduler:
|
||||
num_scheduled = scheduler_output.num_scheduled_tokens.get(req.request_id, 0)
|
||||
return req.num_computed_tokens + num_scheduled >= req.num_prompt_tokens
|
||||
|
||||
def _get_write_blocks(
|
||||
self, req: "Request", block_ids: BlockIds, scheduler_output: SchedulerOutput
|
||||
) -> BlockIds:
|
||||
if self._has_mamba:
|
||||
# Align-mode state slots can move without appending a new block.
|
||||
# Resolve the authoritative table at the final prefill step.
|
||||
block_state = scheduler_output.kv_connector_block_state
|
||||
assert block_state is not None
|
||||
current_blocks = block_state.get_block_ids(req.request_id)
|
||||
assert current_blocks is not None
|
||||
block_ids = current_blocks
|
||||
return self._clamp_to_prompt_blocks(req, block_ids)
|
||||
|
||||
def build_connector_meta(
|
||||
self,
|
||||
scheduler_output: SchedulerOutput,
|
||||
@@ -1274,10 +1329,10 @@ class MoRIIOConnectorScheduler:
|
||||
# reserved speculative-lookahead slack), so new_block_ids is None;
|
||||
# still fall through to the final-chunk check below.
|
||||
new_block_ids = scheduler_output.scheduled_cached_reqs.new_block_ids[i]
|
||||
if new_block_ids is not None:
|
||||
if new_block_ids is not None and not self._has_mamba:
|
||||
existing_blocks = [
|
||||
existing_blocks[g] + new_block_ids[g]
|
||||
for g in range(len(new_block_ids))
|
||||
old + new
|
||||
for old, new in zip(existing_blocks, new_block_ids, strict=True)
|
||||
]
|
||||
self._reqs_need_pending_save[req_id] = (req, existing_blocks)
|
||||
if not self._is_final_prefill_chunk(req, scheduler_output):
|
||||
@@ -1287,7 +1342,9 @@ class MoRIIOConnectorScheduler:
|
||||
kv_params = self._req_kv_params.pop(
|
||||
req_id, req.kv_transfer_params or {}
|
||||
)
|
||||
save_block_ids = self._clamp_to_prompt_blocks(req, existing_blocks)
|
||||
save_block_ids = self._get_write_blocks(
|
||||
req, existing_blocks, scheduler_output
|
||||
)
|
||||
meta.add_new_req(
|
||||
request_id=req_id,
|
||||
local_block_ids=save_block_ids,
|
||||
@@ -1328,7 +1385,7 @@ class MoRIIOConnectorScheduler:
|
||||
if not self._is_final_prefill_chunk(req, scheduler_output):
|
||||
self._reqs_need_pending_save[req_id] = (req, block_ids)
|
||||
continue
|
||||
save_block_ids = self._clamp_to_prompt_blocks(req, block_ids)
|
||||
save_block_ids = self._get_write_blocks(req, block_ids, scheduler_output)
|
||||
meta.add_new_req(
|
||||
request_id=req_id,
|
||||
local_block_ids=save_block_ids,
|
||||
@@ -1425,7 +1482,7 @@ class MoRIIOConnectorScheduler:
|
||||
"""Once a request is finished, determine whether request blocks
|
||||
should be freed now or will be sent asynchronously and freed later.
|
||||
|
||||
``block_ids`` normally contains all cache groups. Hybrid callers from
|
||||
``block_ids`` normally contains all cache groups. Hybrid READ callers from
|
||||
``request_finished_all_groups`` pass a flat attention list and the
|
||||
recurrent-state groups separately in ``mamba_block_groups``.
|
||||
"""
|
||||
@@ -1494,6 +1551,10 @@ class MoRIIOConnectorScheduler:
|
||||
list(cast(list[int], block_ids)),
|
||||
*[list(group) for group in mamba_block_groups],
|
||||
]
|
||||
elif self._has_mamba and self.mode == MoRIIOMode.WRITE:
|
||||
computed_block_ids = self._clamp_to_prompt_blocks(
|
||||
request, cast(BlockIds, block_ids)
|
||||
)
|
||||
elif self._has_mamba:
|
||||
raise MoRIIOError("Hybrid request completion requires all KV cache groups")
|
||||
else:
|
||||
@@ -1660,7 +1721,7 @@ class MoRIIOConnectorWorker:
|
||||
_, mamba_transfer_group_ids = _split_kv_cache_group_kinds(kv_cache_config)
|
||||
self._num_mamba_transfer_groups = len(mamba_transfer_group_ids)
|
||||
self._mamba_payload_index_by_layer = {
|
||||
layer_name: payload_index
|
||||
layer_name: group_id if self.mode == MoRIIOMode.WRITE else payload_index
|
||||
for payload_index, group_id in enumerate(mamba_transfer_group_ids, start=1)
|
||||
for layer_name in kv_cache_config.transfer_groups[group_id].layer_names
|
||||
}
|
||||
@@ -1884,8 +1945,8 @@ class MoRIIOConnectorWorker:
|
||||
request_id: ReqId,
|
||||
transfer_id: TransferId,
|
||||
dst_engine_id: str,
|
||||
local_block_ids: list[int],
|
||||
remote_block_ids: list[int] | None,
|
||||
local_block_ids: BlockIds,
|
||||
remote_block_ids: BlockIds | None,
|
||||
layer_name: str,
|
||||
kv_layer: torch.Tensor,
|
||||
remote_notify_port: int,
|
||||
@@ -1893,6 +1954,7 @@ class MoRIIOConnectorWorker:
|
||||
multi_pod_hosts: list[str],
|
||||
remote_dp_size_local: int,
|
||||
remote_dp_size: int,
|
||||
remote_tp_size: int | None = None,
|
||||
) -> None:
|
||||
"""Schedule a block write operation.
|
||||
|
||||
@@ -1909,6 +1971,7 @@ class MoRIIOConnectorWorker:
|
||||
multi_pod_hosts: List of pod IPs for multi-pod Wide-EP
|
||||
remote_dp_size_local: Per-pod DP size for multi-pod
|
||||
remote_dp_size: Global DP size of the decode instance
|
||||
remote_tp_size: Tensor parallel size advertised by the peer
|
||||
|
||||
"""
|
||||
# synchronization to prevent dirty reads between
|
||||
@@ -1933,6 +1996,7 @@ class MoRIIOConnectorWorker:
|
||||
multi_pod_hosts=multi_pod_hosts,
|
||||
remote_dp_size_local=remote_dp_size_local,
|
||||
remote_dp_size=remote_dp_size,
|
||||
remote_tp_size=remote_tp_size,
|
||||
)
|
||||
self._writer.schedule_write(task)
|
||||
|
||||
@@ -2359,15 +2423,6 @@ class MoRIIOConnectorWorker:
|
||||
self._has_mamba = any(
|
||||
self._is_mamba_layer(layer_name) for layer_name in kv_caches
|
||||
)
|
||||
if self._has_mamba and self.mode == MoRIIOMode.WRITE:
|
||||
# WRITE needs a distinct conv/ssm transfer plan. Reject before
|
||||
# allocating or registering any regions rather than silently
|
||||
# serving a zero recurrent state.
|
||||
raise MoRIIOError(
|
||||
"MoRIIO hybrid (mamba/KDA) transfer is implemented for READ "
|
||||
"mode only; set kv_connector_extra_config.read_mode=true "
|
||||
"for hybrid models"
|
||||
)
|
||||
if self._has_mamba:
|
||||
mamba_specs = [
|
||||
cast(MambaSpec, self.layer_to_spec[layer_name])
|
||||
@@ -3292,6 +3347,9 @@ class MoRIIOConnectorWorker:
|
||||
)
|
||||
|
||||
def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer):
|
||||
if self._has_mamba:
|
||||
# Reject before even an attention-layer write can be queued.
|
||||
self._mamba_tp_ratio(int(meta.tp_size))
|
||||
# Compute per-request values to pass through the task (no shared state).
|
||||
# This avoids race conditions when concurrent requests target different
|
||||
# decode pods - each WriteTask carries its own routing info.
|
||||
@@ -3303,16 +3361,12 @@ class MoRIIOConnectorWorker:
|
||||
if meta.remote_dp_size_local
|
||||
else int(meta.remote_dp_size)
|
||||
)
|
||||
# WRITE does not support HMA, so unwrap blocks into flat lists.
|
||||
# remote_block_ids can itself be empty so need to be careful when unwrapping.
|
||||
local_block_ids = meta.local_block_ids[0]
|
||||
remote_block_ids = meta.remote_block_ids[0] if meta.remote_block_ids else []
|
||||
self.schedule_write_blocks(
|
||||
request_id=req_id,
|
||||
transfer_id=meta.transfer_id,
|
||||
dst_engine_id=meta.remote_engine_id,
|
||||
local_block_ids=local_block_ids,
|
||||
remote_block_ids=remote_block_ids,
|
||||
local_block_ids=meta.local_block_ids,
|
||||
remote_block_ids=meta.remote_block_ids,
|
||||
layer_name=layer_name,
|
||||
kv_layer=kv_layer,
|
||||
remote_notify_port=meta.remote_notify_port,
|
||||
@@ -3320,6 +3374,7 @@ class MoRIIOConnectorWorker:
|
||||
multi_pod_hosts=hosts,
|
||||
remote_dp_size_local=dp_local,
|
||||
remote_dp_size=int(meta.remote_dp_size),
|
||||
remote_tp_size=int(meta.tp_size),
|
||||
)
|
||||
|
||||
def merge_contiguous_blocks(
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import threading
|
||||
import time
|
||||
from collections import OrderedDict, defaultdict
|
||||
from math import ceil
|
||||
from queue import Empty, Queue
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from weakref import ref as weakref_ref
|
||||
@@ -61,6 +62,11 @@ try:
|
||||
except ImportError:
|
||||
logger.error("MoRIIO is not available")
|
||||
|
||||
try:
|
||||
from mori.io import StatusCode
|
||||
except ImportError:
|
||||
StatusCode = None
|
||||
|
||||
"""Write task execution logic for MoRIIO connector."""
|
||||
|
||||
|
||||
@@ -372,11 +378,15 @@ class MoRIIOWriter:
|
||||
task.dst_engine_id
|
||||
)
|
||||
|
||||
# Prepare transfer plan
|
||||
plan = self._prepare_transfer_plan(task, request_info, remote_moriio_meta)
|
||||
|
||||
# Execute transfer
|
||||
transfer_statuses = self._do_layer_write(plan, sessions)
|
||||
if self.worker._is_mamba_layer(task.layer_name):
|
||||
plans = self._prepare_mamba_transfer_plans(task, request_info)
|
||||
else:
|
||||
plans = [
|
||||
self._prepare_transfer_plan(task, request_info, remote_moriio_meta)
|
||||
]
|
||||
transfer_statuses = []
|
||||
for plan in plans:
|
||||
transfer_statuses.extend(self._do_layer_write(plan, sessions))
|
||||
with self._write_state_lock:
|
||||
request_info.transfer_statuses.extend(transfer_statuses)
|
||||
|
||||
@@ -404,10 +414,17 @@ class MoRIIOWriter:
|
||||
key = (task.layer_name, *_get_write_geometry_key(layer_cache))
|
||||
offsets = request_info.transfer_offsets.get(key)
|
||||
if offsets is None:
|
||||
group = self.worker.layer_to_group[task.layer_name]
|
||||
assert request_info.block_ids is not None
|
||||
local_blocks = task.local_block_ids[group]
|
||||
remote_blocks = request_info.block_ids[group]
|
||||
if self.worker._has_mamba:
|
||||
# D can reserve a tail block for the token recomputed locally.
|
||||
remote_blocks = remote_blocks[: len(local_blocks)]
|
||||
offsets = self.worker._compute_block_transfer_offsets(
|
||||
task.layer_name,
|
||||
task.local_block_ids,
|
||||
request_info.block_ids,
|
||||
local_blocks,
|
||||
remote_blocks,
|
||||
remote_moriio_meta,
|
||||
remote_engine_id=task.dst_engine_id,
|
||||
)
|
||||
@@ -430,6 +447,44 @@ class MoRIIOWriter:
|
||||
use_batch=True,
|
||||
)
|
||||
|
||||
def _prepare_mamba_transfer_plans(
|
||||
self, task: WriteTask, request_info: RemoteAllocInfo
|
||||
) -> list[LayerTransferPlan]:
|
||||
"""Plan both registered regions as one scheduled layer write."""
|
||||
assert request_info.block_ids is not None
|
||||
local_slots = self.worker._mamba_blocks_for_layer(
|
||||
task.layer_name, task.local_block_ids
|
||||
)
|
||||
remote_slots = self.worker._mamba_blocks_for_layer(
|
||||
task.layer_name, request_info.block_ids
|
||||
)
|
||||
if len(local_slots) != 1 or len(remote_slots) != 1:
|
||||
raise MoRIIOError(
|
||||
"Hybrid WRITE requires one running-state slot per Mamba group"
|
||||
)
|
||||
local, remote, sizes, n_conv = self.worker._compute_mamba_transfer_offsets(
|
||||
task.layer_name,
|
||||
local_slots,
|
||||
remote_slots,
|
||||
remote_tp_size=task.remote_tp_size,
|
||||
)
|
||||
conv_session, ssm_session = self.worker._region_session_indices(task.layer_name)
|
||||
return [
|
||||
LayerTransferPlan(
|
||||
request_id=task.request_id,
|
||||
transfer_id=task.transfer_id,
|
||||
layer_name=task.layer_name,
|
||||
sess_idx=session,
|
||||
transfer_local_offsets=local[region],
|
||||
transfer_remote_offsets=remote[region],
|
||||
transfer_sizes=sizes[region],
|
||||
)
|
||||
for session, region in (
|
||||
(conv_session, slice(0, n_conv)),
|
||||
(ssm_session, slice(n_conv, None)),
|
||||
)
|
||||
]
|
||||
|
||||
def _do_layer_write(self, plan: LayerTransferPlan, sessions: list) -> list[Any]:
|
||||
"""Perform the actual layer write.
|
||||
|
||||
@@ -438,6 +493,8 @@ class MoRIIOWriter:
|
||||
sessions: List of transfer sessions
|
||||
|
||||
"""
|
||||
if not plan.transfer_sizes:
|
||||
return []
|
||||
if plan.use_batch:
|
||||
return [
|
||||
self.worker.moriio_wrapper.write_remote_data(
|
||||
@@ -721,7 +778,28 @@ class MoRIIOWrapper:
|
||||
if not transfers_to_wait:
|
||||
return
|
||||
|
||||
self._poll_transfers_until_done(transfers_to_wait)
|
||||
wait_all = getattr(self.moriio_engine, "wait_all", None)
|
||||
if wait_all is None or StatusCode is None:
|
||||
self._poll_transfers_until_done(transfers_to_wait)
|
||||
return
|
||||
|
||||
# Zero means a nonblocking progress call in Mori, not an immediate
|
||||
# deadline. Round up so even sub-millisecond waits remain blocking.
|
||||
timeout_ms = ceil(self._transfer_timeout * 1000)
|
||||
result = wait_all(transfers_to_wait, timeout_ms=timeout_ms)
|
||||
if result != StatusCode.SUCCESS:
|
||||
details = [
|
||||
f"{status.Message()} (code={status.Code()})"
|
||||
for status in transfers_to_wait
|
||||
if status.Failed()
|
||||
]
|
||||
if not details and result in (StatusCode.INIT, StatusCode.IN_PROGRESS):
|
||||
details.append(
|
||||
f"RDMA transfer timed out after {self._transfer_timeout:g}s"
|
||||
)
|
||||
raise TransferError(
|
||||
f"MoRIIO batch wait failed ({result}): " + "; ".join(details)
|
||||
)
|
||||
|
||||
def _poll_transfers_until_done(self, transfers_to_wait: list[Any]) -> None:
|
||||
"""Fallback for mori builds without the batched wait.
|
||||
@@ -858,7 +936,12 @@ class MoRIIOWrapper:
|
||||
)
|
||||
return
|
||||
self.done_remote_allocate_req_dict[transfer_id] = RemoteAllocInfo(
|
||||
block_ids=block_notify_list, decode_dp_rank=decode_dp_rank
|
||||
block_ids=(
|
||||
[block_notify_list]
|
||||
if isinstance(block_notify_list[0], int)
|
||||
else block_notify_list
|
||||
),
|
||||
decode_dp_rank=decode_dp_rank,
|
||||
)
|
||||
|
||||
def _handle_write_done_message(self, data: dict):
|
||||
|
||||
Reference in New Issue
Block a user