onnxruntime
1aa82412 - [CUDA] Fall back to FP32 in DSA INT8 TC kernel when rescale can flush a large-value-scale accumulator (#33847)

Commit
2 days ago
[CUDA] Fall back to FP32 in DSA INT8 TC kernel when rescale can flush a large-value-scale accumulator (#33847) ### Description Follow-up to #33802 addressing [this review comment](https://github.com/microsoft/onnxruntime/pull/33802#discussion_r4459726501): the INT8 tensor-core DynamicSparseAttention kernel could silently lose the contribution of earlier tiles when a later tile raises the running maximum, now that the INT8 TC prefill/decode kernels are enabled by default. The kernel rescales the PV accumulator with `alpha = FastExp2(old_max - new_max)`, which uses `ex2.approx.ftz.f32`. When the maximum rises by more than ~87 nats, `alpha` flushes to zero and erases the earlier accumulator. A large `value_scale` (up to the finite range) can make that earlier contribution significant. The existing large-value-scale FP32 fallback only tracked how far the *current tile's* scores sit below the new maximum (`lowest_depth`), so it did not fire in this case. ### Changes | File | Change | |---|---| | `contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh` | When `check_tail_depth` is armed (some `abs(value_scale) > 2^16`), the rescale `run_max - new_max` of a non-empty accumulator now also feeds `lowest_depth`. A rise of more than `kInt8TcTailDepth` (20 nats) makes the CTA take the existing `AttendRowFp32` fallback. Updated the comments accordingly. | | `test/contrib_ops/dynamic_sparse_attention_op_test.cc` | Added `Int8TensorCoreIncreasingMaximumLargeValueScale_CUDA`: candidate 0 scores 0 (value code 1), candidate 1 scores 88 in a later tile, `value_scale = 2^120`; expected output `2^120 * e^-88 / (1 + e^-88) ~= 8.0e-3` (the buggy path returned 0). Covered for prefill and for decode with 1/2/3 tiles per split. | | `docs/contrib_ops/cuda/dynamic_sparse_attention.md` | Documented the additional fallback condition. | The fallback is only armed for `abs(value_scale) > 2^16`, so models with ordinary value scales (e.g. Qwen3.8-Flash-Next, max 0.285) see no change in dispatch or performance. ### Testing - Built `onnxruntime_provider_test` (CUDA, SM80) and ran `DynamicSparseAttentionTest.*` on A100: 74 tests pass. - The new test fails without the kernel change and passes with it. - `lintrunner` (CLANGFORMAT) clean. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading