[None][feat] Add DeepSeek DSpark speculative decoding#15808
[None][feat] Add DeepSeek DSpark speculative decoding#15808longlee0622 wants to merge 2 commits into
Conversation
533cd80 to
c776eb5
Compare
📝 WalkthroughWalkthroughThis PR adds a new DSpark speculative decoding mode across the TensorRT-LLM ChangesDSpark Speculative Decoding
Sequence Diagram(s)sequenceDiagram
participant TargetModel as DeepseekV4DecoderLayer
participant SpecMetadata as DSparkSpecMetadata
participant DSparkWorker
participant DraftModel as DSparkDraftModel
TargetModel->>SpecMetadata: maybe_capture_hidden_states(resolved_residual)
DSparkWorker->>SpecMetadata: prepare(request_ids)
DSparkWorker->>SpecMetadata: get_hidden_states()
DSparkWorker->>DSparkWorker: write_context_windows / back-fill accepted tokens
DSparkWorker->>DraftModel: forward_batched(main_hidden, bonus_token_ids, kv_windows, slots)
DraftModel->>DraftModel: _forward_stage (attention + MoE) per stage
DraftModel->>DraftModel: forward_head (dspark_propose)
DraftModel-->>DSparkWorker: draft tokens, num_proposed
DSparkWorker-->>TargetModel: next_draft_tokens
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (2)
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py (1)
151-210: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCoverage is insufficient for confidence-based truncation in the worker.
These tests only cover slot/window bookkeeping. They never stub
draft_model.forward()/forward_batched()to returnnum_proposed < block_size, so the worker contract that should honor DSpark truncation is still untested. Please add that regression intests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py(and, if you want end-to-end coverage, follow up intests/unittest/_torch/speculative/hw_agnostic/test_dspark.py).As per path instructions, coverage here is insufficient and the feedback should name the target test files explicitly.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py` around lines 151 - 210, The current tests only verify slot/window bookkeeping and do not cover the confidence-based truncation path in the worker. Add a regression test in test_dspark_worker that stubs draft_model.forward() or forward_batched() to return num_proposed less than block_size, then assert the worker honors DSpark truncation behavior; use _lazy_init, prepare, and the draft model call path to locate the right spot. If you want broader coverage, add a matching end-to-end test in test_dspark as well.Source: Path instructions
tests/unittest/_torch/test_model_config.py (1)
173-198: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDSpark config coverage is still missing.
This regression closes the
compress_ratios=Nonegap, but none of the new DSpark validation intensorrt_llm/llmapi/llm_args.pyis covered here. Please add focused cases intests/unittest/_torch/test_model_config.pyor a newtests/unittest/llmapi/test_llm_args.pyfor same-checkpointspeculative_modelhandling and rejectingblock_size != max_draft_len, so the first signal is not the 8-GPU e2e.As per path instructions, "
tests/**: Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM. Keep feedback actionable: suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/test_model_config.py` around lines 173 - 198, Coverage is insufficient for the new DSpark validation in llm_args.py. Add focused tests in tests/unittest/_torch/test_model_config.py or a new tests/unittest/llmapi/test_llm_args.py that exercise the speculative_model same-checkpoint path and verify block_size is rejected when it differs from max_draft_len. Use the relevant entry points ModelConfig.from_pretrained and the llm_args validation logic so the failure is caught in unit tests instead of only in e2e.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/models/modeling_dspark.py`:
- Around line 367-377: The DSpark confidence path is always enabling
`with_markov` even when `build_markov_head()` returns `None`, which later makes
`DSparkConfidenceHead.forward()` assert on missing `prev_embeddings`. Update the
initialization in `DSparkModel` so `with_markov` is derived from whether
`self.markov_head` actually exists (or from `self.markov_rank > 0`), and pass
that value into `DSparkConfidenceHead` instead of hardcoding `True`. This should
keep `dspark_propose()` and `DSparkConfidenceHead` aligned when
`dspark_markov_rank=0`.
In `@tensorrt_llm/_torch/speculative/dspark.py`:
- Around line 339-347: `num_proposed` is being computed by
`draft_model.forward*()` but then discarded, so speculative scheduling still
always queues a full `block_size` block. Update the DSpark flow in `dspark.py`
to capture the second return value from both draft paths and thread it through
`next_draft_tokens` / the verifier so the scheduled draft length is truncated
according to `confidence_threshold`. Use the existing `draft_model.forward*()`,
`next_draft_tokens`, and speculative verification logic to ensure the count
affects runtime behavior, then add a regression test covering truncated draft
blocks.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 5137-5179: Move the DSpark validation into this
`DSparkDecodingConfig` handling block in `llm_args.py` so misconfigurations fail
early: explicitly reject `speculative_model=None` when `speculative_config` is
`DSparkDecodingConfig`, and validate that the resolved `block_size` matches
`max_draft_len` after loading any draft config overrides. Keep the checks close
to the existing `spec_cfg` resolution logic and use the `DSparkDecodingConfig` /
`speculative_model` / `max_draft_len` symbols so the invariant is enforced
before downstream DSpark model construction.
- Around line 2479-2484: The confidence_threshold field in llm_args.py is
currently unconstrained, so it can accept values outside the valid probability
range and cause DSpark truncation misconfiguration. Update the
confidence_threshold definition in the relevant args/model to enforce a [0, 1]
bound using Pydantic field constraints, and keep the existing default and
description intact. Use the confidence_threshold symbol to locate the field and
ensure the validation is declarative rather than handled by a custom validator.
In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py`:
- Around line 214-252: The CUDA graph test only captures the non-persistent
branch, so it misses the production write-through path in
dspark_attention_forward_batched. Update
test_batched_attention_cuda_graph_capture_replay to also exercise persist=True
with a separate cache buffer, using the same CUDA graph capture/replay pattern.
Add an assertion that verifies the expected rolling KV window mutation after
replay, alongside the existing output comparison.
In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark.py`:
- Around line 107-114: Wrap the DSpark speculative decoding section in a
try/finally so the 8-GPU LLM created in spec_llm is always shut down even if
generate() or the output extraction for
spec_out/spec_texts/spec_ids/avg_accepted fails. Keep the existing spec_config,
spec_llm.generate, and shutdown logic, but move spec_llm.shutdown() into the
finally block to prevent leaked model/process-group state.
---
Nitpick comments:
In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py`:
- Around line 151-210: The current tests only verify slot/window bookkeeping and
do not cover the confidence-based truncation path in the worker. Add a
regression test in test_dspark_worker that stubs draft_model.forward() or
forward_batched() to return num_proposed less than block_size, then assert the
worker honors DSpark truncation behavior; use _lazy_init, prepare, and the draft
model call path to locate the right spot. If you want broader coverage, add a
matching end-to-end test in test_dspark as well.
In `@tests/unittest/_torch/test_model_config.py`:
- Around line 173-198: Coverage is insufficient for the new DSpark validation in
llm_args.py. Add focused tests in tests/unittest/_torch/test_model_config.py or
a new tests/unittest/llmapi/test_llm_args.py that exercise the speculative_model
same-checkpoint path and verify block_size is rejected when it differs from
max_draft_len. Use the relevant entry points ModelConfig.from_pretrained and the
llm_args validation logic so the failure is caught in unit tests instead of only
in e2e.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: fec2a059-f520-48da-ac10-e69c8dd54c96
📒 Files selected for processing (22)
tensorrt_llm/_torch/model_config.pytensorrt_llm/_torch/models/modeling_deepseekv4.pytensorrt_llm/_torch/models/modeling_dspark.pytensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/modules/mhc/hyper_connection.pytensorrt_llm/_torch/speculative/dspark.pytensorrt_llm/_torch/speculative/dspark_attention.pytensorrt_llm/_torch/speculative/dspark_draft.pytensorrt_llm/_torch/speculative/dspark_heads.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/__init__.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_weight_schema.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.pytests/unittest/_torch/test_model_config.pytests/unittest/api_stability/references_committed/llm.yaml
| self.markov_head = build_markov_head( | ||
| markov_head_type=getattr(config, "dspark_markov_head_type", "vanilla"), | ||
| vocab_size=config.vocab_size, | ||
| markov_rank=self.markov_rank, | ||
| hidden_size=config.hidden_size, | ||
| ) | ||
| self.confidence_head = DSparkConfidenceHead( | ||
| hidden_size=config.hidden_size, | ||
| markov_rank=self.markov_rank, | ||
| with_markov=True, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Gate with_markov on whether a Markov head actually exists.
build_markov_head() can return None when markov_rank <= 0, but this still constructs DSparkConfidenceHead(..., with_markov=True). Later, dspark_propose() only passes prev_embeddings when markov_head is present, so any confidence_threshold > 0 run with dspark_markov_rank=0 trips the assert prev_embeddings is not None in DSparkConfidenceHead.forward().
Suggested fix
self.confidence_head = DSparkConfidenceHead(
hidden_size=config.hidden_size,
markov_rank=self.markov_rank,
- with_markov=True,
+ with_markov=self.markov_rank > 0,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| self.markov_head = build_markov_head( | |
| markov_head_type=getattr(config, "dspark_markov_head_type", "vanilla"), | |
| vocab_size=config.vocab_size, | |
| markov_rank=self.markov_rank, | |
| hidden_size=config.hidden_size, | |
| ) | |
| self.confidence_head = DSparkConfidenceHead( | |
| hidden_size=config.hidden_size, | |
| markov_rank=self.markov_rank, | |
| with_markov=True, | |
| ) | |
| self.confidence_head = DSparkConfidenceHead( | |
| hidden_size=config.hidden_size, | |
| markov_rank=self.markov_rank, | |
| with_markov=self.markov_rank > 0, | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/models/modeling_dspark.py` around lines 367 - 377, The
DSpark confidence path is always enabling `with_markov` even when
`build_markov_head()` returns `None`, which later makes
`DSparkConfidenceHead.forward()` assert on missing `prev_embeddings`. Update the
initialization in `DSparkModel` so `with_markov` is derived from whether
`self.markov_head` actually exists (or from `self.markov_rank > 0`), and pass
that value into `DSparkConfidenceHead` instead of hardcoding `True`. This should
keep `dspark_propose()` and `DSparkConfidenceHead` aligned when
`dspark_markov_rank=0`.
| toks, _ = draft_model.forward( | ||
| main_hidden, | ||
| bonus, | ||
| start_pos, | ||
| kv_windows=win_slice, | ||
| temperature=0.0, | ||
| confidence_threshold=conf_thr, | ||
| ) | ||
| out_tokens[i] = toks[0].to(torch.int32) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
num_proposed is dropped before the next speculative step is scheduled.
Both draft paths discard the second value returned by draft_model.forward*() and always enqueue the full block_size tokens into next_draft_tokens. That makes confidence_threshold a no-op at runtime: the confidence head computes num_proposed, but the verifier still sees K draft tokens every step, so DSpark's advertised block truncation never takes effect.
This needs to be plumbed through the speculative scheduling/verification path, not just computed in the draft model. Please also add a regression once this is wired.
Also applies to: 422-431
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/speculative/dspark.py` around lines 339 - 347,
`num_proposed` is being computed by `draft_model.forward*()` but then discarded,
so speculative scheduling still always queues a full `block_size` block. Update
the DSpark flow in `dspark.py` to capture the second return value from both
draft paths and thread it through `next_draft_tokens` / the verifier so the
scheduled draft length is truncated according to `confidence_threshold`. Use the
existing `draft_model.forward*()`, `next_draft_tokens`, and speculative
verification logic to ensure the count affects runtime behavior, then add a
regression test covering truncated draft blocks.
| confidence_threshold: float = Field( | ||
| default=0.0, | ||
| description= | ||
| "Static confidence threshold for proposal-length truncation. A draft " | ||
| "position k is dropped when sigmoid(confidence_k) < threshold. 0.0 " | ||
| "disables truncation (propose the full block).") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Bound confidence_threshold to [0, 1].
sigmoid(confidence_k) is a probability. Values outside [0, 1] silently turn DSpark truncation into “always keep” or “always drop,” which is a bad user-facing misconfiguration surface.
As per coding guidelines, "Prefer PositiveInt, NonNegativeInt, NonNegativeFloat, PositiveFloat, or Field(gt=0) for numeric constraints in Pydantic fields instead of custom validators."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/llmapi/llm_args.py` around lines 2479 - 2484, The
confidence_threshold field in llm_args.py is currently unconstrained, so it can
accept values outside the valid probability range and cause DSpark truncation
misconfiguration. Update the confidence_threshold definition in the relevant
args/model to enforce a [0, 1] bound using Pydantic field constraints, and keep
the existing default and description intact. Use the confidence_threshold symbol
to locate the field and ensure the validation is declarative rather than handled
by a custom validator.
Source: Coding guidelines
| @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") | ||
| def test_batched_attention_cuda_graph_capture_replay(): | ||
| """The batched attention captures + replays and matches eager output. | ||
|
|
||
| Proves the path is free of capture-illegal ops (host syncs, dynamic shapes). | ||
| """ | ||
| g = _make_batched_inputs(seed=0) | ||
| dev = "cuda" | ||
| G = g["G"] | ||
| start_pos = torch.tensor(g["start_positions"], dtype=torch.long, device=dev) | ||
| slots = torch.arange(G, dtype=torch.long, device=dev) | ||
| # Static input tensors the graph reads/writes. | ||
| x = g["x"].to(dev) | ||
| main_x = g["main_x"].to(dev) | ||
| cache = g["kv_cache"].to(dev) | ||
| kw = {k: (v.to(dev) if torch.is_tensor(v) else v) for k, v in _attn_kwargs(g).items()} | ||
|
|
||
| def run(persist): | ||
| return dspark_attention_forward_batched( | ||
| x, main_x, start_pos, cache, slots, persist=persist, **kw | ||
| ) | ||
|
|
||
| eager = run(persist=False) | ||
|
|
||
| # Warmup (PyTorch CUDA-graph semantics) then capture on a non-persist call so | ||
| # the comparison isn't perturbed by the window write-through. | ||
| s = torch.cuda.Stream() | ||
| s.wait_stream(torch.cuda.current_stream()) | ||
| with torch.cuda.stream(s): | ||
| for _ in range(3): | ||
| run(persist=False) | ||
| torch.cuda.current_stream().wait_stream(s) | ||
|
|
||
| graph = torch.cuda.CUDAGraph() | ||
| with torch.cuda.graph(graph): | ||
| out = run(persist=False) | ||
| graph.replay() | ||
| torch.cuda.synchronize() | ||
| torch.testing.assert_close(out, eager, rtol=2e-2, atol=2e-2) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Capture the persist=True CUDA-graph path too.
Line 248 captures run(persist=False), but the production DSpark path calls dspark_attention_forward_batched(..., persist=True) when it updates the rolling KV windows. That means tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py never exercises the indexed write-through branch that actually runs under CUDA graphs, so a capture-only regression there can ship unnoticed. Please add a replay assertion for persist=True on a dedicated cache buffer and verify the expected window-row mutation. As per path instructions, "Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py`
around lines 214 - 252, The CUDA graph test only captures the non-persistent
branch, so it misses the production write-through path in
dspark_attention_forward_batched. Update
test_batched_attention_cuda_graph_capture_replay to also exercise persist=True
with a separate cache buffer, using the same CUDA graph capture/replay pattern.
Add an assertion that verifies the expected rolling KV window mutation after
replay, alongside the existing output comparison.
Source: Path instructions
| # DSpark speculative decoding (draft = same checkpoint's mtp.* stages). | ||
| spec_config = DSparkDecodingConfig(max_draft_len=5, speculative_model=model_dir) | ||
| spec_llm = LLM(speculative_config=spec_config, **common) | ||
| spec_out = spec_llm.generate(PROMPTS, sampling) | ||
| spec_texts = [o.outputs[0].text for o in spec_out] | ||
| spec_ids = [list(o.outputs[0].token_ids) for o in spec_out] | ||
| avg_accepted = [o.avg_decoded_tokens_per_iter - 1 for o in spec_out] | ||
| spec_llm.shutdown() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Always shut down the 8-GPU LLM in a finally block.
If generate() or output extraction raises, this test exits before shutdown() and can leave the model/process group alive, which tends to poison the rest of the suite on shared CI nodes.
Suggested fix
- spec_llm = LLM(speculative_config=spec_config, **common)
- spec_out = spec_llm.generate(PROMPTS, sampling)
- spec_texts = [o.outputs[0].text for o in spec_out]
- spec_ids = [list(o.outputs[0].token_ids) for o in spec_out]
- avg_accepted = [o.avg_decoded_tokens_per_iter - 1 for o in spec_out]
- spec_llm.shutdown()
+ spec_llm = LLM(speculative_config=spec_config, **common)
+ try:
+ spec_out = spec_llm.generate(PROMPTS, sampling)
+ spec_texts = [o.outputs[0].text for o in spec_out]
+ spec_ids = [list(o.outputs[0].token_ids) for o in spec_out]
+ avg_accepted = [o.avg_decoded_tokens_per_iter - 1 for o in spec_out]
+ finally:
+ spec_llm.shutdown()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # DSpark speculative decoding (draft = same checkpoint's mtp.* stages). | |
| spec_config = DSparkDecodingConfig(max_draft_len=5, speculative_model=model_dir) | |
| spec_llm = LLM(speculative_config=spec_config, **common) | |
| spec_out = spec_llm.generate(PROMPTS, sampling) | |
| spec_texts = [o.outputs[0].text for o in spec_out] | |
| spec_ids = [list(o.outputs[0].token_ids) for o in spec_out] | |
| avg_accepted = [o.avg_decoded_tokens_per_iter - 1 for o in spec_out] | |
| spec_llm.shutdown() | |
| # DSpark speculative decoding (draft = same checkpoint's mtp.* stages). | |
| spec_config = DSparkDecodingConfig(max_draft_len=5, speculative_model=model_dir) | |
| spec_llm = LLM(speculative_config=spec_config, **common) | |
| try: | |
| spec_out = spec_llm.generate(PROMPTS, sampling) | |
| spec_texts = [o.outputs[0].text for o in spec_out] | |
| spec_ids = [list(o.outputs[0].token_ids) for o in spec_out] | |
| avg_accepted = [o.avg_decoded_tokens_per_iter - 1 for o in spec_out] | |
| finally: | |
| spec_llm.shutdown() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark.py` around lines
107 - 114, Wrap the DSpark speculative decoding section in a try/finally so the
8-GPU LLM created in spec_llm is always shut down even if generate() or the
output extraction for spec_out/spec_texts/spec_ids/avg_accepted fails. Keep the
existing spec_config, spec_llm.generate, and shutdown logic, but move
spec_llm.shutdown() into the finally block to prevent leaked model/process-group
state.
5e177ce to
c8fcd3d
Compare
|
/bot run |
|
PR_Github #57370 [ run ] triggered by Bot. Commit: |
longlee0622
added
the
api-compatible
|
PR_Github #57370 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57435 [ run ] triggered by Bot. Commit: |
|
PR_Github #57435 [ run ] completed with state
|
15169d5 to
fa6461a
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57501 [ run ] triggered by Bot. Commit: |
|
PR_Github #57501 [ run ] completed with state |
| ) | ||
| # Pad entries hold BAD_PAGE_INDEX (-1); clamp before _compute_slot_mappings | ||
| # dereferences them (mirrors base class). | ||
| self.indexer_k_cache_block_offsets.clamp_(min=0) |
There was a problem hiding this comment.
Don't really understand the comment here - in what case can they be negative?
There was a problem hiding this comment.
updated comment to include the background of doing this.
| @@ -0,0 +1,462 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
| # SPDX-License-Identifier: Apache-2.0 | |||
There was a problem hiding this comment.
I did not review the modeling part carefully, but we should make sure we are reusing components from DSV4 where possible
There was a problem hiding this comment.
Thanks for pointing this out.
DSpark already reuses the DeepSeek-V4 decoder layer, attention/MoE/mHC modules, weight loader, embeddings, and LM head. The remaining custom attention execution implements DSpark-specific captured-context rolling windows, non-causal block attention, attention sinks, inverse RoPE, and CUDA-graph-safe batched indexing, so reusing the standard DSV4 attention forward directly would change DSpark semantics.
I did find duplicated checkpoint remapping logic. In 1c951bb016, I factored the common attention/FFN key remapping, routed-MoE scale detection, and MXFP4 tensor conversion into shared DeepSeek-V4 helpers, and updated DSpark to reuse them while retaining only its stage-specific namespace and head mappings.
I also verified the refactor against the previous implementation for FP8 and MXFP4 checkpoints, including compressor fusion, shared/routed experts, MTP keys, and DSpark stage filtering.
| # Keep every MLA reachable when one-model speculative decoding | ||
| # reuses local attention layer indices in a shared registry. |
There was a problem hiding this comment.
Can we add some more elaboration on this comment?
There was a problem hiding this comment.
Sure. I expanded the comment to document the distinction and naming rule of layer ids.
| # Per-rank generation-request counts. External drafters (DSpark) run the | ||
| # draft only over generation requests, so num_seqs (which includes |
There was a problem hiding this comment.
Why do we only run on generation requests?
There was a problem hiding this comment.
As I understand it, DSpark requires the bonus token from main model and the main model hidden states generated from that token as input. prefill request doesn't have the latter.
I've updated the comment to make it clear. BTW, it seems DFlash is doing similar things for gen-only
TensorRT-LLM/tensorrt_llm/_torch/speculative/dflash.py
Lines 474 to 518 in 71d8291
8a67c7f to
24ff931
Compare
Addresses review feedback from mikeiovine on PR NVIDIA#15808: - Remove the DSPARK_FIX_WO_A_ENV gate around the wo_a dequant. The real-MLA path already proved the dequant is correct (cos 1.0 vs the reference), so always dequantize instead of keeping the historical buggy raw fp8-cast-to-bf16 path around. - DSparkDecodingConfig.markov_head_type was never consulted by the model (which read dspark_markov_head_type straight off the HF draft config); wire it through DSparkBlock/DSparkDraftModel as a user override that falls back to the draft checkpoint value, matching the resolution pattern already used for target_layer_ids/block_size/ markov_rank. - DSparkDecodingConfig.enable_confidence_head was unused; it now gates confidence-based draft-length truncation in DSparkWorker. - DSparkDecodingConfig.mask_token_id was resolved from the draft config but never consumed by the model; wire it through the same way as markov_head_type. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Addresses review feedback from mikeiovine on PR NVIDIA#15808: - Remove the DSPARK_FIX_WO_A_ENV gate around the wo_a dequant. The real-MLA path already proved the dequant is correct (cos 1.0 vs the reference), so always dequantize instead of keeping the historical buggy raw fp8-cast-to-bf16 path around. - DSparkDecodingConfig.markov_head_type was never consulted by the model (which read dspark_markov_head_type straight off the HF draft config); wire it through DSparkBlock/DSparkDraftModel as a user override that falls back to the draft checkpoint value, matching the resolution pattern already used for target_layer_ids/block_size/ markov_rank. - DSparkDecodingConfig.enable_confidence_head was unused; it now gates confidence-based draft-length truncation in DSparkWorker. - DSparkDecodingConfig.mask_token_id was resolved from the draft config but never consumed by the model; wire it through the same way as markov_head_type. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
0e6c264 to
3c296d9
Compare
|
/bot run |
|
PR_Github #59391 [ run ] triggered by Bot. Commit: |
|
PR_Github #59391 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59433 [ run ] triggered by Bot. Commit: |
|
PR_Github #59433 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59556 [ run ] triggered by Bot. Commit: |
chienchunhung
left a comment
chienchunhung
left a comment
There was a problem hiding this comment.
Curious: Is there a way to ensure spec_config.target_layer_ids matches the draft checkpoint’s dspark_target_layer_ids? An explicit override could make these shapes incompatible and fail at runtime.
Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Make the checkpoint's dspark_target_layer_ids authoritative for DSpark speculative decoding. When target_layer_ids is unset it is copied from the checkpoint verbatim (order preserved, since main_proj columns are order-dependent). An explicit override that does not exactly match the checkpoint list is now rejected during config validation instead of failing later with a projection shape mismatch (different count) or silently degrading acceptance (same count, different/reordered layers). Add unit coverage for defaulting, matching override, mismatched count, same-count-different-layers, and order-mismatch cases. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
9581c66 to
1df8a86
Compare
Does this help? 1df8a86 |
|
/bot run --disable-fail-fast |
|
PR_Github #59689 [ run ] triggered by Bot. Commit: |
|
PR_Github #59556 [ run ] completed with state |
|
PR_Github #59689 [ run ] completed with state
|
9 participants
Summary
Adds DeepSeek DSpark speculative decoding (the DeepSeek-V4-Pro-DSpark draft) to the PyTorch backend.
DSpark extends the MTP one-model family with three full DeepSeek-V4 draft blocks under the mtp.* namespace. One backbone forward proposes a block of tokens, a lightweight sequential Markov head refines the proposal, and a per-position confidence head can truncate it. Acceptance still uses standard target verification, so greedy correctness is preserved independently of draft quality.
The implementation runs end-to-end on 8×B300 as a one-engine drafter, and the generation draft path can be captured into the target CUDA graph.
What is added
CUDA-graph safety
The default generation draft path is batched and avoids host synchronization and data-dependent shapes. It uses per-request start-position tensors, RoPE gathered from a fixed table, slot-indexed rolling windows, and fixed-width masked top-k. The batched path is numerically checked against the scalar reference path.
No DSpark-specific CUDA-graph flag is required. The eager per-request loop remains only as the fallback for the experimental real-FP8 MLA path selected by TLLM_DSPARK_REAL_MLA.
Validated LLM API configuration
The new full-accuracy LLM API test uses the following topology and backend:
The draft inherits the target MoE backend. The validated high-acceptance path is DEP8 + MegaMoE DeepGEMM. CUTLASS remains a compatible fallback; TRTLLM-Gen blockScaleMoe cannot route this checkpoint's 384-expert / 8-group layout.
Correctness and accuracy
The 8-GPU MoE path is not bit-reproducible across independent runs because of non-associative FP atomics, so the end-to-end test does not require token-for-token equality with a separate no-spec run. The verification invariant is covered by unit tests and numerical goldens.
Acceptance length: DL1-DL7 sweep
Measured offline on 8×B300 with the same TP8 + EP8 + attention DP + MegaMoE DeepGEMM configuration as the LLM API test. The sweep used greedy decoding, 12 prompts per dataset / thinking mode / draft length, up to 512 new tokens, with chunked prefill and CUDA graph disabled.
AL is avg_decoded_tokens_per_iter: committed tokens per target step, including the one guaranteed target token. DL is max_draft_len, so AL is in [1, DL+1]. Each dataset cell below is Thinking-OFF / Thinking-ON.
Key observations:
All 42 reported AL values were validated against the raw log, all seven draft-length runs exited successfully, and the sweep completed successfully.
Notes and limitations
PR checklist