b80fd2b56b
Applied agent_merge_sort.cuh union _TempStorage pattern: cache V2 temporary tensors (tmp_output, exp_sums, max_logits) across decode steps instead of re-allocating each step. agent_merge_sort uses union to share one SMEM block across load_keys/load_items/store_keys/block_merge (serial ops). Our equivalent: module-level dict caches V2 tensors by shape key. For max_num_seqs=1 + 100K context: tmp_output: [1, 24, 200, 256] × 2B = 2.4 MB saved per step exp_sums + max_logits: 38 KB saved per step At ~200 steps/sec: ~480 MB/s saved CUDA malloc bandwidth. Also from weld_vertices.cu: confirmed slot_mapping int32 cast is safe (max 8M slots << int32_max=2.1B). CCCL files: cub/agent/agent_merge_sort.cuh, thrust/examples/weld_vertices.cu