biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Cleanup Step 2: Archive Lineage P code, fix broken imports

Browse Source 
- Move dead dsv4/ modules to dsv4/_archive/ (52 files)
  - model/{dsv4,mtp,layer,layer_schedule}
  - layers/{embedding,attention,ffn,norm} (kept linear,mhc,router,moe,shared_expert,grouped_linear - live)
  - cache/*, kernels/cache/*, kernels/indexer/{csa_indexer,score_topk,compute_valid_lens}
  - kernels/router/{nvfp4_fused_router,dense_router_decode_kernel,dense_router_prefill}
  - ops/{topk,topk_select,rope,router}, loader/{hf_checkpoint,layout_convert}
  - reference/{attention,compressor,csa_attention,moe_pipeline}
  - kernels/compressor/{compress_tail,csa_hca}
- Restore dsv4/ops/{router,custom_ops}.py (needed by live layers)
- Fix dsv4/kernels/{indexer,compressor,attention}/__init__.py (removed broken imports)
- Remove preload_all() from loader.py (dead, referenced nonexistent .cu file)
- Fix loader.py docstring (fused_amax_quantize_nvfp4 → quantize_nvfp4_from_buffer)
- Move broken tests to tests/e2e_archive/
  - test_fused_router, production_values_test, e2e/{one_layer,model_construction,csa_hca}
- vLLM has 0 imports of dsv4 (Step 0 confirmed)
This commit is contained in:
biondizzle
2026-06-02 19:27:07 +00:00
parent 8de47e26ce
commit f3b551956d
55 changed files with 2881 additions and 306 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
dsv4/cache/allocator.py vendored
View File

@@ -1,56 +0,0 @@
"""Fixed-size block allocator for the classical paged KV cache.
One BlockAllocator per layer per "pool kind" (classical / indexer).
Total blocks are sized at engine startup. Blocks are recycled on
request completion.
Cudagraph-safety: allocation can't happen inside a captured graph
(allocation rate is per-request not per-token). The contract is:
- acquire() called between graph captures.
- release() called between graph captures.
- read access (via block table) happens INSIDE captured graphs.
"""
from __future__ import annotations
import torch
class BlockAllocator:
def __init__(
self,
num_total_blocks: int,
device: str = "cuda",
):
self.num_total_blocks = num_total_blocks
self.device = device
# Free-list as a GPU stack: ids[0..top-1] holds free block IDs.
# `top` lives in pinned host memory so we can read it without a
# device sync (it's modified only between graph captures).
self.free_ids = torch.arange(
num_total_blocks, dtype=torch.int32, device=device,
)
self.top_cpu = torch.tensor([num_total_blocks], dtype=torch.int32, pin_memory=True)
@property
def num_free(self) -> int:
return int(self.top_cpu[0])
def acquire(self, n: int) -> torch.Tensor:
"""Return a tensor of `n` block IDs. Called between captures."""
top = int(self.top_cpu[0])
if n > top:
raise RuntimeError(
f"KV cache OOM: requested {n} blocks, {top} available "
f"(of {self.num_total_blocks} total)"
)
new_top = top - n
ids = self.free_ids[new_top:top].clone() # snapshot
self.top_cpu[0] = new_top
return ids
def release(self, ids: torch.Tensor) -> None:
"""Return blocks to the free list. Called between captures."""
n = ids.numel()
top = int(self.top_cpu[0])
self.free_ids[top:top + n] = ids.to(device=self.device)
self.top_cpu[0] = top + n

2
dsv4/cache/block_table.py vendored
View File

@@ -1,2 +0,0 @@
"""Block table for paged KV cache."""
# TODO: Phase 3

162
dsv4/cache/flush.py vendored
View File

@@ -1,162 +0,0 @@
"""In-graph flush orchestration.
Called when tail_len crosses the compression threshold. The actual
compression math is in the csa_hca_compressor kernel; this module
handles the quantize-scatter-write step and the state rotation.
The maybe_flush_* functions always run when their attention type
matches — no host-side `if tail_full` check. The kernels gate
internally via `valid_mask` computed from `tail_len`. This keeps
the call sequence identical across forward passes for cudagraph.
"""
from __future__ import annotations
from typing import Optional
import os
import torch
from torch.utils.cpp_extension import load
from dsv4.cache.schema import LayerCacheSchema, AttentionType
_flush_mod = None
def _get_flush_module():
global _flush_mod
if _flush_mod is not None:
return _flush_mod
kernel_dir = os.path.join(os.path.dirname(__file__), "..", "kernels", "cuda")
_flush_mod = load(
name="flush_write",
sources=[os.path.join(kernel_dir, "flush_write.cu")],
extra_cuda_cflags=["-O3", "--generate-code=arch=compute_100a,code=[sm_100a]"],
verbose=False,
)
return _flush_mod
def maybe_flush_csa(
handle,
schema: LayerCacheSchema,
m: int,
) -> None:
"""For CSA: emit compressed entries for requests whose tail is full.
Steps:
1. Determine which requests have tail_len >= m (valid_mask).
2. Run the CSA compressor on tail buffers.
3. Scatter compressed entry + indexer key into paged pool.
4. Rotate a-stream -> b-stream, clear a-stream.
"""
from dsv4.kernels.compressor import csa_compress_tail
state = handle.state
paged = handle.paged
mod = _get_flush_module()
# Step 1: valid_mask — which requests have a full tail buffer.
# tail_len is [max_requests], request_slots is [B].
tail_lens = state.tail_len[handle.request_slots] # [B]
valid_mask = tail_lens >= m # [B] bool
# If no requests need flushing, short-circuit.
if not valid_mask.any().item():
return
# Step 2: compress the tail.
# The compressor kernel takes the tail buffers and produces
# one compressed entry per request (for those where valid_mask=True).
entry, indexer_key = csa_compress_tail(
tail_ka=state.tail_ka,
tail_za=state.tail_za,
tail_kb=state.tail_kb,
tail_zb=state.tail_zb,
tail_len=state.tail_len,
request_slots=handle.request_slots,
m=m,
)
# entry: [B, head_dim] BF16
# indexer_key: [B, indexer_head_dim] BF16
# Step 3: scatter into the paged pool.
# The flush position for each request = the position of the last
# token in the tail (positions before this forward minus 1 would
# be the wrong reference; we need the tail's last position).
# For the block table lookup, we use the compressed entry index
# derived from positions.
# Use the positions of the requests' current tokens to figure
# out which entry slot to write into.
flush_positions = handle.positions # [tokens] -> need per-request
# For now, derive entry index from the per-request state:
# compressed_entry_idx = sum of all flushes so far for this request.
# This is (positions_of_last_appended_token) // m
# Simplification: use request_slots to look up per-request position.
# The handle's positions are per-token, not per-request.
# We need one position per request = position of the last appended token.
# For a single-token decode, that's just positions[-1] per request.
# For a general case, take the max position per request.
# This is computed by the append kernel (stored in tail_len and the
# actual positions in the tail). For now, use handle.positions
# and scatter by request.
# The kernel resolves slot_in_block from positions internally.
mod.flush_write_csa(
entry, indexer_key, valid_mask, handle.request_slots,
handle.positions[:handle.request_slots.shape[0]], # one pos per request
handle.block_table,
paged.entries_fp8, paged.entries_rope, paged.inv_scale,
paged.indexer_keys_fp4, paged.indexer_scale,
schema.entries_per_block, m, schema.rope_dim,
schema.entry_head_dim, schema.indexer_head_dim,
)
# Step 4: rotate state — current a-stream becomes next b-stream.
mod.csa_rotate_state(
valid_mask, handle.request_slots,
state.tail_ka, state.tail_za, state.tail_kb, state.tail_zb,
state.tail_len, m, schema.entry_head_dim,
)
def maybe_flush_hca(
handle,
schema: LayerCacheSchema,
m_prime: int,
) -> None:
"""For HCA: emit one entry per request whose tail_len >= m'."""
from dsv4.kernels.compressor import hca_compress_tail
state = handle.state
paged = handle.paged
mod = _get_flush_module()
tail_lens = state.tail_len[handle.request_slots]
valid_mask = tail_lens >= m_prime
if not valid_mask.any().item():
return
entry = hca_compress_tail(
tail_ka=state.tail_ka,
tail_za=state.tail_za,
tail_len=state.tail_len,
request_slots=handle.request_slots,
m=m_prime,
)
# entry: [B, head_dim] BF16
mod.flush_write_hca(
entry, valid_mask, handle.request_slots,
handle.positions[:handle.request_slots.shape[0]],
handle.block_table,
paged.entries_fp8, paged.entries_rope, paged.inv_scale,
schema.entries_per_block, m_prime, schema.rope_dim,
schema.entry_head_dim,
)
# Reset tail — no b-stream rotation for HCA.
mod.hca_reset_state(
valid_mask, handle.request_slots,
state.tail_ka, state.tail_za, state.tail_len,
m_prime, schema.entry_head_dim,
)

286
dsv4/cache/handle.py vendored
View File

@@ -1,286 +0,0 @@
"""LayerCacheHandle — typed per-call view onto one layer's cache.
Constructed by KVCacheManager.acquire() once per layer per forward.
Holds tensor references and integer indices; no allocation. Methods
expose the operations AttentionSubBlock needs without exposing the
underlying storage layout.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional, TYPE_CHECKING
import torch
if TYPE_CHECKING:
from dsv4.cache.paged_cache import PagedKVPool
from dsv4.cache.state_cache import StateCachePool
from dsv4.cache.schema import LayerCacheSchema
@dataclass
class LayerCacheHandle:
"""Read/write interface for one layer's cache.
The fields are the resolved indices and tensor refs for THIS call's
batch of requests. AttentionSubBlock never sees raw pool tensors.
"""
# Pool references (shared across handles — never mutated).
paged: Optional["PagedKVPool"]
state: "StateCachePool"
schema: "LayerCacheSchema"
# Per-call indices.
request_slots: torch.Tensor # [batch] int32 — state cache slot per request
positions: torch.Tensor # [tokens] int32 — absolute position per token
request_ids: torch.Tensor # [tokens] int32 — which request each token belongs to
# Block table for the classical pool (None for SWA-only layers).
# Shape: [batch, max_logical_blocks] int32. -1 padding for unused entries.
block_table: Optional[torch.Tensor]
# Number of valid blocks per request (excludes padding).
block_lens: Optional[torch.Tensor]
# ------------------------------------------------------------------
# Properties called by AttentionSubBlock
# ------------------------------------------------------------------
@property
def num_query_heads(self) -> int:
"""Number of query heads (from schema)."""
# The schema doesn't store n_q directly — derive from the config.
# For now, store on the handle at construction.
return self._num_query_heads
@num_query_heads.setter
def num_query_heads(self, value: int):
self._num_query_heads = value
@property
def head_dim(self) -> int:
"""Head dimension (from schema)."""
return self.schema.entry_head_dim
# ------------------------------------------------------------------
# Methods called by AttentionSubBlock
# ------------------------------------------------------------------
def write_swa(
self,
raw_kv: torch.Tensor, # (T, head_dim) BF16
) -> None:
"""Write raw KV into the SWA ring buffer AND tail compression buffer.
Both regions get the same tokens — SWA consumes the last n_win,
the tail accumulates until it can flush.
"""
from dsv4.kernels.cache.append_swa import append_swa_kernel
append_swa_kernel(
raw_kv=raw_kv,
request_slots=self.request_slots,
positions=self.positions,
swa_fp8=self.state.swa_fp8,
swa_rope=self.state.swa_rope,
swa_inv=self.state.swa_inv,
swa_pos=self.state.swa_pos,
swa_head=self.state.swa_head,
rope_dim=self.schema.rope_dim,
)
def flush_compression(
self,
compressed: torch.Tensor, # (T_flush, head_dim) BF16 — newly produced
indexer_keys: Optional[torch.Tensor] = None,
) -> None:
"""Promote pending tail tokens into the classical pool.
Called by the compressor when the tail buffer has enough tokens.
Allocates a new block if the latest block is full.
Block allocation requires going outside the captured graph — in
a fully-captured decode this is rare (once per m or m' tokens),
so we make it explicit. The manager has the contract.
"""
raise NotImplementedError("see kernels/cache/flush_compression.py")
def gather_compressed_kv(
self,
selected_indices: torch.Tensor, # (T, top_k) int64 — from indexer
) -> tuple[torch.Tensor, torch.Tensor]:
"""CSA: gather top-k compressed KV entries into dense BF16 tensors.
Returns:
(k_compressed, v_compressed) each of shape (1, n_comp, head_dim) BF16.
The leading dim=1 is for the single KV head (MQA in DSV4).
"""
assert self.paged is not None, "CSA gather requires paged pool"
from dsv4.kernels.cache.gather import gather_compressed_kv
hd = self.head_dim
rd = self.schema.rope_dim
epb = self.schema.entries_per_block
# selected_indices is int64, gather kernel needs int32
indices_i32 = selected_indices.to(torch.int32)
# block_table for CSA: [batch, max_logical_blocks]
# For per-request gather, use the first request's block_table
# (decode: batch=1, so this is trivial)
if self.block_table.dim() == 1:
bt = self.block_table.unsqueeze(0)
else:
bt = self.block_table
k_out = gather_compressed_kv(
entries_fp8=self.paged.entries_fp8,
entries_rope=self.paged.entries_rope,
inv_scale=self.paged.inv_scale,
topk_indices=indices_i32,
block_table=bt,
entries_per_block=epb,
head_dim=hd,
rope_dim=rd,
)
# k_out: (T, top_k, hd) — for FMHA we need (1, n_comp, hd)
# At decode T=1: squeeze to (top_k, hd) then unsqueeze for KV head dim
n_comp = k_out.shape[1]
k_compressed = k_out.squeeze(0).unsqueeze(0) # (1, n_comp, hd)
# V shares the same storage but is transposed — DSV4 uses K=V for
# the compressed KV (same entries, different projection weights applied
# before compression). For now, return the same gathered tensor.
# TODO: verify if K and V are stored separately or shared.
v_compressed = k_compressed.clone()
return k_compressed, v_compressed
def gather_all_compressed_kv(self) -> tuple[torch.Tensor, torch.Tensor]:
"""HCA: gather ALL compressed KV entries into dense BF16 tensors.
No indexer — dense attention over the short compressed sequence.
Returns:
(k_compressed, v_compressed) each of shape (1, n_comp, head_dim) BF16.
"""
assert self.paged is not None, "HCA gather requires paged pool"
from dsv4.kernels.cache.gather import gather_all_compressed_kv
hd = self.head_dim
rd = self.schema.rope_dim
epb = self.schema.entries_per_block
if self.block_table.dim() == 1:
bt = self.block_table.unsqueeze(0)
bl = self.block_lens.unsqueeze(0) if self.block_lens is not None else None
else:
bt = self.block_table
bl = self.block_lens
if bl is None:
# Default: all blocks valid
bl = torch.full((bt.shape[0],), bt.shape[1], dtype=torch.int32, device=bt.device)
k_out = gather_all_compressed_kv(
entries_fp8=self.paged.entries_fp8,
entries_rope=self.paged.entries_rope,
inv_scale=self.paged.inv_scale,
block_table=bt,
block_lens=bl,
entries_per_block=epb,
head_dim=hd,
rope_dim=rd,
)
# k_out: (batch, total_entries, hd) — for FMHA we need (1, n_comp, hd)
n_comp = k_out.shape[1]
k_compressed = k_out.squeeze(0).unsqueeze(0) # (1, n_comp, hd)
v_compressed = k_compressed.clone()
return k_compressed, v_compressed
def gather_swa_kv(self) -> tuple[torch.Tensor, torch.Tensor]:
"""Gather SWA window entries into dense BF16 tensors.
Returns:
(k_swa, v_swa) each of shape (1, swa_len, head_dim) BF16.
"""
from dsv4.kernels.cache.gather import gather_swa_kv
hd = self.head_dim
rd = self.schema.rope_dim
k_out = gather_swa_kv(
swa_fp8=self.state.swa_fp8,
swa_rope=self.state.swa_rope,
swa_inv=self.state.swa_inv,
swa_pos=self.state.swa_pos,
request_slots=self.request_slots,
head_dim=hd,
rope_dim=rd,
)
# k_out: (batch, n_win, hd) — for FMHA we need (1, swa_len, hd)
k_swa = k_out.squeeze(0).unsqueeze(0) # (1, swa_len, hd)
v_swa = k_swa.clone()
return k_swa, v_swa
def read_swa_view(self) -> "SWAView":
"""Return a typed view of the SWA window for this batch."""
return SWAView(
fp8=self.state.swa_fp8,
rope=self.state.swa_rope,
inv_scale=self.state.swa_inv,
positions=self.state.swa_pos,
head=self.state.swa_head,
slots=self.request_slots,
)
def read_classical_view(self) -> "ClassicalView":
"""Return a typed view of compressed entries for this batch."""
assert self.paged is not None, "SWA-only layers have no classical cache"
return ClassicalView(
entries_fp8=self.paged.entries_fp8,
entries_rope=self.paged.entries_rope,
inv_scale=self.paged.inv_scale,
block_table=self.block_table,
block_lens=self.block_lens,
)
def read_indexer_view(self) -> "IndexerView":
"""CSA-only. Returns FP4 indexer keys with their scales."""
assert self.paged is not None and self.paged.indexer_keys_fp4 is not None
return IndexerView(
keys_fp4=self.paged.indexer_keys_fp4,
scale=self.paged.indexer_scale,
global_scale=self.paged.indexer_global_scale,
block_table=self.block_table,
block_lens=self.block_lens,
)
def __post_init__(self):
# Initialize _num_query_heads (must be set by the manager at construction)
if not hasattr(self, '_num_query_heads'):
self._num_query_heads = 0
# Typed views — simple dataclasses, no logic. The FMHA / indexer / SWA
# kernels accept these to keep their signatures clean.
@dataclass
class SWAView:
fp8: torch.Tensor
rope: torch.Tensor
inv_scale: torch.Tensor
positions: torch.Tensor
head: torch.Tensor
slots: torch.Tensor
@dataclass
class ClassicalView:
entries_fp8: torch.Tensor
entries_rope: torch.Tensor
inv_scale: torch.Tensor
block_table: torch.Tensor
block_lens: torch.Tensor
@dataclass
class IndexerView:
keys_fp4: torch.Tensor
scale: torch.Tensor
global_scale: torch.Tensor
block_table: torch.Tensor
block_lens: torch.Tensor

203
dsv4/cache/manager.py vendored
View File

@@ -1,203 +0,0 @@
"""KVCacheManager — owns all KV cache state for one model instance.
Responsibilities:
- Build per-layer pools and allocators at startup.
- Hand out state-cache slots when requests are admitted.
- Hand out classical blocks when layers need to flush compression.
- Compose LayerCacheHandle for each layer per forward call.
- Reclaim slots and blocks on request completion.
Not on the manager:
- On-disk prefix storage. (Paper §3.5.2 — deferred entirely.)
- Eviction policies. (Single-instance; requests run to completion.)
- Cross-instance coordination.
"""
from __future__ import annotations
from typing import List, Optional, Dict
import torch
from dsv4.model.config import DSV4Config
from dsv4.model.layer_schedule import LayerSpec, AttentionType
from dsv4.cache.schema import LayerCacheSchema, build_schema, compute_block_budget
from dsv4.cache.allocator import BlockAllocator
from dsv4.cache.paged_cache import PagedKVPool
from dsv4.cache.state_cache import StateCachePool
from dsv4.cache.handle import LayerCacheHandle
class KVCacheManager:
def __init__(
self,
config: DSV4Config,
schedule: List[LayerSpec],
max_concurrent_requests: int,
max_context_tokens: int = 1_000_000,
# Per-layer-type block budget. If None, computed from
# max_context_tokens and max_concurrent_requests.
num_blocks_per_csa_layer: Optional[int] = None,
num_blocks_per_hca_layer: Optional[int] = None,
device: str = "cuda",
):
self.config = config
self.schedule = schedule
self.max_concurrent_requests = max_concurrent_requests
self.device = device
# ---- Per-layer schemas ----
self.schemas: Dict[int, LayerCacheSchema] = {
spec.layer_idx: build_schema(config, spec) for spec in schedule
}
# ---- Compute block budgets if not provided ----
if num_blocks_per_csa_layer is None or num_blocks_per_hca_layer is None:
budget = compute_block_budget(config, schedule, max_context_tokens,
max_concurrent_requests)
num_blocks_per_csa_layer = num_blocks_per_csa_layer or budget.get("csa", 0)
num_blocks_per_hca_layer = num_blocks_per_hca_layer or budget.get("hca", 0)
# ---- Per-layer pools ----
# State cache exists for every layer.
self.state_pools: Dict[int, StateCachePool] = {
i: StateCachePool(schema, max_concurrent_requests, device)
for i, schema in self.schemas.items()
}
# Classical paged pool only for compressed layers.
self.paged_pools: Dict[int, Optional[PagedKVPool]] = {}
self.allocators: Dict[int, Optional[BlockAllocator]] = {}
for i, schema in self.schemas.items():
if schema.entries_per_block == 0:
self.paged_pools[i] = None
self.allocators[i] = None
else:
nb = (num_blocks_per_csa_layer
if schema.attn_type == AttentionType.CSA
else num_blocks_per_hca_layer)
self.paged_pools[i] = PagedKVPool(schema, nb, device)
self.allocators[i] = BlockAllocator(nb, device)
# ---- Request state ----
# Slot index per request, into state cache pools (same index in
# every layer). -1 = slot free.
self.request_slot_map: torch.Tensor = torch.full(
(max_concurrent_requests,), -1, dtype=torch.int32, device=device,
)
# Block table per request per layer:
# block_tables[layer_idx][request_slot, logical_block_idx]
# -> physical_block_idx
max_blocks = max_context_tokens // 128 # BLOCK_SIZE_ORIGINAL_TOKENS
self.max_blocks_per_request = max_blocks
self.block_tables: Dict[int, torch.Tensor] = {}
self.block_lens: Dict[int, torch.Tensor] = {}
for i, schema in self.schemas.items():
if schema.entries_per_block > 0:
self.block_tables[i] = torch.full(
(max_concurrent_requests, max_blocks), -1,
dtype=torch.int32, device=device,
)
self.block_lens[i] = torch.zeros(
(max_concurrent_requests,), dtype=torch.int32, device=device,
)
# ------------------------------------------------------------------
# Request lifecycle (called between captured graphs)
# ------------------------------------------------------------------
def admit_request(self) -> int:
"""Allocate a state cache slot. Returns the slot index."""
free = (self.request_slot_map == -1).nonzero(as_tuple=False)
if free.numel() == 0:
raise RuntimeError("max concurrent requests exceeded")
slot = int(free[0])
self.request_slot_map[slot] = slot
return slot
def release_request(self, slot: int) -> None:
"""Return state cache slot and all associated blocks to free lists."""
for layer_idx, alloc in self.allocators.items():
if alloc is None:
continue
table = self.block_tables[layer_idx]
lens = self.block_lens[layer_idx]
valid = int(lens[slot])
if valid > 0:
alloc.release(table[slot, :valid].clone())
lens[slot] = 0
table[slot].fill_(-1)
# Reset state cache slot.
for state in self.state_pools.values():
state.reset_slot(slot)
self.request_slot_map[slot] = -1
# ------------------------------------------------------------------
# Block allocation for compression flush (called between captures)
# ------------------------------------------------------------------
def allocate_block(self, layer_idx: int, request_slot: int) -> int:
"""Allocate one new classical block for a request. Returns block ID."""
alloc = self.allocators[layer_idx]
assert alloc is not None, f"layer {layer_idx} has no classical pool"
block_id = alloc.acquire(1)
bid = int(block_id[0])
# Append to the request's block table.
table = self.block_tables[layer_idx]
lens = self.block_lens[layer_idx]
pos = int(lens[request_slot])
assert pos < self.max_blocks_per_request, "block table overflow"
table[request_slot, pos] = bid
lens[request_slot] = pos + 1
return bid
# ------------------------------------------------------------------
# Per-forward handle construction (called INSIDE captured graph)
# ------------------------------------------------------------------
def acquire(
self,
layer_idx: int,
request_slots: torch.Tensor, # [batch] int32
positions: torch.Tensor, # [tokens] int32
request_ids: torch.Tensor, # [tokens] int32
) -> LayerCacheHandle:
"""Build the LayerCacheHandle for one layer's forward.
No allocation happens here — critical for cudagraph safety.
"""
paged = self.paged_pools[layer_idx]
state = self.state_pools[layer_idx]
if paged is not None:
# Pass the full tensors — no indexing, no allocation.
# The attention kernel indexes by request_slots internally.
block_table = self.block_tables[layer_idx]
block_lens = self.block_lens[layer_idx]
else:
block_table = None
block_lens = None
handle = LayerCacheHandle(
paged=paged,
state=state,
schema=self.schemas[layer_idx],
request_slots=request_slots,
positions=positions,
request_ids=request_ids,
block_table=block_table,
block_lens=block_lens,
)
handle.num_query_heads = self.config.num_query_heads
return handle
# ------------------------------------------------------------------
# Diagnostics
# ------------------------------------------------------------------
def memory_bytes(self) -> int:
"""Total GPU memory used by all pools."""
total = 0
for pool in self.state_pools.values():
total += pool.memory_bytes()
for pool in self.paged_pools.values():
if pool is not None:
total += pool.memory_bytes()
for i, table in self.block_tables.items():
total += table.numel() * table.element_size()
total += self.block_lens[i].numel() * self.block_lens[i].element_size()
return total

91
dsv4/cache/paged_cache.py vendored
View File

@@ -1,91 +0,0 @@
"""Storage for one layer's classical paged KV cache.
Layout per block:
entries: [num_blocks, entries_per_block, head_dim - rope_dim] FP8 (uint8 view)
entries_r: [num_blocks, entries_per_block, rope_dim] BF16
inv_scale: [num_blocks, entries_per_block] FP32
The FP8/BF16 split mirrors paper §2.3.4 ("BF16 for RoPE dims, FP8 for
the rest"). The kernel reads both halves and concatenates in registers.
For CSA layers, a parallel pool stores indexer keys at the same block
granularity — same block ID maps to a block in both pools.
"""
from __future__ import annotations
from typing import Optional
import torch
from dsv4.cache.schema import LayerCacheSchema
class PagedKVPool:
"""Per-layer classical paged KV storage.
Indexed by [physical_block_id, slot_in_block, ...].
Both compressed entries and indexer keys (if applicable) are
indexed by the SAME physical_block_id so a CSA layer's two pools
share the block table.
"""
def __init__(
self,
schema: LayerCacheSchema,
num_blocks: int,
device: str = "cuda",
):
self.schema = schema
self.num_blocks = num_blocks
self.device = device
nb = num_blocks
epb = schema.entries_per_block
hd = schema.entry_head_dim
rd = schema.rope_dim
fp8_dim = hd - rd
# ---- Compressed entries ----
# FP8 stored as uint8 (we view as float8_e4m3fn at read time).
self.entries_fp8 = torch.zeros(
(nb, epb, fp8_dim), dtype=torch.uint8, device=device,
)
# BF16 RoPE'd half — no quantization.
self.entries_rope = torch.zeros(
(nb, epb, rd), dtype=torch.bfloat16, device=device,
)
# Per-entry inverse scale (for FP8 dequant in attention kernel).
self.inv_scale = torch.ones(
(nb, epb), dtype=torch.float32, device=device,
)
# ---- Indexer keys (CSA only) ----
if schema.indexer_entries_per_block > 0:
i_epb = schema.indexer_entries_per_block
i_hd = schema.indexer_head_dim
# Indexer QK is FP4 per paper §2.3.4 — but we store the keys
# post-quant. uint8 = 2 FP4 packed per byte.
self.indexer_keys_fp4 = torch.zeros(
(nb, i_epb, i_hd // 2), dtype=torch.uint8, device=device,
)
# Per-block-vector scale for the FP4 (one E4M3 scalar per
# 16-element group, per the NVFP4 quantization scheme).
self.indexer_scale = torch.ones(
(nb, i_epb, i_hd // 16),
dtype=torch.float8_e4m3fn, device=device,
)
self.indexer_global_scale = torch.ones(
(nb,), dtype=torch.float32, device=device,
)
else:
self.indexer_keys_fp4 = None
self.indexer_scale = None
self.indexer_global_scale = None
def memory_bytes(self) -> int:
"""Total GPU memory used by this pool."""
total = 0
for name in ("entries_fp8", "entries_rope", "inv_scale",
"indexer_keys_fp4", "indexer_scale", "indexer_global_scale"):
t = getattr(self, name)
if t is not None:
total += t.numel() * t.element_size()
return total

83
dsv4/cache/prepare_forward.py vendored
View File

@@ -1,83 +0,0 @@
"""Pre-forward block allocation.
Runs between captured graphs. Computes how many new compressed entries
will be produced by this forward (deterministic from positions), allocates
the required physical blocks, and updates block tables.
After this runs, the captured graph can perform flushes by writing to
already-resolved (request, layer, logical_block) -> physical_block
mappings. No allocation inside the graph.
"""
from __future__ import annotations
from typing import List
import torch
from dsv4.model.layer_schedule import LayerSpec, AttentionType
from dsv4.cache.manager import KVCacheManager
def prepare_forward(
manager: KVCacheManager,
request_slots: torch.Tensor, # [B] state cache slots
positions_before: torch.Tensor, # [B] absolute position BEFORE this forward
positions_after: torch.Tensor, # [B] absolute position AFTER this forward
) -> None:
"""Pre-allocate any blocks that will be needed by flushes in this forward.
Pure CPU/GPU bookkeeping — runs between captures, not in hot path.
For each compressed layer, works out how many flushes happen per
request and allocates blocks to cover them.
"""
for layer_idx, spec in enumerate(manager.schedule):
if spec.attn == AttentionType.SWA:
continue # No classical pool, no flushes.
schema = manager.schemas[layer_idx]
alloc = manager.allocators[layer_idx]
if alloc is None:
continue
m = (manager.config.csa_compression_ratio
if spec.attn == AttentionType.CSA
else manager.config.hca_compression_ratio)
epb = schema.entries_per_block
# How many compressed entries are NEWLY produced per request?
# = floor(positions_after / m) - floor(positions_before / m)
entries_after = (positions_after // m).to(torch.int64)
entries_before = (positions_before // m).to(torch.int64)
new_entries = entries_after - entries_before # [B] int64
# For each request, figure out how many new blocks are needed.
# A block holds `epb` entries. If there are already some entries
# in the current (open) block, they take some slots.
for b in range(request_slots.numel()):
n_new = int(new_entries[b])
if n_new == 0:
continue
req_slot = int(request_slots[b])
# How many entries are already in the current open block?
existing_blocks = int(manager.block_lens[layer_idx][req_slot])
entries_in_open_block = int(entries_before[b]) % epb if existing_blocks > 0 else 0
slots_remaining_in_open = epb - entries_in_open_block if entries_in_open_block > 0 else 0
# How many new blocks do we need?
if entries_in_open_block == 0 and existing_blocks == 0:
# Fresh — no open block yet
blocks_needed = (n_new + epb - 1) // epb
elif slots_remaining_in_open >= n_new:
# Fits in the current open block
blocks_needed = 0
else:
# Need additional blocks beyond the current open one
overflow = n_new - slots_remaining_in_open
blocks_needed = (overflow + epb - 1) // epb
if blocks_needed == 0:
continue
ids = alloc.acquire(blocks_needed)
existing = int(manager.block_lens[layer_idx][req_slot])
manager.block_tables[layer_idx][req_slot, existing:existing + blocks_needed] = ids
manager.block_lens[layer_idx][req_slot] = existing + blocks_needed

125
dsv4/cache/schema.py vendored
View File

@@ -1,125 +0,0 @@
"""Per-layer KV cache shape.
Computed once per layer at engine startup from the LayerSpec. The
schema is what tells the allocator how big each pool slot is and what
sub-regions exist (compressed entries / indexer keys / SWA window /
uncompressed tail).
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
from dsv4.model.config import DSV4Config
from dsv4.model.layer_schedule import LayerSpec, AttentionType
# Block size is invariant for DSV4 — derived from compression ratios.
# lcm(m, m') = lcm(4, 128) = 128 original tokens per block.
# Holds 128/4 = 32 CSA entries OR 128/128 = 1 HCA entry per block.
BLOCK_SIZE_ORIGINAL_TOKENS = 128
@dataclass(frozen=True)
class LayerCacheSchema:
"""Cache layout for one transformer layer.
Fields with `_per_block` are the dimensions of one block in the
classical paged pool. `_per_state_slot` are dimensions of one
request's slot in the state cache.
All sizes are in number of entries — bytes come from the dtypes.
"""
layer_idx: int
attn_type: AttentionType
# ---- Classical paged cache (compressed entries) ----
entries_per_block: int
entry_head_dim: int
rope_dim: int
# ---- Indexer pool (CSA only) ----
indexer_entries_per_block: int
indexer_head_dim: int
# ---- State cache (SWA window + uncompressed tail) ----
swa_window_size: int
# CSA: paper eq.11-12, the i-th flush uses Ca[m*i:m*(i+1)] and
# Cb[m*(i-1):m*i]. After flush, current a-stream becomes next b-stream.
# So we need m entries for current a-stream AND m entries for previous
# b-stream. Total tail = 2*m for CSA.
tail_buffer_size_a: int # m (CSA) or m' (HCA) — current tokens
tail_buffer_size_b: int # m (CSA only) — previous block's a-stream kept as b-input
# Per-token inverse scale storage (for FP8 dequant).
needs_inv_scale: bool = True
@property
def tail_buffer_size(self) -> int:
"""Total tail entries (for backward compat with schema consumers)."""
return self.tail_buffer_size_a + self.tail_buffer_size_b
def build_schema(config: DSV4Config, spec: LayerSpec) -> LayerCacheSchema:
"""Derive cache schema for a single layer from architectural config."""
if spec.attn == AttentionType.CSA:
return LayerCacheSchema(
layer_idx=spec.layer_idx,
attn_type=AttentionType.CSA,
entries_per_block=BLOCK_SIZE_ORIGINAL_TOKENS // config.csa_compression_ratio,
entry_head_dim=config.head_dim,
rope_dim=config.rope_dim,
indexer_entries_per_block=BLOCK_SIZE_ORIGINAL_TOKENS // config.csa_compression_ratio,
indexer_head_dim=config.indexer_head_dim,
swa_window_size=config.sliding_window,
tail_buffer_size_a=config.csa_compression_ratio, # m=4 current
tail_buffer_size_b=config.csa_compression_ratio, # m=4 previous (b-stream)
)
elif spec.attn == AttentionType.HCA:
return LayerCacheSchema(
layer_idx=spec.layer_idx,
attn_type=AttentionType.HCA,
entries_per_block=BLOCK_SIZE_ORIGINAL_TOKENS // config.hca_compression_ratio,
entry_head_dim=config.head_dim,
rope_dim=config.rope_dim,
indexer_entries_per_block=0,
indexer_head_dim=0,
swa_window_size=config.sliding_window,
tail_buffer_size_a=config.hca_compression_ratio, # m'=128 current
tail_buffer_size_b=0, # HCA has no b-stream
)
else: # SWA-only
return LayerCacheSchema(
layer_idx=spec.layer_idx,
attn_type=AttentionType.SWA,
entries_per_block=0,
entry_head_dim=config.head_dim,
rope_dim=config.rope_dim,
indexer_entries_per_block=0,
indexer_head_dim=0,
swa_window_size=config.sliding_window,
tail_buffer_size_a=0,
tail_buffer_size_b=0,
)
def compute_block_budget(
config: DSV4Config,
schedule: list[LayerSpec],
max_context_tokens: int,
max_concurrent_requests: int,
) -> dict[str, int]:
"""Compute per-layer-type block counts for the allocator."""
blocks_per_request = max_context_tokens // BLOCK_SIZE_ORIGINAL_TOKENS
headroom = 1.10
result = {}
for spec in schedule:
if spec.attn == AttentionType.CSA:
key = "csa"
elif spec.attn == AttentionType.HCA:
key = "hca"
else:
continue
total = int(max_concurrent_requests * blocks_per_request * headroom)
result[key] = max(result.get(key, 0), total)
return result

102
dsv4/cache/state_cache.py vendored
View File

@@ -1,102 +0,0 @@
"""State cache: SWA window + uncompressed tail buffer.
One slot per active request. Slot index is fixed for a request's
lifetime — the manager hands out slot indices at request admission
and reclaims them at completion.
Per paper §3.5.1: SWA and tail tokens are state-space-like — they
depend only on the current position, not on a paged history. No
block table; a flat [max_requests, ...] tensor.
CSA b-stream lifecycle (paper eq.11-12):
After a CSA flush, the current a-stream (tail_ka/tail_za) becomes
the next flush's b-stream input (tail_kb/tail_zb). Both are sized
at m entries, not m-1. On first flush, tail_zb is filled with -1e9
so the softmax in the compressor naturally masks out the b-stream
(exp(-inf) = 0).
"""
from __future__ import annotations
import torch
from dsv4.cache.schema import LayerCacheSchema, AttentionType
class StateCachePool:
"""Per-layer state cache (SWA window + uncompressed tail).
Storage layout per slot:
swa_fp8: [n_win, head_dim - rope_dim] FP8 raw KV in window
swa_rope: [n_win, rope_dim] BF16 RoPE'd half
swa_inv: [n_win] FP32 per-token inv scale
swa_pos: [n_win] int32 — absolute position
swa_head: scalar int32 — ring buffer write head
tail_ka: [m_a, head_dim] BF16 — current a-stream tokens
tail_za: [m_a, head_dim] BF16 — current a-stream Z weights
tail_kb: [m_b, head_dim] BF16 — previous a-stream kept as b-input (CSA only)
tail_zb: [m_b, head_dim] BF16 — previous Z b-stream (CSA only, init to -1e9)
tail_len: scalar int32 — how many entries in a-stream are valid
"""
def __init__(
self,
schema: LayerCacheSchema,
max_requests: int,
device: str = "cuda",
):
self.schema = schema
self.max_requests = max_requests
self.device = device
mr = max_requests
nw = schema.swa_window_size
hd = schema.entry_head_dim
rd = schema.rope_dim
fp8 = hd - rd
# SWA window — circular within each slot.
self.swa_fp8 = torch.zeros((mr, nw, fp8), dtype=torch.uint8, device=device)
self.swa_rope = torch.zeros((mr, nw, rd), dtype=torch.bfloat16, device=device)
self.swa_inv = torch.ones((mr, nw), dtype=torch.float32, device=device)
self.swa_pos = torch.full((mr, nw), -1, dtype=torch.int32, device=device)
self.swa_head = torch.zeros((mr,), dtype=torch.int32, device=device)
# Tail buffer — only for compressed layers.
m_a = schema.tail_buffer_size_a # m (CSA) or m' (HCA)
m_b = schema.tail_buffer_size_b # m (CSA only)
if m_a > 0:
self.tail_ka = torch.zeros((mr, m_a, hd), dtype=torch.bfloat16, device=device)
self.tail_za = torch.zeros((mr, m_a, hd), dtype=torch.bfloat16, device=device)
self.tail_len = torch.zeros((mr,), dtype=torch.int32, device=device)
if m_b > 0: # CSA: need b-stream
self.tail_kb = torch.zeros((mr, m_b, hd), dtype=torch.bfloat16, device=device)
# Paper §3.5.1: Z^b padded with -inf at first flush.
# Init to -1e9 so softmax naturally masks b-stream on first flush.
self.tail_zb = torch.full((mr, m_b, hd), -1e9, dtype=torch.bfloat16, device=device)
else:
self.tail_kb = None
self.tail_zb = None
else:
self.tail_ka = self.tail_za = None
self.tail_kb = self.tail_zb = None
self.tail_len = None
def reset_slot(self, slot: int) -> None:
"""Clear a request's state after completion."""
self.swa_pos[slot].fill_(-1)
self.swa_head[slot] = 0
if self.tail_len is not None:
self.tail_len[slot] = 0
# Re-init tail_zb to -1e9 for CSA (paper §3.5.1 first-flush mask)
if self.tail_zb is not None:
self.tail_zb[slot].fill_(-1e9)
def memory_bytes(self) -> int:
"""Total GPU memory used by this pool."""
total = 0
for name in ("swa_fp8", "swa_rope", "swa_inv", "swa_pos", "swa_head",
"tail_ka", "tail_za", "tail_kb", "tail_zb", "tail_len"):
t = getattr(self, name)
if t is not None:
total += t.numel() * t.element_size()
return total