Skip to content

perf(npu): optimize K3 dense/static verify path on A5 - #10

Draft
McZyWu wants to merge 11 commits into
zhaozx-cn:a5-k3-0828from
McZyWu:a5-k3-0828-dense-static-opt-20260903
Draft

McZyWu wants to merge 11 commits into
zhaozx-cn:a5-k3-0828from
McZyWu:a5-k3-0828-dense-static-opt-20260903

Conversation

@McZyWu

@McZyWu McZyWu commented Sep 3, 2026 •

Copy link
Copy Markdown

Summary

This is a new PR against a5-k3-0828 that carries the K3 KDA/DSpark
hot-path port from #8 and adds the dense/static follow-up optimizations.

  • Keep the KDA raw/preactivated gate contract explicit, including K3
    lower_bound handling.
  • Reuse fixed verify metadata and the existing fused KDA verify pieces.
  • Add an opt-in fixed-width [B, T, C] causal-conv path for dense/static
    target verify (SGLANG_NPU_KDA_DENSE_CONV3D=1).
  • Derive one graph-replay-safe int64 cache-index tensor from
    query_start_loc; zero-length graph-padding requests are mapped to -1
    and the same tensor is shared by causal-conv and KDA recurrent.
  • Capture generic DSpark draft graphs at TP-local request batch sizes,
    including B1, while preserving TBO and context-parallel alignment.
  • Keep dense output norm on the existing graph-proven fused path; the
    separate dense-only onorm experiment was not kept because it had no
    measurable graph gain.
  • Add focused tests for DSpark graph-batch alignment and dense KDA padding
    metadata, plus a reproducible full-model TP32/DP1 launcher.

The formatting-only baseline fixes are isolated in commit 313acfa5f and
match the A5 branch's pre-commit run --all-files requirements.

Paired kernel PR

Requires zhaozx-cn/sgl-kernel-npu#3
from branch McZyWu:a5-k3-0828-dense-static-opt-20260903.

The framework passes -1 for graph-padding KDA state rows. The paired
kernel skips recurrence for that sentinel, zeroes the dense output row, and
does not touch the snapshot or persistent state.

Validation

  • Framework unit tests in the CANN 9.0 container: 9 passed.
  • Paired kernel target-verify and dense3d causal-conv tests: passed on NPU.
  • Four-node kernel smoke tests on 209/212/216/217: 2 passed per node.
  • Dense causal-conv microbenchmark (B=4,T=8,C=1152,W=4,state=10):
    20.387 us -> 18.446 us (1.105x) in the representative run; repeats
    were 1.078x-1.104x.
  • Full 93-layer TP32/DP1 current-head run completed target graph capture on
    all 32 ranks (bs=[4,32]) and entered generic draft capture with
    bs=[1,4,32]. 29/32 ranks logged draft-capture completion before an
    external Docker SIGKILL stopped all four shared containers; the other
    three were interrupted in that same second. Dense3d was enabled and no
    KDA shape check, Triton BS conflict, traceback, or assertion occurred.
    End-to-end HTTP replay is pending an idle window.

The full run used SGLANG_RAGGED_VERIFY_MODE=static. Therefore the claimed
dense-path benefit comes from dense3d causal-conv, graph-safe padding, and
local draft graph buckets; compact-only ragged scatter/gather is not counted
as an active optimization in this result.


CI States

Latest PR Test (Base): ❌ Run #33762654951
Latest PR Test (Extra): ❌ Run #33762654713
Latest PR Test (AMD ROCm 7.2): ❌ Run #33762654653

@qybnb

qybnb commented Sep 3, 2026

Copy link
Copy Markdown

2ms

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants