[Performance] SmolVLA in-place DynamicCache crop prevents torch.compile and limits inference to 4.8 Hz

Description

In SmolVLAPolicy.denoise_step (lerobot/policies/smolvla/modeling_smolvla.py and smolvlm_with_expert.py), past_key_values uses Hugging Face's DynamicCache.

During each step of the Euler flow-matching integration (num_steps=10), the action suffix key/value tokens are appended in-place via past_key_values.update(), and then truncated back to the prefix length via past_key_values.crop(prefix_len):

outputs_embeds, _ = self.vlm_with_expert.forward(
    attention_mask=full_att_2d_masks,
    position_ids=position_ids,
    past_key_values=past_key_values,
    inputs_embeds=[None, suffix_embs],
    use_cache=self.config.use_cache,
)
if past_key_values is not None:
    # Self-attention layers append suffix K/V in place; restore the prefix for the next step.
    past_key_values.crop(prefix_len)

Problems Caused:

1. Dynamic Reallocations & Crop Overhead: Calling update() followed by crop(prefix_len) 10 times per inference chunk incurs significant Python and CUDA memory allocation overhead.
2. Incompatible with torch.compile: Because the cache tensors are mutated in-place and dynamically resized every iteration, torch.compile catnot capture CUDA graphs or optimize kernels across iterations (resulting in dynamic shape guard re-evaluations and slower execution).

Proposed Solution

Make the prefix key/value cache static and immutable:

1. After computing the prefix KV cache once in sample_actions(), freeze it into a tuple of static tensors:
   python复制

   past_key_values = tuple((layer.keys, layer.values) for layer in past_key_values.layers)

2. In forward_attn_layer and forward_cross_attn_layer, when past_key_values is a tuple, concatenate the prefix functionally without mutating the stored prefix buffers:
   python复制

   prefix_k, prefix_v = past_key_values[layer_idx]
   key_states = torch.cat([prefix_k, key_states.transpose(1, 2)], dim=2).transpose(1, 2)
   value_states = torch.cat([prefix_v, value_states.transpose(1, 2)], dim=2).transpose(1, 2)

3. Skip past_key_values.crop(prefix_len) when using static tuple KV.
4. Add an optional compile_denoise: bool = False flag in SmolVLAConfig to allow compiling the denoise step with torch.compile(mode="max-autotune").

Benchmark (RTX 4090, lerobot/smolvla_base):

• Eager baseline: 207.28 ms (4.8 Hz)
• Static KV + max-autotune: 49.43 ms (20.2 Hz) — 4.19x speedup (saving 157.85 ms per prediction)
• Bitwise parity: Cosine Similarity = 0.99999

We have tested these changes and have a clean PR ready to link!