$ the-wire · showcase
Inductor hands Triton kernels to the worker pool, RL rollout KV cache persists across turns
By RepoJournal · Filed · About PyTorch · Composed from the cited sources · methodology
Module-level Triton kernels now compile through the same worker pool as string kernels, torchtitan adds an opt-in RDMA microbatch path and cross-turn KV retention for multi-turn rollouts, and tensordict speeds up the nested get/update calls that run once per environment step.
[module-level kernels] inductor: compile module-level Triton kernels on the worker pool (#200491) pytorch/pytorch
Module-level Triton kernels previously skipped the compile worker pool: each hoisted kernel's def built its own CachingAutotuner and compiled it serially in process on first launch, so a cold compile of a 60-block MLP on an H100 took 28.5s against 18.1s for the default string-kernel path. The change records each hoisted kernel's source and hands it to the pool, and the companion fix covers FX g...
[rl] Add an opt-in RDMA fast path for trainer microbatches (#5285) pytorch/torchtitan
Setting async_loop.batcher.rdma_min_tensor_bytes makes the batcher keep tensors at or above that size as Monarch RDMA buffers, with each trainer rank reading only its own column; a 5.2 GB step at dp 2 x tp 2 drops from 9.1-11.1s to 3.3-5.1s on CPU Monarch over TCP on one host. The default stays None, which keeps main's by-value RPC.
[rl] hold each multi-turn session's KV between its turns (5/5) (#5220) pytorch/torchtitan
Opt-in via VLLMGenerator.Config.session_kv_holder = VLLMSessionKVHolder.Config(): between turns of a multi-turn rollout the session has no request in vLLM, so a new request could take its cached blocks and force the next turn to recompute its history. The holder keeps a reference on those blocks until the next turn or the rollout's end, releasing whole ids when memory runs short.
[Performance] Make update into nested tensordicts about twice as fast pytorch/tensordict
On a TensorDict, update took its fast path only when every value of the source was a tensor, so any source containing a sub-tensordict (which every nested env step has) fell through to the general per-key loop at the top level. Catching that case makes update into nested tensordicts about twice as fast, which matters for the exclude/select/update step that runs once per rollout and collector.
[Performance] Make reading a nested entry with td.get up to twice as fast pytorch/tensordict
td.get(('a', 'b'), default) is the idiom for reading an optional nested entry, and PhysicsNeMo's Mesh calls it on every property access; a two-level get now takes 3 or 4 Python calls instead of the previous path, up to twice as fast. The companion fix makes get and pop return the default when an earlier key segment names a tensor, instead of raising ValueError while (key) in td reports False.
[MPS] Move threshold and prelu to Metal, fixing NaN propagation (#200928) pytorch/pytorch
threshold was an MPSGraph op computing x > threshold ? x : value, so NaN inputs were silently replaced by value; the structured elementwise rewrite uses x <= threshold ? value : x and keeps NaN. That also removes the MPS expected failures on the _refs.nn.functional.threshold and prelu test_python_ref suite. Related metal work from the same author makes max/min resolve ties to the first argument...
[Deprecation] Warn that a tuple of floats will be stored as float32, like a list, in TensorDict 0.16 pytorch/tensordict
TensorDict 0.15 now emits a FutureWarning for float tuples: a tuple of floats stores as float64 while the same values as a list, or via torch.tensor, give float32, and the float64 leaf then promotes everything it meets. The switch to float32 lands in 0.16. The long tail is quiet: helion's pallas ordered-carry and pad fixes plus Blackwell MXFP8 lowering and autotuning canonicalization, and CI mi...
Action items