$ the-wire · showcase
Green-context SM splitting, cat-kernel DCE, and a debug map that pinned every exported module
By RepoJournal · Filed · About PyTorch · Composed from the cited sources · methodology
12 people shipped this
Natalia Gimelshein
@ngimel
1 cited
Yuanyuan Chen
@cyyever
1 cited
Chong Gu
@oniononion36
1 cited
Laith Sakka
@laithsakka
1 cited
Ivan Kobzarev
@IvanKobzarev
1 cited
rraminen_amdeng
@rraminen
Huy Do
@huydhn
anijain2305
@anijain2305
jinsooihm
@jinsooihm
SherlockNoMad
@SherlockNoMad
liangel-02
@liangel-02
can-gaa-hou
@can-gaa-hou
PyTorch's CUDA layer gains explicit SM partitioning, ROCm sheds dead cat-kernel instantiations, and torch.export stops retaining the traced graph and its parameters through a debug-only map.
[Lowering Memory] Don't retain the traced graph in the export leak-detection map (#198295) (#198295) pytorch/pytorch
Every torch.export trace populated the module-level _FAKE_TENSOR_ID_TO_PROXY_MAP_FOR_EXPORT in proxy_tensor.py, which permanently retained the traced graph and, through it, the whole module and every parameter; the map is debug-only and read almost never, so exported models now release that memory.
[CUDA] Expose SM partitions and disjoint green-context splitting (#198278) pytorch/pytorch
Independent GreenContext constructors could not guarantee disjoint SM allocations or report their real size, so SMPartition now wraps CUDA-selected resources and GreenContext.split builds contexts with disjoint allocations, for example num_sms=(24, 40), coscheduled_sm_count=(8, 4).
Stop instantiating unreachable cat kernels (#194184) pytorch/pytorch
Under USE_ROCM, parallel_cat cleared isAligned and isInOutAligned at runtime because CatArrayBatchedCopy_contig is faster there, leaving the alignedK_contig and vectorized families unreachable but still emitted: __global__ functions are DCE roots, so a runtime flag cannot drop them and a constexpr bool, as in SoftMax.cu, is required.
[rel-land] allow horizontal fusion of copies from dynamic inputs into concat operation (#195635) (#198294) (#198294) pytorch/pytorch
ConcatKernel.create drops any cat input with a symbolic size from the foreach op list, so Scheduler.create_foreach_nodes never groups the copies and each becomes its own kernel launch, while the same op with static shapes is a single triton_for_fused_* launch; the static-shape optimization is not shape dependent, so the gap is now closed.
[graph_trainer] Cleanup compile mode config, now only aot_fx_trace (#4884) pytorch/torchtitan
GraphTrainer supports only aot_fx_trace, so compile.mode is being removed from the configuration and its call sites; check any torchtitan config or launcher that still passes it. Elsewhere, fused_mla moved to contiguous rope tiles and out-of-place Q rotation, which the author measures at 60.5us versus 14.1us in-place for 4k tokens, and LoRA for GroupedLinear lands with the constraint that rank ...
Action items
- → Drop compile.mode from torchtitan GraphTrainer configs; only aot_fx_trace is supported pytorch/torchtitan [plan]
