feat(sglang): upgrade to v0.5.9 on plain torch base for SLIME support

SLIME's post-training pipeline (Megatron + SGLang) requires SGLang v0.5.9.
This upgrades from v0.4.x and rebases the image onto a plain torch base,
dropping the torch-extras layer (DeepSpeed, Apex, xFormers) that neither
SGLang nor SLIME actually uses.

FlashInfer moves from JIT to v0.6.3 AOT compilation via TVM. sgl-kernel
is now built with scikit-build-core and enables SM100A (Blackwell) and
FP4 support. vLLM and Triton are removed from this image since they are
served by the dedicated vllm-tensorizer image.

Author: abatilo

.github/workflows/sglang.yml

@@ -24,7 +24,8 @@ jobs:
     with:
       image-name: sglang
       folder: sglang
-      tag-suffix: ${{ inputs.tag || '386fabe-nccl-cuda12.8.0-ubuntu22.04-nccl2.25.1-1-torch2.6.0-vision0.21.0-audio2.6.0-abi1' }}
+      tag-suffix: ${{ inputs.tag || 'nccl-cuda12.9.1-ubuntu22.04-nccl2.29.2-1-torch2.10.0-vision0.25.0-audio2.10.0-abi1' }}
       build-args: |
-        BASE_IMAGE=${{ inputs.base-image || 'ghcr.io/coreweave/ml-containers/torch-extras:es-actions-386fabe-nccl-cuda12.8.0-ubuntu22.04-nccl2.25.1-1-torch2.6.0-vision0.21.0-audio2.6.0-abi1'}}
-        ${{ inputs.base-image && 'BASE_IMAGE=' }}${{ inputs.base-image}}
+        BASE_IMAGE=${{ inputs.base-image || 'ghcr.io/coreweave/ml-containers/torch:17ad6db-nccl-cuda12.9.1-ubuntu22.04-nccl2.29.2-1-torch2.10.0-vision0.25.0-audio2.10.0-abi1'}}
+        ${{ inputs.builder-image && format('BUILDER_IMAGE={0}', inputs.builder-image) || '' }}
+

docs/sglang-review-research/01-cuda-architectures-and-flashinfer-aot.md

@@ -0,0 +1,303 @@
+# CUDA Architectures and FlashInfer AOT Compilation
+
+## The Problem
+
+Our sglang build has a line that strips `12.0+PTX` from FlashInfer's
+architecture list, with a comment claiming "nvcc in CUDA 12.x doesn't support
+compute_120." Eta0's review pointed out this is wrong — CUDA 12.8+ fully
+supports compute_120 — and suggested we also enable FlashInfer's communication
+module by installing nvshmem.
+
+To understand why this matters, we need to understand how CUDA code targets
+different GPUs, what FlashInfer's AOT build does, and why multi-GPU
+communication kernels are valuable.
+
+## GPU Compute Capabilities
+
+Every NVIDIA GPU has a **compute capability** (CC) — a version number that
+identifies which instructions it supports. Think of it like a CPU's instruction
+set (x86, ARM, etc.) but for GPUs.
+
+| Compute Capability | Architecture | Example GPUs |
+|---|---|---|
+| 8.0 | Ampere | A100 |
+| 8.6 | Ampere | A40, RTX 3090 |
+| 8.9 | Ada Lovelace | L40, RTX 4090 |
+| 9.0 | Hopper | H100, H200 |
+| 10.0 | Blackwell | B200, GB200 |
+| 12.0 | Blackwell | DGX Spark, RTX PRO 6000 |
+
+Notice that 10.0 and 12.0 are both "Blackwell" but are **different
+sub-architectures** with different instruction sets. B200 (the big data center
+GPU) is CC 10.0. DGX Spark and RTX PRO 6000 are CC 12.0. They share the
+Blackwell marketing name but have different silicon and different capabilities.
+
+### The `a` Suffix
+
+Some capabilities have an `a` variant: `9.0a`, `10.0a`, `12.0a`. The `a` means
+"architecture-specific features" — extra instructions that only exist on that
+exact GPU generation and have no equivalent in the portable PTX instruction set.
+
+For example, `sm_100a` enables MXFP8 (e2m1) matrix operations that only
+Blackwell B200 hardware can execute. Code compiled for `sm_100` (without `a`)
+would miss these instructions. For high-performance inference, the `a` suffix
+matters.
+
+## SASS vs PTX: Two Ways to Compile CUDA
+
+When you compile CUDA code, nvcc can produce two kinds of output:
+
+**SASS (Shader Assembly)** — native machine code for a specific GPU. Fast,
+optimal, but only runs on that exact compute capability.
+
+**PTX (Parallel Thread Execution)** — a portable intermediate representation.
+Runs on any GPU with equal or higher compute capability, but must be
+JIT-compiled to SASS at runtime, which is slower.
+
+```
+Source Code (.cu)
+     |
+     v
+  nvcc compiler
+     |
+     +---> SASS for sm_90a   (native, fast, H100 only)
+     +---> SASS for sm_100a  (native, fast, B200 only)
+     +---> PTX for compute_90 (portable, JIT at runtime)
+```
+
+### The +PTX Notation
+
+In PyTorch's `TORCH_CUDA_ARCH_LIST`, `12.0+PTX` means "compile native SASS for
+sm_120 AND include PTX for compute_120 as a fallback." This is PyTorch-specific
+syntax. The `+PTX` part tells PyTorch's build system to generate an extra
+`-gencode` flag:
+
+```
+12.0+PTX  -->  -gencode arch=compute_120,code=sm_120
+               -gencode arch=compute_120,code=compute_120  (PTX fallback)
+```
+
+Without `+PTX`, you only get the native SASS:
+
+```
+12.0      -->  -gencode arch=compute_120,code=sm_120
+```
+
+This notation is a **PyTorch convention**, not a CUDA standard. Other build
+systems (like FlashInfer's) don't understand it.
+
+## How FlashInfer's AOT Build Works
+
+FlashInfer uses **Ahead-of-Time (AOT) compilation** to pre-compile attention
+kernels for each target GPU architecture. This avoids the startup latency of
+JIT compilation (which can take minutes on first use).
+
+The AOT build is invoked via:
+
+```bash
+FLASHINFER_CUDA_ARCH_LIST="9.0a 10.0a" python3 -m flashinfer.aot
+```
+
+### FlashInfer's Arch Parser
+
+FlashInfer parses `FLASHINFER_CUDA_ARCH_LIST` with this logic:
+
+```python
+for arch in os.environ["FLASHINFER_CUDA_ARCH_LIST"].split(" "):
+    major, minor = arch.split(".")
+    major = int(major)
+    self.TARGET_CUDA_ARCHS.add((int(major), str(minor)))
+```
+
+It splits on spaces, then splits each value on `.` to get a `(major, minor)`
+tuple. The minor part is kept as a string to support suffixes like `"0a"`.
+
+Then it generates nvcc flags:
+
+```python
+f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}"
+```
+
+For `"9.0a"`, this produces:
+```
+-gencode=arch=compute_90a,code=sm_90a
+```
+
+### Why 12.0+PTX Breaks
+
+If you pass `"12.0+PTX"` to FlashInfer:
+
+1. `arch.split(".")` gives `["12", "0+PTX"]`
+2. `minor = "0+PTX"`
+3. The gencode flag becomes: `-gencode=arch=compute_120+PTX,code=sm_120+PTX`
+4. nvcc doesn't understand `compute_120+PTX` and crashes
+
+This is why our build.bash filters it out. **The filtering itself is correct
+and necessary.** The comment explaining it was wrong — it blamed nvcc for not
+supporting compute_120, when the real issue is FlashInfer's parser not
+understanding the `+PTX` suffix.
+
+### The Right Fix: Convert, Don't Strip
+
+Instead of stripping `12.0+PTX` entirely (losing all sm_120 support), we should
+**convert** it to `12.0a` like the vllm-tensorizer image does:
+
+```bash
+# Current (strips 12.0+PTX, no sm_120 kernels):
+FLASHINFER_ARCH_LIST="$(echo "${TORCH_CUDA_ARCH_LIST}" | sed 's/12\.0+PTX//')"
+
+# Better (converts to 12.0a, native sm_120a kernels):
+FLASHINFER_ARCH_LIST="$(echo "${TORCH_CUDA_ARCH_LIST}" \
+  | sed -E 's@\b(9|10|12)\.0\b@\1\.0a@g; s@\+PTX\b@@g')"
+```
+
+The vllm-tensorizer approach handles all architectures uniformly:
+1. Convert bare `.0` to `.0a` (adds architecture-specific features)
+2. Strip `+PTX` globally (FlashInfer only generates native SASS, no PTX)
+
+## What sm_120 Kernels We're Missing
+
+Without sm_120a in FlashInfer's arch list, these kernels are not compiled:
+
+| Kernel | Purpose | Why it Matters |
+|---|---|---|
+| `mla_sm120.cu` | Multi-Latent Attention using CGA and TMA | Dedicated attention kernel for CC 12.0 GPUs |
+| `fp4_quantization_sm120` | FP4 quantization | Native FP4 support on Blackwell |
+| `cutlass_fused_moe_sm120` | Fused MoE dispatch | Efficient expert routing |
+| `gemm_sm120_cutlass_fp4` | FP4 GEMM | Low-precision matrix multiply |
+| `xqa_mla` | XQA MLA backend | sm_120-exclusive MLA implementation |
+
+The MLA kernel (`mla_sm120.cu`) uses **CGA (Cooperative Group Arrays)** and
+**TMA (Tensor Memory Accelerator)** — hardware features exclusive to sm_120a
+that have no PTX equivalent. Without native compilation, these kernels simply
+don't exist. There's no PTX fallback — FlashInfer would fall back to JIT
+compilation from source at runtime (minutes of startup) or use a generic
+implementation that doesn't leverage the hardware.
+
+### sgl-kernel Already Has sm_120a
+
+Importantly, **sgl-kernel handles this correctly already.** It ignores
+`TORCH_CUDA_ARCH_LIST` entirely and manages its own gencode flags through CMake
+options. When CUDA >= 12.8, sgl-kernel automatically adds:
+
+```
+-gencode=arch=compute_120a,code=sm_120a
+```
+
+So the gap is FlashInfer-specific.
+
+## The Communication Module and nvshmem
+
+### What FlashInfer's Comm Module Provides
+
+FlashInfer includes a communication module (`flashinfer.comm`) that provides
+optimized multi-GPU collective operations:
+
+| Feature | What It Does |
+|---|---|
+| TRT-LLM allreduce fusion | Fused allreduce + residual add + RMSNorm in one kernel |
+| vLLM custom allreduce | CUDA IPC-based allreduce for vLLM |
+| NVSHMEM operations | GPU-initiated allreduce, alltoall, barrier |
+| MoE alltoall | Expert-parallel dispatch/combine for MoE models |
+| Multi-Node NVLink | MNNVL allreduce for multi-node setups |
+
+The key innovation is **operation fusion**. Normally, tensor-parallel inference
+does three separate kernel launches per transformer layer:
+
+```
+1. allreduce (collect results from all GPUs)
+2. residual add (add skip connection)
+3. RMSNorm (normalize)
+```
+
+FlashInfer's comm module fuses these into a single kernel:
+
+```
+1. allreduce + residual + RMSNorm (one kernel, one memory pass)
+```
+
+This reduces kernel launch overhead and memory bandwidth usage significantly on
+H100+ GPUs.
+
+### sglang Uses This
+
+sglang v0.5.9 imports `flashinfer.comm` in
+`python/sglang/srt/layers/flashinfer_comm_fusion.py`:
+
+```python
+import flashinfer.comm as comm
+
+# Creates CUDA IPC workspace for fused allreduce
+workspace = comm.trtllm_create_ipc_workspace_for_all_reduce_fusion()
+
+# Performs fused allreduce + residual + RMSNorm
+comm.trtllm_allreduce_fusion(
+    pattern=AllReduceFusionPattern.kARResidualRMSNorm,
+    ...
+)
+```
+
+This is auto-enabled for DeepSeek, GLM, Qwen, and other MoE models when
+running with tensor parallelism (TP > 1) on SM90+ GPUs.
+
+**Our `--add-comm false` silently disables this optimization.** sglang falls
+back to standard NCCL allreduce (three separate kernels instead of one fused
+kernel).
+
+### Why We Had `--add-comm false`
+
+The comm module requires `nvidia-nvshmem-cu12` as a build dependency. When we
+wrote the build script, the base image didn't have nvshmem and we didn't want
+to add it. So we disabled the entire comm module to avoid a build failure.
+
+The vllm-tensorizer image already installs nvshmem in its builder stage:
+
+```dockerfile
+python3 -m pip install --no-cache-dir \
+  "nvidia-nvshmem-cu${CUDA_VERSION%%.*}"
+```
+
+This adds ~145MB to the builder stage download but **does not affect the final
+image size** because it's a multi-stage build — the builder stage is discarded
+after wheel compilation.
+
+### What SLIME Needs
+
+SLIME itself doesn't directly call `flashinfer.comm`. It inherits whatever
+sglang provides. SLIME's RL coordination uses `torch.distributed` (NCCL) for
+weight sync and process group management. However, when SLIME drives sglang for
+rollout inference with TP > 1, the fused allreduce would automatically
+activate — so enabling it benefits SLIME's inference performance.
+
+SLIME does use nvshmem separately for **DeepEP** (expert-parallel MoE
+dispatch), but that's a different nvshmem dependency path, not related to
+FlashInfer's comm module.
+
+## Summary of Required Changes
+
+1. **Fix the comment** — CUDA 12.8+ supports compute_120. The issue is
+   FlashInfer's parser, not nvcc.
+
+2. **Convert `12.0+PTX` to `12.0a`** instead of stripping it — gives FlashInfer
+   native sm_120a Blackwell kernels, matching vllm-tensorizer's approach.
+
+3. **Install `nvidia-nvshmem-cu12`** in the builder stage — build-time only
+   dependency, ~145MB, discarded in final image.
+
+4. **Remove `--add-comm false`** — enables the fused allreduce+RMSNorm
+   optimization for TP deployments on H100+ GPUs.
+
+## Key Concepts to Remember
+
+- **Compute capability** identifies a GPU's instruction set (8.0 = A100, 10.0 =
+  B200, 12.0 = DGX Spark)
+- **SASS** = native GPU machine code, **PTX** = portable intermediate code
+- **`+PTX`** is PyTorch syntax for "include PTX fallback" — not understood by
+  all build systems
+- **AOT** = pre-compile kernels ahead of time (fast startup, no JIT latency)
+- **The `a` suffix** enables architecture-exclusive instructions (no PTX
+  equivalent)
+- **Fused operations** combine multiple kernels into one, reducing memory
+  bandwidth and launch overhead
+- **Multi-stage Docker builds** mean builder dependencies don't bloat the final
+  image