diff --git a/docs/contrib_ops/cuda/dynamic_sparse_attention.md b/docs/contrib_ops/cuda/dynamic_sparse_attention.md index 86ca09455b0a3..5635d7c9b93d5 100644 --- a/docs/contrib_ops/cuda/dynamic_sparse_attention.md +++ b/docs/contrib_ops/cuda/dynamic_sparse_attention.md @@ -47,7 +47,17 @@ the code is only reported through `validation_status`. `ORT_DSA_INT8_VECTOR_STAGING` and `ORT_DSA_SHARED_SCALE_MERGE` are read once per process (cached in function-local statics) when the grouped decode path is first launched, so set them before the first operator invocation. Only `0` and `1` are accepted. Neither changes schemas or cache layouts. -`ORT_DSA_PREFILL_KERNEL` (default `1`) is read on every BF16 call. +`ORT_DSA_PREFILL_KERNEL`, `ORT_DSA_PREFILL_INT8_TC` and `ORT_DSA_DECODE_INT8_TC` are read once at kernel +construction. `ORT_DSA_PREFILL_KERNEL` defaults to `1`. The two INT8 tensor-core switches default to `0`: they are +opt-in, and without them the operator runs the BF16-dequantizing prefill and the CUDA-core grouped decode, bit for +bit as before (see [Rollout](#rollout)). Session options `ep.cuda.dsa_prefill_kernel`, `ep.cuda.dsa_prefill_int8_tc` +and `ep.cuda.dsa_decode_int8_tc` take precedence over their corresponding environment variables. Each of them +accepts exactly the one-character string `0` or `1` (no whitespace, sign or leading zeros: `01`, `-0` and `true` are +rejected); changing a value requires a new session. The loaded prefill and decode specializations must have a PTX +target of SM80+; otherwise dispatch uses the generic fallback. +`ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` (session option `ep.cuda.dsa_decode_int8_tc_tiles_per_split`) is a +tuning and test knob with the same precedence: an integer in `[0, INT_MAX]`, `0` keeps the automatic split +heuristic. Invalid values of any of these options fail session initialization. Grouped decode is used when `1 <= sequence_length <= 8`, `num_heads / kv_num_heads >= 4` (and divisible), `head_size == 256`, 16-bit `T` (FP16/BF16), more than 256 candidates @@ -63,12 +73,123 @@ and 8 otherwise. | | | | `0` | `MergeGroupedDecodeDynamicSparseAttentionKernel`: per-channel `expf` over all splits. | | `ORT_DSA_PREFILL_KERNEL` | `1` | BF16, not grouped decode, `sequence_length >= 8`, `head_size == 256`, `selected_only` from the main cache, `smooth_softmax = 0`, SM80+ with enough opt-in shared memory | `1` | `DynamicSparseAttentionPrefillKernel` tensor-core prefill. | | | | | `0` | Generic fused/split or grouped prefill kernels. | +| `ORT_DSA_PREFILL_INT8_TC` | `0` | The `ORT_DSA_PREFILL_KERNEL` conditions with an INT8 main cache, 16-byte-aligned `present_key` and `present_value`, an 8-byte-aligned `output`, a softmax `scale <= 1` and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel` ([dynamic_sparse_attention_prefill_int8.cuh](../../../onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh)): the gathered K/V rows stay INT8 in shared memory and are converted exactly to FP16 in registers for FP16 `mma` with FP32 accumulation. `key_scale` is folded into the query once per block (`Q' = Q * key_scale`, rounded once to FP16) and `value_scale` is applied to the output once. Scores and the running maximum are kept in the natural-log domain. A softmax `scale > 1` (the default is `1 / sqrt(head_size)`) uses the kernel below, because the rounding bound grows with it. The FP16 probability tile holds `2^15 * p`, so a candidate keeps its contribution down to `p = 2^-40` (plain FP16 flushes below `2^-25`); the row sum stays unscaled FP32. Prefill keeps a single FP16 probability tile, as the BF16-dequantizing kernel does with a coarser BF16 one. A block whose query or value scales leave the validated range is evaluated by an FP32 path (`AttendRowFp32`) instead of being clipped or rounded, which would break the cancellation between large terms; the conditions, their cost and how likely they are are described under [Numerics and fallbacks](#numerics-and-fallbacks). | +| | | | `0` | `DynamicSparseAttentionPrefillKernel`: dequantizes every gathered tile to BF16 in shared memory. | +| `ORT_DSA_DECODE_INT8_TC` | `0` | Grouped decode (conditions above) with BF16, an INT8 main cache, `head_size == 256`, `selected_only` from the main cache (no `local_plus_selected`), `smooth_softmax = 0`, a softmax `scale <= 1`, at most `kMaxGroup` query heads per KV head, 16-byte-aligned `prepared_query`, `present_key`, `present_value`, `key_scale`, `value_scale` and attention workspace, SM80+ and enough opt-in shared memory for `dsa_prefill::kInt8SplitKSharedBytes` (75,264 B) | `1` | `DynamicSparseAttentionInt8TcKernel`: the tensor-core pipeline of the row above, split over the 64-candidate tiles of each (row, KV head). Each CTA of the (row, KV head, split) grid writes an un-normalized partial output, the natural-log row maximum and the row sum in the layout of the grouped decode merge kernel, which is reused unchanged. The probability tile is an FP16 high part plus an FP16 residual (about 22 significant bits) and both multiply the converted V fragments, because the signed sum of large value codes can cancel to a value far below the rounding of one FP16 probability. A partial is written as `(acc * 2^-15) * value_scale`. The split count follows `ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` below. A CTA that meets one of the conditions of [Numerics and fallbacks](#numerics-and-fallbacks) evaluates its candidates in FP32 (`AttendRowFp32`) and publishes an FP32 partial in the same layout, so the fallback also holds for the splits. | +| | | | `0` | The CUDA-core `GroupedDecodeDynamicSparseAttentionKernel`, bit for bit as before. | +| `ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` | `0` | `ORT_DSA_DECODE_INT8_TC=1` and the conditions above | `0` | Automatic: 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 (`rows = batch_size * sequence_length`; resident CTAs = SM count times the CTAs per SM that the kernel's shared-memory and register footprint allow, queried with `cudaOccupancyMaxActiveBlocksPerMultiprocessor`: 396 on an H200). The split count is `ceil(tiles / tiles_per_split)`, at most 32 and never more than the grouped decode workspace was sized for. | +| | | | `N > 0` | A requested `N` tiles per split instead of the automatic choice. The 32-split cap, the workspace size and the even distribution of the tiles over the splits still apply, so the number of tiles of a split can exceed `N` (8192 candidates are 128 tiles: `N = 1` asks for 128 splits and runs 32 splits of 4 tiles). Used by the tests to force many or few splits. | + +The decode time is not flat in the tiles per split: with too few splits the CTAs are long and SMs stay idle, with too +many the grid spills into a second, partly filled wave. On an H200, 8 rows and 2048 candidates (16 pairs of row and KV +head with 32 tiles each), kernel plus merge, +`benchmark_dynamic_sparse_attention.py --rows 8 --tiles-per-split 0 1 2 4 8 16`: + +| 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 at 2 tiles per split). The +previous target of about 384 CTAs is 2.7% and 14.2% slower on the same points: with 396 resident CTAs the two rules mostly +agree on an H200, and the larger shapes gain from the power-of-two choice (8 rows, 8192 candidates: 0.2% against 6.7%). +The rule 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; 1 to 32 rows, 2048 candidates) +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. A device with 108 SMs and two resident CTAs per SM (the 75,264 B of shared memory +allow two on an A100) is expected to choose 4 tiles per split, 8 splits, for 8 rows and 2048 candidates, the faster +setting in an A100 measurement reported for this change (22.1 us against 25.9 us at 16 splits). The A100 has not been +measured with this document's script; `ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` tunes other devices. When `split_count > 32` (floating-point caches only) the generic `MergeDynamicSparseAttentionKernel` is used regardless of `ORT_DSA_SHARED_SCALE_MERGE`. The tests `Int8CacheGroupedDecodeAlternateSwitchesInFreshProcess_CUDA` re-run the INT8 grouped-decode and merge cases in fresh processes with each non-default combination. +### Numerics and fallbacks + +The INT8 tensor-core kernels multiply `Q' = fp16(Q * key_scale)` with the exactly converted INT8 K rows and the FP16 +probability tile with the exactly converted INT8 V rows, in FP16 `mma` with FP32 accumulation. They therefore differ from +the CUDA-core kernels by the rounding of `Q'` and of the probability tile and by the FP32 summation order, so outputs are +close but not bit-identical to the previous kernels. A CTA (one query row and KV head, that is a group of up to 16 query +heads over the row's candidates) that cannot bound that rounding does not clip or round: it evaluates its heads in FP32 +on the CUDA cores (`AttendRowFp32`), which is exact. The decision is uniform across the CTA and the FP32 path starts +from the inputs, so the result of a CTA is either the tensor-core result within the bounds below or the FP32 result. + +| A CTA takes the FP32 path when | Bound that it enforces | RC model (Qwen3.8-Flash-Next DSA layers) | +|---|---|---| +| An element of `Q * key_scale` is not finite or its magnitude is above 65504 | `Q'` fits FP16 | Does not happen | +| For one of its query heads, `eps = scale * 128 * sum_d abs(Q'_d - fp16(Q'_d))` is above `kInt8TcQueryErrorBudget` (0.125) | Rounding `Q'` moves any logit of that head by at most `eps` nats (`abs(K_int8) <= 128`), so a probability by at most a factor `exp(eps)`. A query whose large terms cancel to a small logit exceeds the budget by orders of magnitude, an ordinary query does not | `eps` is 0.015 on average and 0.022 at most per head (table below) | +| A `value_scale` is not finite | The FP16 accumulator result is meaningful | Does not happen | +| Some channel has `abs(value_scale) > 2^16` and a valid candidate of one head is more than 20 nats (`kInt8TcTailDepth`) below the running maximum | A probability below `2^-29` is a subnormal of the `2^15 * p` tile and below `2^-40` it is zero. For `abs(value_scale) <= 2^16` the lost contribution is below `2^-17` of one value code; above it the check decides per CTA instead of sending every CTA to the FP32 path | The largest `value_scale` is 0.285: the check is never armed | + +The scores and the running maximum are kept in the natural-log domain and only the difference to the new maximum is +scaled by `log2(e)` for `ex2`, so the difference of two large, nearby logits is not rounded away. Prefill keeps a single +FP16 probability tile holding `2^15 * p`, as the BF16-dequantizing kernel does with a coarser BF16 tile; decode adds an +FP16 residual tile, because the signed sum of large value codes can cancel to a value far below the `2^-11` rounding of +one FP16 probability. + +The bound `eps` is proportional to the magnitude of the query. Monte Carlo over the 12 DSA layers of the RC model +(their `q_norm` weights and key scales, 4000 rows per layer and KV head, scale `1/16`, unit-RMS Gaussian queries with +partial rotary embedding rounded to BF16, then multiplied by `m`; `m = 1` is what the model produces): + +| `m` | mean `eps` of a head | largest `eps` of a head | CTAs over the budget (12 heads per CTA) | +|---|---|---|---| +| 1 | 0.015 | 0.022 | 0 | +| 2 | 0.031 | 0.045 | 0 | +| 4 | 0.061 | 0.088 | 0 | +| 8 | 0.122 | 0.179 | 69% | +| 16 | 0.245 | 0.356 | 100% | + +The budget is therefore about 8 times the magnitude the model produces. Heavy-tailed queries (Student-t with 4 degrees +of freedom) reach 48% of the CTAs at `m = 8`, queries whose channel variances differ by a log-normal factor (sigma 1) +about 1%. Workloads whose queries are 8 times larger than those of the RC model, or whose cache has +`abs(value_scale) > 2^16` and wide logit spreads, should leave the switches at `0`: the FP32 path is slower than the +previous kernels, as the measurements below show. It stays accurate: with queries 8 times larger the relative RMS error +against FP64 is 1.7e-3 (the previous kernels: 8.5e-3 for prefill, 1.7e-3 for decode). + +Cost of the FP32 path on an H200 (`benchmark_dynamic_sparse_attention.py`; 24 query / 2 KV heads, head size 256, +2048 selected candidates, context 8192; attention-kernel time, previous kernel: BF16-dequantizing prefill / CUDA-core +decode): + +| Data | Rows | Previous kernel | Tensor-core kernel | Tensor core vs previous | +|---|---|---|---|---| +| Queries like the RC model | 128 | 0.173 ms | 0.063 ms | 2.7x faster | +| | 8192 | 7.82 ms | 2.83 ms | 2.8x faster | +| Queries 8 times larger (`--query-std 11.2`, every CTA over the budget) | 128 | 0.172 ms | 3.68 ms | 21x slower | +| | 8192 | 7.82 ms | 135.6 ms | 17x slower | +| | 8 (decode) | 73.6 us | 240 us | 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 | 0.172 ms | 3.61 ms | 21x slower | +| | 512 | 0.577 ms | 3.90 ms | 6.8x slower | +| | 2048 | 2.288 ms | 6.17 ms | 2.7x slower | +| The same value scales, logit standard deviation about 1.3 nats (the cache holds only rows appended by the call): no CTA is deeper than 20 nats | 8192 | 7.78 ms | 2.86 ms | 2.7x faster | +| The same value scales | 8 (decode) | 73.6 us | 13.3 us | 5.6x faster | + +A CTA on the FP32 path needs about 3.6 ms for 2048 candidates, so the kernel time is set by the slowest CTA of a launch +with one wave of CTAs and grows with the fraction of such CTAs for larger launches. A decode CTA covers fewer candidates +(64 to 128 at the automatic split on an H200 for 2048 candidates and 1 to 8 rows), so a candidate 20 nats below the +running maximum is rarer and the decode row of the last group shows no slowdown. With +`ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT=16` (1024 candidates per CTA) the same data takes 1.8 ms instead of 13 us, with +the same RMS error (1.7e-3). The previous design of this kernel sent every CTA to the FP32 path as soon as one channel +had `abs(value_scale) > 2^16` (8192 rows: 2.76 ms became 136 ms), which the check above replaces; it still applies to +queries above the budget, as the second group of rows shows. The unit tests exercise every condition of the table +(`Int8TensorCoreQueryFoldingRoundingError_CUDA`, `Int8TensorCoreDecodeSignedValueCancellation_CUDA`, +`Int8TensorCoreLargeCommonLogit_CUDA`, `Int8TensorCoreTinyValueScale_CUDA`, +`Int8TensorCoreLargeValueScaleRandom_CUDA`, `Int8TensorCoreRandomQueryFolding_CUDA`, +`Int8PrefillQueryBeyondHalfRange_CUDA`, `Int8DecodeQueryBeyondHalfRange_CUDA`) against an FP64 host reference. + +### Rollout + +`ORT_DSA_PREFILL_INT8_TC` and `ORT_DSA_DECODE_INT8_TC` are opt-in and independent, so prefill and decode can be enabled +one at a time. The end-to-end accuracy measured for an earlier revision of these kernels (MMLU-Pro, 2400 items, +pooled difference +0.08 percentage points, 95% confidence interval -0.60 to +0.77) predates the numerics described above +and does not cover the final revision. The end-to-end timings in [Benchmark](#benchmark) are preliminary for the same +reason. Making a switch the default needs paired accuracy runs of the final revision with that switch toggled alone, +including long contexts, and a stated bound on the degradation. 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`). + ### Buffer alignment The wide-load fast paths assume more alignment than a tensor's element type guarantees, and tensors @@ -81,6 +202,8 @@ buffer is misaligned: | INT8 grouped decode vector staging (`ORT_DSA_INT8_VECTOR_STAGING=1`) | `key_scale`, `value_scale`, `auxiliary_key`, `auxiliary_value`: 16 bytes; `present_key`, `present_value`: 8 bytes | Scalar per-channel staging (`kVectorizedStaging=false`). | | Floating-point grouped decode and grouped prefill | Each staged K/V row: 16 bytes (checked per row in the kernel) | Element-wise copy into shared memory. | | Tensor-core BF16 prefill (`ORT_DSA_PREFILL_KERNEL=1`) | `present_key`, `present_value`: 16 bytes (8 for INT8); `key_scale`, `value_scale`: 16 bytes; `output`: 4 bytes | Generic fused/split or grouped prefill kernels. | +| Tensor-core INT8-cache prefill (opt-in, `ORT_DSA_PREFILL_INT8_TC=1`) | In addition to the row above: `present_key`, `present_value`: 16 bytes; `output`: 8 bytes | `DynamicSparseAttentionPrefillKernel`. | +| Tensor-core INT8-cache decode (opt-in, `ORT_DSA_DECODE_INT8_TC=1`) | `prepared_query`, `present_key`, `present_value`, `key_scale`, `value_scale`, attention workspace: 16 bytes | The CUDA-core grouped decode kernel. | Buffers allocated by ONNX Runtime are 256-byte aligned, so only IOBinding with offset buffers takes a fallback. The `*MisalignedBuffers_CUDA` tests cover each case. @@ -96,5 +219,72 @@ only as fallbacks: | `ORT_DSA_INT8_VECTOR_STAGING=0` | about +0.3 ms (+2%) slower | | `ORT_DSA_SHARED_SCALE_MERGE=0` | about +0.05 to +0.1 ms (+0.3 to 0.5%) slower | +The INT8 tensor-core kernels (opt-in, see [Rollout](#rollout)) are measured at the operator level with +[benchmark_dynamic_sparse_attention.py](../../../onnxruntime/test/python/transformers/benchmark_dynamic_sparse_attention.py) +against the kernels that they replace (the BF16-dequantizing prefill and the CUDA-core grouped decode). The script 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 full-attention layers of Qwen3.8-Flash-Next) 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, taken from the 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. +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 RC 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 (`--rc-model` reads the scales of a model's node instead). H200, CUDA 13.0, the automatic tiles per split: + +```bash +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 previous prefill kernel spent 2.8% of its instructions on `mma` and the rest on in-kernel dequantization and index +handling, and the previous grouped decode kernel was instruction-issue and shared-load bound on its per-candidate +serial chain. 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. The tensor-core kernels accumulate Q.K +and P.V in a different FP32 order, so their outputs are close but not bit-identical to the previous kernels. + +End-to-end numbers are preliminary. In the Qwen3.8-Flash-Next packages, an earlier revision of these kernels +(one round of consecutive A/B runs, not alternated) shortened the time to first token by 4.0 to 4.8% at 2048, 8.4 to 9.5% +at 8192 and 7.6 to 9.4% at 32768 prompt tokens, the speculative-decoding cycle (S = 8) by 4.6 to 6.8% at 2K, 8K and 32K +context, and the 8192-row prefill launch from 7.9 to 2.8 ms. They are single A/B rounds without raw timings or +confidence intervals and have not been reproduced for the current revision of the kernels, whose numerics changed (see +[Rollout](#rollout)). A shorter speculative-decoding cycle is not better user-visible throughput by itself: the accepted +tokens per cycle and the time per accepted token were not measured. + +`ORT_DSA_PREFILL_INT8_TC=0` and `ORT_DSA_DECODE_INT8_TC=0` (the defaults) run the previous kernels exactly. The tests +`Int8CachePrefillTensorCoreKernel_CUDA` and `Int8CacheDecodeTensorCoreKernel_CUDA` compare the tensor-core kernels with a +reference over head counts, candidate counts, splits and invalid indices; `Int8DecodeQueryBeyondHalfRange_CUDA` covers +the FP32 fallback of the splits, `Int8TensorCoreDecodeTilesPerSplitCaps_CUDA` the split caps, and +`PrefillSessionOptions_CUDA` and `DecodeSessionOptions_CUDA` the switches, the tuning knob and invalid values. +`Int8TensorCoreLargeSoftmaxScale_CUDA` (the `scale > 1` fallback) and +`Int8TensorCoreTinyProbabilityLargeValueScale_CUDA` (the FP16 probability range) cover prefill and decode, and the tests +listed under [Numerics and fallbacks](#numerics-and-fallbacks) the guards. + The corresponding indexer switch results are in [../packed_sparse_attention_indexer.md](../packed_sparse_attention_indexer.md#benchmark). diff --git a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md index 348131a3db9b3..6d772a36c8397 100644 --- a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md +++ b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md @@ -37,6 +37,31 @@ Read once per process, so set it before starting it. weight inputs (`norm_scale`, `down_weight`, `up_weight`, `inject_weight`) are constant initializers; `0` forces plain launches. +`ProjectedSigmoidGateMulKernel` +([gated_add_impl.cu](../../../onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu)), +which precedes `HcDownKernel` in the Qwen3.8 decode graph, issues +`griddepcontrol.launch_dependents` at its start on `SM90+`. When `HcDownKernel` +is launched with PDL, its prologue (constant-weight prefetch) then overlaps with +the gate-mul kernel, while its reads of the gate-mul output still wait for +`griddepcontrol.wait`. The trigger has no effect for plain launches, changes no +result and has no switch. Together with the same early trigger in the fpA_intB GEMV +kernels (MatMulNBits) it shortened a Qwen3.8-Flash-Next `S = 8` decode cycle by 1.4-1.7% in a +preliminary end-to-end A/B on an H200 (one round, not alternated, no confidence interval). + +`SigmoidF32` evaluates `expf(-|v|)` once and selects the numerator, +`e = expf(-|v|); (v >= 0 ? 1 : e) / (1 + e)`. The textbook two-branch form passes the +same argument to `expf` (`-v` for `v >= 0`, `v` otherwise), so the result is +bit-identical for all 2^32 `float` values (checked exhaustively in a standalone harness, +no fast-math). The branch itself was the cost: in the prefill pre-mix neighbouring +elements have different signs, so nearly every warp executed both sides. For hidden +2560, 4 branches and an 8192-token prefill chunk on an H200, `PreMixVec8Kernel` takes +82.7 µs instead of 112.1 µs per launch and `ScaledSiluKernel` 6.1 µs instead of 7.3 µs. +The TTFT shortened by 0.3-1.1% in a preliminary end-to-end A/B (median of 5 repetitions, +faster in 9 of 9 package and context cells), which is about the size of the roughly 1% +noise between repetitions, so the kernel times are the firmer evidence. The fused decode +kernels call the same function for at most eight rows, and decode time is unchanged within +the noise of 0.3%. + The rank-320 decode test (`HyperConnectionLowRankMixTest.FusedDecodeShapeRank320`) covers 1, 2, 3, 4, 7 and 8 rows. diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 96dcb9de23cae..62423de5863c2 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -510,6 +510,29 @@ static const char* const kOrtSessionOptionsCudaFpAIntBGemvWaveAware = "ep.cuda.f static const char* const kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength = "ep.cuda.gqa_workspace_max_total_sequence_length"; +/// CUDA DynamicSparseAttention kernel switches. Each is resolved once when the kernel is created, so a change needs a new +/// session, and overrides the environment variable named below. The 0/1 switches accept only "0" and "1", and the +/// thresholds a non-negative decimal integer; any other value, including an empty string, fails session initialization. +/// "1" makes the INT8-cache BF16 prefill run FP16 tensor-core MMA on the gathered K/V tiles instead of dequantizing them +/// to BF16 (opt-in; the tensor-core kernel has different numerics). Default "0". Overrides ORT_DSA_PREFILL_INT8_TC. +static const char* const kOrtSessionOptionsCudaDsaPrefillInt8Tc = "ep.cuda.dsa_prefill_int8_tc"; + +/// "0" routes the BF16 prefill through the generic kernels. Default "1". Overrides ORT_DSA_PREFILL_KERNEL. +static const char* const kOrtSessionOptionsCudaDsaPrefillKernel = "ep.cuda.dsa_prefill_kernel"; + +/// "1" makes the INT8-cache BF16 decode (up to 8 query rows) use the split-K tensor-core kernel instead of the +/// CUDA-core grouped kernel (opt-in; the tensor-core kernel accumulates in a different order and does not reproduce +/// the CUDA-core results bit for bit). Default "0". Overrides ORT_DSA_DECODE_INT8_TC. +static const char* const kOrtSessionOptionsCudaDsaDecodeInt8Tc = "ep.cuda.dsa_decode_int8_tc"; + +/// Requested number of 64-candidate tiles each CTA of the split-K tensor-core decode kernel consumes. "0" derives it from +/// the problem size and the device: the smallest power of two with which all CTAs fit one wave of the CTAs that can be +/// resident on the device. A positive value is a target: the number of splits is ceil(tiles / value), limited to 32 and +/// to the split count of the decode workspace, so a small value on a long candidate list is raised and the last split +/// may hold fewer tiles. Default "0". Overrides ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT. +static const char* const kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit = + "ep.cuda.dsa_decode_int8_tc_tiles_per_split"; + /// Reader-only, hardware-neutral kernel dispatch intent shared across operators and EPs. /// Valid values are "auto", "latency", "memory", and "safe"; matching is case-sensitive. /// Unset or empty values map to "auto". Other values cause INVALID_ARGUMENT during diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc index d3e020eb16e2e..c115a8c9db26e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc @@ -105,6 +105,15 @@ DynamicSparseAttention::DynamicSparseAttention(const OpKernelInfo& info) rotary_interleaved_ = ParseBoolAttribute(info, "rotary_interleaved", 0); use_smooth_softmax_ = ParseBoolAttribute(info, "smooth_softmax", 0); auxiliary_kv_shared_ = ParseBoolAttribute(info, "auxiliary_kv_shared", 0); + prefill_kernel_ = ReadCudaSessionSwitch(info, kOrtSessionOptionsCudaDsaPrefillKernel, "ORT_DSA_PREFILL_KERNEL", true); + // The tensor-core INT8-cache kernels are opt-in: their numerics differ from the kernels they replace, so the default + // dispatch stays on the BF16-dequantizing prefill and the CUDA-core grouped decode kernels. + prefill_int8_tc_ = + ReadCudaSessionSwitch(info, kOrtSessionOptionsCudaDsaPrefillInt8Tc, "ORT_DSA_PREFILL_INT8_TC", false); + decode_int8_tc_ = ReadCudaSessionSwitch(info, kOrtSessionOptionsCudaDsaDecodeInt8Tc, "ORT_DSA_DECODE_INT8_TC", false); + decode_int8_tc_tiles_per_split_ = + ReadCudaSessionOption(info, kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, + "ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT", 0, 0, std::numeric_limits::max()); strict_validation_ = ParseEnvironmentVariableWithDefault(kStrictValidation, false); attention_mode_ = ParseAttentionMode( info.GetAttrOrDefault("attention_mode", "selected_only")); @@ -203,6 +212,10 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons using CudaT = typename OrtToCudaType::type; DynamicSparseAttentionData data; + data.prefill_kernel = prefill_kernel_; + data.prefill_int8_tc = prefill_int8_tc_; + data.decode_int8_tc = decode_int8_tc_; + data.decode_int8_tc_tiles_per_split = decode_int8_tc_tiles_per_split_; data.query = reinterpret_cast(query->Data()); data.key = key == nullptr ? nullptr : reinterpret_cast(key->Data()); data.value = value == nullptr ? nullptr : reinterpret_cast(value->Data()); diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h index 8719af8b3cb04..40f8d78d58859 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h @@ -4,6 +4,7 @@ #pragma once #include "core/providers/cuda/cuda_kernel.h" +#include "contrib_ops/cuda/cuda_session_options.h" #include "contrib_ops/cpu/bert/dynamic_sparse_attention_helper.h" namespace onnxruntime { @@ -29,6 +30,10 @@ class DynamicSparseAttention final : public onnxruntime::cuda::CudaKernel { bool use_smooth_softmax_; bool auxiliary_kv_shared_; bool strict_validation_; + bool prefill_kernel_; + bool prefill_int8_tc_; + bool decode_int8_tc_; + int decode_int8_tc_tiles_per_split_; DynamicSparseAttentionMode attention_mode_; DynamicSparseAttentionKvSource selected_kv_source_; }; diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu index a612c7071aab9..eaae0df5d3ab7 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -15,6 +15,7 @@ #include #include "contrib_ops/cuda/bert/dynamic_sparse_attention_prefill.cuh" +#include "contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh" #include "core/common/safeint.h" #include "core/providers/cuda/cuda_common.h" @@ -1806,6 +1807,31 @@ int GetGroupedDecodeSplitCount(const DynamicSparseAttentionParameters& parameter return static_cast(std::min({candidate_splits, occupancy_splits, max_splits})); } +// Split count of the tensor-core INT8 decode kernel: one CTA per (row, KV head, split), each consuming at least one +// 64-candidate tile. The automatic choice is the smallest power-of-two number of tiles per split with which all CTAs +// fit one wave of the device, resident_ctas being the SM count times the CTAs per SM that the kernel's shared-memory +// and register footprint allow (396 on an H200). A partly filled second wave costs more than longer splits do (H200 +// with the occupancy limited to one CTA per SM, 8 rows, 32 tiles: 16 splits take 18.7 us, 8 splits 14.7 us), and a +// power of two spreads the tiles evenly (H200, 16 rows, 32 tiles: 11 splits of up to 3 tiles take 20.5 us, 8 splits +// 18.5 us). The kernel is not flat in the tiles per split (H200, 8 rows, 32 tiles: 1 -> 20.4 us, 2 -> 15.0, 4 -> 14.5, +// 8 -> 19.5, 16 -> 30.9). The count never exceeds the split count the workspace was sized for. A positive +// tiles_per_split_override replaces the automatic choice. +int GetInt8TcDecodeSplitCount(const DynamicSparseAttentionParameters& parameters, int workspace_splits, + int tiles_per_split_override, int resident_ctas) { + const int64_t tiles = (static_cast(parameters.max_selected) + dsa_prefill::kTile - 1) / dsa_prefill::kTile; + const int64_t row_kv_heads = + static_cast(parameters.batch_size) * parameters.sequence_length * parameters.kv_num_heads; + int64_t tiles_per_split = tiles_per_split_override; + if (tiles_per_split <= 0) { + tiles_per_split = 1; + while (tiles_per_split * std::max(resident_ctas, 1) < tiles * row_kv_heads) { + tiles_per_split *= 2; + } + } + const int64_t splits = (tiles + tiles_per_split - 1) / tiles_per_split; + return static_cast(std::max(1, std::min({splits, workspace_splits, 32}))); +} + size_t GetValidationHashEntries(const DynamicSparseAttentionParameters& parameters) { if (parameters.max_selected == 0) { return 0; @@ -2073,9 +2099,10 @@ Status LaunchDynamicSparseAttention( grouped_decode_threads_supported && UseGroupedDecodeAttention(parameters, sizeof(T), max_shared_memory_per_block, data.quantized_cache); bool use_prefill_kernel = false; + bool use_int8_tc_prefill = false; if constexpr (std::is_same_v) { if (!use_grouped_decode) { - const bool enabled = ParseEnvironmentVariableWithDefault("ORT_DSA_PREFILL_KERNEL", 1) != 0; + const bool enabled = data.prefill_kernel; // max_shared_memory_per_block is the default 48 KB limit; the kernel opts in to more and requires SM80+. int device = 0; int cc_major = 0; @@ -2101,10 +2128,84 @@ Status LaunchDynamicSparseAttention( parameters.num_heads % parameters.kv_num_heads == 0 && parameters.num_heads / parameters.kv_num_heads <= dsa_prefill::kMaxGroup && static_cast(optin_shared) >= dsa_prefill::kSharedBytes; + if (use_prefill_kernel) { + cudaFuncAttributes attributes{}; + const auto kernel = data.quantized_cache ? dsa_prefill::DynamicSparseAttentionPrefillKernel + : dsa_prefill::DynamicSparseAttentionPrefillKernel; + CUDA_RETURN_IF_ERROR(cudaFuncGetAttributes(&attributes, kernel)); + use_prefill_kernel = attributes.ptxVersion >= 80; + } + // The tensor-core INT8 variant converts the cached rows in registers; it needs 16-byte cache rows and 8-byte + // output vectors, and a softmax scale of at most kInt8TcMaxScale so that it does not amplify its fp16 rounding + // of Q * key_scale (the BF16-dequantizing kernel above is the fallback). The op validates the scale as finite and + // non-negative. + use_int8_tc_prefill = + use_prefill_kernel && data.quantized_cache && + data.prefill_int8_tc && parameters.scale <= dsa_prefill::kInt8TcMaxScale && + IsAligned(data.present_key, 16) && IsAligned(data.present_value, 16) && IsAligned(data.output, 8) && + static_cast(optin_shared) >= dsa_prefill::kInt8SharedBytes; + if (use_int8_tc_prefill) { + cudaFuncAttributes attributes{}; + CUDA_RETURN_IF_ERROR(cudaFuncGetAttributes(&attributes, dsa_prefill::DynamicSparseAttentionInt8TcKernel)); + use_int8_tc_prefill = attributes.ptxVersion >= 80; + } + } + } + // Decode (sequence_length <= 8) over the INT8 cache: the same tensor-core pipeline split over the candidate tiles of + // each (row, KV head); the partials go through the grouped decode merge kernel. + bool use_int8_tc_decode = false; + int int8_tc_decode_resident_ctas = 0; + if constexpr (std::is_same_v) { + if (use_grouped_decode && data.quantized_cache && data.decode_int8_tc) { + int device = 0; + int cc_major = 0; + int optin_shared = 0; + CUDA_RETURN_IF_ERROR(cudaGetDevice(&device)); + CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device)); + CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&optin_shared, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); + constexpr int kMaxGridDimY = 65535; + // The kernel loads 16-byte query / scale vectors, gathers 16-byte cache chunks and stores 16-byte partials; + // buffers bound through IOBinding can start at any element offset, so misaligned ones use the generic kernel. + // Like the prefill variant it needs a softmax scale of at most kInt8TcMaxScale, which keeps the fp16 rounding of + // Q * key_scale from being amplified. + use_int8_tc_decode = + cc_major >= 8 && parameters.scale <= dsa_prefill::kInt8TcMaxScale && + parameters.head_size == dsa_prefill::kHeadDim && !local_plus_selected && + parameters.selected_kv_source == DynamicSparseAttentionKvSource::kMain && + !parameters.use_smooth_softmax && parameters.max_selected > 0 && parameters.kv_num_heads <= kMaxGridDimY && + parameters.num_heads / parameters.kv_num_heads <= dsa_prefill::kMaxGroup && + IsAligned(data.prepared_query, 16) && IsAligned(data.present_key, 16) && IsAligned(data.present_value, 16) && + IsAligned(data.key_scale, 16) && IsAligned(data.value_scale, 16) && + IsAligned(data.attention_workspace, 16) && + static_cast(optin_shared) >= dsa_prefill::kInt8SplitKSharedBytes; + if (use_int8_tc_decode) { + // A build whose PTX is older than compute_80 holds an empty stub for the kernel. + cudaFuncAttributes attributes{}; + CUDA_RETURN_IF_ERROR(cudaFuncGetAttributes(&attributes, dsa_prefill::DynamicSparseAttentionInt8TcKernel)); + use_int8_tc_decode = attributes.ptxVersion >= 80; + } + if (use_int8_tc_decode) { + // The split count targets one wave of resident CTAs: the occupancy query needs the shared-memory opt-in. + CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( + dsa_prefill::DynamicSparseAttentionInt8TcKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(dsa_prefill::kInt8SplitKSharedBytes))); + int multiprocessors = 0; + int blocks_per_multiprocessor = 0; + CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&multiprocessors, cudaDevAttrMultiProcessorCount, device)); + CUDA_RETURN_IF_ERROR(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &blocks_per_multiprocessor, dsa_prefill::DynamicSparseAttentionInt8TcKernel, dsa_prefill::kThreads, + dsa_prefill::kInt8SplitKSharedBytes)); + int8_tc_decode_resident_ctas = multiprocessors * blocks_per_multiprocessor; + } } } if (use_grouped_decode) { - const int split_count = GetGroupedDecodeSplitCount(parameters, data.quantized_cache); + const int workspace_splits = GetGroupedDecodeSplitCount(parameters, data.quantized_cache); + const int split_count = + use_int8_tc_decode + ? GetInt8TcDecodeSplitCount(parameters, workspace_splits, data.decode_int8_tc_tiles_per_split, + int8_tc_decode_resident_ctas) + : workspace_splits; const int64_t partial_blocks = query_blocks * split_count; const int64_t grouped_blocks = static_cast(parameters.batch_size) * parameters.sequence_length * parameters.kv_num_heads * @@ -2116,9 +2217,27 @@ Status LaunchDynamicSparseAttention( float* partial_max = data.attention_workspace; float* partial_sum = partial_max + partial_blocks; float* partial_output = partial_sum + partial_blocks; + if (use_int8_tc_decode) { + // The partial vectors come first so that they stay 16-byte aligned whatever the number of partial blocks. + partial_output = data.attention_workspace; + partial_max = partial_output + partial_blocks * parameters.head_size; + partial_sum = partial_max + partial_blocks; + } const size_t grouped_shared_bytes = GetGroupedAttentionSharedBytes(parameters, data.quantized_cache ? sizeof(float) : sizeof(T)); - if (data.quantized_cache) { + if (use_int8_tc_decode) { + if constexpr (std::is_same_v) { + const dim3 grid(static_cast(parameters.batch_size * parameters.sequence_length), + static_cast(parameters.kv_num_heads), static_cast(split_count)); + dsa_prefill::DynamicSparseAttentionInt8TcKernel + <<>>( + data.prepared_query, reinterpret_cast(data.present_key), + reinterpret_cast(data.present_value), data.selected_indices, data.selected_counts, + data.seqlens_k, nullptr, partial_max, partial_sum, partial_output, split_count, + parameters.sequence_length, parameters.num_heads, parameters.kv_num_heads, parameters.cache_capacity, + parameters.max_selected, parameters.scale, data.key_scale, data.value_scale); + } + } else if (data.quantized_cache) { static const int vectorized_staging = ParseEnvironmentVariableWithDefault("ORT_DSA_INT8_VECTOR_STAGING", 1); ORT_ENFORCE(vectorized_staging == 0 || vectorized_staging == 1, @@ -2249,7 +2368,18 @@ Status LaunchDynamicSparseAttention( if constexpr (std::is_same_v) { const dim3 grid(static_cast(parameters.batch_size * parameters.sequence_length), static_cast(parameters.kv_num_heads)); - if (data.quantized_cache) { + if (use_int8_tc_prefill) { + CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( + dsa_prefill::DynamicSparseAttentionInt8TcKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(dsa_prefill::kInt8SharedBytes))); + dsa_prefill::DynamicSparseAttentionInt8TcKernel + <<>>( + data.prepared_query, reinterpret_cast(data.present_key), + reinterpret_cast(data.present_value), data.selected_indices, data.selected_counts, + data.seqlens_k, data.output, nullptr, nullptr, nullptr, 1, parameters.sequence_length, + parameters.num_heads, parameters.kv_num_heads, parameters.cache_capacity, parameters.max_selected, + parameters.scale, data.key_scale, data.value_scale); + } else if (data.quantized_cache) { CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( dsa_prefill::DynamicSparseAttentionPrefillKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(dsa_prefill::kSharedBytes))); diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h index baa7e566eec17..73f6e56db4231 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h @@ -50,6 +50,10 @@ struct DynamicSparseAttentionData { T* output = nullptr; float* attention_workspace = nullptr; bool quantized_cache = false; + bool prefill_kernel = true; + bool prefill_int8_tc = false; + bool decode_int8_tc = false; + int decode_int8_tc_tiles_per_split = 0; const float* key_scale = nullptr; const float* value_scale = nullptr; }; diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh new file mode 100644 index 0000000000000..a0982f94a0447 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh @@ -0,0 +1,812 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// INT8-cache variant of the tensor-core prefill kernel of DynamicSparseAttention (selected_only, main KV source, +// bfloat16 query / output, head_size 256). One CTA handles one (query row, KV head); the heads of the GQA group form +// the M dimension of an m16n8k16 MMA (padded to 16). +// +// The same body also serves decode (kSplitK = true): the grid gets a third dimension, each CTA consumes a contiguous +// range of the row's 64-candidate tiles and writes an un-normalized partial result (running maximum in the natural +// logarithm domain, running sum, value-scaled accumulator) in the layout of the grouped decode kernel, which the +// existing merge kernel combines. This replaces the CUDA-core grouped decode kernel, whose per-candidate dot product, +// warp reduction and exponentials made it issue-bound, with the tensor-core pipeline below. +// +// The scales of the INT8 cache are per (kv head, channel) constants, so they leave the tile loop: +// * the key scale is folded into the query once per CTA: Q' = Q * key_scale (rounded once to fp16), S = Q' K_int8. +// Rounding Q' moves the logit of any candidate by at most scale * 128 * sum_d |Q'_d - fp16(Q'_d)| (|K| <= 128). +// The CTA evaluates this bound for every query head and evaluates its heads in fp32 instead (AttendRowFp32) when +// it exceeds kInt8TcQueryErrorBudget, so that a query whose large terms cancel cannot amplify the rounding; the +// same happens when Q' does not fit fp16 or is not finite. The bound includes the flush of products below 2^-25, +// and the dispatcher selects this kernel only for scale <= kInt8TcMaxScale, which keeps the bound far below the +// budget for ordinary data; +// * the value scale is applied to the output once per CTA: O = (P V_int8) * value_scale. A probability below 2^-29 +// (more than kInt8TcTailDepth nats below the running maximum) leaves the normal range of the fp16 probability tile +// and loses precision (below 2^-40 it flushes to zero), which a large value scale could amplify. A CTA with a +// channel |value_scale| > kInt8TcMaxValueScale therefore tracks how far below the running maximum its valid +// candidates are and, only if one of them is deeper than kInt8TcTailDepth, discards the tile loop and evaluates +// its heads in fp32 as well. A value scale at or below kInt8TcMaxValueScale never takes this path. Above it, a CTA +// takes it when one of its heads has a candidate kInt8TcTailDepth nats below its running maximum, about 7 standard +// deviations of Gaussian logits over 2048 candidates: a fraction of the CTAs when the standard deviation is about +// 3 nats, practically none for 1 nat. The documentation lists measurements. +// The gathered int8 rows stay int8 in shared memory (cp.async, double buffered, one tile ahead). A fragment read with +// ldmatrix is converted to fp16 exactly in registers (1024 + (byte ^ 0x80) - 1152, "magic number" conversion) and +// consumed by fp16 MMAs with fp32 accumulation, so the per-tile dequantization pass, its shared memory round trip and +// its barriers of the BF16-dequantizing kernel disappear. The scores stay in registers; the only cross-warp exchanges +// per tile are the row maxima and the fp16 probability tile, which holds kInt8ProbScale * p so that the tail of the +// softmax stays out of the fp16 subnormal range. +// +// Scores and the running maximum are kept in the natural logarithm domain and only the difference to the new maximum +// is scaled by log2(e) for ex2: the difference of two nearby large logits is then exact, which the product with log2(e) +// would round away. Decode (kSplitK) splits the probability tile into an fp16 high part and an fp16 residual (about +// 22 significant bits together), because the signed sum of large value codes can cancel to a value far below the 2^-11 +// rounding of a single fp16 probability; both parts multiply the same converted V fragments. The prefill variant keeps +// the single fp16 tile (as does the BF16-dequantizing kernel, whose bfloat16 tile is coarser still). +// +// The channel (reduction) dimension of S = Q' K^T is permuted so that one 4-byte ldmatrix element of a K row supplies +// both B registers of an MMA k step: slot 2t+e <-> channel 16j + 4t + e and slot 2t+8+e <-> channel 16j + 4t + 2 + e, +// with the same permutation applied to the A fragments of Q'. Likewise ldmatrix.trans on the V tile delivers channels +// 2n and 2n + 1 of two consecutive candidates, i.e. the B registers of two n tiles whose columns are the even / odd +// channels. + +#pragma once + +#include "contrib_ops/cuda/bert/dynamic_sparse_attention_prefill.cuh" + +#include + +namespace onnxruntime { +namespace contrib { +namespace cuda { +namespace dsa_prefill { + +constexpr int kInt8RowStride = kHeadDim + 16; // bytes; 272 = 17 x 16 keeps ldmatrix conflict free +constexpr int kInt8TileBytes = kTile * kInt8RowStride; +constexpr int kInt8QStride = kHeadDim + 8; // fp16 elements of the staged query rows +constexpr int kInt8PStride = kTile + 8; // fp16 elements +constexpr int kIdxStages = 3; + +// The fp16 probability tile holds kInt8ProbScale * p for p in [0, 1]. Without the factor a candidate whose probability +// is below 2^-25 rounds to zero although a large value scale can make its contribution matter; with it the flush point +// drops to 2^-40 and the tile stays below the fp16 maximum. The PV accumulator carries the factor until the epilogue +// removes it; the softmax denominator is accumulated in fp32 without it. A probability below 2^-29 is a subnormal +// fp16 value of the tile (kInt8TcTailDepth nats below the running maximum). When some channel has +// |value_scale| > kInt8TcMaxValueScale (so that |p * V * value_scale| for p <= 2^-40 can exceed 2^-17) and a valid +// candidate is that deep, the CTA takes AttendRowFp32 instead. +constexpr float kInt8ProbScale = 32768.0f; // 2^15 +constexpr float kInt8ProbInverseScale = 1.0f / kInt8ProbScale; +constexpr float kInt8TcMaxValueScale = 65536.0f; // 2^16 +constexpr float kInt8TcTailDepth = 20.0f; // nats; 2^-29 is 20.1 + +// Largest softmax scale served by the kernel below (the op validates it as finite and non-negative). The error +// bound on Q * key_scale rounded to fp16 grows linearly with it, so a larger scale would send ordinary data to the +// fp32 fallback; the dispatcher uses the BF16-dequantizing prefill / CUDA-core decode kernels for it instead. +constexpr float kInt8TcMaxScale = 1.0f; + +// Largest bound on the change of a logit caused by rounding Q * key_scale to fp16 that the tensor-core path accepts: +// scale * 128 * sum_d |Q'_d - fp16(Q'_d)| of the query head (see the top). For a normal-range Q' the bound is at most +// 2^-4 * scale * sum_d |Q'_d|, an eighth of what rounding the cache to INT8 can move the same logit, so ordinary +// queries stay far below the budget, while a query whose large terms cancel to a small logit (the INT8 error of an +// exactly representable cache is zero there) exceeds it by orders of magnitude and is evaluated in fp32. The bound is +// about 0.015 for unit-RMS queries scaled by the q-norm weights and the key scales of the Qwen3.8-Flash-Next DSA +// layers (docs/contrib_ops/cuda/dynamic_sparse_attention.md lists how it grows with the query magnitude). +constexpr float kInt8TcQueryErrorBudget = 0.125f; + +// sK8[2] + sV8[2] (sQ aliases sK8[0] before the tile loop) + candidate indices + row maxima / sums + probabilities. +constexpr size_t kInt8PBytes = kMaxGroup * kInt8PStride * sizeof(__half); +constexpr size_t kInt8SharedBytes = 4 * static_cast(kInt8TileBytes) + kIdxStages * kTile * sizeof(int32_t) + + kMaxGroup * 4 * sizeof(float) + kInt8PBytes; +// The split-K (decode) variant adds the tile of fp16 probability residuals. +constexpr size_t kInt8SplitKSharedBytes = kInt8SharedBytes + kInt8PBytes; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 +__device__ __forceinline__ void CpAsync16Zfill(uint32_t smem_dst, const void* gmem_src, bool valid) { + const int src_size = valid ? 16 : 0; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(smem_dst), "l"(gmem_src), "r"(src_size)); +} + +__device__ __forceinline__ void LdmatrixX4Addr(uint32_t (&r)[4], uint32_t addr) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void LdmatrixX4TransAddr(uint32_t (&r)[4], uint32_t addr) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void MmaF16(float (&c)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) { + asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ float FastExp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;\n" : "=f"(y) : "f"(x)); + return y; +} + +// Two bytes b0, b1 (offset binary u = b ^ 0x80 in [0, 255]) placed in the mantissa of fp16 1024: 1024 + u; subtracting +// 1152 gives the signed value exactly. +__device__ __forceinline__ uint32_t Half2MinusBias(uint32_t x) { + uint32_t r; + asm("sub.rn.f16x2 %0, %1, %2;\n" : "=r"(r) : "r"(x), "r"(0x64806480u)); + return r; +} + +// bytes (0, 1) -> lo, bytes (2, 3) -> hi (two fp16 each). +__device__ __forceinline__ void Int8x4ToHalf2Adjacent(uint32_t w, uint32_t& lo, uint32_t& hi) { + const uint32_t x = w ^ 0x80808080u; + uint32_t a; + uint32_t b; + asm("prmt.b32 %0, %1, %2, 0x5140;\n" : "=r"(a) : "r"(x), "r"(0x64646464u)); + asm("prmt.b32 %0, %1, %2, 0x5342;\n" : "=r"(b) : "r"(x), "r"(0x64646464u)); + lo = Half2MinusBias(a); + hi = Half2MinusBias(b); +} + +// bytes (0, 2) -> even, bytes (1, 3) -> odd (two fp16 each). +__device__ __forceinline__ void Int8x4ToHalf2Strided(uint32_t w, uint32_t& even, uint32_t& odd) { + const uint32_t x = w ^ 0x80808080u; + uint32_t a; + uint32_t b; + asm("prmt.b32 %0, %1, %2, 0x5240;\n" : "=r"(a) : "r"(x), "r"(0x64646464u)); + asm("prmt.b32 %0, %1, %2, 0x5341;\n" : "=r"(b) : "r"(x), "r"(0x64646464u)); + even = Half2MinusBias(a); + odd = Half2MinusBias(b); +} + +// FP32 evaluation of the query heads of one (row, KV head), one warp per head and eight channels per lane, for CTAs +// the tensor-core path cannot serve accurately: Q * key_scale leaves the FP16 range or its rounding error bound exceeds +// kInt8TcQueryErrorBudget (clipping or rounding would break the cancellation between large terms), or a value_scale +// above kInt8TcMaxValueScale meets a probability below the normal range of the FP16 probability tile (the tile flushes +// below 2^-40 and a large value scale can amplify that). `query` and `output` point at the first head of +// the group; `key_scale` and `value_scale` at the KV head. The candidates [begin, end) of `row_indices` are consumed: +// all of them for prefill, the range of the CTA for a decode split. With kSplitK the un-normalized partial (running +// maximum in the natural logarithm domain, running sum, value-scaled accumulator) of block +// partial_block + head * split_count goes to partial_*, as the main path does; otherwise the normalized result goes to +// output. Every candidate is read from global memory once per head and handled serially, so this path is one to two +// orders of magnitude slower than the tensor-core one. +template +__device__ __noinline__ void AttendRowFp32(const __nv_bfloat16* __restrict__ query, const int8_t* __restrict__ key, + const int8_t* __restrict__ value, const float* __restrict__ key_scale, + const float* __restrict__ value_scale, + const int32_t* __restrict__ row_indices, int begin, int end, int total, + int query_position, float scale, int group, + __nv_bfloat16* __restrict__ output, float* __restrict__ partial_max, + float* __restrict__ partial_sum, float* __restrict__ partial_output, + int partial_block, int split_count, int warp, int lane) { + float ks[8]; + float vs[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + ks[j] = key_scale[lane + j * 32]; + vs[j] = value_scale[lane + j * 32]; + } + for (int h = warp; h < group; h += kThreads / 32) { + float q[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + q[j] = __bfloat162float(query[static_cast(h) * kHeadDim + lane + j * 32]); + } + float numerator[8] = {}; + float maximum = -FLT_MAX; + float denominator = 0.0f; + for (int i = begin; i < end; ++i) { + const int index = row_indices[i]; + if (index < 0 || index >= total || index > query_position) { + continue; + } + const int64_t offset = static_cast(index) * kHeadDim + lane; + float dot = 0.0f; +#pragma unroll + for (int j = 0; j < 8; ++j) { + dot += q[j] * (static_cast(key[offset + j * 32]) * ks[j]); + } +#pragma unroll + for (int delta = 16; delta > 0; delta /= 2) { + dot += __shfl_xor_sync(0xffffffffu, dot, delta); + } + const float score = dot * scale; + const float next_max = fmaxf(maximum, score); + const float correction = expf(maximum - next_max); + const float diff = score - next_max; + // Shift by 2^80 before expf when diff < -80 so tiny probabilities that value_scale brings back into range + // do not underflow fp32 before multiplying by vs[j]. + const bool shift_exp = diff < -80.0f; + const float probability = expf(shift_exp ? diff + 80.0f * 0.693147180559945309f : diff); + const float prob_unscale = shift_exp ? 0x1p-80f : 1.0f; + denominator = denominator * correction + (shift_exp ? 0.0f : probability); +#pragma unroll + for (int j = 0; j < 8; ++j) { + const float pv = probability * static_cast(value[offset + j * 32]); + numerator[j] = numerator[j] * correction + pv * (vs[j] * prob_unscale); + } + maximum = next_max; + } + if constexpr (kSplitK) { + const int block = partial_block + h * split_count; + if (lane == 0) { + partial_max[block] = maximum; + partial_sum[block] = denominator; + } +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int d = lane + j * 32; + partial_output[static_cast(block) * kHeadDim + d] = numerator[j]; + } + } else { + const float inverse = denominator == 0.0f ? 0.0f : 1.0f / denominator; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int d = lane + j * 32; + output[static_cast(h) * kHeadDim + d] = __float2bfloat16_rn(numerator[j] * inverse); + } + } + } +} +#endif + +// grid = (batch * sequence_length, kv_num_heads) or, with kSplitK, (batch * sequence_length, kv_num_heads, +// split_count); block = kThreads, dynamic smem = kInt8SharedBytes (kInt8SplitKSharedBytes with kSplitK). +// query: [rows, num_heads, 256] (already normalized / rotated); key, value: int8 [batch, kv_heads, capacity, 256]; +// key_scale, value_scale: float [kv_heads, 256]. The normalized result goes to output; with kSplitK the partials +// go to partial_max / partial_sum [rows * num_heads * split_count] and partial_output [.., 256] instead, indexed +// (row * num_heads + head) * split_count + split (output is unused; partial_output must be 16-byte aligned). +template +__global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKernel( + const __nv_bfloat16* __restrict__ query, + const int8_t* __restrict__ key, + const int8_t* __restrict__ value, + const int32_t* __restrict__ selected_indices, + const int32_t* __restrict__ selected_counts, + const int32_t* __restrict__ seqlens_k, + __nv_bfloat16* __restrict__ output, + float* __restrict__ partial_max, + float* __restrict__ partial_sum, + float* __restrict__ partial_output, + int split_count, + int sequence_length, + int num_heads, + int kv_num_heads, + int capacity, + int max_selected, + float scale, + const float* __restrict__ key_scale, + const float* __restrict__ value_scale) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + extern __shared__ __align__(16) unsigned char smem_raw[]; + unsigned char* sK8 = smem_raw; + unsigned char* sV8 = smem_raw + 2 * kInt8TileBytes; + int32_t* sIdx = reinterpret_cast(smem_raw + 4 * kInt8TileBytes); + float* sM = reinterpret_cast(sIdx + kIdxStages * kTile); + __half* sP = reinterpret_cast<__half*>(sM + kMaxGroup * 4); + __half* sQ = reinterpret_cast<__half*>(smem_raw); + // kSplitK only: fp16 residual of the probability tile, right behind sP (see kInt8SplitKSharedBytes). + [[maybe_unused]] __half* sPResidual = sP + kMaxGroup * kInt8PStride; + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int r0 = lane >> 2; // MMA row of the C / A fragments (and row + 8) + const int t = lane & 3; + const int row = blockIdx.x; + const int kv_head = blockIdx.y; + const int batch = row / sequence_length; + const int s = row - batch * sequence_length; + const int group = num_heads / kv_num_heads; + + const int64_t total_length = static_cast(seqlens_k[batch]) + 1; + const bool valid_length = total_length >= sequence_length && total_length <= capacity; + const int query_position = valid_length ? static_cast(total_length) - sequence_length + s : -1; + const int total = valid_length ? static_cast(total_length) : 0; + const int count = valid_length ? max(0, min(selected_counts[row], max_selected)) : 0; + const int32_t* row_indices = selected_indices + static_cast(row) * max_selected; + const int64_t kv_head_offset = (static_cast(batch) * kv_num_heads + kv_head) * capacity * kHeadDim; + const int8_t* key_base = key + kv_head_offset; + const int8_t* value_base = value + kv_head_offset; + + // The 64-candidate tiles this CTA consumes: all of the row's for prefill, a contiguous range for a decode split. + // A split without tiles only publishes an empty partial (maximum -FLT_MAX, sum 0), which the merge skips. + int tile_begin = 0; + int tiles = (count + kTile - 1) / kTile; + if constexpr (kSplitK) { + const int split = static_cast(blockIdx.z); + const int tiles_per_split = (tiles + split_count - 1) / split_count; + const int first = min(tiles, split * tiles_per_split); + tiles = min(tiles, first + tiles_per_split) - first; + tile_begin = first; + if (tiles == 0) { + if (tid < group) { + const int partial_block = (row * num_heads + kv_head * group + tid) * split_count + split; + partial_max[partial_block] = -FLT_MAX; + partial_sum[partial_block] = 0.0f; + } + return; + } + } + + // Q' = Q * key_scale as fp16; the padding rows are zero so their (ignored) scores stay finite. The whole CTA is + // evaluated in fp32 by AttendRowFp32 instead when Q' of a head does not fit fp16 (including inf / NaN), when the bound + // scale * 128 * sum_d |Q'_d - fp16(Q'_d)| on the rounding error of a logit exceeds kInt8TcQueryErrorBudget for a + // head, or when a value_scale is not finite. A finite value_scale above kInt8TcMaxValueScale only arms the check of + // the probability depth after the tile loop. + const __nv_bfloat16* query_group = query + (static_cast(row) * num_heads + kv_head * group) * kHeadDim; + bool needs_fp32 = false; + bool large_value_scale = false; + { + const int chunk = tid & 31; + const float4 k0 = *reinterpret_cast(key_scale + kv_head * kHeadDim + chunk * 8); + const float4 k1 = *reinterpret_cast(key_scale + kv_head * kHeadDim + chunk * 8 + 4); + if (tid < kHeadDim / 4) { + const float4 v = *reinterpret_cast(value_scale + kv_head * kHeadDim + tid * 4); + needs_fp32 = !(fabsf(v.x) <= FLT_MAX) || !(fabsf(v.y) <= FLT_MAX) || !(fabsf(v.z) <= FLT_MAX) || + !(fabsf(v.w) <= FLT_MAX); + large_value_scale = fabsf(v.x) > kInt8TcMaxValueScale || fabsf(v.y) > kInt8TcMaxValueScale || + fabsf(v.z) > kInt8TcMaxValueScale || fabsf(v.w) > kInt8TcMaxValueScale; + } + // Sum of |Q' - fp16(Q')| over the eight channels of this lane (NaN when Q' is not finite or leaves fp16). + float rounding_error = 0.0f; + auto to_half2 = [&needs_fp32, &rounding_error](float x, float y) { + needs_fp32 |= !(fabsf(x) <= 65504.0f) || !(fabsf(y) <= 65504.0f); + const __half2 h = __floats2half2_rn(x, y); + const float2 rounded = __half22float2(h); + rounding_error += fabsf(x - rounded.x) + fabsf(y - rounded.y); + return h; + }; + const float error_scale = scale * 128.0f; +#pragma unroll + for (int i = 0; i < kMaxGroup / 4; ++i) { + const int r = (tid >> 5) + 4 * i; + uint4 packed = make_uint4(0, 0, 0, 0); + if (r < group) { + rounding_error = 0.0f; + const uint4 q = *reinterpret_cast(query_group + static_cast(r) * kHeadDim + chunk * 8); + const float2 q01 = __bfloat1622float2(*reinterpret_cast(&q.x)); + const float2 q23 = __bfloat1622float2(*reinterpret_cast(&q.y)); + const float2 q45 = __bfloat1622float2(*reinterpret_cast(&q.z)); + const float2 q67 = __bfloat1622float2(*reinterpret_cast(&q.w)); + const __half2 h01 = to_half2(q01.x * k0.x, q01.y * k0.y); + const __half2 h23 = to_half2(q23.x * k0.z, q23.y * k0.w); + const __half2 h45 = to_half2(q45.x * k1.x, q45.y * k1.y); + const __half2 h67 = to_half2(q67.x * k1.z, q67.y * k1.w); + packed.x = *reinterpret_cast(&h01); + packed.y = *reinterpret_cast(&h23); + packed.z = *reinterpret_cast(&h45); + packed.w = *reinterpret_cast(&h67); + // The warp covers all 256 channels of row r (r < group is uniform across it): a butterfly sum gives the + // bound on the logit error of head r in every lane. +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + rounding_error += __shfl_xor_sync(0xffffffffu, rounding_error, offset); + } + needs_fp32 |= !(rounding_error * error_scale <= kInt8TcQueryErrorBudget); + } + *reinterpret_cast(sQ + r * kInt8QStride + chunk * 8) = packed; + } + } + // The condition is uniform across the CTA and no cp.async is in flight yet. + auto attend_fp32 = [&]() { + const float* kv_key_scale = key_scale + kv_head * kHeadDim; + const float* kv_value_scale = value_scale + kv_head * kHeadDim; + if constexpr (kSplitK) { + const int partial_block = (row * num_heads + kv_head * group) * split_count + static_cast(blockIdx.z); + AttendRowFp32(query_group, key_base, value_base, kv_key_scale, kv_value_scale, row_indices, + tile_begin * kTile, min(count, (tile_begin + tiles) * kTile), total, query_position, scale, + group, nullptr, partial_max, partial_sum, partial_output, partial_block, split_count, warp, + lane); + } else { + AttendRowFp32(query_group, key_base, value_base, kv_key_scale, kv_value_scale, row_indices, 0, count, + total, query_position, scale, group, + output + (static_cast(row) * num_heads + kv_head * group) * kHeadDim, nullptr, + nullptr, nullptr, 0, 1, warp, lane); + } + }; + if (__syncthreads_or(needs_fp32)) { + attend_fp32(); + return; + } + + // A fragments of Q' with the channel permutation described at the top: a0 / a2 hold channels 16j + 4t + {0, 1} / + // {2, 3} of row r0, a1 / a3 the same of row r0 + 8. + uint32_t qf[16][4]; +#pragma unroll + for (int j = 0; j < 16; ++j) { + const uint2 lo = *reinterpret_cast(sQ + r0 * kInt8QStride + 16 * j + 4 * t); + const uint2 hi = *reinterpret_cast(sQ + (r0 + 8) * kInt8QStride + 16 * j + 4 * t); + qf[j][0] = lo.x; + qf[j][1] = hi.x; + qf[j][2] = lo.y; + qf[j][3] = hi.y; + } + // sQ aliases sK8[0]: this barrier also tells the CTA whether the probability depth has to be checked (see below). + const bool check_tail_depth = __syncthreads_or(large_value_scale); + + auto resolve_index = [&](int raw) -> int { + return (raw >= 0 && raw < total && raw <= query_position && raw < capacity) ? raw : -1; + }; + + const int first_candidate = tile_begin * kTile; + if ((tid >> 6) < tiles) { + const int tile = tid >> 6; + const int candidate = first_candidate + tile * kTile + (tid & 63); + sIdx[tile * kTile + (tid & 63)] = candidate < count ? resolve_index(row_indices[candidate]) : -1; + } + __syncthreads(); + + const uint32_t k_smem_base = static_cast(__cvta_generic_to_shared(sK8)); + const uint32_t v_smem_base = static_cast(__cvta_generic_to_shared(sV8)); + const uint32_t p_smem_base = static_cast(__cvta_generic_to_shared(sP)); + [[maybe_unused]] const uint32_t p_residual_smem_base = p_smem_base + static_cast(kInt8PBytes); + const int load_chunk = tid & 15; + + // 16 B per lane: 16 lanes per candidate row, 8 rows per pass, 8 passes per tile (K and V). + auto issue_tile = [&](int tile, int stage, int idx_buf) { +#pragma unroll + for (int i = 0; i < kTile / 8; ++i) { + const int cand = i * 8 + (tid >> 4); + const int index = sIdx[idx_buf * kTile + cand]; + const bool valid = index >= 0; + const int64_t source = static_cast(valid ? index : 0) * kHeadDim + load_chunk * 16; + const uint32_t offset = static_cast(stage * kInt8TileBytes + cand * kInt8RowStride + load_chunk * 16); + CpAsync16Zfill(k_smem_base + offset, key_base + source, valid); + CpAsync16Zfill(v_smem_base + offset, value_base + source, valid); + } + }; + + if (tiles > 0) { + issue_tile(0, 0, 0); + } + CpAsyncCommit(); + + float acc_o[8][4]; +#pragma unroll + for (int n = 0; n < 8; ++n) { + acc_o[n][0] = acc_o[n][1] = acc_o[n][2] = acc_o[n][3] = 0.0f; + } + // Softmax state of rows r0 / r0 + 8: the running maximum (natural logarithm domain) is identical in all threads of + // the row, the running sum only covers this thread's columns (merged across threads and warps after the loop). + float run_max0 = -FLT_MAX; + float run_max1 = -FLT_MAX; + float run_sum0 = 0.0f; + float run_sum1 = 0.0f; + // Most negative (score - running maximum) of a valid candidate seen by this thread; only tracked when a value scale + // is large. + float lowest_depth = 0.0f; + constexpr float kLog2e = 1.44269504088896341f; + + // Shared memory byte addresses that do not change from tile to tile. + const uint32_t k_lane_offset = + static_cast((warp * 16 + ((lane >> 3) & 1) * 8 + (lane & 7)) * kInt8RowStride + (lane >> 4) * 16); + const uint32_t v_lane_offset = + static_cast((((lane >> 3) & 1) * 8 + (lane & 7)) * kInt8RowStride + (warp * 4 + (lane >> 4)) * 16); + const uint32_t p_lane_offset = static_cast(((lane & 15) * kInt8PStride + (lane >> 4) * 8) * sizeof(__half)); + + int idx_cur = 0; // tile % 3 + int idx_next = 1 % kIdxStages; // (tile + 1) % 3 + int idx_next2 = 2 % kIdxStages; // (tile + 2) % 3 + for (int tile = 0; tile < tiles; ++tile) { + CpAsyncWait<0>(); // K / V of this tile have landed (this thread's copies) + __syncthreads(); // (1) everyone's copies are visible; the previous tile's smem consumers are done + + const int stage = tile & 1; + if (tile + 1 < tiles) { + issue_tile(tile + 1, stage ^ 1, idx_next); + } + CpAsyncCommit(); + // Raw indices of tile + 2: loaded now, resolved and stored at the end of the iteration (hides the load latency). + int raw_next2 = -1; + const bool stage_next2 = (tile + 2 < tiles) && tid < kTile; + if (stage_next2) { + const int candidate = first_candidate + (tile + 2) * kTile + tid; + raw_next2 = candidate < count ? row_indices[candidate] : -1; + } + + // S = Q' K^T for this warp's 16 candidates (two n tiles of 8). + float acc_s[2][4]; +#pragma unroll + for (int n = 0; n < 2; ++n) { + acc_s[n][0] = acc_s[n][1] = acc_s[n][2] = acc_s[n][3] = 0.0f; + } + { + const uint32_t k_addr = k_smem_base + stage * kInt8TileBytes + k_lane_offset; +#pragma unroll + for (int jj = 0; jj < 8; ++jj) { + uint32_t r[4]; + LdmatrixX4Addr(r, k_addr + jj * 32); + uint32_t b0; + uint32_t b1; + Int8x4ToHalf2Adjacent(r[0], b0, b1); + MmaF16(acc_s[0], qf[2 * jj], b0, b1); + Int8x4ToHalf2Adjacent(r[1], b0, b1); + MmaF16(acc_s[1], qf[2 * jj], b0, b1); + Int8x4ToHalf2Adjacent(r[2], b0, b1); + MmaF16(acc_s[0], qf[2 * jj + 1], b0, b1); + Int8x4ToHalf2Adjacent(r[3], b0, b1); + MmaF16(acc_s[1], qf[2 * jj + 1], b0, b1); + } + } + + // Scores in the natural logarithm domain; invalid candidates are masked. + const int cbase = warp * 16 + 2 * t; + const int2 ia = *reinterpret_cast(sIdx + idx_cur * kTile + cbase); + const int2 ib = *reinterpret_cast(sIdx + idx_cur * kTile + cbase + 8); + const bool va0 = ia.x >= 0; + const bool va1 = ia.y >= 0; + const bool vb0 = ib.x >= 0; + const bool vb1 = ib.y >= 0; + const float s00 = va0 ? acc_s[0][0] * scale : -FLT_MAX; + const float s01 = va1 ? acc_s[0][1] * scale : -FLT_MAX; + const float s02 = vb0 ? acc_s[1][0] * scale : -FLT_MAX; + const float s03 = vb1 ? acc_s[1][1] * scale : -FLT_MAX; + const float s10 = va0 ? acc_s[0][2] * scale : -FLT_MAX; + const float s11 = va1 ? acc_s[0][3] * scale : -FLT_MAX; + const float s12 = vb0 ? acc_s[1][2] * scale : -FLT_MAX; + const float s13 = vb1 ? acc_s[1][3] * scale : -FLT_MAX; + float m0 = fmaxf(fmaxf(s00, s01), fmaxf(s02, s03)); + float m1 = fmaxf(fmaxf(s10, s11), fmaxf(s12, s13)); + m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 1)); + m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 1)); + m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 2)); + m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 2)); + if (t == 0) { + sM[r0 * 4 + warp] = m0; + sM[(r0 + 8) * 4 + warp] = m1; + } + __syncthreads(); // (2) the four warps' row maxima are visible + + const float4 mx0 = *reinterpret_cast(sM + r0 * 4); + const float4 mx1 = *reinterpret_cast(sM + (r0 + 8) * 4); + const float new_max0 = fmaxf(run_max0, fmaxf(fmaxf(mx0.x, mx0.y), fmaxf(mx0.z, mx0.w))); + const float new_max1 = fmaxf(run_max1, fmaxf(fmaxf(mx1.x, mx1.y), fmaxf(mx1.z, mx1.w))); + // The differences to the new maximum are formed in the natural logarithm domain, where nearby large logits are + // still distinct, and only then scaled by log2(e) for ex2. + const float alpha0 = FastExp2((run_max0 - new_max0) * kLog2e); + const float alpha1 = FastExp2((run_max1 - new_max1) * kLog2e); + const float d00 = s00 - new_max0; + const float d01 = s01 - new_max0; + const float d02 = s02 - new_max0; + const float d03 = s03 - new_max0; + const float d10 = s10 - new_max1; + const float d11 = s11 - new_max1; + const float d12 = s12 - new_max1; + const float d13 = s13 - new_max1; + const float p00 = va0 ? FastExp2(d00 * kLog2e) : 0.0f; + const float p01 = va1 ? FastExp2(d01 * kLog2e) : 0.0f; + const float p02 = vb0 ? FastExp2(d02 * kLog2e) : 0.0f; + const float p03 = vb1 ? FastExp2(d03 * kLog2e) : 0.0f; + const float p10 = va0 ? FastExp2(d10 * kLog2e) : 0.0f; + const float p11 = va1 ? FastExp2(d11 * kLog2e) : 0.0f; + const float p12 = vb0 ? FastExp2(d12 * kLog2e) : 0.0f; + const float p13 = vb1 ? FastExp2(d13 * kLog2e) : 0.0f; + run_sum0 = run_sum0 * alpha0 + ((p00 + p01) + (p02 + p03)); + run_sum1 = run_sum1 * alpha1 + ((p10 + p11) + (p12 + p13)); + run_max0 = new_max0; + run_max1 = new_max1; + if (check_tail_depth) { + const float low0 = fminf(fminf(va0 ? d00 : 0.0f, va1 ? d01 : 0.0f), fminf(vb0 ? d02 : 0.0f, vb1 ? d03 : 0.0f)); + const float low1 = fminf(fminf(va0 ? d10 : 0.0f, va1 ? d11 : 0.0f), fminf(vb0 ? d12 : 0.0f, vb1 ? d13 : 0.0f)); + lowest_depth = fminf(lowest_depth, fminf(low0, low1)); + } + { + const float x00 = p00 * kInt8ProbScale; + const float x01 = p01 * kInt8ProbScale; + const float x02 = p02 * kInt8ProbScale; + const float x03 = p03 * kInt8ProbScale; + const float x10 = p10 * kInt8ProbScale; + const float x11 = p11 * kInt8ProbScale; + const float x12 = p12 * kInt8ProbScale; + const float x13 = p13 * kInt8ProbScale; + const __half2 h00 = __floats2half2_rn(x00, x01); + const __half2 h01 = __floats2half2_rn(x02, x03); + const __half2 h10 = __floats2half2_rn(x10, x11); + const __half2 h11 = __floats2half2_rn(x12, x13); + *reinterpret_cast<__half2*>(sP + r0 * kInt8PStride + cbase) = h00; + *reinterpret_cast<__half2*>(sP + r0 * kInt8PStride + cbase + 8) = h01; + *reinterpret_cast<__half2*>(sP + (r0 + 8) * kInt8PStride + cbase) = h10; + *reinterpret_cast<__half2*>(sP + (r0 + 8) * kInt8PStride + cbase + 8) = h11; + if constexpr (kSplitK) { + // What the fp16 tile lost: P = P_hi + P_residual keeps about 22 significant bits. + const float2 f00 = __half22float2(h00); + const float2 f01 = __half22float2(h01); + const float2 f10 = __half22float2(h10); + const float2 f11 = __half22float2(h11); + *reinterpret_cast<__half2*>(sPResidual + r0 * kInt8PStride + cbase) = + __floats2half2_rn(x00 - f00.x, x01 - f00.y); + *reinterpret_cast<__half2*>(sPResidual + r0 * kInt8PStride + cbase + 8) = + __floats2half2_rn(x02 - f01.x, x03 - f01.y); + *reinterpret_cast<__half2*>(sPResidual + (r0 + 8) * kInt8PStride + cbase) = + __floats2half2_rn(x10 - f10.x, x11 - f10.y); + *reinterpret_cast<__half2*>(sPResidual + (r0 + 8) * kInt8PStride + cbase + 8) = + __floats2half2_rn(x12 - f11.x, x13 - f11.y); + } + } + __syncthreads(); // (3) the probability tile (and its residual) is visible + + // The maximum rarely changes after the first tiles: skip the accumulator rescale when the whole warp is unchanged. + if (__any_sync(0xffffffffu, alpha0 != 1.0f || alpha1 != 1.0f)) { +#pragma unroll + for (int n = 0; n < 8; ++n) { + acc_o[n][0] *= alpha0; + acc_o[n][1] *= alpha0; + acc_o[n][2] *= alpha1; + acc_o[n][3] *= alpha1; + } + } + + // O += P V for this warp's 64 output channels (four 16-channel chunks, each an even and an odd n tile). With kSplitK + // the residual tile multiplies the same converted V fragments. + { + const uint32_t v_addr = v_smem_base + stage * kInt8TileBytes + v_lane_offset; +#pragma unroll + for (int kk = 0; kk < kTile / 16; ++kk) { + uint32_t a[4]; + LdmatrixX4Addr(a, p_smem_base + p_lane_offset + kk * 32); + uint32_t a_residual[4] = {0, 0, 0, 0}; + if constexpr (kSplitK) { + LdmatrixX4Addr(a_residual, p_residual_smem_base + p_lane_offset + kk * 32); + } +#pragma unroll + for (int cp = 0; cp < 2; ++cp) { + uint32_t r[4]; + LdmatrixX4TransAddr(r, v_addr + kk * 16 * kInt8RowStride + cp * 32); + uint32_t e0; + uint32_t o0; + uint32_t e1; + uint32_t o1; + // r[0] / r[1]: candidates 0-7 / 8-15 of the k step, chunk 2cp; r[2] / r[3]: the same for chunk 2cp + 1. + Int8x4ToHalf2Strided(r[0], e0, o0); + Int8x4ToHalf2Strided(r[1], e1, o1); + MmaF16(acc_o[4 * cp], a, e0, e1); + MmaF16(acc_o[4 * cp + 1], a, o0, o1); + if constexpr (kSplitK) { + MmaF16(acc_o[4 * cp], a_residual, e0, e1); + MmaF16(acc_o[4 * cp + 1], a_residual, o0, o1); + } + Int8x4ToHalf2Strided(r[2], e0, o0); + Int8x4ToHalf2Strided(r[3], e1, o1); + MmaF16(acc_o[4 * cp + 2], a, e0, e1); + MmaF16(acc_o[4 * cp + 3], a, o0, o1); + if constexpr (kSplitK) { + MmaF16(acc_o[4 * cp + 2], a_residual, e0, e1); + MmaF16(acc_o[4 * cp + 3], a_residual, o0, o1); + } + } + } + } + + if (stage_next2) { + sIdx[idx_next2 * kTile + tid] = resolve_index(raw_next2); + } + idx_cur = idx_next; + idx_next = idx_next2; + idx_next2 = idx_next2 + 1 == kIdxStages ? 0 : idx_next2 + 1; + } + + // A large value scale amplifies probabilities that the fp16 tile represents with a large relative error or not at + // all (below 2^-29). The condition is uniform across the CTA and no cp.async copy is in flight here; the tile loop + // wrote nothing to global memory, so the fp32 path simply starts over from the inputs. + if (check_tail_depth && __syncthreads_or(lowest_depth < -kInt8TcTailDepth)) { + attend_fp32(); + return; + } + + // Merge the per-thread running sums: across the four threads of a row, then across the four warps. + float l0 = run_sum0; + float l1 = run_sum1; + l0 += __shfl_xor_sync(0xffffffffu, l0, 1); + l1 += __shfl_xor_sync(0xffffffffu, l1, 1); + l0 += __shfl_xor_sync(0xffffffffu, l0, 2); + l1 += __shfl_xor_sync(0xffffffffu, l1, 2); + __syncthreads(); + if (t == 0) { + sM[r0 * 4 + warp] = l0; + sM[(r0 + 8) * 4 + warp] = l1; + } + __syncthreads(); + const float4 ls0 = *reinterpret_cast(sM + r0 * 4); + const float4 ls1 = *reinterpret_cast(sM + (r0 + 8) * 4); + const float sum0 = (ls0.x + ls0.y) + (ls0.z + ls0.w); + const float sum1 = (ls1.x + ls1.y) + (ls1.z + ls1.w); + // acc_o carries kInt8ProbScale, which the normalization removes. + const float inv0 = sum0 == 0.0f ? 0.0f : kInt8ProbInverseScale / sum0; + const float inv1 = sum1 == 0.0f ? 0.0f : kInt8ProbInverseScale / sum1; + + if constexpr (kSplitK) { + // Un-normalized partial of this split in the natural logarithm domain, as the merge kernel expects. Thread t owns + // channels 64 * warp + 16 * chunk + 4t + {0, 1, 2, 3} (see below), i.e. one float4 per chunk and row. + const int split = static_cast(blockIdx.z); + const int partial_block0 = (row * num_heads + kv_head * group + r0) * split_count + split; + const int partial_block1 = partial_block0 + 8 * split_count; + if (warp == 0 && t == 0) { + if (r0 < group) { + partial_max[partial_block0] = run_max0; + partial_sum[partial_block0] = sum0; + } + if (r0 + 8 < group) { + partial_max[partial_block1] = run_max1; + partial_sum[partial_block1] = sum1; + } + } +#pragma unroll + for (int chunk = 0; chunk < 4; ++chunk) { + const int dim = warp * 64 + chunk * 16 + t * 4; + // acc_o carries kInt8ProbScale (the partial_sum above does not). Removing it before the multiplication by the + // value scale keeps a tiny scale (or the product) from underflowing in the factor. + const float4 vs = *reinterpret_cast(value_scale + kv_head * kHeadDim + dim); + if (r0 < group) { + *reinterpret_cast(partial_output + static_cast(partial_block0) * kHeadDim + dim) = + make_float4((acc_o[2 * chunk][0] * kInt8ProbInverseScale) * vs.x, + (acc_o[2 * chunk + 1][0] * kInt8ProbInverseScale) * vs.y, + (acc_o[2 * chunk][1] * kInt8ProbInverseScale) * vs.z, + (acc_o[2 * chunk + 1][1] * kInt8ProbInverseScale) * vs.w); + } + if (r0 + 8 < group) { + *reinterpret_cast(partial_output + static_cast(partial_block1) * kHeadDim + dim) = + make_float4((acc_o[2 * chunk][2] * kInt8ProbInverseScale) * vs.x, + (acc_o[2 * chunk + 1][2] * kInt8ProbInverseScale) * vs.y, + (acc_o[2 * chunk][3] * kInt8ProbInverseScale) * vs.z, + (acc_o[2 * chunk + 1][3] * kInt8ProbInverseScale) * vs.w); + } + } + } else { + // Thread t owns channels 64 * warp + 16 * chunk + 4t + {0, 1, 2, 3} + // = (even tile c0, odd tile c0, even tile c1, odd tile c1). +#pragma unroll + for (int chunk = 0; chunk < 4; ++chunk) { + const int dim = warp * 64 + chunk * 16 + t * 4; + const float4 vs = *reinterpret_cast(value_scale + kv_head * kHeadDim + dim); + if (r0 < group) { + __nv_bfloat16* dst = output + (static_cast(row) * num_heads + kv_head * group + r0) * kHeadDim + dim; + const __nv_bfloat162 lo = + __floats2bfloat162_rn(acc_o[2 * chunk][0] * inv0 * vs.x, acc_o[2 * chunk + 1][0] * inv0 * vs.y); + const __nv_bfloat162 hi = + __floats2bfloat162_rn(acc_o[2 * chunk][1] * inv0 * vs.z, acc_o[2 * chunk + 1][1] * inv0 * vs.w); + uint2 packed; + packed.x = *reinterpret_cast(&lo); + packed.y = *reinterpret_cast(&hi); + *reinterpret_cast(dst) = packed; + } + if (r0 + 8 < group) { + __nv_bfloat16* dst = + output + (static_cast(row) * num_heads + kv_head * group + r0 + 8) * kHeadDim + dim; + const __nv_bfloat162 lo = + __floats2bfloat162_rn(acc_o[2 * chunk][2] * inv1 * vs.x, acc_o[2 * chunk + 1][2] * inv1 * vs.y); + const __nv_bfloat162 hi = + __floats2bfloat162_rn(acc_o[2 * chunk][3] * inv1 * vs.z, acc_o[2 * chunk + 1][3] * inv1 * vs.w); + uint2 packed; + packed.x = *reinterpret_cast(&lo); + packed.y = *reinterpret_cast(&hi); + *reinterpret_cast(dst) = packed; + } + } + } +#else + (void)query; + (void)key; + (void)value; + (void)selected_indices; + (void)selected_counts; + (void)seqlens_k; + (void)output; + (void)partial_max; + (void)partial_sum; + (void)partial_output; + (void)split_count; + (void)sequence_length; + (void)num_heads; + (void)kv_num_heads; + (void)capacity; + (void)max_selected; + (void)scale; + (void)key_scale; + (void)value_scale; +#endif +} + +} // namespace dsa_prefill +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu index 82aa1baf8d463..8f6139e456c30 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu @@ -56,6 +56,11 @@ __global__ void ProjectedSigmoidGateMulKernel( const __nv_bfloat16* attention, const __nv_bfloat16* activation, const __nv_bfloat16* weight, const __nv_bfloat16* residual, __nv_bfloat16* output, int64_t input_width, int64_t output_width) { + // Let a programmatic-dependent-launch successor (the hyper-connection down kernel) start and prefetch its constant + // weights while this kernel runs; it reads our output only after griddepcontrol.wait. No-op for plain launches. +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif const int64_t row = blockIdx.x; const float gate = ProjectedRowGate(activation, weight, row, input_width); const int64_t row_base = row * output_width; diff --git a/onnxruntime/contrib_ops/cuda/cuda_session_options.h b/onnxruntime/contrib_ops/cuda/cuda_session_options.h index 93f6ffcb64876..0f2f06bbe0e6d 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_session_options.h +++ b/onnxruntime/contrib_ops/cuda/cuda_session_options.h @@ -13,13 +13,16 @@ namespace onnxruntime::contrib::cuda { -// An explicit configured value takes precedence over the environment. Throws unless the result is an integer in -// [min_value, max_value]. +// An explicit configured value takes precedence over the environment. +inline std::string ResolveCudaSessionOptionText(const std::optional& configured, const char* env, + const std::string& default_text) { + return configured.has_value() ? *configured : ParseEnvironmentVariableWithDefault(env, default_text); +} + +// Integer threshold: throws unless the result is an integer in [min_value, max_value]. inline int ResolveCudaSessionOption(const std::optional& configured, const char* key, const char* env, int default_value, int min_value, int max_value) { - const std::string value = configured.has_value() - ? *configured - : ParseEnvironmentVariableWithDefault(env, std::to_string(default_value)); + const std::string value = ResolveCudaSessionOptionText(configured, env, std::to_string(default_value)); int parsed = 0; const auto result = std::from_chars(value.data(), value.data() + value.size(), parsed); ORT_ENFORCE(result.ec == std::errc{} && result.ptr == value.data() + value.size() && @@ -28,6 +31,14 @@ inline int ResolveCudaSessionOption(const std::optional& configured return parsed; } +// 0/1 switch: throws unless the result is exactly "0" or "1" (an integer parse would also take "00", "01" or "-0"). +inline bool ResolveCudaSessionSwitch(const std::optional& configured, const char* key, const char* env, + bool default_value) { + const std::string value = ResolveCudaSessionOptionText(configured, env, default_value ? "1" : "0"); + ORT_ENFORCE(value == "0" || value == "1", key, " (", env, ") must be exactly \"0\" or \"1\", got '", value, "'."); + return value == "1"; +} + // Call only from kernel construction. Explicit session configuration takes precedence over the environment. inline int ReadCudaSessionOption(const OpKernelInfo& info, const char* key, const char* env, int default_value, int min_value, int max_value) { @@ -35,4 +46,8 @@ inline int ReadCudaSessionOption(const OpKernelInfo& info, const char* key, cons max_value); } +inline bool ReadCudaSessionSwitch(const OpKernelInfo& info, const char* key, const char* env, bool default_value) { + return ResolveCudaSessionSwitch(info.GetConfigOptions().GetConfigEntry(key), key, env, default_value); +} + } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 003012a5a6a1b..a737731f104e3 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -4,10 +4,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -31,6 +33,7 @@ #ifdef USE_CUDA #include "core/framework/to_tensor_proto_element_type.h" #include "core/providers/cuda/cuda_provider_options.h" +#include "core/session/onnxruntime_session_options_config_keys.h" #include "test/common/cuda_op_test_utils.h" // Uses std::system, which is unavailable on iOS, so only CUDA builds (Windows/Linux) include it. #include "test/contrib_ops/fresh_process_test_util.h" @@ -99,6 +102,11 @@ struct DynamicSparseAttentionCase { std::vector expected_output; std::vector expected_present_key; std::vector expected_present_value; + + // |actual - expected| <= absolute + relative * |expected| for the output; a negative relative tolerance keeps the + // default of the output type. + float output_abs_tolerance = 0.005f; + float output_rel_tolerance = -1.0f; }; template @@ -120,7 +128,8 @@ void RunDynamicSparseAttentionCase( const DynamicSparseAttentionCase& c, std::unique_ptr execution_provider, OpTester::ExpectResult expected_result = OpTester::ExpectResult::kExpectSuccess, - const std::string& expected_error = "") { + const std::string& expected_error = "", + const std::vector>& session_configs = {}) { ASSERT_NE(execution_provider, nullptr); const int64_t hidden_size = c.num_heads * c.head_size; @@ -263,11 +272,19 @@ void RunDynamicSparseAttentionCase( if (c.expected_validation_status.has_value()) { tester.AddOutput("validation_status", {1}, {*c.expected_validation_status}); } - tester.SetOutputTolerance(0.005f); + tester.SetOutputTolerance(c.output_abs_tolerance, c.output_rel_tolerance); std::vector> execution_providers; execution_providers.push_back(std::move(execution_provider)); - tester.Run(expected_result, expected_error, {}, nullptr, &execution_providers); + // The options of the default OpTester::Run overload plus the requested session configuration entries. + SessionOptions session_options; + session_options.use_per_session_threads = false; + session_options.session_log_verbosity_level = 1; + session_options.graph_optimization_level = TransformerLevel::Default; + for (const auto& [key, value] : session_configs) { + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(key, value)); + } + tester.Run(session_options, expected_result, expected_error, {}, nullptr, &execution_providers); } DynamicSparseAttentionCase MakeSingleTokenSelectedOnlyCase() { @@ -593,32 +610,50 @@ TEST(DynamicSparseAttentionTest, GroupedDecodeRatio12ReferenceParity_CUDA) { } } -// Selected-only prefill without a softmax sink, the configuration the tensor-core BF16 prefill kernel serves. -DynamicSparseAttentionCase MakeSelectedOnlyPrefillCase(bool int8_cache) { - auto c = int8_cache ? MakeGroupedCacheCase(9, false, false, 2, 20, 2, true) - : MakeGroupedCacheCase(9, false, true, 2, 20, 2); - c.head_sink.clear(); - c.smooth_softmax = 0; +// Selected-only attention without a softmax sink, the configuration the tensor-core kernels serve. The inputs are first +// rounded the way the operator receives them (the query and a floating-point cache as BFloat16; the codes of an INT8 +// cache are exact, its scales stay float), then the scores, the max-subtracted weights, the numerator and the +// denominator are evaluated in double precision, so that neither a large common logit nor a tiny weight overflows or +// underflows the reference. With invalid_entries every fifth selected index is replaced by one that must be skipped +// (negative, in the future of the query row, or outside the cache). +void ComputeSelectedOnlyPrefillExpectedOutput(DynamicSparseAttentionCase& c, bool int8_cache) { + const auto bf16 = [](float value) { return static_cast(static_cast(BFloat16(value))); }; + std::vector scores; + std::vector offsets; + std::vector numerator(c.head_size); for (int row = 0; row < c.batch_size * c.sequence_length; ++row) { const int64_t b = row / c.sequence_length; + const int64_t position = c.seqlens_k[b] + 1 - c.sequence_length + row % c.sequence_length; for (int h = 0; h < c.num_heads; ++h) { const int64_t kv = h / (c.num_heads / c.kv_num_heads); const int64_t out = (row * c.num_heads + h) * c.head_size; - std::vector numerator(c.head_size, 0.0f); - double denominator = 0.0; + scores.clear(); + offsets.clear(); for (int i = 0; i < c.selected_counts[row]; ++i) { const int index = c.selected_indices[row * c.max_selected + i]; + if (index < 0 || index > position || index >= c.cache_sequence_length) continue; const int64_t offset = ((b * c.kv_num_heads + kv) * c.cache_sequence_length + index) * c.head_size; double dot = 0.0; for (int d = 0; d < c.head_size; ++d) { - const float key_scale = int8_cache ? c.key_scale[kv * c.head_size + d] : 1.0f; - dot += c.query[out + d] * c.expected_present_key[offset + d] * key_scale; + const double key = int8_cache ? static_cast(c.expected_present_key[offset + d]) * + static_cast(c.key_scale[kv * c.head_size + d]) + : bf16(c.expected_present_key[offset + d]); + dot += bf16(c.query[out + d]) * key; } - const double weight = std::exp(dot * c.scale); + scores.push_back(dot * static_cast(c.scale)); + offsets.push_back(offset); + } + const double max_score = scores.empty() ? 0.0 : *std::max_element(scores.begin(), scores.end()); + double denominator = 0.0; + std::fill(numerator.begin(), numerator.end(), 0.0); + for (size_t i = 0; i < scores.size(); ++i) { + const double weight = std::exp(scores[i] - max_score); denominator += weight; for (int d = 0; d < c.head_size; ++d) { - const float value_scale = int8_cache ? c.value_scale[kv * c.head_size + d] : 1.0f; - numerator[d] += static_cast(weight * c.expected_present_value[offset + d] * value_scale); + const double value = int8_cache ? static_cast(c.expected_present_value[offsets[i] + d]) * + static_cast(c.value_scale[kv * c.head_size + d]) + : bf16(c.expected_present_value[offsets[i] + d]); + numerator[d] += weight * value; } } for (int d = 0; d < c.head_size; ++d) { @@ -626,6 +661,48 @@ DynamicSparseAttentionCase MakeSelectedOnlyPrefillCase(bool int8_cache) { } } } +} + +DynamicSparseAttentionCase MakeSelectedOnlyPrefillCase(bool int8_cache, int num_heads = 20, int selected_count = 509, + bool invalid_entries = false, int sequence_length = 9) { + auto c = int8_cache ? MakeGroupedCacheCase(sequence_length, false, false, 2, num_heads, 2, true, selected_count) + : MakeGroupedCacheCase(sequence_length, false, true, 2, num_heads, 2, false, selected_count); + c.head_sink.clear(); + c.smooth_softmax = 0; + if (int8_cache) { + // Distinct rows make wrong QK scores and online-softmax weights observable in the output; the larger scale + // spreads the softmax weights enough for score errors to show. The value scales shrink by 2^-4 (exact) so that + // the BF16 rounding of the dequantized operands in the BF16-dequantizing kernel, about half of the output + // tolerance at this magnitude, keeps a margin. + c.scale = 0.5f; + for (float& value_scale : c.value_scale) value_scale *= 0.0625f; + for (float& value : c.value) value *= 0.0625f; + std::mt19937 gen(17); + for (int b = 0; b < c.batch_size; ++b) { + const int64_t past_length = c.seqlens_k[b] + 1 - c.sequence_length; + for (int kv = 0; kv < c.kv_num_heads; ++kv) { + for (int pos = 0; pos < past_length; ++pos) { + for (int d = 0; d < c.head_size; ++d) { + const int64_t offset = ((b * c.kv_num_heads + kv) * c.cache_sequence_length + pos) * c.head_size + d; + c.past_key[offset] = c.expected_present_key[offset] = static_cast(int(gen() % 256) - 128); + c.past_value[offset] = c.expected_present_value[offset] = static_cast(int(gen() % 256) - 128); + } + } + } + } + } + if (invalid_entries) { + for (int row = 0; row < c.batch_size * c.sequence_length; ++row) { + const int64_t b = row / c.sequence_length; + const int64_t position = c.seqlens_k[b] + 1 - c.sequence_length + row % c.sequence_length; + const int32_t invalid_values[] = {-1, static_cast(position) + 1, + static_cast(c.cache_sequence_length) + 5, -7, 1 << 30}; + for (int i = 3; i < c.selected_counts[row]; i += 5) { + c.selected_indices[row * c.max_selected + i] = invalid_values[(i / 5) % 5]; + } + } + } + ComputeSelectedOnlyPrefillExpectedOutput(c, int8_cache); return c; } @@ -641,6 +718,718 @@ TEST(DynamicSparseAttentionTest, Int8CachePrefillBFloat16_CUDA) { } } +// The tensor-core INT8 prefill kernel (ORT_DSA_PREFILL_INT8_TC=1) and the BF16-dequantizing one (=0) over query-head +// group sizes 1, 6, 12 and 16 (padded to the 16-row MMA tile), candidate counts around the 64-candidate tile and +// invalid selected indices. +TEST(DynamicSparseAttentionTest, Int8CachePrefillTensorCoreKernel_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + for (int num_heads : {2, 12, 24, 32}) { + for (int selected_count : {1, 64, 65, 509}) { + for (bool invalid_entries : {false, true}) { + SCOPED_TRACE("num_heads=" + std::to_string(num_heads) + ", selected_count=" + std::to_string(selected_count) + + ", invalid_entries=" + std::to_string(invalid_entries)); + const auto c = MakeSelectedOnlyPrefillCase(true, num_heads, selected_count, invalid_entries); + for (const char* int8_tc_env : {"1", "0"}) { + SCOPED_TRACE(std::string("ORT_DSA_PREFILL_INT8_TC=") + int8_tc_env); + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{"ORT_DSA_PREFILL_KERNEL", "1"}, {"ORT_DSA_PREFILL_INT8_TC", int8_tc_env}}}; + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider()); + } + if (::testing::Test::HasFailure()) return; + } + } + } +} + +std::vector> PrefillKernelConfigs(const char* kernel, const char* int8_tc) { + return {{kOrtSessionOptionsCudaDsaPrefillKernel, kernel}, {kOrtSessionOptionsCudaDsaPrefillInt8Tc, int8_tc}}; +} + +std::vector> DecodeKernelConfigs(const char* int8_tc, + const char* tiles_per_split = "0") { + return {{kOrtSessionOptionsCudaDsaDecodeInt8Tc, int8_tc}, + {kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, tiles_per_split}}; +} + +// A CTA of the tensor-core INT8 kernel whose Q * key_scale does not fit fp16 evaluates its scores in fp32. The query +// below cancels exactly: selected index 0 scores 131072 * 1 + 65536 * -2 = 0 like index 1 (all-zero key), so the output +// is the mean of the two values. Clamping Q * key_scale to the fp16 range would score index 0 at -65504 instead. +TEST(DynamicSparseAttentionTest, Int8PrefillCancellingLargeQuery_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + auto c = MakeSelectedOnlyPrefillCase(true, 2, 2); + std::fill(c.query.begin(), c.query.end(), 0.0f); + std::fill(c.key_scale.begin(), c.key_scale.end(), 1.0f); + std::fill(c.value_scale.begin(), c.value_scale.end(), 1.0f); + std::fill(c.key.begin(), c.key.end(), 0.0f); + std::fill(c.value.begin(), c.value.end(), 0.0f); + std::fill(c.past_key.begin(), c.past_key.end(), 0.0f); + std::fill(c.past_value.begin(), c.past_value.end(), 0.0f); + for (int b = 0; b < c.batch_size; ++b) { + for (int kv = 0; kv < c.kv_num_heads; ++kv) { + const int64_t offset = (b * c.kv_num_heads + kv) * c.cache_sequence_length * c.head_size; + c.past_key[offset] = 1.0f; + c.past_key[offset + 1] = -2.0f; + std::fill_n(c.past_value.begin() + offset, c.head_size, 2.0f); + std::fill_n(c.past_value.begin() + offset + c.head_size, c.head_size, 6.0f); + } + } + c.expected_present_key = c.past_key; + c.expected_present_value = c.past_value; + for (size_t i = 0; i < c.query.size(); i += c.head_size) { + c.query[i] = 131072.0f; + c.query[i + 1] = 65536.0f; + } + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + for (size_t i = 0; i < c.expected_output.size(); i += c.head_size) { + const size_t row = i / (c.num_heads * c.head_size); + ASSERT_EQ(c.expected_output[i], row == 0 ? 0.0f : 4.0f); + } + + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", + PrefillKernelConfigs("1", "1")); +} + +// Multiplying every query by 2^28 and dividing the scale by 2^28 keeps the scores while pushing Q * key_scale of every +// CTA past fp16, so the fp32 path serves all query-head group sizes, candidate counts and invalid indices. A second +// case puts a huge query on a channel whose keys are all zero for a single head, so only some CTAs take the fp32 path. +TEST(DynamicSparseAttentionTest, Int8PrefillQueryBeyondHalfRange_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + const auto configs = PrefillKernelConfigs("1", "1"); + + for (int num_heads : {2, 12, 24, 32}) { + for (int selected_count : {1, 65, 509}) { + SCOPED_TRACE("num_heads=" + std::to_string(num_heads) + ", selected_count=" + std::to_string(selected_count)); + auto c = MakeSelectedOnlyPrefillCase(true, num_heads, selected_count, true); + for (auto& q : c.query) q *= 0x1p28f; + c.scale *= 0x1p-28f; + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", configs); + if (::testing::Test::HasFailure()) return; + } + } + + auto c = MakeSelectedOnlyPrefillCase(true, 24, 65, true); + const int64_t channels = c.kv_num_heads * c.head_size; + constexpr int64_t kZeroChannel = 7; + // Every key of KV head 0 is zero on kZeroChannel, so a huge query there contributes exactly nothing to the scores. + for (int b = 0; b < c.batch_size; ++b) { + for (int64_t pos = 0; pos < c.cache_sequence_length; ++pos) { + const int64_t offset = ((b * c.kv_num_heads) * c.cache_sequence_length + pos) * c.head_size + kZeroChannel; + c.past_key[offset] = c.expected_present_key[offset] = 0.0f; + } + } + for (int64_t row = 0; row < c.batch_size * c.sequence_length; ++row) { + c.key[row * channels + kZeroChannel] = 0.0f; + const int64_t b = row / c.sequence_length; + const int64_t pos = c.seqlens_k[b] + 1 - c.sequence_length + row % c.sequence_length; + c.expected_present_key[((b * c.kv_num_heads) * c.cache_sequence_length + pos) * c.head_size + kZeroChannel] = 0.0f; + if (row % 2 == 1) c.query[(row * c.num_heads + 5) * c.head_size + kZeroChannel] = 1.0e7f; + } + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", configs); +} + +// Two selected candidates of an otherwise empty INT8 cache, so that the scores and the output follow from a few numbers: +// every query head holds `query` on channel 0, the keys hold key_codes on channel 0 and the values hold value_codes on +// every channel. Candidate 0 is at cache position 0 and candidate 1 at second_position (negative: not selected). The +// first row of the case selects nothing and outputs zeros. +struct TwoCandidateInt8Spec { + float query; + float key_scale; + float value_scale; + float scale; + float key_codes[2]; + float value_codes[2]; + int second_position = 1; + // Positive: every selecting row lists this many entries, candidate 0 first, candidate 1 last and -1 in between, which + // spreads the candidates over 64-candidate tiles (and over decode splits). Otherwise the row lists the candidates + // ahead of the -1 padding of the default 512 entries. + int list_length = 0; + // Checks the output within one bfloat16 ulp rather than the default 0.005 + 1% tolerance, for outputs that are tiny + // or whose wrong result differs from the right one by less than the default tolerance. + bool exact = false; + // Replaces the single-channel pattern (query, key codes and scales) before the reference is computed; receives + // second_position. + std::function customize; +}; + +// Sets one INT8 key code of every batch item and KV head. +void SetInt8KeyCode(DynamicSparseAttentionCase& c, int position, int channel, float code) { + for (int b = 0; b < c.batch_size; ++b) { + for (int kv = 0; kv < c.kv_num_heads; ++kv) { + c.past_key[((b * c.kv_num_heads + kv) * c.cache_sequence_length + position) * c.head_size + channel] = code; + } + } +} + +DynamicSparseAttentionCase MakeTwoCandidateInt8Case(const TwoCandidateInt8Spec& spec, int num_heads, + int sequence_length) { + auto c = MakeSelectedOnlyPrefillCase(true, num_heads, 2, false, sequence_length); + const int rows = static_cast(c.batch_size * c.sequence_length); + const bool has_second = spec.second_position >= 0; + if (spec.list_length > 0) c.max_selected = spec.list_length; + c.selected_indices.assign(static_cast(rows * c.max_selected), -1); + for (int row = 0; row < rows; ++row) { + c.selected_counts[row] = row == 0 ? 0 : spec.list_length > 0 ? spec.list_length + : (has_second ? 2 : 1); + if (row == 0) continue; + int32_t* entries = c.selected_indices.data() + static_cast(row * c.max_selected); + entries[0] = 0; + if (has_second) entries[spec.list_length > 0 ? spec.list_length - 1 : 1] = spec.second_position; + } + c.scale = spec.scale; + std::fill(c.query.begin(), c.query.end(), 0.0f); + std::fill(c.key_scale.begin(), c.key_scale.end(), spec.key_scale); + std::fill(c.value_scale.begin(), c.value_scale.end(), spec.value_scale); + std::fill(c.key.begin(), c.key.end(), 0.0f); + std::fill(c.value.begin(), c.value.end(), 0.0f); + std::fill(c.past_key.begin(), c.past_key.end(), 0.0f); + std::fill(c.past_value.begin(), c.past_value.end(), 0.0f); + const int positions[2] = {0, spec.second_position}; + for (int b = 0; b < c.batch_size; ++b) { + for (int kv = 0; kv < c.kv_num_heads; ++kv) { + const int64_t offset = (b * c.kv_num_heads + kv) * c.cache_sequence_length * c.head_size; + for (int i = 0; i < 2; ++i) { + if (positions[i] < 0) continue; + c.past_key[offset + positions[i] * c.head_size] = spec.key_codes[i]; + std::fill_n(c.past_value.begin() + offset + positions[i] * c.head_size, c.head_size, spec.value_codes[i]); + } + } + } + for (size_t i = 0; i < c.query.size(); i += c.head_size) c.query[i] = spec.query; + if (spec.customize) spec.customize(c, spec.second_position); + c.expected_present_key = c.past_key; + c.expected_present_value = c.past_value; + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + if (spec.exact) { + c.output_abs_tolerance = 0.0f; + c.output_rel_tolerance = 0x1p-7f; + } + return c; +} + +// What RunTwoCandidateInt8Case runs: prefill (9 rows per batch item) and decode with each of the sequence lengths. +// The tensor-core decode kernel also runs with each tiles-per-split setting (the other kernels ignore it). +struct TwoCandidateInt8Run { + bool prefill = true; + std::vector decode_sequence_lengths = {3}; + std::vector tiles_per_split = {"0"}; +}; + +// Every row but the first outputs `expected` on every channel. The tensor-core INT8 kernels run both disabled +// (reference fallback) and enabled. +void RunTwoCandidateInt8Case(const TwoCandidateInt8Spec& spec, float expected, const TwoCandidateInt8Run& run = {}) { + std::vector> shapes; // (decode, sequence length) + if (run.prefill) shapes.emplace_back(false, 9); + for (int sequence_length : run.decode_sequence_lengths) shapes.emplace_back(true, sequence_length); + for (const auto& [decode, sequence_length] : shapes) { + const auto c = MakeTwoCandidateInt8Case(spec, 12, sequence_length); + for (size_t i = 0; i < c.expected_output.size(); i += c.head_size) { + const size_t row = i / (c.num_heads * c.head_size); + ASSERT_NEAR(c.expected_output[i], row == 0 ? 0.0f : expected, 1e-6 * std::fabs(expected)); + } + for (const char* int8_tc : {"0", "1"}) { + const std::vector tiles_per_split_values = + decode && std::string(int8_tc) == "1" ? run.tiles_per_split : std::vector{"0"}; + for (const char* tiles_per_split : tiles_per_split_values) { + SCOPED_TRACE(std::string(decode ? "decode" : "prefill") + ", sequence_length=" + + std::to_string(sequence_length) + ", int8_tc=" + int8_tc + + ", tiles_per_split=" + tiles_per_split); + RunDynamicSparseAttentionCase( + c, DefaultCudaExecutionProvider(), OpTester::ExpectResult::kExpectSuccess, "", + decode ? DecodeKernelConfigs(int8_tc, tiles_per_split) : PrefillKernelConfigs("1", int8_tc)); + if (::testing::Test::HasFailure()) return; + } + } + } +} + +// Runs the case with its candidates in one 64-candidate tile and again with a list of 129 entries that puts the +// candidates in the first and the last of three tiles, which the tensor-core decode kernel spreads over up to three +// splits (two of them active) and the prefill kernel consumes in one CTA. +void RunTwoCandidateInt8CaseOverTiles(TwoCandidateInt8Spec spec, float expected, bool prefill = true) { + TwoCandidateInt8Run run; + run.prefill = prefill; + RunTwoCandidateInt8Case(spec, expected, run); + if (::testing::Test::HasFailure()) return; + spec.list_length = 129; + if (spec.second_position >= 0) spec.second_position = 128; + run.tiles_per_split = {"0", "1", "2", "3"}; + RunTwoCandidateInt8Case(spec, expected, run); +} + +// The tensor-core INT8 kernels round Q * key_scale to fp16, which flushes products below 2^-25 to zero; a softmax scale +// above 1 would amplify that error, so those calls use the BF16-dequantizing prefill and CUDA-core decode kernels. +// Every product here is 2^-30 (zero in fp16) and the scale of 2^30 turns it into scores of +-1, so the output is +// 1 / (1 + e^2) rather than the 0.5 an fp16 flush gives. +TEST(DynamicSparseAttentionTest, Int8TensorCoreLargeSoftmaxScale_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + TwoCandidateInt8Spec spec{}; + spec.query = 1.0f; + spec.key_scale = 0x1p-30f; + spec.value_scale = 1.0f; + spec.scale = 0x1p30f; + spec.key_codes[0] = 1.0f; + spec.key_codes[1] = -1.0f; + spec.value_codes[1] = 1.0f; + RunTwoCandidateInt8Case(spec, static_cast(1.0 / (1.0 + std::exp(2.0)))); +} + +// The tensor-core INT8 kernels pass the softmax probabilities to the PV mma as fp16 holding 2^15 * p (the factor leaves +// in the epilogue), so a candidate keeps its contribution down to p = e^-20 (kInt8TcTailDepth) where plain fp16 flushes +// p below 2^-25. A CTA with a value scale above 2^16 (kInt8TcMaxValueScale) and a valid candidate more than +// kInt8TcTailDepth nats below its running maximum evaluates the row/split in fp32 (AttendRowFp32) instead, so that +// probabilities that leave the fp16 range, amplified by an arbitrarily large finite value scale, are preserved as well: +// 1. query = 288 -> score -18, p = e^-18 ~= 2^-25.97 (flushes in unscaled fp16, kept by 2^15 * p on the TC path with +// value_scale = 2^16 and value code 100); +// 2. query = 320 -> score -20, p = e^-20 ~= 2^-28.85, value_scale = 2^28: the depth equals kInt8TcTailDepth, so the +// tensor-core path still serves it; +// 3. query = 480 -> score -30, p = e^-30 ~= 2^-43.28 (flushes even 2^15 * p in fp16), value_scale = 2^60, output +// 2^60 * e^-30 / (1 + e^-30) ~= 107886.05 via the AttendRowFp32 fallback. +TEST(DynamicSparseAttentionTest, Int8TensorCoreTinyProbabilityLargeValueScale_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + { + SCOPED_TRACE("tc_prob_scale_2p16"); + TwoCandidateInt8Spec spec{}; + spec.query = 288.0f; + spec.key_scale = 1.0f; + spec.value_scale = 0x1p16f; + spec.scale = 0.0625f; + spec.key_codes[1] = -1.0f; + spec.value_codes[1] = 100.0f; + const double tail = std::exp(-18.0); + RunTwoCandidateInt8Case(spec, static_cast(100.0 * 0x1p16 * tail / (1.0 + tail))); + if (::testing::Test::HasFailure()) return; + } + { + SCOPED_TRACE("fp32_fallback_2p28"); + TwoCandidateInt8Spec spec{}; + spec.query = 320.0f; + spec.key_scale = 1.0f; + spec.value_scale = 0x1p28f; + spec.scale = 0.0625f; + spec.key_codes[1] = -1.0f; + spec.value_codes[1] = 1.0f; + const double tail = std::exp(-20.0); + RunTwoCandidateInt8Case(spec, static_cast(0x1p28 * tail / (1.0 + tail))); + if (::testing::Test::HasFailure()) return; + } + { + SCOPED_TRACE("fp32_fallback_2p60"); + TwoCandidateInt8Spec spec{}; + spec.query = 480.0f; + spec.key_scale = 1.0f; + spec.value_scale = 0x1p60f; + spec.scale = 0.0625f; + spec.key_codes[1] = -1.0f; + spec.value_codes[1] = 1.0f; + const double tail = std::exp(-30.0); + RunTwoCandidateInt8Case(spec, static_cast(0x1p60 * tail / (1.0 + tail))); + } +} + +// The cases below fix the outputs of the tensor-core INT8 kernels on numerical corners with exact arithmetic: the output +// follows from a few numbers, is checked to one bfloat16 ulp, and is also produced with the candidates in the first and +// the last of three 64-candidate tiles, which the decode kernel spreads over up to three splits (two of them active). + +// Q * key_scale is exact in fp32 here (32768 and -32512 * 129 / 128 = -32766 for a channel pair), and the all-ones key +// of candidate 0 makes the 128 pairs add up to 256, a logit gap of 16 at scale 1/16 against the all-zero key of +// candidate 1. -32766 is not an fp16 number though: rounding it to fp16 gives -32768, cancels every pair and turns the +// gap into 0 (output 0.5 instead of 1 - 1.1e-7). +TEST(DynamicSparseAttentionTest, Int8TensorCoreQueryFoldingRoundingError_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + TwoCandidateInt8Spec spec{}; + spec.key_scale = 1.0f; + spec.value_scale = 1.0f; + spec.scale = 0.0625f; + spec.value_codes[0] = 1.0f; + spec.exact = true; + spec.customize = [](DynamicSparseAttentionCase& c, int) { + for (size_t i = 0; i < c.query.size(); ++i) c.query[i] = i % 2 == 0 ? 32768.0f : -32512.0f; + for (size_t i = 0; i < c.key_scale.size(); ++i) c.key_scale[i] = i % 2 == 0 ? 1.0f : 1.0078125f; + for (int d = 0; d < c.head_size; ++d) SetInt8KeyCode(c, 0, d, 1.0f); + }; + RunTwoCandidateInt8CaseOverTiles(spec, 1.0f); +} + +// Logits of +-2^-13 and values of +-4064 (value code +-127, value scale 32) give 4064 * tanh(2^-13) = 127 / 256. The two +// probabilities, 0.5 +- 2^-14, both round to 0.5 in one fp16 tile (the spacing is 2^-11 above 0.5 and 2^-12 below), so +// the difference of the two values would cancel to 0; the tensor-core decode kernel passes P as a high and a low fp16 +// part to keep it. The prefill kernels round each probability to a single 16-bit tile for the PV mma, as the kernels +// before them did, so only decode is covered. +TEST(DynamicSparseAttentionTest, Int8TensorCoreDecodeSignedValueCancellation_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + TwoCandidateInt8Spec spec{}; + spec.query = 0x1p-9f; + spec.key_scale = 1.0f; + spec.value_scale = 32.0f; + spec.scale = 0.0625f; + spec.key_codes[0] = 1.0f; + spec.key_codes[1] = -1.0f; + spec.value_codes[0] = 127.0f; + spec.value_codes[1] = -127.0f; + spec.exact = true; + RunTwoCandidateInt8CaseOverTiles(spec, 127.0f / 256.0f, /*prefill=*/false); +} + +// Both candidates score 8290560 * scale (query 65280, close to the fp16 maximum, times key code 127 on channel 0), and +// candidate 1 adds 0.5 * scale from a second channel. Softmax weights computed in the log2 domain lose that gap to the +// float spacing of the common logit (0.0625 near 7.5e5 in log2 units, against a gap of 0.045 at scale 1/16) and give +// 0.5 instead of 1 / (1 + exp(-0.5 * scale)); the natural-log domain keeps it, since every logit here is an exact +// float. +TEST(DynamicSparseAttentionTest, Int8TensorCoreLargeCommonLogit_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + for (float scale : {1.0f, 0.0625f}) { + SCOPED_TRACE("scale=" + std::to_string(scale)); + TwoCandidateInt8Spec spec{}; + spec.query = 65280.0f; + spec.key_scale = 1.0f; + spec.value_scale = 1.0f; + spec.scale = scale; + spec.key_codes[0] = 127.0f; + spec.key_codes[1] = 127.0f; + spec.value_codes[1] = 1.0f; + spec.exact = true; + spec.customize = [](DynamicSparseAttentionCase& c, int second_position) { + for (size_t i = 1; i < c.query.size(); i += c.head_size) c.query[i] = 0.5f; + SetInt8KeyCode(c, second_position, 1, 1.0f); + }; + RunTwoCandidateInt8CaseOverTiles(spec, static_cast(1.0 / (1.0 + std::exp(-0.5 * scale)))); + if (::testing::Test::HasFailure()) return; + } +} + +// A single candidate with value code 127 and value scale 2^-135 outputs 127 * 2^-135 = 2.9e-39, a subnormal in float +// that rounds to the bfloat16 subnormal 2^-128; a kernel that folds the value scale into an fp16 operand or multiplies +// it in after an fp16 conversion flushes it to 0. The second layout lists the candidate first among 129 entries, which +// leaves two splits empty with three tiles per row. +TEST(DynamicSparseAttentionTest, Int8TensorCoreTinyValueScale_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + TwoCandidateInt8Spec spec{}; + spec.key_scale = 1.0f; + spec.value_scale = 0x1p-135f; + spec.scale = 0.0625f; + spec.value_codes[0] = 127.0f; + spec.second_position = -1; + spec.exact = true; + RunTwoCandidateInt8CaseOverTiles(spec, 127.0f * 0x1p-135f); +} + +// The tensor-core decode kernel runs at most 32 splits per row (and no more than the workspace holds, 16 for the eight +// rows of a batch item here) however many 64-candidate tiles the list has, so a tiles-per-split setting is a target: +// 8192 entries are 128 tiles, and a setting of 1 requests 128 splits but gets 32 splits of four tiles. Whatever the +// split count, every candidate must reach the merge: the logits are all equal and the values are 2 and 6, so a split +// that is dropped or counted twice gives 2, 6 or something else instead of 4. The lists are 2048 entries (32 tiles, +// exactly the cap for 1 tile per split), 2112 (33 tiles, one over) and 8192, with the second candidate in the last tile. +TEST(DynamicSparseAttentionTest, Int8TensorCoreDecodeTilesPerSplitCaps_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + for (int list_length : {2048, 2112, 8192}) { + SCOPED_TRACE("list_length=" + std::to_string(list_length)); + TwoCandidateInt8Spec spec{}; + spec.key_scale = 1.0f; + spec.value_scale = 1.0f; + spec.scale = 0.0625f; + spec.value_codes[0] = 2.0f; + spec.value_codes[1] = 6.0f; + spec.second_position = 300; + spec.list_length = list_length; + TwoCandidateInt8Run run; + run.prefill = false; + run.decode_sequence_lengths = {1, 8}; + run.tiles_per_split = {"0", "1", "2", "3", "64"}; + RunTwoCandidateInt8Case(spec, 4.0f, run); + if (::testing::Test::HasFailure()) return; + } +} + +// Runs a case built for a prefill of 9 rows per batch item (decode == false) or a decode of 3 with the tensor-core +// kernels; the decode kernel gets the automatic, 1 and 3 tiles per split. The BF16-dequantizing prefill kernel is left +// out: it rounds the dequantized keys and values to bfloat16, which the reference of these cases does not allow for. +void RunInt8TensorCoreCase(const DynamicSparseAttentionCase& c, bool decode) { + const std::vector tiles_per_split_values = + decode ? std::vector{"0", "1", "3"} : std::vector{"0"}; + for (const char* tiles_per_split : tiles_per_split_values) { + SCOPED_TRACE(std::string(decode ? "decode" : "prefill") + ", tiles_per_split=" + tiles_per_split); + RunDynamicSparseAttentionCase( + c, DefaultCudaExecutionProvider(), OpTester::ExpectResult::kExpectSuccess, "", + decode ? DecodeKernelConfigs("1", tiles_per_split) : PrefillKernelConfigs("1", "1")); + if (::testing::Test::HasFailure()) return; + } +} + +// Random rows with value scales between 2^14 and 2^17 (outputs up to 10^7), so that the tail-depth guard of the +// tensor-core kernels sees value scales above 2^16 on random data. Queries scaled by 1, 16 and 32 spread the scores of a +// head over about 0.7, 11 and 22 nats: the first two stay on the tensor-core path, and some CTAs of the third find a +// candidate more than 20 nats below the maximum and take the fp32 fallback. The tolerance is the 0.005 of the unscaled +// fixture scaled with the values. +TEST(DynamicSparseAttentionTest, Int8TensorCoreLargeValueScaleRandom_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + for (float query_multiplier : {1.0f, 16.0f, 32.0f}) { + for (int sequence_length : {9, 3}) { + SCOPED_TRACE("query_multiplier=" + std::to_string(query_multiplier) + + ", sequence_length=" + std::to_string(sequence_length)); + auto c = MakeSelectedOnlyPrefillCase(true, 24, 509, true, sequence_length); + for (float& value_scale : c.value_scale) value_scale *= 0x1p22f; + for (float& value : c.value) value *= 0x1p22f; + for (float& query : c.query) query *= query_multiplier; + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + c.output_abs_tolerance = 0.005f * 0x1p22f; + RunInt8TensorCoreCase(c, sequence_length == 3); + if (::testing::Test::HasFailure()) return; + } + } +} + +// Random INT8 rows whose folded query products Q * key_scale are not fp16 numbers: key scales are random in +// [0.01, 0.03] and queries are drawn from N(0, 0.25^2) (rounded to bfloat16 as the kernels receive them). The other random +// fixtures fold to exact fp16 products, so they never exercise the rounding of the folded query. A query multiplier +// of 1 keeps the accumulated rounding error of a head far below the error budget of the fp16 scores, and 16 pushes it over, +// into the fp32 scores. +TEST(DynamicSparseAttentionTest, Int8TensorCoreRandomQueryFolding_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + std::mt19937 gen(29); + std::uniform_real_distribution key_scale_distribution(0.01f, 0.03f); + std::normal_distribution query_distribution(0.0f, 0.25f); + for (float query_multiplier : {1.0f, 16.0f}) { + for (int num_heads : {12, 24}) { + for (int selected_count : {65, 509}) { + for (int sequence_length : {9, 3}) { + SCOPED_TRACE("query_multiplier=" + std::to_string(query_multiplier) + + ", num_heads=" + std::to_string(num_heads) + + ", selected_count=" + std::to_string(selected_count) + + ", sequence_length=" + std::to_string(sequence_length)); + auto c = MakeSelectedOnlyPrefillCase(true, num_heads, selected_count, true, sequence_length); + const size_t channels = static_cast(c.kv_num_heads * c.head_size); + for (float& key_scale : c.key_scale) key_scale = key_scale_distribution(gen); + // The new tokens are the codes of the fixture times the new scales; expected_present_key keeps the codes. + for (size_t i = 0; i < c.key.size(); ++i) { + c.key[i] = static_cast(static_cast(i % 11) - 5) * c.key_scale[i % channels]; + } + for (float& query : c.query) query = query_multiplier * query_distribution(gen); + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + RunInt8TensorCoreCase(c, sequence_length == 3); + if (::testing::Test::HasFailure()) return; + } + } + } + } +} + +// Explicit ep.cuda.* entries take precedence over the environment, which is consulted only when no entry is set; an +// invalid value from either source fails kernel construction (with an exception, so those cases are skipped in builds +// without exceptions). The switches accept exactly the strings "0" and "1". +TEST(DynamicSparseAttentionTest, PrefillSessionOptions_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + const auto c = MakeSelectedOnlyPrefillCase(true, 12, 65, true); + const auto expect_success = OpTester::ExpectResult::kExpectSuccess; +#ifndef ORT_NO_EXCEPTIONS + const auto expect_failure = OpTester::ExpectResult::kExpectFailure; +#endif + + { + ScopedEnvironmentVariables env{ + EnvVarMap{{"ORT_DSA_PREFILL_KERNEL", "invalid"}, {"ORT_DSA_PREFILL_INT8_TC", "invalid"}}}; + for (const char* kernel : {"0", "1"}) { + for (const char* int8_tc : {"0", "1"}) { + SCOPED_TRACE(std::string("kernel=") + kernel + ", int8_tc=" + int8_tc); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_success, "", + PrefillKernelConfigs(kernel, int8_tc)); + } + } +#ifndef ORT_NO_EXCEPTIONS + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + "ORT_DSA_PREFILL_INT8_TC", + {{kOrtSessionOptionsCudaDsaPrefillKernel, "1"}}); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + "ORT_DSA_PREFILL_KERNEL", + {{kOrtSessionOptionsCudaDsaPrefillInt8Tc, "1"}}); +#endif + } + for (const char* int8_tc : {"0", "1"}) { + ScopedEnvironmentVariables env{EnvVarMap{{"ORT_DSA_PREFILL_KERNEL", "1"}, {"ORT_DSA_PREFILL_INT8_TC", int8_tc}}}; + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_success); + } +#ifndef ORT_NO_EXCEPTIONS + // Only the strings "0" and "1" are switch values, not other spellings of the same numbers. + for (const char* key : {kOrtSessionOptionsCudaDsaPrefillKernel, kOrtSessionOptionsCudaDsaPrefillInt8Tc}) { + for (const char* invalid : {"2", "-1", "", "1x", "true", "00", "01", "-0", "000001", "+1", " 1", "1 "}) { + SCOPED_TRACE(std::string(key) + "='" + invalid + "'"); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, key, + {{key, invalid}}); + } + } + for (const char* invalid : {"01", "-0", "+1"}) { + SCOPED_TRACE(std::string("environment value='") + invalid + "'"); + ScopedEnvironmentVariables env{EnvVarMap{{"ORT_DSA_PREFILL_INT8_TC", invalid}}}; + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + "ORT_DSA_PREFILL_INT8_TC"); + } +#endif +} + +// The split-K tensor-core INT8 decode kernel (ep.cuda.dsa_decode_int8_tc=1) and the CUDA-core grouped decode kernel +// (=0) over 1-8 query rows per batch item (row 0 of the case has no selected candidate), query-head group sizes 4, 6, +// 12 and 16 (padded to the 16-row MMA tile), candidate counts around the 64-candidate tile, invalid selected indices +// and 1, 3 or the automatic number of 64-candidate tiles per split (a split without tiles publishes an empty partial). +TEST(DynamicSparseAttentionTest, Int8CacheDecodeTensorCoreKernel_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + struct Config { + const char* decode_tc; + const char* tiles_per_split; + }; + const Config configs[] = {{"1", "0"}, {"1", "1"}, {"1", "3"}, {"0", "0"}}; + for (int sequence_length : {1, 3, 8}) { + for (int num_heads : {8, 12, 24, 32}) { + for (int selected_count : {1, 65, 509}) { + for (bool invalid_entries : {false, true}) { + SCOPED_TRACE("sequence_length=" + std::to_string(sequence_length) + + ", num_heads=" + std::to_string(num_heads) + + ", selected_count=" + std::to_string(selected_count) + + ", invalid_entries=" + std::to_string(invalid_entries)); + const auto c = MakeSelectedOnlyPrefillCase(true, num_heads, selected_count, invalid_entries, sequence_length); + for (const auto& config : configs) { + SCOPED_TRACE(std::string("dsa_decode_int8_tc=") + config.decode_tc + + ", dsa_decode_int8_tc_tiles_per_split=" + config.tiles_per_split); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", + DecodeKernelConfigs(config.decode_tc, + config.tiles_per_split)); + } + if (::testing::Test::HasFailure()) return; + } + } + } + } +} + +// A decode CTA whose Q * key_scale does not fit fp16 evaluates its range of candidates in fp32 and publishes an fp32 +// partial for the merge kernel. Multiplying every query by 2^28 and dividing the scale by 2^28 keeps the scores while +// pushing every CTA past fp16, for any number of splits; a second case puts a huge query on a channel whose keys are +// all zero for a single head of every other row, so that only some CTAs take the fp32 path. +TEST(DynamicSparseAttentionTest, Int8DecodeQueryBeyondHalfRange_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + for (const char* tiles_per_split : {"0", "1", "3"}) { + const auto configs = DecodeKernelConfigs("1", tiles_per_split); + for (int sequence_length : {1, 3}) { + for (int num_heads : {8, 12, 32}) { + for (int selected_count : {1, 65, 509}) { + SCOPED_TRACE("tiles_per_split=" + std::string(tiles_per_split) + + ", sequence_length=" + std::to_string(sequence_length) + + ", num_heads=" + std::to_string(num_heads) + + ", selected_count=" + std::to_string(selected_count)); + auto c = MakeSelectedOnlyPrefillCase(true, num_heads, selected_count, true, sequence_length); + for (auto& q : c.query) q *= 0x1p28f; + c.scale *= 0x1p-28f; + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", configs); + if (::testing::Test::HasFailure()) return; + } + } + } + + auto c = MakeSelectedOnlyPrefillCase(true, 24, 509, true, 3); + const int64_t channels = c.kv_num_heads * c.head_size; + constexpr int64_t kZeroChannel = 7; + for (int b = 0; b < c.batch_size; ++b) { + for (int64_t pos = 0; pos < c.cache_sequence_length; ++pos) { + const int64_t offset = ((b * c.kv_num_heads) * c.cache_sequence_length + pos) * c.head_size + kZeroChannel; + c.past_key[offset] = c.expected_present_key[offset] = 0.0f; + } + } + for (int64_t row = 0; row < c.batch_size * c.sequence_length; ++row) { + c.key[row * channels + kZeroChannel] = 0.0f; + const int64_t b = row / c.sequence_length; + const int64_t pos = c.seqlens_k[b] + 1 - c.sequence_length + row % c.sequence_length; + c.expected_present_key[((b * c.kv_num_heads) * c.cache_sequence_length + pos) * c.head_size + kZeroChannel] = + 0.0f; + if (row % 2 == 1) c.query[(row * c.num_heads + 5) * c.head_size + kZeroChannel] = 1.0e7f; + } + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", configs); + if (::testing::Test::HasFailure()) return; + } +} + +// Explicit ep.cuda.* entries take precedence over the environment, which is consulted only when no entry is set; an +// invalid value from either source fails kernel construction (with an exception, so those cases are skipped in builds +// without exceptions). +TEST(DynamicSparseAttentionTest, DecodeSessionOptions_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + const auto c = MakeSelectedOnlyPrefillCase(true, 12, 65, true, 3); + const auto expect_success = OpTester::ExpectResult::kExpectSuccess; +#ifndef ORT_NO_EXCEPTIONS + const auto expect_failure = OpTester::ExpectResult::kExpectFailure; +#endif + + { + ScopedEnvironmentVariables env{ + EnvVarMap{{"ORT_DSA_DECODE_INT8_TC", "invalid"}, {"ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT", "invalid"}}}; + for (const char* int8_tc : {"0", "1"}) { + for (const char* tiles_per_split : {"0", "2"}) { + SCOPED_TRACE(std::string("int8_tc=") + int8_tc + ", tiles_per_split=" + tiles_per_split); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_success, "", + DecodeKernelConfigs(int8_tc, tiles_per_split)); + } + } +#ifndef ORT_NO_EXCEPTIONS + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + "(ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT)", + {{kOrtSessionOptionsCudaDsaDecodeInt8Tc, "1"}}); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + "(ORT_DSA_DECODE_INT8_TC)", + {{kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, "0"}}); +#endif + } + for (const char* int8_tc : {"0", "1"}) { + ScopedEnvironmentVariables env{ + EnvVarMap{{"ORT_DSA_DECODE_INT8_TC", int8_tc}, {"ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT", "2"}}}; + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_success); + } +#ifndef ORT_NO_EXCEPTIONS + // Only the strings "0" and "1" are switch values, not other spellings of the same numbers. + for (const char* invalid : {"2", "-1", "", "1x", "true", "00", "01", "-0", "000001", "+1", " 1", "1 "}) { + SCOPED_TRACE(std::string("value='") + invalid + "'"); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + kOrtSessionOptionsCudaDsaDecodeInt8Tc, + {{kOrtSessionOptionsCudaDsaDecodeInt8Tc, invalid}}); + } + for (const char* invalid : {"01", "-0", "+1"}) { + SCOPED_TRACE(std::string("environment value='") + invalid + "'"); + ScopedEnvironmentVariables env{EnvVarMap{{"ORT_DSA_DECODE_INT8_TC", invalid}}}; + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + "ORT_DSA_DECODE_INT8_TC"); + } + // The tiles-per-split threshold is an integer, not a switch. + for (const char* invalid : {"-1", "", "1x", "many"}) { + SCOPED_TRACE(std::string("value='") + invalid + "'"); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, + {{kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, invalid}}); + } +#endif +} + TEST(DynamicSparseAttentionTest, FloatingCacheGroupedSingleSplitWorkspace_CUDA) { for (bool bf16 : {false, true}) { auto cuda_ep = DefaultCudaExecutionProvider(); @@ -2242,6 +3031,11 @@ std::vector MakeOffsetBufferCases(bool int8_cache, bool auxili cases.push_back(MisalignedBuffer("key_scale", &BufferOffsets::key_scale, scale_bytes)); cases.push_back(MisalignedBuffer("value_scale", &BufferOffsets::value_scale, scale_bytes)); all.key_scale = all.value_scale = scale_bytes; + // 8-byte cache and 4-byte output alignment still serve the BF16-dequantizing prefill kernel, but not the 16-byte + // cache copies and 8-byte output stores of the tensor-core INT8 prefill kernel. + cases.push_back(MisalignedBuffer("key_cache_8_byte", &BufferOffsets::key_cache, 8)); + cases.push_back(MisalignedBuffer("value_cache_8_byte", &BufferOffsets::value_cache, 8)); + cases.push_back(MisalignedBuffer("output_4_byte", &BufferOffsets::output, 4)); } if (auxiliary) { cases.push_back(MisalignedBuffer("auxiliary_key", &BufferOffsets::auxiliary_key, half_bytes)); @@ -2493,10 +3287,19 @@ TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeMisalignedBuffers_CUDA) { TEST(DynamicSparseAttentionTest, Int8CachePrefillMisalignedBuffers_CUDA) { if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; - // Aligned buffers use the tensor-core prefill kernel; misaligned ones must fall back to the generic kernels. + // 16-byte cache / 8-byte output aligned buffers use the tensor-core INT8 prefill kernel, 8-byte / 4-byte aligned ones + // the BF16-dequantizing kernel and misaligned ones must fall back to the generic kernels. ExpectMisalignedBuffersMatch(MakeSelectedOnlyPrefillCase(true), true); } +TEST(DynamicSparseAttentionTest, Int8CacheDecodeMisalignedBuffers_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + // 16-byte aligned buffers use the split-K tensor-core INT8 decode kernel; misaligned ones must fall back to the + // CUDA-core grouped decode kernel. + ExpectMisalignedBuffersMatch(MakeSelectedOnlyPrefillCase(true, 24, 509, false, 3), true); +} + TEST(DynamicSparseAttentionTest, FloatingCacheGroupedDecodeMisalignedBuffers_CUDA) { if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; for (bool bf16 : {true, false}) { diff --git a/onnxruntime/test/python/transformers/benchmark_dynamic_sparse_attention.py b/onnxruntime/test/python/transformers/benchmark_dynamic_sparse_attention.py new file mode 100644 index 0000000000000..28942551e5be9 --- /dev/null +++ b/onnxruntime/test/python/transformers/benchmark_dynamic_sparse_attention.py @@ -0,0 +1,422 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +""" +Op-level benchmark of the INT8 KV-cache DynamicSparseAttention CUDA kernels. + +The shapes are those of the 12 full-attention layers of Qwen3.8-Flash-Next: 24 query heads, 2 KV heads, head size 256, +INT8 KV cache with per-channel scales, softmax scale 1/16 and 2048 selected candidates per query row. Each +(rows, context length) case becomes a single-node model, and each kernel configuration becomes its own session, +because the ep.cuda.dsa_* options are read when the kernel is created: + + rows >= 9 (prefill) legacy: ep.cuda.dsa_prefill_int8_tc=0, the BF16-dequantizing kernel + tc: ep.cuda.dsa_prefill_int8_tc=1, the tensor-core kernel + rows <= 8 (decode) legacy: ep.cuda.dsa_decode_int8_tc=0, the CUDA-core grouped kernel + tc: ep.cuda.dsa_decode_int8_tc=1, the split-K tensor-core kernel, once for each + --tiles-per-split value (0 = automatic) + +The time of a configuration is the GPU time of the CUDA kernels one Run launches, measured with the CUPTI-based torch +profiler, so host launch overhead does not enter. "attention" sums the DynamicSparseAttention kernels (including the +decode merge kernel), "total" all GPU activity of the node (cache append, query preparation, ...). A round measures +every configuration once, back to back, and the order alternates between rounds. The speedup of a configuration is the +median over rounds of the paired ratio to the first (legacy) configuration, which cancels clock and temperature drift. +The "rms err" column is the RMS error of the BF16 output relative to the RMS of an FP64 reference, computed over +--reference-rows evenly spaced query rows (0 disables it). It is the error against the exact result, not the +difference to the legacy kernel, which has errors of its own. + +Example (a CUDA plugin EP build needs ORT_TEST_CUDA_PLUGIN_EP=1 and, unless the wheel carries the freshly built +library, ORT_CUDA_PLUGIN_PATH): + + python benchmark_dynamic_sparse_attention.py --rows 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 --json prefill.json + +The key and value scales are synthetic (uniform in the ranges of the Qwen3.8-Flash-Next layers) unless --rc-model names +a model whose DynamicSparseAttention node supplies them. The multipliers below reproduce the slow FP32 fallback of the +tensor-core kernels (queries or value scales far outside the usual range). +""" + +import argparse +import json +import platform +import statistics +import sys + +import numpy as np +import onnx +import torch +from cuda_plugin_ep_helper import CUDA_PLUGIN_EP_NAME, ensure_cuda_plugin_ep_registered +from onnx import TensorProto, helper, numpy_helper + +import onnxruntime as ort + +NUM_HEADS = 24 +KV_HEADS = 2 +HEAD_SIZE = 256 +SOFTMAX_SCALE = 0.0625 +MAX_DECODE_ROWS = 8 +TORCH_TO_ONNX = { + torch.bfloat16: TensorProto.BFLOAT16, + torch.int8: TensorProto.INT8, + torch.int32: TensorProto.INT32, +} + + +def load_rc_scales(model_path, layer): + """Returns the INT8 key and value scales of the layer-th DynamicSparseAttention node of a model.""" + model = onnx.load(model_path, load_external_data=False) + initializers = {initializer.name: initializer for initializer in model.graph.initializer} + nodes = [node for node in model.graph.node if node.op_type == "DynamicSparseAttention"] + if layer >= len(nodes): + raise ValueError(f"{model_path} has {len(nodes)} DynamicSparseAttention nodes, --layer {layer} is out of range") + return tuple( + numpy_helper.to_array(initializers[nodes[layer].input[index]]).astype(np.float32) for index in (18, 19) + ) + + +def make_scales(args): + if args.rc_model: + key_scale, value_scale = load_rc_scales(args.rc_model, args.layer) + else: + rng = np.random.default_rng(args.seed) + key_scale = rng.uniform(0.002, 0.07, KV_HEADS * HEAD_SIZE).astype(np.float32) + value_scale = rng.uniform(0.002, 0.03, KV_HEADS * HEAD_SIZE).astype(np.float32) + return key_scale * np.float32(args.key_scale_multiplier), value_scale * np.float32(args.value_scale_multiplier) + + +def make_model(rows, context, max_selected, key_scale, value_scale): + def value_info(name, element_type, shape): + return helper.make_tensor_value_info(name, element_type, shape) + + cache_shape = [1, KV_HEADS, context, HEAD_SIZE] + graph = helper.make_graph( + [ + helper.make_node( + "DynamicSparseAttention", + [ + "query", + "key", + "value", + "past_key", + "past_value", + "", + "", + "selected_indices", + "selected_counts", + "seqlens_k", + "total_sequence_length", + *[""] * 7, + "key_scale", + "value_scale", + ], + ["output", "present_key", "present_value"], + domain="com.microsoft", + num_heads=NUM_HEADS, + kv_num_heads=KV_HEADS, + scale=SOFTMAX_SCALE, + ) + ], + "dynamic_sparse_attention", + [ + value_info("query", TensorProto.BFLOAT16, [1, rows, NUM_HEADS * HEAD_SIZE]), + value_info("key", TensorProto.BFLOAT16, [1, rows, KV_HEADS * HEAD_SIZE]), + value_info("value", TensorProto.BFLOAT16, [1, rows, KV_HEADS * HEAD_SIZE]), + value_info("past_key", TensorProto.INT8, cache_shape), + value_info("past_value", TensorProto.INT8, cache_shape), + value_info("selected_indices", TensorProto.INT32, [rows, max_selected]), + value_info("selected_counts", TensorProto.INT32, [rows]), + value_info("seqlens_k", TensorProto.INT32, [1]), + ], + [ + value_info("output", TensorProto.BFLOAT16, [1, rows, NUM_HEADS * HEAD_SIZE]), + value_info("present_key", TensorProto.INT8, cache_shape), + value_info("present_value", TensorProto.INT8, cache_shape), + ], + [ + numpy_helper.from_array(np.array(context, dtype=np.int32), "total_sequence_length"), + numpy_helper.from_array(key_scale, "key_scale"), + numpy_helper.from_array(value_scale, "value_scale"), + ], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)] + ) + model.ir_version = 10 + return model.SerializeToString() + + +def make_tensors(rows, context, max_selected, args): + """The query rows are the last `rows` positions of a context of `context` tokens; each selects random positions.""" + generator = torch.Generator(device="cuda") + generator.manual_seed(args.seed) + + def normal(*shape, std): + return torch.randn(*shape, device="cuda", generator=generator) * std + + def cache_codes(): + return normal(1, KV_HEADS, context, HEAD_SIZE, std=args.cache_code_std).clamp(-127, 127).round().to(torch.int8) + + positions = torch.arange(context - rows, context, device="cuda") + counts = (positions + 1).clamp(max=max_selected).to(torch.int32) + selected = torch.empty(rows, max_selected, dtype=torch.int32, device="cuda") + columns = torch.arange(context, device="cuda") + for start in range(0, rows, 1024): + stop = min(start + 1024, rows) + noise = torch.rand(stop - start, context, device="cuda", generator=generator) + noise[columns[None, :] > positions[start:stop, None]] = -1.0 + selected[start:stop] = noise.topk(max_selected, dim=1).indices.to(torch.int32) + selected[torch.arange(max_selected, device="cuda")[None, :] >= counts[:, None]] = -1 + return { + "query": normal(1, rows, NUM_HEADS * HEAD_SIZE, std=args.query_std).to(torch.bfloat16), + "key": normal(1, rows, KV_HEADS * HEAD_SIZE, std=1.0).to(torch.bfloat16), + "value": normal(1, rows, KV_HEADS * HEAD_SIZE, std=1.0).to(torch.bfloat16), + "past_key": cache_codes(), + "past_value": cache_codes(), + "selected_indices": selected, + "selected_counts": counts, + "seqlens_k": torch.tensor([context - 1], dtype=torch.int32, device="cuda"), + } + + +def select_plugin_device(device_id): + """The device of the CUDA plugin EP when ORT_TEST_CUDA_PLUGIN_EP=1 selects it, otherwise None.""" + if not ensure_cuda_plugin_ep_registered(): + return None + for device in ort.get_ep_devices(): + if device.ep_name == CUDA_PLUGIN_EP_NAME: + found = device.ep_options.get("device_id", device.ep_metadata.get("cuda_device_id")) + if found is None or int(found) == device_id: + return device + return None + + +def create_session(model_bytes, options, args): + session_options = ort.SessionOptions() + for key, value in options.items(): + session_options.add_session_config_entry(key, value) + plugin_device = select_plugin_device(args.device) + if plugin_device is not None: + session_options.add_provider_for_devices([plugin_device], {}) + return ort.InferenceSession(model_bytes, session_options) + providers = [("CUDAExecutionProvider", {"device_id": args.device})] + return ort.InferenceSession(model_bytes, session_options, providers=providers) + + +def create_binding(session, tensors, output, args): + binding = session.io_binding() + for name, tensor in tensors.items(): + binding.bind_input( + name, "cuda", args.device, TORCH_TO_ONNX[tensor.dtype], list(tensor.shape), tensor.data_ptr() + ) + # The cache is updated in place, like the model does with its past_key_values buffers. + binding.bind_output("output", "cuda", args.device, TensorProto.BFLOAT16, list(output.shape), output.data_ptr()) + for name, cache in (("present_key", tensors["past_key"]), ("present_value", tensors["past_value"])): + binding.bind_output(name, "cuda", args.device, TensorProto.INT8, list(cache.shape), cache.data_ptr()) + return binding + + +def reference_attention(tensors, key_scale, value_scale, rows_to_check): + """FP64 attention of the given query rows over the cache the operator updated in place (new rows appended).""" + scales = [ + torch.tensor(scale, device="cuda", dtype=torch.float64).view(KV_HEADS, 1, HEAD_SIZE) + for scale in (key_scale, value_scale) + ] + keys = tensors["past_key"][0].double() * scales[0] + values = tensors["past_value"][0].double() * scales[1] + query = tensors["query"][0].double().view(-1, NUM_HEADS, HEAD_SIZE) + group = NUM_HEADS // KV_HEADS + reference = torch.empty(len(rows_to_check), NUM_HEADS, HEAD_SIZE, dtype=torch.float64, device="cuda") + for position, row in enumerate(rows_to_check): + selected = tensors["selected_indices"][row, : int(tensors["selected_counts"][row])].long() + for kv_head in range(KV_HEADS): + heads = slice(kv_head * group, (kv_head + 1) * group) + scores = query[row, heads] @ keys[kv_head, selected].T * SOFTMAX_SCALE + reference[position, heads] = torch.softmax(scores, dim=-1) @ values[kv_head, selected] + return reference + + +def relative_rms_error(result, reference): + return float((result - reference).pow(2).mean().sqrt() / reference.pow(2).mean().sqrt()) + + +def short_kernel_name(name): + head = name.replace("(anonymous namespace)::", "").split("(")[0].removeprefix("void ").strip() + base, separator, template_arguments = head.partition("<") + return base.split("::")[-1] + separator + template_arguments + + +def run_and_profile(session, binding, iterations): + """Returns the GPU time per Run in microseconds: {kernel name: total}, summed over the profiled runs.""" + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as profile: + for _ in range(iterations): + session.run_with_iobinding(binding) + binding.synchronize_outputs() + torch.cuda.synchronize() + times = {} + for event in profile.events(): + if event.device_type == torch.autograd.DeviceType.CUDA: + name = short_kernel_name(event.name) + times[name] = times.get(name, 0.0) + event.device_time / iterations + return times + + +def configurations(rows, args): + if rows <= MAX_DECODE_ROWS: + configs = [("legacy (CUDA-core)", {"ep.cuda.dsa_decode_int8_tc": "0"})] + for tiles in args.tiles_per_split: + options = { + "ep.cuda.dsa_decode_int8_tc": "1", + "ep.cuda.dsa_decode_int8_tc_tiles_per_split": str(tiles), + } + configs.append((f"tc tiles_per_split={'auto' if tiles == 0 else tiles}", options)) + return configs + return [ + ("legacy (BF16 dequantizing)", {"ep.cuda.dsa_prefill_int8_tc": "0"}), + ("tc", {"ep.cuda.dsa_prefill_int8_tc": "1"}), + ] + + +def format_range(values, digits): + return f"{statistics.median(values):.{digits}f} [{min(values):.{digits}f}, {max(values):.{digits}f}]" + + +def run_case(rows, context, args, key_scale, value_scale): + max_selected = min(args.selected, context) + model_bytes = make_model(rows, context, max_selected, key_scale, value_scale) + tensors = make_tensors(rows, context, max_selected, args) + output = torch.empty(1, rows, NUM_HEADS * HEAD_SIZE, dtype=torch.bfloat16, device="cuda") + configs = configurations(rows, args) + iterations = args.iterations or (args.decode_iterations if rows <= MAX_DECODE_ROWS else args.prefill_iterations) + + rows_to_check = sorted({int(row) for row in np.linspace(0, rows - 1, min(rows, args.reference_rows)).round()}) + sessions, bindings, reference, rms_error = [], [], None, [] + for _, options in configs: + session = create_session(model_bytes, options, args) + binding = create_binding(session, tensors, output, args) + for _ in range(args.warmup): + session.run_with_iobinding(binding) + binding.synchronize_outputs() + torch.cuda.synchronize() + if args.reference_rows > 0: + if reference is None: + # The first run has appended the new key/value rows to the cache; later runs rewrite the same values. + reference = reference_attention(tensors, key_scale, value_scale, rows_to_check).view( + len(rows_to_check), -1 + ) + rms_error.append(relative_rms_error(output[0, rows_to_check].float(), reference.float())) + else: + rms_error.append(float("nan")) + sessions.append(session) + bindings.append(binding) + + attention = [[] for _ in configs] + total = [[] for _ in configs] + kernels = [{} for _ in configs] + for round_index in range(args.rounds): + order = range(len(configs)) if round_index % 2 == 0 else reversed(range(len(configs))) + for index in order: + for _ in range(args.warmup): + sessions[index].run_with_iobinding(bindings[index]) + times = run_and_profile(sessions[index], bindings[index], iterations) + attention[index].append(sum(t for name, t in times.items() if "DynamicSparseAttention" in name)) + total[index].append(sum(times.values())) + kernels[index] = times + + kind = "decode" if rows <= MAX_DECODE_ROWS else "prefill" + unit, scale, digits = ("us", 1.0, 2) if kind == "decode" else ("ms", 1e-3, 3) + print( + f"\n{kind}: rows={rows}, context={context}, selected={max_selected}, {iterations} runs x {args.rounds} rounds" + ) + print(f"| configuration | attention {unit} | total {unit} | speedup | rms err |") + print("|---|---|---|---|---|") + results = [] + for index, (label, options) in enumerate(configs): + ratios = [base / value for base, value in zip(attention[0], attention[index], strict=True)] + attention_scaled = [value * scale for value in attention[index]] + total_scaled = [value * scale for value in total[index]] + print( + f"| {label} | {format_range(attention_scaled, digits)} | {format_range(total_scaled, digits)} " + f"| {format_range(ratios, 2)}x | {rms_error[index]:.2e} |" + ) + results.append( + { + "label": label, + "options": options, + "attention_us": attention[index], + "total_us": total[index], + "kernels_us": kernels[index], + "rms_error": rms_error[index], + } + ) + if args.breakdown: + for name, value in sorted(kernels[index].items(), key=lambda item: -item[1])[:5]: + print(f"| {name[:90]} | {value:.2f} us | | | |") + return { + "kind": kind, + "rows": rows, + "context": context, + "selected": max_selected, + "iterations": iterations, + "configurations": results, + } + + +def environment(): + properties = torch.cuda.get_device_properties(torch.cuda.current_device()) + return { + "onnxruntime": ort.__version__, + "onnxruntime_build_info": ort.get_build_info(), + "torch": torch.__version__, + "gpu": properties.name, + "sm_count": properties.multi_processor_count, + "compute_capability": f"{properties.major}.{properties.minor}", + "python": platform.python_version(), + "command": " ".join(sys.argv), + } + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--rows", type=int, nargs="+", default=[8, 128, 8192], help="query rows; <= 8 selects decode") + parser.add_argument("--context", type=int, nargs="+", default=[8192], help="total sequence lengths") + parser.add_argument("--selected", type=int, default=2048, help="selected candidates per row") + parser.add_argument( + "--tiles-per-split", type=int, nargs="+", default=[0], help="decode tc tiles per split (0=auto)" + ) + parser.add_argument("--rounds", type=int, default=7, help="alternating rounds over all configurations") + parser.add_argument("--iterations", type=int, default=None, help="profiled runs per round (default by kind)") + parser.add_argument("--prefill-iterations", type=int, default=5) + parser.add_argument("--decode-iterations", type=int, default=200) + parser.add_argument("--warmup", type=int, default=3, help="unprofiled runs before each measurement") + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--reference-rows", type=int, default=64, help="query rows checked against FP64 (0=skip)") + parser.add_argument("--rc-model", default="", help="model whose DynamicSparseAttention node supplies the scales") + parser.add_argument("--layer", type=int, default=0, help="index of that DynamicSparseAttention node") + parser.add_argument("--query-std", type=float, default=1.4, help="BF16 query values ~ N(0, std^2); 1.4 is typical") + parser.add_argument("--cache-code-std", type=float, default=40.0, help="INT8 cache codes ~ N(0, std^2), clipped") + parser.add_argument("--key-scale-multiplier", type=float, default=1.0) + parser.add_argument("--value-scale-multiplier", type=float, default=1.0, help="1.2e7 gives scales above 2^16") + parser.add_argument("--breakdown", action="store_true", help="also print the largest kernels of each configuration") + parser.add_argument("--json", default="", help="write the raw per-round timings and the environment here") + return parser.parse_args() + + +def main(): + args = parse_args() + torch.cuda.set_device(args.device) + key_scale, value_scale = make_scales(args) + info = environment() + print(json.dumps(info, indent=2)) + cases = [] + for rows in args.rows: + for context in args.context: + if rows > context: + continue + cases.append(run_case(rows, context, args, key_scale, value_scale)) + if args.json: + with open(args.json, "w") as file: + json.dump({"environment": info, "arguments": vars(args), "cases": cases}, file, indent=1) + + +if __name__ == "__main__": + main()