onnxruntime

55ab2dff - [CUDA][PagedAttention] Fuse small-batch sequence-length sanitizer and add session.paged_attention.skip_input_sanitizer (#33597)

[CUDA][PagedAttention] Fuse small-batch sequence-length sanitizer and add session.paged_attention.skip_input_sanitizer (#33597) ### Description - Add single-CTA `FusedSanitizeSequenceLengthsSmallBatchKernel` (`cub::BlockScan<int32_t, 256>`, `workspace_bytes = 0`) in `onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu` for `batch_size < 256`, replacing the 4-kernel `ClampCumulativeSequenceLengths` + `cub::DeviceScan::InclusiveScan(Max)` + `SanitizePastSequenceLengths` + `cub::DeviceScan::InclusiveSum` chain with a single kernel launch while preserving exact bit-identical semantics and keeping the general grid-wide CUB path for `batch_size >= 256`. - Add session configuration key `kOrtSessionOptionsPagedAttentionSkipInputSanitizer` (`"session.paged_attention.skip_input_sanitizer"`, with environment variable fallback `ORT_PAGED_ATTENTION_SKIP_SANITIZER=1`) allowing trusted runtimes that already validate `block_table`, `cumulative_sequence_length`, and `past_seqlens` to bypass `SanitizeBlockTableKernel`, `LaunchSanitizeSequenceLengths`, and `MaskInvalidSequenceOutputs` and compute `cumulative_seqlens_kv` via a single `LaunchGetCumulativeSeqlensKV` launch (`sequence_validity = nullptr`). - Update `docs/contrib_ops/cuda/paged_attention.md` to document both the small-batch fused sanitizer and `session.paged_attention.skip_input_sanitizer`. ### Motivation and Context - In speculative decoding with a 5-layer `PagedAttention` drafter (`DFlash2`, `batch_size = 1`), the 4-kernel sequence-length sanitizer plus `SanitizeBlockTableKernel` and `MaskInvalidSequenceOutputs` added 5 extra kernel launches and 2 host-side `cub::DeviceScan` workspace queries per `PagedAttention` layer (`25` extra kernel launches per drafter step, `~175` launches per 7-step cycle). - **Fused small-batch sanitizer (`skip_input_sanitizer = 0`, default)**: Collapses the 4 sequence-length sanitization launches into 1 single-CTA launch (`1.56 µs/call`, `workspace_bytes = 0`), eliminating 15 kernel launches per drafter step (`~105` launches/cycle) while retaining full defensive input sanitization. - **Skipping input sanitizer (`session.paged_attention.skip_input_sanitizer = "1"`)**: Eliminates `SanitizeBlockTableKernel` (`1.13 µs/call`), `MaskInvalidSequenceOutputs` (`1.08 µs/call`), replaces `FusedSanitizeSequenceLengthsSmallBatchKernel` (`1.56 µs/call`) with `GetCumulativeSeqlensKV` (`1.40 µs/call`), removes 3 scratch buffer allocations per layer, and skips the `FindTokenBatch` binary search in `ReshapeAndCache` (`2.79 µs -> 2.62 µs/call`), saving another 2 kernel launches (`2.54 µs` GPU kernel time) per `PagedAttention` layer (`10` launches per drafter step) and reducing end-to-end Qwen3.5-397B decode cycle latency on H200 by **`0.07–0.16 ms/cycle`** (`14.635 -> 14.550 ms/cycle` at `2K`, `14.748 -> 14.681 ms/cycle` at `8K`, `15.044 -> 14.948 ms/cycle` at `32K`). --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>