完整部署快照

core/framework/kv_cache/ilu_layerwise_layout.h

@@ -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

core/framework/kv_cache/kv_cache_capacity.h

@@ -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

core/framework/kv_cache/kv_cache_estimation.cpp

@@ -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

core/framework/kv_cache/kv_cache_estimation.h

@@ -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

core/framework/kv_cache/kv_cache_estimation_layerwise.cpp

@@ -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

core/framework/kv_cache/kv_cache_estimation_layerwise.h

@@ -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

core/framework/kv_cache/kv_cache_layerwise.cpp

@@ -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

core/framework/kv_cache/kv_cache_layerwise.h

@@ -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

core/framework/kv_cache/layerwise_split_layout.h

@@ -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