[None][chore] Route MHC fused-HC workspace through global Buffers pool

[lancelly]

Summary

Split out the MHC fused-HC workspace pool routing from PR #14053 / #14241 as an independent, opt-in memory optimization.

Activated via the env switch TRTLLM_DSV4_MEM_OPTS=1; default off → behavior is byte-identical to current source.

Change

_FusedHcWorkspaceCache.get (tensorrt_llm/_torch/modules/mhc/mhc_cuda.py) gains an opt-in path that routes the four output buffers (residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur) and the three internal scratch buffers (y_acc_ws, r_acc_ws, done_counter_ws) through get_memory_buffers() so MHC scratch can share backing memory with other layer-transient buffers (e.g., the attention workspace).

The env switch is read once at class definition (cached on _FusedHcWorkspaceCache._USE_POOL) so the lookup doesn't run on every workspace request.

Aliasing invariant

Pool routing only applies for B <= _CACHE_MAX_B (decode-sized shapes). For B > _CACHE_MAX_B (prefill), the code falls through to fresh torch.empty per call, matching the upstream uncached path; this preserves prefill's no-aliasing semantics where layer N's *_prev inputs and layer N+1's *_cur outputs live in distinct tensors.

At decode, upstream's existing per-runner LRU already returns the same tensor across layers, so the kernel is already in-place safe; pool routing inherits that aliasing without widening it.

Files touched

tensorrt_llm/_torch/modules/mhc/mhc_cuda.py | 55 +++++++++++++++++++++++++++++
1 file changed, 55 insertions(+)

Required env

export TRTLLM_DSV4_MEM_OPTS=1

Optional (recommended for parallel-job stability, orthogonal to this PR):

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

[lancelly]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49180 [ run ] triggered by Bot. Commit: 0a8ca52 Link to invocation

[tensorrt-cicd]

PR_Github #49180 [ run ] completed with state SUCCESS. Commit: 0a8ca52
/LLM/main/L0_MergeRequest_PR pipeline #38857 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[lancelly]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49312 [ run ] triggered by Bot. Commit: 0a8ca52 Link to invocation

[tensorrt-cicd]

PR_Github #49312 [ run ] completed with state SUCCESS. Commit: 0a8ca52
/LLM/main/L0_MergeRequest_PR pipeline #38973 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[qiaoxj07]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49396 [ run ] triggered by Bot. Commit: 0a8ca52 Link to invocation

[tensorrt-cicd]

PR_Github #49396 [ run ] completed with state SUCCESS. Commit: 0a8ca52
/LLM/main/L0_MergeRequest_PR pipeline #39046 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[lancelly]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49531 [ run ] triggered by Bot. Commit: dc5666f Link to invocation

[tensorrt-cicd]

PR_Github #49531 [ run ] completed with state SUCCESS. Commit: dc5666f
/LLM/main/L0_MergeRequest_PR pipeline #39162 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49608 [ run ] triggered by Bot. Commit: dc5666f Link to invocation

[tensorrt-cicd]

PR_Github #49608 [ run ] completed with state FAILURE. Commit: dc5666f
/LLM/main/L0_MergeRequest_PR pipeline #39235 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49665 [ run ] triggered by Bot. Commit: dc5666f Link to invocation

[mingyangHao]

Concern: with TRTLLM_DSV4_MEM_OPTS=1, this may break CUDA graph address stability.

The old cache is keyed by (B, ws_ks, m_batches, device). The new pool path uses fixed
buffer names and reserve_buffer=False. If a later larger decode bucket/splitK shape
requests the same name, Buffers.get_buffer() can delete the smaller block and allocate a
new one, invalidating pointers captured by an earlier CUDA graph.

Suggested fix:

1. Include shape/device in the buffer names, e.g. mhc_residual_cur_{device}{B}{n}
   {hidden_size}, mhc_y_acc_ws{device}{ws_ks}{B}{shape_n}, mhc_done_counter_ws{device}
   _{m_batches}.
2. Or keep the existing runner-local LRU for CUDA graph decode buckets, and only use the
   global pool for non-graph eager decode.
3. If using the global pool for graph decode, reserve graph-captured buffers so they
   cannot be deleted/replaced by later larger shapes.

Also add a regression test with TRTLLM_DSV4_MEM_OPTS=1:
capture/replay one decode bucket, allocate/use a larger decode bucket or larger splitK
workspace, then replay the earlier graph and compare with eager.

There is also a device-keying issue: old cache includes device; the new pool path does
not, and memory_buffer_utils allocates on current cuda device. The fix should preserve
per-device separation.

[tensorrt-cicd]

PR_Github #49665 [ run ] completed with state SUCCESS. Commit: dc5666f
/LLM/main/L0_MergeRequest_PR pipeline #39276 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[longlee0622]

We no longer use the ds v4 side branch. Closing this PR. Please prepare PR against main.