[https://nvbugs/6442074][fix] DSparkWorker: rename forward to _forward_impl

[brnguyen2]

d973974 added __init_subclass__ to SpecWorkerBase requiring subclasses to implement _forward_impl instead of overriding forward directly, but did not update DSparkWorker. This causes a TypeError at import time on any build that includes both that commit and the DSpark PR (#15808).

Fix: rename DSparkWorker.forward → DSparkWorker._forward_impl. Callers go through SpecWorkerBase.forward which wraps _forward_impl, so no call sites need updating.

Summary by CodeRabbit

• Refactor
  • Updated internal speculative decoding execution flow without changing observable behavior or generated outputs.

[brnguyen2]

/mbot run

[coderabbitai]

📝 Walkthrough

Walkthrough

DSparkWorker.forward(...) is renamed to DSparkWorker._forward_impl(...) in dspark.py; the method parameters, body, call flow, and returned output remain unchanged.

Changes

DSpark worker entry-point refactor

Layer / File(s)	Summary
Rename DSpark implementation entry point tensorrt_llm/_torch/speculative/dspark.py	Renames the worker method from forward to _forward_impl without modifying its implementation.

Estimated code review effort: 1 (Trivial) | ~2 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#15808: Introduces the overlapping DSpark worker implementation.

Suggested labels: api-compatible

Suggested reviewers: longlee0622, bowenfu

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title matches the change and follows the required ticket/type format.
Description check	✅ Passed	It explains the bug and fix clearly, though the Test Coverage and PR Checklist sections are missing.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[dongfengy]

/bot run --disable-fail-fast

[dongfengy]

Approved. Thanks! This seems to be a merge conflict.

[tensorrt-cicd]

PR_Github #60183 [ run ] triggered by Bot. Commit: 25ca8b9 Link to invocation

[tensorrt-cicd]

PR_Github #60183 [ run ] completed with state FAILURE. Commit: 25ca8b9
/LLM/main/L0_MergeRequest_PR pipeline #48556 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[dongfengy]

Heads-up: this PR's CI (pipeline 48556) failed on one more instance of the same problem — tests/unittest/_torch/speculative/test_rejection_buffers_guard.py (also landed after #16382 branched) has a _Worker(SpecWorkerBase) stub with no _forward_impl, so its 3 tests now die with TypeError: Can't instantiate abstract class. I swept current main: dspark.py (fixed here) and this stub are the only two violations.

Could you add this hunk to this PR? (I can't push to your branch — maintainer-edit denied.)

--- a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py
+++ b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py
@@ -201,6 +201,9 @@ class _Worker(SpecWorkerBase):
     def max_draft_len(self) -> int:
         return K
 
+    def _forward_impl(self, *args: object, **kwargs: object) -> None:
+        raise NotImplementedError
+
 
 def _dispatch_meta(**over):
     base = dict(

Note both fixes need to ride in this PR: a separate PR with only the test fix can't go green either, because its CI merge still contains the unfixed dspark import error. With both hunks in, everything else on main is clean per the sweep.

[dongfengy]

To unblock main fastest we opened #16579 combining this dspark rename (cherry-picked, your authorship preserved) with the second missing fix (_Worker test stub). Whichever lands first works — if #16579 goes in, this one can be closed. Thanks again for the quick catch!

[fredricz-20070104]

Review: LGTM — correct and clean bug fix, safe to merge

This is a minimal (1 file, 1 line) fix for an integration/merge-ordering issue. Verified against the PR branch:

1. Root cause is real — commit d97397437f added an __init_subclass__ contract to SpecWorkerBase (interface.py:1004): subclasses may no longer override forward directly and must implement _forward_impl, otherwise TypeError is raised at class-definition/import time. DSparkWorker (dspark.py:162) still defined forward, so any build containing both commits would crash on import.

2. Fix is correct and complete — renaming DSparkWorker.forward → _forward_impl satisfies the contract. The only call site (modeling_speculative.py:2033) invokes self.spec_worker(input_ids=..., ...) via nn.Module.__call__ → base forward → _forward_impl, all keyword args matching the _forward_impl signature (dspark.py:407). No call-site changes needed.

3. No new bugs introduced — DSparkWorker has no internal self.forward() / super().forward() calls, so the rename does not affect internal logic. It now follows the exact same pattern as the other workers (eagle3, mtp, pard, dflash, sa_worker), which are already on _forward_impl. The only behavioral change is that DSpark now goes through the base forward's try/finally cleanup (_ensure_spec_dec_state_restored), which is a no-op on the success path (interface.py:1044) and only adds a safety net against spec-dec state leakage on failure — exactly what nvbugs/6442074 wants.

Risk: minimal — scoped to the DSpark speculative-decoding path, MERGEABLE, no conflicts. Since any build combining both commits fails at import, the sooner this lands the better.

Recommendation: Just confirm CI passes (especially DSpark-related import / unit tests) before merging — the code itself is fine.

[fredricz-20070104]

Approved.

[chzblych]

Closed in favor of #16579.