|
|
0f539e4855
|
Flush compressor: schema fix, prepare_forward, flush_write kernels, state rotation
Schema fix (paper eq.11-12):
CSA needs m entries for current a-stream AND m entries for previous
b-stream (tail_buffer_size_a=4, tail_buffer_size_b=4). After flush,
current a-stream becomes next flush b-stream input.
HCA: tail_buffer_size_a=128, tail_buffer_size_b=0 (no b-stream).
tail_zb initialized to -1e9 so softmax naturally masks b-stream on
first flush (paper: Z^b padded with -inf, C^b with zeros).
prepare_forward.py:
Runs between captured graphs. Computes new compressed entries from
position delta, pre-allocates blocks before the graph runs.
Deterministic: entries_after - entries_before, ceil to block boundary.
No allocation inside the captured graph.
flush_write.cu — 4 kernels:
flush_write_csa_kernel: BF16 -> FP8 E4M3 quantize + scatter compressed
entry + FP4 NVFP4 indexer key write (16-element groups, E4M3 scale).
One block per request, 128 threads. Amax reduction -> inv_scale.
flush_write_hca_kernel: same minus indexer (no FP4 write).
csa_rotate_state_kernel: after CSA flush, rotate a->b stream,
clear a-stream, reset tail_len.
hca_reset_state_kernel: after HCA flush, clear a-stream, reset tail_len.
flush.py: Python orchestration.
maybe_flush_csa/hca: always runs, kernels gate via valid_mask.
Compressor produces entry, flush kernel quantize-scatters, state
kernel rotates/resets. No host-side branching for cudagraph.
All tests pass on B200:
Schema: CSA tail_a=4 tail_b=4, HCA tail_a=128 tail_b=0
State: tail_zb initialized to -1e9, reset_slot preserves it
prepare_forward: correct block allocation for position transitions
HCA flush write: RoPE exact, FP8 <3.6% error, invalid mask no-op
CSA flush write: RoPE exact, indexer FP4 keys written
CSA state rotation: kb<-ka, zb<-za, ka/za zeroed, tail_len=0
HCA state reset: ka/za zeroed, tail_len=0
|
2026-05-22 00:25:47 +00:00 |
|
|
|
b4d58df620
|
KV Cache: schema, allocator, pools, manager, append_swa kernel
Complete KV cache substrate for DSV4 inference:
schema.py: Per-layer cache shape derived from LayerSpec.
- CSA: 32 entries/block, 32 indexer entries, tail=3
- HCA: 1 entry/block, no indexer, tail=127
- SWA: no classical pool, no tail
- BLOCK_SIZE_ORIGINAL_TOKENS=128 (lcm of compression ratios)
- compute_block_budget() for allocator sizing
allocator.py: Fixed-size block free-list.
- GPU stack with pinned host top pointer
- acquire/release between graph captures only
- OOM raises on exhaustion
paged_cache.py: Per-layer classical KV storage.
- FP8 (uint8) for non-RoPE dims, BF16 for RoPE dims (paper 2.3.4)
- Per-entry inverse scale for FP8 dequant
- FP4 indexer keys for CSA layers (NVFP4 scheme)
- memory_bytes() tracking
state_cache.py: Per-layer SWA window + tail buffer.
- Ring buffer with position tracking (swa_head, swa_pos)
- CSA: dual streams (ka/za/kb/zb) for overlapping compression
- HCA: single stream (ka/za only)
- SWA: no tail buffer
- reset_slot() for request completion
handle.py: LayerCacheHandle — typed per-call view.
- write_swa(), read_swa_view(), read_classical_view(), read_indexer_view()
- No GPU allocation in acquire() — 0 bytes delta (cudagraph safe)
- SWAView/ClassicalView/IndexerView dataclasses for kernel signatures
manager.py: KVCacheManager — owns everything.
- Per-layer schema, pool, and allocator construction
- admit_request()/release_request() lifecycle
- allocate_block() for compression flush
- acquire() returns LayerCacheHandle (zero-alloc)
append_swa.cu: Native kernel for SWA writes.
- One block per token, 128 threads per block
- Warp-level amax reduction, BF16->FP8 E4M3 quantization
- Atomic ring buffer head increment
- FP8/BF16 split write + inv_scale + position metadata
- FP8 round-trip: <3.6% relative error
- RoPE half: exact match (no quantization)
All tests pass on B200:
- Schema correctness for CSA/HCA/SWA
- Allocator acquire/release/OOM
- Pool shapes match architecture spec
- Manager lifecycle (admit/release/recycle/exhaustion)
- Zero-alloc acquire() (cudagraph safe)
- append_swa kernel: positions, RoPE exact, FP8 quality, wrap-around, multi-request isolation
|
2026-05-22 00:08:38 +00:00 |
|
|
|
3fb3c925af
|
Restructure: cutedsl/ -> dsv4/ with proper layering
- Split bridge.py -> ops/quantize.py, ops/layouts.py, ops/gemm_runner.py
- Renamed classes: CuTeDSLNvfp4Linear -> Nvfp4Linear, etc.
- Moved kernel code to dsv4/kernels/ (gemm, attention, compressor, decode, cuda)
- Moved PyTorch bridges to dsv4/ops/
- Moved nn.Module layers to dsv4layers/
- Moved reference implementations to dsv4/reference/
- Moved vendored CUTLASS code to vendored/
- Archived ~190 debug tests to tests/archive/
- Kept ~15 canonical tests in tests/unit/
- Updated all import paths
- Added stubs for future components (model/, cache/, loader/)
- Updated pyproject.toml: dsv4-inference package name
|
2026-05-21 17:30:44 +00:00 |
|
