7ef6402936
KV-1/KV-2/KV-3: NVFP4 compressed KV + FP8 indexer keys
...
Architecture:
- Compressed KV: stored as NVFP4 (E2M1 + E4M3 + FP32 gsa)
- Write path: compress→FP32 → FP32 RoPE → quantize FP32→NVFP4
- Read path: dequant_nvfp4/dequant_nvfp4_selective → BF16 for FMHA
- No BF16 intermediate in the write path
- Indexer keys: stored as FP8_E4M3 (1 byte + per-row scale)
- Write path: compress→FP32 → quantize FP32→FP8_E4M3
- Read path: dequant_fp8_e4m3 → BF16 for scoring
- SWA: remains BF16 (8MB total, fits in L2)
New kernels in kv_quantize.cu:
- compute_amax_gsa_fp32: per-row gsa from FP32 input
- quantize_nvfp4_from_fp32: FP32→NVFP4 with GPU gsa buffer
- quantize_fp8_e4m3_from_fp32: FP32→FP8_E4M3 for indexer keys
- dequant_fp8_e4m3 / dequant_fp8_e4m3_selective: FP8→BF16
- rope_fp32: FP32 GPT-J interleaved RoPE (no BF16)
Proven two-kernel pattern (same as quantize_nvfp4_gpu_fused):
Kernel 1: amax_gsa (GPU-only)
Kernel 2: quantize from buffer (GPU gsa)
No shared memory bugs. No cross-CTA race conditions.
KVCache updated:
- comp_kv_fp4/sf/gsa: NVFP4 storage (3.5× smaller than BF16)
- comp_idx_fp8/scale: FP8_E4M3 storage (1.9× smaller than BF16)
- comp_kv property: dequant NVFP4→BF16 on demand
- comp_kv_selective: dequant only top-k entries (bandwidth savings)
- comp_idx_kv property: dequant FP8→BF16 on demand
Removed: compressor_reduce_quant.cu (buggy single-kernel approach)
2026-06-02 10:00:50 +00:00
f23320b5b2
KV-1/KV-2: Fused compress+NVFP4 quantize kernels + dequant
...
- compressor_reduce_quant.cu: Single-kernel CSA/HCA compress + RMSNorm + NVFP4 quantize.
No intermediate BF16. FP32 → E2M1 + E4M3 + FP32 gsa in one kernel.
Shared memory: ~2.5KB per CTA (FP32 staging + nibble buffer).
- dequant_nvfp4.cu: NVFP4 → BF16 dequantization kernels.
Full dequant (HCA dense gather) and selective dequant (CSA top-k gather).
Single kernel launch per gather operation.
- production_compress.py: Added csa_compress_production_nvfp4() and
hca_compress_production_nvfp4() — production path for KV-1/KV-2.
- loader.py: Preload dequant_nvfp4 and compressor_reduce_quant modules.
- test_kv_compress_quant.py: Unit tests verifying cos >= 0.999
between BF16 reference and NVFP4 round-trip path.
2026-06-02 09:37:53 +00:00
62041b78bf
fix: import torch.utils.cpp_extension explicitly in production_compress
2026-06-01 05:20:44 +00:00
b380028c49
feat: production compressor/indexer — NVFP4 GEMM + CUDA softmax/reduce kernel
...
- New compressor_reduce.cu: CSA/HCA token-level softmax + weighted sum + kv_norm
One block per compressed entry, 128 threads, FP32 accumulation
CSA: overlapping Ca/Cb streams (2m tokens per block)
HCA: single stream (m tokens per block)
Includes apply_kv_norm kernel (unweighted RMSNorm + weight)
- New production_compress.py: Python wrapper for CUDA kernels
- single_shot_inference.py: Compressor/Indexer now use production Nvfp4Linear
for kv_proj, gate_proj, q_b_proj, weights_proj projections
Then CUDA reduce kernel for softmax + weighted sum
No more PyTorch reference nvfp4_linear_ref in compressor/indexer path
2026-06-01 05:18:59 +00:00
9d88769f5f
Wire indexer compute_index_scores_topk + fix compressor imports
...
- indexer/__init__.py: compute_index_scores_topk now calls
run_indexer_score_topk with proper tensor reshaping
- compressor/__init__.py: added torch import, fixed csa_compress_tail
and hca_compress_tail imports for flush.py
- Full flush pipeline now importable end-to-end
2026-05-30 21:19:06 +00:00
daf84524ac
E2/E3: compressor bridge, indexer bridge, flush pipeline wiring
...
- compress_tail.py: PyTorch reference CSA/HCA compression
(token-level softmax over m/m' entries, paper eq. 11-12)
- compressor/__init__.py: csa_compress_and_store, hca_compress_and_store
bridges (compression deferred to flush pipeline)
- indexer/__init__.py: compute_index_scores_topk bridge (NotImplemented)
- Fixed attention.py: removed extra positions arg to write_swa
2026-05-30 21:16:54 +00:00
300dddedc0
E1-E4: gather kernels, handle wiring, rope, sync removal, e2e test
...
E1: LayerCacheHandle now exposes gather_compressed_kv,
gather_all_compressed_kv, gather_swa_kv, num_query_heads, head_dim.
Gather kernels in dsv4/kernels/cuda/gather_swa.cu + gather_kv.cu.
Python wrapper in dsv4/kernels/cache/gather.py.
E2: tests/e2e/test_one_layer.py — SWA path smoke test.
E3: Compressor/indexer __init__.py bridges (NotImplementedError stubs
for CSA/HCA compress_and_store, compute_index_scores_topk).
E4: Removed torch.cuda.synchronize() from fmha_multitile_op.py fast path.
Error checking via C API return code instead.
Also: forward_rope_partial in ops/rope.py (GPT-J interleaved, last 64 dims).
2026-05-30 21:10:26 +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
