完整部署快照
This commit is contained in:
root
348
core/runtime/py_attention_metadata.cpp
Normal file
348
core/runtime/py_attention_metadata.cpp
Normal file
@@ -0,0 +1,348 @@
|
||||
/* 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
|
||||
|
||||
PYBIND11_MODULE(py_attention_metadata, m) {
|
||||
m.doc() = "DP-aware attention metadata (project6, ported from xLLM PR #2258)";
|
||||
project6::register_attention_metadata_views(m);
|
||||
}
|
||||
100
core/runtime/py_attention_metadata.h
Normal file
100
core/runtime/py_attention_metadata.h
Normal 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
|
||||
58
core/runtime/worker_layerwise_init.cpp
Normal file
58
core/runtime/worker_layerwise_init.cpp
Normal file
@@ -0,0 +1,58 @@
|
||||
/* Worker-side helper: compute layer_cache_owned mask.
|
||||
Migrated pattern from upstream worker_impl.cpp lines 490-498. */
|
||||
|
||||
#include "runtime/worker_layerwise_init.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "framework/kv_cache/kv_cache_estimation.h"
|
||||
#include "framework/kv_cache/layerwise_split_layout.h"
|
||||
#include "framework/model/model_args.h"
|
||||
|
||||
namespace xllm {
|
||||
|
||||
std::vector<bool> worker_compute_layer_cache_owned(
|
||||
const std::vector<std::string>& layer_types,
|
||||
int32_t layerwise_split_size,
|
||||
int32_t rank,
|
||||
int64_t num_layers) {
|
||||
CHECK_GE(layerwise_split_size, 1);
|
||||
CHECK_GE(rank, 0);
|
||||
CHECK_GT(num_layers, 0);
|
||||
|
||||
if (layerwise_split_size <= 1) {
|
||||
LOG(INFO) << "[Worker " << rank << "] No layerwise split, all "
|
||||
<< num_layers << " layers owned.";
|
||||
return std::vector<bool>(static_cast<size_t>(num_layers), true);
|
||||
}
|
||||
|
||||
// Build a ModelArgs with the layer_types for upstream build_layer_cache_owned
|
||||
ModelArgs args;
|
||||
args.n_layers(num_layers);
|
||||
args.layer_types(layer_types);
|
||||
|
||||
const LayerwiseSplitLayout layout(
|
||||
/*enabled=*/true,
|
||||
layerwise_split_size,
|
||||
rank % layerwise_split_size);
|
||||
|
||||
auto owned = build_layer_cache_owned(args, layout, num_layers);
|
||||
|
||||
int64_t owned_count = 0;
|
||||
int64_t linear_count = 0;
|
||||
for (int64_t i = 0; i < num_layers; ++i) {
|
||||
if (owned[static_cast<size_t>(i)]) ++owned_count;
|
||||
if (!is_full_attention_layer(args, i)) ++linear_count;
|
||||
}
|
||||
|
||||
LOG(INFO) << "[Worker " << rank << "] Layerwise split (size="
|
||||
<< layerwise_split_size << "): " << owned_count << "/"
|
||||
<< num_layers << " layers owned, " << linear_count
|
||||
<< " linear-attention (always owned).";
|
||||
|
||||
return owned;
|
||||
}
|
||||
|
||||
} // namespace xllm
|
||||
40
core/runtime/worker_layerwise_init.h
Normal file
40
core/runtime/worker_layerwise_init.h
Normal file
@@ -0,0 +1,40 @@
|
||||
/* 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 <string>
|
||||
#include <vector>
|
||||
|
||||
#include "framework/kv_cache/layerwise_split_layout.h"
|
||||
|
||||
namespace xllm {
|
||||
|
||||
/// Worker-side entry: compute the layer_cache_owned mask and validate.
|
||||
///
|
||||
/// For Qwen3.5: the worker receives layerwise_split_size from the master
|
||||
/// and builds its ownership mask. Only full_attention layers participate
|
||||
/// in the split; linear_attention layers are always "owned" (no KV cache).
|
||||
///
|
||||
/// Returns the layer_cache_owned mask, which the vendor vLLM cache engine
|
||||
/// uses to decide which layers get real KV allocations vs scratch.
|
||||
std::vector<bool> worker_compute_layer_cache_owned(
|
||||
const std::vector<std::string>& layer_types,
|
||||
int32_t layerwise_split_size,
|
||||
int32_t rank,
|
||||
int64_t num_layers);
|
||||
|
||||
} // namespace xllm
|
||||
Reference in New Issue
Block a user