Boxing redispatch fast path: 7 项优化 + 功能键正确性修复 - #2
Open
shallitbeso wants to merge 4 commits into
Open
Conversation
Manually merge (not cherry-pick) the 7 boxing glue-layer optimizations
described in the 2026-07-29 MetaX three-route report from private branch
perf/dispatch-tocopy-fastpath onto clean main (torch 2.10 + MACA 3.8.1).
Private branch was torch 2.8 and predated main's kGcu backend, so these
were re-implemented by hand and codegen was re-run under torch 2.10.
Source changes (this commit):
- dispatcher.h: cache resolved fn ptr via atomic; hot path skips the
per-call string build + hash lookup (kGcu branches preserved).
- copy_ops.cc: _to_copy fast path — flagos->(flagos|cuda) boxes and
redispatches to CUDA _to_copy, skipping the intermediate alloc
(guarded with !defined(USE_GCU)).
- register.cc / strided_ops.{cc,h}: register unfold view op (fixes
repeat() corruption / GPTJ rotary crash); hand-register 9 pure
metadata ops (transpose.int, permute, select.int, slice.Tensor,
squeeze, squeeze.dim, unsqueeze, _unsafe_view, detach) to bypass
boxing entirely; drop their old Ascend dispatcher registration.
- device_boxing.h: raw-pointer the boxing guard record
(SmallVector<TensorImpl*,4>) to drop refcount atomics per box/unbox;
add MaterializeForCat single-pass materialize+legacy-empty filter.
- codegen_ops.py: GroupNorm non-contiguous input contiguous() before
the boxing guard (fixes Diffusers UNet2DModel); MANUAL_REGISTERED_OPS
for the 9 metadata ops; cat path uses MaterializeForCat.
- backends_cuda.conf: remove the 9 metadata ops from the op-list source
so codegen no longer emits them.
Tests: 4 new dispatch integration tests (group_norm, device_boxing_guard,
metadata_view, cat_materialization) + codegen unit test cases; 07-29
convention perf benches (bench_infer/train_0729.py, _maca_cat_patch.py).
Codegen products (csrc/aten/generated/*, backends_*flaggems*.conf) are
intentionally left in the working tree, not committed. Perf reports and
build artifacts are gitignored.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…mizations c5e38a6 ported 7 boxing-layer optimizations by editing the codegen source (scripts/codegen_ops.py) and backends_cuda.conf, but left the regenerated codegen products uncommitted in the working tree. This commit checks in those products so csrc/aten/generated/ and the flaggems conf files match the source. Changes (all auto-generated by scripts/codegen_ops.py, FULL CUDA mode): - 9 pure-metadata view ops now registered directly (removed from the generated boxing layer, moved to MANUAL_REGISTERED_OPS + register.cc in c5e38a6): transpose.int, permute, select.int, slice.Tensor, squeeze, squeeze.dim, unsqueeze, _unsafe_view, detach. Their Kernel/Fn/dispatcher entries are dropped from cuda_kernels.cc / ops.h / ops.cc / register.inc, and the corresponding `= cuda` lines removed from the 4 backends_*_flaggems.conf. - cat: use MaterializeForCat() (folds MaterializeToTensorVec + DropLegacyEmptyForCat into one pass). - clamp_.Tensor: drop the redundant optional-Tensor unwrap; box self only. - native_group_norm: force a contiguous input before boxing to CUDA. No behavior change beyond the ported optimizations; products regenerated from the committed codegen source with FLAGOS_CODEGEN_ALL=1. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…keys
The redispatch fast path (at::_ops::<op>::redispatch(DispatchKeySet(CUDA)))
lands directly on the CUDA backend key, bypassing every functionality key
above it. The original gate only covered Autograd/Autocast, silently dropping:
- per-tensor lazy bits: Conjugate / Negative / ZeroTensor
- ambient TLS modes: Functionalize / Python(TLSSnapshot) / Batched / vmap
which produced wrong numerics whenever an input carried a conj/neg/zerotensor
bit or a functionalize/vmap mode was active.
Fix (csrc/aten/device_boxing.h):
- CanBoxingRedispatch() now also rejects the ambient TLS unsafe modes.
- New HasBoxingUnsafeKey(t) detects the per-tensor lazy bits.
- Both boxing guards fold that check while boxing and expose CanRedispatch();
generated kernels gate on guard.CanRedispatch() so only the offending call
falls back to the safe at::<op>(...) path — the hot no_grad decode path
(no lazy bits) still takes the fast path.
Also fix the kill-switch parse (scripts/codegen_ops.py emits it): read
FLAGOS_NO_REDISPATCH as "any value != 0" to match the repo convention
(fallback.cc, caching_device_allocator.cc), so true/on/2 also disable it.
Codegen (scripts/codegen_ops.py) regenerated; the only diff vs the previous
generated cuda_kernels.cc is the gate expression (verified by isolated diff).
Verified: incremental build + link clean; op suite 266 passed / 0 failed;
codegen unit tests 5 passed; conj/neg/plain results bit-identical with the
fast path ON vs OFF (max|ON-OFF| = 0.0); Qwen3 decode e2e output coherent.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
This PR ports a set of boxing-layer performance optimizations onto main, regenerates the CUDA boxing codegen outputs, and fixes correctness gaps in the redispatch fast path by guarding against dropped functionality keys and per-tensor lazy bits. It also moves several pure-metadata view ops to hand-written registrations and adds targeted unit/integration tests to validate the new codegen and metadata/view semantics.
Changes:
- Add a guarded
::redispatch(DispatchKeySet(CUDA))fast path to generated boxing kernels (with SymInt argument widening), plus runtime gating against grad/autocast/TLS modes and per-tensor lazy bits. - Convert several view/metadata-only ops (e.g.,
permute,transpose.int,select.int,slice.Tensor,squeeze*,unsqueeze,_unsafe_view,detach,unfold) to hand-written registrations and exclude them from generated backend configs / generated registration code. - Add unit + integration tests covering codegen special-cases, metadata view aliasing, group_norm layout handling, DeviceBoxingGuard restoration, and cat TensorList materialization semantics.
Reviewed changes
Copilot reviewed 19 out of 23 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| torch_fl/configs/backends_metax_flaggems.conf | Remove pure-metadata ops from generated backend routing list. |
| torch_fl/configs/backends_flaggems.conf | Same exclusion of manual metadata ops from generated backend routing. |
| torch_fl/configs/backends_flaggems_cpp.conf | Same exclusion for C++/flaggems config variant. |
| torch_fl/configs/backends_dcu_flaggems.conf | Same exclusion for DCU flaggems backend config. |
| torch_fl/configs/backends_cuda.conf | Same exclusion for CUDA backend config. |
| scripts/codegen_ops.py | Implement redispatch fast path generation with gating, SymInt widening, TensorList handling, and contiguous-arg specialization. |
| tests/unit/test_codegen_ops.py | Add unit tests for manual-op exclusion and contiguous-arg codegen validation. |
| tests/integration/ops/test_metadata_view_dispatch.py | Validate correctness + aliasing for hand-written metadata view ops. |
| tests/integration/ops/test_group_norm_dispatch.py | Validate group_norm correctness across layouts and backward. |
| tests/integration/ops/test_device_boxing_guard.py | Validate boxing guard restoration and edge cases. |
| tests/integration/ops/test_cat_materialization.py | Validate cat/cat.out TensorList materialization and legacy-empty semantics. |
| tests/perf/bench_train_0729.py | Add training benchmark script for torch_fl vs native mode. |
| tests/perf/bench_infer_0729.py | Add inference/generate benchmark script for torch_fl vs native mode. |
| csrc/aten/device_boxing.h | Add redispatch gating helpers, SymInt adapters, unsafe-key tracking, faster boxing guards, and cat materialization helper. |
| csrc/aten/dispatcher.h | Add resolved function-pointer caching for dispatcher hot path. |
| csrc/aten/strided_ops.h | Add unfold / detach declarations for manual metadata path. |
| csrc/aten/strided_ops.cc | Implement unfold and remove old dispatcher registrations for view ops. |
| csrc/aten/register.cc | Register hand-written metadata/view ops (incl. unfold) under PrivateUse1. |
| csrc/aten/copy_ops.cc | Add CUDA redispatch fast path for _to_copy. |
| csrc/aten/generated/register.inc | Remove generated registrations for ops moved to hand-written path. |
| csrc/aten/generated/ops.h | Remove generated dispatcher declarations for ops moved to hand-written path. |
| csrc/aten/generated/ops.cc | Remove generated dispatcher definitions for ops moved to hand-written path. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+259
to
+277
| if (self.is_privateuseone() && | ||
| (device.is_privateuseone() || device.is_cuda())) { | ||
| const bool return_flagos = device.is_privateuseone(); | ||
| const c10::Device cuda_device(c10::DeviceType::CUDA, device.index()); | ||
| DeviceBoxingGuard guard(self); | ||
| auto result = at::_ops::_to_copy::redispatch( | ||
| c10::DispatchKeySet(c10::DispatchKey::CUDA), | ||
| self, | ||
| ::std::optional<c10::ScalarType>(dtype), | ||
| layout_opt, | ||
| ::std::optional<c10::Device>(cuda_device), | ||
| pin_memory_opt, | ||
| non_blocking, | ||
| memory_format_opt); | ||
| if (return_flagos) { | ||
| UnboxToFlagos(result); | ||
| } | ||
| return result; | ||
| } |
… keys The hand-written _to_copy kernel took the CUDA `::redispatch(DispatchKeySet(CUDA))` fast path unconditionally for flagos->(flagos|cuda), unlike the generated boxing kernels which are gated by device_boxing.h. A bare backend redispatch drops the functionality keys above the backend: Autograd/Autocast, the thread-local Functionalize/Python/vmap modes, and the per-tensor lazy Conjugate/Negative/ZeroTensor bits. It also ignored the FLAGOS_NO_REDISPATCH kill switch. Under active GradMode this silently dropped grad_fn recording and broke training's backward pass. Gate the fast path on `!HasBoxingUnsafeKey(self) && CanBoxingRedispatch()` — the same checks every generated boxing kernel applies. When any such key is active the call falls through to the explicit copy path, which routes through full dispatch (and self.clone() for the identity case), preserving grad/mode semantics and honoring the kill switch. Add tests/integration/ops/test_to_copy_redispatch_gate.py covering the plain dtype-cast fast path (bit-identical to CPU), grad-mode fallback (grad_fn built + backward correct), and the kill-switch path. Verified: op suite 266 passed / 0 failed (pure boxing, backends_cuda.conf); new gate tests 3 passed; dtype-cast max|diff|=0.0 vs CPU reference. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Author
|
Good catch — fixed in 56adc4c. The hand-written if (self.is_privateuseone() &&
(device.is_privateuseone() || device.is_cuda()) &&
!HasBoxingUnsafeKey(self) && CanBoxingRedispatch()) {When grad/autocast, a thread-local Functionalize/Python/vmap mode, a per-tensor Conjugate/Negative/ZeroTensor bit, or the Added Verified on MetaX (pure boxing, torch 2.10 + MACA 3.8.1): op suite 266 passed / 0 failed, new gate tests 3 passed, dtype-cast max|diff| = 0.0 vs CPU reference. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
概述
把私有分支上的 7 个 boxing 层优化移植到干净的
main,重新生成 CUDA boxing codegen 产物,并修复 redispatch 快路径在特定功能键(functionality key)下静默丢弃语义的正确性问题。变更内容(3 个 commit)
perf: port 7 boxing-layer optimizations from private branch onto main将私有性能分支上的 7 项 boxing 层优化手工移植到干净
main,不携带任何本地临时补丁。核心是 redispatch 快路径:在推理(no_grad)热路径上,boxing kernel 通过at::_ops::<op>::redispatch(DispatchKeySet(CUDA), ...)直接命中 CUDA 后端 kernel,跳过内层 PrivateUse1→native 的二次 dispatch 与 trampoline,结果 bit-for-bit 一致。perf: regenerate CUDA boxing codegen products用
scripts/codegen_ops.py重新生成csrc/aten/generated/*,让签入的 codegen 产物与优化后的模板一致。fix: guard boxing redispatch fast path against dropped functionality keys修复快路径的正确性缺口。
::redispatch(DispatchKeySet(CUDA))只携带 CUDA 后端键,会静默跳过位于后端键之上的功能键。现按两种到达方式分别拦截,任一命中就回退到安全的at::<op>(...):Functionalize/Python/PythonTLSSnapshot(torch_dispatch)/Batched/BatchedNestedTensor/FuncTorchBatched(vmap)——在CanBoxingRedispatch()检查tls_local_dispatch_key_set().included_。Conjugate/Negative/ZeroTensor(以及防御性的逐张量 Functionalize/Batched)——由HasBoxingUnsafeKey()检查每个 boxed 输入的 keyset,boxing guard 折叠进CanRedispatch(),使单个输入上的惰性位只降级该次调用。FLAGOS_NO_REDISPATCHkill switch 的解析(!= "0"),对齐仓库既有 env-flag 约定(fallback.cc、caching_device_allocator.cc)。验证
FLAGOS_NO_REDISPATCH)对 conj-bit / neg-view / 普通算子结果max|ON−OFF| = 0.0。明确未包含
[TEMP-DO-NOT-PUSH]MetaX cat 绕过补丁(本地临时,不上传)。.gitignore改动(按要求排除)。tests/perf/ab_redispatch_infer.py、torch_fl/lib_maca软链等本地产物。