Repository navigation
Description
proposing an opt-in circular storage path for the shared streaming attention cache, starting with the Wan recipe's unmasked attention and keys whose RoPE is applied before the cache write. the aim is to remove retained-history copies and preserve the same sink, window and denoising updates.
current path
at 9d98934d9652, BlockKVCache._roll_local_window_left clones and copies the retained keys and values whenever a new chunk overflows the window. for a saturated window of W tokens and a chunk of C tokens, that moves the W - C retained tokens in each key/value buffer once per autoregressive step and layer. this is logical copy work, with actual HBM traffic depending on cache residency and the generated kernels.
Wan21TransformerCache.start owns that advance outside the captured network, and the network overwrites the current chunk throughout denoising. the default Self-Forcing configuration is one concrete consumer with a 21-frame window and three-frame chunks.
proposed boundary
keep storage and chunk lifecycle in core.attention, reuse the fixed-slot storage concepts already present in FixedSlotKVCache, and let the attention owner select whether physical-order reads are valid. slot reservation must stay separate from writes so denoising and finalization keep replacing the same current chunk.
for unmasked dense attention, applying the same permutation to key and value rows preserves the mathematical result. a full circular buffer can therefore go directly to the existing attention backend, with no chronological gather or split-attention requirement. absolute positions determine RoPE before storage, and the sink remains pinned.
- start with whole-chunk sink and window sizes. preserve filling, repeated writes, wraparound and reset semantics, plus independent conditional and unconditional state.
- retain stable K/V allocations and stable device-resident write metadata. captured kernels must read the changing slot from device data, since capturing a Python slice offset would keep writing the same slot on replay.
- keep the existing chronological path for cache-relative RoPE, masks, unsupported layouts and any backend or context-parallel configuration that has not passed parity. the existing public chronological-view contract stays intact.
this is separate from the head-sharded mirror and incremental Ulysses exchange proposed in #739. this local-storage proposal leaves inter-rank traffic and the model's retained history unchanged.
first experiment
a CPU fp64 storage prototype compared current BlockKVCache and circular storage against an independent chronological-history reference across 48 aligned layouts, repeated writes and resets. all 6,912 comparisons retained identical keys and values in logical order, and physical-order dense attention differed by at most 1.34e-15. this covers storage semantics and the permutation property. GPU integration and speed remain unmeasured.
the implementation gate would be matched current/circular runs on one existing Self-Forcing or LingBot preset, then the supported multi-GPU path, covering eager, compile, graph replay, the first eviction, multiple wraps and reset. collect cache-maintenance time, total chunk latency, peak allocated memory and HBM/L2 traffic, with the same weights, noise and control inputs. keep the current path if metadata or write-kernel cost erases the end-to-end gain, or output differences exceed the model's accepted quality comparison.
would this fit the shared cache boundary, and is there existing circular-cache work to build on before implementing the opt-in path?
Activity