[pull] main from pytorch:main#1634
Merged
Merged
Conversation
…#189939) `build-triton-wheel` and `build-vllm-wheel` requested the release-isolated `rel-l-*` runners on **every** ref, so their PR / ciflow runs hang against the now-locked release runner group (its access is restricted to protected refs). Gate the release runners to release refs (`main` / `nightly` / `release/*` / `v*` tags) and fall back to the regular multi-tenant OSDC pool (`l-x86iavx512-48-384` / `l-arm64g3-61-463`) on other refs — matching the `generated-linux-*-binary-manywheel-nightly` workflows. Authored with the assistance of Claude Code. Pull Request resolved: #189939 Approved by: https://github.com/jeanschmidt
…p fetch concurrency (#189605) Checking out PyTorch is a recurring source of CI flakiness. Two robustness issues in the `checkout-pytorch` action. **1. The retry safety net is defeated in container runners.** After a failed first checkout, the action reset the workspace and retried. It did `rm -rf "$GITHUB_WORKSPACE"`, but the workspace *directory* can't be unlinked when its parent is root-owned and the job runs as an unprivileged uid with no sudo (the contents are removable, the dir entry is not). That `rm` failed and, under `set -e`, aborted the composite before the retry checkout could run, turning a transient fetch error into a hard failure. Example: [run 29116512761](https://github.com/pytorch/pytorch/actions/runs/29116512761/job/86448178127) hit `HTTP 408` during a promisor fetch, then died on `rm: cannot remove '/__w/pytorch/pytorch': Permission denied`. Fix: clear the workspace *contents* instead (via `find`, so `.git` is included) -- which is what the retry checkout needs to clone into. **2. Unbounded fetch concurrency.** `fetch.parallel`/`submodule.fetchJobs` were `0` (one job per core); on large runners with 37 recursive submodules that connection burst gets throttled into HTTP 408/429. Cap both at 8. Authored with the assistance of an AI coding assistant. Pull Request resolved: #189605 Approved by: https://github.com/Skylion007
## Why
TVM removed the relay frontend in 0.20, so the tvm backend's relay path only works with old TVM builds and is frozen. Users on it currently get no signal that it is going away.
## How
- Emit a `FutureWarning` on entry to `_tvm_relay_compile`, pointing at the relax frontend
- Document relax tuning via `options={"pipeline": ...}` (added in #189638, which should land first)
- Add `test_tvm_relay_future_warning`, which runs with or without a tvm install
Pull Request resolved: #189639
Approved by: https://github.com/ezyang, https://github.com/jansel
…en aten.cat fails (#189509) ## Summary After PR #181854 auto-enabled `batch_linear_lhs` for XPU inference, models using torchao W8A8 quantization crash during `torch.compile` with: ``` NotImplementedError: Int8Tensor dispatch: func=<OpOverload(op=aten.cat, overload=default)> ``` The root cause is `BatchLinearLHSFusion.fuse()` calling `torch.cat(batch_weights_meta)` on weight tensors that are `Int8Tensor` subclasses lacking `aten.cat` dispatch. ## Fix Replaces the blanket subclass skip with a **capability probe** in `BatchLinearLHSFusion.match()`. Instead of returning `None` for all tensor subclasses, it tries `torch.cat` and only skips fusion if the call actually fails: ```python try: _ = torch.cat([weight_val[:1], weight_val[:1]], dim=0) except (RuntimeError, TypeError): return None ``` - **Without pytorch/ao#4560** (released torchao 0.17.0): probe catches `NotImplementedError` on cat → fusion skipped → no crash - **With pytorch/ao#4560 + aten.permute.default**: probe passes → fusion fires → quantized perf win (2.4%/6.6%) preserved - **Plain torch.Tensor / FakeTensor**: type check bypasses the probe entirely → zero overhead Also applies etafs review suggestion: simplified `bias.meta.get("val", bias.meta.get("example_value"))`. ## Test Added `test_batch_linear_lhs_skips_tensor_subclass_weights` to `test/inductor/test_group_batch_fusion.py`. ## Validation Reproduced on XPU nightly (2.14.0.dev20260713+xpu) with torchao 0.17.0 and pytorch/ao#4560 + permute fix. See [comment](#189509 (comment)) for full results. ## References - Regression-introducing PR: #181854 (auto-enable batch_linear_lhs for XPU) - torchao issue: pytorch/ao#4560 (add cat/transpose/mm/addmm to Int8Tensor) - Reproduction branch: https://github.com/xuhancn/dev_llama_up_gate_fusion_op/tree/reproduce_torchao_int8_fusion Pull Request resolved: #189509 Approved by: https://github.com/etaf, https://github.com/jansel Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
CUDASymmetricMemoryAllocator::alloc allocates via the low-level CUDA VMM API (cuMemCreate) but never set prop.allocFlags.gpuDirectRDMACapable, so symmetric memory buffers were not marked as GPUDirect RDMA capable even on hardware/driver combinations that support it. This mirrors what CUDACachingAllocator.cpp already does for expandable segments. Test Plan: ``` lintrunner -a torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.cu pip install -e . -v --no-build-isolation ``` Authored with the assistance of Claude Code. Pull Request resolved: #189941 Approved by: https://github.com/dsjohns2, https://github.com/kapilsh, https://github.com/d4l3k, https://github.com/fduwjj
The in-tree nccl2 c10d backend (torch/csrc/distributed/c10d/nccl2/) is added to torch_cuda whenever CUDA and distributed are enabled, independent of whether NCCL is available. Some of its files were already wrapped in #ifdef USE_C10D_NCCL but many were not, so a non-NCCL-compliant backend such as libtorch_mtia_simt -- which defines no USE_C10D_NCCL -- fails to compile. Two distinct failures appear in sequence: torch/csrc/distributed/c10d/nccl2/ProcessGroupNCCLUtils.cpp:5:10: fatal error: 'nccl.h' file not found torch/csrc/distributed/c10d/nccl2/CudaApi.hpp:75:7: error: unknown type name 'cudaUserObject_t' The first comes from the files that #include <nccl.h> or use NCCL types. The second comes from CudaApi, an NCCL-free CUDA wrapper that uses CUDA-graph APIs (cudaUserObject_t) not present on the mtia_simt CUDA surface. Because the whole nccl2 backend is only usable with NCCL -- its C++ registration in init.cpp and its Python entry point are both gated on NCCL availability -- the entire directory is gated rather than cherry-picking individual files. This wraps every nccl2 source and header in #ifdef USE_C10D_NCCL / #endif, so the backend compiles to empty translation units when NCCL is not built. Files that already carried the guard are unchanged. Test Plan: Non-NCCL backend (the original failure): building libtorch_mtia_simt now succeeds and the previously broken tests pass: ``` buck2 test fbcode//mtia/cutedsl/tests/python:test_dlpack-mtia_simt fbcode//mtia/cutedsl/tests/python:test_cute_tensor-mtia_simt # Tests finished: Pass 190. Fail 0. Timeout 0. Fatal 0. Skip 0. Omit 0. Infra Failure 0. Build failure 0 ``` With USE_C10D_NCCL defined (normal CUDA build), the nccl2 objects compile unchanged. With USE_C10D_NCCL undefined, every nccl2 file preprocesses to an empty translation unit. lintrunner on the changed files reports no issues. Differential Revision: [D111984293](https://our.internmc.facebook.com/intern/diff/D111984293) Pull Request resolved: #189938 Approved by: https://github.com/kapilsh, https://github.com/dolpm
unload_xpu_triton_pyds() hardcodes two intel-xpu-backend-for-triton internals that no longer exist in current backends: - XPULauncher.mod was removed when the per-kernel launcher was replaced with a shared driver-level launcher (intel/intel-xpu-backend-for-triton#6650) - XPUUtils's singleton attribute was renamed instance -> _instance (intel/intel-xpu-backend-for-triton#6767) Both raise AttributeError during fresh_cache() teardown on Windows+XPU, failing every affected test. Guard both accesses defensively instead of assuming either shape. Note: cleans up all kernel/test-scoped cache artifacts, but a couple of small shared native .pyd files may still linger in the temp dir, since newer backends keep them alive via references outside this function's reach (pre-existing ignore_errors=is_windows() here already tolerates partial cleanup). Fixes: intel/torch-xpu-ops#4038 Fixes: intel/torch-xpu-ops#4037 Pull Request resolved: #189147 Approved by: https://github.com/PatrykWilczewski, https://github.com/Silv3S, https://github.com/guangyey, https://github.com/etaf, https://github.com/jansel
Migrate the linux manywheel nightly build **and** test jobs off EC2 onto OSDC: **Build jobs** → release-isolated ARC runners from pytorch/ci-infra#888 (same r7a/r7g.12xlarge instance types as the old EC2 labels), so release wheel builds don't share a node with CI: - x86: `linux.12xlarge.memory.ephemeral` -> `mt-rel-l-x86iavx512-44-340` - aarch64: `linux.arm64.r7g.12xlarge.memory` -> `mt-rel-l-arm64g3-44-340` **Test jobs** (`manywheel-test`) → regular multi-tenant OSDC CI runners (mapped via `.github/arc.yaml`): - CPU: `linux.4xlarge` -> `mt-l-x86iavx512-16-128` - CUDA: `linux.g4dn.4xlarge.nvidia.gpu` -> `mt-l-x86iavx512-29-115-t4` (T4) - aarch64: `linux.arm64.2xlarge` -> `mt-l-arm64g3-16-62` Change is in the template (`linux_binary_build_workflow.yml.j2`) + regenerated; s390x is unchanged. Build/test jobs already use the GHA `container:` directive (the CUDA test container already gets `--gpus all` via `_binary-test-linux.yml`). New labels registered in `actionlint.yaml`. Authored with the assistance of Claude Code. Pull Request resolved: #189818 Approved by: https://github.com/jeanschmidt, https://github.com/atalman
make_fx's post-dispatch tracing can record detach calls emitted by autograd saved-variable handling while tracing a backward with create_graph=True. When the resulting graph is replayed and users take a second derivative through it, those detach nodes cut gradient flow and can produce silently wrong gradients, as in #175477. Add the existing aten.detach.default -> nop_decomposition mapping to _MakefxTracer's default decomposition table for non-pre-dispatch make_fx. The mapping is only added when the caller did not pass an explicit decomposition table, so callers that request exact detach preservation with an explicit table keep the old behavior. pre_dispatch tracing also keeps recording detach so export/pre-autograd use cases preserve detach semantics. A broader C++ shallow_copy_and_detach aliasing fix was considered in prior attempts, but it changes lower-level autograd dispatch behavior and was too broad for this post-dispatch make_fx issue. This patch keeps the behavior localized to make_fx's default decomposition policy and updates expected graph strings that now see alias instead of detach. Fixes #175477 Generated by my agent Test Plan: - python - <<'PY' ...issue reproducer... PY - python - <<'PY' ...guard torchvision import, run test/test_proxy_tensor.py -k TestSymbolicTracing.test_make_fx_second_order_grad -v... PY - PYTORCH_TEST_WITH_DYNAMO=1 python - <<'PY' ...guard torchvision import, run test/test_proxy_tensor.py -k TestSymbolicTracing.test_make_fx_second_order_grad -v... PY - python test/dynamo/test_activation_checkpointing.py ActivationCheckpointingNonStrictTracerTests.test_checkpoint_traces_through_eager_ac_under_non_strict ActivationCheckpointingNonStrictTracerTests.test_sac_traces_through_eager_ac_under_non_strict ActivationCheckpointingNonStrictTracerTests.test_checkpoint_with_rng_op_under_non_strict ActivationCheckpointingNestedCompileTests.test_checkpoint_recompute_preserves_nested_fx_trace_policy -v - python test/export/test_experimental.py -k TestExperiment.test_joint_basic -v - python - <<'PY' ...guard torchvision import, run test/functorch/test_aotdispatch.py TestAOTExport.test_aot_export_module_joint -v... PY - git diff --check && git diff --cached --check - lintrunner test/test_proxy_tensor.py torch/fx/experimental/proxy_tensor.py test/dynamo/test_activation_checkpointing.py test/export/test_experimental.py test/functorch/test_aotdispatch.py - lintrunner -a (fails on unrelated pre-existing aten/src/ATen/cuda/CUDAGraph.cpp clang-tidy findings) Pull Request resolved: #186845 Approved by: https://github.com/aorenste
AOTAutograd currently detaches every saved view before saving it for backward. The original cycle looks like this from #94990: ```py class TestFn(torch.autograd.Function): staticmethod def forward(ctx, a): b = a + 1 c = b.view(-1) ctx.save_for_backward(c) return b ``` The cycle is: returned output b -> grad_fn / TestFnBackward node -> ctx saved tensors -> saved view c -> c._base / ViewAutogradMeta base -> b We call Tensor.detach() to break the cycle. Tensor.detach() is unfortunately expensive. We avoid calling Tensor.detach() if the view being saved for backward is a graph input. Saving graph input views directly does not create this cycle because autograd has some special logic if the output is the input directly that breaks the reference cycle. Note that there are some better solutions, like actually fixing the root cause of #94990, but we avoid that problem for now. This PR saves ~3us runtime overhead from a flex_attention example (when 3 inputs are views); each detach was ~1us. Test Plan: ``` python test/functorch/test_aotdispatch.py -k test_save_input_view_for_bw_does_not_leak_memory python test/functorch/test_aotdispatch.py -k test_mem_leak_from_save_for_bw python -m py_compile test/functorch/test_aotdispatch.py lintrunner -a ``` Authored with assistance from an AI assistant. Pull Request resolved: #189759 Approved by: https://github.com/ezyang, https://github.com/frgossen
THPFunction_apply only needs borrowed tensor handles for most of its bookkeeping. Store Python tensor inputs as an ArrayRef-compatible SmallVector<const Variable*, 24> and thread those borrowed Variables through output wrapping so the common path avoids constructing a variable_list and bumping TensorImpl refcounts. Tracing and the user JVP callback still materialize a variable_list at the boundary that requires owned tensor handles. On the local one-apply 13-in-2-out no-save instruction-count benchmark, this moved from 54,090 to 51,030 instructions. Pull Request resolved: #189788 Approved by: https://github.com/soulitzer ghstack dependencies: #189577, #189582
This avoids one at::Tensor refcount bump per tensor output. Authored with assistance from an AI assistant. On the local one-apply 13-in-2-out no-save instruction-count benchmark, this moved from 51,030 to 50,802 instructions. Pull Request resolved: #189800 Approved by: https://github.com/Skylion007, https://github.com/soulitzer ghstack dependencies: #189577, #189582, #189788
…s, test_typechecks, test_global to Dynamo (#189458) Copies these five pristine CPython 3.13.5 test files into test/cpython/v3_13/ and adapts them to the CPythonTestCase harness (Dynamo patch header, TestCase base swapped to CPythonTestCase, unittest.main() swapped to run_tests(), plus a .diff recording the delta from upstream). They exercise __init_subclass__/__set_name__ (test_subclassinit), keyword-only argument binding (test_keywordonlyarg), dict comprehensions (test_dictcomps), __instancecheck__/__subclasshook__ (test_typechecks), and the global statement (test_global). test_global overrides setUp. CPythonTestCase setUp installs self.handler (and saves grad/config state) that its tearDown tears back down, so an override that does not chain to super() leaves self.handler unset and every test in the class errors in tearDown; the override now chains to super(). Tests that currently graph-break under nopython/strict are baselined as expected failures under test/dynamo_expected_failures/; delete the corresponding file when the underlying gap is fixed. Test Plan: ``` for t in test_subclassinit test_keywordonlyarg test_dictcomps test_typechecks test_global; do PYTORCH_TEST_WITH_DYNAMO=1 python test/cpython/v3_13/$t.py done ``` All report `OK (skipped=N)` (12, 8, 2, 5, 4 respectively). Authored with the assistance of Claude (Anthropic AI). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Pull Request resolved: #189458 Approved by: https://github.com/hameerabbasi
…out (#189902) ## Summary `_PerDocumentHeadTailLoadBalancer` (used for Context Parallel with packed, variable-length documents) produced **imbalanced** attention work across ranks on mixed-length documents. CP shards the rearranged sequence by handing each rank a **contiguous** chunk (`_context_parallel_shard` -> `distribute_tensor(..., [Shard(seq_dim)])`, i.e. `torch.chunk` semantics). The balancer, however, laid out its rearrange indices **document-major** (all of doc0's per-rank head/tail chunks, then all of doc1's, ...). With mixed lengths, the contiguous cut then falls in the *middle* of a document's block, so some ranks inherit a document's entire expensive causal tail. Outputs stay numerically correct (it is still a valid permutation), which is why the existing correctness tests never caught it — the bug is purely a load imbalance. ### Concrete example (`q_len == kv_len == 16`, `cp_world_size == 2`) Packed documents of length **4** and **12** (doc0 = KV 0-3, doc1 = KV 4-15). Each row is a query's document-causal mask; `-> N` is that query's workload (number of attended keys); the per-rank `sum` is the work landing on that rank after the contiguous Q-shard. **Before (document-major): the cut splits doc1, rank1 gets its whole tail** ``` [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 1 (q0) [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 4 (q3) [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 2 (q1) [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 3 (q2) [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 1 (q4) [0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 2 (q5) [0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 3 (q6) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0] -> 10 (q13) rank0 sum=26 ------------------------------------------------------------ [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0] -> 11 (q14) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -> 12 (q15) [0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0] -> 4 (q7) [0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0] -> 5 (q8) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0] -> 6 (q9) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0] -> 7 (q10) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0] -> 8 (q11) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0] -> 9 (q12) rank1 sum=62 (imbalance 2.38x) ``` **After (rank-major): each rank gets head+tail of every document** ``` [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 1 (q0) [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 4 (q3) [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 1 (q4) [0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 2 (q5) [0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 3 (q6) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0] -> 10 (q13) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0] -> 11 (q14) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -> 12 (q15) rank0 sum=44 ------------------------------------------------------------ [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 2 (q1) [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 3 (q2) [0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0] -> 4 (q7) [0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0] -> 5 (q8) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0] -> 6 (q9) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0] -> 7 (q10) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0] -> 8 (q11) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0] -> 9 (q12) rank1 sum=44 (imbalance 1.00x) ``` The only change is the loop nesting: document-major (document outer, rank inner) vs **rank-major** (rank outer, document inner). Both feed the identical CP shard (permute + contiguous halve); only the rank-major layout makes the halves equal. ### The code did not match its own docstring (`[4, 8, 4]`) `_PerDocumentHeadTailLoadBalancer._generate_indices`'s docstring illustrates the "after load-balancing on 2 devices" result for documents `[4, 8, 4]` (doc0 = KV 0-3, doc1 = KV 4-11, doc2 = KV 12-15). That documented diagram is actually the **rank-major** layout — rank0 owns `{0, 3, 4, 5, 10, 11, 12, 15}`: ``` [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 1 (q0) [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 4 (q3) [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 1 (q4) [0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 2 (q5) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0] -> 7 (q10) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0] -> 8 (q11) [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0] -> 1 (q12) [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1] -> 4 (q15) rank0 sum=28 ------------------------------------------------------------ ... rank1 = {1, 2, 6, 7, 8, 9, 13, 14} rank1 sum=28 ``` But the actual (pre-fix) **document-major** code produced a *different* assignment — rank0 owns `{0, 1, 2, 3, 4, 5, 10, 11}`: ``` [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 1 (q0) [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 2 (q1) [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 3 (q2) [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 4 (q3) [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 1 (q4) [0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> 2 (q5) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0] -> 7 (q10) [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0] -> 8 (q11) rank0 sum=28 ------------------------------------------------------------ ... rank1 = {6, 7, 8, 9, 12, 13, 14, 15} rank1 sum=28 ``` Because the lengths are symmetric (doc0 + doc2 = doc1), the contiguous cut still happens to land on a rank boundary, so **both layouts sum to 28/28**. The totals matched even though the actual per-rank assignment did not match the documented diagram — so the symmetric example "looked fine" and hid the mixed-length bug shown above. After this fix, the code produces exactly the assignment its docstring already illustrates. ## The fix Lay out the rearrange indices **rank-major** so a contiguous per-rank shard gives each rank a head+tail slice of every document. Because head chunk `r` is paired with tail chunk `2W-1-r`, each rank's per-document load is `c*(2W-1) + c*(c+1)` — **independent of the rank** — so the split is *exactly* balanced for any document lengths and world size. The single-sequence `_HeadTailLoadBalancer` already used this rank-major layout; this makes the per-document variant consistent with it (and with its own docstring diagram). The implementation is also vectorized to mirror `_HeadTailLoadBalancer`. The contiguous-shard requirement the layout must satisfy is now documented on the base `_LoadBalancer._generate_indices` interface. ## Test plan New tests in `test/distributed/tensor/test_attention.py`: - `PerDocumentHeadTailLoadBalancerTest` (CPU, non-distributed): exact per-rank balance on the `[4, 12]` regression and on randomized mixed-length documents across world sizes 2/3/4/8, plus a restore-inverts-rearrange check. - `TestSharding.test_context_parallel_shard_per_document_balance` (2-GPU): verifies balance end-to-end through the real `_context_parallel_shard` path. ``` python -m pytest test/distributed/tensor/test_attention.py \ -k "PerDocumentHeadTailLoadBalancerTest or test_context_parallel_shard_per_document_balance" # => 4 passed, 1 skipped python -m pytest test/distributed/tensor/test_attention.py \ -k test_cp_flex_attention_document_mask # existing correctness test, no regression # => 1 passed, 1 skipped ``` The before/after numbers above were reproduced through the real `_context_parallel_shard` path on a gloo CPU mesh (26/62 -> 44/44 for `[4, 12]`). `lintrunner` on both changed files reports no issues. Pull Request resolved: #189902 Approved by: https://github.com/fegin
…pile` support (#186595)" (#189305) Reapplies #186595 after Meta internal revert. Let me know if there are any failures that require fixing from my end. /cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo @jansel @ngimel Pull Request resolved: #189305 Approved by: https://github.com/jansel, https://github.com/ngimel
# Summary
Okay this one is a doozy.
I am going to have a long blurb, maybe copied to every commit, on what we are doing and then add to each commit what step it is accomplishing. At a high level we are trying to add reductions in the epilogue. These reductions need to be `local` per some loose definition of `local`. The stronger definition is that this is dependent on the kernel implementation: we are not using distributed shared memory with clusters yet, the reduction cannot cross CTA boundaries, and we require that the epilogue is fully fused.
Let's go through some examples:
```python
def epi(acc):
# N-axis local reduce: output shape [M, N // group]
out, aux = acc.relu(), acc.float().view(M, -1, group).sum(-1)
return out, aux
def epi(acc):
# M-axis local reduce: output shape [M // group, N]
out, aux = acc.relu(), acc.float().view(-1, group, N).sum(1)
return out, aux
def epi(acc):
# Feed and store the local-reduce value.
grouped = acc.float().view(-1, group, N)
scale = grouped.sum(1, keepdim=True)
out = (grouped * (scale + 1.0)).view(M, N)
return out, scale.view(M // group, N)
```
The last example is both feed-main and a compressed aux store because it returns `scale`. A feed-only version just returns `out`.
Since we are still only dealing with 2-D M/N tiles, we basically have two different forms of local reductions: along the M dimension or along the N dimension. To be local, the group must fit inside the selected CTA tile and the tile must be compatible with that group. The group also has to divide the logical M or N dimension. Groups up to 32 must divide the 32-wide TensorSSA fragment; larger groups must be multiples of 32.
I like this code blurb on how we are doing this:
> FlexGEMM recognizes a narrow local-reduction contract inside the GEMM output tile: an epilogue reshapes the accumulator to expose contiguous groups along M or N, then reduces only that grouped dimension. N-axis groups up to one 32-wide TensorSSA fragment lower as ordinary in-fragment TensorSSA reductions; larger N groups produce TensorSSA partials that QuACK combines physically.
>
> M-axis groups currently always use QuACK's physical lane/warp combine path, even when the group is small enough to fit in one fragment. Inductor owns the FX pattern matching and output contracts; these helpers describe the supported TensorSSA shapes and generated combine/finalize expressions QuACK needs.
This is a little claudy, but it is accurate. One factual tweak from how I was originally thinking about the Blackwell layout: the current local-reduce configs are non-`swap_ab`, but the implementation does not hard-code that one thread owns all of the N values and that the neighboring thread is exactly one row down along M. QuACK derives the M/N lane and warp ownership from the active epilogue tiled-copy TV layout. Operationally we still get the distinction we care about here: fragment-sized N reductions can stay in TensorSSA, while M reductions require lane/warp reduction and may require shared memory when they span M warps.
Okay, so at a high level that's what we are doing, and at the individual level we need a lot of codegen code to get there. Some things we will visit:
- adding reduction support to the CuTeDSL SSA op overrides for in-fragment reductions;
- a multi-step reduction extension to QuACK/codegen to accumulate reductions across epilogue iterations;
- a lot of specific view/reshape expression handling for user local-reduction expressions;
- constraint and edge-case handling, much of it found through fuzzing.
Claudex has a particular style here of named error messages throughout. I think that's fine. I probably would not code it this way from scratch, but I am not really drilling down on that in this stack.
## How are we representing local reductions?
```python
# axis 0 or 1 and the group size, i.e. how local the reduction is
FlexGemmLocalReduceGeometry(group=group, axis=axis)
```
At compile time we bind the FX reduction value to its actual consumers:
```python
FlexGemmOutputLocalReducePlan(
geometry=geometry,
value_node=value_node,
store_node=maybe_store_node,
aux_index=maybe_aux_index,
feeds_main=feeds_main,
)
```
At runtime, a compressed local-reduce aux has `out=local_reduce_out`. That output comes from th FX graph and is allocated before the generated wrapper calls QuACK;
```python
FlexGemmRuntimeLocalReducePlan(
FlexGemmLocalReduceGeometry(group=group, axis=axis),
out=local_reduce_out,
callbacks=maybe_generated_callbacks,
)
```
The callbacks are needed for every M-axis reduction and for N-axis groups larger than 32 (non thread local). A fragment-sized N-axis store still has a runtime plan because it has an output buffer, but it does not need physical callbacks.
If a physical reduction feeds into the main output and is not stored to global memory, then `out=None` and `feeds_main=True`:
```python
FlexGemmRuntimeLocalReducePlan(
FlexGemmLocalReduceGeometry(group=group, axis=axis),
out=None,
callbacks=generated_callbacks,
feeds_main=True,
)
```
Also, N-axis feed-main groups up to 32 stay as ordinary TensorSSA epilogue values, so those do not need a runtime local-reduce plan
The identity/init value lives in the TensorSSA reduction spec. For physical reductions, Inductor generates the `combine_fn` and `finalize_fn`; there is no separate QuACK init-value slot. This is in the cutedsl op overrides
```python
"sum": init_val=0.0, combine="lhs + rhs", finalize="value"
"prod": init_val=1.0, combine="lhs * rhs", finalize="value"
"max": init_val=-inf, combine="cutlass.max(lhs, rhs)", finalize="value"
"min": init_val=inf, combine="cutlass.min(lhs, rhs)", finalize="value"
"mean": init_val=0.0, combine="lhs + rhs", finalize="value / group"
```
`cutlass.max/min` is intentional here because physical `amax`/`amin` need NaN-propagating behavior matching PyTorch; this is not `fmax/fmin`.
## Newish QuACK interface
This is the reduction interface. I have wrestled many times with how much codegen we should own in Inductor. Ideally it is just as much as we need and no more. In that vein, I tried to design specific slots: give QuACK the generated combine and finalize behavior, and QuACK owns the physical accumulation, layout, synchronization, replay, broadcast, and store mechanics.
For the inter-thread/warp reductions we need these functions:
```python
FlexGemmLocalReduceCallbacks(
combine_fn=flex_gemm_epilogue_..._local_reduce_combine_fn,
finalize_fn=flex_gemm_epilogue_..._local_reduce_finalize_fn,
)
```
We add `__cache_key__` on those functions, and runtime turns them into QuACK registry keys:
```python
local_reduce_combine_key, local_reduce_finalize_key = (
register_runtime_local_reduce_callbacks(local_reduce, epilogue_key)
)
gemm_act(
...,
tensor_epilogue_key=epilogue_key,
tensor_epilogue_returns_aux=returns_aux,
**local_reduce_kwargs,
)
```
`local_reduce_kwargs` expands to the new QuACK slots:
```python
{
"tensor_epilogue_returns_local_reduce": quack_local_reduce_out is not None,
"local_reduce_feeds_main": local_reduce.feeds_main,
"local_reduce_out": quack_local_reduce_out,
"local_reduce_group": local_reduce.group,
"local_reduce_axis": local_reduce.axis,
"local_reduce_combine_key": local_reduce_combine_key,
"local_reduce_finalize_key": local_reduce_finalize_key,
}
```
### Which slots are new in this stack?
| QuACK-facing slot | New in this stack? | First active PR/commit that makes it a real boundary kwarg | Notes |
| --- | --- | --- | --- |
| `tensor_epilogue_key` | No | Pre-existing before this stack | Existing generated tensor-epilogue registry key. |
| `tensor_epilogue_returns_aux` | No | Pre-existing before this stack | Existing same-shape aux-output signal. |
| `tensor_epilogue_returns_local_reduce` | Yes | #188469 / `afff9f05d9c` | Distinguishes a compressed local-reduce aux store from ordinary aux. The old analysis-only PR #188468 was folded into this PR. |
| `local_reduce_out` | Yes | #188469 / `afff9f05d9c` | The compressed aux output tensor that QuACK stores into. Runtime adds the leading batch dimension before calling QuACK. |
| `local_reduce_group` | Yes | #188469 / `afff9f05d9c` | The grouped M/N size. #188470 / `915f0630e9d` broadens the valid range for larger groups. |
| `local_reduce_axis` | Yes | #188469 / `afff9f05d9c` | `0` means M-axis grouping and `1` means N-axis grouping. |
| `local_reduce_combine_key` | Yes | #188469 / `afff9f05d9c` | Registry key for the generated physical combine callback. Initially needed for physical compressed aux; #188470 extends the physical range and #188112 uses it for physical feed-main. |
| `local_reduce_finalize_key` | Yes | #188469 / `afff9f05d9c` | Registry key for the generated physical finalize callback, with the same rollout as the combine key. |
| `local_reduce_feeds_main` | Yes | #188112 / `16e34a0ef42` | Adds the second physical consumer mode: use the reduced value inside the main epilogue instead of, or in addition to, storing `local_reduce_out`. The old follow-up PR #189190 was folded into this PR. |
| `tensor_epilogue_scalar_biases` | Yes | #189188 / `6c987bbcf62` | Adds captured `[1, 1]` scalar epilogue arguments for tensorwise quantization. This is not a local-reduce slot, but it is another new QuACK-facing slot in this six-PR stack. |
#188470 does not add a new slot; it changes the meaning and requirements of the existing group and callback slots so larger local-reduce groups can use physical combine/finalize. #188739 adds quant/MX lowering but no new QuACK boundary slot. #189493 is boundary cleanup and adds no user-visible reduction slot.
QuACK also has a tiny registry for generated physical callbacks. We do this so that we can compile in parallel and then cache-hit on the key:
```python
register_local_reduce_fns(
combine_key,
combine_fn,
finalize_key,
finalize_fn,
)
```
## Active stack: what each commit is doing
I had claudex generate this
The actual active stack is six PRs. Oldest is at the bottom in ghstack, but the landing order is:
```text
#188469
-> #188470
-> #188112
-> #188739
-> #189188
-> #189493
```
### 1. #188469 / `afff9f05d9c` — compressed local-reduce aux outputs
This is the real base PR now. It includes the analysis that used to live in #188468, the compiler/output/runtime plans, compressed output allocation, generated callbacks, runtime validation, and the initial QuACK reduction boundary.
### 2. #188470 / `915f0630e9d` — larger local-reduce groups
This extends compressed reductions to N-axis groups larger than one TensorSSA fragment and M-axis groups that span more lanes, warps, or the CTA. It adds the config capability checks and the larger physical combine/finalize paths.
### 3. #188112 / `16e34a0ef42` — feed local reductions into epilogues
This adds feed-main, including the trailing pointwise and axis-1 fragment work that used to be in #189190. M-axis feed-main uses the same-warp physical path through group 32. N-axis groups through 32 stay in TensorSSA and can feed later pointwise/cast operations without a runtime plan.
### 4. #188739 / `da5a3dc63a9` — quant lowerings
This adds the MX E8M0 and NVFP4 E4M3 scale lowerings. MX works for both fragment TensorSSA values and generated physical finalizers, including the zero, saturation, infinity, NaN, and `max_power` cases.
### 5. #189188 / `6c987bbcf62` — scalar captured epilogue args
This adds the `scalar` captured-argument kind for `[1, 1]` read-only tensors, which unlocks tensorwise float8 quantization epilogues.
### 6. #189493 / `0274dde67b4` — simplify QUACK lowering boundaries
This does not add another reduction capability. It separates ordinary FlexGEMM backend dispatch from QUACK-specific lowering, requires the selected canonical QuACK config key, and derives output dtype from the preallocated output.
## Some limitations
Some of these can be removed, but this stack is already insane:
- local reductions currently support dense, non-batched `aten.mm` only; BMM/BAddBMM local reductions are a follow-up;
- feed-main supports M-axis same-warp groups up to 32 and N-axis TensorSSA groups up to 32;
- larger physical groups are currently compressed-store-only;
- local-reduce plans do not compose with `C`, non-default `alpha`, or non-default `beta` yet;
- `swap_ab` local-reduce configs are not supported yet;
- blockwise/2-D reductions are deferred to a follow-up and are not in this stack;
- only one compressed/physical local-reduction contract is supported per epilogue today, although the compressed-output plumbing stays plural for future multiple reductions;
- the grouped reshape must split exactly one original M or N dimension and reduce only that grouped dimension;
- explicit reduction dtype, multi-axis reductions, chained reductions, and cross-CTA final reductions are not supported yet.
Unsupported cases reject during analysis, config selection, or runtime validation rather than silently selecting an incompatible kernel.
## Validation
Each retained commit was tested independently with fresh Inductor and QuACK caches:
| PR | FlexGEMM result |
| --- | ---: |
| #188469 | 163 passed, 3 subtests |
| #188470 | 181 passed, 3 subtests |
| #188112 | 259 passed, 3 subtests |
| #188739 | 268 passed, 3 subtests |
| #189188 | 274 passed, 3 subtests |
| #189493 | 274 passed, 3 subtests |
Final validation also includes the QuACK vendor tests, vendoring reproduction, targeted lint including PYREFLY, and `py_compile`.
Pull Request resolved: #188469
Approved by: https://github.com/mlazos
https://github.com/NVIDIA/cutlass/releases/tag/v4.5.3 Summary: Bump the mirrored CUTLASS submodule pins under `fbcode/caffe2/third_party` and `xplat/caffe2/third_party` from upstream commit `da5e086dab31d63815acafdac9a9c5893b1c69e2` to the CUTLASS `v4.5.3` tag commit `4552152794e8bd3bcfd63cf9b44369e590420dba`. Test Plan: arc lint fbcode/caffe2/third_party/cutlass.submodule.txt xplat/caffe2/third_party/cutlass.submodule.txt Differential Revision: D111177502 Pull Request resolved: #189332 Approved by: https://github.com/Skylion007, https://github.com/ngimel
Update the torch-xpu-ops commit to [intel/torch-xpu-ops@2429ff](intel/torch-xpu-ops@2429ff8). New feature: - Support at::xpu::sleep on XPU to back torch.xpu._sleep Build System: - Align device debug flags (XPU_DEVICE_DEBUG, DEBUG_XPU, DEBUG) with CUDA conventions - Disable disable SYCL-TLA (USE_SYCLTLA OFF) in Debug build to avoid excessive memory usage with debug info Kernel optimization/enabling/fixes: - Use if constexpr for compile-time-constant branches in SGD/LogAddExp/RadixSelect - Replace manual subgroup tree reduction loop with reduce_over_group in Reduce Sum where reduce range < sg_size - Add _foreach_clone XPU implementation - Fix read/write races in topk single-workgroup kernel by adding memory barriers - Fix stale found-flag in radix select findPattern causing incorrect topk results Pull Request resolved: #189864 Approved by: https://github.com/EikanWang, https://github.com/guangyey
…#189958) Summary: Adds a version guard for `cudaEventRecordWithFlags` to fix build failures on ROCm versions before 7.0. The `hipEventRecordWithFlags` API is unavailable in older ROCm releases, so we fall back to the plain `cudaEventRecord` call for compatibility. Test Plan: Verified the build failure was resolved in the CI link provided. The change maintains backward compatibility while allowing newer ROCm versions to use the flags variant. Differential Revision: D112035112 Pull Request resolved: #189958 Approved by: https://github.com/ngimel, https://github.com/fegin
## Summary This PR introduces hardware classification for PyTorch test classes. Tests can now explicitly declare their hardware classification through a `hw_classification` class attribute, and users can filter which categories of tests to run via `--hw-classification`. When `--hw-classification` is not specified, all existing test discovery and execution paths remain unchanged. --- ## Changes ### Hardware classification Added `HardwareClassification` with three categories: - `GENERIC` - `DEVICE_GENERIC` - `CPU`, `CUDA`, `XPU`, `MPS` — specific built-in device types ### Test filtering Added `--hw-classification` support across unittest, pytest, subprocess, parallel, XML, repeat, and `run_test.py` execution paths. Filtering is implemented through: - `HardwareClassificationTestLoader` - `HardwareClassificationPytestPlugin` - `_filter_suite_by_hw_classification()` ### Tests Added unit tests in `test/test_testing.py` covering: - hardware classification filtering - inherited hw classification metadata ### Initial rollout Annotated test classes in `test/profiler/test_memory_profiler.py` to validate the rollout and provide an initial set of classified tests. --- ## Test Plan Verified on CUDA `python test/test_testing.py -v -k 'TestHardwareClassifications'` ```bash ---------------------------------------------------------------------- Ran 4 tests in 0.006s OK ``` Validated the feature using hardware classification annotations in `test/profiler/test_memory_profiler.py`: (34 total test methods: 20 GENERIC + 14 DEVICE_GENERIC) Verified: - `GENERIC` runs 20 tests - `DEVICE_GENERIC` runs 14 tests - combined `GENERIC` + `DEVICE_GENERIC` runs 34 tests - unmatched filters run 0 tests Verified the filter works across the main execution paths: - direct unittest execution - pytest / `--use-pytest` - `python -m pytest` - `run_test.py` forwarding - repeat mode - parallel mode - XML output Commands: ```bash # direct unittest execution python test/profiler/test_memory_profiler.py --hw-classification GENERIC # pytest / --use-pytest python test/profiler/test_memory_profiler.py --use-pytest --hw-classification GENERIC # python -m pytest python -m pytest test/profiler/test_memory_profiler.py --hw-classification GENERIC -v # run_test.py forwarding python test/run_test.py -i profiler/test_memory_profiler --hw-classification GENERIC # repeat mode python test/profiler/test_memory_profiler.py --repeat 2 --hw-classification GENERIC # parallel mode python test/profiler/test_memory_profiler.py --run-parallel 2 --hw-classification GENERIC # XML output python test/profiler/test_memory_profiler.py --save-xml /tmp/test_xml2 --hw-classification GENERIC # unmatched filter (0 tests) python test/profiler/test_memory_profiler.py --hw-classification CUDA # combined filter (all tests) python test/profiler/test_memory_profiler.py --hw-classification GENERIC DEVICE_GENERIC ``` Pull Request resolved: #186918 Approved by: https://github.com/fffrog, https://github.com/albanD, https://github.com/jbschlosser
This commit uses CMake's build instrumentation feature to collect and summarize build metrics. This requires CMake version >= 4.3. The PR is not set up to handle older versions as of this commit. Resolves: #189078 Pull Request resolved: #189455 Approved by: https://github.com/drisspg, https://github.com/atalman
This PR is auto-generated nightly by [this action](https://github.com/pytorch/pytorch/blob/main/.github/workflows/nightly.yml). Update the pinned torchcomms hash. Pull Request resolved: #189957 Approved by: https://github.com/pytorchbot
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 subscribe to this conversation on GitHub.
Already have an account?
Sign in.
19 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.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )