Add RTX 5090 (SM 12.0, 32GB) serve profile: tuned launcher, x86 CI image, compose, README section

Author: am423

.github/workflows/build-rtx5090-image.yml

@@ -0,0 +1,43 @@
+name: build-rtx5090-image
+
+on:
+  push:
+    branches: [main]
+    paths:
+      - "container/Dockerfile.rtx5090"
+      - "serve-profiles/rtx5090/**"
+      - ".github/workflows/build-rtx5090-image.yml"
+  workflow_dispatch:
+
+env:
+  REGISTRY: ghcr.io
+  IMAGE: ghcr.io/${{ github.repository_owner }}/ornith-15-35b-nvfp4-w4a16-sm121-sglang
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: "Set up Docker Buildx"
+        uses: docker/setup-buildx-action@v3
+
+      - name: "Log in to GHCR"
+        uses: docker/login-action@v3
+        with:
+          registry: ${{ env.REGISTRY }}
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: "Build and push tag rtx5090"
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: container/Dockerfile.rtx5090
+          push: true
+          tags: ${{ env.IMAGE }}:rtx5090
+          cache-from: type=gha
+          cache-to: type=gha,mode=max

README.md

@@ -99,11 +99,53 @@ If you omit `MODEL_PATH`, the entrypoint defaults to pulling
 `r0b0tlab/Ornith-1.5-35B-A3B-NVFP4-W4A16` from Hugging Face at boot
 (`HF_HOME` is a named volume, so the download persists across restarts).
 
-The image bakes the **pre-warmed FlashInfer JIT caches** (CUTLASS FP4 GEMM
+The GB10 image bakes the **pre-warmed FlashInfer JIT caches** (CUTLASS FP4 GEMM
 sm121 modules + autotune tables + Triton cache) from the validated runs —
 first boot does zero JIT compilation. On a cold system, first-use JIT with
 the model resident can OOM (see [Known runtime notes](#known-runtime-notes)).
 
+### RTX 5090 profile (Blackwell SM 12.0, 32 GB, x86-64)
+
+A second, x86 image and tuned profile target a single RTX 5090:
+
+```bash
+hf download r0b0tlab/Ornith-1.5-35B-A3B-NVFP4-W4A16 \
+  --local-dir ./models/ornith-15-35b-a3b-nvfp4-w4a16-B
+
+docker compose -f docker-compose.rtx5090.yml up -d
+# or bare:
+docker run --rm --gpus all --ipc host -p 8000:8000 \
+  -v $PWD/models/ornith-15-35b-a3b-nvfp4-w4a16-B:/models/ckpt:ro \
+  -e MODEL_PATH=/models/ckpt \
+  ghcr.io/r0b0tlab/ornith-15-35b-nvfp4-w4a16-sm121-sglang:rtx5090
+```
+
+Differences vs the GB10 profile ([`serve-profiles/rtx5090/serve.sh`](serve-profiles/rtx5090/serve.sh)):
+
+| Knob | GB10 (121 GB unified) | RTX 5090 (32 GB) | Why |
+|---|---|---|---|
+| `--mem-fraction-static` | 0.80 | **0.92** | dedicated 32 GB part; OS doesn't share VRAM |
+| `--max-running-requests` | default (48/105) | **4** | mamba state cache is the binding constraint at 32 GB |
+| `--max-mamba-cache-size` | default (528) | **24** | ~1.4 GB state cache → 4 × 32K contexts fit |
+| context / KV | 32768 × 105 reqs | 32768 × 4 reqs | ~2.5 GB FP8 KV at this occupancy |
+| MTP (EAGLE K=1) | on | on; `MTP=0` env to disable | draft costs ~3.6 GB; disabling reclaims it for KV |
+
+Memory budget ≈ 21.1 GB weights + 3.6 GB draft + 1.1 GB graphs + 1.4 GB
+mamba + 2.5 GB KV ≈ 29.7 GB. First boot on a cold cache JIT-compiles the
+sm_120 CUTLASS FP4 kernels (~10–20 min; runs in host RAM on a discrete-GPU
+box — safe, unlike unified-memory parts). `docker-compose.rtx5090.yml`
+mounts a persistent volume for the caches so subsequent boots are warm.
+If capture OOMs, lower `MAX_RUNNING_REQUESTS`/`CONTEXT_LENGTH` first.
+
+The `rtx5090` tag is built by CI ([workflow](.github/workflows/build-rtx5090-image.yml))
+from [container/Dockerfile.rtx5090](container/Dockerfile.rtx5090), which
+pip-installs the same pinned stack (sglang `5a7b26c63`, torch 2.13.0+cu130,
+triton 3.7.1, flashinfer-python 0.6.17 + cubins) on x86 Ubuntu 24.04.
+**Note:** validated end-to-end on GB10/SM121; the 5090 profile is sized from
+the measured memory ledger and the vendor's SM 12.0 support envelope, but was
+not run on physical 5090 hardware before publication — report issues if
+first-boot JIT behaves differently on sm_120 discrete parts.
+
 ### Serving without the container
 
 SGLang `0.5.6.post3.dev9218+g5a7b26c63`, FlashInfer 0.6.17,
@@ -203,22 +245,26 @@ It scored GSM8K 70.00% vs this recipe's 76.25% and decoded at 38.6 tok/s vs
 
 ```
 ├── README.md
-├── docker-compose.yml        # one-command launch
-├── container/                # Dockerfile, serve.sh entrypoint, build script
+├── docker-compose.yml            # GB10/SM121 one-command launch
+├── docker-compose.rtx5090.yml    # RTX 5090 (SM 12.0, 32 GB) launch
+├── container/                    # GB10 Dockerfile + serve.sh + build script
+│   └── Dockerfile.rtx5090        # x86 CI-built image (same pinned stack)
+├── serve-profiles/
+│   └── rtx5090/serve.sh          # tuned launcher for a single 32 GB card
 ├── eval/
-│   ├── run_quality_set.py    # run the 200-question suite (resumable)
-│   ├── rescore.py            # fixed-scorer summaries from raw rows
-│   ├── answer_extract.py     # GSM8K answer extraction (the fixed scorer)
-│   └── data/                 # quality-200.jsonl + raw rows for 4 runs
+│   ├── run_quality_set.py        # run the 200-question suite (resumable)
+│   ├── rescore.py                # fixed-scorer summaries from raw rows
+│   ├── answer_extract.py         # GSM8K answer extraction (the fixed scorer)
+│   └── data/                     # quality-200.jsonl + raw rows for 4 runs
 ├── quantization/
-│   ├── w4a16-nvfp4-std.yaml  # the shipped recipe (candidate B)
-│   ├── w4a16-expert-4o6.yaml # experts-only 4/6 alternative (candidate A)
-│   ├── quant-cand-custom.py  # CPU-first quantize driver
-│   └── reattach-mtp.py       # BF16 MTP re-attachment
+│   ├── w4a16-nvfp4-std.yaml      # the shipped recipe (candidate B)
+│   ├── w4a16-expert-4o6.yaml     # experts-only 4/6 alternative (candidate A)
+│   ├── quant-cand-custom.py      # CPU-first quantize driver
+│   └── reattach-mtp.py           # BF16 MTP re-attachment
 └── audits/
-    ├── audit_checkpoint.py   # scale pairing / MTP hashes / key closure
-    ├── cosine_probe.py       # NVFP4 dequant cosine vs BF16 source
-    └── results/              # audit JSON outputs
+    ├── audit_checkpoint.py       # scale pairing / MTP hashes / key closure
+    ├── cosine_probe.py           # NVFP4 dequant cosine vs BF16 source
+    └── results/                  # audit JSON outputs
 ```
 
 ## Credits and attribution

container/Dockerfile.rtx5090

@@ -0,0 +1,37 @@
+# x86-64 / RTX 5090 (SM 12.0) image for Ornith-1.5-35B-A3B-NVFP4-W4A16.
+# Pinned to the exact stack the checkpoint was validated with on GB10:
+#   sglang 0.5.6.post3.dev9218+g5a7b26c63 (git), torch 2.13.0+cu130,
+#   triton 3.7.1, flashinfer-python 0.6.17 (+ flashinfer-cubin for sm_120
+#   cubins so FP4 GEMMs JIT without a full CUDA toolkit).
+# Built by CI (no GPU needed to assemble); first boot on the GPU completes
+# any remaining JIT with the model resident-safe (discrete-GPU host RAM).
+FROM ubuntu:24.04
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && apt-get install -y --no-install-recommends \
+      python3.12 python3-venv ninja-build git ca-certificates curl \
+    && rm -rf /var/lib/apt/lists/* \
+    && ln -sf /usr/bin/python3.12 /usr/local/bin/python \
+    && python -m venv /opt/venv
+ENV PATH=/opt/venv/bin:$PATH \
+    HOME=/home/sglang \
+    HF_HOME=/home/sglang/.cache/huggingface
+
+# Pinned runtime (order matters: torch first, then flashinfer, then sglang).
+RUN pip install --no-cache-dir torch==2.13.0 triton==3.7.1 \
+ && pip install --no-cache-dir flashinfer-python==0.6.17 flashinfer-cubin==0.2.7 \
+      --find-links https://flashinfer.ai/whl/flashinfer-cubin/ \
+ || pip install --no-cache-dir flashinfer-python==0.6.17
+RUN pip install --no-cache-dir "sglang @ git+https://github.com/sgl-project/sglang@5a7b26c63" \
+ && pip install --no-cache-dir transformers accelerate
+
+RUN useradd -m -u 1001 sglang && mkdir -p /home/sglang/.cache \
+    && chown -R sglang:sglang /home/sglang /opt/venv
+USER sglang
+WORKDIR /home/sglang
+
+EXPOSE 8000
+HEALTHCHECK --interval=30s --timeout=10s --start-period=1200s --retries=5 \
+  CMD curl -sf http://127.0.0.1:8000/health || exit 1
+
+COPY --chmod=755 serve-profiles/rtx5090/serve.sh /usr/local/bin/serve.sh
+ENTRYPOINT ["/usr/local/bin/serve.sh"]

docker-compose.rtx5090.yml

@@ -0,0 +1,32 @@
+services:
+  ornith-nvfp4-rtx5090:
+    image: ghcr.io/r0b0tlab/ornith-15-35b-nvfp4-w4a16-sm121-sglang:rtx5090
+    container_name: ornith-nvfp4-rtx5090
+    ipc: host
+    shm_size: 16g
+    environment:
+      MODEL_PATH: /models/ornith-15-35b-a3b-nvfp4-w4a16-B
+      CONTEXT_LENGTH: "32768"
+      MEM_FRACTION_STATIC: "0.92"
+      MAX_RUNNING_REQUESTS: "4"
+      MAX_MAMBA_CACHE_SIZE: "24"
+      # MTP: "0" to disable the EAGLE draft (reclaims ~3.6 GB for KV)
+      PORT: "8000"
+    volumes:
+      - ${MODEL_DIR:-./models/ornith-15-35b-a3b-nvfp4-w4a16-B}:/models/ornith-15-35b-a3b-nvfp4-w4a16-B:ro
+      # Persist HF downloads AND the JIT caches (first-boot sm_120 CUTLASS
+      # compile ~10-20 min; warm boots thereafter).
+      - sglang-cache:/home/sglang/.cache
+    ports:
+      - "8000:8000"
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+    restart: unless-stopped
+
+volumes:
+  sglang-cache:

serve-profiles/rtx5090/serve.sh

@@ -0,0 +1,47 @@
+#!/usr/bin/env bash
+# RTX 5090 (Blackwell SM 12.0, 32 GB GDDR7, x86-64) serve profile for
+# Ornith-1.5-35B-A3B-NVFP4-W4A16.
+#
+# Memory budget on one 32 GB card (measured on the validated GB10 runs,
+# scaled to a dedicated-GPU part):
+#   quantized target weights   ~21.1 GB
+#   BF16 MTP draft (EAGLE)      ~3.6 GB   (disable with MTP=0 to reclaim)
+#   CUDA graphs (verify+draft)  ~1.1 GB
+#   mamba state cache (24 slots)~1.4 GB   (hybrid GDN layers own no KV)
+#   FP8 KV for 4x32768 ctx      ~2.5 GB
+#   => ~29.7 GB @ mem-fraction 0.92; lower MAX_RUNNING_REQUESTS/CONTEXT_LENGTH
+#      if you OOM during graph capture.
+#
+# First boot on a cold cache JIT-compiles the FlashInfer CUTLASS FP4 GEMMs
+# for sm_120 (10-20 min, runs in system RAM on a discrete-GPU host — safe,
+# unlike unified-memory parts). Mount /home/sglang/.cache to persist it.
+set -euo pipefail
+MODEL_PATH="${MODEL_PATH:-r0b0tlab/Ornith-1.5-35B-A3B-NVFP4-W4A16}"
+CONTEXT_LENGTH="${CONTEXT_LENGTH:-32768}"
+MEM_FRACTION_STATIC="${MEM_FRACTION_STATIC:-0.92}"
+MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-4}"
+MAX_MAMBA_CACHE_SIZE="${MAX_MAMBA_CACHE_SIZE:-24}"
+PORT="${PORT:-8000}"
+
+MTP_ARGS=(
+  --speculative-algorithm EAGLE
+  --speculative-num-steps 1
+  --speculative-eagle-topk 1
+  --speculative-num-draft-tokens 2
+)
+if [ "${MTP:-1}" = "0" ]; then MTP_ARGS=(); fi
+
+exec python -m sglang.launch_server \
+  --model-path "$MODEL_PATH" \
+  --served-model-name Ornith-1.5-35B-A3B \
+  --trust-remote-code \
+  --attention-backend triton \
+  --moe-runner-backend marlin \
+  --context-length "$CONTEXT_LENGTH" \
+  --kv-cache-dtype fp8_e4m3 \
+  --mem-fraction-static "$MEM_FRACTION_STATIC" \
+  --max-running-requests "$MAX_RUNNING_REQUESTS" \
+  --max-mamba-cache-size "$MAX_MAMBA_CACHE_SIZE" \
+  "${MTP_ARGS[@]}" \
+  --host 0.0.0.0 --port "$PORT" \
+  ${EXTRA_ARGS:-}