完整部署快照
This commit is contained in:
root
389
tests/core/test_layerwise_split_kv_cache.cpp
Normal file
389
tests/core/test_layerwise_split_kv_cache.cpp
Normal file
@@ -0,0 +1,389 @@
|
||||
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
// Test suite for layerwise split KV cache sharding (PR #2260).
|
||||
// Adapted for Iluvatar BI-V100 running Qwen3.6-35B-A3B (Qwen3_5 arch).
|
||||
//
|
||||
// Verified hardware:
|
||||
// 4× BI-V100, Bus-Id 4B-4E, NUMA 1, flat PIX topology (all pairs PIX).
|
||||
// 32768 MiB HBM each, IX-ML 3.2.3, Driver 3.2.1, CUDA 10.2 (CoreX).
|
||||
// Warp size: 64 (ivcore architecture).
|
||||
//
|
||||
// Model: Qwen3.6-35B-A3B
|
||||
// num_attention_heads=16, num_key_value_heads=4, head_dim=256
|
||||
// layer_types: interleaved full_attention + linear_attention
|
||||
// With TP=4: local_q_heads=4, local_kv_heads=1, GQA=4
|
||||
// KV cache shape: key=(n,1,32,16,8) value=(n,1,256,16)
|
||||
|
||||
#include <gflags/gflags.h>
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <numeric>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "config/ilu_hw_constants.h"
|
||||
#include "config/parallel_config_layerwise.h"
|
||||
#include "distributed_runtime/layerwise_split_engine_ext.h"
|
||||
#include "distributed_runtime/layerwise_split_master.h"
|
||||
#include "framework/kv_cache/kv_cache_estimation_layerwise.h"
|
||||
#include "framework/kv_cache/layerwise_split_layout.h"
|
||||
#include "framework/model/model_args.h"
|
||||
#include "framework/kv_cache/kv_cache_estimation.h"
|
||||
#include "framework/kv_cache/ilu_layerwise_layout.h"
|
||||
#include "framework/parallel_state/mapping_ilu.h"
|
||||
#include "runtime/worker_layerwise_init.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace {
|
||||
|
||||
// Matches verified hardware: 4 BI-V100 cards.
|
||||
constexpr int32_t kBIV100WorldSize = 4;
|
||||
// Qwen3.5 model constants.
|
||||
constexpr int64_t kQwen35KVHeads = ilu_hw::kQwen35NumKVHeads; // 4
|
||||
constexpr int64_t kQwen35HeadDim = ilu_hw::kQwen35HeadDim; // 256
|
||||
|
||||
// Helper: build a realistic Qwen3.5 layer_types vector.
|
||||
// The actual pattern alternates full_attention and linear_attention.
|
||||
std::vector<std::string> make_qwen35_layer_types(int64_t num_layers) {
|
||||
std::vector<std::string> types;
|
||||
types.reserve(static_cast<size_t>(num_layers));
|
||||
for (int64_t i = 0; i < num_layers; ++i) {
|
||||
// Qwen3.5 pattern: every other layer is linear_attention
|
||||
types.push_back(i % 2 == 0 ? "full_attention" : "linear_attention");
|
||||
}
|
||||
return types;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TC-01 Layerwise KV allocation correctness (Qwen3.5 parameters)
|
||||
// ---------------------------------------------------------------------------
|
||||
TEST(LayerwiseSplitKV, TC01_AllocationCorrectness) {
|
||||
// 16 full-attention layers (Qwen3.5 has interleaved layers, but the
|
||||
// IluLayerwiseLayout only tracks full-attention layers).
|
||||
const int64_t num_full_attn = 16;
|
||||
|
||||
std::vector<int64_t> per_layer_heads(num_full_attn, kQwen35KVHeads);
|
||||
|
||||
auto layout = compute_ilu_layerwise_layout(
|
||||
num_full_attn, per_layer_heads, kBIV100WorldSize);
|
||||
|
||||
ASSERT_EQ(layout.num_layers(), num_full_attn);
|
||||
|
||||
// Qwen3.5 with TP=4: kv_heads=4 = world_size → each rank gets 1 head.
|
||||
for (int64_t lid = 0; lid < num_full_attn; ++lid) {
|
||||
const auto& spec = layout.layer_spec(lid);
|
||||
EXPECT_EQ(static_cast<int32_t>(spec.assigned_ranks.size()),
|
||||
kBIV100WorldSize);
|
||||
for (int32_t r = 0; r < kBIV100WorldSize; ++r) {
|
||||
EXPECT_EQ(layout.heads_for_rank(r, lid), 1)
|
||||
<< "Each rank should have exactly 1 KV head for Qwen3.5 TP=4";
|
||||
}
|
||||
}
|
||||
|
||||
// Total heads should match.
|
||||
int64_t total = 0;
|
||||
for (int64_t lid = 0; lid < num_full_attn; ++lid)
|
||||
total += layout.layer_spec(lid).total_heads();
|
||||
EXPECT_EQ(total, num_full_attn * kQwen35KVHeads);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TC-02 Memory estimation accuracy
|
||||
// ---------------------------------------------------------------------------
|
||||
TEST(LayerwiseSplitKV, TC02_MemoryEstimation) {
|
||||
const int64_t num_full_attn = 16;
|
||||
const int64_t n_blocks = 68837; // from real runtime log
|
||||
const int64_t block_size = 16; // from real runtime log
|
||||
const int64_t head_dim = kQwen35HeadDim; // 256
|
||||
const int64_t max_tokens = 262144;
|
||||
const int dtype_enum = 5; // float16 (from real runtime)
|
||||
|
||||
std::vector<int64_t> per_layer_heads(num_full_attn, kQwen35KVHeads);
|
||||
|
||||
auto layout = compute_ilu_layerwise_layout(
|
||||
num_full_attn, per_layer_heads, kBIV100WorldSize);
|
||||
|
||||
auto est = estimate_layerwise_kv_memory(
|
||||
layout, n_blocks, block_size, head_dim, max_tokens, dtype_enum,
|
||||
kBIV100WorldSize);
|
||||
|
||||
// With uniform sharding (kv_heads == world_size), layerwise split
|
||||
// doesn't save memory — all ranks own all layers with 1 head each.
|
||||
// Peak should equal uniform in this case.
|
||||
EXPECT_EQ(est.peak_per_rank_bytes, est.uniform_per_rank_bytes);
|
||||
EXPECT_GT(est.peak_per_rank_bytes, 0);
|
||||
EXPECT_GT(est.average_per_rank_bytes, 0);
|
||||
|
||||
// Per-rank breakdown should have exactly 4 entries.
|
||||
EXPECT_EQ(static_cast<int32_t>(est.per_rank_bytes.size()),
|
||||
kBIV100WorldSize);
|
||||
|
||||
// All ranks should have equal memory (symmetric distribution).
|
||||
for (int32_t r = 1; r < kBIV100WorldSize; ++r) {
|
||||
EXPECT_EQ(est.per_rank_bytes[0], est.per_rank_bytes[r]);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TC-03 ILU topology-aware mapping (flat PIX)
|
||||
// ---------------------------------------------------------------------------
|
||||
TEST(LayerwiseSplitKV, TC03_IluTopologyMapping) {
|
||||
const int64_t num_full_attn = 16;
|
||||
std::vector<int64_t> per_layer_heads(num_full_attn, kQwen35KVHeads);
|
||||
|
||||
auto layout = compute_ilu_layerwise_layout(
|
||||
num_full_attn, per_layer_heads, kBIV100WorldSize,
|
||||
IluTopoKind::kFlatPIX);
|
||||
|
||||
ASSERT_EQ(layout.num_layers(), num_full_attn);
|
||||
|
||||
// All layers must have at least one assigned rank.
|
||||
for (int64_t lid = 0; lid < num_full_attn; ++lid) {
|
||||
EXPECT_FALSE(layout.layer_spec(lid).assigned_ranks.empty());
|
||||
}
|
||||
|
||||
// Qwen3.5 kv_heads=4 == world_size=4: all layers on all ranks.
|
||||
for (int64_t lid = 0; lid < num_full_attn; ++lid) {
|
||||
for (int32_t r = 0; r < kBIV100WorldSize; ++r) {
|
||||
EXPECT_TRUE(layout.rank_owns_layer(r, lid));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TC-03b ILU topology with fewer heads than ranks (future model)
|
||||
// ---------------------------------------------------------------------------
|
||||
TEST(LayerwiseSplitKV, TC03b_FewerHeadsThanRanks) {
|
||||
const int64_t num_layers = 16;
|
||||
// Hypothetical model with 2 KV heads per layer (< world_size=4).
|
||||
std::vector<int64_t> per_layer_heads(num_layers, 2);
|
||||
|
||||
auto layout = compute_ilu_layerwise_layout(
|
||||
num_layers, per_layer_heads, kBIV100WorldSize,
|
||||
IluTopoKind::kFlatPIX);
|
||||
|
||||
ASSERT_EQ(layout.num_layers(), num_layers);
|
||||
|
||||
// Each layer should have exactly 2 assigned ranks.
|
||||
for (int64_t lid = 0; lid < num_layers; ++lid) {
|
||||
EXPECT_EQ(static_cast<int64_t>(layout.layer_spec(lid).assigned_ranks.size()), 2);
|
||||
}
|
||||
|
||||
// Round-robin should distribute evenly: 16 layers × 2 ranks = 32
|
||||
// assignments over 4 ranks → ~8 each.
|
||||
std::vector<int32_t> rank_count(kBIV100WorldSize, 0);
|
||||
for (int64_t lid = 0; lid < num_layers; ++lid) {
|
||||
for (auto r : layout.layer_spec(lid).assigned_ranks) {
|
||||
rank_count[r]++;
|
||||
}
|
||||
}
|
||||
int32_t min_c = *std::min_element(rank_count.begin(), rank_count.end());
|
||||
int32_t max_c = *std::max_element(rank_count.begin(), rank_count.end());
|
||||
EXPECT_LE(max_c - min_c, 1)
|
||||
<< "Round-robin should balance assignments across flat PIX topology";
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TC-04 Distributed engine layout propagation
|
||||
// ---------------------------------------------------------------------------
|
||||
TEST(LayerwiseSplitKV, TC04_EnginePropagation) {
|
||||
FLAGS_enable_layerwise_split = true;
|
||||
|
||||
std::vector<int64_t> heads(16, kQwen35KVHeads);
|
||||
|
||||
auto layout = maybe_compute_layerwise_layout(
|
||||
16, heads, kBIV100WorldSize);
|
||||
ASSERT_TRUE(layout.has_value());
|
||||
EXPECT_EQ(layout->num_layers(), 16);
|
||||
|
||||
for (int32_t r = 0; r < kBIV100WorldSize; ++r) {
|
||||
EXPECT_GT(layout->layers_on_rank(r), 0);
|
||||
}
|
||||
|
||||
FLAGS_enable_layerwise_split = false;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TC-05 Worker layer_cache_owned computation (Qwen3.5 layer types)
|
||||
// ---------------------------------------------------------------------------
|
||||
TEST(LayerwiseSplitKV, TC05_WorkerLayerCacheOwned) {
|
||||
const int64_t num_layers = 32;
|
||||
auto layer_types = make_qwen35_layer_types(num_layers);
|
||||
// 16 full_attention + 16 linear_attention
|
||||
|
||||
// rank=0, split_size=2: full-attn layers at even layer_ids (0,2,4,...30)
|
||||
// layer_id % 2 == 0 for all of them → all owned by rank 0.
|
||||
// Use rank=1 to see the split: layer_id%2==1 needed, but full layers are
|
||||
// at even ids → none owned → 16 linear (always owned) + 0 full = 16 owned.
|
||||
auto owned_r0 = worker_compute_layer_cache_owned(
|
||||
layer_types, /*layerwise_split_size=*/2, /*rank=*/0, num_layers);
|
||||
auto owned_r1 = worker_compute_layer_cache_owned(
|
||||
layer_types, /*layerwise_split_size=*/2, /*rank=*/1, num_layers);
|
||||
|
||||
ASSERT_EQ(static_cast<int64_t>(owned_r0.size()), num_layers);
|
||||
ASSERT_EQ(static_cast<int64_t>(owned_r1.size()), num_layers);
|
||||
|
||||
// Linear-attention layers always owned on both ranks.
|
||||
for (int64_t i = 1; i < num_layers; i += 2) {
|
||||
EXPECT_TRUE(owned_r0[static_cast<size_t>(i)]);
|
||||
EXPECT_TRUE(owned_r1[static_cast<size_t>(i)]);
|
||||
}
|
||||
|
||||
// Full-attention layers: rank 0 owns all (even layer_ids % 2 == 0),
|
||||
// rank 1 owns none.
|
||||
int64_t r0_full = 0, r1_full = 0;
|
||||
for (int64_t i = 0; i < num_layers; i += 2) {
|
||||
if (owned_r0[static_cast<size_t>(i)]) ++r0_full;
|
||||
if (owned_r1[static_cast<size_t>(i)]) ++r1_full;
|
||||
}
|
||||
EXPECT_EQ(r0_full, 16); // all full-attn layers owned
|
||||
EXPECT_EQ(r1_full, 0); // none owned
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TC-06 Fallback to uniform when disabled
|
||||
// ---------------------------------------------------------------------------
|
||||
TEST(LayerwiseSplitKV, TC06_FallbackUniform) {
|
||||
FLAGS_enable_layerwise_split = false;
|
||||
|
||||
std::vector<int64_t> heads(16, kQwen35KVHeads);
|
||||
auto layout = maybe_compute_layerwise_layout(
|
||||
16, heads, kBIV100WorldSize);
|
||||
EXPECT_FALSE(layout.has_value());
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TC-06b Worker with split_size=1 (all owned)
|
||||
// ---------------------------------------------------------------------------
|
||||
TEST(LayerwiseSplitKV, TC06b_WorkerNoSplit) {
|
||||
const int64_t num_layers = 32;
|
||||
auto layer_types = make_qwen35_layer_types(num_layers);
|
||||
|
||||
auto owned = worker_compute_layer_cache_owned(
|
||||
layer_types, /*layerwise_split_size=*/1, /*rank=*/0, num_layers);
|
||||
|
||||
ASSERT_EQ(static_cast<int64_t>(owned.size()), num_layers);
|
||||
|
||||
// All layers should be owned when split_size=1.
|
||||
for (int64_t i = 0; i < num_layers; ++i) {
|
||||
EXPECT_TRUE(owned[static_cast<size_t>(i)]);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TC-07 Upstream-compatible LayerwiseSplitLayout API
|
||||
// ---------------------------------------------------------------------------
|
||||
TEST(LayerwiseSplitKV, TC07_UpstreamCompatibleLayout) {
|
||||
const LayerwiseSplitLayout layout(/*enabled=*/true,
|
||||
/*group_size=*/2,
|
||||
/*local_rank=*/0);
|
||||
|
||||
// Round-robin: layer 0 → rank 0 (owns), layer 1 → rank 1 (not owned)
|
||||
EXPECT_TRUE(layout.owns(0));
|
||||
EXPECT_FALSE(layout.owns(1));
|
||||
EXPECT_TRUE(layout.owns(2));
|
||||
EXPECT_FALSE(layout.owns(3));
|
||||
|
||||
// Disabled layout → everything owned.
|
||||
const LayerwiseSplitLayout disabled(/*enabled=*/false,
|
||||
/*group_size=*/2,
|
||||
/*local_rank=*/0);
|
||||
EXPECT_TRUE(disabled.owns(0));
|
||||
EXPECT_TRUE(disabled.owns(1));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TC-08 build_layer_cache_owned with mixed layer types
|
||||
// ---------------------------------------------------------------------------
|
||||
TEST(LayerwiseSplitKV, TC08_BuildLayerCacheOwned) {
|
||||
std::vector<std::string> types = {
|
||||
"full_attention", "linear_attention", "full_attention", "linear_attention",
|
||||
"full_attention", "linear_attention", "full_attention", "linear_attention",
|
||||
};
|
||||
const LayerwiseSplitLayout layout(/*enabled=*/true,
|
||||
/*group_size=*/2,
|
||||
/*local_rank=*/0);
|
||||
|
||||
ModelArgs args;
|
||||
args.n_layers(8);
|
||||
args.layer_types(types);
|
||||
auto owned = build_layer_cache_owned(args, layout, 8);
|
||||
ASSERT_EQ(owned.size(), 8u);
|
||||
|
||||
// Linear-attention layers (indices 1,3,5,7) → always true.
|
||||
EXPECT_TRUE(owned[1]);
|
||||
EXPECT_TRUE(owned[3]);
|
||||
EXPECT_TRUE(owned[5]);
|
||||
EXPECT_TRUE(owned[7]);
|
||||
|
||||
// Full-attention layers (indices 0,2,4,6):
|
||||
// Upstream uses absolute layer_id % group_size, NOT full_attn_idx.
|
||||
// layer 0: 0%2==0 → owned (local_rank=0)
|
||||
// layer 2: 2%2==0 → owned
|
||||
// layer 4: 4%2==0 → owned
|
||||
// layer 6: 6%2==0 → owned
|
||||
// All even layer_ids are owned by rank 0! This is correct upstream behavior.
|
||||
EXPECT_TRUE(owned[0]);
|
||||
EXPECT_TRUE(owned[2]);
|
||||
EXPECT_TRUE(owned[4]);
|
||||
EXPECT_TRUE(owned[6]);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TC-09 Model type validation
|
||||
// ---------------------------------------------------------------------------
|
||||
TEST(LayerwiseSplitKV, TC09_ModelTypeValidation) {
|
||||
EXPECT_TRUE(is_layerwise_split_supported_model("qwen3_5"));
|
||||
EXPECT_TRUE(is_layerwise_split_supported_model("qwen3_5_moe_text"));
|
||||
EXPECT_TRUE(is_layerwise_split_supported_model("deepseek_v32"));
|
||||
EXPECT_TRUE(is_layerwise_split_supported_model("glm_moe_dsa"));
|
||||
EXPECT_FALSE(is_layerwise_split_supported_model("llama"));
|
||||
EXPECT_FALSE(is_layerwise_split_supported_model(""));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TC-10 Block count estimation with layerwise split
|
||||
// ---------------------------------------------------------------------------
|
||||
TEST(LayerwiseSplitKV, TC10_BlockCountEstimation) {
|
||||
// Qwen3.5 per-layer KV cache per block (TP=4, 1 kv head, head_dim=256):
|
||||
// key: 1 * 16 * 16 * 8 * 2 = 4096 bytes
|
||||
// value: 1 * 256 * 16 * 2 = 8192 bytes
|
||||
// total per layer per block: ~12288 bytes (but varies by layout)
|
||||
//
|
||||
// Simplified: use 16384 bytes per layer per block (K+V).
|
||||
const int64_t per_layer_block_bytes = 16384;
|
||||
const int64_t scratch_block_bytes = 16384;
|
||||
const int64_t available_bytes = int64_t{30} * 1024 * 1024 * 1024; // ~30 GiB
|
||||
|
||||
auto blocks_split1 = estimate_layerwise_split_block_count(
|
||||
/*layerwise_split_size=*/1, /*num_full_attn_layers=*/16,
|
||||
per_layer_block_bytes, scratch_block_bytes, available_bytes);
|
||||
|
||||
auto blocks_split2 = estimate_layerwise_split_block_count(
|
||||
/*layerwise_split_size=*/2, /*num_full_attn_layers=*/16,
|
||||
per_layer_block_bytes, scratch_block_bytes, available_bytes);
|
||||
|
||||
// With split_size=2, each rank owns ~8 layers instead of 16,
|
||||
// so it should be able to afford more blocks.
|
||||
EXPECT_GT(blocks_split2, blocks_split1);
|
||||
EXPECT_GT(blocks_split1, 0);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace xllm
|
||||
482
tests/python/test_qwen3_5_parallel.py
Normal file
482
tests/python/test_qwen3_5_parallel.py
Normal 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"])
|
||||
Reference in New Issue
Block a user