diff --git a/tensorrt_llm/_torch/speculative/sa_worker.py b/tensorrt_llm/_torch/speculative/sa_worker.py index 07023fe6d661..9f97c8d69404 100644 --- a/tensorrt_llm/_torch/speculative/sa_worker.py +++ b/tensorrt_llm/_torch/speculative/sa_worker.py @@ -31,6 +31,7 @@ from tensorrt_llm._utils import prefer_pinned +from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.sampler import TorchSampler from .interface import SpecMetadata, SpecWorkerBase from .spec_sampler_base import SampleStateSpec, SpecSamplerBase @@ -180,6 +181,20 @@ def _forward_impl( input_ids, logits, spec_metadata, attn_metadata ) + # Hybrid (SSM/recurrent) models: promote the accepted step's + # recurrent state from the verification scratch buffers into the + # live pools — verification never writes the pools in place (a + # rejected draft would corrupt them). Same call site as the other + # one-engine workers (dflash/eagle3); no-op for pure-attention + # models via the isinstance gate. + num_gens = batch_size - num_contexts + if num_gens > 0 and isinstance(attn_metadata.kv_cache_manager, MambaHybridCacheManager): + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + # Step 3-4: Extend SA and generate next draft tokens using GPU kernel next_draft_tokens = self._generate_draft_tokens( accepted_tokens, num_accepted_tokens, spec_metadata, batch_size, num_contexts diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 133a34e15300..fc91474f07a0 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -50,6 +50,7 @@ l0_h100: - unittest/_torch/sampler -k "not test_speculative_d2h_parity_real_predictor" - unittest/_torch/speculative/test_eagle3.py - unittest/_torch/speculative/test_rejection_buffers_guard.py + - unittest/_torch/speculative/test_sa_hybrid_state_promotion.py - unittest/_torch/speculative/hw_agnostic - unittest/_torch/thop/parallel - unittest/_torch/thop/parallel_hw_agnostic diff --git a/tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py b/tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py new file mode 100644 index 000000000000..17e964d45fbf --- /dev/null +++ b/tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for SAWorker promoting accepted recurrent states on +hybrid (SSM) models. + +Verification writes per-step recurrent states to the cache manager's +speculative scratch buffers, never the live pools; the worker must promote +the accepted step via ``update_mamba_states`` after acceptance (mirroring +the dflash/eagle3 one-engine workers), or hybrid models silently corrupt +their recurrent state under standalone SA speculative decoding. +""" + +from types import SimpleNamespace +from unittest import mock + +import pytest + +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MixedMambaHybridCacheManager +from tensorrt_llm._torch.speculative.sa_worker import SAWorker + + +def _make_worker() -> SAWorker: + spec_config = SimpleNamespace(max_draft_len=4, max_matching_ngram_size=2) + worker = SAWorker(spec_config) + # Stub out everything around the state-promotion call site; the mocked + # sampler return values flow into update_mamba_states unchanged. + worker._execute_guided_decoder_if_present = mock.MagicMock() + worker._sample_and_accept_draft_tokens = mock.MagicMock( + return_value=(mock.sentinel.accepted_tokens, mock.sentinel.num_accepted_tokens) + ) + worker._generate_draft_tokens = mock.MagicMock(return_value=mock.sentinel.next_draft_tokens) + worker._prepare_next_new_tokens = mock.MagicMock(return_value=mock.sentinel.next_new_tokens) + return worker + + +def _make_metadata(kv_cache_manager, num_seqs: int, num_contexts: int): + attn_metadata = SimpleNamespace( + num_seqs=num_seqs, + num_contexts=num_contexts, + kv_cache_manager=kv_cache_manager, + mamba_metadata=SimpleNamespace(state_indices=mock.sentinel.state_indices), + ) + spec_metadata = SimpleNamespace( + runtime_draft_len=4, + batch_indices_cuda=mock.sentinel.batch_indices_cuda, + ) + return attn_metadata, spec_metadata + + +def _run_forward(worker, attn_metadata, spec_metadata): + return worker._forward_impl( + input_ids=mock.sentinel.input_ids, + position_ids=mock.sentinel.position_ids, + hidden_states=mock.sentinel.hidden_states, + logits=mock.sentinel.logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + ) + + +def test_hybrid_manager_promotes_accepted_states(): + worker = _make_worker() + manager = mock.MagicMock(spec=MixedMambaHybridCacheManager) + attn_metadata, spec_metadata = _make_metadata(manager, num_seqs=2, num_contexts=0) + + result = _run_forward(worker, attn_metadata, spec_metadata) + + manager.update_mamba_states.assert_called_once_with( + attn_metadata=attn_metadata, + num_accepted_tokens=mock.sentinel.num_accepted_tokens, + state_indices=mock.sentinel.state_indices, + ) + assert result["new_tokens"] is mock.sentinel.accepted_tokens + + +def test_hybrid_manager_context_only_batch_skips_promotion(): + worker = _make_worker() + manager = mock.MagicMock(spec=MixedMambaHybridCacheManager) + attn_metadata, spec_metadata = _make_metadata(manager, num_seqs=2, num_contexts=2) + + _run_forward(worker, attn_metadata, spec_metadata) + + manager.update_mamba_states.assert_not_called() + + +def test_pure_attention_manager_skips_promotion(): + worker = _make_worker() + manager = mock.MagicMock() # not a MambaHybridCacheManager + attn_metadata, spec_metadata = _make_metadata(manager, num_seqs=2, num_contexts=0) + + _run_forward(worker, attn_metadata, spec_metadata) + + manager.update_mamba_states.assert_not_called() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])