Repository navigation
[CUDA] Optimize CUDA prefill and quantized paged attention - #34117
Merged
Merged
Conversation
Append optional per-token grouped deltas and numeric previous-term masks. Preserve staged rounding, static behavior, raw state capture, and in-place state ownership. Reject the dynamic mode on WebGPU and update schema documentation and focused tests. Validation: CUDA provider and operator test binary build successfully; three CPU schema tests pass. CUDA numerical equivalence, graph replay, legacy regressions, and benchmarks remain pending an isolated GPU.
Keep grouped BF16 block32 tensor-core prefill opt-in, guard rounding boundaries, and recompute ambiguous outputs with baseline arithmetic. Add exact-output and shadow checks and explicit CUDA header rebuild dependencies. Validated on H200 with 11 tests and 306 subtests, full 512-token baseline comparisons across 8K-128K, and zero missed corrections in the qualified shadow checks. Guard windows remain empirical; no universal numerical guarantee or across-context 3.6x performance claim.
Add GQA6 XQA kernel specializations, dispatch, and resource checks while preserving unsupported-format fallbacks. Cover RC geometry, long contexts, graph replay, cache writes, scale modes, and fallback boundaries.
Skip invisible local-window prefix tiles and vectorize H128/GQA6 per-channel gathers while preserving dense offsets, FlashAttention arithmetic, cache writes, and decode dispatch. Add default-on switches with scalar/full-prefix opt-outs, exact parity and dispatch tests, and operator documentation. Validated on H200: 49 quantized tests (169 subtests), 36 native regressions, and 150 measured requests. All 50 INT8 A/B pairs preserve exact 512-token outputs and speculative counters.
Use regular FlashAttention 64-query/64-key tiles for eligible causal global and local INT8 prefill while retaining the existing key tiling and exact tested arithmetic. Add optional CPU epoch/capacity/final-chunk control for bounded dequantized prefill shadows, append-only suffix conversion, and guarded stateless fallbacks. Preserve INT8 decode caches and dispatch. Add operator documentation and exact parity, epoch/reset/release, invalid-control, fallback, and graph replay coverage. Existing H200 validation: 54 quantized tests (186 subtests), 36 native regressions, and 200 measured requests with 100 exact token/counter comparison pairs.
kunal-vaishnavi
changed the title
# Conflicts: # onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved critical issues can cause CUDA misaligned-address failures for valid INT8 cache bindings and silently incorrect CPU results when dynamic convolution inputs are ignored.
5 open issues
What changed in this PR
This PR extends ONNX Runtime’s CUDA inference paths to reduce quantized prefill overhead and support additional BF16 paged-attention configurations.
Changes:
- Adds dynamic grouped two-tap convolution and opt-in grouped INT3 MoE prefill with guarded tensor-core acceleration.
- Extends BF16 INT8/INT4 XQA support and optimizes INT8 prefill gathering, query tiling, and controlled prefix reuse.
- Expands correctness, dispatch, fallback, and lifecycle tests, alongside operator documentation and CUDA build dependencies.
| File | Description |
|---|---|
| onnxruntime/test/python/transformers/test_qmoe_shared_direct.py | Tests grouped INT3 prefill and tensor-core compatibility. |
| onnxruntime/test/python/transformers/test_paged_attention_int4.py | Tests quantized attention dispatch, parity, and shadow-cache lifecycle. |
| onnxruntime/test/contrib_ops/cuda_kernels/softmax_topk_kernel_test.cc | Extends grouped-kernel selection coverage. |
| onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc | Tests dynamic convolution validation, arithmetic, and state aliasing. |
| onnxruntime/core/graph/contrib_ops/bert_defs.cc | Extends convolution and paged-attention schemas. |
| onnxruntime/contrib_ops/webgpu/bert/varlen_causal_conv_with_state.cc | Rejects CUDA-only dynamic convolution inputs. |
| onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h | Adds grouped-prefill and compatibility parameters. |
| onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu | Optimizes grouped kernels and shared-memory dispatch. |
| onnxruntime/contrib_ops/cuda/moe/qmoe_int3_tc.cuh | Adds guarded tensor-core computation and canonical recomputation. |
| onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc | Integrates opt-in grouped prefill and workspace budgeting. |
| onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_spec_dec_bf16_int8_128.cu | Adds BF16 H128 INT8 speculative specialization. |
| onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_spec_dec_bf16_int4_128.cu | Adds BF16 H128 INT4 speculative specialization. |
| onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu | Extends XQA dispatch and resource queries. |
| onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_int4_128.cu | Adds BF16 H128 INT4 single-token specialization. |
| onnxruntime/contrib_ops/cuda/bert/varlen_causal_conv_with_state.h | Declares dynamic convolution parameters. |
| onnxruntime/contrib_ops/cuda/bert/varlen_causal_conv_with_state.cc | Validates and forwards dynamic convolution inputs. |
| onnxruntime/contrib_ops/cuda/bert/varlen_causal_conv_with_state_impl.cu | Implements dynamic grouped convolution. |
| onnxruntime/contrib_ops/cuda/bert/paged_attention.h | Adds optimization flags and synchronized shadow state. |
| onnxruntime/contrib_ops/cuda/bert/paged_attention.cc | Integrates quantized dispatch and prefix-cache lifecycle controls. |
| onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu | Adds window-aware and vectorized INT8 gathering. |
| onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim128_bf16_causal_sm80.cu | Adds the 64-query-row FlashAttention kernel. |
| onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h | Exposes optional tile64 selection. |
| onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc | Validates and dispatches tile64 attention. |
| onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.h | Declares prefill-aware debug reporting. |
| onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.cc | Identifies optimized INT8 prefill in telemetry. |
| onnxruntime/contrib_ops/cuda/bert/attention_data.h | Carries gather, tiling, and prefix-offset settings. |
| docs/OperatorKernels.md | Updates convolution input listings. |
| docs/ContribOperators.md | Updates dynamic convolution attributes and input listings. |
| docs/contrib_ops/cuda/paged_attention.md | Documents quantized paths and shadow-cache lifetime requirements. |
| cmake/onnxruntime_providers_cuda.cmake | Adds explicit MoE header dependencies for incremental builds. |
Review effort: Balanced
Tianlei Wu (tianleiwu)
requested changes
Nov 12, 2026
Tianlei Wu (tianleiwu)
left a comment
Tianlei Wu (tianleiwu)
left a comment
Contributor
There was a problem hiding this comment.
Thanks for the detailed docs and the strong exactness coverage on the INT8 gather paths. I reviewed this for benefit per feature, default-path impact, backward compatibility, GPU generality, and precision. I also ran a couple of micro-benchmarks on an A100 (sm80) to check behavior off H200.
Per-feature benefit
| Feature | Default | Evidence | Assessment |
|---|---|---|---|
| Window-aware + vectorized INT8 gather | on | §8.10: 128K gather 3.01 s → 0.71 s | Clear win. Gating is hard-coded to H128 / kv_heads=8 / GQA6 (see inline) |
| H128/GQA6 BF16 INT8/INT4 XQA (decode + spec-dec) | on for matching shapes | decode table | Plausible. The table is confounded by speculative acceptance: BF16 decode is non-monotonic in context (446→227→195→302→211 tok/s). Please add verify-step latency or acceptance rate |
64-row FA tile (ORT_ENABLE_PAGED_INT8_TILED_PREFILL) |
on, all GPUs | not isolated; H200 only | GPU-specific. A100 global 8K chunks are 3–7% slower. On sm86/sm89 it changes the local-window kernel and is not bit-exact (inline) |
BF16 prefill shadow (prefill_cache_control) |
graph opt-in, env default on | not isolated | Questionable. Upper bound ≈0.71 s of 27.5 s TTFT at 128K vs ~3.76 GB retained, plus a stateful kernel (inline) |
| Grouped INT3 MoE prefill + TC compatibility repair | env opt-in | none in PR | Please share prefill numbers to justify ~400 lines of kernel code and 5 env knobs |
| Dynamic grouped two-tap conv | graph opt-in | functional enabler | Needed, but the prefill kernel is serial per channel: 12–14x slower than token-parallel (inline) |
Default-path impact
- Without
prefill_cache_control, the shadow adds no allocation, lock, or work. Grouped INT3 prefill is env-gated. Dynamic conv only branches onkernel_delta. 👍 - Default-on changes: windowed gather, vectorized gather, and tile64. The first two are bit-exact and tested. tile64 is the risky one (inline).
qmoe_kernels.curelaxes grouped-kernel eligibility for all callers, including the non-opt-inqmoe_offload.ccpath (inline).- The INT3 env vars are parsed on every
ComputeInt3call (inline).
Backward compatibility
- New inputs are optional trailing inputs and new attributes have defaults, so existing models are unaffected. 👍
prefill_cache_controlis a pure hint; other EPs and older graphs keep stateless behavior. 👍- Main gap: the CPU
VarlenCausalConvWithStatekernel silently ignoreskernel_delta/previous_mask. I agree with the existing open thread.OperatorKernels.mdnow advertises those inputs for CPU too. - Design concern:
prefill_cache_controlmakes correctness depend on the caller renewing epochs whenever cache bytes or page mappings change. A stale epoch silently returns attention over stale dequantized K/V. That is a new class of contract for an ORT kernel.
GPU generality
Most fast paths are keyed to a single model geometry: H128 with kv_heads=8 and GQA6 (kHiddenSize = 8 * 128), INT3 block 32, and group-6-only XQA. The window skipping and the 8-wide vectorization are generic and would benefit every INT8 GQA model if gated on head_size % 8 == 0 and alignment instead. tile64 and the TC-repair ULP bound were tuned or validated on one GPU. The shadow's ~3.76 GB at 128K is prohibitive on 12–32 GB consumer GPUs.
Precision
- INT8 gather, vectorized gather, shadow, and tile64 on sm80/sm90 use the same arithmetic (bit-exact). 👍
- The INT3 SIMT grouped path is bit-exact vs the existing SIMT path, and the test asserts equality. 👍 TC repair is heuristic and arch-dependent (inline).
- The dynamic conv spec makes four intermediate roundings to T normative (inline).
- BF16 H128 XQA reuses the fold design. The PER_CHANNEL K scale goes into a BF16 copy of Q, and the V scale is applied to the BF16 output in a second pass. That adds roughly one extra BF16 ULP vs the portable FP32 kernel. The current tolerances cover it. Applying the V scale in the XQA epilogue in FP32 would remove one rounding.
Also agree with the open threads
The INT8 cache uint2 alignment issue (I replied with a shadow-specific caveat), the CPU kernel ignoring the dynamic inputs, the empty-query shadow release, and TensorCore=1 test gating on pre-SM80 devices.
Nits
- Test names embed an internal codename (
rc14, "RC14 four-row ..."). Shape-based names (e.g.h128_gqa6) would age better. - The
OBJECT_DEPENDSadded inonnxruntime_providers_cuda.cmakeshould be unnecessary, since nvcc depfiles already track these includes. If it works around a specific generator issue, a comment would help.
Requesting changes mainly for: tile64 arch gating plus the sm86/sm89 exactness break, the shadow's cross-run arena ownership, and the serial dynamic-conv prefill kernel. Splitting the shadow cache and the INT3 grouped prefill into follow-up PRs with ablation data would also make this much easier to land.
Recognize SM121 in the SM120 kernel family and add a native CUDA plug-in registration smoke utility. Add guarded compact-group scale handling and BF16 group-64 small-matrix tiling controls with regression coverage. Compact-group expansion and non-default tiling remain experimental and disabled by default; neither is promoted by the RC performance results. Validated with focused ARM64 provider builds, native output comparisons, and CPU scale-equivalence checks.
…o kvaishnavi/mai-coder-dflash # Conflicts: # cmake/external/cuda_configuration.cmake
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Remove compact-group scale expansion and the BF16 group-64 small-matrix row/column tile overrides, along with their controls, helper header, and experiment-specific tests. Restore the normal fixed dispatch and original layout tests while retaining checked scale allocation and the BF16 GEMV group-size safety guard.
Expose six session keys for INT8 shadow reuse and INT3 QMoE tuning. Resolve controls during kernel construction with session values overriding the environment, and validate boolean and floating-point values. Keep eligible INT8 prefill gathering and tiling automatic without switches. Default shadow reuse off, update parity and configuration tests, and document GenAI session configuration.
Akshay Sonawane (apsonawane)
left a comment
Akshay Sonawane (apsonawane)
left a comment
Contributor
There was a problem hiding this comment.
Reviewed the current 41-file change at 4751cc1, including attention dispatch/gather/shadow lifetime, BF16 INT4/INT8 XQA wiring, dynamic convolution/schema/provider behavior, grouped INT3 MoE packing and rounding, and related build/test changes.
No additional distinct actionable bugs were substantiated beyond the existing review findings. Instead of opening duplicate threads, I added source-level corroboration and concrete regression guidance inline for:
- P1: unaligned INT8 cache vector loads
- P1: CPU silently ignoring dynamic convolution operands
- P2: empty-query calls not releasing retained shadow state
These correctness issues should be addressed before merging.
Validation: the documentation test file passed all 20 tests and 13 subtests. Host-side checks covered packing/workspace offsets, BF16 rounding, address alignment, and source-level lifetime/provider behavior. CUDA tooling/GPU execution was unavailable; the installed local CPU wheel also lacks this convolution kernel, so this is not a claim of full CUDA or expanded-schema CPU runtime validation. The observed React Native iOS CI failure occurred during npm dependency bootstrap with HTTP 504, before application tests.
Allow explicit CUDA session settings to override the XQA environment switches for PagedAttention. Add dispatch and validation regressions and document the GenAI configuration keys. Verified with nine focused tests and a Release compile probe.
Tianlei Wu (tianleiwu)
requested changes
Nov 13, 2026
Tianlei Wu (tianleiwu)
left a comment
Tianlei Wu (tianleiwu)
left a comment
Contributor
There was a problem hiding this comment.
Thanks for the quick iteration. I re-reviewed at 47a5a0c (7 commits since my previous review at 9d05498).
Addressed
- The BF16 prefill shadow and
prefill_cache_controlare removed cleanly. That resolves the arena-ownership and benefit/cost threads. - The QMoE INT3 knobs and the PagedAttention XQA switches are resolved once in the constructor, with strict validation.
- The tile64 test now checks against the numerical reference.
ContribOperators.mdkeeps the schema descriptions.- The bot threads on empty-query shadow release and docs regeneration are now moot or addressed and can be resolved.
Blocking
- MatMulNBits BF16 INT4 GEMV guard (new, inline). This is a default-path change outside the PR's scope. In the default compact build it disables GEMV for every eligible BF16 INT4 node.
- tile64 has no opt-out (reply on the existing thread).
d4a1cd1removed the env switch, so it is on for every SM but was measured only on H200. - Still open from round 1 (existing threads):
- The CPU
VarlenCausalConvWithStatekernel silently ignoreskernel_delta/previous_mask, whileOperatorKernels.mdlists them under CPU. - The vectorized INT8 gather issues
uint2loads without checking cache-pointer alignment.
- The CPU
Should fix or track as follow-ups (existing threads unless noted)
- Dynamic conv prefill is serial per channel.
- The
TensorCore=1assertions fail when the TC path correctly falls back. - Grouped QMoE eligibility is widened for the offload path.
- The TC-compatibility ULP bound is arch-dependent.
- The
kInt8PrefillTileAlignment/kBlockN invariant, the hard-codedkHiddenSize, and the duplicated FP16 tile64 TU are unchanged. - New inline suggestions:
- size the vectorized gather's grid per row
- keep the diagnostic QMoE knobs env-only
- drop the Windows-only plugin smoke tool
- restore the deleted
needs_dense_kvcomment
- There are still no prefill numbers for the opt-in INT3 grouped path (SIMT grouped vs. TC+repair vs. the default TC prefill).
PR description
- Bullets 5-6 still describe the shadow cache.
- The benchmark notes say the INT8 runs used "tile64 + shadow", so the 64K/128K INT8 TTFT rows may not reflect this head.
- The description doesn't mention:
- the MatMulNBits guard
- the 7 new public
ep.cuda.*session keys - the offload grouped-eligibility change
- the plugin smoke tool
- the
gen_contrib_doc.pycheck
Please update the description and re-run the INT8 TTFT rows.
Reject unsupported CPU dynamic convolution, guard and generalize INT8 gather, preserve architecture-specific tile/GEMV choices, parallelize long dynamic convolution without changing rounding, and retain legacy offload eligibility. Clarify experimental compatibility controls and regenerate docs. Local CUDA and C++ regressions pass; full quantlm before/after runs preserve all tokens and counters with timing changes below 0.3%. SM121 device compilation is verified, not hardware performance qualification.
Contributor
Author
|
From Copilot: Description
Performance PreservationThe tuning boundaries are explicit: commits through The following is a before/after comparison of the review follow-up on H200/SM90, not a BF16-versus-INT8 cache comparison. Both quantized-LM packages use INT8 target KV, BF16 draft KV, and their existing quantized heads. Batch 1; summary prompts; 512 output tokens; one warmup and three measured requests per case; identical GPU, package hashes, and presets. TTFT and decode values are medians. All 30 paired measured generations have exact token and speculative-counter parity, so acceptance changes do not explain the timings.
Absolute TTFT and decode changes are below 0.3% on these runs. This preserves the measured H200 behavior; it is not a new speedup claim or model-quality qualification. Grouped-Prefill AblationH200/SM90, BF16, hidden 1536, intermediate 2688, two experts, top-k 2, block size 32. GPU-resident whole-op wall latency includes casts/routing; prepacking is disabled. Seven samples per cell; medians below in milliseconds. Dispatch was checked separately from timed runs. These are fixture measurements, not model throughput.
Grouped SIMT and TC+repair match the canonical output exactly on these fixtures. Default TC is faster for larger rows but numerically different (relative error about 0.0031-0.0040), so it is not reported as canonical parity. The repair mode uses the package-calibrated 1.5/4 tolerances; its extra absolute-value MMA, workspace clear, and repair launch costs are documented. It remains opt-in and requires application/device-specific qualification. Validation and Limits
Lossy KV quantization and experimental TC compatibility still need separate model-quality qualification. The removed shadow path is not part of any current result above. |
Tianlei Wu (tianleiwu)
left a comment
Tianlei Wu (tianleiwu)
left a comment
Contributor
There was a problem hiding this comment.
Thanks for the thorough follow-up. I re-reviewed at b3be80f and checked each resolved thread against the code, not only the replies.
Verified fixed
- tile64: gated to SM90/SM121 via
UseInt8PrefillTile64ForDevice; other architectures keep their default Flash dispatch. §8.11/§8.12 no longer claim universal exact order. - Gather/tile invariant:
GetPagedKvGatherAlignmentplus thestatic_assertinrun_flash_fwdcover every regular kernel. The gather's window start (local_window_size - 1) matches thewindow_size_leftpassed tomha_varlen_fwd, so Flash never reads a skipped tile. The split-KV traits (kBlockN128/256) aren't covered by the assert, but they can't be reached here because varlen split-KV requires pure decode. - Vectorized INT8 gather: requires 8-byte-aligned K and V base pointers. With
head_size % 8 == 0, everyuint2load anduint4store is then aligned. The generic path matches the scalar gather bit-for-bit (dequant, page indexing, missing pages), and the grid is sized per sequence frommax_kv_len. The offset-binding tests cover this well. - Dynamic conv: the prefix plus token-parallel body matches the serial kernel term for term (taps, group, delta indices, mask, capture positions,
final_state = input[end - 1]). It is alias-safe: onlyMayInplace(4, 1)is declared, the prefix readsinitial_statebefore writingfinal_state, and the body never reads state. - CPU: rejects the dynamic attributes at construction and inputs 6/7 in
Compute.Input<Tensor>()returns null for an out-of-range index, so six-input nodes are unaffected. - QMoE: legacy grouped eligibility keeps the float FC2 accounting. Compact accounting applies only to the opt-in path, where it matches the kernel's
sizeof(T)staging; that staging is numerically neutral. The diagnostic knobs are internal-only, and the ablation table answers the benefit question. - XQA env fallback accepts legacy nonzero values, TC tests tolerate the SIMT fallback, and the smoke tool is removed. The dtype-split tile64 TUs, copyright headers, and
needs_dense_kvcomment are all in.
Remaining (non-blocking)
- SM121 BF16 INT4 GEMV guard (inline): please add the rationale. In the default compact build the
block_size == 64exception can't apply, so the guard turns off BF16 INT4 GEMV on SM121 for every eligible node. - Inline nits: the opt-in TC canary in the QMoE test, and the
ef4575ahash in the docs. - Carried-over nits:
RC14remains inpaged_attention.md§8.10,DynamicBFloat16RC14, and the "RC14 four-row grouped prefill" selector case.- The
OBJECT_DEPENDSaddition inonnxruntime_providers_cuda.cmakestill has no comment explaining why nvcc depfiles aren't enough.
The blocking items from my previous reviews are resolved. CI was still running at review time (27 pass, 55 pending, none failed).
Tianlei Wu (tianleiwu)
approved these changes
Nov 13, 2026
Tianlei Wu (tianleiwu)
enabled auto-merge (squash)
November 13, 2026 03:44
Tianlei Wu (tianleiwu)
added a commit
that referenced
this pull request
Nov 13, 2026
This cherry-picks the following commits for the release: | Commit ID | PR Number | Commit Title | |-----------|-----------|-------------| | 85dcb54 | #32475 | Update Eigen to 5.0.1 + SSE2 fix (broken Clip & Sign) | | 1ea552d | #33976 | [CUDA/CPU] Add hybrid GPU/CPU expert offload for MoEPacked block-packed weights | | 3a8b8f2 | #34066 | [CUDA] Optimize multi-token MatMulPacked GGUF GEMV (M=2..8) and eliminate ReadU32Dev warp divergence | | 64506ea | #34098 | [CUDA] Guard tests requiring SM80 on older GPUs | | 74f861c | #34067 | [MoE] Optimize multi-token (M=4, 8) MoEPacked CPU hybrid cold routes and GPU FC2 decode | | 6aa9b27 | #33963 | [CUDA] Extend MoEPacked for top-k, dispatch and fused MoE calculation | | 8fb1b27 | #33969 | [CUDA] Add Mamba2Verify CUDA Kernel | | cedae42 | #33994 | [CUDA] Update `MoEGroupRouter` for routing to `MoEPacked` | | 158d0c0 | #34002 | [CUDA] Speed up native Bolt packed projections and IQ2 experts | | 38c6dce | #34003 | [CUDA] Reduce native packed decode tile and FC2 loader overhead | | d5704db | #34147 | [CUDA] Optimize MoEPacked, MoEGroupRouter, and MatMulPacked for Nemotron Bolt | | 7cca0e5 | #34126 | Fix Windows junction detection in model_package for non-MSVC standard libraries | | 2bf400a | #34099 | Fix external data validation for Hugging Face Hub shared blob store | | d2f77cd | #34164 | [CUDA] Accelerate Nemotron Bolt large-chunk prefill across contrib ops | | ca7ea48 | #28457 | Add set_default_logger_callback Python API for custom ORT logging | | fca9447 | #34117 | [CUDA] Optimize CUDA prefill and quantized paged attention | | 1804c6f | #34114| [CUDA] Optimize MoEPacked GGUF offload 3-stage prefill pipeline and F16C/Batch4/40-core cold decode | --------- Signed-off-by: Christian Bourjau <christian.bourjau@quantco.com> Signed-off-by: rui-ren <ruiren1225@gmail.com> Co-authored-by: Christian Bourjau <cbourjau@users.noreply.github.com> Co-authored-by: BBaditi <bbaditi@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Rui Ren <ruiren1225@gmail.com> Co-authored-by: Mustapha Jaber <mustjab@users.noreply.github.com> Co-authored-by: Mustapha Jaber <mustjab+microsoft@microsoft.com> Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b Copilot-Session: c8741c5a-3e9e-4860-b846-7d6fb0025040
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
VarlenCausalConvWithStatewith dynamic grouped two-tap convolutions, including per-token coefficient deltas, boundary masking, and schema validation.prefill_cache_controlinput supplies an epoch, capacity, and final-chunk flag to explicitly govern reuse.Motivation and Context
There are existing bottlenecks in INT3 expert computation and repeated quantized-KV gathering, alongside gaps in BF16 quantized-cache speculative attention. These changes provide the CUDA primitives needed to reduce prefill overhead and support lower-memory speculative decoding.
Optimizations are restricted to eligible configurations, with grouped INT3 prefill remaining opt-in. Prefix reuse requires explicit caller-managed lifetimes and trades temporary BF16 workspace for reduced dequantization work. The companion GenAI changes supply the necessary graph lowering and lifecycle controls. Model weights are not requantized; quantized-KV accuracy still requires application-specific qualification.
Benchmarks
Notes:
H200, batch 1, 512 output tokens, one warmup and five measured requests. The below benchmarks compare BF16 KV cache and INT8 KV cache versions of the RCs. TTFT and decode are medians; peak GPU memory includes weights and reserved pools, in decimal GB. INT8 KV cache versions use the optimized tile64 + shadow path; draft KV remains BF16 throughout.
Memory usage is not a fair comparison between ORT GenAI and llama.cpp because ORT GenAI tries to capture as much GPU memory as it can for future conversations. Memory usage is simply noted here for recording purposes.
INT3 RC
INT2 RC
llama.cpp (2-bit)