dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

under test, not sure no errors

Browse Source
This commit is contained in:
DP Migration
2026-09-01 10:24:14 +00:00
parent 8c9d913f3f
commit 94d77cf0b4
15 changed files with 1859 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

343
core/runtime/py_attention_metadata.cpp Normal file
View File

@@ -0,0 +1,343 @@
/* Adapted from xLLM commit 78aa2a85 (PR #2258).
Adds dp_token_counts / dp_is_decode to the pybind11-exported
AttentionMetadataView so Python model executors (Qwen3.5 MoE layers,
decode graph runners) can read per-DP-rank token counts and decide
between padded vs compact all-gather.
Original: xllm/core/runtime/py_attention_metadata.cpp
Scope: Qwen3.5 data-parallel support in project_6.
==============================================================================*/
#include "core/runtime/py_attention_metadata.h"
#include <pybind11/stl.h>
#include <torch/extension.h>
#include <utility>
/*
* NOTE: The upstream xLLM implementation #includes
* "core/framework/model/model_input_params.h"
* "core/layers/common/attention_metadata.h"
* Those headers are part of xLLM's internal C++ framework and are NOT
* open-sourced in project_6. The stub types below satisfy the build so
* the DP-specific logic compiles; the real integration will link against
* the xLLM shared libraries that provide the concrete structs.
*/
namespace project6::layer {
struct ExpandedDecodeMetadata {
bool enabled = false;
torch::Tensor kv_seq_lens;
torch::Tensor block_table;
torch::Tensor paged_kv_indptr;
torch::Tensor paged_kv_indices;
torch::Tensor paged_kv_last_page_len;
torch::Tensor paged_attention_tiling_data;
torch::Tensor kv_seq_lens_host;
std::vector<int32_t> kv_seq_lens_host_vec;
};
struct AttentionMetadata {
torch::Tensor slot_mapping;
torch::Tensor paged_kv_indptr;
torch::Tensor paged_kv_indices;
torch::Tensor paged_kv_last_page_len;
std::optional<torch::Tensor> qo_indptr;
torch::Tensor q_cu_seq_lens;
torch::Tensor kv_cu_seq_lens;
torch::Tensor block_table;
torch::Tensor kv_seq_lens;
torch::Tensor q_seq_lens;
torch::Tensor has_initial_states;
std::vector<int32_t> kv_seq_lens_vec;
std::vector<int32_t> q_seq_lens_vec;
bool is_prefill = false;
bool is_chunked_prefill = false;
ExpandedDecodeMetadata expanded_decode;
};
} // namespace project6::layer
namespace project6 {
/* Minimal stub so the two-arg constructor compiles. */
struct ModelInputParams {
struct {
std::vector<int32_t> raw_dp_global_token_nums;
std::vector<int32_t> dp_global_token_nums;
std::vector<int32_t> dp_is_decode;
} parallel;
struct {
torch::Tensor linear_state_indices;
} embedding;
};
namespace py = pybind11;
// ---------------------------------------------------------------------------
// pybind11 registration
// ---------------------------------------------------------------------------
void register_attention_metadata_views(py::module_& module) {
py::class_<PyExpandedDecodeMetadataView>(module, "ExpandedDecodeMetadataView")
.def_property_readonly("enabled", &PyExpandedDecodeMetadataView::enabled)
.def_property_readonly("kv_seq_lens",
&PyExpandedDecodeMetadataView::kv_seq_lens)
.def_property_readonly("block_table",
&PyExpandedDecodeMetadataView::block_table)
.def_property_readonly("paged_kv_indptr",
&PyExpandedDecodeMetadataView::paged_kv_indptr)
.def_property_readonly("paged_kv_indices",
&PyExpandedDecodeMetadataView::paged_kv_indices)
.def_property_readonly(
"paged_kv_last_page_len",
&PyExpandedDecodeMetadataView::paged_kv_last_page_len)
.def_property_readonly(
"paged_attention_tiling_data",
&PyExpandedDecodeMetadataView::paged_attention_tiling_data)
.def_property_readonly("kv_seq_lens_host",
&PyExpandedDecodeMetadataView::kv_seq_lens_host)
.def_property_readonly(
"kv_seq_lens_host_values",
&PyExpandedDecodeMetadataView::kv_seq_lens_host_values);
py::class_<PyAttentionMetadataView>(module, "AttentionMetadataView")
.def_property_readonly("slot_mapping",
&PyAttentionMetadataView::slot_mapping)
.def_property_readonly("paged_kv_indptr",
&PyAttentionMetadataView::paged_kv_indptr)
.def_property_readonly("paged_kv_indices",
&PyAttentionMetadataView::paged_kv_indices)
.def_property_readonly("paged_kv_last_page_len",
&PyAttentionMetadataView::paged_kv_last_page_len)
.def_property_readonly("qo_indptr", &PyAttentionMetadataView::qo_indptr)
.def_property_readonly("q_cu_seq_lens",
&PyAttentionMetadataView::q_cu_seq_lens)
.def_property_readonly("kv_cu_seq_lens",
&PyAttentionMetadataView::kv_cu_seq_lens)
.def_property_readonly("kv_seq_lens_host",
&PyAttentionMetadataView::kv_seq_lens_host)
.def_property_readonly("kv_seq_lens_host_values",
&PyAttentionMetadataView::kv_seq_lens_host_values)
.def_property_readonly("q_seq_lens_host",
&PyAttentionMetadataView::q_seq_lens_host)
.def_property_readonly("block_table",
&PyAttentionMetadataView::block_table)
.def_property_readonly("kv_seq_lens",
&PyAttentionMetadataView::kv_seq_lens)
.def_property_readonly("linear_state_indices",
&PyAttentionMetadataView::linear_state_indices)
.def_property_readonly("has_initial_state",
&PyAttentionMetadataView::has_initial_state)
/* ---- DP fields (added by PR #2258) ------------------------------ */
.def_property_readonly("dp_token_counts",
&PyAttentionMetadataView::dp_token_counts)
.def_property_readonly("dp_is_decode",
&PyAttentionMetadataView::dp_is_decode)
/* ----------------------------------------------------------------- */
.def_property_readonly("q_seq_lens", &PyAttentionMetadataView::q_seq_lens)
.def_property_readonly("expanded_decode_metadata",
&PyAttentionMetadataView::expanded_decode_metadata)
.def_property_readonly("is_prefill", &PyAttentionMetadataView::is_prefill)
.def_property_readonly("is_chunked_prefill",
&PyAttentionMetadataView::is_chunked_prefill);
}
// ---------------------------------------------------------------------------
// PyExpandedDecodeMetadataView
// ---------------------------------------------------------------------------
PyExpandedDecodeMetadataView::PyExpandedDecodeMetadataView(
std::shared_ptr<layer::AttentionMetadata> metadata)
: metadata_(std::move(metadata)) {}
bool PyExpandedDecodeMetadataView::enabled() const {
return metadata().enabled;
}
py::object PyExpandedDecodeMetadataView::kv_seq_lens() const {
return metadata().kv_seq_lens.defined() ? py::cast(metadata().kv_seq_lens)
: py::none();
}
py::object PyExpandedDecodeMetadataView::block_table() const {
return metadata().block_table.defined() ? py::cast(metadata().block_table)
: py::none();
}
py::object PyExpandedDecodeMetadataView::paged_kv_indptr() const {
return metadata().paged_kv_indptr.defined()
? py::cast(metadata().paged_kv_indptr)
: py::none();
}
py::object PyExpandedDecodeMetadataView::paged_kv_indices() const {
return metadata().paged_kv_indices.defined()
? py::cast(metadata().paged_kv_indices)
: py::none();
}
py::object PyExpandedDecodeMetadataView::paged_kv_last_page_len() const {
return metadata().paged_kv_last_page_len.defined()
? py::cast(metadata().paged_kv_last_page_len)
: py::none();
}
py::object PyExpandedDecodeMetadataView::paged_attention_tiling_data() const {
return metadata().paged_attention_tiling_data.defined()
? py::cast(metadata().paged_attention_tiling_data)
: py::none();
}
py::object PyExpandedDecodeMetadataView::kv_seq_lens_host() const {
return metadata().kv_seq_lens_host.defined()
? py::cast(metadata().kv_seq_lens_host)
: py::none();
}
const std::vector<int32_t>&
PyExpandedDecodeMetadataView::kv_seq_lens_host_values() const {
return metadata().kv_seq_lens_host_vec;
}
const layer::ExpandedDecodeMetadata& PyExpandedDecodeMetadataView::metadata()
const {
return metadata_->expanded_decode;
}
// ---------------------------------------------------------------------------
// PyAttentionMetadataView
// ---------------------------------------------------------------------------
PyAttentionMetadataView::PyAttentionMetadataView(
std::shared_ptr<layer::AttentionMetadata> metadata)
: metadata_(std::move(metadata)),
kv_seq_lens_host_(
make_host_int32_view(metadata_, metadata_->kv_seq_lens_vec)),
q_seq_lens_host_(
make_host_int32_view(metadata_, metadata_->q_seq_lens_vec)) {}
PyAttentionMetadataView::PyAttentionMetadataView(
std::shared_ptr<layer::AttentionMetadata> metadata,
const ModelInputParams& params)
: PyAttentionMetadataView(std::move(metadata)) {
linear_state_indices_ = params.embedding.linear_state_indices;
/* ---- DP fields (added by PR #2258) ---------------------------------- */
dp_token_counts_ = params.parallel.raw_dp_global_token_nums.empty()
? params.parallel.dp_global_token_nums
: params.parallel.raw_dp_global_token_nums;
dp_is_decode_ = params.parallel.dp_is_decode;
/* --------------------------------------------------------------------- */
}
const torch::Tensor& PyAttentionMetadataView::slot_mapping() const {
return metadata_->slot_mapping;
}
const torch::Tensor& PyAttentionMetadataView::paged_kv_indptr() const {
return metadata_->paged_kv_indptr;
}
const torch::Tensor& PyAttentionMetadataView::paged_kv_indices() const {
return metadata_->paged_kv_indices;
}
const torch::Tensor& PyAttentionMetadataView::paged_kv_last_page_len() const {
return metadata_->paged_kv_last_page_len;
}
py::object PyAttentionMetadataView::qo_indptr() const {
if (!metadata_->qo_indptr.has_value() || !metadata_->qo_indptr->defined()) {
return py::none();
}
return py::cast(*metadata_->qo_indptr);
}
py::object PyAttentionMetadataView::q_cu_seq_lens() const {
return optional_tensor(metadata_->q_cu_seq_lens);
}
py::object PyAttentionMetadataView::kv_cu_seq_lens() const {
return optional_tensor(metadata_->kv_cu_seq_lens);
}
py::object PyAttentionMetadataView::kv_seq_lens_host() const {
return optional_tensor(kv_seq_lens_host_);
}
const std::vector<int32_t>& PyAttentionMetadataView::kv_seq_lens_host_values()
const {
return metadata_->kv_seq_lens_vec;
}
py::object PyAttentionMetadataView::block_table() const {
return optional_tensor(metadata_->block_table);
}
py::object PyAttentionMetadataView::kv_seq_lens() const {
return optional_tensor(metadata_->kv_seq_lens);
}
py::object PyAttentionMetadataView::linear_state_indices() const {
return optional_tensor(linear_state_indices_);
}
py::object PyAttentionMetadataView::has_initial_state() const {
return optional_tensor(metadata_->has_initial_states);
}
/* ---- DP fields (added by PR #2258) ------------------------------------ */
const std::vector<int32_t>& PyAttentionMetadataView::dp_token_counts() const {
return dp_token_counts_;
}
const std::vector<int32_t>& PyAttentionMetadataView::dp_is_decode() const {
return dp_is_decode_;
}
/* ----------------------------------------------------------------------- */
py::object PyAttentionMetadataView::q_seq_lens() const {
return optional_tensor(metadata_->q_seq_lens);
}
py::object PyAttentionMetadataView::q_seq_lens_host() const {
return optional_tensor(q_seq_lens_host_);
}
PyExpandedDecodeMetadataView PyAttentionMetadataView::expanded_decode_metadata()
const {
return PyExpandedDecodeMetadataView(metadata_);
}
bool PyAttentionMetadataView::is_prefill() const {
return metadata_->is_prefill;
}
bool PyAttentionMetadataView::is_chunked_prefill() const {
return metadata_->is_chunked_prefill;
}
torch::Tensor PyAttentionMetadataView::make_host_int32_view(
const std::shared_ptr<layer::AttentionMetadata>& metadata,
std::vector<int32_t>& host_vec) {
if (host_vec.empty()) {
return torch::Tensor();
}
std::shared_ptr<layer::AttentionMetadata> owner = metadata;
return torch::from_blob(
host_vec.data(),
{static_cast<int64_t>(host_vec.size())},
[owner = std::move(owner)](void*) mutable { owner.reset(); },
torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU));
}
py::object PyAttentionMetadataView::optional_tensor(
const torch::Tensor& tensor) {
return tensor.defined() ? py::cast(tensor) : py::none();
}
} // namespace project6

100
core/runtime/py_attention_metadata.h Normal file
View File

@@ -0,0 +1,100 @@
/* Adapted from xLLM commit 78aa2a85 (PR #2258).
Adds dp_token_counts / dp_is_decode fields to PyAttentionMetadataView
so the Python attention backend can partition KV cache by DP group.
Original: xllm/core/runtime/py_attention_metadata.h
Scope: Qwen3.5 data-parallel support in project_6.
==============================================================================*/
#pragma once
#include <pybind11/pybind11.h>
#include <torch/torch.h>
#include <cstdint>
#include <memory>
#include <vector>
/* Forward declarations — project_6 keeps these in its own layer namespace. */
namespace project6::layer {
struct AttentionMetadata;
struct ExpandedDecodeMetadata;
} // namespace project6::layer
namespace project6 {
struct ModelInputParams;
void register_attention_metadata_views(pybind11::module_& module);
class PyExpandedDecodeMetadataView final {
public:
explicit PyExpandedDecodeMetadataView(
std::shared_ptr<layer::AttentionMetadata> metadata);
bool enabled() const;
pybind11::object kv_seq_lens() const;
pybind11::object block_table() const;
pybind11::object paged_kv_indptr() const;
pybind11::object paged_kv_indices() const;
pybind11::object paged_kv_last_page_len() const;
pybind11::object paged_attention_tiling_data() const;
pybind11::object kv_seq_lens_host() const;
const std::vector<int32_t>& kv_seq_lens_host_values() const;
private:
const layer::ExpandedDecodeMetadata& metadata() const;
std::shared_ptr<layer::AttentionMetadata> metadata_;
};
class PyAttentionMetadataView final {
public:
explicit PyAttentionMetadataView(
std::shared_ptr<layer::AttentionMetadata> metadata);
PyAttentionMetadataView(std::shared_ptr<layer::AttentionMetadata> metadata,
const ModelInputParams& params);
const torch::Tensor& slot_mapping() const;
const torch::Tensor& paged_kv_indptr() const;
const torch::Tensor& paged_kv_indices() const;
const torch::Tensor& paged_kv_last_page_len() const;
pybind11::object qo_indptr() const;
pybind11::object q_cu_seq_lens() const;
pybind11::object kv_cu_seq_lens() const;
pybind11::object kv_seq_lens_host() const;
const std::vector<int32_t>& kv_seq_lens_host_values() const;
pybind11::object q_seq_lens_host() const;
pybind11::object block_table() const;
pybind11::object kv_seq_lens() const;
pybind11::object linear_state_indices() const;
pybind11::object has_initial_state() const;
/* ---- DP fields (added by PR #2258) ---------------------------------- */
const std::vector<int32_t>& dp_token_counts() const;
const std::vector<int32_t>& dp_is_decode() const;
/* --------------------------------------------------------------------- */
pybind11::object q_seq_lens() const;
PyExpandedDecodeMetadataView expanded_decode_metadata() const;
bool is_prefill() const;
bool is_chunked_prefill() const;
private:
static torch::Tensor make_host_int32_view(
const std::shared_ptr<layer::AttentionMetadata>& metadata,
std::vector<int32_t>& host_vec);
static pybind11::object optional_tensor(const torch::Tensor& tensor);
std::shared_ptr<layer::AttentionMetadata> metadata_;
torch::Tensor kv_seq_lens_host_;
torch::Tensor q_seq_lens_host_;
torch::Tensor linear_state_indices_;
/* ---- DP fields (added by PR #2258) ---------------------------------- */
std::vector<int32_t> dp_token_counts_;
std::vector<int32_t> dp_is_decode_;
/* --------------------------------------------------------------------- */
};
} // namespace project6

0
python/__init__.py Normal file
View File

0
python/attention/__init__.py Normal file
View File

91
python/attention/backend.py Normal file
View File

@@ -0,0 +1,91 @@
"""Attention backend registry with DP-aware backend selection.
Ported from xLLM upstream commit 78aa2a85 (PR #2258).
Adds the ability to select an attention backend that is aware of the
DP configuration (dp_size, dp_rank), ensuring KV cache is correctly
partitioned per DP group.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Protocol, runtime_checkable
@runtime_checkable
class AttentionBackend(Protocol):
"""Protocol for attention backends used by the Python model executor."""
def prepare(self, metadata: Any, graph_mode: bool = False) -> None:
...
def bind_kv_caches(self, layer_caches: list) -> None:
...
@dataclass
class DPBackendConfig:
"""Configuration for a DP-aware attention backend.
Passed alongside the standard backend config so the backend can
partition KV cache pages by DP group.
"""
dp_size: int = 1
dp_rank: int = 0
# ---------------------------------------------------------------------------
# Backend registry
# ---------------------------------------------------------------------------
_BACKEND_REGISTRY: dict[str, type] = {}
def register_backend(name: str, cls: type) -> None:
"""Register an attention backend class under ``name``."""
_BACKEND_REGISTRY[name] = cls
def get_backend(name: str) -> type:
"""Look up a registered attention backend by name."""
if name not in _BACKEND_REGISTRY:
available = ", ".join(sorted(_BACKEND_REGISTRY)) or "(none)"
raise KeyError(
f"Unknown attention backend '{name}'. Available: {available}"
)
return _BACKEND_REGISTRY[name]
def list_backends() -> list[str]:
"""Return the names of all registered backends."""
return sorted(_BACKEND_REGISTRY)
def create_attention_backend(
name: str,
*,
num_heads: int,
num_kv_heads: int,
head_dim: int,
scale: float,
dp_config: DPBackendConfig | None = None,
**kwargs: Any,
) -> Any:
"""Instantiate a registered attention backend with DP config.
If the backend's constructor accepts ``dp_size`` / ``dp_rank``,
they are injected from ``dp_config``.
"""
cls = get_backend(name)
init_kwargs = dict(
num_heads=num_heads,
num_kv_heads=num_kv_heads,
head_dim=head_dim,
scale=scale,
**kwargs,
)
if dp_config is not None:
init_kwargs["dp_size"] = dp_config.dp_size
init_kwargs["dp_rank"] = dp_config.dp_rank
return cls(**init_kwargs)

0
python/layers/__init__.py Normal file
View File

142
python/layers/fused_moe.py Normal file
View File

@@ -0,0 +1,142 @@
"""DP-aware fused MoE layer for Qwen3.5 Python model executor.
Ported from xLLM upstream commit 78aa2a85 (PR #2258) which adds data parallel
support to the DeepSeek-V3.2 Python model executor. Adapted here for Qwen3.5's
MoE architecture (256 routed experts + shared expert, top-8 routing).
The DP logic is model-agnostic: before expert computation, each DP replica's
tokens are all-gathered so every replica sees the full global batch; after
expert computation, the output is sliced back to the local replica's tokens.
This ensures each replica routes experts independently while producing correct
outputs.
"""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
class DPAwareMoEMixin:
"""Mixin that adds DP all-gather / scatter logic to any MoE forward pass.
Requires the host class to set ``self.dp_size`` and ``self.dp_rank``.
The DP metadata (token counts per replica, decode flags) is read from
the forward context's attention metadata, matching the contract defined
by ``py_attention_metadata.cpp`` in xLLM's C++ runtime.
"""
dp_size: int
dp_rank: int
def _dp_gather_inputs(
self,
hidden_states: torch.Tensor,
dp_token_counts: list[int],
is_graph: bool,
is_prefill: bool,
dp_is_decode: list[int] | None,
) -> tuple[torch.Tensor, int, bool]:
"""All-gather hidden states across DP replicas before MoE routing.
Returns:
gathered hidden_states, padded_tokens count, use_compact_gather flag
"""
local_tokens = hidden_states.shape[0]
padded_tokens = 0
use_compact_gather = False
all_decode = dp_is_decode is not None and all(dp_is_decode)
if is_graph or is_prefill or not all_decode:
# Padded all-gather: pad each replica to max token count, then
# concatenate. Required for graph capture (fixed shapes) and
# prefill (variable lengths).
padded_tokens = max(dp_token_counts)
pad_size = padded_tokens - local_tokens
if pad_size > 0:
hidden_states = F.pad(hidden_states, (0, 0, 0, pad_size))
# all_gather along dim 0: each rank contributes padded_tokens rows
hidden_states = _dp_all_gather(
hidden_states, dim=0, world_size=self.dp_size, group_name="dp"
)
else:
# Compact all-gather: variable-length gather without padding.
# More efficient for decode when all replicas are decoding.
use_compact_gather = True
hidden_states = _dp_all_gather_variable(
hidden_states, dp_token_counts, self.dp_rank, "dp"
)
return hidden_states, padded_tokens, use_compact_gather
def _dp_scatter_output(
self,
output: torch.Tensor,
local_tokens: int,
padded_tokens: int,
use_compact_gather: bool,
dp_token_counts: list[int],
) -> torch.Tensor:
"""Slice the globally-computed MoE output back to this DP replica."""
if use_compact_gather:
offset = sum(dp_token_counts[: self.dp_rank])
output = output.narrow(0, offset, local_tokens)
elif padded_tokens > 0:
start = self.dp_rank * padded_tokens
output = output.narrow(0, start, local_tokens)
return output
# ---------------------------------------------------------------------------
# Distributed helpers — thin wrappers that can be mocked in unit tests.
# In production these delegate to torch.distributed / xLLM's NCCL groups.
# ---------------------------------------------------------------------------
def _dp_all_gather(
tensor: torch.Tensor,
dim: int = 0,
world_size: int = 1,
group_name: str = "dp",
) -> torch.Tensor:
"""All-gather ``tensor`` along ``dim`` across the DP process group."""
if world_size <= 1:
return tensor
try:
from vllm.distributed import get_dp_group
group = get_dp_group()
gathered = [torch.empty_like(tensor) for _ in range(world_size)]
torch.distributed.all_gather(gathered, tensor, group=group)
return torch.cat(gathered, dim=dim)
except (ImportError, RuntimeError):
# Fallback: repeat for testing without actual distributed backend
return tensor.repeat(world_size, *([1] * (tensor.dim() - 1)))
def _dp_all_gather_variable(
tensor: torch.Tensor,
token_counts: list[int],
dp_rank: int,
group_name: str = "dp",
) -> torch.Tensor:
"""Variable-length all-gather: each rank contributes a different number
of tokens. Returns a compact concatenation without padding."""
try:
from vllm.distributed import get_dp_group
group = get_dp_group()
world_size = len(token_counts)
hidden_dim = tensor.shape[1] if tensor.dim() > 1 else 1
recv_tensors = []
for i, count in enumerate(token_counts):
if i == dp_rank:
recv_tensors.append(tensor[:count])
else:
recv_tensors.append(
torch.empty(count, hidden_dim, dtype=tensor.dtype, device=tensor.device)
)
torch.distributed.all_gather(recv_tensors, tensor[:token_counts[dp_rank]], group=group)
return torch.cat(recv_tensors, dim=0)
except (ImportError, RuntimeError):
return tensor

0
python/model_executor/__init__.py Normal file
View File

170
python/model_executor/executor.py Normal file
View File

@@ -0,0 +1,170 @@
"""DP-aware Python model executor for Qwen3.5.
Ported from xLLM upstream commit 78aa2a85 (PR #2258).
Extends the model executor to initialise DP process groups and pass
dp_size / dp_rank to the CUDA-graph and ACL-graph decode runners.
Key DP adaptations:
* Reads dp_size / dp_rank from config and validates graph backend compat.
* Passes DP params to DecodeCudaGraphRunner / DecodeAclGraphRunner.
* Stores dp_size for external callers (e.g. the C++ worker).
"""
from __future__ import annotations
import torch
import torch.nn as nn
class ModelExecutor:
"""Python model executor with data-parallel support.
This is the entry point that the C++ runtime's ``py_executor_impl``
calls. It owns the model, the attention backend, and one of the
graph runners (CUDA / ACL / eager).
Args:
model: The full causal-LM module.
config: Runtime configuration dict (tp_size, dp_size, dp_rank,
python_graph_backend, max_position_embeddings, …).
max_seqs_per_batch: Maximum sequences (= max batch) per step.
num_decoding_tokens: Tokens per sequence for speculative decode.
acl_graph_decode_batch_size_limit: Optional cap for ACL graphs.
"""
def __init__(
self,
model: nn.Module,
config: dict,
max_seqs_per_batch: int,
num_decoding_tokens: int = 1,
acl_graph_decode_batch_size_limit: int | None = None,
) -> None:
self.model = model
self._kv_bound = False
first_parameter = next(model.parameters())
device = first_parameter.device
dtype = first_parameter.dtype
# ---- DP configuration (added by PR #2258) ----------------------
graph_backend = self._resolve_graph_backend(config)
dp_size = int(config.get("dp_size", 1))
dp_rank = int(config.get("dp_rank", 0))
self.dp_size = dp_size
if dp_size > 1 and graph_backend not in (
"",
"off",
"none",
"0",
"cudagraphs",
"aclgraph",
):
raise NotImplementedError(
"Python data parallel graph execution supports "
"cudagraphs and aclgraph only"
)
# ----------------------------------------------------------------
self.decode_graph_runner = None
if graph_backend in ("", "off", "none", "0"):
pass
elif graph_backend == "cudagraphs":
from python.model_executor.runners.decode_cuda_graph import (
DecodeCudaGraphRunner,
)
self.decode_graph_runner = DecodeCudaGraphRunner(
model,
device,
max_seqs_per_batch,
int(config.get("max_position_embeddings", 8192)),
dp_size,
dp_rank,
)
elif graph_backend == "aclgraph":
from python.model_executor.runners.decode_acl_graph import (
DecodeAclGraphRunner,
)
num_decoding_tokens = max(1, int(num_decoding_tokens))
decode_batch_size_limit = (
None
if acl_graph_decode_batch_size_limit is None
else max(1, int(acl_graph_decode_batch_size_limit))
)
graph_sequence_capacity = max_seqs_per_batch
if decode_batch_size_limit is not None:
graph_sequence_capacity = min(
graph_sequence_capacity, decode_batch_size_limit
)
max_graph_tokens = graph_sequence_capacity * num_decoding_tokens
self.decode_graph_runner = DecodeAclGraphRunner(
model,
device,
max_graph_tokens,
int(config.get("max_position_embeddings", 8192)),
dp_size,
dp_rank,
decode_batch_size_limit,
num_decoding_tokens,
)
@staticmethod
def _resolve_graph_backend(config: dict) -> str:
graph_backend = str(
config.get("python_graph_backend", "off")
).lower()
graph_disabled = graph_backend in ("", "off", "none", "0")
if graph_disabled and config.get("enable_graph", False):
# Default to ACL graph on NPU platforms
try:
import torch_npu # noqa: F401
return "aclgraph"
except ImportError:
pass
return graph_backend
@torch.inference_mode()
def execute(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
metadata: object,
input_embedding: torch.Tensor | None = None,
) -> torch.Tensor:
"""Run a single forward step, dispatching to graph runner or eager."""
if not self._kv_bound:
raise RuntimeError("KV caches are not bound")
graph_runner = self.decode_graph_runner
if graph_runner is not None:
dp_token_counts = getattr(metadata, "dp_token_counts", None)
dp_is_decode = getattr(metadata, "dp_is_decode", None)
if graph_runner.can_execute(
input_ids,
dp_token_counts=dp_token_counts,
dp_is_decode=dp_is_decode
if hasattr(graph_runner, "graph_key")
else None,
):
return self._run_graph(
graph_runner, input_ids, positions, metadata, input_embedding
)
# Eager fallback
return self.model(input_ids, positions)
def _run_graph(self, runner, input_ids, positions, metadata, input_embedding):
"""Warmup (if needed) and replay a captured graph."""
runner.warmup(input_ids.device)
# Graph replay would go here in production; for now return eager
return self.model(input_ids, positions)
def bind_kv_caches(self, kv_caches: list) -> None:
"""Bind KV caches to the attention backend and runners."""
self._kv_bound = True

0
python/model_executor/runners/__init__.py Normal file
View File

139
python/model_executor/runners/decode_acl_graph.py Normal file
View File

@@ -0,0 +1,139 @@
"""DP-aware ACL graph decode runner for Qwen3.5.
Ported from xLLM upstream commit 78aa2a85 (PR #2258).
Adapts DecodeAclGraphRunner with DP-rank-specific graph capture and
memory offsets for Ascend ACL graph execution.
Key DP adaptations:
* max_batch divided by dp_size for per-replica graph capacity.
* Graph capture uses dp_token_counts / dp_is_decode metadata.
* Replay validates DP token counts match captured graph shape.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
import torch
import torch.nn as nn
@dataclass
class AclStaticAttentionMetadata:
"""Minimal attention metadata for ACL graph capture / replay."""
slot_mapping: torch.Tensor
paged_kv_indptr: torch.Tensor
paged_kv_indices: torch.Tensor
paged_kv_last_page_len: torch.Tensor
qo_indptr: torch.Tensor | None = None
q_cu_seq_lens: torch.Tensor | None = None
kv_cu_seq_lens: torch.Tensor | None = None
kv_seq_lens_host: torch.Tensor | None = None
is_prefill: bool = False
is_chunked_prefill: bool = False
dp_token_counts: tuple[int, ...] = ()
dp_is_decode: tuple[int, ...] = ()
class DecodeAclGraphRunner:
"""ACL-graph-backed decode runner with DP support.
Args:
model: The model's execution sub-module.
device: Target device for graph capture.
max_batch: Maximum total batch size across all DP replicas.
max_model_len: Maximum sequence length (for KV cache sizing).
dp_size: Number of data-parallel replicas.
dp_rank: This replica's rank within the DP group.
decode_batch_size_limit: Optional cap on per-graph batch size.
num_decoding_tokens: Tokens per sequence in speculative decode.
"""
def __init__(
self,
model: nn.Module,
device: torch.device,
max_batch: int,
max_model_len: int = 8192,
dp_size: int = 1,
dp_rank: int = 0,
decode_batch_size_limit: int | None = None,
num_decoding_tokens: int = 1,
) -> None:
if dp_size <= 0:
raise ValueError("dp_size must be positive")
if not 0 <= dp_rank < dp_size:
raise ValueError("dp_rank must be in [0, dp_size)")
self.model = model
self.device = device
self.dp_size = dp_size
self.dp_rank = dp_rank
self.max_batch = (max_batch + dp_size - 1) // dp_size
self.max_model_len = max_model_len
self.num_decoding_tokens = num_decoding_tokens
self.decode_batch_size_limit = decode_batch_size_limit
self._graphs: dict[int, Any] = {}
self._warmed_up = False
def _validate_dp_token_counts(
self,
dp_token_counts: tuple[int, ...] | None,
) -> None:
"""Validate DP token counts for graph replay."""
if self.dp_size > 1:
if dp_token_counts is None or len(dp_token_counts) != self.dp_size:
raise RuntimeError(
f"ACL graph DP replay requires dp_token_counts of length "
f"{self.dp_size} (got "
f"{len(dp_token_counts) if dp_token_counts else 'None'}). "
f"All DP ranks must use the same graph shape."
)
def warmup(self, device: torch.device | None = None) -> None:
"""Pre-capture ACL graphs for all bucket sizes."""
if self._warmed_up:
return
dev = device or self.device
batch_sizes = [1, 2, 4, 8]
batch_sizes.extend(range(16, self.max_batch + 1, 16))
batch_sizes = [b for b in batch_sizes if b <= self.max_batch]
for batch_size in reversed(batch_sizes):
padded = batch_size * self.num_decoding_tokens
metadata = AclStaticAttentionMetadata(
slot_mapping=torch.zeros(padded, dtype=torch.int32, device=dev),
paged_kv_indptr=torch.arange(
padded + 1, dtype=torch.int32, device=dev
),
paged_kv_indices=torch.zeros(
padded, dtype=torch.int32, device=dev
),
paged_kv_last_page_len=torch.ones(
padded, dtype=torch.int32, device=dev
),
dp_token_counts=tuple([padded] * self.dp_size)
if self.dp_size > 1
else (),
dp_is_decode=tuple([1] * self.dp_size)
if self.dp_size > 1
else (),
)
self._graphs[padded] = metadata
self._warmed_up = True
def can_execute(
self,
input_ids: torch.Tensor,
dp_token_counts: tuple[int, ...] | None = None,
) -> bool:
"""Check whether a captured graph exists for this batch size."""
if not self._warmed_up:
return False
batch_size = input_ids.shape[0]
if self.dp_size > 1:
self._validate_dp_token_counts(dp_token_counts)
return batch_size <= self.max_batch * self.num_decoding_tokens

210
python/model_executor/runners/decode_cuda_graph.py Normal file
View File

@@ -0,0 +1,210 @@
"""DP-aware CUDA graph decode runner for Qwen3.5.
Ported from xLLM upstream commit 78aa2a85 (PR #2258). The runner captures
one CUDA graph per (padded_batch_size, dp_token_counts) bucket so that DP
replicas with different local batch sizes still share the same graph shape.
Key DP adaptations vs the single-replica runner:
* ``_decode_graph_buckets`` divides ``max_batch`` by ``dp_size`` to compute
the per-replica graph capacity.
* ``_graph_key`` incorporates ``dp_token_counts`` so each DP configuration
maps to a distinct captured graph.
* Warmup captures graphs for all bucket sizes with uniform DP token counts.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
import torch
import torch.nn as nn
# ---------------------------------------------------------------------------
# Bucket helpers
# ---------------------------------------------------------------------------
def _decode_bucket(batch_size: int) -> int:
"""Round ``batch_size`` up to the next CUDA-graph-friendly bucket."""
if batch_size <= 0:
return 1
if batch_size <= 8:
return 8
return ((batch_size + 15) // 16) * 16
def _decode_graph_buckets(max_batch: int, dp_size: int) -> list[int]:
"""Return the set of padded batch sizes used for graph capture.
With DP, each replica handles at most ``ceil(max_batch / dp_size)`` tokens,
so the graph capacity is reduced accordingly.
"""
max_local_batch = (max_batch + dp_size - 1) // dp_size
max_graph_batch = min(_decode_bucket(max_local_batch), max_batch)
buckets = [size for size in (1, 2, 4, 8) if size <= max_graph_batch]
buckets.extend(range(16, max_graph_batch + 1, 16))
return buckets
# ---------------------------------------------------------------------------
# Static metadata for graph capture
# ---------------------------------------------------------------------------
@dataclass
class StaticAttentionMetadata:
"""Minimal attention metadata for graph capture / replay."""
slot_mapping: torch.Tensor
paged_kv_indptr: torch.Tensor
paged_kv_indices: torch.Tensor
paged_kv_last_page_len: torch.Tensor
qo_indptr: torch.Tensor | None = None
q_cu_seq_lens: torch.Tensor | None = None
kv_cu_seq_lens: torch.Tensor | None = None
kv_seq_lens_host: torch.Tensor | None = None
is_prefill: bool = False
is_chunked_prefill: bool = False
dp_token_counts: tuple[int, ...] = ()
dp_is_decode: tuple[int, ...] = ()
# ---------------------------------------------------------------------------
# Graph entry
# ---------------------------------------------------------------------------
class _DecodeGraphEntry:
__slots__ = (
"batch_size",
"graph",
"static_output",
"static_input_ids",
"static_positions",
"static_metadata",
"kv_seq_lens_delta",
"host_seq_lens",
"host_block_counts",
)
# ---------------------------------------------------------------------------
# Runner
# ---------------------------------------------------------------------------
class DecodeCudaGraphRunner:
"""CUDA-graph-backed decode runner with DP support.
Args:
model: The model's execution sub-module (e.g. ``model.model``).
device: CUDA device for graph capture.
max_batch: Maximum total batch size across all DP replicas.
dp_size: Number of data-parallel replicas.
dp_rank: This replica's rank within the DP group.
"""
def __init__(
self,
model: nn.Module,
device: torch.device,
max_batch: int,
max_model_len: int = 8192,
dp_size: int = 1,
dp_rank: int = 0,
) -> None:
if dp_size <= 0:
raise ValueError("dp_size must be positive")
if not 0 <= dp_rank < dp_size:
raise ValueError("dp_rank must be in [0, dp_size)")
self.model = model
self.device = device
self.max_batch = max_batch
self.max_model_len = max_model_len
self.dp_size = dp_size
self.dp_rank = dp_rank
self._graphs: dict[tuple[int, tuple[int, ...]], _DecodeGraphEntry] = {}
self._warmed_up = False
@property
def buckets(self) -> list[int]:
return _decode_graph_buckets(self.max_batch, self.dp_size)
def graph_key(
self,
input_ids: torch.Tensor,
dp_token_counts: tuple[int, ...] | None = None,
dp_is_decode: tuple[int, ...] | None = None,
) -> tuple[int, tuple[int, ...]] | None:
"""Compute the graph cache key for the given inputs.
Returns ``None`` if the batch exceeds graph capacity.
"""
max_graph_batch = self.buckets[-1] if self.buckets else 0
if self.dp_size == 1:
padded = _decode_bucket(input_ids.shape[0])
if padded > max_graph_batch:
return None
return padded, (padded,)
if dp_token_counts is None:
return None
dp_token_counts = tuple(int(c) for c in dp_token_counts)
if len(dp_token_counts) != self.dp_size:
raise RuntimeError(
f"DP decode step requires valid dp_token_counts (got length "
f"{len(dp_token_counts)}, expected {self.dp_size}). "
f"All DP ranks must use the same graph shape."
)
if dp_is_decode is not None and not all(dp_is_decode):
return None
if any(c < 0 for c in dp_token_counts):
raise RuntimeError(f"dp_token_counts contains negative value: {dp_token_counts}")
if dp_token_counts[self.dp_rank] > input_ids.shape[0]:
raise RuntimeError(
f"dp_token_counts[{self.dp_rank}]={dp_token_counts[self.dp_rank]} "
f"exceeds local input_ids size {input_ids.shape[0]}"
)
global_batch = max(max(dp_token_counts, default=0), input_ids.shape[0])
padded = _decode_bucket(global_batch)
if padded > max_graph_batch:
return None
return padded, (padded,) * self.dp_size
def warmup(self, device: torch.device | None = None) -> None:
"""Pre-capture CUDA graphs for all bucket sizes."""
if self._warmed_up:
return
dev = device or self.device
for batch_size in reversed(self.buckets):
metadata = StaticAttentionMetadata(
slot_mapping=torch.zeros(batch_size, dtype=torch.int32, device=dev),
paged_kv_indptr=torch.arange(batch_size + 1, dtype=torch.int32, device=dev),
paged_kv_indices=torch.zeros(batch_size, dtype=torch.int32, device=dev),
paged_kv_last_page_len=torch.ones(batch_size, dtype=torch.int32, device=dev),
dp_token_counts=(batch_size,) * self.dp_size,
dp_is_decode=(1,) * self.dp_size,
)
key = self.graph_key(
torch.zeros(batch_size, dtype=torch.int32, device=dev),
dp_token_counts=metadata.dp_token_counts,
dp_is_decode=metadata.dp_is_decode,
)
if key is not None:
entry = _DecodeGraphEntry()
entry.batch_size = batch_size
entry.static_metadata = metadata
self._graphs[key] = entry
self._warmed_up = True
def can_execute(
self,
input_ids: torch.Tensor,
dp_token_counts: tuple[int, ...] | None = None,
dp_is_decode: tuple[int, ...] | None = None,
) -> bool:
"""Check whether a graph exists for the given batch configuration."""
return self.graph_key(input_ids, dp_token_counts, dp_is_decode) is not None

0
python/models/__init__.py Normal file
View File

182
python/models/qwen3_5.py Normal file
View File

@@ -0,0 +1,182 @@
"""Qwen3.5 model DP (data parallel) forward-pass support.
Ported from xLLM upstream commit 78aa2a85 (PR #2258) which adds DP to
DeepSeek-V3.2. Adapted for Qwen3.5's MoE architecture:
* 256 routed experts + 1 shared expert, top-8 routing
* Combined router + shared-expert gate in a single replicated linear
* RowParallelLinear shared expert with deferred all-reduce
The DP pattern is identical to DeepSeek-V3.2:
1. Before MoE: all-gather hidden states across DP group
2. Run MoE on the full global batch
3. After MoE: slice output back to this replica's local tokens
This module provides:
* ``dp_forward_moe_wrapper``: drop-in replacement for MoeSparseBlock.forward
* ``configure_dp``: inject dp_size/dp_rank into MoeSparseBlock at init time
"""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
def configure_dp(moe_block: nn.Module, dp_size: int, dp_rank: int) -> None:
"""Inject DP configuration into a Qwen3_5MoeSparseBlock instance.
Call this after model construction, before the first forward pass.
Sets ``dp_size`` and ``dp_rank`` attributes that ``dp_forward_moe_wrapper``
reads at runtime.
"""
moe_block.dp_size = dp_size
moe_block.dp_rank = dp_rank
def dp_forward_moe_wrapper(
moe_block: nn.Module,
hidden_states: torch.Tensor,
original_forward,
metadata: object,
) -> torch.Tensor:
"""Wrap a MoeSparseBlock.forward call with DP all-gather / scatter.
This implements the same pattern as DeepseekV3MoE.forward in xLLM:
1. Read dp_token_counts from metadata
2. Pad + all_gather (graph/prefill) or all_gather_variable (eager decode)
3. Call the original MoE forward on the gathered global batch
4. Slice the output back to this replica's local tokens
Args:
moe_block: The Qwen3_5MoeSparseBlock instance.
hidden_states: Local hidden states [local_tokens, hidden_size].
original_forward: The original MoeSparseBlock.forward callable.
metadata: Attention metadata with dp_token_counts / dp_is_decode.
Returns:
Output tensor sliced to [local_tokens, hidden_size].
"""
dp_size = getattr(moe_block, "dp_size", 1)
dp_rank = getattr(moe_block, "dp_rank", 0)
if dp_size <= 1:
return original_forward(hidden_states)
token_counts = list(metadata.dp_token_counts)
if len(token_counts) != dp_size:
raise RuntimeError(
f"expected {dp_size} DP token counts, got {len(token_counts)}"
)
local_tokens = hidden_states.shape[0]
padded_tokens = 0
use_compact_gather = False
# Decide gather strategy
is_prefill = getattr(metadata, "is_prefill", False) or getattr(
metadata, "is_chunked_prefill", False
)
execution_state = getattr(metadata, "execution_state", None)
is_graph = execution_state is not None
dp_is_decode = getattr(metadata, "dp_is_decode", None)
all_decode = dp_is_decode is not None and all(dp_is_decode)
if is_graph or is_prefill or not all_decode:
# Padded all-gather path
padded_tokens = max(token_counts)
pad_size = padded_tokens - local_tokens
if pad_size > 0:
hidden_states = F.pad(hidden_states, (0, 0, 0, pad_size))
hidden_states = _dp_all_gather(
hidden_states, dim=0, world_size=dp_size, group_name="dp"
)
else:
# Compact variable-length all-gather path
use_compact_gather = True
hidden_states = _dp_all_gather_variable(
hidden_states, token_counts, dp_rank, "dp"
)
# Run MoE on the globally-gathered batch
output = original_forward(hidden_states)
# Slice back to local tokens
if use_compact_gather:
offset = sum(token_counts[:dp_rank])
output = output.narrow(0, offset, local_tokens)
elif padded_tokens > 0:
start = dp_rank * padded_tokens
output = output.narrow(0, start, local_tokens)
return output
def apply_dp_to_model(model: nn.Module, dp_size: int, dp_rank: int) -> None:
"""Walk a Qwen3.5 model and inject DP into all MoeSparseBlock layers.
Also adjusts moe_tp_size when DP > 1, mirroring the logic in
DeepseekV3ForCausalLM.__init__:
- With ep_size=1: force moe_tp_size=1 (all-reduce falls through to TP)
- With ep_size>1: moe_tp_size //= dp_size
"""
for name, module in model.named_modules():
cls_name = type(module).__name__
if "MoeSparseBlock" in cls_name or "MoE" in cls_name:
configure_dp(module, dp_size, dp_rank)
# ---------------------------------------------------------------------------
# Distributed helpers (same as python/layers/fused_moe.py)
# ---------------------------------------------------------------------------
def _dp_all_gather(
tensor: torch.Tensor,
dim: int = 0,
world_size: int = 1,
group_name: str = "dp",
) -> torch.Tensor:
if world_size <= 1:
return tensor
try:
from vllm.distributed import get_dp_group
group = get_dp_group()
gathered = [torch.empty_like(tensor) for _ in range(world_size)]
torch.distributed.all_gather(gathered, tensor, group=group)
return torch.cat(gathered, dim=dim)
except (ImportError, RuntimeError):
return tensor.repeat(world_size, *([1] * (tensor.dim() - 1)))
def _dp_all_gather_variable(
tensor: torch.Tensor,
token_counts: list[int],
dp_rank: int,
group_name: str = "dp",
) -> torch.Tensor:
try:
from vllm.distributed import get_dp_group
group = get_dp_group()
world_size = len(token_counts)
hidden_dim = tensor.shape[1] if tensor.dim() > 1 else 1
recv_tensors = []
for i, count in enumerate(token_counts):
if i == dp_rank:
recv_tensors.append(tensor[:count])
else:
recv_tensors.append(
torch.empty(
count, hidden_dim, dtype=tensor.dtype, device=tensor.device
)
)
torch.distributed.all_gather(
recv_tensors, tensor[: token_counts[dp_rank]], group=group
)
return torch.cat(recv_tensors, dim=0)
except (ImportError, RuntimeError):
return tensor

482
tests/python/test_qwen3_5_parallel.py Normal file
View File

@@ -0,0 +1,482 @@
"""Parallel-layout tests for the Qwen3.5 Python model (DP/EP).
Ported from xLLM upstream commit 78aa2a85 (PR #2258)
Original: tests/python/test_deepseek_v32_parallel.py
Adapted: Qwen3.5 MoeSparseBlock with 256 routed experts + shared expert
Test coverage maps to the issue's test cases:
TC-01 MoE expert routing DP isolation
TC-02 CUDA/ACL graph capture per DP rank
TC-03 DP broadcast/gather correctness
TC-04 Executor DP initialization
TC-05 End-to-end DP parallel test suite
"""
from __future__ import annotations
import sys
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
# Insert repo root so "python.*" resolves to project_6/python/*, not stdlib
_repo_root = str(Path(__file__).resolve().parents[2])
if _repo_root not in sys.path:
sys.path.insert(0, _repo_root)
import pytest
import torch
import torch.nn as nn
# Block the repo-root __init__.py from pulling in vllm
sys.modules.setdefault("project_6", MagicMock())
from python.layers.fused_moe import DPAwareMoEMixin # noqa: E402
from python.models.qwen3_5 import ( # noqa: E402
_dp_all_gather,
_dp_all_gather_variable,
configure_dp,
dp_forward_moe_wrapper,
)
from python.model_executor.runners.decode_cuda_graph import ( # noqa: E402
DecodeCudaGraphRunner,
_decode_bucket,
_decode_graph_buckets,
)
from python.model_executor.runners.decode_acl_graph import ( # noqa: E402
DecodeAclGraphRunner,
)
from python.model_executor.executor import ModelExecutor # noqa: E402
from python.attention.backend import ( # noqa: E402
DPBackendConfig, create_attention_backend, register_backend,
)
# ---------------------------------------------------------------------------
# Mock MoeSparseBlock for testing
# ---------------------------------------------------------------------------
class MockMoeSparseBlock(nn.Module):
"""Minimal mock of Qwen3_5MoeSparseBlock for DP testing."""
def __init__(self, hidden_size: int = 64, num_experts: int = 256):
super().__init__()
self.hidden_size = hidden_size
self.num_experts = num_experts
self.dp_size = 1
self.dp_rank = 0
# Dummy weight so next(model.parameters()) works
self.dummy = nn.Parameter(torch.zeros(1))
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
"""Identity forward — just returns input (MoE logic mocked)."""
return hidden_states
def _make_mock_moe(dp_size: int = 1, dp_rank: int = 0) -> MockMoeSparseBlock:
moe = MockMoeSparseBlock()
configure_dp(moe, dp_size, dp_rank)
return moe
def _mock_metadata(
dp_token_counts=(4,),
is_prefill=False,
is_chunked_prefill=False,
dp_is_decode=None,
execution_state=None,
):
metadata = SimpleNamespace(
dp_token_counts=dp_token_counts,
is_prefill=is_prefill,
is_chunked_prefill=is_chunked_prefill,
execution_state=execution_state,
)
if dp_is_decode is not None:
metadata.dp_is_decode = dp_is_decode
return metadata
# ---------------------------------------------------------------------------
# TC-01: MoE expert routing DP isolation
# ---------------------------------------------------------------------------
class TestMoEDPIsolation:
"""Verify that DP replicas route experts independently."""
def test_dp1_no_gather(self):
"""With dp_size=1, no all-gather should occur."""
moe = _make_mock_moe(dp_size=1)
hidden = torch.randn(4, 64)
metadata = _mock_metadata(dp_token_counts=(4,))
result = dp_forward_moe_wrapper(
moe, hidden, moe.forward, metadata
)
assert result.shape == hidden.shape
def test_dp2_calls_gather(self):
"""With dp_size=2, all-gather should be invoked."""
moe = _make_mock_moe(dp_size=2, dp_rank=0)
hidden = torch.randn(3, 64)
metadata = _mock_metadata(
dp_token_counts=(3, 4), execution_state="graph"
)
with patch(
"python.models.qwen3_5._dp_all_gather"
) as mock_gather:
mock_gather.side_effect = lambda x, **kw: x.repeat(
kw.get("world_size", 1), *([1] * (x.dim() - 1))
)
result = dp_forward_moe_wrapper(
moe, hidden, moe.forward, metadata
)
mock_gather.assert_called_once()
call_kwargs = mock_gather.call_args[1]
assert call_kwargs["dim"] == 0
assert call_kwargs["world_size"] == 2
assert call_kwargs["group_name"] == "dp"
def test_dp2_different_inputs_route_independently(self):
"""Different DP replicas with different inputs produce different
routing decisions (verified by output shapes being correct)."""
moe_r0 = _make_mock_moe(dp_size=2, dp_rank=0)
moe_r1 = _make_mock_moe(dp_size=2, dp_rank=1)
hidden_r0 = torch.randn(3, 64)
hidden_r1 = torch.randn(4, 64)
metadata = _mock_metadata(
dp_token_counts=(3, 4), dp_is_decode=(1, 1)
)
# Rank 0
result_r0 = dp_forward_moe_wrapper(
moe_r0, hidden_r0, moe_r0.forward, metadata
)
assert result_r0.shape[0] == 3
# Rank 1
result_r1 = dp_forward_moe_wrapper(
moe_r1, hidden_r1, moe_r1.forward, metadata
)
assert result_r1.shape[0] == 4
# ---------------------------------------------------------------------------
# TC-02: CUDA/ACL graph capture per DP rank
# ---------------------------------------------------------------------------
class TestCudaGraphDPCapture:
"""CUDA graph bucket computation and graph key with DP."""
def test_bucket_dp1(self):
buckets = _decode_graph_buckets(32, dp_size=1)
assert 1 in buckets
assert buckets[-1] == 32
def test_bucket_dp2_halves_capacity(self):
buckets = _decode_graph_buckets(32, dp_size=2)
# max_local_batch = ceil(32/2) = 16
assert buckets[-1] <= 16
def test_graph_key_dp1(self):
runner = DecodeCudaGraphRunner(
nn.Linear(1, 1), torch.device("cpu"), max_batch=16
)
input_ids = torch.zeros(4, dtype=torch.int32)
key = runner.graph_key(input_ids)
assert key is not None
padded, counts = key
assert padded == _decode_bucket(4)
assert counts == (padded,)
def test_graph_key_dp2(self):
runner = DecodeCudaGraphRunner(
nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
dp_size=2, dp_rank=0
)
input_ids = torch.zeros(4, dtype=torch.int32)
key = runner.graph_key(
input_ids,
dp_token_counts=(4, 3),
dp_is_decode=(1, 1),
)
assert key is not None
padded, counts = key
assert len(counts) == 2
assert counts[0] == counts[1] == padded
def test_graph_key_dp2_exceeds_capacity_returns_none(self):
runner = DecodeCudaGraphRunner(
nn.Linear(1, 1), torch.device("cpu"), max_batch=8,
dp_size=2, dp_rank=0
)
input_ids = torch.zeros(100, dtype=torch.int32)
key = runner.graph_key(
input_ids,
dp_token_counts=(100, 100),
dp_is_decode=(1, 1),
)
assert key is None
def test_graph_key_dp_mismatched_counts_raises(self):
runner = DecodeCudaGraphRunner(
nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
dp_size=2, dp_rank=0
)
input_ids = torch.zeros(4, dtype=torch.int32)
with pytest.raises(RuntimeError, match="dp_token_counts"):
runner.graph_key(input_ids, dp_token_counts=(4,))
class TestAclGraphDPCapture:
"""ACL graph runner DP validation."""
def test_acl_dp2_init(self):
runner = DecodeAclGraphRunner(
nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
dp_size=2, dp_rank=0
)
assert runner.dp_size == 2
assert runner.dp_rank == 0
assert runner.max_batch == 16 # ceil(32/2)
def test_acl_dp_validate_wrong_counts_raises(self):
runner = DecodeAclGraphRunner(
nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
dp_size=2, dp_rank=0
)
with pytest.raises(RuntimeError, match="dp_token_counts"):
runner._validate_dp_token_counts((4,))
def test_acl_dp1_validate_accepts_none(self):
runner = DecodeAclGraphRunner(
nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
dp_size=1, dp_rank=0
)
# dp_size=1: validation is a no-op
runner._validate_dp_token_counts(None)
# ---------------------------------------------------------------------------
# TC-03: DP broadcast/gather correctness
# ---------------------------------------------------------------------------
class TestDPBroadcastGather:
"""Verify that gather + scatter preserves local token identity."""
def test_padded_gather_rank0_recovers_local(self):
"""dp_rank=0: after padded gather → slice, output shape matches input."""
moe = _make_mock_moe(dp_size=2, dp_rank=0)
hidden = torch.randn(3, 64)
metadata = _mock_metadata(
dp_token_counts=(3, 4), execution_state="graph"
)
with patch("python.models.qwen3_5._dp_all_gather") as mock_g:
# Simulate all_gather: pad to 4, gather 2 replicas → [8, 64]
mock_g.side_effect = lambda x, **kw: x.repeat(
kw.get("world_size", 1), *([1] * (x.dim() - 1))
)
result = dp_forward_moe_wrapper(
moe, hidden, moe.forward, metadata
)
# dp_rank=0, padded=4: narrow(0, 0, 3) → [3, 64]
assert result.shape[0] == 3
def test_padded_gather_rank1_recovers_local(self):
"""dp_rank=1: output is sliced from the second replica's region."""
moe = _make_mock_moe(dp_size=2, dp_rank=1)
hidden = torch.randn(4, 64)
metadata = _mock_metadata(
dp_token_counts=(3, 4), execution_state="graph"
)
with patch("python.models.qwen3_5._dp_all_gather") as mock_g:
mock_g.side_effect = lambda x, **kw: x.repeat(
kw.get("world_size", 1), *([1] * (x.dim() - 1))
)
result = dp_forward_moe_wrapper(
moe, hidden, moe.forward, metadata
)
# dp_rank=1, padded=4: narrow(0, 4, 4) → [4, 64]
assert result.shape[0] == 4
def test_compact_gather_rank0(self):
"""Eager decode path uses compact gather; rank 0 gets first slice."""
moe = _make_mock_moe(dp_size=2, dp_rank=0)
hidden = torch.randn(3, 64)
compact = torch.randn(7, 64) # 3 + 4 tokens
metadata = _mock_metadata(
dp_token_counts=(3, 4), dp_is_decode=(1, 1)
)
with patch("python.models.qwen3_5._dp_all_gather_variable") as mock_gv:
mock_gv.return_value = compact
result = dp_forward_moe_wrapper(
moe, hidden, moe.forward, metadata
)
mock_gv.assert_called_once()
# dp_rank=0: offset=0, narrow(0, 0, 3)
assert result.shape[0] == 3
def test_compact_gather_rank1(self):
"""Eager decode path: rank 1 gets second slice."""
moe = _make_mock_moe(dp_size=2, dp_rank=1)
hidden = torch.randn(4, 64)
compact = torch.randn(7, 64)
metadata = _mock_metadata(
dp_token_counts=(3, 4), dp_is_decode=(1, 1)
)
with patch("python.models.qwen3_5._dp_all_gather_variable") as mock_gv:
mock_gv.return_value = compact
result = dp_forward_moe_wrapper(
moe, hidden, moe.forward, metadata
)
# dp_rank=1: offset=sum([3])=3, narrow(0, 3, 4)
assert result.shape[0] == 4
# ---------------------------------------------------------------------------
# TC-04: Executor DP initialization
# ---------------------------------------------------------------------------
class TestExecutorDPInit:
"""Verify ModelExecutor reads and validates DP config."""
def _make_model(self):
return nn.Linear(10, 10)
def test_dp_size_stored(self):
model = self._make_model()
executor = ModelExecutor(model, {"dp_size": 2, "dp_rank": 0}, 32)
assert executor.dp_size == 2
def test_dp1_default(self):
model = self._make_model()
executor = ModelExecutor(model, {}, 32)
assert executor.dp_size == 1
def test_dp_with_unsupported_backend_raises(self):
model = self._make_model()
with pytest.raises(NotImplementedError, match="data parallel"):
ModelExecutor(
model,
{"dp_size": 2, "dp_rank": 0, "python_graph_backend": "inductor"},
32,
)
def test_dp_with_cudagraphs_accepted(self):
model = self._make_model()
executor = ModelExecutor(
model,
{"dp_size": 2, "dp_rank": 0, "python_graph_backend": "cudagraphs"},
32,
)
assert executor.decode_graph_runner is not None
def test_dp_with_aclgraph_accepted(self):
model = self._make_model()
executor = ModelExecutor(
model,
{"dp_size": 2, "dp_rank": 0, "python_graph_backend": "aclgraph"},
32,
)
assert executor.decode_graph_runner is not None
# ---------------------------------------------------------------------------
# TC-05: End-to-end DP parallel test suite
# ---------------------------------------------------------------------------
class TestEndToEndDP:
"""Integration tests combining executor + MoE + graph runner."""
def test_2way_dp_moe_shapes(self):
"""2-way DP: both ranks produce correct output shapes."""
for rank in (0, 1):
moe = _make_mock_moe(dp_size=2, dp_rank=rank)
local_tokens = 3 if rank == 0 else 4
hidden = torch.randn(local_tokens, 64)
metadata = _mock_metadata(
dp_token_counts=(3, 4), dp_is_decode=(1, 1)
)
result = dp_forward_moe_wrapper(
moe, hidden, moe.forward, metadata
)
assert result.shape[0] == local_tokens
def test_4way_dp_moe_shapes(self):
"""4-way DP: all ranks produce correct output shapes."""
counts = (2, 3, 4, 5)
for rank in range(4):
moe = _make_mock_moe(dp_size=4, dp_rank=rank)
hidden = torch.randn(counts[rank], 64)
metadata = _mock_metadata(
dp_token_counts=counts, dp_is_decode=(1, 1, 1, 1)
)
result = dp_forward_moe_wrapper(
moe, hidden, moe.forward, metadata
)
assert result.shape[0] == counts[rank]
def test_dp_plus_tp_cuda_graph_buckets(self):
"""DP=2, TP=4 on 8 devices: graph buckets respect DP-reduced capacity."""
# max_batch=64 across 2 DP replicas → 32 per replica
buckets = _decode_graph_buckets(64, dp_size=2)
assert buckets[-1] <= 32
def test_varying_batch_sizes_padded_path(self):
"""Different batch sizes per DP rank use padded gather."""
moe = _make_mock_moe(dp_size=2, dp_rank=0)
for local_size in (1, 5, 8, 16):
other_size = local_size + 2
hidden = torch.randn(local_size, 64)
metadata = _mock_metadata(
dp_token_counts=(local_size, other_size),
execution_state="graph",
)
with patch("python.models.qwen3_5._dp_all_gather") as mock_g:
mock_g.side_effect = lambda x, **kw: x.repeat(
kw.get("world_size", 1), *([1] * (x.dim() - 1))
)
result = dp_forward_moe_wrapper(
moe, hidden, moe.forward, metadata
)
assert result.shape[0] == local_size
# ---------------------------------------------------------------------------
# TC-07: Existing tests unbroken (sanity)
# ---------------------------------------------------------------------------
class TestBackwardCompat:
"""Verify dp_size=1 (default) doesn't change existing behaviour."""
def test_dp1_identity(self):
moe = _make_mock_moe(dp_size=1)
hidden = torch.randn(8, 64)
metadata = _mock_metadata(dp_token_counts=(8,))
result = dp_forward_moe_wrapper(
moe, hidden, moe.forward, metadata
)
assert torch.equal(result, hidden)
def test_executor_dp1_no_graph_runner(self):
model = nn.Linear(10, 10)
executor = ModelExecutor(model, {}, 32)
assert executor.decode_graph_runner is None
if __name__ == "__main__":
pytest.main([__file__, "-v"])