From 791901f68423061151659458d200148df1d80986 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 19 Jul 2026 08:44:18 -0500 Subject: [PATCH 1/2] [https://nvbugs/6442074][fix] DSparkWorker: rename forward to _forward_impl d97397437f added __init_subclass__ to SpecWorkerBase requiring subclasses to implement _forward_impl instead of overriding forward directly, but did not update DSparkWorker. Rename the method to satisfy the contract. Signed-off-by: Brian Nguyen --- tensorrt_llm/_torch/speculative/dspark.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index f2597f5e9ef3..80425833b0be 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -404,7 +404,7 @@ def _draft_gen_block_batched( ) return block_logits - def forward( + def _forward_impl( self, input_ids, position_ids, From 0dbe52f5263d37fe21f76fd65129895d45f401f4 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Sun, 19 Jul 2026 21:59:55 +0000 Subject: [PATCH 2/2] [https://nvbugs/6442074][fix] Give the rejection-buffers-guard test stub a _forward_impl d97397437f made SpecWorkerBase._forward_impl abstract; the _Worker stub in test_rejection_buffers_guard.py (added after that PR branched) only stubs max_draft_len, so instantiating it now raises TypeError. Same fix as _StubSpecWorker in test_force_accepted_tokens.py. Swept current main for other gaps: dspark.py (fixed by this PR) and this stub are the only two SpecWorkerBase violations. Co-Authored-By: Claude Fable 5 Signed-off-by: Dongfeng Yu --- .../_torch/speculative/test_rejection_buffers_guard.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py index e2cbe8752d7a..a71b9f1e6640 100644 --- 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(