Repository navigation
Fix #1865 - #83
Open
laifenxiawucha wants to merge 1379 commits into
Open
Fix #1865#83laifenxiawucha wants to merge 1379 commits into
laifenxiawucha wants to merge 1379 commits into
Conversation
…zation` (pytorch#186630) Otherwise pytorch#184737 seems to expose a preexisting issue that causes `test_simple_sharing_cuda` to fail authored with codex Pull Request resolved: pytorch#186630 Approved by: https://github.com/ezyang
B200 runners used CUDA 13.0 docker images with Python 3.10, which caused install_cutlass_api to skip installation and left NV Universal GEMM smoke tests untested. Add py3.12 cuda13.0 docker image variants and point B200 workflows at them so cutlass_api can install on the runner. Authored with Cursor Composer 2.5 (composer-2.5-fast). Fixes pytorch#186991 Pull Request resolved: pytorch#186997 Approved by: https://github.com/huydhn
…ytorch#186962) Pull Request resolved: pytorch#186962 Approved by: https://github.com/desertfire
…ytorch#177927) Note: This issue is mentioned here: pytorch#175872 (comment). Not sure if there is an issue tracking it. - Many properties are called for each node (N) in the graph, where each property iterates over IO specs (M) with conditions. This sets up a O(NM) time complexity that is entirely avoidable and can be O(N+M). - inputs_to_parameters - inputs_to_buffers - buffers_to_mutate - parameters_to_mutate - user_inputs_to_mutate - Next fix is a O(to_module^2), where for each module it iterates all modules and calls _is_call_name. This function is expensive and compiles a regex each time. Instead, we can more directly iterate the numbered variants and quit when we run out. This assumes no holes exist in the numbering. The only use of _is_call_name is thus removed so the function is removed as well. **Before fixes:** | Layers | Export (s) | Decompose (s) | |--------|------------|----------------| | 64 | 0.38 | 0.85 | | 128 | 0.55 | 1.67 | | 256 | 1.14 | 3.16 | | 512 | 2.19 | 7.11 | | 1024 | 4.45 | 17.31 | | 2048 | 8.88 | 47.38 | | 4096 | 17.0 | 143.6 | Decompose: Grows super-linearly **After fixes:** | Layers | Export (s) | Decompose (s) | |--------|------------|----------------| | 64 | 0.52 | 0.81 | | 128 | 0.52 | 1.77 | | 256 | 1.28 | 2.88 | | 512 | 2.34 | 5.47 | | 1024 | 4.62 | 10.84 | | 2048 | 8.48 | 20.26 | | 4096 | 16.08 | 40.14 | Decompose: Grows linearly Pull Request resolved: pytorch#177927 Approved by: https://github.com/angelayi
…ytorch#183991) Pull Request resolved: pytorch#183991 Approved by: https://github.com/jeffdaily
…rch#187137) # Motivation pytorch#184735 introduces AOTI lazy autotuning with the dual-wrapper mechanism, which uncovered some new XPU test failures in CI. # Additonal Context Fix pytorch#187124 Fix pytorch#187123 Fix pytorch#187035 Fix pytorch#187035 Fix pytorch#187036 Fix pytorch#187125 Fix pytorch#187121 This needs to be cherry-picked to the release branch. Pull Request resolved: pytorch#187137 Approved by: https://github.com/EikanWang, https://github.com/atalman ghstack dependencies: pytorch#181233
…o driver (pytorch#187232) # Motivation Previously, we didn't handle pyzes import failures gracefully on machines without level zero driver. Without this PR: <img width="510" height="436" alt="image" src="https://github.com/user-attachments/assets/5c6db7ff-ee8d-4ebe-9177-db2ff80916c1" /> With this PR: <img width="530" height="179" alt="image" src="https://github.com/user-attachments/assets/9239ac61-3722-42d7-8411-68f61abd2e0a" /> # Additional Context <img width="1004" height="36" alt="image" src="https://github.com/user-attachments/assets/0035d2c6-2e49-4570-97de-3a4b48ae0575" /> This needs to be cherry-picked to the release branch. Pull Request resolved: pytorch#187232 Approved by: https://github.com/EikanWang, https://github.com/atalman ghstack dependencies: pytorch#181233, pytorch#187137
# Motivation This PR reverts pytorch#182630 to fix pytorch#187277 Root cause is `ext_oneapi_wait_and_throw` will introduce SYCL Graph to sync an invalid queue. # Additional Context This needs to be cherry-picked to the release branch. Why isn't it captured on CI? This issue is only found on BMG (Xe2), and the current CI is on Data Center GPU (Xe). Pull Request resolved: pytorch#187306 Approved by: https://github.com/EikanWang, https://github.com/atalman ghstack dependencies: pytorch#181233, pytorch#187137, pytorch#187232
…ytorch#187314) # Motivation When pytorch#184737 generalize `test_multiprocess.py`, it helps enable `TestMultiprocessingDeviceType` test on XPU. However, XPU doesn't support this test suite because lack of IPC support currently. So we would like to skip this test case and restore XPU CI to green. Fixes pytorch#186585 Code originates from pytorch#186793 Co-authored-by: @xuhancn Pull Request resolved: pytorch#187314 Approved by: https://github.com/EikanWang, https://github.com/atalman ghstack dependencies: pytorch#181233, pytorch#187137, pytorch#187232, pytorch#187306
) `dispatch` and `combine` are adjoint operations: backward of dispatch calls combine, and backward of combine calls dispatch. Single dispatch/combine public API following `nn.Module`/`torch.matmul` conventions: - With `out=(out_tokens, out_weights, out_idx)`: writes to caller-supplied buffers and returns them; no autograd (efficient buffer reuse). - Without `out`: allocates buffers internally, returns them with full autograd support via `_DispatchAutograd` / `_CombineAutograd`. Subclasses implement `_dispatch`/`_combine` as the raw buffer-writing primitives and inherit both modes from the base class. `topk_weights` receives no gradient (routing metadata at a different byte width). Authored with Claude. Pull Request resolved: pytorch#181314 Approved by: https://github.com/kapilsh ghstack dependencies: pytorch#178712
The NCCL EP bindings move out of libtorch_cuda / libtorch_python into their own optional Python
extension, torch._nccl_ep, imported lazily by torch/distributed/_token_switch.py. libtorch_cuda no
longer references ncclEp*, so torch imports with or without EP, and a missing EP build is a clean
ImportError rather than a core-torch load failure.
The extension follows how torch links its (statically embedded, submodule) NCCL: libnccl_ep.a is built
from source and statically linked into the extension, so its ncclEp* are embedded and its nccl* bind to
torch's own NCCL. The runtime JIT header paths are baked in from the submodule build dir. The extension
is fully self-contained -- no libnccl_ep.so, no preload, no nccl4py. (USE_NCCL_EP keeps requiring
NOT USE_SYSTEM_NCCL here; a follow-up enables the USE_SYSTEM_NCCL=ON / wheel path via nccl4py.)
Review order:
1. nccl_ep_pybind.cpp (new): the PYBIND11_MODULE(_nccl_ep) bindings, moved from init.cpp; sets
NCCL_EP_HOME/NCCL_HOME from the baked JIT home at import.
2. init.cpp / build_variables.bzl / caffe2/CMakeLists.txt / Dependencies.cmake: remove the bindings,
drop nccl_ep.cu from torch_cuda, stop linking __caffe2_nccl_ep into torch_cuda.
3. cmake/External/nccl_ep.cmake: build libnccl_ep.a against the submodule NCCL; bake NCCL_EP_JIT_HOME.
4. torch/CMakeLists.txt: the Python_add_library(_nccl_ep) target -- static-link + bake JIT home +
$ORIGIN/lib RPATH.
5. NCCLUtils.hpp: export getNcclDataType (TORCH_API) -- the extension uses it.
6. _token_switch.py: import torch._nccl_ep (self-contained).
7. test_token_switch.py: gate on the extension being importable.
Test Plan:
USE_SYSTEM_NCCL=OFF on 4xH100, with NO nccl4py and NO LD_PRELOAD -- only the CUDA toolkit (nccl-ep JITs
kernels at runtime):
USE_NCCL_EP=1 TORCH_CUDA_ARCH_LIST="9.0" pip install --no-build-isolation -v -e .
CUDA_VISIBLE_DEVICES=0,1,2,3 python test/distributed/test_token_switch.py -v
# Ran 7 tests ... OK; the extension has no libnccl_ep.so NEEDED (statically embedded ncclEp*).
Pull Request resolved: pytorch#187366
Approved by: https://github.com/d4l3k
ghstack dependencies: pytorch#178712, pytorch#181314
… backend. (pytorch#187040)" This reverts commit 29dbb32. Reverted pytorch#187040 on behalf of https://github.com/etaf due to Regression on XPU ([comment](pytorch#187040 (comment)))
Upper bound was originally introduced by pytorch#174631, but `pkg_resources` is not used by `torch.utils.cpp_extensions` Test Plan: Install latest setuptools and all `cpp_extensions` test still work Fixes pytorch#187188 Pull Request resolved: pytorch#187262 Approved by: https://github.com/malfet, https://github.com/atalman Co-authored-by: Nikita Shulga <2453524+malfet@users.noreply.github.com>
…ly, and histc ops (pytorch#179202) Add single_dim_strategy implementations for: - Category G: Pooling ops (batch/channel dim sharding, partial for avg pools) - Category H: Replicate-only ops (SVD, QR, diagonal, tril, triu, upsample, etc.) - Category L: Histc (partial output when min/max specified) New registrations are commented out; old register_op_strategy kept active. Pull Request resolved: pytorch#179202 Approved by: https://github.com/pianpwk ghstack dependencies: pytorch#179068, pytorch#179200, pytorch#179201
…186997)" This reverts commit e8f1ea4. Reverted pytorch#186997 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#186997 (comment)))
… output_mask (pytorch#187383) **Summary:** PR pytorch#187199 tightened single-dim strategy handling for `None` output specs, which exposed latent bugs in several `math_ops` single-dim backward strategies. These strategies were emitting placements for outputs that were disabled by `output_mask`, causing invalid mixed`None`/placement specs on multi-dimensional meshes, as reported in pytorch#187373. Update the affected single-dim strategies to respect `output_mask`: - `native_layer_norm_backward` - `_fused_rms_norm_backward` - `grid_sampler_{2,3}d_backward` - `native_batch_norm_backward` For masked-off outputs, the strategies now emit `None` even when the corresponding input tensor still exists. Input placements are preserved separately, so cases like “weight input exists but d_weight output is masked off” are handled correctly. **Test Cases** 1. pytest /data/users/anshulsi/pytorch/test/distributed/tensor/test_op_strategy.py -k test_layer_norm_bwd_single_dim_strategy 2. pytest /data/users/anshulsi/pytorch/test/distributed/tensor/test_op_strategy.py -k test_rms_norm_bwd_single_dim_strategy Pull Request resolved: pytorch#187383 Approved by: https://github.com/aditvenk
…arams (pytorch#187376) **Summary:** fixes pytorch#186998. FSDP2 mixed precision initializes a param group’s `_orig_dtype` and `_reduce_dtype` during the first forward. If that first forward happens while all params are frozen, there are no trainable params to derive dtype metadata from, so both fields were left as `None`. After later unfreezing, gradient reduction could then fall back to the bf16 compute dtype instead of the requested fp32 reduce dtype. This updates dtype resolution so that when no params are currently trainable, FSDP2 falls back to the group’s floating-point params if their original and reduce dtypes are uniform. Mixed all-frozen groups still preserve the previous behavior by resolving to `None`. **Test Case** 1. pytest test/distributed/_composable/fsdp/test_fully_shard_mixed_precision.py -k test_reduce_dtype_after_frozen_first_forward Pull Request resolved: pytorch#187376 Approved by: https://github.com/weifengpy
Stacked PRs: * pytorch#187273 * pytorch#183838 * __->__#187231 --- --- --- ### [PT2] Preserve symbolic metadata across tracing EP-overlap graph chunking traces symbolic tensor dimensions and scalar captures through ProxyTensor, make_fx, and Regional Inductor. These paths should preserve metadata at the boundary where values become FX proxies instead of rediscovering it later in FlexAttention or another downstream subsystem. ProxyTensor records fresh unbacked symbols on FX nodes via node.meta["unbacked_bindings"]. For nested proxy pytrees, route each binding along its full output path to the proxy that owns the tensor leaf; if the output proxy is already a single Proxy, that proxy owns the whole output and keeps the original path. This preserves provenance for multiple inputs that share token-grid dimensions. ProxyTensor also records meta["val"] for primitive scalar placeholders when those scalars already have FX proxies. This is the scalar analogue of the existing tensor and SymInt placeholder metadata path, and keeps scalar captures visible to subgraph consumers without Flex-specific placeholder patching. Regional Inductor compiles temporary submodule graph modules with symbolic fake metadata. Passing donate_graph_module=True avoids deepcopying graph modules that Regional Inductor already owns; deepcopying symbolic fake tensor metadata can call numel() on symbolic sizes and strides. Fixes pytorch#187230 This PR was authored with assistance from an AI assistant. Test Plan: ```bash python -m pytest test/test_proxy_tensor.py -q -s -k "scalar_placeholder_meta or unbacked_bindings_on_multiple_token_grid_placeholders" ``` ```bash python -m pytest test/dynamo/test_regional_inductor.py -q -s -k "RegionalInductorTests and test_max_autotune_no_cudagraphs" ``` ```bash lintrunner -a ``` Pull Request resolved: pytorch#187231 Approved by: https://github.com/ydwu4
Based on pytorch#180850, this PR updates AArch64 CI toolchains to GCC15. The necessary GCC15 enablement fixes are part of pytorch#187167 and pytorch#185633. Pull Request resolved: pytorch#182016 Approved by: https://github.com/aditew01, https://github.com/Skylion007 Co-authored-by: Aditya Tewari <aditya.tewari@arm.com>
This PR increases the OpenBLAS thread limit from 128 to 256 to address the warning: `OpenBLAS warning: precompiled NUM_THREADS exceeded, adding auxiliary array for thread metadata.` The warning occurs when the number of threads exceeds the precompiled NUM_THREADS value. By raising the limit to 256, this change prevents the warning and avoids the overhead of allocating auxiliary thread metadata when running with higher thread counts. The PR fixes this issue pytorch#174373 ### OpenBLAS `NUM_THREADS=128` vs `NUM_THREADS=256` I compared the focused TorchBench CPU fp32 subset using two PyTorch builds: one built with OpenBLAS `NUM_THREADS=128` and one with `NUM_THREADS=256`. | Model | 128 latency | 256 latency | 256 vs 128 | |---|---:|---:|---:| | nvidia_deeprecommender | 132.822 | 132.983 | +0.12% | | dlrm | 20.616 | 20.613 | -0.02% | | BERT_pytorch | 75.753 | 75.947 | +0.26% | | hf_Roberta_base | 323.586 | 323.764 | +0.06% | | hf_Bert_large | 381.067 | 381.698 | +0.17% | | hf_T5 | 103.504 | 103.720 | +0.21% | | hf_T5_generate | 1278.884 | 1300.140 | +1.66% | | resnet50 | 418.311 | 418.352 | +0.01% | | densenet121 | 651.161 | 651.172 | +0.00% | | timm_efficientnet | 177.996 | 178.226 | +0.13% | Overall, this does not show a meaningful performance improvement from increasing OpenBLAS `NUM_THREADS` from 128 to 256 for this benchmark set. Most differences are within noise. This suggests changing `NUM_THREADS` to 256 is safe from a performance perspective. Pull Request resolved: pytorch#177656 Approved by: https://github.com/fadara01, https://github.com/aditew01, https://github.com/jansel
…7359) This should have been in pytorch#187204; if we don't drop the reference to the source file here, things won't compile when we remove the source files in Kineto. Pull Request resolved: pytorch#187359 Approved by: https://github.com/ryanzhang22
…numel <= 64 on HIP (pytorch#183364) - Default XBLOCK=2 launches too few programs to fill CUs for tiny-x reductions, - and ** under reduction_hint=INNER ** is the only candidate returned. - Force XBLOCK=1 on HIP when xnumel <= 64 so the autotuner has a candidate that can saturate the device. - CUDA path unchanged. Pull Request resolved: pytorch#183364 Approved by: https://github.com/eellison, https://github.com/haoyuz
gcc5.4 to gcc-11 clang-8 to clang-18 C++14 to C++20 Pull Request resolved: pytorch#187395 Approved by: https://github.com/janeyx99
ConstantVariable.is_python_equal blindly called other.as_python_constant(), which raised AsPythonConstantNotImplementedError whenever a dict/set key comparison hit a non-constant key (a user-defined object with a custom __eq__). This diverged from CPython, where dict/set membership and equality go through PyObject_RichCompareBool and run the user-defined __eq__. The fix mirrors PyObject_RichCompareBool: when `other` is not a python constant, instead of forcing as_python_constant(), route the comparison through generic_richcompare_bool so the user-defined __eq__ runs and any exception it raises is observed by the traced program. Test Plan: ``` CUDA_VISIBLE_DEVICES= PYTORCH_TESTING_DEVICE_ONLY_FOR=cpu PYTORCH_TEST_WITH_DYNAMO=1 \ pixi run -w pytorch -e pytorch313 python -m pytest \ test/cpython/v3_13/test_dict.py::DictTest::test_eq -q ``` ``` CUDA_VISIBLE_DEVICES= PYTORCH_TESTING_DEVICE_ONLY_FOR=cpu \ pixi run -w pytorch -e pytorch313 python -m pytest \ test/dynamo/test_dicts.py -k test_cmp_eq_key_raises -q ``` ``` CUDA_VISIBLE_DEVICES= PYTORCH_TESTING_DEVICE_ONLY_FOR=cpu PYTORCH_TEST_WITH_DYNAMO=1 \ pixi run -w pytorch -e pytorch313 python -m pytest \ test/cpython/v3_13/test_dict.py -q ``` All pass; the previously-xfailing CPython313-test_dict-DictTest.test_eq sentinel is removed. This change was authored with the help of an AI assistant. Pull Request resolved: pytorch#186669 Approved by: https://github.com/rtimpe, https://github.com/williamwen42
…184407) ## Summary Port the Windows x86_64 wheel CD build off the legacy `.bat` chain onto a Python pipeline, mirroring the Linux split landed in pytorch#182409. CPU, CUDA, and XPU wheel builds now run through a small set of Python scripts instead of the `build_pytorch.bat -> cpu/cuda/xpu.bat -> internal/setup.bat` chain. The new entry chain on x86_64 wheel builds: ``` binary_windows_build.sh -> windows/internal/vc_install_helper.bat (ensure VS2022 BuildTools) -> windows/set_desired_python.sh (install Python.org Python) -> python windows/build_env_setup.py (vcvarsall + CUDA/XPU env + MAGMA + sccache) -> python windows/build_install_deps.py (pip deps + libuv) -> python windows/build_wheel.py (python -m build --wheel) ``` The three Python steps are driven through the interpreter `set_desired_python.sh` selected (exported as `DESIRED_PYTHON_EXE`), not a bare `python`, so the wheel ABI tag and pip build deps track `DESIRED_PYTHON` for free-threaded builds (see Build-correctness fixes). vcvarsall env capture mirrors `torch.utils.cpp_extension._jit._get_vc_env` from pytorch#180707 (vswhere discovery, `cmd /u /c` for UTF-16LE output, diff against current env so additions layer on top). `internal/cuda_install.bat`, `internal/xpu_install.bat`, and `internal/vc_install_helper.bat` are kept as Python/bash-invoked installers -- same pragmatic precedent as the Linux refactor (`install_cuda.sh` etc.). Heavy installer logic stays in bash/cmd; orchestration moves to Python. `set_desired_python.sh` inlines what `internal/install_python.bat` did (the `bash -> cmd /c` boundary was unreliable for the python.org installer), so the wheel path no longer invokes `install_python.bat`. ## Removing the legacy bat chain The x86_64 wheel build was the last consumer of the legacy `build_pytorch.bat -> {cpu,cuda,xpu}.bat -> internal/setup.bat` chain, so porting it to the Python pipeline lets that chain go. This PR deletes: - `build_pytorch.bat`, `cpu/cuda/xpu.bat`, `setup_build.bat`, and `internal/{build_wheels,check_deps,check_nvtx,check_opts,clean,clone, cuda_config,init,setup}.bat`. These are now orphaned -- nothing in the tree invokes them. The only remaining mentions are provenance comments in the new Python scripts (e.g. "Mirrors the MAGMA block in the legacy build_pytorch.bat"), plus the separate `arm64/build_pytorch.bat` and `win-test-helpers/build_pytorch.bat`, which are different files left untouched. Note on debug libtorch: upstream has already stopped shipping the Windows debug libtorch builds and removed their nightly workflows, and this branch is rebased on top of that -- so it carries no `.github/workflows` or `generate_ci_workflows.py` changes. (Earlier revisions of this PR removed those workflows here; those commits were dropped in the rebase once the upstream removal landed.) The release path is unaffected: release libtorch is still produced by extraction from the release wheel (`extract_libtorch_from_wheel.py`, the same mechanism as on Linux), and those wheels are now built by this Python pipeline. ## Build-correctness fixes (surfaced by the wheel smoke tests) With the builds green, the `wheel-*-test` smoke jobs ran and caught two gaps the legacy bat chain had handled implicitly: - **Free-threaded wheels (`DESIRED_PYTHON=3.14t`).** A free-threaded python.org install ships both `python.exe` and `python3.14t.exe` under the same prefix, so a bare `python` resolved to the regular interpreter and `python -m build` produced a `cp314` wheel that the `cp314t` test job rejected ("not a supported wheel on this platform"). Fix: export the selected interpreter as `DESIRED_PYTHON_EXE` and drive all three pipeline steps through it, so deps and the wheel ABI tag match `DESIRED_PYTHON`. - **XPU MKL availability.** `CMAKE_PREFIX_PATH` was assigned `<sys.prefix>/Library` (the pip `mkl-static`/`mkl-include` prefix) before the per-device setup ran, and the XPU oneAPI `vars.bat` then overwrote it with its own value -- so the XPU build configured BLAS as Eigen and `torch.backends.mkl.is_available()` was False (the smoke test asserts on it). cpu/cuda were unaffected because nothing else touches `CMAKE_PREFIX_PATH` there. Fix: set `CMAKE_PREFIX_PATH` after the device dispatch and prepend the MKL prefix to whatever the device contributed, so both the pip MKL and the oneAPI compiler prefix stay on the search path. ## Python patch pinning `set_desired_python.sh` selects the python.org installer from an explicit per-minor patch table (e.g. 3.13 -> 3.13.14) rather than the `.0` patch the legacy `install_python.bat` used. The `.0` pin froze each minor on its oldest, CVE-laden release; floating to the newest installer was rejected because python.org serves source-only (no-installer) releases for security-only minors, which would require scraping the release index and HEAD-probing candidates on every build. Security-only minors (3.10-3.12) are pinned to their last release that shipped a Windows installer; bugfix minors (3.13, 3.14) point at the current newest installer and are bumped here in review. The table must stay in sync with the matrix in `generated-windows-binary-wheel-nightly.yml`. ## Divergences from the Linux template (pytorch#182409) Where this PR deviates from the Linux pipeline, why: - **No `repair_wheel.py` step.** Windows DLL discovery uses PATH / `AddDllDirectory` rather than RPATH, so there's nothing for patchelf to do. CUDA DLLs already ship inside `torch/lib/`. Wheel-content diff against the latest published nightly cpu wheel confirms parity: identical DLL set in `torch/lib/`. - **`set_desired_python.sh` installs Python from python.org** rather than picking from a baked image (the Windows CI AMI doesn't ship Pythons preinstalled). - **`build_env_setup.py` downloads MAGMA / sccache / randomtemp** during the build rather than baking them into a Docker image. (These downloads now go through `urllib` with default TLS verification, replacing the legacy `curl -k` which disabled cert checking.) - **MKL comes from pip** (`mkl-static` / `mkl-include`), put on `CMAKE_PREFIX_PATH` via the `<sys.prefix>/Library` prefix, rather than the legacy `check_opts.bat` `LIB`/`MKLProductDir` system-MKL prepend. - **No cuda-aarch64 / rocm / cpu-cxx11-abi branches** (not applicable on Windows). - **`build_install_deps.py` hardcodes the pip dep list** (matches what `setup_build.bat` did) rather than reading `requirements*.txt` like Linux. Migrating to `requirements*.txt` is a worthwhile follow-up but out of scope here. The binary-build env knobs the legacy chain set are preserved in `COMMON_BUILD_ENV` (`PYTORCH_BINARY_BUILD`, `TH_BINARY_BUILD`, `BUILD_TEST=0`, `INSTALL_TEST=0`, `MSSdk`, `DISTUTILS_USE_SDK`). ## Out of scope - `windows-arm64` wheels (separate `arm64/build_pytorch.bat` chain, untouched). - Windows pull-CI build (uses a separate `win-test-helpers/build_pytorch.bat`, untouched). - `cpu-s390x` (Linux only). ## Test plan Builds are green across the full x86_64 wheel matrix: - [x] x86_64 cpu wheel build x Python 3.10-3.14 + 3.14t - [x] x86_64 cuda wheel build (12.6, 13.0, 13.2) x the same Python set - [x] x86_64 xpu wheel build x the same Python set - [x] Wheel content verified equivalent to the published nightly cpu wheel (same DLL set in `torch/lib/`) - [x] sccache enabled on CPU + CUDA host compilation; server lifecycle confirmed via build logs - [x] cpu / cuda smoke tests (`wheel-*-test`) green - [x] windows-arm64 wheel + libtorch builds unaffected (path goes through `arm64/build_pytorch.bat`, untouched) Smoke-test status for the two build-correctness fixes above (free-threaded ABI tag, XPU MKL) is being re-validated on the current head; the prior run caught both at the `wheel-*-test` stage and the fixes target exactly those failures. Earlier-run note (resolved): the `cuda13_0` / `cuda13_2` wheel builds had failed at the cuDNN download step because the `cudnn-...-9.23.1.3_cuda13-archive.zip` object was briefly missing from `ossci-windows` (uploaded ~19 min after the build fetched it). It is present and valid now and the cuda13 builds are green; unrelated to this PR. ## Follow-ups (deferred, captured separately) - **Migrate `build_install_deps.py` to read `requirements*.txt`** rather than hardcoded `PIP_PACKAGES`, matching the Linux pattern. - Remaining self-review cleanup items: `CMAKE_GENERATOR` explicit for CPU/XPU, env-handoff fragilities, PATH conversion guard, numpy pin for a future Python 3.15, and bumping the python patch table / keeping it in sync with the wheel matrix. Authored with [Claude Code](https://claude.com/claude-code). Pull Request resolved: pytorch#184407 Approved by: https://github.com/atalman
Reduce the wait time for a long running profiling test that's causing test pollution. Most of the tests in this file are single-shot, no-warmup profiling tests, which makes them super susceptible to test pollution and other weird behavior. It looks like when we have a long wait/sleep duration where no CUDA events are happening after Kineto has been initialized, some events get filtered out during out-of-range checks. This issue goes away when you add a warmup to the polluted test. This makes me think that there's something wrong with the timestamp checks in Kineto under this specific scenario that goes away when you allow a warmup -> active transition, but I haven't had time to debug it yet. For now, the simplest thing to do is to reduce the duration we are idle for. Pull Request resolved: pytorch#187363 Approved by: https://github.com/scotts
…ytorch#186323) (pytorch#186323)" This reverts commit aaa39f0. Reverted pytorch#186323 on behalf of https://github.com/huydhn due to The doc build failure is legit ([comment](pytorch#186323 (comment)))
…ge (pytorch#186173) Summary: Pipeline parallelism issues all send/recv on a single communicator. NCCL-group coalescing (batch_isend_irecv -> one ncclGroupStart/End) makes a single mixed send+recv batch deadlock-free, but it does not remove ordering hazards across separate batches. This adds per-direction PP communicators. When enabled, forward ( r->r+1) and backward (r->r-1) P2P are carried on two separate communicators derived from the stage's group via split_group. Test Plan: - `test_pp_per_direction_pipeline_stage.py` - `test_pp_torchcomms_p2p_deadlock.py` Pull Request resolved: pytorch#186173 Approved by: https://github.com/d4l3k, https://github.com/sanketpurandare
I find the new paragraph makes a big deal for readability for me, esp when I am reviewing PRs. I tested it as an explicit prompt on a couple of PRs that were not authored by me. Pull Request resolved: pytorch#187355 Approved by: https://github.com/drisspg
# Summary
Before this Pr if a user wrote a mask_mod that had dynamic scalars (quite common) we would disable flex/flash.
Example:
```py
def causal_offset(b,h,q,kv):
return q_idx + offset >= kv_idx
```
This would either require a compile kernel per implementation, thats bad / you will spend your whole life recompiling!
With this we can write mask/score mods that are dynamic over python ints and floats instead of needing to recompute. The most common case is for decode attention. Today that is done via 1 off seq_len info but in general this could have been done via generic query and key size or a dynamic integer offset ( handy if you cant derive from input tensor sizes).
We end up lowering to something like;
```Py
@cute.jit
def shifted_causal_mask(batch, head, q_idx, kv_idx, seqlen_info, aux_tensors, aux_scalars):
q_offset = cutlass.Int32(aux_scalars[0])
return kv_idx <= q_idx + q_offset
out = flash_attn_func(
q,
k,
v,
mask_mod=shifted_causal_mask,
aux_scalars=[cutlass.Int32(cache_position)],
)
```
Much of the details were making this work around the existing optimized implementations for generating efficient score and mask mods.
<img width="718" height="323" alt="image" src="https://github.com/user-attachments/assets/2532f174-12e7-45fa-8c39-5b35cdffb6fb" />
The above table also translated to aux_mask_mods. We need to make some updates to the packed analayzer to distinguish between captured tensors -> these neeed the existing logic, and captured scalars -> these will be static over the lanes for vectorization.
Also bumped CI to b17 which has the landed aux_data path in flash
Pull Request resolved: pytorch#186019
Approved by: https://github.com/liangel-02
The base op drives each (parameter, peer) shard with a single CTA, so a bucket of few large shards leaves most of the GPU idle, while both the LSA-push and multimem bandwidth scale with the CTA count. This flattens the work into fixed-size byte tiles (at least `AG_MIN_BYTES_PER_TILE`, the largest shard capped at `AG_MAX_CTAS` tiles) and grid-strides CTAs over the flat tile space, load-balancing shards of any size -- including skewed buckets -- without a per-shard grid dimension. On the push path the destination peer is the fast-varying dimension, so consecutive CTAs keep every outbound NVLink busy. Effective output bandwidth (8x H100, bf16, GB/s); `before` = one CTA per `(parameter, peer)`: | shape | multimem (before -> after) | LSA push (before -> after) | | --- | --- | --- | | 1 split | 171 -> 350 (2.0x) | 76 -> 289 (3.8x) | | 2 splits | 313 -> 352 (1.1x) | 146 -> 289 (2.0x) | | 8 splits | 330 -> 355 (1.1x) | 270 -> 289 (1.1x) | | **1 huge + 8 tiny (skew)** | **180 -> 369 (2.0x)** | **77 -> 301 (3.9x)** | | 64 x 64KB | 245 -> 302 (1.2x) | 218 -> 212 (~parity) | | 256 x 4KB | 58 -> 178 (3.0x) | 93 -> 75 (latency-bound) | ## Test Plan Built and ran on an H100 node with 4 ranks, both paths (multimem and the LSA push fallback via `NCCL_NVLS_ENABLE=0`); `split_sizes` includes a large + tiny case (`[16384, 32]`) that exercises the per-shard tiling. ``` CUDA_VISIBLE_DEVICES=0,1,2,3 \ python test/distributed/test_nccl.py -v -k all_gather_offset CUDA_VISIBLE_DEVICES=0,1,2,3 NCCL_NVLS_ENABLE=0 \ python test/distributed/test_nccl.py -v -k all_gather_offset ``` All 16 cases pass on each path. Pull Request resolved: pytorch#187750 Approved by: https://github.com/weifengpy ghstack dependencies: pytorch#187642
…rch#187752) torch.utils.hipify already maps cublasSetMathMode/cublasGetMathMode but not the cublasMath_t type or its enum values, and the CUBLAS_COMPUTE_* set omits CUBLAS_COMPUTE_16F. Sources using these (for example a HIP-ported extension calling cublasGemmEx with CUBLAS_COMPUTE_16F, or setting a cublasMath_t mode) do not hipify cleanly and need manual per-project aliases. Add the direct name-preserving mappings to CUDA_IDENTIFIER_MAP; every target exists in ROCm hipBLAS (hipblas.h: hipblasMath_t and HIPBLAS_COMPUTE_16F). CUBLAS_TENSOR_OP_MATH maps to HIPBLAS_TENSOR_OP_MATH, the direct analogue (deprecated in hipBLAS). This change was authored with the assistance of Claude, an AI assistant. Pull Request resolved: pytorch#187752 Approved by: https://github.com/jeffdaily, https://github.com/Skylion007 Co-authored-by: Jeff Daily <jeff.daily@amd.com>
…87610) Move the main Linux CI image (.ci/docker/ubuntu/Dockerfile) off conda and onto the same deadsnakes-plus-venv approach already used by the linter images via `install_python.sh`. `install_python.sh` is generalized with a `VENV_PATH` override (defaulting to the existing linter location so those images are unchanged), optional freethreaded interpreter selection, and the DOCS requirements step; the ubuntu image points it at /opt/python-${PYTHON_VERSION}-venv and bakes that venv onto PATH. MAGMA, previously pulled from the conda env, is now installed via the existing install_magma.sh into the CUDA toolkit dir with MAGMA_HOME set. The conda-only packages llvmdev, libpython-static, sysroot, libstdcxx-ng, icu, and cmake are dropped (cmake already comes from requirements-ci.txt; llvmdev/libpython-static were torch::deploy legacy with no remaining references). build.sh mirrors ANACONDA_PYTHON_VERSION into PYTHON_VERSION so the ubuntu image gets a version while the rocm/xpu images keep conda. test.sh drops the sm86 conda-libstdc++ LD_LIBRARY_PATH special case, which is now a no-op. The rocm/xpu images and the pre-built XLA test image still use conda and are intentionally left untouched. Authored-with: Claude (Claude Code) Pull Request resolved: pytorch#187610 Approved by: https://github.com/seemethere, https://github.com/huydhn ghstack dependencies: pytorch#187608, pytorch#187623
…ytorch#178495) ## What does this PR do? Fix a crash when calling vmap(f, out_dims=-1) where f's output does not depend on the vmapped input. ### Root Cause The root cause was a broken shape calculation in the C++ _remove_batch_dim function. The line `expanded_sizes.insert(expanded_sizes.begin() + out_dim, batch_size)` computed an invalid insertion index for out_dims=-1 with input-independent outputs, leading to a corrupted tensor shape. ### Solution This change fixes the issue by adding proper dimension wrapping: Calculate the number of dimensions of expanded_sizes, compute a valid wrapped dimension, insert the batch size at the wrapped index, then expand the tensor with the corrected shape. ``` int64_t ndim = expanded_sizes.size(); int64_t wrapped_out_dim = at::maybe_wrap_dim(out_dim, ndim + 1); expanded_sizes.insert(expanded_sizes.begin() + wrapped_out_dim, batch_size); auto result = self.unsqueeze(wrapped_out_dim).expand_symint(expanded_sizes); ``` ### Related Issue Fixes pytorch#178425 ## Test Cases Added a unit test to verify the fix (in `test/functorch/test_vmap.py/TestVmapAPI`): ``` def test_vmap_out_dims_negative_one_independent_output(self): t = torch.randn(2, 3) t_scalar = torch.randn([]) def f_dep(x): return x def f_ind(x): return t def f_ind_scalar(x): return t_scalar res_dep_neg1 = vmap(f_dep, in_dims=0, out_dims=-1)(t) self.assertEqual(res_dep_neg1.shape, torch.Size([3, 2])) res_ind_neg1 = vmap(f_ind, in_dims=0, out_dims=-1)(torch.zeros(1)) self.assertEqual(res_ind_neg1.shape, torch.Size([2, 3, 1])) res_ind_scalar_neg1 = vmap(f_ind_scalar, in_dims=0, out_dims=-1)(torch.zeros(5)) self.assertEqual(res_ind_scalar_neg1.shape, torch.Size([5])) ``` ``` import torch from torch.func import vmap t = torch.randn(2, 3) # Dependent output: out_dims=-1 works def f_dep(x): return x print(vmap(f_dep, in_dims=0, out_dims=-1)(t).shape) # (3, 2) ✓ # Independent output: out_dims=0 works def f_ind(x): return t print(vmap(f_ind, in_dims=0, out_dims=0)(torch.zeros(1)).shape) # (1, 2, 3) ✓ # Independent output: out_dims=-1 crashes print(vmap(f_ind, in_dims=0, out_dims=-1)(torch.zeros(1)).shape) # expected (2, 3, 1) torch.Size([3, 2]) torch.Size([1, 2, 3]) torch.Size([2, 3, 1]) ``` Pull Request resolved: pytorch#178495 Approved by: https://github.com/kshitij12345, https://github.com/frgossen, https://github.com/zou3519 Co-authored-by: Richard Zou via mergedog <zou3519@gmail.com>
…ytorch#187610)" This reverts commit d10f9af. Reverted pytorch#187610 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#187610 (comment)))
The negative assertion in test_export_strict_narrow_unbacked_expr was checking for a fake tensor cache-era failure instead of export semantics. The underlying storage nbytes path now avoids the symbolic data-dependent comparison before this test runs, so Module(identity) is expected to export successfully both with and without the fake tensor dispatch cache. Update the test to assert the intended successful export directly, while keeping the existing sym_fresh_size case for the expression-to-unbacked-size path. Fixes pytorch#154574 Generated by my agent Test Plan: - TORCH_FAKE_TENSOR_DISPATCH_CACHE=0 pytest test/export/test_export.py::TestDynamismExpression::test_export_strict_narrow_unbacked_expr - pytest test/export/test_export.py::TestDynamismExpression::test_export_strict_narrow_unbacked_expr - lintrunner -a Pull Request resolved: pytorch#185500 Approved by: https://github.com/zou3519
Dynamo decomposed torch._foreach_lerp_ whenever it saw the scalar overload. That helped tensor weights, where the native lowering can force scalar extraction, but it also decomposed ordinary Python scalar weights into _foreach_sub, per-tensor weight materialization, and _foreach_addcmul_. For optimizers such as RAdam this created extra full-size temporaries that pure eager never allocates, so torch.compile(backend="eager") reported a higher peak CUDA allocation. Keep the decomposition only for Tensor weights. Python scalar weights can be represented directly in the FX graph as the native foreach_lerp_ call, which matches eager's lifetime behavior and avoids the additional temporaries while preserving the tensor-weight graph-break fix. Benchmark Results: - CUDA RAdam peak-memory repro before fix: eager [220160, 220160, 220160], compile/eager [256000, 256000, 256000] - CUDA RAdam peak-memory repro after fix: eager [220160, 220160, 220160], compile/eager [220160, 220160, 220160] - CUDA NAdam peak-memory repro after fix: eager [266240, 266240, 266240], compile/eager [266240, 266240, 266240] Test Plan: - python test/dynamo/test_dynamo_decompositions.py -k "foreach_lerp" - python test/dynamo/test_dynamo_decompositions.py - python test/dynamo/test_functions.py -k "foreach_lerp" - CUDA peak-memory reproduction script from issue pytorch#125511 - git diff --check - lintrunner torch/_dynamo/variables/torch.py test/dynamo/test_dynamo_decompositions.py - lintrunner -a (fails on pre-existing clang-tidy findings in untouched C++ files) Fixes pytorch#125511 Generated by my agent Pull Request resolved: pytorch#186452 Approved by: https://github.com/Skylion007, https://github.com/mlazos
…ytorch#187702) Fixes pytorch#135424 ### Summary `vmap`-ing a function that calls `torch.repeat_interleave` over a **batched `repeats` tensor** previously failed with: ``` RuntimeError: ... no batching rule implemented for aten::repeat_interleave.Tensor ... ``` The `repeat_interleave.self_Tensor` / `.self_int` overloads are already decomposed for functorch (`BatchRulesDecompositions.cpp`), but both bottom out in the `repeat_interleave.Tensor(Tensor repeats, *, SymInt? output_size=None)` primitive, which had no batching rule. This PR adds it. ### Why `output_size` is required when `repeats` is batched The output length of `repeat_interleave.Tensor` is `sum(repeats)`, i.e. **data-dependent**. When `repeats` is batched, that sum can differ between batch elements, so `vmap` cannot infer a single static output shape. The new batch rule: * **Unbatched `repeats`** (the pre-existing, already-working case where only the *values* are vmapped and `repeats` is a constant): re-dispatches to the regular implementation, so there is no behavior change and any outer vmap levels are handled correctly. * **Batched `repeats`**: requires `output_size` and builds the gather indices in a shape-static, fully batched way: ``` marks[b, cumsum(repeats[b])] += 1 out[b] = marks[b, :output_size].cumsum(-1) ``` which yields index `i` repeated `repeats[b, i]` times, e.g. `repeats = [4, 0, 8]`, `output_size = 12` → `[0,0,0,0, 2,2,2,2,2,2,2,2]`. A clear error is raised when `output_size` is omitted for a batched `repeats`. ### Repro from the issue (now works with `output_size`) ```python import torch from torch.func import vmap indices = torch.tensor([[4, 0, 8], [2, 2, 8], [4, 2, 6], [2, 4, 6], [0, 4, 8]]) def func(rep): values = torch.arange(1, rep.shape[0] + 1) return torch.repeat_interleave(values, rep, output_size=12) # all rows sum to 12 print(vmap(func)(indices)) ``` ### Test Adds `TestVmapOperators.test_repeat_interleave` covering: the indices overload, the `self_Tensor` overload, the unbatched-`repeats` fallback path, and the missing-`output_size` error. > Note: I was unable to build PyTorch locally, so this relies on CI for build/test validation. The index-construction algorithm was verified against a reference implementation for the issue's inputs and edge cases (empty / single-element / leading & trailing zeros). Pull Request resolved: pytorch#187702 Approved by: https://github.com/Skylion007 Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: PyTorch MergeBot <pytorchmergebot@users.noreply.github.com>
Always rely on `torch/headeronly/win32-headers.h` that hides a lot of global defines that windows.h unnecesserily pollutes address space with Pull Request resolved: pytorch#187673 Approved by: https://github.com/atalman, https://github.com/Skylion007
---- - Remove unused `using opmath_t = c10::metal::opmath_t<T>;` - Initialize threadgroup variable to 0 (it's a no-op, but suppresses the warning) Do not use `::metal::` prefix for unambigous as `using namespace metal;` is declared at the top of the shader Pull Request resolved: pytorch#187753 Approved by: https://github.com/Skylion007
MetaConverter used to memoize a freshly allocated tensor storage as the representative storage for the source storage whenever the converted tensor had matching stride and storage offset. That is not sufficient for `nn.Parameter`s constructed from tensor views: wrapping a view as a Parameter drops the autograd-view relationship, but the Parameter can still share a larger backing storage with other parameters. When `draft_export` enables real tensor propagation, the first offset-zero parameter could therefore memoize only its smaller tensor storage for the larger source storage. A later parameter sharing that storage with a nonzero offset then reused the too-small real storage and failed in `real_tensor.set_` with an out-of-bounds `setStorage` error. Symbolic storage sizes also made the representative-storage check unsafe when Python attempted to bool-cast symbolic comparisons. Fix this by only memoizing the freshly allocated storage when it is statically known to cover the full source storage. Unknown or symbolic comparisons fall back to the existing `meta_storage(s)` plus `set_` path, which allocates and memoizes the full source storage. This keeps the fast path for tensors whose fresh storage is provably representative without assuming that offset-zero non-view Parameters own their whole backing storage. An alternative would be to special-case Parameters created from views, but this invariant belongs in MetaConverter's storage memoization logic: any tensor that is not an autograd view can still have a source storage larger than its freshly allocated tensor storage. Fixes pytorch#168908 Generated by my agent Test Plan: - python test/test_fake_tensor.py -k test_parameter_views_keep_full_storage - python test/test_fake_tensor.py -k test_refake_unbacked_storage_static_shapes - python test/export/test_draft_export.py -k test_shared_storage_parameters_with_offset - git diff --cached --check - lintrunner -a Pull Request resolved: pytorch#185014 Approved by: https://github.com/ezyang
Only model inlined frame globals as module attributes when the globals dict belongs to the registered module. Fall back to dict-backed globals otherwise and graph break cleanly for non-module global writes. Fixes pytorch#181243 Generated by my agent Pull Request resolved: pytorch#184653 Approved by: https://github.com/williamwen42
…85014)" This reverts commit 9d79673. Reverted pytorch#185014 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#185014 (comment)))
…ytorch#187776) Factor the 16-byte-vectorized copy used by `nccl_all_to_all_nd` into a shared header (`symm_mem_copy.cuh`) and use it for the `all_gather_offset` push kernel, which previously copied each tile with a plain `ld_vec`/`st_vec` grid-stride loop. The helper batches `kUnroll=4` loads before issuing the stores, so several loads stay in flight per thread (memory-level parallelism). `all_gather_offset` caps the grid at `AG_MAX_CTAS=32`, so it runs at low occupancy where there are too few warps to hide load latency on their own -- the batched loads are what hide it. The multimem path is unchanged (it uses `ncclMultimemCopy`, not the copy loop). Effective output bandwidth on the push path (8x H100, bf16, GB/s): | shape | plain | unrolled | delta | | --- | --- | --- | --- | | 1 x 16MB | 288 | 337 | +17% | | 2 x 8MB | 288 | 335 | +16% | | 8 x 2MB | 288 | 332 | +15% | | 1 x 64MB + 8 x 8KB (skew) | 299 | 334 | +12% | | 64 x 64KB | 251 | 286 | +14% | | 256 x 4KB | 74 | 65 | -12% | The 256 x 4KB case regresses: each 4KB shard is exactly `blockDim.x` vectors, so the unrolled body never runs and only the extra register pressure remains; that shape is launch-latency-bound and not a realistic bucket. ## Test Plan Built and ran on an H100 node with 4 ranks: `all_gather_offset` (16 cases, both the multimem and `NCCL_NVLS_ENABLE=0` push paths) and `all_to_all_nd` (8 cases) all pass. Push bandwidth A/B measured on 8 ranks with the numbers above. ``` CUDA_VISIBLE_DEVICES=0,1,2,3 \ python test/distributed/test_nccl.py -v -k all_gather_offset CUDA_VISIBLE_DEVICES=0,1,2,3 \ python test/distributed/test_nccl.py -v -k all_to_all_nd ``` Pull Request resolved: pytorch#187776 Approved by: https://github.com/Skylion007 ghstack dependencies: pytorch#187642, pytorch#187750
Support returning lse in flex attention, setup for backward implementation of FA on MPS Pull Request resolved: pytorch#187768 Approved by: https://github.com/Skylion007
…torch#187758) ### Summary `Sequential.__getitem__` is annotated `slice | int -> Sequential | Module`, so every index widens to the full union regardless of the key type: ```python seq = nn.Sequential(nn.Linear(2, 2)) layer = seq[0] # typed Sequential | Module (should be Module) sub = seq[0:1] # typed Sequential | Module (should be Sequential) ``` `ModuleList.__getitem__` already distinguishes these with `@overload` (int -> `Module`, slice -> `ModuleList`). This mirrors that pattern on `Sequential` so int indexing narrows to `Module` and slice indexing narrows to `Sequential`. Runtime behavior is unchanged (the overloads are typing-only stubs; the implementation and its `@_copy_to_script_wrapper` are untouched), consistent with how `ModuleList` already combines `@overload` with the script wrapper. This reduces the over-general indexing typing reported in pytorch#181861. Note it does not fully resolve that issue: the nested-container case there (`model[0][0]`) requires knowing the element type, which is a broader generics discussion for `nn.Sequential`; this PR is scoped to the single-index/slice precision that matches `ModuleList`. ### Test plan Type-checking only; runtime behavior is unchanged. `seq[0]` now type-checks as `Module` and `seq[0:2]` as `Sequential`. Pull Request resolved: pytorch#187758 Approved by: https://github.com/Skylion007
The kernel assigned one CTA's worth of threads to a single row of a peer's slot, so a row narrower than the block left
most threads idle -- a `copy_row_bytes`-wide row used only `copy_row_bytes/16` of the threads, collapsing to a single
thread for 16-byte rows. It also re-resolved the peer pointer with `ncclGetLsaPointer` once per row and ran at 128
threads per CTA.
For col-scatter, this rank copies its own column block out of every source row of a peer's `[rows, G*local_cols]`
matrix; the unit of work is one `copy_row_bytes = local_cols*esize`-wide segment (a "row"):
```
|<-local_cols->| total_cols
+--------------+--------------+-----+--------------+
row 0 | rank 0 | rank 1 | ... | rank G-1 |
+--------------+--------------+-----+--------------+
row 1 | rank 0 | rank 1 | ... | rank G-1 |
+--------------+--------------+-----+--------------+
... | ... | | | |
+--------------+--------------+-----+--------------+
row R-1 | rank 0 | rank 1 | ... | rank G-1 |
+--------------+--------------+-----+--------------+
^^^^^^^^^^^^ this rank's block, copied from every row
```
A "row" is wide or narrow, measured in 16-byte vectors:
```
wide (local_cols=1024 bf16 = 128 vecs): [v0][v1]...[v127] <- all threads busy
narrow (local_cols=8 bf16 = 1 vec ): [v0] <- 1 of 256 threads, rest idle
```
This flattens each peer slot into 16-byte vectors over `(row, col)` and grid-strides every thread over that flat space,
so no thread idles regardless of row width: a wide row's vectors map to consecutive threads (coalesced) and narrow rows
pack many rows across the block. The peer's source base is resolved once per CTA and indexed by arithmetic (the LSA
pointer is offsettable within the window), loads are batched `kUnroll=4` deep before the stores for memory-level
parallelism at this kernel's low occupancy, threads go from 128 to 256, and `ctas_per_slot` is sized by total vectors
rather than rows.
Effective output bandwidth (8x H100, bf16, GB/s):
| shape | before | after | speedup |
| --- | --- | --- | --- |
| 8 x 4M (few wide) | 320 | 349 | 1.1x |
| 64 x 512K | 356 | 352 | ~parity |
| 8192 x 8192 (ulysses) | 204 | 365 | 1.8x |
| 65536 x 1024 | 28 | 362 | 12.9x |
| 512K x 64 (tall narrow) | 1.8 | 109 | 60x |
Narrow rows -- the common case for small per-rank column or row blocks -- go from a few GB/s to saturating, and the
realistic Ulysses shape gains ~1.8x. The `512K x 64` case stays lower because it gathers a 16-byte column slice from
each 128-byte source row (uncoalesced remote reads), but is still ~60x the old path.
## Test Plan
Built and ran on an H100 node with 4 ranks: `all_to_all_nd` (8 cases, both `(scatter, gather)` = `(1, 0)` and `(0, 1)`)
pass. Bandwidth measured on 8 ranks with the numbers above.
```
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python test/distributed/test_nccl.py -v -k all_to_all_nd
```
Pull Request resolved: pytorch#187778
Approved by: https://github.com/Skylion007
ghstack dependencies: pytorch#187642, pytorch#187750, pytorch#187776
…ch#187733) The Python native bmm outer-product override launches a Triton kernel for CUDA inputs shaped like `(B, M, 1) x (B, 1, N)`. Unlike the generated C++ ATen CUDA path, this Python dispatch path does not get an automatic CUDA device guard before launching the kernel. In vLLM multimodal tests, the model hooks can run the relevant rotary embedding buffers and inputs on `cuda:1` while the process current device is still `cuda:0`; Triton then launches on the wrong device/stream and rejects the `cuda:1` pointers with `ValueError: Pointer argument cannot be accessed from Triton (cpu tensor?)`. Guard the Triton call when the input device is not already current, while leaving the common already-current path direct to avoid unnecessary Python context manager overhead. Also require both bmm inputs to be on the same CUDA device before selecting the override, so mismatched CUDA inputs fall back to native bmm and report the normal same-device error instead of reaching the Triton kernel. I considered guarding unconditionally, but a small same-device benchmark showed that the Python context manager adds measurable overhead for small outer-product bmm calls. The conditional guard keeps the fix scoped to the non-current-device case that needs it. Fixes pytorch#187729 Generated by my agent Benchmark Results: - Command: `CUDA_VISIBLE_DEVICES=0 python - <<'PY' ...` benchmarking 7 trials of 1000 `torch.bmm` calls on `(32, 128, 1) x (32, 1, 512)` CUDA tensors with current device matching the inputs. - Before: median 24.210 us/call; raw `[24.34, 24.365, 24.195, 24.206, 24.21, 24.272, 24.142]`. - After: median 24.946 us/call; raw `[25.045, 24.909, 25.008, 25.042, 24.906, 24.946, 24.926]`. Test Plan: - `CUDA_VISIBLE_DEVICES=0,1 python - <<'PY' ...` minimal repro: before the fix failed with the Triton pointer `ValueError`; after the fix produced a `cuda:1` result and preserved current device `0`. - `CUDA_VISIBLE_DEVICES=0,1 python -m pytest test/test_bmm_outer_product.py -q -rs` passed: `11 passed, 10 subtests passed`. - `lintrunner torch/_native/ops/bmm_outer_product/triton_impl.py test/test_bmm_outer_product.py` passed. - `git diff --cached --check` passed. - `lintrunner -a` was run and failed on unrelated pre-existing clang-tidy findings in untouched C++ files. Pull Request resolved: pytorch#187733 Approved by: https://github.com/slayton58
Fix threshold for empty input Pull Request resolved: pytorch#187719 Approved by: https://github.com/malfet
The CPU vectorized atomic_add store path converted the destination index with ops.index_expr and then asserted that the result was already a vector CSE variable. Scatter/index_put backward graphs can vectorize the value lanes while the destination index is zero-stride in the vectorized dimension, leaving the index as a scalar CSE variable. That made otherwise valid compiled backwards fail in C++ codegen before any kernel was emitted. Broadcast scalar CSE indices in the vector atomic_add path before emitting atomic_add_vec. This matches the existing treatment for scalar store values and preserves the current vector-index path unchanged. Fixes pytorch#160391 Generated by my agent Test Plan: python - <<'PY' with the exact issue script from state/160391/issue python - <<'PY' reduced eager/compiled index_put accumulate check python - <<'PY' unittest runner for CPUReproTests.test_index_put_accumulate_atomic_add_vec_scalar_index with torchvision marked absent due local torchvision::nms import failure python - <<'PY' unittest runner for CPUReproTests.test_scatter_using_atomic_add_vec with torchvision marked absent due local torchvision::nms import failure lintrunner -a Pull Request resolved: pytorch#185325 Approved by: https://github.com/frgossen
Fixes: intel/torch-xpu-ops#2815 Add proper handling for in xpu::addmv_out for `mat.numel() = 0` tensors. Make output full zeros if Beta is zero and multiplicate self by beta otherwise. Pull Request resolved: pytorch#174193 Approved by: https://github.com/guangyey, https://github.com/EikanWang Co-authored-by: Sławomir Siwek <slawomir.siwek@intel.com> Co-authored-by: Yu, Guangye <guangye.yu@intel.com>
…sage per process (pytorch#185192) # Motivation This PR introduces a new API, `torch.xpu.list_gpu_processes`, analogous to the existing CUDA counterpart. # Additional Context `zes_process_state_t` defined in pyzes doesn't match the definition in L0 spec, see https://oneapi-src.github.io/level-zero-spec/level-zero/1.11/sysman/api.html#_CPPv4N19zes_process_state_t5stypeE From pyzes 0.1.2, it works. <img width="636" height="114" alt="image" src="https://github.com/user-attachments/assets/676c47a0-2f23-46b5-b51d-e1ae3d843fb9" /> Pull Request resolved: pytorch#185192 Approved by: https://github.com/EikanWang, https://github.com/gujinghui
The export core ATen decomposition table is applied after functionalization, so mutable post-autograd operators cannot appear in the graph where the table is consumed. However, _core_aten_decompositions_post_autograd() used get_decompositions() on operator packets, which expanded packets into all overloads, including in-place and out variants. Those mutable entries then propagated into torch.export.default_decompositions(), and user-provided mutable entries were silently passed through the export decomposition split. Filter mutable OpOverload entries out of the core post-autograd table and add an explicit validation step after CIA entries are separated. Functional CIA ops remain handled by the export CIA override path, while non-functional post-autograd entries now fail early with a clear error. Tests that explicitly build decomposition tables now use export-safe defaults or functional overloads instead of reintroducing mutable packet overloads. Fixes pytorch#135759 Generated by my agent Benchmark Results Measured table construction with timeit.repeat(number=20, repeat=5), comparing main against this patch using the same local build artifacts: main 30326da: - core table entries: 445 - default table entries: 1004 - _core_aten_decompositions_post_autograd(): min 0.000421s, mean 0.000431s - default_decompositions().materialize(): min 0.015557s, mean 0.016068s this patch: - core table entries: 215 - default table entries: 774 - _core_aten_decompositions_post_autograd(): min 0.000523s, mean 0.000536s - default_decompositions().materialize(): min 0.015485s, mean 0.016067s Test Plan - python test/export/test_export.py TestOneOffModelExportResult.test_input_output_no_stacktrace TestExport.test_decomp_item_in_prim_after_decomposition TestExport.test_error_when_passing_mutating_primitive_op -v - python test/export/test_strict_export_v2.py StrictExportV2TestExport.test_decomp_item_in_prim_after_decomposition_strict_export_v2 -v - python test/export/test_export_training_ir_to_run_decomp.py TrainingIRToRunDecompExportTestExport.test_decomp_item_in_prim_after_decomposition_training_ir_to_decomp_strict -v - python test/export/test_serdes.py SerDesExportTestExport.test_decomp_item_in_prim_after_decomposition_serdes_strict -v - python test/test_decomp.py HasDecompTest.test_core_aten_decompositions_are_functional HasDecompTest.test_aten_core_operators -v - lintrunner -a ghstack-source-id: 4a5b73a Pull-Request: pytorch#186287
Lazy torch.compile can run substantial Dynamo and Inductor compiler code inside a user torch.profiler.profile(with_stack=True) region. with_stack enables Kineto's Python tracer, so the profiler records compiler Python frames and later has to materialize a very large Python event tree. In the reported vLLM workload this caused the process to spend unbounded time and memory in profiler post-processing. Add a private autograd binding to pause and resume the active Kineto Python tracer, and use it around Dynamo frame conversion and Inductor backward compilation. This suppresses Python tracing for compiler internals while keeping profiler record-function events such as the compiled region visible. Pause depth is tracked on the active Kineto profiler state, so nested compiler entries and profile_all_threads=True use the same active state rather than a Python process-global counter. I used a profiler-state pause instead of disabling the whole profiler or removing Dynamo record functions because users still need runtime and compiled-region events in the trace; only compiler Python stack tracing is the pathological part. Fixes pytorch#136404 Generated by my agent Test Plan: - ninja -C build torch_python - timeout 600s python test/dynamo/test_profiler.py DynamoProfilerTests.test_disable_profiler_python_tracing_profile_all_threads DynamoProfilerTests.test_profiler_with_stack_skips_compile_python_tracing DynamoProfilerTests.test_profiler_with_stack_skips_compile_fx_backward_python_tracing - timeout 700s python test/dynamo/test_profiler.py - Original-style Inductor repro with torch.profiler.profile(with_stack=True, verbose=True): prof.events() parsed successfully, symbolic_convert.py and trace_rules.py frames were absent, and the compiled-region event was present. - git diff --cached --check - lintrunner -a was run and failed on existing/file-wide clang-tidy findings in torch/csrc/autograd/profiler_kineto.cpp and torch/csrc/autograd/profiler_kineto.h. Benchmark Results: - No-profiler cold compile microbenchmark, 25 iterations of torch._dynamo.reset(); torch.compile(fn, backend="eager")(x). - main: median 18.2047 ms, mean 36.8621 ms, min 17.2573 ms, max 478.9803 ms. - this change: median 18.4466 ms, mean 34.8736 ms, min 17.6654 ms, max 427.2582 ms. - No meaningful regression; both runs are dominated by similar cold-compile outliers. ghstack-source-id: 5e6bda9 Pull-Request: pytorch#186280
Fake tensor and proxy tensor tracing only kept tensor constants with one element. Shape tensors built from Python shape lists, such as torch.as_tensor([[64, 64], [32, 32]]), therefore lost their constant payload even though they are small. Later scalar uses of those values, including torch.linspace(..., steps=H), became data-dependent unbacked symbols and strict export failed with GuardOnDataDependentSymNode. Raise the small tensor constant propagation limit from 1 to 8 in both fake tensor and proxy tensor tracing. This keeps common small shape tensors constant while preserving a conservative cap. I chose this root-cause change over adding linspace or slicing special cases because maintainers pointed at more aggressive constant propagation and the same lost constant payload affects multiple downstream scalar consumers. The higher limit also exposed expected follow-ups: tests that intentionally need dynamic fake values now create non-constant inputs, cache tests that specifically exercise dispatch-cache behavior clear the constant payload, and Dynamo recompilation-count expectations are updated for the improved constant propagation. To avoid a real distributed regression, fake/proxy constant propagation now excludes runtime DeviceMesh coordinate computation so compile-on-one-rank graphs do not bake rank-local literals. Fixes pytorch#136642 Generated by my agent Test Plan: - python test/export/test_export.py TestExport.test_strict_export_constant_props_small_shape_tensor - python test/test_fake_tensor.py FakeTensorTest.test_sparse_spdiags_dynamic_nnz_bound - python test/test_fake_tensor.py FakeTensorConstHandling - python test/test_fake_tensor.py FakeTensorDispatchCache.test_cache_aten_index - python test/test_fake_tensor.py FakeTensorDispatchCache - python test/dynamo/test_repros.py ReproTests.test_convert_boxes_to_pooler_format - python test/dynamo/test_dynamic_shapes.py DynamicShapesReproTests.test_convert_boxes_to_pooler_format_dynamic_shapes - python test/distributed/tensor/test_compile_on_one_rank.py TestCompileOnOneRank.test_compiled_rowwise_embedding_graph_consistency - python test/distributed/test_local_tensor.py TestLocalTensorWorld4.test_runtime_compute_coordinate_on_dim_per_rank_call_style_positional TestLocalTensorWorld4.test_runtime_compute_coordinate_on_dim_per_rank_call_style_kwargs - git diff --check - lintrunner test/test_fake_tensor.py torch/_subclasses/fake_tensor.py torch/fx/experimental/proxy_tensor.py - lintrunner -a Rebase Update Test Plan: - git diff --check - lintrunner torch/_subclasses/fake_tensor.py torch/fx/experimental/proxy_tensor.py torch/distributed/_ops/device_mesh.py test/distributed/test_local_tensor.py - lintrunner -a - attempted targeted Python tests and import sanity check, but Python import hung in this workspace before executing test code Benchmark Results: Synthetic FakeTensorMode benchmark, 1200 iterations creating a 4-element int64 tensor then add and multiply, 9 alternating reps with the limit toggled in-process. Old limit 1 median: 0.561588 s, 467.990 us/iter. New limit 8 median: 1.051484 s, 876.237 us/iter, 1.872x. This intentionally stresses the newly constant-propagated case; the change remains capped at tensors with at most 8 elements. ghstack-source-id: b365b85 Pull-Request: pytorch#186263
Fixes pytorch#186220. TritonTemplate, CuteDSLTemplate, and ExternKernelChoice register themselves into process-global registries at module-init time, guarded by hard asserts that forbid any name that is already present. If the registering module is initialized twice in a single process (e.g. a double-import path, reported on the 2.11.0+cu130 Blackwell build), the second initialization re-runs the same registration and the assert turns a benign duplicate into a fatal crash during import -- before any compilation runs. Relax each guard so that re-registering under an existing name is tolerated as long as it is equivalent to what is already registered (matching template source hash for templates, matching wrapped callable for extern kernels), while a genuine name collision between different templates/kernels still raises. The assignment that follows is unconditional so the registry and the module-level symbol that constructed the instance continue to refer to the same object. This addresses the proximate crash; it does not explain why the cu130 build re-imports the registering module in a single process, which is left to the issue's root-cause investigation. Test Plan: New pure-Python regression tests (no GPU / compilation required), placed in their own file because test_select_algorithm.py is gated behind a big-GPU check that would skip them on the CI most relevant to this import-time bug: ``` python -m pytest test/inductor/test_template_registration.py -q ``` This PR was authored with the assistance of an AI coding assistant. ghstack-source-id: 9d41999 Pull-Request: pytorch#186262
Upstream: pytorch#186272 Generated by AI assistant.
…ied issues Summary of fixes applied (issues pytorch#1787-pytorch#1870 from intel-sandbox/torch-xpu-ops-exp): ## C++ Fixes (aten/c10/torch/csrc) - aten: FunctionalTensorWrapper, NestedTensorImpl, SparseCsrTensorImpl, SparseTensorImpl - aten/native: LinearAlgebra (matmul fold), Loss (BCEWithLogitsLoss validation), RangeFactories (div-by-zero), Gelu NaN, CUDA Gelu - c10: TensorImpl (detach tracing), PyInterpreter (shallow_copy_and_detach) - torch/csrc: PyInterpreter, autograd/engine (thread deferral), autograd/FunctionsManual (batch_norm 3rd order), profiler_kineto (Python tracing toggle), jit/pybind_utils (schema errors with SymInt types), profiler/kineto_shim (PRIVATEUSE1 mapping), utils/tensor_new (FakeTensor meta constructor) ## Python Fixes - torch/_dynamo/symbolic_convert.py, torch/_subclasses/fake_tensor.py - torch/fx/experimental/symbolic_shapes.py - torch/nn/modules/rope.py (RoPE dtype coverage) ## Test Coverage - test/test_autograd.py, test/test_proxy_tensor.py Verified with Intel oneAPI 2025.3 + ninja incremental build on 8x Intel Data Center GPU Max 1100 (PVC). Related: intel-sandbox/torch-xpu-ops-exp#1787 through pytorch#1870
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
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.
Fixes intel-sandbox/torch-xpu-ops-exp#1865. Verified XPU.