Sign inSign up

orthozany/sglang-blackwell

By orthozany

•Updated 8 months ago

Image
1

298

orthozany/sglang-blackwell repository overview

⁠SGLang Blackwell with PCIe Custom Allreduce

Pre-built SGLang image for NVIDIA Blackwell (SM120) GPUs with luke's custom PCIe allreduce kernel and RMSNorm fusion.

⁠Features

  • PCIe oneshot allreduce — low-latency all-GPU-read kernel for small messages (<64KB)
  • Allreduce + RMSNorm fusion — fuses communication with normalization for reduced latency
  • SM120 (RTX PRO 6000 Blackwell) — built with CUDA 13.0.1, full SM120 support
  • FlashInfer with pre-compiled cubins for SM120
  • torch.compile support for CUDA graph optimization

⁠Source

  • SGLang fork: https://github.com/lukealonso/sglang⁠
  • Fork commit: 24994611c ("Move pcie oneshot to it's own flags")
  • Base image: nvidia/cuda:13.0.1-cudnn-devel-ubuntu24.04 (sha256:92a047cf4837)
  • Build args: CUDA_VERSION=13.0.1, GRACE_BLACKWELL=1, BRANCH_TYPE=local

⁠Dockerfile

# Built from luke's SGLang fork with Blackwell support
# https://github.com/lukealonso/sglang @ 24994611c

FROM nvidia/cuda:13.0.1-cudnn-devel-ubuntu24.04

# See full Dockerfile at:
# https://github.com/lukealonso/sglang/blob/main/docker/Dockerfile

# Build command:
# cd <sglang-repo>
# docker build -t orthozany/sglang-blackwell:luke-ar \
#   -f docker/Dockerfile \
#   --build-arg CUDA_VERSION=13.0.1 \
#   --build-arg GRACE_BLACKWELL=1 \
#   --build-arg BUILD_AND_DOWNLOAD_PARALLEL=8 \
#   --build-arg USE_LATEST_SGLANG=0 \
#   --build-arg BRANCH_TYPE=local \
#   .

⁠Usage

⁠K2.5 with Custom Allreduce + torch.compile
docker run --gpus all --ipc host --shm-size 32g \
  -v /path/to/models/moonshotai/Kimi-K2.5:/model:ro \
  -e NCCL_P2P_LEVEL=SYS -e NCCL_IB_DISABLE=1 -e NCCL_PROTO=LL \
  -p 8400:8000 \
  orthozany/sglang-blackwell:luke-ar \
  python3 -m sglang.launch_server \
    --model-path=/model \
    --served-model-name=Kimi-K2.5 \
    --tp=8 \
    --host=0.0.0.0 --port=8000 \
    --mem-fraction-static=0.94 \
    --attention-backend=flashinfer \
    --enable-pcie-oneshot-allreduce \
    --enable-pcie-oneshot-allreduce-fusion \
    --pcie-oneshot-allreduce-max-size=65536 \
    --enable-torch-compile \
    --enable-mixed-chunk \
    --trust-remote-code \
    --reasoning-parser=kimi_k2 \
    --tool-call-parser=kimi_k2
⁠Custom Allreduce Flags
FlagDescription
--enable-pcie-oneshot-allreduceEnable PCIe system-scope barrier allreduce for small messages
--enable-pcie-oneshot-allreduce-fusionFuse allreduce with RMSNorm for reduced kernel launches
--pcie-oneshot-allreduce-max-size=65536Max message size (bytes) for custom AR (falls back to NCCL above this)

⁠Benchmarks (8x RTX PRO 6000 Blackwell, PCIe Gen5 switches)

⁠K2.5 INT4 Marlin — Server-side gen throughput (tok/s)
ContextCustom AR + TCBaselineDelta
Short116106+9%
4K10899+9%
8K9487+8%
16K8680+8%
32K8075+7%
128K6360+5%
⁠Concurrent Throughput
ConcurrencyAggregate tok/s
8469
16720
641,085
1281,852

⁠Requirements

  • NVIDIA Blackwell GPU (SM120) — RTX PRO 6000, RTX 5090, etc.
  • NVIDIA driver 590+ (tested on 595.45.04)
  • PCIe switches recommended for multi-GPU (custom AR benefits from low-latency P2P)

⁠Credits

  • lukealonso — custom PCIe allreduce kernel and SGLang fork
  • SGLang team — base framework

Tag summary

Content type

Image

Digest

sha256:40b7d8306…

Size

11.3 GB

Last updated

8 months ago

docker pull orthozany/sglang-blackwell:luke-ar