完整部署快照

core/distributed_runtime/layerwise_split_engine_ext.cpp

@@ -0,0 +1,63 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Commit: 494f293b5629 · feat · PR #2260  (adapted for Iluvatar BI-V100)
+// Engine-level plumbing:  Both llm_engine and speculative_engine propagate
+// the layerwise layout to workers during initialisation.
+//
+// In the upstream xLLM, this would be edits to llm_engine.cpp (+18 lines)
+// and speculative_engine.cpp (+12 lines).  Here we isolate them in a
+// self-contained compilation unit that the engines call into.
+//
+// The vendor vLLM on BI-V100 uses prebuilt CoreX .so extensions
+// (corex_paged_kv_gather.so, etc.) for the kernel layer.  This module
+// operates at the scheduling layer above — deciding which layers each
+// rank owns for KV cache purposes.
+
+#include "distributed_runtime/layerwise_split_engine_ext.h"
+
+#include <glog/logging.h>
+
+#include <memory>
+#include <optional>
+#include <vector>
+
+#include "framework/kv_cache/layerwise_split_layout.h"
+#include "framework/parallel_state/mapping_ilu.h"
+
+// The flag is declared in parallel_config_layerwise.h / .cpp (sub-task 7).
+DECLARE_bool(enable_layerwise_split);
+
+namespace xllm {
+
+std::optional<IluLayerwiseLayout> maybe_compute_layerwise_layout(
+    int64_t num_layers,
+    const std::vector<int64_t>& per_layer_kv_heads,
+    int32_t world_size) {
+  if (!FLAGS_enable_layerwise_split) {
+    return std::nullopt;
+  }
+
+  LOG(INFO) << "[LayerwiseSplit] Computing layout for " << num_layers
+            << " full-attention layers, world_size=" << world_size;
+
+  // Iluvatar BI-V100: verified 4-card flat PIX topology (ixsmi topo -m).
+  // All pairs connected via single PCIe bridge, equal bandwidth.
+  return compute_ilu_layerwise_layout(
+      num_layers, per_layer_kv_heads, world_size,
+      IluTopoKind::kFlatPIX);
+}
+
+}  // namespace xllm

core/distributed_runtime/layerwise_split_engine_ext.h

@@ -0,0 +1,38 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <cstdint>
+#include <optional>
+#include <vector>
+
+#include "framework/kv_cache/ilu_layerwise_layout.h"
+
+namespace xllm {
+
+/// Called by llm_engine / speculative_engine at startup.
+/// Returns a IluLayerwiseLayout if the feature is enabled, otherwise
+/// std::nullopt (fallback to uniform allocation).
+///
+/// |per_layer_kv_heads|: KV head count for each full-attention layer.
+///   Linear-attention layers are excluded.
+///   For Qwen3.5 TP=4: all entries are 4 (each full-attn layer has 4 KV heads).
+std::optional<IluLayerwiseLayout> maybe_compute_layerwise_layout(
+    int64_t num_layers,
+    const std::vector<int64_t>& per_layer_kv_heads,
+    int32_t world_size);
+
+}  // namespace xllm

core/distributed_runtime/layerwise_split_master.cpp

@@ -0,0 +1,91 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Commit: 494f293b5629 · feat · PR #2260  (adapted for Iluvatar BI-V100)
+// Master-side orchestration: at startup the master reads model config to
+// extract per-layer KV head counts, computes the layout, and stores it
+// for distribution to workers.
+//
+// For Qwen3.5 (Qwen3.6-35B-A3B):
+//   - layer_types is a mix of "full_attention" and "linear_attention"
+//   - Only full_attention layers have KV cache
+//   - Each full_attention layer has the same KV head count (4 for the
+//     35B-A3B variant)
+
+#include "distributed_runtime/layerwise_split_master.h"
+
+#include <glog/logging.h>
+
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "distributed_runtime/layerwise_split_engine_ext.h"
+#include "framework/kv_cache/kv_cache_estimation_layerwise.h"
+#include "framework/kv_cache/layerwise_split_layout.h"
+
+DECLARE_bool(enable_layerwise_split);
+
+namespace xllm {
+
+std::optional<IluLayerwiseLayout> master_compute_layerwise_layout(
+    const std::vector<std::string>& layer_types,
+    int64_t kv_heads_per_full_attn_layer,
+    int32_t world_size,
+    int64_t n_blocks,
+    int64_t block_size,
+    int64_t head_dim,
+    int64_t max_tokens,
+    int dtype_enum) {
+  if (!FLAGS_enable_layerwise_split) {
+    LOG(INFO) << "[LayerwiseSplit] Disabled; using uniform KV sharding.";
+    return std::nullopt;
+  }
+
+  // Count and build per-layer KV head count vector for full-attention
+  // layers only.  Linear-attention layers are skipped.
+  int64_t full_attn_count = 0;
+  for (const auto& lt : layer_types) {
+    if (lt == "full_attention") {
+      ++full_attn_count;
+    }
+  }
+  CHECK_GT(full_attn_count, 0) << "No full_attention layers found.";
+
+  std::vector<int64_t> per_layer_heads(full_attn_count,
+                                        kv_heads_per_full_attn_layer);
+
+  auto layout = maybe_compute_layerwise_layout(
+      full_attn_count, per_layer_heads, world_size);
+
+  if (layout.has_value()) {
+    // Run estimation for logging / capacity planning.
+    auto est = estimate_layerwise_kv_memory(
+        *layout, n_blocks, block_size, head_dim, max_tokens,
+        dtype_enum, world_size);
+
+    LOG(INFO) << "[LayerwiseSplit] Qwen3.5 layout: "
+              << full_attn_count << " full-attn layers, "
+              << (layer_types.size() - full_attn_count) << " linear-attn layers"
+              << " — Peak per-rank KV: "
+              << (est.peak_per_rank_bytes >> 20) << " MiB  (uniform would be "
+              << (est.uniform_per_rank_bytes >> 20) << " MiB, saving "
+              << est.savings_vs_uniform_pct << "%)";
+  }
+
+  return layout;
+}
+
+}  // namespace xllm

core/distributed_runtime/layerwise_split_master.h

@@ -0,0 +1,43 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "framework/kv_cache/ilu_layerwise_layout.h"
+
+namespace xllm {
+
+/// Master-side entry point: compute and log the layerwise layout.
+///
+/// For Qwen3.5 (Qwen3.6-35B-A3B):
+///   |layer_types|: "full_attention" or "linear_attention" per layer.
+///   |kv_heads_per_full_attn_layer|: KV heads for each full-attention layer.
+///   Linear-attention layers are skipped (they use GDN conv+temporal state).
+std::optional<IluLayerwiseLayout> master_compute_layerwise_layout(
+    const std::vector<std::string>& layer_types,
+    int64_t kv_heads_per_full_attn_layer,
+    int32_t world_size,
+    int64_t n_blocks,
+    int64_t block_size,
+    int64_t head_dim,
+    int64_t max_tokens,
+    int dtype_enum);
+
+}  // namespace xllm