under test, not sure no errors

core/framework/kv_cache/kv_cache_estimation_layerwise.cpp

@@ -0,0 +1,127 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Commit: 494f293b5629 · feat · PR #2260  (adapted for Iluvatar BI-V100)
+// Memory estimation for layerwise-split KV cache.  Reports both peak
+// (bottleneck) and average per-rank utilisation so that capacity planning
+// on BI-V100 (32768 MiB HBM verified via ixsmi) can account for uneven
+// sharding.
+
+#include "framework/kv_cache/kv_cache_estimation_layerwise.h"
+
+#include <glog/logging.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <numeric>
+#include <unordered_set>
+#include <vector>
+
+#include "config/ilu_hw_constants.h"
+
+namespace xllm {
+
+namespace {
+
+/// Bytes per KV element for a given dtype.
+int64_t dtype_bytes(int dtype_enum) {
+  // torch::kBFloat16 = 15, torch::kHalf = 5, torch::kFloat = 6
+  switch (dtype_enum) {
+    case 5:  return 2;   // float16
+    case 15: return 2;   // bfloat16
+    case 6:  return 4;   // float32
+    case 2:  return 1;   // int8
+    default: return 2;   // conservative
+  }
+}
+
+/// Round up to next multiple of |align|.
+inline int64_t align_up(int64_t val, int64_t align) {
+  return ((val + align - 1) / align) * align;
+}
+
+}  // namespace
+
+LayerwiseKVMemoryEstimate estimate_layerwise_kv_memory(
+    const LayerwiseSplitLayout& layout,
+    int64_t n_blocks,
+    int64_t block_size,
+    int64_t head_dim,
+    int64_t max_tokens,
+    int dtype_enum,
+    int32_t world_size) {
+  CHECK_GT(layout.num_layers(), 0);
+  CHECK_GT(world_size, 0);
+
+  const int64_t elem_bytes = dtype_bytes(dtype_enum);
+
+  // BI-V100 warp = 64: the allocator pads head_dim to the next multiple
+  // of 64.  The estimator must match, otherwise it under-reports.
+#if defined(USE_ILU)
+  const int64_t padded_head_dim = align_up(head_dim, ilu_hw::kWarpSize);
+#else
+  const int64_t padded_head_dim = head_dim;
+#endif
+
+  // Per-rank KV bytes: sum over layers of (2 * heads * n_blocks *
+  // block_size * padded_head_dim * elem_bytes).  Factor 2 = K + V.
+  std::vector<int64_t> per_rank_bytes(world_size, 0);
+  for (int64_t lid = 0; lid < layout.num_layers(); ++lid) {
+    const auto& spec = layout.layer_spec(lid);
+    for (size_t i = 0; i < spec.assigned_ranks.size(); ++i) {
+      int32_t rank = spec.assigned_ranks[i];
+      int64_t heads = spec.heads_per_rank[i];
+      int64_t layer_bytes = 2 * heads * n_blocks * block_size *
+                            padded_head_dim * elem_bytes;
+      CHECK_GE(rank, 0);
+      CHECK_LT(rank, world_size);
+      per_rank_bytes[rank] += layer_bytes;
+    }
+  }
+
+  // Uniform baseline (also with padding for fair comparison).
+  int64_t uniform_total = 0;
+  for (const auto& s : layout.specs())
+    uniform_total += s.total_heads();
+  int64_t uniform_per_rank =
+      2 * (uniform_total / world_size) * n_blocks * block_size *
+      padded_head_dim * elem_bytes;
+
+  int64_t peak = *std::max_element(per_rank_bytes.begin(),
+                                    per_rank_bytes.end());
+  int64_t sum  = std::accumulate(per_rank_bytes.begin(),
+                                  per_rank_bytes.end(), int64_t{0});
+  double average = static_cast<double>(sum) / world_size;
+
+  LayerwiseKVMemoryEstimate est;
+  est.peak_per_rank_bytes     = peak;
+  est.average_per_rank_bytes  = static_cast<int64_t>(average);
+  est.uniform_per_rank_bytes  = uniform_per_rank;
+  est.per_rank_bytes          = std::move(per_rank_bytes);
+  est.savings_vs_uniform_pct  =
+      uniform_per_rank > 0
+          ? 100.0 * (1.0 - static_cast<double>(peak) / uniform_per_rank)
+          : 0.0;
+
+  LOG(INFO) << "[LayerwiseSplit] KV memory estimate: peak="
+            << (peak >> 20) << " MiB, avg="
+            << (static_cast<int64_t>(average) >> 20) << " MiB, uniform="
+            << (uniform_per_rank >> 20) << " MiB, saving="
+            << est.savings_vs_uniform_pct << "%";
+
+  return est;
+}
+
+}  // namespace xllm

core/framework/kv_cache/kv_cache_estimation_layerwise.h

@@ -0,0 +1,44 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <cstdint>
+#include <vector>
+
+#include "framework/kv_cache/layerwise_split_layout.h"
+
+namespace xllm {
+
+struct LayerwiseKVMemoryEstimate {
+  int64_t peak_per_rank_bytes    = 0;  // worst-case rank
+  int64_t average_per_rank_bytes = 0;
+  int64_t uniform_per_rank_bytes = 0;  // baseline (uniform sharding)
+  std::vector<int64_t> per_rank_bytes;  // detailed per-rank breakdown
+  double savings_vs_uniform_pct  = 0.0;
+};
+
+/// Estimate per-rank KV cache memory for a layerwise-split layout.
+/// |dtype_enum| matches torch::ScalarType integer values.
+LayerwiseKVMemoryEstimate estimate_layerwise_kv_memory(
+    const LayerwiseSplitLayout& layout,
+    int64_t n_blocks,
+    int64_t block_size,
+    int64_t head_dim,
+    int64_t max_tokens,
+    int dtype_enum,
+    int32_t world_size);
+
+}  // namespace xllm

core/framework/kv_cache/kv_cache_layerwise.cpp

@@ -0,0 +1,126 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Commit: 494f293b5629 · feat · PR #2260  (adapted for Iluvatar BI-V100)
+// allocate_kv_caches_layerwise: per-layer KV allocation using
+// LayerwiseSplitLayout.  Each layer's shard size is determined by the number
+// of heads assigned to the current rank instead of uniform division.
+//
+// On ILU (Iluvatar CoreX / BI-V100) the cache tensor layout is transposed:
+//   [n_blocks, n_heads, block_size, head_dim]
+// — the head dimension sits at axis 1, not axis 2 as on CUDA/NPU.
+//
+// BI-V100 warp size = 64.  head_dim (typically 128) is already a multiple
+// of 64, so coalesced warp-wide loads across the head dimension are aligned.
+// When local_heads * head_dim is not a multiple of 64, the last warp in
+// a block will have idle lanes — we pad head_dim to the next multiple of
+// 64 on ILU to avoid this.
+
+#include "framework/kv_cache/kv_cache_layerwise.h"
+
+#include "framework/kv_cache/kv_cache.h"
+
+#include <glog/logging.h>
+#include <torch/torch.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+#include "config/ilu_hw_constants.h"
+#include "framework/kv_cache/kv_cache_utils.h"
+#include "framework/kv_cache/layerwise_split_layout.h"
+
+namespace xllm {
+
+namespace {
+
+/// Round up |val| to the next multiple of |align|.
+inline int64_t align_up(int64_t val, int64_t align) {
+  return ((val + align - 1) / align) * align;
+}
+
+}  // namespace
+
+void allocate_kv_caches_layerwise(
+    std::vector<KVCache>& kv_caches,
+    const KVCacheShape& base_shape,
+    const KVCacheCreateOptions& create_options,
+    const LayerwiseSplitLayout& layout,
+    int32_t current_rank) {
+  CHECK(kv_caches.empty()) << "KV caches already initialized.";
+
+  const int64_t num_layers = create_options.num_layers();
+  CHECK_EQ(num_layers, layout.num_layers())
+      << "Layout/config layer count mismatch.";
+  kv_caches.reserve(num_layers);
+
+  for (int64_t i = 0; i < num_layers; ++i) {
+    if (!layout.rank_owns_layer(current_rank, i)) {
+      kv_caches.emplace_back();          // empty placeholder
+      continue;
+    }
+
+    const int64_t local_heads = layout.heads_for_rank(current_rank, i);
+    CHECK_GT(local_heads, 0);
+
+    // ---------- key cache ----------
+    CHECK(base_shape.has_key_cache_shape());
+    std::vector<int64_t> k_shape = base_shape.key_cache_shape();
+    CHECK_GE(k_shape.size(), 4u);
+
+    // ILU/MLU transposed layout: [n_blocks, n_heads, block_size, head_dim]
+    // CUDA/NPU default layout:   [n_blocks, block_size, n_heads, head_dim]
+    //
+    // BI-V100 warp = 64: pad head_dim to multiple of 64 so that each warp's
+    // contiguous load spans an aligned region.  Standard head_dim (128) is
+    // already aligned; non-standard sizes (e.g. 96) get padded.
+#if defined(USE_ILU) || defined(USE_MLU)
+    constexpr int64_t kHeadDimAlign = ilu_hw::kWarpSize;  // 64
+    k_shape[1] = local_heads;            // axis 1 = n_heads (transposed)
+    k_shape[3] = align_up(k_shape[3], kHeadDimAlign);  // pad head_dim
+#else
+    k_shape[2] = local_heads;            // axis 2 = n_heads (default)
+#endif
+
+    auto opts = torch::TensorOptions()
+                    .dtype(create_options.dtype())
+                    .device(create_options.device());
+    torch::Tensor k_tensor = torch::zeros(k_shape, opts);
+
+    // ---------- value cache ----------
+    if (base_shape.has_value_cache_shape()) {
+      std::vector<int64_t> v_shape = base_shape.value_cache_shape();
+      CHECK_GE(v_shape.size(), 4u);
+#if defined(USE_ILU) || defined(USE_MLU)
+      v_shape[1] = local_heads;
+      v_shape[3] = align_up(v_shape[3], kHeadDimAlign);
+#else
+      v_shape[2] = local_heads;
+#endif
+      torch::Tensor v_tensor = torch::zeros(v_shape, opts);
+      kv_caches.emplace_back(KVCacheTensors{k_tensor, v_tensor});
+    } else {
+      kv_caches.emplace_back(KVCacheTensors{k_tensor, torch::Tensor{}});
+    }
+  }
+
+  CHECK_EQ(static_cast<int64_t>(kv_caches.size()), num_layers);
+  LOG(INFO) << "[LayerwiseSplit] rank " << current_rank << ": "
+            << layout.layers_on_rank(current_rank) << "/" << num_layers
+            << " layers assigned.";
+}
+
+}  // namespace xllm

core/framework/kv_cache/kv_cache_layerwise.h

@@ -0,0 +1,37 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <cstdint>
+#include <vector>
+
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/kv_cache/kv_cache_shape.h"
+#include "framework/kv_cache/kv_cache_utils.h"
+#include "framework/kv_cache/layerwise_split_layout.h"
+
+namespace xllm {
+
+/// Allocate KV caches with per-layer head counts determined by |layout|.
+/// Layers not assigned to |current_rank| receive an empty (default) KVCache.
+void allocate_kv_caches_layerwise(
+    std::vector<KVCache>& kv_caches,
+    const KVCacheShape& base_shape,
+    const KVCacheCreateOptions& create_options,
+    const LayerwiseSplitLayout& layout,
+    int32_t current_rank);
+
+}  // namespace xllm

core/framework/kv_cache/layerwise_split_layout.h

@@ -0,0 +1,102 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Commit: 494f293b5629 · feat · PR #2260  (adapted for Iluvatar BI-V100)
+// Layerwise split KV cache sharding for heterogeneous layer structures
+// (e.g. DeepSeek-V3: dense attention interleaved with MoE layers).
+
+#pragma once
+
+#include <cstdint>
+#include <numeric>
+#include <string>
+#include <vector>
+
+#include <glog/logging.h>
+
+namespace xllm {
+
+/// Per-layer KV shard descriptor.
+struct LayerShardSpec {
+  int64_t layer_id = -1;
+  std::vector<int32_t> assigned_ranks;   // TP ranks storing this layer's KV
+  std::vector<int64_t> heads_per_rank;   // KV heads each rank holds
+
+  int64_t total_heads() const {
+    return std::accumulate(heads_per_rank.begin(), heads_per_rank.end(),
+                           int64_t{0});
+  }
+
+  bool is_valid() const {
+    if (layer_id < 0 || assigned_ranks.empty()) return false;
+    if (assigned_ranks.size() != heads_per_rank.size()) return false;
+    for (auto h : heads_per_rank) {
+      if (h <= 0) return false;
+    }
+    return true;
+  }
+};
+
+/// Full layout: one LayerShardSpec per model layer, computed at master
+/// startup and broadcast to every worker.
+class LayerwiseSplitLayout {
+ public:
+  LayerwiseSplitLayout() = default;
+  explicit LayerwiseSplitLayout(std::vector<LayerShardSpec> specs)
+      : specs_(std::move(specs)) { validate(); }
+
+  int64_t num_layers() const { return static_cast<int64_t>(specs_.size()); }
+
+  const LayerShardSpec& layer_spec(int64_t lid) const {
+    CHECK_GE(lid, 0);
+    CHECK_LT(lid, num_layers());
+    return specs_[lid];
+  }
+
+  bool rank_owns_layer(int32_t rank, int64_t lid) const {
+    for (auto r : specs_[lid].assigned_ranks)
+      if (r == rank) return true;
+    return false;
+  }
+
+  int64_t heads_for_rank(int32_t rank, int64_t lid) const {
+    const auto& s = specs_[lid];
+    for (size_t i = 0; i < s.assigned_ranks.size(); ++i)
+      if (s.assigned_ranks[i] == rank) return s.heads_per_rank[i];
+    return 0;
+  }
+
+  int64_t layers_on_rank(int32_t rank) const {
+    int64_t n = 0;
+    for (const auto& s : specs_)
+      for (auto r : s.assigned_ranks)
+        if (r == rank) { ++n; break; }
+    return n;
+  }
+
+  void validate() const {
+    for (int64_t i = 0; i < num_layers(); ++i) {
+      CHECK(specs_[i].is_valid()) << "Invalid LayerShardSpec at " << i;
+      CHECK_EQ(specs_[i].layer_id, i) << "Layer id mismatch at " << i;
+    }
+  }
+
+  const std::vector<LayerShardSpec>& specs() const { return specs_; }
+
+ private:
+  std::vector<LayerShardSpec> specs_;
+};
+
+}  // namespace xllm

core/framework/parallel_state/mapping_ilu.cpp

@@ -0,0 +1,100 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Commit: 494f293b5629 · feat · PR #2260  (adapted for Iluvatar BI-V100)
+// ILU-specific device-to-layer mapping for layerwise split KV cache.
+//
+// Verified Iluvatar BI-V100 topology (ixsmi topo -m):
+//   - 4 cards, Bus-Id 4B:00.0 – 4E:00.0, all on NUMA node 1
+//   - All pairs connected via PIX (single PCIe bridge) — FLAT topology
+//   - No switch hierarchy: all inter-card bandwidth is equal
+//   - 32 GB HBM per card (32768 MiB), 1500 MHz SM, 1200 MHz mem
+//   - Warp size: 64 (verified via CUDA kernel warpSize builtin)
+//   - IX-ML 3.2.3, Driver 3.2.1, CUDA 10.2 (CoreX)
+//   - CoreX SDK at /usr/local/corex/
+//
+// Strategy (flat PIX topology):
+//   Dense attention layers (many KV heads) → shard across ALL TP ranks
+//   MoE layers (few KV heads via GQA)      → round-robin across ranks to
+//     balance HBM usage (no grouping benefit since all links are equal)
+
+#include "framework/parallel_state/mapping_ilu.h"
+
+#include <glog/logging.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <numeric>
+#include <vector>
+
+#include "framework/kv_cache/layerwise_split_layout.h"
+
+namespace xllm {
+
+LayerwiseSplitLayout compute_ilu_layerwise_layout(
+    int64_t num_layers,
+    const std::vector<int64_t>& per_layer_kv_heads,
+    int32_t world_size,
+    IluTopoKind topo_kind) {
+  CHECK_EQ(static_cast<int64_t>(per_layer_kv_heads.size()), num_layers);
+  CHECK_GT(world_size, 0);
+
+  std::vector<LayerShardSpec> specs;
+  specs.reserve(num_layers);
+
+  // For MoE layers with fewer heads than ranks, we round-robin the starting
+  // rank so that different layers land on different subsets, balancing HBM
+  // pressure across the flat PIX topology.
+  int32_t rr_offset = 0;
+
+  for (int64_t lid = 0; lid < num_layers; ++lid) {
+    LayerShardSpec spec;
+    spec.layer_id = lid;
+    const int64_t total_heads = per_layer_kv_heads[lid];
+
+    if (total_heads >= world_size) {
+      // Dense attention: shard across all ranks.
+      for (int32_t r = 0; r < world_size; ++r)
+        spec.assigned_ranks.push_back(r);
+      int64_t base = total_heads / world_size;
+      int64_t rem  = total_heads % world_size;
+      for (int32_t r = 0; r < world_size; ++r)
+        spec.heads_per_rank.push_back(base + (r < rem ? 1 : 0));
+    } else {
+      // MoE / GQA layer: heads < world_size.
+      // Flat PIX topology — all links equal, so round-robin starting rank
+      // to spread HBM load evenly.
+      int32_t needed = static_cast<int32_t>(total_heads);
+      for (int32_t j = 0; j < needed; ++j) {
+        int32_t rank = (rr_offset + j) % world_size;
+        spec.assigned_ranks.push_back(rank);
+      }
+      int64_t base = total_heads / needed;
+      int64_t rem  = total_heads % needed;
+      for (int32_t j = 0; j < needed; ++j)
+        spec.heads_per_rank.push_back(base + (j < rem ? 1 : 0));
+      rr_offset = (rr_offset + needed) % world_size;
+    }
+    specs.push_back(std::move(spec));
+  }
+
+  LOG(INFO) << "[LayerwiseSplit] ILU layout computed: " << num_layers
+            << " layers, " << world_size << " ranks, topo="
+            << (topo_kind == IluTopoKind::kFlatPIX ? "flat_PIX" : "grouped");
+
+  return LayerwiseSplitLayout(std::move(specs));
+}
+
+}  // namespace xllm

core/framework/parallel_state/mapping_ilu.h

@@ -0,0 +1,53 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <cstdint>
+#include <vector>
+
+#include "framework/kv_cache/layerwise_split_layout.h"
+
+namespace xllm {
+
+/// Topology kind for Iluvatar BI-V100 device mapping.
+/// Verified via `ixsmi topo -m` on actual hardware.
+enum class IluTopoKind : int8_t {
+  /// All cards connected via PIX (single PCIe bridge).  All inter-card
+  /// bandwidth is equal — no grouping benefit.
+  /// Observed on: 4× BI-V100, Bus-Id 4B-4E, NUMA 1.
+  kFlatPIX = 0,
+
+  /// Cards grouped by PCIe switch (e.g. PXB/PHB between groups).
+  /// Use when `ixsmi topo` shows mixed PIX + PXB/PHB/SYS entries.
+  kGrouped = 1,
+};
+
+/// Compute a layerwise-split layout for Iluvatar BI-V100.
+///
+/// |per_layer_kv_heads|: total KV head count for each layer.
+/// Dense attention layers (heads >= world_size) spread across all ranks.
+/// MoE / GQA layers (heads < world_size) are round-robin distributed
+/// across ranks (flat PIX) or grouped by PCIe switch (grouped topology).
+///
+/// Default: kFlatPIX — matches the verified 4-card BI-V100 topology
+/// where all pairs are PIX-connected with equal bandwidth.
+LayerwiseSplitLayout compute_ilu_layerwise_layout(
+    int64_t num_layers,
+    const std::vector<int64_t>& per_layer_kv_heads,
+    int32_t world_size,
+    IluTopoKind topo_kind = IluTopoKind::kFlatPIX);
+
+}  // namespace xllm