Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
39 commits
Select commit Hold shift + click to select a range
2b3d875
feat(ascend): add RoPE kernel
zhangj1an Sep 10, 2026
f63c488
Merge pull request #378 from erfgss/feat/ascend-rope
inaniloquentee Sep 10, 2026
da4044b
[WS1][kernels] Deterministic attention Ascend C kernel
zhangj1an Sep 10, 2026
d4e6b0d
Merge pull request #320 from zhangj1an/feat/ws1-ascend-deterministic-…
zhangj1an Sep 10, 2026
55e24e8
feat(ascend): add prefix-shared attention Ascend C kernel
zhangj1an Sep 10, 2026
6cb5254
Merge pull request #340 from zhangj1an/feat/ascend-prefix-shared-atte…
zhangj1an Sep 10, 2026
06c8655
feat(ascend): add deterministic collective Ascend C kernel
zhangj1an Sep 10, 2026
92960a9
Merge pull request #355 from zhangj1an/feat/ascend-deterministic-coll…
zhangj1an Sep 10, 2026
aef08c1
feat(ascend): add batch-invariant RMSNorm Ascend C operator
zhangj1an Sep 10, 2026
998b5fd
Merge pull request #364 from erfgss/rmsnorm_ascend
zhangj1an Sep 10, 2026
2d761bc
[WS1][Ascend] [Qwen3-8b] Embedding ops
zhangj1an Sep 10, 2026
5b73a5f
Merge pull request #369 from zhangj1an/feat/ascend-deterministic-embe…
zhangj1an Sep 10, 2026
0f66509
[WS1][Ascend] [Qwen3-8b] Fused logp ops
zhangj1an Sep 10, 2026
ce87e78
Merge pull request #370 from zhangj1an/feat/ascend-deterministic-logp
zhangj1an Sep 10, 2026
b3f10da
[skill] add ws1 ascend kernel
zhangj1an Sep 10, 2026
dd0cf3b
[WS1][Ascend] [Qwen3-8b] LM head ops
zhangj1an Sep 10, 2026
fea1364
Merge pull request #371 from zhangj1an/feat/ascend-deterministic-lm-head
zhangj1an Sep 10, 2026
26c751e
[WS1][Ascend] [Qwen3-8b] Fused linear logp ops
zhangj1an Sep 10, 2026
1ebeec4
Merge pull request #372 from zhangj1an/feat/ascend-deterministic-line…
zhangj1an Sep 10, 2026
b2af98b
feat(ascend): add SwiGLU forward and backward kernels
zhangj1an Sep 10, 2026
348bdea
Merge pull request #381 from erfgss/feat/swiGLU_ascend
zhangj1an Sep 10, 2026
01b33bd
feat(ascend): add batch-invariant deterministic GEMM Ascend C kernel
zhangj1an Sep 11, 2026
d9b246e
fix(ascend): make DataCopyPad padding narrowing explicit in det_gemm …
zhangj1an Sep 11, 2026
58e322c
fix(ascend): correct det_gemm data movement, layouts, and tree reference
zhangj1an Sep 11, 2026
8b59417
Merge pull request #405 from zhangj1an/feat/ascend-deterministic-gemm
zhangj1an Sep 12, 2026
a82a52d
[WS1][Ascend] #266 closeout on NPU: ascend_bf16 as a third required p…
zhangj1an Sep 11, 2026
b1a437a
fix(ascend): close the WS1 C2/C4/C8 gaps for ascend_bf16
zhangj1an Sep 11, 2026
04d7bb8
fix(ascend): force the fused-logp kernel path for non-contiguous logits
zhangj1an Sep 12, 2026
6551777
fix(ascend): make deterministic attention padding-invariant (keyBegin…
zhangj1an Sep 12, 2026
8ed1693
fix(ascend): shape-invariant RMSNorm rstd via a fixed-order reduction
zhangj1an Sep 12, 2026
d452e95
fix(ascend): fixed FP32 pairwise tree for the RMSNorm backward reduction
zhangj1an Sep 12, 2026
2993cff
fix(ascend): layout-invariant attention backward via valid-token comp…
zhangj1an Sep 12, 2026
9039fee
fix(ascend): dedicated attention backward kernel for bitwise layout i…
zhangj1an Sep 12, 2026
17b18f1
test(ascend): bf16-reference local deviation for the C10 gate on 64 G…
zhangj1an Sep 12, 2026
1fe03e7
test(ascend): offloaded FP32 reference for the C10 gate on 64 GB HBM
zhangj1an Sep 12, 2026
9daf565
Merge pull request #406 from zhangj1an/feat/ws1-ascend-closeout
Flink-ddd Sep 12, 2026
e946139
Merge branch 'test' into feat/skill-ws1-single-card-kernel
Flink-ddd Sep 12, 2026
c7e2b70
Merge pull request #373 from zhangj1an/feat/skill-ws1-single-card-kernel
Flink-ddd Sep 12, 2026
761f4ad
Merge branch 'main' into test
Flink-ddd Sep 17, 2026
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
243 changes: 243 additions & 0 deletions .claude/skills/ws1-single-card-kernel/SKILL.md
Original file line number Diff line number Diff line change
@@ -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/<area>/<op>.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-<op> upstream/test`.
- PR base = `test`; title format: `[WS1][Ascend] [Qwen3-8b] <Op> 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/<n> -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/<op>_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 <recent-ascend-branch> -- 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/<area>/<op>.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_<OP>` enum member + npu priority map
override (`self._priority_map["npu"]["<op>"] = [ASCEND_..., PYTORCH_...]`).
8. `rl_engine/tests/test_dispatch.py` — npu priority assertion.
9. `tests/test_<op>_ascend.py` — pytest suite (PR 320 style, see below).
10. `docs/operators/<op>.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_<op>_ascend.py` structure:

- Module docstring stating the two orthogonal properties (correctness + batch
invariance).
- `_npu_available()` / `_ascend_kernel_available()` helpers +
`requires_ascend = pytest.mark.skipif(...)`.
- `TestAscend<Op>Correctness`: 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).
- `TestAscend<Op>BatchInvariance`: bitwise (`torch.equal`).
- `TestAscendRegistryDispatch`: `kernel_registry.get_op("<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 <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_<op>_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 + <details> 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.
95 changes: 95 additions & 0 deletions .github/workflows/ws1-chain-npu.yml
Original file line number Diff line number Diff line change
@@ -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
Loading
Loading