[Baiju Meswani (baijumeswani)]

Avoid allocating pinned host pointer buffers in the CUDA Split and Concat fast paths when all partitions are equal-sized and there are at most 32 pointers.

These paths pass a TArray directly as a kernel argument, so the CudaAsyncBuffer allocation and destruction were unnecessary. In a Qwen 3.8 decode trace, their removal reduced pinned allocation/free pairs from 1,635 to 1,355 and reduced time attributed to cudaFreeHost from 4.90 seconds to 11.3 milliseconds.

The GPU pointer-buffer fallback is unchanged for more than 32 pointers and for unequal Split/Concat dimensions.

Motivation

cudaFreeHost can synchronize outstanding device work. Allocating pinned pointer metadata before selecting the inline fast path introduced avoidable decode gaps even though the allocation was never copied to or consumed by the GPU.

Testing

• Release ARM64 CUDA build with CUDA 13.4 for SM120/SM121
• onnxruntime_provider_test --gtest_filter=*ConcatOpTest*:*SplitOperatorTest* (53 tests passed across four suites)
• Equal-sized fast paths and existing fallback-path coverage

[Copilot]

🟢 Approval recommended

The refactoring preserves existing dispatch behavior while removing unnecessary pinned allocations.

Pull request overview

Moves CUDA Split and Concat equal-size fast paths ahead of pinned-buffer allocation.

Changes:

• Passes up to 32 pointers directly through TArray.
• Retains GPU pointer buffers for larger or unequal operations.
• No actionable issues identified.

File summaries

File	Description
split.cc	Avoids pinned buffers for small equal splits.
concat.cc	Avoids pinned buffers for small equal concatenations.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 0
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Tianlei Wu (tianleiwu)]

Reviewed the reordered dispatch in both kernels. The change is functionally equivalent to the previous code and the optimization is well-motivated — CudaAsyncBuffer's constructor eagerly calls AllocCpuPtr() (pinned), so declaring it before the branch really did cost a pinned alloc/free pair per Run() even on the inline path that never touched it.

Correctness checks that pass:

• p.inputs is populated with exactly ctx->InputCount() entries by ConcatBase::PrepareForComputeImpl, so the hoisted concat_sizes loop indexing p.inputs[i] is safe.
• element_bytes / block_size_* are still computed after the p.output_num_elements == 0 early return, so p.output_axis_pitch / Shape()[p.axis] cannot divide by zero.
• The Concat fast path is reachable only when the output is non-empty and all concat dims are equal, which implies every input is non-empty — no null DataRaw() reaches the kernel.
• Split still materializes all output tensors via ctx->Output(i, ...) before the Shape().Size() <= 0 early return, preserving the contract that outputs are allocated for an empty input.
• TArray<T, 32>::TArray(int32_t) ORT_ENFORCEs size <= capacity, and both call sites are guarded by <= 32.
• Hoisting CopyToGpu above the equal/unequal branch is order-safe: all copies are issued on the same compute stream and consumed only by the following launch.

Nice side effect worth calling out: the fast path no longer touches pinned memory, which widens the set of CUDA-Graph-capturable Split/Concat configurations — the same property the existing Split3Inner comment documents.

Suggestions (see inline comments):

1. axis_dimension_input_output_mapping is still allocated and filled on the same_*_size && count > 32 path where it is never consumed. Since removing per-Run() allocations is the point of this PR, moving it into the unequal-size branch is in scope.
2. No test exercises the > 32 fallback, which is the branch whose structure actually changed. split_op_test.cc / concat_op_test.cc top out at a handful of inputs/outputs, so same_*_size == true && count > 32 is unexercised on CUDA. A case with 33 equal-sized inputs (Concat) and 33 equal-sized outputs (Split) would pin both the boundary and the GPU pointer-buffer fallback.

Minor, non-blocking:

• 32 now appears four times across the two files as both the TArray capacity and the dispatch guard. A shared constexpr int kMaxInlinePointerCount = 32; used for both would stop the guard and the capacity from drifting apart and turning an ORT_ENFORCE into a runtime failure.
• Neither file directly includes <algorithm> (for std::all_of) or <vector>; both rely on transitive includes. Pre-existing, but the PR rewrites those exact lines.