Move inference context bookkeeping to CPU with ContextGPUView - #4306

Merged

lmcafee-nvidia merged 33 commits into

NVIDIA:mainfrom

lmcafee-nvidia:context-cpu

May 3, 2026

Commits on Apr 16, 2026

Add NVTX ranges for inference loop stages (baseline commit)

lmcafee-nvidia
and
claude
committed
419ddbf
Move context bookkeeping tensors to CPU with ContextGPUView interface

lmcafee-nvidia
and
claude
committed
3425d29
Defer Mamba GPU ops to transfer step, add transfer_samples_to_cpu()

lmcafee-nvidia
and
claude
committed
4efda46
Move deferred Mamba GPU ops to top of initialize_attention_state()

lmcafee-nvidia
and
claude
committed
3dbc4d4
Move Mamba metadata computation to CPU, H2D copy via load_from_cpu()

lmcafee-nvidia
and
claude
committed
0271285
Move MambaSlotAllocator intermediate tensors to CPU

lmcafee-nvidia
and
claude
committed
0c755bb

Commits on Apr 17, 2026

Update run_dynamic_text_generation_server.py for PR #4225 args init split

root
authored and
root
committed
2e5b588

Commits on Apr 21, 2026

Remove redundant fill_(0) calls in MHAMetadata.reset()

root
authored and
root
committed
a5a1590
Replace per-step NCCL AllReduce with ZMQ in EP batch-dim sync

root
authored and
root
committed
7e715cd
Coalesce transfer_bookkeeping_to_gpu into one H2D copy

root
authored and
root
committed
e15a12a
Coalesce MHA metadata H2D into unified bookkeeping buffer

authored and
root
committed
eab349e
Coalesce Mamba metadata H2D into unified bookkeeping buffer

authored and
root
committed
2f295c6

Commits on Apr 22, 2026

Gate logging-only work behind will_log_this_step in async_forward

authored and
root
committed
149b4f5
Return CPU-side sampled_tokens in step_result to avoid D2H sync

authored and
root
committed
39ffe4d

Commits on Apr 27, 2026

Enable NVTX range helper in inference server entry

authored and
root
committed
811925d

Commits on Apr 29, 2026

Merge remote-tracking branch 'main/main' into context-cpu
lmcafee-nvidia
committed
300f799

Commits on Apr 30, 2026

Move _rewind_kv_cache bookkeeping rewind to CPU

lmcafee-nvidia
and
claude
committed
5d5cb60
Fix _execute_pending_mamba_ops docstring callsite reference

lmcafee-nvidia
and
claude
committed
17d3238
Merge remote-tracking branch 'main/main' into context-cpu
lmcafee-nvidia
committed
f54c52c

Commits on May 1, 2026

Bind MHA state_data and run H2D inside initialize_attention_state

lmcafee-nvidia
and
claude
committed
674955a
Merge remote-tracking branch 'main/main' into context-cpu
lmcafee-nvidia
committed
423d8a7
Restore smallest_non_decode_cuda_graph_size attribute

lmcafee-nvidia
and
claude
committed
34834c0
Clarify coalesced bookkeeping comments
lmcafee-nvidia
committed
578e1f2
Fix dynamic inference CUDA graph replay
lmcafee-nvidia
committed
4bcb61c
Fix docs and install CI failures
lmcafee-nvidia
committed
74d63b5
Fix PyTorch DeviceMesh docs link
lmcafee-nvidia
committed
477d780

Commits on May 2, 2026

Commits on May 3, 2026

Stabilize MTP CUDA graph eager comparison
lmcafee-nvidia
committed
8c91a68
Merge remote-tracking branch 'main/main' into context-cpu
lmcafee-nvidia
committed
1dda0ac
Fix rewind cache test device expectations
lmcafee-nvidia
committed
2bb6d5e