perf(npu): retain K3 hot paths and add opt-in parallel verify gates

[McZyWu]

CANN 9.1 parallel-gate correction (2026-09-10)

With SGLANG_NPU_KDA_VERIFY_PARALLEL_GATES=1, four-host target graph capture failed at the beta slice's .reshape(()): Triton-Ascend 3.2 rejects a zero-dimensional block type. An isolated test of the previous head e5859c8 reproduced the same compile error with parallel gates enabled; the disabled control passed.

This correction extracts scalar beta with a one-element reduction and loads raw gates with flattened addresses before restoring the token/key axes. The second change fixes an independently reproduced numerical error for narrow strided raw gates (K=7, element stride=2) in the original two-dimensional load. Parallel FP32 activation remains inside the KDA kernel. Existing stride/mask behavior, the sequential state updates, and the separate/preactivated paths are retained.

The existing padding test now also covers gamma7, 6 heads, K/V128 and actual NPU graph replay. Replay changes beta inputs and checks output, every snapshot, padding NaNs, disabled snapshot writes, and persistent-state preservation against an independent CPU reference.

Measured validation of d4d0a89: 32 tests passed on A3 Ascend910_9362, CANN9.1, PyTorch2.10.0+cpu, torch_npu2.10.0.post4 and Triton3.2. These include all 3/8-step raw-versus-preactivated cases and all 16-step V-tile32/64 cases, with parallel gates off/on and bounded/unbounded activation. Sixteen passing cases enabled parallel gates. Applicable local pre-commit checks and git diff checks passed.

Known remaining failure: the next case, 16 steps with K/V128, V tile128, unbounded raw gates and parallel gates disabled, fails compilation with UB overflow: 1,583,104 bits required versus 1,572,864 bits available. The previous head e5859c8 reproduces exactly the same failure. This is the chip-local UB compilation limit, not HBM exhaustion. The 55-test suite therefore has 32 passes, 1 failure and 22 unexecuted tests; it is not a complete pass.

A separate working-tile correction is being prepared but is not in this PR head. Its first targeted compilation was externally interrupted by Docker SIGKILL before producing a passing test. Device and full-model validation are paused while other users' jobs run; no other processes will be stopped. Earlier interrupted runs and the current UB failure are retained separately. No full-model all-on acceptance or new performance result is claimed. The reported eager PagedAttentionOperation incident on the other machines also remains unresolved.

This PR retains the K3 DSpark/KDA hot paths and padding correctness changes, and adds an explicit token-vectorized gate experiment for fixed-width KDA verify. It is paired with framework PR31, which defaults to separate FP32 gate activation after the supplied profiles showed a slower original raw-gate specialization.

The branch now incorporates current a5-k3-0828 at 1ec11e08ea3999a0f84f8b78fa57f19d59892ca5, including #5 and #11. This supplies the chunk-KDA prefill operator required by PR31's framework base. The new parallel-gate changes are Python/Triton only; deploying the full paired stack still requires a library built with the current base operators.

Kernel behavior

• kda_target_verify_npu(..., precompute_raw_gates=True) computes FP32 decay and beta over all verify positions before loading the recurrent state. Activation stays inside the same kernel; there is no new global gate buffer or extra kernel launch. State updates and per-position snapshots remain in their original order.
• The experiment requires raw gates, at most 16 verify positions, and Triton-Ascend cann.extract_slice. Unsupported settings fail explicitly. value_block_size=32/64/128 independently selects the V tile; omission retains the original default.
• The raw/preactivated API and default dispatch remain compatible. Preactivated gates cannot apply the K3 lower bound twice. PR31 does not select the old per-token raw-gate path; the old API remains available for existing callers and the comparison benchmark.
• Negative initial-state indices skip undefined padded producer rows, emit exact zero output, leave snapshots untouched, and preserve the persistent source state. Cache slot 0 remains a valid live slot. Snapshot index -1 suppresses snapshot writes independently.
• Retain DSpark local/global top1, ragged input/output and fused output norm, packed-QKV stride handling, fixed-width causal-conv coverage, and test-group/changed-file registration. These paths are not removed or declared ineffective merely because the supplied workload did not exercise them.

Framework selection and A/B

In PR31, before service startup on every rank:

# Reference: standalone FP32 gate activation, original V tile.
export SGLANG_NPU_KDA_VERIFY_PARALLEL_GATES=0
export SGLANG_NPU_KDA_VERIFY_VALUE_BLOCK_SIZE=64

# Candidate: change only this setting and restart to recapture graphs.
export SGLANG_NPU_KDA_VERIFY_PARALLEL_GATES=1

The default for both variables is 0: separate activation and unchanged kernel tiling. The removed legacy SGLANG_NPU_FUSED_KDA_VERIFY_GATES framework switch cannot enable the new experiment.

The supplied top1 launch configuration had SGLANG_DSPARK_FOLDED_PROPOSAL=0, preventing construction of the fused sampler. Use folded proposal 1, folded sampling 0, fused local top1 1, BF16 W2/base logits and TP sharding; the folded result also requires a greedy batch and draft graph replay. The full settings and profiler signatures are documented in benchmark/README_kda_parallel_gates.md and PR31. This optimization changes draft logits communication, not target embedding AllReduce.

python -m pytest tests/python/sgl_kernel_npu/test_kda_target_verify.py -q

# Compare separate activation, original raw fusion, and token-vectorized fusion.
python benchmark/bench_kda_verify_parallel_gates.py \
  --batch 32 --steps 8 --heads 3 --dim 128 \
  --value-blocks 64 --output kda-gates-bv64.json

# Tune V tiling separately after establishing a same-run baseline.
python benchmark/bench_kda_verify_parallel_gates.py \
  --batch 32 --steps 8 --heads 3 --dim 128 \
  --value-blocks 32 64 128 --output kda-gates-tiles.json

The benchmark includes gate preparation, warms compilation before graph capture, checks output and every snapshot before/after replay, and records persistent-state preservation, timing samples, software versions, and loaded kernel identity. Larger V tiles can reduce repeated gate work while increasing local-memory pressure. A faster raw-fused result alone is insufficient: compare against the same-run separate-gate baseline, then measure end-to-end acceptance and TPOT.

Evidence and validation

The historical profiles measured 59.559 us/layer for separate gate preparation plus verify and 75.083 us/layer for original raw-fused verify, an observed difference of +15.524 us/layer. The new specialization has no measured NPU performance result yet. These profiles do not prove the cause of the reported 2 ms/token regression, and do not validate the optional top1/ragged/Conv3D paths.

Earlier CPU validation of head e5859c8, before the CANN9.1 correction above:

• 43 CPU semantic/API tests passed, executing actual kernel bodies with PyTorch tensor/pointer adapters: strided raw gates, packed QKV, grouped key/value heads, 3/8/16 positions, K/V tails, bounded/unbounded activation, V tiles 32/64/128, padding NaNs, valid slot 0, disabled snapshots, and unsupported experiment settings.
• Paired framework: 73 CPU contract tests passed for dispatch, precision/lower-bound placement, metadata reuse, padding, convolution layout, V tiles, and graph buckets.
• Black 24.10.0, isort 5.13.2, Ruff F401/F821, configured spelling checks, Python syntax, shell syntax, workflow YAML parsing, and git diff --check passed for relevant changed files.
• Tests remain registered in the existing kernel groups. The kernel device workflow triggers automatically for main, not this A5 branch; this update does not broaden that policy.

The earlier CPU harness does not compile Triton, execute CANN extensions, or measure graph/UB behavior. The latest CANN9.1 result above includes 32 passing cases and a remaining V-tile128 UB compilation failure. Complete graph replay, full-model acceptance, and performance remain unvalidated.

Before this update, the previous PR3 head fd4e17d recorded CANN 9.0 target-verify/conv/ragged and four-node padding tests in the PR's earlier validation. Those historical results apply to that head, not to the new parallel specialization or the newly synchronized base. They must not be counted as validation of this update.