diff --git a/.claude/skills/ws1-single-card-kernel/SKILL.md b/.claude/skills/ws1-single-card-kernel/SKILL.md
new file mode 100644
index 00000000..2807ef9c
--- /dev/null
+++ b/.claude/skills/ws1-single-card-kernel/SKILL.md
@@ -0,0 +1,243 @@
+---
+name: ws1-single-card-kernel
+description: Use when writing a new WS1 single-card kernel operator in this repo (rl-kernel) - PyTorch golden first, then CUDA, then ROCm, then Ascend; gtest registration; PR with exact pytest/gtest commands and results; deterministic backward when the op needs one. The Ascend section is battle-tested; CUDA/ROCm sections are placeholders.
+---
+
+# WS1 Single-Card Kernel Workflow
+
+Follow this workflow when adding a new operator (rmsnorm / embedding / lm_head / logp /
+fused linear logp / rope / silu / swiglu / attention ...). The per-platform order is
+fixed, as are the registration and PR deliverable requirements.
+
+## Global Workflow (all platforms)
+
+1. **Write the PyTorch golden first**: `rl_engine/kernels/ops/pytorch//.py`,
+ the WS1 ground-truth reference — a hand-written fixed-order fp32 reference (e.g.
+ the hand-written softmax for attention, per-row `torch.mv` for lm_head),
+ deliberately NOT `F.scaled_dot_product_attention` / `torch.matmul` shortcuts whose
+ reduction order is unspecified. Expose both `forward` (dtype path) and
+ `forward_fp32` (golden path).
+2. **CUDA platform** (next section — placeholder for now).
+3. **ROCm platform** (next section — placeholder for now).
+4. **Ascend platform** (see the Ascend section) — **only after CUDA is done**: the
+ Ascend kernel mirrors the CUDA deterministic kernel's reduction contract (e.g.
+ contract v1 in `csrc/cuda/fused_linear_logp_sm90.cu`).
+5. **Register in gtest**: add a platform entry (e.g. `"ascend"`) to the op's
+ `candidate_paths` in `rl_engine/kernels/gtest/operator_specs.py`. Registration
+ itself is the CI gate (`tests/test_ws1_gtest_gpu.py` checks every WS1 op is in the
+ spec).
+6. **PR must report exact commands and results**: give the actual pytest command
+ line, the gtest command line (`scripts/check_operator.py` with full arguments),
+ and the outputs (template below).
+7. **Backward must also be deterministic**: whenever the op needs a backward, the
+ backward must be a deterministic implementation (rules in the Ascend section).
+
+## CUDA
+
+(Placeholder — to be filled in.)
+
+## ROCm
+
+(Placeholder — to be filled in.)
+
+## Ascend (battle-tested workflow)
+
+### Branch and PR conventions
+
+- Branch off `upstream/test` (NOT `main`): `git checkout -b feat/ascend-deterministic- upstream/test`.
+- PR base = `test`; title format: `[WS1][Ascend] [Qwen3-8b] ops` (e.g.
+ `[WS1][Ascend] [Qwen3-8b] Fused logp ops`).
+- Pushing: authenticate `gh` (`gh auth login --with-token`, then
+ `gh auth setup-git`); if direct github.com connectivity is flaky, push through
+ the ghfast proxy (token as the proxy host's userinfo). `gh pr edit --base` hits a
+ GraphQL classic-projects deprecation error — use
+ `gh api -X PATCH repos/RL-Align/RL-Kernel/pulls/ -f base=test` instead.
+- PR description drafts: keep a scratch directory OUTSIDE the repo for
+ `PR_DESCRIPTION_*.md` drafts (template below).
+
+### Implementation checklist (file level)
+
+The complete landing list for a new Ascend op:
+
+1. `csrc/ascend/_ascend.asc` — Ascend C kernel + torch host wrapper. No
+ `PYBIND11_MODULE` (consolidated in npu_module.cpp).
+2. `csrc/ascend/npu_module.cpp` — the single pybind entry declaring and binding all
+ ops. Each `.asc` carrying its own `PYBIND11_MODULE` causes duplicate
+ `PyInit__C_npu` link errors; for an existing `.asc` (e.g. batch_invariant_logp)
+ just drop its `PYBIND11_MODULE` block.
+3. `setup.py` — port the Ascend extension build (bisheng, `**/*.asc` glob,
+ `_find_ascend_home()` exporting `ASCEND_HOME_PATH`/`ASCEND_TOOLKIT_HOME`).
+ Fastest: `git checkout -- setup.py scripts/check_operator.py`.
+4. `rl_engine/_C_npu.pyi` — type stub (black: no blank line between two top-level
+ defs).
+5. `rl_engine/kernels/ops/ascend//.py` — the op wrapper (mirror the CUDA
+ wrapper's surface: `__call__`/`apply`/`forward`/`forward_fp32`, dtype gate,
+ `_NPU_EXT_AVAILABLE` + `hasattr(_C_npu, ...)` check, native fallback path).
+6. `rl_engine/kernels/gtest/operator_specs.py` — the `"ascend"` candidate.
+7. `rl_engine/kernels/registry.py` — `ASCEND_` enum member + npu priority map
+ override (`self._priority_map["npu"][""] = [ASCEND_..., PYTORCH_...]`).
+8. `rl_engine/tests/test_dispatch.py` — npu priority assertion.
+9. `tests/test__ascend.py` — pytest suite (PR 320 style, see below).
+10. `docs/operators/.md` — Ascend row in the Backends table, npu dispatch
+ paragraph, Tests and Implementation Files updates (**keep existing entries**,
+ add only).
+11. `scripts/check_operator.py` — already supports `--device npu` (auto-detect).
+
+Build and smoke test:
+
+```bash
+KERNEL_ALIGN_FORCE_ASCEND=1 pip install -e . --no-build-isolation
+```
+
+### Bitwise-consistency rules (mandatory; must be stated clearly in the PR)
+
+Classify the op BEFORE writing the PR:
+
+- **Copy/lookup ops (elementwise, e.g. embedding)**: the forward is a pure byte
+ move, so it MUST be bitwise-identical to the PyTorch golden — assert with
+ `torch.equal`.
+- **Reduction ops (reduction / logprob, e.g. lm_head, logp, fused linear logp)**:
+ **no independent kernel can be bitwise-identical to the golden** — the golden's
+ reduction order is the private implementation of
+ `torch.mv`/`torch.matmul`/`logsumexp`, fp32 addition is not associative, and two
+ different reduction trees over D=4096 inevitably drift ~1e-4 (the logprob
+ contract's fp32 atol=1e-5 is naturally unmeetable). Practice:
+ - The bitwise guarantee goes to **batch invariance on the NPU**: the same row
+ content across batch 1 vs {2,4,16,300}, different positions, strided blocks
+ (>MAX_BLOCKS), multi-tile shapes, repeated runs — all asserted with
+ `torch.equal`.
+ - Compare against the golden at the existing contract tolerances; state
+ prominently in a blockquote at the top of the PR body WHY bitwise parity is
+ impossible (golden's private reduction order + measured drift numbers).
+- **Never touch tolerances**: `rl_engine/kernels/gtest/tolerance_contract.json` is
+ read-only; look up rows by op_class x dtype.
+
+Known NPU-side golden gotchas (check before writing tests):
+- NPU `torch.mv` **rejects bf16** → golden references must go through the
+ `forward_fp32` paths.
+- The gtest `linear_logp` forward comparison is unwinnable even for the CUDA
+ candidate (the golden's `apply()` accumulates the matmul in the input dtype); CI
+ never executes that candidate, it only checks registration. Do not try to adjust
+ tolerances for it.
+- `torch.argsort(int64, stable=True)` on NPU runs on the AiCpu — a performance
+ warning only, results are correct.
+
+### Ascend C kernel gotchas (each one hit on real hardware)
+
+- **Cross-pipe race on shared UB buffers**: when one UB tile is written by MTE2 and
+ read by MTE3, use the canonical two-queue GM→UB→GM pipeline; the fixed out-queue
+ order is `AllocTensor → EnQue → DeQue → DataCopy → FreeTensor` (the queues
+ provide the MTE2→V / V→MTE3 sync). Do NOT hand-roll `MTE2_MTE3`/`MTE3_MTE2`
+ flags (random data corruption or hangs).
+- **Vector ops need 32B-aligned counts**: UB→UB `DataCopy`, `Cast`, etc. report
+ "VEC supports illegal configurations" for small counts → round the count up to a
+ multiple of `32/sizeof(T)` (over-copy inside UB is harmless; the copy-out writes
+ only the real byte count to GM).
+- **GM scalar reads/writes are unreliable**: `GlobalTensor.GetValue/SetValue` has
+ hardware issues — always read through a 32B `DataCopyPad` window (int64 window =
+ 4 per 32B, fp32 = 8 per 32B), sync with an `MTE2_S` flag before `GetValue`.
+- **`SyncAll` deadlocks**: with more blocks launched than physical cores the
+ cross-core barrier deadlocks — only per-pipe `SetFlag`/`WaitFlag` (V_S, S_V,
+ S_MTE3, MTE3_S, ...).
+- **Strided rows across blocks**: `MAX_BLOCKS=128`,
+ `for (row = GetBlockIdx(); row < N; row += GetBlockNum())`, host side
+ `blockNum = min(N, MAX_BLOCKS)` — each row is processed end-to-end by one block,
+ so the instruction sequence depends only on the shape, never on batch layout or
+ block assignment (the foundation of batch invariance).
+- **Scalar math in the kernel**: the scalar unit has no exp/log → use a padded
+ 8-element vector `Exp`/`Log` (`SetValue → S_V flag → vector op → V_S wait →
+ GetValue`).
+- **Output staging**: `SetValue` into a UB scalar buffer, `S_MTE3` flag, then
+ `DataCopyPad` out to GM; drain with `MTE3_S` after each row so the next row does
+ not overwrite the staging area.
+- **fp16/bf16 output cast**: `Cast(..., RoundMode::CAST_RINT, 32/sizeof(T))` —
+ CAST_RINT is IEEE round-to-nearest, matching CUDA's `static_cast` semantics.
+- **bisheng build**: needs `ASCEND_HOME_PATH` (setup.py exports it automatically);
+ when pip swallows the real compiler error, compile the `.asc` manually with
+ `bisheng` to see it.
+- **const pointers**: kernel-launch GM_ADDR parameters take `uint8_t*` (non-const).
+
+### Backward determinism
+
+Per PR #299 (frank-2077, FFN deterministic backward): **the backward is assembled
+from existing deterministic forward kernels — no new reductions, no fallback to
+cuBLAS/torch.matmul**. Priority order:
+
+1. **Reuse a pure-PyTorch deterministic formula**: when the CUDA op's backward is
+ itself pure PyTorch (e.g. embedding's sorted-segment dweight: stable argsort +
+ unique_consecutive + fixed-order accumulation), the Ascend op reuses the exact
+ same function → bitwise-identical backward.
+2. **Row-local fp32 VJP formulas** (logp / linear_logp / lm_head): compute the VJP
+ in fp32 with torch ops, cast back to the input dtype at the end; no cross-row
+ reduction → batch-layout independent.
+3. **GEMM-shaped backward**: assemble with `det_gemm` forwards
+ (`grad_hidden = det_gemm(grad, W)`, `grad_weight = det_gemm(grad^T, H)`); the
+ wrapper must raise when the det_gemm symbols are missing instead of silently
+ falling back.
+4. **TP scenarios**: mind the shard semantics (PR #299 checklist: gate/up input
+ grads each take one AllReduce, weight grads stay column-parallel shards, down is
+ row-parallel, etc.).
+5. Low-precision gradient comparisons: when both implementations compute the VJP in
+ fp32 and quantize at the end, compare against the quantization-aligned
+ reference (`ref_grad.to(dtype)`) — can be bitwise equal; tolerances only absorb
+ the rare 1-ULP straddle.
+
+### pytest suite conventions (PR 320 style)
+
+`tests/test__ascend.py` structure:
+
+- Module docstring stating the two orthogonal properties (correctness + batch
+ invariance).
+- `_npu_available()` / `_ascend_kernel_available()` helpers +
+ `requires_ascend = pytest.mark.skipif(...)`.
+- `TestAscendCorrectness`: class-level
+ `@pytest.mark.parametrize("dtype", [fp32, bf16, fp16])`; forward vs golden
+ (bitwise for copy ops / contract tolerance for reductions), `forward_fp32`,
+ out-of-range targets, backward, bias (if any).
+- `TestAscendBatchInvariance`: bitwise (`torch.equal`).
+- `TestAscendRegistryDispatch`: `kernel_registry.get_op("", device="npu")`
+ `type(op).__name__` assertion.
+
+Test bugs already hit (check before writing new tests):
+- Under class-level `parametrize`, every method must take the `dtype` argument —
+ move tests that don't into their own class.
+- Batch-comparison tests must **reuse the same weight** (regenerating with the same
+ seed produces a different weight for different batch sizes).
+- Position-invariance tests: pin the same row content
+ (`logits[pos].copy_(base)` + `target[pos] = base_id`).
+- Row-local VJP bitwise assertions: align the `grad_out` rows too
+ (`grad_out[1] = grad_out[0]`).
+
+### PR description template
+
+Structure (mirror the wording of previous Ascend PR descriptions):
+
+```markdown
+## Latest Status [date]
+Ready for review.
+
+## Summary
+- Bitwise-consistency status (prominent blockquote — mandatory for reduction ops)
+- Forward kernel design (which CUDA kernel/contract it mirrors)
+- Wrapper / Backward / Registration / Build
+
+## Files (table, one row per file with Status)
+
+## Test
+# The exact commands that were run:
+export KERNEL_ALIGN_FORCE_ASCEND=1
+pip install -e . --no-build-isolation
+python scripts/check_operator.py --op --candidate ascend --device npu \
+ --dtype {fp32,bf16,fp16} --batch 2 --seq 16 --vocab 257 --normalized-dim 4096 --check-grad
+python -m pytest tests/test__ascend.py -v
+python -m pytest tests/test_batch_invariant_logp.py -q # regression
+python -m pytest rl_engine/tests/test_dispatch.py -q # regression
+
+## Test results (environment line + results table + folded raw output)
+
+## Notes
+```
+
+Must include: the actual test environment (NPU model, CANN version, torch +
+torch_npu versions), per-dtype gtest output and pytest results,
+bitwise-invariance conclusions, regression results, pre-commit status.
diff --git a/.github/workflows/ws1-chain-npu.yml b/.github/workflows/ws1-chain-npu.yml
new file mode 100644
index 00000000..362c18b3
--- /dev/null
+++ b/.github/workflows/ws1-chain-npu.yml
@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# WS1 C10/C11 full Qwen3-8B Dense model-level gate on Ascend NPU (ascend_bf16).
+# Required check: no skip / xfail / synthetic weights / silent fallback.
+#
+# Unlike the CUDA job there is no cloud NPU provider wired up here, so this runs
+# on a self-hosted Ascend runner (Atlas A2 / 910B with CANN + torch_npu) that a
+# maintainer registers with the labels below. Without such a runner the job
+# queues rather than reporting a false pass - a required profile that did not
+# execute is red, never N/A.
+#
+# Security: do not use pull_request_target. Fork PRs never reach the self-hosted
+# runner; a maintainer dispatches the reviewed SHA from a trusted branch.
+
+name: WS1-chain-NPU
+
+on:
+ pull_request:
+ branches: [ main, test ]
+ push:
+ branches: [ main, test ]
+ workflow_dispatch:
+ inputs:
+ source_repository:
+ description: "Public repository containing the reviewed commit (owner/name)"
+ required: true
+ default: "RL-Align/RL-Kernel"
+ type: string
+ source_sha:
+ description: "Exact reviewed 40-character commit SHA to execute on the NPU host"
+ required: true
+ type: string
+
+concurrency:
+ group: ws1-chain-npu-${{ github.ref }}
+ cancel-in-progress: false
+
+permissions:
+ contents: read
+
+jobs:
+ fork-pr-notice:
+ if: github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository
+ runs-on: ubuntu-latest
+ steps:
+ - name: Report required trusted execution
+ run: |
+ echo "Fork code does not run on the self-hosted Ascend runner."
+ echo "A maintainer must dispatch this workflow from a trusted upstream branch."
+ echo "source_repository=${{ github.event.pull_request.head.repo.full_name }}"
+ echo "source_sha=${{ github.event.pull_request.head.sha }}"
+
+ ws1-chain-npu:
+ if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository
+ runs-on: [ self-hosted, linux, ascend-npu ]
+ timeout-minutes: 240
+ env:
+ # Set on the runner: the pinned Qwen3-8B Dense snapshot directory.
+ WS1_WEIGHTS_PATH: ${{ vars.WS1_WEIGHTS_PATH }}
+ RL_KERNEL_REQUIRE_EXT: "1"
+ WS1_WORKFLOW_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
+ steps:
+ - name: Validate trusted dispatch target
+ if: github.event_name == 'workflow_dispatch'
+ env:
+ SOURCE_REPOSITORY: ${{ inputs.source_repository }}
+ SOURCE_SHA: ${{ inputs.source_sha }}
+ run: |
+ [[ "$SOURCE_REPOSITORY" =~ ^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$ ]]
+ [[ "$SOURCE_SHA" =~ ^[0-9a-fA-F]{40}$ ]]
+
+ - name: Checkout reviewed commit
+ uses: actions/checkout@v4
+ with:
+ repository: ${{ github.event_name == 'workflow_dispatch' && inputs.source_repository || github.repository }}
+ ref: ${{ github.event_name == 'workflow_dispatch' && inputs.source_sha || github.event.pull_request.head.sha || github.sha }}
+ persist-credentials: false
+
+ - name: Report Ascend environment
+ run: |
+ python3 -c "import torch, torch_npu; print('torch', torch.__version__, 'torch_npu', torch_npu.__version__)"
+ npu-smi info || true
+
+ - name: Run WS1 Ascend C3-C11 gates
+ run: bash ci/run_ws1_ascend_ci.sh
+
+ - name: Upload C2/C8/C10 JSON
+ if: always()
+ uses: actions/upload-artifact@v4
+ with:
+ name: ws1-closeout-ascend
+ path: |
+ /tmp/ws1-c2-ascend.json
+ /tmp/ws1-c8-ascend.json
+ /tmp/ws1-c10-ascend_bf16.json
+ if-no-files-found: error
diff --git a/benchmarks/benchmark_attention.py b/benchmarks/benchmark_attention.py
index c9509a19..99705a52 100644
--- a/benchmarks/benchmark_attention.py
+++ b/benchmarks/benchmark_attention.py
@@ -1,70 +1,187 @@
-# File: benchmarks/benchmark_attention.py
-import pandas as pd
-import torch
-import triton
-from tabulate import tabulate
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
-from rl_engine.kernels.ops.cuda.attention.prefix_shared_attn import PrefixSharedAttentionOp
+"""Benchmark deterministic standard-softmax attention across backends.
+All backends compute ``softmax(Q K^T * scale + causal mask) @ V`` with a locked,
+per-row reduction order (batch-invariant, no split-K). The comparison here is
+latency across a sequence sweep:
-def run_benchmark():
- bs = 1
- G = 64
- len_q = 512
- dim = 128
+- Native is the pure-PyTorch fp32-accumulating ground-truth reference.
+- Ascend is the CANN two-pass streaming kernel (one AI core block/row); only
+ present when the extension is built with ``KERNEL_ALIGN_FORCE_ASCEND=1``.
+- CUDA is the deterministic op (one CTA/row); only present when the extension
+ is built with ``KERNEL_ALIGN_FORCE_SM90=1`` on an SM90 device.
- len_kvs = [1024, 2048, 4096, 8192, 16384]
+Timing dispatch through the active accelerator (``torch.cuda`` or
+``torch.npu``), so the benchmark runs on CUDA/ROCm/NPU devices.
- print("Benchmarking GRPO Prefix-Shared Attention")
- print(f"Fixed Shapes: Batch={bs}, Group(Response)={G}, Query_Len={len_q}, Head_Dim={dim}\n")
+Usage:
+ python benchmarks/benchmark_attention.py
+ python benchmarks/benchmark_attention.py --backward
+ python benchmarks/benchmark_attention.py --configs "1,8,512;2,8,2048"
+"""
- prefix_shared_sdpa = PrefixSharedAttentionOp()
- results = []
+import argparse
- for len_kv in len_kvs:
- q = torch.randn(bs, G, len_q, dim, dtype=torch.bfloat16, device="cuda")
- k = torch.randn(bs, len_kv, dim, dtype=torch.bfloat16, device="cuda")
- v = torch.randn(bs, len_kv, dim, dtype=torch.bfloat16, device="cuda")
+import torch
+from tabulate import tabulate
- k_exp = k.unsqueeze(1).expand(-1, G, -1, -1).reshape(bs * G, 1, len_kv, dim).contiguous()
- v_exp = v.unsqueeze(1).expand(-1, G, -1, -1).reshape(bs * G, 1, len_kv, dim).contiguous()
- q_res = q.view(bs * G, 1, len_q, dim)
+from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp
+from rl_engine.platforms.device import device_ctx
+from rl_engine.utils.logger import logger
- for _ in range(5):
- _ = torch.nn.functional.scaled_dot_product_attention(q_res, k_exp, v_exp)
- _ = prefix_shared_sdpa(q, k, v)
+_D = 128
+_DEFAULT_KV_HEADS = 8
- native_ms = triton.testing.do_bench(
- lambda: torch.nn.functional.scaled_dot_product_attention(q_res, k_exp, v_exp),
- return_mode="median",
- )
- custom_ms = triton.testing.do_bench(
- lambda: prefix_shared_sdpa(q, k, v), return_mode="median"
- )
+def _accel():
+ """The active accelerator module (torch.npu on Ascend, torch.cuda otherwise)."""
+ if device_ctx.device_type == "npu":
+ return torch.npu
+ return torch.cuda
- speedup = native_ms / custom_ms
- reduction = (native_ms - custom_ms) / native_ms * 100
-
- flops = 4 * bs * G * len_q * len_kv * dim
- native_tflops = (flops / 1e12) / (native_ms / 1000)
- custom_tflops = (flops / 1e12) / (custom_ms / 1000)
-
- results.append(
- {
- "Prompt Len": len_kv,
- "Native (ms)": f"{native_ms:.3f}",
- "RL-Kernel (ms)": f"{custom_ms:.3f}",
- "Native TFLOPS": f"{native_tflops:.1f}",
- "RL-Kernel TFLOPS": f"{custom_tflops:.1f}",
- "Speedup": f"{speedup:.2f}x",
- "Time Saved": f"{reduction:.1f}%",
- }
- )
- df = pd.DataFrame(results)
- print(tabulate(df, headers="keys", tablefmt="pretty", stralign="center", showindex=False))
+def _maybe_ascend_op():
+ """The Ascend C op, or None when unavailable (no NPU / not built)."""
+ if device_ctx.device_type != "npu":
+ return None
+ try:
+ from rl_engine.kernels.ops.ascend.attention.deterministic_attn import (
+ DeterministicAttentionAscendOp,
+ )
+ except (ImportError, RuntimeError):
+ return None
+ return DeterministicAttentionAscendOp()
+
+
+def _maybe_cuda_op():
+ """The CUDA deterministic op, or None when unavailable (no CUDA / not built)."""
+ from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE
+
+ if not (
+ torch.cuda.is_available()
+ and _EXT_AVAILABLE
+ and hasattr(_C, "deterministic_attention_forward")
+ ):
+ return None
+ from rl_engine.kernels.ops.cuda.attention.deterministic_attn import DeterministicAttentionOp
+
+ return DeterministicAttentionOp()
+
+
+# (batch, num_q_heads, seqlen); Hq = 32, Hkv = 8 (Qwen3-style GQA, g = 4).
+DEFAULT_CONFIGS = [
+ (1, 32, 512),
+ (1, 32, 1024),
+ (1, 32, 2048),
+ (4, 32, 2048),
+]
+
+
+def _make_inputs(batch, hq, seq, device, dtype):
+ generator = torch.Generator(device="cpu").manual_seed(0)
+ q = torch.randn(batch, hq, seq, _D, dtype=dtype, generator=generator).to(device)
+ k = torch.randn(batch, _DEFAULT_KV_HEADS, seq, _D, dtype=dtype, generator=generator).to(device)
+ v = torch.randn(batch, _DEFAULT_KV_HEADS, seq, _D, dtype=dtype, generator=generator).to(device)
+ return q, k, v
+
+
+def _time_ms(fn, warmup, iters):
+ acc = _accel()
+ for _ in range(warmup):
+ fn()
+ acc.synchronize()
+ start = acc.Event(enable_timing=True)
+ end = acc.Event(enable_timing=True)
+ start.record()
+ for _ in range(iters):
+ fn()
+ end.record()
+ acc.synchronize()
+ return start.elapsed_time(end) / iters
+
+
+def _forward_closure(op, q, k, v):
+ def run():
+ with torch.no_grad():
+ op(q, k, v, causal=True)
+
+ return run
+
+
+def _forward_backward_closure(op, q, k, v):
+ def run():
+ qq = q.detach().requires_grad_(True)
+ kk = k.detach().requires_grad_(True)
+ vv = v.detach().requires_grad_(True)
+ op(qq, kk, vv, causal=True).sum().backward()
+
+ return run
+
+
+def _bench_table(configs, native_op, other_ops, closure_factory, device, dtype, warmup, iters):
+ label = "fwd" if closure_factory is _forward_closure else "fwd+bwd"
+ rows = []
+ for batch, hq, seq in configs:
+ q, k, v = _make_inputs(batch, hq, seq, device, dtype)
+ n_c = closure_factory(native_op, q, k, v)
+ n_ms = _time_ms(n_c, warmup, iters)
+ row = [f"{batch}x{hq}x{seq}", f"{n_ms:.3f}"]
+ for _name, op in other_ops:
+ if op is None:
+ row += ["-"]
+ continue
+ o_ms = _time_ms(closure_factory(op, q, k, v), warmup, iters)
+ row += [f"{o_ms:.3f}", f"{n_ms / o_ms:.2f}x"]
+ rows.append(row)
+
+ headers = ["shape (B x Hq x S)", f"native {label} ms"]
+ for name, _ in other_ops:
+ headers += [f"{name} {label} ms", "vs native"]
+ logger.info("\n" + tabulate(rows, headers=headers, tablefmt="github"))
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ "--backward", action="store_true", help="also emit the forward+backward table"
+ )
+ parser.add_argument(
+ "--configs",
+ default=";".join(",".join(map(str, c)) for c in DEFAULT_CONFIGS),
+ help="semicolon-separated 'batch,hq,seq' triples",
+ )
+ parser.add_argument("--dtype", choices=("bf16", "fp16"), default="bf16")
+ parser.add_argument("--warmup", type=int, default=5)
+ parser.add_argument("--iters", type=int, default=20)
+ args = parser.parse_args()
+
+ configs = [tuple(int(x) for x in part.split(",")) for part in args.configs.split(";")]
+ dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
+ device = torch.device(device_ctx.device_type)
+
+ native_op = NativeAttentionOp()
+ other_ops = [
+ ("ascend", _maybe_ascend_op()),
+ ("cuda", _maybe_cuda_op()),
+ ]
+
+ _bench_table(
+ configs, native_op, other_ops, _forward_closure, device, dtype, args.warmup, args.iters
+ )
+ if args.backward:
+ _bench_table(
+ configs,
+ native_op,
+ other_ops,
+ _forward_backward_closure,
+ device,
+ dtype,
+ args.warmup,
+ args.iters,
+ )
if __name__ == "__main__":
- run_benchmark()
+ main()
diff --git a/benchmarks/benchmark_rmsnorm.py b/benchmarks/benchmark_rmsnorm.py
index 81325281..5227729c 100644
--- a/benchmarks/benchmark_rmsnorm.py
+++ b/benchmarks/benchmark_rmsnorm.py
@@ -14,14 +14,22 @@
except ImportError:
HAS_CUDA_EXT = False
+try:
+ from rl_engine.kernels.ops.ascend.norm.rmsnorm import RMSNormAscendOp
+
+ HAS_ASCEND_EXT = True
+except (ImportError, OSError, RuntimeError):
+ HAS_ASCEND_EXT = False
+
def bench(fn, x, w, dy, warmup=20, iters=100):
+ sync = torch.npu.synchronize if x.device.type == "npu" else torch.cuda.synchronize
for _ in range(warmup):
x.grad = None
w.grad = None
y = fn(x, w)
y.backward(dy)
- torch.cuda.synchronize()
+ sync()
start = time.time()
for _ in range(iters):
@@ -29,7 +37,7 @@ def bench(fn, x, w, dy, warmup=20, iters=100):
w.grad = None
y = fn(x, w)
y.backward(dy)
- torch.cuda.synchronize()
+ sync()
return (time.time() - start) * 1000.0 / iters
@@ -41,7 +49,7 @@ def main():
args = parser.parse_args()
dtype = torch.float16 if args.dtype == "fp16" else torch.bfloat16
- device = "cuda"
+ device = "npu" if torch.npu.is_available() else "cuda"
T, H = args.T, args.H
torch.manual_seed(0)
@@ -72,6 +80,15 @@ def make_inputs():
else:
print("cuda : skipped, extension is not built")
+ if device == "npu":
+ ascend_op = RMSNormAscendOp() if HAS_ASCEND_EXT else None
+ if ascend_op is not None:
+ x, w = make_inputs()
+ t_asc = bench(lambda a, b: ascend_op(a, b), x, w, dy)
+ print(f"ascend : {t_asc:.4f} ms | speedup vs ref: {t_ref / t_asc:.2f}x")
+ else:
+ print("ascend : skipped, extension is not built")
+
if __name__ == "__main__":
main()
diff --git a/ci/run_ws1_ascend_ci.sh b/ci/run_ws1_ascend_ci.sh
new file mode 100755
index 00000000..b7eb4d23
--- /dev/null
+++ b/ci/run_ws1_ascend_ci.sh
@@ -0,0 +1,143 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# WS1 C3-C11 gate for the Ascend BF16 profile (#266, ascend_bf16).
+#
+# Runs on an Ascend host (Atlas A2 / 910B) with CANN and torch_npu. It is the
+# NPU twin of ci/run_ws1_gtest.sh + ci/run_ws1_chain_gate.sh: same contract,
+# same harnesses, same fail-closed rules. Nothing here may fall back to CPU or
+# to another vendor's kernels - a required profile that cannot run is red.
+#
+# Required:
+# WS1_WEIGHTS_PATH (or QWEN3_8B) pinned Qwen3-8B Dense safetensors snapshot
+# Optional:
+# PY interpreter (default python3)
+# WS1_SKIP_BUILD=1 reuse an already-built rl_engine._C_npu
+
+set -euo pipefail
+
+ROOT="$(cd "$(dirname "$0")/.." && pwd)"
+cd "$ROOT"
+
+PY="${PY:-python3}"
+export RL_KERNEL_REQUIRE_EXT="${RL_KERNEL_REQUIRE_EXT:-1}"
+WEIGHTS_PATH="${WS1_WEIGHTS_PATH:-${QWEN3_8B:-}}"
+
+echo "[ws1-ascend] interpreter=$PY"
+
+if [ "${WS1_SKIP_BUILD:-0}" != "1" ]; then
+ echo "[ws1-ascend] building the Ascend C extension"
+ KERNEL_ALIGN_FORCE_ASCEND=1 "$PY" -m pip install -e . --no-build-isolation --no-deps
+fi
+
+# Fail before any gate if the NPU or the compiled kernels are missing, so a
+# later red cell is never confused with an environment problem.
+"$PY" - <<'PY'
+import sys
+
+from rl_engine.kernels.gtest.accelerator import describe, npu_available, resolve_device
+
+if not npu_available():
+ sys.exit("[ws1-ascend] FATAL: torch_npu reports no available NPU")
+info = describe(resolve_device(None, profile="ascend_bf16"))
+print(f"[ws1-ascend] device={info.device} name={info.name} soc={info.arch_key}")
+
+from rl_engine import _C_npu # noqa: E402
+
+required = (
+ "rmsnorm_ascend",
+ "rope_apply_ascend",
+ "deterministic_attention_ascend",
+ "embedding_ascend",
+ "lm_head_ascend",
+ "fused_logp_ascend",
+ "batch_invariant_logp_ascend",
+ "swiglu_forward",
+ "silu_forward",
+ "det_gemm_ascend_fwd",
+ "det_gemm_rowwise_ascend_fwd_fp32",
+)
+missing = [name for name in required if not hasattr(_C_npu, name)]
+if missing:
+ sys.exit(f"[ws1-ascend] FATAL: _C_npu is missing {missing}; rebuild the extension")
+print(f"[ws1-ascend] all {len(required)} required Ascend entry points are linked")
+PY
+
+echo "[ws1-ascend] CPU-side contract, workload and wiring tests"
+"$PY" -m pytest -q \
+ tests/test_tolerance_contract.py \
+ tests/test_ws1_workload.py \
+ tests/test_four_judgment_matrix.py \
+ tests/test_elementwise_inventory.py \
+ tests/test_ws1_ascend_closeout.py
+
+echo "[ws1-ascend] Ascend operator tests"
+"$PY" -m pytest -q \
+ tests/test_det_gemm_ascend.py \
+ tests/test_silu_ascend.py
+
+echo "[ws1-ascend] C2 runtime candidate evidence"
+"$PY" scripts/ws1_candidate_evidence.py \
+ --profile ascend_bf16 --all --check-grad --emit-json /tmp/ws1-c2-ascend.json
+"$PY" - /tmp/ws1-c2-ascend.json <<'PY'
+import json
+import sys
+
+payload = json.load(open(sys.argv[1], encoding="utf-8"))
+if not payload.get("passed"):
+ failed = [c["case_id"] for c in payload["cases"] if c["runtime_status"] != "passed"]
+ raise SystemExit(f"C2 Ascend runtime evidence failed: {failed}")
+print(f"[ws1-ascend] C2 evidence passed for {len(payload['cases'])} pinned cases")
+PY
+
+echo "[ws1-ascend] C3/C4 smoke (silu)"
+"$PY" scripts/check_forward_invariance.py \
+ --op silu --candidate ascend --backend-profile ascend_bf16
+"$PY" scripts/check_gradient_invariance.py \
+ --op silu --candidate ascend --backend-profile ascend_bf16
+
+echo "[ws1-ascend] C6 direct decode-prefill"
+"$PY" scripts/check_decode_prefill.py --backend-profile ascend_bf16
+echo "[ws1-ascend] C7 stateful KV + generate-rescore"
+"$PY" scripts/check_stateful_kv.py --backend-profile ascend_bf16
+
+C8_OUT="${WS1_C8_JSON:-${TMPDIR:-/tmp}/ws1-c8-ascend.json}"
+export WS1_C8_EVIDENCE_PATH="$C8_OUT"
+echo "[ws1-ascend] C8 four-judgment sweep -> $C8_OUT"
+"$PY" scripts/sweep_ws1_four_judgments.py \
+ --execute --profile ascend_bf16 --json > "$C8_OUT"
+"$PY" - "$C8_OUT" <<'PY'
+import json
+import subprocess
+import sys
+
+payload = json.load(open(sys.argv[1], encoding="utf-8"))
+git_meta = payload.get("git") or {}
+expected = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip()
+if git_meta.get("commit") != expected or git_meta.get("dirty"):
+ raise SystemExit(f"C8 is not from the clean current commit: {git_meta}")
+counts = payload.get("counts") or {}
+if int(counts.get("red", 0)):
+ raise SystemExit(f"C8 contains red rows: {counts}")
+if int(counts.get("green", 0)) == 0:
+ raise SystemExit("C8 artifact has no green cells")
+for cell in payload.get("cells") or []:
+ if cell.get("op_name") == "pack" or cell.get("status") != "green":
+ continue
+ if not cell.get("judgment", "").endswith("invariance"):
+ continue
+ if not cell.get("actual_backend_id") or not cell.get("actual_kernel_config_id"):
+ raise SystemExit(
+ f"invariance cell missing provenance: {cell.get('profile')} {cell.get('op_name')}"
+ )
+print(f"[ws1-ascend] C8 passed counts={counts}")
+PY
+
+if [ -z "$WEIGHTS_PATH" ]; then
+ echo "[ws1-ascend] FATAL: set WS1_WEIGHTS_PATH or QWEN3_8B for the C10/C11 full-model gate"
+ exit 2
+fi
+
+echo "[ws1-ascend] C10/C11 full Qwen3-8B Dense model gate"
+WS1_PROFILES="ascend_bf16" WS1_C8_JSON="$C8_OUT" bash ci/run_ws1_chain_gate.sh
+
+echo "[ws1-ascend] ascend_bf16 passed every required WS1 gate"
diff --git a/ci/run_ws1_chain_gate.sh b/ci/run_ws1_chain_gate.sh
index b8382ba4..c367ab10 100755
--- a/ci/run_ws1_chain_gate.sh
+++ b/ci/run_ws1_chain_gate.sh
@@ -1,7 +1,12 @@
#!/usr/bin/env bash
# SPDX-License-Identifier: Apache-2.0
-# WS1 C10/C11 full Qwen3-8B Dense model-level gate (CUDA BF16 and Triton-on-CUDA BF16).
-# Intended for H20 / H100. Fails closed on skip, xfail, synthetic weights, or silent fallback.
+# WS1 C10/C11 full Qwen3-8B Dense model-level gate.
+# Profiles come from WS1_PROFILES (default: the two CUDA-host profiles). One host
+# has either a GPU or an NPU, so each vendor's job runs its own profiles here:
+# CUDA host : WS1_PROFILES="cuda_bf16 triton_cuda_bf16" (H20 / H100)
+# Ascend host : WS1_PROFILES="ascend_bf16" (Atlas A2 / 910B)
+# C11 closes only when every required profile has passed on its own hardware.
+# Fails closed on skip, xfail, synthetic weights, or silent fallback.
set -euo pipefail
@@ -11,13 +16,14 @@ cd "$ROOT"
PY="${PY:-python3}"
export RL_KERNEL_REQUIRE_EXT="${RL_KERNEL_REQUIRE_EXT:-1}"
WEIGHTS_PATH="${WS1_WEIGHTS_PATH:-${QWEN3_8B:-}}"
+WS1_PROFILES="${WS1_PROFILES:-cuda_bf16 triton_cuda_bf16}"
if [ -z "$WEIGHTS_PATH" ]; then
echo "[ws1-chain] FATAL: set WS1_WEIGHTS_PATH or QWEN3_8B to the pinned Qwen3-8B snapshot"
exit 2
fi
-echo "[ws1-chain] interpreter=$PY weights=$WEIGHTS_PATH"
+echo "[ws1-chain] interpreter=$PY weights=$WEIGHTS_PATH profiles=$WS1_PROFILES"
"$PY" -m pytest -q \
tests/test_kv_consistency.py \
@@ -27,7 +33,11 @@ echo "[ws1-chain] interpreter=$PY weights=$WEIGHTS_PATH"
C8_OUT="${WS1_C8_JSON:-${TMPDIR:-/tmp}/ws1-c8-ci.json}"
export WS1_C8_EVIDENCE_PATH="$C8_OUT"
echo "[ws1-chain] C8 runtime evidence $C8_OUT"
-"$PY" scripts/sweep_ws1_four_judgments.py --execute --json > "$C8_OUT"
+C8_PROFILE_ARGS=()
+for PROFILE in $WS1_PROFILES; do
+ C8_PROFILE_ARGS+=(--profile "$PROFILE")
+done
+"$PY" scripts/sweep_ws1_four_judgments.py --execute "${C8_PROFILE_ARGS[@]}" --json > "$C8_OUT"
"$PY" - "$C8_OUT" <<'PY'
import json
import subprocess
@@ -44,7 +54,7 @@ if int((payload.get("counts") or {}).get("red", 0)):
print(f"[ws1-chain] C8 passed source={git_meta}")
PY
-for PROFILE in cuda_bf16 triton_cuda_bf16; do
+for PROFILE in $WS1_PROFILES; do
OUT="/tmp/ws1-c10-${PROFILE}.json"
echo "[ws1-chain] C10/C11 $PROFILE"
"$PY" scripts/ws1_chain_gate.py \
@@ -166,7 +176,11 @@ for kind in ("lm_head", "rms_norm", "det_gemm", "embedding"):
raise SystemExit(f"{profile} missing runtime backward record for {kind}")
if not event.get("kernel_id"):
raise SystemExit(f"{profile} backward {kind} missing kernel_id")
- family = "triton" if profile.startswith("triton") else "cuda"
+ family = {
+ "cuda_bf16": "cuda",
+ "triton_cuda_bf16": "triton",
+ "ascend_bf16": "ascend",
+ }[profile]
if not event.get("kernel_ids"):
raise SystemExit(f"{profile} backward {kind} missing kernel_ids")
if not event.get("implementation_ids"):
@@ -197,4 +211,4 @@ print(f"[ws1-chain] {profile} passed first_drift={payload.get('first_drift')}")
PY
done
-echo "[ws1-chain] both required profiles passed"
+echo "[ws1-chain] profiles passed: $WS1_PROFILES"
diff --git a/csrc/ascend/activation.asc b/csrc/ascend/activation.asc
new file mode 100644
index 00000000..6a6d3613
--- /dev/null
+++ b/csrc/ascend/activation.asc
@@ -0,0 +1,441 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+// SwiGLU: out = (gate * sigmoid(gate)) * up, with FP32 intermediates.
+// SiLU: out = x * sigmoid(x), the same tile shape with one operand.
+// Fixed elementwise tiles have no reductions or inter-core synchronization.
+
+#include
+#include
+
+#include "kernel_operator.h"
+#include
+#include
+#include "torch_npu/csrc/core/npu/NPUStream.h"
+
+namespace {
+
+constexpr uint32_t TILE_LENGTH = 2048;
+constexpr int64_t MAX_BLOCKS = 32;
+
+template
+class KernelSwiGLU {
+public:
+ __aicore__ inline void Init(AscendC::TPipe* pipe, GM_ADDR gate, GM_ADDR up,
+ GM_ADDR grad, GM_ADDR out, GM_ADDR dUp, int64_t n)
+ {
+ n_ = n;
+ gateGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(gate));
+ upGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(up));
+ outGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(out));
+ if constexpr (Backward) {
+ gradGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(grad));
+ dUpGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(dUp));
+ }
+ // Each segment starts on a 32-byte boundary, including for short tails.
+ // Worst-case UB use: (3 inputs + 2 outputs + 5 FP32 tiles) * 8 KiB = 80 KiB.
+ pipe->InitBuffer(inQueue_, 1, (Backward ? 3 : 2) * TILE_LENGTH * sizeof(T));
+ pipe->InitBuffer(outQueue_, 1, (Backward ? 2 : 1) * TILE_LENGTH * sizeof(T));
+ pipe->InitBuffer(work_, 5 * TILE_LENGTH * sizeof(float));
+ }
+
+ __aicore__ inline void Process()
+ {
+ for (int64_t offset = static_cast(AscendC::GetBlockIdx()) * TILE_LENGTH;
+ offset < n_;
+ offset += static_cast(AscendC::GetBlockNum()) * TILE_LENGTH) {
+ const uint32_t count = static_cast(
+ n_ - offset < TILE_LENGTH ? n_ - offset : TILE_LENGTH);
+ CopyIn(offset, count);
+ Compute(count);
+ CopyOut(offset, count);
+ }
+ }
+
+private:
+ __aicore__ inline void CopyIn(int64_t offset, uint32_t count)
+ {
+ auto input = inQueue_.AllocTensor();
+ AscendC::DataCopyExtParams params{
+ 1, static_cast(count * sizeof(T)), 0, 0, 0};
+ AscendC::DataCopyPadExtParams pad{false, 0, 0, 0};
+ AscendC::DataCopyPad(input, gateGm_[offset], params, pad);
+ AscendC::DataCopyPad(input[TILE_LENGTH], upGm_[offset], params, pad);
+ if constexpr (Backward) {
+ AscendC::DataCopyPad(input[2 * TILE_LENGTH], gradGm_[offset], params, pad);
+ }
+ inQueue_.EnQue(input);
+ }
+
+ __aicore__ inline void ToFloat(AscendC::LocalTensor dst,
+ AscendC::LocalTensor src, uint32_t count)
+ {
+ if constexpr (std::is_same_v) {
+ // UB-to-UB copies require a multiple of 32 bytes. Padding stays in UB.
+ AscendC::DataCopy(dst, src, (count + 7) / 8 * 8);
+ } else {
+ AscendC::Cast(dst, src, AscendC::RoundMode::CAST_NONE, count);
+ }
+ }
+
+ __aicore__ inline void Store(AscendC::LocalTensor dst,
+ AscendC::LocalTensor src, uint32_t count)
+ {
+ AscendC::PipeBarrier();
+ if constexpr (std::is_same_v) {
+ AscendC::DataCopy(dst, src, (count + 7) / 8 * 8);
+ } else {
+ // Round to nearest, ties to even, matching PyTorch dtype conversion.
+ AscendC::Cast(dst, src, AscendC::RoundMode::CAST_RINT, count);
+ }
+ // The caller may reuse src immediately for the next result.
+ AscendC::PipeBarrier();
+ }
+
+ __aicore__ inline void Compute(uint32_t count)
+ {
+ auto input = inQueue_.DeQue(); // MTE2 -> vector synchronization
+ auto output = outQueue_.AllocTensor();
+ auto gate = work_.Get();
+ auto up = gate[TILE_LENGTH];
+ auto grad = gate[2 * TILE_LENGTH];
+ auto sigmoid = gate[3 * TILE_LENGTH];
+ auto tmp = gate[4 * TILE_LENGTH];
+ ToFloat(gate, input, count);
+ ToFloat(up, input[TILE_LENGTH], count);
+ if constexpr (Backward) {
+ ToFloat(grad, input[2 * TILE_LENGTH], count);
+ }
+ AscendC::PipeBarrier();
+
+ AscendC::Muls(sigmoid, gate, -1.0f, count);
+ AscendC::PipeBarrier();
+ AscendC::Exp(sigmoid, sigmoid, count);
+ AscendC::PipeBarrier();
+ AscendC::Adds(sigmoid, sigmoid, 1.0f, count);
+ AscendC::Duplicate(tmp, 1.0f, count);
+ AscendC::PipeBarrier();
+ AscendC::Div(sigmoid, tmp, sigmoid, count);
+ AscendC::PipeBarrier();
+
+ AscendC::Mul(tmp, gate, sigmoid, count);
+ AscendC::PipeBarrier();
+ if constexpr (Backward) {
+ // d_up = grad * (gate * sigmoid(gate)).
+ AscendC::Mul(tmp, grad, tmp, count);
+ Store(output[TILE_LENGTH], tmp, count);
+
+ // d_gate = (grad * up) * (sigmoid * (1 + gate * (1 - sigmoid))).
+ AscendC::Muls(tmp, sigmoid, -1.0f, count);
+ AscendC::PipeBarrier();
+ AscendC::Adds(tmp, tmp, 1.0f, count);
+ AscendC::PipeBarrier();
+ AscendC::Mul(tmp, gate, tmp, count);
+ AscendC::PipeBarrier();
+ AscendC::Adds(tmp, tmp, 1.0f, count);
+ AscendC::PipeBarrier();
+ AscendC::Mul(tmp, sigmoid, tmp, count);
+ AscendC::Mul(up, grad, up, count);
+ AscendC::PipeBarrier();
+ AscendC::Mul(tmp, up, tmp, count);
+ } else {
+ AscendC::Mul(tmp, tmp, up, count);
+ }
+ Store(output, tmp, count);
+ outQueue_.EnQue(output);
+ inQueue_.FreeTensor(input);
+ }
+
+ __aicore__ inline void CopyOut(int64_t offset, uint32_t count)
+ {
+ auto output = outQueue_.DeQue(); // vector -> MTE3 synchronization
+ AscendC::DataCopyExtParams params{
+ 1, static_cast(count * sizeof(T)), 0, 0, 0};
+ AscendC::DataCopyPad(outGm_[offset], output, params);
+ if constexpr (Backward) {
+ AscendC::DataCopyPad(dUpGm_[offset], output[TILE_LENGTH], params);
+ }
+ outQueue_.FreeTensor(output);
+ }
+
+ AscendC::GlobalTensor gateGm_, upGm_, gradGm_, outGm_, dUpGm_;
+ AscendC::TQue inQueue_;
+ AscendC::TQue outQueue_;
+ AscendC::TBuf work_;
+ int64_t n_;
+};
+
+template
+class KernelSiLU {
+public:
+ __aicore__ inline void Init(AscendC::TPipe* pipe, GM_ADDR x, GM_ADDR grad,
+ GM_ADDR out, int64_t n)
+ {
+ n_ = n;
+ xGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x));
+ outGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(out));
+ if constexpr (Backward) {
+ gradGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(grad));
+ }
+ // Worst-case UB use: (2 inputs + 1 output) * 4 KiB + 4 FP32 tiles * 8 KiB.
+ pipe->InitBuffer(inQueue_, 1, (Backward ? 2 : 1) * TILE_LENGTH * sizeof(T));
+ pipe->InitBuffer(outQueue_, 1, TILE_LENGTH * sizeof(T));
+ pipe->InitBuffer(work_, 4 * TILE_LENGTH * sizeof(float));
+ }
+
+ __aicore__ inline void Process()
+ {
+ // Identical strided tiling to SwiGLU: element i is always evaluated by
+ // the same expression regardless of n_ or the launched block count.
+ for (int64_t offset = static_cast(AscendC::GetBlockIdx()) * TILE_LENGTH;
+ offset < n_;
+ offset += static_cast(AscendC::GetBlockNum()) * TILE_LENGTH) {
+ const uint32_t count = static_cast(
+ n_ - offset < TILE_LENGTH ? n_ - offset : TILE_LENGTH);
+ CopyIn(offset, count);
+ Compute(count);
+ CopyOut(offset, count);
+ }
+ }
+
+private:
+ __aicore__ inline void CopyIn(int64_t offset, uint32_t count)
+ {
+ auto input = inQueue_.AllocTensor();
+ AscendC::DataCopyExtParams params{
+ 1, static_cast(count * sizeof(T)), 0, 0, 0};
+ AscendC::DataCopyPadExtParams pad{false, 0, 0, 0};
+ AscendC::DataCopyPad(input, xGm_[offset], params, pad);
+ if constexpr (Backward) {
+ AscendC::DataCopyPad(input[TILE_LENGTH], gradGm_[offset], params, pad);
+ }
+ inQueue_.EnQue(input);
+ }
+
+ __aicore__ inline void ToFloat(AscendC::LocalTensor dst,
+ AscendC::LocalTensor src, uint32_t count)
+ {
+ if constexpr (std::is_same_v) {
+ // UB-to-UB copies require a multiple of 32 bytes. Padding stays in UB.
+ AscendC::DataCopy(dst, src, (count + 7) / 8 * 8);
+ } else {
+ AscendC::Cast(dst, src, AscendC::RoundMode::CAST_NONE, count);
+ }
+ }
+
+ __aicore__ inline void Store(AscendC::LocalTensor dst,
+ AscendC::LocalTensor src, uint32_t count)
+ {
+ AscendC::PipeBarrier();
+ if constexpr (std::is_same_v) {
+ AscendC::DataCopy(dst, src, (count + 7) / 8 * 8);
+ } else {
+ // Round to nearest, ties to even, matching PyTorch dtype conversion.
+ AscendC::Cast(dst, src, AscendC::RoundMode::CAST_RINT, count);
+ }
+ AscendC::PipeBarrier();
+ }
+
+ __aicore__ inline void Compute(uint32_t count)
+ {
+ auto input = inQueue_.DeQue(); // MTE2 -> vector synchronization
+ auto output = outQueue_.AllocTensor();
+ auto x = work_.Get();
+ auto grad = x[TILE_LENGTH];
+ auto sigmoid = x[2 * TILE_LENGTH];
+ auto tmp = x[3 * TILE_LENGTH];
+ ToFloat(x, input, count);
+ if constexpr (Backward) {
+ ToFloat(grad, input[TILE_LENGTH], count);
+ }
+ AscendC::PipeBarrier();
+
+ // sigmoid(x) = 1 / (1 + exp(-x)), the same sequence SwiGLU uses on gate,
+ // so silu(x) is bitwise equal to swiglu(x, ones) on this hardware.
+ AscendC::Muls(sigmoid, x, -1.0f, count);
+ AscendC::PipeBarrier();
+ AscendC::Exp(sigmoid, sigmoid, count);
+ AscendC::PipeBarrier();
+ AscendC::Adds(sigmoid, sigmoid, 1.0f, count);
+ AscendC::Duplicate(tmp, 1.0f, count);
+ AscendC::PipeBarrier();
+ AscendC::Div(sigmoid, tmp, sigmoid, count);
+ AscendC::PipeBarrier();
+
+ if constexpr (Backward) {
+ // dx = grad * (sigmoid * (1 + x * (1 - sigmoid))).
+ AscendC::Muls(tmp, sigmoid, -1.0f, count);
+ AscendC::PipeBarrier();
+ AscendC::Adds(tmp, tmp, 1.0f, count);
+ AscendC::PipeBarrier();
+ AscendC::Mul(tmp, x, tmp, count);
+ AscendC::PipeBarrier();
+ AscendC::Adds(tmp, tmp, 1.0f, count);
+ AscendC::PipeBarrier();
+ AscendC::Mul(tmp, sigmoid, tmp, count);
+ AscendC::PipeBarrier();
+ AscendC::Mul(tmp, grad, tmp, count);
+ } else {
+ AscendC::Mul(tmp, x, sigmoid, count);
+ }
+ Store(output, tmp, count);
+ outQueue_.EnQue(output);
+ inQueue_.FreeTensor(input);
+ }
+
+ __aicore__ inline void CopyOut(int64_t offset, uint32_t count)
+ {
+ auto output = outQueue_.DeQue(); // vector -> MTE3 synchronization
+ AscendC::DataCopyExtParams params{
+ 1, static_cast(count * sizeof(T)), 0, 0, 0};
+ AscendC::DataCopyPad(outGm_[offset], output, params);
+ outQueue_.FreeTensor(output);
+ }
+
+ AscendC::GlobalTensor xGm_, gradGm_, outGm_;
+ AscendC::TQue inQueue_;
+ AscendC::TQue outQueue_;
+ AscendC::TBuf work_;
+ int64_t n_;
+};
+
+template
+__global__ __vector__ void silu_ascend_kernel(
+ GM_ADDR x, GM_ADDR grad, GM_ADDR out, int64_t n)
+{
+ AscendC::TPipe pipe;
+ KernelSiLU op;
+ op.Init(&pipe, x, grad, out, n);
+ op.Process();
+}
+
+template
+__global__ __vector__ void swiglu_ascend_kernel(
+ GM_ADDR gate, GM_ADDR up, GM_ADDR grad, GM_ADDR out, GM_ADDR dUp, int64_t n)
+{
+ AscendC::TPipe pipe;
+ KernelSwiGLU op;
+ op.Init(&pipe, gate, up, grad, out, dUp, n);
+ op.Process();
+}
+
+void CheckInput(const torch::Tensor& tensor, const char* name)
+{
+ TORCH_CHECK(tensor.is_privateuseone(), name, " must be on an NPU device");
+ TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+ TORCH_CHECK(tensor.scalar_type() == at::kHalf || tensor.scalar_type() == at::kBFloat16 ||
+ tensor.scalar_type() == at::kFloat, name, " must be fp16, bf16, or fp32");
+}
+
+void CheckLike(const torch::Tensor& tensor, const torch::Tensor& gate, const char* name)
+{
+ CheckInput(tensor, name);
+ TORCH_CHECK(tensor.device() == gate.device(), name, " must be on the same NPU device as gate");
+ TORCH_CHECK(tensor.sizes() == gate.sizes(), name, " must share shape with gate");
+ TORCH_CHECK(tensor.scalar_type() == gate.scalar_type(), name, " must share dtype with gate");
+}
+
+template
+void Launch(torch::Tensor gate, torch::Tensor up, torch::Tensor grad,
+ torch::Tensor out, torch::Tensor dUp)
+{
+ const int64_t n = gate.numel();
+ if (n == 0) {
+ return;
+ }
+ const uint32_t blocks = static_cast(
+ std::min((n + TILE_LENGTH - 1) / TILE_LENGTH, MAX_BLOCKS));
+ // Flush torch_npu's task queue before launching directly on its current stream.
+ auto stream = c10_npu::getCurrentNPUStream().stream(true);
+ auto gatePtr = reinterpret_cast(gate.mutable_data_ptr());
+ auto upPtr = reinterpret_cast(up.mutable_data_ptr());
+ auto outPtr = reinterpret_cast(out.mutable_data_ptr());
+ uint8_t* gradPtr = nullptr;
+ uint8_t* dUpPtr = nullptr;
+ if constexpr (Backward) {
+ gradPtr = reinterpret_cast(grad.mutable_data_ptr());
+ dUpPtr = reinterpret_cast(dUp.mutable_data_ptr());
+ }
+ if (gate.scalar_type() == at::kHalf) {
+ swiglu_ascend_kernel<<>>(
+ gatePtr, upPtr, gradPtr, outPtr, dUpPtr, n);
+ } else if (gate.scalar_type() == at::kBFloat16) {
+ swiglu_ascend_kernel<<>>(
+ gatePtr, upPtr, gradPtr, outPtr, dUpPtr, n);
+ } else {
+ swiglu_ascend_kernel<<>>(
+ gatePtr, upPtr, gradPtr, outPtr, dUpPtr, n);
+ }
+}
+
+template
+void LaunchSiLU(torch::Tensor x, torch::Tensor grad, torch::Tensor out)
+{
+ const int64_t n = x.numel();
+ if (n == 0) {
+ return;
+ }
+ const uint32_t blocks = static_cast(
+ std::min((n + TILE_LENGTH - 1) / TILE_LENGTH, MAX_BLOCKS));
+ // Flush torch_npu's task queue before launching directly on its current stream.
+ auto stream = c10_npu::getCurrentNPUStream().stream(true);
+ auto xPtr = reinterpret_cast(x.mutable_data_ptr());
+ auto outPtr = reinterpret_cast(out.mutable_data_ptr());
+ uint8_t* gradPtr = nullptr;
+ if constexpr (Backward) {
+ gradPtr = reinterpret_cast(grad.mutable_data_ptr());
+ }
+ if (x.scalar_type() == at::kHalf) {
+ silu_ascend_kernel<<>>(
+ xPtr, gradPtr, outPtr, n);
+ } else if (x.scalar_type() == at::kBFloat16) {
+ silu_ascend_kernel<<>>(
+ xPtr, gradPtr, outPtr, n);
+ } else {
+ silu_ascend_kernel<<>>(
+ xPtr, gradPtr, outPtr, n);
+ }
+}
+
+} // namespace
+
+torch::Tensor swiglu_ascend_forward(torch::Tensor gate, torch::Tensor up)
+{
+ CheckInput(gate, "gate");
+ CheckLike(up, gate, "up");
+ const c10::DeviceGuard guard(gate.device());
+ auto out = at::empty(gate.sizes(), gate.options());
+ Launch(gate, up, {}, out, {});
+ return out;
+}
+
+std::vector swiglu_ascend_backward(
+ torch::Tensor grad, torch::Tensor gate, torch::Tensor up)
+{
+ CheckInput(gate, "gate");
+ CheckLike(up, gate, "up");
+ CheckLike(grad, gate, "grad_out");
+ const c10::DeviceGuard guard(gate.device());
+ auto dGate = at::empty(gate.sizes(), gate.options());
+ auto dUp = at::empty(gate.sizes(), gate.options());
+ Launch(gate, up, grad, dGate, dUp);
+ return {dGate, dUp};
+}
+
+torch::Tensor silu_ascend_forward(torch::Tensor x)
+{
+ CheckInput(x, "x");
+ const c10::DeviceGuard guard(x.device());
+ auto out = at::empty(x.sizes(), x.options());
+ LaunchSiLU(x, {}, out);
+ return out;
+}
+
+torch::Tensor silu_ascend_backward(torch::Tensor grad, torch::Tensor x)
+{
+ CheckInput(x, "x");
+ CheckLike(grad, x, "grad_out");
+ const c10::DeviceGuard guard(x.device());
+ auto dX = at::empty(x.sizes(), x.options());
+ LaunchSiLU(x, grad, dX);
+ return dX;
+}
diff --git a/csrc/ascend/attention/deterministic_attention_ascend.asc b/csrc/ascend/attention/deterministic_attention_ascend.asc
new file mode 100644
index 00000000..4d48e93f
--- /dev/null
+++ b/csrc/ascend/attention/deterministic_attention_ascend.asc
@@ -0,0 +1,600 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+//
+// Deterministic (batch-invariant) standard-softmax attention, Ascend C (CANN)
+// forward kernel.
+//
+// out = softmax(Q K^T * scale + masks) @ V, lse = rowmax + log(sum(exp(s - rowmax)))
+//
+// Mirrors the batch-invariant algorithm of the Triton reference
+// (rl_engine/kernels/ops/triton/attention/standard_attn.py, i.e.
+// rlkernel.attention.deterministic_core.v1) and the CUDA deterministic op
+// (issue #147):
+// - layout : q [B, Hq, Sq, D], k/v [B, Hkv, Skv, D] contiguous, D = 128
+// - masks : causal (upper triangle at offset Skv - Sq + 1) and optional
+// key_padding_mask [B, Skv] bool, True = keep
+// - numerics: all fp32 intermediate; bf16/fp16 inputs are upcast, the output
+// row is cast back once at the end
+//
+// Batch-invariance / no split-K: every (b, q_head, row) is processed
+// end-to-end by exactly one AI-core block, with a fixed 64-key tile size and a
+// fixed two-pass (max, then sum-exp + P.V) reduction order over the key
+// dimension. The instruction sequence for a row depends only on Skv, D and the
+// masks -- never on the batch size, the block the row lands on, or how many
+// blocks were launched (rows are strided across blocks). No second-pass merge
+// of per-split (m, l, u) summaries exists, so the reduction tree is fixed.
+//
+// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by
+// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu.
+
+#include
+#include
+#include
+
+#include "kernel_operator.h"
+
+#include
+
+#include "torch_npu/csrc/core/npu/NPUStream.h"
+
+namespace {
+
+// Fixed head dimension (matches the CUDA deterministic op gate).
+constexpr uint32_t HEAD_DIM = 128;
+// Keys per tile. Fixed for all rows and batch sizes; this is what makes the
+// reduction order batch-invariant (mirrors the Triton reference _BLOCK_N = 64).
+constexpr uint32_t TILE_N = 64;
+// Cap on launched blocks. Work items are strided across blocks, so launching
+// fewer blocks than items is fine and never changes per-row numerics.
+constexpr int64_t MAX_BLOCKS = 512;
+constexpr float NEG_INF = -3.402823466e+38f; // -FLT_MAX, mask sentinel
+// lse of a fully-masked row: true -inf (matches the Triton reference, whose
+// max_score == -inf / log(denom=0) == -inf path writes -inf, not -FLT_MAX).
+constexpr float LSE_INVALID = -std::numeric_limits::infinity();
+
+template
+class KernelDeterministicAttention {
+public:
+ __aicore__ inline KernelDeterministicAttention(AscendC::TPipe* pipe) : pipe_(pipe) {}
+
+ __aicore__ inline void Init(GM_ADDR q,
+ GM_ADDR k,
+ GM_ADDR v,
+ GM_ADDR mask,
+ GM_ADDR out,
+ GM_ADDR lse,
+ int64_t B,
+ int64_t Hq,
+ int64_t Hkv,
+ int64_t Sq,
+ int64_t Skv,
+ float scale,
+ int32_t causal,
+ int32_t hasMask,
+ int32_t outFp32)
+ {
+ B_ = B;
+ Hq_ = Hq;
+ Hkv_ = Hkv;
+ Sq_ = Sq;
+ Skv_ = Skv;
+ scale_ = scale;
+ causal_ = causal;
+ hasMask_ = hasMask;
+ outFp32_ = outFp32 != 0;
+ qGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(q));
+ kGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(k));
+ vGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(v));
+ maskGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(mask));
+ outGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(out));
+ outGmF32_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(out));
+ lseGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(lse));
+
+ // UB budget stays well under 192 KB:
+ // k tile bf16 16 KB + k tile fp32 32 KB + v tile fp32 32 KB
+ // + q/acc/prod/work/scores/scalar/mask ~4 KB.
+ pipe_->InitBuffer(qBufF_, HEAD_DIM * sizeof(float));
+ pipe_->InitBuffer(accBufF_, HEAD_DIM * sizeof(float));
+ pipe_->InitBuffer(prodBufF_, HEAD_DIM * sizeof(float));
+ pipe_->InitBuffer(workBufF_, HEAD_DIM * sizeof(float));
+ pipe_->InitBuffer(kBufT_, TILE_N * HEAD_DIM * sizeof(T));
+ pipe_->InitBuffer(kBufF_, TILE_N * HEAD_DIM * sizeof(float));
+ pipe_->InitBuffer(vBufF_, TILE_N * HEAD_DIM * sizeof(float));
+ pipe_->InitBuffer(scoresBuf_, TILE_N * sizeof(float));
+ // 128 B: mask tile (64 B) + up to 31 B misalignment + 32 B rounding.
+ pipe_->InitBuffer(maskBuf_, 128);
+ // 64 B: floats [0,8) for reduce/log scratch, floats [8,16) for the
+ // output staging slots (both halves 32-byte aligned for DataCopyPad).
+ pipe_->InitBuffer(scalarBuf_, 64);
+
+ // Intra-core pipeline events. NOTE: AscendC::SyncAll() is a cross-core
+ // barrier and deadlocks when more blocks are launched than there are
+ // physical cores (resident blocks wait for unscheduled ones), so all
+ // synchronization here uses per-pipe SetFlag/WaitFlag instead.
+ eventVS_ = pipe_->FetchEventID(AscendC::HardEvent::V_S);
+ eventSV_ = pipe_->FetchEventID(AscendC::HardEvent::S_V);
+ eventMTE2S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE2_S);
+ eventSMTE2_ = pipe_->FetchEventID(AscendC::HardEvent::S_MTE2);
+ eventVMTE2_ = pipe_->FetchEventID(AscendC::HardEvent::V_MTE2);
+ eventSMTE3_ = pipe_->FetchEventID(AscendC::HardEvent::S_MTE3);
+ eventVMTE3_ = pipe_->FetchEventID(AscendC::HardEvent::V_MTE3);
+ eventMTE3S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE3_S);
+ }
+
+ __aicore__ inline void Process()
+ {
+ const int64_t items = B_ * Hq_ * Sq_;
+ for (int64_t item = AscendC::GetBlockIdx(); item < items;
+ item += AscendC::GetBlockNum()) {
+ const int64_t row = item % Sq_;
+ const int64_t qh = (item / Sq_) % Hq_;
+ const int64_t b = item / (Sq_ * Hq_);
+ ProcessRow(b, qh, row);
+ }
+ }
+
+private:
+ __aicore__ inline void LoadQRow(int64_t b, int64_t qh, int64_t row)
+ {
+ const int64_t offset = ((b * Hq_ + qh) * Sq_ + row) * HEAD_DIM;
+ AscendC::LocalTensor qT = kBufT_.Get();
+ AscendC::DataCopyExtParams cp{1, static_cast(HEAD_DIM * sizeof(T)), 0, 0, 0};
+ AscendC::DataCopyPadExtParams pp{false, 0, 0, 0};
+ AscendC::DataCopyPad(qT, qGm_[offset], cp, pp);
+ AscendC::SetFlag(eventMTE2S_);
+ AscendC::WaitFlag(eventMTE2S_);
+ AscendC::Cast(qBufF_.Get(), qT, AscendC::RoundMode::CAST_NONE, HEAD_DIM);
+ }
+
+ __aicore__ inline void LoadKTile(int64_t b, int64_t kvh, int64_t start, uint32_t count)
+ {
+ // kBufT_ staging may hold data the vector pipe is still reading (the
+ // q cast of the first tile or the v cast of the previous tile).
+ AscendC::SetFlag(eventVMTE2_);
+ AscendC::WaitFlag(eventVMTE2_);
+ const int64_t offset = ((b * Hkv_ + kvh) * Skv_ + start) * HEAD_DIM;
+ AscendC::LocalTensor kT = kBufT_.Get();
+ AscendC::DataCopyExtParams cp{1, static_cast(count * HEAD_DIM * sizeof(T)), 0, 0, 0};
+ AscendC::DataCopyPadExtParams pp{false, 0, 0, 0};
+ AscendC::DataCopyPad(kT, kGm_[offset], cp, pp);
+ AscendC::SetFlag(eventMTE2S_);
+ AscendC::WaitFlag(eventMTE2S_);
+ AscendC::Cast(kBufF_.Get(), kT, AscendC::RoundMode::CAST_NONE,
+ count * HEAD_DIM);
+ }
+
+ __aicore__ inline void LoadVTile(int64_t b, int64_t kvh, int64_t start, uint32_t count)
+ {
+ // vBufF_ is still being read by the P . V accumulation of the previous
+ // tile; kBufT_ staging may be in use by the vector pipe as well.
+ AscendC::SetFlag(eventVMTE2_);
+ AscendC::WaitFlag(eventVMTE2_);
+ const int64_t offset = ((b * Hkv_ + kvh) * Skv_ + start) * HEAD_DIM;
+ AscendC::LocalTensor vT = kBufT_.Get();
+ AscendC::DataCopyExtParams cp{1, static_cast(count * HEAD_DIM * sizeof(T)), 0, 0, 0};
+ AscendC::DataCopyPadExtParams pp{false, 0, 0, 0};
+ AscendC::DataCopyPad(vT, vGm_[offset], cp, pp);
+ AscendC::SetFlag(eventMTE2S_);
+ AscendC::WaitFlag(eventMTE2S_);
+ AscendC::Cast(vBufF_.Get(), vT, AscendC::RoundMode::CAST_NONE,
+ count * HEAD_DIM);
+ }
+
+ // Load the mask window covering keys [start, start + count) of batch b.
+ // The window is read from the last 32-byte-aligned GM address at or below
+ // the tile start; returns the byte offset of the tile inside the window.
+ __aicore__ inline uint32_t LoadMaskTile(int64_t b, int64_t start, uint32_t count)
+ {
+ // maskBuf_ holds values the scalar pipe read for the previous tile;
+ // wait for those reads before MTE2 overwrites the window.
+ AscendC::SetFlag(eventSMTE2_);
+ AscendC::WaitFlag(eventSMTE2_);
+ const int64_t base = b * Skv_ + start;
+ const int64_t aligned = base & ~31LL;
+ const uint32_t offset = static_cast(base - aligned);
+ const int64_t remaining = (b + 1) * Skv_ - aligned;
+ uint32_t alignedCount = (offset + count + 31) & ~31u;
+ if (alignedCount > remaining) {
+ alignedCount = static_cast(remaining);
+ }
+ AscendC::LocalTensor m = maskBuf_.Get();
+ AscendC::DataCopyExtParams cp{1, alignedCount, 0, 0, 0};
+ AscendC::DataCopyPadExtParams pp{false, 0, 0, 0};
+ AscendC::DataCopyPad(m, maskGm_[aligned], cp, pp);
+ AscendC::SetFlag(eventMTE2S_);
+ AscendC::WaitFlag(eventMTE2S_);
+ return offset;
+ }
+
+ // scores[j] = scale * (q . k[start + j]) for j in [0, count), with the
+ // causal and key-padding masks applied; masked lanes become NEG_INF.
+ //
+ // Each dot product lands directly in its score lane (ReduceSum's dst
+ // scalar aliases scores[j]), so the vector pipe runs the whole tile
+ // without per-key scalar synchronization; a single V_S wait covers all
+ // lanes before the scalar mask pass.
+ __aicore__ inline void ComputeScores(int64_t b,
+ int64_t row,
+ int64_t start,
+ uint32_t count,
+ uint32_t maskOffset)
+ {
+ AscendC::LocalTensor qRow = qBufF_.Get();
+ AscendC::LocalTensor kTile = kBufF_.Get();
+ AscendC::LocalTensor scores = scoresBuf_.Get();
+ AscendC::LocalTensor prod = prodBufF_.Get();
+ AscendC::LocalTensor maskTile = maskBuf_.Get();
+
+ // The previous tile's mask pass (and the caller's padding loop) write
+ // the score lanes on the scalar pipe; drain them before the vector
+ // ReduceSum targets the same lanes.
+ AscendC::SetFlag(eventSV_);
+ AscendC::WaitFlag(eventSV_);
+ for (uint32_t j = 0; j < count; ++j) {
+ AscendC::Mul(prod, qRow, kTile[j * HEAD_DIM], HEAD_DIM);
+ AscendC::ReduceSum(scores[j], prod, workBufF_.Get(), HEAD_DIM);
+ }
+ WaitVector(); // all dot products visible to the scalar pipe
+ AscendC::Muls(scores, scores, scale_, count);
+ WaitVector(); // scaled lanes visible to the scalar pipe
+
+ const int64_t causalKeep = row + Skv_ - Sq_;
+ for (uint32_t j = 0; j < count; ++j) {
+ float s = scores.GetValue(j);
+ const int64_t jGlobal = start + j;
+ if (causal_ && jGlobal > causalKeep) {
+ s = NEG_INF;
+ }
+ if (hasMask_ && maskTile.GetValue(maskOffset + j) == 0) {
+ s = NEG_INF;
+ }
+ scores.SetValue(j, s);
+ }
+ }
+
+ __aicore__ inline void ProcessRow(int64_t b, int64_t qh, int64_t row)
+ {
+ const int64_t kvh = qh / (Hq_ / Hkv_); // GQA: query head h -> KV head h / g
+ AscendC::LocalTensor scores = scoresBuf_.Get();
+ AscendC::LocalTensor scalar = scalarBuf_.Get();
+
+ LoadQRow(b, qh, row);
+
+ // Left padding shifts the physical positions of the valid keys, and
+ // the softmax denominator (sumExp) is reduced per physical 64-key
+ // tile before the per-tile results are summed -- so the FP32
+ // addition grouping of the valid keys depends on where the padding
+ // sits. Anchor the tiles to the first valid key (keyBegin) so the
+ // valid keys always start at the first position of the first tile
+ // and the denominator (hence the whole output and LSE) is
+ // padding-invariant. The numerator's per-key accumulation already
+ // runs across tiles in a fixed order, and causalKeep keeps its
+ // physical coordinate.
+ int64_t keyBegin = 0;
+ if (hasMask_) {
+ keyBegin = Skv_;
+ for (int64_t base = 0; base < Skv_; base += TILE_N) {
+ const uint32_t winCount =
+ static_cast(base + TILE_N <= Skv_ ? TILE_N : Skv_ - base);
+ const uint32_t offset = LoadMaskTile(b, base, winCount);
+ AscendC::LocalTensor m = maskBuf_.Get();
+ for (uint32_t j = 0; j < winCount; ++j) {
+ if (m.GetValue(offset + j) != 0) {
+ keyBegin = base + j;
+ break;
+ }
+ }
+ if (keyBegin < Skv_) {
+ break;
+ }
+ }
+ }
+ const int64_t tileCount = (Skv_ - keyBegin + TILE_N - 1) / TILE_N;
+
+ // Pass 1: row max with a fixed tile order.
+ float rowMax = NEG_INF;
+ bool anyValid = false;
+ for (int64_t tile = 0; tile < tileCount; ++tile) {
+ const int64_t start = keyBegin + tile * TILE_N;
+ const uint32_t count = TileCount(start);
+ LoadKTile(b, kvh, start, count);
+ uint32_t maskOffset = 0;
+ if (hasMask_) {
+ maskOffset = LoadMaskTile(b, start, count);
+ }
+ ComputeScores(b, row, start, count, maskOffset);
+ for (uint32_t j = count; j < TILE_N; ++j) {
+ scores.SetValue(j, NEG_INF);
+ }
+ AscendC::SetFlag(eventSV_);
+ AscendC::WaitFlag(eventSV_);
+ AscendC::ReduceMax(scalar, scores, workBufF_.Get(), TILE_N, false);
+ WaitVector();
+ const float tileMax = scalar.GetValue(0);
+ if (tileMax > NEG_INF) {
+ anyValid = true;
+ }
+ rowMax = tileMax > rowMax ? tileMax : rowMax;
+ }
+
+ // Pass 2: sum(exp(s - rowMax)) and P . V with the same fixed tile
+ // order (see the batch-invariance note at the top of the file).
+ float sumExp = 0.0f;
+ AscendC::LocalTensor acc = accBufF_.Get();
+ // Real zeroing: accBuf_ starts as uninitialized UB and 0 * inf == NaN.
+ AscendC::Duplicate(acc, 0.0f, HEAD_DIM);
+ if (!anyValid) {
+ // Fully-masked row: exp(s - rowMax) would be exp(0) = 1 for the
+ // masked lanes (NEG_INF - NEG_INF == 0), not exp(-inf) = 0, so the
+ // row is defined as out = 0, lse = -inf -- mirroring the Triton
+ // reference's max_score == -inf / denom > 0 guards.
+ WriteOutputs(b, qh, row, NEG_INF, 0.0f, acc);
+ return;
+ }
+ for (int64_t tile = 0; tile < tileCount; ++tile) {
+ const int64_t start = keyBegin + tile * TILE_N;
+ const uint32_t count = TileCount(start);
+ LoadKTile(b, kvh, start, count);
+ uint32_t maskOffset = 0;
+ if (hasMask_) {
+ maskOffset = LoadMaskTile(b, start, count);
+ }
+ ComputeScores(b, row, start, count, maskOffset);
+ for (uint32_t j = count; j < TILE_N; ++j) {
+ scores.SetValue(j, NEG_INF);
+ }
+ AscendC::SetFlag(eventSV_);
+ AscendC::WaitFlag(eventSV_);
+ AscendC::Adds(scores, scores, -rowMax, TILE_N);
+ AscendC::Exp(scores, scores, TILE_N);
+ WaitVector(); // vector -> scalar visibility; covers the Exp above
+ // Guarantee the masked lanes contribute exactly zero. The vector
+ // Exp of the -FLT_MAX sentinel is not required to flush to 0.0f,
+ // and any residue would shift sumExp (hence the final invDenom)
+ // by an ULP that depends on where the padding sits in the
+ // physical layout -- breaking padding invariance at the FP32
+ // level even though the BF16 outputs round identically.
+ if (causal_ || hasMask_) {
+ const int64_t causalKeep = row + Skv_ - Sq_;
+ AscendC::LocalTensor maskT = maskBuf_.Get();
+ for (uint32_t j = 0; j < count; ++j) {
+ const int64_t jGlobal = start + j;
+ const bool masked = (causal_ && jGlobal > causalKeep) ||
+ (hasMask_ && maskT.GetValue(maskOffset + j) == 0);
+ if (masked) {
+ scores.SetValue(j, 0.0f);
+ }
+ }
+ AscendC::SetFlag(eventSV_);
+ AscendC::WaitFlag(eventSV_);
+ }
+ AscendC::ReduceSum(scalar, scores, workBufF_.Get(), TILE_N);
+ WaitVector(); // vector -> scalar read
+ sumExp += scalar.GetValue(0);
+
+ LoadVTile(b, kvh, start, count);
+ AscendC::LocalTensor vTile = vBufF_.Get();
+ AscendC::LocalTensor prod = prodBufF_.Get();
+ for (uint32_t j = 0; j < count; ++j) {
+ const float pj = scores.GetValue(j);
+ if (pj == 0.0f) {
+ continue;
+ }
+ AscendC::Muls(prod, vTile[j * HEAD_DIM], pj, HEAD_DIM);
+ AscendC::Add(acc, acc, prod, HEAD_DIM);
+ }
+ }
+
+ // out = acc / sumExp. A fully-masked row keeps rowMax == NEG_INF, so
+ // sumExp is 0 and the output row is defined as 0 with lse = -inf
+ // (mirrors the Triton reference's denom > 0 guard).
+ const float invDenom = (sumExp > 0.0f) ? (1.0f / sumExp) : 0.0f;
+ AscendC::Muls(acc, acc, invDenom, HEAD_DIM);
+ WriteOutputs(b, qh, row, rowMax, sumExp, acc);
+ }
+
+ __aicore__ inline void WriteOutputs(int64_t b,
+ int64_t qh,
+ int64_t row,
+ float rowMax,
+ float sumExp,
+ AscendC::LocalTensor acc)
+ {
+ AscendC::LocalTensor scalar = scalarBuf_.Get();
+ float lse = LSE_INVALID;
+ if (rowMax > NEG_INF) {
+ scalar.SetValue(0, sumExp);
+ AscendC::SetFlag(eventSV_); // scalar write -> vector op
+ AscendC::WaitFlag(eventSV_);
+ AscendC::Log(scalar, scalar, 8);
+ WaitVector(); // vector -> scalar read
+ lse = rowMax + scalar.GetValue(0);
+ }
+ // Stage outputs in UB, then DataCopyPad to GM. GlobalTensor.SetValue
+ // is unreliable on hardware (cannbot ascendc-precision-debug
+ // common-traps), so scalar GM stores are not used.
+ scalar.SetValue(0, lse);
+ AscendC::LocalTensor outT = kBufT_.Get();
+ AscendC::LocalTensor outF = kBufF_.Get();
+ AscendC::SetFlag(eventVMTE3_); // vector write -> copy-out
+ AscendC::WaitFlag(eventVMTE3_);
+ AscendC::SetFlag(eventSMTE3_); // scalar write -> copy-out
+ AscendC::WaitFlag(eventSMTE3_);
+ AscendC::DataCopyExtParams p4{1, sizeof(float), 0, 0, 0};
+ AscendC::DataCopyPad(lseGm_[(b * Hq_ + qh) * Sq_ + row], scalar[0], p4);
+ if (outFp32_) {
+ // FP32 output: stage the exact FP32 accumulator (Muls by 1.0 is
+ // an exact same-type copy on A2; Cast with CAST_NONE would emit
+ // nothing). kBufF_ is dead after the final P . V accumulation.
+ AscendC::Muls(outF, acc, 1.0f, HEAD_DIM);
+ AscendC::SetFlag(eventVMTE3_);
+ AscendC::WaitFlag(eventVMTE3_);
+ AscendC::DataCopyExtParams outCp{1, static_cast(HEAD_DIM * sizeof(float)),
+ 0, 0, 0};
+ AscendC::DataCopyPad(outGmF32_[((b * Hq_ + qh) * Sq_ + row) * HEAD_DIM], outF, outCp);
+ } else {
+ AscendC::Cast(outT, acc, AscendC::RoundMode::CAST_ROUND, HEAD_DIM);
+ AscendC::SetFlag(eventVMTE3_);
+ AscendC::WaitFlag(eventVMTE3_);
+ AscendC::DataCopyExtParams outCp{1, static_cast(HEAD_DIM * sizeof(T)), 0, 0, 0};
+ AscendC::DataCopyPad(outGm_[((b * Hq_ + qh) * Sq_ + row) * HEAD_DIM], outT, outCp);
+ }
+ // Drain MTE3 before the next row stages new values into the shared
+ // buffers; the scalar pipe issues all later MTE2 copies in order, so
+ // this wait alone orders them after the copy-outs.
+ AscendC::SetFlag(eventMTE3S_);
+ AscendC::WaitFlag(eventMTE3S_);
+ }
+
+ // Wait until all outstanding vector-pipe results are readable as scalars.
+ __aicore__ inline void WaitVector()
+ {
+ AscendC::SetFlag(eventVS_);
+ AscendC::WaitFlag(eventVS_);
+ }
+
+ __aicore__ inline uint32_t TileCount(int64_t start) const
+ {
+ const int64_t remaining = Skv_ - start;
+ return static_cast(remaining < TILE_N ? remaining : TILE_N);
+ }
+
+ AscendC::TPipe* pipe_;
+ AscendC::GlobalTensor qGm_;
+ AscendC::GlobalTensor kGm_;
+ AscendC::GlobalTensor vGm_;
+ AscendC::GlobalTensor maskGm_;
+ AscendC::GlobalTensor outGm_;
+ AscendC::GlobalTensor outGmF32_;
+ AscendC::GlobalTensor lseGm_;
+ AscendC::TBuf qBufF_;
+ AscendC::TBuf accBufF_;
+ AscendC::TBuf prodBufF_;
+ AscendC::TBuf workBufF_;
+ AscendC::TBuf kBufT_;
+ AscendC::TBuf kBufF_;
+ AscendC::TBuf vBufF_;
+ AscendC::TBuf scoresBuf_;
+ AscendC::TBuf maskBuf_;
+ AscendC::TBuf scalarBuf_;
+ AscendC::TEventID eventVS_;
+ AscendC::TEventID eventSV_;
+ AscendC::TEventID eventMTE2S_;
+ AscendC::TEventID eventSMTE2_;
+ AscendC::TEventID eventVMTE2_;
+ AscendC::TEventID eventSMTE3_;
+ AscendC::TEventID eventVMTE3_;
+ AscendC::TEventID eventMTE3S_;
+ int64_t B_;
+ int64_t Hq_;
+ int64_t Hkv_;
+ int64_t Sq_;
+ int64_t Skv_;
+ float scale_;
+ int32_t causal_;
+ int32_t hasMask_;
+ int32_t outFp32_;
+};
+
+} // namespace
+
+extern "C" __global__ __vector__ void deterministic_attention_ascend_kernel_bf16(
+ GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR mask, GM_ADDR out, GM_ADDR lse,
+ int64_t B, int64_t Hq, int64_t Hkv, int64_t Sq, int64_t Skv,
+ float scale, int32_t causal, int32_t hasMask, int32_t outFp32)
+{
+ AscendC::TPipe pipe;
+ KernelDeterministicAttention op(&pipe);
+ op.Init(q, k, v, mask, out, lse, B, Hq, Hkv, Sq, Skv, scale, causal, hasMask, outFp32);
+ op.Process();
+}
+
+extern "C" __global__ __vector__ void deterministic_attention_ascend_kernel_fp16(
+ GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR mask, GM_ADDR out, GM_ADDR lse,
+ int64_t B, int64_t Hq, int64_t Hkv, int64_t Sq, int64_t Skv,
+ float scale, int32_t causal, int32_t hasMask, int32_t outFp32)
+{
+ AscendC::TPipe pipe;
+ KernelDeterministicAttention op(&pipe);
+ op.Init(q, k, v, mask, out, lse, B, Hq, Hkv, Sq, Skv, scale, causal, hasMask, outFp32);
+ op.Process();
+}
+
+std::vector deterministic_attention_ascend_forward(
+ torch::Tensor q,
+ torch::Tensor k,
+ torch::Tensor v,
+ bool causal,
+ double scale,
+ c10::optional key_padding_mask,
+ bool outFp32)
+{
+ TORCH_CHECK(q.is_privateuseone() && k.is_privateuseone() && v.is_privateuseone(),
+ "q, k, v must be on an NPU device");
+ TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4,
+ "q, k, v must be 4-D [B, H, S, D]");
+ TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(),
+ "q, k, v must be contiguous");
+ TORCH_CHECK(q.scalar_type() == at::kBFloat16 || q.scalar_type() == at::kHalf,
+ "q must be bf16 or fp16");
+ TORCH_CHECK(k.scalar_type() == q.scalar_type() && v.scalar_type() == q.scalar_type(),
+ "q, k, v must share the same dtype");
+ TORCH_CHECK(q.size(3) == HEAD_DIM && k.size(3) == HEAD_DIM && v.size(3) == HEAD_DIM,
+ "head dim D must be 128");
+ TORCH_CHECK(q.size(0) == k.size(0) && q.size(0) == v.size(0),
+ "batch size mismatch between q/k/v");
+ TORCH_CHECK(k.size(1) == v.size(1) && k.size(2) == v.size(2),
+ "k/v must have the same head and key-length layout");
+ TORCH_CHECK(q.size(1) % k.size(1) == 0, "Hq not divisible by Hkv (GQA group)");
+
+ const int64_t B = q.size(0);
+ const int64_t Hq = q.size(1);
+ const int64_t Hkv = k.size(1);
+ const int64_t Sq = q.size(2);
+ const int64_t Skv = k.size(2);
+
+ torch::Tensor mask;
+ bool hasMask = key_padding_mask.has_value() && key_padding_mask->defined();
+ if (hasMask) {
+ mask = key_padding_mask->to(torch::kBool).contiguous();
+ TORCH_CHECK(mask.is_privateuseone(), "key_padding_mask must be on an NPU device");
+ TORCH_CHECK(mask.dim() == 2 && mask.size(0) == B && mask.size(1) == Skv,
+ "key_padding_mask must be [B, Skv]");
+ }
+
+ torch::Tensor out = at::empty(
+ {B, Hq, Sq, HEAD_DIM},
+ q.options().dtype(outFp32 ? at::kFloat : q.scalar_type()));
+ torch::Tensor lse = at::empty({B, Hq, Sq}, q.options().dtype(at::kFloat));
+
+ // stream(true): flush the task queue before launch so the kernel cannot
+ // overtake earlier NPU work; outputs were allocated with at::empty (no
+ // queued initializer) for the same reason.
+ auto aclStream = c10_npu::getCurrentNPUStream().stream(true);
+ const uint32_t blockNum =
+ static_cast(std::min(B * Hq * Sq, MAX_BLOCKS));
+ uint8_t* maskPtr = hasMask ? reinterpret_cast(mask.mutable_data_ptr()) : nullptr;
+
+ if (q.scalar_type() == at::kBFloat16) {
+ deterministic_attention_ascend_kernel_bf16<<>>(
+ reinterpret_cast(q.mutable_data_ptr()),
+ reinterpret_cast(k.mutable_data_ptr()),
+ reinterpret_cast