[longboat2010]

Based on #2120 (feat/npu-gdn-ascendc-ops), which implements the GDN part of the proposal in #2118. fla_npu's 26.10 line splits its API surface: the stable raw ops move to fla_npu.ops.ascendc.*, the legacy torch.ops.npu.* dispatcher registrations are no longer performed on import (they need an explicit fla_npu.load_legacy_torch_ops() call), and some 26.10 wheel variants ship no embedded legacy extension at all. The GDN stack as written in #2120 breaks or silently degrades on such wheels. This PR adapts the two affected modules; behavior on 26.6 is unchanged.

Since #2120 is not merged yet, this PR is stacked on its branch. The delta relative to #2120 is now three commits: the 26.10 compat fixes (two files), the env-switched composite GDN forward for the chunk route (touches chunk_gated_delta_rule_npu.py again), and the runtime-length l2norm kernels.

What (3 files)

1. xtuner/v1/ops/gated_deltanet/chunk_gated_delta_rule_npu.py

   • load_ascendc_ops() best-effort calls fla_npu.load_legacy_torch_ops() when present. On 26.10 wheels without the embedded legacy extension this raises FileNotFoundError — caught and logged once; loading continues via the stable API (fla_npu.ops.ascendc.*), which is what the chunk route actually uses. Legacy-API callers still fail with their own specific hasattr/assert carrying the real cause, instead of this generic loader error.
   • Backend-selection log lines are deduped (_log_once): a 16-layer × N-rank job previously printed the same resolution line once per layer per rank.
   • (second commit) Composite forward arm: the 26.10 wheel ships npu_chunk_gated_delta_rule_fwd — one AscendC call covering the chunk-local cumsum, the WY/KKT scores, the tri-inverse, the state pass and the output projection that the decomposed forward runs as three Triton helpers + three AscendC ops plus six transposes. The composite is layout-native BSND, so the module's T-major q/k/v/g/beta flow through un-transposed and o comes back T-major. The composite backward (npu_chunk_gated_delta_rule_bwd) is Ascend950-only — aclnn rejects it with "ChunkGatedDeltaRuleBwd is supported on Ascend950 only" on the 910 parts this repo trains on (measured) — so the arm pairs the fused forward with the unchanged decomposed backward: the composite's g_cumsum (fp32, in-kernel cumsum) and A ([B, HV, T, chunk] bf16, the solve_tril layout) replace exactly what the backward's save-for-backward used to hold, so backward stays copy-free. Kernel-facing flags match the module contract (use_qk_l2norm_in_kernel=False, use_beta_sigmoid_in_kernel=False, use_exp2=False, disable_recompute=True).

2. xtuner/v1/ops/gated_deltanet/causal_conv1d_npu.py (the bulk of the compat delta)

   • _ensure_conv_ops() rewritten as a dual-backend resolver, decided exactly once per process: on the decoupled runtime (detected via the presence of fla_npu.load_legacy_torch_ops) it imports the stable raw fla_npu.ops.ascendc.npu_causal_conv1d_fn / npu_causal_conv1d_bwd, with a second fallback to the deprecated raw npu_causal_conv1d entry whose host-metadata path differs (tracked by a host_list flag); on 26.6 it registers and uses the legacy torch.ops.npu.npu_causal_conv1d{,_bwd} names as before. XTUNER_NPU_GDN_API no longer selects the conv backend (a note is printed when it is set) — it only governs the chunk GDN op resolution, as its name always suggested.
   • Forward/backward calling conventions are dispatched per backend: stable = kwargs with host-side query_start_loc, NTD-layout backward returning out[0]/out[1]; legacy = positional torch.ops calls with a per-document backward split (26.6 dw accuracy).
   • The startup self-probe is upgraded from isfinite-only to a numeric check against an F.conv1d autograd reference (relative-error threshold 0.15 for fwd/bwd). This catches "finite but wrong" kernels: on the 26.10 dev wheel the stable conv adapter returned all-zeros for all-ones input and NaN-adjacent garbage otherwise — the old probe passed it and the corruption surfaced as NaN loss mid-training. With the numeric probe the fused path is rejected at startup and the unmasked-correction fallback engages cleanly.

3. xtuner/v1/ops/gated_deltanet/l2norm_npu.py (new; resolved by _fused_l2norm() in chunk_gated_delta_rule_npu.py)

   • fla_npu's blocked triton l2norm kernels (the D <= 512 path that serves head_dim 128) declare the row count T as tl.constexpr. Triton inlines it into the binary, so every new row count forces a fresh seconds-long compiler-subprocess run on the calling rank's main thread. With variable-length packing, a megabatch whose padded token count differs by a few tokens makes a handful of ranks hit a never-before-seen row count mid-run and stall the whole job at the next collective — observed as step-locked periodic throughput spikes, reproducible across independent runs.
   • The new module is fla_npu's blocked kernels with T passed as a runtime argument and do_not_specialize=["T"]. Making T runtime is not sufficient on its own: triton still specializes runtime int arguments by divisibility class, and the dead NB = cdiv(rows, 2048) parameter flipped its class between row-count families and compiled a second binary — so the unused parameter is dropped entirely (the kernel bodies never read it). The kernel bodies are otherwise unchanged: the if block_start < T guard and the block-pointer boundary_check already handle arbitrary T, and D/BD/BT stay constexpr (bound to the model's head_dim, constant in training). D > 512 keeps delegating to fla_npu's per-row kernel, which takes no row-count argument at all.
   • _fused_l2norm() resolves XTUNER_GDN_L2NORM: rt (default, the runtime-length kernels) / fla (the previous route) / inline; the startup [l2norm] log line reports which one is live.

Usage

The composite forward arm is env-switched, resolved per chunk_gated_delta_rule call site (slab-split and single-graph paths dispatch identically):

• XTUNER_NPU_GDN_FUSED_FWD unset or 1 (default): FusedFwdChunkGatedDeltaRule — composite forward + decomposed backward, whenever the wheel ships the composite. On wheels without it (26.6 and earlier) the arm falls back to the decomposed chain with a single log line, so a wheel downgrade degrades instead of breaking; under XTUNER_NPU_STRICT_KERNELS=1 only an explicit =1 with a missing composite fail-fasts at wiring time (the silent default-on fallback never raises).
• XTUNER_NPU_GDN_FUSED_FWD=0: ChunkGatedDeltaRule, the decomposed chain — byte-identical to the pre-arm code path.
• When active, the startup log prints [gdn_fused_fwd] fla_npu composite GDN forward active (layout=BSND, bwd decomposed) per rank.

Validation

• Runtime-length l2norm parity: forward and backward compared against fla_npu's kernels and a pure-fp32 decomposition over 10/100/1000 randomized iterations covering both head-count families and the +8-token row-count variant (up to 524416 × 128 rows): the forward is bit-identical to both, the backward is bit-identical to fla_npu; against the fp32 reference chain the backward deviation stays at ~1 bf16 ulp, the same band fla_npu itself sits in (the forward output is saved in bf16 and the backward recomputes from it).
• Runtime-length l2norm, single binary: across row counts 32768/32776/32780/4800 each kernel holds exactly one triton cache entry — fla_npu's constexpr row count yields one binary per value.
• Runtime-length l2norm, end-to-end: 50-step distributed run with the runtime-length kernels active: both binaries compile during warmup on every rank and the training-step window sees zero kernel compiles; the step-locked throughput spikes disappear (the spiked steps drop from ~27 s back to ~17 s) and steady-state step time is flat.
• fla_npu 26.6.0 (the version [Feature] NPU AscendC op stack for GatedDeltaNet (causal conv1d / chunk delta rule / fused gated norm) #2120 pins), dispatcher level, 6 input packs including degenerate boundaries: conv forward bit-exact (0.00000), dx ≤ 0.125, dw ≤ 1.0 (bf16 few-ulp, the same band as the unmasked path) — the legacy route is behaviorally unchanged.
• fla_npu 26.10.0+dev (stable API, no embedded legacy .so): the load_legacy_torch_ops skip is clean and the GDN chunk route resolves via the stable API; the broken stable conv adapter is rejected by the numeric probe at startup and the run degrades to the unmasked-correction conv (error band fwd 0.094 / dx 0.125 / dw 1.0), exit 0.
• Composite-forward parity (Ascend 910, HV=32 / K=V=128, module-faithful inputs — fp32 sum-based q/k l2norm, post-sigmoid bf16 beta, raw fp32 log-decay; varlen with uneven tails, varlen many-doc, and dense): o and all five autograd gradients compare with 0 violations at 2e-2/3e-2 tolerance; max |Δo| ≤ 4.9e-4, max |Δgrad| ≤ 7.8e-3 (bf16 rounding band); the composite's g_cumsum is bit-identical to the Mindspeed chunk_local_cumsum (max |Δ| = 0.0), so the backward's saved mirrors are the same tensors the decomposed chain would have saved.
• End-to-end A/B, three runs (16 Ascend 910 dies, 1 node, Qwen3.5, 8K seq, 50 steps each, profile_step=9, identical config and seed/data — the only delta is the knob): decomposed baseline 30.44 s steady per-step vs 30.10 s and 29.75 s for the fused arm across two independent runs (−1.1% / −2.3%, mean −1.7%; same-arm run-to-run spread 1.2% bounds the noise band, so the effect is direction-consistent and worth ~0.3–0.7 s/step at this shape). The composite backward (the heavier half) remains unavailable off-Ascend950.
• Numerics and memory in training (same three runs, rank 0): reduced_llm_loss at matched steps agrees to bf16 precision across all three arms (spread ≤ 1.7e-4, e.g. 3.26358 / 3.26347 / 3.26364 at step 43; grad_norm likewise), and peak memory is byte-identical on all three arms (max 39.53 GB, reserved 42.76 GB) — the composite forward allocates no extra persistent memory. XTUNER_FAIL_ON_NONFINITE=1 clean for all 50 steps on both fused runs.
• Cumulative tgs (1794.3 decomposed / 1783.3 first fused / 1897.5 second fused) is dominated by outlier steps: the first fused run's 203.6 s profiler dump on the profiled step did not reproduce (115.2 s in the second fused run, in line with the decomposed run's 118.7 s), confirming it as a one-off measurement artifact rather than a cost of the composite path.
• End-to-end (compat commits) (Qwen3.5 dense text, 256 Ascend NPUs = 16 nodes × 16, 8K seq, 10-step benchmark, mbz=1 / gas=2 / gbs=512): GDN resolves [gdn_ascendc] GDN kernel: fla_npu AscendC + [l2norm] fla_npu triton, 10/10 steps, 0 NaN, benchmark success; steady 18–21 s/step at 16384 tokens/rank/step (~772 tok/s per rank instantaneous on the final step), losses consistent across all 256 ranks.
• Strict-kernel mode (XTUNER_NPU_STRICT_KERNELS=1) still fail-fasts when a requested kernel route cannot resolve — no silent degradation in A/B bookkeeping.

Notes

• XTUNER_GDN_L2NORM=fla restores the previous fla_npu l2norm route if a pin is ever needed.
• No changes to non-NPU paths; on fla_npu 26.6 every route resolves exactly as in [Feature] NPU AscendC op stack for GatedDeltaNet (causal conv1d / chunk delta rule / fused gated norm) #2120.
• Enabling the fused conv on 26.10 requires a wheel whose stable conv adapter passes the numeric probe (the current dev wheel does not — the unmasked-correction fallback is used); the chunk GDN route is unaffected and runs on the stable API.