FROZEN_KV_MTP slower than no-spec on Gemma 4 (SWA hybrid) under concurrency

[midakuyo]

Environment

• sm_80 (CMP 170HX, GA100 64GB) · triton attention backend (forced for Gemma 4)
• Model: lokeshe09/gemma-4-31B-it-INT8 (compressed-tensors W8A8)
• Draft: google/gemma-4-31B-it-assistant (official, BF16)
• --page-size 64 · lmsysorg/sglang:latest (v0.5.18)

Symptom

MTP (NEXTN → FROZEN_KV_MTP) is roughly on par with no-spec for a single request,
but drops to less than half of no-spec under concurrency. (For reference, vLLM's
MTP with the same model, draft, and batch 8 gave +77%.)

config	solo tg	batch-8 warm tg (per-request median)
no-spec	34.3	25.5
MTP s1/d2	—	9.4
MTP s2/d3	34.9	11.5
MTP s3/d4	—	12.2

(~2.4–3.3k-token contexts, ~160-token generations, sustained load with responses
appended to history each round. Full raw logs in the gist below.)

Observation 1 — draft dominates scheduler.run_batch

Captured a 6s torch profile (/start_profile) mid-load at batch 8 and aggregated
user_annotation spans (34–36 iterations sampled):

steps	run_batch mean	draft mean (share)	TARGET_VERIFY
1	173.4ms	165.8ms (96%)	0.5ms
2	164.4ms	155.9ms (95%)	0.4ms
3	153.8ms	144.5ms (94%)	0.3ms

GPU utilization stayed at 97–99% against the 175W power cap during load
(1s sampling) — this looks like real GPU work, not a host-side stall. Verify is
under 1ms, essentially free.

Observation 2 — iteration time is depth-independent; tg tracks accept length

steps/draft	accept len (server logs)	batch-8 tg	iteration time (profiler, direct)
1/2	1.65–1.68	9.4	173ms
2/3	1.95–2.15	11.5	164ms
3/4	2.23–2.52	12.2	154ms

Acceptance rates are in a normal range per depth (0.42–0.68), so this doesn't
look like a draft-quality problem. Tripling the depth barely changes iteration
time, which suggests a fixed once-per-iteration cost inside the draft phase
rather than the draft loop itself (tg / accept-len = iteration rate is constant
at 5.2–5.6 iter/s).

Observation 3 — the cost scales with batch × min(context, SWA window)

Batch 8 fixed, context sweep (warm tg medians):

context (tokens)	~220	~670	~1270	~1870	~2460
tg	21.5	14.2	13.4	12.9	12.2

It degrades steeply up to around the window size (1024) and is nearly flat
beyond it. Batch 4 at 2k→4k is also near-flat (21→18.5), so it doesn't look
like a read of the full context. With context fixed at 2048, batch 1/2/4/8
gives 34 / 29.5 / 20.5 / 12.3 — degradation proportional to batch size.

Taken together, this looks like a cost that reassembles the window-sized
sliding-window KV/index state on every iteration. Since the window only
slides by the number of accepted tokens per iteration, an incremental update
seems possible. I did check the two related flags:
--speculative-attention-mode decode made no difference, and
--speculative-draft-window-size is documented as ignored on this path.

Question

Is this a known cost? (#29858 seems to be reworking the same buffers from the
correctness side, and the #23705 adaptive-spec roadmap looks like a mitigation
for it.)

---

This is my first question here. To be honest, GitHub is a bit of an intimidating
place for me, so I wanted to ask in Discussions before filing an issue. I did
look through the Gemma 4 related roadmap, but I may well have misunderstood
parts of it (this is my first time building an AI serving setup).

I had AI assistance during the investigation and testing. That said, every
number comes from saved server logs and profiler traces, and the repro scripts
plus raw logs are in this gist:
https://gist.github.com/midakuyo/4dbe0cbb009d3718b8f1a41f150a5131

I hope this doesn't read as slop. Thank you for reading.

[midakuyo]

#35592 #36372
It seems they may be handling a related issue