Repository navigation
Support Transformers 5.20 cache query offsets in DeepSpeedStaticCache - #8852
Draft
nathon-lee wants to merge 13 commits into
Draft
nathon-lee wants to merge 13 commits into
nathon-lee wants to merge 13 commits into
Conversation
Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Cache continuous-rollout CUDA graph state across matching generate calls. Rebuild the graph when the model parameters or static decode shape changes, and reset static buffers and KV cache rows before reuse. Use eager attention and a preallocated additive causal mask during capture to avoid Transformers capture-time mask allocations. Relax cache-class detection for models that support the cache interface without declaring the private capability flag. Add coverage for eager-attention capture, graph reuse, weight invalidation, and current StaticCache update semantics. Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Use the preallocated additive causal mask to make CUDA Graph capture safe without overriding the model attention implementation. This preserves default SDPA behavior and restores eager/graph token parity for pretrained GPT-2. Update graph-capture coverage to verify the configured attention implementation remains unchanged. Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Populate continuous rollout cache statistics after CUDA Graph decoding so callers do not receive stale eager-path metrics. Cover graph cache capacity, peak usage, decode steps, and trimming metrics. Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Validate fixed-slot CUDA Graph options before capture and cover unsupported decode alignment, cache trimming, and insufficient cache capacity. Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Add CUDA rollout coverage for pretrained Hugging Face models behind DEEPSPEED_PRETRAINED_ROLLOUT_MODEL so regular CI does not download weights. Validate eager and CUDA Graph parity across prompts and verify graph reuse. Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Route dummy writes for non-decoding graph slots to a scratch cache position so refilled prompts retain their prefetched KV state. Add CUDA coverage for a slot refill while another request continues decoding. Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Reject flash_attention_2 before continuous CUDA Graph capture because the fixed replay path requires a four-dimensional additive attention mask. Document the restriction and cover the configuration error. Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Reject models that explicitly lack cache-class support before the continuous CUDA Graph path creates static cache state. Add regression coverage for the graph-specific validation path. Signed-off-by: nathon-lee <leejianwoo@gmail.com>
…raction Keep static_cache_position in _ContinuousGraphState so it remains alive for the lifetime of the captured graph. Add is_current_stream_capturing() to the accelerator interface, with a default False implementation and a CUDA override. Route the static cache capture check through this interface instead of calling torch.cuda directly. Add regression coverage for tensor lifetime, graph replay, capture state, and eager cache bounds checks, and document the new accelerator method. Signed-off-by: nathon-lee <leejianwoo@gmail.com>
…cCache Signed-off-by: nathon-lee <leejianwoo@gmail.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Background
Native Llama continuous rollout on Transformers 5.20.1 failed because
DeepSpeedStaticCachelackedget_query_offset(layer_idx). The query offset must exclude the current write; usingget_seq_length()would be off by one.Changes
Implemented and validated on
fix/static-cache-query-offset:Validation
RTX A4500 (20 GiB), PyTorch 2.9.1+cu128 / CUDA 12.8, FP16:
Actual CUDA capture/replay passed. Randomly initialized tiny Llama rollouts matched native
generate()tokens, attention masks, and response offsets. Both versions skipped the optional pretrained-model test; 4.57.6 also skipped three tests requiring the newer masking API. Modified-file pre-commit checks passed.Run in each Transformers environment:
python -m pytest \ tests/unit/utils/test_static_cache.py \ tests/unit/runtime/rollout/test_hybrid_engine_rollout.py \ tests/unit/runtime/rollout/test_continuous_batching.py \ tests/unit/v1/test_hybrid_engine_rollout_v1.py -qScope and dependencies
Depends on the cache code from #8732, not the benchmark. Changes are limited to compatibility, tests, and documentation; this does not claim full Transformers 5.20 cache API support. Issue 1 depends on this fix in the tested 5.20.1 environment.