[CUDA] Add speculative decoding to paged XQA (#32340)
### Description
- Add paged XQA kernels for multi-token speculative verification at head
size 256 and group size 6, covering INT8, FP8 and native FP16/BF16 paged
KV caches.
- Route metadata-bounded query groups of 2-8 tokens through the new
kernels while preserving the existing one-token XQA ABI and fallbacks.
- Generate the linear causal mask from `cumulative_seqlens_q` on device,
including ragged batches and zero-query requests, without an operator
schema change.
### Motivation and Context
Qwen3.8 DFlash2 verifies seven tokens per target step. The existing
paged XQA backend only handled one query token per sequence, so
verification fell back to the generic path. This adds the narrowly
scoped H256/group-6 specialization needed by that workload.
The implementation reuses the existing 128-token XQA page mapping, folds
per-channel K/V scales into Q/output, sizes workspace for multiple
32-row tiles, and keeps all per-request lengths device-resident for CUDA
graph compatibility.
Eligibility is keyed on the `attention_metadata` query bound rather than
on the aggregate token count: a zero-heavy ragged step can have
`token_count <= batch_size` while still carrying a multi-token sequence,
and the metadata bound is the only replay-safe source under CUDA graph
capture.
### A ptxas miscompile in the packed causal mask
The mask builder was originally written as:
```cpp
const int allowed_bits = max(0, min(32, local_row + 1 - word * 32));
mask[i] = allowed_bits == 32 ? ~uint32_t{0}
: (allowed_bits == 0 ? 0u : ((uint32_t{1} << allowed_bits) - 1u));
```
ptxas (CUDA 13.0.48) folds the `min`/`max` into a single `VIMNMX.RELU`
and then reuses that instruction's clamp predicate for the `== 32` test
with inverted polarity:
```
VIMNMX.RELU R5, P1, P1, R4, 0x20, PT // R5 = max(0, min(x, 32)); P1 = clamp flag
ISETP.GT.AND P0, PT, R5, 0x1f, PT // (b >= 32) -> correct
SEL R9, RZ, 0x1, !P0 // b >= 32 -> correct
SEL R7, RZ, 0x1, !P1 // b == 32 -> folded into VIMNMX predicate, wrong polarity
```
So `clamped == 32` was true for every value, every mask word became
`0xffffffff`, and there was effectively **no intra-block causal mask** —
each draft token could attend to its own future. The emitted PTX is
correct (`setp.eq.s32 %p2, %r8, 32`); only the SASS is wrong. It
reproduces in ~30 lines of standalone CUDA at `-O0` through `-O3` on
both sm_80 and sm_90, and `>= 32`, `== 8` and `== 31` all compile
correctly — only equality against the clamp bound is affected.
The fix is to never compare a clamped value against its own bound.
Symptom before the fix, on Qwen3.8-27B DFlash2 with INT8 KV: MMLU-Pro
0.7913 vs 0.8250 and GPQA-D 0.6616 vs 0.7626.
The per-row error signature is what identified it: the **last** draft
token was exact and the error decayed monotonically toward it, with
magnitude proportional to `1 / past_len` — the signature of a fixed set
of extra keys leaking in, rather than an indexing error.
### Testing
- `onnxruntime_provider_test --gtest_filter='*PagedAttention*'` — 19/19
pass, including new parameterized native FP16 and BF16 speculative
dispatch cases.
- `pytest test_paged_attention.py -k "Speculative or Xqa or xqa or
Metadata"` — 68 pass. The 3 failures are a pre-existing harness
limitation (`TypeError: Got unsupported ScalarType Float8_e4m3fn` when
converting an FP8 torch tensor to numpy) and are unrelated to this
change.
- New `TestPagedAttentionXqaSpeculative` (21 cases): query lengths 2-8,
ragged batches including zero-length requests and `token_count <=
batch_size`, contexts 64-8192, block sizes 128/512, per-channel scales,
native FP16 cache, and an out-of-window bound that must fall back.
- `lintrunner` on all touched files.
These python cases use random K/V and a torch reference. That matters:
the pre-existing C++ fixture fills K/V with near-constant values, so
attending to 250 vs 256 nearly identical keys stays inside tolerance —
it could not have caught the mask bug.
### Quality gates (Qwen3.8-27B DFlash2, NVFP4 weights + INT8 KV,
thinking, MTP-7)
Paired per-question comparison against archived runs on the same harness
and settings.
| Task | This PR | vs XQA off | vs FP16 KV | vs pre-fix mask |
|---|---|---|---|---|
| MMLU-Pro (800) | 661/800 = 82.63% | 659, +2, p=0.81 | 663, -2, p=0.88
| 633, **+28, p=0.0026** |
| GPQA-Diamond (198) | 147/198 = 74.24% | 146, +1, p=1.00 | 148, -1,
p=1.00 | 131, **+16, p=0.0139** |
Statistically indistinguishable from both the XQA-off and FP16-KV
references (exact McNemar), and significantly better than the pre-fix
mask.
### Benchmark (H200, generation 256, speculative depth 7)
Decode throughput versus `ORT_ENABLE_XQA=0`, same build, same script,
same machine.
| INT8 KV cell | XQA | no XQA | speedup | token hash identical |
|---|---|---|---|---|
| ctx 512 / b1 | 156.6 | 156.4 | 1.001x | yes |
| ctx 2048 / b1 | 233.1 | 221.7 | 1.052x | yes |
| ctx 8192 / b1 | 161.2 | 131.3 | 1.228x | yes |
| ctx 8192 / b4 | 341.1 | 248.1 | 1.375x | yes |
| ctx 32768 / b1 | 178.5 | 92.4 | **1.931x** | yes |
INT8 geometric mean 1.155x over 9 cells; the gain grows with context, as
expected for an attention-bound kernel. The FP16-KV arm is flat at
1.001x geomean, which confirms the A/B carries no systematic bias. 7 of
9 INT8 cells are token-for-token identical to the no-XQA arm; the two
that differ also moved acceptance, i.e. trajectory divergence rather
than an engine regression.
### Notes for reviewers
- Native BF16 is compiled and dispatch-tested, but the python parity
harness is FP16-only, so BF16 has no discriminating numeric test. It
shares the FP16 code path with only the element type differing.
- The native kernel needs 165,504 bytes of shared memory at `M_TILESIZE
32` versus 132,736 for INT8 (H200 opt-in limit 232,448), so the existing
shared-memory fallback matters more on that path.
