diff --git a/.codespellrc b/.codespellrc
index 272a6ea66551..20b5855a22bc 100644
--- a/.codespellrc
+++ b/.codespellrc
@@ -1,3 +1,3 @@
[codespell]
-ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, ather, MIS, medias, allready, inout, nd
+ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, ather, MIS, medias, allready, inout, nd, dOut
skip = *.json, *.jsonl, *.patch, *.txt, *.lock
diff --git a/.github/workflows/release-pypi-nightly.yml b/.github/workflows/release-pypi-nightly.yml
index 3ce927df4d13..ff8d031aa9f7 100644
--- a/.github/workflows/release-pypi-nightly.yml
+++ b/.github/workflows/release-pypi-nightly.yml
@@ -48,6 +48,11 @@ jobs:
run: |
pip install build wheel setuptools setuptools-scm
+ # Needed by setuptools-rust to build the bundled native gRPC extension
+ # (rust/sglang-grpc) when `python -m build` builds the sglang wheel.
+ - name: Install protoc
+ run: sudo bash scripts/ci/utils/install_protoc.sh
+
- name: Build wheel
id: build
run: |
diff --git a/.github/workflows/release-whl-kernel.yml b/.github/workflows/release-whl-kernel.yml
index 119b04fa2de8..98f6f24ea3fa 100644
--- a/.github/workflows/release-whl-kernel.yml
+++ b/.github/workflows/release-whl-kernel.yml
@@ -50,6 +50,15 @@ jobs:
runner: arm-kernel-build-node
runs-on: ${{ matrix.runner }}
steps:
+ # Self-hosted build nodes retain the workspace across jobs. Prior builds
+ # leave root-owned artifacts under sgl-kernel/build/ that actions/checkout
+ # cannot remove, causing EACCES on rmdir. Wipe them via a throwaway root
+ # container before checkout recreates the workspace.
+ - name: Clean workspace (remove root-owned files from prior runs)
+ run: |
+ docker run --rm -v "${{ github.workspace }}:/workspace" alpine:3 \
+ sh -c 'rm -rf /workspace/..?* /workspace/.[!.]* /workspace/*' || true
+
- uses: actions/checkout@v4
with:
submodules: "recursive"
@@ -69,11 +78,42 @@ jobs:
BUILD_JOBS: 64
NVCC_THREADS: 8
+ # rename_wheels.sh tags the cu129 wheel's METADATA Version with +cu129
+ # (needed elsewhere — e.g. local install discrimination — see PR #23587),
+ # but PyPI rejects PEP 440 local version labels with HTTP 400. Repack a
+ # PyPI-clean copy with the +cu129 segment stripped into dist-pypi/, and
+ # upload that copy. dist/ is left untouched so the +cu129 wheel still
+ # flows through the upload-artifact -> sgl-project/whl index path below.
+ - name: Strip +cu129 local version for PyPI upload
+ working-directory: sgl-kernel
+ run: |
+ set -eux
+ pip install wheel
+ mkdir -p dist-pypi
+ for w in dist/*.whl; do
+ tmp=$(mktemp -d)
+ python3 -m wheel unpack "$w" --dest "$tmp"
+ unpacked=$(find "$tmp" -mindepth 1 -maxdepth 1 -type d | head -1)
+ info=$(find "$unpacked" -maxdepth 1 -type d -name "*.dist-info" | head -1)
+ meta="$info/METADATA"
+ orig=$(grep '^Version:' "$meta" | head -1 | sed 's/^Version:[[:space:]]*//')
+ new=$(echo "$orig" | sed 's/+cu[0-9]\+$//')
+ if [ "$orig" != "$new" ]; then
+ sed -i "s/^Version:.*/Version: ${new}/" "$meta"
+ old_base=$(basename "$info")
+ new_base="${old_base/${orig}/${new}}"
+ mv "$info" "$(dirname "$info")/${new_base}"
+ fi
+ python3 -m wheel pack "$unpacked" --dest-dir dist-pypi
+ rm -rf "$tmp"
+ done
+ ls -lh dist-pypi/
+
- name: Upload to PyPI
working-directory: sgl-kernel
run: |
pip install twine
- python3 -m twine upload --skip-existing dist/* -u __token__ -p ${{ secrets.PYPI_TOKEN_SGLANG_KERNEL }}
+ python3 -m twine upload --skip-existing dist-pypi/* -u __token__ -p ${{ secrets.PYPI_TOKEN_SGLANG_KERNEL }}
- name: Upload artifacts
uses: actions/upload-artifact@v4
@@ -149,6 +189,15 @@ jobs:
runner: arm-kernel-build-node
runs-on: ${{ matrix.runner }}
steps:
+ # Self-hosted build nodes retain the workspace across jobs. Prior builds
+ # leave root-owned artifacts under sgl-kernel/build/ that actions/checkout
+ # cannot remove, causing EACCES on rmdir. Wipe them via a throwaway root
+ # container before checkout recreates the workspace.
+ - name: Clean workspace (remove root-owned files from prior runs)
+ run: |
+ docker run --rm -v "${{ github.workspace }}:/workspace" alpine:3 \
+ sh -c 'rm -rf /workspace/..?* /workspace/.[!.]* /workspace/*' || true
+
- uses: actions/checkout@v4
with:
submodules: "recursive"
@@ -235,6 +284,15 @@ jobs:
python-version: ["3.10"]
rocm-version: ["700", "720"]
steps:
+ # Self-hosted build nodes retain the workspace across jobs. Prior builds
+ # leave root-owned artifacts under sgl-kernel/build/ that actions/checkout
+ # cannot remove, causing EACCES on rmdir. Wipe them via a throwaway root
+ # container before checkout recreates the workspace.
+ - name: Clean workspace (remove root-owned files from prior runs)
+ run: |
+ docker run --rm -v "${{ github.workspace }}:/workspace" alpine:3 \
+ sh -c 'rm -rf /workspace/..?* /workspace/.[!.]* /workspace/*' || true
+
- uses: actions/checkout@v4
with:
submodules: "recursive"
@@ -370,6 +428,15 @@ jobs:
python-version: ["3.10"]
musa-version: ["43"]
steps:
+ # Self-hosted build nodes retain the workspace across jobs. Prior builds
+ # leave root-owned artifacts under sgl-kernel/build/ that actions/checkout
+ # cannot remove, causing EACCES on rmdir. Wipe them via a throwaway root
+ # container before checkout recreates the workspace.
+ - name: Clean workspace (remove root-owned files from prior runs)
+ run: |
+ docker run --rm -v "${{ github.workspace }}:/workspace" alpine:3 \
+ sh -c 'rm -rf /workspace/..?* /workspace/.[!.]* /workspace/*' || true
+
- uses: actions/checkout@v4
with:
submodules: "recursive"
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 89c93558b5f6..e6a413074540 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -11,7 +11,7 @@ ARG GRACE_BLACKWELL_DEEPEP_BRANCH=gb200_blog_part_2
ARG HOPPER_SBO_DEEPEP_COMMIT=9f2fc4b3182a51044ae7ecb6610f7c9c3258c4d6
ARG DEEPEP_COMMIT=9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee
ARG BUILD_AND_DOWNLOAD_PARALLEL=8
-ARG SGL_KERNEL_VERSION=0.4.1
+ARG SGL_KERNEL_VERSION=0.4.1.post1
ARG SGL_VERSION
ARG USE_LATEST_SGLANG=0
ARG GDRCOPY_VERSION=2.5.1
diff --git a/docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx b/docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx
index e11b3a75a39e..8a26dd402413 100644
--- a/docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx
+++ b/docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx
@@ -29,13 +29,13 @@ tag: NEW
| DeepSeek-V4-Flash |
284B |
13B |
- single-node serving: B200 / GB300 / H200 on 4 GPUs |
+ single-node serving: B200 / GB200 / GB300 / H200 on 4 GPUs |
| DeepSeek-V4-Pro |
1.6T |
49B |
- high-capacity: B200 8 GPU / GB300 4 GPU / H200 16 GPU (2 nodes) |
+ high-capacity: B200 8 GPU / GB200 8 GPU (2 nodes) / GB300 4 GPU / H200 16 GPU (2 nodes) |
@@ -88,6 +88,10 @@ Please refer to the [official SGLang installation guide](../../../docs/get-start
NVIDIA B200 |
lmsysorg/sglang:deepseek-v4-blackwell |
+
+ | NVIDIA GB200 |
+ lmsysorg/sglang:deepseek-v4-grace-blackwell |
+
| NVIDIA GB300 |
lmsysorg/sglang:deepseek-v4-grace-blackwell |
diff --git a/docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx b/docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx
index 621c98091367..2f8c07e75c3e 100644
--- a/docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx
+++ b/docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx
@@ -4,6 +4,7 @@ export const DeepSeekV4Deployment = () => {
//
// Hardware (quantization determined by GPU generation):
// B200 → FP4 weights, Flash TP=4 / Pro TP=8 single-node
+ // GB200 → FP4 weights, Flash TP=4 / Pro TP=8 2-node
// GB300 → FP4 weights, Flash TP=4 / Pro TP=4 single-node
// H200 → FP8 weights, Flash TP=4 / Pro TP=16 2-node
// Model variant → HF slug:
@@ -27,6 +28,7 @@ export const DeepSeekV4Deployment = () => {
items: [
{ id: "b200", label: "B200 (FP4)", default: true },
{ id: "b300", label: "B300 (FP4)", default: false },
+ { id: "gb200", label: "GB200 (FP4)", default: false },
{ id: "gb300", label: "GB300 (FP4)", default: false },
{ id: "h200", label: "H200 (FP8)", default: false },
],
@@ -138,6 +140,8 @@ export const DeepSeekV4Deployment = () => {
"b200|big": { slug: "deepseek-ai/DeepSeek-V4-Pro", tp: 8, multinode: false },
"gb300|small": { slug: "deepseek-ai/DeepSeek-V4-Flash", tp: 4, multinode: false },
"gb300|big": { slug: "deepseek-ai/DeepSeek-V4-Pro", tp: 4, multinode: false },
+ "gb200|small": { slug: "deepseek-ai/DeepSeek-V4-Flash", tp: 4, multinode: false },
+ "gb200|big": { slug: "deepseek-ai/DeepSeek-V4-Pro", tp: 8, multinode: true, nnodes: 2 },
// H200 needs an FP8-only Instruct ckpt (deepseek-ai's Flash/Pro repos ship
// FP4-mixed weights that Hopper can't run). sgl-project publishes FP8
// repackagings for both variants.
@@ -150,6 +154,8 @@ export const DeepSeekV4Deployment = () => {
"b200|big": { tp: 8, multinode: false },
"gb300|small": { tp: 4, multinode: false },
"gb300|big": { tp: 4, multinode: false },
+ "gb200|small": { tp: 4, multinode: false },
+ "gb200|big": { tp: 8, multinode: true, nnodes: 2 },
"h200|small": { tp: 4, multinode: false },
"h200|big": { tp: 16, multinode: true, nnodes: 2 },
};
@@ -178,17 +184,29 @@ export const DeepSeekV4Deployment = () => {
"gb300|small|max-throughput",
"h200|small|cp",
"h200|small|pd-disagg",
+ "h200|big|low-latency",
+ "h200|big|balanced",
+ "h200|big|max-throughput",
"h200|big|pd-disagg",
"gb300|small|cp",
"gb300|big|cp",
"gb300|small|pd-disagg",
"gb300|big|pd-disagg",
+ "gb200|small|low-latency",
+ "gb200|small|balanced",
+ "gb200|small|max-throughput",
+ "gb200|small|cp",
+ "gb200|big|low-latency",
+ "gb200|big|balanced",
+ "gb200|big|max-throughput",
]);
// Recipes whose command is intentionally not yet provided (e.g. blocked by an
// upstream limitation). Showing a minimal placeholder is friendlier to users
// than emitting a commented-out invalid command.
const TBD_RECIPES = new Set([
"h200|big|cp",
+ "gb200|small|pd-disagg",
+ "gb200|big|pd-disagg",
]);
const TBD_PLACEHOLDER = "# to be provided";
const BEING_VERIFIED_NOTE =
@@ -242,18 +260,24 @@ export const DeepSeekV4Deployment = () => {
h200: ["SGLANG_DSV4_FP4_EXPERTS=0"], // allinone _ENV_H200
b200: [], // _ENV_B200 minus NVSHMEM
gb300: [], // _ENV_GB300
+ // GB200 multinode needs NCCL MNNVL for cross-node NVLink communication.
+ gb200: multinode ? ["NCCL_MNNVL_ENABLE=1", "NCCL_CUMEM_ENABLE=1"] : [],
}[hardware];
// Recipe-specific env (matches allinone exactly, taking size into account).
const recipeEnv = [];
if (recipe === "low-latency") {
- // H200 big low-latency has extra dispatch-token cap (allinone line 233).
+ // Big low-latency dispatch-token cap.
if (hardware === "h200" && isBig) {
recipeEnv.push("SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=128");
+ } else if (hardware === "gb200" && isBig) {
+ recipeEnv.push("SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256");
}
} else if (recipe === "balanced") {
if (hardware === "h200") {
- recipeEnv.push("SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256");
+ recipeEnv.push(isBig
+ ? "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=128"
+ : "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256");
} else {
// Blackwell: small=1024, big=256 (allinone ternary).
recipeEnv.push(isBig
@@ -262,7 +286,9 @@ export const DeepSeekV4Deployment = () => {
}
} else if (recipe === "max-throughput") {
if (hardware === "h200") {
- recipeEnv.push("SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256");
+ recipeEnv.push(isBig
+ ? "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=128"
+ : "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256");
} else {
recipeEnv.push(isBig
? "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256"
@@ -293,6 +319,7 @@ export const DeepSeekV4Deployment = () => {
// allinone:
// H200 small: pure TP + MTP_314
// H200 big: DP-attn + DeepEP + MTP_314 + cg=32 max-run=64 + multi-node + mem-frac 0.82
+ // GB200 big: pure TP + multinode + flashinfer_mxfp4 + MTP_314 + mem-frac 0.82 (no DP-attn/DeepEP)
// Blackwell: TP + flashinfer_mxfp4 + MTP_314 + chunked-prefill-size 4096 + autotune-fix
// Big Blackwell additionally: mem-frac 0.82
flags.push(` --tp ${tp}`);
@@ -308,8 +335,8 @@ export const DeepSeekV4Deployment = () => {
flags.push(" --moe-runner-backend flashinfer_mxfp4");
}
if (hardware === "h200" && isBig) {
- flags.push(" --cuda-graph-max-bs 32");
- flags.push(" --max-running-requests 64");
+ flags.push(" --cuda-graph-max-bs 8");
+ flags.push(" --max-running-requests 32");
}
// MTP 3/4
flags.push(" --speculative-algo EAGLE");
@@ -320,7 +347,7 @@ export const DeepSeekV4Deployment = () => {
flags.push(" --chunked-prefill-size 4096");
flags.push(" --disable-flashinfer-autotune");
}
- if (isBig) flags.push(" --mem-fraction-static 0.82");
+ if (isBig) flags.push(" --mem-fraction-static 0.88");
} else if (recipe === "balanced") {
// allinone balanced: TP + DP + DP-attn + DeepEP + MTP_112.
// H200 small: cg=128 max-run=128 | H200 big: cg=128 max-run=128 (same)
@@ -335,8 +362,17 @@ export const DeepSeekV4Deployment = () => {
flags.push(" --speculative-num-steps 1");
flags.push(" --speculative-eagle-topk 1");
flags.push(" --speculative-num-draft-tokens 2");
- if (isBig) flags.push(" --mem-fraction-static 0.82");
- if (hardware === "h200") {
+ if (hardware === "h200" && isBig) {
+ flags.push(" --mem-fraction-static 0.88");
+ } else if (isBig && hardware === "gb200") {
+ flags.push(" --mem-fraction-static 0.78");
+ } else if (isBig) {
+ flags.push(" --mem-fraction-static 0.82");
+ }
+ if (hardware === "h200" && isBig) {
+ flags.push(" --cuda-graph-max-bs 8");
+ flags.push(" --max-running-requests 32");
+ } else if (hardware === "h200") {
flags.push(" --cuda-graph-max-bs 128");
flags.push(" --max-running-requests 128");
} else if (isBig && hardware === "b200") {
@@ -345,6 +381,9 @@ export const DeepSeekV4Deployment = () => {
} else if (isBig && hardware === "gb300") {
flags.push(" --cuda-graph-max-bs 128");
flags.push(" --max-running-requests 256");
+ } else if (isBig && hardware === "gb200") {
+ flags.push(" --cuda-graph-max-bs 64");
+ flags.push(" --max-running-requests 128");
}
// allinone H200 gates DEEPEP_LARGE_SMS_FLAG on !multinode — only H200 big
// is multi-node; all Blackwell cells get the flag unconditionally.
@@ -359,7 +398,13 @@ export const DeepSeekV4Deployment = () => {
flags.push(" --enable-dp-attention");
if (multinode) flags.push(...multiNodeFlags(nnodes));
flags.push(" --moe-a2a-backend deepep");
- if (isBig) flags.push(" --mem-fraction-static 0.82");
+ if (hardware === "h200" && isBig) {
+ flags.push(" --mem-fraction-static 0.88");
+ } else if (isBig && hardware === "gb200") {
+ flags.push(" --mem-fraction-static 0.78");
+ } else if (isBig) {
+ flags.push(" --mem-fraction-static 0.82");
+ }
if (hardware === "h200") {
flags.push(" --cuda-graph-max-bs 128");
flags.push(" --max-running-requests 256");
@@ -369,6 +414,9 @@ export const DeepSeekV4Deployment = () => {
} else if (isBig && hardware === "gb300") {
flags.push(" --cuda-graph-max-bs 128");
flags.push(" --max-running-requests 256");
+ } else if (isBig && hardware === "gb200") {
+ flags.push(" --cuda-graph-max-bs 64");
+ flags.push(" --max-running-requests 256");
}
if (!multinode) flags.push(DEEPEP_LARGE_SMS_FLAG);
} else if (recipe === "cp") {
@@ -416,7 +464,18 @@ export const DeepSeekV4Deployment = () => {
const envAll = [...HW_ENV, ...recipeEnv, ...COMMON_ENV];
const envBlock = envAll.length ? envAll.join(" \\\n") + " \\\n" : "";
const base = `${envBlock}sglang serve \\\n${flags.join(" \\\n")}`;
- const withMultinode = multinode ? prependMultiNodeNote(base, nnodes) : base;
+ // GB200 multinode may need machine-specific NVSHMEM / Gloo env vars;
+ // emit them as commented hints above the env block so users know to check.
+ let cmd = base;
+ if (hardware === "gb200" && multinode) {
+ cmd =
+ `# The following env vars may be needed depending on your cluster:\n` +
+ `# GLOO_SOCKET_IFNAME=\n` +
+ `# NVSHMEM_ENABLE_NIC_PE_MAPPING=1\n` +
+ `# NVSHMEM_HCA_LIST=\n` +
+ cmd;
+ }
+ const withMultinode = multinode ? prependMultiNodeNote(cmd, nnodes) : cmd;
const verifyKey = `${hardware}|${modelSize}|${recipe}`;
if (TBD_RECIPES.has(verifyKey)) return TBD_PLACEHOLDER;
return VERIFIED_RECIPES.has(verifyKey)
@@ -447,14 +506,15 @@ export const DeepSeekV4Deployment = () => {
const specKey = `${hardware}|${modelSize}`;
const { tp: pdTp, multinode, nnodes } = PD_TP_SPEC[specKey];
const slug = HW_SIZE_SPEC[specKey].slug;
- const ibDevice = { h200: "mlx5_0", b200: "mlx5_7", gb300: "" }[hardware];
+ const ibDevice = { h200: "mlx5_0", b200: "mlx5_7", gb300: "", gb200: "" }[hardware];
const isGB300 = hardware === "gb300";
- const isBlackwell = hardware === "b200" || isGB300;
+ const isBlackwell = hardware === "b200" || hardware === "gb200" || isGB300;
const HW_ENV = {
h200: ["SGLANG_DSV4_FP4_EXPERTS=0"],
b200: [],
gb300: [],
+ gb200: [],
}[hardware];
// Whitelist #5: only SGLANG_MOONCAKE_CUSTOM_MEM_POOL kept; MC_FORCE_MNNVL /
// NCCL_MNNVL_ENABLE / NCCL_CUMEM_ENABLE may also be needed depending on the
diff --git a/python/pyproject.toml b/python/pyproject.toml
index 17b2f4bae2cd..bfdcd2762957 100755
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -59,7 +59,7 @@ dependencies = [
"sentencepiece",
"setproctitle",
"flash-attn-4>=4.0.0b9",
- "sglang-kernel==0.4.1",
+ "sglang-kernel==0.4.1.post1",
"soundfile==0.13.1",
"tiktoken",
"timm==1.0.16",
diff --git a/python/run_dsv4.sh b/python/run_dsv4.sh
new file mode 100755
index 000000000000..0a7ece49dd58
--- /dev/null
+++ b/python/run_dsv4.sh
@@ -0,0 +1,45 @@
+#export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+export SGLANG_REASONING_EFFORT=max
+
+export SGLANG_OPT_USE_FUSED_COMPRESS=false #use PyTorch implemented compressor
+export SGLANG_OPT_USE_OLD_COMPRESSOR=true #use old compressor
+export SGLANG_OPT_USE_TILELANG_SWA_PREPARE=false #use old prepare
+export SGLANG_OPT_USE_JIT_KERNEL_FUSED_TOPK=false #use old topk
+export SGLANG_OPT_USE_FUSED_HASH_TOPK=false #AMD: hash_topk JIT needs CUDA toolchain
+
+export SGLANG_HACK_FLASHMLA_BACKEND=torch
+export SGLANG_OPT_DEEPGEMM_HC_PRENORM=false #use old prenorm
+
+export SGLANG_OPT_USE_TILELANG_MHC_PRE=false #use torch hc_pre
+export SGLANG_OPT_USE_TILELANG_MHC_POST=false #use torch hc_post
+
+export SGLANG_ENABLE_THINKING=1
+export SGLANG_USE_AITER=1
+export SGLANG_USE_ROCM700A=1
+export SGLANG_TOPK_TRANSFORM_512_TORCH=1
+export SGLANG_FP8_PAGED_MQA_LOGITS_TORCH=1
+
+export SGLANG_DSV4_FP4_EXPERTS=false
+
+export SGLANG_OPT_DPSK_V4_RADIX=0
+export SGLANG_OPT_USE_OVERLAP_STORE_CACHE=false #non-radix backend has no store_cache method
+export SGLANG_OPT_USE_FUSED_STORE_CACHE=false #fused_store_cache JIT needs CUDA toolchain
+
+export SGLANG_FORCE_TRITON_MOE_FP8=1 # this is required to apply swiglu_limit clamp in fused_moe_triton
+
+python3 -m sglang.launch_server \
+ --model-path /dockerx/data2/models/DeepSeek-V4-Flash-FP8 \
+ --trust-remote-code \
+ --tp 8 \
+ --enable-dp-attention \
+ --disable-radix-cache \
+ --attention-backend compressed \
+ --max-running-request 256 \
+ --page-size 256 \
+ --chunked-prefill-size 8192 \
+ --port 8000 \
+ --disable-shared-experts-fusion \
+ --disable-cuda-graph \
+ --tool-call-parser deepseekv4 \
+ --reasoning-parser deepseek-v4
diff --git a/python/sglang/jit_kernel/activation.py b/python/sglang/jit_kernel/activation.py
index 5756c2f38287..89b28bbf6d92 100644
--- a/python/sglang/jit_kernel/activation.py
+++ b/python/sglang/jit_kernel/activation.py
@@ -38,6 +38,10 @@ def _jit_activation_module(dtype: torch.dtype) -> Module:
extra_cuda_cflags=_fast_math_flags(),
cuda_wrappers=[
("run_activation", f"ActivationKernel<{args}>::run_activation"),
+ (
+ "run_activation_filtered",
+ f"ActivationKernel<{args}>::run_activation_filtered",
+ ),
],
)
@@ -56,30 +60,68 @@ def _run_activation_inplace(
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]
+ 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)
- _run_activation_inplace(op_name, input, out)
+ 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
def silu_and_mul(
- input: torch.Tensor, out: Optional[torch.Tensor] = None
+ 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)
+ return run_activation("silu", input, out, expert_ids, expert_step)
def gelu_and_mul(
- input: torch.Tensor, out: Optional[torch.Tensor] = None
+ 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)
+ return run_activation("gelu", input, out, expert_ids, expert_step)
def gelu_tanh_and_mul(
- input: torch.Tensor, out: Optional[torch.Tensor] = None
+ 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)
+ return run_activation("gelu_tanh", input, out, expert_ids, expert_step)
diff --git a/python/sglang/jit_kernel/benchmark/bench_activation.py b/python/sglang/jit_kernel/benchmark/bench_activation.py
index 2caac5551c8e..3f0ba2f6c85e 100644
--- a/python/sglang/jit_kernel/benchmark/bench_activation.py
+++ b/python/sglang/jit_kernel/benchmark/bench_activation.py
@@ -82,5 +82,76 @@ def f():
return run_benchmark(f, scale=NUM_LAYERS)
+FILTER_OPS = ["silu", "gelu"]
+FILTER_BS = get_benchmark_range(
+ full_range=[64, 256, 1024, 4096, 16384], ci_range=[1024]
+)
+FILTER_DIMS = get_benchmark_range(full_range=[1024, 4096, 8192], ci_range=[4096])
+FILTER_RATIOS = get_benchmark_range(full_range=[0.0, 0.25, 0.5], ci_range=[0.25])
+FILTER_CONFIGS = list(
+ itertools.product(FILTER_OPS, FILTER_DIMS, FILTER_BS, FILTER_RATIOS)
+)
+
+
+def _make_expert_ids(num_tokens: int, skip_ratio: float) -> torch.Tensor:
+ expert_ids = torch.randint(
+ low=0, high=8, size=(num_tokens,), dtype=torch.int32, device=DEFAULT_DEVICE
+ )
+ if skip_ratio > 0:
+ skip = torch.rand(num_tokens, device=DEFAULT_DEVICE) < skip_ratio
+ expert_ids[skip] = -1
+ return expert_ids
+
+
+@triton.testing.perf_report(
+ triton.testing.Benchmark(
+ x_names=["op_name", "dim", "batch_size", "skip_ratio"],
+ x_vals=FILTER_CONFIGS,
+ line_arg="provider",
+ line_vals=["unfiltered", "filtered"],
+ line_names=["JIT (no filter_expert)", "JIT (with expert_ids)"],
+ styles=[("blue", "--"), ("orange", "-")],
+ ylabel="us",
+ plot_name="activation-filter-expert",
+ args={},
+ )
+)
+def benchmark_filter(
+ op_name: str, dim: int, batch_size: int, skip_ratio: float, provider: str
+):
+ x = torch.randn(
+ NUM_LAYERS,
+ batch_size,
+ 2 * dim,
+ dtype=DEFAULT_DTYPE,
+ device=DEFAULT_DEVICE,
+ )
+ out = torch.empty(
+ NUM_LAYERS,
+ batch_size,
+ dim,
+ dtype=DEFAULT_DTYPE,
+ device=DEFAULT_DEVICE,
+ )
+ expert_ids = _make_expert_ids(batch_size, skip_ratio)
+
+ jit_fn = silu_and_mul_jit if op_name == "silu" else gelu_and_mul_jit
+
+ if provider == "unfiltered":
+
+ def f():
+ for i in range(NUM_LAYERS):
+ jit_fn(x[i], out[i])
+
+ else: # filtered
+
+ def f():
+ for i in range(NUM_LAYERS):
+ jit_fn(x[i], out[i], expert_ids=expert_ids, expert_step=1)
+
+ return run_benchmark(f, scale=NUM_LAYERS)
+
+
if __name__ == "__main__":
benchmark.run(print_data=True)
+ benchmark_filter.run(print_data=True)
diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c128.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c128.cuh
new file mode 100644
index 000000000000..d91ad1d0e685
--- /dev/null
+++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c128.cuh
@@ -0,0 +1,524 @@
+#include
+#include
+
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include
+
+#include
+#include
+#include
+
+#include
+
+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;
+ 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
+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;
+ const auto element_size = head_dim * 2;
+ const auto gmem = tile::Memory{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
+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;
+ const auto gmem_in = tile::Memory{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(score[j][i]) + cast(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(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(local_val_max);
+ const auto rescale = expf(local_val_max - global_val_max);
+ const auto global_exp_sum = warp::reduce_sum(local_exp_sum * rescale);
+ const auto final_scale = rescale / global_exp_sum;
+ const auto global_product = warp::reduce_sum(local_product * final_scale);
+ kv_out[local_elem_id] = cast(global_product);
+ }
+}
+
+template
+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(_kv_score_buffer);
+ const auto kv_buf = kv_score_buffer + index * (kElementSize * 128) + split_offset;
+
+ // kv input
+ const auto kv_score_input = static_cast(_kv_score_input);
+ const auto kv_src = kv_score_input + global_bid * kElementSize + split_offset;
+
+ // kv output
+ const auto kv_compressed_output = static_cast(_kv_compressed_output);
+ const auto kv_out = kv_compressed_output + global_bid * kHeadDim + split_offset;
+
+ // score bias (ape)
+ const auto score_bias = static_cast(_score_bias) + split_offset;
+
+ PDLWaitPrimary();
+
+ /// 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();
+}
+
+// compress kernel
+template
+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(_kv_score_input);
+ const auto kv_src = kv_score_input + ragged_id * kElementSize + split_offset;
+
+ // kv output
+ const auto kv_compressed_output = static_cast(_kv_compressed_output);
+ const auto kv_out = kv_compressed_output + ragged_id * kHeadDim + split_offset;
+
+ // score bias (ape)
+ const auto score_bias = static_cast(_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(_kv_score_buffer);
+ const auto kv_buf = kv_score_buffer + index * (kElementSize * 128) + split_offset;
+
+ PDLWaitPrimary();
+
+ // 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();
+}
+
+template
+struct FlashCompress128Kernel {
+ static constexpr auto decode_kernel = flash_c128_decode;
+ template
+ static constexpr auto prefill_kernel = flash_c128_prefill;
+ static constexpr auto prefill_c_kernel = prefill_kernel;
+ static constexpr auto prefill_w_kernel = prefill_kernel;
+ 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 /* UNUSED */) {
+ using namespace host;
+
+ // this should not happen in practice
+ auto B = SymbolicSize{"batch_size"};
+ auto device = SymbolicDevice{};
+ device.set_options();
+
+ TensorMatcher({-1, 128, kHeadDim * 2}) // kv score
+ .with_dtype()
+ .with_device(device)
+ .verify(kv_score_buffer);
+ TensorMatcher({B, kHeadDim * 2}) // kv score input
+ .with_dtype()
+ .with_device(device)
+ .verify(kv_score_input);
+ TensorMatcher({B, kHeadDim}) // kv compressed output
+ .with_dtype()
+ .with_device(device)
+ .verify(kv_compressed_output);
+ TensorMatcher({128, kHeadDim}) // ape
+ .with_dtype()
+ .with_device(device)
+ .verify(ape);
+ TensorMatcher({B}) // indices
+ .with_dtype()
+ .with_device(device)
+ .verify(indices);
+ TensorMatcher({B}) // seq lens
+ .with_dtype()
+ .with_device(device)
+ .verify(seq_lens);
+
+ const auto batch_size = static_cast(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(indices.data_ptr()),
+ .seq_lens = static_cast(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 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();
+
+ TensorMatcher({-1, 128, kHeadDim * 2}) // kv score
+ .with_dtype()
+ .with_device(device_)
+ .verify(kv_score_buffer);
+ TensorMatcher({N, kHeadDim * 2}) // kv score input
+ .with_dtype()
+ .with_device(device_)
+ .verify(kv_score_input);
+ TensorMatcher({N, kHeadDim}) // kv compressed output
+ .with_dtype()
+ .with_device(device_)
+ .verify(kv_compressed_output);
+ TensorMatcher({128, kHeadDim}) // ape
+ .with_dtype()
+ .with_device(device_)
+ .verify(ape);
+ TensorMatcher({B}) // indices
+ .with_dtype()
+ .with_device(device_)
+ .verify(indices);
+ TensorMatcher({X, compress::kPrefillPlanDim}) // compress plan
+ .with_dtype()
+ .with_device(device_)
+ .verify(compress_plan);
+ TensorMatcher({Y, compress::kPrefillPlanDim}) // write plan
+ .with_dtype()
+ .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()
+ .with_device(device_)
+ .verify(load_indices);
+
+ const auto device = device_.unwrap();
+ const auto batch_size = static_cast(B.unwrap());
+ const auto num_q_tokens = static_cast(N.unwrap());
+ const auto num_c = static_cast(X.unwrap());
+ const auto num_w = static_cast(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(indices.data_ptr()),
+ .load_indices = static_cast(load_indices.data_ptr()),
+ .compress_plan = static_cast(compress_plan.data_ptr()),
+ .write_plan = static_cast(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
diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c4.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c4.cuh
new file mode 100644
index 000000000000..145ab1fb081e
--- /dev/null
+++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c4.cuh
@@ -0,0 +1,549 @@
+#include
+#include
+
+#include
+#include
+#include
+#include
+#include
+
+#include
+
+#include
+#include
+#include
+
+#include
+
+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
+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;
+ const auto element_size = head_dim * 4;
+ const auto gmem = tile::Memory::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
+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;
+ const auto gmem_in = tile::Memory::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(0.0f));
+ score[i].fill(cast(kFloatNegInf));
+ }
+ }
+
+ /// NOTE: part 2: safe online softmax + weighted sum
+ using StorageOut = AlignedVector;
+ const auto gmem_out = tile::Memory::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(score[j][i]) + cast(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(kv[j][i]) * exp_score;
+ sum_exp_value += exp_score;
+ }
+
+ result[i] = cast(sum_product / sum_exp_value);
+ }
+
+ gmem_out.store(kv_out, result);
+}
+
+template
+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(_kv_score_buffer);
+
+ // kv input
+ const auto kv_score_input = static_cast(_kv_score_input);
+ const auto kv_src = kv_score_input + global_bid * kElementSize + split_offset;
+
+ // kv output
+ const auto kv_compressed_output = static_cast(_kv_compressed_output);
+ const auto kv_out = kv_compressed_output + global_bid * kHeadDim + split_offset;
+
+ // score bias (ape)
+ const auto score_bias = static_cast(_score_bias) + split_offset;
+
+ PDLWaitPrimary();
+
+ /// 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(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(kv_buf, kv_src, kv_out, score_bias, kHeadDim, seq_len, /*window_size=*/8);
+ }
+ }
+
+ PDLTriggerSecondary();
+}
+
+template
+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(_kv_score_buffer);
+
+ // kv input
+ const auto kv_score_input = static_cast(_kv_score_input);
+ const auto kv_src = kv_score_input + ragged_id * kElementSize + split_offset;
+
+ // kv output
+ const auto kv_compressed_output = static_cast(_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(_score_bias) + split_offset;
+ const auto seq_len = position + 1;
+ const int32_t index = indices[global_bid];
+
+ PDLWaitPrimary();
+
+ 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(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(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(kv_buf, kv_src, kv_out, score_bias, kHeadDim, seq_len, window_len);
+ }
+ }
+
+ PDLTriggerSecondary();
+}
+
+template
+struct FlashCompress4Kernel {
+ template
+ static constexpr auto decode_kernel = flash_c4_decode;
+ template
+ static constexpr auto prefill_kernel = flash_c4_prefill;
+ template
+ static constexpr auto prefill_c_kernel = prefill_kernel;
+ template
+ static constexpr auto prefill_w_kernel = prefill_kernel;
+ 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 extra) {
+ using namespace host;
+
+ // this should not happen in practice
+ auto B = SymbolicSize{"batch_size"};
+ auto device_ = SymbolicDevice{};
+ device_.set_options();
+ 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()
+ .with_device(device_)
+ .verify(kv_score_buffer);
+ TensorMatcher({B, kHeadDim * 4}) // kv score input
+ .with_dtype()
+ .with_device(device_)
+ .verify(kv_score_input);
+ TensorMatcher({B, kHeadDim}) // kv compressed output
+ .with_dtype()
+ .with_device(device_)
+ .verify(kv_compressed_output);
+ TensorMatcher({8, kHeadDim}) // ape
+ .with_dtype()
+ .with_device(device_)
+ .verify(ape);
+ TensorMatcher({B}) // indices
+ .with_dtype()
+ .with_device(device_)
+ .verify(indices);
+ TensorMatcher({B}) // seq lens
+ .with_dtype()
+ .with_device(device_)
+ .verify(seq_lens);
+
+ const auto device = device_.unwrap();
+ const auto batch_size = static_cast(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(indices.data_ptr()),
+ .seq_lens = static_cast(seq_lens.data_ptr()),
+ .extra = static_cast(extra_ptr),
+ .batch_size = batch_size,
+ };
+ const auto kernel = extra_ptr != nullptr ? decode_kernel //
+ : decode_kernel;
+ 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 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();
+ 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()
+ .with_device(device_)
+ .verify(kv_score_buffer);
+ TensorMatcher({N, kHeadDim * 4}) // kv score input
+ .with_dtype()
+ .with_device(device_)
+ .verify(kv_score_input);
+ TensorMatcher({N, kHeadDim}) // kv compressed output
+ .with_dtype()
+ .with_device(device_)
+ .verify(kv_compressed_output);
+ TensorMatcher({8, kHeadDim}) // ape
+ .with_dtype()
+ .with_device(device_)
+ .verify(ape);
+ TensorMatcher({B}) // indices
+ .with_dtype()
+ .with_device(device_)
+ .verify(indices);
+ TensorMatcher({X, compress::kPrefillPlanDim}) // compress plan
+ .with_dtype()
+ .with_device(device_)
+ .verify(compress_plan);
+ TensorMatcher({Y, compress::kPrefillPlanDim}) // write plan
+ .with_dtype()
+ .with_device(device_)
+ .verify(write_plan);
+
+ const auto device = device_.unwrap();
+ const auto batch_size = static_cast(B.unwrap());
+ const auto num_q_tokens = static_cast(N.unwrap());
+ const auto num_c = static_cast(X.unwrap());
+ const auto num_w = static_cast(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(indices.data_ptr()),
+ .extra = static_cast(extra_ptr),
+ .compress_plan = static_cast(compress_plan.data_ptr()),
+ .write_plan = static_cast(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 //
+ : prefill_c_kernel;
+ 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 //
+ : prefill_w_kernel;
+ 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& 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()
+ .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(data_ptr) % 16 == 0, "extra tensor is not properly aligned");
+ }
+ return data_ptr;
+ }
+};
+
+} // namespace
diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/common.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/common.cuh
new file mode 100644
index 000000000000..46acaa9c46b3
--- /dev/null
+++ b/python/sglang/jit_kernel/csrc/deepseek_v4/common.cuh
@@ -0,0 +1,208 @@
+#include
+#include
+
+#include
+
+#include
+
+namespace host::compress {
+
+using PlanResult = tvm::ffi::Tuple;
+
+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(-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(-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();
+ return true;
+ } else {
+ device.set_options();
+ return false;
+ }
+ }();
+ TensorMatcher({N}) // extend_lens and seq_lens
+ .with_dtype()
+ .with_device(device)
+ .verify(extend_lens)
+ .verify(seq_lens);
+ TensorMatcher({M, kPrefillPlanDim}) // compress_plan and write_plan
+ .with_dtype()
+ .with_device(device)
+ .verify(compress_plan)
+ .verify(write_plan);
+
+ const auto params = CompressParams{
+ .compress_plan = static_cast(compress_plan.data_ptr()),
+ .write_plan = static_cast(write_plan.data_ptr()),
+ .seq_lens = static_cast(seq_lens.data_ptr()),
+ .extend_lens = static_cast(extend_lens.data_ptr()),
+ .batch_size = static_cast(N.unwrap()),
+ .num_tokens = static_cast(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
diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope.cuh
new file mode 100644
index 000000000000..d3953578b925
--- /dev/null
+++ b/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope.cuh
@@ -0,0 +1,254 @@
+#include
+#include
+
+#include
+#include
+#include
+#include
+#include
+
+#include
+
+#include
+
+#include
+#include
+
+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
+__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(handle)[work_id];
+ input = static_cast(_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(_input) + work_id * kHeadDim;
+ const auto seq_len = static_cast(handle)[work_id];
+ if (seq_len % compress_ratio != 0) return;
+ position = seq_len - compress_ratio;
+ } else if constexpr (kMode == DefaultForward) {
+ input = static_cast(_input) + work_id * kHeadDim;
+ position = static_cast(handle)[work_id];
+ } else {
+ static_assert(host::dependent_false_v, "Unsupported Mode");
+ }
+
+ using Storage = AlignedVector;
+ __shared__ Storage s_rope_input[kNumWarps][kRopeSize];
+
+ // prefetch freq
+ const auto mem_freq = tile::Memory::warp();
+ const auto freq = mem_freq.load(freqs_cis + position * kRopeDim);
+
+ PDLWaitPrimary();
+
+ // part 1: norm
+ {
+ const auto gmem = tile::Memory::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(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(input_vec[i][j]);
+ const auto fp32_weight = cast(weight_vec[i][j]);
+ input_vec[i][j] = cast(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;
+ const auto mem_elem = tile::Memory::warp();
+ const auto elem = mem_elem.load(s_rope_input[warp_id]);
+ const auto [x_real, x_imag] = cast(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(output));
+ }
+
+ PDLTriggerSecondary();
+}
+
+template
+struct FusedNormRopeKernel {
+ template
+ static constexpr auto fused_kernel = fused_norm_rope;
+
+ 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(_mode);
+
+ auto B = SymbolicSize{"num_q_tokens"};
+ auto N = SymbolicSize{"num_compress_tokens"};
+ auto device_ = SymbolicDevice{};
+ device_.set_options();
+
+ TensorMatcher({B, kHeadDim}) // input
+ .with_dtype()
+ .with_device(device_)
+ .verify(input);
+ TensorMatcher({kHeadDim}) // weight
+ .with_dtype()
+ .with_device(device_)
+ .verify(weight);
+ TensorMatcher({-1, kRopeDim}) // freqs_cis
+ .with_dtype()
+ .with_device(device_)
+ .verify(freqs_cis);
+ switch (mode) {
+ case CompressExtend:
+ TensorMatcher({N, compress::kPrefillPlanDim}) // plan
+ .with_dtype()
+ .with_device(device_)
+ .verify(handle);
+ RuntimeCheck(compress_ratio > 0);
+ break;
+ case CompressDecode:
+ TensorMatcher({N}) // seq_len
+ .with_dtype()
+ .with_device(device_)
+ .verify(handle);
+ RuntimeCheck(compress_ratio > 0);
+ break;
+ case DefaultForward:
+ TensorMatcher({N}) // position
+ .with_dtype()
+ .with_device(device_)
+ .verify(handle);
+ RuntimeCheck(compress_ratio == 0);
+ break;
+ default:
+ Panic("unsupported forward mode: ", static_cast(mode));
+ }
+
+ // launch kernel
+ const auto num_compress_tokens = static_cast(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(freqs_cis.data_ptr()),
+ .compress_ratio = compress_ratio,
+ };
+ const auto num_blocks = div_ceil(num_compress_tokens, kNumWarps);
+ using KernelType = std::decay_t)>;
+ static constexpr KernelType kernel_table[3] = {
+ [static_cast(CompressExtend)] = fused_kernel,
+ [static_cast(CompressDecode)] = fused_kernel,
+ [static_cast(DefaultForward)] = fused_kernel,
+ };
+ const auto kernel = kernel_table[static_cast(mode)];
+ LaunchKernel(num_blocks, kBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(kernel, params);
+ }
+};
+
+} // namespace
diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh
new file mode 100644
index 000000000000..cf422c58bc1e
--- /dev/null
+++ b/python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh
@@ -0,0 +1,137 @@
+#include
+#include
+
+#include
+#include
+#include
+
+#include
+
+#include
+#include
+
+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
+__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();
+
+ 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();
+}
+
+template
+struct HashTopKKernel {
+ static constexpr auto kernel = moe_hash_topk_fused