dylanyunlong/project_6_layerwiseSplit_without_dp
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

under test, not sure no errors

Browse Source
This commit is contained in:
root
2026-09-02 07:03:56 +00:00
commit 43c43b491c
4211 changed files with 1013777 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

343
core/runtime/py_attention_metadata.cpp Normal file
View File

@@ -0,0 +1,343 @@
/* Adapted from xLLM commit 78aa2a85 (PR #2258).
Adds dp_token_counts / dp_is_decode to the pybind11-exported
AttentionMetadataView so Python model executors (Qwen3.5 MoE layers,
decode graph runners) can read per-DP-rank token counts and decide
between padded vs compact all-gather.
Original: xllm/core/runtime/py_attention_metadata.cpp
Scope: Qwen3.5 data-parallel support in project_6.
==============================================================================*/
#include "core/runtime/py_attention_metadata.h"
#include <pybind11/stl.h>
#include <torch/extension.h>
#include <utility>
/*
* NOTE: The upstream xLLM implementation #includes
* "core/framework/model/model_input_params.h"
* "core/layers/common/attention_metadata.h"
* Those headers are part of xLLM's internal C++ framework and are NOT
* open-sourced in project_6. The stub types below satisfy the build so
* the DP-specific logic compiles; the real integration will link against
* the xLLM shared libraries that provide the concrete structs.
*/
namespace project6::layer {
struct ExpandedDecodeMetadata {
bool enabled = false;
torch::Tensor kv_seq_lens;
torch::Tensor block_table;
torch::Tensor paged_kv_indptr;
torch::Tensor paged_kv_indices;
torch::Tensor paged_kv_last_page_len;
torch::Tensor paged_attention_tiling_data;
torch::Tensor kv_seq_lens_host;
std::vector<int32_t> kv_seq_lens_host_vec;
};
struct AttentionMetadata {
torch::Tensor slot_mapping;
torch::Tensor paged_kv_indptr;
torch::Tensor paged_kv_indices;
torch::Tensor paged_kv_last_page_len;
std::optional<torch::Tensor> qo_indptr;
torch::Tensor q_cu_seq_lens;
torch::Tensor kv_cu_seq_lens;
torch::Tensor block_table;
torch::Tensor kv_seq_lens;
torch::Tensor q_seq_lens;
torch::Tensor has_initial_states;
std::vector<int32_t> kv_seq_lens_vec;
std::vector<int32_t> q_seq_lens_vec;
bool is_prefill = false;
bool is_chunked_prefill = false;
ExpandedDecodeMetadata expanded_decode;
};
} // namespace project6::layer
namespace project6 {
/* Minimal stub so the two-arg constructor compiles. */
struct ModelInputParams {
struct {
std::vector<int32_t> raw_dp_global_token_nums;
std::vector<int32_t> dp_global_token_nums;
std::vector<int32_t> dp_is_decode;
} parallel;
struct {
torch::Tensor linear_state_indices;
} embedding;
};
namespace py = pybind11;
// ---------------------------------------------------------------------------
// pybind11 registration
// ---------------------------------------------------------------------------
void register_attention_metadata_views(py::module_& module) {
py::class_<PyExpandedDecodeMetadataView>(module, "ExpandedDecodeMetadataView")
.def_property_readonly("enabled", &PyExpandedDecodeMetadataView::enabled)
.def_property_readonly("kv_seq_lens",
&PyExpandedDecodeMetadataView::kv_seq_lens)
.def_property_readonly("block_table",
&PyExpandedDecodeMetadataView::block_table)
.def_property_readonly("paged_kv_indptr",
&PyExpandedDecodeMetadataView::paged_kv_indptr)
.def_property_readonly("paged_kv_indices",
&PyExpandedDecodeMetadataView::paged_kv_indices)
.def_property_readonly(
"paged_kv_last_page_len",
&PyExpandedDecodeMetadataView::paged_kv_last_page_len)
.def_property_readonly(
"paged_attention_tiling_data",
&PyExpandedDecodeMetadataView::paged_attention_tiling_data)
.def_property_readonly("kv_seq_lens_host",
&PyExpandedDecodeMetadataView::kv_seq_lens_host)
.def_property_readonly(
"kv_seq_lens_host_values",
&PyExpandedDecodeMetadataView::kv_seq_lens_host_values);
py::class_<PyAttentionMetadataView>(module, "AttentionMetadataView")
.def_property_readonly("slot_mapping",
&PyAttentionMetadataView::slot_mapping)
.def_property_readonly("paged_kv_indptr",
&PyAttentionMetadataView::paged_kv_indptr)
.def_property_readonly("paged_kv_indices",
&PyAttentionMetadataView::paged_kv_indices)
.def_property_readonly("paged_kv_last_page_len",
&PyAttentionMetadataView::paged_kv_last_page_len)
.def_property_readonly("qo_indptr", &PyAttentionMetadataView::qo_indptr)
.def_property_readonly("q_cu_seq_lens",
&PyAttentionMetadataView::q_cu_seq_lens)
.def_property_readonly("kv_cu_seq_lens",
&PyAttentionMetadataView::kv_cu_seq_lens)
.def_property_readonly("kv_seq_lens_host",
&PyAttentionMetadataView::kv_seq_lens_host)
.def_property_readonly("kv_seq_lens_host_values",
&PyAttentionMetadataView::kv_seq_lens_host_values)
.def_property_readonly("q_seq_lens_host",
&PyAttentionMetadataView::q_seq_lens_host)
.def_property_readonly("block_table",
&PyAttentionMetadataView::block_table)
.def_property_readonly("kv_seq_lens",
&PyAttentionMetadataView::kv_seq_lens)
.def_property_readonly("linear_state_indices",
&PyAttentionMetadataView::linear_state_indices)
.def_property_readonly("has_initial_state",
&PyAttentionMetadataView::has_initial_state)
/* ---- DP fields (added by PR #2258) ------------------------------ */
.def_property_readonly("dp_token_counts",
&PyAttentionMetadataView::dp_token_counts)
.def_property_readonly("dp_is_decode",
&PyAttentionMetadataView::dp_is_decode)
/* ----------------------------------------------------------------- */
.def_property_readonly("q_seq_lens", &PyAttentionMetadataView::q_seq_lens)
.def_property_readonly("expanded_decode_metadata",
&PyAttentionMetadataView::expanded_decode_metadata)
.def_property_readonly("is_prefill", &PyAttentionMetadataView::is_prefill)
.def_property_readonly("is_chunked_prefill",
&PyAttentionMetadataView::is_chunked_prefill);
}
// ---------------------------------------------------------------------------
// PyExpandedDecodeMetadataView
// ---------------------------------------------------------------------------
PyExpandedDecodeMetadataView::PyExpandedDecodeMetadataView(
std::shared_ptr<layer::AttentionMetadata> metadata)
: metadata_(std::move(metadata)) {}
bool PyExpandedDecodeMetadataView::enabled() const {
return metadata().enabled;
}
py::object PyExpandedDecodeMetadataView::kv_seq_lens() const {
return metadata().kv_seq_lens.defined() ? py::cast(metadata().kv_seq_lens)
: py::none();
}
py::object PyExpandedDecodeMetadataView::block_table() const {
return metadata().block_table.defined() ? py::cast(metadata().block_table)
: py::none();
}
py::object PyExpandedDecodeMetadataView::paged_kv_indptr() const {
return metadata().paged_kv_indptr.defined()
? py::cast(metadata().paged_kv_indptr)
: py::none();
}
py::object PyExpandedDecodeMetadataView::paged_kv_indices() const {
return metadata().paged_kv_indices.defined()
? py::cast(metadata().paged_kv_indices)
: py::none();
}
py::object PyExpandedDecodeMetadataView::paged_kv_last_page_len() const {
return metadata().paged_kv_last_page_len.defined()
? py::cast(metadata().paged_kv_last_page_len)
: py::none();
}
py::object PyExpandedDecodeMetadataView::paged_attention_tiling_data() const {
return metadata().paged_attention_tiling_data.defined()
? py::cast(metadata().paged_attention_tiling_data)
: py::none();
}
py::object PyExpandedDecodeMetadataView::kv_seq_lens_host() const {
return metadata().kv_seq_lens_host.defined()
? py::cast(metadata().kv_seq_lens_host)
: py::none();
}
const std::vector<int32_t>&
PyExpandedDecodeMetadataView::kv_seq_lens_host_values() const {
return metadata().kv_seq_lens_host_vec;
}
const layer::ExpandedDecodeMetadata& PyExpandedDecodeMetadataView::metadata()
const {
return metadata_->expanded_decode;
}
// ---------------------------------------------------------------------------
// PyAttentionMetadataView
// ---------------------------------------------------------------------------
PyAttentionMetadataView::PyAttentionMetadataView(
std::shared_ptr<layer::AttentionMetadata> metadata)
: metadata_(std::move(metadata)),
kv_seq_lens_host_(
make_host_int32_view(metadata_, metadata_->kv_seq_lens_vec)),
q_seq_lens_host_(
make_host_int32_view(metadata_, metadata_->q_seq_lens_vec)) {}
PyAttentionMetadataView::PyAttentionMetadataView(
std::shared_ptr<layer::AttentionMetadata> metadata,
const ModelInputParams& params)
: PyAttentionMetadataView(std::move(metadata)) {
linear_state_indices_ = params.embedding.linear_state_indices;
/* ---- DP fields (added by PR #2258) ---------------------------------- */
dp_token_counts_ = params.parallel.raw_dp_global_token_nums.empty()
? params.parallel.dp_global_token_nums
: params.parallel.raw_dp_global_token_nums;
dp_is_decode_ = params.parallel.dp_is_decode;
/* --------------------------------------------------------------------- */
}
const torch::Tensor& PyAttentionMetadataView::slot_mapping() const {
return metadata_->slot_mapping;
}
const torch::Tensor& PyAttentionMetadataView::paged_kv_indptr() const {
return metadata_->paged_kv_indptr;
}
const torch::Tensor& PyAttentionMetadataView::paged_kv_indices() const {
return metadata_->paged_kv_indices;
}
const torch::Tensor& PyAttentionMetadataView::paged_kv_last_page_len() const {
return metadata_->paged_kv_last_page_len;
}
py::object PyAttentionMetadataView::qo_indptr() const {
if (!metadata_->qo_indptr.has_value() || !metadata_->qo_indptr->defined()) {
return py::none();
}
return py::cast(*metadata_->qo_indptr);
}
py::object PyAttentionMetadataView::q_cu_seq_lens() const {
return optional_tensor(metadata_->q_cu_seq_lens);
}
py::object PyAttentionMetadataView::kv_cu_seq_lens() const {
return optional_tensor(metadata_->kv_cu_seq_lens);
}
py::object PyAttentionMetadataView::kv_seq_lens_host() const {
return optional_tensor(kv_seq_lens_host_);
}
const std::vector<int32_t>& PyAttentionMetadataView::kv_seq_lens_host_values()
const {
return metadata_->kv_seq_lens_vec;
}
py::object PyAttentionMetadataView::block_table() const {
return optional_tensor(metadata_->block_table);
}
py::object PyAttentionMetadataView::kv_seq_lens() const {
return optional_tensor(metadata_->kv_seq_lens);
}
py::object PyAttentionMetadataView::linear_state_indices() const {
return optional_tensor(linear_state_indices_);
}
py::object PyAttentionMetadataView::has_initial_state() const {
return optional_tensor(metadata_->has_initial_states);
}
/* ---- DP fields (added by PR #2258) ------------------------------------ */
const std::vector<int32_t>& PyAttentionMetadataView::dp_token_counts() const {
return dp_token_counts_;
}
const std::vector<int32_t>& PyAttentionMetadataView::dp_is_decode() const {
return dp_is_decode_;
}
/* ----------------------------------------------------------------------- */
py::object PyAttentionMetadataView::q_seq_lens() const {
return optional_tensor(metadata_->q_seq_lens);
}
py::object PyAttentionMetadataView::q_seq_lens_host() const {
return optional_tensor(q_seq_lens_host_);
}
PyExpandedDecodeMetadataView PyAttentionMetadataView::expanded_decode_metadata()
const {
return PyExpandedDecodeMetadataView(metadata_);
}
bool PyAttentionMetadataView::is_prefill() const {
return metadata_->is_prefill;
}
bool PyAttentionMetadataView::is_chunked_prefill() const {
return metadata_->is_chunked_prefill;
}
torch::Tensor PyAttentionMetadataView::make_host_int32_view(
const std::shared_ptr<layer::AttentionMetadata>& metadata,
std::vector<int32_t>& host_vec) {
if (host_vec.empty()) {
return torch::Tensor();
}
std::shared_ptr<layer::AttentionMetadata> owner = metadata;
return torch::from_blob(
host_vec.data(),
{static_cast<int64_t>(host_vec.size())},
[owner = std::move(owner)](void*) mutable { owner.reset(); },
torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU));
}
py::object PyAttentionMetadataView::optional_tensor(
const torch::Tensor& tensor) {
return tensor.defined() ? py::cast(tensor) : py::none();
}
} // namespace project6

100
core/runtime/py_attention_metadata.h Normal file
View File

@@ -0,0 +1,100 @@
/* Adapted from xLLM commit 78aa2a85 (PR #2258).
Adds dp_token_counts / dp_is_decode fields to PyAttentionMetadataView
so the Python attention backend can partition KV cache by DP group.
Original: xllm/core/runtime/py_attention_metadata.h
Scope: Qwen3.5 data-parallel support in project_6.
==============================================================================*/
#pragma once
#include <pybind11/pybind11.h>
#include <torch/torch.h>
#include <cstdint>
#include <memory>
#include <vector>
/* Forward declarations — project_6 keeps these in its own layer namespace. */
namespace project6::layer {
struct AttentionMetadata;
struct ExpandedDecodeMetadata;
} // namespace project6::layer
namespace project6 {
struct ModelInputParams;
void register_attention_metadata_views(pybind11::module_& module);
class PyExpandedDecodeMetadataView final {
public:
explicit PyExpandedDecodeMetadataView(
std::shared_ptr<layer::AttentionMetadata> metadata);
bool enabled() const;
pybind11::object kv_seq_lens() const;
pybind11::object block_table() const;
pybind11::object paged_kv_indptr() const;
pybind11::object paged_kv_indices() const;
pybind11::object paged_kv_last_page_len() const;
pybind11::object paged_attention_tiling_data() const;
pybind11::object kv_seq_lens_host() const;
const std::vector<int32_t>& kv_seq_lens_host_values() const;
private:
const layer::ExpandedDecodeMetadata& metadata() const;
std::shared_ptr<layer::AttentionMetadata> metadata_;
};
class PyAttentionMetadataView final {
public:
explicit PyAttentionMetadataView(
std::shared_ptr<layer::AttentionMetadata> metadata);
PyAttentionMetadataView(std::shared_ptr<layer::AttentionMetadata> metadata,
const ModelInputParams& params);
const torch::Tensor& slot_mapping() const;
const torch::Tensor& paged_kv_indptr() const;
const torch::Tensor& paged_kv_indices() const;
const torch::Tensor& paged_kv_last_page_len() const;
pybind11::object qo_indptr() const;
pybind11::object q_cu_seq_lens() const;
pybind11::object kv_cu_seq_lens() const;
pybind11::object kv_seq_lens_host() const;
const std::vector<int32_t>& kv_seq_lens_host_values() const;
pybind11::object q_seq_lens_host() const;
pybind11::object block_table() const;
pybind11::object kv_seq_lens() const;
pybind11::object linear_state_indices() const;
pybind11::object has_initial_state() const;
/* ---- DP fields (added by PR #2258) ---------------------------------- */
const std::vector<int32_t>& dp_token_counts() const;
const std::vector<int32_t>& dp_is_decode() const;
/* --------------------------------------------------------------------- */
pybind11::object q_seq_lens() const;
PyExpandedDecodeMetadataView expanded_decode_metadata() const;
bool is_prefill() const;
bool is_chunked_prefill() const;
private:
static torch::Tensor make_host_int32_view(
const std::shared_ptr<layer::AttentionMetadata>& metadata,
std::vector<int32_t>& host_vec);
static pybind11::object optional_tensor(const torch::Tensor& tensor);
std::shared_ptr<layer::AttentionMetadata> metadata_;
torch::Tensor kv_seq_lens_host_;
torch::Tensor q_seq_lens_host_;
torch::Tensor linear_state_indices_;
/* ---- DP fields (added by PR #2258) ---------------------------------- */
std::vector<int32_t> dp_token_counts_;
std::vector<int32_t> dp_is_decode_;
/* --------------------------------------------------------------------- */
};
} // namespace project6

73
core/runtime/worker_layerwise_init.cpp Normal file
View File

@@ -0,0 +1,73 @@
/* 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)
// Worker-side helper: after the worker receives its LayerwiseSplitLayout
// from the master, it calls this to allocate per-layer KV caches with
// the correct shard sizes.
#include "runtime/worker_layerwise_init.h"
#include <glog/logging.h>
#include <vector>
#include "framework/kv_cache/kv_cache.h"
#include "framework/kv_cache/kv_cache_layerwise.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 {
bool worker_allocate_layerwise_kv_cache(
std::vector<KVCache>& kv_caches,
const KVCacheShape& kv_cache_shape,
const KVCacheCreateOptions& create_options,
const LayerwiseSplitLayout& layout,
int32_t rank) {
LOG(INFO) << "[Worker " << rank << "] Applying layerwise KV layout: "
<< layout.layers_on_rank(rank) << " layers assigned.";
try {
allocate_kv_caches_layerwise(
kv_caches, kv_cache_shape, create_options, layout, rank);
} catch (const std::exception& e) {
LOG(ERROR) << "[Worker " << rank
<< "] Failed to allocate layerwise KV cache: " << e.what();
return false;
}
// Verify: assigned layers should have non-empty caches.
for (int64_t lid = 0; lid < layout.num_layers(); ++lid) {
bool owns = layout.rank_owns_layer(rank, lid);
bool empty = kv_caches[lid].empty();
if (owns && empty) {
LOG(ERROR) << "[Worker " << rank << "] Layer " << lid
<< " is assigned but KV cache is empty.";
return false;
}
if (!owns && !empty) {
LOG(ERROR) << "[Worker " << rank << "] Layer " << lid
<< " is NOT assigned but KV cache is non-empty.";
return false;
}
}
LOG(INFO) << "[Worker " << rank << "] Layerwise KV cache allocation OK.";
return true;
}
} // namespace xllm

37
core/runtime/worker_layerwise_init.h Normal file
View File

@@ -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 {
/// Worker-side entry: allocate KV caches per the received layout.
/// Returns true on success; false if any verification check fails.
bool worker_allocate_layerwise_kv_cache(
std::vector<KVCache>& kv_caches,
const KVCacheShape& kv_cache_shape,
const KVCacheCreateOptions& create_options,
const LayerwiseSplitLayout& layout,
int32_t rank);
} // namespace xllm