[heslami]

Motivation

An AMD monokernel spec today serves the target's decode or verify step. An MTP draft runs other steps: EAGLE's draft decode (1 token per request) and draft extend (N tokens). So a draft spec has to wrap the model's forward itself, and SGLang neither checks its graphs nor knows which of its steps are routed.

On routed steps the DSA backend also rebuilds metadata on every replay that a monokernel forward never reads: lengths, page tables, the indexer schedule and the kpool write plan. This happens on the verify step, the draft extend and every draft decode step.

The unread-metadata contract (#44601) needs every served key to be routed. Some kernels can't serve every width at every TP size; GLM-5.3-Flash's 8-row sparse attention doesn't fit LDS at TP2, for example. Those specs fall back to the model's forward inside their own, and the backend would then skip metadata that the fallback reads.

Modifications

• kernel_plugins/amd/spec.py:
  • MonoSpec.step_modes(model): the (forward mode, tokens per request) steps a draft spec serves, e.g. [("decode", 1), ("draft_extend", 4)]. A draft spec returning none keeps the normal forward. ShapeKey.forward_mode gains "draft_extend".
  • DeviceLimits.max_rows: TP size → the most rows a step is routed with, when below widths. Wider steps run the model's forward. Every served key is then a routed key.
  • BuildContext.served: the keys routed to the built forward.
  • StepMetadata.DSA_REPLAY_PLAN: lengths, indexer schedule, kpool write plan. StepMetadata.DSA_REPLAY_PAGE_TABLE: the page table, skipped only together with the plan.
• kernel_plugins/amd/binding.py:
  • served keys are collected per step mode and capped by max_rows;
  • shape_key keys DRAFT_EXTEND_V2 steps, and a draft decode replay that plans its metadata before input_ids is set;
  • the attention-backend check runs once per mode;
  • _CaptureCheck checks each graph runner's captures on their own, since a draft model has a draft decode and a draft extend runner.
• eagle_draft_cuda_graph_runner.py, eagle_draft_extend_cuda_graph_runner.py: run the model runner's capture tail hooks after each captured graph, as DecodeCudaGraphRunner does. On a draft model runner those hooks are the AMD and NVIDIA plugin capture checks; the DSpark and DFlash draft-sampler hooks sit on their own workers' runners, not these.
• dsa_backend.py, on a routed replay:
  • both metadata members unread: the single-step replay (decode, verify, draft extend) and the multi-step draft decode replay keep the captured metadata;
  • only the plan unread: a decode replay rebuilds just the page table the attention reads;
  • every other step builds everything as before.
• CPU tests in test_mono_framework.py:
  • a draft spec's step modes are routed by graph key and the target's verify is not;
  • a draft spec with no steps keeps the normal forward;
  • each draft graph runner is checked separately;
  • max_rows at TP4, and not at TP2, leaves 6-row steps to the model's forward, and their metadata is built;
  • the DSA replay keeps the captured metadata with both members unread. With only the plan unread, it rebuilds SGLang's real page table (req_to_token at each page's first token, divided by the page size, at page sizes 4 and 1); negative control: the requests in another order. It builds everything for an unread page table alone, an unserved width, or a page-table-only verify.
  • test_mono_contract.py lists the draft runners' capture_bs.

The draft spec this serves is a GLM-5.3 MTP draft spec for sgl-amd-mono, outside this repo and not yet upstream there. It declares both DSA members, as does the target spec, and builds the page tables it reads from req_to_token in its own step-table launch.

Stacked on #47101 (merge that first; this PR's own change is its last commit): both edit shape_key, and the GLM-5.3 plugin needs both.

Accuracy Tests

The page tables the plugin builds match SGLang's req_to_token[req, 64j] // 64 word for word. That is the plugin's own kernel check, outside this repo.

Accept length at real acceptance (SGLang's default sampling), TP8, 1 request, 8K, two servers per arm in ABBA order, 30 requests each, paired by prompt:

	#47101	this PR	change
accept length	3.229	3.262	+0.032 (se 0.039); A/A +0.011 and −0.013
per-position acceptance	0.885 / 0.724 / 0.574	0.888 / 0.730 / 0.594

test/registered/unit/kernel_plugins/amd (CPU): 24 passed (19 on its base). The neighbouring runner and DSA tests (test_decode_cuda_graph_runner.py, test_eagle_draft_cuda_graph_runner.py, test_dsa_backend_dp_padding.py, test_graph_routing_context.py): 44 passed, as on main.

Speed Tests and Profiling

8× MI355X, TP8, GLM-5.3 FP8 with --enable-amd-monokernel, EAGLE 3-1-4, FP8 KV, 8K. Acceptance is simulated at 2.99. One server per arm, 6 waves, paired by prompt:

• base ([Perf] [AMD] monokernels: Capture a routed step's DSA graph once, without the host read #47101): the GLM-5.3 draft wraps the model's forward;
• this PR: the draft is served by its step modes and both specs declare the DSA metadata unread.

	#47101	this PR	change
MTP cycle, 1 request (ms)	7.69	7.64	−0.7%
MTP cycle, 2 requests (ms)	12.01	11.88	−1.0%

One round, so there is no same-arm A/A pair on this run; at C=1 the change is within what one round resolves. The change is the median per-request log-ratio, requests paired by prompt; at 2 requests only each wave's last-prefilled request is used, since the first one waits out the other's prefill. The trace shows where it comes from.

Time outside the launches per cycle: rank 0, 1 request, CPU+GPU trace (median of 37 cycles, profiled), wall / GPU-busy µs:

	#47101	this PR
draft glue (draft extend and draft decode preparation, metadata)	240 / 219	190 / 168
verify preparation	83 / 65	70 / 52
sampling and accept	195 / 153	197 / 154
total outside the launches	522	461

Setup:

• image lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025 with this branch's python/sglang mounted over the image's;
• zai-org/GLM-5.3 @ ba32f8c1;
• sgl-amd-mono fa6710e with the GLM-5.3 target and MTP draft specs that are not yet upstream in sgl-amd-mono;
• base 9f706f0b5c ([Perf] [AMD] monokernels: Capture a routed step's DSA graph once, without the host read #47101), head 6fc08c61f2.

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.

---

CI States

Latest PR Test (Base): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.

Added: no verify broadcast when every rank's logits are the same

On HIP, a greedy EAGLE verify broadcasts its results from rank 0. Without it, ranks whose draft tokens differ would accept different counts and deadlock the next collective. A monokernel forward that all-gathers its logits over the TP group returns them bitwise equal on every rank, so every rank's draft and verify reach the same result on their own.

The commit makes two changes:

• A spec declares this with MonoSpec.gathered_logits. binding.gathers_logits(model, key) reports it for the routed keys only.
• On HIP with topk=1, the verify, draft-extend and draft-decode steps of the batch size may all be served by declaring specs. When they are and rejection sampling is off, eagle_sample skips the broadcast.

Every other case keeps the broadcast: rejection sampling on (SGLang's HIP default for EAGLE), sampled requests, an undeclared or unrouted step, and CUDA.

The CPU test test_mono_framework.py checks that only routed steps of a declaring spec gather their logits; a width that isn't built, an undeclared spec, and a model without monokernels don't. test/registered/unit/kernel_plugins/amd: 25 passed (24 before).

Development measurements on our GLM-5.3 monokernel stack, at TP4, 1 request, with rejection sampling off:

• the broadcast and the idle after it took 23 µs a cycle, profiled;
• real accept length +0.091 (se 0.047), and every request completed.

The sgl-amd-mono side is one line per spec, gathered_logits=True on the GLM-5.3 target and draft specs, and it follows this PR.

Added: the capture check skips draft runners whose modes the spec doesn't serve

With this PR, the EAGLE draft decode and draft extend runners run the model runner's capture hooks, so the monokernel capture check now sees their graphs.

That would break draft specs that don't declare step modes. A draft spec that wraps the model's forward serves the server's verify key, and none of its draft runners' graphs are of that mode. The check would refuse them ("none of the N graphs EAGLEDraftCudaGraphRunner captured … runs …"), so such a spec could not start with CUDA graphs on. This includes the GLM-5.3 NextN spec in sgl-amd-mono today.

_CaptureCheck now skips a runner none of whose captured graphs is of a forward mode the spec serves:

• Forward-wrapped draft specs keep working as on main.
• A runner of a served mode is still refused if none of its graphs run the monokernels. That covers the target's decode or verify runner, and the draft runners of a spec with step modes.

The CPU test test_mono_framework.py checks that the draft runners of a spec without step modes pass, and that a verify runner whose graphs aren't routed still raises. It fails without the fix. test/registered/unit/kernel_plugins/amd: 26 passed.