[XPU] Add _make_xccl_premul_sum binding for XCCL backend#25
Open
Chao1Han wants to merge 1172 commits into
Open
[XPU] Add _make_xccl_premul_sum binding for XCCL backend#25Chao1Han wants to merge 1172 commits into
Chao1Han wants to merge 1172 commits into
Conversation
Chao1Han
changed the title
pytorchmergebot
force-pushed
the
xccl-premul
branch
from
February 12, 2026 06:00
7ba5e71 to
b7e7e59
Compare
pytorchmergebot
force-pushed
the
xccl-premul
branch
from
March 10, 2026 01:37
b7e7e59 to
b5018b0
Compare
Chao1Han
pushed a commit
that referenced
this pull request
Mar 16, 2026
…nces between x86 vs aarch64 (pytorch#176085) In the test: ``` python test/cpp_extensions/test_libtorch_agnostic.py TestLibtorchAgnosticCUDA.test_std_cuda_check_error_show_cpp_stacktraces_True_cuda ``` it raises an exception when calling `STD_CUDA_CHECK(cudaSetDevice(99999));` which got the expected `CUDA error: invalid device` message. However, the expected string for the C++ stack trace is different between `x86` vs `aarch64` due perhaps in these issues: - pytorch#119905 - pytorch#134387 In the current setup when getting a stack trace string: - x86 contains `C++ CapturedTraceback:` - aarch64 contains `Exception raised from` + `frame #` An example of the full string from an aarch64 system when : ``` AssertionError: 'C++ CapturedTraceback:' not found in 'CUDA error: invalid device ordinal\nGPU device may be out of range, do you have enough GPUs?\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\n\nException raised from test_std_cuda_check_error at /opt/pytorch/pytorch/test/cpp_extensions/libtorch_agn_2_10_extension/csrc/test_std_cuda_check.cu:23 (most recent call first):\nframe #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0xd4 (0xe471ebcd39f4 in /usr/local/lib/python3.12/dist-packages/torch/lib/libc10.so)\nframe #1: <unknown function> + 0x43f998 (0xe471ebdcf998 in /usr/local/lib/python3.12/dist-packages/torch/lib/libc10_cuda.so)\nframe #2: c10::cuda::c10_cuda_check_implementation(int, char const*, char const*, unsigned int, bool) + 0x1bc (0xe471ebdcfc0c in /usr/local/lib/python3.12/dist-packages/torch/lib/libc10_cuda.so)\nframe #3: torch_c10_cuda_check_msg + 0x1c (0xe471ef335c4c in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_cuda.so)\nframe #4: test_std_cuda_check_error() + 0x58 (0xe470cd396678 in /opt/pytorch/pytorch/test/cpp_extensions/libtorch_agn_2_10_extension/install/usr/local/lib/python3.12/dist-packages/libtorch_agn_2_10/_C.so)\nframe #5: c10::BoxedKernel::makeFromFunctor<StableIValueBoxedKernel>(std::unique_ptr<StableIValueBoxedKernel, std::default_delete<StableIValueBoxedKernel> >)::{lambda(c10::OperatorKernel*, c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocator<c10::IValue> >*)#1}::_FUN(c10::OperatorKernel*, c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocator<c10::IValue> >*) + 0x16c (0xe47211cd419c in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_cpu.so)\nframe #6: <unknown function> + 0x61d34bc (0xe47211cf34bc in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_cpu.so)\nframe #7: <unknown function> + 0xe6c324 (0xe4721532c324 in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so)\nframe #8: <unknown function> + 0xe6c7e0 (0xe4721532c7e0 in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so)\nframe #9: <unknown function> + 0xd3907c (0xe472151f907c in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so)\nframe #10: <unknown function> + 0x5ccbf8 (0xe47214a8cbf8 in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so)\nframe #11: /usr/bin/python() [0x504a34]\nframe #12: PyObject_Call + 0x6c (0x4c633c in /usr/bin/python)\nframe #13: _PyEval_EvalFrameDefault + 0x3ea0 (0x568564 in /usr/bin/python)\nframe #14: _PyObject_Call_Prepend + 0xc4 (0x4c5934 in /usr/bin/python)\nframe #15: /usr/bin/python() [0x52a070]\nframe #16: _PyObject_MakeTpCall + 0x78 (0x4c3e58 in /usr/bin/python)\nframe #17: _PyEval_EvalFrameDefault + 0x8a0 (0x564f64 in /usr/bin/python)\nframe #18: PyEval_EvalCode + 0x130 (0x5632b4 in /usr/bin/python)\nframe #19: PyRun_StringFlags + 0xe0 (0x59c330 in /usr/bin/python)\nframe #20: PyRun_SimpleStringFlags + 0x44 (0x67ebc4 in /usr/bin/python)\nframe #21: Py_RunMain + 0x390 (0x68b380 in /usr/bin/python)\nframe #22: Py_BytesMain + 0x28 (0x68ae88 in /usr/bin/python)\nframe #23: <unknown function> + 0x284c4 (0xe47216b084c4 in /lib/aarch64-linux-gnu/libc.so.6)\nframe #24: __libc_start_main + 0x98 (0xe47216b08598 in /lib/aarch64-linux-gnu/libc.so.6)\nframe #25: _start + 0x30 (0x5f6770 in /usr/bin/python)\n\n' To execute this test, run the following from the base repo dir: python test/cpp_extensions/test_libtorch_agnostic.py TestLibtorchAgnosticCUDA.test_std_cuda_check_error_show_cpp_stacktraces_True_cuda ``` Pull Request resolved: pytorch#176085 Approved by: https://github.com/eqy
pytorchmergebot
force-pushed
the
xccl-premul
branch
3 times, most recently
from
March 26, 2026 02:21
bb50206 to
683f443
Compare
pytorchmergebot
force-pushed
the
xccl-premul
branch
from
March 30, 2026 08:03
0eac563 to
e06d66a
Compare
…ytorch#179310) (pytorch#179310) Summary: Remove `!defined(USE_ROCM)` guards from the 3 FP8 code paths in cuSPARSELtOps.cpp, enabling Float8_e4m3fn support for hipSparseLt. hipSparseLt 0.2.7 on MI350X (gfx950) supports FP8 input with FP32 output only. FP16/BF16/FP8 output types are not supported and will produce a clear error message. Changes: - `_cslt_compress`: Allow FP8 dtype on ROCm - `_cslt_sparse_mm_impl`: Allow FP8 input type, default to FP32 output on ROCm - `_cslt_sparse_mm_impl` out_dtype: Restrict to FP32 on ROCm with descriptive error for unsupported types Tested on MI350X (gfx950) with hipsparselt 0.2.7: - FP8 compress: PASS - FP8 sparse_mm with FP32 output: PASS - Correctness (64x64): max diff 0.000732 - Correctness (256x512x128): max diff 0.001709 Test Plan: ## NVIDIA devserver ## ``` buck test fbcode//mode/opt fbcode//caffe2/test:test_sparse_cuda -- 'test_sparse_semi_structured.' ``` https://www.internalfb.com/intern/testinfra/testrun/25614222893211678 ## MI350x ## ``` buck2 run mode/opt-amd-gpu -m rocm70 -m "fbcode//caffe2:hipsparselt[enabled]" fbcode//caffe2/test:test_sparse_cuda -- -r 'test_sparse_semi_structured' ``` ``` - test_cslt_compress_fp8_cuda ... ok — FP8 tensor compression via torch._cslt_compress - test_cslt_sparse_mm_fp8_to_fp32_dense_input_shape0_cuda ... ok — FP8 sparse matmul with float32 output - test_cslt_sparse_mm_fp8_unsupported_out_dtype_rocm_bfloat16_cuda ... ok — verifies unsupported bfloat16 output raises error on ROCm - test_cslt_sparse_mm_fp8_unsupported_out_dtype_rocm_float16_cuda ... ok — verifies unsupported float16 output raises error on ROCm - test_cslt_sparse_mm_fp8_unsupported_out_dtype_rocm_float8_e4m3fn_cuda ... ok — verifies unsupported fp8 output raises error on ROCm ``` Full log: P2265144022 Reviewed By: RandySheriff Differential Revision: D96821901 Pull Request resolved: pytorch#179310 Approved by: https://github.com/Skylion007
``` (pytorch-3048) [pianpwk@devvm3048.dkl0 /data/users/pianpwk/ptclone/pytorch (b1582add)]$ python -m torch.distributed.tensor._ops.strategy_validation --op "_upsample_bilinear2d_aa,grid_sampler_2d,max_pool2d_with_indices_backward,nn.functional.adaptive_avg_pool1d,nn.functional.adaptive_avg_pool2d,nn.functional.grid_sample,nn.functional.interpolate,nn.functional.max_unpool1d,nn.functional.max_unpool2d,nn.functional.max_unpool3d" --max-samples 10 --allow-composite ====================================================================== Summary ====================================================================== Op Correct Incorrect Missing Time -------------------------------------------------------------------------- aten._upsample_bilinear2d_aa 21 0 0 7.6s aten.grid_sampler_2d 30 0 3 118.0s aten.max_pool2d_with_indices_backward 164 0 28 ~8s nn.functional.adaptive_avg_pool1d 66 0 4 13.6s nn.functional.adaptive_avg_pool2d 31 0 2 9.3s nn.functional.grid_sample 30 0 4 126.4s nn.functional.interpolate 280 0 2 83.9s -------------------------------------------------------------------------- Total 622 0 ~43 0 incorrect across the board. The remaining missing rules are: - Upsample/interpolate ops (16): P(max)->P(max), P(min)->P(min) — invalid, interpolation blends values - grid_sampler_2d (3-4): R, S(2)->S(3) (grid spatial sharding) + P(max/min) — invalid - _adaptive_avg_pool2d (2): P(max)->P(max), P(min)->P(min) — invalid - mean.dim (4): P(avg)->P(avg), P(sum)->P(sum), S(2)->P(avg), S(3)->P(avg) — out of scope (decomposition chain) - ones_like (16): P(*)->P(*) / P(*)->S(3) — out of scope (decomposition chain) - max_pool2d_with_indices_backward (28): mixed-Partial combos (e.g. P(max), P(avg), P(min) -> P(max)) — can't be expressed with current framework, likely coincidental ``` Pull Request resolved: pytorch#176991 Approved by: https://github.com/zpcore
…h#179749) ## Summary When a PR updates `.github/ci_commit_pins/torchtitan.txt`, the torchtitan CI workflow should run automatically. Currently it doesn't because the pin file path is not in the `ciflow/torchtitan` labeler config. This adds the pin file path to `.github/labeler.yml` so the labeler auto-applies the `ciflow/torchtitan` label, which triggers the workflow via the existing ciflow tag mechanism. This is problematic PR example: pytorch#178727 ## Test plan - Create a PR that modifies `.github/ci_commit_pins/torchtitan.txt` - Verify the `ciflow/torchtitan` label is auto-applied by the labeler - Verify the `torchtitan-test` workflow is triggered Pull Request resolved: pytorch#179749 Approved by: https://github.com/huydhn
### Summary
When `torch.compile` generates an `fx_graph_runnable` repro script for a graph
containing `graphsafe_run_with_rng_state`, the `fwd_rng_state_N` argument (a
`torch._C.Generator`) falls through to `InputWriter.unsupported()`, producing
`reader.unsupported(...)` in the repro. At load time, `InputReader.unsupported()`
returns `None`, so the repro crashes with
`AttributeError: 'NoneType' object has no attribute 'device'`.
This adds `generator()` methods to `InputWriter`, `InputReader`, and
`NopInputReader` that properly serialize/deserialize Generator arguments, and
detects `torch._C.Generator` in `save_graph_repro` so it routes through the new
path instead of `unsupported`.
Authored with Claude.
### Test plan
- `python test/dynamo/test_after_aot.py TestAfterAot.test_dump_generator`
- `python test/dynamo/test_after_aot.py TestAfterAot.test_graphsafe_rng_repro`
---
E2E Example:
```
import torch
import torch.utils.checkpoint
torch._logging.set_logs(aot_graphs=True)
def fn(q, k, v):
return torch.nn.functional.scaled_dot_product_attention(
q, k, v, dropout_p=0.1, is_causal=True
)
def checkpointed_fn(q, k, v):
return torch.utils.checkpoint.checkpoint(fn, q, k, v, use_reentrant=False)
model = torch.compile(checkpointed_fn)
q = torch.randn(2, 8, 64, 32, device="cuda", dtype=torch.float16, requires_grad=True)
k = torch.randn(2, 8, 64, 32, device="cuda", dtype=torch.float16, requires_grad=True)
v = torch.randn(2, 8, 64, 32, device="cuda", dtype=torch.float16, requires_grad=True)
out = model(q, k, v)
out.sum().backward()
```
fx_graph_runnable, note the **reader.generator('cuda', 0) # fwd_rng_state_0** generated.
Without the PR, it will generate `reader.unsupported('fwd_rng_state_0') # unsupported type for dumping: <class 'torch._C.Generator'>` and error out with
```
device_idx = rng_state.device.index
^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'device'
```
```
.....
torch._higher_order_ops.triton_kernel_wrap.kernel_side_table.reset_table()
from torch.nn import *
class Repro(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
def forward(self, primals_1, primals_2, primals_3, fwd_rng_state_0):
graphsafe_run_with_rng_state = torch.ops.higher_order.graphsafe_run_with_rng_state(torch.ops.aten._scaled_dot_product_efficient_attention.default, primals_1, primals_2, primals_3, None, True, 0.1, True, rng_state = fwd_rng_state_0); fwd_rng_state_0 = None
getitem = graphsafe_run_with_rng_state[0]; graphsafe_run_with_rng_state = None
return (getitem, primals_1, primals_2, primals_3)
def load_args(reader):
buf0 = reader.storage(None, 65536, device=device(type='cuda', index=0), dtype_hint=torch.float16)
reader.tensor(buf0, (2, 8, 64, 32), dtype=torch.float16, is_leaf=True) # primals_1
buf1 = reader.storage(None, 65536, device=device(type='cuda', index=0), dtype_hint=torch.float16)
reader.tensor(buf1, (2, 8, 64, 32), dtype=torch.float16, is_leaf=True) # primals_2
buf2 = reader.storage(None, 65536, device=device(type='cuda', index=0), dtype_hint=torch.float16)
reader.tensor(buf2, (2, 8, 64, 32), dtype=torch.float16, is_leaf=True) # primals_3
reader.generator('cuda', 0) # fwd_rng_state_0
```
Pull Request resolved: pytorch#179657
Approved by: https://github.com/PaulZhang12
…ingFormatVariable (pytorch#179629) Add importlib to BUILTIN_SKIPLIST so all importlib functions become SkipFunctionVariable, then constant-fold `importlib.util.find_spec` and `importlib.metadata.version` when their args are constants. Also teach callable() that StringFormatVariable is not callable. These fixes resolve graph breaks seen in HuggingFace VitsModel. Authored with Claude. Pull Request resolved: pytorch#179629 Approved by: https://github.com/Lucaskabela ghstack dependencies: pytorch#179611
Add return type annotations and parameterize generic types across core torch/fx utility files to satisfy pyrefly's `unannotated-return` and `implicit-any` checks. Authored with Claude. Pull Request resolved: pytorch#179590 Approved by: https://github.com/Skylion007
… for StringFormatVariable (pytorch#179629)" This reverts commit 42e4e00. Reverted pytorch#179629 on behalf of https://github.com/anijain2305 due to missed on william's reveiew ([comment](pytorch#179629 (comment)))
…179028) (pytorch#179028) Summary: ## Summary `aten.arange.start_step` only existed in Inductor's decomposition table (decomposed to `prims.iota`), but was NOT registered in the `lowerings` dict. This caused `KeyError` when any code tried to access `lowerings[aten.arange.start_step]`. This fixes 176 opinfo e2e test failures for `arange` on MTIA, where `arange_int_lowering.py` (D98064707) delegates to `lowerings[aten.arange.start_step]`. ### Root Cause D98064707 added `arange_int_lowering.py` which registers a lowering for `torch.ops.mtia.arange_int` and delegates to `lowerings[aten.arange.start_step]`. However, upstream Inductor only registers `aten.arange` in its **decomposition** table (decomposed to `prims.iota` at `lowering.py:3453`), not in the **lowerings** dict. The unit tests masked this by injecting a mock via `patch.dict(lowerings, {aten.arange.start_step: mock_arange})`. ### Fix Add `register_lowering(aten.arange.start_step)` that delegates to the existing `prims.iota` lowering (which generates `Pointwise.create` with `ops.index_expr`). This makes the key available in the `lowerings` dict for any downstream code. ### Changed Files - `fbcode/caffe2/torch/_inductor/lowering.py` — add `arange_start_step` lowering after `iota` - `xplat/caffe2/torch/_inductor/lowering.py` — mirror - `fbcode/caffe2/test/inductor/test_torchinductor.py` — add `test_arange7` for integer arange via `aten.arange.start_step` - `xplat/caffe2/test/inductor/test_torchinductor.py` — mirror Test Plan: - `test_arange7` in `test_torchinductor.py` — verifies `aten.arange.start_step` with integer dtype - Existing `test_arange3` already calls `aten.arange.start_step` directly — passes with new lowering - MTIA `test_arange_int_lowering` unit tests — delegation to `lowerings[aten.arange.start_step]` no longer raises KeyError - Opinfo e2e: `arange; 1; 0; int, 2` should compile successfully with auto DS Reviewed By: blaine-rister Differential Revision: D98755321 Pull Request resolved: pytorch#179028 Approved by: https://github.com/blaine-rister
)" This reverts commit e189ea7. Reverted pytorch#179195 on behalf of https://github.com/fxdawnn due to Some more internal tests failing that might be related to this PR ([comment](pytorch#179195 (comment)))
…tioner classification (pytorch#179660) When an opaque object (e.g. a DeviceMesh submesh) is reconstructed during tracing via `_try_reconstruct_opaque`, the reconstruct_fn dispatches through a custom op with a Proxy argument. This goes through `Proxy.__torch_function__`, which creates a graph node but does NOT populate `meta["val"]` — unlike the `ProxyTorchDispatchMode.__torch_dispatch__` path, which calls `track_tensor_tree` to set it. Without `meta["val"]`, `is_opaque_node()` returns False for these nodes, causing two downstream failures in the min-cut partitioner (`default_partition`): 1. Nodes tagged `CheckpointPolicy.MUST_SAVE` by activation checkpointing are routed to `saved_values` instead of `saved_opaque_nodes`. The runtime `save_from_forward` then hits: `AssertionError: Expected _get_submesh_1 to be a tensor` 2. Impure opaque nodes hit the same `saved_values` path with the same assertion failure. This is needed to support precompile with regional_inductor in torchtitan. In that setting, DTensor operations produce DeviceMesh submeshes via `_get_submesh` call_function nodes, which are opaque objects derived from the parent mesh (hoisted as a graph input). When activation checkpointing is applied and the partitioner runs, these submesh nodes must be correctly classified as opaque so they go through the opaque save/restore path rather than the tensor save path. Fix: - proxy_tensor.py: Call `set_meta(result, a)` in `_try_reconstruct_opaque` after the reconstruct_fn returns, so the node gets `meta["val"]` set. - partitioners.py: Add `is_opaque_node()` checks in the MUST_SAVE and is_impure branches of `default_partition`, routing opaque nodes to `saved_opaque_nodes` instead of `saved_values`. Tests: - test_reconstruct_fn_sets_meta_val: Verifies that nodes created via reconstruct_fn have meta["val"] set and is_opaque_node() returns True. - test_partitioner_must_save_opaque_node: Verifies that opaque nodes tagged MUST_SAVE are correctly routed through the partitioner without triggering the tensor assertion. Pull Request resolved: pytorch#179660 Approved by: https://github.com/zhxchen17 ghstack dependencies: pytorch#179396
…ng (pytorch#177743) Fix issue 2 discovered in pytorch#177166. elementwise_dtypes was registered as TorchInGraphFunctionVariable via torch._higher_order_ops.out_dtype, causing dynamo to try putting it in the FX graph. Since it returns (dtype, dtype) rather than tensors, this failed with "torch.* op returned non-Tensor". The fix adds a handler that evaluates elementwise_dtypes eagerly on fake tensor metadata during compilation and returns the result as constants. Pull Request resolved: pytorch#177743 Approved by: https://github.com/anijain2305
### Summary Refactors `SetVariable` to inherit directly from `VariableTracker` instead of `ConstDictVariable`, fully decoupling set tracking from dict tracking in Dynamo. ### Motivation Cpython handles dict and set quite differently now, so to more closely follow the cpython paradigm and support transitioning to tp_slot support it was important to decouple the previous inheritence scheme. The inheritance was a known footgun due to incorrect fallback routing possibility from set operations sharing names with dict operations. ### Changes - `sets.py` (new standalone module): `SetVariable` and all subclasses (`OrderedSetVariable`, `FrozensetVariable`, `DictKeySetVariable`, `OrderedSetClassVariable`) now live in a self-contained module with proper imports. `HashableTracker`, `is_hashable`, and `raise_unhashable` also moved here from `dicts.py` since they're part of the shared hashability protocol. - `SetVariable` is now independent of `ConstDictVariable`: All functionality that `SetVariable` previously inherited — `__contains__, len, clone, unpack_var_sequence, has_new_items, is_new_item, is_python_hashable, var_getattr, call_obj_hasattr, install_dict_contains_guard` is ported directly onto `SetVariable`. - All call_method cases that previously fell through to `ConstDictVariable` (`__contains__, __len__, __iter__, __ne__, copy, clear, add, pop, remove, discard`) are now handled by `SetVariable`. Dict-only methods are correctly omitted. - `isinstance` checks updated across 8 files: Each `isinstance(x, ConstDictVariable)` check was individually assessed — some were simplified (removing now-unnecessary `not isinstance(x, SetVariable)` exclusions in `side_effects.py` and `builtin.py`), some had `SetVariable` added to maintain equivalent behavior (side_effects.py codegen, builtin.py bool/not, higher_order_ops.py, object_protocol.py), and dict-only checks were left unchanged. - `dicts.py` cleaned up: Set-related classes, `is_hashable`, and `raise_unhashable` removed (now imported from sets.py). Docstring updated. Import chain is strictly one-directional: dicts.py → sets.py. `SetVariable` still handles its contents as a dict with `None` keys. This was a decision to limit the scope of code changes where possible, but should maintain all functionality. ### Testing The existing dynamo and cpython test coverage of the SetVariable class and Subclasses to support the correctness of the changes. Pull Request resolved: pytorch#179192 Approved by: https://github.com/Lucaskabela, https://github.com/anijain2305
…ytorch#178525) (pytorch#178525) Summary: ## Summary Add Triton codegen support for sort-based operations in Inductor, gated behind `config.triton.decompose_sort_ops`. When enabled, these ops decompose into sort-based Inductor IR instead of falling back to ATen eager, enabling Triton kernel generation via `ir.Sort`: **New lowerings in this diff:** - `aten.sort.stable` / `aten.sort` — decompose into `ir.Sort` directly (sort + argsort) - `aten.median.default` / `aten.median.dim` — decompose into sort + select middle element - `aten.mode.default` — decompose into sort + run-length comparison - `aten.topk.default` — decompose into sort + slice top-k elements - `aten.kthvalue.default` — decompose into sort + select k-th element - `aten.nanmedian.default` / `aten.nanmedian.dim` — registered as fallback-only (data-dependent indexing prevents Triton decomposition) **Key design decisions:** - Each lowering checks `config.triton.decompose_sort_ops` and falls back to ATen if disabled - Fallback handlers use `add_to_fallback_set=False` so they don't register as global make_fallback (allows other backends to provide their own lowerings) - Removed `make_fallback` for `aten.sort`, `aten.sort.stable`, `aten.kthvalue`, `aten.topk`, `aten.nanmedian` since they now have explicit register_lowering - The `ir.Sort` persistent kernel size limit (512 elements) is lifted when `decompose_sort_ops` is enabled (change in ir.py) - Named-dim and out= overloads (e.g., `aten.median.names_dim`, `aten.sort.values`) have upstream decompositions and don't need explicit handling **Changed files:** - `torch/_inductor/lowering.py` — Add lowerings for topk, kthvalue, nanmedian; remove make_fallback entries - `torch/_inductor/ir.py` — Lift sort size limit when decompose_sort_ops is enabled Test Plan: ## Triton codegen tests (GPU-only, in TritonCodeGenTests) Tests moved from CommonTemplate to TritonCodeGenTests (GPU-only) to properly test Triton codegen generation. Each test: 1. Compiles with `decompose_sort_ops=True` 2. Verifies `sort_with_index` appears in generated Triton code (codegen verification) 3. Verifies numerics match eager PyTorch - `test_median_decompose_sort_ops` — median decomposition via sort + select - `test_mode_decompose_sort_ops` — mode decomposition via sort + run-length - `test_topk_decompose_sort_ops` — topk decomposition via sort + slice (largest=True and False) - `test_kthvalue_decompose_sort_ops` — kthvalue decomposition via sort + select ## Changed files - `torch/_inductor/lowering.py` — topk, kthvalue, nanmedian lowerings + removed make_fallback - `torch/_inductor/ir.py` — lift sort size limit when decompose_sort_ops enabled - `test/inductor/test_torchinductor.py` — moved decompose_sort_ops tests to TritonCodeGenTests with codegen verification Reviewed By: blaine-rister Differential Revision: D98238316 Pull Request resolved: pytorch#178525 Approved by: https://github.com/jansel
…torch#179105) Add `_patch_autograd_grad` context manager that patches `autograd.grad` to install stacktrace preservation hooks (for seq_nr propagation) and tag backward nodes with `custom["autograd_backward"]` via `fx_traceback.annotate`. Update the remat pass to detect backward nodes via two mechanisms: - `custom["autograd_backward"]`: set automatically by _patch_autograd_grad - `custom["phase"] == "backward"`: user-level annotation for disambiguation When the user provides `phase` annotations, only those are used. If multiple backward regions are detected without user annotation (e.g. multiple `autograd.grad` calls), the pass errors and asks the user to annotate the real backward with `fx.traceback.annotate({"phase": "backward"})`. Migrate existing tests from the legacy `remat_pass_tag` annotation to the new `phase: backward` annotation. Authored with Claude. ghstack-source-id: 575d00e Pull Request resolved: pytorch#179105 Pull Request resolved: pytorch#179105 Approved by: https://github.com/aorenste ghstack dependencies: pytorch#179043
…79217) The strategy computation for aten.index_put had two remaining issues blocking models like Falcon-7b that use list indexing (e.g. tensor[..., [-1], :]): 1. aten.index_put_.default (in-place variant, generated by autograd in the backward of aten.index.Tensor) was not registered in the single-dim strategy system. 2. _get_num_tensor_inputs counted None children in TupleStrategy, but args_strategy drops them, causing a count mismatch for ops with List[Optional[Tensor]] arguments. The values-dim-to-self-dim mapping fix for consecutive indexed dims was already landed in pytorch#178350; this commit adds the missing op registration, the infra bug fix, updated docstrings, and additional test coverage for None indices, in-place index_put_, non-contiguous indexed dims, and mid-tensor indexing. Authored with Claude. Fixes meta-pytorch/autoparallel#348 Pull Request resolved: pytorch#179217 Approved by: https://github.com/pianpwk
Summary: As titled. Created from CodeHub with https://fburl.com/edit-in-codehub Test Plan: Sandcastle run Differential Revision: D100042625 Pull Request resolved: pytorch#179737 Approved by: https://github.com/Skylion007
This PR ports the NVSHMEM-based symmetric memory collective operations to ROCm by integrating rocSHMEM, enabling the same `torch.ops.symm_mem.*` APIs to work on AMD GPUs. The ported operations include point-to-point (`put`, `get`, `put_with_signal`, `wait_for_signal`), collectives (`broadcast`, `all_to_all`), and the MoE-oriented 2D AllToAllv variants (`all_to_all_vdev`, `all_to_all_vdev_2d`, `all_to_all_vdev_2d_offset`). **Key design decisions** **Separate `rocshmem_extension.cu` instead of `#ifdef` in `nvshmem_extension.cu`.** Rather than threading `#if defined(USE_ROCM)` throughout the existing NVSHMEM file, the ROCm implementation lives in a dedicated `rocshmem_extension.cu`. This was necessary because: **rocSHMEM** requires `-fgpu-rdc` (relocatable device code) for device-side API calls, while the rest of PyTorch's HIP code uses `-fno-gpu-rdc`. A separate compilation unit with its own flags avoids contaminating the global build. The NVSHMEM and rocSHMEM device APIs differ enough (e.g., nvshmemx_collective_launch vs standard hip kernel launches, `nvshmemx_getmem_nbi_block` vs `rocshmem_getmem_nbi_wg`, `warp_size` 32 vs 64) that a shared source would be more `#ifdef` than code. The **hipification layer** handles the host API mapping (nvshmem_* → rocshmem_*) so `NVSHMEMSymmetricMemory.cpp` and `nvshmem_team_manager.hpp` work for both backends with minimal `#ifdefs`. No device-side team dereference **No grid-wide sync → separate offset writeback kernel**. NVSHMEM uses nvshmemx_collective_launch which provides grid-wide synchronization, allowing allToAllV_2d to write output offsets in the same kernel after the data exchange. rocSHMEM launches are standard hip kernel lauch with no grid-wide barrier, so writing offsets inside a multi-block kernel would race with blocks still reading source_offsets. A small single-block `writeOutputOffsets_2d` kernel runs after allToAllV_2d completes to compute and write the output offsets safely. Fixes pytorch#168482 Fixes pytorch#168483 Fixes pytorch#168484 Fixes pytorch#168485 Fixes pytorch#168486 Fixes pytorch#168487 Fixes pytorch#168488 Fixes pytorch#168489 Fixes pytorch#168490 Fixes pytorch#168491 Fixes pytorch#168492 Fixes pytorch#168493 Fixes pytorch#168494 Fixes pytorch#168495 Fixes pytorch#168496 Fixes pytorch#168497 Fixes pytorch#168498 Fixes pytorch#168499 Fixes pytorch#166807 Pull Request resolved: pytorch#173518 Approved by: https://github.com/jeffdaily, https://github.com/kwen2501
…rch#179708) (pytorch#179708) Summary: When sympy expressions contain `sympy.nan` (e.g., in guard code for comparison ops with NaN inputs), the default sympy printer emits bare `nan`, which is not defined in the generated Python or C++ code. This causes `NameError: name 'nan' is not defined` at runtime. This mirrors the existing `_print_Infinity` and `_print_NegativeInfinity` overrides that were already in place: - `ExprPrinter`: raises `NotImplementedError` (safety net for subclasses) - `PythonPrinter`: returns `"math.nan"` - `CppPrinter`: returns `"std::numeric_limits<double>::quiet_NaN()"` Also adds unit tests for NaN, Infinity, and NegativeInfinity printing. Fixes 120 opinfo e2e test failures (20 per op) across 6 comparison ops: eq, ge, gt, le, lt, ne — all test cases involving NaN inputs. Test Plan: - `buck test //caffe2/test/inductor:test_indexing -- ExprPrinterTests` — all 17 tests pass (14 existing + 3 new) - e2e: `opinfo-test "eq; 2; 0; tensor, i8x23; float, nan"` — verifies NaN guard code generation works end-to-end Differential Revision: D99380088 Pull Request resolved: pytorch#179708 Approved by: https://github.com/Skylion007, https://github.com/blaine-rister
…172618) * Enable test_cpu_select_algorithm.py testing * Fix unexpected success by using `xfailIf(IS_ARM64)`. Pull Request resolved: pytorch#172618 Approved by: https://github.com/chuanqi129, https://github.com/jansel
Fixes pytorch#172390, pytorch#171263 Related issue pytorch#168042 This PR re-submits and is based on @galv’s PR pytorch#172396 Re-submit pytorch#173421 Refactor CUDA RNG state tracking from per-generator to per-(generator, capture_id). Previously, RNG capture state was tracked at generator scope. That model is not sufficient for concurrent or nested CUDA graph captures, because multiple captures using the same generator may need to maintain independent RNG timelines. This shows up most clearly with `torch.cond()` under CUDA graphs, where conditional branches create nested child captures. This change introduces explicit per-capture RNG state. Each capture gets its own `CUDAGeneratorCaptureState`, which owns the GPU RNG tensors for that capture and tracks its offset independently. `CUDAGeneratorState` now holds a `capture_states_` map keyed by `CaptureId_t`, and graph capture/replay paths are updated to use `capture_id` when creating and looking up capture state. As part of this refactor, `conditional_rng_snapshots_` is removed, since explicit per-capture state makes that mechanism unnecessary. * add `CUDAGeneratorCaptureState` for per-capture GPU RNG tensors and offset tracking * change `CUDAGeneratorState` to store `capture_states_`, keyed by `CaptureId_t` * create per-capture RNG state in `capture_begin` * use `capture_id` in `capture_end` and replay-time RNG handling * remove `conditional_rng_snapshots_` Resulting state layout: ```cpp struct CUDAGeneratorCaptureState { at::TensorBase rng_state_seed_extragraph_; at::TensorBase rng_state_offset_extragraph_; }; struct CUDAGeneratorState { // Host-side seed and offset uint64_t seed_; uint64_t philox_offset_per_thread_; // Per-capture RNG state ska::flat_hash_map< CaptureId_t, c10::intrusive_ptr<CUDAGeneratorCaptureState>> capture_states_; mutable std::mutex capture_states_mutex_; }; struct CUDAGraph { // Generator states captured by this graph and their whole-graph increments ska::flat_hash_map< c10::intrusive_ptr<at::CUDAGeneratorState>, uint64_t> captured_generator_states_; }; ``` Pull Request resolved: pytorch#176752 Approved by: https://github.com/ngimel
…ler (pytorch#179759) `TestProfiler::test_profiler` is failing in pytorch/kineto#1353. It checks that GPU profiling works by looking for `"gemm"` or `"Cijk"` substrings in kernel names. This is fragile. On the CI server it dispatches to a kernel whose name contains neither substring, causing a test failure. This PR replace the kernel-name check with a straightforward assertion that the profiler captured at least one GPU event, which is what this test actually intends to verify. Pull Request resolved: pytorch#179759 Approved by: https://github.com/scotts
This PR is auto-generated nightly by [this action](https://github.com/pytorch/pytorch/blob/main/.github/workflows/nightly.yml). Update the pinned vllm hash. Pull Request resolved: pytorch#179872 Approved by: https://github.com/pytorchbot Co-authored-by: Huy Do <huydhn@gmail.com>
The `check_public_api_test_fails` function in test.sh uses
`echo "${test_output}" | grep -q` to verify that the canary module was
detected. Under `-o pipefail` (set by the GHA shell), this is a race
condition: `grep -q` exits immediately on first match, but `echo` may
still be writing the large test output (~hundreds of lines of import
tracebacks), receiving SIGPIPE (exit 141). With `pipefail`, the pipeline
returns 141 instead of grep's 0, so the script incorrectly concludes the
canary wasn't found.
Replace the pipe with a here-string (`grep -q ... <<< "$var"`), which
passes content via a temporary file rather than a pipe, eliminating the
SIGPIPE race entirely.
Fixes flaky failures like:
https://github.com/pytorch/pytorch/actions/runs/24216395495/job/70700891103
### Testing
https://github.com/pytorch/pytorch/actions/runs/24262053518/job/70850088090?pr=179999
Authored with Claude.
Pull Request resolved: pytorch#179999
Approved by: https://github.com/malfet
ghstack dependencies: pytorch#179994
…nel gets too large (pytorch#180217) Fold expressions create deeply nested AST nodes proportional to the number of parameter pack elements. For combo kernels with hundreds of arguments, this can exceed clang's expression-nesting limit (256). Array initialization produces a flat initializer list that avoids the depth issue while preserving sequential evaluation order. Also removes the redundant `#include <lazy_triton_compile.h>` from codegen — `cuda.h` (the precompiled header source) already includes it, and the duplicate caused redefinition errors across the PCH boundary. Authored with: Claude Pull Request resolved: pytorch#180217 Approved by: https://github.com/karthickai
…ytorch#177809)" This reverts commit 1a4c456. Reverted pytorch#177809 on behalf of https://github.com/pytorch-auto-revert due to Reverted automatically by pytorch's autorevert, to avoid this behaviour add the tag autorevert: disable ([comment](pytorch#177809 (comment)))
Update Triton hash contains : - AMD Cherry Picks - Plugin extension enablement Pull Request resolved: pytorch#179971 Approved by: https://github.com/malfet, https://github.com/njriasan
This reverts commit eb9c9e7. Reverted pytorch#179872 on behalf of https://github.com/pytorch-auto-revert due to Reverted automatically by pytorch's autorevert, to avoid this behaviour add the tag autorevert: disable ([comment](pytorch#179872 (comment)))
Bumps [pillow](https://github.com/python-pillow/Pillow) from 12.1.1 to 12.2.0. - [Release notes](https://github.com/python-pillow/Pillow/releases) - [Changelog](https://github.com/python-pillow/Pillow/blob/main/CHANGES.rst) - [Commits](python-pillow/Pillow@12.1.1...12.2.0) --- updated-dependencies: - dependency-name: pillow dependency-version: 12.2.0 dependency-type: direct:production ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…id_lookup_table_entry_operation_addmm (pytorch#180316) Skip in-code since pytorch#179955 isn't working as expected: pytorch/test-infra#7963 Pull Request resolved: pytorch#180316 Approved by: https://github.com/jeffdaily Co-authored-by: Pruthvi Madugundu <pruthvigithub@gmail.com>
Summary: As title. We're making aot_compile publicly available to be used so for the completeness reason (and trivial to support) we added a test to ensure aot_eager backend can work well with .aot_compile() Test Plan: pytest test/dynamo/test_aot_compile.py Reviewers: Subscribers: Tasks: Tags: Pull Request resolved: pytorch#180008 Approved by: https://github.com/bobrenjc93
…pytorch#179917) Summary: Now .aot_compile() has been used by vLLM for a while by default, we should make this API more publicly available and have a common ground to share knowledge of the usage for this API. Test Plan: CI https://docs-preview.pytorch.org/pytorch/pytorch/179917/user_guide/torch_compiler/torch.compiler_aot_compile.html <img width="1550" height="1256" alt="Screenshot 2026-04-10 at 11 25 00 AM" src="https://github.com/user-attachments/assets/ee1c31ea-0de7-477f-8df8-d8db83e0dbab" /> Reviewers: Subscribers: Tasks: Tags: Pull Request resolved: pytorch#179917 Approved by: https://github.com/Lucaskabela
As titled Demo, on top of pytorch/kineto#1361 ``` with torch.profiler.profile( activities=[torch.profiler.ProfilerActivity.CUDA], experimental_config=torch._C._profiler._ExperimentalConfig(expose_kineto_event_metadata=True) ) as prof: x = torch.randn(500, 500, device='cuda') y = torch.randn(500, 500, device='cuda') result = torch.matmul(x, y) ``` ``` EventMetadata(registers_per_thread=84, shared_memory=12672, grid=[7, 7, 1], block=[64, 1, 1], priority=4, blocks_per_sm=0.331081, warps_per_sm=0.662162, occupancy=1.0, queued=0, graph_id=0, graph_node_id=0, stream=7, context=1, bytes=None, bandwidth_gb_s=None, collective_name=None, dtype=None, in_msg_nelems=None, out_msg_nelems=None, in_split_size=None, out_split_size=None, global_rank_start=None, global_rank_stride=None, group_size=None, process_group_name=None, process_group_desc=None, group_ranks=None, rank=None, src_rank=None, dst_rank=None, seq=None, is_async=None) ``` Pull Request resolved: pytorch#180100 Approved by: https://github.com/scotts
Replace the closure that iterates over all args checking requires_grad
with a codegen'd function that emits assertions only for the specific
indices where requires_grad was False at compile time. Positions where
requires_grad=True are safe and generate no runtime check.
Fixed tests to ensure at least one input has requires_grad=True so the
autograd dispatch path (where DebugAssertWrapper is applied) is taken.
Mixed requires_grad (arg 0 has grad, arg 1 does not):
def inner_fn(args):
if args[1].requires_grad: raise AssertionError(_msg_1)
return compiled_fn(args)
All requires_grad (no assertions needed):
def inner_fn(args):
return compiled_fn(args)
Some no_grad (args 1 and 2 are non-grad):
def inner_fn(args):
if args[1].requires_grad: raise AssertionError(_msg_1)
if args[2].requires_grad: raise AssertionError(_msg_2)
return compiled_fn(args)
Assert step in isolation (us/call):
| Case | Before (loop over all args) | After (codegen spot-check) | Speedup |
|---|---|---|---|
| 5 args, 1 no-grad | 0.55 us | 0.29 us | 1.9x |
| 10 args, 2 no-grad | 0.89 us | 0.43 us | 2.1x |
| 20 args, 5 no-grad | 1.67 us | 0.87 us | 1.9x |
The codegen eliminates the runtime loop over all args and the per-element
`can_require_grad is None` / `not can_require_grad` branching. Assertions
are only emitted for positions where requires_grad was False at compile
time; positions where requires_grad=True generate no runtime check at all.
Error messages are pre-formatted at compile time and baked in as globals.
Pull Request resolved: pytorch#179599
Approved by: https://github.com/Lucaskabela
…bucketing (pytorch#177132) Standard bucketing cat-s inputs into one buffer before a single RS call. The new "coalesced" mode passes tensors directly to reduce_scatter_tensor_coalesced, which wraps individual ncclReduceScatter calls in a single ncclGroupStart/End block — zero-copy and fewer NCCL kernel launches. This shows improvements on NVLink interconnects. The coalesced mode produces a different FX graph pattern than standard bucketing, which requires updates to overlap scheduling passes: Standard: rs_start -> wait Coalesced: rs_coalesced -> getitem[0] -> wait[0] -> getitem[1] -> wait[1] ... The overlap passes previously assumed a 1:1 mapping between bucketed start and wait nodes. With coalesced mode there is 1 start and N waits (one per original tensor), connected through getitem nodes. _get_collective_node_from_wait handles both patterns by looking through getitem to find the underlying collective. Authored with Claude. Pull Request resolved: pytorch#177132 Approved by: https://github.com/aditvenk
Co-authored-by: Yu, Guangye <guangye.yu@intel.com>
Co-authored-by: Yu, Guangye <guangye.yu@intel.com>
Co-authored-by: Yu, Guangye <guangye.yu@intel.com>
Co-authored-by: Yu, Guangye <guangye.yu@intel.com>
Co-authored-by: Frost Mitchell <frost.mitchell@intel.com>
Co-authored-by: Yu, Guangye <guangye.yu@intel.com>
pytorchmergebot
force-pushed
the
xccl-premul
branch
from
April 15, 2026 01:35
c78da70 to
aff61fb
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
20 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds Python binding for
_make_xccl_premul_sumto supportPREMUL_SUMreduce operation in the XCCL backend (Intel GPU).The
_make_xccl_premul_sumbinding directly reusesmakeNCCLPreMulSumbecause XCCL/NCCL have identical memory layouts.Currently, XCCL backend supports PREMUL_SUM for:
Not yet supported:
Implemented code in intel/torch-xpu-ops#1948, and base test case in intel/torch-xpu-ops#2690.