完整部署快照
This commit is contained in:
42
core/framework/config/parallel_config.cpp
Normal file
42
core/framework/config/parallel_config.cpp
Normal 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
34
core/framework/config/parallel_config.h
Normal 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
84
core/framework/kv_cache/ilu_layerwise_layout.h
Normal 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
44
core/framework/kv_cache/kv_cache_capacity.h
Normal 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
370
core/framework/kv_cache/kv_cache_estimation.cpp
Normal 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
61
core/framework/kv_cache/kv_cache_estimation.h
Normal 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
183
core/framework/kv_cache/kv_cache_estimation_layerwise.cpp
Normal 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
62
core/framework/kv_cache/kv_cache_estimation_layerwise.h
Normal 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
96
core/framework/kv_cache/kv_cache_layerwise.cpp
Normal 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
48
core/framework/kv_cache/kv_cache_layerwise.h
Normal 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
80
core/framework/kv_cache/layerwise_split_layout.h
Normal 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
95
core/framework/model/model_args.h
Normal 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
110
core/framework/parallel_state/mapping_ilu.cpp
Normal 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
57
core/framework/parallel_state/mapping_ilu.h
Normal 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
29
core/framework/parallel_state/parallel_args.h
Normal 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
|
||||
Reference in New Issue
Block a user