biondizzle
300dddedc0
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).
61 lines
2.4 KiB
Python
61 lines
2.4 KiB
Python
"""CSA/HCA compressor — Python API bridge.
|
|
|
|
Wraps the CuTeDSL compressor kernels with the interface that
|
|
AttentionSubBlock expects. The compressor itself is CuTeDSL because
|
|
it doesn't have the FMHA pipeline constraints — pure elementwise
|
|
softmax over m entries, no tensor cores needed.
|
|
|
|
The long-term path is raw CUDA C++ per doctrine, but the CuTeDSL
|
|
compressor is already working and correct. Rewrite only if MLIR
|
|
compilation becomes a blocker.
|
|
"""
|
|
from dsv4.kernels.compressor.csa_hca import (
|
|
launch_csa_compress_projected,
|
|
launch_hca_compress_projected,
|
|
)
|
|
|
|
|
|
def csa_compress_and_store(
|
|
kv_raw: "torch.Tensor", # (T, 4 * head_dim) BF16 — (Ca, Cb, Za, Zb) interleaved
|
|
cache: "LayerCacheHandle", # writes compressed entries to paged pool
|
|
positions: "torch.Tensor", # (T,) int64
|
|
compression_ratio: int = 4, # m=4
|
|
) -> None:
|
|
"""CSA: compress KV entries and store into the classical paged cache.
|
|
|
|
The compressor reads the (Ca, Cb, Za, Zb) streams from kv_raw,
|
|
runs token-level softmax compression (paper eq. 11-12), and writes
|
|
the compressed entries to the cache's paged pool.
|
|
|
|
The b-stream from the previous flush is read from the state cache's
|
|
tail buffer (tail_kb, tail_zb).
|
|
"""
|
|
# TODO: implement the full CSA compression + store path.
|
|
# For now, this is a placeholder that writes raw KV to the tail buffer.
|
|
# The full path needs:
|
|
# 1. Read prev b-stream from state cache
|
|
# 2. Run CuTeDSL compression kernel
|
|
# 3. Write compressed output to paged pool via flush kernel
|
|
# 4. Update tail buffer (a-stream becomes next b-stream)
|
|
raise NotImplementedError(
|
|
"CSA compress_and_store requires the full flush pipeline. "
|
|
"See dsv4/kernels/cuda/flush_write.cu and dsv4/cache/flush.py"
|
|
)
|
|
|
|
|
|
def hca_compress_and_store(
|
|
kv_raw: "torch.Tensor", # (T, 2 * head_dim) BF16 — (C, Z) interleaved
|
|
cache: "LayerCacheHandle", # writes compressed entries to paged pool
|
|
positions: "torch.Tensor", # (T,) int64
|
|
compression_ratio: int = 128, # m'=128
|
|
) -> None:
|
|
"""HCA: compress KV entries and store into the classical paged cache.
|
|
|
|
Same structure as CSA but no b-stream, no overlap, and m'=128
|
|
means compression only fires once per 128 tokens.
|
|
"""
|
|
raise NotImplementedError(
|
|
"HCA compress_and_store requires the full flush pipeline. "
|
|
"See dsv4/kernels/cuda/flush_write.cu and dsv4/cache/flush.py"
|
|
)
|