Lightning-AI / Lightning-AI/lightning-thunder
Integration Test: KV-Cache management with symbolic values
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Create a comprehensive integration test that demonstrates all symbolic value features working together in a realistic transformer KV-cache management scenario.
## Depends On
- [ ] https://github.com/Lightning-AI/lightning-thunder/issues/2735
- [ ] https://github.com/Lightning-AI/lightning-thunder/issues/2736
- [ ] https://github.com/Lightning-AI/lightning-thunder/issues/2737
- [ ] https://github.com/Lightning-AI/lightning-thunder/issues/2738
- [ ] https://github.com/Lightning-AI/lightning-thunder/issues/2741
## Test Case
```py
import torch
import thunder
@thunder.jit
def kv_cache_update(
input_ids: torch.Tensor, # [1, seq_len] - seq_len is symbolic
cumulative_length: int, # symbolic scalar
kv_cache: torch.Tensor, # [1, 8, 1056, 128]
sliding_window: int = 1056,
) -> tuple[torch.Tensor, int, torch.Tensor]:
"""
Realistic KV-cache update demonstrating all symbolic patterns.
This mirrors the pattern from Llama-4 model.
"""
batch_size, seq_len = input_ids.shape
# Pattern 1: Symbolic arithmetic
new_cumulative_length = cumulative_length + seq_len
# Pattern 2: torch.arange with symbolic bounds
cache_position = torch.arange(
cumulative_length,
new_cumulative_length,
device=input_ids.device
)
# Pattern 3: sym_max for sliding window
kv_offset = max(cumulative_length - sliding_window + 1, 0)
# Pattern 4: Symbolic tensor shapes
# Compute query states with symbolic seq_len dimension
hidden_dim = 5120
embedding_weight = torch.randn(
50000, hidden_dim,
device='cuda',
dtype=torch.bfloat16
)
hidden_states = torch.nn.functional.embedding(input_ids, embedding_weight)
# hidden_states.shape = [1, s50, 5120]
# Compute keys for cache update
key_proj = torch.randn(1024, hidden_dim, device='cuda', dtype=torch.bfloat16)
keys = torch.nn.functional.linear(hidden_states, key_proj)
# keys.shape = [1, s50, 1024]
keys = keys.view(batch_size, seq_len, 8, 128).transpose(1, 2)
# keys.shape = [1, 8, s50, 128]
# Update KV-cache using symbolic cache_position
kv_cache = kv_cache.index_copy_(2, cache_position, keys)
return kv_cache, new_cumulative_length, cache_position
kv_cache = torch.zeros(1, 8, 1056, 128, device='cuda', dtype=torch.bfloat16)
cumulative_length = 0
# Prefill phase - long sequence
input_ids_prefill = torch.randint(0, 50000, (1, 1024), device='cuda')
kv_cache, cumulative_length, cache_pos = kv_cache_update(
input_ids_prefill, cumulative_length, kv_cache
)
assert cumulative_length == 1024
assert cache_pos.shape == (1024,)
# Decode phase - single tokens
for _ in range(32):
input_ids_decode = torch.randint(0, 50000, (1, 1), device='cuda')
kv_cache, cumulative_length, cache_pos = kv_cache_update(
input_ids_decode, cumulative_length, kv_cache
)
assert cache_pos.shape == (1,)
assert cumulative_length == 1056
```
## Validation Criteria
[ ] Test compiles successfully with @thunder.jit
[ ] Works with variable sequence lengths without recompilation
[ ] Performance comparable to torch.compile
[ ] Symbolic values tracked correctly throughout
[ ] Generated code is efficient without recompilation
## Performance Target
Should match or exceed torch.compile performance for this pattern.
## Documentation
- Once passing, add to Thunder documentation as an example of:
- Dynamic shape handling
- KV-cache management patterns
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with dependency issues 2735, 2736, 2737, 2738, and 2741, then review the @thunder.jit KV-cache example in this issue. Done means the integration test compiles, handles variable sequence lengths without recompilation, tracks symbolic values correctly, generates efficient code, meets the torch.compile performance target, and is added to the Thunder documentation.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, testing-qa
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100