gahow/agentic-pd-hybrid
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: vendor sglang v0.5.10 snapshot

Browse Source
This commit is contained in:
Gahow Wang
2026-04-24 12:29:36 +00:00
parent 78f0d15221
commit bded08301f
4308 changed files with 1200894 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

203
third_party/sglang/python/pyproject.toml vendored Executable file
View File

@@ -0,0 +1,203 @@
[build-system]
requires = ["setuptools>=61.0", "setuptools-scm>=8.0", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "sglang"
dynamic = ["version"]
description = "SGLang is a fast serving framework for large language models and vision language models."
readme = "README.md"
requires-python = ">=3.10"
license = { file = "LICENSE" }
classifiers = [
"Programming Language :: Python :: 3",
"License :: OSI Approved :: Apache Software License",
]
dependencies = [
"IPython",
"aiohttp",
"apache-tvm-ffi>=0.1.5,<0.2",
"anthropic>=0.20.0",
"blobfile==3.0.0",
"build",
"compressed-tensors",
"cuda-python==12.9",
"decord2 ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')",
"datasets",
"einops",
"fastapi",
"flashinfer_python==0.6.7.post2", # keep it aligned with jit-cache version in Dockerfile
"flashinfer_cubin==0.6.7.post2",
"gguf",
"interegular",
"llguidance>=0.7.11,<0.8.0",
"modelscope",
"msgspec",
"ninja",
"numpy",
"nvidia-cutlass-dsl>=4.4.1",
"nvidia-ml-py",
"openai-harmony==0.0.4",
"openai==2.6.1",
"orjson",
"outlines==0.1.11",
"packaging",
"partial_json_parser",
"pillow",
"prometheus-client>=0.20.0",
"psutil",
"py-spy",
"pybase64",
"pydantic",
"python-multipart",
"pyzmq>=25.1.2",
"quack-kernels>=0.3.0",
"requests",
"scipy",
"sentencepiece",
"setproctitle",
"flash-attn-4>=4.0.0b4",
"sglang-kernel==0.4.1",
"soundfile==0.13.1",
"tiktoken",
"timm==1.0.16",
"torch_memory_saver==0.0.9",
"torch==2.9.1",
"torchao==0.9.0",
"torchaudio==2.9.1",
"torchcodec==0.9.1 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torchcodec 0.9.1 for torch 2.9.x. Not available on Linux ARM.
"av ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')",
"torchvision",
"tqdm",
"mistral_common>=1.9.0",
"transformers==5.3.0",
"uvicorn",
"uvloop",
"watchfiles",
"xgrammar==0.1.32",
"smg-grpc-servicer>=0.5.0",
]
[[tool.uv.index]]
name = "pypi"
url = "https://pypi.org/simple"
default = true
[[tool.uv.index]]
name = "torch-cu129"
url = "https://download.pytorch.org/whl/cu129"
explicit = true
[tool.uv.sources]
torch = [
{ index = "pypi", marker = "platform_machine == 'x86_64'"},
{ index = "torch-cu129", marker = "platform_machine == 'aarch64'"},
]
[project.optional-dependencies]
checkpoint-engine = ["checkpoint-engine==0.1.2"]
runai = ["runai-model-streamer[s3,gcs,azure]>=0.15.7"]
diffusion = [
"PyYAML==6.0.1",
"cloudpickle==3.1.2",
"diffusers==0.37.0",
"imageio==2.36.0",
"imageio-ffmpeg==0.5.1",
"moviepy>=2.0.0",
"opencv-python-headless==4.10.0.84",
"remote-pdb==2.1.0",
"st_attn==0.0.7 ; platform_machine != 'aarch64' and platform_machine != 'arm64'",
"vsa==0.0.4 ; platform_machine != 'aarch64' and platform_machine != 'arm64'",
"runai_model_streamer>=0.15.7",
"cache-dit==1.3.0",
"addict==2.4.0",
"av==16.1.0",
"scikit-image==0.25.2",
"trimesh>=4.0.0",
"xatlas",
]
ray = [
"ray[default]>=2.54.0",
]
tracing = [
"opentelemetry-api",
"opentelemetry-exporter-otlp",
"opentelemetry-exporter-otlp-proto-grpc",
"opentelemetry-sdk",
]
test = [
"accelerate",
"addict",
"bitsandbytes",
"expecttest",
"jsonlines",
"lm-eval[api]>=0.4.9.2",
"matplotlib",
"pandas",
"parameterized",
"peft>=0.18.0",
"polars",
"pytest",
"pytest-cov",
"diff-cover",
"sentence_transformers",
"tabulate",
]
dev = ["sglang[test]"]
all = [
"sglang[diffusion]",
"sglang[tracing]",
]
[tool.uv.extra-build-dependencies]
st-attn = ["torch", "setuptools"]
vsa = ["torch", "setuptools"]
[project.urls]
"Homepage" = "https://github.com/sgl-project/sglang"
"Bug Tracker" = "https://github.com/sgl-project/sglang/issues"
[project.scripts]
sglang = "sglang.cli.main:main"
killall_sglang = "sglang.cli.killall:main"
[tool.setuptools.package-data]
"sglang" = [
"srt/**/*",
"jit_kernel/**/*"
]
[tool.setuptools.packages.find]
exclude = [
"assets*",
"benchmark*",
"docs*",
"dist*",
"playground*",
"scripts*",
"tests*",
]
[tool.wheel]
exclude = [
"assets*",
"benchmark*",
"docs*",
"dist*",
"playground*",
"scripts*",
"tests*",
]
[tool.setuptools_scm]
root = ".."
version_file = "sglang/_version.py"
git_describe_command = ["bash", "-c", "git tag --list --sort=-version:refname 'v*.*.*' | head -1 | xargs git describe --tags --long"]
# Allow editable installs even when .git metadata is not available.
fallback_version = "0.0.0.dev0"

130
third_party/sglang/python/pyproject_cpu.toml vendored Normal file
View File

@@ -0,0 +1,130 @@
# https://docs.sglang.io/platforms/cpu_server.html
[build-system]
requires = ["setuptools>=61.0", "setuptools-scm>=8.0", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "sglang-cpu"
dynamic = ["version"]
description = "SGLang is a fast serving framework for large language models and vision language models."
readme = "README.md"
requires-python = ">=3.10"
license = { file = "LICENSE" }
classifiers = [
"Programming Language :: Python :: 3",
"License :: OSI Approved :: Apache Software License",
]
dependencies = [
"IPython",
"aiohttp",
"anthropic>=0.20.0",
"blobfile==3.0.0",
"build",
"compressed-tensors",
"datasets",
"einops",
"fastapi",
"gguf",
"intel-openmp; platform_machine == 'x86_64'",
"interegular",
"llguidance>=0.7.11,<0.8.0",
"modelscope",
"msgspec",
"ninja",
"numpy",
"openai-harmony==0.0.4",
"openai==2.6.1",
"orjson",
"outlines",
"packaging",
"partial_json_parser",
"pillow",
"prometheus-client>=0.20.0",
"psutil",
"py-spy",
"pybase64",
"pydantic",
"python-multipart",
"pyzmq>=25.1.2",
"requests",
"scipy",
"sentencepiece",
"setproctitle",
"soundfile==0.13.1",
"tabulate",
"tiktoken",
"timm==1.0.16",
"torch==2.9.0",
"torchao==0.14.1",
"torchaudio==2.9.0",
"torchvision==0.24.0",
"tqdm",
"mistral_common>=1.9.0",
"transformers==5.3.0",
"triton==3.5.0",
"uvicorn",
"uvloop",
"xgrammar==0.1.32",
"smg-grpc-servicer>=0.5.0",
]
[project.optional-dependencies]
tracing = [
"opentelemetry-sdk",
"opentelemetry-api",
"opentelemetry-exporter-otlp",
"opentelemetry-exporter-otlp-proto-grpc",
]
test = [
"accelerate",
"expecttest",
"jsonlines",
"matplotlib",
"pandas",
"peft>=0.18.0",
"pytest",
"sentence_transformers",
]
all = []
dev = ["sglang[test]"]
[project.urls]
"Homepage" = "https://github.com/sgl-project/sglang"
"Bug Tracker" = "https://github.com/sgl-project/sglang/issues"
[project.scripts]
sglang = "sglang.cli.main:main"
[tool.setuptools.package-data]
"sglang" = [
"srt/**/*",
"jit_kernel/**/*"
]
[tool.setuptools.packages.find]
exclude = [
"assets*",
"benchmark*",
"docs*",
"dist*",
"playground*",
"scripts*",
"tests*",
]
[tool.wheel]
exclude = [
"assets*",
"benchmark*",
"docs*",
"dist*",
"playground*",
"scripts*",
"tests*",
]
[tool.setuptools_scm]
root = ".."
version_file = "sglang/_version.py"
git_describe_command = ["git", "describe", "--tags", "--long", "--match", "v*"]

149
third_party/sglang/python/pyproject_npu.toml vendored Normal file
View File

@@ -0,0 +1,149 @@
[build-system]
requires = ["setuptools>=61.0", "setuptools-scm>=8.0", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "sglang"
dynamic = ["version"]
description = "SGLang is a fast serving framework for large language models and vision language models."
readme = "README.md"
requires-python = ">=3.10"
license = { file = "LICENSE" }
classifiers = [
"Programming Language :: Python :: 3",
"License :: OSI Approved :: Apache Software License",
]
dependencies = [
"IPython",
"aiohttp",
"anthropic>=0.20.0",
"blobfile==3.0.0",
"av",
"build",
"compressed-tensors",
"datasets",
"einops",
"fastapi",
"gguf",
"hf_transfer",
"huggingface_hub",
"interegular",
"llguidance>=0.7.11,<0.8.0",
"modelscope",
"msgspec",
"ninja",
"numpy",
"openai-harmony==0.0.4",
"openai==2.6.1",
"orjson",
"outlines==0.1.11",
"packaging",
"partial_json_parser",
"pillow",
"prometheus-client>=0.20.0",
"psutil",
"py-spy",
"pybase64",
"pydantic",
"python-multipart",
"pyzmq>=25.1.2",
"requests",
"scipy",
"sentencepiece",
"setproctitle",
"soundfile==0.13.1",
"tiktoken",
"timm==1.0.16",
"torchao==0.9.0",
"tqdm",
"mistral_common>=1.9.0",
"transformers==5.3.0",
"uvicorn",
"uvloop",
"xgrammar==0.1.32",
"smg-grpc-servicer>=0.5.0",
]
[project.optional-dependencies]
checkpoint-engine = ["checkpoint-engine==0.1.2"]
diffusion = [
"PyYAML==6.0.1",
"cloudpickle",
"diffusers==0.37.0",
"imageio==2.36.0",
"imageio-ffmpeg==0.5.1",
"moviepy>=2.0.0",
"opencv-python==4.10.0.84",
"remote-pdb",
"cache-dit==1.2.1",
"addict",
"scikit-image==0.25.2",
"trimesh>=4.0.0",
"xatlas",
]
tracing = [
"opentelemetry-api",
"opentelemetry-exporter-otlp",
"opentelemetry-exporter-otlp-proto-grpc",
"opentelemetry-sdk",
]
test = [
"accelerate",
"expecttest",
"gguf",
"jsonlines",
"matplotlib",
"pandas",
"peft>=0.18.0",
"pytest",
"sentence_transformers",
"tabulate",
]
# https://docs.sglang.io/platforms/ascend_npu.html
srt_npu = []
all_npu = ["sglang[diffusion]"]
dev_npu = ["sglang[all_npu]", "sglang[test]"]
[project.urls]
"Homepage" = "https://github.com/sgl-project/sglang"
"Bug Tracker" = "https://github.com/sgl-project/sglang/issues"
[project.scripts]
sglang = "sglang.cli.main:main"
[tool.setuptools.package-data]
"sglang" = [
"srt/**/*",
"jit_kernel/**/*"
]
[tool.setuptools.packages.find]
exclude = [
"assets*",
"benchmark*",
"docs*",
"dist*",
"playground*",
"scripts*",
"tests*",
]
[tool.wheel]
exclude = [
"assets*",
"benchmark*",
"docs*",
"dist*",
"playground*",
"scripts*",
"tests*",
]
[tool.setuptools_scm]
root = ".."
version_file = "sglang/_version.py"
git_describe_command = ["git", "describe", "--tags", "--long", "--match", "v*"]

212
third_party/sglang/python/pyproject_other.toml vendored Executable file
View File

@@ -0,0 +1,212 @@
[build-system]
requires = ["setuptools>=61.0", "setuptools-scm>=8.0", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "sglang"
dynamic = ["version"]
description = "SGLang is a fast serving framework for large language models and vision language models."
readme = "README.md"
requires-python = ">=3.10"
license = { file = "LICENSE" }
classifiers = [
"Programming Language :: Python :: 3",
"License :: OSI Approved :: Apache Software License",
]
dependencies = ["aiohttp", "requests", "tqdm", "numpy", "IPython", "setproctitle"]
[project.optional-dependencies]
runtime_common = [
"IPython",
"aiohttp",
"anthropic>=0.20.0",
"blobfile==3.0.0",
"av",
"build",
"compressed-tensors",
"datasets",
"einops",
"fastapi",
"gguf",
"interegular",
"llguidance>=0.7.11,<0.8.0",
"modelscope",
"msgspec",
"ninja",
"numpy",
"openai-harmony==0.0.4",
"openai==2.6.1",
"orjson",
"outlines==0.1.11",
"packaging",
"partial_json_parser",
"pillow",
"prometheus-client>=0.20.0",
"psutil",
"py-spy",
"pybase64",
"pydantic",
"python-multipart",
"pyzmq>=25.1.2",
"requests",
"scipy",
"sentencepiece",
"setproctitle",
"soundfile==0.13.1",
"tiktoken",
"timm==1.0.16",
"torchao==0.9.0",
"tqdm",
"mistral_common>=1.9.0",
"transformers==5.3.0",
"uvicorn",
"uvloop",
"xgrammar==0.1.32",
"smg-grpc-servicer>=0.5.0",
]
diffusion_common = [
"PyYAML==6.0.1",
"cloudpickle",
"diffusers==0.37.0",
"imageio==2.36.0",
"imageio-ffmpeg==0.5.1",
"moviepy>=2.0.0",
"opencv-python-headless==4.10.0.84",
"remote-pdb",
"addict",
"scikit-image==0.25.2",
"trimesh>=4.0.0",
"xatlas",
]
tracing = [
"opentelemetry-sdk",
"opentelemetry-api",
"opentelemetry-exporter-otlp",
"opentelemetry-exporter-otlp-proto-grpc",
]
# HIP (Heterogeneous-computing Interface for Portability) for AMD
# => base docker rocm/vllm-dev:20250114, not from public vllm whl
srt_hip = [
"sglang[runtime_common]",
"torch",
"petit_kernel==0.0.2",
"wave-lang==3.8.2",
]
diffusion_hip = [
"sglang[diffusion_common]",
"peft>=0.18.0",
"st_attn==0.0.7",
"vsa==0.0.4",
"runai_model_streamer>=0.15.5",
"cache-dit==1.1.8",
]
# For Intel Gaudi(device : hpu) follow the installation guide
# https://docs.vllm.ai/en/latest/getting_started/gaudi-installation.html
srt_hpu = ["sglang[runtime_common]"]
# https://docs.sglang.io/platforms/mthreads_gpu.md
srt_musa = [
"sglang[runtime_common]",
"torch",
"torch_musa",
"torchada>=0.1.45",
"mthreads-ml-py",
"numpy<2.0",
]
diffusion_musa = [
"sglang[diffusion_common]",
"st_attn==0.0.7",
"vsa==0.0.4",
"runai_model_streamer>=0.15.5",
"cache-dit==1.1.8",
]
# https://docs.sglang.io/platforms/mps.md
srt_mps = [
"sglang[runtime_common]",
"torch==2.9.1",
"torchao==0.9.0",
"torchaudio==2.9.1",
"torchvision",
]
diffusion_mps = [
"sglang[diffusion_common]",
"cloudpickle==3.1.2",
"remote-pdb==2.1.0",
"cache-dit==1.2.3",
"addict==2.4.0",
"av==16.1.0",
"scikit-image==0.25.2",
"trimesh>=4.0.0",
"xatlas",
]
test = [
"accelerate",
"expecttest",
"gguf",
"jsonlines",
"matplotlib",
"pandas",
"peft>=0.18.0",
"pytest",
"sentence_transformers",
"tabulate",
]
all_hip = ["sglang[srt_hip]", "sglang[diffusion_hip]"]
all_hpu = ["sglang[srt_hpu]"]
all_musa = ["sglang[srt_musa]", "sglang[diffusion_musa]"]
all_mps = ["sglang[srt_mps]", "sglang[diffusion_mps]"]
dev_hip = ["sglang[all_hip]", "sglang[test]"]
dev_hpu = ["sglang[all_hpu]", "sglang[test]"]
dev_musa = ["sglang[all_musa]", "sglang[test]"]
dev_mps = ["sglang[all_mps]", "sglang[test]"]
[project.urls]
"Homepage" = "https://github.com/sgl-project/sglang"
"Bug Tracker" = "https://github.com/sgl-project/sglang/issues"
[project.scripts]
sglang = "sglang.cli.main:main"
[tool.setuptools.package-data]
"sglang" = [
"srt/**/*",
"jit_kernel/**/*"
]
[tool.setuptools.packages.find]
exclude = [
"assets*",
"benchmark*",
"docs*",
"dist*",
"playground*",
"scripts*",
"tests*",
]
[tool.wheel]
exclude = [
"assets*",
"benchmark*",
"docs*",
"dist*",
"playground*",
"scripts*",
"tests*",
]
[tool.setuptools_scm]
root = ".."
version_file = "sglang/_version.py"
git_describe_command = ["git", "describe", "--tags", "--long", "--match", "v*"]

139
third_party/sglang/python/pyproject_xpu.toml vendored Normal file
View File

@@ -0,0 +1,139 @@
[build-system]
requires = ["setuptools>=61.0", "setuptools-scm>=8.0", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "sglang"
dynamic = ["version"]
description = "SGLang is a fast serving framework for large language models and vision language models."
readme = "README.md"
requires-python = ">=3.10"
license = { file = "LICENSE" }
classifiers = [
"Programming Language :: Python :: 3",
"License :: OSI Approved :: Apache Software License",
]
dependencies = [
"torch==2.10.0+xpu",
"torchcodec==0.10.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torchcodec does not exist in those systems. torch==2.10.0 on XPU uses 0.10.0
"av ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')",
"torchaudio==2.10.0+xpu",
"torchvision",
"sgl-kernel @ git+https://github.com/sgl-project/sgl-kernel-xpu.git",
"IPython",
"aiohttp",
"anthropic>=0.20.0",
"blobfile==3.0.0",
"build",
"compressed-tensors",
"datasets",
"einops",
"fastapi",
"gguf",
"interegular",
"llguidance>=0.7.11,<0.8.0",
"modelscope",
"msgspec",
"ninja",
"numpy",
"openai-harmony==0.0.4",
"openai==2.6.1",
"orjson",
"outlines==0.1.11",
"packaging",
"partial_json_parser",
"pillow",
"prometheus-client>=0.20.0",
"psutil",
"py-spy",
"pybase64",
"pydantic",
"python-multipart",
"pyzmq>=25.1.2",
"requests",
"scipy",
"sentencepiece",
"setproctitle",
"soundfile==0.13.1",
"tiktoken",
"timm==1.0.16",
"torchao==0.9.0",
"tqdm",
"mistral_common>=1.9.0",
"transformers==5.3.0",
"uvicorn",
"uvloop",
# "xgrammar==0.1.24", , xgrammar depends on CUDA PyTorch and Triton only
"smg-grpc-servicer>=0.5.0",
]
[project.optional-dependencies]
tracing = [
"opentelemetry-api",
"opentelemetry-exporter-otlp",
"opentelemetry-exporter-otlp-proto-grpc",
"opentelemetry-sdk",
]
test = [
"accelerate",
"bitsandbytes",
"expecttest",
"jsonlines",
"lm-eval[api]>=0.4.9.2",
"matplotlib",
"pandas",
"parameterized",
"peft>=0.18.0",
"pytest",
"sentence_transformers",
"tabulate",
]
dev = ["sglang[test]"]
all = [
"sglang[tracing]",
]
[project.urls]
"Homepage" = "https://github.com/sgl-project/sglang"
"Bug Tracker" = "https://github.com/sgl-project/sglang/issues"
[project.scripts]
sglang = "sglang.cli.main:main"
[tool.setuptools.package-data]
"sglang" = [
"srt/**/*",
"jit_kernel/**/*"
]
[tool.setuptools.packages.find]
exclude = [
"assets*",
"benchmark*",
"docs*",
"dist*",
"playground*",
"scripts*",
"tests*",
]
[tool.wheel]
exclude = [
"assets*",
"benchmark*",
"docs*",
"dist*",
"playground*",
"scripts*",
"tests*",
]
[tool.setuptools_scm]
root = ".."
version_file = "sglang/_version.py"
git_describe_command = ["bash", "-c", "git tag --list --sort=-version:refname 'v*.*.*' | head -1 | xargs git describe --tags --long"]
# Allow editable installs even when .git metadata is not available.
fallback_version = "0.0.0.dev0"

18
third_party/sglang/python/sglang/README.md vendored Normal file
View File

@@ -0,0 +1,18 @@
# Code Structure
- `eval`: The evaluation utilities.
- `lang`: The frontend language.
- `multimodal_gen`: Inference framework for accelerated image/video generation.
- `srt`: The backend engine for running local models. (SRT = SGLang Runtime).
- `test`: The test utilities.
- `api.py`: The public APIs.
- `bench_offline_throughput.py`: Benchmark the performance in the offline mode.
- `bench_one_batch.py`: Benchmark the latency of running a single static batch without a server.
- `bench_one_batch_server.py`: Benchmark the latency of running a single batch with a server.
- `bench_serving.py`: Benchmark online serving with dynamic requests.
- `check_env.py`: Check the environment variables and dependencies.
- `global_config.py`: The global configs and constants.
- `launch_server.py`: The entry point for launching a local server.
- `profiler.py`: The profiling entry point to send profile requests.
- `utils.py`: Common utilities.
- `version.py`: Version info.

107
third_party/sglang/python/sglang/__init__.py vendored Normal file
View File

@@ -0,0 +1,107 @@
# SGLang public APIs
# Install stubs early for platforms where certain dependencies are unavailable
# (e.g. macOS/MPS has no triton, and torch.mps lacks Stream / set_device /
# get_device_properties). This must run before any downstream imports.
import sys as _sys
if _sys.platform == "darwin":
try:
import torch as _torch
if _torch.backends.mps.is_available():
from sglang._triton_stub import install as _install_triton_stub
_install_triton_stub()
del _install_triton_stub
from sglang._mps_stub import install as _install_mps_stub
_install_mps_stub()
del _install_mps_stub
del _torch
except ImportError:
pass
del _sys
# Frontend Language APIs
from sglang.global_config import global_config
from sglang.lang.api import (
Engine,
Runtime,
assistant,
assistant_begin,
assistant_end,
flush_cache,
function,
gen,
gen_int,
gen_string,
get_server_info,
image,
select,
separate_reasoning,
set_default_backend,
system,
system_begin,
system_end,
user,
user_begin,
user_end,
video,
)
from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
from sglang.lang.choices import (
greedy_token_selection,
token_length_normalized,
unconditional_likelihood_normalized,
)
# Lazy import some libraries
from sglang.utils import LazyImport
from sglang.version import __version__
Anthropic = LazyImport("sglang.lang.backend.anthropic", "Anthropic")
LiteLLM = LazyImport("sglang.lang.backend.litellm", "LiteLLM")
OpenAI = LazyImport("sglang.lang.backend.openai", "OpenAI")
VertexAI = LazyImport("sglang.lang.backend.vertexai", "VertexAI")
# Runtime Engine APIs
ServerArgs = LazyImport("sglang.srt.server_args", "ServerArgs")
Engine = LazyImport("sglang.srt.entrypoints.engine", "Engine")
__all__ = [
"Engine",
"Runtime",
"assistant",
"assistant_begin",
"assistant_end",
"flush_cache",
"function",
"gen",
"gen_int",
"gen_string",
"get_server_info",
"image",
"select",
"separate_reasoning",
"set_default_backend",
"system",
"system_begin",
"system_end",
"user",
"user_begin",
"user_end",
"video",
"RuntimeEndpoint",
"greedy_token_selection",
"token_length_normalized",
"unconditional_likelihood_normalized",
"ServerArgs",
"Anthropic",
"LiteLLM",
"OpenAI",
"VertexAI",
"global_config",
"__version__",
]

270
third_party/sglang/python/sglang/_mps_stub.py vendored Normal file
View File

@@ -0,0 +1,270 @@
"""Stub implementations for APIs missing from ``torch.mps``.
``torch.mps`` lacks several APIs that ``torch.cuda`` provides (``Stream``,
``set_device``, ``get_device_properties``, …). Rather than scattering
``hasattr`` / ``getattr`` guards throughout the codebase, we monkey-patch
``torch.mps`` once at startup so that generic device-agnostic code paths
just work.
"""
from __future__ import annotations
import functools
from dataclasses import dataclass, field
from typing import Any
class Stream:
"""Minimal stand-in for ``torch.cuda.Stream``.
MPS does not expose user-visible streams. Every method is a no-op so
that code written for CUDA's multi-stream model still runs.
"""
def __init__(self, device: Any = None, priority: int = 0) -> None:
pass
def synchronize(self) -> None:
pass
def wait_stream(self, stream: Any) -> None:
pass
def wait_event(self, event: Any) -> None:
pass
def record_event(self, event: Any = None) -> Any:
return None
def query(self) -> bool:
return True
# context-manager protocol (``with stream:``)
def __enter__(self) -> "Stream":
return self
def __exit__(self, *args: Any) -> None:
pass
class Event:
"""Minimal stand-in for ``torch.cuda.Event``."""
def __init__(self, enable_timing: bool = False) -> None:
pass
def record(self, stream: Any = None) -> None:
pass
def wait(self, stream: Any = None) -> None:
pass
def query(self) -> bool:
return True
def synchronize(self) -> None:
pass
def elapsed_time(self, end_event: Any) -> float:
return 0.0
class StreamContext:
"""Minimal stand-in for ``torch.cuda.StreamContext``."""
def __init__(self, stream: Any = None) -> None:
pass
def __enter__(self) -> "StreamContext":
return self
def __exit__(self, *args: Any) -> None:
pass
_default_stream = Stream()
def current_stream(device: Any = None) -> Stream:
"""Return the default (and only) MPS stream."""
return _default_stream
def stream(s: Any) -> Stream:
"""Return a context manager that is a no-op on MPS."""
return s if s is not None else _default_stream
def set_device(device: Any) -> None: # noqa: ARG001
"""Set the current device. This is a no-op for MPS as it has exactly one device."""
pass
def current_device() -> int:
"""Return the index of the current MPS device (always 0)."""
return 0
def device_count() -> int:
"""Return the number of available MPS devices (always 1)."""
return 1
@dataclass
class _MPSDeviceProperties:
"""Mimics the object returned by ``torch.cuda.get_device_properties``."""
name: str = "Apple MPS"
total_memory: int = 0 # populated at install time
multi_processor_count: int = 0
warp_size: int = 32
is_integrated: bool = True
major: int = 0
minor: int = 0
# Extra attrs some callers inspect
_extra: dict = field(default_factory=dict)
def __getattr__(self, name: str) -> Any:
# Return a safe default for any attribute we didn't anticipate
try:
return self._extra[name]
except KeyError:
return None
_cached_props: _MPSDeviceProperties | None = None
def get_device_properties(device: Any = 0) -> _MPSDeviceProperties: # noqa: ARG001
"""Return the properties of the MPS device. Results are cached after first call."""
global _cached_props
if _cached_props is None:
import psutil
_cached_props = _MPSDeviceProperties(
total_memory=psutil.virtual_memory().total,
)
return _cached_props
class _MPSMemoryTracker:
"""Tracks peak memory values on top of ``torch.mps`` current-value APIs.
* ``memory_allocated`` → ``torch.mps.current_allocated_memory()``
* ``memory_reserved`` → ``torch.mps.driver_allocated_memory()``
* ``max_memory_*`` → high-water marks of the above
"""
def __init__(self) -> None:
self._peak_allocated: int = 0
self._peak_reserved: int = 0
def memory_allocated(self, device: Any = None) -> int: # noqa: ARG002
import torch
val = torch.mps.current_allocated_memory()
if val > self._peak_allocated:
self._peak_allocated = val
return val
def memory_reserved(self, device: Any = None) -> int: # noqa: ARG002
import torch
val = torch.mps.driver_allocated_memory()
if val > self._peak_reserved:
self._peak_reserved = val
return val
def max_memory_allocated(self, device: Any = None) -> int: # noqa: ARG002
self.memory_allocated()
return self._peak_allocated
def max_memory_reserved(self, device: Any = None) -> int: # noqa: ARG002
self.memory_reserved()
return self._peak_reserved
def reset_peak_memory_stats(self, device: Any = None) -> None: # noqa: ARG002
import torch
self._peak_allocated = torch.mps.current_allocated_memory()
self._peak_reserved = torch.mps.driver_allocated_memory()
_memory_tracker = _MPSMemoryTracker()
def _patch_non_blocking() -> None:
"""Force ``non_blocking=False`` for copies targeting the MPS device.
Unlike CUDA, MPS does not guarantee that a subsequent kernel on the same
"stream" will wait for an async host-to-device transfer to finish. Reading
the tensor before the transfer completes yields uninitialised (garbage)
data. Patching ``Tensor.to`` and ``Tensor.copy_`` centrally avoids having
to sprinkle ``non_blocking=not is_mps()`` at every call-site.
"""
import torch
_original_to = torch.Tensor.to
@functools.wraps(_original_to)
def _patched_to(self, *args, **kwargs):
if kwargs.get("non_blocking"):
# Detect target device from positional or keyword args
device = None
if args and isinstance(args[0], (str, torch.device)):
device = torch.device(args[0]) if isinstance(args[0], str) else args[0]
elif "device" in kwargs:
d = kwargs["device"]
device = torch.device(d) if isinstance(d, str) else d
if device is not None and device.type == "mps":
kwargs = {**kwargs, "non_blocking": False}
return _original_to(self, *args, **kwargs)
torch.Tensor.to = _patched_to
_original_copy_ = torch.Tensor.copy_
@functools.wraps(_original_copy_)
def _patched_copy_(self, src, non_blocking=False):
if non_blocking and self.device.type == "mps":
non_blocking = False
return _original_copy_(self, src, non_blocking=non_blocking)
torch.Tensor.copy_ = _patched_copy_
_installed = False
def install() -> None:
"""Patch ``torch.mps`` with the stubs above. Safe to call multiple times."""
global _installed
if _installed:
return
import torch
mps = torch.mps
# Only patch attributes that are actually missing
for name, obj in [
("Stream", Stream),
("StreamContext", StreamContext),
("Event", Event),
("current_stream", current_stream),
("stream", stream),
("set_device", set_device),
("current_device", current_device),
("device_count", device_count),
("get_device_properties", get_device_properties),
("reset_peak_memory_stats", _memory_tracker.reset_peak_memory_stats),
("memory_allocated", _memory_tracker.memory_allocated),
("memory_reserved", _memory_tracker.memory_reserved),
("max_memory_allocated", _memory_tracker.max_memory_allocated),
("max_memory_reserved", _memory_tracker.max_memory_reserved),
]:
if not hasattr(mps, name):
setattr(mps, name, obj)
_patch_non_blocking()
_installed = True

228
third_party/sglang/python/sglang/_triton_stub.py vendored Normal file
View File

@@ -0,0 +1,228 @@
"""
Mock triton module for platforms where triton is not available (e.g., macOS/MPS).
This module provides stub implementations of triton APIs so that modules which
import triton at the top level can be loaded without error. The actual triton
kernels are never executed on these platforms – alternative backends (e.g. SDPA
for MPS) are used instead.
Usage – call ``install()`` **before** any ``import triton`` in the process:
from sglang._triton_stub import install
install()
"""
import sys
import types
class _StubBase:
"""A base class that any mock attribute can safely be subclassed from.
Used when external code does ``class Foo(triton.runtime.KernelInterface):``.
"""
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
class _MockModule(types.ModuleType):
"""A module whose every attribute is itself a ``_MockModule``.
When called (e.g. ``@triton.jit``), it acts as a pass-through decorator so
that kernel *definitions* are syntactically valid even though they will never
be compiled.
"""
def __init__(self, name: str):
super().__init__(name)
self.__path__: list[str] = [] # make it look like a package
self.__package__ = name
self.__file__ = __file__
self._children: dict[str, object] = {}
# Set __spec__ so that importlib.util.find_spec() works on cached modules
import importlib
self.__spec__ = importlib.machinery.ModuleSpec(name, None, is_package=True)
def __getattr__(self, name: str):
"""Handle attribute access by creating and returning a child _MockModule."""
if name.startswith("__") and name.endswith("__"):
raise AttributeError(name)
full = f"{self.__name__}.{name}"
if full in sys.modules:
return sys.modules[full]
# If the name looks like a class (CamelCase / uppercase), return a
# stub class that can be used as a base class for inheritance.
if name[0:1].isupper():
stub_cls = type(name, (_StubBase,), {"__module__": self.__name__})
self._children[name] = stub_cls
return stub_cls
child = _MockModule(full)
sys.modules[full] = child
self._children[name] = child
return child
def __call__(self, *args, **kwargs):
# Direct decorator usage: @triton.jit (receives the function)
if len(args) == 1 and callable(args[0]) and not kwargs:
return args[0]
# Parameterised decorator: @triton.jit(...) → returns a decorator
def _decorator(fn):
return fn
return _decorator
def __instancecheck__(self, instance):
"""Return False for all instance checks against the mock."""
return False
def __contains__(self, item):
"""Return False for all membership checks."""
return False
def __iter__(self):
return iter([])
def __len__(self):
return 0
def __bool__(self):
return False
def __repr__(self):
return f"<triton-stub {self.__name__!r}>"
def _cdiv(a: int, b: int) -> int:
"""Ceiling division – mirrors ``triton.cdiv``."""
return -(a // -b)
def _next_power_of_2(n: int) -> int:
"""Mirrors ``triton.next_power_of_2``."""
return 1 << (n - 1).bit_length() if n > 0 else 1
class _Config:
"""Minimal stand-in for ``triton.Config`` used in ``@triton.autotune``."""
def __init__(self, kwargs=None, num_warps=4, num_stages=2, **extra):
self.kwargs = kwargs or {}
self.num_warps = num_warps
self.num_stages = num_stages
class _TritonFinder:
"""A meta-path finder that intercepts all ``import triton.*`` statements.
When Python encounters ``import triton.backends.compiler``, it walks the
dotted path and tries to import each component. Our mock module's
``__getattr__`` handles *attribute* access, but the import machinery uses
``importlib`` finders, not attribute access, for sub-module resolution.
This finder bridges that gap by creating ``_MockModule`` instances for any
``triton.*`` sub-module that isn't already in ``sys.modules``.
"""
def find_spec(self, fullname, path=None, target=None):
"""PEP 451 meta-path finder for ``triton.*`` sub-modules."""
if fullname == "triton" or fullname.startswith("triton."):
if fullname in sys.modules:
return getattr(sys.modules[fullname], "__spec__", None)
# Create and register the mock so the import machinery finds it
mod = _MockModule(fullname)
sys.modules[fullname] = mod
parts = fullname.rsplit(".", 1)
if len(parts) == 2:
parent_name, child_name = parts
parent = sys.modules.get(parent_name)
if parent is not None:
setattr(parent, child_name, mod)
return mod.__spec__
return None
def _make_mock(name: str) -> _MockModule:
"""Create a ``_MockModule`` and register it in ``sys.modules``."""
mod = _MockModule(name)
sys.modules[name] = mod
return mod
def install() -> None:
"""Register a mock ``triton`` package in *sys.modules*.
This is a no-op if a real ``triton`` is already importable.
"""
if "triton" in sys.modules:
return
# Check whether a real triton exists before installing the stub.
import importlib.util
if importlib.util.find_spec("triton") is not None:
return
# Register the meta-path finder FIRST so that any ``import triton.X``
# during the rest of install() (or later) is handled.
sys.meta_path.insert(0, _TritonFinder())
triton = _make_mock("triton")
triton.__version__ = "3.0.0"
triton.cdiv = _cdiv
triton.next_power_of_2 = _next_power_of_2
triton.Config = _Config
# triton.language (commonly imported as ``tl``)
tl = _make_mock("triton.language")
class _constexpr:
"""Stand-in for ``tl.constexpr`` – works as both annotation and value wrapper."""
def __init__(self, value=None):
self.value = value
def __repr__(self):
return f"constexpr({self.value!r})"
tl.constexpr = _constexpr
triton.language = tl
# triton.language.extra.libdevice
extra = _make_mock("triton.language.extra")
tl.extra = extra
libdevice = _make_mock("triton.language.extra.libdevice")
extra.libdevice = libdevice
# triton.runtime.jit (JITFunction used in isinstance checks)
runtime = _make_mock("triton.runtime")
triton.runtime = runtime
jit_mod = _make_mock("triton.runtime.jit")
class _JITFunction:
"""Dummy so ``isinstance(fn, triton.runtime.jit.JITFunction)`` works."""
pass
jit_mod.JITFunction = _JITFunction
runtime.jit = jit_mod
# triton.runtime.driver (used by fla/utils.py)
driver = _make_mock("triton.runtime.driver")
runtime.driver = driver
# triton.testing
testing = _make_mock("triton.testing")
triton.testing = testing
# triton.tools / triton.tools.tensor_descriptor
tools = _make_mock("triton.tools")
triton.tools = tools
td = _make_mock("triton.tools.tensor_descriptor")
tools.tensor_descriptor = td
# triton.backends / triton.backends.compiler (used by torch._inductor)
backends = _make_mock("triton.backends")
triton.backends = backends
compiler = _make_mock("triton.backends.compiler")
backends.compiler = compiler

543
third_party/sglang/python/sglang/bench_offline_throughput.py vendored Normal file
View File

@@ -0,0 +1,543 @@
"""
Benchmark the throughput in the offline mode.
It accepts server arguments (the same as launch_server.py) and benchmark arguments (the same as bench_serving.py).
# Usage
## Sharegpt dataset with default args
python -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --num-prompts 10
## Random dataset with default args
python -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024
"""
import argparse
import asyncio
import dataclasses
import inspect
import json
import logging
import os
import random
import time
from typing import Dict, List, Optional
import numpy as np
from sglang.benchmark.datasets import DatasetRow, get_dataset
from sglang.benchmark.datasets.random import sample_random_requests
from sglang.benchmark.utils import get_tokenizer, set_ulimit
from sglang.lang.backend.runtime_endpoint import Runtime
from sglang.srt.entrypoints.engine import Engine
from sglang.srt.server_args import ServerArgs
@dataclasses.dataclass
class BenchArgs:
backend: str = "engine"
result_filename: str = ""
dataset_name: str = "sharegpt"
dataset_path: str = ""
num_prompts: int = 1000
sharegpt_output_len: Optional[int] = None
sharegpt_context_len: Optional[int] = None
random_input_len: int = 1024
random_output_len: int = 1024
random_range_ratio: float = 0.0
gsp_num_groups: int = 64
gsp_prompts_per_group: int = 16
gsp_system_prompt_len: int = 2048
gsp_question_len: int = 128
gsp_output_len: int = 256
seed: int = 1
disable_ignore_eos: bool = False
extra_request_body: Optional[str] = None
apply_chat_template: bool = False
profile: bool = False
skip_warmup: bool = False
do_not_exit: bool = False
prompt_suffix: str = ""
return_logprob: bool = False
logprob_start_len: int = -1
@staticmethod
def add_cli_args(parser: argparse.ArgumentParser):
parser.add_argument("--backend", type=str, default=BenchArgs.backend)
parser.add_argument(
"--result-filename", type=str, default=BenchArgs.result_filename
)
parser.add_argument(
"--dataset-name",
type=str,
default="sharegpt",
choices=["sharegpt", "random", "generated-shared-prefix"],
help="Name of the dataset to benchmark on.",
)
parser.add_argument(
"--dataset-path", type=str, default="", help="Path to the dataset."
)
parser.add_argument(
"--num-prompts",
type=int,
default=BenchArgs.num_prompts,
help="Number of prompts to process. Default is 1000.",
)
parser.add_argument(
"--sharegpt-output-len",
type=int,
default=BenchArgs.sharegpt_output_len,
help="Output length for each request. Overrides the output length from the ShareGPT dataset.",
)
parser.add_argument(
"--sharegpt-context-len",
type=int,
default=BenchArgs.sharegpt_context_len,
help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.",
)
parser.add_argument(
"--random-input-len",
type=int,
default=BenchArgs.random_input_len,
help="Number of input tokens per request, used only for random dataset.",
)
parser.add_argument(
"--random-output-len",
type=int,
default=BenchArgs.random_output_len,
help="Number of output tokens per request, used only for random dataset.",
)
parser.add_argument(
"--random-range-ratio",
type=float,
default=BenchArgs.random_range_ratio,
help="Range of sampled ratio of input/output length, "
"used only for random dataset.",
)
parser.add_argument(
"--gsp-num-groups",
type=int,
default=BenchArgs.gsp_num_groups,
help="Number of groups with shared prefix, used"
"only for generate-shared-prefix",
)
parser.add_argument(
"--gsp-prompts-per-group",
type=int,
default=BenchArgs.gsp_prompts_per_group,
help="Number of prompts per group of shared prefix, used"
"only for generate-shared-prefix",
)
parser.add_argument(
"--gsp-system-prompt-len",
type=int,
default=BenchArgs.gsp_system_prompt_len,
help="System prompt length, used" "only for generate-shared-prefix",
)
parser.add_argument(
"--gsp-question-len",
type=int,
default=BenchArgs.gsp_question_len,
help="Question length, used" "only for generate-shared-prefix",
)
parser.add_argument(
"--gsp-output-len",
type=int,
default=BenchArgs.gsp_output_len,
help="Target length in tokens for outputs in generated-shared-prefix dataset",
)
parser.add_argument("--seed", type=int, default=1, help="The random seed.")
parser.add_argument(
"--disable-ignore-eos",
action="store_true",
help="Disable ignore EOS token",
)
parser.add_argument(
"--extra-request-body",
metavar='{"key1": "value1", "key2": "value2"}',
type=str,
default=BenchArgs.extra_request_body,
help="Append given JSON object to the request payload. You can use this to specify"
"additional generate params like sampling params.",
)
parser.add_argument(
"--apply-chat-template",
action="store_true",
help="Apply chat template",
)
parser.add_argument(
"--profile",
action="store_true",
help="Use Torch Profiler. The endpoint must be launched with "
"SGLANG_TORCH_PROFILER_DIR to enable profiler.",
)
parser.add_argument(
"--skip-warmup",
action="store_true",
help="Skip the warmup batches.",
)
parser.add_argument(
"--do-not-exit",
action="store_true",
help="Do not exit the program. This is useful for nsys profile with --duration and --delay.",
)
parser.add_argument(
"--prompt-suffix",
type=str,
default="",
help="Suffix applied to the end of all user prompts, followed by assistant prompt suffix.",
)
parser.add_argument(
"--return-logprob",
action="store_true",
help="Enable returning log probabilities.",
)
parser.add_argument(
"--logprob-start-len",
type=int,
default=-1,
help="Start length for logprob. -1 means only return logprobs for output tokens (default). 0 means return logprobs for all tokens including input.",
)
@classmethod
def from_cli_args(cls, args: argparse.Namespace):
attrs = [attr.name for attr in dataclasses.fields(cls)]
return cls(**{attr: getattr(args, attr) for attr in attrs})
def throughput_test_once(
backend_name: str,
backend,
reqs: List[DatasetRow],
ignore_eos: bool,
extra_request_body: Dict,
profile: bool,
return_logprob: bool = False,
logprob_start_len: int = -1,
):
measurement_results = {
"backend": backend_name,
"successful_requests": len(reqs),
"total_latency": -1,
"total_input_tokens": sum(r.prompt_len for r in reqs),
"total_output_tokens": -1,
"request_throughput": -1,
"input_throughput": -1,
"output_throughput": -1,
"total_throughput": -1,
}
prompt = [r.prompt for r in reqs]
sampling_params = [
{
"temperature": 0,
"max_new_tokens": r.output_len,
"ignore_eos": ignore_eos,
**extra_request_body,
}
for r in reqs
]
if profile:
assert (
"SGLANG_TORCH_PROFILER_DIR" in os.environ
), "Please set SGLANG_TORCH_PROFILER_DIR."
os.makedirs(os.environ["SGLANG_TORCH_PROFILER_DIR"], exist_ok=True)
backend.start_profile()
st = time.perf_counter()
gen_out = backend.generate(
prompt=prompt,
sampling_params=sampling_params,
return_logprob=return_logprob,
logprob_start_len=logprob_start_len,
)
latency = time.perf_counter() - st
if profile:
dir = os.getenv("SGLANG_TORCH_PROFILER_DIR")
known_files = set(os.listdir(dir))
backend.stop_profile()
monitor_trace_file(known_files, dir)
if backend_name == "runtime":
gen_out = json.loads(gen_out)
server_info = backend.get_server_info()
measurement_results["total_latency"] = latency
measurement_results["total_output_tokens"] = sum(
o["meta_info"]["completion_tokens"] for o in gen_out
)
measurement_results["request_throughput"] = (
measurement_results["successful_requests"] / latency
)
measurement_results["input_throughput"] = (
measurement_results["total_input_tokens"] / latency
)
measurement_results["output_throughput"] = (
measurement_results["total_output_tokens"] / latency
)
measurement_results["total_throughput"] = (
measurement_results["total_input_tokens"]
+ measurement_results["total_output_tokens"]
) / latency
if inspect.isawaitable(server_info):
server_info = asyncio.run(server_info)
measurement_results["last_gen_throughput"] = server_info["internal_states"][0][
"last_gen_throughput"
]
return measurement_results
def monitor_trace_file(known_files, directory, interval=1):
print(f"Monitoring {directory} for new trace files...")
while True:
flag = False
time.sleep(interval)
current_files = set(os.listdir(directory))
new_files = current_files - known_files
for new_file in new_files:
new_file_path = os.path.join(directory, new_file)
print(f"New file detected: {new_file}")
previous_size = 0
while True:
try:
current_size = os.path.getsize(new_file_path)
except FileNotFoundError:
print(f"File {new_file} is no longer accessible.")
break
if current_size > previous_size:
previous_size = current_size
else:
flag = True
break
time.sleep(interval)
if flag:
break
def _create_ray_engine_backend(server_args: ServerArgs):
"""Create a RayEngine inside a Ray actor on a placement group.
RayEngine requires a placement group, so we launch it inside a Ray actor
and return a lightweight proxy that forwards calls via ray.get().
"""
import ray
from ray.runtime_env import RuntimeEnv
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
env_vars = {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}
if os.environ.get("HF_TOKEN"):
env_vars["HF_TOKEN"] = os.environ["HF_TOKEN"]
if not ray.is_initialized():
ray.init(runtime_env=RuntimeEnv(env_vars=env_vars))
total_gpus = server_args.tp_size * server_args.pp_size
pg = placement_group([{"CPU": 1, "GPU": total_gpus}], strategy="STRICT_PACK")
ray.get(pg.ready())
@ray.remote
class _EngineActor:
def __init__(self, **kwargs):
from sglang.srt.ray.engine import RayEngine
self.engine = RayEngine(**kwargs)
def call(self, method, **kwargs):
return getattr(self.engine, method)(**kwargs)
actor = _EngineActor.options(
num_cpus=1,
num_gpus=0,
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=pg,
placement_group_bundle_index=0,
),
).remote(**dataclasses.asdict(server_args))
class _Proxy:
"""Forwards method calls to the remote RayEngine actor."""
def generate(self, **kwargs):
return ray.get(actor.call.remote("generate", **kwargs))
def get_server_info(self, **kwargs):
return ray.get(actor.call.remote("get_server_info", **kwargs))
def start_profile(self, **kwargs):
return ray.get(actor.call.remote("start_profile", **kwargs))
def stop_profile(self, **kwargs):
return ray.get(actor.call.remote("stop_profile", **kwargs))
def shutdown(self):
try:
ray.get(actor.call.remote("shutdown"), timeout=60)
except Exception:
pass
try:
ray.util.remove_placement_group(pg)
except Exception:
pass
return _Proxy()
def throughput_test(
server_args: ServerArgs,
bench_args: BenchArgs,
):
if bench_args.backend == "engine":
if server_args.use_ray:
backend = _create_ray_engine_backend(server_args)
else:
backend = Engine(**dataclasses.asdict(server_args))
if not backend:
raise ValueError("Please provide valid engine arguments")
elif bench_args.backend == "runtime":
backend = Runtime(**dataclasses.asdict(server_args))
else:
raise ValueError('Please set backend to either "engine" or "runtime"')
tokenizer_id = server_args.tokenizer_path or server_args.model_path
tokenizer = get_tokenizer(tokenizer_id)
# Set global environments
set_ulimit()
random.seed(bench_args.seed)
np.random.seed(bench_args.seed)
# Parse args
extra_request_body = {}
if bench_args.extra_request_body:
extra_request_body = json.loads(args.extra_request_body)
# Read dataset
input_requests = get_dataset(bench_args, tokenizer)
warmup_requests = sample_random_requests(
input_len=256,
output_len=16,
num_prompts=min(bench_args.num_prompts, 16),
range_ratio=1.0,
tokenizer=tokenizer,
dataset_path=bench_args.dataset_path,
)
# Warm up
if not bench_args.skip_warmup:
logging.info("\nWarmup...")
throughput_test_once(
backend_name=bench_args.backend,
backend=backend,
reqs=warmup_requests,
ignore_eos=not bench_args.disable_ignore_eos,
extra_request_body=extra_request_body,
profile=False,
return_logprob=bench_args.return_logprob,
logprob_start_len=bench_args.logprob_start_len,
)
time.sleep(0.5)
logging.info("\nBenchmark...")
result = throughput_test_once(
backend_name=bench_args.backend,
backend=backend,
reqs=input_requests,
ignore_eos=not bench_args.disable_ignore_eos,
extra_request_body=extra_request_body,
profile=bench_args.profile,
return_logprob=bench_args.return_logprob,
logprob_start_len=bench_args.logprob_start_len,
)
backend.shutdown()
if bench_args.result_filename:
with open(bench_args.result_filename, "a") as fout:
fout.write(json.dumps(result) + "\n")
print(
"\n{s:{c}^{n}}".format(s=" Offline Throughput Benchmark Result ", n=50, c="=")
)
print("{:<40} {:<10}".format("Backend:", result["backend"]))
print("{:<40} {:<10}".format("Successful requests:", result["successful_requests"]))
print("{:<40} {:<10.2f}".format("Benchmark duration (s):", result["total_latency"]))
print("{:<40} {:<10}".format("Total input tokens:", result["total_input_tokens"]))
print(
"{:<40} {:<10}".format("Total generated tokens:", result["total_output_tokens"])
)
print(
"{:<40} {:<10.2f}".format(
"Last generation throughput (tok/s):", result["last_gen_throughput"]
)
)
print(
"{:<40} {:<10.2f}".format(
"Request throughput (req/s):", result["request_throughput"]
)
)
print(
"{:<40} {:<10.2f}".format(
"Input token throughput (tok/s):", result["input_throughput"]
)
)
print(
"{:<40} {:<10.2f}".format(
"Output token throughput (tok/s):", result["output_throughput"]
)
)
print(
"{:<40} {:<10.2f}".format(
"Total token throughput (tok/s):", result["total_throughput"]
)
)
print("=" * 50)
return result
if __name__ == "__main__":
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
BenchArgs.add_cli_args(parser)
args = parser.parse_args()
# handling ModelScope model downloads
if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() in ("true", "1"):
if os.path.exists(args.model_path):
print(f"Using local model path: {args.model_path}")
else:
try:
from modelscope import snapshot_download
print(f"Using ModelScope to download model: {args.model_path}")
# download the model and replace args.model_path
args.model_path = snapshot_download(
args.model_path,
)
print(f"Model downloaded to: {args.model_path}")
except Exception as e:
print(f"ModelScope download failed: {str(e)}")
raise e
server_args = ServerArgs.from_cli_args(args)
bench_args = BenchArgs.from_cli_args(args)
logging.basicConfig(
level=getattr(logging, server_args.log_level.upper()),
format="%(message)s",
)
throughput_test(server_args, bench_args)
while bench_args.do_not_exit:
pass

936
third_party/sglang/python/sglang/bench_one_batch.py vendored Normal file
View File

@@ -0,0 +1,936 @@
"""
Benchmark the latency of running a single static batch without a server.
This script does not launch a server and uses the low-level APIs.
It accepts server arguments (the same as launch_server.py) and benchmark arguments (e.g., batch size, input lengths).
# Usage (latency test)
## with dummy weights:
python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --load-format dummy
## sweep through multiple data points and store (append) the results in a jsonl file:
python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --output-len 32 256 --run-name test_run
## run with profiling:
python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --profile
## run with profiling to custom directory:
export SGLANG_TORCH_PROFILER_DIR=/root/sglang/profile_log
python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 --input-len 256 --profile
## run with CUDA profiler (nsys):
nsys profile --force-overwrite=true -o bench_one_batch python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 --input-len 256 --profile --profile-activities CUDA_PROFILER
# Usage (correctness test):
python -m sglang.bench_one_batch --model-path TinyLlama/TinyLlama-1.1B-Chat-v0.4 --correct
## Reference output (of the correctness test above, can be gpu dependent):
input_ids=[[1, 450, 7483, 310, 3444, 338], [1, 450, 7483, 310, 278, 3303, 13187, 290, 338], [1, 20628, 338, 263, 6575, 1460, 2462, 322, 306, 763]]
prefill logits (first half): tensor([[-10.0312, -9.5000, 0.8931, ..., -4.9414, -3.2422, -3.3633],
[-10.0312, -9.5000, 0.8931, ..., -4.9414, -3.2422, -3.3633],
[ -9.1875, -10.2500, 2.7129, ..., -4.3359, -4.0664, -4.1328]],
device='cuda:0')
prefill logits (final): tensor([[-8.3125, -7.1172, 3.3457, ..., -4.9570, -4.1328, -3.4141],
[-8.9141, -9.0156, 4.1445, ..., -4.9922, -4.4961, -4.0781],
[-9.6328, -9.0547, 4.0195, ..., -5.3047, -4.7148, -4.4570]],
device='cuda:0')
========== Prompt 0 ==========
<s> The capital of France is Paris.
The capital of the United States is Washington, D.C.
========== Prompt 1 ==========
<s> The capital of the United Kindom is London.
The capital of the United Kingdom is London.
The capital of the
========== Prompt 2 ==========
<s> Today is a sunny day and I like to go for a walk in the park.
I'm going to the park
"""
import argparse
import copy
import dataclasses
import itertools
import json
import logging
import multiprocessing
import os
import time
from types import SimpleNamespace
from typing import Optional, Tuple
import numpy as np
import torch
import torch.distributed as dist
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed.parallel_state import destroy_distributed_environment
from sglang.srt.entrypoints.engine import _set_envs_and_config
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.layers.moe import initialize_moe_config
from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config
from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
from sglang.srt.managers.scheduler_dp_attn_mixin import prepare_mlp_sync_batch_raw
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.utils import (
configure_logger,
get_bool_env_var,
kill_process_tree,
maybe_reindex_device_id,
require_mlp_sync,
require_mlp_tp_gather,
set_gpu_proc_affinity,
suppress_other_loggers,
)
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
from sglang.srt.utils.tensor_bridge import use_mlx
def start_profile(profile_activities, profile_record_shapes=False, rank_print=print):
"""
Abstracted function to start profiling based on profile_activities.
Returns profiler object (or None).
"""
if "CUDA_PROFILER" in profile_activities:
try:
torch.cuda.cudart().cudaProfilerStart()
rank_print("CUDA Profiler started (nsys will begin capturing)")
except Exception as e:
rank_print(f"Failed to start CUDA profiler: {e}")
return None
else:
activities = []
if "CPU" in profile_activities:
activities.append(torch.profiler.ProfilerActivity.CPU)
if "GPU" in profile_activities:
activities.append(torch.profiler.ProfilerActivity.CUDA)
if "XPU" in profile_activities:
activities.append(torch.profiler.ProfilerActivity.XPU)
if activities:
profiler = torch.profiler.profile(
activities=activities,
with_stack=True,
record_shapes=profile_record_shapes,
)
profiler.start()
return profiler
return None
def stop_profile(
profiler,
profile_activities,
rank_print=print,
save_trace=False,
trace_filename=None,
stage=None,
):
"""
Abstracted function to stop profiling based on profile_activities.
Optionally saves trace results and prints completion messages.
"""
if "CUDA_PROFILER" in profile_activities:
try:
torch.cuda.cudart().cudaProfilerStop()
rank_print("CUDA Profiler stopped (nsys should dump traces)")
except Exception as e:
rank_print(f"Failed to stop CUDA profiler: {e}")
elif profiler is not None:
profiler.stop()
if save_trace:
if profiler is not None:
if trace_filename:
_save_profile_trace_results(profiler, trace_filename)
stage_desc = f"for {stage}" if stage else ""
rank_print(
f"torch profiler chrome trace {stage_desc} saved to {trace_filename}"
)
if "CUDA_PROFILER" in profile_activities:
rank_print(f"CUDA profiler trace for {stage} completed")
@dataclasses.dataclass
class BenchArgs:
run_name: str = "default"
batch_size: Tuple[int] = (1,)
input_len: Tuple[int] = (1024,)
output_len: Tuple[int] = (16,)
prompt_filename: str = ""
result_filename: str = "result.jsonl"
correctness_test: bool = False
# This is only used for correctness test
cut_len: int = 4
log_decode_step: int = 0
profile: bool = False
profile_record_shapes: bool = False
profile_activities: Tuple[str] = ("CPU", "GPU")
profile_stage: str = "all"
profile_filename_prefix: str = "profile"
profile_start_step: Optional[int] = None
profile_steps: Optional[int] = None
@staticmethod
def add_cli_args(parser: argparse.ArgumentParser):
parser.add_argument("--run-name", type=str, default=BenchArgs.run_name)
parser.add_argument(
"--batch-size", type=int, nargs="+", default=BenchArgs.batch_size
)
parser.add_argument(
"--input-len", type=int, nargs="+", default=BenchArgs.input_len
)
parser.add_argument(
"--output-len", type=int, nargs="+", default=BenchArgs.output_len
)
parser.add_argument(
"--prompt-filename", type=str, default=BenchArgs.prompt_filename
)
parser.add_argument(
"--result-filename", type=str, default=BenchArgs.result_filename
)
parser.add_argument("--correctness-test", action="store_true")
parser.add_argument("--cut-len", type=int, default=BenchArgs.cut_len)
parser.add_argument(
"--log-decode-step",
type=int,
default=BenchArgs.log_decode_step,
help="Log decode latency by step, default is set to zero to disable.",
)
parser.add_argument("--profile", action="store_true", help="Enable profiling.")
parser.add_argument(
"--profile-record-shapes",
action="store_true",
help="Record tensor shapes in profiling results.",
)
parser.add_argument(
"--profile-activities",
type=str,
nargs="+",
default=["CPU", "GPU"],
choices=["CPU", "GPU", "CUDA_PROFILER", "XPU"],
help="Profiler activities: CPU, GPU, XPU, CUDA_PROFILER. If CPU/GPU/XPU, use torch profiler. If CUDA_PROFILER, use CUDA profiler.",
)
parser.add_argument(
"--profile-stage",
type=str,
default=BenchArgs.profile_stage,
choices=["all", "prefill", "decode"],
help="Which stage to profile: all, prefill, or decode only.",
)
parser.add_argument(
"--profile-filename-prefix",
type=str,
default=BenchArgs.profile_filename_prefix,
help="Prefix of the profiling file names. The full profiling result file(s) be "
'"[profile_filename_prefix]_batch[batch_size]_input[input_len]_output[output_len].trace.json.gz"',
)
parser.add_argument(
"--profile-start-step",
type=int,
default=None,
help="Decode step at which to start profiling (0-indexed). If not specified, defaults to output_len // 2.",
)
parser.add_argument(
"--profile-steps",
type=int,
default=None,
help="Number of decode steps to profile starting from profile-start-step. If not specified, profiles only one step.",
)
@classmethod
def from_cli_args(cls, args: argparse.Namespace):
# use the default value's type to cast the args into correct types.
attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)]
result = {}
for attr, attr_type in attrs:
value = getattr(args, attr)
# Handle None values - don't try to cast them
if value is None or attr_type == type(None):
result[attr] = value
else:
result[attr] = attr_type(value)
return cls(**result)
def load_model(server_args, port_args, gpu_id, tp_rank):
suppress_other_loggers()
rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None
moe_ep_rank = tp_rank // (server_args.tp_size // server_args.ep_size)
model_config = ModelConfig.from_server_args(server_args)
runner_kwargs = dict(
model_config=model_config,
mem_fraction_static=server_args.mem_fraction_static,
gpu_id=gpu_id,
tp_rank=tp_rank,
tp_size=server_args.tp_size,
moe_ep_rank=moe_ep_rank,
moe_ep_size=server_args.ep_size,
pp_rank=0,
pp_size=1,
nccl_port=port_args.nccl_port,
server_args=server_args,
)
_use_mlx = use_mlx()
if _use_mlx:
from sglang.srt.hardware_backend.mlx.model_runner_stub import (
MlxModelRunnerStub,
)
model_runner = MlxModelRunnerStub(**runner_kwargs)
else:
model_runner = ModelRunner(**runner_kwargs)
rank_print(f"max_total_num_tokens={model_runner.max_total_num_tokens}")
tokenizer = get_tokenizer(
server_args.tokenizer_path,
tokenizer_mode=server_args.tokenizer_mode,
trust_remote_code=server_args.trust_remote_code,
)
if server_args.tp_size > 1:
dist.barrier()
if _use_mlx:
model_runner = _MlxBenchRunner(model_runner, server_args)
else:
model_runner = _TorchBenchRunner(model_runner)
return model_runner, tokenizer
def prepare_inputs_for_correctness_test(bench_args, tokenizer, custom_prompts):
if custom_prompts:
custom_input_len = len(custom_prompts)
bs = bench_args.batch_size[0]
if custom_input_len > bs:
logging.warning(
f"Custom input size ({custom_input_len}) is larger than batch_size ({bs}). "
f"Using the first {bs} prompts."
)
custom_prompts = custom_prompts[:bs]
prompts = (
custom_prompts
if custom_prompts
else [
"The capital of France is",
"The capital of the United Kindom is",
"Today is a sunny day and I like",
]
)
input_ids = [tokenizer.encode(p) for p in prompts]
sampling_params = SamplingParams(
temperature=0,
max_new_tokens=BenchArgs.output_len,
)
reqs = []
for i in range(len(prompts)):
assert len(input_ids[i]) > bench_args.cut_len
tmp_input_ids = input_ids[i][: bench_args.cut_len]
req = Req(
rid=i,
origin_input_text=prompts[i],
origin_input_ids=tmp_input_ids,
sampling_params=sampling_params,
)
req.fill_ids = req.origin_input_ids
req.logprob_start_len = -1
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
reqs.append(req)
return input_ids, reqs
def prepare_extend_inputs_for_correctness_test(
bench_args, input_ids, reqs, model_runner
):
for i in range(len(reqs)):
req: Req = reqs[i]
req.fill_ids += input_ids[i][bench_args.cut_len :]
if model_runner is not None:
req.prefix_indices = model_runner.req_to_token_pool.req_to_token[
i, : bench_args.cut_len
].to(req.prefix_indices.dtype)
req.logprob_start_len = -1
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
return reqs
def prepare_synthetic_inputs_for_latency_test(
batch_size, input_len, custom_inputs=None
):
input_ids = (
custom_inputs
if custom_inputs
else np.random.randint(0, 10000, (batch_size, input_len), dtype=np.int32)
)
sampling_params = SamplingParams(
temperature=0,
max_new_tokens=BenchArgs.output_len,
)
reqs = []
for i in range(len(input_ids)):
req = Req(
rid=i,
origin_input_text="",
origin_input_ids=list(input_ids[i]),
sampling_params=sampling_params,
)
req.fill_ids = req.origin_input_ids
req.logprob_start_len = -1
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
reqs.append(req)
return reqs
class TreeCacheNamespace(SimpleNamespace):
def supports_swa(self) -> bool:
return False
def supports_mamba(self) -> bool:
return False
def is_chunk_cache(self) -> bool:
return False
def is_tree_cache(self) -> bool:
return not self.is_chunk_cache()
def evict(self, params: EvictParams):
pass
@torch.no_grad
def extend(reqs, model_runner):
# Create dummy tree_cache for benchmarks (no prefix caching, just allocation)
dummy_tree_cache = TreeCacheNamespace(
page_size=model_runner.server_args.page_size,
device=model_runner.device,
token_to_kv_pool_allocator=model_runner.token_to_kv_pool_allocator,
)
batch = ScheduleBatch.init_new(
reqs=reqs,
req_to_token_pool=model_runner.req_to_token_pool,
token_to_kv_pool_allocator=model_runner.token_to_kv_pool_allocator,
tree_cache=dummy_tree_cache,
model_config=model_runner.model_config,
enable_overlap=False,
spec_algorithm=SpeculativeAlgorithm.NONE,
)
batch.prepare_for_extend()
_maybe_prepare_mlp_sync_batch(batch, model_runner)
model_worker_batch = batch.get_model_worker_batch()
forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
logits_output = model_runner.forward(forward_batch).logits_output
next_token_ids = model_runner.sample(logits_output, forward_batch)
return next_token_ids, logits_output.next_token_logits, batch
@torch.no_grad
def decode(input_token_ids, batch, model_runner):
batch.output_ids = input_token_ids
batch.prepare_for_decode()
_maybe_prepare_mlp_sync_batch(batch, model_runner)
model_worker_batch = batch.get_model_worker_batch()
forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
logits_output = model_runner.forward(forward_batch).logits_output
next_token_ids = model_runner.sample(logits_output, forward_batch)
return next_token_ids, logits_output.next_token_logits
def _maybe_prepare_mlp_sync_batch(batch: ScheduleBatch, model_runner):
if require_mlp_sync(model_runner.server_args):
prepare_mlp_sync_batch_raw(
batch,
dp_size=model_runner.server_args.dp_size,
attn_tp_size=get_attention_tp_size(),
attn_cp_size=model_runner.attn_cp_size,
tp_group=model_runner.tp_group,
get_idle_batch=None,
disable_cuda_graph=model_runner.server_args.disable_cuda_graph,
require_mlp_tp_gather=require_mlp_tp_gather(model_runner.server_args),
disable_overlap_schedule=model_runner.server_args.disable_overlap_schedule,
offload_tags=set(),
)
class _TorchBenchRunner:
"""Wraps ModelRunner for the standard PyTorch benchmark path."""
def __init__(self, model_runner):
self.torch_runner = model_runner
def clear(self):
self.torch_runner.req_to_token_pool.clear()
self.torch_runner.token_to_kv_pool_allocator.clear()
def extend(self, reqs):
return extend(reqs, self.torch_runner)
def decode(self, next_token_ids, batch):
return decode(next_token_ids, batch, self.torch_runner)
def cleanup(self, batch):
pass
def synchronize(self):
synchronize(self.torch_runner.device)
def max_batch_size(self, input_len, output_len):
return self.torch_runner.max_total_num_tokens // (input_len + output_len)
class _MlxBenchRunner:
"""Wraps MlxModelRunner for the MLX benchmark path."""
def __init__(self, model_runner, server_args):
from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
self.mlx_runner = MlxModelRunner(
model_path=server_args.model_path,
trust_remote_code=server_args.trust_remote_code,
)
self.fake_torch_runner = model_runner
def clear(self):
self.mlx_runner.clear()
def extend(self, reqs):
req_ids = [str(req.rid) for req in reqs]
token_ids_list = [[int(t) for t in req.fill_ids] for req in reqs]
next_token_ids = self.mlx_runner.prefill_batch(req_ids, token_ids_list)
return torch.tensor(next_token_ids), None, req_ids
def decode(self, next_token_ids, req_ids):
next_token_ids = self.mlx_runner.decode_batch(req_ids)
return torch.tensor(next_token_ids), None
def cleanup(self, batch):
if isinstance(batch, list):
for req_id in batch:
self.mlx_runner.remove_request(req_id)
def synchronize(self):
pass
def max_batch_size(self, input_len, output_len):
return self.fake_torch_runner.max_total_num_tokens // (input_len + output_len)
def _read_prompts_from_file(prompt_file, rank_print):
"""Read custom prompts from the file specified by `--prompt-filename`."""
if not prompt_file:
return []
if not os.path.exists(prompt_file):
rank_print(
f"Custom prompt file {prompt_file} not found. Using default inputs..."
)
return []
with open(prompt_file, "r") as pf:
return pf.readlines()
def _get_torch_profiler_output_dir():
return os.environ.get("SGLANG_TORCH_PROFILER_DIR", "/tmp")
def _create_torch_profiler_filename(
profile_filename_prefix, batch_size, input_len, output_len, stage
):
output_dir = _get_torch_profiler_output_dir()
filename = f"{profile_filename_prefix}_batch{batch_size}_input{input_len}_output{output_len}_{stage}.trace.json.gz"
return os.path.join(output_dir, filename)
def _save_profile_trace_results(profiler, filename):
parent_dir = os.path.dirname(os.path.abspath(filename))
os.makedirs(parent_dir, exist_ok=True)
profiler.export_chrome_trace(filename)
print(
profiler.key_averages(group_by_input_shape=True).table(
sort_by="self_cpu_time_total"
)
)
def correctness_test(
server_args,
port_args,
bench_args,
gpu_id,
tp_rank,
):
# Configure the logger
configure_logger(server_args, prefix=f" TP{tp_rank}")
rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None
# Load the model
model_runner, tokenizer = load_model(server_args, port_args, gpu_id, tp_rank)
# Prepare inputs
custom_prompts = _read_prompts_from_file(bench_args.prompt_filename, rank_print)
input_ids, reqs = prepare_inputs_for_correctness_test(
bench_args, tokenizer, custom_prompts
)
rank_print(f"\n{input_ids=}\n")
if bench_args.cut_len > 0:
# Prefill
next_token_ids, next_token_logits, batch = model_runner.extend(reqs)
rank_print(f"prefill logits (first half): {next_token_logits} \n")
# Prepare extend inputs
torch_runner = getattr(model_runner, "torch_runner", None)
reqs = prepare_extend_inputs_for_correctness_test(
bench_args, input_ids, reqs, torch_runner
)
# Extend (prefill w/ KV cache)
next_token_ids, next_token_logits, batch = model_runner.extend(reqs)
rank_print(f"prefill logits (final): {next_token_logits} \n")
# Decode
output_ids = [input_ids[i] + [next_token_ids[i]] for i in range(len(input_ids))]
for _ in range(bench_args.output_len[0] - 1):
next_token_ids, _ = model_runner.decode(next_token_ids, batch)
next_token_ids_list = next_token_ids.tolist()
for i in range(len(reqs)):
output_ids[i].append(next_token_ids_list[i])
# Clean up
model_runner.cleanup(batch)
# Print output texts
for i in range(len(reqs)):
rank_print(f"========== Prompt {i} ==========")
rank_print(tokenizer.decode(output_ids[i]), "\n")
def synchronize(device):
torch.get_device_module(device).synchronize()
def latency_test_run_once(
run_name,
model_runner,
rank_print,
reqs,
batch_size,
input_len,
output_len,
log_decode_step,
profile,
profile_record_shapes,
profile_activities,
profile_filename_prefix,
profile_stage,
tp_rank,
profile_start_step=None,
profile_steps=None,
):
max_batch_size = model_runner.max_batch_size(input_len, output_len)
if batch_size > max_batch_size:
rank_print(
f"skipping ({batch_size}, {input_len}, {output_len}) due to max batch size limit"
)
return
model_runner.clear()
measurement_results = {
"run_name": run_name,
"batch_size": batch_size,
"input_len": input_len,
"output_len": output_len,
}
tot_latency = 0
profiler = None
enable_profile_prefill = profile and profile_stage in ["all", "prefill"]
if enable_profile_prefill:
profiler = start_profile(
profile_activities,
profile_record_shapes=profile_record_shapes,
rank_print=rank_print,
)
model_runner.synchronize()
tic = time.perf_counter()
next_token_ids, _, batch = model_runner.extend(reqs)
model_runner.synchronize()
prefill_latency = time.perf_counter() - tic
if enable_profile_prefill:
trace_filename = _create_torch_profiler_filename(
profile_filename_prefix, batch_size, input_len, output_len, "prefill"
)
stop_profile(
profiler,
profile_activities,
rank_print=rank_print,
save_trace=True,
trace_filename=trace_filename,
stage="prefill",
)
tot_latency += prefill_latency
throughput = input_len * batch_size / prefill_latency
rank_print(
f"Prefill. latency: {prefill_latency:6.5f} s, throughput: {throughput:9.2f} token/s"
)
measurement_results["prefill_latency"] = prefill_latency
measurement_results["prefill_throughput"] = throughput
decode_latencies = []
# Determine profiling start step and end step
profile_start = (
profile_start_step if profile_start_step is not None else (output_len // 2)
)
profile_end = profile_start + (profile_steps if profile_steps is not None else 1)
enable_profile_decode = profile and profile_stage in ["all", "decode"]
profiler = None
for i in range(output_len - 1):
model_runner.synchronize()
# Start profiler at the specified step
if enable_profile_decode and i == profile_start:
profiler = start_profile(
profile_activities,
profile_record_shapes=profile_record_shapes,
rank_print=rank_print,
)
tic = time.perf_counter()
next_token_ids, _ = model_runner.decode(next_token_ids, batch)
model_runner.synchronize()
latency = time.perf_counter() - tic
# Stop profiler after the specified number of steps
if enable_profile_decode and profiler is not None and i >= profile_end - 1:
trace_filename = _create_torch_profiler_filename(
profile_filename_prefix, batch_size, input_len, output_len, "decode"
)
stop_profile(
profiler,
profile_activities,
rank_print=rank_print,
save_trace=True,
trace_filename=trace_filename,
stage="decode",
)
profiler = None
tot_latency += latency
throughput = batch_size / latency
decode_latencies.append(latency)
if i < 5 or (log_decode_step > 0 and i % log_decode_step == 0):
rank_print(
f"Decode {i}. Batch size: {batch_size}, latency: {latency:6.5f} s, throughput: {throughput:9.2f} token/s"
)
# Record decode timing from 2nd output
if output_len > 1:
med_decode_latency = np.median(decode_latencies)
med_decode_throughput = batch_size / med_decode_latency
rank_print(
f"Decode. median latency: {med_decode_latency:6.5f} s, median throughput: {med_decode_throughput:9.2f} token/s"
)
measurement_results["median_decode_latency"] = med_decode_latency
measurement_results["median_decode_throughput"] = med_decode_throughput
throughput = (input_len + output_len) * batch_size / tot_latency
rank_print(
f"Total. latency: {tot_latency:6.3f} s, throughput: {throughput:9.2f} token/s"
)
measurement_results["total_latency"] = tot_latency
measurement_results["overall_throughput"] = throughput
model_runner.cleanup(batch)
return measurement_results
def latency_test(
server_args,
port_args,
bench_args,
gpu_id,
tp_rank,
):
initialize_moe_config(server_args)
initialize_fp8_gemm_config(server_args)
initialize_fp4_gemm_config(server_args)
# Set CPU affinity
if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"):
set_gpu_proc_affinity(
server_args.pp_size, server_args.tp_size, server_args.nnodes, tp_rank
)
# Configure the logger
configure_logger(server_args, prefix=f" TP{tp_rank}")
rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None
# Load the model
model_runner, tokenizer = load_model(server_args, port_args, gpu_id, tp_rank)
# Prepare inputs for warm up
reqs = prepare_synthetic_inputs_for_latency_test(
bench_args.batch_size[0], bench_args.input_len[0]
)
# Warm up
rank_print("Warmup ...")
latency_test_run_once(
bench_args.run_name,
model_runner,
rank_print,
reqs,
bench_args.batch_size[0],
bench_args.input_len[0],
min(32, bench_args.output_len[0]), # shorter decoding to speed up the warmup
log_decode_step=0,
profile=False,
profile_record_shapes=False,
profile_activities=("CPU", "GPU"),
profile_filename_prefix="",
profile_stage="all",
tp_rank=tp_rank,
profile_start_step=None,
profile_steps=None,
)
rank_print("Benchmark ...")
custom_inputs = _read_prompts_from_file(bench_args.prompt_filename, rank_print)
custom_inputs = [tokenizer.encode(p.strip()) for p in custom_inputs]
custom_input_len = len(custom_inputs)
# Run the sweep
result_list = []
for bs, il, ol in itertools.product(
bench_args.batch_size, bench_args.input_len, bench_args.output_len
):
bs_aligned_inputs = []
if custom_inputs:
if custom_input_len == bs:
bs_aligned_inputs = custom_inputs
elif custom_input_len > bs:
rank_print(
f"Custom input size ({custom_input_len}) is larger than batch_size ({bs}). "
f"Using the first {bs} prompts."
)
bs_aligned_inputs = copy.deepcopy(custom_inputs[:bs])
else:
rank_print(
f"Custom input size ({custom_input_len}) is smaller than batch_size ({bs}). "
f"Pad to the desired batch_size with the last prompt."
)
bs_aligned_inputs = copy.deepcopy(custom_inputs)
bs_aligned_inputs.extend(
[bs_aligned_inputs[-1]] * (bs - custom_input_len)
)
reqs = prepare_synthetic_inputs_for_latency_test(bs, il, bs_aligned_inputs)
ret = latency_test_run_once(
bench_args.run_name,
model_runner,
rank_print,
reqs,
bs,
il,
ol,
bench_args.log_decode_step,
bench_args.profile if tp_rank == 0 else None,
bench_args.profile_record_shapes if tp_rank == 0 else None,
bench_args.profile_activities,
bench_args.profile_filename_prefix,
bench_args.profile_stage,
tp_rank,
bench_args.profile_start_step,
bench_args.profile_steps,
)
if ret is not None:
result_list.append(ret)
# Write results in jsonlines format on rank 0.
if tp_rank == 0 and bench_args.result_filename:
with open(bench_args.result_filename, "a") as fout:
for result in result_list:
fout.write(json.dumps(result) + "\n")
if server_args.tp_size > 1:
destroy_distributed_environment()
def main(server_args, bench_args):
server_args.cuda_graph_max_bs = max(bench_args.batch_size)
_set_envs_and_config(server_args)
if server_args.model_path:
if bench_args.correctness_test:
work_func = correctness_test
else:
work_func = latency_test
else:
raise ValueError(
"Provide --model-path for running the tests or "
"provide --result-filename for plotting the results"
)
port_args = PortArgs.init_new(server_args)
if server_args.tp_size == 1:
work_func(server_args, port_args, bench_args, 0, 0)
else:
workers = []
for tp_rank in range(server_args.tp_size):
with maybe_reindex_device_id(tp_rank) as gpu_id:
proc = multiprocessing.Process(
target=work_func,
args=(
server_args,
port_args,
bench_args,
gpu_id,
tp_rank,
),
)
proc.start()
workers.append(proc)
for proc in workers:
proc.join()
proc.terminate()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
BenchArgs.add_cli_args(parser)
args = parser.parse_args()
server_args = ServerArgs.from_cli_args(args)
bench_args = BenchArgs.from_cli_args(args)
logging.basicConfig(
level=getattr(logging, server_args.log_level.upper()),
format="%(message)s",
)
try:
main(server_args, bench_args)
finally:
if server_args.tp_size != 1:
kill_process_tree(os.getpid(), include_parent=False)

49
third_party/sglang/python/sglang/bench_one_batch_server.py vendored Normal file
View File

@@ -0,0 +1,49 @@
"""
Benchmark the latency of running a single batch with a server.
This script launches a server and uses the HTTP interface.
It accepts server arguments (the same as launch_server.py) and benchmark arguments (e.g., batch size, input lengths).
Usage:
python3 -m sglang.bench_one_batch_server --model meta-llama/Meta-Llama-3.1-8B --batch-size 1 16 64 --input-len 1024 --output-len 8
python3 -m sglang.bench_one_batch_server --model None --base-url http://localhost:30000 --batch-size 16 --input-len 1024 --output-len 8
python3 -m sglang.bench_one_batch_server --model None --base-url http://localhost:30000 --batch-size 16 --input-len 1024 --output-len 8 --show-report --profile --profile-by-stage
python3 -m sglang.bench_one_batch_server --model None --base-url http://localhost:30000 --batch-size 16 --input-len 1024 --output-len 8 --output-path results.json --profile
"""
import argparse
from sglang.srt.server_args import ServerArgs
from sglang.test.bench_one_batch_server_internal import (
BenchArgs,
run_benchmark_internal,
)
from sglang.test.nightly_bench_utils import save_results_as_pydantic_models
def run_benchmark(server_args: ServerArgs, bench_args: BenchArgs):
results, server_info = run_benchmark_internal(server_args, bench_args)
# Save results as pydantic models in the JSON format
if bench_args.pydantic_result_filename:
save_results_as_pydantic_models(
results,
pydantic_result_filename=bench_args.pydantic_result_filename,
model_path=server_args.model_path,
server_args=bench_args.server_args_for_metrics,
)
return results, server_info
if __name__ == "__main__":
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
BenchArgs.add_cli_args(parser)
args = parser.parse_args()
server_args = ServerArgs.from_cli_args(args)
bench_args = BenchArgs.from_cli_args(args)
run_benchmark(server_args, bench_args)

2352
third_party/sglang/python/sglang/bench_serving.py vendored Normal file
View File

File diff suppressed because it is too large Load Diff

0
third_party/sglang/python/sglang/benchmark/__init__.py vendored Normal file
View File

23
third_party/sglang/python/sglang/benchmark/bench_utils.py vendored Normal file
View File

@@ -0,0 +1,23 @@
"""Triton do_bench/do_bench_cudagraph compatible wrapper using flashinfer.testing.bench_gpu_time."""
import numpy as np
from flashinfer.testing import bench_gpu_time
def run_bench(
fn,
use_cuda_graph: bool = True,
quantiles=(0.5, 0.2, 0.8),
warmup_ms: int = 25,
rep_ms: int = 100,
):
"""Returns (ms, min_ms, max_ms) or (median,) when quantiles=None."""
times = bench_gpu_time(
fn=fn,
use_cuda_graph=use_cuda_graph,
dry_run_time_ms=warmup_ms,
repeat_time_ms=rep_ms,
)
if quantiles is None:
return (float(np.median(times)),)
return tuple(float(np.percentile(times, q * 100)) for q in quantiles)

49
third_party/sglang/python/sglang/benchmark/datasets/__init__.py vendored Normal file
View File

@@ -0,0 +1,49 @@
from typing import Dict, Type
from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
from sglang.benchmark.datasets.custom import CustomDataset
from sglang.benchmark.datasets.generated_shared_prefix import (
GeneratedSharedPrefixDataset,
)
from sglang.benchmark.datasets.image import ImageDataset
from sglang.benchmark.datasets.longbench_v2 import LongBenchV2Dataset
from sglang.benchmark.datasets.mmmu import MMMUDataset
from sglang.benchmark.datasets.mooncake import MooncakeDataset
from sglang.benchmark.datasets.openai_dataset import OpenAIDataset
from sglang.benchmark.datasets.random import RandomDataset
from sglang.benchmark.datasets.sharegpt import ShareGPTDataset
DATASET_MAPPING: Dict[str, Type[BaseDataset]] = {
"sharegpt": ShareGPTDataset,
"custom": CustomDataset,
"openai": OpenAIDataset,
# TODO: "random" vs "random-ids" should be a flag (e.g. --random-source=sharegpt|integers),
# not two separate dataset names sharing the same class.
"random": RandomDataset,
"random-ids": RandomDataset,
"generated-shared-prefix": GeneratedSharedPrefixDataset,
"mmmu": MMMUDataset,
"image": ImageDataset,
"mooncake": MooncakeDataset,
"longbench_v2": LongBenchV2Dataset,
}
def get_dataset(args, tokenizer, model_id=None):
dataset_name = args.dataset_name
if dataset_name.startswith("random") and dataset_name not in DATASET_MAPPING:
dataset_name = "random-ids"
if dataset_name not in DATASET_MAPPING:
raise ValueError(f"Unknown dataset: {args.dataset_name}")
dataset_cls = DATASET_MAPPING[dataset_name]
dataset = dataset_cls.from_args(args)
return dataset.load(tokenizer=tokenizer, model_id=model_id)
__all__ = [
"DATASET_MAPPING",
"DatasetRow",
"get_dataset",
]

86
third_party/sglang/python/sglang/benchmark/datasets/common.py vendored Normal file
View File

@@ -0,0 +1,86 @@
import random
from abc import ABC, abstractmethod
from argparse import Namespace
from dataclasses import dataclass
from functools import lru_cache
from typing import Any, Dict, List, Optional
import numpy as np
ASSISTANT_SUFFIX = "Assistant:"
SHAREGPT_REPO_ID = "anon8231489123/ShareGPT_Vicuna_unfiltered"
SHAREGPT_FILENAME = "ShareGPT_V3_unfiltered_cleaned_split.json"
MOONCAKE_DATASET_URL = {
"mooncake": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl",
"conversation": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/conversation_trace.jsonl",
"synthetic": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/synthetic_trace.jsonl",
"toolagent": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/toolagent_trace.jsonl",
}
@dataclass
class DatasetRow:
prompt: Any
prompt_len: int
output_len: int
text_prompt_len: Optional[int] = None
vision_prompt_len: Optional[int] = None
image_data: Optional[List[str]] = None
timestamp: Optional[float] = None
routing_key: Optional[str] = None
extra_request_body: Optional[Dict[str, Any]] = None # Per-request API parameters
def __post_init__(self):
if self.text_prompt_len is None:
self.text_prompt_len = self.prompt_len
if self.vision_prompt_len is None:
self.vision_prompt_len = 0
if self.extra_request_body is None:
self.extra_request_body = {}
@dataclass
class BaseDataset(ABC):
@classmethod
@abstractmethod
def from_args(cls, args: Namespace) -> "BaseDataset": ...
@abstractmethod
def load(
self,
tokenizer: Any,
model_id: Optional[str] = None,
) -> List[DatasetRow]: ...
def compute_random_lens(full_len: int, range_ratio: float, num: int) -> List[int]:
# full_len=0 is valid for embedding benchmarks where no output tokens are generated
if full_len <= 0:
return [0] * num
return np.random.randint(
max(int(full_len * range_ratio), 1),
full_len + 1,
size=num,
).tolist()
@lru_cache(maxsize=1)
def get_available_tokens(tokenizer):
"""Get all available token ids from the tokenizer vocabulary."""
return list(tokenizer.get_vocab().values())
def gen_prompt(tokenizer, token_num):
"""Generate a random prompt of specified token length using tokenizer vocabulary."""
all_available_tokens = get_available_tokens(tokenizer)
selected_tokens = random.choices(all_available_tokens, k=token_num)
return tokenizer.decode(selected_tokens)
def gen_mm_prompt(tokenizer, image_pad_id, token_num):
"""Generate a random prompt of specified token length using tokenizer vocabulary."""
all_available_tokens = list(tokenizer.get_vocab().values())
if image_pad_id:
all_available_tokens.remove(image_pad_id)
selected_tokens = random.choices(all_available_tokens, k=token_num)
return tokenizer.decode(selected_tokens)

147
third_party/sglang/python/sglang/benchmark/datasets/custom.py vendored Normal file
View File

@@ -0,0 +1,147 @@
import json
import os
import random
from argparse import Namespace
from dataclasses import dataclass
from typing import List, Optional
import numpy as np
from transformers import PreTrainedTokenizerBase
from sglang.benchmark.datasets.common import (
ASSISTANT_SUFFIX,
BaseDataset,
DatasetRow,
)
from sglang.benchmark.utils import remove_suffix
@dataclass
class CustomDataset(BaseDataset):
dataset_path: str
num_requests: int
fixed_output_len: Optional[int]
context_len: Optional[int]
prompt_suffix: str
apply_chat_template: bool
@classmethod
def from_args(cls, args: Namespace) -> "CustomDataset":
assert not getattr(args, "tokenize_prompt", False)
return cls(
dataset_path=args.dataset_path,
num_requests=args.num_prompts,
fixed_output_len=args.sharegpt_output_len,
context_len=args.sharegpt_context_len,
prompt_suffix=args.prompt_suffix,
apply_chat_template=args.apply_chat_template,
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
return sample_custom_requests(
dataset_path=self.dataset_path,
num_requests=self.num_requests,
tokenizer=tokenizer,
fixed_output_len=self.fixed_output_len,
context_len=self.context_len,
prompt_suffix=self.prompt_suffix,
apply_chat_template=self.apply_chat_template,
)
def sample_custom_requests(
dataset_path: str,
num_requests: int,
tokenizer: PreTrainedTokenizerBase,
fixed_output_len: Optional[int] = None,
context_len: Optional[int] = None,
prompt_suffix: Optional[str] = "",
apply_chat_template=False,
) -> List[DatasetRow]:
"""
Sample requests from a custom JSONL dataset: supports 'content'/'value' as conversation keys.
"""
if fixed_output_len is not None and fixed_output_len < 4:
raise ValueError("output_len too small")
# Load the dataset
dataset = []
if not os.path.isfile(dataset_path):
raise FileNotFoundError(f"Dataset not found at {dataset_path}")
with open(dataset_path, "r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if line: # skip empty lines
try:
dataset.append(json.loads(line))
except json.JSONDecodeError:
continue # skip lines with JSON errors
# Filter out the conversations with less than 2 turns.
processed_dataset = []
for data in dataset:
convs = data.get("conversations", data.get("conversation", []))
if len(convs) >= 2:
user_turn = convs[0].get("content", convs[0].get("value", ""))
assist_turn = convs[1].get("content", convs[1].get("value", ""))
processed_dataset.append((user_turn, assist_turn))
dataset = processed_dataset
random.shuffle(dataset)
# Filter out sequences that are too long or too short
filtered_dataset: List[DatasetRow] = []
for i in range(len(dataset)):
if len(filtered_dataset) == num_requests:
break
# Tokenize the prompts and completions.
prompt = dataset[i][0]
if prompt_suffix:
prompt = (
remove_suffix(prompt, ASSISTANT_SUFFIX)
+ prompt_suffix
+ ASSISTANT_SUFFIX
)
if apply_chat_template:
prompt = tokenizer.apply_chat_template(
[{"role": "user", "content": prompt}],
add_generation_prompt=True,
tokenize=False,
return_dict=False,
)
if tokenizer.bos_token:
prompt = prompt.replace(tokenizer.bos_token, "")
prompt_token_ids = tokenizer.encode(prompt)
completion = dataset[i][1]
completion_token_ids = tokenizer.encode(completion)
prompt_len = len(prompt_token_ids)
output_len = (
len(completion_token_ids) if fixed_output_len is None else fixed_output_len
)
if prompt_len < 2 or output_len < 2:
# Prune too short sequences.
continue
if context_len and prompt_len + output_len > context_len:
# Prune too long sequences.
continue
filtered_dataset.append(
DatasetRow(
prompt=prompt,
prompt_len=prompt_len,
output_len=output_len,
)
)
print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}")
print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}")
return filtered_dataset

231
third_party/sglang/python/sglang/benchmark/datasets/generated_shared_prefix.py vendored Normal file
View File

@@ -0,0 +1,231 @@
import pickle
import random
import uuid
from argparse import Namespace
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import List
import numpy as np
from tqdm.asyncio import tqdm
from transformers import PreTrainedTokenizerBase
from sglang.benchmark.datasets.common import (
BaseDataset,
DatasetRow,
compute_random_lens,
gen_prompt,
)
@dataclass
class GeneratedSharedPrefixDataset(BaseDataset):
num_groups: int
prompts_per_group: int
system_prompt_len: int
question_len: int
output_len: int
range_ratio: float
seed: int
fast_prepare: bool
send_routing_key: bool
num_turns: int
ordered: bool
@classmethod
def from_args(cls, args: Namespace) -> "GeneratedSharedPrefixDataset":
assert not getattr(args, "tokenize_prompt", False)
return cls(
num_groups=args.gsp_num_groups,
prompts_per_group=args.gsp_prompts_per_group,
system_prompt_len=args.gsp_system_prompt_len,
question_len=args.gsp_question_len,
output_len=args.gsp_output_len,
range_ratio=getattr(args, "gsp_range_ratio", 1.0),
seed=args.seed,
fast_prepare=getattr(args, "gsp_fast_prepare", False),
send_routing_key=getattr(args, "gsp_send_routing_key", False),
num_turns=getattr(args, "gsp_num_turns", 1),
ordered=getattr(args, "gsp_ordered", False),
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
return sample_generated_shared_prefix_requests(
num_groups=self.num_groups,
prompts_per_group=self.prompts_per_group,
system_prompt_len=self.system_prompt_len,
question_len=self.question_len,
output_len=self.output_len,
range_ratio=self.range_ratio,
tokenizer=tokenizer,
seed=self.seed,
send_routing_key=self.send_routing_key,
num_turns=self.num_turns,
fast_prepare=self.fast_prepare,
ordered=self.ordered,
)
def get_gen_prefix_cache_path(
seed: int,
num_groups: int,
prompts_per_group: int,
system_prompt_len: int,
question_len: int,
output_len: int,
tokenizer,
):
"""Create cache directory under ~/.cache/sglang/benchmark"""
cache_dir = Path.home() / ".cache" / "sglang" / "benchmark"
cache_key = (
f"gen_shared_prefix_{seed}_{num_groups}_{prompts_per_group}_"
f"{system_prompt_len}_{question_len}_{output_len}_"
f"{tokenizer.__class__.__name__}.pkl"
)
return cache_dir / cache_key
def sample_generated_shared_prefix_requests(
num_groups: int,
prompts_per_group: int,
system_prompt_len: int,
question_len: int,
output_len: int,
range_ratio: float,
tokenizer: PreTrainedTokenizerBase,
seed: int,
send_routing_key: bool = False,
num_turns: int = 1,
fast_prepare: bool = False,
ordered: bool = False,
) -> List[DatasetRow]:
"""Generate benchmark requests with shared system prompts using random tokens and caching."""
cache_path = get_gen_prefix_cache_path(
seed,
num_groups,
prompts_per_group,
system_prompt_len,
question_len,
output_len,
tokenizer,
)
should_cache = (range_ratio == 1) and not send_routing_key and num_turns == 1
# Try to load from cache first
if cache_path.exists() and should_cache:
print(f"\nLoading cached generated input data from {cache_path}")
with open(cache_path, "rb") as f:
return pickle.load(f)
print(
f"\nGenerating new input data... "
f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=}, {num_turns=})"
)
run_random_str = uuid.uuid4().hex[:8]
run_start_timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
system_prompt_lens = compute_random_lens(
full_len=system_prompt_len,
range_ratio=range_ratio,
num=num_groups,
)
question_lens = np.array(
compute_random_lens(
full_len=question_len,
range_ratio=range_ratio,
num=num_groups * prompts_per_group * num_turns,
)
).reshape(num_groups, prompts_per_group, num_turns)
output_lens = np.array(
compute_random_lens(
full_len=output_len,
range_ratio=range_ratio,
num=num_groups * prompts_per_group,
)
).reshape(num_groups, prompts_per_group)
del system_prompt_len, question_len, output_len
# Generate system prompts for each group
system_prompts = [
gen_prompt(tokenizer, system_prompt_lens[i]) for i in range(num_groups)
]
# Generate questions: shape (num_groups, prompts_per_group, num_turns)
questions = [
[
[
gen_prompt(tokenizer, int(question_lens[g, p, t]))
for t in range(num_turns)
]
for p in range(prompts_per_group)
]
for g in range(num_groups)
]
# Combine system prompts with questions
input_requests = []
total_input_tokens = 0
total_output_tokens = 0
for group_idx in tqdm(range(num_groups), desc="Generating system prompt"):
system_prompt = system_prompts[group_idx]
routing_key = (
f"{run_random_str}_{run_start_timestamp}_{group_idx}"
if send_routing_key
else None
)
for prompt_idx in tqdm(
range(prompts_per_group), desc="Generating questions", leave=False
):
turn_questions = questions[group_idx][prompt_idx]
turn_prompts = [f"{system_prompt}\n\n{turn_questions[0]}"] + turn_questions[
1:
]
full_prompt = turn_prompts[0] if num_turns == 1 else turn_prompts
prompt_len = 1 if fast_prepare else len(tokenizer.encode(turn_prompts[0]))
output_len_val = int(output_lens[group_idx, prompt_idx])
input_requests.append(
DatasetRow(
prompt=full_prompt,
prompt_len=prompt_len,
output_len=output_len_val,
routing_key=routing_key,
)
)
total_input_tokens += prompt_len
total_output_tokens += output_len_val
if not ordered:
random.shuffle(input_requests)
# Print statistics
print(f"\nGenerated shared prefix dataset statistics:")
print(f"Number of groups: {num_groups}")
print(f"Prompts per group: {prompts_per_group}")
print(f"Number of turns: {num_turns}")
print(f"Total prompts: {len(input_requests)}")
if not fast_prepare:
print(f"Total input tokens: {total_input_tokens}")
print(f"Total output tokens: {total_output_tokens}")
print(
f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens"
)
all_questions = [q for group in questions for conv in group for q in conv]
print(
f"Average question length: {sum(len(tokenizer.encode(q)) for q in all_questions) / len(all_questions):.1f} tokens\n"
)
# Save to cache
if should_cache:
cache_path.parent.mkdir(parents=True, exist_ok=True)
print(f"Caching generated input data to {cache_path}")
with open(cache_path, "wb") as f:
pickle.dump(input_requests, f)
return input_requests

291
third_party/sglang/python/sglang/benchmark/datasets/image.py vendored Normal file
View File

@@ -0,0 +1,291 @@
import io
import warnings
from argparse import Namespace
from dataclasses import dataclass
from typing import List, Tuple
import numpy as np
import pybase64
from PIL import Image
from transformers import AutoProcessor
from sglang.benchmark.datasets.common import (
BaseDataset,
DatasetRow,
compute_random_lens,
gen_mm_prompt,
)
from sglang.benchmark.utils import get_processor
@dataclass
class ImageDataset(BaseDataset):
num_requests: int
image_count: int
input_len: int
output_len: int
range_ratio: float
image_content: str
image_format: str
image_resolution: str
backend: str
random_image_count: bool
@classmethod
def from_args(cls, args: Namespace) -> "ImageDataset":
return cls(
num_requests=args.num_prompts,
image_count=args.image_count,
input_len=args.random_input_len,
output_len=args.random_output_len,
range_ratio=args.random_range_ratio,
image_content=args.image_content,
image_format=args.image_format,
image_resolution=args.image_resolution,
backend=args.backend,
random_image_count=args.random_image_count,
)
def load(self, tokenizer=None, model_id=None) -> List[DatasetRow]:
processor = get_processor(model_id)
return sample_image_requests(
num_requests=self.num_requests,
image_count=self.image_count,
input_len=self.input_len,
output_len=self.output_len,
range_ratio=self.range_ratio,
processor=processor,
image_content=self.image_content,
image_format=self.image_format,
image_resolution=self.image_resolution,
backend=self.backend,
random_image_count=self.random_image_count,
)
def parse_image_resolution(image_resolution: str) -> Tuple[int, int]:
"""Parse image resolution into (width, height).
Supports presets '1080p', '720p', '360p' and custom 'heightxwidth' format
(e.g., '1080x1920' means height=1080, width=1920).
"""
resolution_to_size = {
"4k": (3840, 2160),
"1080p": (1920, 1080),
"720p": (1280, 720),
"360p": (640, 360),
}
if image_resolution in resolution_to_size:
return resolution_to_size[image_resolution]
res = image_resolution.strip().lower()
if "x" in res:
parts = res.split("x")
if len(parts) == 2 and parts[0].isdigit() and parts[1].isdigit():
height = int(parts[0])
width = int(parts[1])
if height > 0 and width > 0:
return (width, height)
raise ValueError(
f"Unsupported image resolution: {image_resolution}. "
"Choose from 4k, 1080p, 720p, 360p, or provide custom 'heightxwidth' (e.g., 1080x1920)."
)
def create_mm_data_row(
text_prompt, images: list, images_base64, output_len, processor, backend
):
try:
if type(processor).__name__ == "Phi4MMProcessor":
# <|endoftext10|> is the image token used in the phi-4-multimodal model.
content_items = text_prompt.replace("image 1", "|endoftext10|")
else:
content_items = [
{"type": "image", "image": {"url": image_base64}}
for image_base64 in images_base64
]
content_items.append({"type": "text", "text": text_prompt})
prompt_str = processor.apply_chat_template(
[{"role": "user", "content": content_items}],
add_generation_prompt=True,
tokenize=False,
)
except Exception as e:
# Note (Xinyuan): This is a workaround for an issue where some tokenizers do not support content as a list. (e.g. InternVL)
print(f"Error applying chat template: {e}, fallback to <image> tag")
# Some tokenizers do not support list content; fall back to a placeholder in the text
prompt_str = f"<image>{text_prompt}"
# Calculate total tokens (text + vision)
prompt_len = processor(
text=[prompt_str],
images=images,
padding=False,
return_tensors="pt",
)["input_ids"].numel()
# Calculate text-only tokens
try:
# Create text-only version of the prompt
text_only_prompt = processor.apply_chat_template(
[{"role": "user", "content": text_prompt}],
add_generation_prompt=True,
tokenize=False,
)
text_prompt_len = processor(
text=[text_only_prompt],
padding=False,
return_tensors="pt",
)["input_ids"].numel()
except Exception:
# Fallback: just tokenize the text prompt directly
tokenizer_to_use = (
processor.tokenizer if hasattr(processor, "tokenizer") else processor
)
text_prompt_len = len(tokenizer_to_use.encode(text_prompt))
# Vision tokens = total tokens - text tokens
vision_prompt_len = prompt_len - text_prompt_len
supported_backends = ["sglang", "sglang-native", "sglang-oai-chat"]
if backend not in supported_backends:
raise ValueError(
f"Image dataset only supports backends: {supported_backends}, "
f"got '{backend}'."
)
# sglang-oai-chat: server's chat handler applies chat template, so send raw text.
# sglang/sglang-native: /generate does not apply chat template, so send prompt_str
# which contains image placeholder tokens needed by the multimodal processor.
use_raw_prompt = backend == "sglang-oai-chat"
return DatasetRow(
prompt=text_prompt if use_raw_prompt else prompt_str,
prompt_len=prompt_len,
output_len=output_len,
text_prompt_len=text_prompt_len,
vision_prompt_len=vision_prompt_len,
image_data=images_base64,
)
def sample_image_requests(
num_requests: int,
image_count: int,
input_len: int,
output_len: int,
range_ratio: float,
processor: AutoProcessor,
image_content: str,
image_format: str,
image_resolution: str,
backend: str,
random_image_count: bool = False,
) -> List[DatasetRow]:
"""Generate requests with images.
- If ``random_image_count`` is True, each request includes a random number of images between 1 and ``image_count``.
- If ``random_image_count`` is False, each request includes exactly ``image_count`` images.
- Supported resolutions: 4k (3840x2160), 1080p (1920x1080), 720p (1280x720), 360p (640x360),
or custom 'heightxwidth' (e.g., 1080x1920).
- Text lengths follow the 'random' dataset sampling rule. ``prompt_len``
only counts text tokens and excludes image data.
"""
# Parse resolution (supports presets and 'heightxwidth')
width, height = parse_image_resolution(image_resolution)
# Determine image counts for each request
if random_image_count:
# Random number of images per request
image_counts = np.random.randint(1, image_count + 1, size=num_requests)
total_images = np.sum(image_counts)
else:
# Fixed number of images per request
image_counts = np.full(num_requests, image_count)
total_images = image_count * num_requests
# Check for potentially problematic combinations and warn user
if width * height >= 1920 * 1080 and total_images >= 100:
warnings.warn(
f"High resolution ({width}x{height}) with {total_images} total images "
f"may take a long time. Consider reducing resolution or image count.",
UserWarning,
stacklevel=2,
)
# Sample text lengths
input_lens = compute_random_lens(
full_len=input_len,
range_ratio=range_ratio,
num=num_requests,
)
output_lens = compute_random_lens(
full_len=output_len,
range_ratio=range_ratio,
num=num_requests,
)
def _gen_random_image_data_uri(
width: int = width, height: int = height
) -> Tuple[Image.Image, str, int]:
if image_content == "blank":
# Generate blank white image
arr = np.full((height, width, 3), 255, dtype=np.uint8)
else:
# Generate random colored image
arr = (np.random.rand(height, width, 3) * 255).astype(np.uint8)
img = Image.fromarray(arr)
buf = io.BytesIO()
img.save(buf, format=image_format, quality=85)
encoded = pybase64.b64encode(buf.getvalue()).decode("utf-8")
image_data = f"data:image/{image_format};base64,{encoded}"
image_bytes = len(image_data.encode("utf-8"))
return img, image_data, image_bytes
dataset: List[DatasetRow] = []
total_image_bytes = 0
for i in range(num_requests):
# Get the number of images for this request
request_image_count = int(image_counts[i])
# Generate text prompt
text_prompt = gen_mm_prompt(
processor.tokenizer,
processor.image_token_id if hasattr(processor, "image_token_id") else None,
int(input_lens[i]),
)
# Generate image list
images, images_base64, images_bytes = zip(
*[_gen_random_image_data_uri() for _ in range(request_image_count)]
)
total_image_bytes += sum(images_bytes)
data_row = create_mm_data_row(
text_prompt,
list(images),
list(images_base64),
int(output_lens[i]),
processor,
backend,
)
dataset.append(data_row)
# Print statistics
print(f"#Input tokens: {np.sum([x.prompt_len for x in dataset])}")
print(f"#Output tokens: {np.sum([x.output_len for x in dataset])}")
print(f"#Total images: {total_images}")
if random_image_count:
print(
f"#Images per request: min={np.min(image_counts)}, max={np.max(image_counts)}, mean={np.mean(image_counts):.2f}"
)
else:
print(f"#Images per request: {image_count} (fixed)")
print(
f"\nCreated {len(dataset)} {image_content} {image_format} images with average {total_image_bytes // num_requests} bytes per request"
)
return dataset

104
third_party/sglang/python/sglang/benchmark/datasets/longbench_v2.py vendored Normal file
View File

@@ -0,0 +1,104 @@
import random
from argparse import Namespace
from dataclasses import dataclass
from typing import List, Optional
from transformers import PreTrainedTokenizerBase
from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
LONGBENCH_V2_REPO_ID = "THUDM/LongBench-v2"
LONGBENCH_V2_DEFAULT_OUTPUT_LEN = 10 # answer letter + short explanation
def _format_prompt(example: dict) -> str:
return (
f"{example['context']}\n\n"
f"Question: {example['question']}\n"
f"A. {example['choice_A']}\n"
f"B. {example['choice_B']}\n"
f"C. {example['choice_C']}\n"
f"D. {example['choice_D']}\n"
f"Answer:"
)
@dataclass
class LongBenchV2Dataset(BaseDataset):
dataset_path: str
num_requests: int
fixed_output_len: Optional[int]
context_len: Optional[int]
@classmethod
def from_args(cls, args: Namespace) -> "LongBenchV2Dataset":
return cls(
dataset_path=args.dataset_path,
num_requests=args.num_prompts,
fixed_output_len=args.sharegpt_output_len,
context_len=args.sharegpt_context_len,
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
return sample_longbench_v2_requests(
dataset_path=self.dataset_path,
num_requests=self.num_requests,
tokenizer=tokenizer,
fixed_output_len=self.fixed_output_len,
context_len=self.context_len,
)
def sample_longbench_v2_requests(
dataset_path: str,
num_requests: int,
tokenizer: PreTrainedTokenizerBase,
fixed_output_len: Optional[int] = None,
context_len: Optional[int] = None,
) -> List[DatasetRow]:
output_len = (
fixed_output_len
if fixed_output_len is not None
else LONGBENCH_V2_DEFAULT_OUTPUT_LEN
)
# Load dataset
if dataset_path:
# Local file (parquet or JSON lines)
import pandas as pd
if dataset_path.endswith(".parquet"):
df = pd.read_parquet(dataset_path)
examples = df.to_dict(orient="records")
else:
import json
with open(dataset_path) as f:
examples = [json.loads(line) for line in f if line.strip()]
else:
from datasets import load_dataset
ds = load_dataset(LONGBENCH_V2_REPO_ID, split="train")
examples = list(ds)
random.shuffle(examples)
rows: List[DatasetRow] = []
for example in examples:
if len(rows) >= num_requests:
break
prompt = _format_prompt(example)
prompt_ids = tokenizer(prompt).input_ids
prompt_len = len(prompt_ids)
if context_len is not None and prompt_len + output_len > context_len:
continue
rows.append(
DatasetRow(prompt=prompt, prompt_len=prompt_len, output_len=output_len)
)
return rows

124
third_party/sglang/python/sglang/benchmark/datasets/mmmu.py vendored Normal file
View File

@@ -0,0 +1,124 @@
import io
import random
from argparse import Namespace
from dataclasses import dataclass
from typing import List, Optional
import pybase64
from datasets import load_dataset
from transformers import AutoProcessor, AutoTokenizer
from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
from sglang.benchmark.datasets.image import create_mm_data_row
from sglang.benchmark.utils import get_processor
@dataclass
class MMMUDataset(BaseDataset):
num_requests: int
backend: str
fixed_output_len: Optional[int]
@classmethod
def from_args(cls, args: Namespace) -> "MMMUDataset":
return cls(
num_requests=args.num_prompts,
backend=args.backend,
fixed_output_len=args.random_output_len,
)
def load(self, tokenizer=None, model_id=None) -> List[DatasetRow]:
processor = get_processor(model_id)
return sample_mmmu_requests(
num_requests=self.num_requests,
processor=processor,
backend=self.backend,
fixed_output_len=self.fixed_output_len,
)
def sample_mmmu_requests(
num_requests: int,
processor: AutoProcessor | AutoTokenizer,
backend: str = "sglang",
fixed_output_len: Optional[int] = None,
random_sample: bool = True,
) -> List[DatasetRow]:
"""
Sample requests from the MMMU dataset using HuggingFace datasets.
Args:
num_requests: Number of requests to sample.
fixed_output_len: If provided, use this fixed output length for all requests.
random_sample: Whether to randomly sample or take the first N.
Returns:
List of tuples (prompt, prompt_token_len, output_token_len).
"""
print("Loading MMMU dataset from HuggingFace...")
try:
print("Attempting to load MMMU Math dataset...")
mmmu_dataset = load_dataset("MMMU/MMMU", "Math", split="test")
print(
f"Successfully loaded MMMU Math dataset from HuggingFace with {len(mmmu_dataset)} examples"
)
except Exception as e:
print(f"Failed to load MMMU Math dataset: {e}")
raise ValueError(f"Failed to load MMMU dataset: {e}")
# Sample from the dataset
if len(mmmu_dataset) > num_requests:
if random_sample:
# Random sample
indices = random.sample(range(len(mmmu_dataset)), num_requests)
sample_dataset = mmmu_dataset.select(indices)
else:
# Take first N
sample_dataset = mmmu_dataset.select(
range(min(num_requests, len(mmmu_dataset)))
)
else:
print(f"Dataset has less than {num_requests} examples, using all examples")
sample_dataset = mmmu_dataset
print(f"Selected {len(sample_dataset)} examples for benchmarking")
# Create prompts
filtered_dataset = []
for i, example in enumerate(sample_dataset):
try:
# Extract image_1
image = example.get("image_1")
if image is not None:
if hasattr(image, "save"):
# Convert RGBA images to RGB before encoding
if image.mode == "RGBA":
image = image.convert("RGB")
# Encode image to base64 (save as PNG to support palette/alpha modes)
buffered = io.BytesIO()
image.save(buffered, format="PNG")
img_str = pybase64.b64encode(buffered.getvalue()).decode("utf-8")
image_data = f"data:image/png;base64,{img_str}"
else:
continue
# Extract the question
question = example.get("question")
# Construct the prompt
text_prompt = f"Question: {question}\n\nAnswer: "
output_len = fixed_output_len if fixed_output_len is not None else 256
data_row = create_mm_data_row(
text_prompt, [image], [image_data], output_len, processor, backend
)
filtered_dataset.append(data_row)
except Exception as e:
print(f"Error processing example {i}: {e}")
print(f"\nCreated {len(filtered_dataset)} MMMU prompts")
return filtered_dataset

123
third_party/sglang/python/sglang/benchmark/datasets/mooncake.py vendored Normal file
View File

@@ -0,0 +1,123 @@
import asyncio
import json
import os
import time
from argparse import Namespace
from dataclasses import dataclass
from typing import AsyncGenerator, Dict, List
from transformers import PreTrainedTokenizerBase
from sglang.benchmark.datasets.common import (
MOONCAKE_DATASET_URL,
BaseDataset,
DatasetRow,
)
from sglang.benchmark.utils import download_and_cache_file
@dataclass
class MooncakeDataset(BaseDataset):
dataset_path: str
mooncake_workload: str
num_requests: int
@classmethod
def from_args(cls, args: Namespace) -> "MooncakeDataset":
return cls(
dataset_path=args.dataset_path,
mooncake_workload=args.mooncake_workload,
num_requests=args.num_prompts,
)
def load(self, tokenizer=None, model_id=None) -> List[Dict]:
if not self.dataset_path:
local_path = os.path.join("/tmp", self.mooncake_workload + "_trace.jsonl")
else:
local_path = self.dataset_path
if not os.path.exists(local_path):
download_and_cache_file(
MOONCAKE_DATASET_URL[self.mooncake_workload], local_path
)
with open(local_path, "r") as f:
all_requests_data = [json.loads(line) for line in f if line.strip()]
return all_requests_data[: self.num_requests]
async def get_mooncake_request_over_time(
input_requests: List[Dict],
tokenizer: PreTrainedTokenizerBase,
slowdown_factor: float,
num_rounds: int,
) -> AsyncGenerator[DatasetRow, None]:
"""
An async generator that yields requests based on the timestamps in the Mooncake trace file,
with support for multi-round sessions.
"""
if not input_requests:
return
input_requests.sort(key=lambda r: r["timestamp"])
start_time = time.perf_counter()
trace_start_time_ms = input_requests[0]["timestamp"]
for record in input_requests:
# Calculate when this entire session should start
relative_arrival_time_s = (record["timestamp"] - trace_start_time_ms) / 1000.0
target_arrival_time_s = relative_arrival_time_s * slowdown_factor
current_elapsed_time_s = time.perf_counter() - start_time
sleep_duration_s = target_arrival_time_s - current_elapsed_time_s
if sleep_duration_s > 0:
await asyncio.sleep(sleep_duration_s)
# Once the session starts, generate all rounds for it as a burst
# This simulates a user engaging in a multi-turn conversation
# Base user query constructed from hash_ids
user_query_base = ""
hash_ids = record.get("hash_ids", [])
for hash_id in hash_ids:
user_query_base += f"{hash_id}" + " ".join(
["hi"] * 128
) # Shorter for multi-round
user_query_base += "Tell me a story based on this context."
output_len_per_round = record.get("output_length", 256)
chat_history = []
for i in range(num_rounds):
# Add user query for the current round
chat_history.append(
{"role": "user", "content": f"Round {i + 1}: {user_query_base}"}
)
# Form the full prompt from history
try:
full_prompt_text = tokenizer.apply_chat_template(
chat_history,
tokenize=False,
add_generation_prompt=True,
return_dict=False,
)
except Exception:
full_prompt_text = "\n".join(
[f"{msg['role']}: {msg['content']}" for msg in chat_history]
)
prompt_len = len(tokenizer.encode(full_prompt_text))
yield DatasetRow(
prompt=full_prompt_text,
prompt_len=prompt_len,
output_len=output_len_per_round,
)
# Add a placeholder assistant response for the next round's context
# We use a placeholder because we don't know the real response
placeholder_response = " ".join(["story"] * output_len_per_round)
chat_history.append({"role": "assistant", "content": placeholder_response})

113
third_party/sglang/python/sglang/benchmark/datasets/openai_dataset.py vendored Normal file
View File

@@ -0,0 +1,113 @@
import json
from argparse import Namespace
from dataclasses import dataclass
from typing import List, Optional
import numpy as np
from transformers import PreTrainedTokenizerBase
from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
@dataclass
class OpenAIDataset(BaseDataset):
dataset_path: str
num_requests: int
fixed_output_len: Optional[int]
@classmethod
def from_args(cls, args: Namespace) -> "OpenAIDataset":
return cls(
dataset_path=args.dataset_path,
num_requests=args.num_prompts,
fixed_output_len=args.sharegpt_output_len,
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
return sample_openai_requests(
dataset_path=self.dataset_path,
num_requests=self.num_requests,
tokenizer=tokenizer,
fixed_output_len=self.fixed_output_len,
)
def sample_openai_requests(
dataset_path: str,
num_requests: int,
tokenizer: PreTrainedTokenizerBase,
fixed_output_len: Optional[int] = None,
) -> List[DatasetRow]:
"""
Load OpenAI-compatible chat completion requests from a JSONL file.
Each line should be a JSON object with:
- "messages": list of {"role": str, "content": str}
- "max_tokens": int (used as output_len if fixed_output_len not set)
- "tools": optional list of tool definitions
- "temperature": optional temperature value
- "top_p": optional top_p value
- Other OpenAI API parameters are also extracted and passed through
"""
dataset = []
with open(dataset_path, "r") as f:
for line in f:
if num_requests > 0 and len(dataset) >= num_requests:
break
if line.strip():
try:
dataset.append(json.loads(line))
except json.JSONDecodeError:
# Skip invalid JSON lines
continue
# Fields that should NOT be passed through extra_request_body
# These are either handled separately or are metadata
# max_tokens is excluded because it's handled via output_len -> max_completion_tokens
# max_completion_tokens is also excluded to avoid conflicts
EXCLUDED_FIELDS = {"messages", "max_tokens", "max_completion_tokens", "model"}
filtered_dataset: List[DatasetRow] = []
for data in dataset:
messages = data.get("messages", [])
if not messages:
continue
# Use max_tokens from the request, or fall back to fixed_output_len
output_len = fixed_output_len or data.get("max_tokens", 256)
# Extract extra request body parameters (tools, temperature, top_p, etc.)
extra_body = {k: v for k, v in data.items() if k not in EXCLUDED_FIELDS}
# Calculate prompt length by applying chat template
# This includes the messages but not the tools
prompt_len = len(
tokenizer.apply_chat_template(
messages, tokenize=True, add_generation_prompt=True
)
)
# If tools are present, we need to add their token count
# Tools are sent as part of the request and count toward input tokens
if "tools" in extra_body:
# Encode tools as JSON string to estimate token count
tools_str = json.dumps(extra_body["tools"])
tools_tokens = len(tokenizer.encode(tools_str))
prompt_len += tools_tokens
# Pass messages list directly - bench_serving handles List[Dict] prompts
filtered_dataset.append(
DatasetRow(
prompt=messages,
prompt_len=prompt_len,
output_len=output_len,
extra_request_body=extra_body, # Store per-request parameters
)
)
print(f"Loaded {len(filtered_dataset)} OpenAI-format requests")
print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}")
print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}")
return filtered_dataset

167
third_party/sglang/python/sglang/benchmark/datasets/random.py vendored Normal file
View File

@@ -0,0 +1,167 @@
import json
import random
from argparse import Namespace
from dataclasses import dataclass
from typing import List
import numpy as np
from transformers import PreTrainedTokenizerBase
from sglang.benchmark.datasets.common import (
SHAREGPT_FILENAME,
SHAREGPT_REPO_ID,
BaseDataset,
DatasetRow,
compute_random_lens,
)
from sglang.benchmark.utils import download_and_cache_hf_file, is_file_valid_json
@dataclass
class RandomDataset(BaseDataset):
input_len: int
output_len: int
num_requests: int
range_ratio: float
dataset_path: str
return_text: bool
random_sample: bool
@classmethod
def from_args(cls, args: Namespace) -> "RandomDataset":
return cls(
input_len=args.random_input_len,
output_len=args.random_output_len,
num_requests=args.num_prompts,
range_ratio=args.random_range_ratio,
dataset_path=args.dataset_path,
return_text=not getattr(args, "tokenize_prompt", False),
random_sample=(args.dataset_name == "random"),
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
return sample_random_requests(
input_len=self.input_len,
output_len=self.output_len,
num_prompts=self.num_requests,
range_ratio=self.range_ratio,
tokenizer=tokenizer,
dataset_path=self.dataset_path,
random_sample=self.random_sample,
return_text=self.return_text,
)
def sample_random_requests(
input_len: int,
output_len: int,
num_prompts: int,
range_ratio: float,
tokenizer: PreTrainedTokenizerBase,
dataset_path: str,
random_sample: bool = True,
return_text: bool = True,
) -> List[DatasetRow]:
input_lens = compute_random_lens(
full_len=input_len,
range_ratio=range_ratio,
num=num_prompts,
)
output_lens = compute_random_lens(
full_len=output_len,
range_ratio=range_ratio,
num=num_prompts,
)
if return_text:
# Need to truncate input_len as server encode will add special token.
num_special_tokens = int(tokenizer.num_special_tokens_to_add())
for i in range(num_prompts):
input_lens[i] = max(1, input_lens[i] - num_special_tokens)
if random_sample:
# Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens
# Download sharegpt if necessary
if not is_file_valid_json(dataset_path):
dataset_path = download_and_cache_hf_file(
repo_id=SHAREGPT_REPO_ID,
filename=SHAREGPT_FILENAME,
)
# Load the dataset.
with open(dataset_path) as f:
dataset = json.load(f)
# Filter out the conversations with less than 2 turns.
dataset = [
data
for data in dataset
if len(data.get("conversations", data.get("conversation", []))) >= 2
]
# Only keep the first two turns of each conversation.
dataset = [
(
data.get("conversations", data.get("conversation", []))[0]["value"],
data.get("conversations", data.get("conversation", []))[1]["value"],
)
for data in dataset
]
# Shuffle the dataset.
random.shuffle(dataset)
# Filter out sequences that are too long or too short
input_requests: List[DatasetRow] = []
for data in dataset:
i = len(input_requests)
if i == num_prompts:
break
# Tokenize the prompts and completions.
prompt = data[0]
prompt_token_ids = tokenizer.encode(prompt)
prompt_len = len(prompt_token_ids)
# Skip empty prompt
if prompt_len == 0:
continue
if prompt_len > input_lens[i]:
input_ids = prompt_token_ids[: input_lens[i]]
else:
ratio = (input_lens[i] + prompt_len - 1) // prompt_len
input_ids = (prompt_token_ids * ratio)[: input_lens[i]]
input_content = input_ids
if return_text:
input_content = tokenizer.decode(input_content)
input_requests.append(
DatasetRow(
prompt=input_content,
prompt_len=input_lens[i],
output_len=output_lens[i],
)
)
else:
# Sample token ids from random integers. This can cause some NaN issues.
offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts)
input_requests = []
for i in range(num_prompts):
# Use int() to convert numpy.int64 to native Python int for JSON serialization
input_content = [
int((offsets[i] + i + j) % tokenizer.vocab_size)
for j in range(input_lens[i])
]
if return_text:
input_content = tokenizer.decode(input_content)
input_requests.append(
DatasetRow(
prompt=input_content,
prompt_len=input_lens[i],
output_len=output_lens[i],
)
)
print(f"#Input tokens: {np.sum(input_lens)}")
print(f"#Output tokens: {np.sum(output_lens)}")
return input_requests

151
third_party/sglang/python/sglang/benchmark/datasets/sharegpt.py vendored Normal file
View File

@@ -0,0 +1,151 @@
import json
import random
from argparse import Namespace
from dataclasses import dataclass
from typing import List, Optional
import numpy as np
from transformers import PreTrainedTokenizerBase
from sglang.benchmark.datasets.common import (
ASSISTANT_SUFFIX,
SHAREGPT_FILENAME,
SHAREGPT_REPO_ID,
BaseDataset,
DatasetRow,
)
from sglang.benchmark.utils import (
download_and_cache_hf_file,
is_file_valid_json,
remove_suffix,
)
@dataclass
class ShareGPTDataset(BaseDataset):
dataset_path: str
num_requests: int
fixed_output_len: Optional[int]
context_len: Optional[int]
prompt_suffix: str
apply_chat_template: bool
@classmethod
def from_args(cls, args: Namespace) -> "ShareGPTDataset":
assert not getattr(args, "tokenize_prompt", False)
return cls(
dataset_path=args.dataset_path,
num_requests=args.num_prompts,
fixed_output_len=args.sharegpt_output_len,
context_len=args.sharegpt_context_len,
prompt_suffix=args.prompt_suffix,
apply_chat_template=args.apply_chat_template,
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
return sample_sharegpt_requests(
dataset_path=self.dataset_path,
num_requests=self.num_requests,
tokenizer=tokenizer,
fixed_output_len=self.fixed_output_len,
context_len=self.context_len,
prompt_suffix=self.prompt_suffix,
apply_chat_template=self.apply_chat_template,
)
def sample_sharegpt_requests(
dataset_path: str,
num_requests: int,
tokenizer: PreTrainedTokenizerBase,
fixed_output_len: Optional[int] = None,
context_len: Optional[int] = None,
prompt_suffix: Optional[str] = "",
apply_chat_template=False,
) -> List[DatasetRow]:
if fixed_output_len is not None and fixed_output_len < 4:
raise ValueError("output_len too small")
# Download sharegpt if necessary
if not is_file_valid_json(dataset_path) and dataset_path == "":
dataset_path = download_and_cache_hf_file(
repo_id=SHAREGPT_REPO_ID,
filename=SHAREGPT_FILENAME,
)
# Load the dataset.
with open(dataset_path) as f:
dataset = json.load(f)
# Filter out the conversations with less than 2 turns.
dataset = [
data
for data in dataset
if len(data.get("conversations", data.get("conversation", []))) >= 2
]
# Only keep the first two turns of each conversation.
dataset = [
(
data.get("conversations", data.get("conversation", []))[0]["value"],
data.get("conversations", data.get("conversation", []))[1]["value"],
)
for data in dataset
]
# Shuffle the dataset.
random.shuffle(dataset)
# Filter out sequences that are too long or too short
filtered_dataset: List[DatasetRow] = []
for i in range(len(dataset)):
if len(filtered_dataset) == num_requests:
break
# Tokenize the prompts and completions.
prompt = dataset[i][0]
if prompt_suffix:
prompt = (
remove_suffix(prompt, ASSISTANT_SUFFIX)
+ prompt_suffix
+ ASSISTANT_SUFFIX
)
if apply_chat_template:
prompt = tokenizer.apply_chat_template(
[{"role": "user", "content": prompt}],
add_generation_prompt=True,
tokenize=False,
return_dict=False,
)
if tokenizer.bos_token:
prompt = prompt.replace(tokenizer.bos_token, "")
prompt_token_ids = tokenizer.encode(prompt)
completion = dataset[i][1]
completion_token_ids = tokenizer.encode(completion)
prompt_len = len(prompt_token_ids)
output_len = (
len(completion_token_ids) if fixed_output_len is None else fixed_output_len
)
if prompt_len < 2 or output_len < 2:
# Prune too short sequences.
continue
if context_len and prompt_len + output_len > context_len:
# Prune too long sequences.
continue
filtered_dataset.append(
DatasetRow(
prompt=prompt,
prompt_len=prompt_len,
output_len=output_len,
)
)
print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}")
print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}")
return filtered_dataset

159
third_party/sglang/python/sglang/benchmark/utils.py vendored Normal file
View File

@@ -0,0 +1,159 @@
import json
import os
import resource
from json import JSONDecodeError
from typing import Dict, List, Optional, Union
import requests
from tqdm.asyncio import tqdm
from transformers import (
AutoProcessor,
AutoTokenizer,
PreTrainedTokenizer,
PreTrainedTokenizerFast,
)
def remove_prefix(text: str, prefix: str) -> str:
return text[len(prefix) :] if text.startswith(prefix) else text
def remove_suffix(text: str, suffix: str) -> str:
return text[: -len(suffix)] if text.endswith(suffix) else text
def parse_custom_headers(header_list: List[str]) -> Dict[str, str]:
return {k: v for h in header_list for k, _, v in [h.partition("=")] if k and v}
def get_model(pretrained_model_name_or_path: str) -> str:
if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() == "true":
import huggingface_hub.constants
from modelscope import snapshot_download
model_path = snapshot_download(
model_id=pretrained_model_name_or_path,
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"],
)
return model_path
return pretrained_model_name_or_path
def get_tokenizer(
pretrained_model_name_or_path: str,
) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
assert (
pretrained_model_name_or_path is not None
and pretrained_model_name_or_path != ""
)
if pretrained_model_name_or_path.endswith(
".json"
) or pretrained_model_name_or_path.endswith(".model"):
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
return get_tokenizer(pretrained_model_name_or_path)
if pretrained_model_name_or_path is not None and not os.path.exists(
pretrained_model_name_or_path
):
pretrained_model_name_or_path = get_model(pretrained_model_name_or_path)
return AutoTokenizer.from_pretrained(
pretrained_model_name_or_path, trust_remote_code=True
)
def get_processor(
pretrained_model_name_or_path: str,
) -> AutoProcessor:
assert (
pretrained_model_name_or_path is not None
and pretrained_model_name_or_path != ""
)
if pretrained_model_name_or_path.endswith(
".json"
) or pretrained_model_name_or_path.endswith(".model"):
from sglang.srt.utils.hf_transformers_utils import get_processor
return get_processor(pretrained_model_name_or_path)
if pretrained_model_name_or_path is not None and not os.path.exists(
pretrained_model_name_or_path
):
pretrained_model_name_or_path = get_model(pretrained_model_name_or_path)
return AutoProcessor.from_pretrained(
pretrained_model_name_or_path, trust_remote_code=True
)
def download_and_cache_hf_file(
repo_id: str,
filename: str,
repo_type: str = "dataset",
):
"""Download a file from Hugging Face and cache it locally."""
from huggingface_hub import hf_hub_download
return hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type)
def download_and_cache_file(url: str, filename: Optional[str] = None):
"""Read and cache a file from a url."""
if filename is None:
filename = os.path.join("/tmp", url.split("/")[-1])
# Check if the cache file already exists
if is_file_valid_json(filename):
return filename
print(f"Downloading from {url} to {filename}")
# Stream the response to show the progress bar
response = requests.get(url, stream=True)
response.raise_for_status() # Check for request errors
# Total size of the file in bytes
total_size = int(response.headers.get("content-length", 0))
chunk_size = 1024 # Download in chunks of 1KB
# Use tqdm to display the progress bar
with open(filename, "wb") as f, tqdm(
desc=filename,
total=total_size,
unit="B",
unit_scale=True,
unit_divisor=1024,
) as bar:
for chunk in response.iter_content(chunk_size=chunk_size):
f.write(chunk)
bar.update(len(chunk))
return filename
def is_file_valid_json(path):
if not os.path.isfile(path):
return False
# TODO can fuse into the real file open later
try:
with open(path) as f:
json.load(f)
return True
except JSONDecodeError as e:
print(
f"{path} exists but json loading fails ({e=}), thus treat as invalid file"
)
return False
def set_ulimit(target_soft_limit=65535):
resource_type = resource.RLIMIT_NOFILE
current_soft, current_hard = resource.getrlimit(resource_type)
if current_soft < target_soft_limit:
try:
resource.setrlimit(resource_type, (target_soft_limit, current_hard))
except ValueError as e:
print(f"Fail to set RLIMIT_NOFILE: {e}")

592
third_party/sglang/python/sglang/check_env.py vendored Normal file
View File

@@ -0,0 +1,592 @@
"""Check environment configurations and dependency versions."""
import importlib.metadata
import os
import resource
import subprocess
import sys
from abc import abstractmethod
from collections import OrderedDict, defaultdict
import torch
from sglang.srt.utils import is_hip, is_mps, is_musa, is_npu
def is_cuda_v2():
return torch.version.cuda is not None
# List of packages to check versions
PACKAGE_LIST = [
"sglang",
"sglang-kernel",
"flashinfer_python",
"flashinfer_cubin",
"flashinfer_jit_cache",
"triton",
"transformers",
"torchao",
"numpy",
"aiohttp",
"fastapi",
"huggingface_hub",
"interegular",
"modelscope",
"orjson",
"outlines",
"packaging",
"psutil",
"pydantic",
"python-multipart",
"pyzmq",
"torchao",
"uvicorn",
"uvloop",
"vllm",
"xgrammar",
"openai",
"tiktoken",
"anthropic",
"litellm",
"torchcodec",
]
class BaseEnv:
"""Base class for environment check"""
def __init__(self):
self.package_list = PACKAGE_LIST
@abstractmethod
def get_info(self) -> dict:
"""
Get CUDA-related information if available.
"""
raise NotImplementedError
@abstractmethod
def get_topology(self) -> dict:
raise NotImplementedError
def get_package_versions(self) -> dict:
"""
Get versions of specified packages.
"""
versions = {}
for package in self.package_list:
package_name = package.split("==")[0].split(">=")[0].split("<=")[0]
try:
version = importlib.metadata.version(package_name)
versions[package_name] = version
except ModuleNotFoundError:
versions[package_name] = "Module Not Found"
return versions
def get_device_info(self):
"""
Get information about available GPU devices.
"""
devices = defaultdict(list)
capabilities = defaultdict(list)
for k in range(torch.cuda.device_count()):
devices[torch.cuda.get_device_name(k)].append(str(k))
capability = torch.cuda.get_device_capability(k)
capabilities[f"{capability[0]}.{capability[1]}"].append(str(k))
gpu_info = {}
for name, device_ids in devices.items():
gpu_info[f"GPU {','.join(device_ids)}"] = name
if len(capabilities) == 1:
# All GPUs have the same compute capability
cap, gpu_ids = list(capabilities.items())[0]
gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap
else:
# GPUs have different compute capabilities
for cap, gpu_ids in capabilities.items():
gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap
return gpu_info
def get_hypervisor_vendor(self) -> dict:
try:
output = subprocess.check_output(["lscpu"], text=True)
for line in output.split("\n"):
if "Hypervisor vendor:" in line:
return {"Hypervisor vendor:": line.split(":")[1].strip()}
return {}
except:
return {}
def get_ulimit_soft(self) -> dict:
ulimit_soft, _ = resource.getrlimit(resource.RLIMIT_NOFILE)
return {"ulimit soft": ulimit_soft}
def check_env(self):
"""
Check and print environment information.
"""
env_info = OrderedDict()
env_info["Python"] = sys.version.replace("\n", "")
env_info.update(self.get_info())
env_info["PyTorch"] = torch.__version__
env_info.update(self.get_package_versions())
env_info.update(self.get_topology())
env_info.update(self.get_hypervisor_vendor())
env_info.update(self.get_ulimit_soft())
for k, v in env_info.items():
print(f"{k}: {v}")
class GPUEnv(BaseEnv):
"""Environment checker for Nvidia GPU"""
def get_info(self):
cuda_info = {"CUDA available": torch.cuda.is_available()}
if cuda_info["CUDA available"]:
cuda_info.update(self.get_device_info())
cuda_info.update(self._get_cuda_version_info())
return cuda_info
def _get_cuda_version_info(self):
"""
Get CUDA version information.
"""
from torch.utils.cpp_extension import CUDA_HOME
cuda_info = {"CUDA_HOME": CUDA_HOME}
if CUDA_HOME and os.path.isdir(CUDA_HOME):
cuda_info.update(self._get_nvcc_info())
cuda_info.update(self._get_cuda_driver_version())
return cuda_info
def _get_nvcc_info(self):
"""
Get NVCC version information.
"""
from torch.utils.cpp_extension import CUDA_HOME
try:
nvcc = os.path.join(CUDA_HOME, "bin/nvcc")
nvcc_output = (
subprocess.check_output(f'"{nvcc}" -V', shell=True)
.decode("utf-8")
.strip()
)
return {
"NVCC": nvcc_output[
nvcc_output.rfind("Cuda compilation tools") : nvcc_output.rfind(
"Build"
)
].strip()
}
except subprocess.SubprocessError:
return {"NVCC": "Not Available"}
def _get_cuda_driver_version(self):
"""
Get CUDA driver version.
"""
from sglang.srt.utils.common import get_nvidia_driver_version_str
ver = get_nvidia_driver_version_str()
if ver is None:
return {"CUDA Driver Version": "Not Available"}
return {"CUDA Driver Version": ver}
def get_topology(self):
"""
Get GPU topology information.
"""
try:
result = subprocess.run(
["nvidia-smi", "topo", "-m"],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
check=True,
)
return {
"NVIDIA Topology": (
"\n" + result.stdout if result.returncode == 0 else None
)
}
except subprocess.SubprocessError:
return {}
class HIPEnv(BaseEnv):
"""Environment checker for ROCm/HIP"""
def get_info(self):
cuda_info = {"ROCM available": torch.cuda.is_available()}
if cuda_info["ROCM available"]:
cuda_info.update(self.get_device_info())
cuda_info.update(self._get_cuda_version_info())
return cuda_info
def _get_cuda_version_info(self):
from torch.utils.cpp_extension import ROCM_HOME as ROCM_HOME
cuda_info = {"ROCM_HOME": ROCM_HOME}
if ROCM_HOME and os.path.isdir(ROCM_HOME):
cuda_info.update(self._get_hipcc_info())
cuda_info.update(self._get_rocm_driver_version())
return cuda_info
def _get_hipcc_info(self):
from torch.utils.cpp_extension import ROCM_HOME
try:
hipcc = os.path.join(ROCM_HOME, "bin/hipcc")
hipcc_output = (
subprocess.check_output(f'"{hipcc}" --version', shell=True)
.decode("utf-8")
.strip()
)
return {
"HIPCC": hipcc_output[
hipcc_output.rfind("HIP version") : hipcc_output.rfind("AMD clang")
].strip()
}
except subprocess.SubprocessError:
return {"HIPCC": "Not Available"}
def _get_rocm_driver_version(self):
try:
output = subprocess.check_output(
[
"rocm-smi",
"--showdriverversion",
"--csv",
]
)
versions = set(output.decode().strip().split("\n"))
versions.discard("name, value")
ver = versions.pop()
ver = ver.replace('"Driver version", ', "").replace('"', "")
return {"ROCM Driver Version": ver}
except subprocess.SubprocessError:
return {"ROCM Driver Version": "Not Available"}
def get_topology(self):
try:
result = subprocess.run(
["rocm-smi", "--showtopotype"],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
check=True,
)
return {
"AMD Topology": "\n" + result.stdout if result.returncode == 0 else None
}
except subprocess.SubprocessError:
return {}
class NPUEnv(BaseEnv):
"""Environment checker for Ascend NPU"""
EXTRA_PACKAGE_LIST = [
"torch_npu",
"sgl-kernel-npu",
"deep_ep",
]
def __init__(self):
super().__init__()
self.package_list.extend(NPUEnv.EXTRA_PACKAGE_LIST)
def get_info(self):
cuda_info = {"NPU available": torch.npu.is_available()}
if cuda_info["NPU available"]:
cuda_info.update(self.get_device_info())
cuda_info.update(self._get_cann_version_info())
return cuda_info
def get_device_info(self):
"""
Get information about available NPUs.
Need to override due to torch_npu interface differences.
"""
devices = defaultdict(list)
for k in range(torch.npu.device_count()):
devices[torch.npu.get_device_name(k)].append(str(k))
npu_info = {}
for name, device_ids in devices.items():
npu_info[f"NPU {','.join(device_ids)}"] = name
return npu_info
def _get_cann_version_info(self):
cann_envs = ["ASCEND_TOOLKIT_HOME", "ASCEND_INSTALL_PATH"]
for var in cann_envs:
path = os.environ.get(var)
if path and os.path.exists(path):
CANN_HOME = path
break
else:
default_path = "/usr/local/Ascend/ascend-toolkit/latest"
CANN_HOME = default_path if os.path.exists(default_path) else None
if CANN_HOME:
npu_info = {"CANN_HOME": CANN_HOME}
npu_info.update(self._get_cann_info(CANN_HOME))
npu_info.update(self._get_ascend_driver_version())
return npu_info
else:
return {"CANN_HOME": "Not found"}
def _get_cann_info(self, CANN_HOME: str):
cann_info = {}
cann_version_file = os.path.join(CANN_HOME, "version.cfg")
if os.path.exists(cann_version_file):
with open(cann_version_file, "r", encoding="utf-8") as f:
f.readline() # discard first line comment in version.cfg
cann_info["CANN"] = f.readline().split("[")[1].split("]")[0]
else:
cann_info["CANN"] = "Not Available"
try:
bisheng = os.path.join(CANN_HOME, "compiler/ccec_compiler/bin/bisheng")
bisheng_output = (
subprocess.check_output([bisheng, "--version"]).decode("utf-8").strip()
)
cann_info["BiSheng"] = bisheng_output.split("\n")[0].strip()
except subprocess.SubprocessError:
cann_info["BiSheng"] = "Not Available"
return cann_info
def _get_ascend_driver_version(self):
try:
output = subprocess.check_output(
[
"npu-smi",
"info",
"-t",
"board",
"-i",
"0",
]
)
for line in output.decode().strip().split("\n"):
if "Software Version" in line:
version = line.split(":")[-1].strip()
break
else:
version = "Not Available"
return {"Ascend Driver Version": version}
except subprocess.SubprocessError:
return {"Ascend Driver Version": "Not Available"}
def get_topology(self):
try:
result = subprocess.run(
["npu-smi", "info", "-t", "topo"],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
check=True,
)
return {
"Ascend Topology": (
"\n" + result.stdout if result.returncode == 0 else None
)
}
except subprocess.SubprocessError:
return {}
class MUSAEnv(BaseEnv):
"""Environment checker for MThreads GPU"""
def get_info(self):
musa_info = {"MUSA available": torch.musa.is_available()}
if musa_info["MUSA available"]:
musa_info.update(self.get_device_info())
musa_info.update(self._get_musa_version_info())
return musa_info
def _get_musa_version_info(self):
"""
Get MUSA version information.
"""
from torch_musa.utils.musa_extension import MUSA_HOME
musa_info = {"MUSA_HOME": MUSA_HOME}
if MUSA_HOME and os.path.isdir(MUSA_HOME):
musa_info.update(self._get_mcc_info())
musa_info.update(self._get_musa_driver_version())
return musa_info
def _get_mcc_info(self):
"""
Get MCC version information.
"""
from torch_musa.utils.musa_extension import MUSA_HOME
try:
mcc = os.path.join(MUSA_HOME, "bin/mcc")
mcc_output = (
subprocess.check_output(f'"{mcc}" --version', shell=True)
.decode("utf-8")
.strip()
)
return {
"MCC": mcc_output[
mcc_output.rfind("mcc version") : mcc_output.rfind("Target")
].strip()
}
except subprocess.SubprocessError:
return {"MCC": "Not Available"}
def _get_musa_driver_version(self):
"""
Get MUSA driver version.
"""
try:
output = subprocess.check_output(
[
"mthreads-gmi",
"-q",
],
text=True,
)
driver_version = None
for line in output.splitlines():
if "Driver Version" in line:
driver_version = line.split(":", 1)[1].strip()
break
return {"MUSA Driver Version": driver_version}
except subprocess.SubprocessError:
return {"MUSA Driver Version": "Not Available"}
def get_topology(self):
"""
Get GPU topology information.
"""
try:
result = subprocess.run(
["mthreads-gmi", "topo", "-m"],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
check=True,
)
return {
"MTHREADS Topology": (
"\n" + result.stdout if result.returncode == 0 else None
)
}
except subprocess.SubprocessError:
return {}
class MPSEnv(BaseEnv):
"""Environment checker for Apple Silicon MPS"""
EXTRA_PACKAGE_LIST = ["mlx", "mlx-lm", "mlx-metal"]
def __init__(self):
super().__init__()
self.package_list.extend(MPSEnv.EXTRA_PACKAGE_LIST)
def get_info(self):
import platform
info = {"MPS available": torch.backends.mps.is_available()}
if not info["MPS available"]:
return info
info["macOS Version"] = platform.mac_ver()[0]
try:
info["macOS Build"] = subprocess.check_output(
["sw_vers", "-buildVersion"], text=True
).strip()
except Exception:
info["macOS Build"] = "Not Available"
for label, key in [
("Apple Silicon", "machdep.cpu.brand_string"),
("Unified Memory", "hw.memsize"),
("CPU Cores (Total)", "hw.ncpu"),
]:
try:
info[label] = subprocess.check_output(
["sysctl", "-n", key], text=True
).strip()
except Exception:
info[label] = "Not Available"
try:
mem_bytes = int(info["Unified Memory"])
info["Unified Memory"] = f"{mem_bytes / 1024**3:.1f} GB"
except Exception:
pass
for label, key in [
("CPU Cores (Performance)", "hw.perflevel0.logicalcpu"),
("CPU Cores (Efficiency)", "hw.perflevel1.logicalcpu"),
]:
try:
info[label] = subprocess.check_output(
["sysctl", "-n", key], text=True
).strip()
except Exception:
pass
# Single system_profiler call for both Metal support and GPU cores
info["Metal Support"] = "Not Available"
info["GPU Cores"] = "Not Available"
try:
sp = subprocess.check_output(
["system_profiler", "SPDisplaysDataType"], text=True
)
for line in sp.splitlines():
line = line.strip()
if "Metal Support" in line or "Metal Family" in line:
info["Metal Support"] = line.partition(":")[2].strip()
if "Total Number of Cores" in line:
info["GPU Cores"] = line.partition(":")[2].strip()
except Exception:
pass
return info
def get_topology(self):
return {}
if __name__ == "__main__":
if is_cuda_v2():
env = GPUEnv()
elif is_hip():
env = HIPEnv()
elif is_npu():
env = NPUEnv()
elif is_musa():
env = MUSAEnv()
elif is_mps():
env = MPSEnv()
env.check_env()

0
third_party/sglang/python/sglang/cli/__init__.py vendored Normal file
View File

33
third_party/sglang/python/sglang/cli/generate.py vendored Normal file
View File

@@ -0,0 +1,33 @@
import argparse
from sglang.cli.utils import get_is_diffusion_model, get_model_path
def generate(args, extra_argv):
# If help is requested, show generate subcommand help without requiring --model-path
if any(h in extra_argv for h in ("-h", "--help")):
from sglang.multimodal_gen.runtime.entrypoints.cli.generate import (
add_multimodal_gen_generate_args,
)
parser = argparse.ArgumentParser(description="SGLang Multimodal Generation")
add_multimodal_gen_generate_args(parser)
parser.parse_args(extra_argv)
return
model_path = get_model_path(extra_argv)
is_diffusion_model = get_is_diffusion_model(model_path)
if is_diffusion_model:
from sglang.multimodal_gen.runtime.entrypoints.cli.generate import (
add_multimodal_gen_generate_args,
generate_cmd,
)
parser = argparse.ArgumentParser(description="SGLang Multimodal Generation")
add_multimodal_gen_generate_args(parser)
parsed_args, unknown_args = parser.parse_known_args(extra_argv)
generate_cmd(parsed_args, unknown_args)
else:
raise Exception(
f"Generate subcommand is not yet supported for model: {model_path}"
)

457
third_party/sglang/python/sglang/cli/killall.py vendored Executable file
View File

@@ -0,0 +1,457 @@
#!/usr/bin/env python3
"""Kill SGLang processes on CUDA_VISIBLE_DEVICES GPUs (CI mode only).
Called at the start of every CI job to clean up orphaned processes from
previous (possibly cancelled) runs. Requires SGLANG_IS_IN_CI=true.
For local/non-CI usage, use scripts/killall_sglang.sh instead.
Usage:
python killall.py
Exit codes:
0 - Clean: all target GPUs have <10% memory usage after cleanup
1 - Dirty: GPU memory still >10% after cleanup, indicating stuck processes
or orphaned CUDA contexts that need a container restart
"""
import os
import re
import signal
import subprocess
import sys
import time
from pathlib import Path
# Constants
MEMORY_THRESHOLD_PCT = 10
# Patterns matching SGLang process command lines (equivalent to pgrep -f in killall_sglang.sh)
_SGLANG_PROCESS_PATTERNS = re.compile(
r"sglang::|sglang\.launch_server|sglang\.bench|sglang\.data_parallel|sglang\.srt|sgl_diffusion::"
)
# Boxed output helpers
_LOG_LINES = []
def _log(msg=""):
"""Buffer a line for boxed output."""
_LOG_LINES.append(msg)
def _flush_box(title, status=""):
"""Print all buffered lines inside a box, then clear buffer."""
lines = _LOG_LINES.copy()
_LOG_LINES.clear()
all_text = [title] + ([status] if status else []) + lines
width = max((len(line) for line in all_text), default=40) + 4
width = max(width, 60)
h_bar = "─" * (width - 2)
print(f"\n┌{h_bar}┐")
print(f"│ {title:<{width - 3}}│")
print(f"├{h_bar}┤")
for line in lines:
print(f"│ {line:<{width - 3}}│")
if status:
print(f"├{h_bar}┤")
print(f"│ {status:<{width - 3}}│")
print(f"└{h_bar}┘")
# nvidia-smi helpers
def _run_smi(query, query_type="gpu"):
"""Run nvidia-smi query and return raw CSV lines."""
flag = "--query-gpu" if query_type == "gpu" else "--query-compute-apps"
try:
out = subprocess.check_output(
["nvidia-smi", f"{flag}={query}", "--format=csv,noheader,nounits"],
text=True,
timeout=10,
)
return [line.strip() for line in out.strip().splitlines() if line.strip()]
except (subprocess.SubprocessError, FileNotFoundError):
return []
def _get_smi_version():
"""Return nvidia-smi driver version and GPU name, or None on failure."""
# Inline nvidia-smi query — killall.py runs before pip install, so sglang
# internals may not be importable.
try:
result = subprocess.run(
[
"nvidia-smi",
"--query-gpu=driver_version",
"--format=csv,noheader,nounits",
],
capture_output=True,
text=True,
check=True,
timeout=10,
)
driver = result.stdout.strip().split("\n")[0].strip() or None
except (subprocess.SubprocessError, FileNotFoundError):
driver = None
if driver is None:
return None
try:
out = subprocess.check_output(
["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
text=True,
timeout=10,
)
gpu_name = out.strip().splitlines()[0].strip() if out.strip() else "unknown"
except (subprocess.SubprocessError, FileNotFoundError, IndexError):
gpu_name = "unknown"
return f"driver {driver}, {gpu_name}"
def _get_target_gpus():
"""Return GPU indices from CUDA_VISIBLE_DEVICES, or all visible GPUs.
Note: only numeric indices are supported (e.g. "0,1,2").
UUID-style CUDA_VISIBLE_DEVICES values (e.g. "GPU-d4f1...") are not handled.
"""
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
if cvd is not None and cvd.strip():
return {int(g.strip()) for g in cvd.split(",") if g.strip().isdigit()}
return {int(line) for line in _run_smi("index") if line.isdigit()}
def _get_gpu_pids(gpu_indices):
"""Return PIDs using the specified GPUs (by index)."""
target_uuids = set()
for line in _run_smi("index,uuid"):
parts = line.split(",", 1)
if len(parts) == 2 and parts[0].strip().isdigit():
if int(parts[0].strip()) in gpu_indices:
target_uuids.add(parts[1].strip())
pids = set()
for line in _run_smi("gpu_uuid,pid", query_type="apps"):
parts = line.split(",", 1)
if len(parts) == 2 and parts[0].strip() in target_uuids:
pid = parts[1].strip()
if pid.isdigit():
pids.add(int(pid))
return pids
def _get_gpu_memory(gpu_indices):
"""Query memory usage for target GPUs.
Returns list of (idx, used_mib, total_mib, pct) tuples.
"""
result = []
for line in _run_smi("index,memory.used,memory.total"):
parts = line.split(",")
if len(parts) != 3 or not parts[0].strip().isdigit():
continue
idx = int(parts[0].strip())
if idx not in gpu_indices:
continue
try:
used, total = int(float(parts[1].strip())), int(float(parts[2].strip()))
except ValueError:
continue
pct = used / total * 100 if total > 0 else 0
result.append((idx, used, total, pct))
return result
def _get_dirty_gpus(gpu_indices):
"""Return list of dirty GPU description strings (memory >= threshold)."""
return [
f"GPU {idx} ({pct:.0f}%)"
for idx, _, _, pct in _get_gpu_memory(gpu_indices)
if pct >= MEMORY_THRESHOLD_PCT
]
def _log_gpu_memory(gpu_indices):
"""Log memory usage for all target GPUs and return dirty GPU descriptions."""
dirty = []
for idx, used, total, pct in _get_gpu_memory(gpu_indices):
_log(f" GPU {idx}: {used} MiB / {total} MiB ({pct:.0f}%)")
if pct >= MEMORY_THRESHOLD_PCT:
dirty.append(f"GPU {idx} ({pct:.0f}%)")
return dirty
# /proc helpers
def _read_proc_cmdline(pid):
"""Read /proc/{pid}/cmdline and return as decoded string, or None on failure."""
try:
raw = Path(f"/proc/{pid}/cmdline").read_bytes()
return raw.decode("utf-8", errors="replace").replace("\x00", " ")
except (FileNotFoundError, PermissionError):
return None
def _get_pid_cmdline(pid):
"""Get truncated command line for a PID."""
cmdline = _read_proc_cmdline(pid)
if cmdline is None:
return "<unknown>"
cmdline = cmdline.strip()
return cmdline[:120] + ("..." if len(cmdline) > 120 else "")
def _find_sglang_pids_by_name():
"""Find SGLang process PIDs by command-line pattern matching.
Scans /proc/*/cmdline for patterns matching known SGLang entry points.
Equivalent to: pgrep -f 'sglang::|sglang.launch_server|...'
Safe in shared-GPU containers: without --pid=host, /proc only exposes
processes in our own PID namespace, so this cannot kill other containers.
"""
my_pid = os.getpid()
pids = set()
for entry in Path("/proc").iterdir():
if not entry.name.isdigit():
continue
pid = int(entry.name)
if pid <= 1 or pid == my_pid:
continue
cmdline = _read_proc_cmdline(pid)
if cmdline and _SGLANG_PROCESS_PATTERNS.search(cmdline):
pids.add(pid)
return pids
def _check_pid_namespace(pid):
"""Check if a PID is in our PID namespace. Linux-only via /proc."""
try:
my_ns = os.readlink("/proc/self/ns/pid")
except OSError:
return "unknown (can't read self ns)"
try:
target_ns = os.readlink(f"/proc/{pid}/ns/pid")
except FileNotFoundError:
return f"NOT in our namespace (pid not in /proc, self={my_ns})"
except PermissionError:
return "unknown (no permission to read ns)"
if my_ns == target_ns:
return f"same namespace ({my_ns})"
return f"DIFFERENT namespace (self={my_ns}, target={target_ns})"
def _get_orchestrator_ancestors(pids):
"""Walk process tree upward from PIDs, return ancestors that are test orchestrators.
Linux-only: reads /proc filesystem. Returns empty set on other platforms.
"""
orchestrator_patterns = ["run_suite.py", "run_tests.py"]
ancestors, visited = set(), set()
for pid in pids:
current = pid
while current > 1 and current not in visited:
visited.add(current)
cmdline = _read_proc_cmdline(current)
if cmdline is None:
break
if any(p in cmdline for p in orchestrator_patterns):
ancestors.add(current)
try:
current = int(Path(f"/proc/{current}/stat").read_text().split()[3])
except (FileNotFoundError, PermissionError, IndexError, ValueError):
break
return ancestors
# Kill & diagnostic helpers
def _kill_pids(pids, label="", quiet=False):
"""Send SIGKILL to PIDs, skipping self and init.
Returns dict of {pid: exception_name} for PIDs that could not be killed.
When quiet=True, does not log individual kill results.
"""
my_pid = os.getpid()
pids = {p for p in pids if p != my_pid and p > 1}
if not pids:
return {}
if label and not quiet:
_log(f" Killing {label}:")
failed = {}
for pid in sorted(pids):
try:
os.kill(pid, signal.SIGKILL)
if not quiet:
_log(f" PID {pid}: killed ({_get_pid_cmdline(pid)})")
except (ProcessLookupError, PermissionError) as e:
failed[pid] = type(e).__name__
if not quiet:
_log(f" PID {pid}: failed ({type(e).__name__})")
return failed
def _get_ps_diagnostic():
"""Return ps auxf output filtered for GPU/sglang-related processes."""
try:
out = subprocess.run(["ps", "auxf"], capture_output=True, text=True, timeout=5)
return [
line.strip()[:140]
for line in out.stdout.splitlines()
if any(k in line.lower() for k in ["sglang", "python", "cuda", "gpu"])
][:20]
except (subprocess.SubprocessError, FileNotFoundError):
return []
def _print_diagnostics(unkillable_pids):
"""Print detailed diagnostics after the FAIL box (to stdout, outside box)."""
if unkillable_pids:
print("\n[killall] Diagnostic — unkillable PIDs:")
for pid in sorted(unkillable_pids):
ns_info = _check_pid_namespace(pid)
print(f" PID {pid}: ns: {ns_info}")
ps_lines = _get_ps_diagnostic()
if ps_lines:
print("\n[killall] Diagnostic — processes in this container (ps auxf):")
for line in ps_lines:
print(f" {line}")
else:
print(
"\n[killall] Diagnostic — no sglang/python/gpu processes "
"in this container"
)
# CI mode
def _kill_all_targets(gpu_indices, gpu_pids):
"""Kill all target processes: name-matched, orchestrator ancestors, GPU processes."""
# Kill name-matched SGLang processes (catches processes not visible to nvidia-smi)
name_only = _find_sglang_pids_by_name() - gpu_pids
if name_only:
_kill_pids(name_only, "name-matched SGLang processes")
time.sleep(1)
_log()
# Kill orchestrator ancestors first, then GPU processes (retry once)
if gpu_pids:
_kill_pids(_get_orchestrator_ancestors(gpu_pids), "orchestrator ancestors")
time.sleep(1)
for attempt in range(2):
current_pids = _get_gpu_pids(gpu_indices)
if not current_pids:
break
label = "GPU processes" if attempt == 0 else "stubborn GPU processes"
_kill_pids(current_pids, label)
time.sleep(3)
_log()
def _verify_gpu_clean(gpu_indices):
"""Retry loop: wait for GPUs to become clean.
Returns (dirty_list, unkillable_pids, elapsed_seconds).
"""
max_wait_secs = 100
retry_interval = 10
elapsed = 0
dirty = None
unkillable_pids = {}
while True:
dirty = _get_dirty_gpus(gpu_indices)
remaining_pids = _get_gpu_pids(gpu_indices)
if not dirty:
_log(f"Check at {elapsed}s: GPUs clean")
break
dirty_summary = ", ".join(dirty)
if elapsed >= max_wait_secs:
remaining_info = (
f", {len(remaining_pids)} processes remaining" if remaining_pids else ""
)
_log(f"Check at {elapsed}s: still dirty [{dirty_summary}]{remaining_info}")
break
# Kill remaining processes before waiting (silently for retries)
if remaining_pids:
failed = _kill_pids(remaining_pids, quiet=True)
unkillable_pids.update(failed)
print(
f"[killall] GPUs still dirty at {elapsed}s [{dirty_summary}], "
f"retrying in {retry_interval}s "
f"({elapsed + retry_interval}/{max_wait_secs}s)..."
)
time.sleep(retry_interval)
elapsed += retry_interval
if unkillable_pids:
parts = [f"{p} ({unkillable_pids[p]})" for p in sorted(unkillable_pids)]
_log(f" Unkillable PIDs: {', '.join(parts)}")
return dirty, unkillable_pids, elapsed
def _ci_mode():
"""GPU-scoped kill, abort if GPUs remain dirty."""
gpu_indices = _get_target_gpus()
if not gpu_indices:
_log("No GPUs detected, skipping cleanup")
_flush_box("killall_sglang", status="SKIP")
return 0
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
gpu_list = ", ".join(str(g) for g in sorted(gpu_indices))
smi_info = _get_smi_version()
if smi_info:
_log(f"nvidia-smi: {smi_info}")
if cvd is None or not cvd.strip():
_log(
"WARNING: CUDA_VISIBLE_DEVICES is not set. "
"Falling back to all visible GPUs."
)
_log("This may kill processes from other CI jobs on shared hosts.")
else:
_log(f"CUDA_VISIBLE_DEVICES={cvd}")
_log()
# Log pre-cleanup state
_log("Before cleanup:")
_log_gpu_memory(gpu_indices)
gpu_pids = _get_gpu_pids(gpu_indices)
if not gpu_pids:
_log(" No processes on target GPUs")
else:
_log(f" Processes ({len(gpu_pids)}):")
for pid in sorted(gpu_pids):
_log(f" PID {pid}: {_get_pid_cmdline(pid)}")
_log()
# Kill phase
_kill_all_targets(gpu_indices, gpu_pids)
# Verify phase
dirty, unkillable_pids, elapsed = _verify_gpu_clean(gpu_indices)
if dirty:
_log()
_log("Final GPU memory:")
_log_gpu_memory(gpu_indices)
_log(f"ERROR: memory >={MEMORY_THRESHOLD_PCT}%: {', '.join(dirty)}")
_log(f"Orphaned CUDA contexts after {elapsed}s — container needs restart.")
_flush_box(f"killall_sglang: GPUs [{gpu_list}]", status="FAIL — Aborting CI")
_print_diagnostics(unkillable_pids)
return 1
_flush_box(f"killall_sglang: GPUs [{gpu_list}]", status="PASS — GPUs clean")
return 0
# Entry point
def main():
return _ci_mode()
if __name__ == "__main__":
sys.exit(main())

46
third_party/sglang/python/sglang/cli/main.py vendored Normal file
View File

@@ -0,0 +1,46 @@
import argparse
from sglang.cli.utils import get_git_commit_hash
from sglang.version import __version__
def version(args, extra_argv):
print(f"sglang version: {__version__}")
print(f"git revision: {get_git_commit_hash()[:7]}")
def main():
parser = argparse.ArgumentParser()
# complex sub commands
subparsers = parser.add_subparsers(dest="subcommand", required=True)
subparsers.add_parser(
"serve",
help="Launch an SGLang server.",
add_help=False,
)
subparsers.add_parser(
"generate",
help="Run inference on a multimodal model.",
add_help=False,
)
# simple commands
version_parser = subparsers.add_parser(
"version",
help="Show the version information.",
)
version_parser.set_defaults(func=version)
args, extra_argv = parser.parse_known_args()
if args.subcommand == "serve":
from sglang.cli.serve import serve
serve(args, extra_argv)
elif args.subcommand == "generate":
from sglang.cli.generate import generate
generate(args, extra_argv)
elif args.subcommand == "version":
version(args, extra_argv)

126
third_party/sglang/python/sglang/cli/serve.py vendored Normal file
View File

@@ -0,0 +1,126 @@
# SPDX-License-Identifier: Apache-2.0
import argparse
import logging
import os
from sglang.cli.utils import get_is_diffusion_model, get_model_path
from sglang.srt.utils import kill_process_tree
from sglang.srt.utils.common import suppress_noisy_warnings
suppress_noisy_warnings()
logger = logging.getLogger(__name__)
def _extract_model_type_override(extra_argv):
"""Extract and remove --model-type override from argv."""
model_type = "auto"
filtered_argv = []
i = 0
while i < len(extra_argv):
arg = extra_argv[i]
if arg == "--model-type":
if i + 1 >= len(extra_argv):
raise Exception(
"Error: --model-type requires a value. "
"Valid values are: auto, llm, diffusion."
)
model_type = extra_argv[i + 1]
i += 2
continue
if arg.startswith("--model-type="):
model_type = arg.split("=", 1)[1]
i += 1
continue
filtered_argv.append(arg)
i += 1
if model_type not in ("auto", "llm", "diffusion"):
raise Exception(
f"Error: invalid --model-type '{model_type}'. "
"Valid values are: auto, llm, diffusion."
)
return model_type, filtered_argv
def serve(args, extra_argv):
if any(h in extra_argv for h in ("-h", "--help")):
# Since the server type is determined by the model, and we don't have a model path,
# we can't show the exact help. Instead, we show a general help message and then
# the help for both possible server types.
print(
"Usage: sglang serve --model-path <model-name-or-path> [additional-arguments]\n\n"
"This command can launch either a standard language model server or a diffusion model server.\n"
"The server type is determined by the --model-path.\n"
"Optional override: --model-type {auto,llm,diffusion} "
"(default: auto, fallback to LLM on detection failure)."
)
print("\n--- Help for Standard Language Model Server ---")
from sglang.srt.server_args import prepare_server_args
try:
prepare_server_args(["--help"])
except SystemExit:
pass # argparse --help calls sys.exit
print("\n--- Help for Diffusion Model Server ---")
try:
from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
add_multimodal_gen_serve_args,
)
parser = argparse.ArgumentParser(
prog="sglang serve",
description="SGLang Diffusion Model Serving",
)
add_multimodal_gen_serve_args(parser)
parser.print_help()
except ImportError:
print(
"Diffusion model support is not available. "
'Install with: pip install "sglang[diffusion]"'
)
return
model_type, dispatch_argv = _extract_model_type_override(extra_argv)
model_path = get_model_path(dispatch_argv)
try:
if model_type == "auto":
is_diffusion_model = get_is_diffusion_model(model_path)
if is_diffusion_model:
logger.info("Diffusion model detected")
else:
is_diffusion_model = model_type == "diffusion"
logger.info(
"Dispatch override enabled: --model-type=%s " "(skip auto detection)",
model_type,
)
if is_diffusion_model:
# Logic for Diffusion Models
from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
add_multimodal_gen_serve_args,
execute_serve_cmd,
)
parser = argparse.ArgumentParser(
description="SGLang Diffusion Model Serving"
)
add_multimodal_gen_serve_args(parser)
parsed_args, remaining_argv = parser.parse_known_args(dispatch_argv)
execute_serve_cmd(parsed_args, remaining_argv)
else:
# Logic for Standard Language Models
from sglang.launch_server import run_server
from sglang.srt.server_args import prepare_server_args
server_args = prepare_server_args(dispatch_argv)
run_server(server_args)
finally:
kill_process_tree(os.getpid(), include_parent=False)

146
third_party/sglang/python/sglang/cli/utils.py vendored Normal file
View File

@@ -0,0 +1,146 @@
import json
import logging
import os
import subprocess
from functools import lru_cache
from huggingface_hub import HfApi
from sglang.srt.environ import envs
from sglang.utils import (
has_diffusion_overlay_registry_match,
is_known_non_diffusers_diffusion_model,
load_diffusion_overlay_registry_from_env,
)
logger = logging.getLogger(__name__)
@lru_cache(maxsize=1)
def _load_overlay_registry() -> dict:
return load_diffusion_overlay_registry_from_env()
def _is_overlay_diffusion_model(model_path: str) -> bool:
return has_diffusion_overlay_registry_match(model_path, _load_overlay_registry())
def _is_registered_diffusion_model(model_path: str) -> bool:
try:
from sglang.multimodal_gen.registry import has_registered_diffusion_model_path
except ImportError:
# if diffusion dependencies are not installed
return False
return has_registered_diffusion_model_path(model_path)
def _is_diffusers_model_dir(model_dir: str) -> bool:
"""Check if a local directory contains a valid diffusers model_index.json."""
config_path = os.path.join(model_dir, "model_index.json")
if not os.path.exists(config_path):
return False
with open(config_path) as f:
config = json.load(f)
return "_diffusers_version" in config
def _is_gated_diffusion_repo(repo_id: str) -> bool:
"""Query HF model card metadata to check if a gated repo is a diffusers model."""
try:
info = HfApi().model_info(repo_id)
return getattr(info, "library_name", None) == "diffusers"
except Exception:
return False
def get_is_diffusion_model(model_path: str) -> bool:
"""Detect whether model_path points to a diffusion model.
For local directories, checks the filesystem directly.
For HF/ModelScope model IDs, attempts to fetch only model_index.json.
For gated repos where file download fails, falls back to HF model card
metadata (library_name == "diffusers").
Returns False on any failure (network error, 404, offline mode, etc.)
so that the caller falls through to the standard LLM server path.
"""
if _is_overlay_diffusion_model(model_path):
# short-circuit, if applicable for the overlay mechanism (diffusion-only)
return True
if os.path.isdir(model_path):
if _is_diffusers_model_dir(model_path):
return True
return is_known_non_diffusers_diffusion_model(model_path)
if is_known_non_diffusers_diffusion_model(model_path):
return True
if _is_registered_diffusion_model(model_path):
return True
try:
if envs.SGLANG_USE_MODELSCOPE.get():
from modelscope import model_file_download
file_path = model_file_download(
model_id=model_path, file_path="model_index.json"
)
else:
from huggingface_hub import hf_hub_download
file_path = hf_hub_download(repo_id=model_path, filename="model_index.json")
return _is_diffusers_model_dir(os.path.dirname(file_path))
except Exception as e:
logger.debug("Failed to auto-detect diffusion model for %s: %s", model_path, e)
return False
def get_model_path(extra_argv):
# Find the model_path argument
model_path = None
for i, arg in enumerate(extra_argv):
if arg == "--model-path":
if i + 1 < len(extra_argv):
model_path = extra_argv[i + 1]
break
elif arg.startswith("--model-path="):
model_path = arg.split("=", 1)[1]
break
if model_path is None:
# Fallback for --help or other cases where model-path is not provided
if any(h in extra_argv for h in ["-h", "--help"]):
raise Exception(
"Usage: sglang serve --model-path <model-name-or-path> [additional-arguments]\n\n"
"This command can launch either a standard language model server or a diffusion model server.\n"
"The server type is determined by the --model-path.\n"
)
else:
raise Exception(
"Error: --model-path is required. "
"Please provide the path to the model."
)
return model_path
@lru_cache(maxsize=1)
def get_git_commit_hash() -> str:
try:
commit_hash = os.environ.get("SGLANG_GIT_COMMIT")
if not commit_hash:
commit_hash = (
subprocess.check_output(
["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
)
.strip()
.decode("utf-8")
)
_CACHED_COMMIT_HASH = commit_hash
return commit_hash
except (subprocess.CalledProcessError, FileNotFoundError):
_CACHED_COMMIT_HASH = "N/A"
return "N/A"

191
third_party/sglang/python/sglang/compile_deep_gemm.py vendored Normal file
View File

@@ -0,0 +1,191 @@
"""
Compile DeepGEMM Kernels for a model with specify server arguments
This script launches a server for capturing DeepGEMM calls and then compiles the kernels.
It accepts server arguments (the same as launch_server.py).
Usage:
python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code
"""
import argparse
import dataclasses
import multiprocessing
import os
import time
import requests
from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST
from sglang.srt.entrypoints.http_server import launch_server
from sglang.srt.entrypoints.warmup import warmup
from sglang.srt.environ import envs
from sglang.srt.managers.io_struct import GenerateReqInput
from sglang.srt.managers.tokenizer_manager import TokenizerManager
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import kill_process_tree
multiprocessing.set_start_method("spawn", force=True)
# Reduce warning
envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.set(True)
# Force enable deep gemm
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(True)
# Force enable mha chunked kv for DeepSeek V3 to avoid missing kv_b_proj DeepGEMM case
envs.SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD.set(0)
@dataclasses.dataclass
class CompileArgs:
timeout: int = 3600
@staticmethod
def add_cli_args(parser: argparse.ArgumentParser):
parser.add_argument("--timeout", type=int, default=CompileArgs.timeout)
@classmethod
def from_cli_args(cls, args: argparse.Namespace):
# use the default value's type to cast the args into correct types.
attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)]
return cls(
**{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs}
)
@warmup("compile-deep-gemm")
async def warm_up_compile(
disaggregation_mode: str, tokenizer_manager: TokenizerManager
):
print("\nGenerate warm up request for compiling DeepGEMM...\n")
generate_req_input = GenerateReqInput(
input_ids=[0, 1, 2, 3],
sampling_params={
"temperature": 0.0,
"max_new_tokens": 8,
"ignore_eos": True,
},
)
if disaggregation_mode != "null":
generate_req_input.bootstrap_room = 0
generate_req_input.bootstrap_host = FAKE_BOOTSTRAP_HOST
await tokenizer_manager.generate_request(generate_req_input, None).__anext__()
def launch_server_internal(server_args):
try:
launch_server(server_args)
except Exception as e:
raise e
finally:
kill_process_tree(os.getpid(), include_parent=False)
def launch_server_process_and_send_one_request(
server_args: ServerArgs, compile_args: CompileArgs
):
proc = multiprocessing.Process(target=launch_server_internal, args=(server_args,))
proc.start()
base_url = f"http://{server_args.host}:{server_args.port}"
timeout = compile_args.timeout
start_time = time.perf_counter()
while time.perf_counter() - start_time < timeout:
try:
headers = {
"Content-Type": "application/json; charset=utf-8",
}
if server_args.node_rank == 0:
response = requests.get(f"{base_url}/v1/models", headers=headers)
else:
# This http api is created by launch_dummy_health_check_server for none-rank0 node.
response = requests.get(f"{base_url}/health", headers=headers)
if response.status_code == 200:
# Rank-0 node send a request to sync with other node and then return.
if server_args.node_rank == 0:
payload = {
"input_ids": [0, 1, 2, 3],
"sampling_params": {
"max_new_tokens": 8,
"temperature": 0,
},
}
# In PD mode, include fake bootstrap fields so workers don't assert
if server_args.disaggregation_mode != "null":
payload["bootstrap_host"] = FAKE_BOOTSTRAP_HOST
payload["bootstrap_room"] = 0
response = requests.post(
f"{base_url}/generate",
json=payload,
timeout=600,
)
if response.status_code != 200:
error = response.json()
raise RuntimeError(f"Sync request failed: {error}")
# Other nodes should wait for the exit signal from Rank-0 node.
else:
start_time_waiting = time.perf_counter()
while proc.is_alive():
if time.perf_counter() - start_time_waiting < timeout:
time.sleep(10)
else:
raise TimeoutError("Waiting for main node timeout!")
return proc
except requests.RequestException:
pass
time.sleep(10)
raise TimeoutError(
"DeepGEMM Kernels compilation timeout."
"\n\nFeel free and please restart the command."
)
def refine_server_args(server_args: ServerArgs, compile_args: CompileArgs):
# Disable cuda graph and torch compile to save time
server_args.disable_cuda_graph = True
server_args.enable_torch_compile = False
print(f"Disable CUDA Graph and Torch Compile to save time...")
# Set watchdog timeout to compile_args.timeout because compilation will take a long time
server_args.watchdog_timeout = compile_args.timeout
server_args.warmups = "compile-deep-gemm"
def run_compile(server_args: ServerArgs, compile_args: CompileArgs):
print(
"Begin DeepGEMM Kernels compilation...\n"
"It may take a long time and timeout maybe raised "
"while the compilation is still in progress.\n"
"Just feel free to restart the command "
"until the compilation is fully finished.\n"
)
proc = launch_server_process_and_send_one_request(server_args, compile_args)
print("\nDeepGEMM Kernels compilation finished successfully.")
# Sleep for safety
time.sleep(10)
if proc.is_alive():
# This is the rank0 node.
kill_process_tree(proc.pid)
else:
try:
kill_process_tree(proc.pid)
except Exception:
pass
if __name__ == "__main__":
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
CompileArgs.add_cli_args(parser)
args = parser.parse_args()
server_args = ServerArgs.from_cli_args(args)
compile_args = CompileArgs.from_cli_args(args)
refine_server_args(server_args, compile_args)
run_compile(server_args, compile_args)

315
third_party/sglang/python/sglang/eval/llama3_eval.py vendored Normal file
View File

@@ -0,0 +1,315 @@
# Adapt from https://github.com/fw-ai/llm_eval_meta
import argparse
import asyncio
import os
import pickle
import re
import shutil
from collections import defaultdict
from dataclasses import dataclass
import httpx
import numpy as np
import openai
from datasets import load_dataset
from openai import AsyncOpenAI
from tqdm import tqdm
# Mapping providers to their clients and models
provider_to_models = {
"b10": {
"8b": "meta-llama/Llama-3.1-8B-Instruct",
"70b": "meta-llama/Llama-3.1-70B-Instruct",
"405b": "meta-llama/Llama-3.1-405B-Instruct",
},
"oai": {
"8b": "meta-llama/Llama-3.1-8B-Instruct",
"70b": "meta-llama/Llama-3.1-70B-Instruct",
"405b": "meta-llama/Llama-3.1-405B-Instruct",
},
"sgl": {
"8b": "meta-llama/Llama-3.1-8B-Instruct",
"70b": "meta-llama/Llama-3.1-70B-Instruct",
"405b": "meta-llama/Llama-3.1-405B-Instruct",
},
}
async def fetch_responses(
client, prompt, semaphore, index, provider, model_size, output_dir, max_tokens
):
output_file = os.path.join(output_dir, f"response_{index}.pkl")
if os.path.exists(output_file):
print(f"File {output_file} already exists, skipping.")
return
async with semaphore:
response = await client.completions.create(
model=provider_to_models[provider][model_size],
prompt=prompt,
temperature=0.0,
max_tokens=max_tokens,
)
if isinstance(response, openai.BadRequestError):
with open(output_file, "wb") as f:
pickle.dump("bad_response", f)
assert isinstance(response, openai.types.completion.Completion)
# Save response to a file
with open(output_file, "wb") as f:
pickle.dump(response, f)
TASK_TO_MAX_TOKENS = {
"evals__mmlu__details": 1,
"evals__mmlu__0_shot__cot__details": 1024,
# Official meta uses 1024, but a small % (.05) of questions are answered correctly after relaxing
"evals__mmlu_pro__details": 2048,
"evals__gsm8k__details": 1024,
}
TASK_TO_EVAL_SET = {
"mmlu": "evals__mmlu__details",
"mmlu_cot": "evals__mmlu__0_shot__cot__details",
"mmlu_pro": "evals__mmlu_pro__details",
"gsm8k": "evals__gsm8k__details",
}
class CustomAsyncHTTPXClient(httpx.AsyncClient):
async def send(self, request: httpx.Request, *args, **kwargs) -> httpx.Response:
request.url = httpx.URL(
f"https://model-{os.getenv('MODEL_ID')}.api.baseten.co/development/predict"
)
return await super().send(request, *args, **kwargs)
def get_client(provider):
if provider not in "b10":
if os.getenv("OPENAI_API_KEY") is None:
os.environ["OPENAI_API_KEY"] = "EMPTY"
return {
"oai": AsyncOpenAI(base_url="http://127.0.0.1:8000/v1/"),
"b10": AsyncOpenAI(
api_key=f"Api-Key {os.getenv('OPENAI_API_KEY')}",
base_url=f"https://model-{os.getenv('MODEL_ID')}.api.baseten.co/development/predict",
http_client=CustomAsyncHTTPXClient(),
),
"sgl": AsyncOpenAI(base_url="http://127.0.0.1:30000/v1/"),
}[provider]
# Define the benchmark function
async def benchmark(args):
ds = load_dataset(
"meta-llama/Llama-3.1-405B-Instruct-evals",
f"Llama-3.1-405B-Instruct-{TASK_TO_EVAL_SET[args.task]}",
)
semaphore = asyncio.Semaphore(args.concurrency) # Limit to 16 concurrent tasks
if args.num_examples is None:
args.num_examples = len(ds["latest"]["input_final_prompts"])
prompts = ds["latest"]["input_final_prompts"][: args.num_examples]
# Create the output directory if it does not exist
os.makedirs(args.output_dir, exist_ok=True)
tasks = []
# Create the tasks with tqdm progress bar
max_tokens = TASK_TO_MAX_TOKENS[TASK_TO_EVAL_SET[args.task]]
client = get_client(args.provider)
for idx, prompt in enumerate(tqdm(prompts, desc="Creating tasks")):
tasks.append(
asyncio.create_task(
fetch_responses(
client,
f"<|begin_of_text|>{prompt[0]}",
semaphore,
idx,
args.provider,
args.model_size,
args.output_dir,
max_tokens=max_tokens,
)
)
)
# Run the tasks with tqdm progress bar
for future in tqdm(
asyncio.as_completed(tasks), total=len(tasks), desc="Processing tasks"
):
await future
def get_mmlu_answer(response):
if response is not None:
return response.choices[0].text.lstrip().rstrip().upper().replace(".", "")
return None
def get_mmlu_cot_answer(response):
pattern = r"The best answer is (.+)\.?"
match = re.search(pattern, response.choices[0].text)
if match:
return match.group(1).replace(".", "").replace("*", "")
pattern = r"the best answer is (.+)\.?"
match = re.search(pattern, response.choices[0].text)
if match:
return match.group(1).replace(".", "")
pattern = r"The correct answer is (.+)\.?"
match = re.search(pattern, response.choices[0].text)
if match:
return match.group(1).replace(".", "")
pattern = r"the correct answer is (.+)\.?"
match = re.search(pattern, response.choices[0].text)
if match:
return match.group(1).replace(".", "")
def get_answer_gsm8k(response):
pattern = r"The final answer is (.+)\.?"
match = re.search(pattern, response.choices[0].text)
if match:
s = match.group(1)
for ok_symbol in ["%", "$"]:
s = s.replace(ok_symbol, "")
return s
TASK_TO_ANSWER_EXTRACTOR = {
"evals__mmlu__details": get_mmlu_answer,
"evals__mmlu__0_shot__cot__details": get_mmlu_cot_answer,
"evals__gsm8k__details": get_answer_gsm8k,
"evals__mmlu_pro__details": get_mmlu_cot_answer,
}
def get_dataset_from_task(task, response_path, model_size):
ds_405b = load_dataset(
f"meta-llama/Llama-3.1-405B-Instruct-evals",
f"Llama-3.1-405B-Instruct-{task}",
)
ds_405b_hash_order = [x[0] for x in ds_405b["latest"]["input_final_prompts_hash"]]
if "70b" in model_size or "8b" in model_size:
if "70" in model_size:
ref_model_ds = load_dataset(
f"meta-llama/Llama-3.1-70B-Instruct-evals",
f"Llama-3.1-70B-Instruct-{task}",
)
else:
ref_model_ds = load_dataset(
f"meta-llama/Llama-3.1-8B-Instruct-evals",
f"Llama-3.1-8B-Instruct-{task}",
)
hash_to_row = {}
for row in ref_model_ds["latest"]:
hash_to_row[row["input_final_prompts_hash"][0]] = row
reordered_rows = []
for prompt_hash in ds_405b_hash_order:
reordered_rows.append(hash_to_row[prompt_hash])
ref_model_ds["latest"] = reordered_rows
return ref_model_ds
return ds_405b
def analyze(task, response_path, model_size):
ds = get_dataset_from_task(task, response_path, model_size)
responses = []
total = len(ds["latest"])
for i in range(0, total):
response = pickle.load(
open(os.path.join(response_path, f"response_{i}.pkl"), "rb")
)
responses.append(response)
@dataclass
class Stats:
correct: int = 0
total: int = 0
meta_correct: int = 0
average: float = None
subtask_name_to_stats = defaultdict(lambda: Stats())
for response, ds_row in zip(responses, ds["latest"]):
model_answer = TASK_TO_ANSWER_EXTRACTOR[task](response)
subtask = ds_row["subtask_name"]
is_eval_correct = model_answer in ds_row["input_correct_responses"]
if is_eval_correct:
subtask_name_to_stats[subtask].correct += 1
if ds_row["is_correct"]:
subtask_name_to_stats[subtask].meta_correct += 1
subtask_name_to_stats[subtask].total += 1
micro_stats = Stats()
for subtask, stats in subtask_name_to_stats.items():
stats.average = stats.correct / stats.total
stats.meta_average = stats.meta_correct / stats.total
micro_stats.correct += stats.correct
micro_stats.total += stats.total
micro_stats.meta_correct += stats.meta_correct
micro_stats.average = micro_stats.correct / micro_stats.total
micro_stats.meta_average = micro_stats.meta_correct / micro_stats.total
print("Macro average", np.mean([x.average for x in subtask_name_to_stats.values()]))
print(
"Meta Macro average",
np.mean([x.meta_average for x in subtask_name_to_stats.values()]),
)
print("Micro average", micro_stats.average)
print("Meta Micro average", micro_stats.meta_average)
# Entry point for the script
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Script to run model with specified parameters."
)
parser.add_argument(
"--model-size",
type=str,
default="8b",
help="Size of the model (e.g., 8b or 70b)",
)
parser.add_argument(
"--provider",
type=str,
default="sgl",
help="Provider name (e.g., sgl, oai, b10)",
)
parser.add_argument(
"--task",
type=str,
required=True,
help="Task (e.g., mmlu, mmlu_cot, mmlu_pro, gsm8k)",
)
parser.add_argument(
"--num-examples", type=int, default=None, help="Number of examples to process"
)
parser.add_argument("--concurrency", type=int, default=16)
parser.add_argument(
"--output-dir",
type=str,
default="tmp-output-dir",
help="Directory to save responses",
)
args = parser.parse_args()
asyncio.run(benchmark(args))
analyze(TASK_TO_EVAL_SET[args.task], args.output_dir, args.model_size)
shutil.rmtree("tmp-output-dir", ignore_errors=True)

164
third_party/sglang/python/sglang/eval/loogle_eval.py vendored Normal file
View File

@@ -0,0 +1,164 @@
import argparse
import asyncio
import os
import pickle
from pathlib import Path
from typing import List
import openai
import torch
from bert_score import BERTScorer
from datasets import load_dataset
from tqdm import tqdm
def get_client(api_url: str) -> openai.AsyncOpenAI:
if os.getenv("OPENAI_API_KEY") is None:
os.environ["OPENAI_API_KEY"] = "EMPTY"
return openai.AsyncOpenAI(base_url=api_url)
def get_dataset():
return load_dataset("bigai-nlco/LooGLE", "longdep_qa", split="test")
async def fetch_response(
client: openai.AsyncOpenAI,
context: str,
question: str,
semaphore: asyncio.Semaphore,
index: int,
model: str,
output_dir: Path,
):
output_file = output_dir / f"response_{index}.pkl"
if output_file.exists():
return
prompt = (
"Please answer the question based on the long texts below.\n"
f"{context}\n"
f"Question: {question}\n"
"Answer:"
)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": prompt},
]
async with semaphore:
try:
response = await client.chat.completions.create(
model=model,
messages=messages,
temperature=0.0,
max_tokens=512,
)
except openai.BadRequestError as e:
with open(output_file, "wb") as f:
pickle.dump({"error": str(e)}, f)
return
with open(output_file, "wb") as f:
pickle.dump(response, f)
async def benchmark(args):
dataset = get_dataset()
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
client = get_client(args.api_url)
semaphore = asyncio.Semaphore(args.max_concurrency)
tasks: List[asyncio.Task] = []
for idx, ex in enumerate(dataset):
if idx >= args.num_prompts:
break
tasks.append(
asyncio.create_task(
fetch_response(
client,
ex["context"],
ex["question"],
semaphore,
idx,
args.model,
output_dir,
)
)
)
for _ in tqdm(
asyncio.as_completed(tasks), total=len(tasks), desc="Running benchmark"
):
await _
def analyse(args):
dataset = get_dataset()
output_dir = Path(args.output_dir)
device = "cuda" if torch.cuda.is_available() else "cpu"
scorer = BERTScorer(lang="en", device=device)
hyps: List[str] = []
refs: List[str] = []
for idx, ex in enumerate(tqdm(dataset, desc="Loading responses")):
if idx >= args.num_prompts:
break
pkl_file = output_dir / f"response_{idx}.pkl"
if not pkl_file.exists():
raise FileNotFoundError(pkl_file)
response = pickle.load(open(pkl_file, "rb"))
if isinstance(response, dict) and "error" in response:
continue
hyps.append(response.choices[0].message.content.strip())
refs.append(ex["answer"])
if not hyps:
print("No valid responses to score!")
return
batch_size = 64
all_f1: List[float] = []
for i in tqdm(range(0, len(hyps), batch_size), desc="Scoring batches"):
h_batch = hyps[i : i + batch_size]
r_batch = refs[i : i + batch_size]
_, _, f1_scores = scorer.score(h_batch, r_batch, verbose=False)
all_f1.extend([float(x) for x in f1_scores])
avg = sum(all_f1) / len(all_f1)
print(f"Average BERTScore (F1): {avg:.2%}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Run benchmark and evaluation in one go."
)
parser.add_argument(
"--api-url",
default="http://127.0.0.1:30000/v1",
help="OpenAI‑compatible API base URL",
)
parser.add_argument(
"--model",
default="meta-llama/Llama-4-Maverick-17B-128E-Instruct",
help="Model name or ID, only used for model name",
)
parser.add_argument(
"--max-concurrency", type=int, default=144, help="Maximum concurrent requests"
)
parser.add_argument(
"--output-dir", default="tmp-output-dir", help="Directory for cached responses"
)
parser.add_argument(
"--num-prompts", type=int, default=10000, help="Number of prompts to run"
)
args = parser.parse_args()
asyncio.run(benchmark(args))
analyse(args)

29
third_party/sglang/python/sglang/global_config.py vendored Normal file
View File

@@ -0,0 +1,29 @@
"""Global configurations"""
# FIXME: deprecate this file and move all usage to sglang.srt.environ or sglang.__init__.py
class GlobalConfig:
"""
Store some global constants.
"""
def __init__(self):
# Verbosity level
# 0: do not output anything
# 2: output final text after every run
self.verbosity = 0
# Default backend of the language
self.default_backend = None
# Output tokenization configs
self.skip_special_tokens_in_output = True
self.spaces_between_special_tokens_in_out = True
# Language frontend interpreter optimization configs
self.enable_precache_with_tracing = True
self.enable_parallel_encoding = True
global_config = GlobalConfig()

25
third_party/sglang/python/sglang/jit_kernel/.clang-format vendored Normal file
View File

@@ -0,0 +1,25 @@
BasedOnStyle: Google
IndentWidth: 2
ColumnLimit: 120
AllowShortFunctionsOnASingleLine: Empty
DerivePointerAlignment: false
PointerAlignment: Left
NamespaceIndentation: None
SortIncludes: true
AllowShortLoopsOnASingleLine: false
BinPackParameters: false # Prevents packing parameters in declarations
BinPackArguments: false # Prevents packing arguments in function calls
AlignAfterOpenBracket: AlwaysBreak # Forces a break after the opening parenthesis
AlignOperands: Align # Aligns arguments vertically
PenaltyBreakBeforeFirstCallParameter: 1 # Encourages breaking before the first argument
PenaltyReturnTypeOnItsOwnLine: 100 # Keeps return type with function name
IncludeCategories:
- Regex: '^<sgl_kernel/.*\.h>$'
Priority: 0
- Regex: '^<sgl_kernel/.*/.*>$'
Priority: 2
- Regex: '^<sgl_kernel/.*\.cuh>$'
Priority: 1
- Regex: '^<.*/.*>$'
Priority: 3

0
third_party/sglang/python/sglang/jit_kernel/__init__.py vendored Normal file
View File

99
third_party/sglang/python/sglang/jit_kernel/__main__.py vendored Normal file
View File

@@ -0,0 +1,99 @@
import argparse
import logging
import os
from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path
from sglang.jit_kernel.utils import (
_REGISTERED_DEPENDENCIES,
DEFAULT_INCLUDE,
_get_default_target_flags,
get_jit_cuda_arch,
override_jit_cuda_arch,
)
def generate_clangd():
logger = logging.getLogger()
parser = argparse.ArgumentParser(
description="Generate .clangd file for sglang jit kernel development."
)
parser.add_argument(
"--overwrite",
action="store_true",
help="Overwrite existing .clangd file if it exists.",
)
parser.add_argument(
"--dependencies",
"--dep",
nargs="*",
default=[],
choices=_REGISTERED_DEPENDENCIES.keys(),
help="Extra dependency libraries to include.",
)
parser.add_argument(
"--cuda-target",
"--cuda",
default=None,
type=str,
help="Target architecture to generate compile flags for.",
)
args = parser.parse_args()
dep_include_paths = []
for dep in args.dependencies:
if dep not in _REGISTERED_DEPENDENCIES:
raise ValueError(f"Dependency {dep} is not registered.")
dep_include_paths += _REGISTERED_DEPENDENCIES[dep]()
include_paths = [
*DEFAULT_INCLUDE,
find_include_path(),
find_dlpack_include_path(),
*dep_include_paths,
]
if args.cuda_target:
assert args.cuda_target.count(".") == 1
major, minor = args.cuda_target.split(".")
major, minor = int(major), int(minor)
context = override_jit_cuda_arch(major, minor)
context.__enter__()
else:
arch = get_jit_cuda_arch()
major, minor = arch.major, f"{arch.minor}{arch.suffix}"
assert (
major > 0
), "Cannot detect CUDA architecture, please specify --cuda-target explicitly."
compile_flags = [
"-xcuda",
f"--cuda-gpu-arch=sm_{major}{minor}",
"-Wall",
"-Wextra",
*_get_default_target_flags(),
*[f"-isystem{path}" for path in include_paths],
]
# NOTE: skip these flags because clangd don't recognize them
UNSUPPORTED_FLAGS = {"--expt-relaxed-constexpr"}
compile_flags = [flag for flag in compile_flags if flag not in UNSUPPORTED_FLAGS]
compile_flags_str = ",\n ".join(compile_flags)
clangd_content = f"""
CompileFlags:
Add: [
{compile_flags_str}
]
"""
if os.path.exists(".clangd") and not args.overwrite:
logger.warning(".clangd file already exists, nothing done.")
logger.warning("Use --overwrite to force overwrite the existing .clangd file.")
logger.warning(f"suggested content: {clangd_content}")
else:
with open(".clangd", "w") as f:
f.write(clangd_content)
logger.info(".clangd file generated.")
assert __name__ == "__main__"
logging.basicConfig(level=logging.INFO)
generate_clangd()

28
third_party/sglang/python/sglang/jit_kernel/add_constant.py vendored Normal file
View File

@@ -0,0 +1,28 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_add_constant_module(constant: int) -> Module:
args = make_cpp_args(constant)
return load_jit(
"add_constant",
*args,
cuda_files=["add_constant.cuh"],
cuda_wrappers=[("add_constant", f"add_constant<{args}>")],
)
def add_constant(src: torch.Tensor, constant: int) -> torch.Tensor:
dst = torch.empty_like(src)
module = _jit_add_constant_module(constant)
module.add_constant(dst, src)
return dst

196
third_party/sglang/python/sglang/jit_kernel/all_reduce.py vendored Normal file
View File

@@ -0,0 +1,196 @@
from __future__ import annotations
import enum
from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, cast
import torch
import tvm_ffi
from sglang.jit_kernel.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.kernel_api_logging import debug_kernel_api
class ConfigResult(NamedTuple):
num_blocks: int
num_threads: int
class AllReduceAlgo(enum.Enum):
ONE_SHOT_PUSH = enum.auto()
ONE_SHOT_PULL = enum.auto()
TWO_SHOT_PULL = enum.auto()
def is_push(self) -> bool:
return self == AllReduceAlgo.ONE_SHOT_PUSH
@property
def shot(self) -> int:
return 2 if self == AllReduceAlgo.TWO_SHOT_PULL else 1
if TYPE_CHECKING:
CUSTOM_AR_HANDLE = List[int]
CUSTOM_AR_PAIR = Tuple[int, CUSTOM_AR_HANDLE]
class CustomAllReduceObj:
def __init__(
self,
rank: int,
world_size: int,
pull_buffer_bytes: int,
push_buffer_bytes: int,
graph_input_count: int,
*,
max_pull_blocks: Optional[int] = None,
max_push_blocks: Optional[int] = None,
) -> None:
"""
Create a CustomAllReduceObj instance.
:param rank: The rank of the current process.
:param world_size: The total number of processes in the group.
:param pull_buffer_bytes: The size of the buffer (in bytes) used for pull-based all-reduce.
:param push_buffer_bytes: The size of the buffer (in bytes) used for push-based all-reduce.
:param graph_input_count: The maximum number of inputs in all CUDA graphs.
:param max_pull_blocks: The maximum number of thread blocks to launch for pull-based all-reduce.
If None, it will be determined by the implementation.
:param max_push_blocks: The maximum number of thread blocks to launch for push-based all-reduce.
If None, it will be determined by the implementation.
"""
@property
def world_size(self) -> int: ...
def share_storage(self) -> CUSTOM_AR_HANDLE: ...
def share_graph_inputs(self) -> List[CUSTOM_AR_PAIR]: ...
def post_init(self, handles: List[CUSTOM_AR_HANDLE]) -> None: ...
def register_inputs(self, handles: List[List[CUSTOM_AR_PAIR]]) -> None: ...
def set_cuda_graph_capture(self, is_capturing: bool) -> None: ...
def free(self, tp_cpu_group: torch.distributed.ProcessGroup) -> None: ...
def all_reduce(
self, input: torch.Tensor, algo: AllReduceAlgo
) -> tvm_ffi.Tensor: ...
def config_pull(
self, num_blocks: int = -1, num_threads: int = -1
) -> ConfigResult:
"""
Configure the CUDA kernel's grid and block dimensions.
This provides only the upper bound of the configuration,
and the actual launch configuration may be determined by implementation.
Note that push-based all-reduce can not be configured currently.
:param num_blocks: The maximum number of thread blocks to launch. -1 means no limit.
:param num_threads: The maximum number of threads per block. -1 means no limit.
:return: The previous configuration as a ConfigResult named tuple.
"""
...
@cache_once
def _jit_custom_all_reduce_pull_module(dtype: torch.dtype, world_size: int):
args = make_cpp_args(dtype, world_size, is_arch_support_pdl())
return load_jit(
"custom_all_reduce_pull",
*args,
extra_ldflags=["-lcuda"],
cuda_files=["distributed/custom_all_reduce_pull.cuh"],
cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")],
)
@cache_once
def _jit_custom_all_reduce_push_module(dtype: torch.dtype, world_size: int):
args = make_cpp_args(dtype, world_size, is_arch_support_pdl())
return load_jit(
"custom_all_reduce_push",
*args,
extra_ldflags=["-lcuda"],
cuda_files=["distributed/custom_all_reduce_push.cuh"],
cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")],
)
@cache_once
def get_custom_all_reduce_cls() -> type[CustomAllReduceObj]:
module = load_jit(
"custom_all_reduce_base",
extra_ldflags=["-lcuda"],
cuda_files=["distributed/custom_all_reduce_base.cuh"],
cuda_wrappers=[("register_once", "register_custom_all_reduce")],
)
module.register_once()
device = torch.cuda.current_device()
props = torch.cuda.get_device_properties(device)
NUM_CTA = props.multi_processor_count
MAX_THREADS = 512
@tvm_ffi.register_object("sgl.CustomAllReduce")
class CustomAllReduceObjReal(tvm_ffi.Object):
__slots__ = ("__dict__",)
def __init__(
self,
rank: int,
world_size: int,
pull_buffer_bytes: int,
push_buffer_bytes: int,
graph_input_count: int,
*,
max_pull_blocks: Optional[int] = None,
max_push_blocks: Optional[int] = None,
) -> None:
self.__ffi_init__(
rank,
world_size,
NUM_CTA if max_pull_blocks is None else max_pull_blocks,
NUM_CTA if max_push_blocks is None else max_push_blocks,
pull_buffer_bytes,
push_buffer_bytes,
graph_input_count,
)
self._world_size = world_size
self._pull_config = ConfigResult(NUM_CTA, MAX_THREADS)
self.configure_pull(*self._pull_config) # type: ignore
@property
def world_size(self) -> int:
return self._world_size
@debug_kernel_api
def all_reduce(
self,
input: torch.Tensor,
algo: AllReduceAlgo,
) -> tvm_ffi.Tensor:
compile_fn = (
_jit_custom_all_reduce_push_module
if algo.is_push()
else _jit_custom_all_reduce_pull_module
)
module = compile_fn(input.dtype, self._world_size)
return module.all_reduce(self, input, algo.shot)
def config_pull(
self, num_blocks: int = -1, num_threads: int = -1
) -> ConfigResult:
old_config = self._pull_config
num_blocks = num_blocks if num_blocks != -1 else old_config.num_blocks
num_threads = num_threads if num_threads != -1 else old_config.num_threads
new_config = ConfigResult(num_blocks, num_threads)
if new_config != old_config:
result = ConfigResult(*self.configure_pull(*new_config)) # type: ignore
assert result == self._pull_config
self._pull_config = new_config
return old_config
def free(self, tp_cpu_group: torch.distributed.ProcessGroup) -> None:
self.free_ipc_handles() # type: ignore
torch.distributed.barrier(group=tp_cpu_group)
self.free_storage() # type: ignore
return cast(type["CustomAllReduceObj"], CustomAllReduceObjReal)

38
third_party/sglang/python/sglang/jit_kernel/awq_dequantize.py vendored Normal file
View File

@@ -0,0 +1,38 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_awq_dequantize_module(dtype: torch.dtype) -> Module:
args = make_cpp_args(dtype)
return load_jit(
"awq_dequantize",
*args,
cuda_files=["gemm/awq_dequantize.cuh"],
cuda_wrappers=[("awq_dequantize", f"awq_dequantize<{args}>")],
)
def awq_dequantize(
qweight: torch.Tensor,
scales: torch.Tensor,
qzeros: torch.Tensor,
) -> torch.Tensor:
qweight_rows = qweight.shape[0]
qweight_cols = qweight.shape[1]
output = torch.empty(
(qweight_rows, qweight_cols * 8),
dtype=scales.dtype,
device=scales.device,
)
module = _jit_awq_dequantize_module(scales.dtype)
module.awq_dequantize(output, qweight, scales, qzeros)
return output

59
third_party/sglang/python/sglang/jit_kernel/awq_marlin_repack.py vendored Normal file
View File

@@ -0,0 +1,59 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
from sglang.kernel_api_logging import debug_kernel_api
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_awq_marlin_repack_module() -> Module:
return load_jit(
"awq_marlin_repack",
cuda_files=["gemm/marlin/awq_marlin_repack.cuh"],
cuda_wrappers=[("awq_marlin_repack", "awq_marlin_repack")],
)
@debug_kernel_api
def awq_marlin_repack(
b_q_weight: torch.Tensor,
size_k: int,
size_n: int,
num_bits: int,
) -> torch.Tensor:
tile_size = 16
pack_factor = 32 // num_bits
out = torch.empty(
(size_k // tile_size, size_n * tile_size // pack_factor),
dtype=b_q_weight.dtype,
device=b_q_weight.device,
)
module = _jit_awq_marlin_repack_module()
module.awq_marlin_repack(out, b_q_weight, size_k, size_n, num_bits)
return out
@debug_kernel_api
def awq_marlin_moe_repack(
b_q_weight: torch.Tensor,
perm: torch.Tensor,
size_k: int,
size_n: int,
num_bits: int,
) -> torch.Tensor:
num_experts = b_q_weight.shape[0]
assert size_k % 16 == 0
output = torch.empty(
(num_experts, size_k // 16, size_n * (num_bits // 2)),
device=b_q_weight.device,
dtype=b_q_weight.dtype,
)
for e in range(num_experts):
output[e] = awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
return output

122
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py vendored Normal file
View File

@@ -0,0 +1,122 @@
import itertools
import torch
import triton
import triton.testing
from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
try:
from sgl_kernel import awq_dequantize as aot_awq_dequantize
AOT_AVAILABLE = True
except ImportError:
AOT_AVAILABLE = False
IS_CI = is_in_ci()
if IS_CI:
qweight_row_range = [128]
qweight_cols_range = [16]
else:
qweight_row_range = [128, 256, 512, 1024, 3584]
qweight_cols_range = [16, 32, 64, 128, 448]
configs = list(itertools.product(qweight_row_range, qweight_cols_range))
def check_correctness():
if not AOT_AVAILABLE:
print("sgl_kernel AOT not available, skipping correctness check")
return
qweight_row, qweight_col = 128, 16
device = torch.device("cuda")
qweight = torch.randint(
0,
torch.iinfo(torch.int32).max,
(qweight_row, qweight_col),
dtype=torch.int32,
device=device,
)
group_size = qweight_row
scales_row = qweight_row // group_size
scales_col = qweight_col * 8
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
qzeros = torch.randint(
0,
torch.iinfo(torch.int32).max,
(scales_row, qweight_col),
dtype=torch.int32,
device=device,
)
jit_out = jit_awq_dequantize(qweight, scales, qzeros)
aot_out = aot_awq_dequantize(qweight, scales, qzeros)
torch.cuda.synchronize()
torch.testing.assert_close(jit_out, aot_out, rtol=0, atol=0)
print("Correctness check passed (JIT vs AOT)")
if AOT_AVAILABLE:
line_vals = ["jit", "aot"]
line_names = ["JIT Kernel", "AOT Kernel"]
styles = [("blue", "-"), ("green", "-")]
else:
line_vals = ["jit"]
line_names = ["JIT Kernel"]
styles = [("blue", "-")]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["qweight_row", "qweight_col"],
x_vals=configs,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="awq-dequantize-jit-vs-aot",
args={},
)
)
def benchmark(qweight_row, qweight_col, provider):
device = torch.device("cuda")
qweight = torch.randint(
0,
torch.iinfo(torch.int32).max,
(qweight_row, qweight_col),
dtype=torch.int32,
device=device,
)
group_size = qweight_row
scales_row = qweight_row // group_size
scales_col = qweight_col * 8
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
qzeros = torch.randint(
0,
torch.iinfo(torch.int32).max,
(scales_row, qweight_col),
dtype=torch.int32,
device=device,
)
if provider == "jit":
fn = lambda: jit_awq_dequantize(qweight, scales, qzeros)
elif provider == "aot":
fn = lambda: aot_awq_dequantize(qweight, scales, qzeros)
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
check_correctness()
benchmark.run(print_data=True)

106
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_cast.py vendored Normal file
View File

@@ -0,0 +1,106 @@
import torch
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import (
DEFAULT_DEVICE,
get_benchmark_range,
run_benchmark,
)
from sglang.jit_kernel.cast import downcast_fp8 as downcast_fp8_jit
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=10, suite="stage-b-kernel-benchmark-1-gpu-large")
DEVICE = DEFAULT_DEVICE
DTYPE = torch.bfloat16
# ── Config ranges ──────────────────────────────────────────────────────────────
SL_LIST = get_benchmark_range(
full_range=[4, 16, 64, 256, 512, 1024, 2048],
ci_range=[4, 64],
)
HEAD_DIM_LIST = get_benchmark_range(
full_range=[(8, 128), (32, 128), (8, 256), (32, 256)],
ci_range=[(8, 128)],
)
CONFIGS = [(sl, h, d, sl * 2) for sl in SL_LIST for h, d in HEAD_DIM_LIST]
LINE_VALS = ["jit"]
LINE_NAMES = ["JIT (cast.cuh, 256 threads, 2D grid)"]
STYLES = [("orange", "-")]
# ── Perf report ────────────────────────────────────────────────────────────────
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["input_sl", "head", "dim", "out_sl"],
x_vals=CONFIGS,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="downcast-fp8-jit",
args={},
)
)
def benchmark(input_sl, head, dim, out_sl, provider):
k = torch.randn(input_sl, head, dim, dtype=DTYPE, device=DEVICE)
v = torch.randn(input_sl, head, dim, dtype=DTYPE, device=DEVICE)
k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device=DEVICE)
v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device=DEVICE)
k_scale = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
v_scale = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
loc = torch.arange(input_sl, dtype=torch.int64, device=DEVICE)
fn = lambda: downcast_fp8_jit(k, v, k_out, v_out, k_scale, v_scale, loc)
return run_benchmark(fn)
# ── Bandwidth analysis ─────────────────────────────────────────────────────────
def _report_bandwidth(input_sl, head, dim, dtype):
elem_bytes = torch.finfo(dtype).bits // 8
total_bytes = input_sl * head * dim * (2 * elem_bytes + 2)
k = torch.randn(input_sl, head, dim, dtype=dtype, device=DEVICE)
v = torch.randn(input_sl, head, dim, dtype=dtype, device=DEVICE)
k_out = torch.zeros(input_sl * 2, head, dim, dtype=torch.uint8, device=DEVICE)
v_out = torch.zeros(input_sl * 2, head, dim, dtype=torch.uint8, device=DEVICE)
k_scale = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
v_scale = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
loc = torch.arange(input_sl, dtype=torch.int64, device=DEVICE)
jit_fn = lambda: downcast_fp8_jit(k, v, k_out, v_out, k_scale, v_scale, loc)
jit_ms, _, _ = triton.testing.do_bench(jit_fn, quantiles=[0.5, 0.2, 0.8])
def fmt(ms):
return f"{ms*1000:6.2f}us {total_bytes/(ms*1e-3)/1e9:6.0f}GB/s"
print(f" sl={input_sl:5d} h={head:2d} d={dim:4d}" f" | jit {fmt(jit_ms)}")
def report_bandwidth():
print(f"\n{'='*95}")
print(" JIT (cast.cuh, 256 threads, 2D grid)")
print(f" dtype={DTYPE}, device={DEVICE}")
print(f"{'='*95}")
for sl in [64, 256, 1024, 2048]:
for h, d in [(8, 128), (32, 128), (8, 256), (32, 256)]:
_report_bandwidth(sl, h, d, DTYPE)
print()
if __name__ == "__main__":
benchmark.run(print_data=True)
report_bandwidth()

64
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_clamp_position.py vendored Normal file
View File

@@ -0,0 +1,64 @@
import itertools
import torch
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import (
DEFAULT_DEVICE,
get_benchmark_range,
run_benchmark,
)
from sglang.jit_kernel.clamp_position import clamp_position_cuda
from sglang.srt.utils import get_compiler_backend
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=13, suite="stage-b-kernel-benchmark-1-gpu-large")
SIZE_LIST = get_benchmark_range(
full_range=[2**n for n in range(4, 16)],
ci_range=[256, 4096],
)
configs = list(itertools.product(SIZE_LIST))
def _torch_clamp_position(seq_lens):
return torch.clamp(seq_lens - 1, min=0).to(torch.int64)
_compiled_clamp_position = torch.compile(
_torch_clamp_position, dynamic=True, backend=get_compiler_backend()
)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["size"],
x_vals=configs,
line_arg="provider",
line_vals=["jit", "torch_compile", "torch"],
line_names=["SGL JIT Kernel", "torch.compile", "PyTorch"],
styles=[("blue", "-"), ("green", "-."), ("red", "--")],
ylabel="us",
plot_name="clamp-position-performance",
args={},
)
)
def benchmark(size: int, provider: str):
seq_lens = torch.randint(
0, 10000, (size,), dtype=torch.int64, device=DEFAULT_DEVICE
)
if provider == "jit":
fn = lambda: clamp_position_cuda(seq_lens)
elif provider == "torch_compile":
fn = lambda: _compiled_clamp_position(seq_lens)
else:
fn = lambda: _torch_clamp_position(seq_lens)
return run_benchmark(fn)
if __name__ == "__main__":
benchmark.run(print_data=True)

162
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_concat_mla.py vendored Normal file
View File

@@ -0,0 +1,162 @@
import itertools
import torch
import triton
import triton.testing
from sgl_kernel import concat_mla_absorb_q as aot_absorb_q
from sgl_kernel import concat_mla_k as aot_k
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.jit_kernel.concat_mla import concat_mla_absorb_q as jit_absorb_q
from sglang.jit_kernel.concat_mla import concat_mla_k as jit_k
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(est_time=6, suite="stage-b-kernel-benchmark-1-gpu-large")
IS_CI = is_in_ci()
NUM_LOCAL_HEADS = 128
QK_NOPE_HEAD_DIM = 128
QK_ROPE_HEAD_DIM = 64
K_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM
A_LAST_DIM = 512
B_LAST_DIM = 64
DTYPE = torch.bfloat16
DEVICE = "cuda"
def aot_concat_mla_k(k, k_nope, k_rope):
aot_k(k, k_nope, k_rope)
def jit_concat_mla_k(k, k_nope, k_rope):
jit_k(k, k_nope, k_rope)
def torch_concat_mla_k(k, k_nope, k_rope):
nope_head_dim = k_nope.shape[-1]
k[:, :, :nope_head_dim] = k_nope
k[:, :, nope_head_dim:] = k_rope.expand(-1, k.shape[1], -1)
def aot_concat_mla_absorb_q(a, b):
return aot_absorb_q(a, b)
def jit_concat_mla_absorb_q(a, b):
return jit_absorb_q(a, b)
def torch_concat_mla_absorb_q(a, b, out):
a_last_dim = a.shape[-1]
out[:, :, :a_last_dim] = a
out[:, :, a_last_dim:] = b
if IS_CI:
NUM_TOKENS_VALS = [256, 1024]
else:
NUM_TOKENS_VALS = [256, 512, 1024, 2048, 4096, 8192, 16384, 32768]
K_LINE_VALS = ["aot", "jit", "torch"]
K_LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"]
K_STYLES = [("orange", "-"), ("blue", "--"), ("green", "-.")]
def _create_concat_mla_k_data(num_tokens):
"""Allocate oversized containers and slice to produce non-contiguous tensors."""
k_nope_container = torch.randn(
(num_tokens, NUM_LOCAL_HEADS, QK_NOPE_HEAD_DIM + 128),
dtype=DTYPE,
device=DEVICE,
)
k_nope = k_nope_container[:, :, :QK_NOPE_HEAD_DIM]
k_rope_container = torch.randn(
(num_tokens, 1, 128 + QK_ROPE_HEAD_DIM),
dtype=DTYPE,
device=DEVICE,
)
k_rope = k_rope_container[:, :, -QK_ROPE_HEAD_DIM:]
k = torch.empty(
(num_tokens, NUM_LOCAL_HEADS, K_HEAD_DIM),
dtype=DTYPE,
device=DEVICE,
)
return k, k_nope, k_rope
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["num_tokens"],
x_vals=NUM_TOKENS_VALS,
line_arg="provider",
line_vals=K_LINE_VALS,
line_names=K_LINE_NAMES,
styles=K_STYLES,
ylabel="us",
plot_name="concat-mla-k-performance",
args={},
)
)
def bench_concat_mla_k(num_tokens: int, provider: str):
k, k_nope, k_rope = _create_concat_mla_k_data(num_tokens)
FN_MAP = {
"aot": aot_concat_mla_k,
"jit": jit_concat_mla_k,
"torch": torch_concat_mla_k,
}
fn = lambda: FN_MAP[provider](k, k_nope, k_rope)
return run_benchmark(fn)
if IS_CI:
ABSORB_Q_VALS = list(itertools.product([4, 16], [16]))
else:
ABSORB_Q_VALS = list(itertools.product([1, 4, 8, 16, 32], [1, 8, 32, 128]))
Q_LINE_VALS = ["aot", "jit", "torch"]
Q_LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"]
Q_STYLES = [("orange", "-"), ("blue", "--"), ("green", "-.")]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["dim_0", "dim_1"],
x_vals=ABSORB_Q_VALS,
line_arg="provider",
line_vals=Q_LINE_VALS,
line_names=Q_LINE_NAMES,
styles=Q_STYLES,
ylabel="us",
plot_name="concat-mla-absorb-q-performance",
args={},
)
)
def bench_concat_mla_absorb_q(dim_0: int, dim_1: int, provider: str):
a = torch.randn(dim_0, dim_1, A_LAST_DIM, dtype=DTYPE, device=DEVICE)
b = torch.randn(dim_0, dim_1, B_LAST_DIM, dtype=DTYPE, device=DEVICE)
if provider == "torch":
out = torch.empty(
dim_0, dim_1, A_LAST_DIM + B_LAST_DIM, dtype=DTYPE, device=DEVICE
)
fn = lambda: torch_concat_mla_absorb_q(a, b, out)
else:
FN_MAP = {
"aot": aot_concat_mla_absorb_q,
"jit": jit_concat_mla_absorb_q,
}
fn = lambda: FN_MAP[provider](a, b)
return run_benchmark(fn)
if __name__ == "__main__":
bench_concat_mla_k.run(print_data=True)
bench_concat_mla_absorb_q.run(print_data=True)

384
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_custom_all_reduce.py vendored Normal file
View File

@@ -0,0 +1,384 @@
"""
Benchmark JIT custom all-reduce (v2) vs NCCL vs AOT custom all-reduce (v1).
Usage (torchrun required for multi-GPU):
torchrun --nproc_per_node=2 bench_custom_all_reduce.py
torchrun --nproc_per_node=4 bench_custom_all_reduce.py --dtype float16
torchrun --nproc_per_node=8 bench_custom_all_reduce.py --warmup 10 --iters 100
The script initializes all three backends, then benchmarks each over a sweep
of message sizes. Results are printed as a comparison table on rank 0.
"""
import argparse
import contextlib
import gc
import logging
import os
from math import isnan
from typing import Dict, List, Optional
import torch
import torch.distributed as dist
from sglang.jit_kernel.benchmark.utils import is_in_ci
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(
est_time=120,
suite="stage-b-kernel-benchmark-1-gpu-large",
disabled="requires multi-GPU, self-skips in CI",
)
DTYPE_MAP = {
"float16": torch.float16,
"bfloat16": torch.bfloat16,
"float32": torch.float32,
}
MESSAGE_SIZES_BYTES = [
4 * 1024, # 4K
16 * 1024, # 16K
64 * 1024, # 64K
128 * 1024, # 128K
3 * 64 * 1024, # 192K
4 * 64 * 1024, # 256K
3 * 128 * 1024, # 384K
4 * 128 * 1024, # 512K
5 * 128 * 1024, # 640K
6 * 128 * 1024, # 768K
7 * 128 * 1024, # 896K
1 * 1024 * 1024, # 1M
2 * 1024 * 1024, # 2M
3 * 1024 * 1024, # 2M
4 * 1024 * 1024, # 4M
8 * 1024 * 1024, # 8M
16 * 1024 * 1024, # 16M
32 * 1024 * 1024, # 32M
]
# ---------------------------------------------------------------------------
# Backend wrappers - each exposes a uniform interface:
# .name - display name
# .capture() - context manager for CUDA-graph recording
# .all_reduce() - perform an all-reduce and return the result tensor
# ---------------------------------------------------------------------------
class NCCLAllReduceBackend:
name = "NCCL"
def __init__(self, group: dist.ProcessGroup):
self.group = group
def capture(self, register_input: bool):
return contextlib.nullcontext()
def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
dist.all_reduce(tensor, group=self.group)
return tensor
class AOTAllReduceBackend:
name = "AOT"
def __init__(self, group: dist.ProcessGroup, device: torch.device):
from sglang.srt.distributed.device_communicators.custom_all_reduce import (
CustomAllreduce,
)
max_size = max(MESSAGE_SIZES_BYTES)
self.comm = CustomAllreduce(group, device, max_size=max_size)
if self.comm.disabled:
raise RuntimeError("AOT CustomAllreduce is disabled on this system")
def capture(self, register_input: bool):
return self.comm.capture() # ignore register_input since v1 always requires it
def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
assert self.comm.should_custom_ar(tensor), str(tensor.shape)
return self.comm.custom_all_reduce(tensor)
class JITAllReduceBackend:
name = "JIT"
def __init__(self, group: dist.ProcessGroup, device: torch.device):
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
CustomAllReduceV2,
)
max_size = max(MESSAGE_SIZES_BYTES)
self.comm = CustomAllReduceV2(group, device, max_pull_size=max_size)
if self.comm.disabled:
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
def capture(self, register_input: bool):
return self.comm.capture() if register_input else contextlib.nullcontext()
def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
assert self.comm.should_custom_ar(tensor), str(tensor.shape)
return self.comm.custom_all_reduce(tensor)
class FlashInferAllReduceBackend:
name = "FI"
def __init__(self, group: dist.ProcessGroup, dtype: torch.dtype):
import flashinfer.comm as comm
rank = torch.distributed.get_rank(group=group)
world_size = torch.distributed.get_world_size(group=group)
max_size = max(MESSAGE_SIZES_BYTES)
hidden_dim = min(MESSAGE_SIZES_BYTES) // 2
num_tokens = max_size // hidden_dim
self.comm = comm
self.hidden_dim = hidden_dim
self.workspace = comm.create_allreduce_fusion_workspace(
backend="trtllm",
world_size=world_size,
rank=rank,
max_token_num=num_tokens,
hidden_dim=hidden_dim,
dtype=dtype,
)
def capture(self, *_):
return contextlib.nullcontext()
def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
return self.comm.allreduce_fusion(
input=tensor.view(-1, self.hidden_dim),
workspace=self.workspace,
pattern=self.comm.AllReduceFusionPattern.kAllReduce,
launch_with_pdl=True,
fp32_acc=True,
)
# ---------------------------------------------------------------------------
# Benchmarking helpers
# ---------------------------------------------------------------------------
def parse_args():
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--dtype", choices=DTYPE_MAP.keys(), default="bfloat16")
p.add_argument("--warmup", type=int, default=5)
p.add_argument("--iters", type=int, default=50)
p.add_argument("--no-inplace", dest="register_input", action="store_false")
return p.parse_args()
@torch.inference_mode()
def bench_one(
backend,
inp: torch.Tensor,
warmup: int,
iters: int,
group: dist.ProcessGroup,
register_input: bool,
) -> float:
"""
Run *warmup* iterations of all-reduce first.
Return the average time for *iters* iterations of all-reduce.
"""
dist.barrier(group=group)
for _ in range(warmup):
backend.all_reduce(inp)
torch.cuda.synchronize()
# Capture a CUDA graph with *iters* all-reduce calls.
inp_batch = torch.stack([inp] * 4)
graph = torch.cuda.CUDAGraph()
with backend.capture(register_input):
with torch.cuda.graph(graph):
for i in range(iters):
backend.all_reduce(inp_batch[i % 4])
torch.cuda.synchronize()
# Warm up the graph once.
graph.replay()
# Timed replay.
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
torch.cuda.synchronize()
dist.barrier(group=group)
graph.replay() # make the stream busy
start.record()
graph.replay()
end.record()
torch.cuda.synchronize()
return start.elapsed_time(end) / iters
def bench_sweep(
backend,
sizes_bytes: List[int],
dtype: torch.dtype,
device: torch.device,
warmup: int,
iters: int,
group: dist.ProcessGroup,
register_input: bool,
) -> Dict[int, float]:
"""Benchmark one backend over all message sizes."""
elem_size = torch.tensor([], dtype=dtype).element_size()
results: Dict[int, float] = {}
for sz in sizes_bytes:
numel = sz // elem_size
inp = torch.zeros(numel, dtype=dtype, device=device)
try:
elapsed_ms = bench_one(backend, inp, warmup, iters, group, register_input)
results[sz] = elapsed_ms * 1000 # convert to us per iter
except AssertionError:
results[sz] = float("nan")
return results
# ---------------------------------------------------------------------------
# Result printing
# ---------------------------------------------------------------------------
def print_results(
backends: list,
all_results: Dict[str, Dict[int, float]],
sizes_bytes: List[int],
) -> None:
"""Print a comparison table on rank 0."""
def human_bytes(n: int) -> str:
for suffix, unit in [("M", 1 << 20), ("K", 1 << 10)]:
if n >= unit and n % unit == 0:
return f"{n // unit}{suffix}"
return f"{n}B"
def fmt_us(v: float) -> str:
return f"{v:13.1f}" if not isnan(v) else " n/a"
names = [b.name for b in backends]
nccl_name = "NCCL"
# Header
header_cols = [f"{n:>13}" for n in names]
speedup_cols = [f"{n:>13}/NCCL" for n in names if n != nccl_name]
header = f"{'Size':>8} " + " ".join(header_cols)
for sc in speedup_cols:
header += f" {sc}"
header += " "
print()
print(header)
print("-" * len(header))
# Rows
for sz in sizes_bytes:
row = f"{human_bytes(sz):>8}"
nccl_lat = all_results[nccl_name][sz]
for n in names:
row += f" {fmt_us(all_results[n][sz])}"
for n in names:
if n == nccl_name:
continue
lat = all_results[n][sz]
if not isnan(lat):
row += f" {nccl_lat / lat:17.2f}x"
else:
row += f" {'n/a':>17}"
print(row)
# ---------------------------------------------------------------------------
# Distributed setup
# ---------------------------------------------------------------------------
def init_distributed():
"""Initialize distributed groups using torchrun env vars.
Returns (rank, world_size, device, cpu_group, nccl_group).
"""
import sglang.srt.distributed.parallel_state as ps
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
world_size = int(os.environ.get("WORLD_SIZE", "1"))
rank = local_rank
device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device)
torch.cuda.set_stream(torch.cuda.Stream()) # use a non-default stream
torch.distributed.init_process_group(backend="gloo")
ps._WORLD = coord = ps.init_world_group(
ranks=list(range(world_size)),
local_rank=local_rank,
backend="nccl",
)
cpu_group = coord.cpu_group
nccl_group = coord.device_group
assert nccl_group is not None
return rank, world_size, device, cpu_group, nccl_group
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
logging.basicConfig(level=logging.WARNING)
args = parse_args()
dtype = DTYPE_MAP[args.dtype]
rank, world_size, device, cpu_group, nccl_group = init_distributed()
# Instantiate backends.
backends = [
NCCLAllReduceBackend(nccl_group),
JITAllReduceBackend(cpu_group, device),
]
if world_size in [2, 4, 6, 8]:
backends.insert(1, AOTAllReduceBackend(cpu_group, device))
if world_size in [2, 4, 8]:
backends.append(FlashInferAllReduceBackend(cpu_group, dtype))
# Run benchmarks.
all_results: Dict[str, Dict[int, float]] = {}
torch.cuda.synchronize()
for backend in backends:
if rank == 0:
print(f"Benchmarking {backend.name} ...")
all_results[backend.name] = bench_sweep(
backend,
MESSAGE_SIZES_BYTES,
dtype,
device,
args.warmup,
args.iters,
cpu_group,
args.register_input,
)
# Aggregate across ranks (use max to reflect the slowest rank).
for name in list(all_results):
for sz in MESSAGE_SIZES_BYTES:
val = all_results[name].get(sz)
if val is None:
continue
t = torch.tensor([val], dtype=torch.float64, device=device)
dist.all_reduce(t, op=dist.ReduceOp.MAX, group=nccl_group)
all_results[name][sz] = t.item()
# Print results on rank 0.
if rank == 0:
print_results(backends, all_results, MESSAGE_SIZES_BYTES)
del backends, all_results
gc.collect()
dist.destroy_process_group()
if __name__ == "__main__" and not is_in_ci():
main()

267
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_fused_qknorm_rope.py vendored Normal file
View File

@@ -0,0 +1,267 @@
"""
Benchmark: fused_qknorm_rope JIT vs AOT (sgl_kernel)
Measures throughput (µs) for fused_qk_norm_rope across typical
LLM configurations (head_dim × num_heads × num_tokens).
Run:
python python/sglang/jit_kernel/benchmark/bench_fused_qknorm_rope.py
"""
import itertools
import torch
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.fused_qknorm_rope import (
fused_qk_norm_rope as fused_qk_norm_rope_jit,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=6, suite="stage-b-kernel-benchmark-1-gpu-large")
try:
from sgl_kernel import fused_qk_norm_rope as fused_qk_norm_rope_aot
AOT_AVAILABLE = True
except ImportError:
fused_qk_norm_rope_aot = None
AOT_AVAILABLE = False
# ---------------------------------------------------------------------------
# Benchmark configuration
# ---------------------------------------------------------------------------
NUM_TOKENS_RANGE = get_benchmark_range(
full_range=[1, 64, 256, 1024, 4096],
ci_range=[64, 512],
)
# (head_dim, num_heads_q, num_heads_k, num_heads_v) — typical MoE/dense configs
MODEL_CONFIGS = get_benchmark_range(
full_range=[
(64, 32, 8, 8), # small
(128, 32, 8, 8), # typical (e.g. Qwen3-8B)
(256, 16, 4, 4), # large head_dim
],
ci_range=[(128, 32, 8, 8)],
)
# Real production shapes (self-attention; num_heads_k == num_heads_v == num_heads_q).
# Format: (name, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim, rotary_dim)
PRODUCTION_SHAPES = [
("flux_1024", 4096, 24, 24, 24, 128, 128),
("qwen_image_1024", 4096, 32, 32, 32, 128, 128),
("qwen_image_partial", 4096, 32, 32, 32, 128, 64),
("zimage_1024", 4096, 30, 30, 30, 128, 128),
("batch2_medium", 4096, 24, 24, 24, 128, 128), # B=2, T=2048
]
LINE_VALS = ["jit", "aot"] if AOT_AVAILABLE else ["jit"]
LINE_NAMES = ["JIT (new)", "AOT sgl_kernel"] if AOT_AVAILABLE else ["JIT (new)"]
STYLES = [("blue", "--"), ("orange", "-")] if AOT_AVAILABLE else [("blue", "--")]
# ---------------------------------------------------------------------------
# Benchmark: fused_qk_norm_rope (interleave style, no YaRN)
# ---------------------------------------------------------------------------
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["num_tokens", "head_dim", "num_heads_q", "num_heads_k", "num_heads_v"],
x_vals=[
(nt, hd, nq, nk, nv)
for nt, (hd, nq, nk, nv) in itertools.product(
NUM_TOKENS_RANGE, MODEL_CONFIGS
)
],
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="fused-qknorm-rope-performance",
args={},
)
)
def bench_fused_qknorm_rope(
num_tokens: int,
head_dim: int,
num_heads_q: int,
num_heads_k: int,
num_heads_v: int,
provider: str,
):
device = "cuda"
total_heads = num_heads_q + num_heads_k + num_heads_v
qkv = torch.randn(
(num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device
)
q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device)
common_kwargs = dict(
num_heads_q=num_heads_q,
num_heads_k=num_heads_k,
num_heads_v=num_heads_v,
head_dim=head_dim,
eps=1e-5,
q_weight=q_weight,
k_weight=k_weight,
base=10000.0,
is_neox=False,
position_ids=position_ids,
factor=1.0,
low=1.0,
high=32.0,
attention_factor=1.0,
rotary_dim=head_dim,
)
if provider == "jit":
fn = lambda: fused_qk_norm_rope_jit(qkv.clone(), **common_kwargs)
elif provider == "aot":
fn = lambda: fused_qk_norm_rope_aot(qkv.clone(), **common_kwargs)
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
# ---------------------------------------------------------------------------
# Benchmark: fused_qk_norm_rope — real production shapes (with speedup column)
# ---------------------------------------------------------------------------
def bench_fused_qknorm_rope_production():
device = "cuda"
header = f"{'name':<22} {'tokens':>6} {'nq':>4} {'nk':>4} {'nv':>4} {'hd':>4} {'rdim':>5} {'JIT(us)':>9} {'AOT(us)':>9} {'speedup':>8}"
sep = "-" * len(header)
print("\nfused-qknorm-rope-production-shapes:")
print(sep)
print(header)
print(sep)
for (
name,
num_tokens,
num_heads_q,
num_heads_k,
num_heads_v,
head_dim,
rotary_dim,
) in PRODUCTION_SHAPES:
total_heads = num_heads_q + num_heads_k + num_heads_v
qkv = torch.randn(
(num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device
)
q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device)
common_kwargs = dict(
num_heads_q=num_heads_q,
num_heads_k=num_heads_k,
num_heads_v=num_heads_v,
head_dim=head_dim,
eps=1e-5,
q_weight=q_weight,
k_weight=k_weight,
base=10000.0,
is_neox=False,
position_ids=position_ids,
factor=1.0,
low=1.0,
high=32.0,
attention_factor=1.0,
rotary_dim=rotary_dim,
)
jit_us, _, _ = run_benchmark(
lambda: fused_qk_norm_rope_jit(qkv.clone(), **common_kwargs)
)
if AOT_AVAILABLE:
aot_us, _, _ = run_benchmark(
lambda: fused_qk_norm_rope_aot(qkv.clone(), **common_kwargs)
)
speedup = f"{aot_us / jit_us:.2f}x"
aot_str = f"{aot_us:9.3f}"
else:
aot_str = f"{'N/A':>9}"
speedup = "N/A"
print(
f"{name:<22} {num_tokens:>6} {num_heads_q:>4} {num_heads_k:>4} {num_heads_v:>4}"
f" {head_dim:>4} {rotary_dim:>5} {jit_us:9.3f} {aot_str} {speedup:>8}"
)
print(sep)
# ---------------------------------------------------------------------------
# Quick correctness diff
# ---------------------------------------------------------------------------
def calculate_diff():
if not AOT_AVAILABLE:
print("sgl_kernel not available — skipping AOT diff check")
return
device = "cuda"
print("Correctness diff (JIT vs AOT):")
for head_dim, is_neox in [(64, False), (128, False), (128, True), (256, False)]:
num_tokens = 32
num_heads_q, num_heads_k, num_heads_v = 4, 2, 2
total_heads = num_heads_q + num_heads_k + num_heads_v
qkv = torch.randn(
(num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device
)
q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device)
common = dict(
num_heads_q=num_heads_q,
num_heads_k=num_heads_k,
num_heads_v=num_heads_v,
head_dim=head_dim,
eps=1e-5,
q_weight=q_weight,
k_weight=k_weight,
base=10000.0,
is_neox=is_neox,
position_ids=position_ids,
factor=1.0,
low=1.0,
high=32.0,
attention_factor=1.0,
rotary_dim=head_dim,
)
qkv_jit = qkv.clone()
fused_qk_norm_rope_jit(qkv_jit, **common)
qkv_aot = qkv.clone()
fused_qk_norm_rope_aot(qkv_aot, **common)
match = torch.allclose(qkv_jit.float(), qkv_aot.float(), atol=1e-2, rtol=1e-2)
status = "OK" if match else "MISMATCH"
max_err = (qkv_jit.float() - qkv_aot.float()).abs().max().item()
print(
f" head_dim={head_dim:3d} is_neox={str(is_neox):5s} "
f"max_err={max_err:.2e} [{status}]"
)
if __name__ == "__main__":
calculate_diff()
print()
bench_fused_qknorm_rope.run(print_data=True)
print()
bench_fused_qknorm_rope_production()

119
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_hadamard.py vendored Normal file
View File

@@ -0,0 +1,119 @@
import itertools
import math
from typing import Tuple
import torch
import torch.nn.functional as F
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import (
DEFAULT_DEVICE,
DEFAULT_DTYPE,
get_benchmark_range,
run_benchmark,
)
from sglang.jit_kernel.hadamard import hadamard_transform
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
# AOT kernel: might not be available in all environments.
# This is used for performance baseline comparison.
try:
from sgl_kernel import hadamard_transform as hadamard_transform_aot
AOT_AVAILABLE = True
except Exception:
AOT_AVAILABLE = False
# Naive reference implementation using scipy hadamard matrix.
try:
from scipy.linalg import hadamard
SCIPY_AVAILABLE = True
except ImportError:
SCIPY_AVAILABLE = False
# CI environment uses simplified parameters
batch_sizes = get_benchmark_range(
full_range=[1, 16, 64, 256],
ci_range=[16],
)
dim_range = get_benchmark_range(
full_range=[64, 256, 1024, 4096, 8192, 16384, 32768],
ci_range=[1024],
)
# Naive reference implementation using precomputed scipy hadamard matrix.
def torch_hadamard_transform(x, scale, H, dim, dim_padded):
flat = x.reshape(-1, dim)
if dim != dim_padded:
flat = F.pad(flat, (0, dim_padded - dim))
out = F.linear(flat, H) * scale
return out[..., :dim].reshape(x.shape)
available_providers = ["jit_kernel"]
available_names = ["JIT Kernel"]
available_styles = [("red", "-")]
if AOT_AVAILABLE:
available_providers.insert(0, "aot_kernel")
available_names.insert(0, "AOT Kernel")
available_styles.insert(0, ("green", "-"))
if SCIPY_AVAILABLE:
available_providers.append("naive")
available_names.append("Naive (scipy)")
available_styles.append(("blue", "-"))
configs = list(itertools.product(batch_sizes, dim_range))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size", "dim"],
x_vals=[list(c) for c in configs],
line_arg="provider",
line_vals=available_providers,
line_names=available_names,
styles=available_styles,
ylabel="us",
plot_name="hadamard-transform-performance",
args={},
)
)
def benchmark(batch_size: int, dim: int, provider: str) -> Tuple[float, float, float]:
scale = 1.0 / math.sqrt(dim)
x = torch.randn(batch_size, dim, device=DEFAULT_DEVICE, dtype=DEFAULT_DTYPE)
FN_MAP = {
"jit_kernel": lambda: hadamard_transform(x.clone(), scale=scale),
}
if AOT_AVAILABLE:
FN_MAP["aot_kernel"] = lambda: hadamard_transform_aot(x.clone(), scale=scale)
if SCIPY_AVAILABLE:
# Precompute Hadamard matrix on GPU to avoid CPU-GPU transfer
# during CUDA graph capture.
log_dim = math.ceil(math.log2(dim)) if dim > 0 else 0
dim_padded = 2**log_dim if dim > 0 else 1
H = torch.tensor(
hadamard(dim_padded, dtype=float),
dtype=DEFAULT_DTYPE,
device=DEFAULT_DEVICE,
)
FN_MAP["naive"] = lambda: torch_hadamard_transform(
x.clone(), scale, H, dim, dim_padded
)
fn = FN_MAP[provider]
return run_benchmark(fn)
if __name__ == "__main__":
print("=" * 80)
print("Benchmarking Fast Hadamard Transform")
print("=" * 80)
benchmark.run(print_data=True)

421
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_hicache.py vendored Normal file
View File

@@ -0,0 +1,421 @@
"""Benchmark for HiCache JIT kernel performance.
This benchmark tests the performance of KV cache transfer operations
between GPU and CPU (host pinned memory), comparing:
- SGL AOT Kernel: Pre-compiled transfer_kv kernels from sgl_kernel
- SGL JIT Kernel: JIT-compiled hicache kernels
- PyTorch Indexing: Plain PyTorch index copy
- PyTorch 2 Stream: PyTorch implementation using 2 CUDA streams
Tests cover:
- One Layer: CPU->GPU
- All Layer: GPU->CPU
Note: Uses do_bench instead of do_bench_cudagraph since CUDA graph
capture doesn't support CPU-GPU memory transfers.
"""
import itertools
import os
from dataclasses import dataclass
from typing import Tuple
import torch
import triton
import triton.testing
from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer
from sglang.jit_kernel.benchmark.utils import DEFAULT_QUANTILES, get_benchmark_range
from sglang.jit_kernel.hicache import (
can_use_hicache_jit_kernel,
transfer_hicache_all_layer,
transfer_hicache_one_layer,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=29, suite="stage-b-kernel-benchmark-1-gpu-large")
DISABLE_TORCH = os.environ.get("DISABLE_TORCH", "0") == "1"
PAGE_SIZE = 1
ENABLE_SORT = True
GPU_CACHE_SIZE = 256 * 1024 # 256K tokens on GPU
HOST_CACHE_SIZE = 512 * 1024 # 512K tokens on CPU
NUM_LAYERS = 8
@dataclass(frozen=True)
class HiCacheCache:
k_cache_cuda: torch.Tensor
v_cache_cuda: torch.Tensor
k_cache_host: torch.Tensor
v_cache_host: torch.Tensor
def get_slice(self, num_layers: int, element_size: int) -> "HiCacheCache":
def slice_cuda(t: torch.Tensor) -> torch.Tensor:
needed_cuda = num_layers * GPU_CACHE_SIZE
return t.view(-1, element_size)[:needed_cuda].unflatten(0, (num_layers, -1))
def slice_host(t: torch.Tensor) -> torch.Tensor:
needed_host = num_layers * HOST_CACHE_SIZE
return t.view(-1, element_size)[:needed_host].unflatten(0, (num_layers, -1))
return HiCacheCache(
k_cache_cuda=slice_cuda(self.k_cache_cuda),
v_cache_cuda=slice_cuda(self.v_cache_cuda),
k_cache_host=slice_host(self.k_cache_host),
v_cache_host=slice_host(self.v_cache_host),
)
def gen_indices(
size: int, max_size: int, *, page_size: int = PAGE_SIZE
) -> torch.Tensor:
def align(x: int) -> int:
return (x + page_size - 1) // page_size
assert size <= max_size and max_size % page_size == 0
indices = torch.randperm(align(max_size))[: align(size)]
offsets = torch.arange(page_size)
return (indices[:, None] * page_size + offsets).flatten().cuda()[:size]
def sglang_aot_transfer_one(
k_cache_dst: torch.Tensor,
v_cache_dst: torch.Tensor,
indices_dst: torch.Tensor,
k_cache_src: torch.Tensor,
v_cache_src: torch.Tensor,
indices_src: torch.Tensor,
item_size: int,
) -> None:
"""SGL AOT Kernel for single layer transfer."""
transfer_kv_per_layer(
k_cache_src,
k_cache_dst,
v_cache_src,
v_cache_dst,
indices_src,
indices_dst,
item_size,
)
def sglang_jit_transfer_one(
k_cache_dst: torch.Tensor,
v_cache_dst: torch.Tensor,
indices_dst: torch.Tensor,
k_cache_src: torch.Tensor,
v_cache_src: torch.Tensor,
indices_src: torch.Tensor,
element_dim: int,
) -> None:
"""SGL JIT Kernel for single layer transfer."""
transfer_hicache_one_layer(
k_cache_dst,
v_cache_dst,
indices_dst,
k_cache_src,
v_cache_src,
indices_src,
element_dim=element_dim,
)
def sglang_aot_transfer_all(
k_ptrs_dst: torch.Tensor,
v_ptrs_dst: torch.Tensor,
indices_dst: torch.Tensor,
k_ptrs_src: torch.Tensor,
v_ptrs_src: torch.Tensor,
indices_src: torch.Tensor,
item_size: int,
num_layers: int,
) -> None:
"""SGL AOT Kernel for all layer transfer."""
transfer_kv_all_layer(
k_ptrs_src,
k_ptrs_dst,
v_ptrs_src,
v_ptrs_dst,
indices_src,
indices_dst,
item_size,
num_layers,
)
def sglang_jit_transfer_all(
k_ptrs_dst: torch.Tensor,
v_ptrs_dst: torch.Tensor,
indices_dst: torch.Tensor,
k_ptrs_src: torch.Tensor,
v_ptrs_src: torch.Tensor,
indices_src: torch.Tensor,
stride_bytes: int,
element_size: int,
) -> None:
"""SGL JIT Kernel for all layer transfer."""
transfer_hicache_all_layer(
k_ptrs_dst,
v_ptrs_dst,
indices_dst,
k_ptrs_src,
v_ptrs_src,
indices_src,
kv_cache_src_stride_bytes=stride_bytes,
kv_cache_dst_stride_bytes=stride_bytes,
element_size=element_size,
)
def pytorch_transfer(
k_cache_dst: torch.Tensor,
v_cache_dst: torch.Tensor,
indices_dst_on_dst: torch.Tensor,
k_cache_src: torch.Tensor,
v_cache_src: torch.Tensor,
indices_src_on_src: torch.Tensor,
) -> None:
"""PyTorch indexing baseline."""
dst_device = k_cache_dst.device
k_cache_dst[indices_dst_on_dst] = k_cache_src[indices_src_on_src].to(dst_device)
v_cache_dst[indices_dst_on_dst] = v_cache_src[indices_src_on_src].to(dst_device)
# Benchmark configuration
BS_RANGE = get_benchmark_range(
full_range=[2**n for n in range(0, 16)],
ci_range=[16],
)
ELEMENT_SIZE_RANGE = get_benchmark_range(
full_range=[64, 128, 256, 512, 1024],
ci_range=[1024],
)
LINE_VALS = ["aot", "jit", "torch"]
LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"]
STYLES = [("orange", "-"), ("blue", "--"), ("red", ":")]
CONFIGS = list(itertools.product(ELEMENT_SIZE_RANGE, BS_RANGE))
# =============================================================================
# One Layer Benchmarks
# =============================================================================
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["element_size", "batch_size"],
x_vals=CONFIGS,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="hicache-one-layer-h2d",
args={},
)
)
def benchmark_one_layer_h2d(
element_size: int, batch_size: int, provider: str
) -> Tuple[float, float, float]:
"""One Layer: Host (CPU) -> Device (GPU)."""
global cache
cache_local = cache.get_slice(num_layers=NUM_LAYERS, element_size=element_size)
k_cache_src = cache_local.k_cache_host
v_cache_src = cache_local.v_cache_host
k_cache_dst = cache_local.k_cache_cuda
v_cache_dst = cache_local.v_cache_cuda
torch.manual_seed(batch_size * 65536 + element_size)
indices_src_gpu = gen_indices(batch_size, HOST_CACHE_SIZE)
indices_dst_gpu = gen_indices(batch_size, GPU_CACHE_SIZE)
if ENABLE_SORT:
indices_src_gpu, mapping = indices_src_gpu.sort()
indices_dst_gpu = indices_dst_gpu[mapping]
indices_src_cpu = indices_src_gpu.cpu()
torch.cuda.synchronize()
element_bytes = element_size * k_cache_src.element_size()
FN_MAP = {
"aot": lambda: [
sglang_aot_transfer_one(
k_cache_dst[i],
v_cache_dst[i],
indices_dst_gpu,
k_cache_src[i],
v_cache_src[i],
indices_src_gpu,
element_bytes,
)
for i in range(NUM_LAYERS)
],
"jit": lambda: [
sglang_jit_transfer_one(
k_cache_dst[i],
v_cache_dst[i],
indices_dst_gpu,
k_cache_src[i],
v_cache_src[i],
indices_src_gpu,
element_size,
)
for i in range(NUM_LAYERS)
],
"torch": lambda: [
pytorch_transfer(
k_cache_dst[i],
v_cache_dst[i],
indices_dst_gpu,
k_cache_src[i],
v_cache_src[i],
indices_src_cpu,
)
for i in range(NUM_LAYERS)
],
}
if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes):
return (float("nan"), float("nan"), float("nan"))
if DISABLE_TORCH and provider in ["torch"]:
return (float("nan"), float("nan"), float("nan"))
ms, min_ms, max_ms = triton.testing.do_bench( # type: ignore
FN_MAP[provider], quantiles=DEFAULT_QUANTILES, warmup=5, rep=25
)
return (
1000 * ms / NUM_LAYERS,
1000 * max_ms / NUM_LAYERS,
1000 * min_ms / NUM_LAYERS,
)
# =============================================================================
# All Layer Benchmarks
# =============================================================================
def _create_ptr_tensor(tensors, device="cuda"):
"""Create a tensor of data pointers."""
return torch.tensor(
[t.data_ptr() for t in tensors],
dtype=torch.uint64,
device=device,
)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["element_size", "batch_size"],
x_vals=CONFIGS,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="hicache-all-layer-d2h",
args={},
)
)
def benchmark_all_layer_d2h(
element_size: int, batch_size: int, provider: str
) -> Tuple[float, float, float]:
"""All Layer: Device (GPU) -> Host (CPU)."""
global cache
cache_local = cache.get_slice(num_layers=NUM_LAYERS, element_size=element_size)
k_caches_src = cache_local.k_cache_cuda
v_caches_src = cache_local.v_cache_cuda
k_caches_dst = cache_local.k_cache_host
v_caches_dst = cache_local.v_cache_host
torch.manual_seed(batch_size * 65536 + element_size)
indices_src_gpu = gen_indices(batch_size, GPU_CACHE_SIZE)
indices_dst_gpu = gen_indices(batch_size, HOST_CACHE_SIZE)
if ENABLE_SORT:
indices_dst_gpu, mapping = indices_dst_gpu.sort()
indices_src_gpu = indices_src_gpu[mapping]
indices_dst_cpu = indices_dst_gpu.cpu()
torch.cuda.synchronize()
element_bytes = element_size * k_caches_src.element_size()
k_ptrs_src = _create_ptr_tensor([k_caches_src[i] for i in range(NUM_LAYERS)])
v_ptrs_src = _create_ptr_tensor([v_caches_src[i] for i in range(NUM_LAYERS)])
k_ptrs_dst = _create_ptr_tensor([k_caches_dst[i] for i in range(NUM_LAYERS)])
v_ptrs_dst = _create_ptr_tensor([v_caches_dst[i] for i in range(NUM_LAYERS)])
FN_MAP = {
"aot": lambda: sglang_aot_transfer_all(
k_ptrs_dst,
v_ptrs_dst,
indices_dst_gpu,
k_ptrs_src,
v_ptrs_src,
indices_src_gpu,
element_bytes,
NUM_LAYERS,
),
"jit": lambda: sglang_jit_transfer_all(
k_ptrs_dst,
v_ptrs_dst,
indices_dst_gpu,
k_ptrs_src,
v_ptrs_src,
indices_src_gpu,
element_bytes,
element_bytes,
),
"torch": lambda: [
pytorch_transfer(
k_caches_dst[i],
v_caches_dst[i],
indices_dst_cpu,
k_caches_src[i],
v_caches_src[i],
indices_src_gpu,
)
for i in range(NUM_LAYERS)
],
}
if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes):
return (float("nan"), float("nan"), float("nan"))
if DISABLE_TORCH and provider in ["torch"]:
return (float("nan"), float("nan"), float("nan"))
ms, min_ms, max_ms = triton.testing.do_bench( # type: ignore
FN_MAP[provider], quantiles=DEFAULT_QUANTILES, warmup=5, rep=25
)
return (
1000 * ms / NUM_LAYERS,
1000 * max_ms / NUM_LAYERS,
1000 * min_ms / NUM_LAYERS,
)
if __name__ == "__main__":
MAX_SIZE = max(ELEMENT_SIZE_RANGE)
DEVICE_SHAPE = (NUM_LAYERS * GPU_CACHE_SIZE, MAX_SIZE)
HOST_SHAPE = (NUM_LAYERS * HOST_CACHE_SIZE, MAX_SIZE)
cache = HiCacheCache(
k_cache_cuda=torch.empty(DEVICE_SHAPE, dtype=torch.bfloat16, device="cuda"),
v_cache_cuda=torch.empty(DEVICE_SHAPE, dtype=torch.bfloat16, device="cuda"),
k_cache_host=torch.empty(HOST_SHAPE, dtype=torch.bfloat16, pin_memory=True),
v_cache_host=torch.empty(HOST_SHAPE, dtype=torch.bfloat16, pin_memory=True),
)
print("=" * 60)
print("One Layer: Host -> Device (CPU -> GPU)")
print("=" * 60)
benchmark_one_layer_h2d.run(print_data=True)
print("\n" + "=" * 60)
print("All Layer: Device -> Host (GPU -> CPU) [per-layer avg]")
print("=" * 60)
benchmark_all_layer_d2h.run(print_data=True)

100
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_norm.py vendored Normal file
View File

@@ -0,0 +1,100 @@
import itertools
import torch
import triton
import triton.testing
from flashinfer.norm import fused_add_rmsnorm as fi_fused_add_rmsnorm
from flashinfer.norm import rmsnorm as fi_rmsnorm
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, suite="stage-b-kernel-benchmark-1-gpu-large")
DTYPE = torch.bfloat16
DEVICE = "cuda"
BS_LIST = get_benchmark_range(
full_range=[2**n for n in range(0, 14)],
ci_range=[16, 32],
)
HIDDEN_SIZE_LIST = get_benchmark_range(
full_range=sorted([1536, *range(1024, 8192 + 1, 1024)]),
ci_range=[512, 2048],
)
LINE_VALS = ["flashinfer", "jit"]
LINE_NAMES = ["FlashInfer", "SGL JIT Kernel"]
STYLES = [("blue", "--"), ("green", "-.")]
NUM_LAYERS = 4 # avoid L2 effect
configs_0 = list(itertools.product(HIDDEN_SIZE_LIST + [16384], BS_LIST))
configs_1 = list(itertools.product(HIDDEN_SIZE_LIST, BS_LIST))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["hidden_size", "batch_size"],
x_vals=configs_0,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="rmsnorm-performance",
args={},
)
)
def benchmark_rmsnorm(hidden_size: int, batch_size: int, provider: str):
input = torch.randn(
(NUM_LAYERS, batch_size, hidden_size), dtype=DTYPE, device=DEVICE
)
weight = torch.randn((NUM_LAYERS, hidden_size), dtype=DTYPE, device=DEVICE)
FN_MAP = {"jit": jit_rmsnorm, "flashinfer": fi_rmsnorm}
def f():
fn = FN_MAP[provider]
for i in range(NUM_LAYERS):
fn(input[i], weight[i], out=input[i])
return run_benchmark(f, scale=NUM_LAYERS)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["hidden_size", "batch_size"],
x_vals=configs_1,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="fused-add-rmsnorm-performance",
args={},
)
)
def benchmark_fused_add_rmsnorm(hidden_size: int, batch_size: int, provider: str):
input = torch.randn(
(NUM_LAYERS, batch_size, hidden_size), dtype=DTYPE, device=DEVICE
)
residual = torch.randn_like(input)
weight = torch.randn((NUM_LAYERS, hidden_size), dtype=DTYPE, device=DEVICE)
FN_MAP = {"jit": jit_fused_add_rmsnorm, "flashinfer": fi_fused_add_rmsnorm}
def f():
fn = FN_MAP[provider]
for i in range(NUM_LAYERS):
fn(input[i], residual[i], weight[i])
return run_benchmark(f, scale=NUM_LAYERS)
if __name__ == "__main__":
print("Benchmarking rmsnorm...")
benchmark_rmsnorm.run(print_data=True)
print("Benchmarking fused_add_rmsnorm...")
benchmark_fused_add_rmsnorm.run(print_data=True)

261
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py vendored Normal file
View File

@@ -0,0 +1,261 @@
from __future__ import annotations
import sys
from typing import Any
import torch
import triton
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.nvfp4 import (
cutlass_fp4_group_mm,
scaled_fp4_experts_quant,
scaled_fp4_quant,
)
from sglang.srt.utils import is_sm100_supported
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
_NVFP4_SUPPORTED = is_sm100_supported()
def _round_up(x: int, y: int) -> int:
return ((x + y - 1) // y) * y
def _expert_offsets(m_per_expert: list[int], device: torch.device) -> torch.Tensor:
offsets = [0]
for m in m_per_expert:
offsets.append(offsets[-1] + m)
return torch.tensor(offsets, dtype=torch.int32, device=device)
def _blockscale_offsets(m_per_expert: list[int], device: torch.device) -> torch.Tensor:
offsets = [0]
for m in m_per_expert:
offsets.append(offsets[-1] + _round_up(m, 128))
return torch.tensor(offsets, dtype=torch.int32, device=device)
def _prepare_case(
total_tokens: int, n: int, k: int, num_experts: int, dtype: torch.dtype
) -> dict[str, Any]:
device = torch.device("cuda")
base = total_tokens // num_experts
rem = total_tokens % num_experts
m_per_expert = [base + (1 if i < rem else 0) for i in range(num_experts)]
expert_offsets_full = _expert_offsets(m_per_expert, device)
blockscale_offsets_full = _blockscale_offsets(m_per_expert, device)
a = torch.randn((total_tokens, k), device=device, dtype=dtype) * 0.1
b = torch.randn((num_experts, n, k), device=device, dtype=dtype) * 0.1
a_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32)
for i in range(num_experts):
start = int(expert_offsets_full[i].item())
end = int(expert_offsets_full[i + 1].item())
a_global_scale[i] = (
FLOAT8_E4M3_MAX
* FLOAT4_E2M1_MAX
/ a[start:end].abs().max().to(torch.float32)
)
b_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32)
for i in range(num_experts):
b_global_scale[i] = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b[i].abs().max().to(torch.float32)
)
a_fp4, a_blockscale = scaled_fp4_experts_quant(
a,
a_global_scale,
expert_offsets_full,
blockscale_offsets_full,
topk=1,
)
b_fp4 = torch.empty((num_experts, n, k // 2), device=device, dtype=torch.uint8)
b_blockscale = torch.empty(
(num_experts, _round_up(n, 128), _round_up(k // 16, 4)),
device=device,
dtype=torch.float8_e4m3fn,
)
for i in range(num_experts):
b_fp4_i, b_scale_i = scaled_fp4_quant(b[i], b_global_scale[i])
b_fp4[i].copy_(b_fp4_i)
b_blockscale[i].copy_(b_scale_i)
alphas = (1.0 / (a_global_scale * b_global_scale)).to(torch.float32)
params = {
"ab_strides": torch.full((num_experts,), k, dtype=torch.int64, device=device),
"c_strides": torch.full((num_experts,), n, dtype=torch.int64, device=device),
"problem_sizes": torch.tensor(
[[m, n, k] for m in m_per_expert], dtype=torch.int32, device=device
),
"expert_offsets": expert_offsets_full[:-1].contiguous(),
"blockscale_offsets": blockscale_offsets_full[:-1].contiguous(),
"a_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"b_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"out_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"a_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"b_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"alpha_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"layout_sfa": torch.empty((num_experts, 5), dtype=torch.int64, device=device),
"layout_sfb": torch.empty((num_experts, 5), dtype=torch.int64, device=device),
}
expert_ranges: list[tuple[int, int]] = []
start = 0
for m in m_per_expert:
end = start + m
expert_ranges.append((start, end))
start = end
return {
"a": a,
"b": b,
"a_fp4": a_fp4,
"b_fp4": b_fp4,
"a_blockscale": a_blockscale,
"b_blockscale": b_blockscale,
"alphas": alphas,
"params": params,
"expert_offsets_full": expert_offsets_full,
"expert_ranges": expert_ranges,
"dtype": dtype,
}
def _torch_ref_group_mm(case: dict[str, Any]) -> torch.Tensor:
a = case["a"]
b = case["b"]
dtype = case["dtype"]
expert_ranges = case["expert_ranges"]
total_tokens = a.shape[0]
n = b.shape[1]
out = torch.empty((total_tokens, n), device=a.device, dtype=dtype)
for i, (start, end) in enumerate(expert_ranges):
out[start:end] = torch.matmul(a[start:end], b[i].t())
return out
def _aot_cutlass_fp4_group_mm(case: dict[str, Any]) -> torch.Tensor:
a_fp4 = case["a_fp4"]
b_fp4 = case["b_fp4"]
a_blockscale = case["a_blockscale"]
b_blockscale = case["b_blockscale"]
alphas = case["alphas"]
params = case["params"]
out_dtype = case["dtype"]
out = torch.empty(
(a_fp4.shape[0], b_fp4.shape[1]), device=a_fp4.device, dtype=out_dtype
)
torch.ops.sgl_kernel.cutlass_fp4_group_mm.default(
out,
a_fp4,
b_fp4,
a_blockscale,
b_blockscale,
alphas,
params["ab_strides"],
params["c_strides"],
params["problem_sizes"],
params["expert_offsets"],
params["blockscale_offsets"],
)
return out
def _probe_legacy_aot_group_mm() -> tuple[bool, str]:
if not torch.cuda.is_available():
return False, "CUDA is not available."
if not _NVFP4_SUPPORTED:
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
try:
import sgl_kernel # noqa: F401
except Exception as e:
return False, f"import sgl_kernel failed: {e}"
if not hasattr(torch.ops, "sgl_kernel"):
return False, "torch.ops.sgl_kernel is not registered."
op = getattr(torch.ops.sgl_kernel, "cutlass_fp4_group_mm", None)
if op is None or not hasattr(op, "default"):
return False, "torch.ops.sgl_kernel.cutlass_fp4_group_mm.default is missing."
try:
case = _prepare_case(64, 256, 128, 4, torch.bfloat16)
_aot_cutlass_fp4_group_mm(case)
torch.cuda.synchronize()
except Exception as e:
return False, f"calling AOT grouped_mm op failed: {e}"
return True, ""
_AOT_GROUP_MM_AVAILABLE, _AOT_GROUP_MM_REASON = _probe_legacy_aot_group_mm()
shape_range = get_benchmark_range(
full_range=[(128, 256, 128, 4), (256, 512, 128, 8), (512, 512, 256, 8)],
ci_range=[(128, 256, 128, 4)],
)
line_vals = ["jit"]
line_names = ["JIT NVFP4 MoE GroupMM"]
styles = [("green", "-")]
if _AOT_GROUP_MM_AVAILABLE:
line_vals.append("aot_sgl_kernel")
line_names.append("AOT NVFP4 MoE GroupMM")
styles.append(("orange", "-"))
line_vals.append("torch_ref")
line_names.append("Torch Ref")
styles.append(("blue", "-"))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["total_tokens", "n", "k", "num_experts"],
x_vals=shape_range,
x_log=False,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="nvfp4-blockwise-moe-groupmm-performance",
args={},
)
)
def benchmark(total_tokens, n, k, num_experts, provider):
case = _prepare_case(total_tokens, n, k, num_experts, torch.bfloat16)
if provider == "jit":
fn = lambda: cutlass_fp4_group_mm(
case["a_fp4"],
case["b_fp4"],
case["a_blockscale"],
case["b_blockscale"],
case["alphas"],
case["dtype"],
case["params"],
)
elif provider == "aot_sgl_kernel":
fn = lambda: _aot_cutlass_fp4_group_mm(case)
elif provider == "torch_ref":
fn = lambda: _torch_ref_group_mm(case)
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
if not _NVFP4_SUPPORTED:
print("[skip] NVFP4 blockwise MoE benchmark requires sm100+ with CUDA 12.8+.")
sys.exit(0)
if not _AOT_GROUP_MM_AVAILABLE:
print(
f"[info] legacy AOT grouped_mm baseline unavailable: {_AOT_GROUP_MM_REASON}"
)
benchmark.run(print_data=True)

195
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py vendored Normal file
View File

@@ -0,0 +1,195 @@
from __future__ import annotations
import sys
import torch
import triton
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.nvfp4 import scaled_fp4_quant
from sglang.srt.utils import is_sm100_supported
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
BLOCK_SIZE = 16
_NVFP4_SUPPORTED = is_sm100_supported()
try:
from flashinfer import fp4_quantize as flashinfer_fp4_quantize
except Exception:
flashinfer_fp4_quantize = None
def _torch_ref_quant(input: torch.Tensor, input_global_scale: torch.Tensor):
m, n = input.shape
x = input.view(m, n // BLOCK_SIZE, BLOCK_SIZE)
vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32)
scale = input_global_scale * (vec_max / FLOAT4_E2M1_MAX)
scale = scale.to(torch.float8_e4m3fn).to(torch.float32)
output_scale = torch.where(scale == 0, torch.zeros_like(scale), 1.0 / scale)
scaled_x = x.to(torch.float32) * output_scale
clipped = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n)
rounded = clipped.clone()
rounded[(rounded >= 0.0) & (rounded <= 0.25)] = 0.0
rounded[(rounded > 0.25) & (rounded < 0.75)] = 0.5
rounded[(rounded >= 0.75) & (rounded <= 1.25)] = 1.0
rounded[(rounded > 1.25) & (rounded < 1.75)] = 1.5
rounded[(rounded >= 1.75) & (rounded <= 2.5)] = 2.0
rounded[(rounded > 2.5) & (rounded < 3.5)] = 3.0
rounded[(rounded >= 3.5) & (rounded <= 5.0)] = 4.0
rounded[rounded > 5.0] = 6.0
# This baseline intentionally keeps work on GPU but does not pack to uint8.
return rounded, scale
def _aot_scaled_fp4_quant(input: torch.Tensor, input_global_scale: torch.Tensor):
m, n = input.shape
output = torch.empty((m, n // 2), device=input.device, dtype=torch.uint8)
rounded_m = ((m + 128 - 1) // 128) * 128
scale_n = n // BLOCK_SIZE
rounded_n = ((scale_n + 4 - 1) // 4) * 4
output_scale = torch.empty(
(rounded_m, rounded_n // 4), device=input.device, dtype=torch.int32
)
torch.ops.sgl_kernel.scaled_fp4_quant.default(
output, input, output_scale, input_global_scale
)
return output, output_scale.view(torch.float8_e4m3fn)
def _probe_legacy_aot_quant() -> tuple[bool, str]:
if not torch.cuda.is_available():
return False, "CUDA is not available."
if not _NVFP4_SUPPORTED:
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
try:
import sgl_kernel # noqa: F401
except Exception as e:
return False, f"import sgl_kernel failed: {e}"
if not hasattr(torch.ops, "sgl_kernel"):
return False, "torch.ops.sgl_kernel is not registered."
op = getattr(torch.ops.sgl_kernel, "scaled_fp4_quant", None)
if op is None or not hasattr(op, "default"):
return False, "torch.ops.sgl_kernel.scaled_fp4_quant.default is missing."
try:
x = torch.randn((16, 64), dtype=torch.bfloat16, device="cuda")
global_scale = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.abs(x).max().to(torch.float32)
)
_aot_scaled_fp4_quant(x, global_scale)
torch.cuda.synchronize()
except Exception as e:
return False, f"calling AOT quant op failed: {e}"
return True, ""
_AOT_QUANT_AVAILABLE, _AOT_QUANT_REASON = _probe_legacy_aot_quant()
def _probe_flashinfer_quant() -> tuple[bool, str]:
if flashinfer_fp4_quantize is None:
return False, "import flashinfer.fp4_quantize failed."
if not torch.cuda.is_available():
return False, "CUDA is not available."
if not _NVFP4_SUPPORTED:
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
try:
x = torch.randn((16, 64), dtype=torch.bfloat16, device="cuda")
global_scale = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.abs(x).max().to(torch.float32)
)
flashinfer_fp4_quantize(
x,
global_scale,
BLOCK_SIZE, # sf_vec_size
False, # use_ue8m0
True, # is_sf_swizzled_layout
)
torch.cuda.synchronize()
except Exception as e:
return False, f"calling flashinfer.fp4_quantize failed: {e}"
return True, ""
_FLASHINFER_QUANT_AVAILABLE, _FLASHINFER_QUANT_REASON = _probe_flashinfer_quant()
shape_range = get_benchmark_range(
full_range=[(128, 2048), (512, 4096), (1024, 4096), (2048, 8192)],
ci_range=[(128, 2048)],
)
line_vals = []
line_names = []
styles = []
if _FLASHINFER_QUANT_AVAILABLE:
line_vals.append("flashinfer")
line_names.append("FlashInfer FP4 Quant")
styles.append(("purple", "-"))
line_vals.append("jit")
line_names.append("JIT NVFP4 Quant")
styles.append(("green", "-"))
if _AOT_QUANT_AVAILABLE:
line_vals.append("aot_sgl_kernel")
line_names.append("AOT NVFP4 Quant")
styles.append(("orange", "-"))
line_vals.append("torch_ref")
line_names.append("Torch Ref")
styles.append(("blue", "-"))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["m", "n"],
x_vals=shape_range,
x_log=False,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="nvfp4-quant-performance",
args={},
)
)
def benchmark(m, n, provider):
x = torch.randn((m, n), dtype=torch.bfloat16, device="cuda")
tensor_amax = torch.abs(x).max().to(torch.float32)
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
if provider == "jit":
fn = lambda: scaled_fp4_quant(x, global_scale)
elif provider == "flashinfer":
fn = lambda: flashinfer_fp4_quantize(
x,
global_scale,
BLOCK_SIZE, # sf_vec_size
False, # use_ue8m0
True, # is_sf_swizzled_layout
)
elif provider == "aot_sgl_kernel":
fn = lambda: _aot_scaled_fp4_quant(x, global_scale)
elif provider == "torch_ref":
fn = lambda: _torch_ref_quant(x, global_scale)
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
if not _NVFP4_SUPPORTED:
print("[skip] NVFP4 quant benchmark requires sm100+ with CUDA 12.8+.")
sys.exit(0)
if not _FLASHINFER_QUANT_AVAILABLE:
print(
f"[info] flashinfer quant baseline unavailable: {_FLASHINFER_QUANT_REASON}"
)
if not _AOT_QUANT_AVAILABLE:
print(f"[info] legacy AOT quant baseline unavailable: {_AOT_QUANT_REASON}")
benchmark.run(print_data=True)

187
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py vendored Normal file
View File

@@ -0,0 +1,187 @@
from __future__ import annotations
import sys
import torch
import triton
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
from sglang.srt.utils import is_sm100_supported, is_sm120_supported
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
BLOCK_SIZE = 16
_NVFP4_SUPPORTED = is_sm100_supported() or is_sm120_supported()
K_E2M1_TO_FLOAT = [
0.0,
0.5,
1.0,
1.5,
2.0,
3.0,
4.0,
6.0,
0.0,
-0.5,
-1.0,
-1.5,
-2.0,
-3.0,
-4.0,
-6.0,
]
def _dequantize_to_fp16(
tensor_fp4: torch.Tensor, tensor_sf: torch.Tensor, global_scale: torch.Tensor
):
m, packed_k = tensor_fp4.shape
k = packed_k * 2
flat = tensor_fp4.flatten()
high = (flat & 0xF0) >> 4
low = flat & 0x0F
f_h = torch.tensor([K_E2M1_TO_FLOAT[x] for x in high], device=tensor_fp4.device)
f_l = torch.tensor([K_E2M1_TO_FLOAT[x] for x in low], device=tensor_fp4.device)
val = torch.stack((f_l, f_h), dim=-1).reshape(m, k)
rounded_m = ((m + 128 - 1) // 128) * 128
scale_n = k // BLOCK_SIZE
rounded_n = ((scale_n + 4 - 1) // 4) * 4
sf = tensor_sf.view(torch.float8_e4m3fn)
tmp = torch.reshape(sf, (1, rounded_m // 128, rounded_n // 4, 32, 4, 4))
tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
scale = torch.reshape(tmp, (rounded_m, rounded_n))[:m, :scale_n].to(torch.float32)
scale = scale / global_scale
return (val.view(m, scale_n, BLOCK_SIZE) * scale.unsqueeze(-1)).reshape(m, k)
def _aot_cutlass_scaled_fp4_mm(
a: torch.Tensor,
b: torch.Tensor,
block_scale_a: torch.Tensor,
block_scale_b: torch.Tensor,
alpha: torch.Tensor,
out_dtype: torch.dtype,
) -> torch.Tensor:
out = torch.empty((a.shape[0], b.shape[0]), dtype=out_dtype, device=a.device)
torch.ops.sgl_kernel.cutlass_scaled_fp4_mm.default(
out, a, b, block_scale_a, block_scale_b, alpha
)
return out
def _probe_legacy_aot_scaled_mm() -> tuple[bool, str]:
if not torch.cuda.is_available():
return False, "CUDA is not available."
if not _NVFP4_SUPPORTED:
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
try:
import sgl_kernel # noqa: F401
except Exception as e:
return False, f"import sgl_kernel failed: {e}"
if not hasattr(torch.ops, "sgl_kernel"):
return False, "torch.ops.sgl_kernel is not registered."
op = getattr(torch.ops.sgl_kernel, "cutlass_scaled_fp4_mm", None)
if op is None or not hasattr(op, "default"):
return False, "torch.ops.sgl_kernel.cutlass_scaled_fp4_mm.default is missing."
try:
m, n, k = 16, 32, 64
a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda")
b = torch.randn((n, k), dtype=torch.bfloat16, device="cuda")
a_global_scale = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(a.flatten(), dim=-1)
).to(torch.float32)
b_global_scale = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(b.flatten(), dim=-1)
).to(torch.float32)
alpha = 1.0 / (a_global_scale * b_global_scale)
a_fp4, a_sf = scaled_fp4_quant(a, a_global_scale)
b_fp4, b_sf = scaled_fp4_quant(b, b_global_scale)
_aot_cutlass_scaled_fp4_mm(a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16)
torch.cuda.synchronize()
except Exception as e:
return False, f"calling AOT scaled_mm op failed: {e}"
return True, ""
_AOT_SCALED_MM_AVAILABLE, _AOT_SCALED_MM_REASON = _probe_legacy_aot_scaled_mm()
shape_range = get_benchmark_range(
full_range=[(128, 4096, 4096), (512, 4096, 4096), (1024, 8192, 4096)],
ci_range=[(128, 4096, 4096)],
)
line_vals = ["jit"]
line_names = ["JIT NVFP4 GEMM"]
styles = [("green", "-")]
if _AOT_SCALED_MM_AVAILABLE:
line_vals.append("aot_sgl_kernel")
line_names.append("AOT NVFP4 GEMM")
styles.append(("orange", "-"))
line_vals.append("torch_ref")
line_names.append("Torch Ref")
styles.append(("blue", "-"))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["m", "n", "k"],
x_vals=shape_range,
x_log=False,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="nvfp4-scaled-mm-performance",
args={},
)
)
def benchmark(m, n, k, provider):
a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda")
b = torch.randn((n, k), dtype=torch.bfloat16, device="cuda")
a_global_scale = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(a.flatten(), dim=-1)
).to(torch.float32)
b_global_scale = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(b.flatten(), dim=-1)
).to(torch.float32)
alpha = 1.0 / (a_global_scale * b_global_scale)
a_fp4, a_sf = scaled_fp4_quant(a, a_global_scale)
b_fp4, b_sf = scaled_fp4_quant(b, b_global_scale)
if provider == "jit":
fn = lambda: cutlass_scaled_fp4_mm(
a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16
)
elif provider == "aot_sgl_kernel":
fn = lambda: _aot_cutlass_scaled_fp4_mm(
a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16
)
elif provider == "torch_ref":
a_ref = _dequantize_to_fp16(a_fp4, a_sf, a_global_scale)
b_ref = _dequantize_to_fp16(b_fp4, b_sf, b_global_scale)
fn = lambda: torch.matmul(a_ref, b_ref.t())
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
if not _NVFP4_SUPPORTED:
print("[skip] NVFP4 scaled_mm benchmark requires sm100/sm120 with CUDA 12.8+.")
sys.exit(0)
if not _AOT_SCALED_MM_AVAILABLE:
print(
f"[info] legacy AOT scaled_mm baseline unavailable: {_AOT_SCALED_MM_REASON}"
)
benchmark.run(print_data=True)

120
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py vendored Normal file
View File

@@ -0,0 +1,120 @@
from typing import Optional, Tuple
import torch
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
try:
from vllm import _custom_ops as ops
VLLM_AVAILABLE = True
except ImportError:
ops = None
VLLM_AVAILABLE = False
try:
from sglang.srt.utils import is_hip
_is_hip = is_hip()
except ImportError:
_is_hip = False
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
def vllm_scaled_fp8_quant(
input: torch.Tensor,
scale: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
if not VLLM_AVAILABLE:
return sglang_scaled_fp8_quant(input, scale)
return ops.scaled_fp8_quant(input, scale)
def sglang_scaled_fp8_quant(
input: torch.Tensor,
scale: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
fp8_type_: torch.dtype = torch.float8_e4m3fn
output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
is_static = True
if scale is None:
scale = torch.zeros(1, device=input.device, dtype=torch.float32)
is_static = False
per_tensor_quant_fp8(input, output, scale, is_static)
return output, scale
def calculate_diff(batch_size: int, seq_len: int):
device = torch.device("cuda")
x = torch.rand((batch_size, seq_len), dtype=torch.bfloat16, device=device)
if not VLLM_AVAILABLE:
print("vLLM not available, skipping comparison")
return
vllm_out, vllm_scale = vllm_scaled_fp8_quant(x)
sglang_out, sglang_scale = sglang_scaled_fp8_quant(x)
vllm_out = vllm_out.to(torch.float32)
sglang_out = sglang_out.to(torch.float32)
triton.testing.assert_close(vllm_out, sglang_out, rtol=1e-3, atol=1e-3)
triton.testing.assert_close(vllm_scale, sglang_scale, rtol=1e-3, atol=1e-3)
# Benchmark configuration
element_range = get_benchmark_range(
full_range=[2**n for n in range(10, 20)],
ci_range=[16384],
)
if VLLM_AVAILABLE:
line_vals = ["vllm", "sglang"]
line_names = ["VLLM", "SGL Kernel"]
styles = [("blue", "-"), ("green", "-")]
else:
line_vals = ["sglang"]
line_names = ["SGL Kernel"]
styles = [("green", "-")]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["element_count"],
x_vals=element_range,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="per-tensor-quant-fp8-performance",
args={},
)
)
def benchmark(element_count, provider):
dtype = torch.float16
device = torch.device("cuda")
x = torch.randn(element_count, 4096, device=device, dtype=dtype)
if provider == "vllm":
fn = lambda: vllm_scaled_fp8_quant(x.clone())
elif provider == "sglang":
fn = lambda: sglang_scaled_fp8_quant(x.clone())
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
calculate_diff(batch_size=4, seq_len=4096)
benchmark.run(print_data=True)

287
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit.py vendored Normal file
View File

@@ -0,0 +1,287 @@
import itertools
from typing import Any, Dict, List
import torch
import triton
from sgl_kernel.test_utils import create_per_token_group_quant_test_data
from sglang.jit_kernel.benchmark.utils import get_benchmark_range
from sglang.jit_kernel.per_token_group_quant_8bit import (
per_token_group_quant_8bit as sglang_per_token_group_quant_8bit,
)
from sglang.srt.layers.quantization.fp8_kernel import (
create_per_token_group_quant_fp8_output_scale,
)
from sglang.srt.layers.quantization.fp8_kernel import (
per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
)
from sglang.srt.utils import is_hip
from sglang.srt.utils.bench_utils import bench_kineto
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(est_time=13, suite="stage-b-kernel-benchmark-1-gpu-large")
IS_CI = is_in_ci()
_is_hip = is_hip()
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
NUM_TESTS = 30 if IS_CI else 300
GROUP_SIZE_RANGE = [128]
DST_DTYPE_RANGE = [fp8_type_]
# ---- GEMM-like branch (num_ranks=None) ----
NUM_TOKENS_RANGE_GEMM = get_benchmark_range(
full_range=[1, 4, 16, 64, 256, 768, 2048, 8192, 16384],
ci_range=[768],
)
HIDDEN_DIM_RANGE_GEMM = [1536, 7168, 16384]
NUM_RANKS_RANGE_GEMM = [None]
FLAGS_GEMM_FULL: List[Dict[str, Any]] = [
dict(
column_major_scales=False,
scale_tma_aligned=False,
scale_ue8m0=False,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=False,
scale_ue8m0=False,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=False,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
]
FLAGS_GEMM_CI: List[Dict[str, Any]] = [
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
]
FLAGS_RANGE_GEMM = get_benchmark_range(
full_range=FLAGS_GEMM_FULL, ci_range=FLAGS_GEMM_CI
)
CONFIGS_GEMM = list(
itertools.product(
NUM_TOKENS_RANGE_GEMM,
HIDDEN_DIM_RANGE_GEMM,
GROUP_SIZE_RANGE,
NUM_RANKS_RANGE_GEMM,
DST_DTYPE_RANGE,
FLAGS_RANGE_GEMM,
)
)
# ---- MoE-like / multi-rank branch (hidden_dim=2048, num_ranks in {8,16,32,48}) ----
NUM_TOKENS_RANGE_MOE = get_benchmark_range(
full_range=[1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8],
ci_range=[768 * 8],
)
HIDDEN_DIM_RANGE_MOE = [2048]
NUM_RANKS_RANGE_MOE = get_benchmark_range(
full_range=[8, 16, 32, 48],
ci_range=[48],
)
FLAGS_MOE: List[Dict[str, Any]] = [
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="balanced",
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="imbalanced",
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="extreme",
),
]
FLAGS_RANGE_MOE = get_benchmark_range(full_range=FLAGS_MOE, ci_range=FLAGS_MOE)
CONFIGS_MOE = list(
itertools.product(
NUM_TOKENS_RANGE_MOE,
HIDDEN_DIM_RANGE_MOE,
GROUP_SIZE_RANGE,
NUM_RANKS_RANGE_MOE,
DST_DTYPE_RANGE,
FLAGS_RANGE_MOE,
)
)
# ---- Final configs ----
CONFIGS = CONFIGS_GEMM + CONFIGS_MOE
LINE_VALS = ["triton", "sglang"]
LINE_NAMES = ["Triton (Inaccurate)", "SGL Kernel"]
STYLES = [("blue", "-"), ("green", "-")]
def _flatten_to_2d(t: torch.Tensor) -> torch.Tensor:
"""Reshape a tensor with 3+ dims to 2D by merging all leading dims."""
if t.ndim <= 2:
return t
return t.reshape(-1, t.shape[-1])
def _make_sglang_bench_fn(
x: torch.Tensor,
group_size: int,
dst_dtype: torch.dtype,
flags: dict,
):
"""
Adapter that pre-allocates output tensors and returns a zero-arg callable
matching the JIT kernel's signature.
The JIT kernel does not support fuse_silu_and_mul, so when enabled we
pre-compute silu+mul on the input. bench_kineto only times the kernel
matching the given name, so the pre-processing is not included.
The JIT kernel expects 2D tensors, so any higher-dimensional inputs
(e.g. from masked_layout_mode) are flattened to 2D.
"""
fuse_silu_and_mul = flags.get("fuse_silu_and_mul", False)
column_major_scales = flags.get("column_major_scales", False)
scale_tma_aligned = flags.get("scale_tma_aligned", False)
scale_ue8m0 = flags.get("scale_ue8m0", False)
# JIT kernel does not support fuse_silu_and_mul; pre-compute it
if fuse_silu_and_mul:
half = x.shape[-1] // 2
x_input = torch.nn.functional.silu(x[..., :half]) * x[..., half:]
else:
x_input = x
# JIT kernel expects 2D (num_tokens, hidden_dim); flatten if needed
x_input = _flatten_to_2d(x_input.contiguous())
out_shape = x_input.shape
output_q = torch.empty(out_shape, device=x.device, dtype=dst_dtype)
fp8_max = torch.finfo(dst_dtype).max
fp8_min = -fp8_max
output_s = create_per_token_group_quant_fp8_output_scale(
x_shape=out_shape,
device=x.device,
group_size=group_size,
column_major_scales=column_major_scales,
scale_tma_aligned=scale_tma_aligned,
scale_ue8m0=scale_ue8m0,
)
def _run():
sglang_per_token_group_quant_8bit(
input=x_input,
output_q=output_q,
output_s=output_s,
group_size=group_size,
eps=1e-10,
fp8_min=fp8_min,
fp8_max=fp8_max,
scale_ue8m0=scale_ue8m0,
)
return _run
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=[
"num_tokens",
"hidden_dim",
"group_size",
"num_ranks",
"dst_dtype",
"flags",
],
x_vals=CONFIGS,
line_arg="provider",
line_vals=LINE_VALS,
# Triton has multi kernels and we only report the time for the core one
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="per-token-group-quant-8bit-performance",
args={},
)
)
def benchmark(
num_tokens, hidden_dim, group_size, num_ranks, dst_dtype, flags, provider
):
print(
f"Testing: {num_tokens=} {hidden_dim=} {group_size=} {num_ranks=} {dst_dtype=} {flags=} {provider=}"
)
x, masked_m = create_per_token_group_quant_test_data(
num_tokens=num_tokens, hidden_dim=hidden_dim, num_ranks=num_ranks, flags=flags
)
if provider == "triton":
fn = triton_per_token_group_quant_8bit
kernel_names = "_per_token_group_quant_8bit|_silu_and_mul_post_quant_kernel"
bench_fn = lambda: fn(
x=x,
masked_m=masked_m,
group_size=group_size,
dst_dtype=dst_dtype,
**{k: v for k, v in flags.items() if k not in ["masked_layout_mode"]},
)
elif provider == "sglang":
kernel_names = "per_token_group_quant_8bit_kernel"
bench_fn = _make_sglang_bench_fn(
x=x,
group_size=group_size,
dst_dtype=dst_dtype,
flags=flags,
)
else:
raise ValueError(f"Unknown provider: {provider}")
time_s = bench_kineto(bench_fn, kernel_names=kernel_names, num_tests=NUM_TESTS)
return time_s * 1e6
if __name__ == "__main__":
benchmark.run(print_data=True)

140
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_qknorm.py vendored Normal file
View File

@@ -0,0 +1,140 @@
import itertools
import torch
import triton
import triton.testing
from sgl_kernel import rmsnorm
from sglang.jit_kernel.benchmark.utils import (
DEFAULT_DEVICE,
DEFAULT_DTYPE,
get_benchmark_range,
run_benchmark,
)
from sglang.jit_kernel.norm import fused_inplace_qknorm
from sglang.srt.utils import get_current_device_stream_fast
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=10, suite="stage-b-kernel-benchmark-1-gpu-large")
alt_stream = torch.cuda.Stream()
def sglang_aot_qknorm(
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
) -> None:
head_dim = q.shape[-1]
q = q.view(-1, head_dim)
k = k.view(-1, head_dim)
current_stream = get_current_device_stream_fast()
alt_stream.wait_stream(current_stream)
rmsnorm(q, q_weight, out=q)
with torch.cuda.stream(alt_stream):
rmsnorm(k, k_weight, out=k)
current_stream.wait_stream(alt_stream)
def sglang_jit_qknorm(
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
) -> None:
fused_inplace_qknorm(q, k, q_weight, k_weight)
def flashinfer_qknorm(
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
) -> None:
from flashinfer import rmsnorm
rmsnorm(q, q_weight, out=q)
rmsnorm(k, k_weight, out=k)
@torch.compile()
def torch_impl_qknorm(
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
eps: float = 1e-6,
) -> None:
q_mean = q.float().pow(2).mean(dim=-1, keepdim=True)
k_mean = k.float().pow(2).mean(dim=-1, keepdim=True)
q_norm = (q_mean + eps).rsqrt()
k_norm = (k_mean + eps).rsqrt()
q.copy_(q.float() * q_norm * q_weight.float())
k.copy_(k.float() * k_norm * k_weight.float())
BS_RANGE = get_benchmark_range(
full_range=[2**n for n in range(0, 14)],
ci_range=[16],
)
GQA_RANGE = get_benchmark_range(
full_range=[4, 8],
ci_range=[4],
)
KV_HEAD_RANGE = get_benchmark_range(
full_range=[1, 2, 4, 8],
ci_range=[1],
)
HEAD_DIM_RANGE = get_benchmark_range(
full_range=[128, 256, 512, 1024],
ci_range=[128],
)
LINE_VALS = ["aot", "jit", "flashinfer", "torch"]
LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "FlashInfer", "PyTorch"]
STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")]
configs = list(itertools.product(HEAD_DIM_RANGE, GQA_RANGE, KV_HEAD_RANGE, BS_RANGE))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["head_dim", "GQA", "num_kv_heads", "batch_size"],
x_vals=configs,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="qknorm-performance",
args={},
)
)
def benchmark(
head_dim: int, GQA: int, num_kv_heads: int, batch_size: int, provider: str
):
num_qo_heads = GQA * num_kv_heads
q = torch.randn(
(batch_size, num_qo_heads, head_dim), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
)
k = torch.randn(
(batch_size, num_kv_heads, head_dim), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
)
q_weight = torch.randn(head_dim, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE)
k_weight = torch.randn(head_dim, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE)
FN_MAP = {
"aot": sglang_aot_qknorm,
"jit": sglang_jit_qknorm,
"flashinfer": flashinfer_qknorm,
"torch": torch_impl_qknorm,
}
fn = lambda: FN_MAP[provider](q, k, q_weight, k_weight)
return run_benchmark(fn)
if __name__ == "__main__":
benchmark.run(print_data=True)

123
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py vendored Normal file
View File

@@ -0,0 +1,123 @@
import itertools
from typing import Tuple
import torch
import triton
import triton.testing
from sgl_kernel import rmsnorm
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads
from sglang.srt.utils import get_current_device_stream_fast
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(est_time=12, suite="stage-b-kernel-benchmark-1-gpu-large")
IS_CI = is_in_ci()
alt_stream = torch.cuda.Stream()
def sglang_jit_qknorm_across_heads(
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
) -> None:
fused_inplace_qknorm_across_heads(q, k, q_weight, k_weight)
def sglang_aot_qknorm_across_heads(
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
) -> None:
current_stream = get_current_device_stream_fast()
alt_stream.wait_stream(current_stream)
rmsnorm(q, q_weight, out=q)
with torch.cuda.stream(alt_stream):
rmsnorm(k, k_weight, out=k)
current_stream.wait_stream(alt_stream)
def flashinfer_qknorm_across_heads(
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
) -> None:
from flashinfer import rmsnorm
rmsnorm(q, q_weight, out=q)
rmsnorm(k, k_weight, out=k)
@torch.compile()
def torch_impl_qknorm_across_heads(
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
eps: float = 1e-6,
) -> None:
q_mean = q.float().pow(2).mean(dim=-1, keepdim=True)
k_mean = k.float().pow(2).mean(dim=-1, keepdim=True)
q_norm = (q_mean + eps).rsqrt()
k_norm = (k_mean + eps).rsqrt()
q.copy_(q.float() * q_norm * q_weight.float())
k.copy_(k.float() * k_norm * k_weight.float())
DTYPE = torch.bfloat16
DEVICE = "cuda"
if IS_CI:
BS_RANGE = [16]
HIDDEN_DIM_RANGE = [1024]
else:
BS_RANGE = [2**n for n in range(0, 14)]
HIDDEN_DIM_RANGE = [512, 1024, 2048, 4096, 8192]
LINE_VALS = ["jit", "aot", "flashinfer", "torch"]
LINE_NAMES = ["SGL JIT Kernel", "SGL AOT Kernel", "FlashInfer", "PyTorch"]
STYLES = [("blue", "-"), ("orange", "--"), ("green", "-."), ("red", ":")]
configs = list(itertools.product(BS_RANGE, HIDDEN_DIM_RANGE))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size", "hidden_dim"],
x_vals=configs,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="qknorm-across-heads-performance",
args={},
)
)
def benchmark(
batch_size: int, hidden_dim: int, provider: str
) -> Tuple[float, float, float]:
q = torch.randn((batch_size, hidden_dim), dtype=DTYPE, device=DEVICE)
k = torch.randn((batch_size, hidden_dim), dtype=DTYPE, device=DEVICE)
q_weight = torch.randn(hidden_dim, dtype=DTYPE, device=DEVICE)
k_weight = torch.randn(hidden_dim, dtype=DTYPE, device=DEVICE)
FN_MAP = {
"jit": sglang_jit_qknorm_across_heads,
"aot": sglang_aot_qknorm_across_heads,
"flashinfer": flashinfer_qknorm_across_heads,
"torch": torch_impl_qknorm_across_heads,
}
fn = lambda: FN_MAP[provider](q, k, q_weight, k_weight)
return run_benchmark(fn)
if __name__ == "__main__":
benchmark.run(print_data=True)

239
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_renorm.py vendored Normal file
View File

@@ -0,0 +1,239 @@
import itertools
import sgl_kernel
import torch
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
def torch_top_k_renorm_probs(probs, top_k):
"""Vectorized PyTorch implementation of top-k renormalization."""
batch_size, vocab_size = probs.shape
# Handle scalar or tensor k
if isinstance(top_k, int):
k_val = min(max(top_k, 1), vocab_size)
# Get top-k indices for all batches at once
_, topk_indices = torch.topk(probs, k_val, dim=1, largest=True)
# Create mask: batch_size x vocab_size
mask = torch.zeros_like(probs)
mask.scatter_(1, topk_indices, 1.0)
# Vectorized renormalization
masked_probs = probs * mask
renorm_probs = masked_probs / (masked_probs.sum(dim=1, keepdim=True) + 1e-10)
return renorm_probs
else:
# Variable k per batch - need to handle separately
renorm_probs = torch.zeros_like(probs)
for i in range(batch_size):
k_val = min(max(top_k[i].item(), 1), vocab_size)
_, topk_indices = torch.topk(probs[i], k_val, largest=True)
mask = torch.zeros_like(probs[i])
mask[topk_indices] = 1.0
masked_probs = probs[i] * mask
renorm_probs[i] = masked_probs / (masked_probs.sum() + 1e-10)
return renorm_probs
def torch_top_p_renorm_probs(probs, top_p, eps=1e-5):
"""Vectorized PyTorch implementation of top-p renormalization."""
batch_size, vocab_size = probs.shape
# Handle scalar or tensor p
if isinstance(top_p, float):
p_val = top_p
# Vectorized implementation for uniform top_p
# Sort probs in descending order
sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=1)
cumsum_probs = torch.cumsum(sorted_probs, dim=1)
# Find cutoff: where cumsum exceeds top_p
cutoff_mask = cumsum_probs <= p_val
# Keep at least one token (the highest prob)
cutoff_mask[:, 0] = True
# Create mask in original order
mask = torch.zeros_like(probs)
mask.scatter_(1, sorted_indices, cutoff_mask.float())
# Vectorized renormalization
masked_probs = probs * mask
renorm_probs = masked_probs / (masked_probs.sum(dim=1, keepdim=True) + eps)
return renorm_probs
else:
# Variable p per batch - need to handle separately
renorm_probs = torch.zeros_like(probs)
for i in range(batch_size):
p_val = top_p[i].item()
sorted_prob, indices = torch.sort(probs[i], descending=False)
cdf = torch.cumsum(sorted_prob, dim=-1)
mask = torch.zeros(vocab_size, dtype=torch.float32, device=probs.device)
mask.scatter_(0, indices, (cdf >= (1 - p_val) - eps).float())
masked_probs = probs[i] * mask
renorm_probs[i] = masked_probs / (masked_probs.sum() + eps)
return renorm_probs
def calculate_diff_top_k_renorm(batch_size, vocab_size, k):
"""Compare Torch reference and SGLang kernel for top-k renorm correctness."""
torch.manual_seed(42)
device = torch.device("cuda")
pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32)
torch_output = torch_top_k_renorm_probs(probs, top_k_tensor)
sglang_output = sgl_kernel.top_k_renorm_prob(probs, top_k_tensor)
torch.testing.assert_close(torch_output, sglang_output, rtol=1e-3, atol=1e-3)
def calculate_diff_top_p_renorm(batch_size, vocab_size, p):
"""Compare Torch reference and SGLang kernel for top-p renorm correctness."""
torch.manual_seed(42)
device = torch.device("cuda")
pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
top_p_tensor = torch.full((batch_size,), p, device=device, dtype=torch.float32)
torch_output = torch_top_p_renorm_probs(probs, top_p_tensor)
sglang_output = sgl_kernel.top_p_renorm_prob(probs, top_p_tensor)
torch.testing.assert_close(torch_output, sglang_output, rtol=1e-3, atol=1e-3)
# Parameter space - simplified for CI
if is_in_ci():
batch_size_range = [16]
vocab_size_range = [111]
k_range = [10]
p_range = [0.5]
else:
batch_size_range = [16, 64, 128]
vocab_size_range = [111, 32000, 128256]
k_range = [10, 100, 500]
p_range = [0.1, 0.5, 0.9]
configs_k = list(itertools.product(batch_size_range, vocab_size_range, k_range))
configs_p = list(itertools.product(batch_size_range, vocab_size_range, p_range))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size", "vocab_size", "k"],
x_vals=configs_k,
line_arg="provider",
line_vals=["torch", "sglang"],
line_names=["Torch Reference", "SGL Kernel"],
styles=[("red", "-"), ("green", "-")],
ylabel="us",
plot_name="top-k-renorm-probs-performance",
args={},
)
)
def benchmark_top_k_renorm(batch_size, vocab_size, k, provider):
# Skip invalid configurations
if k >= vocab_size:
return float("nan"), float("nan"), float("nan")
torch.manual_seed(42)
device = torch.device("cuda")
pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32)
if provider == "torch":
fn = lambda: torch_top_k_renorm_probs(probs.clone(), top_k_tensor)
elif provider == "sglang":
fn = lambda: sgl_kernel.top_k_renorm_prob(probs.clone(), top_k_tensor)
return run_benchmark_no_cudagraph(fn)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size", "vocab_size", "p"],
x_vals=configs_p,
line_arg="provider",
line_vals=["torch", "sglang"],
line_names=["Torch Reference", "SGL Kernel"],
styles=[("red", "-"), ("blue", "-")],
ylabel="us",
plot_name="top-p-renorm-probs-performance",
args={},
)
)
def benchmark_top_p_renorm(batch_size, vocab_size, p, provider):
torch.manual_seed(42)
device = torch.device("cuda")
pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
top_p_tensor = torch.full((batch_size,), p, device=device, dtype=torch.float32)
if provider == "torch":
fn = lambda: torch_top_p_renorm_probs(probs.clone(), top_p_tensor)
elif provider == "sglang":
fn = lambda: sgl_kernel.top_p_renorm_prob(probs.clone(), top_p_tensor)
return run_benchmark_no_cudagraph(fn)
if __name__ == "__main__":
print("=" * 60)
print("Running correctness checks...")
print("=" * 60)
# Correctness checks - simplified for CI
if is_in_ci():
test_configs_k = [configs_k[0]] if configs_k else [(16, 111, 10)]
test_configs_p = [configs_p[0]] if configs_p else [(16, 111, 0.5)]
else:
test_configs_k = configs_k[:3] # Test first 3 configs
test_configs_p = configs_p[:3]
print("\n1. Testing top_k_renorm_probs...")
for cfg in test_configs_k:
batch_size, vocab_size, k = cfg
if k < vocab_size: # Skip invalid configs
calculate_diff_top_k_renorm(batch_size, vocab_size, k)
print(
f" ✓ Passed: batch_size={batch_size}, vocab_size={vocab_size}, k={k}"
)
print("\n2. Testing top_p_renorm_probs...")
for cfg in test_configs_p:
calculate_diff_top_p_renorm(*cfg)
batch_size, vocab_size, p = cfg
print(f" ✓ Passed: batch_size={batch_size}, vocab_size={vocab_size}, p={p}")
print("\n" + "=" * 60)
print("All correctness checks passed!")
print("=" * 60)
print("\n" + "=" * 60)
print("Starting performance benchmarks...")
print("=" * 60)
print("\n1. Benchmarking top_k_renorm_probs...")
benchmark_top_k_renorm.run(print_data=True)
print("\n2. Benchmarking top_p_renorm_probs...")
benchmark_top_p_renorm.run(print_data=True)
print("\n" + "=" * 60)
print("Benchmarking complete!")
print("=" * 60)

72
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_resolve_future_token_ids.py vendored Normal file
View File

@@ -0,0 +1,72 @@
import itertools
import torch
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import (
DEFAULT_DEVICE,
get_benchmark_range,
run_benchmark,
)
from sglang.jit_kernel.resolve_future_token_ids import resolve_future_token_ids_cuda
from sglang.srt.utils import get_compiler_backend
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=10, suite="stage-b-kernel-benchmark-1-gpu-large")
SIZE_LIST = get_benchmark_range(
full_range=[2**n for n in range(4, 16)], # 16 … 32K elements
ci_range=[256, 4096],
)
configs = list(itertools.product(SIZE_LIST))
def _torch_resolve(input_ids, future_map):
input_ids[:] = torch.where(
input_ids < 0,
future_map[torch.clamp(-input_ids, min=0)],
input_ids,
)
_compiled_resolve = torch.compile(
_torch_resolve, dynamic=True, backend=get_compiler_backend()
)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["size"],
x_vals=configs,
line_arg="provider",
line_vals=["jit", "torch_compile", "torch"],
line_names=["SGL JIT Kernel", "torch.compile", "PyTorch"],
styles=[("blue", "-"), ("green", "-."), ("red", "--")],
ylabel="us",
plot_name="resolve-future-token-ids-performance",
args={},
)
)
def benchmark(size: int, provider: str):
map_size = 8192
future_map = torch.randint(
0, 50000, (map_size,), dtype=torch.int64, device=DEFAULT_DEVICE
)
input_ids = torch.randint(
-map_size + 1, 50000, (size,), dtype=torch.int64, device=DEFAULT_DEVICE
)
if provider == "jit":
fn = lambda: resolve_future_token_ids_cuda(input_ids.clone(), future_map)
elif provider == "torch_compile":
fn = lambda: _compiled_resolve(input_ids.clone(), future_map)
else:
fn = lambda: _torch_resolve(input_ids.clone(), future_map)
return run_benchmark(fn)
if __name__ == "__main__":
benchmark.run(print_data=True)

308
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_rope.py vendored Normal file
View File

@@ -0,0 +1,308 @@
import itertools
import torch
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import (
DEFAULT_DEVICE,
DEFAULT_DTYPE,
get_benchmark_range,
run_benchmark,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=6, suite="stage-b-kernel-benchmark-1-gpu-large")
MAX_SEQ_LEN = 131072
ROPE_BASE = 10000.0
ROPE_DIM = 128
CACHE_SIZE = 1024 * 1024
def create_cos_sin_cache(
rotary_dim: int = ROPE_DIM,
max_position: int = MAX_SEQ_LEN,
base: float = ROPE_BASE,
) -> torch.Tensor:
inv_freq = 1.0 / (
base
** (
torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=DEFAULT_DEVICE)
/ rotary_dim
)
)
t = torch.arange(max_position, dtype=torch.float32, device=DEFAULT_DEVICE)
freqs = torch.einsum("i,j->ij", t, inv_freq)
cos = freqs.cos()
sin = freqs.sin()
return torch.cat((cos, sin), dim=-1)
# Pre-build the cache once
COS_SIN_CACHE = create_cos_sin_cache()
# ---------------------------------------------------------------------------
# RoPE-only provider implementations
# ---------------------------------------------------------------------------
def flashinfer_rope(
q: torch.Tensor,
k: torch.Tensor,
positions: torch.Tensor,
is_neox: bool,
) -> None:
from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
head_size = q.shape[-1]
apply_rope_with_cos_sin_cache_inplace(
positions=positions,
query=q.view(q.shape[0], -1),
key=k.view(k.shape[0], -1),
head_size=head_size,
cos_sin_cache=COS_SIN_CACHE,
is_neox=is_neox,
)
def sglang_pos_enc_rope(
q: torch.Tensor,
k: torch.Tensor,
positions: torch.Tensor,
is_neox: bool,
) -> None:
from sglang.jit_kernel.rope import rotary_embedding_with_key
head_size = q.shape[-1]
rotary_embedding_with_key(
positions=positions,
query=q.view(q.shape[0], -1),
key=k.view(k.shape[0], -1),
head_size=head_size,
cos_sin_cache=COS_SIN_CACHE,
is_neox=is_neox,
)
def sglang_fused_rope(
q: torch.Tensor,
k: torch.Tensor,
positions: torch.Tensor,
is_neox: bool,
) -> None:
from sglang.jit_kernel.rope import apply_rope_inplace
apply_rope_inplace(q, k, COS_SIN_CACHE, positions, is_neox=is_neox)
# ---------------------------------------------------------------------------
# RoPE + KV cache store provider implementations
# ---------------------------------------------------------------------------
def jit_rope_then_store(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
positions: torch.Tensor,
out_loc: torch.Tensor,
is_neox: bool,
) -> None:
from sglang.jit_kernel.kvcache import store_cache
from sglang.jit_kernel.rope import apply_rope_inplace
head_size = q.shape[-1]
row_dim = k.shape[-2] * head_size
apply_rope_inplace(
positions=positions,
q=q,
k=k,
rope_dim=head_size,
cos_sin_cache=COS_SIN_CACHE,
is_neox=is_neox,
)
store_cache(
k.view(-1, row_dim),
v.view(-1, row_dim),
k_cache,
v_cache,
out_loc,
)
def jit_fused_rope_store(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
positions: torch.Tensor,
out_loc: torch.Tensor,
is_neox: bool,
) -> None:
from sglang.jit_kernel.rope import apply_rope_inplace_with_kvcache
apply_rope_inplace_with_kvcache(
q, k, v, k_cache, v_cache, COS_SIN_CACHE, positions, out_loc, is_neox=is_neox
)
# ---------------------------------------------------------------------------
# Benchmark configuration (shared)
# ---------------------------------------------------------------------------
BS_RANGE = get_benchmark_range(
full_range=[2**n for n in range(0, 16)],
ci_range=[16],
)
QK_HEAD_RANGE = get_benchmark_range(
full_range=[(8, 1), (16, 2), (32, 8)],
ci_range=[(16, 2)],
)
QK_HEAD_RANGE = [f"{q},{k}" for q, k in QK_HEAD_RANGE]
IS_NEOX_RANGE = get_benchmark_range(
full_range=[True, False],
ci_range=[True],
)
# ---------------------------------------------------------------------------
# Benchmark 1: RoPE only
# ---------------------------------------------------------------------------
ROPE_LINE_VALS = ["flashinfer", "jit_pos_enc", "jit_fused_rope"]
ROPE_LINE_NAMES = [
"FlashInfer",
"SGL JIT PosEnc",
"SGL JIT Fused RoPE",
]
ROPE_STYLES = [("green", "-."), ("red", "-"), ("blue", "--")]
rope_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["num_q_k_heads", "is_neox", "batch_size"],
x_vals=rope_configs,
line_arg="provider",
line_vals=ROPE_LINE_VALS,
line_names=ROPE_LINE_NAMES,
styles=ROPE_STYLES,
ylabel="us",
plot_name="rope-performance",
args={},
)
)
def benchmark(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str):
qo, kv = num_q_k_heads.split(",")
num_qo_heads = int(qo)
num_kv_heads = int(kv)
q = torch.randn(
(batch_size, num_qo_heads, ROPE_DIM),
dtype=DEFAULT_DTYPE,
device=DEFAULT_DEVICE,
)
k = torch.randn(
(batch_size, num_kv_heads, ROPE_DIM),
dtype=DEFAULT_DTYPE,
device=DEFAULT_DEVICE,
)
seed = batch_size << 16 | num_qo_heads << 8 | num_kv_heads << 4 | is_neox
torch.random.manual_seed(seed)
positions = torch.randint(
MAX_SEQ_LEN, (batch_size,), device=DEFAULT_DEVICE, dtype=torch.int64
)
torch.cuda.synchronize()
FN_MAP = {
"flashinfer": flashinfer_rope,
"jit_pos_enc": sglang_pos_enc_rope,
"jit_fused_rope": sglang_fused_rope,
}
fn = lambda: FN_MAP[provider](q, k, positions, is_neox)
return run_benchmark(fn)
# ---------------------------------------------------------------------------
# Benchmark 2: RoPE + KV cache store
# ---------------------------------------------------------------------------
STORE_LINE_VALS = ["jit_rope_then_store", "jit_fused_store"]
STORE_LINE_NAMES = [
"SGL JIT RoPE + Store",
"SGL JIT Fused RoPE + Store",
]
STORE_STYLES = [("red", "-"), ("blue", "--")]
store_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["num_q_k_heads", "is_neox", "batch_size"],
x_vals=store_configs,
line_arg="provider",
line_vals=STORE_LINE_VALS,
line_names=STORE_LINE_NAMES,
styles=STORE_STYLES,
ylabel="us",
plot_name="rope-store-performance",
args={},
)
)
def benchmark_store(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str):
qo, kv = num_q_k_heads.split(",")
num_qo_heads = int(qo)
num_kv_heads = int(kv)
q = torch.randn(
(batch_size, num_qo_heads, ROPE_DIM),
dtype=DEFAULT_DTYPE,
device=DEFAULT_DEVICE,
)
k = torch.randn(
(batch_size, num_kv_heads, ROPE_DIM),
dtype=DEFAULT_DTYPE,
device=DEFAULT_DEVICE,
)
v = torch.randn(
(batch_size, num_kv_heads, ROPE_DIM),
dtype=DEFAULT_DTYPE,
device=DEFAULT_DEVICE,
)
row_size = num_kv_heads * ROPE_DIM
k_cache = torch.zeros(
CACHE_SIZE, row_size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
)
v_cache = torch.zeros(
CACHE_SIZE, row_size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
)
out_loc = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE, dtype=torch.int64)[
:batch_size
]
seed = batch_size << 16 | num_qo_heads << 8 | num_kv_heads << 4 | is_neox
torch.random.manual_seed(seed)
positions = torch.randint(
MAX_SEQ_LEN, (batch_size,), device=DEFAULT_DEVICE, dtype=torch.int64
)
torch.cuda.synchronize()
FN_MAP = {
"jit_rope_then_store": jit_rope_then_store,
"jit_fused_store": jit_fused_rope_store,
}
fn = lambda: FN_MAP[provider](
q, k, v, k_cache, v_cache, positions, out_loc, is_neox
)
return run_benchmark(fn)
if __name__ == "__main__":
print("Running RoPE performance benchmark...")
benchmark.run(print_data=True)
print("\nRunning RoPE + KV cache store performance benchmark...")
benchmark_store.run(print_data=True)

133
third_party/sglang/python/sglang/jit_kernel/benchmark/bench_store_cache.py vendored Normal file
View File

@@ -0,0 +1,133 @@
import itertools
from typing import Tuple
import torch
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import (
DEFAULT_DEVICE,
DEFAULT_DTYPE,
DEFAULT_QUANTILES,
get_benchmark_range,
)
from sglang.jit_kernel.kvcache import store_cache
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=9, suite="stage-b-kernel-benchmark-1-gpu-large")
def sglang_jit_store_cache(
k: torch.Tensor,
v: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
indices: torch.Tensor,
) -> None:
store_cache(k, v, k_cache, v_cache, indices)
@torch.compile()
def torch_compile_store_cache(
k: torch.Tensor,
v: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
indices: torch.Tensor,
) -> None:
k_cache[indices] = k
v_cache[indices] = v
alt_stream = torch.cuda.Stream()
def torch_streams_store_cache(
k: torch.Tensor,
v: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
indices: torch.Tensor,
) -> None:
current_stream = torch.cuda.current_stream()
alt_stream.wait_stream(current_stream)
k_cache[indices] = k
with torch.cuda.stream(alt_stream):
v_cache[indices] = v
current_stream.wait_stream(alt_stream)
NUM_LAYERS = 8
CACHE_SIZE = 2 * 1024 * 1024 // NUM_LAYERS
BS_RANGE = get_benchmark_range(
full_range=[2**n for n in range(0, 15)],
ci_range=[16],
)
ITEM_SIZE = get_benchmark_range(
full_range=[64, 128, 256, 512, 1024],
ci_range=[1024],
)
LINE_VALS = ["jit", "torch_compile", "torch_streams"]
LINE_NAMES = ["SGL JIT Kernel", "PyTorch Compile", "PyTorch 2 Stream"]
STYLES = [("blue", "--"), ("red", ":"), ("green", "-.")]
X_NAMES = ["item_size", "batch_size"]
CONFIGS = list(itertools.product(ITEM_SIZE, BS_RANGE))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=X_NAMES,
x_vals=CONFIGS,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="store-kvcache-performance",
args={},
)
)
def benchmark(
batch_size: int, item_size: int, provider: str
) -> Tuple[float, float, float]:
k = torch.randn(
(NUM_LAYERS, batch_size, item_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
)
v = torch.randn(
(NUM_LAYERS, batch_size, item_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
)
k_cache = torch.randn(
(NUM_LAYERS, CACHE_SIZE, item_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
)
v_cache = torch.randn(
(NUM_LAYERS, CACHE_SIZE, item_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
)
indices = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE)[:batch_size]
torch.cuda.synchronize()
FN_MAP = {
"jit": sglang_jit_store_cache,
"torch_compile": torch_compile_store_cache,
"torch_streams": torch_streams_store_cache,
}
def fn():
impl = FN_MAP[provider]
for i in range(NUM_LAYERS):
impl(k[i], v[i], k_cache[i], v_cache[i], indices)
# Custom time calculation: divide by NUM_LAYERS
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
fn, quantiles=DEFAULT_QUANTILES
)
return (
1000 * ms / NUM_LAYERS,
1000 * max_ms / NUM_LAYERS,
1000 * min_ms / NUM_LAYERS,
)
if __name__ == "__main__":
benchmark.run(print_data=True)

363
third_party/sglang/python/sglang/jit_kernel/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py vendored Normal file
View File

@@ -0,0 +1,363 @@
import argparse
import csv
import json
import os
import re
import statistics
from pathlib import Path
from typing import Any, Callable
import flashinfer
import sgl_kernel
import torch
from sglang.jit_kernel.benchmark.utils import DEFAULT_DTYPE
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(
est_time=120,
suite="stage-b-kernel-benchmark-1-gpu-large",
disabled="standalone diffusion NVFP4 benchmark",
)
SCRIPT_DIR = Path(__file__).resolve().parent
REPO_ROOT = (
Path(os.environ["SGLANG_NVFP4_REPO_ROOT"])
if os.environ.get("SGLANG_NVFP4_REPO_ROOT")
else Path(__file__).resolve().parents[5]
)
DEFAULT_OUTPUT_DIR = REPO_ROOT / "outputs" / "nvfp4_benchmarks"
DEFAULT_SHAPE_LIBRARY = SCRIPT_DIR / "diffusion_nvfp4_shapes.json"
DTYPE = DEFAULT_DTYPE
WARMUP = 8
ITERS = 20
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
METHODS = ("cutlass", "flashinfer_auto", "flashinfer_cudnn")
def benchmark_provider(
fn: Callable[[], torch.Tensor],
warmup: int = WARMUP,
iters: int = ITERS,
) -> tuple[float, float, float]:
for _ in range(warmup):
y = fn()
del y
torch.cuda.synchronize()
times_ms: list[float] = []
for _ in range(iters):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
y = fn()
end.record()
end.synchronize()
times_ms.append(start.elapsed_time(end))
del y
return statistics.median(times_ms), max(times_ms), min(times_ms)
def make_global_scale(x: torch.Tensor) -> torch.Tensor:
max_abs = torch.amax(x.abs()).clamp_min_(1e-6)
return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / max_abs).to(torch.float32)
def build_quantized_inputs(
m: int,
n: int,
k: int,
device: torch.device,
seed: int,
) -> dict[str, Any]:
assert k % 16 == 0, f"NVFP4 requires k % 16 == 0, got k={k}"
gen = torch.Generator(device=device)
gen.manual_seed(seed)
x = torch.randn((m, k), device=device, dtype=DTYPE, generator=gen)
w = torch.randn((n, k), device=device, dtype=DTYPE, generator=gen)
x_global_scale = make_global_scale(x)
w_global_scale = make_global_scale(w)
alpha = (1.0 / (x_global_scale * w_global_scale)).to(torch.float32)
x_fp4, x_sf = flashinfer.fp4_quantize(x, x_global_scale)
w_fp4, w_sf = flashinfer.fp4_quantize(w, w_global_scale)
if x_sf.dtype == torch.uint8:
x_sf = x_sf.view(torch.float8_e4m3fn)
if w_sf.dtype == torch.uint8:
w_sf = w_sf.view(torch.float8_e4m3fn)
return {
"x_fp4": x_fp4,
"w_fp4": w_fp4,
"x_sf": x_sf,
"w_sf": w_sf,
"alpha": alpha,
}
def make_shape_id(
model: str, shape_kind: str, prefix: str, m: int, n: int, k: int
) -> str:
prefix_slug = re.sub(r"[^a-zA-Z0-9]+", "_", prefix).strip("_")
return f"{model}_{shape_kind}_{prefix_slug}_{m}x{n}x{k}"
def load_shape_cases(shape_library: Path) -> list[dict[str, Any]]:
payload = json.loads(shape_library.read_text(encoding="utf-8"))
if not isinstance(payload, dict) or not payload:
raise RuntimeError(
f"Expected a non-empty model->shape list mapping in {shape_library}."
)
rows: list[dict[str, Any]] = []
for model, shapes in payload.items():
if not isinstance(shapes, list):
raise RuntimeError(
f"Expected {model} to map to a list of shapes in {shape_library}."
)
for shape in shapes:
m, n, k = (int(x) for x in shape["shape"])
count = int(shape["count"])
shape_kind = str(shape.get("kind", "actual_runtime_linear"))
prefix = str(shape.get("prefix", ""))
rows.append(
{
"shape_id": make_shape_id(model, shape_kind, prefix, m, n, k),
"source_model": model,
"shape_kind": shape_kind,
"runtime_prefix": prefix,
"m": m,
"n": n,
"k": k,
"count": count,
"approx_flops": 2 * m * n * k * count,
}
)
if not rows:
raise RuntimeError(f"No shapes found in {shape_library}.")
return rows
def split_csv_arg(text: str | None) -> set[str]:
if text is None or not text.strip():
return set()
return {item.strip() for item in text.split(",") if item.strip()}
def select_shape_cases(
rows: list[dict[str, Any]],
*,
models: set[str],
shape_kinds: set[str],
top_k: int,
rank_by: str,
) -> list[dict[str, Any]]:
filtered = [
row
for row in rows
if (not models or row["source_model"] in models)
and (not shape_kinds or row["shape_kind"] in shape_kinds)
]
key = "approx_flops" if rank_by == "flops" else "count"
return sorted(filtered, key=lambda row: int(row[key]), reverse=True)[:top_k]
def write_csv(rows: list[dict[str, Any]], output_path: Path) -> None:
with output_path.open("w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(
f,
fieldnames=[
"shape_id",
"source_model",
"shape_kind",
"runtime_prefix",
"m",
"n",
"k",
"count",
"approx_flops",
"method",
"median_ms",
"min_ms",
"max_ms",
"tflops",
],
)
writer.writeheader()
writer.writerows(rows)
def write_markdown(rows: list[dict[str, Any]], output_path: Path) -> None:
shape_rows = []
seen_shape_ids = set()
for row in rows:
if row["shape_id"] in seen_shape_ids:
continue
seen_shape_ids.add(row["shape_id"])
shape_rows.append(row)
lines: list[str] = []
lines.append("# Diffusion NVFP4 Scaled MM Benchmark")
lines.append("")
lines.append("## Shape Cases")
lines.append("")
lines.append("| Shape ID | Model | Shape Kind | Calls | Shape `(M,N,K)` | Prefix |")
lines.append("|---|---|---|---:|---|---|")
for row in shape_rows:
lines.append(
f"| {row['shape_id']} | {row['source_model']} | {row['shape_kind']} | {row['count']} | `({row['m']}, {row['n']}, {row['k']})` | `{row['runtime_prefix']}` |"
)
lines.append("")
for shape_row in shape_rows:
shape_id = shape_row["shape_id"]
scoped = [row for row in rows if row["shape_id"] == shape_id]
lines.append(f"## {shape_id}")
lines.append("")
lines.append("| Method | Median ms | TFLOPS |")
lines.append("|---|---:|---:|")
for row in sorted(scoped, key=lambda item: float(item["median_ms"])):
lines.append(
f"| {row['method']} | {float(row['median_ms']):.4f} | {float(row['tflops']):.1f} |"
)
lines.append("")
output_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
def run_shape_suite(shape_cases: list[dict[str, Any]]) -> list[dict[str, Any]]:
device = torch.device("cuda")
rows: list[dict[str, Any]] = []
for idx, shape in enumerate(shape_cases):
m = int(shape["m"])
n = int(shape["n"])
k = int(shape["k"])
quantized = build_quantized_inputs(m, n, k, device, seed=idx)
metadata = {
"shape_id": str(shape["shape_id"]),
"source_model": str(shape["source_model"]),
"shape_kind": str(shape["shape_kind"]),
"runtime_prefix": str(shape["runtime_prefix"]),
"m": m,
"n": n,
"k": k,
"count": int(shape["count"]),
"approx_flops": int(shape["approx_flops"]),
}
providers: dict[str, Callable[[], torch.Tensor]] = {
"cutlass": lambda: sgl_kernel.cutlass_scaled_fp4_mm(
quantized["x_fp4"],
quantized["w_fp4"],
quantized["x_sf"],
quantized["w_sf"],
quantized["alpha"],
DTYPE,
),
"flashinfer_auto": lambda: flashinfer.mm_fp4(
quantized["x_fp4"],
quantized["w_fp4"].T,
quantized["x_sf"],
quantized["w_sf"].T,
quantized["alpha"],
DTYPE,
backend="auto",
),
"flashinfer_cudnn": lambda: flashinfer.mm_fp4(
quantized["x_fp4"],
quantized["w_fp4"].T,
quantized["x_sf"],
quantized["w_sf"].T,
quantized["alpha"],
DTYPE,
backend="cudnn",
),
}
for method in METHODS:
median_ms, max_ms, min_ms = benchmark_provider(providers[method])
rows.append(
{
**metadata,
"method": method,
"median_ms": median_ms,
"min_ms": min_ms,
"max_ms": max_ms,
"tflops": (2 * m * n * k) / (median_ms / 1e3) / 1e12,
}
)
return rows
def main() -> None:
parser = argparse.ArgumentParser(
description="Benchmark diffusion NVFP4 GEMM backends on the captured diffusion shape library."
)
parser.add_argument(
"--models",
help="Comma-separated source_model filter. Default: all models in the JSON shape library.",
)
parser.add_argument(
"--shape-kinds",
help="Comma-separated shape_kind filter. Default: benchmark every shape kind in the JSON shape library.",
)
parser.add_argument(
"--top-k",
type=int,
default=64,
help="Benchmark the top-k shapes after filtering and ranking.",
)
parser.add_argument(
"--rank-by",
choices=["flops", "count"],
default="flops",
help="How to rank shapes before selecting top-k.",
)
parser.add_argument(
"--output-dir",
default=str(DEFAULT_OUTPUT_DIR),
help="Directory for CSV/Markdown outputs.",
)
args = parser.parse_args()
if is_in_ci():
print("Skipping bench_diffusion_nvfp4_scaled_mm.py in CI")
return
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required for NVFP4 scaled mm benchmarks.")
if not DEFAULT_SHAPE_LIBRARY.exists():
raise RuntimeError(
f"Shape library not found at {DEFAULT_SHAPE_LIBRARY}. "
"Commit or copy the generated diffusion_nvfp4_shapes.json first."
)
shape_cases = load_shape_cases(DEFAULT_SHAPE_LIBRARY)
selected_shapes = select_shape_cases(
shape_cases,
models=split_csv_arg(args.models),
shape_kinds=split_csv_arg(args.shape_kinds),
top_k=args.top_k,
rank_by=args.rank_by,
)
if not selected_shapes:
raise RuntimeError("No shapes matched the requested filters.")
rows = run_shape_suite(selected_shapes)
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
csv_path = output_dir / "diffusion_nvfp4_scaled_mm.csv"
md_path = output_dir / "diffusion_nvfp4_scaled_mm_summary.md"
write_csv(rows, csv_path)
write_markdown(rows, md_path)
print(f"Wrote {csv_path}")
print(f"Wrote {md_path}")
if __name__ == "__main__":
main()

138
third_party/sglang/python/sglang/jit_kernel/benchmark/diffusion/bench_fused_norm_scale_shift.py vendored Normal file
View File

@@ -0,0 +1,138 @@
# Benchmarks SGLang fused layernorm/rmsnorm scale shift kernels
# 1. fused_norm_scale_shift
# 2. fused_scale_residual_norm_scale_shift
import itertools
from typing import Tuple
import torch
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
from sglang.multimodal_gen.runtime.layers.layernorm import (
LayerNormScaleShift,
RMSNormScaleShift,
ScaleResidualLayerNormScaleShift,
ScaleResidualRMSNormScaleShift,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(
est_time=17,
suite="stage-b-kernel-benchmark-1-gpu-large",
disabled="Temporarily skipped to unblock flashinfer upgrade. Ref: https://github.com/sgl-project/sglang/actions/runs/23735552939/job/69139238979?pr=21422",
)
if is_in_ci():
B_RANGE, S_RANGE, D_RANGE = [1], [128], [1024]
else:
B_RANGE, S_RANGE, D_RANGE = [1], [128, 1024, 4096], [1024, 3072, 4096]
NORM_TYPE_RANGE = ["layer", "rms"]
AFFINE_RANGE = [True, False]
DTYPE = torch.bfloat16
DEVICE = "cuda"
EPS = 1e-5
LINE_VALS = ["native", "cuda"]
LINE_NAMES = ["SGLang Native", "SGLang Fused"]
STYLES = [("red", "-"), ("blue", "--")]
config = list(
itertools.product(B_RANGE, S_RANGE, D_RANGE, NORM_TYPE_RANGE, AFFINE_RANGE)
)
def preprocess_layer(layer, affine: bool, D: int, DTYPE: torch.dtype):
if affine:
weight = torch.randn(D, dtype=DTYPE, device=DEVICE)
bias = torch.randn(D, dtype=DTYPE, device=DEVICE)
with torch.no_grad():
layer.norm.weight.copy_(weight)
if hasattr(layer.norm, "bias"):
layer.norm.bias.copy_(bias)
layer.requires_grad_(False)
return layer.to(DEVICE)
# ============================================================================
# Benchmark 1: fused_norm_scale_shift
# ============================================================================
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["B", "S", "D", "norm_type", "affine"],
x_vals=config,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="fused_norm_scale_shift",
args={},
)
)
def bench_fused_norm_scale_shift(
B: int, S: int, D: int, norm_type, affine: bool, provider: str
) -> Tuple[float, float, float]:
x = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
scale = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
shift = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
if norm_type == "layer":
layer = LayerNormScaleShift(D, EPS, affine, dtype=DTYPE)
else:
layer = RMSNormScaleShift(D, EPS, affine, dtype=DTYPE)
layer = preprocess_layer(layer, affine, D, DTYPE)
if provider == "native":
fn = lambda: layer.forward_native(x, shift, scale)
else:
fn = lambda: layer.forward_cuda(x, shift, scale)
return run_benchmark_no_cudagraph(fn)
# ============================================================================
# Benchmark 2: fused_scale_residual_norm_scale_shift
# ============================================================================
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["B", "S", "D", "norm_type", "affine"],
x_vals=config,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="fused_scale_residual_norm_scale_shift",
args={},
)
)
def bench_fused_scale_residual_norm_scale_shift(
B: int, S: int, D: int, norm_type, affine: bool, provider: str
) -> Tuple[float, float, float]:
residual = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
x = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
scale = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
shift = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
gate = torch.randn(B, 1, D, dtype=DTYPE, device=DEVICE)
if norm_type == "layer":
layer = ScaleResidualLayerNormScaleShift(D, EPS, affine, dtype=DTYPE).to(DEVICE)
else:
layer = ScaleResidualRMSNormScaleShift(D, EPS, affine, dtype=DTYPE).to(DEVICE)
layer = preprocess_layer(layer, affine, D, DTYPE)
if provider == "native":
fn = lambda: layer.forward_native(residual, x, gate, shift, scale)
else:
fn = lambda: layer.forward_cuda(residual, x, gate, shift, scale)
return run_benchmark_no_cudagraph(fn)
if __name__ == "__main__":
print(f"\n{'='*80}")
print("Benchmark: fused_norm_scale_shift")
print(f"{'='*80}\n")
bench_fused_norm_scale_shift.run(print_data=True)
print(f"\n{'='*80}")
print("Benchmark: fused_scale_residual_norm_scale_shift")
print(f"{'='*80}\n")
bench_fused_scale_residual_norm_scale_shift.run(print_data=True)

753
third_party/sglang/python/sglang/jit_kernel/benchmark/diffusion/bench_norm_impls.py vendored Normal file
View File

@@ -0,0 +1,753 @@
import argparse
import csv
import functools
import importlib
import math
import os
import statistics
import subprocess
import sys
from pathlib import Path
from typing import Callable
import torch
import torch.nn.functional as F
from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE
from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn
from sglang.jit_kernel.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm
from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
from sglang.jit_kernel.utils import KERNEL_PATH
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(
est_time=120,
suite="stage-b-kernel-benchmark-1-gpu-large",
disabled="self-skips in CI, standalone tool",
)
os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
REPO_ROOT = KERNEL_PATH.parents[2]
THIRD_PARTY_ROOT = REPO_ROOT / "third_party"
FLAGGEMS_REPO = "https://github.com/flagos-ai/FlagGems.git"
QUACK_REPO = "https://github.com/Dao-AILab/quack.git"
TORCH_LN = "torch.nn.LayerNorm"
SGL_RMS = "sglang.RMSNorm.forward_cuda"
SGL_FUSED = "sgl_kernel.fused_add_rmsnorm"
SGL_LN = "sglang.LayerNormScaleShift"
SGL_RES_LN = "sglang.ScaleResidualLayerNormScaleShift"
SGL_LN_PAIR = f"{SGL_LN} / {SGL_RES_LN}"
MOVA_LN_MIX = f"{TORCH_LN} / {SGL_LN_PAIR}"
ACTUAL_DIFFUSION_GROUPS: list[
tuple[str, str, list[tuple[str, str, tuple[int, ...], str]]]
] = [
(
"qwen",
"1 GPU",
[
("qwen_ln_4096x3072", "layernorm", (1, 4096, 3072), SGL_LN_PAIR),
("qwen_ln_26x3072", "layernorm", (1, 26, 3072), SGL_LN_PAIR),
("qwen_ln_6x3072", "layernorm", (1, 6, 3072), SGL_LN_PAIR),
("qwen_rms_26x3584", "rmsnorm", (1, 26, 3584), SGL_RMS),
("qwen_rms_6x3584", "rmsnorm", (1, 6, 3584), SGL_RMS),
],
),
(
"qwen-edit",
"1 GPU",
[
("qwen_edit_ln_200x3072", "layernorm", (1, 200, 3072), SGL_LN_PAIR),
("qwen_edit_ln_203x3072", "layernorm", (1, 203, 3072), SGL_LN_PAIR),
("qwen_edit_ln_8308x3072", "layernorm", (1, 8308, 3072), TORCH_LN),
("qwen_edit_rms_200x3584", "rmsnorm", (1, 200, 3584), SGL_RMS),
("qwen_edit_rms_203x3584", "rmsnorm", (1, 203, 3584), SGL_RMS),
],
),
(
"flux",
"1 GPU",
[
("flux_ln_77x768", "layernorm", (1, 77, 768), TORCH_LN),
("flux_ln_512x3072", "layernorm", (1, 512, 3072), TORCH_LN),
("flux_ln_4096x3072", "layernorm", (1, 4096, 3072), TORCH_LN),
("flux_ln_4608x3072", "layernorm", (1, 4608, 3072), TORCH_LN),
("flux_rms_512x4096", "rmsnorm", (1, 512, 4096), SGL_RMS),
],
),
(
"flux2",
"1 GPU",
[
("flux2_ln_512x6144", "layernorm", (1, 512, 6144), TORCH_LN),
("flux2_ln_4096x6144", "layernorm", (1, 4096, 6144), TORCH_LN),
("flux2_ln_4608x6144", "layernorm", (1, 4608, 6144), TORCH_LN),
("flux2_rms_4608x48x128", "rmsnorm", (1, 4608, 48, 128), SGL_RMS),
],
),
(
"zimage",
"1 GPU",
[
("zimage_ln_4128x3840", "layernorm", (1, 4128, 3840), TORCH_LN),
("zimage_rms_32x3840", "rmsnorm", (1, 32, 3840), SGL_RMS),
("zimage_rms_4096x3840", "rmsnorm", (1, 4096, 3840), SGL_RMS),
("zimage_rms_4128x3840", "rmsnorm", (1, 4128, 3840), SGL_RMS),
("zimage_rms_32x2560", "rmsnorm", (32, 2560), SGL_RMS),
],
),
(
"wan-ti2v",
"1 GPU",
[
("wan_ti2v_ln_17850x3072", "layernorm", (1, 17850, 3072), SGL_LN_PAIR),
("wan_ti2v_rms_17850x3072", "rmsnorm", (1, 17850, 3072), SGL_RMS),
("wan_ti2v_rms_512x3072", "rmsnorm", (1, 512, 3072), SGL_RMS),
("wan_ti2v_rms_512x4096", "rmsnorm", (1, 512, 4096), SGL_RMS),
],
),
(
"hunyuanvideo",
"1 GPU",
[
("hunyuan_ln_46x768", "layernorm", (1, 46, 768), TORCH_LN),
("hunyuan_ln_45x3072", "layernorm", (1, 45, 3072), SGL_LN_PAIR),
("hunyuan_ln_27030x3072", "layernorm", (1, 27030, 3072), SGL_LN_PAIR),
("hunyuan_ln_27075x3072", "layernorm", (1, 27075, 3072), SGL_LN),
("hunyuan_rms_140x4096", "rmsnorm", (1, 140, 4096), SGL_RMS),
("hunyuan_rms_45x24x128", "rmsnorm", (1, 45, 24, 128), SGL_RMS),
("hunyuan_rms_27030x24x128", "rmsnorm", (1, 27030, 24, 128), SGL_RMS),
("hunyuan_rms_27075x24x128", "rmsnorm", (1, 27075, 24, 128), SGL_RMS),
("hunyuan_fused_add_140x4096", "fused_add_rmsnorm", (140, 4096), SGL_FUSED),
],
),
(
"mova-720p",
"4 GPU, ulysses=4, ring=1",
[
("mova_ln_101x1536", "layernorm", (1, 101, 1536), MOVA_LN_MIX),
("mova_ln_403x1536", "layernorm", (1, 403, 1536), TORCH_LN),
("mova_ln_44100x5120", "layernorm", (1, 44100, 5120), MOVA_LN_MIX),
("mova_ln_176400x5120", "layernorm", (1, 176400, 5120), SGL_LN),
("mova_rms_101x1536", "rmsnorm", (1, 101, 1536), SGL_RMS),
("mova_rms_101x5120", "rmsnorm", (1, 101, 5120), SGL_RMS),
("mova_rms_44100x1536", "rmsnorm", (1, 44100, 1536), SGL_RMS),
("mova_rms_44100x5120", "rmsnorm", (1, 44100, 5120), SGL_RMS),
("mova_rms_512x1536", "rmsnorm", (1, 512, 1536), SGL_RMS),
("mova_rms_512x4096", "rmsnorm", (1, 512, 4096), SGL_RMS),
("mova_rms_512x5120", "rmsnorm", (1, 512, 5120), SGL_RMS),
],
),
]
ACTUAL_DIFFUSION_SHAPES: list[dict[str, object]] = [
{
"shape_id": shape_id,
"model": model,
"gpu_config": gpu_config,
"op": op,
"input_shape": list(input_shape),
"source_impl": source_impl,
}
for model, gpu_config, cases in ACTUAL_DIFFUSION_GROUPS
for shape_id, op, input_shape, source_impl in cases
]
def effective_rows_from_shape(input_shape: list[int]) -> int:
rows = 1
for dim in input_shape[:-1]:
rows *= dim
return rows
def ensure_repo(repo_name: str, repo_url: str) -> Path:
repo_path = THIRD_PARTY_ROOT / repo_name
if repo_path.exists():
return repo_path
repo_path.parent.mkdir(parents=True, exist_ok=True)
subprocess.run(
["git", "clone", "--depth", "1", repo_url, str(repo_path)],
check=True,
cwd=REPO_ROOT,
)
return repo_path
def ensure_python_dep(module_name: str, package_name: str | None = None) -> None:
package_name = package_name or module_name
try:
importlib.import_module(module_name)
except ModuleNotFoundError:
subprocess.run(
[sys.executable, "-m", "pip", "install", package_name],
check=True,
)
def dtype_from_name(name: str) -> torch.dtype:
mapping = {
"bf16": torch.bfloat16,
"bfloat16": torch.bfloat16,
"fp16": torch.float16,
"float16": torch.float16,
"fp32": torch.float32,
"float32": torch.float32,
}
return mapping[name]
def dtype_name(dtype: torch.dtype) -> str:
mapping = {
torch.bfloat16: "bf16",
torch.float16: "fp16",
torch.float32: "fp32",
}
return mapping[dtype]
def normalize_hidden_sizes(text: str) -> list[int]:
return [int(x) for x in text.split(",") if x]
def normalize_dtypes(text: str) -> list[torch.dtype]:
return [dtype_from_name(x.strip()) for x in text.split(",") if x.strip()]
def prewarm(fn: Callable[[], object], iters: int = 3) -> None:
for _ in range(iters):
fn()
torch.cuda.synchronize()
def benchmark_provider(
fn: Callable[[], object],
setup_fn: Callable[[], None] | None = None,
warmup: int = 10,
rep: int = 30,
) -> tuple[float, float, float]:
for _ in range(warmup):
if setup_fn is not None:
setup_fn()
fn()
torch.cuda.synchronize()
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
times_us: list[float] = []
for _ in range(rep):
if setup_fn is not None:
setup_fn()
start_event.record()
fn()
end_event.record()
end_event.synchronize()
times_us.append(start_event.elapsed_time(end_event) * 1000.0)
return statistics.median(times_us), max(times_us), min(times_us)
def geometric_mean(values: list[float]) -> float:
if not values:
return float("nan")
return math.exp(sum(math.log(v) for v in values) / len(values))
@functools.cache
def load_flaggems():
ensure_python_dep("sqlalchemy")
ensure_repo("FlagGems", FLAGGEMS_REPO)
src_root = THIRD_PARTY_ROOT / "FlagGems" / "src"
if str(src_root) not in sys.path:
sys.path.insert(0, str(src_root))
from flag_gems.fused.fused_add_rms_norm import fused_add_rms_norm
from flag_gems.ops.layernorm import layer_norm
from flag_gems.ops.rms_norm import rms_norm
return rms_norm, layer_norm, fused_add_rms_norm
@functools.cache
def load_quack():
repo_path = ensure_repo("quack", QUACK_REPO)
try:
quack_rmsnorm = importlib.import_module("quack.rmsnorm")
except ModuleNotFoundError:
subprocess.run(
[sys.executable, "-m", "pip", "install", "-e", str(repo_path)],
check=True,
)
quack_rmsnorm = importlib.import_module("quack.rmsnorm")
return quack_rmsnorm.rmsnorm_fwd, quack_rmsnorm.layernorm_fwd
def build_rmsnorm_providers(dtype: torch.dtype, batch_size: int, hidden_size: int):
import flashinfer.norm as flashinfer_norm
import sgl_kernel
x = torch.randn((batch_size, hidden_size), device=DEFAULT_DEVICE, dtype=dtype)
weight = torch.randn(hidden_size, device=DEFAULT_DEVICE, dtype=dtype)
jit_out = torch.empty_like(x)
sgl_out = torch.empty_like(x)
flashinfer_out = torch.empty_like(x)
flaggems_rms_norm, _, _ = load_flaggems()
quack_rmsnorm_fwd, _ = load_quack()
providers = {
"pytorch": lambda: F.rms_norm(x, (hidden_size,), weight, 1e-6),
"sgl_kernel": lambda: sgl_kernel.rmsnorm(x, weight, eps=1e-6, out=sgl_out),
"flashinfer": lambda: flashinfer_norm.rmsnorm(
x, weight, eps=1e-6, out=flashinfer_out
),
"jit_rmsnorm": lambda: jit_rmsnorm(x, weight, jit_out, 1e-6),
"quack": lambda: quack_rmsnorm_fwd(x, weight, eps=1e-6),
"triton_rms_norm_fn": lambda: rms_norm_fn(
x, weight, bias=None, residual=None, eps=1e-6
),
"flaggems": lambda: flaggems_rms_norm(x, (hidden_size,), weight, 1e-6),
}
if hidden_size <= 128:
providers["triton_one_pass"] = lambda: triton_one_pass_rms_norm(x, weight, 1e-6)
return providers
def build_fused_add_rmsnorm_providers(
dtype: torch.dtype, batch_size: int, hidden_size: int
):
import flashinfer.norm as flashinfer_norm
import sgl_kernel
base_x = torch.randn((batch_size, hidden_size), device=DEFAULT_DEVICE, dtype=dtype)
base_residual = torch.randn_like(base_x)
weight = torch.randn(hidden_size, device=DEFAULT_DEVICE, dtype=dtype)
x = base_x.clone()
residual = base_residual.clone()
def reset():
x.copy_(base_x)
residual.copy_(base_residual)
_, _, flaggems_fused_add_rms_norm = load_flaggems()
quack_rmsnorm_fwd, _ = load_quack()
def pytorch_impl():
out = x + residual
return F.rms_norm(out, (hidden_size,), weight, 1e-6)
providers = {
"pytorch": (pytorch_impl, reset),
"sgl_kernel": (
lambda: sgl_kernel.fused_add_rmsnorm(x, residual, weight, eps=1e-6),
reset,
),
"flashinfer": (
lambda: flashinfer_norm.fused_add_rmsnorm(x, residual, weight, eps=1e-6),
reset,
),
"jit_fused_add_rmsnorm": (
lambda: jit_fused_add_rmsnorm(x, residual, weight, 1e-6),
reset,
),
"quack": (
lambda: quack_rmsnorm_fwd(x, weight, residual=residual, eps=1e-6),
reset,
),
"flaggems": (
lambda: flaggems_fused_add_rms_norm(
x, residual, (hidden_size,), weight, 1e-6
),
reset,
),
}
return providers
def build_layernorm_providers(dtype: torch.dtype, batch_size: int, hidden_size: int):
import flashinfer.norm as flashinfer_norm
x = torch.randn((batch_size, hidden_size), device=DEFAULT_DEVICE, dtype=dtype)
weight = torch.randn(hidden_size, device=DEFAULT_DEVICE, dtype=dtype)
bias = torch.randn(hidden_size, device=DEFAULT_DEVICE, dtype=dtype)
flashinfer_weight = torch.randn(
hidden_size, device=DEFAULT_DEVICE, dtype=torch.float32
)
flashinfer_bias = torch.randn(
hidden_size, device=DEFAULT_DEVICE, dtype=torch.float32
)
triton_out = torch.empty_like(x)
_, flaggems_layer_norm, _ = load_flaggems()
_, quack_layernorm_fwd = load_quack()
providers = {
"pytorch": lambda: F.layer_norm(x, (hidden_size,), weight, bias, 1e-6),
"triton_norm_infer": lambda: norm_infer(
x, weight, bias, eps=1e-6, is_rms_norm=False, out=triton_out
),
"flashinfer": lambda: flashinfer_norm.layernorm(
x, flashinfer_weight, flashinfer_bias, 1e-6
),
"quack": lambda: quack_layernorm_fwd(
x, flashinfer_weight, flashinfer_bias, 1e-6
),
"flaggems": lambda: flaggems_layer_norm(x, (hidden_size,), weight, bias)[0],
}
return providers
def maybe_benchmark(
op_name: str,
provider_name: str,
fn: Callable[[], object],
rows: list[dict[str, object]],
dtype: torch.dtype,
batch_size: int,
hidden_size: int,
reset: Callable[[], None] | None = None,
metadata: dict[str, object] | None = None,
) -> None:
metadata = metadata or {}
try:
median_us, max_us, min_us = benchmark_provider(fn, reset)
rows.append(
{
"op": op_name,
"provider": provider_name,
"dtype": dtype_name(dtype),
"batch_size": batch_size,
"hidden_size": hidden_size,
"median_us": median_us,
"min_us": min_us,
"max_us": max_us,
"status": "ok",
"error": "",
**metadata,
}
)
except Exception as exc: # pragma: no cover - benchmark failures are data
rows.append(
{
"op": op_name,
"provider": provider_name,
"dtype": dtype_name(dtype),
"batch_size": batch_size,
"hidden_size": hidden_size,
"median_us": "",
"min_us": "",
"max_us": "",
"status": "unsupported",
"error": str(exc),
**metadata,
}
)
def write_csv(rows: list[dict[str, object]], output_path: Path) -> None:
output_path.parent.mkdir(parents=True, exist_ok=True)
with output_path.open("w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(
f,
fieldnames=[
"op",
"provider",
"dtype",
"batch_size",
"hidden_size",
"median_us",
"min_us",
"max_us",
"shape_id",
"source_model",
"source_gpu_config",
"source_input_shape",
"source_impl",
"status",
"error",
],
)
writer.writeheader()
writer.writerows(rows)
def write_markdown(rows: list[dict[str, object]], output_path: Path) -> None:
lines: list[str] = []
lines.append("# Norm Benchmark Summary")
lines.append("")
actual_shape_rows = [row for row in rows if row.get("shape_id")]
if actual_shape_rows:
seen: set[tuple[str, str, str, str, str, str]] = set()
lines.append("## Diffusion Shape Cases")
lines.append("")
lines.append(
"| Shape ID | Op | Model | GPU Config | Input Shape | Source Impl |"
)
lines.append("|---|---|---|---|---|---|")
for row in actual_shape_rows:
key = (
str(row.get("shape_id", "")),
str(row.get("op", "")),
str(row.get("source_model", "")),
str(row.get("source_gpu_config", "")),
str(row.get("source_input_shape", "")),
str(row.get("source_impl", "")),
)
if key in seen:
continue
seen.add(key)
lines.append(
f"| {key[0]} | {key[1]} | {key[2]} | {key[3]} | `{key[4]}` | {key[5]} |"
)
lines.append("")
for op_name in ("rmsnorm", "fused_add_rmsnorm", "layernorm"):
for dtype in sorted({row["dtype"] for row in rows}):
scoped = [
row
for row in rows
if row["op"] == op_name
and row["dtype"] == dtype
and row["status"] == "ok"
]
if not scoped:
continue
provider_to_values: dict[str, list[float]] = {}
provider_to_speedups: dict[str, list[float]] = {}
by_shape: dict[tuple[str, int, int], dict[str, float]] = {}
for row in scoped:
provider = str(row["provider"])
value = float(row["median_us"])
provider_to_values.setdefault(provider, []).append(value)
shape = (
str(row.get("shape_id", "")),
int(row["batch_size"]),
int(row["hidden_size"]),
)
by_shape.setdefault(shape, {})[provider] = value
for shape, perf in by_shape.items():
if "pytorch" not in perf:
continue
baseline = perf["pytorch"]
for provider, value in perf.items():
provider_to_speedups.setdefault(provider, []).append(
baseline / value
)
lines.append(f"## {op_name} ({dtype})")
lines.append("")
lines.append(
"| Provider | Geomean Speedup vs PyTorch | Median Latency (us) | Win Count |"
)
lines.append("|---|---:|---:|---:|")
wins: dict[str, int] = {}
for perf in by_shape.values():
best_provider = min(perf, key=perf.get)
wins[best_provider] = wins.get(best_provider, 0) + 1
for provider in sorted(provider_to_values):
geomean_speedup = geometric_mean(provider_to_speedups.get(provider, []))
median_latency = statistics.median(provider_to_values[provider])
win_count = wins.get(provider, 0)
lines.append(
f"| {provider} | {geomean_speedup:.3f}x | {median_latency:.2f} | {win_count} |"
)
lines.append("")
output_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
def run_suite(
hidden_sizes: list[int],
batch_sizes: list[int],
dtypes: list[torch.dtype],
ops: list[str],
) -> list[dict[str, object]]:
rows: list[dict[str, object]] = []
for dtype in dtypes:
for batch_size in batch_sizes:
for hidden_size in hidden_sizes:
if "rmsnorm" in ops:
rms_providers = build_rmsnorm_providers(
dtype, batch_size, hidden_size
)
for provider_name, fn in rms_providers.items():
maybe_benchmark(
"rmsnorm",
provider_name,
fn,
rows,
dtype,
batch_size,
hidden_size,
)
if "fused_add_rmsnorm" in ops:
fused_providers = build_fused_add_rmsnorm_providers(
dtype, batch_size, hidden_size
)
for provider_name, provider in fused_providers.items():
fn, reset = provider
maybe_benchmark(
"fused_add_rmsnorm",
provider_name,
fn,
rows,
dtype,
batch_size,
hidden_size,
reset,
)
if "layernorm" in ops:
layernorm_providers = build_layernorm_providers(
dtype, batch_size, hidden_size
)
for provider_name, fn in layernorm_providers.items():
maybe_benchmark(
"layernorm",
provider_name,
fn,
rows,
dtype,
batch_size,
hidden_size,
)
return rows
def run_shape_suite(
shape_cases: list[dict[str, object]],
dtypes: list[torch.dtype],
) -> list[dict[str, object]]:
rows: list[dict[str, object]] = []
for case in shape_cases:
op_name = str(case["op"])
input_shape = [int(x) for x in case["input_shape"]]
batch_size = effective_rows_from_shape(input_shape)
hidden_size = input_shape[-1]
metadata = {
"shape_id": str(case["shape_id"]),
"source_model": str(case["model"]),
"source_gpu_config": str(case["gpu_config"]),
"source_input_shape": str(input_shape),
"source_impl": str(case["source_impl"]),
}
for dtype in dtypes:
if op_name == "rmsnorm":
providers = build_rmsnorm_providers(dtype, batch_size, hidden_size)
for provider_name, fn in providers.items():
maybe_benchmark(
op_name,
provider_name,
fn,
rows,
dtype,
batch_size,
hidden_size,
metadata=metadata,
)
elif op_name == "fused_add_rmsnorm":
providers = build_fused_add_rmsnorm_providers(
dtype, batch_size, hidden_size
)
for provider_name, provider in providers.items():
fn, reset = provider
maybe_benchmark(
op_name,
provider_name,
fn,
rows,
dtype,
batch_size,
hidden_size,
reset,
metadata=metadata,
)
elif op_name == "layernorm":
providers = build_layernorm_providers(dtype, batch_size, hidden_size)
for provider_name, fn in providers.items():
maybe_benchmark(
op_name,
provider_name,
fn,
rows,
dtype,
batch_size,
hidden_size,
metadata=metadata,
)
else:
raise ValueError(f"Unsupported op in shape preset: {op_name}")
return rows
def main() -> None:
parser = argparse.ArgumentParser(
description="Benchmark RMSNorm/LayerNorm implementations across providers."
)
parser.add_argument(
"--hidden-sizes",
default="64,128,256,512,1024,2048,4096,8192,16384",
help="Comma-separated hidden sizes.",
)
parser.add_argument(
"--batch-sizes",
default="1,16,128,1024",
help="Comma-separated batch sizes.",
)
parser.add_argument(
"--dtypes",
default="bf16,fp16",
help="Comma-separated dtypes: bf16, fp16, fp32.",
)
parser.add_argument(
"--output-dir",
default=str(REPO_ROOT / "outputs" / "norm_benchmarks"),
help="Directory for CSV/Markdown outputs.",
)
parser.add_argument(
"--ops",
default="rmsnorm,fused_add_rmsnorm,layernorm",
help="Comma-separated ops to benchmark.",
)
parser.add_argument(
"--shape-preset",
choices=["grid", "diffusion-actual"],
default="grid",
help="Use the default grid sweep or the captured diffusion workload shapes.",
)
args = parser.parse_args()
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required for norm benchmarks.")
hidden_sizes = normalize_hidden_sizes(args.hidden_sizes)
batch_sizes = normalize_hidden_sizes(args.batch_sizes)
dtypes = normalize_dtypes(args.dtypes)
ops = [op.strip() for op in args.ops.split(",") if op.strip()]
if args.shape_preset == "diffusion-actual":
shape_cases = [case for case in ACTUAL_DIFFUSION_SHAPES if case["op"] in ops]
rows = run_shape_suite(shape_cases, dtypes)
else:
rows = run_suite(hidden_sizes, batch_sizes, dtypes, ops)
output_dir = Path(args.output_dir)
csv_path = output_dir / "norm_impls.csv"
md_path = output_dir / "norm_impls_summary.md"
write_csv(rows, csv_path)
write_markdown(rows, md_path)
print(f"Wrote {csv_path}")
print(f"Wrote {md_path}")
if __name__ == "__main__":
if is_in_ci():
print("Skipping bench_norm_impls.py in CI")
sys.exit(0)
main()

190
third_party/sglang/python/sglang/jit_kernel/benchmark/diffusion/bench_qknorm_rope.py vendored Normal file
View File

@@ -0,0 +1,190 @@
from dataclasses import dataclass
from typing import Tuple
import torch
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import (
DEFAULT_DEVICE,
DEFAULT_DTYPE,
get_benchmark_range,
run_benchmark_no_cudagraph,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=13, suite="stage-b-kernel-benchmark-1-gpu-large")
MAX_SEQ_LEN = 131072
ROPE_BASE = 10000.0
@dataclass(frozen=True)
class CaseSpec:
name: str
batch_size: int
num_tokens: int
num_heads: int
head_dim: int
rope_dim: int
is_neox: bool
BENCH_CASES = (
CaseSpec("flux_1024", 1, 4096, 24, 128, 128, False),
CaseSpec("qwen_image_1024", 1, 4096, 32, 128, 128, False),
CaseSpec("qwen_image_partial", 1, 4096, 32, 128, 64, False),
# Z-Image-Turbo default 1024x1024 config: dim=3840, num_heads=30 -> head_dim=128.
CaseSpec("zimage_1024", 1, 4096, 30, 128, 128, False),
CaseSpec("batch2_medium", 2, 2048, 24, 128, 128, False),
)
CASE_BY_NAME = {case.name: case for case in BENCH_CASES}
CASE_NAMES = get_benchmark_range(
full_range=[case.name for case in BENCH_CASES],
ci_range=[case.name for case in BENCH_CASES],
)
LINE_VALS = ["split", "fused"]
LINE_NAMES = ["JIT QKNorm + FlashInfer RoPE", "SGL JIT Fused QKNorm+RoPE"]
STYLES = [("red", "-"), ("blue", "--")]
def create_cos_sin_cache(
rotary_dim: int,
max_position: int = MAX_SEQ_LEN,
base: float = ROPE_BASE,
) -> torch.Tensor:
inv_freq = 1.0 / (
base
** (
torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=DEFAULT_DEVICE)
/ rotary_dim
)
)
t = torch.arange(max_position, dtype=torch.float32, device=DEFAULT_DEVICE)
freqs = torch.einsum("i,j->ij", t, inv_freq)
return torch.cat((freqs.cos(), freqs.sin()), dim=-1)
def make_inputs(case: CaseSpec) -> dict[str, torch.Tensor | bool]:
seed = (
case.batch_size * 1_000_003
+ case.num_tokens * 8191
+ case.num_heads * 127
+ case.head_dim * 17
+ case.rope_dim
)
generator = torch.Generator(device=DEFAULT_DEVICE)
generator.manual_seed(seed)
return {
"q": torch.randn(
case.batch_size * case.num_tokens,
case.num_heads,
case.head_dim,
device=DEFAULT_DEVICE,
dtype=DEFAULT_DTYPE,
generator=generator,
),
"k": torch.randn(
case.batch_size * case.num_tokens,
case.num_heads,
case.head_dim,
device=DEFAULT_DEVICE,
dtype=DEFAULT_DTYPE,
generator=generator,
),
"q_weight": torch.randn(
case.head_dim,
device=DEFAULT_DEVICE,
dtype=DEFAULT_DTYPE,
generator=generator,
),
"k_weight": torch.randn(
case.head_dim,
device=DEFAULT_DEVICE,
dtype=DEFAULT_DTYPE,
generator=generator,
),
"positions": torch.randint(
0,
MAX_SEQ_LEN,
(case.batch_size * case.num_tokens,),
device=DEFAULT_DEVICE,
dtype=torch.int64,
generator=generator,
),
"cos_sin_cache": create_cos_sin_cache(case.rope_dim),
"is_neox": case.is_neox,
}
def clone_inputs(
inputs: dict[str, torch.Tensor | bool],
) -> dict[str, torch.Tensor | bool]:
out: dict[str, torch.Tensor | bool] = {}
for key, value in inputs.items():
out[key] = value.clone() if isinstance(value, torch.Tensor) else value
return out
def split_qknorm_rope(inputs: dict[str, torch.Tensor | bool]) -> None:
from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
from sglang.jit_kernel.norm import fused_inplace_qknorm
q = inputs["q"]
k = inputs["k"]
q_weight = inputs["q_weight"]
k_weight = inputs["k_weight"]
positions = inputs["positions"]
cos_sin_cache = inputs["cos_sin_cache"]
is_neox = bool(inputs["is_neox"])
fused_inplace_qknorm(q, k, q_weight, k_weight)
apply_rope_with_cos_sin_cache_inplace(
positions=positions,
query=q.view(q.shape[0], -1),
key=k.view(k.shape[0], -1),
head_size=q.shape[-1],
cos_sin_cache=cos_sin_cache,
is_neox=is_neox,
)
def fused_qknorm_rope(inputs: dict[str, torch.Tensor | bool]) -> None:
from sglang.jit_kernel.diffusion.qknorm_rope import fused_inplace_qknorm_rope
fused_inplace_qknorm_rope(
inputs["q"],
inputs["k"],
inputs["q_weight"],
inputs["k_weight"],
inputs["cos_sin_cache"],
inputs["positions"],
is_neox=bool(inputs["is_neox"]),
rope_dim=inputs["cos_sin_cache"].shape[-1],
)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["case_name"],
x_vals=CASE_NAMES,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="diffusion-qknorm-rope-performance",
args={},
)
)
def benchmark(case_name: str, provider: str) -> Tuple[float, float, float]:
case = CASE_BY_NAME[case_name]
inputs = make_inputs(case)
fn = split_qknorm_rope if provider == "split" else fused_qknorm_rope
return run_benchmark_no_cudagraph(lambda: fn(inputs))
if __name__ == "__main__":
print("Running diffusion qknorm + rope performance benchmark...")
benchmark.run(print_data=True)

184
third_party/sglang/python/sglang/jit_kernel/benchmark/diffusion/bench_qwen_image_modulation.py vendored Normal file
View File

@@ -0,0 +1,184 @@
from typing import Tuple
import torch
import triton.testing
from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
from sglang.jit_kernel.diffusion.triton.norm import norm_infer
from sglang.jit_kernel.diffusion.triton.scale_shift import (
fuse_layernorm_scale_shift_gate_select01_kernel,
fuse_residual_layernorm_scale_shift_gate_select01_kernel,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(est_time=13, suite="stage-b-kernel-benchmark-1-gpu-large")
if is_in_ci():
B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072]
else:
B_RANGE, S_RANGE, D_RANGE = [1, 2], [128, 512, 2048], [1024, 1536, 3072]
DTYPE = torch.bfloat16
DEVICE = "cuda"
EPS = 1e-6
LINE_VALS = ["split", "fused"]
LINE_NAMES = ["Triton Norm + Torch Select", "Fused Triton"]
STYLES = [("red", "-"), ("blue", "--")]
CONFIG = [(b, s, d) for b in B_RANGE for s in S_RANGE for d in D_RANGE]
def _make_common_inputs(batch_size: int, seq_len: int, hidden_size: int):
x = torch.randn(batch_size, seq_len, hidden_size, dtype=DTYPE, device=DEVICE)
weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE)
bias = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE)
index = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.int32, device=DEVICE)
scale0 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
shift0 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
gate0 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
scale1 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
shift1 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
gate1 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
return x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1
def _apply_select01_modulation(
x: torch.Tensor,
scale0: torch.Tensor,
shift0: torch.Tensor,
gate0: torch.Tensor,
scale1: torch.Tensor,
shift1: torch.Tensor,
gate1: torch.Tensor,
index: torch.Tensor,
):
idx = index.bool().unsqueeze(-1)
scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
return x * (1 + scale) + shift, gate
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["B", "S", "D"],
x_vals=CONFIG,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="qwen_image_layernorm_scale_shift_gate_select01",
args={},
)
)
def bench_layernorm_scale_shift_gate_select01(
B: int, S: int, D: int, provider: str
) -> Tuple[float, float, float]:
x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1 = (
_make_common_inputs(B, S, D)
)
if provider == "split":
def fn():
normalized = norm_infer(
x.view(-1, x.shape[-1]),
weight,
bias,
eps=EPS,
is_rms_norm=False,
).view_as(x)
return _apply_select01_modulation(
normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
)
else:
def fn():
return fuse_layernorm_scale_shift_gate_select01_kernel(
x,
weight=weight,
bias=bias,
scale0=scale0,
shift0=shift0,
gate0=gate0,
scale1=scale1,
shift1=shift1,
gate1=gate1,
index=index,
eps=EPS,
)
return run_benchmark_no_cudagraph(fn)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["B", "S", "D"],
x_vals=CONFIG,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="qwen_image_residual_layernorm_scale_shift_gate_select01",
args={},
)
)
def bench_residual_layernorm_scale_shift_gate_select01(
B: int, S: int, D: int, provider: str
) -> Tuple[float, float, float]:
x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1 = (
_make_common_inputs(B, S, D)
)
residual = torch.randn_like(x)
residual_gate = torch.randn_like(x)
if provider == "split":
def fn():
residual_out = residual + residual_gate * x
normalized = norm_infer(
residual_out.view(-1, residual_out.shape[-1]),
weight,
bias,
eps=EPS,
is_rms_norm=False,
).view_as(residual_out)
return _apply_select01_modulation(
normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
)
else:
def fn():
return fuse_residual_layernorm_scale_shift_gate_select01_kernel(
x,
residual=residual,
residual_gate=residual_gate,
weight=weight,
bias=bias,
scale0=scale0,
shift0=shift0,
gate0=gate0,
scale1=scale1,
shift1=shift1,
gate1=gate1,
index=index,
eps=EPS,
)
return run_benchmark_no_cudagraph(fn)
if __name__ == "__main__":
print(f"\n{'=' * 80}")
print("Benchmark: qwen_image layernorm + scale_shift_gate_select01")
print(f"{'=' * 80}\n")
bench_layernorm_scale_shift_gate_select01.run(print_data=True)
print(f"\n{'=' * 80}")
print("Benchmark: qwen_image residual + layernorm + scale_shift_gate_select01")
print(f"{'=' * 80}\n")
bench_residual_layernorm_scale_shift_gate_select01.run(print_data=True)

2676
third_party/sglang/python/sglang/jit_kernel/benchmark/diffusion/diffusion_nvfp4_shapes.json vendored Normal file
View File

File diff suppressed because it is too large Load Diff

48
third_party/sglang/python/sglang/jit_kernel/benchmark/utils.py vendored Normal file
View File

@@ -0,0 +1,48 @@
"""Common utilities for jit_kernel benchmark files."""
from typing import Callable, List, Sequence, Tuple
import torch
import triton.testing
from sglang.utils import is_in_ci
# Common constants
DEFAULT_DTYPE = torch.bfloat16
DEFAULT_DEVICE = "cuda"
DEFAULT_QUANTILES = [0.5, 0.2, 0.8]
def get_benchmark_range(full_range: List, ci_range: List) -> List:
"""Return appropriate benchmark range based on CI environment."""
return ci_range if is_in_ci() else full_range
def run_benchmark(
fn: Callable,
quantiles: Sequence[float] = (),
scale: float = 1.0,
) -> Tuple[float, float, float]:
"""Execute benchmark using CUDA graph and return times in microseconds.
Args:
fn: Function to benchmark
quantiles: Quantiles for timing measurements [median, min, max]
scale: Scale the result down (usually num_layers).
Returns:
Tuple of (median_us, max_us, min_us)
"""
quantiles = list(quantiles or DEFAULT_QUANTILES)
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms / scale, 1000 * max_ms / scale, 1000 * min_ms / scale
def run_benchmark_no_cudagraph(
fn: Callable,
quantiles: Sequence[float] = (),
scale: float = 1.0,
) -> Tuple[float, float, float]:
quantiles = list(quantiles or DEFAULT_QUANTILES)
ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
return 1000 * ms / scale, 1000 * max_ms / scale, 1000 * min_ms / scale

52
third_party/sglang/python/sglang/jit_kernel/cast.py vendored Normal file
View File

@@ -0,0 +1,52 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_cast_module(dtype: torch.dtype) -> Module:
args = make_cpp_args(dtype)
return load_jit(
"cast",
*args,
cuda_files=["elementwise/cast.cuh"],
cuda_wrappers=[("downcast_fp8", f"downcast_fp8<{args}>")],
)
def downcast_fp8(
k: torch.Tensor,
v: torch.Tensor,
k_out: torch.Tensor,
v_out: torch.Tensor,
k_scale: torch.Tensor,
v_scale: torch.Tensor,
loc: torch.Tensor,
mult: int = 1,
offset: int = 0,
) -> None:
"""Fused downcast of KV cache tensors from bf16/fp16 to fp8 (E4M3).
Scales each value by the inverse of its per-tensor scale, clamps to the
fp8 representable range [-448, 448], then converts to fp8 storage.
Args:
k: [input_sl, head, dim] bf16/fp16 CUDA tensor
v: [input_sl, head, dim] bf16/fp16 CUDA tensor
k_out: [out_sl, head, dim] uint8 CUDA tensor (fp8 storage)
v_out: [out_sl, head, dim] uint8 CUDA tensor (fp8 storage)
k_scale: [1] float32 CUDA tensor, scale for k
v_scale: [1] float32 CUDA tensor, scale for v
loc: [input_sl] int64 CUDA tensor, destination sequence indices
mult: stride multiplier for output index (default 1)
offset: offset added to output index (default 0)
"""
module = _jit_cast_module(k.dtype)
module.downcast_fp8(k, v, k_out, v_out, k_scale, v_scale, loc, mult, offset)

35
third_party/sglang/python/sglang/jit_kernel/clamp_position.py vendored Normal file
View File

@@ -0,0 +1,35 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_clamp_position_module(dtype: torch.dtype) -> Module:
"""Compile and cache the JIT clamp_position module for a given dtype."""
args = make_cpp_args(dtype)
return load_jit(
"clamp_position",
*args,
cuda_files=["elementwise/clamp_position.cuh"],
cuda_wrappers=[
("clamp_position", f"ClampPosition<{args}>::run"),
],
)
def clamp_position_cuda(seq_lens: torch.Tensor) -> torch.Tensor:
"""Compute positions = clamp(seq_lens - 1, min=0) on CUDA.
Supported dtypes: torch.int32, torch.int64.
"""
dst = torch.empty_like(seq_lens)
module = _jit_clamp_position_module(seq_lens.dtype)
module.clamp_position(dst, seq_lens)
return dst

65
third_party/sglang/python/sglang/jit_kernel/concat_mla.py vendored Normal file
View File

@@ -0,0 +1,65 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_concat_mla_k_module() -> Module:
return load_jit(
"concat_mla_k",
cuda_files=["elementwise/concat_mla.cuh"],
cuda_wrappers=[("concat_mla_k", "ConcatMlaKKernel::run")],
)
@cache_once
def _jit_concat_mla_absorb_q_module() -> Module:
return load_jit(
"concat_mla_absorb_q",
cuda_files=["elementwise/concat_mla.cuh"],
cuda_wrappers=[("concat_mla_absorb_q", "ConcatMlaAbsorbQKernel::run")],
)
def concat_mla_k(k: torch.Tensor, k_nope: torch.Tensor, k_rope: torch.Tensor) -> None:
"""
Concatenate k_nope and k_rope into k for MLA (Multi-head Latent Attention).
This kernel efficiently broadcasts k_rope across all heads while copying
k_nope values directly.
Args:
k: Output tensor of shape [num_tokens, num_heads=128, k_head_dim=192], dtype=bfloat16
k_nope: Input tensor of shape [num_tokens, num_heads=128, nope_head_dim=128], dtype=bfloat16
k_rope: Input tensor of shape [num_tokens, 1, rope_head_dim=64], dtype=bfloat16
"""
module = _jit_concat_mla_k_module()
module.concat_mla_k(k, k_nope, k_rope)
def concat_mla_absorb_q(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""
Concatenate tensors a and b for MLA absorbed Q computation.
Args:
a: Input tensor of shape [dim_0, dim_1, a_last_dim], dtype=bfloat16
b: Input tensor of shape [dim_0, dim_1, b_last_dim], dtype=bfloat16
Returns:
Output tensor of shape [dim_0, dim_1, a_last_dim + b_last_dim], dtype=bfloat16
"""
out = torch.empty(
(*a.shape[:-1], a.shape[-1] + b.shape[-1]),
dtype=a.dtype,
device=a.device,
)
module = _jit_concat_mla_absorb_q_module()
module.concat_mla_absorb_q(a, b, out)
return out

59
third_party/sglang/python/sglang/jit_kernel/csrc/add_constant.cuh vendored Normal file
View File

@@ -0,0 +1,59 @@
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
#include <sgl_kernel/utils.h> // For div_ceil, RuntimeCheck
#include <sgl_kernel/utils.cuh> // For LaunchKernel
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <cstddef>
#include <cstdint>
namespace {
template <int32_t kConstant>
__global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < length) {
dst[idx] = src[idx] + kConstant;
}
}
constexpr size_t kBlockSize = 256;
// You can also use struct with static method as an alternative
template <int32_t kConstant>
void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
using namespace host;
// 1. Validate input tensors
SymbolicSize N = {"num_elements"};
SymbolicDevice device_;
TensorMatcher({N}) // 1D tensor, must be contiguous
.with_dtype<int32_t>() // must be int32
.with_device<kDLCUDA>(device_) // must be on CUDA device
.verify(dst) // check tensor dst
.verify(src); // check tensor src
// 2. Extract required parameters, prepare for kernel launch
const size_t num_elements = N.unwrap();
const size_t grid_size = div_ceil(num_elements, kBlockSize);
const DLDevice device = device_.unwrap();
[[maybe_unused]] // optional, can be omitted
const size_t dynamic_smem = 0;
[[maybe_unused]] // optional, LaunchKernel can auto determine stream from device
const cudaStream_t stream = LaunchKernel::resolve_device(device);
// some extra runtime checks using host::RuntimeCheck
RuntimeCheck(num_elements > 0, "We only support non-empty tensors, got num_elements = ", num_elements);
// 3. Launch the kernel. Error code will be automatically checked.
LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)(
// kernel function
add_constant_kernel<kConstant>,
// kernel arguments
static_cast<int32_t*>(dst.data_ptr()),
static_cast<int32_t*>(src.data_ptr()),
num_elements);
}
} // namespace

246
third_party/sglang/python/sglang/jit_kernel/csrc/diffusion/qknorm_rope.cuh vendored Normal file
View File

@@ -0,0 +1,246 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <dlpack/dlpack.h>
#include <cstdint>
#include <type_traits>
namespace {
struct QKNormRopeParams {
void* __restrict__ q_ptr;
void* __restrict__ k_ptr; // pre-offset by -num_qo_heads * head_stride_bytes
const void* __restrict__ q_weight_ptr;
const void* __restrict__ k_weight_ptr;
const void* __restrict__ cos_sin_cache_ptr;
const void* __restrict__ positions;
int64_t q_stride_bytes;
int64_t k_stride_bytes;
int64_t head_stride_bytes;
uint32_t num_qo_heads;
uint32_t num_kv_heads;
uint32_t num_tokens;
float eps;
};
constexpr uint32_t kThreadsPerBlock = 256;
constexpr uint32_t kWarpsPerBlock = kThreadsPerBlock / device::kWarpThreads;
template <uint32_t kLaneCount>
constexpr uint32_t active_mask() {
static_assert(kLaneCount <= device::kWarpThreads, "active_mask lane count must not exceed warp size");
if constexpr (kLaneCount == device::kWarpThreads) {
return 0xffffffffu;
} else {
return (1u << kLaneCount) - 1u;
}
}
SGL_DEVICE float load_cache_value(const float* ptr, int64_t idx) {
#ifdef USE_ROCM
return ptr[idx];
#else
return __ldg(ptr + idx);
#endif
}
template <int64_t kHeadDim, int64_t kRopeDim, bool kIsNeox, bool kUsePDL, typename DType, typename IdType>
__global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__ params) {
using namespace device;
static_assert(std::is_same_v<DType, fp16_t> || std::is_same_v<DType, bf16_t>);
static_assert(kHeadDim <= 256, "Only warp-level fused qknorm+rope is supported");
static_assert(kHeadDim % kWarpThreads == 0, "head_dim must be divisible by warp size");
constexpr uint32_t kElemsPerThread = kHeadDim / kWarpThreads;
constexpr uint32_t kVecSize = kElemsPerThread / 2;
constexpr uint32_t kRotaryLanes = kRopeDim / kElemsPerThread;
constexpr uint32_t kHalfRotaryLanes = kRotaryLanes / 2;
constexpr uint32_t kActiveMask = active_mask<kRotaryLanes>();
constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float);
static_assert(kElemsPerThread % 2 == 0, "Each lane must own an even number of elements");
static_assert(kRopeDim > 0 && kRopeDim <= kHeadDim, "Invalid rope dimension");
static_assert(kRopeDim % kElemsPerThread == 0, "rope_dim must align with per-lane vector width");
static_assert(
!kIsNeox || (kRotaryLanes >= 2 && ((kRotaryLanes & (kRotaryLanes - 1)) == 0)),
"NeoX fused qknorm+rope requires rotary lane count to be a power of 2");
using Packed = packed_t<DType>;
using Storage = AlignedVector<Packed, kVecSize>;
const auto& [q_ptr, k_ptr, q_weight_ptr, k_weight_ptr, cos_sin_cache_ptr, positions, q_stride_bytes, k_stride_bytes, head_stride_bytes, num_qo_heads, num_kv_heads, num_tokens, eps] =
params;
const uint32_t lane_id = threadIdx.x % kWarpThreads;
const uint32_t warp_id = threadIdx.x / kWarpThreads;
const uint32_t start_worker_id = blockIdx.x * kWarpsPerBlock + warp_id;
const uint32_t num_workers = gridDim.x * kWarpsPerBlock;
const uint32_t num_qk_heads = num_qo_heads + num_kv_heads;
const uint32_t num_works = num_qk_heads * num_tokens;
PDLWaitPrimary<kUsePDL>();
for (uint32_t idx = start_worker_id; idx < num_works; idx += num_workers) {
const uint32_t token_id = idx / num_qk_heads;
const uint32_t head_id = idx % num_qk_heads;
const bool load_q = head_id < num_qo_heads;
const void* input = load_q ? pointer::offset(q_ptr, token_id * q_stride_bytes, head_id * head_stride_bytes)
: pointer::offset(k_ptr, token_id * k_stride_bytes, head_id * head_stride_bytes);
const void* weight_ptr = load_q ? q_weight_ptr : k_weight_ptr;
auto input_vec = load_as<Storage>(input, lane_id);
const auto weight_vec = load_as<Storage>(weight_ptr, lane_id);
float elems[kElemsPerThread];
float sum_of_squares = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < kVecSize; ++j) {
const auto [x0, x1] = cast<fp32x2_t>(input_vec[j]);
elems[2 * j] = x0;
elems[2 * j + 1] = x1;
sum_of_squares += x0 * x0 + x1 * x1;
}
sum_of_squares = warp::reduce_sum(sum_of_squares);
const float norm_factor = math::rsqrt(sum_of_squares / static_cast<float>(kHeadDim) + eps);
#pragma unroll
for (uint32_t j = 0; j < kVecSize; ++j) {
const auto [w0, w1] = cast<fp32x2_t>(weight_vec[j]);
elems[2 * j] *= norm_factor * w0;
elems[2 * j + 1] *= norm_factor * w1;
}
if constexpr (kIsNeox) {
if (lane_id < kRotaryLanes) {
const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
const auto cos_ptr = static_cast<const float*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
const auto sin_ptr = cos_ptr + kRopeDim / 2;
#pragma unroll
for (uint32_t i = 0; i < kElemsPerThread; ++i) {
float swapped = __shfl_xor_sync(kActiveMask, elems[i], kHalfRotaryLanes);
if (lane_id < kHalfRotaryLanes) {
swapped = -swapped;
}
int dim_idx = static_cast<int>(lane_id * kElemsPerThread + i);
dim_idx = (dim_idx * 2) % kRopeDim;
const int half_idx = dim_idx / 2;
const float cos = load_cache_value(cos_ptr, half_idx);
const float sin = load_cache_value(sin_ptr, half_idx);
elems[i] = elems[i] * cos + swapped * sin;
}
}
} else {
if (lane_id < kRotaryLanes) {
const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
const auto cos_ptr = static_cast<const float*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
const auto sin_ptr = cos_ptr + kRopeDim / 2;
#pragma unroll
for (uint32_t i = 0; i < kElemsPerThread; i += 2) {
const float x = elems[i];
const float y = elems[i + 1];
const int half_idx = static_cast<int>(lane_id * kElemsPerThread + i) / 2;
const float cos = load_cache_value(cos_ptr, half_idx);
const float sin = load_cache_value(sin_ptr, half_idx);
elems[i] = x * cos - y * sin;
elems[i + 1] = y * cos + x * sin;
}
}
}
#pragma unroll
for (uint32_t j = 0; j < kVecSize; ++j) {
input_vec[j] = cast<Packed, fp32x2_t>({elems[2 * j], elems[2 * j + 1]});
}
store_as<Storage>(const_cast<void*>(input), input_vec, lane_id);
}
PDLTriggerSecondary<kUsePDL>();
}
template <int64_t kHeadDim, int64_t kRopeDim, bool kIsNeox, bool kUsePDL, typename DType>
struct QKNormRopeKernel {
static_assert(kHeadDim <= 256, "Only head_dim <= 256 is supported");
template <typename IdType>
static constexpr auto kernel = fused_qknorm_rope_warp<kHeadDim, kRopeDim, kIsNeox, kUsePDL, DType, IdType>;
static void
run(const tvm::ffi::TensorView q,
const tvm::ffi::TensorView k,
const tvm::ffi::TensorView q_weight,
const tvm::ffi::TensorView k_weight,
const tvm::ffi::TensorView cos_sin_cache,
const tvm::ffi::TensorView positions,
float eps) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto Q = SymbolicSize{"num_qo_heads"};
auto K = SymbolicSize{"num_kv_heads"};
auto D = SymbolicSize{"head_dim"};
auto R = SymbolicSize{"rope_dim"};
auto Dq = SymbolicSize{"q_stride"};
auto Dk = SymbolicSize{"k_stride"};
auto Dd = SymbolicSize{"head_stride"};
auto device = SymbolicDevice{};
auto id_type = SymbolicDType{};
D.set_value(kHeadDim);
R.set_value(kRopeDim);
device.set_options<kDLCUDA>();
TensorMatcher({N, Q, D}).with_strides({Dq, Dd, 1}).with_dtype<DType>().with_device(device).verify(q);
TensorMatcher({N, K, D}).with_strides({Dk, Dd, 1}).with_dtype<DType>().with_device(device).verify(k);
TensorMatcher({D}).with_dtype<DType>().with_device(device).verify(q_weight).verify(k_weight);
TensorMatcher({-1, R}).with_dtype<float>().with_device(device).verify(cos_sin_cache);
TensorMatcher({N}).with_dtype<int32_t, int64_t>(id_type).with_device(device).verify(positions);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
const auto q_stride_bytes = static_cast<int64_t>(Dq.unwrap() * sizeof(DType));
const auto k_stride_bytes = static_cast<int64_t>(Dk.unwrap() * sizeof(DType));
const auto head_stride_bytes = static_cast<int64_t>(Dd.unwrap() * sizeof(DType));
const int64_t k_offset = static_cast<int64_t>(num_qo_heads) * head_stride_bytes;
const auto params = QKNormRopeParams{
.q_ptr = q.data_ptr(),
.k_ptr = pointer::offset(k.data_ptr(), -k_offset),
.q_weight_ptr = q_weight.data_ptr(),
.k_weight_ptr = k_weight.data_ptr(),
.cos_sin_cache_ptr = cos_sin_cache.data_ptr(),
.positions = positions.data_ptr(),
.q_stride_bytes = q_stride_bytes,
.k_stride_bytes = k_stride_bytes,
.head_stride_bytes = head_stride_bytes,
.num_qo_heads = num_qo_heads,
.num_kv_heads = num_kv_heads,
.num_tokens = num_tokens,
.eps = eps,
};
const auto is_int32 = id_type.is_type<int32_t>();
const auto selected_kernel = is_int32 ? kernel<int32_t> : kernel<int64_t>;
const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
static const uint32_t kOccupancyTable[2] = {
runtime::get_blocks_per_sm(kernel<int32_t>, kThreadsPerBlock),
runtime::get_blocks_per_sm(kernel<int64_t>, kThreadsPerBlock),
};
const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM;
const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens;
const auto needed_blocks = div_ceil(num_works, kWarpsPerBlock);
const auto num_blocks = std::min(max_blocks, needed_blocks);
LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap()).enable_pdl(kUsePDL)(selected_kernel, params);
}
};
} // namespace

150
third_party/sglang/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh vendored Normal file
View File

@@ -0,0 +1,150 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cuda_runtime.h>
#include <type_traits>
namespace {
template <bool kFlipSinToCos, typename TIn>
__global__ void timestep_embedding_kernel(
const TIn* __restrict__ t_ptr,
float* __restrict__ output_ptr,
int dim,
float neg_log_max_period,
float scale,
int batch_size) {
int row_idx = static_cast<int>(blockIdx.x * blockDim.y + threadIdx.y);
if (row_idx >= batch_size) {
return;
}
float t_val = device::cast<float>(t_ptr[row_idx]);
float* output_batch_base_ptr = output_ptr + row_idx * dim;
int half_dim = dim / 2;
int thread_offset = static_cast<int>(threadIdx.x);
while (thread_offset * 4 < half_dim) {
float4* top_half;
float4* bottom_half;
if constexpr (!kFlipSinToCos) {
bottom_half = reinterpret_cast<float4*>(output_batch_base_ptr + thread_offset * 4);
top_half = reinterpret_cast<float4*>(output_batch_base_ptr + half_dim + thread_offset * 4);
} else {
top_half = reinterpret_cast<float4*>(output_batch_base_ptr + thread_offset * 4);
bottom_half = reinterpret_cast<float4*>(output_batch_base_ptr + half_dim + thread_offset * 4);
}
float4 vals;
vals.x = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 0));
vals.y = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 1));
vals.z = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 2));
vals.w = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 3));
float4 cos_vals;
cos_vals.x = device::math::cos(vals.x);
cos_vals.y = device::math::cos(vals.y);
cos_vals.z = device::math::cos(vals.z);
cos_vals.w = device::math::cos(vals.w);
*top_half = cos_vals;
float4 sin_vals;
sin_vals.x = device::math::sin(vals.x);
sin_vals.y = device::math::sin(vals.y);
sin_vals.z = device::math::sin(vals.z);
sin_vals.w = device::math::sin(vals.w);
*bottom_half = sin_vals;
thread_offset += static_cast<int>(blockDim.x);
}
}
template <typename TIn>
inline void launch_timestep_embedding(
const tvm::ffi::TensorView t,
const tvm::ffi::TensorView output,
int dim,
bool flip_sin_to_cos,
float downscale_freq_shift,
float scale,
int max_period) {
using namespace host;
const int batch_size = static_cast<int>(t.shape()[0]);
const int half_dim = dim / 2;
constexpr int kMaxThreadsPerBlock = 1024;
constexpr int kMinThreadsPerBlock = 128;
const int num_threads_per_row = std::min(kMaxThreadsPerBlock, half_dim / 4);
const int num_rows = (kMinThreadsPerBlock + num_threads_per_row - 1) / num_threads_per_row;
dim3 grid((batch_size + num_rows - 1) / num_rows);
dim3 block(num_threads_per_row, num_rows);
const float neg_log_max_period =
std::log(static_cast<float>(max_period)) * (-1.0f) / (static_cast<float>(half_dim) - downscale_freq_shift);
const DLDevice device = output.device();
if (flip_sin_to_cos) {
LaunchKernel(grid, block, device)(
timestep_embedding_kernel<true, TIn>,
static_cast<const TIn*>(t.data_ptr()),
static_cast<float*>(output.data_ptr()),
dim,
neg_log_max_period,
scale,
batch_size);
} else {
LaunchKernel(grid, block, device)(
timestep_embedding_kernel<false, TIn>,
static_cast<const TIn*>(t.data_ptr()),
static_cast<float*>(output.data_ptr()),
dim,
neg_log_max_period,
scale,
batch_size);
}
}
template <typename TIn>
void timestep_embedding(
tvm::ffi::TensorView input,
tvm::ffi::TensorView output,
int dim,
bool flip_sin_to_cos,
float downscale_freq_shift,
float scale,
int max_period) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto D = SymbolicSize{"dim"};
auto device = SymbolicDevice{};
TensorMatcher({B}) // input
.with_strides({1})
.with_dtype<TIn>()
.template with_device<kDLCUDA>(device)
.verify(input);
TensorMatcher({B, D}).with_strides({D, 1}).with_dtype<float>().template with_device<kDLCUDA>(device).verify(output);
RuntimeCheck(D.unwrap() == dim, "Output dim mismatch: ", D.unwrap(), " vs ", dim);
RuntimeCheck(dim % 8 == 0, "dim must align to 8, got ", dim);
launch_timestep_embedding<TIn>(input, output, dim, flip_sin_to_cos, downscale_freq_shift, scale, max_period);
}
} // namespace

27
third_party/sglang/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh vendored Normal file
View File

@@ -0,0 +1,27 @@
#include <sgl_kernel/ffi.h>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/distributed/custom_all_reduce.cuh>
#include <cstdint>
#include <cstring>
inline void register_custom_all_reduce() {
namespace refl = tvm::ffi::reflection;
using Class = host::distributed::CustomAllReduceBase;
refl::ObjectDef<Class>()
.def(refl::init<uint32_t, uint32_t, uint32_t, uint32_t, int64_t, int64_t, int64_t>(), "__init__")
.def("share_storage", &Class::share_storage)
.def("share_graph_inputs", &Class::share_graph_inputs)
.def("post_init", &Class::post_init)
.def("register_inputs", &Class::register_inputs)
.def("set_cuda_graph_capture", &Class::set_cuda_graph_capture)
.def("free_ipc_handles", &Class::free_ipc_handles)
.def("free_storage", &Class::free_storage)
.def("configure_pull", &Class::configure_pull);
}

205
third_party/sglang/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh vendored Normal file
View File

@@ -0,0 +1,205 @@
// Partially migrated from AOT kernel:
// https://github.com/sgl-project/sglang/blob/v0.5.9/sgl-kernel/csrc/allreduce/custom_all_reduce.cu
// Which was originally adapted from:
// https://github.com/vllm-project/vllm/blob/v0.8.2/csrc/custom_all_reduce.cu
// We redesign the controller interface to minimize control plane traffic,
// and fuse the reduce-scatter and broadcast in the 2-shot all reduce
#include <sgl_kernel/ffi.h>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/distributed/common.cuh>
#include <sgl_kernel/distributed/custom_all_reduce.cuh>
#include <bit>
#include <cstdint>
#include <cstring>
namespace {
using device::distributed::PullController;
using host::distributed::AllReduceData;
using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef;
struct AllReduceParams {
void* __restrict__ output;
uint32_t rank;
uint32_t num_items; // NOTE: support at most 4G, but that's too much
};
[[maybe_unused]]
SGL_DEVICE void prefetch_uniform_ptr(const void* ptr) {
asm volatile("prefetchu.L1 [%0];" ::"l"(ptr) : "memory");
}
#define CUSTOM_AR_KERNEL __global__ __launch_bounds__(1024, 1)
template <bool kBroadcast, typename DType, uint32_t kNumGPU>
SGL_DEVICE void all_reduce_impl(const AllReduceParams& params, DType* (&input)[kNumGPU]) {
using namespace device;
constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
using DType2 = packed_t<DType>;
using Storage = AlignedVector<DType2, kVecSize>;
const auto& [output, rank, num_items] = params;
for (auto i = blockIdx.x;; i += gridDim.x) {
const auto offset = i * blockDim.x + threadIdx.x;
if (offset * kVecSize * 2 >= num_items) break;
Storage storage[kNumGPU];
#pragma unroll
for (uint32_t i = 0; i < kNumGPU; ++i) {
storage[i].load(input[i], offset);
}
const Storage result = distributed::reduce_impl(storage);
if constexpr (kBroadcast) {
#pragma unroll
for (uint32_t i = 0; i < kNumGPU; ++i) {
result.store(input[i], offset);
}
} else {
result.store(output, offset);
}
}
}
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
CUSTOM_AR_KERNEL void all_reduce_one_shot_kernel(
const AllReduceData* __restrict__ data,
const AllReduceParams __grid_constant__ params,
const PullController __grid_constant__ ctrl) {
/// NOTE: we assume the data array is ready before the previous kernel
DType* input[kNumGPU];
prefetch_uniform_ptr(data);
#pragma unroll
for (uint32_t i = 0; i < kNumGPU; ++i)
input[i] = static_cast<DType*>(data->input[i]);
device::PDLWaitPrimary<kUsePDL>();
ctrl.sync</*kFence=*/0, /*kStart=*/1>(params.rank, kNumGPU);
all_reduce_impl</*kBroadcast=*/false>(params, input);
device::PDLTriggerSecondary<kUsePDL>();
ctrl.sync</*kFence=*/0, /*kStart=*/0>(params.rank, kNumGPU);
}
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
CUSTOM_AR_KERNEL void all_reduce_two_shot_kernel(
const AllReduceData* __restrict__ data,
const AllReduceParams __grid_constant__ params,
const PullController __grid_constant__ ctrl) {
// get the range of this rank
using device::kWarpThreads, device::div_ceil;
prefetch_uniform_ptr(data);
DType* input[kNumGPU];
#pragma unroll
for (uint32_t i = 0; i < kNumGPU; ++i)
input[i] = static_cast<DType*>(data->input[i]);
constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
const uint32_t num_items = params.num_items;
const uint32_t total_vec = num_items / (kVecSize * 2); // must be divisible here
const uint32_t vec_per_rank = div_ceil(div_ceil(total_vec, kNumGPU), kWarpThreads) * kWarpThreads;
const uint32_t local_vec_start = min(params.rank * vec_per_rank, total_vec);
const uint32_t local_vec_finish = min(local_vec_start + vec_per_rank, total_vec);
const uint32_t local_start = local_vec_start * kVecSize * 2;
const uint32_t local_length = (local_vec_finish - local_vec_start) * kVecSize * 2;
const auto local_params = AllReduceParams{
.output = nullptr, // this is not used for 2-shot all reduce
.rank = params.rank,
.num_items = local_length,
};
#pragma unroll
for (uint32_t i = 0; i < kNumGPU; ++i)
input[i] += local_start;
device::PDLWaitPrimary<kUsePDL>();
ctrl.sync</*kFence=*/0, /*kStart=*/1>(params.rank, kNumGPU);
all_reduce_impl</*kBroadcast=*/true>(local_params, input);
device::PDLTriggerSecondary<kUsePDL>();
ctrl.sync</*kFence=*/1, /*kStart=*/0>(params.rank, kNumGPU);
}
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
struct CustomAllReducePull : public CustomAllReduceBase {
static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
static constexpr auto one_shot_kernel = all_reduce_one_shot_kernel<DType, kNumGPU, kUsePDL>;
static constexpr auto two_shot_kernel = all_reduce_two_shot_kernel<DType, kNumGPU, kUsePDL>;
static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs");
tvm::ffi::Tensor all_reduce(tvm::ffi::Tensor input, int shot) {
using namespace host;
const bool use_2shot = (shot == 2);
const auto device = input.device();
const auto input_ptr = input.data_ptr();
const auto buffer_ptr = get_pull_buffer(m_storage);
const auto num_items_int64 = input.numel();
const auto num_items = static_cast<uint32_t>(num_items_int64);
const auto items_per_block = m_cta_size * kVecSize * 2;
const auto needed_blocks = div_ceil(num_items, items_per_block);
const auto num_blocks = std::min(needed_blocks, m_num_cta);
const auto kernel = use_2shot ? two_shot_kernel : one_shot_kernel;
// only 1-shot + graph capture need extra output buffer
const auto output = (m_is_graph_capturing && !use_2shot) ? ffi::empty_like(input) : input;
const auto params = AllReduceParams{
.output = use_2shot ? nullptr : output.data_ptr(),
.rank = m_rank,
.num_items = num_items,
};
RuntimeCheck(input.IsContiguous(), "Input tensor must be contiguous");
RuntimeCheck(m_num_gpu == kNumGPU, "Mismatch GPU count");
RuntimeCheck(shot == 1 || shot == 2, "Invalid shot count: ", shot);
RuntimeCheck(device.device_type == kDLCUDA, "Only CUDA device is supported");
RuntimeCheck(is_type<DType>(input.dtype()), "Input dtype mismatch");
RuntimeCheck(std::bit_cast<intptr_t>(input_ptr) % 16 == 0, "Input pointer is not properly aligned");
RuntimeCheck(m_pull_ctrl.has_value(), "Controller is not initialized");
RuntimeCheck(static_cast<int64_t>(num_items) == num_items_int64, "Number of items exceeds 4G limit");
const auto& ctrl = *m_pull_ctrl;
const auto stream = LaunchKernel::resolve_device(device);
auto launch = LaunchKernel{num_blocks, m_cta_size, stream};
launch.enable_pdl(kUsePDL);
const auto check_capturing = [&] {
if (!m_is_graph_capturing) return false; // override to avoid cudaRT call overhead
cudaStreamCaptureStatus status;
RuntimeDeviceCheck(cudaStreamIsCapturing(stream, &status));
return status == cudaStreamCaptureStatusActive;
};
if (check_capturing()) {
// no-op if not really capturing, we're in a dummy run
const auto data_ptr = allocate_graph_capture_input(input_ptr);
/// NOTE: we assume when the graph is replayed, the data_ptr should be ready
launch(kernel, data_ptr, params, ctrl);
} else {
// 1.copy the input to the buffer
const auto input_bytes = static_cast<int64_t>(sizeof(DType) * num_items);
RuntimeCheck(input_bytes <= m_pull_buffer_bytes, "Input is too large, num items: ", num_items);
RuntimeDeviceCheck(cudaMemcpyAsync(buffer_ptr, input_ptr, input_bytes, cudaMemcpyDeviceToDevice, stream));
// 2. launch the all reduce kernel
const auto data_ptr = get_data_ptr(); // use default buffer
launch(kernel, data_ptr, params, ctrl);
if (use_2shot) { // 3. copy the reduced result back to the output, because 2-shot doesn't write to output
RuntimeDeviceCheck(cudaMemcpyAsync(input_ptr, buffer_ptr, input_bytes, cudaMemcpyDeviceToDevice, stream));
}
}
return output;
}
};
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
tvm::ffi::Tensor custom_all_reduce(CustomAllReduceRef obj, tvm::ffi::Tensor input, int shot) {
using Impl = CustomAllReducePull<DType, kNumGPU, kUsePDL>;
return static_cast<Impl&>(*obj.get()).all_reduce(input, shot);
}
} // namespace

253
third_party/sglang/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh vendored Normal file
View File

@@ -0,0 +1,253 @@
// Partially adapted from:
// https://github.com/flashinfer-ai/flashinfer/blob/v0.6.4/include/flashinfer/comm/trtllm_allreduce_fusion.cuh
// We simplify the lamport design and minimize the ring buffer count (from 3 -> 2)
#include <sgl_kernel/ffi.h>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/distributed/common.cuh>
#include <sgl_kernel/distributed/custom_all_reduce.cuh>
#include <cstdint>
#include <cstring>
namespace {
using device::distributed::PushController;
using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef;
struct AllReducePushData {
void* __restrict__ buffer[device::distributed::kMaxNumGPU];
const void* input;
void* output;
uint32_t rank;
uint32_t num_items;
uint32_t buffer_bytes;
uint32_t epoch_bytes;
};
#define CUSTOM_AR_KERNEL __global__ __launch_bounds__(1024, 1)
template <typename T>
struct fp_trait {};
// TODO: support more dtypes
template <>
struct fp_trait<bf16_t> {
using type = uint16_t;
[[maybe_unused]]
static constexpr uint16_t pos_zero = 0x0000u;
[[maybe_unused]]
static constexpr uint16_t neg_zero = 0x8000u;
};
template <>
struct fp_trait<fp16_t> {
using type = uint16_t;
[[maybe_unused]]
static constexpr uint16_t pos_zero = 0x0000u;
[[maybe_unused]]
static constexpr uint16_t neg_zero = 0x8000u;
};
template <>
struct fp_trait<float> {
using type = uint32_t;
[[maybe_unused]]
static constexpr uint32_t pos_zero = 0x00000000u;
[[maybe_unused]]
static constexpr uint32_t neg_zero = 0x80000000u;
};
template <typename DType>
SGL_DEVICE void clear_pos_zero(DType& val) {
using Trait = fp_trait<DType>;
const auto ptr = reinterpret_cast<typename Trait::type*>(&val);
if (*ptr == Trait::pos_zero) *ptr = Trait::neg_zero;
}
template <typename DType>
SGL_DEVICE bool is_pos_zero(const DType& val) {
using Trait = fp_trait<DType>;
const auto ptr = reinterpret_cast<const typename Trait::type*>(&val);
return *ptr == Trait::pos_zero;
}
template <typename DType>
SGL_DEVICE DType get_pos_zero() {
using Trait = fp_trait<DType>;
const auto value = Trait::pos_zero;
return *reinterpret_cast<const DType*>(&value);
}
template <typename T>
SGL_DEVICE void ld_global_volatile_16B(T& x, const void* addr, int64_t offset) {
static_assert(alignof(T) == 16 && sizeof(T) == 16);
addr = device::pointer::offset<T>(addr, offset);
uint4 val;
asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];"
: "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
: "l"(addr));
x = *reinterpret_cast<const T*>(&val);
}
template <typename T>
SGL_DEVICE void st_global_volatile_16B(const T& x, void* addr, int64_t offset) {
static_assert(alignof(T) == 16 && sizeof(T) == 16);
const uint4 val = *reinterpret_cast<const uint4*>(&x);
addr = device::pointer::offset<T>(addr, offset);
asm volatile(
"st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), "l"(addr));
}
template <typename DType, uint32_t kNumGPU>
SGL_DEVICE void push_impl(DType* (&push_buf)[kNumGPU], const void* data, uint32_t num_items) {
using namespace device;
constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
using Storage = AlignedVector<packed_t<DType>, kVecSize>;
for (auto i = blockIdx.x;; i += gridDim.x) {
const auto offset = i * blockDim.x + threadIdx.x;
if (offset * kVecSize * 2 >= num_items) break;
Storage vec;
vec.load(data, offset);
#pragma unroll
for (uint32_t j = 0; j < kVecSize; ++j) {
clear_pos_zero(vec[j].x);
clear_pos_zero(vec[j].y);
}
#pragma unroll
for (uint32_t i = 0; i < kNumGPU; ++i) {
st_global_volatile_16B(vec, push_buf[i], offset);
}
}
}
template <typename DType, uint32_t kNumGPU>
SGL_DEVICE void poll_impl(DType* (&poll_buf)[kNumGPU], void* data, uint32_t num_items) {
using namespace device;
constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
using Storage = AlignedVector<packed_t<DType>, kVecSize>;
for (auto i = blockIdx.x;; i += gridDim.x) {
const auto offset = i * blockDim.x + threadIdx.x;
if (offset * kVecSize * 2 >= num_items) break;
Storage storage[kNumGPU];
while (true) {
bool has_pos_zero = false;
#pragma unroll
for (uint32_t i = 0; i < kNumGPU; ++i) {
ld_global_volatile_16B(storage[i], poll_buf[i], offset);
#pragma unroll
for (auto j = 0; j < kVecSize; ++j) {
has_pos_zero |= is_pos_zero(storage[i][j].x);
has_pos_zero |= is_pos_zero(storage[i][j].y);
}
}
if (!has_pos_zero) break;
}
const Storage result = distributed::reduce_impl(storage);
result.store(data, offset);
Storage pos_zeros;
pos_zeros.fill({get_pos_zero<DType>(), get_pos_zero<DType>()});
#pragma unroll
for (uint32_t i = 0; i < kNumGPU; ++i) {
pos_zeros.store(poll_buf[i], offset);
}
}
}
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
CUSTOM_AR_KERNEL void all_reduce_one_shot_push_kernel(
const AllReducePushData __grid_constant__ params, //
const PushController __grid_constant__ ctrl) {
using namespace device;
const auto [buffer, input, output, rank, num_items, buffer_bytes, epoch_bytes] = params;
PDLWaitPrimary<kUsePDL>();
// Phase 1: Push data from input to all ranks' buffers
const auto epoch_offset = ctrl.epoch() * epoch_bytes;
DType* push_buf[kNumGPU];
#pragma unroll
for (uint32_t i = 0; i < kNumGPU; ++i) {
push_buf[i] = static_cast<DType*>(pointer::offset(buffer[i], rank * buffer_bytes, epoch_offset));
}
push_impl(push_buf, input, num_items);
PDLTriggerSecondary<kUsePDL>();
// Phase 2: Poll local data
DType* poll_buf[kNumGPU];
#pragma unroll
for (uint32_t i = 0; i < kNumGPU; ++i) {
poll_buf[i] = static_cast<DType*>(pointer::offset(buffer[rank], i * buffer_bytes, epoch_offset));
}
poll_impl(poll_buf, output, num_items);
ctrl.exit();
}
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
struct CustomAllReducePush : public CustomAllReduceBase {
static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs");
tvm::ffi::Tensor all_reduce(tvm::ffi::Tensor input, int shot) {
using namespace host;
const auto device = input.device();
const auto input_ptr = input.data_ptr();
const auto num_items_int64 = input.numel();
const auto num_items = static_cast<uint32_t>(num_items_int64);
const auto num_blocks = m_max_num_cta_push; // must be constant to ensure correctness
const auto num_threads = [&] {
for (const auto t : {128u, 256u, 512u}) {
if (t * num_blocks * 2 * kVecSize >= num_items) return t;
}
return 1024u;
}();
const auto output = input;
AllReducePushData params;
for (uint32_t i = 0; i < kNumGPU; ++i) {
params.buffer[i] = get_push_buffer(m_peer_storage[i]);
}
params.input = input_ptr;
params.output = input_ptr;
params.rank = m_rank;
params.num_items = num_items;
params.buffer_bytes = m_push_buffer_bytes;
params.epoch_bytes = kNumGPU * params.buffer_bytes;
RuntimeCheck(input.IsContiguous(), "Input must be contiguous");
RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch");
RuntimeCheck(device.device_type == kDLCUDA, "Only CUDA device is supported");
RuntimeCheck(is_type<DType>(input.dtype()), "Input dtype mismatch");
RuntimeCheck(std::bit_cast<intptr_t>(input_ptr) % 16 == 0, "Input pointer is not properly aligned");
RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized");
RuntimeCheck(shot == 1, "Push all-reduce only supports 1-shot, got: ", shot);
RuntimeCheck(static_cast<int64_t>(num_items) == num_items_int64, "Number of items exceeds 4G limit");
const auto input_bytes = static_cast<int64_t>(sizeof(DType) * num_items_int64);
RuntimeCheck(input_bytes <= m_push_buffer_bytes, "Input is too large, num items: ", num_items);
const auto kernel = all_reduce_one_shot_push_kernel<DType, kNumGPU, kUsePDL>;
LaunchKernel(num_blocks, num_threads, device) //
.enable_pdl(kUsePDL)(kernel, params, *m_push_ctrl);
return output;
}
};
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
tvm::ffi::Tensor custom_all_reduce(CustomAllReduceRef obj, tvm::ffi::Tensor input, int shot) {
using Impl = CustomAllReducePush<DType, kNumGPU, kUsePDL>;
return static_cast<Impl&>(*obj.get()).all_reduce(input, shot);
}
} // namespace

137
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/cast.cuh vendored Normal file
View File

@@ -0,0 +1,137 @@
#pragma once
// Optimized cast kernel: fixed 256 threads, scaled out via 2D grid.
// Each thread handles exactly one float4 (kVecSize fp16/bf16 elements).
// No per-thread loop — pure grid scaling for any head*dim.
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/type.cuh> // For dtype_trait fp8 specialization
#include <sgl_kernel/utils.cuh> // For LaunchKernel
#include <sgl_kernel/vec.cuh> // For AlignedVector
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <cstdint>
namespace {
constexpr int kBlockSize = 256;
template <typename T>
__global__ void fused_downcast_kernel(
const T* __restrict__ cache_k,
const T* __restrict__ cache_v,
const float* __restrict__ k_scale,
const float* __restrict__ v_scale,
fp8_e4m3_t* __restrict__ output_k,
fp8_e4m3_t* __restrict__ output_v,
const int input_num_tokens,
const int head,
const int dim,
const T max_fp8,
const T min_fp8,
const int64_t mult,
const int64_t offset,
const int64_t* __restrict__ loc) {
using namespace device;
constexpr int kVecSize = 16 / sizeof(T);
using vec_t = AlignedVector<T, kVecSize>;
using out_vec_t = AlignedVector<fp8_e4m3_t, kVecSize>;
const int token_idx = blockIdx.x;
const int vec_idx = blockIdx.y * kBlockSize + threadIdx.x;
const int num_vecs = head * dim / kVecSize;
if (token_idx >= input_num_tokens || vec_idx >= num_vecs) return;
T k_scale_inv = static_cast<T>(1.f) / cast<T>(k_scale[0]);
T v_scale_inv = static_cast<T>(1.f) / cast<T>(v_scale[0]);
auto clamp = [&](T val) { return val > max_fp8 ? max_fp8 : (min_fp8 > val ? min_fp8 : val); };
const int out_seq_idx = loc[token_idx];
const T* in_k_base = cache_k + token_idx * head * dim;
const T* in_v_base = cache_v + token_idx * head * dim;
fp8_e4m3_t* out_k_base = output_k + (out_seq_idx * mult + offset) * head * dim;
fp8_e4m3_t* out_v_base = output_v + (out_seq_idx * mult + offset) * head * dim;
vec_t k_vec, v_vec;
k_vec.load(in_k_base, vec_idx);
v_vec.load(in_v_base, vec_idx);
out_vec_t out_k, out_v;
#pragma unroll
for (int j = 0; j < kVecSize; j++) {
out_k[j] = cast<fp8_e4m3_t>(clamp(k_vec[j] * k_scale_inv));
out_v[j] = cast<fp8_e4m3_t>(clamp(v_vec[j] * v_scale_inv));
}
out_k.store(out_k_base, vec_idx);
out_v.store(out_v_base, vec_idx);
}
template <typename T>
void downcast_fp8(
tvm::ffi::TensorView k,
tvm::ffi::TensorView v,
tvm::ffi::TensorView k_out,
tvm::ffi::TensorView v_out,
tvm::ffi::TensorView k_scale,
tvm::ffi::TensorView v_scale,
tvm::ffi::TensorView loc,
int64_t mult,
int64_t offset) {
using namespace host;
auto input_num_tokens = SymbolicSize{"input_num_tokens"};
auto head = SymbolicSize{"head"};
auto dim = SymbolicSize{"dim"};
auto output_num_tokens = SymbolicSize{"out_sl"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({input_num_tokens, head, dim}).with_dtype<T>().with_device(device).verify(k);
TensorMatcher({input_num_tokens, head, dim}).with_dtype<T>().with_device(device).verify(v);
TensorMatcher({output_num_tokens, head, dim}).with_dtype<uint8_t>().with_device(device).verify(k_out);
TensorMatcher({output_num_tokens, head, dim}).with_dtype<uint8_t>().with_device(device).verify(v_out);
TensorMatcher({1}).with_dtype<float>().with_device(device).verify(k_scale);
TensorMatcher({1}).with_dtype<float>().with_device(device).verify(v_scale);
TensorMatcher({input_num_tokens}).with_dtype<int64_t>().with_device(device).verify(loc);
const int num_tokens = static_cast<int>(input_num_tokens.unwrap());
const int h = static_cast<int>(head.unwrap());
const int d = static_cast<int>(dim.unwrap());
constexpr int kVecSize = 16 / sizeof(T);
const int num_vecs = h * d / kVecSize;
const int grid_y = (num_vecs + kBlockSize - 1) / kBlockSize;
dim3 grid(num_tokens, grid_y);
dim3 block(kBlockSize);
const T max_fp8 = static_cast<T>(kFP8E4M3Max);
const T min_fp8 = static_cast<T>(-kFP8E4M3Max);
LaunchKernel(grid, block, device.unwrap())(
fused_downcast_kernel<T>,
static_cast<const T*>(k.data_ptr()),
static_cast<const T*>(v.data_ptr()),
static_cast<const float*>(k_scale.data_ptr()),
static_cast<const float*>(v_scale.data_ptr()),
static_cast<fp8_e4m3_t*>(k_out.data_ptr()),
static_cast<fp8_e4m3_t*>(v_out.data_ptr()),
num_tokens,
h,
d,
max_fp8,
min_fp8,
mult,
offset,
static_cast<const int64_t*>(loc.data_ptr()));
}
} // namespace

54
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/clamp_position.cuh vendored Normal file
View File

@@ -0,0 +1,54 @@
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
#include <sgl_kernel/utils.h> // For div_ceil
#include <sgl_kernel/utils.cuh> // For LaunchKernel
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <cstddef>
#include <cstdint>
namespace {
template <typename T>
__global__ void clamp_position_kernel(T* __restrict__ dst, const T* __restrict__ seq_lens, size_t n) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
T val = seq_lens[idx] - 1;
dst[idx] = val < 0 ? 0 : val;
}
}
constexpr size_t kBlockSize = 256;
template <typename T>
struct ClampPosition {
static void run(tvm::ffi::TensorView dst, tvm::ffi::TensorView seq_lens) {
using namespace host;
SymbolicSize N = {"num_elements"};
SymbolicDevice device_;
device_.set_options<kDLCUDA, kDLROCM>();
TensorMatcher({N}) //
.with_dtype<T>()
.with_device(device_)
.verify(dst)
.verify(seq_lens);
const size_t num_elements = N.unwrap();
if (num_elements == 0) return;
const size_t grid_size = div_ceil(num_elements, kBlockSize);
const DLDevice device = device_.unwrap();
LaunchKernel(grid_size, kBlockSize, device)(
clamp_position_kernel<T>,
static_cast<T*>(dst.data_ptr()),
static_cast<const T*>(seq_lens.data_ptr()),
num_elements);
}
};
} // namespace

325
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/concat_mla.cuh vendored Normal file
View File

@@ -0,0 +1,325 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <tvm/ffi/container/tensor.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
namespace {
// ======================= Memory Utilities =======================
// Adapted from DeepEP: https://github.com/deepseek-ai/DeepEP/blob/main/csrc/kernels/utils.cuh
SGL_DEVICE int get_lane_id() {
int lane_id;
asm("mov.s32 %0, %laneid;" : "=r"(lane_id));
return lane_id;
}
SGL_DEVICE void st_na_global_v1(const int* ptr, int v) {
asm volatile("st.global.L1::no_allocate.s32 [%0], %1;" ::"l"(ptr), "r"(v) : "memory");
}
SGL_DEVICE void st_na_global_v2(const int2* ptr, const int2& v) {
asm volatile("st.global.L1::no_allocate.v2.s32 [%0], {%1, %2};" ::"l"(ptr), "r"(v.x), "r"(v.y) : "memory");
}
SGL_DEVICE int ld_na_global_v1(const int* ptr) {
int r;
asm volatile("ld.global.nc.L1::no_allocate.s32 %0, [%1];" : "=r"(r) : "l"(ptr));
return r;
}
SGL_DEVICE int2 ld_na_global_v2(const int2* ptr) {
int2 r;
asm volatile("ld.global.nc.L1::no_allocate.v2.s32 {%0, %1}, [%2];" : "=r"(r.x), "=r"(r.y) : "l"(ptr));
return r;
}
SGL_DEVICE void prefetch_L2(const void* p) {
#if defined(ENABLE_L2_PREFETCH)
asm volatile("prefetch.global.L2 [%0];" ::"l"(p));
#endif
}
// ======================= concat_mla_k Kernel =======================
constexpr int NUM_LOCAL_HEADS = 128;
constexpr int QK_NOPE_HEAD_DIM = 128;
constexpr int QK_ROPE_HEAD_DIM = 64;
constexpr int K_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM;
constexpr int HEAD_CHUNK_SIZE = 16;
constexpr int NUM_HEAD_CHUNKS = NUM_LOCAL_HEADS / HEAD_CHUNK_SIZE;
__global__ void concat_mla_k_kernel(
bf16_t* __restrict__ k,
const bf16_t* __restrict__ k_nope,
const bf16_t* __restrict__ k_rope,
const int num_tokens,
const int64_t k_stride_0,
const int k_stride_1,
const int64_t k_nope_stride_0,
const int k_nope_stride_1,
const int64_t k_rope_stride_0) {
const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
const int token_id = flat_warp_id / NUM_HEAD_CHUNKS;
const int head_chunk_id = flat_warp_id % NUM_HEAD_CHUNKS;
const int lane_id = get_lane_id();
if (token_id >= num_tokens) return;
using NopeVec = int2; // 8B/thread, 32 threads = 256B/row
using RopeVec = int; // 4B/thread, 32 threads = 128B/row
static_assert(sizeof(NopeVec) * 32 == QK_NOPE_HEAD_DIM * sizeof(bf16_t), "nope vec mismatch");
static_assert(sizeof(RopeVec) * 32 == QK_ROPE_HEAD_DIM * sizeof(bf16_t), "rope vec mismatch");
const int head_row0 = head_chunk_id * HEAD_CHUNK_SIZE;
const int2* __restrict__ nope_src =
reinterpret_cast<const int2*>(k_nope + token_id * k_nope_stride_0 + head_row0 * k_nope_stride_1) + lane_id;
int2* __restrict__ nope_dst = reinterpret_cast<int2*>(k + token_id * k_stride_0 + head_row0 * k_stride_1) + lane_id;
int* __restrict__ rope_dst =
reinterpret_cast<int*>(k + token_id * k_stride_0 + head_row0 * k_stride_1 + QK_NOPE_HEAD_DIM) + lane_id;
const int nope_src_stride_v = (k_nope_stride_1 >> 2); // int2 covers 4 bf16
const int nope_dst_stride_v = (k_stride_1 >> 2);
const int rope_dst_stride_v = (k_stride_1 >> 1); // int covers 2 bf16
const int* rope_base = reinterpret_cast<const int*>(k_rope + token_id * k_rope_stride_0);
const RopeVec rope_val = ld_na_global_v1(rope_base + lane_id);
prefetch_L2(nope_src);
NopeVec cur = ld_na_global_v2(nope_src);
#pragma unroll
for (int i = 0; i < HEAD_CHUNK_SIZE; ++i) {
NopeVec next;
if (i + 1 < HEAD_CHUNK_SIZE) {
const int2* next_src = nope_src + nope_src_stride_v;
prefetch_L2(next_src);
next = ld_na_global_v2(next_src);
}
st_na_global_v2(nope_dst, cur);
st_na_global_v1(rope_dst, rope_val);
nope_src += nope_src_stride_v;
nope_dst += nope_dst_stride_v;
rope_dst += rope_dst_stride_v;
cur = next;
}
}
struct ConcatMlaKKernel {
static void run(tvm::ffi::TensorView k, tvm::ffi::TensorView k_nope, tvm::ffi::TensorView k_rope) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto H = SymbolicSize{"num_heads"};
auto D = SymbolicSize{"k_head_dim"};
auto D_nope = SymbolicSize{"nope_head_dim"};
auto D_rope = SymbolicSize{"rope_head_dim"};
auto S0_k = SymbolicSize{"k_stride_0"};
auto S1_k = SymbolicSize{"k_stride_1"};
auto S0_k_nope = SymbolicSize{"k_nope_stride_0"};
auto S1_k_nope = SymbolicSize{"k_nope_stride_1"};
auto S0_k_rope = SymbolicSize{"k_rope_stride_0"};
auto device = SymbolicDevice{};
// Set known fixed values
H.set_value(NUM_LOCAL_HEADS);
D.set_value(K_HEAD_DIM);
D_nope.set_value(QK_NOPE_HEAD_DIM);
D_rope.set_value(QK_ROPE_HEAD_DIM);
// Verify k: [num_tokens, num_heads, k_head_dim]
TensorMatcher({N, H, D}).with_strides({S0_k, S1_k, 1}).with_dtype<bf16_t>().with_device<kDLCUDA>(device).verify(k);
// Verify k_nope: [num_tokens, num_heads, nope_head_dim]
TensorMatcher({N, H, D_nope})
.with_strides({S0_k_nope, S1_k_nope, 1})
.with_dtype<bf16_t>()
.with_device<kDLCUDA>(device)
.verify(k_nope);
// Verify k_rope: [num_tokens, 1, rope_head_dim]
TensorMatcher({N, 1, D_rope})
.with_strides({S0_k_rope, -1, 1})
.with_dtype<bf16_t>()
.with_device<kDLCUDA>(device)
.verify(k_rope);
// Check alignment
RuntimeCheck(reinterpret_cast<uintptr_t>(k.data_ptr()) % 16 == 0, "Tensor k must be 16-byte aligned");
RuntimeCheck(reinterpret_cast<uintptr_t>(k_nope.data_ptr()) % 16 == 0, "Tensor k_nope must be 16-byte aligned");
RuntimeCheck(reinterpret_cast<uintptr_t>(k_rope.data_ptr()) % 16 == 0, "Tensor k_rope must be 16-byte aligned");
const int num_tokens = static_cast<int>(N.unwrap());
constexpr int num_warps_per_block = 32;
const int grid_size = div_ceil(num_tokens * NUM_HEAD_CHUNKS, num_warps_per_block);
const int block_size = num_warps_per_block * 32;
LaunchKernel(grid_size, block_size, device.unwrap())(
concat_mla_k_kernel,
static_cast<bf16_t*>(k.data_ptr()),
static_cast<const bf16_t*>(k_nope.data_ptr()),
static_cast<const bf16_t*>(k_rope.data_ptr()),
num_tokens,
S0_k.unwrap(),
static_cast<int>(S1_k.unwrap()),
S0_k_nope.unwrap(),
static_cast<int>(S1_k_nope.unwrap()),
S0_k_rope.unwrap());
}
};
// ======================= concat_mla_absorb_q Kernel =======================
constexpr int A_LAST_DIM = 512;
constexpr int B_LAST_DIM = 64;
constexpr int OUT_LAST_DIM = A_LAST_DIM + B_LAST_DIM;
__global__ void concat_mla_absorb_q_kernel(
bf16_t* a,
bf16_t* b,
bf16_t* out,
const int num_items,
const int dim_1,
const int64_t a_stride_0,
const int a_stride_1,
const int64_t b_stride_0,
const int b_stride_1,
const int64_t out_stride_0,
const int out_stride_1) {
const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
const int lane_id = get_lane_id();
const int idx_0 = flat_warp_id / dim_1;
const int idx_1 = flat_warp_id % dim_1;
if (flat_warp_id >= num_items) {
return;
}
using ABufType = int4;
constexpr int A_NUM_UNROLL = 2;
static_assert(sizeof(ABufType) * A_NUM_UNROLL == A_LAST_DIM * sizeof(a[0]) / 32);
ABufType a_buf[A_NUM_UNROLL];
using BBufType = int;
constexpr int B_NUM_UNROLL = 1;
static_assert(sizeof(BBufType) * B_NUM_UNROLL == B_LAST_DIM * sizeof(b[0]) / 32);
BBufType b_buf;
{
const BBufType* base_addr = reinterpret_cast<BBufType*>(b + idx_0 * b_stride_0 + idx_1 * b_stride_1);
b_buf = *(base_addr + lane_id);
}
#pragma unroll
for (int i = 0; i < A_NUM_UNROLL; ++i) {
const ABufType* base_addr = reinterpret_cast<ABufType*>(a + idx_0 * a_stride_0 + idx_1 * a_stride_1);
a_buf[i] = *(base_addr + i * 32 + lane_id);
}
{
BBufType* base_addr = reinterpret_cast<BBufType*>(out + idx_0 * out_stride_0 + idx_1 * out_stride_1 + A_LAST_DIM);
*(base_addr + lane_id) = b_buf;
}
#pragma unroll
for (int i = 0; i < A_NUM_UNROLL; ++i) {
ABufType* base_addr = reinterpret_cast<ABufType*>(out + idx_0 * out_stride_0 + idx_1 * out_stride_1);
*(base_addr + i * 32 + lane_id) = a_buf[i];
}
}
struct ConcatMlaAbsorbQKernel {
static void run(tvm::ffi::TensorView a, tvm::ffi::TensorView b, tvm::ffi::TensorView out) {
using namespace host;
auto N0_a = SymbolicSize{"a_dim_0"};
auto N1_a = SymbolicSize{"a_dim_1"};
auto D_a = SymbolicSize{"a_last_dim"};
auto N0_b = SymbolicSize{"b_dim_0"};
auto N1_b = SymbolicSize{"b_dim_1"};
auto D_b = SymbolicSize{"b_last_dim"};
auto N0_out = SymbolicSize{"out_dim_0"};
auto N1_out = SymbolicSize{"out_dim_1"};
auto D_out = SymbolicSize{"out_last_dim"};
auto S0_a = SymbolicSize{"a_stride_0"};
auto S1_a = SymbolicSize{"a_stride_1"};
auto S0_b = SymbolicSize{"b_stride_0"};
auto S1_b = SymbolicSize{"b_stride_1"};
auto S0_out = SymbolicSize{"out_stride_0"};
auto S1_out = SymbolicSize{"out_stride_1"};
auto device = SymbolicDevice{};
// Set known fixed values
D_a.set_value(A_LAST_DIM);
D_b.set_value(B_LAST_DIM);
D_out.set_value(OUT_LAST_DIM);
// Verify a: [dim_0, dim_1, A_LAST_DIM]
TensorMatcher({N0_a, N1_a, D_a})
.with_strides({S0_a, S1_a, 1})
.with_dtype<bf16_t>()
.with_device<kDLCUDA>(device)
.verify(a);
// Verify b: [dim_0, dim_1, B_LAST_DIM]
TensorMatcher({N0_b, N1_b, D_b})
.with_strides({S0_b, S1_b, 1})
.with_dtype<bf16_t>()
.with_device<kDLCUDA>(device)
.verify(b);
// Verify out: [dim_0, dim_1, OUT_LAST_DIM]
TensorMatcher({N0_out, N1_out, D_out})
.with_strides({S0_out, S1_out, 1})
.with_dtype<bf16_t>()
.with_device<kDLCUDA>(device)
.verify(out);
// Check alignment
RuntimeCheck(reinterpret_cast<uintptr_t>(a.data_ptr()) % 16 == 0, "Tensor a must be 16-byte aligned");
RuntimeCheck(reinterpret_cast<uintptr_t>(b.data_ptr()) % 16 == 0, "Tensor b must be 16-byte aligned");
RuntimeCheck(reinterpret_cast<uintptr_t>(out.data_ptr()) % 16 == 0, "Tensor out must be 16-byte aligned");
// Verify dimensions match: a.size(0) * a.size(1) == b.size(0) * b.size(1)
RuntimeCheck(
N0_a.unwrap() * N1_a.unwrap() == N0_b.unwrap() * N1_b.unwrap(),
"Dimension mismatch: a.size(0) * a.size(1) must equal b.size(0) * b.size(1)");
RuntimeCheck(N1_a.unwrap() == N1_b.unwrap(), "Dimension mismatch: a.size(1) must equal b.size(1)");
const int num_items = static_cast<int>(N0_a.unwrap() * N1_a.unwrap());
const int dim_1 = static_cast<int>(N1_a.unwrap());
constexpr int num_warps_per_block = 32;
const int grid_size = div_ceil(num_items, num_warps_per_block);
const int block_size = num_warps_per_block * 32;
LaunchKernel(grid_size, block_size, device.unwrap())(
concat_mla_absorb_q_kernel,
static_cast<bf16_t*>(a.data_ptr()),
static_cast<bf16_t*>(b.data_ptr()),
static_cast<bf16_t*>(out.data_ptr()),
num_items,
dim_1,
S0_a.unwrap(),
static_cast<int>(S1_a.unwrap()),
S0_b.unwrap(),
static_cast<int>(S1_b.unwrap()),
S0_out.unwrap(),
static_cast<int>(S1_out.unwrap()));
}
};
} // namespace

182
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh vendored Normal file
View File

@@ -0,0 +1,182 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <cooperative_groups/reduce.h>
#include <tvm/ffi/container/tensor.h>
#include <cooperative_groups.h>
#include <type_traits>
namespace {
template <typename T, int VEC_SIZE_IN_BYTE>
struct VecTypeTrait;
template <>
struct VecTypeTrait<bf16_t, 16> {
using packed_t = packed_t<bf16_t>;
using vec_t = device::AlignedVector<packed_t, 4>;
};
template <>
struct VecTypeTrait<fp16_t, 16> {
using packed_t = packed_t<fp16_t>;
using vec_t = device::AlignedVector<packed_t, 4>;
};
template <>
struct VecTypeTrait<bf16_t, 32> {
using packed_t = packed_t<bf16_t>;
using vec_t = device::AlignedVector<packed_t, 8>;
};
template <>
struct VecTypeTrait<fp16_t, 32> {
using packed_t = packed_t<fp16_t>;
using vec_t = device::AlignedVector<packed_t, 8>;
};
template <typename packed_t>
SGL_DEVICE packed_t rms(packed_t& val, packed_t& weight, float rsqrt_square_sum) {
float2 valf = device::cast<fp32x2_t, packed_t>(val);
float2 weightf = device::cast<fp32x2_t, packed_t>(weight);
return device::cast<packed_t, fp32x2_t>(
make_float2(valf.x * weightf.x * rsqrt_square_sum, valf.y * weightf.y * rsqrt_square_sum));
}
template <typename T, int VEC_SIZE_IN_BYTE>
__global__ void fused_add_rmsnorm_reg_kernel(
T* __restrict__ input, T* __restrict__ residual, const T* __restrict__ weight, int vec_hidden_size, float eps) {
constexpr int inner_loop = VEC_SIZE_IN_BYTE == 16 ? 4 : 8;
__shared__ float shared_memory[32]; // Used for CTA reduce
using vec_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::vec_t;
using packed_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::packed_t;
vec_t v; // Save input
vec_t v_res; // Save residual
vec_t v_weight; // Save weight
vec_t v_out; // Save output
auto token_id = blockIdx.x;
float2 acc_square = make_float2(0.0f, 0.0f); // Sum of squares for each thread
if (threadIdx.x < vec_hidden_size) {
// Compute address
vec_t* p = reinterpret_cast<vec_t*>(input) + token_id * vec_hidden_size;
vec_t* p_res = reinterpret_cast<vec_t*>(residual) + token_id * vec_hidden_size;
const vec_t* p_weight = reinterpret_cast<const vec_t*>(weight);
// Load data
v = p[threadIdx.x];
v_res = p_res[threadIdx.x];
v_weight = p_weight[threadIdx.x];
for (int i = 0; i < inner_loop; i++) {
float2 val = device::cast<fp32x2_t, packed_t>(v[i]);
float2 res = device::cast<fp32x2_t, packed_t>(v_res[i]);
float2 inp_res = make_float2(val.x + res.x, val.y + res.y);
acc_square.x += inp_res.x * inp_res.x;
acc_square.y += inp_res.y * inp_res.y;
v[i] = device::cast<packed_t, fp32x2_t>(inp_res);
}
// Store inp+res to residual
p_res[threadIdx.x] = v;
}
// CTA Reduce
// Step 0: Warp Reduce
auto cg_warp = cooperative_groups::tiled_partition<32>(cooperative_groups::this_thread_block());
float warp_sum = cooperative_groups::reduce(cg_warp, acc_square.x + acc_square.y, cooperative_groups::plus<float>());
float* buffer = shared_memory;
if (threadIdx.x % 32 == 0) {
buffer[threadIdx.x / 32] = warp_sum; // Write warp_sum to buffer
}
// Step 1: CTA Reduce
__syncthreads();
if (threadIdx.x < 32) {
float cta_sum = cooperative_groups::reduce(
cg_warp, (threadIdx.x < blockDim.x / 32) ? buffer[threadIdx.x] : 0.0f, cooperative_groups::plus<float>());
buffer[threadIdx.x] =
rsqrtf(eps + cta_sum * (1.0f / static_cast<float>(vec_hidden_size * (VEC_SIZE_IN_BYTE / sizeof(T)))));
}
__syncthreads();
// Compute RMSNorm
if (threadIdx.x < vec_hidden_size) {
float rsqrt_square_sum = buffer[threadIdx.x / 32]; // Read rsqrt from Shared Memory(Broadcast)
for (int i = 0; i < inner_loop; i++) {
v_out[i] = rms(v[i], v_weight[i], rsqrt_square_sum);
}
vec_t* p_out = reinterpret_cast<vec_t*>(input) + token_id * vec_hidden_size;
p_out[threadIdx.x] = v_out;
}
}
template <typename DType>
struct FusedAddRMSNormKernel {
static void
run(const tvm::ffi::TensorView input,
const tvm::ffi::TensorView residual,
const tvm::ffi::TensorView weight,
float eps) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto D = SymbolicSize{"hidden_size"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({N, D}) // input
.with_strides({D, 1})
.with_dtype<DType>()
.with_device(device)
.verify(input);
TensorMatcher({D}) // weight
.with_dtype<DType>()
.with_device(device)
.verify(weight);
TensorMatcher({N, D}) // residual
.with_strides({D, 1})
.with_dtype<DType>()
.with_device(device)
.verify(residual);
int hidden_size = static_cast<int>(D.unwrap());
if (hidden_size <= (device::kMaxVecBytes == 32 ? 12288 : 8192)) {
int elements_in_vec = device::kMaxVecBytes / sizeof(DType);
int vec_hidden_size = hidden_size / elements_in_vec;
uint threads = (vec_hidden_size + 31) / 32 * 32;
// Runtime check
host::RuntimeCheck(
hidden_size % elements_in_vec == 0,
"hidden_size",
hidden_size,
" can not align to elements_in_vec ",
elements_in_vec);
// Launch kernel
auto kernel = fused_add_rmsnorm_reg_kernel<DType, device::kMaxVecBytes>;
LaunchKernel(static_cast<uint>(N.unwrap()), threads, device.unwrap())
.enable_pdl(false)(
kernel,
reinterpret_cast<DType*>(input.data_ptr()),
reinterpret_cast<DType*>(residual.data_ptr()),
reinterpret_cast<DType*>(weight.data_ptr()),
vec_hidden_size,
eps);
} else {
host::RuntimeCheck(false, "Large hidden_sizes are not supported for now.");
}
}
};
} // namespace

722
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh vendored Normal file
View File

@@ -0,0 +1,722 @@
/*
* Fused metadata copy kernel for NSA backend CUDA graph replay.
* JIT-compiled version for python/sglang/jit_kernel.
*
* OVERVIEW:
* This kernel fuses multiple tensor copy operations (cache_seqlens, cu_seqlens_k,
* page_table, nsa metadata, and optional FlashMLA metadata) into single kernel
* launches, significantly reducing kernel launch overhead and improving CUDA
* graph replay performance during inference.
*
* PERFORMANCE BENEFITS:
* - Single kernel launch vs. multiple separate copies (3-10x faster)
* - Optimized memory coalescing and SM utilization
* - __grid_constant__ parameter passing via constant memory
* - Especially beneficial in CUDA graph replay scenarios
*
* DESIGN:
* - Unified kernel supporting all forward modes (DECODE, TARGET_VERIFY, DRAFT_EXTEND)
* - Structured parameter passing (SourcePointers/DestinationPointers) for clarity
* - Template parameters (HAS_REAL_PAGE_TABLE, HAS_FLASHMLA) for compile-time optimization
* - Multi-backend variant copies to 3 destinations in one kernel (for speculative decoding)
*
* USAGE:
* This header is included by JIT compilation system. The FusedMetadataCopyKernel
* and FusedMetadataCopyMultiKernel wrapper structs provide the Python-accessible interface.
*/
#pragma once
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <tvm/ffi/container/tensor.h>
#include <algorithm> // for std::min
#include <cuda_runtime.h>
// Forward mode enum (must match Python ForwardMode in sglang/srt/layers/attention/nsa_backend.py)
enum ForwardModeEnum { DECODE = 0, TARGET_VERIFY = 1, DRAFT_EXTEND = 2 };
/**
* Source pointers for metadata copy operations.
* Groups all source tensor pointers for cleaner parameter passing.
* Some pointers may be nullptr depending on forward mode and feature flags.
*/
struct SourcePointers {
const int32_t* __restrict__ cache_seqlens; // [bs] sequence lengths in cache
const int32_t* __restrict__ cu_seqlens_k; // [bs+1] cumulative sequence lengths
const int32_t* __restrict__ page_indices; // page table indices
const int32_t* __restrict__ nsa_cache_seqlens; // NSA-specific cache lengths
const int32_t* __restrict__ seqlens_expanded; // expanded sequence lengths (TARGET_VERIFY/DRAFT_EXTEND only)
const int32_t* __restrict__ nsa_cu_seqlens_k; // NSA cumulative sequence lengths
const int32_t* __restrict__ real_page_table; // optional real page table
const int32_t* __restrict__ flashmla_num_splits; // optional FlashMLA split counts
const int32_t* __restrict__ flashmla_metadata; // optional FlashMLA metadata
};
/**
* Destination pointers for metadata copy operations.
* Groups all destination tensor pointers for cleaner parameter passing.
* Layout matches SourcePointers for consistency.
*/
struct DestinationPointers {
int32_t* __restrict__ cache_seqlens; // [bs] sequence lengths in cache
int32_t* __restrict__ cu_seqlens_k; // [bs+1] cumulative sequence lengths
int32_t* __restrict__ page_table_1; // page table (note: different name from source)
int32_t* __restrict__ nsa_cache_seqlens; // NSA-specific cache lengths
int32_t* __restrict__ seqlens_expanded; // expanded sequence lengths (TARGET_VERIFY/DRAFT_EXTEND only)
int32_t* __restrict__ nsa_cu_seqlens_k; // NSA cumulative sequence lengths
int32_t* __restrict__ real_page_table; // optional real page table
int32_t* __restrict__ flashmla_num_splits; // optional FlashMLA split counts
int32_t* __restrict__ flashmla_metadata; // optional FlashMLA metadata
};
/**
* Parameter structure for single-backend fused metadata copy kernel.
* Passed via __grid_constant__ for efficient constant memory access.
*/
struct FusedMetadataCopyParams {
SourcePointers src; // Source tensor pointers
DestinationPointers dst; // Destination tensor pointers
// Kernel parameters
int forward_mode; // 0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND
int bs; // Batch size
int max_len; // Max length for DECODE mode
int max_seqlen_k; // Max sequence length for TARGET_VERIFY/DRAFT_EXTEND
int seqlens_expanded_size; // Size of expanded sequence lengths
int page_indices_rows; // Number of rows in page_indices
int page_table_1_stride; // Stride for page_table_1
int real_page_table_cols; // Columns in real_page_table
int real_page_table_dst_stride; // Stride for destination real_page_table
int flashmla_metadata_size; // Size of FlashMLA metadata
};
/**
* Parameter structure for multi-backend fused metadata copy kernel.
* Enables copying from one source to three destinations in a single kernel launch.
* Used for speculative decoding with multiple draft backends.
*/
struct FusedMetadataCopyMultiParams {
SourcePointers src; // Source pointers (shared across all backends)
DestinationPointers dst0; // Backend 0 destination pointers
DestinationPointers dst1; // Backend 1 destination pointers
DestinationPointers dst2; // Backend 2 destination pointers
// Kernel parameters
int bs; // Batch size
int max_len; // Max length (DECODE mode only)
int seqlens_expanded_size; // Size of expanded sequence lengths
int page_table_1_stride; // Stride for page_table_1
int real_page_table_cols; // Columns in real_page_table
int real_page_table_dst_stride; // Stride for destination real_page_table
int flashmla_metadata_size; // Size of FlashMLA metadata
};
/**
* Unified kernel for all forward modes (DECODE, TARGET_VERIFY, DRAFT_EXTEND).
* Uses runtime branches for mode selection, with template parameters for
* compile-time optimization of optional features.
*
* DESIGN:
* - Runtime branches (forward_mode) handle mode-specific logic
* - Template parameters (HAS_*) eliminate unused feature code at compile time
* - Structured parameters (SourcePointers/DestinationPointers) passed via constant memory
*
* Used by FusedMetadataCopyKernel for single-backend metadata copy.
*
* @tparam HAS_REAL_PAGE_TABLE Compile-time flag for real_page_table support
* @tparam HAS_FLASHMLA Compile-time flag for FlashMLA metadata support
*/
template <bool HAS_REAL_PAGE_TABLE, bool HAS_FLASHMLA>
__global__ void fused_metadata_copy_kernel(const FusedMetadataCopyParams __grid_constant__ params) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int total_threads = gridDim.x * blockDim.x;
// Unpack parameters for readability
const auto& src = params.src;
const auto& dst = params.dst;
const int forward_mode = params.forward_mode;
const int bs = params.bs;
const int max_len = params.max_len;
const int max_seqlen_k = params.max_seqlen_k;
const int seqlens_expanded_size = params.seqlens_expanded_size;
const int page_indices_rows = params.page_indices_rows;
const int page_table_1_stride = params.page_table_1_stride;
const int real_page_table_cols = params.real_page_table_cols;
const int real_page_table_dst_stride = params.real_page_table_dst_stride;
const int flashmla_metadata_size = params.flashmla_metadata_size;
// Copy cache_seqlens (bs elements) - common to all modes
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
dst.cache_seqlens[i] = src.cache_seqlens[i];
}
// Copy cu_seqlens_k (skip first element) - common to all modes
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
dst.cu_seqlens_k[i + 1] = src.cu_seqlens_k[i + 1];
}
// Branch 1: page_table copy (different dimensions per mode)
if (forward_mode == 0) { // DECODE
int page_table_elements = bs * max_len;
#pragma unroll 4
for (int i = tid; i < page_table_elements; i += total_threads) {
int row = i / max_len;
int col = i % max_len;
dst.page_table_1[row * page_table_1_stride + col] = src.page_indices[i];
}
} else { // TARGET_VERIFY or DRAFT_EXTEND
int page_table_elements = page_indices_rows * max_seqlen_k;
#pragma unroll 4
for (int i = tid; i < page_table_elements; i += total_threads) {
int row = i / max_seqlen_k;
int col = i % max_seqlen_k;
dst.page_table_1[row * page_table_1_stride + col] = src.page_indices[i];
}
}
// Branch 2: seqlens_expanded copy (only for TARGET_VERIFY/DRAFT_EXTEND)
if (forward_mode != 0) { // TARGET_VERIFY or DRAFT_EXTEND
#pragma unroll 4
for (int i = tid; i < seqlens_expanded_size; i += total_threads) {
dst.seqlens_expanded[i] = src.seqlens_expanded[i];
}
}
// Branch 3: NSA metadata copy (different loop sizes per mode)
if (forward_mode == 0) { // DECODE
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
dst.nsa_cache_seqlens[i] = src.nsa_cache_seqlens[i];
}
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
dst.nsa_cu_seqlens_k[i + 1] = src.nsa_cu_seqlens_k[i + 1];
}
} else { // TARGET_VERIFY or DRAFT_EXTEND
#pragma unroll 4
for (int i = tid; i < seqlens_expanded_size; i += total_threads) {
dst.nsa_cache_seqlens[i] = src.nsa_cache_seqlens[i];
}
#pragma unroll 4
for (int i = tid; i < seqlens_expanded_size; i += total_threads) {
dst.nsa_cu_seqlens_k[i + 1] = src.nsa_cu_seqlens_k[i + 1];
}
}
// Copy real page table - compile-time branch
if constexpr (HAS_REAL_PAGE_TABLE) {
int real_table_elements = (forward_mode == 0 ? bs : page_indices_rows) * real_page_table_cols;
#pragma unroll 2
for (int i = tid; i < real_table_elements; i += total_threads) {
int row = i / real_page_table_cols;
int col = i % real_page_table_cols;
dst.real_page_table[row * real_page_table_dst_stride + col] =
src.real_page_table[row * real_page_table_cols + col];
}
}
// Branch 4: FlashMLA metadata copy (different sizes per mode)
if constexpr (HAS_FLASHMLA) {
int flashmla_size = (forward_mode == 0) ? (bs + 1) : (seqlens_expanded_size + 1);
if (forward_mode == 0) {
#pragma unroll 8
for (int i = tid; i < flashmla_size; i += total_threads) {
dst.flashmla_num_splits[i] = src.flashmla_num_splits[i];
}
} else {
#pragma unroll 4
for (int i = tid; i < flashmla_size; i += total_threads) {
dst.flashmla_num_splits[i] = src.flashmla_num_splits[i];
}
}
#pragma unroll 2
for (int i = tid; i < flashmla_metadata_size; i += total_threads) {
dst.flashmla_metadata[i] = src.flashmla_metadata[i];
}
}
}
/**
* Multi-backend kernel for DECODE mode.
* Copies from one source to THREE destinations in a single kernel launch.
*
* PERFORMANCE: 3x faster than three separate kernel launches due to:
* - Reduced kernel launch overhead (1 launch instead of 3)
* - Improved memory coalescing (source read once, written to 3 destinations)
* - Better instruction-level parallelism
*
* Used by FusedMetadataCopyMultiKernel for speculative decoding scenarios.
*
* @tparam HAS_REAL_PAGE_TABLE Compile-time flag for real_page_table support
* @tparam HAS_FLASHMLA Compile-time flag for FlashMLA metadata support
*/
template <bool HAS_REAL_PAGE_TABLE, bool HAS_FLASHMLA>
__global__ void fused_metadata_copy_multi_kernel(const FusedMetadataCopyMultiParams __grid_constant__ params) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int total_threads = gridDim.x * blockDim.x;
// Unpack parameters for readability
const auto& src = params.src;
const auto& dst0 = params.dst0;
const auto& dst1 = params.dst1;
const auto& dst2 = params.dst2;
const int bs = params.bs;
const int max_len = params.max_len;
const int seqlens_expanded_size = params.seqlens_expanded_size;
const int page_table_1_stride = params.page_table_1_stride;
const int real_page_table_cols = params.real_page_table_cols;
const int real_page_table_dst_stride = params.real_page_table_dst_stride;
const int flashmla_metadata_size = params.flashmla_metadata_size;
// Copy cache_seqlens to all 3 backends
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
int32_t val = src.cache_seqlens[i];
dst0.cache_seqlens[i] = val;
dst1.cache_seqlens[i] = val;
dst2.cache_seqlens[i] = val;
}
// Copy cu_seqlens_k to all 3 backends (skip first element)
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
int32_t val = src.cu_seqlens_k[i + 1];
dst0.cu_seqlens_k[i + 1] = val;
dst1.cu_seqlens_k[i + 1] = val;
dst2.cu_seqlens_k[i + 1] = val;
}
// DECODE mode: copy page_table_1 to all 3 backends
int page_table_elements = bs * max_len;
#pragma unroll 4
for (int i = tid; i < page_table_elements; i += total_threads) {
int row = i / max_len;
int col = i % max_len;
int32_t val = src.page_indices[i];
dst0.page_table_1[row * page_table_1_stride + col] = val;
dst1.page_table_1[row * page_table_1_stride + col] = val;
dst2.page_table_1[row * page_table_1_stride + col] = val;
}
// Copy nsa_cache_seqlens to all 3 backends
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
int32_t val = src.nsa_cache_seqlens[i];
dst0.nsa_cache_seqlens[i] = val;
dst1.nsa_cache_seqlens[i] = val;
dst2.nsa_cache_seqlens[i] = val;
}
// Copy NSA cu_seqlens to all 3 backends
#pragma unroll 8
for (int i = tid; i < bs; i += total_threads) {
int32_t val = src.nsa_cu_seqlens_k[i + 1];
dst0.nsa_cu_seqlens_k[i + 1] = val;
dst1.nsa_cu_seqlens_k[i + 1] = val;
dst2.nsa_cu_seqlens_k[i + 1] = val;
}
// Copy real page table to all 3 backends
if (src.real_page_table != nullptr && dst0.real_page_table != nullptr) {
int real_table_elements = bs * real_page_table_cols;
#pragma unroll 2
for (int i = tid; i < real_table_elements; i += total_threads) {
int row = i / real_page_table_cols;
int col = i % real_page_table_cols;
int src_idx = row * real_page_table_cols + col;
int dst_idx = row * real_page_table_dst_stride + col;
int32_t val = src.real_page_table[src_idx];
dst0.real_page_table[dst_idx] = val;
dst1.real_page_table[dst_idx] = val;
dst2.real_page_table[dst_idx] = val;
}
}
// Copy FlashMLA metadata to all 3 backends
if constexpr (HAS_FLASHMLA) {
int flashmla_size = bs + 1;
#pragma unroll 8
for (int i = tid; i < flashmla_size; i += total_threads) {
int32_t val = src.flashmla_num_splits[i];
dst0.flashmla_num_splits[i] = val;
dst1.flashmla_num_splits[i] = val;
dst2.flashmla_num_splits[i] = val;
}
#pragma unroll 2
for (int i = tid; i < flashmla_metadata_size; i += total_threads) {
int32_t val = src.flashmla_metadata[i];
dst0.flashmla_metadata[i] = val;
dst1.flashmla_metadata[i] = val;
dst2.flashmla_metadata[i] = val;
}
}
}
// ============================================================================
// Host-side launcher wrappers for JIT compilation
// ============================================================================
namespace {
// Launch configuration constants
constexpr int THREADS_PER_BLOCK = 256;
constexpr int MAX_GRID_SIZE = 1024; // Limit to prevent excessive resource usage
/**
* Helper function to extract a typed data pointer from a TensorView.
* Performs runtime type checking and returns the properly cast pointer.
*
* @tparam T The expected element type (e.g., int32_t)
* @param tensor The TensorView to extract the pointer from
* @param name The name of the tensor (for error reporting)
* @return Typed pointer to the tensor data
*/
template <typename T>
inline const T* unwrap_data_ptr(const tvm::ffi::TensorView& tensor, const char* name) {
using namespace host;
if (tensor.data_ptr()) {
RuntimeCheck(is_type<T>(tensor.dtype()), "Tensor ", name, " must have dtype int32");
}
return static_cast<const T*>(tensor.data_ptr());
}
/**
* Helper function to extract a typed mutable data pointer from a TensorView.
* Performs runtime type checking and returns the properly cast pointer.
*
* @tparam T The expected element type (e.g., int32_t)
* @param tensor The TensorView to extract the pointer from
* @param name The name of the tensor (for error reporting)
* @return Typed mutable pointer to the tensor data
*/
template <typename T>
inline T* unwrap_data_ptr_mut(const tvm::ffi::TensorView& tensor, const char* name) {
using namespace host;
if (tensor.data_ptr()) {
RuntimeCheck(is_type<T>(tensor.dtype()), "Tensor ", name, " must have dtype int32");
}
return static_cast<T*>(tensor.data_ptr());
}
/**
* Helper function to extract a typed data pointer from an Optional TensorView.
* Returns nullptr if the optional has no value, otherwise performs type checking.
*
* @tparam T The expected element type (e.g., int32_t)
* @param optional_tensor The Optional TensorView to extract the pointer from
* @param name The name of the tensor (for error reporting)
* @return Typed pointer to the tensor data, or nullptr if optional has no value
*/
template <typename T>
inline const T*
unwrap_optional_data_ptr(const tvm::ffi::Optional<tvm::ffi::TensorView>& optional_tensor, const char* name) {
using namespace host;
if (!optional_tensor.has_value()) {
return nullptr;
}
const auto& tensor = optional_tensor.value();
RuntimeCheck(is_type<T>(tensor.dtype()), "Tensor ", name, " must have dtype int32");
return static_cast<const T*>(tensor.data_ptr());
}
/**
* Helper function to extract a typed mutable data pointer from an Optional TensorView.
* Returns nullptr if the optional has no value, otherwise performs type checking.
*
* @tparam T The expected element type (e.g., int32_t)
* @param optional_tensor The Optional TensorView to extract the pointer from
* @param name The name of the tensor (for error reporting)
* @return Typed mutable pointer to the tensor data, or nullptr if optional has no value
*/
template <typename T>
inline T*
unwrap_optional_data_ptr_mut(const tvm::ffi::Optional<tvm::ffi::TensorView>& optional_tensor, const char* name) {
using namespace host;
if (!optional_tensor.has_value()) {
return nullptr;
}
const auto& tensor = optional_tensor.value();
RuntimeCheck(is_type<T>(tensor.dtype()), "Tensor ", name, " must have dtype int32");
return static_cast<T*>(tensor.data_ptr());
}
/**
* Calculate kernel launch configuration.
*
* @param total_work Total number of work items
* @param threads_per_block Threads per block (default: THREADS_PER_BLOCK)
* @return Grid dimension for kernel launch
*/
inline dim3 get_launch_config(int total_work, int threads_per_block = THREADS_PER_BLOCK) {
int num_blocks = (total_work + threads_per_block - 1) / threads_per_block;
// Limit grid size to prevent excessive resource usage while ensuring coverage
num_blocks = std::min(num_blocks, MAX_GRID_SIZE);
return dim3(num_blocks);
}
/**
* JIT wrapper for single-backend fused metadata copy kernel.
*
* This struct provides a unified interface for launching the fused metadata copy
* kernel with different forward modes. It constructs the parameter struct and
* launches the unified kernel.
*
* IMPLEMENTATION:
* - Extracts raw pointers from TensorView objects
* - Constructs FusedMetadataCopyParams with nested SourcePointers/DestinationPointers
* - Calculates grid configuration based on maximum work size
* - Launches fused_metadata_copy_kernel with __grid_constant__ parameters
*
* @tparam FORWARD_MODE Forward mode: 0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND
* @tparam HAS_REAL_PAGE_TABLE Whether real_page_table tensors are present
* @tparam HAS_FLASHMLA Whether FlashMLA metadata tensors are present
*/
template <int FORWARD_MODE, bool HAS_REAL_PAGE_TABLE, bool HAS_FLASHMLA>
struct FusedMetadataCopyKernel {
static_assert(
FORWARD_MODE >= 0 && FORWARD_MODE <= 2,
"FORWARD_MODE must be 0 (DECODE), 1 (TARGET_VERIFY), or 2 (DRAFT_EXTEND)");
static void
run(const tvm::ffi::TensorView cache_seqlens_src,
const tvm::ffi::TensorView cu_seqlens_k_src,
const tvm::ffi::TensorView page_indices_src,
const tvm::ffi::TensorView nsa_cache_seqlens_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> seqlens_expanded_src,
const tvm::ffi::TensorView nsa_cu_seqlens_k_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_src,
const tvm::ffi::TensorView cache_seqlens_dst,
const tvm::ffi::TensorView cu_seqlens_k_dst,
const tvm::ffi::TensorView page_table_1_dst,
const tvm::ffi::TensorView nsa_cache_seqlens_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> seqlens_expanded_dst,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst,
int bs,
int max_len,
int max_seqlen_k,
int seqlens_expanded_size) {
using namespace host;
// Build parameter struct with nested source/destination pointers
// unwrap_data_ptr and unwrap_optional_data_ptr perform dtype validation
const auto params = FusedMetadataCopyParams{
.src =
{
.cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"),
.cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"),
.page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"),
.nsa_cache_seqlens = unwrap_data_ptr<int32_t>(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"),
.seqlens_expanded = unwrap_optional_data_ptr<int32_t>(seqlens_expanded_src, "seqlens_expanded_src"),
.nsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"),
.real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"),
.flashmla_num_splits =
unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"),
.flashmla_metadata = unwrap_optional_data_ptr<int32_t>(flashmla_metadata_src, "flashmla_metadata_src"),
},
.dst =
{
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst, "cache_seqlens_dst"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst, "cu_seqlens_k_dst"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst, "page_table_1_dst"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst, "nsa_cache_seqlens_dst"),
.seqlens_expanded = unwrap_optional_data_ptr_mut<int32_t>(seqlens_expanded_dst, "seqlens_expanded_dst"),
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst, "nsa_cu_seqlens_k_dst"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst, "real_page_table_dst"),
.flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst, "flashmla_num_splits_dst"),
.flashmla_metadata =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_metadata_dst, "flashmla_metadata_dst"),
},
.forward_mode = FORWARD_MODE,
.bs = bs,
.max_len = max_len,
.max_seqlen_k = max_seqlen_k,
.seqlens_expanded_size = seqlens_expanded_size,
.page_indices_rows = static_cast<int>(page_indices_src.shape()[0]),
.page_table_1_stride = static_cast<int>(page_table_1_dst.shape()[1]),
.real_page_table_cols =
real_page_table_src.has_value() ? static_cast<int>(real_page_table_src.value().shape()[1]) : 0,
.real_page_table_dst_stride =
real_page_table_dst.has_value() ? static_cast<int>(real_page_table_dst.value().stride(0)) : 0,
.flashmla_metadata_size =
flashmla_metadata_src.has_value() ? static_cast<int>(flashmla_metadata_src.value().numel()) : 0,
};
// Calculate grid configuration
int max_elements = std::max(
{bs,
params.page_indices_rows * max_seqlen_k,
seqlens_expanded_size,
HAS_FLASHMLA ? (seqlens_expanded_size + 1) : 0,
HAS_FLASHMLA ? params.flashmla_metadata_size : 0});
dim3 grid = get_launch_config(max_elements);
dim3 block(THREADS_PER_BLOCK);
DLDevice device = cache_seqlens_src.device();
// Launch unified kernel with params struct
host::LaunchKernel(grid, block, device)(fused_metadata_copy_kernel<HAS_REAL_PAGE_TABLE, HAS_FLASHMLA>, params);
}
};
/**
* JIT wrapper for multi-backend fused metadata copy kernel.
*
* This kernel optimizes the common case where metadata needs to be copied from
* one source to THREE destination backends in a single kernel launch. This is
* 3x faster than launching three separate kernels due to:
* - Reduced kernel launch overhead (1 launch instead of 3)
* - Improved memory coalescing (source read once, written to 3 destinations)
* - Better GPU occupancy and instruction-level parallelism
*
* USAGE: Primarily for speculative decoding with multiple draft models, where
* the same source metadata needs to be replicated to multiple backend contexts.
*
* LIMITATION: Currently only supports DECODE mode, which is the most frequently
* used mode in speculative decoding scenarios.
*
* IMPLEMENTATION:
* - Constructs FusedMetadataCopyMultiParams with 1 SourcePointers + 3 DestinationPointers
* - Launches fused_metadata_copy_multi_kernel with __grid_constant__ parameters
*
* @tparam HAS_REAL_PAGE_TABLE Whether real_page_table tensors are present
* @tparam HAS_FLASHMLA Whether FlashMLA metadata tensors are present
*/
template <bool HAS_REAL_PAGE_TABLE, bool HAS_FLASHMLA>
struct FusedMetadataCopyMultiKernel {
static void
run(const tvm::ffi::TensorView cache_seqlens_src,
const tvm::ffi::TensorView cu_seqlens_k_src,
const tvm::ffi::TensorView page_indices_src,
const tvm::ffi::TensorView nsa_cache_seqlens_src,
const tvm::ffi::TensorView nsa_cu_seqlens_k_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_src,
const tvm::ffi::TensorView cache_seqlens_dst0,
const tvm::ffi::TensorView cu_seqlens_k_dst0,
const tvm::ffi::TensorView page_table_1_dst0,
const tvm::ffi::TensorView nsa_cache_seqlens_dst0,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst0,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst0,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst0,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst0,
const tvm::ffi::TensorView cache_seqlens_dst1,
const tvm::ffi::TensorView cu_seqlens_k_dst1,
const tvm::ffi::TensorView page_table_1_dst1,
const tvm::ffi::TensorView nsa_cache_seqlens_dst1,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst1,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst1,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst1,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst1,
const tvm::ffi::TensorView cache_seqlens_dst2,
const tvm::ffi::TensorView cu_seqlens_k_dst2,
const tvm::ffi::TensorView page_table_1_dst2,
const tvm::ffi::TensorView nsa_cache_seqlens_dst2,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst2,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst2,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst2,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst2,
int bs,
int max_len,
int seqlens_expanded_size) {
using namespace host;
// Build parameter struct with nested source/destination pointers
// unwrap_data_ptr and unwrap_optional_data_ptr perform dtype validation
const auto params = FusedMetadataCopyMultiParams{
.src =
{
.cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"),
.cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"),
.page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"),
.nsa_cache_seqlens = unwrap_data_ptr<int32_t>(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"),
.seqlens_expanded = nullptr, // Not used in multi-backend DECODE mode
.nsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"),
.real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"),
.flashmla_num_splits =
unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"),
.flashmla_metadata = unwrap_optional_data_ptr<int32_t>(flashmla_metadata_src, "flashmla_metadata_src"),
},
.dst0 =
{
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst0, "cache_seqlens_dst0"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst0, "cu_seqlens_k_dst0"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst0, "page_table_1_dst0"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst0, "nsa_cache_seqlens_dst0"),
.seqlens_expanded = nullptr,
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst0, "nsa_cu_seqlens_k_dst0"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst0, "real_page_table_dst0"),
.flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst0, "flashmla_num_splits_dst0"),
.flashmla_metadata =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_metadata_dst0, "flashmla_metadata_dst0"),
},
.dst1 =
{
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst1, "cache_seqlens_dst1"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst1, "cu_seqlens_k_dst1"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst1, "page_table_1_dst1"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst1, "nsa_cache_seqlens_dst1"),
.seqlens_expanded = nullptr,
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst1, "nsa_cu_seqlens_k_dst1"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst1, "real_page_table_dst1"),
.flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst1, "flashmla_num_splits_dst1"),
.flashmla_metadata =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_metadata_dst1, "flashmla_metadata_dst1"),
},
.dst2 =
{
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst2, "cache_seqlens_dst2"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst2, "cu_seqlens_k_dst2"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst2, "page_table_1_dst2"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst2, "nsa_cache_seqlens_dst2"),
.seqlens_expanded = nullptr,
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst2, "nsa_cu_seqlens_k_dst2"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst2, "real_page_table_dst2"),
.flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst2, "flashmla_num_splits_dst2"),
.flashmla_metadata =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_metadata_dst2, "flashmla_metadata_dst2"),
},
.bs = bs,
.max_len = max_len,
.seqlens_expanded_size = seqlens_expanded_size,
.page_table_1_stride = static_cast<int>(page_table_1_dst0.shape()[1]),
.real_page_table_cols =
real_page_table_src.has_value() ? static_cast<int>(real_page_table_src.value().shape()[1]) : 0,
.real_page_table_dst_stride =
real_page_table_dst0.has_value() ? static_cast<int>(real_page_table_dst0.value().stride(0)) : 0,
.flashmla_metadata_size =
flashmla_metadata_src.has_value() ? static_cast<int>(flashmla_metadata_src.value().numel()) : 0,
};
dim3 grid = get_launch_config(bs * max_len);
dim3 block(THREADS_PER_BLOCK);
DLDevice device = cache_seqlens_src.device();
// Launch multi-backend kernel with params struct
host::LaunchKernel(grid, block, device)(
fused_metadata_copy_multi_kernel<HAS_REAL_PAGE_TABLE, HAS_FLASHMLA>, params);
}
};
} // namespace

346
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/fused_qknorm_rope.cuh vendored Normal file
View File

@@ -0,0 +1,346 @@
/*
* Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// Adapted from
// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <tvm/ffi/container/tensor.h>
#include <cmath>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
namespace {
// ---------------------------------------------------------------------------
// YaRN-aware frequency computation
//
// When factor == 1.0, reduces to standard RoPE: base^(-2*half_dim/rotary_dim)
// When factor != 1.0, blends interpolated and extrapolated frequencies.
// ---------------------------------------------------------------------------
template <bool yarn>
__device__ inline float compute_freq(float base, int rotary_dim, int half_dim, float factor, float low, float high) {
float freq = powf(base, -2.0f * half_dim / static_cast<float>(rotary_dim));
if constexpr (yarn) {
float inv_freq_extrapolation = freq;
float inv_freq_interpolation = freq / factor;
float high_adj = high;
if (fabsf(low - high_adj) <= 1e-6f) {
high_adj += 0.001f;
}
float linear_func = (static_cast<float>(half_dim) - low) / (high_adj - low);
float ramp_func = fminf(fmaxf(linear_func, 0.0f), 1.0f);
float inv_freq_extrapolation_factor = 1.0f - ramp_func;
freq = inv_freq_interpolation * (1.0f - inv_freq_extrapolation_factor) +
inv_freq_extrapolation * inv_freq_extrapolation_factor;
}
return freq;
}
// ---------------------------------------------------------------------------
// Fused QK-Norm + RoPE kernel
//
// Each warp processes one (token, head) pair.
// head_dim: compile-time head dimension (64, 128, or 256)
// interleave: true → interleave / GPT-J style RoPE (!is_neox)
// false → NeoX style RoPE (is_neox)
// ---------------------------------------------------------------------------
// interleave (GPT-J) pairs (2k,2k+1) share the same freq/theta,
// so sin/cos is computed once per pair and copied to the odd element,
// halving powf + __sincosf calls vs a naive per-element approach.
template <int head_dim, bool interleave, bool yarn>
__global__ void fusedQKNormRopeKernel(
__nv_bfloat16* qkv, // [num_tokens, (nq+nk+nv)*head_dim], in-place
int const num_heads_q,
int const num_heads_k,
int const num_heads_v,
float const eps,
__nv_bfloat16 const* q_weight, // [head_dim]
__nv_bfloat16 const* k_weight, // [head_dim]
float const base,
int const* position_ids, // [num_tokens]
int const num_tokens,
float factor,
float low,
float high,
float attention_factor,
int const rotary_dim) {
int const warpsPerBlock = blockDim.x / 32;
int const warpId = threadIdx.x / 32;
int const laneId = threadIdx.x % 32;
int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
int const total_qk_heads = num_heads_q + num_heads_k;
int const tokenIdx = globalWarpIdx / total_qk_heads;
int const localHeadIdx = globalWarpIdx % total_qk_heads;
if (tokenIdx >= num_tokens) return;
bool const isQ = localHeadIdx < num_heads_q;
int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q;
int const num_heads = num_heads_q + num_heads_k + num_heads_v;
static_assert(head_dim % (32 * 2) == 0, "head_dim must be divisible by 64 (each warp handles one head)");
constexpr int numElemsPerThread = head_dim / 32;
float elements[numElemsPerThread];
using vec_T = device::AlignedVector<bf16_t, numElemsPerThread>;
// Compute flat offset of this warp's head in qkv
int offsetWarp;
if (isQ) {
offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim;
} else {
offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim + headIdx * head_dim;
}
int offsetThread = offsetWarp + laneId * numElemsPerThread;
// -------------------------------------------------------------------
// Load and compute sum-of-squares for RMSNorm
// -------------------------------------------------------------------
float sumOfSquares = 0.0f;
{
vec_T vec;
vec.load(qkv + offsetThread);
for (int i = 0; i < numElemsPerThread; i++) {
float val = device::cast<float>(vec[i]);
sumOfSquares += val * val;
elements[i] = val;
}
}
sumOfSquares = device::warp::reduce_sum(sumOfSquares);
// -------------------------------------------------------------------
// Apply RMSNorm
// -------------------------------------------------------------------
float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps);
{
vec_T wvec;
wvec.load((isQ ? q_weight : k_weight) + offsetThread - offsetWarp);
for (int i = 0; i < numElemsPerThread; i++) {
elements[i] *= rms_rcp * device::cast<float>(wvec[i]);
}
}
// -------------------------------------------------------------------
// Apply RoPE to the first rotary_dim elements
// -------------------------------------------------------------------
float pos_id = static_cast<float>(position_ids[tokenIdx]);
int const rotary_lanes = rotary_dim / numElemsPerThread;
bool const applyRotary = (laneId < rotary_lanes);
if (applyRotary) {
if constexpr (interleave) {
// Pairs (2k, 2k+1) share the same half_dim → same freq/theta.
// numElemsPerThread is always even (head_dim/32, head_dim in {64,128,256}),
// so we step by 2 and handle both elements of each pair per iteration.
//
// freq follows a geometric series across pairs: freq[k] = freq[0] * ratio^k,
// where ratio = base^(-2/rotary_dim). Pre-compute both outside the loop to
// replace all but the first powf call with a single multiply per iteration.
//
// sin/cos are applied immediately to e0/e1, eliminating the elements2,
// cos_vals, sin_vals intermediate arrays and reducing register pressure.
int const half_dim_start = laneId * numElemsPerThread / 2;
float freq = powf(base, -2.0f * static_cast<float>(half_dim_start) / static_cast<float>(rotary_dim));
float const freq_ratio = powf(base, -2.0f / static_cast<float>(rotary_dim));
for (int i = 0; i < numElemsPerThread; i += 2) {
float e0 = elements[i];
float e1 = elements[i + 1];
float f = freq;
if constexpr (yarn) {
int half_dim = half_dim_start + i / 2;
float inv_freq_interpolation = freq / factor;
float high_adj = (fabsf(low - high) <= 1e-6f) ? high + 0.001f : high;
float linear_func = (static_cast<float>(half_dim) - low) / (high_adj - low);
float ramp_func = fminf(fmaxf(linear_func, 0.0f), 1.0f);
float extrap_factor = 1.0f - ramp_func;
f = inv_freq_interpolation * (1.0f - extrap_factor) + freq * extrap_factor;
}
float s, c;
__sincosf(pos_id * f, &s, &c);
elements[i] = (e0 * c - e1 * s) * attention_factor;
elements[i + 1] = (e1 * c + e0 * s) * attention_factor;
freq *= freq_ratio;
}
} else {
// NeoX style: first and second halves of the rotary region are paired
float elements2[numElemsPerThread];
float cos_vals[numElemsPerThread];
float sin_vals[numElemsPerThread];
__syncwarp();
int const half_rotary_lanes = rotary_lanes / 2;
// Avoid UB from (1u << 32) when rotary_lanes == 32
unsigned int active_mask = 0xffffffffu >> (32 - rotary_lanes);
for (int i = 0; i < numElemsPerThread; i++) {
elements2[i] = __shfl_xor_sync(active_mask, elements[i], half_rotary_lanes);
if (laneId < half_rotary_lanes) {
elements2[i] = -elements2[i];
}
int dim_idx = laneId * numElemsPerThread + i;
// Remap so that both halves use the same set of frequencies
dim_idx = (dim_idx * 2) % rotary_dim;
int half_dim = dim_idx / 2;
float freq = compute_freq<yarn>(base, rotary_dim, half_dim, factor, low, high);
float theta = pos_id * freq;
__sincosf(theta, &sin_vals[i], &cos_vals[i]);
}
__syncwarp();
for (int i = 0; i < numElemsPerThread; i++) {
elements[i] = (elements[i] * cos_vals[i] + elements2[i] * sin_vals[i]) * attention_factor;
}
}
}
// -------------------------------------------------------------------
// Store (all elements: rotated + pass-through normalized)
// -------------------------------------------------------------------
{
vec_T vec;
for (int i = 0; i < numElemsPerThread; i++) {
vec[i] = device::cast<bf16_t>(elements[i]);
}
vec.store(qkv + offsetThread);
}
}
// ---------------------------------------------------------------------------
// Host-side tvm-ffi entry point
// ---------------------------------------------------------------------------
void fused_qk_norm_rope(
tvm::ffi::TensorView qkv, // [num_tokens, (nq+nk+nv)*head_dim] bf16
tvm::ffi::TensorView q_weight, // [head_dim] bf16
tvm::ffi::TensorView k_weight, // [head_dim] bf16
tvm::ffi::TensorView position_ids, // [num_tokens] int32
int num_heads_q,
int num_heads_k,
int num_heads_v,
int head_dim,
float eps,
float base,
int is_neox, // 0 = interleave style, 1 = NeoX style
float factor,
float low,
float high,
float attention_factor,
int rotary_dim) {
using namespace host;
RuntimeCheck(qkv.device().device_type == kDLCUDA, "qkv must be a CUDA tensor");
RuntimeCheck(qkv.is_contiguous(), "qkv must be contiguous");
RuntimeCheck(qkv.dtype().code == kDLBfloat && qkv.dtype().bits == 16, "qkv must be bfloat16");
RuntimeCheck(qkv.ndim() == 2, "qkv must be 2D: [num_tokens, (nq+nk+nv)*head_dim]");
RuntimeCheck(q_weight.is_contiguous(), "q_weight must be contiguous");
RuntimeCheck(q_weight.dtype().code == kDLBfloat && q_weight.dtype().bits == 16, "q_weight must be bfloat16");
RuntimeCheck(
q_weight.ndim() == 1 && static_cast<int>(q_weight.size(0)) == head_dim, "q_weight must be 1D of size head_dim");
RuntimeCheck(k_weight.is_contiguous(), "k_weight must be contiguous");
RuntimeCheck(k_weight.dtype().code == kDLBfloat && k_weight.dtype().bits == 16, "k_weight must be bfloat16");
RuntimeCheck(
k_weight.ndim() == 1 && static_cast<int>(k_weight.size(0)) == head_dim, "k_weight must be 1D of size head_dim");
RuntimeCheck(position_ids.device().device_type == kDLCUDA, "position_ids must be a CUDA tensor");
RuntimeCheck(position_ids.is_contiguous(), "position_ids must be contiguous");
RuntimeCheck(position_ids.dtype().code == kDLInt && position_ids.dtype().bits == 32, "position_ids must be int32");
RuntimeCheck(position_ids.ndim() == 1, "position_ids must be 1D: [num_tokens]");
int num_tokens = static_cast<int>(qkv.size(0));
int total_heads = num_heads_q + num_heads_k + num_heads_v;
RuntimeCheck(
static_cast<int>(qkv.size(1)) == total_heads * head_dim, "qkv.size(1) must equal (nq + nk + nv) * head_dim");
RuntimeCheck(static_cast<int>(position_ids.size(0)) == num_tokens, "position_ids must have num_tokens elements");
static_assert(
JIT_HEAD_DIM == 64 || JIT_HEAD_DIM == 128 || JIT_HEAD_DIM == 256, "JIT_HEAD_DIM must be 64, 128, or 256");
static_assert(JIT_INTERLEAVE == 0 || JIT_INTERLEAVE == 1, "JIT_INTERLEAVE must be 0 or 1");
static_assert(JIT_YARN == 0 || JIT_YARN == 1, "JIT_YARN must be 0 or 1");
RuntimeCheck(head_dim == JIT_HEAD_DIM, "head_dim mismatch with JIT-compiled kernel");
int numElemsPerThread = head_dim / 32;
RuntimeCheck(rotary_dim % numElemsPerThread == 0, "rotary_dim must be divisible by (head_dim / 32)");
bool neox = static_cast<bool>(is_neox);
if (neox) {
// NeoX uses __shfl_xor_sync which requires half_rotary_lanes to be a power of 2
int rotary_lanes = rotary_dim / numElemsPerThread;
int half_rotary_lanes = rotary_lanes / 2;
bool is_pow2 = (half_rotary_lanes >= 1) && ((half_rotary_lanes & (half_rotary_lanes - 1)) == 0);
RuntimeCheck(is_pow2, "half_rotary_lanes must be a power of 2 for NeoX style RoPE");
}
bool interleave = !neox;
RuntimeCheck(interleave == static_cast<bool>(JIT_INTERLEAVE), "interleave mismatch with JIT-compiled kernel");
bool use_yarn = (factor != 1.0f);
RuntimeCheck(use_yarn == static_cast<bool>(JIT_YARN), "yarn mismatch with JIT-compiled kernel");
cudaStream_t stream = LaunchKernel::resolve_device(qkv.device());
constexpr int blockSize = 256;
int warpsPerBlock = blockSize / 32;
int totalQKHeads = num_heads_q + num_heads_k;
int totalWarps = num_tokens * totalQKHeads;
int gridSize = div_ceil(totalWarps, warpsPerBlock);
auto* qkv_ptr = reinterpret_cast<__nv_bfloat16*>(qkv.data_ptr());
auto const* qw_ptr = reinterpret_cast<__nv_bfloat16 const*>(q_weight.data_ptr());
auto const* kw_ptr = reinterpret_cast<__nv_bfloat16 const*>(k_weight.data_ptr());
auto const* pos_ptr = reinterpret_cast<int const*>(position_ids.data_ptr());
fusedQKNormRopeKernel<JIT_HEAD_DIM, static_cast<bool>(JIT_INTERLEAVE), static_cast<bool>(JIT_YARN)>
<<<gridSize, blockSize, 0, stream>>>(
qkv_ptr,
num_heads_q,
num_heads_k,
num_heads_v,
eps,
qw_ptr,
kw_ptr,
base,
pos_ptr,
num_tokens,
factor,
low,
high,
attention_factor,
rotary_dim);
}
} // namespace

201
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh vendored Normal file
View File

@@ -0,0 +1,201 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <cstdint>
namespace {
struct StoreKVCacheParams {
const void* __restrict__ k;
const void* __restrict__ v;
void* __restrict__ k_cache;
void* __restrict__ v_cache;
const void* __restrict__ indices;
int64_t stride_k_bytes;
int64_t stride_v_bytes;
int64_t stride_cache_bytes;
int64_t stride_indices;
uint32_t batch_size;
};
constexpr uint32_t kNumWarps = 4;
constexpr uint32_t kThreadsPerBlock = kNumWarps * device::kWarpThreads;
/**
* \brief Use a single warp to copy key and value data from source to destination.
* Each thread in the warp copies a portion of the data in a coalesced manner.
* \tparam kElementBytes The size of each key/value element in bytes.
* \param k_src Pointer to the source key data.
* \param v_src Pointer to the source value data.
* \param k_dst Pointer to the destination key data.
* \param v_dst Pointer to the destination value data.
*/
template <int64_t kElementBytes>
SGL_DEVICE void copy_kv_warp(
const void* __restrict__ k_src,
const void* __restrict__ v_src,
void* __restrict__ k_dst,
void* __restrict__ v_dst) {
using namespace device;
constexpr int64_t kAlignment = (kElementBytes % (16 * kWarpThreads) == 0) ? 16
: kElementBytes % (8 * kWarpThreads) == 0 ? 8
: kElementBytes % (4 * kWarpThreads) == 0 ? 4
: kElementBytes % 4 == 0 ? 4
: 0;
static_assert(kAlignment > 0, "Element size must be multiple of 4 bytes");
using vec_t = AlignedStorage<uint32_t, kAlignment / 4>;
constexpr auto kLoopBytes = sizeof(vec_t) * kWarpThreads;
constexpr auto kLoopCount = kElementBytes / kLoopBytes;
const auto gmem = tile::Memory<vec_t>::warp();
#pragma unroll kLoopCount
for (int64_t i = 0; i < kLoopCount; ++i) {
const auto k = gmem.load(k_src, i);
const auto v = gmem.load(v_src, i);
gmem.store(k_dst, k, i);
gmem.store(v_dst, v, i);
}
// handle the epilogue if any
if constexpr (kLoopCount * kLoopBytes < kElementBytes) {
if (gmem.in_bound(kElementBytes / sizeof(vec_t), kLoopCount)) {
const auto k = gmem.load(k_src, kLoopCount);
const auto v = gmem.load(v_src, kLoopCount);
gmem.store(k_dst, k, kLoopCount);
gmem.store(v_dst, v, kLoopCount);
}
}
}
/**
* \brief Kernel to store key-value pairs into the KV cache.
* Each element is split into multiple parts to allow parallel memory copy.
* \tparam kElementBytes The size of each key/value element in bytes.
* \tparam kSplit The number of warps that handle each element.
* \tparam kUsePDL Whether to use PDL feature.
* \tparam T The data type of the indices (`int32_t` or `int64_t`).
*/
template <int64_t kElementBytes, int kSplit, bool kUsePDL, typename T>
__global__ void store_kvcache(const __grid_constant__ StoreKVCacheParams params) {
using namespace device;
constexpr auto kSplitSize = kElementBytes / kSplit;
const uint32_t warp_id = blockIdx.x * kNumWarps + threadIdx.x / kWarpThreads;
const uint32_t item_id = warp_id / kSplit;
const uint32_t split_id = warp_id % kSplit;
const auto& [
k_input, v_input, k_cache, v_cache, indices, // ptr
stride_k, stride_v, stride_cache, stride_indices, batch_size // size
] = params;
if (item_id >= batch_size) return;
const auto index_ptr = static_cast<const T*>(indices) + item_id * stride_indices;
PDLWaitPrimary<kUsePDL>();
const auto index = *index_ptr;
const auto k_src = pointer::offset(k_input, item_id * stride_k, split_id * kSplitSize);
const auto v_src = pointer::offset(v_input, item_id * stride_v, split_id * kSplitSize);
const auto k_dst = pointer::offset(k_cache, index * stride_cache, split_id * kSplitSize);
const auto v_dst = pointer::offset(v_cache, index * stride_cache, split_id * kSplitSize);
copy_kv_warp<kSplitSize>(k_src, v_src, k_dst, v_dst);
PDLTriggerSecondary<kUsePDL>();
}
template <int64_t kElementBytes, bool kUsePDL>
struct StoreKVCacheKernel {
static_assert(kElementBytes > 0 && kElementBytes % 4 == 0);
template <int kSplit, typename T>
static constexpr auto store_kernel = store_kvcache<kElementBytes, kSplit, kUsePDL, T>;
template <typename T>
static auto get_kernel(const int num_split) {
using namespace host;
// only apply split optimization when element size is aligned
if constexpr (kElementBytes % (4 * 128) == 0) {
if (num_split == 4) return store_kernel<4, T>;
}
if constexpr (kElementBytes % (2 * 128) == 0) {
if (num_split == 2) return store_kernel<2, T>;
}
if (num_split == 1) return store_kernel<1, T>;
Panic("Unsupported num_split {} for element size {}", num_split, kElementBytes);
}
static void
run(const tvm::ffi::TensorView k,
const tvm::ffi::TensorView v,
const tvm::ffi::TensorView k_cache,
const tvm::ffi::TensorView v_cache,
const tvm::ffi::TensorView indices,
const int num_split) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto D = SymbolicSize{"element_size"};
auto KS = SymbolicSize{"k_stride"};
auto VS = SymbolicSize{"v_stride"};
auto S = SymbolicSize{"cache_stride"};
auto I = SymbolicSize{"indices_stride"};
auto dtype = SymbolicDType{};
auto device = SymbolicDevice{};
auto indice_dtype = SymbolicDType{};
device.set_options<kDLCUDA, kDLROCM>();
TensorMatcher({B, D}) //
.with_strides({KS, 1})
.with_dtype(dtype)
.with_device(device)
.verify(k);
TensorMatcher({B, D}) //
.with_strides({VS, 1})
.with_dtype(dtype)
.with_device(device)
.verify(v);
TensorMatcher({-1, D}) //
.with_strides({S, 1})
.with_dtype(dtype)
.with_device(device)
.verify(k_cache)
.verify(v_cache);
TensorMatcher({B}) //
.with_strides({I})
.with_dtype<int32_t, int64_t>(indice_dtype)
.with_device(device)
.verify(indices);
const int64_t dtype_size = dtype_bytes(dtype.unwrap());
const uint32_t num_elements = static_cast<uint32_t>(B.unwrap());
RuntimeCheck(kElementBytes == dtype_size * D.unwrap());
const auto params = StoreKVCacheParams{
.k = k.data_ptr(),
.v = v.data_ptr(),
.k_cache = k_cache.data_ptr(),
.v_cache = v_cache.data_ptr(),
.indices = indices.data_ptr(),
.stride_k_bytes = KS.unwrap() * dtype_size,
.stride_v_bytes = VS.unwrap() * dtype_size,
.stride_cache_bytes = S.unwrap() * dtype_size,
.stride_indices = I.unwrap(),
.batch_size = static_cast<uint32_t>(B.unwrap()),
};
// select kernel and update num_split if needed
const auto use_int32 = indice_dtype.is_type<int32_t>();
const auto kernel = use_int32 ? get_kernel<int32_t>(num_split) : get_kernel<int64_t>(num_split);
const auto num_blocks = div_ceil(num_elements * num_split, kNumWarps);
LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace

313
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/pos_enc.cuh vendored Normal file
View File

@@ -0,0 +1,313 @@
// Adapted from
// https://github.com/vllm-project/vllm/blob/014ece97c7aa49084a1119dca792af081a18dbc1/csrc/pos_encoding_kernels.cu
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <tvm/ffi/container/tensor.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
namespace {
template <typename scalar_t, bool IS_NEOX>
inline __device__ void apply_token_rotary_embedding(
scalar_t* __restrict__ arr,
const scalar_t* __restrict__ cos_ptr,
const scalar_t* __restrict__ sin_ptr,
int rot_offset,
int embed_dim) {
int x_index, y_index;
scalar_t cos, sin;
if (IS_NEOX) {
// GPT-NeoX style rotary embedding.
x_index = rot_offset;
y_index = embed_dim + rot_offset;
cos = SGLANG_LDG(cos_ptr + x_index);
sin = SGLANG_LDG(sin_ptr + x_index);
} else {
// GPT-J style rotary embedding.
x_index = 2 * rot_offset;
y_index = 2 * rot_offset + 1;
cos = SGLANG_LDG(cos_ptr + x_index / 2);
sin = SGLANG_LDG(sin_ptr + x_index / 2);
}
const scalar_t x = arr[x_index];
const scalar_t y = arr[y_index];
arr[x_index] = x * cos - y * sin;
arr[y_index] = y * cos + x * sin;
}
template <typename scalar_t, bool IS_NEOX>
inline __device__ void apply_rotary_embedding(
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
// head_size] or [num_tokens, num_heads,
// head_size]
scalar_t* __restrict__ key, // nullptr or
// [batch_size, seq_len, num_kv_heads,
// head_size] or [num_tokens, num_kv_heads,
// head_size]
const scalar_t* cache_ptr,
const int head_size,
const int num_heads,
const int num_kv_heads,
const int rot_dim,
const int token_idx,
const int64_t query_stride,
const int64_t key_stride,
const int64_t head_stride) {
const int embed_dim = rot_dim / 2;
const scalar_t* cos_ptr = cache_ptr;
const scalar_t* sin_ptr = cache_ptr + embed_dim;
const int nq = num_heads * embed_dim;
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
const int head_idx = i / embed_dim;
const int64_t token_head = token_idx * query_stride + head_idx * head_stride;
const int rot_offset = i % embed_dim;
apply_token_rotary_embedding<scalar_t, IS_NEOX>(query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
}
if (key != nullptr) {
const int nk = num_kv_heads * embed_dim;
for (int i = threadIdx.x; i < nk; i += blockDim.x) {
const int head_idx = i / embed_dim;
const int64_t token_head = token_idx * key_stride + head_idx * head_stride;
const int rot_offset = i % embed_dim;
apply_token_rotary_embedding<scalar_t, IS_NEOX>(key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
}
}
}
template <typename scalar_t, bool IS_NEOX>
__global__ void rotary_embedding_kernel(
const int64_t* __restrict__ positions, // [batch_size, seq_len] or
// [num_tokens]
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
// head_size] or [num_tokens, num_heads,
// head_size]
scalar_t* __restrict__ key, // nullptr or
// [batch_size, seq_len, num_kv_heads,
// head_size] or [num_tokens, num_kv_heads,
// head_size]
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //
// 2]
const int rot_dim,
const int64_t query_stride,
const int64_t key_stride,
const int64_t head_stride,
const int num_heads,
const int num_kv_heads,
const int head_size) {
// Each thread block is responsible for one token.
const int token_idx = blockIdx.x;
int64_t pos = positions[token_idx];
const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
apply_rotary_embedding<scalar_t, IS_NEOX>(
query,
key,
cache_ptr,
head_size,
num_heads,
num_kv_heads,
rot_dim,
token_idx,
query_stride,
key_stride,
head_stride);
}
// Helper struct to launch kernel
template <typename scalar_t, bool IS_NEOX>
void launch_kernel(
const int64_t* positions_data_ptr,
void* query_ptr,
void* key_ptr,
const void* cos_sin_cache_ptr,
int rot_dim,
int64_t query_stride,
int64_t key_stride,
int64_t head_stride,
int num_heads,
int num_kv_heads,
int head_size,
dim3 grid,
dim3 block,
const cudaStream_t stream) {
rotary_embedding_kernel<scalar_t, IS_NEOX><<<grid, block, 0, stream>>>(
positions_data_ptr,
static_cast<scalar_t*>(query_ptr),
static_cast<scalar_t*>(key_ptr),
static_cast<const scalar_t*>(cos_sin_cache_ptr),
rot_dim,
query_stride,
key_stride,
head_stride,
num_heads,
num_kv_heads,
head_size);
};
// Helper macro to reduce repetition
#define DISPATCH_DTYPE(DTYPE_CODE, DTYPE_BITS, IS_NEOX, ...) \
if (DTYPE_CODE == kDLFloat && DTYPE_BITS == 32) { \
launch_kernel<float, IS_NEOX>(__VA_ARGS__); \
} else if (DTYPE_CODE == kDLFloat && DTYPE_BITS == 16) { \
launch_kernel<half, IS_NEOX>(__VA_ARGS__); \
} else if (DTYPE_CODE == kDLBfloat && DTYPE_BITS == 16) { \
launch_kernel<nv_bfloat16, IS_NEOX>(__VA_ARGS__); \
} else { \
RuntimeCheck( \
false, "Unsupported data type for rotary embedding. Only float32, float16, and bfloat16 are supported."); \
}
// Helper function to dispatch based on data type
template <bool IS_NEOX>
void dispatch_by_dtype(
const int64_t* positions_data_ptr,
DLDataType query_dtype,
void* query_ptr,
void* key_ptr,
void* cos_sin_cache_ptr,
int rot_dim,
int64_t query_stride,
int64_t key_stride,
int64_t head_stride,
int num_heads,
int num_kv_heads,
int head_size,
dim3 grid,
dim3 block,
const cudaStream_t stream) {
using namespace host;
DISPATCH_DTYPE(
query_dtype.code,
query_dtype.bits,
IS_NEOX,
positions_data_ptr,
query_ptr,
key_ptr,
cos_sin_cache_ptr,
rot_dim,
query_stride,
key_stride,
head_stride,
num_heads,
num_kv_heads,
head_size,
grid,
block,
stream);
}
struct RotaryEmbeddingKernel {
static void
run(tvm::ffi::TensorView positions, // [batch_size, seq_len] or [num_tokens]
tvm::ffi::TensorView query, // [batch_size, seq_len, num_heads * head_size] or
// [num_tokens, num_heads * head_size] or
// [batch_size, seq_len, num_heads, head_size] or
// [num_tokens, num_heads, head_size]
tvm::ffi::Optional<tvm::ffi::TensorView> key,
// null or
// [batch_size, seq_len, num_kv_heads * head_size] or
// [num_tokens, num_kv_heads * head_size] or
// [batch_size, seq_len, num_heads, head_size] or
// [num_tokens, num_heads, head_size]
int64_t head_size,
tvm::ffi::TensorView cos_sin_cache, // [max_position, rot_dim]
bool is_neox) {
using namespace host;
// num_tokens = batch_size * seq_len
int64_t num_tokens = positions.numel();
int32_t positions_ndim = positions.ndim();
// Make sure num_tokens dim is consistent across positions, query, and key
RuntimeCheck(
positions_ndim == 1 || positions_ndim == 2, "positions must have shape [num_tokens] or [batch_size, seq_len]");
if (positions_ndim == 1) {
RuntimeCheck(
query.size(0) == positions.size(0) && (!key.has_value() || key.value().size(0) == positions.size(0)),
"query, key and positions must have the same number of tokens");
}
if (positions_ndim == 2) {
RuntimeCheck(
query.size(0) == positions.size(0) && (!key.has_value() || key.value().size(0) == positions.size(0)) &&
query.size(1) == positions.size(1) && (!key.has_value() || key.value().size(1) == positions.size(1)),
"query, key and positions must have the same batch_size and seq_len");
}
// Make sure head_size is valid for query and key
// hidden_size = num_heads * head_size
int query_hidden_size = query.numel() / num_tokens;
int key_hidden_size = key.has_value() ? key.value().numel() / num_tokens : 0;
RuntimeCheck(query_hidden_size % head_size == 0);
RuntimeCheck(key_hidden_size % head_size == 0);
// Make sure query and key have consistent number of heads
int num_heads = query_hidden_size / head_size;
int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads;
RuntimeCheck(num_heads % num_kv_heads == 0);
int rot_dim = cos_sin_cache.size(1);
int seq_dim_idx = positions_ndim - 1;
int64_t query_stride = query.stride(seq_dim_idx);
int64_t key_stride = key.has_value() ? key.value().stride(seq_dim_idx) : 0;
// Determine head stride: for [*, heads, head_size] use stride of last dim;
// for flat [*, heads*head_size], heads blocks are contiguous of size
// head_size
int query_ndim = query.dim();
int64_t head_stride = (query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size;
dim3 grid(num_tokens);
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
auto device = query.device();
const cudaStream_t stream = LaunchKernel::resolve_device(device);
auto positions_data_ptr = static_cast<const int64_t*>(positions.data_ptr());
if (is_neox) {
dispatch_by_dtype<true>(
positions_data_ptr,
query.dtype(),
query.data_ptr(),
key.has_value() ? key.value().data_ptr() : nullptr,
cos_sin_cache.data_ptr(),
rot_dim,
query_stride,
key_stride,
head_stride,
num_heads,
num_kv_heads,
head_size,
grid,
block,
stream);
} else {
dispatch_by_dtype<false>(
positions_data_ptr,
query.dtype(),
query.data_ptr(),
key.has_value() ? key.value().data_ptr() : nullptr,
cos_sin_cache.data_ptr(),
rot_dim,
query_stride,
key_stride,
head_stride,
num_heads,
num_kv_heads,
head_size,
grid,
block,
stream);
}
}
};
} // namespace

257
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh vendored Normal file
View File

@@ -0,0 +1,257 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/impl/norm.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <cstdint>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <type_traits>
namespace {
struct QKNormParams {
void* __restrict__ q;
void* __restrict__ k; // k is offset by (-num_qo_heads * head_dim) elements
int64_t q_stride;
int64_t k_stride;
uint32_t num_qo_heads;
uint32_t num_kv_heads;
float eps;
const void* __restrict__ q_weight;
const void* __restrict__ k_weight;
uint32_t num_tokens;
};
constexpr uint32_t kWarpsPerBlock = 4;
constexpr uint32_t kThreadsPerBlock = kWarpsPerBlock * device::kWarpThreads;
// Warp-level kernel for head_dim <= 256
template <int64_t kHeadDim, bool kUsePDL, typename Float>
__global__ void fused_qknorm_warp(const QKNormParams __grid_constant__ params) {
using namespace device;
using Storage = norm::StorageType<Float, kHeadDim>;
static_assert(sizeof(Float) == 2, "Only support FP16/BF16");
const auto& [q, k, q_stride, k_stride, num_qo_heads, num_kv_heads, eps, q_weight, k_weight, num_tokens] = params;
const auto num_blks = gridDim.x;
const auto num_workers = num_blks * kWarpsPerBlock;
const auto num_q_and_k_heads = num_qo_heads + num_kv_heads;
const auto num_works = num_q_and_k_heads * num_tokens;
const auto start_worker_id = blockIdx.x * kWarpsPerBlock + threadIdx.x / kWarpThreads;
const auto gmem = tile::Memory<Storage>::warp();
PDLWaitPrimary<kUsePDL>(); // wait for primary kernel
for (auto idx = start_worker_id; idx < num_works; idx += num_workers) {
const int64_t token_id = idx / num_q_and_k_heads;
const int64_t head_id = idx % num_q_and_k_heads;
const auto load_q = head_id < num_qo_heads;
const auto input = load_q ? pointer::offset(q, 2 * (token_id * q_stride + head_id * kHeadDim))
: pointer::offset(k, 2 * (token_id * k_stride + head_id * kHeadDim));
const auto weight = load_q ? q_weight : k_weight;
const auto input_vec = gmem.load(input);
const auto weight_vec = gmem.load(weight);
const auto output_vec = norm::apply_norm_warp<kHeadDim>(input_vec, weight_vec, eps);
gmem.store(input, output_vec);
}
PDLTriggerSecondary<kUsePDL>(); // launch secondary kernel
}
// For CTA level, used for head_dim > 256 (512,1024)
template <int64_t kHeadDim, bool kUsePDL, typename Float>
__global__ void fused_qknorm_cta(const QKNormParams __grid_constant__ params) {
using namespace device;
using Storage = norm::StorageType<Float, kHeadDim>;
constexpr auto kNumThreads = host::norm::get_cta_threads<Float, kHeadDim>();
constexpr auto kNumWarps = kNumThreads / kWarpThreads;
static_assert(sizeof(Float) == 2, "Only support FP16/BF16");
const auto& [q, k, q_stride, k_stride, num_qo_heads, num_kv_heads, eps, q_weight, k_weight, num_tokens] = params;
const auto num_q_and_k_heads = num_qo_heads + num_kv_heads;
const auto num_works = num_q_and_k_heads * num_tokens;
const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
__shared__ float smem[norm::kSmemBufferSize];
PDLWaitPrimary<kUsePDL>(); // wait for primary kernel
for (auto idx = blockIdx.x; idx < num_works; idx += gridDim.x) {
const int64_t token_id = idx / num_q_and_k_heads;
const int64_t head_id = idx % num_q_and_k_heads;
const auto load_q = head_id < num_qo_heads;
const auto input = load_q ? pointer::offset(q, 2 * (token_id * q_stride + head_id * kHeadDim))
: pointer::offset(k, 2 * (token_id * k_stride + head_id * kHeadDim));
const auto weight = load_q ? q_weight : k_weight;
const auto input_vec = gmem.load(input);
const auto weight_vec = gmem.load(weight);
const auto output_vec = norm::apply_norm_cta<kHeadDim>(input_vec, weight_vec, eps, smem, kNumWarps);
gmem.store(input, output_vec);
}
PDLTriggerSecondary<kUsePDL>(); // launch secondary kernel
}
// Warp-level kernel struct for head_dim <= 256
template <int64_t kHeadDim, bool kUsePDL, typename DType>
struct QKNormKernelWarp {
static_assert(std::is_same_v<DType, fp16_t> || std::is_same_v<DType, bf16_t>);
static_assert(!host::norm::should_use_cta<DType, kHeadDim>(), "Use QKNormKernelCTA for head_dim > 256");
static constexpr auto kernel = fused_qknorm_warp<kHeadDim, kUsePDL, DType>;
static void
run(const tvm::ffi::TensorView q,
const tvm::ffi::TensorView k,
const tvm::ffi::TensorView q_weight,
const tvm::ffi::TensorView k_weight,
float eps) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto Q = SymbolicSize{"num_qo_heads"};
auto K = SymbolicSize{"num_kv_heads"};
auto D = SymbolicSize{"head_dim"};
auto Sq = SymbolicSize{"q_stride"};
auto Sk = SymbolicSize{"k_stride"};
auto device = SymbolicDevice{};
D.set_value(kHeadDim);
device.set_options<kDLCUDA>();
TensorMatcher({N, Q, D}) // q input
.with_strides({Sq, D, 1})
.with_dtype<DType>()
.with_device(device)
.verify(q);
TensorMatcher({N, K, D}) // k input
.with_strides({Sk, D, 1})
.with_dtype<DType>()
.with_device(device)
.verify(k);
TensorMatcher({D}) // weight
.with_dtype<DType>()
.with_device(device)
.verify(q_weight)
.verify(k_weight);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
// NOTE: we offset the k here to reduce computation cost in the kernel
const auto params = QKNormParams{
.q = q.data_ptr(),
.k = pointer::offset(k.data_ptr(), -2 * static_cast<int64_t>(num_qo_heads) * kHeadDim),
.q_stride = static_cast<int64_t>(Sq.unwrap()),
.k_stride = static_cast<int64_t>(Sk.unwrap()),
.num_qo_heads = num_qo_heads,
.num_kv_heads = num_kv_heads,
.eps = eps,
.q_weight = q_weight.data_ptr(),
.k_weight = k_weight.data_ptr(),
.num_tokens = num_tokens,
};
static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kThreadsPerBlock);
static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
// choose kernel based on dtype
const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens;
const auto needed_blocks = div_ceil(num_works, kWarpsPerBlock);
// we use persistent kernel, which limit the number of blocks to reduce overhead
const auto num_blocks = std::min(kNumSM * max_occupancy, needed_blocks);
LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
// This goes with fused_qknorm_cta
template <int64_t kHeadDim, bool kUsePDL, typename DType>
struct QKNormKernelCTA {
static_assert(std::is_same_v<DType, fp16_t> || std::is_same_v<DType, bf16_t>);
static_assert(host::norm::should_use_cta<DType, kHeadDim>(), "Use QKNormKernelWarp for head_dim <= 256");
static constexpr auto kernel = fused_qknorm_cta<kHeadDim, kUsePDL, DType>;
static constexpr auto kNumThreads = host::norm::get_cta_threads<DType, kHeadDim>();
static void
run(const tvm::ffi::TensorView q,
const tvm::ffi::TensorView k,
const tvm::ffi::TensorView q_weight,
const tvm::ffi::TensorView k_weight,
float eps) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto Q = SymbolicSize{"num_qo_heads"};
auto K = SymbolicSize{"num_kv_heads"};
auto D = SymbolicSize{"head_dim"};
auto Sq = SymbolicSize{"q_stride"};
auto Sk = SymbolicSize{"k_stride"};
auto device = SymbolicDevice{};
D.set_value(kHeadDim);
device.set_options<kDLCUDA>();
TensorMatcher({N, Q, D}) // q input
.with_strides({Sq, D, 1})
.with_dtype<DType>()
.with_device(device)
.verify(q);
TensorMatcher({N, K, D}) // k input
.with_strides({Sk, D, 1})
.with_dtype<DType>()
.with_device(device)
.verify(k);
TensorMatcher({D}) // weight
.with_dtype<DType>()
.with_device(device)
.verify(q_weight)
.verify(k_weight);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
// NOTE: we offset the k here to reduce computation cost in the kernel
const auto params = QKNormParams{
.q = q.data_ptr(),
.k = pointer::offset(k.data_ptr(), -2 * static_cast<int64_t>(num_qo_heads) * kHeadDim),
.q_stride = static_cast<int64_t>(Sq.unwrap()),
.k_stride = static_cast<int64_t>(Sk.unwrap()),
.num_qo_heads = num_qo_heads,
.num_kv_heads = num_kv_heads,
.eps = eps,
.q_weight = q_weight.data_ptr(),
.k_weight = k_weight.data_ptr(),
.num_tokens = num_tokens,
};
static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kNumThreads);
static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens;
// we use persistent kernel, which limit the number of blocks to reduce overhead
const auto num_blocks = std::min<uint32_t>(num_works, max_occupancy * kNumSM);
LaunchKernel(num_blocks, kNumThreads, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
// Unified dispatch: select warp or CTA kernel based on head_dim
template <int64_t kHeadDim, bool kUsePDL, typename DType>
using QKNormKernel = std::conditional_t<
host::norm::should_use_cta<DType, kHeadDim>(),
QKNormKernelCTA<kHeadDim, kUsePDL, DType>,
QKNormKernelWarp<kHeadDim, kUsePDL, DType>>;
} // namespace

179
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh vendored Normal file
View File

@@ -0,0 +1,179 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <cooperative_groups/reduce.h>
#include <tvm/ffi/container/tensor.h>
#include <cooperative_groups.h>
#include <type_traits>
namespace {
template <typename T, int VEC_SIZE_IN_BYTE>
struct VecTypeTrait;
template <>
struct VecTypeTrait<bf16_t, 16> {
using packed_t = packed_t<bf16_t>;
using vec_t = device::AlignedVector<packed_t, 4>;
};
template <>
struct VecTypeTrait<fp16_t, 16> {
using packed_t = packed_t<fp16_t>;
using vec_t = device::AlignedVector<packed_t, 4>;
};
template <>
struct VecTypeTrait<bf16_t, 32> {
using packed_t = packed_t<bf16_t>;
using vec_t = device::AlignedVector<packed_t, 8>;
};
template <>
struct VecTypeTrait<fp16_t, 32> {
using packed_t = packed_t<fp16_t>;
using vec_t = device::AlignedVector<packed_t, 8>;
};
template <typename packed_t>
SGL_DEVICE packed_t rms(const packed_t& val, const packed_t& weight, float rsqrt_square_sum) {
float2 valf = device::cast<fp32x2_t, packed_t>(val);
float2 weightf = device::cast<fp32x2_t, packed_t>(weight);
return device::cast<packed_t, fp32x2_t>(
make_float2(valf.x * weightf.x * rsqrt_square_sum, valf.y * weightf.y * rsqrt_square_sum));
}
template <typename T, int VEC_SIZE_IN_BYTE>
__global__ void qknorm_across_heads_reg_kernel(
T* __restrict__ q,
T* __restrict__ k,
const T* __restrict__ q_weight,
const T* __restrict__ k_weight,
int vec_hidden_size,
float eps) {
constexpr int inner_loop = VEC_SIZE_IN_BYTE == 16 ? 4 : 8;
__shared__ float shared_memory[32];
using vec_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::vec_t;
using packed_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::packed_t;
vec_t v_data;
vec_t v_weight;
const int warp_id = threadIdx.x >> 5;
const int lane_id = threadIdx.x & 31;
const int warp_count = (blockDim.x + 31) >> 5;
const float inv_hidden_size = 1.0f / static_cast<float>(vec_hidden_size * (VEC_SIZE_IN_BYTE / sizeof(T)));
const bool is_q = blockIdx.y == 0;
const auto token_id = blockIdx.x;
float2 acc_square = make_float2(0.0f, 0.0f);
vec_t* data = reinterpret_cast<vec_t*>(is_q ? q : k) + token_id * vec_hidden_size;
const vec_t* weight = reinterpret_cast<const vec_t*>(is_q ? q_weight : k_weight);
if (threadIdx.x < vec_hidden_size) {
v_data = data[threadIdx.x];
v_weight = weight[threadIdx.x];
for (int i = 0; i < inner_loop; i++) {
float2 val = device::cast<fp32x2_t, packed_t>(v_data[i]);
acc_square.x += val.x * val.x;
acc_square.y += val.y * val.y;
}
}
auto cg_warp = cooperative_groups::tiled_partition<32>(cooperative_groups::this_thread_block());
float* buffer = shared_memory;
float warp_sum = cooperative_groups::reduce(cg_warp, acc_square.x + acc_square.y, cooperative_groups::plus<float>());
if (lane_id == 0) {
buffer[warp_id] = warp_sum;
}
__syncthreads();
if (threadIdx.x < 32) {
float cta_sum = cooperative_groups::reduce(
cg_warp, (threadIdx.x < warp_count) ? buffer[threadIdx.x] : 0.0f, cooperative_groups::plus<float>());
if (threadIdx.x == 0) {
buffer[0] = rsqrtf(eps + cta_sum * inv_hidden_size);
}
}
__syncthreads();
if (threadIdx.x < vec_hidden_size) {
float rsqrt_val = buffer[0];
for (int i = 0; i < inner_loop; i++) {
v_data[i] = rms(v_data[i], v_weight[i], rsqrt_val);
}
data[threadIdx.x] = v_data;
}
}
template <typename DType>
struct QKNormAcrossHeadsKernel {
static void
run(const tvm::ffi::TensorView q,
const tvm::ffi::TensorView k,
const tvm::ffi::TensorView q_weight,
const tvm::ffi::TensorView k_weight,
float eps) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto D = SymbolicSize{"hidden_size"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({N, D}) // q
.with_strides({D, 1})
.with_dtype<DType>()
.with_device(device)
.verify(q);
TensorMatcher({N, D}) // k
.with_strides({D, 1})
.with_dtype<DType>()
.with_device(device)
.verify(k);
TensorMatcher({D}) // q_weight
.with_dtype<DType>()
.with_device(device)
.verify(q_weight);
TensorMatcher({D}) // k_weight
.with_dtype<DType>()
.with_device(device)
.verify(k_weight);
int hidden_size = static_cast<int>(D.unwrap());
if (hidden_size <= (device::kMaxVecBytes == 32 ? 12288 : 8192)) {
int elements_in_vec = device::kMaxVecBytes / sizeof(DType);
int vec_hidden_size = hidden_size / elements_in_vec;
uint threads = (vec_hidden_size + 31) / 32 * 32;
// Runtime check
host::RuntimeCheck(
hidden_size % elements_in_vec == 0,
"hidden_size",
hidden_size,
" can not align to elements_in_vec ",
elements_in_vec);
auto kernel = qknorm_across_heads_reg_kernel<DType, device::kMaxVecBytes>;
LaunchKernel(dim3(static_cast<uint>(N.unwrap()), 2), threads, device.unwrap())
.enable_pdl(false)(
kernel,
reinterpret_cast<DType*>(q.data_ptr()),
reinterpret_cast<DType*>(k.data_ptr()),
reinterpret_cast<DType*>(q_weight.data_ptr()),
reinterpret_cast<DType*>(k_weight.data_ptr()),
vec_hidden_size,
eps);
} else {
host::RuntimeCheck(false, "Large hidden_sizes are not supported for now.");
}
}
};
} // namespace

57
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/resolve_future_token_ids.cuh vendored Normal file
View File

@@ -0,0 +1,57 @@
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
#include <sgl_kernel/utils.h> // For RuntimeCheck, div_ceil
#include <sgl_kernel/utils.cuh> // For LaunchKernel
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <cstddef>
#include <cstdint>
namespace {
template <typename T>
__global__ void resolve_future_token_ids_kernel(T* __restrict__ input_ids, const T* __restrict__ future_map, size_t n) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
T val = input_ids[idx];
if (val < 0) {
T key = -val;
if (key < 0) key = 0; // clamp for overflow
input_ids[idx] = future_map[key];
}
}
}
constexpr size_t kBlockSize = 256;
template <typename T>
struct ResolveFutureTokenIds {
static void run(tvm::ffi::TensorView input_ids, tvm::ffi::TensorView future_map) {
using namespace host;
SymbolicSize N = {"num_tokens"};
SymbolicSize M = {"map_size"};
SymbolicDevice device_;
device_.set_options<kDLCUDA, kDLROCM>();
TensorMatcher({N}).with_dtype<T>().with_device(device_).verify(input_ids);
TensorMatcher({M}).with_dtype<T>().with_device(device_).verify(future_map);
const size_t num_tokens = N.unwrap();
if (num_tokens == 0) return;
const size_t grid_size = div_ceil(num_tokens, kBlockSize);
const DLDevice device = device_.unwrap();
LaunchKernel(grid_size, kBlockSize, device)(
resolve_future_token_ids_kernel<T>,
static_cast<T*>(input_ids.data_ptr()),
static_cast<const T*>(future_map.data_ptr()),
num_tokens);
}
};
} // namespace

361
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh vendored Normal file
View File

@@ -0,0 +1,361 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/impl/norm.cuh>
#include <tvm/ffi/container/tensor.h>
namespace {
struct RMSNormParams {
const void* input;
const void* __restrict__ weight;
void* output;
int64_t input_stride;
int64_t output_stride;
uint32_t num_tokens;
float eps;
};
template <int64_t kDim, bool kUsePDL, typename Float>
__global__ void rmsnorm_cta(const RMSNormParams __grid_constant__ params) {
using namespace device;
using Storage = norm::StorageType<Float, kDim>;
constexpr auto kNumThreads = host::norm::get_cta_threads<Float, kDim>();
constexpr auto kNumWarps = kNumThreads / kWarpThreads;
const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params;
const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
__shared__ float smem[norm::kSmemBufferSize];
PDLWaitPrimary<kUsePDL>(); // wait for primary kernel
for (uint32_t i = blockIdx.x; i < num_tokens; i += gridDim.x) {
const auto input_ptr = pointer::offset<Float>(input, i * input_stride);
const auto output_ptr = pointer::offset<Float>(output, i * output_stride);
const auto input_vec = gmem.load(input_ptr);
const auto weight_vec = gmem.load(weight_ptr);
const auto output_vec = norm::apply_norm_cta<kDim>(input_vec, weight_vec, eps, smem, kNumWarps);
gmem.store(output_ptr, output_vec);
}
PDLTriggerSecondary<kUsePDL>(); // launch secondary kernel
}
// Pre-Blackwell: 16B vector, each thread loads/stores twice
template <int64_t kDim, bool kUsePDL, typename Float>
__global__ __launch_bounds__(kDim / 16) void rmsnorm_cta_double(const RMSNormParams __grid_constant__ params) {
using namespace device;
using Float2 = packed_t<Float>;
using Storage = AlignedVector<Float2, 4>;
constexpr auto kNumThreads = kDim / 16;
constexpr auto kNumWarps = kNumThreads / kWarpThreads;
const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params;
const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
__shared__ float smem[32];
PDLWaitPrimary<kUsePDL>();
const auto input_ptr = pointer::offset<Float>(input, blockIdx.x * input_stride);
const auto output_ptr = pointer::offset<Float>(output, blockIdx.x * output_stride);
const auto input_first = gmem.load(input_ptr, 0);
const auto input_second = gmem.load(input_ptr, 1);
const auto weight_first = gmem.load(weight_ptr, 0);
const auto weight_second = gmem.load(weight_ptr, 1);
float sum_of_squares = 0.0f;
#pragma unroll
for (auto j = 0u; j < 4u; ++j) {
const auto [x, y] = cast<fp32x2_t>(input_first[j]);
sum_of_squares += x * x + y * y;
}
#pragma unroll
for (auto j = 0u; j < 4u; ++j) {
const auto [x, y] = cast<fp32x2_t>(input_second[j]);
sum_of_squares += x * x + y * y;
}
sum_of_squares = warp::reduce_sum(sum_of_squares);
const auto warp_id = threadIdx.x / kWarpThreads;
smem[warp_id] = sum_of_squares;
__syncthreads();
if (warp_id == 0) {
const auto tx = threadIdx.x;
const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
sum_of_squares = warp::reduce_sum(local_sum);
smem[tx] = math::rsqrt(sum_of_squares / kDim + eps);
}
__syncthreads();
const float norm_factor = smem[warp_id];
Storage output_first, output_second;
#pragma unroll
for (auto j = 0u; j < 4u; ++j) {
const auto [ix, iy] = cast<fp32x2_t>(input_first[j]);
const auto [wx, wy] = cast<fp32x2_t>(weight_first[j]);
output_first[j] = cast<Float2>(fp32x2_t{ix * norm_factor * wx, iy * norm_factor * wy});
}
#pragma unroll
for (auto j = 0u; j < 4u; ++j) {
const auto [ix, iy] = cast<fp32x2_t>(input_second[j]);
const auto [wx, wy] = cast<fp32x2_t>(weight_second[j]);
output_second[j] = cast<Float2>(fp32x2_t{ix * norm_factor * wx, iy * norm_factor * wy});
}
gmem.store(output_ptr, output_first, 0);
gmem.store(output_ptr, output_second, 1);
PDLTriggerSecondary<kUsePDL>();
}
// Blackwell: 32B vector, each thread loads/stores once
template <int64_t kDim, bool kUsePDL, typename Float>
__global__ __launch_bounds__(kDim / 16) void rmsnorm_cta_wide(const RMSNormParams __grid_constant__ params) {
using namespace device;
using Float2 = packed_t<Float>;
using Storage = AlignedVector<Float2, 8>;
constexpr auto kNumThreads = kDim / 16;
constexpr auto kNumWarps = kNumThreads / kWarpThreads;
const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params;
const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
__shared__ float smem[32];
PDLWaitPrimary<kUsePDL>();
const auto input_ptr = pointer::offset<Float>(input, blockIdx.x * input_stride);
const auto output_ptr = pointer::offset<Float>(output, blockIdx.x * output_stride);
const auto input_vec = gmem.load(input_ptr);
const auto weight_vec = gmem.load(weight_ptr);
float sum_of_squares = 0.0f;
#pragma unroll
for (auto j = 0u; j < 8u; ++j) {
const auto [x, y] = cast<fp32x2_t>(input_vec[j]);
sum_of_squares += x * x + y * y;
}
sum_of_squares = warp::reduce_sum(sum_of_squares);
const auto warp_id = threadIdx.x / kWarpThreads;
smem[warp_id] = sum_of_squares;
__syncthreads();
if (warp_id == 0) {
const auto tx = threadIdx.x;
const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
sum_of_squares = warp::reduce_sum(local_sum);
smem[tx] = math::rsqrt(sum_of_squares / kDim + eps);
}
__syncthreads();
const float norm_factor = smem[warp_id];
Storage output_vec;
#pragma unroll
for (auto j = 0u; j < 8u; ++j) {
const auto [ix, iy] = cast<fp32x2_t>(input_vec[j]);
const auto [wx, wy] = cast<fp32x2_t>(weight_vec[j]);
output_vec[j] = cast<Float2>(fp32x2_t{ix * norm_factor * wx, iy * norm_factor * wy});
}
gmem.store(output_ptr, output_vec);
PDLTriggerSecondary<kUsePDL>();
}
template <int64_t kDim, bool kUsePDL, typename Float>
__global__ void rmsnorm_warp(const RMSNormParams __grid_constant__ params) {
using namespace device;
using Storage = norm::StorageType<Float, kDim>;
const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params;
const auto gmem = tile::Memory<Storage>::warp();
PDLWaitPrimary<kUsePDL>(); // wait for primary kernel
for (uint32_t i = blockIdx.x; i < num_tokens; i += gridDim.x) {
const auto input_ptr = pointer::offset<Float>(input, i * input_stride);
const auto output_ptr = pointer::offset<Float>(output, i * output_stride);
const auto input_vec = gmem.load(input_ptr);
const auto weight_vec = gmem.load(weight_ptr);
const auto output_vec = norm::apply_norm_warp<kDim>(input_vec, weight_vec, eps);
gmem.store(output_ptr, output_vec);
}
PDLTriggerSecondary<kUsePDL>(); // launch secondary kernel
}
template <int64_t kDim, bool kUsePDL, typename DType>
struct RMSNormWarpKernel {
static_assert(host::norm::is_config_supported<DType, kDim>(), "Unsupported norm configuration");
static_assert(kDim <= 256, "Use RMSNormKernel for hidden sizes > 256");
static constexpr auto kernel = rmsnorm_warp<kDim, kUsePDL, DType>;
static void
run(const tvm::ffi::TensorView input,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView output,
float eps) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto D = SymbolicSize{"hidden_size"};
auto SI = SymbolicSize{"input_stride"};
auto SO = SymbolicSize{"output_stride"};
auto device = SymbolicDevice{};
D.set_value(kDim);
device.set_options<kDLCUDA>();
TensorMatcher({N, D}) // input
.with_strides({SI, 1})
.with_dtype<DType>()
.with_device(device)
.verify(input);
TensorMatcher({D}) // weight
.with_dtype<DType>()
.with_device(device)
.verify(weight);
TensorMatcher({N, D}) // output
.with_strides({SO, 1})
.with_dtype<DType>()
.with_device(device)
.verify(output);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto params = RMSNormParams{
.input = input.data_ptr(),
.weight = weight.data_ptr(),
.output = output.data_ptr(),
.input_stride = SI.unwrap(),
.output_stride = SO.unwrap(),
.num_tokens = num_tokens,
.eps = eps,
};
static constexpr uint32_t kNumThreads = device::kWarpThreads;
static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kNumThreads);
static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
const auto num_blocks = std::min<uint32_t>(num_tokens, max_occupancy * kNumSM);
LaunchKernel(num_blocks, kNumThreads, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
template <int64_t kDim, bool kUsePDL, typename DType>
struct RMSNormKernel {
static_assert(host::norm::should_use_cta<DType, kDim>(), "Hidden size invalid for RMSNorm");
static constexpr auto kernel = rmsnorm_cta<kDim, kUsePDL, DType>;
static void
run(const tvm::ffi::TensorView input,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView output,
float eps) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto D = SymbolicSize{"hidden_size"};
auto SI = SymbolicSize{"input_stride"};
auto SO = SymbolicSize{"output_stride"};
auto device = SymbolicDevice{};
D.set_value(kDim);
device.set_options<kDLCUDA>();
TensorMatcher({N, D}) // input
.with_strides({SI, 1})
.with_dtype<DType>()
.with_device(device)
.verify(input);
TensorMatcher({D}) // weight
.with_dtype<DType>()
.with_device(device)
.verify(weight);
TensorMatcher({N, D}) // output
.with_strides({SO, 1})
.with_dtype<DType>()
.with_device(device)
.verify(output);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto params = RMSNormParams{
.input = input.data_ptr(),
.weight = weight.data_ptr(),
.output = output.data_ptr(),
.input_stride = SI.unwrap(),
.output_stride = SO.unwrap(),
.num_tokens = num_tokens,
.eps = eps,
};
static constexpr auto kNumThreads = norm::get_cta_threads<DType, kDim>();
static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kNumThreads);
static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
const auto num_blocks = std::min<uint32_t>(num_tokens, max_occupancy * kNumSM);
LaunchKernel(num_blocks, kNumThreads, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
template <int64_t kDim, bool kUsePDL, typename DType>
struct RMSNormHalfKernel {
static_assert(kDim % 512 == 0 && sizeof(DType) == 2);
#if SGL_ARCH_BLACKWELL_OR_GREATER
static constexpr auto kernel = rmsnorm_cta_wide<kDim, kUsePDL, DType>;
#else
static constexpr auto kernel = rmsnorm_cta_double<kDim, kUsePDL, DType>;
#endif
static constexpr auto kBlockSize = static_cast<uint32_t>(kDim / 16);
static void
run(const tvm::ffi::TensorView input,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView output,
float eps) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto D = SymbolicSize{"hidden_size"};
auto SI = SymbolicSize{"input_stride"};
auto SO = SymbolicSize{"output_stride"};
auto device = SymbolicDevice{};
D.set_value(kDim);
device.set_options<kDLCUDA>();
TensorMatcher({N, D}) // input
.with_strides({SI, 1})
.with_dtype<DType>()
.with_device(device)
.verify(input);
TensorMatcher({D}) // weight
.with_dtype<DType>()
.with_device(device)
.verify(weight);
TensorMatcher({N, D}) // output
.with_strides({SO, 1})
.with_dtype<DType>()
.with_device(device)
.verify(output);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto params = RMSNormParams{
.input = input.data_ptr(),
.weight = weight.data_ptr(),
.output = output.data_ptr(),
.input_stride = SI.unwrap(),
.output_stride = SO.unwrap(),
.num_tokens = num_tokens,
.eps = eps,
};
LaunchKernel(num_tokens, kBlockSize, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace

470
third_party/sglang/python/sglang/jit_kernel/csrc/elementwise/rope.cuh vendored Normal file
View File

@@ -0,0 +1,470 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <dlpack/dlpack.h>
#include <numeric>
namespace {
struct FusedRopeParams {
void* __restrict__ q_ptr;
void* __restrict__ k_ptr; // NOTE: this k is pre-offset in host code to reduce computation in kernel
const void* __restrict__ cos_sin_cache_ptr;
const void* __restrict__ positions;
int64_t q_stride_bytes;
int64_t k_stride_bytes;
int64_t head_stride_bytes;
uint32_t num_qo_heads;
uint32_t num_kv_heads;
uint32_t num_tokens;
};
struct FusedRopeStoreParams {
FusedRopeParams base_params;
void* v_ptr;
void* __restrict__ k_cache;
void* __restrict__ v_cache;
const void* __restrict__ out_loc;
int64_t v_stride_bytes;
int64_t cache_stride_bytes;
};
constexpr uint32_t kBlockSize = 128;
[[maybe_unused]]
constexpr auto next_pow2(uint32_t target, uint32_t factor = 1) {
uint32_t power = 1;
while (power * factor < target)
power *= 2;
return power;
}
template <bool kIsNeox, int64_t kRopeDim, bool kUsePDL, typename DType, typename IdType, uint32_t kWorkThreads>
__global__ void fused_rope_kernel(const __grid_constant__ FusedRopeParams params) {
using namespace device;
constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float);
constexpr int64_t kVecSize = next_pow2(kRopeDim, (2 * kWorkThreads * (1 + kIsNeox)));
using DType2 = packed_t<DType>;
using InputStorage = AlignedVector<DType2, kVecSize>;
constexpr int64_t kDimPerThread = kVecSize * 2 * (1 + kIsNeox);
constexpr uint32_t kLaneCount = kRopeDim / kDimPerThread;
static_assert(kRopeDim % kDimPerThread == 0 && kLaneCount <= kWorkThreads);
const auto &[
q, k, cos_sin_cache_ptr, positions, // pointers
q_stride_bytes, k_stride_bytes, head_stride_bytes, // strides
num_qo_heads, num_kv_heads, num_tokens // dimensions
] = params;
const auto num_blks = gridDim.x;
constexpr auto kWorkersPerBlock = kBlockSize / kWorkThreads;
const auto num_workers = num_blks * kWorkersPerBlock;
const auto num_q_and_k_heads = num_qo_heads + num_kv_heads;
const auto num_works = num_q_and_k_heads * num_tokens;
const auto start_worker_id = (blockIdx.x * kBlockSize + threadIdx.x) / kWorkThreads;
const auto cos_cache_ptr = cos_sin_cache_ptr;
const auto sin_cache_ptr = pointer::offset(cos_sin_cache_ptr, kCosSinStrideBytes / 2);
uint32_t lane_id = threadIdx.x % kWorkThreads;
if constexpr (kLaneCount < kWorkThreads) {
if (lane_id >= kLaneCount) return;
}
PDLWaitPrimary<kUsePDL>();
for (auto idx = start_worker_id; idx < num_works; idx += num_workers) {
const int64_t token_id = idx / num_q_and_k_heads;
const int64_t head_id = idx % num_q_and_k_heads;
const auto pos = static_cast<const IdType*>(positions)[token_id];
const auto load_q = head_id < num_qo_heads;
const auto input_ = load_q ? pointer::offset(q, token_id * q_stride_bytes) //
: pointer::offset(k, token_id * k_stride_bytes);
const auto input = pointer::offset(input_, head_id * head_stride_bytes);
const auto cos_ptr = pointer::offset(cos_cache_ptr, pos * kCosSinStrideBytes);
const auto sin_ptr = pointer::offset(sin_cache_ptr, pos * kCosSinStrideBytes);
if constexpr (kIsNeox) {
using CacheStorage = AlignedVector<fp32x2_t, kVecSize>;
const auto input_x = input;
const auto input_y = pointer::offset(input, (kRopeDim / 2) * sizeof(DType));
auto input_vec_x = load_as<InputStorage>(input_x, lane_id);
auto input_vec_y = load_as<InputStorage>(input_y, lane_id);
const auto cos_pair = load_as<CacheStorage>(cos_ptr, lane_id);
const auto sin_pair = load_as<CacheStorage>(sin_ptr, lane_id);
#pragma unroll
for (int64_t j = 0; j < kVecSize; ++j) {
const auto [x0, x1] = cast<fp32x2_t>(input_vec_x[j]);
const auto [y0, y1] = cast<fp32x2_t>(input_vec_y[j]);
const auto [cos_0, cos_1] = cos_pair[j];
const auto [sin_0, sin_1] = sin_pair[j];
const auto out_x0 = x0 * cos_0 - y0 * sin_0;
const auto out_y0 = x0 * sin_0 + y0 * cos_0;
const auto out_x1 = x1 * cos_1 - y1 * sin_1;
const auto out_y1 = x1 * sin_1 + y1 * cos_1;
input_vec_x[j] = cast<DType2, fp32x2_t>({out_x0, out_x1});
input_vec_y[j] = cast<DType2, fp32x2_t>({out_y0, out_y1});
}
store_as<InputStorage>(input_x, input_vec_x, lane_id);
store_as<InputStorage>(input_y, input_vec_y, lane_id);
} else {
using CacheStorage = AlignedVector<float, kVecSize>;
auto input_vec = load_as<InputStorage>(input, lane_id);
const auto cos_vec = load_as<CacheStorage>(cos_ptr, lane_id);
const auto sin_vec = load_as<CacheStorage>(sin_ptr, lane_id);
#pragma unroll
for (int64_t j = 0; j < kVecSize; ++j) {
const auto [x, y] = cast<fp32x2_t>(input_vec[j]);
const auto cos = cos_vec[j];
const auto sin = sin_vec[j];
const auto out_x = x * cos - y * sin;
const auto out_y = x * sin + y * cos;
input_vec[j] = cast<DType2, fp32x2_t>({out_x, out_y});
}
store_as<InputStorage>(input, input_vec, lane_id);
}
}
PDLTriggerSecondary<kUsePDL>();
}
template <bool kIsNeox, int64_t kRopeDim, bool kUsePDL, typename DType, typename IdType, uint32_t kWorkThreads>
__global__ void fused_rope_store_kernel(const __grid_constant__ FusedRopeStoreParams params) {
using namespace device;
constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float);
constexpr int64_t kVecSize = kRopeDim / (2 * kWorkThreads * (1 + kIsNeox));
using DType2 = packed_t<DType>;
using InputStorage = AlignedVector<DType2, kVecSize>;
constexpr int64_t kDimPerThread = kVecSize * 2 * (1 + kIsNeox);
static_assert(kRopeDim == kDimPerThread * kWorkThreads);
const auto& [base_params, v_ptr, k_cache, v_cache, out_loc, v_stride_bytes, cache_stride_bytes] = params;
const auto &[
q, k, cos_sin_cache_ptr, positions, // pointers
q_stride_bytes, k_stride_bytes, head_stride_bytes, // strides
num_qo_heads, num_kv_heads, num_tokens // dimensions
] = base_params;
const auto num_blks = gridDim.x;
constexpr auto kWorkersPerBlock = kBlockSize / kWorkThreads;
const auto num_workers = num_blks * kWorkersPerBlock;
const auto num_q_and_k_heads = num_qo_heads + num_kv_heads;
const auto num_works = num_q_and_k_heads * num_tokens;
const auto num_extra_works = num_kv_heads * num_tokens; // rope works + v store works
const auto start_worker_id = (blockIdx.x * kBlockSize + threadIdx.x) / kWorkThreads;
const auto lane_id = threadIdx.x % kWorkThreads;
const auto cos_cache_ptr = cos_sin_cache_ptr;
const auto sin_cache_ptr = pointer::offset(cos_sin_cache_ptr, kCosSinStrideBytes / 2);
auto idx = start_worker_id;
PDLWaitPrimary<kUsePDL>();
// in this case, head_dim = rope_dim must be true
__builtin_assume(head_stride_bytes == kRopeDim * sizeof(DType));
for (; idx < num_works; idx += num_workers) {
const int64_t token_id = idx / num_q_and_k_heads;
const int64_t head_id = idx % num_q_and_k_heads;
const auto pos = static_cast<const IdType*>(positions)[token_id];
const auto loc = static_cast<const IdType*>(out_loc)[token_id];
const auto load_q = head_id < num_qo_heads;
const auto input_ = load_q ? pointer::offset(q, token_id * q_stride_bytes) //
: pointer::offset(k, token_id * k_stride_bytes);
const auto input = pointer::offset(input_, head_id * head_stride_bytes);
const auto cos_ptr = pointer::offset(cos_cache_ptr, pos * kCosSinStrideBytes);
const auto sin_ptr = pointer::offset(sin_cache_ptr, pos * kCosSinStrideBytes);
if constexpr (kIsNeox) {
using CacheStorage = AlignedVector<fp32x2_t, kVecSize>;
const auto input_x = input;
const auto input_y = pointer::offset(input, (kRopeDim / 2) * sizeof(DType));
auto input_vec_x = load_as<InputStorage>(input_x, lane_id);
auto input_vec_y = load_as<InputStorage>(input_y, lane_id);
const auto cos_pair = load_as<CacheStorage>(cos_ptr, lane_id);
const auto sin_pair = load_as<CacheStorage>(sin_ptr, lane_id);
#pragma unroll
for (int64_t j = 0; j < kVecSize; ++j) {
const auto [x0, x1] = cast<fp32x2_t>(input_vec_x[j]);
const auto [y0, y1] = cast<fp32x2_t>(input_vec_y[j]);
const auto [cos_0, cos_1] = cos_pair[j];
const auto [sin_0, sin_1] = sin_pair[j];
const auto out_x0 = x0 * cos_0 - y0 * sin_0;
const auto out_y0 = x0 * sin_0 + y0 * cos_0;
const auto out_x1 = x1 * cos_1 - y1 * sin_1;
const auto out_y1 = x1 * sin_1 + y1 * cos_1;
input_vec_x[j] = cast<DType2, fp32x2_t>({out_x0, out_x1});
input_vec_y[j] = cast<DType2, fp32x2_t>({out_y0, out_y1});
}
store_as<InputStorage>(input, input_vec_x, lane_id);
const auto input_y_out = pointer::offset(input, (kRopeDim / 2) * sizeof(DType));
store_as<InputStorage>(input_y_out, input_vec_y, lane_id);
if (!load_q) {
const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes);
store_as<InputStorage>(k_out, input_vec_x, lane_id);
const auto k_out_y = pointer::offset(k_out, (kRopeDim / 2) * sizeof(DType));
store_as<InputStorage>(k_out_y, input_vec_y, lane_id);
}
} else {
using CacheStorage = AlignedVector<float, kVecSize>;
auto input_vec = load_as<InputStorage>(input, lane_id);
const auto cos_vec = load_as<CacheStorage>(cos_ptr, lane_id);
const auto sin_vec = load_as<CacheStorage>(sin_ptr, lane_id);
#pragma unroll
for (int64_t j = 0; j < kVecSize; ++j) {
const auto [x, y] = cast<fp32x2_t>(input_vec[j]);
const auto cos = cos_vec[j];
const auto sin = sin_vec[j];
const auto out_x = x * cos - y * sin;
const auto out_y = x * sin + y * cos;
input_vec[j] = cast<DType2, fp32x2_t>({out_x, out_y});
}
store_as<InputStorage>(input, input_vec, lane_id);
if (!load_q) {
const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes);
store_as<InputStorage>(k_out, input_vec, lane_id);
}
}
}
__syncwarp(); // to avoid warp divergence
idx -= num_works;
for (; idx < num_extra_works; idx += num_workers) {
using VStorage = AlignedVector<DType, kRopeDim / kWorkThreads>;
const int64_t token_id = idx / num_kv_heads;
const int64_t head_id = idx % num_kv_heads;
const auto loc = static_cast<const IdType*>(out_loc)[token_id];
const auto input = pointer::offset(v_ptr, token_id * v_stride_bytes, head_id * head_stride_bytes);
const auto input_vec = load_as<VStorage>(input, lane_id);
const auto output = pointer::offset(v_cache, loc * cache_stride_bytes, head_id * head_stride_bytes);
store_as<VStorage>(output, input_vec, lane_id);
}
PDLTriggerSecondary<kUsePDL>();
}
template <bool kIsNeox, int64_t kRopeDim, bool kUsePDL, typename DType>
struct FusedRopeKernel {
static constexpr uint32_t kDimPerThread = std::gcd(16 / sizeof(DType), kRopeDim);
static constexpr uint32_t kWorkThreads = next_pow2(kRopeDim, kDimPerThread);
static constexpr bool kSupportFused = kWorkThreads * kDimPerThread == kRopeDim;
static_assert(kRopeDim % kDimPerThread == 0);
static_assert(kBlockSize % kWorkThreads == 0);
template <typename IdType>
static constexpr auto _kernel_0 = fused_rope_kernel<kIsNeox, kRopeDim, kUsePDL, DType, IdType, kWorkThreads>;
template <typename IdType>
static constexpr auto _kernel_1 = fused_rope_store_kernel<kIsNeox, kRopeDim, kUsePDL, DType, IdType, kWorkThreads>;
static auto get_num_sm(DLDevice device) {
static const auto kNumSM = host::runtime::get_sm_count(device.device_id);
return kNumSM;
}
static void
run(const tvm::ffi::TensorView q,
const tvm::ffi::TensorView k,
const tvm::ffi::TensorView cos_sin_cache,
const tvm::ffi::TensorView positions) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto Q = SymbolicSize{"num_qo_heads"};
auto K = SymbolicSize{"num_kv_heads"};
auto D = SymbolicSize{"rope_dim"};
auto Dq = SymbolicSize{"q_stride"};
auto Dk = SymbolicSize{"k_stride"};
auto Dd = SymbolicSize{"head_stride"};
auto device = SymbolicDevice{};
auto id_type = SymbolicDType{};
D.set_value(kRopeDim);
device.set_options<kDLCUDA>();
TensorMatcher({N, Q, D}) // q input
.with_strides({Dq, Dd, 1})
.with_dtype<DType>()
.with_device(device)
.verify(q);
TensorMatcher({N, K, D}) // k input
.with_strides({Dk, Dd, 1})
.with_dtype<DType>()
.with_device(device)
.verify(k);
TensorMatcher({-1, D}) // cos_sin_cache
.with_dtype<float>()
.with_device(device)
.verify(cos_sin_cache);
TensorMatcher({N}) // positions
.with_dtype<int32_t, int64_t>(id_type)
.with_device(device)
.verify(positions);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
const auto q_stride_bytes = static_cast<int64_t>(Dq.unwrap() * sizeof(DType));
const auto k_stride_bytes = static_cast<int64_t>(Dk.unwrap() * sizeof(DType));
const auto head_stride_bytes = static_cast<int64_t>(Dd.unwrap() * sizeof(DType));
// NOTE: we offset the k here to reduce computation cost in the kernel
const int64_t k_offset = static_cast<int64_t>(num_qo_heads) * head_stride_bytes;
const auto params = FusedRopeParams{
.q_ptr = q.data_ptr(),
.k_ptr = pointer::offset(k.data_ptr(), -k_offset),
.cos_sin_cache_ptr = cos_sin_cache.data_ptr(),
.positions = positions.data_ptr(),
.q_stride_bytes = q_stride_bytes,
.k_stride_bytes = k_stride_bytes,
.head_stride_bytes = head_stride_bytes,
.num_qo_heads = num_qo_heads,
.num_kv_heads = num_kv_heads,
.num_tokens = num_tokens,
};
const auto is_int32 = id_type.is_type<int32_t>();
const auto kernel = is_int32 ? _kernel_0<int32_t> : _kernel_0<int64_t>;
const uint32_t kNumSM = get_num_sm(device.unwrap());
static const uint32_t kOccupancyTable[2] = {
runtime::get_blocks_per_sm(_kernel_0<int32_t>, kBlockSize),
runtime::get_blocks_per_sm(_kernel_0<int64_t>, kBlockSize),
};
const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM;
const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens;
const auto needed_blocks = div_ceil(num_works, (kBlockSize / kWorkThreads));
const auto num_blocks = std::min(max_blocks, needed_blocks);
LaunchKernel(num_blocks, kBlockSize, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
static void run_fused(
const tvm::ffi::TensorView q,
const tvm::ffi::TensorView k,
const tvm::ffi::TensorView v,
const tvm::ffi::TensorView k_cache,
const tvm::ffi::TensorView v_cache,
const tvm::ffi::TensorView cos_sin_cache,
const tvm::ffi::TensorView positions,
const tvm::ffi::TensorView out_loc) {
if constexpr (kSupportFused) {
return _run_fused_impl(q, k, v, k_cache, v_cache, cos_sin_cache, positions, out_loc);
} else {
host::Panic("Fused rope + store is not supported for rope_dim ", kRopeDim);
}
}
static void _run_fused_impl(
const tvm::ffi::TensorView q,
const tvm::ffi::TensorView k,
const tvm::ffi::TensorView v,
const tvm::ffi::TensorView k_cache,
const tvm::ffi::TensorView v_cache,
const tvm::ffi::TensorView cos_sin_cache,
const tvm::ffi::TensorView positions,
const tvm::ffi::TensorView out_loc) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto Q = SymbolicSize{"num_qo_heads"};
auto K = SymbolicSize{"num_kv_heads"};
auto D = SymbolicSize{"rope_dim"};
auto R = SymbolicSize{"row_size"};
auto Dq = SymbolicSize{"q_stride"};
auto Dk = SymbolicSize{"k_stride"};
auto Dv = SymbolicSize{"v_stride"};
auto Dd = SymbolicSize{"head_stride"};
auto Dc = SymbolicSize{"cache_stride"};
auto device = SymbolicDevice{};
auto id_type = SymbolicDType{};
D.set_value(kRopeDim);
device.set_options<kDLCUDA>();
TensorMatcher({N, Q, D}) // q input
.with_strides({Dq, Dd, 1})
.with_dtype<DType>()
.with_device(device)
.verify(q);
TensorMatcher({N, K, D}) // k input
.with_strides({Dk, Dd, 1})
.with_dtype<DType>()
.with_device(device)
.verify(k);
TensorMatcher({N, K, D}) // v input
.with_strides({Dv, Dd, 1})
.with_dtype<DType>()
.with_device(device)
.verify(v);
TensorMatcher({-1, D}) // cos_sin_cache
.with_dtype<float>()
.with_device(device)
.verify(cos_sin_cache);
TensorMatcher({N}) // positions, out_loc
.with_dtype<int32_t, int64_t>(id_type)
.with_device(device)
.verify(positions)
.verify(out_loc);
TensorMatcher({-1, R}) // k_cache
.with_strides({Dc, 1})
.with_dtype<DType>()
.with_device(device)
.verify(k_cache)
.verify(v_cache);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
const auto q_stride_bytes = static_cast<int64_t>(Dq.unwrap() * sizeof(DType));
const auto k_stride_bytes = static_cast<int64_t>(Dk.unwrap() * sizeof(DType));
const auto head_stride = Dd.unwrap();
const auto row_dim = R.unwrap();
const auto head_stride_bytes = static_cast<int64_t>(Dd.unwrap() * sizeof(DType));
RuntimeCheck(kRopeDim == head_stride, "rope_dim ", kRopeDim, " should = head_stride ", head_stride);
RuntimeCheck(num_kv_heads * kRopeDim == row_dim, "invalid kvcache");
// NOTE: we offset the k here to reduce computation cost in the kernel
const int64_t k_offset = static_cast<int64_t>(num_qo_heads) * head_stride_bytes;
const auto params = FusedRopeParams{
.q_ptr = q.data_ptr(),
.k_ptr = pointer::offset(k.data_ptr(), -k_offset),
.cos_sin_cache_ptr = cos_sin_cache.data_ptr(),
.positions = positions.data_ptr(),
.q_stride_bytes = q_stride_bytes,
.k_stride_bytes = k_stride_bytes,
.head_stride_bytes = head_stride_bytes,
.num_qo_heads = num_qo_heads,
.num_kv_heads = num_kv_heads,
.num_tokens = num_tokens,
};
const auto v_stride_bytes = static_cast<int64_t>(Dv.unwrap() * sizeof(DType));
const auto cache_stride_bytes = static_cast<int64_t>(Dc.unwrap() * sizeof(DType));
const auto store_params = FusedRopeStoreParams{
.base_params = params,
.v_ptr = v.data_ptr(),
.k_cache = pointer::offset(k_cache.data_ptr(), -k_offset),
.v_cache = v_cache.data_ptr(),
.out_loc = out_loc.data_ptr(),
.v_stride_bytes = v_stride_bytes,
.cache_stride_bytes = cache_stride_bytes,
};
const auto is_int32 = id_type.is_type<int32_t>();
const auto kernel = is_int32 ? _kernel_1<int32_t> : _kernel_1<int64_t>;
const uint32_t kNumSM = get_num_sm(device.unwrap());
static const uint32_t kOccupancyTable[2] = {
runtime::get_blocks_per_sm(_kernel_1<int32_t>, kBlockSize),
runtime::get_blocks_per_sm(_kernel_1<int64_t>, kBlockSize),
};
const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM;
// rope works for q+k heads, plus v store works for kv heads
const auto num_total_works = (num_qo_heads + 2 * num_kv_heads) * num_tokens;
const auto needed_blocks = div_ceil(num_total_works, (kBlockSize / kWorkThreads));
const auto num_blocks = std::min(max_blocks, needed_blocks);
LaunchKernel(num_blocks, kBlockSize, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, store_params);
}
};
} // namespace

197
third_party/sglang/python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py vendored Normal file
View File

@@ -0,0 +1,197 @@
from pathlib import Path
import numpy as np
# From https://en.wikipedia.org/wiki/Paley_construction (construction II for q = 5)
had_12_paley = """
+-++++++++++
--+-+-+-+-+-
+++-++----++
+---+--+-++-
+++++-++----
+-+---+--+-+
++--+++-++--
+--++---+--+
++----+++-++
+--+-++---+-
++++----+++-
+-+--+-++---
"""
# From http://neilsloane.com/hadamard/
had_12 = """
+-----------
++-+---+++-+
+++-+---+++-
+-++-+---+++
++-++-+---++
+++-++-+---+
++++-++-+---
+-+++-++-+--
+--+++-++-+-
+---+++-++-+
++---+++-++-
+-+---+++-++
"""
had_20_will = """
+----+----++--++-++-
-+----+---+++---+-++
--+----+---+++-+-+-+
---+----+---+++++-+-
----+----++--++-++-+
-+++++-----+--+++--+
+-+++-+---+-+--+++--
++-++--+---+-+--+++-
+++-+---+---+-+--+++
++++-----++--+-+--++
--++-+-++-+-----++++
---++-+-++-+---+-+++
+---++-+-+--+--++-++
++---++-+----+-+++-+
-++---++-+----+++++-
-+--+--++-+----+----
+-+-----++-+----+---
-+-+-+---+--+----+--
--+-+++------+----+-
+--+--++------+----+
"""
had_28_will = """
+------++----++-+--+-+--++--
-+-----+++-----+-+--+-+--++-
--+-----+++---+-+-+----+--++
---+-----+++---+-+-+-+--+--+
----+-----+++---+-+-+++--+--
-----+-----++++--+-+--++--+-
------++----++-+--+-+--++--+
--++++-+-------++--+++-+--+-
---++++-+-----+-++--+-+-+--+
+---+++--+----++-++--+-+-+--
++---++---+----++-++--+-+-+-
+++---+----+----++-++--+-+-+
++++--------+-+--++-++--+-+-
-++++--------+++--++--+--+-+
-+-++-++--++--+--------++++-
+-+-++--+--++--+--------++++
-+-+-++--+--++--+----+---+++
+-+-+-++--+--+---+---++---++
++-+-+-++--+------+--+++---+
-++-+-+-++--+------+-++++---
+-++-+---++--+------+-++++--
-++--++-+-++-+++----++------
+-++--++-+-++-+++-----+-----
++-++---+-+-++-+++-----+----
-++-++-+-+-+-+--+++-----+---
--++-++++-+-+----+++-----+--
+--++-+-++-+-+----+++-----+-
++--++-+-++-+-+----++------+
"""
had_40_tpal = """
+-------------------+-------------------
++-++----+-+-++++--+++-++----+-+-++++--+
+++-++----+-+-++++--+++-++----+-+-++++--
+-++-++----+-+-++++-+-++-++----+-+-++++-
+--++-++----+-+-+++++--++-++----+-+-++++
++--++-++----+-+-+++++--++-++----+-+-+++
+++--++-++----+-+-+++++--++-++----+-+-++
++++--++-++----+-+-+++++--++-++----+-+-+
+++++--++-++----+-+-+++++--++-++----+-+-
+-++++--++-++----+-++-++++--++-++----+-+
++-++++--++-++----+-++-++++--++-++----+-
+-+-++++--++-++----++-+-++++--++-++----+
++-+-++++--++-++----++-+-++++--++-++----
+-+-+-++++--++-++---+-+-+-++++--++-++---
+--+-+-++++--++-++--+--+-+-++++--++-++--
+---+-+-++++--++-++-+---+-+-++++--++-++-
+----+-+-++++--++-+++----+-+-++++--++-++
++----+-+-++++--++-+++----+-+-++++--++-+
+++----+-+-++++--++-+++----+-+-++++--++-
+-++----+-+-++++--+++-++----+-+-++++--++
+--------------------+++++++++++++++++++
++-++----+-+-++++--+--+--++++-+-+----++-
+++-++----+-+-++++-----+--++++-+-+----++
+-++-++----+-+-++++--+--+--++++-+-+----+
+--++-++----+-+-++++-++--+--++++-+-+----
++--++-++----+-+-+++--++--+--++++-+-+---
+++--++-++----+-+-++---++--+--++++-+-+--
++++--++-++----+-+-+----++--+--++++-+-+-
+++++--++-++----+-+------++--+--++++-+-+
+-++++--++-++----+-+-+----++--+--++++-+-
++-++++--++-++----+---+----++--+--++++-+
+-+-++++--++-++----+-+-+----++--+--++++-
++-+-++++--++-++------+-+----++--+--++++
+-+-+-++++--++-++----+-+-+----++--+--+++
+--+-+-++++--++-++---++-+-+----++--+--++
+---+-+-++++--++-++--+++-+-+----++--+--+
+----+-+-++++--++-++-++++-+-+----++--+--
++----+-+-++++--++-+--++++-+-+----++--+-
+++----+-+-++++--++----++++-+-+----++--+
+-++----+-+-++++--++-+--++++-+-+----++--
"""
header = """
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// This file is auto-generated. See "code_gen.py"\n
#pragma once
"""
template = """
__device__ __forceinline__ void hadamard_mult_thread_{N}(float x[{N}]) {
float out[{N}];
{code}
#pragma unroll
for (int i = 0; i < {N}; i++) { x[i] = out[i]; }
}
"""
def string_to_array(string):
# Convert strings of + and - to bool arrays
string = string.strip().replace("+", "1").replace("-", "-1").split()
return np.stack(
[
np.fromstring(" ".join(string[i]), dtype=np.int32, sep=" ")
for i in range(len(string))
]
)
def array_code_gen(arr):
N = arr.shape[0]
assert arr.shape[0] == arr.shape[1]
out = []
for i in range(N):
out.append(
f"out[{i}] = "
+ " ".join([f"{'+' if arr[i, j] == 1 else '-'} x[{j}]" for j in range(N)])
+ ";"
)
return template.replace("{N}", str(N)).replace("{code}", "\n ".join(out))
def main():
output_dir = Path(__file__).parent / "fast_hadamard_transform_special.h"
output_dir.write_text(
header
+ array_code_gen(string_to_array(had_12_paley))
+ array_code_gen(string_to_array(had_20_will))
+ array_code_gen(string_to_array(had_28_will))
+ array_code_gen(string_to_array(had_40_tpal))
)
if __name__ == "__main__":
main()

24
third_party/sglang/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.h vendored Normal file
View File

@@ -0,0 +1,24 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Copied from https://github.com/sgl-project/fast-hadamard-transform
#pragma once
////////////////////////////////////////////////////////////////////////////////////////////////////
struct HadamardParamsBase {
using index_t = int64_t;
int batch, dim, log_N;
index_t x_batch_stride;
index_t out_batch_stride;
float scale;
// Common data pointers.
void* __restrict__ x_ptr;
void* __restrict__ out_ptr;
};

214
third_party/sglang/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_common.h vendored Normal file
View File

@@ -0,0 +1,214 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Copied from https://github.com/sgl-project/fast-hadamard-transform
#pragma once
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#define FULL_MASK 0xffffffff
////////////////////////////////////////////////////////////////////////////////////////////////////
struct uint8 {
uint4 u;
uint4 v;
};
template <int BYTES>
struct BytesToType {};
template <>
struct BytesToType<32> {
using Type = uint8;
static_assert(sizeof(Type) == 32);
};
template <>
struct BytesToType<16> {
using Type = uint4;
static_assert(sizeof(Type) == 16);
};
template <>
struct BytesToType<8> {
using Type = uint64_t;
static_assert(sizeof(Type) == 8);
};
template <>
struct BytesToType<4> {
using Type = uint32_t;
static_assert(sizeof(Type) == 4);
};
template <>
struct BytesToType<2> {
using Type = uint16_t;
static_assert(sizeof(Type) == 2);
};
template <>
struct BytesToType<1> {
using Type = uint8_t;
static_assert(sizeof(Type) == 1);
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
struct SumOp {
__device__ inline T operator()(T const& x, T const& y) {
return x + y;
}
};
template <int THREADS>
struct Allreduce {
static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
template <typename T, typename Operator>
static __device__ inline T run(T x, Operator& op) {
constexpr int OFFSET = THREADS / 2;
x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));
return Allreduce<OFFSET>::run(x, op);
}
};
template <>
struct Allreduce<2> {
template <typename T, typename Operator>
static __device__ inline T run(T x, Operator& op) {
x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));
return x;
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
// https://stackoverflow.com/questions/35311711/whats-the-right-way-to-compute-integral-base-2-logarithms-at-compile-time
constexpr int cilog2(int val) {
return val > 0 ? 1 + cilog2(val >> 1) : -1;
}
////////////////////////////////////////////////////////////////////////////////////////////////////
template <int kLogN, int kNChunks>
__device__ __forceinline__ void hadamard_mult_thread(float x[kNChunks][1 << kLogN]) {
constexpr int N = 1 << kLogN;
#pragma unroll
for (int i = 0; i < kLogN; ++i) {
const int stride = 1 << i;
#pragma unroll
for (int j = 0; j < N / 2; ++j) {
const int lo = j & (stride - 1);
const int idx = (j - lo) * 2 + lo;
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
const float a = x[c][idx];
const float b = x[c][idx + stride];
x[c][idx] = a + b;
x[c][idx + stride] = a - b;
}
}
}
}
template <int kLogWarpSize, int kStepStart, int kNChunks, int kNItems>
__device__ __forceinline__ void hadamard_mult_warp(float x[kNChunks][kNItems]) {
constexpr int N = 1 << kLogWarpSize;
int lane_id = threadIdx.x % N;
#pragma unroll
for (int step = kStepStart; step < kLogWarpSize; ++step) {
const int lane_mask = 1 << step;
const float sign = (lane_id & lane_mask) ? -1.f : 1.f;
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
float x_val_other = __shfl_xor_sync(FULL_MASK, x[c][i], lane_mask);
x[c][i] = sign * x[c][i] + x_val_other;
}
}
}
}
////////////////////////////////////////////////////////////////////////////////////////////////////
template <int kNChunks, int kNElts, typename input_t>
inline __device__ void load_input(input_t* x, float x_vals[kNChunks][kNElts], int dim) {
using vec_t = typename BytesToType<sizeof(input_t) * kNElts>::Type;
input_t x_vals_load[kNChunks][kNElts] = {0};
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
if ((c * blockDim.x + threadIdx.x) * kNElts < dim) {
reinterpret_cast<vec_t*>(x_vals_load)[c] = reinterpret_cast<const vec_t*>(x)[c * blockDim.x + threadIdx.x];
}
}
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
#pragma unroll
for (int i = 0; i < kNElts; ++i) {
x_vals[c][i] = float(x_vals_load[c][i]);
}
}
}
template <int kNChunks, int kNElts, typename output_t>
inline __device__ void store_output(output_t* out, float out_vals[kNChunks][kNElts], int dim, float scale = 1.f) {
using vec_t = typename BytesToType<sizeof(output_t) * kNElts>::Type;
output_t out_vals_store[kNChunks][kNElts];
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
#pragma unroll
for (int i = 0; i < kNElts; ++i) {
out_vals_store[c][i] = out_vals[c][i] * scale;
}
}
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
if ((c * blockDim.x + threadIdx.x) * kNElts < dim) {
reinterpret_cast<vec_t*>(out)[c * blockDim.x + threadIdx.x] = reinterpret_cast<const vec_t*>(out_vals_store)[c];
}
}
}
////////////////////////////////////////////////////////////////////////////////////////////////////
// Pre=true means the exchange before the hadamard_mult_warp, Pre=false means after.
template <int kNChunks, int kChunksPerExchange, int kNElts, int kWarpSize, int kNWarps, bool Pre, typename vec_t>
inline __device__ void exchange_smem_pre(float x_vals[kNChunks][kNElts], vec_t* smem) {
constexpr int kNThreads = kWarpSize * kNWarps;
constexpr int kNExchangePerVec = kNElts / (sizeof(vec_t) / sizeof(float));
const int warp_id = threadIdx.x / kWarpSize;
const int lane_id = threadIdx.x % kWarpSize;
const int row_t = threadIdx.x % kNWarps;
const int col_t = threadIdx.x / kNWarps;
// We use the XOR swizzle trick (new_col = col ^ row) to avoid / reduce smem bank conflicts.
#pragma unroll
for (int c0 = 0; c0 < kNChunks / kChunksPerExchange; ++c0) {
__syncthreads();
#pragma unroll
for (int c1 = 0; c1 < kChunksPerExchange; ++c1) {
#pragma unroll
for (int r = 0; r < kNExchangePerVec; ++r) {
smem
[(c1 * kNExchangePerVec + r) * kNThreads +
(Pre ? warp_id * kWarpSize + lane_id ^ warp_id : row_t * kWarpSize + col_t ^ row_t)] =
reinterpret_cast<vec_t*>(x_vals[c0 * kChunksPerExchange + c1])[r];
}
}
__syncthreads();
#pragma unroll
for (int c1 = 0; c1 < kChunksPerExchange; ++c1) {
#pragma unroll
for (int r = 0; r < kNExchangePerVec; ++r) {
reinterpret_cast<vec_t*>(x_vals[c0 * kChunksPerExchange + c1])[r] = smem
[(c1 * kNExchangePerVec + r) * kNThreads +
(Pre ? row_t * kWarpSize + col_t ^ row_t : warp_id * kWarpSize + lane_id ^ warp_id)];
}
}
}
}

298
third_party/sglang/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_special.h vendored Normal file
View File

@@ -0,0 +1,298 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Copied from https://github.com/sgl-project/fast-hadamard-transform
// This file is auto-generated. See "code_gen.py"
#pragma once
__device__ __forceinline__ void hadamard_mult_thread_12(float x[12]) {
float out[12];
out[0] = +x[0] - x[1] + x[2] + x[3] + x[4] + x[5] + x[6] + x[7] + x[8] + x[9] + x[10] + x[11];
out[1] = -x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] - x[11];
out[2] = +x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11];
out[3] = +x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] + x[7] - x[8] + x[9] + x[10] - x[11];
out[4] = +x[0] + x[1] + x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11];
out[5] = +x[0] - x[1] + x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] + x[9] - x[10] + x[11];
out[6] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11];
out[7] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] + x[11];
out[8] = +x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] + x[7] + x[8] - x[9] + x[10] + x[11];
out[9] = +x[0] - x[1] - x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11];
out[10] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11];
out[11] = +x[0] - x[1] + x[2] - x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11];
#pragma unroll
for (int i = 0; i < 12; i++) {
x[i] = out[i];
}
}
__device__ __forceinline__ void hadamard_mult_thread_20(float x[20]) {
float out[20];
out[0] = +x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11] - x[12] - x[13] +
x[14] + x[15] - x[16] + x[17] + x[18] - x[19];
out[1] = -x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] - x[9] + x[10] + x[11] + x[12] - x[13] -
x[14] - x[15] + x[16] - x[17] + x[18] + x[19];
out[2] = -x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] - x[9] - x[10] + x[11] + x[12] + x[13] -
x[14] + x[15] - x[16] + x[17] - x[18] + x[19];
out[3] = -x[0] - x[1] - x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] + x[13] +
x[14] + x[15] + x[16] - x[17] + x[18] - x[19];
out[4] = -x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] +
x[14] - x[15] + x[16] + x[17] - x[18] + x[19];
out[5] = -x[0] + x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] - x[13] +
x[14] + x[15] + x[16] - x[17] - x[18] + x[19];
out[6] = +x[0] - x[1] + x[2] + x[3] + x[4] - x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] -
x[14] + x[15] + x[16] + x[17] - x[18] - x[19];
out[7] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] + x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] -
x[14] - x[15] + x[16] + x[17] + x[18] - x[19];
out[8] = +x[0] + x[1] + x[2] - x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] - x[13] +
x[14] - x[15] - x[16] + x[17] + x[18] + x[19];
out[9] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] -
x[14] + x[15] - x[16] - x[17] + x[18] + x[19];
out[10] = -x[0] - x[1] + x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] + x[10] - x[11] - x[12] - x[13] -
x[14] - x[15] + x[16] + x[17] + x[18] + x[19];
out[11] = -x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] - x[13] -
x[14] + x[15] - x[16] + x[17] + x[18] + x[19];
out[12] = +x[0] - x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] - x[11] + x[12] - x[13] -
x[14] + x[15] + x[16] - x[17] + x[18] + x[19];
out[13] = +x[0] + x[1] - x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] -
x[14] + x[15] + x[16] + x[17] - x[18] + x[19];
out[14] = -x[0] + x[1] + x[2] - x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] +
x[14] + x[15] + x[16] + x[17] + x[18] - x[19];
out[15] = -x[0] + x[1] - x[2] - x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] - x[11] - x[12] - x[13] -
x[14] + x[15] - x[16] - x[17] - x[18] - x[19];
out[16] = +x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] - x[13] -
x[14] - x[15] + x[16] - x[17] - x[18] - x[19];
out[17] = -x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] - x[13] -
x[14] - x[15] - x[16] + x[17] - x[18] - x[19];
out[18] = -x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] + x[13] -
x[14] - x[15] - x[16] - x[17] + x[18] - x[19];
out[19] = +x[0] - x[1] - x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] +
x[14] - x[15] - x[16] - x[17] - x[18] + x[19];
#pragma unroll
for (int i = 0; i < 20; i++) {
x[i] = out[i];
}
}
__device__ __forceinline__ void hadamard_mult_thread_28(float x[28]) {
float out[28];
out[0] = +x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] +
x[14] - x[15] + x[16] - x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] + x[25] - x[26] -
x[27];
out[1] = -x[0] + x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] -
x[14] + x[15] - x[16] + x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] - x[24] + x[25] + x[26] -
x[27];
out[2] = -x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11] - x[12] - x[13] +
x[14] - x[15] + x[16] - x[17] + x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] + x[26] +
x[27];
out[3] = -x[0] - x[1] - x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] + x[11] - x[12] - x[13] -
x[14] + x[15] - x[16] + x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] - x[25] - x[26] +
x[27];
out[4] = -x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11] + x[12] - x[13] -
x[14] - x[15] + x[16] - x[17] + x[18] - x[19] + x[20] + x[21] + x[22] - x[23] - x[24] + x[25] - x[26] -
x[27];
out[5] = -x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] + x[12] + x[13] +
x[14] - x[15] - x[16] + x[17] - x[18] + x[19] - x[20] - x[21] + x[22] + x[23] - x[24] - x[25] + x[26] -
x[27];
out[6] = -x[0] - x[1] - x[2] - x[3] - x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] + x[13] -
x[14] + x[15] - x[16] - x[17] + x[18] - x[19] + x[20] - x[21] - x[22] + x[23] + x[24] - x[25] - x[26] +
x[27];
out[7] = -x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] + x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] -
x[14] + x[15] + x[16] - x[17] - x[18] + x[19] + x[20] + x[21] - x[22] + x[23] - x[24] - x[25] + x[26] -
x[27];
out[8] = -x[0] - x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] - x[12] - x[13] +
x[14] - x[15] + x[16] + x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] + x[24] - x[25] - x[26] +
x[27];
out[9] = +x[0] - x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] +
x[14] + x[15] - x[16] + x[17] + x[18] - x[19] - x[20] + x[21] - x[22] + x[23] - x[24] + x[25] - x[26] -
x[27];
out[10] = +x[0] + x[1] - x[2] - x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11] - x[12] - x[13] -
x[14] + x[15] + x[16] - x[17] + x[18] + x[19] - x[20] - x[21] + x[22] - x[23] + x[24] - x[25] + x[26] -
x[27];
out[11] = +x[0] + x[1] + x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] - x[13] -
x[14] - x[15] + x[16] + x[17] - x[18] + x[19] + x[20] - x[21] - x[22] + x[23] - x[24] + x[25] - x[26] +
x[27];
out[12] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] +
x[14] - x[15] - x[16] + x[17] + x[18] - x[19] + x[20] + x[21] - x[22] - x[23] + x[24] - x[25] + x[26] -
x[27];
out[13] = -x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] + x[13] +
x[14] + x[15] - x[16] - x[17] + x[18] + x[19] - x[20] - x[21] + x[22] - x[23] - x[24] + x[25] - x[26] +
x[27];
out[14] = -x[0] + x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] - x[13] +
x[14] - x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] + x[24] + x[25] + x[26] -
x[27];
out[15] = +x[0] - x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] -
x[14] + x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] - x[22] - x[23] + x[24] + x[25] + x[26] +
x[27];
out[16] = -x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] + x[13] -
x[14] - x[15] + x[16] - x[17] - x[18] - x[19] - x[20] + x[21] - x[22] - x[23] - x[24] + x[25] + x[26] +
x[27];
out[17] = +x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] - x[11] - x[12] + x[13] -
x[14] - x[15] - x[16] + x[17] - x[18] - x[19] - x[20] + x[21] + x[22] - x[23] - x[24] - x[25] + x[26] +
x[27];
out[18] = +x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] + x[11] - x[12] - x[13] -
x[14] - x[15] - x[16] - x[17] + x[18] - x[19] - x[20] + x[21] + x[22] + x[23] - x[24] - x[25] - x[26] +
x[27];
out[19] = -x[0] + x[1] + x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] + x[12] - x[13] -
x[14] - x[15] - x[16] - x[17] - x[18] + x[19] - x[20] + x[21] + x[22] + x[23] + x[24] - x[25] - x[26] -
x[27];
out[20] = +x[0] - x[1] + x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] -
x[14] - x[15] - x[16] - x[17] - x[18] - x[19] + x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] -
x[27];
out[21] = -x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] - x[12] + x[13] +
x[14] + x[15] - x[16] - x[17] - x[18] - x[19] + x[20] + x[21] - x[22] - x[23] - x[24] - x[25] - x[26] -
x[27];
out[22] = +x[0] - x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] - x[13] +
x[14] + x[15] + x[16] - x[17] - x[18] - x[19] - x[20] - x[21] + x[22] - x[23] - x[24] - x[25] - x[26] -
x[27];
out[23] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] -
x[14] + x[15] + x[16] + x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] - x[26] -
x[27];
out[24] = -x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] -
x[14] - x[15] + x[16] + x[17] + x[18] - x[19] - x[20] - x[21] - x[22] - x[23] + x[24] - x[25] - x[26] -
x[27];
out[25] = -x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] + x[10] - x[11] + x[12] - x[13] -
x[14] - x[15] - x[16] + x[17] + x[18] + x[19] - x[20] - x[21] - x[22] - x[23] - x[24] + x[25] - x[26] -
x[27];
out[26] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] + x[13] -
x[14] - x[15] - x[16] - x[17] + x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] + x[26] -
x[27];
out[27] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] - x[11] + x[12] - x[13] +
x[14] - x[15] - x[16] - x[17] - x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] - x[26] +
x[27];
#pragma unroll
for (int i = 0; i < 28; i++) {
x[i] = out[i];
}
}
__device__ __forceinline__ void hadamard_mult_thread_40(float x[40]) {
float out[40];
out[0] = +x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] -
x[14] - x[15] - x[16] - x[17] - x[18] - x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] - x[26] -
x[27] - x[28] - x[29] - x[30] - x[31] - x[32] - x[33] - x[34] - x[35] - x[36] - x[37] - x[38] - x[39];
out[1] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] +
x[14] + x[15] + x[16] - x[17] - x[18] + x[19] + x[20] + x[21] - x[22] + x[23] + x[24] - x[25] - x[26] -
x[27] - x[28] + x[29] - x[30] + x[31] - x[32] + x[33] + x[34] + x[35] + x[36] - x[37] - x[38] + x[39];
out[2] = +x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] +
x[14] + x[15] + x[16] + x[17] - x[18] - x[19] + x[20] + x[21] + x[22] - x[23] + x[24] + x[25] - x[26] -
x[27] - x[28] - x[29] + x[30] - x[31] + x[32] - x[33] + x[34] + x[35] + x[36] + x[37] - x[38] - x[39];
out[3] = +x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] -
x[14] + x[15] + x[16] + x[17] + x[18] - x[19] + x[20] - x[21] + x[22] + x[23] - x[24] + x[25] + x[26] -
x[27] - x[28] - x[29] - x[30] + x[31] - x[32] + x[33] - x[34] + x[35] + x[36] + x[37] + x[38] - x[39];
out[4] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] +
x[14] - x[15] + x[16] + x[17] + x[18] + x[19] + x[20] - x[21] - x[22] + x[23] + x[24] - x[25] + x[26] +
x[27] - x[28] - x[29] - x[30] - x[31] + x[32] - x[33] + x[34] - x[35] + x[36] + x[37] + x[38] + x[39];
out[5] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] -
x[14] + x[15] - x[16] + x[17] + x[18] + x[19] + x[20] + x[21] - x[22] - x[23] + x[24] + x[25] - x[26] +
x[27] + x[28] - x[29] - x[30] - x[31] - x[32] + x[33] - x[34] + x[35] - x[36] + x[37] + x[38] + x[39];
out[6] = +x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] +
x[14] - x[15] + x[16] - x[17] + x[18] + x[19] + x[20] + x[21] + x[22] - x[23] - x[24] + x[25] + x[26] -
x[27] + x[28] + x[29] - x[30] - x[31] - x[32] - x[33] + x[34] - x[35] + x[36] - x[37] + x[38] + x[39];
out[7] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] + x[10] - x[11] - x[12] - x[13] -
x[14] + x[15] - x[16] + x[17] - x[18] + x[19] + x[20] + x[21] + x[22] + x[23] - x[24] - x[25] + x[26] +
x[27] - x[28] + x[29] + x[30] - x[31] - x[32] - x[33] - x[34] + x[35] - x[36] + x[37] - x[38] + x[39];
out[8] = +x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] -
x[14] - x[15] + x[16] - x[17] + x[18] - x[19] + x[20] + x[21] + x[22] + x[23] + x[24] - x[25] - x[26] +
x[27] + x[28] - x[29] + x[30] + x[31] - x[32] - x[33] - x[34] - x[35] + x[36] - x[37] + x[38] - x[39];
out[9] = +x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] -
x[14] - x[15] - x[16] + x[17] - x[18] + x[19] + x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] -
x[27] + x[28] + x[29] - x[30] + x[31] + x[32] - x[33] - x[34] - x[35] - x[36] + x[37] - x[38] + x[39];
out[10] = +x[0] + x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] - x[11] + x[12] + x[13] -
x[14] - x[15] - x[16] - x[17] + x[18] - x[19] + x[20] + x[21] - x[22] + x[23] + x[24] + x[25] + x[26] -
x[27] - x[28] + x[29] + x[30] - x[31] + x[32] + x[33] - x[34] - x[35] - x[36] - x[37] + x[38] - x[39];
out[11] = +x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] + x[13] +
x[14] - x[15] - x[16] - x[17] - x[18] + x[19] + x[20] - x[21] + x[22] - x[23] + x[24] + x[25] + x[26] +
x[27] - x[28] - x[29] + x[30] + x[31] - x[32] + x[33] + x[34] - x[35] - x[36] - x[37] - x[38] + x[39];
out[12] = +x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] +
x[14] + x[15] - x[16] - x[17] - x[18] - x[19] + x[20] + x[21] - x[22] + x[23] - x[24] + x[25] + x[26] +
x[27] + x[28] - x[29] - x[30] + x[31] + x[32] - x[33] + x[34] + x[35] - x[36] - x[37] - x[38] - x[39];
out[13] = +x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] - x[11] + x[12] + x[13] -
x[14] + x[15] + x[16] - x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] + x[24] - x[25] + x[26] +
x[27] + x[28] + x[29] - x[30] - x[31] + x[32] + x[33] - x[34] + x[35] + x[36] - x[37] - x[38] - x[39];
out[14] = +x[0] - x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] +
x[14] - x[15] + x[16] + x[17] - x[18] - x[19] + x[20] - x[21] - x[22] + x[23] - x[24] + x[25] - x[26] +
x[27] + x[28] + x[29] + x[30] - x[31] - x[32] + x[33] + x[34] - x[35] + x[36] + x[37] - x[38] - x[39];
out[15] = +x[0] - x[1] - x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] + x[10] + x[11] - x[12] - x[13] +
x[14] + x[15] - x[16] + x[17] + x[18] - x[19] + x[20] - x[21] - x[22] - x[23] + x[24] - x[25] + x[26] -
x[27] + x[28] + x[29] + x[30] + x[31] - x[32] - x[33] + x[34] + x[35] - x[36] + x[37] + x[38] - x[39];
out[16] = +x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] - x[13] -
x[14] + x[15] + x[16] - x[17] + x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] + x[25] - x[26] +
x[27] - x[28] + x[29] + x[30] + x[31] + x[32] - x[33] - x[34] + x[35] + x[36] - x[37] + x[38] + x[39];
out[17] = +x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] -
x[14] - x[15] + x[16] + x[17] - x[18] + x[19] + x[20] + x[21] - x[22] - x[23] - x[24] - x[25] + x[26] -
x[27] + x[28] - x[29] + x[30] + x[31] + x[32] + x[33] - x[34] - x[35] + x[36] + x[37] - x[38] + x[39];
out[18] = +x[0] + x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] + x[13] +
x[14] - x[15] - x[16] + x[17] + x[18] - x[19] + x[20] + x[21] + x[22] - x[23] - x[24] - x[25] - x[26] +
x[27] - x[28] + x[29] - x[30] + x[31] + x[32] + x[33] + x[34] - x[35] - x[36] + x[37] + x[38] - x[39];
out[19] = +x[0] - x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] +
x[14] + x[15] - x[16] - x[17] + x[18] + x[19] + x[20] - x[21] + x[22] + x[23] - x[24] - x[25] - x[26] -
x[27] + x[28] - x[29] + x[30] - x[31] + x[32] + x[33] + x[34] + x[35] - x[36] - x[37] + x[38] + x[39];
out[20] = +x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] -
x[14] - x[15] - x[16] - x[17] - x[18] - x[19] - x[20] + x[21] + x[22] + x[23] + x[24] + x[25] + x[26] +
x[27] + x[28] + x[29] + x[30] + x[31] + x[32] + x[33] + x[34] + x[35] + x[36] + x[37] + x[38] + x[39];
out[21] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] +
x[14] + x[15] + x[16] - x[17] - x[18] + x[19] - x[20] - x[21] + x[22] - x[23] - x[24] + x[25] + x[26] +
x[27] + x[28] - x[29] + x[30] - x[31] + x[32] - x[33] - x[34] - x[35] - x[36] + x[37] + x[38] - x[39];
out[22] = +x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] +
x[14] + x[15] + x[16] + x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] + x[26] +
x[27] + x[28] + x[29] - x[30] + x[31] - x[32] + x[33] - x[34] - x[35] - x[36] - x[37] + x[38] + x[39];
out[23] = +x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] -
x[14] + x[15] + x[16] + x[17] + x[18] - x[19] - x[20] + x[21] - x[22] - x[23] + x[24] - x[25] - x[26] +
x[27] + x[28] + x[29] + x[30] - x[31] + x[32] - x[33] + x[34] - x[35] - x[36] - x[37] - x[38] + x[39];
out[24] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] +
x[14] - x[15] + x[16] + x[17] + x[18] + x[19] - x[20] + x[21] + x[22] - x[23] - x[24] + x[25] - x[26] -
x[27] + x[28] + x[29] + x[30] + x[31] - x[32] + x[33] - x[34] + x[35] - x[36] - x[37] - x[38] - x[39];
out[25] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] -
x[14] + x[15] - x[16] + x[17] + x[18] + x[19] - x[20] - x[21] + x[22] + x[23] - x[24] - x[25] + x[26] -
x[27] - x[28] + x[29] + x[30] + x[31] + x[32] - x[33] + x[34] - x[35] + x[36] - x[37] - x[38] - x[39];
out[26] = +x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] +
x[14] - x[15] + x[16] - x[17] + x[18] + x[19] - x[20] - x[21] - x[22] + x[23] + x[24] - x[25] - x[26] +
x[27] - x[28] - x[29] + x[30] + x[31] + x[32] + x[33] - x[34] + x[35] - x[36] + x[37] - x[38] - x[39];
out[27] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] + x[10] - x[11] - x[12] - x[13] -
x[14] + x[15] - x[16] + x[17] - x[18] + x[19] - x[20] - x[21] - x[22] - x[23] + x[24] + x[25] - x[26] -
x[27] + x[28] - x[29] - x[30] + x[31] + x[32] + x[33] + x[34] - x[35] + x[36] - x[37] + x[38] - x[39];
out[28] = +x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] -
x[14] - x[15] + x[16] - x[17] + x[18] - x[19] - x[20] - x[21] - x[22] - x[23] - x[24] + x[25] + x[26] -
x[27] - x[28] + x[29] - x[30] - x[31] + x[32] + x[33] + x[34] + x[35] - x[36] + x[37] - x[38] + x[39];
out[29] = +x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] -
x[14] - x[15] - x[16] + x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] - x[24] - x[25] + x[26] +
x[27] - x[28] - x[29] + x[30] - x[31] - x[32] + x[33] + x[34] + x[35] + x[36] - x[37] + x[38] - x[39];
out[30] = +x[0] + x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] - x[11] + x[12] + x[13] -
x[14] - x[15] - x[16] - x[17] + x[18] - x[19] - x[20] - x[21] + x[22] - x[23] - x[24] - x[25] - x[26] +
x[27] + x[28] - x[29] - x[30] + x[31] - x[32] - x[33] + x[34] + x[35] + x[36] + x[37] - x[38] + x[39];
out[31] = +x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] + x[13] +
x[14] - x[15] - x[16] - x[17] - x[18] + x[19] - x[20] + x[21] - x[22] + x[23] - x[24] - x[25] - x[26] -
x[27] + x[28] + x[29] - x[30] - x[31] + x[32] - x[33] - x[34] + x[35] + x[36] + x[37] + x[38] - x[39];
out[32] = +x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] +
x[14] + x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] + x[22] - x[23] + x[24] - x[25] - x[26] -
x[27] - x[28] + x[29] + x[30] - x[31] - x[32] + x[33] - x[34] - x[35] + x[36] + x[37] + x[38] + x[39];
out[33] = +x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] - x[11] + x[12] + x[13] -
x[14] + x[15] + x[16] - x[17] - x[18] - x[19] - x[20] + x[21] - x[22] + x[23] - x[24] + x[25] - x[26] -
x[27] - x[28] - x[29] + x[30] + x[31] - x[32] - x[33] + x[34] - x[35] - x[36] + x[37] + x[38] + x[39];
out[34] = +x[0] - x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] +
x[14] - x[15] + x[16] + x[17] - x[18] - x[19] - x[20] + x[21] + x[22] - x[23] + x[24] - x[25] + x[26] -
x[27] - x[28] - x[29] - x[30] + x[31] + x[32] - x[33] - x[34] + x[35] - x[36] - x[37] + x[38] + x[39];
out[35] = +x[0] - x[1] - x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] + x[10] + x[11] - x[12] - x[13] +
x[14] + x[15] - x[16] + x[17] + x[18] - x[19] - x[20] + x[21] + x[22] + x[23] - x[24] + x[25] - x[26] +
x[27] - x[28] - x[29] - x[30] - x[31] + x[32] + x[33] - x[34] - x[35] + x[36] - x[37] - x[38] + x[39];
out[36] = +x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] - x[13] -
x[14] + x[15] + x[16] - x[17] + x[18] + x[19] - x[20] + x[21] + x[22] + x[23] + x[24] - x[25] + x[26] -
x[27] + x[28] - x[29] - x[30] - x[31] - x[32] + x[33] + x[34] - x[35] - x[36] + x[37] - x[38] - x[39];
out[37] = +x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] -
x[14] - x[15] + x[16] + x[17] - x[18] + x[19] - x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] +
x[27] - x[28] + x[29] - x[30] - x[31] - x[32] - x[33] + x[34] + x[35] - x[36] - x[37] + x[38] - x[39];
out[38] = +x[0] + x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] + x[13] +
x[14] - x[15] - x[16] + x[17] + x[18] - x[19] - x[20] - x[21] - x[22] + x[23] + x[24] + x[25] + x[26] -
x[27] + x[28] - x[29] + x[30] - x[31] - x[32] - x[33] - x[34] + x[35] + x[36] - x[37] - x[38] + x[39];
out[39] = +x[0] - x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] +
x[14] + x[15] - x[16] - x[17] + x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] + x[25] + x[26] +
x[27] - x[28] + x[29] - x[30] + x[31] - x[32] - x[33] - x[34] - x[35] + x[36] + x[37] - x[38] - x[39];
#pragma unroll
for (int i = 0; i < 40; i++) {
x[i] = out[i];
}
}

482
third_party/sglang/python/sglang/jit_kernel/csrc/fast-hadamard-transform/hadamard_jit.cuh vendored Normal file
View File

@@ -0,0 +1,482 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
#pragma once
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <tvm/ffi/container/tensor.h>
#include "fast_hadamard_transform.h"
#include "fast_hadamard_transform_common.h"
#include "fast_hadamard_transform_special.h"
#include "static_switch.h"
#include <algorithm>
#include <cstdint>
#include <cstring>
namespace {
using ::bf16_t;
using ::fp16_t;
using ::HadamardParamsBase;
constexpr inline int ceil_log2(int val) {
int log = 0;
int p = 1;
while (p < val) {
p <<= 1;
++log;
}
return log;
}
template <int kNThreads_, int kLogN_, typename input_t_>
struct FastHadamardKernelTraits {
using input_t = input_t_;
static constexpr int kNThreads = kNThreads_;
static constexpr int kLogN = kLogN_;
static constexpr int N = 1 << kLogN;
static constexpr int kNBytes = sizeof(input_t);
static_assert(kNBytes == 2 || kNBytes == 4);
static constexpr int kNElts = kNBytes == 4 ? 4 : 8;
static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t);
using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
static constexpr int kNChunks = N / (kNElts * kNThreads);
static constexpr int kSmemExchangeSize = (N * 4) < (32 * 1024) ? (N * 4) : (32 * 1024);
static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize;
static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4);
static constexpr int kSmemSize = kSmemExchangeSize;
};
template <int kNThreads_, int kLogN_, int kMultiple, int kMaxDim, int kMaxSmem, typename input_t_>
struct FastHadamardMNKernelTraits {
using input_t = input_t_;
static constexpr int kNThreads = kNThreads_;
static constexpr int kLogN = kLogN_;
static constexpr int N = (1 << kLogN) * kMultiple;
static_assert(N <= kMaxDim);
static constexpr int kNBytes = sizeof(input_t);
static_assert(kNBytes == 2 || kNBytes == 4);
static constexpr int kNElts = 4;
static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t);
using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
static constexpr int kNChunks = N / (kNElts * kNThreads);
static_assert(kNChunks == kMultiple);
static constexpr int kSmemExchangeSize = (N * 4) < kMaxSmem ? (N * 4) : kMaxSmem;
static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize;
static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4);
static constexpr int kSmemSize = kSmemExchangeSize;
};
template <int kNThreads_, int kLogN_, typename input_t_>
using FastHadamard12NTraits = FastHadamardMNKernelTraits<kNThreads_, kLogN_, 12, 12 * 1024, 24 * 1024, input_t_>;
template <int kNThreads_, int kLogN_, typename input_t_>
using FastHadamard20NTraits = FastHadamardMNKernelTraits<kNThreads_, kLogN_, 20, 20 * 1024, 40 * 1024, input_t_>;
template <int kNThreads_, int kLogN_, typename input_t_>
using FastHadamard28NTraits = FastHadamardMNKernelTraits<kNThreads_, kLogN_, 28, 28 * 1024, 28 * 1024, input_t_>;
template <int kNThreads_, int kLogN_, typename input_t_>
using FastHadamard40NTraits = FastHadamardMNKernelTraits<kNThreads_, kLogN_, 40, 40 * 1024, 40 * 1024, input_t_>;
template <int kNChunks>
SGL_DEVICE void hadamard_mult_thread_chunk_12(float x[kNChunks][12]) {
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
hadamard_mult_thread_12(x[c]);
}
}
template <int kNChunks>
SGL_DEVICE void hadamard_mult_thread_chunk_20(float x[kNChunks][20]) {
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
hadamard_mult_thread_20(x[c]);
}
}
template <int kNChunks>
SGL_DEVICE void hadamard_mult_thread_chunk_28(float x[kNChunks][28]) {
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
hadamard_mult_thread_28(x[c]);
}
}
template <int kNChunks>
SGL_DEVICE void hadamard_mult_thread_chunk_40(float x[kNChunks][40]) {
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
hadamard_mult_thread_40(x[c]);
}
}
template <typename Ktraits>
__global__ __launch_bounds__(Ktraits::kNThreads) void fast_hadamard_transform_kernel(HadamardParamsBase params) {
constexpr int kNThreads = Ktraits::kNThreads;
constexpr int kNElts = Ktraits::kNElts;
constexpr int kNExchangePerVec = Ktraits::kNExchangePerVec;
constexpr int kNChunks = Ktraits::kNChunks;
using input_t = typename Ktraits::input_t;
using vec_t = typename Ktraits::vec_t;
constexpr int kLogNElts = cilog2(Ktraits::kNElts);
static_assert(1 << kLogNElts == kNElts, "kNElts must be a power of 2");
constexpr int kWarpSize = kNThreads < 32 ? kNThreads : 32;
constexpr int kLogWarpSize = cilog2(kWarpSize);
static_assert(1 << kLogWarpSize == kWarpSize, "Warp size must be a power of 2");
constexpr int kNWarps = kNThreads / kWarpSize;
constexpr int kLogNWarps = cilog2(kNWarps);
static_assert(1 << kLogNWarps == kNWarps, "kNWarps must be a power of 2");
constexpr int kChunksPerExchange = Ktraits::kSmemExchangeSize / (sizeof(vec_t) * kNExchangePerVec * kNThreads);
static_assert(kChunksPerExchange * sizeof(vec_t) * kNExchangePerVec * kNThreads == Ktraits::kSmemExchangeSize);
constexpr int kNExchanges = kNChunks / kChunksPerExchange;
static_assert(kNExchanges * kChunksPerExchange == kNChunks);
extern __shared__ char smem_[];
vec_t* smem_exchange = reinterpret_cast<vec_t*>(smem_);
const int batch_id = static_cast<int>(blockIdx.x);
input_t* x = reinterpret_cast<input_t*>(params.x_ptr) + batch_id * params.x_batch_stride;
input_t* out = reinterpret_cast<input_t*>(params.out_ptr) + batch_id * params.out_batch_stride;
float x_vals[kNChunks][kNElts];
load_input<kNChunks, kNElts, input_t>(x, x_vals, params.dim);
hadamard_mult_thread<kLogNElts, kNChunks>(x_vals);
hadamard_mult_warp<kLogWarpSize, 0, kNChunks, kNElts>(x_vals);
if constexpr (kNWarps > 1) {
exchange_smem_pre<kNChunks, kChunksPerExchange, kNElts, kWarpSize, kNWarps, true, vec_t>(x_vals, smem_exchange);
hadamard_mult_warp<kLogNWarps, 0, kNChunks, kNElts>(x_vals);
exchange_smem_pre<kNChunks, kChunksPerExchange, kNElts, kWarpSize, kNWarps, false, vec_t>(x_vals, smem_exchange);
}
if constexpr (kNChunks > 1) {
float x_vals_transposed[kNElts][kNChunks];
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
#pragma unroll
for (int i = 0; i < kNElts; ++i) {
x_vals_transposed[i][c] = x_vals[c][i];
}
}
if constexpr (kNChunks == 12) {
hadamard_mult_thread_chunk_12<kNElts>(x_vals_transposed);
} else if constexpr (kNChunks == 20) {
hadamard_mult_thread_chunk_20<kNElts>(x_vals_transposed);
} else if constexpr (kNChunks == 28) {
hadamard_mult_thread_chunk_28<kNElts>(x_vals_transposed);
} else if constexpr (kNChunks == 40) {
hadamard_mult_thread_chunk_40<kNElts>(x_vals_transposed);
} else {
constexpr int kLogNChunks = cilog2(kNChunks);
static_assert(1 << kLogNChunks == kNChunks, "kNChunks must be a power of 2");
hadamard_mult_thread<kLogNChunks, kNElts>(x_vals_transposed);
}
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
#pragma unroll
for (int i = 0; i < kNElts; ++i) {
x_vals[c][i] = x_vals_transposed[i][c];
}
}
}
store_output<kNChunks, kNElts, input_t>(out, x_vals, params.dim, params.scale);
}
template <typename Ktraits>
inline void set_max_dynamic_smem() {
constexpr int kSmemSize = Ktraits::kSmemSize;
if constexpr (kSmemSize >= 48 * 1024) {
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
host::RuntimeDeviceCheck(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
}
}
template <typename Ktraits>
inline void launch_kernel(HadamardParamsBase& params, DLDevice device) {
constexpr int kSmemSize = Ktraits::kSmemSize;
set_max_dynamic_smem<Ktraits>();
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
host::LaunchKernel(dim3(params.batch), dim3(Ktraits::kNThreads), device, kSmemSize)(kernel, params);
host::RuntimeDeviceCheck();
}
template <int kNThreads, int kLogN, typename input_t>
inline void fast_hadamard_transform_launch(HadamardParamsBase& params, DLDevice device) {
using Ktraits = FastHadamardKernelTraits<kNThreads, kLogN, input_t>;
launch_kernel<Ktraits>(params, device);
}
template <typename input_t>
inline void fast_hadamard_transform_cuda(HadamardParamsBase& params, DLDevice device) {
if (params.log_N == 3) {
fast_hadamard_transform_launch<1, 3, input_t>(params, device);
} else if (params.log_N == 4) {
fast_hadamard_transform_launch<2, 4, input_t>(params, device);
} else if (params.log_N == 5) {
fast_hadamard_transform_launch<4, 5, input_t>(params, device);
} else if (params.log_N == 6) {
fast_hadamard_transform_launch<8, 6, input_t>(params, device);
} else if (params.log_N == 7) {
fast_hadamard_transform_launch<16, 7, input_t>(params, device);
} else if (params.log_N == 8) {
fast_hadamard_transform_launch<32, 8, input_t>(params, device);
} else if (params.log_N == 9) {
fast_hadamard_transform_launch<32, 9, input_t>(params, device);
} else if (params.log_N == 10) {
fast_hadamard_transform_launch<128, 10, input_t>(params, device);
} else if (params.log_N == 11) {
fast_hadamard_transform_launch<256, 11, input_t>(params, device);
} else if (params.log_N == 12) {
fast_hadamard_transform_launch<256, 12, input_t>(params, device);
} else if (params.log_N == 13) {
fast_hadamard_transform_launch<256, 13, input_t>(params, device);
} else if (params.log_N == 14) {
fast_hadamard_transform_launch<256, 14, input_t>(params, device);
} else if (params.log_N == 15) {
fast_hadamard_transform_launch<256, 15, input_t>(params, device);
} else {
host::Panic("fast_hadamard_transform: unsupported log_N=", params.log_N);
}
}
template <int kNThreads, int kLogN, typename input_t>
inline void fast_hadamard_transform_12N_launch(HadamardParamsBase& params, DLDevice device) {
using Ktraits = FastHadamard12NTraits<kNThreads, kLogN, input_t>;
launch_kernel<Ktraits>(params, device);
}
template <typename input_t>
inline void fast_hadamard_transform_12N_cuda(HadamardParamsBase& params, DLDevice device) {
if (params.log_N == 2) {
fast_hadamard_transform_12N_launch<1, 2, input_t>(params, device);
} else if (params.log_N == 3) {
fast_hadamard_transform_12N_launch<2, 3, input_t>(params, device);
} else if (params.log_N == 4) {
fast_hadamard_transform_12N_launch<4, 4, input_t>(params, device);
} else if (params.log_N == 5) {
fast_hadamard_transform_12N_launch<8, 5, input_t>(params, device);
} else if (params.log_N == 6) {
fast_hadamard_transform_12N_launch<16, 6, input_t>(params, device);
} else if (params.log_N == 7) {
fast_hadamard_transform_12N_launch<32, 7, input_t>(params, device);
} else if (params.log_N == 8) {
fast_hadamard_transform_12N_launch<64, 8, input_t>(params, device);
} else if (params.log_N == 9) {
fast_hadamard_transform_12N_launch<128, 9, input_t>(params, device);
} else if (params.log_N == 10) {
fast_hadamard_transform_12N_launch<256, 10, input_t>(params, device);
} else {
host::Panic("fast_hadamard_transform_12N: unsupported log_N=", params.log_N);
}
}
template <int kNThreads, int kLogN, typename input_t>
inline void fast_hadamard_transform_20N_launch(HadamardParamsBase& params, DLDevice device) {
using Ktraits = FastHadamard20NTraits<kNThreads, kLogN, input_t>;
launch_kernel<Ktraits>(params, device);
}
template <typename input_t>
inline void fast_hadamard_transform_20N_cuda(HadamardParamsBase& params, DLDevice device) {
if (params.log_N == 2) {
fast_hadamard_transform_20N_launch<1, 2, input_t>(params, device);
} else if (params.log_N == 3) {
fast_hadamard_transform_20N_launch<2, 3, input_t>(params, device);
} else if (params.log_N == 4) {
fast_hadamard_transform_20N_launch<4, 4, input_t>(params, device);
} else if (params.log_N == 5) {
fast_hadamard_transform_20N_launch<8, 5, input_t>(params, device);
} else if (params.log_N == 6) {
fast_hadamard_transform_20N_launch<16, 6, input_t>(params, device);
} else if (params.log_N == 7) {
fast_hadamard_transform_20N_launch<32, 7, input_t>(params, device);
} else if (params.log_N == 8) {
fast_hadamard_transform_20N_launch<64, 8, input_t>(params, device);
} else if (params.log_N == 9) {
fast_hadamard_transform_20N_launch<128, 9, input_t>(params, device);
} else if (params.log_N == 10) {
fast_hadamard_transform_20N_launch<256, 10, input_t>(params, device);
} else {
host::Panic("fast_hadamard_transform_20N: unsupported log_N=", params.log_N);
}
}
template <int kNThreads, int kLogN, typename input_t>
inline void fast_hadamard_transform_28N_launch(HadamardParamsBase& params, DLDevice device) {
using Ktraits = FastHadamard28NTraits<kNThreads, kLogN, input_t>;
launch_kernel<Ktraits>(params, device);
}
template <typename input_t>
inline void fast_hadamard_transform_28N_cuda(HadamardParamsBase& params, DLDevice device) {
if (params.log_N == 2) {
fast_hadamard_transform_28N_launch<1, 2, input_t>(params, device);
} else if (params.log_N == 3) {
fast_hadamard_transform_28N_launch<2, 3, input_t>(params, device);
} else if (params.log_N == 4) {
fast_hadamard_transform_28N_launch<4, 4, input_t>(params, device);
} else if (params.log_N == 5) {
fast_hadamard_transform_28N_launch<8, 5, input_t>(params, device);
} else if (params.log_N == 6) {
fast_hadamard_transform_28N_launch<16, 6, input_t>(params, device);
} else if (params.log_N == 7) {
fast_hadamard_transform_28N_launch<32, 7, input_t>(params, device);
} else if (params.log_N == 8) {
fast_hadamard_transform_28N_launch<64, 8, input_t>(params, device);
} else if (params.log_N == 9) {
fast_hadamard_transform_28N_launch<128, 9, input_t>(params, device);
} else if (params.log_N == 10) {
fast_hadamard_transform_28N_launch<256, 10, input_t>(params, device);
} else {
host::Panic("fast_hadamard_transform_28N: unsupported log_N=", params.log_N);
}
}
template <int kNThreads, int kLogN, typename input_t>
inline void fast_hadamard_transform_40N_launch(HadamardParamsBase& params, DLDevice device) {
using Ktraits = FastHadamard40NTraits<kNThreads, kLogN, input_t>;
launch_kernel<Ktraits>(params, device);
}
template <typename input_t>
inline void fast_hadamard_transform_40N_cuda(HadamardParamsBase& params, DLDevice device) {
if (params.log_N == 2) {
fast_hadamard_transform_40N_launch<1, 2, input_t>(params, device);
} else if (params.log_N == 3) {
fast_hadamard_transform_40N_launch<2, 3, input_t>(params, device);
} else if (params.log_N == 4) {
fast_hadamard_transform_40N_launch<4, 4, input_t>(params, device);
} else if (params.log_N == 5) {
fast_hadamard_transform_40N_launch<8, 5, input_t>(params, device);
} else if (params.log_N == 6) {
fast_hadamard_transform_40N_launch<16, 6, input_t>(params, device);
} else if (params.log_N == 7) {
fast_hadamard_transform_40N_launch<32, 7, input_t>(params, device);
} else if (params.log_N == 8) {
fast_hadamard_transform_40N_launch<64, 8, input_t>(params, device);
} else if (params.log_N == 9) {
fast_hadamard_transform_40N_launch<128, 9, input_t>(params, device);
} else if (params.log_N == 10) {
fast_hadamard_transform_40N_launch<256, 10, input_t>(params, device);
} else {
host::Panic("fast_hadamard_transform_40N: unsupported log_N=", params.log_N);
}
}
inline void set_hadamard_params(
HadamardParamsBase& params,
int64_t batch,
int64_t dim,
int64_t multiple,
const tvm::ffi::TensorView x,
const tvm::ffi::TensorView out,
float scale) {
std::memset(&params, 0, sizeof(params));
params.batch = static_cast<int>(batch);
params.dim = static_cast<int>(dim);
params.log_N = ceil_log2(static_cast<int>(dim / multiple));
params.x_ptr = const_cast<void*>(x.data_ptr());
params.out_ptr = const_cast<void*>(out.data_ptr());
params.x_batch_stride = x.stride(0);
params.out_batch_stride = out.stride(0);
params.scale = scale;
}
template <int kMultiple, typename DType>
inline void run_hadamard(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) {
using namespace host;
auto N = SymbolicSize{"batch"};
auto D = SymbolicSize{"dim"};
auto SX = SymbolicSize{"x_batch_stride"};
auto SO = SymbolicSize{"out_batch_stride"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({N, D}).with_strides({SX, 1}).with_dtype<DType>().with_device(device).verify(x);
TensorMatcher({N, D}).with_strides({SO, 1}).with_dtype<DType>().with_device(device).verify(out);
const int64_t batch = N.unwrap();
const int64_t dim = D.unwrap();
RuntimeCheck(dim % kMultiple == 0, "hadamard: dim must be divisible by ", kMultiple);
HadamardParamsBase params;
set_hadamard_params(params, batch, dim, kMultiple, x, out, scale);
if constexpr (kMultiple == 1) {
RuntimeCheck(dim % 8 == 0, "fast_hadamard_transform only supports hidden dim divisible by 8");
RuntimeCheck(dim <= 32768, "fast_hadamard_transform only supports hidden dim <= 32768");
fast_hadamard_transform_cuda<DType>(params, device.unwrap());
} else if constexpr (kMultiple == 12) {
RuntimeCheck(dim % (4 * 12) == 0, "fast_hadamard_transform_12N only supports hidden dim divisible by 48");
RuntimeCheck(dim <= 12 * 1024, "fast_hadamard_transform_12N only supports hidden dim <= 12288");
fast_hadamard_transform_12N_cuda<DType>(params, device.unwrap());
} else if constexpr (kMultiple == 20) {
RuntimeCheck(dim % (4 * 20) == 0, "fast_hadamard_transform_20N only supports hidden dim divisible by 80");
RuntimeCheck(dim <= 20 * 1024, "fast_hadamard_transform_20N only supports hidden dim <= 20480");
fast_hadamard_transform_20N_cuda<DType>(params, device.unwrap());
} else if constexpr (kMultiple == 28) {
RuntimeCheck(dim % (4 * 28) == 0, "fast_hadamard_transform_28N only supports hidden dim divisible by 112");
RuntimeCheck(dim <= 28 * 1024, "fast_hadamard_transform_28N only supports hidden dim <= 28672");
fast_hadamard_transform_28N_cuda<DType>(params, device.unwrap());
} else if constexpr (kMultiple == 40) {
RuntimeCheck(dim % (4 * 40) == 0, "fast_hadamard_transform_40N only supports hidden dim divisible by 160");
RuntimeCheck(dim <= 40 * 1024, "fast_hadamard_transform_40N only supports hidden dim <= 40960");
fast_hadamard_transform_40N_cuda<DType>(params, device.unwrap());
} else {
Panic("Unsupported multiple");
}
}
template <typename DType>
struct HadamardKernel {
static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) {
run_hadamard<1, DType>(x, out, scale);
}
};
template <typename DType>
struct Hadamard12NKernel {
static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) {
run_hadamard<12, DType>(x, out, scale);
}
};
template <typename DType>
struct Hadamard20NKernel {
static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) {
run_hadamard<20, DType>(x, out, scale);
}
};
template <typename DType>
struct Hadamard28NKernel {
static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) {
run_hadamard<28, DType>(x, out, scale);
}
};
template <typename DType>
struct Hadamard40NKernel {
static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) {
run_hadamard<40, DType>(x, out, scale);
}
};
} // namespace

27
third_party/sglang/python/sglang/jit_kernel/csrc/fast-hadamard-transform/static_switch.h vendored Normal file
View File

@@ -0,0 +1,27 @@
// Copied from https://github.com/sgl-project/fast-hadamard-transform
// Inspired by https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h
#pragma once
/// @param COND - a boolean expression to switch by
/// @param CONST_NAME - a name given for the constexpr bool variable.
/// @param ... - code to execute for true and false
///
/// Usage:
/// ```
/// BOOL_SWITCH(flag, BoolConst, [&] {
/// some_function<BoolConst>(...);
/// });
/// ```
#define BOOL_SWITCH(COND, CONST_NAME, ...) \
[&] { \
if (COND) { \
static constexpr bool CONST_NAME = true; \
return __VA_ARGS__(); \
} else { \
static constexpr bool CONST_NAME = false; \
return __VA_ARGS__(); \
} \
}()

Some files were not shown because too many files have changed in this diff Show More