完整部署快照

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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