feat(ascend): Qwen3-0.6B training at 0.82x torch_npu on real 910 - #53
Merged
Conversation
Port the codegen approach from 2.13 onto this branch: generate all 71 CUDA boxing ops from native_functions.yaml via torchgen, replacing hand-written per-op .cu/.cc + structured_ops. - Bring over scripts/codegen_ops.py + generated/, device_boxing.h, dispatcher.h, refactored register.cc, external-libtorch scripts/docs, skill - Delete 123 hand-written kernel files superseded by codegen - Fix test_dispatch_log_bmm_out_flagos_default marker (cuda -> flaggems) Generated products are byte-identical to 2.13 (same native_functions.yaml schema for these 71 ops; ARRAYREF_OPS needs no change). codegen validated on torch 2.11 torchgen. End-to-end CPU-only + external libtorch_cuda.so verification deferred until download.pytorch.org (the +cpu wheel source) is reachable again. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ration Expand codegen from the hand-listed 71-op conf to the full CUDA dispatch set with automatic filtering. backends_cuda.conf now routes 1824 ops to the boxing CUDA kernel (zero hand-written kernels). Enumeration predicate is a strict superset of the original 71-op conf (0 missing): direct CUDA kernel union structured_delegate-with-CUDA-target union composite_explicit_autograd. Ops the templates cannot express are auto-skipped and fall back to cpu_fallback, so coverage only grows. Key changes: - codegen_ops.py: FLAGOS_CODEGEN_ALL full-enumeration mode; authoritative torchgen predicates (part_of_structured_group -> IListRef vs ArrayRef, use_const_ref_for_mutable_tensors -> mutable out param signature), root_name for ATen/ops headers, leading/trailing-underscore name disambiguation, compute-factory CUDA-device redirect (fixes randn 0-dim garbage), try/except auto-skip in all mode. - codegen_skip_ops.txt: ~211 ops the templates cannot express (multi-out, exotic signatures, dunder shifts, const-ref out variants). - test_full_cuda_coverage.py: 46 sampling tests across unary/binary/ reduction/shape/factory/foreach with CPU cross-check + dispatch routing assertions + randn 0-dim regression guard. - Fix silu_backward / nll_loss_backward tests to build inputs on CPU then move to device, instead of relying on cross-device same-seed randn (which no longer matches now that randn correctly uses the CUDA RNG). Full op suite: 311 passed, 0 failed, 64 skipped, 3 xpassed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The flagos (PrivateUse1) CachingDeviceAllocator maintained its own block pool, so compute-factory ops (torch.randn/mm/etc.) — generated as CUDA boxing kernels that allocate straight on the CUDA device — bypassed it entirely. flagos memory_stats saw only `at::empty` allocations (~3% of the real footprint) and empty_cache could not release the boxed-kernel memory. Since flagos and CUDA share the same physical GPU memory (boxing only relabels the device, no copy), route ALL flagos allocation through the same c10::cuda::CUDACachingAllocator that boxed kernels already use. Now empty + boxed outputs live in one pool: memory stats reflect the true footprint, empty_cache works, and OOM-retry is unified. - DeviceMemoryInterface: add opt-in caching-delegation API (provides_caching / caching_alloc / caching_free / caching_empty_cache / caching_record_stream / caching_get_stats / caching_reset_peak_stats), default off so Ascend and other backends keep the self-built block pool. - CudaDeviceMemory: implement the API over CUDACachingAllocator (raw_alloc / raw_delete / emptyCache / getDeviceStats / resetPeakStats), with a lazy once_flag init(cudaGetDeviceCount()) — no Python lazy-init. - CachingDeviceAllocator: allocate / empty_cache / record_stream / get_stats / reset_stats branch to the backend when provides_caching(); add delegated_deleter for the delegation path. - Extract AllocatorStats into allocator_stats.h to break the include cycle between device_memory_interface.h and caching_device_allocator.h. - CMake: define C10_CUDA_NO_CMAKE_CONFIGURE_FILE (non-Ascend) since the CPU torch wheel ships c10/cuda headers but not cuda_cmake_macros.h. Still links torch_cpu only; CUDA symbols resolve at runtime from the preloaded libtorch_cuda.so (CPU-torch + external-libtorch scheme unchanged). Verified: torch.randn now tracked (allocated_bytes 0 -> 1048576), empty_cache releases reserved (86MB -> 20MB). Regression green — ops 426 passed, allocator 11 passed (test_allocator.py reverted to torch.randn), Qwen3 infer/train + fallback_trace 9 passed.
Generalize three codegen templates in scripts/codegen_ops.py so operators that previously fell through to cpu_fallback are boxed directly: - gen_out_variant now supports N mutable Tensor& outputs by calling at::<base>_outf(...) in faithful schema order (outs last), matching the generated signature. Recovers ~121 multi-out ops (sort, topk, svd, native_batch_norm.out). - New foreach_out category: materialize + box every TensorList including out, call _outf, void return. Recovers 75 _foreach_*.out optimizer ops. - New vector_return category: split/unbind returning std::vector<Tensor>, unboxed via UnboxTensorVecToFlagos. Remove the multi_out skip in enumerate_all_cuda_ops and the corresponding 80 entries from codegen_skip_ops.txt. Rename local 'result' to '_ret' in the tuple out-variant branch to avoid collision with out-params named 'result'. Fix with_cuda_libtorch.sh to add .libtorch_cuda_assets to LD_LIBRARY_PATH so SVD's lazy dlopen of libtorch_cuda_linalg.so resolves. Generates 2024 ops (up from ~1900 effective). Regression green: ops 311 passed, full CUDA coverage 46, core+Qwen3 infer/train 124 passed.
… TensorList-out gen_inplace now branches on torchgen variants: function-only inplace ops (silu_/gelu_/celu_/leaky_relu_/threshold_/hardtanh_/mish_/hardsigmoid_/ hardswish_/elu_/embedding_renorm_/rrelu_with_noise_) call the free function at::op_(self,...) instead of the nonexistent Tensor method self.op_(). gen_foreach_out now handles a single mutable Tensor& out with return (cat.out/stack.out/_stack.out/block_diag.out/_chunk_cat.out) and boxes every mutable Tensor& (not just the out list), branching return shape on ret_type. Tuple-returning multi-out RNN out-variants (_cudnn_rnn.out/_lstm_mps.out/ miopen_rnn.out) added to skip list — they were latent -Wreturn-type bugs. 2043 ops generated (was 2024). Regression: ops 311 passed, coverage 46, core+Qwen3 124 — no regressions.
The torch-2.11 migration (285f52c) unified per-op headers into generated/ops.h but only adapted the CUDA backend, leaving the Ascend backend on the deleted per-op-header layout so it no longer compiled. - Point all ascend + flagos python_wrapper kernels at generated/ops.h - Rename _softmax / sum.dim_IntList dispatchers to the codegen names (PrivSoftmaxFn/priv_softmax_dispatcher, SumDimIntlistFn/sum_dim_intlist_dispatcher) - Split mm/bmm into functional + out variants over a shared aclnn helper to match the new MmFn/MmOutFn (BmmFn/BmmOutFn) signatures - Implement AscendDeviceMemory (ACL-backed DeviceMemoryInterface) so the flagos block-pool allocator works on Ascend; previously GetCachingAllocator hit TORCH_CHECK(false) and .to("flagos:0") threw before any kernel ran Verified on Ascend 910 (torch 2.11.0+cpu, CANN): build+install+import OK; mm/bmm/add/mul/cos/sin/neg/abs/silu/rsqrt/softmax/sum.dim/mean.dim match CPU.
Add a category-based aclnn codegen for the Ascend backend. Unlike the CUDA codegen (which emits one-line at::op boxing bodies and lets PyTorch marshal everything), aclnn needs per-op knowledge the aten schema does not carry (API name, arg marshaling, output allocation), so generation is driven by category templates + an aten->aclnn mapping table. - scripts/codegen_ascend.py: reuses codegen_ops.py:schema_to_cpp_name so symbol names match the dispatcher decls already in generated/ops.h; only fills the Backend::kAscend slot. Validates each aclnn symbol via nm on libopapi.so before emitting; skips unmapped/missing ops with a warning. - csrc/aten/backends/ascend/generated/ascend_kernels.cc: 8 unary ops (sqrt/exp/tanh/sigmoid/reciprocal/log/floor/ceil). Path auto-excluded from non-ascend builds by the existing CMake glob. - backends_ascend.conf: route the 8 new ops to ascend. - docs: design (ascend_aclnn_codegen.md), NPU plan (ascend_npu_plan.md), route-rejection record (cpu_torch_external_libtorch_npu.md), and the standalone feasibility prototype. Verified on Ascend 910: all 8 ops match CPU reference, max_err <= 1.2e-7.
Extend the category-based aclnn codegen from unary-only to six categories, driven by a CATEGORIES dict (one kernel-body template each) + an OPS dict mapping each op to (category, aclnn-name override): - unary (28): erf/erfc/expm1/log2/log10/log1p/round/trunc/frac/sign/relu/ cosh/sinh/asin/atan/asinh/acosh/atanh/logical_not/bitwise_not + the 8 from P1 - binary (7): div.Tensor/pow.Tensor_Tensor/atan2/maximum/minimum/ bitwise_or.Tensor/bitwise_xor.Tensor - binary_alpha (1): sub.Tensor - binary_cmp (7, bool out): eq/ne/gt/lt/ge.Tensor + logical_and/logical_or - binary_scalar_alpha (2): add.Scalar/sub.Scalar - binary_scalar_cmp (6, bool out): eq/ne/gt/lt/ge/le.Scalar Device-coercion fix in the shared binary prologue: torch.sub(x, 3.0) and similar tensor-op-python-scalar forms lower to aten::<op>.Tensor (not .Scalar) with the scalar packed as a CPU scalar tensor. The prologue now coerces the other operand to self's device (other.to(self.options()) when not already on the flagos device), mirroring the handwritten add.cc; coercing dtype alone left CPU storage to be read as an NPU device address, producing all-nan output. Both operands are expanded+materialized to the broadcast shape since aclnn does not always broadcast. Symbol validation via nm on libopapi.so auto-excludes ops without the aclnn symbol or dispatcher (square/isnan/isfinite). Verified on Ascend 910: all 51 ops match CPU reference (unary max_err <= 4.4e-5, binary <= 4.7e-6, comparisons exact).
Add three reduce-shaped categories to the aclnn codegen, taking the
generated Ascend kernel set from 51 to 55 ops:
reduce_dims amax/amin -- (Tensor, IntArrayRef dim, bool keepdim),
same dtype. Reuses the handwritten sum.cc dim logic:
wrap negative dims, empty list = reduce all, drop (or
set to 1 with keepdim) each reduced dim high-to-low.
reduce_dim_bool any.dim -- single int64 dim, bool out. aclnnAny takes a
dim list, so the dim is wrapped into a one-element vec.
cumsum (Tensor, int64 dim, optional dtype) -- same-shape scan.
Reduce ops are heterogeneous (each aclnn reduce has its own arg layout),
so there is no single "reduce" template; each sub-shape is its own
category. The long tail (max.dim/min.dim tuple return, var/std/norm with
correction/p args, argmax/argmin/prod/logsumexp with no aclnn symbol in
this CANN) is left for later, bespoke handling.
All 4 new ops verified on Ascend 910 vs CPU reference across single-dim,
dim-list, negative-dim, all-reduce, and keepdim variants (8/8 subtests).
Add 11 new categories / 22 ops to the Ascend aclnn generator: - unary_bool (isinf), unary_scalar (leaky_relu/clamp_min/clamp_max/ fmod.Scalar), unary_two_scalar (softplus/threshold), unary_int (tril/triu), unary_dims (flip) - addcmul/addcdiv (3-tensor broadcast + Scalar value) - pow_scalar_tensor (pow.Scalar: Scalar self, Tensor exponent) - reduce_max_dim (max.dim/min.dim: tuple(values, int64 indices)) - cumprod (separate from cumsum: aclnnCumprod takes dim as aclScalar*) - act_backward (tanh_backward/sigmoid_backward), threshold_backward -- first training-oriented backward ops - grow binary (fmod.Tensor/floor_divide) and binary_cmp (logical_xor) All 22 verified on Ascend 910 vs CPU reference. Candidates were probed against both the dispatcher decls in generated/ops.h and the aclnn symbols in libopapi.so before templating; ops with no symbol or bespoke args (var/std/norm, argmax, gelu string_view, matmul family) are left long-tail and auto-skipped by the generator's nm validation.
…91 ops)
Extend the category codegen toward the training/inference main line: +14 ops
across 9 new/grown categories, all verified on Ascend 910 vs CPU reference.
New categories:
- elu (alpha/scale/input_scale)
- loss (mse_loss; reduction None=elementwise, Mean/Sum=scalar)
- cummax_cummin (tuple values+int64 indices, same shape)
- aminmax (tuple min+max, optional dim)
- prod (scalar out, optional dtype)
- gemm_addmm / gemm_baddbmm (beta/alpha + cubeMathType)
- mv / dot
Grown: unary_scalar += celu/softshrink/hardshrink;
unary_two_scalar += hardtanh.
smooth_l1_loss is intentionally left long-tail: its aclnn signature takes a
by-value float beta, and EXEC_ASCEND_CMD marshals args through a fully
variadic function-pointer typedef, which is unsafe for a bare float on
aarch64 (beta arrived as 0 -> pure L1 output). Scalars wrapped as aclScalar*
or int64 are varargs-safe; a raw float is not.
Replace the hardcoded 5-entry FLAGGEMS_PYTHON_MAP with automatic discovery
from flag_gems._FULL_CONFIG (433 ops) plus safety filtering, and generate
per-category kFlagOsPython kernels through a schema-driven generic IValue caller.
Discovery + safety gates (scripts/codegen_ops.py):
- arity gate (hard red line): exclude ops where the flag_gems positional
count differs from the aten arg count (gems silently drops trailing scalar
args like add.Tensor's alpha / addcmul's value → wrong results).
- type gate: every arg type must be in the generic caller's supported set;
ScalarType excluded (IValue stores it as a plain int, indistinguishable).
- categories: functional_pure / inplace / tuple_return / out_variant.
- FLAGGEMS_PYTHON_SKIP holds 8 convergence holdouts: mm.out (required out
kwarg the generic caller can't supply) and 7 ops with an unconditional
device.type=="cuda" assert flagos PrivateUse1 tensors can't satisfy
(maximum, minimum, 5 upsample variants).
Result: 150 functional_pure + 57 inplace + 21 tuple_return + 7 out_variant
= 235 ops routed to flagos_python. backends_flaggems.conf now auto-generated.
Generic caller (python_op_caller.{h,cc}): IValueToPython covers Tensor/int/
double/bool/None/str/Scalar/IntList/DoubleList/BoolList/TensorList;
CallPythonOp_Generic uses BuildPyArgs; add CallPythonOp_GenericTuple for
tuple-returning ops; GetFunc resolves dotted module.func qualnames.
Cleanup: delete stale hand-written flagos wrappers superseded by codegen.
Verification: numerical spot-check 23/23 (err <= 3e-5); flaggems_python
correctness 27 passed; CUDA-direct path 330 passed / 45 skipped / 3 xpassed
(matches baseline, no degradation).
…ally (235 -> 253) The out_variant safety gate assumed flag_gems never accepts the aten `out` tensor, so ops whose gems function signature is (…non_out, out) failed the arity check (npos == non_out+out, not == non_out) and were excluded. Add an out_variant_gemsout category: when gems npos == #non_out + #out args, pass the aten out tensor(s) positionally to gems (which writes into them in place) and return the aten out arg. Distinguished from mm.out, whose gems out is a required keyword-only arg the positional caller can't supply (stays in FLAGGEMS_PYTHON_SKIP). Recovers: atan2/bmm/cosh/div/exp/expm1/fmin/hardsigmoid/i0/log10/logaddexp/ pixel_unshuffle/reflection_pad1d/reflection_pad2d/replication_pad1d/softshrink/ special_i0e/where.self .out variants. Verified: 10/10 numerical spot-checks via dispatcher, flaggems_python 27 passed, CUDA-direct 330 passed / 45 skipped / 3 xpassed (no degradation).
…(98 ops) Add 4 backbone categories (6 ops), all verified on Ascend 910 vs CPU: - gelu / gelu_backward: use aclnnGeluV2 (int64 approximate) + aclnnGeluBackwardV2 (char* approximate). v1 aclnnGelu hardcodes the tanh approximation, but PyTorch's default is approximate="none" (erf form, used by qwen3 et al); V2 selects 0="none"/1="tanh" so both modes match (err 1.8e-07 / 1.2e-07). - _log_softmax: mirrors handwritten softmax.cc (aclnnLogSoftmax(self,dim,out)). - _softmax_backward_data / _log_softmax_backward_data: aclnnSoftmaxBackward / aclnnLogSoftmaxBackward (grad_output, output, dim, grad_input), out dtype = input_dtype.
…-> 257) gems funcs addcdiv/round/scatter/scatter_ have extra positional params with defaults beyond the aten args (out=None, decimals=0, reduce=None). Relax the arity gate to admit npos > ncall when every extra gems param is strictly trailing with a default, guarded against the reordering trap (gems gather inserts out=None mid-signature -> aten sparse_grad would be misrouted into the out slot; correctly excluded).
…ops) Add 5 ops across 5 categories, all verified on Ascend 910 vs CPU: - addmv / addr: gemm-family completion. aclnnAddmv arg order is (self,mat,vec, alpha,beta) -- alpha before beta, opposite of addmm. addr has no cubeMathType. - binary_cross_entropy (+optional weight), _backward, and _with_logits (+optional pos_weight). Optional tensors marshal via value_or(Tensor()) -> AclTensorWrapper nullptr, which aclnn treats as absent. Left out (probed but not shipped): addbmm (hf32 cube accumulation over the batch dim inflates rel-err to ~1e-2 vs ~1e-4 for a single addmm) and native_batch_norm (aclnnBatchNorm returns ACLNN_ERR_INNER_NULLPTR on 4D NCHW input; 2D N,C works). Both deferred to the conv/pool bespoke batch.
Pre-codegen "seed" kernels (abs/cos/add.Tensor/mul.Scalar/where/softmax/sum/
mean etc.) had bodies expressible by codegen categories -- several byte-identical
to the templates. Migrate them into scripts/codegen_ascend.py and delete the
handwritten .cc files, per the rule: anything a codegen category can express
goes through codegen; handwrite only genuine bespoke ops.
- Reuse existing categories: abs/acos/cos/sin/neg/rsqrt/silu (unary),
mul.Tensor/bitwise_and.Tensor (binary), add.Tensor (binary_alpha),
pow.Tensor_Scalar (unary_scalar).
- Add 7 new categories: binary_scalar (mul.Scalar->aclnnMuls, div.Scalar->
aclnnDivs -- headers absent, marshaling from handwritten refs),
act_backward_self (silu_backward: grad+self), where (aclnnSWhere 3-tensor),
softmax_fwd (_softmax, half_to_float), reduce_all (all), reduce_sum_dtype
(sum.dim_IntList), reduce_mean_dtype (mean.dim via aclnnMeanV2).
Kept handwritten (SKIP={le.Tensor,mm,bmm}): le (aclnnLe absent, needs runtime
version probe), mm/bmm (also register out-variants codegen doesn't emit),
factories, TensorList/SymInt ops, embedding, nll_loss.
Handwritten kAscend regs 36->17, codegen 103->122, total unchanged. Verified
all 19 migrated ops vs CPU on Ascend 910 incl. CPU-scalar coercion,
half_to_float, dtype promotion, broadcast -- zero numeric regression.
Add a boxing build mode for the MetaX backend that reuses the generated CUDA boxing kernels (csrc/aten/generated/cuda_kernels.cc, host g++, no mxcc/nvcc) which dispatch PrivateUse1 -> CUDA into maca's libtorch_cuda.so, instead of the hand-written mxcc .cu kernels under backends/metax/. - CMakeLists.txt / setup.py: FLAGOS_METAX_BOXING=1 sets METAX_KERNEL OFF while keeping the MetaX SDK runtime + cu-bridge headers path. - csrc/CMakeLists.txt, torch_fl/csrc/CMakeLists.txt: define USE_MACA=1 for metax, required because maca's torch headers are a hard fork gated on it (C10_WARP_SIZE, Context.h allow_tf32_cudnn hit static_assert(0) otherwise). - Regenerate generated/* + backends_cuda.conf against maca 2.10 torch (2035 ops; +cudnn_convolution_bias_fused[.out], -upstream-only ldexp/ _foreach_powsum/miopen_ctc_loss/_flash_attention_forward.quantized). Verified on 8x MetaX C550: import + factory ops + add/mul/relu/softmax/ sum/exp/mm/bmm/copy all match CPU reference (matmul matches native maca-cuda exactly; ~1e-3 vs CPU is MetaX TF32 precision).
Add convolution (fwd+bwd) and 2D pooling to the aclnn codegen. These are the first ops needing an explicit output-shape formula since aclnn requires the output pre-allocated -- each template carries a small shape helper. Categories: - adaptive_avg_pool2d, avg_pool2d, max_pool2d_with_indices (tuple w/ int64 idx) - convolution (non-transposed), convolution_backward (3-tuple + output_mask) Two infra additions in op_api_common.h: - AclTensorWrapper gains an optional aclFormat param (default ND). avg_pool2d/ adaptive_avg_pool2d/convolution reject ND 4-D tensors (GetWorkspaceSize ret 161002); they need NCHW/NCL/NCDHW per rank. max_pool2d does NOT care -- the requirement is per-aclnn, so the param defaults to ND and is opt-in. - AclBoolArrayWrapper for convolution_backward's output_mask[3]. conv cubeMathType=0 (KEEP_DTYPE); type 1 loses ~2.5e-3 on the cube unit. Verified vs CPU on Ascend 910: conv fwd/bwd err~1e-6 incl stride/padding/ dilation/grouped; pooling exact-to-1e-7 incl ceil_mode and stride defaults; prior migrated ops regression-clean. codegen 122->127 ops, 52 categories.
…> 288) Many gems funcs declare aten's trailing args keyword-only (sum(inp,*,dtype), add(A,B,*,alpha), gelu(self,*,approximate), var(x,dim,*,correction,keepdim)). The positional-only generic caller couldn't pass them, so non-default values were silently dropped -> these were wrongly excluded as arity_short. Add CallPythonOp_GenericKw/GenericKwTuple: forward the trailing aten args by NAME via a PyKwarg vector. dtype (ScalarType) is tagged is_dtype so the caller converts the int payload to a torch.dtype (IValue can't distinguish ScalarType from int); absent optionals set is_none -> Python None. Discovery matches each trailing aten arg to a gems keyword-only param by name (reject on mismatch) and gates types via _FLAGGEMS_KWARG_OK. Recovers 31 ops incl add/sub/addmm/addmv/addr/addcmul (alpha/beta/value), sum/mean/prod/cumsum (dtype), gelu (approximate), var/std/var_mean (correction), isin (invert), sort.stable. Spot-checked all kwarg types numerically (err<=1e-5, non-default alpha/dtype/correction now correctly applied). Blocked: Generator?/ Layout?/Device? args and name-mismatches (multinomial/var.dim/_grouped_mm).
…loop) Complete the CNN training closure on top of conv/pool: add max_pool2d backward, native_batch_norm forward and backward. conv + pool + bn now cover a full CNN forward/backward pass. Categories: - max_pool2d_with_indices_backward -> grad_input - native_batch_norm -> (out, save_mean, save_invstd) - native_batch_norm_backward -> (grad_input, grad_weight, grad_bias) Two per-op quirks (each read from the aclnn header @param notes): - max_pool2d FWD ignores format and emits int64 indices, but the BWD kernel requires NCHW format AND int32 indices -- fwd/bwd are not symmetric. The bwd template casts indices to int32 and tags NCHW. - batch_norm's save_invstd uses a different definition than PyTorch CPU (~0.18 apart), but this does not affect training: the backward consumes the same NPU save_invstd and all three grads match CPU to <=4e-6. running_mean/ var are passed non-const (updated in-place). This overturns the earlier "native_batch_norm left out (561103)" note -- that was marshaling, not a real limitation. Verified vs CPU on Ascend 910: max_pool bwd 1e-7; bn fwd out 2e-7, bn bwd grads <=4e-6, bn eval exact; conv/pool/elementwise regression-clean. codegen 127->130 ops, 55 categories.
Round out the backward coverage so conv, avg/adaptive pool, layer_norm and group_norm all have both forward and backward -- covers the remaining norm/ pool grads for CNN and transformer training. Categories: - avg_pool2d_backward, _adaptive_avg_pool2d_backward -> grad_input - native_layer_norm_backward -> (grad_input, grad_weight, grad_bias) - native_group_norm_backward -> (grad_input, grad_gamma, grad_beta) Notes: - avg_pool2d/adaptive_avg_pool2d backward also require NCHW format, matching their forwards. - norm backwards feed the forward's mean/rstd straight through and are self- consistent (no save_invstd-style mismatch: layer/group norm expose rstd, not invstd). - aclnn names: native_layer_norm_backward maps to aclnnLayerNormBackward (aclnnNativeLayerNormBackward does not exist); group_norm likewise. Verified vs CPU on Ascend 910, all <=1e-6 (avg_pool bwd 3e-8, adaptive exact, ln/gn grads <=1e-6); conv/max_pool/elementwise regression-clean. codegen 130->134 ops, 59 categories.
…3 ops (288 -> 291) The positional caller can't carry a ScalarType (IValue stores it as a plain int), so ops taking dtype/input_dtype as a positional aten arg were excluded by the type gate (type_gate_dtype: 4 ops). Now that the kwarg path tags is_dtype and converts to torch.dtype by name, promote any ScalarType positional arg into a by-name kwarg -- safe when gems accepts it by name (not positional-only) and the ScalarType args form a strict suffix (guard rejects middle-dtype reordering). Recovers _softmax_backward_data, _log_softmax_backward_data (input_dtype), linalg_vector_norm (dtype). Spot-checked: backward err<=2e-7, vector_norm f32 ord=1/2 correct. _safe_softmax -> SKIP (gems device assert rejects PrivateUse1, same class as maximum/minimum). vector_norm dtype=f64 raises gems' own NotImplementedError (gems limitation, not a forwarding bug -- f32 path correct). Regressions green: cuda-direct 330 passed, flaggems_python 27 passed.
Fill the remaining Transformer op gaps. A coverage probe showed embedding, gelu, silu, bmm, baddbmm and softmax were already handled (handwritten or earlier codegen), so the real gaps were masking and indexing. Categories: - masked_fill.Scalar, masked_fill.Tensor -> broadcast(self, mask) - gather -> index shape; index_select -> self shape w/ dim -> index.numel() Note on out-of-place via inplace aclnn: aclnn only ships inplace masked_fill (aclnnInplaceMaskedFillScalar/Tensor). Implementing the out-of-place aten op means copy-then-fill, but self.clone() routes through empty_like which is NOT registered for the ascend backend. Allocate via apply_tensor_without_format + out.copy_(self.expand(...)) instead -- applies to any copy-then-mutate kernel. SDPA/flash-attention is fused with a bespoke multi-tensor signature and is left to a dedicated batch. Verified vs CPU on Ascend 910: all exact (err=0) incl broadcast mask, 3D gather, negative dims, vocab-size index_select; core ops regression-clean. codegen 134->138 ops, 63 categories.
…ace/logspace), +10 (291 -> 301) Factory ops don't take input tensors -- gems generates the tensor itself. New CallPythonOp_Factory injects device=flagos (so gems' internal torch.empty hits OUR allocator -> PrivateUse1 tensor, no CUDA round-trip, no recursion), layout=strided (gems eye/randperm reject layout=None), pin_memory=None, and forwards the aten dtype. discover_flaggems_ops gains a factory branch that strips the TensorOptions fields and passes only shape/scalar positionals; requires gems to accept dtype/layout/device by name (kwonly on every factory). Recovers 10 ops: arange (+.start/.start_step), eye (+.m), full, linspace, logspace, ones, zeros. Spot-checked all exact (err=0) vs CPU incl f64 dtype. Two correctness guards: - arange: gems defaults dtype=None to int64 unconditionally, but aten infers float when any of start/end/step is floating -> arange(0.,3.,.5) was silently wrong ([0,0,1,1,2,2]). Kernel now replicates aten's rule and passes an explicit dtype. Verified arange float now exact. - rand/randn/randperm SKIP: gems reaches default_generators[device], but the PrivateUse1 device has none (IndexError); randperm asserts an int dtype. Same root cause as the Generator? blocked group -- can't express per-device gen. Regressions green: cuda-direct 330 passed, flaggems_python 27 passed.
- Investigated aclnnFlashAttentionScore/Grad API semantics - Confirmed: inputLayout="BNSD", softmaxMax/Sum shape [B,N,S,8] float32 - Confirmed: causal (sparseMode=3, preTokens=INT32_MAX, nextTokens=0) vs full (sparseMode=0, preTokens=65536, nextTokens=65536) - Key blocker: PyTorch logsumexp [B,N,S] vs aclnn (softmaxMax, softmaxSum) [B,N,S,8] shape mismatch; logsumexp=log(Sum)+Max but 8-way tiling collapse unclear - Conclusion: SDPA is bespoke multi-day task (autograd Function wrapper, ctx save/restore, causal/dropout/mask testing), deferred to dedicated pass Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add handwritten kernel for _scaled_dot_product_efficient_attention (forward only). Wraps aclnnFlashAttentionScore with BNSD layout, handles logsumexp mapping from aclnn's softmaxMax+softmaxSum [B,N,S,8] to PyTorch's [B,N,S] via [:,:,:,0] indexing. Key findings: - attenMask semantics: true=MASK_OUT (opposite of docs), false=KEEP - Causal attention: triu(ones, diagonal=1) masks future positions - Verified: non-causal err=3.34e-06, causal err=7.15e-07 vs CPU Backward NOT implemented: aclnnFlashAttentionScoreGrad needs separate softmaxMax and softmaxSum, but PyTorch's autograd only saves single logsumexp (log addition not invertible). Forward-only covers inference; training needs architectural work. Files: - csrc/aten/backends/ascend/scaled_dot_product_attention.cc (new, 120 lines) - csrc/CMakeLists.txt (add to ascend sources) - test_sdpa_ascend.py (verification script) - docs/ascend_aclnn_codegen.md (document SDPA + attenMask/logsumexp pitfalls) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Add _scaled_dot_product_efficient_attention (forward) - Add _scaled_dot_product_efficient_attention_backward (backward) - Wrap aclnnFlashAttentionScore and aclnnFlashAttentionScoreGrad - Implement activation checkpointing for backward (recompute to get softmaxMax/Sum) - Support causal and non-causal attention - Add dropout=0 constraint (aclnn dropout needs explicit mask handling) - Tests pass: forward+backward, causal mask scenarios Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Move the flaggems-vs-cuda choice from compile time + external LD_PRELOAD to a pure runtime env switch, so one wheel serves both paths. - Default FLAGGEMS_PYTHON option ON: compile both CUDA boxing kernels and flaggems_python kernels into libtorch_fl.so unconditionally (MetaX keeps it opt-in). No more FLAGGEMS_PYTHON=ON at build time. - Bundle .libtorch_cuda_assets/*.so* into torch_fl/lib and ctypes-preload them (nvidia deps -> torch cpu libs -> cuda libs) before import torch, to satisfy the CUDAHooks hard constraint. FLAGOS_DISABLE_CUDA_ASSETS skips. - FLAGOS_USE_FLAGGEMS selects backends_flaggems.conf vs backends_cuda.conf at import; FLAGOS_OP_<name> per-op overrides still apply. Both confs and the nvidia-*-cu12 runtime deps (CUDA builds) are now packaged. Verified on a clean shell (no wrapper, auto-preload): cuda-path ops 330 passed/45 skipped/3 xpassed; FLAGOS_USE_FLAGGEMS=1 flaggems_python 27 passed; all three switch levels confirmed.
…> 307) Add two codegen paths to the FlagGems Python bridge: - *_like factory (zeros_like/ones_like/full_like): new CallPythonOp_LikeFactory injects device=flagos/layout=strided/memory_format=None/pin_memory=None and forwards dtype; gems reads shape/device from the source tensor positional. - random in-place (uniform_/exponential_/bernoulli_.float): new CallPythonOp_RandomInplace injects a module-level CUDA torch.Generator as the generator kwarg. gems only reads philox seed+offset from it (randoms computed in the Triton kernel writing into the flagos tensor), so the generator device need not match. A CUDA generator is required: a flagos/CPU generator's 5056-byte MT19937 state fails gems' 16-byte CUDA-philox state unpack. Serialized by the GIL; offset advances per call via gems' set_state. discover_flaggems_ops gains like_factory / random_inplace branches; random ops are whitelisted (_FLAGGEMS_RANDOM_INPLACE) since normal_ is also inplace+Generator? but hardcodes generator=None internally and can't be routed. Explicitly skipped with reasons: normal_/normal.* (gems drops the generator), rand/randn/randperm/rand_like/randn_like (no generator param -> empty PrivateUse1 default_generators), multinomial (name-mismatch + generator). Also add docs/flaggems_no_dispatcher_analysis.md: the 88 no-dispatcher ops all already execute correctly (composite_implicit decomposition / fallback / manual registration), so they are not a functional gap and bulk routing would drop autograd. Validated: like ops err=0 incl dtype override; random ops distribution-correct with distinct streams across calls. Regression unchanged (cuda 330/45/3, flaggems_python 27).
Reuse PyTorch's CUDA boxing kernels on MetaX with a stock torch==X.Y.Z+cpu wheel by symlinking the active wheel's torch/lib core .so to the MetaX torch wheel's copies, so the process loads the MetaX C++ runtime (at::maca::* fork) instead of the upstream one. A runtime hook in torch_fl/__init__.py runs ensure_maca_libtorch_links() BEFORE import torch (afterwards libc10 is already mapped and relinking is too late). Pure ctypes preloading does not work: the official _C.so / libtorch_python.so carry an $ORIGIN RUNPATH that pulls upstream libc10 back in by full path, double-loading it and crashing on duplicate caffe2 static init. Symlinking makes the physical files RUNPATH resolves to be the MetaX ones. Gated on FLAGOS_METAX_BOXING=1; idempotent; backs up originals to torch/lib/_orig_backup/ (restore_original_libtorch() reverts); no-op when the active torch already IS the MetaX wheel. MetaX torch/lib discovered via FLAGOS_MACA_TORCH_LIB env or conda-env scan. Verified on torch 2.10.0+cpu: full_cuda_coverage 45/46 (the 1 fail is addmm TF32 rounding), per-op suite 310 passed / 45 skipped / 3 xpassed (8 fails are all flaggems-backend or test-case issues, not op wiring).
# Conflicts: # torch_fl/__init__.py
Drop cudnn_convolution_bias_fused (2.11-only op absent in 2.10). Regen matches the 2.10 aten schema; flaggems_python and CUDA paths verified against the 2.11 baseline (27 passed; 330 passed/45 skipped/3 xpassed).
Package the MetaX-forked libtorch C++ .so (~1.1G) inside the wheel under torch_fl/lib_maca/ so target machines need only the official torch+cpu wheel plus this wheel plus the /opt/maca driver runtime -- no separate MetaX torch wheel required. - setup.py: package_data bundles lib_maca/*.so* and all backends*.conf (boxing modes select backends_cuda.conf via FLAGOS_BACKEND_CONFIG); version gains a +metax local segment (FLAGOS_WHEEL_LOCAL overridable). - pyproject.toml: mark version dynamic so setup.py's computed +metax tag is not overridden by a static [project] version. - _metax_libtorch_link.py: _discover_maca_torch_lib prefers the bundled lib_maca/ over env var / conda scan; symlinks the stock wheel's torch/lib to the bundled forked libtorch at import. - scripts/bundle_maca_libtorch.sh: copy the 8 forked libtorch .so and patchelf their RPATH to $ORIGIN + /opt/maca for the target runtime. - .gitignore: ignore torch_fl/lib_maca/ (build artifact). Verified end-to-end from a fresh wheel install in a clean env (official torch 2.10.0+cpu, no MetaX torch wheel, no LD_LIBRARY_PATH): flagos compute on MetaX GPU works, torch/lib auto-symlinks to the installed lib_maca, libmcblas loads from /opt/maca; op coverage 45/46 (only addmm TF32 rounding fails).
Add a MetaX Self-Contained Wheel (CUDA boxing) section covering the FLAGOS_METAX_BOXING=1 path: how to build the wheel (bdist_wheel + bundle_maca_libtorch.sh + repackage), the ~1.1G size / distribution tradeoff, and how to install and run on a clean target (official torch+cpu + this wheel + /opt/maca, no torch+metax wheel, no manual LD_LIBRARY_PATH). Note the two MetaX build modes in the runtime notes and warn that FLAGOS_USE_FLAGGEMS=1 must not be used with the boxing wheel (flagos_python backend is not compiled -> backend not registered).
…wheel Point readers to https://developer.metax-tech.com/softnova (SoftNova) for the MACA SDK (driver + cu-bridge + mxcc/cucc) and the torch+metax wheel, noting login is required and versions must match the card/driver/Python. Referenced from the top-level prerequisites, the MetaX source-build prerequisites, and the boxing-wheel build step.
We no longer use the hand-written mxcc/cucc kernel build (METAX_KERNEL=ON + torch+metax + Triton). Remove that 'Build from Source (MetaX Platform)' section and promote the self-contained CUDA boxing wheel to be the single MetaX build path (renamed to 'Build from Source (MetaX Platform)'). - Fold the MetaX developer-portal (SoftNova) SDK / torch+metax wheel download note into the boxing build steps (needed to build, not run). - Rewrite the 'Two build modes' runtime note to describe only the boxing wheel and fix its now-stale section anchor.
varargs unary in-place (7): asinh_/sinh_/log1p_/digamma_/sgn_/hardswish_/ logit_. gems wraps these as (*args, **kwargs) so inspect.signature can't recover arity; add _FLAGGEMS_ARITY_OVERRIDE so the aten schema supplies the authoritative positional count. Whitelist only holds simple elementwise ops verified to run + match CPU (maxdiff <= 1e-6) and drop no kwarg. rng (6): rand/randn (factory), rand_like/randn_like (like_factory), randperm (factory; no generator arg in this torch schema), multinomial (new rng_dropgen category dropping the trailing Generator?). Unblocked by _patch_flaggems_philox() in torch_fl/__init__.py, which monkeypatches gems' philox_backend_seed_offset to fall back to a held CUDA generator when torch.cuda.default_generators is empty (CPU-torch + cuda shim). One patch covers all 6 rng ops; no caller C++ change needed. Excluded: i0_/zero/zero.out hit a hardcoded tensor.is_cuda assert in the gems kernel (flagos is PrivateUse1, never true) so they stay in FLAGGEMS_PYTHON_SKIP; normal_/normal.* hardcode generator=None upstream (can't thread our generator). Verified: 15/15 numeric spot-checks pass; regressions unchanged from baseline (cuda-path 330 passed/45 skipped/3 xpassed, flaggems_python 27 passed). Fresh regen also drops the now-stale cudnn_convolution_bias_fused[.out] from backends_cuda.conf (absent from the torch 2.10 schema).
Route the high-level SDPA API to the aclnnFlashAttentionScore kernel instead of the math decomposition path, and register the view ops its pre/post-processing needs. - Register _fused_sdp_choice_stub DispatchStub for PrivateUse1 returning efficient_attention (2). PyTorch's scaled_dot_product_attention selects its fused backend via this C++ stub (gated by is_device_supported), not the aten-op-level _fused_sdp_choice. Macro must live in namespace at::native. - Add view/metadata ops for kAscend: transpose.int, permute, select.int, slice.Tensor, squeeze, squeeze.dim, unsqueeze, _unsafe_view, detach. Implemented via at::native::<fn> (not tensor member methods, which re-dispatch through PrivateUse1 and recurse -> segfault). select.int uses select_symint; explicit _native.h includes avoid int64->Dimname overload mis-resolution. Verified end-to-end on Ascend 910 NPU: forward, causal, autograd backward all pass; CPU-reference relative error ~0.0004 (fp16).
Add aclnn codegen for the in-place fill primitives: - zero_ -> aclnnInplaceZero - fill_.Scalar -> aclnnInplaceFillScalar - fill_.Tensor -> aclnnInplaceFillTensor Before this, zero_/fill_ had no device implementation, so the handwritten factory ops (zeros/ones_like/new_ones/scalar_tensor) fell back to flaggems/CPU for their internal .zero_()/.fill_() calls -- a hidden h2d path. Now the whole factory chain runs device-side aclnn. The factory ops themselves stay handwritten (their device/dtype inference is not expressible by codegen), but the fill work is pushed down to aclnn. Also move the SDPA + view-op conf entries above the codegen marker: they were appended after the "# --- generated by codegen_ascend.py ---" marker, which codegen truncates on every run, so a regen would silently drop them. Handwritten conf entries must live before the marker. Verified on Ascend 910: zero_/fill_.Scalar/fill_.Tensor and all four factory ops match CPU exactly (diff 0). SDPA + view-op tests still pass.
…3 (142 -> 145) Migrate three single-aclnn-call kernels from handwritten to codegen: - embedding -> aclnnEmbedding - embedding_dense_backward -> aclnnEmbeddingDenseBackward - constant_pad_nd -> aclnnConstantPadNd Each was a straight aclnn call with deterministic output-shape logic, so the handwritten .cc bodies map verbatim into codegen templates. Delete the three .cc files (globbed by CMake, no explicit list to update) and move their conf entries from the handwritten section into the generated block. Handwritten kAscend registrations drop 30 -> 27; codegen 142 -> 145. Verified on Ascend 910: all three match CPU exactly (diff 0), including embedding backward via autograd and multi-dim constant padding. Prior in-place/factory and view-op tests still pass.
…degen Extend codegen_ascend.py with three new structural capabilities so more handwritten kAscend kernels can be expressed as codegen templates: - .out variants: T_MATMUL / T_MATMUL_OUT generate functional + .out pairs (mm/mm.out via aclnnMm, bmm/bmm.out via aclnnBatchMatMul). Reusable pattern for any op whose .out kernel writes into a caller-shaped out&. - TensorList: T_CAT marshals at::ITensorListRef via aclCreateTensorList (aclnnCat), filtering numel==0 tensors; does not aclDestroyTensorList. - factory ops: T_ZEROS/T_SCALAR_TENSOR/T_ONES_LIKE/T_NEW_ONES build TensorOptions on-host + at::empty then fill via device-side zero_/fill_. New NO_ACLNN_CATEGORIES set skips the libopapi symbol guard for kernels that issue no direct aclnn call. Deletes 7 handwritten .cc (mm/bmm/cat/zeros/scalar_tensor/ones_like/new_ones). Handwritten kAscend regs 16 -> 7, codegen 145 -> 154, total unchanged. Verified vs CPU on NPU (test_matmul_codegen.py, test_cat_codegen.py, test_inplace_fill.py): all pass, matmul f32 within hf32-cube tolerance.
…S_USE_FLAGGEMS) Match the CUDA single-wheel model on MetaX: compile the FlagGems Python path (flagos_python backend) alongside the CUDA boxing kernels and select between them purely at runtime via FLAGOS_USE_FLAGGEMS, instead of forcing it off. - setup.py: drop the metax FLAGGEMS_PYTHON=OFF force; default ON like CUDA (only the C++ FlagGems path FLAGGEMS_KERNEL stays off). FLAGGEMS_PYTHON=0 still available for a slim pure-boxing build. - torch_fl/__init__.py: add explicit MetaX branch in _patch_flaggems_codegen_config (GEMS_VENDOR=metax + patch_torch_cuda_for_metax) before the ascend fallback, fixing metax wrongly getting GEMS_VENDOR=ascend; _select_backend_config picks backends_metax_flaggems.conf when FLAGOS_USE_FLAGGEMS=1 + FLAGOS_METAX_BOXING=1. - _metax_compat.py: add stream/availability/manual_seed shims and _patch_triton_do_bench (wall-clock) so FlagGems Triton kernels run on the CPU-frozen torch wheel against maca libtorch_cuda.so. - scripts/codegen_ops.py: generate backends_metax_flaggems.conf, routing the ops triton-metax cannot run (mm/bmm/mean.dim) and flag_gems device-guarded ops (mul, embedding_dense_backward, etc.) back to cuda boxing. - tests/integration/conftest.py: skip forcing backends_metax.conf in boxing mode (mxcc backend not compiled) so torch_fl's own config selection applies. - tests/integration/ops/conftest.py: skip @mark.metax tests in boxing mode (no metax backend to dispatch to). Verified: 255 passed / 147 skipped; only non-skip failure is the pre-existing out-of-scope cat empty-1d-tensor boxing bug.
Walks the transformers generate() path op-by-op until inference runs coherently (~1.8 tok/s) on a pure-aclnn C++ backend, and takes the ops parity suite from 269 failed/31 passed to 6 failed/294 passed. Bespoke handwritten kernels (generation path): - topk/sort/scatter/multinomial via aclnnTopk/Sort/Scatter/Multinomial - arange/argmax/isin/lift_fresh - rng.cc: device randn/rand/randint/randint.low via aclnnInplace Normal/Uniform/Random, seeded from the default CPU generator's random64() so the ops suite's on-device torch.randn(device=DEVICE) inputs work. Codegen (scripts/codegen_ascend.py): - rsub.Scalar -> binary_scalar_alpha (aclnnRsubs) - sum/max/min full-reduce -> reduce_sum_all / reduce_minmax_all - any -> reduce_all; ones/empty_like/full/full_like factory ops SDPA GQA fix (scaled_dot_product_attention.cc): expand kv heads using kv's own S_kv, not query's S -- during decode query S==1 but kv S==full context, so using query S gave an expand size mismatch. View ops (strided_ops): t, unbind.int (pure metadata, via at::native::). Runtime config: torch_fl auto-selects the ascend conf on a /dev/davinci* box; FLAGOS_USE_FLAGGEMS=1 opts into the FlagGems Triton path. Also patch triton-ascend npu_utils.cpp for the CANN 9.0.0 rtLimitType_t enum name.
Host-side dispatch optimizations for eager decode on real 910, keeping the
op set identical to torch_npu (no fusion in the measured path).
- aten::empty fast path: skip the DeviceGuard registry round-trip when the
device is unchanged; drop the ptr_to_block_ side map + its mutex from the
caching allocator by stashing Block* in the DataPtr context (5.25 -> 2.04
us/call, the single biggest win)
- repeatable aclOpExecutor cache in op_api_common.h (ExecAscendCached):
owns its aclTensors, rebinds addresses on hit, reuses the workspace tensor
- codegen: cached categories for elementwise/unary-scalar/reductions/softmax
plus a CPU-scalar fast path routing T+float through aclnnAdds/Muls/... to
avoid a per-call H2D copy
- register aten::matmul on AutogradPrivateUse1 -> aclnnMatmul, collapsing
mm/bmm/view churn and matching torch_npu's op counts exactly
- isin: compute on device instead of a triple CPU round-trip (this was the
hottest op in the generate() loop)
- _to_copy: cache the aclnnCast executor
Inference 13.4 -> 24.82 tok/s (0.89x torch_npu); training 507.9 tok/s (0.69x).
Adds tests/perf/e2e_qwen3_{infer,train}_ascend.py as the comparison harness.
…9x torch_npu Port the Ascend/aclnn backend onto flagos/main, which already carries the CUDA schema-codegen migration, Tsingmicro and Hygon DCU backends. Kernels: on-device strided copy (aclnnInplaceCopy), SDPA fwd+bwd via aclnnFlashAttentionScore/Grad, fused matmul, rms_norm, masked_select, sort/topk/scatter/multinomial/argmax/arange/isin, plus RNG and cast. Host-side dispatch work, which is where the remaining gap lives (both backends are host-bound at this model size): - async dispatch: drop the per-op aclrtSynchronizeStream; the root cause of the earlier stream mismatch was GetDefaultAclStream() being a per-.so inline static, fixed by a single FLAGOS_EXPORT definition - caching allocator: stash Block* as the DataPtr context so free is O(1) with no side map or lock - empty.memory_format: skip the DeviceGuard when the requested device is already current (~2.8us/call on the decode path) matmul is registered on AutogradPrivateUse1, not plain PrivateUse1. Since aten::matmul is CompositeImplicitAutograd, a concrete PrivateUse1 kernel makes autograd bind backward to aten::matmul_backward, which has no PrivateUse1 path and decays to CPU, crashing training. The wrapper takes the fused kernel only when no input requires grad and otherwise falls through to the composite mm/bmm decomposition, which has working autograd. Measured on 910 (Qwen3-0.6B, seq 128, batch 1, eager, same model and seed): inference 23.8-24.3 tok/s vs torch_npu 28.3 (~0.84x), training 507.9 tok/s vs 738.3 (~0.69x). Loss curves match (2.937->0.324 vs 3.14->0.32).
Upstream flagos-ai#22 made flag_gems + triton>=3.5.1 hard runtime deps for every accelerator except dcu. On Ascend the `triton` module is supplied by triton-ascend, which is installed out of band and has no PyPI release satisfying triton>=3.5.1, so `pip install -e .` pulls stock triton over it. Every Triton entry point then fails with "0 active drivers" and the torch_npu-shim patch from scripts/patch_triton_ascend.py is gone. Ascend needs the same carve-out as dcu, for the same reason: the accelerator provides its own triton and PyPI's NVIDIA-targeted wheel is the wrong artifact. flag_gems imports in the Python layer are ImportError-guarded, so omitting the dep is safe.
…r aclnn's 50-entry cap Adds the 7 _foreach_* Ascend kernels the AdamW foreach path needs (_foreach_mul_/add_.Scalar, _foreach_lerp_.Scalar, _foreach_addcmul_.Scalar, _foreach_sqrt, _foreach_div_/addcdiv_.ScalarList), plus stack, mean, clamp, clamp.Tensor and bitwise_and_/or_/xor_.Tensor. CANN's aclnnForeach* kernels only process the FIRST 50 entries of an aclTensorList. The ScalarList variants at least error past that (561002/161002); Mul/Add/Addcmul/Lerp/Sqrt instead return success and leave entries >= 50 UNTOUCHED, so the bug is silent -- _foreach_lerp_ over 200 tensors "succeeded" with 150 of them never written. The cap is on the entry count alone: measured identical for numel 8..65536 and fp16/fp32/bf16. AdamW passes 310 tensors for Qwen3-0.6B, so each foreach template is split into a <Kernel>Chunk doing the aclnn call and a wrapper that slices the lists into sub-50 chunks. TensorList and ArrayRef<Scalar> are both ArrayRefs, so the slicing is free, and elementwise semantics make it exact. Note for anyone probing other aclnn list limits: an exception-based search reports "no limit up to 4096" for the silent ops. The values have to be compared against CPU entry by entry. aclnnForeachAddcdivScalarList's scalars param is a device aclTensor whose dtype must match self (float32 scalars against fp16 inputs return 161002), so fp16 addcdiv lands a few ulp off CPU, which keeps the divisor at full precision. AdamW holds optimizer state in fp32, so this does not reach the training path. tests/integration/ops/test_foreach_dispatch.py covers all 7 ops at lengths that straddle the boundary (51/60/128/310) and asserts on every entry, plus AdamW(foreach=True) against foreach=False. Verified it fails (85 cases) when the chunk size is raised to 128. e2e_qwen3_train_ascend.py no longer hard-pins foreach=False; it takes --foreach/--no-foreach, applied to both backends. Measured on 910 (Qwen3-0.6B, seq 128, batch 1, eager, card 10): training 621.0 tok/s vs torch_npu 827.1 (0.75x, up from 0.69x at foreach=False's 507.9), loss curve unchanged at 2.937 -> 0.324. Inference unaffected at 24.5 tok/s. 169 foreach cases pass; ops/ suite and the qwen3 infer/train integration tests show no regressions.
aten::matmul is CompositeImplicitAutograd, so claiming a fused aclnnMatmul
kernel on PrivateUse1 stops the mm/bmm/view decomposition -- and with it the
sub-op graph autograd was relying on. The op then binds its real derivative,
aten::matmul_backward, which the backend must supply. Until now that was
sidestepped by taking the fused path only when !requires_grad: inference got
one aclnnMatmul, training kept the decomposition.
Close the gap the way torch_npu does, by generating the missing autograd layer
rather than hand-rolling it. scripts/codegen_autograd.py drives torchgen's own
emit_body() -- the generator behind PyTorch's in-tree VariableType_N.cpp -- to
produce a VariableType::matmul on AutogradPrivateUse1 that builds
MatmulBackward0 and redispatches to the fused kernel. Only that thin layer is
generated: the backward node classes already ship in libtorch, so unlike
torch_npu we do not regenerate Functions.h/ADInplaceOrView/python bindings.
Adding an op is one entry in AUTOGRAD_OPS.
matmul_backward itself is implemented with two cached aclnnMatmul calls.
op-plugin's MatmulBackwardKernelNpuOpApi.cpp was the starting reference and has
two real bugs, deliberately not reproduced:
* 2-D x N-D reshapes grad to {M, -1}. mat2^T flattens to (B*N, K), so its row
index is the pair (b, n); grad must carry the same pair as its column
index, which needs M permuted to the front first. The plain reshape pairs
(m, n) against (b, n) and silently mixes batches.
* Only *leading* singleton batch dims are stripped, so any interior or
trailing broadcast returns a wrong-shaped gradient -- (2,1,3,4) x (2,5,4,6)
yields a (2,5,3,4) grad for a (2,1,3,4) input. Replaced with at::sum_to
onto the promoted operand shape, a no-op when nothing was broadcast.
Guarded by USE_ASCEND throughout: other backends have no fused kernel, keep
PyTorch's decomposition, and never bind a matmul_backward they cannot service.
Qwen3-0.6B training, batch 1 x seq 128, on a real 910:
backward dispatches 5466 -> 4144 (torch_npu 4002)
matmul_backward n/a -> 253 (torch_npu 253)
backward self-CPU 117.1 -> 109.9 ms/step
throughput 566 -> 674 tok/s, 0.69x -> 0.82x torch_npu
Verified against a float64 CPU reference -- the kernel runs hf32 cube math, so
absolute fp32 comparison misreads ~1e-4 precision as a correctness bug. The new
test carries 30 cases and every shape in it catches a real defect in the
op-plugin rules on one side or the other.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Brings the Ascend NPU line up to date with main (18 commits: MUSA backend,
GCU topsaten kernels, FlagGems C++ dispatch, RNG generator injection, DCU/
MetaX fixes, CI rework).
Eight conflict hunks across six files, all resolved by keeping both sides:
csrc/aten/common.h Backend enum gained kGcu (main) and kUncached (here,
the Dispatcher per-op cache sentinel). Kept both,
with kUncached last so real backends stay contiguous.
csrc/aten/copy_ops.cc The _to_copy dtype-cast branch. Main added MUSA's
on-device mudnn cast plus GCU/MUSA to the no-CUDA
list; this line added Ascend's on-device aclnnCast.
Kept the MUSA fast path, the Ascend fast path, and
the CPU round-trip as the shared fallback.
scripts/codegen_ops.py One template lost a variable per side: {holder_lines}
(optional-tensor boxing) here, {inject} (RNG
generator) on main. Both restored, matching the
layout the file's other templates already use.
setup.py Main refactored the triton-dep skip list into
_vendor_supplies_triton(); ascend was missing from
it, so it moved into the helper rather than being
re-hardcoded at the call site.
torch_fl/__init__.py Backend-conf selection: main added the FlagGems C++
path, this line added Ascend /dev/davinci detection.
Independent -- Ascend returns early, the C++ branch
joins the conf-name chain.
csrc/aten/register.cc Pure main-side additions (MUSA registration block).
Verified on a real 910: codegen_ops.py reproduces all four generated files
byte-for-byte, ruff check/format clean, ACCELERATOR=ascend builds, 37 tests
pass (qwen3 train + infer + matmul backward), Qwen3-0.6B training holds at
686 tok/s with loss converging 2.94 -> 0.32.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
CUDA CI failed at `import torch_fl` with
ImportError: libtorch_fl.so: undefined symbol:
_ZN2at6native6flagos18MatmulKernelAscendERKNS_6TensorES4_
Two separate instances of the same mistake, both introduced by this branch.
1. register.cc referenced MatmulKernelAscend behind a *runtime* check
(GetBackendForOp("matmul") == kAscend) with only the forward declaration
guarded asymmetrically -- the backward decl was inside #if defined(USE_ASCEND),
the forward one was not. A runtime branch does not remove a link-time
reference, and backends/ascend/matmul.cc is not compiled without USE_ASCEND.
A shared library links fine with undefined symbols and only fails at dlopen,
which is why "Build wheel (CUDA)" passed and the failure surfaced at import.
Fix: guard both the declarations and the call site at compile time.
2. copy_ops.cc / contiguous_ops.cc call ascend::StridedCopy and ascend::DtypeCast
from #else branches that cover TsingMicro, GCU and MUSA-without-mudnn as well
as Ascend, but included ascend_copy.h only under #ifdef USE_ASCEND. Those
platforms failed to *compile* ("'ascend' has not been declared"); no CI builds
them, so it stayed hidden. Fix: ascend_copy.h now supplies inline no-op
fallbacks for non-Ascend builds (defined, not just declared, so nothing is
left undefined at load), and is included unconditionally. The no-ops report
"unavailable" and callers take the CPU round-trip they already implement.
Verified: register.cc compiled without USE_ASCEND has no undefined
MatmulKernelAscend reference (nm -u); copy_ops.cc and contiguous_ops.cc compile
clean under each of USE_TSINGMICRO / USE_GCU / USE_MUSA and with no macro at
all. On real 910: 30/30 matmul-backward, 175/175 foreach + conf-consistency,
3/3 Qwen3 training.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this is
The Ascend NPU line, brought up to date with
mainand ready to merge. It takesthe backend from "builds" to Qwen3-0.6B training and inference on a real 910,
with training at 0.82x torch_npu.
Everything here is on hardware-verified ground: every number below was measured
on a real 910 in this repo, not estimated.
How ops get added
Three codegen scripts, no hand-written glue:
codegen_ops.pycodegen_ascend.pycodegen_autograd.pyVariableType::<op>on AutogradPrivateUse1212 aclnn kernels total (186 generated + 26 hand-written for shapes the templates
can't express:
masked_select's numel buffer, SDPA's FlashAttentionScore,multinomial/sort/topk).codegen_autograd.pyis the new piece.aten::matmulis CompositeImplicitAutograd,so claiming a fused
aclnnMatmulkernel stops the mm/bmm/view decomposition — andwith it the sub-op graph autograd relied on. Rather than hand-roll the replacement,
it drives torchgen's own
emit_body()(the generator behind PyTorch's in-treeVariableType_N.cpp) so the autograd bookkeeping matches upstream exactly. Onlythat thin layer is generated — the backward node classes already ship in libtorch,
so unlike torch_npu we don't regenerate
Functions.h/ADInplaceOrView/pythonbindings. Adding an op is one entry in
AUTOGRAD_OPS.Performance (Qwen3-0.6B, batch 1 x seq 128, real 910)
Training dispatches per step, vs the baseline:
matmul_backward x253now matches torch_npu exactly;mm394,t788 andview816 are gone from backward.Getting there: 18.9 -> 686 tok/s training, 1.67 -> 8.64 inference. The steps that
mattered were on-device
InplaceCopy(3.7x, killed a D2H->CPU->H2D round-trip inGQA
repeat_kv), an async-dispatch fix (GetDefaultAclStream()was a per-.soinline static, so drain and enqueue hit different streams), on-device
aclnnCast,chunked
_foreach_*(aclnn silently drops list entries past 50 — no error), andthis PR's fused matmul.
A note on op-plugin
matmul_backwardstarted from op-plugin'sMatmulBackwardKernelNpuOpApi.cpp.It has two real bugs, deliberately not reproduced:
{M, -1}. Sincemat2^Tflattens to
(B*N, K), its row index is the pair(b,n)— grad must carry thesame pair as its column index, which needs
Mpermuted to the front first.Verified: true
grad_a[0,:2] = [-2.2095, 3.4045], op-plugin gives[-0.8153, 4.0754].trailing broadcast returns a wrong-shaped gradient —
(2,1,3,4) x (2,5,4,6)yields a
(2,5,3,4)grad for a(2,1,3,4)input. Replaced withat::sum_toonto the promoted operand shape (a no-op when nothing broadcast, so the
transformer path pays nothing).
The algorithm was validated in float64 across 87 case/mask combinations and a
3763-pair randomized shape sweep before being committed to C++.
Merge conflicts
Eight hunks across six files (
common.h,copy_ops.cc,codegen_ops.py,setup.py,torch_fl/__init__.py,register.cc), all resolved by keeping bothsides — see the merge commit for the per-file rationale. Two worth flagging:
copy_ops.cc: main added MUSA's on-device cast, this line added Ascend'saclnnCast. Kept both fast paths plus the CPU round-trip as shared fallback.setup.py: main refactored the triton-skip list into_vendor_supplies_triton()but
ascendwas missing from it. Moved into the helper rather thanre-hardcoding at the call site — without it pip installs stock triton over
triton-ascend and every Triton entry point dies with "0 active drivers".
Verification
ACCELERATOR=ascendbuilds clean;ruff check/ruff format --checkpasscodegen_ops.pyreproduces all four generated files byte-for-byte;codegen_autograd.pylikewisetest_matmul_backward_dispatch.py: 30 cases. Compared against a float64reference — the kernel runs hf32 cube math, so absolute fp32 comparison misreads
~1e-4 precision as a correctness bug. Every shape in it catches a real defect in
the op-plugin rules on one side or the other.
variable_type.ccsyntax-checked withUSE_ASCENDundefined (it's glob-collectedinto every build): compiles inert, so CUDA/MetaX/GCU/MUSA are untouched.
Not addressed
The remaining 0.82x gap is not matmul — it's per-op executor overhead. torch_npu
caches
aclOpExecutorkeyed by op+shape and only swaps tensor addresses on reuse(
AclSetTensorAddr/SetRepeatable/SetPTACacheHashKey). That's the next lever,and it's what the 0.32x inference number is bounded by too: inference is a
per-token stream of small ops where fixed host cost dominates, whereas training
amortizes it over larger compute.
🤖 Generated with Claude Code