🚧 [WIP][llm][npu] Split mixed chunked-prefill FIA into prefill/decode calls on Ascend 950

[TallMessiWu]

Summary

On Ascend 950 (A5), a mixed chunked-prefill batch currently goes through a single
npu_fused_infer_attention_score (FIA) call that has to cover two very different
shapes at once: long causal prefill queries and one-token decode queries. The
kernel tiles for the worst case, so the decode part pays the prefill's cost.

This PR splits that into two independent FIA calls — one for the prefill span and
one for the decode span — preserving SGLang's [prefill][decode] token order.

The design follows vllm-ascend #11948.
Note the batch layouts are mirrored: vllm-ascend reorders to [decode][prefill]
via split_decodes_and_prefills, while SGLang's mix_with_running() produces
[prefill][decode], so the slicing and the cumulative-length rebasing differ.
The 9.4% reported there is a single-FIA-call measurement (4360us → 3949us on one
batch), not a serving result; nothing about it is carried over here.

Dependencies and current validation

Updated on 2026-10-10 with ordinary merges; base remains main. Branch head: 01bb8dc685. Includes official main c55790d38b.

Contains the shared fixes from #32745 at cf48b477db; Ascend 950 Qwen3.5 requires the paired target wheel from sgl-project/sgl-kernel-npu#638. #36426 is a separate validation dependency. Review the FIA increment against the refreshed #32745 branch.

Current-head checks: 70 attention tests executed in isolation against the actual Python definitions with CPU tensors / mocked NPU calls, 8 graph-policy cases (TP/DP, mixed/pure, split enabled/disabled), 8 isolated Gemma tests, shared partial-scale probes, and applicable pre-commit hooks passed. These are CPU/static checks, not 910/950 operator, model-load, warmup, accuracy, or performance validation. Full SRT test collection was unavailable in this environment because its serving dependencies are not installed. Target hardware validation remains required.

Remains Draft. GitHub reports viewerCanReopen=false; the fork head is updated but reopening needs repository permissions. Hardware measurements below predate this refresh and must be rerun before drawing conclusions about this head.

Changes

Split path

• AscendAttnBackend.forward_mixed() picks between one FIA call and two.
• _forward_fia_mixed_split() issues the prefill call over query[:n_prefill]
  with block_table[:n_prefill_reqs], then the decode call over the remainder.
  The decode span's cumulative query lengths are rebased to start at 0, which the
  TND layout requires; both calls read the full paged KV cache and select their
  own pages through their block-table slice.
• _should_split_fia_mixed() gates on: split enabled, FIA in use, non-MLA, no
  topk_indices, AttentionType.DECODER, not cross-attention, no sliding
  window, GQA head counts divisible, and boundary metadata present.

Mixed-batch boundary plumbing

• ScheduleBatch.mix_with_running() records mixed_num_prefill_reqs /
  mixed_num_prefill_tokens before merge_batch() appends the decode
  requests, then sets ForwardMode.MIXED.
• ForwardBatch carries both fields; _set_fia_mixed_split_metadata() validates
  them against seq_lens_list_cumsum and fails loudly on an inconsistent
  boundary rather than silently computing the wrong attention.
• TBO sub-batches explicitly reset the boundary to None — their token ranges do
  not match the parent batch, so they keep the single-call path.

Keeping ForwardMode.MIXED alive until attention dispatch

Two places used to erase MIXED before it reached forward_mixed():

• Hybrid-model wrapper. For hybrid GDN models (Qwen3.5), ModelRunner
  holds an AscendHybridLinearAttnBackend, not the AscendAttnBackend —
  the latter is its full_attn_backend child. HybridLinearAttnBackend.forward()
  overrides the base dispatch and had no MIXED branch, so mixed batches fell
  through to forward_extend() and the split never ran. Added the branch plus a
  forward_mixed() that routes full-attn layers to the child while the
  linear/GDN side keeps its extend path.
• Prefill CUDA graph. Graph replay rewrites MIXED to EXTEND to match a graph
  captured with EXTEND, and that rewrite reaches attention dispatch.
  The runner's common prefer_eager_mixed_prefill policy declines these mixed
  batches in both schedule-time DP votes and forward-time graph eligibility;
  pure prefill batches retain the graph. This mirrors vllm-ascend,
  where ACL graphs only capture DecodeOnly and a mixed batch is always eager.

enable_fia_mixed_split is declared on AttentionBackend (default False) so
dispatch sites read it directly instead of through a defensive getattr, and
wrappers mirror it from their full-attn child.

Gating

Control	Default	Effect
SGLANG_NPU_FIA_MIXED_SPLIT	True	Master switch for the split
supports_fia_mixed_split()	—	SoC probe; true only on Ascend 950 (get_soc_version() == 260)
ASCEND_USE_FIA	—	The split only applies to the FIA path

On every other SoC and on the non-FIA path this is a no-op.

Key NPU APIs used

API	Purpose
torch.ops.npu.npu_fused_infer_attention_score	Both the prefill and the decode call; TND layout, sparse_mode=3, paged KV via block_table + block_size=page_size
torch_npu.npu.get_soc_version()	Restrict the split to Ascend 950

Files Changed

Area	Change
hardware_backend/npu/attention/ascend_backend.py	Split path, profitability gate, boundary metadata validation, ascend.fia_mixed.* profile ranges
hardware_backend/npu/utils.py	supports_fia_mixed_split() SoC probe
layers/attention/hybrid_linear_attn_backend.py	MIXED dispatch + forward_mixed() for hybrid GDN wrappers
layers/attention/base_attn_backend.py	Declare enable_fia_mixed_split
model_executor/runner/prefill_cuda_graph_runner.py	Use the common mixed-prefill policy for DP votes and forward dispatch
managers/schedule_batch.py, model_executor/forward_batch_info.py	Record and carry the prefill/decode boundary
batch_overlap/two_batch_overlap.py	Reset the boundary on TBO sub-batches
environ.py	SGLANG_NPU_FIA_MIXED_SPLIT
Tests	Backend unit tests, boundary out-of-place tests, TBO marker test, manual A5 test

Accuracy Tests

• GSM8K on Ascend A5, Qwen3.5-27B, TP=1, --enable-mixed-chunk,
  --chunked-prefill-size 4096, 100 concurrent requests, greedy decode:

  Split	GSM8K
  off	91.05
  on	90.22

  The concurrency is what makes this the relevant test: at that batch size the
  mixed batches carry many concurrent decode requests, so the multi-request
  slicing of block_table / seq_lens_cpu_int and the rebased cumulative-length
  list are exercised on every step, not just the single-decode-request case.

  The 0.83-point gap is inside this setup's resolution. Splitting one FIA call
  into two changes the bf16 reduction order, one flipped argmax early in a
  chain-of-thought answer redirects the rest of it, and GSM8K's own binomial
  standard error at this size is ~0.8 points per side.

• Operator-level max absolute difference across all 40 shapes in the sweep
  below: exactly 0 for head_dim=256, 4.883e-4 for head_dim=128 — bf16
  rounding order, no systematic drift, and identical whether the prefill span
  is 1024 or 3968.

Speed Tests and Profiling

Operator microbenchmark (Ascend 950)

Direct FIA calls, not a serving measurement. One rank, simulated TP by dividing
head counts, 10 warmup + 200 timed iterations, paired per unique local attention
shape.

The shape is taken from a mixed batch this branch actually produces
(Qwen3.5-27B, TP=1, --chunked-prefill-size 4096, 32 concurrent requests):
3968 prefill tokens against 23 decode requests, kv_len=8192,
block_size=128, mask_size=2048 — the last one being what
AscendAttnMaskBuilder hardwires, independent of the prefill span, because
sparse_mode=3 takes a compressed causal template.

Aggregation: one paired median gain per unique local shape within each TP
Test: exact one-sided Wilcoxon signed-rank test
H0: gain is centered at <= 5%   Ha: gain is centered at > 5%
alpha=0.05, p-value=1.90386e-06, W+=729.0, n=40
Configurations above target: 29/40
Worst configuration: TP=8 Qwen3.5-0.8B/2B, gain=-11.06%
Verdict: PASS

Four shapes regress, and they share a signature -- head_dim=256 with at most
three local query heads:

TP	Model(s)	Q/KV	HD	Gain
8	Qwen3.5-0.8B, 2B	1/1	256	-11.06%
8	Qwen3.5-27B	3/1	256	-3.19%
4	Qwen3.5-0.8B, 2B	2/1	256	-1.66%
8	Qwen3.5-4B, 9B, 35B-A3B	2/1	256	-1.62%

head_dim=128 never regresses, not even at 2 local query heads (+4.87%), so
this is not a plain head-count threshold -- the wider head interacts with it.

Full per-configuration table (40 unique local shapes)

TP | Model(s)                              | Q/KV |  HD | OFF p50 ms | ON p50 ms |    Gain |  Max diff
---+---------------------------------------+------+-----+------------+-----------+---------+----------
 1 | Qwen3.5-0.8B,Qwen3.5-2B               |  8/2 | 256 |      1.320 |     1.173 | +11.29% |         0
 1 | Qwen3.5-4B,Qwen3.5-9B                 | 16/4 | 256 |      2.744 |     2.358 | +14.16% |         0
 1 | Qwen3.5-27B                           | 24/4 | 256 |      3.774 |     3.181 | +15.71% |         0
 1 | Qwen3.5-35B-A3B                       | 16/2 | 256 |      2.496 |     1.909 | +23.43% |         0
 1 | Qwen3.5-122B-A10B,Qwen3.5-397B-A17B   | 32/2 | 256 |      4.298 |     3.611 | +15.79% |         0
 1 | Qwen3-0.6B,Qwen3-1.7B                 | 16/8 | 128 |      1.293 |     1.279 |  +1.18% | 0.0004883
 1 | Qwen3-4B,Qwen3-8B                     | 32/8 | 128 |      2.050 |     1.864 |  +9.16% | 0.0004883
 1 | Qwen3-14B                             | 40/8 | 128 |      2.567 |     2.163 | +15.92% | 0.0004883
 1 | Qwen3-32B                             | 64/8 | 128 |      4.026 |     3.122 | +22.14% | 0.0004883
 1 | Qwen3-30B-A3B                         | 32/4 | 128 |      2.070 |     1.616 | +22.01% | 0.0004883
 1 | Qwen3-235B-A22B                       | 64/4 | 128 |      3.955 |     3.023 | +23.57% | 0.0004883
 2 | Qwen3.5-0.8B,Qwen3.5-2B               |  4/1 | 256 |      0.731 |     0.709 |  +3.08% |         0
 2 | Qwen3.5-4B,Qwen3.5-9B                 |  8/2 | 256 |      1.334 |     1.178 | +11.72% |         0
 2 | Qwen3.5-27B                           | 12/2 | 256 |      1.874 |     1.592 | +15.13% |         0
 2 | Qwen3.5-35B-A3B                       |  8/1 | 256 |      1.323 |     1.060 | +19.97% |         0
 2 | Qwen3.5-122B-A10B,Qwen3.5-397B-A17B   | 16/1 | 256 |      2.555 |     1.993 | +21.82% |         0
 2 | Qwen3-0.6B,Qwen3-1.7B                 |  8/4 | 128 |      0.715 |     0.705 |  +1.66% | 0.0004883
 2 | Qwen3-4B,Qwen3-8B                     | 16/4 | 128 |      1.174 |     1.001 | +14.94% | 0.0004883
 2 | Qwen3-14B                             | 20/4 | 128 |      1.389 |     1.153 | +17.33% | 0.0004883
 2 | Qwen3-32B                             | 32/4 | 128 |      2.045 |     1.598 | +21.97% | 0.0004883
 2 | Qwen3-30B-A3B                         | 16/2 | 128 |      1.170 |     0.861 | +26.59% | 0.0004883
 2 | Qwen3-235B-A22B                       | 32/2 | 128 |      2.044 |     1.547 | +24.31% | 0.0004883
 4 | Qwen3.5-0.8B,Qwen3.5-2B               |  2/1 | 256 |      0.485 |     0.493 |  -1.66% |         0
 4 | Qwen3.5-4B,Qwen3.5-9B,Qwen3.5-35B-A3B |  4/1 | 256 |      0.713 |     0.693 |  +2.86% |         0
 4 | Qwen3.5-27B                           |  6/1 | 256 |      0.975 |     0.858 | +12.11% |         0
 4 | Qwen3.5-122B-A10B,Qwen3.5-397B-A17B   |  8/1 | 256 |      1.304 |     1.044 | +19.94% |         0
 4 | Qwen3-0.6B,Qwen3-1.7B                 |  4/2 | 128 |      0.413 |     0.406 |  +1.67% | 0.0004883
 4 | Qwen3-4B,Qwen3-8B                     |  8/2 | 128 |      0.644 |     0.557 | +13.51% | 0.0004883
 4 | Qwen3-14B                             | 10/2 | 128 |      0.764 |     0.632 | +17.59% | 0.0004883
 4 | Qwen3-32B                             | 16/2 | 128 |      1.172 |     0.860 | +26.78% | 0.0004883
 4 | Qwen3-30B-A3B                         |  8/1 | 128 |      0.635 |     0.496 | +21.88% | 0.0004883
 4 | Qwen3-235B-A22B                       | 16/1 | 128 |      1.165 |     0.880 | +24.54% | 0.0004883
 8 | Qwen3.5-0.8B,Qwen3.5-2B               |  1/1 | 256 |      0.363 |     0.402 | -11.06% |         0
 8 | Qwen3.5-4B,Qwen3.5-9B,Qwen3.5-35B-A3B |  2/1 | 256 |      0.484 |     0.493 |  -1.62% |         0
 8 | Qwen3.5-27B                           |  3/1 | 256 |      0.590 |     0.609 |  -3.19% |         0
 8 | Qwen3.5-122B-A10B,Qwen3.5-397B-A17B   |  4/1 | 256 |      0.714 |     0.693 |  +3.03% |         0
 8 | Qwen3-0.6B,Qwen3-1.7B                 |  2/1 | 128 |      0.263 |     0.250 |  +4.87% | 0.0004883
 8 | Qwen3-4B,Qwen3-8B,Qwen3-30B-A3B       |  4/1 | 128 |      0.375 |     0.307 | +18.22% | 0.0004883
 8 | Qwen3-14B                             |  5/1 | 128 |      0.445 |     0.355 | +20.07% | 0.0004883
 8 | Qwen3-32B,Qwen3-235B-A22B             |  8/1 | 128 |      0.637 |     0.496 | +22.04% | 0.0004883

Gain vs prefill span

Qwen3.5-27B, TP=1, 23 decode requests, kv_len=8192:

prefill tokens	single FIA	split FIA	gain
1024	2.124 ms	1.499 ms	+29.5%
2048	2.782 ms	2.250 ms	+19.2%
3072	3.342 ms	2.750 ms	+17.7%
3968	3.794 ms	3.176 ms	+16.3%
6144	4.441 ms	3.906 ms	+12.0%
8192	4.673 ms	4.153 ms	+11.1%

Positive across the whole range. The gain shrinks as the prefill span grows,
since the span increasingly dominates its own tiling cost and the decode
requests distort it proportionally less. The span is capped by
--chunked-prefill-size, not by the request's input length, so this axis is a
deployment choice rather than a property of the traffic.

No end-to-end numbers are claimed

Qwen3.5 is a hybrid GDN model: layers_block_type marks every
full_attention_interval-th layer as full attention (4 here), so only a quarter
of the layers reach FIA at all. Measured over three runs per side on a workload
where mixed batches are 89% of all scheduler steps, the deltas on throughput,
TTFT, TPOT, ITL and E2E latency all landed within +/-0.4%, against a
run-to-run spread of 1.4-5.3%. That is consistent with the arithmetic — a 16%
kernel win on ~8% of a step's time is ~1.4% — and it is below what this setup
can resolve.

Reporting a serving number from that would be reporting noise, so none is given.
An architecture where every layer is full attention would see proportionally
more of the kernel gain.

Open questions

1. Does the gate need a shape threshold? _should_split_fia_mixed() has no
   head-count condition, so the four regressing shapes above take the slower
   path. tp_q_head_num >= 4 when qk_head_dim >= 256 would cover all four,
   at the cost of also excluding one +2.86% shape (TP=4, 4/1, HD=256). Whether
   that is worth a hardware-shaped condition in the gate is a judgement call I
   would rather have a reviewer make -- the alternative is leaving it and
   accepting a regression on 1-3-head TP=8 deployments of small Qwen3.5 models.
2. Interaction with the prefill CUDA graph. Declining mixed batches trades
   the graph for the split on those steps. On the tested A5 configuration the
   prefill graph is auto-disabled (cuda_graph_config resolves
   prefill.backend='disabled'), so the trade is currently free — that stops
   holding once it is enabled.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #32805837805
Latest PR Test (Extra): ❌ Run #32805837633
Latest PR Test (AMD ROCm 7.2): ❌ Run #32805837824

[github-actions]

Thanks @TallMessiWu. Closing this because it is still a draft and has not been updated in 38 days.

Reopen it if the work is still relevant.

Some directories moved recently, so an older branch may need retargeting:
sgl-kernel/ -> python/sglang/kernels/aot/, python/sglang/jit_kernel/
-> python/sglang/kernels/jit/, docs/ -> docs/docs/ (.mdx),
bench_serving.py -> benchmark/serving.py, test/srt/ -> test/registered/.