Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/CODEOWNERS
Original file line number Diff line number Diff line change
Expand Up @@ -6,3 +6,4 @@
**/*deepseek_v41* @YAMY1234
**/*qwen3_5* @YAMY1234
**/*kimi_k3* @leejnau
**/*glm5_3* @nvrohanv
4 changes: 4 additions & 0 deletions NOTICE
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,10 @@ decode step of DeepSeek-V4.1 (the layer loop, the mHC boundaries, the attention,
MoE, Engram and the head) follows SGLang's own model code (Apache License 2.0,
Copyright SGLang Team). Each adapted file names its origin in its header.

The stock verify step of GLM-5.3 (the layer loop, the residual steps, the attention,
the MoE and the logits) is adapted from SGLang's own model code in the same way
(Apache License 2.0, Copyright SGLang Team).

Portions are adapted from SGLang (https://github.com/sgl-project/sglang): the stock
verify step of Qwen3.5 (the layer loop, the GatedDeltaNet, the attention and the MoE)
follows SGLang's own model code (Apache License 2.0, Copyright SGLang Team). Qwen3.5's
Expand Down
55 changes: 52 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,21 @@ The package loads as an SGLang [general plugin](https://github.com/sgl-project/s
| DeepSeek-V4.1-Flash | GB300 (SM103) | 4 | DSpark target verify steps of 1 to 384 rows (batch 1 to 64 at 6 rows a request) |
| Qwen3.5-397B-A17B (NVFP4) | GB300 (SM103) | 8 (two trays, one NVLink domain) | MTP target verify steps of 7 rows (batch 1) |
| Kimi K3 (packed MXFP4) | GB300 (SM103) | 8 | One-token decode steps of 1–8, 16, 32 and 64 rows |
| GLM-5.3, GLM-5.2 (NVFP4) | B300 / GB300 (SM103), Rubin (SM107) | 8 (one node, or two trays of one NVLink domain) | MTP target verify steps of 4 to 16 rows (batch 1 to 4 at 4 rows a request, batch 1 to 2 at 6); the NextN draft step in its captured graphs |

For DeepSeek-V4.1-Flash and Qwen3.5, larger verify steps, prefill and extend run
SGLang's standard path. Their draft specs provide the DSpark proposal and the
Qwen3.5 EAGLE entries, respectively. For Kimi K3, the plugin covers the 93-layer
language model (24 MLA and 69 KDA layers) at the listed widths; every other call
runs the standard path.
runs the standard path. For GLM-5.3, larger verify steps, prefill and extend run
SGLang's standard path, and the NextN draft step runs on its draft spec.

GLM-5.2 has GLM-5.3's architecture and the same value of every config field the
GLM-5.3 specs pin (`HF_CONFIG` in `glm5_3/__init__.py`), so the same specs serve it,
and what this README says of GLM-5.3 holds for GLM-5.2. On `nvidia/GLM-5.2-NVFP4`
(8 x B300) the megakernel engages, a test-only step-by-step comparison with SGLang's
own target verify and draft steps shows no difference, and GSM8K-200 (strict match)
scores 0.975 with the megakernel and 0.960 without it.

The Kimi K3 adapter is experimental. On the pinned GB300/TP8 configuration
below, every listed CUDA-graph width and padding boundary passed dummy-weight
Expand Down Expand Up @@ -70,6 +79,19 @@ These slots change the arithmetic (numeric: each sums in another order than SGLa

The sampling and acceptance after the verify, and the GatedDeltaNet extend of prefill, stay SGLang's own code. In an SGLang that has the switches, `SGLANG_ENABLE_EAGLE_POST_VERIFY_CUDA_GRAPH=1` runs the sampling and acceptance in CUDA graphs, and `SGLANG_ENABLE_GDN_PREFILL_GRAPH_EXTEND=1` captures the GatedDeltaNet extend inside the breakable prefill CUDA graphs.

GLM-5.3 (`glm5_3/`) follows the same layout: the executor's layer loop runs the hops between the sublayers (`boundary/`), the attention (`attention/`) and the MoE (`moe/`), then the logits (`logits/`), each with a stock path that reproduces SGLang's verify step bit for bit, and CuTe DSL kernels (`sm103/`) where a slot serves its width:

| Slot | Rows | Kernel |
| --- | --- | --- |
| Attention (`attention`) | 4–16 | The q_a / kv_a projection on a small-M GEMM (with the indexer's key / head-gate projection, up to 8 rows), the MLA prologue (norms, q_b, W_UK absorption, RoPE and FP8 query, the MLA and index-K cache rows), the indexer query, the sparse MLA over the fused top-k page table and the value up-projection; the indexer's top-k stays SGLang's |
| Routed and shared experts (`moe`) | 4–16 | One small-M NVFP4 engine launch for the routed experts and the BF16 shared expert, for 1, 2 or 3–4 groups of up to 4 rows (a request's rows with 4 draft tokens; 6 rows run as 4 + 2, 12 as 4 + 4 + 4) |
| All-reduce + residual + RMSNorm (`hops`) | 4–16 | One hop per tensor-parallel sum over symmetric memory: `o_proj` with the router on MoE layers, the MoE finalize folded into the MLP-output all-reduce, and the final norm |
| Shared-top-k q_a / kv_a (`qkv_a_small_m_shared`) | 4–8 | The small-M projection on the layers that reuse the previous layer's top-k (one packed weight copy each) |
| Deferred input norm (`dnorm`) | 4–8 | Those layers' input RMSNorm applied in the projection's epilogue from the fold hop's block sums (moves the rounding points; its tolerances are in `glm5_3/spec.py`) |
| NextN draft step (`draft`) | every captured draft-decode and draft-extend graph | The draft's embedding gather + eh norms in one kernel on the replicated table (any row count), the column-sharded `eh_proj` (1–4 rows, and 8, 12 and 16 with 4 draft tokens or 6 and 12 with 6; gathered by the hops' all-reduce on the target's mailbox), the distributed top-1 in place of the logits all-gather (up to 8 rows), and the draft layer's MLA prologue (up to 16 rows); each kernel's other row counts run SGLang's code |

With the megakernel on, all six slots run; none needs to be listed. `SGLANG_NVIDIA_MEGAKERNEL_SLOTS` (comma-separated) overrides that list for tests and bisecting only; a list without a slot another needs (`dnorm` needs `qkv_a_small_m_shared` and `hops`, `qkv_a_small_m_shared` needs `attention`) stops the server at startup, naming it. Slot `draft` needs embedding replication, which the flag turns on (SGLang sets `SGLANG_ENABLE_EMBED_REPLICATION` when it is unset), greedy draft acceptance and an SGLang whose EAGLE draft graph runners run a draft spec's capture check; the server stops at startup, naming each one that is missing (an explicit `SGLANG_ENABLE_EMBED_REPLICATION=0` too).

### Design

- **Built at load.** Buffers, kernels and tactics for every covered width are allocated and compiled when the model is loaded, before the KV cache is sized. Nothing is compiled or allocated inside a step.
Expand All @@ -85,13 +107,14 @@ The sampling and acceptance after the verify, and the GatedDeltaNet extend of pr

### Requirements

- An SGLang with the megakernels framework (`sglang.kernels.plugin_kernels.nvidia`, or `sglang.srt.megakernels` before sgl-project/sglang#45389; `--enable-nvidia-megakernel`); Qwen3.5 also needs `sglang.srt.kernel_plugins.nvidia.eagle_binding`
- An SGLang with the megakernels framework (`sglang.kernels.plugin_kernels.nvidia`, or `sglang.srt.megakernels` before sgl-project/sglang#45389; `--enable-nvidia-megakernel`); Qwen3.5 also needs `sglang.srt.kernel_plugins.nvidia.eagle_binding`; GLM-5.3 also needs its DSA binding view (`ForwardView.dsa` in `sglang/srt/kernel_plugins/nvidia/binding.py`, the SGLang PR filed with this model), and refuses at startup without it; with that PR the flag also turns on the embedding replication slot `draft` needs (`FLAG_SETS_ENV` in `sglang/kernels/plugin_kernels/nvidia/spec.py`)
- SGLang's CUDA 13 image: FlashInfer (its low-latency block-scaled GEMM as in 0.7.0 to 0.7.2), CuTe DSL (`nvidia-cutlass-dsl`) and the CUDA Python bindings
- DeepSeek-V4.1-Flash: 4 x GB300
- Qwen3.5: 8 x GB300, two trays of one NVLink domain
- Kimi K3: 8 x GB300 (SM103), BF16, context length 8192, TP8 and DP1,
FP8 E4M3 KV cache with page size 64, `trtllm_mla` for prefill and decode,
`flashinfer_mxfp4` with default precision, and full decode CUDA graphs
- GLM-5.3, GLM-5.2: 8 x B300, or 8 GPUs over two trays of one NVLink domain (GB300 NVL72, Rubin VR NVL72)

### Installation

Expand Down Expand Up @@ -139,6 +162,30 @@ python -m sglang.launch_server \
--enable-nvidia-megakernel
```

GLM-5.3 or GLM-5.2 (NVFP4) on 8 x B300, with an NVFP4 GLM-5.3 or GLM-5.2 checkpoint (for example `nvidia/GLM-5.2-NVFP4`) whose config has the shapes in `glm5_3/__init__.py` (`HF_CONFIG`; the spec checks them at startup):

```bash
python -m sglang.launch_server \
--model-path <NVFP4 GLM-5.3 or GLM-5.2 checkpoint> \
--trust-remote-code \
--tp-size 8 \
--kv-cache-dtype fp8_e4m3 --page-size 64 \
--attention-backend dsa --dsa-decode-backend trtllm \
--moe-runner-backend flashinfer_trtllm \
--disable-shared-experts-fusion \
--speculative-algorithm EAGLE --speculative-num-steps 3 \
--speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \
--enable-nvidia-megakernel
```

Every GLM-5.3 slot runs with `--enable-nvidia-megakernel` alone; no slot list and no environment variable are needed. The GLM-5.3 NextN draft step (slot `draft`) needs embedding replication, which the flag turns on: SGLang sets `SGLANG_ENABLE_EMBED_REPLICATION` when it is unset and logs it, and keeps a value you set (`0` stops the server at startup, naming it). The draft step also needs greedy draft acceptance (no `--speculative-use-rejection-sampling`, SGLang's default); without it the server stops at startup, naming it. SGLang's defaults give the rest of the environment the spec needs; a setup it does not serve stops at startup, naming each reason. Serving with HiCache (`--enable-hierarchical-cache`) is supported: the covered steps read only the running requests' device KV pages.

On two trays (GB300 NVL72, Rubin VR NVL72), launch the same command on each tray with `--nnodes 2 --node-rank <0|1> --dist-init-addr <tray-0 host>:<port>`. The covered steps need SGLang's FlashInfer all-reduce + RMSNorm fusion at verify sizes, which needs NVLink multicast; on nodes without it, set `SGLANG_NVIDIA_MEGAKERNEL_GLM_ALLOW_NO_AR_FUSION=1`: the covered steps then run anyway, the stock residual steps on SGLang's unfused all-reduce and slot `hops` on its own symmetric-memory all-reduce.

The GLM-5.3 megakernel serves MTP with 3 draft steps (`--speculative-num-steps 3 --speculative-num-draft-tokens 4`, as above) or 5 (`--speculative-num-steps 5 --speculative-num-draft-tokens 6`); any other draft length stops the server at startup, naming the served values.

The GLM-5.3 megakernel covers verify steps of up to 16 rows: 1 to 4 requests with 4 draft tokens, 1 or 2 with 6; larger steps run SGLang's path while its derived weight copies (about 20 GiB per GPU, logged at startup) still take KV-cache memory, so it helps little on high-concurrency servers, and the server warns when `--max-running-requests` is set above 8.

## Repository Structure

```text
Expand All @@ -157,6 +204,7 @@ sgl-nvidia-mega/
│ │ ├── logits/ # the head and the logits
│ │ └── sm103/ # kernels and the build-time runner; no SGLang imports
│ ├── qwen3_5/ # spec, plan, runtime, boundary, GDN, attention, MoE and SM103 kernels
│ ├── glm5_3/ # GLM-5.3 and GLM-5.2: the same layout (spec, plan, runtime, boundary, attention, moe, logits, sm103)
│ └── kimi_k3/ # the Kimi K3 spec, adapter, executor and CuTe DSL kernels
├── benchmark/ # the static decode table behind the speed numbers
└── test/
Expand All @@ -175,7 +223,8 @@ pytest test/e2e # DeepSeek-V4.1-Flash: 4 x GB300, with SGLANG_NVI

Qwen3.5 runs TP8, which spans two GB300 trays: its e2e tests read a server launched on them through `SGLANG_NVIDIA_MEGAKERNEL_E2E_URL` and the files its ranks log to (`SGLANG_NVIDIA_MEGAKERNEL_E2E_LOGS`, comma-separated).

Kimi K3 GPU and end-to-end suites are not yet part of these commands.
GLM-5.3 and Kimi K3 GPU and end-to-end suites are not yet part of these commands.
The GLM-5.3 CPU config and SGLang contract checks are included in `test/unit`.

## Contributing

Expand Down
41 changes: 41 additions & 0 deletions sgl_nvidia_mega/glm5_3/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
"""GLM-5.3 (NVFP4) megakernel: its spec, the setups it serves (``spec``), the step
groups per width (``plan``), the build and the executor of a covered MTP verify step
(``runtime``), the step code per package (``attention``, ``boundary``, ``moe``,
``logits``), and the CuTe DSL kernels (``sm103``; they also serve SM107, Rubin).

GLM-5.2 has the same architecture and the same value of every config field the specs
pin (``HF_CONFIG``), so they serve it as well."""

MODEL_TYPE = "glm_moe_dsa"
# One tensor-parallel group of 8 ranks; the hidden size the kernels are built for
TP_SIZE = 8
HIDDEN = 6144

# GLM-5.3's shapes and routing, which the kernels are built for (GLM-5.2's are the same)
HF_CONFIG = dict(
hidden_size=HIDDEN,
num_attention_heads=64,
q_lora_rank=2048,
kv_lora_rank=512,
qk_rope_head_dim=64,
qk_nope_head_dim=192,
v_head_dim=256,
index_n_heads=32,
index_head_dim=128,
index_topk=2048,
n_routed_experts=256,
num_experts_per_tok=8,
n_shared_experts=1,
# the routing the o_proj hop's router and the MoE engine build in: sigmoid scores
# with the correction bias, one expert group, the top-8 weights renormalized
n_group=1,
topk_group=1,
norm_topk_prob=True,
scoring_func="sigmoid",
topk_method="noaux_tc",
moe_intermediate_size=2048,
first_k_dense_replace=3,
num_hidden_layers=78,
)
5 changes: 5 additions & 0 deletions sgl_nvidia_mega/glm5_3/attention/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
"""The verify attention of a GLM-5.3 layer: the stock step (``step``), which
reproduces SGLang's attention bit for bit, with the host-only facts the steps
need, and the kernels of slot ``attention`` (``fused``)."""
Loading
Loading