diff --git a/README.md b/README.md index 61b06e9d..03c2eb25 100644 --- a/README.md +++ b/README.md @@ -169,15 +169,40 @@ rm -rf ~/.triton/cache/ #### 4. Verify Installation +Two runtime gotchas on Ascend: + +- **Import order:** `import torch_fl` **before** `import flag_gems` — torch_fl installs the `torch.npu` shim and sets `GEMS_VENDOR=ascend` that FlagGems reads at its own import time. +- **libstdc++:** FlagGems pulls in `sqlalchemy`→`_sqlite3`, which needs `CXXABI_1.3.15`. If the system `libstdc++.so.6` is older, preload conda's: `export LD_PRELOAD=$CONDA_PREFIX/lib/libstdc++.so.6`. + ```bash +export LD_PRELOAD=$CONDA_PREFIX/lib/libstdc++.so.6 # if system libstdc++ is old python -c " -import torch_fl +import torch_fl, flag_gems print('device count:', torch_fl.flagos.device_count()) -print('FlagGems enabled:', torch_fl.is_flaggems_enabled()) -print('registered ops:', len(torch_fl.get_registered_ops())) +print('flag_gems:', flag_gems.__version__) +" +``` + +Enable the FlagGems Triton path at runtime. On an Ascend NPU box (detected via +`/dev/davinci*`) `torch_fl` auto-selects the ascend config — no need to set +`FLAGOS_BACKEND_CONFIG` by hand: + +```bash +# Pure aclnn C++ backend (default): no env needed -> backends_ascend.conf +# FlagGems Triton where triton-ascend runs: FLAGOS_USE_FLAGGEMS=1 +# -> backends_ascend_flagos_py.conf +FLAGOS_USE_FLAGGEMS=1 FLAGOS_LOG_DISPATCH=1 python -c " +import torch, torch_fl, flag_gems +x = torch.randn(64, 64).to('flagos:0') +print('abs matches CPU:', torch.allclose(torch.abs(x).cpu(), x.cpu().abs())) " +# expect: [flagos dispatch] abs -> flagos_python ``` +> Ops that triton-ascend cannot compile are routed back to the `ascend` aclnn +> kernel in `backends_ascend_flagos_py.conf` (annotated per op). FlagGems is +> optional on Ascend — without it, leave `FLAGOS_USE_FLAGGEMS` unset. + #### 5. Run Tests ```bash diff --git a/csrc/aten/backends/ascend/arange.cc b/csrc/aten/backends/ascend/arange.cc new file mode 100644 index 00000000..89389ff3 --- /dev/null +++ b/csrc/aten/backends/ascend/arange.cc @@ -0,0 +1,90 @@ +// Copyright (c) 2026, BAAI. All rights reserved. + +#include "../../generated/ops.h" +#include +#include +#include "op_preparation.h" +#include "op_api_common.h" + +namespace at::native::flagos { + +namespace { + +// Compute the number of elements in arange(start, end, step), matching +// PyTorch's reference (aten/src/ATen/native/RangeFactories.cpp): ceil for +// integral dtypes, and a fudge-factor guard for floating point. +int64_t ArangeSize(const at::Scalar& start, const at::Scalar& end, + const at::Scalar& step, at::ScalarType dtype) { + if (c10::isIntegralType(dtype, /*includeBool=*/false)) { + int64_t s = start.toLong(), e = end.toLong(), st = step.toLong(); + TORCH_CHECK(st != 0, "arange: step must be nonzero"); + if ((st > 0 && e < s) || (st < 0 && e > s)) return 0; + // ceil division that also works for negative step. + return (e - s + st - (st > 0 ? 1 : -1)) / st; + } + double s = start.toDouble(), e = end.toDouble(), st = step.toDouble(); + TORCH_CHECK(st != 0, "arange: step must be nonzero"); + double n = std::ceil((e - s) / st); + return n < 0 ? 0 : static_cast(n); +} + +} // namespace + +// arange.start_step(Scalar start, Scalar end, Scalar step, ScalarType?, ...) -> Tensor +at::Tensor ArangeStartStepKernelAscend( + const at::Scalar& start, const at::Scalar& end, const at::Scalar& step, + ::std::optional dtype, ::std::optional layout, + ::std::optional device, ::std::optional pin_memory) { + namespace ascend = at::native::flagos::ascend; + + // Default dtype: long if all args integral, else the default float type + // (mirrors torch's arange type-promotion for the common cases used by + // transformers' cache_position = arange(...)). + at::ScalarType out_dtype = dtype.value_or( + (start.isIntegral(false) && end.isIntegral(false) && step.isIntegral(false)) + ? at::kLong + : at::typeMetaToScalarType(c10::get_default_dtype())); + + auto options = at::TensorOptions() + .dtype(out_dtype) + .layout(layout.value_or(at::kStrided)) + .device(device.value_or(at::Device(at::kPrivateUse1, 0))) + .pinned_memory(pin_memory.value_or(false)); + + int64_t n = ArangeSize(start, end, step, out_dtype); + auto out = ascend::OpPreparation::apply_tensor_without_format({n}, options); + if (n == 0) return out; + + ascend::AclScalarWrapper acl_start(start, out_dtype); + ascend::AclScalarWrapper acl_end(end, out_dtype); + ascend::AclScalarWrapper acl_step(step, out_dtype); + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD(aclnnArange, acl_start.get(), acl_end.get(), acl_step.get(), + acl_out.get()); + return out; +} + +// arange.start(Scalar start, Scalar end, ...) -> step defaults to 1. +at::Tensor ArangeStartKernelAscend( + const at::Scalar& start, const at::Scalar& end, + ::std::optional dtype, ::std::optional layout, + ::std::optional device, ::std::optional pin_memory) { + return ArangeStartStepKernelAscend(start, end, at::Scalar(1), dtype, layout, + device, pin_memory); +} + +// arange(Scalar end, ...) -> start defaults to 0, step to 1. +at::Tensor ArangeKernelAscend( + const at::Scalar& end, + ::std::optional dtype, ::std::optional layout, + ::std::optional device, ::std::optional pin_memory) { + return ArangeStartStepKernelAscend(at::Scalar(0), end, at::Scalar(1), dtype, + layout, device, pin_memory); +} + +REGISTER_IMPL_TO_DISPATCHER(ArangeStartStepFn, arange_start_step_dispatcher, Backend::kAscend, ArangeStartStepKernelAscend) +REGISTER_IMPL_TO_DISPATCHER(ArangeStartFn, arange_start_dispatcher, Backend::kAscend, ArangeStartKernelAscend) +REGISTER_IMPL_TO_DISPATCHER(ArangeFn, arange_dispatcher, Backend::kAscend, ArangeKernelAscend) + +} // namespace at::native::flagos diff --git a/csrc/aten/backends/ascend/argmax.cc b/csrc/aten/backends/ascend/argmax.cc new file mode 100644 index 00000000..da4a520f --- /dev/null +++ b/csrc/aten/backends/ascend/argmax.cc @@ -0,0 +1,55 @@ +// Copyright (c) 2026, BAAI. All rights reserved. + +#include "../../generated/ops.h" +#include +#include "op_preparation.h" +#include "op_api_common.h" + +namespace at::native::flagos { + +// argmax(Tensor self, int? dim=None, bool keepdim=False) -> Tensor +// +// aclnnArgMax requires a concrete reduction dim, so when dim is nullopt we +// flatten to 1-D and reduce over axis 0 (matching torch's global-argmax +// semantics). The output is int64 (torch always returns Long indices). +at::Tensor ArgmaxKernelAscend(const at::Tensor& self, + ::std::optional dim, bool keepdim) { + namespace ascend = at::native::flagos::ascend; + + at::Tensor input; + int64_t reduce_dim; + if (dim.has_value()) { + input = self; + reduce_dim = dim.value(); + } else { + // Global argmax: flatten, reduce dim 0, keepdim is ignored by torch here + // (result is a 0-d scalar unless keepdim was requested on the flat view). + input = self.reshape({-1}); + reduce_dim = 0; + } + + // Compute output shape: drop (or keep as size-1) the reduced dim. + std::vector out_sizes; + int64_t ndim = input.dim(); + int64_t d = reduce_dim < 0 ? reduce_dim + ndim : reduce_dim; + for (int64_t i = 0; i < ndim; ++i) { + if (i == d) { + if (keepdim) out_sizes.push_back(1); + } else { + out_sizes.push_back(input.size(i)); + } + } + + auto out = ascend::OpPreparation::apply_tensor_without_format( + out_sizes, input.options().dtype(at::kLong)); + + ascend::AclTensorWrapper acl_self(input); + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD(aclnnArgMax, acl_self.get(), d, keepdim, acl_out.get()); + return out; +} + +REGISTER_IMPL_TO_DISPATCHER(ArgmaxFn, argmax_dispatcher, Backend::kAscend, ArgmaxKernelAscend) + +} // namespace at::native::flagos diff --git a/csrc/aten/backends/ascend/ascend_copy.cc b/csrc/aten/backends/ascend/ascend_copy.cc new file mode 100644 index 00000000..4922a1d5 --- /dev/null +++ b/csrc/aten/backends/ascend/ascend_copy.cc @@ -0,0 +1,71 @@ +// Copyright (c) 2026, BAAI. All rights reserved. + +#include "ascend_copy.h" + +#include +#include "op_api_common.h" + +namespace at::native::flagos::ascend { + +bool StridedCopy(const at::Tensor& dst, const at::Tensor& src) { + if (!dst.defined() || !src.defined()) { + return false; + } + if (!dst.is_privateuseone() || !src.is_privateuseone()) { + return false; + } + if (dst.numel() == 0) { + return true; // nothing to copy + } + + // aclnnInplaceCopy(selfRef, src): writes src into selfRef, honoring the + // strides/offset recorded on each aclTensor. AclTensorWrapper preserves the + // tensor's sizes/strides/offset, so a non-contiguous src is copied correctly + // into the (contiguous) dst without a host round-trip. + AclTensorWrapper dst_wrap(dst); + AclTensorWrapper src_wrap(src); + EXEC_ASCEND_CMD(aclnnInplaceCopy, + const_cast(dst_wrap.get()), + src_wrap.get()); + return true; +} + +at::Tensor DtypeCast(const at::Tensor& src, at::ScalarType dtype) { + if (!src.defined() || !src.is_privateuseone()) { + return {}; + } + // aclnnCast expects a dense input; make src contiguous first (cheap, and the + // callers in _to_copy already pass a contiguous tensor). + at::Tensor src_c = src.is_contiguous() ? src : src.contiguous(); + at::Tensor out = at::empty(src_c.sizes(), src_c.options().dtype(dtype)); + if (src_c.numel() == 0) { + return out; + } + + // aclnnCast(self, dtype, out): converts self to the given aclDataType + // on-device. Route through the repeatable-executor cache: RMSNorm emits two + // fp16<->fp32 casts per layer (285/step) at fixed decode shapes, so the + // GetWorkspaceSize + aclCreateTensor build cost is paid once per shape. The + // target aclDataType is baked into the executor at build time, so it must be + // part of the cache key (folded in via SigHasher::val below). + const aclDataType acl_dtype = ToAclDataType(dtype); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + SigHasher hsh; + hsh.tensor(src_c); + hsh.tensor(out); + hsh.val(static_cast(acl_dtype)); + ExecAscendCached( + "aclnnCast", "aclnnCastGetWorkspaceSize", + opApiFuncAddr, getWsFuncAddr, hsh.h, + {&src_c}, {&out}, + [&](GwsFunc gws, + std::vector& in, + std::vector& out_t, + uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_dtype, out_t[0].acl_tensor, pws, pex); + }); + return out; +} + +} // namespace at::native::flagos::ascend diff --git a/csrc/aten/backends/ascend/ascend_copy.h b/csrc/aten/backends/ascend/ascend_copy.h new file mode 100644 index 00000000..84cece94 --- /dev/null +++ b/csrc/aten/backends/ascend/ascend_copy.h @@ -0,0 +1,50 @@ +// Copyright (c) 2026, BAAI. All rights reserved. +// +// On-device strided copy for the Ascend backend. Lets the platform-neutral +// copy_/clone/contiguous paths avoid the CPU round-trip (device->host strided +// copy->device) that dominates GQA repeat_kv clones in Qwen3 inference. + +#pragma once + +#include + +namespace at::native::flagos::ascend { + +#if defined(USE_ASCEND) + +// Copy `src` into `dst` entirely on the NPU via aclnnInplaceCopy, which handles +// differing strides/offsets and dtype casts on-device. `dst` must be an +// allocated PrivateUse1 tensor with matching sizes; `src` may be non-contiguous. +// Returns true on success. Callers use the return value to fall back to the CPU +// round-trip if the on-device path is unavailable. +bool StridedCopy(const at::Tensor& dst, const at::Tensor& src); + +// Cast `src` (a contiguous PrivateUse1 tensor) to `dtype` entirely on the NPU +// via aclnnCast, returning a freshly-allocated contiguous PrivateUse1 tensor. +// Replaces the D2H -> CPU cast -> H2D round-trip in _to_copy's Ascend dtype +// path, which dominated HF RMSNorm (two fp16<->fp32 casts per layer). Returns +// an undefined tensor if the on-device path is unavailable (caller falls back). +at::Tensor DtypeCast(const at::Tensor& src, at::ScalarType dtype); + +#else + +// Non-Ascend builds: the shared copy_/clone/contiguous paths in copy_ops.cc and +// contiguous_ops.cc reach these from an #else branch that covers TsingMicro, +// GCU and MUSA-without-mudnn as well as Ascend. Those platforms have no aclnn, +// so provide inline no-ops that report "unavailable" and let the caller take +// the CPU round-trip it already implements as the fallback. +// +// These must be defined (not just declared): a .so links with undefined symbols +// and only fails at dlopen, so a bare declaration would produce a wheel that +// imports fine on Ascend and dies with "undefined symbol" everywhere else. +inline bool StridedCopy(const at::Tensor&, const at::Tensor&) { + return false; +} + +inline at::Tensor DtypeCast(const at::Tensor&, at::ScalarType) { + return at::Tensor(); +} + +#endif + +} // namespace at::native::flagos::ascend diff --git a/csrc/aten/backends/ascend/generated/ascend_kernels.cc b/csrc/aten/backends/ascend/generated/ascend_kernels.cc index 00168016..22b5ffef 100644 --- a/csrc/aten/backends/ascend/generated/ascend_kernels.cc +++ b/csrc/aten/backends/ascend/generated/ascend_kernels.cc @@ -9,6 +9,7 @@ #include "../../../generated/ops.h" #include +#include #include #include #include @@ -24,10 +25,16 @@ at::Tensor SqrtKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnSqrt, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnSqrt", "aclnnSqrtGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -38,10 +45,16 @@ at::Tensor ExpKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnExp, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnExp", "aclnnExpGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -52,10 +65,16 @@ at::Tensor TanhKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnTanh, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnTanh", "aclnnTanhGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -66,10 +85,16 @@ at::Tensor SigmoidKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnSigmoid, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnSigmoid", "aclnnSigmoidGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -80,10 +105,16 @@ at::Tensor ReciprocalKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnReciprocal, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnReciprocal", "aclnnReciprocalGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -94,10 +125,16 @@ at::Tensor LogKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnLog, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnLog", "aclnnLogGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -108,10 +145,16 @@ at::Tensor FloorKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnFloor, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnFloor", "aclnnFloorGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -122,10 +165,16 @@ at::Tensor CeilKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnCeil, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnCeil", "aclnnCeilGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -136,10 +185,16 @@ at::Tensor ErfKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnErf, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnErf", "aclnnErfGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -150,10 +205,16 @@ at::Tensor ErfcKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnErfc, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnErfc", "aclnnErfcGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -164,10 +225,16 @@ at::Tensor Expm1KernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnExpm1, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnExpm1", "aclnnExpm1GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -178,10 +245,16 @@ at::Tensor Log2KernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnLog2, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnLog2", "aclnnLog2GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -192,10 +265,16 @@ at::Tensor Log10KernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnLog10, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnLog10", "aclnnLog10GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -206,10 +285,16 @@ at::Tensor Log1pKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnLog1p, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnLog1p", "aclnnLog1pGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -220,10 +305,16 @@ at::Tensor RoundKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnRound, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnRound", "aclnnRoundGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -234,10 +325,16 @@ at::Tensor TruncKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnTrunc, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnTrunc", "aclnnTruncGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -248,10 +345,16 @@ at::Tensor FracKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnFrac, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnFrac", "aclnnFracGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -262,10 +365,16 @@ at::Tensor SignKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnSign, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnSign", "aclnnSignGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -276,10 +385,16 @@ at::Tensor ReluKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnRelu, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnRelu", "aclnnReluGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -290,10 +405,16 @@ at::Tensor CoshKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnCosh, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnCosh", "aclnnCoshGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -304,10 +425,16 @@ at::Tensor SinhKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnSinh, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnSinh", "aclnnSinhGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -318,10 +445,16 @@ at::Tensor AsinKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnAsin, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnAsin", "aclnnAsinGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -332,10 +465,16 @@ at::Tensor AtanKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnAtan, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnAtan", "aclnnAtanGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -346,10 +485,16 @@ at::Tensor AsinhKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnAsinh, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnAsinh", "aclnnAsinhGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -360,10 +505,16 @@ at::Tensor AcoshKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnAcosh, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnAcosh", "aclnnAcoshGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -374,10 +525,16 @@ at::Tensor AtanhKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnAtanh, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnAtanh", "aclnnAtanhGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -388,10 +545,16 @@ at::Tensor LogicalNotKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnLogicalNot, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnLogicalNot", "aclnnLogicalNotGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -402,10 +565,16 @@ at::Tensor BitwiseNotKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnBitwiseNot, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnBitwiseNot", "aclnnBitwiseNotGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -416,10 +585,16 @@ at::Tensor AbsKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnAbs, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnAbs", "aclnnAbsGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -430,10 +605,16 @@ at::Tensor AcosKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnAcos, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnAcos", "aclnnAcosGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -444,10 +625,16 @@ at::Tensor CosKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnCos, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnCos", "aclnnCosGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -458,10 +645,16 @@ at::Tensor SinKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnSin, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnSin", "aclnnSinGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -472,10 +665,16 @@ at::Tensor NegKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnNeg, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnNeg", "aclnnNegGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -486,10 +685,16 @@ at::Tensor RsqrtKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnRsqrt, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnRsqrt", "aclnnRsqrtGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -500,10 +705,16 @@ at::Tensor SiluKernelAscend(const at::Tensor& self) { auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnSilu, acl_self.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "aclnnSilu", "aclnnSiluGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -511,21 +722,50 @@ REGISTER_IMPL_TO_DISPATCHER(SiluFn, silu_dispatcher, Backend::kAscend, SiluKerne at::Tensor DivTensorKernelAscend(const at::Tensor& self, const at::Tensor& other) { namespace ascend = at::native::flagos::ascend; + if (self.is_privateuseone() && !other.is_privateuseone() && other.numel() == 1) { + at::Scalar sc = other.item(); + auto out = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options()); + ascend::AclScalarWrapper acl_sc(sc, self.scalar_type()); + static void* sOpAddr = nullptr; static void* sWsAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + { double sv = sc.toDouble(); hsh.val(sv); } + ascend::ExecAscendCached( + "aclnnDivs", "aclnnDivsGetWorkspaceSize", sOpAddr, sWsAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_sc.get(), out_t[0].acl_tensor, pws, pex); + }); + return out; + } auto result_dtype = self.scalar_type(); auto other_c = other.is_privateuseone() ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnDiv, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnDiv", "aclnnDivGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -533,21 +773,50 @@ REGISTER_IMPL_TO_DISPATCHER(DivTensorFn, div_tensor_dispatcher, Backend::kAscend at::Tensor MulTensorKernelAscend(const at::Tensor& self, const at::Tensor& other) { namespace ascend = at::native::flagos::ascend; + if (self.is_privateuseone() && !other.is_privateuseone() && other.numel() == 1) { + at::Scalar sc = other.item(); + auto out = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options()); + ascend::AclScalarWrapper acl_sc(sc, self.scalar_type()); + static void* sOpAddr = nullptr; static void* sWsAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + { double sv = sc.toDouble(); hsh.val(sv); } + ascend::ExecAscendCached( + "aclnnMuls", "aclnnMulsGetWorkspaceSize", sOpAddr, sWsAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_sc.get(), out_t[0].acl_tensor, pws, pex); + }); + return out; + } auto result_dtype = self.scalar_type(); auto other_c = other.is_privateuseone() ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnMul, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnMul", "aclnnMulGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -560,16 +829,28 @@ at::Tensor BitwiseAndTensorKernelAscend(const at::Tensor& self, const at::Tensor ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnBitwiseAndTensor, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnBitwiseAndTensor", "aclnnBitwiseAndTensorGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -582,16 +863,28 @@ at::Tensor PowTensorTensorKernelAscend(const at::Tensor& self, const at::Tensor& ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnPowTensorTensor, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnPowTensorTensor", "aclnnPowTensorTensorGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -604,16 +897,28 @@ at::Tensor Atan2KernelAscend(const at::Tensor& self, const at::Tensor& other) { ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnAtan2, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnAtan2", "aclnnAtan2GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -626,16 +931,28 @@ at::Tensor MaximumKernelAscend(const at::Tensor& self, const at::Tensor& other) ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnMaximum, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnMaximum", "aclnnMaximumGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -648,16 +965,28 @@ at::Tensor MinimumKernelAscend(const at::Tensor& self, const at::Tensor& other) ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnMinimum, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnMinimum", "aclnnMinimumGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -670,16 +999,28 @@ at::Tensor BitwiseOrTensorKernelAscend(const at::Tensor& self, const at::Tensor& ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnBitwiseOrTensor, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnBitwiseOrTensor", "aclnnBitwiseOrTensorGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -692,16 +1033,28 @@ at::Tensor BitwiseXorTensorKernelAscend(const at::Tensor& self, const at::Tensor ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnBitwiseXorTensor, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnBitwiseXorTensor", "aclnnBitwiseXorTensorGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -709,22 +1062,53 @@ REGISTER_IMPL_TO_DISPATCHER(BitwiseXorTensorFn, bitwise_xor_tensor_dispatcher, B at::Tensor SubTensorKernelAscend(const at::Tensor& self, const at::Tensor& other, const at::Scalar& alpha) { namespace ascend = at::native::flagos::ascend; + if (self.is_privateuseone() && !other.is_privateuseone() && other.numel() == 1) { + at::Scalar sc = other.item(); + auto out = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options()); + ascend::AclScalarWrapper acl_sc(sc, self.scalar_type()); + ascend::AclScalarWrapper acl_alpha_s(alpha, self.scalar_type()); + static void* sOpAddr = nullptr; static void* sWsAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + { double sv = sc.toDouble(); hsh.val(sv); double av = alpha.toDouble(); hsh.val(av); } + ascend::ExecAscendCached( + "aclnnSubs", "aclnnSubsGetWorkspaceSize", sOpAddr, sWsAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_sc.get(), acl_alpha_s.get(), out_t[0].acl_tensor, pws, pex); + }); + return out; + } auto result_dtype = self.scalar_type(); auto other_c = other.is_privateuseone() ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); ascend::AclScalarWrapper acl_alpha(alpha, result_dtype); - ascend::AclTensorWrapper acl_out(out); - EXEC_ASCEND_CMD(aclnnSub, acl_self.get(), acl_other.get(), acl_alpha.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + { double av = alpha.toDouble(); hsh.val(av); } + ascend::ExecAscendCached( + "aclnnSub", "aclnnSubGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, acl_alpha.get(), out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -732,22 +1116,53 @@ REGISTER_IMPL_TO_DISPATCHER(SubTensorFn, sub_tensor_dispatcher, Backend::kAscend at::Tensor AddTensorKernelAscend(const at::Tensor& self, const at::Tensor& other, const at::Scalar& alpha) { namespace ascend = at::native::flagos::ascend; + if (self.is_privateuseone() && !other.is_privateuseone() && other.numel() == 1) { + at::Scalar sc = other.item(); + auto out = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options()); + ascend::AclScalarWrapper acl_sc(sc, self.scalar_type()); + ascend::AclScalarWrapper acl_alpha_s(alpha, self.scalar_type()); + static void* sOpAddr = nullptr; static void* sWsAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + { double sv = sc.toDouble(); hsh.val(sv); double av = alpha.toDouble(); hsh.val(av); } + ascend::ExecAscendCached( + "aclnnAdds", "aclnnAddsGetWorkspaceSize", sOpAddr, sWsAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_sc.get(), acl_alpha_s.get(), out_t[0].acl_tensor, pws, pex); + }); + return out; + } auto result_dtype = self.scalar_type(); auto other_c = other.is_privateuseone() ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); ascend::AclScalarWrapper acl_alpha(alpha, result_dtype); - ascend::AclTensorWrapper acl_out(out); - EXEC_ASCEND_CMD(aclnnAdd, acl_self.get(), acl_other.get(), acl_alpha.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + { double av = alpha.toDouble(); hsh.val(av); } + ascend::ExecAscendCached( + "aclnnAdd", "aclnnAddGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, acl_alpha.get(), out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -790,16 +1205,28 @@ at::Tensor EqTensorKernelAscend(const at::Tensor& self, const at::Tensor& other) ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options().dtype(at::kBool)); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnEqTensor, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnEqTensor", "aclnnEqTensorGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -812,16 +1239,28 @@ at::Tensor NeTensorKernelAscend(const at::Tensor& self, const at::Tensor& other) ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options().dtype(at::kBool)); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnNeTensor, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnNeTensor", "aclnnNeTensorGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -834,16 +1273,28 @@ at::Tensor GtTensorKernelAscend(const at::Tensor& self, const at::Tensor& other) ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options().dtype(at::kBool)); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnGtTensor, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnGtTensor", "aclnnGtTensorGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -856,16 +1307,28 @@ at::Tensor LtTensorKernelAscend(const at::Tensor& self, const at::Tensor& other) ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options().dtype(at::kBool)); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnLtTensor, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnLtTensor", "aclnnLtTensorGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -878,16 +1341,28 @@ at::Tensor GeTensorKernelAscend(const at::Tensor& self, const at::Tensor& other) ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options().dtype(at::kBool)); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnGeTensor, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnGeTensor", "aclnnGeTensorGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -900,16 +1375,28 @@ at::Tensor LogicalAndKernelAscend(const at::Tensor& self, const at::Tensor& othe ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options().dtype(at::kBool)); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnLogicalAnd, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnLogicalAnd", "aclnnLogicalAndGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -922,16 +1409,28 @@ at::Tensor LogicalOrKernelAscend(const at::Tensor& self, const at::Tensor& other ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options().dtype(at::kBool)); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnLogicalOr, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnLogicalOr", "aclnnLogicalOrGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -969,6 +1468,22 @@ at::Tensor SubScalarKernelAscend(const at::Tensor& self, const at::Scalar& other REGISTER_IMPL_TO_DISPATCHER(SubScalarFn, sub_scalar_dispatcher, Backend::kAscend, SubScalarKernelAscend) +at::Tensor RsubScalarKernelAscend(const at::Tensor& self, const at::Scalar& other, const at::Scalar& alpha) { + namespace ascend = at::native::flagos::ascend; + auto out = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options()); + + ascend::AclTensorWrapper acl_self(self); + ascend::AclScalarWrapper acl_other(other, self.scalar_type()); + ascend::AclScalarWrapper acl_alpha(alpha, self.scalar_type()); + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD(aclnnRsubs, acl_self.get(), acl_other.get(), acl_alpha.get(), acl_out.get()); + return out; +} + +REGISTER_IMPL_TO_DISPATCHER(RsubScalarFn, rsub_scalar_dispatcher, Backend::kAscend, RsubScalarKernelAscend) + at::Tensor EqScalarKernelAscend(const at::Tensor& self, const at::Scalar& other) { namespace ascend = at::native::flagos::ascend; auto out = ascend::OpPreparation::apply_tensor_without_format( @@ -1141,11 +1656,14 @@ REGISTER_IMPL_TO_DISPATCHER(AnyDimFn, any_dim_dispatcher, Backend::kAscend, AnyD at::Tensor CumsumKernelAscend(const at::Tensor& self, int64_t dim, ::std::optional dtype) { namespace ascend = at::native::flagos::ascend; int64_t d = dim < 0 ? dim + self.dim() : dim; - auto out_dtype = dtype.value_or(self.scalar_type()); + auto out_dtype = dtype.value_or( + at::isIntegralType(self.scalar_type(), /*includeBool=*/true) + ? at::kLong : self.scalar_type()); + auto in = self.scalar_type() == out_dtype ? self : self.to(out_dtype); auto out = ascend::OpPreparation::apply_tensor_without_format( - self.sizes(), self.options().dtype(out_dtype)); + in.sizes(), in.options().dtype(out_dtype)); - ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_self(in); ascend::AclTensorWrapper acl_out(out); aclDataType acl_dtype = ascend::ToAclDataType(out_dtype); @@ -1158,11 +1676,14 @@ REGISTER_IMPL_TO_DISPATCHER(CumsumFn, cumsum_dispatcher, Backend::kAscend, Cumsu at::Tensor CumprodKernelAscend(const at::Tensor& self, int64_t dim, ::std::optional dtype) { namespace ascend = at::native::flagos::ascend; int64_t d = dim < 0 ? dim + self.dim() : dim; - auto out_dtype = dtype.value_or(self.scalar_type()); + auto out_dtype = dtype.value_or( + at::isIntegralType(self.scalar_type(), /*includeBool=*/true) + ? at::kLong : self.scalar_type()); + auto in = self.scalar_type() == out_dtype ? self : self.to(out_dtype); auto out = ascend::OpPreparation::apply_tensor_without_format( - self.sizes(), self.options().dtype(out_dtype)); + in.sizes(), in.options().dtype(out_dtype)); - ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_self(in); ascend::AclScalarWrapper acl_dim(at::Scalar(d), at::kLong); ascend::AclTensorWrapper acl_out(out); aclDataType acl_dtype = ascend::ToAclDataType(out_dtype); @@ -1191,12 +1712,19 @@ at::Tensor LeakyReluKernelAscend(const at::Tensor& self, const at::Scalar& s) { namespace ascend = at::native::flagos::ascend; auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - - ascend::AclTensorWrapper acl_self(self); ascend::AclScalarWrapper acl_s(s, self.scalar_type()); - ascend::AclTensorWrapper acl_out(out); - EXEC_ASCEND_CMD(aclnnLeakyRelu, acl_self.get(), acl_s.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + { double sv = s.toDouble(); hsh.val(sv); } + ascend::ExecAscendCached( + "aclnnLeakyRelu", "aclnnLeakyReluGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_s.get(), out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -1206,12 +1734,19 @@ at::Tensor ClampMinKernelAscend(const at::Tensor& self, const at::Scalar& s) { namespace ascend = at::native::flagos::ascend; auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - - ascend::AclTensorWrapper acl_self(self); ascend::AclScalarWrapper acl_s(s, self.scalar_type()); - ascend::AclTensorWrapper acl_out(out); - EXEC_ASCEND_CMD(aclnnClampMin, acl_self.get(), acl_s.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + { double sv = s.toDouble(); hsh.val(sv); } + ascend::ExecAscendCached( + "aclnnClampMin", "aclnnClampMinGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_s.get(), out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -1221,12 +1756,19 @@ at::Tensor ClampMaxKernelAscend(const at::Tensor& self, const at::Scalar& s) { namespace ascend = at::native::flagos::ascend; auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - - ascend::AclTensorWrapper acl_self(self); ascend::AclScalarWrapper acl_s(s, self.scalar_type()); - ascend::AclTensorWrapper acl_out(out); - EXEC_ASCEND_CMD(aclnnClampMax, acl_self.get(), acl_s.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + { double sv = s.toDouble(); hsh.val(sv); } + ascend::ExecAscendCached( + "aclnnClampMax", "aclnnClampMaxGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_s.get(), out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -1236,12 +1778,19 @@ at::Tensor FmodScalarKernelAscend(const at::Tensor& self, const at::Scalar& s) { namespace ascend = at::native::flagos::ascend; auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - - ascend::AclTensorWrapper acl_self(self); ascend::AclScalarWrapper acl_s(s, self.scalar_type()); - ascend::AclTensorWrapper acl_out(out); - EXEC_ASCEND_CMD(aclnnFmodScalar, acl_self.get(), acl_s.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + { double sv = s.toDouble(); hsh.val(sv); } + ascend::ExecAscendCached( + "aclnnFmodScalar", "aclnnFmodScalarGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_s.get(), out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -1251,12 +1800,19 @@ at::Tensor PowTensorScalarKernelAscend(const at::Tensor& self, const at::Scalar& namespace ascend = at::native::flagos::ascend; auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - - ascend::AclTensorWrapper acl_self(self); ascend::AclScalarWrapper acl_s(s, self.scalar_type()); - ascend::AclTensorWrapper acl_out(out); - EXEC_ASCEND_CMD(aclnnPowTensorScalar, acl_self.get(), acl_s.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + { double sv = s.toDouble(); hsh.val(sv); } + ascend::ExecAscendCached( + "aclnnPowTensorScalar", "aclnnPowTensorScalarGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_s.get(), out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -1391,16 +1947,28 @@ at::Tensor FmodTensorKernelAscend(const at::Tensor& self, const at::Tensor& othe ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnFmodTensor, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnFmodTensor", "aclnnFmodTensorGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -1413,16 +1981,28 @@ at::Tensor FloorDivideKernelAscend(const at::Tensor& self, const at::Tensor& oth ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options()); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnFloorDivide, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnFloorDivide", "aclnnFloorDivideGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -1435,16 +2015,28 @@ at::Tensor LogicalXorKernelAscend(const at::Tensor& self, const at::Tensor& othe ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; auto out = ascend::OpPreparation::apply_tensor_without_format( out_shape, self.options().dtype(at::kBool)); - ascend::AclTensorWrapper acl_self(self_b); - ascend::AclTensorWrapper acl_other(other_b); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnLogicalXor, acl_self.get(), acl_other.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "aclnnLogicalXor", "aclnnLogicalXorGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self_b, &other_b}, {&out}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -1559,12 +2151,19 @@ at::Tensor CeluKernelAscend(const at::Tensor& self, const at::Scalar& s) { namespace ascend = at::native::flagos::ascend; auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - - ascend::AclTensorWrapper acl_self(self); ascend::AclScalarWrapper acl_s(s, self.scalar_type()); - ascend::AclTensorWrapper acl_out(out); - EXEC_ASCEND_CMD(aclnnCelu, acl_self.get(), acl_s.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + { double sv = s.toDouble(); hsh.val(sv); } + ascend::ExecAscendCached( + "aclnnCelu", "aclnnCeluGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_s.get(), out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -1574,12 +2173,19 @@ at::Tensor SoftshrinkKernelAscend(const at::Tensor& self, const at::Scalar& s) { namespace ascend = at::native::flagos::ascend; auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); + ascend::AclScalarWrapper acl_s(s, self.scalar_type()); - ascend::AclTensorWrapper acl_self(self); - ascend::AclScalarWrapper acl_s(s, self.scalar_type()); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnSoftshrink, acl_self.get(), acl_s.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + { double sv = s.toDouble(); hsh.val(sv); } + ascend::ExecAscendCached( + "aclnnSoftshrink", "aclnnSoftshrinkGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_s.get(), out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -1589,12 +2195,19 @@ at::Tensor HardshrinkKernelAscend(const at::Tensor& self, const at::Scalar& s) { namespace ascend = at::native::flagos::ascend; auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options()); - - ascend::AclTensorWrapper acl_self(self); ascend::AclScalarWrapper acl_s(s, self.scalar_type()); - ascend::AclTensorWrapper acl_out(out); - EXEC_ASCEND_CMD(aclnnHardshrink, acl_self.get(), acl_s.get(), acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + { double sv = s.toDouble(); hsh.val(sv); } + ascend::ExecAscendCached( + "aclnnHardshrink", "aclnnHardshrinkGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_s.get(), out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -1848,6 +2461,45 @@ at::Tensor CatKernelAscend(const at::ITensorListRef& tensors, int64_t dim) { REGISTER_IMPL_TO_DISPATCHER(CatFn, cat_dispatcher, Backend::kAscend, CatKernelAscend) +at::Tensor StackKernelAscend(at::TensorList tensors, int64_t dim) { + namespace ascend = at::native::flagos::ascend; + TORCH_CHECK(!tensors.empty(), "stack: expected a non-empty list of tensors"); + + auto& first = tensors[0]; + int64_t out_ndim = first.dim() + 1; + if (dim < 0) dim += out_ndim; + + std::vector out_sizes(first.sizes().begin(), first.sizes().end()); + out_sizes.insert(out_sizes.begin() + dim, static_cast(tensors.size())); + + auto out = ascend::OpPreparation::apply_tensor_without_format( + out_sizes, first.options()); + + std::vector wrappers; + wrappers.reserve(tensors.size()); + for (const auto& t : tensors) { + wrappers.emplace_back(t); + } + + std::vector acl_tensors; + acl_tensors.reserve(tensors.size()); + for (auto& w : wrappers) { + acl_tensors.push_back(w.get()); + } + + aclTensorList* tensor_list = aclCreateTensorList( + acl_tensors.data(), acl_tensors.size()); + + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD(aclnnStack, tensor_list, dim, acl_out.get()); + + (void)tensor_list; // aclTensor* owned by wrappers; do not aclDestroyTensorList + return out; +} + +REGISTER_IMPL_TO_DISPATCHER(StackFn, stack_dispatcher, Backend::kAscend, StackKernelAscend) + at::Tensor ZerosKernelAscend(at::IntArrayRef size, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory) { auto options = at::TensorOptions() .dtype(dtype.value_or(at::kFloat)) @@ -1861,6 +2513,19 @@ at::Tensor ZerosKernelAscend(at::IntArrayRef size, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory) { + auto options = at::TensorOptions() + .dtype(dtype.value_or(at::kFloat)) + .layout(layout.value_or(at::kStrided)) + .device(device.value_or(at::Device(at::kPrivateUse1, 0))) + .pinned_memory(pin_memory.value_or(false)); + auto result = at::empty(size, options); + result.fill_(1); + return result; +} + +REGISTER_IMPL_TO_DISPATCHER(OnesFn, ones_dispatcher, Backend::kAscend, OnesKernelAscend) + at::Tensor ScalarTensorKernelAscend(const at::Scalar& s, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory) { auto options = at::TensorOptions() .dtype(dtype.value_or(at::ScalarType::Float)) @@ -1891,6 +2556,68 @@ at::Tensor OnesLikeKernelAscend(const at::Tensor& self, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory, ::std::optional memory_format) { + auto options = at::TensorOptions() + .dtype(dtype.value_or(self.scalar_type())) + .layout(layout.value_or(self.layout())) + .device(device.value_or(self.device())) + .pinned_memory(pin_memory.value_or(false)); + auto fmt = memory_format.value_or(at::MemoryFormat::Contiguous); + if (fmt == at::MemoryFormat::Preserve) { + fmt = self.suggest_memory_format(); + } + auto result = at::empty(self.sizes(), options, fmt); + result.zero_(); + return result; +} + +REGISTER_IMPL_TO_DISPATCHER(ZerosLikeFn, zeros_like_dispatcher, Backend::kAscend, ZerosLikeKernelAscend) + +at::Tensor EmptyLikeKernelAscend(const at::Tensor& self, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory, ::std::optional memory_format) { + auto options = at::TensorOptions() + .dtype(dtype.value_or(self.scalar_type())) + .layout(layout.value_or(self.layout())) + .device(device.value_or(self.device())) + .pinned_memory(pin_memory.value_or(false)); + auto fmt = memory_format.value_or(at::MemoryFormat::Preserve); + if (fmt == at::MemoryFormat::Preserve) { + fmt = self.suggest_memory_format(); + } + return at::empty(self.sizes(), options, fmt); +} + +REGISTER_IMPL_TO_DISPATCHER(EmptyLikeFn, empty_like_dispatcher, Backend::kAscend, EmptyLikeKernelAscend) + +at::Tensor FullKernelAscend(at::IntArrayRef size, const at::Scalar& fill, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory) { + auto options = at::TensorOptions() + .dtype(dtype.value_or(at::kFloat)) + .layout(layout.value_or(at::kStrided)) + .device(device.value_or(at::Device(at::kPrivateUse1, 0))) + .pinned_memory(pin_memory.value_or(false)); + auto result = at::empty(size, options); + result.fill_(fill); + return result; +} + +REGISTER_IMPL_TO_DISPATCHER(FullFn, full_dispatcher, Backend::kAscend, FullKernelAscend) + +at::Tensor FullLikeKernelAscend(const at::Tensor& self, const at::Scalar& fill, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory, ::std::optional memory_format) { + auto options = at::TensorOptions() + .dtype(dtype.value_or(self.scalar_type())) + .layout(layout.value_or(self.layout())) + .device(device.value_or(self.device())) + .pinned_memory(pin_memory.value_or(false)); + auto fmt = memory_format.value_or(at::MemoryFormat::Preserve); + if (fmt == at::MemoryFormat::Preserve) { + fmt = self.suggest_memory_format(); + } + auto result = at::empty(self.sizes(), options, fmt); + result.fill_(fill); + return result; +} + +REGISTER_IMPL_TO_DISPATCHER(FullLikeFn, full_like_dispatcher, Backend::kAscend, FullLikeKernelAscend) + at::Tensor NewOnesKernelAscend(const at::Tensor& self, at::IntArrayRef size, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory) { auto options = at::TensorOptions() .dtype(dtype.value_or(self.scalar_type())) @@ -2328,6 +3055,446 @@ at::Tensor& FillInplaceTensorKernelAscend(at::Tensor& self, const at::Tensor& va REGISTER_IMPL_TO_DISPATCHER(FillInplaceTensorFn, fill_inplace_tensor_dispatcher, Backend::kAscend, FillInplaceTensorKernelAscend) +at::Tensor& AddInplaceTensorKernelAscend(at::Tensor& self, const at::Tensor& other, const at::Scalar& alpha) { + namespace ascend = at::native::flagos::ascend; + auto other_c = other.is_privateuseone() + ? (other.scalar_type() == self.scalar_type() ? other : other.to(self.scalar_type())) + : other.to(self.options()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_other(other_c); + ascend::AclScalarWrapper acl_alpha(alpha, self.scalar_type()); + EXEC_ASCEND_CMD(aclnnInplaceAdd, const_cast(acl_self.get()), acl_other.get(), + acl_alpha.get()); + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(AddInplaceTensorFn, add_inplace_tensor_dispatcher, Backend::kAscend, AddInplaceTensorKernelAscend) + +at::Tensor& AddInplaceScalarKernelAscend(at::Tensor& self, const at::Scalar& other, const at::Scalar& alpha) { + namespace ascend = at::native::flagos::ascend; + ascend::AclTensorWrapper acl_self(self); + ascend::AclScalarWrapper acl_other(other, self.scalar_type()); + ascend::AclScalarWrapper acl_alpha(alpha, self.scalar_type()); + EXEC_ASCEND_CMD(aclnnInplaceAdds, const_cast(acl_self.get()), acl_other.get(), + acl_alpha.get()); + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(AddInplaceScalarFn, add_inplace_scalar_dispatcher, Backend::kAscend, AddInplaceScalarKernelAscend) + +at::Tensor& MulInplaceTensorKernelAscend(at::Tensor& self, const at::Tensor& other) { + namespace ascend = at::native::flagos::ascend; + auto other_c = other.is_privateuseone() + ? (other.scalar_type() == self.scalar_type() ? other : other.to(self.scalar_type())) + : other.to(self.options()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_other(other_c); + EXEC_ASCEND_CMD(aclnnInplaceMul, const_cast(acl_self.get()), acl_other.get()); + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(MulInplaceTensorFn, mul_inplace_tensor_dispatcher, Backend::kAscend, MulInplaceTensorKernelAscend) + +at::Tensor& MulInplaceScalarKernelAscend(at::Tensor& self, const at::Scalar& other) { + namespace ascend = at::native::flagos::ascend; + ascend::AclTensorWrapper acl_self(self); + ascend::AclScalarWrapper acl_other(other, self.scalar_type()); + EXEC_ASCEND_CMD(aclnnInplaceMuls, const_cast(acl_self.get()), acl_other.get()); + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(MulInplaceScalarFn, mul_inplace_scalar_dispatcher, Backend::kAscend, MulInplaceScalarKernelAscend) + +at::Tensor& DivInplaceTensorKernelAscend(at::Tensor& self, const at::Tensor& other) { + namespace ascend = at::native::flagos::ascend; + auto other_c = other.is_privateuseone() + ? (other.scalar_type() == self.scalar_type() ? other : other.to(self.scalar_type())) + : other.to(self.options()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_other(other_c); + EXEC_ASCEND_CMD(aclnnInplaceDiv, const_cast(acl_self.get()), acl_other.get()); + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(DivInplaceTensorFn, div_inplace_tensor_dispatcher, Backend::kAscend, DivInplaceTensorKernelAscend) + +at::Tensor& BitwiseAndInplaceTensorKernelAscend(at::Tensor& self, const at::Tensor& other) { + namespace ascend = at::native::flagos::ascend; + auto other_c = other.is_privateuseone() + ? (other.scalar_type() == self.scalar_type() ? other : other.to(self.scalar_type())) + : other.to(self.options()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_other(other_c); + EXEC_ASCEND_CMD(aclnnInplaceBitwiseAndTensor, const_cast(acl_self.get()), acl_other.get()); + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(BitwiseAndInplaceTensorFn, bitwise_and_inplace_tensor_dispatcher, Backend::kAscend, BitwiseAndInplaceTensorKernelAscend) + +at::Tensor& BitwiseOrInplaceTensorKernelAscend(at::Tensor& self, const at::Tensor& other) { + namespace ascend = at::native::flagos::ascend; + auto other_c = other.is_privateuseone() + ? (other.scalar_type() == self.scalar_type() ? other : other.to(self.scalar_type())) + : other.to(self.options()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_other(other_c); + EXEC_ASCEND_CMD(aclnnInplaceBitwiseOrTensor, const_cast(acl_self.get()), acl_other.get()); + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(BitwiseOrInplaceTensorFn, bitwise_or_inplace_tensor_dispatcher, Backend::kAscend, BitwiseOrInplaceTensorKernelAscend) + +at::Tensor& BitwiseXorInplaceTensorKernelAscend(at::Tensor& self, const at::Tensor& other) { + namespace ascend = at::native::flagos::ascend; + auto other_c = other.is_privateuseone() + ? (other.scalar_type() == self.scalar_type() ? other : other.to(self.scalar_type())) + : other.to(self.options()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_other(other_c); + EXEC_ASCEND_CMD(aclnnInplaceBitwiseXorTensor, const_cast(acl_self.get()), acl_other.get()); + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(BitwiseXorInplaceTensorFn, bitwise_xor_inplace_tensor_dispatcher, Backend::kAscend, BitwiseXorInplaceTensorKernelAscend) + +at::Tensor& AddcmulInplaceKernelAscend(at::Tensor& self, const at::Tensor& tensor1, const at::Tensor& tensor2, const at::Scalar& value) { + namespace ascend = at::native::flagos::ascend; + auto t1 = tensor1.scalar_type() == self.scalar_type() ? tensor1 : tensor1.to(self.scalar_type()); + auto t2 = tensor2.scalar_type() == self.scalar_type() ? tensor2 : tensor2.to(self.scalar_type()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_t1(t1); + ascend::AclTensorWrapper acl_t2(t2); + ascend::AclScalarWrapper acl_value(value, self.scalar_type()); + EXEC_ASCEND_CMD(aclnnInplaceAddcmul, const_cast(acl_self.get()), acl_t1.get(), + acl_t2.get(), acl_value.get()); + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(AddcmulInplaceFn, addcmul_inplace_dispatcher, Backend::kAscend, AddcmulInplaceKernelAscend) + +at::Tensor& AddcdivInplaceKernelAscend(at::Tensor& self, const at::Tensor& tensor1, const at::Tensor& tensor2, const at::Scalar& value) { + namespace ascend = at::native::flagos::ascend; + auto t1 = tensor1.scalar_type() == self.scalar_type() ? tensor1 : tensor1.to(self.scalar_type()); + auto t2 = tensor2.scalar_type() == self.scalar_type() ? tensor2 : tensor2.to(self.scalar_type()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_t1(t1); + ascend::AclTensorWrapper acl_t2(t2); + ascend::AclScalarWrapper acl_value(value, self.scalar_type()); + EXEC_ASCEND_CMD(aclnnInplaceAddcdiv, const_cast(acl_self.get()), acl_t1.get(), + acl_t2.get(), acl_value.get()); + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(AddcdivInplaceFn, addcdiv_inplace_dispatcher, Backend::kAscend, AddcdivInplaceKernelAscend) + +at::Tensor& SqrtInplaceKernelAscend(at::Tensor& self) { + namespace ascend = at::native::flagos::ascend; + ascend::AclTensorWrapper acl_self(self); + EXEC_ASCEND_CMD(aclnnInplaceSqrt, const_cast(acl_self.get())); + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(SqrtInplaceFn, sqrt_inplace_dispatcher, Backend::kAscend, SqrtInplaceKernelAscend) + +at::Tensor& LerpInplaceScalarKernelAscend(at::Tensor& self, const at::Tensor& end, const at::Scalar& weight) { + namespace ascend = at::native::flagos::ascend; + auto end_c = end.scalar_type() == self.scalar_type() ? end : end.to(self.scalar_type()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_end(end_c); + ascend::AclScalarWrapper acl_weight(weight, self.scalar_type()); + EXEC_ASCEND_CMD(aclnnInplaceLerps, const_cast(acl_self.get()), acl_end.get(), + acl_weight.get()); + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(LerpInplaceScalarFn, lerp_inplace_scalar_dispatcher, Backend::kAscend, LerpInplaceScalarKernelAscend) + +static void ForeachMulInplaceScalarKernelAscendChunk(at::TensorList self, const at::Scalar& scalar) { + namespace ascend = at::native::flagos::ascend; + + std::vector wrappers; + wrappers.reserve(self.size()); + for (const auto& t : self) { + wrappers.emplace_back(t); + } + std::vector acl_tensors; + acl_tensors.reserve(self.size()); + for (auto& w : wrappers) { + acl_tensors.push_back(w.get()); + } + aclTensorList* tensor_list = aclCreateTensorList(acl_tensors.data(), acl_tensors.size()); + // aic-ops-info: ForeachMulScalar/ForeachAddScalar's `scalar` dtype tracks x's + // EXCEPT bf16 x, which requires a float32 scalar (no bf16 scalar entry). + auto scalar_dtype = self[0].scalar_type() == at::kBFloat16 ? at::kFloat : self[0].scalar_type(); + ascend::AclScalarWrapper acl_scalar(scalar, scalar_dtype); + + EXEC_ASCEND_CMD(aclnnForeachMulScalarV2, tensor_list, acl_scalar.get(), tensor_list); + + (void)tensor_list; // aclTensor* owned by wrappers; do not aclDestroyTensorList +} + +void ForeachMulInplaceScalarKernelAscend(at::TensorList self, const at::Scalar& scalar) { + TORCH_CHECK(!self.empty(), "foreach_mul_inplace_scalar_dispatcher: expected a non-empty list of tensors"); + for (size_t off = 0; off < self.size(); off += 32) { + size_t n = std::min(32, self.size() - off); + ForeachMulInplaceScalarKernelAscendChunk(self.slice(off, n), scalar); + } +} + +REGISTER_IMPL_TO_DISPATCHER(ForeachMulInplaceScalarFn, foreach_mul_inplace_scalar_dispatcher, Backend::kAscend, ForeachMulInplaceScalarKernelAscend) + +static void ForeachAddInplaceScalarKernelAscendChunk(at::TensorList self, const at::Scalar& scalar) { + namespace ascend = at::native::flagos::ascend; + + std::vector wrappers; + wrappers.reserve(self.size()); + for (const auto& t : self) { + wrappers.emplace_back(t); + } + std::vector acl_tensors; + acl_tensors.reserve(self.size()); + for (auto& w : wrappers) { + acl_tensors.push_back(w.get()); + } + aclTensorList* tensor_list = aclCreateTensorList(acl_tensors.data(), acl_tensors.size()); + // aic-ops-info: ForeachMulScalar/ForeachAddScalar's `scalar` dtype tracks x's + // EXCEPT bf16 x, which requires a float32 scalar (no bf16 scalar entry). + auto scalar_dtype = self[0].scalar_type() == at::kBFloat16 ? at::kFloat : self[0].scalar_type(); + ascend::AclScalarWrapper acl_scalar(scalar, scalar_dtype); + + EXEC_ASCEND_CMD(aclnnForeachAddScalarV2, tensor_list, acl_scalar.get(), tensor_list); + + (void)tensor_list; // aclTensor* owned by wrappers; do not aclDestroyTensorList +} + +void ForeachAddInplaceScalarKernelAscend(at::TensorList self, const at::Scalar& scalar) { + TORCH_CHECK(!self.empty(), "foreach_add_inplace_scalar_dispatcher: expected a non-empty list of tensors"); + for (size_t off = 0; off < self.size(); off += 32) { + size_t n = std::min(32, self.size() - off); + ForeachAddInplaceScalarKernelAscendChunk(self.slice(off, n), scalar); + } +} + +REGISTER_IMPL_TO_DISPATCHER(ForeachAddInplaceScalarFn, foreach_add_inplace_scalar_dispatcher, Backend::kAscend, ForeachAddInplaceScalarKernelAscend) + +static void ForeachLerpInplaceScalarKernelAscendChunk(at::TensorList self, at::TensorList tensors1, const at::Scalar& weight) { + namespace ascend = at::native::flagos::ascend; + + std::vector self_w, t1_w; + self_w.reserve(self.size()); + t1_w.reserve(tensors1.size()); + for (const auto& t : self) self_w.emplace_back(t); + for (const auto& t : tensors1) t1_w.emplace_back(t); + + std::vector self_ptrs, t1_ptrs; + self_ptrs.reserve(self.size()); + t1_ptrs.reserve(tensors1.size()); + for (auto& w : self_w) self_ptrs.push_back(w.get()); + for (auto& w : t1_w) t1_ptrs.push_back(w.get()); + + aclTensorList* self_list = aclCreateTensorList(self_ptrs.data(), self_ptrs.size()); + aclTensorList* t1_list = aclCreateTensorList(t1_ptrs.data(), t1_ptrs.size()); + // aic-ops-info: ForeachLerpScalar's `weight` is ALWAYS float32, regardless + // of x1/x2's dtype (unlike mul_/add_.Scalar, which track x except for bf16). + ascend::AclScalarWrapper acl_weight(weight, at::kFloat); + + EXEC_ASCEND_CMD(aclnnForeachLerpScalar, self_list, t1_list, acl_weight.get(), self_list); + + (void)self_list; (void)t1_list; // owned by *_w; do not aclDestroyTensorList +} + +void ForeachLerpInplaceScalarKernelAscend(at::TensorList self, at::TensorList tensors1, const at::Scalar& weight) { + TORCH_CHECK(!self.empty(), "foreach_lerp_inplace_scalar_dispatcher: expected a non-empty list of tensors"); + TORCH_CHECK(self.size() == tensors1.size(), "foreach_lerp_inplace_scalar_dispatcher: tensor lists must match in length"); + for (size_t off = 0; off < self.size(); off += 32) { + size_t n = std::min(32, self.size() - off); + ForeachLerpInplaceScalarKernelAscendChunk(self.slice(off, n), tensors1.slice(off, n), weight); + } +} + +REGISTER_IMPL_TO_DISPATCHER(ForeachLerpInplaceScalarFn, foreach_lerp_inplace_scalar_dispatcher, Backend::kAscend, ForeachLerpInplaceScalarKernelAscend) + +static void ForeachAddcmulInplaceScalarKernelAscendChunk(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Scalar& value) { + namespace ascend = at::native::flagos::ascend; + + std::vector self_w, t1_w, t2_w; + self_w.reserve(self.size()); + t1_w.reserve(tensor1.size()); + t2_w.reserve(tensor2.size()); + for (const auto& t : self) self_w.emplace_back(t); + for (const auto& t : tensor1) t1_w.emplace_back(t); + for (const auto& t : tensor2) t2_w.emplace_back(t); + + std::vector self_ptrs, t1_ptrs, t2_ptrs; + self_ptrs.reserve(self.size()); + t1_ptrs.reserve(tensor1.size()); + t2_ptrs.reserve(tensor2.size()); + for (auto& w : self_w) self_ptrs.push_back(w.get()); + for (auto& w : t1_w) t1_ptrs.push_back(w.get()); + for (auto& w : t2_w) t2_ptrs.push_back(w.get()); + + aclTensorList* self_list = aclCreateTensorList(self_ptrs.data(), self_ptrs.size()); + aclTensorList* t1_list = aclCreateTensorList(t1_ptrs.data(), t1_ptrs.size()); + aclTensorList* t2_list = aclCreateTensorList(t2_ptrs.data(), t2_ptrs.size()); + // aic-ops-info: ForeachAddcmulScalar's `scalar` dtype tracks x EXCEPT bf16 x, + // which requires a float32 scalar (same rule as mul_/add_.Scalar). + auto value_dtype = self[0].scalar_type() == at::kBFloat16 ? at::kFloat : self[0].scalar_type(); + ascend::AclScalarWrapper acl_value(value, value_dtype); + + EXEC_ASCEND_CMD(aclnnForeachAddcmulScalarV2, self_list, t1_list, t2_list, acl_value.get(), self_list); + + (void)self_list; (void)t1_list; (void)t2_list; // owned by *_w +} + +void ForeachAddcmulInplaceScalarKernelAscend(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Scalar& value) { + TORCH_CHECK(!self.empty(), "foreach_addcmul_inplace_scalar_dispatcher: expected a non-empty list of tensors"); + TORCH_CHECK(self.size() == tensor1.size() && self.size() == tensor2.size(), + "foreach_addcmul_inplace_scalar_dispatcher: tensor lists must match in length"); + for (size_t off = 0; off < self.size(); off += 32) { + size_t n = std::min(32, self.size() - off); + ForeachAddcmulInplaceScalarKernelAscendChunk(self.slice(off, n), tensor1.slice(off, n), tensor2.slice(off, n), value); + } +} + +REGISTER_IMPL_TO_DISPATCHER(ForeachAddcmulInplaceScalarFn, foreach_addcmul_inplace_scalar_dispatcher, Backend::kAscend, ForeachAddcmulInplaceScalarKernelAscend) + +static void ForeachSqrtKernelAscendChunk(at::TensorList self, at::TensorList outs) { + namespace ascend = at::native::flagos::ascend; + + + std::vector in_w, out_w; + in_w.reserve(self.size()); + out_w.reserve(outs.size()); + for (const auto& t : self) in_w.emplace_back(t); + for (const auto& t : outs) out_w.emplace_back(t); + + std::vector in_ptrs, out_ptrs; + in_ptrs.reserve(self.size()); + out_ptrs.reserve(outs.size()); + for (auto& w : in_w) in_ptrs.push_back(w.get()); + for (auto& w : out_w) out_ptrs.push_back(w.get()); + + aclTensorList* in_list = aclCreateTensorList(in_ptrs.data(), in_ptrs.size()); + aclTensorList* out_list = aclCreateTensorList(out_ptrs.data(), out_ptrs.size()); + + EXEC_ASCEND_CMD(aclnnForeachSqrt, in_list, out_list); + + (void)in_list; (void)out_list; // owned by in_w/out_w +} + +::std::vector ForeachSqrtKernelAscend(at::TensorList self) { + TORCH_CHECK(!self.empty(), "foreach_sqrt_dispatcher: expected a non-empty list of tensors"); + std::vector outs; + outs.reserve(self.size()); + for (const auto& t : self) outs.push_back(at::empty_like(t)); + for (size_t off = 0; off < self.size(); off += 32) { + size_t n = std::min(32, self.size() - off); + ForeachSqrtKernelAscendChunk(self.slice(off, n), at::TensorList(outs).slice(off, n)); + } + return outs; +} + +REGISTER_IMPL_TO_DISPATCHER(ForeachSqrtFn, foreach_sqrt_dispatcher, Backend::kAscend, ForeachSqrtKernelAscend) + +static void ForeachDivInplaceScalarlistKernelAscendChunk(at::TensorList self, at::ArrayRef scalars) { + namespace ascend = at::native::flagos::ascend; + + std::vector wrappers; + wrappers.reserve(self.size()); + for (const auto& t : self) wrappers.emplace_back(t); + std::vector acl_tensors; + acl_tensors.reserve(self.size()); + for (auto& w : wrappers) acl_tensors.push_back(w.get()); + aclTensorList* tensor_list = aclCreateTensorList(acl_tensors.data(), acl_tensors.size()); + + // aic-ops-info: ForeachDivScalarList's `scalars` is ALWAYS float32, + // regardless of x's dtype (same rule as ForeachLerpScalar's weight). + std::vector scalar_wrappers; + scalar_wrappers.reserve(scalars.size()); + for (size_t i = 0; i < scalars.size(); ++i) { + scalar_wrappers.emplace_back(scalars[i], at::kFloat); + } + std::vector acl_scalars; + acl_scalars.reserve(scalar_wrappers.size()); + for (auto& sw : scalar_wrappers) acl_scalars.push_back(sw.get()); + aclScalarList* scalar_list = aclCreateScalarList(acl_scalars.data(), acl_scalars.size()); + + EXEC_ASCEND_CMD(aclnnForeachDivScalarList, tensor_list, scalar_list, tensor_list); + + aclDestroyScalarList(scalar_list); + (void)tensor_list; // aclTensor* owned by wrappers; do not aclDestroyTensorList +} + +void ForeachDivInplaceScalarlistKernelAscend(at::TensorList self, at::ArrayRef scalars) { + TORCH_CHECK(!self.empty(), "foreach_div_inplace_scalarlist_dispatcher: expected a non-empty list of tensors"); + TORCH_CHECK(self.size() == scalars.size(), "foreach_div_inplace_scalarlist_dispatcher: scalars must match tensor list length"); + for (size_t off = 0; off < self.size(); off += 32) { + size_t n = std::min(32, self.size() - off); + ForeachDivInplaceScalarlistKernelAscendChunk(self.slice(off, n), scalars.slice(off, n)); + } +} + +REGISTER_IMPL_TO_DISPATCHER(ForeachDivInplaceScalarlistFn, foreach_div_inplace_scalarlist_dispatcher, Backend::kAscend, ForeachDivInplaceScalarlistKernelAscend) + +static void ForeachAddcdivInplaceScalarlistKernelAscendChunk(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, at::ArrayRef scalars) { + namespace ascend = at::native::flagos::ascend; + + std::vector self_w, t1_w, t2_w; + self_w.reserve(self.size()); + t1_w.reserve(tensor1.size()); + t2_w.reserve(tensor2.size()); + for (const auto& t : self) self_w.emplace_back(t); + for (const auto& t : tensor1) t1_w.emplace_back(t); + for (const auto& t : tensor2) t2_w.emplace_back(t); + + std::vector self_ptrs, t1_ptrs, t2_ptrs; + self_ptrs.reserve(self.size()); + t1_ptrs.reserve(tensor1.size()); + t2_ptrs.reserve(tensor2.size()); + for (auto& w : self_w) self_ptrs.push_back(w.get()); + for (auto& w : t1_w) t1_ptrs.push_back(w.get()); + for (auto& w : t2_w) t2_ptrs.push_back(w.get()); + + aclTensorList* self_list = aclCreateTensorList(self_ptrs.data(), self_ptrs.size()); + aclTensorList* t1_list = aclCreateTensorList(t1_ptrs.data(), t1_ptrs.size()); + aclTensorList* t2_list = aclCreateTensorList(t2_ptrs.data(), t2_ptrs.size()); + + // aclnnForeachAddcdivScalarList's "scalars" param is a plain device aclTensor + // (1-D, one element per list entry), NOT an aclScalarList -- unlike div's + // ScalarList variant. Materialize scalars on host in self[0]'s dtype, then + // move to device once. The dtype MUST match self (a float32 scalars tensor + // against fp16 inputs returns 161002), which costs up to 1 ulp versus CPU, + // where the divisor stays a full-precision Scalar. + at::Tensor scalars_cpu = at::empty({static_cast(scalars.size())}, + at::TensorOptions().dtype(self[0].scalar_type())); + AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, self[0].scalar_type(), + "foreach_addcdiv_inplace_scalarlist_dispatcher_scalars", [&] { + auto* ptr = scalars_cpu.data_ptr(); + for (size_t i = 0; i < scalars.size(); ++i) { + ptr[i] = scalars[i].to(); + } + }); + at::Tensor scalars_dev = scalars_cpu.to(self[0].device()); + ascend::AclTensorWrapper acl_scalars(scalars_dev); + + EXEC_ASCEND_CMD(aclnnForeachAddcdivScalarList, self_list, t1_list, t2_list, acl_scalars.get(), self_list); + + (void)self_list; (void)t1_list; (void)t2_list; // owned by *_w +} + +void ForeachAddcdivInplaceScalarlistKernelAscend(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, at::ArrayRef scalars) { + TORCH_CHECK(!self.empty(), "foreach_addcdiv_inplace_scalarlist_dispatcher: expected a non-empty list of tensors"); + TORCH_CHECK(self.size() == tensor1.size() && self.size() == tensor2.size() && self.size() == scalars.size(), + "foreach_addcdiv_inplace_scalarlist_dispatcher: tensor/scalar lists must match in length"); + for (size_t off = 0; off < self.size(); off += 32) { + size_t n = std::min(32, self.size() - off); + ForeachAddcdivInplaceScalarlistKernelAscendChunk(self.slice(off, n), tensor1.slice(off, n), tensor2.slice(off, n), + scalars.slice(off, n)); + } +} + +REGISTER_IMPL_TO_DISPATCHER(ForeachAddcdivInplaceScalarlistFn, foreach_addcdiv_inplace_scalarlist_dispatcher, Backend::kAscend, ForeachAddcdivInplaceScalarlistKernelAscend) + at::Tensor EmbeddingKernelAscend(const at::Tensor& weight, const at::Tensor& indices, int64_t padding_idx, bool scale_grad_by_freq, bool sparse) { namespace ascend = at::native::flagos::ascend; auto out_sizes = indices.sizes().vec(); @@ -2622,10 +3789,17 @@ at::Tensor PrivSoftmaxKernelAscend(const at::Tensor& self, int64_t dim, bool hal auto out = ascend::OpPreparation::apply_tensor_without_format( self.sizes(), self.options().dtype(out_dtype)); - ascend::AclTensorWrapper acl_self(self); - ascend::AclTensorWrapper acl_out(out); - - EXEC_ASCEND_CMD(aclnnSoftmax, acl_self.get(), dim, acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); hsh.val(dim); + { int8_t h2f = half_to_float ? 1 : 0; hsh.val(h2f); } + ascend::ExecAscendCached( + "aclnnSoftmax", "aclnnSoftmaxGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, dim, out_t[0].acl_tensor, pws, pex); + }); return out; } @@ -2650,6 +3824,25 @@ at::Tensor AllKernelAscend(const at::Tensor& self) { REGISTER_IMPL_TO_DISPATCHER(AllFn, all_dispatcher, Backend::kAscend, AllKernelAscend) +at::Tensor AnyKernelAscend(const at::Tensor& self) { + namespace ascend = at::native::flagos::ascend; + auto input = self.contiguous().reshape({-1}); + auto out = ascend::OpPreparation::apply_tensor_without_format( + {}, self.options().dtype(at::kBool)); + + ascend::AclTensorWrapper acl_self(input); + ascend::AclTensorWrapper acl_out(out); + + int64_t dim_val = 0; + std::vector dims{dim_val}; + ascend::AclIntArrayWrapper acl_dim(dims); + + EXEC_ASCEND_CMD(aclnnAny, acl_self.get(), acl_dim.get(), false, acl_out.get()); + return out; +} + +REGISTER_IMPL_TO_DISPATCHER(AnyFn, any_dispatcher, Backend::kAscend, AnyKernelAscend) + at::Tensor SumDimIntlistKernelAscend(const at::Tensor& self, at::OptionalIntArrayRef dim, bool keepdim, std::optional dtype) { namespace ascend = at::native::flagos::ascend; auto out_dtype = dtype.has_value() ? dtype.value() : self.scalar_type(); @@ -2674,12 +3867,73 @@ at::Tensor SumDimIntlistKernelAscend(const at::Tensor& self, at::OptionalIntArra ascend::AclIntArrayWrapper acl_dim(norm_dims); aclDataType acl_dtype = ascend::ToAclDataType(out_dtype); - EXEC_ASCEND_CMD(aclnnReduceSum, acl_self.get(), acl_dim.get(), keepdim, acl_dtype, acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + for (int64_t d : norm_dims) hsh.val(d); + hsh.val(keepdim); + { int32_t dtk = static_cast(acl_dtype); hsh.val(dtk); } + ascend::ExecAscendCached( + "aclnnReduceSum", "aclnnReduceSumGetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_dim.get(), keepdim, acl_dtype, out_t[0].acl_tensor, pws, pex); + }); return out; } REGISTER_IMPL_TO_DISPATCHER(SumDimIntlistFn, sum_dim_intlist_dispatcher, Backend::kAscend, SumDimIntlistKernelAscend) +at::Tensor SumKernelAscend(const at::Tensor& self, std::optional dtype) { + namespace ascend = at::native::flagos::ascend; + // Integral/bool inputs promote to int64 when no dtype given (matches torch). + at::ScalarType out_dtype = dtype.has_value() + ? dtype.value() + : (c10::isIntegralType(self.scalar_type(), /*includeBool=*/true) + ? at::kLong : self.scalar_type()); + int64_t ndim = self.dim(); + std::vector norm_dims; + for (int64_t d = 0; d < ndim; ++d) norm_dims.push_back(d); + auto out = ascend::OpPreparation::apply_tensor_without_format( + {}, self.options().dtype(out_dtype)); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_out(out); + ascend::AclIntArrayWrapper acl_dim(norm_dims); + aclDataType acl_dtype = ascend::ToAclDataType(out_dtype); + + EXEC_ASCEND_CMD(aclnnReduceSum, acl_self.get(), acl_dim.get(), false, acl_dtype, acl_out.get()); + return out; +} + +REGISTER_IMPL_TO_DISPATCHER(SumFn, sum_dispatcher, Backend::kAscend, SumKernelAscend) + +at::Tensor MaxKernelAscend(const at::Tensor& self) { + namespace ascend = at::native::flagos::ascend; + auto out = ascend::OpPreparation::apply_tensor_without_format( + {}, self.options()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD(aclnnMax, acl_self.get(), acl_out.get()); + return out; +} + +REGISTER_IMPL_TO_DISPATCHER(MaxFn, max_dispatcher, Backend::kAscend, MaxKernelAscend) + +at::Tensor MinKernelAscend(const at::Tensor& self) { + namespace ascend = at::native::flagos::ascend; + auto out = ascend::OpPreparation::apply_tensor_without_format( + {}, self.options()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD(aclnnMin, acl_self.get(), acl_out.get()); + return out; +} + +REGISTER_IMPL_TO_DISPATCHER(MinFn, min_dispatcher, Backend::kAscend, MinKernelAscend) + at::Tensor MeanDimKernelAscend(const at::Tensor& self, at::OptionalIntArrayRef dim, bool keepdim, std::optional dtype) { namespace ascend = at::native::flagos::ascend; auto out_dtype = dtype.has_value() ? dtype.value() : self.scalar_type(); @@ -2704,12 +3958,82 @@ at::Tensor MeanDimKernelAscend(const at::Tensor& self, at::OptionalIntArrayRef d ascend::AclIntArrayWrapper acl_dim(norm_dims); auto acl_dtype = static_cast(ascend::ToAclDataType(out_dtype)); - EXEC_ASCEND_CMD(aclnnMeanV2, acl_self.get(), acl_dim.get(), keepdim, acl_dtype, acl_out.get()); + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + for (int64_t d : norm_dims) hsh.val(d); + hsh.val(keepdim); + hsh.val(acl_dtype); + ascend::ExecAscendCached( + "aclnnMeanV2", "aclnnMeanV2GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self}, {&out}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, acl_dim.get(), keepdim, acl_dtype, out_t[0].acl_tensor, pws, pex); + }); return out; } REGISTER_IMPL_TO_DISPATCHER(MeanDimFn, mean_dim_dispatcher, Backend::kAscend, MeanDimKernelAscend) +at::Tensor MeanKernelAscend(const at::Tensor& self, std::optional dtype) { + namespace ascend = at::native::flagos::ascend; + at::ScalarType out_dtype = dtype.has_value() ? dtype.value() : self.scalar_type(); + int64_t ndim = self.dim(); + std::vector norm_dims; + for (int64_t d = 0; d < ndim; ++d) norm_dims.push_back(d); + auto out = ascend::OpPreparation::apply_tensor_without_format( + {}, self.options().dtype(out_dtype)); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_out(out); + ascend::AclIntArrayWrapper acl_dim(norm_dims); + aclDataType acl_dtype = ascend::ToAclDataType(out_dtype); + + EXEC_ASCEND_CMD(aclnnMean, acl_self.get(), acl_dim.get(), false, acl_dtype, acl_out.get()); + return out; +} + +REGISTER_IMPL_TO_DISPATCHER(MeanFn, mean_dispatcher, Backend::kAscend, MeanKernelAscend) + +at::Tensor ClampKernelAscend(const at::Tensor& self, const ::std::optional& min, const ::std::optional& max) { + namespace ascend = at::native::flagos::ascend; + auto out = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options()); + + ascend::AclTensorWrapper acl_self(self); + ascend::AclScalarWrapper acl_min = min.has_value() + ? ascend::AclScalarWrapper(min.value(), self.scalar_type()) + : ascend::AclScalarWrapper(); + ascend::AclScalarWrapper acl_max = max.has_value() + ? ascend::AclScalarWrapper(max.value(), self.scalar_type()) + : ascend::AclScalarWrapper(); + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD(aclnnClamp, acl_self.get(), acl_min.get(), acl_max.get(), acl_out.get()); + return out; +} + +REGISTER_IMPL_TO_DISPATCHER(ClampFn, clamp_dispatcher, Backend::kAscend, ClampKernelAscend) + +at::Tensor ClampTensorKernelAscend(const at::Tensor& self, const ::std::optional& min, const ::std::optional& max) { + namespace ascend = at::native::flagos::ascend; + auto out_shape = self.sizes().vec(); + if (min.has_value()) out_shape = at::infer_size(out_shape, min.value().sizes()); + if (max.has_value()) out_shape = at::infer_size(out_shape, max.value().sizes()); + auto out = ascend::OpPreparation::apply_tensor_without_format( + out_shape, self.options()); + + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_min(min.value_or(at::Tensor())); + ascend::AclTensorWrapper acl_max(max.value_or(at::Tensor())); + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD(aclnnClampTensor, acl_self.get(), acl_min.get(), acl_max.get(), acl_out.get()); + return out; +} + +REGISTER_IMPL_TO_DISPATCHER(ClampTensorFn, clamp_tensor_dispatcher, Backend::kAscend, ClampTensorKernelAscend) + at::Tensor PrivAdaptiveAvgPool2dKernelAscend(const at::Tensor& self, at::IntArrayRef output_size) { namespace ascend = at::native::flagos::ascend; auto out_shape = self.sizes().vec(); @@ -2850,8 +4174,21 @@ ::std::tuple ConvolutionBackwardKernelAscend input.sizes(), input.options()); auto grad_weight = ascend::OpPreparation::apply_tensor_without_format( weight.sizes(), weight.options()); - std::vector bias_shape = bias_sizes.has_value() - ? bias_sizes.value().vec() : std::vector{weight.size(0)}; + // grad_bias is always allocated and passed, even when output_mask[2] is + // false (aclnn writes nothing to it then). But its shape must still be + // valid: aclnnConvolutionBackward rejects an empty biasSizes, or one whose + // product is 0, with 161002 (ACLNN_ERR_PARAM_INVALID). For bias=None + // autograd hands us [0] (and an empty list is possible too), so in either + // case substitute the real bias length [Cout] = weight.size(0). + std::vector bias_shape = std::vector{weight.size(0)}; + if (bias_sizes.has_value() && !bias_sizes.value().empty()) { + const auto bs = bias_sizes.value(); + int64_t numel = 1; + for (auto d : bs) { numel *= d; } + if (numel > 0) { + bias_shape = bs.vec(); + } + } auto grad_bias = ascend::OpPreparation::apply_tensor_without_format( bias_shape, weight.options()); diff --git a/csrc/aten/backends/ascend/isin.cc b/csrc/aten/backends/ascend/isin.cc new file mode 100644 index 00000000..512560a1 --- /dev/null +++ b/csrc/aten/backends/ascend/isin.cc @@ -0,0 +1,57 @@ +// Copyright (c) 2026, BAAI. All rights reserved. + +#include "../../generated/ops.h" +#include +#include +#include +#include +#include +#include + +namespace at::native::flagos { + +// isin.Tensor_Tensor(Tensor elements, Tensor test_elements, *, bool assume_unique=False, bool invert=False) -> Tensor +// +// CANN has no aclnnIsIn kernel. The previous implementation computed on CPU by +// copying both inputs D2H and the result H2D -- three transfers per call, each +// forcing a stream sync. That is catastrophically slow when HF generate() calls +// isin on a VOCAB-SIZED elements tensor every decode step (measured 4.5 ms/call +// on a (151936,) input: ~2.4 MB copied per step, ~14 ms/token, the single +// largest cost in the generate loop, dwarfing all model.forward ops). +// +// Instead compute entirely on-device: isin(elements, test) is +// (elements.unsqueeze(-1) == test_elements).any(-1). We special-case the common +// tiny test_elements by OR-ing per-value equalities to avoid materializing the +// (numel x |test|) broadcast for a large elements tensor. All ops (eq, any, +// logical_or) are registered aclnn kernels, so no host round-trip occurs. +at::Tensor IsinTensorTensorKernelAscend(const at::Tensor& elements, + const at::Tensor& test_elements, + bool assume_unique, bool invert) { + (void)assume_unique; // no fast-path distinction on device + const int64_t n_test = test_elements.numel(); + + at::Tensor result; + if (n_test == 0) { + // Nothing to match: all-false (or all-true when inverted). + result = at::zeros(elements.sizes(), elements.options().dtype(at::kBool)); + } else { + // OR together elements == test_elements[i] for each test value, all on + // device (test_flat[i] is a 0-dim device tensor -> eq.Tensor). The test set + // is tiny in practice (eos/pad ids), so this stays cheap and avoids + // materializing a (elements.numel() x n_test) broadcast intermediate. + auto test_flat = test_elements.reshape({n_test}); + result = elements == test_flat[0]; + for (int64_t i = 1; i < n_test; ++i) { + result = at::logical_or(result, elements == test_flat[i]); + } + } + + if (invert) { + result = result.logical_not(); + } + return result; +} + +REGISTER_IMPL_TO_DISPATCHER(IsinTensorTensorFn, isin_tensor_tensor_dispatcher, Backend::kAscend, IsinTensorTensorKernelAscend) + +} // namespace at::native::flagos diff --git a/csrc/aten/backends/ascend/lift_fresh.cc b/csrc/aten/backends/ascend/lift_fresh.cc new file mode 100644 index 00000000..bc4bc7ae --- /dev/null +++ b/csrc/aten/backends/ascend/lift_fresh.cc @@ -0,0 +1,22 @@ +// Copyright (c) 2026, BAAI. All rights reserved. + +#include "../../generated/ops.h" +#include + +namespace at::native::flagos { + +// lift_fresh(Tensor(a) self) -> Tensor(a) +// +// A functionalization primitive: it marks a freshly-created tensor (typically +// from torch.tensor(scalar, device=...)) as safe to alias without a defensive +// copy. Semantically it is the identity -- CUDA/CPU both return `self` +// unchanged. transformers' generate() calls it via +// torch.tensor(bos_token_id, device='flagos') in _prepare_special_tokens, so +// the Ascend backend needs it registered even though there is no aclnn kernel. +at::Tensor LiftFreshKernelAscend(const at::Tensor& self) { + return self; +} + +REGISTER_IMPL_TO_DISPATCHER(LiftFreshFn, lift_fresh_dispatcher, Backend::kAscend, LiftFreshKernelAscend) + +} // namespace at::native::flagos diff --git a/csrc/aten/backends/ascend/masked_select.cc b/csrc/aten/backends/ascend/masked_select.cc new file mode 100644 index 00000000..cdf93274 --- /dev/null +++ b/csrc/aten/backends/ascend/masked_select.cc @@ -0,0 +1,51 @@ +// Copyright (c) 2026, BAAI. All rights reserved. + +#include "../../generated/ops.h" +#include +#include +#include "op_preparation.h" +#include "op_api_common.h" + +namespace at::native::flagos { + +// masked_select(Tensor self, Tensor mask) -> Tensor +// +// Returns a 1-D tensor of the elements of `self` where `mask` is true, in +// row-major order. The output length is data-dependent (the number of true +// entries), so it cannot be expressed by the shape-formula codegen and lives +// here as a bespoke kernel. +// +// self and mask broadcast against each other (PyTorch semantics). aclnn's +// aclnnMaskedSelect requires a pre-sized output buffer, so we first materialise +// the broadcast mask, count its true entries on host (one device->host sync via +// .item()), allocate the 1-D output, then run the kernel. +at::Tensor MaskedSelectKernelAscend(const at::Tensor& self, const at::Tensor& mask) { + namespace ascend = at::native::flagos::ascend; + + // Broadcast self and mask to a common shape (aclnn wants matching, contiguous + // buffers). infer_size gives the broadcasted shape. + auto bshape = at::infer_size(self.sizes(), mask.sizes()); + auto self_b = self.expand(bshape).contiguous(); + auto mask_b = mask.expand(bshape).contiguous(); + + // Count true elements: sum the bool mask (promoted to int64) and read to host. + int64_t count = mask_b.to(at::kLong).sum().item(); + + // aclnnMaskedSelect requires the output buffer pre-sized to the full number of + // broadcast elements (it writes `count` entries then reports the used length + // via workspace metadata). Allocate numel, run, then narrow to `count`. + int64_t numel = self_b.numel(); + auto out_full = ascend::OpPreparation::apply_tensor_without_format( + {numel}, self.options()); + + ascend::AclTensorWrapper acl_self(self_b); + ascend::AclTensorWrapper acl_mask(mask_b); + ascend::AclTensorWrapper acl_out(out_full); + + EXEC_ASCEND_CMD(aclnnMaskedSelect, acl_self.get(), acl_mask.get(), acl_out.get()); + return out_full.narrow(0, 0, count); +} + +REGISTER_IMPL_TO_DISPATCHER(MaskedSelectFn, masked_select_dispatcher, Backend::kAscend, MaskedSelectKernelAscend) + +} // namespace at::native::flagos diff --git a/csrc/aten/backends/ascend/matmul.cc b/csrc/aten/backends/ascend/matmul.cc new file mode 100644 index 00000000..df277f08 --- /dev/null +++ b/csrc/aten/backends/ascend/matmul.cc @@ -0,0 +1,239 @@ +// Copyright (c) 2026, BAAI. All rights reserved. +// +// Direct aten::matmul interception for the Ascend backend via aclnnMatmul. +// +// PyTorch's aten::matmul is CompositeImplicitAutograd and normally decomposes +// into mm + bmm + view operations before reaching PrivateUse1. torch_npu +// intercepts it at the aten::matmul level (254 aten.matmul.default/step vs +// torch_fl's mm 197 + bmm 57 + view churn 423). By registering here we +// eliminate the ~5ms/step view churn and collapse 254 ops to one aclnnMatmul +// call each, matching torch_npu's operator path exactly. +// +// Registration: register.cc TORCH_LIBRARY_IMPL(aten, PrivateUse1) adds +// m.impl("matmul", WrapperMatmul) +// which routes to this kernel when GetBackendForOp("matmul") == kAscend; +// non-Ascend backends keep PyTorch's composite decomposition by calling +// at::native::matmul directly, and register on AutogradPrivateUse1 instead so +// the autograd key never binds matmul_backward for them. +// +// Owning aten::matmul means autograd can no longer record the decomposed +// sub-ops, so it binds the op's real derivative, aten::matmul_backward. That is +// implemented below (MatmulBackwardKernelAscend) and reached via the generated +// AutogradPrivateUse1 kernel in csrc/aten/generated/variable_type.cc. Without +// it, training would decay to CPU. + +#include "../../generated/ops.h" +#include +#include +#include +#include "op_preparation.h" +#include "op_api_common.h" + +namespace at::native::flagos { + +// Compute the output shape for aten::matmul following NumPy/PyTorch semantics. +// aclnnMatmul handles all input dimensionalities ("for any shape mat multiply"). +static std::vector matmul_output_shape( + const at::Tensor& a, const at::Tensor& b) { + int64_t da = a.dim(), db = b.dim(); + // 1-D cases + if (da == 1 && db == 1) return {}; // dot -> scalar + if (da == 1 && db == 2) return {b.size(1)}; // (K,)x(K,N)->(N,) + if (da == 2 && db == 1) return {a.size(0)}; // (M,K)x(K,)->(M,) + if (da == 2 && db == 2) return {a.size(0), b.size(1)}; // mm + // N-D batched: treat 1-D inputs as row/col vector, broadcast batch dims, + // output M×N from last two dims of each input. + auto a_sz = a.sizes().vec(); + auto b_sz = b.sizes().vec(); + bool a_1d = (da == 1), b_1d = (db == 1); + if (a_1d) a_sz.insert(a_sz.begin(), 1); + if (b_1d) b_sz.push_back(1); + int64_t na = static_cast(a_sz.size()); + int64_t nb = static_cast(b_sz.size()); + int64_t n = std::max(na, nb); + std::vector out; + for (int64_t i = 0; i < n - 2; ++i) { + int64_t ai = na - n + i, bi = nb - n + i; + int64_t sa = (ai >= 0) ? a_sz[ai] : 1; + int64_t sb = (bi >= 0) ? b_sz[bi] : 1; + out.push_back(sa == 1 ? sb : sa); + } + out.push_back(a_sz[na - 2]); // M + out.push_back(b_sz[nb - 1]); // N + if (a_1d) out.erase(out.end() - 2); + if (b_1d) out.pop_back(); + return out; +} + +at::Tensor MatmulKernelAscend(const at::Tensor& self, + const at::Tensor& other) { + namespace ascend = at::native::flagos::ascend; + int8_t cube_math_type = ascend::OpPreparation::get_cube_math_type(true); + auto out = ascend::OpPreparation::apply_tensor_without_format( + matmul_output_shape(self, other), self.options()); + + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; + hsh.tensor(self); + hsh.tensor(other); + hsh.val(cube_math_type); + ascend::ExecAscendCached( + "aclnnMatmul", "aclnnMatmulGetWorkspaceSize", + opApiFuncAddr, getWsFuncAddr, hsh.h, + {&self, &other}, {&out}, + [&](ascend::GwsFunc gws, + std::vector& in, + std::vector& out_t, + uint64_t* pws, aclOpExecutor** pex) { + return gws(in[0].acl_tensor, in[1].acl_tensor, + out_t[0].acl_tensor, cube_math_type, pws, pex); + }); + return out; +} + +// --- aten::matmul_backward --- +// +// d/dself = grad @ other^T and d/dother = self^T @ grad, but only after undoing +// the shape normalization aten::matmul applied in the forward pass, which is +// where all the subtlety is. Three things have to be undone: +// +// * 1-D operands were promoted to a row/column vector, so grad must be +// unsqueezed on the matching side to line the contraction back up; +// * batch dims were broadcast, so the raw gradient can be larger than the +// operand and has to be summed back down to the operand's shape; +// * when one side is 2-D and the other batched, matmul folded the batch dims +// into the contraction. Reproducing that fold turns the gradient into one +// 2-D matmul instead of a batched one -- cheaper, and it performs the sum +// over the batch dims implicitly. +// +// op-plugin's MatmulBackwardKernelNpuOpApi.cpp was the starting reference (it +// is what aclnnMatmul is known to be driven with), but two of its shape rules +// are wrong and are deliberately not reproduced here; see the comments at the +// fold branch and at squeeze_broadcast_batch_dims. +// +// Each branch issues one aclnnMatmul through the cached executor path, so +// backward costs two fused calls instead of the mm/bmm/view chain the composite +// decomposition produced. + +// Sum a raw gradient back down to `shape`. +// +// matmul broadcasts batch dims, so d(out)/d(operand) is shaped like the +// *broadcast* operand and every dim the forward pass expanded has to be summed +// away. op-plugin instead squeezes leading size-1 dims off the operand before +// the matmul, which only coincidentally works when the broadcasting is a pure +// prefix: for an interior singleton such as (2,1,3,4) x (2,5,4,6) it leaves the +// dim in place and returns a (2,5,3,4) gradient for a (2,1,3,4) input. Summing +// after the fact is both correct in general and cheap (a no-op when no dim was +// broadcast, which is the common case, so the transformer path pays nothing). +static at::Tensor sum_to_shape(at::Tensor grad, at::IntArrayRef shape) { + if (grad.sizes() == shape) { + return grad; + } + return at::sum_to(std::move(grad), shape); +} + +static at::Tensor matmul_mat1_backward(const at::Tensor& self, + const at::Tensor& other, + const at::Tensor& grad_output) { + at::Tensor mat1 = self; + at::Tensor mat2 = other; + at::Tensor grad = grad_output; + + // 1-D operands were promoted to vectors by matmul; match that on grad. + if (mat2.dim() == 1) { + mat2 = mat2.unsqueeze(-1); + grad = grad.unsqueeze(-1); + } + if (mat1.dim() == 1) { + mat1 = mat1.unsqueeze(0); + grad = grad.unsqueeze(-2); + } + // Target the *promoted* shape, not self's: a 1-D self is still a row vector + // at this point and only gets flattened back at the end of the kernel. + const auto target = mat1.sizes().vec(); + + if (mat1.dim() == 2 && mat2.dim() > 2) { + // self is 2-D against a batched other, so grad is [B..., M, N] and the sum + // over B that the gradient needs can be folded into a single 2-D matmul: + // mat2^T is [B..., N, K] flattened to (B*N, K), i.e. its row index is the + // pair (b, n). grad must carry the *same* pair as its column index, so M + // has to be permuted to the front before flattening. op-plugin reshapes + // grad to {M, -1} directly, which pairs (m, n) columns against (b, n) rows + // and silently mixes batches; that is a bug, not a convention. + std::vector perm; + perm.reserve(grad.dim()); + perm.push_back(grad.dim() - 2); // M + for (int64_t i = 0; i < grad.dim() - 2; ++i) { + perm.push_back(i); // B... + } + perm.push_back(grad.dim() - 1); // N + const int64_t m = grad.size(-2); + mat2 = mat2.transpose(-2, -1); + mat2 = mat2.reshape({-1, mat2.size(-1)}); + grad = grad.permute(perm).contiguous().reshape({m, -1}); + // The flattened contraction already summed over B, so this lands on + // target directly and needs no further reduction. + return MatmulKernelAscend(grad, mat2).reshape(target); + } + return sum_to_shape(MatmulKernelAscend(grad, mat2.transpose(-2, -1)), target); +} + +static at::Tensor matmul_mat2_backward(const at::Tensor& self, + const at::Tensor& other, + const at::Tensor& grad_output) { + at::Tensor mat1 = self; + at::Tensor mat2 = other; + at::Tensor grad = grad_output; + + if (mat2.dim() == 1) { + mat2 = mat2.unsqueeze(-1); + grad = grad.unsqueeze(-1); + } + if (mat1.dim() == 1) { + mat1 = mat1.unsqueeze(0); + grad = grad.unsqueeze(-2); + } + const auto target = mat2.sizes().vec(); + + if (mat2.dim() == 2 && mat1.dim() > 2) { + // Mirror of the fold above. Here matmul flattened self's batch dims into + // the M axis, so both operands flatten to 2-D with the same (b, m) row + // index and no permute is needed -- B and M are already adjacent and in + // order on both sides. + at::Tensor lhs = mat1.reshape({-1, mat1.size(-1)}); + at::Tensor rhs = grad.reshape({-1, grad.size(-1)}); + return MatmulKernelAscend(lhs.transpose(-2, -1), rhs).reshape(target); + } + return sum_to_shape(MatmulKernelAscend(mat1.transpose(-2, -1), grad), target); +} + +std::tuple MatmulBackwardKernelAscend( + const at::Tensor& grad, + const at::Tensor& self, + const at::Tensor& other, + ::std::array mask) { + if (!grad.defined()) { + return std::make_tuple(at::Tensor(), at::Tensor()); + } + + at::Tensor self_grad, other_grad; + if (mask[1]) { + other_grad = matmul_mat2_backward(self, other, grad); + } + if (mask[0]) { + self_grad = matmul_mat1_backward(self, other, grad); + } + + // The 1-D promotions above leave a stray dim on the gradient of a 1-D + // operand; strip it so each gradient matches its operand's shape. + if (self.dim() == 1 && self_grad.defined() && self_grad.dim() != 1) { + self_grad = self_grad.reshape(self.sizes()); + } + if (other.dim() == 1 && other_grad.defined() && other_grad.dim() != 1) { + other_grad = other_grad.reshape(other.sizes()); + } + return std::make_tuple(self_grad, other_grad); +} + +} // namespace at::native::flagos diff --git a/csrc/aten/backends/ascend/multinomial.cc b/csrc/aten/backends/ascend/multinomial.cc new file mode 100644 index 00000000..7776b141 --- /dev/null +++ b/csrc/aten/backends/ascend/multinomial.cc @@ -0,0 +1,52 @@ +// Copyright (c) 2026, BAAI. All rights reserved. + +#include "../../generated/ops.h" +#include +#include +#include +#include "op_preparation.h" +#include "op_api_common.h" + +namespace at::native::flagos { + +// multinomial(Tensor self, int num_samples, bool replacement=False, *, +// Generator? generator=None) -> Tensor +// +// aclnnMultinomial(self, numsamples, replacement, seed, offset, out). Output is +// int64 sampled indices with the sample dim replaced by num_samples ([N] input +// -> [num_samples]; [B, N] -> [B, num_samples]). transformers' _sample() calls +// this to draw the next token. We pull a 64-bit seed from the default CPU +// generator so successive calls differ; offset is left at 0. +at::Tensor MultinomialKernelAscend(const at::Tensor& self, int64_t num_samples, + bool replacement, + ::std::optional generator) { + namespace ascend = at::native::flagos::ascend; + + auto out_shape = self.sizes().vec(); + out_shape.back() = num_samples; + auto out = ascend::OpPreparation::apply_tensor_without_format( + out_shape, self.options().dtype(at::kLong)); + + // Derive a seed. Prefer the supplied generator, else the default one. + at::Generator gen = generator.has_value() + ? generator.value() + : at::detail::getDefaultCPUGenerator(); + int64_t seed; + { + std::lock_guard lock(gen.mutex()); + seed = static_cast( + at::check_generator(gen)->random64()); + } + int64_t offset = 0; + + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD(aclnnMultinomial, acl_self.get(), num_samples, replacement, + seed, offset, acl_out.get()); + return out; +} + +REGISTER_IMPL_TO_DISPATCHER(MultinomialFn, multinomial_dispatcher, Backend::kAscend, MultinomialKernelAscend) + +} // namespace at::native::flagos diff --git a/csrc/aten/backends/ascend/op_api_common.h b/csrc/aten/backends/ascend/op_api_common.h index ad61303f..51efa8e3 100644 --- a/csrc/aten/backends/ascend/op_api_common.h +++ b/csrc/aten/backends/ascend/op_api_common.h @@ -4,12 +4,18 @@ #include #include "runtime/accelerator/ascend/acl_stream.h" +#include "runtime/allocator/caching_device_allocator.h" #include #include #include #include +#include +#include #include +#include +#include +#include #include #include @@ -36,11 +42,33 @@ inline aclDataType ToAclDataType(at::ScalarType type) { struct AclTensorWrapper { aclTensor* acl_tensor = nullptr; - // aclCreateTensor may store pointers to these arrays, so they must outlive - // the aclTensor* and the executor that references it. - std::vector sizes_; - std::vector strides_; - std::vector storage_dims_; + // aclCreateTensor stores *pointers* to the shape/stride/storage arrays rather + // than copying them, but those pointers only need to stay valid until + // aclnnGetWorkspaceSize has run: that call bakes the shapes into the + // executor, after which the arrays are dead (verified empirically — freeing + // and poisoning them between GetWorkspaceSize and the execute call still + // yields correct output). EXEC_ASCEND_CMD calls GetWorkspaceSize while this + // wrapper is still in scope, so the arrays can live INLINE in the wrapper + // (stack) instead of the old heap-allocate-and-leak scheme. Inline storage + // covers the common case (ndim <= kInlineDims); larger ranks fall back to a + // heap buffer freed in the destructor. This removes 3 heap allocs per tensor + // (9 per binary op) that previously leaked on every dispatch. + static constexpr int kInlineDims = 8; + int64_t sizes_inl_[kInlineDims]; + int64_t strides_inl_[kInlineDims]; + int64_t storage_dim_ = 0; + int64_t* sizes_ = nullptr; // -> sizes_inl_ or heap + int64_t* strides_ = nullptr; // -> strides_inl_ or heap + uint64_t ndim_ = 0; + bool heap_ = false; + // Saved so a moved-into wrapper can rebuild its aclTensor pointing at its OWN + // inline buffers (aclCreateTensor stores pointers into sizes_/strides_, which + // move with the object; without rebuild they would dangle to the source's + // buffers). Only used on the vector-storage path (cat/index). + aclDataType dtype_ = ACL_FLOAT; + int64_t offset_ = 0; + aclFormat format_ = ACL_FORMAT_ND; + void* storage_ptr_ = nullptr; // `fmt` overrides the aclFormat. Defaults to ACL_FORMAT_ND; pass e.g. // ACL_FORMAT_NCHW for ops (avg_pool2d, conv) that reject ND 4-D inputs. @@ -52,34 +80,80 @@ struct AclTensorWrapper { auto sz = tensor.sizes(); auto st = tensor.strides(); - sizes_.assign(sz.begin(), sz.end()); - strides_.assign(st.begin(), st.end()); - - int64_t offset = tensor.storage_offset(); - aclDataType dtype = ToAclDataType(tensor.scalar_type()); - aclFormat format = fmt; + ndim_ = static_cast(sz.size()); + if (ndim_ <= static_cast(kInlineDims)) { + sizes_ = sizes_inl_; + strides_ = strides_inl_; + } else { + sizes_ = new int64_t[ndim_]; + strides_ = new int64_t[ndim_]; + heap_ = true; + } + std::copy(sz.begin(), sz.end(), sizes_); + std::copy(st.begin(), st.end(), strides_); - int64_t storage_size = static_cast( + offset_ = tensor.storage_offset(); + dtype_ = ToAclDataType(tensor.scalar_type()); + format_ = fmt; + storage_dim_ = static_cast( tensor.storage().nbytes() / tensor.element_size()); - storage_dims_ = {storage_size}; - - void* storage_ptr = const_cast(tensor.storage().data()); + storage_ptr_ = const_cast(tensor.storage().data()); acl_tensor = aclCreateTensor( - sizes_.data(), - static_cast(sizes_.size()), - dtype, - strides_.data(), - offset, - format, - storage_dims_.data(), - static_cast(storage_dims_.size()), - storage_ptr); + sizes_, ndim_, dtype_, strides_, offset_, format_, + &storage_dim_, static_cast(1), storage_ptr_); } - // ACL executor caches references to aclTensor objects passed to - // GetWorkspaceSize. Destroying them here causes use-after-free. - ~AclTensorWrapper() = default; + // The shape/stride arrays are only consumed by GetWorkspaceSize (already run + // by the time this destructor fires at end of the EXEC_ASCEND_CMD scope), so + // it is safe to release everything here. Destroy the aclTensor and free the + // heap fallback (inline storage needs no free). + ~AclTensorWrapper() { + if (acl_tensor) { + aclDestroyTensor(acl_tensor); + } + if (heap_) { + delete[] sizes_; + delete[] strides_; + } + } + + // Move constructor: needed so wrappers can live in std::vector (cat/index). + // The aclTensor stores pointers into sizes_/strides_; for inline storage those + // buffers move WITH the object to a new address, so we rebuild the aclTensor + // to point at the destination's own inline buffers. Heap storage can just + // transfer the pointer. The source is left empty (acl_tensor=nullptr) so its + // destructor is a no-op. + AclTensorWrapper(AclTensorWrapper&& o) noexcept { + ndim_ = o.ndim_; + storage_dim_ = o.storage_dim_; + dtype_ = o.dtype_; offset_ = o.offset_; format_ = o.format_; + storage_ptr_ = o.storage_ptr_; + heap_ = o.heap_; + if (heap_) { + sizes_ = o.sizes_; + strides_ = o.strides_; + acl_tensor = o.acl_tensor; // still points at the (unmoved) heap buffers + } else { + std::copy(o.sizes_inl_, o.sizes_inl_ + ndim_, sizes_inl_); + std::copy(o.strides_inl_, o.strides_inl_ + ndim_, strides_inl_); + sizes_ = sizes_inl_; + strides_ = strides_inl_; + // Rebuild: the source's aclTensor referenced the source's inline buffers. + if (o.acl_tensor) { + aclDestroyTensor(o.acl_tensor); + acl_tensor = aclCreateTensor( + sizes_, ndim_, dtype_, strides_, offset_, format_, + &storage_dim_, static_cast(1), storage_ptr_); + } + } + o.acl_tensor = nullptr; + o.heap_ = false; + } + + AclTensorWrapper(const AclTensorWrapper&) = delete; + AclTensorWrapper& operator=(const AclTensorWrapper&) = delete; + AclTensorWrapper& operator=(AclTensorWrapper&&) = delete; const aclTensor* get() const { return acl_tensor; } }; @@ -88,6 +162,25 @@ inline aclrtStream GetCurrentAclStream() { return GetDefaultAclStream(); } +// Defer reuse of an aclnn scratch-workspace block until the given stream has +// finished the op that consumes it. Records an event on the stream via the +// caching allocator so free_block holds the block back until the event fires. +// No-op when the caching allocator is disabled (passthrough alloc frees the +// block straight back to the device, which cannot be reused before the sync +// on the next host-visible read). +inline void RecordWorkspaceStream(const at::Tensor& workspace, aclrtStream stream) { + if (!c10::flagos::CachingDeviceAllocator::is_enabled()) { + // Passthrough allocator aclrtFree's the block on host immediately, so it + // could be remalloc'd and overwritten before the kernel drains. Without a + // caching pool to defer reuse, the only safe option is a full sync. + aclrtSynchronizeStream(stream); + return; + } + auto* alloc = c10::flagos::GetCachingAllocator(); + alloc->record_stream(workspace.storage().data_ptr(), + reinterpret_cast(stream)); +} + inline void* GetOpApiLibHandle() { static void* handle = []() -> void* { void* h = dlopen("libopapi.so", RTLD_NOW | RTLD_GLOBAL); @@ -148,6 +241,11 @@ struct AclScalarWrapper { at::BFloat16 bf16; } value_storage; + // Absent-optional form: leaves acl_scalar null. aclnn ops that take an + // optional scalar (e.g. aclnnClamp's clipValueMin/Max) read a null pointer + // as "not supplied", which is what an empty std::optional must map to. + AclScalarWrapper() = default; + AclScalarWrapper(const at::Scalar& scalar, at::ScalarType dtype) { // aclCreateScalar stores a pointer to the value, so we must keep it alive switch (dtype) { @@ -244,6 +342,282 @@ struct AclTensorListWrapper { const aclTensorList* get() const { return acl_list; } }; +// ========================================================================== +// Repeatable-executor cache (the torch_npu-parity fast path). +// +// For eager decode the op set and shapes are constant, so the aclOpExecutor +// built by aclnnGetWorkspaceSize can be reused across steps: cache it keyed +// by (op, tensor signatures, scalar bytes), then on a hit only rebind the +// tensor data addresses (aclSetInput/OutputTensorAddr) and execute -- skipping +// both GetWorkspaceSize and aclCreateTensor. Verified on CANN 9.0.0: +// * aclSetAclOpExecutorRepeatable(ex) returns 0 for aclnnMul/aclnnAdd. +// * input/output tensor addresses are indexed over TENSORS ONLY (interleaved +// scalars consume no index) and inputs vs outputs index separately. +// * REBIND REQUIRES THE ORIGINAL aclTensor OBJECTS: passing a freshly-created +// aclTensor (even same shape) yields wrong output. So the cache OWNS the +// aclTensors that built the executor and reuses them on every hit. +// * scalars are baked into the executor at GetWorkspaceSize -> they are NOT +// rebindable and MUST be part of the cache key (alpha, etc.). +// * storage_offset is baked into the aclTensor -> also part of the key; the +// rebind address is the storage BASE ptr (offset applied internally). +// ========================================================================== + +// The op-specific GetWorkspaceSize entry, called through a variadic pointer +// (same ABI contract as EXEC_ASCEND_CMD: int64/bool/pointer args are fine, +// by-value float/double are NOT -- pass scalars as aclScalar*). +typedef int (*GwsFunc)(...); + +inline void GetRebindFuncs(void*& set_repeatable, void*& set_in_addr, + void*& set_out_addr) { + void* h = GetOpBaseLibHandle(); + if (!set_repeatable) set_repeatable = dlsym(h, "aclSetAclOpExecutorRepeatable"); + if (!set_in_addr) set_in_addr = dlsym(h, "aclSetInputTensorAddr"); + if (!set_out_addr) set_out_addr = dlsym(h, "aclSetOutputTensorAddr"); +} + +struct CachedExecKey { + const char* api = nullptr; // static per-call-site string ptr (unique id) + uint64_t sig = 0; // 64-bit hash of tensor sigs + scalar bytes + bool operator==(const CachedExecKey& o) const { + return api == o.api && sig == o.sig; + } +}; + +struct CachedExecKeyHash { + size_t operator()(const CachedExecKey& k) const { + return std::hash()(static_cast(k.api)) ^ + (static_cast(k.sig) * 0x9E3779B97F4A7C15ULL); + } +}; + +struct CachedExecEntry { + aclOpExecutor* executor = nullptr; + uint64_t workspace_size = 0; + // These OWN the aclTensors the executor is bound to. Never moved after the + // executor is built (reserve() below prevents vector realloc), so the + // aclTensor* the executor holds stay valid for the cache's lifetime. The + // entry itself only ever moves via unordered_map node relocation, which + // moves the vector's heap buffer pointer but not the elements, so the + // AclTensorWrapper objects (and their inline buffers) never physically move. + std::vector in_tensors; + std::vector out_tensors; + // Cached scratch workspace (allocated once on first use, reused on every + // hit). Safe because all ops for a given entry execute on the same stream + // serially -- by the time call N+1 reaches device, call N has already + // consumed and released the workspace. Eliminates at::empty + record_stream + // overhead (~6-25 us) on every cache hit for ops with workspace_size > 0. + at::Tensor workspace_tensor; +}; + +// 64-bit FNV-1a over the shape/stride/dtype/offset of each tensor arg plus the +// raw bytes of any scalar args. Collisions would return a wrong executor, so a +// strong 64-bit hash is used (matches torch_npu's own hash-keyed PTA cache). +struct SigHasher { + uint64_t h = 1469598103934665603ULL; + void bytes(const void* p, size_t n) { + const uint8_t* b = static_cast(p); + for (size_t i = 0; i < n; ++i) { h ^= b[i]; h *= 1099511628211ULL; } + } + template void val(const T& v) { bytes(&v, sizeof(T)); } + void tensor(const at::Tensor& t) { + if (!t.defined()) { uint8_t z = 0; bytes(&z, 1); return; } + int64_t nd = t.dim(); + val(nd); + for (auto s : t.sizes()) val(s); + for (auto s : t.strides()) val(s); + int64_t off = t.storage_offset(); + val(off); + int32_t dt = static_cast(t.scalar_type()); + val(dt); + } +}; + +inline std::unordered_map& +GetExecCache() { + static thread_local + std::unordered_map cache; + return cache; +} + +// --- optional hit/miss stats (FLAGOS_CACHE_STATS=1) -------------------------- +inline bool CacheStatsEnabled() { + static const bool on = [] { + const char* e = std::getenv("FLAGOS_CACHE_STATS"); + return e && e[0] == '1'; + }(); + return on; +} + +struct CacheStat { uint64_t hit = 0; uint64_t miss = 0; }; + +inline std::unordered_map& GetCacheStats() { + static std::unordered_map s; + return s; +} + +inline void RecordCacheStat(const char* api, bool hit) { + auto& s = GetCacheStats()[api]; + if (hit) ++s.hit; else ++s.miss; + static bool registered = [] { + std::atexit([] { + fprintf(stderr, "\n=== FLAGOS exec-cache stats (api: hit/miss) ===\n"); + for (auto& kv : GetCacheStats()) { + uint64_t tot = kv.second.hit + kv.second.miss; + fprintf(stderr, " %-28s hit=%-8llu miss=%-6llu (%.1f%% hit)\n", + kv.first, (unsigned long long)kv.second.hit, + (unsigned long long)kv.second.miss, + tot ? 100.0 * kv.second.hit / tot : 0.0); + } + }); + return true; + }(); + (void)registered; +} + +// Cached-executor dispatch. `build` calls the op-specific GetWorkspaceSize via +// the variadic func ptr, interleaving any scalar args (captured by reference) +// in the correct positions, using the OWNED input/output aclTensors passed to +// it. On a miss the executor is built + marked repeatable + stored; on a hit +// only the tensor addresses are rebound. `inputs`/`outputs` list the at::Tensor +// args in tensor-index order (inputs and outputs indexed separately). +template +void ExecAscendCached(const char* api_name, const char* ws_name, + void*& opApiFuncAddr, void*& getWsFuncAddr, + uint64_t sig, + std::initializer_list inputs, + std::initializer_list outputs, + BuildFn&& build) { + GetApiFunc(api_name, ws_name, opApiFuncAddr, getWsFuncAddr); + TORCH_CHECK(opApiFuncAddr && getWsFuncAddr, + "Failed to load symbols for ", api_name, ": ", dlerror()); + + static void* setRepeatableAddr = nullptr; + static void* setInAddrAddr = nullptr; + static void* setOutAddrAddr = nullptr; + GetRebindFuncs(setRepeatableAddr, setInAddrAddr, setOutAddrAddr); + + typedef int (*SetRepeatableFunc)(aclOpExecutor*); + typedef int (*SetAddrFunc)(aclOpExecutor*, size_t, aclTensor*, void*); + auto setRepeatable = reinterpret_cast(setRepeatableAddr); + auto setInAddr = reinterpret_cast(setInAddrAddr); + auto setOutAddr = reinterpret_cast(setOutAddrAddr); + + auto acl_stream = GetCurrentAclStream(); + auto& cache = GetExecCache(); + CachedExecKey key{api_name, sig}; + + auto it = cache.find(key); + // Optional hit/miss instrumentation (FLAGOS_CACHE_STATS=1): prints per-op + // hit/miss counts at process exit. Zero cost when the env var is unset. + if (CacheStatsEnabled()) RecordCacheStat(api_name, it != cache.end()); + // Cap resident executors: eager decode uses a bounded shape set, but a long + // varying-shape workload (e.g. prefill over growing sequence lengths) could + // otherwise accumulate executors without bound. Past the cap, run uncached + // (the executor for this shape is built, used once, and destroyed) so the + // resident set stays fixed while still serving the hot decode shapes. + static constexpr size_t kMaxCachedExecutors = 4096; + if (it == cache.end() && cache.size() >= kMaxCachedExecutors) { + auto gws = reinterpret_cast(getWsFuncAddr); + std::vector in_t, out_t; + in_t.reserve(inputs.size()); + for (const at::Tensor* t : inputs) in_t.emplace_back(*t); + out_t.reserve(outputs.size()); + for (const at::Tensor* t : outputs) out_t.emplace_back(*t); + uint64_t ws = 0; aclOpExecutor* ex = nullptr; + int ret = build(gws, in_t, out_t, &ws, &ex); + TORCH_CHECK(ret == 0, api_name, "GetWorkspaceSize failed, ret=", ret); + void* wsa = nullptr; at::Tensor wst; + if (ws > 0) { + wst = at::empty({static_cast(ws)}, + at::TensorOptions().dtype(at::kByte).device(at::kPrivateUse1)); + wsa = wst.data_ptr(); + } + typedef int (*ExecFunc)(void*, uint64_t, aclOpExecutor*, aclrtStream); + auto ef = reinterpret_cast(opApiFuncAddr); + int er = ef(wsa, ws, ex, acl_stream); + TORCH_CHECK(er == 0, api_name, " execution failed, ret=", er); + if (ws > 0) RecordWorkspaceStream(wst, acl_stream); + return; + } + if (it == cache.end()) { + // ---- miss: build executor bound to owned tensors, mark repeatable ---- + auto res = cache.emplace(key, CachedExecEntry{}); + CachedExecEntry& e = res.first->second; + e.in_tensors.reserve(inputs.size()); + for (const at::Tensor* t : inputs) e.in_tensors.emplace_back(*t); + e.out_tensors.reserve(outputs.size()); + for (const at::Tensor* t : outputs) e.out_tensors.emplace_back(*t); + + auto gws = reinterpret_cast(getWsFuncAddr); + uint64_t ws = 0; + aclOpExecutor* ex = nullptr; + int ret = build(gws, e.in_tensors, e.out_tensors, &ws, &ex); + TORCH_CHECK(ret == 0, api_name, "GetWorkspaceSize failed, ret=", ret); + + // Best-effort: mark reusable. If the op cannot be made repeatable, fall + // back to running it once uncached (executor stays valid for this call). + if (setRepeatable) { + int rr = setRepeatable(ex); + if (rr != 0) { + // Not repeatable: run once, then drop from cache to avoid rebinding a + // non-repeatable executor on a later hit. + void* wsa = nullptr; at::Tensor wst; + if (ws > 0) { + wst = at::empty({static_cast(ws)}, + at::TensorOptions().dtype(at::kByte).device(at::kPrivateUse1)); + wsa = wst.data_ptr(); + } + typedef int (*ExecFunc)(void*, uint64_t, aclOpExecutor*, aclrtStream); + auto ef = reinterpret_cast(opApiFuncAddr); + int er = ef(wsa, ws, ex, acl_stream); + TORCH_CHECK(er == 0, api_name, " execution failed, ret=", er); + if (ws > 0) RecordWorkspaceStream(wst, acl_stream); + cache.erase(key); + return; + } + } + e.executor = ex; + e.workspace_size = ws; + it = res.first; + } else { + // ---- hit: rebind the owned tensors' data addresses to current storage ---- + size_t i = 0; + for (const at::Tensor* t : inputs) { + void* addr = const_cast(t->storage().data()); + int r = setInAddr(it->second.executor, i, it->second.in_tensors[i].acl_tensor, addr); + TORCH_CHECK(r == 0, api_name, " aclSetInputTensorAddr failed, ret=", r); + ++i; + } + size_t j = 0; + for (const at::Tensor* t : outputs) { + void* addr = const_cast(t->storage().data()); + int r = setOutAddr(it->second.executor, j, it->second.out_tensors[j].acl_tensor, addr); + TORCH_CHECK(r == 0, api_name, " aclSetOutputTensorAddr failed, ret=", r); + ++j; + } + } + + // ---- execute (both paths): cached workspace, no per-hit allocation ---- + // Allocate workspace once on the first (miss) call and reuse it on every + // subsequent hit. Safety: all ops sharing this entry execute on the same + // single default ACL stream. The stream guarantees serial device execution, + // so call N's kernel has finished consuming the workspace before call N+1's + // kernel starts -- no concurrent access, no stream-record needed on hits. + CachedExecEntry& e = it->second; + void* workspace_addr = nullptr; + if (e.workspace_size > 0) { + if (!e.workspace_tensor.defined()) { + e.workspace_tensor = at::empty({static_cast(e.workspace_size)}, + at::TensorOptions().dtype(at::kByte).device(at::kPrivateUse1)); + } + workspace_addr = e.workspace_tensor.data_ptr(); + } + typedef int (*ExecFunc)(void*, uint64_t, aclOpExecutor*, aclrtStream); + auto executeFunc = reinterpret_cast(opApiFuncAddr); + int exec_ret = executeFunc(workspace_addr, e.workspace_size, e.executor, acl_stream); + TORCH_CHECK(exec_ret == 0, api_name, " execution failed, ret=", exec_ret); +} + } // namespace at::native::flagos::ascend #define EXEC_ASCEND_CMD(aclnn_api, ...) \ @@ -283,6 +657,24 @@ struct AclTensorListWrapper { workspace_addr, workspace_size, executor, acl_stream); \ TORCH_CHECK(exec_ret == 0, #aclnn_api " execution failed, ret=", \ exec_ret); \ - aclrtSynchronizeStream(acl_stream); \ + /* No per-op aclrtSynchronizeStream: ops enqueue asynchronously on the \ + * shared default stream (FIFO on-device), overlapping host dispatch with \ + * device compute. Correctness is preserved by (1) draining the default \ + * stream before any host-visible read (D2H/H2D/D2D memcpy; see \ + * runtime/accelerator/ascend/memory.cc), and (2) stream-ordering the \ + * scratch workspace below. \ + * \ + * The workspace tensor is freed on the host as soon as this scope ends, \ + * returning its block to the caching pool. Under async dispatch the \ + * kernel may still be reading that scratch when a later op reuses the \ + * block, corrupting results. record_stream defers the block's reuse \ + * until the default stream has passed this point, which fixes the race \ + * without a full sync. Inputs/outputs need no such guard: they stay \ + * live (referenced by the producing/consuming ops) and are only read \ + * back via the drained memcpy path. */ \ + if (workspace_size > 0) { \ + at::native::flagos::ascend::RecordWorkspaceStream( \ + workspace_tensor, acl_stream); \ + } \ } while (false) diff --git a/csrc/aten/backends/ascend/rms_norm.cc b/csrc/aten/backends/ascend/rms_norm.cc new file mode 100644 index 00000000..68dff344 --- /dev/null +++ b/csrc/aten/backends/ascend/rms_norm.cc @@ -0,0 +1,84 @@ +// Copyright (c) 2026, BAAI. All rights reserved. +// +// Fused RMSNorm forward for the Ascend backend via aclnnRmsNorm. Intercepts +// aten::_fused_rms_norm (a CompositeImplicitAutograd op force-included by +// scripts/codegen_ops.py FORCE_INCLUDE_OPS) so HF's Qwen3RMSNorm — which +// decomposes into ~6 elementwise ops + 2 dtype casts per layer — collapses to a +// single device kernel. The HF module must call F.rms_norm to route here (a +// small monkey-patch in the inference script); F.rms_norm -> aten::rms_norm -> +// aten::_fused_rms_norm -> this kernel. + +#include "../../generated/ops.h" +#include +#include "op_preparation.h" +#include "op_api_common.h" + +namespace at::native::flagos { + +std::tuple PrivFusedRmsNormKernelAscend( + const at::Tensor& input, + at::IntArrayRef normalized_shape, + const std::optional& weight, + std::optional eps) { + namespace ascend = at::native::flagos::ascend; + + const double epsilon = eps.value_or(1e-6); + const int64_t norm_ndim = static_cast(normalized_shape.size()); + TORCH_CHECK(norm_ndim >= 1 && input.dim() >= norm_ndim, + "_fused_rms_norm: invalid normalized_shape for input of dim ", + input.dim()); + + // aclnnRmsNorm needs a dense input. + at::Tensor x = input.is_contiguous() ? input : input.contiguous(); + + // gamma is required by aclnn; synthesize ones matching normalized_shape when + // weight is absent. Cast to x's dtype (aclnnRmsNorm wants matching dtypes). + at::Tensor gamma; + if (weight.has_value() && weight.value().defined()) { + gamma = weight.value(); + if (gamma.scalar_type() != x.scalar_type()) { + gamma = gamma.to(x.scalar_type()); + } + if (!gamma.is_contiguous()) { + gamma = gamma.contiguous(); + } + } else { + gamma = at::ones(normalized_shape, x.options()); + } + + at::Tensor output = ascend::OpPreparation::apply_tensor_without_format( + x.sizes(), x.options()); + + // rstd (reciprocal std) has the input's leading dims with the normalized dims + // collapsed to 1, in float32 (CANN + torch both produce fp32 rstd). We keep + // it for the returned tuple even though inference discards it. + std::vector rstd_shape(x.sizes().begin(), x.sizes().end()); + for (int64_t i = 0; i < norm_ndim; ++i) { + rstd_shape[x.dim() - 1 - i] = 1; + } + at::Tensor rstd = ascend::OpPreparation::apply_tensor_without_format( + rstd_shape, x.options().dtype(at::kFloat)); + + ascend::AclTensorWrapper acl_x(x); + ascend::AclTensorWrapper acl_gamma(gamma); + ascend::AclTensorWrapper acl_out(output); + ascend::AclTensorWrapper acl_rstd(rstd); + + // aclnnRmsNorm(x, gamma, epsilon, yOut, rstdOut) + EXEC_ASCEND_CMD(aclnnRmsNorm, + acl_x.get(), + acl_gamma.get(), + epsilon, + const_cast(acl_out.get()), + const_cast(acl_rstd.get())); + + return std::make_tuple(output, rstd); +} + +REGISTER_IMPL_TO_DISPATCHER( + PrivFusedRmsNormFn, + priv_fused_rms_norm_dispatcher, + Backend::kAscend, + PrivFusedRmsNormKernelAscend) + +} // namespace at::native::flagos diff --git a/csrc/aten/backends/ascend/rng.cc b/csrc/aten/backends/ascend/rng.cc new file mode 100644 index 00000000..8cddcc62 --- /dev/null +++ b/csrc/aten/backends/ascend/rng.cc @@ -0,0 +1,99 @@ +// Copyright (c) 2026, BAAI. All rights reserved. + +#include "../../generated/ops.h" +#include +#include +#include +#include "op_preparation.h" +#include "op_api_common.h" + +namespace at::native::flagos { + +namespace { + +// Draw a fresh 64-bit seed from the default CPU generator so successive RNG +// calls decorrelate. offset is left at 0 (aclnn advances its own state). +int64_t next_seed() { + auto gen = at::detail::getDefaultCPUGenerator(); + std::lock_guard lock(gen.mutex()); + return static_cast( + at::check_generator(gen)->random64()); +} + +at::Tensor make_empty(at::IntArrayRef size, ::std::optional dtype, + ::std::optional layout, + ::std::optional device, + ::std::optional pin_memory) { + auto options = at::TensorOptions() + .dtype(dtype.value_or(at::kFloat)) + .layout(layout.value_or(at::kStrided)) + .device(device.value_or(at::Device(at::kPrivateUse1, 0))) + .pinned_memory(pin_memory.value_or(false)); + return at::empty(size, options); +} + +} // namespace + +// randn(int[] size, *, ScalarType?, Layout?, Device?, bool? pin_memory) +// -> Tensor of N(0, 1) samples. aclnnInplaceNormal(selfRef, mean, std, seed, offset). +at::Tensor RandnKernelAscend(at::IntArrayRef size, + ::std::optional dtype, + ::std::optional layout, + ::std::optional device, + ::std::optional pin_memory) { + namespace ascend = at::native::flagos::ascend; + auto out = make_empty(size, dtype, layout, device, pin_memory); + ascend::AclTensorWrapper acl_out(out); + EXEC_ASCEND_CMD(aclnnInplaceNormal, const_cast(acl_out.get()), + 0.0f, 1.0f, next_seed(), static_cast(0)); + return out; +} + +// rand(int[] size, ...) -> Tensor of U[0, 1) samples. +// aclnnInplaceUniform(selfRef, from, to, seed, offset). +at::Tensor RandKernelAscend(at::IntArrayRef size, + ::std::optional dtype, + ::std::optional layout, + ::std::optional device, + ::std::optional pin_memory) { + namespace ascend = at::native::flagos::ascend; + auto out = make_empty(size, dtype, layout, device, pin_memory); + ascend::AclTensorWrapper acl_out(out); + EXEC_ASCEND_CMD(aclnnInplaceUniform, const_cast(acl_out.get()), + 0.0, 1.0, static_cast(next_seed()), + static_cast(0)); + return out; +} + +// randint.low(int low, int high, int[] size, ...) -> Tensor of ints in [low, high). +// aclnnInplaceRandom(selfRef, from, to, seed, offset). +at::Tensor RandintLowKernelAscend(int64_t low, int64_t high, at::IntArrayRef size, + ::std::optional dtype, + ::std::optional layout, + ::std::optional device, + ::std::optional pin_memory) { + namespace ascend = at::native::flagos::ascend; + // randint defaults to int64 output when no dtype is given. + auto out = make_empty(size, dtype.value_or(at::kLong), layout, device, pin_memory); + ascend::AclTensorWrapper acl_out(out); + EXEC_ASCEND_CMD(aclnnInplaceRandom, const_cast(acl_out.get()), + low, high, next_seed(), static_cast(0)); + return out; +} + +// randint(int high, int[] size, ...) -> ints in [0, high). Delegates to the +// low overload with low=0. +at::Tensor RandintKernelAscend(int64_t high, at::IntArrayRef size, + ::std::optional dtype, + ::std::optional layout, + ::std::optional device, + ::std::optional pin_memory) { + return RandintLowKernelAscend(0, high, size, dtype, layout, device, pin_memory); +} + +REGISTER_IMPL_TO_DISPATCHER(RandnFn, randn_dispatcher, Backend::kAscend, RandnKernelAscend) +REGISTER_IMPL_TO_DISPATCHER(RandFn, rand_dispatcher, Backend::kAscend, RandKernelAscend) +REGISTER_IMPL_TO_DISPATCHER(RandintFn, randint_dispatcher, Backend::kAscend, RandintKernelAscend) +REGISTER_IMPL_TO_DISPATCHER(RandintLowFn, randint_low_dispatcher, Backend::kAscend, RandintLowKernelAscend) + +} // namespace at::native::flagos diff --git a/csrc/aten/backends/ascend/scaled_dot_product_attention.cc b/csrc/aten/backends/ascend/scaled_dot_product_attention.cc index e313535b..6db7fb1a 100644 --- a/csrc/aten/backends/ascend/scaled_dot_product_attention.cc +++ b/csrc/aten/backends/ascend/scaled_dot_product_attention.cc @@ -52,16 +52,39 @@ PrivScaledDotProductEfficientAttentionKernelAscend( // Input validation TORCH_CHECK(query.dim() == 4, "query must be 4D [B, N, S, D]"); - TORCH_CHECK(query.sizes() == key.sizes(), "query and key must have same shape"); - TORCH_CHECK(query.sizes() == value.sizes(), "query and value must have same shape"); + TORCH_CHECK(key.dim() == 4 && value.dim() == 4, "key/value must be 4D [B, N, S, D]"); TORCH_CHECK(query.is_privateuseone(), "SDPA Ascend: inputs must be on NPU"); TORCH_CHECK(dropout_p == 0.0, "SDPA Ascend: dropout not yet supported (aclnn requires explicit mask handling)"); int64_t B = query.size(0); - int64_t N = query.size(1); // num_heads + int64_t N = query.size(1); // num query heads int64_t S = query.size(2); // seq_len int64_t D = query.size(3); // head_dim + // Grouped-query / multi-query attention: key and value may carry fewer heads + // than the query (Qwen3 uses num_kv_heads < num_attention_heads). PyTorch's + // SDPA repeats the kv heads (repeat_kv) before the math path; the aclnn flash + // kernel expects q/k/v with matching head counts, so replicate each kv head + // N/N_kv times along dim 1 to match the query. Contiguous so the aclnn tensor + // wrapper sees a dense [B, N, S, D] buffer. + // key/value carry their own seq_len (S_kv), which differs from the query's S + // during incremental decode (query S==1, kv S==full context). Expand only the + // head dim, preserving each tensor's own seq_len. + at::Tensor key_eff = key; + at::Tensor value_eff = value; + int64_t N_kv = key.size(1); + int64_t S_kv = key.size(2); + if (N_kv != N) { + TORCH_CHECK(N_kv > 0 && N % N_kv == 0, + "SDPA Ascend GQA: query heads (", N, ") must be a multiple of kv heads (", N_kv, ")"); + int64_t repeat = N / N_kv; + // [B, N_kv, S_kv, D] -> [B, N_kv, repeat, S_kv, D] -> [B, N, S_kv, D] + key_eff = key.unsqueeze(2).expand({B, N_kv, repeat, S_kv, D}).reshape({B, N, S_kv, D}).contiguous(); + value_eff = value.unsqueeze(2).expand({B, N_kv, repeat, S_kv, D}).reshape({B, N, S_kv, D}).contiguous(); + } + TORCH_CHECK(key_eff.size(1) == N && value_eff.size(1) == N, + "GQA expand failed to match query head count"); + // Compute scale (default: 1/sqrt(D)) double scale_value = scale.value_or(1.0 / std::sqrt(static_cast(D))); @@ -107,8 +130,8 @@ PrivScaledDotProductEfficientAttentionKernelAscend( // Prepare aclnn arguments AclTensorWrapper q_wrap(query); - AclTensorWrapper k_wrap(key); - AclTensorWrapper v_wrap(value); + AclTensorWrapper k_wrap(key_eff); + AclTensorWrapper v_wrap(value_eff); AclTensorWrapper mask_wrap(is_causal ? atten_mask : at::Tensor()); AclTensorWrapper drop_mask_wrap(drop_mask); AclTensorWrapper softmax_max_wrap(softmax_max); diff --git a/csrc/aten/backends/ascend/scatter.cc b/csrc/aten/backends/ascend/scatter.cc new file mode 100644 index 00000000..d95b2d87 --- /dev/null +++ b/csrc/aten/backends/ascend/scatter.cc @@ -0,0 +1,36 @@ +// Copyright (c) 2026, BAAI. All rights reserved. + +#include "../../generated/ops.h" +#include +#include "op_preparation.h" +#include "op_api_common.h" + +namespace at::native::flagos { + +// scatter.src(Tensor self, int dim, Tensor index, Tensor src) -> Tensor +// +// Out-of-place scatter: out = self.clone(); out.scatter_(dim, index, src). +// aclnnScatter(self, dim, index, src, reduce, out) with reduce=0 (replace). +// transformers' TopPLogitsWarper uses scatter() to unsort the removal mask. +at::Tensor ScatterSrcKernelAscend(const at::Tensor& self, int64_t dim, + const at::Tensor& index, const at::Tensor& src) { + namespace ascend = at::native::flagos::ascend; + + int64_t d = dim < 0 ? dim + self.dim() : dim; + // aclnnScatter writes the full result to `out`; seed it with self so entries + // not covered by index retain their original values. + auto out = self.clone(); + + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_index(index); + ascend::AclTensorWrapper acl_src(src); + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD(aclnnScatter, acl_self.get(), d, acl_index.get(), + acl_src.get(), static_cast(0), acl_out.get()); + return out; +} + +REGISTER_IMPL_TO_DISPATCHER(ScatterSrcFn, scatter_src_dispatcher, Backend::kAscend, ScatterSrcKernelAscend) + +} // namespace at::native::flagos diff --git a/csrc/aten/backends/ascend/sort.cc b/csrc/aten/backends/ascend/sort.cc new file mode 100644 index 00000000..a5c32bd4 --- /dev/null +++ b/csrc/aten/backends/ascend/sort.cc @@ -0,0 +1,52 @@ +// Copyright (c) 2026, BAAI. All rights reserved. + +#include "../../generated/ops.h" +#include +#include "op_preparation.h" +#include "op_api_common.h" + +namespace at::native::flagos { + +namespace { + +std::tuple SortImpl( + const at::Tensor& self, bool stable, int64_t dim, bool descending) { + namespace ascend = at::native::flagos::ascend; + + int64_t d = dim < 0 ? dim + self.dim() : dim; + // sort preserves the full shape; values keep dtype, indices are int64. + auto values = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options()); + auto indices = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options().dtype(at::kLong)); + + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_values(values); + ascend::AclTensorWrapper acl_indices(indices); + + EXEC_ASCEND_CMD(aclnnSort, acl_self.get(), stable, d, descending, + acl_values.get(), acl_indices.get()); + return std::make_tuple(values, indices); +} + +} // namespace + +// sort(Tensor self, int dim=-1, bool descending=False) -> (values, indices) +std::tuple SortKernelAscend( + const at::Tensor& self, int64_t dim, bool descending) { + return SortImpl(self, /*stable=*/false, dim, descending); +} + +// sort.stable(Tensor self, *, bool? stable, int dim=-1, bool descending=False) +// -> (values, indices). transformers' TopPLogitsWarper calls torch.sort(), +// which resolves to this overload on recent torch. +std::tuple SortStableKernelAscend( + const at::Tensor& self, ::std::optional stable, int64_t dim, + bool descending) { + return SortImpl(self, stable.value_or(false), dim, descending); +} + +REGISTER_IMPL_TO_DISPATCHER(SortFn, sort_dispatcher, Backend::kAscend, SortKernelAscend) +REGISTER_IMPL_TO_DISPATCHER(SortStableFn, sort_stable_dispatcher, Backend::kAscend, SortStableKernelAscend) + +} // namespace at::native::flagos diff --git a/csrc/aten/backends/ascend/topk.cc b/csrc/aten/backends/ascend/topk.cc new file mode 100644 index 00000000..cfb0caed --- /dev/null +++ b/csrc/aten/backends/ascend/topk.cc @@ -0,0 +1,40 @@ +// Copyright (c) 2026, BAAI. All rights reserved. + +#include "../../generated/ops.h" +#include +#include "op_preparation.h" +#include "op_api_common.h" + +namespace at::native::flagos { + +// topk(Tensor self, int k, int dim=-1, bool largest=True, bool sorted=True) +// -> (Tensor values, Tensor indices) +// +// aclnnTopk(self, k, dim, largest, sorted, valuesOut, indicesOut). Used by +// transformers' TopKLogitsWarper during sampling. Output shape equals the +// input with the reduced dim resized to k; indices are int64. +std::tuple TopkKernelAscend( + const at::Tensor& self, int64_t k, int64_t dim, bool largest, bool sorted) { + namespace ascend = at::native::flagos::ascend; + + int64_t d = dim < 0 ? dim + self.dim() : dim; + auto out_shape = self.sizes().vec(); + out_shape[d] = k; + + auto values = ascend::OpPreparation::apply_tensor_without_format( + out_shape, self.options()); + auto indices = ascend::OpPreparation::apply_tensor_without_format( + out_shape, self.options().dtype(at::kLong)); + + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_values(values); + ascend::AclTensorWrapper acl_indices(indices); + + EXEC_ASCEND_CMD(aclnnTopk, acl_self.get(), k, d, largest, sorted, + acl_values.get(), acl_indices.get()); + return std::make_tuple(values, indices); +} + +REGISTER_IMPL_TO_DISPATCHER(TopkFn, topk_dispatcher, Backend::kAscend, TopkKernelAscend) + +} // namespace at::native::flagos diff --git a/csrc/aten/common.h b/csrc/aten/common.h index 4935a066..a3114439 100644 --- a/csrc/aten/common.h +++ b/csrc/aten/common.h @@ -17,7 +17,20 @@ namespace at::native::flagos { // Backend selector for unified op wrappers. // Determines which physical backend impl() dispatches to. -enum class Backend { kCuda, kFlagOs, kFlagOsPython, kAscend, kMusa, kMetax, kTsingMicro, kGcu }; +// kUncached is a sentinel used by Dispatcher's per-op backend cache; it is +// never stored in the BackendTable and never returned by GetBackendForOp. +// Keep it last so the real backends stay contiguous. +enum class Backend { + kCuda, + kFlagOs, + kFlagOsPython, + kAscend, + kMusa, + kMetax, + kTsingMicro, + kGcu, + kUncached +}; // Returns the backend for a given op name, loaded once from config file at startup. // Config file path: $FLAGOS_BACKEND_CONFIG or torch_fl/configs/backends.conf diff --git a/csrc/aten/contiguous_ops.cc b/csrc/aten/contiguous_ops.cc index 6977436d..e6155b36 100644 --- a/csrc/aten/contiguous_ops.cc +++ b/csrc/aten/contiguous_ops.cc @@ -10,6 +10,10 @@ #include #include #include "device_boxing.h" +// Included unconditionally: the #else branches below cover TsingMicro, GCU and +// MUSA-without-mudnn as well as Ascend, and this header supplies inline no-op +// fallbacks for those platforms. +#include "backends/ascend/ascend_copy.h" #if defined(FLAGOS_MUSA_KERNEL) #include "backends/musa/mudnn_common.h" @@ -41,25 +45,29 @@ at::Tensor contiguous( DeviceBoxingGuard guard(self, result); at::native::copy_(result, self, false); #else - // Ascend: no CUDA runtime, fall back to CPU round-trip. - size_t storage_size = self.storage().nbytes(); - at::Tensor storage_cpu = at::empty( - {static_cast(storage_size)}, - at::TensorOptions().dtype(at::kByte).device(at::kCPU)); - Memcpy(storage_cpu.data_ptr(), self.storage().data(), storage_size, MemcpyDeviceToHost); - - at::Tensor cpu_view = at::empty({0}, self.options().device(at::kCPU)); - cpu_view.set_( - storage_cpu.storage(), - self.storage_offset(), - self.sizes(), - self.strides()); - - auto cpu_contig = at::empty(self.sizes(), self.options().device(at::kCPU).memory_format(memory_format)); - cpu_contig.copy_(cpu_view); - - size_t nbytes = cpu_contig.numel() * cpu_contig.element_size(); - Memcpy(result.data_ptr(), cpu_contig.data_ptr(), nbytes, MemcpyHostToDevice); + // Ascend: copy the strided source into the contiguous result on-device + // via aclnnInplaceCopy. Falls back to a CPU round-trip only if that path + // is unavailable. + if (!ascend::StridedCopy(result, self)) { + size_t storage_size = self.storage().nbytes(); + at::Tensor storage_cpu = at::empty( + {static_cast(storage_size)}, + at::TensorOptions().dtype(at::kByte).device(at::kCPU)); + Memcpy(storage_cpu.data_ptr(), self.storage().data(), storage_size, MemcpyDeviceToHost); + + at::Tensor cpu_view = at::empty({0}, self.options().device(at::kCPU)); + cpu_view.set_( + storage_cpu.storage(), + self.storage_offset(), + self.sizes(), + self.strides()); + + auto cpu_contig = at::empty(self.sizes(), self.options().device(at::kCPU).memory_format(memory_format)); + cpu_contig.copy_(cpu_view); + + size_t nbytes = cpu_contig.numel() * cpu_contig.element_size(); + Memcpy(result.data_ptr(), cpu_contig.data_ptr(), nbytes, MemcpyHostToDevice); + } #endif } @@ -108,7 +116,12 @@ at::Tensor clone( // torch_musa implements for strided/dtype-casting copies on device. auto result = at::empty( self.sizes(), self.options().memory_format(memory_format)); - result.copy_(self); + // On-device strided copy (aclnnInplaceCopy) instead of result.copy_(self), + // which would bounce through a CPU round-trip. This is the Qwen3 GQA + // repeat_kv hotspot (~59% of inference time before this change). + if (!ascend::StridedCopy(result, self)) { + result.copy_(self); + } return result; #endif } diff --git a/csrc/aten/copy_ops.cc b/csrc/aten/copy_ops.cc index e616f0d2..431cde29 100644 --- a/csrc/aten/copy_ops.cc +++ b/csrc/aten/copy_ops.cc @@ -13,6 +13,10 @@ #include #include #include "device_boxing.h" +// Included unconditionally: the #else branches below cover TsingMicro, GCU and +// MUSA-without-mudnn as well as Ascend, and this header supplies inline no-op +// fallbacks for those platforms. +#include "backends/ascend/ascend_copy.h" #if defined(FLAGOS_MUSA_KERNEL) #include "backends/musa/mudnn_common.h" @@ -84,49 +88,53 @@ at::Tensor _copy_from( DeviceBoxingGuard guard(self, dst); at::native::copy_(const_cast(dst), self, false); #else - // Ascend: no CUDA runtime, fall back to CPU round-trip. - at::Tensor self_contig = self.is_contiguous() - ? self - : at::native::flagos::contiguous(self, c10::MemoryFormat::Contiguous); - size_t nbytes = self_contig.numel() * self_contig.element_size(); - at::Tensor cpu_src = - at::empty(self_contig.sizes(), self_contig.options().device(at::kCPU)); - if (nbytes > 0) { - Memcpy( - cpu_src.data_ptr(), - self_contig.data_ptr(), - nbytes, - MemcpyDeviceToHost); - } - size_t dst_storage_nbytes = dst.storage().nbytes(); - at::Tensor cpu_dst_storage = at::empty( - {static_cast(dst_storage_nbytes)}, - dst.options().device(at::kCPU).dtype(at::kByte)); - int64_t dst_storage_offset_bytes = - dst.storage_offset() * static_cast(dst.element_size()); - char* dst_storage_base = - static_cast(dst.data_ptr()) - dst_storage_offset_bytes; - if (dst_storage_nbytes > 0) { - Memcpy( - cpu_dst_storage.data_ptr(), - dst_storage_base, - dst_storage_nbytes, - MemcpyDeviceToHost); - } + // Ascend: copy on-device via aclnnInplaceCopy, which honors both src and + // dst strides/offset and casts dtype. Avoids the CPU round-trip below. + if (!ascend::StridedCopy(dst, self)) { + // Fallback: CPU round-trip (device->host, strided copy on CPU, host->device). + at::Tensor self_contig = self.is_contiguous() + ? self + : at::native::flagos::contiguous(self, c10::MemoryFormat::Contiguous); + size_t nbytes = self_contig.numel() * self_contig.element_size(); + at::Tensor cpu_src = + at::empty(self_contig.sizes(), self_contig.options().device(at::kCPU)); + if (nbytes > 0) { + Memcpy( + cpu_src.data_ptr(), + self_contig.data_ptr(), + nbytes, + MemcpyDeviceToHost); + } + size_t dst_storage_nbytes = dst.storage().nbytes(); + at::Tensor cpu_dst_storage = at::empty( + {static_cast(dst_storage_nbytes)}, + dst.options().device(at::kCPU).dtype(at::kByte)); + int64_t dst_storage_offset_bytes = + dst.storage_offset() * static_cast(dst.element_size()); + char* dst_storage_base = + static_cast(dst.data_ptr()) - dst_storage_offset_bytes; + if (dst_storage_nbytes > 0) { + Memcpy( + cpu_dst_storage.data_ptr(), + dst_storage_base, + dst_storage_nbytes, + MemcpyDeviceToHost); + } - at::Tensor cpu_dst = at::empty({0}, dst.options().device(at::kCPU)); - cpu_dst.set_( - cpu_dst_storage.storage(), - dst.storage_offset(), - dst.sizes(), - dst.strides()); - at::native::copy_(cpu_dst, cpu_src, false); - if (dst_storage_nbytes > 0) { - Memcpy( - dst_storage_base, - cpu_dst_storage.data_ptr(), - dst_storage_nbytes, - MemcpyHostToDevice); + at::Tensor cpu_dst = at::empty({0}, dst.options().device(at::kCPU)); + cpu_dst.set_( + cpu_dst_storage.storage(), + dst.storage_offset(), + dst.sizes(), + dst.strides()); + at::native::copy_(cpu_dst, cpu_src, false); + if (dst_storage_nbytes > 0) { + Memcpy( + dst_storage_base, + cpu_dst_storage.data_ptr(), + dst_storage_nbytes, + MemcpyHostToDevice); + } } #endif } @@ -308,28 +316,38 @@ at::Tensor _to_copy( musa_ops::MudnnCopy(self_contig, result); #elif defined(USE_ASCEND) || defined(USE_TSINGMICRO) || defined(USE_GCU) || \ defined(USE_MUSA) - // Ascend / TsingMicro: no CUDA runtime, fall back to CPU round-trip for dtype cast. - size_t nbytes = self_contig.numel() * self_contig.element_size(); - at::Tensor cpu_tensor = - at::empty(self_contig.sizes(), self_contig.options().device(at::kCPU)); - if (nbytes > 0) { - Memcpy( - cpu_tensor.data_ptr(), - self_contig.data_ptr(), - nbytes, - MemcpyDeviceToHost); - } - cpu_tensor = cpu_tensor.to(dtype); - result = at::empty( - cpu_tensor.sizes(), - cpu_tensor.options().device(c10::Device(c10::kPrivateUse1, device_index))); - size_t result_nbytes = cpu_tensor.numel() * cpu_tensor.element_size(); - if (result_nbytes > 0) { - Memcpy( - result.data_ptr(), - cpu_tensor.data_ptr(), - result_nbytes, - MemcpyHostToDevice); + // No CUDA runtime on these backends, so the CUDA TensorIterator cast + // below is unavailable. +#ifdef USE_ASCEND + // Ascend casts on-device via aclnnCast, avoiding the D2H->CPU->H2D + // round-trip that dominated HF RMSNorm (two fp16<->fp32 casts per layer). + result = ascend::DtypeCast(self_contig, dtype); +#endif + if (!result.defined()) { + // Fallback: CPU round-trip when no on-device cast is available + // (TsingMicro / GCU / MUSA, or an Ascend dtype pair aclnnCast rejects). + size_t nbytes = self_contig.numel() * self_contig.element_size(); + at::Tensor cpu_tensor = + at::empty(self_contig.sizes(), self_contig.options().device(at::kCPU)); + if (nbytes > 0) { + Memcpy( + cpu_tensor.data_ptr(), + self_contig.data_ptr(), + nbytes, + MemcpyDeviceToHost); + } + cpu_tensor = cpu_tensor.to(dtype); + result = at::empty( + cpu_tensor.sizes(), + cpu_tensor.options().device(c10::Device(c10::kPrivateUse1, device_index))); + size_t result_nbytes = cpu_tensor.numel() * cpu_tensor.element_size(); + if (result_nbytes > 0) { + Memcpy( + result.data_ptr(), + cpu_tensor.data_ptr(), + result_nbytes, + MemcpyHostToDevice); + } } #else // CUDA platform: use DeviceBoxingGuard + CUDA TensorIterator copy kernel diff --git a/csrc/aten/dispatcher.h b/csrc/aten/dispatcher.h index cf99237f..7aeea27e 100644 --- a/csrc/aten/dispatcher.h +++ b/csrc/aten/dispatcher.h @@ -60,7 +60,20 @@ class Dispatcher { template decltype(auto) operator()(Args&&... args) const { - return DispatchAs(op_name_, std::forward(args)...); + // Hot path: the op name is fixed (op_name_) and the backend routing is + // immutable once the config is loaded, so resolve it once and cache. This + // avoids constructing a std::string from op_name_ and hashing it in the + // BackendTable on EVERY op call — measured as a significant per-op cost in + // the Ascend eager decode loop (thousands of ops/token). + Backend backend = cached_backend_; + if (__builtin_expect(backend == Backend::kUncached, 0)) { + backend = GetBackendForOp(op_name_); + cached_backend_ = backend; + } + LogDispatch(op_name_, backend); + auto fn = GetFn(backend); + TORCH_CHECK(fn, op_name_, ": backend not registered"); + return fn(std::forward(args)...); } template @@ -109,6 +122,10 @@ class Dispatcher { } const char* op_name_ = nullptr; + // Per-op backend cache for the hot operator() path (see comment there). + // mutable: operator() is const but memoizes on first call. Benign data race + // under concurrent first-use — all threads compute the same immutable value. + mutable Backend cached_backend_ = Backend::kUncached; FnPtr cuda_fn_ = nullptr; FnPtr flagos_fn_ = nullptr; FnPtr flagos_python_fn_ = nullptr; diff --git a/csrc/aten/empty.cc b/csrc/aten/empty.cc index 58e1e7a0..2a0d187e 100644 --- a/csrc/aten/empty.cc +++ b/csrc/aten/empty.cc @@ -27,7 +27,20 @@ at::Tensor empty_memory_format( TORCH_CHECK( !c10::pinned_memory_or_default(pin_memory_opt), "Pin memory can only be on CPU"); - const c10::DeviceGuard device_guard(device); + // The caching allocator resolves the current device itself (via + // aclrtGetDevice) and allocates there, so a DeviceGuard is only needed to + // switch the ambient device when the requested device differs from the + // current one. Constructing an unconditional c10::DeviceGuard here costs + // ~2.8us/call on the decode hot path (measured) because its ctor/dtor route + // through the guard registry; skipping it when the device already matches + // removes that cost for the overwhelmingly common single-device case while + // preserving multi-device correctness. + int cur_device = -1; + ::GetDevice(&cur_device); + std::optional device_guard; + if (device.has_index() && device.index() != cur_device) { + device_guard.emplace(device); + } constexpr c10::DispatchKeySet pu1_dks(c10::DispatchKey::PrivateUse1); auto allocator = at::GetAllocator(at::kPrivateUse1); return at::detail::empty_generic( diff --git a/csrc/aten/generated/cuda_kernels.cc b/csrc/aten/generated/cuda_kernels.cc index 7e20774f..46acf9e7 100644 --- a/csrc/aten/generated/cuda_kernels.cc +++ b/csrc/aten/generated/cuda_kernels.cc @@ -121,6 +121,7 @@ #include #include #include +#include #include #include #include @@ -3929,6 +3930,15 @@ ::std::tuple return result; } +::std::tuple PrivFusedRmsNormKernelCuda(const at::Tensor & input, at::IntArrayRef normalized_shape, const ::std::optional & weight, ::std::optional eps) { + at::Tensor weight_t = weight.has_value() ? *weight : at::Tensor(); + DeviceBoxingGuard guard(input, weight_t); + auto result = at::_fused_rms_norm(input, normalized_shape, weight, eps); + UnboxToFlagos(std::get<0>(result)); + UnboxToFlagos(std::get<1>(result)); + return result; +} + ::std::tuple PrivFusedRmsNormBackwardKernelCuda(const at::Tensor & grad_out, const at::Tensor & input, at::IntArrayRef normalized_shape, const at::Tensor & rstd, const ::std::optional & weight, ::std::array output_mask) { at::Tensor weight_t = weight.has_value() ? *weight : at::Tensor(); DeviceBoxingGuard guard(grad_out, input, rstd, weight_t); @@ -16984,6 +16994,7 @@ REGISTER_IMPL_TO_DISPATCHER(PrivFusedDropoutOutFn, priv_fused_dropout_out_dispat REGISTER_IMPL_TO_DISPATCHER(PrivFusedMovingAvgObsFqHelperFn, priv_fused_moving_avg_obs_fq_helper_dispatcher, Backend::kCuda, PrivFusedMovingAvgObsFqHelperKernelCuda) REGISTER_IMPL_TO_DISPATCHER(PrivFusedMovingAvgObsFqHelperOutFn, priv_fused_moving_avg_obs_fq_helper_out_dispatcher, Backend::kCuda, PrivFusedMovingAvgObsFqHelperOutKernelCuda) REGISTER_IMPL_TO_DISPATCHER(PrivFusedMovingAvgObsFqHelperFunctionalFn, priv_fused_moving_avg_obs_fq_helper_functional_dispatcher, Backend::kCuda, PrivFusedMovingAvgObsFqHelperFunctionalKernelCuda) +REGISTER_IMPL_TO_DISPATCHER(PrivFusedRmsNormFn, priv_fused_rms_norm_dispatcher, Backend::kCuda, PrivFusedRmsNormKernelCuda) REGISTER_IMPL_TO_DISPATCHER(PrivFusedRmsNormBackwardFn, priv_fused_rms_norm_backward_dispatcher, Backend::kCuda, PrivFusedRmsNormBackwardKernelCuda) REGISTER_IMPL_TO_DISPATCHER(PrivFusedSgdOutFn, priv_fused_sgd_out_dispatcher, Backend::kCuda, PrivFusedSgdOutKernelCuda) REGISTER_IMPL_TO_DISPATCHER(PrivFusedSgdTensorLrOutFn, priv_fused_sgd_tensor_lr_out_dispatcher, Backend::kCuda, PrivFusedSgdTensorLrOutKernelCuda) diff --git a/csrc/aten/generated/ops.cc b/csrc/aten/generated/ops.cc index f76968a4..adf404f9 100644 --- a/csrc/aten/generated/ops.cc +++ b/csrc/aten/generated/ops.cc @@ -347,6 +347,7 @@ ADD_IMPL_TO_DISPATCHER(PrivFusedDropoutOutFn, priv_fused_dropout_out_dispatcher, ADD_IMPL_TO_DISPATCHER(PrivFusedMovingAvgObsFqHelperFn, priv_fused_moving_avg_obs_fq_helper_dispatcher, "_fused_moving_avg_obs_fq_helper") ADD_IMPL_TO_DISPATCHER(PrivFusedMovingAvgObsFqHelperOutFn, priv_fused_moving_avg_obs_fq_helper_out_dispatcher, "_fused_moving_avg_obs_fq_helper.out") ADD_IMPL_TO_DISPATCHER(PrivFusedMovingAvgObsFqHelperFunctionalFn, priv_fused_moving_avg_obs_fq_helper_functional_dispatcher, "_fused_moving_avg_obs_fq_helper_functional") +ADD_IMPL_TO_DISPATCHER(PrivFusedRmsNormFn, priv_fused_rms_norm_dispatcher, "_fused_rms_norm") ADD_IMPL_TO_DISPATCHER(PrivFusedRmsNormBackwardFn, priv_fused_rms_norm_backward_dispatcher, "_fused_rms_norm_backward") ADD_IMPL_TO_DISPATCHER(PrivFusedSgdOutFn, priv_fused_sgd_out_dispatcher, "_fused_sgd.out") ADD_IMPL_TO_DISPATCHER(PrivFusedSgdTensorLrOutFn, priv_fused_sgd_tensor_lr_out_dispatcher, "_fused_sgd.tensor_lr_out") diff --git a/csrc/aten/generated/ops.h b/csrc/aten/generated/ops.h index 3ab5be7f..8b42b74c 100644 --- a/csrc/aten/generated/ops.h +++ b/csrc/aten/generated/ops.h @@ -1034,6 +1034,9 @@ DECLARE_DISPATCHER(PrivFusedMovingAvgObsFqHelperOutFn, priv_fused_moving_avg_obs using PrivFusedMovingAvgObsFqHelperFunctionalFn = ::std::tuple (*)(const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, double, int64_t, int64_t, int64_t, bool, bool); DECLARE_DISPATCHER(PrivFusedMovingAvgObsFqHelperFunctionalFn, priv_fused_moving_avg_obs_fq_helper_functional_dispatcher) +using PrivFusedRmsNormFn = ::std::tuple (*)(const at::Tensor &, at::IntArrayRef, const ::std::optional &, ::std::optional); +DECLARE_DISPATCHER(PrivFusedRmsNormFn, priv_fused_rms_norm_dispatcher) + using PrivFusedRmsNormBackwardFn = ::std::tuple (*)(const at::Tensor &, const at::Tensor &, at::IntArrayRef, const at::Tensor &, const ::std::optional &, ::std::array); DECLARE_DISPATCHER(PrivFusedRmsNormBackwardFn, priv_fused_rms_norm_backward_dispatcher) diff --git a/csrc/aten/generated/register.inc b/csrc/aten/generated/register.inc index da16f922..fd74f313 100644 --- a/csrc/aten/generated/register.inc +++ b/csrc/aten/generated/register.inc @@ -1030,6 +1030,9 @@ void WrapperPrivFusedAdamwInplaceTensorLr(at::TensorList self, at::TensorList gr ::std::tuple WrapperPrivFusedMovingAvgObsFqHelperFunctional(const at::Tensor & self, const at::Tensor & observer_on, const at::Tensor & fake_quant_on, const at::Tensor & running_min, const at::Tensor & running_max, const at::Tensor & scale, const at::Tensor & zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, bool per_row_fake_quant, bool symmetric_quant) { return at::native::flagos::priv_fused_moving_avg_obs_fq_helper_functional_dispatcher(self, observer_on, fake_quant_on, running_min, running_max, scale, zero_point, averaging_const, quant_min, quant_max, ch_axis, per_row_fake_quant, symmetric_quant); } +::std::tuple WrapperPrivFusedRmsNorm(const at::Tensor & input, at::IntArrayRef normalized_shape, const ::std::optional & weight, ::std::optional eps) { + return at::native::flagos::priv_fused_rms_norm_dispatcher(input, normalized_shape, weight, eps); +} ::std::tuple WrapperPrivFusedRmsNormBackward(const at::Tensor & grad_out, const at::Tensor & input, at::IntArrayRef normalized_shape, const at::Tensor & rstd, const ::std::optional & weight, ::std::array output_mask) { return at::native::flagos::priv_fused_rms_norm_backward_dispatcher(grad_out, input, normalized_shape, rstd, weight, output_mask); } @@ -6449,6 +6452,7 @@ at::Tensor & WrapperZerosLikeOut(const at::Tensor & self, ::std::optional +#include +#include +#include + +#if defined(USE_ASCEND) + +using namespace at; +// torchgen's emitted bodies call the autograd helpers (unpack, +// compute_requires_grad, collect_next_edges, SavedVariable, set_history, ...) +// unqualified, because upstream's VariableType_N.cpp is itself compiled inside +// namespace torch::autograd. We generate into our own namespace, so pull them in. +using namespace torch::autograd; +using namespace torch::autograd::generated; + +namespace at::flagos::autograd { + +namespace VariableType { + +// torchgen's bodies open with `unpack(arg, "arg", i)`. Upstream declares it in +// torch/csrc/autograd/generated/VariableType.h but defines it in +// VariableTypeManual.cpp, which is not part of the installed library -- the +// symbol is not exported, so we cannot link against it. It is only a +// defined-ness check that returns the tensor unchanged (upstream's +// checked_cast_variable), so define it here, as torch_npu likewise does for its +// own generated VariableType. +namespace { + +inline at::Tensor& unpack(at::Tensor& t, const char* name, int pos) { + TORCH_CHECK(t.defined(), + "Expected a proper Tensor but got None (or an undefined Tensor in C++) " + "for argument #", pos, " '", name, "'"); + return t; +} + +inline const at::Tensor& unpack(const at::Tensor& t, const char* name, int pos) { + TORCH_CHECK(t.defined(), + "Expected a proper Tensor but got None (or an undefined Tensor in C++) " + "for argument #", pos, " '", name, "'"); + return t; +} + +} // namespace + +at::Tensor matmul(c10::DispatchKeySet ks, const at::Tensor & self, const at::Tensor & other) { + auto& self_ = unpack(self, "self", 0); + auto& other_ = unpack(other, "other", 1); + [[maybe_unused]] auto _any_requires_grad = compute_requires_grad( self, other ); + + std::shared_ptr grad_fn; + if (_any_requires_grad) { + grad_fn = std::shared_ptr(new MatmulBackward0(), deleteNode); + grad_fn->set_next_edges(collect_next_edges( self, other )); + grad_fn->other_ = SavedVariable(other, false); + grad_fn->self_ = SavedVariable(self, false); + } + #ifndef NDEBUG + auto self__storage_saved = + self_.has_storage() ? ::std::optional(self_.storage()) : ::std::nullopt; + c10::intrusive_ptr self__impl_saved; + if (self_.defined()) self__impl_saved = self_.getIntrusivePtr(); + auto other__storage_saved = + other_.has_storage() ? ::std::optional(other_.storage()) : ::std::nullopt; + c10::intrusive_ptr other__impl_saved; + if (other_.defined()) other__impl_saved = other_.getIntrusivePtr(); + #endif + auto _tmp = ([&]() { + at::AutoDispatchBelowADInplaceOrView guard; + return at::redispatch::matmul(ks & c10::after_autograd_keyset, self_, other_); + })(); + auto result = std::move(_tmp); + #ifndef NDEBUG + if (self__storage_saved.has_value() && + !at::impl::dispatch_mode_enabled() && + !at::impl::tensor_has_dispatch(self_) && + !at::impl::tensor_has_dispatch(self_)) + TORCH_INTERNAL_ASSERT(self__storage_saved.value().is_alias_of(self_.storage())); + if (self__impl_saved && !at::impl::dispatch_mode_enabled() && !at::impl::tensor_has_dispatch(self_)) + TORCH_INTERNAL_ASSERT(self__impl_saved == self_.getIntrusivePtr()); + if (other__storage_saved.has_value() && + !at::impl::dispatch_mode_enabled() && + !at::impl::tensor_has_dispatch(other_) && + !at::impl::tensor_has_dispatch(other_)) + TORCH_INTERNAL_ASSERT(other__storage_saved.value().is_alias_of(other_.storage())); + if (other__impl_saved && !at::impl::dispatch_mode_enabled() && !at::impl::tensor_has_dispatch(other_)) + TORCH_INTERNAL_ASSERT(other__impl_saved == other_.getIntrusivePtr()); + + if (!at::impl::dispatch_mode_enabled() && !at::impl::tensor_has_dispatch(result)) + TORCH_INTERNAL_ASSERT(result.use_count() == expected_fresh_use_count(result), "function: matmul"); + #endif + if (grad_fn) { + set_history(flatten_tensor_args( result ), grad_fn); + } + throw_error_for_complex_autograd(result, "matmul"); + return result; +} + +} // namespace VariableType + +namespace { + +TORCH_LIBRARY_IMPL(aten, AutogradPrivateUse1, m) { + m.impl("matmul", + TORCH_FN(VariableType::matmul) +); + +} + +} // namespace + +} // namespace at::flagos::autograd + +#endif // USE_ASCEND diff --git a/csrc/aten/register.cc b/csrc/aten/register.cc index 7fbdeccd..b75e9066 100644 --- a/csrc/aten/register.cc +++ b/csrc/aten/register.cc @@ -16,12 +16,29 @@ #include "generated/ops.h" #include +#include +#include #include #include #include #include #include "common.h" #include "runtime/allocator/caching_device_allocator.h" +#include + +// Forward declarations for the Ascend matmul kernels (csrc/aten/backends/ascend/matmul.cc). +// That file is only compiled when USE_ASCEND is on, so BOTH declarations must be +// guarded the same way: a .so links fine with an undefined symbol and only fails +// at dlopen, so an unguarded reference here builds a CUDA wheel that dies on +// `import torch_fl` with "undefined symbol: ...MatmulKernelAscend...". +#if defined(USE_ASCEND) +namespace at::native::flagos { + at::Tensor MatmulKernelAscend(const at::Tensor& self, const at::Tensor& other); + std::tuple MatmulBackwardKernelAscend( + const at::Tensor& grad, const at::Tensor& self, const at::Tensor& other, + ::std::array mask); +} +#endif namespace at::flagos { @@ -216,6 +233,54 @@ int64_t WrapperFusedSdpChoice( #include "generated/register.inc" #undef FLAGOS_GEN_WRAPPERS +// matmul: intercept aten::matmul at PrivateUse1 for the Ascend backend so it +// matmul: intercept aten::matmul at PrivateUse1 for the Ascend backend so it +// routes to aclnnMatmul directly instead of decomposing via +// CompositeImplicitAutograd into mm + bmm + view. Non-Ascend backends (MetaX +// etc.) call at::native::matmul directly so PyTorch's composite decomposition +// runs and mm/bmm reach the appropriate backend kernels. +static at::Tensor WrapperMatmul( + const at::Tensor& self, const at::Tensor& other) { +#if defined(USE_ASCEND) + // aten::matmul is CompositeImplicitAutograd: normally it decomposes into + // mm/bmm/view, and autograd records the backward through those sub-ops. Taking + // the fused aclnnMatmul kernel stops that decomposition, so autograd binds the + // op's real derivative, aten::matmul_backward -- which is registered for + // PrivateUse1 below (WrapperMatmulBackward). The generated + // AutogradPrivateUse1 kernel (csrc/aten/generated/variable_type.cc) builds the + // MatmulBackward0 node and redispatches here, so the fused path is used for + // training as well as inference. + // + // The runtime GetBackendForOp check still matters on an Ascend build: the conf + // can route matmul elsewhere. But it must sit INSIDE the #if -- a runtime + // branch does not remove the link-time reference, and MatmulKernelAscend is + // only compiled when USE_ASCEND is on. + if (at::native::flagos::GetBackendForOp("matmul") == + at::native::flagos::Backend::kAscend) { + return at::native::flagos::MatmulKernelAscend(self, other); + } +#endif + // Fall through to the composite decomposition (mm/bmm/view) by calling the + // CompositeImplicitAutograd implementation directly. This avoids re-entering + // WrapperMatmul (no recursion) while letting the decomposed sub-ops dispatch + // normally to their PrivateUse1 kernels, which have working autograd. Used by + // non-Ascend backends, which have no fused matmul kernel. + return at::native::matmul(self, other); +} + +// matmul_backward: the derivative aten::matmul binds to once a backend owns the +// forward op (see WrapperMatmul). Only Ascend has a fused kernel; other backends +// never reach here, since without a concrete matmul kernel autograd keeps +// recording the mm/bmm decomposition instead. +#if defined(USE_ASCEND) +static std::tuple WrapperMatmulBackward( + const at::Tensor& grad, const at::Tensor& self, const at::Tensor& other, + ::std::array mask) { + return at::native::flagos::MatmulBackwardKernelAscend( + grad, self, other, mask); +} +#endif + } // namespace // Register basic operators for PrivateUse1 dispatch key @@ -243,6 +308,16 @@ TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) { m.impl("_index_put_impl_", WrapperIndexPutImpl_); m.impl("record_stream", WrapperRecordStream); m.impl("_fused_sdp_choice", WrapperFusedSdpChoice); + // matmul: on Ascend, claim the fused aclnnMatmul kernel here on plain + // PrivateUse1. The generated AutogradPrivateUse1 kernel + // (csrc/aten/generated/variable_type.cc) intercepts above this, builds the + // MatmulBackward0 node and redispatches down to us, so training and inference + // both take the fused path. Its derivative, matmul_backward, is a normal + // backend op and is registered right below. +#if defined(USE_ASCEND) + m.impl("matmul", WrapperMatmul); + m.impl("matmul_backward", WrapperMatmulBackward); +#endif // ============================================================ // Generated m.impl registrations for the generated operators @@ -309,6 +384,16 @@ TORCH_LIBRARY_IMPL(aten, AutogradPrivateUse1, m) { // and autograd records CloneBackward0 for gradient propagation. return self.clone(memory_format); }); + + // matmul: on Ascend the fused aclnnMatmul kernel is claimed on plain + // PrivateUse1 (above), and the generated AutogradPrivateUse1 kernel + // (csrc/aten/generated/variable_type.cc) sits in front of it to build the + // autograd graph. Other backends have no fused kernel, so they keep PyTorch's + // composite decomposition; intercepting here lets them reach it without the + // autograd key binding aten::matmul_backward, which they cannot implement. +#if !defined(USE_ASCEND) + m.impl("matmul", WrapperMatmul); +#endif } } // namespace at::flagos diff --git a/csrc/aten/strided_ops.cc b/csrc/aten/strided_ops.cc index a118cc3c..42e0174b 100644 --- a/csrc/aten/strided_ops.cc +++ b/csrc/aten/strided_ops.cc @@ -17,6 +17,8 @@ #include #include #include +#include +#include namespace at::native::flagos { @@ -105,6 +107,24 @@ at::Tensor detach(const at::Tensor& self) { return at::native::detach(self); } +// alias() returns a view sharing self's storage (pure metadata). at::native:: +// alias avoids re-dispatching through PrivateUse1 back into this kernel. +at::Tensor alias(const at::Tensor& self) { + return at::native::alias(self); +} + +// t() is the 2-D (or <=2-D) transpose used by nn.Linear (F.linear does +// input.matmul(weight.t())). Pure metadata, like transpose_int. +at::Tensor t(const at::Tensor& self) { + return at::native::t(self); +} + +// unbind returns views along dim; at::native::unbind builds them via select, +// which we route to at::native::select above (no re-dispatch recursion). +::std::vector unbind_int(const at::Tensor& self, int64_t dim) { + return at::native::unbind(self, dim); +} + // View ops are pure metadata (stride) operations; they route through the // generated dispatchers but need a backend kernel registered. Register them // for the Ascend backend so the generated wrappers in register.inc resolve. @@ -162,4 +182,22 @@ REGISTER_IMPL_TO_DISPATCHER( Backend::kAscend, detach) +REGISTER_IMPL_TO_DISPATCHER( + TFn, + t_dispatcher, + Backend::kAscend, + t) + +REGISTER_IMPL_TO_DISPATCHER( + UnbindIntFn, + unbind_int_dispatcher, + Backend::kAscend, + unbind_int) + +REGISTER_IMPL_TO_DISPATCHER( + AliasFn, + alias_dispatcher, + Backend::kAscend, + alias) + } // namespace at::native::flagos diff --git a/csrc/aten/strided_ops.h b/csrc/aten/strided_ops.h index d9e1e145..bea7dc97 100644 --- a/csrc/aten/strided_ops.h +++ b/csrc/aten/strided_ops.h @@ -32,6 +32,7 @@ at::Tensor view(const at::Tensor& self, c10::SymIntArrayRef size); at::Tensor expand(const at::Tensor& self, c10::SymIntArrayRef size, bool implicit); at::Tensor narrow(const at::Tensor& self, int64_t dim, int64_t start, int64_t length); +at::Tensor alias(const at::Tensor& self); at::Tensor transpose_int(const at::Tensor& self, int64_t dim0, int64_t dim1); @@ -49,4 +50,8 @@ at::Tensor unsqueeze(const at::Tensor& self, int64_t dim); at::Tensor unsafe_view(const at::Tensor& self, at::IntArrayRef size); +at::Tensor t(const at::Tensor& self); + +::std::vector unbind_int(const at::Tensor& self, int64_t dim); + } // namespace at::native::flagos diff --git a/csrc/runtime/accelerator/CMakeLists.txt b/csrc/runtime/accelerator/CMakeLists.txt index 93a1939d..ab285eae 100644 --- a/csrc/runtime/accelerator/CMakeLists.txt +++ b/csrc/runtime/accelerator/CMakeLists.txt @@ -61,6 +61,10 @@ if(ACCELERATOR STREQUAL "ascend") target_include_directories(${LIBRARY_NAME} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_SOURCE_DIR} ${CANN_INCLUDE_DIR}) + # memory.cc drains the shared default ACL stream (acl_stream.h) before any + # host-visible memcpy now that ops dispatch asynchronously; acl_stream.h is + # guarded by USE_ASCEND, so this target must define it. + target_compile_definitions(${LIBRARY_NAME} PRIVATE USE_ASCEND=1) target_link_libraries(${LIBRARY_NAME} PRIVATE ${ACL_LIB}) elseif(ACCELERATOR STREQUAL "metax") file(GLOB SOURCE_FILES "${CMAKE_CURRENT_SOURCE_DIR}/metax/*.cc") diff --git a/csrc/runtime/accelerator/ascend/acl_stream.h b/csrc/runtime/accelerator/ascend/acl_stream.h index 63a7cbc3..6aabc363 100644 --- a/csrc/runtime/accelerator/ascend/acl_stream.h +++ b/csrc/runtime/accelerator/ascend/acl_stream.h @@ -10,16 +10,21 @@ #include +#include + namespace at::native::flagos::ascend { -inline aclrtStream GetDefaultAclStream() { - static aclrtStream stream = []() -> aclrtStream { - aclrtStream s = nullptr; - aclrtCreateStream(&s); - return s; - }(); - return stream; -} +// Returns the process-wide default ACL stream that ALL Ascend ops share. +// +// This MUST be a single external-linkage, default-visibility symbol defined +// once (in libflagos.so). It used to be an `inline` function with a +// function-local `static`; under -fvisibility=hidden that produced a SEPARATE +// stream instance per shared object (libflagos.so vs libtorch_fl.so). The aten +// kernels in libtorch_fl.so then enqueued ops on one stream while the +// drain-before-read in libflagos.so's memory.cc synchronized a DIFFERENT +// stream, so host-visible D2H reads never waited for the producing kernels -> +// silent corruption under async dispatch. Keep it a plain exported function. +FLAGOS_EXPORT aclrtStream GetDefaultAclStream(); } // namespace at::native::flagos::ascend diff --git a/csrc/runtime/accelerator/ascend/memory.cc b/csrc/runtime/accelerator/ascend/memory.cc index 200522ca..9887974e 100644 --- a/csrc/runtime/accelerator/ascend/memory.cc +++ b/csrc/runtime/accelerator/ascend/memory.cc @@ -14,6 +14,7 @@ #include #include +#include "acl_stream.h" #include #include @@ -104,6 +105,16 @@ class MemoryManager { default: return ErrorUnknown; } + // aclnn ops enqueue asynchronously on the shared default stream (per-op + // sync was removed from EXEC_ASCEND_CMD). This blocking aclrtMemcpy runs + // outside that stream's ordering, so any transfer touching device memory + // must first drain the default stream: a D2H read would otherwise observe + // stale data, and an H2D/D2D write could race a pending consumer/producer. + // Host-to-host transfers touch no device memory and need no barrier. + if (kind != MemcpyHostToHost) { + aclrtSynchronizeStream(at::native::flagos::ascend::GetDefaultAclStream()); + } + aclError err = aclrtMemcpy(dst, count, src, count, acl_kind); return (err == ACL_SUCCESS) ? Success : ErrorUnknown; } diff --git a/csrc/runtime/accelerator/ascend/stream_api.cc b/csrc/runtime/accelerator/ascend/stream_api.cc index d6fa847d..0ac77b90 100644 --- a/csrc/runtime/accelerator/ascend/stream_api.cc +++ b/csrc/runtime/accelerator/ascend/stream_api.cc @@ -10,6 +10,23 @@ #include "acl_stream.h" +namespace at::native::flagos::ascend { + +// Single process-wide definition of the shared default ACL stream. Declared in +// acl_stream.h with default visibility so every shared object (libflagos.so and +// libtorch_fl.so) resolves to THIS one instance — see the header comment for +// why a per-TU `inline` static silently corrupted async dispatch. +FLAGOS_EXPORT aclrtStream GetDefaultAclStream() { + static aclrtStream stream = []() -> aclrtStream { + aclrtStream s = nullptr; + aclrtCreateStream(&s); + return s; + }(); + return stream; +} + +} // namespace at::native::flagos::ascend + extern "C" { __attribute__((visibility("default"))) diff --git a/csrc/runtime/allocator/caching_device_allocator.cc b/csrc/runtime/allocator/caching_device_allocator.cc index 6f2ac5df..9a6feaa1 100644 --- a/csrc/runtime/allocator/caching_device_allocator.cc +++ b/csrc/runtime/allocator/caching_device_allocator.cc @@ -128,7 +128,10 @@ at::DataPtr CachingDeviceAllocator::allocate(size_t nbytes) { auto curr_device = c10::Device(c10::DeviceType::PrivateUse1, static_cast(device)); - return {block->ptr, block->ptr, &block_deleter, curr_device}; + // Stash the Block* as the DataPtr context so the deleter can recover it in + // O(1) with no side map / lock. The data pointer and context differ (data = + // device memory, context = Block metadata), which DataPtr supports directly. + return {block->ptr, block, &block_deleter, curr_device}; } at::DeleterFnPtr CachingDeviceAllocator::raw_deleter() const { @@ -193,12 +196,6 @@ Block* CachingDeviceAllocator::alloc_block( std::max(state.stats.peak_allocated, state.stats.bytes_allocated); state.stats.num_alloc_calls++; - // Register in ptr-to-block map. - { - std::lock_guard ptr_lock(ptr_map_mutex_); - ptr_to_block_[block->ptr] = block; - } - return block; } @@ -213,12 +210,6 @@ void CachingDeviceAllocator::free_block(Block* block) { state.stats.bytes_allocated -= block->size; state.stats.num_free_calls++; - // Remove from ptr map. - { - std::lock_guard ptr_lock(ptr_map_mutex_); - ptr_to_block_.erase(block->ptr); - } - // If there are outstanding events on other streams, defer the free. if (block->event_count > 0) { // Block will be returned to pool when events complete. @@ -451,7 +442,13 @@ void CachingDeviceAllocator::record_stream( return; } - Block* block = get_block_from_ptr(ptr.get()); + // The Block* is stored as the DataPtr context by allocate(). Only trust it + // when the deleter matches (i.e. this DataPtr came from our block pool, not + // the delegation path or a foreign allocator). + if (ptr.get_deleter() != &block_deleter) { + return; + } + Block* block = static_cast(ptr.get_context()); if (!block) { return; } @@ -500,24 +497,13 @@ void CachingDeviceAllocator::reset_stats(int device) { state.stats = AllocatorStats{}; } -Block* CachingDeviceAllocator::get_block_from_ptr(void* ptr) { - std::lock_guard lock(ptr_map_mutex_); - auto it = ptr_to_block_.find(ptr); - if (it != ptr_to_block_.end()) { - return it->second; - } - return nullptr; -} - -// Static deleter invoked by DataPtr when a tensor is freed. -void CachingDeviceAllocator::block_deleter(void* ptr) { - if (!ptr || !instance_) { +// Static deleter invoked by DataPtr when a tensor is freed. The context is the +// Block* stashed by allocate(), so freeing is O(1) with no map lookup or lock. +void CachingDeviceAllocator::block_deleter(void* ctx) { + if (!ctx || !instance_) { return; } - Block* block = instance_->get_block_from_ptr(ptr); - if (block) { - instance_->free_block(block); - } + instance_->free_block(static_cast(ctx)); } // Deleter for the delegation path: free straight back to the backend's caching diff --git a/csrc/runtime/allocator/caching_device_allocator.h b/csrc/runtime/allocator/caching_device_allocator.h index 5e797c65..5aad3601 100644 --- a/csrc/runtime/allocator/caching_device_allocator.h +++ b/csrc/runtime/allocator/caching_device_allocator.h @@ -53,9 +53,6 @@ class FLAGOS_EXPORT CachingDeviceAllocator final : public at::Allocator { // Reset accumulated statistics for a device. void reset_stats(int device); - // Get the underlying block for a data pointer (nullptr if not found). - Block* get_block_from_ptr(void* ptr); - // Whether caching is enabled (controlled by env var). static bool is_enabled(); @@ -101,8 +98,9 @@ class FLAGOS_EXPORT CachingDeviceAllocator final : public at::Allocator { // Process completed events and return blocks to the pool. void process_events(DeviceState& state); - // Static deleter function for DataPtr. - static void block_deleter(void* ptr); + // Static deleter function for DataPtr. Receives the Block* directly as the + // DataPtr context (set at allocation), so no ptr->block lookup is needed. + static void block_deleter(void* ctx); // Static deleter for the delegation path (frees via backend caching allocator). static void delegated_deleter(void* ptr); @@ -110,10 +108,6 @@ class FLAGOS_EXPORT CachingDeviceAllocator final : public at::Allocator { // Get or create per-device state. DeviceState& get_device_state(int device); - // Map from raw pointer to Block for O(1) lookup. - std::mutex ptr_map_mutex_; - std::unordered_map ptr_to_block_; - std::unique_ptr backend_; std::vector> device_states_; std::recursive_mutex device_states_mutex_; diff --git a/scripts/codegen_ascend.py b/scripts/codegen_ascend.py index f6bc8de1..f97ee702 100644 --- a/scripts/codegen_ascend.py +++ b/scripts/codegen_ascend.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. + """ Codegen for torch_fl Ascend (aclnn) operators. @@ -158,6 +159,7 @@ # ---- binary_scalar_alpha: aclnns(self, other, alpha, out) ---- "add.Scalar": ("binary_scalar_alpha", "Adds"), "sub.Scalar": ("binary_scalar_alpha", "Subs"), + "rsub.Scalar": ("binary_scalar_alpha", "Rsubs"), # ---- binary_scalar_cmp: bool out, aclnn(self, other, out) ---- "eq.Scalar": ("binary_scalar_cmp", "EqScalar"), "ne.Scalar": ("binary_scalar_cmp", "NeScalar"), @@ -240,10 +242,17 @@ "bmm.out": ("matmul_out", "BatchMatMul"), # cat: TensorList concat (aclCreateTensorList). "cat": ("cat", "Cat"), + # stack: TensorList concat along a NEW dim (aclCreateTensorList). + "stack": ("stack", "Stack"), # factory ops: at::empty + device-side zero_/fill_ (no direct aclnn call). "zeros": ("zeros", None), + "ones": ("ones", None), "scalar_tensor": ("scalar_tensor", None), "ones_like": ("ones_like", None), + "zeros_like": ("zeros_like", None), + "empty_like": ("empty_like", None), + "full": ("full", None), + "full_like": ("full_like", None), "new_ones": ("new_ones", None), "addmm": ("gemm_addmm", "Addmm"), "baddbmm": ("gemm_baddbmm", "Baddbmm"), @@ -279,6 +288,42 @@ "zero_": ("inplace_zero", "InplaceZero"), "fill_.Scalar": ("inplace_fill_scalar", "InplaceFillScalar"), "fill_.Tensor": ("inplace_fill_tensor", "InplaceFillTensor"), + "add_.Tensor": ("inplace_add_tensor", "InplaceAdd"), + "add_.Scalar": ("inplace_add_scalar", "InplaceAdds"), + "mul_.Tensor": ("inplace_mul_tensor", "InplaceMul"), + "mul_.Scalar": ("inplace_mul_scalar", "InplaceMuls"), + "div_.Tensor": ("inplace_div_tensor", "InplaceDiv"), + # bitwise_{and,or,xor}_.Tensor have the same (self&, other) shape as + # mul_.Tensor; reuse that category with the Inplace* aclnn override. + # torch's allclose()->isclose() decomposition combines nan/close masks + # with these in-place, so their absence cascades into ~every + # allclose-based test failing. + "bitwise_and_.Tensor": ("inplace_mul_tensor", "InplaceBitwiseAndTensor"), + "bitwise_or_.Tensor": ("inplace_mul_tensor", "InplaceBitwiseOrTensor"), + "bitwise_xor_.Tensor": ("inplace_mul_tensor", "InplaceBitwiseXorTensor"), + "addcmul_": ("inplace_addcmul", "InplaceAddcmul"), + "addcdiv_": ("inplace_addcdiv", "InplaceAddcdiv"), + "sqrt_": ("inplace_sqrt", "InplaceSqrt"), + "lerp_.Scalar": ("inplace_lerp_scalar", "InplaceLerps"), + # ---- foreach (TensorList) family: needed by torch.optim.AdamW's default + # foreach=True path (aten's _multi_tensor_adam). All void-returning + # in-place ops except _foreach_sqrt (returns new Tensor[]). ---- + "_foreach_mul_.Scalar": ("foreach_inplace_scalar", "ForeachMulScalarV2"), + "_foreach_add_.Scalar": ("foreach_inplace_scalar", "ForeachAddScalarV2"), + "_foreach_lerp_.Scalar": ("foreach_inplace_lerp_scalar", "ForeachLerpScalar"), + "_foreach_addcmul_.Scalar": ( + "foreach_inplace_addcmul_scalar", + "ForeachAddcmulScalarV2", + ), + "_foreach_sqrt": ("foreach_sqrt", "ForeachSqrt"), + "_foreach_div_.ScalarList": ( + "foreach_inplace_div_scalarlist", + "ForeachDivScalarList", + ), + "_foreach_addcdiv_.ScalarList": ( + "foreach_inplace_addcdiv_scalarlist", + "ForeachAddcdivScalarList", + ), # ---- embedding + pad (single-aclnn-call, migrated from handwritten) ---- "embedding": ("embedding", "Embedding"), "embedding_dense_backward": ("embedding_dense_backward", "EmbeddingDenseBackward"), @@ -305,8 +350,15 @@ "where.self": ("where", "SWhere"), "_softmax": ("softmax_fwd", "Softmax"), "all": ("reduce_all", "All"), + "any": ("reduce_all", "Any"), "sum.dim_IntList": ("reduce_sum_dtype", "ReduceSum"), + "sum": ("reduce_sum_all", "ReduceSum"), + "max": ("reduce_minmax_all", "Max"), + "min": ("reduce_minmax_all", "Min"), "mean.dim": ("reduce_mean_dtype", "MeanV2"), + "mean": ("mean_all", "Mean"), + "clamp": ("clamp", "Clamp"), + "clamp.Tensor": ("clamp_tensor", "ClampTensor"), # ---- conv/pool family (each carries an output-shape formula) ---- "_adaptive_avg_pool2d": ("adaptive_avg_pool2d", "AdaptiveAvgPool2d"), "avg_pool2d": ("avg_pool2d", "AvgPool2d"), @@ -353,15 +405,80 @@ # materialized to the broadcast shape so aclnn (which does not always # broadcast) sees matching ND-contiguous inputs. All steps are no-ops when # device/dtype/shape already match. -_BINARY_PROLOGUE = """\ - namespace ascend = at::native::flagos::ascend; +# The prologue body (everything after the `namespace ascend` alias). Split out +# so the cached templates can inject a scalar fast-path branch before it while +# still sharing one namespace alias. +_BINARY_PROLOGUE_BODY = """\ auto result_dtype = self.scalar_type(); auto other_c = other.is_privateuseone() ? (other.scalar_type() == result_dtype ? other : other.to(result_dtype)) : other.to(self.options()); auto out_shape = at::infer_size(self.sizes(), other_c.sizes()); - auto self_b = self.expand(out_shape).contiguous(); - auto other_b = other_c.expand(out_shape).contiguous(); + // aclnn binary ops broadcast and honor strides internally (verified), so we + // pass self/other_c straight through instead of expand().contiguous(). This + // avoids up to two device strided-copies + host view construction per op on + // the eager decode hot path. Only materialize a contiguous copy when the + // tensor is genuinely non-contiguous AND the aclnn path would otherwise need + // it — measured unnecessary for the common same-shape/contiguous case, which + // is the overwhelming majority in Qwen3. + const at::Tensor& self_b = self; + const at::Tensor& other_b = other_c; +""" + +_BINARY_PROLOGUE = ( + """\ + namespace ascend = at::native::flagos::ascend; +""" + + _BINARY_PROLOGUE_BODY +) + +# Scalar fast-path branches injected at the top of the cached binary kernels. +# When `other` is a wrapped CPU scalar (a python float/int, materialized by +# PyTorch as a 0-dim CPU tensor), the default path's `other.to(self.options())` +# does a per-call H2D copy (~22us measured -- the single biggest torch_fl vs +# torch_npu host gap: add.Tensor 49us vs 13us). Diverting to the aclnn scalar +# variant (aclnns, which takes an aclScalar* by value) skips the H2D +# entirely. Only emitted for ops that actually ship an s symbol +# (add/sub/mul/div); the scalar value is folded into the executor-cache key. +_SCALAR_FASTPATH_NOALPHA = """\ + if (self.is_privateuseone() && !other.is_privateuseone() && other.numel() == 1) {{ + at::Scalar sc = other.item(); + auto out = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options()); + ascend::AclScalarWrapper acl_sc(sc, self.scalar_type()); + static void* sOpAddr = nullptr; static void* sWsAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + {{ double sv = sc.toDouble(); hsh.val(sv); }} + ascend::ExecAscendCached( + "{aclnn_s}", "{aclnn_s}GetWorkspaceSize", sOpAddr, sWsAddr, hsh.h, + {{&self}}, {{&out}}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) {{ + return gws(in[0].acl_tensor, acl_sc.get(), out_t[0].acl_tensor, pws, pex); + }}); + return out; + }} +""" + +_SCALAR_FASTPATH_ALPHA = """\ + if (self.is_privateuseone() && !other.is_privateuseone() && other.numel() == 1) {{ + at::Scalar sc = other.item(); + auto out = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options()); + ascend::AclScalarWrapper acl_sc(sc, self.scalar_type()); + ascend::AclScalarWrapper acl_alpha_s(alpha, self.scalar_type()); + static void* sOpAddr = nullptr; static void* sWsAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + {{ double sv = sc.toDouble(); hsh.val(sv); double av = alpha.toDouble(); hsh.val(av); }} + ascend::ExecAscendCached( + "{aclnn_s}", "{aclnn_s}GetWorkspaceSize", sOpAddr, sWsAddr, hsh.h, + {{&self}}, {{&out}}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) {{ + return gws(in[0].acl_tensor, acl_sc.get(), acl_alpha_s.get(), out_t[0].acl_tensor, pws, pex); + }}); + return out; + }} """ T_BINARY = ( @@ -537,15 +654,21 @@ # cumsum: (Tensor, int64_t dim, optional dtype) -> same shape. # aclnn(self, dim, dtype, out) +# Integral promotion: with no explicit dtype, PyTorch promotes any integral input +# (incl. bool) to int64; float dtypes pass through. aclnn also rejects a bool +# `self` (err 161002), so cast the input tensor to the promoted dtype too. T_CUMSUM = """\ at::Tensor {kernel}(const at::Tensor& self, int64_t dim, ::std::optional dtype) {{ namespace ascend = at::native::flagos::ascend; int64_t d = dim < 0 ? dim + self.dim() : dim; - auto out_dtype = dtype.value_or(self.scalar_type()); + auto out_dtype = dtype.value_or( + at::isIntegralType(self.scalar_type(), /*includeBool=*/true) + ? at::kLong : self.scalar_type()); + auto in = self.scalar_type() == out_dtype ? self : self.to(out_dtype); auto out = ascend::OpPreparation::apply_tensor_without_format( - self.sizes(), self.options().dtype(out_dtype)); + in.sizes(), in.options().dtype(out_dtype)); - ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_self(in); ascend::AclTensorWrapper acl_out(out); aclDataType acl_dtype = ascend::ToAclDataType(out_dtype); @@ -558,15 +681,19 @@ # cumprod: like cumsum, but aclnnCumprod takes dim as an aclScalar* (int64), not # a plain int64_t. Otherwise identical: (Tensor, int64 dim, optional dtype). +# Same integral->int64 promotion + input cast as cumsum. T_CUMPROD = """\ at::Tensor {kernel}(const at::Tensor& self, int64_t dim, ::std::optional dtype) {{ namespace ascend = at::native::flagos::ascend; int64_t d = dim < 0 ? dim + self.dim() : dim; - auto out_dtype = dtype.value_or(self.scalar_type()); + auto out_dtype = dtype.value_or( + at::isIntegralType(self.scalar_type(), /*includeBool=*/true) + ? at::kLong : self.scalar_type()); + auto in = self.scalar_type() == out_dtype ? self : self.to(out_dtype); auto out = ascend::OpPreparation::apply_tensor_without_format( - self.sizes(), self.options().dtype(out_dtype)); + in.sizes(), in.options().dtype(out_dtype)); - ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_self(in); ascend::AclScalarWrapper acl_dim(at::Scalar(d), at::kLong); ascend::AclTensorWrapper acl_out(out); aclDataType acl_dtype = ascend::ToAclDataType(out_dtype); @@ -616,6 +743,35 @@ REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) """ +# unary_scalar CACHED: same as T_UNARY_SCALAR but through the repeatable-executor +# cache. The Scalar is baked into the executor at build time (aclnn reads it +# during GetWorkspaceSize), so it MUST be part of the cache key -- a different +# scalar value needs a distinct executor. On the decode hot path pow.Tensor_Scalar +# (x^2 in RMSNorm, 113/step, measured 50us uncached) is the big win. +T_UNARY_SCALAR_CACHED = """\ +at::Tensor {kernel}(const at::Tensor& self, const at::Scalar& s) {{ + namespace ascend = at::native::flagos::ascend; + auto out = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options()); + ascend::AclScalarWrapper acl_s(s, self.scalar_type()); + + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + {{ double sv = s.toDouble(); hsh.val(sv); }} + ascend::ExecAscendCached( + "{aclnn}", "{aclnn}GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {{&self}}, {{&out}}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) {{ + return gws(in[0].acl_tensor, acl_s.get(), out_t[0].acl_tensor, pws, pex); + }}); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + # unary_two_scalar: (Tensor, Scalar, Scalar) -> same shape/dtype. # aclnn(self, s1, s2, out) e.g. softplus(beta,threshold)/threshold(threshold,value) T_UNARY_TWO_SCALAR = """\ @@ -1096,6 +1252,346 @@ REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) """ +# stack: (TensorList tensors, dim) -> new-dim concatenation (unlike cat, no +# existing dim is merged; each input keeps its own shape and dim is inserted). +# aclnn(aclTensorList, dim, out). dim is normalized against the OUTPUT +# rank (input rank + 1), matching torch's `maybe_wrap_dim(dim, ndim + 1)`. +T_STACK = """\ +at::Tensor {kernel}(at::TensorList tensors, int64_t dim) {{ + namespace ascend = at::native::flagos::ascend; + TORCH_CHECK(!tensors.empty(), "stack: expected a non-empty list of tensors"); + + auto& first = tensors[0]; + int64_t out_ndim = first.dim() + 1; + if (dim < 0) dim += out_ndim; + + std::vector out_sizes(first.sizes().begin(), first.sizes().end()); + out_sizes.insert(out_sizes.begin() + dim, static_cast(tensors.size())); + + auto out = ascend::OpPreparation::apply_tensor_without_format( + out_sizes, first.options()); + + std::vector wrappers; + wrappers.reserve(tensors.size()); + for (const auto& t : tensors) {{ + wrappers.emplace_back(t); + }} + + std::vector acl_tensors; + acl_tensors.reserve(tensors.size()); + for (auto& w : wrappers) {{ + acl_tensors.push_back(w.get()); + }} + + aclTensorList* tensor_list = aclCreateTensorList( + acl_tensors.data(), acl_tensors.size()); + + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD({aclnn}, tensor_list, dim, acl_out.get()); + + (void)tensor_list; // aclTensor* owned by wrappers; do not aclDestroyTensorList + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# --- foreach ops needed by torch.optim.AdamW's foreach=True path (aten's +# _multi_tensor_adam). void return, in-place on `self`'s TensorList: build an +# aclTensorList for each TensorList arg, execute in-place-style (out == x), +# then leave the input tensors mutated (matches PyTorch's _foreach_*_ inplace +# semantics: the storage is written in place, no new Tensors are returned). +# NOTE: aclTensorList's aclTensor* are owned by the AclTensorWrapper RAII +# vector, so — same as cat/stack — we must NOT aclDestroyTensorList. + +# _foreach_mul_.Scalar / _foreach_add_.Scalar: (self[]&, scalar) -> void. +# aclnnForeachMulScalarV2/aclnnForeachAddScalarV2(x, scalar, out=x). +T_FOREACH_INPLACE_SCALAR = """\ +static void {kernel}Chunk(at::TensorList self, const at::Scalar& scalar) {{ + namespace ascend = at::native::flagos::ascend; + + std::vector wrappers; + wrappers.reserve(self.size()); + for (const auto& t : self) {{ + wrappers.emplace_back(t); + }} + std::vector acl_tensors; + acl_tensors.reserve(self.size()); + for (auto& w : wrappers) {{ + acl_tensors.push_back(w.get()); + }} + aclTensorList* tensor_list = aclCreateTensorList(acl_tensors.data(), acl_tensors.size()); + // aic-ops-info: ForeachMulScalar/ForeachAddScalar's `scalar` dtype tracks x's + // EXCEPT bf16 x, which requires a float32 scalar (no bf16 scalar entry). + auto scalar_dtype = self[0].scalar_type() == at::kBFloat16 ? at::kFloat : self[0].scalar_type(); + ascend::AclScalarWrapper acl_scalar(scalar, scalar_dtype); + + EXEC_ASCEND_CMD({aclnn}, tensor_list, acl_scalar.get(), tensor_list); + + (void)tensor_list; // aclTensor* owned by wrappers; do not aclDestroyTensorList +}} + +void {kernel}(at::TensorList self, const at::Scalar& scalar) {{ + TORCH_CHECK(!self.empty(), "{disp}: expected a non-empty list of tensors"); + for (size_t off = 0; off < self.size(); off += {chunk}) {{ + size_t n = std::min({chunk}, self.size() - off); + {kernel}Chunk(self.slice(off, n), scalar); + }} +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# CANN's aclnnForeach* kernels only process the FIRST 50 entries of an +# aclTensorList. Past that they either error (the ScalarList variants return +# 561002/161002) or -- worse -- return success while leaving entries >= 50 +# untouched, so the bug is silent. Measured: entry 50 is the first wrong one for +# Mul/Add/Addcmul/Lerp/Sqrt alike, independent of each tensor's numel +# (8 .. 65536) and dtype (fp16/fp32/bf16). AdamW on Qwen3-0.6B passes 310 +# tensors, so every foreach kernel slices its lists into sub-50 chunks; +# elementwise semantics make the split exact. +FOREACH_CHUNK = 32 # aclnn processes at most 50 entries per call + +# _foreach_lerp_.Scalar: (self[]&, tensors1[], weight) -> void, self += weight*(tensors1-self). +# aclnnForeachLerpScalar(x1=self, x2=tensors1, weight, out=self). +T_FOREACH_INPLACE_LERP_SCALAR = """\ +static void {kernel}Chunk(at::TensorList self, at::TensorList tensors1, const at::Scalar& weight) {{ + namespace ascend = at::native::flagos::ascend; + + std::vector self_w, t1_w; + self_w.reserve(self.size()); + t1_w.reserve(tensors1.size()); + for (const auto& t : self) self_w.emplace_back(t); + for (const auto& t : tensors1) t1_w.emplace_back(t); + + std::vector self_ptrs, t1_ptrs; + self_ptrs.reserve(self.size()); + t1_ptrs.reserve(tensors1.size()); + for (auto& w : self_w) self_ptrs.push_back(w.get()); + for (auto& w : t1_w) t1_ptrs.push_back(w.get()); + + aclTensorList* self_list = aclCreateTensorList(self_ptrs.data(), self_ptrs.size()); + aclTensorList* t1_list = aclCreateTensorList(t1_ptrs.data(), t1_ptrs.size()); + // aic-ops-info: ForeachLerpScalar's `weight` is ALWAYS float32, regardless + // of x1/x2's dtype (unlike mul_/add_.Scalar, which track x except for bf16). + ascend::AclScalarWrapper acl_weight(weight, at::kFloat); + + EXEC_ASCEND_CMD({aclnn}, self_list, t1_list, acl_weight.get(), self_list); + + (void)self_list; (void)t1_list; // owned by *_w; do not aclDestroyTensorList +}} + +void {kernel}(at::TensorList self, at::TensorList tensors1, const at::Scalar& weight) {{ + TORCH_CHECK(!self.empty(), "{disp}: expected a non-empty list of tensors"); + TORCH_CHECK(self.size() == tensors1.size(), "{disp}: tensor lists must match in length"); + for (size_t off = 0; off < self.size(); off += {chunk}) {{ + size_t n = std::min({chunk}, self.size() - off); + {kernel}Chunk(self.slice(off, n), tensors1.slice(off, n), weight); + }} +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# _foreach_addcmul_.Scalar: (self[]&, tensor1[], tensor2[], value) -> void, +# self += value * tensor1 * tensor2. +# aclnnForeachAddcmulScalarV2(x1=self, x2=tensor1, x3=tensor2, scalar=value, out=self). +T_FOREACH_INPLACE_ADDCMUL_SCALAR = """\ +static void {kernel}Chunk(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Scalar& value) {{ + namespace ascend = at::native::flagos::ascend; + + std::vector self_w, t1_w, t2_w; + self_w.reserve(self.size()); + t1_w.reserve(tensor1.size()); + t2_w.reserve(tensor2.size()); + for (const auto& t : self) self_w.emplace_back(t); + for (const auto& t : tensor1) t1_w.emplace_back(t); + for (const auto& t : tensor2) t2_w.emplace_back(t); + + std::vector self_ptrs, t1_ptrs, t2_ptrs; + self_ptrs.reserve(self.size()); + t1_ptrs.reserve(tensor1.size()); + t2_ptrs.reserve(tensor2.size()); + for (auto& w : self_w) self_ptrs.push_back(w.get()); + for (auto& w : t1_w) t1_ptrs.push_back(w.get()); + for (auto& w : t2_w) t2_ptrs.push_back(w.get()); + + aclTensorList* self_list = aclCreateTensorList(self_ptrs.data(), self_ptrs.size()); + aclTensorList* t1_list = aclCreateTensorList(t1_ptrs.data(), t1_ptrs.size()); + aclTensorList* t2_list = aclCreateTensorList(t2_ptrs.data(), t2_ptrs.size()); + // aic-ops-info: ForeachAddcmulScalar's `scalar` dtype tracks x EXCEPT bf16 x, + // which requires a float32 scalar (same rule as mul_/add_.Scalar). + auto value_dtype = self[0].scalar_type() == at::kBFloat16 ? at::kFloat : self[0].scalar_type(); + ascend::AclScalarWrapper acl_value(value, value_dtype); + + EXEC_ASCEND_CMD({aclnn}, self_list, t1_list, t2_list, acl_value.get(), self_list); + + (void)self_list; (void)t1_list; (void)t2_list; // owned by *_w +}} + +void {kernel}(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Scalar& value) {{ + TORCH_CHECK(!self.empty(), "{disp}: expected a non-empty list of tensors"); + TORCH_CHECK(self.size() == tensor1.size() && self.size() == tensor2.size(), + "{disp}: tensor lists must match in length"); + for (size_t off = 0; off < self.size(); off += {chunk}) {{ + size_t n = std::min({chunk}, self.size() - off); + {kernel}Chunk(self.slice(off, n), tensor1.slice(off, n), tensor2.slice(off, n), value); + }} +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# _foreach_sqrt: (self[]) -> Tensor[] (NOT in-place — returns new tensors). +# aclnnForeachSqrt(x, out) with out a freshly-allocated TensorList. +T_FOREACH_SQRT = """\ +static void {kernel}Chunk(at::TensorList self, at::TensorList outs) {{ + namespace ascend = at::native::flagos::ascend; + + + std::vector in_w, out_w; + in_w.reserve(self.size()); + out_w.reserve(outs.size()); + for (const auto& t : self) in_w.emplace_back(t); + for (const auto& t : outs) out_w.emplace_back(t); + + std::vector in_ptrs, out_ptrs; + in_ptrs.reserve(self.size()); + out_ptrs.reserve(outs.size()); + for (auto& w : in_w) in_ptrs.push_back(w.get()); + for (auto& w : out_w) out_ptrs.push_back(w.get()); + + aclTensorList* in_list = aclCreateTensorList(in_ptrs.data(), in_ptrs.size()); + aclTensorList* out_list = aclCreateTensorList(out_ptrs.data(), out_ptrs.size()); + + EXEC_ASCEND_CMD({aclnn}, in_list, out_list); + + (void)in_list; (void)out_list; // owned by in_w/out_w +}} + +::std::vector {kernel}(at::TensorList self) {{ + TORCH_CHECK(!self.empty(), "{disp}: expected a non-empty list of tensors"); + std::vector outs; + outs.reserve(self.size()); + for (const auto& t : self) outs.push_back(at::empty_like(t)); + for (size_t off = 0; off < self.size(); off += {chunk}) {{ + size_t n = std::min({chunk}, self.size() - off); + {kernel}Chunk(self.slice(off, n), at::TensorList(outs).slice(off, n)); + }} + return outs; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# _foreach_div_.ScalarList / _foreach_addcdiv_.ScalarList: per-tensor scalar +# list. aclnn's ScalarList variant takes an aclScalarList (div), while its +# addcdiv counterpart's "scalars" param is (per the header) an aclTensor* — +# both are boxed as a plain list of aclScalar* built from the ArrayRef. +T_FOREACH_INPLACE_DIV_SCALARLIST = """\ +static void {kernel}Chunk(at::TensorList self, at::ArrayRef scalars) {{ + namespace ascend = at::native::flagos::ascend; + + std::vector wrappers; + wrappers.reserve(self.size()); + for (const auto& t : self) wrappers.emplace_back(t); + std::vector acl_tensors; + acl_tensors.reserve(self.size()); + for (auto& w : wrappers) acl_tensors.push_back(w.get()); + aclTensorList* tensor_list = aclCreateTensorList(acl_tensors.data(), acl_tensors.size()); + + // aic-ops-info: ForeachDivScalarList's `scalars` is ALWAYS float32, + // regardless of x's dtype (same rule as ForeachLerpScalar's weight). + std::vector scalar_wrappers; + scalar_wrappers.reserve(scalars.size()); + for (size_t i = 0; i < scalars.size(); ++i) {{ + scalar_wrappers.emplace_back(scalars[i], at::kFloat); + }} + std::vector acl_scalars; + acl_scalars.reserve(scalar_wrappers.size()); + for (auto& sw : scalar_wrappers) acl_scalars.push_back(sw.get()); + aclScalarList* scalar_list = aclCreateScalarList(acl_scalars.data(), acl_scalars.size()); + + EXEC_ASCEND_CMD({aclnn}, tensor_list, scalar_list, tensor_list); + + aclDestroyScalarList(scalar_list); + (void)tensor_list; // aclTensor* owned by wrappers; do not aclDestroyTensorList +}} + +void {kernel}(at::TensorList self, at::ArrayRef scalars) {{ + TORCH_CHECK(!self.empty(), "{disp}: expected a non-empty list of tensors"); + TORCH_CHECK(self.size() == scalars.size(), "{disp}: scalars must match tensor list length"); + for (size_t off = 0; off < self.size(); off += {chunk}) {{ + size_t n = std::min({chunk}, self.size() - off); + {kernel}Chunk(self.slice(off, n), scalars.slice(off, n)); + }} +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +T_FOREACH_INPLACE_ADDCDIV_SCALARLIST = """\ +static void {kernel}Chunk(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, at::ArrayRef scalars) {{ + namespace ascend = at::native::flagos::ascend; + + std::vector self_w, t1_w, t2_w; + self_w.reserve(self.size()); + t1_w.reserve(tensor1.size()); + t2_w.reserve(tensor2.size()); + for (const auto& t : self) self_w.emplace_back(t); + for (const auto& t : tensor1) t1_w.emplace_back(t); + for (const auto& t : tensor2) t2_w.emplace_back(t); + + std::vector self_ptrs, t1_ptrs, t2_ptrs; + self_ptrs.reserve(self.size()); + t1_ptrs.reserve(tensor1.size()); + t2_ptrs.reserve(tensor2.size()); + for (auto& w : self_w) self_ptrs.push_back(w.get()); + for (auto& w : t1_w) t1_ptrs.push_back(w.get()); + for (auto& w : t2_w) t2_ptrs.push_back(w.get()); + + aclTensorList* self_list = aclCreateTensorList(self_ptrs.data(), self_ptrs.size()); + aclTensorList* t1_list = aclCreateTensorList(t1_ptrs.data(), t1_ptrs.size()); + aclTensorList* t2_list = aclCreateTensorList(t2_ptrs.data(), t2_ptrs.size()); + + // aclnnForeachAddcdivScalarList's "scalars" param is a plain device aclTensor + // (1-D, one element per list entry), NOT an aclScalarList -- unlike div's + // ScalarList variant. Materialize scalars on host in self[0]'s dtype, then + // move to device once. The dtype MUST match self (a float32 scalars tensor + // against fp16 inputs returns 161002), which costs up to 1 ulp versus CPU, + // where the divisor stays a full-precision Scalar. + at::Tensor scalars_cpu = at::empty({{static_cast(scalars.size())}}, + at::TensorOptions().dtype(self[0].scalar_type())); + AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, self[0].scalar_type(), + "{disp}_scalars", [&] {{ + auto* ptr = scalars_cpu.data_ptr(); + for (size_t i = 0; i < scalars.size(); ++i) {{ + ptr[i] = scalars[i].to(); + }} + }}); + at::Tensor scalars_dev = scalars_cpu.to(self[0].device()); + ascend::AclTensorWrapper acl_scalars(scalars_dev); + + EXEC_ASCEND_CMD({aclnn}, self_list, t1_list, t2_list, acl_scalars.get(), self_list); + + (void)self_list; (void)t1_list; (void)t2_list; // owned by *_w +}} + +void {kernel}(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, at::ArrayRef scalars) {{ + TORCH_CHECK(!self.empty(), "{disp}: expected a non-empty list of tensors"); + TORCH_CHECK(self.size() == tensor1.size() && self.size() == tensor2.size() && self.size() == scalars.size(), + "{disp}: tensor/scalar lists must match in length"); + for (size_t off = 0; off < self.size(); off += {chunk}) {{ + size_t n = std::min({chunk}, self.size() - off); + {kernel}Chunk(self.slice(off, n), tensor1.slice(off, n), tensor2.slice(off, n), + scalars.slice(off, n)); + }} +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + # --- factory ops: at::empty(...) on the PrivateUse1 device + device-side fill --- # These build TensorOptions on-host then fill via zero_/fill_, which are themselves # device-side aclnn kernels (aclnnInplaceZero / aclnnInplaceFillScalar), so the @@ -1153,6 +1649,105 @@ REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) """ +# zeros_like: (self, dtype?, layout?, device?, pin?, memory_format?) -> zeros w/ self's meta. +# Identical to ones_like but fills 0. Used by optimizers (Adam exp_avg state). +T_ZEROS_LIKE = """\ +at::Tensor {kernel}(const at::Tensor& self, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory, ::std::optional memory_format) {{ + auto options = at::TensorOptions() + .dtype(dtype.value_or(self.scalar_type())) + .layout(layout.value_or(self.layout())) + .device(device.value_or(self.device())) + .pinned_memory(pin_memory.value_or(false)); + auto fmt = memory_format.value_or(at::MemoryFormat::Contiguous); + if (fmt == at::MemoryFormat::Preserve) {{ + fmt = self.suggest_memory_format(); + }} + auto result = at::empty(self.sizes(), options, fmt); + result.zero_(); + return result; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# empty_like: (self, dtype?, layout?, device?, pin?, memory_format?) -> uninit tensor +# with self's meta. Same shape as ones_like but no fill_ (contents undefined). +# FlagGems' pointwise_dynamic allocates its outputs via torch.empty_like, so this +# must exist on the ascend backend for any op routed to flagos_python. +T_EMPTY_LIKE = """\ +at::Tensor {kernel}(const at::Tensor& self, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory, ::std::optional memory_format) {{ + auto options = at::TensorOptions() + .dtype(dtype.value_or(self.scalar_type())) + .layout(layout.value_or(self.layout())) + .device(device.value_or(self.device())) + .pinned_memory(pin_memory.value_or(false)); + auto fmt = memory_format.value_or(at::MemoryFormat::Preserve); + if (fmt == at::MemoryFormat::Preserve) {{ + fmt = self.suggest_memory_format(); + }} + return at::empty(self.sizes(), options, fmt); +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# full: (IntArrayRef size, Scalar fill, dtype?, layout?, device?, pin?) -> filled tensor. +# Default dtype: if a fill value is integral and no dtype given, torch uses long; +# but transformers' generate always passes an explicit dtype, and value_or(kFloat) +# matches zeros/ones behaviour, so keep it simple and consistent with T_ZEROS. +T_FULL = """\ +at::Tensor {kernel}(at::IntArrayRef size, const at::Scalar& fill, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory) {{ + auto options = at::TensorOptions() + .dtype(dtype.value_or(at::kFloat)) + .layout(layout.value_or(at::kStrided)) + .device(device.value_or(at::Device(at::kPrivateUse1, 0))) + .pinned_memory(pin_memory.value_or(false)); + auto result = at::empty(size, options); + result.fill_(fill); + return result; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# full_like: (self, Scalar fill, dtype?, layout?, device?, pin?, memory_format?) -> self-shaped, filled. +T_FULL_LIKE = """\ +at::Tensor {kernel}(const at::Tensor& self, const at::Scalar& fill, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory, ::std::optional memory_format) {{ + auto options = at::TensorOptions() + .dtype(dtype.value_or(self.scalar_type())) + .layout(layout.value_or(self.layout())) + .device(device.value_or(self.device())) + .pinned_memory(pin_memory.value_or(false)); + auto fmt = memory_format.value_or(at::MemoryFormat::Preserve); + if (fmt == at::MemoryFormat::Preserve) {{ + fmt = self.suggest_memory_format(); + }} + auto result = at::empty(self.sizes(), options, fmt); + result.fill_(fill); + return result; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# ones: (IntArrayRef size, dtype?, layout?, device?, pin?) -> tensor of ones. +# Same shape as T_ZEROS; fill_(1) instead of zero_(). transformers' generate() +# uses torch.ones(batch_size, device=...) for unfinished_sequences bookkeeping. +T_ONES = """\ +at::Tensor {kernel}(at::IntArrayRef size, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory) {{ + auto options = at::TensorOptions() + .dtype(dtype.value_or(at::kFloat)) + .layout(layout.value_or(at::kStrided)) + .device(device.value_or(at::Device(at::kPrivateUse1, 0))) + .pinned_memory(pin_memory.value_or(false)); + auto result = at::empty(size, options); + result.fill_(1); + return result; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + # new_ones: (self, IntArrayRef size, dtype?, layout?, device?, pin?) -> ones w/ self's meta. T_NEW_ONES = """\ at::Tensor {kernel}(const at::Tensor& self, at::IntArrayRef size, ::std::optional dtype, ::std::optional layout, ::std::optional device, ::std::optional pin_memory) {{ @@ -1513,6 +2108,56 @@ REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) """ +# clamp: (self, Scalar? min, Scalar? max) -> Tensor, self's shape/dtype. +# aclnnClamp(self, clipValueMin, clipValueMax, out). Either bound may be +# absent (torch allows min=None or max=None, just not both); AclScalarWrapper's +# default ctor leaves the acl_scalar null, which aclnn reads as "not supplied". +T_CLAMP = """\ +at::Tensor {kernel}(const at::Tensor& self, const ::std::optional& min, const ::std::optional& max) {{ + namespace ascend = at::native::flagos::ascend; + auto out = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options()); + + ascend::AclTensorWrapper acl_self(self); + ascend::AclScalarWrapper acl_min = min.has_value() + ? ascend::AclScalarWrapper(min.value(), self.scalar_type()) + : ascend::AclScalarWrapper(); + ascend::AclScalarWrapper acl_max = max.has_value() + ? ascend::AclScalarWrapper(max.value(), self.scalar_type()) + : ascend::AclScalarWrapper(); + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD({aclnn}, acl_self.get(), acl_min.get(), acl_max.get(), acl_out.get()); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# clamp.Tensor: (self, Tensor? min, Tensor? max) -> Tensor, broadcast shape. +# aclnnClampTensor(self, minT, maxT, out). AclTensorWrapper already maps an +# undefined at::Tensor to a null aclTensor*, matching an absent bound. +T_CLAMP_TENSOR = """\ +at::Tensor {kernel}(const at::Tensor& self, const ::std::optional& min, const ::std::optional& max) {{ + namespace ascend = at::native::flagos::ascend; + auto out_shape = self.sizes().vec(); + if (min.has_value()) out_shape = at::infer_size(out_shape, min.value().sizes()); + if (max.has_value()) out_shape = at::infer_size(out_shape, max.value().sizes()); + auto out = ascend::OpPreparation::apply_tensor_without_format( + out_shape, self.options()); + + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_min(min.value_or(at::Tensor())); + ascend::AclTensorWrapper acl_max(max.value_or(at::Tensor())); + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD({aclnn}, acl_self.get(), acl_min.get(), acl_max.get(), acl_out.get()); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + # softmax_fwd: _softmax(self, int64 dim, bool half_to_float) -> same shape. # aclnn(self, dim, out). half_to_float promotes the output dtype to float. T_SOFTMAX_FWD = """\ @@ -1532,6 +2177,37 @@ REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) """ +# softmax_fwd CACHED: same as T_SOFTMAX_FWD but through the repeatable-executor +# cache. `dim` is baked into the executor at GetWorkspaceSize, so it MUST be in +# the key. half_to_float only changes the output dtype, which is already part of +# the out-tensor signature, but fold it in too for safety. Uncached softmax was +# measured at a FLAT ~38us/call regardless of shape (pure GetWorkspaceSize + +# aclCreateTensor build cost) vs ~14us on torch_npu; the decode attention shape +# is fixed so caching drops it to the aclnn-execute floor. +T_SOFTMAX_FWD_CACHED = """\ +at::Tensor {kernel}(const at::Tensor& self, int64_t dim, bool half_to_float) {{ + namespace ascend = at::native::flagos::ascend; + auto out_dtype = half_to_float ? at::kFloat : self.scalar_type(); + auto out = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options().dtype(out_dtype)); + + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); hsh.val(dim); + {{ int8_t h2f = half_to_float ? 1 : 0; hsh.val(h2f); }} + ascend::ExecAscendCached( + "{aclnn}", "{aclnn}GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {{&self}}, {{&out}}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) {{ + return gws(in[0].acl_tensor, dim, out_t[0].acl_tensor, pws, pex); + }}); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + # reduce_all: all(self) -> bool scalar over ALL elements. aclnnAll reduces along # a dim list, so flatten to 1-D and reduce dim=0 to a 0-d bool out. # aclnn(self_flat, dim_list, keepdim=false, out) @@ -1601,6 +2277,108 @@ """ ) +# reduce_sum_dtype CACHED. dims/keepdim/out_dtype are baked into the executor at +# build (aclnnReduceSum reads them during GetWorkspaceSize), so all three go in +# the cache key alongside the input tensor signature. +T_REDUCE_SUM_DTYPE_CACHED = ( + """\ +at::Tensor {kernel}(const at::Tensor& self, at::OptionalIntArrayRef dim, bool keepdim, std::optional dtype) {{ +""" + + _REDUCE_DTYPE_PROLOGUE + + """\ + aclDataType acl_dtype = ascend::ToAclDataType(out_dtype); + + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + for (int64_t d : norm_dims) hsh.val(d); + hsh.val(keepdim); + {{ int32_t dtk = static_cast(acl_dtype); hsh.val(dtk); }} + ascend::ExecAscendCached( + "{aclnn}", "{aclnn}GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {{&self}}, {{&out}}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) {{ + return gws(in[0].acl_tensor, acl_dim.get(), keepdim, acl_dtype, out_t[0].acl_tensor, pws, pex); + }}); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" +) + +# reduce_sum_all: sum(self, ScalarType? dtype) -> full reduction to a 0-d tensor. +# Reuses aclnnReduceSum over every axis with keepdim=false. transformers' +# fast_all() calls tensor.sum() on the causal-mask bool tensor. +T_REDUCE_SUM_ALL = """\ +at::Tensor {kernel}(const at::Tensor& self, std::optional dtype) {{ + namespace ascend = at::native::flagos::ascend; + // Integral/bool inputs promote to int64 when no dtype given (matches torch). + at::ScalarType out_dtype = dtype.has_value() + ? dtype.value() + : (c10::isIntegralType(self.scalar_type(), /*includeBool=*/true) + ? at::kLong : self.scalar_type()); + int64_t ndim = self.dim(); + std::vector norm_dims; + for (int64_t d = 0; d < ndim; ++d) norm_dims.push_back(d); + auto out = ascend::OpPreparation::apply_tensor_without_format( + {{}}, self.options().dtype(out_dtype)); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_out(out); + ascend::AclIntArrayWrapper acl_dim(norm_dims); + aclDataType acl_dtype = ascend::ToAclDataType(out_dtype); + + EXEC_ASCEND_CMD({aclnn}, acl_self.get(), acl_dim.get(), false, acl_dtype, acl_out.get()); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# mean_all: mean(self, ScalarType? dtype) -> full reduction to a 0-d tensor. +# aclnnMean(self, dims=all, keepdim=false, aclDataType, out). Unlike sum, +# mean does NOT integer-promote (undefined for int/bool without a given +# dtype); out_dtype defaults to self's own float dtype. +T_MEAN_ALL = """\ +at::Tensor {kernel}(const at::Tensor& self, std::optional dtype) {{ + namespace ascend = at::native::flagos::ascend; + at::ScalarType out_dtype = dtype.has_value() ? dtype.value() : self.scalar_type(); + int64_t ndim = self.dim(); + std::vector norm_dims; + for (int64_t d = 0; d < ndim; ++d) norm_dims.push_back(d); + auto out = ascend::OpPreparation::apply_tensor_without_format( + {{}}, self.options().dtype(out_dtype)); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_out(out); + ascend::AclIntArrayWrapper acl_dim(norm_dims); + aclDataType acl_dtype = ascend::ToAclDataType(out_dtype); + + EXEC_ASCEND_CMD({aclnn}, acl_self.get(), acl_dim.get(), false, acl_dtype, acl_out.get()); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# reduce_minmax_all: max(self) / min(self) -> 0-d tensor over ALL elements. +# aclnn(self, out); out keeps self's dtype. transformers' generate() +# loop calls unfinished_sequences.max() to test the stop condition. +T_REDUCE_MINMAX_ALL = """\ +at::Tensor {kernel}(const at::Tensor& self) {{ + namespace ascend = at::native::flagos::ascend; + auto out = ascend::OpPreparation::apply_tensor_without_format( + {{}}, self.options()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_out(out); + + EXEC_ASCEND_CMD({aclnn}, acl_self.get(), acl_out.get()); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + # reduce_mean_dtype: mean.dim(self, int[]? dim, keepdim, ScalarType? dtype). # aclnnMeanV2(self, dims, keepdim, int32 dtype, out) -- MeanV2 for CANN 8.5. T_REDUCE_MEAN_DTYPE = ( @@ -1619,6 +2397,37 @@ """ ) +# reduce_mean_dtype CACHED. Like the sum variant: dims/keepdim/dtype baked into +# the executor, so keyed on all three. mean.dim is 113/step in RMSNorm variance +# (measured 57us uncached) -- one of the largest remaining host lines. +T_REDUCE_MEAN_DTYPE_CACHED = ( + """\ +at::Tensor {kernel}(const at::Tensor& self, at::OptionalIntArrayRef dim, bool keepdim, std::optional dtype) {{ +""" + + _REDUCE_DTYPE_PROLOGUE + + """\ + auto acl_dtype = static_cast(ascend::ToAclDataType(out_dtype)); + + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + for (int64_t d : norm_dims) hsh.val(d); + hsh.val(keepdim); + hsh.val(acl_dtype); + ascend::ExecAscendCached( + "{aclnn}", "{aclnn}GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {{&self}}, {{&out}}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) {{ + return gws(in[0].acl_tensor, acl_dim.get(), keepdim, acl_dtype, out_t[0].acl_tensor, pws, pex); + }}); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" +) + # ========================================================================== # conv / pool family. These need an explicit output-shape formula (aclnn wants # the output pre-allocated), so each carries a small shape helper in its body. @@ -1879,6 +2688,129 @@ REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) """ +# add_.Tensor: (self&, other, alpha) -> self&, in-place self += alpha*other. +# aclnnInplaceAdd(selfRef, other, alpha). other may broadcast against self and +# is coerced to self's device/dtype (mirrors the out-of-place add prologue). +T_INPLACE_ADD_TENSOR = """\ +at::Tensor& {kernel}(at::Tensor& self, const at::Tensor& other, const at::Scalar& alpha) {{ + namespace ascend = at::native::flagos::ascend; + auto other_c = other.is_privateuseone() + ? (other.scalar_type() == self.scalar_type() ? other : other.to(self.scalar_type())) + : other.to(self.options()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_other(other_c); + ascend::AclScalarWrapper acl_alpha(alpha, self.scalar_type()); + EXEC_ASCEND_CMD({aclnn}, const_cast(acl_self.get()), acl_other.get(), + acl_alpha.get()); + return self; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# mul_.Tensor: (self&, other) -> self&, in-place self *= other. +# aclnnInplaceMul(selfRef, other). other coerced to self device/dtype. +T_INPLACE_MUL_TENSOR = """\ +at::Tensor& {kernel}(at::Tensor& self, const at::Tensor& other) {{ + namespace ascend = at::native::flagos::ascend; + auto other_c = other.is_privateuseone() + ? (other.scalar_type() == self.scalar_type() ? other : other.to(self.scalar_type())) + : other.to(self.options()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_other(other_c); + EXEC_ASCEND_CMD({aclnn}, const_cast(acl_self.get()), acl_other.get()); + return self; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# div_.Tensor: (self&, other) -> self&, in-place self /= other. +# aclnnInplaceDiv(selfRef, other). Same shape as mul_.Tensor. +T_INPLACE_DIV_TENSOR = T_INPLACE_MUL_TENSOR + +# mul_.Scalar: (self&, other) -> self&, in-place self *= scalar. +# aclnnInplaceMuls(selfRef, aclScalar). +T_INPLACE_MUL_SCALAR = """\ +at::Tensor& {kernel}(at::Tensor& self, const at::Scalar& other) {{ + namespace ascend = at::native::flagos::ascend; + ascend::AclTensorWrapper acl_self(self); + ascend::AclScalarWrapper acl_other(other, self.scalar_type()); + EXEC_ASCEND_CMD({aclnn}, const_cast(acl_self.get()), acl_other.get()); + return self; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# add_.Scalar: (self&, other, alpha) -> self&, in-place self += alpha*scalar. +# aclnnInplaceAdds(selfRef, otherScalar, alphaScalar). +T_INPLACE_ADD_SCALAR = """\ +at::Tensor& {kernel}(at::Tensor& self, const at::Scalar& other, const at::Scalar& alpha) {{ + namespace ascend = at::native::flagos::ascend; + ascend::AclTensorWrapper acl_self(self); + ascend::AclScalarWrapper acl_other(other, self.scalar_type()); + ascend::AclScalarWrapper acl_alpha(alpha, self.scalar_type()); + EXEC_ASCEND_CMD({aclnn}, const_cast(acl_self.get()), acl_other.get(), + acl_alpha.get()); + return self; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# addcmul_ / addcdiv_: (self&, tensor1, tensor2, value) -> self&, in-place +# self += value * (tensor1 {{*,/}} tensor2). +# aclnn(selfRef, tensor1, tensor2, value). tensor1/tensor2 coerced to +# self's dtype when they live on device. +T_INPLACE_ADDCMUL = """\ +at::Tensor& {kernel}(at::Tensor& self, const at::Tensor& tensor1, const at::Tensor& tensor2, const at::Scalar& value) {{ + namespace ascend = at::native::flagos::ascend; + auto t1 = tensor1.scalar_type() == self.scalar_type() ? tensor1 : tensor1.to(self.scalar_type()); + auto t2 = tensor2.scalar_type() == self.scalar_type() ? tensor2 : tensor2.to(self.scalar_type()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_t1(t1); + ascend::AclTensorWrapper acl_t2(t2); + ascend::AclScalarWrapper acl_value(value, self.scalar_type()); + EXEC_ASCEND_CMD({aclnn}, const_cast(acl_self.get()), acl_t1.get(), + acl_t2.get(), acl_value.get()); + return self; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +T_INPLACE_ADDCDIV = T_INPLACE_ADDCMUL + +# lerp_.Scalar: (self&, end, weight) -> self&, in-place self += weight*(end-self). +# aclnnInplaceLerps(selfRef, end, weightScalar). end coerced to self dtype. +T_INPLACE_LERP_SCALAR = """\ +at::Tensor& {kernel}(at::Tensor& self, const at::Tensor& end, const at::Scalar& weight) {{ + namespace ascend = at::native::flagos::ascend; + auto end_c = end.scalar_type() == self.scalar_type() ? end : end.to(self.scalar_type()); + ascend::AclTensorWrapper acl_self(self); + ascend::AclTensorWrapper acl_end(end_c); + ascend::AclScalarWrapper acl_weight(weight, self.scalar_type()); + EXEC_ASCEND_CMD({aclnn}, const_cast(acl_self.get()), acl_end.get(), + acl_weight.get()); + return self; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +# sqrt_: (self&) -> self&, in-place self = sqrt(self). aclnnInplaceSqrt(selfRef). +T_INPLACE_SQRT = """\ +at::Tensor& {kernel}(at::Tensor& self) {{ + namespace ascend = at::native::flagos::ascend; + ascend::AclTensorWrapper acl_self(self); + EXEC_ASCEND_CMD({aclnn}, const_cast(acl_self.get())); + return self; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + # embedding: (weight, indices, padding_idx, scale_grad_by_freq, sparse) -> Tensor. # aclnnEmbedding(weight, indices, out) uses only weight+indices; the trailing # three args are ignored by aclnn. Output = indices.sizes() + [weight.size(1)]. @@ -2260,8 +3192,21 @@ input.sizes(), input.options()); auto grad_weight = ascend::OpPreparation::apply_tensor_without_format( weight.sizes(), weight.options()); - std::vector bias_shape = bias_sizes.has_value() - ? bias_sizes.value().vec() : std::vector{{weight.size(0)}}; + // grad_bias is always allocated and passed, even when output_mask[2] is + // false (aclnn writes nothing to it then). But its shape must still be + // valid: aclnnConvolutionBackward rejects an empty biasSizes, or one whose + // product is 0, with 161002 (ACLNN_ERR_PARAM_INVALID). For bias=None + // autograd hands us [0] (and an empty list is possible too), so in either + // case substitute the real bias length [Cout] = weight.size(0). + std::vector bias_shape = std::vector{{weight.size(0)}}; + if (bias_sizes.has_value() && !bias_sizes.value().empty()) {{ + const auto bs = bias_sizes.value(); + int64_t numel = 1; + for (auto d : bs) {{ numel *= d; }} + if (numel > 0) {{ + bias_shape = bs.vec(); + }} + }} auto grad_bias = ascend::OpPreparation::apply_tensor_without_format( bias_shape, weight.options()); @@ -2289,6 +3234,150 @@ REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) """ +# ========================================================================== +# Cached (repeatable-executor) variants of the hot pure-tensor categories. +# +# These mirror the plain templates but route through ascend::ExecAscendCached, +# which caches the aclOpExecutor keyed by (op, tensor signature, scalar value). +# On a cache hit (constant shapes -- the eager decode steady state) it skips +# aclnnGetWorkspaceSize + aclCreateTensor and only rebinds the tensor data +# addresses, matching torch_npu's per-op host cost. Only categories whose aclnn +# call is purely (tensors..., [scalars baked into key], out) are cached; scalars +# are folded into the key because they are baked into the executor and are NOT +# rebindable (verified on CANN 9.0.0). The `build` lambda replays the exact +# GetWorkspaceSize arg order on a miss using the cache-owned aclTensors. +# ========================================================================== + +T_UNARY_CACHED = """\ +at::Tensor {kernel}(const at::Tensor& self) {{ + namespace ascend = at::native::flagos::ascend; + auto out = ascend::OpPreparation::apply_tensor_without_format( + self.sizes(), self.options()); + + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self); + ascend::ExecAscendCached( + "{aclnn}", "{aclnn}GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {{&self}}, {{&out}}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) {{ + return gws(in[0].acl_tensor, out_t[0].acl_tensor, pws, pex); + }}); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" + +T_BINARY_CACHED = ( + """\ +at::Tensor {kernel}(const at::Tensor& self, const at::Tensor& other) {{ + namespace ascend = at::native::flagos::ascend; +{scalar_fastpath}""" + + _BINARY_PROLOGUE_BODY + + """\ + auto out = ascend::OpPreparation::apply_tensor_without_format( + out_shape, self.options()); + + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "{aclnn}", "{aclnn}GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {{&self_b, &other_b}}, {{&out}}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) {{ + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }}); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" +) + +T_BINARY_ALPHA_CACHED = ( + """\ +at::Tensor {kernel}(const at::Tensor& self, const at::Tensor& other, const at::Scalar& alpha) {{ + namespace ascend = at::native::flagos::ascend; +{scalar_fastpath}""" + + _BINARY_PROLOGUE_BODY + + """\ + auto out = ascend::OpPreparation::apply_tensor_without_format( + out_shape, self.options()); + ascend::AclScalarWrapper acl_alpha(alpha, result_dtype); + + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + {{ double av = alpha.toDouble(); hsh.val(av); }} + ascend::ExecAscendCached( + "{aclnn}", "{aclnn}GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {{&self_b, &other_b}}, {{&out}}, + [&](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) {{ + return gws(in[0].acl_tensor, in[1].acl_tensor, acl_alpha.get(), out_t[0].acl_tensor, pws, pex); + }}); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" +) + +T_BINARY_CMP_CACHED = ( + """\ +at::Tensor {kernel}(const at::Tensor& self, const at::Tensor& other) {{ +""" + + _BINARY_PROLOGUE + + """\ + auto out = ascend::OpPreparation::apply_tensor_without_format( + out_shape, self.options().dtype(at::kBool)); + + static void* opApiFuncAddr = nullptr; + static void* getWsFuncAddr = nullptr; + ascend::SigHasher hsh; hsh.tensor(self_b); hsh.tensor(other_b); + ascend::ExecAscendCached( + "{aclnn}", "{aclnn}GetWorkspaceSize", opApiFuncAddr, getWsFuncAddr, hsh.h, + {{&self_b, &other_b}}, {{&out}}, + [](ascend::GwsFunc gws, std::vector& in, + std::vector& out_t, uint64_t* pws, aclOpExecutor** pex) {{ + return gws(in[0].acl_tensor, in[1].acl_tensor, out_t[0].acl_tensor, pws, pex); + }}); + return out; +}} + +REGISTER_IMPL_TO_DISPATCHER({fn}, {disp}, Backend::kAscend, {kernel}) +""" +) + +# Map each cacheable category to its cached template. Gated by the env var +# FLAGOS_EXEC_CACHE (default ON); set FLAGOS_EXEC_CACHE=0 to regenerate the +# plain uncached kernels (bisection / correctness fallback). +CACHED_CATEGORIES = { + "unary": T_UNARY_CACHED, + "binary": T_BINARY_CACHED, + "binary_alpha": T_BINARY_ALPHA_CACHED, + "binary_cmp": T_BINARY_CMP_CACHED, + "unary_scalar": T_UNARY_SCALAR_CACHED, + "reduce_sum_dtype": T_REDUCE_SUM_DTYPE_CACHED, + "reduce_mean_dtype": T_REDUCE_MEAN_DTYPE_CACHED, + "softmax_fwd": T_SOFTMAX_FWD_CACHED, +} + +# Maps a Tensor-Tensor binary op to its aclnn scalar variant (s) for the +# CPU-scalar fast path. The value is the aclnn base name (without "aclnn"); the +# variant kind ("noalpha"/"alpha") selects which fast-path template to inject. +# Only ops whose s symbol exists in libopapi.so are listed; the presence +# check at codegen time is a hard gate on top of this map. +SCALAR_VARIANT = { + "mul.Tensor": ("Muls", "noalpha"), + "div.Tensor": ("Divs", "noalpha"), + "add.Tensor": ("Adds", "alpha"), + "sub.Tensor": ("Subs", "alpha"), +} + CATEGORIES = { "unary": T_UNARY, "binary": T_BINARY, @@ -2322,6 +3411,7 @@ "matmul": T_MATMUL, "matmul_out": T_MATMUL_OUT, "cat": T_CAT, + "stack": T_STACK, "mv": T_MV, "dot": T_DOT, "bce": T_BCE, @@ -2336,9 +3426,14 @@ "binary_scalar": T_BINARY_SCALAR, "act_backward_self": T_ACT_BACKWARD_SELF, "where": T_WHERE, + "clamp": T_CLAMP, + "clamp_tensor": T_CLAMP_TENSOR, "softmax_fwd": T_SOFTMAX_FWD, "reduce_all": T_REDUCE_ALL, "reduce_sum_dtype": T_REDUCE_SUM_DTYPE, + "reduce_sum_all": T_REDUCE_SUM_ALL, + "mean_all": T_MEAN_ALL, + "reduce_minmax_all": T_REDUCE_MINMAX_ALL, "reduce_mean_dtype": T_REDUCE_MEAN_DTYPE, "adaptive_avg_pool2d": T_ADAPTIVE_AVG_POOL2D, "avg_pool2d": T_AVG_POOL2D, @@ -2359,19 +3454,61 @@ "inplace_zero": T_INPLACE_ZERO, "inplace_fill_scalar": T_INPLACE_FILL_SCALAR, "inplace_fill_tensor": T_INPLACE_FILL_TENSOR, + "inplace_add_tensor": T_INPLACE_ADD_TENSOR, + "inplace_add_scalar": T_INPLACE_ADD_SCALAR, + "inplace_mul_tensor": T_INPLACE_MUL_TENSOR, + "inplace_mul_scalar": T_INPLACE_MUL_SCALAR, + "inplace_div_tensor": T_INPLACE_DIV_TENSOR, + "inplace_addcmul": T_INPLACE_ADDCMUL, + "inplace_addcdiv": T_INPLACE_ADDCDIV, + "inplace_sqrt": T_INPLACE_SQRT, + "inplace_lerp_scalar": T_INPLACE_LERP_SCALAR, + "foreach_inplace_scalar": T_FOREACH_INPLACE_SCALAR, + "foreach_inplace_lerp_scalar": T_FOREACH_INPLACE_LERP_SCALAR, + "foreach_inplace_addcmul_scalar": T_FOREACH_INPLACE_ADDCMUL_SCALAR, + "foreach_sqrt": T_FOREACH_SQRT, + "foreach_inplace_div_scalarlist": T_FOREACH_INPLACE_DIV_SCALARLIST, + "foreach_inplace_addcdiv_scalarlist": T_FOREACH_INPLACE_ADDCDIV_SCALARLIST, "embedding": T_EMBEDDING, "embedding_dense_backward": T_EMBEDDING_DENSE_BACKWARD, "constant_pad_nd": T_CONSTANT_PAD_ND, "zeros": T_ZEROS, + "ones": T_ONES, "scalar_tensor": T_SCALAR_TENSOR, "ones_like": T_ONES_LIKE, + "zeros_like": T_ZEROS_LIKE, + "empty_like": T_EMPTY_LIKE, + "full": T_FULL, + "full_like": T_FULL_LIKE, "new_ones": T_NEW_ONES, } # Categories whose kernels do NOT issue a direct aclnn call (they build tensors # on-host and fill via zero_/fill_, which are themselves device-side aclnn ops). # The symbol-validation guard is skipped for these; their OPS override is unused. -NO_ACLNN_CATEGORIES = {"zeros", "scalar_tensor", "ones_like", "new_ones"} +NO_ACLNN_CATEGORIES = { + "zeros", + "ones", + "scalar_tensor", + "ones_like", + "zeros_like", + "empty_like", + "full", + "full_like", + "new_ones", +} + +# Categories whose template splits its TensorList args into chunks to stay under +# the CANN per-kernel list-length cap (see FOREACH_CHUNK above). Maps the +# category to the chunk size substituted into the template's {chunk} slot. +FOREACH_CHUNKED_CATEGORIES = { + "foreach_inplace_scalar": FOREACH_CHUNK, + "foreach_inplace_lerp_scalar": FOREACH_CHUNK, + "foreach_inplace_addcmul_scalar": FOREACH_CHUNK, + "foreach_sqrt": FOREACH_CHUNK, + "foreach_inplace_div_scalarlist": FOREACH_CHUNK, + "foreach_inplace_addcdiv_scalarlist": FOREACH_CHUNK, +} FILE_HEADER = """\ // Copyright (c) 2026, BAAI. All rights reserved. @@ -2385,6 +3522,7 @@ #include "../../../generated/ops.h" #include +#include #include #include #include @@ -2442,9 +3580,14 @@ def main(): syms = symbols(libopapi_path()) + # Repeatable-executor cache: on for the cacheable categories unless disabled. + exec_cache = os.environ.get("FLAGOS_EXEC_CACHE", "1") != "0" + bodies = [] covered = [] # (op, aclnn, category) skipped = [] # (op, reason) + cached_ops = [] # ops emitted with the cached template + scalar_fastpath_ops = [] # ops that got the CPU-scalar diversion for op, (cat, override) in OPS.items(): if args.category != "all" and cat != args.category: @@ -2460,9 +3603,33 @@ def main(): continue fn, disp = schema_to_cpp_name(op) kernel = fn[:-2] + "KernelAscend" # SqrtFn -> SqrtKernelAscend - bodies.append( - CATEGORIES[cat].format(kernel=kernel, aclnn=acl, fn=fn, disp=disp) - ) + template = CATEGORIES[cat] + fmt = dict(kernel=kernel, aclnn=acl, fn=fn, disp=disp) + if cat in FOREACH_CHUNKED_CATEGORIES: + fmt["chunk"] = FOREACH_CHUNKED_CATEGORIES[cat] + if exec_cache and cat in CACHED_CATEGORIES: + template = CACHED_CATEGORIES[cat] + cached_ops.append(op) + # binary/binary_alpha cached templates carry a {scalar_fastpath} + # slot. Fill it with the CPU-scalar diversion when the op has an + # aclnn scalar variant present in libopapi.so; otherwise leave empty. + if cat in ("binary", "binary_alpha"): + sf = "" + if exec_cache and op in SCALAR_VARIANT: + sname, kind = SCALAR_VARIANT[op] + acl_s = "aclnn" + sname + if syms is None or ( + acl_s in syms and acl_s + "GetWorkspaceSize" in syms + ): + tmpl = ( + _SCALAR_FASTPATH_ALPHA + if kind == "alpha" + else _SCALAR_FASTPATH_NOALPHA + ) + sf = tmpl.format(aclnn_s=acl_s) + scalar_fastpath_ops.append(op) + fmt["scalar_fastpath"] = sf + bodies.append(template.format(**fmt)) covered.append((op, acl, cat)) OUT_CC.parent.mkdir(parents=True, exist_ok=True) @@ -2470,6 +3637,16 @@ def main(): # Report grouped by category. print(f"[gen] {OUT_CC.relative_to(REPO)} ({len(covered)} kernels)") + if exec_cache: + print( + f" [exec-cache] ON for {len(cached_ops)} op(s): {', '.join(cached_ops)}" + ) + if scalar_fastpath_ops: + print( + f" [scalar-fastpath] {len(scalar_fastpath_ops)} op(s): {', '.join(scalar_fastpath_ops)}" + ) + else: + print(" [exec-cache] OFF (FLAGOS_EXEC_CACHE=0)") by_cat = {} for op, acl, cat in covered: by_cat.setdefault(cat, []).append((op, acl)) diff --git a/scripts/codegen_autograd.py b/scripts/codegen_autograd.py new file mode 100644 index 00000000..76eccbac --- /dev/null +++ b/scripts/codegen_autograd.py @@ -0,0 +1,310 @@ +#!/usr/bin/env python3 +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Codegen for AutogradPrivateUse1 kernels. + +Why this exists +--------------- +`aten::matmul` (and friends) are CompositeImplicitAutograd: PyTorch has no +backend kernel for them, it decomposes them into mm/bmm/view and lets autograd +record the *sub-ops*. A backend that wants to own the fused op (one aclnnMatmul +instead of mm + bmm + view churn) hits a wall: registering a concrete kernel on +plain PrivateUse1 stops the decomposition, so autograd stops seeing sub-ops and +instead binds the op's real derivative -- `aten::matmul_backward` -- which the +backend must then implement. Without it, training decays to CPU and crashes. + +That is why csrc/aten/register.cc used to take the fused path only when +`!requires_grad`: inference got one aclnnMatmul, training fell back to the +decomposition (+816 forward and +1464 backward dispatches/step vs torch_npu on +Qwen3-0.6B). + +The fix is the one torch_npu uses (its codegen/autograd/, which in turn reuses +PyTorch's own torchgen): generate a `VariableType::` kernel that builds the +proper autograd node, and register it on **AutogradPrivateUse1**. That kernel +creates e.g. MatmulBackward0, calls `set_history`, and redispatches below +autograd to the backend's fused kernel. Backward then goes through +`aten::matmul_backward`, which we implement once with aclnn. + +What we generate vs what torch_npu generates +-------------------------------------------- +torch_npu regenerates the whole autograd stack (Functions.h/cpp, +ADInplaceOrViewType, python bindings) because it also adds *custom* ops with +*custom* derivative formulas. We only ever re-own ops that already exist in +PyTorch's derivatives.yaml, so the backward node classes (MatmulBackward0, ...) +are already compiled into libtorch and declared in the shipped +torch/csrc/autograd/generated/Functions.h. We therefore generate only the thin +VariableType layer and link against libtorch's node classes. + +The function bodies come verbatim from torchgen's own `emit_body()` -- the same +code that produces PyTorch's in-tree VariableType_N.cpp -- so the autograd +bookkeeping (saved variables, version counters, fw-grad, view/inplace handling) +is exactly what upstream does, not a hand-rolled imitation. + +Reads: + - AUTOGRAD_OPS below (the ops to re-own) + - PyTorch derivatives.yaml + native_functions.yaml (via torchgen) + +Generates (into csrc/aten/generated/): + - variable_type.cc VariableType:: definitions + TORCH_LIBRARY_IMPL(aten, + AutogradPrivateUse1) registrations + +Usage: + python scripts/codegen_autograd.py +""" + +import os +import re +import sys +from pathlib import Path + +try: + import torchgen + from torchgen.api import cpp + from torchgen.api.autograd import match_differentiability_info + from torchgen.context import native_function_manager + from torchgen.gen import parse_native_yaml + from torchgen.packaged.autograd.gen_inplace_or_view_type import ( + METHOD_DEFINITION, + gen_formals, + use_derived, + ) + from torchgen.packaged.autograd.gen_trace_type import type_wrapper_name + from torchgen.packaged.autograd.gen_variable_type import ( + emit_body, + gen_wrapper_registration, + ) + from torchgen.packaged.autograd.load_derivatives import load_derivatives +except ImportError as e: # pragma: no cover + print( + f"Error: torchgen not found ({e}). Install torch>=2.0 and pyyaml.", + file=sys.stderr, + ) + raise SystemExit(1) + + +# Ops to re-own on AutogradPrivateUse1. +# +# Add an op here only when the backend registers a *fused* kernel for it on +# PrivateUse1 AND the op is CompositeImplicitAutograd. The op's derivative +# (e.g. matmul -> matmul_backward) must then have a PrivateUse1 kernel, or +# backward will fall back to CPU. Ops whose backward is itself composite over +# already-registered ops need no extra work. +AUTOGRAD_OPS = [ + "matmul", +] + +# Preprocessor guard wrapped around the registrations. The ops above are re-owned +# only because the Ascend backend has a fused kernel for them; on other backends +# no such kernel exists, PyTorch's composite decomposition is still what runs, and +# claiming the autograd key would bind a derivative (matmul_backward) that has no +# kernel there. Compiling the registrations out keeps those backends untouched. +REGISTRATION_GUARD = "USE_ASCEND" + + +# torchgen's emitted body carries two things we must strip. +# +# 1) The JVP/forward-AD branch calls `run_jit_decomposition_with_args_for_jvp`, +# which re-enters the *composite* decomposition through the JIT. That defeats +# the whole point (we registered a fused kernel to avoid the decomposition) +# and drags in JIT decomposition machinery. We keep only the else-branch +# (the plain redispatch). Consequence: forward-mode AD is not supported for +# these ops on this backend; reverse-mode (what training uses) is unaffected. +# torch_npu strips the same pattern for the same reason. +_JIT_DECOMP_RE = re.compile( + r"if \(\(.*?\)\) \{.*?static c10::OperatorName full_name\(" + r"\"aten::.*?\", .*?\);\n.*?" + r"return impl::run_jit_decomposition_with_args_for_jvp<.*?>" + r"\(\".*?\", \*opt_op, ks, .*?\);\n\s*\} else \{\n\s*(.*?)\n\s*\}", + re.DOTALL, +) + +# 2) A debug-only assert that the result storage is uniquely owned. Our fused +# kernels may return a tensor that shares storage with a cache entry (the +# executor cache in op_api_common.h owns its tensors), so this NDEBUG-only +# check can trip spuriously. Dropped, as torch_npu does. +_USE_COUNT_RE = re.compile( + r"if \(\S+\.has_storage\(\) && !at::impl::dispatch_mode_enabled\(\) && " + r"!at::impl::tensor_has_dispatch\(\S+\)\) \{\s+TORCH_INTERNAL_ASSERT\(" + r"\S+\.storage\(\)\.use_count\(\) == 1, \"function: \S+\"\);\s+\}", + re.DOTALL, +) + +FILE_HEADER = """\ +// Copyright 2026 FlagOS Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// @generated by scripts/codegen_autograd.py -- DO NOT EDIT. +// +// AutogradPrivateUse1 kernels for CompositeImplicitAutograd ops that a backend +// re-owns with a fused kernel. Each function body is produced by torchgen's own +// emit_body(), the same generator behind PyTorch's in-tree VariableType_N.cpp, +// so the autograd bookkeeping matches upstream exactly. The backward node +// classes (MatmulBackward0, ...) come from libtorch; we only add the thin +// VariableType layer that builds them. +// +// Registering here (AutogradPrivateUse1) rather than PrivateUse1 is what lets +// the fused kernel run while autograd still records a proper graph: this kernel +// creates the grad_fn, then redispatches below autograd to the backend kernel. +// See scripts/codegen_autograd.py for the full rationale. + +#include +#include +#include +#include + +#if defined({guard}) + +using namespace at; +// torchgen's emitted bodies call the autograd helpers (unpack, +// compute_requires_grad, collect_next_edges, SavedVariable, set_history, ...) +// unqualified, because upstream's VariableType_N.cpp is itself compiled inside +// namespace torch::autograd. We generate into our own namespace, so pull them in. +using namespace torch::autograd; +using namespace torch::autograd::generated; + +namespace at::flagos::autograd {{ + +namespace VariableType {{ + +// torchgen's bodies open with `unpack(arg, "arg", i)`. Upstream declares it in +// torch/csrc/autograd/generated/VariableType.h but defines it in +// VariableTypeManual.cpp, which is not part of the installed library -- the +// symbol is not exported, so we cannot link against it. It is only a +// defined-ness check that returns the tensor unchanged (upstream's +// checked_cast_variable), so define it here, as torch_npu likewise does for its +// own generated VariableType. +namespace {{ + +inline at::Tensor& unpack(at::Tensor& t, const char* name, int pos) {{ + TORCH_CHECK(t.defined(), + "Expected a proper Tensor but got None (or an undefined Tensor in C++) " + "for argument #", pos, " '", name, "'"); + return t; +}} + +inline const at::Tensor& unpack(const at::Tensor& t, const char* name, int pos) {{ + TORCH_CHECK(t.defined(), + "Expected a proper Tensor but got None (or an undefined Tensor in C++) " + "for argument #", pos, " '", name, "'"); + return t; +}} + +}} // namespace +""" + +FILE_FOOTER = """\ +}} // namespace VariableType + +namespace {{ + +TORCH_LIBRARY_IMPL(aten, AutogradPrivateUse1, m) {{ +{registrations}}} + +}} // namespace + +}} // namespace at::flagos::autograd +""" + + +def main() -> int: + repo_root = Path(__file__).parent.parent + out_dir = repo_root / "csrc/aten/generated" + out_dir.mkdir(exist_ok=True) + + root = Path(torchgen.__file__).parent + native_yaml = str(root / "packaged/ATen/native/native_functions.yaml") + tags_yaml = str(root / "packaged/ATen/native/tags.yaml") + derivatives_yaml = str(root / "packaged/autograd/derivatives.yaml") + + print("Loading derivatives and native functions via torchgen...") + infos, _ = load_derivatives(derivatives_yaml, native_yaml, tags_yaml) + native_funcs = parse_native_yaml(native_yaml, tags_yaml).native_functions + fns = match_differentiability_info(native_funcs, infos) + + wanted = set(AUTOGRAD_OPS) + definitions: list[str] = [] + registrations: list[str] = [] + seen: set[str] = set() + + for fn in fns: + name = str(fn.func.func.name) + if name not in wanted: + continue + if fn.info is None: + print( + f" WARNING: {name} has no derivative info; skipping", file=sys.stderr + ) + continue + if not use_derived(fn): + print( + f" WARNING: {name} is not a derived-type function; skipping", + file=sys.stderr, + ) + continue + + with native_function_manager(fn.func): + body = emit_body(fn, "Default") + definition = METHOD_DEFINITION.substitute( + return_type=cpp.returns_type(fn.func.func.returns).cpp_type(), + type_wrapper_name=type_wrapper_name(fn.func), + type_definition_body=body, + formals=gen_formals(fn.func), + ) + definition = _JIT_DECOMP_RE.sub(r"\1", definition) + definition = _USE_COUNT_RE.sub("", definition) + definitions.append(definition) + registrations.append(gen_wrapper_registration(fn.func, "Default")) + seen.add(name) + node = ", ".join(sorted({i.op for i in fn.info.values() if i.op})) + print(f" {name}: autograd node {node}") + + missing = wanted - seen + if missing: + print(f"Error: no derivative found for {sorted(missing)}", file=sys.stderr) + return 1 + + out = out_dir / "variable_type.cc" + body = "\n".join(definitions) + regs = "".join(f" {r}\n" for r in registrations) + text = ( + FILE_HEADER.format(guard=REGISTRATION_GUARD) + + "\n" + + body + + "\n" + + FILE_FOOTER.format(registrations=regs) + + f"\n#endif // {REGISTRATION_GUARD}\n" + ) + out.write_text(text) + print( + f"Generated {out.relative_to(repo_root)} " + f"({len(definitions)} kernel(s), {os.path.getsize(out)} bytes)" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/codegen_ops.py b/scripts/codegen_ops.py index 96df2b80..75ac26ad 100644 --- a/scripts/codegen_ops.py +++ b/scripts/codegen_ops.py @@ -136,6 +136,19 @@ def cuda_supported(func, funcs, cuda_index): return False +# Ops that are CompositeImplicitAutograd (so normally decomposed above our +# dispatch key and skipped) but that we WANT to intercept with a fused backend +# kernel. Registering a PrivateUse1 kernel for these overrides the composite +# decomposition (verified: F.rms_norm and aten._fused_rms_norm both land on the +# PrivateUse1 impl). The backend kernel is hand-written (aclnnRmsNorm) since the +# tuple(output, rstd) + normalized_shape semantics no codegen category expresses. +# Without this, HF's Qwen3RMSNorm decomposes into ~6 elementwise ops + 2 dtype +# casts per layer (the eager decode hot path). +FORCE_INCLUDE_OPS = { + "_fused_rms_norm", +} + + def enumerate_all_cuda_ops(nf, funcs, cuda_index): """ Returns (kept_ops, skipped) where kept_ops is the list of op-name strings to @@ -146,12 +159,19 @@ def enumerate_all_cuda_ops(nf, funcs, cuda_index): composite_implicit ops are excluded up front: PyTorch decomposes them ABOVE our dispatch key into leaf ops we already box, so registering them is both unnecessary and risky. structured_delegate ops survive that exclusion. + Ops in FORCE_INCLUDE_OPS bypass both the cuda_supported and composite checks + so a hand-written backend kernel can intercept them. """ kept = [] skipped = defaultdict(list) for func in nf.native_functions: op = str(func.func.name) + if op in FORCE_INCLUDE_OPS: + if op not in MANUAL_REGISTERED_OPS: + kept.append(op) + continue + if not cuda_supported(func, funcs, cuda_index): continue diff --git a/scripts/patch_triton_ascend.py b/scripts/patch_triton_ascend.py index a4c22a71..b5e3f8ac 100644 --- a/scripts/patch_triton_ascend.py +++ b/scripts/patch_triton_ascend.py @@ -191,6 +191,29 @@ def patch_utils(triton_path): return patch_file(fp, replacements) +def patch_npu_utils(triton_path): + """Patch backends/ascend/npu_utils.cpp for CANN 9.0.0 enum names. + + triton-ascend 3.2.0's npu_utils.cpp references rtLimitType_t enumerators + from a newer CANN release. CANN 9.0.0 (rt_external_base.h) names the SIMT + per-warp stack limit RT_LIMIT_TYPE_SIMT_STACK_SIZE, not the newer + RT_LIMIT_TYPE_SIMT_WARP_STACK_SIZE, so the JIT compile of npu_utils.cpp + fails with "could not convert brace-enclosed initializer list". Map the + "WARP_STACK_SIZE" key onto the enumerator that CANN 9.0.0 actually + provides. Idempotent: the newer name only ever appears here. + """ + fp = os.path.join(triton_path, "backends", "ascend", "npu_utils.cpp") + + replacements = [ + ( + "rtLimitType_t::RT_LIMIT_TYPE_SIMT_WARP_STACK_SIZE", + "rtLimitType_t::RT_LIMIT_TYPE_SIMT_STACK_SIZE", + ), + ] + + return patch_file(fp, replacements) + + def main(): parser = argparse.ArgumentParser( description="Patch triton-ascend for torch_fl compatibility" @@ -212,12 +235,15 @@ def main(): ) sys.exit(1) - print("\n[1/2] Patching backends/ascend/driver.py ...") + print("\n[1/3] Patching backends/ascend/driver.py ...") patch_driver(triton_path) - print("\n[2/2] Patching backends/ascend/utils.py ...") + print("\n[2/3] Patching backends/ascend/utils.py ...") patch_utils(triton_path) + print("\n[3/3] Patching backends/ascend/npu_utils.cpp (CANN 9.0.0 enum) ...") + patch_npu_utils(triton_path) + print("\nDone. triton-ascend is now compatible with torch_fl.") print("NOTE: Clear triton kernel cache if you had previously compiled kernels:") print(" rm -rf ~/.triton/cache/") diff --git a/setup.py b/setup.py index 5589b097..9e06fd34 100644 --- a/setup.py +++ b/setup.py @@ -705,13 +705,17 @@ def _vendor_supplies_triton() -> bool: NVIDIA-targeted wheel must not be pulled in as a dependency. - ACCELERATOR=dcu: DTK ships its own Triton (and builds pure-boxing). + - ACCELERATOR=ascend: `triton` is provided by triton-ascend, installed out + of band (it has no PyPI release satisfying `triton>=3.5.1`). Declaring the + dep makes pip install stock triton over triton-ascend, after which any + Triton entry point dies with "0 active drivers". - PPU (PPU_SDK present): the vendor Triton lives on a private index and is versioned 3.x+ (e.g. 3.5.0+v0.2.0.ppu2.1.0), which does not satisfy a `triton>=3.5.1` pin; its sdist is also a download shim that pip cannot always build. Install it manually, then `pip install --no-deps` this package. See "Build from Source (PPU Platform)" in the README. """ - if ACCELERATOR == "dcu": + if ACCELERATOR in ("dcu", "ascend"): return True return bool(os.environ.get("PPU_SDK") or os.environ.get("PPU_HOME")) diff --git a/tests/integration/ops/test_foreach_dispatch.py b/tests/integration/ops/test_foreach_dispatch.py new file mode 100644 index 00000000..9fffc249 --- /dev/null +++ b/tests/integration/ops/test_foreach_dispatch.py @@ -0,0 +1,152 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +_foreach_* TensorList dispatch tests (the AdamW foreach=True path). + +The list lengths here are the point of the test. CANN's aclnnForeach* kernels +only process the first 50 entries of an aclTensorList: past that they either +error or -- worse -- return success and leave the remaining tensors UNTOUCHED. +The Ascend kernels therefore slice their lists into sub-50 chunks, so every case +below runs a length that straddles a chunk boundary and asserts on EVERY entry, +not just the first few. A regression that drops the chunking is silent unless +the tail entries are checked. + +Usage: + pytest tests/integration/ops/test_foreach_dispatch.py -v +""" + +import pytest +import torch +import torch_fl # noqa: F401 + + +DEVICE = "flagos:0" + +# 51 and 60 cross aclnn's raw 50-entry cap; 310 is what AdamW passes for +# Qwen3-0.6B and crosses many chunk boundaries. +LENGTHS = [1, 31, 32, 33, 51, 60, 128, 310] +DTYPES = [torch.float32, torch.float16, torch.bfloat16] +TOL = { + torch.float32: dict(rtol=1e-4, atol=1e-5), + torch.float16: dict(rtol=1e-2, atol=1e-2), + torch.bfloat16: dict(rtol=5e-2, atol=5e-2), +} + + +def _lists(n, dtype, seed, numel=8): + """n CPU tensors of varying shape + their flagos copies.""" + g = torch.Generator().manual_seed(seed) + cpu = [ + ((torch.rand(1 + (i % 3), numel, generator=g) + 0.5).to(dtype)) + for i in range(n) + ] + return cpu, [t.to(DEVICE) for t in cpu] + + +def _assert_all_close(got, ref, dtype): + """Compare every entry -- a truncating kernel only differs in the tail.""" + assert len(got) == len(ref) + for i, (g, r) in enumerate(zip(got, ref)): + torch.testing.assert_close( + g.cpu().float(), + r.float(), + msg=lambda m, i=i: f"entry {i}: {m}", + **TOL[dtype], + ) + + +@pytest.mark.parametrize("n", LENGTHS) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.anyplatform +class TestForeachAgainstCpu: + """Each _foreach_* Ascend kernel vs the CPU implementation, entry by entry.""" + + def test_foreach_mul_scalar(self, n, dtype): + cpu, dev = _lists(n, dtype, 1) + torch._foreach_mul_(cpu, 0.9) + torch._foreach_mul_(dev, 0.9) + _assert_all_close(dev, cpu, dtype) + + def test_foreach_add_scalar(self, n, dtype): + cpu, dev = _lists(n, dtype, 2) + torch._foreach_add_(cpu, 1.0) + torch._foreach_add_(dev, 1.0) + _assert_all_close(dev, cpu, dtype) + + def test_foreach_sqrt(self, n, dtype): + cpu, dev = _lists(n, dtype, 3) + _assert_all_close(torch._foreach_sqrt(dev), torch._foreach_sqrt(cpu), dtype) + + def test_foreach_lerp_scalar(self, n, dtype): + cpu, dev = _lists(n, dtype, 4) + # distinct end points: lerp_ toward a copy of self is a no-op and would + # hide a truncating kernel entirely. + end_cpu, end_dev = _lists(n, dtype, 5) + torch._foreach_lerp_(cpu, end_cpu, 0.3) + torch._foreach_lerp_(dev, end_dev, 0.3) + _assert_all_close(dev, cpu, dtype) + + def test_foreach_addcmul_scalar(self, n, dtype): + cpu, dev = _lists(n, dtype, 6) + t1_cpu, t1_dev = _lists(n, dtype, 7) + t2_cpu, t2_dev = _lists(n, dtype, 8) + torch._foreach_addcmul_(cpu, t1_cpu, t2_cpu, 0.1) + torch._foreach_addcmul_(dev, t1_dev, t2_dev, 0.1) + _assert_all_close(dev, cpu, dtype) + + def test_foreach_div_scalarlist(self, n, dtype): + cpu, dev = _lists(n, dtype, 9) + scalars = [0.5 + 0.01 * i for i in range(n)] + torch._foreach_div_(cpu, scalars) + torch._foreach_div_(dev, scalars) + _assert_all_close(dev, cpu, dtype) + + def test_foreach_addcdiv_scalarlist(self, n, dtype): + cpu, dev = _lists(n, dtype, 10) + t1_cpu, t1_dev = _lists(n, dtype, 11) + t2_cpu, t2_dev = _lists(n, dtype, 12) + scalars = [0.5 + 0.01 * i for i in range(n)] + torch._foreach_addcdiv_(cpu, t1_cpu, t2_cpu, scalars) + torch._foreach_addcdiv_(dev, t1_dev, t2_dev, scalars) + _assert_all_close(dev, cpu, dtype) + + +class TestForeachAdamW: + """AdamW(foreach=True) must match AdamW(foreach=False) on the same model.""" + + @pytest.mark.anyplatform + def test_adamw_foreach_matches_single_tensor(self): + # >50 params so the optimizer's TensorLists cross a chunk boundary. + torch.manual_seed(0) + shapes = [(8, 8)] * 40 + [(16,)] * 40 + + def build(): + torch.manual_seed(0) + return [torch.nn.Parameter(torch.randn(*s, device=DEVICE)) for s in shapes] + + def run(foreach): + params = build() + opt = torch.optim.AdamW(params, lr=1e-2, foreach=foreach) + for _ in range(3): + for p in params: + p.grad = torch.ones_like(p) * 0.1 + opt.step() + opt.zero_grad() + return [p.detach().cpu().clone() for p in params] + + for i, (a, b) in enumerate(zip(run(True), run(False))): + torch.testing.assert_close( + a, b, rtol=1e-4, atol=1e-5, msg=lambda m, i=i: f"param {i}: {m}" + ) diff --git a/tests/integration/ops/test_matmul_backward_dispatch.py b/tests/integration/ops/test_matmul_backward_dispatch.py new file mode 100644 index 00000000..09d3689b --- /dev/null +++ b/tests/integration/ops/test_matmul_backward_dispatch.py @@ -0,0 +1,173 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +matmul forward + backward shape/value tests. + +On Ascend, aten::matmul is claimed as a fused aclnnMatmul kernel instead of +decomposing into mm/bmm/view. That means autograd binds the op's real +derivative, aten::matmul_backward, which the backend implements by hand -- and +that kernel has to reproduce every shape rule aten::matmul applies in the +forward pass: 1-D promotion, batch broadcasting, and the 2-D/N-D fold. + +Those rules are exactly where a hand-written backward goes wrong silently: the +gradient still has a plausible shape, so only a value comparison catches it. +(op-plugin's reference kernel, which this one started from, mixes batches for +2-D x N-D and returns a wrong-shaped gradient for interior broadcasts. Both are +covered below.) + +Everything is compared against a float64 CPU reference and scored relatively, +so fp32/hf32 accumulation over a large K is not mistaken for a shape bug. + +Usage: + pytest tests/integration/ops/test_matmul_backward_dispatch.py -v +""" + +import pytest +import torch +import torch_fl # noqa: F401 + + +DEVICE = "flagos:0" + +# Relative to the float64 reference. The Ascend kernel runs with +# ALLOW_FP32_DOWN_PRECISION (hf32) cube math, which lands around 1e-4. +TOL = 2e-3 + +# (a_shape, b_shape, id) -- one per shape family the kernel branches on. +SHAPE_CASES = [ + ((5,), (5,), "1d_x_1d_dot"), + ((3, 4), (4,), "2d_x_1d"), + ((4,), (4, 6), "1d_x_2d"), + ((3, 4), (4, 6), "2d_x_2d_mm"), + ((2, 3, 4), (2, 4, 6), "3d_x_3d_bmm"), + ((2, 5, 3, 4), (2, 5, 4, 6), "4d_x_4d"), + ((2, 3, 4), (4,), "3d_x_1d"), + ((4,), (2, 4, 6), "1d_x_3d"), + ((2, 3, 4, 5), (5,), "4d_x_1d"), + ((4,), (2, 3, 4, 6), "1d_x_4d"), + # The 2-D/N-D fold: matmul collapses batch dims into the contraction, and + # the backward has to reproduce that pairing rather than a plain reshape. + ((2, 3, 4), (4, 6), "3d_x_2d_fold"), + ((3, 4), (2, 4, 6), "2d_x_3d_fold"), + ((1, 128, 1024), (1024, 512), "qwen_like_fold"), + # Batch broadcasting: the raw gradient is shaped like the *broadcast* + # operand and must be summed back down. Leading, interior and trailing + # singletons all have to work, not just a leading prefix. + ((1, 3, 4), (2, 4, 6), "broadcast_leading_a"), + ((1, 1, 3, 4), (3, 4, 6), "broadcast_leading_singletons"), + ((2, 1, 3, 4), (2, 5, 4, 6), "broadcast_interior_a"), + ((2, 5, 3, 4), (2, 1, 4, 6), "broadcast_interior_b"), + ((2, 3, 4), (1, 4, 6), "broadcast_leading_b"), + ((3, 4), (1, 4, 6), "broadcast_2d_x_singleton_batch"), + ((5, 3, 4), (1, 4, 6), "broadcast_b_batch"), + ((2, 1, 1, 3, 4), (7, 5, 4, 6), "broadcast_5d_multi"), + ((7, 5, 3, 4), (1, 1, 4, 6), "broadcast_trailing_singleton"), +] + + +def _rel_err(got: torch.Tensor, ref: torch.Tensor) -> float: + """Max abs error relative to the reference's magnitude.""" + return (got.double() - ref).abs().max().item() / max(ref.abs().max().item(), 1e-12) + + +def _run(a_shape, b_shape, requires_grad=(True, True)): + """matmul fwd+bwd on device and on a float64 CPU reference.""" + torch.manual_seed(0) + a_ref = torch.randn(a_shape, dtype=torch.float64, requires_grad=requires_grad[0]) + b_ref = torch.randn(b_shape, dtype=torch.float64, requires_grad=requires_grad[1]) + a_dev = a_ref.detach().float().to(DEVICE).requires_grad_(requires_grad[0]) + b_dev = b_ref.detach().float().to(DEVICE).requires_grad_(requires_grad[1]) + + out_ref = torch.matmul(a_ref, b_ref) + out_dev = torch.matmul(a_dev, b_dev) + + torch.manual_seed(1) + grad = torch.randn(out_ref.shape, dtype=torch.float64) + out_ref.backward(grad) + out_dev.backward(grad.float().to(DEVICE)) + + return (a_ref, b_ref, out_ref), (a_dev, b_dev, out_dev) + + +@pytest.mark.anyplatform +@pytest.mark.parametrize( + "a_shape,b_shape", + [(a, b) for a, b, _ in SHAPE_CASES], + ids=[i for _, _, i in SHAPE_CASES], +) +def test_matmul_forward_backward_matches_cpu(a_shape, b_shape): + ref, dev = _run(a_shape, b_shape) + (a_ref, b_ref, out_ref), (a_dev, b_dev, out_dev) = ref, dev + + assert out_dev.shape == out_ref.shape + # Shape first: a wrong-shaped gradient is a different (and more severe) bug + # than a wrong-valued one, and asserting it separately says which occurred. + assert a_dev.grad.shape == a_ref.grad.shape, "grad_self shape mismatch" + assert b_dev.grad.shape == b_ref.grad.shape, "grad_other shape mismatch" + + assert _rel_err(out_dev.cpu(), out_ref.detach()) < TOL, "forward value mismatch" + assert _rel_err(a_dev.grad.cpu(), a_ref.grad) < TOL, "grad_self value mismatch" + assert _rel_err(b_dev.grad.cpu(), b_ref.grad) < TOL, "grad_other value mismatch" + + +@pytest.mark.anyplatform +@pytest.mark.parametrize( + "a_shape,b_shape", + [((2, 3, 4), (4, 6)), ((3, 4), (2, 4, 6)), ((1, 3, 4), (5, 4, 6))], +) +@pytest.mark.parametrize("side", ["self", "other"]) +def test_matmul_backward_grad_input_mask(a_shape, b_shape, side): + """Only the requested side gets a gradient (exercises grad_input_mask).""" + mask = (side == "self", side == "other") + (a_ref, b_ref, _), (a_dev, b_dev, _) = _run(a_shape, b_shape, requires_grad=mask) + + if side == "self": + assert b_dev.grad is None, "other must not get a gradient" + assert _rel_err(a_dev.grad.cpu(), a_ref.grad) < TOL + else: + assert a_dev.grad is None, "self must not get a gradient" + assert _rel_err(b_dev.grad.cpu(), b_ref.grad) < TOL + + +@pytest.mark.anyplatform +def test_matmul_records_autograd_graph(): + """The fused kernel must still build a real autograd node, not detach. + + Claiming a CompositeImplicitAutograd op on PrivateUse1 silently drops the + graph unless an AutogradPrivateUse1 kernel re-creates it, which is what + csrc/aten/generated/variable_type.cc exists to do. + """ + a = torch.randn(2, 3, 4, device=DEVICE, requires_grad=True) + b = torch.randn(4, 6, device=DEVICE, requires_grad=True) + out = torch.matmul(a, b) + assert out.grad_fn is not None, "matmul produced no grad_fn" + out.sum().backward() + assert a.grad is not None and b.grad is not None + + +@pytest.mark.anyplatform +def test_matmul_backward_through_chain(): + """Gradients flow through a matmul that is not the last op in the graph.""" + torch.manual_seed(0) + w_ref = torch.randn(4, 6, dtype=torch.float64, requires_grad=True) + x_ref = torch.randn(2, 3, 4, dtype=torch.float64, requires_grad=True) + w_dev = w_ref.detach().float().to(DEVICE).requires_grad_(True) + x_dev = x_ref.detach().float().to(DEVICE).requires_grad_(True) + + (torch.matmul(x_ref, w_ref) * 2.0).sum().backward() + (torch.matmul(x_dev, w_dev) * 2.0).sum().backward() + + assert _rel_err(x_dev.grad.cpu(), x_ref.grad) < TOL + assert _rel_err(w_dev.grad.cpu(), w_ref.grad) < TOL diff --git a/tests/integration/test_qwen3_train.py b/tests/integration/test_qwen3_train.py index e49119e1..5ad1a8f0 100644 --- a/tests/integration/test_qwen3_train.py +++ b/tests/integration/test_qwen3_train.py @@ -86,8 +86,11 @@ def ctx(request): print(f" Parameters: {total:.2f}M total, {trainable:.2f}M trainable") print(f" Load time: {time.time() - t0:.2f}s") + # foreach=False: the ascend backend has no fused _foreach_* TensorList + # kernels, so route AdamW through the single-tensor path (add_/addcmul_/ + # addcdiv_/sqrt), which the backend does implement. optimizer = torch.optim.AdamW( - [p for p in model.parameters() if p.requires_grad], lr=lr + [p for p in model.parameters() if p.requires_grad], lr=lr, foreach=False ) dataset = DummyTextDataset(tokenizer, num_samples=100, max_length=seq_len) dataloader = DataLoader( diff --git a/tests/perf/e2e_qwen3_infer_ascend.py b/tests/perf/e2e_qwen3_infer_ascend.py new file mode 100644 index 00000000..8101cbff --- /dev/null +++ b/tests/perf/e2e_qwen3_infer_ascend.py @@ -0,0 +1,180 @@ +""" +End-to-end Qwen3 inference benchmark on Ascend 910, comparing backends. + +Two backends share one identical measurement harness (same model, same prompt, +same fixed token count, same warmup/round counts) so the numbers are directly +comparable: + + --backend torch_fl torch_fl + aclnn C++ kernels (device flagos:0) + --backend torch_npu Huawei torch_npu baseline (device npu:0) + +Usage: + # aclnn path (env ascend_p0_210) + ACCELERATOR=ascend python tests/perf/e2e_qwen3_infer_ascend.py \ + --backend torch_fl --model /tmp/Qwen3-0.6B --tokens 64 + + # torch_npu baseline (env torch_npu_210) + python tests/perf/e2e_qwen3_infer_ascend.py \ + --backend torch_npu --model /tmp/Qwen3-0.6B --tokens 64 +""" + +import argparse +import time + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + + +def setup_backend(backend): + """Import the backend module, return (device_str, synchronize_fn).""" + if backend == "torch_fl": + import torch_fl + + torch_fl.flagos.set_device(0) + return "flagos:0", torch_fl.flagos.synchronize + elif backend == "torch_npu": + import torch_npu # noqa: F401 + + torch.npu.set_device(0) + return "npu:0", torch.npu.synchronize + raise ValueError(f"unknown backend {backend}") + + +def main(): + args = parse_args() + device, synchronize = setup_backend(args.backend) + + print(f"Backend: {args.backend}") + print(f"Device: {device}") + print(f"PyTorch version: {torch.__version__}") + print() + + # Load model + print("Loading model...") + t0 = time.time() + tokenizer = AutoTokenizer.from_pretrained(args.model) + model = AutoModelForCausalLM.from_pretrained( + args.model, torch_dtype=torch.float16, device_map="cpu" + ) + model = model.to(device) + model.eval() + # Force eager attention so both backends run the same math path. + model.config._attn_implementation = "eager" + + # Optional: route Qwen3RMSNorm through F.rms_norm so it lands on a single + # fused kernel (aclnnRmsNorm on torch_fl/Ascend) instead of HF's ~6 + # elementwise ops + 2 dtype casts. Applied to BOTH backends so the + # comparison stays fair (torch_npu also gets its fused rms_norm path). + if args.fuse_rmsnorm: + from transformers.models.qwen3 import modeling_qwen3 as _m + + def _fused_forward(self, hidden_states): + return torch.nn.functional.rms_norm( + hidden_states, + (hidden_states.shape[-1],), + self.weight, + self.variance_epsilon, + ) + + _m.Qwen3RMSNorm.forward = _fused_forward + print("RMSNorm: fused (F.rms_norm)") + else: + print("RMSNorm: HF default (decomposed)") + print(f"Model loaded in {time.time() - t0:.2f}s") + print("Attention: eager") + print() + + # Prepare input + text = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "Give me a short introduction to large language model.", + } + ], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + inputs = tokenizer([text], return_tensors="pt").to(device) + input_len = inputs["input_ids"].shape[1] + print(f"Input tokens: {input_len}") + print(f"Output tokens: {args.tokens} (fixed, greedy)") + print(f"Warmup rounds: {args.warmup_rounds}, Benchmark rounds: {args.rounds}") + print() + + gen_kwargs = dict( + **inputs, + max_new_tokens=args.tokens, + min_new_tokens=args.tokens, # force exact token count + do_sample=False, # greedy decoding + temperature=None, + top_p=None, + top_k=None, + ) + + # Warmup + print("Warmup...") + for i in range(args.warmup_rounds): + synchronize() + t0 = time.perf_counter() + with torch.no_grad(): + _ = model.generate(**gen_kwargs) + synchronize() + print(f" Round {i + 1}: {time.perf_counter() - t0:.3f}s") + print() + + # Benchmark + print(f"Benchmarking ({args.rounds} rounds)...") + round_times = [] + for i in range(args.rounds): + synchronize() + t0 = time.perf_counter() + with torch.no_grad(): + output = model.generate(**gen_kwargs) + synchronize() + elapsed = time.perf_counter() - t0 + + new_tokens = output.shape[1] - input_len + tps = new_tokens / elapsed + round_times.append(elapsed) + print(f" Round {i + 1}: {elapsed:.3f}s, {new_tokens} tokens, {tps:.2f} tok/s") + + round_times.sort() + median_time = round_times[len(round_times) // 2] + min_time = round_times[0] + max_time = round_times[-1] + median_tps = args.tokens / median_time + + print(f"\n=== E2E Inference Results ({args.backend}) ===") + print(f"Tokens generated: {args.tokens} (greedy, fixed)") + print(f"Median: {median_time:.3f}s ({median_tps:.2f} tok/s)") + print(f"Min: {min_time:.3f}s ({args.tokens / min_time:.2f} tok/s)") + print(f"Max: {max_time:.3f}s ({args.tokens / max_time:.2f} tok/s)") + print(f"Spread: {(max_time - min_time) / median_time * 100:.1f}%") + print(f"Time per token: {median_time / args.tokens * 1000:.2f}ms") + + +def parse_args(): + parser = argparse.ArgumentParser( + description="E2E Qwen3 inference benchmark (Ascend)" + ) + parser.add_argument("--backend", choices=["torch_fl", "torch_npu"], required=True) + parser.add_argument("--model", default="/tmp/Qwen3-0.6B", help="Path to model") + parser.add_argument( + "--tokens", type=int, default=64, help="Exact number of new tokens to generate" + ) + parser.add_argument( + "--rounds", type=int, default=5, help="Benchmark rounds (take median)" + ) + parser.add_argument( + "--fuse-rmsnorm", + action="store_true", + help="Route Qwen3RMSNorm through F.rms_norm (fused kernel on both backends)", + ) + parser.add_argument("--warmup-rounds", type=int, default=3, help="Warmup rounds") + return parser.parse_args() + + +if __name__ == "__main__": + main() diff --git a/tests/perf/e2e_qwen3_train_ascend.py b/tests/perf/e2e_qwen3_train_ascend.py new file mode 100644 index 00000000..9f07fce0 --- /dev/null +++ b/tests/perf/e2e_qwen3_train_ascend.py @@ -0,0 +1,215 @@ +""" +End-to-end Qwen3 training benchmark on Ascend 910, comparing backends. + +One identical harness (same model, seed, batch/seq, optimizer, step count) runs +against two backends so step-time and throughput are directly comparable: + + --backend torch_fl torch_fl + aclnn C++ kernels (device flagos:0) + --backend torch_npu Huawei torch_npu baseline (device npu:0) + +AdamW defaults to its foreach path on both backends (pass --no-foreach for the +single-tensor path); the setting is applied identically either way so the +optimizer contributes the same op mix to both measurements. + +Usage: + ACCELERATOR=ascend python tests/perf/e2e_qwen3_train_ascend.py \ + --backend torch_fl --model /tmp/Qwen3-0.6B --steps 10 + + python tests/perf/e2e_qwen3_train_ascend.py \ + --backend torch_npu --model /tmp/Qwen3-0.6B --steps 10 +""" + +import argparse +import os +import random +import sys +import time + +import numpy as np +import torch +from torch.utils.data import DataLoader +from transformers import AutoModelForCausalLM, AutoTokenizer + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "common")) +from dummy_dataset import DummyTextDataset # noqa: E402 + + +def setup_backend(backend): + if backend == "torch_fl": + import torch_fl + + torch_fl.flagos.set_device(0) + return "flagos:0", torch_fl.flagos.synchronize + elif backend == "torch_npu": + import torch_npu # noqa: F401 + + torch.npu.set_device(0) + return "npu:0", torch.npu.synchronize + raise ValueError(f"unknown backend {backend}") + + +def set_seed(seed: int = 42): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + +def main(): + parser = argparse.ArgumentParser( + description="E2E Qwen3 training benchmark (Ascend)" + ) + parser.add_argument("--backend", choices=["torch_fl", "torch_npu"], required=True) + parser.add_argument("--model", default="/tmp/Qwen3-0.6B", help="Path to model") + parser.add_argument("--steps", type=int, default=10, help="Benchmark steps") + parser.add_argument("--warmup-steps", type=int, default=3, help="Warmup steps") + parser.add_argument("--batch-size", type=int, default=1, help="Batch size") + parser.add_argument("--seq-len", type=int, default=128, help="Sequence length") + parser.add_argument("--lr", type=float, default=1e-5, help="Learning rate") + parser.add_argument( + "--foreach", + action=argparse.BooleanOptionalAction, + default=True, + help="AdamW foreach path (--no-foreach for the single-tensor path)", + ) + args = parser.parse_args() + + device, synchronize = setup_backend(args.backend) + set_seed(42) + + print(f"Backend: {args.backend}") + print(f"Device: {device}") + print(f"PyTorch version: {torch.__version__}") + print(f"Batch size: {args.batch_size}, Seq len: {args.seq_len}") + print() + + print("[1] Loading model...") + t0 = time.time() + tokenizer = AutoTokenizer.from_pretrained(args.model) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + model = AutoModelForCausalLM.from_pretrained( + args.model, + torch_dtype=torch.float32, + device_map="cpu", + attn_implementation="eager", + ) + model = model.to(device) + model.train() + print(f" Load time: {time.time() - t0:.2f}s") + + # Freeze unused parameters (embedding tie etc.) so grads match across backends. + dummy = torch.randint(0, 1000, (1, 32), device=device) + with torch.enable_grad(): + out = model(input_ids=dummy, use_cache=False) + out.logits.sum().backward() + unused = [] + for name, param in model.named_parameters(): + if param.grad is None: + param.requires_grad = False + unused.append(name) + else: + param.grad = None + print(f" Frozen {len(unused)} unused parameters") + + synchronize() + total = sum(p.numel() for p in model.parameters()) / 1e6 + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6 + print(f" Parameters: {total:.2f}M total, {trainable:.2f}M trainable") + print() + + # Both backends get the same foreach setting so the optimizer contributes the + # same op mix to each measurement. + optimizer = torch.optim.AdamW( + [p for p in model.parameters() if p.requires_grad], + lr=args.lr, + foreach=args.foreach, + ) + print(f" Optimizer: AdamW(foreach={args.foreach})") + dataset = DummyTextDataset(tokenizer, num_samples=100, max_length=args.seq_len) + dataloader = DataLoader( + dataset, batch_size=args.batch_size, shuffle=False, drop_last=True + ) + tokens_per_step = args.batch_size * args.seq_len + + def run_step(batch): + input_ids = batch["input_ids"].to(device) + attention_mask = batch["attention_mask"].to(device) + labels = batch["labels"].to(device) + outputs = model( + input_ids=input_ids, + attention_mask=attention_mask, + labels=labels, + use_cache=False, + ) + loss = outputs.loss + loss.backward() + optimizer.step() + optimizer.zero_grad() + return loss.item() + + print(f"[2] Warmup ({args.warmup_steps} steps)...") + data_iter = iter(dataloader) + for i in range(args.warmup_steps): + try: + batch = next(data_iter) + except StopIteration: + data_iter = iter(dataloader) + batch = next(data_iter) + synchronize() + t0 = time.perf_counter() + loss = run_step(batch) + synchronize() + elapsed = time.perf_counter() - t0 + print( + f" Step {i + 1}: loss={loss:.4f}, time={elapsed:.2f}s, " + f"{tokens_per_step / elapsed:.1f} tok/s" + ) + print() + + print(f"[3] Benchmarking ({args.steps} steps)...") + step_times = [] + step_losses = [] + for i in range(args.steps): + try: + batch = next(data_iter) + except StopIteration: + data_iter = iter(dataloader) + batch = next(data_iter) + synchronize() + t0 = time.perf_counter() + loss = run_step(batch) + synchronize() + elapsed = time.perf_counter() - t0 + step_times.append(elapsed) + step_losses.append(loss) + print( + f" Step {i + 1}: loss={loss:.4f}, time={elapsed:.3f}s, " + f"{tokens_per_step / elapsed:.1f} tok/s" + ) + + step_times_sorted = sorted(step_times) + median_time = step_times_sorted[len(step_times_sorted) // 2] + min_time = step_times_sorted[0] + max_time = step_times_sorted[-1] + median_tps = tokens_per_step / median_time + + print(f"\n=== E2E Training Results ({args.backend}) ===") + print(f"Model: {args.model}") + print(f"Batch size: {args.batch_size}, Seq len: {args.seq_len}") + print(f"Tokens per step: {tokens_per_step}") + print(f"Steps: {args.steps}") + print(f"Median step time: {median_time:.3f}s ({median_tps:.1f} tok/s)") + print(f"Min: {min_time:.3f}s ({tokens_per_step / min_time:.1f} tok/s)") + print(f"Max: {max_time:.3f}s ({tokens_per_step / max_time:.1f} tok/s)") + print(f"Spread: {(max_time - min_time) / median_time * 100:.1f}%") + print(f"Time per token: {median_time / tokens_per_step * 1000:.2f}ms") + print() + print("=== Loss Trend ===") + print(f"First loss: {step_losses[0]:.4f}") + print(f"Last loss: {step_losses[-1]:.4f}") + print(f"Avg loss: {sum(step_losses) / len(step_losses):.4f}") + + +if __name__ == "__main__": + main() diff --git a/torch_fl/__init__.py b/torch_fl/__init__.py index 472dc9b4..ef33157a 100644 --- a/torch_fl/__init__.py +++ b/torch_fl/__init__.py @@ -53,6 +53,13 @@ def _select_backend_config() -> None: (liboperators.so, no GIL), the remainder fall back to flagos_python. Only valid when torch_fl was built with FLAGGEMS_KERNEL=ON. + On an Ascend NPU box (detected via /dev/davinci*), the ACL C++ backend is the + only usable one, so the choice is instead: + + * FLAGOS_USE_FLAGGEMS=1 -> backends_ascend_flagos_py.conf (FlagGems Triton + where triton-ascend can run, else ascend aclnn) + * unset / 0 -> backends_ascend.conf (pure aclnn C++) + The MetaX flaggems conf mirrors backends_flaggems.conf but routes the ops triton-metax cannot run (mm/bmm/mean.dim) back to the cuda boxing kernel (maca libtorch_cuda) instead of flagos_python. The DCU one does the same for @@ -94,6 +101,35 @@ def _select_backend_config() -> None: "FALSE", ) metax_boxing = os.environ.get("FLAGOS_METAX_BOXING", "0") == "1" + + conf_dir = os.path.join(os.path.dirname(__file__), "configs") + + # Ascend builds compile the ACL C++ backend (Backend::kAscend), not the CUDA + # boxing kernels, so the cuda/flaggems confs (which route ops to `cuda`) can + # never apply. Since every wheel ships all backends*.conf files, the conf set + # can't distinguish the build; use the runtime hardware signal instead. An + # Ascend NPU exposes /dev/davinci* device nodes -- their presence means this + # is an Ascend box, where the only usable routing is the ascend conf. (A CUDA + # build could not run here anyway, so this never mis-fires on a CUDA host.) + ascend_default = os.path.join(conf_dir, "backends_ascend.conf") + ascend_flaggems = os.path.join(conf_dir, "backends_ascend_flagos_py.conf") + try: + is_ascend_build = os.path.exists(ascend_default) and any( + name.startswith("davinci") for name in os.listdir("/dev") + ) + except OSError: + is_ascend_build = False + + if is_ascend_build: + conf_path = ( + ascend_flaggems + if (use_flaggems and os.path.exists(ascend_flaggems)) + else ascend_default + ) + if os.path.exists(conf_path): + os.environ["FLAGOS_BACKEND_CONFIG"] = conf_path + return + if use_flaggems_cpp: conf_name = "backends_flaggems_cpp.conf" elif use_flaggems and metax_boxing: @@ -351,6 +387,65 @@ def _check_privateuse1_unclaimed() -> None: _registered_ops = [] +def _patch_flaggems_philox(): + """Give FlagGems' RNG ops a working default generator on flagos. + + gems' rand/randn/rand_like/randn_like/randperm/multinomial (and any op that + draws randomness without an explicit generator) call + ``philox_backend_seed_offset(increment)`` with no generator, which then + reaches for ``torch_device_fn.default_generators[current_device()]``. Under + the nvidia branch torch_device_fn is torch.cuda, whose ``default_generators`` + is an EMPTY tuple on a CPU-torch wheel + cuda shim -> IndexError, crashing + every generator-less gems RNG op. + + We hold one module-level CUDA ``torch.Generator`` and monkeypatch + ``philox_backend_seed_offset`` so a None generator with an empty + default_generators falls back to it. gems reads only the philox seed+offset + from the generator and ``set_state``'s the advanced offset back, so the one + shared generator yields distinct streams across calls. The GIL serializes + the set_state (every gems call holds it), so no extra locking is needed. + Seeded from ``torch.initial_seed()`` so ``torch.manual_seed(...)`` before the + first RNG op is honoured. + + No-op / best-effort: wrapped in try/except so a missing flag_gems or a + version without this symbol degrades silently (those ops just stay broken, + same as before). Only meaningful on the nvidia branch (empty cuda + default_generators); ascend/metax have their own generators. + """ + try: + import sys + + import torch + from flag_gems.utils import random_utils + + _fallback = torch.Generator(device="cuda") + _fallback.manual_seed(torch.initial_seed()) + _orig = random_utils.philox_backend_seed_offset + + def _patched(increment, generator=None): + if ( + generator is None + and len(random_utils.torch_device_fn.default_generators) == 0 + ): + generator = _fallback + return _orig(increment, generator=generator) + + # rand.py etc. do `from ..utils.random_utils import + # philox_backend_seed_offset` at import time, binding the name into their + # own module namespace -- patching random_utils alone would not reach + # those local bindings. Rebind the name in every flag_gems module that + # exported it (plus the canonical location). + for mod in list(sys.modules.values()): + name = getattr(mod, "__name__", "") + if name.startswith("flag_gems") and hasattr( + mod, "philox_backend_seed_offset" + ): + mod.philox_backend_seed_offset = _patched + random_utils.philox_backend_seed_offset = _patched + except Exception: + pass + + def _patch_flaggems_codegen_config(): """ Configure FlagGems' vendor + torch.cuda shim for the flagos device. diff --git a/torch_fl/configs/backends_ascend.conf b/torch_fl/configs/backends_ascend.conf index 39afb9e0..b6a5179d 100644 --- a/torch_fl/configs/backends_ascend.conf +++ b/torch_fl/configs/backends_ascend.conf @@ -26,6 +26,7 @@ mm = ascend mm.out = ascend bmm = ascend bmm.out = ascend +matmul = ascend cat = ascend add.Tensor = ascend mul.Tensor = ascend @@ -51,6 +52,13 @@ sum.dim_IntList = ascend slice_backward = ascend nll_loss_forward = ascend nll_loss_backward = ascend +# bespoke handwritten kernels (csrc/aten/backends/ascend/): generation path +lift_fresh = ascend +arange = ascend +arange.start = ascend +arange.start_step = ascend +argmax = ascend +isin.Tensor_Tensor = ascend abs = ascend acos = ascend @@ -66,6 +74,22 @@ squeeze.dim = ascend unsqueeze = ascend _unsafe_view = ascend detach = ascend +t = ascend +unbind.int = ascend +alias = ascend +topk = ascend +sort = ascend +sort.stable = ascend +scatter.src = ascend +multinomial = ascend +randn = ascend +rand = ascend +randint = ascend +randint.low = ascend +masked_select = ascend +# fused RMSNorm via aclnnRmsNorm (intercepts aten::_fused_rms_norm; HF must call +# F.rms_norm to route here). Collapses HF's ~6 ops + 2 casts per layer. +_fused_rms_norm = ascend # --- generated by codegen_ascend.py --- sqrt = ascend @@ -115,6 +139,7 @@ logical_and = ascend logical_or = ascend add.Scalar = ascend sub.Scalar = ascend +rsub.Scalar = ascend eq.Scalar = ascend ne.Scalar = ascend gt.Scalar = ascend @@ -157,6 +182,12 @@ cummax = ascend cummin = ascend aminmax = ascend prod = ascend +stack = ascend +ones = ascend +zeros_like = ascend +empty_like = ascend +full = ascend +full_like = ascend addmm = ascend baddbmm = ascend mv = ascend @@ -177,6 +208,25 @@ index_select = ascend zero_ = ascend fill_.Scalar = ascend fill_.Tensor = ascend +add_.Tensor = ascend +add_.Scalar = ascend +mul_.Tensor = ascend +mul_.Scalar = ascend +div_.Tensor = ascend +bitwise_and_.Tensor = ascend +bitwise_or_.Tensor = ascend +bitwise_xor_.Tensor = ascend +addcmul_ = ascend +addcdiv_ = ascend +sqrt_ = ascend +lerp_.Scalar = ascend +_foreach_mul_.Scalar = ascend +_foreach_add_.Scalar = ascend +_foreach_lerp_.Scalar = ascend +_foreach_addcmul_.Scalar = ascend +_foreach_sqrt = ascend +_foreach_div_.ScalarList = ascend +_foreach_addcdiv_.ScalarList = ascend embedding = ascend embedding_dense_backward = ascend constant_pad_nd = ascend @@ -190,6 +240,13 @@ gelu_backward = ascend _log_softmax = ascend _softmax_backward_data = ascend _log_softmax_backward_data = ascend +any = ascend +sum = ascend +max = ascend +min = ascend +mean = ascend +clamp = ascend +clamp.Tensor = ascend _adaptive_avg_pool2d = ascend avg_pool2d = ascend max_pool2d_with_indices = ascend diff --git a/torch_fl/configs/backends_ascend_flagos_py.conf b/torch_fl/configs/backends_ascend_flagos_py.conf index d8f5b0eb..8e973adf 100644 --- a/torch_fl/configs/backends_ascend_flagos_py.conf +++ b/torch_fl/configs/backends_ascend_flagos_py.conf @@ -34,6 +34,7 @@ le.Tensor = flagos_python mean.dim = ascend # FlagGems non-inner dim path uses CUDA context, fails on ascend mm = ascend # FlagGems uses SPLIT_K kwarg not supported by triton-ascend mm.out = ascend # same as mm +matmul = ascend # route aten::matmul directly to aclnnMatmul, bypasses view churn mul.Tensor = flagos_python mul.Scalar = ascend # excluded from codegen flaggems-python discovery (arity/type gate) neg = flagos_python @@ -56,6 +57,27 @@ embedding = ascend embedding_dense_backward = ascend new_ones = ascend ones_like = ascend +empty_like = ascend # FlagGems pointwise_dynamic allocates outputs via torch.empty_like scalar_tensor = ascend zeros = ascend +full = ascend +full_like = ascend slice_backward = ascend +# generation path (transformers .generate): bespoke ascend kernels, no FlagGems op +lift_fresh = ascend +arange = ascend +arange.start = ascend +arange.start_step = ascend +argmax = ascend +isin.Tensor_Tensor = ascend + +topk = ascend +sort = ascend +sort.stable = ascend +scatter.src = ascend +multinomial = ascend +randn = ascend +rand = ascend +randint = ascend +randint.low = ascend +masked_select = ascend diff --git a/torch_fl/configs/backends_cuda.conf b/torch_fl/configs/backends_cuda.conf index ed4d1f4c..a3e642d5 100644 --- a/torch_fl/configs/backends_cuda.conf +++ b/torch_fl/configs/backends_cuda.conf @@ -361,6 +361,7 @@ _fused_dropout.out = cuda _fused_moving_avg_obs_fq_helper = cuda _fused_moving_avg_obs_fq_helper.out = cuda _fused_moving_avg_obs_fq_helper_functional = cuda +_fused_rms_norm = cuda _fused_rms_norm_backward = cuda _fused_sgd.out = cuda _fused_sgd.tensor_lr_out = cuda diff --git a/torch_fl/configs/backends_flaggems.conf b/torch_fl/configs/backends_flaggems.conf index bd94fd77..a542f98d 100644 --- a/torch_fl/configs/backends_flaggems.conf +++ b/torch_fl/configs/backends_flaggems.conf @@ -365,6 +365,7 @@ _fused_dropout.out = cuda _fused_moving_avg_obs_fq_helper = cuda _fused_moving_avg_obs_fq_helper.out = cuda _fused_moving_avg_obs_fq_helper_functional = cuda +_fused_rms_norm = cuda _fused_rms_norm_backward = cuda _fused_sgd.out = cuda _fused_sgd.tensor_lr_out = cuda diff --git a/torch_fl/configs/backends_metax_flaggems.conf b/torch_fl/configs/backends_metax_flaggems.conf index 0c137893..54da3551 100644 --- a/torch_fl/configs/backends_metax_flaggems.conf +++ b/torch_fl/configs/backends_metax_flaggems.conf @@ -363,6 +363,7 @@ _fused_dropout.out = cuda _fused_moving_avg_obs_fq_helper = cuda _fused_moving_avg_obs_fq_helper.out = cuda _fused_moving_avg_obs_fq_helper_functional = cuda +_fused_rms_norm = cuda _fused_rms_norm_backward = cuda _fused_sgd.out = cuda _fused_sgd.tensor_lr_out = cuda diff --git a/torch_fl/flagos/__init__.py b/torch_fl/flagos/__init__.py index 74810d4b..a8b447f8 100644 --- a/torch_fl/flagos/__init__.py +++ b/torch_fl/flagos/__init__.py @@ -109,34 +109,43 @@ def _lazy_init(): _original_getitem = torch.Tensor.__getitem__ + _Tensor = torch.Tensor + _full_slice = slice(None, None, None) + _aten_index = torch.ops.aten.index.Tensor + def _patched_getitem(self, indices): + # Fast path: the workaround only applies to a tuple of indices that + # contains at least one Tensor. Anything else (the vast majority of + # __getitem__ calls, e.g. x[:, -1:]) returns immediately, avoiding the + # device property access and any tuple scan. + if type(indices) is not tuple: + return _original_getitem(self, indices) + + has_tensor = False + for idx in indices: + if isinstance(idx, _Tensor): + has_tensor = True + break + if not has_tensor: + return _original_getitem(self, indices) + # Only patch for our device if self.device.type not in ("privateuseone", "flagos"): return _original_getitem(self, indices) - # Handle tuple of indices with at least one tensor - if isinstance(indices, tuple): - has_tensor = any(isinstance(idx, torch.Tensor) for idx in indices) - if has_tensor: - # Convert to list for aten.index.Tensor - indices_list = [] - for idx in indices: - if isinstance(idx, slice): - if idx == slice(None, None, None): - indices_list.append(None) - else: - # Non-trivial slice — fall back to original - return _original_getitem(self, indices) - elif isinstance(idx, torch.Tensor): - indices_list.append(idx) - else: - # Other types (int, etc.) — fall back to original - return _original_getitem(self, indices) - - # Use aten.index.Tensor which works correctly - return torch.ops.aten.index.Tensor(self, indices_list) - - return _original_getitem(self, indices) + # Convert to list for aten.index.Tensor + indices_list = [] + for idx in indices: + if isinstance(idx, _Tensor): + indices_list.append(idx) + elif idx is _full_slice or idx == _full_slice: + indices_list.append(None) + else: + # Non-trivial slice / int / other — fall back to original + return _original_getitem(self, indices) + + # Use aten.index.Tensor which works correctly + return _aten_index(self, indices_list) torch.Tensor.__getitem__ = _patched_getitem