Skip to content
Merged
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
28 changes: 28 additions & 0 deletions blog/2026-10-18-sglang-vendor-kernel-runtime.md
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,34 @@ Low-latency inference offers many opportunities for optimizations tailored to sp

### AMD Low-Latency Kernels

[`sgl-amd-mono`](https://github.com/sgl-project/sgl-amd-mono) provides model-specific kernels for AMD Instinct™ MI350 Series GPUs through `--monokernel-backend amd`. Its first implementation supports **DeepSeek-V4.1-Flash with DSpark speculative decoding on MI350X and MI355X (TP2 or TP4)**. It reuses the weights, KV cache and serving scheduler managed by SGLang, and runs both the target verify step and the DSpark draft forward on its own kernels.

At low concurrency, SGLang's standard path runs each decoder layer as dozens of small kernels and two tensor-parallel all-reduces. With a draft block size of five, one request produces six rows in a target verify step. At one or two concurrent requests, one draft-and-verify cycle runs about 1,300 kernels in about 9.1 ms on MI350X. Three costs dominate at this size:

- **Launch overhead:** HIP graphs remove the CPU's launch work, but each small kernel still costs 2-4 µs on the GPU.
- **Idle boundaries:** Each kernel must finish before the next starts, so compute units idle while one kernel drains and the next ramps up, and weight loading pauses.
- **Communication stops:** Each layer waits twice for an all-reduce across GPUs, with no computation running.

**Two persistent launches per layer.** The plugin replaces this chain with mono kernels: persistent FlyDSL kernels that each run a large part of a layer in one launch, with one workgroup resident on each of the GPU's 256 compute units. **K1** applies the mHC update, computes the query and key-value projections and writes the sliding-window KV cache. **K2** runs sparse attention, the output projections, the second mHC update, the router, the shared and routed experts, and both all-reduces. On the eight layers with a compressed attention index, SGLang's compressor runs between the launches, and a third launch, **K2i**, scores the FP4 index keys for SGLang's top-k. The layer splits at attention because attention reads KV entries written by every workgroup of K1. Each of the GPU's eight XCDs has its own L2 cache, so a launch boundary is the cheapest way to make those writes visible across the GPU.

**Overlapping stages inside a launch.** Each value a stage produces is stored together with a tag for the current decode step, so the next stage, on the same GPU or another one, starts on a value as soon as its tag matches. No separate flag or reset between steps is needed. Weights that do not depend on routing load before their inputs arrive: K1's projection weights before the norm finishes, and in K2 the mHC mix weights and the shared expert's weights before the data they multiply. Routed-expert weights load once the router has chosen the experts.

**Bringing the all-reduce into the kernel.** As each output-projection tile is computed, K2 pushes it to every GPU's memory, and each GPU sums the pieces in rank order while other workgroups are still producing tiles ([push and reduce stages](https://github.com/sgl-project/sgl-amd-mono/blob/e952b221aafc05ef5adb84d2caf942f7283c1f0b/sgl_amd_mono/deepseek_v41/gfx950/stages/seam.py)). The MoE output is reduced the same way. Every cross-GPU wait has a deadline below SGLang's watchdog timeout; a missed deadline stops the launch on every GPU instead of hanging the server.

![Schematic timeline of one DeepSeek-V4.1 layer at batch size 1. The standard path runs dozens of small kernels with an idle gap at every launch boundary and two all-reduces that stop computation. With the AMD plugin, the layer runs as two FlyDSL launches, K1 and K2, where weights that do not depend on routing load early and partial sums go to peer GPUs while computation continues.](/images/blog/sglang-vendor-kernel-runtime/amd-mono-layer.svg)

*One DeepSeek-V4.1 layer at batch size 1 on SGLang's standard path and with the AMD plugin. Numbered markers pair each cost of the standard path with the change that removes it. The timeline is schematic; it does not represent measured durations.*

The supported configuration uses decode graphs. Kernels compile and buffers are allocated at model initialization, and the graph records the selected path for each covered shape. Steps wider than 48 rows (more than eight requests at block size 5), prefill and unsupported configurations use SGLang's original forward.

**AgentX results.** The comparison runs DeepSeek-V4.1-Flash on four MI355X GPUs with TP4 and DSpark block size 5. It uses the in-tree prototype of these kernels and compares it with SGLang's AgentX baseline on the same recipe. With the plugin, P90 interactivity improves by 50%, 56% and 45% at one, four and eight concurrent AgentX sessions, and total throughput per GPU improves by 23% at one session.

![Reported AgentX P90 interactivity for DeepSeek-V4.1-Flash with DSpark on four MI355X GPUs, comparing SGLang with and without the AMD plugin](/images/blog/sglang-vendor-kernel-runtime/amd-agentx-interactivity.svg)

*DeepSeek-V4.1-Flash · 4 × MI355X · TP4 · DSpark block size 5 · AgentX trace replay. "Plugin off" is SGLang's AgentX baseline; "plugin on" adds `--monokernel-backend amd`. Speculative acceptance is fixed at 3.51 tokens per verify step in both arms. P90 interactivity is the inverse of the 90th-percentile inter-token latency.*

Accuracy is reported separately. The report gives GSM8K scores of 0.911 without the plugin and 0.901 with it (1,319 questions, TP4), and 60 of 60 on needle-in-a-haystack retrieval at 32K, 64K and 120K tokens. These are the observed results; the one-point GSM8K difference does not establish an accuracy change or numerical equivalence.

### NVIDIA Low-Latency Kernels

## Plugin System: Design and Integration
Expand Down
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading