[Spec Decode][Perf] Replicate DSpark Markov head across TP ranks - #49731
Merged
mgoin merged 3 commits intoJul 29, 2026
Conversation
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: mgoin <mgoin64@gmail.com>
mgoin
marked this pull request as ready for review
July 24, 2026 19:17
mgoin
added
speculative-decoding
ready
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
benchislett
self-requested a review
July 27, 2026 14:34
benchislett
reviewed
Jul 27, 2026
benchislett
reviewed
Jul 27, 2026
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: mgoin <mgoin64@gmail.com>
benchislett
reviewed
Jul 29, 2026
| self.update_param_tp_status() | ||
|
|
||
| def update_param_tp_status(self): | ||
| for param in self.parameters(): |
Member
Author
There was a problem hiding this comment.
This is for compat with weight reloading, see #48025. This is just mirroring LinearBase now that we have disable_tp in this class
benchislett
approved these changes
Jul 29, 2026
aoshen02
pushed a commit
to zllion/vllm
that referenced
this pull request
Aug 1, 2026
…m-project#49731) Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
linfeng-yuan
pushed a commit
to vllm-project/vllm-ascend
that referenced
this pull request
Aug 3, 2026
### What this PR does / why we need it? | Files | Upstream vLLM change | vllm-ascend adaptation | |-------|---------------------|------------------------| | `vllm_ascend/__init__.py` | [vllm#48841](vllm-project/vllm#48841) — added `from triton.experimental import gluon`, `from triton.experimental.gluon import language as gl` and `from triton.language.core import _aggregate` to `triton_utils/__init__.py`, requiring Triton 3.6+ API absent from `triton-ascend 3.2.1` | Pre-register `triton.experimental.gluon` / `.language` as `sys.modules` stubs; stub `triton.language.core._aggregate`. Gated on `os.getenv("VLLM_VERSION") != "0.26.0"`, mirroring `vllm_version_is` env var path. | | `tests/ut/patch/platform/test_patch_structured_output.py` | [vllm#49665](vllm-project/vllm#49665) — changed `VLLMValidationError` base class from `ValueError` to `VLLMClientError`, breaking `pytest.raises(ValueError)` assertions | Updated 3 assertions from `pytest.raises(ValueError, ...)` to `pytest.raises(VLLMValidationError, ...)` | | `tests/e2e/pull_request/one_card/test_guided_decoding.py` | [vllm#49665](vllm-project/vllm#49665) — same `VLLMValidationError` base-class change; on 0.26.0 the upstream validation may still raise `ValueError` | Version-gated assertion: `pytest.raises(ValueError, ...)` on 0.26.0 vs `pytest.raises(VLLMValidationError, ...)` on main | | `vllm_ascend/ops/vocab_parallel_embedding.py`<br>`vllm_ascend/_310p/ops/vocab_parallel_embedding.py` | [vllm#49731](vllm-project/vllm#49731) — added `*, disable_tp: bool = False` keyword arg to `ParallelLMHead.__init__()` and `VocabParallelEmbedding.__init__()` | Version-gated `__init__` with `vllm_version_is("0.26.0")`: else branch accepts `disable_tp` kwarg and forwards it (Ascend manages TP via `lmhead_tp_enable()`) | | `vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py`<br>`vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py` | [vllm#26593](vllm-project/vllm#26593) — `KVConnectorBase_V1.__init__()` sets `self._kv_transfer_config`, required by `requires_kv_delivery` property | Version-gated `__init__`: else branch sets `self._kv_transfer_config` on vllm main | | `vllm_ascend/ops/fused_moe/fused_moe.py` | [vllm#50089](vllm-project/vllm#50089) — added `fused_output_is_reduced` kwarg to `_maybe_reduce_shared_expert_output()` and `output_is_reduced` kwarg to `_maybe_reduce_final_output()` | Version-gated both methods: else branch accepts new kwargs but ignores them (Ascend handles reduction independently); `trunc_size` may now be `None` | | `vllm_ascend/patch/worker/patch_distributed.py` | [vllm#47288](vllm-project/vllm#47288) — added `use_all2all: bool = False` param to `GroupCoordinator.__init__()` | Version-gated via `_IS_VLLM_026` env var check (avoids vllm-ascend import cycle). Else branch stores `self.use_all2all` (unused, no all2all on Ascend) | | `vllm_ascend/distributed/device_communicators/npu_communicator.py` | [vllm#47288](vllm-project/vllm#47288) — added `use_all2all: bool = False` param to `DeviceCommunicatorBase.__init__()` | Version-gated `NPUCommunicator.__init__` with `vllm_version_is("0.26.0")`: else branch accepts and forwards `use_all2all` (NPU keeps the no-op `_NpuAll2AllManager`) | | `vllm_ascend/worker/npu_input_batch.py` | [vllm#48018](vllm-project/vllm#48018) — added `use_replayssm: bool` kwarg (Mamba replay-SSM); [vllm#40996](vllm-project/vllm#40996) — added `slot_mapping_modes` kwarg (DCP hybrid attention) | Accepts both kwargs unconditionally; stores `self.use_replayssm` and `self.slot_mapping_modes` only on main (`not vllm_version_is("0.26.0")`) for interface alignment | | `vllm_ascend/worker/v2/spec_decode/eagle/aclgraph.py` | [vllm#49364](vllm-project/vllm#49364) — renamed `skip_attn` parameter to `full_cudagraph` in cudagraph mode logic | Version-gated call: 0.26.0 keeps `skip_attn=(cg_mode != PIECEWISE)`; main passes `full_cudagraph=(cg_mode != PIECEWISE)` | | `vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py`<br>`vllm_ascend/distributed/weight_transfer/__init__.py` | [vllm#48042](vllm-project/vllm#48042) + [vllm#48981](vllm-project/vllm#48981) — replaced static `IPCTrainerSendWeightsArgs`/`trainer_send_weights` with a stateful `IPCTrainerWeightTransferEngine` driven by `WeightTransferTrainerFactory.trainer_init(...).send_weights()`; `packed` moved onto trainer init-info; per-round `update_info` slimmed | Whole module version-gated: 0.26.0 keeps the static `NPUIPCWeightTransferEngine`; main defines stateful `NPUIPCTrainerWeightTransferEngine` (`_send_unpacked` instance method, `packed` on `NPUIPCWeightTransferInitInfo`, no-arg `npu_generate_uuid()`, `is_sender`/`_all_gather_and_merge_handles`). `register_engine()` also registers the trainer engine in `WeightTransferTrainerFactory` on main only | | `examples/rl/rlhf_http_npu_ipc.py` | [vllm#48981](vllm-project/vllm#48981) — ported RL examples to `WeightTransferTrainerFactory.trainer_init(...)` + `engine.send_weights()` with `HTTPVLLMWeightSyncClient`/`ModuleSource` | Version-gated example: 0.26.0 uses static `NPUIPCWeightTransferEngine.trainer_send_weights`; main uses the stateful trainer engine path | | `tests/e2e/pull_request/one_card/test_npu_ipc_weight_transfer.py` | [vllm#48981](vllm-project/vllm#48981) — IPC weight-transfer E2E now exercised via the stateful trainer engine | Version-gated test: 0.26.0 keeps the manual `_post(start/finish)` + static `trainer_send_weights`; main drives `WeightTransferTrainerFactory.trainer_init(...).send_weights()` (engine owns the lifecycle) | | `tests/ut/distributed/weight_transfer/test_npu_ipc_engine.py` | [vllm#48981](vllm-project/vllm#48981) — trainer-side IPC APIs changed from static methods to instance methods on the stateful engine | Version-gated tests (`IS_VLLM_026`): `_send_unpacked`, `packed` placement, `npu_generate_uuid` arity, and `start/finish_weight_update` no-op behavior asserted per version | | `vllm_ascend/worker/v2/spec_decode/dflash/speculator.py` | [vllm#50000](vllm-project/vllm#50000) — `_prepare_dflash_inputs_kernel` gained `temperature`/`seeds` sampling params for probabilistic draft sampling | Version-gated `_prepare_dflash_inputs_kernel_ascend`: 0.26.0 branch keeps the old signature; main branch adds the four `temperature`/`seeds` pointers and the corresponding stores to stay aligned with upstream while keeping Ascend's own kernel | | `.github/vllm-main-verified.commit` | — | Updated verified main commit hash to `0351e9aa1fdf1a51329d1906881528dfe61fc88e` | ### Does this PR introduce _any_ user-facing change? ### How was this patch tested? - vLLM version: v0.26.0 - vLLM main: vllm-project/vllm@d02df74 --------- Signed-off-by: hfadzxy <starmoon_zhang@163.com>
pranavthakur0-0
pushed a commit
to pranavthakur0-0/vllm
that referenced
this pull request
Aug 4, 2026
…m-project#49731) Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
itej89
pushed a commit
to itej89/vllm
that referenced
this pull request
Aug 4, 2026
…m-project#49731) Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Tej Kiran <kiran.tej@amd.com>
aditi-amd
pushed a commit
to aditi-amd/vllm
that referenced
this pull request
Aug 4, 2026
…m-project#49731) Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: root <root@smci355-ccs-aus-m02-09.cs-aus.dcgpu>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Replicate the DSpark Markov embedding and projection across TP ranks, removing an all-reduce and full-vocabulary gather per draft position.
Benchmark
Qwen3-14B DSpark on B300, batch size 1, 32 input / 128 output tokens, 7 speculative tokens, dummy weights, 3 warmups and 10 measured iterations:
Testing
Searched open DSpark/Markov-head PRs; #47584 optimizes a separate draft LM-head path and does not overlap. AI assistance was used; the submitter reviewed the change and results.