[bugfix] avoid recomputing Flash-decode RoPE cache

[taking-lying-flat]

Summary

• populate the Flash-decode RoPE cache only when a maximum sequence-length entry is missing
• reuse the cached cos/sin tensors on every subsequent decode step
• avoid repeated RoPE kernels and transient CUDA allocations on cache hits

Root cause

dict.setdefault(key, default) evaluates default before checking whether key already exists. Therefore this code recomputed a complete cos/sin pair on every decode step, even on a cache hit:

self.rotary_pos_emb_cache.setdefault(
    max_sequence_length,
    self.rotary_pos_emb.get_cos_sin(max_sequence_length),
)

The returned value still came from the dictionary, so the cache appeared to work functionally, but the newly computed tensors were discarded after consuming compute and temporary GPU memory.

Scope

GPTModel._get_rotary_pos_emb() runs before the decoder and passes the resulting RoPE tensors into the decoder. Consequently, the redundant get_cos_sin() call happens once per model decode step, not once per decoder layer. Decoder depth increases the model/cache memory baseline, but does not multiply this particular call count.

CUDA memory benchmark

The benchmark models the full text-cache topology of Qwen3.5-4B instead of measuring RoPE in isolation:

• 32 decoder layers: 24 linear-attention layers and 8 full-attention layers
• batch size 1, BF16, 4 KV heads, head dimension 256
• per full-attention layer: K/V cache shaped [1, 4, sequence_length, 256]
• per linear-attention layer: BF16 convolution state plus FP32 recurrent state
• RoPE dimension 64 (head_dim=256, partial_rotary_factor=0.25)
• RoPE cache populated once, followed by 200 warm decode cache hits

Environment: NVIDIA RTX A1000 Laptop GPU (4 GiB), PyTorch 2.13.0+cu130, Megatron-Core 0.18.2.

Sequence length	32-layer cache	Resident RoPE cache	Old peak allocated	Fixed peak allocated	Peak reduction	get_cos_sin() calls	RoPE path / step
16K	561.5 MiB	4.0 MiB	571.5 MiB	565.5 MiB	6.0 MiB	200 -> 0	0.0607 -> 0.0003 ms
32K	1073.5 MiB	8.0 MiB	1093.5 MiB	1081.5 MiB	12.0 MiB	200 -> 0	0.1126 -> 0.0003 ms
64K	2097.5 MiB	16.0 MiB	2137.5 MiB	2113.5 MiB	24.0 MiB	200 -> 0	0.2148 -> 0.0003 ms

At 32K and 64K, the old path also increased peak CUDA reserved memory by 20 MiB; the fixed path added 0 MiB. The absolute allocation reduction scales with sequence length. Its percentage of total process memory depends on model weights, parallelism, KV-cache layout, and other runtime allocations.

This is a cache-focused Transformer simulation, not an end-to-end throughput benchmark: model weights, attention kernels, and unrelated decode activations are intentionally excluded so the allocation caused by this code path remains measurable. The timing column measures only the RoPE cache path.

Validation

• python -m py_compile src/mcore_bridge/model/gpt_model.py
• git diff --check upstream/main...HEAD
• 32-layer CUDA cache simulation comparing the old and fixed paths over 200 warm decode steps