[Tianlei Wu (tianleiwu)]

Description

This PR adds tensor-core CUDA kernels for the INT8-KV-cache DynamicSparseAttention of Qwen3.8-Flash-Next (12 full-attention layers, 24 query / 2 KV heads, head size 256): a prefill kernel and a split-K decode kernel. It also triggers the programmatic-dependent-launch (PDL) successor early in ProjectedSigmoidGateMulKernel and makes the sigmoid of HyperConnectionLowRankMix branch free.

It is the second half of #33689, which was split by operator: #33689 keeps the tensor-core MatMulNBits (fpA_intB) and QMoE / MoE GEMM kernels, and this PR carries the DynamicSparseAttention, gate-mul and hyper-connection changes. The two PRs are independent. The only file both add is cuda_session_options.h, and this PR's version is a superset of the one in #33689 (see "Related PRs and review guide").

Both INT8 tensor-core kernels are opt-in. ep.cuda.dsa_prefill_int8_tc / ORT_DSA_PREFILL_INT8_TC and ep.cuda.dsa_decode_int8_tc / ORT_DSA_DECODE_INT8_TC default to 0, which runs the previous kernels bit for bit, and the two switches are independent. The end-to-end accuracy and speed results were measured for an earlier revision of the kernels, before the numerics changes of the review, so they are preliminary. The kernels stay off by default until the final revision is qualified end to end (see "Accuracy").

At the operator level on an H200 (benchmark_dynamic_sparse_attention.py, added by this PR) the tensor-core decode kernel is 3.3x (1 row) to 5.6x (8 rows) faster than the CUDA-core kernel it replaces, and the prefill kernel is 2.7-3.0x faster than the BF16-dequantizing kernel. Against an FP64 reference the decode error stays within 3% of the previous kernel and the prefill error is 1.6-2.5 times smaller. A CTA whose inputs break the error bounds of the FP16 path is evaluated on a slower FP32 CUDA-core path, which can take longer than the previous kernels (see "Numerics"); the Qwen3.8-Flash-Next layers stay far from its conditions. In earlier-revision end-to-end runs the kernels lowered TTFT by 4.0-9.5% and the S = 8 decode cycle by 4.6-6.8% (preliminary), and the branch-free sigmoid lowered TTFT by 0.3-1.1%, which is about the size of the repeat noise.

Summary of Changes

1. DynamicSparseAttention: tensor-core prefill for the INT8 KV cache

File	Change
onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh (new)	DynamicSparseAttentionInt8TcKernel<kSplitK> (kSplitK = false for prefill): grid (rows, KV heads), 128 threads, 64-candidate tiles, GQA group <= 16 padded to the m16n8k16 M dimension. The gathered K/V rows stay INT8 in shared memory (cp.async, double buffered, 3-deep index ring) and the ldmatrix fragments are converted exactly to FP16 in registers, for FP16 mma with FP32 accumulation. key_scale is folded into the query once per CTA (Q' = fp16(Q * key_scale)) and value_scale is applied to the output once, after the accumulator has been divided by 2^15 (so a tiny value_scale cannot underflow). Scores and the running maximum stay in the natural-log domain, and only the difference to the new maximum is scaled by log2(e) for ex2. The FP16 probability tile holds p x 2^15 and the row sum stays unscaled FP32. A CTA that cannot bound the FP16 rounding evaluates its query heads in FP32 on the CUDA cores instead (AttendRowFp32; the decision is uniform per CTA and the FP32 path starts from the inputs, so nothing is clipped or rounded): an element of Q * key_scale is not finite or above 65504; for one head eps = scale * 128 * sum_d abs(Q'_d - fp16(Q'_d)), which bounds the shift of a logit by the rounding of Q', exceeds 0.125 nats (kInt8TcQueryErrorBudget); a value_scale is not finite; or some channel has abs(value_scale) > 2^16 and a valid candidate is more than 20 nats (kInt8TcTailDepth) below the running maximum (a probability below 2^-29 is subnormal in the p x 2^15 tile). The kernel is selected only for scale <= 1 (kInt8TcMaxScale; the op validates the scale as finite and non-negative, and the default 1 / sqrt(head_size) is 0.0625 here), because the rounding bound grows with the scale; a larger scale takes the BF16-dequantizing kernel.
onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu	Dispatch to the new kernel when its switch is on, the softmax scale is at most 1, the buffers are aligned and the loaded PTX target is at least SM80 (cudaFuncGetAttributes, otherwise an SM75 PTX image would run an empty kernel stub); DynamicSparseAttentionPrefillKernel<true> stays as the fallback.
onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.{h,cc}, dynamic_sparse_attention_impl.h	The four ep.cuda.dsa_* options are resolved once at kernel construction (session option, else environment variable, else default) and passed to the dispatch as parameters. The two INT8 tensor-core switches default to off.
onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc	The expected output of the INT8 cases is an FP64 reference computed from the inputs as the kernel receives them (BF16 query, INT8 codes, scales) with max-subtracted weights, rounded to BF16 once. Int8CachePrefillTensorCoreKernel_CUDA sweeps query heads {2, 12, 24, 32} x selected candidates {1, 64, 65, 509} x invalid indices with the tensor-core kernel on and off, on random INT8 rows with per-channel scales. A two-candidate fixture (two selected candidates of an otherwise empty cache, so that the output follows from a few numbers) checks the output to one BF16 ulp, each case once with both candidates in one 64-candidate tile and once with them in the first and the last of three tiles, which the decode kernel spreads over up to three splits (two of them active): Int8TensorCoreQueryFoldingRoundingError_CUDA (an in-range query whose FP16 rounding cancels the logit gap: 0.5 instead of 1), Int8TensorCoreLargeCommonLogit_CUDA (a large common logit, 8290560 at scale 1) and Int8TensorCoreTinyValueScale_CUDA (value_scale = 2^-135, a BF16 subnormal result). Int8TensorCoreRandomQueryFolding_CUDA uses random rows whose folded query products are not FP16 numbers, on both sides of the error budget, and Int8TensorCoreLargeValueScaleRandom_CUDA random rows with value scales of 2^14 to 2^17 and logit spreads of 0.7, 11 and 22 nats for the tail-depth check. The FP32 paths are also covered by Int8PrefillCancellingLargeQuery_CUDA and Int8PrefillQueryBeyondHalfRange_CUDA (four head counts and three candidate counts with invalid entries and the query scaled by 2^28, which sends every CTA to the FP32 path, plus a case where only some CTAs take it). PrefillSessionOptions_CUDA checks that options override the environment, that invalid values fail (the failure cases are not compiled with ORT_NO_EXCEPTIONS) and that only the strings 0 and 1 are accepted. Misaligned-buffer cases (8-byte K/V, 4-byte output) cover the fallback to the BF16-dequantizing kernel. Int8TensorCoreLargeSoftmaxScale_CUDA (scale = 2^30 with key_scale = 2^-30: the scale <= 1 gate) and Int8TensorCoreTinyProbabilityLargeValueScale_CUDA (value_scale = 2^28 with a candidate of probability about 2^-29: the probability scaling) run prefill and then decode on a two-candidate case with an analytic result.

2. DynamicSparseAttention: split-K tensor-core decode for the INT8 KV cache

File	Change
onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh	The kernel body of section 1 with kSplitK = true: the grid gets a third dimension, every CTA consumes a contiguous range of the row's 64-candidate tiles and writes an un-normalized partial result (running maximum in the natural-log domain, running sum, value-scaled accumulator) in the layout of the grouped decode kernel. It replaces the CUDA-core GroupedDecodeDynamicSparseAttentionKernel, which was issue bound on its per-candidate dot product, warp reduction and exponentials. Decode passes the probability tile as an FP16 high part plus an FP16 residual (about 22 significant bits), because the signed sum of large value codes can cancel to a value far below the rounding of one FP16 probability. A CTA that meets one of the conditions of section 1 evaluates its candidates in FP32 (AttendRowFp32<true>) and publishes an FP32 partial in the same layout, so the fallback holds for the splits too. A partial is written as (acc * 2^-15) * value_scale, so the merge kernel is unchanged.
onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu	GetInt8TcDecodeSplitCount: the smallest power-of-two number of 64-candidate tiles per split for which all ceil(max_selected / 64) * rows * kv_heads tiles fit one wave of resident CTAs (SM count x cudaOccupancyMaxActiveBlocksPerMultiprocessor of the kernel's shared-memory and register footprint: 396 on an H200). The split count is ceil(tiles / tiles_per_split), at most 32 and never more than the workspace was sized for. A positive ep.cuda.dsa_decode_int8_tc_tiles_per_split replaces the automatic value; it is a target, because the cap, the workspace size and the even distribution of a row's tiles still apply. The partial vectors come first in the workspace so that they stay 16-byte aligned. MergeGroupedDecodeDynamicSparseAttentionKernel combines the partials unchanged. Routing on top of UseGroupedDecodeAttention is listed under "Gates and switches"; everything else keeps the CUDA-core kernel.
onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc	Int8CacheDecodeTensorCoreKernel_CUDA: 1, 3 and 8 query rows x 8, 12, 24 and 32 query heads (group sizes 4, 6, 12 and 16, padded to the 16-row MMA tile) x selected candidates {1, 65, 509} around the 64-candidate tile x invalid indices, with the automatic split count, 1 and 3 tiles per split (a split without tiles publishes an empty partial) and the kernel switched off. Int8TensorCoreDecodeSignedValueCancellation_CUDA (probabilities 0.5 +- 2^-14 with values +-4064: a single FP16 tile would cancel them to 0) and Int8TensorCoreDecodeTilesPerSplitCaps_CUDA (lists of 2048, 2112 and 8192 candidates, that is 32, 33 and 128 tiles, with the automatic split count and 1, 2, 3 and 64 tiles per split, 1 and 8 rows: whatever number of splits results, every candidate reaches the merge; 8192 candidates with 1 tile per split get 32 splits of 4 tiles). Int8DecodeQueryBeyondHalfRange_CUDA multiplies every query by 2^28 so that every CTA takes the FP32 path for any split count, plus a case where only some CTAs do. DecodeSessionOptions_CUDA checks the switch, the tuning knob, the precedence over the environment and the rejection of invalid values. Int8CacheDecodeMisalignedBuffers_CUDA: 16-byte-aligned buffers use the tensor-core kernel, misaligned ones fall back to the CUDA-core kernel. The two scale tests of section 1 run the decode kernel as well.

3. Early PDL trigger in ProjectedSigmoidGateMulKernel

File	Change
onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu	ProjectedSigmoidGateMulKernel issues griddepcontrol.launch_dependents first (__CUDA_ARCH__ >= 900).

The kernel feeds the hyper-connection down kernel, which is launched with PDL and reads the gate-mul output only after griddepcontrol.wait; its constant-weight prefetch now overlaps with the gate-mul kernel. There is no effect for plain launches, no change of any result and no switch. The same trigger in the fpA_intB GEMV kernels is part of #33689; together they shorten an S = 8 decode cycle by 1.4-1.7% (preliminary, see "Performance").

4. HyperConnectionLowRankMix: branch-free sigmoid

File	Change
onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu	SigmoidF32 evaluates e = expf(-fabsf(v)) once and selects the numerator, (v >= 0 ? 1 : e) / (1 + e), instead of branching between 1 / (1 + expf(-v)) and expf(v) / (1 + expf(v)). Both branches already pass -fabsf(v) to expf, so the result is bit-identical. In the prefill pre-mix neighbouring elements have different signs, so nearly every warp executed both sides of the old branch.

5. Session options (shared by 1 and 2)

File	Change
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h	Four ep.cuda.* keys: dsa_prefill_int8_tc, dsa_prefill_kernel, dsa_decode_int8_tc and dsa_decode_int8_tc_tiles_per_split. The descriptions state the opt-in defaults and the exact "0" / "1" grammar, and that tiles per split is a target subject to the split cap, the workspace size and the balancing of the row's tiles.
onnxruntime/contrib_ops/cuda/cuda_session_options.h (new)	ReadCudaSessionOption (integer threshold) and ReadCudaSessionSwitch (exactly "0" or "1": "00", "01", "-0" and "+1" are rejected): an explicit session option wins, otherwise the environment variable (an empty variable counts as unset), otherwise the default. The whole string is parsed and range-checked, and an invalid value fails kernel construction. #33689 adds the same file for its own switches (see "Related PRs and review guide").

6. Documentation and benchmark

File	Change
docs/contrib_ops/cuda/dynamic_sparse_attention.md	The ep.cuda.dsa_* options next to their environment variables with the opt-in defaults, the dispatch-table rows and alignment fallbacks of the INT8-cache prefill and split-K decode kernels, the tiles-per-split rule with its sweep and regret statistics, "Numerics and fallbacks" (the conditions of the FP32 path, what the model does, the measured cost of the path), "Rollout" and the operator-level benchmark.
docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md	The early griddepcontrol.launch_dependents trigger of the kernel in front of HcDownKernel and the branch-free SigmoidF32.
onnxruntime/test/python/transformers/benchmark_dynamic_sparse_attention.py (new)	Operator-level benchmark: a single-node DynamicSparseAttention model per case (the shapes of the 12 layers; --rc-model reads the scales of a model's node), one session per kernel configuration, the GPU time of the kernels of one Run from the CUPTI-based torch profiler, alternating rounds, and the RMS error against an FP64 reference. --query-std, --key-scale-multiplier and --value-scale-multiplier reach the FP32 paths, and --tiles-per-split sweeps the decode split rule.

Gates and switches

Every switch is available as an environment variable and as a CUDA session option (ep.cuda.*, declared in onnxruntime_session_options_config_keys.h). It is resolved once when the kernel is created, so a change needs a new session, and an explicit session option takes precedence over the environment variable. The 0/1 switches accept exactly the strings 0 and 1, the tuning knob a non-negative decimal integer; any other value (including an empty session-option string) fails session initialization. ORT_DSA_PREFILL_KERNEL already exists on main (read on every call; any non-zero integer enabled it) and now follows the same rules.

Kernel	Used when	Environment variable / session option
INT8-KV DSA prefill	BF16 query / output, INT8 main cache, the prefill kernel path (sequence_length >= 8, not the grouped decode path), head_size = 256, at most 16 query heads per KV head, selected-only from the main cache (no local_plus_selected), no smooth softmax, scale <= 1, max_selected > 0, SM80+, 16-byte-aligned present_key, present_value and scales, 8-byte-aligned output, enough opt-in shared memory (72,960 B)	ORT_DSA_PREFILL_INT8_TC=1 / ep.cuda.dsa_prefill_int8_tc turns it on (default 0: the BF16-dequantizing kernel, which has no FP16 range limit); ORT_DSA_PREFILL_KERNEL=0 / ep.cuda.dsa_prefill_kernel turns off the whole tensor-core prefill family (default 1)
INT8-KV DSA decode	BF16, INT8 main cache, grouped decode path (1 <= sequence_length <= 8, 4 to 16 query heads per KV head, max_selected > 256), head_size = 256, selected-only from the main cache, no smooth softmax, scale <= 1, 16-byte-aligned query, present_key, present_value, scales and attention workspace, SM80+, enough opt-in shared memory (75,264 B)	ORT_DSA_DECODE_INT8_TC=1 / ep.cuda.dsa_decode_int8_tc turns it on (default 0: the CUDA-core grouped decode kernel, bit for bit as before); ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT / ep.cuda.dsa_decode_int8_tc_tiles_per_split (default 0 = automatic; a requested tiles-per-split target, see section 2)
Early PDL trigger	SM 9.0+, PDL successor	none (no effect for plain launches)
Branch-free sigmoid	always	none (bit-identical)

The tensor-core kernels also require a loaded PTX target of at least SM80. If only older PTX is available (for example a compute_75 image that the driver JIT-compiles for a newer GPU), their __CUDA_ARCH__ >= 800 bodies are empty stubs that would silently leave the output unwritten, so the dispatch checks cudaFuncGetAttributes(...).ptxVersion >= 80 and otherwise uses the generic kernels.

Numerics

• The branch-free sigmoid is bit-identical to the branchy one for all 2^32 float bit patterns (exhaustive check in a standalone harness, same nvcc flags, no fast-math).
• The early PDL trigger changes no result (generated tokens identical in 45 / 45 decode runs).
• The INT8 tensor-core kernels are close to, but not bit-identical with, the previous kernels: they round Q' and the probability tile to FP16 and sum in a different FP32 order. The switches' default 0 keeps the previous kernels bit for bit. Against an FP64 reference (the BF16 rounding of the output alone gives about 1.5e-3 relative RMS), decode is at 1.6-1.7e-3 with both kernels (the tensor-core kernel is up to 3% higher), and prefill at 1.7e-3 against 2.7-4.3e-3 for the BF16-dequantizing kernel (tables under "Performance").
• The cases found in review, and how the kernels handle them:

Case	Handling	Test
A finite Q * key_scale whose FP16 rounding changes the logits (products 32768 and -32766 round to 32768 and -32768, the gap of 16 becomes 0: 0.5 instead of 1)	The per-head bound eps above 0.125 nats sends the CTA to the FP32 path	Int8TensorCoreQueryFoldingRoundingError_CUDA, Int8TensorCoreRandomQueryFolding_CUDA
Decode: two probabilities 0.5 +- 2^-14 round to the same FP16 value and the signed values cancel (the output 127/256 = 0.496 became 0)	The decode tile is an FP16 high part plus an FP16 residual (prefill keeps one tile, as the BF16 kernel does with a coarser one)	Int8TensorCoreDecodeSignedValueCancellation_CUDA
Scores of 8290560 and 8290560.5 (a large common logit) converted to the log2 domain before the maximum is subtracted lose their gap to float rounding (0.5 instead of 0.62 at scale 1 and 0.508 at scale 1/16)	Natural-log scores and maximum; only the difference is scaled by log2(e)	Int8TensorCoreLargeCommonLogit_CUDA
value_scale = 2^-135: value_scale * 2^-15 underflows before the accumulator is applied (BF16 subnormal result became 0)	Divide the accumulator by 2^15 first, then multiply by value_scale	Int8TensorCoreTinyValueScale_CUDA
abs(value_scale) > 2^16 with a probability below 2^-29 (subnormal in the FP16 tile)	The tail-depth check decides per CTA, instead of sending every CTA to the FP32 path as soon as one channel exceeds 2^16	Int8TensorCoreLargeValueScaleRandom_CUDA

• Cost of the FP32 path (H200, benchmark_dynamic_sparse_attention.py; 24 query / 2 KV heads, 2048 selected candidates, context 8192; attention-kernel time, previous kernel against tensor-core kernel):

Data	Rows	Previous	Tensor core	Tensor core vs previous
Queries like the model's	128 / 8192	0.173 / 7.82 ms	0.063 / 2.83 ms	2.7x / 2.8x faster
Queries 8 times larger (--query-std 11.2): every CTA is over the error budget	128 / 8192 / 8 (decode)	0.172 ms / 7.82 ms / 73.6 us	3.68 ms / 135.6 ms / 240 us	21x / 17x / 3.3x slower
abs(value_scale) 24k to 360k (--value-scale-multiplier 1.2e7), logit standard deviation about 2.4 nats: 10 to 12% of the CTAs have a candidate deeper than 20 nats	128 / 512 / 2048	0.172 / 0.577 / 2.288 ms	3.61 / 3.90 / 6.17 ms	21x / 6.8x / 2.7x slower
The same value scales, logit standard deviation about 1.3 nats: no CTA is deeper than 20 nats	8192 / 8 (decode)	7.78 ms / 73.6 us	2.86 ms / 13.3 us	2.7x / 5.6x faster

A CTA on the FP32 path needs about 3.6 ms for 2048 candidates, so the time of a launch with one wave of CTAs is set by its slowest CTA. The fallback keeps the accuracy (relative RMS 1.7e-3 against FP64 with queries 8 times larger; the previous prefill kernel: 8.5e-3). The first group of rows is the Qwen3.8-Flash-Next case: the largest value_scale of its 12 layers is 0.285 (the tail-depth check is never armed) and a Monte Carlo over their q_norm weights and key scales gives a per-head eps of 0.015 on average and 0.022 at most, against the budget of 0.125, which corresponds to queries about 8 times larger than the model produces. Workloads with such queries, or with abs(value_scale) > 2^16 and wide logit spreads, should leave the switches at 0. The FP32 path of a CTA is documented in the new "Numerics and fallbacks" section of docs/contrib_ops/cuda/dynamic_sparse_attention.md.

• Cost of the numerics fixes on the fast path (H200, the same script; separate processes against the previous head of this PR, ceed60271c, not paired; attention-kernel time): prefill at 128 / 2048 / 8192 rows 0.062 / 0.815 / 2.760 ms against 0.063 / 0.832 / 2.830 ms (+1.6% to +2.5%); decode with 8 rows 12.23 us against 13.2-13.3 us (+8-9%) and with 1 row 8.84 us against 9.46 us (+7%). The previous head sent every CTA to the FP32 path when one channel had abs(value_scale) > 2^16 (8192 rows: 2.76 ms became 136 ms).

Testing

In-tree CUDA 13.0 Release build, sm_90a, one H200 (SM 9.0), on this branch (onnxruntime_provider_test; this checkout needs -Donnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=OFF because gemm_tactic_cache_test.cc does not link otherwise (undefined reference to onnxruntime::g_host), which is unrelated to this PR):

Command (--gtest_filter=)	Result
*DynamicSparseAttention*:*HyperConnection*	103 run, 103 passed, 0 failed (73 DynamicSparseAttentionTest, 13 HyperConnectionLowRankMixTest, 17 HyperConnectionOpsTest); about 100 s
*GatedAdd*:*ProjectedSigmoid*:*SigmoidGate*:*GateMul*	14 run, 14 passed, 0 failed (12 ContribOpGatedAddTest, including ProjectedGateBFloat16 and ProjectedGateK1Regression, and 2 EngramOpsTest)

• The review round adds 7 DSA tests (Int8TensorCoreQueryFoldingRoundingError_CUDA, Int8TensorCoreDecodeSignedValueCancellation_CUDA, Int8TensorCoreLargeCommonLogit_CUDA, Int8TensorCoreTinyValueScale_CUDA, Int8TensorCoreDecodeTilesPerSplitCaps_CUDA, Int8TensorCoreLargeValueScaleRandom_CUDA, Int8TensorCoreRandomQueryFolding_CUDA), moves the INT8 reference to FP64 with max-subtracted weights, and tightens the two session-option tests ("00", "01", "-0", "000001", "+1", " 1", "1 ", "", "2", "-1", "1x" and "true" are rejected for the switches, and "01", "-0" and "+1" for the environment variables).
• Red-green against the previous head of this PR (ceed60271c): the final test binary was run against the provider built from that head (ORT_UNIT_TEST_MAIN_DYNAMIC_PLUGIN_EP_CONFIG_JSON selects the library) for the 13 tests that concern the review. 7 fail on it: the five numerics tests Int8TensorCoreQueryFoldingRoundingError_CUDA, Int8TensorCoreDecodeSignedValueCancellation_CUDA, Int8TensorCoreLargeCommonLogit_CUDA, Int8TensorCoreTinyValueScale_CUDA and Int8TensorCoreRandomQueryFolding_CUDA, and PrefillSessionOptions_CUDA and DecodeSessionOptions_CUDA (the stricter parser). The other 6 pass on it: Int8TensorCoreDecodeTilesPerSplitCaps_CUDA and Int8TensorCoreLargeValueScaleRandom_CUDA are coverage tests (the split caps; the tail-depth check against FP64 on random data), and Int8TensorCoreTinyProbabilityLargeValueScale_CUDA, Int8TensorCoreLargeSoftmaxScale_CUDA, Int8PrefillQueryBeyondHalfRange_CUDA and Int8DecodeQueryBeyondHalfRange_CUDA are the guards of the earlier rounds. All 73 pass with the final provider.
• Under nsys the session-option tests launch the intended kernels on the final tree. PrefillSessionOptions_CUDA: the split / merge kernels for dsa_prefill_kernel=0, DynamicSparseAttentionPrefillKernel<true> for dsa_prefill_int8_tc=0 and DynamicSparseAttentionInt8TcKernel<false> for =1 (two launches of each). DecodeSessionOptions_CUDA: GroupedDecodeDynamicSparseAttentionKernel for dsa_decode_int8_tc=0 and DynamicSparseAttentionInt8TcKernel<true> for =1 (three launches of each), each followed by MergeGroupedDecodeDynamicSparseAttentionKernel.
• lintrunner (clang-format, ruff, ruff-format) is clean on all 13 changed files. dynamic_sparse_attention_impl.cu, gated_add_impl.cu and hyper_connection_low_rank_mix_impl.cu also compile for sm_75, sm_80, sm_86 and sm_120 with the flags of the build.
• The PTX guards (cudaFuncGetAttributes(...).ptxVersion >= 80) have no unit test, because they need an executable that carries only compute_75 PTX. On the H200 (the driver JIT-compiles the PTX for SM 9.0) such an executable reports ptxVersion = 75 for DynamicSparseAttentionPrefillKernel<true> and <false>, DynamicSparseAttentionInt8TcKernel<false> and DynamicSparseAttentionInt8TcKernel<true>, so dispatch rejects them; built for sm_90a the same kernels report ptxVersion = 90.
• CI: the new DSA tests need a BF16-capable device (SM80+) and skip otherwise. They need 72,960 B (prefill) and 75,264 B (decode) of opt-in shared memory for the tensor-core kernels, which sm_86 has, so they run the tensor-core paths there.

Checklist

• Tests added/updated and passing
• No breaking API changes. By default the INT8-KV attention behaves as before (both tensor-core kernels are opt-in; the CUDA-core decode kernel and the BF16-dequantizing prefill kernel run bit for bit as before). Behavior changes: ORT_DSA_PREFILL_KERNEL is now resolved once per kernel and strictly 0/1, and the tensor-core kernels, when enabled, are not bit-identical to the previous kernels
• Documentation updated (if applicable)

Motivation and Context

At 8192 tokens per prefill chunk the INT8-KV prefill kernel took 12-13% of the chunk with HMMA at only 2.8% of its instructions: it dequantized every gathered K/V tile to BF16 in shared memory before the BF16 MMAs. In decode (S = 8, 12 query heads per KV head) the CUDA-core INT8 kernel took 70.3 us plus 5.3 us for the merge per layer at 2K to 32K context, 6.9% of a decode cycle (ncu: issue and shared-load bound on the per-candidate serial chain). In the prefill pre-mix of the hyper-connections, the data-dependent branch of SigmoidF32 made nearly every warp execute both of its paths.

Performance on NVIDIA H200 (SM 9.0a, CUDA 13.0)

Operator level (final revision)

benchmark_dynamic_sparse_attention.py builds a single-node DynamicSparseAttention model per case (24 query / 2 KV heads, head size 256, INT8 cache, scale 1/16, 2048 selected candidates per row: the shapes of the 12 layers) and one session per kernel configuration, because the ep.cuda.dsa_* options are read when the kernel is created. The time is the GPU time of the kernels that one Run launches (CUPTI-based torch profiler, so host launch overhead does not enter): "attention" sums the DynamicSparseAttention kernels including the decode merge kernel, "node" all GPU activity of the node (cache append, query preparation, ...). Each round measures every configuration once, the order alternates between rounds, the tables show the median of 7 rounds, and the speedup is the median of the paired ratios of the rounds. 3 warmup runs per measurement, 200 profiled runs per round for decode and 5 for prefill. The error is the RMS error of the BF16 output relative to the RMS of an FP64 reference over 64 evenly spaced rows. The data are synthetic: BF16 queries with standard deviation 1.4 (the queries of the model after q_norm have a per-head RMS of 1.3 to 1.6), INT8 cache codes with standard deviation 40, and key and value scales drawn from the ranges of that model. "Previous" is the kernel that the switch replaces (the CUDA-core grouped decode, the BF16-dequantizing prefill):

python benchmark_dynamic_sparse_attention.py --rows 1 8 --context 2048 8192 32768 --tiles-per-split 0 1 2 4 8 16
python benchmark_dynamic_sparse_attention.py --rows 128 2048 8192 --context 8192 32768

Decode (sequence_length <= 8):

Rows	Context	Attention, previous	Attention, tensor core	Speedup	Node, previous	Node, tensor core	RMS error, previous / tensor core
1	2048	31.1 us	9.5 us	3.30x	34.1 us	12.5 us	1.59e-3 / 1.62e-3
1	8192	31.2 us	9.5 us	3.30x	34.2 us	12.5 us	1.55e-3 / 1.57e-3
1	32768	31.2 us	9.5 us	3.30x	34.1 us	12.5 us	1.67e-3 / 1.72e-3
8	2048	73.6 us	13.2 us	5.58x	76.7 us	16.5 us	1.64e-3 / 1.69e-3
8	8192	73.6 us	13.3 us	5.57x	76.8 us	16.6 us	1.62e-3 / 1.66e-3
8	32768	73.7 us	13.2 us	5.57x	77.0 us	16.6 us	1.68e-3 / 1.73e-3

Prefill:

Rows	Context	Attention, previous	Attention, tensor core	Speedup	Node, previous	Node, tensor core	RMS error, previous / tensor core
128	8192	0.173 ms	0.063 ms	2.72x	0.179 ms	0.070 ms	4.19e-3 / 1.71e-3
128	32768	0.191 ms	0.064 ms	2.99x	0.199 ms	0.072 ms	4.25e-3 / 1.70e-3
2048	8192	2.287 ms	0.832 ms	2.75x	2.350 ms	0.894 ms	4.25e-3 / 1.71e-3
2048	32768	2.294 ms	0.832 ms	2.75x	2.356 ms	0.895 ms	4.21e-3 / 1.70e-3
8192	8192	7.818 ms	2.830 ms	2.76x	8.057 ms	3.069 ms	2.69e-3 / 1.66e-3
8192	32768	8.877 ms	3.200 ms	2.77x	9.116 ms	3.439 ms	4.29e-3 / 1.71e-3

The decode time does not depend on the context length (the selected candidates are gathered). The error includes the rounding of the BF16 output, whose floor is about 1.5e-3 relative to the RMS: both decode kernels and the tensor-core prefill are near it, and the BF16-dequantizing prefill, which rounds the dequantized keys and values and the probabilities to BF16, has 1.6 to 2.5 times the error of the tensor-core prefill.

Decode time against the tiles per split (8 rows, 2048 candidates, kernel plus merge, --tiles-per-split):

Tiles per split	1	2 (automatic)	4	8	16
Splits, CTAs	32, 512	16, 256	8, 128	4, 64	2, 32
Time	18.8 us	13.2 us	12.7 us	17.7 us	29.2 us

The automatic rule is a heuristic, not the optimum. Over 20 shapes on the H200 (1, 2, 4 and 8 rows, 512 to 8192 candidates, context 8192, each swept over 1, 2, 3, 4, 6, 8 and 16 tiles per split) it is on average 1.8% and at most 14.8% slower than the best swept value (4 rows and 2048 candidates: 11.9 us against 10.4 us); the previous target of about 384 CTAs is 2.7% and 14.2% slower on the same points. It matters when fewer CTAs are resident: with the kernel's shared memory padded on the H200 so that one or two CTAs fit an SM (132 or 264 resident CTAs; an ad hoc harness that is not part of this change) the rule is on average 0.4% (at most 2.4%) and 3.1% (13.4%) slower than the best swept value, where the 384-CTA target is 23.8% (44.9%) and 14.4% (38.4%) slower. An A100 (108 SMs, two resident CTAs per SM by the 75,264 B of shared memory) is expected to choose 4 tiles per split, 8 splits, for 8 rows and 2048 candidates, the faster setting in the A100 measurement reported in the review (22.1 us against 25.9 us at 16 splits). That choice is derived from the rule, not measured: the A100 has not been run with this script. ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT tunes other devices.

The other kernels of this PR, per launch (H200, nsys / CUDA events; unchanged since the earlier revision):

Kernel	Before	After
PreMixVec8Kernel per launch (hidden 2560, 4 branches, 8192-token chunk; 96 launches per chunk)	112.1 us	82.7 us
PreMixVec8Kernel, the same launch in the model (nsys, 8192-token chunk, all_int4)	161.2 us	119.2 us
ScaledSiluKernel	7.3 us	6.1 us

End to end (preliminary)

These end-to-end numbers were measured for an earlier revision of the kernels, before the numerics changes of the review, and have not been reproduced for the final revision: the stack that they need (the model packages and runtime, and #33689 together with this PR on top of a main that contains #32126, #33650 and #33714) is not available for a final-head run. Each row is an A/B of consecutive providers on the same packages, one round unless noted (not alternated), without raw timings and without confidence intervals. They show the direction of the effect, not a qualified gain. A shorter cycle is not the same as better user-visible throughput, and accepted tokens per cycle and time per accepted output token were not measured for any of these runs.

Measured on the Qwen3.8-Flash-Next packages (hidden = 2560, intermediate = 640, 512 experts, top-10; all_int4 is the all-INT4 package, top16 and top08 are the two INT3-mixed packages), B = 1, DFlash2 S = 8 decode. Decode is in ms per cycle.

TTFT of the INT8-KV prefill kernel in seconds (one round; decode is unchanged):

Package	2048	8192	32768
all_int4	0.1985 -> 0.1889 (-4.8%)	0.7122 -> 0.6524 (-8.4%)	3.0210 -> 2.7782 (-8.0%)
top16	0.2155 -> 0.2068 (-4.0%)	0.7339 -> 0.6714 (-8.5%)	3.0969 -> 2.8608 (-7.6%)
top08	0.2216 -> 0.2121 (-4.3%)	0.7478 -> 0.6768 (-9.5%)	3.1504 -> 2.8540 (-9.4%)

Decode of the split-K INT8-KV kernel, ms per cycle at 2K / 8K / 32K context:

Package	2K	8K	32K
all_int4	11.672 -> 11.017 (-5.6%)	11.962 -> 11.143 (-6.8%)	12.157 -> 11.391 (-6.3%)
top16	12.148 -> 11.381 (-6.3%)	12.353 -> 11.608 (-6.0%)	12.629 -> 11.953 (-5.4%)
top08	12.216 -> 11.505 (-5.8%)	12.474 -> 11.675 (-6.4%)	12.661 -> 12.082 (-4.6%)

Branch-free sigmoid, TTFT in ms (median of 5 repetitions; 9 / 9 cells faster, first tokens identical). Decode is within +-0.3%, because the fused decode kernels call the function for at most 8 rows. The per-cell gains of 0.3-1.1% are about the size of the repeat noise between repeats of one stack (about 1%); the kernel-level PreMixVec8Kernel time above is the stronger evidence:

Package	2048	8192	32768
all_int4	174.1 -> 173.3 (-0.45%)	574.7 -> 572.9 (-0.31%)	2450.1 -> 2439.2 (-0.45%)
top16	191.3 -> 190.5 (-0.42%)	596.3 -> 592.1 (-0.70%)	2560.2 -> 2531.7 (-1.11%)
top08	196.2 -> 194.7 (-0.79%)	601.1 -> 597.2 (-0.65%)	2573.5 -> 2554.2 (-0.75%)

The early PDL trigger of the gate-mul kernel and the one of the fpA_intB GEMV kernels (#33689) together shorten a cycle by 1.4-1.7% (a joint result of two PRs). For the whole stack that contains both PRs (plus #32126, #33650 and the package configuration), the median over 4 rounds against the previous release candidate is -32.8 / -36.0 / -36.9% TTFT at 8192 tokens, -32.3 / -35.2 / -36.1% at 32768 tokens, and -27.6 / -24.6 / -24.2% decode ms per cycle at 2K (all_int4 / top16 / top08); these combined-stack numbers cannot be attributed to a single PR.

Accuracy

The sigmoid change is bit-exact and the early PDL trigger changes no result.

At the operator level (final revision) the error against an FP64 reference is in the tables above: decode 1.6-1.7e-3 with both kernels (the tensor-core kernel is up to 3% higher) and prefill 1.7e-3 against 2.7-4.3e-3 for the BF16-dequantizing kernel, with the BF16 rounding of the output alone at about 1.5e-3. The unit tests compare against an FP64 reference with max-subtracted weights and fix the numerical corners found in review to one BF16 ulp.

End to end, the final revision has not been qualified, and that is why both tensor-core kernels are opt-in. The one end-to-end accuracy gate is for an earlier revision. The split-K decode kernel changes the generated tokens (a different FP32 accumulation order; for 263 of 800 all_int4 MMLU-Pro prompts), so it was gated with MMLU-Pro: 800 questions x 3 packages = 2400 items, greedy decoding through the full engine, max_new = 8192. Together with the dequantize + cuBLAS route of #33689, it gave 668 / 678 / 667 correct answers against 665 / 676 / 670 for the previous stack, pooled +0.08 pp (95% CI -0.60 .. +0.77 pp, exact McNemar p = 0.91); no package differs at the 5% level. The first-round stack, which contains the INT8-KV prefill kernel, was -0.67 pp pooled against the previous stack (95% CI -1.36 .. +0.03 pp, p = 0.076). These runs measure the accumulated effect of all kernels of the stack, not of one PR, and a drift of up to about 1.4 pp cannot be excluded with 2400 items; a non-significant difference does not establish non-inferiority. They predate the numerics of the final revision (the natural-log domain, the Q' rounding bound, the decode probability residual, the value-scale ordering and the tail-depth check), and so does a later check of commit 8 (engine_behaviours.py 27 / 27 on the three packages; TTFT and decode within the noise against the stack before commit 8). They do not cover the final revision.

Making a switch the default needs the final head qualified with the rest of the stack held fixed and prefill and decode toggled separately: per-package paired accuracy with confidence intervals against a stated degradation bound, including long contexts, exact build and model identities, and for speculative decode the accepted tokens per cycle and the time per accepted token. Packages that qualify the kernels enable them in their session options (ep.cuda.dsa_prefill_int8_tc and ep.cuda.dsa_decode_int8_tc set to 1).

Related PRs and review guide

• [CUDA] Tensor-core QMoE INT4/INT3 decode, INT8 MatMulNBits GEMV and large-M dequantize + cuBLAS GEMM #33689 is the sibling PR (tensor-core MatMulNBits / fpA_intB and QMoE / MoE GEMM kernels). The changes are divided by operator and have no kernel in common. Both PRs add cuda_session_options.h, and the two versions now differ: this PR's is a superset. It adds ResolveCudaSessionOptionText, ResolveCudaSessionSwitch (exactly "0" or "1") and ReadCudaSessionSwitch, and ReadCudaSessionOption (the integer parser, with unchanged behavior) now calls the text helper. A trial merge of the two heads (git merge-tree) has this add/add conflict in that one file and no other, and the two edits of onnxruntime_session_options_config_keys.h are in different places. Whichever PR lands second takes this PR's version of the file; the switches of [CUDA] Tensor-core QMoE INT4/INT3 decode, INT8 MatMulNBits GEMV and large-M dequantize + cuBLAS GEMM #33689 can then adopt ReadCudaSessionSwitch, which is a change in [CUDA] Tensor-core QMoE INT4/INT3 decode, INT8 MatMulNBits GEMV and large-M dequantize + cuBLAS GEMM #33689.
• [CUDA] Add SwiGLU contrib operator and SwiGLU fusion #32126 (SwiGLU operator and fusion), [CUDA] Add RouterGateUpProj contrib operator #33650 (RouterGateUpProj) and [CUDA] QMoE: select the 64-row CTA tile for repacked INT3 prefill when experts average >= 24 rows #33714 (QMoE INT3 prefill tile) have merged into main since this PR was opened. This branch does not touch their files, and a trial merge of this head with the current main (git merge-tree) is clean. The decode graph measured above contains all three.
• The PR has 13 files and 14 commits: commit 1 (early PDL trigger of the gate-mul kernel), commit 2 (INT8-KV prefill kernel), commit 3 (review follow-up: the FP32 fallback for out-of-range queries, the ep.cuda.dsa_* session options with their tests), commit 4 (fixes MSVC narrowing warnings in the DSA tests), commit 5 (split-K decode kernel and its options and tests), commit 6 (branch-free sigmoid), commit 7 (docs), commit 8 (review follow-up: the scale <= 1 condition, the 2^15 probability scaling, ORT_NO_EXCEPTIONS guards in the option tests and two regression tests), commit 9 (ceed60271c, value-scale gate with an FP32 fallback for all CTAs, replaced by the tail-depth check of commit 10), commit 10 (1bf2eb17cb, review fixes of the kernels, the split count, the options and the opt-in defaults), commit 11 (bd24bb8503, tests), commit 12 (2dedd0c4ff, docs and the benchmark script), commit 13 (e8b3ad2232, the end-to-end timings in the hyper-connection doc are labeled preliminary) and commit 14 (d60b0f4b63, the DSA doc states that the end-to-end timings have no acceptance measurements). dynamic_sparse_attention_prefill_int8.cuh is the main file to review (812 lines).
• The review threads of [CUDA] Tensor-core QMoE INT4/INT3 decode, INT8 MatMulNBits GEMV and large-M dequantize + cuBLAS GEMM #33689 on the prefill kernel and its test (clamping of the scaled query, identical K/V rows in the sweep) were answered there with the fixes of commit 3; its other threads concern the QMoE and dense INT8 kernels that stay in [CUDA] Tensor-core QMoE INT4/INT3 decode, INT8 MatMulNBits GEMV and large-M dequantize + cuBLAS GEMM #33689. The review threads of this PR on the numerics (the rounding of the folded query, the decode probability rounding, the log2 domain, the tiny value scale and the value-scale cliff), the tests (FP64 reference, active splits, exact checks), the split target, the switch grammar, the tiles-per-split wording and the rollout and end-to-end claims are answered by commits 10 to 14.

[Copilot]

🟡 Changes recommended

FP16 underflow can corrupt valid attention results in both prefill and decode, while invalid-option tests can abort the test process in exception-disabled CUDA builds.

3 open issues

• FP16 underflow corrupts attention results · New
• FP16 probability rounding erases scaled value contributions · New
• Guard CUDA invalid-option tests in no-exceptions builds · New

What changed in this PR

This PR accelerates CUDA DynamicSparseAttention with INT8-cache tensor-core prefill and split-K decode kernels, alongside targeted gate-mul and hyper-connection optimizations. It contains the attention-focused changes split from #33689.

Changes:

• Adds tensor-core attention kernels with dispatch checks, FP32 fallback for oversized queries, and configurable kernel switches.
• Enables an early programmatic-dependent-launch trigger and replaces the hyper-connection sigmoid with a branch-free expression.
• Adds numerical, configuration, and alignment tests, plus dispatch and performance documentation.

File	Description
onnxruntime/​test/​contrib_ops/​dynamic_sparse_attention_op_test.cc	Adds kernel parity, range, session-option, and misalignment coverage.
onnxruntime/​contrib_ops/​cuda/​cuda_session_options.h	Adds validated session-option parsing with environment fallback.
onnxruntime/​contrib_ops/​cuda/​bert/​hyper_connection_low_rank_mix_impl.cu	Makes the sigmoid calculation branch-free.
onnxruntime/​contrib_ops/​cuda/​bert/​gated_add_impl.cu	Triggers dependent kernels early on SM90+.
onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention.h	Stores per-kernel dispatch settings.
onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention.cc	Resolves settings during construction and passes them to dispatch.
onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention_prefill_int8.cuh	Implements shared INT8-cache prefill and split-K decode kernels.
onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention_impl.h	Adds dispatch configuration fields.
onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention_impl.cu	Adds eligibility checks, split selection, workspace layout, and launches.
include/​onnxruntime/​core/​session/​onnxruntime_session_options_config_keys.h	Declares four CUDA attention configuration keys.
docs/​contrib_ops/​cuda/​hyper_connection_low_rank_mix.md	Documents launch overlap and sigmoid optimization.
docs/​contrib_ops/​cuda/​dynamic_sparse_attention.md	Documents new kernels, switches, fallbacks, and performance.

_{Review effort: Balanced}

[Akshay Sonawane (apsonawane)]

Follow-up review at 204d6593c0eb57f0a4fc9f8bcec19a3a2d2c9e62

The existing inline threads are resolved. The current source includes the scale <= 1 dispatch restriction, the 2^15 probability scaling, and the exception-disabled test guards. However, the probability change moves the underflow boundary rather than eliminating the reported class of error.

Remaining correctness issue: finite value scales can still amplify probabilities discarded by FP16

The probability conversion rounds p * 32768 to FP16. Both prefill and decode eligibility permit this without a value-scale numerical eligibility check.

A direct extension of the previous tiny-probability regression:

• Head size 256; one nonzero prepared-query channel equal to 480.
• Unit key scales; two selected keys with that channel {0, -1}, in the same candidate tile/split.
• Softmax scale 1/16, so the relative scores are {0, -30}.
• Value codes {0, 1} in the observed output channel and finite FP32 value_scale = 2^60.

Here p = exp(-30) ~= 9.357623e-14, and FP16(p * 2^15) = 0. The reference output is

2^60 * exp(-30) / (1 + exp(-30)) ~= 107886.0475

which is finite and representable in BF16, whereas the tensor-core PV calculation loses the only nonzero value contribution. The query itself fits FP16, and the softmax scale is below the new dispatch limit, so neither existing safeguard handles this case.

Please preserve these contributions with range-safe arithmetic or a fallback based on numerical eligibility. Increasing a fixed probability multiplier alone cannot cover the unrestricted finite value-scale domain. Extend the current tiny-probability regression for both prefill and split-K decode, keeping the two valid keys in the same tile/split, and compare TC-enabled versus the range-safe reference with a nonzero expected result.

Evidence/limits: I reproduced the FP16 rounding and expected output with a standalone host arithmetic check and inspected the current dispatch/conversion source. I did not execute this case on CUDA or run a provider build; this is a source/host-arithmetic finding, not a GPU reproduction or an approval.

[Tianlei Wu (tianleiwu)]

Addressed in ceed60271c:

• In DynamicSparseAttentionInt8TcKernel (dynamic_sparse_attention_prefill_int8.cuh), the CTA prologue now also checks !(fabsf(value_scale[d]) <= kInt8TcMaxValueScale) (kInt8TcMaxValueScale = 65536.0f = 2^16) alongside the Q * key_scale range check and routes the CTA through AttendRowFp32<kSplitK> (for both prefill and split-K decode) whenever any channel exceeds 2^16. Within |value_scale| <= 2^16, any probability flushed below 2^-40 by the 2^15-scaled FP16 probability tile changes p * V * value_scale by at most 2^-40 * 128 * 2^16 = 2^-17.
• AttendRowFp32<kSplitK> now folds value_scale[d] into the per-candidate numerator update with a 2^80 exponent shift when score - next_max < -80, preserving tiny probabilities across the unrestricted finite value_scale domain.
• Extended DynamicSparseAttentionTest.Int8TensorCoreTinyProbabilityLargeValueScale_CUDA to test both prefill and split-K decode with TC-enabled (1) and TC-disabled (0) across:
  1. query = 288, scale = 1/16, value_scale = 2^16, value_codes = {0, 100} (p = e^-18 ~= 2^-25.97, exercising the 2^15 * p FP16 scaling on the TC path),
  2. query = 320, scale = 1/16, value_scale = 2^28, value_codes = {0, 1} (p = e^-20),
  3. query = 480, scale = 1/16, value_scale = 2^60, value_codes = {0, 1} (p = e^-30 ~= 2^-43.28, expected 2^60 * e^-30 / (1 + e^-30) ~= 107886.05).
• Verified all 66 DynamicSparseAttentionTest.* tests pass on A100 and register/spill counts of DynamicSparseAttentionInt8TcKernel<{false,true}> remain unchanged (149 / 153 registers, 0 stack).

[Baiju Meswani (baijumeswani)]

Summary

This PR adds INT8-KV tensor-core prefill and split-K decode kernels for CUDA DynamicSparseAttention, shares their implementation, adds session/environment switches and fallbacks, triggers the gate-mul PDL successor early, and makes the hyper-connection sigmoid branch-free.

What looks good

The shared kernel design, aligned workspace layout, reuse of the existing merge kernel, dispatch guards, and expanded boundary/configuration tests are sensible. Independent A100/CUDA 13.0 measurements support a real kernel speedup: prefill improved from about 18.52 to 6.53 ms for 8192 rows and 2048 candidates, and decode plus merge improved from about 141.04 to 25.86 us for eight rows and 2048 candidates. The sigmoid comparison had zero mismatches across all 2^32 float bit patterns without fast-math. A separate CUDA review exercised 140 cases with memcheck, racecheck, synccheck and initcheck, with no errors, and compiled the kernels for SM75/80/86/90/120. The PTX-target guard was also checked directly. PDL ordering looks sound by inspection; it was not executed on Hopper.

What must change before approval

I would not approve this head yet. The inline comments describe four GPU-reproduced numerical regressions: finite query folding loses score differences, FP16 probability rounding loses signed-value cancellation in decode, early log2 conversion loses a representable logit difference, and the decode epilogue loses tiny value scales. The first and last independently confirm the earlier review findings; the other two are additional issues. The value-scale safety fallback also made a valid prefill workload about 20.8x slower than the previous kernel, so the default rollout needs a performance-safe solution without removing the numerical safeguard.

Please fix these cases and add focused tests with an input-precision-aware, stable FP64 reference. After the fixes, qualify the exact final head with isolated prefill/decode switch comparisons, paired model-accuracy results, repeated H200 timings, and accepted-token throughput. The current combined-stack MMLU-Pro results and earlier one-round timing tables do not establish this PR's final-head accuracy or the exact end-to-end percentages independently. This is an evidence gap, not a claim that the reported H200 gains are false.

The A100 split heuristic and the two option-documentation mismatches are smaller, nonblocking concerns. In the tested decode case, automatic TC remains 5.46x faster than the old decoder, although eight splits are about 17% faster than the automatic sixteen.

Evidence and limits

This consolidates three independent reviews covering numerics, CUDA safety/portability, and performance/design. Numerical reproductions directly launched the actual head kernels and compared them with the unchanged prefill and exact-base grouped-decode/merge kernels. Common repro settings are B=1, head size 256, 12 query/2 KV heads, aligned BF16 query/output and INT8 caches, selected-only/main, no rotary, Q/K normalization or smooth softmax, capacity 520 and max_selected 512; use S=9 for prefill and S=3 for decode. Unless a comment says otherwise, two valid past positions are selected. All other channels are zero except where a repeated pattern is specified. Performance probes used 24 query/2 KV heads; clocks were not locked. No full ORT provider suite, H200/Qwen model evaluation, or Hopper PDL execution was run. These are kernel-level results, not independent reproductions of the full-model percentages.