[ATOM SGL] MTP Spec decode

[ZhiweiYan-96]

Motivation

Enable DeepSeek-V4 MTP speculative decoding in the ATOM SGLang plugin path.

Technical Details

• Export DeepSeek-V4 target hidden states through the SGLang wrapper so the MTP draft model can consume the pre-hc_head hidden states.
• Add a DeepseekV4ForCausalLMNextN SGLang wrapper that creates the DeepSeek-V4 MTP draft model through ATOM.
• Share the target model’s embedding and LM head weights with the MTP draft model.
• Extend the DeepSeek-V4 SGLang KV-cache bridge to support both target layers and MTP draft layers.
• Route SGLang’s DeepSeek-V4 speculative draft attention setup through the ATOM plugin path.

Test Result

DeepSeek-V4-Pro

Acceptance rate(MTP=3):

Submission Checklist

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

[Copilot]

Pull request overview

Enable DeepSeek-V4 MTP speculative decoding in the ATOM SGLang plugin path by adding an SGLang-compatible NextN draft wrapper and extending the DeepSeek-V4 proxy KV-cache/metadata bridge to support speculative + CUDA-graph workflows.

Changes:

• Add DeepseekV4ForCausalLMNextN SGLang wrapper backed by ATOM’s DeepseekV4MTP, including shared embedding/LM-head weight tying for draft runs.
• Extend DeepSeek-V4 SGLang bridge/backends to publish and reuse ATOM V4 graph/attention metadata across decode/prefill/graph capture paths, including MTP layer support.
• Patch SGLang speculative factories / CUDA-graph eligibility to route DeepSeek-V4 draft backends through ATOM plugin mode.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
atom/plugin/sglang/runtime/forward_context.py	Adds V4 graph-metadata capture slicing and more robust metadata discovery paths.
atom/plugin/sglang/models/deepseek_v4_nextn_wrapper.py	New SGLang entry wrapper that instantiates ATOM DeepseekV4MTP and shares embed/head weights for NextN draft.
atom/plugin/sglang/models/base_model_wrapper.py	Routes DeepSeek-V4 logits through LogitsProcessor consistently; exposes embed/head for weight sharing.
atom/plugin/sglang/deepseek_v4_bridge.py	Extends proxy KV binding to MTP blocks and updates decode-graph metadata construction/indexer metadata.
atom/plugin/sglang/attention_backend/deepseek_v4_backend.py	Publishes ATOM metadata into ForwardBatch, supports multi-step draft backend expectations, and supports old/new graph hooks.
atom/plugin/register.py	Monkeypatches SGLang speculative backend factories and CUDA-graph behavior to force DeepSeek-V4 through ATOM shims.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.