Add KVPress support to KerasHub - #2941
Conversation
Introduces keras_hub.press with a KVCachePress base class plus KnormPress, RandomPress, and StreamingLLMPress implementations for evicting less-important tokens from a KV cache to bound memory usage during generation.
CausalLM.compile() takes a new `press` argument, and a `_compress_layer_cache()` helper lets subclasses compress each layer's KV cache right after it is seeded during prefill, one layer at a time, rather than only after the whole cache has been assembled.
GPT2CausalLM.call_with_cache() now runs each transformer layer's freshly-seeded cache through CausalLM._compress_layer_cache() during prefill, so a `press` passed to compile() actually shrinks the prompt cache before generation starts.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces KV cache compression to KerasHub, adding a base KVCachePress class and concrete strategies (KnormPress, RandomPress, and StreamingLLMPress). It integrates this compression mechanism into CausalLM tasks, specifically updating GPT2CausalLM and LlamaCausalLM to decouple physical cache indices from logical position indices for correct position and rotary embeddings. The review feedback focuses on enhancing the robustness of the implementation under JAX/JIT compilation and tracing, specifically by using ops.shape for dynamic dimensions, avoiding symbolic batch sizes in ops.reshape, and optimizing memory usage on accelerators by casting indices to int32 instead of int64.
| reserve_shape = list(compressed.shape) | ||
| reserve_shape[3] = reserve | ||
| compressed = ops.concatenate( | ||
| [compressed, ops.zeros(reserve_shape, dtype=compressed.dtype)], | ||
| axis=3, | ||
| ) |
There was a problem hiding this comment.
Using list(compressed.shape) to construct reserve_shape is fragile because the batch dimension can be dynamic (i.e., None) during tracing or compilation. Passing a shape containing None to ops.zeros will raise a ValueError. Instead, construct the shape dynamically using ops.shape(compressed) to ensure robustness across all backends.
| reserve_shape = list(compressed.shape) | |
| reserve_shape[3] = reserve | |
| compressed = ops.concatenate( | |
| [compressed, ops.zeros(reserve_shape, dtype=compressed.dtype)], | |
| axis=3, | |
| ) | |
| compressed_shape = ops.shape(compressed) | |
| reserve_shape = ( | |
| compressed_shape[0], | |
| compressed_shape[1], | |
| compressed_shape[2], | |
| reserve, | |
| compressed_shape[4], | |
| compressed_shape[5], | |
| ) | |
| compressed = ops.concatenate( | |
| [compressed, ops.zeros(reserve_shape, dtype=compressed.dtype)], | |
| axis=3, | |
| ) |
| batch_size, num_layers = cache.shape[0], cache.shape[1] | ||
| num_heads = cache.shape[4] | ||
| keep_len = keep_indices.shape[-1] |
There was a problem hiding this comment.
Accessing dimensions directly via cache.shape is fragile when the batch dimension is dynamic (i.e., None) during tracing or compilation. Use ops.shape(cache) to dynamically retrieve the shape, ensuring robustness across all backends.
| batch_size, num_layers = cache.shape[0], cache.shape[1] | |
| num_heads = cache.shape[4] | |
| keep_len = keep_indices.shape[-1] | |
| cache_shape = ops.shape(cache) | |
| batch_size = cache_shape[0] | |
| num_layers = cache_shape[1] | |
| num_heads = cache_shape[4] | |
| keep_len = ops.shape(keep_indices)[-1] |
| indices = ops.cast(indices, "int64") | ||
| indices = ops.broadcast_to( | ||
| indices, (batch_size, num_layers, 2, keep_len, num_heads, 1) | ||
| ) |
There was a problem hiding this comment.
Casting indices to int64 is less efficient and consumes more memory on accelerators. Casting to int32 is sufficient for sequence length indices and reduces memory overhead.
| indices = ops.cast(indices, "int64") | |
| indices = ops.broadcast_to( | |
| indices, (batch_size, num_layers, 2, keep_len, num_heads, 1) | |
| ) | |
| indices = ops.cast(indices, "int32") | |
| indices = ops.broadcast_to( | |
| indices, (batch_size, num_layers, 2, keep_len, num_heads, 1) | |
| ) |
| shape = ops.shape(keys) | ||
| batch_size, num_layers, seq_len, num_heads = ( | ||
| shape[0], | ||
| shape[1], | ||
| shape[2], | ||
| shape[3], | ||
| ) |
There was a problem hiding this comment.
Using ops.shape(keys) to retrieve static dimensions like num_layers, seq_len, and num_heads can cause tracing issues in JAX/JIT when passed to ops.reshape. Retrieve the dynamic batch size using ops.shape(keys)[0] and the static dimensions using keys.shape to ensure compatibility with JIT compilation.
batch_size = ops.shape(keys)[0]
num_layers = keys.shape[1]
seq_len = keys.shape[2]
num_heads = keys.shape[3]| ops.ones_like(is_kept, dtype="float32"), | ||
| ops.full_like(ops.cast(is_kept, "float32"), _MASKED_SCORE), | ||
| ) | ||
| score = ops.reshape(score, (batch_size, 1, 1, seq_len)) |
There was a problem hiding this comment.
Avoid passing a symbolic batch_size tensor to ops.reshape as it can cause tracing issues under JAX/JIT. Use -1 for the dynamic batch dimension instead.
| score = ops.reshape(score, (batch_size, 1, 1, seq_len)) | |
| score = ops.reshape(score, (-1, 1, 1, seq_len)) |
Wires the same per-layer compression path used by GPT2CausalLM into LlamaCausalLM.call_with_cache(). LlamaAttention and LlamaTransformerDecoder gain a `rotary_start_index` argument so that rotary position embeddings still line up correctly once a layer's cache has been shrunk by a press.
Adds the two slot-mapping assertions and the test_kv_cache_compression_retains_prompt_through_decoding regression test that gpt2_causal_lm_test.py already had, so LlamaCausalLM gets the same coverage for compressed-cache decoding.
Adds keras_hub/src/kv_press/serialization_test.py, mirroring samplers/serialization_test.py, to cover keras_hub.press.get()'s string, dict, and instance identifier paths. Adds a test_serialization to each concrete press's own test file via the existing run_serialization_test() helper, covering that press's get_config()/from_config() and keras.saving round trip.
64fcb83 to
e30ced8
Compare
1 participant
Description of the change
Adds
keras_hub.press, KVPress support for KerasHub — a lightweight port of the eviction ideas behind [NVIDIA's KVPress (https://github.com/NVIDIA/kvpress), raised in #2075. Duringgenerate(), aKVCachePressscores every cached token and evicts the lowest-scoring ones right after each transformer layer's KV cache is produced, before the next layer runs. Compressing per layer, rather than once over the fully assembled, all-layers cache, bounds the transient memory overlap between a layer's old and compressed cache to a single layer instead of paying that cost across the whole stack.CausalLM.compile()gets a newpressargument that any subclass can opt into. This PR wires it up forGPT2CausalLMandLlamaCausalLM.Example Usage
Available presses
KnormPressRandomPressStreamingLLMPressn_sink"sink" tokens plus the most recent tokensCurrent limitation
Compression during
generate()requires the PyTorch backend. Sizing the compressed buffer correctly depends on the real (non-padding) prompt length, which is only known at runtime, but the buffer shape must be static at trace time underjax.jit/XLA-compiledtf.functions. Rather than silently fall back to the padded length (which would let the decode loop overwrite the very tokens that were just retained),compress()raisesNotImplementedErroron JAX/TensorFlow whenever apadding_maskis supplied.compression_ratio=0.0stays a true no-op on every backend.Testing
KVCachePressbase + built-in pressesget_config/from_configandkeras.savinground-tripsGPT2CausalLM/LlamaCausalLMcompression_ratio=0.0matching an uncompressed baseline exactly25 new tests total. Ran under
KERAS_BACKEND=torch; all pass.Colab Reference
This notebook adapts NVIDIA's
speed_and_memorybenchmark notebook. It confirms peak memory decreases monotonically as the compression ratio increases, matching what the reference kvpress benchmark shows for its own models.Reference
Checklist