mirror of
https://github.com/pytorch/pytorch.git
synced 2026-10-28 18:56:51 +08:00
main
79
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 PyTorch MergeBot
|
21756fd7f3 |
Revert "[COW] track lifetimes of cow tensors via explicit completion events, harden materialization (#200800)"
This reverts commit
|
||
|
Simon Layton
|
948343523e |
[native_aot] bmm_outer_product: first Triton AOT op, block-size bucket dispatch (#196508)
Human Notes: Use the new shiny triton AOT toolchain to embed an op. > AI-assisted description (Claude Code and Codex): > > Adds embedded CUDA kernels for outer-product `torch.bmm` (`K == 1`) on Hopper and Blackwell. Four kernels per target cover float32 and bfloat16 with `N >= 128`: `BLOCK_M=32` for `M <= 96`, and `BLOCK_M=64` above that. Masks handle M and N tails. > > AOT and JIT share the instrumented kernel in `triton_kernels.py`, with the AOT builder beside it. The prerequisite exporter change accepts the kernel object directly. The declaration remains importable without Triton. > > Dispatch guards enforce layout, alignment, and int32 limits while preserving ATen's dtype and overlap errors. > > Test Plan: > > ```bash > spin develop > python test/python_native/test_native_aot_bmm.py -v > python -m unittest tools.test.test_native_aot_tools > spin quicklint -- -a > ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/196508 Approved by: https://github.com/albanD ghstack dependencies: #200446 |
||
|
Tomasz SzulistandClaude Opus 5.5
|
484ec6f0a5 |
[Test] Refactor test_dlpack.py and enable on XPU (#195722)
## Linked issue or supporting maintainer Part of #142029 ## BC-breaking? No ## Summary This PR updates `test/test_dlpack.py` to reduce repeated device-specific test coverage and enable XPU coverage where tests are backend-agnostic. ## Agent details ### What changed - Moved the device-independent tests out of the device template into a plain `TestTorchDlPack`, so each runs once instead of once per backend; the device-generic tests stay in `TestTorchDlPackDevice`. - Switched the copy, NumPy cross-device, stream and read-only wrapper tests from `onlyCUDA` / `onlyOn(["xpu", "cuda"])` to `@onlyAccelerator`, with `@skipMPS` on the four stream tests MPS cannot run. - Dropped the stale XPU skips for intel/torch-xpu-ops#3074. - Ported the two cross-stream exchange tests (`test_dlpack_conversion_with_diff_streams` and its narrow-precision variant) to `torch.Stream` and `Stream.native_handle`, and skipped them on XPU as a feature gap (intel/torch-xpu-ops#5641); the five that test the CUDA/ROCm stream numbers stay `@onlyCUDA`. - Added hardware classification metadata to all test classes. - Test-only change; no product code was modified. ### Test coverage delta | Metric | Before | After | |---|---:|---:| | XPU tests collected | 161 | 140 | | XPU passed / skipped | 117 / 44 | 115 / 25 | | CPU tests collected | 161 | 140 | | CPU passed / skipped | 113 / 48 | 92 / 48 | | Device-independent tests collected | 9 | 30 | | Device-independent passed / skipped | 7 / 2 | 28 / 2 | | Total tests collected | 331 | 310 | | Total passed / skipped | 237 / 94 | 235 / 75 | No tests failed. The 21 tests moved to `TestTorchDlPack` now run once instead of once per device, which accounts for the drop in collected tests. On XPU, `test_copy`, `test_needs_copy_error`, `test_dlpack_exchange_api` and `test_dlpack_exchange_api_sliced` now pass, and the 18 XPU skips of the two cross-stream exchange tests point at intel/torch-xpu-ops#5641 instead of `onlyCUDA`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/195722 Approved by: https://github.com/pbielak, https://github.com/etaf, https://github.com/jansel Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> |
||
|
Natalia Gimelshein
|
e506c0307c |
[COW] track lifetimes of cow tensors via explicit completion events, harden materialization (#200800)
## Summary (human written only) This PR is another approach to handling materialization of lazy clones. Instead of previously used cudaMemcpy it 1) allocates the materialized copy storage on the original stream where _lazy_clone was called, to have the same storage behavior 2) runs the copy on allocation stream, we are guaranteed that the data is correct on allocations so no sync to the current stream is needed 3) syncs to the current stream so that it has correct data To enable "stealing" of the last remaining reference, this approach requires users to mark the last use of any of lazy-cloned siblings `torch._C._cow_record_done(tensor)`. If we determine that we have the last remaining reference of the lazy clone, and all previous uses have a recorded event, we are free to reuse the last remaining reference, after syncing to all the recorded events. This way, we are syncing to the earliest possible point on other streams after the data is no longer needed, so the stalls will be minimal. <!-- Add your own summary and explain the relevance of the agent details below. --> ## Agent details (if applicable) AI-generated description (Codex), covering both commits on this branch. CUDA COW materialization used a synchronous cudaMemcpy on the legacy default stream. That copy was unordered with respect to nonblocking operation streams and was absent from CUDA graph replay, allowing replay to read stale data. Enqueue materializing copies on the original allocations reuse stream. Allocate replacements on the original allocation's reuse stream and establish event dependencies in the direction of current op stream. The same event operations provide eager ordering and captured graph dependencies, without recordStream or host synchronization. Retain replaced COW references with their logical storage to protect reads submitted before materialization. Add `torch._C._cow_record_done(tensor)` so callers can explicitly complete those prior uses. The caller first joins all relevant uses, including aliases and earlier pointer accesses, onto the current stream. Completion records an event and releases retained references. An active COW reference remains alive until destruction or materialization; later data access cancels its completion. Each sibling has its own completion record, and renewed completion preserves its earlier event dependencies. When every departed sibling completed explicitly, the last sibling can steal the original allocation after waiting on the recorded events, without fencing unrelated later work on the allocation stream. Releasing retained references also enables subsequent lazy clones to share the new allocation. Callers that omit completion retain the conservative copying behavior. Capture boundaries remain explicit: completion and shared materialization must stay within the context's capture. Creating COW storage for an external graph input during capture is rejected. Private-pool outputs cloned after capture, unsupported sharing streams, and allocators without allocation metadata use eager copies. Unique materialization preserves the allocation's address. Compatibility impact is limited to explicit CUDA `_lazy_clone()` users: unsupported capture crossings now raise errors, and some lazy clones fall back to eager copies. CPU behavior is unchanged. Deferred materialization can introduce dependencies at the later operation, so this does not promise the exact stream behavior of an eager `clone()` at the original `_lazy_clone()` call. Validation: built on GB200 with CUDA 13.1. Passed 16 C++ COW tests, 43 native CUDA COW tests, 96 existing tensor lazy-clone tests, and two cudaMallocAsync fallback tests (157 passed total). Tests cover independent readers, copies and stealing on different streams, allocator reuse, completion renewal, chained clones, capture boundaries, repeated graph replay with changing inputs, and reduce-overhead compilation. Targeted `spin lint` and pre-commit `lintrunner -a` passed. Pull Request resolved: https://github.com/pytorch/pytorch/pull/200800 Approved by: https://github.com/ezyang |
||
|
PyTorch MergeBot
|
2f8c573ca4 |
Revert "[COW] track lifetimes of cow tensors via explicit completion events, harden materialization (#200800)"
This reverts commit
|
||
|
Natalia Gimelshein
|
20f82b619b |
[COW] track lifetimes of cow tensors via explicit completion events, harden materialization (#200800)
## Summary (human written only) This PR is another approach to handling materialization of lazy clones. Instead of previously used cudaMemcpy it 1) allocates the materialized copy storage on the original stream where _lazy_clone was called, to have the same storage behavior 2) runs the copy on allocation stream, we are guaranteed that the data is correct on allocations so no sync to the current stream is needed 3) syncs to the current stream so that it has correct data To enable "stealing" of the last remaining reference, this approach requires users to mark the last use of any of lazy-cloned siblings `torch._C._cow_record_done(tensor)`. If we determine that we have the last remaining reference of the lazy clone, and all previous uses have a recorded event, we are free to reuse the last remaining reference, after syncing to all the recorded events. This way, we are syncing to the earliest possible point on other streams after the data is no longer needed, so the stalls will be minimal. <!-- Add your own summary and explain the relevance of the agent details below. --> ## Agent details (if applicable) AI-generated description (Codex), covering both commits on this branch. CUDA COW materialization used a synchronous cudaMemcpy on the legacy default stream. That copy was unordered with respect to nonblocking operation streams and was absent from CUDA graph replay, allowing replay to read stale data. Enqueue materializing copies on the original allocations reuse stream. Allocate replacements on the original allocation's reuse stream and establish event dependencies in the direction of current op stream. The same event operations provide eager ordering and captured graph dependencies, without recordStream or host synchronization. Retain replaced COW references with their logical storage to protect reads submitted before materialization. Add `torch._C._cow_record_done(tensor)` so callers can explicitly complete those prior uses. The caller first joins all relevant uses, including aliases and earlier pointer accesses, onto the current stream. Completion records an event and releases retained references. An active COW reference remains alive until destruction or materialization; later data access cancels its completion. Each sibling has its own completion record, and renewed completion preserves its earlier event dependencies. When every departed sibling completed explicitly, the last sibling can steal the original allocation after waiting on the recorded events, without fencing unrelated later work on the allocation stream. Releasing retained references also enables subsequent lazy clones to share the new allocation. Callers that omit completion retain the conservative copying behavior. Capture boundaries remain explicit: completion and shared materialization must stay within the context's capture. Creating COW storage for an external graph input during capture is rejected. Private-pool outputs cloned after capture, unsupported sharing streams, and allocators without allocation metadata use eager copies. Unique materialization preserves the allocation's address. Compatibility impact is limited to explicit CUDA `_lazy_clone()` users: unsupported capture crossings now raise errors, and some lazy clones fall back to eager copies. CPU behavior is unchanged. Deferred materialization can introduce dependencies at the later operation, so this does not promise the exact stream behavior of an eager `clone()` at the original `_lazy_clone()` call. Validation: built on GB200 with CUDA 13.1. Passed 16 C++ COW tests, 43 native CUDA COW tests, 96 existing tensor lazy-clone tests, and two cudaMallocAsync fallback tests (157 passed total). Tests cover independent readers, copies and stealing on different streams, allocator reuse, completion renewal, chained clones, capture boundaries, repeated graph replay with changing inputs, and reduce-overhead compilation. Targeted `spin lint` and pre-commit `lintrunner -a` passed. Pull Request resolved: https://github.com/pytorch/pytorch/pull/200800 Approved by: https://github.com/ezyang |
||
|
Simon Layton
|
fa62a650d0 |
[_native][8/8] Serve the ordered sum/prod overrides with the shared kernel's inner-tree order (#195062)
Human Notes: We finally get there! Cut over the existing inner-tree implementations to the new stack. Codex: Route the existing ordered `aten::sum` and `aten::prod` overrides through the shared fixed-DAG kernel. Eligibility, TensorIterator canonicalization, keepdim, and direct/out overload handling remain in the adapter. Unsupported dtypes and geometry continue to fall through to ATen. Nonzero input row strides and unit-stride outputs use the shared kernel. Noncompact rows use a dynamic leading stride; compactness and alignment key the plan, with underaligned rows selecting the bit-neutral unstaged path. Strided outputs retain the legacy same-order implementation. H100 uses a targeted native-suite list; B200's existing dynamic discovery covers the same suites without scheduling them twice. Labeling covers both implementation and test-only changes. End-to-end tests assert both reference bits and that the ordered route actually ran. Test Plan: ``` spin quicklint PYTORCH_SUM_INNER_TREE=1 python test/python_native/test_sum_cutedsl.py python test/python_native/test_sum_inner_tree_plan.py python test/python_native/test_inner_tree_order.py ``` Authored with assistance from an AI coding agent. Pull Request resolved: https://github.com/pytorch/pytorch/pull/195062 Approved by: https://github.com/ngimel |
||
|
Simon Layton
|
c21b8bbee2 |
Fix topk test pollution of CuTeDSL's GPU target (#200670)
## Linked issue or supporting maintainer Fixes #200665 Fixes #200577 Fixes #200626 Fixes #200706 ## BC-breaking? > No. ## Summary (human written only) Fix multiple test issues caused by a test-ordering issue & chosen arch for cuteDSL ## Agent details (AI-generated) > The no-device builder test's mocked GPU queries make CuTeDSL cache its > sm_100a fallback, breaking subsequent tests on GB300. Restore the cached > architecture after the test so cold no-device compilation stays covered > without affecting subsequent tests. > > Test Plan: > > Running the full file exercises the builder and numerical tests in one > process: all 7 passed on B200. > > ```bash > python test/python_native/test_topk_radix_kernel.py -v > ``` > > All 3 builder tests also pass with GPUs hidden: > > ```bash > CUDA_VISIBLE_DEVICES='' python test/python_native/test_topk_radix_kernel.py TestRadixKernelBuilder -v > ``` > > Verified on [GB300 CI](https://github.com/pytorch/pytorch/actions/runs/38425476816/job/115435448950): > all 7 radix tests passed, with no skips or retries. > > Authored with assistance from an AI coding assistant. Pull Request resolved: https://github.com/pytorch/pytorch/pull/200670 Approved by: https://github.com/ngimel |
||
|
Driss Guessous
|
60de3b7f33 |
[quack] Stop relying on the process-wide mixed_constexpr_if AST patch (#200710)
Summary: Vendored quack relied on upstream's `quack.dsl.mixed_constexpr_if`, which monkey-patches the global CuTe `DSLPreprocessor` so `if const_expr(S) and/or D:` folds the static term at trace time. The install is guarded by a class-level flag shared by every quack copy. When another copy loads first, it sets that flag, torch's install is skipped, and torch's quack kernels are traced without the rewrite. Inductor's sm100 `X @ X.T` -> quack symmetric GEMM path then fails with `NameError: cannot access free variable 'split_idx'` in `epilogue_split_k`, which killed GB300 GRPO jobs. The unscoped patch also rewrites every non-quack CuTe kernel traced in the process, such as FA4. Fix this in the source instead of at runtime: a new vendoring patch (`0010-explicit-constexpr-if.patch`) rewrites the seven mixed tests into explicit `if const_expr(S): ... elif D: ...` nesting. That is the same AST the rewrite produced, so codegen is unchanged. The patch also drops the `mixed_constexpr_if` import from `dsl/__init__.py`, and `vendor.sh` no longer vendors the module. Torch's quack no longer touches global DSL state, so import order with other quack copies does not matter. `test_quack_vendor` gains a check that the vendored tree has no leading-`const_expr` BoolOp `if` tests. Without the rewrite, those would silently lower to dynamic ifs. Supersedes D124861837 (mslk-side flag rename) and D124692177 (torch-side scoped patch). Test Plan: GB200 repro: preset the upstream flag to simulate mslk loading first, then run `torch._vendor.quack.gemm_interface.gemm_symmetric` on a 4096x4096 bf16 `x @ x.T`: - before, foreign copy first: `NameError: cannot access free variable 'split_idx'` (matches the GB300 job) - before, clean: OK - after, foreign copy first: OK, max abs err ~16 vs fp32 reference (bf16 rounding) - after, clean: OK `vendor.sh --check` reproduces the committed tree from the pinned upstream SHA plus the patch series. Differential Revision: D124906630 Pull Request resolved: https://github.com/pytorch/pytorch/pull/200710 Approved by: https://github.com/eellison |
||
|
Simon Layton
|
6a65f81736 |
[native] Include native operator sources in QuACK's cache fingerprint (#200057)
Human Notes: Fixes fairly longstanding cache-invalidation bug w/Quack RMSNorm cache invalidation. Flagged at https://github.com/pytorch/pytorch/pull/197008#issuecomment-6021427647 and possibly elsewhere. Codex: > Include `torch/_native/ops` in QuACK's cache fingerprint so native operator changes invalidate cached kernels. Preserve the registration when re-vendoring. > > Move the fingerprint test from the FlexGEMM suite to `test/python_native/test_quack_vendor.py` and generalize it to cover both source directories. > > Test Plan: > > ```sh > python test/python_native/test_quack_vendor.py TestQuackCache -v > bash tools/vendoring/quack/vendor.sh --check > ``` > > AI-assisted. Pull Request resolved: https://github.com/pytorch/pytorch/pull/200057 Approved by: https://github.com/pytorchgreenlight, https://github.com/drisspg |
||
|
Simon Layton
|
69b019cbf8 |
[native-aot] Model CUDA architecture families explicitly (#198229)
Human notes: My attempt at #197663 - add the concept of CUDA arch families, and try to build out explicitly as "the widest support that fulfils arch requests" -- use `a` suffix arch when requested, then the family `f` suffix, then no-suffix. Implementation clarification: TORCH_CUDA_ARCH_LIST states the devices a build supports; its suffixes do not select native-AOT compile targets. Each op's required ARCHS determines its candidates. Selection uses the widest compatible candidate, prefers `f` over plain when their coverage is equal, and retains deliberate exact-only compilation when an op declares only a target such as `sm_100a`. Codex: Make CUDA compatibility declarative and shared by export, generated C++ routing, and Python coverage. The family table records known device members explicitly, including SM100f support for SM100, SM103, and SM107, while plain targets forward to later known minors in the same major and `a` targets remain exact-only. Top-k now AOT-compiles only for SM90 and SM100f. A build's supported-device list is mapped through each declaration independently, so devices without a compatible top-k AOT target retain the JIT route. Unknown future build-list entries are ignored rather than failing the build. Generated routing derives its device gate from the targets actually shipped. For declarations with cpp_covers, codegen conjoins the shipped-target and ABI gates with that predicate. Top-k uses this path. Declarations may still omit cpp_covers and retain the existing Python covered_axes fallback; this change does not generate or call a second runtime coverage op. Review the architecture model in torchgen/native_aot_decl.py first, followed by the exporter, generated C++ routing, and manifest integration. The tests use an invented SM6 family to establish that the mechanism is not special-cased to SM10x. Test Plan: ``` /home/dev/pytorch/.venv/bin/spin lint /home/dev/pytorch/.venv/bin/lintrunner -a /home/dev/pytorch/.venv/bin/python test/python_native/test_aot_manifest.py /home/dev/pytorch/.venv/bin/python tools/test/test_native_aot.py /home/dev/pytorch/.venv/bin/python tools/test/test_native_aot_tools.py TestInt32SizeGate TestInt32GateTypeClassifier TestAotSourceGeneration TestDeclarationArchs TestEndToEndGeneration TestArchScopedGeneration /home/dev/pytorch/.venv/bin/python tools/test/test_native_aot_tools.py TestSidecarFieldValidation ``` Authored with assistance from an AI coding assistant. Pull Request resolved: https://github.com/pytorch/pytorch/pull/198229 Approved by: https://github.com/eqy |
||
|
Benedykt Bela
|
867370d452 |
[Test] Add HardwareClassification to test_cutedsl_smoketest.py and move device-generalized tests to test_dlpack.py (#195575)
### Changes - Add `HardwareClassification` to all test classes in `test/python_native/test_cutedsl_smoketest.py`. - Add `CUDA` suffixes to cuda-specific test classes. - Move generalized version of cuda-specific test case `test_readonly_wrapper_bare_capsule_is_valid` from `test_cutedsl_smoketest.py` into device-generic test class in `test_dlpack.py`. The cuda-specific test file `test_cutedsl_smoketest.py` has only one test that doesn't depend on cuda-specific CuteDSL and can be made generic. It is `test_readonly_wrapper_bare_capsule_is_valid` test function. To mark the whole test class as `HardwareClassification.CUDA` I generalized it and moved it to the `test_dlpack.py` not to loose the features test coverage. Pull Request resolved: https://github.com/pytorch/pytorch/pull/195575 Approved by: https://github.com/pbielak, https://github.com/etaf, https://github.com/jansel |
||
|
Simon Layton
|
4c3912140d |
[native] Optimize RMSNorm backward reduction and dispatch (#197011)
Human Notes:
Perf tuning.
Codex:
> Tune the RMSNorm backward persistent grid by normalized width, element size,
> and row count. For more than 32 partial rows, use the shared reduction with work
> distributed across rows within and between warps; retain the original geometry
> for smaller partials.
>
> Weight-only calls also compute an unused input gradient. Decline unprofitable
> cases to ATen before allocating native buffers, using measured thresholds.
> Review the launch policy and reduction together, then the eligibility guard and
> dispatch-boundary tests. Keep representative persistent-grid and gradient cases
> while removing the measured redundant test combinations.
>
> Performance
> -----------
> Measured on NVIDIA B200 with CUDA 13.2 on 2026-10-05 after rebasing onto
>
|
||
|
Simon Layton
|
145b18c1c7 |
[native] Reduce RMSNorm JIT host overhead (#197009)
Human Notes:
Be a bit cleverer w.r.t. caching, other launch-y things.
Codex:
> Cache shape-dependent RMSNorm eligibility calculations and reuse tensors whose
> shapes already match the launcher. Allocate normalization statistics as (M, 1),
> matching the JIT fake descriptor, so the common forward/backward path avoids
> flattening and reconstructing views. Hardware caching already comes from the
> shared utilities used by the preceding commit.
>
> Review the eligibility caches, then the adapter and matching statistics
> descriptor. Add autograd and CUDA graph coverage for the changed tensor lifetime
> and layout behavior.
>
> Performance
> -----------
> Measured on NVIDIA B200 with CUDA 13.2 on 2026-10-05 after rebasing onto
>
|
||
|
Simon Layton
|
1be50d8e09 |
[native] Share RMSNorm forward and backward kernel launchers (#197008)
Human Notes:
Rework the launchers, reduce repeated code.
Codex:
> Replace the RMSNorm adapter's per-call CuTe conversions and ATen weight-gradient
> sum/cast with general native CuTeDSL launchers. Compile against fake tensor
> descriptors and pass read-only tensors through TVM FFI for every eligible width,
> including mixed gradient dtypes. Remove the older vendored launch path.
>
> Review the kernel specification and launch policy, then the adapter. Reuse the
> shared dtype, hardware, capability, and launch utilities from main. Hoist tensor
> preparation, const pointer access, and partial-row reduction into shared modules.
> Type the new interfaces and add the missing Tensor.const_data_ptr stub definition.
> The tests cover layouts, alignment, copy-on-write, output masks, fake tensors,
> and fallback routing without duplicating the later persistent-grid cases.
>
> Performance
> -----------
> Measured on NVIDIA B200 with CUDA 13.2 on 2026-10-05 after rebasing onto
>
|
||
|
Simon Layton
|
09d2900a9f |
[native_aot] Track embedded coverage per op and device (#199296)
Human Notes: Add ability to skip AOT on a "AOT didn't build for my device" or a "if my build didn't include this op for AOT" basis. Inspired by @alband's review on #196508 Also fix up some infra issues exposed during review of #196508 Codex: > Python coverage could remove a call from JIT routing even when the build had > not embedded kernels for that op. Generate per-op architecture metadata from > the exported sidecars and require an op/device match before the Python > fallback claims coverage. The optional C++ coverage predicate retains its > existing fast path. > > Review the generated registration and runtime availability check first, then > the shared skipIfNoNativeAot(op, device=...) decorator and its top-k callers. > Tests now check what this build actually embedded for the requested device, > so a declaration supporting an architecture does not imply the build contains > kernels for it. > > Test Plan: > > ```bash > agent_space/tooling-venv/bin/python -m unittest tools.test.test_native_aot_tools > source agent_space/build-env.sh > .venv/bin/python test/python_native/test_aot_manifest.py > .venv/bin/python test/python_native/test_native_aot.py -v > source /home/dev/pytorch/.venv/bin/activate > /home/dev/pytorch/.venv/bin/spin quicklint -- -a > ``` > > The intermediate tree passed 370 tooling tests, 25 coverage tests, 18 top-k > tests, and lintrunner -a. Runtime tests reused the existing B200 AOT build. > > Authored with an AI assistant (Codex). Pull Request resolved: https://github.com/pytorch/pytorch/pull/199296 Approved by: https://github.com/albanD ghstack dependencies: #196507 |
||
|
Simon Layton
|
59cfff53ce |
[_native][7/8] Stage the inner-tree order's mid-band through smem, in tiles (#195061)
Human Notes: Inner tree optimization time! It was correct, now make it fast (or less slow). After this set of optimizations we're close to parity with unordered reductions, but not quite there. I believe this is as close as we can get - the only remaining change is the butterfly reduction across lanes in a warp, and for fixed ordering we must have them. Sad face. Codex: Stage coalesced global loads in shared-memory tiles so each lane can fold a contiguous run and finish with one butterfly. This removes repeated lane butterflies without changing the inner-tree DAG. The implementation uses one buffer and refills it only after folding, avoiding the large-M write-after-read race and limiting shared-memory occupancy. Staging is restricted to aligned, single-batch mid-band plans; ragged and multi-batch plans retain the register fold. Bitwise tests cover staged identity padding, and routing tests ensure the optimization actually runs. Test Plan: ``` spin quicklint python test/python_native/test_inner_tree_order.py python test/python_native/test_kernel_rowtile.py ``` Authored with assistance from an AI coding agent. Pull Request resolved: https://github.com/pytorch/pytorch/pull/195061 Approved by: https://github.com/ngimel ghstack dependencies: #195052, #195056, #195051, #195059 |
||
|
Simon Layton
|
6f790d4343 |
[_native][6/8] Add the inner-tree fold order, with the static-fragment datapath it needs (#195059)
Human Notes: Previously all datapaths use runtime trip counts, which allows a single kernel to service all N values without recompile. We will be adding in fixed-ordering (inner-tree) reductions which instead will require compile-time trip counts. This functionality is added as an optional thing when required. Codex: Add an opt-in reduction order whose addition DAG depends only on N and reproduces the existing ATen inner-tree sum/prod bits. Static fragments, streaming carries, ascending lane butterflies, and N-specific launch plans preserve that DAG while reusing the shared load, projection, and store body. This fixed order is the stack's basis for eager/Inductor bitwise parity. The default rolled folds remain runtime-shaped and continue sharing kernels by vector class. The fixed order compiles per N, covers all reference plan shapes, and rejects explicit requests it cannot honor rather than silently using another order. Golden hashes and differential tests pin all plan shapes, four dtypes, sum/prod, signed zero, batch independence, routing, and multi-field traits. Review the static datapath in `tile.py`, the plans in `kernel_rowtile.py`, then the host and device tests. Test Plan: ``` spin quicklint python test/python_native/test_tile_datapath.py python test/python_native/test_inner_tree_order.py python test/python_native/test_kernel_rowtile.py ``` Authored with assistance from an AI coding agent. Pull Request resolved: https://github.com/pytorch/pytorch/pull/195059 Approved by: https://github.com/ngimel ghstack dependencies: #195052, #195056, #195051 |
||
|
Simon Layton
|
b1e0adad7e |
[_native][5/8] Add general, column, and cross-CTA reductions (#195051)
Human Notes: Complete axis coverage and add the multi-CTA path for large rows. Codex: Complete axis coverage on the shared reduction body. The general driver represents TensorIterator geometry as runtime extent/stride pairs, allowing one structural kernel to serve arbitrary dimensions and layouts. The column driver vectorizes kept-axis outputs and splits tall reduced axes into partials. The cross-CTA driver serves few-row, large-N reductions with two fused launches. Exact divisor splits use the optimized row body; the general ragged split preserves coverage and global indices for prime lengths. The dispatcher now routes row, column, cross-CTA, reduce-all, and arbitrary-axis cases without adding a second general device kernel. Review `tile.py` and `kernel_general.py` before `kernel_coltile.py` and `kernel_xcta.py`. Test Plan: ``` spin quicklint python test/python_native/test_kernel_general.py python test/python_native/test_kernel_xcta.py python test/python_native/test_kernel_rowtile.py python test/python_native/test_kernel_coltile.py ``` Authored with assistance from an AI coding agent. Pull Request resolved: https://github.com/pytorch/pytorch/pull/195051 Approved by: https://github.com/drisspg ghstack dependencies: #195052, #195056 |
||
|
Simon Layton
|
cc5b03bf85 |
[_native][4/8] Serve narrow rows on the tile body, TMA-staged where that wins (#195056)
Human Notes: Optimization for tall-and-narrow reductions - large M, small N. Potentially use TMA for these cases where that has been measured to be faster. Codex: Add a one-thread-per-row configuration for tall, narrow reductions, avoiding the idle lanes imposed by a full-warp row mapping. For wider narrow rows, TMA stages whole rows into shared memory and rotates reads to avoid bank conflicts. The measured row-count and chunk gates retain direct loads where they are faster. TMA additionally requires SM90+, fp32, a power-of-two row length, and a lane stride at the measured direct-load cliff. The shared body validates those requirements even when callers force the option. Test Plan: ``` spin quicklint python test/python_native/test_kernel_rowtile.py ``` Authored with assistance from an AI coding agent. Pull Request resolved: https://github.com/pytorch/pytorch/pull/195056 Approved by: https://github.com/drisspg ghstack dependencies: #195052 |
||
|
Simon Layton
|
d183654bda |
[_native][3/8] Add shared tile reduction and the row driver (#195052)
Human Notes: Add the shared reduction body and the common contiguous-row fast path. Codex: Add `tile.TileReduce`, the shared reduction datapath responsible for vector loads, rolled folds, cross-thread merging, projection, and raw partial stores. The row driver owns launch policy, geometry, and plan caching, and shares one compiled kernel across each runtime vector class. The body supports multiple accumulator fields and one or two projected outputs. It is structured for later axes without introducing an axis-specific device kernel that a later commit must replace. Review `tile.py` before `kernel_rowtile.py` and its tests. Test Plan: ``` spin quicklint python test/python_native/test_kernel_rowtile.py ``` Authored with assistance from an AI coding agent. Pull Request resolved: https://github.com/pytorch/pytorch/pull/195052 Approved by: https://github.com/drisspg, https://github.com/ngimel |
||
|
Pearu Peterson
|
e42b20ef0c |
torch/_native: less per-call work in the chunked linear_cross_entropy loop (#196134)
Three changes independent of the kernel's own arithmetic: `grad_linear_weight`
is no longer zeroed before the loop, because the first chunk writes it outright;
the kernel emits each row's loss contribution instead of the two statistics it
was built from, so the per-chunk loss arithmetic disappears; and the kernel's
staging default becomes `(512, 8)`. Together they are about 6% on top of the
kernel path on an H100 and about 10% on a B200, with every gradient bitwise
unchanged.
> None of the three changes what the loop computes. The first two remove work
> that existed only because of how the loop was written; the third changes only
> how the kernel schedules the work it already does.
>
> `grad_linear_weight` was allocated with `zeros` and then accumulated into,
> so the first chunk read back a (C, F) buffer that was known to be zero, after
> a (C, F) fill that nothing consumed. The first chunk now writes it with
> `mm(out=)` and later chunks accumulate as before, which drops a fixed
> `2 * C * F` bytes per call -- independent of the batch, the chunk size and the
> chunk count, so it reads as a share of call length: 4.8% at N=1024, 1.0% at
> N=16384. The accumulator is still zeroed when there is no first chunk to
> write it, which is the empty-batch early return.
>
> The kernel emitted `lse` and the target logit, and the caller then formed
> `scale * (lse - target_logit)` and reduced it, per chunk, on (Bc,) data:
> four launches moving 2 KB each. Both values are already in registers where
> the row's statistics are finished, so the kernel forms that product itself
> and writes one vector; the caller keeps one slot per row of the call and
> reduces once, after the loop, with `sum(out=loss)`. Nothing else read the two
> statistics, so they are gone from the kernel's signature. That slot is the one
> buffer here that is not a chunk, so the forward-only path -- which forms the
> loss chunk by chunk and never reads it -- does not allocate it.
>
> The staging default moves from `(512, 4)` to `(512, 8)` -- eight column tiles
> staged per barrier in pass 2 rather than four. `tiles_per_stage` does not
> reach pass 1's reduction, so this is a scheduling change and nothing else.
> Measured by calling the kernel directly from a captured CUDA graph, so that no
> GEMM time is in the number, over seven `(Bc, V)` chunk shapes -- `V` in
> {4096, 8192, 16384, 32000, 65536} at `Bc = 4096`, plus `Bc` in {1024, 2048} at
> `V = 32000` -- in both dtypes on an H100 and a B200, so 28 points: `(512, 8)`
> is faster at every one, by 1.4% to 12.2%, and uses less energy at every one,
> by 3.3% to 8.5%.
>
> Energy is worth reporting separately because it does not follow from the time.
> Per-call joules come from NVML's total-energy counter read over windows
> aligned to its updates, with the idle gaps subtracted at measured idle power.
> Where the board is not power-limited the faster setting draws 2-4% more power,
> so it saves less energy than time -- on B200 fp16, 9.1-12.2% faster for
> 6.9-8.5% less energy. Where the board is at its limit the two coincide: B200
> bf16 at `V >= 32000` runs both settings at ~1000 W, and there the energy win
> equals the time win. On H100 bf16 the faster setting also draws less power,
> 673 W against 691 W at `V=65536`, so it saves more energy (4.0%) than time
> (1.4%). Neither metric prefers `(512, 4)` at any point measured.
>
> That kernel is 4-25% of a chunked call and the rest is three cuBLAS GEMMs, so
> the retune is worth 0.1-2.9% end to end, largest where `in_features` is small
> relative to `num_classes`.
>
> Measured through the op's 12-shape sweep, paired per point across three
> repetitions -- the two states differ only in two Python files, so the arms are
> swapped in place at each point and a pair is measured seconds apart in one
> session.
>
> | | H100 bf16 | H100 fp16 | B200 bf16 | B200 fp16 |
> | ---------------------------------- | --------- | --------- | --------- | --------- |
> | geomean | 1.058 | 1.065 | 1.095 | 1.099 |
> | min | 1.015 | 1.025 | 1.029 | 1.024 |
> | max | 1.143 | 1.160 | 1.377 | 1.390 |
> | above 1.0 in all three repetitions | 11/12 | 12/12 | 12/12 | 12/12 |
>
> The maxima are at `in_features=1024` on the B200 and `V=4096` on the H100,
> where the fixed costs dominate; the minima are at `in_features=8192` and
> N=16384, where the surrounding GEMMs are largest. The single point not above
> 1.0 in every repetition is H100 bf16 at N=16384: 1.037 over a 0.992-1.042
> spread.
>
> Peak memory moves by at most the one new buffer, the (N,) fp32 loss terms:
> 0 to +48 kB measured, never more than `N * 4` bytes and sometimes nothing at
> all where the allocator reuses a freed block.
>
> The loss's summation order changes -- one reduction over N rather than a
> partial sum per chunk -- so the value is deterministic but not guaranteed
> bit-identical to a per-chunk accumulation. Measured against the previous path
> at four shapes spanning one to four chunks, in both dtypes, `grad_input`,
> `grad_linear_weight`, `grad_linear_bias` and the loss all came back bitwise
> identical.
>
> Test Plan:
>
> The kernel's isolation test asserts the loss term against a torch reference
> rather than against `lse` and the target logit separately. That pins the
> difference only: the row max cancels out of the term, and out of `g`, so what
> those assertions see of it is its effect on the row sum and not its value.
> Its other job, keeping `exp` in range, is pinned by
> `test_a_lone_high_column_pins_the_row_max`, where the row's maximum sits
> alone in one thread's column and a partial lost at either reduction stage
> takes the row out of fp32. The uniform-logit test became a stronger invariant
> in the process (`term == scale * log(V)`, whatever the logit value), and the
> aliased-buffer test still pins that the target logit is read before any write
> can occupy its bytes, since the term is now the only path by which that logit
> reaches an output.
>
> ```
> python test/python_native/test_lce_fused_grad_logits_kernel.py
> python test/python_native/test_linear_cross_entropy_override.py
> python test/test_nn.py -k linear_cross_entropy
> OPINFO_RESTRICT_TO_DSL=cutedsl python test/test_ops.py -k linear_cross_entropy
> lintrunner -m origin/main
> ```
>
> All pass on an H100 (sm_90) and a B200 (sm_100): 55 kernel tests, 31 override
> tests, 226 test_nn linear_cross_entropy tests, and the DSL-restricted test_ops
> run over its 5 OpInfo entries.
---
_🤖 Drafted by Claude Code (an AI agent) and reviewed & approved by pearu._
Pull Request resolved: https://github.com/pytorch/pytorch/pull/196134
Approved by: https://github.com/slayton58
ghstack dependencies: #195829
|
||
|
Pearu Peterson
|
8ccd6e087b |
torch/_native: one CuTeDSL kernel for the chunked linear_cross_entropy backward (#195829)
Per chunk the chunked path is three cuBLAS GEMM calls plus a set of pointwise passes over the `(Bc, V)` logits buffer. This commit collects those passes into one CuTeDSL kernel which also writes the dense gradient of logits into the logits buffer itself, so the `index_add_` scatter and one buffer disappear. Measured on an H100, the call takes 0.61-0.87 of eager's time at a lower peak memory; on a B200, 0.30-0.78. The kernel is verified on sm_80, sm_90 and sm_100, with its block knobs tuned on H100. Accuracy is level or better in fp16; in bf16 the weight gradient's error runs 1.02x to 1.10x of eager's at the same chunk size, from the same accumulation count. > The `torch_nn::_linear_cross_entropy_batch_chunked` override stops delegating > and computes the loss and gradients itself for the slice its `cond` claims. > Per chunk: one cuBLAS matmul for the logits, one CuTeDSL kernel, and two more > matmuls for the gradients. > > The kernel is what the change is about. Eager walks the `(Bc, V)` logits buffer > five more times after the matmul -- row max, subtract, gather, `exp_`, row sum > -- and then scatters the one-hot correction into `grad_linear_weight` with > `index_add_`. Here one launch reads the raw logits twice (an online > max-and-sumexp pass, then an apply pass) and writes the dense > gradient-of-logits `g[n, v] = s_n * (softmax(z)[n, v] - [v == T_hat_n])`, > emitting the two `(Bc,)` statistics the loss needs. Forming `g` densely is what > makes both parameter gradients plain matmuls -- `grad_input = g @ W`, > `grad_linear_weight = g^T @ X` -- so the scatter disappears: the backward is > deterministic where eager's was not, and the separate `scatter_add` override > stops being reachable from this path. > > `g` is written into the logits buffer's own storage rather than a second > buffer, so a chunk costs one `(Bc, V)` buffer. That needs ordering, because > `g[n, j]` occupies the bytes of `z[n, j / 2]`: pass 2 stages a group of column > tiles in registers, synchronizes, and only then writes them, which is safe > because a group's writes reach at most halfway into the columns it just read. > The alternative that reaches the same footprint -- keeping the logits at low > precision, as eager does for fp16 -- was measured and rejected for bf16: bf16 > stores `z` with absolute error ~|z| * 2^-9 while the softmax depends on > DIFFERENCES of z, so gradients degrade once |z| reaches the tens, which is the > regime a trained head operates in. At fp16 the buffer follows eager down to two > bytes per element, since that is eager's own choice there and matching its > footprint is the compatibility contract; the exponential and the scaled > gradient still stay in fp32 registers, so this path is strictly more accurate > than eager's at an identical footprint. > > Review `cutedsl_impl.py` first: it holds the eligibility predicate and the > implementation, and it replaces the delegating implementation the routing commit > installed. Eligibility is shapes, dtypes and device only, so it is safe under > FakeTensor tracing; an input outside it never enters the override at all and the > router falls back, which keeps "the kernel ran" observable rather than hidden > behind an internal delegation. Its architecture condition is the DSL runtime's > floor, which is also where this path's own arithmetic becomes available -- the > logits matmul asks cuBLAS for an fp32 output from low-precision inputs, rejected > below 8.0 -- and it lives in one predicate that `cond` and the tests both call. > Holding out for a measurement on every capability would keep the kernel off > hardware it runs on perfectly well; what is unmeasured below sm_90 is > performance, not correctness. > > The implementation re-runs every check the op would have run -- an override > installed at a backend key replaces the body that ran them, and the checks below > it too -- by calling the op's own check functions, and it takes the per-row > target correction and gradient weight the same way. Nothing about the op's > contract is restated here; the previous commit moved each of those pieces out > so this one could call them. > > Then `fused_grad_logits_kernel.py`: one block per row, because a row's own > maximum and sum must be complete before any element of it can be written, and a > block is the largest unit that can share them without a grid-wide barrier. > > Its two shape knobs -- the block width, and how many column tiles pass 2 stages > per barrier -- are keyword arguments rather than constants, so a tuner or a > caller with measured shapes can choose them without editing the file, and each > combination compiles its own kernel. The defaults were measured on an H100 over > a 3 x 5 grid at 20 chunk shapes, which the comment beside them summarizes: no > setting wins everywhere, (512, 4) is never the worst and at most 1.39x behind > the per-shape best, and since the chunked loop spends ~85% of its time in > cuBLAS at the shapes where that share can be measured cleanly, the best setting > at the loop's own chunk shapes is worth under 2% end to end. Their legal ranges > are enforced, and one of them is a correctness limit rather than a preference > -- above 32 warps per row the cross-warp combine would silently drop the > surplus partials. > > Not covered here. The forward-only path (no gradient requested) has no `g` to > write and keeps eager's statistics. `reduction='none'` keeps its own op and > implementation. Probability targets, label smoothing, `acc_policy='accurate'` > and fp32 input fall back through the router. One accuracy trade is worth > naming: the dense `g` puts O(1) target-row contributions and O(1/V) off-target > ones through a single matmul where eager keeps the scales apart, which costs > ~14x on `grad_linear_weight` in isolation -- visible as ~15% at fp16 with > `reduction='sum'`, invisible at bf16 where other error sources dominate. It > buys the deleted scatter, the determinism and the speed. > > What the kernel buys, measured on an H100 through the op's own 12-shape sweep, > in two comparisons. > > Against the op's own chunked path, at the same shapes and the same chunk size > -- identical work, so this isolates the kernel: in bf16 the call takes > 0.61-0.87 of eager's time (1.14-1.63x faster) at 0.77-0.97 of eager's peak > memory; in fp16, 0.63-0.89 of the time at 0.86-1.00 of the peak. > > Accuracy splits by dtype. In fp16 the override is at least as accurate as > eager everywhere: `grad_linear_weight` error is equal, 0.996-1.001 of > eager's, and `grad_input` error is below it, 0.824-0.975. In bf16 the weight > gradient is the one that gives ground -- 1.018-1.100 of eager's error, > largest at `num_classes=16384` -- while `grad_input` stays level at > 1.000-1.021. The dense `g^T @ X` accumulates in a different order than > eager's `index_add_` scatter, which is where that comes from. > > The three matmuls are the same shapes at the same chunk size in both paths, > so the whole win is in what surrounds them. Per chunk, eager makes five > passes over the (Bc, V) logits buffer -- row max, subtract, exp, row sum, scale > -- and then scatters the one-hot correction into grad_linear_weight with > index_add_. Here one kernel reads the buffer twice and writes the dense > gradient-of-logits once, and the dense form makes the scatter part of an > ordinary matmul: less traffic and about a third of the launches (611 -> 195 per > call at D=1024). The gain tracks how much of eager's time those passes were: > 1.56x where the matmuls are cheap relative to the buffer (D=1024), 1.14x where > they dominate (D=8192). Memory: one (Bc, V) fp32 buffer per chunk, with the > gradient written into the same storage, replaces eager's logits buffer plus its > downcast copy and its (Bc, D) scratch. > > Against liger-kernel 0.8.2, at constant memory: for each shape the override's > chunk size is set so its peak matches liger's where that is reachable > (0.997-1.025 of liger's peak at those points; elsewhere the whole batch fits > one chunk and the peak stays at 0.48-0.87 of liger's). At that budget the > override runs 1.25-6x faster than liger in bf16 (0.17-0.80 of its time) and > 1.25-7x in fp16 (0.14-0.80), and 1.2-4x faster than eager at its default chunk > size. The extra factor over the same-shape comparison is the memory itself: the > override needs one (Bc, V) buffer per chunk where liger materializes more, so > at equal peak it runs a chunk several times larger -- fewer chunks, fewer re- > accumulations of grad_linear_weight, matmuls with a longer inner dimension. The > chunk size stays under the caller's control; this comparison shows what the > design returns when given liger's memory. > > This kernel's module imports the CuTeDSL runtime at import time, so the two > tests that import every module under `torch/` -- > `test_public_bindings::test_modules_can_be_imported` and > `test_testing::TestImports::test_circular_dependencies` -- need it in their > optional-runtime ignore lists, where the other `_native` DSL kernels already > sit. Both entries name the module rather than the package, so `cutedsl_impl` > and the package `__init__` stay checked -- they import without the runtime > and should keep having to. > > The new test files need the CuTeDSL runtime, so the B200 smoke job is the only > place they can run at all -- nothing else collects them. That job already > discovers every `test/python_native/test_*.py`, so they are picked up without > being named. Added to it explicitly: `test_nn -k linear_cross_entropy`, the > op's own accuracy harness. It compares fp16/bf16 against an fp64 reference > with tolerances the op's author calibrated, and it runs in every CUDA job > already -- but only where the DSL is installed does it measure this kernel > instead of eager, so this job is where it earns its keep. The `-k` lines have > to be named because `test_ops` and `test_nn` sit outside what that discovery > reaches. > > Test Plan: > > The numbers above come from the out-of-tree benchmark harness, one process per > configuration, over the same 12 shapes in both dtypes; the commands below are > the in-tree correctness checks. > > ``` > python test/python_native/test_lce_fused_grad_logits_kernel.py > python test/python_native/test_linear_cross_entropy_override.py > python test/python_native/test_override_declarations.py > python test/python_native/test_registry.py > python test/test_nn.py -k linear_cross_entropy > python test/test_ops.py -k linear_cross_entropy > spin lint > ``` > > All of the above pass -- 54 kernel tests, 30 override tests -- including the > two upstream batteries: `test_nn.py` runs 226 linear_cross_entropy tests and > `test_ops.py` 191, both clean. Tests that > need the kernel to be eligible skip where it is not, so a machine below its > architecture floor reports skips rather than failures. > > The kernel's numerics are checked three ways. Against a torch expression of its > contract, on its own, for three output dtypes across shapes that cross the block > width and the staging group, including both aliasing ratios -- an fp32 buffer > with narrower gradients, where a write lands halfway back, and the fp16 buffer > where it lands on the column just read -- and including logits large enough > that an unshifted exponential would overflow. End to end by the op's own > fp64-referenced harness, which this change routes into the job that has the > DSL. And two identical runs are required to produce bit-identical gradients, > with `index_add_` asserted never to run, so the determinism the deleted scatter > bought is pinned rather than asserted in prose. > > The op-level tests cover the paths a caller can reach rather than the kernel > alone: the forward-only branch, which reimplements the loss with eager ops and > which any `no_grad` call takes; `reduction='sum'`, a class weight and an > `ignore_index` that fires, which are the four `_neg_weight_target` branches the > row scale is built from; the empty batch, where the one allocation left > unwritten has to be zeroed; and an out-of-range target, which traps as > eager's index assert does, checked in a child process because the trap takes > the CUDA context with it. --- _🤖 Drafted by Claude Code (an AI agent) and reviewed & approved by pearu._ Pull Request resolved: https://github.com/pytorch/pytorch/pull/195829 Approved by: https://github.com/slayton58 |
||
|
Warren Deng
|
d602395970 |
[torch._native] Fall back when the Triton driver is unavailable (#197695)
# Summary Outer-product `bmm` dispatches to Triton through `torch._native`. However, driver setup can fail when Python headers, a host compiler, or another prerequisite from Triton is unavailable even if `triton` is importable. In this case, the user hits failure instead of falling back to aten implementation. This is the failure reported in #196977. # Change Add an optional `try_initialize_runtime` hook to native override registration. The router calls it after the predicate matches and before invoking the implementation. A `False` result continues to the next override or the captured ATen kernel. Conditional Triton overrides use a cached `_try_initialize_driver` hook backed by `triton.runtime.driver.active`. Unconditional overrides do not accept runtime initialization hooks, and graph construction does not run them. This keeps import and nonmatching operations lazy while leaving implementation, kernel-compilation, and launch errors visible. # Test plan ```bash python test/python_native/test_registry.py -k runtime_initialization_falls_through -v python test/python_native/test_registry.py -k fake_tensor_shape_inference -v python test/python_native/test_registry.py -k unconditional_override_cond_none -v python test/python_native/test_registry.py -k unconditional_override_rejects_runtime_initialization -v python test/python_native/test_registry.py -k override_node_dataclass -v python test/python_native/test_registry.py -k register_op_override_basic -v python test/python_native/test_native_dsl_ops.py -k triton_driver_initialization_is_cached -v python test/python_native/test_native_dsl_ops.py -k version_skip_env_var_overrides -v python test/python_native/test_decomp.py -k export_ -v ``` All focused tests passed. Fresh-cache missing-header and invalid-compiler runs used ATen; a healthy cold start used Triton. The cached hook measured 47.8 ns, and repeated fallback measured 30.3 μs versus 15.9 μs for direct ATen. Authored with assistance from Codex Fixes #196977 Pull Request resolved: https://github.com/pytorch/pytorch/pull/197695 Approved by: https://github.com/slayton58 |
||
|
Pearu Peterson
|
55f1d787ee |
torch/_native: OpInfo coverage for the chunked linear_cross_entropy overrides (#195334)
The overrides should be tested by the existing OpInfo machinery rather than by hand-written tests, so we add `cutedsl` variants of the two chunked ops whose samples are the eager generators' samples filtered by the op's own eligibility predicate. These variants exist only where the CuTeDSL runtime is installed, so we also add them to the B200 smoke job, which is the only job that has the runtime. > `torch/_native/README.md` names an OpInfo entry in `dsl_ops_by_dsl` as the > preferred way to test an override, so that the existing `test_ops` / > `test_meta` / gradient / functorch infrastructure applies. This adds one > variant per chunked op, `cutedsl` and `cutedsl_none`. > > The samples are the eager generators' whole space, filtered rather than > hand-picked, so the variants track those generators instead of drifting from > them. `_cutedsl_linear_cross_entropy_eligible` is the filter, and it asks > `F.linear_cross_entropy`'s own predicate -- extracted in the parent commit for > this caller -- whether a chunked op is reached, so the selection cannot drift > from the dispatch it is predicting. Both variants share the filter, because > `reduction` selects *which* of the two ops runs rather than whether either does, > and each generator already fixes that. Nothing an op receives is filtered out. > > Sharing the predicate removes the risk of the filter drifting, but not of the > filter handing it the wrong pieces of a `SampleInput`, nor of the predicate > disagreeing with what the call actually does. So > `test_sample_filter_matches_the_gate` checks both across each generator's full > space, in both grad modes, and that the filter admits something. The signal it > uses is that the shared accumulator runs if and only if a chunked op does -- the > reference path uses `linear` + `cross_entropy` instead -- so its call count > answers exactly the question the predicate answers, with no dispatcher > introspection. > > Two details there are load-bearing. The mirror is evaluated in the same grad > mode as the call it predicts, because one of its clauses reads > `torch.is_grad_enabled()` and evaluating it under a different mode would make > the two disagree by construction. And both grad modes are covered, since > eligibility genuinely differs between them -- a grad-requiring probability > target goes to the reference path -- and that clause is unreachable otherwise. > > Exceptions from the call are not tolerated. On supported hardware no sample of > either generator raises, in either grad mode, so catching them could only hide > a regression: an override that started throwing would be counted and skipped > rather than reported. > > Everything else in the two entries is copied from the eager `chunked` and > `chunked_none` variants, because the behaviour is the op's rather than the > override's: dtypes, `gradcheck_nondet_tol`, the COW materialization allow-lists, > the inherited autograd limits, and the JIT / Inductor / decomposition skips. > What differs is only what the override imposes -- CUDA only, CuTeDSL required, > and the filtered samples. `cutedsl_none` additionally takes > `check_batched_grad=False`, `check_batched_gradgrad=False` and the > vmap-over-grad expected failure, mirroring `chunked_none`, whose backward > recomputes from saved inputs. > > Three declarations live in the test files rather than the OpInfo, next to the > eager variants' equivalents: the `test_fn_gradgrad` xfails in > test_ops_gradients.py, and the `TestCompositeCompliance::test_backward` skip in > test_ops.py that `chunked_none` already carries (its recompute-in-backward does > an in-place `mul_` that is not composite compliant). > > The three variants' skip lists overlapped by fourteen declarations -- the > limits any implementation of the chunked op inherits, whatever computes the > loss: no Inductor lowering, an options dataclass JIT cannot script, > `test_cow_input`'s unspecified launch failure, and autograd support stopping at > a first-order VJP. Those move into `_lce_chunked_op_skips`, which the eager > `chunked` / `chunked_none` entries and both `cutedsl` entries now build on, so > the inherited set has one definition instead of three. > > What stays per variant is what actually differs, and that distinction matters > more than it looks: an over-applied `expectedFailure` reports an UNEXPECTED > SUCCESS, so it is not safe to inherit one. `test_jvp`, > `test_normalize_operator_exhaustive` and `test_vmap_autograd_grad` are > therefore declared only where the failure was observed -- the first two on the > eager variants, the third on `chunked_none` alone -- and the MPS accuracy and > NNC entries stay with the eager variants, which are the ones that run there. > > The shared generator also gains two large-logit shapes. Every sample it built > until now takes its tensors from the shared linear generator, which tops out > around |z| = 20 with a median of 6 -- and accuracy on this op can depend on |z| > well past that, because an implementation may hold the logits at low precision > (eager does for fp16) and such a buffer stores `z` with absolute error > proportional to `|z|` while the softmax depends on DIFFERENCES of `z`. The two > shapes raise the magnitude two ways that separate those effects. One puts a > uniform `linear_bias` over deliberately mild logits -- narrow input and weight > ranges, so the spread stays a few units while the bias carries |z| into the > hundreds -- which leaves the softmax mathematically unchanged and isolates > storage rounding. The other uses a wide weight range with the target at the > argmax, which is genuinely peaked. Both use class-index targets, like the rest > of this generator's chunked samples. They are > emitted in whatever flavour the call asks for -- chunked options for the chunked > variants, `reduction='none'` for that variant, none for the reference -- so > every variant sees them, including the eager ones, whose fp16 path has the same > exposure and nothing else that exercises it. > > The variants also need a job that collects them. They are constructed only when > the CuTeDSL runtime is importable at a version registration accepts -- a variant > built where `cu.register_op_override` drops every registration would run eager > under a `cutedsl` label -- and smoke_b200 is the one job that has it -- > and its include list held no test that iterates `op_db`, so until now the > `cutedsl` and `cutedsl_none` entries ran nowhere in CI, and neither did the > routing test. Both join that list. `OPINFO_RESTRICT_TO_DSL=cutedsl` narrows > `op_db` to the DSL entries and `-k` narrows to this op, so the run is the two > variants and nothing else: 27 tests in 37s on an H100. Keeping it to this op > means a failure there points at this stack rather than at another DSL override. > > Two more suites join it for the same reason. The variants' expectations are not > confined to `test_ops`: `test_ops_gradients.py` carries two xfails and the > shared skip tuple names nine `TestOperators` entries, and an expectation that > never runs rots into an unexpected success unnoticed. The gradient run is > restricted the same way; functorch's is selected with `-k "linear_cross_entropy > and cutedsl"` instead, because restricting `op_db` trips `opsToleranceOverride`, > which asserts that every op it names is present. > > #193606, open at the time of writing, replaces that kind of enumeration with > discovery of `test/python_native/**`. When it lands, the `--include > python_native/test_linear_cross_entropy_override` line here becomes redundant > and should be deleted; the `OPINFO_RESTRICT_TO_DSL=cutedsl ... --include > test_ops` line must STAY, because the OpInfo variants live in `test_ops`, which > that discovery does not reach, and they exist only where the runtime does. Both > patches edit `test_python_smoke_b200()` around the same anchor -- this one > immediately before the `test_linalg -k "mm or addmv"` line, #193606 immediately > after it -- so whichever lands second wants a manual rebase there rather than a > clean automatic one. > > Test Plan: > > Verified on an H100 (sm_90). Each declaration above was added only after > observing the corresponding failure there -- not predicted -- which is what > kept `test_vmap_autograd_grad` off the `cutedsl` variant, where it passes. > > The B200 job itself can only run in CI, so what was verified locally is that > its four new invocations select what they claim: 7 tests for the routing file, > 27 for the DSL-restricted `test_ops` run, 2 passed and 2 xfailed for the > gradient run, and 2 passed and 20 xfailed for functorch's. > > The large-logit samples were run through every battery that consumes this > generator, on an H100: `test_ops`, `test_ops_gradients` (its fp64 gradcheck > legs, where a peaked softmax was the plausible risk), `test_meta`, `test_nn`, > `test_ops` restricted to the cutedsl variants, and functorch's `test_ops` -- > all clean, and no unexpected successes among the 60 expected failures there. > > The shared-tuple extraction is behaviour-preserving by construction and was > checked as such: every variant's resolved skip set was dumped before and after > the move -- test class, test name, device, dtypes, active_if and decorator kind > -- and the two dumps are identical, so no variant gained or lost a > declaration. > > ``` > OPINFO_RESTRICT_TO_DSL=cutedsl python -m pytest test/test_ops.py -q -k linear_cross_entropy > OPINFO_RESTRICT_TO_DSL=cutedsl python -m pytest test/test_meta.py -q -k linear_cross_entropy > OPINFO_RESTRICT_TO_DSL=cutedsl python -m pytest test/test_decomp.py -q -k linear_cross_entropy > OPINFO_RESTRICT_TO_DSL=cutedsl python -m pytest test/test_ops_fwd_gradients.py -q -k linear_cross_entropy > OPINFO_RESTRICT_TO_DSL=cutedsl python -m pytest test/inductor/test_torchinductor_opinfo.py -q -k linear_cross_entropy > python -m pytest test/test_ops_gradients.py -q -k "linear_cross_entropy and cutedsl" > python -m pytest test/functorch/test_ops.py -q -k "linear_cross_entropy and cutedsl" > python -m pytest test/python_native/test_linear_cross_entropy_override.py -q > python -m pytest test/test_ops.py test/test_meta.py -q -k linear_cross_entropy > python -m pytest test/python_native/ -q > python scripts/lintrunner.py > ``` > > Restricted to the new variants: 27 passed and 47 skipped in test_ops, 26 and 78 > in test_meta, 4 xfailed and 8 skipped in fwd-gradients, with the decomposition > (16) and Inductor (12) suites skipped as declared. Unrestricted but filtered to > them: 2 passed and 2 xfailed in gradients, 2 passed and 20 xfailed in functorch. > The routing tests pass (7). Across the whole `test/python_native/` directory 287 > pass, and the 8 failures are pre-existing and unrelated: six subtest failures > plus `test_prod_override_engaged` in test_sum_cutedsl.py, which needs > `PYTORCH_SUM_INNER_TREE=1` for the sum/prod overrides to register, and > `test_import_does_not_load_cutlass`, which holds only when its file runs alone. > > Unrestricted, against the same checkout without this commit: test_ops 90 -> 117 > passed, test_meta 75 -> 101. > > The filter agrees with the dispatch on every sample of all four > generator/grad-mode combinations, with 92, 52, 52 and 28 of 161, 161, 91 and 91 > samples eligible. --- _🤖 Drafted by Claude Code (an AI agent) and reviewed & approved by pearu._ Pull Request resolved: https://github.com/pytorch/pytorch/pull/195334 Approved by: https://github.com/slayton58 ghstack dependencies: #194509, #194510, #195331, #195332 |
||
|
Pearu Peterson
|
717d7a288d |
torch/_native: route the chunked linear_cross_entropy ops (#195332)
This adds the `torch._native` override for the two chunked ops together with its eligibility condition, while the implementation still delegates to eager. Keeping the routing in its own commit leaves the kernel commit about the kernel. > `linear_cross_entropy` puts a whole fused loss behind one dispatcher symbol: an > `(N, F) x (F, C)` logits matmul, a row-shifted softmax, the cross-entropy > reduction, and three gradient matmuls, chunked over the batch. Per chunk the > eager loop walks the `(B, C)` logits buffer about seven times -- matmul write, > row max, subtract, gather, `exp_`, row sum, scale -- before the gradient > matmuls read it again. Collapsing those passes is what a DSL kernel is for. > > This installs overrides for both chunked ops -- the scalar-reduction > `_linear_cross_entropy_batch_chunked` and its `..._no_reduction` sibling -- and > nothing else. Both impls delegate to the op's own accumulator, so routing > through them is behaviourally indistinguishable from not routing through them, > and both conds gate nothing: they are the hooks a kernel will narrow to the > shapes it implements. No performance change is expected or claimed. Keeping the > wiring in its own commit means a later numeric or performance regression is > attributable to the kernel rather than to the dispatch. > > Both ops are registered from one `_OVERRIDES` table, which the registration > loop and the tests both read, so adding a third override is a row rather than an > edit in three places. > > Two things are worth knowing to review it. > > Registration has an ordering constraint. Overrides are installed at the end of > `import torch`, but the module defining the `torch_nn::` ops cannot be imported > from `torch/nn/modules/__init__.py` -- `torch.library` does not exist that > early -- so it is imported lazily on first use and the ops are absent when the > registry runs. `register_linear_cross_entropy_overrides()` therefore imports it > before registering, after asking the same question `cu.register_op_override` > asks -- runtime present, not disabled, and at a known-good version -- so the > import is only paid where an override can actually be installed. > > Installing at a backend dispatch key means the op's body never executes, so the > checks at the top of it would be skipped. `_batch_chunked_impl` calls > `_check_batch_chunked_grad_flags` for that reason -- the function the parent > commit extracted for exactly this caller. The no_reduction op needs no such > call: it takes no `compute_*_grad` flags, and its probability-target guard sits > in its backward, registered above this dispatch key and so not displaced. The > checks *inside* the accumulator are inherited only because the impls delegate to > it; a kernel will have to re-enforce those too, which a comment at each call > site records. > > `torch_nn` joins `_ALLOWED_LIB_SYMBOLS` here rather than in the registry > commit, since this is the change that registers on it and tests it. > > Unlike `aten`, whose ops the build guarantees, these ops exist only after > `torch/nn/modules/linear_cross_entropy.py` has executed, so the binding between > a registration's string and its referent is a runtime property whose failures > appear only where the DSL is installed -- never on CPU or standard GPU CI. The > module therefore declares `_NAMESPACE`, `_DEFINING_MODULE` and `_OVERRIDES`, > the registrar loops over that same table, and > test/python_native/test_override_declarations.py imports each defining module > and resolves every declared symbol. It needs neither a GPU nor the DSL, so it > runs on every shard: a one-token typo in `_OVERRIDES` fails it. Loud failure at > import remains the contract for anything that slips past. > > Test Plan: > > Two tests, because there are two things to test. The impls are > behaviour-preserving, so comparing their output to the ops' compares the ops to > themselves; what is not self-evident is that the overrides are installed at all, > given the lazily-imported module they target. One test asserts `import torch` > left every entry in `_OVERRIDES` registered with a live cutedsl node and a live > router. The other exports an `nn.LinearCrossEntropyLoss` call per reduction and > applies the registry's decomposition table: the override's own > `_native::<node_id>` op appears in the graph exactly when the call routes to it. > `mean` reaches the scalar-reduction override and `none` the no_reduction one, so > those two cover both. Nodes are matched on `target.namespace` and compared > against the node ids the registry minted, since one op symbol is a prefix of the > other. Exactly one such node is expected: there is one call, and the `_native` > node is opaque, so overrides of the aten ops inside it cannot add nodes. > > An eager call leaves no trace, since the impls are behaviour-preserving and the > router logs nothing per call, so the graph is the only place the choice is > observable. > > Both tests skip on those same three conditions. An installed CuTeDSL whose > version is not known-good registers nothing, and a test that looked for the > overrides anyway would fail where it should skip. > > ``` > python -m pytest test/python_native/test_linear_cross_entropy_override.py -q > python -m pytest test/python_native/ -q > python -m pytest test/test_nn.py -q -k linear_cross_entropy > CUDA_VISIBLE_DEVICES=0 python -m pytest test/test_cuda.py::TestCuda::test_lazy_init -q > python scripts/lintrunner.py > ``` > > `test_lazy_init` covers the added import: registration must not initialize CUDA, > since that poisons `fork` for anything importing torch. --- _🤖 Drafted by Claude Code (an AI agent) and reviewed & approved by pearu._ Pull Request resolved: https://github.com/pytorch/pytorch/pull/195332 Approved by: https://github.com/slayton58, https://github.com/Neilblaze ghstack dependencies: #194509, #194510, #195331 |
||
|
Pearu Peterson
|
28355dd593 |
torch/_native: drop aten from the registry's now-generic internals (#194510)
This is a follow-up rename. After the previous commit the registry internals are namespace-agnostic, so the `aten` names in them are misleading. No behavior change. > Follow-up to the parent commit, which taught the override registry to work on > any namespace but deliberately left every identifier alone so the functional > diff stayed reviewable. The result is a handful of helpers whose names claim > more than they do: `_install_aten_override` installs into whichever namespace it > is handed, and `_aten_schema_tail` clones the schema of any op. > > Renaming and comment work only, no behavior change. `_install_aten_override`, > `_destroy_aten_override`, `_resolve_aten_overload`, `_aten_schema_tail`, > `_aten_override_libs` and the router closure's `_aten_overload` argument lose > the `aten`, and the comments and docstrings that explained routing, teardown and > the Dynamo shortcut in terms of aten now describe them in terms of the > overridden op. `test_decomp.py`'s fixture comments follow: > `_destroy_live_libs` destroys overrides in every namespace, not only the aten > ones that class registers, which is what its docstring and its call site now > say. The missing-op diagnostic also gains the hint that the module defining the > namespace may simply not be imported yet. > > Kept separate from its parent so the functional change reviews without the > rename noise. It can still be squashed into that parent on request, but no > longer dropped: later commits in this stack bind to the new names, importing > `_resolve_overload` in the override drift-guard test and `_override_libs` in > the routing test. > > Test Plan: > > ``` > python -m pytest test/python_native/test_registry.py test/python_native/test_decomp.py -q > python -m pytest test/python_native/ -q > python -m pytest test/test_scatter_gather_ops.py -q > python scripts/lintrunner.py > ``` --- _🤖 Drafted by Claude Code (an AI agent) and reviewed & approved by pearu._ Pull Request resolved: https://github.com/pytorch/pytorch/pull/194510 Approved by: https://github.com/slayton58, https://github.com/Neilblaze ghstack dependencies: #194509 |
||
|
Pearu Peterson
|
4109cbc8d0 |
torch/_native: allow overrides on namespaces other than aten (#194509)
`torch._native` could override `aten` ops only. Since `linear_cross_entropy`
lives in `torch/nn` and its chunked ops are registered in the `torch_nn`
namespace, the registry has to accept any namespace, which is what this commit
does.
> The override registry could only install kernels on the `aten` namespace:
> `register_op_override` rejected any other `lib_symbol`, and the internals
> hardcoded `aten` when resolving the op, cloning its schema, capturing the
> fallback kernel, and creating the IMPL library. That rules out overriding ops
> defined in Python via `torch.library.custom_op`, which live in namespaces of
> their own (e.g. `torch_nn`).
>
> Overriding such an op needs nothing new from the dispatcher. A
> `custom_op`-defined op registers its implementation at
> `CompositeExplicitAutograd`, so an override installed at a backend key takes
> precedence there, `torch.library.get_kernel` resolves that composite kernel as
> the fallback, and the op's own `Autograd` kernel still runs above the router.
>
> What changes is that the namespace becomes part of an override's identity:
> `_OverrideNode` carries `lib_symbol`, the graph key widens from
> `(op_symbol, dispatch_key)` to `(lib_symbol, op_symbol, dispatch_key)`, and
> every helper that used to assume `aten` takes the namespace as an argument.
> Identifiers are deliberately left alone so this diff stays small even though
> several helper names now claim more than they do; a follow-up commit does that
> renaming, and nothing depends on it landing.
>
> `_set_mask_unconditional` iterates `_graphs` and calls
> `_cleanup_and_reregister_graph`, so it takes the namespace too.
>
> The namespace check is replaced by an allowlist rather than dropped, so support
> stays opt-in and each addition arrives with its own tests. Being open-ended
> would be wrong in three ways: `_native` is the registry's own namespace, holding
> the opaque `_native::<node_id>` ops that carry the override impls, so a router
> installed there would route into itself; some namespaces have no backend kernel
> to capture as a fallback at all (`static_runtime` has none); and many hold only
> metadata or plumbing (`profiler`, `streams`, `mempool`, `export`), where a
> kernel would have nothing to accelerate. `_ALLOWED_LIB_SYMBOLS` starts as
> `{aten}` and carries a comment covering those constraints; `torch_nn` is added
> by the commit that registers the first override on it, under the same policy.
>
> `UserOrderingFn` keeps its `(op_symbol, dispatch_key, graph)` signature. Adding
> a namespace parameter would break every existing user ordering function and the
> `TORCH_PYTHON_NATIVE_USER_GRAPH_ORDER_FN` contract, so a function that needs to
> tell namespaces apart reads `node.lib_symbol` from the graph it is handed.
>
> One constraint is worth knowing before using this: overrides are installed at
> the end of `import torch`, so the op has to be in the dispatcher by then. A
> namespace defined by a lazily-imported module must be imported from the
> caller's `register_to_dispatch()` first, or cloning the schema fails; the error
> message names that as the likely cause.
>
> Both test classes destroy only the override-lib entries that are not the
> handle their snapshot holds for the key. On a machine with a DSL installed
> that map is not empty at teardown -- `import torch` leaves 13 production
> overrides live (scatter_add, topk, _fused_rms_norm, ...) -- so destroying
> every entry and reinstating the snapshot's handles would strip the dispatcher
> of kernels the registry still lists as installed, for the rest of the process.
> `pytest test/python_native/ -q`, this commit's own test-plan command, runs
> several files in one process, where anything sorting after test_registry.py
> would inherit that state.
>
> Teardown restores `_native_decomp_overrides` as well: installing an override
> records a compile router there in addition to the dispatcher entry, and
> `native_decomp_table()` merges that dict, so a router keyed on an op that dies
> with a test's scoped library would outlive the test and shadow later work
> through the decomposition table.
>
> An override a test installs over one already at that key needs more than the
> snapshot. `_install_override` destroys the library it replaces, so the
> snapshot's handle is dead and that override has to be rebuilt from the saved
> graph -- and the rebuild captures whatever kernel is live at that moment as its
> fallback. Teardown therefore destroys everything that is not the snapshot's own
> handle for its key, the replacement included, before `_restore_override_libs`
> puts the snapshot back minus the dead entries and `_register_all_overrides()`
> rebuilds them. A replacement left registered becomes the restored override's
> fallback, so every later call in the process that the restored override
> declines lands in a test's impl.
> `test_teardown_leaves_no_decomposition_router` and
> `test_teardown_restores_the_fallback_of_a_replaced_override` pin the two
> halves, both asserting against the state teardown restores, since neither leak
> is observable until the instance is gone.
>
> `test_same_op_name_in_two_namespaces_is_independent` never built the collision
> its name claims: `<ns>::twice` and `aten::mul` share no op symbol, so the
> per-op-symbol bucket -- what the widened key actually puts at risk -- was never
> populated by two namespaces at once. A second throwaway namespace now defines
> the same `twice` symbol, and two tests replace it: one asserting each namespace
> routes to its own override while both sit in one bucket, and one pinning what
> the namespace-blind filters do, since same-symbol collisions across namespaces
> become possible for the first time here.
>
> Two namespace-blind surfaces were only observable, never stated. The string
> filters keep their behaviour and gain documentation: `deregister_op_overrides`
> and `reenable_op_overrides` match bare op symbols in every namespace and
> silently ignore a qualified "ns::op" string. Disabling a symbol everywhere is
> a real use -- a kill switch -- so the semantics stay, pinned by a test. The
> README's ordering section says the same for user ordering functions, which are
> called once per namespace with identical `(op_symbol, dispatch_key)`
> arguments.
>
> Reporting is the one surface that changes behaviour instead.
> `get_dsl_operations` merged same-named ops from different namespaces into a
> single list of bare symbols, which cannot be resolved back to an op; it now
> takes `lib_symbol` (default "aten", so existing callers see what they saw) and
> reports one namespace per call. Merged reporting has no application that
> survives contact with a collision, and this commit's own tests treat
> same-symbol collisions across namespaces as first-class. The public wrapper in
> torch/backends/python_native takes the same argument -- its docstring is what
> the rendered docs show, and it never mentioned the merging.
>
> Test Plan:
>
> New `TestRegistryNonAtenNamespace` exercises a throwaway namespace whose only
> kernel is composite, mirroring the shape of a `custom_op` definition: the
> namespace reaches both the graph key and the node, `cond=True` routes to the
> impl, `cond=False` falls through to the composite kernel, an emptied graph
> restores the original kernel, two namespaces stay independent, and an op that
> is missing from the dispatcher fails loudly instead of silently registering
> nothing. The class patches the allowlist to admit its throwaway namespaces, and
> a separate test confirms that a namespace off the allowlist is rejected at
> registration time, using `_native` as the example since it can never be
> allowlisted.
>
> ```
> python -m pytest test/python_native/test_registry.py test/python_native/test_decomp.py -q
> TORCH_DISABLE_NATIVE_JIT=1 python -m pytest test/python_native/test_registry.py -q
> python -m pytest test/python_native/ -q
> python -m pytest test/test_scatter_gather_ops.py -q
> python -m pytest test/test_bmm_outer_product.py -q
> python -m pytest test/test_linalg.py -q -k polar
> python -m pytest test/test_foreach.py -q -k "native or foreach_mm"
> CUDA_VISIBLE_DEVICES=0 python -m pytest test/test_cuda.py::TestCuda::test_lazy_init -q
> python scripts/lintrunner.py
> ```
>
> The fixture's invariants only have state to exercise where production overrides
> are live at teardown, so `test_registry.py` (46 tests) and `test_decomp.py` (8)
> also run on an H100 with the CuTeDSL runtime installed, together with
> `python -m pytest test/python_native/ -q`.
---
_🤖 Drafted by Claude Code (an AI agent) and reviewed & approved by pearu._
Pull Request resolved: https://github.com/pytorch/pytorch/pull/194509
Approved by: https://github.com/slayton58
|
||
|
Nichols A. RomeroandCursor
|
f2be4ec90b |
[inductor][eager] Resolve which distribution provides the importable Triton module (#195022)
**Context:** split out of #194131, and addresses @slayton58's comment that the lookup returns `triton_rocm`'s version instead of `triton`'s and that "we should be careful for the event that" the two do not align. ### The problem Native DSL ops gate on the Triton version, read from distribution metadata. `version("triton")` resolves by distribution *name* and never checks which files that distribution installed: - Normalization never strips a suffix, so `triton-rocm` and `fbtriton` miss. The ops then silently do not register on an install where `import triton` works. - A provider swap leaves the previous dist-info behind, still reporting a version for a module it no longer owns (`TRITON_DISTRIBUTIONS` in `tools/torchtlx/dev.py`). ### The fix Collect versions from known distribution names before scanning `sys.path`. A single candidate is used directly. When multiple distributions report versions, compare the hash recorded for the resolved `triton` module with its bytes on disk; scan only if the known candidates remain unresolved. - `find_spec` locates the module without importing it. - Unlisted names still resolve through the scan, and names are compared using metadata normalization. - A missing RECORD or hash and an editable-install shim RECORD are undecidable; an empty RECORD is evidence of non-ownership. - One unreadable provider does not discard providers listed after it. Irreducible ties warn and preserve the prior preference for `triton`. A lone candidate is intentionally not verified, preserving `main`'s behavior and avoiding RECORD work on the common path. Consequently, a lone stale known name can still hide an unlisted provider. Also fixed: malformed metadata could raise during `import torch` and was logged only at `debug`; it is now caught and reported at `warning`. ### Cost Against `main`'s single `version("triton")`, the common path measured **0.27 ms -> 0.76 ms**, once per process behind `functools.cache`. A known same-path collision resolves in about **1.7 ms**. The full 419 ms `sys.path` scan is reserved for an unknown provider or unresolved ambiguity. ### Test Plan ``` python test/python_native/test_native_dsl_ops.py python test/test_public_bindings.py TestPublicBindings.test_correct_module_names lintrunner torch/_native/triton_utils.py test/python_native/test_native_dsl_ops.py ``` Tests cover known and unlisted names, same-path RECORD collisions with different hashes, editable installs, missing/empty metadata, unreadable providers, normalized names, ambiguity fallback, and the version gate boundaries. The collision was also reproduced in isolated virtual environments with two actual wheels installed in both orders; the lookup selected the wheel whose bytes were live each time. Made with [Cursor](https://cursor.com) Pull Request resolved: https://github.com/pytorch/pytorch/pull/195022 Approved by: https://github.com/jansel Co-authored-by: Cursor <cursoragent@cursor.com> |
||
|
Simon Layton
|
96600b6cc8 |
[Native DSL] Share topk policy across JIT and AOT (#196599)
Human Notes: Match JIT and AOT conditionals, unify expected coverage. Codex: Make the torch-free AOT declaration the shared source of truth for topk dtype, capability, K, N, and work-rung selection. JIT now imports that policy, gains the measured SM90 and bfloat16 regimes, and uses the same dynamic-N kernel variants as AOT. Replace the exact-N AOT grid with runtime-N radix kernels and one K=16 register branch ladder. The two one-shape register outliers, K=16/N=2048 and K=32/N=256, remain JIT-only because folding N=2048 into the ladder regressed full-wave latency by up to 72%. On B200, K=16/N=1024 is also an exact JIT-only rung because the shared ladder is 10.91% slower there. This reduces the embedded manifest from 48 to 29 kernels per architecture, or 96 to 58 across SM90 and SM100; the B200 exact rung does not add an AOT kernel. The deterministic dynamic-vector tail policy is selected at compile time inside the same manifest point. SM90 keeps the fixed four-iteration bound, while SM100 uses the runtime tail extent. Fixed-vector small-N and K=512 rungs retain their bounded work, and K=1024 retains its two-CTA launch bound. This gives each architecture its measured best aggregate policy without increasing the AOT kernel count. The B200 matrix below uses one full wave (M=148), direct int64 outputs, and the median of seven interleaved measurements. Positive deltas are slower and negative deltas are faster than an otherwise identical full-static kernel: | B200 selected regime | Median delta vs full static | Range | | --- | ---: | ---: | | FP32 register ladder, N=64-512 | +1.44% | -2.65% to +2.41% | | FP32 register exact JIT, N=1024 | -0.08% | -0.08% | | FP32 radix, small fixed vector, nondeterministic | +4.60% | +3.71% to +5.59% | | FP32 radix, small fixed vector, deterministic | +4.67% | +3.08% to +5.63% | | FP32 radix, dynamic vector, nondeterministic | +4.47% | -0.68% to +6.12% | | FP32 radix, architecture tail, deterministic | +2.82% | -0.90% to +4.46% | | FP32 radix, K=512 fixed vector | +3.08% | +1.85% to +3.35% | | BF16 radix, dynamic vector, nondeterministic | +5.63% | +1.03% to +6.71% | | BF16 radix, architecture tail, deterministic | +3.42% | -0.57% to +4.55% | | BF16 radix, K=512 fixed vector | +1.63% | +0.87% to +1.96% | For the dynamic-vector deterministic matrix, forcing the SM90 fixed tail on B200 raised the median delta from 3.2% to 4.8% for FP32 and from 3.5% to 5.2% for BF16, with worst cases above 10%. Exact tail residues were faster but would multiply artifacts by residue count; runtime vector extents were also slower than the retained fixed-vector rungs. Eligibility checks previously called `data_ptr()`, which materializes copy-on-write storage even when topk subsequently rejects the case. Use `const_data_ptr()` for both JIT and AOT checks. AOT dispatch also reuses its preloaded device properties and keeps the N-divisibility requirement on radix branches rather than the shared register prelude. Embedded H100 performance versus ATen at full-wave M: | Regime | Speedup vs ATen | | --- | ---: | | FP32 register | 1.63x to 2.96x | | FP32 radix | 1.61x to 3.53x | | BF16, K=512/1024 | 1.74x to 2.31x | AOT coverage is intentionally a strict subset of JIT eligibility for the exact register cases. Routing tests exercise AOT with JIT disabled, JIT with AOT present, the shared register ladder, the B200 exact rung, architecture-tail and dynamic work rungs, both determinism modes, bfloat16 functional and out variants, direct int64 indices, and copy-on-write inputs. Test Plan: ```bash TORCH_CUDA_ARCH_LIST='9.0;10.0' USE_CUDA=1 spin develop python test/python_native/test_native_aot.py python test/python_native/test_topk_cutedsl.py python test/python_native/test_topk_cutedsl_eligibility.py python test/python_native/test_topk_radix_kernel.py python tools/test/test_native_aot.py python tools/test/test_native_aot_tools.py spin lint ``` Authored with assistance from an AI coding assistant. Pull Request resolved: https://github.com/pytorch/pytorch/pull/196599 Approved by: https://github.com/ngimel ghstack dependencies: #196598 |
||
|
Simon Layton
|
c06a37dcbc |
[Native DSL] Make topk kernels N-dynamic (#196598)
Human notes: Remove static-`N` dependence in the kernels - move to dynamic/rungs as necessary. Codex: Replace exact-N radix compilation with runtime N and bounded work rungs. Scalar tails use a fixed maximum of four iterations, small vector regimes use a fixed maximum vector count, and deterministic K=512 retains its measured fixed-work rung. The register kernel uses runtime code branches for the K=16 N=64-1024 ladder while keeping occupancy outliers exact. The alternatives were materially slower: a fully runtime deterministic tail loop regressed K=1024 by 11-20%, and capacity-padding register ladders regressed by up to 3.5x. The selected bounded-work approach stays close to exact specialization: | Approach | Latency delta vs exact-N | | --- | ---: | | Register branch ladder, N=64-1024 | 0% to +9.3% | | Fixed max scalar tail | 0% to +5.5% | | Fully runtime scalar tail, deterministic K=1024 | +11% to +20% | | Fixed max vector count for small N | +0.8% to +4.9% | The same lower-commit policy was remeasured on B200 at one full wave (M=148). Each result is the median of seven interleaved measurements; positive deltas are slower than an otherwise identical exact-N kernel: | B200 regime | Median delta vs full static | Worst delta | | --- | ---: | ---: | | FP32 register ladder, N=64-1024 | +2.30% | +10.91% | | FP32 radix, nondeterministic | +4.58% | +6.03% | | FP32 radix, deterministic fixed tail | +4.82% | +9.97% | | FP32 radix, deterministic K=512 fixed vector | +3.04% | +3.25% | | BF16 radix, nondeterministic | +5.54% | +6.74% | | BF16 radix, deterministic fixed tail | +5.81% | +10.11% | | BF16 radix, deterministic K=512 fixed vector | +1.58% | +1.96% | Combining K=16, N=2048 into the runtime ladder regressed full-wave inputs by up to 72%, so it remains an exact JIT rung. Kernels now write int64 indices directly instead of launching a separate widening operation. Runtime-N vector loops use their runtime extent directly, while only fixed-work rungs retain an active-iteration guard. Static-N AOT callers continue to fold exact tail work. Offline AOT build workers do not have a CUDA driver. Setting `min_blocks_per_mp=2` made CUTLASS DSL 4.6.2 query device shared-memory attributes while compiling the K=1024 launch bounds. A neutral preferred shared-memory carveout avoids that device query while preserving `nvvm.minctasm` and its register cap. Builder tests cover the static lower-commit contract and compile this path with device queries forced to fail. Test Plan: ```bash TORCH_DISABLE_NATIVE_AOT=1 python test/python_native/test_topk_cutedsl.py python test/python_native/test_topk_radix_kernel.py python test/python_native/test_topk_cutedsl.py python test/python_native/test_native_aot.py python test/python_native/test_topk_cutedsl_eligibility.py spin lint ``` Authored with assistance from an AI coding assistant. Pull Request resolved: https://github.com/pytorch/pytorch/pull/196598 Approved by: https://github.com/ngimel |
||
|
Simon Layton
|
3e9e74c046 |
[_native][2/8] Add shared CuTeDSL and reduction machinery (#195050)
Human Notes: Shared infrastructure we'll build on for the rest of the stack * How we define an operator in a reasonably general manner -- per-element transforms, reduction specifics, etc. * How we will launch a given generated kernel, wrap read-only inputs, etc. * How to query the launch device and derive useful quantities for the eventual kernel definitions/launches. Centralise all of this logic and re-use for everything that comes afterwards. Codex: Add the reduction trait protocol, fake operand and launch helpers, hardware capabilities, dtype mappings, and an instrumented plan cache used by the remaining stack. Generic CuTeDSL helpers live in `torch._native.cutedsl`; reduction traits live with their consumers in `torch._native.ops.reductions`. Traits separate leaf conversion, accumulator combination, and serial reduction so the same operation can use rolled or fixed fold orders. Welford merging mirrors ATen's zero-count identity handling, and butterfly widths reject invalid non-power-of-two specializations. The B200 job relies on its existing dynamic native-suite discovery instead of scheduling these tests twice. CuTeDSL-dependent modules are imported lazily, keeping `import torch` DSL-free. Review `traits.py` first, then `launch.py`, `plan_cache.py`, and `hw_caps.py`. Test Plan: ``` spin quicklint python test/python_native/test_traits.py python test/python_native/test_hw_caps.py python test/python_native/test_instrumentation.py ``` Authored with assistance from an AI coding agent. Pull Request resolved: https://github.com/pytorch/pytorch/pull/195050 Approved by: https://github.com/drisspg ghstack dependencies: #195049 |
||
|
Simon Layton
|
7961344fe1 |
[_native][1/8] Add DSL-agnostic native-op utils (utils/) (#195049)
Human Notes: Basic shared functionality used throughout the stack. Canonicalise it here so equivalent implementations do not drift. Codex: Add non-raising tensor capability predicates and a lazy module proxy for native-op families. Registration and condition evaluation can use these helpers without importing a DSL runtime, while each family retains its own hardware allowlist and geometry checks. Centralizing the common predicates prevents equivalent implementations from drifting. No operators are registered in this commit. Test Plan: ``` spin quicklint python test/python_native/test_native_utils.py ``` Authored with assistance from an AI coding agent. Pull Request resolved: https://github.com/pytorch/pytorch/pull/195049 Approved by: https://github.com/karthickai, https://github.com/njriasan |
||
|
Simon Layton
|
85ae6d195e |
[native_aot] topk: first AOT declaration (fp32/bf16 radix kernels) + e2e tests (#190899)
Human Notes:
Finally do a thing! Add `topk` as an AOT-provided override in libtorch_cuda.
Claude:
Declares a 48-point grid (fp32/bf16 x N{2048..16384} x K{64,128,256}
x both determinism modes) over the radix-select kernel. The det
specialization (stable prefix-sum gather + lex (ord, -idx) sort) is
bit-exact vs aten including tie resolution, so det-mode calls are
served AOT too; the dispatch chain keys on
deterministicAlgorithms(). The kernel gains
in_dtype and index_dtype compile parameters; the AOT route writes
int64 indices directly (no widening copy). build(spec) lives beside
the JIT wrappers, so both routes compile the same class. fp16 and
off-grid shapes keep their JIT eligibility, as do calls the prelude
declines: a non-last dim, largest/sorted false, sub-full-wave M, or a
base pointer less aligned than the kernels were compiled for.
ARCHS is stated rather than defaulted, and covers Hopper as well as
Blackwell: measured on an H100, the AOT kernels beat aten on 48/48 of
the grid, median 2.16x (min 1.26x at M=4096/N=2048, max 2.99x at
M=256/N=4096 fp32). Both spellings of each capability are listed --
archs_from_cuda_arch_list maps 9.0 and 9.0a to different nvcc targets,
and a wheel built with the plain list would otherwise export trees
generation then filters out.
CI: a native_aot test config (test.sh) asserts the wheel actually
embeds the kernels -- a silently artifact-free wheel would make every
routing test vacuously pass on the JIT/aten fallbacks -- then runs the
python_native suites. One build+test pair per capability in
native-aot.yml: cuda-arch-list 9.0a tested on linux.aws.h100 and
10.0a on linux.dgx.b200, both built on a GPU-less builder
(path-triggered + nightly). test_uncovered_fp32_served_by_jit_layer
skips below sm_100, since the JIT topk override is gated there while
these AOT kernels are not. This lands with the first op so the AOT path
is exercised in CI from this commit onward.
covered_axes() and cpp_covers() decline a dim other than the last, largest=False
and sorted=False, matching the stub's gates. Coverage wider than the stub is not
wrong -- the stub declines and stock aten serves -- but it costs those calls
their JIT route, since coverage subtracts them from it (review: #190899).
Authored with an AI assistant (Claude Code).
Test Plan:
```
.venv/bin/python test/python_native/test_native_aot.py
# incl. det-mode routing + bit-exact-with-ties tests
# (test_deterministic_mode_routes_to_aot_bit_exact,
# test_deterministic_ties_bit_exact)
.venv/bin/python test/python_native/test_topk_radix_kernel.py
.venv/bin/python test/python_native/test_topk_cutedsl.py
```
On an H100, against the kernels a `spin develop` embedded:
```
.venv/bin/python test/python_native/test_native_aot.py # 11 tests, 1 skipped
.venv/bin/python test/python_native/test_aot_manifest.py # 18 tests
.venv/bin/python test/python_native/test_registry.py # 35 tests
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/190899
Approved by: https://github.com/MatthiasKohl, https://github.com/ngimel
ghstack dependencies: #196747, #196757
|
||
|
Simon Layton
|
3502231bc9 |
[python_native] Skip QuACK import tests on Python 3.10 (#196757)
Human Notes:
Necessary to make 3.10 builds not fail CI when AOT ops come in.
Note: We could change schema behavior to be consistent across 3.10 onwards, but
that feels like a much bigger risk than this..
Codex:
Installing CuTeDSL on Python 3.10 made the vendored QuACK import tests execute
there for the first time. QuACK's custom-op signatures use IntEnum defaults,
whose symbolic Python 3.10 string form produces invalid inferred dispatcher
schemas before any test body can run.
This is independent of native-AOT and its Python 3.10 support. Skip only the
QuACK import and operator-registration tests below Python 3.11; the patch-series
and vendoring reproducibility tests remain enabled. Normalizing IntEnum defaults
in torch.library would also make these imports work, but would change core
Python 3.10 behavior for an unrelated compatibility check.
Authored with assistance from an AI coding assistant.
Test Plan:
```
python - <<'PY'
import runpy
import sys
from unittest import mock
sys.path.pop(0)
import torch
with mock.patch.object(sys, "version_info", (3, 10)):
module = runpy.run_path("test/python_native/test_quack_vendor.py")
quack = module["TestQuackVendor"]
patches = module["TestQuackFlexGemmPatches"]
assert quack.__unittest_skip__
assert quack.__unittest_skip_why__ == "QuACK custom-op imports require Python 3.11 or newer"
assert not getattr(patches, "__unittest_skip__", False)
PY
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/196757
Approved by: https://github.com/ngimel
ghstack dependencies: #196747
|
||
|
drisspg
|
4d2bd99cd0 |
flex_gemm: Inductor-owned QUACK autotuning (#196174)
## Human Note Prior to the rebase on epimod we had inductor own tuning (via a patch in quack). In the base pr we let quack own autotning and it by default it is much slower than inductor. This switches us back. But one thing we gave up (intentionally) is the verification of configs and let epi-mod(quack) handle this. We keep the handiling there. We dropped our patch for autotuning and now we are adding back a modified version for the new api. This pr adds a way to pass in a representative of the problem ModProblem and not need to allocate tensors -> `legal_mod_configs`. We dropped the owned config list and this pr adds it back to heruistics. We now use this to create the usual multi lowering choices with different configs that have been verified `legal`. This drops tuning time down to 1/10 the time. And if in the future we want to update the input tensors ot be more representative input data ( have seen this affect optimal tuning choices) form this e.g. the input to the mm can have different distrubitons of 0,1 and be more or less power limited. ## Agent note The EpiMod cutover handed GemmConfig selection to QuACK: Inductor emitted one template choice and the generated kernel called `epimod(tuned=...)`, so QuACK's own tuner swept every legal config at first call through a second forkserver pool, and Inductor's precompile pool and autotune caches were bypassed. Cold tuned compiles regressed to about two minutes for a plain mm+relu. Inductor now owns selection while QuACK keeps legality. The lowering renders one unpinned probe of the template and runs its module on FakeTensors under `select_flex_gemm_configs`, so the exact EpiMod the kernel will build reports its legal GemmConfigs through QuACK's own `_legal_mod_configs` (the list form of its untuned selection, added in patch 0001). Each candidate becomes a pinned template choice (`config=` replaces `tuned=`) for `autotune_select_algorithm`; untuned calls pin only QuACK's default. The default search space is the measured SM100 priority order the previous FlexGEMM heuristics carried; other capabilities and `max_autotune_gemm_search_space=EXHAUSTIVE` benchmark every legal config. Each choice's kernel is compiled in-process at precompile time (serialized: Inductor precompiles from several threads), and a single untuned choice compiles during lowering rather than at first call. QuACK's async compile pool is not used: it ships EpiMods to workers via cloudpickle, which PyTorch does not depend on, and the fused route must not change behavior based on optional packages (a subprocess test asserts neither `cloudpickle` nor a pip `quack` is ever imported). Inductor benchmarks while `torch.compiler.is_compiling()` is still true, so `EpiMod.__call__` gains `compile_dispatch=False` for generated runtime callers instead of mutating the process-global compiler flag. This change was prepared with assistance from an AI coding tool. ## Performance Cold `tuned=True` compile+first call for `flex_gemm(mm, relu)` at 2048x2048x512 on a GB200 (caches disabled, QuACK object cache cleared): | configuration | candidates | seconds | |-----------------------------------|-----------:|--------:| | QuACK-owned tuner (before) | 196 | 131.5 | | Inductor-owned, EXHAUSTIVE | 196 | 124.3 | | Inductor-owned, default (this PR) | 12 | 12.2 | ## Test Plan ```bash python -m pytest test/inductor/test_flex_gemm.py test/python_native/test_quack_vendor.py test/inductor/test_symmetric_mm.py -q TORCHINDUCTOR_COMPILE_THREADS=1 python -m pytest test/inductor/test_flex_gemm.py -q -k tuned tools/vendoring/quack/vendor.sh --check ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/196174 Approved by: https://github.com/eellison |
||
|
PyTorch MergeBot
|
32d114b449 |
Revert "[native_aot] topk: first AOT declaration (fp32/bf16 radix kernels) + e2e tests (#190899)"
This reverts commit
|
||
|
drisspg
|
c0d27b73b4 |
flex_gemm: route QUACK through modern EpiMod (#192662)
The previous FlexGEMM integration carried a second QuACK runtime stack inside PyTorch: `gemm_act`, configuration mirrors, callback registries, template heuristics, and a long operation-specific vendoring series. Feature work had to cross two implementations, and QuACK could not own configuration legality, compilation, caching, or launch. Route the QUACK backend through QuACK's modern EpiMod frontend while keeping captured FX as the expression IR. PyTorch owns graph semantics, output planning, allocation, and graph-level fallbacks (`torch/_inductor/kernel/flex_gemm/`: `fx_cutedsl_codegen.py` plans QuACK-owned stores and emits a TensorSSA EpiMod source; `lowering.py`, `template.py`, `runtime.py` render and call it). The concrete FlexGEMM EpiOps are PyTorch-owned too (`quack_ops/`: `GroupedMainStore` and the grouped local reductions subclass QuACK's `EpiOp`); QuACK owns the generic EpiOp protocol and same-ABI kernel selection. Every generated callback is TensorSSA: the main body, the accumulator prepass, and the grouped-reduction finalizers all run once per register fragment through one emitter (`FlexGemmCuteDSLOpOverrides`), so there is no scalar/F2 code path; finalizers may read captured scalar operands (`finalize_operands`), while row/col/tile captures remain rejected, and integer inline-asm results keep their dtype until the store. The FX analysis stays in the shared `gemm_epilogue` / `gemm_epilogue_analysis` modules (`GemmEpilogueGraph`, `Normalized*`, `GemmLocalReduceAnalysis`, `GemmOutputPlan`), evolved to the EpiMod contract rather than copied into FlexGEMM. Output layouts are caller-owned descriptors carrying shape, CuTeDSL callbacks, legality, and carrier validation, which QuACK fingerprints and executes without a registry. Grouped and quantized main outputs, dense and caller-layout reductions, bool and NVFP4 storage, swap_ab, and inline assembly compose through that boundary. The public GEMM scope stays mm/addmm/bmm/baddbmm; block-scaled `F.scaled_mm` lands above this commit. Body GEMMs tagged by the HOP are excluded from `pad_mm` so the captured GEMM contract survives ragged shapes. The vendored copy moves to upstream QuACK v0.6.4 (`4709411169dcc3dc4e23f8f32f385eb4b6871d9d`) plus the hook-only series in `tools/vendoring/quack/flex_gemm_patches/` (0001-0003 here: fragmentwise TensorSSA main/prepass callbacks with an apply port and sink operands, `supports_config`/`supports_swap_ab` legality, `keep_tensorless`, `output_n`/ `store_tile_shape_mn`, caller-owned callback fingerprinting), rendered by `tools/vendoring/quack/vendor.sh`; the legacy GEMM implementation and its old patches are removed, and the QuACK-side tests for the hooks stay out of the series for upstreaming. Review order: `tools/vendoring/quack/`, then `torch/_higher_order_ops/flex_gemm.py`, `gemm_epilogue_analysis.py`, `lowering.py`, `fx_cutedsl_codegen.py`, `quack_ops/`, `runtime.py`/`template.py`, then tests. This change was prepared with assistance from an AI coding tool. ```bash tools/vendoring/quack/vendor.sh --check python -m pytest test/inductor/test_flex_gemm.py test/python_native/test_quack_vendor.py test/inductor/test_symmetric_mm.py -q ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/192662 Approved by: https://github.com/eellison |
||
|
Simon Layton
|
122777aa13 |
[native_aot] topk: first AOT declaration (fp32/bf16 radix kernels) + e2e tests (#190899)
Human Notes:
Finally do a thing! Add `topk` as an AOT-provided override in libtorch_cuda.
Claude:
Declares a 48-point grid (fp32/bf16 x N{2048..16384} x K{64,128,256}
x both determinism modes) over the radix-select kernel. The det
specialization (stable prefix-sum gather + lex (ord, -idx) sort) is
bit-exact vs aten including tie resolution, so det-mode calls are
served AOT too; the dispatch chain keys on
deterministicAlgorithms(). The kernel gains
in_dtype and index_dtype compile parameters; the AOT route writes
int64 indices directly (no widening copy). build(spec) lives beside
the JIT wrappers, so both routes compile the same class. fp16 and
off-grid shapes keep their JIT eligibility, as do calls the prelude
declines: a non-last dim, largest/sorted false, sub-full-wave M, or a
base pointer less aligned than the kernels were compiled for.
ARCHS is stated rather than defaulted, and covers Hopper as well as
Blackwell: measured on an H100, the AOT kernels beat aten on 48/48 of
the grid, median 2.16x (min 1.26x at M=4096/N=2048, max 2.99x at
M=256/N=4096 fp32). Both spellings of each capability are listed --
archs_from_cuda_arch_list maps 9.0 and 9.0a to different nvcc targets,
and a wheel built with the plain list would otherwise export trees
generation then filters out.
CI: a native_aot test config (test.sh) asserts the wheel actually
embeds the kernels -- a silently artifact-free wheel would make every
routing test vacuously pass on the JIT/aten fallbacks -- then runs the
python_native suites. One build+test pair per capability in
native-aot.yml: cuda-arch-list 9.0a tested on linux.aws.h100 and
10.0a on linux.dgx.b200, both built on a GPU-less builder
(path-triggered + nightly). test_uncovered_fp32_served_by_jit_layer
skips below sm_100, since the JIT topk override is gated there while
these AOT kernels are not. This lands with the first op so the AOT path
is exercised in CI from this commit onward.
Authored with an AI assistant (Claude Code).
Test Plan:
```
.venv/bin/python test/python_native/test_native_aot.py
# incl. det-mode routing + bit-exact-with-ties tests
# (test_deterministic_mode_routes_to_aot_bit_exact,
# test_deterministic_ties_bit_exact)
.venv/bin/python test/python_native/test_topk_radix_kernel.py
.venv/bin/python test/python_native/test_topk_cutedsl.py
```
On an H100, against the kernels a `spin develop` embedded:
```
.venv/bin/python test/python_native/test_native_aot.py # 11 tests, 1 skipped
.venv/bin/python test/python_native/test_aot_manifest.py # 18 tests
.venv/bin/python test/python_native/test_registry.py # 35 tests
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/190899
Approved by: https://github.com/MatthiasKohl, https://github.com/ngimel
ghstack dependencies: #190898, #196070
|
||
|
Simon Layton
|
f5ad4e7b5a |
[native_aot] runtime: router-level AOT coverage, C++ covers fast path, embedded detection, disabled() masking (#190898)
Human Notes:
Optimization of "is-this-covered-by-AOT" check by allowing a C++
implementation to be defined.
Allow control of AOT'd kernels at runtime where appropriate:
* If this is a sparsely-populated optimization (i.e. AOT covers
a few specific cases, and expects to fall back otherwise, allow opt-out
* If this is intended to have 100% coverage and be the only valid
implementation, do **not** allow opt-out
Claude:
torch/_native/aot_manifest.py loads declarations into coverage
objects; the eager router checks coverage once per call ahead of its
cond chain, so covered calls decline the JIT route and fall through
to the aten wrapper's stub. When the AOT library ships a cpp_covers
predicate (torch.ops._native_aot.covers_<id>), coverage prefers it
over the Python matching.
AOT kernels are always linked into libtorch_cuda (no separate
library): _native_aot_embedded() detects them in-process via the
_native_aot covers-op namespace. Two distinct off states, deliberately
different:
* Artifacts ABSENT (stage 2 never ran): coverage still declines
covered calls, the stub has no kernel, and they land on STOCK aten.
Artifact presence is static per process and coverage is a
declaration-time statement ("these points ship AOT"), not a runtime
probe; checking per call would tax every covered call for a
transient packaging condition. The correctness contract holds
(stock aten serves), only the acceleration source changes; tested
explicitly (test_covered_call_correct_without_aot_lib).
* Context switch OFF (set_aot_enabled(False), also flipped by
TORCH_DISABLE_NATIVE_AOT=1 and masked by
python_native.<dsl>.disabled()): coverage itself returns False, so
covered calls KEEP their JIT route rather than declining into a
stub that will not fire -- masking must not silently lose both
accelerated routes (test_masked_switch_restores_jit_coverage).
Inside disabled() the JIT layer is off too, so those blocks still
see stock aten end to end.
Unconditional overrides sit outside both of those, on both routes. An
override whose declaration says its kernels ARE the op's implementation
is exempt from every user-facing filter (registry.check_enabled) -- the
same rule the generated C++ gate follows in #190897 -- so
set_aot_enabled(False) and disabled() cannot change what such an op
computes. The one thing that reaches them is the private
_set_mask_unconditional() behind torch._native._unconditional_masked(),
which exists so a numerics test can still obtain stock aten values: a
reference silently computed through the override under test is worse
than no reference. It rebuilds every router rather than only setting the
flag, because check_enabled runs while a graph is being registered, not
per call -- so flipping the flag alone would leave the already-installed
routers in place and the mask would appear to do nothing.
Authored with an AI assistant (Claude Code).
Test Plan:
```
python test/python_native/test_aot_manifest.py
python test/python_native/test_registry.py
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/190898
Approved by: https://github.com/MatthiasKohl, https://github.com/ngimel
|
||
|
lirui927
|
13818df097 |
[Eager] Add FlyDSL topk native override (#193548)
### Motivation
Add FlyDSL-based register and radix-select TopK kernels as native eager overrides, on top of the FlyDSL infrastructure in the PR below.
### Supported cases
FP32 tensors on gfx950, contiguous input, reduction over the last dimension, and `largest=True, sorted=True`. N-D inputs are logically flattened to `(M, N)`.
Both functional and `out=` variants are supported. Read-only inputs use `ConstTensorWrapper`, allowing copy-on-write tensors to dispatch without materialization. Unsupported shapes, architectures, dtypes, layouts, and options retain ATen behavior.
The register kernel handles small fixed K values and is reproducible, using `(value desc, index asc)` ordering for ties, which may differ from ATen. The radix-select kernel handles continuous K ranges, padding non-power-of-two K values to the next complete bitonic sorting network. For finite inputs, its deterministic mode preserves ATen tie ordering. NaN payload and index choices may differ from ATen because NaNs are canonicalized to a single ordinal.
### When FlyDSL acceleration kicks in
The input must be on gfx950 with at least one row per compute unit (`M >= 256` on MI355X).
| kernel | K | N |
|:---|:---|:---|
| register | `{2, 4, 8, 16}` | power of two, `1024 <= N <= 8192` |
| radix | `64 <= K <= 256` | `8192 <= N <= 32768` |
| radix | `257 <= K <= 383` | `16384 <= N <= 32768` |
| radix | `384 <= K <= 831` | `32768 <= N <= 131072` |
| radix | `832 <= K <= 1024` | `32768 <= N <= 262144` |
### Caching
FlyDSL TopK uses two cache layers:
#### PyTorch in-process cache
- Caches the `CompiledFunction` returned by `flyc.compile()`.
- **Register key:** **`(N, K, arch, backend, device_index)`**
- **Radix key:** **`(N, K, deterministic, arch, backend, device_index)`**
- `backend` distinguishes compilation backends, while `device_index` distinguishes callables bound to different HIP device contexts.
#### FlyDSL compiled-artifact cache
- **Register specialization key:** **`(N, K, arch, backend)`**
- **Radix specialization key:** **`(N, K, deterministic, arch, backend)`**
- Caches persistent compilation artifacts to avoid repeating lower-level compilation.
- `device_index` does not affect generated code, so devices with the same architecture can reuse compiled artifacts. Each device only needs to load the artifact and create its own `CompiledFunction`.
A PyTorch cache hit bypasses `flyc.compile()` entirely. The FlyDSL compilation cache is consulted only when the PyTorch cache misses.
### Test plan
```bash
python test/python_native/test_topk_flydsl.py
OPINFO_RESTRICT_TO_DSL=flydsl python test/test_ops.py -k topk
lintrunner -a
```
The focused FlyDSL TopK suite passes 43 tests.
### Performance methodology
Latency is measured with CUDA events after 20 warmup iterations and over 100 timed iterations. Each result is the median of three runs. The ATen baseline runs with FlyDSL disabled; FlyDSL invokes the same `torch.topk` operation with the override enabled.
```python
def bench(fn, iters=100, warmup=20):
for _ in range(warmup):
fn()
torch.cuda.synchronize()
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(iters):
fn()
end.record()
torch.cuda.synchronize()
return start.elapsed_time(end) / iters * 1000.0 # microseconds
clear_topk_cache()
fly_v, fly_i = torch.topk(x, k, dim=-1) # compile
fly_us = bench(lambda: torch.topk(x, k, dim=-1))
with pn.flydsl.disabled():
aten_us = bench(lambda: torch.topk(x, k, dim=-1))
speedup = aten_us / fly_us
```
### Performance
GPU: AMD 256-CU MI355X (gfx950).
Register kernel:
<table>
<thead>
<tr>
<th rowspan="2">M</th>
<th rowspan="2">N</th>
<th rowspan="2">K</th>
<th colspan="3">non-deterministic</th>
<th colspan="3">deterministic</th>
</tr>
<tr>
<th>ATen (us)</th>
<th>FlyDSL (us)</th>
<th>speedup</th>
<th>ATen (us)</th>
<th>FlyDSL (us)</th>
<th>speedup</th>
</tr>
</thead>
<tbody>
<tr><td align="right">256</td><td align="right">1024</td><td align="right">8</td><td align="right">33.5</td><td align="right">18.6</td><td align="right">1.80x</td><td align="right">33.6</td><td align="right">26.7</td><td align="right">1.25x</td></tr>
<tr><td align="right">256</td><td align="right">8192</td><td align="right">16</td><td align="right">55.9</td><td align="right">28.0</td><td align="right">2.00x</td><td align="right">57.4</td><td align="right">31.5</td><td align="right">1.82x</td></tr>
<tr><td align="right">4096</td><td align="right">1024</td><td align="right">8</td><td align="right">72.1</td><td align="right">18.6</td><td align="right">3.87x</td><td align="right">72.8</td><td align="right">26.4</td><td align="right">2.76x</td></tr>
<tr><td align="right">4096</td><td align="right">2048</td><td align="right">2</td><td align="right">102.6</td><td align="right">18.6</td><td align="right">5.52x</td><td align="right">103.1</td><td align="right">26.6</td><td align="right">3.88x</td></tr>
<tr><td align="right">4096</td><td align="right">4096</td><td align="right">8</td><td align="right">147.5</td><td align="right">27.0</td><td align="right">5.45x</td><td align="right">147.1</td><td align="right">30.9</td><td align="right">4.77x</td></tr>
<tr><td align="right">4096</td><td align="right">8192</td><td align="right">16</td><td align="right">240.3</td><td align="right">67.3</td><td align="right">3.57x</td><td align="right">240.2</td><td align="right">71.3</td><td align="right">3.37x</td></tr>
<tr><td align="right">16384</td><td align="right">1024</td><td align="right">2</td><td align="right">224.4</td><td align="right">18.6</td><td align="right">12.04x</td><td align="right">224.3</td><td align="right">26.5</td><td align="right">8.46x</td></tr>
<tr><td align="right">16384</td><td align="right">1024</td><td align="right">8</td><td align="right">225.4</td><td align="right">40.8</td><td align="right">5.52x</td><td align="right">224.9</td><td align="right">43.5</td><td align="right">5.17x</td></tr>
<tr><td align="right">16384</td><td align="right">4096</td><td align="right">8</td><td align="right">536.9</td><td align="right">89.0</td><td align="right">6.04x</td><td align="right">541.4</td><td align="right">90.6</td><td align="right">5.98x</td></tr>
<tr><td align="right">16384</td><td align="right">8192</td><td align="right">2</td><td align="right">1070.2</td><td align="right">96.2</td><td align="right">11.13x</td><td align="right">1069.4</td><td align="right">100.2</td><td align="right">10.67x</td></tr>
</tbody>
</table>
Radix-select kernel:
<table>
<thead>
<tr>
<th rowspan="2">K range</th>
<th rowspan="2">M</th>
<th rowspan="2">N</th>
<th rowspan="2">K</th>
<th colspan="3">non-deterministic</th>
<th colspan="3">deterministic</th>
</tr>
<tr>
<th>ATen (us)</th>
<th>FlyDSL (us)</th>
<th>speedup</th>
<th>ATen (us)</th>
<th>FlyDSL (us)</th>
<th>speedup</th>
</tr>
</thead>
<tbody>
<tr><td>64-256</td><td align="right">256</td><td align="right">8192</td><td align="right">64</td><td align="right">56.8</td><td align="right">51.2</td><td align="right">1.11x</td><td align="right">57.7</td><td align="right">58.5</td><td align="right">0.99x</td></tr>
<tr><td></td><td align="right">2048</td><td align="right">8192</td><td align="right">64</td><td align="right">207.6</td><td align="right">118.0</td><td align="right">1.76x</td><td align="right">207.1</td><td align="right">131.3</td><td align="right">1.58x</td></tr>
<tr><td></td><td align="right">4096</td><td align="right">8192</td><td align="right">64</td><td align="right">242.5</td><td align="right">188.0</td><td align="right">1.29x</td><td align="right">242.4</td><td align="right">210.4</td><td align="right">1.15x</td></tr>
<tr><td></td><td align="right">256</td><td align="right">32768</td><td align="right">256</td><td align="right">124.4</td><td align="right">96.5</td><td align="right">1.29x</td><td align="right">124.4</td><td align="right">118.6</td><td align="right">1.05x</td></tr>
<tr><td></td><td align="right">2048</td><td align="right">32768</td><td align="right">256</td><td align="right">725.7</td><td align="right">283.3</td><td align="right">2.56x</td><td align="right">719.3</td><td align="right">343.1</td><td align="right">2.10x</td></tr>
<tr><td></td><td align="right">4096</td><td align="right">32768</td><td align="right">256</td><td align="right">1536.3</td><td align="right">539.7</td><td align="right">2.85x</td><td align="right">1535.1</td><td align="right">672.7</td><td align="right">2.28x</td></tr>
<tr><td>257-383</td><td align="right">256</td><td align="right">16384</td><td align="right">257</td><td align="right">86.6</td><td align="right">60.9</td><td align="right">1.42x</td><td align="right">86.9</td><td align="right">70.4</td><td align="right">1.23x</td></tr>
<tr><td></td><td align="right">2048</td><td align="right">16384</td><td align="right">257</td><td align="right">378.4</td><td align="right">210.5</td><td align="right">1.80x</td><td align="right">378.5</td><td align="right">237.8</td><td align="right">1.59x</td></tr>
<tr><td></td><td align="right">4096</td><td align="right">16384</td><td align="right">257</td><td align="right">536.5</td><td align="right">400.5</td><td align="right">1.34x</td><td align="right">535.1</td><td align="right">451.8</td><td align="right">1.18x</td></tr>
<tr><td></td><td align="right">256</td><td align="right">32768</td><td align="right">350</td><td align="right">125.3</td><td align="right">76.9</td><td align="right">1.63x</td><td align="right">125.1</td><td align="right">91.6</td><td align="right">1.37x</td></tr>
<tr><td></td><td align="right">2048</td><td align="right">32768</td><td align="right">350</td><td align="right">734.3</td><td align="right">304.6</td><td align="right">2.41x</td><td align="right">722.9</td><td align="right">349.2</td><td align="right">2.07x</td></tr>
<tr><td></td><td align="right">4096</td><td align="right">32768</td><td align="right">350</td><td align="right">1538.8</td><td align="right">579.1</td><td align="right">2.66x</td><td align="right">1536.8</td><td align="right">657.5</td><td align="right">2.34x</td></tr>
<tr><td>384-831</td><td align="right">256</td><td align="right">32768</td><td align="right">390</td><td align="right">125.4</td><td align="right">77.8</td><td align="right">1.61x</td><td align="right">125.3</td><td align="right">91.8</td><td align="right">1.36x</td></tr>
<tr><td></td><td align="right">2048</td><td align="right">32768</td><td align="right">390</td><td align="right">722.6</td><td align="right">306.1</td><td align="right">2.36x</td><td align="right">721.8</td><td align="right">350.5</td><td align="right">2.06x</td></tr>
<tr><td></td><td align="right">4096</td><td align="right">32768</td><td align="right">390</td><td align="right">1542.9</td><td align="right">581.0</td><td align="right">2.66x</td><td align="right">1541.4</td><td align="right">660.2</td><td align="right">2.33x</td></tr>
<tr><td></td><td align="right">256</td><td align="right">131072</td><td align="right">768</td><td align="right">315.5</td><td align="right">162.4</td><td align="right">1.94x</td><td align="right">315.3</td><td align="right">193.3</td><td align="right">1.63x</td></tr>
<tr><td></td><td align="right">2048</td><td align="right">131072</td><td align="right">768</td><td align="right">2256.9</td><td align="right">1284.6</td><td align="right">1.76x</td><td align="right">2255.7</td><td align="right">1503.3</td><td align="right">1.50x</td></tr>
<tr><td></td><td align="right">4096</td><td align="right">131072</td><td align="right">768</td><td align="right">4374.7</td><td align="right">2558.1</td><td align="right">1.71x</td><td align="right">4341.5</td><td align="right">2993.0</td><td align="right">1.45x</td></tr>
<tr><td>832-1024</td><td align="right">256</td><td align="right">32768</td><td align="right">833</td><td align="right">126.3</td><td align="right">73.8</td><td align="right">1.71x</td><td align="right">126.2</td><td align="right">85.0</td><td align="right">1.49x</td></tr>
<tr><td></td><td align="right">2048</td><td align="right">32768</td><td align="right">833</td><td align="right">735.3</td><td align="right">565.4</td><td align="right">1.30x</td><td align="right">733.5</td><td align="right">632.3</td><td align="right">1.16x</td></tr>
<tr><td></td><td align="right">4096</td><td align="right">32768</td><td align="right">833</td><td align="right">1550.8</td><td align="right">1129.3</td><td align="right">1.37x</td><td align="right">1550.4</td><td align="right">1256.4</td><td align="right">1.23x</td></tr>
<tr><td></td><td align="right">256</td><td align="right">262144</td><td align="right">1024</td><td align="right">599.6</td><td align="right">280.6</td><td align="right">2.14x</td><td align="right">606.2</td><td align="right">341.6</td><td align="right">1.77x</td></tr>
<tr><td></td><td align="right">2048</td><td align="right">262144</td><td align="right">1024</td><td align="right">3957.5</td><td align="right">2221.5</td><td align="right">1.78x</td><td align="right">3924.0</td><td align="right">2738.2</td><td align="right">1.43x</td></tr>
</tbody>
</table>
Across the sampled shapes (non-deterministic / deterministic):
- Register kernel: **1.80x-12.04x** / **1.25x-10.67x**
- Radix `K=64-256`: **1.11x-2.85x** / **0.99x-2.28x**
- Radix `K=257-383`: **1.34x-2.66x** / **1.18x-2.34x**
- Radix `K=384-831`: **1.61x-2.66x** / **1.36x-2.33x**
- Radix `K=832-1024`: **1.30x-2.14x** / **1.16x-1.77x**
Co-authored-by: lirui927 ruili@amd.com
Co-authored-by: Cursor cursoragent@cursor.com
Pull Request resolved: https://github.com/pytorch/pytorch/pull/193548
Approved by: https://github.com/jgong5, https://github.com/slayton58
Co-authored-by: lirui927 <Rui.Li@amd.com>
|
||
|
lirui927
|
14fb8257d3 |
[Eager] Add FlyDSL native RMSNorm forward override (#191447)
Add a FlyDSL-based fused RMSNorm forward kernel as a native op override, on top
of the backend in the PR below.
### Supported cases
fp16, bf16 and fp32 tensors on gfx950, contiguous input and weight, 1-D `normalized_shape`, and
`weight` of shape `(N,)` with matching dtype and device. N-D inputs are logically flattened to
`(M, N)`. `eps` is passed through at runtime and does not participate in the cache key. The
kernel returns the forward result and the FP32 `rstd` that ATen's backward consumes.
### When FlyDSL acceleration kicks in
gfx950, non-copy-on-write contiguous input and weight, one normalized dimension, non-negative
`eps`, and:
| N | minimum rows M |
|---|---:|
| `4096 <= N < 8192` | 8192 |
| `8192 <= N < 16384` | 4096 |
| `16384 <= N <= 114688` | 2048 |
### Test plan:
- lintrunner on all changed files
- python test/python_native/test_flydsl_rmsnorm_fwd.py
- OPINFO_RESTRICT_TO_DSL=flydsl python test/test_ops.py -k rms_norm
### Performance methodology
Latency is measured with CUDA events after 10 warmup iterations and over 50 timed iterations. The
ATen baseline runs with FlyDSL disabled; the FlyDSL measurement invokes the same operator with the
override enabled. Every row below was confirmed to actually dispatch, by checking the compile
cache call count advanced, and its output and `rstd` were compared against the ATen baseline.
Each row is a single run. Repeating the smallest shape ten times spread its speedup over
1.16x-1.26x.
```python
def bench(fn, iters=50, warmup=10):
for _ in range(warmup):
fn()
torch.cuda.synchronize()
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(iters):
fn()
end.record()
torch.cuda.synchronize()
return start.elapsed_time(end) / iters
fwd = lambda: torch.ops.aten._fused_rms_norm(x, [N], weight, eps)
with pn.flydsl.disabled():
aten_ms = bench(fwd)
flydsl_ms = bench(fwd)
speedup = aten_ms / flydsl_ms
```
### Performance
GPU: AMD 256-CU MI355X (gfx950). FlyDSL 0.3.0.
| dtype | M | N | ATen (us) | FlyDSL (us) | speedup |
|:---|---:|---:|---:|---:|---:|
| fp16 | 8192 | 4096 | 26.7 | 22.9 | 1.16x |
| fp16 | 4096 | 8192 | 30.1 | 22.6 | 1.33x |
| fp16 | 2048 | 16384 | 31.0 | 23.1 | 1.34x |
| fp16 | 2048 | 65536 | 173.2 | 112.3 | 1.54x |
| fp16 | 2048 | 114688 | 303.2 | 258.8 | 1.17x |
| fp16 | 8192 | 4097 | 96.7 | 26.4 | 3.66x |
| fp16 | 4096 | 8193 | 82.7 | 26.0 | 3.18x |
| fp16 | 2048 | 16385 | 91.9 | 25.9 | 3.55x |
| bf16 | 8192 | 4096 | 26.7 | 22.2 | 1.20x |
| bf16 | 2048 | 16384 | 30.9 | 22.4 | 1.38x |
| bf16 | 8192 | 4097 | 96.4 | 26.4 | 3.65x |
| bf16 | 2048 | 16385 | 73.7 | 25.8 | 2.86x |
| fp32 | 8192 | 4096 | 56.3 | 40.1 | 1.41x |
| fp32 | 4096 | 8192 | 58.5 | 40.7 | 1.44x |
| fp32 | 2048 | 16384 | 57.7 | 44.1 | 1.31x |
| fp32 | 2048 | 32768 | 154.1 | 105.5 | 1.46x |
| fp32 | 2048 | 65536 | 321.8 | 210.3 | 1.53x |
| fp32 | 2048 | 114688 | 564.7 | 483.8 | 1.17x |
| fp32 | 8192 | 4097 | 91.8 | 41.6 | 2.20x |
| fp32 | 4096 | 8193 | 82.5 | 42.3 | 1.95x |
| fp32 | 2048 | 16385 | 78.0 | 43.9 | 1.78x |
| fp32 | 2048 | 32769 | 212.7 | 105.0 | 2.03x |
Aligned shapes improve by 1.2x to 1.5x, with both ends of the band -- `N` = 4096 and
`N` = 114688 -- sitting near 1.2x. Off-by-one `N` improves by up to 3.7x.
Unsupported shapes retain ATen behavior.
Co-authored-by: lirui927 <ruili@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/191447
Approved by: https://github.com/jgong5, https://github.com/slayton58
Co-authored-by: lirui927 <Rui.Li@amd.com>
|
||
|
Nick Riasanovsky
|
bdf0f03732 |
Avoid sum router overhead when inner-tree reduction is off (#193858) (#193858)
Summary: Replace the per-call `os.getenv` rollout gate with a PyTorch config and only register the sum/prod CuTeDSL dispatcher overrides when the feature is enabled at import time. The default-disabled production path therefore has no Python router or environment lookup on each reduction call. Update the dedicated tests and strict-numerics benchmark to enable the rollout before importing `torch`, and use the CuTeDSL disable context when an ATen baseline is needed. Follow-up to D115338972. Test Plan: `arc f` on all changed files `arc lint` on all changed files `buck2 build fbcode//caffe2:gen_zipped_modules` Verified all fbcode/xplat mirror pairs are byte-identical Reviewed By: karthickai Differential Revision: D116334589 Pull Request resolved: https://github.com/pytorch/pytorch/pull/193858 Approved by: https://github.com/karthickai |
||
|
drisspg
|
411c2b5856 |
ci: align FlashAttention and CuTeDSL dependencies (#193646)
## Human Note
split up big ghstack rebase pr into the smaller chunk that rebases quack so we can get ci in better spot
## Agent note
FlashAttention 4 b17 leaves its QuACK and CuTeDSL dependencies open-ended. The B200 setup then
installed a separate QuACK revision and later pinned CuTeDSL independently, so the final environment
depended on resolver and installation order.
Keep the existing FlashAttention version and pin released QuACK 0.6.4, whose package metadata selects
CuTeDSL 4.6.2. Install that set in one resolver transaction, align the standalone CuTeDSL installer and
native version gate, and run the CuTeDSL smoke test after CUTLASS API installation so it checks the final
runtime. Keeping b17 avoids bundling an unrelated FlashAttention runtime upgrade into this dependency
fix.
This change was prepared with assistance from an AI coding tool.
## Test Plan
The B200 GPU smoke suite was not run locally.
```bash
bash -n .ci/pytorch/common_utils.sh
bash -n .ci/pytorch/test.sh
python -m py_compile torch/_native/cutedsl_utils.py test/python_native/test_cutedsl_smoketest.py
git diff --check
lintrunner -a
bash <<'SH'
set -euo pipefail
source .ci/pytorch/common_utils.sh
pip_install() { printf "pip_install %s\n" "$*"; }
unset DESIRED_CUDA CUDA_VERSION BUILD_ENVIRONMENT
cpu=$(install_flash_attn_cute)
[[ "$cpu" == *"flash-attn-4==4.0.0b17 quack-kernels==0.6.4 apache-tvm-ffi==0.1.11"* ]]
DESIRED_CUDA=13.2
cuda13=$(install_flash_attn_cute)
[[ "$cuda13" == *"flash-attn-4[cu13]==4.0.0b17 quack-kernels==0.6.4 apache-tvm-ffi==0.1.11"* ]]
SH
python - <<'PY'
import json
import urllib.request
from packaging.requirements import Requirement
from packaging.version import Version
selected = {
"nvidia-cutlass-dsl": Version("4.6.2"),
"quack-kernels": Version("0.6.4"),
"apache-tvm-ffi": Version("0.1.11"),
}
expected = {
"flash-attn-4": {"nvidia-cutlass-dsl", "quack-kernels", "apache-tvm-ffi"},
"quack-kernels": {"nvidia-cutlass-dsl", "apache-tvm-ffi"},
}
for package, version in (("flash-attn-4", "4.0.0b17"), ("quack-kernels", "0.6.4")):
with urllib.request.urlopen(f"https://pypi.org/pypi/{package}/{version}/json") as response:
requirements = json.load(response)["info"]["requires_dist"] or []
checked = set()
for raw_requirement in requirements:
requirement = Requirement(raw_requirement)
if requirement.name not in selected:
continue
if requirement.marker is not None and not requirement.marker.evaluate({"extra": "cu13"}):
continue
assert selected[requirement.name] in requirement.specifier, raw_requirement
checked.add(requirement.name)
assert expected[package] <= checked, (package, expected[package] - checked)
PY
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/193646
Approved by: https://github.com/malfet, https://github.com/slayton58
ghstack dependencies: #192661
|
||
|
Simon Layton
|
f6a3fb82a0 |
Fix native eager router recursion under a folded Dynamo flag (#191507)
`eager_router` diverts to the plain aten overload when `torch.compiler.is_dynamo_compiling()` is true, so a Dynamo trace records the aten op instead of the router's eager-only predicates and boxed fallback (#186354). That predicate cannot tell tracing from executing. Its Python body is `return False`; Dynamo constant-folds the call to `True` at trace time (`tracing_state_functions` in `_dynamo/variables/torch.py`). A frame carrying the folded constant can still run eagerly, and then `_aten_overload(...)` re-enters the dispatcher from the top, lands back in the same router, and recurses: RecursionError: maximum recursion depth exceeded test_ops.py::TestCommonCUDA::test_out_warning_scatter_add_cuda (PYTORCH_TEST_WITH_INDUCTOR=1, native scatter_add override installed) A thread-local re-entrancy guard breaks the cycle. The outer call still takes the shortcut, so trace-time behavior is unchanged; a re-entrant call falls through to eager dispatch. Rejected alternative: keying the branch on fake inputs instead. Measured, it reintroduces the graph breaks #186388 removed, because `is_fake` lives in `torch/_subclasses`, which is in Dynamo's MOD_SKIPLIST. Against test_device_context_matmul_avoids_native_bmm_router_graph_break: 5 new breaks with the fake check, 0 with the guard. Also fixes three tests that self-skipped whenever `import cutlass` failed, and so run for the first time on an image that installs nvidia-cutlass-dsl, on hardware they were not written for: - test_cutedsl_add_e2e / _autotune: ROCm images match the wheel's x86_64-Linux marker but have no NVIDIA driver, so cuInit fails with cudaErrorInsufficientDriver. Skip on torch.version.hip; the rest of the class is mock-based. - test_precompile_serializes_rmsnorm_tuned_tensor_kwargs: the quack autotuner imports triton (absent on CPU images) before the test's own CUDA skip is reached. - TestScatterAddOverrideConds: its accepts-cases assert the dispatch predicates fire, but both kernels gate on sm_90+, so on pre-sm_90 CUDA (trunk's L4 runners) they are uniformly False and the class tests nothing. Gate on SM90OrLater. TestScatterAddOverrideCorrectness stays ungated, since correctness-on-fallback is what pre-sm_90 should verify. Authored with an AI assistant (Claude Code). Test Plan: ``` python test/python_native/test_registry.py TestRegistryRuntime python test/dynamo/test_repros.py -k native_bmm_router PYTORCH_TEST_WITH_INDUCTOR=1 python test/test_ops.py \ TestCommonCUDA.test_out_warning_scatter_add_cuda lintrunner --take PYFMT,RUFF,FLAKE8 <changed files> ``` The new regression test fails with RecursionError on unmodified main and passes with the guard. Its `cond` also asserts the guard is held while it runs, so the test fails if the router reaches `cond` without going through the shortcut at all -- checked by forcing the shortcut's COW clause false. TestRegistryRuntime passes in full (23 tests); the six failures reported here earlier were a stale local build and no longer reproduce. Pull Request resolved: https://github.com/pytorch/pytorch/pull/191507 Approved by: https://github.com/drisspg |
||
|
lirui927
|
35eadc3dfd |
Add FlyDSL as a backend for native op overrides (#191446)
Add FlyDSL as a native op backend in torch/_native, following the existing CuTeDSL/Triton/Helion integration pattern. This introduces the runtime availability gate, a JIT compile cache, and an instrumentation entry point. It does not register any operator implementation -- the RMSNorm override is the PR above #191447. Requires FlyDSL 0.3.0; older or missing installs report unavailable and everything falls back to ATen. Test plan: - lintrunner on all changed files - python test/python_native/test_flydsl_cache.py - python test/python_native/test_dsl_registry.py - python test/python_native/test_flydsl_utils.py - python test/python_native/test_instrumentation.py - python test/python_native/test_native_dsl_ops.py Co-authored-by: lirui927 <ruili@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/191446 Approved by: https://github.com/slayton58, https://github.com/jgong5 Co-authored-by: lirui927 <Rui.Li@amd.com> |
||
|
Trevor Mathisen
|
e8ad0c997f |
topk: restrict CuTeDSL override to SM100 and newer (#192647) (#192647) (#192647) (#192647)
Summary: The CuTeDSL `topk` thresholds were tuned on B200, but `_eligible()` admitted older CUDA architectures where `aten` is faster. # This diff: - Require compute capability major 10 or newer before selecting the CuTeDSL override. - Keep A100 and Hopper on `aten::topk`; preserve the optimized path on B200 and newer GPUs. - Pin the dispatch boundary with a CPU-only SM80/SM90/SM100 eligibility test. Test Plan: **Validation** - `test/python_native/test_topk_cutedsl_eligibility.py`: 2 passed under the internal Python test runner, covering SM80 and SM90 fallback plus SM100 eligibility. - `test/test_sort_and_select.py -k topk`: 31 passed under the internal Python test runner. - GitHub CI will run OSS public-bindings and lintrunner coverage. - Tested internal flows and runtime reduced 6.5x on A100 Pulled By: Trolann Trolann Reviewed By: ezrakhuzadi Differential Revision: D115327835 Pull Request resolved: https://github.com/pytorch/pytorch/pull/192647 Approved by: https://github.com/slayton58 |
||
|
karthickai
|
6eb6744c0e |
[Eager] Add a CuTeDSL inner-tree prod reduction override (#192882)
Extends the inner-tree reduction engine (shared `reductions` module) to
`aten::prod`, so `x.prod(dim=...)` is bit-for-bit deterministic in the same
inner-tree order as `sum` when `PYTORCH_SUM_INNER_TREE` is set. `prod` (and
`nansum`) were the flagged followups to the inner-tree `sum` work.
sum and prod share the entire inner-tree engine, tiling plan, and eligibility;
only two things differ, threaded through as a small `_Reduction(op, identity)`
descriptor:
- combiner: `+` (sum) vs `*` (prod)
- identity: `0.0` (sum) vs `1.0` (prod) -- pads ragged tails and initializes
accumulators. Prod must pad with 1.0 (`x * 1 == x`); a 0.0 pad would collapse
every row to 0.
`sum` keeps identical behavior via the `_SUM` defaults. The override registers
`prod.dim_int` / `prod.int_out`, reusing the reduction-agnostic eligibility
gate. fp32/fp64 are the bitwise-validated set; fp16/bf16 are functionally
supported (not part of the bitwise contract).
Compile instrumentation now labels each compile per-reduction (`aten::sum` vs
`aten::prod`) via a callable `op` label. Because the label is evaluated in the
instrumentation `finally` block, a failing label callback falls back to
`<op-label-error>` -- it never changes the wrapped compile's result or the
original exception.
Test Plan:
- `python test/python_native/test_sum_cutedsl.py` -- 22/22 pass.
- prod matches ATen (product-appropriate tolerance) across the
multirow/looped/two-kernel paths x fp32/fp64.
- `test_prod_override_engaged` proves inner-tree prod differs bit-for-bit
from the default ATen order (so the override actually ran).
- `test_prod_low_precision_matches_aten` (fp16/bf16) spies on
`inner_tree_prod_into` and asserts it was called -- proving the CuTeDSL
kernel runs rather than silently falling back to ATen.
- `out=` is bitwise-equal to the functional path; strided-outer and
looped-partial-block confirm the 1.0 ragged-tail padding; determinism;
int/complex fall through to ATen.
- Sum regression: `test_bitwise` (27 shapes x 4 dtypes, SHA-pinned) and
`test_cross_warp_order_sensitive_hash_sm100` hashes are unchanged.
- `python test/python_native/test_instrumentation.py` -- 35/35 pass, incl.
callable-`op` label + label-failure-preserves-result/exception tests.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/192882
Approved by: https://github.com/njriasan
ghstack dependencies: #192881
|
||
|
karthickai
|
edb45b1aa6 |
[Eager] Move the CuTeDSL inner-tree sum override into a shared reductions module (#192881)
Pure rename, no behavior change. `torch/_native/ops/sum/` hosts the CuTeDSL inner-tree reduction engine (kernel + tiling plan + eligibility), which is already op-agnostic and will host additional reductions (prod, mean, nansum) next. Rename it to `torch/_native/ops/reductions/` so the module name reflects the family it owns, mirroring the existing `norm/` family module. The four module files move byte-for-byte (git detects pure renames); the only edits are import-path updates in the seven importers: the `ops` registry, the inductor `ir.py` inner-tree-plan import, and five tests. Test Plan: - `python test/python_native/test_sum_cutedsl.py` -- 11/11 pass, including the SHA-pinned `test_bitwise` (27 shapes x 4 dtypes) and `test_cross_warp_order_sensitive_hash_sm100`. The hashes are unchanged, which proves the rename alters no bits. - `import torch` + `torch.sum(x, 1)` (PYTORCH_SUM_INNER_TREE=1) routes through the new `reductions` module; the inductor `inner_tree_plan` import resolves. - No remaining `torch._native.ops.sum` references in the tree. Pull Request resolved: https://github.com/pytorch/pytorch/pull/192881 Approved by: https://github.com/drisspg |
||
|
Nick Riasanovsky
|
0754e4e682 |
[Eager] [Bitwise-Equivalence] Implement an INNER_TREE sum reduction (#182986) (#182986)
Summary: Adds the triton-style inner tree reduction kernel for sum in the caffe2/ATen CUDA backend, including partitioned launch variants, tests, and the output/input offset handling from D104676317 so TensorIterator output strides and strided outer input rows are preserved. Test Plan: arc lint -a fbcode/caffe2/aten/src/ATen/native/cuda/ReduceSumProdKernel.cu xplat/caffe2/aten/src/ATen/native/cuda/ReduceSumProdKernel.cu fbcode/caffe2/test/test_reductions.py xplat/caffe2/test/test_reductions.py buck2 test fbcode//caffe2/test:torch_cuda_cuda -- TestInnerTreeSum Differential Revision: D100024904 Pull Request resolved: https://github.com/pytorch/pytorch/pull/182986 Approved by: https://github.com/karthickai |