fix(qsa): restore SM121 correctness with Humanize and Kernel Design Agent

[BBuf]

Motivation

Follow-up to #36806; fixes the remaining SM121 / GB10 part of #36537.

QSA decode needs different paths on SM120 and SM121:

• Exact SM120 is numerically correct with FlashInfer's TRT-LLM paged decode.
• On SM121, that path silently corrupts long-context decode. The 120k, 190k,
  and 210k reproductions all returned 32/32 token id 0 (!) while the server
  remained healthy and returned HTTP 200.
• Excluding SM121 avoids the corruption, but the existing FA4 CuTe varlen
  implementation does not compile for QSA's packed one-query shape on GB10
  (expects coord and shape of view are weakly congruent).

Disabling FlashInfer autotune only makes the incorrect SM121 kernel start; it
does not fix its numerics.

Changes

• Add one optimized Triton kernel for the exact Qwen3.8 packed-QSA decode
  contract on SM121: BF16, D=256, 12:1 GQA, TP1 24Q/2KV or TP2 12Q/1KV,
  one query per sequence, 1 <= bs <= 128, and selected KV capacity <=2055.
• Route SM121 directly to this kernel. There is no opt-in switch and no second
  generic SM121 attention kernel in this PR; unsupported calls fail with a
  descriptive contract error.
• Keep SM100/SM120 on the TRT-LLM sparse-decode path and leave the existing
  FA2/FA4 selection unchanged on other architectures.
• Register the kernel and preserve its KDA-1.5 provenance using the integration
  structure from [Kernel] Add KDA NVFP4 GEMM for Qwen3.x on SM120 #36865.
• Add reference, architecture-gating, integration, CUDA Graph replay, complete
  validated batch coverage, and repeated atomic-state tests.

The kernel reads live sequence lengths on device, so captured graphs can replay
with changed cu_seqlens without host readback.

Agent workflow

Humanize provided the iterative implementation and review workflow used for
this change. The Kernel Design Agent ran through KDA-1.5 with Codex and Kimi K3
and produced the winning SM121 QSA kernel. This PR fixes SM121 QSA correctness
and integrates that single optimized Triton kernel. Humanize and the Kernel
Design Agent had distinct roles: Humanize structured the workflow, while the
Kernel Design Agent produced the kernel selected for integration.

Optimized kernel

The original correct Triton implementation launched one CTA per query head, so
the 12 query heads sharing one KV head repeatedly loaded the same K/V rows. The
optimized kernel groups those 12 heads into one CTA, uses BF16 tensor-core QK/PV
products with FP32 online-softmax state, and splits long KV rows across CTAs.
The final split performs a stable FP32 merge and resets its device counter in
the same launch.

Provenance:

• Task: radixark/KDA-1.5#4
  at 414ce456e14ae8546f77d9356d2c4d955c5bb7f1
• Winning submission: b4181149c8884ddb
• Exact submitted source SHA256:
  4f9977f88abfea4393a2add3a2c9255699f7e13b981dbc1a976b024b3b00e909
• Coverage: 15 real TP1/TP2 tensors captured from one- and two-DGX-Spark runs

The packaged kernel passes 15/15 real replays against both the pinned
production output and an independent FP32 reference (relative-L2 <= 2e-3),
with exactly one CUDA activity per invocation. Its all/large/small geomean over
the original correct Triton implementation is 2.0702x / 1.6951x / 2.3653x.
The final source also passed 150,000 consecutive launches with every counter
reset to zero.

An additional GB10 sweep covers both TP topologies, every batch size from 1 to
16, and short plus saturated selected-KV rows: 64/64 passed, maximum
relative L2 0.002422, and speedup 1.41x-5.09x over the original Triton
implementation.

Batch	Geomean speedup	Minimum speedup
1	2.72x	1.71x
2	2.31x	1.76x
3	2.02x	1.79x
4	1.93x	1.63x
5	1.82x	1.56x
6	1.67x	1.41x
7	1.77x	1.68x
8	1.72x	1.57x
9	1.73x	1.61x
10	1.67x	1.54x
11	1.64x	1.48x
12	1.61x	1.48x
13	1.79x	1.52x
14	1.77x	1.51x
15	1.82x	1.55x
16	1.83x	1.58x

The follow-up extended-batch sweep covers both TP topologies, batch sizes
17/24/32/48/64/96/128, and short, medium, plus saturated KV: 42/42 passed,
maximum relative L2 0.002410, 4.48x geomean speedup, and 1.58x minimum
speedup. No case regressed. The worst-case TP1 scratch allocation at the
largest qualified batch is 32.3 MiB.

Batch	Geomean speedup	Minimum speedup
17	1.86x	1.58x
24	1.94x	1.70x
32	2.53x	1.82x
48	4.39x	2.64x
64	8.37x	6.19x
96	10.15x	8.96x
128	10.66x	9.75x

These are isolated-kernel measurements against the original correct Triton
implementation. End-to-end serving remains intentionally measured at low
concurrency because a single GB10 cannot sustain medium/high-concurrency model
serving without changing the workload.

End-to-end performance and accuracy

Hardware: one NVIDIA GB10 / SM121 DGX Spark, CUDA 13.0, PyTorch 2.13.

Model: RadixArk/Qwen3.8-Flash-Next-NVFP4 TP1 with ModelOpt NVFP4,
FlashInfer CUTLASS FP4/MoE, exact FP8 PLE from NVMe, 262,144 context, and NEXTN
(3 steps, top-k 1, 4 draft tokens). The benchmark compared the original
correct Triton implementation with the final KDA kernel under otherwise
identical server flags.

Three-round low-concurrency A/B, 12/12 successful requests per round and the
same 45,946 input plus 1,667 output tokens:

Concurrency	Metric	Original Triton	KDA	Change
1	Total throughput (tok/s)	644.51	673.19	+4.45%
1	Mean E2E latency (ms)	6,173.50	5,899.09	-4.45%
1	Mean TTFT (ms)	1,914.83	1,753.04	-8.45%
4	Total throughput (tok/s)	1,068.37	1,111.08	+4.00%
4	Mean E2E latency (ms)	13,423.72	13,080.11	-2.56%
4	Mean TTFT (ms)	3,207.84	3,073.10	-4.20%

Accuracy A/B uses the first 50 official GSM8K test examples, five-shot Chat API
prompting, Qwen thinking, greedy decoding, concurrency 4, and a 2048-token
output limit:

Route	Score	Failed examples
Original Triton	49/50 (98%)	index 7
KDA	49/50 (98%)	index 7

The original correctness validation also passes 20/20 sequential requests,
4/4 concurrent requests, structured tool parsing, and exact NIAH retrieval at
120k, 190k, and 210k tokens with zero token-id-0 output.

Full commands, TPOT/ITL results, replay tables, and the output-limit sensitivity
check are in this validation comment.

Validation

• pre-commit run --files ...: pass
• compileall: pass
• git diff --check: pass
• pytest -q test/registered/kernels/test_qsa.py test/registered/kernels/test_kda_qsa_sm121.py -s:
  73 passed on SM121 after deleting the old implementation
• CUDA Graph replay with changed device-side cu_seqlens: pass
• Real tensor replay: 15/15, production output + FP32 reference
• Synthetic baseline differential: 106/106 cases across bs 1-128
• Complete low-concurrency sweep: 64/64, TP1/TP2 x bs 1-16 x short/long KV
• Extended-batch sweep: 42/42, TP1/TP2 x bs 17-128 x short/medium/long KV
• State stress: 150,000/150,000 launches, counters zero

The Base/Extra/AMD workflow failures below are the repository rebase gate, not
test failures in this change: CI currently requires main@1ec20fd, while this
PR intentionally targets qwen4-main-squashed, so the test jobs were skipped.

---

CI States

Latest PR Test (Base): ❌ Run #33226754199
Latest PR Test (Extra): ❌ Run #33226754138
Latest PR Test (AMD ROCm 7.2): ❌ Run #33226754194

[BBuf]

KDA SM121 QSA integration update

The final PR now contains one SM121 packed-QSA attention kernel, integrated
using the same package/registry/provenance structure as #36865:

• Added the KDA-optimized package under
  python/sglang/kernels/kda_kernels/qwen38_qsa_sm121.
• Registered attention.kda_qwen38_qsa_sm121 in the unified kernel registry.
• SM121 now routes directly to its flash-attn-compatible adapter; the temporary
  generic qsa/sm121_varlen.py implementation and opt-in environment variable
  have been removed.
• The complete qualified contract is CUDA SM121, BF16, D=256, 12:1 GQA,
  TP1 24Q/2KV or TP2 12Q/1KV, every bs from 1 through 128, and selected-KV
  capacity <=2055. Unsupported calls raise a descriptive error.
• Added registered tests for TP1/TP2 geometries, low- and extended-batch
  boundaries, CUDA Graph replay with changed device-side cu_seqlens,
  integration with QSA compaction, contract rejection, and repeated
  split-counter reset.

The net PR diff therefore adds exactly one @triton.jit QSA attention kernel.

Provenance and kernel result

• KDA task: radixark/KDA-1.5#4 at
  414ce456e14ae8546f77d9356d2c4d955c5bb7f1
• Winning submission: b4181149c8884ddb
• Byte-exact submitted source SHA256:
  4f9977f88abfea4393a2add3a2c9255699f7e13b981dbc1a976b024b3b00e909
• Captures: 15 real tensors from RadixArk/Qwen3.8-Flash-Next-NVFP4, covering
  one-Spark TP1 and two-Spark TP2 rank 0/rank 1, q_rows 1/3/4/12, and short plus
  saturated selected-KV rows.

Replay	Correct	All geomean	Large	Small
KDA authoritative winner	15/15	2.1156x	1.7066x	2.4414x
SGLang packaged integration	15/15	2.0702x	1.6951x	2.3653x
Independent second GB10	15/15	2.1216x	1.6958x	2.4634x

Every replay row passes both the pinned production-output check and an
independent FP32 reference with relative-L2 <=2e-3, and the timed path contains
exactly one CUDA activity. The final source also passed all 15 rows x 10,000
sequential launches (150,000 total), with every cached atomic counter returning
to zero.

After removing the generic implementation, I also swept the shapes not present
in the original capture set: both TP topologies x every bs from 1 to 16 x
short and saturated KV, for 64/64 passing cases. Maximum relative L2 versus
the original correct Triton implementation was 0.002422; observed speedup was
1.41x-5.09x, with no slower case.

I then widened the benchmark scratch and tested both TP topologies at
bs=17/24/32/48/64/96/128, with short, medium, and saturated KV rows. All
42/42 cases passed with maximum relative L2 0.002410; no case regressed:

Batch	Geomean speedup	Minimum speedup
17	1.86x	1.58x
24	1.94x	1.70x
32	2.53x	1.82x
48	4.39x	2.64x
64	8.37x	6.19x
96	10.15x	8.96x
128	10.66x	9.75x

The production limit is now the largest tested batch, 128. TP1's worst-case
scratch allocation at that boundary is 32.3 MiB. A new bs=17, TP1, long-KV
stress test crosses the old scratch boundary and verifies every split counter
returns to zero.

One-DGX-Spark end-to-end A/B

Hardware: one NVIDIA GB10 / SM121 DGX Spark, CUDA 13.0, PyTorch 2.13.

Model: RadixArk/Qwen3.8-Flash-Next-NVFP4 revision
7b719225242aacd3dbd3f9407468c2ee9a9d2594, TP1, ModelOpt NVFP4,
FlashInfer CUTLASS FP4/MoE, exact FP8 PLE from NVMe, 262,144 context, and NEXTN
with 3 steps / top-k 1 / 4 draft tokens. The A/B compared the original correct
Triton implementation with the final KDA kernel; all other server flags were
identical.

python3 -m sglang.bench_serving \
  --backend sglang-oai \
  --model qwen3.8-flash-next \
  --dataset-name random \
  --random-input-len 8192 \
  --random-output-len 256 \
  --random-range-ratio 0 \
  --num-prompts 12 \
  --max-concurrency ${CONCURRENCY} \
  --request-rate inf \
  --seed 1234 \
  --flush-cache \
  --disable-tqdm

CONCURRENCY is 1 or 4. Each number below is the mean of three rounds. Every
round completed 12/12 requests with the same 45,946 input tokens and 1,667
generated tokens. --random-range-ratio 0 means 8192 is the input-length upper
bound rather than a fixed per-request length.

Concurrency	Metric	Original Triton	KDA	Change
1	Total throughput (tok/s)	644.51	673.19	+4.45%
1	Mean E2E latency (ms)	6,173.50	5,899.09	-4.45%
1	Mean TTFT (ms)	1,914.83	1,753.04	-8.45%
1	Mean TPOT (ms)	30.05	29.76	-0.98%
1	Mean ITL (ms)	30.97	30.17	-2.56%
4	Total throughput (tok/s)	1,068.37	1,111.08	+4.00%
4	Mean E2E latency (ms)	13,423.72	13,080.11	-2.56%
4	Mean TTFT (ms)	3,207.84	3,073.10	-4.20%
4	Mean TPOT (ms)	74.22	74.78	+0.76%
4	Mean ITL (ms)	74.36	72.75	-2.16%

The concurrency-4 TPOT movement is a small regression/noise point; total
throughput, E2E latency, TTFT, and ITL all improve in the same three-round A/B.

End-to-end accuracy A/B

I used the first 50 official GSM8K test examples, five-shot prompting, Chat API,
Qwen thinking enabled, greedy decoding, concurrency 4, and a 2048-token output
limit:

python3 -m sglang.test.run_eval \
  --host 127.0.0.1 --port 8892 \
  --model qwen3.8-flash-next \
  --eval-name gsm8k --api chat \
  --num-examples 50 --num-threads 4 --num-shots 5 \
  --max-tokens 2048 --temperature 0 \
  --thinking-mode qwen-3 \
  --gsm8k-data-path gsm8k-test.jsonl

Route	Score	Failed examples
Original Triton	49/50 (98%)	index 7
KDA	49/50 (98%)	index 7

The same example failed in both reports. I also checked the sensitivity of a
1024-token cap: Triton scored 47/50, while two KDA repeats scored 44/50 and
49/50. Those additional failures were responses cut off mid-reasoning at the
token limit, so the 2048-token matched A/B without those truncations is the
primary accuracy result.

The original correctness validation also passes 20/20 sequential requests,
4/4 concurrent requests, structured tool parsing, and exact NIAH retrieval at
120k, 190k, and 210k tokens with zero token-id-0 output.

Validation after single-kernel refactor

pytest -q test/registered/kernels/test_qsa.py \
  test/registered/kernels/test_kda_qsa_sm121.py -s
73 passed, 17 warnings in 6.85s

All pre-commit hooks pass, including Black, isort, Ruff, codespell, and the
registered-test CI registry checks. compileall and git diff --check also
pass.

Server-side speculative configuration

The benchmark command above is the client command only. The server was explicitly launched with the checkpoint's embedded MTP/NEXTN path:

--speculative-algorithm EAGLE \
--speculative-draft-model-path /persistent/qwen38-model \
--speculative-num-steps 3 \
--speculative-eagle-topk 1 \
--speculative-num-draft-tokens 4

This is EAGLE/NEXTN MTP, not SGLang's separate DSPARK speculative algorithm. The draft path is the same target checkpoint and is redundant for this embedded-MTP architecture. Without --speculative-algorithm EAGLE (or the NEXTN alias), speculative decoding is off by default. Therefore the E2E table in this comment is specifically an MTP-on result.

[BBuf]

Two-DGX-Spark / TP2 follow-up

I completed the deferred two-Spark validation on two NVIDIA GB10 / SM121
systems with TP2. Both servers used the same model revision, server flags,
random_seed=1234, Mamba-cache size (41 slots), request set, and eventually the
same warmed SGLANG_CACHE_DIR for the common Triton/Inductor kernels.

Server-side speculative configuration

The serving benchmark command below is only the benchmark client. The server
was explicitly launched with the checkpoint's embedded MTP/NEXTN path:

--speculative-algorithm EAGLE \
--speculative-draft-model-path /persistent/qwen38-model \
--speculative-num-steps 3 \
--speculative-eagle-topk 1 \
--speculative-num-draft-tokens 4

This is EAGLE/NEXTN MTP, not SGLang's separate DSPARK speculative
algorithm. --speculative-draft-model-path was the same target checkpoint and
is redundant for Qwen4ExpForConditionalGeneration, because its one-layer
model.mtp.* drafter is embedded in the target checkpoint. Without
--speculative-algorithm EAGLE (or its NEXTN alias), SGLang defaults to no
speculative decoding.

Correctness

• The matched Chat API GSM8K run scored 19/20 for original Triton and 19/20
  for KDA (five-shot, Qwen thinking, greedy, concurrency 4).
• A fixed 8192-token input with 128 forced decode tokens produced identical
  output token IDs and identical text for the two implementations.
• This is in addition to the existing 15/15 real-tensor replay, including both
  TP2 ranks, and the FP32-reference checks documented above.

Matched 8K decode trace

I captured ten decode forwards per implementation with the same 8192-token
input. Each trace contains 150 packed-QSA calls.

Metric	Original Triton	KDA	Change
QSA mean per call	114.897 us	25.914 us	4.43x faster
QSA total, 150 calls	17.235 ms	3.887 ms	-77.45%
GPU-kernel total excluding NCCL AllReduce/AllGather	534.739 ms	525.315 ms	-1.76%

The QSA saving is 13.348 ms over ten forwards, or about 1.335 ms per decode
forward. This is the directly attributable TP2 benefit of the PR kernel.

The trace also explains why whole-server TP2 numbers need care: NCCL AllReduce
alone was 641.855 ms in the baseline capture and 526.492 ms in the KDA capture,
a 115 ms difference that is much larger than the QSA optimization itself.

Low-concurrency serving result

Every serving round completed 12/12 requests with the same 45,946 input tokens
and 1,667 output tokens. After sharing the warmed common-kernel cache, the
results were still mixed:

Concurrency	Metric	Original Triton	KDA round 1	KDA round 2
1	Total throughput (tok/s)	593.95	560.23	540.22
1	Mean TPOT (ms)	31.58	31.54	36.88
1	MTP accept length	3.632	3.566	3.331
4	Total throughput (tok/s)	942.81	974.56	882.10
4	Mean TPOT (ms)	91.07	86.04	98.46
4	MTP accept length	3.575	3.818	3.705

The two KDA rounds ran on the same already-warm server but still moved by
6-10% as MTP acceptance and cross-node communication changed. Concurrency 4
beats the baseline in one round and trails it in the next. Therefore I do
not claim a stable TP2 end-to-end speedup or hide the negative rounds. The
defensible conclusion is:

• TP2 correctness passes at kernel, token-output, and sampled task-accuracy
  levels.
• The new QSA kernel itself is 4.43x faster on a matched real 8K decode trace
  and saves about 1.335 ms per model decode forward.
• End-to-end TP2 serving on this pair is communication/speculation dominated
  and has no stable attributable gain in the current measurements.

The one-DGX-Spark three-round E2E result in the previous comment remains the
cleaner whole-model performance signal because it does not contain the TP2
cross-node collective variance.

[hashd1ve]

Independent validation on one DGX Spark (GB10, sm_121, TP=1), on top of the day-0 image with the single-Spark recipe (NVFP4 routed experts, dense path in FP8, exact FP8 PLE table mmap'd from NVMe, NEXTN 3/1/4, 262,144 context):

Kernel in isolation, at the model's real QSA decode shape (24 query heads / 2 KV heads / head dim 256 — the differential cases in this PR stop at head dim 192), against a PyTorch fp32 reference, BF16 and FP16, packed lengths 1–2051, batch 1–7: max abs error 0.001; CUDA-graph replay with changed device-side cu_seqlens (including a zero-length row): ok; 0.104 / 0.176 / 0.341 ms per call at batch 1 / 4 / 7 (2051 rows each).

End to end, qwen_sparse_attn_backend.py with the trtllm gate left at sm100 and this PR's hunk, kernel mounted at qsa/sm121_varlen.py:

prompt tokens	needle (mid-depth), 4 runs each	decode after the prompt
120,103	4/4 exact	35–41 tok/s
190,103	4/4 exact	30–39 tok/s
210,103	4/4 exact	38–48 tok/s

No token id 0 anywhere, /health 200 after each request, min free memory 5 GB, PSI full avg10 max 2.9 (chunked-prefill 4096). 13-case code suite unchanged vs the previous build. For contrast, my earlier build with the widened gate (→ XQA on this device) is what #36556/#36806 measured as corrupt from ~120k; I never saw it because my long-context runs were n=1.

Scripts and logs: https://github.com/hashd1ve/qwen38-flash-next-one-dgx-spark (patch 2 replaced with this PR's fallback, README documents the correction).

[hasso5703]

Field data from a single DGX Spark (GB10), in case it helps the review.

Our installer for this box (hasso5703/dgx-spark-qwen38, v1.5.2 and later) ships the Triton sm121_varlen kernel from the 2026-08-28 revision of this PR, vendored verbatim on top of the qwen38flashnext image, with the is_sm121() route and without the widened trtllm gate. Today on that machine, through the served endpoint: an hour of 39 consecutive prompts at 60k, 100k and 120k tokens with no cache flush between them, then single prompts at 135k, 150k, 166k, 171k and 177k tokens, every one with exact needle retrieval (the answer is a passphrase planted at the depth), and canaries passing after every boot. NEXTN decode on that path stays at 34 to 42 tok/s. With the widened gate instead of this route, the same hardware produced runs of token id 0 deep in long contexts, as hashd1ve reported above.