[Tianlei Wu (tianleiwu)]

Description

CUDA Concat's ragged path (inputs of unequal size along the concat axis) uploaded four CudaAsyncBuffers to the device on every Run, one of which is sized by the output axis dimension. This change recovers the input index by scanning prefix sums instead, which makes it possible to pass the per-input metadata by value in TArrays — matching what the equal-size path already does.

The ragged by-value path is dispatched before any CudaAsyncBuffer is created, like the equal-size path after #32410. So for up to 32 inputs, Concat allocates no pinned host memory at all. That matters under CUDA graph capture: cudaHostAlloc is not allowed while a stream is capturing, and the failed allocation surfaced as a misleading alloc failed (seen with a DFlash2 drafter graph that concatenates [1, H] and [N - 1, H] slices).

The result is also a large reduction in host-to-device traffic for graphs that concatenate many small tensors, which is common in decode-time LLM workloads.

Summary of Changes

File	Change
onnxruntime/core/providers/cuda/tensor/concat.cc	For up to 32 inputs, dispatch both equal-size and ragged inputs by value before the pinned pointer buffer is allocated; keep the buffer path for wider concats.
onnxruntime/core/providers/cuda/tensor/concat_impl.cu	Add the by-value ragged kernel; recover the input index by scanning prefix sums rather than reading the output-sized mapping table.
onnxruntime/core/providers/cuda/tensor/concat_impl.h	Declare the by-value launcher.
onnxruntime/test/providers/cpu/tensor/concat_op_test.cc	Ragged CUDA tests at 32 inputs (by-value) and 33 inputs (buffer fallback).
onnxruntime/test/python/onnxruntime_test_python_cudagraph.py	Ragged Concat on axis 0 and axis 1 under CUDA graph capture, with the inputs updated before each replay.

Notes:

• The existing buffer-based path is retained unchanged for concats wider than 32 inputs, so nothing regresses for large fan-in.
• Passing parameters by value leaves no host staging buffer for a captured memcpy node to re-read on replay.

Testing

• concat_op_test.cc: CudaRagged32InputsUsesByValueMetadata, CudaRagged33InputsUsesGpuMetadataFallback.
• onnxruntime_test_python_cudagraph.py: test_concat_unequal_sizes_cuda_graph_replay (all 6 tests in the file pass on H200, CUDA 13.0).
• Kernel output is bit-exact against NumPy across fp16/fp32/int64/int8, with zero-length inputs and 32/33/40-input ragged cases, with and without CUDA graph.
• End-to-end on DeepSeek-V4-Flash (8xH200, tensor parallel, CUDA graph): MMLU-Pro (800 samples) unchanged at 0.67000.
• End-to-end on a Qwen3.8-27B DFlash2 model with the drafter's CUDA graph enabled: previously failed at the drafter's ragged Concat on engine step 2; now runs cleanly (54/54 draft tokens accepted).

Motivation and Context

Measured on a DeepSeek-V4-Flash decode step (8xH200, tensor parallel, CUDA graph):

Metric	Before	After
Host-to-device copies per step	417	1
Step time	20.2 ms	18.8 ms
Batch-1 decode throughput	49.49 tok/s	53.21 tok/s

The saving is 0.46 ms of transfer plus 1.21 ms of the stalls that trailed those transfers. The benefit is general: any model whose graph concatenates a moderate number of tensors per step pays this cost today.

Checklist

• Tests added/updated
• No breaking changes
• Documentation updated (not applicable — no schema or API change)

[Copilot]

Pull request overview

Optimizes CUDA Concat for ragged inputs by passing metadata directly to kernels when there are at most 32 inputs.

Changes:

• Adds a by-value ragged Concat path.
• Recovers input indices by scanning prefix sums.
• Retains the existing buffer path for larger fan-in.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
concat.cc	Selects and prepares the by-value path.
concat_impl.cu	Implements the new CUDA kernel and launcher.
concat_impl.h	Declares the new launcher overload.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Baiju Meswani (baijumeswani)]

Thanks for the fix. I tested a Qwen 3.8 DFlash2 model with CUDA graphs enabled: generation now works, and the drafter runs without errors. The inline notes are suggestions to clarify the supported range and make it easier to keep this working; they are not blockers for this model.