dylanyunlong/project_6_layerwiseSplit_without_dp
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:
root
2026-09-02 07:03:56 +00:00
commit 43c43b491c
4211 changed files with 1013777 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
core/CMakeLists.txt Normal file
View File

@@ -0,0 +1,45 @@
# Copyright 2026 The xLLM Authors. All Rights Reserved.
# Commit: 494f293b5629 · feat · PR #2260 (adapted for Iluvatar BI-V100)
cmake_minimum_required(VERSION 3.18)
# ---------- layerwise split KV cache library ----------
add_library(layerwise_split_kv STATIC
framework/kv_cache/kv_cache_layerwise.cpp
framework/kv_cache/kv_cache_estimation_layerwise.cpp
framework/parallel_state/mapping_ilu.cpp
distributed_runtime/layerwise_split_engine_ext.cpp
distributed_runtime/layerwise_split_master.cpp
runtime/worker_layerwise_init.cpp
config/parallel_config_layerwise.cpp
)
target_include_directories(layerwise_split_kv PUBLIC
${CMAKE_CURRENT_SOURCE_DIR}/..
${CMAKE_CURRENT_SOURCE_DIR}
)
target_link_libraries(layerwise_split_kv PUBLIC
gflags
glog::glog
torch
)
# Iluvatar BI-V100 build: define USE_ILU
if(USE_ILU)
target_compile_definitions(layerwise_split_kv PUBLIC USE_ILU)
endif()
# ---------- tests ----------
if(BUILD_TESTING)
add_executable(test_layerwise_split
${CMAKE_CURRENT_SOURCE_DIR}/../tests/core/test_layerwise_split_kv_cache.cpp
)
target_link_libraries(test_layerwise_split PRIVATE
layerwise_split_kv
GTest::gtest_main
gflags
glog::glog
)
add_test(NAME LayerwiseSplitKVTests COMMAND test_layerwise_split)
endif()

91
core/config/ilu_hw_constants.h Normal file
View File

@@ -0,0 +1,91 @@
/* 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.
==============================================================================*/
// Iluvatar BI-V100 hardware constants.
// ALL values verified by on-device probing — do NOT change without re-probing.
//
// Probing environment:
// Machine: cc-adc62d1c-476c-4ee4-9647-0c011c0b6d70-0
// Cards: 4× Iluvatar BI-V100
// Bus-Id: 4B:00.0, 4C:00.0, 4D:00.0, 4E:00.0
// NUMA: node 1, CPU affinity 16-31,80-95
// Topology: flat PIX (all pairs via single PCIe bridge, equal BW)
// IX-ML: 3.2.3
// Driver: 3.2.1
// CUDA ver: 10.2 (CoreX compatibility layer)
// SDK path: /usr/local/corex/
//
// Probing commands used:
// ixsmi -L → card count, names, UUIDs
// ixsmi topo -m → PIX/PXB/PHB/SYS topology matrix
// ixsmi -q -d MEMORY → HBM capacity per card
// ixsmi (default) → SM clock, mem clock, TDP
// debug_warpsize.py → warp size via CUDA kernel (warpSize builtin)
// torch.cuda.get_device_properties() → partial (warp_size N/A on CoreX)
#pragma once
#include <cstdint>
namespace xllm {
namespace ilu_hw {
// ---------- Core compute ----------
/// Warp size: 64 threads (NOT 32 like NVIDIA).
/// Verified via: CUDA kernel `warpSize` builtin → 64.
/// torch.cuda.get_device_properties(0).warp_size returns N/A on CoreX.
/// This affects all warp-level primitives: __shfl, __ballot, reductions, etc.
constexpr int32_t kWarpSize = 64;
/// SM clock: 1500 MHz (from ixsmi).
constexpr int32_t kSmClockMHz = 1500;
/// Memory clock: 1200 MHz (from ixsmi).
constexpr int32_t kMemClockMHz = 1200;
// ---------- Memory ----------
/// HBM per card: 32768 MiB (from ixsmi -q -d MEMORY).
constexpr int64_t kHbmPerCardMiB = 32768;
constexpr int64_t kHbmPerCardBytes = kHbmPerCardMiB * int64_t{1024} * 1024;
/// Baseline HBM usage (driver/runtime overhead): ~257 MiB observed idle.
constexpr int64_t kHbmBaselineUsageMiB = 257;
// ---------- Topology ----------
/// Number of cards in the verified configuration.
constexpr int32_t kVerifiedCardCount = 4;
/// Topology kind: all pairs are PIX (single PCIe bridge, equal bandwidth).
/// No NVLink, no HCCS mesh, no multi-switch hierarchy.
/// If deploying on a different BI-V100 server with PXB/PHB/SYS links,
/// use IluTopoKind::kGrouped instead.
constexpr bool kFlatTopology = true;
// ---------- TDP ----------
/// TDP per card: 250W (from ixsmi Pwr cap).
constexpr int32_t kTdpWatts = 250;
// ---------- Software ----------
/// CUDA compatibility version exposed by CoreX SDK.
constexpr int32_t kCudaMajor = 10;
constexpr int32_t kCudaMinor = 2;
} // namespace ilu_hw
} // namespace xllm

31
core/config/parallel_config_layerwise.cpp Normal file
View File

@@ -0,0 +1,31 @@
/* 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.
==============================================================================*/
// Commit: 494f293b5629 · feat · PR #2260 (adapted for Iluvatar BI-V100)
// gflag definition for enabling/disabling layerwise split KV cache.
//
// Usage:
// --enable_layerwise_split=true (enable the feature)
// --enable_layerwise_split=false (default — uniform sharding, no change)
#include <gflags/gflags.h>
DEFINE_bool(enable_layerwise_split, false,
"Enable layerwise-split KV cache sharding. When true, each "
"layer's KV cache is independently sharded across a configurable "
"subset of TP ranks, allowing dense attention layers to spread "
"across all ranks while MoE layers (few KV heads, GQA) "
"concentrate on fewer ranks. Requires a heterogeneous-layer "
"model (e.g. DeepSeek-V3). Default: false (uniform sharding).");

20
core/config/parallel_config_layerwise.h Normal file
View File

@@ -0,0 +1,20 @@
/* 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.
==============================================================================*/
#pragma once
#include <gflags/gflags.h>
DECLARE_bool(enable_layerwise_split);

70
core/distributed_runtime/layerwise_split_engine_ext.cpp Normal file
View File

@@ -0,0 +1,70 @@
/* 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.
==============================================================================*/
// Commit: 494f293b5629 · feat · PR #2260 (adapted for Iluvatar BI-V100)
// Engine-level plumbing: Both llm_engine and speculative_engine propagate
// the layerwise layout to workers during initialisation.
//
// In the upstream xLLM, this would be edits to llm_engine.cpp (+18 lines)
// and speculative_engine.cpp (+12 lines). Here we isolate them in a
// self-contained compilation unit that the engines call into.
#include "distributed_runtime/layerwise_split_engine_ext.h"
#include <glog/logging.h>
#include <memory>
#include <optional>
#include <vector>
#include "common/global_flags.h"
#include "framework/kv_cache/layerwise_split_layout.h"
#include "framework/parallel_state/mapping_ilu.h"
// The flag is declared in parallel_config.cpp / global_flags.h (sub-task 7).
DECLARE_bool(enable_layerwise_split);
namespace xllm {
std::optional<LayerwiseSplitLayout> maybe_compute_layerwise_layout(
int64_t num_layers,
const std::vector<int64_t>& per_layer_kv_heads,
int32_t world_size) {
if (!FLAGS_enable_layerwise_split) {
return std::nullopt;
}
LOG(INFO) << "[LayerwiseSplit] Computing layout for " << num_layers
<< " layers, world_size=" << world_size;
#if defined(USE_ILU)
// Iluvatar BI-V100: verified 4-card flat PIX topology (ixsmi topo -m).
// All pairs connected via single PCIe bridge, equal bandwidth.
return compute_ilu_layerwise_layout(
num_layers, per_layer_kv_heads, world_size,
IluTopoKind::kFlatPIX);
#elif defined(USE_NPU)
// Ascend NPU: would use mapping_npu.cpp (not this adaptation).
LOG(WARNING) << "[LayerwiseSplit] NPU path not compiled in this build.";
return std::nullopt;
#else
// Generic CUDA fallback: flat topology (all ranks equidistant).
return compute_ilu_layerwise_layout(
num_layers, per_layer_kv_heads, world_size,
IluTopoKind::kFlatPIX);
#endif
}
} // namespace xllm

34
core/distributed_runtime/layerwise_split_engine_ext.h Normal file
View File

@@ -0,0 +1,34 @@
/* 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.
==============================================================================*/
#pragma once
#include <cstdint>
#include <optional>
#include <vector>
#include "framework/kv_cache/layerwise_split_layout.h"
namespace xllm {
/// Called by llm_engine / speculative_engine at startup.
/// Returns a LayerwiseSplitLayout if the feature is enabled, otherwise
/// std::nullopt (fallback to uniform allocation).
std::optional<LayerwiseSplitLayout> maybe_compute_layerwise_layout(
int64_t num_layers,
const std::vector<int64_t>& per_layer_kv_heads,
int32_t world_size);
} // namespace xllm

77
core/distributed_runtime/layerwise_split_master.cpp Normal file
View File

@@ -0,0 +1,77 @@
/* 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.
==============================================================================*/
// Commit: 494f293b5629 · feat · PR #2260 (adapted for Iluvatar BI-V100)
// Master-side orchestration: at startup the master reads model_args to
// extract per-layer KV head counts, computes the layout, and stores it
// for distribution to workers.
#include "distributed_runtime/layerwise_split_master.h"
#include <glog/logging.h>
#include <optional>
#include <vector>
#include "distributed_runtime/layerwise_split_engine_ext.h"
#include "framework/kv_cache/kv_cache_estimation_layerwise.h"
#include "framework/kv_cache/layerwise_split_layout.h"
DECLARE_bool(enable_layerwise_split);
namespace xllm {
std::optional<LayerwiseSplitLayout> master_compute_layerwise_layout(
int64_t num_layers,
int64_t dense_kv_heads,
int64_t moe_kv_heads,
int64_t first_moe_layer,
int32_t world_size,
int64_t n_blocks,
int64_t block_size,
int64_t head_dim,
int64_t max_tokens,
int dtype_enum) {
if (!FLAGS_enable_layerwise_split) {
LOG(INFO) << "[LayerwiseSplit] Disabled; using uniform KV sharding.";
return std::nullopt;
}
// Build per-layer KV head count vector.
// Layers [0, first_moe_layer) are dense attention; the rest are MoE.
std::vector<int64_t> per_layer_heads(num_layers);
for (int64_t i = 0; i < num_layers; ++i) {
per_layer_heads[i] = (i < first_moe_layer) ? dense_kv_heads : moe_kv_heads;
}
auto layout = maybe_compute_layerwise_layout(
num_layers, per_layer_heads, world_size);
if (layout.has_value()) {
// Run estimation for logging / capacity planning.
auto est = estimate_layerwise_kv_memory(
*layout, n_blocks, block_size, head_dim, max_tokens,
dtype_enum, world_size);
LOG(INFO) << "[LayerwiseSplit] Peak per-rank KV: "
<< (est.peak_per_rank_bytes >> 20) << " MiB (uniform would be "
<< (est.uniform_per_rank_bytes >> 20) << " MiB, saving "
<< est.savings_vs_uniform_pct << "%)";
}
return layout;
}
} // namespace xllm

40
core/distributed_runtime/layerwise_split_master.h Normal file
View File

@@ -0,0 +1,40 @@
/* 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.
==============================================================================*/
#pragma once
#include <cstdint>
#include <optional>
#include "framework/kv_cache/layerwise_split_layout.h"
namespace xllm {
/// Master-side entry point: compute and log the layerwise layout.
/// |first_moe_layer|: index of the first MoE layer (layers before it are
/// dense attention with |dense_kv_heads|).
std::optional<LayerwiseSplitLayout> master_compute_layerwise_layout(
int64_t num_layers,
int64_t dense_kv_heads,
int64_t moe_kv_heads,
int64_t first_moe_layer,
int32_t world_size,
int64_t n_blocks,
int64_t block_size,
int64_t head_dim,
int64_t max_tokens,
int dtype_enum);
} // namespace xllm

127
core/framework/kv_cache/kv_cache_estimation_layerwise.cpp Normal file
View File

@@ -0,0 +1,127 @@
/* 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.
==============================================================================*/
// Commit: 494f293b5629 · feat · PR #2260 (adapted for Iluvatar BI-V100)
// Memory estimation for layerwise-split KV cache. Reports both peak
// (bottleneck) and average per-rank utilisation so that capacity planning
// on BI-V100 (32768 MiB HBM verified via ixsmi) can account for uneven
// sharding.
#include "framework/kv_cache/kv_cache_estimation_layerwise.h"
#include <glog/logging.h>
#include <algorithm>
#include <cstdint>
#include <numeric>
#include <unordered_set>
#include <vector>
#include "config/ilu_hw_constants.h"
namespace xllm {
namespace {
/// Bytes per KV element for a given dtype.
int64_t dtype_bytes(int dtype_enum) {
// torch::kBFloat16 = 15, torch::kHalf = 5, torch::kFloat = 6
switch (dtype_enum) {
case 5: return 2; // float16
case 15: return 2; // bfloat16
case 6: return 4; // float32
case 2: return 1; // int8
default: return 2; // conservative
}
}
/// Round up to next multiple of |align|.
inline int64_t align_up(int64_t val, int64_t align) {
return ((val + align - 1) / align) * align;
}
} // namespace
LayerwiseKVMemoryEstimate estimate_layerwise_kv_memory(
const LayerwiseSplitLayout& layout,
int64_t n_blocks,
int64_t block_size,
int64_t head_dim,
int64_t max_tokens,
int dtype_enum,
int32_t world_size) {
CHECK_GT(layout.num_layers(), 0);
CHECK_GT(world_size, 0);
const int64_t elem_bytes = dtype_bytes(dtype_enum);
// BI-V100 warp = 64: the allocator pads head_dim to the next multiple
// of 64. The estimator must match, otherwise it under-reports.
#if defined(USE_ILU)
const int64_t padded_head_dim = align_up(head_dim, ilu_hw::kWarpSize);
#else
const int64_t padded_head_dim = head_dim;
#endif
// Per-rank KV bytes: sum over layers of (2 * heads * n_blocks *
// block_size * padded_head_dim * elem_bytes). Factor 2 = K + V.
std::vector<int64_t> per_rank_bytes(world_size, 0);
for (int64_t lid = 0; lid < layout.num_layers(); ++lid) {
const auto& spec = layout.layer_spec(lid);
for (size_t i = 0; i < spec.assigned_ranks.size(); ++i) {
int32_t rank = spec.assigned_ranks[i];
int64_t heads = spec.heads_per_rank[i];
int64_t layer_bytes = 2 * heads * n_blocks * block_size *
padded_head_dim * elem_bytes;
CHECK_GE(rank, 0);
CHECK_LT(rank, world_size);
per_rank_bytes[rank] += layer_bytes;
}
}
// Uniform baseline (also with padding for fair comparison).
int64_t uniform_total = 0;
for (const auto& s : layout.specs())
uniform_total += s.total_heads();
int64_t uniform_per_rank =
2 * (uniform_total / world_size) * n_blocks * block_size *
padded_head_dim * elem_bytes;
int64_t peak = *std::max_element(per_rank_bytes.begin(),
per_rank_bytes.end());
int64_t sum = std::accumulate(per_rank_bytes.begin(),
per_rank_bytes.end(), int64_t{0});
double average = static_cast<double>(sum) / world_size;
LayerwiseKVMemoryEstimate est;
est.peak_per_rank_bytes = peak;
est.average_per_rank_bytes = static_cast<int64_t>(average);
est.uniform_per_rank_bytes = uniform_per_rank;
est.per_rank_bytes = std::move(per_rank_bytes);
est.savings_vs_uniform_pct =
uniform_per_rank > 0
? 100.0 * (1.0 - static_cast<double>(peak) / uniform_per_rank)
: 0.0;
LOG(INFO) << "[LayerwiseSplit] KV memory estimate: peak="
<< (peak >> 20) << " MiB, avg="
<< (static_cast<int64_t>(average) >> 20) << " MiB, uniform="
<< (uniform_per_rank >> 20) << " MiB, saving="
<< est.savings_vs_uniform_pct << "%";
return est;
}
} // namespace xllm

44
core/framework/kv_cache/kv_cache_estimation_layerwise.h Normal file
View File

@@ -0,0 +1,44 @@
/* 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.
==============================================================================*/
#pragma once
#include <cstdint>
#include <vector>
#include "framework/kv_cache/layerwise_split_layout.h"
namespace xllm {
struct LayerwiseKVMemoryEstimate {
int64_t peak_per_rank_bytes = 0; // worst-case rank
int64_t average_per_rank_bytes = 0;
int64_t uniform_per_rank_bytes = 0; // baseline (uniform sharding)
std::vector<int64_t> per_rank_bytes; // detailed per-rank breakdown
double savings_vs_uniform_pct = 0.0;
};
/// Estimate per-rank KV cache memory for a layerwise-split layout.
/// |dtype_enum| matches torch::ScalarType integer values.
LayerwiseKVMemoryEstimate estimate_layerwise_kv_memory(
const LayerwiseSplitLayout& layout,
int64_t n_blocks,
int64_t block_size,
int64_t head_dim,
int64_t max_tokens,
int dtype_enum,
int32_t world_size);
} // namespace xllm

126
core/framework/kv_cache/kv_cache_layerwise.cpp Normal file
View File

@@ -0,0 +1,126 @@
/* 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.
==============================================================================*/
// Commit: 494f293b5629 · feat · PR #2260 (adapted for Iluvatar BI-V100)
// allocate_kv_caches_layerwise: per-layer KV allocation using
// LayerwiseSplitLayout. Each layer's shard size is determined by the number
// of heads assigned to the current rank instead of uniform division.
//
// On ILU (Iluvatar CoreX / BI-V100) the cache tensor layout is transposed:
// [n_blocks, n_heads, block_size, head_dim]
// — the head dimension sits at axis 1, not axis 2 as on CUDA/NPU.
//
// BI-V100 warp size = 64. head_dim (typically 128) is already a multiple
// of 64, so coalesced warp-wide loads across the head dimension are aligned.
// When local_heads * head_dim is not a multiple of 64, the last warp in
// a block will have idle lanes — we pad head_dim to the next multiple of
// 64 on ILU to avoid this.
#include "framework/kv_cache/kv_cache_layerwise.h"
#include "framework/kv_cache/kv_cache.h"
#include <glog/logging.h>
#include <torch/torch.h>
#include <algorithm>
#include <cstdint>
#include <vector>
#include "config/ilu_hw_constants.h"
#include "framework/kv_cache/kv_cache_utils.h"
#include "framework/kv_cache/layerwise_split_layout.h"
namespace xllm {
namespace {
/// Round up |val| to the next multiple of |align|.
inline int64_t align_up(int64_t val, int64_t align) {
return ((val + align - 1) / align) * align;
}
} // namespace
void allocate_kv_caches_layerwise(
std::vector<KVCache>& kv_caches,
const KVCacheShape& base_shape,
const KVCacheCreateOptions& create_options,
const LayerwiseSplitLayout& layout,
int32_t current_rank) {
CHECK(kv_caches.empty()) << "KV caches already initialized.";
const int64_t num_layers = create_options.num_layers();
CHECK_EQ(num_layers, layout.num_layers())
<< "Layout/config layer count mismatch.";
kv_caches.reserve(num_layers);
for (int64_t i = 0; i < num_layers; ++i) {
if (!layout.rank_owns_layer(current_rank, i)) {
kv_caches.emplace_back(); // empty placeholder
continue;
}
const int64_t local_heads = layout.heads_for_rank(current_rank, i);
CHECK_GT(local_heads, 0);
// ---------- key cache ----------
CHECK(base_shape.has_key_cache_shape());
std::vector<int64_t> k_shape = base_shape.key_cache_shape();
CHECK_GE(k_shape.size(), 4u);
// ILU/MLU transposed layout: [n_blocks, n_heads, block_size, head_dim]
// CUDA/NPU default layout: [n_blocks, block_size, n_heads, head_dim]
//
// BI-V100 warp = 64: pad head_dim to multiple of 64 so that each warp's
// contiguous load spans an aligned region. Standard head_dim (128) is
// already aligned; non-standard sizes (e.g. 96) get padded.
#if defined(USE_ILU) || defined(USE_MLU)
constexpr int64_t kHeadDimAlign = ilu_hw::kWarpSize; // 64
k_shape[1] = local_heads; // axis 1 = n_heads (transposed)
k_shape[3] = align_up(k_shape[3], kHeadDimAlign); // pad head_dim
#else
k_shape[2] = local_heads; // axis 2 = n_heads (default)
#endif
auto opts = torch::TensorOptions()
.dtype(create_options.dtype())
.device(create_options.device());
torch::Tensor k_tensor = torch::zeros(k_shape, opts);
// ---------- value cache ----------
if (base_shape.has_value_cache_shape()) {
std::vector<int64_t> v_shape = base_shape.value_cache_shape();
CHECK_GE(v_shape.size(), 4u);
#if defined(USE_ILU) || defined(USE_MLU)
v_shape[1] = local_heads;
v_shape[3] = align_up(v_shape[3], kHeadDimAlign);
#else
v_shape[2] = local_heads;
#endif
torch::Tensor v_tensor = torch::zeros(v_shape, opts);
kv_caches.emplace_back(KVCacheTensors{k_tensor, v_tensor});
} else {
kv_caches.emplace_back(KVCacheTensors{k_tensor, torch::Tensor{}});
}
}
CHECK_EQ(static_cast<int64_t>(kv_caches.size()), num_layers);
LOG(INFO) << "[LayerwiseSplit] rank " << current_rank << ": "
<< layout.layers_on_rank(current_rank) << "/" << num_layers
<< " layers assigned.";
}
} // namespace xllm

37
core/framework/kv_cache/kv_cache_layerwise.h Normal file
View File

@@ -0,0 +1,37 @@
/* 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.
==============================================================================*/
#pragma once
#include <cstdint>
#include <vector>
#include "framework/kv_cache/kv_cache.h"
#include "framework/kv_cache/kv_cache_shape.h"
#include "framework/kv_cache/kv_cache_utils.h"
#include "framework/kv_cache/layerwise_split_layout.h"
namespace xllm {
/// Allocate KV caches with per-layer head counts determined by |layout|.
/// Layers not assigned to |current_rank| receive an empty (default) KVCache.
void allocate_kv_caches_layerwise(
std::vector<KVCache>& kv_caches,
const KVCacheShape& base_shape,
const KVCacheCreateOptions& create_options,
const LayerwiseSplitLayout& layout,
int32_t current_rank);
} // namespace xllm

102
core/framework/kv_cache/layerwise_split_layout.h Normal file
View File

@@ -0,0 +1,102 @@
/* 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.
==============================================================================*/
// Commit: 494f293b5629 · feat · PR #2260 (adapted for Iluvatar BI-V100)
// Layerwise split KV cache sharding for heterogeneous layer structures
// (e.g. DeepSeek-V3: dense attention interleaved with MoE layers).
#pragma once
#include <cstdint>
#include <numeric>
#include <string>
#include <vector>
#include <glog/logging.h>
namespace xllm {
/// Per-layer KV shard descriptor.
struct LayerShardSpec {
int64_t layer_id = -1;
std::vector<int32_t> assigned_ranks; // TP ranks storing this layer's KV
std::vector<int64_t> heads_per_rank; // KV heads each rank holds
int64_t total_heads() const {
return std::accumulate(heads_per_rank.begin(), heads_per_rank.end(),
int64_t{0});
}
bool is_valid() const {
if (layer_id < 0 || assigned_ranks.empty()) return false;
if (assigned_ranks.size() != heads_per_rank.size()) return false;
for (auto h : heads_per_rank) {
if (h <= 0) return false;
}
return true;
}
};
/// Full layout: one LayerShardSpec per model layer, computed at master
/// startup and broadcast to every worker.
class LayerwiseSplitLayout {
public:
LayerwiseSplitLayout() = default;
explicit LayerwiseSplitLayout(std::vector<LayerShardSpec> specs)
: specs_(std::move(specs)) { validate(); }
int64_t num_layers() const { return static_cast<int64_t>(specs_.size()); }
const LayerShardSpec& layer_spec(int64_t lid) const {
CHECK_GE(lid, 0);
CHECK_LT(lid, num_layers());
return specs_[lid];
}
bool rank_owns_layer(int32_t rank, int64_t lid) const {
for (auto r : specs_[lid].assigned_ranks)
if (r == rank) return true;
return false;
}
int64_t heads_for_rank(int32_t rank, int64_t lid) const {
const auto& s = specs_[lid];
for (size_t i = 0; i < s.assigned_ranks.size(); ++i)
if (s.assigned_ranks[i] == rank) return s.heads_per_rank[i];
return 0;
}
int64_t layers_on_rank(int32_t rank) const {
int64_t n = 0;
for (const auto& s : specs_)
for (auto r : s.assigned_ranks)
if (r == rank) { ++n; break; }
return n;
}
void validate() const {
for (int64_t i = 0; i < num_layers(); ++i) {
CHECK(specs_[i].is_valid()) << "Invalid LayerShardSpec at " << i;
CHECK_EQ(specs_[i].layer_id, i) << "Layer id mismatch at " << i;
}
}
const std::vector<LayerShardSpec>& specs() const { return specs_; }
private:
std::vector<LayerShardSpec> specs_;
};
} // namespace xllm

100
core/framework/parallel_state/mapping_ilu.cpp Normal file
View File

@@ -0,0 +1,100 @@
/* 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.
==============================================================================*/
// Commit: 494f293b5629 · feat · PR #2260 (adapted for Iluvatar BI-V100)
// ILU-specific device-to-layer mapping for layerwise split KV cache.
//
// Verified Iluvatar BI-V100 topology (ixsmi topo -m):
// - 4 cards, Bus-Id 4B:00.0 – 4E:00.0, all on NUMA node 1
// - All pairs connected via PIX (single PCIe bridge) — FLAT topology
// - No switch hierarchy: all inter-card bandwidth is equal
// - 32 GB HBM per card (32768 MiB), 1500 MHz SM, 1200 MHz mem
// - Warp size: 64 (verified via CUDA kernel warpSize builtin)
// - IX-ML 3.2.3, Driver 3.2.1, CUDA 10.2 (CoreX)
// - CoreX SDK at /usr/local/corex/
//
// Strategy (flat PIX topology):
// Dense attention layers (many KV heads) → shard across ALL TP ranks
// MoE layers (few KV heads via GQA) → round-robin across ranks to
// balance HBM usage (no grouping benefit since all links are equal)
#include "framework/parallel_state/mapping_ilu.h"
#include <glog/logging.h>
#include <algorithm>
#include <cstdint>
#include <numeric>
#include <vector>
#include "framework/kv_cache/layerwise_split_layout.h"
namespace xllm {
LayerwiseSplitLayout compute_ilu_layerwise_layout(
int64_t num_layers,
const std::vector<int64_t>& per_layer_kv_heads,
int32_t world_size,
IluTopoKind topo_kind) {
CHECK_EQ(static_cast<int64_t>(per_layer_kv_heads.size()), num_layers);
CHECK_GT(world_size, 0);
std::vector<LayerShardSpec> specs;
specs.reserve(num_layers);
// For MoE layers with fewer heads than ranks, we round-robin the starting
// rank so that different layers land on different subsets, balancing HBM
// pressure across the flat PIX topology.
int32_t rr_offset = 0;
for (int64_t lid = 0; lid < num_layers; ++lid) {
LayerShardSpec spec;
spec.layer_id = lid;
const int64_t total_heads = per_layer_kv_heads[lid];
if (total_heads >= world_size) {
// Dense attention: shard across all ranks.
for (int32_t r = 0; r < world_size; ++r)
spec.assigned_ranks.push_back(r);
int64_t base = total_heads / world_size;
int64_t rem = total_heads % world_size;
for (int32_t r = 0; r < world_size; ++r)
spec.heads_per_rank.push_back(base + (r < rem ? 1 : 0));
} else {
// MoE / GQA layer: heads < world_size.
// Flat PIX topology — all links equal, so round-robin starting rank
// to spread HBM load evenly.
int32_t needed = static_cast<int32_t>(total_heads);
for (int32_t j = 0; j < needed; ++j) {
int32_t rank = (rr_offset + j) % world_size;
spec.assigned_ranks.push_back(rank);
}
int64_t base = total_heads / needed;
int64_t rem = total_heads % needed;
for (int32_t j = 0; j < needed; ++j)
spec.heads_per_rank.push_back(base + (j < rem ? 1 : 0));
rr_offset = (rr_offset + needed) % world_size;
}
specs.push_back(std::move(spec));
}
LOG(INFO) << "[LayerwiseSplit] ILU layout computed: " << num_layers
<< " layers, " << world_size << " ranks, topo="
<< (topo_kind == IluTopoKind::kFlatPIX ? "flat_PIX" : "grouped");
return LayerwiseSplitLayout(std::move(specs));
}
} // namespace xllm

53
core/framework/parallel_state/mapping_ilu.h Normal file
View File

@@ -0,0 +1,53 @@
/* 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.
==============================================================================*/
#pragma once
#include <cstdint>
#include <vector>
#include "framework/kv_cache/layerwise_split_layout.h"
namespace xllm {
/// Topology kind for Iluvatar BI-V100 device mapping.
/// Verified via `ixsmi topo -m` on actual hardware.
enum class IluTopoKind : int8_t {
/// All cards connected via PIX (single PCIe bridge). All inter-card
/// bandwidth is equal — no grouping benefit.
/// Observed on: 4× BI-V100, Bus-Id 4B-4E, NUMA 1.
kFlatPIX = 0,
/// Cards grouped by PCIe switch (e.g. PXB/PHB between groups).
/// Use when `ixsmi topo` shows mixed PIX + PXB/PHB/SYS entries.
kGrouped = 1,
};
/// Compute a layerwise-split layout for Iluvatar BI-V100.
///
/// |per_layer_kv_heads|: total KV head count for each layer.
/// Dense attention layers (heads >= world_size) spread across all ranks.
/// MoE / GQA layers (heads < world_size) are round-robin distributed
/// across ranks (flat PIX) or grouped by PCIe switch (grouped topology).
///
/// Default: kFlatPIX — matches the verified 4-card BI-V100 topology
/// where all pairs are PIX-connected with equal bandwidth.
LayerwiseSplitLayout compute_ilu_layerwise_layout(
int64_t num_layers,
const std::vector<int64_t>& per_layer_kv_heads,
int32_t world_size,
IluTopoKind topo_kind = IluTopoKind::kFlatPIX);
} // namespace xllm

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

73
core/runtime/worker_layerwise_init.cpp Normal file
View File

@@ -0,0 +1,73 @@
/* 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.
==============================================================================*/
// Commit: 494f293b5629 · feat · PR #2260 (adapted for Iluvatar BI-V100)
// Worker-side helper: after the worker receives its LayerwiseSplitLayout
// from the master, it calls this to allocate per-layer KV caches with
// the correct shard sizes.
#include "runtime/worker_layerwise_init.h"
#include <glog/logging.h>
#include <vector>
#include "framework/kv_cache/kv_cache.h"
#include "framework/kv_cache/kv_cache_layerwise.h"
#include "framework/kv_cache/kv_cache_shape.h"
#include "framework/kv_cache/kv_cache_utils.h"
#include "framework/kv_cache/layerwise_split_layout.h"
namespace xllm {
bool worker_allocate_layerwise_kv_cache(
std::vector<KVCache>& kv_caches,
const KVCacheShape& kv_cache_shape,
const KVCacheCreateOptions& create_options,
const LayerwiseSplitLayout& layout,
int32_t rank) {
LOG(INFO) << "[Worker " << rank << "] Applying layerwise KV layout: "
<< layout.layers_on_rank(rank) << " layers assigned.";
try {
allocate_kv_caches_layerwise(
kv_caches, kv_cache_shape, create_options, layout, rank);
} catch (const std::exception& e) {
LOG(ERROR) << "[Worker " << rank
<< "] Failed to allocate layerwise KV cache: " << e.what();
return false;
}
// Verify: assigned layers should have non-empty caches.
for (int64_t lid = 0; lid < layout.num_layers(); ++lid) {
bool owns = layout.rank_owns_layer(rank, lid);
bool empty = kv_caches[lid].empty();
if (owns && empty) {
LOG(ERROR) << "[Worker " << rank << "] Layer " << lid
<< " is assigned but KV cache is empty.";
return false;
}
if (!owns && !empty) {
LOG(ERROR) << "[Worker " << rank << "] Layer " << lid
<< " is NOT assigned but KV cache is non-empty.";
return false;
}
}
LOG(INFO) << "[Worker " << rank << "] Layerwise KV cache allocation OK.";
return true;
}
} // namespace xllm

37
core/runtime/worker_layerwise_init.h Normal file
View File

@@ -0,0 +1,37 @@
/* 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.
==============================================================================*/
#pragma once
#include <cstdint>
#include <vector>
#include "framework/kv_cache/kv_cache.h"
#include "framework/kv_cache/kv_cache_shape.h"
#include "framework/kv_cache/kv_cache_utils.h"
#include "framework/kv_cache/layerwise_split_layout.h"
namespace xllm {
/// Worker-side entry: allocate KV caches per the received layout.
/// Returns true on success; false if any verification check fails.
bool worker_allocate_layerwise_kv_cache(
std::vector<KVCache>& kv_caches,
const KVCacheShape& kv_cache_shape,
const KVCacheCreateOptions& create_options,
const LayerwiseSplitLayout& layout,
int32_t rank);
} // namespace xllm