perf+fix: _to_copy/boxing hot-path opts + unfold view op & group-norm layout fixes - #51
Merged
zhaoyinglia merged 4 commits intoAug 5, 2026
Merged
Conversation
shallitbeso
changed the title
shallitbeso
force-pushed
the
perf/dispatch-copy-boxing
branch
2 times, most recently
from
August 4, 2026 03:06
81239f9 to
6cbbfca
Compare
_to_copy is one of the hottest ops in Qwen decode (FP32<->FP16 casts in RMSNorm and attention). The flagos(PrivateUse1)->CUDA branch allocated an intermediate contiguous tensor, memcpy'd device-to-device, then ran a separate .to(dtype) cast -- an extra allocation and an extra pass. Since flagos and CUDA share the same GPU memory, box self to CUDA and redispatch straight to the native CUDA _to_copy with an explicit CUDA DispatchKeySet (at::_ops::_to_copy::redispatch). That skips re-entering the dispatcher from the top (no chance of routing back to PrivateUse1) and lets the native kernel read self's strides + cast dtype in one on-device pass, allocating the result directly on CUDA. self is unboxed back to PrivateUse1 on DeviceBoxingGuard teardown, matching the boxing pattern used by the generated CUDA kernels (e.g. PrivToCopyOutKernelCuda). Guarded so only platforms with a CUDA runtime take the fast path; USE_ASCEND/TSINGMICRO/GCU/MUSA keep the original contiguous + Memcpy path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
DeviceBoxingGuard and TensorListBoxingGuard recorded each boxed tensor in a std::vector<at::Tensor>, so every box/unbox paid an intrusive refcount atomic (fetch_add on push, fetch_sub on teardown) per tensor. Qwen decode boxes/unboxes tens of tensors per token, so this stacks up on the hot path. Record raw c10::TensorImpl* instead and unbox by calling the new SetTensorImplDevice directly (extracted from SetTensorDevice), skipping the owning-Tensor round trip entirely -- zero refcount atomics on box or unbox. - SmallVector<TensorImpl*, 4>: the common case (<=4 tensor operands) stays on the stack, no heap allocation. - The guard now takes a forwarding pack with a static_assert that every argument is an lvalue reference, rejecting temporary (rvalue) tensor handles at compile time so a recorded impl can never dangle. - CPU-scalar / genuine-CUDA / undefined-tensor skipping and exception-safe restore semantics are unchanged; unboxing still routes through PyTorch's _change_backend_component_keys rather than poking dispatch-key bits directly. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Tensor.repeat() calls aten::unfold internally to build a strided view before copy_-ing into it. flagos had no PrivateUse1 impl for the plain unfold view op (only unfold_backward / unfold_copy.out existed). Because view ops cannot fall back to CPU -- storage can't be shared across devices -- PyTorch emitted a warning and returned an uninitialized shell tensor. repeat() then read garbage. GPTJ's rotary embedding uses repeat() to build gather indices; corrupted indices led to out-of-bounds gather -> illegal memory access -> crash. unfold is pure stride computation (no GPU kernel): delegate to at::native::unfold, which computes the new size/stride and calls as_strided (already registered for flagos). Registered via direct m.impl in the always-on block so it applies to every backend. Verified on MetaX (t210-box): unfold/repeat now numerically correct, the view-fallback warning is gone, and a tiny GPTJ forward runs cleanly (cosine 1.0 vs CPU, rel err 0.05% -- fp32 drift, not corruption). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Diffusers UNet2DModel feeds GroupNorm channels-last (NHWC) or strided (strides>1) inputs. The backend native_group_norm CUDA kernel requires standard-contiguous layout (is_contiguous() == true); the native CUDA path normalizes the layout before the kernel, but the flagos boxing wrapper passed `input` straight through, raising: Expected X.is_contiguous(memory_format) to be true Add a generic _CONTIGUOUS_TENSOR_ARGS_BY_OP map in codegen naming, per op, which Tensor args must be forced contiguous before boxing. The generated kernel now emits: at::Tensor input_contiguous = input.is_contiguous() ? input : input.contiguous(); DeviceBoxingGuard guard(input_contiguous, weight_t, bias_t); auto result = at::native_group_norm(input_contiguous, weight, bias, ...); Design points: - The contiguous() must run before DeviceBoxingGuard: boxing rewrites the tensor's device metadata, so the copy has to happen while it still lives on flagos. - is_contiguous() short-circuits to a no-op when already contiguous, so the only cost is on the exact non-contiguous inputs that would crash. - Wired into gen_functional_pure / gen_out_variant / gen_tuple_return but currently applied only to native_group_norm.input -- precise, scoped. Verified on MetaX (t210-box): channels_last / strided GroupNorm forward + backward now numerically match CPU (max diff ~2e-7), a Diffusers-style channels_last Conv->GroupNorm->SiLU ResnetBlock runs cleanly (cosine 0.99999988), and sort/layernorm/topk (same generators) are unaffected. codegen unit checks pass, including the contiguous-before-guard ordering. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
shallitbeso
force-pushed
the
perf/dispatch-copy-boxing
branch
from
August 5, 2026 04:55
6cbbfca to
f6aca33
Compare
shallitbeso
changed the title
lvyufeng
approved these changes
Aug 5, 2026
zhaoyinglia
approved these changes
Aug 5, 2026
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.
概述
一组 PrivateUse1（flagos）后端改动：两项热路径性能优化 + 两项正确性修复，每项一个独立 commit。
性能优化
1.
perf(copy):_to_copy在 CUDA/MetaX 上走直接 CUDA redispatch 快速路径flagos(PrivateUse1) 与 CUDA 共享同一块 GPU 显存。flagos→CUDA 分支原本要分配中间 contiguous tensor + 手动 Memcpy。改为 box
self到 CUDA 后，用at::_ops::_to_copy::redispatch(DispatchKeySet(CUDA), ...)直接重派发到原生 CUDA kernel，一次 on-device pass 完成读 stride + dtype 转换并直接在 CUDA 上分配结果，省掉中间拷贝。仅 CUDA-runtime 平台走快速路径；USE_ASCEND / TSINGMICRO / GCU / MUSA 保持原有 contiguous + Memcpy 路径。2.
perf(boxing): boxing guard 用裸 TensorImpl* 记录已 box 的 tensorDeviceBoxingGuard/TensorListBoxingGuard记录的 boxed tensor 从std::vector<at::Tensor>改为c10::SmallVector<c10::TensorImpl*, 4>，消除每次 box/unbox 的 intrusive refcount 原子操作。新增static_assert+ 转发引用约束，在编译期拒绝对临时（rvalue）tensor 构造 guard，避免悬垂指针。unbox 直接调用SetTensorImplDevice。正确性修复
3.
fix(flagos): 注册 unfold 视图算子，修复 repeat 数据损坏Tensor.repeat()内部调aten::unfold构造 strided 视图再copy_进去。flagos 之前只注册了unfold_backward/unfold_copy.out，没有注册普通的unfold视图算子。视图算子不能 fallback 到 CPU（存储无法跨设备共享），于是 PyTorch 只打个 warning 就返回一个内容未初始化的空壳 tensor，repeat()读到垃圾值。GPTJ 的 rotary embedding 用repeat()生成 gather 索引，索引损坏 → 越界 gather → 非法访存崩溃。unfold是纯 stride 计算，无 GPU kernel：委托at::native::unfold（计算新 size/stride 后调已注册的as_strided）。通过 always-on 的m.impl直接注册，对所有后端生效。改动文件：csrc/aten/register.cc、csrc/aten/strided_ops.cc、csrc/aten/strided_ops.h。4.
fix(flagos): GroupNorm 输入在 boxing 前整理为连续布局Diffusers UNet2DModel 会给 GroupNorm 喂 channels-last (NHWC) 或 strided（strides>1）输入。后端
native_group_normCUDA kernel 要求标准连续布局（is_contiguous() == true），原生 CUDA 路径会在调 kernel 前自动整理，但 flagos boxing wrapper 漏了这步，直接报Expected X.is_contiguous(memory_format) to be true。在 codegen 里加了通用机制
_CONTIGUOUS_TENSOR_ARGS_BY_OP，指定哪些算子的哪些参数需在 boxing 前整理为连续。布局整理必须在DeviceBoxingGuard之前完成（boxing 会改写 device 元数据）；is_contiguous()为 true 时零开销拷贝，只在实际非连续输入才触发contiguous()。目前仅应用于native_group_norm.input，影响范围精确可控。改动文件：scripts/codegen_ops.py、csrc/aten/generated/cuda_kernels.cc。验证
在容器
torch-fl-zhanglh-fulltest（MACA 3.8.1.3 + torch 2.10）以 boxing 模式（ACCELERATOR=metax FLAGOS_METAX_BOXING=1）完整构建通过，import torch_fl冒烟通过。两处性能改动均被copy_ops.cc编译单元传递包含并编译链接成功。两项正确性修复均在同一环境（
t210-box）复现原始故障并验证修复：unfold/repeat/2D unfold 数值正确；小 GPTJ 前向从崩溃变为跑通（cosine 1.0 vs CPU，相对误差 0.05%，纯 fp32 漂移）。Conv→GroupNorm→SiLUResnetBlock 跑通（cosine 0.99999988）；sort/layernorm/topk（共用 codegen generator）无回归；codegen 单元检查全过（含 contiguous-before-guard 顺序不变量）。未删除任何已有算子或接口。
🤖 Generated with Claude Code