229 lines
8.5 KiB
Rust
229 lines
8.5 KiB
Rust
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
//
|
|
// 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
|
|
//
|
|
// http://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.
|
|
|
|
#![cfg(feature = "block-manager")]
|
|
|
|
use super::*;
|
|
use pyo3::PyResult;
|
|
|
|
mod block;
|
|
mod block_list;
|
|
mod dlpack;
|
|
mod layer;
|
|
|
|
/// Add bingings from this crate to the provided module
|
|
pub fn add_to_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
|
m.add_class::<layer::Layer>()?;
|
|
m.add_class::<block::Block>()?;
|
|
m.add_class::<block_list::BlockList>()?;
|
|
m.add_class::<BlockManager>()?;
|
|
Ok(())
|
|
}
|
|
|
|
#[pyclass]
|
|
pub struct BlockManager {
|
|
inner: Arc<dynamo_llm::block_manager::ReferenceBlockManager>,
|
|
// TODO: Metadata should be stored in the block manager?
|
|
dtype: dynamo_llm::common::dtype::DType,
|
|
device_id: usize,
|
|
}
|
|
|
|
#[pymethods]
|
|
impl BlockManager {
|
|
#[new]
|
|
#[pyo3(signature = (worker_id, num_layer, outer_dim, page_size, inner_dim, dtype=None, host_num_blocks=None, device_num_blocks=None, device_id=0))]
|
|
fn new(
|
|
worker_id: u64,
|
|
num_layer: usize,
|
|
outer_dim: usize,
|
|
page_size: usize,
|
|
inner_dim: usize,
|
|
dtype: Option<String>,
|
|
host_num_blocks: Option<usize>,
|
|
device_num_blocks: Option<usize>,
|
|
device_id: usize,
|
|
) -> PyResult<Self> {
|
|
let mut config = dynamo_llm::block_manager::KvBlockManagerConfig::builder().runtime(
|
|
dynamo_llm::block_manager::KvManagerRuntimeConfig::builder()
|
|
.worker_id(worker_id)
|
|
.build()
|
|
.map_err(to_pyerr)?,
|
|
);
|
|
let mut model_config = dynamo_llm::block_manager::KvManagerModelConfig::builder()
|
|
.num_layers(num_layer)
|
|
.outer_dim(outer_dim)
|
|
.page_size(page_size)
|
|
.inner_dim(inner_dim);
|
|
let mut dtype_ = dynamo_llm::common::dtype::DType::FP16; // Default in block_manager config
|
|
if let Some(dtype_str) = dtype {
|
|
dtype_ = match dtype_str.as_str() {
|
|
"fp8" | "FP8" => dynamo_llm::common::dtype::DType::FP8,
|
|
"fp16" | "FP16" => dynamo_llm::common::dtype::DType::FP16,
|
|
"bf16" | "BF16" => dynamo_llm::common::dtype::DType::BF16,
|
|
"fp32" | "FP32" => dynamo_llm::common::dtype::DType::FP32,
|
|
"u8" | "U8" => dynamo_llm::common::dtype::DType::U8,
|
|
"u16" | "U16" => dynamo_llm::common::dtype::DType::U16,
|
|
"u32" | "U32" => dynamo_llm::common::dtype::DType::U32,
|
|
"u64" | "U64" => dynamo_llm::common::dtype::DType::U64,
|
|
"i8" | "I8" => dynamo_llm::common::dtype::DType::I8,
|
|
"i16" | "I16" => dynamo_llm::common::dtype::DType::I16,
|
|
"i32" | "I32" => dynamo_llm::common::dtype::DType::I32,
|
|
"i64" | "I64" => dynamo_llm::common::dtype::DType::I64,
|
|
_ => {
|
|
return Err(pyo3::exceptions::PyValueError::new_err(format!(
|
|
"Unsupported dtype: {}",
|
|
dtype_str
|
|
)))
|
|
}
|
|
};
|
|
}
|
|
model_config = model_config.dtype(dtype_.clone());
|
|
config = config.model(model_config.build().map_err(to_pyerr)?);
|
|
if let Some(host_num_blocks) = host_num_blocks {
|
|
config = config.host_layout(
|
|
dynamo_llm::block_manager::KvManagerLayoutConfig::builder()
|
|
.num_blocks(host_num_blocks)
|
|
.allocator(
|
|
dynamo_llm::block_manager::storage::PinnedAllocator::new()
|
|
.map_err(to_pyerr)?,
|
|
)
|
|
.build()
|
|
.map_err(to_pyerr)?,
|
|
);
|
|
}
|
|
if let Some(device_num_blocks) = device_num_blocks {
|
|
config = config.device_layout(
|
|
dynamo_llm::block_manager::KvManagerLayoutConfig::builder()
|
|
.num_blocks(device_num_blocks)
|
|
.allocator(
|
|
dynamo_llm::block_manager::storage::DeviceAllocator::new(device_id)
|
|
.map_err(to_pyerr)?,
|
|
)
|
|
.build()
|
|
.map_err(to_pyerr)?,
|
|
);
|
|
}
|
|
let config = config.build().map_err(to_pyerr)?;
|
|
let tokio_runtime = pyo3_async_runtimes::tokio::get_runtime();
|
|
Ok(BlockManager {
|
|
inner: Arc::from(
|
|
tokio_runtime
|
|
.block_on(async {
|
|
dynamo_llm::block_manager::ReferenceBlockManager::new(config)
|
|
})
|
|
.map_err(to_pyerr)?,
|
|
),
|
|
dtype: dtype_,
|
|
device_id: device_id,
|
|
})
|
|
}
|
|
|
|
fn allocate_host_blocks_blocking(&self, count: usize) -> PyResult<block_list::BlockList> {
|
|
let blocks = self
|
|
.inner
|
|
.host()
|
|
.ok_or_else(|| {
|
|
pyo3::exceptions::PyRuntimeError::new_err("Host allocator not available")
|
|
})?
|
|
.allocate_blocks_blocking(count)
|
|
.map_err(to_pyerr)?;
|
|
// Wrap each block in an enum accounting for Pinned & Device block
|
|
let blocks = blocks
|
|
.into_iter()
|
|
.map(|b| block::BlockType::Pinned(b))
|
|
.collect();
|
|
Ok(block_list::BlockList::from_rust(
|
|
blocks,
|
|
self.dtype.clone(),
|
|
self.device_id,
|
|
))
|
|
}
|
|
|
|
#[pyo3(signature = (count))]
|
|
fn allocate_host_blocks<'py>(
|
|
&self,
|
|
py: Python<'py>,
|
|
count: usize,
|
|
) -> PyResult<Bound<'py, PyAny>> {
|
|
let inner = self.inner.clone();
|
|
let dtype = self.dtype.clone();
|
|
let device_id = self.device_id;
|
|
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
|
let blocks = inner
|
|
.host()
|
|
.ok_or_else(|| {
|
|
pyo3::exceptions::PyRuntimeError::new_err("Host allocator not available")
|
|
})?
|
|
.allocate_blocks(count)
|
|
.await
|
|
.map_err(to_pyerr)?;
|
|
// Wrap each block in an enum accounting for Pinned & Device block
|
|
let blocks = blocks
|
|
.into_iter()
|
|
.map(|b| block::BlockType::Pinned(b))
|
|
.collect();
|
|
Ok(block_list::BlockList::from_rust(blocks, dtype, device_id))
|
|
})
|
|
}
|
|
|
|
fn allocate_device_blocks_blocking(&self, count: usize) -> PyResult<block_list::BlockList> {
|
|
let blocks = self
|
|
.inner
|
|
.device()
|
|
.ok_or_else(|| {
|
|
pyo3::exceptions::PyRuntimeError::new_err("Device allocator not available")
|
|
})?
|
|
.allocate_blocks_blocking(count)
|
|
.map_err(to_pyerr)?;
|
|
// Wrap each block in an enum accounting for Pinned & Device block
|
|
let blocks = blocks
|
|
.into_iter()
|
|
.map(|b| block::BlockType::Device(b))
|
|
.collect();
|
|
Ok(block_list::BlockList::from_rust(
|
|
blocks,
|
|
self.dtype.clone(),
|
|
self.device_id,
|
|
))
|
|
}
|
|
|
|
#[pyo3(signature = (count))]
|
|
fn allocate_device_blocks<'py>(
|
|
&self,
|
|
py: Python<'py>,
|
|
count: usize,
|
|
) -> PyResult<Bound<'py, PyAny>> {
|
|
let inner = self.inner.clone();
|
|
let dtype = self.dtype.clone();
|
|
let device_id = self.device_id;
|
|
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
|
let blocks = inner
|
|
.device()
|
|
.ok_or_else(|| {
|
|
pyo3::exceptions::PyRuntimeError::new_err("Device allocator not available")
|
|
})?
|
|
.allocate_blocks(count)
|
|
.await
|
|
.map_err(to_pyerr)?;
|
|
// Wrap each block in an enum accounting for Pinned & Device block
|
|
let blocks = blocks
|
|
.into_iter()
|
|
.map(|b| block::BlockType::Device(b))
|
|
.collect();
|
|
Ok(block_list::BlockList::from_rust(blocks, dtype, device_id))
|
|
})
|
|
}
|
|
}
|