Integration Test: KV-Cache management with symbolic values

[IvanYashchuk]

🚀 Feature

Create a comprehensive integration test that demonstrates all symbolic value features working together in a realistic transformer KV-cache management scenario.

Depends On

• Support symbolic integer scalars #2735
• Support symbolic arithmetic operations #2736
• Support torch.sym_max/sym_min #2737
• Support symbolic values in torch.arange and other tensor operations #2738
• Support symbolic tensor shapes and dynamic dimensions #2741

Test Case

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

[shino16]

On main fb989d4 with nvFuser dc844187 on NVIDIA/Fuser#5648, the above script runs with no recompilation with cache="symbolic values".

[shino16]

NVIDIA/Fuser#5648 has been merged, but it looks like we need more patch.

  File "thunder.computation_2", line 23, in computation
  File "/opt/pytorch/lightning-thunder/thunder/executors/nvfuserex_impl.py", line 567, in __call__
    return fd.execute(
           ^^^^^^^^^^^
  File "/opt/pytorch/nvfuser/python/nvfuser_direct/__init__.py", line 369, in execute
    return self.fec.execute(
           ^^^^^^^^^^^^^^^^^
RuntimeError: Expected (stride)==(0) . Found 1 vs 0. Expecting an expanded dimension on dimension 1
Exception raised from validateAllocationSizesAndStrides at /opt/pytorch/nvfuser/csrc/tensor_metadata.cpp:253 (most recent call first):
frame #0: nvfuser::nvfCheckFail(char const*, char const*, long, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) + 0x94 (0x7bdb80145403 in /opt/pytorch/nvfuser/python/nvfuser_direct/../build/libnvfuser_codegen.so)
frame #1: <unknown function> + 0x17c3245 (0x7bdb80bc3245 in /opt/pytorch/nvfuser/python/nvfuser_direct/../build/libnvfuser_codegen.so)

I made a PR NVIDIA/Fuser#5669 to fix this. EDIT: This was irrelevant, the issue persisted.

[shino16]

This is caused by the following FusionDefinition.

def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[1, -1], contiguity=[None, None], dtype=DataType.Int, is_cpu=False)
    tv1 = fd.define_tensor(shape=[50000, 5120], contiguity=[True, True], dtype=DataType.BFloat16, is_cpu=False)
    tv2 = fd.ops.reshape(tv0, new_shape=[1])
    tv3 = fd.ops.broadcast(tv2, is_broadcast_dim=[False, True])
    tv4 = fd.ops.index_select(tv1, tv3, dim=0)
    c1 = fd.ops.size(tv0, dim=1)
    tv5 = fd.ops.reshape(tv4, new_shape=[1, c1, 5120])
    fd.add_output(tv5)

It was executed with inputs inputs[0] shape [1, 1] and stride [1. 1], and inputs[1] shape [50000, 5120] and stride [5120, 1].

Minimized repro:

from nvfuser_direct import FusionDefinition, DataType
import torch

def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[1, -1], contiguity=[None, None], dtype=DataType.Int, is_cpu=False)
    fd.add_output(tv0)

with FusionDefinition() as fd:
    nvfuser_fusion(fd)

inputs = [
    torch.randint(0, 50000, (1, 1), device='cuda'),
]
outputs = fd.execute(inputs)

This can be fixed by this patch on nvFuser.

diff --git a/csrc/tensor_metadata.cpp b/csrc/tensor_metadata.cpp
index 42a037211..89d83a97e 100644
--- a/csrc/tensor_metadata.cpp
+++ b/csrc/tensor_metadata.cpp
@@ -250,11 +250,15 @@ void validateAllocationSizesAndStrides(
     if (alloc_id->isBroadcast()) {
       NVF_CHECK(!contiguity[domain_index].has_value());
       if (alloc_id->hasExpandedExtent()) {
-        NVF_CHECK_EQ(
-            stride,
-            0,
-            "Expecting an expanded dimension on dimension ",
-            dim_index);
+        // When the runtime size after materialization is 1, a
+        // non-zero stride is harmless
+        if (size != 1) {
+          NVF_CHECK_EQ(
+              stride,
+              0,
+              "Expecting an expanded dimension on dimension ",
+              dim_index);
+        }
       }
       continue;
     }

This also aligns with how we decide the contiguity option to fd.define_tensor.

lightning-thunder/thunder/executors/nvfuserex_impl.py

Lines 473 to 475 in fb989d4

	The contiguity is represented by True, False and None. True represents
	dimensions that are contiguous, and False represents dimensions that are not
	contiguous, and None represents stride-0 or size-1 dimensions.

Or, we could change the way we compute contiguity option to align with nvfuser's real behavior.

[shino16]

Made a PR to nvFuser. NVIDIA/Fuser#5677

Now the update script (below) runs with no error!

Details

from functools import partial
import torch
import thunder

@partial(thunder.jit, cache="symbolic values")
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

assert thunder.cache_misses(kv_cache_update) == 1