perf(npu): optimize K3 DSpark greedy replay and verify hot paths

[McZyWu]

Kimi-K3 DSpark on NPU spends avoidable time in its greedy graph tail and in copies of merged KDA projection views. This PR gives DP1 greedy requests a separate graph, connects the generic BF16 VanillaMarkov head to vocab-sharded W2 computation and local top1, removes a duplicate draft QKV projection, and lets the KDA gate/norm consumers read the merged projection layout directly. Mixed/stochastic requests retain the full-logits sampling path and independent per-step random noise.

Greedy graph and W2 computation

Each DP1 draft batch-size bucket has greedy and sampling variants in the existing graph pool. Host sampling metadata chooses the variant; input updates finish before that variant replays. Sampling keeps staged [batch, gamma, vocab] noise and corrected logits. The configured speculative TP synchronization sites remain in both variants.

For the unquantized BF16 generic VanillaMarkov head, enabling SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD=1 and SGLANG_DSPARK_FUSED_LOCAL_TOP1=1 now activates the greedy optimization. At TP32/vocab163840, each rank projects W2 against 5120 vocabulary rows, keeps its LM-head logits local, and exchanges (value, global token id) candidate pairs. The full W2 checkpoint parameter stays available to stochastic/eager fallback; this is compute sharding, not a reduction in stored checkpoint memory. The existing BF16 parameter is used without changing its precision.

The generic head rounds base + bias to BF16 before argmax. Its new selector preserves that rounding, first-token tie behavior, vocabulary padding and NaN selection. The DSv4 selector uses a different FP32-add contract and is not substituted here. Every Markov step still executes the caller's configured TP broadcast after candidate selection. Unsupported head classes, quantized/FP32 weights, DP/CP layouts and vocabulary partitions retain the existing path. AUTO folded sampling is not disabled by this optimization.

Startup reports DSpark VanillaMarkov NPU greedy graph: W2 compute sharded over TP32, BF16 local top1 enabled; sampling retains full logits. In a greedy replay, expect seven 32×256 @ 256×5120 W2 products and _vanilla_local_top1_kernel / _vanilla_global_top1_kernel, replacing the full-vocabulary W2/add/argmax tail and the large base-logit all-gather.

DFlashAttention.forward() also reuses its already computed QKV tensor in forward_prepare_npu(). The NPU preparation previously projected the same hidden states again and discarded the first result. The five-layer draft now performs one QKV projection per layer.

Retained KDA and scheduler corrections

The NPU gated norm reads noncontiguous token/head strides directly. A small beta conversion kernel combines the strided BF16 read with a contiguous FP32 store; the original NPU sigmoid remains. The standalone FP32 KDA gate contract, one application of the lower bound, metadata reuse and valid cache slot zero are preserved. The serving launcher defaults to a 64-wide value tile; explicit launch-script values still take precedence. Optional parallel gates remain a separate experiment paired with kernel PR3. Kernel PR10 separately restores the original recurrence launch defaults and optimizes narrow shared SiTU.

The PR27 scheduler prerequisites and DP1 local-finalization correction remain included. Greedy graph selection does not turn off folded proposals or folded sampling. CUDA behavior and stochastic proposal semantics remain unchanged.

Validation

• 39 CPU/interface tests pass for the new sharded head and fallback predicates, generic-model attachment, BF16 Markov chain equality, duplicate QKV removal, existing folded sampling, graph variant/update ordering, and graph-bucket alignment.

• Two new native A3 NPUGraph tests pass for strided local inputs, BF16 rounding ties, vocabulary tails, infinities/NaNs, global first-id ties, and updated inputs on replay.

• A four-host, 32-NPU synthetic head test passes five changed-input replays at BS32/gamma7/vocab163840/rank256/hidden7168. All 1120 token IDs per rank match the original path (35840 comparisons across ranks). Every test process exited successfully; peak rank-0 allocated memory was 393485824 bytes.

• The previously validated KDA view increment passed 87 CPU/interface tests and four native gate/norm tests. Those attention files are unchanged by this head increment.

• Black, isort, Ruff F401/F821/UP037, syntax, registered-test/entry checks and git diff --check pass.

The four-host A3 head benchmark includes the LM projection, vocabulary/candidate communication, all seven Markov steps and the original per-step broadcast. With torch/torch_npu 2.10, CANN 9.0 and HCCL AIV expansion, five alternating rounds of 40 graph replays produced these medians of maximum-rank event time:

| Synthetic TP32 head | Median per proposal |

|---|---:|

| Original full-vocabulary path | 1.927142 ms |

| Sharded greedy path | 0.871268 ms |

| Reduction | 1.055873 ms / 54.8% |

This is an A3 component benchmark using torch.distributed HCCL collectives, not the user's A5 end-to-end TPOT or a full-checkpoint acceptance test. Synthetic token parity does not establish full-model acceptance parity. At an unchanged acceptance length of 4.3–4.4, a 1.056 ms per-proposal saving would correspond to about 0.24 ms/token; it does not demonstrate a 2 ms TPOT reduction.

Integration

Keep folded proposal/sampling enabled, keep the chunk-KDA precision fix, and restart to recapture graphs after updating. The W2 and local-top1 switches above are already present in the supplied launch script; no new switch is required. The newest changes do not edit kimi_k3.py, add kernel PR dependencies, or overlap the changed files of the other selected framework PR16/20/24/25/26/27. This does not certify an unseen private integrated tree byte-for-byte. Existing conflict resolutions can be retained by applying only the new incremental commits.

Context RoPE coverage update

Commit 9be7207f2e65f5c62022f5987de8d62dec2ec7dd restores the NPU 3D input contract for generic DFlash KV-only RoPE and DSpark stacked context RoPE. It retains the PR12 context-path adaptation without restoring its older sampling implementation and does not touch kimi_k3.py or MLA cache/FIA code.

Correction: this change does not require kernel PR13. The installed Ascend Triton compiler already supports non-power-of-two head counts. The original kernel passed ten NPU configurations including stacked heads, partial RoPE, strides and graph input updates. A direct original-stacked versus masked-kernel A/B found the proposed mask specialization slightly slower in all six tested shapes; kernel PR13 is closed and removed from integration.

The earlier comparison of five per-layer calls versus one stacked call only measured the benefit of stacking, not a benefit of replacing the original kernel. No extra TPOT improvement is assigned to that kernel change. The actual framework 3D methods are tested with the original kernel, including changing positions/data and 2/3/5 stacked layers. Fully fused context-cache paths may bypass this fallback.

Generic VanillaMarkov greedy W2 sharding is already provided by VanillaMarkovGreedyShard in this PR. Sampling/eager full-logits behavior is preserved; absence of the older configure_tp_shard name is not evidence that the greedy optimization is missing.

NPU graph embedding opt-out correction

Commit 1a4066caef30aa7d67d0e1f5216d0dc5134b4119 makes SGLANG_DSPARK_EMBED_IN_GRAPH=0 effective in NPU graph capture. The proposer already computes embeddings eagerly for this setting, but capture previously ignored those inputs whenever the model exposed forward_embed, so the graph embedded again. Capture now consumes the staged buffer. The NPU preplanned replay branch also stages DSpark inputs, as the normal load path already does.

The default enabled path, W2/local-top1 selection, folded sampling, KDA code and kimi_k3.py are unchanged by this increment. This provides a working isolation/fallback switch; it is not an established fix for the reported A5 TP32 AIV AllReduce MTE incident.

The previous source fails two new regression cases (capture ignores the disabled switch; preplanned staging leaves the buffer stale). The corrected source passes six CPU cases with the NPU case skipped, and seven cases on an A3 NPU. The native graph test replays four changed eager inputs through the production-selected static view. It uses 224x16 tensors, no HCCL and no full model; peak tensor allocation is 51,712 bytes. This does not reproduce or validate the user's four-machine A5 failure. Black/isort, selected Ruff checks and diff checks pass.

---

CI States

Latest PR Test (Base): ❌ Run #34387507438
Latest PR Test (Extra): ❌ Run #34387507206
Latest PR Test (AMD ROCm 7.2): ❌ Run #34387507446