feat(ascend): Qwen3-0.6B training at 0.82x torch_npu on real 910

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

csrc/aten/backends/ascend/arange.cc

@@ -0,0 +1,90 @@
+// Copyright (c) 2026, BAAI. All rights reserved.
+
+#include "../../generated/ops.h"
+#include <ATen/core/Tensor.h>
+#include <cmath>
+#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<int64_t>(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<at::ScalarType> dtype, ::std::optional<at::Layout> layout,
+    ::std::optional<at::Device> device, ::std::optional<bool> 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<at::ScalarType> dtype, ::std::optional<at::Layout> layout,
+    ::std::optional<at::Device> device, ::std::optional<bool> 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<at::ScalarType> dtype, ::std::optional<at::Layout> layout,
+    ::std::optional<at::Device> device, ::std::optional<bool> 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

csrc/aten/backends/ascend/argmax.cc

@@ -0,0 +1,55 @@
+// Copyright (c) 2026, BAAI. All rights reserved.
+
+#include "../../generated/ops.h"
+#include <ATen/core/Tensor.h>
+#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<int64_t> 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<int64_t> 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

csrc/aten/backends/ascend/ascend_copy.cc

@@ -0,0 +1,71 @@
+// Copyright (c) 2026, BAAI. All rights reserved.
+
+#include "ascend_copy.h"
+
+#include <ATen/core/Tensor.h>
+#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<aclTensor*>(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<int32_t>(acl_dtype));
+  ExecAscendCached(
+      "aclnnCast", "aclnnCastGetWorkspaceSize",
+      opApiFuncAddr, getWsFuncAddr, hsh.h,
+      {&src_c}, {&out},
+      [&](GwsFunc gws,
+          std::vector<AclTensorWrapper>& in,
+          std::vector<AclTensorWrapper>& 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

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 <ATen/core/Tensor.h>
+
+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