yangzhuxinyzxandClaude Opus 5.5
899e9ed331
Runs DeepSeek-V4.1-Flash-1017 (GSQ-RCO IQ2_XXS GGUF + Engram IQ4_NL + DSpark5 draft) on 4x V100-32G (TP4, one x16 PCIe uplink, 62 GB host RAM) with expert offloading. All new paths are lossless and switchable. Prefill (G1-G7): - G1 activation memory: per-tile MoE scratch + FP32 accumulator, in-place expand_hc / large all-reduce, windowed hc_transition, pynccl all_gather, empty_cache before big passes; chunks up to 16K. - G2 cold-expert staging ring issued ahead across tiles and layers. - G3 routed down projection split by output rows (no duplicated TP boundary blocks, -2.25 GiB/GPU) with an IPC peer gather for decode. - G5 Engram readers, adaptive indexer tile, persistent activation grid. - G6 routed-only staging for short prompts. - G7 same-session prefix cache (scheduler guard + worker restore_prefix). Long context (G4): - Window layers keep 256-row private rings. - Packed KV (VLLM_V41_PACKED_KV): FP4 + E4M3 compressed latents in a shared host memfd plus a 32K GPU hot ring, FP4 index keys on GPU, CUDA decode index scores and segment-parallel exact top-k. Enables 512K. - RoPE rotates on the fly past its table; C2/Engram/guard limits and the DSpark draft length cap are lifted to the target length. Results vs phase 1: TTFT short 5.6 -> 0.8-1.1 s, 8K 27.7 -> 5.8 s, 15K 44.0 -> 11.1 s; 128K 110 s and 500K 651 s (needle passes); follow-up turns hit the prefix cache in ~1-1.7 s; fresh-text decode 57.8 -> 53.5 ms/round. Tokens at <=16K are identical with packed KV. Tooling: benchmarks/deepseek_v41/ gains the v41k SM70 extension source and build script, serve configs (phase 2: 16K, phase 3: 512K + prefix cache + tools + thinking), launch/validation tools, GPU micro-tests and HANDOFF.md. V4.1 switches are registered in vllm/envs.py and read through vllm.envs. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
933 lines
40 KiB
C++
933 lines
40 KiB
C++
// SPDX-License-Identifier: Apache-2.0
|
|
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
//
|
|
// Host lane of the routed-expert GPU cache.
|
|
//
|
|
// Every MoE layer of a decode step publishes its routed expert ids into a
|
|
// host-mapped ring (one entry per layer call). This thread consumes the ring
|
|
// in order and keeps per-layer LRU recency. When an entry reports experts
|
|
// without a GPU slot, it plans the copies: a slot per missing expert (a spare
|
|
// one, or the least recently used expert's, written back first when the host
|
|
// has no copy of it), plus evictions that restore the spare slots. The plan is
|
|
// written into a host-mapped reply that the model stream polls; the model
|
|
// stream then performs the copies itself and updates its slot table.
|
|
//
|
|
// The lane never issues GPU work on the decode path. A release that had to
|
|
// pass through a CUDA stream could sit in the same hardware queue behind the
|
|
// very wait it releases (streams share a few hardware queues), which
|
|
// deadlocks. Results never depend on residency: a slot holds the same bytes
|
|
// as the expert's original bank row.
|
|
|
|
#include <torch/extension.h>
|
|
|
|
#include <cuda_runtime.h>
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <atomic>
|
|
#include <chrono>
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <memory>
|
|
#include <mutex>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <thread>
|
|
#include <vector>
|
|
|
|
#include "gguf_shared_expert_dma.h"
|
|
#include "gguf_shared_expert_lane.h"
|
|
#include "gguf_shared_expert_jobs.h"
|
|
#include "gguf_shared_engram.h"
|
|
|
|
#if defined(__x86_64__) || defined(__i386__)
|
|
#include <immintrin.h>
|
|
#endif
|
|
|
|
namespace {
|
|
|
|
constexpr int kShards = 3;
|
|
constexpr int kPrefillSwapBatch = 8;
|
|
constexpr int kEntryHeader = 8; // seq, layer, count, miss, reserved...
|
|
constexpr int kReplyHeader = 8; // seq, ops, failed, reserved...
|
|
constexpr int kOpInts = 5; // slot, load host slot, writeback host slot,
|
|
// expert now in slot, expert evicted
|
|
|
|
enum Stat : int {
|
|
kEntries = 0,
|
|
kRequests,
|
|
kLoaded,
|
|
kEvicted,
|
|
kWritebacks,
|
|
kServeNs,
|
|
kMaxServeNs,
|
|
kSkipped,
|
|
kErrors,
|
|
kPrefills,
|
|
kPublished,
|
|
kConsumed,
|
|
kStats
|
|
};
|
|
|
|
inline void cpu_relax() {
|
|
#if defined(__x86_64__) || defined(__i386__)
|
|
_mm_pause();
|
|
#endif
|
|
}
|
|
|
|
inline int64_t now_ns() {
|
|
return std::chrono::duration_cast<std::chrono::nanoseconds>(
|
|
std::chrono::steady_clock::now().time_since_epoch())
|
|
.count();
|
|
}
|
|
|
|
void check(cudaError_t status, const char* what) {
|
|
if (status != cudaSuccess) {
|
|
throw std::runtime_error(std::string("expert cache: ") + what + ": " +
|
|
cudaGetErrorString(status));
|
|
}
|
|
}
|
|
|
|
// Only independent copies belong in one batch. Callers separate incoming
|
|
// saves, victim write-backs and GPU commits to preserve host-slot aliases.
|
|
class PrefillCopies {
|
|
public:
|
|
static constexpr int kCapacity = kPrefillSwapBatch * kShards;
|
|
|
|
void add(void* destination, const void* source, size_t bytes,
|
|
cudaMemcpyKind kind) {
|
|
TORCH_CHECK(count_ < kCapacity, "expert cache: copy batch is full");
|
|
destinations_[count_] = destination;
|
|
sources_[count_] = const_cast<void*>(source);
|
|
bytes_[count_] = bytes;
|
|
kinds_[count_++] = kind;
|
|
}
|
|
|
|
void submit(cudaStream_t stream) {
|
|
#if CUDART_VERSION >= 12080
|
|
// Legacy default streams and captured preparation retain the established
|
|
// individual-copy route. The serving prefill streams are explicit.
|
|
if (count_ > 1 && stream != nullptr) {
|
|
cudaStreamCaptureStatus capture;
|
|
check(cudaStreamIsCapturing(stream, &capture), "query copy capture");
|
|
if (capture == cudaStreamCaptureStatusNone) {
|
|
cudaMemcpyAttributes attributes{};
|
|
attributes.srcAccessOrder = cudaMemcpySrcAccessOrderStream;
|
|
size_t first = 0, failed = 0;
|
|
auto status = cudaMemcpyBatchAsync(
|
|
destinations_.data(), sources_.data(), bytes_.data(), count_,
|
|
&attributes, &first, 1, &failed, stream);
|
|
if (status == cudaSuccess) return;
|
|
if (status != cudaErrorNotSupported)
|
|
check(status, "batch prefill copies");
|
|
}
|
|
}
|
|
#endif
|
|
for (int i = 0; i < count_; ++i)
|
|
check(cudaMemcpyAsync(destinations_[i], sources_[i], bytes_[i], kinds_[i],
|
|
stream),
|
|
"copy prefill bytes");
|
|
}
|
|
|
|
bool full() const { return count_ == kCapacity; }
|
|
|
|
private:
|
|
std::array<void*, kCapacity> destinations_{}, sources_{};
|
|
std::array<size_t, kCapacity> bytes_{};
|
|
std::array<cudaMemcpyKind, kCapacity> kinds_{};
|
|
int count_ = 0;
|
|
};
|
|
|
|
class ExpertCacheLane {
|
|
public:
|
|
ExpertCacheLane(int64_t device, int64_t layers, int64_t experts,
|
|
int64_t slots, int64_t host_slots, int64_t spare,
|
|
std::vector<int64_t> slot_bytes,
|
|
std::vector<int64_t> device_banks,
|
|
std::vector<int64_t> host_banks, int64_t ring,
|
|
int64_t ring_entries, int64_t entry_ints, int64_t published,
|
|
int64_t reply, int64_t max_ops, torch::Tensor slot_of,
|
|
torch::Tensor host_slot_of)
|
|
: device_(device),
|
|
layers_(layers),
|
|
experts_(experts),
|
|
slots_(slots),
|
|
host_slots_(host_slots),
|
|
spare_(spare),
|
|
max_ops_(max_ops),
|
|
slot_bytes_(std::move(slot_bytes)),
|
|
device_banks_(std::move(device_banks)),
|
|
host_banks_(std::move(host_banks)),
|
|
ring_(reinterpret_cast<int32_t*>(ring)),
|
|
ring_entries_(ring_entries),
|
|
entry_ints_(entry_ints),
|
|
published_(reinterpret_cast<int64_t*>(published)),
|
|
reply_(reinterpret_cast<int32_t*>(reply)),
|
|
slot_of_(slot_of.data_ptr<int32_t>()),
|
|
host_slot_of_(host_slot_of.data_ptr<int32_t>()) {
|
|
TORCH_CHECK(slot_bytes_.size() == size_t(layers * kShards) &&
|
|
device_banks_.size() == size_t(layers * kShards) &&
|
|
host_banks_.size() == size_t(layers * kShards),
|
|
"expert cache: bad bank tables");
|
|
TORCH_CHECK(slot_of.numel() == layers * experts &&
|
|
host_slot_of.numel() == layers * experts,
|
|
"expert cache: bad residency tables");
|
|
TORCH_CHECK(entry_ints_ > kEntryHeader && ring_entries_ > 0,
|
|
"expert cache: bad ring");
|
|
TORCH_CHECK(max_ops_ >= entry_ints_ - kEntryHeader + spare_,
|
|
"expert cache: reply too small");
|
|
expert_of_slot_.assign(layers * slots, -1);
|
|
host_owner_.assign(layers * host_slots, -1);
|
|
last_use_.assign(layers * experts, 0);
|
|
clock_.assign(layers, 0);
|
|
free_slots_.resize(layers);
|
|
free_host_.resize(layers);
|
|
for (int64_t l = 0; l < layers; ++l) {
|
|
for (int64_t e = 0; e < experts; ++e) {
|
|
int32_t s = slot_of_[l * experts + e];
|
|
int32_t h = host_slot_of_[l * experts + e];
|
|
if (s >= 0) {
|
|
TORCH_CHECK(s < slots && expert_of_slot_[l * slots + s] < 0,
|
|
"expert cache: duplicate slot");
|
|
expert_of_slot_[l * slots + s] = e;
|
|
} else {
|
|
TORCH_CHECK(h >= 0, "expert cache: expert without storage");
|
|
}
|
|
if (h >= 0) {
|
|
TORCH_CHECK(h < host_slots && host_owner_[l * host_slots + h] < 0,
|
|
"expert cache: duplicate host slot");
|
|
host_owner_[l * host_slots + h] = e;
|
|
}
|
|
}
|
|
for (int64_t s = slots - 1; s >= 0; --s)
|
|
if (expert_of_slot_[l * slots + s] < 0)
|
|
free_slots_[l].push_back(static_cast<int32_t>(s));
|
|
for (int64_t h = host_slots - 1; h >= 0; --h)
|
|
if (host_owner_[l * host_slots + h] < 0)
|
|
free_host_[l].push_back(static_cast<int32_t>(h));
|
|
}
|
|
for (auto& value : stats_) value.store(0);
|
|
check(cudaHostAlloc(reinterpret_cast<void**>(&prefill_maps_),
|
|
layers * experts * sizeof(int32_t),
|
|
cudaHostAllocPortable),
|
|
"alloc prefill maps");
|
|
check(cudaHostAlloc(reinterpret_cast<void**>(&residency_maps_),
|
|
2 * layers * experts * sizeof(int32_t),
|
|
cudaHostAllocPortable),
|
|
"alloc residency maps");
|
|
consumed_ = __atomic_load_n(published_, __ATOMIC_ACQUIRE);
|
|
consumed_seen_.store(consumed_);
|
|
thread_ = std::thread([this] { run(); });
|
|
}
|
|
|
|
~ExpertCacheLane() { stop(); }
|
|
|
|
void stop() {
|
|
stop_.store(true, std::memory_order_release);
|
|
if (thread_.joinable()) thread_.join();
|
|
if (prefill_maps_ != nullptr) {
|
|
cudaFreeHost(prefill_maps_);
|
|
prefill_maps_ = nullptr;
|
|
}
|
|
if (residency_maps_ != nullptr) {
|
|
cudaFreeHost(residency_maps_);
|
|
residency_maps_ = nullptr;
|
|
}
|
|
}
|
|
|
|
std::vector<int64_t> stats() const {
|
|
std::vector<int64_t> out(kStats);
|
|
for (int i = 0; i < kStats; ++i) out[i] = stats_[i].load();
|
|
out[kPublished] = __atomic_load_n(published_, __ATOMIC_ACQUIRE);
|
|
out[kConsumed] = consumed_seen_.load();
|
|
return out;
|
|
}
|
|
|
|
// Copy layer `layer`'s complete bank into `dst` (one per shard, `capacity`
|
|
// rows each): the slot bank first, then every non-resident expert, on
|
|
// `stream`. Writes the expert -> row map to `map_device`. The caller orders
|
|
// `stream` after the decode steps that last changed this layer.
|
|
int64_t prefill(int64_t layer, std::vector<int64_t> dst, int64_t capacity,
|
|
int64_t map_device, int64_t stream_ptr) {
|
|
TORCH_CHECK(layer >= 0 && layer < layers_ && dst.size() == kShards,
|
|
"expert cache: bad prefill request");
|
|
auto stream = reinterpret_cast<cudaStream_t>(stream_ptr);
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
check(cudaSetDevice(static_cast<int>(device_)), "set device");
|
|
for (int shard = 0; shard < kShards; ++shard) {
|
|
check(cudaMemcpyAsync(reinterpret_cast<void*>(dst[shard]),
|
|
bank(layer, shard), slots_ * bytes(layer, shard),
|
|
cudaMemcpyDeviceToDevice, stream),
|
|
"prefill slot bank");
|
|
}
|
|
int32_t* map = prefill_maps_ + layer * experts_;
|
|
int64_t cold = 0;
|
|
for (int64_t e = 0; e < experts_; ++e) {
|
|
int32_t s = slot_of_[layer * experts_ + e];
|
|
if (s >= 0) {
|
|
map[e] = s;
|
|
continue;
|
|
}
|
|
int64_t row = slots_ + cold++;
|
|
TORCH_CHECK(row < capacity, "expert cache: prefill bank too small");
|
|
int32_t h = host_slot_of_[layer * experts_ + e];
|
|
for (int shard = 0; shard < kShards; ++shard) {
|
|
check(cudaMemcpyAsync(reinterpret_cast<uint8_t*>(dst[shard]) +
|
|
row * bytes(layer, shard),
|
|
host(layer, shard, h), bytes(layer, shard),
|
|
cudaMemcpyHostToDevice, stream),
|
|
"prefill cold expert");
|
|
}
|
|
map[e] = static_cast<int32_t>(row);
|
|
}
|
|
check(cudaMemcpyAsync(reinterpret_cast<void*>(map_device), map,
|
|
experts_ * sizeof(int32_t), cudaMemcpyHostToDevice,
|
|
stream),
|
|
"prefill map");
|
|
stats_[kPrefills]++;
|
|
return cold;
|
|
}
|
|
|
|
// Copy layer `layer`'s residency to `map_device` on `stream`: int32
|
|
// [2, experts], the GPU slot of each expert (-1 when not resident), then
|
|
// its host slot (-1 when the host has no copy). The prefill transcode
|
|
// reads every expert from where this snapshot says it lives. The caller
|
|
// orders `stream` after the decode steps that last changed this layer.
|
|
// Returns the host slots that hold non-resident experts: [0, extent).
|
|
int64_t residency(int64_t layer, int64_t map_device, int64_t stream_ptr) {
|
|
TORCH_CHECK(layer >= 0 && layer < layers_, "expert cache: bad layer");
|
|
auto stream = reinterpret_cast<cudaStream_t>(stream_ptr);
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
int32_t* map = residency_maps_ + 2 * layer * experts_;
|
|
std::memcpy(map, slot_of_ + layer * experts_, experts_ * sizeof(int32_t));
|
|
std::memcpy(map + experts_, host_slot_of_ + layer * experts_,
|
|
experts_ * sizeof(int32_t));
|
|
check(cudaSetDevice(static_cast<int>(device_)), "set device");
|
|
check(cudaMemcpyAsync(reinterpret_cast<void*>(map_device), map,
|
|
2 * experts_ * sizeof(int32_t),
|
|
cudaMemcpyHostToDevice, stream),
|
|
"residency map");
|
|
stats_[kPrefills]++;
|
|
int64_t extent = 0;
|
|
for (int64_t e = 0; e < experts_; ++e)
|
|
if (map[e] < 0) extent = std::max<int64_t>(extent, map[experts_ + e] + 1);
|
|
return extent;
|
|
}
|
|
|
|
// Stage only host slots used by the most recent residency snapshot. Keep
|
|
// their offsets unchanged: the GPU source map still addresses those slots.
|
|
// The caller orders the copy after decode write-backs and the preceding
|
|
// cold transcode, exactly as for a full-arena copy.
|
|
void stage_cold(int64_t layer, int64_t shard, int64_t destination,
|
|
int64_t stream_ptr) {
|
|
stage_cold_impl(layer, shard, destination, stream_ptr, nullptr);
|
|
}
|
|
|
|
void stage_cold_selected(int64_t layer, int64_t shard, int64_t destination,
|
|
int64_t stream_ptr, torch::Tensor selected) {
|
|
TORCH_CHECK(
|
|
selected.device().is_cpu() && selected.scalar_type() == torch::kBool &&
|
|
selected.is_contiguous() && selected.dim() == 1 &&
|
|
selected.numel() == experts_,
|
|
"expert cache: selected experts must be contiguous CPU bool[E]");
|
|
stage_cold_impl(layer, shard, destination, stream_ptr,
|
|
selected.data_ptr<bool>());
|
|
}
|
|
|
|
// Eager prefill only: prior decode requests and this layer's raw readers
|
|
// must have completed. The caller owns pinned plans until their stream
|
|
// completes, then uses the ordinary GPU remap to commit the copied plan.
|
|
// This is called by the owner, never by the GPU-polled decode lane.
|
|
int64_t admit_prefill(int64_t layer, torch::Tensor ids,
|
|
torch::Tensor host_plan, torch::Tensor device_plan,
|
|
torch::Tensor temporary, int64_t stream_ptr) {
|
|
TORCH_CHECK(layer >= 0 && layer < layers_, "expert cache: bad layer");
|
|
TORCH_CHECK(ids.device().is_cpu() && ids.scalar_type() == torch::kInt32 &&
|
|
ids.dim() == 1 && ids.is_contiguous() &&
|
|
ids.numel() <= slots_ - spare_,
|
|
"expert cache: prefill ids must fit resident CPU int32 rows");
|
|
const int64_t plan_size = kReplyHeader + kOpInts * max_ops_;
|
|
TORCH_CHECK(host_plan.device().is_cpu() && host_plan.is_pinned() &&
|
|
host_plan.scalar_type() == torch::kInt32 &&
|
|
host_plan.is_contiguous() && host_plan.numel() >= plan_size,
|
|
"expert cache: prefill plan must be pinned CPU int32");
|
|
TORCH_CHECK(
|
|
device_plan.is_cuda() && device_plan.device().index() == device_ &&
|
|
device_plan.scalar_type() == torch::kInt32 &&
|
|
device_plan.is_contiguous() && device_plan.numel() >= plan_size,
|
|
"expert cache: prefill device plan has wrong storage");
|
|
int64_t required = 0;
|
|
for (int s = 0; s < kShards; ++s)
|
|
required = std::max(required, bytes(layer, s));
|
|
TORCH_CHECK(temporary.is_cuda() && temporary.device().index() == device_ &&
|
|
temporary.scalar_type() == torch::kUInt8 &&
|
|
temporary.is_contiguous() && temporary.numel() >= required,
|
|
"expert cache: prefill swap temporary is too small");
|
|
const int32_t* requested = ids.data_ptr<int32_t>();
|
|
for (int64_t i = 0; i < ids.numel(); ++i)
|
|
TORCH_CHECK(requested[i] >= 0 && requested[i] < experts_,
|
|
"expert cache: invalid prefill expert id");
|
|
const int64_t published = __atomic_load_n(published_, __ATOMIC_ACQUIRE);
|
|
while (consumed_seen_.load(std::memory_order_acquire) < published) {
|
|
TORCH_CHECK(
|
|
!stop_.load() && !__atomic_load_n(reply_ + 2, __ATOMIC_ACQUIRE),
|
|
"expert cache: prefill admission on a stopped/failed lane");
|
|
cpu_relax();
|
|
}
|
|
auto stream = reinterpret_cast<cudaStream_t>(stream_ptr);
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
TORCH_CHECK(!stop_.load() && !__atomic_load_n(reply_ + 2, __ATOMIC_ACQUIRE),
|
|
"expert cache: prefill admission on a stopped/failed lane");
|
|
check(cudaSetDevice(static_cast<int>(device_)), "set device");
|
|
const uint64_t now = ++clock_[layer];
|
|
uint64_t* last = last_use_.data() + layer * experts_;
|
|
for (int64_t i = 0; i < ids.numel(); ++i) last[requested[i]] = now;
|
|
plan_layer(layer, requested, ids.numel(), now);
|
|
const int64_t n = ops_.size() / kOpInts;
|
|
TORCH_CHECK(n <= max_ops_, "expert cache: plan exceeds reply capacity");
|
|
int32_t* plan = host_plan.data_ptr<int32_t>();
|
|
std::fill(plan, plan + kReplyHeader, 0);
|
|
plan[0] = static_cast<int32_t>(n);
|
|
std::copy(ops_.begin(), ops_.end(), plan + kReplyHeader);
|
|
int64_t expert_bytes = 0;
|
|
for (int s = 0; s < kShards; ++s) expert_bytes += bytes(layer, s);
|
|
if (temporary.numel() >= expert_bytes) {
|
|
const int64_t batch_size = std::min<int64_t>(
|
|
kPrefillSwapBatch, temporary.numel() / expert_bytes);
|
|
for (int64_t first = 0; first < n; first += batch_size) {
|
|
PrefillCopies save, writeback, admit;
|
|
int64_t offset = 0;
|
|
for (int64_t op = first; op < std::min(n, first + batch_size); ++op) {
|
|
const int32_t slot = ops_[op * kOpInts];
|
|
const int32_t load = ops_[op * kOpInts + 1];
|
|
const int32_t back = ops_[op * kOpInts + 2];
|
|
for (int s = 0; s < kShards; ++s) {
|
|
const int64_t size = bytes(layer, s);
|
|
auto* device = static_cast<uint8_t*>(bank(layer, s)) + slot * size;
|
|
auto* scratch = temporary.data_ptr<uint8_t>() + offset;
|
|
if (load >= 0 && back >= 0)
|
|
save.add(scratch, host(layer, s, load), size,
|
|
cudaMemcpyHostToDevice);
|
|
if (back >= 0)
|
|
writeback.add(host(layer, s, back), device, size,
|
|
cudaMemcpyDeviceToHost);
|
|
if (load >= 0)
|
|
admit.add(device, back >= 0 ? scratch : host(layer, s, load),
|
|
size,
|
|
back >= 0 ? cudaMemcpyDeviceToDevice
|
|
: cudaMemcpyHostToDevice);
|
|
offset += size;
|
|
}
|
|
}
|
|
save.submit(stream);
|
|
writeback.submit(stream);
|
|
admit.submit(stream);
|
|
}
|
|
} else {
|
|
for (int64_t op = 0; op < n; ++op) {
|
|
const int32_t slot = ops_[op * kOpInts];
|
|
const int32_t load = ops_[op * kOpInts + 1];
|
|
const int32_t back = ops_[op * kOpInts + 2];
|
|
for (int s = 0; s < kShards; ++s) {
|
|
const int64_t size = bytes(layer, s);
|
|
auto* device = static_cast<uint8_t*>(bank(layer, s)) + slot * size;
|
|
// Incoming host data may share its host slot with the write-back.
|
|
// Save it before overwriting that slot, exactly as _apply_plan does.
|
|
if (load >= 0 && back >= 0)
|
|
check(cudaMemcpyAsync(temporary.data_ptr(), host(layer, s, load),
|
|
size, cudaMemcpyHostToDevice, stream),
|
|
"save incoming");
|
|
if (back >= 0)
|
|
check(cudaMemcpyAsync(host(layer, s, back), device, size,
|
|
cudaMemcpyDeviceToHost, stream),
|
|
"write back");
|
|
if (load >= 0)
|
|
check(cudaMemcpyAsync(
|
|
device,
|
|
back >= 0 ? temporary.data_ptr() : host(layer, s, load),
|
|
size,
|
|
back >= 0 ? cudaMemcpyDeviceToDevice
|
|
: cudaMemcpyHostToDevice,
|
|
stream),
|
|
"admit expert");
|
|
}
|
|
}
|
|
}
|
|
check(cudaMemcpyAsync(device_plan.data_ptr(), plan,
|
|
(kReplyHeader + n * kOpInts) * sizeof(int32_t),
|
|
cudaMemcpyHostToDevice, stream),
|
|
"copy prefill plan");
|
|
return n;
|
|
}
|
|
|
|
private:
|
|
void stage_cold_impl(int64_t layer, int64_t shard, int64_t destination,
|
|
int64_t stream_ptr, const bool* selected) {
|
|
TORCH_CHECK(layer >= 0 && layer < layers_ && shard >= 0 && shard < kShards,
|
|
"expert cache: bad cold stage bank");
|
|
auto stream = reinterpret_cast<cudaStream_t>(stream_ptr);
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
const int32_t* map = residency_maps_ + 2 * layer * experts_;
|
|
std::vector<uint8_t> occupied(host_slots_, 0);
|
|
for (int64_t e = 0; e < experts_; ++e) {
|
|
if (selected != nullptr && !selected[e]) continue;
|
|
if (map[e] >= 0) continue;
|
|
const int32_t h = map[experts_ + e];
|
|
TORCH_CHECK(h >= 0 && h < host_slots_,
|
|
"expert cache: non-resident expert has no host slot");
|
|
occupied[h] = 1;
|
|
}
|
|
check(cudaSetDevice(static_cast<int>(device_)), "set device");
|
|
const int64_t stride = bytes(layer, shard);
|
|
PrefillCopies copies;
|
|
for (int64_t begin = 0; begin < host_slots_;) {
|
|
if (!occupied[begin]) {
|
|
++begin;
|
|
continue;
|
|
}
|
|
int64_t end = begin + 1;
|
|
while (end < host_slots_ && occupied[end]) ++end;
|
|
copies.add(reinterpret_cast<uint8_t*>(destination) + begin * stride,
|
|
host(layer, shard, begin), (end - begin) * stride,
|
|
cudaMemcpyHostToDevice);
|
|
if (copies.full()) {
|
|
copies.submit(stream);
|
|
copies = PrefillCopies{};
|
|
}
|
|
begin = end;
|
|
}
|
|
copies.submit(stream);
|
|
}
|
|
|
|
int64_t bytes(int64_t layer, int shard) const {
|
|
return slot_bytes_[layer * kShards + shard];
|
|
}
|
|
void* bank(int64_t layer, int shard) const {
|
|
return reinterpret_cast<void*>(device_banks_[layer * kShards + shard]);
|
|
}
|
|
uint8_t* host(int64_t layer, int shard, int64_t h) const {
|
|
return reinterpret_cast<uint8_t*>(host_banks_[layer * kShards + shard]) +
|
|
h * bytes(layer, shard);
|
|
}
|
|
|
|
void run() {
|
|
try {
|
|
int64_t idle_since = now_ns();
|
|
while (!stop_.load(std::memory_order_acquire)) {
|
|
int64_t published = __atomic_load_n(published_, __ATOMIC_ACQUIRE);
|
|
if (published == consumed_) {
|
|
if (now_ns() - idle_since > 5'000'000) {
|
|
std::this_thread::sleep_for(std::chrono::microseconds(50));
|
|
} else {
|
|
cpu_relax();
|
|
}
|
|
continue;
|
|
}
|
|
if (published - consumed_ > ring_entries_ - 2) {
|
|
// The model stream never runs ahead of an unserved miss, so only
|
|
// recency information is lost by skipping old entries.
|
|
stats_[kSkipped] += published - 1 - consumed_;
|
|
consumed_ = published - 1;
|
|
}
|
|
while (consumed_ < published) {
|
|
try {
|
|
serve(consumed_);
|
|
} catch (const std::exception& error) {
|
|
fail(consumed_, error.what());
|
|
}
|
|
++consumed_;
|
|
}
|
|
consumed_seen_.store(consumed_, std::memory_order_release);
|
|
idle_since = now_ns();
|
|
}
|
|
} catch (const std::exception& error) {
|
|
stats_[kErrors]++;
|
|
fprintf(stderr, "expert cache lane stopped: %s\n", error.what());
|
|
}
|
|
}
|
|
|
|
// A request that cannot be planned still has to release the model stream:
|
|
// answer with no copies and a sticky failure flag, which every later wait
|
|
// also honours. Residency state may be inconsistent after this.
|
|
void fail(int64_t seq, const char* what) {
|
|
stats_[kErrors]++;
|
|
fprintf(stderr, "expert cache lane: entry %lld failed: %s\n",
|
|
static_cast<long long>(seq), what);
|
|
__atomic_store_n(reply_ + 2, 1, __ATOMIC_RELEASE);
|
|
__atomic_store_n(reply_ + 1, 0, __ATOMIC_RELEASE);
|
|
__atomic_store_n(reply_, static_cast<int32_t>(seq), __ATOMIC_RELEASE);
|
|
}
|
|
|
|
void serve(int64_t seq) {
|
|
const int32_t* entry = ring_ + (seq % ring_entries_) * entry_ints_;
|
|
// The entry's tag is released after its body; wait for it rather than
|
|
// dropping a request the model stream is blocked on.
|
|
// A hit is tagged ~seq and published without fences: its body may still
|
|
// be in flight, which only affects recency, and it never needs a plan.
|
|
const int32_t want = static_cast<int32_t>(seq), hit = ~want;
|
|
int32_t tag = __atomic_load_n(entry, __ATOMIC_ACQUIRE);
|
|
for (int64_t deadline = now_ns() + 1'000'000'000;
|
|
tag != want && tag != hit && now_ns() < deadline;) {
|
|
cpu_relax();
|
|
tag = __atomic_load_n(entry, __ATOMIC_ACQUIRE);
|
|
}
|
|
int32_t layer = __atomic_load_n(entry + 1, __ATOMIC_ACQUIRE);
|
|
int32_t count = __atomic_load_n(entry + 2, __ATOMIC_ACQUIRE);
|
|
const int32_t miss =
|
|
tag == want ? __atomic_load_n(entry + 3, __ATOMIC_ACQUIRE) : 0;
|
|
if (tag == hit) {
|
|
// Only recency uses a hit's body; ignore a stale header.
|
|
if (layer < 0 || layer >= layers_) return;
|
|
count = std::max<int32_t>(
|
|
0, std::min<int32_t>(count, entry_ints_ - kEntryHeader));
|
|
}
|
|
if ((tag != want && tag != hit) || layer < 0 || layer >= layers_ ||
|
|
count < 0 || count > entry_ints_ - kEntryHeader) {
|
|
char what[96];
|
|
snprintf(what, sizeof(what), "bad entry (tag %d layer %d count %d)", tag,
|
|
layer, count);
|
|
throw std::runtime_error(what);
|
|
}
|
|
const int32_t* ids = entry + kEntryHeader;
|
|
stats_[kEntries]++;
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
uint64_t now = ++clock_[layer];
|
|
uint64_t* last = last_use_.data() + layer * experts_;
|
|
for (int32_t i = 0; i < count; ++i) {
|
|
int32_t e = ids[i];
|
|
if (e >= 0 && e < experts_) last[e] = now;
|
|
}
|
|
if (!miss) return;
|
|
if (__atomic_load_n(reply_ + 2, __ATOMIC_ACQUIRE))
|
|
throw std::runtime_error("lane already failed");
|
|
int64_t start = now_ns();
|
|
plan_layer(layer, ids, count, now);
|
|
int64_t n = static_cast<int64_t>(ops_.size()) / kOpInts;
|
|
TORCH_CHECK(n <= max_ops_, "expert cache: plan exceeds reply capacity");
|
|
std::memcpy(reply_ + kReplyHeader, ops_.data(),
|
|
ops_.size() * sizeof(int32_t));
|
|
__atomic_store_n(reply_ + 1, static_cast<int32_t>(n), __ATOMIC_RELEASE);
|
|
__atomic_store_n(reply_, static_cast<int32_t>(seq), __ATOMIC_RELEASE);
|
|
stats_[kRequests]++;
|
|
int64_t served = now_ns() - start;
|
|
stats_[kServeNs] += served;
|
|
if (served > stats_[kMaxServeNs].load()) stats_[kMaxServeNs].store(served);
|
|
}
|
|
|
|
void plan_layer(int32_t layer, const int32_t* ids, int64_t count,
|
|
uint64_t now) {
|
|
int32_t* slot_of = slot_of_ + layer * experts_;
|
|
ops_.clear();
|
|
for (int32_t i = 0; i < count; ++i) {
|
|
int32_t e = ids[i];
|
|
if (e < 0 || e >= experts_ || slot_of[e] >= 0) continue;
|
|
int32_t s, writeback = -1, victim = -1;
|
|
int32_t load = host_slot_of_[layer * experts_ + e];
|
|
if (free_slots_[layer].empty()) {
|
|
s = evict(layer, now, &writeback, &victim, e);
|
|
TORCH_CHECK(s >= 0, "expert cache: no evictable slot");
|
|
} else {
|
|
s = free_slots_[layer].back();
|
|
free_slots_[layer].pop_back();
|
|
}
|
|
ops_.insert(ops_.end(), {s, load, writeback, e, victim});
|
|
slot_of[e] = s;
|
|
expert_of_slot_[layer * slots_ + s] = e;
|
|
stats_[kLoaded]++;
|
|
}
|
|
// Restore the spare slots in the same plan; these slots are not loaded
|
|
// before a later request, so their write-backs need no ordering here.
|
|
// Best effort: a later request retries when the host arena is short.
|
|
while (static_cast<int64_t>(free_slots_[layer].size()) < spare_) {
|
|
int32_t writeback = -1, victim = -1;
|
|
int32_t s = evict(layer, now, &writeback, &victim, -1);
|
|
if (s < 0) break;
|
|
ops_.insert(ops_.end(), {s, -1, writeback, -1, victim});
|
|
free_slots_[layer].push_back(s);
|
|
}
|
|
}
|
|
|
|
// Evict the least recently used expert not routed by the current entry and
|
|
// return its slot (-1 when nothing can be evicted). Reports the host slot it
|
|
// must be written back to (-1 when the host keeps a copy) and the evicted
|
|
// expert. With the host arena full, the victim swaps places with `incoming`
|
|
// (the expert loaded into its slot by the same op): the copy kernel reads
|
|
// both sides of each word before writing either.
|
|
int32_t evict(int32_t layer, uint64_t now, int32_t* writeback,
|
|
int32_t* victim_out, int32_t incoming) {
|
|
const uint64_t* last = last_use_.data() + layer * experts_;
|
|
int32_t* owners = expert_of_slot_.data() + layer * slots_;
|
|
int32_t victim_slot = -1;
|
|
uint64_t best = UINT64_MAX;
|
|
for (int64_t s = 0; s < slots_; ++s) {
|
|
int32_t e = owners[s];
|
|
if (e < 0 || last[e] >= now) continue;
|
|
if (last[e] < best) {
|
|
best = last[e];
|
|
victim_slot = static_cast<int32_t>(s);
|
|
}
|
|
}
|
|
if (victim_slot < 0) return -1;
|
|
int32_t victim = owners[victim_slot];
|
|
int32_t* host_slot = host_slot_of_ + layer * experts_;
|
|
*writeback = -1;
|
|
if (host_slot[victim] < 0) {
|
|
int32_t h = take_host_slot(layer, now);
|
|
if (h < 0 && incoming >= 0 && host_slot[incoming] >= 0) {
|
|
h = host_slot[incoming];
|
|
host_slot[incoming] = -1;
|
|
}
|
|
if (h < 0) return -1;
|
|
host_slot[victim] = h;
|
|
host_owner_[layer * host_slots_ + h] = victim;
|
|
*writeback = h;
|
|
stats_[kWritebacks]++;
|
|
}
|
|
slot_of_[layer * experts_ + victim] = -1;
|
|
owners[victim_slot] = -1;
|
|
*victim_out = victim;
|
|
stats_[kEvicted]++;
|
|
return victim_slot;
|
|
}
|
|
|
|
// A free host slot, or the host copy of the most recently used resident
|
|
// expert (the least likely to be evicted next), or -1. Copies of experts
|
|
// routed by the current entry are never taken: this plan may still load
|
|
// them.
|
|
int32_t take_host_slot(int32_t layer, uint64_t now) {
|
|
if (!free_host_[layer].empty()) {
|
|
int32_t h = free_host_[layer].back();
|
|
free_host_[layer].pop_back();
|
|
return h;
|
|
}
|
|
const uint64_t* last = last_use_.data() + layer * experts_;
|
|
int32_t* owners = host_owner_.data() + layer * host_slots_;
|
|
int32_t chosen = -1;
|
|
uint64_t best = 0;
|
|
for (int64_t h = 0; h < host_slots_; ++h) {
|
|
int32_t e = owners[h];
|
|
if (e < 0 || slot_of_[layer * experts_ + e] < 0 || last[e] >= now)
|
|
continue;
|
|
if (chosen < 0 || last[e] >= best) {
|
|
best = last[e];
|
|
chosen = static_cast<int32_t>(h);
|
|
}
|
|
}
|
|
if (chosen < 0) return -1;
|
|
host_slot_of_[layer * experts_ + owners[chosen]] = -1;
|
|
owners[chosen] = -1;
|
|
return chosen;
|
|
}
|
|
|
|
int64_t device_, layers_, experts_, slots_, host_slots_, spare_, max_ops_;
|
|
std::vector<int64_t> slot_bytes_, device_banks_, host_banks_;
|
|
int32_t* ring_;
|
|
int64_t ring_entries_, entry_ints_;
|
|
int64_t* published_;
|
|
int32_t* reply_;
|
|
int32_t* slot_of_;
|
|
int32_t* host_slot_of_;
|
|
std::vector<int32_t> expert_of_slot_, host_owner_, ops_;
|
|
std::vector<uint64_t> last_use_, clock_;
|
|
std::vector<std::vector<int32_t>> free_slots_, free_host_;
|
|
int32_t* prefill_maps_ = nullptr;
|
|
int32_t* residency_maps_ = nullptr;
|
|
int64_t consumed_ = 0;
|
|
std::atomic<int64_t> consumed_seen_{0};
|
|
std::atomic<bool> stop_{false};
|
|
std::atomic<int64_t> stats_[kStats];
|
|
std::mutex mutex_;
|
|
std::thread thread_;
|
|
};
|
|
|
|
int64_t host_register(int64_t pointer, int64_t bytes) {
|
|
check(cudaHostRegister(reinterpret_cast<void*>(pointer), bytes,
|
|
cudaHostRegisterPortable | cudaHostRegisterMapped),
|
|
"register host arena");
|
|
return pointer;
|
|
}
|
|
|
|
void host_unregister(int64_t pointer) {
|
|
cudaHostUnregister(reinterpret_cast<void*>(pointer));
|
|
}
|
|
|
|
// Device view of registered host memory (no copy). vLLM's
|
|
// get_cuda_view_from_cpu_tensor copies tensors it does not see as pinned,
|
|
// which includes cudaHostRegister'ed memory.
|
|
torch::Tensor mapped_view(torch::Tensor host, int64_t device) {
|
|
void* pointer = nullptr;
|
|
check(cudaHostGetDevicePointer(&pointer, host.data_ptr(), 0),
|
|
"map host arena");
|
|
return torch::from_blob(
|
|
pointer, host.sizes(), host.strides(), [host](void*) {},
|
|
host.options().device(torch::kCUDA, static_cast<int>(device)));
|
|
}
|
|
|
|
// Asynchronous host -> device copy by the copy engine. `src` may be
|
|
// cudaHostRegister'ed memory, which torch does not see as pinned (its copy
|
|
// would stage synchronously).
|
|
void copy_h2d(int64_t dst, int64_t src, int64_t bytes, int64_t stream_ptr) {
|
|
check(cudaMemcpyAsync(reinterpret_cast<void*>(dst),
|
|
reinterpret_cast<const void*>(src), bytes,
|
|
cudaMemcpyHostToDevice,
|
|
reinterpret_cast<cudaStream_t>(stream_ptr)),
|
|
"copy host to device");
|
|
}
|
|
|
|
// Batched strided host->device copies (one Python call per prefill tile).
|
|
void copy_h2d_2d(const std::vector<int64_t>& dst,
|
|
const std::vector<int64_t>& src,
|
|
const std::vector<int64_t>& dpitch,
|
|
const std::vector<int64_t>& spitch,
|
|
const std::vector<int64_t>& width,
|
|
const std::vector<int64_t>& height, int64_t stream_ptr) {
|
|
const size_t n = dst.size();
|
|
if (src.size() != n || dpitch.size() != n || spitch.size() != n ||
|
|
width.size() != n || height.size() != n)
|
|
throw std::invalid_argument("copy_h2d_2d: mismatched descriptor lists");
|
|
auto stream = reinterpret_cast<cudaStream_t>(stream_ptr);
|
|
for (size_t i = 0; i < n; ++i)
|
|
check(
|
|
cudaMemcpy2DAsync(reinterpret_cast<void*>(dst[i]), dpitch[i],
|
|
reinterpret_cast<const void*>(src[i]), spitch[i],
|
|
width[i], height[i], cudaMemcpyHostToDevice, stream),
|
|
"copy host to device 2d");
|
|
}
|
|
|
|
} // namespace
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
using namespace vllm::gguf;
|
|
m.def("shared_engram_control_layout", [](uint64_t maximum) {
|
|
return std::vector<uint64_t>{SharedEngramReader::control_words(maximum),
|
|
SharedEngramReader::channel_words(maximum),
|
|
SharedEngramReader::rank_words(maximum)};
|
|
});
|
|
pybind11::class_<SharedEngramTable>(m, "SharedEngramTable")
|
|
.def(pybind11::init([](int fd, uint64_t offset, uint64_t rows) {
|
|
return SharedEngramTable{fd, offset, rows};
|
|
}));
|
|
pybind11::class_<SharedEngramReader>(m, "SharedEngramReader")
|
|
.def(pybind11::init([](uint64_t data, uint64_t bytes, uint64_t control,
|
|
uint64_t maximum, uint64_t row_bytes,
|
|
std::vector<SharedEngramTable> tables,
|
|
uint64_t timeout_ns) {
|
|
return std::make_unique<SharedEngramReader>(
|
|
reinterpret_cast<uint8_t*>(data), bytes,
|
|
reinterpret_cast<uint64_t*>(control), maximum, row_bytes,
|
|
std::move(tables), timeout_ns);
|
|
}))
|
|
.def("stop", &SharedEngramReader::stop,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>())
|
|
.def("error", &SharedEngramReader::error);
|
|
pybind11::class_<SharedExpertDmaLane>(m, "SharedExpertDmaLane")
|
|
.def(pybind11::init<uint64_t, uint64_t, uint64_t, int, unsigned, unsigned,
|
|
std::vector<std::vector<int64_t>>, uint64_t, unsigned,
|
|
unsigned, unsigned, unsigned, unsigned, uint64_t>())
|
|
.def("stop", &SharedExpertDmaLane::stop,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>())
|
|
.def("error", &SharedExpertDmaLane::error)
|
|
.def("completed", &SharedExpertDmaLane::completed);
|
|
pybind11::class_<SharedExpertJobLane>(m, "SharedExpertJobLane")
|
|
.def(pybind11::init<uint64_t, uint64_t, int, unsigned,
|
|
std::vector<std::vector<int64_t>>,
|
|
std::vector<std::vector<int64_t>>, uint64_t>())
|
|
.def("stop", &SharedExpertJobLane::stop,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>())
|
|
.def("error", &SharedExpertJobLane::error)
|
|
.def("completed", &SharedExpertJobLane::completed);
|
|
m.def("shared_expert_hint_layout", [] {
|
|
return std::vector<uint64_t>{kSharedHintStart, kSharedHintHeader,
|
|
kSharedHintEntries, kSharedHintWords,
|
|
kSharedHintTokens, kSharedHintCands};
|
|
});
|
|
m.def("shared_expert_job_layout", [] {
|
|
return std::vector<uint64_t>{kSharedJobStart, kSharedJobHeader,
|
|
kSharedJobEntries, kSharedJobWords};
|
|
});
|
|
m.def("shared_expert_control_layout", [] {
|
|
return std::vector<uint64_t>{kSharedControlWords, kSharedRankStart,
|
|
kSharedRankWords, kSharedOpStart,
|
|
kSharedOpWords, kSharedMaxCopies};
|
|
});
|
|
m.def("shared_expert_hit_layout", [] {
|
|
return std::vector<uint64_t>{kSharedHitStart, kSharedHitHeader,
|
|
kSharedHitEntries, kSharedHitWords,
|
|
kSharedTotalWords};
|
|
});
|
|
pybind11::class_<SharedExpertLayer>(m, "SharedExpertLayer")
|
|
.def(
|
|
pybind11::init([](uint64_t bytes, int32_t experts, int32_t resident) {
|
|
return SharedExpertLayer{bytes, experts, resident};
|
|
}));
|
|
pybind11::class_<SharedExpertCopy>(m, "SharedExpertCopy")
|
|
.def_readonly("slot", &SharedExpertCopy::slot)
|
|
.def_readonly("incoming", &SharedExpertCopy::incoming)
|
|
.def_readonly("victim", &SharedExpertCopy::victim)
|
|
.def_readonly("load_offset", &SharedExpertCopy::load_offset)
|
|
.def_readonly("save_offset", &SharedExpertCopy::save_offset);
|
|
pybind11::class_<SharedExpertPlanner>(m, "SharedExpertPlanner")
|
|
.def(
|
|
pybind11::init([](uint64_t pointer, uint64_t bytes,
|
|
std::vector<SharedExpertLayer> layers,
|
|
unsigned ranks, unsigned routes, uint64_t reserve) {
|
|
return std::make_unique<SharedExpertPlanner>(
|
|
reinterpret_cast<uint8_t*>(pointer), bytes, std::move(layers),
|
|
ranks, routes, reserve);
|
|
}))
|
|
.def("plan", &SharedExpertPlanner::plan,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>())
|
|
.def("acknowledge", &SharedExpertPlanner::acknowledge)
|
|
.def("ready", &SharedExpertPlanner::ready)
|
|
.def("finish", &SharedExpertPlanner::finish,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>())
|
|
.def("slot", &SharedExpertPlanner::slot)
|
|
.def("cold_offset", &SharedExpertPlanner::cold_offset)
|
|
.def("acquire", &SharedExpertPlanner::acquire)
|
|
.def("release", &SharedExpertPlanner::release)
|
|
.def("stop", &SharedExpertPlanner::stop);
|
|
pybind11::class_<SharedExpertCoordinator>(m, "SharedExpertCoordinator")
|
|
.def(pybind11::init([](uint64_t arena, uint64_t bytes, uint64_t control,
|
|
std::vector<SharedExpertLayer> layers,
|
|
unsigned ranks, unsigned routes, uint64_t reserve,
|
|
uint64_t timeout) {
|
|
return std::make_unique<SharedExpertCoordinator>(
|
|
reinterpret_cast<uint8_t*>(arena), bytes,
|
|
reinterpret_cast<uint64_t*>(control), std::move(layers), ranks,
|
|
routes, reserve, timeout);
|
|
}))
|
|
.def("stop", &SharedExpertCoordinator::stop,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>())
|
|
.def("error", &SharedExpertCoordinator::error)
|
|
.def("slot", &SharedExpertCoordinator::slot)
|
|
.def("cold_offset", &SharedExpertCoordinator::cold_offset)
|
|
.def("acquire", &SharedExpertCoordinator::acquire)
|
|
.def("release", &SharedExpertCoordinator::release);
|
|
pybind11::class_<ExpertCacheLane>(m, "ExpertCacheLane")
|
|
.def(pybind11::init<int64_t, int64_t, int64_t, int64_t, int64_t, int64_t,
|
|
std::vector<int64_t>, std::vector<int64_t>,
|
|
std::vector<int64_t>, int64_t, int64_t, int64_t,
|
|
int64_t, int64_t, int64_t, torch::Tensor,
|
|
torch::Tensor>())
|
|
.def("stop", &ExpertCacheLane::stop,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>())
|
|
.def("stats", &ExpertCacheLane::stats)
|
|
.def("prefill", &ExpertCacheLane::prefill,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>())
|
|
.def("residency", &ExpertCacheLane::residency,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>())
|
|
.def("stage_cold", &ExpertCacheLane::stage_cold,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>())
|
|
.def("stage_cold_selected", &ExpertCacheLane::stage_cold_selected,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>())
|
|
.def("admit_prefill", &ExpertCacheLane::admit_prefill,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>());
|
|
m.def("host_register", &host_register);
|
|
m.def("host_unregister", &host_unregister);
|
|
m.def("mapped_view", &mapped_view);
|
|
m.def("copy_h2d", ©_h2d);
|
|
m.def("copy_h2d_2d", ©_h2d_2d,
|
|
pybind11::call_guard<pybind11::gil_scoped_release>());
|
|
}
|