[atom-vllm]: enable DSpark speculative decoding for DeepSeek-V4-Flash-0731

[peizhang56]

Motivation

Closes #2136.

DeepSeek-V4-Flash-0731 runs with dspark in ATOM native mode but not through the vLLM plugin. This PR brings the plugin to parity: DSpark as a draft method via --speculative-config '{"method":"dspark"}', plus four correctness fixes on the existing V4 plugin path that block the model regardless of spec decode.

Technical Details

Model. DeepseekV4Model.forward optionally returns per-layer aux hidden states for EAGLE3/DSpark/DFlash drafts, averaged over the mHC dim; DeepseekV4ForCausalLM PCP-all-gathers them outside the compile boundary. The aux branch sits inside @support_torch_compile, so aux_hidden_state_layers is baked into the traced graph and must be set at load time: from --level 2 up, the dispatcher replays code object 0 without evaluating guards, so later mutation is silently ignored. Noted in a COMPILE BOUNDARY comment in atom/models/deepseek_v4.py.

KV cache. The draft needs a second KV group (sliding-window MLA, block 64) beside the V4 proxy's block 128; dspark_draft_kv_patch registers it lazily and exposes it for pickle via PEP-562 __getattr__. V4's SWA ring is not keyed by a vLLM block, so a prefix-cache hit reads a stale window; every hit is rolled back by max(win_with_spec, index_topk) tokens and the tail re-forwarded, converted per group since the block sizes differ.

KNOWN ISSUE: the index_topk term is empirical (512 on V4-Flash vs a 128-token window), and the same floor reproduces with prefix caching off, so the real defect is in the sparse indexer. This rollback is a workaround; happy to track it separately.

Correctness fixes (independent of DSpark, latent on the existing V4 path):

• VocabParallelEmbedding.forward used a raw F.embedding on the tp_size == 1 branch, which GPU-faults on the -1 spec placeholder async scheduling can emit. Now replicated_embedding, which already masks it; its comment says VocabParallelEmbedding routes through the masked op, true only of the TP>1 branch.
• Async scheduling drops the sampled-token sync, so pinned staging buffers can tear; gated with a depth-1 torch.cuda.Event.
• Idle DP ranks reach build() via execute_dummy_batch() with seq_len == query_len; synthesized a decode context.
• Piecewise cudagraph modes are demoted on the plugin path, a demotion vLLM skips when it compiles nothing.

Perf. Request ids reach ATOM metadata builders through a thread-local instead of a per-step block_table[:, 0] D2H copy.

Test Plan

New: test_vllm_dspark_draft_kv_patch.py, test_vllm_v4_async_scheduling_guards.py. Extended: test_vllm_026_compat.py, test_vllm_kimi_k3.py.

# 3 sglang files cannot be collected here (no sglang module), which aborts the whole dir
python -m pytest tests/plugin -q \
  --ignore=tests/plugin/test_sglang_gdn_target_verify.py \
  --ignore=tests/plugin/test_sglang_gdn_verify_batched_ssm.py \
  --ignore=tests/plugin/test_sglang_model_arch_metadata.py
black --check . && ruff check .

Run on this branch (6fef3b9d) and its base (412f5bfe) to separate regressions from pre-existing failures. Accuracy: DeepSeek-V4-Flash-0731, 8x MI355X (gfx950), TP8, GSM8K 5-shot, 200 questions, via lm_eval local-completions against a live server, with a spec-off control and a third arm at temperature 1.0 to exercise the stochastic rejection path.

Test Result

base 412f5bfe : 37 failed, 174 passed
this branch   : 36 failed, 209 passed
new failures  : none

All 36 are pre-existing and unrelated (27 sglang, 9 rtpllm), failing identically on base. black --check . clean; ruff on the touched files is 24 findings vs 33 at base.

GSM8K exact_match (strict):

arm	value	stderr
DSpark off (control), greedy	0.965	+/- 0.013
DSpark on, greedy	0.960	+/- 0.014
DSpark on, temperature 1.0	0.925	+/- 0.019

DSpark costs 0.005 against the control, inside one stderr, and both greedy arms clear the 0.94 bar the nightly DeepSeek-V4-Pro TP8 case uses. The temperature arm differs from greedy on 197/200 questions, confirming the stochastic verifier is engaged: a draft that ignored the requested distribution would return the greedy text and score identically.

Prefix-cache sweep on V4-Flash-0731: rollback 128/256/384 -> 1/6 correct, 512 -> 6/6, and 512 holds at 2K/6K/17K prefixes.

Not run: throughput benchmarks. .github/workflows/atom-vllm-test.yaml is gated behind ci:full or an approving review, so it needs a maintainer to trigger. No CI case passes --speculative-config today, so none exercises dspark; glad to add one.

Submission Checklist

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

[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 2137 --add-label <label>

[ZhangLirong-amd]

can we do this in plugin bridge? Since aux_hidden_state we configure it in model_runner in ATOM native, not in model.