[kkHuang-amd]

Motivation

At decode batch sizes a DeepSeek-V4.1-Flash layer is dozens of small launches plus two all-reduces, and on MI355X the launch gaps and reduce latency dominate the step. This PR adds a model-agnostic framework for persistent fused "mono" kernels in SGLang, and uses it to run every DeepSeek-V4.1-Flash decoder layer (and the DSpark draft stages) as two launches on gfx950.

Modifications

• One switch: --layer-backend {module,mono} (default module). With mono, decode and target-verify steps the model's spec supports run on its kernels; every other step (prefill, unsupported widths) runs the normal module path. Unsupported deployments are refused at startup with every reason listed, never silently degraded.
• Framework in python/sglang/srt/layers/mono/ (no model code): spec contract, registry, startup guards, per-step routing, peer memory, check mode and fail-stop plumbing.
• Kernels in python/sglang/kernels/ops/mono/gfx950/: a shared device library (common/, adapted from ROCm/ATOM) and the DeepSeek-V4.1 kernels (deepseek_v41/). The model glue lives next to the model in srt/models/deepseek_common/amd/mono/.
• DeepSeek-V4.1-Flash: all 40 layers as K1 (attention seam + front) and K2 (attention back + FFN seam + MoE, TP all-reduce inside the kernel), on the production KV layouts; the 8 indexer layers run SGLang's compressor between K1 and K2 and score with a persistent FP4 scorer (K2i, from ATOM's K2a); DSpark draft stages on their own runner; the per-step host sync in DSpark decode is removed. A MiniMax-M3 spec skeleton is included (refuses until its kernels land) to keep the contract honest for a second model.

How the framework fits into SGLang

flowchart TD
    A["--layer-backend mono"] --> B["config hook<br/>registry: HF architecture -> MonoSpec<br/>spec.refuse_config + feature guards"]
    B -->|any reason| R["refuse at startup<br/>(all reasons listed)"]
    B --> C["model load (normal path)"]
    C --> D["first forward: attach()<br/>spec.refuse_model -> spec.pack (weights as views)<br/>-> spec.build per width (before graph capture)"]
    D --> E["each step: runtime.begin_step(forward_batch)<br/>routing: decode / verify, width, spec.eligible"]
    E -->|not routed| M["every layer: module path"]
    E -->|routed| F["each layer i: runtime.routed(i)"]
    F -->|Route| K["model glue calls the spec's kernel<br/>same inputs / outputs as the module code"]
    F -->|None| M2["layer i: module path"]

Loading

A step is decided once for all layers, from the CUDA-graph key (step type, rows), so a captured graph and its replays always take the same path. Kernel layers and module layers can alternate inside one step because a kernel takes and returns exactly what the model's own layer code does at that point.

sequenceDiagram
    participant Model as model.forward
    participant RT as MonoRuntime
    participant Spec as model's MonoSpec / kernels
    Model->>RT: begin_step(forward_batch)
    RT-->>Model: Decision(step or reason)
    loop layer i
        Model->>RT: routed(i)
        alt Route(mode, kernel, step)
            Model->>Spec: kernel.launch...(layer inputs)
            Spec-->>Model: layer outputs
        else None
            Model->>Model: module path
        end
    end

Loading

Layout

python/sglang/srt/layers/mono/          framework (model-agnostic)
  spec.py        MonoSpec, LayerMode, Feature, MonoContext, MonoStep, LayerKernel
  registry.py    HF architecture -> "module:MONO_SPEC" (config code never imports the model)
  guards.py      --layer-backend handling, default-deny feature list, MonoRefused
  routing.py     per-step decision (decode / verify, width, rows per request)
  layer.py       MonoRuntime: attach, begin_step, routed, per-width builds
  check.py       SGLANG_MONO_CHECK=1: each routed layer vs the module path (eager)
  failstop.py, runtime/   peer memory, mailboxes, step fences, consensus, timelines
python/sglang/kernels/ops/mono/gfx950/
  common/        shared FlyDSL device helpers (one copy for every model)
  deepseek_v41/  K1 / K2 / MoE stages, K2i indexer scorer, runner
python/sglang/srt/models/deepseek_common/amd/mono/   DeepSeek-V4.1 spec + glue
python/sglang/srt/models/minimax_m3_mono/            MiniMax-M3 spec skeleton

DeepSeek-V4.1 layer on mono

flowchart LR
    P["previous layer's<br/>pending FFN post"] --> K1["K1<br/>attention mHC seam, norm,<br/>wqkv_a, q/kv norm, KV insert, wq_b"]
    K1 --> X{"indexer layer?"}
    X -->|yes| I["SGLang compressor<br/>+ K2i FP4 scorer, top-k"]
    X -->|no| K2
    I --> K2["K2<br/>sparse attention, wo_a / wo_b,<br/>TP reduce, FFN mHC seam,<br/>router, shared + routed experts,<br/>MoE TP reduce"]
    K2 --> N["pending FFN post<br/>-> next layer's K1"]

Loading

Adding another model

1. Kernels: put them under kernels/ops/mono/<arch>/<model>/ and build on common/ (a CPU test fails if a model directory redefines a shared helper).
2. Spec: write a MonoSpec next to the model (see srt/models/minimax_m3_mono/spec.py for a minimal one):
   • architectures, layers(hf_config) -> {layer id: LayerMode.WHOLE | FFN}, widths (rows per launch, built before capture), supported_tp
   • supports: the features the kernels handle (e.g. Feature.SPEC_VERIFY); anything else in use is refused
   • refuse_config(server_args, hf_config) / refuse_model(model): return reasons, empty when fine
   • pack(model, ctx): weights as views of the loaded tensors (or a one-time conversion), plus the per-layer modes
   • build(ctx, packed, width) -> LayerKernel with launch(...) and health()
   • optional: eligible(ctx, step) for per-step limits, state_bindings for per-layer cache readiness (HiCache)
3. Register the architecture in registry.py as an import path.
4. Glue in the model: call attach once the weights are final (DeepSeek-V4.1 does it on the first forward, before any capture), begin_step(forward_batch) at the top of forward, and per layer route = mono.routed(i): launch the kernel when it returns a route, otherwise run the module code unchanged.
5. Validate: CPU tests in test/registered/mono/, SGLANG_MONO_CHECK=1 with --disable-cuda-graph (logs each routed layer's difference from the module path), then GSM8K and a long-context check.

Accuracy Tests

GSM8K (1319, TP4, --layer-backend mono): 0.901 (module path 0.911); TP2 0.908; NIAH 32k/64k/120k 60/60.

Speed Tests and Profiling

AgentX (DSpark, TP4), TTT/GPU / P90 interactivity, vs the SGLang AgentX CI baseline (module path, same recipe): c1 6,019 / 406.5 -> 7,401 / 610.3, c4 8,807 / 331.6 -> 9,561 / 516.0, c8 15,362 / 282.2 -> 16,357 / 409.9, c16 28,106 / 214.6 -> 29,078 / 272.5.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #38399664249
Latest PR Test (Extra): ❌ Run #38399663724
Latest PR Test (AMD ROCm 10): ❌ Run #38399664574

[amd-pedghazi]

Hi @kkHuang-amd, thanks for sharing this. This PR and #43497 both port the DeepSeek-V4.1 mono decode kernels from vllm-project/vllm#60397 to SGLang on gfx950, so I want to link them to avoid duplicate work.

• About 88% of the kernel code in [AMD] DeepSeek-V4.1: opt-in mono decode launches on gfx950 #43497 also appears here. Almost all of these shared lines (99%) come from [ROCm][Perf][DSv4.1] Mono decode layer for DeepSeek-V4.1-Flash on MI355X vllm-project/vllm#60397, so it is the same source, ported twice.
• The SGLang side is different. [AMD] DeepSeek-V4.1: opt-in mono decode launches on gfx950 #43497 (open since Oct 10) runs the FFN launch on all 40 layers and the attention-front launch on 30 layers, in decode and DSpark verify steps of up to 48 rows. It is behind SGLANG_ROCM_MONO_DECODE / SGLANG_ROCM_MONO_K1, with startup checks and serving results on the 10-10 image. This PR adds a general framework with --layer-backend, the attention back (K2), the indexer and the DSpark draft stages.
• RFC [RFC] Monokernel framework for SGLang #43949 lists [AMD] DeepSeek-V4.1: opt-in mono decode launches on gfx950 #43497 as the DeepSeek-V4.1 step.

Could we agree on one path? One option: the framework part of this PR follows RFC #43949, #43497 moves onto it as the DeepSeek-V4.1 PR, and your K2, indexer and draft stages come on top. Happy to discuss here or offline.

[kkHuang-amd]

Thanks @amd-pedghazi, agreed that we should not carry two ports of the same kernels. Happy to converge on RFC #43949.

What we plan to change on our side:

• Move the framework out of srt/layers/mono/ into kernels/ops/mono/ with the RFC's hooks.py / spec.py shape, and keep SGLang's side to the hook calls. The DeepSeek-V4.1 glue moves next to its kernels.
• Attach in precompile_kernels_after_loading() (today it happens on the first forward), so all memory is allocated before KV sizing, and add the load_tensor hook for weights that loading rewrites.
• Add the spec fields we are missing: numerics, per-ISA devices, requires. The A/B env vars become SGLANG_TEST_MONO_*.
• Bounded waits with error words and fail-stop in the DeepSeek-V4.1 kernels, plus the fault-injection test.
• Split the PR: framework (no kernels), the shared HIP runtime, the DeepSeek-V4.1 kernels, and the DSpark draft stages / host-sync change separately, since draft layers are a non-goal of the RFC.

For DeepSeek-V4.1 itself, two options, either is fine with us:

1. [AMD] DeepSeek-V4.1: opt-in mono decode launches on gfx950 #43497 lands on the framework as the DeepSeek-V4.1 PR (FFN launch on 40 layers, attention front on 30), and our K2, the indexer layers (compressor between the launches + an FP4 scorer) and the draft stages come on top as follow-ups.
2. The whole-layer version (K1 + K2 on all 40 layers, indexer layers included) is the DeepSeek-V4.1 PR, and the FFN-only / front-only splits stay as spans for layers that cannot go whole.

It may be easiest to decide with numbers from both on the same image and recipe; we can run that.

A few things from our side that might be worth folding into the RFC:

• --layer-backend {module,mono} instead of --enable-monokernel: same single switch, but it leaves room for other layer backends later.
• A check mode (SGLANG_TEST_MONO_CHECK=1, eager only) that runs the normal path on copies of each routed layer's inputs and logs the difference per output. It has been the fastest way to localize a numerics issue.
• Per-layer state readiness in the spec (state_bindings), so a spec can support HiCache instead of refusing it.

Glad to sync offline as well to settle the split and the order.

[kevin-mii]

Thanks @kkHuang-amd, that plan lines up with RFC #43949.

For reference, the plugin shape is implemented on kevin-mii/sglang amd/monokernel-dsv41, with all 40 DeepSeek-V4.1 layers. Could you model this PR's srt/ change on it? This is everything it touches in srt/ (39 lines, 7 files):

File	What
arg_groups/fields/exec_.py	the --enable-monokernel flag
arg_groups/pipeline.py, arg_groups/resolution_hooks.py	run_hook(handle_monokernel, ...) and its whitelist entry
environ.py	SGLANG_TEST_MONO_FAULT_INJECT
models/deepseek_v4.py	mono.load_tensor, mono.attach, mono.route
models/deepseek_common/amd/deepseek_v4_hip.py	passes the plan to the layer boundary
models/deepseek_common/amd/deepseek_v4_fused_mhc.py	mono.kernel(plan, layer_id, "layer_in"): each layer runs whole

Everything else (the framework, the model glue, your FP4 indexer scorer) belongs in the plugin. srt/ changes that matter with monokernels off, like the DSpark host-sync removal and the simulated-acceptance fold, would be their own PRs.

Three runtime points beyond your list, all from the RFC's rules:

• Build every width at attach and refuse startup if a build fails. Today each width compiles on its first step, and a failed build sends that width to the normal path.
• Drop the ATOM code the DeepSeek-V4.1 path never uses: everything in srt/layers/mono/runtime/ except consensus and peer_memory (about 590 lines, including the step fence nothing calls), plus gfx950/common/device/select.py and gfx950/common/plan/check.py (about 190). The host health() / failstop path can go too: nothing calls it, and once the kernels' waits have deadlines nothing on the host needs to check.
• Compare every rank's routing plan at attach, as runtime.attach does on the branch.

One heads-up: on the branch, the RFC's hooks.py is now two modules. guards.py holds the startup refusals, capture sizes and deadlines; runtime.py holds load_tensor, attach, route and kernel. The package re-exports the five calls, so model code does from sglang.kernels.ops import mono and calls mono.route(...).

Everything else looks good.

[billishyahao]

I am adapting #43481 into either this mono framework or rfc #43949 . I didn't observe any blockers. But I think that we still need to eliminate most of the kernel related change from srt/layers/mono/runtime/ e.g. runtime/abi.py specifically. It would be good to drop it and keep KernelAbi on the kernel side. Also mono kernels are under active development so we need to consider it as experimental feature as well.

[kkHuang-amd]

Superseded by a 4-part split following the RFC #43949 review: #44405 (DSpark host sync), #44406 (monokernel framework and shared HIP runtime), #44407 (DeepSeek-V4.1 decode layers), #44408 (DSpark draft stages). Closing this one.