Enable index_share_for_mtp_iteration in native MTP EagleProposer.

[fanxingran]

Wire GLM-5.2 draft index_share_for_mtp_iteration like vLLM: step 0 runs the MTP indexer, steps 1+ reuse sparse_kv_indices_buffer via skip_topk and compact_topk_indices. Gated on method=mtp, DSA index_topk, and the config flag so other draft backends are unchanged.

Motivation

GLM-5.2-MXFP4 sets index_share_for_mtp_iteration: true in config.json. Within one MTP draft iteration, step 0 should run the DSA indexer once; steps 1+ should reuse the same top-k indices from sparse_kv_indices_buffer instead of re-scoring the full KV cache.

vLLM already implements this in the MTP proposer. ATOM native openai_server did not, so each of the 3 draft steps ran a full indexer pass. This PR aligns ATOM draft behavior with vLLM and the model config (correctness / parity), not only a micro-optimization.

Scope: draft propose only. Target verify is unchanged (long-ISL gap vs vLLM is mostly verify + acceptance; separate follow-up).

Goals:

• Honor index_share_for_mtp_iteration for GLM-5.2 MTP (n=3)
• Reduce redundant draft indexer work (3× → 1× per propose)
• No weight, algorithm, or verify-path changes; other draft backends unchanged

---

Technical Details

Changed files

File	Change
atom/models/deepseek_mtp.py	Add set_skip_topk / compact_topk_indices on DeepSeekMultiTokenPredictor; delegate from DeepSeekMTP
atom/spec_decode/eagle.py	Gate + drive skip/compact in EagleProposer.propose()
tests/test_mtp_index_share.py	Unit tests for helpers and gate logic

deepseek_mtp.py

• DeepSeekMultiTokenPredictor.set_skip_topk(skip) — sets self_attn.skip_topk on MTP layers with an indexer
• DeepSeekMultiTokenPredictor.compact_topk_indices(slot_ids) — copies top-k rows from last_token_indices into the front of sparse_kv_indices_buffer
• DeepSeekMTP delegates both methods

eagle.py

In __init__, detect _share_mtp_indices and log when enabled.

In propose() loop:

• Step 0, before forward: set_skip_topk(False)
• Step 0, after forward: set_skip_topk(True) + compact_topk_indices(last_token_indices)
• Steps 1+: indexer skipped via skip_topk=True

Enable gates (all required)

1. speculative_config.method == "mtp"
2. draft_hf_config.index_share_for_mtp_iteration == True
3. DSA model (index_topk in config)
4. DeepSeekMultiTokenPredictor.set_skip_topk exists

Eagle3, Qwen MTP, DSpark, and models without the flag are unaffected.

---

Test Plan

Unit tests

python -m pytest tests/test_mtp_index_share.py -v

Expected: 8 passed.

Server smoke test

Start the GLM-5.2 MXFP4 MTP server per recipes/GLM-5.md (MI355 FP4 MTP). Confirm log line (per TP rank):

MTP draft index_share_for_mtp_iteration enabled: step 0 computes indexer top-k, steps 1+ reuse the buffer.

Benchmark (fair A/B methodology)

Use atom.benchmarks.benchmark_serving with fixed lengths (random-range-ratio=1.0), 160 prompts, closed-loop (request-rate=inf), --ignore-eos, OSL=1024, CONC=16 — same as internal GLM-5.2 MTP quick bench.

Important: Compare pre-PR vs post-PR on the same machine/session, with a clean MTP server and a valid pre-PR baseline (main without this patch). An earlier internal run used a contaminated 10240 baseline (~633 out tok/s); that number is invalid and must not be used to evaluate this PR.

Suggested points: ISL=1024, 8192, 10240.

Acceptance: accept_length / acceptance_rate within ~1pp of baseline.

---

Test Result

Item	Value
ATOM version	0.1.6rc1.dev161+g00266a01f (editable install with this PR)
Model	GLM-5.2-MXFP4, MTP n=3, TP=4, port 8055
P0 log	MTP draft index_share_for_mtp_iteration enabled (×4 TP ranks)

Unit tests: 8/8 passed

Fair A/B (same container, 2026-07-24)

Pre-PR = main without this patch (083220 / 083420 / 083658). Post-PR = same session after patch + server restart (090112 / 090310 / 090546).

ISL	out tok/s (before → after)	Δ	TPOT mean (before → after)	Δ	accept_len	accept_rate
1024	1483.5 → 1555.0	+4.8%	10.11 ms → 9.69 ms	−4.2%	2.97 → 3.05	65.7% → 68.3%
8192	1150.3 → 1156.2	+0.5%	12.74 ms → 12.82 ms	+0.7%	3.04 → 3.07	68.1% → 69.0%
10240	1065.4 → 1069.1	+0.3%	13.77 ms → 13.65 ms	−1.2%	3.07 → 3.10	68.9% → 70.1%

Interpretation

• End-to-end gains at 8192/10240 are ~0–1% on throughput / TPOT — expected because verify dominates E2E; P0 only removes 2 of 3 draft indexer passes per propose.
• 1024 shows larger variance (+4.8% out); treat 8192/10240 as the stable read for long-context draft share.
• No acceptance regression vs valid baseline.
• vs vLLM at 10240 (1521 out tok/s), gap remains (+42%); this PR does not claim to close that gap.

Retracted claim: Pre-PR 10240 633 out tok/s / 23 ms TPOT and post-PR +70% were driven by a wrong baseline (stale server / environment), not by this patch alone. Do not use those numbers in review.

Internal artifacts (optional for reviewers): glm5.2/mxfp4/p0_draft_index_share/bench_dev161/ + P0_压测结果_dev161.md on our fork bench branch.

---

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

---

Suggested follow-up comment (paste on the PR thread)

Benchmark update (fair A/B): We re-ran pre/post on the same machine with a valid pre-PR baseline (~1065 out tok/s @ ISL=10240, not 633). End-to-end delta is ~0–1% at 8192/10240; the PR value is primarily vLLM/config parity on draft indexer sharing, with no accept regression. Earlier +70% @ 10240 in the description is retracted. Updated tables are in the PR body.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[fanxingran]

Hi @whx-sjtu @zejunchen-zejun @ZLkanyo009 — mind taking a look at #1682 when you get a chance?

Quick version: GLM-5.2 has index_share_for_mtp_iteration in config; vLLM already shares draft indexer top-k across MTP steps, but our native MTP path was running indexer 3× per propose. This PR wires the same skip/compact behavior — draft only, verify untouched. Three files, gated so other spec-decode backends stay the same.

Checks: test_mtp_index_share.py 8/8, server log shows index_share_for_mtp_iteration enabled.

Fair A/B on the same box (valid pre-PR baseline): 1024/8192/10240 end-to-end is roughly flat (~0–3%), accept looks fine. Main win here is parity with vLLM + config, not a big bench bump. (I’ve dropped the old 10240 +70% numbers — bad baseline.)

Details + tables are in the PR description. Happy to rerun with recipe/CI settings if you want. Thanks!

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1682 --add-label <label>

[fanxingran]

@junhaha666 @zejunchen-zejun — friendly ping on #1682 🙏

Same story as above, shorter: MTP draft indexer sharing to match vLLM when the flag is on. Unit tests pass, smoke log OK.

I also refreshed the benchmark write-up — fair pre/post on one machine; ~0–3% at 1024/8192/10240, no accept regression. Please ignore my earlier ping about +70% @ 10240; PR body is updated.

Let me know if anything’s unclear or you want a different bench setup. Thanks!