iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[3/N] HiSparse: host-resident sparse-MLA decode hot-buffering (#53781)

Browse Source 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Signed-off-by: Matthew Bonanni <matthew@neuralmagic.com>
Signed-off-by: Nicolò Lucchesi <nicolo.lucchesi@mistral.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@mistral.ai>
Co-authored-by: OpenAI Codex <noreply@openai.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
This commit is contained in:
Matthew Bonanni
2026-09-12 00:36:47 -04:00
committed by GitHub
co-authored by OpenAI Codex Lucas Wilkinson Lucas Wilkinson Nicolò Lucchesi OpenAI Codex Claude Opus 5 Lucas Wilkinson
parent 120ec4ebd2
commit d43bb2f37f
124 changed files with 14262 additions and 906 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+1
View File
@@ -476,6 +476,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
if(VLLM_GPU_LANG STREQUAL "CUDA")
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/hisparse_kernels.cu"
"csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu"
"csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu"
"csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu")
csrc/libtorch_stable/cache_kernels.cu
+70 -5
View File
@@ -1458,7 +1458,11 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache_page(
const int32_t block_size, const int32_t total_tokens,
const int64_t block_table_stride, const int64_t cache_block_stride,
const int64_t cache_entry_stride, const int64_t dst_entry_stride,
const int32_t* __restrict__ seq_starts) {
const int32_t* __restrict__ seq_starts,
const uint8_t* __restrict__ host_cache,
const int32_t* __restrict__ host_row_ids,
const int32_t* __restrict__ device_row_ids,
const int64_t host_entry_stride) {
constexpr int32_t warps_per_cta = 16;
__shared__ GatherPageTask page;
__shared__ int32_t physical_block;
@@ -1484,8 +1488,22 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache_page(
page_token < page.page_token_end; page_token += warps_per_cta) {
const int32_t output_token =
page.output_token_begin + page_token - page.page_token_begin;
const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride +
page_token * cache_entry_stride;
const int32_t source_row = physical_block * block_size + page_token;
const uint8_t* token_ptr;
if (device_row_ids == nullptr) {
token_ptr = src_cache + physical_block * cache_block_stride +
page_token * cache_entry_stride;
} else {
const int32_t device_row = device_row_ids[source_row];
if (device_row >= 0) {
token_ptr = src_cache + (device_row / block_size) * cache_block_stride +
(device_row % block_size) * cache_entry_stride;
} else {
token_ptr =
host_cache +
static_cast<int64_t>(host_row_ids[source_row]) * host_entry_stride;
}
}
__nv_bfloat16* dst_ptr = dst + output_token * dst_entry_stride;
gather_and_upconvert_fp8_token(token_ptr, dst_ptr, lane_id);
}
@@ -1694,7 +1712,10 @@ void cp_gather_and_upconvert_fp8_kv_cache(
torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
torch::stable::Tensor const& workspace_starts, // [BATCH]
int64_t batch_size,
std::optional<torch::stable::Tensor> seq_starts = std::nullopt) {
std::optional<torch::stable::Tensor> seq_starts = std::nullopt,
std::optional<torch::stable::Tensor> host_cache = std::nullopt,
std::optional<torch::stable::Tensor> host_row_ids = std::nullopt,
std::optional<torch::stable::Tensor> device_row_ids = std::nullopt) {
torch::stable::accelerator::DeviceGuard device_guard(
src_cache.get_device_index());
const cudaStream_t stream = get_current_cuda_stream();
@@ -1714,6 +1735,12 @@ void cp_gather_and_upconvert_fp8_kv_cache(
"seq_starts must be int32");
}
const bool has_host_rows = host_cache.has_value();
STD_TORCH_CHECK(
has_host_rows == host_row_ids.has_value() &&
has_host_rows == device_row_ids.has_value(),
"host_cache, host_row_ids, and device_row_ids must be provided together");
STD_TORCH_CHECK(src_cache.device() == dst.device(),
"src_cache and dst must be on the same device");
STD_TORCH_CHECK(src_cache.device() == block_table.device(),
@@ -1724,6 +1751,43 @@ void cp_gather_and_upconvert_fp8_kv_cache(
STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
"src_cache and seq_starts must be on the same device");
}
const uint8_t* host_cache_ptr = nullptr;
const int32_t* host_row_ids_ptr = nullptr;
const int32_t* device_row_ids_ptr = nullptr;
int64_t host_entry_stride = 0;
if (has_host_rows) {
auto const& host = host_cache.value();
auto const& host_rows = host_row_ids.value();
auto const& device_rows = device_row_ids.value();
STD_TORCH_CHECK(
host.device().is_cpu() &&
host.scalar_type() == torch::headeronly::ScalarType::Byte &&
host.dim() == 2 && host.is_contiguous() &&
host.size(1) == src_cache.size(2),
"host_cache must be a contiguous uint8 CPU row matrix matching "
"src_cache");
cudaPointerAttributes attributes{};
const cudaError_t pointer_status =
cudaPointerGetAttributes(&attributes, host.const_data_ptr());
if (pointer_status != cudaSuccess) {
cudaGetLastError();
}
STD_TORCH_CHECK(
pointer_status == cudaSuccess && attributes.type == cudaMemoryTypeHost,
"host_cache must be CUDA-accessible pinned memory");
STD_TORCH_CHECK(
host_rows.is_cuda() && device_rows.is_cuda() &&
host_rows.scalar_type() == torch::headeronly::ScalarType::Int &&
device_rows.scalar_type() == torch::headeronly::ScalarType::Int &&
host_rows.is_contiguous() && device_rows.is_contiguous() &&
host_rows.numel() == device_rows.numel() && host_rows.dim() <= 2 &&
device_rows.dim() <= 2,
"host/device row maps must be matching contiguous CUDA int32 tensors");
host_cache_ptr = reinterpret_cast<const uint8_t*>(host.const_data_ptr());
host_row_ids_ptr = host_rows.const_data_ptr<int32_t>();
device_row_ids_ptr = device_rows.const_data_ptr<int32_t>();
host_entry_stride = host.stride(0);
}
auto dtype = src_cache.scalar_type();
STD_TORCH_CHECK(
dtype == torch::headeronly::ScalarType::Byte || // uint8
@@ -1766,7 +1830,8 @@ void cp_gather_and_upconvert_fp8_kv_cache(
workspace_starts.const_data_ptr<int32_t>(),
static_cast<int32_t>(batch_size), block_size, total_tokens,
block_table_stride, cache_block_stride, cache_entry_stride,
dst_entry_stride, seq_starts_ptr);
dst_entry_stride, seq_starts_ptr, host_cache_ptr, host_row_ids_ptr,
device_row_ids_ptr, host_entry_stride);
}
void cp_gather_and_upconvert_nvfp4_kv_cache(
csrc/libtorch_stable/hisparse_kernels.cu
+1226
View File
File diff suppressed because it is too large Load Diff
csrc/libtorch_stable/ops.h
+52 -1
View File
@@ -661,6 +661,54 @@ void concat_and_cache_mla_grouped(
std::optional<torch::stable::Tensor> kv_scales,
const std::string& kv_cache_dtype);
#ifndef USE_ROCM
// HiSparse kernels use raw PTX in the row copy; CUDA-only.
void hisparse_resolve_residency(
torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
torch::stable::Tensor const& hot_block_table,
torch::stable::Tensor const& global_indices,
torch::stable::Tensor& hot_indices,
torch::stable::Tensor& device_global_indices,
torch::stable::Tensor& lru_slots,
std::optional<torch::stable::Tensor> const& request_state_indices,
int64_t region_stride,
std::optional<torch::stable::Tensor> const& miss_mask,
std::optional<torch::stable::Tensor> const& attention_indices,
int64_t attention_block_stride,
std::optional<torch::stable::Tensor> const& request_ids,
std::optional<torch::stable::Tensor> const& source_block_table,
int64_t source_block_size,
std::optional<torch::stable::Tensor> const& resolved_global_indices,
std::optional<torch::stable::Tensor> const& valid_counts,
std::optional<torch::stable::Tensor> const& swap_host_physical_rows,
std::optional<torch::stable::Tensor> const& swap_device_physical_rows,
std::optional<torch::stable::Tensor> const& swap_counts,
std::optional<torch::stable::Tensor> const& resident_block_table,
int64_t resident_block_size, int64_t resident_null_block);
void hisparse_invalidate_written_slots(
torch::stable::Tensor& device_global_indices,
torch::stable::Tensor const& request_state_indices,
torch::stable::Tensor const& req_id_per_token,
torch::stable::Tensor const& written_slots);
void hisparse_gather_plan(
torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
torch::stable::Tensor const& global_indices,
torch::stable::Tensor const& hot_indices,
torch::stable::Tensor const& miss_mask,
std::optional<torch::stable::Tensor> const& request_state_indices,
std::optional<torch::stable::Tensor> const& attention_indices,
int64_t attention_block_stride);
void hisparse_gather_compact(torch::stable::Tensor const& host_cache,
torch::stable::Tensor& hot_cache,
torch::stable::Tensor const& miss_global_indices,
torch::stable::Tensor const& miss_hot_indices,
torch::stable::Tensor const& miss_counts);
#endif // !USE_ROCM
// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
void concat_and_cache_mla_rope_fused(
torch::stable::Tensor& positions, torch::stable::Tensor& q_pe,
@@ -704,7 +752,10 @@ void cp_gather_and_upconvert_fp8_kv_cache(
torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
torch::stable::Tensor const& workspace_starts, // [BATCH]
int64_t batch_size,
std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
std::optional<torch::stable::Tensor> seq_starts = std::nullopt,
std::optional<torch::stable::Tensor> host_cache = std::nullopt,
std::optional<torch::stable::Tensor> host_row_ids = std::nullopt,
std::optional<torch::stable::Tensor> device_row_ids = std::nullopt);
// Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace
void cp_gather_and_upconvert_nvfp4_kv_cache(
csrc/libtorch_stable/torch_bindings.cpp
+62 -1
View File
@@ -987,6 +987,57 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
" Tensor? kv_scales=None,"
" str kv_cache_dtype='auto') -> ()");
#ifndef USE_ROCM
ops.def(
"hisparse_resolve_residency(Tensor host_cache,"
" Tensor! hot_cache,"
" Tensor hot_block_table,"
" Tensor global_indices,"
" Tensor! hot_indices,"
" Tensor! device_global_indices,"
" Tensor! lru_slots,"
" Tensor? request_state_indices,"
" int region_stride,"
" Tensor(a!)? miss_mask=None,"
" Tensor(b!)? attention_indices=None,"
" int attention_block_stride=0,"
" Tensor? request_ids=None,"
" Tensor? source_block_table=None,"
" int source_block_size=0,"
" Tensor(c!)? resolved_global_indices=None,"
" Tensor(d!)? valid_counts=None,"
" Tensor(e!)? swap_host_physical_rows=None,"
" Tensor(f!)? swap_device_physical_rows=None,"
" Tensor(g!)? swap_counts=None,"
" Tensor? resident_block_table=None,"
" int resident_block_size=0,"
" int resident_null_block=0) -> ()");
ops.def(
"hisparse_invalidate_written_slots(Tensor! device_global_indices,"
" Tensor request_state_indices,"
" Tensor req_id_per_token,"
" Tensor written_slots) -> ()");
ops.def(
"hisparse_gather_plan(Tensor host_cache,"
" Tensor! hot_cache,"
" Tensor global_indices,"
" Tensor hot_indices,"
" Tensor miss_mask,"
" Tensor? request_state_indices,"
" Tensor(a!)? attention_indices=None,"
" int attention_block_stride=0) -> ()");
ops.def(
"hisparse_gather_compact(Tensor host_cache,"
" Tensor! hot_cache,"
" Tensor miss_global_indices,"
" Tensor miss_hot_indices,"
" Tensor miss_counts) -> ()");
#endif // !USE_ROCM
// Rotate Q and K, then write to kv cache for MLA
ops.def(
"concat_and_cache_mla_rope_fused("
@@ -1023,7 +1074,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
ops.def(
"cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
"Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? "
"seq_starts) -> ()");
"seq_starts, Tensor? host_cache=None, Tensor? host_row_ids=None, Tensor? "
"device_row_ids=None) -> ()");
ops.def(
"cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, "
@@ -1091,6 +1143,15 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) {
ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla));
ops.impl("concat_and_cache_mla_grouped",
TORCH_BOX(&concat_and_cache_mla_grouped));
#ifndef USE_ROCM
ops.impl("hisparse_resolve_residency",
TORCH_BOX(&hisparse_resolve_residency));
ops.impl("hisparse_invalidate_written_slots",
TORCH_BOX(&hisparse_invalidate_written_slots));
ops.impl("hisparse_gather_plan", TORCH_BOX(&hisparse_gather_plan));
ops.impl("hisparse_gather_compact", TORCH_BOX(&hisparse_gather_compact));
#endif // !USE_ROCM
ops.impl("concat_and_cache_mla_rope_fused",
TORCH_BOX(&concat_and_cache_mla_rope_fused));
ops.impl("convert_fp8", TORCH_BOX(&convert_fp8));
docs/design/hisparse.md
+244
View File
@@ -0,0 +1,244 @@
# HiSparse local KV offload architecture
Status: experimental
## The short version
There are three different jobs:
1. The normal KV cache system manages GPU block pools and tables.
2. `HiSparseCoordinator` manages logical host blocks, source-prefix identity, and
host/GPU residency transitions.
3. `HiSparseConnector` carries residency work between scheduler and worker;
`HiSparseWorker` coordinates transfers, while per-cache
`HiSparseRuntime` objects own host/hot views and GPU replacement state.
Neither worker-side object allocates or frees logical blocks. HMA provides the
GPU allocation shared by the resident and hot groups; it does not manage CPU
memory or KV identity.
```text
request ────► HiSparseCoordinator ── source blocks + residency policy
│
├──► KV cache manager ── resident/hot GPU leases (HMA)
│
└──► HiSparseConnector ── host bytes + copies + GPU LRU
```
This is a local KV connector. When P/D or another offload connector is also
configured, `MultiConnector` composes it with `HiSparseConnector`.
`host_pool_gib` is configured on `HiSparseConnector` and is the usable
host-cache capacity per data-parallel replica, not a node-wide memory budget.
Tensor-parallel ranks hold replicated views of that logical cache. Those views
may use private per-rank backing or one shared
physical allocation without changing the configured capacity. Physical host
memory consumption is therefore topology- and implementation-dependent. The
realized capacity may be slightly smaller because the budget is rounded down
to complete host blocks.
## Ownership
| Thing | Owner | What “owner” means |
| --- | --- | --- |
| HiSparse source and prefix identity | `HiSparseCoordinator` | maps tokens to logical host blocks |
| Resident GPU block leases | normal KV cache manager | allocates and frees HMA blocks |
| Resident block tables | normal KV cache manager | tells attention where resident pages are |
| Residency transitions | `HiSparseCoordinator` | plans spill-before-free transactions |
| Logical host block allocation | `HiSparseCoordinator` | owns the separate CPU block pool and its lifecycle |
| Pinned host-pool lifecycle | `HiSparseWorker` | worker-wide backing and teardown |
| Per-cache host view and hot contents | `HiSparseRuntime` | binds host/hot storage and fills cache-manager-provided hot leases |
| Hot row map and LRU | `HiSparseRuntime` | resolves hits and chooses victims on GPU |
| Resident-cache route | `HiSparseCacheHandle` | exposes resident or host/hot resolution to attention |
| Sparse attention | attention backend | consumes a device cache and physical row IDs |
| HMA | allocator | provides GPU capacity; owns no KV meaning |
The key distinction is logical allocation versus contents. `HiSparseCoordinator`
owns host block IDs and their request/prefix associations. `HiSparseWorker` and
its per-cache runtimes own the corresponding bytes. The normal cache manager
sees only device pools.
The source group has `block_pool_id=None`; device-pool consumers must narrow it
before indexing, so host ownership cannot masquerade as a numeric GPU pool.
## Code boundary
```text
scheduler process worker process
HiSparseConnector HiSparseConnector
└─ HiSparseCoordinator └─ HiSparseWorker
│ │
│ connector metadata ├─ host bytes
│ - page transfers ├─ copy scheduling
│ - block-table replacements └─ per-layer hot state
└───────────────────────────────────────────────►│
◄──────── connector worker metadata ─────────────┘
enqueued and completed transfer IDs
```
The command travels in `kv_connector_metadata`; transfer updates return in
`KVConnectorOutput.kv_connector_worker_meta`. The model runner does not
interpret page transfers. Enqueue acknowledgements let the scheduler release
source leases in stream order; completion acknowledgements publish the copied
host pages.
## Resident device pages
Resident pages are intentionally outside `HiSparseRuntime`.
KV-cache initialization binds cache-manager allocations to the attention-facing
`HiSparseCacheHandle` before constructing `HiSparseWorker`. That same
handle's runtime retains the resident source index needed by a transfer plan.
There is no second resident object or registration wrapper.
```text
KV cache setup
│
├─ bind resident allocation ──► HiSparseCacheHandle
│ cache + block table + slot mapping
│
├─ bind host/hot allocation ──► HiSparseRuntime
│ host + hot + GPU LRU
│
└─ register cache handles ────► HiSparseWorker
step-level transfers
HiSparseWorker registers the same HiSparseCacheHandle objects directly
```
Attention construction links each layer to the most recent layer that actually
owns an indexer. This releases a follower's duplicate LRU tensors before GPU
memory profiling. Cache binding only attaches storage; it does not infer
semantic groups from the physical packed-tensor order. The construction cursor
is discarded with the worker's pinned state.
Every HiSparse decode batch uses the same fused resolver. It checks resident
pages first, then hot rows, then pinned host memory. A resident hit exits inside
the kernel before hot-LRU lookup or host copying; there is no framework-level
residency route or separate CUDA graph. No CPU decision is added to the decode
path. The resolver consumes the existing graph-stable request mapping from
attention metadata; neither the worker nor individual cache handles keep a
duplicate mapping.
Speculative decoding resolves and consumes each verification step in order.
Each step receives distinct replayable plan rows while sharing the request's
hot-cache state, so a later step cannot reuse a hot row before an earlier step
has consumed it.
## P/D import target
The decoder chooses the landing target once per request from the normal cache
admission calculation. If the complete imported prefix fits the device pools,
NIXL transfers it directly into resident GPU pages. Otherwise, if the fixed
host-backed GPU footprint and host source blocks fit, the request imports into
the host tier. There is no context-length threshold or other heuristic, and a
request waiting for capacity retains its choice across admission retries.
A host import reads through a bounded decoder-GPU staging pool before copying
into registered host memory. Pages needed immediately are mirrored into their
resident destinations during that copy. Both landing targets then use the same
fused decode resolver described above.
## Indexer KV offloading
HiSparse does not keep a private CPU copy of indexer KV. The indexer remains a
normal prefix-cacheable GPU cache group. If `OffloadingConnector` is configured
with HiSparse, it stores and restores that group through the generic KV
offloading path; HiSparse continues to own only the sparse MLA host tier.
The two prefix sources can have different hit lengths. When the HiSparse host
prefix extends beyond the GPU-resident indexer prefix, the scheduler asks
`OffloadingConnector` to restore only the missing indexer suffix, capped at the
host prefix boundary. If that suffix is unavailable, all groups fall back to
the shorter prefix they share. NIXL P/D transfers continue to place indexer KV
directly in its GPU group.
## Spill transaction
A resident block cannot be reused until its contents have been handed to the
worker.
```text
HiSparseCoordinator HiSparseWorker
│ │
│ pin source and destination leases │
│── SparseKVPageTransfer ─────────────►│
│ │ enqueue GPU-to-host copy
│◄── enqueued transfer ID ────────────│
│ replace resident table entry │
│ release resident lease to HMA │
│ │ copy reaches its event
│◄── completed transfer ID ───────────│
│ mark host page valid │
│ release destination host lease │
```
“Enqueued” means the copy has entered the worker stream. Stream ordering makes
it safe to reuse the resident GPU block for later work, but the host page is
not yet published. “Completed” means the worker has observed the copy's event;
only then does the coordinator publish the host page for prefix reuse and
release its destination lease. A host-write event separately protects direct
CPU readers from writes already queued on the accelerator.
The worker transfer contains only its transfer ID and physical copy
coordinates. Request identity and logical page state remain in the scheduler.
## Hot lookup and LRU
The NVIDIA CUDA path keeps replacement entirely on the accelerator:
```text
top-K logical positions
│
▼
resident page? ── yes ──► resident physical row
│ no
▼
hot row? ──────── yes ──► existing hot physical row + update GPU LRU
│ no
▼
choose GPU LRU victim ──► copy pinned host row ──► hot physical row
```
ROCm is not currently supported because the fused HiSparse cache operations are
implemented only by CUDA kernels. A future platform-specific worker may provide
the same command, output, and cache-resolution boundaries.
## Main classes
| Class | Inherits / implements | Responsibility |
| --- | --- | --- |
| `HiSparseCoordinator` | plain scheduler component | host allocation, source prefixes, resident leases, and spill state machine |
| `HiSparseConnector` | `KVConnectorBase_V1`, `SupportsHMA` | scheduler/worker metadata and lifecycle boundary |
| `HiSparseResidentManager` | `SingleTypeKVCacheManager` | normal block-pool bookkeeping with host-backed holes |
| `PagedCacheView` | immutable data object | shared resident/hot HMA tensor binding |
| `HiSparseWorker` | connector-owned worker component | worker-wide transfer scheduling and host-pool lifecycle |
| `HiSparseRuntime` | plain worker-owned component | per-cache host/hot tensors, GPU LRU, and fused resolution |
| `HiSparseCacheHandle` | plain attention component | resident view and fused cache resolution |
| `SparseKVOffloadCommand` | dataclass | opaque scheduler-to-worker work |
## Performance invariants
- Resident hits bypass hot-LRU lookup and host copies inside the fused resolver.
- Hot lookup, victim selection, and LRU updates stay on the GPU.
- A hot miss still copies directly from registered pinned host memory.
- Top-K resolution stays inside the attention invocation and remains graph
capturable.
- Compatible layers still share one miss plan.
- Index-sharing followers release their private LRU state before memory sizing.
- Indexer KV is untouched by HiSparse unless a generic KV offloader is configured.
- Resident and hot leases can still share one packed HMA allocation.
- No device scalar readback or CPU/device metadata round trip is added.
- The abstraction wraps the fused kernel; it does not add another kernel
launch.
- When HiSparse is disabled, the scheduler does not construct an offload
command or empty update table.
## What remains platform-specific
The command/result and attention-layer boundaries can be shared. The host
allocator, copy implementation, hot layout, and replacement policy should stay
platform-specific. NVIDIA uses the current accelerator LRU and fused host/hot
kernel. ROCm is not currently supported; AMD or other accelerator backends can
implement their own worker without forcing NVIDIA's policy into the shared
boundary.
tests/evals/gsm8k/gsm8k_eval.py
+7
View File
@@ -364,6 +364,12 @@ def main() -> None:
type=int,
help="Maximum number of concurrent requests",
)
parser.add_argument(
"--request-timeout-seconds",
type=float,
default=600,
help="Timeout for each request, including time waiting for a connection",
)
parser.add_argument("--save-results", type=str, help="Save results to JSON file")
args = parser.parse_args()
@@ -377,6 +383,7 @@ def main() -> None:
temperature=args.temperature,
seed=args.seed,
max_concurrency=args.max_concurrency,
request_timeout_seconds=args.request_timeout_seconds,
)
# Print results to terminal
tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py
+8 -4
View File
@@ -217,7 +217,8 @@ def test_flashinfer_mla_dense_decode_unified_slot_view():
assert (out_ref - out_view).abs().max().item() == 0.0
def test_flashmla_fp8_sparse_decode_unified_slot_view():
@pytest.mark.parametrize("next_n", [1, 2, 4])
def test_flashmla_fp8_sparse_decode_unified_slot_view(next_n: int):
"""FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a
unified-slot block-major view bit-identically to a contiguous fp8_ds_mla
cache, with finite nonzero output."""
@@ -240,7 +241,10 @@ def test_flashmla_fp8_sparse_decode_unified_slot_view():
n_layers = 3
layer = 1
q = torch.randn(batch, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
q = (
torch.randn(batch, next_n, h_q, head_dim, device=dev, dtype=torch.bfloat16)
* 0.1
)
# Structurally valid fp8 ds_mla payload: 512B fp8 + 16B f32 scales + 128B
# bf16 rope (random bytes corrupt the scale region and yield NaNs).
@@ -273,13 +277,13 @@ def test_flashmla_fp8_sparse_decode_unified_slot_view():
# Sparse indices: each batch uses its own disjoint blocks.
blocks_per_batch = num_blocks // batch
idx = torch.full((batch, 1, topk), -1, device=dev, dtype=torch.int32)
idx = torch.full((batch, next_n, topk), -1, device=dev, dtype=torch.int32)
for b in range(batch):
slots: list[int] = []
for blk in range(b * blocks_per_batch, (b + 1) * blocks_per_batch):
slots.extend(blk * page + off for off in range(page))
slots_t = torch.tensor(slots[:topk], device=dev, dtype=torch.int32)
idx[b, 0, : slots_t.numel()] = slots_t
idx[b, :, : slots_t.numel()] = slots_t
def run(kc):
meta, num_splits = fm.get_mla_metadata()
tests/kernels/test_cp_gather_fp8.py
+42
View File
@@ -316,6 +316,48 @@ def test_cp_gather_fp8_shuffled_blocks():
assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])
def test_cp_gather_fp8_mixes_resident_and_pinned_host_rows():
seq_lens = [7, 5]
block_size = 4
(
cache,
block_table,
workspace_starts,
num_reqs,
total_tokens,
expected,
) = _build_test_case(seq_lens, block_size)
device_cache = cache.clone()
host_cache = cache.cpu().pin_memory().view(-1, ENTRY_BYTES)
device_rows = torch.arange(host_cache.shape[0], dtype=torch.int32, device="cuda")
use_device = device_rows % 2 == 0
host_rows = device_rows.clone()
device_rows.masked_fill_(~use_device, -1)
host_cache[use_device.cpu()] = 0
device_cache.view(-1, ENTRY_BYTES)[~use_device] = 0
dst = torch.empty(
total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda"
)
ops.cp_gather_and_upconvert_fp8_kv_cache(
device_cache,
dst,
block_table,
workspace_starts,
num_reqs,
host_cache=host_cache,
host_row_ids=host_rows,
device_row_ids=device_rows,
)
torch.testing.assert_close(
dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2
)
assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])
@pytest.mark.parametrize(
"gather_seq_lens,seq_starts",
[
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+70
View File
@@ -234,6 +234,76 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str
assert (topk == -1).all(), "topk buffer not cleared on indexer layer"
def test_fused_norm_rope_packed_indexer_block_stride():
"""Indexer writes use their HMA block stride and independent slot map."""
torch.manual_seed(7)
dev = "cuda"
num_tokens = block_size = 2
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
cos_sin = make_cos_sin(32, ROPE_DIM, dev)
idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4
packed_block_stride = block_size * idx_row + 64
backing = torch.zeros(2 * packed_block_stride, device=dev, dtype=torch.uint8)
idx_cache = torch.as_strided(
backing,
(2, block_size, idx_row),
(packed_block_stride, idx_row, 1),
)
mla_cache = torch.zeros(
1,
block_size,
KV_LORA + ROPE_DIM,
device=dev,
dtype=torch.bfloat16,
)
topk = torch.zeros(num_tokens, 8, device=dev, dtype=torch.int32)
K.fused_norm_rope(
pos,
q_c,
qw,
EPS,
kv_c,
kvw,
EPS,
k_pe,
cos_sin,
ik,
ikw,
ikb,
EPS,
cos_sin,
topk,
slot_mapping=torch.arange(num_tokens, device=dev, dtype=torch.int64),
indexer_slot_mapping=torch.arange(
block_size, block_size + num_tokens, device=dev, dtype=torch.int64
),
indexer_k_cache=idx_cache,
mla_kv_cache=mla_cache,
has_indexer=True,
)
ik_ref = rope(layer_norm(ik, ikw, ikb), pos, cos_sin, interleave=False)
q_ref, s_ref = ue8m0_quant(ik_ref)
packed = idx_cache[1].reshape(-1)
values = (
packed[: block_size * INDEX_HEAD_DIM].view(FP8).view(block_size, INDEX_HEAD_DIM)
)
scales = packed[block_size * INDEX_HEAD_DIM :].view(torch.float32)
assert_fp8(values, q_ref, "packed indexer-K fp8")
torch.testing.assert_close(scales, s_ref, rtol=0, atol=0)
assert (backing[block_size * idx_row : packed_block_stride] == 0).all()
def test_fused_norm_rope_materializes_qk_without_local_cache_slots():
"""DCP non-owner ranks still need valid query shards for query AllGather."""
torch.manual_seed(7)
tests/model_executor/layers/test_mla_short_prefill_indexer.py
+50
View File
@@ -7,6 +7,7 @@ import pytest
import torch
import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer
import vllm.models.deepseek_v32.attention as deepseek_attention
from vllm.config import CUDAGraphMode
from vllm.models.deepseek_v32 import attention as deepseek_v32_attention
from vllm.models.deepseek_v32.attention import DeepseekV32Attention
@@ -16,6 +17,51 @@ INDEXER_LAYER = "model.layers.0.self_attn.indexer.k_cache"
MLA_LAYER = "model.layers.0.self_attn.attn"
def test_sparse_attention_refreshes_batch_state_inside_eager_segment(
monkeypatch: pytest.MonkeyPatch,
) -> None:
metadata = SimpleNamespace(num_actual_tokens=1)
class BatchStateRefreshed(Exception):
pass
def refresh_batch_state(attn_metadata):
assert attn_metadata is metadata
raise BatchStateRefreshed
layer = SimpleNamespace(
indexer=None,
skip_topk=True,
layer_name=MLA_LAYER,
impl=SimpleNamespace(
prepare_for_batch=refresh_batch_state,
record_logical_topk_ready=lambda: None,
),
)
monkeypatch.setattr(
deepseek_attention,
"get_attention_context",
lambda layer_name: (metadata, None, torch.empty(1), None),
)
with pytest.raises(BatchStateRefreshed):
DeepseekV32Attention._sparse_indexer_and_attn(
layer,
torch.empty(1, dtype=torch.long),
torch.empty(1, 1),
torch.empty(1, 1, 1),
torch.empty(1, 1, 1),
None,
None,
None,
None,
None,
torch.empty(1, 1, 1),
torch.empty(1, 1, 1),
torch.empty(1, 1),
)
def make_indexer_metadata(
*,
num_decodes: int = 0,
@@ -246,6 +292,10 @@ def test_deepseek_v32_dispatches_selected_mha(
_fp8_query=fp8_query,
_use_sparse_mha=lambda _: True,
rotary_emb=lambda _positions, q: (q + 1, None),
impl=SimpleNamespace(
record_logical_topk_ready=lambda: None,
prepare_for_batch=lambda _: None,
),
forward_impl=record_forward_impl,
)
q_nope = torch.randn(2, 1, 2)
tests/model_executor/test_flashinfer_autotune_warmup.py
+13 -2
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from inspect import signature
from types import SimpleNamespace
from typing import Any
from unittest.mock import Mock, call, patch
@@ -76,8 +77,16 @@ def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated():
assert token_counts == (32,)
def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run():
@pytest.mark.parametrize("skip_attn", [False, True])
def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn):
from vllm.v1.worker.gpu.model_runner import GPUModelRunner as V2Runner
from vllm.v1.worker.gpu_model_runner import GPUModelRunner as V1Runner
runner = _make_runner([])
dummy_run_signature = signature((V2Runner if skip_attn else V1Runner)._dummy_run)
runner._dummy_run.side_effect = lambda **kwargs: dummy_run_signature.bind(
None, **kwargs
)
max_buckets = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 8192)
deferred_buckets = (1, 2, 4, 8, 16, 32, 64, 128)
@@ -93,7 +102,7 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run():
) as get_buckets,
patch("vllm.utils.flashinfer.autotune") as autotune,
):
_run_flashinfer_autotune_dummy_runs(runner)
_run_flashinfer_autotune_dummy_runs(runner, skip_attn=skip_attn)
assert get_buckets.call_args_list == [call(8192), call(128)]
assert autotune.call_args_list == [
@@ -106,11 +115,13 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run():
skip_eplb=True,
is_profile=True,
randomize_inputs=True,
**({"skip_attn": True} if skip_attn else {}),
),
call(
num_tokens=128,
skip_eplb=True,
is_profile=True,
randomize_inputs=True,
**({"skip_attn": True} if skip_attn else {}),
),
]
tests/test_config.py
+66 -2
View File
@@ -19,10 +19,12 @@ import vllm.config.vllm as vllm_config_module
import vllm.envs as envs
from vllm.compilation.backends import VllmBackend
from vllm.config import (
AttentionConfig,
CacheConfig,
CompilationConfig,
DeviceConfig,
EngramConfig,
HiSparseConfig,
KernelConfig,
KVTransferConfig,
ModelConfig,
@@ -279,6 +281,62 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected):
assert envs.VLLM_USE_V2_MODEL_RUNNER is expected
def test_hisparse_requires_v2_model_runner():
config = object.__new__(VllmConfig)
config.attention_config = AttentionConfig(hisparse_config=HiSparseConfig())
with patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", None):
assert config.use_v2_model_runner
with (
patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", False),
pytest.raises(ValueError, match="requires Model Runner V2"),
):
_ = config.use_v2_model_runner
def test_hisparse_rejects_decode_context_parallelism(monkeypatch):
monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
monkeypatch.setattr(current_platform, "device_count", lambda: 2)
with pytest.raises(ValueError, match="decode context parallelism"):
VllmConfig(
attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
parallel_config=ParallelConfig(
tensor_parallel_size=2,
decode_context_parallel_size=2,
),
)
def test_hisparse_rejects_pipeline_parallelism(monkeypatch):
monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
monkeypatch.setattr(current_platform, "device_count", lambda: 2)
with pytest.raises(ValueError, match="pipeline parallelism"):
VllmConfig(
attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
parallel_config=ParallelConfig(pipeline_parallel_size=2),
)
def test_hisparse_rejects_disabled_hybrid_kv_cache_manager(monkeypatch):
monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True)
with pytest.raises(ValueError, match="requires the hybrid KV cache manager"):
VllmConfig(
attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
scheduler_config=SchedulerConfig(
max_model_len=2048,
is_encoder_decoder=False,
disable_hybrid_kv_cache_manager=True,
),
)
def test_hisparse_rejects_non_cuda(monkeypatch):
monkeypatch.setattr(current_platform, "is_cuda", lambda: False)
with pytest.raises(ValueError, match="requires NVIDIA CUDA"):
VllmConfig(attention_config=AttentionConfig(hisparse_config=HiSparseConfig()))
def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch):
"""ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled
MRV1 paths and off breakable cudagraphs by default."""
@@ -305,7 +363,8 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch):
# (warning_once args must be hashable for its lru_cache).
monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)
config = SimpleNamespace(
model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"])
model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"]),
attention_config=AttentionConfig(),
)
assert VllmConfig.use_v2_model_runner.fget(config) is False
finally:
@@ -385,6 +444,7 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph(
)
config = SimpleNamespace(
model_config=model_config,
attention_config=AttentionConfig(),
speculative_config=SimpleNamespace(method="mtp") if with_mtp else None,
parallel_config=SimpleNamespace(prefill_context_parallel_size=1),
compilation_config=CompilationConfig(
@@ -897,7 +957,10 @@ def test_models_default_to_v2_model_runner(model_config, expected, monkeypatch):
monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)
monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True)
monkeypatch.setattr(current_platform, "is_rocm", lambda: False)
config = SimpleNamespace(model_config=model_config)
config = SimpleNamespace(
model_config=model_config,
attention_config=AttentionConfig(),
)
config._get_v2_model_runner_unsupported_features = lambda: []
assert VllmConfig.use_v2_model_runner.fget(config) is expected
@@ -2787,6 +2850,7 @@ def test_draft_sample_method_gumbel_is_rejected():
def _watermarked_vllm_config() -> VllmConfig:
config = object.__new__(VllmConfig)
config.watermark_config = WatermarkConfig(key=42)
config.attention_config = AttentionConfig()
config.speculative_config = None
return config
tests/v1/attention/test_backend_per_kind.py
+54 -2
View File
@@ -1,16 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for per-KV-group attention backend selection (backend_per_kind)."""
"""Tests for attention configuration and per-KV-group backend selection."""
from types import SimpleNamespace
import pytest
from vllm.config.attention import AttentionConfig
from vllm.config.attention import AttentionConfig, HiSparseConfig
from vllm.model_executor.layers.attention.attention import (
_largest_kernel_block_within,
)
from vllm.v1.attention.backend import AttentionType, MultipleOf
from vllm.v1.attention.backends.registry import AttentionBackendEnum
from vllm.v1.attention.selector import get_attn_spec_kind
from vllm.v1.hisparse.runtime import ResolvedHiSparseConfig
from vllm.v1.kv_cache_interface import KVCacheSpecKind
@@ -86,3 +89,52 @@ def test_backend_per_kind_rejects_unknown_kind():
def test_backend_per_kind_defaults_empty():
assert AttentionConfig().backend_per_kind == {}
def test_hisparse_device_buffer_size_boundaries():
vllm_config = SimpleNamespace(
attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
speculative_config=None,
)
resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
assert resolved is not None
assert resolved.device_buffer_size == 256
vllm_config.attention_config.hisparse_config = HiSparseConfig(
device_buffer_size=127
)
with pytest.raises(ValueError, match="expected at least 128"):
ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
vllm_config.attention_config.hisparse_config = HiSparseConfig(
device_buffer_size=32768
)
resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
assert resolved is not None
assert resolved.device_buffer_size == 32768
vllm_config.attention_config.hisparse_config = HiSparseConfig(
device_buffer_size=32769
)
with pytest.raises(ValueError, match="int16 slot-index limit"):
ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
def test_hisparse_device_buffer_covers_speculative_window():
vllm_config = SimpleNamespace(
attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
speculative_config=SimpleNamespace(
num_speculative_tokens=3,
parallel_drafting=False,
),
)
resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
assert resolved is not None
assert resolved.device_buffer_size == 5 * 128
vllm_config.attention_config.hisparse_config = HiSparseConfig(
device_buffer_size=4 * 128 - 1
)
with pytest.raises(ValueError, match="expected at least 512"):
ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
tests/v1/attention/test_mla_backends.py
+1
View File
@@ -200,6 +200,7 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype(
layer = SimpleNamespace(
kv_cache_dtype=cache_dtype,
head_size=576,
indexer=None,
non_causal_multi_token_decode=False,
sliding_window=None,
)
tests/v1/attention/test_mla_context_chunks.py
+23
View File
@@ -55,6 +55,7 @@ def build_chunked_context(
def test_sparse_context_lengths_do_not_force_gpu_sync(monkeypatch):
"""Sparse MLA must pin computed context lengths before the H2D copy."""
builder = SimpleNamespace(
_context_lens_cpu=torch.empty(2, dtype=torch.int32, pin_memory=True),
chunked_prefill_workspace=torch.empty((2048, 1)),
chunked_prefill_workspace_size=1024,
kv_cache_spec=SimpleNamespace(block_size=BLOCK_SIZE),
@@ -113,6 +114,7 @@ def test_chunks_gather_every_context_row_exactly_once(context_lens, workspace_si
query_lens = [4] * len(context_lens)
metadata = build_chunked_context(context_lens, query_lens, workspace_size)
assert metadata is not None
assert metadata.context_lens.tolist() == context_lens
gathered: dict[int, list[tuple[int, int]]] = {}
previous_request_start = -1
@@ -253,6 +255,7 @@ def test_dcp_chunks_fit_the_per_rank_row_budget():
dcp_local_block_size=interleave,
)
assert metadata is not None
assert metadata.context_lens.tolist() == context_lens
virtual_block_size = interleave * dcp_world_size
local_cursor: dict[int, int] = {}
@@ -330,3 +333,23 @@ def test_dcp_reorg_uses_each_chunks_local_starts():
)
torch.testing.assert_close(reorganized, torch.cat(expected))
def test_chunked_context_metadata_owns_context_lengths():
"""Reusing CPU scratch must not overwrite an earlier batch's context lengths."""
context_lens = torch.tensor([16, 32], dtype=torch.int32)
metadata = build_mla_chunked_context_metadata(
context_lens_cpu=context_lens,
prefill_query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32),
chunked_prefill_workspace=torch.empty((64, 1)),
chunked_prefill_workspace_size=64,
block_size=BLOCK_SIZE,
align_chunk_to_block=True,
device=torch.device("cpu"),
dcp_world_size=1,
dcp_local_block_size=1,
dcp_virtual_block_size=1,
)
context_lens.zero_()
assert metadata is not None
assert metadata.context_lens.tolist() == [16, 32]
tests/v1/attention/test_sparse_mla_backends.py
+1834 -49
View File
File diff suppressed because it is too large Load Diff
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+6 -67
View File
@@ -41,7 +41,9 @@ def test_connector_without_divergent_hit_support_uses_common_lookup():
manager = MagicMock()
manager.get_computed_blocks.return_value = (common_blocks, 0, 0)
scheduler = SimpleNamespace(
connector=SimpleNamespace(supports_divergent_local_hybrid_hits=False),
connector=SimpleNamespace(
supports_divergent_local_hybrid_hits=False,
),
kv_cache_manager=manager,
)
@@ -61,7 +63,9 @@ def test_capable_connector_uses_divergent_partial_hit_lookup():
True,
)
scheduler = SimpleNamespace(
connector=SimpleNamespace(supports_divergent_local_hybrid_hits=True),
connector=SimpleNamespace(
supports_divergent_local_hybrid_hits=True,
),
kv_cache_manager=manager,
)
@@ -922,71 +926,6 @@ def test_connector_finish_registers_partial_tail_before_cleanup():
scheduler.kv_cache_manager.remove_skipped_blocks.assert_called_once()
def test_block_pool_touch_pins_released_cow_target():
"""The connector can rescue an offered CoW target after its step-scoped
retention is released by using the bound BlockPool's touch method."""
hash_block_size = 2
block_size = 2 * hash_block_size
kv_cache_config = KVCacheConfig(
num_blocks=24,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec(
["full"],
FullAttentionSpec(
block_size=hash_block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
),
KVCacheGroupSpec(
["mamba"],
MambaSpec(
block_size=block_size,
shapes=(1, 1),
dtypes=(torch.float32,),
mamba_cache_mode="align",
),
),
],
)
manager = make_kv_cache_manager(
kv_cache_config=kv_cache_config,
max_model_len=8192,
enable_caching=True,
hash_block_size=hash_block_size,
)
req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
req0.num_computed_tokens = 6
req0.append_output_token_ids([3])
assert manager.allocate_slots(req0, 1) is not None
# Retention released before the drain (defer_block_free=False ordering).
_copies, retained = manager.take_kv_cache_block_copies()
manager.block_pool.free_blocks(retained)
offloads = drain_boundary_state_offloads(manager)
((_group_id, block_id, boundary_tokens),) = offloads["0"]
assert boundary_tokens == 6
cow_block = manager.block_pool.blocks[block_id]
assert cow_block.ref_cnt == 0
assert block_id in _free_block_ids(manager)
manager.block_pool.touch([cow_block])
assert cow_block.ref_cnt == 1
assert block_id not in _free_block_ids(manager)
# The connector-pinned block is out of the free queue: draining every free block
# neither trips the allocator's ref_cnt assert nor hands it out.
new_blocks = manager.block_pool.get_new_blocks(
manager.block_pool.get_num_free_blocks()
)
assert block_id not in {b.block_id for b in new_blocks}
def test_boundary_state_offload_dropped_when_request_freed_before_drain():
"""A hand-off recorded in the same scheduling pass as the request's death
must not be drained: its release hook has already run, so draining would
tests/v1/core/test_async_scheduler.py
+27 -1
View File
@@ -13,7 +13,7 @@ from vllm.v1.request import RequestStatus
from vllm.v1.structured_output import StructuredOutputGrammar
from vllm.v1.utils import ConstantList
from .utils import create_requests, create_scheduler
from .utils import create_requests, create_scheduler, mock_kv
pytestmark = pytest.mark.cpu_test
@@ -145,6 +145,32 @@ def test_abort():
assert req.num_output_tokens == abort_order_copy.index(i)
def test_connector_metadata_precedes_async_placeholder_advance(monkeypatch):
"""Mirroring must see earlier unresolved outputs, excluding the current step."""
scheduler = create_scheduler(
async_scheduling=True,
use_kv_connector=mock_kv(matched_tokens=0, is_async=False),
)
(request,) = create_requests(num_requests=1, num_tokens=4, max_tokens=4)
scheduler.add_request(request)
assert scheduler.connector is not None
build_connector_meta = scheduler.connector.build_connector_meta
observed_placeholders = []
def record_placeholders(scheduler_output):
observed_placeholders.append(request.num_output_placeholders)
return build_connector_meta(scheduler_output)
monkeypatch.setattr(
scheduler.connector, "build_connector_meta", record_placeholders
)
scheduler.schedule()
scheduler.schedule()
assert observed_placeholders == [0, 1]
def test_preempt():
scheduler = create_scheduler(async_scheduling=True)
requests = create_requests(num_requests=10, max_tokens=20)
tests/v1/core/test_contiguous_kv_packing.py
+1
View File
@@ -75,6 +75,7 @@ def _mock_vllm_config(layout: str | None):
config.cache_config = CacheConfig()
config.cache_config.num_gpu_blocks_override = None
config.cache_config.kv_cache_layout = layout
config.attention_config.hisparse_config = None
return config
tests/v1/core/test_kv_cache_utils.py
+144
View File
@@ -22,6 +22,7 @@ from vllm.config import (
SchedulerConfig,
VllmConfig,
)
from vllm.config.attention import HiSparseConfig
from vllm.config.kv_events import KVEventsConfig
from vllm.lora.request import LoRARequest
from vllm.multimodal.inputs import (
@@ -52,10 +53,16 @@ from vllm.v1.core.kv_cache_utils import (
make_block_hash_with_group_id,
tensor_data,
)
from vllm.v1.hisparse.layout import (
create_hisparse_layout,
get_hisparse_gpu_memory_usage,
)
from vllm.v1.kv_cache_interface import (
ChunkedLocalAttentionSpec,
FullAttentionSpec,
HiddenStateCacheSpec,
HiSparseHotSpec,
HiSparseResidentSpec,
KpoolTailSpec,
KVCacheConfig,
KVCacheGroupSpec,
@@ -68,18 +75,139 @@ from vllm.v1.kv_cache_interface import (
SinkFullAttentionSpec,
SlidingWindowMLASpec,
SlidingWindowSpec,
SparseCacheRole,
UniformTypeKVCacheSpecs,
get_kv_cache_spec_kind,
get_kv_cache_spec_sliding_window,
is_full_attention_spec,
iter_layer_specs,
)
from vllm.v1.kv_cache_layout import KVCacheLayout
from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats
from vllm.v1.request import Request
pytestmark = pytest.mark.cpu_test
@pytest.mark.parametrize("gpu_block_size", [32, 64])
def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size):
specs = {
"model.layers.0.self_attn": MLAAttentionSpec(
block_size=gpu_block_size,
num_kv_heads=1,
head_size=576,
dtype=torch.bfloat16,
is_index_group_leader=True,
),
"model.layers.0.self_attn.indexer": MLAAttentionSpec(
block_size=gpu_block_size,
num_kv_heads=1,
head_size=128,
dtype=torch.bfloat16,
cache_role=SparseCacheRole.INDEXER,
),
}
group_spec = UniformTypeKVCacheSpecs.from_specs(specs)
assert group_spec is not None
group = KVCacheGroupSpec(list(specs), group_spec)
config = SimpleNamespace(
attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
model_config=SimpleNamespace(
hf_config=SimpleNamespace(index_topk=128),
max_model_len=gpu_block_size,
),
parallel_config=SimpleNamespace(decode_context_parallel_size=1),
cache_config=SimpleNamespace(
num_gpu_blocks_override=7,
prefix_cache_retention_interval=None,
get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC,
),
)
indexer_spec = specs["model.layers.0.self_attn.indexer"]
assert get_hisparse_gpu_memory_usage(config, [group]) == (
indexer_spec.max_memory_usage_bytes(config)
)
monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30)
cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
config, [group], available_memory=2**30
)
assert cache_config.num_blocks == 7
assert cache_config.hisparse_host_num_blocks is not None
assert cache_config.hisparse_host_num_blocks > 7
host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups
assert host_group.host_resident
assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups])
host_layers = set(host_group.layer_names)
for tensor in cache_config.kv_cache_tensors:
assert all(
(name in host_layers) == tensor.host_resident for name in tensor.layers
)
host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident)
host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]]
assert host_tensor.block_stride == host_spec.page_size_bytes
assert host_tensor.layer_stride == (
host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks
)
assert host_group.kv_cache_spec.block_size == gpu_block_size
assert indexer_group.kv_cache_spec.block_size == gpu_block_size
host_specs = host_group.kv_cache_spec.kv_cache_specs
gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs
assert set(host_specs) == {"model.layers.0.self_attn"}
assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"}
assert indexer_group.kv_cache_spec.prefix_cacheable
assert host_group.enable_kv_transfer
auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups]
assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs)
assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs)
assert all(
spec.block_size == gpu_block_size
for spec in auxiliary_specs
if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec))
)
scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes(
cache_config,
SimpleNamespace(
cache_config=SimpleNamespace(
block_size=16,
enable_prefix_caching=True,
prefix_match_unit=None,
),
parallel_config=SimpleNamespace(decode_context_parallel_size=1),
kv_transfer_config=object(),
),
)
assert scheduler_block_size == hash_block_size == gpu_block_size
def test_hisparse_rejects_deepseek_v4():
full_specs = {
"model.layers.0.attn": MLAAttentionSpec(
block_size=256,
num_kv_heads=1,
head_size=512,
dtype=torch.uint8,
model_version="deepseek_v4",
)
}
full_uniform = UniformTypeKVCacheSpecs.from_specs(full_specs)
assert full_uniform is not None
group = KVCacheGroupSpec(list(full_specs), full_uniform)
config = SimpleNamespace(
attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
model_config=SimpleNamespace(hf_config=SimpleNamespace(index_topk=512)),
cache_config=SimpleNamespace(num_gpu_blocks_override=7),
)
with pytest.raises(ValueError, match="does not support DeepSeek V4"):
create_hisparse_layout(
config,
[group],
host_budget=2**30,
)
@pytest.fixture(autouse=True)
def _auto_init_hash_fn(request):
hash_fn: Callable
@@ -3848,6 +3976,22 @@ def test_iter_layer_specs_returns_group_members():
assert list(iter_layer_specs(wrapped)) == [full, mla]
def test_wrapped_mamba_group_requires_block_zeroing():
mamba = MambaSpec(
block_size=4,
shapes=((4, 1),),
dtypes=(torch.float32,),
)
wrapped = UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs={"mamba": mamba})
config = KVCacheConfig(
num_blocks=4,
kv_cache_tensors=[],
kv_cache_groups=[KVCacheGroupSpec(["mamba"], wrapped)],
)
assert config.needs_kv_cache_zeroing
def _spec_decode_grouping_config(method="dspark", model_type=None):
"""Grouping config with an EAGLE-family speculative method enabled."""
return SimpleNamespace(
tests/v1/core/test_prefix_caching.py
+1096 -1
View File
File diff suppressed because it is too large Load Diff
tests/v1/core/test_scheduler.py
+69
View File
@@ -19,7 +19,12 @@ from vllm.config import (
SpeculativeConfig,
VllmConfig,
)
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
HiSparseConnector,
HiSparseConnectorScheduler,
)
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector
from vllm.multimodal.inputs import (
MultiModalFeatureSpec,
MultiModalKwargsItem,
@@ -34,6 +39,7 @@ from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
from vllm.v1.core.sched.scheduler import Scheduler
from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs
from vllm.v1.engine import FinishReason
from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
@@ -148,6 +154,30 @@ def test_get_num_unfinished_requests():
assert scheduler.get_num_unfinished_requests() == len(requests) - i - 1
def _bind_hisparse_connector(scheduler):
connector = object.__new__(HiSparseConnector)
connector.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False)
from tests.v1.core.test_prefix_caching import make_hisparse_kv_cache_manager
coordinator = get_hisparse_coordinator(make_hisparse_kv_cache_manager(16, 16))
connector.connector_scheduler.bind_coordinator(coordinator)
composite = object.__new__(MultiConnector)
composite._connectors = (connector,)
scheduler.connector = composite
return coordinator
def test_pending_hisparse_spill_keeps_scheduler_alive():
"""A final host spill must complete after the last request finishes."""
scheduler = create_scheduler()
coordinator = _bind_hisparse_connector(scheduler)
pending = coordinator.pending_spills
pending[0] = Mock()
assert scheduler.has_requests()
pending.clear()
assert not scheduler.has_requests()
@pytest.mark.parametrize(
"enable_prefix_caching, prompt_logprobs",
[
@@ -1181,6 +1211,45 @@ def test_preempt_during_execution():
assert requests[1].output_token_ids[0] == 42
def test_pending_hisparse_reclamation_defers_preemption(monkeypatch):
"""Pending reclamation must stall rather than preempt its owning request."""
scheduler = create_scheduler(enable_prefix_caching=False)
request = create_requests(num_requests=1, num_tokens=16, block_size=16)[0]
scheduler.add_request(request)
scheduler_output = scheduler.schedule()
scheduler.update_from_output(
scheduler_output,
ModelRunnerOutput(
req_ids=[request.request_id],
req_id_to_index={request.request_id: 0},
sampled_token_ids=[[0]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
),
)
monkeypatch.setattr(
scheduler.kv_cache_manager, "allocate_slots", Mock(return_value=None)
)
coordinator = _bind_hisparse_connector(scheduler)
monkeypatch.setattr(
scheduler.connector, "build_connector_meta", Mock(return_value=None)
)
monkeypatch.setattr(
coordinator,
"has_pending_reclamation",
Mock(return_value=True),
)
deferred_output = scheduler.schedule()
assert not deferred_output.num_scheduled_tokens
assert scheduler.running == [request]
assert request.status == RequestStatus.RUNNING
assert request.num_preemptions == 0
def test_prefix_cache_query_not_inflated_by_connector_defer():
"""The GPU prefix-cache query is recorded at admission, so a request the
connector defers several times is counted once, not once per retry."""
tests/v1/e2e/general/test_hisparse.py
+225
View File
@@ -0,0 +1,225 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import prometheus_client
import pytest
import torch
from tests.conftest import VllmRunner
from tests.utils import fork_new_process_for_each_test
from vllm.config import AttentionConfig, HiSparseConfig, KVTransferConfig
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
HiSparseConnector,
)
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
HiSparseConnectorWorker,
)
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
MultiConnector,
)
from vllm.platforms import current_platform
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
MODEL = "deepseek-ai/DeepSeek-V3.2"
HF_OVERRIDES = {
"num_hidden_layers": 8,
"hidden_size": 256,
"intermediate_size": 512,
"num_attention_heads": 8,
"num_key_value_heads": 1,
"n_routed_experts": 8,
"num_experts_per_tok": 2,
"index_topk": 128,
}
def _shrink_config(config):
overrides = HF_OVERRIDES.copy()
if any("MTP" in architecture for architecture in config.architectures):
overrides.pop("num_hidden_layers")
config.update(overrides)
return config
def _num_hisparse_spills(runner: VllmRunner) -> int:
client = runner.llm.llm_engine.engine_core
coordinator = get_hisparse_coordinator(
client.engine_core.scheduler.kv_cache_manager
)
return coordinator.next_spill_id
def _offload_load_bytes() -> float:
return sum(
sample.value
for metric in prometheus_client.REGISTRY.collect()
if metric.name == "vllm:kv_offload_load_bytes"
for sample in metric.samples
if sample.name == "vllm:kv_offload_load_bytes_total"
)
def _get_hisparse_worker(runner: VllmRunner) -> HiSparseConnectorWorker:
engine_core = runner.llm.llm_engine.engine_core.engine_core
model_runner = engine_core.model_executor.driver_worker.worker.model_runner
connector = model_runner.kv_connector.kv_connector
if isinstance(connector, MultiConnector):
connector = next(
child
for child in connector.sub_connectors
if isinstance(child, HiSparseConnector)
)
assert isinstance(connector, HiSparseConnector)
assert connector.connector_worker is not None
return connector.connector_worker
@pytest.mark.skipif(
not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA"
)
@pytest.mark.parametrize(
"with_offloading", [False, True], ids=["standalone", "offload"]
)
@fork_new_process_for_each_test
def test_hisparse_spill_and_prefix_restore(
monkeypatch: pytest.MonkeyPatch,
vllm_runner: type[VllmRunner],
with_offloading: bool,
):
"""Spilled prefixes restore and FULL-graph decode writes reach host KV.
A regression omitted attention metadata before FULL graph replay, so decode
completed normally while its newly written KV rows were never copied to host.
"""
capability = current_platform.get_device_capability()
if capability is None or capability.major < 9:
pytest.skip("Sparse MLA requires Hopper or newer")
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip")
target = [1000 + i % 64 for i in range(257)]
pressure = [
[2000 + request_idx * 128 + i % 64 for i in range(257)]
for request_idx in range(4)
]
hisparse_connector = {
"kv_connector": "HiSparseConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {"host_pool_gib": 1},
}
kv_transfer_config = KVTransferConfig(**hisparse_connector)
if with_offloading:
kv_transfer_config = KVTransferConfig(
kv_connector="MultiConnector",
kv_role="kv_both",
kv_connector_extra_config={
"connectors": [
hisparse_connector,
{
"kv_connector": "OffloadingConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
},
]
},
)
with vllm_runner(
MODEL,
load_format="dummy",
hf_overrides=_shrink_config,
attention_config=AttentionConfig(
hisparse_config=HiSparseConfig(
device_buffer_size=512,
)
),
kv_transfer_config=kv_transfer_config,
block_size=64,
max_model_len=320,
max_num_batched_tokens=1024,
max_num_seqs=4,
num_gpu_blocks_override=128,
disable_log_stats=False,
enable_chunked_prefill=True,
enable_prefix_caching=True,
speculative_config={"method": "mtp", "num_speculative_tokens": 3},
compilation_config={
"cudagraph_mode": "FULL_AND_PIECEWISE",
"cudagraph_capture_sizes": [1, 4],
},
) as runner:
worker = _get_hisparse_worker(runner)
engine_core = runner.llm.llm_engine.engine_core.engine_core
model_runner = engine_core.model_executor.driver_worker.worker.model_runner
full_graph_calls = 0
full_replay_pending = False
original_run_fullgraph = model_runner.cudagraph_manager.run_fullgraph
def record_fullgraph(batch_desc):
nonlocal full_graph_calls, full_replay_pending
full_graph_calls += 1
full_replay_pending = True
return original_run_fullgraph(batch_desc)
monkeypatch.setattr(
model_runner.cudagraph_manager, "run_fullgraph", record_fullgraph
)
original_finish_forward = worker.finish_forward
verified_full_replay_rows = 0
copied_nonzero_kv = False
def finish_forward():
nonlocal full_replay_pending, verified_full_replay_rows
nonlocal copied_nonzero_kv
try:
original_finish_forward()
if not full_replay_pending:
return
# Synchronize only the test's inspection, not finish_forward itself.
with gpu_sync_allowed():
torch.accelerator.synchronize()
assert worker._row_mirrors
for layer_index, cache in enumerate(worker.cache_handles):
source_index = cache.runtime.resident_source_index
for mirror in worker._row_mirrors:
source_row = mirror.source_starts[source_index]
destination_row = mirror.destination_start
num_rows = mirror.num_rows
source_block, source_offset = divmod(
source_row, worker.kernel_block_size
)
with gpu_sync_allowed():
gpu_rows = worker.resident_caches[layer_index][
source_block, source_offset : source_offset + num_rows
].cpu()
host_rows = worker.host_caches[layer_index][
destination_row : destination_row + num_rows
]
assert torch.equal(host_rows, gpu_rows)
copied_nonzero_kv |= bool(torch.count_nonzero(gpu_rows).item())
verified_full_replay_rows += num_rows
finally:
full_replay_pending = False
monkeypatch.setattr(worker, "finish_forward", finish_forward)
expected = runner.generate_greedy([target], max_tokens=8)
assert full_graph_calls > 0
assert verified_full_replay_rows > 0
assert copied_nonzero_kv
runner.generate_greedy(pressure, max_tokens=8)
assert _num_hisparse_spills(runner) > 0
load_bytes = _offload_load_bytes()
actual = runner.generate_greedy([target], max_tokens=8)
runner.generate_greedy([[42]], max_tokens=1)
assert actual == expected
if with_offloading:
assert _offload_load_bytes() > load_bytes
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+31 -1
View File
@@ -21,6 +21,7 @@ from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
HiddenStateCacheSpec,
KVCacheConfig,
KVCacheGroupRole,
KVCacheGroupSpec,
KVCacheSpec,
KVCacheTensor,
@@ -336,6 +337,36 @@ def test_worker_kv_bytes_preserves_tensor_layout(packed: bool):
assert offloading_config.cache.blocks_per_chunk == 2
def test_hisparse_offloads_only_indexer_group():
source = KVCacheGroupSpec(
["source"],
_full_attention_spec(),
host_resident=True,
)
indexer = KVCacheGroupSpec(
["indexer"],
_full_attention_spec(),
role=KVCacheGroupRole.HISPARSE_INDEXER,
)
kv_cache_config = KVCacheConfig(
num_blocks=4,
kv_cache_tensors=[],
kv_cache_groups=[source, indexer],
hisparse_host_num_blocks=4,
)
config = _make_vllm_config()
offloading_config = build_offloading_config(config, kv_cache_config)
scheduler_config = SchedulerOffloadConfig.from_spec(
MockOffloadingSpec(offloading_config), config, kv_cache_config
)
assert [
(group.group_id, group.layer_names) for group in offloading_config.groups
] == [(1, ("indexer",))]
assert [group.group_idx for group in scheduler_config.kv_group_configs] == [1]
def test_zero_blocks_skips_tensor_layout_validation():
kv_cache_config = _make_sizing_kv_cache_config(packed=False)
kv_cache_config.num_blocks = 0
@@ -792,7 +823,6 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
False,
id="uniform-uncertifiable-inner",
),
pytest.param([], False, id="attention-free"),
],
)
def test_canonical_layout_gate(kv_cache_groups, certified):
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+45 -2
View File
@@ -98,6 +98,14 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state):
)
manager = CPUOffloadingManager(num_chunks=100)
spec = SimpleNamespace(
config=SimpleNamespace(
groups=tuple(
SimpleNamespace(
group_id=i, tokens_per_block=group.kv_cache_spec.block_size
)
for i, group in enumerate(groups)
)
),
tokens_per_block=(256, 64, 8),
tokens_per_hash=8,
blocks_per_chunk=4,
@@ -303,7 +311,24 @@ def test_aligned_boundary_store_uses_exact_source_with_partial_tail():
assert src_spec.block_indices == [0, 0]
def test_aligned_boundary_store_flushes_before_cow_destination_reuse():
def test_aligned_boundary_store_maps_sparse_group_id_to_dense_transfer_slot():
scheduler = _make_partial_tail_scheduler()
indexer_config = scheduler.config.kv_group_configs[1]
scheduler.config = scheduler.config._replace(kv_group_configs=(indexer_config,))
_make_partial_tail_request(scheduler)
scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
jobs = scheduler._build_aligned_boundary_store_jobs({"req": [(1, 98, 16)]})
[job] = jobs.values()
assert job.src_spec.group_sizes == [1]
assert job.src_spec.block_indices == [0]
@pytest.mark.parametrize("cow_reuse", [False, True])
def test_aligned_boundary_store_flushes_before_block_reuse(cow_reuse):
scheduler = _make_partial_tail_scheduler()
_make_partial_tail_request(scheduler)
scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (
@@ -320,7 +345,11 @@ def test_aligned_boundary_store_flushes_before_cow_destination_reuse():
[job_id] = meta.store_jobs
output = SchedulerOutput.make_empty()
output.kv_cache_block_copies = [KVCacheBlockCopy(98, 99)]
if cow_reuse:
output.kv_cache_block_copies = [KVCacheBlockCopy(98, 99)]
else:
output.scheduled_cached_reqs.req_ids = ["req"]
output.scheduled_cached_reqs.new_block_ids = [([], [99])]
meta = scheduler.build_connector_meta(output)
assert meta.jobs_to_flush == {job_id}
@@ -384,6 +413,20 @@ def test_partial_lookup_returns_exact_boundary_and_group_load_keys():
assert req_status.partial_tail_boundary is None
def test_lookup_cap_stops_at_authoritative_prefix_boundary():
scheduler = _make_partial_tail_scheduler()
request = _make_partial_tail_request(scheduler)
request.skip_reading_prefix_cache = False
scheduler.manager.lookup.return_value = LookupResult.HIT
tokens, load_async = scheduler.get_num_new_matched_tokens(
request, 0, max_num_new_tokens=20
)
assert (tokens, load_async) == (20, True)
assert scheduler._req_status["req"].partial_tail_boundary == 20
def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary():
"""A recurrent group may hold unhashed blocks inside the computed prefix.
tests/v1/kv_connector/unit/test_hisparse_connector.py
+116
View File
@@ -0,0 +1,116 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
import torch
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import (
ExampleConnector,
)
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
HiSparseConnector,
HiSparseConnectorMetadata,
HiSparseConnectorScheduler,
)
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
HiSparseConnectorWorker,
)
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector
from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
from vllm.v1.hisparse.runtime import HiSparseCacheHandle
from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
@pytest.mark.parametrize("nested", [False, True])
def test_cache_manager_binding_preserves_hisparse_and_legacy_pool_hooks(nested):
"""Composite binding must reach HiSparse and existing pool-only connectors."""
hisparse = object.__new__(HiSparseConnector)
hisparse._role = KVConnectorRole.SCHEDULER
hisparse.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False)
legacy = object.__new__(ExampleConnector)
legacy.bind_gpu_block_pool = MagicMock()
connector = object.__new__(MultiConnector)
connector._connectors = [hisparse, legacy]
if nested:
parent = object.__new__(MultiConnector)
parent._connectors = [connector]
connector = parent
from tests.v1.core.test_prefix_caching import make_hisparse_kv_cache_manager
manager = make_hisparse_kv_cache_manager(16, 16)
connector.bind_kv_cache_manager(manager)
assert hisparse.connector_scheduler.coordinator is get_hisparse_coordinator(manager)
legacy.bind_gpu_block_pool.assert_called_once_with(manager.block_pool)
def test_hisparse_requires_block_outermost_device_layout():
assert HiSparseConnector.get_required_kvcache_layout(MagicMock()) == "BLHNC"
def test_no_forward_enqueues_deferred_hisparse_transfers():
"""A zero-token step must still enqueue deferred post-forward transfers."""
connector = object.__new__(ActiveKVConnector)
connector._disabled = False
connector.pre_forward = MagicMock()
connector.finish_forward = MagicMock()
connector.post_forward = MagicMock(return_value=None)
scheduler_output = SimpleNamespace(finished_req_ids=set())
connector.no_forward(scheduler_output)
connector.pre_forward.assert_called_once_with(scheduler_output)
connector.finish_forward.assert_called_once_with()
def test_full_graph_step_prepares_host_mirror_outside_model():
"""Graph replay must restore host-mirror state cleared at step start."""
runtime = SimpleNamespace(
is_group_leader=True,
eager_host_mirror=True,
begin_forward=MagicMock(),
invalidate_written_slots=MagicMock(),
)
handle = HiSparseCacheHandle(runtime)
handle.mirror_slot_mapping = torch.tensor([4, 5])
worker = object.__new__(HiSparseConnectorWorker)
worker.cache_layer_names = ["layer"]
worker.cache_handles = [handle]
worker._group_leaders = (("layer", handle),)
worker._per_layer_mirrored = set()
worker._submitted_mirror_layers = set()
worker.is_host_writer = True
worker._enqueue_row_dma = MagicMock()
worker.start_step = MagicMock(
side_effect=lambda *_args, **_kwargs: worker._clear_forward_mirror_state()
)
connector = object.__new__(HiSparseConnector)
connector.connector_worker = worker
connector._get_connector_metadata = MagicMock(
return_value=HiSparseConnectorMetadata(None, (), (), {}, True)
)
req_id_per_token = torch.tensor([0, 1], dtype=torch.int32)
attn_metadata = SimpleNamespace(
num_actual_tokens=2,
num_decode_tokens=2,
num_reqs=2,
max_query_len=1,
req_id_per_token=req_id_per_token,
)
connector.start_load_kv(
SimpleNamespace(),
request_state_indices=None,
request_ids=[],
attn_metadata={"layer": attn_metadata},
)
worker._enqueue_host_mirror()
worker._enqueue_row_dma.assert_called_once_with((0,), ready_event=None)
runtime.invalidate_written_slots.assert_called_once()
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+85
View File
@@ -3297,7 +3297,9 @@ def _make_bare_worker(
# path; everything flows through the coordinator).
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheTensor,
)
worker.disk_offload_buffer_budget_bytes = None
@@ -3311,6 +3313,18 @@ def _make_bare_worker(
block_size=block_size, num_kv_heads=8, head_size=64, dtype=None
)
group = KVCacheGroupSpec(["layer0", "__cross_layer__"], spec)
worker._kv_cache_config = KVCacheConfig(
num_blocks=num_gpu_blocks,
kv_cache_tensors=[
KVCacheTensor(
size=0,
layers=["layer0", "__cross_layer__"],
layer_stride=0,
block_stride=0,
)
],
kv_cache_groups=[group],
)
worker._kv_cache_groups = [group]
worker.pcp_size = 1
worker.dcp_size = 1
@@ -4059,6 +4073,77 @@ def test_register_kv_caches_shared_storage(layout: KVCacheLayout):
worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes)
def test_register_kv_caches_uses_transfer_group_memory_domain():
"""Derived GPU caches must not change host-source transfer addresses."""
from vllm.v1.kv_cache_interface import (
KVCacheConfig,
KVCacheGroupRole,
KVCacheGroupSpec,
KVCacheTensor,
)
host_num_blocks = 4
gpu_num_blocks = 2
page_size = 32
source_spec = FullAttentionSpec(
block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float16
)
indexer_spec = FullAttentionSpec(
block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float16
)
source_group = KVCacheGroupSpec(
["source"],
source_spec,
host_resident=True,
)
indexer_group = KVCacheGroupSpec(
["indexer"],
indexer_spec,
enable_kv_transfer=False,
role=KVCacheGroupRole.HISPARSE_INDEXER,
)
config = KVCacheConfig(
num_blocks=gpu_num_blocks,
kv_cache_tensors=[
KVCacheTensor(
size=host_num_blocks * page_size,
layers=["source"],
layer_stride=host_num_blocks * page_size,
block_stride=page_size,
host_resident=True,
),
KVCacheTensor(
size=gpu_num_blocks * page_size,
layers=["indexer"],
layer_stride=gpu_num_blocks * page_size,
block_stride=page_size,
),
],
kv_cache_groups=[source_group, indexer_group],
hisparse_host_num_blocks=host_num_blocks,
)
worker = _make_bare_worker(num_gpu_blocks=gpu_num_blocks)
worker._kv_cache_config = config
worker._kv_cache_groups = list(config.transfer_groups)
worker.token_dbs = [
ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), block_size=16)
]
source = torch.zeros(host_num_blocks, page_size, dtype=torch.uint8)
indexer = torch.zeros(gpu_num_blocks, page_size, dtype=torch.uint8)
_register_with_mocked_threads(
worker,
{"source": source, "indexer": indexer},
)
db = worker.token_dbs[0]
assert db.kv_caches_base_addr == [source.data_ptr()]
assert db.block_len == [page_size]
worker.store.register_buffer.assert_called_once_with(
source.data_ptr(), source.nbytes
)
def test_register_kv_caches_separate_head_groups():
# LHBNC gives each head group its own region; the K/V split doubles heads.
layout = KVCacheLayout.LHBNC
tests/v1/kv_connector/unit/test_multi_connector.py
+6 -6
View File
@@ -304,10 +304,10 @@ def test_multi_example_connector_consistency():
events = get_connector_events()
storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])
storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])
# Initial events bind the block pool, query completion counts, and exchange
# Initial events bind the cache manager, query completion counts, and exchange
# handshake metadata before the request is enqueued.
assert storage1_scheduler_events[:7] == [
"bind_gpu_block_pool",
"bind_kv_cache_manager",
"get_finished_count",
"set_xfer_handshake_metadata_pp_aware",
"on_new_request",
@@ -315,8 +315,8 @@ def test_multi_example_connector_consistency():
"update_state_after_alloc num_blocks=[7] 0",
"build_connector_meta",
]
# First three events are from initialization. During generate(), layer hooks
# run before the deferred load is started after the forward pass.
# First three events are from initialization. Layer hooks run before the
# deferred load starts after the forward pass.
expected_worker_prefix = [
"register_kv_caches",
"set_host_xfer_buffer_ops",
@@ -328,12 +328,12 @@ def test_multi_example_connector_consistency():
]
for connector_name in ("storage1-WORKER", "storage2-WORKER"):
worker_events = events[connector_name]
assert worker_events[:7] == expected_worker_prefix
assert worker_events[: len(expected_worker_prefix)] == expected_worker_prefix
assert worker_events.index("start_load_kv") > worker_events.index(
"save_kv_layer"
)
assert storage2_scheduler_events[:7] == [
"bind_gpu_block_pool",
"bind_kv_cache_manager",
"get_finished_count",
"set_xfer_handshake_metadata_pp_aware",
"on_new_request",
tests/v1/kv_connector/unit/test_nixl_connector.py
+140 -38
View File
@@ -29,7 +29,10 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (
get_current_attn_backend,
)
from vllm.distributed.kv_transfer.kv_connector.v1 import nixl
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
KVConnectorRole,
KVConnectorTransferResults,
)
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
MultiKVConnectorStats,
@@ -518,6 +521,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
slot_size_bytes = 4096
self.slot_size_per_layer = [slot_size_bytes]
self.block_len_per_layer = [slot_size_bytes * self.block_size]
self.num_blocks = self.kv_cache_config.num_blocks
self.num_regions = 1
self.block_stride_per_layer = list(self.block_len_per_layer)
self.region_num_blocks = [self.num_blocks]
@@ -635,6 +639,12 @@ class TestNixlHandshake:
if pcp_rank > 0:
assert connector.get_finished(set()) == (set(), set())
worker.get_transfer_results = MagicMock(
return_value=KVConnectorTransferResults(finished_sending={"sent"})
)
results = connector.get_transfer_results(set())
assert results.finished_sending == ({"sent"} if pcp_rank == 0 else set())
@patch(
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
FakeNixlWrapper,
@@ -1195,6 +1205,7 @@ class TestNixlHandshake:
ssm_sizes=(0, 0),
attn_backend_name=worker.backend_name,
physical_blocks_per_logical_kv_block=1,
region_num_blocks=None,
)
assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [
@@ -1435,9 +1446,9 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
metadata = NixlConnectorMetadata()
metadata.add_new_req_to_recv(
request_id=request_id,
local_block_ids=([1, 2, 3],),
local_block_ids=([0],),
kv_transfer_params={
"remote_block_ids": ([4, 5, 6],),
"remote_block_ids": ([0],),
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
"remote_request_id": f"prefill-{request_id}",
"remote_host": "localhost",
@@ -1482,6 +1493,10 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
assert stats_after_reset is None
@patch(
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
FakeNixlWrapper,
)
def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist_init):
"""reqs_to_send deadlines are stamped with the scheduler process's
perf_counter, whose epoch differs across processes and (by boot-time
@@ -1906,9 +1921,10 @@ def test_mixed_memory_local_descriptors_split_by_memory_type():
assert handle == 22
assert worker._dram_src_handles_by_block_size[worker.block_size] == 11
assert [block[2] for block in blocks] == [0, 0, 3, 3]
assert [
memory_types = [
call.args[1] for call in worker.nixl_wrapper.get_xfer_descs.call_args_list
] == ["DRAM", "VRAM"]
]
assert memory_types == ["DRAM", "VRAM"]
def test_mixed_memory_read_notifies_after_both_transfers_finish():
@@ -1919,6 +1935,7 @@ def test_mixed_memory_read_notifies_after_both_transfers_finish():
worker._recving_metadata = {"request": MagicMock()}
worker._recving_transfers = defaultdict(list)
worker._pending_recv_notifs = {}
worker._failed_inflight_recvs = set()
worker.xfer_stats = MagicMock()
worker.nixl_wrapper = MagicMock()
worker.nixl_wrapper.make_prepped_xfer.side_effect = [101, 102]
@@ -1959,6 +1976,7 @@ def test_mixed_memory_read_failure_does_not_notify_producer():
worker._recving_metadata = {"request": MagicMock()}
worker._recving_transfers = {"request": [101, 102]}
worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]}
worker._failed_inflight_recvs = set()
worker._is_hma_required = True
worker._failed_recv_reqs = MagicMock()
worker._log_failure = MagicMock() # type: ignore[method-assign]
@@ -2871,9 +2889,11 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
FailingNixlWrapper,
)
def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init):
"""Test that transfer setup failures mark blocks invalid
and return via get_finished."""
@pytest.mark.parametrize("is_hma", [False, True])
def test_transfer_setup_failure_returns_finished(
default_vllm_config, dist_init, is_hma
):
"""Setup failures report the request; only non-HMA reports block IDs."""
vllm_config = create_vllm_config()
connector = NixlConnector(
@@ -2883,6 +2903,7 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
vllm_config, connector.engine_id, hand_shake_latency=0
)
connector.connector_worker.nixl_wrapper.fail_transfer_setup = True
connector.connector_worker._is_hma_required = is_hma
request_id = "test_transfer_fail"
metadata = NixlConnectorMetadata()
@@ -2912,13 +2933,106 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
time.sleep(0.1)
connector.start_load_kv(dummy_ctx)
# check that blocks were marked invalid
results = connector.get_transfer_results(finished_req_ids=set())
assert request_id in results.finished_recving
assert results.failed_recving == {request_id}
invalid_blocks = connector.get_block_ids_with_load_errors()
assert invalid_blocks == {7, 8, 9}
assert invalid_blocks == (set() if is_hma else {7, 8, 9})
# ensure request appears in get_finished
class _ScriptedXferWrapper(FakeNixlWrapper):
"""Scripts per-handle xfer states; forbids releasing in-flight handles."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# handle -> successive states; the last one repeats.
self.xfer_states: dict[int, list[str]] = {}
self.released: list[int] = []
def check_xfer_state(self, handle: int) -> str:
states = self.xfer_states[handle]
return states.pop(0) if len(states) > 1 else states[0]
def release_xfer_handle(self, handle: int) -> None:
# A posted-but-unfinished transfer cannot be aborted: releasing it
# leaves the RDMA READ armed. Production code must never do this.
assert self.xfer_states.get(handle, ["DONE"])[0] != "PROC", (
f"released in-flight handle {handle}"
)
self.released.append(handle)
def _make_split_read_connector(vllm_config, request_id, states):
"""Seed a request with scripted in-flight xfer handles + recv metadata."""
connector = NixlConnector(
vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
)
connector.connector_worker = FakeNixlConnectorWorker(
vllm_config, connector.engine_id, hand_shake_latency=0
)
worker = connector.connector_worker
wrapper = _ScriptedXferWrapper("agent")
worker.nixl_wrapper = wrapper
metadata = NixlConnectorMetadata()
metadata.add_new_req_to_recv(
request_id=request_id,
local_block_ids=([7, 8, 9],),
kv_transfer_params={
"remote_block_ids": ([10, 11, 12],),
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
"remote_request_id": f"prefill-{request_id}",
"remote_host": "localhost",
"remote_port": 1234,
"remote_tp_size": 1,
},
)
worker._recving_metadata[request_id] = metadata.reqs_to_recv[request_id]
wrapper.xfer_states = dict(states)
worker._recving_transfers[request_id] = list(states)
return connector, worker, wrapper
@patch(
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
FakeNixlWrapper,
)
def test_split_read_failure_defers_report_until_last_handle(
default_vllm_config, dist_init
):
"""One half of a split (mixed DRAM/VRAM) read failing must not report the
request — nor invalidate its blocks — while the sibling xfer is still in
flight: a posted READ cannot be aborted and would DMA into blocks the
scheduler could free and reuse. The report happens exactly once, when the
last handle is terminal."""
request_id = "split_read_partial_failure"
err_handle, live_handle = 11, 22
connector, worker, wrapper = _make_split_read_connector(
create_vllm_config(),
request_id,
{err_handle: ["ERR"], live_handle: ["PROC", "PROC", "DONE"]},
)
# Poll 1: one half fails; the sibling is in flight -> nothing reported.
_, done_recving = connector.get_finished(finished_req_ids=set())
assert request_id in done_recving
assert done_recving == set()
assert connector.get_block_ids_with_load_errors() == set()
assert request_id in worker._recving_metadata
assert wrapper.released == [err_handle]
# Poll 2: sibling still in flight.
_, done_recving = connector.get_finished(finished_req_ids=set())
assert done_recving == set()
# Poll 3: sibling terminal -> reported exactly once, blocks invalidated.
_, done_recving = connector.get_finished(finished_req_ids=set())
assert done_recving == {request_id}
assert connector.get_block_ids_with_load_errors() == {7, 8, 9}
assert request_id not in worker._recving_metadata
assert wrapper.released == [err_handle, live_handle]
# Poll 4: nothing left; no double report.
_, done_recving = connector.get_finished(finished_req_ids=set())
assert done_recving == set()
@patch(
@@ -2938,7 +3052,7 @@ def test_failed_request_skips_kv_postprocessing(
default_vllm_config, dist_init, failure_mode
):
"""Test that failed requests skip KV sync and post-processing in
get_finished().
get_transfer_results().
This is the core safety behavior: when a KV transfer fails at any stage,
the request must still appear in done_recving (so the scheduler can apply
@@ -3023,11 +3137,10 @@ def test_failed_request_skips_kv_postprocessing(
patch.object(worker, "sync_recved_kv_to_device") as mock_sync,
patch.object(worker, "post_process_device_kv_on_receive") as mock_postprocess,
):
_, done_recving = connector.get_finished(finished_req_ids=set())
results = connector.get_transfer_results(finished_req_ids=set())
# The failed request must appear in done_recving so the scheduler
# can handle it (e.g., trigger recompute via kv_load_failure_policy).
assert request_id in done_recving
assert request_id in results.finished_recving
assert results.failed_recving == {request_id}
# Critical: KV sync and post-processing must NOT have been called
# since no valid KV data was received for the failed request.
@@ -3049,19 +3162,7 @@ def test_failed_request_skips_kv_postprocessing(
def test_handles_failing_in_separate_polls_do_not_kill_the_engine(
default_vllm_config, dist_init
):
"""A request whose handles fail in different polls must not crash the engine.
One transfer handle is created per remote rank, and when a peer goes away
they do not all fail in the same poll: one errors while another is still
PROC. The first failure reports the request via _failed_recv_reqs and
get_finished() pops its metadata; when the remaining handles fail in a
later poll, the request must be cleaned up without being reported again.
Reporting twice kills the EngineCore: the scheduler's assert in
_update_from_kv_xfer_finished only expects a finished recv for a request
still waiting for KVs, having moved this one out of
WAITING_FOR_REMOTE_KVS to recompute locally after the first report.
"""
"""Failures across polls report once, after the last writer is terminal."""
vllm_config = create_vllm_config()
connector = NixlConnector(
vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
@@ -3107,20 +3208,21 @@ def test_handles_failing_in_separate_polls_do_not_kill_the_engine(
):
_, done_recving = connector.get_finished(finished_req_ids=set())
assert request_id in done_recving
assert request_id not in worker._recving_metadata
assert request_id not in done_recving
assert request_id in worker._recving_metadata
assert worker._recving_transfers[request_id] == [second]
assert connector.get_block_ids_with_load_errors() == set()
# Poll 2: the second handle fails too. The request was already reported,
# so it must not be reported a second time: the scheduler has since moved
# it out of WAITING_FOR_REMOTE_KVS to recompute locally, and asserts on a
# finished recv for a request in that state. Its remaining handles are
# still released and the transfer entry removed.
# Only the last terminal handle permits the scheduler to reuse the blocks.
with patch.object(worker.nixl_wrapper, "check_xfer_state", return_value="ERR"):
_, done_recving = connector.get_finished(finished_req_ids=set())
assert request_id not in done_recving
assert done_recving == {request_id}
assert request_id not in worker._recving_transfers
assert request_id not in worker._recving_metadata
assert connector.get_block_ids_with_load_errors() == {1, 2, 3}
_, done_recving = connector.get_finished(finished_req_ids=set())
assert done_recving == set()
def _set_test_speculative_config(
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+188 -28
View File
@@ -5,16 +5,31 @@
import gc
from unittest.mock import MagicMock, patch
import msgspec
import pytest
import torch
from tests.v1.attention.utils import MockMambaBuilder
from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig
from vllm.config import KVTransferConfig, set_current_vllm_config
from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
NixlAgentMetadata,
)
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
NixlConnectorWorker,
)
from vllm.v1.core.single_type_kv_cache_manager import (
FullAttentionManager,
SlidingWindowManager,
)
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheTensor,
MLAAttentionSpec,
)
from .utils import (
create_request,
@@ -272,26 +287,24 @@ def test_logical_to_kernel_block_ids_with_hma():
@pytest.mark.cpu_test
@pytest.mark.parametrize(
"is_rocm,has_mamba,use_host_buffer,done_recving,failed_recving,expected_syncs",
"is_rocm,has_mamba,use_host_buffer,done_recving,expected_syncs",
[
(True, True, False, {"req"}, set(), 1),
(False, True, False, {"req"}, set(), 0),
(True, False, False, {"req"}, set(), 0),
(True, True, True, {"req"}, set(), 0),
(True, True, False, set(), set(), 0),
(True, True, False, {"req"}, {"req"}, 0),
(True, True, False, {"req"}, 1),
(False, False, False, {"req"}, 0),
(False, True, False, {"req"}, 0),
(True, True, True, {"req"}, 0),
(True, True, False, set(), 0),
],
)
def test_sync_device_after_mamba_recv_gates(
def test_sync_device_after_direct_recv_gates(
monkeypatch,
is_rocm,
has_mamba,
use_host_buffer,
done_recving,
failed_recving,
expected_syncs,
):
"""Only direct-GPU Mamba receives on ROCm need a device fence."""
"""Only backends requiring visibility fences synchronize direct receives."""
from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
NixlConnectorWorker,
@@ -309,7 +322,7 @@ def test_sync_device_after_mamba_recv_gates(
lambda: sync_calls.append(True),
)
worker._sync_device_after_mamba_recv(done_recving, failed_recving)
worker._sync_device_after_direct_recv(done_recving)
assert len(sync_calls) == expected_syncs
@@ -399,6 +412,8 @@ def test_read_blocks_for_req_expands_remote_ids(
)
remote_engine_id = "remote-engine"
worker.region_group_ids = list(range(len(resolved_types)))
worker.dst_region_group_ids = {remote_engine_id: list(worker.region_group_ids)}
worker.transfer_topo = MagicMock()
# tp_ratio not exercised (all_source_ranks is empty so no reads run),
@@ -888,6 +903,45 @@ def _make_mock_worker_for_desc_ids(
return worker
@pytest.mark.cpu_test
def test_shared_nixl_region_broadcasts_every_group_block_table():
"""A shared region must receive every group's block table exactly once."""
worker = _make_mock_worker_for_desc_ids(
num_regions=3,
has_mamba=False,
group_spec_types=(FullAttentionSpec,) * 3,
block_len_per_layer=[100] * 3,
)
worker.region_group_ids = [0, bw._SHARED_REGION_GROUP_ID, 2]
worker._uses_region_group_mapping = True
result = worker._compute_desc_ids(
block_ids=([1], [2], [3]),
dst_num_blocks=10,
block_size_ratio=None,
physical_blocks_per_logical=1,
region_num_blocks=[10, 10, 10],
)
assert result.tolist() == [1, 11, 12, 13, 23]
@pytest.mark.cpu_test
def test_nixl_region_sort_preserves_pipeline_layer_order():
"""Numeric layer order prevents PP peers from pairing different regions."""
names = [
"model.layers.10.self_attn",
"model.layers.2.self_attn",
"model.layers.1.self_attn",
]
assert sorted(names, key=bw._region_sort_key) == [
"model.layers.1.self_attn",
"model.layers.2.self_attn",
"model.layers.10.self_attn",
]
@pytest.mark.cpu_test
def test_get_block_descs_ids_hybrid_ssm():
"""Test _compute_desc_ids uses per-group strides for hybrid
@@ -1097,10 +1151,7 @@ def test_post_process_zeroes_untransferred_tail():
worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache}
fa_group = MagicMock(layer_names=["attn.0"])
ssm_group = MagicMock(layer_names=["mamba.0"])
worker.kv_cache_config = MagicMock(
kv_cache_groups=[fa_group, ssm_group],
transfer_groups=[fa_group, ssm_group],
)
worker.kv_cache_config = MagicMock(transfer_groups=[fa_group, ssm_group])
# The cached property filters mamba layers out of the permuted caches.
attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker)
assert len(attn_caches) == 1 and attn_caches[0] is attn_cache
@@ -1188,6 +1239,11 @@ def _make_fake_kv_cache_manager():
manager = object.__new__(KVCacheManager)
manager.coordinator = MagicMock()
manager.kv_cache_config = MagicMock()
manager.kv_cache_config.kv_cache_groups = (
MagicMock(host_resident=False),
MagicMock(host_resident=False),
)
manager.coordinator.single_type_managers = (
_FakeSingleTypeManager(True, 16, [10, 11, 12, 13, 14, 15]), # attention
_FakeSingleTypeManager(False, 16, [20, 21, 22, 23, 24, 25]), # mamba
@@ -1723,16 +1779,6 @@ def test_exchange_clipped_blocks_ssm_positional_states():
assert clipped == ([1, 2, 3], [0, 5, 6, 7])
@pytest.mark.cpu_test
def test_exchange_clipped_blocks_excludes_nontransfer_groups():
"""Scheduler block tables must match the worker's registered regions."""
sched = make_nixl_scheduler()
sched.kv_cache_config = make_kv_cache_config(block_size=16, swa_enabled=True)
sched.kv_cache_config.kv_cache_groups[1].enable_kv_transfer = False
assert sched.get_exchange_clipped_blocks(([1, 2], [9])) == ([1, 2],)
# ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ─────────────────────
@@ -1786,6 +1832,95 @@ def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4):
)
@pytest.mark.cpu_test
@pytest.mark.parametrize("kernel_block_size", [16, 8])
def test_nixl_keeps_device_block_count_with_hisparse_host_pool(kernel_block_size):
host_num_blocks = 4
gpu_num_blocks = 9
spec = MLAAttentionSpec(
block_size=16,
num_kv_heads=1,
head_size=8,
dtype=torch.float16,
)
kv_cache_config = KVCacheConfig(
num_blocks=gpu_num_blocks,
hisparse_host_num_blocks=host_num_blocks,
kv_cache_tensors=[
KVCacheTensor(
size=host_num_blocks * spec.page_size_bytes,
layers=["mla.host"],
layer_stride=host_num_blocks * spec.page_size_bytes,
block_stride=spec.page_size_bytes,
host_resident=True,
),
KVCacheTensor(
size=gpu_num_blocks * spec.page_size_bytes,
layers=["mla.device"],
layer_stride=gpu_num_blocks * spec.page_size_bytes,
block_stride=spec.page_size_bytes,
),
],
kv_cache_groups=[
KVCacheGroupSpec(
["mla.host"],
spec,
host_resident=True,
),
KVCacheGroupSpec(["mla.device"], spec),
],
)
vllm_config = create_vllm_config(block_size=16)
vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
fake_backend = MagicMock()
fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size]
fake_backend.get_name.return_value = "FLASHMLA"
fake_backend.full_cls_name.return_value = "fake.FLASHMLA"
fake_backend.get_kv_cache_shape.return_value = (1, kernel_block_size, 1, 1)
fake_platform = MagicMock()
fake_platform.device_type = "cuda"
fake_platform.get_nixl_memory_type.return_value = "VRAM"
physical_per_logical = spec.block_size // kernel_block_size
physical_page_bytes = spec.page_size_bytes // physical_per_logical
host_cache = torch.zeros(
host_num_blocks * physical_per_logical, physical_page_bytes, dtype=torch.uint8
)
device_cache = torch.zeros(
gpu_num_blocks * physical_per_logical, physical_page_bytes, dtype=torch.uint8
)
with (
patch.object(bw, "NixlWrapper"),
patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
patch.object(bw, "current_platform", fake_platform),
set_current_vllm_config(vllm_config),
):
worker = NixlConnectorWorker(vllm_config, "test-engine", kv_cache_config)
worker.use_mla = True
worker.nixl_wrapper.get_agent_metadata.return_value = b"metadata"
worker.register_kv_caches({"mla.host": host_cache, "mla.device": device_cache})
assert worker.num_blocks == gpu_num_blocks * physical_per_logical
assert dict(zip(worker.region_group_ids, worker.region_num_blocks)) == {
0: host_num_blocks * physical_per_logical,
1: gpu_num_blocks * physical_per_logical,
}
caches = {"mla.host": host_cache, "mla.device": device_cache}
expected_addrs = [
caches[name].data_ptr() + block * stride
for name, stride, count in zip(
worker.region_names,
worker.block_stride_per_layer,
worker.region_num_blocks,
strict=True,
)
for block in range(count)
]
assert worker.src_blocks_data[:, 0].tolist() == expected_addrs
@pytest.mark.cpu_test
def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
"""Hybrid MLA+KDA registration: HMA tensors shared by both layer types
@@ -1803,6 +1938,9 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
kv_cache_config = _make_hybrid_mla_kv_cache_config()
unified_page = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes
vllm_config = create_vllm_config(block_size=12)
# The engine rewrites this global value to the smallest scheduler-group
# block size. NIXL must retain the transferable groups' 12-token geometry.
vllm_config.cache_config.block_size = 4
# kv_buffer_device defaults to the *real* platform's device type, which on
# a CPU-only test host would make this a host-buffer worker: host xfer
# buffers are per-layer, so the HMA shared tensors would not be
@@ -1833,7 +1971,13 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
worker.use_mla = True # opt-125m test config is not MLA; force the flag
worker.nixl_wrapper.get_agent_metadata.return_value = b"fake-agent-metadata"
tensors = [torch.zeros(4 * unified_page, dtype=torch.uint8) for _ in range(2)]
kernel_page = unified_page // 3
block_stride = 2 * kernel_page
backing = torch.zeros(12, block_stride, dtype=torch.uint8)
tensors = [
backing[:, :kernel_page],
backing[:, kernel_page:],
]
# KDA layer first per tensor: exercises the dual-purpose flag merge.
worker.register_kv_caches(
{
@@ -1849,6 +1993,8 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
# 12-token logical blocks over the 4-token MLA kernel block.
assert worker._physical_blocks_per_logical_kv_block == 3
assert worker.block_size == 4 and worker.num_blocks == 12
assert worker.region_group_ids == [bw._SHARED_REGION_GROUP_ID] * 2
assert worker.block_stride_per_layer == [block_stride, block_stride]
# Both shared tensors are dual-purpose: their FA view is MLA even though
# a KDA layer registered them first.
assert worker._region_is_mla == [True, True]
@@ -1861,8 +2007,20 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
# Mamba descs: 2 regions x (3 conv sub-projections + 1 ssm) x 4 blocks.
assert worker.src_blocks_data.shape == (24 + 32, 3)
fa_descs = worker.src_blocks_data[:24]
assert fa_descs[1][0] - fa_descs[0][0] == unified_page // 3
expected_addrs = [
tensor.data_ptr() + block * block_stride
for tensor in tensors
for block in range(12)
]
assert fa_descs[:, 0].tolist() == expected_addrs
assert all(size == unified_page // 3 for size in fa_descs[:, 1])
worker.nixl_wrapper.register_memory.assert_called_once()
metadata = msgspec.msgpack.decode(
worker.xfer_handshake_metadata.agent_metadata_bytes,
type=NixlAgentMetadata,
)
assert metadata.block_strides == [block_stride, block_stride]
assert metadata.region_num_blocks == [12, 12]
@pytest.mark.cpu_test
@@ -1909,6 +2067,8 @@ def test_push_write_hybrid_mla_replicates_attention():
worker.dst_xfer_side_handles = {engine_id: {0: 100, 1: 101}}
worker.src_xfer_handles_by_tp_ratio = {(-2, 4): [200, 201]}
worker.src_xfer_handles_by_block_size = {4: 300}
worker.region_group_ids = [0, 1]
worker.dst_region_group_ids = {engine_id: [0, 1]}
worker._sending_transfers = defaultdict(list)
worker._sending_transfers_lock = threading.Lock()
worker.kv_cache_config = _make_hybrid_mla_kv_cache_config()
tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+19 -8
View File
@@ -186,11 +186,16 @@ def test_overlaid_transfer_groups_share_region_geometry():
worker.host_buffer_kv_cache_layout = "NHD"
worker._physical_blocks_per_logical_kv_block = 1
worker._logical_num_blocks = num_blocks
worker.region_mem_types = []
worker.region_group_ids = []
worker.region_mem_types = []
worker._mixed_mem_types = False
worker.region_names = []
worker.region_num_blocks = []
worker._mixed_mem_types = False
worker._desc_is_dram_by_block_size = {}
worker._desc_pos_by_block_size = {}
worker._dram_src_handles_by_block_size = {}
worker._region_is_mla = []
worker.block_len_per_layer = []
worker.block_stride_per_layer = []
@@ -304,8 +309,6 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo
# would not be deduplicated. Pin it to the faked device type.
vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
from unittest.mock import MagicMock
fake_backend = MagicMock()
fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size]
fake_backend.get_name.return_value = "FLASHMLA"
@@ -552,6 +555,16 @@ def _make_remote_meta(
)
def _register_remote_agents(worker, metadata, tp_size):
"""Mirror the async handshake callback that publishes prepared agents."""
worker._remote_agents[metadata.engine_id] = {
(0, rank): worker.add_remote_agent(
metadata, remote_tp_rank=rank, remote_tp_size=tp_size
)
for rank in range(tp_size)
}
def _owned_byte_ranges(worker, group_logical_ids):
"""Byte ranges owned by a request: for each HMA region tensor, every
logical block id of every group maps to one unified page."""
@@ -609,8 +622,7 @@ def test_hetero_ppl_multi_read_writes_stay_within_request_blocks():
remote_num_logical=12,
remote_ssm_sizes=(24, 32),
)
for rank in (0, 1):
worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=2)
_register_remote_agents(worker, meta_r, 2)
# Request B: 17 matched tokens. Local: 2 logical blocks (24 tok
# capacity); remote: 16 prefilled tokens -> 2 remote logical blocks.
@@ -626,7 +638,7 @@ def test_hetero_ppl_multi_read_writes_stay_within_request_blocks():
"remote_block_ids": remote_ids,
"remote_engine_id": "remote-engine",
"remote_request_id": "prefill-req-b",
"remote_host": "localhost",
"remote_host": "remote-host",
"remote_port": 1234,
"tp_size": 2,
},
@@ -709,8 +721,7 @@ def _run_hetero_case(
remote_num_logical=max(2 * n_remote + 4, 8),
remote_ssm_sizes=(48 // tp_size, 64 // tp_size),
)
for rank in range(tp_size):
worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=tp_size)
_register_remote_agents(worker, meta_r, tp_size)
# Sparse ids so neighbors exist between the request's blocks.
local_attn = [2 * i + 1 for i in range(n_local)]
@@ -726,7 +737,7 @@ def _run_hetero_case(
"remote_block_ids": remote_ids,
"remote_engine_id": "remote-engine",
"remote_request_id": "prefill-req-b",
"remote_host": "localhost",
"remote_host": "remote-host",
"remote_port": 1234,
"tp_size": tp_size,
},
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+54 -22
View File
@@ -15,7 +15,7 @@ requiring a real NIXL agent or network:
* The worker matches D registrations against P finished blocks (both
scenario directions) and forwards non-PUSH_REG NIXL notifs to the main
thread's ``_get_new_notifs``.
* ``get_finished`` enqueues evictions for the writer.
* ``get_transfer_results`` enqueues evictions for the writer.
"""
from __future__ import annotations
@@ -32,6 +32,9 @@ from unittest.mock import MagicMock, patch
import msgspec
import pytest
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
KVConnectorTransferResults,
)
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
PUSH_REG_NOTIF_PREFIX,
NixlAgentMetadata,
@@ -347,6 +350,8 @@ class _StubWriterWorker(NixlPushConnectorWorker):
# Base worker fields touched by start_load_kv / _get_new_notifs.
w._recving_metadata = {}
w._pending_recv_notifs = {}
w._failed_inflight_recvs = set()
w._recving_transfers = defaultdict(list)
w._reqs_to_process = set()
w._reqs_to_send = {}
@@ -742,21 +747,19 @@ class TestPushWriterNotifs:
assert notified == set()
assert request_id in w._recving_transfers
def test_get_finished_evicts_completed_state(self):
"""``get_finished`` should enqueue evictions and wake the writer."""
def test_get_transfer_results_evicts_completed_state(self):
"""Transfer completion should enqueue evictions and wake the writer."""
w = _StubWriterWorker.fresh()
# Stub the base ``get_finished`` to return one done_sending entry.
# Patch via the MRO's parent class.
with patch.object(
NixlPushConnectorWorker.__mro__[1],
"get_finished",
return_value=({"req-done"}, set()),
"get_transfer_results",
return_value=KVConnectorTransferResults(finished_sending={"req-done"}),
):
done_sending, done_recving = w.get_finished()
results = w.get_transfer_results()
assert "req-done" in done_sending
assert done_recving == set()
assert results.finished_sending == {"req-done"}
assert results.finished_recving == set()
# Eviction enqueued for the writer.
evicted = []
while True:
@@ -928,19 +931,21 @@ class TestPushWriterNegative:
assert len(w._pending_d_registrations) == 1
assert w.start_push_calls == []
def test_get_finished_enqueues_eviction_for_each_done_request(self):
"""``get_finished`` must enqueue an eviction for every request
def test_get_transfer_results_enqueues_each_completed_request(self):
"""``get_transfer_results`` must enqueue every completed request
in ``done_sending`` so the writer can drop stale matching state.
Unlike the happy-path test, this verifies the *cardinality*: N
completed requests -> N evictions, in order."""
w = _StubWriterWorker.fresh()
with patch.object(
NixlPushConnectorWorker.__mro__[1],
"get_finished",
return_value=({"req-1", "req-2", "req-3"}, set()),
"get_transfer_results",
return_value=KVConnectorTransferResults(
finished_sending={"req-1", "req-2", "req-3"}
),
):
done_sending, _ = w.get_finished()
assert done_sending == {"req-1", "req-2", "req-3"}
results = w.get_transfer_results()
assert results.finished_sending == {"req-1", "req-2", "req-3"}
evicted: list[str] = []
while True:
@@ -950,23 +955,47 @@ class TestPushWriterNegative:
break
assert sorted(evicted) == ["req-1", "req-2", "req-3"]
def test_get_finished_with_no_completions_does_not_enqueue_eviction(self):
def test_empty_transfer_results_do_not_enqueue_eviction(self):
"""If there's nothing newly done, no eviction should be enqueued.
The wake event IS still set because ``get_finished`` always wakes
the writer to drain notifs."""
w = _StubWriterWorker.fresh()
with patch.object(
NixlPushConnectorWorker.__mro__[1],
"get_finished",
return_value=(set(), set()),
"get_transfer_results",
return_value=KVConnectorTransferResults(),
):
done_sending, done_recving = w.get_finished()
assert done_sending == set()
assert done_recving == set()
results = w.get_transfer_results()
assert results.finished_sending == set()
assert results.finished_recving == set()
assert w._evict_finished_inbox.qsize() == 0
# Wake set so the writer drains NIXL notifs even when idle.
assert w._push_writer_wake.is_set()
def test_failed_push_transfer_finishes_sending(self):
w = _StubWriterWorker.fresh()
w.nixl_wrapper = MagicMock()
w.nixl_wrapper.check_xfer_state.return_value = "ERR"
w.xfer_stats = MagicMock()
w._log_failure = MagicMock()
w._sending_transfers["req-failed"] = [17]
w._reqs_to_send["req-failed"] = time.perf_counter() + 30
w._reqs_to_process.add("req-failed")
with patch.object(
NixlPushConnectorWorker.__mro__[1],
"get_transfer_results",
return_value=KVConnectorTransferResults(),
):
results = w.get_transfer_results()
assert results.finished_sending == {"req-failed"}
assert results.finished_recving == set()
assert "req-failed" not in w._sending_transfers
assert "req-failed" not in w._reqs_to_send
assert "req-failed" not in w._reqs_to_process
w.nixl_wrapper.release_xfer_handle.assert_called_once_with(17)
def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog):
"""A completion notif for a request the worker doesn't know
about should be logged but not crash."""
@@ -1160,6 +1189,7 @@ class TestPushWriterMlaReplication:
)
}
w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids
w.dst_region_group_ids[engine_id] = [0]
w.dst_xfer_side_handles = {engine_id: {r: 1000 + r for r in d_ranks}}
w.src_xfer_handles_by_block_size = {16: 2000}
w._remote_agents = {engine_id: {(0, r): f"agent-{r}" for r in d_ranks}}
@@ -1240,6 +1270,8 @@ class TestPushWriterMlaReplication:
)
}
w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids
w.region_group_ids = [0, 1, 2, 3]
w.dst_region_group_ids[engine_id] = [0, 1, 2, 3]
w.dst_xfer_side_handles = {engine_id: {0: 1000, 1: 1001}}
w.src_xfer_handles_by_tp_ratio = {(-2, 16): [2000, 2001]}
tests/v1/kv_connector/unit/test_output_aggregator.py
+11 -1
View File
@@ -15,12 +15,14 @@ class DummyModelRunnerOutput(ModelRunnerOutput):
finished_sending: set[str] | None = None,
finished_recving: set[str] | None = None,
invalid_block_ids: set[int] | None = None,
failed_recving: set[str] | None = None,
expected_finished_count: int = 0,
):
self.kv_connector_output = KVConnectorOutput(
finished_sending=finished_sending,
finished_recving=finished_recving,
invalid_block_ids=invalid_block_ids or set(),
failed_recving=failed_recving or set(),
expected_finished_count=expected_finished_count,
)
@@ -73,7 +75,9 @@ def test_aggregate_workers_output():
output1 = DummyModelRunnerOutput(invalid_block_ids={3, 4})
output2 = DummyModelRunnerOutput(
finished_recving={"req2"}, invalid_block_ids={4, 5}
finished_recving={"req2"},
invalid_block_ids={4, 5},
failed_recving={"req3"},
)
aggregated = aggregator.aggregate([output1, output2])
@@ -83,6 +87,12 @@ def test_aggregate_workers_output():
assert aggregated.finished_sending is None
assert aggregated.finished_recving == {"req2"}
assert aggregated.invalid_block_ids == {3, 4, 5}
assert not aggregated.failed_recving
output1 = DummyModelRunnerOutput(finished_recving={"req3"})
output2 = DummyModelRunnerOutput(finished_recving={"req3"})
aggregated = aggregator.aggregate([output1, output2])
assert aggregated.kv_connector_output.failed_recving == {"req3"}
def test_aggregate_workers_output_with_expected_finished_count():
tests/v1/kv_connector/unit/utils.py
+1
View File
@@ -539,6 +539,7 @@ def make_nixl_scheduler(
sched._heartbeat_interval = kv_lease_duration // 6
# Fields touched by build_connector_meta / request_finished:
sched._reqs_need_recv = {}
sched._hisparse_host_blocks_to_recv = {}
sched._reqs_need_send = {}
sched._reqs_in_batch = set()
sched._reqs_not_processed = set()
tests/v1/kv_offload/test_factory.py
+3 -3
View File
@@ -62,7 +62,7 @@ def _make_offloading_config(
normalized_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use
if groups is None:
groups = (OffloadingGroupConfig(16, ("layer",)),)
groups = (OffloadingGroupConfig(16, ("layer",), 0),)
return OffloadingConfig(
groups=groups,
@@ -493,8 +493,8 @@ def test_offloading_spec_has_replicated_layout_default():
def test_offloading_spec_uses_normalized_chunk_geometry():
groups = (
OffloadingGroupConfig(12, ("full_layer",)),
OffloadingGroupConfig(16, ("mla_layer",)),
OffloadingGroupConfig(12, ("full_layer",), 0),
OffloadingGroupConfig(16, ("mla_layer",), 1),
)
spec = _create_spec(
groups=groups,
tests/v1/kv_offload/test_file_mapper.py
+4 -1
View File
@@ -26,8 +26,11 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
OffloadingGroupConfig(
tokens_per_block=tokens_per_block,
layer_names=(layer_name,),
group_id=group_id,
)
for group_id, (tokens_per_block, layer_name) in enumerate(
kwargs.get("groups", ())
)
for tokens_per_block, layer_name in kwargs.get("groups", ())
),
worker_kv_bytes_per_block=0,
enable_kv_cache_events=False,
tests/v1/simple_kv_offload/test_worker.py
+6
View File
@@ -400,6 +400,8 @@ def test_register_mixed_page_sizes_in_one_cache_group(monkeypatch):
vllm_config.cache_config = CacheConfig()
vllm_config.cache_config.kv_cache_layout = layout.name
vllm_config.cache_config.num_gpu_blocks_override = None
vllm_config.attention_config.hisparse_config = None
vllm_config.kv_transfer_config = None
pages = [spec.page_size_bytes for spec in specs.values()]
num_blocks = 4
@@ -458,6 +460,8 @@ def test_register_mixed_page_sizes_odd_block_counts(monkeypatch, rank_blocks):
vllm_config.cache_config = CacheConfig()
vllm_config.cache_config.kv_cache_layout = layout.name
vllm_config.cache_config.num_gpu_blocks_override = None
vllm_config.attention_config.hisparse_config = None
vllm_config.kv_transfer_config = None
pages = [spec.page_size_bytes for spec in specs.values()]
kv_cache_config = get_kv_cache_config_from_groups(
@@ -507,6 +511,8 @@ def test_mixed_page_byte_placement_is_dcp_invariant():
vllm_config.cache_config.kv_cache_layout = KVCacheLayout.LBNHC.name
vllm_config.cache_config.num_gpu_blocks_override = None
vllm_config.parallel_config.decode_context_parallel_size = dcp_size
vllm_config.attention_config.hisparse_config = None
vllm_config.kv_transfer_config = None
page_bytes = sum(spec.page_size_bytes for spec in specs.values())
config = get_kv_cache_config_from_groups(vllm_config, [group], page_bytes * 4)
tests/v1/test_kv_cache_spec_registry.py
+15
View File
@@ -16,6 +16,7 @@ from vllm.v1.core.single_type_kv_cache_manager import (
ChunkedLocalAttentionManager,
CrossAttentionManager,
FullAttentionManager,
HiSparseSourceManager,
MambaManager,
SingleTypeKVCacheManager,
SinkFullAttentionManager,
@@ -27,6 +28,7 @@ from vllm.v1.kv_cache_interface import (
CrossAttentionSpec,
FullAttentionSpec,
HiddenStateCacheSpec,
KVCacheGroupRole,
KVCacheSpec,
KVCacheSpecKind,
MambaSpec,
@@ -176,6 +178,19 @@ class TestKVCacheSpecRegistry:
is spec_uniform_base_map[spec_cls]
)
@pytest.mark.parametrize("role", list(KVCacheGroupRole))
def test_mla_manager_selection_by_role(self, role):
"""Only the source role overrides ordinary MLA manager selection."""
expected = (
HiSparseSourceManager
if role == KVCacheGroupRole.HISPARSE_SOURCE
else FullAttentionManager
)
assert (
KVCacheSpecRegistry.get_manager_class(make_spec(MLAAttentionSpec), role)
is expected
)
@pytest.mark.parametrize("spec_cls", list(spec_manager_map))
def test_custom_spec_register(self, spec_cls):
"""A decorated custom spec resolves to the declared manager."""
tests/v1/worker/test_attn_utils.py
+126 -1
View File
@@ -13,15 +13,18 @@ import pytest
import torch
from tests.v1.attention.utils import dense_kv_cache_views
from vllm.v1.attention.backend import AttentionCGSupport
from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf
from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
from vllm.v1.hisparse.binding import allocate_hisparse_kv_caches
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
HiSparseResidentSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheLayout,
KVCacheTensor,
MLAAttentionSpec,
SparseCacheRole,
compute_layout_strides,
)
from vllm.v1.worker.gpu import attn_utils
@@ -36,6 +39,60 @@ from vllm.v1.worker.utils import (
)
@pytest.mark.parametrize(
("enabled", "block_size", "main_sizes", "indexer_sizes", "expected"),
[
(True, 256, [64], [64], 64),
(True, 64, [32, 64], [16, 32], 32),
(True, 64, [MultipleOf(16)], [32], 32),
(True, 64, [64], [32], None),
(False, 256, [64], [64], 256),
],
)
def test_get_kv_cache_spec_resolves_hisparse_block_size(
monkeypatch, enabled, block_size, main_sizes, indexer_sizes, expected
):
"""Resolve shared MLA geometry before planning; leave other specs alone."""
specs = {
"main": MLAAttentionSpec(
block_size=block_size, num_kv_heads=1, head_size=576, dtype=torch.bfloat16
),
"indexer": MLAAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=128,
dtype=torch.bfloat16,
cache_role=SparseCacheRole.INDEXER,
),
"dense": FullAttentionSpec(
block_size=block_size, num_kv_heads=1, head_size=128, dtype=torch.bfloat16
),
}
layers = {}
for name, sizes in zip(specs, [main_sizes, indexer_sizes, [block_size]]):
backend = SimpleNamespace(
customize_spec=AttentionBackend.customize_spec,
get_supported_kernel_block_sizes=lambda sizes=sizes: sizes,
)
layers[name] = SimpleNamespace(
get_kv_cache_spec=lambda _, spec=specs[name]: spec,
get_attn_backend=lambda backend=backend: backend,
)
monkeypatch.setattr(attn_utils, "get_layers_from_vllm_config", lambda *_: layers)
config = SimpleNamespace(
attention_config=SimpleNamespace(hisparse_config=object() if enabled else None)
)
if expected is None:
with pytest.raises(ValueError, match="supported by every sparse"):
attn_utils.get_kv_cache_spec(config)
return
resolved = attn_utils.get_kv_cache_spec(config)
assert resolved["main"].block_size == resolved["indexer"].block_size == expected
assert resolved["dense"] is specs["dense"]
assert all(spec.block_size == block_size for spec in specs.values())
class _FakeMetadataBuilder:
def __init__(self, support: AttentionCGSupport):
self.support = support
@@ -390,3 +447,71 @@ def test_copy_kv_cache_blocks_with_virtual_block_splitting(
torch.testing.assert_close(
cache[dst_start + physical_idx], expected[layer_idx][physical_idx]
)
def test_allocate_hisparse_kv_caches_host_pool_and_view_less_specs():
"""Host tensors get their own backing; view-less specs keep the raw one."""
spec = FullAttentionSpec(
block_size=2, num_kv_heads=1, head_size=4, dtype=torch.float32
)
page = spec.page_size_bytes
resident_spec = HiSparseResidentSpec(block_size=2, page_size=page)
device_size = 4 * page
config = KVCacheConfig(
num_blocks=4,
hisparse_host_num_blocks=3,
kv_cache_tensors=[
KVCacheTensor(
size=3 * page,
layers=["source"],
layer_stride=3 * page,
block_stride=page,
host_resident=True,
),
KVCacheTensor(
size=device_size,
layers=["indexer"],
layer_stride=device_size,
block_stride=page,
),
KVCacheTensor(
size=device_size,
layers=["resident"],
layer_stride=device_size,
block_stride=page,
),
],
kv_cache_groups=[
KVCacheGroupSpec(["source"], spec, host_resident=True),
KVCacheGroupSpec(["indexer"], spec),
KVCacheGroupSpec(["resident"], resident_spec),
],
)
host_buffers: list[torch.Tensor] = []
def host_allocator(size: int) -> torch.Tensor:
host_buffers.append(torch.zeros(size, dtype=torch.int8))
return host_buffers[-1]
caches = allocate_hisparse_kv_caches(
config,
torch.device("cpu"),
KVCacheLayout.LBHNC,
[2, 2, 2],
SimpleNamespace(allocate=host_allocator),
)
assert len(config.kv_cache_tensors) == 3
assert [buf.numel() for buf in host_buffers] == [3 * page]
assert caches["source"].shape[0] == 3
assert (
caches["source"].untyped_storage().data_ptr()
== host_buffers[0].untyped_storage().data_ptr()
)
assert caches["indexer"].shape[0] == 4
backing = caches["resident"]
assert backing.dtype == torch.int8 and backing.numel() >= device_size
assert (
backing.untyped_storage().data_ptr()
== caches["indexer"].untyped_storage().data_ptr()
)
tests/v1/worker/test_dsv4_packed_zeroer_geometry.py
+1
View File
@@ -134,6 +134,7 @@ def test_overlaid_zeroer_dedups_segments_with_max_span():
vllm_config.cache_config = CacheConfig()
vllm_config.cache_config.num_gpu_blocks_override = None
vllm_config.cache_config.kv_cache_layout = "BLHNC"
vllm_config.attention_config.hisparse_config = None
config = get_kv_cache_config_from_groups(vllm_config, groups, 8 * 1024 * 1024)
views = allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout.BLHNC, None)
buf_ptr = views["g1.big"].data_ptr()
tests/v1/worker/test_gpu_autoregressive_speculator.py
+2
View File
@@ -351,6 +351,8 @@ def test_multi_step_decode_replays_captured_graph_as_expected(
speculator = object.__new__(_TestSpeculator)
speculator.num_speculative_steps = 4
speculator.current_draft_step = torch.tensor(0)
speculator.slot_mapping_observer = None
speculator.host_mirror_forward_observer = None
speculator.input_buffers = SimpleNamespace(
positions=torch.arange(2),
query_start_loc=torch.arange(3),
tests/v1/worker/test_gpu_kv_connector.py
+23 -3
View File
@@ -5,9 +5,11 @@ from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
import vllm.v1.worker.gpu.kv_connector as kv_connector_module
from vllm.config import KVTransferConfig
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorTransferResults
from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
@@ -18,9 +20,9 @@ def _make_connector(
backend = Mock()
backend.handle_preemptions.side_effect = lambda _: events.append("handle")
backend.bind_connector_metadata.side_effect = lambda _: events.append("bind")
backend.start_load_kv.side_effect = lambda _: events.append("start")
backend.start_load_kv.side_effect = lambda *_args, **_kwargs: events.append("start")
backend.wait_for_save.side_effect = lambda: events.append("wait")
backend.get_finished.side_effect = lambda _: (set(), set())
backend.get_transfer_results.return_value = KVConnectorTransferResults()
backend.get_block_ids_with_load_errors.return_value = set()
backend.get_kv_connector_stats.return_value = None
backend.get_kv_connector_kv_cache_events.return_value = None
@@ -61,7 +63,15 @@ def test_load_start_phase(
connector = _make_connector(monkeypatch, events)
output = _scheduler_output(has_sync_kv_loads)
connector.pre_forward(output) # type: ignore[arg-type]
request_indices = torch.tensor([3, 1])
request_ids = ["first", "second"]
attn_metadata = {"layer": object()}
connector.pre_forward( # type: ignore[arg-type]
output,
request_state_indices=request_indices,
request_ids=request_ids,
attn_metadata=attn_metadata,
)
assert events == (
["handle", "bind", "start"] if has_sync_kv_loads else ["handle", "bind"]
)
@@ -69,6 +79,16 @@ def test_load_start_phase(
connector.post_forward(set())
assert events == ["handle", "bind", "start", "wait", "clear"]
kwargs = connector.kv_connector.start_load_kv.call_args.kwargs
assert kwargs["request_state_indices"] is request_indices
assert kwargs["request_ids"] is request_ids
assert kwargs["attn_metadata"] is attn_metadata
# A subsequent step without a forward must not reuse the prior batch.
connector.no_forward(_scheduler_output(False)) # type: ignore[arg-type]
assert connector.kv_connector.start_load_kv.call_count == 2
assert connector.kv_connector.start_load_kv.call_args.kwargs == {}
def test_no_forward_starts_deferred_load_once(monkeypatch: pytest.MonkeyPatch):
events: list[str] = []
tests/v1/worker/test_gpu_model_runner.py
+4
View File
@@ -312,6 +312,10 @@ def test_select_common_block_size_uses_largest_shared_int():
assert selected_size == 64
def test_select_common_block_size_without_active_backends_uses_manager_size():
assert select_common_block_size(256, []) == 256
def test_select_common_block_size_accepts_rocm_sparse_block_size_16(monkeypatch):
monkeypatch.setattr(current_platform, "is_rocm", lambda: True)
tests/v1/worker/test_gpu_model_runner_v2.py
+1
View File
@@ -237,6 +237,7 @@ def _make_capture_runner(captured: bool) -> GPUModelRunner:
runner.attn_groups = None
runner.kv_cache_config = None
runner.use_aux_hidden_state_outputs = False
runner.kv_connector = model_runner_module.NO_OP_KV_CONNECTOR
return runner
tests/v1/worker/test_gpu_warmup_blocks.py
+1
View File
@@ -89,6 +89,7 @@ def _make_runner(
return SimpleNamespace(
num_speculative_steps=num_spec_steps,
adaptive_verification=None,
rejection_sampler=None,
decode_query_len=num_spec_steps + 1,
is_pooling_model=False,
is_encoder_decoder=False,
tests/v1/worker/test_kv_cache_allocation_scope.py
+31 -3
View File
@@ -5,8 +5,10 @@ from contextlib import AbstractContextManager
from types import SimpleNamespace
from typing import Any, cast
import pytest
import torch
import vllm.v1.hisparse.binding as hisparse_binding
import vllm.v1.worker.gpu.attn_utils as attn_utils
import vllm.v1.worker.gpu_model_runner as gpu_model_runner
from vllm.v1.worker.gpu_worker import Worker
@@ -26,7 +28,8 @@ class _AllocationScope(AbstractContextManager):
self.active = False
def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
@pytest.mark.parametrize("hisparse", [False, True])
def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch, hisparse) -> None:
scope = _AllocationScope()
kv_caches = {"layer": torch.empty(0)}
@@ -38,23 +41,48 @@ def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
assert not scope.active
monkeypatch.setattr(attn_utils, "allocate_kv_cache", allocate)
monkeypatch.setattr(attn_utils, "bind_kv_cache", bind)
monkeypatch.setattr(attn_utils, "bind_kv_cache_to_layers", bind)
monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda config: {})
hisparse_bindings = []
if hisparse:
host_pool = SimpleNamespace()
def bind_hisparse(**kwargs):
hisparse_bindings.append(kwargs)
assert scope.active
assert kwargs["kv_caches"] is kv_caches
assert kwargs["host_pool"] is host_pool
return [
SimpleNamespace(
view=SimpleNamespace(cache=torch.empty(1, 1, 1), block_size=1),
runtime=SimpleNamespace(),
)
]
monkeypatch.setattr(hisparse_binding, "HiSparseHostPool", lambda: host_pool)
monkeypatch.setattr(hisparse_binding, "allocate_hisparse_kv_caches", allocate)
monkeypatch.setattr(hisparse_binding, "bind_hisparse_kv_caches", bind_hisparse)
config = SimpleNamespace(
attention_config=SimpleNamespace(
hisparse_config=object() if hisparse else None
),
scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=1),
cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: None),
model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="test")),
)
result = attn_utils.init_kv_cache(
[],
{},
SimpleNamespace(kv_cache_groups=[]),
torch.device("cpu"),
[],
config,
kv_cache_allocation_context=scope,
block_tables=SimpleNamespace(),
)
assert len(hisparse_bindings) == int(hisparse)
assert result is kv_caches
assert not scope.active
tests/v1/worker/test_mixed_warmup_gate.py
+7 -1
View File
@@ -35,10 +35,15 @@ def test_mixed_warmup_skipped_for_single_seq(max_num_reqs):
def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_warmup):
"""Startup must warm fixed drafts before calibration and retain its manager."""
manager = SimpleNamespace(cost_tables=None)
runner = SimpleNamespace(adaptive_verification=manager)
rejection_sampler = SimpleNamespace(enable_adaptive_verification=True)
runner = SimpleNamespace(
adaptive_verification=manager,
rejection_sampler=rejection_sampler,
)
def run_steps(model_runner, execute, sample):
assert model_runner.adaptive_verification is None
assert not model_runner.rejection_sampler.enable_adaptive_verification
if fail_warmup:
raise RuntimeError("warmup failed")
@@ -50,3 +55,4 @@ def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_
warmup.warmup_kernels(runner, _fail, _fail)
assert runner.adaptive_verification is manager
assert manager.cost_tables is None
assert rejection_sampler.enable_adaptive_verification
tests/v1/worker/test_utils.py
+891 -3
View File
@@ -1,13 +1,897 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections import deque
from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import MagicMock, call
import numpy as np
import pytest
import torch
import vllm.v1.attention.backends.mla.index_group as index_group_module
import vllm.v1.hisparse.runtime as hisparse_runtime_module
from vllm.config import CUDAGraphMode
from vllm.config.mamba import MambaBackendEnum, MambaConfig
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import (
worker as hisparse_worker_module,
)
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
HiSparseConnectorWorker,
_flatten_row_mirrors,
_select_written_row_mirrors,
_SlotMappingStaging,
)
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
from vllm.v1.worker.utils import bind_kv_cache
from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror
from vllm.v1.worker.utils import (
bind_kv_cache,
bind_kv_cache_to_layers,
copy_kv_cache_blocks_inplace,
)
def _make_hisparse_worker() -> HiSparseConnectorWorker:
worker = object.__new__(HiSparseConnectorWorker)
worker.refines_row_mirrors = False
worker._slot_mapping_staging = None
worker._row_mirror_num_rows = 0
worker._per_layer_mirrored = set()
worker._submitted_mirror_layers = set()
worker._pending_dma_descriptors = deque()
worker._dma_free_descriptors = []
worker.host_write_events = (MagicMock(), MagicMock())
worker.host_write_event = worker.host_write_events[1]
worker._next_host_write_event = 0
worker.dma_stream = None
worker.shared_host_region = None
worker._completed_host_copy_dst_ids = []
worker._slot_mapping_staging = None
return worker
def test_hisparse_row_mirrors_follow_runner_request_order():
first = SparseKVRowMirror((1,), 10, 1)
second = SparseKVRowMirror((2,), 20, 1)
row_mirrors = {"first": (first,), "second": (second,)}
assert _flatten_row_mirrors(row_mirrors, ["second", "first"]) == (
second,
first,
)
def test_hisparse_gpu_slots_select_written_rows_from_async_envelope():
"""Rejected speculative rows must never become durable host KV."""
candidates = (
SparseKVRowMirror((100, 500), 1000, 4),
SparseKVRowMirror((20, 300), 2000, 4),
)
source_slots = np.array([102, 103, -1, 21, 22, 23], dtype=np.int64)
assert _select_written_row_mirrors(candidates, source_slots, 0) == (
SparseKVRowMirror((102, 502), 1002, 2),
SparseKVRowMirror((21, 301), 2001, 3),
)
def test_hisparse_written_rows_preserve_scheduler_page_boundaries():
candidates = (
SparseKVRowMirror((0,), 100, 4),
SparseKVRowMirror((4,), 104, 4),
)
assert _select_written_row_mirrors(
candidates, np.array([2, 3, 4, 5], dtype=np.int64), 0
) == (
SparseKVRowMirror((2,), 102, 2),
SparseKVRowMirror((4,), 104, 2),
)
def test_hisparse_written_rows_without_host_destinations_are_skipped():
assert (
_select_written_row_mirrors((), np.array([100, 101], dtype=np.int64), 0) == ()
)
def test_hisparse_appends_reference_slots_within_a_mirror_phase(monkeypatch):
"""Separate context and query writes must share one mirror phase."""
worker = _make_hisparse_worker()
worker.is_host_writer = True
worker.refines_row_mirrors = True
state = _SlotMappingStaging(
MagicMock(), MagicMock(), torch.empty(4, dtype=torch.int64)
)
worker._slot_mapping_staging = state
handle = SimpleNamespace(runtime=SimpleNamespace(resident_source_index=1))
worker.cache_layer_names = ["layer"]
worker.cache_handles = [handle]
worker._layer_mirror_callbacks = (MagicMock(),)
worker._row_mirror_num_rows = 1
worker._submitted_mirror_layers = set()
main_stream = MagicMock()
monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: main_stream)
monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext())
handle.slot_mapping = torch.tensor([7, 8], dtype=torch.int64)
worker._stage_row_mirror_mapping(2)
handle.slot_mapping = torch.tensor([9, 10], dtype=torch.int64)
worker._stage_row_mirror_mapping(2)
torch.testing.assert_close(state.slots, torch.tensor([7, 8, 9, 10]))
assert state.stream.wait_stream.call_args_list == [call(main_stream)] * 2
assert state.event.record.call_args_list == [call(state.stream)] * 2
assert state.num_tokens == 4
assert state.source_index == 1
@pytest.mark.parametrize(
("mode", "expected_calls"),
[(CUDAGraphMode.PIECEWISE, 1), (CUDAGraphMode.FULL, 0)],
)
def test_hisparse_submits_layer_mirror_at_replayed_attention_boundary(
monkeypatch, mode, expected_calls
):
"""Piecewise replay must pipeline DMA; FULL replay uses the batch fallback."""
handle = object.__new__(hisparse_runtime_module.HiSparseCacheHandle)
handle.dummy_batch = False
handle.decode_batch = False
handle.submit_layer_mirror = MagicMock()
monkeypatch.setattr(
hisparse_runtime_module,
"get_forward_context",
lambda: SimpleNamespace(cudagraph_runtime_mode=mode),
)
handle.finish_kv_update()
assert handle.submit_layer_mirror.call_count == expected_calls
def test_copy_cpu_kv_cache_logical_blocks_ignores_storage_padding():
backing = torch.full((10, 2, 3), -1, dtype=torch.float32)
cache = backing[1:9]
cache[2:4] = 7
cache[6:8] = 11
copy_kv_cache_blocks_inplace(
[cache],
num_blocks=4,
kv_cache_block_copies=[
KVCacheBlockCopy(1, 0),
KVCacheBlockCopy(3, 2),
],
)
torch.testing.assert_close(cache[0:2], torch.full_like(cache[0:2], 7))
torch.testing.assert_close(cache[4:6], torch.full_like(cache[4:6], 11))
assert (backing[0] == -1).all()
assert (backing[9] == -1).all()
@pytest.mark.parametrize("shared,rank", [(False, 0), (True, 0), (True, 1)])
@pytest.mark.parametrize("has_copies", [False, True])
def test_hisparse_host_copy_waits_for_writes_on_writer_only(
monkeypatch, shared, rank, has_copies
):
"""Host CoW must read completed writes; shared non-writers only barrier."""
worker = _make_hisparse_worker()
cache = torch.zeros((2, 1, 1))
worker.host_caches = [cache]
worker.host_num_blocks = 2
worker.shared_host_region = object() if shared else None
event = MagicMock()
event.synchronize.side_effect = lambda: cache[0].fill_(7)
barrier = MagicMock()
monkeypatch.setattr(
hisparse_worker_module, "get_tensor_model_parallel_rank", lambda: rank
)
monkeypatch.setattr(
hisparse_worker_module, "get_tp_group", lambda: SimpleNamespace(barrier=barrier)
)
copies = [KVCacheBlockCopy(0, 1)] if has_copies else []
worker._copy_host_blocks(copies, event)
writes = has_copies and (not shared or rank == 0)
assert cache[1].item() == (7 if writes else 0)
assert event.synchronize.call_count == int(writes)
assert barrier.call_count == int(shared and has_copies)
assert worker._completed_host_copy_dst_ids == ([1] if has_copies else [])
def test_hisparse_worker_updates_request_state_mapping_in_place(monkeypatch):
worker = _make_hisparse_worker()
worker.request_state_indices = torch.arange(4, dtype=torch.int32)
worker._pending_invalid_block_ids = [5]
invalidations = []
worker.invalidate_blocks = lambda blocks, states: invalidations.append(
(blocks.copy(), states.clone())
)
original_ptr = worker.request_state_indices.data_ptr()
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
worker.set_request_state_indices(torch.tensor([3, 1], dtype=torch.int32))
assert worker.request_state_indices.data_ptr() == original_ptr
assert worker.request_state_indices.tolist() == [3, 1, -1, -1]
assert len(invalidations) == 1
assert invalidations[0][0] == [5]
torch.testing.assert_close(
invalidations[0][1], torch.tensor([3, 1], dtype=torch.int32)
)
assert worker._pending_invalid_block_ids == []
def test_hisparse_pre_forward_transfer_builds_page_descriptors():
worker = _make_hisparse_worker()
worker.is_host_writer = True
worker.kernel_block_size = 2
source = torch.empty((3, 2, 4), dtype=torch.uint8)
destination = torch.empty((6, 4), dtype=torch.uint8)
worker.resident_caches = (source,)
worker.host_caches = (destination,)
worker.cache_handles = [
SimpleNamespace(runtime=SimpleNamespace(resident_source_index=0))
]
worker._dma_free_descriptors = []
worker._submit_dma_descriptors = MagicMock()
worker._enqueue_transfers([SparseKVPageTransfer(7, 2, (1,), after_forward=False)])
descriptors, count = worker._submit_dma_descriptors.call_args.args
assert count == 1
assert descriptors.src[:count].tolist() == [
source.data_ptr() + source.stride(0) * source.element_size()
]
assert descriptors.dst[:count].tolist() == [destination.data_ptr() + 4 * 4]
assert descriptors.sizes[:count].tolist() == [2 * 4]
assert worker._submit_dma_descriptors.call_args.kwargs["transfer_ids"] == (7,)
def test_hisparse_eager_mirror_records_transfer_without_page_copy():
worker = _make_hisparse_worker()
worker.cache_handles = [
SimpleNamespace(runtime=SimpleNamespace(eager_host_mirror=True))
]
worker._record_transfer_completion = MagicMock()
worker._enqueue_transfers = MagicMock()
transfers = [SparseKVPageTransfer(7, 2, (1,), after_forward=False)]
worker._submit_transfers(transfers)
worker._record_transfer_completion.assert_called_once_with(transfers)
worker._enqueue_transfers.assert_not_called()
def test_hisparse_lazy_mirror_copies_transfer_pages():
worker = _make_hisparse_worker()
worker.cache_handles = [
SimpleNamespace(runtime=SimpleNamespace(eager_host_mirror=False))
]
worker._record_transfer_completion = MagicMock()
worker._enqueue_transfers = MagicMock()
transfers = [SparseKVPageTransfer(7, 2, (1,), after_forward=False)]
worker._submit_transfers(transfers)
worker._enqueue_transfers.assert_called_once_with(transfers)
worker._record_transfer_completion.assert_not_called()
def test_hisparse_dma_row_mirror_builds_descriptors(monkeypatch):
worker = _make_hisparse_worker()
worker.is_host_writer = True
worker.kernel_block_size = 2
source = torch.empty((3, 2, 4), dtype=torch.uint8)
destination = torch.empty((6, 4), dtype=torch.uint8)
worker.resident_caches = (source,)
worker.host_caches = (destination,)
worker.cache_handles = [
SimpleNamespace(
runtime=SimpleNamespace(resident_source_index=0), decode_batch=True
)
]
worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0"))
worker.dma_stream = MagicMock()
worker.host_write_event = MagicMock()
worker._dma_free_descriptors = []
worker._pending_dma_descriptors = deque()
worker._pending_transfer_events = deque()
worker._enqueued_transfer_ids = []
worker._dma_submitted = False
current_stream = MagicMock()
event = MagicMock()
swap_blocks_batch = MagicMock()
monkeypatch.setattr(
hisparse_worker_module, "current_stream", lambda: current_stream
)
monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext())
monkeypatch.setattr(torch, "Event", lambda: event)
monkeypatch.setattr(
hisparse_worker_module.ops, "swap_blocks_batch", swap_blocks_batch
)
worker._set_row_mirrors((SparseKVRowMirror((3,), 4, 1),))
worker._enqueue_row_dma(range(1))
src_ptrs, dst_ptrs, sizes = swap_blocks_batch.call_args.args
assert src_ptrs.tolist() == [source.data_ptr() + 3 * 4]
assert dst_ptrs.tolist() == [destination.data_ptr() + 4 * 4]
assert sizes.tolist() == [4]
worker.dma_stream.wait_stream.assert_called_once_with(current_stream)
worker.host_write_event.record.assert_called_once_with(worker.dma_stream)
event.record.assert_called_once_with(worker.dma_stream)
assert worker._dma_submitted
def test_hisparse_row_dma_uses_resident_spans():
worker = _make_hisparse_worker()
source = torch.empty((2, 2, 4), dtype=torch.uint8)
destination = torch.empty((12, 4), dtype=torch.uint8)
worker.is_host_writer = True
worker.kernel_block_size = 2
worker.resident_caches = (source,)
worker.host_caches = (destination,)
worker.cache_handles = [
SimpleNamespace(
runtime=SimpleNamespace(resident_source_index=0), decode_batch=True
)
]
worker._set_row_mirrors(
(
SparseKVRowMirror((0,), 4, 2),
SparseKVRowMirror((2,), 9, 1),
)
)
worker._dma_free_descriptors = []
worker._submit_dma_descriptors = MagicMock()
worker._enqueue_row_dma(range(1))
descriptors, count = worker._submit_dma_descriptors.call_args.args
assert count == 2
assert descriptors.src[:count].tolist() == [
source.data_ptr(),
source.data_ptr() + 2 * 4,
]
assert descriptors.dst[:count].tolist() == [
destination.data_ptr() + 4 * 4,
destination.data_ptr() + 9 * 4,
]
assert descriptors.sizes[:count].tolist() == [2 * 4, 4]
def test_hisparse_finish_forward_mirrors_all_layers_once(monkeypatch):
dst_slots = torch.tensor([7, 8, 9], dtype=torch.int64)
req_ids = torch.tensor([0, 1, 2], dtype=torch.int32)
leader = SimpleNamespace(
eager_host_mirror=True,
is_group_leader=True,
invalidate_written_slots=MagicMock(),
)
follower = SimpleNamespace(
eager_host_mirror=True,
is_group_leader=False,
invalidate_written_slots=MagicMock(),
)
handles = [
SimpleNamespace(
runtime=runtime,
decode_batch=True,
host_mirror_required=True,
num_actual_tokens=3,
num_decode_tokens=2,
req_id_per_token=req_ids,
mirror_slot_mapping=dst_slots,
)
for runtime in (leader, follower)
]
worker = _make_hisparse_worker()
worker.is_host_writer = True
worker.cache_handles = handles
worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0"))
worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 3),))
worker._dma_submitted = False
worker._per_layer_mirrored = set()
worker._submitted_mirror_layers = set()
worker._post_forward_transfers = []
worker._enqueue_row_dma = MagicMock()
worker.host_write_event = MagicMock()
worker._forward_ready_event = MagicMock()
current_stream = MagicMock()
monkeypatch.setattr(
hisparse_worker_module, "current_stream", lambda: current_stream
)
worker.finish_forward()
worker._forward_ready_event.record.assert_called_once_with()
worker._enqueue_row_dma.assert_called_once_with(
(0, 1), ready_event=worker._forward_ready_event
)
leader.invalidate_written_slots.assert_called_once()
torch.testing.assert_close(
leader.invalidate_written_slots.call_args.args[0], dst_slots
)
follower.invalidate_written_slots.assert_not_called()
worker.host_write_event.record.assert_called_once_with(current_stream)
def test_hisparse_finish_forward_does_not_repeat_per_layer_mirrors():
slots = torch.tensor([7, 8], dtype=torch.int64)
runtime = SimpleNamespace(
eager_host_mirror=False,
is_group_leader=False,
invalidate_written_slots=MagicMock(),
)
handles = [
SimpleNamespace(
runtime=runtime,
decode_batch=False,
host_mirror_required=True,
num_actual_tokens=2,
num_decode_tokens=0,
req_id_per_token=torch.empty(0, dtype=torch.int32),
mirror_slot_mapping=slots,
)
for _ in range(2)
]
worker = _make_hisparse_worker()
worker.is_host_writer = True
worker.cache_handles = handles
worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 2),))
worker._per_layer_mirrored = {0, 1}
worker._submitted_mirror_layers = {0, 1}
worker._enqueue_row_dma = MagicMock()
worker._enqueue_host_mirror()
worker._enqueue_row_dma.assert_not_called()
def test_hisparse_finish_forward_submits_lazy_post_forward_transfer(monkeypatch):
runtime = SimpleNamespace(eager_host_mirror=False)
worker = _make_hisparse_worker()
worker.is_host_writer = True
worker.cache_handles = [SimpleNamespace(runtime=runtime, num_actual_tokens=0)]
transfer = SparseKVPageTransfer(7, 2, (1,), after_forward=True)
worker._post_forward_transfers = [transfer]
worker._forward_ready_event = MagicMock()
worker._enqueue_host_mirror = MagicMock()
worker._submit_transfers = MagicMock()
worker._dma_submitted = False
worker._submitted_mirror_layers = set()
worker._finish_mirror_phase = MagicMock()
worker._release_completed_dma_descriptors = MagicMock()
stream = MagicMock()
monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
worker.finish_forward()
worker._finish_mirror_phase.assert_called_once_with(worker._forward_ready_event)
worker._submit_transfers.assert_called_once_with([transfer])
worker.host_write_event.record.assert_called_once_with(stream)
def test_hisparse_prefill_mirrors_source_groups_and_flushes_partial_group():
slots = torch.tensor([7, 8], dtype=torch.int64)
source_indices = [0, 0, 1, 1, 1, 1, 2]
handles = [
SimpleNamespace(
runtime=SimpleNamespace(
eager_host_mirror=False,
is_group_leader=False,
resident_source_index=source_index,
),
decode_batch=False,
host_mirror_required=layer_index < 5,
num_actual_tokens=2 if layer_index < 5 else 0,
num_decode_tokens=0,
req_id_per_token=torch.empty(0, dtype=torch.int32),
mirror_slot_mapping=slots,
)
for layer_index, source_index in enumerate(source_indices)
]
worker = _make_hisparse_worker()
worker.is_host_writer = True
worker.cache_handles = handles
worker._set_row_mirrors((SparseKVRowMirror((0, 0, 0), 7, 2),))
worker._per_layer_mirrored = set()
worker._submitted_mirror_layers = set()
worker._layer_ready_events = tuple(MagicMock() for _ in handles)
worker._enqueue_row_dma = MagicMock()
for layer_index in range(5):
worker._enqueue_layer_mirror(layer_index)
worker._enqueue_row_dma.assert_called_once_with(
(0, 1), ready_event=worker._layer_ready_events[1]
)
worker._enqueue_host_mirror(ready_event=worker._layer_ready_events[-1])
assert worker._enqueue_row_dma.call_args_list[1].args == ((2, 3, 4),)
assert worker._enqueue_row_dma.call_args_list[1].kwargs == {
"ready_event": worker._layer_ready_events[-1]
}
def test_hisparse_finish_forward_rejects_partial_per_layer_mirror():
slots = torch.tensor([7, 8], dtype=torch.int64)
runtime = SimpleNamespace(
eager_host_mirror=False,
is_group_leader=False,
invalidate_written_slots=MagicMock(),
)
handles = [
SimpleNamespace(
runtime=runtime,
decode_batch=False,
host_mirror_required=True,
num_actual_tokens=2,
num_decode_tokens=0,
req_id_per_token=torch.empty(0, dtype=torch.int32),
mirror_slot_mapping=slots,
)
for _ in range(2)
]
worker = _make_hisparse_worker()
worker.is_host_writer = True
worker.cache_handles = handles
worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 2),))
worker._per_layer_mirrored = {0}
worker._submitted_mirror_layers = set()
worker._enqueue_row_dma = MagicMock()
with pytest.raises(RuntimeError, match="did not mirror every active layer"):
worker._enqueue_host_mirror()
def test_hisparse_finish_forward_orders_next_forward_after_dma(monkeypatch):
current_stream = MagicMock()
worker = _make_hisparse_worker()
worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0"))
worker.is_host_writer = True
worker._dma_submitted = True
worker.host_write_event = MagicMock()
worker.cache_handles = [
SimpleNamespace(
runtime=SimpleNamespace(eager_host_mirror=True), num_actual_tokens=0
)
]
worker._post_forward_transfers = []
worker._pending_dma_descriptors = deque()
worker._dma_free_descriptors = []
worker._forward_ready_event = MagicMock()
worker._finish_mirror_phase = MagicMock()
monkeypatch.setattr(
hisparse_worker_module, "current_stream", lambda: current_stream
)
worker.finish_forward()
current_stream.wait_event.assert_called_once_with(worker.host_write_event)
assert not worker._dma_submitted
def test_hisparse_finish_forward_mirrors_standalone_mtp_cache(monkeypatch):
dst_slots = torch.tensor([7, 8], dtype=torch.int64)
runtime = SimpleNamespace(
eager_host_mirror=True,
is_group_leader=False,
invalidate_written_slots=MagicMock(),
)
inactive = SimpleNamespace(
runtime=runtime,
decode_batch=False,
host_mirror_required=False,
num_actual_tokens=0,
num_decode_tokens=0,
req_id_per_token=None,
mirror_slot_mapping=dst_slots,
)
mtp = SimpleNamespace(
runtime=runtime,
decode_batch=True,
host_mirror_required=True,
num_actual_tokens=2,
num_decode_tokens=2,
req_id_per_token=torch.tensor([0, 1], dtype=torch.int32),
mirror_slot_mapping=dst_slots,
)
worker = _make_hisparse_worker()
worker.is_host_writer = True
worker.cache_handles = [inactive, mtp]
worker._set_row_mirrors((SparseKVRowMirror((0,), 7, 2),))
worker._per_layer_mirrored = set()
worker._submitted_mirror_layers = set()
worker._enqueue_row_dma = MagicMock()
worker._enqueue_host_mirror()
worker._enqueue_row_dma.assert_called_once_with((1,), ready_event=None)
def test_hisparse_shared_host_reader_skips_mirror(monkeypatch):
"""A non-writer TP rank must not mirror rows into the shared host pool."""
dst_slots = torch.tensor([7, 8], dtype=torch.int64)
leader = SimpleNamespace(
eager_host_mirror=True,
is_group_leader=True,
invalidate_written_slots=MagicMock(),
)
handle = SimpleNamespace(
runtime=leader,
decode_batch=True,
host_mirror_required=True,
num_actual_tokens=2,
num_decode_tokens=1,
req_id_per_token=torch.tensor([0], dtype=torch.int32),
mirror_slot_mapping=dst_slots,
)
worker = _make_hisparse_worker()
worker.is_host_writer = False
worker.cache_handles = [handle]
worker._set_row_mirrors((SparseKVRowMirror((0,), 7, 2),))
worker._per_layer_mirrored = set()
worker._enqueue_row_dma = MagicMock()
worker._enqueue_host_mirror()
worker._enqueue_row_dma.assert_not_called()
leader.invalidate_written_slots.assert_called_once()
def test_hisparse_shared_host_reader_skips_transfer_completion():
"""A non-writer TP rank must not acknowledge host writes."""
worker = _make_hisparse_worker()
worker.is_host_writer = False
worker.kernel_block_size = 1
worker._enqueued_transfer_ids = []
worker._pending_transfer_events = []
worker._record_transfer_completion([SparseKVPageTransfer(1, 2, (3,), True)])
assert worker._enqueued_transfer_ids == []
assert worker._pending_transfer_events == []
def test_hisparse_writer_records_transfer_completion_after_dma(monkeypatch):
worker = _make_hisparse_worker()
worker.is_host_writer = True
worker.dma_stream = MagicMock()
worker._enqueued_transfer_ids = []
worker._pending_transfer_events = deque()
event = MagicMock()
monkeypatch.setattr(torch, "Event", lambda: event)
worker._record_transfer_completion(
[
SparseKVPageTransfer(3, 2, (1,), False),
SparseKVPageTransfer(7, 4, (5,), True),
]
)
event.record.assert_called_once_with(worker.dma_stream)
assert worker._enqueued_transfer_ids == [3, 7]
assert worker._pending_transfer_events == deque([(event, (3, 7))])
@pytest.mark.parametrize("is_host_writer", [False, True])
def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer):
worker = _make_hisparse_worker()
worker.is_host_writer = is_host_writer
worker.hot_backing = SimpleNamespace(device=torch.device("cuda:1"))
host_write_events = (MagicMock(), MagicMock())
worker.host_write_events = host_write_events
worker.host_write_event = host_write_events[1]
worker._next_host_write_event = 0
worker.host_caches = ()
worker.host_num_blocks = 1
worker._post_forward_transfers = []
worker._pending_invalid_block_ids = []
worker.cache_handles = [
SimpleNamespace(
runtime=SimpleNamespace(eager_host_mirror=True),
decode_batch=False,
host_mirror_required=False,
num_actual_tokens=0,
num_decode_tokens=0,
req_id_per_token=None,
)
]
worker._per_layer_mirrored = set()
worker._submitted_mirror_layers = set()
worker._layer_mirror_callbacks = ()
stream = MagicMock()
monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
worker.start_step(
SimpleNamespace(
host_block_copies=[],
command=None,
source_block_ids=[],
row_mirrors={},
all_context_pages_resident=True,
row_mirrors_from_resident=False,
),
None,
)
worker.start_step(
SimpleNamespace(
host_block_copies=[],
command=None,
source_block_ids=[],
row_mirrors={},
all_context_pages_resident=True,
row_mirrors_from_resident=False,
),
None,
)
assert stream.wait_event.call_args_list == [
call(host_write_events[1]),
call(host_write_events[0]),
]
assert worker.host_write_event is host_write_events[1]
@pytest.mark.parametrize("tp_rank, expected_copies", [(0, 1), (1, 0)])
def test_hisparse_shared_host_block_copy_has_one_writer(
monkeypatch, tp_rank, expected_copies
):
worker = _make_hisparse_worker()
worker.shared_host_region = object()
worker.host_caches = (torch.empty(4, 1),)
worker.host_num_blocks = 4
previous_event = MagicMock()
copies = (KVCacheBlockCopy(0, 1),)
copy_blocks = MagicMock()
tp_group = MagicMock()
monkeypatch.setattr(
hisparse_worker_module, "get_tensor_model_parallel_rank", lambda: tp_rank
)
monkeypatch.setattr(hisparse_worker_module, "get_tp_group", lambda: tp_group)
monkeypatch.setattr(
hisparse_worker_module, "copy_kv_cache_blocks_inplace", copy_blocks
)
worker._copy_host_blocks(copies, previous_event)
assert copy_blocks.call_count == expected_copies
tp_group.barrier.assert_called_once_with()
assert worker.take_completed_host_copies() == [1]
assert worker.take_completed_host_copies() == []
def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch):
handle = SimpleNamespace(
runtime=SimpleNamespace(eager_host_mirror=True),
decode_batch=True,
num_actual_tokens=2,
num_decode_tokens=2,
req_id_per_token=torch.tensor([0, 1]),
mirror_slot_mapping=torch.tensor([4, 5]),
)
worker = _make_hisparse_worker()
worker.is_host_writer = True
worker.hot_backing = SimpleNamespace(device=torch.device("cpu"))
worker.host_write_event = MagicMock()
worker._forward_ready_event = MagicMock()
worker.host_caches = ()
worker.host_num_blocks = 1
worker.cache_handles = [handle]
worker._per_layer_mirrored = set()
worker._submitted_mirror_layers = set()
worker._layer_mirror_callbacks = (MagicMock(),)
worker._post_forward_transfers = []
worker._pending_invalid_block_ids = []
worker._enqueue_host_mirror = MagicMock(wraps=worker._enqueue_host_mirror)
stream = MagicMock()
monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
worker.start_step(
SimpleNamespace(
host_block_copies=[],
command=None,
source_block_ids=[],
row_mirrors={},
all_context_pages_resident=True,
row_mirrors_from_resident=False,
),
None,
)
worker.finish_forward()
worker._enqueue_host_mirror.assert_called_once_with(worker._forward_ready_event)
assert handle.num_actual_tokens == 0
torch.testing.assert_close(handle.mirror_slot_mapping, torch.tensor([4, 5]))
def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatch):
"""Followers must not allocate duplicate runtime state before profiling."""
config = SimpleNamespace(
scheduler_config=SimpleNamespace(
max_num_seqs=2,
max_num_batched_tokens=2,
async_scheduling=False,
),
speculative_config=None,
kv_transfer_config=None,
)
resolved = hisparse_runtime_module.ResolvedHiSparseConfig(
top_k=4,
device_buffer_size=8,
)
monkeypatch.setattr(hisparse_runtime_module, "_has_hisparse_ops", lambda: True)
monkeypatch.setattr(
hisparse_runtime_module.ResolvedHiSparseConfig,
"from_vllm_config",
classmethod(lambda cls, vllm_config, model_top_k: resolved),
)
shared_states: list[object] = []
streams: list[object] = []
def create_shared_state(_device, _max_rows, _top_k):
shared_states.append(object())
return shared_states[-1]
def create_stream(_device):
streams.append(object())
return streams[-1]
monkeypatch.setattr(
hisparse_runtime_module, "_create_shared_topk_state", create_shared_state
)
monkeypatch.setattr(hisparse_runtime_module, "_create_copy_stream", create_stream)
monkeypatch.setattr(index_group_module, "_create_side_stream", lambda _: object())
monkeypatch.setattr(index_group_module, "_create_event", lambda: object())
index_group_builder = index_group_module.SparseMLAIndexGroupBuilder(
torch.empty((2, 4), dtype=torch.int32)
)
def make_cache_handle(is_leader: bool):
index_group, _ = index_group_builder.register_layer(is_leader)
cache_handle = hisparse_runtime_module.create_hisparse_cache_handle(
config,
model_top_k=4,
is_index_group_leader=is_leader,
row_width=8,
kv_dtype=torch.float32,
index_group=index_group,
device="cpu",
)
assert cache_handle is not None
return cache_handle
first_leader = make_cache_handle(True)
first_index_group = index_group_builder.current_group
first_follower = make_cache_handle(False)
second_leader = make_cache_handle(True)
second_index_group = index_group_builder.current_group
second_follower = make_cache_handle(False)
assert first_index_group is not None
assert second_index_group is not None
assert first_follower.runtime.index_group is first_leader.runtime.index_group
assert second_follower.runtime.index_group is second_leader.runtime.index_group
assert first_leader.runtime.index_group is not second_leader.runtime.index_group
assert first_leader.runtime.index_group.copy_stream is first_index_group.side_stream
assert (
first_leader.runtime.index_group.logical_topk_ready
is first_index_group.logical_topk_ready
)
assert (
second_leader.runtime.index_group.copy_stream is second_index_group.side_stream
)
assert (
second_leader.runtime.index_group.logical_topk_ready
is second_index_group.logical_topk_ready
)
assert len(shared_states) == 2
assert streams == []
class _TestReplaySSMMixer(MambaMixer2):
@@ -30,7 +914,8 @@ def _packed_replayssm_cache(num_blocks: int) -> torch.Tensor:
return torch.full((num_blocks, 1, 1, 80), 0, dtype=torch.int8)
def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group():
@pytest.mark.parametrize("layers_only", [False, True])
def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(layers_only):
mixers = [_TestReplaySSMMixer() for _ in range(3)]
layer_names = [f"layers.{i}.mixer" for i in range(3)]
ctx = dict(zip(layer_names, mixers))
@@ -45,7 +930,10 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group():
SimpleNamespace(layer_names=[layer_names[1]]),
]
bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups)
if layers_only:
bind_kv_cache_to_layers(kv_cache, ctx, kv_cache_groups=kv_cache_groups)
else:
bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups)
assert (
mixers[0]._replayssm_ring_start.data_ptr()
vllm/_custom_ops.py
+15 -1
View File
@@ -3008,6 +3008,9 @@ def cp_gather_and_upconvert_fp8_kv_cache(
workspace_starts: torch.Tensor,
batch_size: int,
seq_starts: torch.Tensor | None = None,
host_cache: torch.Tensor | None = None,
host_row_ids: torch.Tensor | None = None,
device_row_ids: torch.Tensor | None = None,
) -> None:
"""Gather and upconvert FP8 KV cache to BF16 workspace.
@@ -3018,9 +3021,20 @@ def cp_gather_and_upconvert_fp8_kv_cache(
workspace_starts: Workspace start offsets [num_reqs]
batch_size: Number of requests
seq_starts: Optional source sequence offsets [num_reqs]
host_cache: Optional pinned host rows used for non-resident entries
host_row_ids: Host source row for each remapped cache row
device_row_ids: Device source row, or -1, for each remapped cache row
"""
torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache(
src_cache, dst, block_table, workspace_starts, batch_size, seq_starts
src_cache,
dst,
block_table,
workspace_starts,
batch_size,
seq_starts,
host_cache,
host_row_ids,
device_row_ids,
)
vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py
+4 -1
View File
@@ -35,7 +35,10 @@ def fused_rope_unified_mla_kv_cache_update_impl(
layer_name: LayerNameType,
) -> torch.Tensor:
layer_name = _resolve_layer_name(layer_name)
attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(
layer_name
)
attn_layer.prepare_kv_cache_update(attn_metadata)
if layer_slot_mapping is not None:
ops.concat_and_cache_mla_rope_fused(
positions,
vllm/config/__init__.py
+2 -1
View File
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.config.attention import AttentionConfig
from vllm.config.attention import AttentionConfig, HiSparseConfig
from vllm.config.cache import CacheConfig
from vllm.config.compilation import (
CompilationConfig,
@@ -68,6 +68,7 @@ from vllm.config.weight_transfer import WeightTransferConfig
__all__ = [
# From vllm.config.attention
"AttentionConfig",
"HiSparseConfig",
# From vllm.config.cache
"CacheConfig",
# From vllm.config.compilation
vllm/config/attention.py
+24 -1
View File
@@ -4,7 +4,7 @@
from dataclasses import field
from typing import Any, Literal
from pydantic import field_validator
from pydantic import Field, field_validator
from vllm.config.utils import config
from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
@@ -14,6 +14,25 @@ IndexerKVDType = Literal["auto", "bf16", "fp8", "mxfp4", "nvfp4"]
MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"]
@config
class HiSparseConfig:
"""Configuration for HiSparse sparse-MLA KV offloading."""
device_buffer_size: int | None = Field(default=None, gt=0)
"""Total per-request GPU hot-buffer rows, including the newest-token slot.
Defaults to one top-k per decode query plus one top-k of LRU slack. The
physical allocation is rounded up to the GPU cache block size selected
from the active backends.
"""
eager_host_mirror: bool = True
"""Mirror decode-written KV rows to the host pool during the forward so
page spills complete without moving data. When disabled, decode rows stay
resident-only and evicted pages are copied to host at spill time. Prefill
rows are always mirrored during the forward."""
@config
class AttentionConfig:
"""Configuration for attention mechanisms in vLLM."""
@@ -67,6 +86,10 @@ class AttentionConfig:
indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel
support in the backend."""
hisparse_config: HiSparseConfig | None = None
"""HiSparse host-resident KV configuration. Setting this enables experimental
Model Runner V2-only HiSparse sparse-MLA decode hot-buffering."""
use_non_causal: bool = False
"""Whether to use non-causal (bidirectional) attention."""
vllm/config/kv_transfer.py
+34
View File
@@ -13,6 +13,40 @@ KVConsumer = Literal["kv_consumer", "kv_both"]
KVRole = Literal[KVProducer, KVConsumer]
def hisparse_host_pool_gib(
kv_transfer_config: "KVTransferConfig | None",
) -> float | None:
if kv_transfer_config is None:
return None
if kv_transfer_config.kv_connector == "MultiConnector":
connectors = kv_transfer_config.kv_connector_extra_config.get("connectors", [])
else:
connectors = [
{
"kv_connector": kv_transfer_config.kv_connector,
"kv_connector_extra_config": (
kv_transfer_config.kv_connector_extra_config
),
}
]
entries = [
connector.get("kv_connector_extra_config", {})
for connector in connectors
if connector.get("kv_connector") == "HiSparseConnector"
]
if len(entries) > 1:
raise ValueError("Only one HiSparseConnector may be configured")
if not entries:
return None
host_pool_gib = entries[0].get("host_pool_gib")
if host_pool_gib is None:
raise ValueError("HiSparseConnector requires host_pool_gib")
host_pool_gib = float(host_pool_gib)
if host_pool_gib <= 0:
raise ValueError("HiSparseConnector host_pool_gib must be positive")
return host_pool_gib
def kv_buffer_device_default_factory() -> str:
from vllm.platforms import current_platform
vllm/config/speculative.py
+7
View File
@@ -1919,6 +1919,13 @@ class SpeculativeConfig:
def uses_draft_model(self) -> bool:
return self.method == "draft_model"
def uses_draft_kv_cache(self) -> bool:
return (
self.use_eagle()
or self.uses_draft_model()
or self.uses_extract_hidden_states()
)
def uses_extract_hidden_states(self) -> bool:
return self.method == "extract_hidden_states"
vllm/config/vllm.py
+51
View File
@@ -666,6 +666,14 @@ class VllmConfig:
@property
def use_v2_model_runner(self) -> bool:
if self.attention_config.hisparse_config is not None:
if envs.VLLM_USE_V2_MODEL_RUNNER is False:
raise ValueError(
"HiSparse requires Model Runner V2; remove "
"VLLM_USE_V2_MODEL_RUNNER=0."
)
return True
if getattr(self, "watermark_config", None) is not None:
if envs.VLLM_USE_V2_MODEL_RUNNER is False:
logger.info_once(
@@ -1559,6 +1567,39 @@ class VllmConfig:
self._maybe_disable_dynamic_sd_for_data_parallel()
self._maybe_override_dynamic_sd_cudagraph_mode()
if self.attention_config.hisparse_config is not None:
if not current_platform.is_cuda():
raise ValueError("HiSparse currently requires NVIDIA CUDA.")
if self.parallel_config.pipeline_parallel_size > 1:
raise ValueError("HiSparse does not support pipeline parallelism.")
if self.parallel_config.decode_context_parallel_size > 1:
raise ValueError(
"HiSparse does not support decode context parallelism."
)
if self.model_config is not None and not hasattr(
self.model_config.hf_config, "index_topk"
):
raise ValueError(
"HiSparse is only supported for DSA models with index_topk."
)
if self.kv_transfer_config is not None and (
self.kv_transfer_config.kv_connector
not in (
None,
"NixlConnector",
"MooncakeStoreConnector",
"MultiConnector",
)
):
logger.warning(
"HiSparse host-resident KV is configured with connector "
"%s. NixlConnector (GPU-staged host imports) and "
"MooncakeStoreConnector (shared-store offload) are the "
"validated paths; other connectors are treated as "
"debug/fallback paths.",
self.kv_transfer_config.kv_connector,
)
self._normalize_piecewise_cudagraph_mode(
breakable_cudagraph_enabled=breakable_cudagraph_enabled
)
@@ -1937,6 +1978,16 @@ class VllmConfig:
# Default to enable HMA if not explicitly disabled by user or logic above.
self.scheduler_config.disable_hybrid_kv_cache_manager = False
if (
self.attention_config.hisparse_config is not None
and self.scheduler_config.disable_hybrid_kv_cache_manager
):
raise ValueError(
"HiSparse requires the hybrid KV cache manager; remove "
"--disable-hybrid-kv-cache-manager or use connectors that "
"support HMA."
)
if self.compilation_config.debug_dump_path:
self.compilation_config.debug_dump_path = (
self.compilation_config.debug_dump_path.absolute().expanduser()
vllm/distributed/kv_transfer/kv_connector/factory.py
+6
View File
@@ -197,6 +197,12 @@ KVConnectorFactory.register_connector(
"MultiConnector",
)
KVConnectorFactory.register_connector(
"HiSparseConnector",
"vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector",
"HiSparseConnector",
)
KVConnectorFactory.register_connector(
"MoRIIOConnector",
"vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector",
vllm/distributed/kv_transfer/kv_connector/utils.py
+6
View File
@@ -60,6 +60,7 @@ class KVOutputAggregator:
# [req_id -> n_remaining_workers]
self._recv_remaining_count = dict[str, int]()
self._send_remaining_count = dict[str, int]()
self._failed_recving_pending = set[str]()
self._expected_finished_count = expected_finished_count
@classmethod
@@ -156,6 +157,10 @@ class KVOutputAggregator:
combined_kv_cache_events.increment_workers(1)
invalid_block_ids |= kv_output.invalid_block_ids
self._failed_recving_pending |= kv_output.failed_recving
failed_recving = self._failed_recving_pending & finished_recving
self._failed_recving_pending -= failed_recving
# select output of the worker specified by output_rank
output = outputs[output_rank]
@@ -168,6 +173,7 @@ class KVOutputAggregator:
kv_cache_events=combined_kv_cache_events or None,
kv_connector_worker_meta=aggregated_kv_connector_worker_meta or None,
invalid_block_ids=invalid_block_ids,
failed_recving=failed_recving,
expected_finished_count=self._expected_finished_count,
)
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+49 -6
View File
@@ -34,8 +34,8 @@ The class provides the following primitives:
save_kv_layer() - starts saving KV for layer i (maybe async)
wait_for_save() - blocks until all saves are done
get_finished() - called with ids of finished requests, returns
ids of requests that have completed async sending/recving.
get_transfer_results() - returns async send/receive completions and
receive failures in one snapshot.
build_connector_worker_meta() - builds metadata to be sent
back to the scheduler-side connector
"""
@@ -43,6 +43,7 @@ The class provides the following primitives:
import enum
from abc import ABC, abstractmethod
from collections.abc import Callable, Iterable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Literal
import torch
@@ -63,7 +64,7 @@ if TYPE_CHECKING:
)
from vllm.forward_context import ForwardContext
from vllm.v1.core.block_pool import BlockPool
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.request import Request
@@ -82,6 +83,19 @@ CopyBlocksOp = Callable[
logger = init_logger(__name__)
@dataclass
class KVConnectorTransferResults:
"""Asynchronous transfer completions from one worker snapshot.
Failed receives also appear in ``finished_recving`` so the scheduler can
release the request from its transfer wait state.
"""
finished_sending: set[str] = field(default_factory=set)
finished_recving: set[str] = field(default_factory=set)
failed_recving: set[str] = field(default_factory=set)
class SupportsHMA(ABC):
"""
The class that indicates the corresponding connector supports hybrid memory
@@ -211,6 +225,7 @@ class KVConnectorBase_V1(ABC):
else:
raise ValueError("kv_transfer_config must be set for KVConnectorBase_V1")
self._kv_cache_config = kv_cache_config
self._kv_cache_manager: KVCacheManager | None = None
self._role = role
@property
@@ -278,6 +293,15 @@ class KVConnectorBase_V1(ABC):
"""
return
# TODO(NickLucche): group model-runner lifecycle hooks in the interface.
def finish_forward(self) -> None:
"""Notify the connector that the model no longer reads this step's KV."""
return
def reset_capture_state(self) -> None:
"""Reset worker state mutated while capturing CUDA graphs."""
return
def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):
"""
Handle preempted requests or evicted blocks BEFORE they are overwritten.
@@ -368,6 +392,16 @@ class KVConnectorBase_V1(ABC):
"""
return None, None
def get_transfer_results(
self, finished_req_ids: set[str]
) -> KVConnectorTransferResults:
"""Return completed sends, receives, and receive failures together."""
finished_sending, finished_recving = self.get_finished(finished_req_ids)
return KVConnectorTransferResults(
finished_sending=set(finished_sending or ()),
finished_recving=set(finished_recving or ()),
)
def get_block_ids_with_load_errors(self) -> set[int]:
"""
Get the set of block IDs that failed to load.
@@ -380,9 +414,9 @@ class KVConnectorBase_V1(ABC):
- Applies to both sync- and async-loading requests.
- Async loading: failed blocks may be reported in any forward pass
up to and including the pass where the request ID is returned by
`get_finished()`. Even if failures occur, the request must still
be reported via `get_finished()`, and the failed block IDs must
appear here no later than that same pass.
`get_transfer_results()`. Even if failures occur, the request
must still be reported as finished receiving, and the failed
block IDs must appear here no later than that same pass.
- Sync loading: failed blocks should be reported in the forward
pass in which they are detected.
"""
@@ -436,6 +470,11 @@ class KVConnectorBase_V1(ABC):
# Scheduler-side methods
# ==============================
def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None:
"""Bind the scheduler's cache manager after it has been constructed."""
self._kv_cache_manager = kv_cache_manager
self.bind_gpu_block_pool(kv_cache_manager.block_pool)
def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
"""
Bind the GPU block pool to the connector for per-GPU block status tracking.
@@ -583,6 +622,10 @@ class KVConnectorBase_V1(ABC):
"""
return ()
def has_pending_block_frees(self) -> bool:
"""Whether pending transfers can release blocks instead of preemption."""
return False
def has_pending_push_work(self) -> bool:
"""Return True if the connector has push-mode work that requires
the engine main loop to keep stepping (e.g. a P-side request whose
vllm/distributed/kv_transfer/kv_connector/v1/hisparse/__init__.py
+2
View File
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py
+329
View File
@@ -0,0 +1,329 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Connector boundary for HiSparse scheduler and worker state."""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import torch
from vllm.config import VllmConfig
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
KVConnectorBase_V1,
KVConnectorMetadata,
KVConnectorRole,
KVConnectorWorkerMetadata,
SupportsHMA,
)
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
HiSparseConnectorWorker,
)
from vllm.v1.attention.backend import AttentionMetadata
from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
from vllm.v1.hisparse.types import SparseKVOffloadCommand, SparseKVRowMirror
from vllm.v1.outputs import KVConnectorOutput
if TYPE_CHECKING:
from vllm.forward_context import ForwardContext
from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
from vllm.v1.hisparse.coordinator import HiSparseCoordinator
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.request import Request
@dataclass
class HiSparseConnectorMetadata(KVConnectorMetadata):
command: SparseKVOffloadCommand | None
host_block_copies: tuple[KVCacheBlockCopy, ...]
source_block_ids: tuple[int, ...]
row_mirrors: dict[str, tuple[SparseKVRowMirror, ...]]
all_context_pages_resident: bool
@dataclass
class HiSparseConnectorWorkerMetadata(KVConnectorWorkerMetadata):
enqueued_transfer_counts: dict[int, int]
completed_transfer_counts: dict[int, int]
completed_host_copy_dst_ids: tuple[int, ...] = ()
def aggregate(self, other: KVConnectorWorkerMetadata) -> KVConnectorWorkerMetadata:
assert isinstance(other, HiSparseConnectorWorkerMetadata)
def add_counts(first: dict[int, int], second: dict[int, int]) -> dict[int, int]:
combined = first.copy()
for transfer_id, count in second.items():
combined[transfer_id] = combined.get(transfer_id, 0) + count
return combined
return HiSparseConnectorWorkerMetadata(
enqueued_transfer_counts=add_counts(
self.enqueued_transfer_counts, other.enqueued_transfer_counts
),
completed_transfer_counts=add_counts(
self.completed_transfer_counts, other.completed_transfer_counts
),
# Every worker runs the same copies, so a union is the ack.
completed_host_copy_dst_ids=tuple(
{*self.completed_host_copy_dst_ids, *other.completed_host_copy_dst_ids}
),
)
class HiSparseConnectorScheduler:
def __init__(
self,
*,
async_speculative: bool,
draft_kv_lookahead: int = 0,
) -> None:
self.coordinator: HiSparseCoordinator | None = None
self.async_speculative = async_speculative
self.draft_kv_lookahead = draft_kv_lookahead
self.requests: dict[str, Request] = {}
def bind_coordinator(self, coordinator: HiSparseCoordinator) -> None:
assert self.coordinator is None
self.coordinator = coordinator
def build_connector_meta(
self, scheduler_output: SchedulerOutput
) -> HiSparseConnectorMetadata:
assert self.coordinator is not None
# HiSparse rebinds worker state every step, so the load must start
# before the forward rather than being deferred to post-forward.
scheduler_output.has_sync_kv_loads = True
scheduler_output.block_table_updates = (
self.coordinator.take_block_table_updates() or None
)
command = self.coordinator.build_offload_command()
host_block_copies = self.coordinator.take_host_block_copies()
source_group_id = self.coordinator.host_group_id
assert source_group_id is not None
source_block_ids = [
block_id
for request in scheduler_output.scheduled_new_reqs
for block_id in request.block_ids[source_group_id]
]
for new_block_ids in scheduler_output.scheduled_cached_reqs.new_block_ids:
if new_block_ids is not None:
source_block_ids.extend(new_block_ids[source_group_id])
num_computed_tokens = {
request.req_id: request.num_computed_tokens
for request in scheduler_output.scheduled_new_reqs
}
num_computed_tokens.update(
zip(
scheduler_output.scheduled_cached_reqs.req_ids,
scheduler_output.scheduled_cached_reqs.num_computed_tokens,
)
)
scheduled_requests = tuple(
(
request_id,
num_computed_tokens[request_id],
scheduled_count,
)
for request_id, scheduled_count in (
scheduler_output.num_scheduled_tokens.items()
)
)
row_mirrors = {}
for request_id, scheduled_start, scheduled_count in scheduled_requests:
mirror_start = scheduled_start
if self.async_speculative:
mirror_start = max(
0,
scheduled_start - self.requests[request_id].num_output_placeholders,
)
row_mirrors[request_id] = self.coordinator.build_row_mirrors(
(
(
request_id,
mirror_start,
scheduled_count
+ scheduled_start
- mirror_start
+ self.draft_kv_lookahead,
),
),
)
return HiSparseConnectorMetadata(
command,
host_block_copies,
tuple(source_block_ids),
row_mirrors,
self.coordinator.all_context_pages_resident(scheduled_requests),
)
def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
assert self.coordinator is not None
for request_id in connector_output.finished_recving or ():
self.coordinator.finish_host_import(
request_id, failed=request_id in connector_output.failed_recving
)
metadata = connector_output.kv_connector_worker_meta
if metadata is None:
return
assert isinstance(metadata, HiSparseConnectorWorkerMetadata)
self.coordinator.release_completed_host_copies(
metadata.completed_host_copy_dst_ids
)
self.coordinator.update_spills(
metadata.enqueued_transfer_counts,
metadata.completed_transfer_counts,
)
class HiSparseConnector(KVConnectorBase_V1, SupportsHMA):
"""Join the scheduler coordinator to the worker's transfer engine."""
@classmethod
def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str:
return "BLHNC"
def __init__(
self,
vllm_config: VllmConfig,
role: KVConnectorRole,
kv_cache_config: KVCacheConfig,
) -> None:
super().__init__(vllm_config, role, kv_cache_config)
if kv_cache_config.hisparse_host_num_blocks is None:
raise ValueError("HiSparseConnector requires a HiSparse host pool")
self.connector_scheduler: HiSparseConnectorScheduler | None = None
self.connector_worker: HiSparseConnectorWorker | None = None
if role == KVConnectorRole.SCHEDULER:
speculative_config = vllm_config.speculative_config
self.connector_scheduler = HiSparseConnectorScheduler(
async_speculative=bool(
vllm_config.scheduler_config.async_scheduling
and speculative_config is not None
),
draft_kv_lookahead=vllm_config.num_lookahead_tokens,
)
elif role == KVConnectorRole.WORKER:
self.connector_worker = HiSparseConnectorWorker(
vllm_config, kv_cache_config
)
else:
raise ValueError(f"Unsupported KV connector role: {role}")
def bind_kv_cache_manager(self, kv_cache_manager: KVCacheManager) -> None:
assert self.connector_scheduler is not None
self.connector_scheduler.bind_coordinator(
get_hisparse_coordinator(kv_cache_manager)
)
@property
def requires_kv_delivery(self) -> bool:
return False
def has_pending_push_work(self) -> bool:
assert self.connector_scheduler is not None
assert self.connector_scheduler.coordinator is not None
return self.connector_scheduler.coordinator.has_pending_work()
def has_pending_block_frees(self) -> bool:
assert self.connector_scheduler is not None
assert self.connector_scheduler.coordinator is not None
return self.connector_scheduler.coordinator.has_pending_reclamation()
def finish_forward(self) -> None:
assert self.connector_worker is not None
self.connector_worker.finish_forward()
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
assert self.connector_worker is not None
self.connector_worker.register_kv_caches(kv_caches)
def reset_capture_state(self) -> None:
assert self.connector_worker is not None
self.connector_worker.reset_hot_state()
def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
assert self.connector_worker is not None
metadata = self._get_connector_metadata()
assert isinstance(metadata, HiSparseConnectorMetadata)
request_state_indices = kwargs.get("request_state_indices")
assert request_state_indices is None or isinstance(
request_state_indices, torch.Tensor
)
request_ids = kwargs.get("request_ids")
assert request_ids is None or isinstance(request_ids, list)
attn_metadata = kwargs.get("attn_metadata")
assert attn_metadata is None or isinstance(attn_metadata, dict)
self.connector_worker.start_step(
metadata,
request_state_indices,
request_ids,
num_tokens=int(kwargs.get("num_tokens") or 0),
)
self.connector_worker.prepare_forward(attn_metadata)
def wait_for_layer_load(self, layer_name: str) -> None:
return
def save_kv_layer(
self,
layer_name: str,
kv_layer: torch.Tensor,
attn_metadata: AttentionMetadata,
**kwargs: Any,
) -> None:
return
def wait_for_save(self) -> None:
return
def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:
assert self.connector_worker is not None
enqueued, completed = self.connector_worker.take_transfer_updates()
host_copies = self.connector_worker.take_completed_host_copies()
if not enqueued and not completed and not host_copies:
return None
return HiSparseConnectorWorkerMetadata(
enqueued_transfer_counts={transfer_id: 1 for transfer_id in enqueued},
completed_transfer_counts={transfer_id: 1 for transfer_id in completed},
completed_host_copy_dst_ids=tuple(host_copies),
)
def shutdown(self) -> None:
if self.connector_worker is not None:
self.connector_worker.shutdown()
def get_num_new_matched_tokens(
self, request: Request, num_computed_tokens: int
) -> tuple[int | None, bool]:
return 0, False
def update_state_after_alloc(
self,
request: Request,
blocks: KVCacheBlocks,
num_external_tokens: int,
) -> None:
assert self.connector_scheduler is not None
self.connector_scheduler.requests[request.request_id] = request
def build_connector_meta(
self, scheduler_output: SchedulerOutput
) -> KVConnectorMetadata:
assert self.connector_scheduler is not None
return self.connector_scheduler.build_connector_meta(scheduler_output)
def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
assert self.connector_scheduler is not None
self.connector_scheduler.update_connector_output(connector_output)
def request_finished_all_groups(
self,
request: Request,
block_ids: tuple[list[int], ...],
) -> tuple[bool, dict[str, Any] | None]:
assert self.connector_scheduler is not None
self.connector_scheduler.requests.pop(request.request_id, None)
return False, None
vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py
+832
View File
@@ -0,0 +1,832 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""HiSparse worker-side host/hot state and data movement."""
from __future__ import annotations
from collections import deque
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from functools import partial
from typing import TYPE_CHECKING, Any, NamedTuple
import numpy as np
import torch
from vllm import _custom_ops as ops
from vllm.config import VllmConfig
from vllm.distributed.parallel_state import (
get_tensor_model_parallel_rank,
get_tp_group,
)
from vllm.utils.torch_utils import current_stream
from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX
from vllm.v1.hisparse.runtime import HiSparseCacheHandle, release_pinned_state
from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror
from vllm.v1.kv_cache_interface import (
HiSparseHotSpec,
KVCacheConfig,
)
from vllm.v1.worker.utils import copy_kv_cache_blocks_inplace
if TYPE_CHECKING:
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
HiSparseConnectorMetadata,
)
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
class _DMADescriptors(NamedTuple):
src: torch.Tensor
dst: torch.Tensor
sizes: torch.Tensor
src_np: np.ndarray
dst_np: np.ndarray
sizes_np: np.ndarray
@dataclass
class _SlotMappingStaging:
stream: torch.Stream
event: torch.Event
slots: torch.Tensor
candidates: tuple[SparseKVRowMirror, ...] = ()
num_tokens: int = 0
source_index: int = 0
def _allocate_dma_descriptors(size: int) -> _DMADescriptors:
src, dst, sizes = (torch.empty(size, dtype=torch.int64) for _ in range(3))
return _DMADescriptors(
src,
dst,
sizes,
src.numpy(),
dst.numpy(),
sizes.numpy(),
)
def _get_hisparse_cache(
forward_context: dict[str, Any], layer_name: str
) -> HiSparseCacheHandle:
attention_layer = forward_context[layer_name]
hisparse_cache = attention_layer.hisparse_cache
assert hisparse_cache is not None
return hisparse_cache
def _flatten_row_mirrors(
row_mirrors: Mapping[str, tuple[SparseKVRowMirror, ...]],
request_ids: Sequence[str] | None,
) -> tuple[SparseKVRowMirror, ...]:
ordered_ids = row_mirrors if request_ids is None else request_ids
return tuple(
mirror
for request_id in ordered_ids
for mirror in row_mirrors.get(request_id, ())
)
def _select_written_row_mirrors(
candidates: tuple[SparseKVRowMirror, ...],
source_slots: np.ndarray,
source_index: int,
) -> tuple[SparseKVRowMirror, ...]:
"""Select and coalesce GPU-written rows from a scheduler-owned envelope."""
source_slots = source_slots[source_slots >= 0]
if source_slots.size == 0 or not candidates:
return ()
starts = np.asarray([mirror.source_starts for mirror in candidates], dtype=np.int64)
counts = np.fromiter((mirror.num_rows for mirror in candidates), dtype=np.int64)
destinations = np.fromiter(
[mirror.destination_start for mirror in candidates], dtype=np.int64
)
if source_index >= starts.shape[1]:
raise RuntimeError("HiSparse row-mirror source index is out of range.")
source_starts = starts[:, source_index]
order = np.argsort(source_starts)
matches = np.searchsorted(source_starts[order], source_slots, side="right") - 1
candidate_ids = order[np.maximum(matches, 0)]
offsets = source_slots - source_starts[candidate_ids]
valid = (matches >= 0) & (offsets < counts[candidate_ids])
if not np.all(valid):
# A page the scheduler could not map has no host destination, so its
# rows are simply not mirrorable this step.
source_slots = source_slots[valid]
candidate_ids = candidate_ids[valid]
offsets = offsets[valid]
if source_slots.size == 0:
return ()
sources = starts[candidate_ids] + offsets[:, None]
destinations = destinations[candidate_ids] + offsets
boundaries = (
np.flatnonzero(
(candidate_ids[1:] != candidate_ids[:-1])
| (np.diff(destinations) != 1)
| np.any(np.diff(sources, axis=0) != 1, axis=1)
)
+ 1
)
run_starts = np.concatenate(([0], boundaries))
run_ends = np.concatenate((boundaries, [destinations.size]))
return tuple(
SparseKVRowMirror(
source_starts=tuple(int(value) for value in sources[start]),
destination_start=int(destinations[start]),
num_rows=int(end - start),
)
for start, end in zip(run_starts, run_ends, strict=True)
)
def _create_hisparse_host_events(
shared_host_region: SharedOffloadRegion | None,
is_host_writer: bool,
device: torch.device,
) -> tuple[torch.Event, torch.Event]:
if shared_host_region is None:
return torch.Event(), torch.Event()
events: tuple[torch.Event, torch.Event] | None = None
ipc_handles = None
if is_host_writer:
events = (
torch.cuda.Event(interprocess=True),
torch.cuda.Event(interprocess=True),
)
stream = current_stream()
for event in events:
event.record(stream)
ipc_handles = tuple(event.ipc_handle() for event in events)
ipc_handles = get_tp_group().broadcast_object(ipc_handles, src=0)
if events is None:
events = (
torch.cuda.Event.from_ipc_handle(device, ipc_handles[0]),
torch.cuda.Event.from_ipc_handle(device, ipc_handles[1]),
)
return events
class HiSparseConnectorWorker:
"""Own HiSparse host/hot state and execute its worker-side transfers."""
def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None:
self.vllm_config = vllm_config
self.kv_cache_config = kv_cache_config
self._initialized = False
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
forward_context = self.vllm_config.compilation_config.static_forward_context
cache_handles: list[HiSparseCacheHandle] = []
cache_layer_names: list[str] = []
for group in self.kv_cache_config.kv_cache_groups:
if not isinstance(group.kv_cache_spec, HiSparseHotSpec):
continue
for cache_name in group.layer_names:
assert cache_name.endswith(HISPARSE_HOT_SUFFIX)
layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)]
cache_handles.append(_get_hisparse_cache(forward_context, layer_name))
cache_layer_names.append(layer_name)
if not cache_handles:
raise RuntimeError("HiSparse connector found no hot-cache handles.")
hot_backings: dict[int, torch.Tensor] = {}
registered_host_pools: dict[int, torch.Tensor] = {}
for cache in cache_handles:
hot_backing = cache.runtime.hot_backing
hot_backings[hot_backing.untyped_storage().data_ptr()] = hot_backing
registered_pool = cache.runtime.registered_host_pool
registered_host_pools[registered_pool.data_ptr()] = registered_pool
if len(hot_backings) != 1:
raise RuntimeError("HiSparse hot tensors must share one GPU backing.")
hot_backing = next(iter(hot_backings.values()))
pinned_host_pools = list(registered_host_pools.values())
resident = cache_handles[0].view
assert resident is not None
host_num_blocks = self.kv_cache_config.hisparse_host_num_blocks
assert host_num_blocks is not None
try:
self.initialize(
cache_handles,
cache_layer_names,
hot_backing,
self.vllm_config.scheduler_config.max_num_seqs,
host_num_blocks,
hot_backing.device,
pinned_host_pools,
)
except Exception:
release_pinned_state(
[cache.runtime for cache in cache_handles],
pinned_host_pools,
)
raise
def initialize(
self,
cache_handles: list[HiSparseCacheHandle],
cache_layer_names: list[str],
hot_backing: torch.Tensor,
max_num_reqs: int,
host_num_blocks: int,
device: torch.device,
pinned_host_pools: list[torch.Tensor],
*,
shared_host_region: SharedOffloadRegion | None = None,
is_host_writer: bool = True,
) -> None:
if self._initialized:
raise RuntimeError("HiSparse connector worker is already initialized.")
resident = cache_handles[0].view
assert resident is not None
self.kernel_block_size = resident.block_size
self.host_num_blocks = host_num_blocks
self.pinned_host_pools = pinned_host_pools
self.shared_host_region = shared_host_region
self.is_host_writer = is_host_writer
self.dma_stream = (
torch.cuda.Stream(device=device) if self.is_host_writer else None
)
self._slot_mapping_staging = None
if self.is_host_writer:
max_mirror_rows = (
self.vllm_config.scheduler_config.max_num_batched_tokens
+ max_num_reqs * (self.vllm_config.num_lookahead_tokens + 1)
)
self._slot_mapping_staging = _SlotMappingStaging(
stream=torch.Stream(device=device),
event=torch.Event(),
slots=torch.empty(
max_mirror_rows,
dtype=torch.int64,
pin_memory=True,
),
)
self._dma_free_descriptors: list[_DMADescriptors] = []
self._pending_dma_descriptors: deque[tuple[torch.Event, _DMADescriptors]] = (
deque()
)
self._dma_submitted = False
self._per_layer_mirrored: set[int] = set()
self._submitted_mirror_layers: set[int] = set()
self._layer_ready_events = tuple(torch.Event() for _ in cache_handles)
self._forward_ready_event = torch.Event()
self._set_row_mirrors(())
self.cache_layer_names = cache_layer_names
self._group_leaders = tuple(
(layer_name, cache)
for layer_name, cache in zip(cache_layer_names, cache_handles, strict=True)
if cache.runtime.is_group_leader
)
self.host_write_events = _create_hisparse_host_events(
shared_host_region, is_host_writer, device
)
self.host_write_event = self.host_write_events[1]
self._next_host_write_event = 0
self.cache_handles = cache_handles
self.leader_runtimes = [cache.runtime for _, cache in self._group_leaders]
request_state_indices = {
indices.data_ptr(): indices
for cache in cache_handles
if (indices := cache.runtime.request_state_indices) is not None
}
if len(request_state_indices) != 1:
raise RuntimeError(
"HiSparse runtimes must share one request-state mapping."
)
self.request_state_indices = next(iter(request_state_indices.values()))
if self.request_state_indices.numel() != max_num_reqs:
raise RuntimeError(
"HiSparse request-state mapping does not match max_num_seqs."
)
self.hot_backing = hot_backing
self._pending_invalid_block_ids: list[int] = []
# Destination block ids of host copies this worker has run.
self._completed_host_copy_dst_ids: list[int] = []
self._post_forward_transfers: list[SparseKVPageTransfer] = []
self._enqueued_transfer_ids: list[int] = []
self._pending_transfer_events: deque[tuple[torch.Event, tuple[int, ...]]] = (
deque()
)
self._init_dma()
if self.is_host_writer:
for layer_index, handle in enumerate(cache_handles):
handle.submit_layer_mirror = partial(
self._enqueue_layer_mirror, layer_index
)
self._initialized = True
def set_request_state_indices(self, indices: torch.Tensor) -> None:
if indices.numel() > self.request_state_indices.numel():
raise ValueError(
"HiSparse request-state mapping exceeds max_num_seqs: "
f"{indices.numel()} > {self.request_state_indices.numel()}."
)
if torch.cuda.is_current_stream_capturing():
return
# Attention indexes persistent request state by input-batch row. Refresh
# that indirection after every batch compaction or reorder.
self.request_state_indices.fill_(-1)
self.request_state_indices[: indices.numel()].copy_(indices)
if self._pending_invalid_block_ids:
self.invalidate_blocks(self._pending_invalid_block_ids, indices)
self._pending_invalid_block_ids.clear()
def _init_dma(self) -> None:
host_caches = tuple(cache.runtime.host_cache for cache in self.cache_handles)
resident_caches = []
for cache in self.cache_handles:
assert cache.view is not None and cache.slot_mapping is not None
resident_caches.append(cache.view.cache)
self.host_caches = host_caches
self.resident_caches = tuple(resident_caches)
for resident_cache, host_cache in zip(resident_caches, host_caches):
row_bytes = resident_cache.shape[-1] * resident_cache.element_size()
if (
resident_cache.ndim != 3
or resident_cache.shape[1] != self.kernel_block_size
or resident_cache.stride(1) * resident_cache.element_size() != row_bytes
):
raise RuntimeError("HiSparse DMA requires contiguous resident rows.")
if (
host_cache.ndim != 2
or not host_cache.is_contiguous()
or host_cache.shape[1] * host_cache.element_size() != row_bytes
):
raise RuntimeError("HiSparse DMA requires contiguous host rows.")
def start_step(
self,
metadata: HiSparseConnectorMetadata,
request_state_indices: torch.Tensor | None,
request_ids: list[str] | None = None,
num_tokens: int = 0,
) -> None:
self._stage_row_mirror_mapping(num_tokens)
previous_host_write_event = self.host_write_event
self.host_write_event = self.host_write_events[self._next_host_write_event]
self._next_host_write_event ^= 1
current_stream().wait_event(previous_host_write_event)
self._release_completed_dma_descriptors()
mirrors = _flatten_row_mirrors(metadata.row_mirrors, request_ids)
if self._slot_mapping_staging is not None:
self._slot_mapping_staging.candidates = mirrors
self._set_row_mirrors(mirrors)
self._dma_submitted = False
self._clear_forward_mirror_state()
for handle in self.cache_handles:
handle.all_context_pages_resident = metadata.all_context_pages_resident
handle.mirror_from_resident = True
self._copy_host_blocks(metadata.host_block_copies, previous_host_write_event)
transfers = (
metadata.command.page_transfers if metadata.command is not None else []
)
self._post_forward_transfers = [
transfer for transfer in transfers if transfer.after_forward
]
self._submit_transfers(
[transfer for transfer in transfers if not transfer.after_forward]
)
self._pending_invalid_block_ids.extend(metadata.source_block_ids)
if request_state_indices is not None:
self.set_request_state_indices(request_state_indices)
def _clear_forward_mirror_state(self) -> None:
self._per_layer_mirrored.clear()
self._submitted_mirror_layers.clear()
for handle in self.cache_handles:
handle.decode_batch = False
handle.host_mirror_required = False
handle.num_actual_tokens = 0
handle.num_decode_tokens = 0
handle.req_id_per_token = None
def prepare_forward(self, attn_metadata: Mapping[str, Any] | None) -> None:
if attn_metadata is None:
return
for layer_name, handle in self._group_leaders:
metadata = attn_metadata.get(layer_name)
if metadata is not None:
handle.prepare_group_for_batch(metadata)
def _stage_row_mirror_mapping(self, num_tokens: int) -> None:
"""Snapshot the rows this forward will write, off the compute stream.
The resident slot mapping is a persistent view bound at registration
time, so the rows are already staged by the time the forward launches.
"""
state = self._slot_mapping_staging
if state is None or not num_tokens:
return
handle = self.cache_handles[0]
slots = handle.slot_mapping
assert slots is not None
source_index = handle.runtime.resident_source_index
start = state.num_tokens
end = start + num_tokens
if end > state.slots.shape[0]:
raise ValueError(
"HiSparse row mapping exceeds staging capacity: "
f"{end} > {state.slots.shape[0]}."
)
main_stream = current_stream()
state.stream.wait_stream(main_stream)
with torch.cuda.stream(state.stream):
state.slots[start:end].copy_(
slots[:num_tokens],
non_blocking=True,
)
state.event.record(state.stream)
state.num_tokens = end
state.source_index = source_index
def _resolve_row_mirrors(self, state: _SlotMappingStaging) -> None:
"""Compute the final row mirrors from the staged GPU slot mapping."""
self._set_row_mirrors(
_select_written_row_mirrors(
state.candidates,
state.slots[: state.num_tokens].numpy(),
state.source_index,
)
)
state.num_tokens = 0
def _copy_host_blocks(
self,
host_block_copies: Sequence[KVCacheBlockCopy],
previous_host_write_event: torch.Event,
) -> None:
if not host_block_copies:
return
self._completed_host_copy_dst_ids.extend(
copy.dst_block_id for copy in host_block_copies
)
if self.shared_host_region is None or get_tensor_model_parallel_rank() == 0:
if self.host_caches:
previous_host_write_event.synchronize()
copy_kv_cache_blocks_inplace(
self.host_caches,
self.host_num_blocks,
host_block_copies,
)
if self.shared_host_region is not None:
get_tp_group().barrier()
def invalidate_blocks(
self, block_ids: list[int], request_state_indices: torch.Tensor
) -> None:
"""Invalidate recycled host slots in this worker's leader runtimes."""
if not block_ids:
return
device = self.cache_handles[0].runtime.device
staging = torch.tensor(block_ids, dtype=torch.int32, pin_memory=True)
blocks = staging.to(device, dtype=torch.int32, non_blocking=True)
offsets = torch.arange(self.kernel_block_size, dtype=torch.int32, device=device)
slots = (blocks[:, None] * self.kernel_block_size + offsets[None, :]).flatten()
sorted_slots = torch.sort(slots).values
state_indices = request_state_indices.to(device=device, dtype=torch.long)
for runtime in self.leader_runtimes:
runtime.invalidate_sorted_slots(sorted_slots, state_indices)
def reset_hot_state(self) -> None:
for runtime in self.leader_runtimes:
runtime.reset_hot_state()
def _release_completed_dma_descriptors(self) -> None:
pending = self._pending_dma_descriptors
while pending and pending[0][0].query():
_, descriptors = pending.popleft()
self._dma_free_descriptors.append(descriptors)
def _acquire_dma_descriptors(self, size: int) -> _DMADescriptors:
for index, descriptors in enumerate(self._dma_free_descriptors):
if descriptors.src.numel() >= size:
return self._dma_free_descriptors.pop(index)
return _allocate_dma_descriptors(size)
def _submit_dma_descriptors(
self,
descriptors: _DMADescriptors,
descriptor_count: int,
transfer_ids: tuple[int, ...] = (),
ready_event: torch.Event | None = None,
) -> None:
stream = self.dma_stream
assert stream is not None
if ready_event is None:
stream.wait_stream(current_stream())
else:
stream.wait_event(ready_event)
completion_event = torch.Event()
with torch.cuda.stream(stream):
ops.swap_blocks_batch(
descriptors.src[:descriptor_count],
descriptors.dst[:descriptor_count],
descriptors.sizes[:descriptor_count],
)
self.host_write_event.record(stream)
completion_event.record(stream)
self._pending_dma_descriptors.append((completion_event, descriptors))
self._dma_submitted = True
if transfer_ids:
self._pending_transfer_events.append((completion_event, transfer_ids))
self._enqueued_transfer_ids.extend(transfer_ids)
def _set_row_mirrors(self, mirrors: tuple[SparseKVRowMirror, ...]) -> None:
self._row_mirrors = mirrors
self._row_mirror_destination_starts = np.fromiter(
(mirror.destination_start for mirror in mirrors),
dtype=np.int64,
count=len(mirrors),
)
self._row_mirror_counts = np.fromiter(
(mirror.num_rows for mirror in mirrors),
dtype=np.int64,
count=len(mirrors),
)
if mirrors:
self._row_mirror_source_starts = np.asarray(
[mirror.source_starts for mirror in mirrors], dtype=np.int64
)
if self._row_mirror_source_starts.ndim != 2:
raise RuntimeError("HiSparse DMA source mappings must be rectangular.")
else:
self._row_mirror_source_starts = np.empty((0, 0), dtype=np.int64)
self._row_mirror_num_rows = int(self._row_mirror_counts.sum())
def _enqueue_row_dma(
self, layer_indices: Sequence[int], ready_event: torch.Event | None = None
) -> None:
if (
not layer_indices
or not self._row_mirror_num_rows
or not self.is_host_writer
):
return
mirrors = self._row_mirrors
num_layers = len(layer_indices)
descriptor_count = len(mirrors) * num_layers
descriptors = self._acquire_dma_descriptors(descriptor_count)
destination_starts = self._row_mirror_destination_starts
row_counts = self._row_mirror_counts
for descriptor_offset, layer_index in enumerate(layer_indices):
cache = self.cache_handles[layer_index]
source_index = cache.runtime.resident_source_index
if source_index >= self._row_mirror_source_starts.shape[1]:
raise RuntimeError("HiSparse row DMA source index is out of range.")
source_rows = self._row_mirror_source_starts[:, source_index]
source = self.resident_caches[layer_index]
destination = self.host_caches[layer_index]
row_bytes = source.shape[-1] * source.element_size()
if (
source.stride(1) * source.element_size() != row_bytes
or destination.shape[1] * destination.element_size() != row_bytes
):
raise RuntimeError("HiSparse row DMA requires contiguous rows.")
source_blocks, source_row_offsets = np.divmod(
source_rows, self.kernel_block_size
)
source_out_of_range = (
np.any(source_blocks < 0)
or np.any(source_blocks >= source.shape[0])
or np.any(source_row_offsets + row_counts > self.kernel_block_size)
)
if source_out_of_range:
raise RuntimeError("HiSparse row DMA source is out of range.")
if np.any(destination_starts < 0) or np.any(
destination_starts + row_counts > destination.shape[0]
):
raise RuntimeError("HiSparse row DMA index is out of range.")
descriptor_slice = slice(descriptor_offset, descriptor_count, num_layers)
descriptors.src_np[descriptor_slice] = (
source.data_ptr()
+ source_blocks * source.stride(0) * source.element_size()
+ source_row_offsets * row_bytes
)
descriptors.dst_np[descriptor_slice] = (
destination.data_ptr() + destination_starts * row_bytes
)
descriptors.sizes_np[descriptor_slice] = row_counts * row_bytes
self._submit_dma_descriptors(
descriptors, descriptor_count, ready_event=ready_event
)
def _enqueue_layer_mirror(self, layer_index: int) -> None:
handle = self.cache_handles[layer_index]
if not handle.host_mirror_required:
return
if layer_index in self._per_layer_mirrored:
raise RuntimeError(f"HiSparse layer {layer_index} mirrored twice.")
state = self._slot_mapping_staging
if state is not None and state.num_tokens:
# The staging stream waited on the compute stream before the
# forward launched, so this only drains pre-forward work that is
# already queued ahead of the layer kernels: no GPU bubble.
state.event.synchronize()
self._resolve_row_mirrors(state)
self._per_layer_mirrored.add(layer_index)
next_layer = layer_index + 1
if (
next_layer < len(self.cache_handles)
and self.cache_handles[next_layer].runtime.resident_source_index
== handle.runtime.resident_source_index
):
return
ready_event = self._layer_ready_events[layer_index]
ready_event.record()
pending_layers = tuple(
sorted(self._per_layer_mirrored - self._submitted_mirror_layers)
)
self._enqueue_row_dma(
pending_layers,
ready_event=ready_event,
)
self._submitted_mirror_layers.update(pending_layers)
def _record_transfer_completion(
self, transfers: list[SparseKVPageTransfer]
) -> None:
if not transfers or not self.is_host_writer:
return
stream = self.dma_stream
assert stream is not None
completion_event = torch.Event()
completion_event.record(stream)
transfer_ids = tuple(transfer.transfer_id for transfer in transfers)
self._pending_transfer_events.append((completion_event, transfer_ids))
self._enqueued_transfer_ids.extend(transfer_ids)
def _submit_transfers(self, transfers: list[SparseKVPageTransfer]) -> None:
if self.cache_handles[0].runtime.eager_host_mirror:
self._record_transfer_completion(transfers)
else:
self._enqueue_transfers(transfers)
def _enqueue_transfers(self, transfers: list[SparseKVPageTransfer]) -> None:
if not transfers or not self.is_host_writer:
return
num_layers = len(self.cache_handles)
descriptor_count = len(transfers) * num_layers
descriptors = self._acquire_dma_descriptors(descriptor_count)
destination_rows = np.fromiter(
(
transfer.destination_block_id * self.kernel_block_size
for transfer in transfers
),
dtype=np.int64,
count=len(transfers),
)
source_blocks_by_transfer = np.asarray(
[transfer.source_block_ids for transfer in transfers], dtype=np.int64
)
if source_blocks_by_transfer.ndim != 2:
raise RuntimeError(
"HiSparse spill DMA source mappings must be rectangular."
)
for layer_index, cache in enumerate(self.cache_handles):
source_index = cache.runtime.resident_source_index
if source_index >= source_blocks_by_transfer.shape[1]:
raise RuntimeError("HiSparse spill DMA source index is out of range.")
source_blocks = source_blocks_by_transfer[:, source_index]
source = self.resident_caches[layer_index]
destination = self.host_caches[layer_index]
if np.any(source_blocks < 0) or np.any(source_blocks >= source.shape[0]):
raise RuntimeError("HiSparse spill DMA source is out of range.")
if np.any(destination_rows < 0) or np.any(
destination_rows + self.kernel_block_size > destination.shape[0]
):
raise RuntimeError("HiSparse spill DMA destination is out of range.")
row_bytes = source.shape[-1] * source.element_size()
descriptor_slice = slice(layer_index, descriptor_count, num_layers)
descriptors.src_np[descriptor_slice] = (
source.data_ptr()
+ source_blocks * source.stride(0) * source.element_size()
)
descriptors.dst_np[descriptor_slice] = (
destination.data_ptr() + destination_rows * row_bytes
)
descriptors.sizes_np[descriptor_slice] = self.kernel_block_size * row_bytes
self._submit_dma_descriptors(
descriptors,
descriptor_count,
transfer_ids=tuple(transfer.transfer_id for transfer in transfers),
)
def _enqueue_host_mirror(
self,
ready_event: torch.Event | None = None,
) -> None:
active = [
(index, handle)
for index, handle in enumerate(self.cache_handles)
if handle.num_actual_tokens != 0
]
if not active:
return
def mirror_key(handle: HiSparseCacheHandle) -> tuple:
slots = handle.mirror_slot_mapping
return (
handle.num_actual_tokens,
handle.num_decode_tokens,
handle.decode_batch,
handle.host_mirror_required,
handle.runtime.eager_host_mirror,
None if slots is None else slots.data_ptr(),
)
keys = {mirror_key(handle) for _, handle in active}
if len(keys) > 1:
raise RuntimeError(
f"HiSparse cache layers disagree on mirror metadata: {list(keys)}."
)
cache = active[0][1]
if not cache.host_mirror_required:
return
dst_slots = cache.mirror_slot_mapping
if dst_slots is None:
raise RuntimeError("HiSparse host mirror has no source slot mapping.")
num_rows = min(cache.num_actual_tokens, dst_slots.numel())
if num_rows == 0:
return
if self.is_host_writer:
expected_layers = {index for index, _ in active}
if self._per_layer_mirrored and self._per_layer_mirrored != expected_layers:
raise RuntimeError(
"HiSparse per-layer DMA did not mirror every active layer: "
f"expected {sorted(expected_layers)}, got "
f"{sorted(self._per_layer_mirrored)}."
)
pending_layers = tuple(
sorted(expected_layers - self._submitted_mirror_layers)
)
if pending_layers:
self._enqueue_row_dma(pending_layers, ready_event=ready_event)
self._submitted_mirror_layers.update(pending_layers)
assert cache.req_id_per_token is not None
for _, handle in active:
if handle.runtime.is_group_leader:
handle.runtime.invalidate_written_slots(
dst_slots[:num_rows],
cache.req_id_per_token[:num_rows],
)
def _finish_mirror_phase(self, ready_event: torch.Event | None = None) -> None:
state = self._slot_mapping_staging
if state is not None and state.num_tokens:
# Measured free: blocking here costs no throughput (26.7 vs 27.6
# gen tok/s) and lets every step mirror exactly the written rows.
state.event.synchronize()
self._resolve_row_mirrors(state)
self._enqueue_host_mirror(ready_event)
self._clear_forward_mirror_state()
def finish_forward(self) -> None:
compute_stream = current_stream()
self._forward_ready_event.record()
self._finish_mirror_phase(self._forward_ready_event)
transfers = self._post_forward_transfers
self._post_forward_transfers = []
self._submit_transfers(transfers)
if self.is_host_writer:
if self._dma_submitted:
compute_stream.wait_event(self.host_write_event)
self._dma_submitted = False
else:
self.host_write_event.record(compute_stream)
self._release_completed_dma_descriptors()
def take_completed_host_copies(self) -> list[int]:
"""Drain host copies this worker has enqueued for this step."""
completed = self._completed_host_copy_dst_ids
self._completed_host_copy_dst_ids = []
return completed
def take_transfer_updates(self) -> tuple[list[int], list[int]]:
enqueued = self._enqueued_transfer_ids
self._enqueued_transfer_ids = []
completed: list[int] = []
while self._pending_transfer_events:
event, transfer_ids = self._pending_transfer_events[0]
if not event.query():
break
self._pending_transfer_events.popleft()
completed.extend(transfer_ids)
return enqueued, completed
def shutdown(self) -> None:
if not self._initialized:
return
if self._slot_mapping_staging is not None:
self._slot_mapping_staging.stream.synchronize()
if self.dma_stream is not None:
self.dma_stream.synchronize()
release_pinned_state(
[cache.runtime for cache in self.cache_handles], self.pinned_host_pools
)
self._initialized = False
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
+3 -3
View File
@@ -702,9 +702,9 @@ class LMCacheMPConnectorUpstream(KVConnectorBase_V1):
- Applies to both sync- and async-loading requests.
- Async loading: failed blocks may be reported in any forward pass
up to and including the pass where the request ID is returned by
`get_finished()`. Even if failures occur, the request must still
be reported via `get_finished()`, and the failed block IDs must
appear here no later than that same pass.
`get_transfer_results()`. Even if failures occur, the request
must still be reported as finished receiving, and the failed
block IDs must appear here no later than that same pass.
- Sync loading: failed blocks should be reported in the forward
pass in which they are detected.
"""
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+17 -15
View File
@@ -635,6 +635,7 @@ class MooncakeConnectorScheduler:
):
self.vllm_config = vllm_config
self.block_size = vllm_config.cache_config.block_size
self.kv_cache_config = kv_cache_config
assert vllm_config.kv_transfer_config
self.is_kv_producer: bool = (
@@ -649,7 +650,7 @@ class MooncakeConnectorScheduler:
not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
and any(
not isinstance(g.kv_cache_spec, FullAttentionSpec)
for g in kv_cache_config.kv_cache_groups
for g in kv_cache_config.transfer_groups
)
)
# GDN is represented as a MambaSpec in vLLM. This Mooncake MambaSpec
@@ -670,7 +671,7 @@ class MooncakeConnectorScheduler:
(g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
if isinstance(g.kv_cache_spec, SlidingWindowSpec)
else (0, self.block_size)
for g in kv_cache_config.kv_cache_groups
for g in kv_cache_config.transfer_groups
]
# cdiv(n_tokens, block_size) gives blocks/window; add 1 to
# conservatively account for boundary overlap.
@@ -684,11 +685,17 @@ class MooncakeConnectorScheduler:
block_ids: tuple[list[int], ...] | list[list[int]],
) -> list[list[int]]:
"""Clip per-group block IDs to sliding window size."""
if len(block_ids) == 0 or not self._is_hma_required:
return list(block_ids)
if len(block_ids) == 0:
return []
if len(block_ids) == len(self.kv_cache_config.transfer_group_ids):
selected = list(block_ids)
else:
selected = list(self.kv_cache_config.select_transfer_block_ids(block_ids))
if len(selected) == 0 or not self._is_hma_required:
return selected
return [
blocks[-self.blocks_per_sw[i] :] if self.blocks_per_sw[i] > 0 else blocks
for i, blocks in enumerate(block_ids)
for i, blocks in enumerate(selected)
]
def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:
@@ -1046,18 +1053,13 @@ class MooncakeConnectorWorker:
self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size}
self._layer_specs: dict[str, KVCacheSpec] = {}
for group in kv_cache_config.kv_cache_groups:
for group in kv_cache_config.transfer_groups:
group_spec = group.kv_cache_spec
specs_by_layer = getattr(group_spec, "kv_cache_specs", {})
for layer_name in group.layer_names:
self._layer_specs[layer_name] = specs_by_layer.get(
layer_name, group_spec
)
self._layer_group_indices: dict[str, int] = {
layer: group_index
for group_index, group in enumerate(kv_cache_config.kv_cache_groups)
for layer in group.layer_names
}
self.transfer_topo = TransferTopology(
tp_rank=self.tp_rank,
tp_size=self.tp_size,
@@ -1441,7 +1443,7 @@ class MooncakeConnectorWorker:
block_arange = np.arange(self._physical_blocks_per_logical_kv_block).reshape(
1, -1
)
group_specs = self.kv_cache_config.kv_cache_groups
group_specs = self.kv_cache_config.transfer_groups
return [
BlockTable.map_to_kernel_blocks(
np.array(group),
@@ -1494,7 +1496,7 @@ class MooncakeConnectorWorker:
local_block_ids_by_group: list[list[int]] = []
remote_block_ids_by_group: list[list[int]] = []
has_block_error = False
group_specs = self.kv_cache_config.kv_cache_groups
group_specs = self.kv_cache_config.transfer_groups
for group_index, (local_group, remote_group) in enumerate(
zip(send_meta.local_block_ids, remote_block_ids_per_group)
):
@@ -1750,7 +1752,7 @@ class MooncakeConnectorWorker:
self.registered_layer_names.append(layer_name)
self.registered_layer_indices.append(layer_index)
self.registered_group_indices.append(
self._layer_group_indices[layer_name]
self.kv_cache_config.transfer_group_index_by_layer[layer_name]
)
storage = cache.untyped_storage()
storage_addr = storage.data_ptr()
@@ -2121,7 +2123,7 @@ class MooncakeConnectorWorker:
) -> list[TransferRegion]:
if not group_indices:
group_indices = [
self._layer_group_indices.get(layer_name, 0)
self.kv_cache_config.transfer_group_index_by_layer.get(layer_name, 0)
for layer_name in layer_names
]
return _expand_transfer_regions(
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+13 -4
View File
@@ -3,18 +3,20 @@
"""External-store cache-hit coordinator for MooncakeStoreConnector."""
from collections.abc import Sequence
from typing import cast
from typing import NamedTuple, cast
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
chunk_hashes_for_block_size,
)
from vllm.utils.math_utils import cdiv
from vllm.v1.core.block_pool import BlockPool
from vllm.v1.core.kv_cache_coordinator import SpecGroup
from vllm.v1.core.kv_cache_utils import (
BlockHash,
KVCacheBlock,
)
from vllm.v1.core.single_type_kv_cache_manager import (
SingleTypeKVCacheManager,
)
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheGroupSpec,
@@ -25,6 +27,13 @@ from vllm.v1.kv_cache_interface import (
from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
class StoreSpecGroup(NamedTuple):
spec: KVCacheSpec
group_ids: list[int]
manager_cls: type[SingleTypeKVCacheManager]
use_eagle: bool
class ExternalCachedBlockPool:
"""Duck-typed BlockPool backed by a ``(group_id, hash)`` exists set."""
@@ -121,7 +130,7 @@ class MooncakeStoreCoordinator:
"""Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but
dispatches via spec_manager_map (we don't allocate managers).
"""
attention_groups: list[SpecGroup] = []
attention_groups: list[StoreSpecGroup] = []
for i, g in enumerate(self.kv_cache_groups):
# Skip groups that opt out of prefix caching (e.g. GLM-5.3-Flash
# kpool tail): per-request scratch, never shareable, so they must
@@ -143,7 +152,7 @@ class MooncakeStoreCoordinator:
break
else:
attention_groups.append(
SpecGroup(spec, [i], manager_cls, g.is_eagle_group)
StoreSpecGroup(spec, [i], manager_cls, g.is_eagle_group)
)
# Full attention first (matches upstream convergence ordering).
attention_groups.sort(key=lambda g: not isinstance(g.spec, FullAttentionSpec))
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+134 -17
View File
@@ -67,6 +67,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import
from vllm.logger import init_logger
from vllm.utils.math_utils import cdiv
from vllm.utils.network_utils import get_ip, make_zmq_socket
from vllm.utils.torch_utils import is_non_overlapping_and_dense
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
from vllm.v1.core.kv_cache_utils import (
BlockHash,
@@ -1909,21 +1910,31 @@ class MooncakeStoreWorker:
logger.warning("No KV caches to offload.")
return
# Resolve each entry to a representative tensor for storage
# deduplication. For attention layers the value is already a tensor;
# for Mamba layers it is a list of tensors that all share the same
# underlying raw storage, so we take the first one.
def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
assert isinstance(v, torch.Tensor | list)
return v if isinstance(v, torch.Tensor) else v[0]
assert self.cache_config.num_gpu_blocks is not None
self.num_blocks = self.cache_config.num_gpu_blocks
use_group_regions = self._kv_cache_config.hisparse_host_num_blocks is not None
seen_storage_ptrs: set[int] = set()
cache_tensors: list[torch.Tensor] = []
for cache in kv_caches.values():
cache = group_kernel_blocks(cache, self.num_blocks)
cache_tensors.append(cache)
cache_storage = cache.untyped_storage()
base_addr = cache_storage.data_ptr()
region_len = cache_storage.nbytes()
if base_addr not in seen_storage_ptrs:
if not use_group_regions:
seen_storage_ptrs: set[int] = set()
cache_tensors = [
group_kernel_blocks(_repr_tensor(cache), self.num_blocks)
for cache in kv_caches.values()
]
for cache in cache_tensors:
cache_storage = cache.untyped_storage()
base_addr = cache_storage.data_ptr()
if base_addr in seen_storage_ptrs:
continue
seen_storage_ptrs.add(base_addr)
region_len = cache_storage.nbytes()
ret = self.store.register_buffer(base_addr, region_len)
if ret != 0:
logger.error(
@@ -1932,17 +1943,110 @@ class MooncakeStoreWorker:
region_len,
ret,
)
for db in self.token_dbs:
db.store_layout.register_kv_caches(cache_tensors, self.num_blocks)
registered_buffers: dict[int, int] = {}
num_host_segments = 0
num_device_segments = 0
for group, db in zip(self._kv_cache_groups, self.token_dbs, strict=True):
if not use_group_regions:
continue
missing_layers = [
layer_name
for layer_name in group.layer_names
if layer_name not in kv_caches
]
if missing_layers:
raise ValueError(
f"Missing transferable KV caches for layers: {missing_layers}"
)
group_caches = [
(layer_name, kv_caches[layer_name]) for layer_name in group.layer_names
]
seen_group_ptrs: set[int] = set()
addrs: list[int] = []
block_lens: list[int] = []
for layer_name, value in group_caches:
cache = _repr_tensor(value)
is_host_resident = group.host_resident
if is_host_resident:
num_blocks = self._kv_cache_config.hisparse_host_num_blocks
assert num_blocks is not None
else:
num_blocks = self._kv_cache_config.num_blocks
cache = group_kernel_blocks(cache, num_blocks)
cache_storage = cache.untyped_storage()
if is_host_resident:
base_addr = cache.data_ptr()
region_len = cache.nbytes
else:
base_addr = cache_storage.data_ptr()
region_len = cache_storage.nbytes()
if region_len % num_blocks != 0:
raise ValueError(
f"KV cache region for {layer_name} has size {region_len}, "
f"which is not divisible by its {num_blocks} blocks."
)
registered_len = registered_buffers.get(base_addr)
if registered_len is None:
ret = self.store.register_buffer(base_addr, region_len)
if ret != 0:
mem_kind = "host" if is_host_resident else "device"
raise RuntimeError(
"Mooncake register_buffer failed for addr "
f"{base_addr:#x} len {region_len} ({mem_kind}): {ret}"
)
registered_buffers[base_addr] = region_len
elif registered_len != region_len:
raise ValueError(
f"KV cache views at {base_addr:#x} expose inconsistent "
f"region sizes: {registered_len} and {region_len}."
)
regions: list[tuple[int, int]] = []
if not is_non_overlapping_and_dense(cache[0]):
for head_idx in range(cache.shape[1]):
head_cache = cache[:, head_idx]
assert is_non_overlapping_and_dense(head_cache[0])
regions.append(
(
head_cache.data_ptr(),
head_cache.stride(0) * head_cache.element_size(),
)
)
elif cache.stride(0) * cache.element_size() * num_blocks == region_len:
regions.append((base_addr, region_len // num_blocks))
else:
regions.append(
(cache.data_ptr(), cache.stride(0) * cache.element_size())
)
for region_addr, block_len in regions:
if region_addr in seen_group_ptrs:
continue
seen_group_ptrs.add(region_addr)
addrs.append(region_addr)
block_lens.append(block_len)
if is_host_resident:
num_host_segments += 1
else:
num_device_segments += 1
db.set_kv_caches_base_addr(addrs)
db.set_block_len(block_lens)
logger.info(
"Registered KV caches: num_groups=%d, num_tensors=%d, num_blocks=%d",
"Registered KV caches: num_groups=%d, num_segments=%d (host=%d, device=%d)",
len(self.token_dbs),
len(cache_tensors),
self.num_blocks,
num_host_segments + num_device_segments,
num_host_segments,
num_device_segments,
)
for db in self.token_dbs:
db.store_layout.register_kv_caches(cache_tensors, self.num_blocks)
# Start transfer threads
if self.can_put:
ready_event_sending = threading.Event()
@@ -2316,6 +2420,19 @@ class MooncakeStoreWorker:
if store is None:
return
self.store = None
deadline = time.monotonic() + 5.0
for thread in (self.kv_send_thread, *self.kv_recv_threads):
if thread is None:
continue
while thread.request_queue.unfinished_tasks and time.monotonic() < deadline:
time.sleep(0.05)
if thread.request_queue.unfinished_tasks:
logger.warning(
"Mooncake store %s still has %d in-flight requests at "
"close; closing anyway.",
thread.name,
thread.request_queue.unfinished_tasks,
)
try:
store.close()
except Exception as e:
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+44 -22
View File
@@ -17,6 +17,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
KVConnectorHandshakeMetadata,
KVConnectorMetadata,
KVConnectorRole,
KVConnectorTransferResults,
KVConnectorWorkerMetadata,
SupportsHMA,
)
@@ -35,7 +36,7 @@ if TYPE_CHECKING:
from vllm.distributed.kv_events import KVCacheEvent
from vllm.forward_context import ForwardContext
from vllm.v1.core.block_pool import BlockPool
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.request import Request
@@ -208,6 +209,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
# Propagated from scheduler to worker side via the connector metadata.
self._extra_async_saves: dict[str, int] = {}
@property
def sub_connectors(self) -> list[KVConnectorBase_V1]:
return list(self._connectors)
@property
def supports_divergent_local_hybrid_hits(self) -> bool:
return bool(self._connectors) and all(
@@ -248,6 +253,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
for c in self._connectors:
c.register_kv_caches(kv_caches)
def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None:
for connector in self._connectors:
connector.bind_kv_cache_manager(kv_cache_manager)
def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
for c in self._connectors:
c.bind_gpu_block_pool(gpu_block_pool)
@@ -292,6 +301,14 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
for c in self._connectors:
c.start_load_kv(forward_context, **kwargs)
def finish_forward(self) -> None:
for c in self._connectors:
c.finish_forward()
def reset_capture_state(self) -> None:
for c in self._connectors:
c.reset_capture_state()
def wait_for_layer_load(self, layer_name: str) -> None:
for c in self._connectors:
c.wait_for_layer_load(layer_name)
@@ -313,29 +330,31 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
def get_finished(
self, finished_req_ids: set[str]
) -> tuple[set[str] | None, set[str] | None]:
finished_sending: set[str] = set()
finished_recving: set[str] = set()
for c in self._connectors:
sending, recving = c.get_finished(finished_req_ids)
if not recving and not sending:
continue
# Aggregate finished recving request ids.
finished_recving.update(recving or ())
# Aggregate finished sending request ids - only include
# once we've drained the "extra" count (for cases where
# more than one connector is async-saving the same request).
for req_id in sending or ():
results = self.get_transfer_results(finished_req_ids)
return (
results.finished_sending or None,
results.finished_recving or None,
)
def get_transfer_results(
self, finished_req_ids: set[str]
) -> KVConnectorTransferResults:
results = KVConnectorTransferResults()
for connector in self._connectors:
child_results = connector.get_transfer_results(finished_req_ids)
results.finished_recving.update(child_results.finished_recving)
results.failed_recving.update(child_results.failed_recving)
for req_id in child_results.finished_sending:
extra_pending = self._extra_async_saves.get(req_id)
if extra_pending is None:
finished_sending.add(req_id)
continue
assert extra_pending > 0
if extra_pending == 1:
del self._extra_async_saves[req_id]
results.finished_sending.add(req_id)
else:
self._extra_async_saves[req_id] = extra_pending - 1
return finished_sending or None, finished_recving or None
assert extra_pending > 0
if extra_pending == 1:
del self._extra_async_saves[req_id]
else:
self._extra_async_saves[req_id] = extra_pending - 1
return results
def get_block_ids_with_load_errors(self) -> set[int]:
agg_block_ids: set[int] = set()
@@ -563,6 +582,9 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
for c in self._connectors:
yield from c.take_events()
def has_pending_block_frees(self) -> bool:
return any(c.has_pending_block_frees() for c in self._connectors)
def has_pending_push_work(self) -> bool:
return any(c.has_pending_push_work() for c in self._connectors)
@@ -630,7 +652,7 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
return MultiKVConnectorStats(data=reconstructed_data)
def get_kv_connector_stats(self) -> MultiKVConnectorStats | None:
def get_kv_connector_stats(self) -> KVConnectorStats | None:
# Group connector stats by connector type.
stats_by_connector: MultiKVConnectorStats | None = None
for c in self._connectors:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+1 -4
View File
@@ -93,10 +93,7 @@ class NixlBaseConnectorScheduler:
for g in kv_cache_config.transfer_groups
)
)
self._has_mamba = any(
isinstance(g.kv_cache_spec, MambaSpec)
for g in kv_cache_config.transfer_groups
)
self._has_mamba = kv_cache_config.has_mamba_layers
logger.info("Initializing NIXL Scheduler %s", engine_id)
if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+218 -136
View File
@@ -2,8 +2,10 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Base worker-side logic for the NIXL connector."""
import contextlib
import itertools
import logging
import math
import os
import queue
import threading
@@ -31,7 +33,10 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (
kv_postprocess_blksize_on_receive,
kv_postprocess_layout_on_receive,
)
from vllm.distributed.kv_transfer.kv_connector.v1.base import CopyBlocksOp
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
CopyBlocksOp,
KVConnectorTransferResults,
)
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
GET_META_MSG,
@@ -186,23 +191,28 @@ class NixlBaseConnectorWorker:
num_blocks = dst_num_blocks
if block_size_ratio is not None:
num_blocks = int(num_blocks * block_size_ratio)
num_regions = self.num_regions or len(self.block_len_per_layer)
if region_num_blocks is None:
region_num_blocks = [num_blocks] * self.num_regions
region_num_blocks = [num_blocks] * num_regions
elif block_size_ratio is not None:
region_num_blocks = [
int(count * block_size_ratio) for count in region_num_blocks
]
assert len(region_num_blocks) == self.num_regions
if region_num_blocks:
num_regions = len(region_num_blocks)
assert len(region_num_blocks) == num_regions
region_offsets = np.cumsum([0, *region_num_blocks[:-1]])
num_fa_descs = sum(region_num_blocks)
# All-attention fast path: single vectorized broadcast.
# All-attention fast path.
if num_ssm_regions == 0:
if region_group_ids is None:
region_group_ids = self.region_group_ids
if uses_region_group_mapping is None:
uses_region_group_mapping = self._uses_region_group_mapping
assert len(region_group_ids) == self.num_regions
if not region_group_ids and num_regions == 1:
region_group_ids = [0]
assert len(region_group_ids) == num_regions
if uses_region_group_mapping is None:
uses_region_group_mapping = len(set(region_group_ids)) > 1
region_group_ids_array = np.asarray(region_group_ids, dtype=np.int32)
@@ -397,8 +407,6 @@ class NixlBaseConnectorWorker:
# Config.
self.vllm_config = vllm_config
# mypy will complain on re-assignment otherwise.
self.block_size: int = cast(int, vllm_config.cache_config.block_size)
if vllm_config.kv_transfer_config is None:
raise ValueError("kv_transfer_config must be set for NixlConnector")
@@ -420,6 +428,17 @@ class NixlBaseConnectorWorker:
)
self.kv_cache_config = kv_cache_config
transfer_block_sizes = [
group.kv_cache_spec.block_size
for group in kv_cache_config.transfer_groups
if get_representative_spec_type(group.kv_cache_spec)
not in (MambaSpec, CircularBufferSpec)
]
self.block_size = (
math.lcm(*transfer_block_sizes)
if transfer_block_sizes
else cast(int, vllm_config.cache_config.block_size)
)
# Per-layer specs, unwrapping UniformTypeKVCacheSpecs group wrappers.
self._layer_specs: dict[str, KVCacheSpec] = {}
for group in kv_cache_config.transfer_groups:
@@ -599,12 +618,15 @@ class NixlBaseConnectorWorker:
# Number of NIXL regions. Currently one region per cache
# (so 1 per layer for MLA, otherwise 2 per layer)
self.num_regions = 0
self.region_mem_types: list[str] = []
self.region_group_ids: list[int] = []
self._uses_region_group_mapping = False
self.region_names: list[str] = []
self.region_num_blocks: list[int] = []
self.region_mem_types: list[str] = []
self._mixed_mem_types = False
self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {}
self._desc_pos_by_block_size: dict[int, np.ndarray] = {}
self._dram_src_handles_by_block_size: dict[int, int] = {}
# PP>1 (push mode): this worker holds a contiguous layer slice and
# transfers into the matching sub-range of a PP=1 remote's regions.
@@ -644,9 +666,6 @@ class NixlBaseConnectorWorker:
self.dst_region_group_ids: dict[EngineId, list[int]] = {}
self.dst_uses_region_group_mapping: dict[EngineId, bool] = {}
self.dst_region_mem_types: dict[EngineId, list[str]] = {}
self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {}
self._desc_pos_by_block_size: dict[int, np.ndarray] = {}
self._dram_src_handles_by_block_size: dict[int, int] = {}
self._registered_descs: list[Any] = []
# In progress transfers.
@@ -667,6 +686,7 @@ class NixlBaseConnectorWorker:
# Uses Queue for thread-safe cross-thread coordination with the
# background handshake thread, matching the _ready_requests pattern.
self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()
self._failed_inflight_recvs: set[ReqId] = set()
self._pending_recv_notifs: dict[ReqId, list[tuple[str, bytes]]] = {}
# Handshake metadata of this worker for NIXL transfers.
@@ -690,7 +710,6 @@ class NixlBaseConnectorWorker:
"engine_ttl", 3600.0
)
self.block_size = vllm_config.cache_config.block_size
self.model_config = vllm_config.model_config
self.use_mla = self.model_config.use_mla
@@ -1240,7 +1259,7 @@ class NixlBaseConnectorWorker:
self.use_host_buffer,
)
registration_ranges: dict[tuple[int, str], tuple[int, int, int, str]] = {}
registration_ranges: dict[tuple[int, str], tuple[int, int, int]] = {}
region_mem_types: list[str] = []
seen_base_addresses: list[int] = []
self._ssm_region_indices = []
@@ -1280,10 +1299,16 @@ class NixlBaseConnectorWorker:
# K and V are packed into the content dim, so each attention layer is a
# single NIXL region whose block transfers as one unit. Mamba layers instead
# register separate conv/ssm sub-regions (see `_build_mamba_local`).
# P and D may allocate equivalent transferable layers in different
# cache-group orders. Keep their region lists aligned without putting
# layers.10 before layers.2, which would break PP region slicing.
layer_names = (
xfer_buffers
if self._is_csa_linear
else sorted(xfer_buffers, key=_region_sort_key)
else sorted(
xfer_buffers,
key=_region_sort_key,
)
)
for layer_name in layer_names:
cache = xfer_buffers[layer_name]
@@ -1299,6 +1324,7 @@ class NixlBaseConnectorWorker:
layer_name,
)
continue
group_index = self.kv_cache_config.transfer_group_index_by_layer[layer_name]
# `layer_spec.page_size_bytes` only accounts for logical page_size, that is
# the page_size assuming constant `self._logical_num_blocks`.
physical_page_size = (
@@ -1307,17 +1333,48 @@ class NixlBaseConnectorWorker:
else layer_spec.page_size_bytes
// self._physical_blocks_per_logical_kv_block
)
group_id = self.kv_cache_config.transfer_group_index_by_layer[layer_name]
group = self.kv_cache_config.transfer_groups[group_index]
if group.host_resident:
logical_num_blocks = self.kv_cache_config.hisparse_host_num_blocks
assert logical_num_blocks is not None
else:
logical_num_blocks = self.kv_cache_config.num_blocks
group_id = group_index
num_blocks = (
self._logical_num_blocks
logical_num_blocks
if isinstance(layer_spec, MambaSpec)
else self.num_blocks
else logical_num_blocks * self._physical_blocks_per_logical_kv_block
)
base_addr = cache.data_ptr()
is_mla_region = isinstance(
layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
)
logger.debug(
"Registering layer %s with cache shape: %s", layer_name, cache.shape
)
storage = cache.untyped_storage()
storage_addr = storage.data_ptr()
is_host_resident = group.host_resident
if cache.device.type == "cpu":
mem_type = "DRAM"
region_device_id = 0
else:
mem_type = self.nixl_memory_type
region_device_id = max(cache.get_device(), 0)
self.device_id = region_device_id
if is_host_resident:
registration_base = cache.data_ptr()
registration_len = cache.nbytes
else:
registration_base = storage_addr
registration_len = storage.nbytes()
storage_key = (registration_base, mem_type)
if storage_key not in registration_ranges:
registration_ranges[storage_key] = (
registration_base,
registration_base + registration_len,
region_device_id,
)
if isinstance(layer_spec, KpoolTailSpec):
compressed_owner = compressed_region_owners.get(cache.data_ptr())
@@ -1345,22 +1402,6 @@ class NixlBaseConnectorWorker:
)
continue
if cache.device.type == "cpu":
mem_type = "DRAM"
region_device_id = 0
else:
mem_type = self.nixl_memory_type
region_device_id = max(cache.get_device(), 0)
self.device_id = region_device_id
registration_ranges.setdefault(
(storage_addr, mem_type),
(storage_addr, storage.nbytes(), region_device_id, ""),
)
is_mla_region = isinstance(
layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
)
if isinstance(layer_spec, MambaSpec):
physical_ratio = self._physical_blocks_per_logical_kv_block
block_len = physical_page_size // physical_ratio
@@ -1381,14 +1422,13 @@ class NixlBaseConnectorWorker:
block_stride = physical_page_size
else:
block_stride = cache.stride(0) * cache.element_size()
storage_is_block_major = num_blocks * block_stride == storage.nbytes()
# A layer whose [H, N, C] interior is dense addresses its own page
# as one chunk. Otherwise the block's whole row is the only
# contiguous transfer unit.
hnc_contiguous = (
cache.ndim == 4
and cache.stride(2) == cache.shape[3]
and cache.stride(1) == cache.shape[2] * cache.shape[3]
storage_is_block_major = num_blocks * block_stride == registration_len
# A view with a contiguous page can own a strided region even
# when sibling layers share the same backing allocation.
page_contiguous = (
cache.ndim > 1
and cache[0].is_contiguous()
and cache[0].nbytes == physical_page_size
)
virtual_transfer_pages = _uses_dense_virtual_transfer_pages(
layer_spec, cache, physical_page_size, num_blocks
@@ -1399,13 +1439,14 @@ class NixlBaseConnectorWorker:
(cache.data_ptr(), physical_page_size, physical_page_size)
]
elif storage_is_block_major and (
(packed_storage and is_mla_region)
or (not hnc_contiguous and not self._is_csa_linear)
not page_contiguous
and ((packed_storage and is_mla_region) or not self._is_csa_linear)
):
# Packed MLA layouts transfer the complete storage row.
storage_block_len = storage.nbytes() // num_blocks
# TODO(Lucas): handle TP slicing for packed_storage; for now
# restrict to MLA (DSv4) where kv is replicated.
storage_block_len = registration_len // num_blocks
region_specs = [
(storage_addr, storage_block_len, storage_block_len)
(registration_base, storage_block_len, storage_block_len)
]
elif storage_is_block_major:
region_specs = [
@@ -1440,6 +1481,10 @@ class NixlBaseConnectorWorker:
region_index = seen_base_addresses.index(base_addr)
assert region_mem_types[region_index] == mem_type
self._region_is_mla[region_index] |= is_mla_region
if is_mla_region:
self.block_len_per_layer[region_index] = block_len
self.block_stride_per_layer[region_index] = block_stride
self.region_num_blocks[region_index] = num_blocks
if self.region_group_ids[region_index] != group_id:
self.region_group_ids[region_index] = _SHARED_REGION_GROUP_ID
else:
@@ -1450,8 +1495,8 @@ class NixlBaseConnectorWorker:
self.region_group_ids.append(group_id)
self.region_names.append(layer_name)
self.region_num_blocks.append(num_blocks)
region_mem_types.append(mem_type)
self._region_is_mla.append(is_mla_region)
region_mem_types.append(mem_type)
if isinstance(layer_spec, MambaSpec):
if layer_spec.tp_replicated:
@@ -1496,7 +1541,6 @@ class NixlBaseConnectorWorker:
== len(self.region_group_ids)
== len(self.region_names)
== len(self.region_num_blocks)
== len(region_mem_types)
)
# Descriptor ids must be region-ordered, matching the remote side.
self._scratch_region_indices.sort()
@@ -1521,16 +1565,23 @@ class NixlBaseConnectorWorker:
# Total local FA descriptors (boundary between FA and mamba descs).
self.num_descs = sum(self.region_num_blocks)
self._mixed_mem_types = len(set(region_mem_types)) > 1
if self._mixed_mem_types:
assert self.use_mla and not self._has_mamba, (
"Mixed-device KV registration is only supported for MLA "
"models without Mamba layers."
)
for mem_type in sorted(set(region_mem_types)):
ranges = [
registration
for (
_,
registration_mem_type,
), registration in registration_ranges.items()
if registration_mem_type == mem_type
ranges_for_mem_type = [
(start, end - start, device_id, "")
for (_, cache_mem_type), (
start,
end,
device_id,
) in registration_ranges.items()
if cache_mem_type == mem_type
]
descs = self.nixl_wrapper.get_reg_descs(ranges, mem_type)
descs = self.nixl_wrapper.get_reg_descs(ranges_for_mem_type, mem_type)
self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
self._registered_descs.append(descs)
@@ -1854,7 +1905,11 @@ class NixlBaseConnectorWorker:
if self._mixed_mem_types:
desc_is_dram = np.concatenate(
[
np.full(count * block_size_ratio, mem_type == "DRAM", dtype=bool)
np.full(
count * block_size_ratio,
mem_type == "DRAM",
dtype=bool,
)
for mem_type, count in zip(
self.region_mem_types,
self.region_num_blocks,
@@ -1864,24 +1919,29 @@ class NixlBaseConnectorWorker:
)
assert len(desc_is_dram) == len(blocks_data)
desc_pos = np.empty(len(desc_is_dram), dtype=np.int64)
dram_indices = np.flatnonzero(desc_is_dram)
device_indices = np.flatnonzero(~desc_is_dram)
desc_pos[dram_indices] = np.arange(len(dram_indices), dtype=np.int64)
desc_pos[device_indices] = np.arange(len(device_indices), dtype=np.int64)
dram_idx = np.where(desc_is_dram)[0]
vram_idx = np.where(~desc_is_dram)[0]
desc_pos[dram_idx] = np.arange(len(dram_idx), dtype=np.int64)
desc_pos[vram_idx] = np.arange(len(vram_idx), dtype=np.int64)
self._desc_is_dram_by_block_size[block_size] = desc_is_dram
self._desc_pos_by_block_size[block_size] = desc_pos
blocks_data[desc_is_dram, 2] = 0
dram_descs = self.nixl_wrapper.get_xfer_descs(
blocks_data[desc_is_dram], "DRAM"
)
# DRAM descriptors are registered under CPU device 0.
blocks_data = [
(addr, length, 0) if is_dram else (addr, length, dev)
for (addr, length, dev), is_dram in zip(
blocks_data, desc_is_dram, strict=True
)
]
dram_blocks = [blocks_data[i] for i in dram_idx]
vram_blocks = [blocks_data[i] for i in vram_idx]
dram_descs = self.nixl_wrapper.get_xfer_descs(dram_blocks, "DRAM")
self._dram_src_handles_by_block_size[block_size] = (
self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs)
)
device_descs = self.nixl_wrapper.get_xfer_descs(
blocks_data[~desc_is_dram], self.nixl_memory_type
)
descs = self.nixl_wrapper.get_xfer_descs(vram_blocks, self.nixl_memory_type)
return (
self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", device_descs),
self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs),
blocks_data,
)
@@ -1968,10 +2028,6 @@ class NixlBaseConnectorWorker:
]
nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end]
nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end]
if nixl_agent_meta.region_num_blocks is not None:
nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[
start:end
]
if nixl_agent_meta.region_group_ids is not None:
nixl_agent_meta.region_group_ids = nixl_agent_meta.region_group_ids[
start:end
@@ -1982,6 +2038,10 @@ class NixlBaseConnectorWorker:
nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[
start:end
]
if nixl_agent_meta.region_num_blocks is not None:
nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[
start:end
]
### Register remote engine in TransferTopology (idempotent).
assert self.transfer_topo is not None
@@ -2026,11 +2086,12 @@ class NixlBaseConnectorWorker:
nixl_agent_meta.region_num_blocks
or [nixl_agent_meta.num_blocks] * num_remote_regions
)
self.dst_region_group_ids[engine_id] = nixl_agent_meta.region_group_ids or (
remote_region_group_ids = nixl_agent_meta.region_group_ids or (
self.region_group_ids
if len(self.region_group_ids) == num_remote_regions
else [0] * num_remote_regions
)
self.dst_region_group_ids[engine_id] = remote_region_group_ids
self.dst_uses_region_group_mapping[engine_id] = (
len(set(self.dst_region_group_ids[engine_id])) > 1
)
@@ -2138,8 +2199,6 @@ class NixlBaseConnectorWorker:
mamba = self._build_mamba_remote(nixl_agent_meta, tp_ratio, transfer_info)
blocks_data = np.concatenate([blocks_data, mamba])
# A mixed local destination is split into multiple READs below. The
# producer remains a single descriptor list in this prerequisite.
remote_mem_types = set(self.dst_region_mem_types[engine_id])
if len(remote_mem_types) != 1:
raise NotImplementedError(
@@ -2540,9 +2599,10 @@ class NixlBaseConnectorWorker:
indices=indices,
)
def get_finished(self) -> tuple[set[str], set[str]]:
def get_transfer_results(self) -> KVConnectorTransferResults:
"""
Get requests that are done sending or recving on this specific worker.
Get transfers that completed on this specific worker.
The scheduler process (via the MultiprocExecutor) will use this output
to track which workers are done.
"""
@@ -2564,7 +2624,6 @@ class NixlBaseConnectorWorker:
# Add failed requests to done_recving for scheduler tracking
# (blocks are already marked invalid, scheduler will handle recompute)
done_recving.update(failed_recv_reqs)
if len(done_sending) > 0 or len(done_recving) > 0:
logger.debug(
"Rank %s, get_finished: %s requests done sending "
@@ -2577,6 +2636,7 @@ class NixlBaseConnectorWorker:
block_ids_for_blocksize_post_process = defaultdict(list)
block_ids_for_heterogeneous_attn_post_process = list[list[int]]()
direct_device_recving = set[str]()
for req_id in done_recving:
# clean up metadata for completed requests
meta = self._recving_metadata.pop(req_id, None)
@@ -2594,6 +2654,8 @@ class NixlBaseConnectorWorker:
if self.use_host_buffer:
self.sync_recved_kv_to_device(req_id, meta)
direct_device_recving.add(req_id)
# Post processing for heteroblocksize/layout, and for blocks the
# transfer clipped. The latter happens either at remote-block
# granularity (block_size_ratio > 1) or at kernel-block
@@ -2642,7 +2704,7 @@ class NixlBaseConnectorWorker:
for block_ids in block_ids_for_heterogeneous_attn_post_process:
self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids)
self._sync_device_after_mamba_recv(done_recving, failed_recv_reqs)
self._sync_device_after_direct_recv(direct_device_recving)
# Handle timeout to avoid stranding blocks on remote.
now = time.perf_counter()
@@ -2664,20 +2726,21 @@ class NixlBaseConnectorWorker:
del self._reqs_to_send[req_id]
done_sending.add(req_id)
return done_sending, done_recving
return KVConnectorTransferResults(
finished_sending=done_sending,
finished_recving=done_recving,
failed_recving=failed_recv_reqs,
)
def _sync_device_after_mamba_recv(
self,
done_recving: set[str],
failed_recv_reqs: set[str],
) -> None:
"""Synchronize ROCm direct-GPU Mamba receives before model execution."""
if (
not current_platform.is_rocm()
or not self._has_mamba
or self.use_host_buffer
or not (done_recving - failed_recv_reqs)
):
def get_finished(self) -> tuple[set[str], set[str]]:
"""Compatibility wrapper for the legacy completion API."""
results = self.get_transfer_results()
return results.finished_sending, results.finished_recving
def _sync_device_after_direct_recv(self, done_recving: set[str]) -> None:
"""Make direct NIXL writes visible before model execution."""
requires_sync = current_platform.is_rocm() and self._has_mamba
if self.use_host_buffer or not done_recving or not requires_sync:
return
torch.accelerator.synchronize()
@@ -2731,11 +2794,11 @@ class NixlBaseConnectorWorker:
self.nixl_wrapper.release_xfer_handle(handle)
elif xfer_state == "PROC":
in_progress.append(handle)
continue
else:
self._log_failure(
failure_type="transfer_failed",
msg="Marking blocks as invalid",
msg="Deferring request completion until its last "
"xfer is terminal",
req_id=req_id,
xfer_state=xfer_state,
)
@@ -2743,21 +2806,20 @@ class NixlBaseConnectorWorker:
except Exception as e:
self._log_failure(
failure_type="transfer_exception",
msg="Marking blocks as invalid",
msg="Handle is unpollable; treating it as terminal",
req_id=req_id,
error=e,
)
self._handle_failed_transfer(req_id, handle)
if not in_progress:
# Only report request as completed when all transfers are done.
# A request failed in an earlier poll was already reported via
# _failed_recv_reqs and its metadata popped by get_finished();
# don't report it again, just drop the remaining handles.
if req_id in self._recving_metadata:
done_req_ids.add(req_id)
self._send_pending_recv_notifs(req_id)
del transfers[req_id]
done_req_ids.add(req_id)
if req_id in self._failed_inflight_recvs:
self._failed_inflight_recvs.remove(req_id)
self._report_failed_recv(req_id)
else:
self._send_pending_recv_notifs(req_id)
else:
transfers[req_id] = in_progress
return done_req_ids
@@ -2771,22 +2833,25 @@ class NixlBaseConnectorWorker:
req_id: The request ID.
handle: The transfer handle.
"""
# (multi-read) One handle is created per remote rank, and they do not
# all fail in the same _pop_done_transfers poll. The request is
# reported failed on the first one, which pops its metadata in
# get_finished(); on the later failures only the handle cleanup is left.
# TODO (NickLucche) handle failed transfer for HMA.
# A split READ's notification is sent only after every handle succeeds.
# A sibling READ may still be writing these blocks. Retain metadata
# and defer invalidation until every handle is terminal.
self._pending_recv_notifs.pop(req_id, None)
if (meta := self._recving_metadata.get(req_id)) is not None:
if not self._is_hma_required:
self._invalid_block_ids.put(set(meta.local_block_ids[0]))
self._failed_recv_reqs.put(req_id)
if req_id in self._recving_transfers:
self._failed_inflight_recvs.add(req_id)
else:
self._report_failed_recv(req_id)
if handle is not None:
self.nixl_wrapper.release_xfer_handle(handle)
self.xfer_stats.record_failed_transfer()
def _report_failed_recv(self, req_id: str) -> None:
if (meta := self._recving_metadata.get(req_id)) is not None:
if not self._is_hma_required:
self._invalid_block_ids.put(set(meta.local_block_ids[0]))
self._failed_recv_reqs.put(req_id)
def _send_pending_recv_notifs(self, req_id: str) -> None:
"""Send notifications deferred by split DRAM/VRAM reads."""
for agent_name, notif_id in self._pending_recv_notifs.pop(req_id, []):
try:
self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)
@@ -2909,10 +2974,10 @@ class NixlBaseConnectorWorker:
return block_ids
block_arange = np.arange(0, ratio).reshape(1, -1)
# Mamba blocks have no logical<>physical discrepancy (block-size=1)
group_specs = self.kv_cache_config.transfer_groups
physical_block_ids = []
for i, group in enumerate(block_ids):
if _is_ssm_spec(get_representative_spec_type(group_specs[i].kv_cache_spec)):
spec = self.kv_cache_config.transfer_groups[i].kv_cache_spec
if _is_ssm_spec(get_representative_spec_type(spec)):
physical_block_ids.append(group)
else:
physical_block_ids.append(
@@ -3215,9 +3280,44 @@ class NixlBaseConnectorWorker:
)
def __del__(self):
self.shutdown()
with contextlib.suppress(Exception):
self.shutdown()
def shutdown(self):
def _finish_shutdown(self) -> None:
self._recving_transfers.clear()
try:
for handle in self.src_xfer_handles_by_block_size.values():
self.nixl_wrapper.release_dlist_handle(handle)
for handles in self.src_xfer_handles_by_tp_ratio.values():
for handle in handles:
self.nixl_wrapper.release_dlist_handle(handle)
for handles in self._dram_src_handles_by_tp_ratio.values():
for handle in handles:
self.nixl_wrapper.release_dlist_handle(handle)
for handle in self._dram_src_handles_by_block_size.values():
self.nixl_wrapper.release_dlist_handle(handle)
except Exception:
logger.exception("NIXL dlist-handle release failed at shutdown.")
self.src_xfer_handles_by_block_size.clear()
self.src_xfer_handles_by_tp_ratio.clear()
self._dram_src_handles_by_tp_ratio.clear()
self._dram_src_handles_by_block_size.clear()
try:
for engine_id in list(self._remote_agents):
self._cleanup_remote_engine(engine_id, log_eviction=False)
except Exception:
logger.exception("NIXL remote-engine cleanup failed at shutdown.")
try:
for desc in self._registered_descs:
self.nixl_wrapper.deregister_memory(desc)
finally:
self._registered_descs.clear()
# Drop cache references before their owners release registered
# host memory; handshake futures may outlive model-runner shutdown.
self.device_kv_caches = {}
self.host_xfer_buffers = {}
def shutdown(self) -> None:
"""Shutdown the connector worker."""
if not hasattr(self, "_handshake_initiation_executor"):
# error happens during init, no need to shutdown
@@ -3227,22 +3327,4 @@ class NixlBaseConnectorWorker:
for handle in handles:
self.nixl_wrapper.release_xfer_handle(handle)
self._recving_transfers.clear()
for handle in self.src_xfer_handles_by_block_size.values():
self.nixl_wrapper.release_dlist_handle(handle)
self.src_xfer_handles_by_block_size.clear()
for handles in self.src_xfer_handles_by_tp_ratio.values():
for handle in handles:
self.nixl_wrapper.release_dlist_handle(handle)
self.src_xfer_handles_by_tp_ratio.clear()
for handles in self._dram_src_handles_by_tp_ratio.values():
for handle in handles:
self.nixl_wrapper.release_dlist_handle(handle)
self._dram_src_handles_by_tp_ratio.clear()
for handle in self._dram_src_handles_by_block_size.values():
self.nixl_wrapper.release_dlist_handle(handle)
self._dram_src_handles_by_block_size.clear()
for engine_id in list(self._remote_agents):
self._cleanup_remote_engine(engine_id, log_eviction=False)
for desc in self._registered_descs:
self.nixl_wrapper.deregister_memory(desc)
self._registered_descs.clear()
self._finish_shutdown()
vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
+13
View File
@@ -26,6 +26,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
KVConnectorHandshakeMetadata,
KVConnectorMetadata,
KVConnectorRole,
KVConnectorTransferResults,
SupportsHMA,
)
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
@@ -234,6 +235,18 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
assert self.connector_worker is not None
return self.connector_worker.get_finished()
def get_transfer_results(
self, finished_req_ids: set[str]
) -> KVConnectorTransferResults:
assert self.connector_worker is not None
results = self.connector_worker.get_transfer_results()
if (
self.kv_transfer_config.kv_role == "kv_producer"
and self.connector_worker.pcp_rank > 0
):
results.finished_sending.clear()
return results
def get_block_ids_with_load_errors(self) -> set[int]:
"""Get block IDs that failed to load via NIXL."""
assert self.connector_worker is not None
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+3 -1
View File
@@ -305,7 +305,9 @@ class NixlConnectorMetadata(KVConnectorMetadata):
local_num_computed_blocks: tuple[int, ...] = (),
):
req = self._add_new_req(
local_block_ids, kv_transfer_params, local_num_computed_blocks
local_block_ids,
kv_transfer_params,
local_num_computed_blocks,
)
req.remote = RemoteMeta(
block_ids=kv_transfer_params["remote_block_ids"],
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+17 -11
View File
@@ -166,11 +166,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
dcp_active = self.dcp_size > 1 or remote_info.remote_dcp_size > 1
local_block_ids = meta.local_physical_block_ids
local_region_groups = getattr(self, "region_group_ids", [])
remote_region_groups = getattr(self, "dst_region_group_ids", {}).get(
engine_id, local_region_groups
)
if local_region_groups != remote_region_groups:
remote_region_groups = self.dst_region_group_ids[engine_id]
local_region_groups = self.region_group_ids or remote_region_groups
groups_differ = local_region_groups != remote_region_groups
if groups_differ:
if not self.use_mla or self._has_mamba:
raise NotImplementedError(
"Different NIXL cache-group layouts are only supported for "
@@ -292,6 +291,9 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
spec.remote_rank
]
# Once a read routes the request to failure reporting, the
# scheduler may free and reuse its blocks, so no sibling READs
# may be posted (and P must not be notified).
if not self._read_blocks(
read_spec=spec,
request_id=req_id,
@@ -329,6 +331,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
"""
Post a READ point-to-point xfer request from a single local worker to
a single remote worker.
Returns True when the read was posted (or was unnecessary), False
when the request was routed to failure reporting — the caller must
not post further transfers for it.
"""
assert self.transfer_topo is not None
remote_rank = read_spec.remote_rank
@@ -416,11 +422,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
dst_num_blocks=self.dst_num_blocks[dst_engine_id],
block_size_ratio=None,
physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
region_num_blocks=self.dst_region_num_blocks[dst_engine_id],
region_num_blocks=(self.dst_region_num_blocks.get(dst_engine_id) or None),
region_group_ids=(
list(range(self.num_regions))
if read_spec.block_ids_by_region
else self.dst_region_group_ids[dst_engine_id]
else (self.dst_region_group_ids.get(dst_engine_id) or None)
),
uses_region_group_mapping=(
self.num_regions > 1
@@ -433,11 +439,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
dst_num_blocks=self.dst_num_blocks[self.engine_id],
block_size_ratio=block_size_ratio,
physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
region_num_blocks=self.dst_region_num_blocks[self.engine_id],
region_num_blocks=(self.dst_region_num_blocks.get(self.engine_id) or None),
region_group_ids=(
list(range(self.num_regions))
if read_spec.block_ids_by_region
else self.region_group_ids
else (self.region_group_ids or None)
),
uses_region_group_mapping=(
self.num_regions > 1
@@ -537,11 +543,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
self._pending_recv_notifs.setdefault(request_id, []).append(
(notif_agent, notif_id)
)
for index, handle in enumerate(handles):
for i, handle in enumerate(handles):
try:
self.nixl_wrapper.transfer(handle)
except Exception:
for unstarted in handles[index:]:
for unstarted in handles[i:]:
self.nixl_wrapper.release_xfer_handle(unstarted)
raise
self._recving_transfers[request_id].append(handle)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+18 -5
View File
@@ -41,6 +41,9 @@ from typing import TYPE_CHECKING, Any
import msgspec
from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
KVConnectorTransferResults,
)
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
NixlBaseConnectorWorker,
)
@@ -145,7 +148,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
self._push_writer_thread.start()
logger.info("nixl-push-writer thread started (rank=%d)", self.tp_rank)
def shutdown(self):
def shutdown(self) -> None:
self._push_writer_stop.set()
# Unblock the writer if it's waiting in the no-active-state branch.
self._push_writer_wake.set()
@@ -522,7 +525,14 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
)
remote_block_ids = meta.remote.block_ids
local_block_ids = meta.local_physical_block_ids
num_groups = len(local_block_ids)
local_region_groups = self.region_group_ids
remote_region_groups = self.dst_region_group_ids[engine_id]
groups_differ = local_region_groups != remote_region_groups
if groups_differ:
raise NotImplementedError(
"NixlPushConnector does not support different producer and "
"consumer cache-group layouts"
)
# MLA latent is replicated across D's TP ranks: the tp-mapping
# collapses it to one rank (fine for reads), but push must WRITE every
@@ -537,6 +547,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
else:
write_ranks = list(plan.all_source_ranks)
num_groups = len(local_block_ids)
def group_ids(block_ids: BlockIds, rank: int) -> BlockIds:
return [
list(block_ids[g])
@@ -777,13 +789,14 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
self._reqs_to_send.pop(req_id, None)
return notified_req_ids
def get_finished(self) -> tuple[set[str], set[str]]:
def get_transfer_results(self) -> KVConnectorTransferResults:
# Engine main thread asking for completions: also wake the writer
# so it gets a chance to drain NIXL notifs (heartbeats, completion
# notifs, late PUSH_REGs) even if it had been parked.
self._push_writer_wake.set()
done_sending, done_recving = super().get_finished()
results = super().get_transfer_results()
done_sending = results.finished_sending
# ``_pop_done_transfers`` mutates ``_sending_transfers``; the
# writer thread also appends to it, so guard the pop.
@@ -803,4 +816,4 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
if done_sending:
self._push_writer_wake.set()
return done_sending, done_recving
return results
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+29 -2
View File
@@ -12,6 +12,7 @@ from vllm.v1.core.kv_cache_utils import (
from vllm.v1.kv_cache_interface import (
AttentionSpec,
FullAttentionSpec,
KVCacheGroupRole,
KVCacheSpec,
MLAAttentionSpec,
SlidingWindowMLASpec,
@@ -31,6 +32,16 @@ if TYPE_CHECKING:
from vllm.v1.kv_cache_interface import KVCacheConfig
def get_offloading_group_ids(kv_cache_config: "KVCacheConfig") -> tuple[int, ...]:
if kv_cache_config.hisparse_host_num_blocks is None:
return tuple(range(len(kv_cache_config.kv_cache_groups)))
return tuple(
group_id
for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
if group.role is KVCacheGroupRole.HISPARSE_INDEXER
)
def build_offloading_config(
vllm_config: "VllmConfig",
kv_cache_config: "KVCacheConfig",
@@ -43,15 +54,22 @@ def build_offloading_config(
engine_id = kv_transfer_config.engine_id
parallel_config = vllm_config.parallel_config
selected_groups = tuple(
(group_id, kv_cache_config.kv_cache_groups[group_id])
for group_id in get_offloading_group_ids(kv_cache_config)
)
if not selected_groups:
raise ValueError("KV offloading found no eligible cache groups.")
groups = tuple(
OffloadingGroupConfig(
group_id=group_id,
tokens_per_block=resolve_dcp_kv_block_size(
group.kv_cache_spec,
parallel_config.decode_context_parallel_size,
),
layer_names=tuple(group.layer_names),
)
for group in kv_cache_config.kv_cache_groups
for group_id, group in selected_groups
)
_, tokens_per_hash = resolve_kv_cache_block_sizes(kv_cache_config, vllm_config)
@@ -103,11 +121,20 @@ def build_offloading_config(
)
worker_kv_bytes_per_block = 0
if kv_cache_config.num_blocks > 0 and kv_cache_config.kv_cache_tensors:
all_groups_selected = len(selected_groups) == len(kv_cache_config.kv_cache_groups)
if (
all_groups_selected
and kv_cache_config.num_blocks > 0
and kv_cache_config.kv_cache_tensors
):
# Every KVCacheTensor describes placement within the same backing allocation,
# so its size is the total, not a per-tensor share.
total_gpu_kv_bytes = kv_cache_config.kv_cache_tensors[0].size
worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks
elif kv_cache_config.num_blocks > 0:
worker_kv_bytes_per_block = sum(
group.kv_cache_spec.page_size_bytes for _, group in selected_groups
)
single_group_spec = (
kv_cache_config.kv_cache_groups[0].kv_cache_spec
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+70 -33
View File
@@ -152,8 +152,9 @@ def resolve_mamba_align_size(
such groups agree on the same value.
"""
mamba_align_size: int | None = None
for idx, tokens_per_block in enumerate(spec.tokens_per_block):
kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec
for group in spec.config.groups:
tokens_per_block = group.tokens_per_block
kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec
if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode in (
"align",
"all",
@@ -188,8 +189,9 @@ class SchedulerOffloadConfig(NamedTuple):
# each segment can never serve a load hit. Relevant for hybrid
# architectures like DeepSeek V4 (MLA + SWA groups).
full_attn_tokens_per_chunk: set[int] = set()
for idx, tokens_per_block in enumerate(spec.tokens_per_block):
kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec
for group in spec.config.groups:
tokens_per_block = group.tokens_per_block
kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec
sw = get_sliding_window_size_in_chunks(
kv_spec, tokens_per_block * spec.blocks_per_chunk
)
@@ -214,9 +216,9 @@ class SchedulerOffloadConfig(NamedTuple):
retention_interval = vllm_config.cache_config.prefix_cache_retention_interval
eagle_groups = {
idx
for idx, g in enumerate(kv_cache_config.kv_cache_groups)
if g.is_eagle_group
group.group_id
for group in spec.config.groups
if kv_cache_config.kv_cache_groups[group.group_id].is_eagle_group
}
use_eagle_block_drop = (
@@ -249,8 +251,10 @@ class SchedulerOffloadConfig(NamedTuple):
)
kv_group_configs_list: list[GroupOffloadConfig] = []
for idx, tokens_per_block in enumerate(spec.tokens_per_block):
kv_cache_group = kv_cache_config.kv_cache_groups[idx]
for group in spec.config.groups:
group_id = group.group_id
tokens_per_block = group.tokens_per_block
kv_cache_group = kv_cache_config.kv_cache_groups[group_id]
kv_spec = kv_cache_group.kv_cache_spec
sw = get_sliding_window_size_in_chunks(
kv_spec, tokens_per_block * spec.blocks_per_chunk
@@ -261,7 +265,7 @@ class SchedulerOffloadConfig(NamedTuple):
)
kv_group_configs_list.append(
GroupOffloadConfig(
group_idx=idx,
group_idx=group_id,
tokens_per_block=tokens_per_block,
tokens_per_chunk=tokens_per_block * spec.blocks_per_chunk,
hashes_per_chunk=(
@@ -272,7 +276,7 @@ class SchedulerOffloadConfig(NamedTuple):
kv_cache_spec=kv_spec,
manager_cls=manager_cls,
kv_event_group_spec=get_offloading_event_group_spec(kv_cache_group),
is_eagle_group=idx in eagle_groups,
is_eagle_group=group_id in eagle_groups,
requires_cow_source=(
isinstance(kv_spec, MambaSpec)
and kv_spec.mamba_cache_mode == "align"
@@ -408,9 +412,10 @@ class RequestOffloadState:
if new_block_id_groups is None:
return
assert len(new_block_id_groups) == len(self.group_states)
for group_state, new_blocks in zip(self.group_states, new_block_id_groups):
group_state.block_ids.extend(new_blocks)
for group_config, group_state in zip(
self.config.kv_group_configs, self.group_states
):
group_state.block_ids.extend(new_block_id_groups[group_config.group_idx])
def storable_chunks(
self,
@@ -541,11 +546,11 @@ class OffloadingConnectorScheduler:
full_attention_groups: list[int] = []
sliding_window_groups: list[int] = []
for group_config in self.config.kv_group_configs:
for config_idx, group_config in enumerate(self.config.kv_group_configs):
if group_config.sliding_window_size_in_chunks is None:
full_attention_groups.append(group_config.group_idx)
full_attention_groups.append(config_idx)
else:
sliding_window_groups.append(group_config.group_idx)
sliding_window_groups.append(config_idx)
# sort sliding window groups by window size in decreasing order
def _sliding_window_sort_key(i: int) -> int:
@@ -737,7 +742,11 @@ class OffloadingConnectorScheduler:
req_status.req_context,
)
def _lookup_complete_chunks(self, req_status: RequestOffloadState) -> int | None:
def _lookup_complete_chunks(
self,
req_status: RequestOffloadState,
max_num_new_tokens: int | None = None,
) -> int | None:
"""
Find how many tokens beyond num_locally_computed_tokens can be loaded.
@@ -748,6 +757,10 @@ class OffloadingConnectorScheduler:
"""
num_computed_tokens = req_status.num_locally_computed_tokens
max_hit_size_tokens: int = req_status.req.num_tokens
if max_num_new_tokens is not None:
max_hit_size_tokens = min(
max_hit_size_tokens, num_computed_tokens + max_num_new_tokens
)
if self._sliding_window_groups:
# the last prompt token has to be recomputed to get the logprobs
# for sliding window attention, we must reduce by 1 to make sure
@@ -929,8 +942,12 @@ class OffloadingConnectorScheduler:
hash_idx = boundary_tokens // self.config.tokens_per_hash - 1
return make_offload_key(request.block_hashes[hash_idx], group_idx)
def _lookup(self, req_status: RequestOffloadState) -> int | None:
complete_hit = self._lookup_complete_chunks(req_status)
def _lookup(
self,
req_status: RequestOffloadState,
max_num_new_tokens: int | None = None,
) -> int | None:
complete_hit = self._lookup_complete_chunks(req_status, max_num_new_tokens)
req_status.partial_tail_boundary = None
if complete_hit is None or not self.config.supports_partial_tail:
return complete_hit
@@ -939,9 +956,10 @@ class OffloadingConnectorScheduler:
complete_boundary = local_tokens + complete_hit
tokens_per_hash = self.config.tokens_per_hash
block_end = complete_boundary + self._partial_tail_block_size
max_boundary = round_down(
min(req_status.req.num_prompt_tokens - 1, block_end - 1), tokens_per_hash
)
max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1)
if max_num_new_tokens is not None:
max_boundary = min(max_boundary, local_tokens + max_num_new_tokens)
max_boundary = round_down(max_boundary, tokens_per_hash)
if max_boundary <= complete_boundary:
return complete_hit
@@ -990,7 +1008,10 @@ class OffloadingConnectorScheduler:
self._req_status[request.request_id] = req_status
def get_num_new_matched_tokens(
self, request: Request, num_computed_tokens: int
self,
request: Request,
num_computed_tokens: int,
max_num_new_tokens: int | None = None,
) -> tuple[int | None, bool]:
"""
Get number of new tokens that can be loaded beyond the
@@ -1030,7 +1051,7 @@ class OffloadingConnectorScheduler:
num_hit_tokens = 0
else:
lookup_start = time.monotonic()
num_hit_tokens = self._lookup(req_status)
num_hit_tokens = self._lookup(req_status, max_num_new_tokens)
self._connector_stats.observe_histogram(
_ConnectorMetricName.LOOKUP_SYNC_DELAY,
time.monotonic() - lookup_start,
@@ -1065,11 +1086,11 @@ class OffloadingConnectorScheduler:
# per group
group_sizes: list[int] = []
block_indices: list[int] = []
for group_config, group_state, group_blocks in zip(
for group_config, group_state in zip(
self.config.kv_group_configs,
req_status.group_states,
blocks.blocks,
):
group_blocks = blocks.blocks[group_config.group_idx]
self._current_batch_allocated_block_ids.update(
block.block_id for block in group_blocks if block.block_id != 0
)
@@ -1184,7 +1205,8 @@ class OffloadingConnectorScheduler:
for grp_idx in self._sliding_window_groups
)
req_status.update_block_id_groups(new_block_id_groups)
for new_blocks in new_block_id_groups:
for group_config in self.config.kv_group_configs:
new_blocks = new_block_id_groups[group_config.group_idx]
for bid in new_blocks:
if bid != 0:
self._current_batch_allocated_block_ids.add(bid)
@@ -1218,6 +1240,10 @@ class OffloadingConnectorScheduler:
) -> dict[int, TransferJob]:
store_jobs: dict[int, TransferJob] = {}
num_groups = len(self.config.kv_group_configs)
config_idx_by_group = {
config.group_idx: idx
for idx, config in enumerate(self.config.kv_group_configs)
}
for req_id, entries in handoffs.items():
req_status = self._req_status.get(req_id)
if req_status is None:
@@ -1225,7 +1251,10 @@ class OffloadingConnectorScheduler:
req = req_status.req
max_boundary = self._calc_num_offloadable_tokens(req_status, req.num_tokens)
for group_idx, block_id, boundary in entries:
group_config = self.config.kv_group_configs[group_idx]
config_idx = config_idx_by_group.get(group_idx)
if config_idx is None:
continue
group_config = self.config.kv_group_configs[config_idx]
if (
block_id == 0
or boundary > max_boundary
@@ -1254,9 +1283,11 @@ class OffloadingConnectorScheduler:
fenced_block_ids=[block_id],
)
group_sizes = [0] * num_groups
group_sizes[group_idx] = 1
group_sizes[config_idx] = 1
block_indices = [0] * num_groups
block_indices[group_idx] = boundary // group_config.tokens_per_block - 1
block_indices[config_idx] = (
boundary // group_config.tokens_per_block - 1
)
store_jobs[job_id] = TransferJob(
req_id=req_id,
src_spec=GPULoadStoreSpec(
@@ -1300,6 +1331,12 @@ class OffloadingConnectorScheduler:
assert len(boundaries) == 1
boundary = boundaries.pop()
req = req_status.req
group_states = {
group.group_idx: state
for group, state in zip(
self.config.kv_group_configs, req_status.group_states
)
}
max_boundary = min(
req.num_prompt_tokens,
req_status.max_offload_tokens or req.num_prompt_tokens,
@@ -1315,7 +1352,7 @@ class OffloadingConnectorScheduler:
block_idx = boundary // self._partial_tail_block_size
if any(
group.group_idx not in self._cow_source_groups
and block_idx >= len(req_status.group_states[group.group_idx].block_ids)
and block_idx >= len(group_states[group.group_idx].block_ids)
for group in self.config.kv_group_configs
):
continue
@@ -1326,7 +1363,7 @@ class OffloadingConnectorScheduler:
block_ids = [
cow_blocks[group.group_idx]
if group.group_idx in self._cow_source_groups
else req_status.group_states[group.group_idx].block_ids[block_idx]
else group_states[group.group_idx].block_ids[block_idx]
for group in self.config.kv_group_configs
]
assert all(block_id != 0 for block_id in block_ids)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+16 -8
View File
@@ -14,6 +14,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
OffloadingWorkerMetadata,
ReqId,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
get_offloading_group_ids,
)
from vllm.logger import init_logger
from vllm.v1.kv_cache_interface import (
AttentionSpec,
@@ -65,7 +68,10 @@ class OffloadingConnectorWorker:
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
kv_cache_config = self.kv_cache_config
num_blocks = kv_cache_config.num_blocks
selected_group_ids = get_offloading_group_ids(kv_cache_config)
selected_groups = tuple(
kv_cache_config.kv_cache_groups[group_id] for group_id in selected_group_ids
)
mappings = derive_canonical_mappings(
self.vllm_config, kv_cache_config, kv_caches
)
@@ -76,7 +82,9 @@ class OffloadingConnectorWorker:
unpadded_page_size_bytes: dict[str, int] = {}
# layer_name -> size of page in bytes
page_size_bytes: dict[str, int] = {}
for kv_cache_group in kv_cache_config.kv_cache_groups:
for kv_cache_group in selected_groups:
assert not kv_cache_group.host_resident
num_blocks = kv_cache_config.num_blocks
group_layer_names = kv_cache_group.layer_names
group_kv_cache_spec = kv_cache_group.kv_cache_spec
if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs):
@@ -131,8 +139,11 @@ class OffloadingConnectorWorker:
),
None,
)
if packed_layer_name is not None:
if packed_layer_name is not None and len(selected_groups) == len(
kv_cache_config.kv_cache_groups
):
(tensor,) = tensors_per_block[packed_layer_name]
num_blocks = tensor.shape[0]
block_stride = tensor.stride(0)
packed_tensor = tensor.as_strided(
(num_blocks, block_stride),
@@ -142,10 +153,7 @@ class OffloadingConnectorWorker:
self._init_worker(
CanonicalKVCaches(
[CanonicalKVCacheTensor(packed_tensor, block_stride)],
[
[CanonicalKVCacheRef(0, block_stride)]
for _ in kv_cache_config.kv_cache_groups
],
[[CanonicalKVCacheRef(0, block_stride)] for _ in selected_groups],
)
)
return
@@ -192,7 +200,7 @@ class OffloadingConnectorWorker:
)
group_data_refs: list[list[CanonicalKVCacheRef]] = []
for kv_cache_group in kv_cache_config.kv_cache_groups:
for kv_cache_group in selected_groups:
group_refs: list[CanonicalKVCacheRef] = []
for layer_name in kv_cache_group.layer_names:
group_refs += block_data_refs[layer_name]
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+46 -1
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable
from functools import cached_property
from typing import Any
import torch
@@ -25,6 +26,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
build_offloading_config,
get_offloading_group_ids,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
OffloadingConnectorStats,
@@ -38,6 +40,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
)
from vllm.forward_context import ForwardContext
from vllm.v1.attention.backend import AttentionMetadata
from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.kv_cache_interface import KVCacheConfig
@@ -47,6 +50,25 @@ from vllm.v1.request import Request
class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
@cached_property
def _bounding_group_ids(self) -> tuple[int, ...]:
"""Prefix-cacheable groups this connector does not offload.
Never offer tokens for a range some group cannot cover: past such a
group's own cached prefix its KV would be left unwritten.
"""
offloaded = set(get_offloading_group_ids(self._kv_cache_config))
return tuple(
group_id
for group_id, group in enumerate(self._kv_cache_config.kv_cache_groups)
if group.kv_cache_spec.prefix_cacheable and group_id not in offloaded
)
@property
def scheduler(self) -> OffloadingConnectorScheduler:
assert self.connector_scheduler is not None
return self.connector_scheduler
@property
def requires_kv_delivery(self) -> bool:
# Runs as kv_both, but is a best-effort cache: a dropped save is just a
@@ -138,9 +160,30 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
) -> tuple[int | None, bool]:
assert self.connector_scheduler is not None
return self.connector_scheduler.get_num_new_matched_tokens(
request, num_computed_tokens
request,
num_computed_tokens,
max_num_new_tokens=self._max_loadable_tokens(request, num_computed_tokens),
)
def _max_loadable_tokens(
self, request: "Request", num_computed_tokens: int
) -> int | None:
"""How far past ``num_computed_tokens`` a load may reach, if bounded.
Bounded by the deepest prefix the groups this connector does not
offload already hold, since nothing refills them beyond it.
"""
if not self._bounding_group_ids:
return None
assert self._kv_cache_manager is not None
coordinator = self._kv_cache_manager.coordinator
assert isinstance(coordinator, HybridKVCacheCoordinator)
_, per_group_hits = coordinator.find_longest_cache_hit_per_group(
request.block_hashes, request.num_tokens - 1
)
bound = min(per_group_hits[group_id] for group_id in self._bounding_group_ids)
return max(0, bound - num_computed_tokens)
def update_state_after_alloc(
self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
):
@@ -185,6 +228,8 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
@classmethod
def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str | None:
if vllm_config.attention_config.hisparse_config is not None:
return "BLHNC"
return "LBHNC"
def reset_cache(self) -> bool | None:
vllm/distributed/kv_transfer/kv_transfer_state.py
+3 -5
View File
@@ -78,13 +78,11 @@ def ensure_kv_transfer_initialized(
global _KV_CONNECTOR_AGENT
if vllm_config.kv_transfer_config is None:
if _KV_CONNECTOR_AGENT is not None:
return
if (
vllm_config.kv_transfer_config.is_kv_transfer_instance
and _KV_CONNECTOR_AGENT is None
):
kv_transfer_config = vllm_config.kv_transfer_config
if kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance:
# NIXL P/D requires an interleave_size equal to block_size.
vllm_config.adjust_dcp_kv_cache_interleave_size(kv_cache_config)
_sync_engine_id_across_tp(vllm_config)
vllm/model_executor/layers/attention/mla_attention.py
+100 -31
View File
@@ -207,7 +207,7 @@ from collections.abc import Callable
from dataclasses import dataclass, replace
from enum import Enum
from math import lcm
from typing import ClassVar, Generic, TypeVar, cast
from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast
import numpy as np
import torch
@@ -284,6 +284,7 @@ from vllm.v1.attention.backend import (
CommonAttentionMetadata,
MLAAttentionImpl,
)
from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
from vllm.v1.attention.backends.mla.prefill import (
MLAPrefillBackend,
get_mla_prefill_backend,
@@ -309,6 +310,11 @@ from vllm.v1.kv_cache_interface import (
get_kv_quant_mode,
)
if TYPE_CHECKING:
from vllm.v1.attention.backends.mla.index_group import (
SparseMLAIndexGroupBuilder,
)
logger = init_logger(__name__)
_FP8_DTYPE = current_platform.fp8_dtype()
@@ -433,6 +439,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
use_sparse: bool = False,
indexer: object | None = None,
topk_indices_buffer: torch.Tensor | None = None,
index_group_builder: "SparseMLAIndexGroupBuilder | None" = None,
non_causal_multi_token_decode: bool = False,
sliding_window: int | None = None,
prefill_backend_cls: type[MLAPrefillBackend] | None = None,
@@ -538,9 +545,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):
# explicitly so backbone "skip" layers (indexer=None) still find it.
if use_sparse:
extra_impl_args["topk_indices_buffer"] = topk_indices_buffer
if index_group_builder is not None:
extra_impl_args["index_group_builder"] = index_group_builder
impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
self.impl = impl_cls( # type: ignore[assignment] # impl_cls always returns an MLAAttentionImpl subclass
impl = impl_cls(
num_heads=self.num_heads,
head_size=self.head_size,
scale=self.scale,
@@ -562,6 +571,13 @@ class MLAAttention(nn.Module, AttentionLayerBase):
indexer=indexer,
**extra_impl_args,
)
self.impl = impl # type: ignore[assignment]
index_group = getattr(impl, "index_group", None)
self.hisparse_cache = (
index_group.cache(cast(Any, impl).index_group_index)
if isinstance(index_group, HiSparseMLAIndexGroup)
else None
)
self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False)
# MLA reads this weight directly to build W_UK/W_UV, so a backend must
@@ -723,6 +739,61 @@ class MLAAttention(nn.Module, AttentionLayerBase):
)
return self._chunked_prefill_workspace_size
def update_kv_cache(
self,
kv_c_normed: torch.Tensor,
k_pe: torch.Tensor,
kv_cache: torch.Tensor,
slot_mapping: torch.Tensor | None,
attn_metadata: "MLACommonMetadata | None",
kv_cache_dtype: str,
k_scale: torch.Tensor,
) -> None:
cache = self.hisparse_cache
if slot_mapping is None or (cache is not None and cache.dummy_batch):
return
kv_c_normed, k_pe, slot_mapping = maybe_gather_mla_latent_cache_inputs(
kv_c_normed,
k_pe,
slot_mapping,
attn_metadata.num_decode_tokens if attn_metadata is not None else None,
self.use_pcp,
)
assert slot_mapping is not None
if cache is not None:
kv_cache, slot_mapping, num_rows = cache.write_target(
kv_c_normed.shape[0], slot_mapping.numel()
)
kv_c_normed = kv_c_normed[:num_rows]
k_pe = k_pe[:num_rows]
self.impl.do_kv_cache_update( # type: ignore[attr-defined]
kv_c_normed,
k_pe,
kv_cache,
slot_mapping,
kv_cache_dtype,
k_scale,
)
if cache is not None:
mirror_target = cache.mirror_write_target(kv_c_normed.shape[0])
if mirror_target is not None:
mirror_cache, mirror_slots = mirror_target
self.impl.do_kv_cache_update( # type: ignore[attr-defined]
kv_c_normed,
k_pe,
mirror_cache,
mirror_slots,
kv_cache_dtype,
k_scale,
)
def prepare_kv_cache_update(
self, attn_metadata: "MLACommonMetadata | None"
) -> None:
cache = self.hisparse_cache
if cache is not None and cache.runtime.is_group_leader:
cache.prepare_group_for_batch(attn_metadata)
def forward(
self,
q: torch.Tensor,
@@ -750,25 +821,18 @@ class MLAAttention(nn.Module, AttentionLayerBase):
f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
)
layer_slot_mapping = slot_mapping.get(self.layer_name)
kv_for_cache, kpe_for_cache, layer_slot_mapping = (
maybe_gather_mla_latent_cache_inputs(
kv_c_normed,
k_pe,
layer_slot_mapping,
attn_metadata.num_decode_tokens
if attn_metadata is not None
else None,
self.use_pcp,
)
)
self.impl.do_kv_cache_update( # type: ignore[attr-defined]
kv_for_cache,
kpe_for_cache,
self.prepare_kv_cache_update(attn_metadata)
self.update_kv_cache(
kv_c_normed,
k_pe,
self_kv_cache,
layer_slot_mapping,
attn_metadata,
self.kv_cache_dtype,
self._k_scale,
)
if self.hisparse_cache is not None:
self.hisparse_cache.finish_kv_update()
output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
self.forward_impl(
q,
@@ -1113,6 +1177,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
return output_padded
def _use_sparse_mha(self, attn_metadata: "MLACommonMetadata") -> bool:
if self.hisparse_cache is not None:
return False
prefill = attn_metadata.prefill
if prefill is None:
return False
@@ -1296,6 +1362,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
)
return MLAAttentionSpec(
**common_kwargs,
is_index_group_leader=self.indexer is not None,
non_causal_multi_token_decode=self.non_causal_multi_token_decode,
)
@@ -1349,19 +1416,14 @@ def unified_mla_kv_cache_update(
attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(
layer_name
)
attn_layer.prepare_kv_cache_update(attn_metadata)
if layer_slot_mapping is not None:
kv_c_normed, k_pe, layer_slot_mapping = maybe_gather_mla_latent_cache_inputs(
kv_c_normed,
k_pe,
layer_slot_mapping,
attn_metadata.num_decode_tokens if attn_metadata is not None else None,
attn_layer.use_pcp,
)
attn_layer.impl.do_kv_cache_update( # type: ignore[attr-defined]
attn_layer.update_kv_cache(
kv_c_normed,
k_pe,
kv_cache,
layer_slot_mapping,
attn_metadata,
kv_cache_dtype,
k_scale,
)
@@ -1409,6 +1471,8 @@ def unified_mla_attention_with_output(
del kv_cache_dummy_dep
layer_name = _resolve_layer_name(layer_name)
attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
if layer.hisparse_cache is not None:
layer.hisparse_cache.finish_kv_update()
layer.forward_impl(
q,
kv_c_normed,
@@ -1501,6 +1565,7 @@ class MLACommonBackend(AttentionBackend):
def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":
"""Per-token-head modes pack an inline fp32 scale pair after the
latent data (single-sided: ``head_size_v == 0`` for MLA)."""
spec = super().customize_spec(spec)
mode = spec.kv_quant_mode
if spec.state_content_bytes is not None or not mode.is_per_token_head:
return spec
@@ -2090,7 +2155,11 @@ def build_mla_chunked_context_metadata(
local_token_offset += num_local_tokens
seq_lens_cpu = _flat_int32(seq_lens_flat)
starts = _flat_int32(starts_flat).to(device, non_blocking=True)
starts_and_context_lens = _flat_int32(starts_flat + context_lens).to(
device, non_blocking=True
)
starts = starts_and_context_lens[: len(starts_flat)]
context_lens_gpu = starts_and_context_lens[len(starts_flat) :]
cu_seq_lens = _flat_int32(cu_seq_lens_flat).to(device, non_blocking=True)
cu_seqlens_q = _flat_int32(cu_seqlens_q_flat).to(device, non_blocking=True)
token_to_seq = _flat_int32(np.concatenate(token_to_seq_parts)).to(
@@ -2145,7 +2214,7 @@ def build_mla_chunked_context_metadata(
chunks.append(chunk)
return MLACommonPrefillMetadata.ChunkedContextMetadata(
context_lens=context_lens_cpu.to(device, non_blocking=True),
context_lens=context_lens_gpu,
workspace=chunked_prefill_workspace,
chunks=chunks,
context_lens_list=context_lens,
@@ -2835,12 +2904,12 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
output_lse = None
for chunk in chunked_context.chunks:
toks = chunk.num_context_tokens
block_table = prefill_metadata.block_table[chunk.request_slice]
chunk_block_table = prefill_metadata.block_table[chunk.request_slice]
if self.kv_cache_dtype == "fp8_ds_mla":
ops.cp_gather_and_upconvert_fp8_kv_cache(
src_cache=kv_c_and_k_pe_cache,
dst=workspace[:toks],
block_table=block_table,
block_table=chunk_block_table,
workspace_starts=chunk.cu_seq_lens,
batch_size=chunk.num_requests,
seq_starts=chunk.starts,
@@ -2853,7 +2922,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
ops.gather_mla_context_cache_cpu(
src_cache=kv_c_and_k_pe_cache,
dst=workspace[:toks],
block_table=block_table,
block_table=chunk_block_table,
starts=chunk.starts,
cu_seq_lens=chunk.cu_seq_lens,
)
@@ -2861,7 +2930,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
ops.gather_and_maybe_dequant_cache(
src_cache=kv_c_and_k_pe_cache,
dst=workspace,
block_table=block_table,
block_table=chunk_block_table,
cu_seq_lens=chunk.cu_seq_lens,
token_to_seq=chunk.token_to_seq,
num_tokens=toks,
@@ -2874,7 +2943,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
ops.cp_gather_cache(
src_cache=kv_c_and_k_pe_cache,
dst=workspace[:toks],
block_table=block_table,
block_table=chunk_block_table,
cu_seq_lens=chunk.cu_seq_lens,
batch_size=chunk.num_requests,
seq_starts=chunk.starts,
vllm/model_executor/layers/attention/sparse_mla_attention.py
+166 -25
View File
@@ -3,13 +3,15 @@
"""Shared MHA implementation and metadata builder for sparse MLA backends."""
import math
from dataclasses import dataclass
from shutil import which
from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast
from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar
import numpy as np
import torch
from vllm import _custom_ops as ops
from vllm.config import get_current_vllm_config
from vllm.distributed import (
get_dcp_group,
get_tensor_model_parallel_world_size,
@@ -17,6 +19,7 @@ from vllm.distributed import (
from vllm.logger import init_logger
from vllm.model_executor.layers.attention.mla_attention import (
MLACommonBaseImpl,
MLACommonDecodeMetadata,
MLACommonMetadata,
MLACommonPrefillMetadata,
accumulate_mla_context_chunk,
@@ -33,11 +36,19 @@ from vllm.utils.torch_utils import (
is_quantized_kv_cache,
np_to_pinned_tensor,
)
from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder
from vllm.v1.attention.backend import AttentionMetadataBuilder
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
from vllm.v1.attention.backends.mla.index_group import (
SparseMLAIndexGroup,
SparseMLAIndexGroupBuilder,
)
from vllm.v1.attention.backends.utils import split_decodes_and_prefills
from vllm.v1.attention.ops.dcp import MLADCPManager
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
from vllm.v1.hisparse.runtime import (
HiSparsePrefillStagingPlan,
build_hisparse_prefill_staging_plan,
)
if TYPE_CHECKING:
from vllm.config import VllmConfig
@@ -47,8 +58,6 @@ if TYPE_CHECKING:
logger = init_logger(__name__)
T = TypeVar("T", bound=AttentionMetadata)
GLOBAL_TOPK_MASK_MAX_BYTES = 128 * 1024 * 1024 # 128 MiB
@@ -116,9 +125,45 @@ def _is_masked_mha_available(
return fa_version == 4 and not is_quantized_kv_cache(kv_cache_dtype)
def _use_dense_mha_prefill(
vllm_config: "VllmConfig", prefill_max_seq_len: int, topk_tokens: int
) -> bool:
attention_config = vllm_config.attention_config
return (
attention_config.hisparse_config is None
and not attention_config.sparse_mla_force_mqa
and prefill_max_seq_len <= topk_tokens
)
@dataclass
class SparseMLAPrefillMetadata(MLACommonPrefillMetadata):
host_staging_plan: HiSparsePrefillStagingPlan | None = None
@dataclass(kw_only=True)
class SparseMLACommonMetadata(MLACommonMetadata[MLACommonDecodeMetadata]):
block_table: torch.Tensor
req_id_per_token: torch.Tensor
seq_lens: torch.Tensor | None = None
block_size: int = 64
topk_tokens: int = 2048
num_decodes: int = 0
num_prefills: int = 0
num_decode_tokens: int = 0
decode_max_query_len: int = 0
prefill_max_seq_len: int = 0
prefill: SparseMLAPrefillMetadata | None = None
cp_kv_cache_interleave_size: int = 1
T = TypeVar("T", bound=SparseMLACommonMetadata)
class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
metadata_cls: type[T]
require_uniform_decodes: ClassVar[bool] = False
hisparse_supports_multi_token_decode: ClassVar[bool] = False
def __init__(
self,
@@ -138,6 +183,12 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
dtype=torch.int32,
device=device,
)
self._context_lens_cpu = torch.empty(
vllm_config.scheduler_config.max_num_seqs,
dtype=torch.int32,
device="cpu",
pin_memory=PIN_MEMORY,
)
parallel_config = vllm_config.parallel_config
self.use_pcp = parallel_config.prefill_context_parallel_size > 1
try:
@@ -197,6 +248,22 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
layer_prefill_backend.clone() if layer_prefill_backend is not None else None
)
def _init_reorder_batch_threshold(
self,
reorder_batch_threshold: int | None = 1,
supports_spec_as_decode: bool = False,
supports_dcp_with_varlen: bool = False,
) -> None:
if self.vllm_config.attention_config.hisparse_config is not None:
reorder_batch_threshold = 1
if not self.hisparse_supports_multi_token_decode:
supports_spec_as_decode = False
super()._init_reorder_batch_threshold(
reorder_batch_threshold,
supports_spec_as_decode,
supports_dcp_with_varlen,
)
@staticmethod
def determine_chunked_prefill_workspace_size(vllm_config: "VllmConfig") -> int:
scheduler_config = vllm_config.scheduler_config
@@ -245,9 +312,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
assert seq_lens_cpu is not None
context_lens_cpu = torch.empty(
num_prefills, dtype=seq_lens_cpu.dtype, pin_memory=PIN_MEMORY
)
context_lens_cpu = self._context_lens_cpu[:num_prefills]
torch.subtract(
seq_lens_cpu[num_decodes : num_decodes + num_prefills],
prefill_query_lens_cpu,
@@ -289,34 +354,66 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
prefill_max_query_len,
prefill_query_lens_cpu,
) = self._build_prefill_fields(common_attn_metadata, num_decodes, num_prefills)
decode_max_query_len = 0
if num_decodes:
if common_attn_metadata.max_logits_per_req is not None:
decode_max_query_len = common_attn_metadata.max_logits_per_req
elif num_prefills == 0:
decode_max_query_len = common_attn_metadata.max_query_len
else:
decode_query_lens_cpu = torch.diff(
common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
)
decode_max_query_len = int(decode_query_lens_cpu.max().item())
prefill_max_seq_len = 0
prefill: MLACommonPrefillMetadata | None = None
prefill: SparseMLAPrefillMetadata | None = None
if num_prefills > 0 and self._prefill_backend is not None:
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
assert seq_lens_cpu is not None
prefill_max_seq_len = int(
seq_lens_cpu[num_decodes : num_decodes + num_prefills].max().item()
)
prefill = MLACommonPrefillMetadata(
block_table=common_attn_metadata.block_table_tensor[num_decodes:, ...],
block_table = common_attn_metadata.block_table_tensor[num_decodes:, ...]
chunked_context = self._build_chunked_context_fields(
common_attn_metadata,
num_decodes,
num_prefills,
prefill_query_lens_cpu,
)
staging_plan = None
if self.vllm_config.attention_config.hisparse_config is not None:
prefill_seq_lens_cpu = seq_lens_cpu[
num_decodes : num_decodes + num_prefills
]
staging_block_capacity = int(
(
(prefill_seq_lens_cpu + self.kv_cache_spec.block_size - 1)
// self.kv_cache_spec.block_size
).sum()
)
staging_plan = build_hisparse_prefill_staging_plan(
block_table,
common_attn_metadata.seq_lens[num_decodes:],
self.kv_cache_spec.block_size,
staging_block_capacity,
)
if chunked_context is not None:
block_table = staging_plan.block_table
prefill = SparseMLAPrefillMetadata(
block_table=block_table,
query_start_loc=prefill_query_start_loc,
max_query_len=prefill_max_query_len,
query_lens_cpu=prefill_query_lens_cpu,
chunked_context=self._build_chunked_context_fields(
common_attn_metadata,
num_decodes,
num_prefills,
prefill_query_lens_cpu,
),
chunked_context=chunked_context,
q_data_type=self.model_config.dtype,
output_dtype=self.model_config.dtype,
prefill_backend=self._prefill_backend,
use_dense_mha=(
prefill_max_seq_len <= self.topk_tokens
and not self.vllm_config.attention_config.sparse_mla_force_mqa
use_dense_mha=_use_dense_mha_prefill(
self.vllm_config, prefill_max_seq_len, self.topk_tokens
),
topk_mask_workspace=self.topk_mask_workspace,
host_staging_plan=staging_plan,
)
self._prefill_backend.prepare_metadata(prefill)
@@ -335,6 +432,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
num_decodes=num_decodes,
num_prefills=num_prefills,
num_decode_tokens=num_decode_tokens,
decode_max_query_len=decode_max_query_len,
prefill_max_seq_len=prefill_max_seq_len,
prefill=prefill,
cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
@@ -560,6 +658,7 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
kv_b_proj: "ColumnParallelLinear",
indexer: object | None = None,
topk_indices_buffer: torch.Tensor | None = None,
index_group_builder: SparseMLAIndexGroupBuilder | None = None,
q_pad_num_heads: int | None = None,
) -> None:
super().__init__(
@@ -577,6 +676,21 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
)
self.init_topk_indices_buffer(indexer, topk_indices_buffer)
self.index_group: SparseMLAIndexGroup | None = None
self.index_group_index = 0
if index_group_builder is None and self.topk_indices_buffer is not None:
index_group_builder = SparseMLAIndexGroupBuilder(self.topk_indices_buffer)
if index_group_builder is not None:
vllm_config = get_current_vllm_config()
self.index_group, self.index_group_index = (
index_group_builder.register_layer(
indexer is not None,
vllm_config,
head_size=head_size,
kv_cache_dtype=kv_cache_dtype,
)
)
self._use_flashinfer_concat_mla_k = (
has_flashinfer()
and which("ninja") is not None
@@ -593,6 +707,33 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
kv_cache_dtype=kv_cache_dtype,
)
def record_logical_topk_ready(self) -> None:
if self.index_group is not None:
self.index_group.set_logical_topk_ready(self.index_group_index)
def prepare_for_batch(self, attn_metadata: T | None) -> None:
if self.index_group is not None:
self.index_group.prepare_for_batch(self.index_group_index, attn_metadata)
def _convert_logical_to_physical_topk(
self,
logical_topk_indices: torch.Tensor,
attn_metadata: Any,
*,
block_stride_rows: int | None,
return_valid_counts: bool,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
group = self.index_group
assert group is not None
assert self.dcp_world_size == 1
return group.convert_logical_to_physical_topk(
self.index_group_index,
logical_topk_indices,
attn_metadata,
block_stride_rows=block_stride_rows,
return_valid_counts=return_valid_counts,
)
@staticmethod
def masked_mha_workspace_fits(prefill: MLACommonPrefillMetadata) -> bool:
"""Whether this prefill batch's top-k masks fit the workspace."""
@@ -853,8 +994,10 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
output: torch.Tensor,
output_scale: torch.Tensor | None = None,
) -> None:
prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined]
topk_tokens = attn_metadata.topk_tokens # type: ignore[attr-defined]
prefill_metadata = attn_metadata.prefill
assert prefill_metadata is not None
prefill_max_seq_len = attn_metadata.prefill_max_seq_len
topk_tokens = attn_metadata.topk_tokens
force_dense = getattr(self, "_sparse_mla_force_dense_mha", False)
force_masked = getattr(self, "_sparse_mla_force_masked_mha", False)
if force_dense or (prefill_max_seq_len <= topk_tokens and not force_masked):
@@ -863,7 +1006,7 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
kv_c_normed,
k_pe,
kv_c_and_k_pe_cache,
cast(MLACommonMetadata, attn_metadata),
attn_metadata,
k_scale,
output,
output_scale,
@@ -871,13 +1014,11 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
assert output_scale is None
assert self.masked_mha_available
prefill_metadata = attn_metadata.prefill # type: ignore[attr-defined]
assert prefill_metadata is not None
assert prefill_metadata.query_lens_cpu is not None
assert self.topk_indices_buffer is not None
q_lens = prefill_metadata.query_lens_cpu.tolist()
num_decode_tokens = attn_metadata.num_decode_tokens # type: ignore[attr-defined]
num_decode_tokens = attn_metadata.num_decode_tokens
topk_all = self.topk_indices_buffer[
num_decode_tokens : num_decode_tokens + q.shape[0]
]
vllm/model_executor/layers/mla.py
+10
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
from typing import TYPE_CHECKING
import torch
@@ -11,6 +12,11 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.models.common.ops import fused_q_kv_rmsnorm
from vllm.platforms import current_platform
if TYPE_CHECKING:
from vllm.v1.attention.backends.mla.index_group import (
SparseMLAIndexGroupBuilder,
)
@dataclass
class MLAModules:
@@ -30,6 +36,7 @@ class MLAModules:
topk_indices_buffer: torch.Tensor | None
indexer_rotary_emb: torch.nn.Module | None = None
g_proj: torch.nn.Module | None = None
index_group_builder: "SparseMLAIndexGroupBuilder | None" = None
# --8<-- [start:multi_head_latent_attention]
@@ -128,6 +135,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
use_sparse=self.is_sparse,
indexer=self.indexer,
topk_indices_buffer=mla_modules.topk_indices_buffer,
index_group_builder=mla_modules.index_group_builder,
non_causal_multi_token_decode=non_causal_multi_token_decode,
)
indexer_op = getattr(self.indexer, "indexer_op", None)
@@ -223,6 +231,8 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
if self.indexer and self.is_sparse and not self.skip_topk:
self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb)
if self.is_sparse:
self.mla_attn.impl.record_logical_topk_ready() # type: ignore[attr-defined]
if llama_4_scaling is not None:
q *= llama_4_scaling
vllm/model_executor/layers/sparse_attn_indexer.py
+1 -1
View File
@@ -425,7 +425,6 @@ def sparse_attn_indexer(
quant_block_size,
scale_fmt,
)
# The indexer and main MLA may classify the same short extend differently
# because they use independent decode thresholds. Only the main MLA route
# can determine whether the top-k indices will be consumed.
@@ -955,6 +954,7 @@ class SparseAttnIndexer(CustomOp):
self.dcp_rank,
self.dcp_world_size,
self.cp_kv_cache_interleave_size,
False,
candidate_blocks=self.candidate_blocks,
candidate_block_size=self.candidate_block_size,
candidate_write=self.candidate_write,
vllm/model_executor/models/deepseek_v2.py
+24 -1
View File
@@ -103,10 +103,14 @@ from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm.utils.torch_utils import direct_register_custom_op
from vllm.v1.attention.backend import AttentionBackend, AttentionType
from vllm.v1.attention.backends.mla.index_group import (
SparseMLAIndexGroupBuilder,
get_sparse_mla_index_group_max_rows,
)
from vllm.v1.attention.backends.mla.indexer import (
DeepseekV32IndexerBackend,
)
from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, SparseCacheRole
from .interfaces import (
MixtureOfExperts,
@@ -656,6 +660,7 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
num_kv_heads=1,
head_size=self.head_dim,
dtype=self.dtype,
cache_role=SparseCacheRole.INDEXER,
) # Only has one vector instead of K + V
def forward(self): ...
@@ -1012,6 +1017,7 @@ class DeepseekV2MLAAttention(nn.Module):
quant_config: QuantizationConfig | None = None,
prefix: str = "",
topk_indices_buffer: torch.Tensor | None = None,
index_group_builder: SparseMLAIndexGroupBuilder | None = None,
input_size: int | None = None,
reduce_results: bool = True,
non_causal_multi_token_decode: bool = False,
@@ -1199,6 +1205,7 @@ class DeepseekV2MLAAttention(nn.Module):
indexer_rotary_emb=self.indexer_rope_emb,
is_sparse=self.is_v32,
topk_indices_buffer=topk_indices_buffer,
index_group_builder=index_group_builder,
)
self.mla_attn = MultiHeadLatentAttentionWrapper(
@@ -1242,6 +1249,7 @@ class DeepseekV2DecoderLayer(nn.Module):
prefix: str,
config: DeepseekV2Config | None = None,
topk_indices_buffer: torch.Tensor | None = None,
index_group_builder: SparseMLAIndexGroupBuilder | None = None,
) -> None:
super().__init__()
@@ -1290,6 +1298,11 @@ class DeepseekV2DecoderLayer(nn.Module):
and parallel_config.pipeline_parallel_size == 1
and is_moe_layer
)
attn_kwargs = (
{"index_group_builder": index_group_builder}
if attn_cls is DeepseekV2MLAAttention
else {}
)
self.self_attn = attn_cls(
vllm_config=vllm_config,
config=config,
@@ -1306,6 +1319,7 @@ class DeepseekV2DecoderLayer(nn.Module):
prefix=f"{prefix}.self_attn",
topk_indices_buffer=topk_indices_buffer,
reduce_results=not self.use_sequence_parallel_moe,
**attn_kwargs,
)
if is_moe_layer:
@@ -1429,6 +1443,14 @@ class DeepseekV2Model(nn.Module):
)
else:
topk_indices_buffer = None
index_group_builder = (
SparseMLAIndexGroupBuilder(
topk_indices_buffer,
get_sparse_mla_index_group_max_rows(vllm_config),
)
if topk_indices_buffer is not None
else None
)
if get_pp_group().is_first_rank:
self.embed_tokens = VocabParallelEmbedding(
@@ -1445,6 +1467,7 @@ class DeepseekV2Model(nn.Module):
vllm_config=vllm_config,
prefix=prefix,
topk_indices_buffer=topk_indices_buffer,
index_group_builder=index_group_builder,
),
prefix=f"{prefix}.layers",
)
vllm/model_executor/models/glm4_moe_lite.py
+20
View File
@@ -67,6 +67,10 @@ from vllm.model_executor.models.glm4_moe import (
)
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm.v1.attention.backends.mla.index_group import (
SparseMLAIndexGroupBuilder,
get_sparse_mla_index_group_max_rows,
)
from .interfaces import SupportsLoRA, SupportsPP
from .utils import (
@@ -108,6 +112,7 @@ class Glm4MoeLiteDecoderLayer(nn.Module):
prefix: str,
config: "Glm4MoeLiteConfig | None" = None,
topk_indices_buffer: torch.Tensor | None = None,
index_group_builder: SparseMLAIndexGroupBuilder | None = None,
) -> None:
super().__init__()
@@ -136,6 +141,11 @@ class Glm4MoeLiteDecoderLayer(nn.Module):
attn_cls = Glm4MoeLiteMLAAttention
else:
attn_cls = Glm4MoeLiteAttention
attn_kwargs: dict[str, typing.Any] = (
{"index_group_builder": index_group_builder}
if attn_cls is Glm4MoeLiteMLAAttention
else {}
)
self.self_attn = attn_cls(
vllm_config=vllm_config,
@@ -152,6 +162,7 @@ class Glm4MoeLiteDecoderLayer(nn.Module):
quant_config=quant_config,
prefix=f"{prefix}.self_attn",
topk_indices_buffer=topk_indices_buffer,
**attn_kwargs,
)
if (
@@ -234,6 +245,14 @@ class Glm4MoeLiteModel(nn.Module):
)
else:
topk_indices_buffer = None
index_group_builder = (
SparseMLAIndexGroupBuilder(
topk_indices_buffer,
get_sparse_mla_index_group_max_rows(vllm_config),
)
if topk_indices_buffer is not None
else None
)
if get_pp_group().is_first_rank:
self.embed_tokens = VocabParallelEmbedding(
@@ -252,6 +271,7 @@ class Glm4MoeLiteModel(nn.Module):
config=config,
prefix=prefix,
topk_indices_buffer=topk_indices_buffer,
index_group_builder=index_group_builder,
),
prefix=f"{prefix}.layers",
)
vllm/model_executor/models/glm4_moe_lite_mtp.py
+13
View File
@@ -51,6 +51,10 @@ from vllm.model_executor.model_loader.weight_utils import (
)
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm.v1.attention.backends.mla.index_group import (
SparseMLAIndexGroupBuilder,
get_sparse_mla_index_group_max_rows,
)
from .glm4_moe_lite import (
Glm4MixtureOfExperts,
@@ -108,6 +112,14 @@ class Glm4MoeLiteMultiTokenPredictorLayer(nn.Module):
)
else:
topk_indices_buffer = None
index_group_builder = (
SparseMLAIndexGroupBuilder(
topk_indices_buffer,
get_sparse_mla_index_group_max_rows(vllm_config),
)
if topk_indices_buffer is not None
else None
)
self.shared_head = SharedHead(
config=config, prefix=prefix, quant_config=quant_config
@@ -117,6 +129,7 @@ class Glm4MoeLiteMultiTokenPredictorLayer(nn.Module):
prefix=prefix,
config=self.config,
topk_indices_buffer=topk_indices_buffer,
index_group_builder=index_group_builder,
)
def forward(
vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py
+29
View File
@@ -47,6 +47,35 @@ _FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS = {
_SPARSE_MLA_MIXED_WARMUP_TOKENS = 16
def autotune_hisparse_flashinfer_attention(runner: "GPUModelRunner") -> None:
"""Autotune each HiSparse FlashInfer sparse-MLA configuration."""
from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
FlashInferMLASparseImpl,
)
tuned: set[tuple[object, ...]] = set()
for layer in runner.vllm_config.compilation_config.static_forward_context.values():
impl = getattr(layer, "impl", None)
if not isinstance(impl, FlashInferMLASparseImpl):
continue
if getattr(layer, "hisparse_cache", None) is None:
continue
if impl.topk_indices_buffer is None:
continue
key = (
impl.kv_cache_dtype,
impl.num_heads,
impl.qk_nope_head_dim,
impl.qk_rope_head_dim,
impl.kv_lora_rank,
impl.topk_indices_buffer.shape[1],
)
if key in tuned:
continue
impl.autotune_hisparse_decode(layer)
tuned.add(key)
def _attention_backend_name(backend: object) -> str | None:
get_name = getattr(backend, "get_name", None)
if get_name is None:
vllm/model_executor/warmup/kernel_warmup.py
+14 -2
View File
@@ -25,6 +25,7 @@ from vllm.model_executor.warmup.flashinfer_autotune_cache import (
write_flashinfer_autotune_cache,
)
from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
autotune_hisparse_flashinfer_attention,
deepseek_v4_sparse_mla_attention_warmup,
flashinfer_sparse_mla_decode_autotune_warmup,
)
@@ -353,9 +354,12 @@ def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, ..
return tuple(dict.fromkeys(token_counts))
def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None:
def _run_flashinfer_autotune_dummy_runs(
runner: "GPUModelRunner", *, skip_attn: bool = False
) -> None:
import vllm.utils.flashinfer as fi_utils
dummy_run_kwargs = {"skip_attn": True} if skip_attn else {}
for num_tokens in _flashinfer_autotune_token_counts(runner):
tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens)
logger.info(
@@ -369,6 +373,7 @@ def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None:
skip_eplb=True,
is_profile=True,
randomize_inputs=True,
**dummy_run_kwargs,
)
@@ -431,7 +436,14 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
torch.inference_mode(),
fi_utils.autotune(tune_mode=True, **autotune_kwargs),
):
_run_flashinfer_autotune_dummy_runs(runner)
hisparse_enabled = (
runner.vllm_config.attention_config.hisparse_config is not None
)
if hisparse_enabled:
# HiSparse hot-buffer attention is bounded by decode batch
# size, not the prefill-sized batch used for the full model.
autotune_hisparse_flashinfer_attention(runner)
_run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled)
replayssm_autotune_warmup(runner)
_autotune_kimi_k3_kda_qkvg(runner.get_model())
finally:
vllm/models/deepseek_v32/attention.py
+36 -3
View File
@@ -42,6 +42,9 @@ from vllm.v1.attention.ops.pcp import (
if TYPE_CHECKING:
from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
from vllm.v1.attention.backends.mla.index_group import (
SparseMLAIndexGroupBuilder,
)
class DeepseekV32Indexer(nn.Module):
@@ -126,6 +129,7 @@ class DeepseekV32Attention(MLAAttention):
prefix: str,
topk_indices_buffer: torch.Tensor | None = None,
attn_backend: "type | None" = None,
index_group_builder: "SparseMLAIndexGroupBuilder | None" = None,
) -> None:
quant_config = vllm_config.quant_config
cache_config = vllm_config.cache_config
@@ -209,6 +213,7 @@ class DeepseekV32Attention(MLAAttention):
use_sparse=True,
indexer=indexer,
topk_indices_buffer=topk_indices_buffer,
index_group_builder=index_group_builder,
attn_backend=attn_backend,
)
@@ -304,6 +309,14 @@ class DeepseekV32Attention(MLAAttention):
slot_mapping = forward_context.slot_mapping
assert isinstance(slot_mapping, dict)
mla_slot = slot_mapping.get(self.layer_name)
indexer_slot = (
slot_mapping.get(self.indexer.k_cache.prefix)
if self.indexer is not None
else None
)
hisparse_cache = self.hisparse_cache
layer_attn_metadata, _, _, _ = get_attention_context(self.layer_name)
self.impl.prepare_for_batch(layer_attn_metadata)
if self.indexer is not None and not self.skip_topk:
has_indexer = True
@@ -331,12 +344,17 @@ class DeepseekV32Attention(MLAAttention):
mla_k_scale = None
indexer_k_cache = None
mla_slot = None
indexer_slot = None
else:
mla_kv_cache = self.kv_cache
mla_kv_cache = None if hisparse_cache is not None else self.kv_cache
mla_k_scale = self._k_scale
kv_c_out = torch.empty_like(kv_c)
k_pe_out = torch.empty_like(k_pe)
if self.use_pcp or hisparse_cache is None:
kv_c_out = torch.empty_like(kv_c)
k_pe_out = torch.empty_like(k_pe)
else:
kv_c_out = torch.empty_like(kv_c)
k_pe_out = torch.empty_like(k_pe)
q_c = fused_norm_rope(
positions,
q_c,
@@ -354,6 +372,7 @@ class DeepseekV32Attention(MLAAttention):
indexer_k_rope_cos_sin_cache,
self.topk_indices_buffer,
slot_mapping=mla_slot,
indexer_slot_mapping=indexer_slot,
indexer_k_cache=indexer_k_cache,
mla_kv_cache=mla_kv_cache,
mla_kv_cache_dtype=self.kv_cache_dtype,
@@ -365,6 +384,18 @@ class DeepseekV32Attention(MLAAttention):
index_k_out=index_k_out,
)
if hisparse_cache is not None and mla_slot is not None:
assert kv_c_out is not None and k_pe_out is not None
self.update_kv_cache(
kv_c_out,
k_pe_out,
self.kv_cache,
mla_slot,
layer_attn_metadata,
self.kv_cache_dtype,
self._k_scale,
)
q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)
q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
ql_nope = torch.bmm(q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1)
@@ -459,6 +490,7 @@ class DeepseekV32Attention(MLAAttention):
),
skip_topk_buffer_clear=True,
)
self.impl.record_logical_topk_ready() # type: ignore[attr-defined]
attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context(
self.layer_name
@@ -467,6 +499,7 @@ class DeepseekV32Attention(MLAAttention):
output.zero_()
return
attn_metadata = cast("MLACommonMetadata", attn_metadata)
self.impl.prepare_for_batch(attn_metadata)
if self.use_pcp:
assert kv_c is not None and k_pe is not None
vllm/models/deepseek_v32/common/kernels.py
+12 -7
View File
@@ -68,7 +68,7 @@ def _fp8_quant_and_cache_write(
kv_cache_ptr,
kv_cache_scale_ptr,
cache_block_size,
cache_stride,
cache_block_stride,
offsets,
HEAD_DIM: tl.constexpr,
):
@@ -76,7 +76,7 @@ def _fp8_quant_and_cache_write(
block_idx = slot_idx // cache_block_size
block_offset = slot_idx % cache_block_size
block_start = block_idx * cache_block_size * cache_stride
block_start = block_idx * cache_block_stride
tl.store(
kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets,
@@ -129,8 +129,9 @@ def _fused_norm_rope_kernel(
index_k_out_ptr,
index_k_out_stride,
INDEX_K_HALF_ROT_DIM: tl.constexpr,
# Cache params (shared by indexer K and MLA)
# Cache params
slot_mapping_ptr,
indexer_slot_mapping_ptr,
# Index K FP8 cache
indexer_cache_ptr,
indexer_cache_scale_ptr,
@@ -400,8 +401,8 @@ def _fused_norm_rope_kernel(
)
# PCP inserts index K after gathering; other paths write it directly.
if indexer_cache_ptr is not None and slot_mapping_ptr is not None:
slot_idx = tl.load(slot_mapping_ptr + tok_idx)
if indexer_cache_ptr is not None and indexer_slot_mapping_ptr is not None:
slot_idx = tl.load(indexer_slot_mapping_ptr + tok_idx)
_fp8_quant_and_cache_write(
result,
index_k_mask,
@@ -431,8 +432,9 @@ def fused_norm_rope(
index_k_layer_norm_eps: float,
index_k_rope_cos_sin_cache: torch.Tensor | None,
topk_indices_buffer: torch.Tensor,
# Cache params for fused writes (single slot_mapping for both caches)
# Cache params for fused writes
slot_mapping: torch.Tensor | None = None,
indexer_slot_mapping: torch.Tensor | None = None,
indexer_k_cache: torch.Tensor | None = None,
mla_kv_cache: torch.Tensor | None = None,
mla_kv_cache_dtype: str = "auto",
@@ -476,9 +478,11 @@ def fused_norm_rope(
# --- Indexer K cache setup ---
if indexer_k_cache is not None:
assert slot_mapping is not None
if indexer_slot_mapping is None:
indexer_slot_mapping = slot_mapping
idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32)
idx_cache_block_size = indexer_k_cache.shape[1]
idx_cache_stride = indexer_k_cache.shape[2]
idx_cache_stride = indexer_k_cache.stride(0)
if indexer_k_cache.dtype == torch.uint8:
indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn)
else:
@@ -578,6 +582,7 @@ def fused_norm_rope(
index_k_rope_cos_sin_cache.shape[-1] // 2,
# Cache params
slot_mapping,
indexer_slot_mapping,
indexer_k_cache,
idx_cache_scale_view,
idx_cache_block_size,
vllm/models/deepseek_v32/nvidia/model.py
+11
View File
@@ -46,6 +46,10 @@ from vllm.models.common.ops.sequence_parallel import (
)
from vllm.models.deepseek_v32.attention import DeepseekV32Attention
from vllm.sequence import IntermediateTensors
from vllm.v1.attention.backends.mla.index_group import (
SparseMLAIndexGroupBuilder,
get_sparse_mla_index_group_max_rows,
)
from .glm52_low_latency_gemm import enable_glm52_low_latency_gemm
@@ -57,6 +61,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
prefix: str,
config=None,
topk_indices_buffer: torch.Tensor | None = None,
index_group_builder: SparseMLAIndexGroupBuilder | None = None,
) -> None:
super().__init__()
@@ -80,6 +85,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
config=config,
prefix=f"{prefix}.self_attn",
topk_indices_buffer=topk_indices_buffer,
index_group_builder=index_group_builder,
)
if (
@@ -189,6 +195,10 @@ class DeepseekV32Model(torch.nn.Module):
dtype=torch.int32,
device=self.device,
)
index_group_builder = SparseMLAIndexGroupBuilder(
self.topk_indices_buffer,
get_sparse_mla_index_group_max_rows(vllm_config),
)
if get_pp_group().is_first_rank:
self.embed_tokens = make_input_embedding(
@@ -209,6 +219,7 @@ class DeepseekV32Model(torch.nn.Module):
vllm_config=vllm_config,
prefix=prefix,
topk_indices_buffer=self.topk_indices_buffer,
index_group_builder=index_group_builder,
),
prefix=f"{prefix}.layers",
)
vllm/models/deepseek_v32/nvidia/mtp.py
+9
View File
@@ -48,6 +48,10 @@ from vllm.models.common.ops.sequence_parallel import (
from vllm.models.deepseek_v32.common.kernels import fused_eh_norm
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm.v1.attention.backends.mla.index_group import (
SparseMLAIndexGroupBuilder,
get_sparse_mla_index_group_max_rows,
)
from .glm52_low_latency_gemm import (
build_glm52_plan,
@@ -80,6 +84,10 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
dtype=torch.int32,
device=current_platform.device_type,
)
index_group_builder = SparseMLAIndexGroupBuilder(
topk_indices_buffer,
get_sparse_mla_index_group_max_rows(vllm_config),
)
self.shared_head = SharedHead(
config=config, prefix=prefix, quant_config=quant_config
)
@@ -88,6 +96,7 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
prefix,
config=config,
topk_indices_buffer=topk_indices_buffer,
index_group_builder=index_group_builder,
)
def forward(
vllm/models/minimax_m3/common/indexer.py
+2
View File
@@ -55,6 +55,7 @@ from vllm.v1.kv_cache_interface import (
AttentionSpec,
KVCacheSpec,
MLAAttentionSpec,
SparseCacheRole,
)
logger = init_logger(__name__)
@@ -148,6 +149,7 @@ class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
num_kv_heads=1,
head_size=self.head_dim,
dtype=self.dtype,
cache_role=SparseCacheRole.INDEXER,
)
def forward(self) -> None: ...
vllm/v1/attention/backend.py
+5
View File
@@ -22,6 +22,7 @@ if TYPE_CHECKING:
from vllm.model_executor.layers.linear import ColumnParallelLinear
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
from vllm.platforms.interface import DeviceCapability
from vllm.v1.hisparse.runtime import HiSparseCacheHandle
from vllm.v1.kv_cache_interface import (
AttentionSpec,
KVCacheLayout,
@@ -838,6 +839,9 @@ class AttentionImplBase(ABC, Generic[T]):
def process_weights_after_loading(self, act_dtype: torch.dtype):
pass
def prepare_for_batch(self, attn_metadata: T | None) -> None:
"""Prepare implementation-specific state for the current batch."""
class AttentionImpl(AttentionImplBase[T], Generic[T]):
"""Standard attention implementation with forward method."""
@@ -956,6 +960,7 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):
class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
"""MLA attention implementation with forward_mqa and forward_mha methods."""
hisparse_cache: "HiSparseCacheHandle | None" = None
supports_pcp: bool = True
@abstractmethod
vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
+103 -24
View File
@@ -8,9 +8,9 @@ import torch
from vllm.config import VllmConfig
from vllm.config.cache import CacheDType
from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
from vllm.model_executor.layers.attention.sparse_mla_attention import (
SparseMLACommonImpl,
SparseMLACommonMetadata,
SparseMLACommonMetadataBuilder,
)
from vllm.platforms.interface import DeviceCapability
@@ -18,11 +18,11 @@ from vllm.v1.attention.backend import (
AttentionBackend,
AttentionCGSupport,
AttentionLayer,
AttentionMetadata,
MLAAttentionImpl,
MultipleOf,
)
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
from vllm.v1.attention.backends.mla.sparse_utils import (
flat_kv_row_view,
triton_convert_req_index_to_global_index,
@@ -106,26 +106,8 @@ class FlashAttnMLASparseBackend(AttentionBackend):
@dataclass
class FlashAttnMLASparseMetadata(AttentionMetadata):
num_reqs: int
max_query_len: int
max_seq_len: int
num_actual_tokens: int
query_start_loc: torch.Tensor
slot_mapping: torch.Tensor
block_table: torch.Tensor
req_id_per_token: torch.Tensor
seq_lens: torch.Tensor
block_size: int = 64
topk_tokens: int = 2048
num_decodes: int = 0
num_prefills: int = 0
num_decode_tokens: int = 0
prefill_max_seq_len: int = 0
prefill: MLACommonPrefillMetadata | None = None
cp_kv_cache_interleave_size: int = 1
class FlashAttnMLASparseMetadata(SparseMLACommonMetadata):
pass
class FlashAttnMLASparseMetadataBuilder(
@@ -214,6 +196,76 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
assert self.topk_indices_buffer is not None
topk_indices = self.topk_indices_buffer[:num_actual_toks]
index_group = self.index_group
if isinstance(index_group, HiSparseMLAIndexGroup):
num_decode_tokens = attn_metadata.num_decode_tokens
outputs = []
if num_decode_tokens:
physical_topk, valid_counts = (
index_group.convert_decode_logical_to_physical_topk(
self.index_group_index,
topk_indices[:num_decode_tokens],
attn_metadata,
return_valid_counts=True,
)
)
outputs.append(
self._run_mqa_kernel(
q_nope[:num_decode_tokens],
q_rope[:num_decode_tokens],
index_group.physical_kv_cache(self.index_group_index).view(
kv_c_and_k_pe_cache.dtype
),
physical_topk,
valid_counts,
attn_metadata.block_size,
)
)
if num_decode_tokens < num_actual_toks:
cache = index_group.cache(self.index_group_index)
if num_decode_tokens == 0 and cache.all_context_pages_resident:
physical_topk, valid_counts = (
index_group.convert_logical_to_physical_topk(
self.index_group_index,
topk_indices,
attn_metadata,
block_stride_rows=None,
return_valid_counts=True,
)
)
prefill_cache = index_group.physical_kv_cache(
self.index_group_index
).view(kv_c_and_k_pe_cache.dtype)
else:
prefill_cache, block_table, req_ids = (
index_group.stage_prefill_rows(
self.index_group_index,
kv_c_and_k_pe_cache,
attn_metadata,
)
)
physical_topk, valid_counts = (
triton_convert_req_index_to_global_index(
req_ids,
block_table,
topk_indices[num_decode_tokens:],
BLOCK_SIZE=attn_metadata.block_size,
NUM_TOPK_TOKENS=topk_indices.shape[1],
return_valid_counts=True,
)
)
outputs.append(
self._run_mqa_kernel(
q_nope[num_decode_tokens:],
q_rope[num_decode_tokens:],
prefill_cache,
physical_topk,
valid_counts,
attn_metadata.block_size,
)
)
return torch.cat(outputs) if len(outputs) > 1 else outputs[0], None
kv_rows, block_stride_rows = flat_kv_row_view(
kv_c_and_k_pe_cache, attn_metadata.block_size
)
@@ -226,9 +278,36 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
NUM_TOPK_TOKENS=topk_indices.shape[1],
return_valid_counts=True,
)
return (
self._run_mqa_kernel(
q_nope,
q_rope,
kv_rows,
topk_indices,
valid_counts,
attn_metadata.block_size,
cache_is_flat=True,
),
None,
)
def _run_mqa_kernel(
self,
q_nope: torch.Tensor,
q_rope: torch.Tensor,
kv_cache: torch.Tensor,
topk_indices: torch.Tensor,
valid_counts: torch.Tensor,
block_size: int,
*,
cache_is_flat: bool = False,
) -> torch.Tensor:
kv_rows = (
kv_cache if cache_is_flat else flat_kv_row_view(kv_cache, block_size)[0]
)
cu_seqlens_q = torch.arange(
0, num_actual_toks + 1, dtype=torch.int32, device=q_rope.device
0, q_rope.shape[0] + 1, dtype=torch.int32, device=q_rope.device
)
k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1)
v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1)
@@ -247,4 +326,4 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
causal=True,
fa_version=3,
)
return out, None
return out
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+154 -41
View File
@@ -11,22 +11,23 @@ from vllm import envs
from vllm.config import VllmConfig
from vllm.config.cache import CacheDType
from vllm.logger import init_logger
from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
from vllm.model_executor.layers.attention.sparse_mla_attention import (
SparseMLACommonImpl,
SparseMLACommonMetadata,
SparseMLACommonMetadataBuilder,
)
from vllm.platforms import current_platform
from vllm.platforms.interface import DeviceCapability
from vllm.utils.torch_utils import is_quantized_kv_cache
from vllm.v1.attention.backend import (
AttentionBackend,
AttentionCGSupport,
AttentionLayer,
AttentionMetadata,
AttentionType,
MLAAttentionImpl,
MultipleOf,
)
from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
from vllm.v1.attention.backends.mla.sparse_utils import (
flat_kv_row_view,
triton_convert_req_index_to_global_index,
@@ -224,33 +225,9 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):
@dataclass
class FlashInferMLASparseMetadata(AttentionMetadata):
class FlashInferMLASparseMetadata(SparseMLACommonMetadata):
"""Attention metadata for FlashInfer MLA Sparse backend."""
num_reqs: int
max_query_len: int
max_seq_len: int
num_actual_tokens: int
# Query start locations
query_start_loc: torch.Tensor
slot_mapping: torch.Tensor
block_table: torch.Tensor
req_id_per_token: torch.Tensor
# Sequence lengths for all requests (context + query)
seq_lens: torch.Tensor
num_decodes: int
num_prefills: int
num_decode_tokens: int
prefill_max_seq_len: int = 0
prefill: MLACommonPrefillMetadata | None = None
# Sparse-specific
block_size: int = 64
topk_tokens: int = 2048
cp_kv_cache_interleave_size: int = 1
class FlashInferMLASparseMetadataBuilder(
SparseMLACommonMetadataBuilder[FlashInferMLASparseMetadata]
@@ -286,6 +263,7 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde
"""Metadata builder for the SM100 TRT-LLM sparse MLA kernel."""
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
hisparse_supports_multi_token_decode: ClassVar[bool] = True
def _build_req_id_per_token(
self,
@@ -402,6 +380,78 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
assert self.topk_indices_buffer is not None
topk_indices = self.topk_indices_buffer[:num_actual_toks]
self._prepare_mqa_kernel(layer, q.device)
index_group = self.index_group
if isinstance(index_group, HiSparseMLAIndexGroup):
num_decode_tokens = attn_metadata.num_decode_tokens
decode_out: torch.Tensor | None = None
decode_lse: torch.Tensor | None = None
if num_decode_tokens > 0:
physical_topk, valid_counts = (
index_group.convert_decode_logical_to_physical_topk(
self.index_group_index,
topk_indices[:num_decode_tokens],
attn_metadata,
return_valid_counts=True,
)
)
decode_out, decode_lse = self._run_mqa_kernel(
q[:num_decode_tokens],
index_group.physical_kv_cache(self.index_group_index).view(
kv_c_and_k_pe_cache.dtype
),
physical_topk,
valid_counts,
)
if num_decode_tokens == num_actual_toks:
return decode_out, decode_lse
cache = index_group.cache(self.index_group_index)
if num_decode_tokens == 0 and cache.all_context_pages_resident:
physical_topk, valid_counts = (
index_group.convert_logical_to_physical_topk(
self.index_group_index,
topk_indices,
attn_metadata,
block_stride_rows=None,
return_valid_counts=True,
)
)
return self._run_mqa_kernel(
q,
index_group.physical_kv_cache(self.index_group_index).view(
kv_c_and_k_pe_cache.dtype
),
physical_topk,
valid_counts,
)
prefill_cache, block_table, req_ids = index_group.stage_prefill_rows(
self.index_group_index, kv_c_and_k_pe_cache, attn_metadata
)
prefill_indices, prefill_lens = triton_convert_req_index_to_global_index(
req_ids,
block_table,
topk_indices[num_decode_tokens:],
BLOCK_SIZE=attn_metadata.block_size,
NUM_TOPK_TOKENS=topk_indices.shape[1],
return_valid_counts=True,
)
prefill_out, prefill_lse = self._run_mqa_kernel(
q[num_decode_tokens:],
prefill_cache,
prefill_indices,
prefill_lens,
)
if decode_out is None:
return prefill_out, prefill_lse
output = torch.cat((decode_out, prefill_out))
if decode_lse is None:
return output, None
assert prefill_lse is not None
return output, torch.cat((decode_lse, prefill_lse))
_, block_stride_rows = flat_kv_row_view(
kv_c_and_k_pe_cache, attn_metadata.block_size
)
@@ -420,18 +470,27 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
return_valid_counts=True,
)
else:
topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index(
attn_metadata.req_id_per_token[:num_actual_toks],
attn_metadata.block_table,
topk_indices_physical, seq_lens = self._convert_logical_to_physical_topk(
topk_indices,
BLOCK_SIZE=attn_metadata.block_size,
BLOCK_STRIDE_ROWS=block_stride_rows,
NUM_TOPK_TOKENS=topk_indices.shape[1],
attn_metadata,
block_stride_rows=block_stride_rows,
return_valid_counts=True,
)
return self._run_mqa_kernel(
q,
kv_c_and_k_pe_cache,
topk_indices_physical,
seq_lens,
)
def _prepare_mqa_kernel(
self,
layer: AttentionLayer,
device: torch.device,
) -> None:
if self._workspace_buffer is None:
self._workspace_buffer = _get_workspace_buffer(q.device)
self._workspace_buffer = _get_workspace_buffer(device)
if self.bmm1_scale is None:
self.bmm1_scale = self.scale
@@ -442,15 +501,71 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
if is_quantized_kv_cache(self.kv_cache_dtype):
self.bmm2_scale *= layer._k_scale_float
def autotune_hisparse_decode(self, layer: AttentionLayer) -> None:
"""Autotune the largest legal HiSparse decode batch."""
assert isinstance(self.index_group, HiSparseMLAIndexGroup)
cache = self.index_group.cache(self.index_group_index)
assert self.topk_indices_buffer is not None
runtime = cache.runtime
kv_cache = runtime.hot.attention_cache
num_tokens = runtime.max_num_reqs
topk_tokens = self.topk_indices_buffer.shape[1]
self._prepare_mqa_kernel(layer, kv_cache.device)
q_dtype = (
current_platform.fp8_dtype()
if is_quantized_kv_cache(self.kv_cache_dtype)
else kv_cache.dtype
)
q = torch.zeros(
(
num_tokens,
self.num_heads,
self.kv_lora_rank + self.qk_rope_head_dim,
),
dtype=q_dtype,
device=kv_cache.device,
)
topk_indices = (
torch.arange(
topk_tokens,
dtype=torch.int32,
device=kv_cache.device,
)
.expand(num_tokens, -1)
.contiguous()
)
seq_lens = torch.full(
(num_tokens,),
topk_tokens,
dtype=torch.int32,
device=kv_cache.device,
)
self._run_mqa_kernel(q, kv_cache, topk_indices, seq_lens)
def _run_mqa_kernel(
self,
q: torch.Tensor,
kv_cache: torch.Tensor,
topk_indices: torch.Tensor,
seq_lens: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor | None]:
assert self._workspace_buffer is not None
assert self.bmm1_scale is not None
assert self.bmm2_scale is not None
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
kv_cache = kv_cache.view(q.dtype)
# Single-token sparse decode. trtllm-gen requires the q_len_per_request
# dim, but the sparse attention mask is fully per-token (each query token
# carries its own top-k index row), so unsqueeze is sufficient and
# correct. The MTP/multi-token q_len grouping is a perf-only layout and is
# deferred until MTP is validated end-to-end for this backend.
query = q.unsqueeze(1)
block_tables = topk_indices_physical.unsqueeze(1)
block_tables = topk_indices.unsqueeze(1)
seq_lens_arg = seq_lens
# page_table width = topk buffer width, which kpool widens past
@@ -460,7 +575,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
# count), so the -1 padding slots past seq_lens are never attended to.
# Use the actual buffer width instead of the fixed topk_tokens, which
# mismatches the page_table when index_kpool > 1.
sparse_topk_capacity = topk_indices_physical.shape[1]
sparse_topk_capacity = topk_indices.shape[1]
extra_kwargs: dict[str, torch.Tensor] = {}
empty_rows: torch.Tensor | None = None
@@ -473,14 +588,12 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
# valid indices into a contiguous prefix, which is what the kernel
# requires of the page table.
empty_rows = seq_lens == 0
topk_indices_physical[:, 0] = topk_indices_physical[:, 0].masked_fill(
empty_rows, 0
)
topk_indices[:, 0] = topk_indices[:, 0].masked_fill(empty_rows, 0)
extra_kwargs["sparse_mla_top_k_lens"] = seq_lens.clamp(min=1)
kernel_out = trtllm_batch_decode_with_kv_cache_mla(
query=query,
kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
kv_cache=kv_cache.unsqueeze(1),
workspace_buffer=self._workspace_buffer,
qk_nope_head_dim=self.qk_nope_head_dim,
kv_lora_rank=self.kv_lora_rank,
@@ -506,7 +619,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
if lse is not None:
lse = self._normalize_lse(lse, out.shape[0], out.shape[1])
if empty_rows is None and lse is not None:
empty_rows = (topk_indices_physical == -1).all(dim=-1)
empty_rows = (topk_indices == -1).all(dim=-1)
if empty_rows is not None:
out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)
if lse is not None:
vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
+106 -25
View File
@@ -7,17 +7,14 @@ from typing import TYPE_CHECKING, cast
import torch
from vllm.model_executor.layers.attention.sparse_mla_attention import (
SharedTopkIndicesBuffer,
)
from vllm.v1.attention.backend import (
AttentionLayer,
AttentionType,
MLAAttentionImpl,
SparseMLACommonImpl,
)
from vllm.v1.attention.backend import AttentionLayer, AttentionType
from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
FlashInferMLASparseMetadata,
_get_workspace_buffer,
)
from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
from vllm.v1.attention.backends.mla.sparse_utils import (
triton_convert_req_index_to_global_index,
)
@@ -32,9 +29,7 @@ def _kv_scale_format_for_model(model_type: str | None) -> str:
return "pow2_fp32"
class FlashInferMLASparseSM120Impl(
MLAAttentionImpl[FlashInferMLASparseMetadata], SharedTopkIndicesBuffer
):
class FlashInferMLASparseSM120Impl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
"""SM120 FlashInfer sparse-MLA implementation."""
is_sparse = True
@@ -65,20 +60,28 @@ class FlashInferMLASparseSM120Impl(
"FLASHINFER_MLA_SPARSE_SM120 only supports decoder self-attention"
)
self.num_heads = num_heads
self.head_size = head_size
self.scale = float(scale)
self.num_kv_heads = num_kv_heads
self.kv_cache_dtype = kv_cache_dtype
if self.kv_cache_dtype != "fp8_ds_mla":
if kv_cache_dtype != "fp8_ds_mla":
raise NotImplementedError(
"FLASHINFER_MLA_SPARSE_SM120 requires the packed fp8_ds_mla "
f"KV cache layout; got kv_cache_dtype={kv_cache_dtype!r}."
)
self.kv_lora_rank: int = mla_args["kv_lora_rank"]
self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"]
self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
topk_indices_buffer = mla_args.pop("topk_indices_buffer", None)
super().__init__(
num_heads,
head_size,
scale,
num_kv_heads,
alibi_slopes,
sliding_window,
kv_cache_dtype,
logits_soft_cap,
attn_type,
kv_sharing_target_layer_name,
indexer=indexer,
topk_indices_buffer=topk_indices_buffer,
**mla_args,
)
from vllm.config import get_current_vllm_config
vllm_config = get_current_vllm_config()
@@ -89,9 +92,6 @@ class FlashInferMLASparseSM120Impl(
)
self.kv_scale_format = _kv_scale_format_for_model(model_type)
# Skip-topk layers are built with indexer=None and get the shared
# buffer via mla_args instead (cf. FLASHMLA_SPARSE).
self.init_topk_indices_buffer(indexer, mla_args.get("topk_indices_buffer"))
from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
if not has_flashinfer_sparse_mla_sm120():
@@ -118,6 +118,70 @@ class FlashInferMLASparseSM120Impl(
assert self.topk_indices_buffer is not None
topk_indices = self.topk_indices_buffer[:num_actual_toks]
index_group = self.index_group
if isinstance(index_group, HiSparseMLAIndexGroup):
num_decode_tokens = attn_metadata.num_decode_tokens
outputs = []
if num_decode_tokens:
topk_indices_physical = cast(
torch.Tensor,
index_group.convert_decode_logical_to_physical_topk(
self.index_group_index,
topk_indices[:num_decode_tokens],
attn_metadata,
return_valid_counts=False,
),
)
outputs.append(
self._run_mqa_kernel(
q[:num_decode_tokens],
index_group.physical_kv_cache(self.index_group_index),
topk_indices_physical,
)
)
if num_decode_tokens < num_actual_toks:
cache = index_group.cache(self.index_group_index)
if num_decode_tokens == 0 and cache.all_context_pages_resident:
topk_indices_physical = cast(
torch.Tensor,
index_group.convert_logical_to_physical_topk(
self.index_group_index,
topk_indices,
attn_metadata,
block_stride_rows=None,
return_valid_counts=False,
),
)
prefill_cache = index_group.physical_kv_cache(
self.index_group_index
)
else:
prefill_cache, block_table, req_ids = (
index_group.stage_prefill_rows(
self.index_group_index,
kv_c_and_k_pe_cache,
attn_metadata,
)
)
topk_indices_physical = cast(
torch.Tensor,
triton_convert_req_index_to_global_index(
req_ids,
block_table,
topk_indices[num_decode_tokens:],
BLOCK_SIZE=attn_metadata.block_size,
NUM_TOPK_TOKENS=topk_indices.shape[1],
),
)
outputs.append(
self._run_mqa_kernel(
q[num_decode_tokens:],
prefill_cache,
topk_indices_physical,
)
)
output = torch.cat(outputs) if len(outputs) > 1 else outputs[0]
return output, None
topk_indices_physical = cast(
torch.Tensor,
@@ -129,6 +193,22 @@ class FlashInferMLASparseSM120Impl(
NUM_TOPK_TOKENS=topk_indices.shape[1],
),
)
return (
self._run_mqa_kernel(
q,
kv_c_and_k_pe_cache,
topk_indices_physical,
),
None,
)
def _run_mqa_kernel(
self,
q: torch.Tensor,
kv_cache: torch.Tensor,
topk_indices_physical: torch.Tensor,
) -> torch.Tensor:
num_actual_toks = q.shape[0]
output = q.new_empty(
(num_actual_toks, self.num_heads, self.kv_lora_rank),
@@ -142,20 +222,21 @@ class FlashInferMLASparseSM120Impl(
flashinfer_trtllm_batch_decode_with_kv_cache_mla,
)
sparse_capacity = topk_indices_physical.shape[1]
out = flashinfer_trtllm_batch_decode_with_kv_cache_mla(
query=q.unsqueeze(1),
kv_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(1),
kv_cache=kv_cache.view(torch.uint8).unsqueeze(1),
workspace_buffer=self._workspace_buffer,
qk_nope_head_dim=self.qk_nope_head_dim,
kv_lora_rank=self.kv_lora_rank,
qk_rope_head_dim=self.qk_rope_head_dim,
block_tables=topk_indices_physical.unsqueeze(1),
seq_lens=None,
max_seq_len=attn_metadata.topk_tokens,
max_seq_len=sparse_capacity,
out=output.unsqueeze(1),
bmm1_scale=self.scale,
bmm2_scale=1.0,
sparse_mla_top_k=attn_metadata.topk_tokens,
sparse_mla_top_k=sparse_capacity,
kv_scale_format=self.kv_scale_format,
)
return out.squeeze(1), None
return out.squeeze(1)
vllm/v1/attention/backends/mla/flashmla_sparse.py
+225 -91
View File
@@ -9,23 +9,23 @@ from vllm import _custom_ops as ops
from vllm.config import VllmConfig, get_current_vllm_config
from vllm.config.cache import CacheDType
from vllm.logger import init_logger
from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
from vllm.model_executor.layers.attention.sparse_mla_attention import (
SparseMLACommonImpl,
SparseMLACommonMetadata,
SparseMLACommonMetadataBuilder,
)
from vllm.platforms import current_platform
from vllm.platforms.interface import DeviceCapability
from vllm.utils.platform_utils import num_compute_units
from vllm.utils.torch_utils import is_quantized_kv_cache
from vllm.utils.torch_utils import current_stream, is_quantized_kv_cache
from vllm.v1.attention.backend import (
AttentionBackend,
AttentionCGSupport,
AttentionLayer,
AttentionMetadata,
CommonAttentionMetadata,
MultipleOf,
)
from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
from vllm.v1.attention.backends.mla.sparse_utils import (
flat_kv_row_view,
triton_convert_req_index_to_global_index,
@@ -172,28 +172,7 @@ class FlashMLASparseBackend(AttentionBackend):
@dataclass
class FlashMLASparseMetadata(AttentionMetadata):
num_reqs: int
max_query_len: int
max_seq_len: int
num_actual_tokens: int # Number of tokens excluding padding.
query_start_loc: torch.Tensor
slot_mapping: torch.Tensor
block_table: torch.Tensor
req_id_per_token: torch.Tensor
block_size: int = 64
topk_tokens: int = 2048
num_decodes: int = 0
num_prefills: int = 0
num_decode_tokens: int = 0
seq_lens: torch.Tensor | None = None
prefill_max_seq_len: int = 0
prefill: MLACommonPrefillMetadata | None = None
cp_kv_cache_interleave_size: int = 1
class FlashMLASparseMetadata(SparseMLACommonMetadata):
@dataclass
class FP8KernelMetadata:
scheduler_metadata: FlashMLASchedMeta
@@ -233,6 +212,7 @@ class FlashMLASparseMetadata(AttentionMetadata):
req_start_idx: int
workspace_starts: torch.Tensor
chunk_tot_seqlen: int
seq_lens: torch.Tensor | None = None
chunks: list[Chunk]
@@ -263,6 +243,7 @@ class FlashMLASparseMetadataBuilder(
):
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
require_uniform_decodes: ClassVar[bool] = True
hisparse_supports_multi_token_decode: ClassVar[bool] = True
metadata_cls = FlashMLASparseMetadata
def __init__(
@@ -283,6 +264,9 @@ class FlashMLASparseMetadataBuilder(
)
else:
threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256)
self.use_hisparse = vllm_config.attention_config.hisparse_config is not None
if self.use_hisparse:
threshold = 1
# Varlen decodes are safe under DCP: causality comes from the
# indexer's top-k indices, not from the kernel metadata.
self._init_reorder_batch_threshold(
@@ -347,7 +331,9 @@ class FlashMLASparseMetadataBuilder(
device=device,
)
self.fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL
self.fp8_use_mixed_batch = (
self.num_heads < MIN_HEADS_FOR_BF16_PREFILL and not self.use_hisparse
)
if parallel_config.decode_context_parallel_size > 1:
if parallel_config.dcp_comm_backend != "ag_rs":
@@ -385,37 +371,21 @@ class FlashMLASparseMetadataBuilder(
def _build_fp8_mixed_decode_prefill(
self,
common_attn_metadata: CommonAttentionMetadata,
) -> "FlashMLASparseMetadata.FP8KernelMetadata":
) -> FlashMLASparseMetadata.FP8KernelMetadata:
"""Build FP8 metadata treating MQA tokens as one batch.
The scheduler initializes lazily from the runtime query shape, which may
be the full batch or only decodes when prefills use dense MHA. This avoids
the BF16 prefill kernel's head-padding overhead at high TP.
"""
num_tokens = common_attn_metadata.num_actual_tokens
# Use padded head count since that's what the kernel will see
padded_heads = self.fp8_decode_padded_heads
# Build metadata for all tokens as a single batch
scheduler_metadata, _ = get_mla_metadata(
cache_seqlens=self.topk_tokens_tensor[:1], # Single batch
num_q_tokens_per_head_k=num_tokens * padded_heads,
topk=self.topk_tokens,
num_heads_q=padded_heads,
num_heads_k=1,
is_fp8_kvcache=True,
)
fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata(
scheduler_metadata, _ = get_mla_metadata()
return FlashMLASparseMetadata.FP8KernelMetadata(
scheduler_metadata=scheduler_metadata,
cache_lens=self.max_model_len_tensor[:1],
dummy_block_table=self.dummy_block_table[:1],
)
return fp8_metadata
def _build_fp8_separate_prefill_decode(
self,
common_attn_metadata: CommonAttentionMetadata,
@@ -520,6 +490,9 @@ class FlashMLASparseMetadataBuilder(
chunk_block_table = common_attn_metadata.block_table_tensor[
num_decodes + chunk_start : num_decodes + chunk_end
]
chunk_seq_lens = common_attn_metadata.seq_lens[
num_decodes + chunk_start : num_decodes + chunk_end
]
prefill_chunks.append(
FP8Meta.Prefill.Chunk(
@@ -528,6 +501,7 @@ class FlashMLASparseMetadataBuilder(
req_start_idx=chunk_start,
workspace_starts=chunk_workspace_starts,
chunk_tot_seqlen=chunk_tot_seqlen,
seq_lens=chunk_seq_lens,
)
)
@@ -569,9 +543,7 @@ class FlashMLASparseMetadataBuilder(
metadata.fp8_use_mixed_batch = self.fp8_use_mixed_batch
if self.use_fp8_kv_cache:
if self.fp8_use_mixed_batch:
metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill(
common_attn_metadata
)
metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill()
else:
metadata.fp8_extra_metadata = self._build_fp8_separate_prefill_decode(
common_attn_metadata, metadata
@@ -679,30 +651,83 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
topk_indices: torch.Tensor,
attn_metadata: FlashMLASparseMetadata,
actual_num_heads: int,
) -> tuple[torch.Tensor, torch.Tensor]:
# Convert per-request indices to global slots (decode) or workspace
# offsets (prefill). req_id_per_token covers the whole batch; slice it
# to the MQA tokens (q may exclude prefill tokens routed to dense MHA).
) -> tuple[torch.Tensor, torch.Tensor | None]:
index_group = self.index_group
if isinstance(index_group, HiSparseMLAIndexGroup):
cache = index_group.cache(self.index_group_index)
else:
cache = None
block_table = attn_metadata.block_table
# req_id_per_token covers the whole batch; slice it to the MQA tokens
# (q may exclude prefill tokens routed to dense MHA).
req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]]
decode_out: torch.Tensor | None = None
if cache is not None:
assert isinstance(index_group, HiSparseMLAIndexGroup)
num_decode_tokens = attn_metadata.num_decode_tokens
if num_decode_tokens > 0:
decode_topk, decode_lengths = (
index_group.convert_decode_logical_to_physical_topk(
self.index_group_index,
topk_indices[:num_decode_tokens],
attn_metadata,
return_valid_counts=True,
)
)
decode_out, _ = self._bf16_flash_mla_kernel(
q[:num_decode_tokens],
index_group.physical_kv_cache(self.index_group_index),
decode_topk,
decode_lengths,
actual_num_heads,
)
if num_decode_tokens == q.shape[0]:
return decode_out, None
q = q[num_decode_tokens:]
topk_indices = topk_indices[num_decode_tokens:]
kv_c_and_k_pe_cache, block_table, req_id_per_token = (
index_group.stage_prefill_rows(
self.index_group_index, kv_c_and_k_pe_cache, attn_metadata
)
)
# Convert per-request indices to global slots (decode) or workspace offsets.
kv_rows, block_stride_rows = flat_kv_row_view(
kv_c_and_k_pe_cache, attn_metadata.block_size
)
topk_indices, topk_length = triton_convert_req_index_to_global_index(
attn_metadata.req_id_per_token[: topk_indices.shape[0]],
attn_metadata.block_table,
topk_indices,
BLOCK_SIZE=attn_metadata.block_size,
BLOCK_STRIDE_ROWS=block_stride_rows,
NUM_TOPK_TOKENS=topk_indices.shape[1],
return_valid_counts=True,
decode_only = (
attn_metadata.num_decode_tokens
== attn_metadata.num_actual_tokens
== topk_indices.shape[0]
)
uses_host_cache = isinstance(index_group, HiSparseMLAIndexGroup)
if not uses_host_cache and decode_only:
topk_indices, topk_length = self._convert_logical_to_physical_topk(
topk_indices,
attn_metadata,
block_stride_rows=block_stride_rows,
return_valid_counts=True,
)
else:
topk_indices, topk_length = triton_convert_req_index_to_global_index(
req_id_per_token,
block_table,
topk_indices,
BLOCK_SIZE=attn_metadata.block_size,
BLOCK_STRIDE_ROWS=block_stride_rows,
NUM_TOPK_TOKENS=topk_indices.shape[1],
return_valid_counts=True,
)
return self._bf16_flash_mla_kernel(
attn_out, lse = self._bf16_flash_mla_kernel(
q,
kv_rows,
topk_indices,
topk_length,
actual_num_heads,
)
if decode_out is None:
return attn_out, lse
return torch.cat([decode_out, attn_out], dim=0), None
def _forward_fp8_kv_separate_prefill_decode(
self,
@@ -721,6 +746,28 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
"FP8 sparse MLA expects either the decode subset or the full batch"
)
decode_topk: torch.Tensor | None = None
index_group = self.index_group
uses_host_cache = isinstance(index_group, HiSparseMLAIndexGroup)
if uses_host_cache and num_decode_tokens > 0:
decode_topk = topk_indices[:num_decode_tokens]
prefill_ready = None
if num_prefill_tokens > 0 and uses_host_cache:
assert fp8_metadata.prefill is not None
first_chunk = fp8_metadata.prefill.chunks[0]
assert isinstance(index_group, HiSparseMLAIndexGroup)
prefill_ready = index_group.gather_fp8_prefill(
self.index_group_index,
kv_c_and_k_pe_cache,
self.prefill_bf16_workspace[: first_chunk.chunk_tot_seqlen],
first_chunk.block_table,
first_chunk.workspace_starts,
len(first_chunk.block_table),
attn_metadata,
first_chunk.req_start_idx,
)
prefill_request_ids = None
prefill_workspace_starts = None
has_prefill_workspace = False
@@ -736,17 +783,26 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
# For BF16 cache: always use global cache slots (no workspace)
# prefill_workspace_starts has been adjusted in-place per chunk so
# prefill indices automatically come out chunk-local
topk_indices, topk_length = triton_convert_req_index_to_global_index(
attn_metadata.req_id_per_token[: topk_indices.shape[0]],
attn_metadata.block_table,
topk_indices,
BLOCK_SIZE=attn_metadata.block_size,
NUM_TOPK_TOKENS=topk_indices.shape[1],
HAS_PREFILL_WORKSPACE=has_prefill_workspace,
prefill_workspace_request_ids=prefill_request_ids,
prefill_workspace_starts=prefill_workspace_starts,
return_valid_counts=True,
)
topk_length = None
if num_prefill_tokens == 0 and not uses_host_cache:
topk_indices, topk_length = self._convert_logical_to_physical_topk(
topk_indices,
attn_metadata,
block_stride_rows=None,
return_valid_counts=True,
)
elif num_prefill_tokens > 0:
topk_indices, topk_length = triton_convert_req_index_to_global_index(
attn_metadata.req_id_per_token[: topk_indices.shape[0]],
attn_metadata.block_table,
topk_indices,
BLOCK_SIZE=attn_metadata.block_size,
NUM_TOPK_TOKENS=topk_indices.shape[1],
HAS_PREFILL_WORKSPACE=has_prefill_workspace,
prefill_workspace_request_ids=prefill_request_ids,
prefill_workspace_starts=prefill_workspace_starts,
return_valid_counts=True,
)
fp8_metadata = attn_metadata.fp8_extra_metadata
assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode)
@@ -755,6 +811,16 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
q: torch.Tensor,
topk_indices: torch.Tensor,
) -> torch.Tensor:
assert fp8_metadata.decode is not None
if uses_host_cache:
return self._host_backed_fp8_decode(
q,
topk_indices,
attn_metadata,
fp8_metadata.decode.kernel_metadata,
num_decodes,
fp8_metadata.decode.decode_query_len,
)
# Reshape q: (num_decode_tokens, num_heads, head_dim)
# -> (num_decodes, seq_len, num_heads, head_dim)
q = reshape_query_for_spec_decode(q, num_decodes)
@@ -762,7 +828,6 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
# Reshape topk_indices: (num_decode_tokens, topk)
# -> (num_decodes, seq_len, topk)
topk_indices = topk_indices.view(num_decodes, seq_len, -1)
assert fp8_metadata.decode is not None
attn_out, _ = self._fp8_flash_mla_kernel(
q=q,
kv_c_and_k_pe_cache=kv_c_and_k_pe_cache,
@@ -776,7 +841,9 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
# Pure decode: direct call without allocation
if num_decode_tokens > 0 and num_prefill_tokens == 0:
assert fp8_metadata.decode is not None
attn_out = _fp8_decode(q, topk_indices)
attn_out = _fp8_decode(
q, decode_topk if decode_topk is not None else topk_indices
)
else:
# Mixed or pure prefill: allocate output tensor
attn_out = q.new_empty(
@@ -788,13 +855,30 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
if num_decode_tokens > 0:
attn_out[:num_decode_tokens] = _fp8_decode(
q[:num_decode_tokens],
topk_indices[:num_decode_tokens],
decode_topk
if decode_topk is not None
else topk_indices[:num_decode_tokens],
)
assert fp8_metadata.prefill is not None
for chunk in fp8_metadata.prefill.chunks:
for chunk_index, chunk in enumerate(fp8_metadata.prefill.chunks):
chunk_workspace = self.prefill_bf16_workspace[: chunk.chunk_tot_seqlen]
if self.kv_cache_dtype == "fp8_ds_mla":
if uses_host_cache and chunk_index > 0:
assert isinstance(index_group, HiSparseMLAIndexGroup)
prefill_ready = index_group.gather_fp8_prefill(
self.index_group_index,
kv_c_and_k_pe_cache,
chunk_workspace,
chunk.block_table,
chunk.workspace_starts,
len(chunk.block_table),
attn_metadata,
chunk.req_start_idx,
)
if uses_host_cache:
assert prefill_ready is not None
current_stream().wait_event(prefill_ready)
elif self.kv_cache_dtype == "fp8_ds_mla":
ops.cp_gather_and_upconvert_fp8_kv_cache(
kv_c_and_k_pe_cache,
chunk_workspace,
@@ -813,6 +897,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
chunk_q = q[chunk.tokens_slice]
chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice]
assert topk_length is not None
chunk_topk_length = topk_length[chunk.tokens_slice]
attn_out[chunk.tokens_slice], _ = self._bf16_flash_mla_kernel(
@@ -839,6 +924,17 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
The lse is only returned when DCP needs it, otherwise None.
"""
assert attn_metadata.fp8_extra_metadata is not None
assert isinstance(
attn_metadata.fp8_extra_metadata,
FlashMLASparseMetadata.FP8KernelMetadata,
)
fp8_metadata = attn_metadata.fp8_extra_metadata
block_table = attn_metadata.block_table
# req_id_per_token covers the whole batch; slice it to the MQA tokens
# (q may exclude prefill tokens routed to dense MHA).
req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]]
if self.dcp_world_size > 1:
# The indexer emits global token ids; keep this rank's shard and
# convert to local slots. compact_valid_to_front=False keeps the
@@ -846,8 +942,8 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
# empty-row neutralization below relies on. req_id is sliced to
# topk_indices rows (the converter grids from req_id).
topk_indices = triton_filter_and_convert_dcp_index(
attn_metadata.req_id_per_token[: topk_indices.shape[0]],
attn_metadata.block_table,
req_id_per_token,
block_table,
topk_indices,
dcp_size=self.dcp_world_size,
dcp_rank=self.dcp_rank,
@@ -859,19 +955,26 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
else:
# Convert per-request indices to global slots (decode) or workspace
# offsets (prefill).
topk_indices = triton_convert_req_index_to_global_index(
attn_metadata.req_id_per_token[: topk_indices.shape[0]],
attn_metadata.block_table,
topk_indices,
BLOCK_SIZE=attn_metadata.block_size,
NUM_TOPK_TOKENS=topk_indices.shape[1],
decode_only = (
attn_metadata.num_decode_tokens
== attn_metadata.num_actual_tokens
== topk_indices.shape[0]
)
assert attn_metadata.fp8_extra_metadata is not None
assert isinstance(
attn_metadata.fp8_extra_metadata, FlashMLASparseMetadata.FP8KernelMetadata
)
fp8_metadata = attn_metadata.fp8_extra_metadata
if decode_only:
topk_indices = self._convert_logical_to_physical_topk(
topk_indices,
attn_metadata,
block_stride_rows=None,
return_valid_counts=False,
)
else:
topk_indices = triton_convert_req_index_to_global_index(
req_id_per_token,
block_table,
topk_indices,
BLOCK_SIZE=attn_metadata.block_size,
NUM_TOPK_TOKENS=topk_indices.shape[1],
)
_attn_out, _lse = self._fp8_flash_mla_kernel(
q=q.unsqueeze(0), # unsqueeze to add batch_dim: (T, H, D) -> (1, T, H, D)
@@ -937,6 +1040,37 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
return out, lse
def _host_backed_fp8_decode(
self,
q: torch.Tensor,
topk_indices: torch.Tensor,
attn_metadata: FlashMLASparseMetadata,
kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata,
num_decodes: int,
decode_query_len: int,
) -> torch.Tensor:
assert isinstance(self.index_group, HiSparseMLAIndexGroup)
physical_topk = self.index_group.convert_decode_logical_to_physical_topk(
self.index_group_index,
topk_indices,
attn_metadata,
return_valid_counts=False,
num_decodes=num_decodes,
decode_query_len=decode_query_len,
)
assert isinstance(physical_topk, torch.Tensor)
q = reshape_query_for_spec_decode(q, num_decodes)
physical_topk = physical_topk.view(num_decodes, q.shape[1], -1)
output, _ = self._fp8_flash_mla_kernel(
q=q,
kv_c_and_k_pe_cache=self.index_group.physical_kv_cache(
self.index_group_index
),
topk_indices=physical_topk,
kernel_metadata=kernel_metadata,
)
return reshape_attn_output_for_spec_decode(output)
def _bf16_flash_mla_kernel(
self,
q: torch.Tensor,
vllm/v1/attention/backends/mla/index_group.py
+503
View File
@@ -0,0 +1,503 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
import torch
from vllm import _custom_ops as ops
from vllm.config import VllmConfig
from vllm.utils.torch_utils import current_stream
from vllm.v1.attention.backends.mla.sparse_utils import (
triton_convert_req_index_to_global_index,
)
from vllm.v1.hisparse.runtime import (
FP8_DS_MLA_ROW_BYTES,
HiSparseCacheHandle,
HiSparsePrefillStagingPlan,
create_hisparse_cache_handle,
)
def _create_side_stream(device: torch.device) -> torch.Stream:
return torch.Stream(device=device)
def _create_event() -> torch.Event:
return torch.Event()
@dataclass
class SparseMLAIndexGroup:
"""Layers that consume one sparse-indexer result."""
logical_topk_indices: torch.Tensor
physical_topk_indices: torch.Tensor
valid_topk_counts: torch.Tensor
request_ids: torch.Tensor
side_stream: torch.Stream
logical_topk_ready: torch.Event
physical_topk_ready: torch.Event
has_indexer: bool
num_layers: int = 0
def register_layer(
self,
vllm_config: VllmConfig | None = None,
*,
head_size: int | None = None,
kv_cache_dtype: str | None = None,
) -> int:
layer_index = self.num_layers
self.num_layers += 1
return layer_index
def set_logical_topk_ready(self, layer_index: int) -> None:
if layer_index == 0 and self.has_indexer:
self.logical_topk_ready.record(current_stream())
def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None:
pass
def gather_fp8_prefill(
self,
layer_index: int,
source_cache: torch.Tensor,
dst: torch.Tensor,
block_table: torch.Tensor,
workspace_starts: torch.Tensor,
batch_size: int,
attn_metadata: Any,
request_start: int,
) -> torch.Event | None:
ops.cp_gather_and_upconvert_fp8_kv_cache(
source_cache,
dst,
block_table,
workspace_starts,
batch_size,
)
return None
def convert_logical_to_physical_topk(
self,
layer_index: int,
logical_topk_indices: torch.Tensor,
attn_metadata: Any,
*,
block_stride_rows: int | None,
return_valid_counts: bool,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""Convert one index result once for all layers in this group."""
return self._convert_once(
layer_index,
logical_topk_indices,
attn_metadata.req_id_per_token[: logical_topk_indices.shape[0]],
attn_metadata.block_table,
attn_metadata.block_size,
block_stride_rows=block_stride_rows,
return_valid_counts=return_valid_counts,
)
def _convert_once(
self,
layer_index: int,
logical_topk_indices: torch.Tensor,
req_id_per_token: torch.Tensor,
block_table: torch.Tensor,
block_size: int,
*,
block_stride_rows: int | None,
return_valid_counts: bool,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
num_tokens = logical_topk_indices.shape[0]
if num_tokens > self.physical_topk_indices.shape[0]:
return triton_convert_req_index_to_global_index(
req_id_per_token,
block_table,
logical_topk_indices,
BLOCK_SIZE=block_size,
BLOCK_STRIDE_ROWS=block_stride_rows,
NUM_TOPK_TOKENS=logical_topk_indices.shape[1],
return_valid_counts=return_valid_counts,
)
physical_topk_indices = self.physical_topk_indices[:num_tokens]
valid_topk_counts = self.valid_topk_counts[:num_tokens]
compute_stream = current_stream()
if layer_index == 0:
if self.has_indexer:
self.side_stream.wait_event(self.logical_topk_ready)
else:
self.side_stream.wait_stream(compute_stream)
with self.side_stream:
triton_convert_req_index_to_global_index(
req_id_per_token,
block_table,
logical_topk_indices,
BLOCK_SIZE=block_size,
BLOCK_STRIDE_ROWS=block_stride_rows,
NUM_TOPK_TOKENS=logical_topk_indices.shape[1],
return_valid_counts=return_valid_counts,
out=physical_topk_indices,
valid_counts_out=(
valid_topk_counts if return_valid_counts else None
),
)
self.physical_topk_ready.record(self.side_stream)
compute_stream.wait_event(self.physical_topk_ready)
if return_valid_counts:
return physical_topk_indices, valid_topk_counts
return physical_topk_indices
@dataclass
class HiSparseMLAIndexGroup(SparseMLAIndexGroup):
"""Sparse index group with per-layer host-backed KV caches."""
caches: list[HiSparseCacheHandle] = field(default_factory=list)
hisparse_group: Any | None = None
prefill_stream: torch.Stream | None = None
prefill_ready_events: list[torch.Event] = field(default_factory=list)
def register_layer(
self,
vllm_config: VllmConfig | None = None,
*,
head_size: int | None = None,
kv_cache_dtype: str | None = None,
) -> int:
assert vllm_config is not None
assert head_size is not None
assert kv_cache_dtype is not None
layer_index = super().register_layer()
if self.prefill_stream is None:
self.prefill_stream = _create_side_stream(self.logical_topk_indices.device)
if kv_cache_dtype == "fp8_ds_mla":
row_width = FP8_DS_MLA_ROW_BYTES
kv_dtype = torch.uint8
else:
from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
row_width = head_size
kv_dtype = kv_cache_dtype_str_to_dtype(
kv_cache_dtype, vllm_config.model_config
)
cache = create_hisparse_cache_handle(
vllm_config,
self.logical_topk_indices.shape[1],
is_index_group_leader=layer_index == 0,
row_width=row_width,
kv_dtype=kv_dtype,
index_group=self,
)
assert cache is not None
self.caches.append(cache)
cache.index_group_caches = self.caches
self.prefill_ready_events.append(_create_event())
return layer_index
def cache(self, layer_index: int) -> HiSparseCacheHandle:
return self.caches[layer_index]
def physical_kv_cache(self, layer_index: int) -> torch.Tensor:
cache = self.cache(layer_index)
return cache.runtime.hot.attention_cache
def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None:
if layer_index == 0:
self.cache(0).prepare_group_for_batch(attn_metadata)
def convert_logical_to_physical_topk(
self,
layer_index: int,
logical_topk_indices: torch.Tensor,
attn_metadata: Any,
*,
block_stride_rows: int | None,
return_valid_counts: bool,
req_id_per_token: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
cache = self.cache(layer_index)
num_tokens = logical_topk_indices.shape[0]
if req_id_per_token is None:
req_id_per_token = attn_metadata.req_id_per_token[:num_tokens]
if num_tokens > self.physical_topk_indices.shape[0]:
# Prefill-sized batches do not fit the decode residency workspace.
# Non-resident prefills are staged before reaching this path.
assert cache.all_context_pages_resident
leader = self.cache(0)
assert leader.view is not None and leader.block_table is not None
return self._convert_once(
layer_index,
logical_topk_indices,
req_id_per_token,
leader.block_table,
leader.view.block_size,
block_stride_rows=leader.view.attention_block_stride,
return_valid_counts=return_valid_counts,
)
source_block_table = cache.source_block_table
assert source_block_table is not None
return cache.swap_in(
req_id_per_token,
block_table=source_block_table,
logical_topk_indices=logical_topk_indices,
block_size=attn_metadata.block_size,
return_valid_counts=return_valid_counts,
)
def convert_decode_logical_to_physical_topk(
self,
layer_index: int,
logical_topk_indices: torch.Tensor,
attn_metadata: Any,
*,
return_valid_counts: bool,
num_decodes: int | None = None,
decode_query_len: int | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
num_tokens = logical_topk_indices.shape[0]
if num_decodes is None:
num_decodes = attn_metadata.num_decodes
if decode_query_len is None:
decode_query_len = attn_metadata.decode_max_query_len
if decode_query_len == 1:
return self.convert_logical_to_physical_topk(
layer_index,
logical_topk_indices,
attn_metadata,
block_stride_rows=None,
return_valid_counts=return_valid_counts,
req_id_per_token=self.request_ids[:num_decodes],
)
if num_tokens != num_decodes * decode_query_len:
query_start_loc = attn_metadata.query_start_loc[: num_decodes + 1]
physical_topk_indices = self.physical_topk_indices[: num_tokens + 1]
valid_topk_counts = self.valid_topk_counts[: num_tokens + 1]
request_ids = self.request_ids[:num_decodes]
cache = self.cache(layer_index)
source_block_table = cache.source_block_table
assert source_block_table is not None
for step in range(decode_query_len):
token_indices = query_start_loc[:-1] + step
active = (token_indices < query_start_loc[1:]) & (
token_indices < num_tokens
)
output_indices = torch.where(
active, token_indices, torch.full_like(token_indices, num_tokens)
).long()
token_indices = token_indices.clamp(0, num_tokens - 1).long()
step_topk = logical_topk_indices.index_select(
0, token_indices
).masked_fill(~active.unsqueeze(1), -1)
step_result = cache.swap_in(
request_ids,
block_table=source_block_table,
logical_topk_indices=step_topk,
block_size=attn_metadata.block_size,
return_valid_counts=return_valid_counts,
)
if layer_index == 0:
if return_valid_counts:
step_indices, step_counts = step_result
valid_topk_counts.index_copy_(0, output_indices, step_counts)
else:
step_indices = step_result
physical_topk_indices.index_copy_(0, output_indices, step_indices)
physical_topk_indices = physical_topk_indices[:num_tokens]
if return_valid_counts:
return physical_topk_indices, valid_topk_counts[:num_tokens]
return physical_topk_indices
assert num_tokens == num_decodes * decode_query_len
logical_topk_by_request = logical_topk_indices.view(
num_decodes, decode_query_len, -1
)
physical_topk_by_request = self.physical_topk_indices[:num_tokens].view(
num_decodes, decode_query_len, -1
)
valid_topk_by_request = self.valid_topk_counts[:num_tokens].view(
num_decodes, decode_query_len
)
request_ids = self.request_ids[:num_decodes]
for step in range(decode_query_len):
cache = self.cache(layer_index)
source_block_table = cache.source_block_table
assert source_block_table is not None
cache.swap_in(
request_ids,
block_table=source_block_table,
logical_topk_indices=logical_topk_by_request[:, step],
block_size=attn_metadata.block_size,
return_valid_counts=return_valid_counts,
attention_indices_out=(
physical_topk_by_request[:, step] if layer_index == 0 else None
),
valid_counts_out=(
valid_topk_by_request[:, step]
if layer_index == 0 and return_valid_counts
else None
),
)
physical_topk_indices = physical_topk_by_request.view(num_tokens, -1)
if return_valid_counts:
return physical_topk_indices, valid_topk_by_request.view(num_tokens)
return physical_topk_indices
def stage_prefill_rows(
self, layer_index: int, kv_cache: torch.Tensor, attn_metadata: Any
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
cache = self.cache(layer_index)
prefill = attn_metadata.prefill
staging_plan = prefill.host_staging_plan if prefill is not None else None
assert staging_plan is not None
resident_cache = None
if cache.view is not None and cache.block_table is not None:
staging_plan.ensure_gpu_sources(
cache.block_table[attn_metadata.num_decodes :],
cache.view.block_size,
)
resident_cache = cache.view.cache
staged_cache = cache.runtime.gather_prefill_cache(
kv_cache,
staging_plan,
resident_cache=resident_cache,
)
req_ids = attn_metadata.req_id_per_token[attn_metadata.num_decode_tokens :]
if attn_metadata.num_decodes > 0:
req_ids = req_ids - attn_metadata.num_decodes
return staged_cache, staging_plan.block_table, req_ids
def _prefill_gather_plan(
self, layer_index: int, attn_metadata: Any
) -> HiSparsePrefillStagingPlan:
cache = self.cache(layer_index)
prefill = attn_metadata.prefill
plan = prefill.host_staging_plan if prefill is not None else None
assert plan is not None
assert cache.view is not None and cache.block_table is not None
plan.ensure_gpu_sources(
cache.block_table[attn_metadata.num_decodes :],
cache.view.block_size,
)
assert plan.gpu_row_ids is not None
return plan
def gather_fp8_prefill(
self,
layer_index: int,
source_cache: torch.Tensor,
dst: torch.Tensor,
block_table: torch.Tensor,
workspace_starts: torch.Tensor,
batch_size: int,
attn_metadata: Any,
request_start: int,
) -> torch.Event | None:
cache = self.cache(layer_index)
plan = self._prefill_gather_plan(layer_index, attn_metadata)
assert cache.view is not None and plan.gpu_row_ids is not None
row_width = source_cache.shape[-1]
prefill_stream = self.prefill_stream
assert prefill_stream is not None
block_table = plan.block_table[request_start : request_start + batch_size]
prefill_stream.wait_stream(current_stream())
with prefill_stream:
ops.cp_gather_and_upconvert_fp8_kv_cache(
cache.view.cache,
dst,
block_table,
workspace_starts,
batch_size,
host_cache=source_cache.view(-1, row_width),
host_row_ids=plan.row_ids,
device_row_ids=plan.gpu_row_ids,
)
ready = self.prefill_ready_events[layer_index]
ready.record(prefill_stream)
return ready
class SparseMLAIndexGroupBuilder:
"""Assign sparse MLA layers to their index-producing layer."""
def __init__(
self, logical_topk_indices: torch.Tensor, max_decode_rows: int | None = None
) -> None:
self.logical_topk_indices = logical_topk_indices
self.max_decode_rows = (
logical_topk_indices.shape[0]
if max_decode_rows is None
else max_decode_rows
)
self.current_group: SparseMLAIndexGroup | None = None
def register_layer(
self,
is_index_producing_layer: bool,
vllm_config: VllmConfig | None = None,
*,
head_size: int | None = None,
kv_cache_dtype: str | None = None,
) -> tuple[SparseMLAIndexGroup, int]:
if is_index_producing_layer or self.current_group is None:
group_cls = (
HiSparseMLAIndexGroup
if vllm_config is not None
and vllm_config.attention_config.hisparse_config is not None
else SparseMLAIndexGroup
)
workspace_rows = self.max_decode_rows
physical_topk_indices = torch.empty(
(workspace_rows + 1, self.logical_topk_indices.shape[1]),
dtype=self.logical_topk_indices.dtype,
device=self.logical_topk_indices.device,
)
self.current_group = group_cls(
logical_topk_indices=self.logical_topk_indices,
physical_topk_indices=physical_topk_indices,
valid_topk_counts=torch.empty(
workspace_rows + 1,
dtype=torch.int32,
device=self.logical_topk_indices.device,
),
request_ids=torch.arange(
workspace_rows,
dtype=torch.int32,
device=self.logical_topk_indices.device,
),
side_stream=_create_side_stream(self.logical_topk_indices.device),
logical_topk_ready=_create_event(),
physical_topk_ready=_create_event(),
has_indexer=is_index_producing_layer,
)
group = self.current_group
group_index = group.register_layer(
vllm_config,
head_size=head_size,
kv_cache_dtype=kv_cache_dtype,
)
return group, group_index
def get_sparse_mla_index_group_max_rows(vllm_config: VllmConfig) -> int:
max_query_len = 1
speculative_config = vllm_config.speculative_config
if (
speculative_config is not None
and speculative_config.num_speculative_tokens is not None
):
max_query_len += speculative_config.num_speculative_tokens * (
2 if speculative_config.parallel_drafting else 1
)
scheduler_config = vllm_config.scheduler_config
return min(
scheduler_config.max_num_batched_tokens,
scheduler_config.max_num_seqs * max_query_len,
)
vllm/v1/attention/backends/mla/sparse_utils.py
+28 -6
View File
@@ -407,6 +407,8 @@ def triton_convert_req_index_to_global_index(
prefill_workspace_request_ids: torch.Tensor | None = None,
prefill_workspace_starts: torch.Tensor | None = None,
return_valid_counts: bool = False,
out: torch.Tensor | None = None,
valid_counts_out: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""
out[token_id, indice_id] =
@@ -462,16 +464,36 @@ def triton_convert_req_index_to_global_index(
# contiguous prefix [0, valid_count) and leaves the tail unwritten, so
# pre-fill -1 there. flash_mla_sparse_fwd then bounds attention to
# [:topk_length] == exactly the valid set (no dropped tokens).
if return_valid_counts:
out = torch.full_like(token_indices_c, -1)
if out is None:
out = (
torch.full_like(token_indices_c, -1)
if return_valid_counts
else torch.empty_like(token_indices_c)
)
else:
out = torch.empty_like(token_indices_c)
assert out.dtype == token_indices_c.dtype
assert out.device == token_indices_c.device
assert out.shape == token_indices_c.shape
assert out.is_contiguous()
if return_valid_counts:
out.fill_(-1)
valid_counts: torch.Tensor | None = None
if return_valid_counts:
# Zero-init only matters for the atomic accumulation path.
alloc = torch.empty if single_tile else torch.zeros
valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)
if valid_counts_out is None:
# Zero-init only matters for the atomic accumulation path.
alloc = torch.empty if single_tile else torch.zeros
valid_counts = alloc(
num_tokens, dtype=torch.int32, device=token_indices.device
)
else:
assert valid_counts_out.dtype == torch.int32
assert valid_counts_out.device == token_indices.device
assert valid_counts_out.shape == (num_tokens,)
assert valid_counts_out.is_contiguous()
valid_counts = valid_counts_out
if not single_tile:
valid_counts.zero_()
# Prepare prefill pointers
if HAS_PREFILL_WORKSPACE:
vllm/v1/core/block_pool.py
+51 -5
View File
@@ -1,6 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable, Sequence
from collections.abc import Callable, Iterable, Sequence
from typing import Any
from vllm.distributed.kv_events import (
@@ -157,6 +157,7 @@ class BlockPool:
actual block size can be a multiple of hash_block_size.
enable_kv_cache_events: Whether to enable kv cache events.
metrics_collector: Optional metrics collector for tracking block residency.
medium: Storage medium reported in KV cache events.
"""
def __init__(
@@ -166,14 +167,16 @@ class BlockPool:
hash_block_size: int,
enable_kv_cache_events: bool = False,
metrics_collector: KVCacheMetricsCollector | None = None,
medium: str = MEDIUM_GPU,
):
assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0
self.num_gpu_blocks = num_gpu_blocks
self.medium = medium
self.enable_caching = enable_caching
self.hash_block_size = hash_block_size
# All kv-cache blocks.
self.blocks: list[KVCacheBlock] = [
KVCacheBlock(idx) for idx in range(num_gpu_blocks)
KVCacheBlock(idx, pool=self) for idx in range(num_gpu_blocks)
]
# Free block queue that constructs and manipulates a doubly linked
# list of free blocks (including eviction candidates when caching is
@@ -194,6 +197,9 @@ class BlockPool:
self.kv_event_queue: list[KVCacheEvent] = []
self.metrics_collector = metrics_collector
# Callbacks for blocks released with ``unpin_blocks`` whose contents
# are still being read until the pool reuses them.
self._reuse_watchers: dict[int, Callable[[KVCacheBlock], None]] = {}
def get_cached_block(
self, block_hash: BlockHash, kv_cache_group_ids: list[int]
@@ -364,7 +370,7 @@ class BlockPool:
token_ids=request.all_token_ids[start_token_idx:end_token_idx],
block_size=block_size,
lora_id=request.lora_request.adapter_id if request.lora_request else None,
medium=MEDIUM_GPU,
medium=self.medium,
lora_name=request.lora_request.name if request.lora_request else None,
extra_keys=extra_keys_list if extra_keys_list else None,
group_idx=kv_cache_group_id,
@@ -543,7 +549,7 @@ class BlockPool:
lora_id=request.lora_request.adapter_id
if request.lora_request
else None,
medium=MEDIUM_GPU,
medium=self.medium,
lora_name=request.lora_request.name
if request.lora_request
else None,
@@ -610,7 +616,7 @@ class BlockPool:
self.kv_event_queue.append(
BlockRemoved(
block_hashes=[maybe_convert_block_hash(get_block_hash(block_hash))],
medium=MEDIUM_GPU,
medium=self.medium,
group_idx=get_group_id(block_hash),
)
)
@@ -671,6 +677,9 @@ class BlockPool:
ret: list[KVCacheBlock] = self.free_block_queue.popleft_n(num_blocks)
if self._reuse_watchers:
self._notify_reuse(ret)
# In order to only iterate the list once, we duplicated code a bit
if self.enable_caching:
for block in ret:
@@ -687,6 +696,33 @@ class BlockPool:
self.metrics_collector.on_block_allocated(block)
return ret
def _notify_reuse(self, blocks: Iterable[KVCacheBlock]) -> None:
for block in blocks:
watcher = self._reuse_watchers.pop(block.block_id, None)
if watcher is not None:
watcher(block)
def unpin_blocks(
self,
blocks: Iterable[KVCacheBlock],
on_reuse: Callable[[KVCacheBlock], None],
) -> None:
"""Release references to blocks that stay readable until reused.
The blocks become last-resort eviction candidates regardless of prefix
caching: they join the tail of the free queue and count as free.
``on_reuse`` fires when ``get_new_blocks`` hands a block out (or the
cache is reset), so the caller can stop reading it.
"""
released: list[KVCacheBlock] = []
for block in blocks:
assert block.ref_cnt > 0 and not block.is_null
block.ref_cnt -= 1
self._reuse_watchers[block.block_id] = on_reuse
if block.ref_cnt == 0:
released.append(block)
self.free_block_queue.append_n(released)
def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool:
"""
If a block is cached in `cached_block_hash_to_block`, we reset its hash
@@ -742,7 +778,11 @@ class BlockPool:
# Identify blocks with hash (LRU cache) and without it (never match APC)
blocks_to_evict_last = []
blocks_to_evict_first = []
other_pools: dict[BlockPool, list[KVCacheBlock]] = {}
for block in ordered_blocks:
if block.pool is not None and block.pool is not self:
other_pools.setdefault(block.pool, []).append(block)
continue
block.ref_cnt -= 1
if block.ref_cnt == 0 and not block.is_null:
if block.block_hash is None or not self.enable_caching:
@@ -756,6 +796,8 @@ class BlockPool:
self.free_block_queue.prepend_n(blocks_to_evict_first)
# Blocks to reuse last are appended to the end of the free queue.
self.free_block_queue.append_n(blocks_to_evict_last)
for pool, blocks in other_pools.items():
pool.free_blocks(blocks)
def evict_blocks(self, block_ids: set[int]) -> None:
"""evict blocks from the prefix cache by their block IDs.
@@ -797,6 +839,10 @@ class BlockPool:
# Remove all hashes so that no new blocks will hit.
self.cached_block_hash_to_block = BlockHashToBlockMap()
self.cached_block_hashes_by_block.clear()
if self._reuse_watchers:
self._notify_reuse(
[self.blocks[block_id] for block_id in list(self._reuse_watchers)]
)
# Remove all hashes from all blocks.
for block in self.blocks:

Some files were not shown because too many files have changed in this diff Show More