Repository navigation
Fix arena configuration handling and optimize CUDA GEMV and prefill - #33381
Open
Akshay Sonawane (apsonawane) wants to merge 4 commits into
Open
Akshay Sonawane (apsonawane) wants to merge 4 commits into
Akshay Sonawane (apsonawane) wants to merge 4 commits into
Conversation
Port standalone allocator fixes from the original Qwen optimization branch without its operator or model dependencies. Expose the CUDA arena pointer option, copy its configuration into provider-owned storage, and avoid splitting exact-size free chunks when the dead-byte limit is zero. Add allocation and pointer-option regression coverage.
Raise the split limit and vectorized occupancy target for the existing generic GEMV implementation. Cover higher split counts, workspace sizing, FP16/BF16 outputs, and repeated CUDA graph replay.
Port the temporal prefill optimization onto main without changing the operator schema or adding the Qwen branch state-update-active input. Preserve state aliasing with a separate final-state launch and retain the general fallback. Test threshold boundaries, capture clamping, all floating-point types, optional outputs, and aliased two-call continuation.
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| #pragma unroll | ||
| for (int k = 0; k < 4; ++k) { | ||
| sum += weights[k] * to_float(ReadStateOrInput( | ||
| input, channel_state, 0, channels, c, pad, t - pad + k * dilation)); |
Contributor
There was a problem hiding this comment.
Suggested change
| input, channel_state, 0, channels, c, pad, t - pad + k * dilation)); | |
| input, channel_state, 0, channels, c, pad, t - pad + k * dilation)); |
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
CUDA scheduling and alias-sensitive state updates warrant human review beyond the reported H200 validation.
0 open findings
What changed in this PR
Improves arena configuration handling and CUDA kernel parallelism without changing operator schemas.
Changes:
- Fixes exact-size arena allocation and CUDA arena-configuration ownership.
- Increases GEMV split-K parallelism and adds tiled causal-convolution prefill.
- Adds regression coverage for configuration lifetime, graph replay, and state aliasing.
| File | Description |
|---|---|
| onnxruntime/test/shared_lib/test_inference.cc | Tests CUDA arena pointer options and configuration lifetime. |
| onnxruntime/test/providers/cuda/test_cases/matmul_small_n_gemv_test.cc | Tests higher split counts, workspace sizing, and graph replay. |
| onnxruntime/test/framework/bfc_arena_test.cc | Tests exact-size allocation with zero dead bytes. |
| onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc | Tests tiled prefill boundaries and aliased continuation. |
| onnxruntime/core/session/provider_bridge_ort.cc | Supports arena pointer options and rejects unknown update keys. |
| onnxruntime/core/providers/cuda/math/matmul_small_n_gemv.cu | Raises split-K and occupancy targets. |
| onnxruntime/core/providers/cuda/cuda_provider_factory.cc | Copies arena configuration into owned storage. |
| onnxruntime/core/providers/cuda/cuda_execution_provider_info.h | Retains shared ownership of arena configuration. |
| onnxruntime/core/framework/bfc_arena.cc | Prevents splitting exact-size free chunks. |
| onnxruntime/contrib_ops/cuda/bert/varlen_causal_conv_with_state_impl.cu | Adds tiled prefill with a separate final-state update. |
| include/onnxruntime/core/session/onnxruntime_c_api.h | Documents arena pointer-option support and copy timing. |
🧠 Review effort: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
This branch has not been deployed
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.
Summary
Fix allocator edge cases and improve parallelism in existing CUDA GEMV and causal-convolution kernels. These changes are based directly on
mainand do not introduce model-specific operators or schema changes.Allocator fixes
max_dead_bytes_per_chunkis zero.default_memory_arena_cfgthrough the CUDA pointer-provider-option APIs.CUDA kernel improvements
VarlenCausalConvWithStateprefill workloads.Validation
Built the runtime, CUDA provider, and native test runners with CUDA 13.3 for H200. The focused build disabled unrelated optional FP4 QMoE, FpA-IntB GEMM, and fused-attention kernel families.
Total: 71 tests passed.
Convolution coverage includes dispatch boundaries, capture-count clamping, FP32/FP16/BF16, optional state updates, and aliased two-call continuation. GEMV coverage includes higher split counts and CUDA graph replay.