Repository navigation
[CUDA] Give each plugin EP session its own device arena - #32807
Merged
Tianlei Wu (tianleiwu) merged 4 commits intoSep 26, 2026
Merged
Conversation
The CUDA plugin EP factory returned one BFC device arena per GPU to every CreateAllocator caller. A captured CUDA graph keeps reading and writing the intermediate chunks its session freed back to the arena after capture, so a second session (or the environment allocator) could be handed those chunks and overwrite them between replays. With two graph-captured sessions on one device (target + DFlash2 drafter) this silently corrupted generation. Create a device arena per CreateAllocator call, matching the bundled CUDA EP, and track quarantine/abandon state per arena. Stream run-end reset and undrained-release quarantine now visit every arena on the device.
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Environment allocator sharing still permits unsafe cross-session graph reuse, and allocator publication is not exception-safe.
Get a fresh assessment by requesting another Copilot review.
Review effort: Balanced
Findings: 1
Open (3)
What changed in this PR
Creates per-caller CUDA device arenas to prevent cross-session memory reuse from corrupting CUDA graph replay.
Changes:
- Replaces the shared device arena with independently tracked arenas.
- Updates allocator lifecycle and stream cleanup handling.
- Adds arena-isolation coverage and documentation.
| File | Description |
|---|---|
cuda_ep_factory.h |
Defines per-device arena collections and state. |
cuda_ep_factory.cc |
Creates, releases, resets, and quarantines individual arenas. |
cuda_plugin_arena_test.cc |
Tests allocator isolation across sessions and environment. |
cuda_plugin_ep_design.md |
Documents allocator ownership. |
cuda_graph_for_cuda_plugin.md |
Explains graph-safety requirements. |
arena_allocator_migration_design.md |
Updates the arena migration design. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Baiju Meswani (baijumeswani)
left a comment
Baiju Meswani (baijumeswani)
left a comment
Contributor
There was a problem hiding this comment.
I found one remaining CUDA graph safety gap and three smaller issues. Details are inline.
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Undrained streams can still reference untagged arena allocations that the new conditional abandonment logic may free.
Get a fresh assessment by requesting another Copilot review.
Review effort: Balanced
Findings: 1
Open (2)
Baiju Meswani (baijumeswani)
approved these changes
Sep 26, 2026
Tianlei Wu (tianleiwu)
deleted the
tlwu/20260925/cuda-plugin-session-arena
branch
September 26, 2026 05:57
Tianlei Wu (tianleiwu)
added a commit
that referenced
this pull request
Sep 29, 2026
…Attention (#32852) ### Description Replace the bias-only rank-one GEMMs in CUDA `Gemm` and `com.microsoft::DecoderAttention` with direct bias broadcasting on the consuming CUDA stream. The main cuBLAS GEMMs remain unchanged. - Add a broadcast helper for scalar, row, column, and matrix bias, with 64-bit element counts/offsets and a fast 32-bit division path for smaller outputs. - Remove both operators' calls to `GetConstOnes` and delete the now-unused bundled/plugin constant-ones cache infrastructure. - Correct the tightly coupled zero-K Gemm path to apply `beta` when the main GEMM is skipped. - Add FP32/FP64/FP16/BF16 Gemm broadcast, coefficient, transpose, zero-K, and empty-output coverage. - Extend DecoderAttention coverage to FP32/FP16 across self/cross-attention and cache paths. - Add bundled/plugin CUDA regressions for concurrent sessions with increasing dimensions and multiple graph IDs replayed after/while another session runs larger shapes. Tests disable CPU fallback. **Validation** - Formatting and whitespace checks passed. - Secret scans passed. - Supplemental read-only code review found no significant issues. - Built the CUDA 13.0 plugin EP in Release mode for SM80 and ran CUDA-only Gemm and DecoderAttention workloads on an NVIDIA A100-SXM4-80GB. CPU fallback was disabled and profile-based placement checks passed. - Automated review was unavailable because its configured model could not be loaded. CodeQL skipped analysis because its database was too large; this is not a clean security-scan result. Suggested provider-test filter: `CudaGemmBiasBroadcastTest/*.*:DecoderAttentionTest/*.*:CudaGemmBiasSessionTest.*` in both bundled and plugin CUDA builds. ### Performance Compared exact base `925940051a` with implementation head `2ad4765479` using identically configured CUDA 13.0 plugin EP builds on an A100-SXM4-80GB. Later commits only add tests and remove the unused cache infrastructure without changing the measured execution path. Each result is the median of five order-balanced process trials with 100 warmups and 500 CUDA-event iterations per trial, device-resident I/O, and strict plugin-library selection. | FP16 workload | Base median | PR median | Change | |---|---:|---:|---:| | Gemm row bias, `M=1, N=4096, K=0` | 50.176 us | 40.960 us | **-18.4%** | | Gemm row bias, `M=32, N=4096, K=64` | 63.488 us | 57.344 us | **-9.7%** | | Gemm row bias, `M=512, N=4096, K=4096` | 164.864 us | 156.672 us | **-5.0%** | | Gemm column bias, `M=4096, N=32, K=64` | 63.488 us | 57.344 us | **-9.7%** | | DecoderAttention decode, batch 1, hidden 768 | 167.936 us | 160.768 us | **-4.3%** | | DecoderAttention prompt, sequence 128, hidden 768 | 179.200 us | 169.984 us | **-5.1%** | Nsight Systems attributes the improvement to replacing rank-one cuBLAS dispatches with the direct broadcast kernel: - Zero-K Gemm bias GPU time decreased from 3.002 us to 2.390 us (-20.4%). - DecoderAttention's two affected bias operations decreased from 5.507 us to 4.741 us (-13.9%); total kernel time decreased by 2.1% because the projection GEMMs and attention work dominate. - The changed path is primarily launch/host-dispatch limited for these shapes. The broadcast itself becomes memory-bound at larger output sizes. No statistically credible regression was observed. `K=1`, scalar-bias, matrix-bias, and DecoderAttention batch-8 differences were within trial variance. With verified CUDA Graph replay, most gains narrowed to approximately 0-1% because graph replay already amortizes cuBLAS dispatch overhead; the large row-bias case retained a 1.0% improvement. DecoderAttention graph capture failed symmetrically on base and head in this standalone harness, so no graph-mode DecoderAttention latency comparison is reported. ### Motivation and Context Fixes #32847. Alternative to #32848. The merged #32807 isolates session BFC arenas, but the shared ones buffer bypasses those arenas and allocates directly through CUDA. Concurrent initialization/growth can expose inconsistent or uninitialized storage; later growth can invalidate an address retained by a CUDA graph even with serial session execution. Direct bias broadcasting eliminates the dependency at both consuming operators, avoiding a persistent shared cache, cross-stream initialization synchronization, and graph-lifetime management. More extensive bias/GEMM or bias/transpose fusion is intentionally deferred. Analysis and trade-offs: #32847 (comment) --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com> Co-authored-by: skottmckay <979079+skottmckay@users.noreply.github.com> Co-authored-by: Tianlei Wu <tlwu@microsoft.com>
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
The CUDA Plugin EP factory handed one BFC device arena per GPU to every
CreateAllocatorcaller: every session, plus the environment's shared allocator. After a CUDA graph is captured, the session frees the graph's intermediate chunks back to the arena, but the captured graph still reads and writes those addresses on every replay. Because the arena was shared, another session could be given those chunks and overwrite them between replays.With two graph-captured sessions on one device (a Qwen3.8-27B target plus a DFlash2 drafter in onnxruntime-genai), this silently corrupted generation. MMLU-Pro dropped 5.5 pp and speculative acceptance fell from 87.7% to 69.8%. The bundled CUDA EP is not affected because it creates a separate arena for each session.
This PR gives each
CreateAllocatorcall for device memory its own arena, which matches the bundled EP. It replaces the temporary graph-disable gate in #32801.Changes
cuda_ep_factory.h/.ccdevice_arenaand its refcount withstd::vector<DeviceArena> device_arenas. Each entry tracks its ownhas_quarantine/abandonedstate.CreateAllocatorImplcreates a new arena on every call and now honors that session's arena options, which the shared arena used to ignore with a warning.ReleaseAllocatorImpldestroys only the matching arena, or leaks it on purpose if it is quarantined or abandoned. Stream run-end reset and quarantine/abandon after an undrained release now visit every arena on the device. For abandon, this is the same scope the single shared arena had before.GetDeviceArenaForDevice()had no callers and is removed.cuda_plugin_arena_test.ccDeviceAllocator_IsSessionScopedtest: an allocation from session A must not show up in session B's arena stats or the environment allocator's stats. It fails on the current shared-arena plugin (NumAllocs3 vs 2) and passes with this change.docs/cuda_plugin_ep/*.mdThe pinned arena and the CUDA mempool allocator are still shared per device and are unchanged.
Motivation and Context
Root-cause isolation (MMLU-Pro, 200 questions, 8,192 max new tokens, DFlash2 drafts=6, concurrency 8)
Toggling only the arena sharing, with the same binary and both graphs on, flips the result. That identifies the shared arena as the cause.
This PR (graphs on for both sessions)
Performance
Batch 1, 32,768 prompt tokens, 1,024 generated tokens, DFlash2 draft width 6, 3 fresh-process repetitions, all arms run at the same time on separate H200s:
Re-enabling graphs safely recovers the 2.8% decode loss from the gate. Peak memory is 64 MiB below the bundled EP.
Tests
CudaPluginArenaTest.*,CudaPluginUserStreamGraphTest.*,CudaPluginDeviceDiscoveryTest.*: 51/51 passed on an H200 (Release, SM90).clang-format --dry-run --Werrorandgit diff --checkpass.