[Jiminator]

Motivation

Beam's router multiplies BF16 hidden states by an FP32 gate in each of its 49 MoE layers. #45256 runs it with the tiny FP32 GEMM, the GEMV and the cuBLAS BF16x3 GEMM. This part adds the cuTile BF16x3 router GEMM from #44747 and makes it Beam's default router on Hopper and SM10x Blackwell.

It's optional: without it, Beam serves on the cuBLAS router. It needs the tileiras compiler, which is #45290. Position 7 of 7 in the merge order below. It is based on #45290's branch (beam-stack-deps-tileiras), so its diff shows only this part.

Modifications

• kernels/ops/gemm/router_gemm_cutile.py: BF16 activations × FP32 router weights → FP32 logits. It uses a three-part BF16 split of the weight prepared at load time, and combines K tiles with explicit FP32 rounding. Small batches use split-K plus a reduction.
• models/beam.py:
  • BeamRouterLinearMethod refreshes the split weight in place after weight loading, so captured graphs see updated weights.
  • Above the tiny FP32 GEMM's 16 tokens, the router uses cuTile.
  • FP8 and NVFP4 checkpoints also keep the tiny FP32 GEMM when cuTile is on.
• Shared splitter ownership: Beam and the cuBLAS wrapper import split_fp32_to_bf16x3 from ops.quantization.bf16x3_split, registered as quantization.split_fp32_to_bf16x3. The old inline splitter, GEMM inventory entry and shared split scratch are removed. This branch carries the canonical module from [Feat] kernels: Beam model support #45254 until its later restack; the existing out=parts use is unchanged.
• SGLANG_BEAM_ROUTER_BACKEND chooses cutile (default), cublas or torch.
• The kernel is registered in the kernel inventory as gemm.router_gemm_bf16x3, with a new cutile KernelBackend provenance value (and the matching forward_cutile stub in BaseFusedOp), for SM90 through SM10x.

This part adds no test code. The recorded GB300 serving runs below belong to earlier revisions.

This part changes 9 files (+340 / -64) relative to its #45290 base branch.

Validation

The splitter migration in 4a8c0cde40 has not been reverified, as requested. No tests, lint, pre-commit, import/compilation smoke, GPU checks or benchmarks were run for that update. CuTile arithmetic and model dispatch are unchanged; this was a targeted ownership migration, not a full stack restack.

Historical results

The results below come from earlier revisions, before 4a8c0cde40. They do not establish verification of the latest update.

• Synced with main: the branches fast-forward from their previous heads and merge into main b57b5e2286 without conflicts ([Feat] parser: Add Beam reasoning, tool calls and response templates (5/7) #45257 merges that main, which includes [Feat] parser: Enrich chat_parsing streaming events #40478; the others are on 6efdd33d07).
• No upstream code is copied into the stack any more. The cherry-picks of [Fix] unified-memory: enabling trtllm_mha with decode graphs and speculative decoding #42930, [Fix] swa: Retain a whole window behind insertion pages #43168 and [Perf] unified-memory: Free SWA segments without a host sync #43170 (merged), [unified-memory] Stop allocation eviction on the joint FULL/SWA fit #43169 (closed) and [Fix] nccl: Disable graph buffer registration whenever a captured graph can hold a collective #43216 (open) are reverted, and [Feat] parser: Add checkpoint-driven response-template adapters #40479 and [Fix] parser: Accept whole-number floats for integer tool arguments #43941 come in as merges of those PRs in [Feat] parser: Add Beam reasoning, tool calls and response templates (5/7) #45257 ([Feat] parser: Enrich chat_parsing streaming events #40478 has since merged into main).
• On 4×GB300 with the public reflection/Beam-1.0-501B-A23B-NVFP4 checkpoint and the cookbook GB300 NVFP4 low-latency recipe, serving main 6efdd33d07 + [Fix] nccl: Disable graph buffer registration whenever a captured graph can hold a collective #43216, [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 + this stack:
  • With [Chore] deps: Add nvidia-cuda-tileiras for the optional cuTile router #45290 and [Perf] beam: Add optional cuTile BF16x3 router GEMM #45361: GSM8K 96.44–96.59% over three server runs (sgl-eval 0.1.3, 1,319 questions, temperature 0), with zero request errors. Chat reasoning, streaming, a chat tool call, Responses text, a Responses function call, a full Responses shell round trip and /tokenize all pass. A profiler trace shows the cuTile router kernel running.
  • Without [Chore] deps: Add nvidia-cuda-tileiras for the optional cuTile router #45290 and [Perf] beam: Add optional cuTile BF16x3 router GEMM #45361: GSM8K 96.51% in both of two server runs, with zero request errors, and the same checks pass.
• The Beam E2E test in [Feat] parser: Add Beam reasoning, tool calls and response templates (5/7) #45257 passes on two 4×GB300 nodes (prefill TP4 on one, decode TP4 on the other, Mooncake over multi-node NVLink, in place of the test's single-node 8-GPU layout): GSM8K 96.51–96.89% at temperature 0.7 over three runs, decode accept length 4.53–4.55, and single-request accept length 4.92 at about 750 tok/s.
• Main's existing GPU tests for the code the stack changes pass unedited on GB300: TRT-LLM MHA (dense, graph metadata, encoder-only, stride parity), SWA FlashInfer, tiny GEMM, collectives, Triton MoE and the DFlash buffer and kernel tests. pre-commit passes.
• On GB300 the registered kernel matches an FP64 reference at Beam's router shape with a mean absolute error of 1.6–2.7e-7 for 1 to 2,048 tokens over two runs, and importing sglang.kernels.ops does not import cuda.tile.
• Earlier measurements from [Feat] beam: Add native Beam serving support #44747, unchanged by this sync:
  • Router latency per call at Beam's shape (N=256, K=5120, CUDA graphs), against the cuBLAS path: the same up to 16 tokens; from 32 to 2,048 tokens, 2–5 µs faster on GB300 and 4–9 µs faster on H200.
  • Mean absolute error against FP64: 1.6–2.8e-7 for cuTile, against 1.9–6.2e-6 for cuBLAS BF16x3.
• tileiras must share its major.minor with nvidia-cuda-nvcc and nvidia-nvvm. cuda-tile rejects mismatched wheels, and the router then fails at startup.

Merge order

Each part is based on the one before it, so its diff shows only its own changes. Merge bottom-up; after each part merges, the next part is retargeted to main and main is merged into it. #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 (merged)
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 (merged)
3. [Feat] kernels: Beam model support #45254 — Kernels
4. [Feat] beam: Add Beam model, runtime support and checkpoints (4/7) #45256 — Attention, memory and DFlash DP support, and the Beam model ([Feat] srt: Add attention, memory and DFlash DP support for Beam (4/8) #45255 folded in; also after [Fix] nccl: Disable graph buffer registration whenever a captured graph can hold a collective #43216)
5. [Feat] parser: Add Beam reasoning, tool calls and response templates (5/7) #45257 — Beam reasoning, tool calls, response templates and the Beam E2E test (merged into [Feat] beam: Add Beam model, runtime support and checkpoints (4/7) #45256)
6. [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
7. [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.

---

CI States

Latest PR Test (Base): ❌ Run #40183433366
Latest PR Test (Extra): ❌ Run #40183433293
Latest PR Test (AMD ROCm 10): ❌ Run #40183433398

[mintlify]

Preview deployment for your docs. Learn more about Mintlify Previews.

Project	Status	Preview	Updated
sglang-doc	🟢 Ready	View Preview	Oct 21, 2026, 7:56 PM

💡 Tip: Enable Automations to automatically generate PRs for you.