[heslami]

Motivation

On HIP, DSA's decode and verify CUDA graphs come in a dense and a sparse variant, and each replay picks one from the batch's longest request (DsaGraphVariants.select). A speculative verify keeps no host copy of the lengths, so the choice reads them from the device with seq_lens.max().item(). The host waits there for the draft to finish before it can prepare and launch the verify graph, and the GPU idles meanwhile.

An AMD monokernel forward runs none of SGLang's attention, so both variants of a batch size it serves run the same launches. The choice buys nothing on those steps. Skipping it takes 354 µs of work outside the launches out of a GLM-5.3 MTP cycle at TP8, most of it GPU idle while the host waits.

Modifications

• kernel_plugins/amd/spec.py: StepMetadata.ATTENTION_GRAPH_VARIANT, which a spec lists in unread_metadata when every step it serves runs the same launches in each variant.
• kernel_plugins/amd/binding.py: unread_capture_sizes(model, metadata, ...) returns the captured batch sizes whose graph key the monokernel forward serves and whose spec declares the metadata unread. shape_key now builds its key in _shape_key(mode, bs, rows) so both use one rule.
• decode_cuda_graph_runner.py: when the backend has attention graph variants and nothing else has set attention_variant_sizes, the routed sizes leave that set. Those graphs are captured once, without variants, and their replay skips the choice. This is the mechanism the NVIDIA megakernel's verify variants already use. Unrouted sizes, and every server without such a spec, capture and select as before.
• CPU test test_routed_graphs_are_captured_without_attention_variants:
  • routed sizes are captured once and resolve without calling select;
  • an unrouted size still selects;
  • negative controls: a spec that does not declare the variant unread, a model without monokernels, and a verify graph on a forward that routes only decode steps all keep every size;
  • sizes another plugin set are kept.
    It fails on main.

The plugin that declares the variant unread is a GLM-5.3 spec for sgl-amd-mono, outside this repo and not yet upstream there; its change is one line, adding StepMetadata.ATTENTION_GRAPH_VARIANT to the target spec's unread_metadata.

Accuracy Tests

Numerics are unchanged: a routed step runs the same launches whichever graph replays, and every other step selects as on main.

test/registered/unit/kernel_plugins/amd (CPU): 19 passed (18 on main). 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 context. Acceptance is simulated at 2.99. One server per arm, 6 waves, paired by prompt:

• main: the GLM-5.3 spec without the declaration;
• this PR: the same spec plus the one-line declaration.

	main	this PR	change
MTP cycle, 1 request (ms)	7.98	7.69	−3.6%
MTP cycle, 2 requests (ms)	12.29	12.01	−2.3%

One round, so there is no same-arm A/A pair on this run; 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.

Where the time outside the launches goes, per cycle: rank 0, 1 request, CPU+GPU trace (median of 36–37 cycles, profiled), wall µs:

	main	this PR
verify preparation, up to the verify graph (the host waits on the read here)	329	83
GPU idle before the verify graph	19	0
sampling and accept after the verify graph	281	195
draft glue	244	240
total outside the launches	876	522

The launches themselves are unchanged (6.85 ms of target and 0.43–0.44 ms of draft launches a cycle in both).

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; the draft wraps the model's forward;
• base: SGLang main ec3365de18, head 9f706f0b5c.

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): ❌ Run #40725958351
Latest PR Test (Extra): ❌ Run #40725958127
Latest PR Test (AMD ROCm 10): ❌ Run #40725957991