fix: GPU tensor allocation fall back to host memory

[mattwittwer]

What does the PR do?

Fixes silent output corruption when a Python model returns a GPU tensor and Triton's CUDA memory pool is exhausted. In that case the response allocator falls back to the pinned host pool and hands the backend a CPU_PINNED buffer; the Python backend only handled MEMORY_CPU on that path, so the tensor never reached the Triton-provided buffer and the response went out with whatever the recycled pinned block last held (zeros or another request's data), with no error. Pairs with the QA-side change in server#8992.

Root cause

• PbMemory::CopyBuffer has no GPU -> CPU_PINNED case and fell through to cudaMemcpyHostToDevice. On older CUDA this produced the failed to copy data: invalid argument reported in server#7148; on CUDA 13 the copy silently succeeds into the shared-memory intermediate.
• PbMemory::Create left the parent-side data pointer null for a CPU_PINNED intermediate.
• Both copy-back loops in python_be.cc (ResponseSendTBD and ProcessRequests) matched neither the CPU nor the GPU branch for CPU_PINNED, so the staged tensor was never copied into the Triton buffer. This is the live cause on current stacks.

--pinned-memory-pool-byte-size=0 masked the bug because the fallback then lands on plain MEMORY_CPU, which was handled.

Fix

• Treat CPU_PINNED as host memory when choosing the cudaMemcpyKind in PbMemory::CopyBuffer (src/pb_memory.cc).
• Point the intermediate's data pointer into shared memory for every non-GPU type in PbMemory::Create (src/pb_memory.cc).
• Copy CPU_PINNED intermediates back to the Triton buffer in both copy-back loops (src/python_be.cc).

No behaviour change on the GPU or plain-CPU paths.

Checklist

• PR title reflects the change and is of format <commit_type>: <Title>
• Changes are described in the pull request.
• Related issues are referenced.
• Populated github labels field
• Added test plan and verified test passes.
• Verified that the PR passes existing CI.
• Verified copyright is correct on all changed files.
• Added succinct git squash message before merging.
• All template sections are filled out.
• Optional: Additional screenshots for behavior/output changes with before/after.

Commit Type:

• fix

Related PRs:

Where should the reviewer start?

Test plan:

• CI Pipeline ID: [71702518]

Caveats:

Background

Related Issues:

• Fixes client silent failure - E0422 05:03:24.145960 1 pb_stub.cc:402] An error occurred while trying to load GPU buffers in the Python backend stub: failed to copy data: invalid argument server#7148

[greptile-apps]

Confidence Score: 5/5

[High risk] Changes GPU memory allocation and tensor copy logic.

The PR appears safe to merge, though the pinned-host output paths still lack regression tests.

Findings

1.  Pinned fallback lacks tests ▶

Summary

The PR handles CPU-pinned memory as host memory when copying tensor data and delivering ordinary or decoupled responses. Since the previous review, the only change updates the copyright year in src/pb_memory.cc.

_{Reviews (2) · Last reviewed commit: "Update copyright"}