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

完整部署快照

Browse Source
This commit is contained in:
root
2026-09-20 02:20:26 +00:00
commit ce553c0dc8
5836 changed files with 1631123 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
core/CMakeLists.txt Normal file
View File

@@ -0,0 +1,63 @@
# Copyright 2026 The xLLM Authors. All Rights Reserved.
# Layerwise split KV cache — migrated from upstream PR #2260
# Adapted for Iluvatar BI-V100 + Qwen3.5 (Qwen3.6-35B-A3B)
cmake_minimum_required(VERSION 3.18)
project(layerwise_split_kv_cache LANGUAGES CXX)
set(CMAKE_CXX_STANDARD 17)
find_package(PkgConfig REQUIRED)
pkg_check_modules(GLOG REQUIRED IMPORTED_TARGET libglog)
find_package(GTest QUIET)
# ---------- layerwise split KV cache library ----------
add_library(layerwise_split_kv STATIC
# Upstream-migrated files
framework/config/parallel_config.cpp
framework/kv_cache/kv_cache_estimation.cpp
# BI-V100 adaptation files
config/parallel_config_layerwise.cpp
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
)
target_include_directories(layerwise_split_kv PUBLIC
${CMAKE_CURRENT_SOURCE_DIR}/..
${CMAKE_CURRENT_SOURCE_DIR}
)
target_link_libraries(layerwise_split_kv PUBLIC
gflags
PkgConfig::GLOG
)
# Iluvatar BI-V100 build: define USE_ILU
if(USE_ILU)
target_compile_definitions(layerwise_split_kv PUBLIC USE_ILU)
endif()
# ---------- tests ----------
if(BUILD_TESTING)
if(NOT GTest_FOUND)
find_library(GTEST_LIB gtest)
find_library(GTEST_MAIN_LIB gtest_main)
endif()
add_executable(test_layerwise_split
${CMAKE_CURRENT_SOURCE_DIR}/../tests/core/test_layerwise_split_kv_cache.cpp
)
if(GTest_FOUND)
target_link_libraries(test_layerwise_split PRIVATE
layerwise_split_kv GTest::gtest_main gflags PkgConfig::GLOG)
else()
target_link_libraries(test_layerwise_split PRIVATE
layerwise_split_kv ${GTEST_LIB} ${GTEST_MAIN_LIB} gflags PkgConfig::GLOG pthread)
endif()
add_test(NAME LayerwiseSplitKVTests COMMAND test_layerwise_split)
endif()

19
core/common/macros.h Normal file
View File

@@ -0,0 +1,19 @@
/* Lightweight PROPERTY macro for project_6.
Adapted from xllm/core/common/macros.h — stripped reflection. */
#pragma once
#include <utility>
#define PROPERTY(T, property) \
public: \
[[nodiscard]] const T& property() const& noexcept { return property##_; } \
[[nodiscard]] T& property() & noexcept { return property##_; } \
auto property(const T& value) & -> decltype(*this) { \
property##_ = value; \
return *this; \
} \
auto property(T&& value) & -> decltype(*this) { \
property##_ = std::move(value); \
return *this; \
} \
T property##_

126
core/config/ilu_hw_constants.h 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.
==============================================================================*/
// 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;
// ---------- KV cache tensor layout (verified from runtime logs) ----------
//
// key_cache shape = (n_blocks, n_kv_heads_local, block_size, head_dim/x, x)
// example = (68837, 1, 32, 16, 8)
// value_cache shape = (n_blocks, n_kv_heads_local, head_dim, block_size)
// example = (68837, 1, 256, 16)
//
// This is the transposed ("block-major") layout used by the vendor vLLM
// for Iluvatar BI-V100. Key cache splits head_dim into two factors
// (head_dim/x, x) where x=8 for fp16 (128-bit / sizeof(fp16)).
// Value cache transposes block_size to the last axis.
//
// Source: runtime log with Qwen3.6-35B-A3B (Qwen3_5 arch), TP=4:
// [BI100 PAGED_ATTN] decode_dispatch ... query=(1, 4, 256)
// key_cache=(68837, 1, 32, 16, 8) value_cache=(68837, 1, 256, 16)
// block_tables=(1, 2841) required_blocks=2841 threshold=32768
//
// The x factor for the key cache head_dim split:
// x = 128 / sizeof(dtype_bits) → 128/16=8 for fp16, 128/16=8 for bf16
constexpr int32_t kKeyCacheHeadDimSplitFp16 = 8;
/// Default block size used by the vendor vLLM paged attention.
constexpr int32_t kDefaultBlockSize = 16;
// ---------- Qwen3.5 model constants (Qwen3.6-35B-A3B) ----------
// From the model config and verified runtime shapes:
// num_attention_heads=16, num_key_value_heads=4, head_dim=256
// With TP=4: local_q_heads=4, local_kv_heads=1, GQA ratio=4
// layer_types: interleaved full_attention and linear_attention
// Linear attention layers (GDN/DeltaNet) do NOT use KV cache.
constexpr int32_t kQwen35NumAttentionHeads = 16;
constexpr int32_t kQwen35NumKVHeads = 4;
constexpr int32_t kQwen35HeadDim = 256;
constexpr int32_t kQwen35GQARatio = kQwen35NumAttentionHeads / kQwen35NumKVHeads;
// ---------- Software ----------
/// CUDA compatibility version exposed by CoreX SDK.
constexpr int32_t kCudaMajor = 10;
constexpr int32_t kCudaMinor = 2;
} // namespace ilu_hw
} // namespace xllm

40
core/config/parallel_config_layerwise.cpp 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.
==============================================================================*/
// Commit: 494f293b5629 · feat · PR #2260 (adapted for Iluvatar BI-V100)
// gflag definitions for layerwise split KV cache.
//
// Usage:
// --enable_layerwise_split=true --layerwise_split_size=2
//
// For Qwen3.5 (Qwen3.6-35B-A3B) on 4× BI-V100 with TP=4:
// layerwise_split_size=2 means each rank owns ~half the full-attention
// layers' KV cache, reducing peak per-rank KV memory.
#include <gflags/gflags.h>
DEFINE_bool(enable_layerwise_split, false,
"Enable layerwise-split KV cache sharding. When true, each "
"full-attention layer's KV cache ownership is assigned to a "
"subset of TP ranks via round-robin. Linear-attention layers "
"(GDN/DeltaNet) are unaffected — they use conv+temporal state, "
"not KV cache. Default: false (uniform sharding).");
DEFINE_int32(layerwise_split_size, 1,
"Layer-owner KV cache group size inside each attention TP group. "
"1 disables layerwise split; values > 1 shard persistent KV by "
"layer owner and enable layerwise-split communication. The value "
"must divide attention TP size. Matches upstream xLLM "
"parallel_config.layerwise_split_size.");

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

@@ -0,0 +1,25 @@
/* 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>
// Bool flag for enabling layerwise split (used by engine_ext / master).
DECLARE_bool(enable_layerwise_split);
// Int32 flag matching upstream xLLM's parallel_config.
// 1 = disabled (default), >1 = group size for layerwise split.
DECLARE_int32(layerwise_split_size);

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

@@ -0,0 +1,63 @@
/* 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.
//
// The vendor vLLM on BI-V100 uses prebuilt CoreX .so extensions
// (corex_paged_kv_gather.so, etc.) for the kernel layer. This module
// operates at the scheduling layer above — deciding which layers each
// rank owns for KV cache purposes.
#include "distributed_runtime/layerwise_split_engine_ext.h"
#include <glog/logging.h>
#include <memory>
#include <optional>
#include <vector>
#include "framework/kv_cache/layerwise_split_layout.h"
#include "framework/parallel_state/mapping_ilu.h"
// The flag is declared in parallel_config_layerwise.h / .cpp (sub-task 7).
DECLARE_bool(enable_layerwise_split);
namespace xllm {
std::optional<IluLayerwiseLayout> 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
<< " full-attention layers, world_size=" << world_size;
// 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);
}
} // namespace xllm

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

@@ -0,0 +1,38 @@
/* 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/ilu_layerwise_layout.h"
namespace xllm {
/// Called by llm_engine / speculative_engine at startup.
/// Returns a IluLayerwiseLayout if the feature is enabled, otherwise
/// std::nullopt (fallback to uniform allocation).
///
/// |per_layer_kv_heads|: KV head count for each full-attention layer.
/// Linear-attention layers are excluded.
/// For Qwen3.5 TP=4: all entries are 4 (each full-attn layer has 4 KV heads).
std::optional<IluLayerwiseLayout> maybe_compute_layerwise_layout(
int64_t num_layers,
const std::vector<int64_t>& per_layer_kv_heads,
int32_t world_size);
} // namespace xllm

91
core/distributed_runtime/layerwise_split_master.cpp 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.
==============================================================================*/
// Commit: 494f293b5629 · feat · PR #2260 (adapted for Iluvatar BI-V100)
// Master-side orchestration: at startup the master reads model config to
// extract per-layer KV head counts, computes the layout, and stores it
// for distribution to workers.
//
// For Qwen3.5 (Qwen3.6-35B-A3B):
// - layer_types is a mix of "full_attention" and "linear_attention"
// - Only full_attention layers have KV cache
// - Each full_attention layer has the same KV head count (4 for the
// 35B-A3B variant)
#include "distributed_runtime/layerwise_split_master.h"
#include <glog/logging.h>
#include <optional>
#include <string>
#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<IluLayerwiseLayout> master_compute_layerwise_layout(
const std::vector<std::string>& layer_types,
int64_t kv_heads_per_full_attn_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;
}
// Count and build per-layer KV head count vector for full-attention
// layers only. Linear-attention layers are skipped.
int64_t full_attn_count = 0;
for (const auto& lt : layer_types) {
if (lt == "full_attention") {
++full_attn_count;
}
}
CHECK_GT(full_attn_count, 0) << "No full_attention layers found.";
std::vector<int64_t> per_layer_heads(full_attn_count,
kv_heads_per_full_attn_layer);
auto layout = maybe_compute_layerwise_layout(
full_attn_count, 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] Qwen3.5 layout: "
<< full_attn_count << " full-attn layers, "
<< (layer_types.size() - full_attn_count) << " linear-attn layers"
<< " — 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

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

@@ -0,0 +1,43 @@
/* 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 <string>
#include <vector>
#include "framework/kv_cache/ilu_layerwise_layout.h"
namespace xllm {
/// Master-side entry point: compute and log the layerwise layout.
///
/// For Qwen3.5 (Qwen3.6-35B-A3B):
/// |layer_types|: "full_attention" or "linear_attention" per layer.
/// |kv_heads_per_full_attn_layer|: KV heads for each full-attention layer.
/// Linear-attention layers are skipped (they use GDN conv+temporal state).
std::optional<IluLayerwiseLayout> master_compute_layerwise_layout(
const std::vector<std::string>& layer_types,
int64_t kv_heads_per_full_attn_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

42
core/framework/config/parallel_config.cpp Normal file
View File

@@ -0,0 +1,42 @@
/* Migrated from xllm/core/framework/config/parallel_config.cpp
Stripped json — uses gflags directly. */
#include "framework/config/parallel_config.h"
#include <gflags/gflags.h>
DEFINE_int32(dp_size, 1, "Data parallel size.");
DEFINE_int32(ep_size, 1, "Expert parallel size for MoE model.");
DEFINE_int32(cp_size, 1, "Context parallel size.");
DEFINE_int32(
layerwise_split_size,
1,
"Layer-owner KV cache group size inside each attention TP group. "
"1 disables layerwise split; values > 1 shard persistent KV by layer owner "
"and enable layerwise-split communication. The value must divide attention "
"TP size.");
DEFINE_int32(kv_split_size, 1, "KV-cache split width.");
DEFINE_int64(tp_size, 1, "Tensor parallelism size.");
DEFINE_string(communication_backend, "nccl", "Communication backend.");
DEFINE_bool(enable_multi_stream_parallel, false,
"Enable multi-stream parallel.");
namespace xllm {
void ParallelConfig::from_flags() {
dp_size(FLAGS_dp_size);
ep_size(FLAGS_ep_size);
cp_size(FLAGS_cp_size);
layerwise_split_size(FLAGS_layerwise_split_size);
kv_split_size(FLAGS_kv_split_size);
tp_size(FLAGS_tp_size);
communication_backend(FLAGS_communication_backend);
enable_multi_stream_parallel(FLAGS_enable_multi_stream_parallel);
}
ParallelConfig& ParallelConfig::get_instance() {
static ParallelConfig config;
return config;
}
} // namespace xllm

34
core/framework/config/parallel_config.h Normal file
View File

@@ -0,0 +1,34 @@
/* Migrated from xllm/core/framework/config/parallel_config.h
Stripped json / option_category — project_6 uses gflags directly. */
#pragma once
#include <cstdint>
#include <string>
#include "common/macros.h"
namespace xllm {
class ParallelConfig final {
public:
ParallelConfig() = default;
static ParallelConfig& get_instance();
void from_flags();
PROPERTY(int32_t, dp_size) = 1;
PROPERTY(int32_t, ep_size) = 1;
PROPERTY(int32_t, cp_size) = 1;
PROPERTY(int32_t, layerwise_split_size) = 1;
PROPERTY(int32_t, kv_split_size) = 1;
PROPERTY(int64_t, tp_size) = 1;
PROPERTY(std::string, communication_backend) = "nccl";
PROPERTY(bool, enable_multi_stream_parallel) = false;
[[nodiscard]] int32_t kv_split_size_effective() const noexcept {
return kv_split_size_ > 0 ? kv_split_size_ : cp_size_;
}
};
} // namespace xllm

84
core/framework/kv_cache/ilu_layerwise_layout.h Normal file
View File

@@ -0,0 +1,84 @@
/* BI-V100-specific detailed per-layer shard descriptors.
This is the 20% adaptation on top of upstream LayerwiseSplitLayout. */
#pragma once
#include <cstdint>
#include <numeric>
#include <vector>
#include <glog/logging.h>
namespace xllm {
/// Per-layer KV shard descriptor for BI-V100.
struct LayerShardSpec {
int64_t layer_id = -1;
std::vector<int32_t> assigned_ranks;
std::vector<int64_t> heads_per_rank;
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 BI-V100-specific layout: one LayerShardSpec per model layer.
class IluLayerwiseLayout {
public:
IluLayerwiseLayout() = default;
explicit IluLayerwiseLayout(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

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

@@ -0,0 +1,44 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
Migrated from xllm/core/framework/kv_cache/kv_cache_capacity.h */
#pragma once
#include <cstdint>
#include "common/macros.h"
namespace xllm {
class KVCacheCapacity final {
public:
PROPERTY(int64_t, n_blocks) = 0;
PROPERTY(int64_t, cache_size_in_bytes) = 0;
PROPERTY(int64_t, block_size) = 0;
PROPERTY(int64_t, slot_size) = 0;
// for index cache
PROPERTY(int64_t, index_slot_size) = 0;
PROPERTY(int64_t, num_indexer_layers) = 0;
PROPERTY(bool, enable_indexer_cache_quant) = false;
// for kv cache quantization scale cache
PROPERTY(int64_t, scale_slot_size) = 0;
// for linear attention
PROPERTY(int64_t, linear_slot_size) = 0;
PROPERTY(int64_t, linear_cache_size_in_bytes) = 0;
PROPERTY(int64_t, linear_conv_state_len) = 0;
PROPERTY(int64_t, linear_ssm_checkpoint_stride) = 1;
PROPERTY(int64_t, n_layers) = 0;
PROPERTY(int64_t, num_linear_state_blocks) = 0;
PROPERTY(int64_t, num_full_attention_layers) = 0;
PROPERTY(int64_t, num_linear_attention_layers) = 0;
// DeepSeek V4 uses separate block pools for sliding-window and compressed
// caches. These fields are only meaningful for deepseek_v4.
PROPERTY(int64_t, swa_count) = 0;
PROPERTY(int64_t, c4_count) = 0;
PROPERTY(int64_t, c128_count) = 0;
};
} // namespace xllm

370
core/framework/kv_cache/kv_cache_estimation.cpp Normal file
View File

@@ -0,0 +1,370 @@
/* Migrated from xllm/core/framework/kv_cache/kv_cache_estimation.cpp (729 lines)
[BI-V100 20%] Changes:
- Removed torch dependency (dtype_size via switch on int32_t enum)
- Removed Platform::supports_dsa_* calls (no DSA indexer on BI-V100)
- Removed DeepSeek V4 estimation (not target model)
- Removed MLU/NPU conditional compilation
- Kept layerwise_split_block_count, build_layer_cache_owned,
estimate_layerwise_split_block_count, init_standard_counts,
estimate_kv_cache_capacity — the scheduling-layer functions */
#include "framework/kv_cache/kv_cache_estimation.h"
#include <algorithm>
#include <cstddef>
#include <limits>
#include <vector>
#include <glog/logging.h>
#include "framework/kv_cache/layerwise_split_layout.h"
#include "framework/model/model_args.h"
namespace xllm {
namespace {
// [BI-V100 20%] Replaces torch::elementSize(dtype) — no libtorch needed.
int64_t dtype_size_from_enum(int32_t dtype_enum) {
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;
}
}
// Upstream: kv_cache_dtype_size
int64_t kv_cache_dtype_size(const std::string& kv_cache_dtype,
int64_t model_dtype_size) {
if (kv_cache_dtype == "auto") {
return model_dtype_size;
}
if (kv_cache_dtype == "int8") {
return 1;
}
if (kv_cache_dtype == "fp8_e4m3" || kv_cache_dtype == "fp8_e5m2") {
return 1;
}
return model_dtype_size;
}
// Upstream: kv_slot_size
int64_t kv_slot_size(const ModelArgs& model_args,
const KVCacheEstimateOptions& options,
int64_t cache_dtype_size) {
if (model_args.enable_mla()) {
return cache_dtype_size *
(model_args.kv_lora_rank() + model_args.qk_rope_head_dim());
}
return 2 * cache_dtype_size * model_args.head_dim() *
options.n_local_kv_heads;
}
// Upstream: scale_slot_size
int64_t scale_slot_size(const ModelArgs& model_args,
const KVCacheEstimateOptions& options) {
if (options.kv_cache_dtype == "auto") {
return 0;
}
if (model_args.enable_mla()) {
return sizeof(float);
}
return 2 * sizeof(float) * options.n_local_kv_heads;
}
// Upstream: standard_full_cache_block_size_in_bytes
int64_t standard_full_cache_block_size_in_bytes(
const KVCacheCapacity& kv_cache_cap) {
const int64_t full_attention_layers =
std::max<int64_t>(kv_cache_cap.num_full_attention_layers(), 1);
const int64_t indexer_layers = kv_cache_cap.num_indexer_layers();
CHECK_GE(indexer_layers, 0) << "num_indexer_layers must be non-negative";
CHECK_LE(indexer_layers, full_attention_layers)
<< "num_indexer_layers cannot exceed full-attention layers";
const int64_t logical_block_bytes =
kv_cache_cap.block_size() *
(full_attention_layers *
(kv_cache_cap.slot_size() + kv_cache_cap.scale_slot_size()) +
indexer_layers * kv_cache_cap.index_slot_size());
CHECK_GT(logical_block_bytes, 0) << "logical block bytes must be positive";
return logical_block_bytes;
}
// Upstream: layerwise_split_block_count (lines 126-186)
int64_t layerwise_split_block_count(const ModelArgs& model_args,
int32_t layerwise_split_size,
const KVCacheCapacity& kv_cache_cap,
int64_t available_bytes,
int64_t additional_block_bytes) {
const int64_t num_layers = kv_cache_cap.n_layers();
const std::vector<bool> indexer_layer_mask =
resolve_indexer_cache_enabled_layers(model_args, num_layers);
std::vector<int64_t> layer_bytes(static_cast<size_t>(num_layers), 0);
bool any_indexer_layer = false;
for (int64_t layer_id = 0; layer_id < num_layers; ++layer_id) {
if (!is_full_attention_layer(model_args, layer_id)) {
continue;
}
const bool has_indexer =
kv_cache_cap.index_slot_size() > 0 &&
(indexer_layer_mask.empty() ||
indexer_layer_mask[static_cast<size_t>(layer_id)]);
any_indexer_layer = any_indexer_layer || has_indexer;
const int64_t bytes =
kv_cache_cap.block_size() *
(kv_cache_cap.slot_size() + kv_cache_cap.scale_slot_size() +
(has_indexer ? kv_cache_cap.index_slot_size() : 0));
layer_bytes[static_cast<size_t>(layer_id)] = bytes;
}
// Scratch matches an owned layer, including indexer when any layer has one.
const int64_t scratch_bytes_per_block =
kv_cache_cap.block_size() *
(kv_cache_cap.slot_size() + kv_cache_cap.scale_slot_size() +
(any_indexer_layer ? kv_cache_cap.index_slot_size() : 0));
CHECK_GT(scratch_bytes_per_block, 0);
int64_t common_block_count = std::numeric_limits<int64_t>::max();
for (int32_t split_rank = 0; split_rank < layerwise_split_size;
++split_rank) {
const LayerwiseSplitLayout layout(
/*enabled=*/true, layerwise_split_size, split_rank);
const std::vector<bool> layer_cache_owned =
build_layer_cache_owned(model_args, layout, num_layers);
int64_t owned_bytes = 0;
for (int64_t layer_id = 0; layer_id < num_layers; ++layer_id) {
if (layer_cache_owned[static_cast<size_t>(layer_id)]) {
owned_bytes += layer_bytes[static_cast<size_t>(layer_id)];
}
}
if (owned_bytes == 0) {
continue;
}
const int64_t per_block_bytes =
owned_bytes + scratch_bytes_per_block + additional_block_bytes;
common_block_count =
std::min(common_block_count, available_bytes / per_block_bytes);
}
CHECK_NE(common_block_count, std::numeric_limits<int64_t>::max())
<< "No layerwise split rank owns a model layer.";
CHECK_GT(common_block_count, 0) << "No memory for one layerwise split block.";
return common_block_count;
}
// Upstream: enable_qwen3_5_spec_verify
bool enable_qwen3_5_spec_verify(const ModelArgs& model_args,
const KVCacheEstimateOptions& options) {
return options.num_speculative_tokens > 0 && !options.is_draft_engine &&
is_qwen3_5_target_model_type(model_args.model_type());
}
// Upstream: linear_slot_size
int64_t linear_slot_size(const ModelArgs& model_args,
const KVCacheEstimateOptions& options,
int64_t dtype_size) {
if (model_args.linear_num_value_heads() <= 0) {
return 0;
}
const int64_t num_speculative_tokens =
enable_qwen3_5_spec_verify(model_args, options)
? options.num_speculative_tokens
: 0;
const int64_t head_k_dim = model_args.linear_key_head_dim();
const int64_t head_v_dim = model_args.linear_value_head_dim();
// [BI-V100 20%] ssm_dtype defaults to model dtype on BI-V100
const int64_t ssm_dtype_size = dtype_size;
const int64_t linear_ssm_slot_size =
ssm_dtype_size * options.n_local_linear_v_heads * head_k_dim * head_v_dim;
const int64_t linear_conv_state_len =
model_args.linear_conv_kernel_dim() - 1 + num_speculative_tokens;
const int64_t linear_conv_slot_size =
dtype_size *
(head_k_dim * options.n_local_linear_k_heads * 2 +
head_v_dim * options.n_local_linear_v_heads) *
linear_conv_state_len;
return linear_conv_slot_size +
linear_ssm_slot_size * (num_speculative_tokens + 1);
}
constexpr int64_t kPaddingLinearStateBlocks = 2;
// Upstream: calculate_linear_state_blocks (simplified — no prefix cache auto-sizing)
int64_t calculate_linear_state_blocks(int64_t cache_size_in_bytes,
int64_t num_linear_attention_layers,
int64_t linear_slot_sz,
int64_t full_cache_block_size_in_bytes,
int64_t max_seqs_per_batch,
int64_t max_linear_state_cache_slots,
bool enable_prefix_cache) {
if (num_linear_attention_layers <= 0 || linear_slot_sz <= 0) {
return kPaddingLinearStateBlocks;
}
const int64_t linear_bytes_per_block =
num_linear_attention_layers * linear_slot_sz;
CHECK_GT(linear_bytes_per_block, 0);
int64_t max_blocks =
(cache_size_in_bytes - 1) / linear_bytes_per_block;
const int64_t balanced =
(cache_size_in_bytes +
kPaddingLinearStateBlocks * full_cache_block_size_in_bytes) /
(linear_bytes_per_block + full_cache_block_size_in_bytes);
max_blocks = std::min(max_blocks, balanced);
max_blocks = std::max(max_blocks, kPaddingLinearStateBlocks);
if (max_linear_state_cache_slots > 0) {
const int64_t requested =
max_linear_state_cache_slots + kPaddingLinearStateBlocks;
CHECK_LE(requested, max_blocks)
<< "max_linear_state_cache_slots requires " << requested
<< " linear-state blocks, but only " << max_blocks << " fit.";
return requested;
}
if (!enable_prefix_cache) {
const int64_t live =
max_seqs_per_batch + kPaddingLinearStateBlocks;
return std::max<int64_t>(std::min(live, max_blocks),
kPaddingLinearStateBlocks);
}
// Auto-size with prefix cache (upstream ratio 0.9)
constexpr double kRatio = 0.9;
const double frac = kRatio / (1.0 + kRatio);
int64_t auto_blocks = std::max<int64_t>(
static_cast<int64_t>(cache_size_in_bytes * frac / linear_bytes_per_block),
kPaddingLinearStateBlocks);
return std::min(auto_blocks, max_blocks);
}
// Upstream: init_standard_counts (lines 483-550)
void init_standard_counts(const ModelArgs& model_args,
const KVCacheEstimateOptions& options,
KVCacheCapacity* kv_cache_cap) {
for (int64_t layer_id = 0; layer_id < kv_cache_cap->n_layers(); ++layer_id) {
if (is_full_attention_layer(model_args, layer_id)) {
++kv_cache_cap->num_full_attention_layers();
} else {
++kv_cache_cap->num_linear_attention_layers();
}
}
// [BI-V100 20%] No DSA indexer on BI-V100 — indexer_layers stays 0.
const int64_t full_cache_block_size_in_bytes =
standard_full_cache_block_size_in_bytes(*kv_cache_cap);
kv_cache_cap->num_linear_state_blocks(
calculate_linear_state_blocks(kv_cache_cap->cache_size_in_bytes(),
kv_cache_cap->num_linear_attention_layers(),
kv_cache_cap->linear_slot_size(),
full_cache_block_size_in_bytes,
options.max_seqs_per_batch,
options.max_linear_state_cache_slots,
options.enable_prefix_cache));
kv_cache_cap->linear_cache_size_in_bytes(
kv_cache_cap->num_linear_attention_layers() *
kv_cache_cap->num_linear_state_blocks() *
kv_cache_cap->linear_slot_size());
const int64_t available_full_cache_size_in_bytes =
kv_cache_cap->cache_size_in_bytes() -
kv_cache_cap->linear_cache_size_in_bytes();
if (kv_cache_cap->linear_slot_size() > 0) {
CHECK_GT(kv_cache_cap->cache_size_in_bytes(),
kv_cache_cap->linear_cache_size_in_bytes())
<< "failed to reserve linear state cache";
}
CHECK_GT(available_full_cache_size_in_bytes, 0)
<< "no memory left for full-attention kv cache";
if (options.layerwise_split_size > 1) {
kv_cache_cap->n_blocks(
layerwise_split_block_count(model_args,
options.layerwise_split_size,
*kv_cache_cap,
available_full_cache_size_in_bytes,
/*additional_block_bytes=*/0));
} else {
kv_cache_cap->n_blocks(available_full_cache_size_in_bytes /
full_cache_block_size_in_bytes);
}
CHECK_GT(kv_cache_cap->n_blocks(), 0) << "no n_blocks for kv cache";
}
} // namespace
// [BI-V100 20%] No DSA indexer on BI-V100 — returns empty mask.
std::vector<bool> resolve_indexer_cache_enabled_layers(
const ModelArgs& /*model_args*/,
int64_t /*num_cache_layers*/) {
return {};
}
// Upstream: build_layer_cache_owned (lines 563-574)
std::vector<bool> build_layer_cache_owned(const ModelArgs& model_args,
const LayerwiseSplitLayout& layout,
int64_t num_layers) {
std::vector<bool> layer_cache_owned;
layer_cache_owned.reserve(static_cast<size_t>(num_layers));
for (int64_t layer_id = 0; layer_id < num_layers; ++layer_id) {
layer_cache_owned.emplace_back(
!is_full_attention_layer(model_args, layer_id) ||
layout.owns(layer_id));
}
return layer_cache_owned;
}
// Upstream: estimate_layerwise_split_block_count (lines 576-587)
int64_t estimate_layerwise_split_block_count(
const ModelArgs& model_args,
int32_t layerwise_split_size,
const KVCacheCapacity& kv_cache_cap,
int64_t available_bytes,
int64_t additional_block_bytes) {
return layerwise_split_block_count(model_args,
layerwise_split_size,
kv_cache_cap,
available_bytes,
additional_block_bytes);
}
// Upstream: estimate_kv_cache_capacity (lines 589-641, stripped DSV4 path)
KVCacheCapacity estimate_kv_cache_capacity(
const ModelArgs& model_args,
const KVCacheEstimateOptions& options) {
KVCacheCapacity kv_cache_cap;
kv_cache_cap
.cache_size_in_bytes(
std::max(options.cache_size_in_bytes, static_cast<int64_t>(0)))
.block_size(options.block_size);
CHECK_GT(kv_cache_cap.cache_size_in_bytes(), 0)
<< "Available kv cache size must be greater than 0";
// [BI-V100 20%] dtype_size via enum, not torch
const int64_t dtype_size = dtype_size_from_enum(options.dtype_enum);
const int64_t cache_dtype_size =
kv_cache_dtype_size(options.kv_cache_dtype, dtype_size);
kv_cache_cap.slot_size(kv_slot_size(model_args, options, cache_dtype_size))
.scale_slot_size(scale_slot_size(model_args, options))
.linear_slot_size(linear_slot_size(model_args, options, dtype_size))
.n_layers(model_args.n_layers())
.block_size(options.block_size);
const int64_t num_speculative_tokens =
enable_qwen3_5_spec_verify(model_args, options)
? options.num_speculative_tokens
: 0;
kv_cache_cap.linear_conv_state_len(model_args.linear_conv_kernel_dim() - 1 +
num_speculative_tokens);
kv_cache_cap.linear_ssm_checkpoint_stride(num_speculative_tokens + 1);
init_standard_counts(model_args, options, &kv_cache_cap);
return kv_cache_cap;
}
} // namespace xllm

61
core/framework/kv_cache/kv_cache_estimation.h Normal file
View File

@@ -0,0 +1,61 @@
/* Migrated from xllm/core/framework/kv_cache/kv_cache_estimation.h
[BI-V100 20%] Replaced torch::ScalarType with int32_t dtype_enum
to avoid libtorch dependency in scheduling layer. */
#pragma once
#include <cstdint>
#include <string>
#include <vector>
#include "framework/kv_cache/kv_cache_capacity.h"
#include "framework/kv_cache/layerwise_split_layout.h"
namespace xllm {
class ModelArgs;
struct KVCacheEstimateOptions {
int32_t dtype_enum = 15; // bfloat16; upstream uses torch::ScalarType
std::string kv_cache_dtype = "auto";
std::string indexer_cache_dtype = "auto";
int64_t cache_size_in_bytes = 0;
int64_t block_size = 0;
int64_t world_size = 1;
int64_t n_local_kv_heads = 0;
int64_t n_local_linear_k_heads = 0;
int64_t n_local_linear_v_heads = 0;
int64_t max_seqs_per_batch = 0;
int64_t num_speculative_tokens = 0;
int64_t max_tokens_per_batch = 0;
int64_t max_linear_state_cache_slots = 0;
bool is_draft_engine = false;
bool enable_prefix_cache = false;
int32_t layerwise_split_size = 1;
const ModelArgs* draft_model_args = nullptr;
const KVCacheEstimateOptions* draft_options = nullptr;
};
std::vector<bool> resolve_indexer_cache_enabled_layers(
const ModelArgs& model_args,
int64_t num_cache_layers);
// Linear-attention layers stay owned on every rank. Full-attention layers
// follow LayerwiseSplitLayout.
std::vector<bool> build_layer_cache_owned(const ModelArgs& model_args,
const LayerwiseSplitLayout& layout,
int64_t num_layers);
// Common block count across a layerwise split group.
int64_t estimate_layerwise_split_block_count(
const ModelArgs& model_args,
int32_t layerwise_split_size,
const KVCacheCapacity& kv_cache_cap,
int64_t available_bytes,
int64_t additional_block_bytes);
KVCacheCapacity estimate_kv_cache_capacity(
const ModelArgs& model_args,
const KVCacheEstimateOptions& options);
} // namespace xllm

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

@@ -0,0 +1,183 @@
/* 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.
//
// BI-V100 verified: 32768 MiB HBM per card (ixsmi -q -d MEMORY),
// baseline usage ~257 MiB idle. Available for KV cache ≈ 32511 MiB
// minus model weights and activations.
//
// Qwen3.5 (Qwen3.6-35B-A3B) KV cache per full-attention layer per rank:
// key: n_blocks * 1 * 32 * 16 * 8 * 2 bytes (fp16) = n_blocks * 8192 B
// value: n_blocks * 1 * 256 * 16 * 2 bytes (fp16) = n_blocks * 8192 B
// total per layer per rank: n_blocks * 16384 bytes
//
// With layerwise split (group_size=2), each rank owns ~half the
// full-attention layers, so peak KV memory drops by ~50% minus
// scratch overhead.
#include "framework/kv_cache/kv_cache_estimation_layerwise.h"
#include "framework/kv_cache/layerwise_split_layout.h"
#include <glog/logging.h>
#include <algorithm>
#include <cstdint>
#include <limits>
#include <numeric>
#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 IluLayerwiseLayout& 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: pad head_dim to the next multiple of 64
// so each warp's contiguous load spans an aligned region.
// Qwen3.5 head_dim=256 → 256/64=4 warps, already aligned.
#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;
}
int64_t estimate_layerwise_split_block_count(
int32_t layerwise_split_size,
int64_t num_full_attn_layers,
int64_t per_layer_block_bytes,
int64_t scratch_block_bytes,
int64_t available_bytes) {
CHECK_GE(layerwise_split_size, 1);
CHECK_GT(num_full_attn_layers, 0);
CHECK_GT(per_layer_block_bytes, 0);
CHECK_GT(scratch_block_bytes, 0);
// Each split rank owns ceil(num_full_attn_layers / group_size) layers.
// The block count is limited by the rank with the most owned layers.
int64_t common_block_count = std::numeric_limits<int64_t>::max();
for (int32_t split_rank = 0; split_rank < layerwise_split_size;
++split_rank) {
const LayerwiseSplitLayout layout(
/*enabled=*/true, layerwise_split_size, split_rank);
int64_t owned_bytes = 0;
for (int64_t lid = 0; lid < num_full_attn_layers; ++lid) {
if (layout.owns(lid)) {
owned_bytes += per_layer_block_bytes;
}
}
if (owned_bytes == 0) {
continue;
}
const int64_t per_block_bytes = owned_bytes + scratch_block_bytes;
common_block_count =
std::min(common_block_count, available_bytes / per_block_bytes);
}
if (common_block_count == std::numeric_limits<int64_t>::max()) {
common_block_count = available_bytes / scratch_block_bytes;
}
LOG(INFO) << "[LayerwiseSplit] Block count estimate: "
<< common_block_count << " blocks (split_size="
<< layerwise_split_size << ", attn_layers="
<< num_full_attn_layers << ")";
return common_block_count;
}
} // namespace xllm

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

@@ -0,0 +1,62 @@
/* 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/ilu_layerwise_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.
///
/// This uses the IluLayerwiseLayout (BI-V100-specific detailed layout)
/// to compute exact per-rank memory accounting for the transposed
/// block-major tensor layout:
/// key: n_blocks * n_heads * block_size * head_dim * elem_bytes
/// value: n_blocks * n_heads * head_dim * block_size * elem_bytes
/// (both are equivalent in total bytes; the axes are just reordered)
///
/// |dtype_enum| matches torch::ScalarType integer values.
LayerwiseKVMemoryEstimate estimate_layerwise_kv_memory(
const IluLayerwiseLayout& layout,
int64_t n_blocks,
int64_t block_size,
int64_t head_dim,
int64_t max_tokens,
int dtype_enum,
int32_t world_size);
/// Estimate per-rank block count using the upstream-compatible
/// LayerwiseSplitLayout (round-robin ownership).
/// Returns the minimum block count that any rank can afford.
int64_t estimate_layerwise_split_block_count(
int32_t layerwise_split_size,
int64_t num_full_attn_layers,
int64_t per_layer_block_bytes,
int64_t scratch_block_bytes,
int64_t available_bytes);
} // namespace xllm

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

@@ -0,0 +1,96 @@
/* 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 the
// layer_cache_owned mask.
//
// On BI-V100 the KV cache tensor layout is transposed ("block-major"):
// key_cache: (n_blocks, n_kv_heads_local, block_size, head_dim/x, x)
// value_cache: (n_blocks, n_kv_heads_local, head_dim, block_size)
//
// Verified from runtime log with Qwen3.6-35B-A3B TP=4:
// key_cache=(68837, 1, 32, 16, 8) value_cache=(68837, 1, 256, 16)
// query=(1, 4, 256) → 4 q-heads, 1 kv-head per rank, head_dim=256
//
// This file implements the scheduling/allocation logic only. The actual
// tensor creation is done by the vendor vLLM cache engine (which calls
// into prebuilt corex_*.so for the kernel layer). This module decides
// WHICH layers get real allocations vs scratch placeholders.
//
// BI-V100 warp size = 64. head_dim (256 for Qwen3.5) is a multiple
// of 64, so coalesced warp-wide loads across the head dimension are aligned.
// When head_dim is not a multiple of 64, the allocator pads to the next
// multiple to avoid idle warp lanes.
#include "framework/kv_cache/kv_cache_layerwise.h"
#include <glog/logging.h>
#include <algorithm>
#include <cstdint>
#include <vector>
#include "config/ilu_hw_constants.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(
const std::vector<bool>& layer_cache_owned,
int64_t num_layers,
int32_t current_rank) {
CHECK_EQ(static_cast<int64_t>(layer_cache_owned.size()), num_layers)
<< "layer_cache_owned size must match num_layers.";
int64_t owned_count = 0;
int64_t scratch_count = 0;
for (int64_t i = 0; i < num_layers; ++i) {
if (layer_cache_owned[static_cast<size_t>(i)]) {
++owned_count;
} else {
++scratch_count;
}
}
// Scheduling decision log — the actual tensor allocation is done by
// the vendor vLLM cache engine (CacheEngine._allocate_kv_cache).
// This module only decides the ownership mask.
LOG(INFO) << "[LayerwiseSplit] rank " << current_rank << ": "
<< owned_count << "/" << num_layers
<< " layers owned, " << scratch_count << " scratch.";
// Validate head_dim alignment for BI-V100 warp size.
// Qwen3.5 head_dim=256, which is 256/64=4 warps — perfectly aligned.
constexpr int64_t kHeadDim = ilu_hw::kQwen35HeadDim;
constexpr int64_t kWarpAlign = ilu_hw::kWarpSize;
const int64_t padded = align_up(kHeadDim, kWarpAlign);
if (padded != kHeadDim) {
LOG(WARNING) << "[LayerwiseSplit] head_dim=" << kHeadDim
<< " not aligned to warp_size=" << kWarpAlign
<< ", padded to " << padded;
}
}
} // namespace xllm

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

@@ -0,0 +1,48 @@
/* 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 {
/// Allocate KV caches with per-layer ownership determined by |layout|.
/// Layers not owned by |current_rank| receive an empty (placeholder) KVCache.
///
/// |layer_cache_owned|: precomputed ownership mask from build_layer_cache_owned.
/// - true → this rank allocates real KV cache tensors for this layer.
/// - false → this rank uses a shared scratch (zero-sized) placeholder.
///
/// For Qwen3.5: linear_attention layers always have owned=true but their
/// cache slot is unused (GDN uses conv+temporal state, not KV cache).
/// Only full_attention layers participate in layerwise split.
///
/// BI-V100 KV cache tensor layout (from verified runtime logs):
/// key_cache: (n_blocks, n_kv_heads_local, block_size, head_dim/x, x)
/// value_cache: (n_blocks, n_kv_heads_local, head_dim, block_size)
/// where x = 128 / sizeof_dtype_bits (=8 for fp16/bf16)
///
/// This is the "block-major" transposed layout used by the vendor vLLM
/// on Iluvatar BI-V100. head dimension sits at axis 1.
void allocate_kv_caches_layerwise(
const std::vector<bool>& layer_cache_owned,
int64_t num_layers,
int32_t current_rank);
} // namespace xllm

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

@@ -0,0 +1,80 @@
/* 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://www.apache.org/licenses/LICENSE-2.0
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 <glog/logging.h>
#include <cstdint>
#include <string>
namespace xllm {
// [BI-V100 20%] Added qwen3_5 / qwen3_5_moe_text to supported models.
[[nodiscard]] inline bool is_layerwise_split_supported_model(
const std::string& model_type) noexcept {
return model_type == "deepseek_v32" || model_type == "glm_moe_dsa" ||
model_type == "qwen3_5" || model_type == "qwen3_5_moe_text";
}
inline void validate_layerwise_split_size_config(int32_t layerwise_split_size) {
CHECK_GE(layerwise_split_size, 1)
<< "layerwise_split_size must be >= 1, got " << layerwise_split_size;
}
inline void validate_layerwise_split_enablement(int32_t layerwise_split_size,
int32_t attn_tp_size,
const std::string& model_type) {
if (layerwise_split_size <= 1) {
return;
}
CHECK(is_layerwise_split_supported_model(model_type))
<< "layerwise_split_size > 1 is only supported for deepseek_v32, "
"glm_moe_dsa, qwen3_5, and qwen3_5_moe_text, got "
<< model_type;
CHECK_EQ(attn_tp_size % layerwise_split_size, 0)
<< "attention tp size (" << attn_tp_size
<< ") must be divisible by layerwise_split_size (" << layerwise_split_size
<< ").";
}
class LayerwiseSplitLayout final {
public:
LayerwiseSplitLayout(bool enabled, int32_t group_size, int32_t local_rank)
: enabled_(enabled), group_size_(group_size), local_rank_(local_rank) {
CHECK_GT(group_size_, 0) << "Layerwise split group size must be positive.";
CHECK(local_rank_ >= 0 && local_rank_ < group_size_)
<< "Layerwise split local rank must be in [0, group_size).";
}
[[nodiscard]] int32_t owner_rank(int64_t layer_id) const {
return static_cast<int32_t>(layer_id % group_size_);
}
[[nodiscard]] bool owns(int64_t layer_id) const {
return !enabled_ || owner_rank(layer_id) == local_rank_;
}
[[nodiscard]] bool enabled() const { return enabled_; }
[[nodiscard]] int32_t group_size() const { return group_size_; }
[[nodiscard]] int32_t local_rank() const { return local_rank_; }
private:
bool enabled_ = false;
int32_t group_size_ = 1;
int32_t local_rank_ = 0;
};
} // namespace xllm

95
core/framework/model/model_args.h Normal file
View File

@@ -0,0 +1,95 @@
/* Minimal ModelArgs for project_6 layerwise split.
Migrated from xllm/core/framework/model/model_args.h — only the fields
and functions used by kv_cache_estimation.cpp's layerwise_split code path.
Full ModelArgs in upstream is 861 lines; we keep only the interface that
the layerwise split scheduling layer needs. */
#pragma once
#include <algorithm>
#include <cstdint>
#include <optional>
#include <string>
#include <string_view>
#include <vector>
#include "common/macros.h"
namespace xllm {
class ModelArgs {
PROPERTY(std::string, model_type);
PROPERTY(int64_t, n_layers) = 0;
PROPERTY(int64_t, head_dim) = 0;
PROPERTY(int64_t, n_heads) = 0;
PROPERTY(std::optional<int64_t>, n_kv_heads);
PROPERTY(int64_t, hidden_size) = 0;
// Linear attention (GDN/DeltaNet) fields — Qwen3.5
PROPERTY(int64_t, full_attention_interval) = 1;
PROPERTY(std::vector<std::string>, layer_types);
PROPERTY(int64_t, linear_num_value_heads) = 0;
PROPERTY(int64_t, linear_key_head_dim) = 0;
PROPERTY(int64_t, linear_value_head_dim) = 0;
PROPERTY(int64_t, linear_conv_kernel_dim) = 0;
PROPERTY(std::string, mamba_ssm_dtype);
// MLA (DeepSeek)
PROPERTY(bool, enable_mla) = false;
PROPERTY(int64_t, kv_lora_rank) = 0;
PROPERTY(int64_t, qk_rope_head_dim) = 0;
// DSA indexer
PROPERTY(int64_t, index_n_heads) = 0;
PROPERTY(int64_t, index_head_dim) = 0;
// Sliding window
PROPERTY(int64_t, window_size) = 0;
PROPERTY(int64_t, max_seq_len) = 0;
// DeepSeek V4 compress ratios
PROPERTY(std::vector<int32_t>, compress_ratios);
// Speculative
PROPERTY(int64_t, num_nextn_predict_layers) = 0;
};
// Qwen hybrid models may describe full-attention layers explicitly via
// layer_types or implicitly via full_attention_interval.
inline bool is_full_attention_layer(const ModelArgs& args, int64_t layer_id) {
const auto& hybrid_layer_types = args.layer_types();
if (layer_id >= 0 &&
layer_id < static_cast<int64_t>(hybrid_layer_types.size())) {
const auto& layer_type = hybrid_layer_types[layer_id];
return layer_type == "full_attention" || layer_type == "attention";
}
int32_t attention_interval = args.full_attention_interval();
if (attention_interval <= 1) {
return true;
}
return (layer_id + 1) % attention_interval == 0;
}
inline bool has_linear_attention_layers(const ModelArgs& args) {
const auto& hybrid_layer_types = args.layer_types();
if (!hybrid_layer_types.empty()) {
return std::any_of(hybrid_layer_types.begin(),
hybrid_layer_types.end(),
[](const std::string& layer_type) {
return layer_type != "full_attention" &&
layer_type != "attention";
});
}
return args.full_attention_interval() > 1;
}
// Closed set by design: a new target variant must be enumerated here rather
// than matched by a "qwen3_5_" prefix.
inline bool is_qwen3_5_target_model_type(std::string_view model_type) {
return model_type == "qwen3_5" || model_type == "qwen3_5_moe" ||
model_type == "qwen3_5_text" || model_type == "qwen3_5_moe_text";
}
} // namespace xllm

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

@@ -0,0 +1,110 @@
/* 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: GPU0-GPU3, Bus-Id 4B:00.0 – 4E:00.0, all on NUMA node 1
// - All pairs connected via PIX (single PCIe bridge) — FLAT topology
// - No NVLink, no HCCS mesh, no multi-switch hierarchy
// - 32 GB HBM per card (32768 MiB), baseline ~257 MiB
// - 1500 MHz SM, 1200 MHz mem
// - CPU affinity: 16-31,80-95
// - Warp size: 64 (verified via BI-V150 docs, same ivcore architecture)
// - IX-ML 3.2.3, Driver 3.2.1, CUDA 10.2 (CoreX)
// - CoreX SDK at /usr/local/corex/
//
// Model: Qwen3.6-35B-A3B (Qwen3_5 MoE architecture)
// - Interleaved full_attention + linear_attention layers
// - Only full_attention layers have KV cache
// - num_kv_heads=4, with TP=4: local_kv_heads=1 per rank per layer
// - head_dim=256, GQA ratio=4
//
// Strategy (flat PIX topology):
// All full-attention layers have kv_heads=4 >= world_size=4,
// so every layer shards across ALL ranks (each gets 1 head).
// If kv_heads < world_size, round-robin starting rank to spread
// HBM load (no grouping benefit since all PIX 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/ilu_layerwise_layout.h"
#include "framework/kv_cache/layerwise_split_layout.h"
namespace xllm {
IluLayerwiseLayout 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 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 case (Qwen3.5 full_attention with TP=4, 4 heads → 1 each):
// 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 {
// Fewer heads than ranks: round-robin subset.
// (Not the case for Qwen3.5 with TP=4, but needed for other models.)
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 IluLayerwiseLayout(std::move(specs));
}
} // namespace xllm

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

@@ -0,0 +1,57 @@
/* 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/ilu_layerwise_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 IluLayerwiseLayout for Iluvatar BI-V100.
///
/// |per_layer_kv_heads|: total KV head count for each full-attention layer.
/// For Qwen3.5 with TP=4: each full_attention layer has 4 total KV heads,
/// so local_kv_heads = 4/4 = 1 per rank.
/// Linear-attention layers are NOT included in this vector — they have no
/// KV cache.
///
/// Dense attention layers (heads >= world_size) spread across ALL TP ranks.
/// Layers with fewer heads than ranks are round-robin distributed.
///
/// Default: kFlatPIX — matches the verified 4-card BI-V100 topology
/// where all pairs are PIX-connected with equal bandwidth.
IluLayerwiseLayout 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

29
core/framework/parallel_state/parallel_args.h Normal file
View File

@@ -0,0 +1,29 @@
/* Migrated from xllm/core/framework/parallel_state/parallel_args.h (251 lines)
Stripped ProcessGroup, mapping_data, vendor comm — scheduling layer only. */
#pragma once
#include <cstdint>
#include "common/macros.h"
namespace xllm {
class ParallelArgs {
public:
ParallelArgs() = default;
ParallelArgs(int32_t rank, int32_t world_size)
: rank_(rank), world_size_(world_size) {}
ParallelArgs(int32_t rank, int32_t world_size, int32_t dp_size,
int32_t cp_size, int32_t layerwise_split_size)
: rank_(rank), world_size_(world_size), dp_size_(dp_size),
cp_size_(cp_size), layerwise_split_size_(layerwise_split_size) {}
PROPERTY(int32_t, rank) = 0;
PROPERTY(int32_t, world_size) = 1;
PROPERTY(int32_t, dp_size) = 1;
PROPERTY(int32_t, cp_size) = 1;
PROPERTY(int32_t, layerwise_split_size) = 1;
};
} // namespace xllm

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

@@ -0,0 +1,348 @@
/* 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
PYBIND11_MODULE(py_attention_metadata, m) {
m.doc() = "DP-aware attention metadata (project6, ported from xLLM PR #2258)";
project6::register_attention_metadata_views(m);
}

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

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

@@ -0,0 +1,58 @@
/* Worker-side helper: compute layer_cache_owned mask.
Migrated pattern from upstream worker_impl.cpp lines 490-498. */
#include "runtime/worker_layerwise_init.h"
#include <glog/logging.h>
#include <string>
#include <vector>
#include "framework/kv_cache/kv_cache_estimation.h"
#include "framework/kv_cache/layerwise_split_layout.h"
#include "framework/model/model_args.h"
namespace xllm {
std::vector<bool> worker_compute_layer_cache_owned(
const std::vector<std::string>& layer_types,
int32_t layerwise_split_size,
int32_t rank,
int64_t num_layers) {
CHECK_GE(layerwise_split_size, 1);
CHECK_GE(rank, 0);
CHECK_GT(num_layers, 0);
if (layerwise_split_size <= 1) {
LOG(INFO) << "[Worker " << rank << "] No layerwise split, all "
<< num_layers << " layers owned.";
return std::vector<bool>(static_cast<size_t>(num_layers), true);
}
// Build a ModelArgs with the layer_types for upstream build_layer_cache_owned
ModelArgs args;
args.n_layers(num_layers);
args.layer_types(layer_types);
const LayerwiseSplitLayout layout(
/*enabled=*/true,
layerwise_split_size,
rank % layerwise_split_size);
auto owned = build_layer_cache_owned(args, layout, num_layers);
int64_t owned_count = 0;
int64_t linear_count = 0;
for (int64_t i = 0; i < num_layers; ++i) {
if (owned[static_cast<size_t>(i)]) ++owned_count;
if (!is_full_attention_layer(args, i)) ++linear_count;
}
LOG(INFO) << "[Worker " << rank << "] Layerwise split (size="
<< layerwise_split_size << "): " << owned_count << "/"
<< num_layers << " layers owned, " << linear_count
<< " linear-attention (always owned).";
return owned;
}
} // namespace xllm

40
core/runtime/worker_layerwise_init.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 <string>
#include <vector>
#include "framework/kv_cache/layerwise_split_layout.h"
namespace xllm {
/// Worker-side entry: compute the layer_cache_owned mask and validate.
///
/// For Qwen3.5: the worker receives layerwise_split_size from the master
/// and builds its ownership mask. Only full_attention layers participate
/// in the split; linear_attention layers are always "owned" (no KV cache).
///
/// Returns the layer_cache_owned mask, which the vendor vLLM cache engine
/// uses to decide which layers get real KV allocations vs scratch.
std::vector<bool> worker_compute_layer_cache_owned(
const std::vector<std::string>& layer_types,
int32_t layerwise_split_size,
int32_t rank,
int64_t num_layers);
} // namespace xllm