[Jiminator]

Motivation

Split #44747 into reviewable parts. This is the attention, memory and DFlash DP part, position 4 of 8 in the merge order below. Like every part it targets main; its branch also carries #45254, so its diff includes it.

Modifications

Adds the shared runtime support Beam relies on:

• TRT-LLM MHA: mixed decode with FP8 KV, the fmha_v2 prefill workspace and page copies, unified page-table bounds.
• FlashInfer ragged prefill for sliding windows and cached prefixes.
• Unified hybrid SWA memory and SWA HiCache.
• DFlash with attention DP.
• FlashInfer MegaMoE all-to-all under attention DP, with an optional decode-only workspace (SGLANG_FLASHINFER_MEGAMOE_DECODE_MAX_TOKENS_PER_RANK).
• PD token-ID logprob streaming, the speculative metrics fallback, and the NCCL graph-buffer registration guard.
• The tests for all of the above.

This part still carries the cherry-picks made on the beam branch:

• [Fix] unified-memory: enabling trtllm_mha with decode graphs and speculative decoding #42930 (7 commits), [Fix] swa: Retain a whole window behind insertion pages #43168 and [Perf] unified-memory: Free SWA segments without a host sync #43170, all merged upstream since. Each was picked before its final revision.
• [unified-memory] Stop allocation eviction on the joint FULL/SWA fit #43169, which was closed in favor of [Fix] unified-memory: Bound joint reclaim and preserve cached prefixes #39294.
• The commits of [Fix] nccl: Disable graph buffer registration whenever a captured graph can hold a collective #43216, still open.

They will be reverted or dropped when the parts are synced with main.

This part changes 44 files (+3,115 / -226) beyond the parts its branch carries. The split commit carries the 18 Co-authored-by trailers of the Beam commits that touched these files.

Validation

• Cut from [Feat] beam: Add native Beam serving support #44747's head (f7473673dd) without modifying the beam branch, which stays the reference. Every branch starts from [Feat] beam: Add native Beam serving support #44747's merge base with main (5bcaf0f8ed, Oct 13); the parts have not yet been synced with current main.
• Conflicts when merged into main 28448a86d4: 11 files including those of [Feat] kernels: Beam model support #45254; this part adds kv_cache_hook.py, unified_hybrid_swa.py, unified_radix_cache.py and four unified-memory tests ([Fix] unified-memory: Bound joint reclaim and preserve cached prefixes #39294, [Fix] unified-memory: enabling trtllm_mha with decode graphs and speculative decoding #42930, [Perf] unified-memory: Free SWA segments without a host sync #43170, [Feat] unified-memory: Fuse drafts with sliding-window and recurrent-state layers #43326, [Perf] unified-memory: Reduce allocation and SWA retirement overhead #44970, [Perf] attention: Capture translated FA4 verify metadata for DSPARK #44973).
• Coverage: [Chore] docker: Add Beam CUDA image (1/8) #45253, [Perf] beam: Add optional cuTile BF16x3 router GEMM #45361 and [Feat] parser: Add Beam reasoning, tool calls and response templates (5/7) #45257 merged, plus the one-line [Chore] deps: Add nvidia-cuda-tileiras for the optional cuTile router #45290 and [Chore] deps: Bump the OpenAI SDK to 2.8.0 for Responses shell tools (2/8) #45291 changes, give exactly the [Feat] beam: Add native Beam serving support #44747 tree, except that streamed LoRA (25 files from sglang-miles) and Beam LoRA serving (10 files plus the LoRA hunks in 3 files) are left out, the five platform pyprojects keep OpenAI 2.6.1, and docker/beam/beam.Dockerfile is added. No other path differs.
• No part except [Chore] deps: Add nvidia-cuda-tileiras for the optional cuTile router #45290 and [Chore] deps: Bump the OpenAI SDK to 2.8.0 for Responses shell tools (2/8) #45291 changes a pyproject, and none changes docker/Dockerfile. environ.py is split by hunk. Every changed Python file in this part parses.
• Not run for this split: pre-commit, registered-test registry parsing and the CPU and GPU suites. Intermediate parts make no claim of standalone Beam serving validation; the recorded accuracy and performance results are in [Feat] beam: Add native Beam serving support #44747.

Merge order

Every part targets main. #44747 stays open, unchanged, as the full reference. Merge in this order:

1. [Chore] docker: Add Beam CUDA image (1/8) #45253 — Beam CUDA Docker image
2. [Chore] deps: Bump the OpenAI SDK to 2.8.0 for Responses shell tools (2/8) #45291 — OpenAI SDK 2.8.0, needed by the Responses shell tools in [Feat] parser: Add Beam reasoning, tool calls and response templates (5/7) #45257
3. [Feat] kernels: Beam model support #45254 — Kernels
4. [Feat] srt: Add attention, memory and DFlash DP support for Beam (4/8) #45255 — Attention, memory and DFlash DP support
5. [Feat] beam: Add Beam model, runtime support and checkpoints (4/7) #45256 — Beam model, checkpoints and tests
6. [Feat] parser: Add Beam reasoning, tool calls and response templates (5/7) #45257 — Beam reasoning, tool calls and response templates (also after [Feat] parser: Enrich chat_parsing streaming events #40478, [Fix] parser: Accept whole-number floats for integer tool arguments #43941 and [Feat] parser: Add checkpoint-driven response-template adapters #40479)
7. [Chore] deps: Add nvidia-cuda-tileiras for the optional cuTile router #45290 — nvidia-cuda-tileiras, needed by [Perf] beam: Add optional cuTile BF16x3 router GEMM #45361
8. [Perf] beam: Add optional cuTile BF16x3 router GEMM #45361 — cuTile BF16x3 router GEMM (optional: without it Beam uses the cuBLAS router)

#45292 (Dynamo and grpcio-tools pins for the Beam image) is not in the merge order for now.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): 🚫 Run #39646405301
Latest PR Test (Extra): 🚫 Run #39646405451
Latest PR Test (AMD ROCm 10): ❌ Run #39646405429

[Jiminator]

Folded into #45256, which now carries both the shared runtime support and the Beam model (4/7 in the merge order), as Muse Spark's runtime and model landed together. Closing this one.