wehub/sgl-project--sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
PR Test (NPU) / check-changes (push) Has been cancelled
Details
PR Test (NPU) / pr-gate (push) Has been cancelled
Details
PR Test (NPU) / set-image-config (push) Has been cancelled
Details
PR Test (NPU) / stage-b-test-1-npu-a2 (0) (push) Has been cancelled
Details
PR Test (NPU) / stage-b-test-1-npu-a2 (1) (push) Has been cancelled
Details
PR Test (NPU) / stage-b-test-2-npu-a2 (0) (push) Has been cancelled
Details
PR Test (NPU) / stage-b-test-2-npu-a2 (1) (push) Has been cancelled
Details
PR Test (NPU) / stage-b-test-4-npu-a3 (push) Has been cancelled
Details
PR Test (NPU) / stage-b-test-16-npu-a3 (push) Has been cancelled
Details
PR Test (NPU) / multimodal-gen-test-1-npu-a3 (push) Has been cancelled
Details
PR Test (NPU) / multimodal-gen-test-2-npu-a3 (push) Has been cancelled
Details
PR Test (Arm64) / pr-gate (push) Has been cancelled
Details
PR Test (Arm64) / check-changes (push) Has been cancelled
Details
PR Test (Arm64) / build-test (push) Has been cancelled
Details
PR Test (sgl-router) / gate (push) Has been cancelled
Details
PR Test (sgl-router) / tier-1 — lint (push) Has been cancelled
Details
PR Test (sgl-router) / tier-2 — build + test (push) Has been cancelled
Details
PR Test (sgl-router) / tier-3 — docker (placeholder) (push) Has been cancelled
Details
PR Test (sgl-router) / tier-3 — k8s integration (push) Has been cancelled
Details
PR Test (sgl-router) / tier-3 — e2e (push) Has been cancelled
Details
PR Test (sgl-router) / finish (push) Has been cancelled
Details
PR Test (NPU) / single-node-poc (map[name:qwen3_6_27b_w8a8_1p_in64k_out1k_50ms runner:linux-aarch64-a3-2 test_case:test/registered/ascend/performance/qwen3_6_27b/test_npu_qwen3_6_27b_w8a8_1p_in64k_out1k_50ms.py test_type:perf]) (push) Has been cancelled
Details
PR Test (NPU) / pr-test-npu-finish (push) Has been cancelled
Details
PR Test (Xeon) / pr-gate (push) Has been cancelled
Details
PR Test (Xeon) / check-changes (push) Has been cancelled
Details
PR Test (Xeon) / build-test (, xeon-gnr, base-b-test-cpu) (push) Has been cancelled
Details
PR Test (XPU) / check-changes (push) Has been cancelled
Details
PR Test (XPU) / pr-gate (push) Has been cancelled
Details
PR Test (XPU) / stage-a-test-1-gpu-xpu (push) Has been cancelled
Details
PR Test (XPU) / wait-for-stage-a (push) Has been cancelled
Details
PR Test (XPU) / stage-b-test-1-gpu-xpu (push) Has been cancelled
Details
PR Test (XPU) / finish (push) Has been cancelled
Details
CI Model Inventory / build-inventory (push) Has been cancelled
Details
Lint / lint (push) Has been cancelled
Details
PR Benchmark (SMG Components) / Benchmark Compilation Check (push) Has been cancelled
Details
PR Benchmark (SMG Components) / Benchmark - Manual Policy (push) Has been cancelled
Details
PR Benchmark (SMG Components) / Benchmark - Request Processing (push) Has been cancelled
Details
PR Benchmark (SMG Components) / Benchmark Summary (push) Has been cancelled
Details
PR Test (SMG) / build-wheel (push) Has been cancelled
Details
Release SGLang Model Gateway to PyPI / build on windows (x86_64 - auto) (push) Has been cancelled
Details
Release SGLang Model Gateway to PyPI / build on macos (x86_64 - auto) (push) Has been cancelled
Details
PR Test (SMG) / python-unit-tests (push) Has been cancelled
Details
PR Test (SMG) / unit-tests (push) Has been cancelled
Details
PR Test (SMG) / benchmarks (push) Has been cancelled
Details
PR Test (SMG) / chat-completions (push) Has been cancelled
Details
PR Test (SMG) / chat-completions-4gpu (push) Has been cancelled
Details
PR Test (SMG) / e2e (push) Has been cancelled
Details
PR Test (SMG) / docker-build-test (push) Has been cancelled
Details
PR Test (SMG) / k8s-integration (push) Has been cancelled
Details
PR Test (SMG) / finish (push) Has been cancelled
Details
PR Test (SMG) / summarize-benchmarks (push) Has been cancelled
Details
Release SGLang Model Gateway Docker Image / publish (push) Has been cancelled
Details
Release SGLang Model Gateway to PyPI / build on macos (aarch64 - auto) (push) Has been cancelled
Details
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - auto) (push) Has been cancelled
Details
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - auto) (push) Has been cancelled
Details
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - musllinux_1_1) (push) Has been cancelled
Details
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - musllinux_1_1) (push) Has been cancelled
Details
Release SGLang Model Gateway to PyPI / Build SDist (push) Has been cancelled
Details
Release SGLang Model Gateway to PyPI / Upload to PyPI (push) Has been cancelled
Details
Release SGLang Kernels / build-cu129-matrix (aarch64, 12.9, 3.10, arm-kernel-build-node) (push) Has been cancelled
Details
Release SGLang Kernels / build-cu129-matrix (x86_64, 12.9, 3.10, x64-kernel-build-node) (push) Has been cancelled
Details
Release SGLang Kernels / release-cu129 (push) Has been cancelled
Details
Release SGLang Kernels / build-cu130-matrix (aarch64, 13.0, 3.10, arm-kernel-build-node) (push) Has been cancelled
Details
Release SGLang Kernels / build-cu130-matrix (x86_64, 13.0, 3.10, x64-kernel-build-node) (push) Has been cancelled
Details
Release SGLang Kernels / release-cu130 (push) Has been cancelled
Details
Release SGLang Kernels / build-rocm-matrix (3.10, 700) (push) Has been cancelled
Details
Release SGLang Kernels / build-rocm-matrix (3.10, 720) (push) Has been cancelled
Details
Release SGLang Kernels / release-rocm700 (push) Has been cancelled
Details
Release SGLang Kernels / release-rocm720 (push) Has been cancelled
Details
Release SGLang Kernels / build-musa43 (43, 3.10) (push) Has been cancelled
Details
Release SGLang Kernels / release-musa43 (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:38:16 +08:00
commit 94057c3d3e
7152 changed files with 2120455 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/pyproject.toml
Executable
+231
View File
@@ -0,0 +1,231 @@
[build-system]
requires = ["setuptools>=61.0", "setuptools-rust>=1.10", "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",
]
# Please keep dependency lists in this file sorted alphabetically by package name.
dependencies = [
"aiohttp",
"anthropic>=0.20.0",
"apache-tvm-ffi==0.1.11",
"av ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')",
"blobfile==3.0.0",
"build",
"compressed-tensors",
"cuda-python>=13.0",
"datasets",
"decord2 ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')",
"distro",
"easydict", # Required by remote model code (e.g. DeepSeek-OCR) loaded via trust_remote_code; validated by transformers 5.4+ check_imports
"einops",
"fastapi",
"flash-attn-4==4.0.0b15",
"flashinfer_python[cu13]==0.6.14", # keep it aligned with jit-cache version in Dockerfile
"gguf",
"interegular",
"IPython",
"kernels>=0.14.1,<0.15",
"llguidance>=0.7.11,<0.8.0",
"mistral_common>=1.11.5",
"modelscope",
"msgspec",
"ninja",
"numpy",
"nvidia-cutlass-dsl[cu13]==4.5.2",
"nvidia-mathdx==25.6.0",
"nvidia-ml-py",
"openai==2.6.1",
"openai-harmony==0.0.4",
"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.4.1",
"requests",
"scipy",
"sentencepiece",
"setproctitle",
"sgl-deep-gemm==0.1.4.post1",
"sglang-kernel==0.4.4",
"smg-grpc-servicer>=0.5.0",
"soundfile==0.13.1",
"tiktoken",
"tilelang==0.1.11",
"timm==1.0.16",
"tokenspeed_mla==0.1.7",
"torch==2.11.0",
"torch_memory_saver>=0.0.9.post1",
"torchao==0.17.0",
"torchaudio==2.11.0",
"torchcodec==0.11.1 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torchcodec 0.11.1 for torch 2.11.x (0.10 is ABI-incompatible: references the pre-2.11 c10::MessageLogger ctor signature). Not available on Linux ARM.
"torchvision",
"tqdm",
"transformers==5.12.1",
"uvicorn",
"uvloop",
"watchfiles",
"xgrammar==0.2.1",
"zstandard",
]
[[tool.uv.index]]
name = "pypi"
url = "https://pypi.org/simple"
default = true
[project.optional-dependencies]
checkpoint-engine = ["checkpoint-engine==0.1.2"]
runai = ["runai-model-streamer[s3,gcs,azure]>=0.15.7"]
diffusion = [
"addict==2.4.0",
"av==16.1.0",
"cache-dit==1.3.0",
"cloudpickle==3.1.2",
"diffusers==0.37.0",
"imageio==2.36.0",
"imageio-ffmpeg==0.5.1",
"moviepy>=2.0.0",
"msgpack",
"nvidia-modelopt",
"opencv-python-headless==4.10.0.84",
"PyYAML==6.0.1",
"remote-pdb==2.1.0",
"runai_model_streamer>=0.15.7",
"scikit-image==0.25.2",
"st_attn==0.0.7 ; platform_machine != 'aarch64' and platform_machine != 'arm64'",
"trimesh>=4.0.0",
"vsa==0.0.4 ; platform_machine != 'aarch64' and platform_machine != 'arm64'",
"websockets",
"xatlas",
]
ray = [
"ray[default]>=2.55.1",
]
tracing = [
"opentelemetry-api",
"opentelemetry-exporter-otlp",
"opentelemetry-exporter-otlp-proto-grpc",
"opentelemetry-sdk",
]
http2 = [
"granian>=2.6.0",
]
fastokens = [
"fastokens>=0.1.1,<0.2.0",
]
test = [
"accelerate",
"addict",
"auto-round>=0.13.1",
"bitsandbytes",
"pymupdf",
"diff-cover",
"expecttest",
"granian>=2.6.0",
"jsonlines",
"lm-eval[api]>=0.4.9.2",
"matplotlib",
# Pin sgl-eval to a git SHA: upgrading changes zero-shot \boxed{} grading, so
# re-baseline MODEL_SCORE_THRESHOLDS in test_text_models_gsm8k_eval.py first.
# antlr4 4.9.3 is forced because latex2sympy2_extended raises ImportError on
# 4.7.x, and an older transitive pin can win during install.
"antlr4-python3-runtime==4.9.3",
"sgl-eval @ git+https://github.com/sgl-project/sgl-eval.git@b2a2703c42cae379bbcb8b7ff092df6601a61694",
"pandas",
"parameterized",
"peft>=0.18.0",
"polars",
"pytest",
"pytest-cov",
"sentence_transformers",
"sglang[fastokens]",
"tabulate",
]
dev = ["sglang[test]"]
all = [
"sglang[diffusion]",
"sglang[http2]",
"sglang[tracing]",
]
[tool.uv.extra-build-dependencies]
st-attn = ["setuptools", "torch"]
vsa = ["setuptools", "torch"]
[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/**/*",
"multimodal_gen/apps/realtime_webui/**/*"
]
[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 = ["python3", "python/tools/get_version_tag.py"]
# Allow editable installs even when .git metadata is not available.
fallback_version = "0.0.0.dev0"
[[tool.setuptools-rust.ext-modules]]
target = "sglang.srt.grpc._core"
path = "../rust/sglang-grpc/Cargo.toml"
binding = "PyO3"
[tool.kernels.dependencies]
"kernels-community/sgl-flash-attn3" = 1
python/pyproject_cpu.toml
+156
View File
@@ -0,0 +1,156 @@
# 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",
]
# Please keep dependency lists in this file sorted alphabetically by package name.
dependencies = [
"aiohttp",
"anthropic>=0.20.0",
"blobfile==3.0.0",
"build",
"compressed-tensors",
"datasets",
"easydict",
"einops",
"fastapi",
"gguf",
"intel-openmp; platform_machine == 'x86_64'",
"interegular",
"IPython",
"llguidance>=0.7.11,<0.8.0",
"mistral_common>=1.11.5",
"modelscope",
"msgspec",
"ninja",
"numpy",
"openai==2.6.1",
"openai-harmony==0.0.4",
"orjson",
"outlines",
"packaging",
"partial_json_parser",
"pillow",
"prometheus-client>=0.20.0",
"psutil",
"py-spy",
"pybase64",
"pydantic",
"python-multipart",
"pytest",
"pyzmq>=25.1.2",
"requests",
"scipy",
"sentencepiece",
"setproctitle",
"smg-grpc-servicer>=0.5.0",
"soundfile==0.13.1",
"tabulate",
"tiktoken",
"timm==1.0.16",
"torch==2.12.0",
"torchao==0.17.0",
"torchaudio==2.11.0",
"torchvision==0.27.0",
"tqdm",
"transformers==5.12.1",
"triton==3.7.0",
"uvicorn",
"uvloop",
"xgrammar==0.2.1",
"zstandard",
]
[project.optional-dependencies]
diffusion = [
"addict==2.4.0",
"av==16.1.0",
"cache-dit==1.3.0",
"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",
"PyYAML==6.0.1",
"remote-pdb==2.1.0",
"runai_model_streamer>=0.15.5",
"scikit-image==0.25.2",
"st_attn==0.0.7 ; platform_machine != 'aarch64' and platform_machine != 'arm64'",
"trimesh>=4.0.0",
"vsa==0.0.4 ; platform_machine != 'aarch64' and platform_machine != 'arm64'",
"xatlas",
]
tracing = [
"opentelemetry-api",
"opentelemetry-exporter-otlp",
"opentelemetry-exporter-otlp-proto-grpc",
"opentelemetry-sdk",
]
test = [
"accelerate",
"pymupdf",
"expecttest",
"jsonlines",
"matplotlib",
"pandas",
"peft>=0.18.0",
"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 = ["python3", "python/tools/get_version_tag.py"]
# Allow editable installs even when .git metadata is not available.
fallback_version = "0.0.0.dev0"
python/pyproject_npu.toml
+155
View File
@@ -0,0 +1,155 @@
[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",
]
# Please keep dependency lists in this file sorted alphabetically by package name.
dependencies = [
"aiohttp",
"anthropic>=0.20.0",
"av",
"blobfile==3.0.0",
"build",
"compressed-tensors",
"datasets",
"decord2 ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')",
"easydict",
"einops",
"fastapi",
"gguf",
"hf_transfer",
"huggingface_hub",
"interegular",
"IPython",
"llguidance>=0.7.11,<0.8.0",
"mistral_common>=1.11.5",
"modelscope",
"msgspec",
"ninja",
"numpy",
"openai==2.6.1",
"openai-harmony==0.0.4",
"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",
"smg-grpc-servicer>=0.5.0",
"soundfile==0.13.1",
"tiktoken",
"timm==1.0.16",
"torchao==0.9.0",
"tqdm",
"transformers==5.12.1",
"uvicorn",
"uvloop",
"xgrammar==0.2.1",
]
[project.optional-dependencies]
checkpoint-engine = ["checkpoint-engine==0.1.2"]
diffusion = [
"addict",
"cache-dit==1.3.5",
"cloudpickle",
"diffusers==0.37.0",
"imageio==2.36.0",
"imageio-ffmpeg==0.5.1",
"moviepy>=2.0.0",
"opencv-python==4.10.0.84",
"PyYAML==6.0.1",
"remote-pdb",
"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",
"pymupdf",
"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 = ["python3", "python/tools/get_version_tag.py"]
# Allow editable installs even when .git metadata is not available.
fallback_version = "0.0.0.dev0"
python/pyproject_other.toml
Executable
+227
View File
@@ -0,0 +1,227 @@
[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",
]
# Please keep dependency lists in this file sorted alphabetically by package name.
dependencies = ["aiohttp", "IPython", "numpy", "requests", "setproctitle", "tqdm"]
[project.optional-dependencies]
runtime_common = [
"aiohttp",
"anthropic>=0.20.0",
"apache-tvm-ffi",
"av",
"blobfile==3.0.0",
"build",
"compressed-tensors",
"datasets",
"easydict",
"einops",
"fastapi",
"gguf",
"interegular",
"IPython",
"llguidance>=0.7.11,<0.8.0",
"mistral_common>=1.11.5",
"modelscope",
"msgspec",
"ninja",
"numpy",
"openai==2.6.1",
"openai-harmony==0.0.4",
"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",
"smg-grpc-servicer>=0.5.0",
"soundfile==0.13.1",
"tiktoken",
"timm==1.0.16",
"torchao==0.9.0",
"tqdm",
"transformers==5.12.1",
"uvicorn",
"uvloop",
"xgrammar==0.2.1",
]
diffusion_common = [
"addict",
"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",
"PyYAML==6.0.1",
"remote-pdb",
"scikit-image==0.25.2",
"trimesh>=4.0.0",
"xatlas",
]
tracing = [
"opentelemetry-api",
"opentelemetry-exporter-otlp",
"opentelemetry-exporter-otlp-proto-grpc",
"opentelemetry-sdk",
]
# HIP (Heterogeneous-computing Interface for Portability) for AMD
# => base docker rocm/vllm-dev:20250114, not from public vllm whl
srt_hip = [
# Pin to 0.15.0: 0.16.0 needs torch>=2.10 (incompatible with ROCm torch
# 2.9.1). An open-ended `<0.16.0` made pip backtrack into an unbuildable
# ancient setuptools sdist; an exact pin keeps the resolver converging.
"compressed-tensors==0.15.0",
"petit_kernel==0.0.2",
"sglang[runtime_common]",
"torch",
"wave-lang==3.8.2",
]
diffusion_hip = [
"cache-dit==1.3.0",
"peft>=0.18.0,<0.19.0", # Pin to <0.19.0 due to torchao incompatibility
"runai_model_streamer>=0.15.5",
"sglang[diffusion_common]",
"st_attn==0.0.7",
"vsa==0.0.4",
]
# 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 = [
"deep-gemm>=0.1.3",
"flash_attn_3>=0.1.4",
"mate>=0.2.0",
"mthreads-ml-py",
"numpy<2.0",
"sglang[runtime_common]",
"torch",
"torch_musa",
"torchada>=0.1.68",
]
diffusion_musa = [
"cache-dit==1.1.8",
"runai_model_streamer>=0.15.5",
"sglang[diffusion_common]",
"st_attn==0.0.7",
"vsa==0.0.4",
]
# https://docs.sglang.io/platforms/mps.md
srt_mps = [
"mlx",
"mlx-lm",
"sglang[runtime_common]",
"torch==2.11.0",
"torchao==0.9.0",
"torchaudio==2.11.0",
"torchvision",
]
diffusion_mps = [
"addict==2.4.0",
"av==16.1.0",
"cache-dit==1.2.3",
"cloudpickle==3.1.2",
"remote-pdb==2.1.0",
"scikit-image==0.25.2",
"sglang[diffusion_common]",
"trimesh>=4.0.0",
"xatlas",
]
test = [
"accelerate",
"pymupdf",
"expecttest",
"gguf",
"jsonlines",
"matplotlib",
"pandas",
"peft>=0.18.0,<0.19.0", # Pin to <0.19.0 due to torchao incompatibility
"pytest",
"sentence_transformers",
"tabulate",
]
all_hip = ["sglang[diffusion_hip]", "sglang[srt_hip]", "sglang[tracing]"]
all_hpu = ["sglang[srt_hpu]"]
all_musa = ["sglang[diffusion_musa]", "sglang[srt_musa]"]
all_mps = ["sglang[diffusion_mps]", "sglang[srt_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 = ["python3", "python/tools/get_version_tag.py"]
# Allow editable installs even when .git metadata is not available.
fallback_version = "0.0.0.dev0"
python/pyproject_xpu.toml
+163
View File
@@ -0,0 +1,163 @@
[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",
]
# Please keep dependency lists in this file sorted alphabetically by package name.
dependencies = [
"addict",
"aiohttp",
"anthropic>=0.20.0",
"av ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')",
"blobfile==3.0.0",
"build",
"compressed-tensors",
"datasets",
"easydict",
"einops",
"fastapi",
"gguf",
"interegular",
"IPython",
"llguidance>=0.7.11,<0.8.0",
"mistral_common>=1.11.5",
"modelscope",
"msgspec",
"ninja",
"numpy",
"openai==2.6.1",
"openai-harmony==0.0.4",
"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",
"sgl-kernel @ git+https://github.com/sgl-project/sgl-kernel-xpu.git",
"smg-grpc-servicer>=0.5.0",
"soundfile==0.13.1",
"tiktoken",
"timm==1.0.16",
"torch==2.12.0+xpu",
"torchao==0.17.0+xpu",
"torchaudio==2.11.0+xpu",
"torchcodec==0.12.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torch==2.12.0 on XPU uses torchcodec 0.12.0
"torchvision==0.27.0+xpu",
"tqdm",
"transformers==5.12.1",
"uvicorn",
"uvloop",
# "xgrammar==0.2.1", xgrammar depends on CUDA PyTorch and Triton only
]
[project.optional-dependencies]
diffusion = [
"addict==2.4.0",
"av==16.1.0",
"cache-dit==1.3.0",
"cloudpickle==3.1.2",
"diffusers==0.38.0",
"imageio==2.36.0",
"imageio-ffmpeg==0.5.1",
"moviepy>=2.0.0",
"opencv-python==4.10.0.84",
"PyYAML==6.0.1",
"remote-pdb==2.1.0",
"runai_model_streamer>=0.15.5",
"scikit-image==0.25.2",
"st_attn==0.0.7 ; platform_machine != 'aarch64' and platform_machine != 'arm64'",
"trimesh>=4.0.0",
"xatlas",
]
tracing = [
"opentelemetry-api",
"opentelemetry-exporter-otlp",
"opentelemetry-exporter-otlp-proto-grpc",
"opentelemetry-sdk",
]
test = [
"accelerate",
"bitsandbytes",
"pymupdf",
"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[diffusion]",
"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 = ["python3", "python/tools/get_version_tag.py"]
# Allow editable installs even when .git metadata is not available.
fallback_version = "0.0.0.dev0"
python/setup.py
+92
View File
@@ -0,0 +1,92 @@
"""sglang build hooks.
SGLANG_BUILD_RUST_EXTS controls which Rust extensions are built:
- unset or "all": build every declared Rust extension (the default).
- "none": build no Rust extensions.
- comma-separated names: build only extensions whose target matches one of the
given (case-insensitive) substrings, e.g. "grpc" matches
"sglang.srt.grpc._core".
This is a build-time environment variable, so it is read directly from
os.environ instead of sglang.srt.environ, which is not available until after the
package has been built.
"""
import os
from setuptools import setup
try:
from setuptools_rust import build_rust
except ModuleNotFoundError as exc:
if exc.name != "setuptools_rust":
raise
# Alternate platform pyprojects do not declare Rust extensions.
build_rust = None
_BUILD_RUST_EXTS_ENV = "SGLANG_BUILD_RUST_EXTS"
def _selected_rust_extensions(declared):
"""Return the Rust extensions selected by SGLANG_BUILD_RUST_EXTS.
`ext.name` is the fully-qualified target (e.g. "sglang.srt.grpc._core") for
the string-target declarations in pyproject.toml, so comma-separated names
are matched as case-insensitive substrings of it.
"""
declared = list(declared)
raw = os.environ.get(_BUILD_RUST_EXTS_ENV)
if raw is None:
return declared
spec = raw.strip().lower()
# An empty or whitespace-only value is treated as unset (build everything).
if not spec or spec == "all":
return declared
if spec == "none":
return []
tokens = [token.strip() for token in spec.split(",")]
if not all(tokens):
raise ValueError(
f"{_BUILD_RUST_EXTS_ENV}={raw!r} has an empty item; unset it or use "
"'all', 'none', or a comma-separated list of extension names"
)
matched = set()
unmatched = []
for token in tokens:
hits = {ext.name for ext in declared if token in ext.name.lower()}
if hits:
matched |= hits
else:
unmatched.append(token)
if unmatched:
declared_names = sorted(ext.name for ext in declared)
raise ValueError(
f"{_BUILD_RUST_EXTS_ENV} matched no declared Rust extension for: "
f"{unmatched}; declared extensions are {declared_names}"
)
return [ext for ext in declared if ext.name in matched]
if build_rust is not None:
class BuildRust(build_rust):
"""Build only the Rust extensions selected by SGLANG_BUILD_RUST_EXTS."""
def run(self) -> None:
rust_extensions = _selected_rust_extensions(self.extensions or [])
self.extensions = rust_extensions
self.distribution.rust_extensions = rust_extensions
if not rust_extensions:
return
super().run()
_cmdclass = {"build_rust": BuildRust}
else:
_cmdclass = {}
setup(cmdclass=_cmdclass)
python/sglang/README.md
+18
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.
python/sglang/__init__.py
+116
View File
@@ -0,0 +1,116 @@
# 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 platform as _platform
import sys as _sys
if _sys.platform == "darwin" and _platform.machine() == "arm64":
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 _platform
del _sys
from sglang.srt.utils.hf_transformers_patches import apply_all as _apply_hf_patches
_apply_hf_patches()
del _apply_hf_patches
# 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")
Crusoe = LazyImport("sglang.lang.backend.crusoe", "Crusoe")
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",
"Crusoe",
"LiteLLM",
"OpenAI",
"VertexAI",
"global_config",
"__version__",
]
python/sglang/_mps_stub.py
+270
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
python/sglang/_triton_stub.py
+228
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
python/sglang/auto_benchmark.py
+82
View File
@@ -0,0 +1,82 @@
import argparse
from sglang.auto_benchmark_lib import (
SUPPORTED_DATASETS,
convert_dataset,
run_auto_benchmark,
validate_dataset,
)
def add_dataset_args(parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--kind",
required=True,
choices=sorted(SUPPORTED_DATASETS),
help="Dataset kind: sharegpt, custom, random, or generated-shared-prefix.",
)
parser.add_argument(
"--path",
default="",
help="Dataset file path. Leave empty for sharegpt auto-download.",
)
parser.add_argument("--tokenizer", required=True)
parser.add_argument("--model", default=None)
parser.add_argument("--num-prompts", type=int, default=1000)
parser.add_argument("--output-len", type=int, default=None)
parser.add_argument("--context-len", type=int, default=None)
parser.add_argument("--prompt-suffix", type=str, default="")
parser.add_argument("--apply-chat-template", action="store_true")
parser.add_argument("--random-input-len", type=int, default=1024)
parser.add_argument("--random-output-len", type=int, default=256)
parser.add_argument("--random-range-ratio", type=float, default=0.0)
parser.add_argument("--gsp-num-groups", type=int, default=64)
parser.add_argument("--gsp-prompts-per-group", type=int, default=16)
parser.add_argument("--gsp-system-prompt-len", type=int, default=2048)
parser.add_argument("--gsp-question-len", type=int, default=128)
parser.add_argument("--gsp-output-len", type=int, default=256)
parser.add_argument("--gsp-range-ratio", type=float, default=1.0)
parser.add_argument("--gsp-fast-prepare", action="store_true")
parser.add_argument("--gsp-send-routing-key", action="store_true")
parser.add_argument("--gsp-num-turns", type=int, default=1)
parser.add_argument("--gsp-ordered", action="store_true")
parser.add_argument("--seed", type=int, default=1)
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="SGLang auto benchmark utilities.")
subparsers = parser.add_subparsers(dest="command", required=True)
run_parser = subparsers.add_parser(
"run", help="Run auto benchmark from YAML config."
)
run_parser.add_argument("--config", required=True)
convert_parser = subparsers.add_parser(
"convert",
help="Prepare sharegpt/custom/random/generated-shared-prefix data into canonical autobench JSONL.",
)
add_dataset_args(convert_parser)
convert_parser.add_argument("--output", required=True)
validate_parser = subparsers.add_parser(
"validate", help="Validate a canonical autobench JSONL dataset."
)
validate_parser.add_argument("--dataset-path", required=True)
validate_parser.add_argument("--tokenizer", required=True)
return parser
def main() -> None:
args = build_parser().parse_args()
if args.command == "run":
run_auto_benchmark(args.config)
elif args.command == "convert":
convert_dataset(args)
elif args.command == "validate":
validate_dataset(args)
if __name__ == "__main__":
main()
python/sglang/auto_benchmark_lib.py
+1985
View File
File diff suppressed because it is too large Load Diff
python/sglang/bench_offline_throughput.py
+23
View File
@@ -0,0 +1,23 @@
"""Deprecated import path for ``sglang.benchmark.offline_throughput``.
``python -m sglang.bench_offline_throughput`` and
``from sglang.bench_offline_throughput import ...`` still work, but the
implementation now lives in ``sglang.benchmark.offline_throughput``.
Update references to the new path.
"""
import warnings
from sglang.benchmark.offline_throughput import * # noqa: F401,F403
from sglang.benchmark.offline_throughput import cli_main
warnings.warn(
"`sglang.bench_offline_throughput` is deprecated and will be removed in a "
"future release; use `sglang.benchmark.offline_throughput` instead "
"(e.g. `python -m sglang.benchmark.offline_throughput`).",
FutureWarning,
stacklevel=1,
)
if __name__ == "__main__":
cli_main()
python/sglang/bench_one_batch.py
+22
View File
@@ -0,0 +1,22 @@
"""Deprecated import path for ``sglang.benchmark.one_batch``.
``python -m sglang.bench_one_batch`` and ``from sglang.bench_one_batch import ...``
still work, but the implementation now lives in ``sglang.benchmark.one_batch``.
Update references to the new path.
"""
import warnings
from sglang.benchmark.one_batch import * # noqa: F401,F403
from sglang.benchmark.one_batch import cli_main
warnings.warn(
"`sglang.bench_one_batch` is deprecated and will be removed in a future "
"release; use `sglang.benchmark.one_batch` instead "
"(e.g. `python -m sglang.benchmark.one_batch`).",
FutureWarning,
stacklevel=1,
)
if __name__ == "__main__":
cli_main()
python/sglang/bench_one_batch_server.py
+23
View File
@@ -0,0 +1,23 @@
"""Deprecated import path for ``sglang.benchmark.one_batch_server``.
``python -m sglang.bench_one_batch_server`` and
``from sglang.bench_one_batch_server import ...`` still work, but the
implementation now lives in ``sglang.benchmark.one_batch_server``.
Update references to the new path.
"""
import warnings
from sglang.benchmark.one_batch_server import * # noqa: F401,F403
from sglang.benchmark.one_batch_server import cli_main
warnings.warn(
"`sglang.bench_one_batch_server` is deprecated and will be removed in a "
"future release; use `sglang.benchmark.one_batch_server` instead "
"(e.g. `python -m sglang.benchmark.one_batch_server`).",
FutureWarning,
stacklevel=1,
)
if __name__ == "__main__":
cli_main()
python/sglang/bench_serving.py
+22
View File
@@ -0,0 +1,22 @@
"""Deprecated import path for ``sglang.benchmark.serving``.
``python -m sglang.bench_serving`` and ``from sglang.bench_serving import ...``
still work, but the implementation now lives in ``sglang.benchmark.serving``.
Update references to the new path.
"""
import warnings
from sglang.benchmark.serving import * # noqa: F401,F403
from sglang.benchmark.serving import cli_main
warnings.warn(
"`sglang.bench_serving` is deprecated and will be removed in a future "
"release; use `sglang.benchmark.serving` instead "
"(e.g. `python -m sglang.benchmark.serving`).",
FutureWarning,
stacklevel=1,
)
if __name__ == "__main__":
cli_main()
python/sglang/benchmark/__init__.py
View File
python/sglang/benchmark/bench_utils.py
+23
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)
python/sglang/benchmark/datasets/__init__.py
+55
View File
@@ -0,0 +1,55 @@
from typing import Dict, Type
from sglang.benchmark.datasets.agentic_trace import AgenticTraceDataset
from sglang.benchmark.datasets.autobench import AutoBenchmarkDataset
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
from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset
DATASET_MAPPING: Dict[str, Type[BaseDataset]] = {
"agentic-trace": AgenticTraceDataset,
"autobench": AutoBenchmarkDataset,
"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,
"speed-bench": SpeedBenchDataset,
}
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",
]
python/sglang/benchmark/datasets/agentic_trace.py
+114
View File
@@ -0,0 +1,114 @@
import json
import os
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
# Per-turn output length when --sharegpt-output-len is not given; matches the
# ~220-token average assistant reply of OpenHands-style agentic traces.
DEFAULT_AGENTIC_OUTPUT_LEN = 220
@dataclass
class AgenticTraceDataset(BaseDataset):
"""Multi-turn agentic trace loader (e.g. OpenHands / SWE-smith traces).
Expects a trace JSON of the shape::
{
"metadata": {...},
"conversations": [
[ # one conversation == a list of turns
{"messages": [{"role": "system", ...}, {"role": "user", ...}],
"prompt_tokens": 73821},
{"messages": [{"role": "user", ...}], "prompt_tokens": 74894},
...
],
...
]
}
Each turn's ``messages`` holds only the new non-assistant messages for that
turn. One conversation becomes one :class:`DatasetRow` whose ``prompt`` is
the list of per-turn message deltas; ``bench_serving`` detects this shape as
multi-turn and replays each conversation round by round, feeding the
server's real assistant reply back into the next round's history.
Use with a chat backend (``--backend sglang-oai-chat``).
"""
dataset_path: str
num_requests: int
fixed_output_len: Optional[int]
offset: int
max_turns: Optional[int]
@classmethod
def from_args(cls, args: Namespace) -> "AgenticTraceDataset":
return cls(
dataset_path=args.dataset_path,
num_requests=args.num_prompts,
fixed_output_len=args.sharegpt_output_len,
offset=args.dataset_offset,
max_turns=args.agentic_max_turns,
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
if not os.path.isfile(self.dataset_path):
raise FileNotFoundError(f"Dataset not found at {self.dataset_path}")
with open(self.dataset_path, "r", encoding="utf-8") as f:
data = json.load(f)
conversations = data.get("conversations", [])
if not conversations:
raise ValueError(f"No 'conversations' found in {self.dataset_path}.")
offset = self.offset % len(conversations)
if offset:
conversations = conversations[offset:] + conversations[:offset]
output_len = self.fixed_output_len or DEFAULT_AGENTIC_OUTPUT_LEN
filtered_dataset: List[DatasetRow] = []
for conversation in conversations:
if self.num_requests > 0 and len(filtered_dataset) >= self.num_requests:
break
prompt = [turn["messages"] for turn in conversation if turn.get("messages")]
if self.max_turns:
prompt = prompt[: self.max_turns]
if not prompt:
continue
# Informational only: multi-turn replay ignores per-row prompt_len.
prompt_len = int(conversation[0].get("prompt_tokens", 0))
filtered_dataset.append(
DatasetRow(
prompt=prompt,
prompt_len=prompt_len,
output_len=output_len,
)
)
if not filtered_dataset:
raise ValueError(
f"No usable conversations loaded from {self.dataset_path}."
)
num_turns = [len(row.prompt) for row in filtered_dataset]
print(
f"#Conversations: {len(filtered_dataset)} "
f"(offset={offset}, turns/conv min={min(num_turns)} "
f"max={max(num_turns)} avg={np.mean(num_turns):.1f})"
)
print(f"#Output tokens per turn: {output_len}")
return filtered_dataset
python/sglang/benchmark/datasets/autobench.py
+299
View File
@@ -0,0 +1,299 @@
import json
from argparse import Namespace
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
from transformers import PreTrainedTokenizerBase
from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
AUTOBENCH_RESERVED_FIELDS = {
"prompt",
"messages",
"prompt_origin",
"output_len",
"max_tokens",
"max_completion_tokens",
"completion_tokens",
"prompt_len",
"text_prompt_len",
"vision_prompt_len",
"image_data",
"timestamp",
"routing_key",
"metadata",
"extra_request_body",
"param_send",
}
def _load_json_if_needed(value: Any) -> Any:
if not isinstance(value, str):
return value
value = value.strip()
if not value:
return value
if value[0] not in "[{":
return value
try:
return json.loads(value)
except json.JSONDecodeError:
return value
def _normalize_messages(messages: Any) -> Optional[List[Dict[str, Any]]]:
messages = _load_json_if_needed(messages)
if not isinstance(messages, list) or not messages:
return None
if not all(isinstance(message, dict) for message in messages):
return None
normalized = []
for message in messages:
if "role" not in message:
return None
content = message.get("content")
if content is None:
return None
normalized.append({"role": message["role"], "content": content})
return normalized
def _normalize_legacy_system_content(
system_prompt: Any, content_list: Any
) -> Optional[List[Dict[str, Any]]]:
if not isinstance(content_list, list) or not content_list:
return None
messages: List[Dict[str, Any]] = []
if system_prompt:
messages.append({"role": "system", "content": str(system_prompt)})
turns = [str(item) for item in content_list]
# In the old auto_benchmark helpers, an even number of items usually means the
# last assistant reply is present and should be removed before benchmarking.
if len(turns) % 2 == 0:
turns = turns[:-1]
if not turns:
return None
for index, turn in enumerate(turns):
role = "user" if index % 2 == 0 else "assistant"
messages.append({"role": role, "content": turn})
return messages
def _normalize_prompt(row: Dict[str, Any]) -> Tuple[Any, str]:
prompt = row.get("prompt")
messages = row.get("messages")
prompt_origin = row.get("prompt_origin")
if messages is not None:
normalized = _normalize_messages(messages)
if normalized is not None:
return normalized, "messages"
if prompt is not None:
prompt = _load_json_if_needed(prompt)
if isinstance(prompt, list) and prompt and isinstance(prompt[0], dict):
normalized = _normalize_messages(prompt)
if normalized is not None:
return normalized, "messages"
if (
isinstance(prompt, list)
and prompt
and all(isinstance(item, str) for item in prompt)
):
return prompt, "multi_turn"
if (
isinstance(prompt, list)
and prompt
and all(
isinstance(item, list)
and item
and all(
isinstance(m, dict) and "role" in m and "content" in m for m in item
)
for item in prompt
)
):
# Multi-turn with N messages per round (e.g. tool observations).
return prompt, "multi_turn"
if (
isinstance(prompt, list)
and prompt
and all(isinstance(item, int) for item in prompt)
):
return prompt, "token_ids"
if isinstance(prompt, str) and prompt:
return prompt, "prompt"
if prompt_origin is not None:
normalized = _normalize_messages(prompt_origin)
if normalized is not None:
return normalized, "messages"
if "system" in row and "content" in row:
normalized = _normalize_legacy_system_content(
row.get("system"), row.get("content")
)
if normalized is not None:
return normalized, "messages"
raise ValueError("Unsupported auto benchmark row: missing prompt/messages")
def _estimate_prompt_lens(
prompt: Any,
prompt_kind: str,
tokenizer: PreTrainedTokenizerBase,
row: Dict[str, Any],
) -> Tuple[int, int, int]:
if row.get("prompt_len") is not None:
prompt_len = int(row["prompt_len"])
text_prompt_len = int(row.get("text_prompt_len", prompt_len))
vision_prompt_len = int(row.get("vision_prompt_len", 0))
return prompt_len, text_prompt_len, vision_prompt_len
if prompt_kind == "messages":
text_prompt_len = len(
tokenizer.apply_chat_template(
prompt, tokenize=True, add_generation_prompt=True
)
)
vision_prompt_len = 0
return text_prompt_len, text_prompt_len, vision_prompt_len
if prompt_kind == "prompt":
prompt_len = len(tokenizer.encode(prompt, add_special_tokens=False))
return prompt_len, prompt_len, 0
if prompt_kind == "token_ids":
prompt_len = len(prompt)
return prompt_len, prompt_len, 0
# Multi-turn prompt lists are handled specially by the serving benchmark and do not
# contribute reliable static prompt lengths.
return 0, 0, 0
def _collect_extra_request_body(row: Dict[str, Any]) -> Dict[str, Any]:
extra: Dict[str, Any] = {}
param_send = row.get("param_send")
if param_send is not None:
parsed = _load_json_if_needed(param_send)
if isinstance(parsed, dict):
extra.update(parsed)
for key, value in row.items():
if key not in AUTOBENCH_RESERVED_FIELDS:
extra[key] = value
explicit_extra = row.get("extra_request_body")
explicit_extra = _load_json_if_needed(explicit_extra)
if isinstance(explicit_extra, dict):
extra.update(explicit_extra)
return extra
def serialize_dataset_row_to_autobench(
row: DatasetRow, metadata: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
record: Dict[str, Any] = {
"prompt": row.prompt,
"output_len": row.output_len,
}
if row.prompt_len:
record["prompt_len"] = row.prompt_len
if row.text_prompt_len not in (None, row.prompt_len):
record["text_prompt_len"] = row.text_prompt_len
if row.vision_prompt_len:
record["vision_prompt_len"] = row.vision_prompt_len
if row.image_data:
record["image_data"] = row.image_data
if row.timestamp is not None:
record["timestamp"] = row.timestamp
if row.routing_key is not None:
record["routing_key"] = row.routing_key
if row.extra_request_body:
record["extra_request_body"] = row.extra_request_body
if metadata:
record["metadata"] = metadata
return record
@dataclass
class AutoBenchmarkDataset(BaseDataset):
dataset_path: str
num_requests: int
fixed_output_len: Optional[int]
@classmethod
def from_args(cls, args: Namespace) -> "AutoBenchmarkDataset":
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_autobench_requests(
dataset_path=self.dataset_path,
num_requests=self.num_requests,
tokenizer=tokenizer,
fixed_output_len=self.fixed_output_len,
)
def sample_autobench_requests(
dataset_path: str,
num_requests: int,
tokenizer: PreTrainedTokenizerBase,
fixed_output_len: Optional[int] = None,
) -> List[DatasetRow]:
dataset: List[DatasetRow] = []
with open(dataset_path, "r", encoding="utf-8") as f:
for line in f:
if num_requests > 0 and len(dataset) >= num_requests:
break
line = line.strip()
if not line:
continue
row = json.loads(line)
prompt, prompt_kind = _normalize_prompt(row)
prompt_len, text_prompt_len, vision_prompt_len = _estimate_prompt_lens(
prompt, prompt_kind, tokenizer, row
)
output_len = fixed_output_len or row.get("output_len")
output_len = output_len or row.get("max_tokens")
output_len = output_len or row.get("max_completion_tokens")
output_len = output_len or row.get("completion_tokens")
output_len = int(output_len or 256)
dataset.append(
DatasetRow(
prompt=prompt,
prompt_len=prompt_len,
output_len=output_len,
text_prompt_len=text_prompt_len,
vision_prompt_len=vision_prompt_len,
image_data=row.get("image_data"),
timestamp=row.get("timestamp"),
routing_key=row.get("routing_key"),
extra_request_body=_collect_extra_request_body(row),
)
)
print(f"Loaded {len(dataset)} auto benchmark requests")
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])}")
return dataset
python/sglang/benchmark/datasets/common.py
+101
View File
@@ -0,0 +1,101 @@
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 valid token ids from the tokenizer vocabulary."""
return [
token_id
for token_id in tokenizer.get_vocab().values()
if isinstance(token_id, int)
]
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)
@lru_cache(maxsize=1)
def get_available_multimodal_text_tokens(tokenizer, image_pad_id):
"""Get valid token ids for synthetic multimodal text prompts."""
excluded_token_ids = set(getattr(tokenizer, "all_special_ids", []) or [])
if image_pad_id is not None:
excluded_token_ids.add(image_pad_id)
return [
token_id
for token_id in get_available_tokens(tokenizer)
if token_id not in excluded_token_ids
]
def gen_mm_prompt(tokenizer, image_pad_id, token_num):
"""Generate a random prompt of specified token length using tokenizer vocabulary."""
all_available_tokens = get_available_multimodal_text_tokens(tokenizer, image_pad_id)
selected_tokens = random.choices(all_available_tokens, k=token_num)
return tokenizer.decode(selected_tokens)
python/sglang/benchmark/datasets/custom.py
+147
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
python/sglang/benchmark/datasets/generated_shared_prefix.py
+328
View File
@@ -0,0 +1,328 @@
import math
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, Optional
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,
)
def _zipf_group_probs(num_groups: int, alpha: float) -> np.ndarray:
"""Rank-based Zipf probability vector with rank starting at 1.
weight(rank) = 1 / rank ** alpha (rank in 1..num_groups)
probability(rank) = weight(rank) / sum_over_all_ranks(weight)
The returned array has length num_groups; element i corresponds to
group index i (rank i + 1), so group 0 is the hottest.
"""
if num_groups <= 0:
raise ValueError(f"num_groups must be > 0, got {num_groups}")
ranks = np.arange(1, num_groups + 1, dtype=np.float64)
weights = 1.0 / (ranks**alpha)
return weights / weights.sum()
@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
group_distribution: str = "uniform"
zipf_alpha: Optional[float] = None
@classmethod
def from_args(cls, args: Namespace) -> "GeneratedSharedPrefixDataset":
assert not getattr(args, "tokenize_prompt", False)
group_distribution = getattr(args, "gsp_group_distribution", "uniform")
zipf_alpha = getattr(args, "gsp_zipf_alpha", None)
# Defensive validation for in-process callers that construct a
# Namespace by hand and bypass the argparse boundary in
# serving.py. The CLI hook enforces the same rules first.
if group_distribution not in ("uniform", "zipf"):
raise ValueError(
f"--gsp-group-distribution must be 'uniform' or 'zipf', "
f"got {group_distribution!r}"
)
if group_distribution == "zipf":
if zipf_alpha is None:
raise ValueError(
"--gsp-group-distribution=zipf requires --gsp-zipf-alpha "
"(a finite float > 0)"
)
if not math.isfinite(zipf_alpha) or zipf_alpha <= 0:
raise ValueError(
f"--gsp-zipf-alpha must be a finite float > 0, got {zipf_alpha!r}"
)
elif zipf_alpha is not None:
raise ValueError(
"--gsp-zipf-alpha is only meaningful with "
"--gsp-group-distribution=zipf; remove --gsp-zipf-alpha "
"or set --gsp-group-distribution=zipf"
)
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),
group_distribution=group_distribution,
zipf_alpha=zipf_alpha,
)
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,
group_distribution=self.group_distribution,
zipf_alpha=self.zipf_alpha,
)
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,
group_distribution: str = "uniform",
zipf_alpha: Optional[float] = None,
):
"""Create cache directory under ~/.cache/sglang/benchmark.
The uniform-mode filename is preserved exactly as before so existing
on-disk caches remain valid. Non-default sampling modes get an extra
suffix encoding the parameters that affect the cached payload.
"""
cache_dir = Path.home() / ".cache" / "sglang" / "benchmark"
suffix = ""
if group_distribution != "uniform":
suffix = f"_{group_distribution}_{zipf_alpha}"
cache_key = (
f"gen_shared_prefix_{seed}_{num_groups}_{prompts_per_group}_"
f"{system_prompt_len}_{question_len}_{output_len}{suffix}_"
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,
group_distribution: str = "uniform",
zipf_alpha: Optional[float] = None,
) -> List[DatasetRow]:
"""Generate benchmark requests with shared system prompts using random tokens and caching.
When group_distribution is "uniform" (default), each group receives exactly
prompts_per_group requests; behavior matches the legacy generator.
When group_distribution is "zipf", each request's group is sampled by rank
with probability 1/rank**zipf_alpha / sum_k(1/k**zipf_alpha); rank starts at
1 and group index 0 is the hottest. Sampling uses an isolated
numpy.random.default_rng(seed) so the shared question/system-prompt pool
stays byte-identical to uniform mode for the same seed and other args.
Zipf mode is cached on disk under a distinct key per (group_distribution,
zipf_alpha) value.
"""
cache_path = get_gen_prefix_cache_path(
seed,
num_groups,
prompts_per_group,
system_prompt_len,
question_len,
output_len,
tokenizer,
group_distribution=group_distribution,
zipf_alpha=zipf_alpha,
)
# range_ratio != 1 / num_turns > 1 perturb the payload but are not in the
# cache key; send_routing_key embeds a per-run uuid + timestamp that is
# meaningless to cache. Bypass for these pre-existing reasons only.
should_cache = range_ratio == 1 and not send_routing_key and num_turns == 1
if should_cache and cache_path.exists():
print(f"\nLoading cached generated input data from {cache_path}")
with open(cache_path, "rb") as f:
return pickle.load(f)
if not should_cache:
print(f"\nCache bypassed ({range_ratio=}, {send_routing_key=}, {num_turns=})")
print(
f"\nGenerating new input data... "
f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=}, {num_turns=}, {group_distribution=}, {zipf_alpha=})"
)
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
system_prompts = [
gen_prompt(tokenizer, system_prompt_lens[i]) for i in range(num_groups)
]
# 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)
]
# Per-slot group assignment. Uniform mode is the identity assignment
# [0,0,...,1,1,...,N-1,N-1]; zipf mode samples from the rank distribution
# using an isolated RNG so the module-level random / numpy.random state
# that compute_random_lens / gen_prompt rely on is never perturbed -- this
# keeps the system-prompt and question pool byte-identical to uniform mode
# for the same seed and other args.
total_slots = num_groups * prompts_per_group
if group_distribution == "uniform":
assignment = np.repeat(np.arange(num_groups), prompts_per_group)
else: # "zipf"
rng = np.random.default_rng(seed)
probs = _zipf_group_probs(num_groups, zipf_alpha)
assignment = rng.choice(num_groups, size=total_slots, replace=True, p=probs)
input_requests = []
total_input_tokens = 0
total_output_tokens = 0
for slot_idx, sampled_g in enumerate(
tqdm(assignment, desc="Generating shared-prefix prompts")
):
# src_(g,p) walks the question pool in uniform-enumeration order, so
# per-slot question text is reproducibly identical across modes.
src_g, src_p = divmod(slot_idx, prompts_per_group)
sampled_g = int(sampled_g)
system_prompt = system_prompts[sampled_g]
routing_key = (
f"{run_random_str}_{run_start_timestamp}_{sampled_g}"
if send_routing_key
else None
)
turn_questions = questions[src_g][src_p]
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[src_g, src_p])
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(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"Group distribution: {group_distribution}")
if group_distribution == "zipf":
print(f"Zipf alpha: {zipf_alpha}")
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"
)
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
python/sglang/benchmark/datasets/image.py
+381
View File
@@ -0,0 +1,381 @@
import io
import warnings
from argparse import Namespace
from dataclasses import dataclass
from typing import List, Optional, 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 parse_random_image_resolution(
image_resolution: str,
) -> Optional[Tuple[Tuple[int, int], Tuple[int, int]]]:
"""Parse ``random:<min_h>x<min_w>-<max_h>x<max_w>`` image bounds.
Returns ``None`` for fixed resolutions. The returned dimensions are
``(width, height)`` pairs, matching :func:`parse_image_resolution`.
"""
prefix = "random:"
if not image_resolution.strip().lower().startswith(prefix):
return None
bounds = image_resolution.strip()[len(prefix) :].split("-", maxsplit=1)
if len(bounds) != 2:
raise ValueError(
"Random image resolution must be 'random:<min_h>x<min_w>-"
"<max_h>x<max_w>', for example 'random:256x256-1024x1024'."
)
min_width, min_height = parse_image_resolution(bounds[0])
max_width, max_height = parse_image_resolution(bounds[1])
if min_width > max_width or min_height > max_height:
raise ValueError("Random image resolution minimum cannot exceed maximum.")
return (min_width, min_height), (max_width, max_height)
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
if type(processor).__name__ == "MiniCPMOProcessor":
prompt_str = f"(<image>./</image>){text_prompt}"
else:
prompt_str = f"<image>{text_prompt}"
# Calculate total tokens (text + vision)
if type(processor).__name__ == "KimiK25Processor":
medias = [{"type": "image", "image": img} for img in images]
prompt_len = processor(
text=prompt_str,
medias=medias,
return_tensors="pt",
)["input_ids"].numel()
elif type(processor).__name__ == "VLChatProcessor":
prompt_len = processor(
prompt=prompt_str,
images=images,
force_batchify=False,
)["input_ids"].numel()
elif type(processor).__name__ == "DeepseekVLV2Processor":
result = processor(
conversations=prompt_str,
images=images,
inference_mode=True,
)
prompt_len = result.input_ids.numel()
else:
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",
"vllm-chat",
]
if backend not in supported_backends:
raise ValueError(
f"Image dataset only supports backends: {supported_backends}, "
f"got '{backend}'."
)
# OpenAI chat handlers apply the chat template and receive images separately, so
# send the raw text. /generate does not apply a chat template, so it needs
# prompt_str, which contains the multimodal processor's image placeholders.
use_raw_prompt = backend in ("sglang-oai-chat", "vllm-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), custom ``heightxwidth`` (e.g.,
1080x1920), or ``random:<min_h>x<min_w>-<max_h>x<max_w>``.
- Text lengths follow the 'random' dataset sampling rule. ``prompt_len``
only counts text tokens and excludes image data.
"""
random_resolution_bounds = parse_random_image_resolution(image_resolution)
if random_resolution_bounds is None:
width, height = parse_image_resolution(image_resolution)
min_width = max_width = width
min_height = max_height = height
else:
(min_width, min_height), (max_width, max_height) = random_resolution_bounds
# 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 max_width * max_height >= 1920 * 1080 and total_images >= 100:
warnings.warn(
f"High resolution (up to {max_width}x{max_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() -> Tuple[Image.Image, str, int, Tuple[int, int]]:
if random_resolution_bounds is None:
width, height = min_width, min_height
else:
width = np.random.randint(min_width, max_width + 1)
height = np.random.randint(min_height, max_height + 1)
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, (width, height)
dataset: List[DatasetRow] = []
total_image_bytes = 0
all_image_sizes: list[Tuple[int, int]] = []
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 if hasattr(processor, "tokenizer") else processor,
processor.image_token_id if hasattr(processor, "image_token_id") else None,
int(input_lens[i]),
)
# Generate image list
images, images_base64, images_bytes, image_sizes = zip(
*[_gen_random_image_data_uri() for _ in range(request_image_count)]
)
total_image_bytes += sum(images_bytes)
all_image_sizes.extend(image_sizes)
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)")
if random_resolution_bounds is not None:
widths, heights = zip(*all_image_sizes)
print(
"#Image resolution: "
f"min={min(widths)}x{min(heights)}, "
f"max={max(widths)}x{max(heights)}, "
f"mean={np.mean(widths):.1f}x{np.mean(heights):.1f}"
)
# Detailed token breakdown (derived from dataset + input_lens)
text_prompt_lens = np.array([r.text_prompt_len for r in dataset])
vision_prompt_lens = np.array([r.vision_prompt_len for r in dataset])
text_prompt_overheads = text_prompt_lens - input_lens
stat_fields = [
("Raw text prompt tokens (without overhead)", input_lens),
("Text prompt tokens (with chat template)", text_prompt_lens),
("Text prompt overhead", text_prompt_overheads),
("Vision tokens", vision_prompt_lens),
]
print("\n=== Token Breakdown (per request avg / total) ===")
for label, vals in stat_fields:
print(f" {label}: avg={np.mean(vals):.1f}, total={np.sum(vals)}")
print(
f"\nCreated {len(dataset)} {image_content} {image_format} images with average {total_image_bytes // num_requests} bytes per request"
)
return dataset
python/sglang/benchmark/datasets/longbench_v2.py
+104
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
python/sglang/benchmark/datasets/mmmu.py
+124
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
python/sglang/benchmark/datasets/mooncake.py
+123
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})
python/sglang/benchmark/datasets/openai_dataset.py
+113
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 - the serving benchmark 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
python/sglang/benchmark/datasets/random.py
+167
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
python/sglang/benchmark/datasets/sharegpt.py
+151
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
python/sglang/benchmark/datasets/speed_bench.py
+102
View File
@@ -0,0 +1,102 @@
"""SPEED-Bench (nvidia/SPEED-Bench) dataset for the SGLang serving benchmark.
Reads the pre-downloaded throughput_1k JSONL produced by prepare_speed_bench.sh
(or equivalent), optionally filtering by category (low_entropy / mixed /
high_entropy) and fixing the output length.
CLI args consumed:
--dataset-path Path to the local JSONL file.
--speed-bench-category Category filter: low_entropy | mixed | high_entropy
(default: all categories).
--speed-bench-output-len Fixed number of output tokens per request (default: 512).
--num-prompts Number of requests to sample (capped by available rows).
"""
import json
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
@dataclass
class SpeedBenchDataset(BaseDataset):
dataset_path: str
category: Optional[str]
output_len: int
num_requests: int
@classmethod
def from_args(cls, args: Namespace) -> "SpeedBenchDataset":
if not args.dataset_path:
raise ValueError(
"--dataset-path must point to the SPEED-Bench JSONL file "
"(run prepare_speed_bench.sh to generate it)."
)
return cls(
dataset_path=args.dataset_path,
category=getattr(args, "speed_bench_category", None) or None,
output_len=getattr(args, "speed_bench_output_len", 512),
num_requests=args.num_prompts,
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
unique_prompts = []
with open(self.dataset_path, encoding="utf-8") as f:
for line in f:
row = json.loads(line)
if self.category and row.get("category") != self.category:
continue
# turns is a list of strings; use the first user turn as the prompt
turns = row.get("turns", [])
if not turns:
continue
unique_prompts.append(turns[0])
if not unique_prompts:
raise ValueError(
f"No rows found in {self.dataset_path}"
+ (f" for category={self.category}" if self.category else "")
)
# Tokenize unique prompts once to avoid redundant work
unique_dataset_rows: List[DatasetRow] = []
for prompt_text in unique_prompts:
# Apply chat template to match vllm bench behaviour
try:
prompt_ids = tokenizer.apply_chat_template(
[{"role": "user", "content": prompt_text}],
add_generation_prompt=True,
tokenize=True,
)
prompt = tokenizer.decode(prompt_ids)
except Exception:
prompt_ids = tokenizer.encode(prompt_text)
prompt = prompt_text
unique_dataset_rows.append(
DatasetRow(
prompt=prompt,
prompt_len=len(prompt_ids),
output_len=self.output_len,
)
)
# Sample (with replacement if needed); shuffle oversampled rows for
# a realistic request distribution
if self.num_requests <= len(unique_dataset_rows):
dataset_rows = random.sample(unique_dataset_rows, self.num_requests)
else:
dataset_rows = unique_dataset_rows * (
self.num_requests // len(unique_dataset_rows) + 1
)
dataset_rows = dataset_rows[: self.num_requests]
random.shuffle(dataset_rows)
return dataset_rows
python/sglang/benchmark/dspark_sps_profiler.py
+1577
View File
File diff suppressed because it is too large Load Diff
python/sglang/benchmark/dspark_sts_fit.py
+227
View File
@@ -0,0 +1,227 @@
from __future__ import annotations
import argparse
import glob
import logging
import math
from pathlib import Path
from typing import Optional
import torch
from sglang.srt.speculative.dspark_components.dspark_sts import (
DSparkStsCalibration,
)
logger = logging.getLogger(__name__)
_EPS_PROB = 1e-8
def default_temperature_grid() -> torch.Tensor:
return torch.logspace(math.log10(0.1), math.log10(10.0), steps=41)
def expected_calibration_error(
*,
probs: torch.Tensor,
targets: torch.Tensor,
num_bins: int,
) -> float:
probs = probs.reshape(-1).to(torch.float64).clamp(_EPS_PROB, 1.0 - _EPS_PROB)
targets = targets.reshape(-1).to(torch.float64)
total = probs.numel()
if total == 0:
return float("nan")
bin_index = (probs * num_bins).long().clamp_(0, num_bins - 1)
count = torch.zeros(num_bins, dtype=torch.float64)
pred_sum = torch.zeros(num_bins, dtype=torch.float64)
target_sum = torch.zeros(num_bins, dtype=torch.float64)
count.scatter_add_(0, bin_index, torch.ones_like(probs))
pred_sum.scatter_add_(0, bin_index, probs)
target_sum.scatter_add_(0, bin_index, targets)
denom = count.clamp_min(1.0)
bin_error = (pred_sum / denom - target_sum / denom).abs()
return float((bin_error * count).sum().item() / total)
def fit_sts_temperatures(
*,
logits: torch.Tensor,
prefix_mask: torch.Tensor,
grid: torch.Tensor,
num_bins: int = 15,
) -> dict[str, list[float]]:
logits = logits.to(torch.float64)
prefix_mask = prefix_mask.to(torch.float64)
num_samples, gamma = logits.shape
if num_samples == 0:
raise ValueError("fit_sts_temperatures requires at least one sample.")
grid_values = grid.to(torch.float64).tolist()
temperatures: list[float] = []
ece_before: list[float] = []
ece_after: list[float] = []
survival_at_one = torch.ones(num_samples, dtype=torch.float64)
survival_fitted = torch.ones(num_samples, dtype=torch.float64)
for position in range(gamma):
position_logits = logits[:, position]
position_target = prefix_mask[:, position]
survival_at_one = survival_at_one * torch.sigmoid(position_logits)
ece_before.append(
expected_calibration_error(
probs=survival_at_one,
targets=position_target,
num_bins=num_bins,
)
)
best_temperature = grid_values[0]
best_survival = survival_fitted * torch.sigmoid(
position_logits / best_temperature
)
best_ece = expected_calibration_error(
probs=best_survival, targets=position_target, num_bins=num_bins
)
for temperature in grid_values[1:]:
candidate_survival = survival_fitted * torch.sigmoid(
position_logits / temperature
)
candidate_ece = expected_calibration_error(
probs=candidate_survival,
targets=position_target,
num_bins=num_bins,
)
if candidate_ece < best_ece:
best_ece = candidate_ece
best_temperature = temperature
best_survival = candidate_survival
temperatures.append(float(best_temperature))
ece_after.append(float(best_ece))
survival_fitted = best_survival
return {
"temperatures": temperatures,
"ece_before": ece_before,
"ece_after": ece_after,
}
def load_collected_shards(*, data_glob: str) -> tuple[torch.Tensor, torch.Tensor]:
shard_paths = sorted(glob.glob(data_glob))
if not shard_paths:
raise ValueError(f"No STS data shards matched {data_glob!r}.")
logits_shards: list[torch.Tensor] = []
prefix_mask_shards: list[torch.Tensor] = []
shard_gamma: Optional[int] = None
for shard_path in shard_paths:
shard = torch.load(shard_path, map_location="cpu")
shard_logits = shard["logits"]
shard_prefix_mask = shard["prefix_mask"]
if shard_logits.shape != shard_prefix_mask.shape:
raise ValueError(
f"Shard {shard_path!r} logits / prefix_mask shape mismatch: "
f"{tuple(shard_logits.shape)} vs {tuple(shard_prefix_mask.shape)}."
)
if shard_gamma is None:
shard_gamma = int(shard_logits.shape[1])
elif int(shard_logits.shape[1]) != shard_gamma:
raise ValueError(
f"Shard {shard_path!r} gamma {int(shard_logits.shape[1])} disagrees "
f"with earlier shards' gamma {shard_gamma}."
)
logits_shards.append(shard_logits)
prefix_mask_shards.append(shard_prefix_mask)
return torch.cat(logits_shards, dim=0), torch.cat(prefix_mask_shards, dim=0)
def fit(
*,
data_glob: str,
out: Path,
num_bins: int = 15,
gamma: Optional[int] = None,
) -> None:
logits, prefix_mask = load_collected_shards(data_glob=data_glob)
resolved_gamma = int(logits.shape[1])
if gamma is not None and gamma != resolved_gamma:
raise ValueError(
f"Collected shards have gamma={resolved_gamma} but --gamma={gamma}."
)
num_samples = int(logits.shape[0])
result = fit_sts_temperatures(
logits=logits,
prefix_mask=prefix_mask,
grid=default_temperature_grid(),
num_bins=num_bins,
)
calibration = DSparkStsCalibration(
temperatures=result["temperatures"],
dataset=data_glob,
num_samples=num_samples,
ece_before=result["ece_before"],
ece_after=result["ece_after"],
)
out.write_text(calibration.to_json(), encoding="utf-8")
print(
f"Fit STS temperatures over {num_samples} samples (gamma={resolved_gamma}) "
f"-> {out}"
)
print("pos temperature ece_before ece_after")
for position in range(resolved_gamma):
print(
f"{position:>3} {result['temperatures'][position]:>11.4f} "
f"{result['ece_before'][position]:>10.4f} "
f"{result['ece_after'][position]:>9.4f}"
)
def main() -> None:
logging.basicConfig(level=logging.INFO)
parser = argparse.ArgumentParser(
description="Fit DSpark Sequential Temperature Scaling (STS) calibration "
"temperatures from collected confidence shards."
)
parser.add_argument(
"--data-glob",
required=True,
help="Glob of collected .pt shards, each a dict with [n, gamma] "
"'logits' and 'prefix_mask' tensors.",
)
parser.add_argument(
"--out",
required=True,
type=Path,
help="Output STS calibration JSON path.",
)
parser.add_argument(
"--num-bins",
type=int,
default=15,
help="Number of equal-width ECE bins.",
)
parser.add_argument(
"--gamma",
type=int,
default=None,
help="Optional gamma override to validate the shards against.",
)
args = parser.parse_args()
fit(
data_glob=args.data_glob,
out=args.out,
num_bins=args.num_bins,
gamma=args.gamma,
)
if __name__ == "__main__":
main()
python/sglang/benchmark/endpoint.py
+127
View File
@@ -0,0 +1,127 @@
"""Connection target for HTTP benchmark scripts.
Owns the launch-vs-connect decision in one place: a benchmark only needs a base
URL, which comes either from a server we launch or one already running.
"""
import dataclasses
import multiprocessing
import os
import time
from typing import Callable, Optional
import requests
from sglang.srt.entrypoints.http_server import launch_server
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import kill_process_tree
from sglang.srt.utils.network import resolve_base_url
DEFAULT_TIMEOUT = 600
# Field defaults of ServerArgs, used to detect when --host/--port were set
# explicitly (and would be silently ignored in connect mode).
_SERVER_ARGS_DEFAULTS = {f.name: f.default for f in dataclasses.fields(ServerArgs)}
def server_is_up(base_url: str, timeout: float = DEFAULT_TIMEOUT) -> bool:
"""Return True if a server answers /v1/models with 200 at base_url."""
try:
headers = {
"Content-Type": "application/json; charset=utf-8",
}
response = requests.get(
f"{base_url}/v1/models", headers=headers, timeout=timeout
)
return response.status_code == 200
except requests.RequestException:
return False
def _launch_server_target(launch_server_func: Callable, server_args: ServerArgs):
try:
launch_server_func(server_args)
except Exception as e:
raise e
finally:
kill_process_tree(os.getpid(), include_parent=False)
def launch_or_reuse_server(launch_server_func: Callable, server_args: ServerArgs):
base_url = resolve_base_url("", server_args.host, server_args.port)
# Reuse an already-running server instead of forking a second one onto the
# occupied port, where it would orphan, compete for the GPU, and OOM.
if server_is_up(base_url, timeout=5):
print(
f"WARNING: reusing the server already running at {base_url} "
f"(--model and server-launch args ignored). Pass --base-url to silence."
)
return None, base_url
proc = multiprocessing.Process(
target=_launch_server_target,
args=(
launch_server_func,
server_args,
),
)
proc.start()
start_time = time.time()
while time.time() - start_time < DEFAULT_TIMEOUT:
# Fail fast if the server died during startup (e.g. OOM).
if not proc.is_alive():
raise RuntimeError(
f"Server process exited during startup (exit code "
f"{proc.exitcode}); see the traceback above for the cause."
)
if server_is_up(base_url):
return proc, base_url
time.sleep(10)
# Timed out: kill the half-started server so it does not linger as an orphan.
kill_process_tree(proc.pid)
raise TimeoutError("Server failed to start within the timeout period.")
@dataclasses.dataclass
class BenchEndpoint:
"""A base URL plus the lifecycle of any server we launched to back it.
``close()`` tears down a launched server; for a connected one it is a no-op.
"""
base_url: str
_proc: Optional[multiprocessing.Process] = None
def close(self) -> None:
if self._proc is not None:
kill_process_tree(self._proc.pid)
self._proc = None
def acquire_endpoint(
server_args: ServerArgs,
base_url: str = "",
launch_server_func: Callable = launch_server,
) -> BenchEndpoint:
"""Resolve the benchmark target -- the single launch-vs-connect decision.
base_url given: connect to it (server_args is ignored). base_url empty:
launch a server from server_args. Caller must close() the result.
"""
if base_url:
ignored = [
f"--{name}"
for name in ("host", "port")
if getattr(server_args, name) != _SERVER_ARGS_DEFAULTS[name]
]
if ignored:
print(
f"WARNING: --base-url is set, so {' / '.join(ignored)} (and other "
f"launch args) are ignored; benchmarking the server at {base_url}."
)
return BenchEndpoint(base_url=base_url)
proc, url = launch_or_reuse_server(launch_server_func, server_args)
return BenchEndpoint(base_url=url, _proc=proc)
python/sglang/benchmark/offline_throughput.py
+586
View File
@@ -0,0 +1,586 @@
"""
Benchmark the throughput in the offline mode.
It accepts server arguments (the same as launch_server.py) and benchmark arguments (the same as serving.py).
# Usage
## Sharegpt dataset with default args
python -m sglang.benchmark.offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --num-prompts 10
## Random dataset with default args
python -m sglang.benchmark.offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024
## Random dataset with profiling args
SGLANG_TORCH_PROFILER_DIR=/tmp python -m sglang.benchmark.offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --dataset-name random --random-input 128 --random-output 128 --num-prompts 4 --max-running-requests 4 --profile-steps 3 --profile --profile-activities "CPU" "XPU"
"""
import argparse
import asyncio
import dataclasses
import inspect
import json
import logging
import os
import random
import time
from typing import Dict, List, Optional, Tuple
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 = 42
disable_ignore_eos: bool = False
extra_request_body: Optional[str] = None
apply_chat_template: bool = False
profile: bool = False
profile_activities: Tuple[str] = ("CPU", "GPU")
profile_steps: Optional[int] = None
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=42, 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(
"--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-steps",
type=int,
default=None,
help="Number of steps to profile. If not specified, profiles all steps.",
)
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,
profile_activities=None,
profile_steps=None,
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)
known_files = None
backend.start_profile(
num_steps=profile_steps,
activities=profile_activities,
)
if profile_steps:
dir = os.getenv("SGLANG_TORCH_PROFILER_DIR")
known_files = set(os.listdir(dir))
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")
if not profile_steps:
known_files = set(os.listdir(dir))
# With --profile-steps the scheduler auto-stops mid-run after N steps, so
# a second stop here raises "not in progress"; a run shorter than N steps
# never hit the target and still needs this explicit stop. Either way we
# must stop before monitor_trace_file, which loops forever waiting for a
# trace that would otherwise never be finalized.
try:
backend.stop_profile()
except RuntimeError:
pass
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(bench_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,
profile_activities=bench_args.profile_activities,
profile_steps=bench_args.profile_steps,
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
def cli_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
if __name__ == "__main__":
cli_main()
python/sglang/benchmark/one_batch.py
+1023
View File
File diff suppressed because it is too large Load Diff
python/sglang/benchmark/one_batch_server.py
+1253
View File
File diff suppressed because it is too large Load Diff
python/sglang/benchmark/serving.py
+2703
View File
File diff suppressed because it is too large Load Diff
python/sglang/benchmark/utils.py
+158
View File
@@ -0,0 +1,158 @@
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 != ""
)
from sglang.srt.utils.hf_transformers_utils import (
get_processor as _srt_get_processor,
)
if not pretrained_model_name_or_path.endswith(
(".json", ".model")
) and not os.path.exists(pretrained_model_name_or_path):
pretrained_model_name_or_path = get_model(pretrained_model_name_or_path)
return _srt_get_processor(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}")
python/sglang/check_env.py
+595
View File
@@ -0,0 +1,595 @@
"""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, "tools/bisheng_compiler/bin/bisheng")
if not os.path.isfile(bisheng):
# Check path for old CANN version
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()
python/sglang/cli/__init__.py
View File
python/sglang/cli/generate.py
+33
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}"
)
python/sglang/cli/killall.py
Executable
+457
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::|sglang serve"
)
# 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())
python/sglang/cli/main.py
+46
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)
python/sglang/cli/serve.py
+143
View File
@@ -0,0 +1,143 @@
# 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 _normalize_positional_model_path(extra_argv):
"""Allow `sglang serve <model>` while preserving existing flag parsing."""
if extra_argv and not extra_argv[0].startswith("-"):
return ["--model-path", extra_argv[0], *extra_argv[1:]], True
return extra_argv, False
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-name-or-path> [additional-arguments]\n"
" or: 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
from sglang.srt.plugins import load_plugins
load_plugins()
model_type, dispatch_argv = _extract_model_type_override(extra_argv)
dispatch_argv, positional_model_path = _normalize_positional_model_path(
dispatch_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)
if positional_model_path:
parsed_args._sglang_explicit_arg_names = {"model_path"}
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)
python/sglang/cli/utils.py
+146
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 in ("--model-path", "--model"):
if i + 1 < len(extra_argv):
model_path = extra_argv[i + 1]
break
elif arg.startswith("--model-path=") or arg.startswith("--model="):
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"
python/sglang/compile_deep_gemm.py
+224
View File
@@ -0,0 +1,224 @@
"""
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.model_executor.cuda_graph_config import Backend, Phase
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")
server_args = tokenizer_manager.server_args
dp_size = server_args.dp_size
base_ids = [0, 1, 2, 3]
sampling_params = {
"temperature": 0.0,
"max_new_tokens": 8,
"ignore_eos": True,
}
if disaggregation_mode != "null":
input_ids = [list(base_ids) for _ in range(dp_size)]
generate_req_input = GenerateReqInput(
input_ids=input_ids,
sampling_params=sampling_params,
)
generate_req_input.bootstrap_host = [FAKE_BOOTSTRAP_HOST] * dp_size
generate_req_input.bootstrap_room = [
i * (2**63 // dp_size) + (i % server_args.tp_size) for i in range(dp_size)
]
else:
input_ids = (
base_ids if dp_size == 1 else [list(base_ids) for _ in range(dp_size)]
)
generate_req_input = GenerateReqInput(
input_ids=input_ids,
sampling_params=sampling_params,
)
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:
dp_size = server_args.dp_size
base_ids = [0, 1, 2, 3]
payload = {
"sampling_params": {
"max_new_tokens": 8,
"temperature": 0,
},
}
if server_args.disaggregation_mode != "null":
payload["input_ids"] = [list(base_ids) for _ in range(dp_size)]
payload["bootstrap_host"] = [FAKE_BOOTSTRAP_HOST] * dp_size
payload["bootstrap_room"] = [
i * (2**63 // dp_size) + (i % server_args.tp_size)
for i in range(dp_size)
]
else:
payload["input_ids"] = (
base_ids
if dp_size == 1
else [list(base_ids) for _ in range(dp_size)]
)
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. Writes after
# ServerArgs.__post_init__ don't propagate to cuda_graph_config via the
# legacy disable_cuda_graph field, so flip both phases directly.
server_args.cuda_graph_config[Phase.DECODE].backend = Backend.DISABLED
server_args.cuda_graph_config[Phase.PREFILL].backend = Backend.DISABLED
print(f"Disable CUDA Graph and Torch Compile to save time...")
# Watchdog timeout follows compile_args.timeout because compilation takes long.
server_args.override(
"compile_deep_gemm.refine_server_args",
enable_torch_compile=False,
watchdog_timeout=compile_args.timeout,
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)
python/sglang/eval/llama3_eval.py
+315
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)
python/sglang/eval/loogle_eval.py
+164
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)
python/sglang/global_config.py
+29
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()
python/sglang/jit_kernel/.clang-format
+25
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
python/sglang/jit_kernel/__init__.py
View File
python/sglang/jit_kernel/__main__.py
+99
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()
python/sglang/jit_kernel/activation.py
+168
View File
@@ -0,0 +1,168 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.jit_kernel.utils import (
cache_once,
get_jit_cuda_arch,
is_arch_support_pdl,
is_hip_runtime,
load_jit,
make_cpp_args,
)
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
from tvm_ffi.module import Module
def _fast_math_flags() -> list[str]:
# Mirrors sgl-kernel's CMake policy: fast-math on SM90, precise on
# SM100+ (Blackwell needs bit-exact expf), off on HIP (clang rejects).
if is_hip_runtime():
return []
if get_jit_cuda_arch().major >= 10:
return []
return ["--use_fast_math"]
@cache_once
def _jit_activation_module(dtype: torch.dtype) -> Module:
args = make_cpp_args(dtype, is_arch_support_pdl())
return load_jit(
"activation",
*args,
cuda_files=["elementwise/activation.cuh"],
extra_cuda_cflags=_fast_math_flags(),
cuda_wrappers=[
("run_activation", f"ActivationKernel<{args}>::run_activation"),
(
"run_activation_filtered",
f"ActivationKernel<{args}>::run_activation_filtered",
),
(
"run_unary_activation",
f"ActivationKernel<{args}>::run_unary_activation",
),
],
)
SUPPORTED_ACTIVATIONS = {"silu", "gelu", "gelu_tanh"}
SUPPORTED_UNARY_ACTIVATIONS = {"relu2"}
@register_custom_op(mutates_args=["out"])
def _run_activation_inplace(
op_name: str, input: torch.Tensor, out: torch.Tensor
) -> None:
hidden_size = input.shape[-1] // 2
module = _jit_activation_module(input.dtype)
input_2d = input.view(-1, hidden_size * 2)
out_2d = out.view(-1, hidden_size)
module.run_activation(input_2d, out_2d, op_name)
@register_custom_op(mutates_args=["out"])
def _run_activation_filtered_inplace(
op_name: str,
input: torch.Tensor,
out: torch.Tensor,
expert_ids: torch.Tensor,
expert_step: int,
) -> None:
hidden_size = input.shape[-1] // 2
module = _jit_activation_module(input.dtype)
input_2d = input.view(-1, hidden_size * 2)
out_2d = out.view(-1, hidden_size)
module.run_activation_filtered(input_2d, out_2d, expert_ids, expert_step, op_name)
def run_activation(
op_name: str,
input: torch.Tensor,
out: Optional[torch.Tensor],
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
"""Apply ``op_name`` activation followed by element-wise multiplication.
When ``expert_ids`` is provided, output rows are skipped for tokens whose
routed expert id is ``-1``. ``expert_step`` is 1 for per-token routing and
``BLOCK_SIZE_M`` for sorted/TMA routing — i.e. ``expert_ids[token_id //
expert_step]`` is consulted before computing each row.
"""
assert op_name in SUPPORTED_ACTIVATIONS, f"Unsupported activation: {op_name}"
hidden_size = input.shape[-1] // 2
if out is None:
out = input.new_empty(*input.shape[:-1], hidden_size)
if expert_ids is None:
_run_activation_inplace(op_name, input, out)
else:
_run_activation_filtered_inplace(op_name, input, out, expert_ids, expert_step)
return out
@register_custom_op(mutates_args=["out"])
def _run_unary_activation_inplace(
op_name: str, input: torch.Tensor, out: torch.Tensor
) -> None:
last = input.shape[-1]
module = _jit_activation_module(input.dtype)
module.run_unary_activation(input.view(-1, last), out.view(-1, last), op_name)
def run_unary_activation(
op_name: str,
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Apply a standalone (non-gated) element-wise activation: ``out = act(input)``.
Unlike :func:`run_activation`, there is no gate/up split — ``input`` and
``out`` share the same shape.
"""
assert (
op_name in SUPPORTED_UNARY_ACTIVATIONS
), f"Unsupported unary activation: {op_name}"
if out is None:
out = torch.empty_like(input)
_run_unary_activation_inplace(op_name, input, out)
return out
def relu2(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Squared ReLU: ``out = max(0, input) ** 2`` (element-wise)."""
return run_unary_activation("relu2", input, out)
def silu_and_mul(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
return run_activation("silu", input, out, expert_ids, expert_step)
def gelu_and_mul(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
return run_activation("gelu", input, out, expert_ids, expert_step)
def gelu_tanh_and_mul(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
expert_ids: Optional[torch.Tensor] = None,
expert_step: int = 1,
) -> torch.Tensor:
return run_activation("gelu_tanh", input, out, expert_ids, expert_step)
python/sglang/jit_kernel/add_constant.py
+28
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
python/sglang/jit_kernel/all_reduce.py
+243
View File
@@ -0,0 +1,243 @@
from __future__ import annotations
import enum
from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, cast
import torch
import tvm_ffi
from tvm_ffi import Module
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 get_graph_capture_bases(
self,
) -> Tuple[List[Tuple[int, int]], List[List[int]], List[int]]: ...
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) -> Module:
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) -> Module:
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 _jit_fused_parallel_qknorm_module(
dtype: torch.dtype, world_size: int, q_dim: int, k_dim: int
) -> Module:
args = make_cpp_args(dtype, world_size, q_dim, k_dim, is_arch_support_pdl())
cls_name = f"FusedParallelQKNormAcrossHead<{args}>"
return load_jit(
"tp_qknorm",
*args,
extra_ldflags=["-lcuda"],
cuda_files=["distributed/tp_qknorm.cuh"],
cuda_wrappers=[
("fused_parallel_qknorm", f"{cls_name}::run"),
("get_max_occupancy", f"{cls_name}::get_max_occupancy"),
],
)
@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:
max_pull_blocks = NUM_CTA if max_pull_blocks is None else max_pull_blocks
max_push_blocks = NUM_CTA if max_push_blocks is None else max_push_blocks
self.__ffi_init__(
rank,
world_size,
max_pull_blocks,
max_push_blocks,
pull_buffer_bytes,
push_buffer_bytes,
graph_input_count,
)
self._world_size = world_size
self._pull_config = ConfigResult(min(NUM_CTA, max_pull_blocks), MAX_THREADS)
if max_pull_blocks > 0: # special case: cannot configure 0 blocks
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)
def get_fused_parallel_qknorm_max_occupancy(
dtype: torch.dtype, world_size: int, q_dim: int, k_dim: int
) -> int:
module = _jit_fused_parallel_qknorm_module(dtype, world_size, q_dim, k_dim)
return module.get_max_occupancy()
def fused_parallel_qknorm(
custom_ar: CustomAllReduceObj,
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
eps: float = 1e-6,
) -> None:
world_size = custom_ar.world_size
q_dim = q.shape[-1] * world_size
k_dim = k.shape[-1] * world_size
module = _jit_fused_parallel_qknorm_module(q.dtype, world_size, q_dim, k_dim)
module.fused_parallel_qknorm(custom_ar, q, k, q_weight, k_weight, eps)
python/sglang/jit_kernel/awq_dequantize.py
+38
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
python/sglang/jit_kernel/awq_marlin_repack.py
+59
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
python/sglang/jit_kernel/benchmark/kv_canary/utils.py
+404
View File
@@ -0,0 +1,404 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Callable
import torch
from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES, RealKvSource
BS_AXIS: list[int] = [1, 4, 32, 128, 256, 1024]
PREFIX_AXIS: list[int] = [0, 128, 1024, 4096, 10240, 16384]
EXTEND_LEN_AXIS: list[int] = [128, 512, 4096, 16384]
POOL_AXIS: list[str] = ["full", "swa_window_128"]
REAL_KV_AXIS: list[str] = ["none", "small_1src", "med_2src", "max_4src"]
HASH_MODE_AXIS: list[str] = ["none", "partial", "all"]
SWA_WINDOW: int = 128
RING_CAPACITY: int = 256
MAX_EXTEND_TOKENS_PER_FORWARD: int = 4096
@dataclass(frozen=True, slots=True, kw_only=True)
class BenchCase:
scenario: str
bs: int
prefix_len: int
mode: str
extend_len: int
pool_kind: str
real_kv_kind: str
hash_mode: str
@property
def case_id(self) -> str:
return (
f"{self.scenario}_bs{self.bs}_prefix{self.prefix_len}_{self.mode}{self.extend_len}"
f"_{self.pool_kind}_rkv{self.real_kv_kind}_hash{self.hash_mode}"
)
def _case(
*,
scenario: str,
bs: int,
prefix_len: int,
mode: str,
extend_len: int,
pool_kind: str,
real_kv_kind: str = "none",
hash_mode: str = "none",
) -> BenchCase:
return BenchCase(
scenario=scenario,
bs=bs,
prefix_len=prefix_len,
mode=mode,
extend_len=extend_len,
pool_kind=pool_kind,
real_kv_kind=real_kv_kind,
hash_mode=hash_mode,
)
def _is_realistic_extend_case(case: BenchCase) -> bool:
if case.mode != "extend":
return True
return case.bs * case.extend_len <= MAX_EXTEND_TOKENS_PER_FORWARD
def _dedupe_cases(cases: list[BenchCase]) -> list[BenchCase]:
seen: set[str] = set()
result: list[BenchCase] = []
for case in cases:
if case.case_id in seen:
continue
seen.add(case.case_id)
result.append(case)
return result
def build_fast_matrix_cases() -> list[BenchCase]:
return _dedupe_cases(
[
_case(
scenario="smoke_decode_empty",
bs=1,
prefix_len=0,
mode="decode",
extend_len=1,
pool_kind="full",
),
_case(
scenario="small_extend_batch",
bs=32,
prefix_len=4096,
mode="extend",
extend_len=128,
pool_kind="full",
),
_case(
scenario="e2e_decode_steady",
bs=256,
prefix_len=4096,
mode="decode",
extend_len=1,
pool_kind="full",
),
_case(
scenario="decode_large_batch_short_prefix",
bs=1024,
prefix_len=1024,
mode="decode",
extend_len=1,
pool_kind="full",
),
_case(
scenario="e2e_prefill_chunk_first",
bs=1,
prefix_len=0,
mode="extend",
extend_len=4096,
pool_kind="full",
),
_case(
scenario="e2e_prefill_chunk_mid",
bs=1,
prefix_len=8192,
mode="extend",
extend_len=4096,
pool_kind="full",
),
_case(
scenario="e2e_prefill_chunk_last",
bs=1,
prefix_len=12288,
mode="extend",
extend_len=4096,
pool_kind="full",
),
_case(
scenario="e2e_decode_tail",
bs=1,
prefix_len=5120,
mode="decode",
extend_len=1,
pool_kind="full",
),
_case(
scenario="swa_decode_long_prefix",
bs=128,
prefix_len=10240,
mode="decode",
extend_len=1,
pool_kind="swa_window_128",
),
_case(
scenario="small_extend_single_req",
bs=1,
prefix_len=128,
mode="extend",
extend_len=128,
pool_kind="full",
),
_case(
scenario="medium_extend_chunk",
bs=4,
prefix_len=1024,
mode="extend",
extend_len=512,
pool_kind="full",
),
_case(
scenario="decode_mid_batch",
bs=128,
prefix_len=4096,
mode="decode",
extend_len=1,
pool_kind="full",
),
_case(
scenario="e2e_prefill_chunk_second",
bs=1,
prefix_len=4096,
mode="extend",
extend_len=4096,
pool_kind="full",
),
_case(
scenario="swa_decode_short_prefix",
bs=256,
prefix_len=128,
mode="decode",
extend_len=1,
pool_kind="swa_window_128",
),
_case(
scenario="swa_decode_tail",
bs=4,
prefix_len=10240,
mode="decode",
extend_len=1,
pool_kind="swa_window_128",
),
_case(
scenario="small_extend_batch_hash",
bs=32,
prefix_len=4096,
mode="extend",
extend_len=128,
pool_kind="full",
real_kv_kind="small_1src",
hash_mode="partial",
),
_case(
scenario="e2e_prefill_chunk_hash",
bs=1,
prefix_len=12288,
mode="extend",
extend_len=4096,
pool_kind="full",
real_kv_kind="med_2src",
hash_mode="all",
),
_case(
scenario="e2e_decode_steady_hash",
bs=256,
prefix_len=4096,
mode="decode",
extend_len=1,
pool_kind="full",
real_kv_kind="max_4src",
hash_mode="all",
),
_case(
scenario="swa_decode_long_prefix_hash",
bs=128,
prefix_len=10240,
mode="decode",
extend_len=1,
pool_kind="swa_window_128",
real_kv_kind="med_2src",
hash_mode="partial",
),
_case(
scenario="smoke_decode_empty_hash",
bs=1,
prefix_len=0,
mode="decode",
extend_len=1,
pool_kind="full",
real_kv_kind="small_1src",
hash_mode="all",
),
]
)
def build_full_matrix_cases() -> list[BenchCase]:
"""Full matrix plus targeted e2e points.
Extend cases are pruned to a maximum token chunk per forward because the scheduler chunks long
prefills; for example, a 4096-token extend is represented as ``bs=1``, not ``bs=32``.
"""
fast = build_fast_matrix_cases()
fast_keys = {c.case_id for c in fast}
full: list[BenchCase] = list(fast)
for bs in BS_AXIS:
for prefix_len in PREFIX_AXIS:
for pool_kind in POOL_AXIS:
for mode, extend_len in (
("decode", 1),
*(("extend", e) for e in EXTEND_LEN_AXIS),
):
case = _case(
scenario="matrix",
bs=bs,
prefix_len=prefix_len,
mode=mode,
extend_len=extend_len,
pool_kind=pool_kind,
)
if not _is_realistic_extend_case(case):
continue
if case.case_id in fast_keys:
continue
full.append(case)
fast_base_points = [
(c.bs, c.prefix_len, c.mode, c.extend_len, c.pool_kind)
for c in fast
if c.real_kv_kind == "none" and c.hash_mode == "none"
]
for bs, prefix_len, mode, extend_len, pool_kind in fast_base_points:
for hash_mode in HASH_MODE_AXIS:
if hash_mode == "none":
continue
for real_kv_kind in REAL_KV_AXIS:
if real_kv_kind == "none":
continue
case = _case(
scenario="fold_matrix",
bs=bs,
prefix_len=prefix_len,
mode=mode,
extend_len=extend_len,
pool_kind=pool_kind,
real_kv_kind=real_kv_kind,
hash_mode=hash_mode,
)
if not _is_realistic_extend_case(case):
continue
if case.case_id in fast_keys:
continue
full.append(case)
fast_keys.add(case.case_id)
return full
def cases_to_x_vals(
cases: list[BenchCase],
) -> list[tuple[str, int, int, str, int, str, str, str]]:
return [
(
c.scenario,
c.bs,
c.prefix_len,
c.mode,
c.extend_len,
c.pool_kind,
c.real_kv_kind,
c.hash_mode,
)
for c in cases
]
def _one_real_kv_source(
*, num_slots: int, num_bytes: int, read_bytes: int, device: torch.device
) -> RealKvSource:
tensor = torch.zeros(max(1, num_slots), num_bytes, dtype=torch.uint8, device=device)
return RealKvSource(
tensor=tensor,
page_size=1,
num_bytes_per_token=num_bytes,
read_bytes=read_bytes,
)
def make_real_kv_sources(
*, kind: str, num_slots: int, device: torch.device
) -> tuple[RealKvSource, ...]:
"""Map a ``real_kv_kind`` axis label to a tuple of ``RealKvSource`` configs.
Byte-volume ladder (none -> small_1src -> med_2src -> max_4src) so the bench exposes the
``real_kv_fold_sources`` PARTIAL/ALL cost gradient. ``max_4src`` hits the
``consts.MAX_REAL_KV_SOURCES = 4`` ABI ceiling.
"""
if kind == "none":
return ()
if kind == "small_1src":
return (
_one_real_kv_source(
num_slots=num_slots, num_bytes=16, read_bytes=16, device=device
),
)
if kind == "med_2src":
return tuple(
_one_real_kv_source(
num_slots=num_slots, num_bytes=32, read_bytes=16, device=device
)
for _ in range(2)
)
if kind == "max_4src":
return tuple(
_one_real_kv_source(
num_slots=num_slots, num_bytes=64, read_bytes=32, device=device
)
for _ in range(4)
)
raise ValueError(f"kv-canary bench: unknown real_kv_kind {kind!r}")
def naive_slot_copy_fn(*, total: int, device: torch.device) -> Callable[[], None]:
n_slots = max(total, 1)
payload = torch.zeros(n_slots, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device)
sink = torch.zeros_like(payload)
indices = torch.arange(n_slots, device=device, dtype=torch.int64) % sink.shape[0]
def baseline() -> None:
sink.index_copy_(0, indices, payload)
return baseline
def naive_cumsum_fn(*, bs: int, device: torch.device) -> Callable[[], None]:
counts = torch.zeros(max(bs, 1), dtype=torch.int32, device=device)
def baseline() -> None:
torch.cumsum(counts, dim=0)
return baseline
python/sglang/jit_kernel/benchmark/marker.py
+554
View File
@@ -0,0 +1,554 @@
import contextlib
import inspect
import itertools
import math
import os
from typing import (
Any,
Callable,
ContextManager,
Dict,
Generic,
Iterable,
List,
Literal,
NamedTuple,
Optional,
Tuple,
TypeAlias,
TypeVar,
)
import torch
from sglang.jit_kernel.utils import cache_once
from sglang.utils import is_in_ci
F = TypeVar("F", bound=Callable[..., "BenchResult"])
Metric: TypeAlias = "float | Literal['avg']"
BENCH_CONFIG: TypeAlias = "List[Tuple[Tuple[str, ...], List[Tuple[Any, ...]]]]"
UNIT_SCALE = {"us": 1e-6, "ms": 1e-3, "s": 1.0}
TYPE_LIST = (bool, int, float, str, torch.dtype, torch.device, None.__class__)
DISABLE_LOG_BANDWIDTH = os.environ.get("SGLANG_KERNEL_DISABLE_LOG_BANDWIDTH") == "1"
__all__ = [
"BenchResult",
"BenchSkip",
"Benchmark",
"benchmark",
"parametrize",
"do_bench",
"skip",
]
class BenchSkip(Exception):
pass
def skip(reason: str):
raise BenchSkip(reason)
@cache_once
def _get_benchmark_stream(device_id: int) -> torch.cuda.Stream:
return torch.cuda.Stream(device=device_id)
def _clone_recursive(in_: Any) -> Any:
if isinstance(in_, torch.Tensor):
return in_.clone()
elif isinstance(in_, (list, tuple)):
return type(in_)(_clone_recursive(x) for x in in_)
elif isinstance(in_, dict):
return {k: _clone_recursive(v) for k, v in in_.items()}
elif isinstance(in_, TYPE_LIST):
return in_
# NOTE: avoid silent error
raise ValueError(f"unsupported type: {type(in_)}")
def _get_nbytes_recursive(in_: Any) -> int:
if isinstance(in_, torch.Tensor):
return in_.nbytes
elif isinstance(in_, (list, tuple)):
return sum(_get_nbytes_recursive(x) for x in in_)
elif isinstance(in_, dict):
return sum(_get_nbytes_recursive(v) for v in in_.values())
elif isinstance(in_, TYPE_LIST):
return 0
# NOTE: avoid silent error
raise ValueError(f"unsupported type: {type(in_)}")
def _process_metrics(times: list[float], metrics: tuple[Metric, ...]) -> list[float]:
results: list[float] = []
times = sorted(x / 1000 for x in times) # convert to seconds and sort
for metric in metrics:
if metric == "avg":
results.append(sum(times) / len(times))
else:
assert 0 <= metric <= 1, f"invalid metric: {metric}"
which = min(int(len(times) * metric), len(times) - 1)
results.append(times[which])
return results
@cache_once
def _get_l2_cache_size() -> int:
device = torch.cuda.current_device()
props = torch.cuda.get_device_properties(device)
return props.L2_cache_size
_L2_SAFE_RATIO = 5
def _get_flush_l2_buffer() -> torch.Tensor:
"""Get a buffer sized to flush the L2 cache when accessed."""
device = torch.device(f"cuda:{torch.cuda.current_device()}")
l2_size = _get_l2_cache_size()
safe_size = int(l2_size * _L2_SAFE_RATIO)
return torch.empty(safe_size, device=device, dtype=torch.uint8)
def _calculate_rotation_count(nbytes: int, min_rotations: int = 2) -> int:
"""
Adapted from flashinfer benchmark utility:
https://github.com/flashinfer-ai/flashinfer/blob/c5a2b06edae4fa2bfd2ae25eed16eb565c70513f/flashinfer/testing/utils.py
Calculate the number of buffer copies needed to ensure cold L2 cache.
The function uses conservative thresholds to account for:
- LRU eviction being gradual (not all data evicted when capacity exceeded)
- Cache associativity effects (some data may persist in non-conflicting sets)
- Hardware prefetching behavior
Returns 1 (no rotation needed) only when tensor size substantially exceeds
L2 cache, ensuring cache effects are truly negligible.
Args:
tensors: List of tensors to consider for rotation (must be on GPU).
device: Device for L2 cache query (None for current device).
min_rotations: Minimum number of rotations when rotation is needed.
Returns:
Number of buffer copies needed (1 means no rotation needed).
"""
l2_size = _get_l2_cache_size()
safe_cache_threshold = l2_size * _L2_SAFE_RATIO
if nbytes <= 0 or nbytes >= safe_cache_threshold:
return 1 # No tensors to rotate
# Conservative formula: ensure between any two uses of the same buffer,
# we've accessed enough data to fully flush L2 with margin
# Using safe_cache_threshold ensures we account for all cache effects
num_rotations = math.ceil(safe_cache_threshold / nbytes) + 1
return max(min_rotations, num_rotations)
class BenchResult(NamedTuple):
metrics: Tuple[Metric, ...]
times: List[float] # in seconds
memory_footprint: Optional[int]
class Table:
"""Aligned text table with `|` section separators and `=`/`-` rules."""
SEP = " | "
def __init__(self) -> None:
self._headers: List[str] = []
self._mins: List[int] = []
self._pads: List[int] = []
self._aligns: List[str] = []
self._seps: set = set()
self._rows: List[List[str]] = []
@staticmethod
def format_latency(r: float) -> str:
if math.isnan(r):
return "N/A"
length = len(str(int(r)))
if length < 5:
return f"{r:.4f}"
# decrease number of the digits
digits = max(0, 4 - (length - 5))
return f"{r:.{digits}f}"
@staticmethod
def format_bandwidth(b: float) -> str:
if math.isnan(b):
return "N/A"
return f"{b:.2f}"
def col(
self,
header: str = "",
*,
min_width: int = 10,
pad: int = 2,
align: str = ">",
) -> None:
self._headers.append(header)
self._mins.append(min_width)
self._pads.append(pad)
self._aligns.append(align)
def sep(self) -> None:
self._seps.add(len(self._headers))
def row(self, *cells: Any) -> None:
assert len(cells) == len(self._headers)
self._rows.append([str(c) for c in cells])
def print(self) -> None:
widths = [
max(max(len(c) + p for c in [h, *(r[i] for r in self._rows)]), mw)
for i, (h, mw, p) in enumerate(zip(self._headers, self._mins, self._pads))
]
total = sum(widths) + len(self.SEP) * len(self._seps)
def fmt(cells: List[str]) -> str:
parts: List[str] = []
for i, (cell, w, a) in enumerate(zip(cells, widths, self._aligns)):
if i in self._seps:
parts.append(self.SEP)
parts.append(f"{cell:{a}{w}}")
return "".join(parts)
print("=" * total)
print(fmt(self._headers))
print("-" * total)
for r in self._rows:
print(fmt(r))
print("=" * total)
class Benchmark(Generic[F]):
def __init__(self, fn: F, line_arg: str, line_vals: List[Any], *, unit: str):
assert unit in UNIT_SCALE and len(set(line_vals)) == len(line_vals) > 0
self._fn = fn
self._line_arg = line_arg
self._line_vals = line_vals
self._unit = unit
self._configs: BENCH_CONFIG = []
self._fn_params = inspect.signature(fn).parameters
self._unit_scale = UNIT_SCALE[unit]
assert line_arg in self._fn_params, (
f"line_arg {line_arg!r} is not a parameter of {fn.__name__}; "
f"available: {list(self._fn_params)}"
)
self._seen_args = {line_arg}
def add_config(self, names: Tuple[str, ...], vals: List[Tuple[Any, ...]]) -> None:
"""Prepend a parametrize axis. Validates that names are real parameters
of the benchmark fn, and rejects duplicates / collisions with line_arg."""
assert len(names) > 0, "parametrize: must provide at least one name"
for name in names:
assert name in self._fn_params, (
f"parametrize name {name!r} is not a parameter of "
f"{self._fn.__name__}; available: {list(self._fn_params)}"
)
assert (
name not in self._seen_args
), f"parametrize name {name!r} is already used"
self._seen_args.add(name)
self._configs.insert(0, (names, vals))
def _collect_results(self) -> Tuple[List[List[float]], List[List[float]], bool]:
axis_names = [n for n, _ in self._configs]
axis_vals = [v for _, v in self._configs]
results: List[List[float]] = []
bandwidth_results: List[List[float]] = []
should_log_bandwidth = False
for system in self._line_vals:
latencies: List[float] = []
bandwidths: List[float] = []
for combo in itertools.product(*axis_vals):
kwargs: Dict[str, Any] = {self._line_arg: system}
for names, values in zip(axis_names, combo):
kwargs.update(zip(names, values))
try:
result = self._fn(**kwargs)
except BenchSkip:
latencies.append(float("nan"))
if not DISABLE_LOG_BANDWIDTH:
bandwidths.append(float("nan"))
continue
latencies.append(result.times[0] / self._unit_scale)
if not DISABLE_LOG_BANDWIDTH and result.memory_footprint is not None:
should_log_bandwidth = True
bandwidths.append(
result.memory_footprint / (1024**3) / result.times[0]
)
results.append(latencies)
bandwidth_results.append(bandwidths)
return results, bandwidth_results, should_log_bandwidth
def run(self) -> None:
# Pre-check: every required fn param must be covered.
flat_names = [n for names, _ in self._configs for n in names]
kinds = (
inspect.Parameter.POSITIONAL_OR_KEYWORD,
inspect.Parameter.KEYWORD_ONLY,
)
missing = {
n
for n, p in self._fn_params.items()
if p.default is inspect.Parameter.empty and p.kind in kinds
} - (set(flat_names) | {self._line_arg})
assert not missing, (
f"parameters not parametrized for {self._fn.__name__}: "
f"{sorted(missing)}"
)
results, bandwidths, should_log_bw = self._collect_results()
table = Table()
table.col(min_width=0, pad=0, align="<") # id column (tight, left-aligned)
for name in flat_names:
table.col(name)
table.sep()
for system in self._line_vals:
table.col(f"{system}({self._unit})", min_width=15)
if should_log_bw:
table.sep()
for system in self._line_vals:
table.col(f"{system}(GB/s)", min_width=15)
axis_vals = [v for _, v in self._configs]
for row_id, combo in enumerate(itertools.product(*axis_vals)):
cells: List[Any] = [row_id]
cells.extend(v for vt in combo for v in vt)
cells.extend(table.format_latency(r[row_id]) for r in results)
if should_log_bw:
cells.extend(table.format_bandwidth(r[row_id]) for r in bandwidths)
table.row(*cells)
table.print()
def benchmark(line_arg: str, line_vals: List[Any], *, unit: str = "us"):
def decorator(fn: F) -> Benchmark[F]:
return Benchmark(fn, line_arg, line_vals, unit=unit)
return decorator
def parametrize(names: str, vals: List[Any], ci_vals: Optional[List[Any]] = None):
"""Add a parametrize axis. Pytest-style:
- Single name: `parametrize("dim", [1024, 4096])`
- Multiple names (correlated):
`parametrize("h,d", [(1, 64), (2, 128)])`
For multi-name axes, each value must be a tuple/list of matching length.
"""
name_tuple = tuple(n.strip() for n in names.split(","))
assert all(name_tuple), f"parametrize: empty name in {names!r}"
arity = len(name_tuple)
def _normalize(vs: List[Any]) -> List[Tuple[Any, ...]]:
if arity == 1:
return [(v,) for v in vs]
out: List[Tuple[Any, ...]] = []
for v in vs:
assert isinstance(
v, (tuple, list)
), f"parametrize: multi-name values must be tuples, got {v!r}"
t = tuple(v)
assert (
len(t) == arity
), f"parametrize: each value must have length {arity}, got {t!r}"
out.append(t)
return out
def decorator(bench: Benchmark[F]) -> Benchmark[F]:
chosen = ci_vals if (ci_vals is not None and is_in_ci()) else vals
bench.add_config(name_tuple, _normalize(chosen))
return bench
return decorator
def _do_bench_internal_graph(
fn: Callable,
replay_iters: int,
input_args: Tuple[Any, ...],
input_kwargs: Dict[str, Any],
graph_clone_args: Iterable[int],
graph_clone_kwargs: Iterable[str],
graph_context: ContextManager,
sync_multigpu_fn: Callable[[], Any],
) -> List[float]:
result: List[float] = []
stream = torch.cuda.current_stream()
empty_tensor = _get_flush_l2_buffer()
# only count the cloned tensors for rotation count
nbytes = sum(_get_nbytes_recursive(input_args[i]) for i in graph_clone_args)
nbytes += sum(_get_nbytes_recursive(input_kwargs[k]) for k in graph_clone_kwargs)
rotate_count = min(_calculate_rotation_count(nbytes), 100)
loop_count = math.ceil(100 / rotate_count) * rotate_count
input_args_list = [input_args] * rotate_count
input_kwargs_list = [input_kwargs] * rotate_count
graph_clone_args = set(graph_clone_args)
graph_clone_kwargs = set(graph_clone_kwargs)
graph = torch.cuda.CUDAGraph()
# NOTE: we rotate the buffer here to avoid L2 cache effect
for i in range(1, rotate_count):
input_args_list[i] = tuple(
(
_clone_recursive(input_args[j])
if j in graph_clone_args
else input_args[j]
)
for j in range(len(input_args))
)
input_kwargs_list[i] = dict(
(k, (_clone_recursive(v) if k in graph_clone_kwargs else v))
for k, v in input_kwargs.items()
)
with graph_context:
with torch.cuda.graph(graph, stream=stream):
for i in range(loop_count):
args = input_args_list[i % rotate_count]
kwargs = input_kwargs_list[i % rotate_count]
fn(*args, **kwargs)
# warm up the graph once
graph.replay()
# then replay the graph and measure the time
tic = torch.cuda.Event(enable_timing=True)
toc = torch.cuda.Event(enable_timing=True)
for _ in range(max(replay_iters // loop_count, 10)):
empty_tensor.zero_() # cold the L2 cache
sync_multigpu_fn() # sync GPU before each iteration for precise timing
tic.record(stream)
graph.replay()
toc.record(stream)
stream.synchronize()
result.append(tic.elapsed_time(toc) / loop_count)
return result
def do_bench(
fn: Callable,
*,
input_args: Tuple[Any, ...] = (),
input_kwargs: Dict[str, Any] = {},
use_cuda_graph: bool = True,
warmup_iters: int = 50,
replay_iters: int = 1000,
metrics: Tuple[Metric, ...] = (0.5, "avg"),
stream: torch.cuda.Stream | None = None,
# NOTE: should only clone the read args to avoid L2 cache effect in cuda graph
graph_clone_args: Iterable[int] | Literal["all"] | None = "all",
graph_clone_kwargs: Iterable[str] | Literal["all"] | None = "all",
# NOTE: for memory-bandwidth profiling
disable_log_bandwidth: bool = DISABLE_LOG_BANDWIDTH,
memory_args: Iterable[Any] | Literal["all"] | None = "all",
memory_output: Iterable[Any] | Literal["out"] | None = "out",
extra_memory_args: Iterable[Any] | None = None,
extra_memory_footprint: int = 0,
graph_context_fn: Optional[Callable[[], ContextManager]] = None,
sync_multigpu_fn: Optional[Callable[[], Any]] = None,
) -> BenchResult:
"""
Benchmark a function using CUDA graph or naive loop.
Adapted from flashinfer benchmark utility:
https://github.com/flashinfer-ai/flashinfer/blob/c5a2b06edae4fa2bfd2ae25eed16eb565c70513f/flashinfer/testing/utils.py
:param fn: Function to benchmark
:param input_args: Positional arguments to pass to the function
:param input_kwargs: Keyword arguments to pass to the function
:param use_cuda_graph: Whether to use CUDA graph for benchmarking
:param warmup_iters: Number of warm-up iterations to run before benchmarking
:param replay_iters: Number of iterations to run for benchmarking
:param metrics: Metrics to compute from the timing results (quantiles in [0, 1] or "avg")
:param stream: CUDA stream to use for benchmarking (if None, a new stream will be created)
:param graph_clone_args: Indices of input_args to clone for each iteration.
Only the read args need to be cloned to avoid L2 cache effect.
:param graph_clone_kwargs: Keys of input_kwargs to clone for each iteration.
Only the read args need to be cloned to avoid L2 cache effect.
:param disable_log_bandwidth: Whether to disable logging memory bandwidth in the profile report.
:param memory_args: Optional sequence of arguments to calculate total memory footprint.
Used for memory bandwidth estimation in the profile report.
:param memory_output: Arguments whose output memory should be included in the memory footprint.
:param extra_memory_args: Additional arguments to consider for memory footprint calculation.
:param extra_memory_footprint: Additional memory footprint to consider.
This is typically used when the load/store bytes is dynamic.
:param graph_context_fn: A callable returning a context manager that wraps the cuda graph capture.
:param sync_multigpu_fn: A callable to synchronize multiple GPUs before each iteration. For precise
benchmark number in multi-GPU benchmark, it should be some synchronization
primitive on GPU side (not on CPU side).
"""
# first warmup the function
device_id = torch.cuda.current_device()
if stream is None:
stream = _get_benchmark_stream(device_id)
old_current_stream = torch.cuda.current_stream(device_id)
result: List[float] = []
sync_multigpu_fn = sync_multigpu_fn or (lambda: None)
with torch.cuda.device(device_id), torch.cuda.stream(stream):
stream.wait_stream(old_current_stream)
sync_multigpu_fn()
for _ in range(warmup_iters):
fn(*input_args, **input_kwargs)
if use_cuda_graph:
# NOTE: by default, reduce all the CPU-side overhead
if graph_clone_args == "all":
graph_clone_args = range(len(input_args))
elif graph_clone_args is None:
graph_clone_args = []
if graph_clone_kwargs == "all":
graph_clone_kwargs = input_kwargs.keys()
elif graph_clone_kwargs is None:
graph_clone_kwargs = []
graph_context = (
graph_context_fn()
if graph_context_fn is not None
else contextlib.nullcontext()
)
result = _do_bench_internal_graph(
fn,
replay_iters,
input_args,
input_kwargs,
graph_clone_args,
graph_clone_kwargs,
graph_context,
sync_multigpu_fn,
)
else:
# NOTE: no cuda graph, naive loop
tic = torch.cuda.Event(enable_timing=True)
toc = torch.cuda.Event(enable_timing=True)
empty_tensor = _get_flush_l2_buffer()
for _ in range(max(replay_iters, 10)):
empty_tensor.zero_() # cold the L2 cache
sync_multigpu_fn()
tic.record(stream)
fn(*input_args, **input_kwargs)
toc.record(stream)
stream.synchronize()
result.append(tic.elapsed_time(toc))
stream.synchronize()
result = _process_metrics(result, metrics)
memory_footprint = None
if not disable_log_bandwidth:
if memory_args == "all":
memory_args = input_args + tuple(input_kwargs.values())
if memory_output == "out":
memory_output = fn(*input_args, **input_kwargs)
memory_footprint = extra_memory_footprint
memory_footprint += _get_nbytes_recursive(extra_memory_args)
memory_footprint += _get_nbytes_recursive(memory_args)
memory_footprint += _get_nbytes_recursive(memory_output)
return BenchResult(metrics, result, memory_footprint)
python/sglang/jit_kernel/benchmark/utils.py
+108
View File
@@ -0,0 +1,108 @@
"""Common utilities for jit_kernel benchmark files."""
from typing import Callable, List, Optional, Sequence, Tuple
import torch
import triton.testing
from sglang.jit_kernel.mp import multigpu_launch
from sglang.utils import is_in_ci
def multigpu_bench_main(
name: str,
file: str,
num_gpus: Sequence[int],
main_fn: Callable[[], None],
*,
pre_launch_fn: Optional[Callable[[List[int]], None]] = None,
timeout: Optional[int] = None,
) -> None:
"""cudalib-style multi-GPU benchmark entry point.
Drop this at the bottom of a benchmark file::
multigpu_bench_main(
name=__name__,
file=__file__,
num_gpus=range(2, 9),
main_fn=benchmark.run,
)
Mirrors :func:`multigpu_pytest_main` but invokes a caller-supplied function
instead of pytest. ``main_fn`` is expected to return ``None`` on success;
any exception propagates as a non-zero exit. Pass ``--num-gpu 2,4`` on the
command line to override ``num_gpus``.
``pre_launch_fn`` (kw-only) runs once in the outer process before any
torchrun child starts, receiving the runnable world sizes. Use it for
parallel JIT precompilation so torchrun children hit a warm disk cache.
``timeout`` (kw-only, seconds) bounds each per-world-size torchrun
invocation. Defaults to ``None`` (wait indefinitely) since benchmark sweeps
can legitimately run long; set it to fail fast on a hung worker.
"""
def inner() -> int:
main_fn()
return 0
return multigpu_launch(
name,
file,
num_gpus,
env_key="_IS_BENCH_MULTIGPU_SGLANG_JIT_KERNEL",
inner=inner,
kind="benchmark",
pre_launch_fn=pre_launch_fn,
timeout=timeout,
)
# Common constants
DEFAULT_DTYPE = torch.bfloat16
DEFAULT_DEVICE = "cuda"
DEFAULT_QUANTILES = [0.5, 0.2, 0.8]
def create_empty(*shape: int, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE):
return torch.empty(shape, dtype=dtype, device=device)
def create_random(*shape: int, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE):
return torch.randn(shape, dtype=dtype, device=device)
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
python/sglang/jit_kernel/clamp_position.py
+35
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
python/sglang/jit_kernel/concat_mla.py
+65
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
python/sglang/jit_kernel/csrc/add_constant.cuh
+101
View File
@@ -0,0 +1,101 @@
#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 <sgl_kernel/vec.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <cstddef>
#include <cstdint>
namespace {
constexpr size_t kBlockSize = 256;
constexpr size_t kVectorizedMinElements = 1 << 20;
constexpr size_t kVectorBytes = device::kMaxVecBytes;
static_assert(kVectorBytes % sizeof(int32_t) == 0, "Vector byte width must contain whole int32_t elements");
constexpr size_t kElementsPerVector = kVectorBytes / sizeof(int32_t);
template <typename Vector>
bool is_aligned_for_vector(const int32_t* ptr) {
return reinterpret_cast<uintptr_t>(ptr) % alignof(Vector) == 0;
}
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;
}
}
template <int32_t kConstant, size_t kElementsPerVector>
__global__ void add_constant_vectorized_kernel(int32_t* dst, const int32_t* src, size_t length) {
using Vector = device::AlignedVector<int32_t, kElementsPerVector>;
const size_t work_idx = blockIdx.x * blockDim.x + threadIdx.x;
const size_t vector_count = length / kElementsPerVector;
const size_t tail_start = vector_count * kElementsPerVector;
if (work_idx < vector_count) {
auto values = device::load_as<Vector>(src, work_idx);
#pragma unroll
for (size_t i = 0; i < kElementsPerVector; ++i) {
values[i] += kConstant;
}
device::store_as<Vector>(dst, values, work_idx);
} else {
const size_t tail_idx = tail_start + work_idx - vector_count;
if (tail_idx < length) {
dst[tail_idx] = src[tail_idx] + kConstant;
}
}
}
// 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<kDLGPU>(device_) // must be on GPU device (CUDA or ROCm)
.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 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);
const auto* src_ptr = static_cast<const int32_t*>(src.data_ptr());
auto* dst_ptr = static_cast<int32_t*>(dst.data_ptr());
using Vector = device::AlignedVector<int32_t, kElementsPerVector>;
const bool is_vector_aligned = is_aligned_for_vector<Vector>(src_ptr) && is_aligned_for_vector<Vector>(dst_ptr);
// 3. Launch the kernel. Error code will be automatically checked.
if (num_elements >= kVectorizedMinElements && is_vector_aligned) {
const size_t vector_count = num_elements / kElementsPerVector;
const size_t tail_count = num_elements - vector_count * kElementsPerVector;
const size_t work_items = vector_count + tail_count;
const size_t grid_size = div_ceil(work_items, kBlockSize);
LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)(
add_constant_vectorized_kernel<kConstant, kElementsPerVector>, dst_ptr, src_ptr, num_elements);
} else {
const size_t grid_size = div_ceil(num_elements, kBlockSize);
LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)(
add_constant_kernel<kConstant>, dst_ptr, src_ptr, num_elements);
}
}
} // namespace
python/sglang/jit_kernel/csrc/attention/fixup_zero_kv.cuh
+124
View File
@@ -0,0 +1,124 @@
#pragma once
// Fixup kernel for TRT-LLM ragged attention zero-KV rows.
// For sequences with kv_len == 0, forces out=0 and lse=-inf.
// 2D grid: (blocks_per_seq, batch_size). Y-dim early-exits for non-zero KV.
// Uses vectorised float4 stores for bandwidth efficiency.
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.cuh>
#include <cstdint>
namespace {
constexpr int kFixupBlockSize = 256;
// -- vectorised zero-fill helpers ------------------------------------------
// Zero-fill `n` elements of type T starting at `ptr`, using float4 stores.
// `ptr` must be 16-byte aligned (guaranteed by PyTorch allocator).
template <typename T>
__device__ __forceinline__ void vec_zero_fill(T* ptr, int n) {
constexpr int kVec = 16 / sizeof(T); // elements per float4
const int n_vec = n / kVec; // full vectors
float4* dst4 = reinterpret_cast<float4*>(ptr);
const float4 z4 = make_float4(0.f, 0.f, 0.f, 0.f);
for (int i = threadIdx.x; i < n_vec; i += blockDim.x) {
dst4[i] = z4;
}
// tail elements
const int tail_start = n_vec * kVec;
for (int i = tail_start + threadIdx.x; i < n; i += blockDim.x) {
ptr[i] = static_cast<T>(0);
}
}
// Fill `n` float elements with -inf using float4 stores.
__device__ __forceinline__ void vec_neginf_fill(float* ptr, int n) {
constexpr int kVec = 4; // float4 = 4 floats
const int n_vec = n / kVec;
float4* dst4 = reinterpret_cast<float4*>(ptr);
const float ninf = -INFINITY;
const float4 inf4 = make_float4(ninf, ninf, ninf, ninf);
for (int i = threadIdx.x; i < n_vec; i += blockDim.x) {
dst4[i] = inf4;
}
const int tail_start = n_vec * kVec;
for (int i = tail_start + threadIdx.x; i < n; i += blockDim.x) {
ptr[i] = ninf;
}
}
// -- main kernel -----------------------------------------------------------
template <typename OutT>
__global__ void fixup_zero_kv_rows_kernel(
OutT* __restrict__ out,
float* __restrict__ lse,
const int32_t* __restrict__ kv_lens,
const int32_t* __restrict__ cum_seq_lens,
const int out_stride,
const int lse_stride) {
const int seq_idx = blockIdx.y;
if (kv_lens[seq_idx] > 0) return;
const int tok_start = cum_seq_lens[seq_idx];
const int tok_end = cum_seq_lens[seq_idx + 1];
const int num_tokens = tok_end - tok_start;
if (num_tokens <= 0) return;
// blockIdx.x selects a token within this sequence.
const int tok = tok_start + blockIdx.x;
if (tok >= tok_end) return;
// Each block handles one token: zero out[tok] and set lse[tok] = -inf.
vec_zero_fill(out + tok * out_stride, out_stride);
vec_neginf_fill(lse + tok * lse_stride, lse_stride);
}
// -- host launcher ---------------------------------------------------------
template <typename OutT>
void fixup_zero_kv_rows(
tvm::ffi::TensorView out,
tvm::ffi::TensorView lse,
tvm::ffi::TensorView kv_lens,
tvm::ffi::TensorView cum_seq_lens,
int64_t max_seq_len) {
using namespace host;
auto batch_size = SymbolicSize{"batch_size"};
auto total_tokens = SymbolicSize{"total_tokens"};
auto num_heads = SymbolicSize{"num_heads"};
auto v_head_dim = SymbolicSize{"v_head_dim"};
auto batch_size_plus_1 = SymbolicSize{"batch_size_plus_1"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({total_tokens, num_heads, v_head_dim}).with_dtype<OutT>().with_device(device).verify(out);
TensorMatcher({total_tokens, num_heads}).with_dtype<float>().with_device(device).verify(lse);
TensorMatcher({batch_size}).with_dtype<int32_t>().with_device(device).verify(kv_lens);
TensorMatcher({batch_size_plus_1}).with_dtype<int32_t>().with_device(device).verify(cum_seq_lens);
const int bs = static_cast<int>(batch_size.unwrap());
const int nh = static_cast<int>(num_heads.unwrap());
const int vd = static_cast<int>(v_head_dim.unwrap());
// Grid: one block per (token, sequence). X = max tokens in any seq.
const int blocks_x = static_cast<int>(max_seq_len);
dim3 grid(blocks_x, bs);
dim3 block(kFixupBlockSize);
LaunchKernel(grid, block, device.unwrap())(
fixup_zero_kv_rows_kernel<OutT>,
static_cast<OutT*>(out.data_ptr()),
static_cast<float*>(lse.data_ptr()),
static_cast<const int32_t*>(kv_lens.data_ptr()),
static_cast<const int32_t*>(cum_seq_lens.data_ptr()),
nh * vd,
nh);
}
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh
+429
View File
@@ -0,0 +1,429 @@
// DeepSeek-V3.2 only.
//
// DSA indexer K kernels: single-head LayerNorm (not RMS), ropes the leading
// kRopeDim dims, and fp8-quantizes the rotated activations. V3.2 drops the
// Hadamard incoherence rotation; it is logit-preserving (see main_norm_rope.cuh).
//
// Independent of the wk + weights_proj GEMM fusion (dsa_indexer.py): `k_input`
// here is the non-contiguous wk slice kw[:, :head_dim] read via
// k_input_stride_batch (no copy).
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <tvm/ffi/container/tensor.h>
#include <bit>
#include <cstdint>
namespace {
using deepseek_v4::fp8::pack_fp8;
constexpr uint32_t kFusedKIndexerBlockSize = 128;
constexpr uint32_t kFusedKIndexerNumWarps = kFusedKIndexerBlockSize / device::kWarpThreads;
#define K_INDEXER_KERNEL __global__ __launch_bounds__(kFusedKIndexerBlockSize, 16)
template <int64_t kRopeDim>
SGL_DEVICE device::AlignedVector<float, 4>
load_rope_first_cos_sin(const float* __restrict__ cos_sin_cache, int32_t lane_id) {
constexpr int64_t kHalfRopeDim = kRopeDim / 2;
const int32_t pair0 = lane_id * 2;
const int32_t pair1 = pair0 + 1;
device::AlignedVector<float, 4> freq;
freq[0] = cos_sin_cache[pair0];
freq[1] = cos_sin_cache[kHalfRopeDim + pair0];
freq[2] = cos_sin_cache[pair1];
freq[3] = cos_sin_cache[kHalfRopeDim + pair1];
return freq;
}
// Indexer K: LayerNorm + RoPE -> bf16.
struct FusedKIndexerNormRopeParams {
const void* __restrict__ k_input; // (B, 128) DType
void* __restrict__ k_out; // (B, 128) DType
const float* __restrict__ weight; // (128,) fp32 -- LayerNorm gamma
const float* __restrict__ bias; // (128,) fp32 -- LayerNorm beta
const float* __restrict__ cos_sin_cache; // (max_pos, 64) fp32 [cos..., sin...]
const void* __restrict__ positions; // (B,) PosT
// Row stride for `k_input` in elements (caller passes the wk slice directly).
int64_t k_input_stride_batch;
uint32_t batch_size;
float eps;
};
template <typename DType, typename PosT, bool kUsePDL>
K_INDEXER_KERNEL void fused_k_indexer_norm_rope(const __grid_constant__ FusedKIndexerNormRopeParams params) {
using namespace device;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
constexpr int64_t kVecSize = 4;
constexpr uint32_t kRopeSize = kRopeDim / kVecSize; // = 16
static_assert(kHeadDim == kWarpThreads * kVecSize);
static_assert(kRopeDim == kWarpThreads * 2);
static_assert(kRopeSize <= kWarpThreads);
using Storage = AlignedVector<DType, kVecSize>;
using Float4 = AlignedVector<float, kVecSize>;
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
const auto work_id = blockIdx.x * kFusedKIndexerNumWarps + warp_id;
const bool is_rope_lane = lane_id < kRopeSize;
if (work_id >= params.batch_size) return;
const auto input_ptr = static_cast<const DType*>(params.k_input) + work_id * params.k_input_stride_batch;
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[work_id]);
const auto cos_sin_cache = params.cos_sin_cache + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float4 data, freq, gamma, beta;
// part 1: LayerNorm
{
Storage input_vec;
input_vec.load(input_ptr, lane_id);
gamma.load(params.weight, lane_id);
beta.load(params.bias, lane_id);
if (is_rope_lane) freq = load_rope_first_cos_sin<kRopeDim>(cos_sin_cache, lane_id);
float sum = 0.0f;
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = cast<float>(input_vec[i]);
sum += data[i];
}
const float mean = warp::reduce_sum(sum) / kHeadDim;
float var = 0.0f;
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const float centered = data[i] - mean;
var += centered * centered;
}
const float inv_std = math::rsqrt(warp::reduce_sum(var) / kHeadDim + params.eps);
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = (data[i] - mean) * inv_std * gamma[i] + beta[i];
}
}
// part 2: rope on rope lanes
if (is_rope_lane) {
const auto x_real = data[0];
const auto x_imag = data[1];
const auto y_real = data[2];
const auto y_imag = data[3];
const auto fxr = freq[0];
const auto fxi = freq[1];
const auto fyr = freq[2];
const auto fyi = freq[3];
data[0] = x_real * fxr - x_imag * fxi;
data[1] = x_real * fxi + x_imag * fxr;
data[2] = y_real * fyr - y_imag * fyi;
data[3] = y_real * fyi + y_imag * fyr;
}
PDLTriggerSecondary<kUsePDL>();
{
Storage out_vec;
#pragma unroll
for (int i = 0; i < kVecSize; ++i)
out_vec[i] = cast<DType>(data[i]);
auto out_row = static_cast<DType*>(params.k_out) + work_id * kHeadDim;
out_vec.store(out_row, lane_id);
}
}
template <typename DType, bool kUsePDL>
struct FusedKIndexerNormRopeKernel {
template <typename PosT>
static constexpr auto kernel = fused_k_indexer_norm_rope<DType, PosT, kUsePDL>;
static void forward(
const tvm::ffi::TensorView k_input,
const tvm::ffi::TensorView k_out,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView bias,
const tvm::ffi::TensorView cos_sin_cache,
const tvm::ffi::TensorView positions,
double eps) {
using namespace host;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
auto B = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B, kHeadDim}) //
.with_strides({-1, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(k_input);
TensorMatcher({B, kHeadDim}) //
.with_strides({kHeadDim, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(k_out);
TensorMatcher({kHeadDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(weight);
TensorMatcher({kHeadDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(bias);
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(cos_sin_cache);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
.with_device(device_)
.verify(positions);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
if (batch_size == 0) return;
const auto params = FusedKIndexerNormRopeParams{
.k_input = k_input.data_ptr(),
.k_out = k_out.data_ptr(),
.weight = static_cast<const float*>(weight.data_ptr()),
.bias = static_cast<const float*>(bias.data_ptr()),
.cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
.positions = positions.data_ptr(),
.k_input_stride_batch = k_input.stride(0),
.batch_size = batch_size,
.eps = static_cast<float>(eps),
};
const auto num_blocks = div_ceil(batch_size, kFusedKIndexerNumWarps);
const auto k_int32 = kernel<int32_t>;
const auto k_int64 = kernel<int64_t>;
const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
LaunchKernel(num_blocks, kFusedKIndexerBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(k, params);
}
};
// Indexer K + fused store: LayerNorm + RoPE + fp8 quant + paged store in one
// launch. Page layout matches fused_store_index_cache.cuh: each page is
// 132*page_size bytes (128*page_size fp8 keys, then 4*page_size fp32 scales).
struct FusedKIndexerNormRopeStoreParams {
const void* __restrict__ k_input; // (B, 128) DType
void* __restrict__ cache; // (num_pages, 132*page_size) uint8
const void* __restrict__ indices; // (B,) int64 -- out_cache_loc
const float* __restrict__ weight; // (128,) fp32 -- LayerNorm gamma
const float* __restrict__ bias; // (128,) fp32 -- LayerNorm beta
const float* __restrict__ cos_sin_cache; // (max_pos, 64) fp32 [cos..., sin...]
const void* __restrict__ positions; // (B,) PosT
// Row stride for `k_input` (caller passes the non-contiguous wk slice directly).
int64_t k_input_stride_batch;
uint32_t batch_size;
float eps;
};
template <typename DType, typename PosT, bool kUsePDL, int32_t kPageBits>
K_INDEXER_KERNEL void fused_k_indexer_norm_rope_store(const __grid_constant__ FusedKIndexerNormRopeStoreParams params) {
using namespace device;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
constexpr int64_t kVecSize = 4;
constexpr uint32_t kRopeSize = kRopeDim / kVecSize; // = 16
constexpr int64_t kPageBytes = 132ll << kPageBits;
static_assert(kHeadDim == kWarpThreads * kVecSize);
static_assert(kRopeDim == kWarpThreads * 2);
static_assert(kRopeSize <= kWarpThreads);
using Storage = AlignedVector<DType, kVecSize>;
using Float4 = AlignedVector<float, kVecSize>;
using OutStorage = AlignedVector<fp8x2_e4m3_t, 2>; // 4 fp8 / lane
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
const auto work_id = blockIdx.x * kFusedKIndexerNumWarps + warp_id;
const bool is_rope_lane = lane_id < kRopeSize;
if (work_id >= params.batch_size) return;
const auto input_ptr = static_cast<const DType*>(params.k_input) + work_id * params.k_input_stride_batch;
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[work_id]);
const auto cos_sin_cache = params.cos_sin_cache + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float4 data, freq, gamma, beta;
// part 1: LayerNorm
{
Storage input_vec;
input_vec.load(input_ptr, lane_id);
gamma.load(params.weight, lane_id);
beta.load(params.bias, lane_id);
if (is_rope_lane) freq = load_rope_first_cos_sin<kRopeDim>(cos_sin_cache, lane_id);
float sum = 0.0f;
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = cast<float>(input_vec[i]);
sum += data[i];
}
const float mean = warp::reduce_sum(sum) / kHeadDim;
float var = 0.0f;
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const float centered = data[i] - mean;
var += centered * centered;
}
const float inv_std = math::rsqrt(warp::reduce_sum(var) / kHeadDim + params.eps);
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = (data[i] - mean) * inv_std * gamma[i] + beta[i];
}
}
// part 2: rope on rope lanes
if (is_rope_lane) {
const auto x_real = data[0];
const auto x_imag = data[1];
const auto y_real = data[2];
const auto y_imag = data[3];
const auto fxr = freq[0];
const auto fxi = freq[1];
const auto fyr = freq[2];
const auto fyi = freq[3];
data[0] = x_real * fxr - x_imag * fxi;
data[1] = x_real * fxi + x_imag * fxr;
data[2] = y_real * fyr - y_imag * fyi;
data[3] = y_real * fyi + y_imag * fyr;
}
PDLTriggerSecondary<kUsePDL>();
// part 3: fp8 act-quant + paged store. Round through bf16 first so the fp8
// scale matches the un-fused path.
#pragma unroll
for (int i = 0; i < kVecSize; ++i)
data[i] = cast<float>(cast<DType>(data[i]));
float local_max = math::abs(data[0]);
#pragma unroll
for (int i = 1; i < kVecSize; ++i)
local_max = math::max(local_max, math::abs(data[i]));
const auto abs_max = warp::reduce_max(local_max);
const auto scale = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
const auto inv_scale = 1.0f / scale;
const auto index = static_cast<const int64_t*>(params.indices)[work_id];
const int32_t page = static_cast<int32_t>(index >> kPageBits);
const int32_t offset = static_cast<int32_t>(index & ((1 << kPageBits) - 1));
const auto page_ptr = static_cast<uint8_t*>(params.cache) + page * kPageBytes;
const auto value_ptr = page_ptr + offset * kHeadDim;
const auto scale_ptr = page_ptr + (kHeadDim << kPageBits) + offset * 4;
OutStorage result;
result[0] = pack_fp8(data[0] * inv_scale, data[1] * inv_scale);
result[1] = pack_fp8(data[2] * inv_scale, data[3] * inv_scale);
reinterpret_cast<OutStorage*>(value_ptr)[lane_id] = result;
if (lane_id == 0) *reinterpret_cast<float*>(scale_ptr) = scale;
}
template <typename DType, bool kUsePDL, uint32_t kPageSize>
struct FusedKIndexerNormRopeStoreKernel {
static constexpr int32_t kPageBits = std::countr_zero(kPageSize);
static constexpr int64_t kPageBytes = 132ll * kPageSize;
static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
template <typename PosT>
static constexpr auto kernel = fused_k_indexer_norm_rope_store<DType, PosT, kUsePDL, kPageBits>;
static void forward(
const tvm::ffi::TensorView k_input,
const tvm::ffi::TensorView cache,
const tvm::ffi::TensorView indices,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView bias,
const tvm::ffi::TensorView cos_sin_cache,
const tvm::ffi::TensorView positions,
double eps) {
using namespace host;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
auto B = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B, kHeadDim}) //
.with_strides({-1, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(k_input);
TensorMatcher({-1, -1}) //
.with_strides({kPageBytes, 1})
.with_dtype<uint8_t>()
.with_device(device_)
.verify(cache);
TensorMatcher({B}) //
.with_dtype<int64_t>()
.with_device(device_)
.verify(indices);
TensorMatcher({kHeadDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(weight);
TensorMatcher({kHeadDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(bias);
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(cos_sin_cache);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
.with_device(device_)
.verify(positions);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
if (batch_size == 0) return;
const auto params = FusedKIndexerNormRopeStoreParams{
.k_input = k_input.data_ptr(),
.cache = cache.data_ptr(),
.indices = indices.data_ptr(),
.weight = static_cast<const float*>(weight.data_ptr()),
.bias = static_cast<const float*>(bias.data_ptr()),
.cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
.positions = positions.data_ptr(),
.k_input_stride_batch = k_input.stride(0),
.batch_size = batch_size,
.eps = static_cast<float>(eps),
};
const auto num_blocks = div_ceil(batch_size, kFusedKIndexerNumWarps);
const auto k_int32 = kernel<int32_t>;
const auto k_int64 = kernel<int64_t>;
const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
LaunchKernel(num_blocks, kFusedKIndexerBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(k, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/c128.cuh
+522
View File
@@ -0,0 +1,522 @@
#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 <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/compress.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/object.h>
#include <cstdint>
namespace {
using Plan128 = device::compress::PrefillPlan;
using IndiceT = int32_t;
/// \brief Each thread will handle this many elements (split along head_dim)
constexpr int32_t kTileElements = 2;
/// \brief Each warp will handle this many elements (split along 128)
constexpr int32_t kElementsPerWarp = 8;
constexpr uint32_t kNumWarps = 128 / kElementsPerWarp;
constexpr uint32_t kBlockSize = device::kWarpThreads * kNumWarps;
/// \brief Need to reduce register usage to increase occupancy
#define C128_KERNEL __global__ __launch_bounds__(kBlockSize, 2)
struct Compress128DecodeParams {
/**
* \brief Shape: `[num_indices, 128, head_dim * 2]` \n
* last dimension layout:
* | kv current | score current |
*/
void* __restrict__ kv_score_buffer;
/** \brief Shape: `[batch_size, head_dim * 2]` */
const void* __restrict__ kv_score_input;
/** \brief Shape: `[batch_size, head_dim]` */
void* __restrict__ kv_compressed_output;
/** \brief Shape: `[128, head_dim]` (called `ape`) */
const void* __restrict__ score_bias;
/** \brief Shape: `[batch_size, ]`*/
const IndiceT* __restrict__ indices;
/** \brief Shape: `[batch_size, ]` */
const IndiceT* __restrict__ seq_lens;
/** \NOTE: `batch_size` <= `num_indices` */
uint32_t batch_size;
};
struct Compress128PrefillParams {
/**
* \brief Shape: `[num_indices, 128, head_dim * 2]` \n
* last dimension layout:
* | kv current | score current |
*/
void* __restrict__ kv_score_buffer;
/** \brief Shape: `[batch_size, head_dim * 2]` */
const void* __restrict__ kv_score_input;
/** \brief Shape: `[batch_size, head_dim]` */
void* __restrict__ kv_compressed_output;
/** \brief Shape: `[128, head_dim]` (called `ape`) */
const void* __restrict__ score_bias;
/** \brief Shape: `[batch_size, ]`*/
const IndiceT* __restrict__ indices;
/** \brief Shape: `[batch_size, ]`*/
const int32_t* __restrict__ load_indices;
/** \brief The following part is plan info. */
const Plan128* __restrict__ compress_plan;
const Plan128* __restrict__ write_plan;
uint32_t num_compress;
uint32_t num_write;
};
struct Compress128SharedBuffer {
using Storage = device::AlignedVector<float, kTileElements>;
Storage data[kNumWarps][device::kWarpThreads + 1]; // padding to avoid bank conflict
SGL_DEVICE Storage& operator()(uint32_t warp_id, uint32_t lane_id) {
return data[warp_id][lane_id];
}
SGL_DEVICE float& operator()(uint32_t warp_id, uint32_t lane_id, uint32_t tile_id) {
return data[warp_id][lane_id][tile_id];
}
};
template <typename T>
SGL_DEVICE void c128_write(
T* kv_score_buf, //
const T* kv_score_src,
const int64_t head_dim,
const int32_t write_pos,
const uint32_t lane_id) {
using namespace device;
using Storage = AlignedVector<T, kTileElements>;
const auto element_size = head_dim * 2;
const auto gmem = tile::Memory<Storage>{lane_id, kWarpThreads};
kv_score_buf += write_pos * element_size;
/// NOTE: Layout | [0] = kv | [1] = score |
Storage kv_score[2];
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
kv_score[i] = gmem.load(kv_score_src + head_dim * i);
}
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
gmem.store(kv_score_buf + head_dim * i, kv_score[i]);
}
}
template <typename InFloat, typename OutFloat>
SGL_DEVICE void c128_forward(
const InFloat* kv_score_buf,
const InFloat* kv_score_src,
OutFloat* kv_out,
const InFloat* score_bias,
const int64_t head_dim,
const int32_t window_len,
const uint32_t warp_id,
const uint32_t lane_id) {
using namespace device;
const auto element_size = head_dim * 2;
const auto score_offset = head_dim;
/// NOTE: part 1: load kv + score
using StorageIn = AlignedVector<InFloat, kTileElements>;
const auto gmem_in = tile::Memory<StorageIn>{lane_id, kWarpThreads};
StorageIn kv[kElementsPerWarp];
StorageIn score[kElementsPerWarp];
StorageIn bias[kElementsPerWarp];
const int32_t warp_offset = warp_id * kElementsPerWarp;
#pragma unroll
for (int32_t i = 0; i < 8; ++i) {
const int32_t j = i + warp_offset;
bias[i] = gmem_in.load(score_bias + j * head_dim);
}
#pragma unroll
for (int32_t i = 0; i < kElementsPerWarp; ++i) {
const int32_t j = i + warp_offset;
const InFloat* src;
__builtin_assume(j < 128);
if (j < window_len) {
src = kv_score_buf + j * element_size;
} else {
/// NOTE: k in [-127, 0]. We'll load from the ragged `kv_score_src`
const int32_t k = j - 127;
src = kv_score_src + k * element_size;
}
kv[i] = gmem_in.load(src);
score[i] = gmem_in.load(src + score_offset);
}
/// NOTE: part 2: safe online softmax + weighted sum
using TmpStorage = typename Compress128SharedBuffer::Storage;
__shared__ Compress128SharedBuffer s_local_val_max;
__shared__ Compress128SharedBuffer s_local_exp_sum;
__shared__ Compress128SharedBuffer s_local_product;
TmpStorage tmp_val_max;
TmpStorage tmp_exp_sum;
TmpStorage tmp_product;
#pragma unroll
for (int32_t i = 0; i < kTileElements; ++i) {
float score_fp32[kElementsPerWarp];
#pragma unroll
for (int32_t j = 0; j < kElementsPerWarp; ++j) {
score_fp32[j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
}
float max_value = score_fp32[0];
float sum_exp_value = 0.0f;
#pragma unroll
for (int32_t j = 1; j < kElementsPerWarp; ++j) {
const auto fp32_score = score_fp32[j];
max_value = fmaxf(max_value, fp32_score);
}
float sum_product = 0.0f;
#pragma unroll
for (int32_t j = 0; j < 8; ++j) {
const auto fp32_score = score_fp32[j];
const auto exp_score = expf(fp32_score - max_value);
sum_product += cast<float>(kv[j][i]) * exp_score;
sum_exp_value += exp_score;
}
tmp_val_max[i] = max_value;
tmp_exp_sum[i] = sum_exp_value;
tmp_product[i] = sum_product;
}
// naturally aligned, so no bank conflict
s_local_val_max(warp_id, lane_id) = tmp_val_max;
s_local_exp_sum(warp_id, lane_id) = tmp_exp_sum;
s_local_product(warp_id, lane_id) = tmp_product;
__syncthreads();
/// NOTE: part 3: online softmax
/// NOTE: We have `kTileElements * kWarpThreads * kNumWarps` values to reduce
/// each reduce will consume `kNumWarps` threads (use partial warp reduction)
constexpr uint32_t kReductionCount = kTileElements * kWarpThreads * kNumWarps;
constexpr uint32_t kIteration = kReductionCount / kBlockSize;
#pragma unroll
for (uint32_t i = 0; i < kIteration; ++i) {
/// NOTE: Range `[0, kTileElements * kWarpThreads * kNumWarps)`
const uint32_t j = i * kBlockSize + warp_id * kWarpThreads + lane_id;
/// NOTE: Range `[0, kNumWarps)`
const uint32_t local_warp_id = j % kNumWarps;
/// NOTE: Range `[0, kTileElements * kWarpThreads)`
const uint32_t local_elem_id = j / kNumWarps;
/// NOTE: Range `[0, kTileElements)`
const uint32_t local_tile_id = local_elem_id % kTileElements;
/// NOTE: Range `[0, kWarpThreads)`
const uint32_t local_lane_id = local_elem_id / kTileElements;
/// NOTE: each warp will access the whole tile (all `kTileElements`)
/// and for different lanes, the memory access only differ in `local_warp_id`
/// so there's no bank conflict in shared memory access.
static_assert(kTileElements * kNumWarps == kWarpThreads, "TODO: support other configs");
const auto local_val_max = s_local_val_max(local_warp_id, local_lane_id, local_tile_id);
const auto local_exp_sum = s_local_exp_sum(local_warp_id, local_lane_id, local_tile_id);
const auto local_product = s_local_product(local_warp_id, local_lane_id, local_tile_id);
const auto global_val_max = warp::reduce_max<kNumWarps>(local_val_max);
const auto rescale = expf(local_val_max - global_val_max);
const auto global_exp_sum = warp::reduce_sum<kNumWarps>(local_exp_sum * rescale);
const auto final_scale = rescale / global_exp_sum;
const auto global_product = warp::reduce_sum<kNumWarps>(local_product * final_scale);
kv_out[local_elem_id] = cast<OutFloat>(global_product);
}
}
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodeParams params) {
using namespace device;
constexpr int64_t kTileDim = kTileElements * kWarpThreads; // 64
constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
constexpr int64_t kElementSize = kHeadDim * 2;
static_assert(kHeadDim % kTileDim == 0, "Head dim must be multiple of tile dim");
const auto& [
_kv_score_buffer, _kv_score_input, _kv_compressed_output, _score_bias, // kv score
indices, seq_lens, batch_size // decode info
] = params;
const uint32_t warp_id = threadIdx.x / kWarpThreads;
const uint32_t lane_id = threadIdx.x % kWarpThreads;
const uint32_t global_bid = blockIdx.x / kNumSplit; // batch id
const uint32_t global_sid = blockIdx.x % kNumSplit; // split id
if (global_bid >= batch_size) return;
const int32_t index = indices[global_bid];
const int32_t seq_len = seq_lens[global_bid];
const int64_t split_offset = global_sid * kTileDim;
// kv score
const auto kv_score_buffer = static_cast<InFloat*>(_kv_score_buffer);
const auto kv_buf = kv_score_buffer + index * (kElementSize * 128) + split_offset;
// kv input
const auto kv_score_input = static_cast<const InFloat*>(_kv_score_input);
const auto kv_src = kv_score_input + global_bid * kElementSize + split_offset;
// kv output
const auto kv_compressed_output = static_cast<OutFloat*>(_kv_compressed_output);
const auto kv_out = kv_compressed_output + global_bid * kHeadDim + split_offset;
// score bias (ape)
const auto score_bias = static_cast<const InFloat*>(_score_bias) + split_offset;
PDLWaitPrimary<kUsePDL>();
/// NOTE: the write must be visible to the subsequent c128_forward,
/// so only the last warp can write to HBM
/// In addition, `position` = `seq_len - 1`. To avoid underflow, we use `seq_len + 127`
if (warp_id == kNumWarps - 1) {
c128_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/(seq_len + 127) % 128, lane_id);
}
if (seq_len % 128 == 0) {
c128_forward(kv_buf, kv_src, kv_out, score_bias, kHeadDim, /*window_len=*/128, warp_id, lane_id);
}
PDLTriggerSecondary<kUsePDL>();
}
// compress kernel
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kWrite, bool kUsePDL>
C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
using namespace device;
constexpr int64_t kTileDim = kTileElements * kWarpThreads; // 64
constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
constexpr int64_t kElementSize = kHeadDim * 2;
static_assert(kHeadDim % kTileDim == 0, "Head dim must be multiple of tile dim");
const auto& [
_kv_score_buffer, _kv_score_input, _kv_compressed_output, _score_bias, // kv score
indices, load_indices, compress_plan, write_plan, num_compress, num_write // prefill plan
] = params;
const uint32_t warp_id = threadIdx.x / kWarpThreads;
const uint32_t lane_id = threadIdx.x % kWarpThreads;
uint32_t global_id;
if constexpr (kWrite) {
// for write kernel, we use global warp_id to dispatch work
global_id = (blockIdx.x * blockDim.x + threadIdx.x) / kWarpThreads;
} else {
// for compress kernel, we use block id to dispatch work
global_id = blockIdx.x; // block id
}
const uint32_t global_pid = global_id / kNumSplit; // plan id
const uint32_t global_sid = global_id % kNumSplit; // split id
/// NOTE: compiler can optimize this if-else at compile time
const auto num_plans = kWrite ? num_write : num_compress;
const auto plan_ptr = kWrite ? write_plan : compress_plan;
if (global_pid >= num_plans) return;
const auto& [ragged_id, global_bid, position, window_len] = plan_ptr[global_pid];
const auto indices_ptr = kWrite ? indices : load_indices;
const int64_t split_offset = global_sid * kTileDim;
// kv input
const auto kv_score_input = static_cast<const InFloat*>(_kv_score_input);
const auto kv_src = kv_score_input + ragged_id * kElementSize + split_offset;
// kv output
const auto kv_compressed_output = static_cast<OutFloat*>(_kv_compressed_output);
const auto kv_out = kv_compressed_output + ragged_id * kHeadDim + split_offset;
// score bias (ape)
const auto score_bias = static_cast<const InFloat*>(_score_bias) + split_offset;
if (ragged_id == 0xFFFFFFFF) [[unlikely]]
return;
const int32_t index = indices_ptr[global_bid];
// kv score
const auto kv_score_buffer = static_cast<InFloat*>(_kv_score_buffer);
const auto kv_buf = kv_score_buffer + index * (kElementSize * 128) + split_offset;
PDLWaitPrimary<kUsePDL>();
// only responsible for the compress part
if constexpr (kWrite) {
c128_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/position % 128, lane_id);
} else {
c128_forward(kv_buf, kv_src, kv_out, score_bias, kHeadDim, window_len, warp_id, lane_id);
}
PDLTriggerSecondary<kUsePDL>();
}
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
struct FlashCompress128Kernel {
static constexpr auto decode_kernel = flash_c128_decode<kHeadDim, InFloat, OutFloat, kUsePDL>;
template <bool kWrite>
static constexpr auto prefill_kernel = flash_c128_prefill<kHeadDim, InFloat, OutFloat, kWrite, kUsePDL>;
static constexpr auto prefill_c_kernel = prefill_kernel</*kWrite=*/false>;
static constexpr auto prefill_w_kernel = prefill_kernel</*kWrite=*/true>;
static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
static constexpr uint32_t kWriteBlockSize = 128;
static constexpr uint32_t kWarpsPerWriteBlock = kWriteBlockSize / device::kWarpThreads;
static void run_decode(
const tvm::ffi::TensorView kv_score_buffer,
const tvm::ffi::TensorView kv_score_input,
const tvm::ffi::TensorView kv_compressed_output,
const tvm::ffi::TensorView ape,
const tvm::ffi::TensorView indices,
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::Optional<tvm::ffi::TensorView> /* UNUSED */) {
using namespace host;
// this should not happen in practice
auto B = SymbolicSize{"batch_size"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({-1, 128, kHeadDim * 2}) // kv score
.with_dtype<InFloat>()
.with_device(device)
.verify(kv_score_buffer);
TensorMatcher({B, kHeadDim * 2}) // kv score input
.with_dtype<InFloat>()
.with_device(device)
.verify(kv_score_input);
TensorMatcher({B, kHeadDim}) // kv compressed output
.with_dtype<OutFloat>()
.with_device(device)
.verify(kv_compressed_output);
TensorMatcher({128, kHeadDim}) // ape
.with_dtype<InFloat>()
.with_device(device)
.verify(ape);
TensorMatcher({B}) // indices
.with_dtype<IndiceT>()
.with_device(device)
.verify(indices);
TensorMatcher({B}) // seq lens
.with_dtype<IndiceT>()
.with_device(device)
.verify(seq_lens);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto params = Compress128DecodeParams{
.kv_score_buffer = kv_score_buffer.data_ptr(),
.kv_score_input = kv_score_input.data_ptr(),
.kv_compressed_output = kv_compressed_output.data_ptr(),
.score_bias = ape.data_ptr(),
.indices = static_cast<const IndiceT*>(indices.data_ptr()),
.seq_lens = static_cast<const IndiceT*>(seq_lens.data_ptr()),
.batch_size = batch_size,
};
const uint32_t num_blocks = batch_size * kNumSplit;
LaunchKernel(num_blocks, kBlockSize, device.unwrap()) //
.enable_pdl(kUsePDL)(decode_kernel, params);
}
static void run_prefill(
const tvm::ffi::TensorView kv_score_buffer,
const tvm::ffi::TensorView kv_score_input,
const tvm::ffi::TensorView kv_compressed_output,
const tvm::ffi::TensorView ape,
const tvm::ffi::TensorView indices,
const tvm::ffi::TensorView compress_plan,
const tvm::ffi::TensorView write_plan,
const tvm::ffi::Optional<tvm::ffi::TensorView> extra) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto N = SymbolicSize{"num_q_tokens"};
auto X = SymbolicSize{"compress_tokens"};
auto Y = SymbolicSize{"write_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({-1, 128, kHeadDim * 2}) // kv score
.with_dtype<InFloat>()
.with_device(device_)
.verify(kv_score_buffer);
TensorMatcher({N, kHeadDim * 2}) // kv score input
.with_dtype<InFloat>()
.with_device(device_)
.verify(kv_score_input);
TensorMatcher({N, kHeadDim}) // kv compressed output
.with_dtype<OutFloat>()
.with_device(device_)
.verify(kv_compressed_output);
TensorMatcher({128, kHeadDim}) // ape
.with_dtype<InFloat>()
.with_device(device_)
.verify(ape);
TensorMatcher({B}) // indices
.with_dtype<IndiceT>()
.with_device(device_)
.verify(indices);
TensorMatcher({X, compress::kPrefillPlanDim}) // compress plan
.with_dtype<compress::PrefillPlanTensorDtype>()
.with_device(device_)
.verify(compress_plan);
TensorMatcher({Y, compress::kPrefillPlanDim}) // write plan
.with_dtype<compress::PrefillPlanTensorDtype>()
.with_device(device_)
.verify(write_plan);
// might be needed for prefill write
const auto load_indices = extra.value_or(indices);
TensorMatcher({B}) // [read_positions]
.with_dtype<IndiceT>()
.with_device(device_)
.verify(load_indices);
const auto device = device_.unwrap();
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
const auto num_c = static_cast<uint32_t>(X.unwrap());
const auto num_w = static_cast<uint32_t>(Y.unwrap());
const auto params = Compress128PrefillParams{
.kv_score_buffer = kv_score_buffer.data_ptr(),
.kv_score_input = kv_score_input.data_ptr(),
.kv_compressed_output = kv_compressed_output.data_ptr(),
.score_bias = ape.data_ptr(),
.indices = static_cast<const IndiceT*>(indices.data_ptr()),
.load_indices = static_cast<const IndiceT*>(load_indices.data_ptr()),
.compress_plan = static_cast<const Plan128*>(compress_plan.data_ptr()),
.write_plan = static_cast<const Plan128*>(write_plan.data_ptr()),
.num_compress = num_c,
.num_write = num_w,
};
RuntimeCheck(num_q_tokens >= batch_size, "num_q_tokens must be >= batch_size");
RuntimeCheck(num_q_tokens >= std::max(num_c, num_w), "invalid prefill plan");
constexpr auto kBlockSize_C = kBlockSize;
constexpr auto kBlockSize_W = kWriteBlockSize;
if (const auto num_c_blocks = num_c * kNumSplit) {
LaunchKernel(num_c_blocks, kBlockSize_C, device) //
.enable_pdl(kUsePDL)(prefill_c_kernel, params);
}
if (const auto num_w_blocks = div_ceil(num_w * kNumSplit, kWarpsPerWriteBlock)) {
LaunchKernel(num_w_blocks, kBlockSize_W, device) //
.enable_pdl(kUsePDL)(prefill_w_kernel, params);
}
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/c128_online.cuh
+726
View File
@@ -0,0 +1,726 @@
#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 <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/compress.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/container/tuple.h>
#include <tvm/ffi/object.h>
#include <algorithm>
#include <cfloat>
#include <cstdint>
namespace device::compress {
/// \brief Plan entry for online compress 128 prefill.
/// Each entry describes a contiguous segment of tokens that lies inside a
/// single 128-chunk. Multiple segments can map to the same batch id when the
/// extend tokens span chunk boundaries.
///
/// **Layout compatibility:** the field order/types match `PrefillPlan` so that
/// downstream kernels (e.g. `fused_norm_rope` in `CompressExtend` mode) can
/// consume the compress_plan tensor as-if it were a `PrefillPlan` tensor --
/// they only read `ragged_id` and `position`, both of which carry identical
/// semantics here (the LAST token of the segment in q-ragged and global
/// coordinates respectively).
///
/// Note that `window_len` here means "number of real tokens in this segment"
/// (1..128), which differs from `PrefillPlan::window_len`. Downstream kernels
/// that share the tensor MUST NOT read it under that name.
struct alignas(16) OnlinePrefillPlan {
/// \brief Ragged-q position of the LAST token in this segment.
/// Equal to `segment_start_ragged + window_len - 1`.
uint32_t ragged_id;
/// \brief Index into the `indices` / `load_indices` arrays.
uint32_t batch_id;
/// \brief Global position of the LAST token in this segment.
/// For compress plans, `position % 128 == 127` (chunk-closing); for write
/// plans, `position % 128 < 127`.
uint32_t position;
/// \brief Number of real tokens in this segment (1..128).
/// The first segment token sits at `position - window_len + 1` (global) and
/// at `ragged_id - window_len + 1` (ragged).
uint32_t window_len;
};
static_assert(alignof(OnlinePrefillPlan) == alignof(PrefillPlan));
static_assert(sizeof(OnlinePrefillPlan) == sizeof(PrefillPlan));
} // namespace device::compress
namespace host::compress {
using device::compress::OnlinePrefillPlan;
using OnlinePrefillPlanTensorDtype = uint8_t;
inline constexpr int64_t kOnlinePrefillPlanDim = 16;
static_assert(alignof(OnlinePrefillPlan) == sizeof(OnlinePrefillPlan));
static_assert(sizeof(OnlinePrefillPlan) == kOnlinePrefillPlanDim * sizeof(OnlinePrefillPlanTensorDtype));
} // namespace host::compress
namespace {
using OnlinePlan = device::compress::OnlinePrefillPlan;
using IndiceT = int32_t;
/// \brief Need to reduce register usage to increase occupancy
struct Compress128OnlineDecodeParams {
/** \brief Shape: `[num_indices, 1, head_dim * 3 (max, sum, kv) ]` \n */
void* __restrict__ kv_score_buffer;
/** \brief Shape: `[batch_size, head_dim * 2]` */
const void* __restrict__ kv_score_input;
/** \brief Shape: `[batch_size, head_dim]` */
void* __restrict__ kv_compressed_output;
/** \brief Shape: `[128, head_dim]` (called `ape`) */
const void* __restrict__ score_bias;
/** \brief Shape: `[batch_size, ]`*/
const IndiceT* __restrict__ indices;
/** \brief Shape: `[batch_size, ]` */
const IndiceT* __restrict__ seq_lens;
/** \NOTE: `batch_size` <= `num_indices` */
uint32_t batch_size;
};
/// \brief Need to reduce register usage to increase occupancy
struct Compress128OnlinePrefillParams {
/** \brief Shape: `[num_indices, 1, head_dim * 3 (max, sum, kv) ]` \n */
void* __restrict__ kv_score_buffer;
/** \brief Shape: `[num_q_tokens, head_dim * 2]` */
const void* __restrict__ kv_score_input;
/** \brief Shape: `[num_q_tokens, head_dim]` */
void* __restrict__ kv_compressed_output;
/** \brief Shape: `[128, head_dim]` (called `ape`) */
const void* __restrict__ score_bias;
/** \brief Shape: `[batch_size, ]`*/
const IndiceT* __restrict__ indices;
/** \brief Shape: `[batch_size, ]`*/
const IndiceT* __restrict__ load_indices;
/// \brief Plan for segments that close a chunk (write to `kv_compressed_output`).
/// Shape: `[num_compress, 16]` (uint8).
const OnlinePlan* __restrict__ compress_plan;
/// \brief Plan for the trailing partial segment of each batch (write back to
/// `kv_score_buffer`). Shape: `[num_write, 16]` (uint8).
const OnlinePlan* __restrict__ write_plan;
uint32_t num_compress;
uint32_t num_write;
};
// 4 elements per thread, kHeadDim / 4 threads per block
template <int64_t kHeadDim, bool kUsePDL>
__global__ void flash_c128_online_decode(const __grid_constant__ Compress128OnlineDecodeParams params) {
using namespace device;
constexpr uint32_t kVecSize = 4;
constexpr uint32_t kBlockSize = kHeadDim / kVecSize;
using Vec = AlignedVector<float, kVecSize>;
const auto gmem = tile::Memory<Vec>::cta(kBlockSize);
const auto batch_id = blockIdx.x;
const auto index = params.indices[batch_id];
const auto seq_len = params.seq_lens[batch_id];
const auto kv_score_buffer = static_cast<float*>(params.kv_score_buffer);
const auto kv_buf = kv_score_buffer + index * (kHeadDim * 3);
const auto kv_score_input = static_cast<const float*>(params.kv_score_input);
const auto kv_src = kv_score_input + batch_id * (kHeadDim * 2);
/// NOTE: kv_score_buffer layout is [max, sum, kv] (slot 0 / 1 / 2). Reads,
/// writes, and the prefill kernel must all agree on this order.
const auto max_score_vec = gmem.load(kv_buf, 0);
const auto sum_score_vec = gmem.load(kv_buf, 1);
const auto old_kv_vec = gmem.load(kv_buf, 2);
/// NOTE: kv_score_input layout is | kv | score | (head_dim each), matching
/// the offline c128 kernel and the online prefill kernel.
const auto new_kv_vec = gmem.load(kv_src, 0);
const auto new_score_raw_vec = gmem.load(kv_src, 1);
/// NOTE: the new token sits at global position `seq_len - 1`, so its
/// position inside the 128-chunk is `(seq_len - 1) % 128`. The previous
/// `seq_len % 128` was off by one (`bias[127]` vs `bias[0]`, etc.).
const auto pos_in_chunk = (seq_len - 1) % 128;
const auto bias_vec = gmem.load(params.score_bias, pos_in_chunk);
Vec out_kv_vec;
Vec out_max_vec;
Vec out_sum_vec;
if (pos_in_chunk != 0) {
// Mid-chunk: combine prior partial state with the new token via online softmax.
#pragma unroll
for (uint32_t i = 0; i < 4; ++i) {
const auto old_max = max_score_vec[i];
const auto old_kv = old_kv_vec[i];
const auto new_score = new_score_raw_vec[i] + bias_vec[i];
const auto new_kv = new_kv_vec[i];
const auto new_max = fmax(old_max, new_score);
const auto old_sum = sum_score_vec[i] * expf(old_max - new_max);
const auto new_exp = expf(new_score - new_max);
const auto new_sum = old_sum + new_exp;
out_kv_vec[i] = (old_kv * old_sum + new_kv * new_exp) / new_sum;
out_max_vec[i] = new_max;
out_sum_vec[i] = new_sum;
}
} else {
// First token of a new 128-chunk: initialize state with this token alone.
#pragma unroll
for (uint32_t i = 0; i < 4; ++i) {
out_kv_vec[i] = new_kv_vec[i];
out_max_vec[i] = new_score_raw_vec[i] + bias_vec[i];
out_sum_vec[i] = 1.0f; // exp(score - max) with max == score
}
}
if (pos_in_chunk == 127) {
// Chunk just closed: emit the compressed kv. No need to update the buffer
// -- the next chunk's first token will overwrite it.
const auto kv_out = static_cast<float*>(params.kv_compressed_output) + batch_id * kHeadDim;
gmem.store(kv_out, out_kv_vec);
} else {
// Otherwise persist the running [max, sum, kv] state for the next step.
gmem.store(kv_buf, out_max_vec, 0);
gmem.store(kv_buf, out_sum_vec, 1);
gmem.store(kv_buf, out_kv_vec, 2);
}
}
constexpr int32_t kTileElements = 2; // split (along head-dim)
/// \brief Each warp will handle this many elements (split along softmax-128)
constexpr int32_t kElementsPerWarp = 8;
constexpr uint32_t kNumWarps = 128 / kElementsPerWarp;
constexpr uint32_t kPrefillBlockSize = device::kWarpThreads * kNumWarps;
using PrefillStorage = device::AlignedVector<float, kTileElements>;
struct Compress128SharedBuffer {
using Storage = device::AlignedVector<float, 4>;
Storage data[kNumWarps][device::kWarpThreads + 1]; // padding to avoid bank conflict
SGL_DEVICE Storage& operator()(uint32_t warp_id, uint32_t lane_id) {
return data[warp_id][lane_id];
}
SGL_DEVICE float& operator()(uint32_t warp_id, uint32_t lane_id, uint32_t tile_id) {
return data[warp_id][lane_id][tile_id];
}
};
template <bool kNeedData>
SGL_DEVICE void c128_prefill_forward(
const PrefillStorage (&kv)[kElementsPerWarp],
const PrefillStorage (&score)[kElementsPerWarp],
float* kv_out,
float* max_out,
float* sum_out,
const uint32_t warp_id,
const uint32_t lane_id) {
using namespace device;
/// NOTE: part 2: safe online softmax + weighted sum
using TmpStorage = typename Compress128SharedBuffer::Storage;
__shared__ Compress128SharedBuffer s_local_val_max;
__shared__ Compress128SharedBuffer s_local_exp_sum;
__shared__ Compress128SharedBuffer s_local_product;
TmpStorage tmp_val_max;
TmpStorage tmp_exp_sum;
TmpStorage tmp_product;
#pragma unroll
for (int32_t i = 0; i < kTileElements; ++i) {
float score_fp32[kElementsPerWarp];
#pragma unroll
for (int32_t j = 0; j < kElementsPerWarp; ++j) {
score_fp32[j] = score[j][i];
}
float max_value = score_fp32[0];
float sum_exp_value = 0.0f;
#pragma unroll
for (int32_t j = 1; j < kElementsPerWarp; ++j) {
const auto fp32_score = score_fp32[j];
max_value = fmaxf(max_value, fp32_score);
}
float sum_product = 0.0f;
#pragma unroll
for (int32_t j = 0; j < 8; ++j) {
const auto fp32_score = score_fp32[j];
const auto exp_score = expf(fp32_score - max_value);
sum_product += cast<float>(kv[j][i]) * exp_score;
sum_exp_value += exp_score;
}
tmp_val_max[i] = max_value;
tmp_exp_sum[i] = sum_exp_value;
tmp_product[i] = sum_product;
}
// naturally aligned, so no bank conflict
s_local_val_max(warp_id, lane_id) = tmp_val_max;
s_local_exp_sum(warp_id, lane_id) = tmp_exp_sum;
s_local_product(warp_id, lane_id) = tmp_product;
__syncthreads();
/// NOTE: part 3: online softmax
/// NOTE: We have `kTileElements * kWarpThreads * kNumWarps` values to reduce
/// each reduce will consume `kNumWarps` threads (use partial warp reduction)
constexpr uint32_t kReductionCount = kTileElements * kWarpThreads * kNumWarps;
constexpr uint32_t kIteration = kReductionCount / kPrefillBlockSize;
#pragma unroll
for (uint32_t i = 0; i < kIteration; ++i) {
/// NOTE: Range `[0, kTileElements * kWarpThreads * kNumWarps)`
const uint32_t j = i * kPrefillBlockSize + warp_id * kWarpThreads + lane_id;
/// NOTE: Range `[0, kNumWarps)`
const uint32_t local_warp_id = j % kNumWarps;
/// NOTE: Range `[0, kTileElements * kWarpThreads)`
const uint32_t local_elem_id = j / kNumWarps;
/// NOTE: Range `[0, kTileElements)`
const uint32_t local_tile_id = local_elem_id % kTileElements;
/// NOTE: Range `[0, kWarpThreads)`
const uint32_t local_lane_id = local_elem_id / kTileElements;
/// NOTE: each warp will access the whole tile (all `kTileElements`)
/// and for different lanes, the memory access only differ in `local_warp_id`
/// so there's no bank conflict in shared memory access.
static_assert(kTileElements * kNumWarps == kWarpThreads, "TODO: support other configs");
const auto local_val_max = s_local_val_max(local_warp_id, local_lane_id, local_tile_id);
const auto local_exp_sum = s_local_exp_sum(local_warp_id, local_lane_id, local_tile_id);
const auto local_product = s_local_product(local_warp_id, local_lane_id, local_tile_id);
const auto global_val_max = warp::reduce_max<kNumWarps>(local_val_max);
const auto rescale = expf(local_val_max - global_val_max);
const auto global_exp_sum = warp::reduce_sum<kNumWarps>(local_exp_sum * rescale);
const auto final_scale = rescale / global_exp_sum;
const auto global_product = warp::reduce_sum<kNumWarps>(local_product * final_scale);
kv_out[local_elem_id] = global_product;
if constexpr (kNeedData) {
max_out[local_elem_id] = global_val_max;
sum_out[local_elem_id] = global_exp_sum;
}
}
if constexpr (kNeedData) __syncthreads();
}
/// \brief Sentinel score for padded positions in a 128-segment.
/// Must be finite so that `score - max` never produces NaN even when an
/// entire warp has only padded positions.
constexpr float kPadScore = -FLT_MAX;
/// \brief Online compress 128 prefill. Two passes share this body:
/// - `kWrite=false` (compress pass): handles segments that close a chunk.
/// May load prior partial state from the buffer, but never writes to it,
/// so concurrent blocks can read the same slot without racing.
/// - `kWrite=true` (write pass): handles the trailing partial segment of each
/// batch. Each batch contributes at most one such plan, so concurrent blocks
/// touch disjoint buffer slots.
///
/// The two passes MUST run as separate kernel launches (in stream order) so
/// that all reads in pass 1 finish before any writes in pass 2 start.
template <int64_t kHeadDim, bool kWrite, bool kUsePDL>
__global__ __launch_bounds__(kPrefillBlockSize, 2) //
void flash_c128_online_prefill(const __grid_constant__ Compress128OnlinePrefillParams params) {
using namespace device;
constexpr int64_t kTileDim = kTileElements * kWarpThreads; // 64
constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
static_assert(kHeadDim % kTileDim == 0, "Head dim must be multiple of tile dim");
/// NOTE: the compiler folds the if-else at compile time.
const auto num_plans = kWrite ? params.num_write : params.num_compress;
const auto plan_ptr = kWrite ? params.write_plan : params.compress_plan;
const uint32_t global_id = blockIdx.x;
const uint32_t global_pid = global_id / kNumSplit; // plan id
const uint32_t global_sid = global_id % kNumSplit; // split id
if (global_pid >= num_plans) return;
const auto [ragged_id, batch_id, position, window_len] = plan_ptr[global_pid];
if (ragged_id == 0xFFFFFFFFu) [[unlikely]]
return;
const uint32_t warp_id = threadIdx.x / kWarpThreads;
const uint32_t lane_id = threadIdx.x % kWarpThreads;
const int32_t split_offset = global_sid * kTileDim; // int32 is enough
const auto kv_score_buffer = static_cast<float*>(params.kv_score_buffer);
const auto kv_score_input = static_cast<const float*>(params.kv_score_input);
const auto kv_compressed_output = static_cast<float*>(params.kv_compressed_output);
const auto score_bias_base = static_cast<const float*>(params.score_bias);
constexpr int64_t kElementSize = kHeadDim * 2; // | kv | score |
const uint32_t chunk_offset = (position % 128u) + 1u - window_len;
const uint32_t window_end = chunk_offset + window_len; // exclusive, in [1, 128]
const int32_t segment_start = ragged_id - (position % 128u); // can be negative, but safe
const int32_t load_index = chunk_offset != 0 ? params.load_indices[batch_id] : -1;
const int32_t store_index = kWrite ? params.indices[batch_id] : -1;
PDLWaitPrimary<kUsePDL>();
// 2 * 8 = 16 register per elem. in theory we should consume 48 register here
PrefillStorage kv[kElementsPerWarp];
PrefillStorage score[kElementsPerWarp];
PrefillStorage bias[kElementsPerWarp];
const auto warp_offset = warp_id * kElementsPerWarp;
#pragma unroll
for (uint32_t i = 0; i < kElementsPerWarp; ++i) {
const uint32_t j = i + warp_offset;
if (j >= chunk_offset && j < window_end) {
const auto kv_src_ptr = kv_score_input + (segment_start + j) * kElementSize + split_offset;
const auto score_src_ptr = kv_src_ptr + kHeadDim;
const auto bias_src_ptr = score_bias_base + j * kHeadDim + split_offset;
kv[i].load(kv_src_ptr, lane_id);
score[i].load(score_src_ptr, lane_id);
bias[i].load(bias_src_ptr, lane_id);
}
}
#pragma unroll
for (uint32_t i = 0; i < kElementsPerWarp; ++i) {
const uint32_t j = i + warp_offset;
const bool is_valid = (j >= chunk_offset && j < window_end);
#pragma unroll
for (uint32_t ii = 0; ii < kTileElements; ++ii) {
score[i][ii] = is_valid ? score[i][ii] + bias[i][ii] : kPadScore;
/// NOTE: must zero out kv on padded slots -- `c128_prefill_forward`
/// computes `kv * exp_score` where `exp_score = expf(-FLT_MAX - max) ??? 0`,
/// and IEEE-754 makes `NaN * 0 = NaN` / `+-inf * 0 = NaN`. An
/// uninitialized register can hold a NaN/inf bit pattern, so without
/// this reset a single padded warp can poison the whole softmax.
kv[i][ii] = is_valid ? kv[i][ii] : 0.0f;
}
}
__shared__ alignas(16) float seg_kv[kTileDim];
__shared__ alignas(16) float seg_max[kTileDim];
__shared__ alignas(16) float seg_sum[kTileDim];
c128_prefill_forward<true>(kv, score, seg_kv, seg_max, seg_sum, warp_id, lane_id);
PDLTriggerSecondary<kUsePDL>();
if (warp_id == 0) {
PrefillStorage out_kv_vec, out_max_vec, out_sum_vec;
out_kv_vec.load(seg_kv, lane_id);
out_max_vec.load(seg_max, lane_id);
out_sum_vec.load(seg_sum, lane_id);
if (chunk_offset != 0) {
/// NOTE: load (max, sum, kv) of the in-progress chunk for this index.
/// `load_indices` may differ from `indices` when the prior partial state
/// lives on a different slot than the slot we ultimately write to.
const auto buf_load = kv_score_buffer + load_index * (kHeadDim * 3) + split_offset;
PrefillStorage buf_max_vec, buf_sum_vec, buf_kv_vec;
buf_max_vec.load(buf_load + 0 * kHeadDim, lane_id);
buf_sum_vec.load(buf_load + 1 * kHeadDim, lane_id);
buf_kv_vec.load(buf_load + 2 * kHeadDim, lane_id);
#pragma unroll
for (uint32_t ii = 0; ii < kTileElements; ++ii) {
const float m1 = buf_max_vec[ii];
const float s1 = buf_sum_vec[ii];
const float k1 = buf_kv_vec[ii];
const float m2 = out_max_vec[ii];
const float s2 = out_sum_vec[ii];
const float k2 = out_kv_vec[ii];
const float new_max = fmaxf(m1, m2);
const float new_s1 = s1 * expf(m1 - new_max);
const float new_s2 = s2 * expf(m2 - new_max);
const float new_sum = new_s1 + new_s2;
const float new_kv = (k1 * new_s1 + k2 * new_s2) / new_sum;
out_max_vec[ii] = new_max;
out_sum_vec[ii] = new_sum;
out_kv_vec[ii] = new_kv;
}
}
if constexpr (kWrite) {
const auto buf_store = kv_score_buffer + store_index * (kHeadDim * 3) + split_offset;
reinterpret_cast<PrefillStorage*>(buf_store + 0 * kHeadDim)[lane_id] = out_max_vec;
reinterpret_cast<PrefillStorage*>(buf_store + 1 * kHeadDim)[lane_id] = out_sum_vec;
reinterpret_cast<PrefillStorage*>(buf_store + 2 * kHeadDim)[lane_id] = out_kv_vec;
} else {
const auto out_ptr = kv_compressed_output + ragged_id * kHeadDim + split_offset;
reinterpret_cast<PrefillStorage*>(out_ptr)[lane_id] = out_kv_vec;
}
}
}
template <int64_t kHeadDim, bool kUsePDL>
struct FlashCompress128OnlineKernel {
static constexpr auto decode_kernel = flash_c128_online_decode<kHeadDim, kUsePDL>;
template <bool kWrite>
static constexpr auto prefill_kernel = flash_c128_online_prefill<kHeadDim, kWrite, kUsePDL>;
static constexpr auto prefill_c_kernel = prefill_kernel</*kWrite=*/false>;
static constexpr auto prefill_w_kernel = prefill_kernel</*kWrite=*/true>;
static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
static constexpr uint32_t kDecodeBlockSize = kHeadDim / 4;
static void run_decode(
const tvm::ffi::TensorView kv_score_buffer,
const tvm::ffi::TensorView kv_score_input,
const tvm::ffi::TensorView kv_compressed_output,
const tvm::ffi::TensorView ape,
const tvm::ffi::TensorView indices,
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::Optional<tvm::ffi::TensorView> /* UNUSED */) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer (max, sum, kv)
.with_dtype<float>()
.with_device(device)
.verify(kv_score_buffer);
TensorMatcher({B, kHeadDim * 2}) // kv score input
.with_dtype<float>()
.with_device(device)
.verify(kv_score_input);
TensorMatcher({B, kHeadDim}) // kv compressed output
.with_dtype<float>()
.with_device(device)
.verify(kv_compressed_output);
TensorMatcher({128, kHeadDim}) // ape
.with_dtype<float>()
.with_device(device)
.verify(ape);
TensorMatcher({B}).with_dtype<IndiceT>().with_device(device).verify(indices);
TensorMatcher({B}).with_dtype<IndiceT>().with_device(device).verify(seq_lens);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto params = Compress128OnlineDecodeParams{
.kv_score_buffer = kv_score_buffer.data_ptr(),
.kv_score_input = kv_score_input.data_ptr(),
.kv_compressed_output = kv_compressed_output.data_ptr(),
.score_bias = ape.data_ptr(),
.indices = static_cast<const IndiceT*>(indices.data_ptr()),
.seq_lens = static_cast<const IndiceT*>(seq_lens.data_ptr()),
.batch_size = batch_size,
};
LaunchKernel(batch_size, kDecodeBlockSize, device.unwrap()) //
.enable_pdl(kUsePDL)(decode_kernel, params);
}
static void run_prefill(
const tvm::ffi::TensorView kv_score_buffer,
const tvm::ffi::TensorView kv_score_input,
const tvm::ffi::TensorView kv_compressed_output,
const tvm::ffi::TensorView ape,
const tvm::ffi::TensorView indices,
const tvm::ffi::TensorView compress_plan,
const tvm::ffi::TensorView write_plan,
const tvm::ffi::Optional<tvm::ffi::TensorView> extra) {
using namespace host;
using host::compress::kOnlinePrefillPlanDim;
using host::compress::OnlinePrefillPlanTensorDtype;
auto B = SymbolicSize{"batch_size"};
auto N = SymbolicSize{"num_q_tokens"};
auto X = SymbolicSize{"compress_tokens"};
auto Y = SymbolicSize{"write_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer (max, sum, kv) ??? 2D
.with_dtype<float>()
.with_device(device_)
.verify(kv_score_buffer);
TensorMatcher({N, kHeadDim * 2}) // kv score input
.with_dtype<float>()
.with_device(device_)
.verify(kv_score_input);
TensorMatcher({N, kHeadDim}) // kv compressed output
.with_dtype<float>()
.with_device(device_)
.verify(kv_compressed_output);
TensorMatcher({128, kHeadDim}) // ape
.with_dtype<float>()
.with_device(device_)
.verify(ape);
TensorMatcher({B}) // indices
.with_dtype<IndiceT>()
.with_device(device_)
.verify(indices);
TensorMatcher({X, kOnlinePrefillPlanDim}) // compress plan
.with_dtype<OnlinePrefillPlanTensorDtype>()
.with_device(device_)
.verify(compress_plan);
TensorMatcher({Y, kOnlinePrefillPlanDim}) // write plan
.with_dtype<OnlinePrefillPlanTensorDtype>()
.with_device(device_)
.verify(write_plan);
/// NOTE: `extra` is `load_indices`. When the previous partial state lives
/// on a slot different from the destination slot (e.g. paged buffers), the
/// caller must supply this; otherwise it defaults to `indices`.
const auto load_indices = extra.value_or(indices);
TensorMatcher({B}).with_dtype<IndiceT>().with_device(device_).verify(load_indices);
const auto device = device_.unwrap();
const auto num_c = static_cast<uint32_t>(X.unwrap());
const auto num_w = static_cast<uint32_t>(Y.unwrap());
const auto params = Compress128OnlinePrefillParams{
.kv_score_buffer = kv_score_buffer.data_ptr(),
.kv_score_input = kv_score_input.data_ptr(),
.kv_compressed_output = kv_compressed_output.data_ptr(),
.score_bias = ape.data_ptr(),
.indices = static_cast<const IndiceT*>(indices.data_ptr()),
.load_indices = static_cast<const IndiceT*>(load_indices.data_ptr()),
.compress_plan = static_cast<const OnlinePlan*>(compress_plan.data_ptr()),
.write_plan = static_cast<const OnlinePlan*>(write_plan.data_ptr()),
.num_compress = num_c,
.num_write = num_w,
};
/// NOTE: pass 1 reads the buffer (for the first segment of each batch
/// that started mid-chunk) and writes only to `kv_compressed_output`.
/// Pass 2 then writes the trailing partial state of each batch back to
/// the buffer. Stream serialization between the two launches enforces
/// read-before-write on shared buffer slots.
if (const auto num_c_blocks = num_c * kNumSplit) {
LaunchKernel(num_c_blocks, kPrefillBlockSize, device) //
.enable_pdl(kUsePDL)(prefill_c_kernel, params);
}
if (const auto num_w_blocks = num_w * kNumSplit) {
LaunchKernel(num_w_blocks, kPrefillBlockSize, device) //
.enable_pdl(kUsePDL)(prefill_w_kernel, params);
}
}
};
} // namespace
namespace host::compress {
using OnlinePlanResult = tvm::ffi::Tuple<uint32_t, uint32_t>;
struct OnlinePrefillCompressParams {
OnlinePrefillPlan* __restrict__ compress_plan;
OnlinePrefillPlan* __restrict__ write_plan;
const int64_t* __restrict__ seq_lens;
const int64_t* __restrict__ extend_lens;
uint32_t batch_size;
uint32_t num_tokens;
};
/// \brief Build the compress + write plans for online compress 128 prefill.
///
/// Each batch's `[prefix_len, prefix_len + extend_len)` range is split at
/// 128-aligned boundaries. Every resulting segment falls into one of:
/// - **compress**: closes a 128-chunk (`chunk_offset + window_len == 128`).
/// These plans only read the buffer (when starting mid-chunk) and write the
/// compressed kv to `kv_compressed_output`.
/// - **write**: trailing partial of the batch (`chunk_offset + window_len < 128`).
/// May read the buffer and always writes the new partial state back to it.
/// Each batch produces at most one such plan.
///
/// The two plans MUST be dispatched as separate kernel launches in stream
/// order so that pass-1 reads of a buffer slot complete before any pass-2
/// write of the same slot.
inline OnlinePlanResult plan_online_prefill_host(const OnlinePrefillCompressParams& params, const bool use_cuda_graph) {
const auto& [compress_plan, write_plan, seq_lens, extend_lens, batch_size, num_tokens] = params;
uint32_t counter = 0;
uint32_t compress_count = 0;
uint32_t write_count = 0;
for (const auto i : irange(batch_size)) {
const uint32_t seq_len = static_cast<uint32_t>(seq_lens[i]);
const uint32_t extend_len = static_cast<uint32_t>(extend_lens[i]);
RuntimeCheck(0 < extend_len && extend_len <= seq_len);
const uint32_t prefix_len = seq_len - extend_len;
const uint32_t end_pos = prefix_len + extend_len;
/// NOTE: split the extend range into per-128-chunk segments. Each segment
/// stays inside one chunk, so the kernel can decide load/store from
/// `chunk_offset` and `window_len` alone.
uint32_t pos = prefix_len;
while (pos < end_pos) {
const uint32_t chunk_start = (pos / 128u) * 128u;
const uint32_t seg_end = std::min(end_pos, chunk_start + 128u); // exclusive
const uint32_t seg_len = seg_end - pos;
const uint32_t chunk_off = pos - chunk_start;
/// NOTE: store last-token coordinates so that downstream consumers
/// (e.g. `fused_norm_rope`) can read `ragged_id` and `position` with the
/// same semantics as `PrefillPlan`. The segment start is recoverable as
/// `ragged_id - window_len + 1` and `position - window_len + 1`.
const uint32_t last_pos = seg_end - 1;
const uint32_t last_ragged = counter + (last_pos - prefix_len);
const auto plan = OnlinePrefillPlan{
.ragged_id = last_ragged,
.batch_id = i,
.position = last_pos,
.window_len = seg_len,
};
if (chunk_off + seg_len == 128u) {
// full chunk, must be complete, maybe read the buffer, no write
RuntimeCheck(compress_count < num_tokens);
compress_plan[compress_count++] = plan;
} else {
// last chunk, must be incomplete, maybe read the buffer, must write
RuntimeCheck(write_count < num_tokens);
write_plan[write_count++] = plan;
}
pos = seg_end;
}
counter += extend_len;
}
RuntimeCheck(counter == num_tokens, "input size ", counter, " != num_q_tokens ", num_tokens);
if (!use_cuda_graph) return OnlinePlanResult{compress_count, write_count};
/// NOTE: pad both plans with sentinel entries so cuda-graph runs always see
/// the same number of blocks. The kernel skips plans whose `ragged_id` is -1.
constexpr auto kInvalid = static_cast<uint32_t>(-1);
constexpr auto kInvalidPlan = OnlinePrefillPlan{kInvalid, kInvalid, kInvalid, kInvalid};
for (const auto i : irange(compress_count, num_tokens)) {
compress_plan[i] = kInvalidPlan;
}
for (const auto i : irange(write_count, num_tokens)) {
write_plan[i] = kInvalidPlan;
}
return OnlinePlanResult{num_tokens, num_tokens};
}
inline OnlinePlanResult plan_online_prefill(
const tvm::ffi::TensorView extend_lens,
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::TensorView compress_plan,
const tvm::ffi::TensorView write_plan,
const bool use_cuda_graph) {
auto N = SymbolicSize{"batch_size"};
auto M = SymbolicSize{"num_tokens"};
auto device = SymbolicDevice{};
/// NOTE: only host (CPU/cuda-host) planning is implemented for now. The
device.set_options<kDLCPU, kDLCUDAHost>();
TensorMatcher({N}) //
.with_dtype<int64_t>()
.with_device(device)
.verify(extend_lens)
.verify(seq_lens);
TensorMatcher({M, kOnlinePrefillPlanDim}) //
.with_dtype<OnlinePrefillPlanTensorDtype>()
.with_device(device)
.verify(compress_plan)
.verify(write_plan);
const auto params = OnlinePrefillCompressParams{
.compress_plan = static_cast<OnlinePrefillPlan*>(compress_plan.data_ptr()),
.write_plan = static_cast<OnlinePrefillPlan*>(write_plan.data_ptr()),
.seq_lens = static_cast<const int64_t*>(seq_lens.data_ptr()),
.extend_lens = static_cast<const int64_t*>(extend_lens.data_ptr()),
.batch_size = static_cast<uint32_t>(N.unwrap()),
.num_tokens = static_cast<uint32_t>(M.unwrap()),
};
return plan_online_prefill_host(params, use_cuda_graph);
}
} // namespace host::compress
namespace {
[[maybe_unused]]
constexpr auto& plan_compress_online_prefill = host::compress::plan_online_prefill;
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh
+876
View File
@@ -0,0 +1,876 @@
#include <sgl_kernel/ffi.h>
#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 <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/compress_v2.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/container/tuple.h>
#include <algorithm>
#include <cfloat>
#include <cstdint>
#include <cstdlib>
#include <cstring>
namespace {
using PlanD = device::compress::DecodePlan;
using PlanC = device::compress::CompressPlan;
// ---------------------------------------------------------------------------
// Decode kernel: 1 token / batch. Each block handles one batch.
// 4 elements per thread -> kBlockSize = head_dim / 4.
// ---------------------------------------------------------------------------
struct Compress128OnlineDecodeParams {
void* __restrict__ kv_score_buffer; // [num_slots, 1, head_dim * 3]
const void* __restrict__ kv_score_input; // [batch_size, head_dim * 2]
void* __restrict__ kv_compressed_output; // [batch_size, head_dim]
const void* __restrict__ score_bias; // [128, head_dim]
const PlanD* __restrict__ plan_d;
uint32_t batch_size;
};
template <int64_t kHeadDim, bool kUsePDL>
__global__ void flash_c128_online_decode_v2(const __grid_constant__ Compress128OnlineDecodeParams params) {
using namespace device;
constexpr uint32_t kVecSize = 4;
constexpr uint32_t kBlockSize = kHeadDim / kVecSize;
using Vec = AlignedVector<float, kVecSize>;
const auto gmem = tile::Memory<Vec>::cta(kBlockSize);
const auto batch_id = blockIdx.x;
if (batch_id >= params.batch_size) return;
// Wait for the plan-finalize kernel to publish `plan.read_page_0 / write_loc`
// before reading the plan. The plan kernel runs on the same stream and does
// NOT issue a PDL trigger, so launching this kernel with PDL means our
// pre-wait global reads can race with the plan kernel's writes.
PDLWaitPrimary<kUsePDL>();
const auto plan = params.plan_d[batch_id];
const auto pos_in_chunk = (plan.seq_len - 1) % 128;
const auto kv_score_buffer = static_cast<float*>(params.kv_score_buffer);
const auto kv_score_input = static_cast<const float*>(params.kv_score_input);
const auto kv_load_buf = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3);
const auto kv_store_buf = kv_score_buffer + plan.write_loc * (kHeadDim * 3);
const auto kv_src = kv_score_input + batch_id * (kHeadDim * 2);
// Buffer layout: [max | sum | kv] (slot 0 / 1 / 2 of the head_dim*3 row).
const auto new_kv_vec = gmem.load(kv_src, 0);
const auto new_score_raw_vec = gmem.load(kv_src, 1);
const auto bias_vec = gmem.load(params.score_bias, pos_in_chunk);
Vec out_kv_vec;
Vec out_max_vec;
Vec out_sum_vec;
if (pos_in_chunk != 0) {
// Mid-chunk: combine prior partial state with the new token.
const auto max_score_vec = gmem.load(kv_load_buf, 0);
const auto sum_score_vec = gmem.load(kv_load_buf, 1);
const auto old_kv_vec = gmem.load(kv_load_buf, 2);
#pragma unroll
for (uint32_t i = 0; i < kVecSize; ++i) {
const auto old_max = max_score_vec[i];
const auto old_kv = old_kv_vec[i];
const auto new_score = new_score_raw_vec[i] + bias_vec[i];
const auto new_kv = new_kv_vec[i];
const auto new_max = fmaxf(old_max, new_score);
const auto old_sum = sum_score_vec[i] * expf(old_max - new_max);
const auto new_exp = expf(new_score - new_max);
const auto new_sum = old_sum + new_exp;
out_kv_vec[i] = (old_kv * old_sum + new_kv * new_exp) / new_sum;
out_max_vec[i] = new_max;
out_sum_vec[i] = new_sum;
}
} else {
// First token of a new chunk: state == this token alone.
#pragma unroll
for (uint32_t i = 0; i < kVecSize; ++i) {
out_kv_vec[i] = new_kv_vec[i];
out_max_vec[i] = new_score_raw_vec[i] + bias_vec[i];
out_sum_vec[i] = 1.0f;
}
}
if (pos_in_chunk == 127) {
// Chunk just closed: emit compressed kv, no buffer update.
const auto kv_out = static_cast<float*>(params.kv_compressed_output) + batch_id * kHeadDim;
gmem.store(kv_out, out_kv_vec);
} else {
gmem.store(kv_store_buf, out_max_vec, 0);
gmem.store(kv_store_buf, out_sum_vec, 1);
gmem.store(kv_store_buf, out_kv_vec, 2);
}
}
// ---------------------------------------------------------------------------
// Prefill kernel: 1 segment / block. Two passes (compress + write) share the
// kernel template, parameterized by `kWrite`.
// 16 warps per block; each warp handles 8 of the 128 chunk positions.
// ---------------------------------------------------------------------------
constexpr int32_t kTileElements = 2; // split along head-dim
constexpr int32_t kElementsPerWarp = 8; // split along the 128-chunk
constexpr uint32_t kNumWarps = 128 / kElementsPerWarp;
constexpr uint32_t kPrefillBlockSize = device::kWarpThreads * kNumWarps;
using PrefillStorage = device::AlignedVector<float, kTileElements>;
struct Compress128OnlinePrefillParams {
void* __restrict__ kv_score_buffer; // [num_slots, 1, head_dim * 3]
const void* __restrict__ kv_score_input; // [num_q_tokens, head_dim * 2]
void* __restrict__ kv_compressed_output; // [num_compress, head_dim]
const void* __restrict__ score_bias; // [128, head_dim]
const PlanC* __restrict__ plan_c; // close-chunk segments
const PlanC* __restrict__ plan_w; // trailing partial segments
uint32_t num_compress;
uint32_t num_write;
};
struct Compress128SharedBuffer {
using Storage = device::AlignedVector<float, 4>;
Storage data[kNumWarps][device::kWarpThreads + 1]; // +1 to avoid bank conflict
SGL_DEVICE Storage& operator()(uint32_t warp_id, uint32_t lane_id) {
return data[warp_id][lane_id];
}
SGL_DEVICE float& operator()(uint32_t warp_id, uint32_t lane_id, uint32_t tile_id) {
return data[warp_id][lane_id][tile_id];
}
};
/// \brief Sentinel score for padded positions in a 128-segment.
constexpr float kPadScore = -FLT_MAX;
[[maybe_unused]]
SGL_DEVICE void c128_prefill_segment_softmax(
const PrefillStorage (&kv)[kElementsPerWarp],
const PrefillStorage (&score)[kElementsPerWarp],
float* seg_kv,
float* seg_max,
float* seg_sum,
const uint32_t warp_id,
const uint32_t lane_id) {
using namespace device;
// Per-warp running state (max, sum, kv) for kTileElements head-dim slots.
using TmpStorage = typename Compress128SharedBuffer::Storage;
__shared__ Compress128SharedBuffer s_local_val_max;
__shared__ Compress128SharedBuffer s_local_exp_sum;
__shared__ Compress128SharedBuffer s_local_product;
TmpStorage tmp_val_max;
TmpStorage tmp_exp_sum;
TmpStorage tmp_product;
#pragma unroll
for (int32_t i = 0; i < kTileElements; ++i) {
float score_fp32[kElementsPerWarp];
#pragma unroll
for (int32_t j = 0; j < kElementsPerWarp; ++j) {
score_fp32[j] = score[j][i];
}
float max_value = score_fp32[0];
#pragma unroll
for (int32_t j = 1; j < kElementsPerWarp; ++j) {
max_value = fmaxf(max_value, score_fp32[j]);
}
float sum_exp_value = 0.0f;
float sum_product = 0.0f;
#pragma unroll
for (int32_t j = 0; j < kElementsPerWarp; ++j) {
const auto exp_score = expf(score_fp32[j] - max_value);
sum_product += kv[j][i] * exp_score;
sum_exp_value += exp_score;
}
tmp_val_max[i] = max_value;
tmp_exp_sum[i] = sum_exp_value;
tmp_product[i] = sum_product;
}
// Aligned writes (no bank conflict thanks to `+1` padding).
s_local_val_max(warp_id, lane_id) = tmp_val_max;
s_local_exp_sum(warp_id, lane_id) = tmp_exp_sum;
s_local_product(warp_id, lane_id) = tmp_product;
__syncthreads();
// Cross-warp reduction. Same recipe as c128_online.cuh: each block-thread
// pair reduces a (tile_id, lane_id) slot using a kNumWarps-wide warp shuffle.
constexpr uint32_t kReductionCount = kTileElements * kWarpThreads * kNumWarps;
constexpr uint32_t kIteration = kReductionCount / kPrefillBlockSize;
static_assert(kTileElements * kNumWarps == kWarpThreads, "TODO: support other configs");
#pragma unroll
for (uint32_t i = 0; i < kIteration; ++i) {
const uint32_t j = i * kPrefillBlockSize + warp_id * kWarpThreads + lane_id;
const uint32_t local_warp_id = j % kNumWarps;
const uint32_t local_elem_id = j / kNumWarps;
const uint32_t local_tile_id = local_elem_id % kTileElements;
const uint32_t local_lane_id = local_elem_id / kTileElements;
const auto local_val_max = s_local_val_max(local_warp_id, local_lane_id, local_tile_id);
const auto local_exp_sum = s_local_exp_sum(local_warp_id, local_lane_id, local_tile_id);
const auto local_product = s_local_product(local_warp_id, local_lane_id, local_tile_id);
const auto global_val_max = warp::reduce_max<kNumWarps>(local_val_max);
const auto rescale = expf(local_val_max - global_val_max);
const auto global_exp_sum = warp::reduce_sum<kNumWarps>(local_exp_sum * rescale);
const auto final_scale = rescale / global_exp_sum;
const auto global_product = warp::reduce_sum<kNumWarps>(local_product * final_scale);
seg_kv[local_elem_id] = global_product;
seg_max[local_elem_id] = global_val_max;
seg_sum[local_elem_id] = global_exp_sum;
}
__syncthreads();
}
/// \brief Online compress 128 prefill v2.
///
/// `kWrite=false` (compress pass): handles segments that close a 128-chunk.
/// Reads optional prior state from `read_page_0` (-1 = none), emits compressed
/// kv to `kv_compressed_output[plan_id]` (compact).
/// `kWrite=true` (write pass) : handles trailing partial segments.
/// Reads optional prior state from `read_page_1` (-1 = fallback to
/// `read_page_0`), writes new running state to `read_page_0`.
template <int64_t kHeadDim, bool kWrite, bool kUsePDL>
__global__ __launch_bounds__(kPrefillBlockSize, 2) //
void flash_c128_online_prefill_v2(const __grid_constant__ Compress128OnlinePrefillParams params) {
using namespace device;
constexpr int64_t kTileDim = kTileElements * kWarpThreads; // 64
constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
static_assert(kHeadDim % kTileDim == 0);
// Compile-time fold to the right plan list.
const auto num_plans = kWrite ? params.num_write : params.num_compress;
const auto plan_ptr = kWrite ? params.plan_w : params.plan_c;
const uint32_t global_id = blockIdx.x;
const uint32_t global_pid = global_id / kNumSplit;
const uint32_t global_sid = global_id % kNumSplit;
if (global_pid >= num_plans) return;
const uint32_t warp_id = threadIdx.x / kWarpThreads;
const uint32_t lane_id = threadIdx.x % kWarpThreads;
const int32_t split_offset = global_sid * kTileDim;
// The previous kernel (plan-finalize stage 1) does NOT issue a PDL trigger,
// so PDLWaitPrimary effectively waits for stage 1 to complete. Read the plan
// AFTER the wait so the freshly-written `read_page_0` (= state-pool slot) is
// visible. Reading it before the wait is a real race -- with PDL enabled the
// kernel can begin executing before stage 1's stores propagate, and we'd see
// the stage-0 batch_id placeholder in `read_page_0` instead of the slot.
PDLWaitPrimary<kUsePDL>();
const auto plan = plan_ptr[global_pid];
if (plan.is_invalid()) [[unlikely]]
return;
const auto kv_score_buffer = static_cast<float*>(params.kv_score_buffer);
const auto kv_score_input = static_cast<const float*>(params.kv_score_input);
const auto kv_compressed_output = static_cast<float*>(params.kv_compressed_output);
const auto score_bias_base = static_cast<const float*>(params.score_bias);
constexpr int64_t kElementSize = kHeadDim * 2; // | kv | score |
// `j` below is a chunk-local offset. Convert it to the ragged-input row by
// anchoring on the last token in this segment: ragged_id - pos_in_chunk_end + 1 + j.
const uint32_t window_len = plan.buffer_len;
const uint32_t position = plan.seq_len - 1;
const uint32_t pos_in_chunk_end = (position % 128u) + 1u; // exclusive, in [1, 128]
const uint32_t chunk_offset = pos_in_chunk_end - window_len; // in [0, 127]
const int32_t chunk_start_ragged = static_cast<int32_t>(plan.ragged_id) - static_cast<int32_t>(pos_in_chunk_end) + 1;
// --- Stage 1: load kv / score / bias for this warp's 8 chunk positions.
PrefillStorage kv[kElementsPerWarp];
PrefillStorage score[kElementsPerWarp];
PrefillStorage bias[kElementsPerWarp];
const uint32_t warp_offset = warp_id * kElementsPerWarp;
#pragma unroll
for (uint32_t i = 0; i < kElementsPerWarp; ++i) {
const uint32_t j = i + warp_offset;
if (j >= chunk_offset && j < pos_in_chunk_end) {
const int32_t ragged_id = chunk_start_ragged + static_cast<int32_t>(j);
const auto kv_src_ptr = kv_score_input + ragged_id * kElementSize + split_offset;
const auto score_src_ptr = kv_src_ptr + kHeadDim;
const auto bias_src_ptr = score_bias_base + j * kHeadDim + split_offset;
kv[i].load(kv_src_ptr, lane_id);
score[i].load(score_src_ptr, lane_id);
bias[i].load(bias_src_ptr, lane_id);
}
}
// --- Stage 2: pad invalid positions. score = -FLT_MAX, kv = 0 (so that
// kv * exp(score-max) ??? 0 / 0 cleanly without producing NaN/inf).
#pragma unroll
for (uint32_t i = 0; i < kElementsPerWarp; ++i) {
const uint32_t j = i + warp_offset;
const bool is_valid = (j >= chunk_offset && j < pos_in_chunk_end);
#pragma unroll
for (uint32_t ii = 0; ii < kTileElements; ++ii) {
score[i][ii] = is_valid ? score[i][ii] + bias[i][ii] : kPadScore;
kv[i][ii] = is_valid ? kv[i][ii] : 0.0f;
}
}
// --- Stage 3: warp-tile online softmax over the 128-position chunk.
__shared__ alignas(16) float seg_kv[kTileDim];
__shared__ alignas(16) float seg_max[kTileDim];
__shared__ alignas(16) float seg_sum[kTileDim];
c128_prefill_segment_softmax(kv, score, seg_kv, seg_max, seg_sum, warp_id, lane_id);
PDLTriggerSecondary<kUsePDL>();
// --- Stage 4: warp 0 folds with prior partial state (if any) and writes.
if (warp_id == 0) {
PrefillStorage out_kv_vec, out_max_vec, out_sum_vec;
out_kv_vec.load(seg_kv, lane_id);
out_max_vec.load(seg_max, lane_id);
out_sum_vec.load(seg_sum, lane_id);
const int32_t read_page = plan.read_page_1 >= 0 ? plan.read_page_1 : plan.read_page_0;
if (chunk_offset != 0 && read_page >= 0) {
// Combine with prior partial state for this slot.
const auto buf_load = kv_score_buffer + read_page * (kHeadDim * 3) + split_offset;
PrefillStorage buf_max_vec, buf_sum_vec, buf_kv_vec;
buf_max_vec.load(buf_load + 0 * kHeadDim, lane_id);
buf_sum_vec.load(buf_load + 1 * kHeadDim, lane_id);
buf_kv_vec.load(buf_load + 2 * kHeadDim, lane_id);
#pragma unroll
for (uint32_t ii = 0; ii < kTileElements; ++ii) {
const float m1 = buf_max_vec[ii];
const float s1 = buf_sum_vec[ii];
const float k1 = buf_kv_vec[ii];
const float m2 = out_max_vec[ii];
const float s2 = out_sum_vec[ii];
const float k2 = out_kv_vec[ii];
const float new_max = fmaxf(m1, m2);
const float new_s1 = s1 * expf(m1 - new_max);
const float new_s2 = s2 * expf(m2 - new_max);
const float new_sum = new_s1 + new_s2;
const float new_kv = (k1 * new_s1 + k2 * new_s2) / new_sum;
out_max_vec[ii] = new_max;
out_sum_vec[ii] = new_sum;
out_kv_vec[ii] = new_kv;
}
}
if constexpr (kWrite) {
// For trailing-partial segments the load and store slots collapse to the
// segment's own chunk slot (the request keeps a single in-progress
// chunk's running state at any time), so we reuse `read_page_0`.
const auto buf_store = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3) + split_offset;
reinterpret_cast<PrefillStorage*>(buf_store + 0 * kHeadDim)[lane_id] = out_max_vec;
reinterpret_cast<PrefillStorage*>(buf_store + 1 * kHeadDim)[lane_id] = out_sum_vec;
reinterpret_cast<PrefillStorage*>(buf_store + 2 * kHeadDim)[lane_id] = out_kv_vec;
} else {
// Compact output: one row per compress plan, indexed by `global_pid`.
const auto out_ptr = kv_compressed_output + global_pid * kHeadDim + split_offset;
reinterpret_cast<PrefillStorage*>(out_ptr)[lane_id] = out_kv_vec;
}
}
}
// ---------------------------------------------------------------------------
// Host wrapper: matches the c128_v2 / c4_v2 host API style (run_decode /
// run_prefill methods on a kernel-class template). We only expose `kHeadDim`
// + `kUsePDL`; the dtype is fixed to fp32 for the online state pool.
// ---------------------------------------------------------------------------
template <int64_t kHeadDim, bool kUsePDL>
struct FlashCompress128OnlineKernel {
static constexpr auto decode_kernel = flash_c128_online_decode_v2<kHeadDim, kUsePDL>;
template <bool kWrite>
static constexpr auto prefill_kernel = flash_c128_online_prefill_v2<kHeadDim, kWrite, kUsePDL>;
static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
static constexpr uint32_t kDecodeBlockSize = kHeadDim / 4;
static void run_decode(
const tvm::ffi::TensorView kv_score_buffer,
const tvm::ffi::TensorView kv_score_input,
const tvm::ffi::TensorView kv_compressed_output,
const tvm::ffi::TensorView ape,
const tvm::ffi::TensorView plan_d_) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer (max, sum, kv)
.with_dtype<float>()
.with_device(device_)
.verify(kv_score_buffer);
TensorMatcher({B, kHeadDim * 2}) // kv score input
.with_dtype<float>()
.with_device(device_)
.verify(kv_score_input);
TensorMatcher({B, kHeadDim}) // kv compressed output (sparse by batch_id)
.with_dtype<float>()
.with_device(device_)
.verify(kv_compressed_output);
TensorMatcher({128, kHeadDim}) // ape
.with_dtype<float>()
.with_device(device_)
.verify(ape);
const auto plan_d = compress::verify_plan_d(plan_d_, B, device_);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
if (batch_size == 0) return;
const auto params = Compress128OnlineDecodeParams{
.kv_score_buffer = kv_score_buffer.data_ptr(),
.kv_score_input = kv_score_input.data_ptr(),
.kv_compressed_output = kv_compressed_output.data_ptr(),
.score_bias = ape.data_ptr(),
.plan_d = plan_d,
.batch_size = batch_size,
};
LaunchKernel(batch_size, kDecodeBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(decode_kernel, params);
}
static void run_prefill(
const tvm::ffi::TensorView kv_score_buffer,
const tvm::ffi::TensorView kv_score_input,
const tvm::ffi::TensorView kv_compressed_output,
const tvm::ffi::TensorView ape,
const tvm::ffi::TensorView plan_c_,
const tvm::ffi::TensorView plan_w_) {
using namespace host;
auto N = SymbolicSize{"num_q_tokens"};
auto C = SymbolicSize{"num_c_plans"};
auto W = SymbolicSize{"num_w_plans"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer
.with_dtype<float>()
.with_device(device_)
.verify(kv_score_buffer);
TensorMatcher({N, kHeadDim * 2}) // kv score input (ragged)
.with_dtype<float>()
.with_device(device_)
.verify(kv_score_input);
TensorMatcher({C, kHeadDim}) // kv compressed output (compact, by plan_c index)
.with_dtype<float>()
.with_device(device_)
.verify(kv_compressed_output);
TensorMatcher({128, kHeadDim}) // ape
.with_dtype<float>()
.with_device(device_)
.verify(ape);
// Both compress and write segments use PlanC layout. Stage 1 stores the
// committed-bank load slot in read_page_1 and the write slot in read_page_0.
const auto plan_c = compress::verify_plan_c(plan_c_, C, device_);
const auto plan_w = compress::verify_plan_c(plan_w_, W, device_);
const auto device = device_.unwrap();
const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
const auto num_c = static_cast<uint32_t>(C.unwrap());
const auto num_w = static_cast<uint32_t>(W.unwrap());
RuntimeCheck(num_q_tokens >= num_w, "invalid prefill plan: num_q < num_w");
const auto params = Compress128OnlinePrefillParams{
.kv_score_buffer = kv_score_buffer.data_ptr(),
.kv_score_input = kv_score_input.data_ptr(),
.kv_compressed_output = kv_compressed_output.data_ptr(),
.score_bias = ape.data_ptr(),
.plan_c = plan_c,
.plan_w = plan_w,
.num_compress = num_c,
.num_write = num_w,
};
// The two passes MUST be serialized in stream order: pass 1 reads slots
// that pass 2 may write to; running them in parallel would race.
if (const auto num_c_blocks = num_c * kNumSplit) {
LaunchKernel(num_c_blocks, kPrefillBlockSize, device) //
.enable_pdl(kUsePDL)(prefill_kernel</*kWrite=*/false>, params);
}
if (const auto num_w_blocks = num_w * kNumSplit) {
LaunchKernel(num_w_blocks, kPrefillBlockSize, device) //
.enable_pdl(kUsePDL)(prefill_kernel</*kWrite=*/true>, params);
}
}
};
} // namespace
// ===========================================================================
// Plan builders. Mirrors the offline v2 pattern (`c_plan.cuh`):
// - Decode: a single GPU kernel reads seq_lens / req_to_token /
// req_pool_indices on device and emits the final PlanD tensor in one go.
// - Prefill: stage 0 (host, on CPU pinned memory) splits each batch's
// extend range into per-chunk segments and emits PlanC entries with the
// batch_id stashed in `read_page_0` as a placeholder. Stage 1 is a tiny
// GPU kernel that finalizes `read_page_0` to `req_to_token[rid][chunk_start]`,
// so the slot tensors never leave GPU memory. The online state pool keeps
// a single in-progress chunk per request, so each segment's load and
// store slot collapse to one value (the slot for the segment's own chunk).
// For online-c128 MTP, stage 1 keeps that write slot in `read_page_0` and
// stores the committed-bank load slot in `read_page_1`.
// ===========================================================================
namespace host::compress {
using device::compress::CompressPlan;
using device::compress::DecodePlan;
// ---------------------------------------------------------------------------
// Decode plan builder.
// ---------------------------------------------------------------------------
struct OnlineDecodePlanParams {
DecodePlan* __restrict__ plan_d;
const int64_t* __restrict__ seq_lens;
const int64_t* __restrict__ req_pool_indices;
const int32_t* __restrict__ req_to_token;
int64_t stride_r2t;
int32_t state_slot_offset;
uint32_t batch_size;
};
__global__ void plan_c128_online_decode_kernel(const OnlineDecodePlanParams params) {
const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= params.batch_size) return;
const auto seq_len = static_cast<uint32_t>(params.seq_lens[idx]);
const auto rid = params.req_pool_indices[idx];
const int32_t slot = static_cast<int32_t>(rid) + params.state_slot_offset;
params.plan_d[idx] = DecodePlan{
.seq_len = seq_len,
.write_loc = slot,
.read_page_0 = slot,
.read_page_1 = -1,
};
}
/// \brief Build the decode plan tensor. Caller (Python) pre-allocates
/// `plan_d_dev` as a `(batch_size, 16)` device uint8 tensor; this routine
/// only fills it. See `plan_online_prefill` for the rationale (avoid
/// `ffi::empty` + dlpack roundtrip / PyTorch caching-allocator stream
/// tracking issue that surfaces as IMA in unrelated downstream kernels).
inline void plan_online_decode(
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::TensorView req_pool_indices,
const tvm::ffi::TensorView req_to_token,
const tvm::ffi::TensorView plan_d_dev_,
const int32_t state_slot_offset) {
auto B = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
auto seq_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int64_t>(seq_dtype)
.with_device(device_)
.verify(seq_lens);
TensorMatcher({B}) //
.with_dtype<int64_t>()
.with_device(device_)
.verify(req_pool_indices);
TensorMatcher({-1, -1}) //
.with_dtype<int32_t>()
.with_device(device_)
.verify(req_to_token);
TensorMatcher({B, sizeof(DecodePlan)}) //
.with_dtype<uint8_t>()
.with_device(device_)
.verify(plan_d_dev_);
RuntimeCheck(state_slot_offset >= 0);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
if (batch_size == 0) return;
const auto device = device_.unwrap();
constexpr uint32_t kBlockSize = 256;
const uint32_t num_blocks = host::div_ceil(batch_size, kBlockSize);
const auto stride_r2t = req_to_token.stride(0);
const auto params = OnlineDecodePlanParams{
.plan_d = static_cast<DecodePlan*>(plan_d_dev_.data_ptr()),
.seq_lens = static_cast<const int64_t*>(seq_lens.data_ptr()),
.req_pool_indices = static_cast<const int64_t*>(req_pool_indices.data_ptr()),
.req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
.stride_r2t = stride_r2t,
.state_slot_offset = state_slot_offset,
.batch_size = batch_size,
};
LaunchKernel(num_blocks, kBlockSize, device)(plan_c128_online_decode_kernel, params);
}
// ---------------------------------------------------------------------------
// Prefill plan builder: host stage 0 + GPU stage 1.
// ---------------------------------------------------------------------------
struct OnlinePrefillStage0Params {
CompressPlan* __restrict__ plan_c;
CompressPlan* __restrict__ plan_w;
const int64_t* __restrict__ seq_lens;
const int64_t* __restrict__ extend_lens;
uint32_t batch_size;
uint32_t num_q_tokens;
};
inline std::tuple<uint32_t, uint32_t> _plan_prefill_partial(const OnlinePrefillStage0Params& p) {
uint32_t counter = 0;
uint32_t compress_count = 0;
uint32_t write_count = 0;
for (const auto i : irange(p.batch_size)) {
const uint32_t seq_len = static_cast<uint32_t>(p.seq_lens[i]);
const uint32_t extend_len = static_cast<uint32_t>(p.extend_lens[i]);
RuntimeCheck(0 < extend_len && extend_len <= seq_len);
const uint32_t prefix_len = seq_len - extend_len;
const uint32_t end_pos = prefix_len + extend_len;
uint32_t pos = prefix_len;
while (pos < end_pos) {
const uint32_t chunk_start = (pos / 128u) * 128u;
const uint32_t seg_end = std::min(end_pos, chunk_start + 128u); // exclusive
const uint32_t seg_len = seg_end - pos;
const uint32_t chunk_off = pos - chunk_start;
const uint32_t last_pos = seg_end - 1;
const uint32_t last_ragged = counter + (last_pos - prefix_len);
RuntimeCheck(last_ragged < (1u << 16), "PlanC.ragged_id is uint16; ragged ", last_ragged, " overflows");
RuntimeCheck(seg_len <= 128u);
// Stash batch_id in `read_page_0` for stage 1 to translate. A
// chunk-aligned segment never loads, so we still need stage 1 to fill
// a slot in -- the kernel keys the load on `chunk_offset != 0`.
const auto plan = CompressPlan{
.seq_len = last_pos + 1u,
.ragged_id = static_cast<uint16_t>(last_ragged),
.buffer_len = static_cast<uint16_t>(seg_len),
.read_page_0 = static_cast<int32_t>(i), // batch_id placeholder
.read_page_1 = -1, // filled by stage 1 with committed-bank slot
};
if (chunk_off + seg_len == 128u) {
// close-chunk segment
RuntimeCheck(compress_count < p.num_q_tokens);
p.plan_c[compress_count++] = plan;
} else {
// trailing partial segment
RuntimeCheck(write_count < p.num_q_tokens);
p.plan_w[write_count++] = plan;
}
pos = seg_end;
}
counter += extend_len;
}
RuntimeCheck(counter == p.num_q_tokens, "input size ", counter, " != num_q_tokens ", p.num_q_tokens);
return std::tuple<uint32_t, uint32_t>{compress_count, write_count};
}
struct OnlinePrefillStage1Params {
CompressPlan* __restrict__ plan_c;
CompressPlan* __restrict__ plan_w;
const int64_t* __restrict__ req_pool_indices; // (batch_size,)
const int32_t* __restrict__ req_to_token; // (num_reqs, max_tokens)
int64_t stride_r2t;
int32_t state_slot_offset;
uint32_t num_c;
uint32_t num_w;
};
__global__ void plan_c128_online_prefill_kernel(const OnlinePrefillStage1Params params) {
const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t total = params.num_c + params.num_w;
if (idx >= total) return;
const bool is_compress = idx < params.num_c;
CompressPlan* const plan_ptr = is_compress ? &params.plan_c[idx] : &params.plan_w[idx - params.num_c];
auto plan = *plan_ptr;
if (plan.is_invalid()) return;
const auto batch_id = plan.read_page_0;
const auto rid = params.req_pool_indices[batch_id];
const int32_t main_slot = static_cast<int32_t>(rid);
plan.read_page_0 = main_slot + params.state_slot_offset;
plan.read_page_1 = main_slot;
*plan_ptr = plan;
}
using OnlinePrefillPlan = tvm::ffi::Tuple<uint32_t, uint32_t>;
inline OnlinePrefillPlan plan_online_prefill(
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::TensorView extend_lens,
const tvm::ffi::TensorView req_pool_indices,
const tvm::ffi::TensorView req_to_token,
const tvm::ffi::TensorView plan_c_pin,
const tvm::ffi::TensorView plan_w_pin,
const tvm::ffi::TensorView plan_c_dev_,
const tvm::ffi::TensorView plan_w_dev_,
const int32_t state_slot_offset,
const bool use_cuda_graph) {
auto B = SymbolicSize{"batch_size"};
auto N = SymbolicSize{"num_q_tokens"};
auto cpu = SymbolicDevice{};
auto device_ = SymbolicDevice{};
cpu.set_options<kDLCPU, kDLCUDAHost>();
device_.set_options<kDLCUDA>();
TensorMatcher({B}) //
.with_dtype<int64_t>()
.with_device(cpu)
.verify(seq_lens)
.verify(extend_lens);
TensorMatcher({B}) //
.with_dtype<int64_t>()
.with_device(device_)
.verify(req_pool_indices);
TensorMatcher({-1, -1}) //
.with_dtype<int32_t>()
.with_device(device_)
.verify(req_to_token);
TensorMatcher({N, sizeof(CompressPlan)}) //
.with_dtype<uint8_t>()
.with_device(cpu)
.verify(plan_c_pin)
.verify(plan_w_pin);
TensorMatcher({N, sizeof(CompressPlan)}) //
.with_dtype<uint8_t>()
.with_device(device_)
.verify(plan_c_dev_)
.verify(plan_w_dev_);
RuntimeCheck(state_slot_offset >= 0);
const auto stage0_params = OnlinePrefillStage0Params{
.plan_c = static_cast<CompressPlan*>(plan_c_pin.data_ptr()),
.plan_w = static_cast<CompressPlan*>(plan_w_pin.data_ptr()),
.seq_lens = static_cast<const int64_t*>(seq_lens.data_ptr()),
.extend_lens = static_cast<const int64_t*>(extend_lens.data_ptr()),
.batch_size = static_cast<uint32_t>(B.unwrap()),
.num_q_tokens = static_cast<uint32_t>(N.unwrap()),
};
// Debug instrumentation: SGLANG_DEBUG_C128_ONLINE_GUARD=1 wraps stage 0
// with redzone + post-write magic-check on the pin buffers, plus a strict
// upper-bound check on `batch_size` and `num_q_tokens`. If stage 0 has a
// CPU OOB this trips a clear panic at the offending byte instead of a
// delayed CUDA IMA from corrupted heap memory.
static const bool kGuard = []() {
const char* v = std::getenv("SGLANG_DEBUG_C128_ONLINE_GUARD");
return v != nullptr && v[0] == '1';
}();
if (kGuard) {
RuntimeCheck(stage0_params.batch_size <= 65536u, "batch_size out of bound: ", stage0_params.batch_size);
RuntimeCheck(stage0_params.num_q_tokens <= 65536u, "num_q_tokens out of bound: ", stage0_params.num_q_tokens);
// Stamp the pin buffers with 0xAB so we can detect any byte still 0xAB
// beyond what stage 0 should have written (= OOB never reached, that's fine)
// or any byte BEYOND num_q_tokens*16 written to (= true OOB into
// adjacent allocation).
auto* pc = static_cast<uint8_t*>(plan_c_pin.data_ptr());
auto* pw = static_cast<uint8_t*>(plan_w_pin.data_ptr());
const auto bytes = static_cast<size_t>(N.unwrap()) * sizeof(CompressPlan);
std::memset(pc, 0xAB, bytes);
std::memset(pw, 0xAB, bytes);
}
const auto [num_c, num_w] = _plan_prefill_partial(stage0_params);
const auto num_c_padded = use_cuda_graph ? static_cast<uint32_t>(N.unwrap()) : num_c;
const auto num_w_padded = use_cuda_graph ? static_cast<uint32_t>(N.unwrap()) : num_w;
if (kGuard) {
// Verify stage 0 wrote ONLY to the [0, num_c*16) and [0, num_w*16) prefix.
auto* pc = static_cast<const uint8_t*>(plan_c_pin.data_ptr());
auto* pw = static_cast<const uint8_t*>(plan_w_pin.data_ptr());
const auto end_c = static_cast<size_t>(num_c) * sizeof(CompressPlan);
const auto end_w = static_cast<size_t>(num_w) * sizeof(CompressPlan);
const auto pin_bytes = static_cast<size_t>(N.unwrap()) * sizeof(CompressPlan);
for (size_t k = end_c; k < pin_bytes; ++k) {
RuntimeCheck(
pc[k] == 0xAB,
"GUARD: plan_c_pin OOB write at byte ",
k,
" (num_c=",
num_c,
", num_q_tokens=",
N.unwrap(),
")");
}
for (size_t k = end_w; k < pin_bytes; ++k) {
RuntimeCheck(
pw[k] == 0xAB,
"GUARD: plan_w_pin OOB write at byte ",
k,
" (num_w=",
num_w,
", num_q_tokens=",
N.unwrap(),
")");
}
}
const auto device = device_.unwrap();
// Out-params pre-allocated by Python. Cast to typed pointers for use.
auto* const plan_c_dev_ptr = static_cast<CompressPlan*>(plan_c_dev_.data_ptr());
auto* const plan_w_dev_ptr = static_cast<CompressPlan*>(plan_w_dev_.data_ptr());
if (use_cuda_graph) {
const auto kInvalidPlan = CompressPlan::invalid();
auto* const plan_c_pin_ptr = static_cast<CompressPlan*>(plan_c_pin.data_ptr());
auto* const plan_w_pin_ptr = static_cast<CompressPlan*>(plan_w_pin.data_ptr());
for (const auto i : irange(num_c, num_c_padded)) {
plan_c_pin_ptr[i] = kInvalidPlan;
}
for (const auto i : irange(num_w, num_w_padded)) {
plan_w_pin_ptr[i] = kInvalidPlan;
}
}
if (const auto total = num_c_padded + num_w_padded) {
const auto stream = LaunchKernel::resolve_device(device);
// SGLANG_DEBUG_C128_ONLINE_SYNC_H2D=1 forces a synchronous H2D copy.
static const bool kSyncH2D = []() {
const char* v = std::getenv("SGLANG_DEBUG_C128_ONLINE_SYNC_H2D");
return v != nullptr && v[0] == '1';
}();
// SGLANG_DEBUG_C128_ONLINE_NO_H2D=1 skips the H2D copy entirely (debug only).
static const bool kNoH2D = []() {
const char* v = std::getenv("SGLANG_DEBUG_C128_ONLINE_NO_H2D");
return v != nullptr && v[0] == '1';
}();
const auto copy_to_device = [stream](void* dst, void* src, int64_t count) {
if (kNoH2D) return;
const auto bytes = count * sizeof(CompressPlan);
if (kSyncH2D) {
RuntimeDeviceCheck(::cudaMemcpy(dst, src, bytes, ::cudaMemcpyHostToDevice));
} else {
RuntimeDeviceCheck(::cudaMemcpyAsync(dst, src, bytes, ::cudaMemcpyHostToDevice, stream));
}
};
if (num_c_padded) copy_to_device(plan_c_dev_ptr, plan_c_pin.data_ptr(), num_c_padded);
if (num_w_padded) copy_to_device(plan_w_dev_ptr, plan_w_pin.data_ptr(), num_w_padded);
const auto stage1_params = OnlinePrefillStage1Params{
.plan_c = plan_c_dev_ptr,
.plan_w = plan_w_dev_ptr,
.req_pool_indices = static_cast<const int64_t*>(req_pool_indices.data_ptr()),
.req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
.stride_r2t = req_to_token.stride(0),
.state_slot_offset = state_slot_offset,
.num_c = num_c_padded,
.num_w = num_w_padded,
};
constexpr uint32_t kBlockSize = 128;
const auto num_blocks = host::div_ceil(total, kBlockSize);
LaunchKernel(num_blocks, kBlockSize, device)(plan_c128_online_prefill_kernel, stage1_params);
}
return OnlinePrefillPlan{num_c_padded, num_w_padded};
}
} // namespace host::compress
namespace {
[[maybe_unused]]
constexpr auto& plan_compress_128_online_decode = host::compress::plan_online_decode;
[[maybe_unused]]
constexpr auto& plan_compress_128_online_prefill = host::compress::plan_online_prefill;
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh
+512
View File
@@ -0,0 +1,512 @@
/**
* \brief Here's some dimension info for the main buffer used in C128 prefill and decode.
*
* kv_buffer: [num_indices, 128, head_dim * 2]
* - last dimension layout: | kv | score |
* kv_input: [batch_size, head_dim * 2]
* kv_output: [batch_size, head_dim]
* score_bias (ape): [128, head_dim]
* plan_c/plan_w: [variable length]
*
* For prefill, batch_size = num_q_tokens
*/
#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 <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/compress_v2.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/object.h>
#include <cfloat>
#include <cstdint>
#include <type_traits>
namespace {
using PlanD = device::compress::DecodePlan;
using PlanC = device::compress::CompressPlan;
using PlanW = device::compress::WritePlan;
/// \brief Each thread will handle this many elements (split along head_dim)
constexpr int32_t kTileElements = 2;
/// \brief Each warp will handle this many elements (split along 128)
constexpr int32_t kElementsPerWarp = 8;
constexpr uint32_t kNumWarps = 128 / kElementsPerWarp;
constexpr uint32_t kBlockSize = device::kWarpThreads * kNumWarps;
constexpr uint32_t kWriteBlockSize = 128; // one warp per write
/// \brief Need to reduce register usage to increase occupancy
#define C128_KERNEL __global__ __launch_bounds__(kBlockSize, 2)
#define WRITE_KERNEL __global__ __launch_bounds__(kWriteBlockSize, 16)
struct Compress128DecodeParams {
void* __restrict__ kv_buffer;
const void* __restrict__ kv_input;
void* __restrict__ kv_output;
const void* __restrict__ score_bias;
const PlanD* __restrict__ plan_d;
uint32_t batch_size;
};
struct Compress128PrefillParams {
void* __restrict__ kv_buffer;
const void* __restrict__ kv_input;
void* __restrict__ kv_output;
const void* __restrict__ score_bias;
const PlanC* __restrict__ plan_c;
const PlanW* __restrict__ plan_w;
uint32_t num_compress;
uint32_t num_write;
};
struct Compress128SharedBuffer {
using Storage = device::AlignedVector<float, kTileElements>;
Storage data[kNumWarps][device::kWarpThreads + 1]; // padding to avoid bank conflict
SGL_DEVICE Storage& operator()(uint32_t warp_id, uint32_t lane_id) {
return data[warp_id][lane_id];
}
SGL_DEVICE float& operator()(uint32_t warp_id, uint32_t lane_id, uint32_t tile_id) {
return data[warp_id][lane_id][tile_id];
}
};
template <int64_t kHeadDim_>
struct C128Trait {
static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64
static constexpr int64_t kHeadDim = kHeadDim_;
static constexpr int64_t kScoreOffset = kHeadDim;
static constexpr int64_t kElementSize = kHeadDim * 2;
static constexpr int64_t kPageElementSize = 128 * kElementSize; // page size = 128
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
static_assert(kHeadDim % kTileDim == 0);
};
template <typename Trait, bool kUsePDL, typename BufferFloat, typename InputFloat, typename OutFloat>
SGL_DEVICE void c128_forward(
const BufferFloat* kv_buf, // [128n, 128n + 127]
const InputFloat* kv_src, // ragged pointer at position = 128n + 127
OutFloat* kv_out,
const InputFloat* score_bias,
const int32_t buffer_len) {
using namespace device;
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
/// NOTE: part 1: load kv + score
using StorageIn = AlignedVector<InputFloat, kTileElements>;
const auto gmem_in = tile::Memory<StorageIn>{lane_id, kWarpThreads};
StorageIn kv[kElementsPerWarp];
StorageIn score[kElementsPerWarp];
StorageIn bias[kElementsPerWarp];
const int32_t warp_offset = warp_id * kElementsPerWarp;
#pragma unroll
for (int32_t i = 0; i < 8; ++i) {
const int32_t j = i + warp_offset;
bias[i] = gmem_in.load(score_bias + j * Trait::kHeadDim);
}
const auto kv_start = kv_src - 127 * Trait::kElementSize; // point to start
if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
#pragma unroll
for (int32_t i = 0; i < kElementsPerWarp; ++i) {
const int32_t j = i + warp_offset;
__builtin_assume(j < 128);
const auto src = j < buffer_len ? kv_buf : kv_start;
kv[i] = gmem_in.load(src + j * Trait::kElementSize);
score[i] = gmem_in.load(src + j * Trait::kElementSize + Trait::kScoreOffset);
}
} else { // mixed dtype
using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
const auto gmem_buffer = tile::Memory<StorageBuffer>{lane_id, kWarpThreads};
#pragma unroll
for (int32_t i = 0; i < kElementsPerWarp; ++i) {
const int32_t j = i + warp_offset;
__builtin_assume(j < 128);
if (j < buffer_len) {
const auto src = kv_buf + j * Trait::kElementSize;
const auto kv_tmp = gmem_buffer.load(src);
const auto score_tmp = gmem_buffer.load(src + Trait::kScoreOffset);
#pragma unroll
for (int32_t k = 0; k < kTileElements; ++k) {
kv[i][k] = cast<InputFloat>(kv_tmp[k]);
score[i][k] = cast<InputFloat>(score_tmp[k]);
}
} else {
const auto src = kv_start + j * Trait::kElementSize;
kv[i] = gmem_in.load(src);
score[i] = gmem_in.load(src + Trait::kScoreOffset);
}
}
}
/// NOTE: part 2: safe online softmax + weighted sum
using TmpStorage = typename Compress128SharedBuffer::Storage;
__shared__ Compress128SharedBuffer s_local_val_max;
__shared__ Compress128SharedBuffer s_local_exp_sum;
__shared__ Compress128SharedBuffer s_local_product;
TmpStorage tmp_val_max;
TmpStorage tmp_exp_sum;
TmpStorage tmp_product;
float score_fp32[kTileElements][kElementsPerWarp];
// convert to fp32 and apply bias first
#pragma unroll
for (int32_t i = 0; i < kTileElements; ++i) {
#pragma unroll
for (int32_t j = 0; j < kElementsPerWarp; ++j) {
score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
}
}
#pragma unroll
for (int32_t i = 0; i < kTileElements; ++i) {
const auto& score = score_fp32[i];
float max_value = score[0];
float sum_exp_value = 0.0f;
#pragma unroll
for (int32_t j = 1; j < kElementsPerWarp; ++j) {
const auto fp32_score = score[j];
max_value = fmaxf(max_value, fp32_score);
}
float sum_product = 0.0f;
#pragma unroll
for (int32_t j = 0; j < 8; ++j) {
const auto fp32_score = score[j];
const auto exp_score = expf(fp32_score - max_value);
sum_product += cast<float>(kv[j][i]) * exp_score;
sum_exp_value += exp_score;
}
tmp_val_max[i] = max_value;
tmp_exp_sum[i] = sum_exp_value;
tmp_product[i] = sum_product;
}
// naturally aligned, so no bank conflict
s_local_val_max(warp_id, lane_id) = tmp_val_max;
s_local_exp_sum(warp_id, lane_id) = tmp_exp_sum;
s_local_product(warp_id, lane_id) = tmp_product;
__syncthreads();
/// NOTE: part 3: online softmax
/// NOTE: We have `kTileElements * kWarpThreads * kNumWarps` values to reduce
/// each reduce will consume `kNumWarps` threads (use partial warp reduction)
constexpr uint32_t kReductionCount = kTileElements * kWarpThreads * kNumWarps;
constexpr uint32_t kIteration = kReductionCount / kBlockSize;
PDLTriggerSecondary<kUsePDL>();
#pragma unroll
for (uint32_t i = 0; i < kIteration; ++i) {
/// NOTE: Range `[0, kTileElements * kWarpThreads * kNumWarps)`
const uint32_t j = i * kBlockSize + warp_id * kWarpThreads + lane_id;
/// NOTE: Range `[0, kNumWarps)`
const uint32_t local_warp_id = j % kNumWarps;
/// NOTE: Range `[0, kTileElements * kWarpThreads)`
const uint32_t local_elem_id = j / kNumWarps;
/// NOTE: Range `[0, kTileElements)`
const uint32_t local_tile_id = local_elem_id % kTileElements;
/// NOTE: Range `[0, kWarpThreads)`
const uint32_t local_lane_id = local_elem_id / kTileElements;
/// NOTE: each warp will access the whole tile (all `kTileElements`)
/// and for different lanes, the memory access only differ in `local_warp_id`
/// so there's no bank conflict in shared memory access.
static_assert(kTileElements * kNumWarps == kWarpThreads, "TODO: support other configs");
const auto local_val_max = s_local_val_max(local_warp_id, local_lane_id, local_tile_id);
const auto local_exp_sum = s_local_exp_sum(local_warp_id, local_lane_id, local_tile_id);
const auto local_product = s_local_product(local_warp_id, local_lane_id, local_tile_id);
const auto global_val_max = warp::reduce_max<kNumWarps>(local_val_max);
const auto rescale = expf(local_val_max - global_val_max);
const auto global_exp_sum = warp::reduce_sum<kNumWarps>(local_exp_sum * rescale);
const auto final_scale = rescale / global_exp_sum;
const auto global_product = warp::reduce_sum<kNumWarps>(local_product * final_scale);
kv_out[local_elem_id] = cast<OutFloat>(global_product);
}
}
template <typename Trait, typename BufferFloat, typename InputFloat>
SGL_DEVICE void c128_write_decode(BufferFloat* kv_buf, const InputFloat* kv_src) {
using namespace device;
using StorageInput = AlignedVector<InputFloat, kTileElements>;
const auto gmem_input = tile::Memory<StorageInput>::warp();
StorageInput data[2];
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
data[i] = gmem_input.load(kv_src + Trait::kHeadDim * i);
}
if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
gmem_input.store(kv_buf + Trait::kHeadDim * i, data[i]);
}
} else {
using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
StorageBuffer data_cast[2];
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
#pragma unroll
for (int32_t j = 0; j < kTileElements; ++j) {
data_cast[i][j] = cast<BufferFloat>(data[i][j]);
}
gmem_buffer.store(kv_buf + Trait::kHeadDim * i, data_cast[i]);
}
}
}
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodeParams params) {
using namespace device;
using Trait = C128Trait<kHeadDim>;
const uint32_t warp_id = threadIdx.x / kWarpThreads;
const uint32_t global_bid = blockIdx.x / Trait::kNumSplit; // batch id
const uint32_t global_sid = blockIdx.x % Trait::kNumSplit; // split id
const int64_t split_offset = global_sid * Trait::kTileDim;
if (global_bid >= params.batch_size) return;
const auto plan = params.plan_d[global_bid];
const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InputFloat*>(params.score_bias) + split_offset;
const auto kv_src = kv_input + global_bid * Trait::kElementSize;
const auto kv_out = kv_output + global_bid * Trait::kHeadDim;
const auto kv_buf = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
const auto kv_dst = kv_buffer + plan.write_loc * Trait::kElementSize;
PDLWaitPrimary<kUsePDL>();
// the write warp must match the load warp in the following `c128_forward`
if (warp_id == kNumWarps - 1) {
c128_write_decode<Trait, BufferFloat, InputFloat>(kv_dst, kv_src);
}
if (plan.write_loc % 128 == 127) {
c128_forward<Trait, kUsePDL, BufferFloat, InputFloat, OutFloat>(kv_buf, kv_src, kv_out, score_bias, 128);
}
}
// compress kernel
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
using namespace device;
using Trait = C128Trait<kHeadDim>;
const uint32_t global_pid = blockIdx.x / Trait::kNumSplit; // plan id
const uint32_t global_sid = blockIdx.x % Trait::kNumSplit; // split id
const int64_t split_offset = global_sid * Trait::kTileDim;
if (global_pid >= params.num_compress) return;
const auto plan = params.plan_c[global_pid];
const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InputFloat*>(params.score_bias) + split_offset;
if (plan.is_invalid()) return;
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
// Compact output: one row per compress plan, indexed by `global_pid`.
const auto kv_out = kv_output + global_pid * Trait::kHeadDim;
const auto kv_buf = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
PDLWaitPrimary<kUsePDL>();
c128_forward<Trait, kUsePDL, BufferFloat, InputFloat, OutFloat>(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len);
}
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
WRITE_KERNEL void write_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
using namespace device;
using Trait = C128Trait<kHeadDim>;
using StorageInput = AlignedVector<InputFloat, kTileElements>;
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
const uint32_t global_pid = global_wid / Trait::kNumSplit; // plan id
const uint32_t global_sid = global_wid % Trait::kNumSplit; // split id
// split the contiguous `kHeadDim * 2` into `kNumSplit` tiles
// each warp handles 1 contiguous tile (in contrast, decode handle the strided head_dim)
const int64_t split_offset = global_sid * (Trait::kTileDim * 2);
if (global_pid >= params.num_write) return;
const auto plan = params.plan_w[global_pid];
const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
if (plan.is_invalid()) return;
// each warp will handle a contiguous region
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
const auto kv_buf = kv_buffer + plan.write_loc * Trait::kElementSize;
const auto gmem_input = tile::Memory<StorageInput>::warp();
PDLWaitPrimary<kUsePDL>();
StorageInput data[2];
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
data[i] = gmem_input.load(kv_src, i);
}
if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
PDLTriggerSecondary<kUsePDL>();
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
gmem_input.store(kv_buf, data[i], i);
}
} else {
using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
StorageBuffer data_cast[2];
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
#pragma unroll
for (int32_t j = 0; j < kTileElements; ++j) {
data_cast[i][j] = cast<BufferFloat>(data[i][j]);
}
}
PDLTriggerSecondary<kUsePDL>();
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
gmem_buffer.store(kv_buf, data_cast[i], i);
}
}
}
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
struct FlashCompress128Kernel {
static constexpr auto decode_kernel = flash_c128_decode<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_c_kernel = flash_c128_prefill<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_w_kernel = write_c128_prefill<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
using Trait = C128Trait<kHeadDim>;
static void run_decode(
const tvm::ffi::TensorView kv_buffer,
const tvm::ffi::TensorView kv_input,
const tvm::ffi::TensorView kv_output,
const tvm::ffi::TensorView ape,
const tvm::ffi::TensorView plan_d_) {
using namespace host;
auto N = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLGPU>();
TensorMatcher({-1, 128, Trait::kElementSize}) // kv score
.with_dtype<BufferFloat>()
.with_device(device_)
.verify(kv_buffer);
TensorMatcher({N, Trait::kElementSize}) // kv score input
.with_dtype<InputFloat>()
.with_device(device_)
.verify(kv_input);
TensorMatcher({N, kHeadDim}) // kv compressed output
.with_dtype<OutFloat>()
.with_device(device_)
.verify(kv_output);
TensorMatcher({128, kHeadDim}) // ape
.with_dtype<InputFloat>()
.with_device(device_)
.verify(ape);
const auto plan_d = compress::verify_plan_d(plan_d_, N, device_);
const auto batch_size = static_cast<uint32_t>(N.unwrap());
const auto params = Compress128DecodeParams{
.kv_buffer = kv_buffer.data_ptr(),
.kv_input = kv_input.data_ptr(),
.kv_output = kv_output.data_ptr(),
.score_bias = ape.data_ptr(),
.plan_d = plan_d,
.batch_size = batch_size,
};
const uint32_t num_blocks = batch_size * kNumSplit;
LaunchKernel(num_blocks, kBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(decode_kernel, params);
}
static void run_prefill(
const tvm::ffi::TensorView kv_buffer,
const tvm::ffi::TensorView kv_input,
const tvm::ffi::TensorView kv_output,
const tvm::ffi::TensorView ape,
const tvm::ffi::TensorView plan_c_,
const tvm::ffi::TensorView plan_w_) {
using namespace host;
auto N = SymbolicSize{"num_q_tokens"};
auto C = SymbolicSize{"num_c_plans"};
auto W = SymbolicSize{"num_w_plans"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLGPU>();
TensorMatcher({-1, 128, Trait::kElementSize}) // kv score
.with_dtype<BufferFloat>()
.with_device(device_)
.verify(kv_buffer);
TensorMatcher({N, Trait::kElementSize}) // kv score input (ragged)
.with_dtype<InputFloat>()
.with_device(device_)
.verify(kv_input);
TensorMatcher({C, kHeadDim}) // kv compressed output (compact)
.with_dtype<OutFloat>()
.with_device(device_)
.verify(kv_output);
TensorMatcher({128, kHeadDim}) // ape
.with_dtype<InputFloat>()
.with_device(device_)
.verify(ape);
const auto plan_c = compress::verify_plan_c(plan_c_, C, device_);
const auto plan_w = compress::verify_plan_w(plan_w_, W, device_);
const auto device = device_.unwrap();
const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
const auto num_c = static_cast<uint32_t>(C.unwrap());
const auto num_w = static_cast<uint32_t>(W.unwrap());
const auto params = Compress128PrefillParams{
.kv_buffer = kv_buffer.data_ptr(),
.kv_input = kv_input.data_ptr(),
.kv_output = kv_output.data_ptr(),
.score_bias = ape.data_ptr(),
.plan_c = plan_c,
.plan_w = plan_w,
.num_compress = num_c,
.num_write = num_w,
};
RuntimeCheck(num_q_tokens >= num_w, "invalid prefill plan: num_q < num_w");
if (const auto num_c_blocks = num_c * kNumSplit) {
constexpr auto kBlockSize_C = kBlockSize;
LaunchKernel(num_c_blocks, kBlockSize_C, device) //
.enable_pdl(kUsePDL)(prefill_c_kernel, params);
}
constexpr uint32_t kWarpsPerWriteBlock = kWriteBlockSize / device::kWarpThreads;
if (const auto num_w_blocks = div_ceil(num_w * kNumSplit, kWarpsPerWriteBlock)) {
constexpr auto kBlockSize_W = kWriteBlockSize;
LaunchKernel(num_w_blocks, kBlockSize_W, device) //
.enable_pdl(kUsePDL)(prefill_w_kernel, params);
}
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/c4.cuh
+549
View File
@@ -0,0 +1,549 @@
#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 <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/compress.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/object.h>
#include <cstdint>
namespace {
using Plan4 = device::compress::PrefillPlan;
using IndiceT = int32_t;
/// \brief Each thread will handle this many elements (split along head_dim)
constexpr int kTileElements = 4;
/// \brief Need to improve register usage to reduce latency
#define C4_KERNEL __global__ __launch_bounds__(128, 4)
enum class PageMode {
RingBuffer = 8,
Page4Align = 4,
};
struct alignas(16) C4IndexBundle {
int32_t load_first_page;
int32_t load_second_page;
int32_t write_first_page;
int32_t last_position;
};
struct Compress4DecodeParams {
/**
* \brief Shape: `[num_indices, 8, head_dim * 4]` \n
* last dimension layout:
* | kv overlap | kv | score overlap | score |
*/
void* __restrict__ kv_score_buffer;
/** \brief Shape: `[batch_size, head_dim * 4]` */
const void* __restrict__ kv_score_input;
/** \brief Shape: `[batch_size, head_dim]` */
void* __restrict__ kv_compressed_output;
/** \brief Shape: `[8, head_dim]` (called `ape`) */
const void* __restrict__ score_bias;
/** \brief Shape: `[batch_size, ]`*/
const IndiceT* __restrict__ indices;
/** \brief Shape: `[batch_size, ]` */
const IndiceT* __restrict__ seq_lens;
/** \brief Shape: `[batch_size, 1]` */
const int32_t* __restrict__ extra;
/** \NOTE: `batch_size` <= `num_indices` */
uint32_t batch_size;
};
struct Compress4PrefillParams {
/**
* \brief Shape: `[num_indices, 8, head_dim * 4]` \n
* last dimension layout:
* | kv overlap | kv | score overlap | score |
*/
void* __restrict__ kv_score_buffer;
/** \brief Shape: `[num_q_tokens, head_dim * 4]` */
const void* __restrict__ kv_score_input;
/** \brief Shape: `[num_q_tokens, head_dim]` */
void* __restrict__ kv_compressed_output;
/** \brief Shape: `[8, head_dim]` (called `ape`) */
const void* __restrict__ score_bias;
/** \brief Shape: `[batch_size, ]`*/
const IndiceT* __restrict__ indices;
/** \brief Shape: `[batch_size, 4]` */
const C4IndexBundle* __restrict__ extra;
/** \brief The following part is plan info. */
const Plan4* __restrict__ compress_plan;
const Plan4* __restrict__ write_plan;
uint32_t num_compress;
uint32_t num_write;
};
template <typename T>
SGL_DEVICE void c4_write(
T* kv_score_buf, //
const T* kv_score_src,
const int64_t head_dim,
const int32_t write_pos) {
using namespace device;
using Storage = AlignedVector<T, kTileElements>;
const auto element_size = head_dim * 4;
const auto gmem = tile::Memory<Storage>::warp();
kv_score_buf += write_pos * element_size;
/// NOTE: Layout | [0] = kv overlap | [1] = kv | [2] = score overlap | [3] = score |
Storage kv_score[4];
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
kv_score[i] = gmem.load(kv_score_src + head_dim * i);
}
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
gmem.store(kv_score_buf + head_dim * i, kv_score[i]);
}
}
template <bool kPaged, typename InFloat, typename OutFloat>
SGL_DEVICE void c4_forward(
const InFloat* kv_score_buf,
const InFloat* kv_score_src,
OutFloat* kv_out,
const InFloat* score_bias,
const int64_t head_dim,
const int32_t seq_len,
const int32_t window_len,
[[maybe_unused]] const InFloat* kv_score_overlap_buf = nullptr) {
using namespace device;
const auto element_size = head_dim * 4;
const auto score_offset = head_dim * 2;
const auto overlap_stride = head_dim;
/// NOTE: part 1: load kv + score
using StorageIn = AlignedVector<InFloat, kTileElements>;
const auto gmem_in = tile::Memory<StorageIn>::warp();
StorageIn kv[8];
StorageIn score[8];
StorageIn bias[8];
#pragma unroll
for (int32_t i = 0; i < 8; ++i) {
bias[i] = gmem_in.load(score_bias + i * head_dim);
}
#pragma unroll
for (int32_t i = 0; i < 8; ++i) {
const bool is_overlap = i < 4;
const InFloat* src;
if (i < window_len) {
/// NOTE: `seq_len` must be a multiple of 4 here
if constexpr (kPaged) {
const auto kv_score_ptr = is_overlap ? kv_score_overlap_buf : kv_score_buf;
const int32_t k = i % 4;
src = kv_score_ptr + k * element_size;
} else {
const int32_t k = (seq_len + i) % 8;
src = kv_score_buf + k * element_size;
}
} else {
/// NOTE: k in [-7, 0]. We'll load from the ragged `kv_score_src`
const int32_t k = i - 7;
src = kv_score_src + k * element_size;
}
src += (is_overlap ? 0 : overlap_stride);
kv[i] = gmem_in.load(src);
score[i] = gmem_in.load(src + score_offset);
}
if (seq_len == 4) {
[[unlikely]];
constexpr float kFloatNegInf = -1e9f;
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
kv[i].fill(cast<InFloat>(0.0f));
score[i].fill(cast<InFloat>(kFloatNegInf));
}
}
/// NOTE: part 2: safe online softmax + weighted sum
using StorageOut = AlignedVector<OutFloat, kTileElements>;
const auto gmem_out = tile::Memory<StorageOut>::warp();
StorageOut result;
#pragma unroll
for (int32_t i = 0; i < kTileElements; ++i) {
float score_fp32[8];
#pragma unroll
for (int32_t j = 0; j < 8; ++j) {
score_fp32[j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
}
float max_value = score_fp32[0];
float sum_exp_value = 0.0f;
#pragma unroll
for (int32_t j = 1; j < 8; ++j) {
const auto fp32_score = score_fp32[j];
max_value = fmaxf(max_value, fp32_score);
}
float sum_product = 0.0f;
#pragma unroll
for (int32_t j = 0; j < 8; ++j) {
const auto fp32_score = score_fp32[j];
const auto exp_score = expf(fp32_score - max_value);
sum_product += cast<float>(kv[j][i]) * exp_score;
sum_exp_value += exp_score;
}
result[i] = cast<OutFloat>(sum_product / sum_exp_value);
}
gmem_out.store(kv_out, result);
}
template <int64_t kHeadDim, typename InFloat, typename OutFloat, PageMode kMode, bool kUsePDL>
C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams params) {
using namespace device;
constexpr int64_t kTileDim = kTileElements * kWarpThreads; // 128
constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
constexpr int64_t kElementSize = kHeadDim * 4; // `* 4` due to overlap transform + score
static_assert(kHeadDim % kTileDim == 0, "Head dim must be multiple of tile dim");
const auto& [
_kv_score_buffer, _kv_score_input, _kv_compressed_output, _score_bias, // kv score
indices, seq_lens, extra, batch_size // decode info
] = params;
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
const uint32_t global_bid = global_wid / kNumSplit; // batch id
const uint32_t global_sid = global_wid % kNumSplit; // split id
if (global_bid >= batch_size) return;
const int32_t index = indices[global_bid];
const int32_t seq_len = seq_lens[global_bid];
const int64_t split_offset = global_sid * kTileDim;
// kv score
const auto kv_score_buffer = static_cast<InFloat*>(_kv_score_buffer);
// kv input
const auto kv_score_input = static_cast<const InFloat*>(_kv_score_input);
const auto kv_src = kv_score_input + global_bid * kElementSize + split_offset;
// kv output
const auto kv_compressed_output = static_cast<OutFloat*>(_kv_compressed_output);
const auto kv_out = kv_compressed_output + global_bid * kHeadDim + split_offset;
// score bias (ape)
const auto score_bias = static_cast<const InFloat*>(_score_bias) + split_offset;
PDLWaitPrimary<kUsePDL>();
/// NOTE: `position` = `seq_len - 1`. To avoid underflow, we use `seq_len + page_size - 1`
if constexpr (kMode == PageMode::Page4Align) {
const auto index_prev = extra[global_bid];
const auto kv_buf = kv_score_buffer + index * (kElementSize * 4) + split_offset;
c4_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/(seq_len + 3) % 4);
if (seq_len % 4 == 0) {
const auto kv_overlap = kv_buf + (index_prev - index) * (kElementSize * 4);
c4_forward<true>(kv_buf, kv_src, kv_out, score_bias, kHeadDim, seq_len, 8, kv_overlap);
}
} else {
static_assert(kMode == PageMode::RingBuffer, "Unsupported PageMode");
const auto kv_buf = kv_score_buffer + index * (kElementSize * 8) + split_offset;
c4_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/(seq_len + 7) % 8);
if (seq_len % 4 == 0) {
c4_forward<false>(kv_buf, kv_src, kv_out, score_bias, kHeadDim, seq_len, /*window_size=*/8);
}
}
PDLTriggerSecondary<kUsePDL>();
}
template <int64_t kHeadDim, typename InFloat, typename OutFloat, PageMode kMode, bool kWrite, bool kUsePDL>
C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
using namespace device;
constexpr int64_t kTileDim = kTileElements * kWarpThreads; // 128
constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
constexpr int64_t kElementSize = kHeadDim * 4; // `* 4` due to overlap transform + score
static_assert(kHeadDim % kTileDim == 0, "Head dim must be multiple of tile dim");
const auto& [
_kv_score_buffer, _kv_score_input, _kv_compressed_output, _score_bias, // kv score
indices, extra, compress_plan, write_plan, num_compress, num_write // prefill plan
] = params;
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
const uint32_t global_pid = global_wid / kNumSplit; // plan id
const uint32_t global_sid = global_wid % kNumSplit; // split id
/// NOTE: compiler can optimize this if-else at compile time
const auto num_plans = kWrite ? num_write : num_compress;
const auto plan_ptr = kWrite ? write_plan : compress_plan;
if (global_pid >= num_plans) return;
const auto& [ragged_id, global_bid, position, window_len] = plan_ptr[global_pid];
const int64_t split_offset = global_sid * kTileDim;
// kv score
const auto kv_score_buffer = static_cast<InFloat*>(_kv_score_buffer);
// kv input
const auto kv_score_input = static_cast<const InFloat*>(_kv_score_input);
const auto kv_src = kv_score_input + ragged_id * kElementSize + split_offset;
// kv output
const auto kv_compressed_output = static_cast<OutFloat*>(_kv_compressed_output);
const auto kv_out = kv_compressed_output + ragged_id * kHeadDim + split_offset;
if (ragged_id == 0xFFFFFFFF) [[unlikely]]
return;
// score bias (ape)
const auto score_bias = static_cast<const InFloat*>(_score_bias) + split_offset;
const auto seq_len = position + 1;
const int32_t index = indices[global_bid];
PDLWaitPrimary<kUsePDL>();
if constexpr (kMode == PageMode::Page4Align) {
const auto write_second_page = index;
const auto [load_first_page, load_second_page, write_first_page, last_pos] = extra[global_bid];
if constexpr (kWrite) {
int32_t index;
if (position < static_cast<uint32_t>(last_pos)) {
index = write_first_page;
} else {
index = write_second_page;
}
const auto kv_buf = kv_score_buffer + index * (kElementSize * 4) + split_offset;
c4_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/position % 4);
} else {
int32_t index_overlap, index_normal;
if (window_len <= 4) {
index_overlap = load_second_page;
index_normal = load_second_page; // not used
} else {
index_overlap = load_first_page;
index_normal = load_second_page;
}
const auto kv_buf = kv_score_buffer + index_normal * (kElementSize * 4) + split_offset;
const auto kv_overlap = kv_score_buffer + index_overlap * (kElementSize * 4) + split_offset;
c4_forward<true>(kv_buf, kv_src, kv_out, score_bias, kHeadDim, seq_len, window_len, kv_overlap);
}
} else {
static_assert(kMode == PageMode::RingBuffer, "Unsupported PageMode");
const auto kv_buf = kv_score_buffer + index * (kElementSize * 8) + split_offset;
if constexpr (kWrite) {
c4_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/position % 8);
} else {
c4_forward<false>(kv_buf, kv_src, kv_out, score_bias, kHeadDim, seq_len, window_len);
}
}
PDLTriggerSecondary<kUsePDL>();
}
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
struct FlashCompress4Kernel {
template <PageMode kMode>
static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, InFloat, OutFloat, kMode, kUsePDL>;
template <PageMode kMode, bool kWrite>
static constexpr auto prefill_kernel = flash_c4_prefill<kHeadDim, InFloat, OutFloat, kMode, kWrite, kUsePDL>;
template <PageMode kMode>
static constexpr auto prefill_c_kernel = prefill_kernel<kMode, /*kWrite=*/false>;
template <PageMode kMode>
static constexpr auto prefill_w_kernel = prefill_kernel<kMode, /*kWrite=*/true>;
static constexpr uint32_t kBlockSize = 128;
static constexpr uint32_t kTileDim = kTileElements * device::kWarpThreads;
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
static constexpr uint32_t kWarpsPerBlock = kBlockSize / device::kWarpThreads;
using Self = FlashCompress4Kernel;
static void run_decode(
const tvm::ffi::TensorView kv_score_buffer,
const tvm::ffi::TensorView kv_score_input,
const tvm::ffi::TensorView kv_compressed_output,
const tvm::ffi::TensorView ape,
const tvm::ffi::TensorView indices,
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::Optional<tvm::ffi::TensorView> extra) {
using namespace host;
// this should not happen in practice
auto B = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
const auto extra_ptr = _get_extra_pointer(B, device_, extra);
const auto page_size = extra_ptr != nullptr ? 4 : 8;
TensorMatcher({-1, page_size, kHeadDim * 4}) // kv score
.with_dtype<InFloat>()
.with_device(device_)
.verify(kv_score_buffer);
TensorMatcher({B, kHeadDim * 4}) // kv score input
.with_dtype<InFloat>()
.with_device(device_)
.verify(kv_score_input);
TensorMatcher({B, kHeadDim}) // kv compressed output
.with_dtype<OutFloat>()
.with_device(device_)
.verify(kv_compressed_output);
TensorMatcher({8, kHeadDim}) // ape
.with_dtype<InFloat>()
.with_device(device_)
.verify(ape);
TensorMatcher({B}) // indices
.with_dtype<IndiceT>()
.with_device(device_)
.verify(indices);
TensorMatcher({B}) // seq lens
.with_dtype<IndiceT>()
.with_device(device_)
.verify(seq_lens);
const auto device = device_.unwrap();
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto params = Compress4DecodeParams{
.kv_score_buffer = kv_score_buffer.data_ptr(),
.kv_score_input = kv_score_input.data_ptr(),
.kv_compressed_output = kv_compressed_output.data_ptr(),
.score_bias = ape.data_ptr(),
.indices = static_cast<const IndiceT*>(indices.data_ptr()),
.seq_lens = static_cast<const IndiceT*>(seq_lens.data_ptr()),
.extra = static_cast<const int32_t*>(extra_ptr),
.batch_size = batch_size,
};
const auto kernel = extra_ptr != nullptr ? decode_kernel<PageMode::Page4Align> //
: decode_kernel<PageMode::RingBuffer>;
const uint32_t num_blocks = div_ceil(batch_size * kNumSplit, kWarpsPerBlock);
LaunchKernel(num_blocks, kBlockSize, device) //
.enable_pdl(kUsePDL)(kernel, params);
}
static void run_prefill(
const tvm::ffi::TensorView kv_score_buffer,
const tvm::ffi::TensorView kv_score_input,
const tvm::ffi::TensorView kv_compressed_output,
const tvm::ffi::TensorView ape,
const tvm::ffi::TensorView indices,
const tvm::ffi::TensorView compress_plan,
const tvm::ffi::TensorView write_plan,
const tvm::ffi::Optional<tvm::ffi::TensorView> extra) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto N = SymbolicSize{"num_q_tokens"};
auto X = SymbolicSize{"compress_tokens"};
auto Y = SymbolicSize{"write_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
const auto extra_ptr = _get_extra_pointer(B, device_, extra, /*is_prefill=*/true);
const auto page_size = extra_ptr != nullptr ? 4 : 8;
TensorMatcher({-1, page_size, kHeadDim * 4}) // kv score
.with_dtype<InFloat>()
.with_device(device_)
.verify(kv_score_buffer);
TensorMatcher({N, kHeadDim * 4}) // kv score input
.with_dtype<InFloat>()
.with_device(device_)
.verify(kv_score_input);
TensorMatcher({N, kHeadDim}) // kv compressed output
.with_dtype<OutFloat>()
.with_device(device_)
.verify(kv_compressed_output);
TensorMatcher({8, kHeadDim}) // ape
.with_dtype<InFloat>()
.with_device(device_)
.verify(ape);
TensorMatcher({B}) // indices
.with_dtype<IndiceT>()
.with_device(device_)
.verify(indices);
TensorMatcher({X, compress::kPrefillPlanDim}) // compress plan
.with_dtype<compress::PrefillPlanTensorDtype>()
.with_device(device_)
.verify(compress_plan);
TensorMatcher({Y, compress::kPrefillPlanDim}) // write plan
.with_dtype<compress::PrefillPlanTensorDtype>()
.with_device(device_)
.verify(write_plan);
const auto device = device_.unwrap();
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
const auto num_c = static_cast<uint32_t>(X.unwrap());
const auto num_w = static_cast<uint32_t>(Y.unwrap());
const auto params = Compress4PrefillParams{
.kv_score_buffer = kv_score_buffer.data_ptr(),
.kv_score_input = kv_score_input.data_ptr(),
.kv_compressed_output = kv_compressed_output.data_ptr(),
.score_bias = ape.data_ptr(),
.indices = static_cast<const IndiceT*>(indices.data_ptr()),
.extra = static_cast<const C4IndexBundle*>(extra_ptr),
.compress_plan = static_cast<const Plan4*>(compress_plan.data_ptr()),
.write_plan = static_cast<const Plan4*>(write_plan.data_ptr()),
.num_compress = num_c,
.num_write = num_w,
};
RuntimeCheck(num_q_tokens >= batch_size, "num_q_tokens must be >= batch_size");
RuntimeCheck(num_q_tokens >= std::max(num_c, num_w), "invalid prefill plan");
if (const auto num_c_blocks = div_ceil(num_c * kNumSplit, kWarpsPerBlock)) {
const auto c_kernel = extra_ptr != nullptr ? prefill_c_kernel<PageMode::Page4Align> //
: prefill_c_kernel<PageMode::RingBuffer>;
LaunchKernel(num_c_blocks, kBlockSize, device) //
.enable_pdl(kUsePDL)(c_kernel, params);
}
if (const auto num_w_blocks = div_ceil(num_w * kNumSplit, kWarpsPerBlock)) {
const auto w_kernel = extra_ptr != nullptr ? prefill_w_kernel<PageMode::Page4Align> //
: prefill_w_kernel<PageMode::RingBuffer>;
LaunchKernel(num_w_blocks, kBlockSize, device) //
.enable_pdl(kUsePDL)(w_kernel, params);
}
}
// some auxiliary functions
private:
static const void* _get_extra_pointer(
host::SymbolicSize& B, // batch_size
host::SymbolicDevice& device,
const tvm::ffi::Optional<tvm::ffi::TensorView>& extra,
bool is_prefill = false) {
// only have value when using page-aligned mode
if (!extra.has_value()) return nullptr;
const auto& extra_tensor = extra.value();
/// NOTE: the metadata layout is different for prefill and decode:
/// for prefill, last 4 are:
/// load overlap | load normal | write overlap | last written page
/// for decode, last 1 is the write (also load) overlap
host::TensorMatcher({B, is_prefill ? 4 : 1}) // extra tensor
.with_dtype<int32_t>()
.with_device(device)
.verify(extra_tensor);
const auto data_ptr = extra_tensor.data_ptr();
host::RuntimeCheck(data_ptr != nullptr, "extra tensor data ptr is null");
if (is_prefill) {
static_assert(alignof(C4IndexBundle) == 16);
host::RuntimeCheck(std::bit_cast<uintptr_t>(data_ptr) % 16 == 0, "extra tensor is not properly aligned");
}
return data_ptr;
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh
+491
View File
@@ -0,0 +1,491 @@
/**
* \brief Here's some dimension info for the main buffer used in C4 prefill and decode.
*
* kv_buffer: [num_indices, 8, head_dim * 4]
* - last dimension layout: | kv overlap | kv | score overlap | score |
* kv_input: [batch_size, head_dim * 4]
* kv_output: [batch_size, head_dim]
* score_bias (ape): [8, head_dim]
* plan_c/plan_w: [variable length]
*
* For prefill, batch_size = num_q_tokens
*/
#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 <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/compress_v2.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/object.h>
#include <cfloat>
#include <cstdint>
#include <type_traits>
namespace {
using PlanD = device::compress::DecodePlan;
using PlanC = device::compress::CompressPlan;
using PlanW = device::compress::WritePlan;
/// \brief Each thread will handle this many elements (split along head_dim)
constexpr int32_t kTileElements = 4;
/// \brief Need to improve register usage to reduce latency
#define C4_KERNEL __global__ __launch_bounds__(128, 4)
#define WRITE_KERNEL __global__ __launch_bounds__(128, 16)
struct Compress4DecodeParams {
void* __restrict__ kv_buffer;
const void* __restrict__ kv_input;
void* __restrict__ kv_output;
const void* __restrict__ score_bias;
const PlanD* __restrict__ plan_d;
uint32_t batch_size;
};
struct Compress4PrefillParams {
void* __restrict__ kv_buffer;
const void* __restrict__ kv_input;
void* __restrict__ kv_output;
const void* __restrict__ score_bias;
const PlanC* __restrict__ plan_c;
const PlanW* __restrict__ plan_w;
uint32_t num_compress;
uint32_t num_write;
};
template <int64_t kHeadDim_>
struct C4Trait {
static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 128
static constexpr int64_t kHeadDim = kHeadDim_;
static constexpr int64_t kOverlapOffset = kHeadDim;
static constexpr int64_t kScoreOffset = kHeadDim * 2;
static constexpr int64_t kElementSize = kHeadDim * 4;
static constexpr int64_t kPageElementSize = 4 * kElementSize; // page size = 4
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
static_assert(kHeadDim % kTileDim == 0);
};
template <typename Trait, bool kUsePDL, typename BufferFloat, typename InputFloat, typename OutFloat>
SGL_DEVICE void c4_forward(
const BufferFloat* kv_buf_0, // overlap [4n - 4, 4n - 1]
const BufferFloat* kv_buf_1, // normal [4n + 0, 4n + 3]
const InputFloat* kv_src, // ragged pointer at position = 4n + 3
OutFloat* kv_out,
const InputFloat* score_bias,
const bool should_overlap,
const int32_t buffer_len) {
using namespace device;
using StorageIn = AlignedVector<InputFloat, kTileElements>;
const auto gmem_in = tile::Memory<StorageIn>::warp();
StorageIn kv[8];
StorageIn score[8];
StorageIn bias[8];
#pragma unroll
for (int32_t i = 0; i < 8; ++i) {
bias[i] = gmem_in.load(score_bias + i * Trait::kHeadDim);
}
if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
if (should_overlap) {
const auto kv_start = kv_src - 7 * Trait::kElementSize; // point to start
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
const auto src = i < buffer_len ? kv_buf_0 : kv_start;
const auto base = src + i * Trait::kElementSize;
kv[i] = gmem_in.load(base);
score[i] = gmem_in.load(base + Trait::kScoreOffset);
}
} else {
[[unlikely]];
constexpr float kFloatNegInf = -FLT_MAX;
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
kv[i].fill(cast<InputFloat>(0.0f));
score[i].fill(cast<InputFloat>(kFloatNegInf));
}
}
const auto kv_start = kv_src - 3 * Trait::kElementSize; // point to start
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
const auto src = i + 4 < buffer_len ? kv_buf_1 : kv_start;
const auto base = src + i * Trait::kElementSize + Trait::kOverlapOffset;
kv[i + 4] = gmem_in.load(base);
score[i + 4] = gmem_in.load(base + Trait::kScoreOffset);
}
} else { // mixed dtype
using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
const auto kv_start_0 = kv_src - 7 * Trait::kElementSize; // point to start
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
if (should_overlap && i < buffer_len) {
const auto base = kv_buf_0 + i * Trait::kElementSize;
const auto kv_tmp = gmem_buffer.load(base);
const auto score_tmp = gmem_buffer.load(base + Trait::kScoreOffset);
#pragma unroll
for (int32_t j = 0; j < kTileElements; ++j) {
kv[i][j] = cast<InputFloat>(kv_tmp[j]);
score[i][j] = cast<InputFloat>(score_tmp[j]);
}
} else if (should_overlap) {
const auto base = kv_start_0 + i * Trait::kElementSize;
kv[i] = gmem_in.load(base);
score[i] = gmem_in.load(base + Trait::kScoreOffset);
} else {
[[unlikely]];
constexpr float kFloatNegInf = -FLT_MAX;
kv[i].fill(cast<InputFloat>(0.0f));
score[i].fill(cast<InputFloat>(kFloatNegInf));
}
}
const auto kv_start = kv_src - 3 * Trait::kElementSize; // point to start
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
if (i + 4 < buffer_len) {
const auto base = kv_buf_1 + i * Trait::kElementSize + Trait::kOverlapOffset;
const auto kv_tmp = gmem_buffer.load(base);
const auto score_tmp = gmem_buffer.load(base + Trait::kScoreOffset);
#pragma unroll
for (int32_t j = 0; j < kTileElements; ++j) {
kv[i + 4][j] = cast<InputFloat>(kv_tmp[j]);
score[i + 4][j] = cast<InputFloat>(score_tmp[j]);
}
} else {
const auto base = kv_start + i * Trait::kElementSize + Trait::kOverlapOffset;
kv[i + 4] = gmem_in.load(base);
score[i + 4] = gmem_in.load(base + Trait::kScoreOffset);
}
}
}
/// NOTE: part 2: safe online softmax + weighted sum
using StorageOut = AlignedVector<OutFloat, kTileElements>;
const auto gmem_out = tile::Memory<StorageOut>::warp();
StorageOut result;
// consume 32 fp registers
float score_fp32[kTileElements][8];
// convert to fp32 and apply bias first
#pragma unroll
for (int32_t i = 0; i < kTileElements; ++i) {
#pragma unroll
for (int32_t j = 0; j < 8; ++j) {
score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
}
}
#pragma unroll
for (int32_t i = 0; i < kTileElements; ++i) {
const auto& score = score_fp32[i];
float max_value = score[0];
float sum_exp_value = 0.0f;
#pragma unroll
for (int32_t j = 1; j < 8; ++j) {
const auto fp32_score = score[j];
max_value = fmaxf(max_value, fp32_score);
}
float sum_product = 0.0f;
#pragma unroll
for (int32_t j = 0; j < 8; ++j) {
const auto fp32_score = score[j];
const auto exp_score = expf(fp32_score - max_value);
sum_product += cast<float>(kv[j][i]) * exp_score;
sum_exp_value += exp_score;
}
result[i] = cast<OutFloat>(sum_product / sum_exp_value);
}
// overlap the store with the next iteration's load
PDLTriggerSecondary<kUsePDL>();
gmem_out.store(kv_out, result);
}
template <typename Trait, typename BufferFloat, typename InputFloat>
SGL_DEVICE void c4_write_decode(BufferFloat* kv_buf, const InputFloat* kv_src) {
using namespace device;
using StorageInput = AlignedVector<InputFloat, kTileElements>;
const auto gmem_input = tile::Memory<StorageInput>::warp();
StorageInput data[4];
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
data[i] = gmem_input.load(kv_src + Trait::kHeadDim * i);
}
if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
gmem_input.store(kv_buf + Trait::kHeadDim * i, data[i]);
}
} else {
using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
StorageBuffer data_cast[4];
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
#pragma unroll
for (int32_t j = 0; j < kTileElements; ++j) {
data_cast[i][j] = cast<BufferFloat>(data[i][j]);
}
gmem_buffer.store(kv_buf + Trait::kHeadDim * i, data_cast[i]);
}
}
}
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams params) {
using namespace device;
using Trait = C4Trait<kHeadDim>;
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
const uint32_t global_bid = global_wid / Trait::kNumSplit; // batch id
const uint32_t global_sid = global_wid % Trait::kNumSplit; // split id
const int64_t split_offset = global_sid * Trait::kTileDim;
if (global_bid >= params.batch_size) return;
const auto plan = params.plan_d[global_bid];
const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InputFloat*>(params.score_bias) + split_offset;
const auto kv_src = kv_input + global_bid * Trait::kElementSize;
const auto kv_out = kv_output + global_bid * Trait::kHeadDim;
const auto kv_buf_0 = kv_buffer + plan.read_page_0 * Trait::kPageElementSize;
const auto kv_buf_1 = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
const auto kv_dst = kv_buffer + plan.write_loc * Trait::kElementSize;
PDLWaitPrimary<kUsePDL>();
c4_write_decode<Trait, BufferFloat, InputFloat>(kv_dst, kv_src);
if (plan.seq_len % 4 == 0) {
const auto need_overlap = plan.seq_len > 4;
c4_forward<Trait, kUsePDL, BufferFloat, InputFloat, OutFloat>(
kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8);
}
}
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
using namespace device;
using Trait = C4Trait<kHeadDim>;
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
const uint32_t global_pid = global_wid / Trait::kNumSplit; // plan id
const uint32_t global_sid = global_wid % Trait::kNumSplit; // split id
const int64_t split_offset = global_sid * Trait::kTileDim;
if (global_pid >= params.num_compress) return;
const auto plan = params.plan_c[global_pid];
const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InputFloat*>(params.score_bias) + split_offset;
if (plan.is_invalid()) return;
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
// Compact output: one row per compress plan, indexed by `global_pid`.
const auto kv_out = kv_output + global_pid * Trait::kHeadDim;
const auto kv_buf_0 = kv_buffer + plan.read_page_0 * Trait::kPageElementSize;
const auto kv_buf_1 = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
const bool need_overlap = plan.seq_len > 4;
PDLWaitPrimary<kUsePDL>();
c4_forward<Trait, kUsePDL, BufferFloat, InputFloat, OutFloat>(
kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len);
}
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
WRITE_KERNEL void write_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
using namespace device;
using Trait = C4Trait<kHeadDim>;
using StorageInput = AlignedVector<InputFloat, kTileElements>;
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
const uint32_t global_pid = global_wid / Trait::kNumSplit; // plan id
const uint32_t global_sid = global_wid % Trait::kNumSplit; // split id
// split the contiguous `kHeadDim * 4` into `kNumSplit` tiles
// each warp handles 1 contiguous tile (in contrast, decode handle the strided head_dim)
const int64_t split_offset = global_sid * (Trait::kTileDim * 4);
if (global_pid >= params.num_write) return;
const auto plan = params.plan_w[global_pid];
const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
if (plan.is_invalid()) return;
// each warp will handle a contiguous region
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
const auto kv_buf = kv_buffer + plan.write_loc * Trait::kElementSize;
const auto gmem_input = tile::Memory<StorageInput>::warp();
PDLWaitPrimary<kUsePDL>();
StorageInput data[4];
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
data[i] = gmem_input.load(kv_src, i);
}
if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
PDLTriggerSecondary<kUsePDL>();
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
gmem_input.store(kv_buf, data[i], i);
}
} else {
using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
StorageBuffer data_cast[4];
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
#pragma unroll
for (int32_t j = 0; j < kTileElements; ++j) {
data_cast[i][j] = cast<BufferFloat>(data[i][j]);
}
}
PDLTriggerSecondary<kUsePDL>();
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
gmem_buffer.store(kv_buf, data_cast[i], i);
}
}
}
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
struct FlashCompress4Kernel {
static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_c_kernel = flash_c4_prefill<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_w_kernel = write_c4_prefill<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
static constexpr uint32_t kBlockSize = 128;
static constexpr uint32_t kTileDim = kTileElements * device::kWarpThreads;
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
static constexpr uint32_t kWarpsPerBlock = kBlockSize / device::kWarpThreads;
using Trait = C4Trait<kHeadDim>;
static void run_decode(
const tvm::ffi::TensorView kv_buffer,
const tvm::ffi::TensorView kv_input,
const tvm::ffi::TensorView kv_output,
const tvm::ffi::TensorView ape,
const tvm::ffi::TensorView plan_d_) {
using namespace host;
auto N = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLGPU>();
TensorMatcher({-1, 4, Trait::kElementSize}) // kv score
.with_dtype<BufferFloat>()
.with_device(device_)
.verify(kv_buffer);
TensorMatcher({N, Trait::kElementSize}) // kv score input
.with_dtype<InputFloat>()
.with_device(device_)
.verify(kv_input);
TensorMatcher({N, kHeadDim}) // kv compressed output
.with_dtype<OutFloat>()
.with_device(device_)
.verify(kv_output);
TensorMatcher({8, kHeadDim}) // ape
.with_dtype<InputFloat>()
.with_device(device_)
.verify(ape);
const auto plan_d = compress::verify_plan_d(plan_d_, N, device_);
const auto batch_size = static_cast<uint32_t>(N.unwrap());
const auto params = Compress4DecodeParams{
.kv_buffer = kv_buffer.data_ptr(),
.kv_input = kv_input.data_ptr(),
.kv_output = kv_output.data_ptr(),
.score_bias = ape.data_ptr(),
.plan_d = plan_d,
.batch_size = batch_size,
};
const uint32_t num_blocks = div_ceil(batch_size * kNumSplit, kWarpsPerBlock);
LaunchKernel(num_blocks, kBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(decode_kernel, params);
}
static void run_prefill(
const tvm::ffi::TensorView kv_buffer,
const tvm::ffi::TensorView kv_input,
const tvm::ffi::TensorView kv_output,
const tvm::ffi::TensorView ape,
const tvm::ffi::TensorView plan_c_,
const tvm::ffi::TensorView plan_w_) {
using namespace host;
auto N = SymbolicSize{"num_q_tokens"};
auto C = SymbolicSize{"num_c_plans"};
auto W = SymbolicSize{"num_w_plans"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLGPU>();
TensorMatcher({-1, 4, Trait::kElementSize}) // kv score
.with_dtype<BufferFloat>()
.with_device(device_)
.verify(kv_buffer);
TensorMatcher({N, Trait::kElementSize}) // kv score input (ragged)
.with_dtype<InputFloat>()
.with_device(device_)
.verify(kv_input);
TensorMatcher({C, kHeadDim}) // kv compressed output (compact)
.with_dtype<OutFloat>()
.with_device(device_)
.verify(kv_output);
TensorMatcher({8, kHeadDim}) // ape
.with_dtype<InputFloat>()
.with_device(device_)
.verify(ape);
const auto plan_c = compress::verify_plan_c(plan_c_, C, device_);
const auto plan_w = compress::verify_plan_w(plan_w_, W, device_);
const auto device = device_.unwrap();
const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
const auto num_c = static_cast<uint32_t>(C.unwrap());
const auto num_w = static_cast<uint32_t>(W.unwrap());
const auto params = Compress4PrefillParams{
.kv_buffer = kv_buffer.data_ptr(),
.kv_input = kv_input.data_ptr(),
.kv_output = kv_output.data_ptr(),
.score_bias = ape.data_ptr(),
.plan_c = plan_c,
.plan_w = plan_w,
.num_compress = num_c,
.num_write = num_w,
};
RuntimeCheck(num_q_tokens >= num_w, "invalid prefill plan: num_q < num_w");
if (const auto num_c_blocks = div_ceil(num_c * kNumSplit, kWarpsPerBlock)) {
LaunchKernel(num_c_blocks, kBlockSize, device) //
.enable_pdl(kUsePDL)(prefill_c_kernel, params);
}
if (const auto num_w_blocks = div_ceil(num_w * kNumSplit, kWarpsPerBlock)) {
LaunchKernel(num_w_blocks, kBlockSize, device) //
.enable_pdl(kUsePDL)(prefill_w_kernel, params);
}
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh
+868
View File
@@ -0,0 +1,868 @@
#include <sgl_kernel/ffi.h>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/compress_v2.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tuple.h>
#include <cstdint>
#include <limits>
namespace host::compress {
constexpr auto kDLUInt8 = DLDataType{.code = kDLUInt, .bits = 8, .lanes = 1};
using PlanC = CompressPlan;
using PlanW = WritePlan;
using PlanD = DecodePlan;
using RID_T = int64_t;
using R2T_T = int32_t;
using F2S_T = int64_t;
using IDX_T = int64_t;
/// NOTE: for the internal use, we pack the ragged and batch id, since both not exceed 65536
SGL_DEVICE __host__ PlanW pack_w(uint32_t ragged_id, uint32_t batch_id, int32_t seq_len) {
return {static_cast<uint32_t>(ragged_id | batch_id << 16), seq_len};
}
/// NOTE: for the internal use, we pack the ragged and batch id, since both not exceed 65536
SGL_DEVICE uint2 unpack_w(PlanW plan) {
return {static_cast<uint16_t>(plan.ragged_id), static_cast<uint16_t>(plan.ragged_id >> 16)};
}
struct Prefill0Params {
PlanC* plan_c;
PlanW* plan_w;
const IDX_T* seq_lens_ptr; // [batch_size]
const IDX_T* extend_lens_ptr; // [batch_size]
uint32_t batch_size;
uint32_t num_q_tokens;
int32_t compress_ratio;
int32_t swa_page_size;
int32_t mtp_pad;
};
struct Prefill1Params {
PlanC* plan_c;
PlanW* plan_w;
const RID_T* rid_ptr; // [batch_size]
const R2T_T* r2t_ptr; // [num_reqs, stride_r2t]
const F2S_T* f2s_ptr; // [num_full_slots], full_loc -> swa_loc
int64_t stride_r2t;
uint32_t num_c;
uint32_t num_w;
uint32_t num_c_padded;
uint32_t num_w_padded;
uint32_t num_work;
int32_t swa_page_size;
int32_t ring_size;
int32_t compress_ratio;
};
struct DecodeParams {
PlanD* plan_d;
const RID_T* rid_ptr; // [batch_size]
const R2T_T* r2t_ptr; // [num_reqs, stride_r2t]
const F2S_T* f2s_ptr; // [num_full_slots], full_loc -> swa_loc
const IDX_T* seq_ptr; // [batch_size]
int64_t stride_r2t;
uint32_t batch_size;
int32_t swa_page_size;
int32_t ring_size;
int32_t compress_ratio;
};
struct Prefill1ParamsLegacy {
PlanC* plan_c;
PlanW* plan_w;
const RID_T* rid_ptr; // [batch_size]
uint32_t num_c;
uint32_t num_w;
uint32_t num_c_padded;
uint32_t num_w_padded;
uint32_t num_work;
int32_t compress_ratio;
};
struct DecodeParamsLegacy {
PlanD* plan_d;
const RID_T* rid_ptr; // [batch_size]
const IDX_T* seq_ptr; // [batch_size]
uint32_t batch_size;
int32_t compress_ratio;
};
inline constexpr uint32_t kMaxPrefillBatchSize = 1024;
SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) {
static_assert(device::kWarpThreads == 32);
#pragma unroll
for (uint32_t offset = 1; offset < 32; offset *= 2) {
#ifndef USE_ROCM
uint32_t n = __shfl_up_sync(device::kFullMask, val, offset);
#else
uint32_t n = __shfl_up(val, offset, 32);
#endif
if (lane_id >= offset) val += n;
}
return val;
}
/// Warp-wide max/min for integer types. `device::warp::reduce_max` routes through
/// `dtype_trait<T>::max` which is only specialized for FP types.
SGL_DEVICE uint32_t warp_reduce_max_u32(uint32_t val) {
#pragma unroll
for (uint32_t mask = 16; mask > 0; mask >>= 1) {
#ifndef USE_ROCM
val = max(val, __shfl_xor_sync(device::kFullMask, val, mask, 32));
#else
val = max(val, __shfl_xor(val, mask, 32));
#endif
}
return val;
}
SGL_DEVICE uint32_t warp_reduce_min_u32(uint32_t val) {
#pragma unroll
for (uint32_t mask = 16; mask > 0; mask >>= 1) {
#ifndef USE_ROCM
val = min(val, __shfl_xor_sync(device::kFullMask, val, mask, 32));
#else
val = min(val, __shfl_xor(val, mask, 32));
#endif
}
return val;
}
__global__ __launch_bounds__(1024, 1) //
void plan_compress_prefill_kernel0(const Prefill0Params params) {
using namespace device;
const auto tx = threadIdx.x;
const auto block_size = kMaxPrefillBatchSize;
constexpr auto kNumWarps = kMaxPrefillBatchSize / kWarpThreads;
const auto cr = params.compress_ratio;
const auto sps = params.swa_page_size;
const bool is_overlap = (cr == 4);
const int32_t window_size = cr * (is_overlap ? 2 : 1);
alignas(128) __shared__ uint32_t counter_c;
alignas(128) __shared__ uint32_t counter_w;
__shared__ int32_t s_seq_len[kMaxPrefillBatchSize];
__shared__ int32_t s_prefix_len[kMaxPrefillBatchSize];
__shared__ uint32_t warp_max[kNumWarps];
__shared__ uint32_t warp_min[kNumWarps];
__shared__ uint32_t s_max_extend;
__shared__ uint32_t s_min_extend;
const auto lane_id = tx % kWarpThreads;
const auto warp_id = tx / kWarpThreads;
// === Stage A: load per-batch fields, init shared scratch ===
int32_t seq_len = 0, extend_len = 0, prefix_len = 0;
if (tx < params.batch_size) {
seq_len = static_cast<int32_t>(params.seq_lens_ptr[tx]);
extend_len = static_cast<int32_t>(params.extend_lens_ptr[tx]);
prefix_len = seq_len - extend_len;
s_seq_len[tx] = seq_len;
s_prefix_len[tx] = prefix_len;
}
if (tx == 0) {
counter_c = 0;
counter_w = 0;
}
if (tx < kNumWarps) {
warp_max[tx] = 0;
warp_min[tx] = 0xFFFFFFFFu;
}
// === Stage B: min/max(extend_len) for MTP-uniform detection ===
// For min, treat threads outside `batch_size` as +inf so they don't pull the min down.
const uint32_t e_for_max = static_cast<uint32_t>(extend_len);
const uint32_t e_for_min = (tx < params.batch_size) ? e_for_max : 0xFFFFFFFFu;
warp_max[warp_id] = warp_reduce_max_u32(e_for_max);
warp_min[warp_id] = warp_reduce_min_u32(e_for_min);
__syncthreads();
if (warp_id == 0) {
s_max_extend = warp_reduce_max_u32(warp_max[lane_id]);
s_min_extend = warp_reduce_min_u32(warp_min[lane_id]);
}
__syncthreads();
const auto num_q = params.num_q_tokens;
// MTP-uniform: every batch shares the same small extend_len `E`, so we can decompose
// a global token id `k` into (batch_id, j) = (k / E, k % E) and skip the per-batch loop.
const bool is_mtp_extend = (s_min_extend == s_max_extend) && (s_max_extend > 0) && (s_max_extend <= 32);
// === Stage C: emit valid plans, slot allocation via shared-mem atomicAdd ===
if (is_mtp_extend) {
// Path 1: token-driven. Each global token id maps to exactly one (batch_id, j).
const uint32_t E = s_max_extend;
// num_q is the padded buffer size (graph bucket), not the work size: cap the
// loop at the real token count so batch_id = k / E stays < batch_size on an
// underfilled replay; Stage D pads [counter, num_q) with invalid.
const uint32_t num_real_q = params.batch_size * E;
for (uint32_t k = tx; k < num_real_q; k += block_size) {
const uint32_t batch_id = k / E;
const uint32_t j = k % E;
const int32_t pl = s_prefix_len[batch_id];
const int32_t sl = s_seq_len[batch_id];
const int32_t position = pl + static_cast<int32_t>(j);
const uint32_t ragged_id = k;
if ((position + 1) % cr == 0) {
const int32_t buffer_len = window_size - min(static_cast<int32_t>(j) + 1, window_size);
const uint32_t out_idx = atomicAdd(&counter_c, 1u);
params.plan_c[out_idx] = {
.seq_len = static_cast<uint32_t>(position + 1),
.ragged_id = static_cast<uint16_t>(ragged_id),
.buffer_len = static_cast<uint16_t>(buffer_len),
.read_page_0 = -1,
.read_page_1 = static_cast<int32_t>(batch_id),
};
}
const int32_t last_c_pos = (sl / cr) * cr;
const int32_t first_w_pos = min(last_c_pos - (is_overlap ? cr : 0), sl - params.mtp_pad);
bool do_write = position >= first_w_pos;
if (!do_write && is_overlap) do_write = (position % sps) >= (sps - cr);
if (do_write) {
const uint32_t out_idx = atomicAdd(&counter_w, 1u);
params.plan_w[out_idx] = pack_w(ragged_id, batch_id, position + 1);
}
}
} else {
// Path 2: general prefill (long extend_len). Iterate batches in an outer loop;
// the whole block sweeps each batch's tokens in parallel.
uint32_t base_e = 0;
for (uint32_t batch_id = 0; batch_id < params.batch_size; ++batch_id) {
const int32_t pl = s_prefix_len[batch_id];
const int32_t sl = s_seq_len[batch_id];
const int32_t el = sl - pl;
const int32_t last_c_pos = (sl / cr) * cr;
const int32_t first_w_pos = min(last_c_pos - (is_overlap ? cr : 0), sl - params.mtp_pad);
for (int32_t j = static_cast<int32_t>(tx); j < el; j += static_cast<int32_t>(block_size)) {
const int32_t position = pl + j;
const uint32_t ragged_id = base_e + static_cast<uint32_t>(j);
if ((position + 1) % cr == 0) {
const int32_t buffer_len = window_size - min(j + 1, window_size);
const uint32_t out_idx = atomicAdd(&counter_c, 1u);
params.plan_c[out_idx] = {
.seq_len = static_cast<uint32_t>(position + 1),
.ragged_id = static_cast<uint16_t>(ragged_id),
.buffer_len = static_cast<uint16_t>(buffer_len),
.read_page_0 = -1,
.read_page_1 = static_cast<int32_t>(batch_id),
};
}
bool do_write = position >= first_w_pos;
if (!do_write && is_overlap) do_write = (position % sps) >= (sps - cr);
if (do_write) {
const uint32_t out_idx = atomicAdd(&counter_w, 1u);
params.plan_w[out_idx] = pack_w(ragged_id, static_cast<uint32_t>(batch_id), position + 1);
}
}
base_e += static_cast<uint32_t>(el);
}
}
__syncthreads();
// === Stage D: pad [counter_c, num_q) / [counter_w, num_q) with invalid ===
const auto total_c = counter_c;
const auto total_w = counter_w;
for (uint32_t k = total_c + tx; k < num_q; k += block_size) {
params.plan_c[k] = PlanC::invalid();
}
for (uint32_t k = total_w + tx; k < num_q; k += block_size) {
params.plan_w[k] = PlanW::invalid();
}
}
/// NOTE: stage 1
__global__ void plan_compress_prefill_kernel_1(const Prefill1Params params) {
const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= params.num_work) return;
auto plan_c = idx < params.num_c ? params.plan_c[idx] : PlanC::invalid();
auto plan_w = idx < params.num_w ? params.plan_w[idx] : PlanW::invalid();
const auto compute_loc = [&](int32_t swa_loc) {
const auto swa_page = swa_loc / params.swa_page_size;
const auto ring_offset = swa_loc % params.ring_size;
return swa_page * params.ring_size + ring_offset;
};
const auto compute_c128_loc = [&](int64_t rid, int32_t position) {
return static_cast<int32_t>(rid * params.ring_size + position % params.ring_size);
};
if (!plan_c.is_invalid()) { // 1. in bound. 2. not masked
if (plan_c.buffer_len > 0) {
const auto batch_id = plan_c.read_page_1;
const auto rid = params.rid_ptr[batch_id];
const auto mapping = params.r2t_ptr + rid * params.stride_r2t;
// `seq_len` should be ratio-aligned here
const auto position_1 = static_cast<int32_t>(plan_c.seq_len - 1);
// only used for c4, harmless for c128
const auto position_0 = max(position_1 - params.compress_ratio, 0);
if (params.compress_ratio == 128) {
plan_c.read_page_0 = compute_c128_loc(rid, position_0) / 128;
plan_c.read_page_1 = compute_c128_loc(rid, position_1) / 128;
} else {
const auto raw_loc_0 = mapping[position_0];
const auto raw_loc_1 = mapping[position_1];
const auto state_loc_0 = params.f2s_ptr[raw_loc_0];
const auto state_loc_1 = params.f2s_ptr[raw_loc_1];
plan_c.read_page_0 = compute_loc(state_loc_0) / params.compress_ratio;
plan_c.read_page_1 = compute_loc(state_loc_1) / params.compress_ratio;
}
params.plan_c[idx] = plan_c;
}
} else if (idx < params.num_c_padded) {
params.plan_c[idx] = PlanC::invalid();
}
if (!plan_w.is_invalid()) { // 1. in bound. 2. not masked
const auto [ragged_id, batch_id] = unpack_w(plan_w);
const auto rid = params.rid_ptr[batch_id];
const auto mapping = params.r2t_ptr + rid * params.stride_r2t;
// `seq_len` (`write_loc`) may not be aligned here
const auto position = static_cast<int32_t>(plan_w.write_loc - 1);
plan_w.ragged_id = ragged_id;
if (params.compress_ratio == 128) {
plan_w.write_loc = compute_c128_loc(rid, position);
} else {
const auto raw_loc = mapping[position];
plan_w.write_loc = compute_loc(params.f2s_ptr[raw_loc]);
}
params.plan_w[idx] = plan_w;
} else if (idx < params.num_w_padded) {
params.plan_w[idx] = PlanW::invalid();
}
}
__global__ void plan_compress_decode_kernel(const DecodeParams params) {
const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= params.batch_size) return;
const auto rid = params.rid_ptr[idx];
const auto mapping = params.r2t_ptr + rid * params.stride_r2t;
const auto compute_loc = [&](int32_t swa_loc) {
const auto swa_page = swa_loc / params.swa_page_size;
const auto ring_offset = swa_loc % params.ring_size;
return swa_page * params.ring_size + ring_offset;
};
const auto compute_c128_loc = [&](int64_t rid, int32_t position) {
return static_cast<int32_t>(rid * params.ring_size + position % params.ring_size);
};
const auto seq_len = static_cast<int32_t>(params.seq_ptr[idx]);
const auto position_1 = static_cast<int32_t>(seq_len - 1);
const auto position_0 = max(position_1 - params.compress_ratio, 0);
int32_t write_loc;
int32_t read_page_0;
int32_t read_page_1;
if (params.compress_ratio == 128) {
write_loc = compute_c128_loc(rid, position_1);
read_page_0 = compute_c128_loc(rid, position_0) / 128;
read_page_1 = compute_c128_loc(rid, position_1) / 128;
} else {
const auto raw_loc_0 = mapping[position_0];
const auto raw_loc_1 = mapping[position_1];
const auto state_loc_0 = params.f2s_ptr[raw_loc_0];
const auto state_loc_1 = params.f2s_ptr[raw_loc_1];
write_loc = static_cast<int32_t>(compute_loc(state_loc_1));
read_page_0 = static_cast<int32_t>(compute_loc(state_loc_0) / params.compress_ratio);
read_page_1 = static_cast<int32_t>(write_loc / params.compress_ratio);
}
params.plan_d[idx] = {
.seq_len = static_cast<uint32_t>(seq_len),
.write_loc = write_loc,
.read_page_0 = read_page_0,
.read_page_1 = read_page_1,
};
}
__global__ void plan_compress_prefill_legacy_kernel(const Prefill1ParamsLegacy params) {
const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= params.num_work) return;
auto plan_c = idx < params.num_c ? params.plan_c[idx] : PlanC::invalid();
auto plan_w = idx < params.num_w ? params.plan_w[idx] : PlanW::invalid();
/// Per-request ring buffer slot translation:
/// - c4: page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % 4
/// - c128: page = rid; slot = rid * 128 + position % 128
const auto legacy_compute_page = [&](int32_t rid, int32_t position) {
if (params.compress_ratio == 4) return rid * 2 + ((position / 4) & 1);
return rid; // c128
};
const auto legacy_compute_loc = [&](int32_t rid, int32_t position) {
const auto remainder = position % params.compress_ratio;
return legacy_compute_page(rid, position) * params.compress_ratio + remainder;
};
if (!plan_c.is_invalid()) {
const auto batch_id = plan_c.read_page_1;
const auto rid = static_cast<int32_t>(params.rid_ptr[batch_id]);
// `seq_len` is ratio-aligned for compress events
const auto position_1 = static_cast<int32_t>(plan_c.seq_len) - 1;
const auto position_0 = max(position_1 - params.compress_ratio, 0);
plan_c.read_page_0 = legacy_compute_page(rid, position_0);
plan_c.read_page_1 = legacy_compute_page(rid, position_1);
params.plan_c[idx] = plan_c;
} else if (idx < params.num_c_padded) {
params.plan_c[idx] = PlanC::invalid();
}
if (!plan_w.is_invalid()) {
const auto [ragged_id, batch_id] = unpack_w(plan_w);
const auto rid = static_cast<int32_t>(params.rid_ptr[batch_id]);
// `write_loc` carries (position + 1) at this stage; may not be ratio-aligned
const auto position = static_cast<int32_t>(plan_w.write_loc) - 1;
plan_w.ragged_id = ragged_id;
plan_w.write_loc = legacy_compute_loc(rid, position);
params.plan_w[idx] = plan_w;
} else if (idx < params.num_w_padded) {
params.plan_w[idx] = PlanW::invalid();
}
}
__global__ void plan_compress_decode_legacy_kernel(const DecodeParamsLegacy params) {
const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= params.batch_size) return;
/// Per-request ring buffer slot translation:
/// - c4: page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % 4
/// - c128: page = rid; slot = rid * 128 + position % 128
const auto legacy_compute_page = [&](int32_t rid, int32_t position) {
if (params.compress_ratio == 4) return rid * 2 + ((position / 4) & 1);
return rid; // c128
};
const auto legacy_compute_loc = [&](int32_t rid, int32_t position) {
const auto remainder = position % params.compress_ratio;
return legacy_compute_page(rid, position) * params.compress_ratio + remainder;
};
const auto rid = static_cast<int32_t>(params.rid_ptr[idx]);
const auto seq_len = static_cast<int32_t>(params.seq_ptr[idx]);
const auto position_1 = seq_len - 1;
const auto position_0 = max(position_1 - params.compress_ratio, 0);
const int32_t write_loc = legacy_compute_loc(rid, position_1);
const int32_t read_page_0 = legacy_compute_page(rid, position_0);
const int32_t read_page_1 = legacy_compute_page(rid, position_1);
params.plan_d[idx] = {
.seq_len = static_cast<uint32_t>(seq_len),
.write_loc = write_loc,
.read_page_0 = read_page_0,
.read_page_1 = read_page_1,
};
}
using PrefillPlan = tvm::ffi::Tuple<tvm::ffi::Tensor, tvm::ffi::Tensor>;
/**
* \brief Build c4/c128 prefill plan tensors. CPU-resident.
* Inputs (all CPU-resident):
* @param req_pool_indices `[batch_size]` int64_t
* @param req_to_token `[num_reqs, max_tokens_per_req]` int64_t
* @param full_to_state `[full_cache_size]` int64_t. For c4 this maps
* full loc -> SWA loc; ignored for c128, whose
* state slot is request-scoped.
* @param seq_lens `[batch_size]` int64
* @param extend_lens `[batch_size]` int64
* @param compress_plan `[num_q_tokens, 16]` uint8 (output)
* @param write_plan `[num_q_tokens, 8]` uint8 (output)
* @param compress_ratio 4 for c4, 128 for c128
* @param use_cuda_graph Whether the plans will be used with cuda graph (affects padding)
* @return (compress plan tensor, write plan tensor)
*/
inline PrefillPlan plan_compress_prefill(
const tvm::ffi::TensorView req_pool_indices, // GPU
const tvm::ffi::TensorView req_to_token, // GPU
const tvm::ffi::TensorView full_to_state, // GPU
const tvm::ffi::TensorView seq_lens, // CPU/GPU
const tvm::ffi::TensorView extend_lens, // CPU/GPU
const tvm::ffi::TensorView pin_buffer, // CPU
const uint32_t num_q_tokens,
const int32_t compress_ratio,
const int32_t swa_page_size,
const int32_t ring_size,
const bool use_cuda_graph) {
auto B = SymbolicSize{"batch_size"};
auto N = SymbolicSize{"num_q_tokens"};
auto cpu_or_gpu = SymbolicDevice{};
auto device_ = SymbolicDevice{};
cpu_or_gpu.set_options<kDLCPU, kDLGPU>();
device_.set_options<kDLGPU>();
TensorMatcher({B}) //
.with_dtype<RID_T>()
.with_device(device_)
.verify(req_pool_indices);
TensorMatcher({-1, -1}) //
.with_dtype<R2T_T>()
.with_device(device_)
.verify(req_to_token);
TensorMatcher({-1}) //
.with_dtype<F2S_T>()
.with_device(device_)
.verify(full_to_state);
TensorMatcher({B}) //
.with_dtype<IDX_T>()
.with_device(cpu_or_gpu)
.verify(seq_lens)
.verify(extend_lens);
TensorMatcher({-1}) //
.with_dtype<uint8_t>()
.with_device<kDLCPU>()
.verify(pin_buffer);
const bool is_overlap = (compress_ratio == 4);
const int32_t window_size = compress_ratio * (is_overlap ? 2 : 1);
const auto seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr());
const auto ext_ptr = static_cast<const IDX_T*>(extend_lens.data_ptr());
const auto rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr());
const auto r2t_ptr = static_cast<const R2T_T*>(req_to_token.data_ptr());
const auto f2s_ptr = static_cast<const F2S_T*>(full_to_state.data_ptr());
const auto batch_size = static_cast<uint32_t>(B.unwrap());
constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
// `swa_page_size` >= `ring_size` >= `compress_ratio`
RuntimeCheck(swa_page_size % ring_size == 0 && ring_size % compress_ratio == 0);
const auto device = device_.unwrap();
const auto stream = LaunchKernel::resolve_device(device);
constexpr int32_t kMaxMTPDraftTokens = 4;
const auto mtp_pad = std::min(ring_size - compress_ratio, kMaxMTPDraftTokens);
if (cpu_or_gpu.unwrap().device_type == kDLGPU) {
// GPU input path: kernel0 builds the (CPU-loop-equivalent) plan metadata directly
// on device, padding to num_q_tokens with invalid; kernel_1 then finalizes the
// SWA-translated read/write locations. Used for MTP / cuda-graph capture where
// a host sync would be expensive.
RuntimeCheck(batch_size <= kMaxPrefillBatchSize, "GPU plan only support batch size up to ", kMaxPrefillBatchSize);
auto C = ffi::empty({num_q_tokens, sizeof(PlanC)}, kDLUInt8, device);
auto W = ffi::empty({num_q_tokens, sizeof(PlanW)}, kDLUInt8, device);
const auto params0 = Prefill0Params{
.plan_c = static_cast<PlanC*>(C.data_ptr()),
.plan_w = static_cast<PlanW*>(W.data_ptr()),
.seq_lens_ptr = seq_ptr,
.extend_lens_ptr = ext_ptr,
.batch_size = batch_size,
.num_q_tokens = num_q_tokens,
.compress_ratio = compress_ratio,
.swa_page_size = swa_page_size,
.mtp_pad = mtp_pad,
};
LaunchKernel(1, kMaxPrefillBatchSize, device)(plan_compress_prefill_kernel0, params0);
// kernel_1 sees the already-padded buffers, so num_c == num_w == num_padded == num_q_tokens.
const auto params1 = Prefill1Params{
.plan_c = static_cast<PlanC*>(C.data_ptr()),
.plan_w = static_cast<PlanW*>(W.data_ptr()),
.rid_ptr = rid_ptr,
.r2t_ptr = r2t_ptr,
.f2s_ptr = f2s_ptr,
.stride_r2t = req_to_token.stride(0),
.num_c = num_q_tokens,
.num_w = num_q_tokens,
.num_c_padded = num_q_tokens,
.num_w_padded = num_q_tokens,
.num_work = num_q_tokens,
.swa_page_size = swa_page_size,
.ring_size = ring_size,
.compress_ratio = compress_ratio,
};
const auto block_size_1 = 256;
const auto num_blocks_1 = div_ceil(params1.num_work, block_size_1);
LaunchKernel(num_blocks_1, block_size_1, device)(plan_compress_prefill_kernel_1, params1);
return PrefillPlan{std::move(C), std::move(W)};
}
// CPU input path: only here do we need the pinned scratch buffer.
const auto pin_buffer_bytes = static_cast<size_t>(pin_buffer.numel()) * sizeof(uint8_t);
RuntimeCheck(pin_buffer_bytes >= num_q_tokens * (sizeof(PlanC) + sizeof(PlanW)));
const auto plan_c_ptr = reinterpret_cast<PlanC*>(pin_buffer.data_ptr());
const auto plan_w_ptr = reinterpret_cast<PlanW*>(plan_c_ptr + num_q_tokens);
uint32_t counter = 0;
uint32_t counter_c = 0;
uint32_t counter_w = 0;
const auto should_compress = [=](int32_t position) { return (position + 1) % compress_ratio == 0; };
for (const auto i : irange(batch_size)) {
const int32_t seq_len = seq_ptr[i];
const int32_t extend_len = ext_ptr[i];
const int32_t prefix_len = seq_len - extend_len;
const int32_t last_c_pos = seq_len / compress_ratio * compress_ratio;
const int32_t first_w_pos = last_c_pos - (is_overlap ? compress_ratio : 0);
RuntimeCheck(0 < extend_len && extend_len <= seq_len);
const auto should_write = [=](int32_t position) {
if (position >= first_w_pos) return true;
return is_overlap && position % swa_page_size >= (swa_page_size - compress_ratio);
};
for (const auto j : irange(extend_len)) {
const int32_t position = prefix_len + j;
const int32_t ragged_id = counter + j;
if (should_compress(position)) {
const auto buffer_len = window_size - std::min(j + 1, window_size);
plan_c_ptr[counter_c++] = {
.seq_len = static_cast<uint32_t>(position + 1),
.ragged_id = static_cast<uint16_t>(ragged_id),
.buffer_len = static_cast<uint16_t>(buffer_len),
// to be filled by kernel
.read_page_0 = -1,
.read_page_1 = static_cast<int32_t>(i),
};
}
if (should_write(position)) {
plan_w_ptr[counter_w++] = pack_w(ragged_id, i, position + 1);
}
}
counter += extend_len;
}
RuntimeCheck(counter == num_q_tokens);
const auto copy_to_device = [stream](void* cuda_ptr, auto* host_ptr, size_t count) {
const auto size_bytes = count * sizeof(*host_ptr);
RuntimeDeviceCheck(cudaMemcpyAsync(cuda_ptr, host_ptr, size_bytes, cudaMemcpyHostToDevice, stream));
};
const auto num_c_padded = use_cuda_graph ? num_q_tokens : counter_c;
const auto num_w_padded = use_cuda_graph ? num_q_tokens : counter_w;
auto C = ffi::empty({num_c_padded, sizeof(PlanC)}, kDLUInt8, device);
auto W = ffi::empty({num_w_padded, sizeof(PlanW)}, kDLUInt8, device);
copy_to_device(C.data_ptr(), plan_c_ptr, counter_c);
copy_to_device(W.data_ptr(), plan_w_ptr, counter_w);
const auto params = Prefill1Params{
.plan_c = static_cast<PlanC*>(C.data_ptr()),
.plan_w = static_cast<PlanW*>(W.data_ptr()),
.rid_ptr = rid_ptr,
.r2t_ptr = r2t_ptr,
.f2s_ptr = f2s_ptr,
.stride_r2t = req_to_token.size(1),
.num_c = counter_c,
.num_w = counter_w,
.num_c_padded = num_c_padded,
.num_w_padded = num_w_padded,
.num_work = std::max(num_c_padded, num_w_padded),
.swa_page_size = swa_page_size,
.ring_size = ring_size,
.compress_ratio = compress_ratio,
};
const auto block_size = 256;
const auto num_blocks = div_ceil(params.num_work, block_size);
LaunchKernel(num_blocks, block_size, device)(plan_compress_prefill_kernel_1, params);
return PrefillPlan{std::move(C), std::move(W)};
}
inline tvm::ffi::Tensor plan_compress_decode(
const tvm::ffi::TensorView req_pool_indices, // GPU
const tvm::ffi::TensorView req_to_token, // GPU
const tvm::ffi::TensorView full_to_state, // GPU
const tvm::ffi::TensorView seq_lens, // CPU/GPU
const int32_t compress_ratio,
const int32_t swa_page_size,
const int32_t ring_size) {
auto B = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLGPU>();
TensorMatcher({B}) //
.with_dtype<RID_T>()
.with_device(device_)
.verify(req_pool_indices);
TensorMatcher({-1, -1}) //
.with_dtype<R2T_T>()
.with_device(device_)
.verify(req_to_token);
TensorMatcher({-1}) //
.with_dtype<F2S_T>()
.with_device(device_)
.verify(full_to_state);
TensorMatcher({B}) //
.with_dtype<IDX_T>()
.with_device(device_)
.verify(seq_lens);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto device = device_.unwrap();
auto D = ffi::empty({batch_size, sizeof(PlanD)}, kDLUInt8, device);
const auto params = DecodeParams{
.plan_d = static_cast<PlanD*>(D.data_ptr()),
.rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr()),
.r2t_ptr = static_cast<const R2T_T*>(req_to_token.data_ptr()),
.f2s_ptr = static_cast<const F2S_T*>(full_to_state.data_ptr()),
.seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr()),
.stride_r2t = req_to_token.size(1),
.batch_size = batch_size,
.swa_page_size = swa_page_size,
.ring_size = ring_size,
.compress_ratio = compress_ratio,
};
const auto block_size = 256;
const auto num_blocks = div_ceil(batch_size, block_size);
LaunchKernel(num_blocks, block_size, device)(plan_compress_decode_kernel, params);
return D;
}
/**
* \brief Build c4/c128 prefill plan tensors for the legacy non-paged ring
* buffer. Uses only `req_pool_indices` to derive ring slots:
* - c4 (overlap): each request occupies 2 contiguous pages (8 token slots)
* - c128: each request occupies 1 page (128 token slots)
*
* Inputs:
* @param req_pool_indices `[batch_size]` int64 (GPU)
* @param seq_lens `[batch_size]` int64 (CPU)
* @param extend_lens `[batch_size]` int64 (CPU)
* @param pin_buffer pinned scratch (CPU uint8)
* @return (compress plan tensor, write plan tensor)
*/
inline PrefillPlan plan_compress_prefill_legacy(
const tvm::ffi::TensorView req_pool_indices, // GPU
const tvm::ffi::TensorView seq_lens, // CPU
const tvm::ffi::TensorView extend_lens, // CPU
const tvm::ffi::TensorView pin_buffer, // CPU
const uint32_t num_q_tokens,
const int32_t compress_ratio,
const bool use_cuda_graph) {
auto B = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLGPU>();
TensorMatcher({B}) //
.with_dtype<RID_T>()
.with_device(device_)
.verify(req_pool_indices);
TensorMatcher({B}) //
.with_dtype<IDX_T>()
.with_device<kDLCPU>()
.verify(seq_lens)
.verify(extend_lens);
TensorMatcher({-1}) //
.with_dtype<uint8_t>()
.with_device<kDLCPU>()
.verify(pin_buffer);
const auto pin_buffer_bytes = static_cast<size_t>(pin_buffer.numel()) * sizeof(uint8_t);
RuntimeCheck(pin_buffer_bytes >= num_q_tokens * (sizeof(PlanC) + sizeof(PlanW)));
const auto plan_c_ptr = reinterpret_cast<PlanC*>(pin_buffer.data_ptr());
const auto plan_w_ptr = reinterpret_cast<PlanW*>(plan_c_ptr + num_q_tokens);
const bool is_overlap = (compress_ratio == 4);
const auto seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr());
const auto ext_ptr = static_cast<const IDX_T*>(extend_lens.data_ptr());
const auto rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr());
const auto window_size = compress_ratio * (is_overlap ? 2 : 1);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
uint32_t counter = 0;
uint32_t counter_c = 0;
uint32_t counter_w = 0;
const auto should_compress = [=](int32_t position) { return (position + 1) % compress_ratio == 0; };
for (const auto i : irange(batch_size)) {
const int32_t seq_len = seq_ptr[i];
const int32_t extend_len = ext_ptr[i];
const int32_t prefix_len = seq_len - extend_len;
const int32_t last_c_pos = seq_len / compress_ratio * compress_ratio;
const int32_t first_w_pos = last_c_pos - (is_overlap ? compress_ratio : 0);
RuntimeCheck(0 < extend_len && extend_len <= seq_len);
const auto should_write = [=](int32_t position) { return position >= first_w_pos; };
for (const auto j : irange(extend_len)) {
const int32_t position = prefix_len + j;
const int32_t ragged_id = counter + j;
if (should_compress(position)) {
const auto buffer_len = window_size - std::min(j + 1, window_size);
plan_c_ptr[counter_c++] = {
.seq_len = static_cast<uint32_t>(position + 1),
.ragged_id = static_cast<uint16_t>(ragged_id),
.buffer_len = static_cast<uint16_t>(buffer_len),
// to be filled by kernel
.read_page_0 = -1,
.read_page_1 = static_cast<int32_t>(i),
};
}
if (should_write(position)) {
plan_w_ptr[counter_w++] = pack_w(ragged_id, i, position + 1);
}
}
counter += extend_len;
}
RuntimeCheck(counter == num_q_tokens);
const auto device = device_.unwrap();
const auto stream = LaunchKernel::resolve_device(device);
const auto copy_to_device = [stream](void* cuda_ptr, auto* host_ptr, size_t count) {
const auto size_bytes = count * sizeof(*host_ptr);
RuntimeDeviceCheck(cudaMemcpyAsync(cuda_ptr, host_ptr, size_bytes, cudaMemcpyHostToDevice, stream));
};
const auto num_c_padded = use_cuda_graph ? num_q_tokens : counter_c;
const auto num_w_padded = use_cuda_graph ? num_q_tokens : counter_w;
auto C = ffi::empty({num_c_padded, sizeof(PlanC)}, kDLUInt8, device);
auto W = ffi::empty({num_w_padded, sizeof(PlanW)}, kDLUInt8, device);
copy_to_device(C.data_ptr(), plan_c_ptr, counter_c);
copy_to_device(W.data_ptr(), plan_w_ptr, counter_w);
const auto params = Prefill1ParamsLegacy{
.plan_c = static_cast<PlanC*>(C.data_ptr()),
.plan_w = static_cast<PlanW*>(W.data_ptr()),
.rid_ptr = rid_ptr,
.num_c = counter_c,
.num_w = counter_w,
.num_c_padded = num_c_padded,
.num_w_padded = num_w_padded,
.num_work = std::max(num_c_padded, num_w_padded),
.compress_ratio = compress_ratio,
};
const auto block_size = 256;
const auto num_blocks = div_ceil(params.num_work, block_size);
if (num_blocks > 0) {
LaunchKernel(num_blocks, block_size, device)(plan_compress_prefill_legacy_kernel, params);
}
return PrefillPlan{std::move(C), std::move(W)};
}
inline tvm::ffi::Tensor plan_compress_decode_legacy(
const tvm::ffi::TensorView req_pool_indices, // GPU
const tvm::ffi::TensorView seq_lens, // GPU
const int32_t compress_ratio) {
auto B = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLGPU>();
TensorMatcher({B}) //
.with_dtype<RID_T>()
.with_device(device_)
.verify(req_pool_indices);
TensorMatcher({B}) //
.with_dtype<IDX_T>()
.with_device(device_)
.verify(seq_lens);
RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto device = device_.unwrap();
auto D = ffi::empty({batch_size, sizeof(PlanD)}, kDLUInt8, device);
const auto params = DecodeParamsLegacy{
.plan_d = static_cast<PlanD*>(D.data_ptr()),
.rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr()),
.seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr()),
.batch_size = batch_size,
.compress_ratio = compress_ratio,
};
const auto block_size = 256;
const auto num_blocks = div_ceil(batch_size, block_size);
LaunchKernel(num_blocks, block_size, device)(plan_compress_decode_legacy_kernel, params);
return D;
}
} // namespace host::compress
using namespace host::compress; // expose binding
python/sglang/jit_kernel/csrc/deepseek_v4/common.cuh
+208
View File
@@ -0,0 +1,208 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/deepseek_v4/compress.cuh>
#include <dlpack/dlpack.h>
namespace host::compress {
using PlanResult = tvm::ffi::Tuple<uint32_t, uint32_t>;
struct CompressParams {
PrefillPlan* __restrict__ compress_plan;
PrefillPlan* __restrict__ write_plan;
const int64_t* __restrict__ seq_lens;
const int64_t* __restrict__ extend_lens;
uint32_t batch_size;
uint32_t num_tokens;
uint32_t compress_ratio;
bool is_overlap;
};
inline constexpr uint32_t kBlockSize = 1024;
#define PLAN_KERNEL __global__ __launch_bounds__(kBlockSize, 1) inline
PLAN_KERNEL void plan_prefill_cuda(const __grid_constant__ CompressParams params) {
const auto &[
compress_plan, write_plan, seq_lens, extend_lens, // pointers
batch_size, num_tokens, compress_ratio, is_overlap // values
] = params;
__shared__ uint32_t compress_counter;
__shared__ uint32_t write_counter;
uint32_t batch_id = 0;
uint32_t counter = 0;
uint32_t extend_len = extend_lens[0];
const auto tid = threadIdx.x;
if (tid == 0) {
compress_counter = 0;
write_counter = 0;
}
__syncthreads();
for (uint32_t i = tid; i < num_tokens; i += blockDim.x) {
const uint32_t ragged_id = i;
uint32_t j = ragged_id - counter;
while (j >= extend_len) {
j -= extend_len;
batch_id += 1;
if (batch_id >= batch_size) [[unlikely]]
break;
counter += extend_len;
extend_len = extend_lens[batch_id];
}
if (batch_id >= batch_size) [[unlikely]]
break;
const uint32_t seq_len = seq_lens[batch_id];
const uint32_t extend_len = extend_lens[batch_id];
const uint32_t prefix_len = seq_len - extend_len;
const uint32_t ratio = compress_ratio * (1 + is_overlap);
const uint32_t window_len = j + 1 < ratio ? ratio - (j + 1) : 0;
const uint32_t position = prefix_len + j;
const auto plan = PrefillPlan{
.ragged_id = ragged_id,
.batch_id = batch_id,
.position = position,
.window_len = window_len,
};
const uint32_t start_write_pos = [seq_len, compress_ratio, is_overlap] {
const uint32_t pos = seq_len / compress_ratio * compress_ratio;
if (!is_overlap) return pos;
return pos >= compress_ratio ? pos - compress_ratio : 0;
}();
if ((position + 1) % compress_ratio == 0) {
const auto write_pos = atomicAdd(&compress_counter, 1);
compress_plan[write_pos] = plan;
}
if (position >= start_write_pos) {
const auto write_pos = atomicAdd(&write_counter, 1);
write_plan[write_pos] = plan;
}
}
__syncthreads();
constexpr auto kInvalid = static_cast<uint32_t>(-1);
const auto kInvalidPlan = PrefillPlan{kInvalid, kInvalid, kInvalid, kInvalid};
const auto compress_count = compress_counter;
const auto write_count = write_counter;
for (uint32_t i = compress_count + tid; i < num_tokens; i += blockDim.x) {
compress_plan[i] = kInvalidPlan;
}
for (uint32_t i = write_count + tid; i < num_tokens; i += blockDim.x) {
write_plan[i] = kInvalidPlan;
}
}
inline PlanResult plan_prefill_host(const CompressParams& params, const bool use_cuda_graph) {
const auto &[
compress_ptr, write_ptr, seq_lens_ptr, extend_lens_ptr, // pointers
batch_size, num_tokens, compress_ratio, is_overlap // values
] = params;
uint32_t counter = 0;
uint32_t compress_counter = 0;
uint32_t write_counter = 0;
const auto ratio = compress_ratio * (1 + is_overlap);
for (const auto i : irange(batch_size)) {
const uint32_t seq_len = seq_lens_ptr[i];
const uint32_t extend_len = extend_lens_ptr[i];
const uint32_t prefix_len = seq_len - extend_len;
RuntimeCheck(0 < extend_len && extend_len <= seq_len);
/// NOTE: `start_write_pos` must be a multiple of `compress_ratio`
const uint32_t start_write_pos = [seq_len, compress_ratio, is_overlap] {
const uint32_t pos = seq_len / compress_ratio * compress_ratio;
if (!is_overlap) return pos;
/// NOTE: to avoid unsigned integer underflow, don't use `pos - compress_ratio`
return pos >= compress_ratio ? pos - compress_ratio : 0;
}();
/// NOTE: `position` is within [prefix_len, seq_len)
for (const auto j : irange(extend_len)) {
const uint32_t position = prefix_len + j;
const auto plan = PrefillPlan{
.ragged_id = counter + j,
.batch_id = i,
.position = position,
.window_len = ratio - std::min(j + 1, ratio),
};
RuntimeCheck(plan.is_valid(compress_ratio, is_overlap), "Internal error!");
if ((position + 1) % compress_ratio == 0) {
compress_ptr[compress_counter++] = plan;
}
if (position >= start_write_pos) {
write_ptr[write_counter++] = plan;
}
}
counter += extend_len;
}
RuntimeCheck(counter == num_tokens, "input size ", counter, " != num_q_tokens ", num_tokens);
if (!use_cuda_graph) return PlanResult{compress_counter, write_counter};
constexpr auto kInvalid = static_cast<uint32_t>(-1);
constexpr auto kInvalidPlan = PrefillPlan{kInvalid, kInvalid, kInvalid, kInvalid};
for (const auto i : irange(compress_counter, num_tokens)) {
compress_ptr[i] = kInvalidPlan;
}
for (const auto i : irange(write_counter, num_tokens)) {
write_ptr[i] = kInvalidPlan;
}
return PlanResult{num_tokens, num_tokens};
}
inline PlanResult plan_prefill(
const tvm::ffi::TensorView extend_lens,
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::TensorView compress_plan,
const tvm::ffi::TensorView write_plan,
const uint32_t compress_ratio,
const bool is_overlap, // for overlap transform, we have to keep 1 more extra window
const bool use_cuda_graph) {
auto N = SymbolicSize{"batch_size"};
auto M = SymbolicSize{"num_tokens"};
auto device = SymbolicDevice{};
const bool is_cuda = [&] {
if (extend_lens.device().device_type == kDLCUDA) {
device.set_options<kDLCUDA>();
return true;
} else {
device.set_options<kDLCPU, kDLCUDAHost>();
return false;
}
}();
TensorMatcher({N}) // extend_lens and seq_lens
.with_dtype<int64_t>()
.with_device(device)
.verify(extend_lens)
.verify(seq_lens);
TensorMatcher({M, kPrefillPlanDim}) // compress_plan and write_plan
.with_dtype<PrefillPlanTensorDtype>()
.with_device(device)
.verify(compress_plan)
.verify(write_plan);
const auto params = CompressParams{
.compress_plan = static_cast<PrefillPlan*>(compress_plan.data_ptr()),
.write_plan = static_cast<PrefillPlan*>(write_plan.data_ptr()),
.seq_lens = static_cast<const int64_t*>(seq_lens.data_ptr()),
.extend_lens = static_cast<const int64_t*>(extend_lens.data_ptr()),
.batch_size = static_cast<uint32_t>(N.unwrap()),
.num_tokens = static_cast<uint32_t>(M.unwrap()),
.compress_ratio = compress_ratio,
.is_overlap = is_overlap,
};
if (!is_cuda) return plan_prefill_host(params, use_cuda_graph);
/// NOTE: cuda kernel plan is naturally compatible with cuda graph
LaunchKernel(1, kBlockSize, device.unwrap())(plan_prefill_cuda, params);
return PlanResult{params.num_tokens, params.num_tokens};
}
} // namespace host::compress
namespace {
[[maybe_unused]]
constexpr auto& plan_compress_prefill = host::compress::plan_prefill;
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh
+169
View File
@@ -0,0 +1,169 @@
// DeepSeek-V4 wo_a activation quantization for DeepGEMM fp8_einsum.
//
// This is intentionally narrower than the generic per_token_group_quant_8bit_v2
// kernel: input is a [T, G, D] view with contiguous hidden groups, output_q is
// contiguous [T, G, D], group_size is fixed to 128, scales are fp32 UE8M0
// power-of-two values, and output_s is a logical [T, G, D/128] view backed by
// group-major [G, T, D/128] storage.
//
// The generic kernel cannot read the strided DSV4 view while producing
// contiguous [T, G, D] codes and group-major scales without an extra full-tensor
// copy.
#include <sgl_kernel/tensor.h> // TensorMatcher, SymbolicSize/Device
#include <sgl_kernel/utils.h> // RuntimeCheck
#include <sgl_kernel/utils.cuh> // fp8 aliases, PDL helpers
#include <sgl_kernel/warp.cuh> // warp::reduce_max
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh> // UE8M0 and FP8 helpers
#include <tvm/ffi/container/tensor.h> // tvm::ffi::TensorView
#include <cstdint>
#include <cuda_fp8.h>
namespace {
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::inv_scale_ue8m0;
using deepseek_v4::fp8::pack_fp8;
constexpr float LOCAL_ABSMAX_ABS = 1e-10f;
constexpr uint32_t GROUP_SIZE = 128;
constexpr uint32_t THREADS_PER_GROUP = 8;
constexpr uint32_t SUBWARPS_PER_BLOCK = 16;
constexpr uint32_t INPUT_VEC_NUM_BYTES = 32;
constexpr uint32_t INPUT_INT4_SIZE = INPUT_VEC_NUM_BYTES / sizeof(int4);
template <int THREADS_PER_SUBWARP>
SGL_DEVICE float GroupReduceMax(float val) {
static_assert(
(THREADS_PER_SUBWARP & (THREADS_PER_SUBWARP - 1)) == 0 && THREADS_PER_SUBWARP <= 16 && THREADS_PER_SUBWARP >= 1,
"THREADS_PER_SUBWARP must be 1, 2, 4, 8, or 16");
// Tail subwarps can be inactive at the bounds check, so reduce with only the
// current subgroup's lanes rather than a full-warp mask.
constexpr device::warp::mask_t kSub = (device::warp::mask_t{1} << THREADS_PER_SUBWARP) - 1;
const device::warp::mask_t mask = kSub << (THREADS_PER_SUBWARP * ((threadIdx.x % 32) / THREADS_PER_SUBWARP));
return device::warp::reduce_max<THREADS_PER_SUBWARP>(val, mask);
}
template <typename T, bool kUsePDL>
__global__ void fp8_wo_a_group_major_quant_ue8m0_kernel(
const T* __restrict__ input,
fp8_e4m3_t* __restrict__ output_q,
float* __restrict__ output_s,
int64_t total_scale_groups,
int64_t num_tokens,
int hidden_dim_groups,
int num_outer_groups,
int64_t input_stride_t) {
device::PDLWaitPrimary<kUsePDL>();
const int64_t subwarp_id = threadIdx.x / THREADS_PER_GROUP;
const int lane_id = threadIdx.x % THREADS_PER_GROUP;
const int64_t group_id = static_cast<int64_t>(blockIdx.x) * SUBWARPS_PER_BLOCK + subwarp_id;
if (group_id < total_scale_groups) {
const int hidden_group = group_id % hidden_dim_groups;
const int64_t token_outer = group_id / hidden_dim_groups;
const int outer_idx = token_outer % num_outer_groups;
const int64_t token_idx = token_outer / num_outer_groups;
constexpr uint32_t INPUT_VEC_SIZE = INPUT_VEC_NUM_BYTES / sizeof(T);
static_assert(INPUT_VEC_SIZE * THREADS_PER_GROUP == GROUP_SIZE);
const int64_t input_group_start_offset =
token_idx * input_stride_t + outer_idx * GROUP_SIZE * hidden_dim_groups + hidden_group * GROUP_SIZE;
const int64_t output_group_start_offset = group_id * GROUP_SIZE;
int4 input_int4[INPUT_INT4_SIZE];
T* input_vec = reinterpret_cast<T*>(input_int4);
#pragma unroll
for (uint32_t j = 0; j < INPUT_INT4_SIZE; ++j) {
input_int4[j] = reinterpret_cast<const int4*>(input + input_group_start_offset + lane_id * INPUT_VEC_SIZE)[j];
}
float local_absmax = LOCAL_ABSMAX_ABS;
#pragma unroll
for (uint32_t j = 0; j < INPUT_VEC_SIZE; ++j) {
const float val = static_cast<float>(input_vec[j]);
local_absmax = fmaxf(local_absmax, fabsf(val));
}
local_absmax = GroupReduceMax<THREADS_PER_GROUP>(local_absmax);
constexpr float kFp8MaxInv = 1.0f / kFP8E4M3Max;
const int32_t scale_ue8m0 = cast_to_ue8m0(local_absmax * kFp8MaxInv);
const float y_scale = inv_scale_ue8m0(scale_ue8m0);
const float y_scale_inv = __uint_as_float(static_cast<uint32_t>(scale_ue8m0) << 23);
int4 output_buf;
auto* output_buf_ptr = reinterpret_cast<fp8x2_e4m3_t*>(&output_buf);
#pragma unroll
for (uint32_t j = 0; j < INPUT_VEC_SIZE; j += 2) {
output_buf_ptr[j / 2] =
pack_fp8(static_cast<float>(input_vec[j]) * y_scale, static_cast<float>(input_vec[j + 1]) * y_scale);
}
*reinterpret_cast<int4*>(output_q + output_group_start_offset + lane_id * INPUT_VEC_SIZE) = output_buf;
if (lane_id == 0) {
output_s[(outer_idx * num_tokens + token_idx) * hidden_dim_groups + hidden_group] = y_scale_inv;
}
}
device::PDLTriggerSecondary<kUsePDL>();
}
template <typename T, bool kUsePDL>
struct FP8WoAGroupMajorQuantUE8M0Kernel {
static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) {
using namespace host;
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
auto TSize = SymbolicSize{"num_tokens"};
auto GSize = SymbolicSize{"num_outer_groups"};
auto DSize = SymbolicSize{"hidden_dim"};
auto SSize = SymbolicSize{"hidden_dim_groups"};
TensorMatcher({TSize, GSize, DSize}).with_strides({-1, DSize, 1}).with_dtype<T>().with_device(device).verify(input);
TensorMatcher({TSize, GSize, DSize}).with_dtype<fp8_e4m3_t>().with_device(device).verify(output_q);
TensorMatcher({GSize, TSize, SSize}).with_dtype<float>().with_device(device).verify(output_s);
const auto num_tokens = TSize.unwrap();
const auto num_outer_groups = GSize.unwrap();
const auto hidden_dim = DSize.unwrap();
const auto hidden_dim_groups = SSize.unwrap();
const auto input_stride_t = input.stride(0);
constexpr int64_t kInputAlignElements = sizeof(int4) / sizeof(T);
RuntimeCheck(hidden_dim % GROUP_SIZE == 0, "hidden_dim must be divisible by 128");
RuntimeCheck(hidden_dim_groups == hidden_dim / GROUP_SIZE, "output_s hidden dim mismatch");
RuntimeCheck(
reinterpret_cast<uintptr_t>(input.data_ptr()) % sizeof(int4) == 0,
"input base pointer must be 16-byte aligned");
RuntimeCheck(
num_tokens <= 1 || input_stride_t % kInputAlignElements == 0,
"input token stride must preserve 16-byte vector-load alignment");
const int64_t total_scale_groups = num_tokens * num_outer_groups * hidden_dim_groups;
if (total_scale_groups == 0) return;
const auto grid = dim3((total_scale_groups + SUBWARPS_PER_BLOCK - 1) / SUBWARPS_PER_BLOCK);
const auto block = dim3(SUBWARPS_PER_BLOCK * THREADS_PER_GROUP);
host::LaunchKernel(grid, block, device.unwrap())
.enable_pdl(kUsePDL)(
fp8_wo_a_group_major_quant_ue8m0_kernel<T, kUsePDL>,
static_cast<const T*>(input.data_ptr()),
static_cast<fp8_e4m3_t*>(output_q.data_ptr()),
static_cast<float*>(output_s.data_ptr()),
total_scale_groups,
static_cast<int64_t>(num_tokens),
static_cast<int>(hidden_dim_groups),
static_cast<int>(num_outer_groups),
static_cast<int64_t>(input_stride_t));
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope.cuh
+254
View File
@@ -0,0 +1,254 @@
#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 <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/compress.cuh>
#include <tvm/ffi/container/tensor.h>
#include <cstdint>
#include <type_traits>
namespace {
using Plan = device::compress::PrefillPlan;
/// \brief common block size for memory-bound kernel
constexpr uint32_t kBlockSize = 128;
constexpr uint32_t kNumWarps = kBlockSize / device::kWarpThreads;
struct FusedNormRopeParams {
void* __restrict__ input;
const void* __restrict__ weight;
float eps;
uint32_t num_works;
const void* __restrict__ handle;
const float* __restrict__ freqs_cis;
uint32_t compress_ratio;
};
enum class ForwardMode {
CompressExtend = 0,
CompressDecode = 1,
DefaultForward = 2,
};
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, ForwardMode kMode, bool kUsePDL>
__global__ void fused_norm_rope(const __grid_constant__ FusedNormRopeParams params) {
using namespace device;
using enum ForwardMode;
constexpr int64_t kMaxVecSize = 16 / sizeof(DType);
constexpr int64_t kVecSize = std::min(kMaxVecSize, kHeadDim / kWarpThreads);
constexpr int64_t kLocalSize = kHeadDim / (kWarpThreads * kVecSize);
constexpr int64_t kRopeVecSize = kRopeDim / (kWarpThreads * 2);
constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
static_assert(kHeadDim % (kWarpThreads * kVecSize) == 0);
static_assert(kLocalSize * kVecSize * kWarpThreads == kHeadDim);
static_assert(kRopeDim % (kWarpThreads * 2) == 0);
static_assert(kRopeDim % (kVecSize * kLocalSize) == 0);
static_assert(kRopeSize <= kWarpThreads);
static_assert(kRopeVecSize == 1, "only support rope dim = 64");
const auto& [
_input, _weight, eps, num_works, // norm
handle, freqs_cis, compress_ratio // rope
] = params;
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
const auto work_id = blockIdx.x * kNumWarps + warp_id;
if (work_id >= num_works) return;
DType* input;
int32_t position;
if constexpr (kMode == CompressExtend) {
const auto plan = static_cast<const Plan*>(handle)[work_id];
input = static_cast<DType*>(_input) + plan.ragged_id * kHeadDim;
position = plan.position + 1 - compress_ratio;
if (plan.ragged_id == 0xFFFFFFFF) [[unlikely]]
return;
} else if constexpr (kMode == CompressDecode) {
input = static_cast<DType*>(_input) + work_id * kHeadDim;
const auto seq_len = static_cast<const int32_t*>(handle)[work_id];
if (seq_len % compress_ratio != 0) return;
position = seq_len - compress_ratio;
} else if constexpr (kMode == DefaultForward) {
input = static_cast<DType*>(_input) + work_id * kHeadDim;
position = static_cast<const int64_t*>(handle)[work_id];
} else {
static_assert(host::dependent_false_v<DType>, "Unsupported Mode");
}
using Storage = AlignedVector<DType, kVecSize>;
__shared__ Storage s_rope_input[kNumWarps][kRopeSize];
// prefetch freq
const auto mem_freq = tile::Memory<fp32x2_t>::warp();
const auto freq = mem_freq.load(freqs_cis + position * kRopeDim);
PDLWaitPrimary<kUsePDL>();
// part 1: norm
{
const auto gmem = tile::Memory<Storage>::warp();
Storage input_vec[kLocalSize];
Storage weight_vec[kLocalSize];
#pragma unroll
for (int i = 0; i < kLocalSize; ++i) {
input_vec[i] = gmem.load(input, i);
}
#pragma unroll
for (int i = 0; i < kLocalSize; ++i) {
weight_vec[i] = gmem.load(_weight, i);
}
float sum_of_squares = 0.0f;
#pragma unroll
for (int i = 0; i < kLocalSize; ++i) {
#pragma unroll
for (int j = 0; j < kVecSize; ++j) {
const auto fp32_input = cast<float>(input_vec[i][j]);
sum_of_squares += fp32_input * fp32_input;
}
}
sum_of_squares = warp::reduce_sum(sum_of_squares);
const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + eps);
#pragma unroll
for (int i = 0; i < kLocalSize; ++i) {
#pragma unroll
for (int j = 0; j < kVecSize; ++j) {
const auto fp32_input = cast<float>(input_vec[i][j]);
const auto fp32_weight = cast<float>(weight_vec[i][j]);
input_vec[i][j] = cast<DType>(fp32_input * norm_factor * fp32_weight);
}
}
const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
#pragma unroll
for (int i = 0; i < kLocalSize; ++i) {
if (i == kLocalSize - 1 && is_rope_lane) {
const auto rope_id = lane_id - (kWarpThreads - kRopeSize);
s_rope_input[warp_id][rope_id] = input_vec[i];
} else {
gmem.store(input, input_vec[i], i);
}
}
__syncwarp();
}
// part 2: rope
{
// mem elem = DType x 2
using DTypex2_t = packed_t<DType>;
const auto mem_elem = tile::Memory<DTypex2_t>::warp();
const auto elem = mem_elem.load(s_rope_input[warp_id]);
const auto [x_real, x_imag] = cast<fp32x2_t>(elem);
const auto [freq_real, freq_imag] = freq;
const fp32x2_t output = {
x_real * freq_real - x_imag * freq_imag,
x_real * freq_imag + x_imag * freq_real,
};
mem_elem.store(input + (kHeadDim - kRopeDim), cast<DTypex2_t>(output));
}
PDLTriggerSecondary<kUsePDL>();
}
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, bool kUsePDL>
struct FusedNormRopeKernel {
template <ForwardMode kMode>
static constexpr auto fused_kernel = fused_norm_rope<DType, kHeadDim, kRopeDim, kMode, kUsePDL>;
static void forward(
const tvm::ffi::TensorView input,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView handle,
const tvm::ffi::TensorView freqs_cis,
int32_t _mode,
float eps,
uint32_t compress_ratio) {
using namespace host;
using enum ForwardMode;
const auto mode = static_cast<ForwardMode>(_mode);
auto B = SymbolicSize{"num_q_tokens"};
auto N = SymbolicSize{"num_compress_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B, kHeadDim}) // input
.with_dtype<DType>()
.with_device(device_)
.verify(input);
TensorMatcher({kHeadDim}) // weight
.with_dtype<DType>()
.with_device(device_)
.verify(weight);
TensorMatcher({-1, kRopeDim}) // freqs_cis
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
switch (mode) {
case CompressExtend:
TensorMatcher({N, compress::kPrefillPlanDim}) // plan
.with_dtype<compress::PrefillPlanTensorDtype>()
.with_device(device_)
.verify(handle);
RuntimeCheck(compress_ratio > 0);
break;
case CompressDecode:
TensorMatcher({N}) // seq_len
.with_dtype<int32_t>()
.with_device(device_)
.verify(handle);
RuntimeCheck(compress_ratio > 0);
break;
case DefaultForward:
TensorMatcher({N}) // position
.with_dtype<int64_t>()
.with_device(device_)
.verify(handle);
RuntimeCheck(compress_ratio == 0);
break;
default:
Panic("unsupported forward mode: ", static_cast<int>(mode));
}
// launch kernel
const auto num_compress_tokens = static_cast<uint32_t>(N.unwrap());
if (num_compress_tokens == 0) return;
const auto params = FusedNormRopeParams{
.input = input.data_ptr(),
.weight = weight.data_ptr(),
.eps = eps,
.num_works = num_compress_tokens,
.handle = handle.data_ptr(),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.compress_ratio = compress_ratio,
};
const auto num_blocks = div_ceil(num_compress_tokens, kNumWarps);
using KernelType = std::decay_t<decltype(fused_norm_rope<DType, kHeadDim, kRopeDim, CompressExtend, kUsePDL>)>;
static constexpr KernelType kernel_table[3] = {
[static_cast<int>(CompressExtend)] = fused_kernel<CompressExtend>,
[static_cast<int>(CompressDecode)] = fused_kernel<CompressDecode>,
[static_cast<int>(DefaultForward)] = fused_kernel<DefaultForward>,
};
const auto kernel = kernel_table[static_cast<int>(mode)];
LaunchKernel(num_blocks, kBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh
+663
View File
@@ -0,0 +1,663 @@
#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 <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/compress_v2.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <tvm/ffi/container/tensor.h>
#include <cstdint>
namespace {
using PlanC = device::compress::CompressPlan;
using PlanD = device::compress::DecodePlan;
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::inv_scale_ue8m0;
using deepseek_v4::fp8::pack_fp8;
SGL_DEVICE uint8_t quant_fp4_e2m1(float x) {
const float ax = fminf(fabsf(x), 6.0f);
uint8_t idx = 0;
idx += ax > 0.25f;
idx += ax > 0.75f;
idx += ax > 1.25f;
idx += ax > 1.75f;
idx += ax > 2.5f;
idx += ax > 3.5f;
idx += ax > 5.0f;
if (x < 0.0f && idx != 0) idx |= 0x8;
return idx;
}
constexpr uint32_t kBlockSize = 256;
constexpr uint32_t kNumWarps = kBlockSize / device::kWarpThreads;
struct FusedNormRopeStoreParams {
void* __restrict__ input;
const void* __restrict__ handle; // plan decode / compress
const void* __restrict__ weight;
const float* __restrict__ freqs_cis;
const int64_t* __restrict__ out_loc;
uint8_t* __restrict__ kvcache;
float eps;
uint32_t compress_ratio;
uint32_t num_tokens;
};
enum class ForwardMode : bool {
CompressExtend = 0,
CompressDecode = 1,
};
#define INDEXER_KERNEL __global__ __launch_bounds__(kBlockSize, 8)
#define FLASHMLA_KERNEL __global__ __launch_bounds__(kBlockSize, 8)
// ----------------------------------------------------------------------------
// Indexer variant: kHeadDim = 128, 1 token per *warp* (8 tokens per block).
// Each warp's 32 lanes cover the full 128-elem head_dim (kVecSize = 4 each).
// Cache layout: 132 bytes/token (128 fp8 nope + 4 fp32 scale).
// ----------------------------------------------------------------------------
template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL>
INDEXER_KERNEL void fused_norm_rope_indexer(const __grid_constant__ FusedNormRopeStoreParams params) {
using namespace device;
using enum ForwardMode;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
constexpr int64_t kVecSize = 4;
constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
constexpr int64_t kPageBytes = 132ll << kPageBits;
static_assert(kHeadDim == kWarpThreads * kVecSize);
static_assert(kRopeDim == kWarpThreads * 2);
static_assert(kRopeSize <= kWarpThreads);
using Storage = AlignedVector<DType, kVecSize>;
using Float4 = AlignedVector<float, kVecSize>;
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
const auto work_id = blockIdx.x * kNumWarps + warp_id;
// Lanes whose 4-elem pack lies in the rope tail (= last `kRopeSize` packs).
const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
if (work_id >= params.num_tokens) return;
const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim;
int32_t position;
int64_t out_loc;
if constexpr (kMode == CompressExtend) {
const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
if (plan.is_invalid()) return;
position = plan.seq_len - params.compress_ratio;
out_loc = params.out_loc[plan.ragged_id];
} else if constexpr (kMode == CompressDecode) {
const auto plan = static_cast<const PlanD*>(params.handle)[work_id];
if (plan.seq_len % params.compress_ratio != 0) return;
position = plan.seq_len - params.compress_ratio;
out_loc = params.out_loc[work_id];
} else {
static_assert(host::dependent_false_v<DType>, "Unsupported Mode");
}
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float4 data, freq;
// part 1: norm
{
Storage input_vec, weight_vec;
input_vec.load(input, lane_id);
weight_vec.load(params.weight, lane_id);
if (is_rope_lane) freq.load(freqs_cis, lane_id - (kWarpThreads - kRopeSize));
float sum_of_squares = 0.0f;
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const auto fp32_input = cast<float>(input_vec[i]);
sum_of_squares += fp32_input * fp32_input;
}
sum_of_squares = warp::reduce_sum(sum_of_squares);
const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const auto fp32_input = cast<float>(input_vec[i]);
const auto fp32_weight = cast<float>(weight_vec[i]);
data[i] = fp32_input * norm_factor * fp32_weight;
}
}
// part 2: rope (rope-lane only, 4 elems per lane = 2 (real, imag) pairs)
if (is_rope_lane) {
const auto x_real = data[0];
const auto x_imag = data[1];
const auto y_real = data[2];
const auto y_imag = data[3];
const auto freq_x_real = freq[0];
const auto freq_x_imag = freq[1];
const auto freq_y_real = freq[2];
const auto freq_y_imag = freq[3];
data[0] = x_real * freq_x_real - x_imag * freq_x_imag;
data[1] = x_real * freq_x_imag + x_imag * freq_x_real;
data[2] = y_real * freq_y_real - y_imag * freq_y_imag;
data[3] = y_real * freq_y_imag + y_imag * freq_y_real;
}
// part 3: hadamard transform
{
// Stage 1: butterfly (data[0], data[1]) and (data[2], data[3]).
{
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
data[0] = a0 + a1;
data[1] = a0 - a1;
data[2] = a2 + a3;
data[3] = a2 - a3;
}
// Stage 2: butterfly (data[0], data[2]) and (data[1], data[3]).
{
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
data[0] = a0 + a2;
data[1] = a1 + a3;
data[2] = a0 - a2;
data[3] = a1 - a3;
}
// Stages 3..7: cross-lane butterflies. Lower-lane (mask bit clear) keeps
// the sum, upper-lane (mask bit set) keeps the difference. shfl_xor is
// unsynchronized across early-returned lanes, but invalid-plan returns
// happen above for *all* lanes of a warp (work_id is warp-uniform), so
// the warp is intact here.
#pragma unroll
for (uint32_t mask = 1; mask < kWarpThreads; mask <<= 1) {
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
#ifndef USE_ROCM
const float other = __shfl_xor_sync(kFullMask, data[i], mask, kWarpThreads);
#else
const float other = __shfl_xor(data[i], mask, kWarpThreads);
#endif
data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other);
}
}
const float kHadamardScale = math::rsqrt(static_cast<float>(kHeadDim));
#pragma unroll
for (int i = 0; i < kVecSize; ++i)
data[i] *= kHadamardScale;
}
// part 4: per-warp UE8M0 quant + store. The whole warp emits one fp8 group
// (= 128 elements) plus a single fp32 scale, matching the indexer cache
// layout (`fused_store_indexer_cache`).
{
using OutStorage = AlignedVector<fp8x2_e4m3_t, 2>;
float local_max = math::abs(data[0]);
#pragma unroll
for (int i = 1; i < kVecSize; ++i) {
local_max = math::max(local_max, math::abs(data[i]));
}
const auto abs_max = warp::reduce_max(local_max);
const auto scale = fmaxf(1e-4f, abs_max) / kFP8E4M3Max;
const auto inv_scale = 1.0f / scale;
const int64_t page = out_loc >> kPageBits;
const int64_t offset = out_loc & ((1 << kPageBits) - 1);
const auto page_ptr = params.kvcache + page * kPageBytes;
const auto value_ptr = page_ptr + offset * 128;
const auto scale_ptr = page_ptr + (128 << kPageBits) + offset * 4;
OutStorage result;
result[0] = pack_fp8(data[0] * inv_scale, data[1] * inv_scale);
result[1] = pack_fp8(data[2] * inv_scale, data[3] * inv_scale);
PDLTriggerSecondary<kUsePDL>();
result.store(value_ptr, lane_id);
// The single fp32 scale is identical across all lanes -- write from any lane.
if (lane_id == 0) reinterpret_cast<float*>(scale_ptr)[0] = scale;
}
}
template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL>
INDEXER_KERNEL void fused_norm_rope_indexer_fp4(const __grid_constant__ FusedNormRopeStoreParams params) {
using namespace device;
using enum ForwardMode;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
constexpr int64_t kVecSize = 4;
constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
constexpr int64_t kPageBytes = 68ll << kPageBits;
static_assert(kHeadDim == kWarpThreads * kVecSize);
static_assert(kRopeDim == kWarpThreads * 2);
static_assert(kRopeSize <= kWarpThreads);
using Storage = AlignedVector<DType, kVecSize>;
using Float4 = AlignedVector<float, kVecSize>;
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
const auto work_id = blockIdx.x * kNumWarps + warp_id;
const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
if (work_id >= params.num_tokens) return;
const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim;
int32_t position;
int64_t out_loc;
if constexpr (kMode == CompressExtend) {
const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
if (plan.is_invalid()) return;
position = plan.seq_len - params.compress_ratio;
out_loc = params.out_loc[plan.ragged_id];
} else if constexpr (kMode == CompressDecode) {
const auto plan = static_cast<const PlanD*>(params.handle)[work_id];
if (plan.seq_len % params.compress_ratio != 0) return;
position = plan.seq_len - params.compress_ratio;
out_loc = params.out_loc[work_id];
} else {
static_assert(host::dependent_false_v<DType>, "Unsupported Mode");
}
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float4 data, freq;
{
Storage input_vec, weight_vec;
input_vec.load(input, lane_id);
weight_vec.load(params.weight, lane_id);
if (is_rope_lane) freq.load(freqs_cis, lane_id - (kWarpThreads - kRopeSize));
float sum_of_squares = 0.0f;
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const auto fp32_input = cast<float>(input_vec[i]);
sum_of_squares += fp32_input * fp32_input;
}
sum_of_squares = warp::reduce_sum(sum_of_squares);
const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const auto fp32_input = cast<float>(input_vec[i]);
const auto fp32_weight = cast<float>(weight_vec[i]);
data[i] = fp32_input * norm_factor * fp32_weight;
}
}
if (is_rope_lane) {
const auto x_real = data[0];
const auto x_imag = data[1];
const auto y_real = data[2];
const auto y_imag = data[3];
const auto freq_x_real = freq[0];
const auto freq_x_imag = freq[1];
const auto freq_y_real = freq[2];
const auto freq_y_imag = freq[3];
data[0] = x_real * freq_x_real - x_imag * freq_x_imag;
data[1] = x_real * freq_x_imag + x_imag * freq_x_real;
data[2] = y_real * freq_y_real - y_imag * freq_y_imag;
data[3] = y_real * freq_y_imag + y_imag * freq_y_real;
}
{
{
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
data[0] = a0 + a1;
data[1] = a0 - a1;
data[2] = a2 + a3;
data[3] = a2 - a3;
}
{
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
data[0] = a0 + a2;
data[1] = a1 + a3;
data[2] = a0 - a2;
data[3] = a1 - a3;
}
#pragma unroll
for (uint32_t mask = 1; mask < kWarpThreads; mask <<= 1) {
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
#ifndef USE_ROCM
const float other = __shfl_xor_sync(kFullMask, data[i], mask, kWarpThreads);
#else
const float other = __shfl_xor(data[i], mask, kWarpThreads);
#endif
data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other);
}
}
const float kHadamardScale = math::rsqrt(static_cast<float>(kHeadDim));
#pragma unroll
for (int i = 0; i < kVecSize; ++i)
data[i] *= kHadamardScale;
}
{
float local_max = math::abs(data[0]);
#pragma unroll
for (int i = 1; i < kVecSize; ++i) {
local_max = math::max(local_max, math::abs(data[i]));
}
local_max = warp::reduce_max<8>(local_max);
const auto scale_raw = fmaxf(1e-4f, local_max) / 6.0f;
const auto scale_ue8m0 = static_cast<uint8_t>(cast_to_ue8m0(scale_raw));
const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
const uint8_t packed0 = quant_fp4_e2m1(data[0] * inv_scale) | (quant_fp4_e2m1(data[1] * inv_scale) << 4);
const uint8_t packed1 = quant_fp4_e2m1(data[2] * inv_scale) | (quant_fp4_e2m1(data[3] * inv_scale) << 4);
const uint16_t packed = static_cast<uint16_t>(packed0) | (static_cast<uint16_t>(packed1) << 8);
const int64_t page = out_loc >> kPageBits;
const int64_t offset = out_loc & ((1 << kPageBits) - 1);
const auto page_ptr = params.kvcache + page * kPageBytes;
const auto value_ptr = page_ptr + offset * 64;
const auto scale_ptr = page_ptr + (64 << kPageBits) + offset * 4;
PDLTriggerSecondary<kUsePDL>();
reinterpret_cast<uint16_t*>(value_ptr)[lane_id] = packed;
if ((lane_id & 7) == 0) static_cast<uint8_t*>(scale_ptr)[lane_id >> 3] = scale_ue8m0;
}
}
// ----------------------------------------------------------------------------
// FlashMLA variant: kHeadDim = 512, 1 token per *block* (256 threads).
// Each thread loads kVecSize=2 BF16, so 256 threads cover the full 512 elems.
// Cache layout: 584 bytes/token = 448 fp8 nope + 64 (=32 bf16x2) rope + 8 scale.
// ----------------------------------------------------------------------------
template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL, bool kBf16Store = false>
FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormRopeStoreParams params) {
using namespace device;
using enum ForwardMode;
constexpr int64_t kHeadDim = 512;
constexpr int64_t kRopeDim = 64;
constexpr int64_t kVecSize = 2;
// Last warp owns the rope tail. The remaining 7 warps each emit one
// 64-element fp8 group (own UE8M0 scale).
constexpr uint32_t kRopeWarp = kNumWarps - 1;
// kBf16Store: write the whole head_dim as plain BF16 (no fp8 / no scale) into a
// [num_slots, head_dim] bf16 cache (page_size==1) at row out_loc
constexpr int64_t kPageBytes =
kBf16Store ? ((kHeadDim * 2ll) << kPageBits) : host::div_ceil(584ll << kPageBits, 576) * 576;
static_assert(kHeadDim == kBlockSize * kVecSize);
static_assert(kRopeDim == kWarpThreads * kVecSize);
static_assert(kHeadDim - kRopeDim == kRopeWarp * kWarpThreads * kVecSize);
using Storage = AlignedVector<DType, kVecSize>;
using Float2 = AlignedVector<float, kVecSize>;
const auto tx = threadIdx.x;
const auto warp_id = tx / kWarpThreads;
const auto lane_id = tx % kWarpThreads;
const auto work_id = blockIdx.x;
if (work_id >= params.num_tokens) return;
const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim;
int32_t position;
int64_t out_loc;
if constexpr (kMode == CompressExtend) {
const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
if (plan.is_invalid()) return;
position = plan.seq_len - params.compress_ratio;
out_loc = params.out_loc[plan.ragged_id];
} else if constexpr (kMode == CompressDecode) {
const auto plan = static_cast<const PlanD*>(params.handle)[work_id];
if (plan.seq_len % params.compress_ratio != 0) return;
position = plan.seq_len - params.compress_ratio;
out_loc = params.out_loc[work_id];
} else {
static_assert(host::dependent_false_v<DType>, "Unsupported Mode");
}
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float2 data, freq;
// part 1: norm. Each thread owns one 2-elem pack (`tx`-th pack of input).
// Sum of squares is reduced across the whole block via per-warp partials.
{
__shared__ float partial_sums[kNumWarps];
Storage input_vec, weight_vec;
input_vec.load(input, tx);
weight_vec.load(params.weight, tx);
if (warp_id == kRopeWarp) freq.load(freqs_cis, lane_id);
float sum_of_squares = 0.0f;
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const auto fp32_input = cast<float>(input_vec[i]);
sum_of_squares += fp32_input * fp32_input;
}
const auto warp_sum = warp::reduce_sum(sum_of_squares);
if (lane_id == 0) partial_sums[warp_id] = warp_sum;
__syncthreads();
// Replicate the per-warp partial sums to a full warp and reduce. Every
// lane-group of `kNumWarps` lanes ends up with the global sum.
sum_of_squares = warp::reduce_sum<kNumWarps>(partial_sums[lane_id % kNumWarps]);
const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const auto fp32_input = cast<float>(input_vec[i]);
const auto fp32_weight = cast<float>(weight_vec[i]);
data[i] = fp32_input * norm_factor * fp32_weight;
}
}
const int64_t page = out_loc >> kPageBits;
const int64_t offset = out_loc & ((1 << kPageBits) - 1);
const auto page_ptr = params.kvcache + page * kPageBytes;
const auto value_ptr = page_ptr + offset * (kBf16Store ? (kHeadDim * 2) : 576);
PDLTriggerSecondary<kUsePDL>();
// part 2: rope on the rope warp (BF16 store), or per-warp FP8 quant + store.
if constexpr (kBf16Store) {
Float2 d = data;
if (warp_id == kRopeWarp) {
const auto x_real = data[0];
const auto x_imag = data[1];
const auto freq_real = freq[0];
const auto freq_imag = freq[1];
d[0] = x_real * freq_real - x_imag * freq_imag;
d[1] = x_real * freq_imag + x_imag * freq_real;
}
reinterpret_cast<bf16x2_t*>(value_ptr)[tx] = cast<bf16x2_t>(fp32x2_t{d[0], d[1]});
} else if (warp_id == kRopeWarp) {
// Each rope-warp lane owns exactly one (real, imag) pair within the rope
// tail. Apply rotation, downcast to BF16, write to the slot's rope region.
const auto x_real = data[0];
const auto x_imag = data[1];
const auto freq_real = freq[0];
const auto freq_imag = freq[1];
data[0] = x_real * freq_real - x_imag * freq_imag;
data[1] = x_real * freq_imag + x_imag * freq_real;
const auto result = cast<bf16x2_t>(fp32x2_t{data[0], data[1]});
const auto rope_ptr = value_ptr + 448;
reinterpret_cast<bf16x2_t*>(rope_ptr)[lane_id] = result;
} else {
// Non-rope warp: per-warp UE8M0 group (64 elems -> 64 fp8 + 1 scale byte).
// BF16 round-trip to match the precision of the non-fused path
// (which goes through quant_to_nope_fp8_rope_bf16_pack_triton with bf16 input).
const auto x = cast<float>(cast<bf16_t>(data[0]));
const auto y = cast<float>(cast<bf16_t>(data[1]));
const auto abs_max = warp::reduce_max(fmaxf(fabs(x), fabs(y)));
const auto scale_raw = fmaxf(1e-4f, abs_max) / kFP8E4M3Max;
const auto scale_ue8m0 = cast_to_ue8m0(scale_raw);
const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
const auto result = pack_fp8(x * inv_scale, y * inv_scale);
const auto scale_ptr = page_ptr + (576 << kPageBits) + offset * 8;
reinterpret_cast<fp8x2_e4m3_t*>(value_ptr)[tx] = result;
// All lanes in this warp produce the same scale byte; let lane 0 publish.
if (lane_id == 0) static_cast<uint8_t*>(scale_ptr)[warp_id] = scale_ue8m0;
}
}
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, bool kUsePDL, bool kBf16Store = false>
struct FusedNormRopeKernel {
static constexpr int32_t kLogPageSize = std::countr_zero(kPageSize);
static constexpr bool kIsIndexer = (kHeadDim == 128);
static_assert(!(kIsIndexer && kBf16Store), "bf16 store only for flashmla head_dim=512");
static constexpr int64_t kIndexerBytes = 132 * kPageSize;
static constexpr int64_t kFlashMLABytes = host::div_ceil(584 * kPageSize, 576) * 576;
static constexpr int64_t kBf16Bytes = kHeadDim * 2 * kPageSize; // plain bf16 cache
static constexpr int64_t kPageBytes = kBf16Store ? kBf16Bytes : (kIsIndexer ? kIndexerBytes : kFlashMLABytes);
/// TODO: Let's fix the config for now.
static_assert(kRopeDim == 64 && (kHeadDim == 128 || kHeadDim == 512));
static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
template <ForwardMode kMode>
static constexpr auto select_kernel() {
if constexpr (kIsIndexer) {
return fused_norm_rope_indexer<DType, kMode, kLogPageSize, kUsePDL>;
} else {
return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kUsePDL, kBf16Store>;
}
}
template <ForwardMode kMode>
static constexpr auto select_fp4_kernel() {
static_assert(kIsIndexer, "FP4 fused store is only defined for the indexer");
return fused_norm_rope_indexer_fp4<DType, kMode, kLogPageSize, kUsePDL>;
}
static void forward(
const tvm::ffi::TensorView input,
const tvm::ffi::TensorView plan,
const tvm::ffi::TensorView weight,
const float eps,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView out_loc,
const tvm::ffi::TensorView kvcache,
const bool is_decode,
const uint32_t compress_ratio) {
using namespace host;
using enum ForwardMode;
const auto mode = static_cast<ForwardMode>(is_decode);
auto N = SymbolicSize{"num_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLGPU>();
TensorMatcher({N, kHeadDim}) // input
.with_dtype<DType>()
.with_device(device_)
.verify(input);
TensorMatcher({kHeadDim}) // weight
.with_dtype<DType>()
.with_device(device_)
.verify(weight);
TensorMatcher({-1, kRopeDim}) // freqs_cis
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
TensorMatcher({-1}) // out_loc
.with_dtype<int64_t>()
.with_device(device_)
.verify(out_loc);
TensorMatcher({-1, -1}) // cache
.with_strides({kPageBytes, 1})
.with_dtype<uint8_t>()
.with_device(device_)
.verify(kvcache);
switch (mode) {
case CompressExtend:
compress::verify_plan_c(plan, N, device_);
RuntimeCheck(out_loc.size(0) >= N.unwrap());
break;
case CompressDecode:
compress::verify_plan_d(plan, N, device_);
RuntimeCheck(out_loc.size(0) == N.unwrap());
break;
}
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
if (num_tokens == 0) return;
const auto params = FusedNormRopeStoreParams{
.input = input.data_ptr(),
.handle = plan.data_ptr(),
.weight = weight.data_ptr(),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.out_loc = static_cast<const int64_t*>(out_loc.data_ptr()),
.kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
.eps = eps,
.compress_ratio = compress_ratio,
.num_tokens = num_tokens,
};
// Indexer packs `kNumWarps` tokens per block (warp-major); FlashMLA uses
// a whole block per token (cta-major sum-reduce over head_dim=512).
const uint32_t num_blocks = kIsIndexer ? div_ceil(num_tokens, kNumWarps) : num_tokens;
const auto device = device_.unwrap();
const auto kernel = mode == CompressExtend ? select_kernel<CompressExtend>() : select_kernel<CompressDecode>();
LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
}
static void forward_fp4(
const tvm::ffi::TensorView input,
const tvm::ffi::TensorView plan,
const tvm::ffi::TensorView weight,
const float eps,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView out_loc,
const tvm::ffi::TensorView kvcache,
const bool is_decode,
const uint32_t compress_ratio) {
using namespace host;
using enum ForwardMode;
static_assert(kIsIndexer, "FP4 fused store is only defined for the indexer");
constexpr int64_t kFp4PageBytes = 68 * kPageSize;
const auto mode = static_cast<ForwardMode>(is_decode);
auto N = SymbolicSize{"num_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({N, kHeadDim}).with_dtype<DType>().with_device(device_).verify(input);
TensorMatcher({kHeadDim}).with_dtype<DType>().with_device(device_).verify(weight);
TensorMatcher({-1, kRopeDim}).with_dtype<float>().with_device(device_).verify(freqs_cis);
TensorMatcher({-1}).with_dtype<int64_t>().with_device(device_).verify(out_loc);
TensorMatcher({-1, -1}).with_strides({kFp4PageBytes, 1}).with_dtype<uint8_t>().with_device(device_).verify(kvcache);
switch (mode) {
case CompressExtend:
compress::verify_plan_c(plan, N, device_);
RuntimeCheck(out_loc.size(0) >= N.unwrap());
break;
case CompressDecode:
compress::verify_plan_d(plan, N, device_);
RuntimeCheck(out_loc.size(0) == N.unwrap());
break;
}
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
if (num_tokens == 0) return;
const auto params = FusedNormRopeStoreParams{
.input = input.data_ptr(),
.handle = plan.data_ptr(),
.weight = weight.data_ptr(),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.out_loc = static_cast<const int64_t*>(out_loc.data_ptr()),
.kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
.eps = eps,
.compress_ratio = compress_ratio,
.num_tokens = num_tokens,
};
const uint32_t num_blocks = div_ceil(num_tokens, kNumWarps);
const auto device = device_.unwrap();
const auto kernel =
mode == CompressExtend ? select_fp4_kernel<CompressExtend>() : select_fp4_kernel<CompressDecode>();
LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh
+214
View File
@@ -0,0 +1,214 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/warp.cuh>
#include <tvm/ffi/container/tensor.h>
#include <cmath>
#include <cstdint>
namespace {
[[maybe_unused]]
SGL_DEVICE float act_sqrt_softplus(float x) {
const float softplus = fmaxf(x, 0.0f) + log1pf(expf(-fabsf(x)));
return sqrtf(softplus);
}
struct MoEHashTopKParams {
const float* __restrict__ router_logits;
const int64_t* __restrict__ input_id;
const int32_t* __restrict__ tid2eid;
int32_t* __restrict__ topk_ids;
float* __restrict__ topk_weights;
uint32_t num_tokens;
uint32_t topk;
uint32_t num_routed_experts;
uint32_t num_shared_experts;
float routed_scaling_factor;
};
template <auto Fn, bool kUsePDL>
__global__ void moe_hash_topk_fused(const MoEHashTopKParams __grid_constant__ params) {
using namespace device;
const auto& [
router_logits, input_id, tid2eid, topk_ids, topk_weights, // pointers
num_tokens, topk, num_routed_experts, num_shared_experts, routed_scaling_factor] =
params;
const uint32_t topk_fused = topk + num_shared_experts;
const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t warp_id = tid / kWarpThreads;
const uint32_t lane_id = tid % kWarpThreads;
if (warp_id >= num_tokens) return;
// we can safely prefetch the token id
const auto token_id = input_id[warp_id];
PDLWaitPrimary<kUsePDL>();
float routed_weight = 0.0f;
int32_t expert_id = 0;
if (lane_id < topk) {
expert_id = tid2eid[token_id * topk + lane_id];
routed_weight = Fn(router_logits[warp_id * num_routed_experts + expert_id]);
}
const auto routed_sum = device::warp::reduce_sum(routed_weight);
if (lane_id < topk_fused) {
const bool is_shared = lane_id >= topk;
const auto output_offset = warp_id * topk_fused + lane_id;
topk_ids[output_offset] = is_shared ? num_routed_experts + lane_id - topk : expert_id;
topk_weights[output_offset] = is_shared ? 1.0f / routed_scaling_factor : routed_weight / routed_sum;
}
PDLTriggerSecondary<kUsePDL>();
}
struct TopKParams {
int32_t* __restrict__ topk_ids;
// Exactly one is active: ntn_ptr == nullptr means use ntn_value.
const int32_t* __restrict__ ntn_ptr;
int32_t ntn_value;
int64_t stride;
uint32_t topk;
uint32_t num_tokens;
};
__global__ void mask_topk_ids_padded_region(const TopKParams __grid_constant__ params) {
const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t warp_id = tid / device::kWarpThreads;
const uint32_t lane_id = tid % device::kWarpThreads;
if (warp_id >= params.num_tokens || lane_id >= params.topk) return;
device::PDLWaitPrimary<true>();
const uint32_t num = (params.ntn_ptr != nullptr) //
? static_cast<uint32_t>(params.ntn_ptr[0])
: static_cast<uint32_t>(params.ntn_value);
if (warp_id >= num) params.topk_ids[warp_id * params.stride + lane_id] = -1;
device::PDLTriggerSecondary<true>();
}
template <auto Fn, bool kUsePDL>
struct HashTopKKernel {
static constexpr auto kernel = moe_hash_topk_fused<Fn, kUsePDL>;
static void
run(const tvm::ffi::TensorView router_logits,
const tvm::ffi::TensorView input_id,
const tvm::ffi::TensorView tid2eid,
const tvm::ffi::TensorView topk_weights,
const tvm::ffi::TensorView topk_ids,
float routed_scaling_factor) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto E = SymbolicSize{"num_routed_experts"};
auto K = SymbolicSize{"topk_fused"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({N, E}) //
.with_dtype<float>()
.with_device(device)
.verify(router_logits);
TensorMatcher({N}) //
.with_dtype<int64_t>()
.with_device(device)
.verify(input_id);
TensorMatcher({-1, -1}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(tid2eid);
TensorMatcher({N, K}) //
.with_dtype<float>()
.with_device(device)
.verify(topk_weights);
TensorMatcher({N, K}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(topk_ids);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto topk_fused = static_cast<uint32_t>(K.unwrap());
const auto topk = static_cast<uint32_t>(tid2eid.size(1));
const auto shared_experts = topk_fused - topk;
RuntimeCheck(topk <= topk_fused, "HashTopKKernel requires topk <= topk_fused");
RuntimeCheck(topk_fused <= device::kWarpThreads, "HashTopKKernel requires topk_fused <= warp size");
const auto params = MoEHashTopKParams{
.router_logits = static_cast<const float*>(router_logits.data_ptr()),
.input_id = static_cast<const int64_t*>(input_id.data_ptr()),
.tid2eid = static_cast<const int32_t*>(tid2eid.data_ptr()),
.topk_ids = static_cast<int32_t*>(topk_ids.data_ptr()),
.topk_weights = static_cast<float*>(topk_weights.data_ptr()),
.num_tokens = num_tokens,
.topk = topk,
.num_routed_experts = static_cast<uint32_t>(E.unwrap()),
.num_shared_experts = shared_experts,
.routed_scaling_factor = routed_scaling_factor,
};
const auto kBlockSize = 128u;
const auto kNumWarps = kBlockSize / device::kWarpThreads;
const auto num_blocks = div_ceil(num_tokens, kNumWarps);
LaunchKernel(num_blocks, kBlockSize, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
// TODO this may not be related to *hash* topk, thus may move
struct MaskKernel {
static constexpr auto kernel = mask_topk_ids_padded_region;
static void run(tvm::ffi::TensorView topk_ids, tvm::ffi::TensorView num_token_non_padded) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto K = SymbolicSize{"topk"};
auto D = SymbolicSize{"stride"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({N, K}) //
.with_strides({D, 1})
.with_dtype<int32_t>()
.with_device(device)
.verify(topk_ids);
RuntimeCheck(num_token_non_padded.numel() == 1, "num_token_non_padded should be a scalar");
RuntimeCheck(K.unwrap() <= device::kWarpThreads, "MaskKernel requires topk <= warp size");
const int32_t* ntn_ptr = nullptr;
int32_t ntn_value = 0;
const auto ntn_dev = num_token_non_padded.device().device_type;
if (ntn_dev == kDLCUDA) {
RuntimeCheck(is_type<int32_t>(num_token_non_padded.dtype()), "num_token_non_padded on CUDA must be int32");
ntn_ptr = static_cast<const int32_t*>(num_token_non_padded.data_ptr());
} else if (ntn_dev == kDLCPU) {
if (is_type<int32_t>(num_token_non_padded.dtype())) {
ntn_value = *static_cast<const int32_t*>(num_token_non_padded.data_ptr());
} else if (is_type<int64_t>(num_token_non_padded.dtype())) {
ntn_value = static_cast<int32_t>(*static_cast<const int64_t*>(num_token_non_padded.data_ptr()));
} else {
RuntimeCheck(false, "num_token_non_padded on CPU must be int32 or int64");
}
} else {
RuntimeCheck(false, "num_token_non_padded must be on CPU or CUDA");
}
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto params = TopKParams{
.topk_ids = static_cast<int32_t*>(topk_ids.data_ptr()),
.ntn_ptr = ntn_ptr,
.ntn_value = ntn_value,
.stride = static_cast<int64_t>(D.unwrap()),
.topk = static_cast<uint32_t>(K.unwrap()),
.num_tokens = num_tokens,
};
const auto kBlockSize = 128u;
const auto kNumWarps = kBlockSize / device::kWarpThreads;
const auto num_blocks = div_ceil(num_tokens, kNumWarps);
LaunchKernel(num_blocks, kBlockSize, device.unwrap()) //
.enable_pdl(true)(kernel, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh
+882
View File
@@ -0,0 +1,882 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <tvm/ffi/container/tensor.h>
#include <bit>
#include <cstdint>
namespace {
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::inv_scale_ue8m0;
using deepseek_v4::fp8::pack_fp8;
SGL_DEVICE uint8_t quant_fp4_e2m1(float x) {
const float ax = fminf(fabsf(x), 6.0f);
uint8_t idx = 0;
idx += ax > 0.25f;
idx += ax > 0.75f;
idx += ax > 1.25f;
idx += ax > 1.75f;
idx += ax > 2.5f;
idx += ax > 3.5f;
idx += ax > 5.0f;
if (x < 0.0f && idx != 0) idx |= 0x8;
return idx;
}
// 4 warps per block: warp-per-(token, head) work-item dispatch (Q kernel).
constexpr uint32_t kFusedQBlockSize = 128;
constexpr uint32_t kFusedQNumWarps = kFusedQBlockSize / device::kWarpThreads;
// 8 warps per block: block-per-token work-item dispatch (K kernel).
constexpr uint32_t kFusedKBlockSize = 256;
constexpr uint32_t kFusedKNumWarps = kFusedKBlockSize / device::kWarpThreads;
#define Q_KERNEL __global__ __launch_bounds__(kFusedQBlockSize, 16)
#define K_KERNEL __global__ __launch_bounds__(kFusedKBlockSize, 8)
template <int64_t kRopeDim>
SGL_DEVICE device::AlignedVector<float, 4>
load_rope_first_cos_sin(const float* __restrict__ cos_sin_cache, int32_t lane_id) {
constexpr int64_t kHalfRopeDim = kRopeDim / 2;
const int32_t pair0 = lane_id * 2;
const int32_t pair1 = pair0 + 1;
device::AlignedVector<float, 4> freq;
freq[0] = cos_sin_cache[pair0];
freq[1] = cos_sin_cache[kHalfRopeDim + pair0];
freq[2] = cos_sin_cache[pair1];
freq[3] = cos_sin_cache[kHalfRopeDim + pair1];
return freq;
}
// ============================================================================
// Q kernel: warp-per-(token, head) rmsnorm-self + RoPE + write to q_out.
// ============================================================================
struct FusedQNormRopeParams {
const void* __restrict__ q_input; // (B, num_q_heads, kHeadDim) DType
void* __restrict__ q_output; // (B, num_q_heads, kHeadDim) DType
const float* __restrict__ freqs_cis; // (max_pos, kRopeDim) fp32 (re/im interleaved)
const void* __restrict__ positions; // (B,) PosT
int64_t q_input_stride_batch;
int64_t q_output_stride_batch;
uint32_t batch_size;
uint32_t num_q_heads;
float eps;
};
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, typename PosT, bool kUsePDL>
Q_KERNEL void fused_q_norm_rope(const __grid_constant__ FusedQNormRopeParams params) {
using namespace device;
constexpr int64_t kMaxVecSize = 16 / sizeof(DType);
constexpr int64_t kVecSize = std::min(kMaxVecSize, kHeadDim / kWarpThreads);
constexpr int64_t kLocalSize = kHeadDim / (kWarpThreads * kVecSize);
constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
static_assert(kHeadDim % (kWarpThreads * kVecSize) == 0);
static_assert(kLocalSize * kVecSize * kWarpThreads == kHeadDim);
static_assert(kRopeDim % kVecSize == 0);
static_assert(kRopeSize <= kWarpThreads);
static_assert(kRopeDim == kWarpThreads * 2, "1 (real, imag) pair per lane");
using Storage = AlignedVector<DType, kVecSize>;
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id;
const uint32_t total_works = params.batch_size * params.num_q_heads;
if (work_id >= total_works) return;
const uint32_t batch_id = work_id / params.num_q_heads;
const uint32_t head_id = work_id % params.num_q_heads;
const auto input_ptr =
static_cast<const DType*>(params.q_input) + batch_id * params.q_input_stride_batch + head_id * kHeadDim;
const auto output_ptr =
static_cast<DType*>(params.q_output) + batch_id * params.q_output_stride_batch + head_id * kHeadDim;
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[batch_id]);
__shared__ Storage s_rope[kFusedQNumWarps][kRopeSize];
// Prefetch this lane's freq pair before the PDL gate so the wait happens
// outside the dependency chain on `position`.
const auto mem_freq = tile::Memory<fp32x2_t>{lane_id, kWarpThreads};
PDLWaitPrimary<kUsePDL>();
// part 1: rmsnorm-self (no weight).
const auto gmem = tile::Memory<Storage>{lane_id, kWarpThreads};
Storage input_vec[kLocalSize];
#pragma unroll
for (int i = 0; i < kLocalSize; ++i) {
input_vec[i] = gmem.load(input_ptr, i);
}
const auto freq = mem_freq.load(params.freqs_cis + position * kRopeDim);
float sum_of_squares = 0.0f;
#pragma unroll
for (int i = 0; i < kLocalSize; ++i) {
#pragma unroll
for (int j = 0; j < kVecSize; ++j) {
const auto x = cast<float>(input_vec[i][j]);
sum_of_squares += x * x;
}
}
sum_of_squares = warp::reduce_sum(sum_of_squares);
const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
#pragma unroll
for (int i = 0; i < kLocalSize; ++i) {
#pragma unroll
for (int j = 0; j < kVecSize; ++j) {
const auto x = cast<float>(input_vec[i][j]);
input_vec[i][j] = cast<DType>(x * norm_factor);
}
}
// Stash the rope tail (last kRopeSize lanes' last tile) into shared memory;
// write nope tiles to gmem directly.
const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
#pragma unroll
for (int i = 0; i < kLocalSize; ++i) {
if (i == kLocalSize - 1 && is_rope_lane) {
const auto rope_id = lane_id - (kWarpThreads - kRopeSize);
s_rope[warp_id][rope_id] = input_vec[i];
} else {
gmem.store(output_ptr, input_vec[i], i);
}
}
__syncwarp();
PDLTriggerSecondary<kUsePDL>();
// part 2: RoPE on all 32 lanes -- one (real, imag) bf16x2 pair per lane.
using DType2 = packed_t<DType>;
const auto mem_elem = tile::Memory<DType2>{lane_id, kWarpThreads};
const auto elem = mem_elem.load(s_rope[warp_id]);
const auto [x_real, x_imag] = cast<fp32x2_t>(elem);
const auto [freq_real, freq_imag] = freq;
const fp32x2_t rotated = {
x_real * freq_real - x_imag * freq_imag,
x_real * freq_imag + x_imag * freq_real,
};
mem_elem.store(output_ptr + (kHeadDim - kRopeDim), cast<DType2>(rotated));
}
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, bool kUsePDL>
struct FusedQNormRopeKernel {
template <typename PosT>
static constexpr auto kernel = fused_q_norm_rope<DType, kHeadDim, kRopeDim, PosT, kUsePDL>;
static void forward(
const tvm::ffi::TensorView q_input,
const tvm::ffi::TensorView q_output,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView positions,
float eps) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto H = SymbolicSize{"num_q_heads"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B, H, kHeadDim}) //
.with_strides({-1, kHeadDim, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(q_input);
TensorMatcher({B, H, kHeadDim}) //
.with_strides({-1, kHeadDim, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(q_output);
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
.with_device(device_)
.verify(positions);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto num_q_heads = static_cast<uint32_t>(H.unwrap());
if (batch_size == 0) return;
const auto params = FusedQNormRopeParams{
.q_input = q_input.data_ptr(),
.q_output = q_output.data_ptr(),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.positions = positions.data_ptr(),
.q_input_stride_batch = q_input.stride(0),
.q_output_stride_batch = q_output.stride(0),
.batch_size = batch_size,
.num_q_heads = num_q_heads,
.eps = eps,
};
const auto total_works = batch_size * num_q_heads;
const auto num_blocks = div_ceil(total_works, kFusedQNumWarps);
const auto k_int32 = kernel<int32_t>;
const auto k_int64 = kernel<int64_t>;
const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
LaunchKernel(num_blocks, kFusedQBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(k, params);
}
};
// ============================================================================
// K kernel: block-per-token rmsnorm (with kv_weight) + RoPE + FlashMLA store.
// ============================================================================
struct FusedKNormRopeFlashMLAParams {
const void* __restrict__ kv; // (B, kHeadDim) DType
const void* __restrict__ kv_weight; // (kHeadDim,) DType
const float* __restrict__ freqs_cis; // (max_pos, kRopeDim) fp32
const void* __restrict__ positions; // (B,) PosT
const int32_t* __restrict__ out_loc; // (B,) int32 -> cache slot id
uint8_t* __restrict__ kvcache; // (npages, kPageBytes) uint8
// Row stride for `kv` in elements. Required because the upstream caller often
// passes `qkv_a[..., q_lora_rank:]`, a non-contiguous slice whose stride[0]
// equals `q_lora_rank + kHeadDim` rather than `kHeadDim`.
int64_t kv_stride_batch;
uint32_t batch_size;
float eps;
};
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, typename PosT, int32_t kPageBits, bool kUsePDL>
K_KERNEL void fused_k_norm_rope_flashmla(const __grid_constant__ FusedKNormRopeFlashMLAParams params) {
using namespace device;
constexpr int64_t kVecSize = 2;
constexpr uint32_t kRopeWarp = kFusedKNumWarps - 1;
constexpr int64_t kPageBytes = host::div_ceil(584ll << kPageBits, 576) * 576;
static_assert(kHeadDim == kFusedKBlockSize * kVecSize);
static_assert(kRopeDim == kWarpThreads * kVecSize);
static_assert(kHeadDim - kRopeDim == kRopeWarp * kWarpThreads * kVecSize);
using Storage = AlignedVector<DType, kVecSize>;
using Float2 = AlignedVector<float, kVecSize>;
const auto tx = threadIdx.x;
const auto warp_id = tx / kWarpThreads;
const auto lane_id = tx % kWarpThreads;
const auto work_id = blockIdx.x;
if (work_id >= params.batch_size) return;
const auto input_ptr = static_cast<const DType*>(params.kv) + work_id * params.kv_stride_batch;
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[work_id]);
const auto out_loc = params.out_loc[work_id];
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float2 data, freq;
// part 1: norm. Each thread owns one 2-elem pack (the `tx`-th).
// Sum-of-squares is reduced block-wide via per-warp partials.
{
__shared__ float partial_sums[kFusedKNumWarps];
Storage input_vec, weight_vec;
input_vec.load(input_ptr, tx);
weight_vec.load(params.kv_weight, tx);
if (warp_id == kRopeWarp) freq.load(freqs_cis, lane_id);
float sum_of_squares = 0.0f;
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const auto x = cast<float>(input_vec[i]);
sum_of_squares += x * x;
}
const auto warp_sum = warp::reduce_sum(sum_of_squares);
if (lane_id == 0) partial_sums[warp_id] = warp_sum;
__syncthreads();
// Replicate the per-warp partial sums onto all lanes of one warp and
// reduce. Every group of `kBlockItemNumWarps` lanes ends up with the
// global sum.
sum_of_squares = warp::reduce_sum<kFusedKNumWarps>(partial_sums[lane_id % kFusedKNumWarps]);
const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const auto x = cast<float>(input_vec[i]);
const auto w = cast<float>(weight_vec[i]);
data[i] = x * norm_factor * w;
}
}
// A negative out_loc marks a slot with no KV write target (e.g. the -1
// sentinel from the full->SWA translation for out-of-window tokens or
// padded rows); skip the row instead of writing out of bounds. Checked
// here, not at the load, so the out_loc prefetch overlaps the norm above.
if (out_loc < 0) return;
const int32_t page = out_loc >> kPageBits;
const int32_t offset = out_loc & ((1 << kPageBits) - 1);
const auto page_ptr = params.kvcache + page * kPageBytes;
const auto value_ptr = page_ptr + offset * 576;
PDLTriggerSecondary<kUsePDL>();
// part 2: rope on warp 7 (BF16 store), per-warp UE8M0 quant + store on warps 0..6.
if (warp_id == kRopeWarp) {
const auto x_real = data[0];
const auto x_imag = data[1];
const auto freq_real = freq[0];
const auto freq_imag = freq[1];
data[0] = x_real * freq_real - x_imag * freq_imag;
data[1] = x_real * freq_imag + x_imag * freq_real;
const auto result = cast<bf16x2_t>(fp32x2_t{data[0], data[1]});
const auto rope_ptr = value_ptr + 448;
reinterpret_cast<bf16x2_t*>(rope_ptr)[lane_id] = result;
} else {
const auto x = data[0];
const auto y = data[1];
const auto abs_max = warp::reduce_max(fmaxf(fabs(x), fabs(y)));
const auto scale_raw = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
const auto scale_ue8m0 = cast_to_ue8m0(scale_raw);
const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
const auto result = pack_fp8(x * inv_scale, y * inv_scale);
const auto scale_ptr = page_ptr + (576 << kPageBits) + offset * 8;
reinterpret_cast<fp8x2_e4m3_t*>(value_ptr)[tx] = result;
if (lane_id == 0) static_cast<uint8_t*>(scale_ptr)[warp_id] = scale_ue8m0;
}
}
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, bool kUsePDL>
struct FusedKNormRopeFlashMLAKernel {
static constexpr int32_t kLogPageSize = std::countr_zero(kPageSize);
static constexpr int64_t kPageBytes = host::div_ceil(584 * kPageSize, 576) * 576;
static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
static_assert(1 << kLogPageSize == kPageSize);
static_assert(kHeadDim == 512 && kRopeDim == 64, "FlashMLA layout requires (512, 64)");
template <typename PosT>
static constexpr auto kernel = fused_k_norm_rope_flashmla<DType, kHeadDim, kRopeDim, PosT, kLogPageSize, kUsePDL>;
static void forward(
const tvm::ffi::TensorView kv,
const tvm::ffi::TensorView kv_weight,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView positions,
const tvm::ffi::TensorView out_loc,
const tvm::ffi::TensorView kvcache,
float eps) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B, kHeadDim}) //
.with_strides({-1, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(kv);
TensorMatcher({kHeadDim}) //
.with_dtype<DType>()
.with_device(device_)
.verify(kv_weight);
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
.with_device(device_)
.verify(positions);
TensorMatcher({B}) //
.with_dtype<int32_t>()
.with_device(device_)
.verify(out_loc);
TensorMatcher({-1, -1}) //
.with_strides({kPageBytes, 1})
.with_dtype<uint8_t>()
.with_device(device_)
.verify(kvcache);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
if (batch_size == 0) return;
const auto params = FusedKNormRopeFlashMLAParams{
.kv = kv.data_ptr(),
.kv_weight = kv_weight.data_ptr(),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.positions = positions.data_ptr(),
.out_loc = static_cast<const int32_t*>(out_loc.data_ptr()),
.kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
.kv_stride_batch = kv.stride(0),
.batch_size = batch_size,
.eps = eps,
};
const auto k_int32 = kernel<int32_t>;
const auto k_int64 = kernel<int64_t>;
const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
LaunchKernel(batch_size, kFusedKBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(k, params);
}
};
// ============================================================================
// Indexer Q kernel: warp-per-(token, head) RoPE + Hadamard + fp8 act-quant.
// ============================================================================
struct FusedQIndexerRopeHadamardQuantParams {
const void* __restrict__ q_input; // (B, num_heads, 128) DType
void* __restrict__ q_fp8; // (B, num_heads, 128) fp8_e4m3
// weights_out[b, h] = weight[b, h] * weight_scale * q_scale[b, h].
// q_scale is computed internally and not exposed -- the only consumer of
// it is `weights_out`.
const void* __restrict__ weight; // (B, num_heads) DType
float* __restrict__ weights_out; // (B, num_heads) fp32 (== (B, H, 1) flat)
float weight_scale; // scalar c4_indexer.weight_scale
// Template-dependent layout:
// kRopeFirst=false: (max_pos, 64) fp32 interleaved [cos0, sin0, ...]
// kRopeFirst=true : (max_pos, 64) fp32 halves [cos..., sin...]
const float* __restrict__ rope_cache;
const void* __restrict__ positions; // (B,) PosT
// Row stride for `weight` (caller passes the non-contiguous wk slice directly).
int64_t weight_stride_batch;
uint32_t batch_size;
uint32_t num_heads;
};
template <typename DType, typename PosT, bool kUsePDL, bool kRopeFirst = false, bool kHadamard = true>
Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQIndexerRopeHadamardQuantParams params) {
using namespace device;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
constexpr int64_t kVecSize = 4;
constexpr uint32_t kRopeSize = kRopeDim / kVecSize; // = 16
static_assert(kHeadDim == kWarpThreads * kVecSize);
static_assert(kRopeDim == kWarpThreads * 2);
static_assert(kRopeSize <= kWarpThreads);
using Storage = AlignedVector<DType, kVecSize>;
using Float4 = AlignedVector<float, kVecSize>;
using OutStorage = AlignedVector<fp8x2_e4m3_t, 2>; // 4 fp8 / lane
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id;
// V4 ropes the trailing kRopeDim dims (kRopeFirst=false); V3.2 ropes the
// leading kRopeDim dims (kRopeFirst=true). Select the owning lanes per layout.
const bool is_rope_lane = kRopeFirst ? (lane_id < kRopeSize) : (lane_id >= kWarpThreads - kRopeSize);
const uint32_t total_works = params.batch_size * params.num_heads;
if (work_id >= total_works) return;
const uint32_t batch_id = work_id / params.num_heads;
const auto input_ptr = static_cast<const DType*>(params.q_input) + work_id * kHeadDim;
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[batch_id]);
const auto rope_cache = params.rope_cache + position * kRopeDim;
// Lane 0 prefetches the weight scalar for this (token, head) work item.
// Weight is (B, num_heads) DType; we need one scalar per warp -- offload
// the load to lane 0 only. The multiply + store happens once the q_scale
// is known (part 4).
PDLWaitPrimary<kUsePDL>();
Float4 data, freq;
const uint32_t head_id = work_id - batch_id * params.num_heads;
const auto weight_val =
cast<float>(static_cast<const DType*>(params.weight)[batch_id * params.weight_stride_batch + head_id]);
// part 1: load (no norm). Each lane owns a 4-elem pack.
{
Storage input_vec;
input_vec.load(input_ptr, lane_id);
if (is_rope_lane) {
if constexpr (kRopeFirst) {
freq = load_rope_first_cos_sin<kRopeDim>(rope_cache, lane_id);
} else {
freq.load(rope_cache, lane_id - (kWarpThreads - kRopeSize));
}
}
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = cast<float>(input_vec[i]);
}
}
// part 2: rope on rope lanes only (4 elems / lane = 2 (real, imag) pairs).
if (is_rope_lane) {
const auto x_real = data[0];
const auto x_imag = data[1];
const auto y_real = data[2];
const auto y_imag = data[3];
const auto fxr = freq[0];
const auto fxi = freq[1];
const auto fyr = freq[2];
const auto fyi = freq[3];
data[0] = x_real * fxr - x_imag * fxi;
data[1] = x_real * fxi + x_imag * fxr;
data[2] = y_real * fyr - y_imag * fyi;
data[3] = y_real * fyi + y_imag * fyr;
}
PDLTriggerSecondary<kUsePDL>();
// part 3: 128-point Hadamard (2 local stages + 5 cross-lane shfl_xor stages).
// Same recipe as `fused_norm_rope_indexer`; see comments there for the
// butterfly invariants and the early-return safety argument. V3.2 omits the
// rotation (kHadamard=false): it is logit-preserving (H orthonormal, applied
// to both q and k), so dropping it only trades fp8 quant accuracy.
if constexpr (kHadamard) {
{
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
data[0] = a0 + a1;
data[1] = a0 - a1;
data[2] = a2 + a3;
data[3] = a2 - a3;
}
{
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
data[0] = a0 + a2;
data[1] = a1 + a3;
data[2] = a0 - a2;
data[3] = a1 - a3;
}
#pragma unroll
for (uint32_t mask = 1; mask < kWarpThreads; mask <<= 1) {
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const float other = __shfl_xor_sync(0xFFFFFFFFu, data[i], mask, kWarpThreads);
data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other);
}
}
const float kHadamardScale = math::rsqrt(static_cast<float>(kHeadDim));
#pragma unroll
for (int i = 0; i < kVecSize; ++i)
data[i] *= kHadamardScale;
}
{
float local_max = math::abs(data[0]);
#pragma unroll
for (int i = 1; i < kVecSize; ++i) {
local_max = math::max(local_max, math::abs(data[i]));
}
const auto abs_max = warp::reduce_max(local_max);
const auto scale = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
const auto inv_scale = 1.0f / scale;
OutStorage result;
result[0] = pack_fp8(data[0] * inv_scale, data[1] * inv_scale);
result[1] = pack_fp8(data[2] * inv_scale, data[3] * inv_scale);
// q_fp8 row pointer: 128 fp8 / row = 32 OutStorage / row, one per lane.
auto out_row = static_cast<uint8_t*>(params.q_fp8) + work_id * kHeadDim;
result.store(out_row, lane_id);
params.weights_out[work_id] = weight_val * params.weight_scale * scale;
}
}
template <typename DType, bool kUsePDL, bool kRopeFirst = false, bool kHadamard = true>
struct FusedQIndexerRopeHadamardQuantKernel {
template <typename PosT>
static constexpr auto kernel = fused_q_indexer_rope_hadamard_quant<DType, PosT, kUsePDL, kRopeFirst, kHadamard>;
static void forward(
const tvm::ffi::TensorView q_input,
const tvm::ffi::TensorView q_fp8,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView weights_out,
double weight_scale,
const tvm::ffi::TensorView rope_cache,
const tvm::ffi::TensorView positions) {
using namespace host;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
auto B = SymbolicSize{"batch_size"};
auto H = SymbolicSize{"num_heads"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
// Caller path is `wq_b(q_lora).view(-1, H, D)` -> contiguous; the kernel
// assumes a flat `(B*H, kHeadDim)` layout for both q_input and q_fp8.
// Pin the head/innermost strides; assert the batch stride below.
TensorMatcher({B, H, kHeadDim}) //
.with_strides({-1, kHeadDim, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(q_input);
TensorMatcher({B, H, kHeadDim}) //
.with_strides({-1, kHeadDim, 1})
.with_dtype<fp8_e4m3_t>()
.with_device(device_)
.verify(q_fp8);
TensorMatcher({B, H}) //
.with_strides({-1, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(weight);
TensorMatcher({B, H, 1}) //
.with_dtype<float>()
.with_device(device_)
.verify(weights_out);
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(rope_cache);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
.with_device(device_)
.verify(positions);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto num_heads = static_cast<uint32_t>(H.unwrap());
if (batch_size == 0) return;
// The kernel computes row pointers as `base + work_id * kHeadDim`, so
// both inputs must be contiguous in (batch, head, elem) order.
const int64_t expected_batch_stride = static_cast<int64_t>(num_heads) * kHeadDim;
RuntimeCheck(
q_input.stride(0) == expected_batch_stride,
"q_input must be contiguous (B, H, kHeadDim); got stride[0]=",
q_input.stride(0));
RuntimeCheck(
q_fp8.stride(0) == expected_batch_stride,
"q_fp8 must be contiguous (B, H, kHeadDim); got stride[0]=",
q_fp8.stride(0));
const auto params = FusedQIndexerRopeHadamardQuantParams{
.q_input = q_input.data_ptr(),
.q_fp8 = q_fp8.data_ptr(),
.weight = weight.data_ptr(),
.weights_out = static_cast<float*>(weights_out.data_ptr()),
.weight_scale = static_cast<float>(weight_scale),
.rope_cache = static_cast<const float*>(rope_cache.data_ptr()),
.positions = positions.data_ptr(),
.weight_stride_batch = weight.stride(0),
.batch_size = batch_size,
.num_heads = num_heads,
};
const auto total_works = batch_size * num_heads;
const auto num_blocks = div_ceil(total_works, kFusedQNumWarps);
const auto k_int32 = kernel<int32_t>;
const auto k_int64 = kernel<int64_t>;
const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
LaunchKernel(num_blocks, kFusedQBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(k, params);
}
};
struct FusedQIndexerRopeHadamardFp4QuantParams {
const void* __restrict__ q_input;
void* __restrict__ q_fp4;
int32_t* __restrict__ q_sf;
const void* __restrict__ weight;
float* __restrict__ weights_out;
float weight_scale;
const float* __restrict__ freqs_cis;
const void* __restrict__ positions;
uint32_t batch_size;
uint32_t num_heads;
};
template <typename DType, typename PosT, bool kUsePDL>
Q_KERNEL void
fused_q_indexer_rope_hadamard_fp4_quant(const __grid_constant__ FusedQIndexerRopeHadamardFp4QuantParams params) {
using namespace device;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
constexpr int64_t kVecSize = 4;
constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
static_assert(kHeadDim == kWarpThreads * kVecSize);
static_assert(kRopeDim == kWarpThreads * 2);
static_assert(kRopeSize <= kWarpThreads);
using Storage = AlignedVector<DType, kVecSize>;
using Float4 = AlignedVector<float, kVecSize>;
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id;
const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
const uint32_t total_works = params.batch_size * params.num_heads;
if (work_id >= total_works) return;
const uint32_t batch_id = work_id / params.num_heads;
const auto input_ptr = static_cast<const DType*>(params.q_input) + work_id * kHeadDim;
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[batch_id]);
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float4 data, freq;
const auto weight_val = cast<float>(static_cast<const DType*>(params.weight)[work_id]);
{
Storage input_vec;
input_vec.load(input_ptr, lane_id);
if (is_rope_lane) freq.load(freqs_cis, lane_id - (kWarpThreads - kRopeSize));
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = cast<float>(input_vec[i]);
}
}
if (is_rope_lane) {
const auto x_real = data[0];
const auto x_imag = data[1];
const auto y_real = data[2];
const auto y_imag = data[3];
const auto fxr = freq[0];
const auto fxi = freq[1];
const auto fyr = freq[2];
const auto fyi = freq[3];
data[0] = x_real * fxr - x_imag * fxi;
data[1] = x_real * fxi + x_imag * fxr;
data[2] = y_real * fyr - y_imag * fyi;
data[3] = y_real * fyi + y_imag * fyr;
#pragma unroll
for (int i = 0; i < kVecSize; ++i)
data[i] = cast<float>(cast<DType>(data[i]));
}
PDLTriggerSecondary<kUsePDL>();
{
{
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
data[0] = a0 + a1;
data[1] = a0 - a1;
data[2] = a2 + a3;
data[3] = a2 - a3;
}
{
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
data[0] = a0 + a2;
data[1] = a1 + a3;
data[2] = a0 - a2;
data[3] = a1 - a3;
}
#pragma unroll
for (uint32_t mask = 1; mask < kWarpThreads; mask <<= 1) {
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const float other = __shfl_xor_sync(0xFFFFFFFFu, data[i], mask, kWarpThreads);
data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other);
}
}
const float kHadamardScale = math::rsqrt(static_cast<float>(kHeadDim));
#pragma unroll
for (int i = 0; i < kVecSize; ++i)
data[i] *= kHadamardScale;
#pragma unroll
for (int i = 0; i < kVecSize; ++i)
data[i] = cast<float>(cast<DType>(data[i]));
}
{
float local_max = math::abs(data[0]);
#pragma unroll
for (int i = 1; i < kVecSize; ++i) {
local_max = math::max(local_max, math::abs(data[i]));
}
local_max = warp::reduce_max<8>(local_max);
const auto scale_raw = fmaxf(1e-4f, local_max) / 6.0f;
const auto scale_ue8m0 = static_cast<uint8_t>(cast_to_ue8m0(scale_raw));
const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
const uint8_t packed0 = quant_fp4_e2m1(data[0] * inv_scale) | (quant_fp4_e2m1(data[1] * inv_scale) << 4);
const uint8_t packed1 = quant_fp4_e2m1(data[2] * inv_scale) | (quant_fp4_e2m1(data[3] * inv_scale) << 4);
const uint16_t packed = static_cast<uint16_t>(packed0) | (static_cast<uint16_t>(packed1) << 8);
auto out_row = static_cast<uint8_t*>(params.q_fp4) + work_id * (kHeadDim / 2);
reinterpret_cast<uint16_t*>(out_row)[lane_id] = packed;
if ((lane_id & 7) == 0) {
reinterpret_cast<uint8_t*>(params.q_sf + work_id)[lane_id >> 3] = scale_ue8m0;
}
params.weights_out[work_id] = weight_val * params.weight_scale;
}
}
template <typename DType, bool kUsePDL>
struct FusedQIndexerRopeHadamardFp4QuantKernel {
template <typename PosT>
static constexpr auto kernel = fused_q_indexer_rope_hadamard_fp4_quant<DType, PosT, kUsePDL>;
static void forward(
const tvm::ffi::TensorView q_input,
const tvm::ffi::TensorView q_fp4,
const tvm::ffi::TensorView q_sf,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView weights_out,
double weight_scale,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView positions) {
using namespace host;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
constexpr int64_t kFp4Dim = kHeadDim / 2;
auto B = SymbolicSize{"batch_size"};
auto H = SymbolicSize{"num_heads"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B, H, kHeadDim})
.with_strides({-1, kHeadDim, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(q_input);
TensorMatcher({B, H, kFp4Dim})
.with_strides({-1, kFp4Dim, 1})
.with_dtype<int8_t>()
.with_device(device_)
.verify(q_fp4);
TensorMatcher({B, H}).with_dtype<int32_t>().with_device(device_).verify(q_sf);
TensorMatcher({B, H}).with_dtype<DType>().with_device(device_).verify(weight);
TensorMatcher({B, H, 1}).with_dtype<float>().with_device(device_).verify(weights_out);
TensorMatcher({-1, kRopeDim}).with_dtype<float>().with_device(device_).verify(freqs_cis);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}).with_dtype<int32_t, int64_t>(pos_dtype).with_device(device_).verify(positions);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto num_heads = static_cast<uint32_t>(H.unwrap());
if (batch_size == 0) return;
const int64_t expected_q_stride = static_cast<int64_t>(num_heads) * kHeadDim;
const int64_t expected_fp4_stride = static_cast<int64_t>(num_heads) * kFp4Dim;
RuntimeCheck(q_input.stride(0) == expected_q_stride, "q_input must be contiguous");
RuntimeCheck(q_fp4.stride(0) == expected_fp4_stride, "q_fp4 must be contiguous");
RuntimeCheck(q_sf.stride(0) == static_cast<int64_t>(num_heads) && q_sf.stride(1) == 1, "q_sf must be contiguous");
const auto params = FusedQIndexerRopeHadamardFp4QuantParams{
.q_input = q_input.data_ptr(),
.q_fp4 = q_fp4.data_ptr(),
.q_sf = static_cast<int32_t*>(q_sf.data_ptr()),
.weight = weight.data_ptr(),
.weights_out = static_cast<float*>(weights_out.data_ptr()),
.weight_scale = static_cast<float>(weight_scale),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.positions = positions.data_ptr(),
.batch_size = batch_size,
.num_heads = num_heads,
};
const auto total_works = batch_size * num_heads;
const auto num_blocks = div_ceil(total_works, kFusedQNumWarps);
const auto k_int32 = kernel<int32_t>;
const auto k_int64 = kernel<int64_t>;
const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
LaunchKernel(num_blocks, kFusedQBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(k, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh
+221
View File
@@ -0,0 +1,221 @@
#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 <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <cstdint>
#include <cuda_fp8.h>
namespace {
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::pack_fp8;
struct MegaMoEPreDispatchParams {
const bf16_t* __restrict__ x; // [num_tokens, hidden]
const int32_t* __restrict__ topk_idx; // [num_tokens, top_k]
const float* __restrict__ topk_weights; // [num_tokens, top_k]
fp8_e4m3_t* __restrict__ buf_x; // [padded_max, hidden]
int32_t* __restrict__ buf_x_sf; // contiguous int32 [P, G/4]; see layout comment
int64_t* __restrict__ buf_topk_idx; // [padded_max, top_k]
float* __restrict__ buf_topk_weights; // [padded_max, top_k]
uint32_t num_tokens;
uint32_t padded_max;
uint32_t hidden;
uint32_t num_groups; // hidden / group_size
uint32_t top_k;
};
// kGroupSize must match sglang_per_token_group_quant_fp8_ue8m0(group_size=).
template <uint32_t kGroupSize, bool kUsePDL>
__global__ __launch_bounds__(1024, 2) void //
mega_moe_pre_dispatch_kernel(const MegaMoEPreDispatchParams __grid_constant__ params) {
using namespace device;
constexpr uint32_t kVecElems = 8; // 8 bf16 = 16B load per thread
static_assert(kGroupSize % kVecElems == 0, "group_size must be a multiple of 8");
constexpr uint32_t kThreadsPerGroup = kGroupSize / kVecElems;
using InputVec = AlignedVector<bf16x2_t, kVecElems / 2>;
using OutputVec = AlignedVector<fp8x2_e4m3_t, kVecElems / 2>;
const uint32_t bid = blockIdx.x;
const uint32_t tid = threadIdx.x;
PDLWaitPrimary<kUsePDL>();
if (bid < params.num_tokens) {
// ---- Quantize path: one CTA per valid token ----
const uint32_t token_id = bid;
const auto token_in = params.x + static_cast<uint64_t>(token_id) * params.hidden;
const auto token_out = params.buf_x + static_cast<uint64_t>(token_id) * params.hidden;
InputVec in_vec;
in_vec.load(token_in, tid);
float local_max = 0.0f;
float vals[kVecElems];
#pragma unroll
for (uint32_t i = 0; i < kVecElems / 2; ++i) {
const auto [v0, v1] = cast<fp32x2_t>(in_vec[i]);
vals[2 * i + 0] = v0;
vals[2 * i + 1] = v1;
local_max = fmaxf(local_max, fmaxf(fabsf(v0), fabsf(v1)));
}
// Absmax across the kThreadsPerGroup threads that cover one group.
local_max = warp::reduce_max<kThreadsPerGroup>(local_max);
const float absmax = fmaxf(local_max, 1e-10f);
const float raw_scale = absmax / math::FP8_E4M3_MAX;
const uint32_t ue8m0_exp = cast_to_ue8m0(raw_scale);
// 2^-ue8m0_exp as fp32 (equivalent to 1 / __uint_as_float(ue8m0 << 23)).
const float inv_scale = __uint_as_float((127u + 127u - ue8m0_exp) << 23);
OutputVec out_vec;
#pragma unroll
for (uint32_t i = 0; i < kVecElems / 2; ++i) {
out_vec[i] = pack_fp8(vals[2 * i + 0] * inv_scale, vals[2 * i + 1] * inv_scale);
}
out_vec.store(token_out, tid);
// One thread per group writes its UE8M0 byte into the contiguous
// row-major int32-packed layout: byte address = t*num_groups + g
// (see layout comment at the top of the file).
const uint32_t group_id = tid / kThreadsPerGroup;
const uint32_t within_group_id = tid % kThreadsPerGroup;
if (within_group_id == 0 && group_id < params.num_groups) {
const uint32_t byte_off = token_id * params.num_groups + group_id;
reinterpret_cast<uint8_t*>(params.buf_x_sf)[byte_off] = static_cast<uint8_t>(ue8m0_exp);
}
// Copy this token's topk row (no alignment assumptions; top_k is small).
if (tid < params.top_k) {
const uint32_t off = token_id * params.top_k + tid;
params.buf_topk_idx[off] = params.topk_idx[off];
params.buf_topk_weights[off] = params.topk_weights[off];
}
} else {
// ---- Pad path: trailing blocks fill [num_tokens, padded_max) with (-1, 0) ----
const uint32_t copy_bid = bid - params.num_tokens;
const uint32_t pad_base = params.num_tokens * params.top_k;
const uint32_t slot = pad_base + copy_bid * blockDim.x + tid;
const uint32_t total_slots = params.padded_max * params.top_k;
if (slot < total_slots) {
params.buf_topk_idx[slot] = -1;
params.buf_topk_weights[slot] = 0.0f;
}
}
PDLTriggerSecondary<kUsePDL>();
}
// ---- Host wrapper
// ------------------------------------------------------------------------------------------------------------------------
template <int64_t kGroupSize, bool kUsePDL>
struct MegaMoEPreDispatchKernel {
static_assert(kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, "unsupported group_size");
static constexpr auto kernel = mega_moe_pre_dispatch_kernel<static_cast<uint32_t>(kGroupSize), kUsePDL>;
static void
run(const tvm::ffi::TensorView x,
const tvm::ffi::TensorView topk_idx,
const tvm::ffi::TensorView topk_weights,
const tvm::ffi::TensorView buf_x,
const tvm::ffi::TensorView buf_x_sf,
const tvm::ffi::TensorView buf_topk_idx,
const tvm::ffi::TensorView buf_topk_weights) {
using namespace host;
auto device = SymbolicDevice{};
auto M = SymbolicSize{"num_tokens"};
auto P = SymbolicSize{"padded_max"};
auto H = SymbolicSize{"hidden"};
auto K = SymbolicSize{"top_k"};
auto G4 = SymbolicSize{"num_groups_div_4"};
device.set_options<kDLCUDA>();
TensorMatcher({M, H}) // input x
.with_dtype<bf16_t>()
.with_device(device)
.verify(x);
TensorMatcher({M, K}) // topk_idx
.with_dtype<int32_t>()
.with_device(device)
.verify(topk_idx);
TensorMatcher({M, K}) // topk_weights
.with_dtype<float>()
.with_device(device)
.verify(topk_weights);
// DeepGEMM versions expose this fp8 dispatch buffer either as raw int8
// storage or as torch.float8_e4m3fn; the kernel writes fp8 bytes in both.
TensorMatcher({P, H}) // buf.x
.with_dtype<int8_t, fp8_e4m3_t>()
.with_device(device)
.verify(buf_x);
// buf.x_sf is the contiguous row-major int32 view from DeepGEMM's mega
// symm buffer (DeepGEMM/csrc/apis/mega.hpp): shape (P, G/4), strides
// (G/4, 1). No explicit strides required -> TensorMatcher enforces
// is_contiguous().
TensorMatcher({P, G4}) // buf_x_sf
.with_dtype<int32_t>()
.with_device(device)
.verify(buf_x_sf);
TensorMatcher({P, K}) // buf.topk_idx
.with_dtype<int64_t>()
.with_device(device)
.verify(buf_topk_idx);
TensorMatcher({P, K}) // buf.topk_weights
.with_dtype<float>()
.with_device(device)
.verify(buf_topk_weights);
const auto num_tokens = static_cast<uint32_t>(M.unwrap());
const auto padded_max = static_cast<uint32_t>(P.unwrap());
const auto hidden = static_cast<uint32_t>(H.unwrap());
const auto top_k = static_cast<uint32_t>(K.unwrap());
const auto num_groups_div_4 = static_cast<uint32_t>(G4.unwrap());
RuntimeCheck(num_tokens <= padded_max, "num_tokens must not exceed padded_max");
RuntimeCheck(hidden % kGroupSize == 0, "hidden must be a multiple of group_size");
const auto num_groups = hidden / static_cast<uint32_t>(kGroupSize);
RuntimeCheck(num_groups == num_groups_div_4 * 4u, "num_groups must be a multiple of 4");
RuntimeCheck(hidden % 8u == 0, "hidden must be a multiple of 8 (16B bf16 loads)");
const auto num_threads = hidden / 8u;
RuntimeCheck(num_threads <= 1024, "hidden too large for single-block-per-row quant");
RuntimeCheck(num_threads >= top_k, "top_k must fit into one quant CTA");
const auto pad_slots = (padded_max - num_tokens) * top_k;
const uint32_t num_pad_blocks = pad_slots == 0 ? 0u : ((pad_slots + num_threads - 1u) / num_threads);
const auto num_total_blocks = num_tokens + num_pad_blocks;
const auto params = MegaMoEPreDispatchParams{
.x = static_cast<const bf16_t*>(x.data_ptr()),
.topk_idx = static_cast<const int32_t*>(topk_idx.data_ptr()),
.topk_weights = static_cast<const float*>(topk_weights.data_ptr()),
.buf_x = static_cast<fp8_e4m3_t*>(buf_x.data_ptr()),
.buf_x_sf = static_cast<int32_t*>(buf_x_sf.data_ptr()),
.buf_topk_idx = static_cast<int64_t*>(buf_topk_idx.data_ptr()),
.buf_topk_weights = static_cast<float*>(buf_topk_weights.data_ptr()),
.num_tokens = num_tokens,
.padded_max = padded_max,
.hidden = hidden,
.num_groups = num_groups,
.top_k = top_k,
};
if (num_total_blocks == 0) return;
LaunchKernel(num_total_blocks, num_threads, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh
+387
View File
@@ -0,0 +1,387 @@
#pragma once
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <cstdint>
namespace {
SGL_DEVICE int64_t clamp_accept_len(int64_t delta, int64_t max_accept) {
if (delta < 0) return 0;
return delta < max_accept ? delta : max_accept;
}
template <typename TSeq, typename TReq>
struct OnlineC128MTPWritePrefixParams {
const float* __restrict__ kv_score_input;
const TSeq* __restrict__ seq_lens;
const TReq* __restrict__ req_pool_indices;
const int32_t* __restrict__ req_to_token;
const float* __restrict__ ape;
float* __restrict__ state;
int64_t kv_score_stride_b;
int64_t req_to_token_stride_b;
int64_t ape_stride_r;
int64_t state_stride_b;
int64_t layer_bs;
int64_t num_verify_tokens;
int64_t state_slot_stride;
};
template <typename TSeq, typename TReq>
struct OnlineC128MTPMarkPendingParams {
const TSeq* __restrict__ seq_lens;
const TReq* __restrict__ req_pool_indices;
int64_t* __restrict__ pending_seq_lens;
int64_t bs;
int64_t max_num_reqs;
};
template <typename TSeq, typename TReq>
struct OnlineC128MTPCommitPendingParams {
const TSeq* __restrict__ cur_seq_lens;
const TReq* __restrict__ cur_req_pool_indices;
const int32_t* __restrict__ req_to_token;
const int64_t* __restrict__ pending_seq_lens;
float* __restrict__ state;
int64_t cur_bs;
int64_t req_to_token_stride_b;
int64_t state_stride_b;
int64_t num_verify_tokens;
int64_t state_slot_stride;
int64_t max_num_reqs;
};
__global__ void online_c128_mtp_clear_all_pending_kernel(int64_t* pending_seq_lens, int64_t max_num_reqs) {
const int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
if (idx < max_num_reqs) pending_seq_lens[idx] = -1;
}
template <typename TSeq, typename TReq>
__global__ void online_c128_mtp_mark_pending_kernel(const OnlineC128MTPMarkPendingParams<TSeq, TReq> params) {
const int64_t bid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
if (bid >= params.bs) return;
const int64_t req = static_cast<int64_t>(params.req_pool_indices[bid]);
if (req >= 0 && req < params.max_num_reqs) {
params.pending_seq_lens[req] = static_cast<int64_t>(params.seq_lens[bid]);
}
}
template <int64_t kHeadDim, typename TSeq, typename TReq>
__global__ void online_c128_mtp_commit_pending_kernel(const OnlineC128MTPCommitPendingParams<TSeq, TReq> params) {
const int64_t bid = static_cast<int64_t>(blockIdx.x);
if (bid >= params.cur_bs) return;
const int64_t req = static_cast<int64_t>(params.cur_req_pool_indices[bid]);
if (req < 0 || req >= params.max_num_reqs) return;
const int64_t old_seq = params.pending_seq_lens[req];
if (old_seq < 0) return;
const int64_t cur_seq = static_cast<int64_t>(params.cur_seq_lens[bid]);
const int64_t accept = clamp_accept_len(cur_seq - old_seq, params.num_verify_tokens);
if (accept <= 0) return;
const int64_t final_seq = old_seq + accept;
if ((final_seq & 127) == 0) return;
const int64_t slot = req;
const float* const src = params.state + (slot + accept * params.state_slot_stride) * params.state_stride_b;
float* const dst = params.state + slot * params.state_stride_b;
for (int64_t d = static_cast<int64_t>(threadIdx.x); d < kHeadDim * 3; d += blockDim.x) {
dst[d] = src[d];
}
}
template <int64_t kHeadDim, typename TSeq, typename TReq>
__global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePrefixParams<TSeq, TReq> params) {
const int64_t bid = static_cast<int64_t>(blockIdx.x);
if (bid >= params.layer_bs) return;
const int64_t seq_before = static_cast<int64_t>(params.seq_lens[bid]);
const int64_t req_idx = static_cast<int64_t>(params.req_pool_indices[bid]);
const int64_t start_pos = seq_before & 127;
const bool has_partial = seq_before > 0 && start_pos != 0;
int64_t init_slot = 0;
if (has_partial) {
init_slot = req_idx;
}
const int64_t d = static_cast<int64_t>(threadIdx.x);
float run_max = 0.0f;
float run_sum = 0.0f;
float run_kv = 0.0f;
if (has_partial) {
const float* const init = params.state + init_slot * params.state_stride_b;
run_max = init[d];
run_sum = init[kHeadDim + d];
run_kv = init[kHeadDim * 2 + d];
}
constexpr int kMaxVerifyTokens = 8;
float kv_steps[kMaxVerifyTokens];
float score_steps[kMaxVerifyTokens];
#pragma unroll
for (int step = 0; step < kMaxVerifyTokens; ++step) {
if (step >= params.num_verify_tokens) break;
const int64_t pos = (start_pos + step) & 127;
const float* const kv = params.kv_score_input + (bid * params.num_verify_tokens + step) * params.kv_score_stride_b;
kv_steps[step] = kv[d];
score_steps[step] = kv[kHeadDim + d] + params.ape[pos * params.ape_stride_r + d];
}
#pragma unroll
for (int step = 0; step < kMaxVerifyTokens; ++step) {
if (step >= params.num_verify_tokens) break;
const int64_t pos = (start_pos + step) & 127;
const float kv_step = kv_steps[step];
const float score_step = score_steps[step];
if (pos == 0) {
run_kv = kv_step;
run_max = score_step;
run_sum = 1.0f;
} else {
const float new_max = fmaxf(run_max, score_step);
const float old_sum_scaled = run_sum * __expf(run_max - new_max);
const float new_exp = __expf(score_step - new_max);
const float new_sum = old_sum_scaled + new_exp;
run_kv = (run_kv * old_sum_scaled + kv_step * new_exp) / new_sum;
run_max = new_max;
run_sum = new_sum;
}
const int64_t final_seq = seq_before + step + 1;
if ((final_seq & 127) != 0) {
const int64_t slot = req_idx + (step + 1) * params.state_slot_stride;
float* const out = params.state + slot * params.state_stride_b;
out[d] = run_max;
out[kHeadDim + d] = run_sum;
out[kHeadDim * 2 + d] = run_kv;
}
if (pos == 127) {
run_kv = 0.0f;
run_max = 0.0f;
run_sum = 0.0f;
}
}
}
template <int64_t kHeadDim, typename TSeq, typename TReq>
struct OnlineC128MTPWritePrefixKernel {
static void launch(
tvm::ffi::TensorView kv_score_input,
tvm::ffi::TensorView seq_lens,
tvm::ffi::TensorView req_pool_indices,
tvm::ffi::TensorView req_to_token,
tvm::ffi::TensorView ape,
tvm::ffi::TensorView state,
int64_t layer_bs,
int64_t num_verify_tokens,
int64_t state_slot_stride,
DLDevice device) {
using namespace host;
const auto params = OnlineC128MTPWritePrefixParams<TSeq, TReq>{
.kv_score_input = static_cast<const float*>(kv_score_input.data_ptr()),
.seq_lens = static_cast<const TSeq*>(seq_lens.data_ptr()),
.req_pool_indices = static_cast<const TReq*>(req_pool_indices.data_ptr()),
.req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
.ape = static_cast<const float*>(ape.data_ptr()),
.state = static_cast<float*>(state.data_ptr()),
.kv_score_stride_b = kv_score_input.stride(0),
.req_to_token_stride_b = req_to_token.stride(0),
.ape_stride_r = ape.stride(0),
.state_stride_b = state.stride(0),
.layer_bs = layer_bs,
.num_verify_tokens = num_verify_tokens,
.state_slot_stride = state_slot_stride,
};
static_assert(kHeadDim == 512, "online c128 MTP write-prefix only supports head_dim=512");
constexpr uint32_t kThreads = static_cast<uint32_t>(kHeadDim);
LaunchKernel(static_cast<uint32_t>(layer_bs), kThreads, device)(
online_c128_mtp_write_prefix_kernel<kHeadDim, TSeq, TReq>, params);
}
static void
run(tvm::ffi::TensorView kv_score_input,
tvm::ffi::TensorView seq_lens,
tvm::ffi::TensorView req_pool_indices,
tvm::ffi::TensorView req_to_token,
tvm::ffi::TensorView ape,
tvm::ffi::TensorView state,
int64_t layer_bs,
int64_t num_verify_tokens,
int64_t state_slot_stride) {
using namespace host;
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({-1, kHeadDim * 2}).with_dtype<float>().with_device(device).verify(kv_score_input);
TensorMatcher({-1}).with_dtype<TSeq>().with_device(device).verify(seq_lens);
TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(req_pool_indices);
TensorMatcher({-1, -1}).with_dtype<int32_t>().with_device(device).verify(req_to_token);
TensorMatcher({128, kHeadDim}).with_dtype<float>().with_device(device).verify(ape);
TensorMatcher({-1, kHeadDim * 3}).with_dtype<float>().with_device(device).verify(state);
if (layer_bs <= 0) return;
RuntimeCheck(num_verify_tokens > 0 && num_verify_tokens <= 8, "unsupported num_verify_tokens=", num_verify_tokens);
RuntimeCheck(state_slot_stride > 0, "state_slot_stride must be positive");
RuntimeCheck(layer_bs <= seq_lens.shape()[0], "layer_bs exceeds seq_lens rows");
RuntimeCheck(layer_bs <= req_pool_indices.shape()[0], "layer_bs exceeds req_pool_indices rows");
RuntimeCheck(layer_bs * num_verify_tokens <= kv_score_input.shape()[0], "kv_score_input is too small");
launch(
kv_score_input,
seq_lens,
req_pool_indices,
req_to_token,
ape,
state,
layer_bs,
num_verify_tokens,
state_slot_stride,
device.unwrap());
}
};
template <int64_t kHeadDim, typename TSeq, typename TReq>
struct OnlineC128MTPMarkPendingKernel {
static void launch(
tvm::ffi::TensorView seq_lens,
tvm::ffi::TensorView req_pool_indices,
tvm::ffi::TensorView pending_seq_lens,
int64_t bs,
int64_t max_num_reqs,
DLDevice device) {
using namespace host;
const auto params = OnlineC128MTPMarkPendingParams<TSeq, TReq>{
.seq_lens = static_cast<const TSeq*>(seq_lens.data_ptr()),
.req_pool_indices = static_cast<const TReq*>(req_pool_indices.data_ptr()),
.pending_seq_lens = static_cast<int64_t*>(pending_seq_lens.data_ptr()),
.bs = bs,
.max_num_reqs = max_num_reqs,
};
constexpr uint32_t kThreads = 256;
const uint32_t clear_blocks = host::div_ceil(static_cast<uint32_t>(max_num_reqs), kThreads);
LaunchKernel(clear_blocks, kThreads, device)(
online_c128_mtp_clear_all_pending_kernel, params.pending_seq_lens, max_num_reqs);
const uint32_t mark_blocks = host::div_ceil(static_cast<uint32_t>(bs), kThreads);
LaunchKernel(mark_blocks, kThreads, device)(online_c128_mtp_mark_pending_kernel<TSeq, TReq>, params);
}
static void
run(tvm::ffi::TensorView seq_lens,
tvm::ffi::TensorView req_pool_indices,
tvm::ffi::TensorView pending_seq_lens,
int64_t bs,
int64_t max_num_reqs) {
using namespace host;
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({-1}).with_dtype<TSeq>().with_device(device).verify(seq_lens);
TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(req_pool_indices);
TensorMatcher({-1}).with_dtype<int64_t>().with_device(device).verify(pending_seq_lens);
if (bs <= 0) return;
RuntimeCheck(bs <= seq_lens.shape()[0], "bs exceeds seq_lens rows");
RuntimeCheck(bs <= req_pool_indices.shape()[0], "bs exceeds req_pool_indices rows");
RuntimeCheck(max_num_reqs <= pending_seq_lens.shape()[0], "max_num_reqs exceeds pending rows");
launch(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap());
}
};
template <int64_t kHeadDim, typename TSeq, typename TReq>
struct OnlineC128MTPCommitPendingKernel {
static void launch(
tvm::ffi::TensorView cur_seq_lens,
tvm::ffi::TensorView cur_req_pool_indices,
tvm::ffi::TensorView req_to_token,
tvm::ffi::TensorView pending_seq_lens,
tvm::ffi::TensorView state,
int64_t cur_bs,
int64_t num_verify_tokens,
int64_t state_slot_stride,
int64_t max_num_reqs,
DLDevice device) {
using namespace host;
const auto params = OnlineC128MTPCommitPendingParams<TSeq, TReq>{
.cur_seq_lens = static_cast<const TSeq*>(cur_seq_lens.data_ptr()),
.cur_req_pool_indices = static_cast<const TReq*>(cur_req_pool_indices.data_ptr()),
.req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
.pending_seq_lens = static_cast<const int64_t*>(pending_seq_lens.data_ptr()),
.state = static_cast<float*>(state.data_ptr()),
.cur_bs = cur_bs,
.req_to_token_stride_b = req_to_token.stride(0),
.state_stride_b = state.stride(0),
.num_verify_tokens = num_verify_tokens,
.state_slot_stride = state_slot_stride,
.max_num_reqs = max_num_reqs,
};
constexpr uint32_t kThreads = 256;
LaunchKernel(static_cast<uint32_t>(cur_bs), kThreads, device)(
online_c128_mtp_commit_pending_kernel<kHeadDim, TSeq, TReq>, params);
}
static void
run(tvm::ffi::TensorView cur_seq_lens,
tvm::ffi::TensorView cur_req_pool_indices,
tvm::ffi::TensorView req_to_token,
tvm::ffi::TensorView pending_seq_lens,
tvm::ffi::TensorView state,
int64_t cur_bs,
int64_t num_verify_tokens,
int64_t state_slot_stride,
int64_t max_num_reqs) {
using namespace host;
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({-1}).with_dtype<TSeq>().with_device(device).verify(cur_seq_lens);
TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(cur_req_pool_indices);
TensorMatcher({-1, -1}).with_dtype<int32_t>().with_device(device).verify(req_to_token);
TensorMatcher({-1}).with_dtype<int64_t>().with_device(device).verify(pending_seq_lens);
TensorMatcher({-1, kHeadDim * 3}).with_dtype<float>().with_device(device).verify(state);
if (cur_bs <= 0) return;
RuntimeCheck(num_verify_tokens > 0 && num_verify_tokens <= 8, "unsupported num_verify_tokens=", num_verify_tokens);
RuntimeCheck(state_slot_stride > 0, "state_slot_stride must be positive");
RuntimeCheck(cur_bs <= cur_seq_lens.shape()[0], "cur_bs exceeds seq_lens rows");
RuntimeCheck(cur_bs <= cur_req_pool_indices.shape()[0], "cur_bs exceeds req rows");
RuntimeCheck(max_num_reqs <= pending_seq_lens.shape()[0], "max_num_reqs exceeds pending rows");
launch(
cur_seq_lens,
cur_req_pool_indices,
req_to_token,
pending_seq_lens,
state,
cur_bs,
num_verify_tokens,
state_slot_stride,
max_num_reqs,
device.unwrap());
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/paged_mqa_metadata.cuh
+119
View File
@@ -0,0 +1,119 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/warp.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
namespace {
constexpr uint32_t kBlockSize = 1024;
constexpr uint32_t kSplitKV = 256; // const for both SM90 and SM100
struct MetadataParams {
/// NOTE: batch_size > 0
uint32_t batch_size;
uint32_t num_sm;
const uint32_t* __restrict__ context_lens;
uint32_t* __restrict__ schedule_metadata;
bool use_smem = true;
};
__global__ __launch_bounds__(kBlockSize, 1) //
void smxx_paged_mqa_logits_metadata(const MetadataParams params) {
using namespace device;
extern __shared__ uint32_t s_length[];
static constexpr auto kNumWarps = kBlockSize / kWarpThreads;
static_assert(kNumWarps == kWarpThreads);
const auto tx = threadIdx.x;
const auto lane_id = tx % kWarpThreads;
const auto warp_id = tx / kWarpThreads;
__shared__ uint32_t s_warp_sum[kNumWarps];
uint32_t local_sum = 0;
for (uint32_t i = tx; i < params.batch_size; i += kBlockSize) {
const auto length = params.context_lens[i];
local_sum += (length + kSplitKV - 1) / kSplitKV;
if (params.use_smem) s_length[i] = length;
}
s_warp_sum[warp_id] = warp::reduce_sum(local_sum);
__syncthreads();
const auto global_sum = warp::reduce_sum(s_warp_sum[lane_id]);
if (lane_id != 0) return;
const auto length_ptr = params.use_smem ? s_length : params.context_lens;
const auto avg = global_sum / params.num_sm;
const auto ret = global_sum % params.num_sm;
uint32_t q = 0;
uint32_t num_work = (length_ptr[0] + kSplitKV - 1) / kSplitKV;
uint32_t sum_work = num_work;
for (auto i = warp_id; i <= params.num_sm; i += kNumWarps) {
const auto target = i * avg + min(i, ret);
while (sum_work <= target) {
if (++q >= params.batch_size) break;
num_work = (length_ptr[q] + kSplitKV - 1) / kSplitKV;
sum_work += num_work;
}
if (q >= params.batch_size) {
params.schedule_metadata[2 * i + 0] = params.batch_size;
params.schedule_metadata[2 * i + 1] = 0;
} else {
// sum > target && (sum - length) <= target
params.schedule_metadata[2 * i + 0] = q;
params.schedule_metadata[2 * i + 1] = target - (sum_work - num_work);
}
}
}
template <auto* f, size_t kMaxDynamicSMEM>
void setup_kernel_smem_once(host::DebugInfo where = {}) {
[[maybe_unused]]
static const auto result = [] {
const auto fptr = std::bit_cast<const void*>(f);
return ::cudaFuncSetAttribute(fptr, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM);
}();
host::RuntimeDeviceCheck(result, where);
}
struct IndexerMetadataKernel {
static constexpr auto kMaxBatchSizeInSmem = 16384 * 2; // 128 KB smeme
static void run(tvm::ffi::TensorView seq_lens, tvm::ffi::TensorView metadata) {
using namespace host;
auto N = SymbolicSize{"batch_size"};
auto M = SymbolicSize{"num_sm"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({N}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(seq_lens);
TensorMatcher({M, 2}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(metadata);
const auto batch_size = static_cast<uint32_t>(N.unwrap());
const auto num_sm = static_cast<uint32_t>(M.unwrap()) - 1;
RuntimeCheck(num_sm <= 1024);
const auto use_smem = batch_size <= kMaxBatchSizeInSmem;
const auto params = MetadataParams{
.batch_size = batch_size,
.num_sm = num_sm,
.context_lens = static_cast<uint32_t*>(seq_lens.data_ptr()),
.schedule_metadata = static_cast<uint32_t*>(metadata.data_ptr()),
.use_smem = use_smem,
};
constexpr auto kernel = smxx_paged_mqa_logits_metadata;
setup_kernel_smem_once<kernel, (kMaxBatchSizeInSmem + 1) * sizeof(uint32_t)>();
const auto smem = use_smem ? (batch_size + 1) * sizeof(uint32_t) : 0;
LaunchKernel(1, kBlockSize, device.unwrap(), smem)(kernel, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/rope.cuh
+169
View File
@@ -0,0 +1,169 @@
#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 <tvm/ffi/container/tensor.h>
#include <cstdint>
namespace {
using DType = bf16_t;
constexpr int64_t kRopeDim = 64;
constexpr uint32_t kBlockSize = 128;
constexpr uint32_t kNumWarps = kBlockSize / device::kWarpThreads;
struct FusedQKRopeParams {
void* __restrict__ q;
void* __restrict__ k;
const float* __restrict__ freqs_cis;
const void* __restrict__ positions;
int64_t q_stride_batch;
int64_t k_stride_batch;
int64_t q_stride_head;
int64_t k_stride_head;
uint32_t num_q_heads;
uint32_t num_k_heads;
uint32_t batch_size;
};
template <bool kUsePDL, bool kInverse, typename IndexType>
__global__ __launch_bounds__(kBlockSize, 16) //
void deepseek_rope_kernel(const __grid_constant__ FusedQKRopeParams param) {
using namespace device;
using DType2 = packed_t<DType>;
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
const auto global_warp_id = blockIdx.x * kNumWarps + warp_id;
const auto& [
q, k, freqs_cis, positions, //
q_stride_batch, k_stride_batch, q_stride_head, k_stride_head, //
num_q_heads, num_k_heads, batch_size
] = param;
const auto num_total_heads = num_q_heads + num_k_heads;
const auto head_id = global_warp_id % num_total_heads;
const auto batch_id = global_warp_id / num_total_heads;
if (batch_id >= batch_size) return;
const auto position = static_cast<const IndexType*>(positions)[batch_id];
const auto is_q = head_id < num_q_heads;
const auto local_head = is_q ? head_id : (head_id - num_q_heads);
const auto stride_batch = is_q ? q_stride_batch : k_stride_batch;
const auto stride_head = is_q ? q_stride_head : k_stride_head;
const auto base_ptr = is_q ? q : k;
const auto input = static_cast<DType2*>(pointer::offset(base_ptr, batch_id * stride_batch, local_head * stride_head));
const auto freq_ptr = reinterpret_cast<const fp32x2_t*>(freqs_cis + position * kRopeDim);
const auto [f_real, f_imag] = freq_ptr[lane_id];
PDLWaitPrimary<kUsePDL>();
const auto data = input[lane_id];
const auto [x_real, x_imag] = cast<fp32x2_t>(data);
fp32x2_t output;
if constexpr (kInverse) {
// (a + bi) * (c - di) = (ac + bd) + (bc - ad)i
output = {
x_real * f_real + x_imag * f_imag,
x_imag * f_real - x_real * f_imag,
};
} else {
// (a + bi) * (c + di) = (ac - bd) + (ad + bc)i
output = {
x_real * f_real - x_imag * f_imag,
x_real * f_imag + x_imag * f_real,
};
}
input[lane_id] = cast<DType2>(output);
PDLTriggerSecondary<kUsePDL>();
}
template <bool kUsePDL>
struct FusedQKRopeKernel {
// 4 kernel variants: {forward, inverse} x {int32, int64}
static constexpr auto kernel_fwd_i32 = deepseek_rope_kernel<kUsePDL, false, int32_t>;
static constexpr auto kernel_fwd_i64 = deepseek_rope_kernel<kUsePDL, false, int64_t>;
static constexpr auto kernel_inv_i32 = deepseek_rope_kernel<kUsePDL, true, int32_t>;
static constexpr auto kernel_inv_i64 = deepseek_rope_kernel<kUsePDL, true, int64_t>;
static void forward(
const tvm::ffi::TensorView q,
const tvm::ffi::Optional<tvm::ffi::TensorView> k,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView positions,
bool inverse) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto Q = SymbolicSize{"num_q_heads"};
auto K = SymbolicSize{"num_k_heads"};
constexpr auto D = kRopeDim;
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B, Q, D}) //
.with_strides({-1, -1, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(q);
if (k.has_value()) {
TensorMatcher({B, K, D}) //
.with_strides({-1, -1, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(k.value());
} else {
K.set_value(0);
}
TensorMatcher({-1, D}) //
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
.with_device(device_)
.verify(positions);
const bool pos_i32 = pos_dtype.is_type<int32_t>();
const auto batch_size = static_cast<uint32_t>(B.unwrap());
if (batch_size == 0) return;
const auto num_q_heads = static_cast<uint32_t>(Q.unwrap());
const auto num_k_heads = static_cast<uint32_t>(K.unwrap());
const auto num_total_heads = num_q_heads + num_k_heads;
const auto total_warps = batch_size * num_total_heads;
const auto num_blocks = div_ceil(total_warps, kNumWarps);
const auto elem_size = static_cast<int64_t>(sizeof(DType));
const auto params = FusedQKRopeParams{
.q = q.data_ptr(),
.k = k ? k.value().data_ptr() : nullptr,
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.positions = positions.data_ptr(),
.q_stride_batch = q.stride(0) * elem_size,
.k_stride_batch = k ? k.value().stride(0) * elem_size : 0,
.q_stride_head = q.stride(1) * elem_size,
.k_stride_head = k ? k.value().stride(1) * elem_size : 0,
.num_q_heads = num_q_heads,
.num_k_heads = num_k_heads,
.batch_size = batch_size,
};
// dispatch: {inverse} x {pos_i32}
using KernelType = decltype(kernel_fwd_i32);
const KernelType kernel =
inverse ? (pos_i32 ? kernel_inv_i32 : kernel_inv_i64) : (pos_i32 ? kernel_fwd_i32 : kernel_fwd_i64);
LaunchKernel(num_blocks, kBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh
+540
View File
@@ -0,0 +1,540 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <cstdint>
#include <cuda_fp8.h>
#include <type_traits>
namespace {
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::pack_fp8;
struct SiluMulQuantVarlenParams {
const bf16_t* __restrict__ input;
fp8_e4m3_t* __restrict__ output;
float* __restrict__ output_scale;
const int32_t* __restrict__ masked_m;
float swiglu_limit; // only read when kApplySwigluLimit=true
int64_t hidden_dim;
uint32_t num_tokens;
uint32_t num_experts;
};
constexpr uint32_t kMaxExperts = 256;
struct alignas(16) CTAWork {
uint32_t expert_id;
uint32_t expert_token_id;
bool valid;
};
SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) {
static_assert(device::kWarpThreads == 32);
#pragma unroll
for (uint32_t offset = 1; offset < 32; offset *= 2) {
uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset);
if (lane_id >= offset) val += n;
}
return val;
}
template <bool kApplySwigluLimit, bool kPrecise = true, typename DType2>
SGL_DEVICE fp32x2_t silu_and_mul(DType2 gate, DType2 up, float limit) {
using namespace device;
// refer to as implementation. TL;DR: must clamp in bf16
// https://github.com/deepseek-ai/DeepGEMM/blob/7f2a703ed51ac1f7af07f5e1453b2d3267d37d50/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh#L984-L997
if constexpr (kApplySwigluLimit) {
static_assert(std::is_same_v<DType2, bf16x2_t>);
gate = __hmin2(gate, {limit, limit});
up = __hmax2(up, {-limit, -limit});
up = __hmin2(up, {limit, limit});
}
const auto [g0, g1] = cast<fp32x2_t>(gate);
const auto [u0, u1] = cast<fp32x2_t>(up);
const auto silu0 = g0 / (1.0f + __expf(-g0));
const auto silu1 = g1 / (1.0f + __expf(-g1));
const float val0 = silu0 * u0;
const float val1 = silu1 * u1;
if constexpr (kPrecise) { // I don't know if we should enable this?
return {val0, val1};
} else {
return cast<fp32x2_t>(cast<bf16x2_t>(fp32x2_t{val0, val1}));
}
}
[[maybe_unused]]
SGL_DEVICE CTAWork get_work(const SiluMulQuantVarlenParams& params) {
// Preconditions:
// 1. blockDim.x >= params.num_experts
// 2. params.num_experts <= kMaxExperts
using namespace device;
static_assert(kWarpThreads == 32);
static __shared__ uint32_t s_warp_sum[32];
static __shared__ CTAWork result;
result.valid = false;
const uint32_t tx = threadIdx.x;
const uint32_t lane_id = tx % kWarpThreads;
const uint32_t warp_id = tx / kWarpThreads;
const uint32_t val = tx < params.num_experts ? params.masked_m[tx] : 0u;
// Per-warp inclusive scan of masked_m.
const uint32_t warp_inclusive = warp_inclusive_sum(lane_id, val);
const uint32_t warp_exclusive = warp_inclusive - val;
// Write each warp total.
if (lane_id == kWarpThreads - 1) s_warp_sum[warp_id] = warp_inclusive;
__syncthreads();
const auto tmp_val = lane_id < warp_id ? s_warp_sum[lane_id] : 0u;
const auto prefix_exclusive = warp::reduce_sum(tmp_val) + warp_exclusive;
const auto bx = blockIdx.x;
if (prefix_exclusive <= bx && bx < prefix_exclusive + val) {
result = {tx, bx - prefix_exclusive, true};
}
__syncthreads();
return result;
}
template <bool kScaleUE8M0, bool kTransposed, bool kSwizzle, bool kUsePDL, bool kApplySwigluLimit>
__global__ __launch_bounds__(1024, 2) void // maximize occupancy
silu_mul_quant_varlen_kernel(const SiluMulQuantVarlenParams __grid_constant__ params) {
using namespace device;
constexpr uint32_t kGroupSize = 128u;
constexpr uint32_t kWorkThreads = 16u;
// each thread will handle 8 elements
using InputVec = AlignedVector<bf16x2_t, 4>;
using OutputVec = AlignedVector<fp8x2_e4m3_t, 4>;
static_assert(8 * kWorkThreads == 128, "Invalid tiling");
static_assert(!(kTransposed && !kScaleUE8M0), "transposed layout only supports ue8m0");
const auto [expert_id, token_id, valid] = get_work(params);
if (!valid) return;
const auto work_id = threadIdx.x / kWorkThreads;
const auto offset = expert_id * params.num_tokens + token_id;
const auto input = params.input + offset * params.hidden_dim * 2;
const auto output = params.output + offset * params.hidden_dim;
[[maybe_unused]]
const auto output_scale = [&] {
const auto num_groups = params.hidden_dim / kGroupSize;
if constexpr (kTransposed) {
const auto base = reinterpret_cast<uint8_t*>(params.output_scale);
// Physical layout is [E, G//4, N] int32. Each int32 packs 4 consecutive
// group scales for the same token, so the byte address is:
// expert_offset + (group/4)*N*4 + token*4 + group%4
return base + expert_id * num_groups * params.num_tokens + (work_id / 4u) * (params.num_tokens * 4u) +
token_id * 4u + (work_id % 4u);
} else {
return params.output_scale + offset * num_groups + work_id;
}
}();
PDLWaitPrimary<kUsePDL>();
InputVec gate_vec, up_vec;
if constexpr (kSwizzle) {
// gran=8 interleaved: every 16-element chunk on the N axis is
// [gate[0..7], up[0..7]]. Each thread handles 8 consecutive output
// elements, so its gate chunk lives at vec index 2*threadIdx.x and its
// up chunk at 2*threadIdx.x+1.
gate_vec.load(input, threadIdx.x * 2);
up_vec.load(input, threadIdx.x * 2 + 1);
} else {
gate_vec.load(input, threadIdx.x);
up_vec.load(input, threadIdx.x + blockDim.x);
}
float local_max = 0.0f;
float results[8];
#pragma unroll
for (uint32_t i = 0; i < 4; ++i) {
const auto [x, y] = silu_and_mul<kApplySwigluLimit>(gate_vec[i], up_vec[i], params.swiglu_limit);
results[2 * i + 0] = x;
results[2 * i + 1] = y;
local_max = fmaxf(local_max, fmaxf(fabsf(x), fabsf(y)));
}
local_max = warp::reduce_max<kWorkThreads>(local_max);
const float absmax = fmaxf(local_max, 1e-10f);
float scale;
uint32_t ue8m0_exp;
if constexpr (kScaleUE8M0) {
const float raw_scale = absmax / math::FP8_E4M3_MAX;
ue8m0_exp = cast_to_ue8m0(raw_scale);
scale = __uint_as_float(ue8m0_exp << 23);
} else {
scale = absmax / math::FP8_E4M3_MAX;
}
const auto inv_scale = 1.0f / scale;
OutputVec out_vec;
#pragma unroll
for (uint32_t i = 0; i < 4; ++i) {
const float scaled_val0 = results[2 * i + 0] * inv_scale;
const float scaled_val1 = results[2 * i + 1] * inv_scale;
out_vec[i] = pack_fp8(scaled_val0, scaled_val1);
}
PDLTriggerSecondary<kUsePDL>();
out_vec.store(output, threadIdx.x);
if constexpr (kTransposed) {
*output_scale = ue8m0_exp;
} else {
*output_scale = scale;
}
}
struct SiluAndMulClampParams {
const void* __restrict__ input;
void* __restrict__ output;
float swiglu_limit;
};
template <typename DType, bool kUsePDL>
__global__ __launch_bounds__(1024, 2) void // maximize occupancy
silu_mul_clamp_kernel(const SiluAndMulClampParams __grid_constant__ params) {
using namespace device;
static_assert(sizeof(DType) == 2, "only fp16/bf16 supported");
using DType2 = packed_t<DType>;
constexpr auto kVecSize = 16 / sizeof(DType);
static_assert(kVecSize % 2 == 0 && kVecSize > 0);
using Vec = AlignedVector<DType2, kVecSize / 2>;
const auto bid = blockIdx.x;
const auto tile = tile::Memory<Vec>::cta();
const float limit = params.swiglu_limit;
PDLWaitPrimary<kUsePDL>();
const auto gate = tile.load(params.input, bid * 2 + 0);
const auto up = tile.load(params.input, bid * 2 + 1);
Vec out;
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
out[i] = cast<DType2>(silu_and_mul<true>(cast<bf16x2_t>(gate[i]), cast<bf16x2_t>(up[i]), limit));
}
tile.store(params.output, out, bid);
PDLTriggerSecondary<kUsePDL>();
}
// ---- Host wrapper
// ------------------------------------------------------------------------------------------------------------------------
template <int64_t kGroupSize, bool kScaleUE8M0, bool kSwizzle, bool kUsePDL, bool kApplySwigluLimit>
struct SiluAndMulMaskedPostQuantKernel {
static_assert(kGroupSize == 128);
static constexpr auto kernel_normal =
silu_mul_quant_varlen_kernel<kScaleUE8M0, false, kSwizzle, kUsePDL, kApplySwigluLimit>;
static constexpr auto kernel_transposed =
silu_mul_quant_varlen_kernel<true, true, kSwizzle, kUsePDL, kApplySwigluLimit>;
static void
run(const tvm::ffi::TensorView input,
const tvm::ffi::TensorView output,
const tvm::ffi::TensorView output_scale,
const tvm::ffi::TensorView masked_m,
const uint32_t topk,
const bool transposed,
const double swiglu_limit) {
using namespace host;
auto device = SymbolicDevice{};
auto E = SymbolicSize{"num_experts"};
auto T = SymbolicSize{"num_tokens_padded"};
auto D = SymbolicSize{"hidden_dim x 2"};
auto N = SymbolicSize{"hidden_dim"};
auto G = SymbolicSize{"num_groups"};
device.set_options<kDLCUDA>();
TensorMatcher({E, T, D}) // input
.with_dtype<bf16_t>()
.with_device(device)
.verify(input);
TensorMatcher({E, T, N}) // output
.with_dtype<fp8_e4m3_t>()
.with_device(device)
.verify(output);
if (!transposed) {
TensorMatcher({E, T, G}) //
.with_dtype<fp32_t>()
.with_device(device)
.verify(output_scale);
} else {
RuntimeCheck(kScaleUE8M0, "transposed layout only supports scale_ue8m0=true");
auto G_ = SymbolicSize{"G // 4"};
TensorMatcher({E, G_, T}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(output_scale);
G.set_value(G_.unwrap() * 4);
}
TensorMatcher({E}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(masked_m);
const auto num_experts = static_cast<uint32_t>(E.unwrap());
const auto num_tokens = static_cast<uint32_t>(T.unwrap());
const auto num_groups = static_cast<uint32_t>(G.unwrap());
const auto hidden_dim = N.unwrap();
RuntimeCheck(D.unwrap() == 2 * hidden_dim, "invalid dimension");
RuntimeCheck(hidden_dim % kGroupSize == 0);
RuntimeCheck(num_experts <= kMaxExperts, "num_experts exceeds maximum (256)");
RuntimeCheck(num_groups * kGroupSize == hidden_dim, "invalid num_groups");
const auto params = SiluMulQuantVarlenParams{
.input = static_cast<const bf16_t*>(input.data_ptr()),
.output = static_cast<fp8_e4m3_t*>(output.data_ptr()),
.output_scale = static_cast<float*>(output_scale.data_ptr()),
.masked_m = static_cast<const int32_t*>(masked_m.data_ptr()),
.swiglu_limit = static_cast<float>(swiglu_limit),
.hidden_dim = hidden_dim,
.num_tokens = num_tokens,
.num_experts = num_experts,
};
const auto num_threads = hidden_dim / 8;
RuntimeCheck(num_threads % device::kWarpThreads == 0);
RuntimeCheck(num_threads >= num_experts);
const auto kernel = transposed ? kernel_transposed : kernel_normal;
LaunchKernel(num_tokens * topk, num_threads, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
template <typename DType, bool kUsePDL>
struct SiluAndMulClampKernel {
static constexpr auto kernel = silu_mul_clamp_kernel<DType, kUsePDL>;
static void run(const tvm::ffi::TensorView input, const tvm::ffi::TensorView output, const double swiglu_limit) {
using namespace host;
auto device = SymbolicDevice{};
auto M = SymbolicSize{"num_tokens"};
auto D = SymbolicSize{"gate_up_dim"}; // 2 * out_dim
auto H = SymbolicSize{"out_dim"};
device.set_options<kDLCUDA>();
TensorMatcher({M, D}) // input (gate || up)
.with_dtype<DType>()
.with_device(device)
.verify(input);
TensorMatcher({M, H}) // output
.with_dtype<DType>()
.with_device(device)
.verify(output);
RuntimeCheck(D.unwrap() == 2 * H.unwrap(), "input last dim must be 2 * output last dim");
constexpr uint32_t kVecSize = 16 / sizeof(DType);
const auto out_dim = static_cast<uint32_t>(H.unwrap());
const auto num_tokens = static_cast<uint32_t>(M.unwrap());
RuntimeCheck(out_dim % kVecSize == 0, "out_dim must be divisible by vector size");
const auto num_threads = out_dim / kVecSize;
RuntimeCheck(num_threads <= 1024, "out_dim too large for single-block-per-row launch");
const auto params = SiluAndMulClampParams{
.input = input.data_ptr(),
.output = output.data_ptr(),
.swiglu_limit = static_cast<float>(swiglu_limit),
};
LaunchKernel(num_tokens, num_threads, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
struct SiluMulQuantContigParams {
const bf16_t* __restrict__ input;
fp8_e4m3_t* __restrict__ output;
float* __restrict__ output_scale;
float swiglu_limit; // only read when kApplySwigluLimit=true
int64_t hidden_dim;
uint32_t num_tokens;
uint32_t scale_row_stride_int32; // only used when kTransposed=true
};
template <bool kScaleUE8M0, bool kTransposed, bool kSwizzle, bool kUsePDL, bool kApplySwigluLimit>
__global__ __launch_bounds__(1024, 2) void // maximize occupancy
silu_mul_quant_contig_kernel(const SiluMulQuantContigParams __grid_constant__ params) {
using namespace device;
constexpr uint32_t kGroupSize = 128u;
constexpr uint32_t kWorkThreads = 16u;
using InputVec = AlignedVector<bf16x2_t, 4>;
using OutputVec = AlignedVector<fp8x2_e4m3_t, 4>;
static_assert(8 * kWorkThreads == 128, "Invalid tiling");
static_assert(!(kTransposed && !kScaleUE8M0), "transposed layout only supports ue8m0");
const auto token_id = blockIdx.x;
const auto work_id = threadIdx.x / kWorkThreads;
const auto input = params.input + token_id * params.hidden_dim * 2;
const auto output = params.output + token_id * params.hidden_dim;
[[maybe_unused]]
const auto output_scale = [&] {
const auto num_groups = params.hidden_dim / kGroupSize;
if constexpr (kTransposed) {
// Physical layout is (G//4_pad, M_pad) int32; each int32 packs 4
// consecutive UE8M0 exponents for the same token. Byte address:
// (work_id / 4) * M_pad * 4 + token * 4 + (work_id % 4).
const auto base = reinterpret_cast<uint8_t*>(params.output_scale);
return base + (work_id / 4u) * (params.scale_row_stride_int32 * 4u) + token_id * 4u + (work_id % 4u);
} else {
return params.output_scale + token_id * num_groups + work_id;
}
}();
PDLWaitPrimary<kUsePDL>();
InputVec gate_vec, up_vec;
if constexpr (kSwizzle) {
gate_vec.load(input, threadIdx.x * 2);
up_vec.load(input, threadIdx.x * 2 + 1);
} else {
gate_vec.load(input, threadIdx.x);
up_vec.load(input, threadIdx.x + blockDim.x);
}
float local_max = 0.0f;
float results[8];
#pragma unroll
for (uint32_t i = 0; i < 4; ++i) {
const auto [x, y] = silu_and_mul<kApplySwigluLimit>(gate_vec[i], up_vec[i], params.swiglu_limit);
results[2 * i + 0] = x;
results[2 * i + 1] = y;
local_max = fmaxf(local_max, fmaxf(fabsf(x), fabsf(y)));
}
local_max = warp::reduce_max<kWorkThreads>(local_max);
const float absmax = fmaxf(local_max, 1e-10f);
float scale;
uint32_t ue8m0_exp;
if constexpr (kScaleUE8M0) {
const float raw_scale = absmax / math::FP8_E4M3_MAX;
ue8m0_exp = cast_to_ue8m0(raw_scale);
scale = __uint_as_float(ue8m0_exp << 23);
} else {
scale = absmax / math::FP8_E4M3_MAX;
}
const auto inv_scale = 1.0f / scale;
OutputVec out_vec;
#pragma unroll
for (uint32_t i = 0; i < 4; ++i) {
const float scaled_val0 = results[2 * i + 0] * inv_scale;
const float scaled_val1 = results[2 * i + 1] * inv_scale;
out_vec[i] = pack_fp8(scaled_val0, scaled_val1);
}
PDLTriggerSecondary<kUsePDL>();
out_vec.store(output, threadIdx.x);
if constexpr (kTransposed) {
*output_scale = ue8m0_exp;
} else {
*output_scale = scale;
}
}
template <int64_t kGroupSize, bool kScaleUE8M0, bool kSwizzle, bool kUsePDL, bool kApplySwigluLimit>
struct SiluAndMulContigPostQuantKernel {
static_assert(kGroupSize == 128);
static constexpr auto kernel_normal =
silu_mul_quant_contig_kernel<kScaleUE8M0, false, kSwizzle, kUsePDL, kApplySwigluLimit>;
static constexpr auto kernel_transposed =
silu_mul_quant_contig_kernel<true, true, kSwizzle, kUsePDL, kApplySwigluLimit>;
static void
run(const tvm::ffi::TensorView input,
const tvm::ffi::TensorView output,
const tvm::ffi::TensorView output_scale,
const bool transposed,
const double swiglu_limit) {
using namespace host;
auto device = SymbolicDevice{};
auto M = SymbolicSize{"num_tokens"};
auto D = SymbolicSize{"hidden_dim x 2"};
auto N = SymbolicSize{"hidden_dim"};
auto G = SymbolicSize{"num_groups"};
device.set_options<kDLCUDA>();
TensorMatcher({M, D}) // input (gate/up, natural or gran=8 interleaved on last dim)
.with_dtype<bf16_t>()
.with_device(device)
.verify(input);
TensorMatcher({M, N}) // fp8 output
.with_dtype<fp8_e4m3_t>()
.with_device(device)
.verify(output);
const auto hidden_dim = N.unwrap();
RuntimeCheck(D.unwrap() == 2 * hidden_dim, "invalid dimension");
RuntimeCheck(hidden_dim % kGroupSize == 0);
const auto num_groups = static_cast<uint32_t>(hidden_dim / kGroupSize);
uint32_t scale_row_stride_int32 = 0;
if (!transposed) {
G.set_value(num_groups);
TensorMatcher({M, G}) // (M, G) fp32 natural row-major
.with_dtype<fp32_t>()
.with_device(device)
.verify(output_scale);
} else {
RuntimeCheck(kScaleUE8M0, "transposed layout only supports scale_ue8m0=true");
RuntimeCheck(num_groups % 4 == 0, "transposed layout requires num_groups % 4 == 0");
auto G_ = SymbolicSize{"G // 4"};
G_.set_value(num_groups / 4);
auto M_pad = SymbolicSize{"M padded"};
TensorMatcher({M, G_}) // `.transpose(-1,-2)[:M,:]` view of (G//4_pad, M_pad) int32
.with_strides({int64_t{1}, M_pad}) // col-major transposed
.with_dtype<int32_t>()
.with_device(device)
.verify(output_scale);
scale_row_stride_int32 = static_cast<uint32_t>(M_pad.unwrap());
}
const auto num_tokens = static_cast<uint32_t>(M.unwrap());
const auto params = SiluMulQuantContigParams{
.input = static_cast<const bf16_t*>(input.data_ptr()),
.output = static_cast<fp8_e4m3_t*>(output.data_ptr()),
.output_scale = static_cast<float*>(output_scale.data_ptr()),
.swiglu_limit = static_cast<float>(swiglu_limit),
.hidden_dim = hidden_dim,
.num_tokens = num_tokens,
.scale_row_stride_int32 = scale_row_stride_int32,
};
const auto num_threads = hidden_dim / 8;
RuntimeCheck(num_threads % device::kWarpThreads == 0);
const auto kernel = transposed ? kernel_transposed : kernel_normal;
LaunchKernel(num_tokens, num_threads, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/store.cuh
+205
View File
@@ -0,0 +1,205 @@
#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 <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <bit>
#include <cstdint>
#include <cuda_fp8.h>
namespace {
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::inv_scale_ue8m0;
using deepseek_v4::fp8::pack_fp8;
struct FusedStoreCacheParam {
const void* __restrict__ input;
void* __restrict__ cache;
const void* __restrict__ indices;
uint32_t num_tokens;
};
template <typename Float, typename IndicesT, uint32_t kPageBits, bool kUsePDL>
__global__ void fused_store_flashmla_cache(const __grid_constant__ FusedStoreCacheParam param) {
using namespace device;
/// NOTE: 584 = 576 + 8
constexpr int64_t kPageBytes = host::div_ceil(584 << kPageBits, 576) * 576;
// each warp handles 64 elements, 8 warps, each block handles 1 row
const auto& [input, cache, indices, num_tokens] = param;
const uint32_t bid = blockIdx.x;
const uint32_t tid = threadIdx.x;
const uint32_t wid = tid / 32;
PDLWaitPrimary<kUsePDL>();
// prefetch the index
const auto index = static_cast<const IndicesT*>(indices)[bid];
// always load the value from input (don't store if invalid)
using Float2 = packed_t<Float>;
const auto elems = static_cast<const Float2*>(input)[tid + bid * 256];
if (wid != 7) {
const auto [x, y] = cast<fp32x2_t>(elems);
const auto abs_max = warp::reduce_max(fmaxf(fabs(x), fabs(y)));
const auto scale_raw = fmaxf(1e-4f, abs_max) / kFP8E4M3Max;
const auto scale_ue8m0 = cast_to_ue8m0(scale_raw);
const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
const auto result = pack_fp8(x * inv_scale, y * inv_scale);
const int32_t page = index >> kPageBits;
const int32_t offset = index & ((1 << kPageBits) - 1);
const auto page_ptr = pointer::offset(cache, page * kPageBytes);
const auto value_ptr = pointer::offset(page_ptr, offset * 576);
const auto scale_ptr = pointer::offset(page_ptr, 576 << kPageBits, offset * 8);
static_cast<fp8x2_e4m3_t*>(value_ptr)[tid] = result;
static_cast<uint8_t*>(scale_ptr)[wid] = scale_ue8m0;
} else {
const auto result = cast<bf16x2_t>(elems);
const int32_t page = index >> kPageBits;
const int32_t offset = index & ((1 << kPageBits) - 1);
const auto page_ptr = pointer::offset(cache, page * kPageBytes);
const auto value_ptr = pointer::offset(page_ptr, offset * 576, 448);
static_cast<bf16x2_t*>(value_ptr)[tid - 7 * 32] = result;
}
PDLTriggerSecondary<kUsePDL>();
}
template <typename Float, typename IndicesT, uint32_t kPageBits, bool kUsePDL>
__global__ void fused_store_indexer_cache(const __grid_constant__ FusedStoreCacheParam param) {
using namespace device;
/// NOTE: 132 = 128 + 4
constexpr int64_t kPageBytes = 132 << kPageBits;
// each warp handles 128 elements, 1 warp, each block handles multiple rows
const auto& [input, cache, indices, num_tokens] = param;
const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x;
const auto global_wid = global_tid / 32;
const auto lane_id = threadIdx.x % 32;
if (global_wid >= num_tokens) return;
PDLWaitPrimary<kUsePDL>();
// prefetch the index
const auto index = static_cast<const IndicesT*>(indices)[global_wid];
// always load the value from input (don't store if invalid)
using Float2 = packed_t<Float>;
using InStorage = AlignedVector<Float2, 2>;
using OutStorage = AlignedVector<fp8x2_e4m3_t, 2>;
const auto elems = static_cast<const InStorage*>(input)[global_tid];
const auto [x0, x1] = cast<fp32x2_t>(elems[0]);
const auto [y0, y1] = cast<fp32x2_t>(elems[1]);
const auto local_max = fmaxf(fmaxf(fabs(x0), fabs(x1)), fmaxf(fabs(y0), fabs(y1)));
const auto abs_max = warp::reduce_max(local_max);
// use normal fp32 scale
const auto scale = fmaxf(1e-4f, abs_max) / kFP8E4M3Max;
const auto inv_scale = 1.0f / scale;
const int32_t page = index >> kPageBits;
const int32_t offset = index & ((1 << kPageBits) - 1);
const auto page_ptr = pointer::offset(cache, page * kPageBytes);
const auto value_ptr = pointer::offset(page_ptr, offset * 128);
const auto scale_ptr = pointer::offset(page_ptr, 128 << kPageBits, offset * 4);
OutStorage result;
result[0] = pack_fp8(x0 * inv_scale, x1 * inv_scale);
result[1] = pack_fp8(y0 * inv_scale, y1 * inv_scale);
static_cast<OutStorage*>(value_ptr)[lane_id] = result;
static_cast<float*>(scale_ptr)[0] = scale;
PDLTriggerSecondary<kUsePDL>();
}
template <typename Float, typename IndicesT, uint32_t kPageSize, bool kUsePDL>
struct FusedStoreCacheFlashMLAKernel {
static constexpr int32_t kLogSize = std::countr_zero(kPageSize);
static constexpr int64_t kPageBytes = host::div_ceil(584 * kPageSize, 576) * 576;
static constexpr auto kernel = fused_store_flashmla_cache<Float, IndicesT, kLogSize, kUsePDL>;
static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
static_assert(1 << kLogSize == kPageSize);
static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView cache, tvm::ffi::TensorView indices) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({N, 512}) // input
.with_dtype<Float>()
.with_device(device_)
.verify(input);
TensorMatcher({-1, -1}) // cache
.with_strides({kPageBytes, 1})
.with_dtype<uint8_t>()
.with_device(device_)
.verify(cache);
TensorMatcher({N}) // indices
.with_dtype<IndicesT>()
.with_device(device_)
.verify(indices);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto params = FusedStoreCacheParam{
.input = input.data_ptr(),
.cache = cache.data_ptr(),
.indices = indices.data_ptr(),
.num_tokens = num_tokens,
};
const auto kBlockSize = 256;
const auto num_blocks = num_tokens;
LaunchKernel(num_blocks, kBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(kernel, params);
}
};
template <typename Float, typename IndicesT, uint32_t kPageSize, bool kUsePDL>
struct FusedStoreCacheIndexerKernel {
static constexpr int32_t kLogSize = std::countr_zero(kPageSize);
static constexpr int64_t kPageBytes = 132 * kPageSize;
static constexpr auto kernel = fused_store_indexer_cache<Float, IndicesT, kLogSize, kUsePDL>;
static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
static_assert(1 << kLogSize == kPageSize);
static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView cache, tvm::ffi::TensorView indices) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({N, 128}) // input
.with_dtype<Float>()
.with_device(device_)
.verify(input);
TensorMatcher({-1, -1}) // cache
.with_strides({kPageBytes, 1})
.with_dtype<uint8_t>()
.with_device(device_)
.verify(cache);
TensorMatcher({N}) // indices
.with_dtype<IndicesT>()
.with_device(device_)
.verify(indices);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto params = FusedStoreCacheParam{
.input = input.data_ptr(),
.cache = cache.data_ptr(),
.indices = indices.data_ptr(),
.num_tokens = num_tokens,
};
const auto kBlockSize = 128;
const auto num_blocks = div_ceil(num_tokens * 32, kBlockSize);
LaunchKernel(num_blocks, kBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/topk_v1.cuh
+340
View File
@@ -0,0 +1,340 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <bit>
#include <cstdint>
namespace {
#ifndef SGL_TOPK
#define SGL_TOPK 512
#endif
constexpr uint32_t kTopK = SGL_TOPK;
constexpr uint32_t kTopKBlockSize = SGL_TOPK;
constexpr uint32_t kSMEM = 16 * 1024 * sizeof(uint32_t); // 64KB (bytes)
struct TopKParams {
const float* __restrict__ scores;
const int32_t* __restrict__ seq_lens;
const int32_t* __restrict__ page_table;
int32_t* __restrict__ page_indices;
int32_t* __restrict__ raw_indices; // optional: output raw abs position indices before page transform
const int64_t score_stride;
const int64_t page_table_stride;
uint32_t page_bits;
};
SGL_DEVICE uint8_t convert_to_uint8(float x) {
__half h = __float2half_rn(x);
uint16_t bits = __half_as_ushort(h);
uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits) : static_cast<uint16_t>(bits | 0x8000);
return static_cast<uint8_t>(key >> 8);
}
SGL_DEVICE uint32_t convert_to_uint32(float x) {
uint32_t bits = __float_as_uint(x);
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
}
SGL_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table, uint32_t i, uint32_t page_bits) {
const uint32_t mask = (1u << page_bits) - 1u;
return (page_table[i >> page_bits] << page_bits) | (i & mask);
}
[[maybe_unused]]
SGL_DEVICE void naive_transform(
const float* __restrict__, // unused
const int32_t* __restrict__ page_table,
int32_t* __restrict__ indices,
int32_t* __restrict__ raw_indices, // optional: output raw abs position indices
const uint32_t length,
const uint32_t page_bits) {
static_assert(kTopK <= kTopKBlockSize);
if (const auto tx = threadIdx.x; tx < length) {
indices[tx] = page_to_indices(page_table, tx, page_bits);
if (raw_indices != nullptr) {
raw_indices[tx] = tx;
}
} else if (kTopK == kTopKBlockSize || tx < kTopK) {
indices[tx] = -1; // fill invalid indices to -1
if (raw_indices != nullptr) {
raw_indices[tx] = -1;
}
}
}
[[maybe_unused]]
SGL_DEVICE void radix_topk(const float* __restrict__ input, int32_t* __restrict__ output, const uint32_t length) {
constexpr uint32_t RADIX = 256;
constexpr uint32_t BLOCK_SIZE = kTopKBlockSize;
constexpr uint32_t SMEM_INPUT_SIZE = kSMEM / (2 * sizeof(int32_t));
alignas(128) __shared__ uint32_t _s_histogram_buf[2][RADIX + 32];
alignas(128) __shared__ uint32_t s_counter;
alignas(128) __shared__ uint32_t s_threshold_bin_id;
alignas(128) __shared__ uint32_t s_num_input[2];
alignas(128) __shared__ int32_t s_last_remain;
extern __shared__ uint32_t s_input_idx[][kSMEM / (2 * sizeof(int32_t))];
const uint32_t tx = threadIdx.x;
uint32_t remain_topk = kTopK;
auto& s_histogram = _s_histogram_buf[0];
const auto run_cumsum = [&] {
#pragma unroll 8
for (int32_t i = 0; i < 8; ++i) {
static_assert(1 << 8 == RADIX);
if (tx < RADIX) {
const auto j = 1 << i;
const auto k = i & 1;
auto value = _s_histogram_buf[k][tx];
if (tx + j < RADIX) {
value += _s_histogram_buf[k][tx + j];
}
_s_histogram_buf[k ^ 1][tx] = value;
}
__syncthreads();
}
};
// stage 1: 8bit coarse histogram
if (tx < RADIX + 1) s_histogram[tx] = 0;
__syncthreads();
for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) {
const auto bin = convert_to_uint8(input[idx]);
::atomicAdd(&s_histogram[bin], 1);
}
__syncthreads();
run_cumsum();
if (tx < RADIX && s_histogram[tx] > remain_topk && s_histogram[tx + 1] <= remain_topk) {
s_threshold_bin_id = tx;
s_num_input[0] = 0;
s_counter = 0;
}
__syncthreads();
const auto threshold_bin = s_threshold_bin_id;
remain_topk -= s_histogram[threshold_bin + 1];
if (remain_topk == 0) {
for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) {
const uint32_t bin = convert_to_uint8(input[idx]);
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
output[pos] = idx;
}
}
__syncthreads();
return;
} else {
__syncthreads();
if (tx < RADIX + 1) {
s_histogram[tx] = 0;
}
__syncthreads();
for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) {
const float raw_input = input[idx];
const uint32_t bin = convert_to_uint8(raw_input);
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
output[pos] = idx;
} else if (bin == threshold_bin) {
const auto pos = ::atomicAdd(&s_num_input[0], 1);
if (pos < SMEM_INPUT_SIZE) {
[[likely]] s_input_idx[0][pos] = idx;
const auto bin = convert_to_uint32(raw_input);
const auto sub_bin = (bin >> 24) & 0xFF;
::atomicAdd(&s_histogram[sub_bin], 1);
}
}
}
__syncthreads();
}
// stage 2: refine with 8bit radix passes
#pragma unroll 4
for (int round = 0; round < 4; ++round) {
const auto r_idx = round % 2;
// clip here to prevent overflow
const auto raw_num_input = s_num_input[r_idx];
const auto num_input = raw_num_input < SMEM_INPUT_SIZE ? raw_num_input : SMEM_INPUT_SIZE;
run_cumsum();
if (tx < RADIX && s_histogram[tx] > remain_topk && s_histogram[tx + 1] <= remain_topk) {
s_threshold_bin_id = tx;
s_num_input[r_idx ^ 1] = 0;
s_last_remain = remain_topk - s_histogram[tx + 1];
}
__syncthreads();
const auto threshold_bin = s_threshold_bin_id;
remain_topk -= s_histogram[threshold_bin + 1];
if (remain_topk == 0) {
for (uint32_t i = tx; i < num_input; i += BLOCK_SIZE) {
const auto idx = s_input_idx[r_idx][i];
const auto offset = 24 - round * 8;
const auto bin = (convert_to_uint32(input[idx]) >> offset) & 0xFF;
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
output[pos] = idx;
}
}
__syncthreads();
break;
} else {
__syncthreads();
if (tx < RADIX + 1) {
s_histogram[tx] = 0;
}
__syncthreads();
for (uint32_t i = tx; i < num_input; i += BLOCK_SIZE) {
const auto idx = s_input_idx[r_idx][i];
const auto raw_input = input[idx];
const auto offset = 24 - round * 8;
const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF;
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
output[pos] = idx;
} else if (bin == threshold_bin) {
if (round == 3) {
const auto pos = ::atomicAdd(&s_last_remain, -1);
if (pos > 0) {
output[kTopK - pos] = idx;
}
} else {
const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1);
if (pos < SMEM_INPUT_SIZE) {
/// NOTE: (dark) fuse the histogram computation here
[[likely]] s_input_idx[r_idx ^ 1][pos] = idx;
const auto bin = convert_to_uint32(raw_input);
const auto sub_bin = (bin >> (offset - 8)) & 0xFF;
::atomicAdd(&s_histogram[sub_bin], 1);
}
}
}
}
__syncthreads();
}
}
}
template <bool kUsePDL>
__global__ void topk_transform_kernel(const __grid_constant__ TopKParams params) {
const auto &[
scores, seq_lens, page_table, page_indices, raw_indices, // pointers
score_stride, page_table_stride, page_bits // sizes
] = params;
const uint32_t work_id = blockIdx.x;
/// NOTE: dangerous prefetch seq_len before PDL wait
const uint32_t seq_len = seq_lens[work_id];
const auto score_ptr = scores + work_id * score_stride;
const auto page_ptr = page_table + work_id * page_table_stride;
const auto indices_ptr = page_indices + work_id * kTopK;
const auto raw_indices_ptr = raw_indices != nullptr ? raw_indices + work_id * kTopK : nullptr;
device::PDLWaitPrimary<kUsePDL>();
if (seq_len <= kTopK) {
naive_transform(score_ptr, page_ptr, indices_ptr, raw_indices_ptr, seq_len, page_bits);
} else {
__shared__ int32_t s_topk_indices[kTopK];
radix_topk(score_ptr, s_topk_indices, seq_len);
static_assert(kTopK <= kTopKBlockSize);
const auto tx = threadIdx.x;
if (kTopK == kTopKBlockSize || tx < kTopK) {
indices_ptr[tx] = page_to_indices(page_ptr, s_topk_indices[tx], page_bits);
if (raw_indices_ptr != nullptr) {
raw_indices_ptr[tx] = s_topk_indices[tx];
}
}
}
device::PDLTriggerSecondary<kUsePDL>();
}
template <auto* f, size_t kMaxDynamicSMEM>
void setup_kernel_smem_once(host::DebugInfo where = {}) {
[[maybe_unused]]
static const auto result = [] {
const auto fptr = std::bit_cast<const void*>(f);
return ::cudaFuncSetAttribute(fptr, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM);
}();
host::RuntimeDeviceCheck(result, where);
}
template <bool kUsePDL>
struct TopKKernel {
static constexpr auto kernel = topk_transform_kernel<kUsePDL>;
static void transform(
const tvm::ffi::TensorView scores,
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::TensorView page_table,
const tvm::ffi::TensorView page_indices,
const uint32_t page_size,
const tvm::ffi::Optional<tvm::ffi::TensorView> raw_indices) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto S = SymbolicSize{"score_stride"};
auto P = SymbolicSize{"page_table_stride"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({B, -1}) // strided scores
.with_strides({S, 1})
.with_dtype<float>()
.with_device(device)
.verify(scores);
TensorMatcher({B}) // seq_lens, must be contiguous
.with_dtype<int32_t>()
.with_device(device)
.verify(seq_lens);
TensorMatcher({B, -1}) // strided page table
.with_strides({P, 1})
.with_dtype<int32_t>()
.with_device(device)
.verify(page_table);
TensorMatcher({B, kTopK}) // output, must be contiguous
.with_dtype<int32_t>()
.with_device(device)
.verify(page_indices);
int32_t* raw_indices_ptr = nullptr;
if (raw_indices.has_value()) {
TensorMatcher({B, kTopK}) // optional raw indices output, must be contiguous
.with_dtype<int32_t>()
.with_device(device)
.verify(raw_indices.value());
raw_indices_ptr = static_cast<int32_t*>(raw_indices.value().data_ptr());
}
RuntimeCheck(std::has_single_bit(page_size), "page_size must be power of 2");
const auto page_bits = static_cast<uint32_t>(std::countr_zero(page_size));
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto params = TopKParams{
.scores = static_cast<float*>(scores.data_ptr()),
.seq_lens = static_cast<int32_t*>(seq_lens.data_ptr()),
.page_table = static_cast<int32_t*>(page_table.data_ptr()),
.page_indices = static_cast<int32_t*>(page_indices.data_ptr()),
.raw_indices = raw_indices_ptr,
.score_stride = S.unwrap(),
.page_table_stride = P.unwrap(),
.page_bits = page_bits,
};
constexpr auto kSMEM_ = kSMEM + sizeof(int32_t); // align up a little
setup_kernel_smem_once<kernel, kSMEM_>();
LaunchKernel(batch_size, kTopKBlockSize, device.unwrap(), kSMEM_).enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh
+458
View File
@@ -0,0 +1,458 @@
/**
* \file topk_v2.cuh
* \brief TopK kernel for DeepSeek v4.
* Adapted from
* 1:
* https://github.com/vllm-project/vllm/blob/a8c6ee9b787d273916206a29b77feebadb80c368/csrc/persistent_topk.cuh
* 2:
* https://github.com/flashinfer-ai/flashinfer/blob/c2b4db2b1a84448d802f0e6ac445243312bd6a4c/include/flashinfer/topk.cuh
* DarkSharpness never took a detailed look at these 2 implementation, but his claude code did.
* So we add credit to the reference implementations.
*/
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/deepseek_v4/topk_impl.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <bit>
#include <cstdint>
#include <iterator>
namespace {
namespace impl = device::topk;
using impl::TopKProblem;
using Register2 = impl::TopKRegister<2>; // <= 8192, register-resident, 1 read
using Register4 = impl::TopKRegister<4>; // <= 16384, register-resident, 1 read
using Streaming = impl::TopKStreaming;
using Cluster = impl::TopKCluster<8>;
constexpr uint32_t kBlockSize = impl::TopKConfig::kBlockSize;
constexpr uint32_t kOccupancy = impl::TopKConfig::kOccupancy;
constexpr uint32_t kMaxTopK = impl::TopKConfig::kMaxTopK;
constexpr uint32_t kClusterSize = Cluster::kClusterSize;
constexpr uint32_t kReg2MaxSeqLen = Register2::kMaxSeqLen; // 8192
constexpr uint32_t kReg4MaxSeqLen = Register4::kMaxSeqLen; // 16384
#define TOPK_KERNEL __global__ __launch_bounds__(kBlockSize, kOccupancy)
#define CLUSTER_TOPK_KERNEL TOPK_KERNEL __cluster_dims__(1, kClusterSize, 1)
constexpr uint32_t kClusterFloor = 65536;
constexpr uint32_t kClusterMaxBatch = 512;
constexpr uint32_t kNumPersistentClusters = 15 * kOccupancy;
/// Metadata tensor rows (each 8 B / 2 int32). Row 0 is the global plan result;
/// rows 1..N are the (batch_id, seq_len) of items routed to the cluster pool.
struct alignas(8) GlobalMetadata {
uint32_t cluster_threshold;
uint32_t num_cluster_items; // N = number of items routed to the cluster pool
};
struct alignas(8) PlanItem {
uint32_t batch_id;
uint32_t seq_len;
};
static_assert(sizeof(GlobalMetadata) == 2 * sizeof(int32_t) && sizeof(PlanItem) == sizeof(GlobalMetadata));
struct TopKLaunchParams {
const float* __restrict__ scores;
const int32_t* __restrict__ seq_lens;
const int32_t* __restrict__ page_table;
int32_t* __restrict__ page_indices;
int32_t* __restrict__ raw_indices; // optional raw (pre-transform) indices output; nullptr if unused
const PlanItem* __restrict__ metadata; // [0]=GlobalMetadata, [1+i]=PlanItem
int64_t score_stride;
int64_t page_table_stride;
uint32_t topk;
uint32_t page_bits;
uint32_t cluster_floor; // seq_len > this routes to the cluster path (batch-aware, host-set)
SGL_DEVICE const GlobalMetadata& global() const {
return *reinterpret_cast<const GlobalMetadata*>(metadata);
}
SGL_DEVICE uint32_t cluster_threshold() const {
return global().cluster_threshold;
}
SGL_DEVICE const PlanItem& item(uint32_t i) const {
return metadata[1 + i];
}
SGL_DEVICE int32_t* get_output_ptr(uint32_t batch_id) const {
return page_indices + batch_id * static_cast<int64_t>(topk);
}
SGL_DEVICE TopKProblem problem(uint32_t batch_id, uint32_t seq_len) const {
const auto k = static_cast<int64_t>(topk);
return TopKProblem{
.in = scores + batch_id * score_stride,
.out = page_indices + batch_id * k,
.raw_out = raw_indices != nullptr ? raw_indices + batch_id * k : nullptr,
.page_table = page_table + batch_id * page_table_stride,
.topk = topk,
.seq_len = seq_len,
.page_bits = page_bits,
};
}
SGL_DEVICE TopKProblem problem(uint32_t batch_id) const {
return this->problem(batch_id, static_cast<uint32_t>(seq_lens[batch_id]));
}
};
/**
* \brief Persistent cluster kernel for the long items. It will handle long inputs.
* The short items are handled by the separate topk_kernel.
*/
template <bool kPDL>
CLUSTER_TOPK_KERNEL void topk_persistent_cluster_kernel(const __grid_constant__ TopKLaunchParams params) {
device::enable_smem_spilling();
__shared__ impl::MaxSmem<Cluster::Smem> smem;
const uint32_t num_cluster_items = params.global().num_cluster_items;
device::PDLWaitPrimary<kPDL>();
device::PDLTriggerSecondary<kPDL>();
#pragma unroll 1
for (uint32_t w = blockIdx.x; w < num_cluster_items; w += kNumPersistentClusters) {
const auto it = params.item(w);
const auto problem = params.problem(it.batch_id, it.seq_len);
Cluster::forward<false>(problem, &smem);
__syncthreads();
}
}
template <typename F>
SGL_DEVICE void for_each_item(uint32_t topk, const F& f) {
constexpr uint32_t kNumElems = kMaxTopK / kBlockSize;
#pragma unroll
for (uint32_t i = 0; i < kNumElems; ++i) {
if (const auto tx = i * kBlockSize + threadIdx.x; tx < topk) {
__builtin_assume(tx < kMaxTopK);
f(tx, i);
}
}
}
template <bool kPDL>
SGL_DEVICE void trivial_transform(const TopKProblem& problem) {
device::PDLWaitPrimary<kPDL>();
device::PDLTriggerSecondary<kPDL>();
for_each_item(problem.topk, [&](uint32_t tx, uint32_t) {
problem.transform_output(tx, tx < problem.seq_len ? static_cast<int32_t>(tx) : -1);
});
}
SGL_DEVICE void problem_transform(TopKProblem& problem, int32_t* output_ptr) {
static_assert(kMaxTopK % kBlockSize == 0);
constexpr uint32_t kNumElems = kMaxTopK / kBlockSize;
int32_t source_index[kNumElems];
for_each_item(problem.topk, [&](uint32_t tx, uint32_t i) { source_index[i] = problem.out[tx]; });
problem.out = output_ptr;
for_each_item(problem.topk, [&](uint32_t tx, uint32_t i) { problem.transform_output(tx, source_index[i]); });
}
/**
* \brief Main kernel for the short items and epilogue of long items.
* \tparam kPDL whether to use PDL to synchronize with the cluster kernel (if any)
* \tparam kLevel:
* - Level 0: max_seq_len <= 8192 -> trivial + register<2>
* - Level 1: max_seq_len <= 16384 -> trivial + register<4>
* - Level 2: max_seq_len <= cluster_floor -> trivial + register<4> + streaming
* - Level 3: max_seq_len > cluster_floor -> + epilogue process of cluster path
*/
template <bool kPDL, int kLevel>
TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKLaunchParams params) {
device::enable_smem_spilling();
auto problem = params.problem(blockIdx.x);
constexpr uint32_t kU32Max = std::numeric_limits<uint32_t>::max();
__shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
if (problem.seq_len <= problem.topk) return trivial_transform<kPDL>(problem);
__shared__ int32_t topk_indices[kMaxTopK];
problem.out = topk_indices;
constexpr bool kHandleCluster = (kLevel == 3);
// non-trivial path: dispatch based on level and seq_len
const auto cluster_threshold = kHandleCluster ? params.cluster_threshold() : kU32Max;
if constexpr (kLevel == 0) {
__builtin_assume(problem.seq_len <= kReg2MaxSeqLen);
Register2::forward<kPDL>(problem, &smem);
} else if constexpr (kLevel == 1) {
__builtin_assume(problem.seq_len <= kReg4MaxSeqLen);
Register4::forward<kPDL>(problem, &smem); // max_seq_len <= 16384 guarantees seq <= 16384
} else {
static_assert(kLevel == 2 || kLevel == 3, "we only support level = 0,1,2,3 now");
// if using cluster, we can delay the PDL wait
constexpr bool kPDLEarly = kPDL && !kHandleCluster;
constexpr bool kPDLFinal = kPDL && kHandleCluster;
if (problem.seq_len <= kReg4MaxSeqLen) {
Register4::forward<kPDLEarly>(problem, &smem);
} else if (problem.seq_len <= cluster_threshold) {
Streaming::forward<kPDLEarly>(problem, &smem);
} else { // cluster path do nothing here
problem.out = params.get_output_ptr(blockIdx.x);
}
device::PDLWaitPrimary<kPDLFinal>();
}
// page-table transform pass (gathers kept out of the hot scatter loop),
// then trigger the dependent kernel only after the full output is written.
device::PDLTriggerSecondary<kPDL>();
__syncthreads();
problem_transform(problem, params.get_output_ptr(blockIdx.x));
}
template <bool kPDL>
CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKLaunchParams params) {
device::enable_smem_spilling();
auto problem = params.problem(blockIdx.x);
__shared__ impl::MaxSmem<Streaming::Smem, Cluster::Smem> smem;
if (problem.seq_len <= problem.topk) return trivial_transform<kPDL>(problem);
__shared__ int32_t topk_indices[kMaxTopK];
problem.out = topk_indices;
// randomly elect one worker rank to avoid workload imbalance
const auto worker_rank = blockIdx.x % kClusterSize;
// for small batch, we will fuse in the cluster case
if (problem.seq_len <= kReg4MaxSeqLen) {
if (blockIdx.y == worker_rank) Register4::forward<kPDL>(problem, &smem);
} else if (problem.seq_len <= params.cluster_floor) {
if (blockIdx.y == worker_rank) Streaming::forward<kPDL>(problem, &smem);
} else {
auto cluster = cooperative_groups::this_cluster();
problem.out = cluster.map_shared_rank(topk_indices, worker_rank);
Cluster::forward<kPDL>(problem, &smem); // write to peer's output shared memory
cluster.sync();
}
device::PDLWaitPrimary<kPDL>();
__syncthreads();
if (blockIdx.y == worker_rank) problem_transform(problem, params.get_output_ptr(blockIdx.x));
}
// --- Plan: choose cluster_threshold from the seq_len distribution -----------
__global__ __launch_bounds__(kBlockSize, 1) void topk_plan(
const uint32_t* __restrict__ seq_lens,
PlanItem* __restrict__ metadata, // [0]=GlobalMetadata, [1+i]=PlanItem
const uint32_t batch_size,
const uint32_t static_cluster_threshold) {
// Candidate (threshold T_j, cap_j) pairs, T strictly increasing. The plan lowers
// cluster_threshold to T_j while #(items with seq_len > T_j) <= cap_j, so cap_j
// bounds how many long items go to the persistent pool. The pool runs N items in
// ceil(N / kNumPersistentClusters) waves; the longer the seq the more waves pay
// off (streaming a single block over a long item is very slow), so cap_j is the
// measured cluster-vs-streaming crossover (B200, occ2) and GROWS with T -- a flat
// cap = pool size only fits the shortest (~98K, one-wave) bucket. (Plan is tunable.)
struct Pair {
uint32_t threshold;
uint32_t max_batch_size;
};
constexpr Pair kCandidates[] = {
{65536, 30}, // (65536,98304]: ~1 pool wave, streams beyond 30
{98304, 48}, // (98304,131072]
{131072, 60}, // (131072,196608]
{196608, 80}, // (196608,262144]
{262144, 112}, // (262144,393216]
{393216, 128}, // (393216,inf): longest -- worth many pool waves; a top
// threshold here lets overloaded ~280-393K batches still stream
};
constexpr uint32_t kNumCandidates = std::size(kCandidates);
static_assert(kCandidates[0].threshold == kClusterFloor);
__shared__ uint32_t s_counts[kNumCandidates];
__shared__ uint32_t s_threshold;
__shared__ uint32_t s_count;
const auto tx = threadIdx.x;
if (tx < kNumCandidates) s_counts[tx] = 0;
if (tx == 0) s_count = 0;
__syncthreads();
if (static_cluster_threshold > 0) {
if (tx == 0) s_threshold = static_cluster_threshold;
} else {
for (uint32_t i = tx; i < batch_size; i += kBlockSize) {
const uint32_t sl = seq_lens[i];
uint32_t count = 0;
#pragma unroll
for (uint32_t j = 0; j < kNumCandidates; ++j) {
count += (sl > kCandidates[j].threshold ? 1 : 0);
}
if (count > 0) atomicAdd(&s_counts[count - 1], 1);
}
__syncthreads();
if (tx == 0) {
uint32_t accum = 0;
uint32_t chosen = kCandidates[kNumCandidates - 1].threshold;
#pragma unroll
for (uint32_t i = 0; i < kNumCandidates; ++i) {
const auto j = kNumCandidates - 1 - i;
accum += s_counts[j]; // # items with seq_len > kCandidates[j].threshold
if (accum > kCandidates[j].max_batch_size) break;
chosen = kCandidates[j].threshold;
}
s_threshold = chosen;
}
}
__syncthreads();
const auto cluster_threshold = max(s_threshold, kClusterFloor);
// Compact items with seq_len > threshold into metadata[1..N]: their batch ids
// are the work list the persistent cluster pool fetches.
for (uint32_t i = tx; i < batch_size; i += kBlockSize) {
const uint32_t sl = seq_lens[i];
if (sl > cluster_threshold) {
const auto pos = atomicAdd(&s_count, 1);
metadata[1 + pos] = {i, sl};
}
}
__syncthreads();
if (tx == 0) {
auto* g = reinterpret_cast<GlobalMetadata*>(metadata);
*g = {.cluster_threshold = cluster_threshold, .num_cluster_items = s_count};
}
}
struct TopKKernel {
static void plan( //
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::TensorView metadata,
const uint32_t static_cluster_threshold) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto Bp1 = SymbolicSize{"batch_size_plus_1"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B}) // seq_lens
.with_dtype<int32_t>()
.with_device(device_)
.verify(seq_lens);
TensorMatcher({Bp1, 2}) // metadata: [0]=GlobalMetadata, [1..N]=PlanItem(batch_id, seq_len)
.with_dtype<int32_t>()
.with_device(device_)
.verify(metadata);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape");
const auto device = device_.unwrap();
LaunchKernel(1, kBlockSize, device)( //
topk_plan,
static_cast<const uint32_t*>(seq_lens.data_ptr()),
static_cast<PlanItem*>(metadata.data_ptr()),
batch_size,
static_cluster_threshold);
}
static void transform(
const tvm::ffi::TensorView scores,
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::TensorView page_table,
const tvm::ffi::TensorView page_indices,
const uint32_t page_size,
const tvm::ffi::TensorView metadata,
const tvm::ffi::Optional<tvm::ffi::TensorView> raw_indices) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto Bp1 = SymbolicSize{"batch_size_plus_1"};
auto L = SymbolicSize{"max_seq_len"};
auto S = SymbolicSize{"score_stride"};
auto P = SymbolicSize{"page_table_stride"};
auto K = SymbolicSize{"topk"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B, L}) // score
.with_strides({S, 1})
.with_dtype<float>()
.with_device(device_)
.verify(scores);
TensorMatcher({B}) // seq_lens
.with_dtype<int32_t>()
.with_device(device_)
.verify(seq_lens);
TensorMatcher({B, -1}) // page_table
.with_strides({P, 1})
.with_dtype<int32_t>()
.with_device(device_)
.verify(page_table);
TensorMatcher({B, K}) // page_indices
.with_dtype<int32_t>()
.with_device(device_)
.verify(page_indices);
TensorMatcher({Bp1, 2}) // metadata: [0]=GlobalMetadata, [1..N]=PlanItem(batch_id, seq_len)
.with_dtype<int32_t>()
.with_device(device_)
.verify(metadata);
int32_t* raw_indices_ptr = nullptr;
if (raw_indices.has_value()) {
TensorMatcher({B, K}).with_dtype<int32_t>().with_device(device_).verify(raw_indices.value());
raw_indices_ptr = static_cast<int32_t*>(raw_indices.value().data_ptr());
}
RuntimeCheck(std::has_single_bit(page_size), "page_size must be power of 2");
RuntimeCheck(S.unwrap() % 4 == 0, "score_stride must be a multiple of 4 (16-byte vectorized load)");
RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape");
const auto topk = static_cast<uint32_t>(K.unwrap());
RuntimeCheck(topk > 0 && topk <= kMaxTopK, "topk must be in (0, 2048]");
const auto page_bits = static_cast<uint32_t>(std::countr_zero(page_size));
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto max_seq_len = static_cast<uint32_t>(L.unwrap());
const auto device = device_.unwrap();
// The fused kernel runs one 8-block cluster per batch element, and B200 fits one
// wave of exactly 15 such clusters (occ2). For batch <= 15 it stays latency-bound,
// so the 8-way split beats streaming from a much lower seq (measured crossover
// ~36-40K); batch 16 spills into a 2nd wave (+25%) and keeps the 64K floor.
// The floor is chosen on the host per launch.
constexpr uint32_t kClusterFloorSmall = 32768;
constexpr uint32_t kSmallBatchLowFloor = 15;
const auto params = TopKLaunchParams{
.scores = static_cast<const float*>(scores.data_ptr()),
.seq_lens = static_cast<const int32_t*>(seq_lens.data_ptr()),
.page_table = static_cast<const int32_t*>(page_table.data_ptr()),
.page_indices = static_cast<int32_t*>(page_indices.data_ptr()),
.raw_indices = raw_indices_ptr,
.metadata = static_cast<const PlanItem*>(metadata.data_ptr()),
.score_stride = S.unwrap(),
.page_table_stride = P.unwrap(),
.topk = topk,
.page_bits = page_bits,
.cluster_floor = (batch_size <= kSmallBatchLowFloor) ? kClusterFloorSmall : kClusterFloor,
};
const bool use_cluster = (max_seq_len > params.cluster_floor) && (batch_size <= kClusterMaxBatch);
constexpr bool kUsePDL = true;
if (use_cluster) {
if (batch_size <= kNumPersistentClusters) {
LaunchKernel({batch_size, kClusterSize}, kBlockSize, device)
.config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
.launch(topk_small_batch_kernel<kUsePDL>, params);
} else {
const uint32_t num_clusters = std::min(batch_size, kNumPersistentClusters);
LaunchKernel({num_clusters, kClusterSize}, kBlockSize, device)
.config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
.launch(topk_persistent_cluster_kernel<kUsePDL>, params);
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/3>, params);
}
} else if (max_seq_len <= kReg2MaxSeqLen) {
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/0>, params);
} else if (max_seq_len <= kReg4MaxSeqLen) {
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/1>, params);
} else {
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/2>, params);
}
}
};
} // namespace
python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh
+253
View File
@@ -0,0 +1,253 @@
// Native CUDA fast path for Cosmos3 VAE causal-Conv3D cat/pad copy.
//
// The op writes the output of:
// pad(cat(cache_x, x, dim=T), (Wl, Wr, Ht, Hb, Dl - cache_t, Dr))
// for 5D NCTHW tensors. It is a memory-bound copy/zero-fill kernel and is only
// entered for contiguous CUDA tensors; unsupported cases fall back to Triton in
// the Python caller.
//
// Developed with MIT HAN Lab Kernel Design Agents:
// https://github.com/mit-han-lab/kernel-design-agents
#pragma once
#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 <sgl_kernel/vec.cuh> // For device::AlignedVector
#include <cstdint>
namespace sglang_causal_conv3d_cat_pad {
namespace {
constexpr int kBlockSize = 256;
template <typename ET, int kVec>
__global__ void __launch_bounds__(kBlockSize) cat_pad_flat_kernel(
const ET* __restrict__ x,
const ET* __restrict__ cache,
ET* __restrict__ out,
int64_t total_vecs,
int64_t channels,
int64_t t_size,
int64_t h_size,
int64_t w_size,
int64_t cache_t,
int64_t out_t,
int64_t out_h,
int64_t out_w,
int64_t pad_d_left,
int64_t pad_h_top,
int64_t pad_w_left) {
using Pack = device::AlignedVector<ET, kVec>;
const int64_t nthreads = static_cast<int64_t>(gridDim.x) * blockDim.x;
for (int64_t vid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; vid < total_vecs; vid += nthreads) {
int64_t base = vid * kVec;
int64_t ow = base % out_w;
int64_t tmp = base / out_w;
int64_t oh = tmp % out_h;
tmp /= out_h;
int64_t od = tmp % out_t;
tmp /= out_t;
int64_t oc = tmp % channels;
int64_t ob = tmp / channels;
int64_t ih = oh - pad_h_top;
int64_t src_t = od - pad_d_left;
bool interior = ih >= 0 && ih < h_size && src_t >= 0 && src_t < cache_t + t_size;
const ET* src = nullptr;
if (interior) {
if (src_t < cache_t) {
src = cache + (((ob * channels + oc) * cache_t + src_t) * h_size + ih) * w_size;
} else {
src = x + (((ob * channels + oc) * t_size + (src_t - cache_t)) * h_size + ih) * w_size;
}
}
Pack pack;
#pragma unroll
for (int i = 0; i < kVec; ++i) {
ET value = ET(0);
if (interior) {
const int64_t iw = ow - pad_w_left;
if (iw >= 0 && iw < w_size) {
value = SGLANG_LDG(src + iw);
}
}
pack[i] = value;
if (++ow == out_w) {
ow = 0;
if (++oh == out_h) {
oh = 0;
if (++od == out_t) {
od = 0;
if (++oc == channels) {
oc = 0;
++ob;
}
}
}
ih = oh - pad_h_top;
src_t = od - pad_d_left;
interior = ih >= 0 && ih < h_size && src_t >= 0 && src_t < cache_t + t_size;
if (interior) {
if (src_t < cache_t) {
src = cache + (((ob * channels + oc) * cache_t + src_t) * h_size + ih) * w_size;
} else {
src = x + (((ob * channels + oc) * t_size + (src_t - cache_t)) * h_size + ih) * w_size;
}
} else {
src = nullptr;
}
}
}
pack.store(out, vid);
}
}
template <typename ET, int kVec>
void launch_cat_pad_flat(
const void* x,
const void* cache,
void* out,
int64_t total,
int64_t channels,
int64_t t_size,
int64_t h_size,
int64_t w_size,
int64_t cache_t,
int64_t out_t,
int64_t out_h,
int64_t out_w,
int64_t depth_left,
int64_t pad_h_top,
int64_t pad_w_left,
DLDevice device) {
const int64_t total_vecs = total / kVec;
const uint32_t grid = static_cast<uint32_t>(host::div_ceil(total_vecs, static_cast<int64_t>(kBlockSize)));
host::LaunchKernel(grid, kBlockSize, device)(
cat_pad_flat_kernel<ET, kVec>,
static_cast<const ET*>(x),
static_cast<const ET*>(cache),
static_cast<ET*>(out),
total_vecs,
channels,
t_size,
h_size,
w_size,
cache_t,
out_t,
out_h,
out_w,
depth_left,
pad_h_top,
pad_w_left);
}
} // namespace
template <typename T>
struct CausalConv3dCatPadKernel {
static void
run(tvm::ffi::TensorView out,
tvm::ffi::TensorView x,
tvm::ffi::TensorView cache,
int64_t pad_w_left,
int64_t pad_w_right,
int64_t pad_h_top,
int64_t pad_h_bottom,
int64_t pad_d_left,
int64_t pad_d_right) {
using namespace host;
auto bsz = SymbolicSize{"batch"};
auto channels = SymbolicSize{"channels"};
auto t_size = SymbolicSize{"t_size"};
auto h_size = SymbolicSize{"h_size"};
auto w_size = SymbolicSize{"w_size"};
auto cache_t = SymbolicSize{"cache_t"};
auto out_t = SymbolicSize{"out_t"};
auto out_h = SymbolicSize{"out_h"};
auto out_w = SymbolicSize{"out_w"};
auto device = SymbolicDevice{};
device.set_options<kDLGPU>();
TensorMatcher({bsz, channels, t_size, h_size, w_size})
.with_dtype<T>()
.template with_device<kDLGPU>(device)
.verify(x);
TensorMatcher({bsz, channels, cache_t, h_size, w_size})
.with_dtype<T>()
.template with_device<kDLGPU>(device)
.verify(cache);
TensorMatcher({bsz, channels, out_t, out_h, out_w})
.with_dtype<T>()
.template with_device<kDLGPU>(device)
.verify(out);
const int64_t depth_left = pad_d_left - cache_t.unwrap();
RuntimeCheck(depth_left >= 0, "pad_d_left must be >= cache_t");
RuntimeCheck(pad_d_right == 0, "pad_d_right must be 0");
RuntimeCheck(pad_w_left == pad_w_right, "width padding must be symmetric");
RuntimeCheck(pad_h_top == pad_h_bottom, "height padding must be symmetric");
RuntimeCheck(out_t.unwrap() == t_size.unwrap() + cache_t.unwrap() + depth_left + pad_d_right, "out_t mismatch");
RuntimeCheck(out_h.unwrap() == h_size.unwrap() + pad_h_top + pad_h_bottom, "out_h mismatch");
RuntimeCheck(out_w.unwrap() == w_size.unwrap() + pad_w_left + pad_w_right, "out_w mismatch");
const int64_t total = bsz.unwrap() * channels.unwrap() * out_t.unwrap() * out_h.unwrap() * out_w.unwrap();
if (total == 0) {
return;
}
constexpr int kVec = 16 / sizeof(T);
RuntimeCheck(total % kVec == 0, "output element count must be divisible by vector width");
RuntimeCheck(reinterpret_cast<uintptr_t>(out.data_ptr()) % 16 == 0, "output pointer must be 16-byte aligned");
if constexpr (sizeof(T) == 2) {
launch_cat_pad_flat<uint16_t, kVec>(
x.data_ptr(),
cache.data_ptr(),
out.data_ptr(),
total,
channels.unwrap(),
t_size.unwrap(),
h_size.unwrap(),
w_size.unwrap(),
cache_t.unwrap(),
out_t.unwrap(),
out_h.unwrap(),
out_w.unwrap(),
depth_left,
pad_h_top,
pad_w_left,
device.unwrap());
} else {
launch_cat_pad_flat<uint32_t, kVec>(
x.data_ptr(),
cache.data_ptr(),
out.data_ptr(),
total,
channels.unwrap(),
t_size.unwrap(),
h_size.unwrap(),
w_size.unwrap(),
cache_t.unwrap(),
out_t.unwrap(),
out_h.unwrap(),
out_w.unwrap(),
depth_left,
pad_h_top,
pad_w_left,
device.unwrap());
}
}
};
} // namespace sglang_causal_conv3d_cat_pad
python/sglang/jit_kernel/csrc/diffusion/ltx2_qknorm_split_rope.cuh
+276
View File
@@ -0,0 +1,276 @@
// CUDA fast path for LTX2 Q/K RMSNorm + split RoPE.
//
// Developed with MIT HAN Lab Kernel Design Agents:
// https://github.com/mit-han-lab/kernel-design-agents
//
// This mirrors the LTX2 eager oracle: RMSNorm and split RoPE both run in
// fp32, rounding to bf16 only once at the final attention input.
#pragma once
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
#include <sgl_kernel/utils.h> // For RuntimeCheck
#include <sgl_kernel/utils.cuh> // For LaunchKernel and CUDA dtype aliases
#include <cstdint>
#include <cuda_bf16.h>
namespace sglang_ltx2_qknorm_split_rope {
namespace {
constexpr int kThreads = 128;
inline const char* data_ptr(const tvm::ffi::TensorView& t) {
return static_cast<const char*>(t.data_ptr()) + t.byte_offset();
}
inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) {
return static_cast<char*>(t.data_ptr()) + t.byte_offset();
}
SGL_DEVICE float compute_rstd(
const bf16_t* __restrict__ xrow,
int64_t hidden_size,
float eps,
int tid,
int lane,
int warp_id,
float* warp_sum,
float* s_rstd) {
float local = 0.f;
const int64_t n_vec = hidden_size >> 2;
for (int64_t i = tid; i < n_vec; i += kThreads) {
const int64_t base = i << 2;
const float v0 = __bfloat162float(xrow[base + 0]);
const float v1 = __bfloat162float(xrow[base + 1]);
const float v2 = __bfloat162float(xrow[base + 2]);
const float v3 = __bfloat162float(xrow[base + 3]);
local = fmaf(v0, v0, local);
local = fmaf(v1, v1, local);
local = fmaf(v2, v2, local);
local = fmaf(v3, v3, local);
}
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
local += __shfl_down_sync(0xffffffffu, local, offset);
}
if (lane == 0) {
warp_sum[warp_id] = local;
}
__syncthreads();
if (tid == 0) {
const float total = (warp_sum[0] + warp_sum[2]) + (warp_sum[1] + warp_sum[3]);
*s_rstd = rsqrtf(total / static_cast<float>(hidden_size) + eps);
}
__syncthreads();
return *s_rstd;
}
SGL_DEVICE float norm_value(float x, float weight, float rstd) {
return weight * (rstd * x);
}
SGL_DEVICE void rope_pair(float x0, float x1, float cos, float sin, float& y0, float& y1) {
const float p0 = x0 * cos;
const float p1 = x1 * cos;
y0 = fmaf(-sin, x1, p0);
y1 = fmaf(sin, x0, p1);
}
__global__ void ltx2_qknorm_split_rope_kernel(
const bf16_t* __restrict__ x,
const bf16_t* __restrict__ cos,
const bf16_t* __restrict__ sin,
const bf16_t* __restrict__ weight,
bf16_t* __restrict__ out,
float eps,
int64_t seq_len,
int64_t num_heads,
int64_t head_dim,
int64_t stride_cos_b,
int64_t stride_cos_h,
int64_t stride_cos_t,
int64_t stride_sin_b,
int64_t stride_sin_h,
int64_t stride_sin_t) {
const int64_t row = static_cast<int64_t>(blockIdx.x);
const int64_t batch = row / seq_len;
const int64_t token = row - batch * seq_len;
const int64_t hidden_size = num_heads * head_dim;
const int64_t half_dim = head_dim >> 1;
const auto* __restrict__ xrow = x + row * hidden_size;
auto* __restrict__ outrow = out + row * hidden_size;
const int tid = threadIdx.x + threadIdx.y * 32;
const int lane = threadIdx.x;
const int warp_id = threadIdx.y;
__shared__ float warp_sum[4];
__shared__ float s_rstd;
const float rstd = compute_rstd(xrow, hidden_size, eps, tid, lane, warp_id, warp_sum, &s_rstd);
const int64_t num_pairs = num_heads * half_dim;
for (int64_t pair = tid; pair < num_pairs; pair += kThreads) {
const int64_t head = pair / half_dim;
const int64_t offset = pair - head * half_dim;
const int64_t idx0 = head * head_dim + offset;
const int64_t idx1 = idx0 + half_dim;
const float n0 = norm_value(__bfloat162float(xrow[idx0]), __bfloat162float(weight[idx0]), rstd);
const float n1 = norm_value(__bfloat162float(xrow[idx1]), __bfloat162float(weight[idx1]), rstd);
const int64_t cos_offset = batch * stride_cos_b + head * stride_cos_h + token * stride_cos_t + offset;
const int64_t sin_offset = batch * stride_sin_b + head * stride_sin_h + token * stride_sin_t + offset;
float y0;
float y1;
rope_pair(n0, n1, __bfloat162float(cos[cos_offset]), __bfloat162float(sin[sin_offset]), y0, y1);
outrow[idx0] = __float2bfloat16_rn(y0);
outrow[idx1] = __float2bfloat16_rn(y1);
}
}
inline void launch_one(
const tvm::ffi::TensorView& x,
const tvm::ffi::TensorView& cos,
const tvm::ffi::TensorView& sin,
const tvm::ffi::TensorView& weight,
const tvm::ffi::TensorView& out,
float eps,
int64_t num_rows,
int64_t seq_len,
int64_t num_heads,
int64_t head_dim,
int64_t stride_cos_b,
int64_t stride_cos_h,
int64_t stride_cos_t,
int64_t stride_sin_b,
int64_t stride_sin_h,
int64_t stride_sin_t,
DLDevice device) {
if (num_rows == 0) {
return;
}
host::RuntimeCheck(num_rows <= static_cast<int64_t>(UINT32_MAX), "LTX2 QKNorm split-RoPE grid is too large");
host::LaunchKernel(dim3(static_cast<uint32_t>(num_rows)), dim3(32, 4), device)(
ltx2_qknorm_split_rope_kernel,
reinterpret_cast<const bf16_t*>(data_ptr(x)),
reinterpret_cast<const bf16_t*>(data_ptr(cos)),
reinterpret_cast<const bf16_t*>(data_ptr(sin)),
reinterpret_cast<const bf16_t*>(data_ptr(weight)),
reinterpret_cast<bf16_t*>(mutable_data_ptr(out)),
eps,
seq_len,
num_heads,
head_dim,
stride_cos_b,
stride_cos_h,
stride_cos_t,
stride_sin_b,
stride_sin_h,
stride_sin_t);
}
} // namespace
struct LTX2QKNormSplitRopeKernel {
static void
run(tvm::ffi::TensorView q_out,
tvm::ffi::TensorView k_out,
tvm::ffi::TensorView q,
tvm::ffi::TensorView q_cos,
tvm::ffi::TensorView q_sin,
tvm::ffi::TensorView q_weight,
tvm::ffi::TensorView k,
tvm::ffi::TensorView k_cos,
tvm::ffi::TensorView k_sin,
tvm::ffi::TensorView k_weight,
double eps,
int64_t num_heads,
int64_t head_dim) {
using namespace host;
RuntimeCheck(num_heads > 0, "num_heads must be positive");
RuntimeCheck(head_dim > 0, "head_dim must be positive");
RuntimeCheck(head_dim % 2 == 0, "head_dim must be even");
const int64_t hidden_size = num_heads * head_dim;
RuntimeCheck(hidden_size % 4 == 0, "hidden size must be divisible by 4");
auto batch = SymbolicSize{"batch"};
auto q_seq_len = SymbolicSize{"q_seq_len"};
auto k_seq_len = SymbolicSize{"k_seq_len"};
auto heads = SymbolicSize{"num_heads"};
auto half_dim = SymbolicSize{"half_dim"};
auto device = SymbolicDevice{};
heads.set_value(num_heads);
half_dim.set_value(head_dim / 2);
device.set_options<kDLCUDA>();
TensorMatcher({batch, q_seq_len, hidden_size}).with_dtype<bf16_t>().with_device(device).verify(q).verify(q_out);
TensorMatcher({batch, k_seq_len, hidden_size}).with_dtype<bf16_t>().with_device(device).verify(k).verify(k_out);
TensorMatcher({hidden_size}).with_dtype<bf16_t>().with_device(device).verify(q_weight);
TensorMatcher({hidden_size}).with_dtype<bf16_t>().with_device(device).verify(k_weight);
TensorMatcher({batch, heads, q_seq_len, half_dim})
.with_strides({-1, -1, -1, 1})
.with_dtype<bf16_t>()
.with_device(device)
.verify(q_cos);
TensorMatcher({batch, heads, q_seq_len, half_dim})
.with_strides({-1, -1, -1, 1})
.with_dtype<bf16_t>()
.with_device(device)
.verify(q_sin);
TensorMatcher({batch, heads, k_seq_len, half_dim})
.with_strides({-1, -1, -1, 1})
.with_dtype<bf16_t>()
.with_device(device)
.verify(k_cos);
TensorMatcher({batch, heads, k_seq_len, half_dim})
.with_strides({-1, -1, -1, 1})
.with_dtype<bf16_t>()
.with_device(device)
.verify(k_sin);
const int64_t batch_size = batch.unwrap();
const DLDevice dl_device = device.unwrap();
launch_one(
q,
q_cos,
q_sin,
q_weight,
q_out,
static_cast<float>(eps),
batch_size * q_seq_len.unwrap(),
q_seq_len.unwrap(),
num_heads,
head_dim,
q_cos.stride(0),
q_cos.stride(1),
q_cos.stride(2),
q_sin.stride(0),
q_sin.stride(1),
q_sin.stride(2),
dl_device);
launch_one(
k,
k_cos,
k_sin,
k_weight,
k_out,
static_cast<float>(eps),
batch_size * k_seq_len.unwrap(),
k_seq_len.unwrap(),
num_heads,
head_dim,
k_cos.stride(0),
k_cos.stride(1),
k_cos.stride(2),
k_sin.stride(0),
k_sin.stride(1),
k_sin.stride(2),
dl_device);
}
};
} // namespace sglang_ltx2_qknorm_split_rope
python/sglang/jit_kernel/csrc/diffusion/norm_scale_shift.cuh
+216
View File
@@ -0,0 +1,216 @@
// Minimal native-CUDA fast path for Qwen-Image diffusion norm-scale-shift.
//
// Supported shape family:
// - bf16 activations, B == 1, hidden dim == 3072
// - layer norm only, no affine weight/bias
// - scale/shift are bf16 row-broadcast tensors ([D], [1,D], or [1,1,D])
// - optional residual path uses a bf16 row-broadcast gate
//
// All other public-op inputs fall back to the existing CuTe-DSL implementation
// from the Python dispatcher.
//
// Developed with MIT HAN Lab Kernel Design Agents:
// https://github.com/mit-han-lab/kernel-design-agents
#pragma once
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
#include <sgl_kernel/math.cuh> // For device::math::rsqrt
#include <sgl_kernel/utils.cuh> // For SGL_DEVICE, bf16_t, LaunchKernel
#include <sgl_kernel/vec.cuh> // For AlignedVector
#include <sgl_kernel/warp.cuh> // For warp::reduce_sum
#include <cstdint>
namespace sglang_norm_scale_shift {
namespace {
constexpr int kHidden = 3072;
constexpr int kVecElems = 16; // 32B/thread for bf16 on Blackwell.
constexpr int kThreads = kHidden / kVecElems;
constexpr int kWarps = kThreads / device::kWarpThreads;
constexpr float kInvHidden = 1.0f / float(kHidden);
static_assert(kThreads == 192);
static_assert(kWarps == 6);
struct QwenImageNormParams {
void* y;
void* res_out;
const void* x;
const void* residual;
const void* gate;
const void* scale;
const void* shift;
float eps;
};
SGL_DEVICE float cta_reduce_sum(float v, int warp, int lane, float* scratch) {
v = device::warp::reduce_sum(v);
if (lane == 0) {
scratch[warp] = v;
}
__syncthreads();
if (warp == 0) {
float a = lane < kWarps ? scratch[lane] : 0.0f;
a = device::warp::reduce_sum(a);
if (lane == 0) {
scratch[kWarps] = a;
}
}
__syncthreads();
return scratch[kWarps];
}
template <bool kHasResidual>
__global__ void qwen_image_norm_scale_shift_kernel(const QwenImageNormParams __grid_constant__ params) {
using namespace device;
using Vec = AlignedVector<bf16_t, kVecElems>;
const int row = blockIdx.x;
const int tid = threadIdx.x;
const int lane = tid & int(kWarpThreads - 1);
const int warp = tid >> 5;
const int row_offset = row * kHidden;
const int elem_offset = tid * kVecElems;
__shared__ float scratch_a[kWarps + 1];
__shared__ float scratch_b[kWarps + 1];
Vec xv;
xv.load(static_cast<const bf16_t*>(params.x) + row_offset + elem_offset);
float v[kVecElems];
#pragma unroll
for (int i = 0; i < kVecElems; ++i) {
v[i] = static_cast<float>(xv[i]);
}
if constexpr (kHasResidual) {
Vec gv;
Vec rv;
Vec ro;
gv.load(static_cast<const bf16_t*>(params.gate) + elem_offset);
rv.load(static_cast<const bf16_t*>(params.residual) + row_offset + elem_offset);
#pragma unroll
for (int i = 0; i < kVecElems; ++i) {
const bf16_t rounded = static_cast<bf16_t>(v[i] * static_cast<float>(gv[i]) + static_cast<float>(rv[i]));
ro[i] = rounded;
v[i] = static_cast<float>(rounded);
}
ro.store(static_cast<bf16_t*>(params.res_out) + row_offset + elem_offset);
}
float sum = 0.0f;
#pragma unroll
for (int i = 0; i < kVecElems; ++i) {
sum += v[i];
}
const float mean = cta_reduce_sum(sum, warp, lane, scratch_a) * kInvHidden;
float var_sum = 0.0f;
#pragma unroll
for (int i = 0; i < kVecElems; ++i) {
const float d = v[i] - mean;
var_sum += d * d;
}
const float var = cta_reduce_sum(var_sum, warp, lane, scratch_b) * kInvHidden;
const float factor = math::rsqrt(var + params.eps);
Vec scv;
Vec shv;
Vec yv;
scv.load(static_cast<const bf16_t*>(params.scale) + elem_offset);
shv.load(static_cast<const bf16_t*>(params.shift) + elem_offset);
#pragma unroll
for (int i = 0; i < kVecElems; ++i) {
const float norm = static_cast<float>(static_cast<bf16_t>((v[i] - mean) * factor));
yv[i] = static_cast<bf16_t>(norm * (1.0f + static_cast<float>(scv[i])) + static_cast<float>(shv[i]));
}
yv.store(static_cast<bf16_t*>(params.y) + row_offset + elem_offset);
}
inline uint32_t verify_qwen_geometry(host::SymbolicSize& num_rows) {
using namespace host;
RuntimeCheck(num_rows.unwrap() > 0, "num_rows must be positive");
RuntimeCheck(num_rows.unwrap() <= int64_t(UINT32_MAX), "num_rows out of range");
return static_cast<uint32_t>(num_rows.unwrap());
}
} // namespace
struct QwenImageNormScaleShiftKernel {
static void
run(tvm::ffi::TensorView y,
tvm::ffi::TensorView x,
tvm::ffi::TensorView scale,
tvm::ffi::TensorView shift,
double eps) {
using namespace host;
auto N = SymbolicSize{"num_rows"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({N, kHidden}).with_dtype<bf16_t>().with_device(device).verify(x).verify(y);
TensorMatcher({kHidden}).with_dtype<bf16_t>().with_device(device).verify(scale).verify(shift);
const uint32_t grid = verify_qwen_geometry(N);
const auto params = QwenImageNormParams{
.y = y.data_ptr(),
.res_out = nullptr,
.x = x.data_ptr(),
.residual = nullptr,
.gate = nullptr,
.scale = scale.data_ptr(),
.shift = shift.data_ptr(),
.eps = static_cast<float>(eps),
};
LaunchKernel(grid, kThreads, device.unwrap())(qwen_image_norm_scale_shift_kernel<false>, params);
}
};
struct QwenImageScaleResidualNormScaleShiftKernel {
static void
run(tvm::ffi::TensorView y,
tvm::ffi::TensorView res_out,
tvm::ffi::TensorView residual,
tvm::ffi::TensorView x,
tvm::ffi::TensorView gate,
tvm::ffi::TensorView scale,
tvm::ffi::TensorView shift,
double eps) {
using namespace host;
auto N = SymbolicSize{"num_rows"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({N, kHidden})
.with_dtype<bf16_t>()
.with_device(device)
.verify(x)
.verify(residual)
.verify(y)
.verify(res_out);
TensorMatcher({kHidden}).with_dtype<bf16_t>().with_device(device).verify(gate).verify(scale).verify(shift);
const uint32_t grid = verify_qwen_geometry(N);
const auto params = QwenImageNormParams{
.y = y.data_ptr(),
.res_out = res_out.data_ptr(),
.x = x.data_ptr(),
.residual = residual.data_ptr(),
.gate = gate.data_ptr(),
.scale = scale.data_ptr(),
.shift = shift.data_ptr(),
.eps = static_cast<float>(eps),
};
LaunchKernel(grid, kThreads, device.unwrap())(qwen_image_norm_scale_shift_kernel<true>, params);
}
};
} // namespace sglang_norm_scale_shift
python/sglang/jit_kernel/csrc/diffusion/qknorm_rope.cuh
+246
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
python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh
+317
View File
@@ -0,0 +1,317 @@
// CUDA fast path for diffusion residual-gate elementwise updates.
//
// Implements:
// out = residual + update * gate
//
// The production shapes come from LTX-2.3 HQ residual/gate updates. This is
// intentionally narrow: contiguous residual/update/out tensors, with either a
// full contiguous gate or a row-broadcast [1, 1, D] gate.
//
// Developed with MIT HAN Lab Kernel Design Agents:
// https://github.com/mit-han-lab/kernel-design-agents
#pragma once
#include <sgl_kernel/tensor.h> // For host dtype helpers and TensorView metadata
#include <sgl_kernel/utils.h> // For RuntimeCheck and div_ceil
#include <sgl_kernel/type.cuh> // For dtype_trait conversions
#include <sgl_kernel/utils.cuh> // For LaunchKernel and CUDA dtype aliases
#include <sgl_kernel/vec.cuh> // For device::AlignedVector
#include <cstdint>
namespace sglang_residual_gate_add {
namespace {
constexpr int kBlockSize = 256;
constexpr int kBcastRowsPerBlock = 4;
constexpr int kBcastColsVecPerBlock = 256;
constexpr int64_t kMaxGrid = 65535;
enum class GateMode : int { kFull = 0, kBcastRow = 1 };
inline const char* data_ptr(const tvm::ffi::TensorView& t) {
return static_cast<const char*>(t.data_ptr()) + t.byte_offset();
}
inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) {
return static_cast<char*>(t.data_ptr()) + t.byte_offset();
}
inline bool aligned16(const void* p) {
return (reinterpret_cast<uintptr_t>(p) & 0xF) == 0;
}
inline int64_t numel(const tvm::ffi::TensorView& t) {
int64_t n = 1;
for (int i = 0; i < t.ndim(); ++i) {
n *= t.size(i);
}
return n;
}
inline int64_t grid_for(int64_t total) {
int64_t grid = host::div_ceil(total, static_cast<int64_t>(kBlockSize));
if (grid < 1) {
grid = 1;
}
if (grid > kMaxGrid) {
grid = kMaxGrid;
}
return grid;
}
inline bool is_dense_contiguous(const tvm::ffi::TensorView& t) {
int64_t expected = 1;
for (int i = t.ndim() - 1; i >= 0; --i) {
if (t.size(i) == 1) {
continue;
}
if (t.stride(i) != expected) {
return false;
}
expected *= t.size(i);
}
return true;
}
template <typename T>
inline void check_dtype(const tvm::ffi::TensorView& t) {
host::RuntimeCheck(host::is_type<T>(t.dtype()), "unexpected dtype for residual_gate_add");
}
template <typename T>
__device__ __forceinline__ float to_float(T v) {
return static_cast<float>(v);
}
template <>
__device__ __forceinline__ float to_float<fp16_t>(fp16_t v) {
return __half2float(v);
}
template <>
__device__ __forceinline__ float to_float<bf16_t>(bf16_t v) {
return __bfloat162float(v);
}
template <typename T>
__device__ __forceinline__ T residual_gate_value(T residual, T update, T gate) {
const T product = dtype_trait<T>::from(to_float(update) * to_float(gate));
return dtype_trait<T>::from(to_float(residual) + to_float(product));
}
template <typename T, int kVec>
__global__ void residual_gate_add_vec_kernel(
const T* __restrict__ residual,
const T* __restrict__ update,
const T* __restrict__ gate,
T* __restrict__ out,
int64_t n_vec) {
using Vec = device::AlignedVector<T, kVec>;
const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
for (int64_t v = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; v < n_vec; v += stride) {
Vec r, u, g, o;
r.load(residual, v);
u.load(update, v);
g.load(gate, v);
#pragma unroll
for (int i = 0; i < kVec; ++i) {
o[i] = residual_gate_value(r[i], u[i], g[i]);
}
o.store(out, v);
}
}
template <typename T, int kVec>
__global__ void residual_gate_add_bcast_row_tile_kernel(
const T* __restrict__ residual,
const T* __restrict__ update,
const T* __restrict__ gate,
T* __restrict__ out,
int64_t rows,
int64_t row_vec) {
using Vec = device::AlignedVector<T, kVec>;
const int64_t col_vec = static_cast<int64_t>(blockIdx.x) * kBcastColsVecPerBlock + threadIdx.x;
if (col_vec >= row_vec) {
return;
}
Vec g;
g.load(gate, col_vec);
// Grid-stride over row tiles so the launch stays valid even when the number
// of row tiles exceeds the gridDim.y hardware limit.
const int64_t row_tile_stride = static_cast<int64_t>(gridDim.y) * kBcastRowsPerBlock;
for (int64_t row_base = static_cast<int64_t>(blockIdx.y) * kBcastRowsPerBlock; row_base < rows;
row_base += row_tile_stride) {
#pragma unroll
for (int row_offset = 0; row_offset < kBcastRowsPerBlock; ++row_offset) {
const int64_t row = row_base + row_offset;
if (row < rows) {
const int64_t v = row * row_vec + col_vec;
Vec r, u, o;
r.load(residual, v);
u.load(update, v);
#pragma unroll
for (int i = 0; i < kVec; ++i) {
o[i] = residual_gate_value(r[i], u[i], g[i]);
}
o.store(out, v);
}
}
}
}
template <typename T, GateMode kGate>
__global__ void residual_gate_add_scalar_kernel(
const T* __restrict__ residual,
const T* __restrict__ update,
const T* __restrict__ gate,
T* __restrict__ out,
int64_t begin,
int64_t total,
int64_t D) {
const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
for (int64_t i = begin + static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < total; i += stride) {
const T gate_value = kGate == GateMode::kFull ? gate[i] : SGLANG_LDG(gate + (i % D));
out[i] = residual_gate_value(residual[i], update[i], gate_value);
}
}
template <typename T>
inline void launch_residual_gate_add(
const tvm::ffi::TensorView& out,
const tvm::ffi::TensorView& residual,
const tvm::ffi::TensorView& update,
const tvm::ffi::TensorView& gate,
GateMode mode) {
const int64_t total = numel(residual);
if (total == 0) {
return;
}
const int64_t D = residual.size(residual.ndim() - 1);
const T* residual_ptr = reinterpret_cast<const T*>(data_ptr(residual));
const T* update_ptr = reinterpret_cast<const T*>(data_ptr(update));
const T* gate_ptr = reinterpret_cast<const T*>(data_ptr(gate));
T* out_ptr = reinterpret_cast<T*>(mutable_data_ptr(out));
constexpr int kVec = 16 / sizeof(T);
const bool vec_ok = aligned16(residual_ptr) && aligned16(update_ptr) && aligned16(gate_ptr) && aligned16(out_ptr) &&
(D % kVec == 0) && (mode == GateMode::kBcastRow || total % kVec == 0);
int64_t done = 0;
if (vec_ok) {
const int64_t n_vec = total / kVec;
const int64_t row_vec = D / kVec;
if (mode == GateMode::kFull) {
host::LaunchKernel(static_cast<uint32_t>(grid_for(n_vec)), kBlockSize, out.device())(
residual_gate_add_vec_kernel<T, kVec>, residual_ptr, update_ptr, gate_ptr, out_ptr, n_vec);
} else {
const int64_t rows = total / D;
const int64_t col_blocks = host::div_ceil(row_vec, static_cast<int64_t>(kBcastColsVecPerBlock));
const int64_t row_tiles = host::div_ceil(rows, static_cast<int64_t>(kBcastRowsPerBlock));
const int64_t row_blocks = row_tiles > kMaxGrid ? kMaxGrid : row_tiles;
host::LaunchKernel(
dim3(static_cast<uint32_t>(col_blocks), static_cast<uint32_t>(row_blocks)),
dim3(kBcastColsVecPerBlock),
out.device())(
residual_gate_add_bcast_row_tile_kernel<T, kVec>, residual_ptr, update_ptr, gate_ptr, out_ptr, rows, row_vec);
}
done = n_vec * kVec;
}
if (done < total) {
if (mode == GateMode::kFull) {
host::LaunchKernel(static_cast<uint32_t>(grid_for(total - done)), kBlockSize, out.device())(
residual_gate_add_scalar_kernel<T, GateMode::kFull>,
residual_ptr,
update_ptr,
gate_ptr,
out_ptr,
done,
total,
D);
} else {
host::LaunchKernel(static_cast<uint32_t>(grid_for(total - done)), kBlockSize, out.device())(
residual_gate_add_scalar_kernel<T, GateMode::kBcastRow>,
residual_ptr,
update_ptr,
gate_ptr,
out_ptr,
done,
total,
D);
}
}
}
template <typename T>
inline GateMode validate_residual_gate_add(
const tvm::ffi::TensorView& out,
const tvm::ffi::TensorView& residual,
const tvm::ffi::TensorView& update,
const tvm::ffi::TensorView& gate) {
check_dtype<T>(out);
check_dtype<T>(residual);
check_dtype<T>(update);
check_dtype<T>(gate);
host::RuntimeCheck(residual.device().device_type == kDLCUDA, "residual must be CUDA");
host::RuntimeCheck(update.device().device_type == kDLCUDA, "update must be CUDA");
host::RuntimeCheck(gate.device().device_type == kDLCUDA, "gate must be CUDA");
host::RuntimeCheck(out.device().device_type == kDLCUDA, "out must be CUDA");
host::RuntimeCheck(
residual.device().device_id == update.device().device_id &&
residual.device().device_id == gate.device().device_id &&
residual.device().device_id == out.device().device_id,
"residual/update/gate/out must be on the same CUDA device");
host::RuntimeCheck(residual.ndim() >= 2, "residual must be at least 2D");
host::RuntimeCheck(update.ndim() == residual.ndim(), "update rank must match residual");
host::RuntimeCheck(out.ndim() == residual.ndim(), "out rank must match residual");
for (int i = 0; i < residual.ndim(); ++i) {
host::RuntimeCheck(update.size(i) == residual.size(i), "update shape must match residual");
host::RuntimeCheck(out.size(i) == residual.size(i), "out shape must match residual");
}
host::RuntimeCheck(is_dense_contiguous(residual), "residual must be contiguous");
host::RuntimeCheck(is_dense_contiguous(update), "update must be contiguous");
host::RuntimeCheck(is_dense_contiguous(out), "out must be contiguous");
host::RuntimeCheck(is_dense_contiguous(gate), "gate must be contiguous");
host::RuntimeCheck(data_ptr(out) != data_ptr(residual), "out must not alias residual");
host::RuntimeCheck(data_ptr(out) != data_ptr(update), "out must not alias update");
host::RuntimeCheck(data_ptr(out) != data_ptr(gate), "out must not alias gate");
const int D_dim = residual.ndim() - 1;
const int row_dim = residual.ndim() - 2;
host::RuntimeCheck(gate.ndim() == residual.ndim(), "gate rank must match residual");
host::RuntimeCheck(gate.size(D_dim) == residual.size(D_dim), "gate last dim must match residual");
bool full_gate = true;
for (int i = 0; i < residual.ndim(); ++i) {
full_gate = full_gate && gate.size(i) == residual.size(i);
}
if (full_gate) {
return GateMode::kFull;
}
host::RuntimeCheck(gate.size(row_dim) == 1, "broadcast gate row dim must be 1");
for (int i = 0; i < D_dim; ++i) {
host::RuntimeCheck(gate.size(i) == 1, "broadcast gate leading dims must be 1");
}
return GateMode::kBcastRow;
}
} // namespace
template <typename T>
struct ResidualGateAddKernel {
static void
run(tvm::ffi::TensorView out, tvm::ffi::TensorView residual, tvm::ffi::TensorView update, tvm::ffi::TensorView gate) {
const GateMode mode = validate_residual_gate_add<T>(out, residual, update, gate);
launch_residual_gate_add<T>(out, residual, update, gate, mode);
}
};
} // namespace sglang_residual_gate_add
python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh
+154
View File
@@ -0,0 +1,154 @@
#pragma once
#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 <sgl_kernel/vec.cuh> // For device::AlignedVector
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cuda_runtime.h>
#include <type_traits>
namespace sglang_timestep_embedding {
namespace {
constexpr int kVec = 4; // 16B float vector store
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) {
using Vec = device::AlignedVector<float, kVec>;
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 * kVec < half_dim) {
// !flip: output is [sin | cos]; flip: output is [cos | sin].
float* cos_dst;
float* sin_dst;
if constexpr (!kFlipSinToCos) {
sin_dst = output_batch_base_ptr + thread_offset * kVec;
cos_dst = output_batch_base_ptr + half_dim + thread_offset * kVec;
} else {
cos_dst = output_batch_base_ptr + thread_offset * kVec;
sin_dst = output_batch_base_ptr + half_dim + thread_offset * kVec;
}
Vec cos_vec;
Vec sin_vec;
#pragma unroll
for (int i = 0; i < kVec; ++i) {
const float angle =
scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * kVec + i));
cos_vec[i] = device::math::cos(angle);
sin_vec[i] = device::math::sin(angle);
}
cos_vec.store(cos_dst);
sin_vec.store(sin_dst);
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);
}
}
} // namespace
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 sglang_timestep_embedding
python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh
+30
View File
@@ -0,0 +1,30 @@
#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("get_graph_capture_ptrs", &Class::get_graph_capture_ptrs)
.def("get_graph_capture_bases", &Class::get_graph_capture_bases)
.def("register_peer_mapped_inputs", &Class::register_peer_mapped_inputs)
.def("free_ipc_handles", &Class::free_ipc_handles)
.def("free_storage", &Class::free_storage)
.def("configure_pull", &Class::configure_pull);
}
python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh
+205
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 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);
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, input_bytes);
/// 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
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
python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh
+253
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
python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh
+325
View File
@@ -0,0 +1,325 @@
// Adapted from https://github.com/NVIDIA/TensorRT-LLM/pull/12163
// We reuse the custom all reduce push buffer in SGLang
#include <sgl_kernel/ffi.h>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/math.cuh>
#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 <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 ParallelQKNormParams {
void* __restrict__ buffer[device::distributed::kMaxNumGPU];
void* q_ptr;
void* k_ptr;
const void* __restrict__ q_weight;
const void* __restrict__ k_weight;
int64_t q_stride_bytes;
int64_t k_stride_bytes;
float eps;
uint32_t rank;
uint32_t num_tokens;
uint32_t epoch_bytes;
uint32_t num_clean_up_count = 0;
};
template <typename T>
SGL_DEVICE void ld_global_volatile_8B(T& x, const void* addr, int64_t offset) {
static_assert(alignof(T) == 8 && sizeof(T) == 8);
addr = device::pointer::offset<T>(addr, offset);
uint2 val;
asm volatile("ld.volatile.global.v2.b32 {%0, %1}, [%2];" : "=r"(val.x), "=r"(val.y) : "l"(addr));
x = *reinterpret_cast<const T*>(&val);
}
template <typename T>
SGL_DEVICE void st_global_volatile_8B(const T& x, void* addr, int64_t offset) {
static_assert(alignof(T) == 8 && sizeof(T) == 8);
const uint2 val = *reinterpret_cast<const uint2*>(&x);
addr = device::pointer::offset<T>(addr, offset);
asm volatile("st.volatile.global.v2.b32 [%2], {%0, %1};" ::"r"(val.x), "r"(val.y), "l"(addr));
}
[[maybe_unused]]
SGL_DEVICE float sync_float(float x) {
return __shfl_sync(0xffffffffu, x, 0);
}
[[maybe_unused]]
constexpr auto next_pow_of_2(uint32_t x) {
uint32_t y = 1;
while (y < x)
y *= 2;
return y;
}
template <typename DType_, uint32_t kNumGPU_, int64_t kQDim_, int64_t kKDim_, bool kUsePDL_>
struct KernelTrait {
// rename the arguments to avoid confusion with the template parameters
using DType = DType_;
static constexpr uint32_t kNumGPU = kNumGPU_;
static constexpr int64_t kQDim = kQDim_;
static constexpr int64_t kKDim = kKDim_;
static constexpr bool kUsePDL = kUsePDL_;
static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
static constexpr int64_t kLocalQDim = kQDim / kNumGPU;
static constexpr int64_t kLocalKDim = kKDim / kNumGPU;
static constexpr uint32_t kNumQThreads = kLocalQDim / (kVecSize * 2);
static constexpr uint32_t kNumKThreads = kLocalKDim / (kVecSize * 2);
static constexpr uint32_t kNumQWarps = kNumQThreads / device::kWarpThreads;
static constexpr uint32_t kNumKWarps = host::div_ceil(kNumKThreads, device::kWarpThreads);
static constexpr uint32_t kBlockSize = (kNumQWarps + kNumKWarps) * device::kWarpThreads;
static constexpr uint32_t kOccupancy = 2048 / kBlockSize;
using DType2 = packed_t<DType>;
using Storage = device::AlignedVector<DType2, kVecSize>;
static_assert(std::has_single_bit(kNumGPU), "must be pow of 2");
static_assert(kQDim % kNumGPU == 0);
static_assert(kKDim % kNumGPU == 0);
static_assert(kLocalQDim % (kVecSize * 2) == 0);
static_assert(kLocalKDim % (kVecSize * 2) == 0);
static_assert(kNumQThreads % device::kWarpThreads == 0);
static_assert(kBlockSize <= 1024);
static_assert(sizeof(Storage) == 16 && alignof(Storage) == 16);
static_assert(kOccupancy * kBlockSize <= 2048);
};
template <typename Trait>
__global__ __launch_bounds__(Trait::kBlockSize, Trait::kOccupancy) void parallel_qknorm_across_head(
const ParallelQKNormParams __grid_constant__ params, const PushController __grid_constant__ ctrl) {
using namespace device;
// each cta will handle exactly 1 token
using Storage = typename Trait::Storage;
using DType2 = typename Trait::DType2;
const auto &[
buffer, q_ptr, k_ptr, q_weight, k_weight, q_stride_bytes, k_stride_bytes, //
eps, rank, num_tokens, epoch_bytes, num_clean_up_count
] = params;
using Package = AlignedVector<float, 2>;
constexpr uint32_t kNumGPU = Trait::kNumGPU;
constexpr uint32_t kNumQReduce = next_pow_of_2(Trait::kNumQWarps);
constexpr uint32_t kNumKReduce = next_pow_of_2(Trait::kNumKWarps);
__shared__ float smem_qk[Trait::kNumQWarps + Trait::kNumKWarps];
__shared__ float scale_q;
__shared__ float scale_k;
const auto tx = threadIdx.x;
const auto bx = blockIdx.x;
/// NOTE: this can hint compiler to optimize `is_valid` out when not needed
constexpr uint32_t kActiveThreads = Trait::kNumQThreads + Trait::kNumKThreads;
const auto is_valid = Trait::kBlockSize == kActiveThreads || tx < kActiveThreads;
const auto smem_q = smem_qk + 0;
const auto smem_k = smem_qk + Trait::kNumQWarps;
const auto load_q = tx < Trait::kNumQThreads;
const auto offset = load_q ? tx : tx - Trait::kNumQThreads;
const auto input_ptr = load_q ? q_ptr : k_ptr;
const auto weight_ptr = load_q ? q_weight : k_weight;
const auto input_stride_bytes = load_q ? q_stride_bytes : k_stride_bytes;
PDLWaitPrimary<Trait::kUsePDL>();
PDLTriggerSecondary<Trait::kUsePDL>();
if (bx >= num_tokens) {
[[unlikely]];
// In this case, we use the last few blocks to clean up other controllers
const auto start = (bx - num_tokens) * blockDim.x + threadIdx.x;
const auto stride = (gridDim.x - num_tokens) * blockDim.x;
for (uint32_t i = start; i < num_clean_up_count; i += stride)
ctrl.exit_unsafe(num_tokens + i);
return;
}
const auto epoch_offset = ctrl.epoch() * epoch_bytes; // only for comm
__builtin_assume(bx < num_tokens); // since we have `bx >= num_tokens`
Storage next_input;
void* input_i_ptr = pointer::offset(input_ptr, bx * input_stride_bytes);
if (is_valid) next_input.load(input_i_ptr, offset);
for (uint32_t i = bx; i < num_tokens; i += gridDim.x) {
// Stage 1. local reduce (warp-level)
Storage local_input;
{
float local_sum = 0.0;
if (is_valid) {
local_input = next_input;
#pragma unroll
for (uint32_t j = 0; j < Trait::kVecSize; ++j) {
const auto [x, y] = cast<fp32x2_t>(local_input[j]);
local_sum += x * x + y * y;
}
}
smem_qk[threadIdx.x / kWarpThreads] = warp::reduce_sum(local_sum);
}
// Stage 2. block reduce + push to peer ranks + poll from local rank
__syncthreads();
Storage local_weight;
const auto input_next_ptr = pointer::offset(input_i_ptr, gridDim.x * input_stride_bytes);
/**
* NOTE: Prefetch to hide the latency.
* This brings around 20% of performance gain in large batches
* The P2P communication is mainly latency bound, so during this waiting period,
* We can let some data loading transparently in the background.
*/
if (is_valid) {
local_weight.load(weight_ptr, offset);
if (i + gridDim.x < num_tokens) next_input.load(input_next_ptr, offset);
}
if (tx < kWarpThreads) {
const auto local_sum_q = tx < Trait::kNumQWarps ? smem_q[tx] : 0.0f;
const auto local_sum_k = tx < Trait::kNumKWarps ? smem_k[tx] : 0.0f;
const auto sum_q = sync_float(warp::reduce_sum<kNumQReduce>(local_sum_q));
const auto sum_k = sync_float(warp::reduce_sum<kNumKReduce>(local_sum_k));
if (tx < kNumGPU) { // push a float2 pack to the peer
Package sum_q_k;
/// NOTE: eps should be scaled down by kNumGPU from host side
/// we add here to ensure that the sum is never zero
sum_q_k[0] = sum_q + eps;
sum_q_k[1] = sum_k + eps;
const auto push_ptr = pointer::offset(buffer[tx], epoch_offset);
st_global_volatile_8B(sum_q_k, push_ptr, i * kNumGPU + rank);
const auto poll_ptr = pointer::offset(buffer[rank], epoch_offset);
while (true) {
ld_global_volatile_8B(sum_q_k, poll_ptr, i * kNumGPU + tx);
if (sum_q_k[0] != 0.0f && sum_q_k[1] != 0.0f) break;
}
constexpr uint32_t kActiveMask = (1 << kNumGPU) - 1;
const auto global_sum_q = warp::reduce_sum<kNumGPU>(sum_q_k[0], kActiveMask);
const auto global_sum_k = warp::reduce_sum<kNumGPU>(sum_q_k[1], kActiveMask);
scale_q = math::rsqrt(global_sum_q / static_cast<float>(Trait::kQDim));
scale_k = math::rsqrt(global_sum_k / static_cast<float>(Trait::kKDim));
Package zeros;
zeros.fill(0.0f);
zeros.store(poll_ptr, i * kNumGPU + tx);
}
}
__syncthreads();
const auto scale = load_q ? scale_q : scale_k;
if (is_valid) {
#pragma unroll
for (uint32_t j = 0; j < Trait::kVecSize; ++j) {
const auto fp32_input = cast<fp32x2_t>(local_input[j]);
const auto fp32_weight = cast<fp32x2_t>(local_weight[j]);
const auto scaled_x = fp32_input.x * scale * fp32_weight.x;
const auto scaled_y = fp32_input.y * scale * fp32_weight.y;
local_input[j] = cast<DType2>(fp32x2_t{scaled_x, scaled_y});
}
local_input.store(input_i_ptr, offset);
}
input_i_ptr = input_next_ptr;
}
ctrl.exit();
}
template <typename DType, uint32_t kNumGPU, int64_t kQDim, int64_t kKDim, bool kUsePDL>
struct FusedParallelQKNormAcrossHead : public CustomAllReduceBase {
using Trait = KernelTrait<DType, kNumGPU, kQDim, kKDim, kUsePDL>;
static constexpr auto kernel = parallel_qknorm_across_head<Trait>;
static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs");
void _run(
const tvm::ffi::Tensor q,
const tvm::ffi::Tensor k,
const tvm::ffi::Tensor q_weight,
const tvm::ffi::Tensor k_weight,
const float eps // passed in unscaled
) {
using namespace host;
constexpr auto Q = Trait::kLocalQDim;
constexpr auto K = Trait::kLocalKDim;
auto N = SymbolicSize{"num_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({N, Q}) // q
.with_strides({-1, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(q);
TensorMatcher({N, K}) // k
.with_strides({-1, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(k);
TensorMatcher({Q}) // q_weight
.with_dtype<DType>()
.with_device(device_)
.verify(q_weight);
TensorMatcher({K}) // k_weight
.with_dtype<DType>()
.with_device(device_)
.verify(k_weight);
const auto device = device_.unwrap();
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
// use at most `world_size` blocks to clean up,
// this is based on the observation that occupancy is usually linear
// with respect to the world size
const bool need_clean = num_tokens < m_max_num_cta_push;
const auto num_clean = need_clean ? (m_max_num_cta_push - num_tokens) : 0;
const auto num_blocks = need_clean ? num_tokens + div_ceil(num_clean, Trait::kBlockSize) //
: m_max_num_cta_push; //
const auto num_threads = Trait::kBlockSize;
RuntimeCheck(num_blocks <= m_max_num_cta_push, "internal error");
ParallelQKNormParams params;
for (uint32_t i = 0; i < kNumGPU; ++i) {
params.buffer[i] = get_push_buffer(m_peer_storage[i]);
}
params.q_ptr = q.data_ptr();
params.k_ptr = k.data_ptr();
params.q_weight = q_weight.data_ptr();
params.k_weight = k_weight.data_ptr();
params.q_stride_bytes = q.stride(0) * sizeof(DType);
params.k_stride_bytes = k.stride(0) * sizeof(DType);
params.eps = eps / kNumGPU; // scale down eps by number of GPUs
params.rank = m_rank;
params.num_tokens = num_tokens;
params.epoch_bytes = m_push_buffer_bytes;
params.num_clean_up_count = num_clean;
const auto needed_buffer_bytes = static_cast<int64_t>(num_tokens) * 2 * sizeof(float);
RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch");
RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized");
RuntimeCheck(std::bit_cast<intptr_t>(params.q_ptr) % 16 == 0, "q pointer is not properly aligned");
RuntimeCheck(std::bit_cast<intptr_t>(params.k_ptr) % 16 == 0, "k pointer is not properly aligned");
RuntimeCheck(std::bit_cast<intptr_t>(params.q_weight) % 16 == 0, "q_weight pointer is not properly aligned");
RuntimeCheck(std::bit_cast<intptr_t>(params.k_weight) % 16 == 0, "k_weight pointer is not properly aligned");
RuntimeCheck(needed_buffer_bytes <= m_push_buffer_bytes, "Push buffer is too small");
LaunchKernel(num_blocks, num_threads, device) //
.enable_pdl(kUsePDL)(kernel, params, *m_push_ctrl);
}
static uint32_t get_max_occupancy() {
return host::runtime::get_blocks_per_sm(kernel, Trait::kBlockSize);
}
static void
run(CustomAllReduceRef obj,
const tvm::ffi::Tensor q,
const tvm::ffi::Tensor k,
const tvm::ffi::Tensor q_weight,
const tvm::ffi::Tensor k_weight,
const float eps) {
using Self = FusedParallelQKNormAcrossHead;
return static_cast<Self*>(obj.get())->_run(q, k, q_weight, k_weight, eps);
}
};
} // namespace
python/sglang/jit_kernel/csrc/dsa/fused_store_index_cache.cuh
+124
View File
@@ -0,0 +1,124 @@
#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 <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <bit>
#include <cstdint>
#include <cuda_fp8.h>
namespace {
struct FusedStoreCacheParam {
const void* __restrict__ input;
void* __restrict__ cache;
const void* __restrict__ indices;
uint32_t num_tokens;
};
[[maybe_unused]]
SGL_DEVICE float fp8_e4m3_clip(float val) {
namespace math = device::math;
return math::max(math::min(val, kFP8E4M3Max), -kFP8E4M3Max);
}
[[maybe_unused]]
SGL_DEVICE fp8x2_e4m3_t pack_fp8(float x, float y) {
return fp8x2_e4m3_t{fp32x2_t{fp8_e4m3_clip(x), fp8_e4m3_clip(y)}};
}
template <typename KeyT, typename IndicesT, uint32_t kPageBits, bool kUsePDL>
__global__ void fused_store_indexer_cache(const __grid_constant__ FusedStoreCacheParam param) {
using namespace device;
/// NOTE: 132 = 128 + 4
constexpr int64_t kPageBytes = 132 << kPageBits;
// each warp handles 128 elements, each block handles multiple rows
const auto& [input, cache, indices, num_tokens] = param;
const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x;
const auto global_wid = global_tid / 32;
const auto lane_id = threadIdx.x % 32;
if (global_wid >= num_tokens) return;
PDLWaitPrimary<kUsePDL>(); // wait for primary kernel
// prefetch the index
const auto index = static_cast<const IndicesT*>(indices)[global_wid];
// always load the value from input (don't store if invalid)
using KeyT2 = packed_t<KeyT>;
using InStorage = AlignedVector<KeyT2, 2>;
using OutStorage = AlignedVector<fp8x2_e4m3_t, 2>;
const auto elems = static_cast<const InStorage*>(input)[global_tid];
const auto [x0, x1] = cast<fp32x2_t>(elems[0]);
const auto [y0, y1] = cast<fp32x2_t>(elems[1]);
const auto local_max = fmaxf(fmaxf(fabs(x0), fabs(x1)), fmaxf(fabs(y0), fabs(y1)));
const auto abs_max = warp::reduce_max(local_max);
// use normal fp32 scale
const auto scale = fmaxf(1e-4f, abs_max) / kFP8E4M3Max;
const auto inv_scale = 1.0f / scale;
const int32_t page = index >> kPageBits;
const int32_t offset = index & ((1 << kPageBits) - 1);
const auto page_ptr = pointer::offset(cache, page * kPageBytes);
const auto value_ptr = pointer::offset(page_ptr, offset * 128);
const auto scale_ptr = pointer::offset(page_ptr, 128 << kPageBits, offset * 4);
OutStorage result;
result[0] = pack_fp8(x0 * inv_scale, x1 * inv_scale);
result[1] = pack_fp8(y0 * inv_scale, y1 * inv_scale);
static_cast<OutStorage*>(value_ptr)[lane_id] = result;
static_cast<float*>(scale_ptr)[0] = scale;
PDLTriggerSecondary<kUsePDL>(); // launch secondary kernel
}
template <typename KeyT, typename IndicesT, uint32_t kPageSize, bool kUsePDL>
struct FusedStoreCacheIndexerKernel {
static constexpr int32_t kLogSize = std::countr_zero(kPageSize);
/// NOTE: 132 = 128 + 4 (128 represent K and 4 represent scale)
static constexpr int64_t kPageBytes = 132 * kPageSize;
static constexpr auto kernel = fused_store_indexer_cache<KeyT, IndicesT, kLogSize, kUsePDL>;
static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
static_assert(1 << kLogSize == kPageSize);
static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView cache, tvm::ffi::TensorView indices) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({N, 128}) // input
.with_dtype<KeyT>()
.with_device(device_)
.verify(input);
TensorMatcher({-1, -1}) // cache
.with_strides({kPageBytes, 1})
.with_dtype<uint8_t>()
.with_device(device_)
.verify(cache);
TensorMatcher({N}) // indices
.with_dtype<IndicesT>()
.with_device(device_)
.verify(indices);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto params = FusedStoreCacheParam{
.input = input.data_ptr(),
.cache = cache.data_ptr(),
.indices = indices.data_ptr(),
.num_tokens = num_tokens,
};
const auto kBlockSize = 128;
const auto num_blocks = div_ceil(num_tokens * 32, kBlockSize);
LaunchKernel(num_blocks, kBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/dsa/kpool_topk_transform.cuh
+440
View File
@@ -0,0 +1,440 @@
/**
* @NOTE: The radix top-k core (fast_topk_cuda_tl_impl) is adapted from
* https://github.com/tile-ai/tilelang/blob/main/examples/deepseek_v32/topk_selector.py
* and was previously shipped as an AOT sgl-kernel op (fast_kpool_topk_transform_fused).
* It is re-implemented here as a lightweight JIT kernel for the NSA kpool indexer:
* select pool groups at pool granularity, expand each group to `pool_size` token
* indices, and optionally transform those indices through a page table or ragged offset.
*
* The pool-level top-k value is a compile-time constant injected via -DSGL_GROUP_TOPK.
*/
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice, is_type
#include <sgl_kernel/utils.h> // For RuntimeCheck, RuntimeDeviceCheck
#include <sgl_kernel/utils.cuh> // For LaunchKernel, type aliases
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <bit>
#include <cstddef>
#include <cstdint>
#include <cuda_fp16.h>
namespace {
#ifndef C10_LIKELY
#define C10_LIKELY(expr) (__builtin_expect(static_cast<bool>(expr), 1))
#endif
#ifndef SGL_GROUP_TOPK
#define SGL_GROUP_TOPK 256
#endif
// Compile-time pool-level top-k (number of groups selected per row).
inline constexpr int kGroupTopK = SGL_GROUP_TOPK;
inline constexpr int kThreadsPerBlock = 1024;
// Reduced from 128KB to 32KB to improve occupancy.
// Each radix pass needs at most ~K candidates in the threshold bin,
// so 4K entries per round (2 rounds = 8K entries = 32KB) is sufficient.
inline constexpr std::size_t kSmem = 8 * 1024 * sizeof(uint32_t); // 32KB (bytes)
struct FastTopKParams {
const float* __restrict__ input; // [B, input_stride]
const int32_t* __restrict__ row_starts; // [B] or nullptr
int32_t* __restrict__ indices; // unused here (kept for layout parity)
const int32_t* __restrict__ lengths; // [B]
int64_t input_stride;
};
__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t {
__half h = __float2half_rn(x);
uint16_t bits = __half_as_ushort(h);
uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits) : static_cast<uint16_t>(bits | 0x8000);
return static_cast<uint8_t>(key >> 8);
}
__device__ __forceinline__ auto convert_to_uint32(float x) -> uint32_t {
uint32_t bits = __float_as_uint(x);
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
}
template <int K>
__device__ void
fast_topk_cuda_tl_impl(const float* __restrict__ input, int* __restrict__ index, int row_start, int length) {
// An optimized topk kernel copied from tilelang kernel
// We assume length > K here, or it will crash
int topk = K;
constexpr auto BLOCK_SIZE = 1024;
constexpr auto RADIX = 256;
constexpr auto SMEM_INPUT_SIZE = kSmem / (2 * sizeof(int));
alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
alignas(128) __shared__ int s_counter;
alignas(128) __shared__ int s_threshold_bin_id;
alignas(128) __shared__ int s_num_input[2];
auto& s_histogram = s_histogram_buf[0];
// allocate for two rounds
extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE];
const int tx = threadIdx.x;
// stage 1: 8bit coarse histogram
if (tx < RADIX + 1) s_histogram[tx] = 0;
__syncthreads();
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
const auto bin = convert_to_uint8(input[idx + row_start]);
::atomicAdd(&s_histogram[bin], 1);
}
__syncthreads();
const auto run_cumsum = [&] {
#pragma unroll 8
for (int i = 0; i < 8; ++i) {
static_assert(1 << 8 == RADIX);
if (C10_LIKELY(tx < RADIX)) {
const auto j = 1 << i;
const auto k = i & 1;
auto value = s_histogram_buf[k][tx];
if (tx < RADIX - j) {
value += s_histogram_buf[k][tx + j];
}
s_histogram_buf[k ^ 1][tx] = value;
}
__syncthreads();
}
};
run_cumsum();
if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
s_threshold_bin_id = tx;
s_num_input[0] = 0;
s_counter = 0;
}
__syncthreads();
const auto threshold_bin = s_threshold_bin_id;
topk -= s_histogram[threshold_bin + 1];
if (topk == 0) {
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
const auto bin = static_cast<int>(convert_to_uint8(input[idx + row_start]));
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
index[pos] = idx;
}
}
__syncthreads();
return;
} else {
__syncthreads();
if (tx < RADIX + 1) {
s_histogram[tx] = 0;
}
__syncthreads();
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
const auto raw_input = input[idx + row_start];
const auto bin = static_cast<int>(convert_to_uint8(raw_input));
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
index[pos] = idx;
} else if (bin == threshold_bin) {
const auto pos = ::atomicAdd(&s_num_input[0], 1);
/// NOTE: (dark) fuse the histogram computation here
if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) {
s_input_idx[0][pos] = idx;
const auto bin = convert_to_uint32(raw_input);
const auto sub_bin = (bin >> 24) & 0xFF;
::atomicAdd(&s_histogram[sub_bin], 1);
}
}
}
__syncthreads();
}
// stage 2: refine with 8bit radix passes
#pragma unroll 4
for (int round = 0; round < 4; ++round) {
__shared__ int s_last_remain;
const auto r_idx = round % 2;
// clip here to prevent overflow
const auto _raw_num_input = s_num_input[r_idx];
const auto num_input = (_raw_num_input < int(SMEM_INPUT_SIZE)) ? _raw_num_input : int(SMEM_INPUT_SIZE);
run_cumsum();
if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
s_threshold_bin_id = tx;
s_num_input[r_idx ^ 1] = 0;
s_last_remain = topk - s_histogram[tx + 1];
}
__syncthreads();
const auto threshold_bin = s_threshold_bin_id;
topk -= s_histogram[threshold_bin + 1];
if (topk == 0) {
for (int i = tx; i < num_input; i += BLOCK_SIZE) {
const auto idx = s_input_idx[r_idx][i];
const auto offset = 24 - round * 8;
const auto bin = (convert_to_uint32(input[idx + row_start]) >> offset) & 0xFF;
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
index[pos] = idx;
}
}
__syncthreads();
break;
} else {
__syncthreads();
if (tx < RADIX + 1) {
s_histogram[tx] = 0;
}
__syncthreads();
for (int i = tx; i < num_input; i += BLOCK_SIZE) {
const auto idx = s_input_idx[r_idx][i];
const auto raw_input = input[idx + row_start];
const auto offset = 24 - round * 8;
const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF;
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
index[pos] = idx;
} else if (bin == threshold_bin) {
if (round == 3) {
const auto pos = ::atomicAdd(&s_last_remain, -1);
if (pos > 0) {
index[K - pos] = idx;
}
} else {
const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1);
if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) {
/// NOTE: (dark) fuse the histogram computation here
s_input_idx[r_idx ^ 1][pos] = idx;
const auto bin = convert_to_uint32(raw_input);
const auto sub_bin = (bin >> (offset - 8)) & 0xFF;
::atomicAdd(&s_histogram[sub_bin], 1);
}
}
}
}
__syncthreads();
}
}
}
__device__ __forceinline__ int32_t transform_kpool_token(
int32_t raw_token,
const int32_t* __restrict__ page_table_entry,
const int32_t* __restrict__ topk_indices_offset,
int32_t offset) {
if (page_table_entry != nullptr) {
return page_table_entry[raw_token];
}
if (topk_indices_offset != nullptr) {
return raw_token + offset;
}
return raw_token;
}
template <int K>
__global__ __launch_bounds__(kThreadsPerBlock) void kpool_topk_transform_kernel(
const __grid_constant__ FastTopKParams params,
int32_t* __restrict__ dst_token_indices,
const int64_t dst_stride,
const int32_t pool_size,
const int32_t token_topk,
const int32_t out_cols,
const int32_t* __restrict__ page_table,
const int64_t page_table_stride,
const int32_t* __restrict__ topk_indices_offset,
const int32_t* __restrict__ seq_lens) {
const auto& [input, row_starts, _, lengths, input_stride] = params;
const auto bid = static_cast<uint64_t>(blockIdx.x);
const auto tid = threadIdx.x;
const auto row_start = row_starts == nullptr ? 0 : row_starts[bid];
const auto length = lengths[bid];
const auto score = input + bid * input_stride;
const auto dst = dst_token_indices + bid * dst_stride;
const auto page_table_entry = page_table == nullptr ? nullptr : page_table + bid * page_table_stride;
const auto offset = topk_indices_offset == nullptr ? 0 : topk_indices_offset[bid];
const bool append_tail = seq_lens != nullptr;
const auto full_pool_token_len = length * pool_size;
const auto history_len = full_pool_token_len < token_topk ? full_pool_token_len : token_topk;
const auto tail_count = append_tail ? seq_lens[bid] % pool_size : 0;
if (length <= K) {
for (int col = tid; col < out_cols; col += kThreadsPerBlock) {
if (col < history_len) {
const auto group_rank = col / pool_size;
const auto slot = col % pool_size;
const auto raw_token = group_rank * pool_size + slot;
dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
} else if (append_tail && col < history_len + tail_count) {
const auto raw_token = length * pool_size + (col - history_len);
dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
} else {
dst[col] = -1;
}
}
return;
}
__shared__ int s_indices[K];
fast_topk_cuda_tl_impl<K>(score, s_indices, row_start, length);
for (int col = tid; col < out_cols; col += kThreadsPerBlock) {
if (col < history_len) {
const auto group_rank = col / pool_size;
const auto group_id = s_indices[group_rank];
const auto slot = col % pool_size;
const auto raw_token = group_id * pool_size + slot;
dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
} else if (append_tail && col < history_len + tail_count) {
const auto raw_token = length * pool_size + (col - history_len);
dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
} else {
dst[col] = -1;
}
}
}
template <auto* f, std::size_t kMaxDynamicSMEM>
void setup_kernel_smem_once(host::DebugInfo where = {}) {
[[maybe_unused]]
static const auto result = [] {
const auto fptr = std::bit_cast<const void*>(f);
return ::cudaFuncSetAttribute(fptr, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM);
}();
host::RuntimeDeviceCheck(result, where);
}
template <typename T>
const T* optional_data_ptr(const tvm::ffi::Optional<tvm::ffi::TensorView>& opt) {
if (!opt.has_value()) {
return nullptr;
}
return static_cast<const T*>(opt.value().data_ptr());
}
struct KpoolTopKTransformKernel {
static constexpr auto kernel = kpool_topk_transform_kernel<kGroupTopK>;
// Pool-level radix top-k for the NSA kpool indexer.
// score : [B, S] strided float32 scores (one score per pool group)
// lengths : [B] int32 valid group count per row
// dst_token_indices : [B, out_cols] int32 output token indices (contiguous)
// pool_size : tokens per pool group
// page_table (opt) : [B, P] strided int32 raw-token -> real-token map
// topk_indices_offset : [B] int32 per-row offset added to raw tokens (ragged)
// row_starts (opt) : [B] int32 score row start offsets
// seq_lens (opt) : [B] int32 sequence lengths; enables tail append
static void transform(
const tvm::ffi::TensorView score,
const tvm::ffi::TensorView lengths,
const tvm::ffi::TensorView dst_token_indices,
const int64_t pool_size,
const tvm::ffi::Optional<tvm::ffi::TensorView> page_table_opt,
const tvm::ffi::Optional<tvm::ffi::TensorView> topk_indices_offset_opt,
const tvm::ffi::Optional<tvm::ffi::TensorView> row_starts_opt,
const tvm::ffi::Optional<tvm::ffi::TensorView> seq_lens_opt) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto S = SymbolicSize{"score_stride"};
auto out_cols_sym = SymbolicSize{"out_cols"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({B, -1}) // strided scores
.with_strides({S, 1})
.with_dtype<float>()
.with_device(device)
.verify(score);
TensorMatcher({B}) // lengths, contiguous int32
.with_dtype<int32_t>()
.with_device(device)
.verify(lengths);
TensorMatcher({B, out_cols_sym}) // output, contiguous int32
.with_dtype<int32_t>()
.with_device(device)
.verify(dst_token_indices);
RuntimeCheck(pool_size > 1, "pool_size must be > 1, got ", pool_size);
RuntimeCheck(
!(page_table_opt.has_value() && topk_indices_offset_opt.has_value()),
"page_table and topk_indices_offset are mutually exclusive");
const auto out_cols = static_cast<int32_t>(out_cols_sym.unwrap());
const auto tail_cols = seq_lens_opt.has_value() ? static_cast<int32_t>(pool_size) - 1 : 0;
RuntimeCheck(out_cols > tail_cols, "dst_token_indices columns ", out_cols, " must exceed tail ", tail_cols);
const auto token_topk = out_cols - tail_cols;
RuntimeCheck(token_topk % static_cast<int32_t>(pool_size) == 0, "token_topk must be a multiple of pool_size");
RuntimeCheck(
token_topk / static_cast<int32_t>(pool_size) == kGroupTopK,
"this module is built for group_topk=",
kGroupTopK,
" but got ",
token_topk / static_cast<int32_t>(pool_size));
const auto batch_size = static_cast<uint32_t>(B.unwrap());
int64_t page_table_stride = 0;
const int32_t* page_table_ptr = nullptr;
if (page_table_opt.has_value()) {
auto P = SymbolicSize{"page_table_stride"};
TensorMatcher({B, -1}) // strided page table
.with_strides({P, 1})
.with_dtype<int32_t>()
.with_device(device)
.verify(page_table_opt.value());
page_table_ptr = static_cast<const int32_t*>(page_table_opt.value().data_ptr());
page_table_stride = static_cast<int64_t>(P.unwrap());
}
if (topk_indices_offset_opt.has_value()) {
TensorMatcher({B}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(topk_indices_offset_opt.value());
}
if (row_starts_opt.has_value()) {
TensorMatcher({B}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(row_starts_opt.value());
}
if (seq_lens_opt.has_value()) {
TensorMatcher({B}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(seq_lens_opt.value());
}
const auto params = FastTopKParams{
.input = static_cast<const float*>(score.data_ptr()),
.row_starts = optional_data_ptr<int32_t>(row_starts_opt),
.indices = nullptr,
.lengths = static_cast<const int32_t*>(lengths.data_ptr()),
.input_stride = static_cast<int64_t>(S.unwrap()),
};
setup_kernel_smem_once<kernel, kSmem>();
LaunchKernel(batch_size, kThreadsPerBlock, device.unwrap(), kSmem)(
kernel,
params,
static_cast<int32_t*>(dst_token_indices.data_ptr()),
static_cast<int64_t>(dst_token_indices.strides()[0]),
static_cast<int32_t>(pool_size),
token_topk,
out_cols,
page_table_ptr,
page_table_stride,
optional_data_ptr<int32_t>(topk_indices_offset_opt),
optional_data_ptr<int32_t>(seq_lens_opt));
}
};
} // namespace

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