|
|
f3b551956d
|
Cleanup Step 2: Archive Lineage P code, fix broken imports
- 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)
|
2026-06-02 19:27:07 +00:00 |
|
|
|
faf92b30ad
|
E1: Wire LayerCacheHandle gather methods + CUDA gather kernels
- gather_compressed_kv: CSA top-k gather via existing gather_kv.cu
- gather_all_compressed_kv: HCA dense gather via new gather_all_compressed_kernel
- gather_swa_kv: SWA ring buffer gather via new gather_swa_kernel
- Added gather_swa.cu with both SWA + all-compressed gather kernels
- Added gather.py Python wrapper (torch.utils.cpp_extension JIT)
- Updated handle.py: added schema field, num_query_heads/head_dim properties
- Updated manager.py: passes schema + num_query_heads to handle
All gather kernels: FP8→BF16 dequant + BF16 RoPE concat in single launch.
Output: dense BF16 tensors ready for FMHA consumption.
|
2026-05-30 21:09:21 +00:00 |
|
|
|
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 |
|
