From 04f7afcc8d65ff8e1691a8fb728da7c3f22e5fb9 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Tue, 18 Aug 2026 10:48:20 -0700 Subject: [PATCH 1/5] [TRTLLM-15035][test] Wire Kimi K3 spec-dec and suffix-automaton tests into L0 CI - Relocate tests/torch/speculative/test_suffix_automaton.py to tests/unittest/_torch/speculative/: the test-db runner only routes unittest/-prefixed entries (tests/integration/defs/conftest.py), so the old location could never be listed. Pure rename, no content changes. - List the suffix-automaton kernel tests in l0_b200.yml (pre-merge, 1 GPU). - Wire test_kimi_k3_specdec.py::test_kimi_k3_sa_specdec_logits_parity (orphaned since #17327) into l0_gb300_multi_gpus.yml (post-merge, 4 GPUs). It resolves /Kimi-K3 and skips cleanly when absent. - Wire the CPU-only test_kimi_k3_disagg_parity_selftest into l0_cpu.yml. - Fix a stale comment in l0_gb300_multi_gpus.yml (the KDA parity unit tests never load a checkpoint) and a dead sbatch filename in the kimi_k3_disagg_parity.py docstring. Signed-off-by: Michal Guzek --- tests/integration/defs/kimi_k3_disagg_parity.py | 2 +- tests/integration/test_lists/test-db/l0_b200.yml | 3 +++ tests/integration/test_lists/test-db/l0_cpu.yml | 2 ++ .../integration/test_lists/test-db/l0_gb300_multi_gpus.yml | 7 +++++-- .../_torch}/speculative/test_suffix_automaton.py | 0 5 files changed, 11 insertions(+), 3 deletions(-) rename tests/{torch => unittest/_torch}/speculative/test_suffix_automaton.py (100%) diff --git a/tests/integration/defs/kimi_k3_disagg_parity.py b/tests/integration/defs/kimi_k3_disagg_parity.py index f327899870a2..fbcb44279ada 100644 --- a/tests/integration/defs/kimi_k3_disagg_parity.py +++ b/tests/integration/defs/kimi_k3_disagg_parity.py @@ -39,7 +39,7 @@ REAL RECIPE (aggregated vs disagg proxy): # reference: aggregated DEP16 deployment (as in - # examples/kimi_k3/run_gsm8k_kimi_k3.sbatch, but served) + # examples/kimi_k3/run_eval_kimi_k3.sbatch, but served) trtllm-serve $KIMI_K3_CKPT --backend pytorch --port 8000 ... # candidate: ctx + gen workers behind the disagg proxy diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 4a217b4294c5..690f35527dfa 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -178,6 +178,9 @@ l0_b200: - unittest/_torch/speculative/test_eagle3.py - unittest/_torch/speculative/test_dspark_cute_dsl_attention.py - unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py + # On-device suffix-automaton spec-dec kernel tests (GPU-only: CUDA graph + # capture, no skip guards; relocated from tests/torch/ which test-db cannot list). + - unittest/_torch/speculative/test_suffix_automaton.py - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/visual_gen/kernels/parallel - unittest/_torch/thop/serial diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index c0071ffaf48e..c98564ad58c5 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -23,6 +23,8 @@ l0_cpu: - disaggregated/test_aiperf_gate.py::test_gate_disabled_paths_not_affected - disaggregated/test_aiperf_gate.py::test_was_cancelled_metadata_fallback - disaggregated/test_aiperf_gate.py::test_non_request_records_excluded_from_denominator + # Kimi K3 disagg parity harness self-test (comparison logic only, no GPUs). + - test_kimi_k3_specdec.py::test_kimi_k3_disagg_parity_selftest - unittest/_torch/attention - unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py - unittest/_torch/distributed diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index ddc4bf45cbe3..dc196c7c3e80 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -36,8 +36,7 @@ l0_gb300_multi_gpus: - unittest/_torch/modeling -k "modeling_llama" - unittest/_torch/modeling -k "modeling_gpt_oss" # Kimi K3 speculative decoding: exact kernel-level verify parity (random - # weights, no checkpoint) + truncated-model e2e SA logits parity - # (auto-skips when the Kimi K3 checkpoint is not staged). + # weights, no checkpoint). - unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py # Fused KDA multi-token verify (trtllm::kda_mtp_decode): kernel-level # parity vs CPU golden + FLA sequential, and runtime-level fused-vs- @@ -51,6 +50,10 @@ l0_gb300_multi_gpus: # Kimi K3 MoE routing shape on the TRTLLM backend (random weights, no # checkpoint; the case moved out of the skipped TestMoeFp4 class). - unittest/_torch/moe/test_moe_backend.py -k "kimi_k3_shape" + # Kimi K3 truncated-model (first 4 layers) e2e SA spec-dec logits parity; + # resolves /Kimi-K3 and skips cleanly when the checkpoint + # is not staged. + - test_kimi_k3_specdec.py::test_kimi_k3_sa_specdec_logits_parity TIMEOUT (40) - unittest/_torch/multi_gpu_modeling -k "deepseek" - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=FLASHINFER-torch_compile=True] diff --git a/tests/torch/speculative/test_suffix_automaton.py b/tests/unittest/_torch/speculative/test_suffix_automaton.py similarity index 100% rename from tests/torch/speculative/test_suffix_automaton.py rename to tests/unittest/_torch/speculative/test_suffix_automaton.py From 34c8214a01c97fa2a786809d1b1cff64f2071994 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Tue, 18 Aug 2026 13:53:35 -0700 Subject: [PATCH 2/5] [TRTLLM-15035][test] Address review feedback on the suffix-automaton tests - Add the NVIDIA SPDX header and a module-level CUDA skipif guard so a plain 'pytest tests/unittest/' run on a GPU-less machine skips instead of failing. - Route every SuffixAutomatonManager construction through a make_manager fixture that shuts managers down on teardown: a mid-test assertion failure no longer leaks pinned host buffers / GPU workspace into subsequent tests in the same process. - Assert results in test_manager_extend (token 6 is unseen in the context: no match, zeroed draft - the extend_ngram no-match convention) and after CUDA-graph replay in test_cuda_graph_capture (a match must exist once warmups appended prior 6s; exact values not pinned since they depend on the executed-extend count). - Add coverage for the global_pool_size < max_num_requests ValueError. - Replace the hand-maintained __main__ test registry with pytest.main([__file__]) so direct execution cannot drift. - l0_cpu-listed disagg parity selftest: a missing harness is now a hard failure instead of a silent skip. Declined review suggestions (rationale in the PR discussion): parametrizing the multi-case extend_ngram tests, and removing SAConfig.threshold (the field is consumed by the SA-enhancer product path via eagle3/pard/mtp). Signed-off-by: Michal Guzek --- .../integration/defs/test_kimi_k3_specdec.py | 8 +- .../speculative/test_suffix_automaton.py | 254 +++++++----------- 2 files changed, 106 insertions(+), 156 deletions(-) diff --git a/tests/integration/defs/test_kimi_k3_specdec.py b/tests/integration/defs/test_kimi_k3_specdec.py index df8eb2fa2151..d4170a5b543c 100644 --- a/tests/integration/defs/test_kimi_k3_specdec.py +++ b/tests/integration/defs/test_kimi_k3_specdec.py @@ -81,11 +81,9 @@ def test_kimi_k3_disagg_parity_selftest(): Comparison logic only: canned responses, no servers or GPUs. """ script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kimi_k3_disagg_parity.py") - if not os.path.exists(script): - pytest.skip( - "kimi_k3_disagg_parity.py harness not present on this branch " - "(ships with the disagg parity PR)" - ) + # Hard failure, not a skip: this test is CI-listed (l0_cpu), so a moved or + # renamed harness must surface as a regression instead of a silent skip. + assert os.path.exists(script), f"kimi_k3_disagg_parity.py harness missing at {script}" result = subprocess.run( [sys.executable, script, "--self-test"], capture_output=True, text=True, timeout=120 ) diff --git a/tests/unittest/_torch/speculative/test_suffix_automaton.py b/tests/unittest/_torch/speculative/test_suffix_automaton.py index ee5a7f03b365..74ed7a48a910 100644 --- a/tests/unittest/_torch/speculative/test_suffix_automaton.py +++ b/tests/unittest/_torch/speculative/test_suffix_automaton.py @@ -1,8 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 """Unit tests for suffix automaton speculative decoding. Tests the native CUDA kernel implementation. """ +import pytest import torch from tensorrt_llm._torch.speculative.suffix_automaton import ( @@ -10,21 +13,43 @@ SuffixAutomatonManager, ) # noqa: I001 +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="suffix automaton tests require a CUDA device" +) + + +@pytest.fixture +def make_manager(): + """Build SuffixAutomatonManager instances that are shut down on teardown. + + A mid-test assertion failure must not leak the manager's pinned host + buffers and GPU workspace into subsequent tests in the same process. + """ + managers = [] + + def _make(*args, **kwargs): + manager = SuffixAutomatonManager(*args, **kwargs) + managers.append(manager) + return manager + + yield _make + for manager in managers: + manager.shutdown() + class TestSuffixAutomatonManager: """Tests for SuffixAutomatonManager class.""" - def test_manager_creation(self): + def test_manager_creation(self, make_manager): """Test manager creation.""" config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) assert manager is not None - manager.shutdown() - def test_manager_add_remove(self): + def test_manager_add_remove(self, make_manager): """Test adding and removing requests.""" config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) # Add requests manager.add_request(0, [1, 2, 3, 4, 5]) @@ -38,12 +63,10 @@ def test_manager_add_remove(self): assert 0 not in manager._request_to_slot assert 1 in manager._request_to_slot - manager.shutdown() - - def test_manager_extend(self): + def test_manager_extend(self, make_manager): """Test extend operation.""" config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) # Add request with repeating pattern context_tokens = [1, 2, 3, 4, 5, 1, 2, 3] @@ -62,19 +85,33 @@ def test_manager_extend(self): request_ids, accepted_tokens, num_accepted_tokens, max_draft_len ) - print(f"match_len: {match_len}") - print(f"draft_tokens: {draft_tokens}") + # Token 6 never occurs in the context, so the extended sequence's + # suffix has no earlier occurrence: no match and a zeroed draft + # (same convention as the extend_ngram no-match cases below). + match_len_val = match_len[0].item() + assert match_len_val == 0, f"Expected no match for unseen token, got {match_len_val}" + draft_list = draft_tokens[0, :max_draft_len].cpu().tolist() + assert draft_list == [0, 0, 0, 0], f"Expected zeroed draft, got {draft_list}" - manager.shutdown() + def test_global_pool_size_below_max_requests_rejected(self, make_manager): + """global_pool_size smaller than max_num_requests must be rejected.""" + config = SAConfig( + max_seq_len=1024, + max_slots=16, + enable_global_pool=True, + global_pool_size=8, + ) + with pytest.raises(ValueError, match="must be >="): + make_manager(config, max_num_requests=16) class TestCUDAGraphCompatibility: """Tests for CUDA graph compatibility.""" - def test_cuda_graph_capture(self): + def test_cuda_graph_capture(self, make_manager): """Test that native extend works during CUDA graph capture.""" config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) # Add request context_tokens = [1, 2, 3, 4, 5, 1, 2, 3] @@ -103,20 +140,25 @@ def test_cuda_graph_capture(self): # Replay g.replay() - print("CUDA graph capture succeeded with native kernel") - print(f"match_len: {match_len}") - print(f"draft_tokens: {draft_tokens}") - - manager.shutdown() + # Every executed extend() appends token 6 to the SA state (the + # warmup calls ran eagerly; capture records without executing), so + # by the replay the state holds earlier 6s and the lookup must find + # a match. Exact values are not pinned because they depend on how + # many extends executed; the load-bearing checks are that capture + + # replay complete and produce sane outputs. + match_len_val = match_len[0].item() + assert match_len_val >= 1, f"Expected a match after replay, got {match_len_val}" + assert draft_tokens.shape[1] >= max_draft_len + assert draft_tokens.dtype == torch.int32 class TestExtendNgram: """Tests for extend_ngram() batched method - CUDA graph compatible.""" - def test_extend_ngram_longest_match(self): + def test_extend_ngram_longest_match(self, make_manager): """Test extend_ngram with longest match mode (max_ngram_size=-1).""" config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) # Case 1: context_tokens=[0,1,2,1,2], extend with token 1 # New sequence: [0, 1, 2, 1, 2, 1] @@ -150,14 +192,12 @@ def test_extend_ngram_longest_match(self): draft_list = draft_tokens[0, :max_draft_len].cpu().tolist() assert draft_list[0] == 2, f"Expected first draft token 2, got {draft_list}" - manager.shutdown() - # Case 2: context_tokens=[0, 1, 2, 3, 1, 2, 4, 1], extend with token 2 # New sequence: [0, 1, 2, 3, 1, 2, 4, 1, 2] # Longest suffix match in context: [1, 2] at positions 1-2 or 4-5 → match_len=2 # Kernel uses leftmost match (1-2): continuation is token at position 3 → draft [3, 1, 2, 4] config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) context_tokens = [0, 1, 2, 3, 1, 2, 4, 1] manager.add_request(0, context_tokens) @@ -184,14 +224,12 @@ def test_extend_ngram_longest_match(self): draft_list = draft_tokens[0, :max_draft_len].cpu().tolist() assert draft_list == [3, 1, 2, 4], f"Expected draft [3, 1, 2, 4], got {draft_list}" - manager.shutdown() - # Case 3: context_tokens=[1, 1, 1, 1], extend with token 1 # New sequence: [1, 1, 1, 1, 1] # Longest suffix match in context: [1, 1, 1, 1] at positions 0-3 → match_len=4 # Only one draft token (kernel yields token at position after match; rest zero-padded) config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) context_tokens = [1, 1, 1, 1] manager.add_request(0, context_tokens) @@ -218,14 +256,12 @@ def test_extend_ngram_longest_match(self): draft_list = draft_tokens[0, :max_draft_len].cpu().tolist() assert draft_list == [1, 0, 0, 0], f"Expected draft [1, 0, 0, 0], got {draft_list}" - manager.shutdown() - # Case 4: context_tokens=[0, 1, 2, 3], extend with token 2 # New sequence: [0, 1, 2, 3, 2] # Longest suffix match in context: [2] at position 2 → match_len=1 # Continuation after match: tokens at 3 and from extended seq → draft [3, 2, 0, 0] config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) context_tokens = [0, 1, 2, 3] manager.add_request(0, context_tokens) @@ -252,12 +288,10 @@ def test_extend_ngram_longest_match(self): draft_list = draft_tokens[0, :max_draft_len].cpu().tolist() assert draft_list == [3, 2, 0, 0], f"Expected draft [3, 2, 0, 0], got {draft_list}" - manager.shutdown() - - def test_extend_ngram_fixed_size(self): + def test_extend_ngram_fixed_size(self, make_manager): """Test extend_ngram with fixed-size ngram matching.""" config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) # Case 1: context_tokens=[0, 1, 2, 3, 1, 2], extend with token 3 # New sequence: [0, 1, 2, 3, 1, 2, 3] @@ -289,14 +323,12 @@ def test_extend_ngram_fixed_size(self): draft_list = draft_tokens[0, :max_draft_len].cpu().tolist() assert draft_list == [1, 2, 3, 0], f"Expected draft [1, 2, 3, 0], got {draft_list}" - manager.shutdown() - # Case 2: context_tokens=[0, 1, 2, 3, 1, 2, 4, 1], extend with token 2 # New sequence: [0, 1, 2, 3, 1, 2, 4, 1, 2] # With max_ngram_size=3: 3-gram [4, 1, 2] not in context; 2-gram [1, 2] matches at 1-2 (leftmost) # match_len=2, continuation → draft [3, 1, 2, 4] config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) context_tokens = [0, 1, 2, 3, 1, 2, 4, 1] manager.add_request(0, context_tokens) @@ -323,14 +355,12 @@ def test_extend_ngram_fixed_size(self): draft_list = draft_tokens[0, :max_draft_len].cpu().tolist() assert draft_list == [3, 1, 2, 4], f"Expected draft [3, 1, 2, 4], got {draft_list}" - manager.shutdown() - # Case 3: context_tokens=[0, 2, 3, 4, 5, 1, 2, 3, 4, 6, 1, 2, 3], extend with token 4 # New sequence: [0, 2, 3, 4, 5, 1, 2, 3, 4, 6, 1, 2, 3, 4] # With max_ngram_size=3: 3-gram [2, 3, 4] matches at positions 1-3 (leftmost) # match_len=3, continuation after 1-3 → draft [5, 1, 2, 3] config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) context_tokens = [0, 2, 3, 4, 5, 1, 2, 3, 4, 6, 1, 2, 3] manager.add_request(0, context_tokens) @@ -357,14 +387,12 @@ def test_extend_ngram_fixed_size(self): draft_list = draft_tokens[0, :max_draft_len].cpu().tolist() assert draft_list == [5, 1, 2, 3], f"Expected draft [5, 1, 2, 3], got {draft_list}" - manager.shutdown() - # Case 4: context_tokens=[1, 2, 1, 2], extend with token 1 # New sequence: [1, 2, 1, 2, 1] # With max_ngram_size=3: 3-gram [2, 1, 2] matches at positions 1-3 (leftmost) # match_len=3; continuation from match start when no token after match → draft [2, 1, 0, 0] config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) context_tokens = [1, 2, 1, 2] manager.add_request(0, context_tokens) @@ -391,15 +419,13 @@ def test_extend_ngram_fixed_size(self): draft_list = draft_tokens[0, :max_draft_len].cpu().tolist() assert draft_list == [2, 1, 0, 0], f"Expected draft [2, 1, 0, 0], got {draft_list}" - manager.shutdown() - - def test_extend_ngram_no_match(self): + def test_extend_ngram_no_match(self, make_manager): """Test extend_ngram when no match exists (moved from longest_match case 5).""" # context_tokens=[0, 1, 2, 3], extend with token 4 # New sequence: [0, 1, 2, 3, 4] # Token 4 not in context → no suffix match, match_len=0, draft is zero-padded config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) context_tokens = [0, 1, 2, 3] manager.add_request(0, context_tokens) @@ -425,11 +451,9 @@ def test_extend_ngram_no_match(self): draft_list = draft_tokens[0, :max_draft_len].cpu().tolist() assert draft_list == [0, 0, 0, 0], f"Expected draft [0, 0, 0, 0], got {draft_list}" - manager.shutdown() - # Same no-match scenario with max_ngram_size=3: context [0, 1, 2, 3], extend token 4 config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) context_tokens = [0, 1, 2, 3] manager.add_request(0, context_tokens) @@ -457,12 +481,10 @@ def test_extend_ngram_no_match(self): draft_list = draft_tokens[0, :max_draft_len].cpu().tolist() assert draft_list == [0, 0, 0, 0], f"Expected draft [0, 0, 0, 0], got {draft_list}" - manager.shutdown() - - def test_extend_ngram_batch(self): + def test_extend_ngram_batch(self, make_manager): """Test extend_ngram with multiple requests in batch.""" config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) # Add multiple requests with different patterns # Request 0: [1, 2, 3, 4, 5, 1, 2, 3] + [4] -> match [1,2,3,4] at pos 0 @@ -514,12 +536,10 @@ def test_extend_ngram_batch(self): # Request 2: no match assert match_lens[2] == 0, f"Request 2: expected match_len=0, got {match_lens[2]}" - manager.shutdown() - - def test_extend_ngram_cuda_graph(self): + def test_extend_ngram_cuda_graph(self, make_manager): """Test that extend_ngram works with CUDA graph capture.""" config = SAConfig(max_seq_len=1024, max_slots=16) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) # Add request with repeating pattern context_tokens = [1, 2, 3, 4, 5, 1, 2, 3] @@ -567,16 +587,14 @@ def test_extend_ngram_cuda_graph(self): draft_list = draft_tokens[0, :max_draft_len].cpu().tolist() assert draft_list == [5, 1, 2, 3], f"Expected [5, 1, 2, 3], got {draft_list}" - manager.shutdown() - class TestExtendGlobal: """Tests for extend_global() — cross-request pattern sharing.""" - def test_extend_global_cross_request_match(self): + def test_extend_global_cross_request_match(self, make_manager): """Request B finds a pattern from Request A's context.""" config = SAConfig(max_seq_len=1024, max_slots=16, enable_global_pool=True) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) # Request 0: context has [1, 2, 3, 4, 5] manager.add_request(0, [1, 2, 3, 4, 5]) @@ -616,12 +634,10 @@ def test_extend_global_cross_request_match(self): assert draft_1[0] == 5, f"Expected continuation starting with 5, got {draft_1}" assert draft_1[1] == 6, f"Expected second draft token 6, got {draft_1}" - manager.shutdown() - - def test_extend_global_prefers_own_slot(self): + def test_extend_global_prefers_own_slot(self, make_manager): """When match lengths are equal, prefer the requesting SA's own slot.""" config = SAConfig(max_seq_len=1024, max_slots=16, enable_global_pool=True) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) # Request 0: [1, 2, 3, 100, 1, 2] — has [1, 2] with continuation [3, 100, ...] manager.add_request(0, [1, 2, 3, 100, 1, 2]) @@ -658,12 +674,10 @@ def test_extend_global_prefers_own_slot(self): draft_1 = draft_tokens[1].cpu().tolist() assert draft_1[0] == 200, f"Request 1 should use own slot continuation (200), got {draft_1}" - manager.shutdown() - - def test_extend_global_no_match(self): + def test_extend_global_no_match(self, make_manager): """No match across any SA returns match_len=0.""" config = SAConfig(max_seq_len=1024, max_slots=16, enable_global_pool=True) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) manager.add_request(0, [1, 2, 3]) manager.add_request(1, [4, 5, 6]) @@ -697,12 +711,10 @@ def test_extend_global_no_match(self): draft_1 = draft_tokens[1, :max_draft_len].cpu().tolist() assert draft_1 == [0] * max_draft_len, f"Expected zeroed draft for request 1, got {draft_1}" - manager.shutdown() - - def test_extend_global_active_slot_mask(self): + def test_extend_global_active_slot_mask(self, make_manager): """Removed requests should not be searchable via the active slot mask.""" config = SAConfig(max_seq_len=1024, max_slots=16, enable_global_pool=True) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) # Request 0 has pattern [1, 2, 3, 4, 5] manager.add_request(0, [1, 2, 3, 4, 5]) @@ -740,12 +752,10 @@ def test_extend_global_active_slot_mask(self): f"got {match_len[0].item()}" ) - manager.shutdown() - - def test_extend_global_single_request(self): + def test_extend_global_single_request(self, make_manager): """Global search with a single request behaves like local search.""" config = SAConfig(max_seq_len=1024, max_slots=16, enable_global_pool=True) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) context_tokens = [0, 1, 2, 1, 2] manager.add_request(0, context_tokens) @@ -774,12 +784,10 @@ def test_extend_global_single_request(self): f"Expected continuation 2, got {draft_tokens[0, 0].item()}" ) - manager.shutdown() - - def test_extend_global_cuda_graph(self): + def test_extend_global_cuda_graph(self, make_manager): """Test that extend_global works with CUDA graph capture.""" config = SAConfig(max_seq_len=1024, max_slots=16, enable_global_pool=True) - manager = SuffixAutomatonManager(config, max_num_requests=16) + manager = make_manager(config, max_num_requests=16) manager.add_request(0, [1, 2, 3, 4, 5, 1, 2, 3]) @@ -822,13 +830,11 @@ def test_extend_global_cuda_graph(self): match_len_val = match_len[0].item() assert match_len_val >= 1, f"Expected match after CUDA graph replay, got {match_len_val}" - manager.shutdown() - class TestRetainedPool: """Tests for retained slot pool (completed requests stay searchable).""" - def test_retained_slot_is_searchable(self): + def test_retained_slot_is_searchable(self, make_manager): """Completed request's SA stays searchable by active requests.""" # pool_size=4 > max_num_requests=2 → retention capacity of 2 config = SAConfig( @@ -837,7 +843,7 @@ def test_retained_slot_is_searchable(self): enable_global_pool=True, global_pool_size=4, ) - manager = SuffixAutomatonManager(config, max_num_requests=2) + manager = make_manager(config, max_num_requests=2) # Request A: context has [1, 2, 3, 4, 5] manager.add_request(0, [1, 2, 3, 4, 5]) @@ -882,9 +888,7 @@ def test_retained_slot_is_searchable(self): draft_c = draft_tokens[1].cpu().tolist() assert draft_c[0] == 5, f"Expected continuation token 5 from A, got {draft_c}" - manager.shutdown() - - def test_eviction_fifo_order(self): + def test_eviction_fifo_order(self, make_manager): """Oldest retained slot is evicted first when pool is full.""" # pool_size=4, max_batch=2 → 2 retained slot capacity config = SAConfig( @@ -893,7 +897,7 @@ def test_eviction_fifo_order(self): enable_global_pool=True, global_pool_size=4, ) - manager = SuffixAutomatonManager(config, max_num_requests=2) + manager = make_manager(config, max_num_requests=2) # Initial: free=[0,1,2,3], active={}, retained={} manager.add_request(0, [1, 2, 3]) @@ -925,9 +929,7 @@ def test_eviction_fifo_order(self): retained_rids = list(manager._retained_slots.values()) assert retained_rids == [1], f"Expected B (rid=1) retained, got {retained_rids}" - manager.shutdown() - - def test_active_never_evicted(self): + def test_active_never_evicted(self, make_manager): """Active (in-flight) requests must never be evicted.""" # pool_size=2 = max_batch=2 → 0 retained capacity → no retention config = SAConfig( @@ -936,7 +938,7 @@ def test_active_never_evicted(self): enable_global_pool=True, global_pool_size=2, ) - manager = SuffixAutomatonManager(config, max_num_requests=2) + manager = make_manager(config, max_num_requests=2) manager.add_request(0, [1, 2, 3]) manager.add_request(1, [4, 5, 6]) @@ -947,16 +949,14 @@ def test_active_never_evicted(self): assert len(manager._retained_slots) == 0 assert len(manager._free_slots) == 1 - manager.shutdown() - - def test_no_retention_when_global_pool_disabled(self): + def test_no_retention_when_global_pool_disabled(self, make_manager): """With global pool off, remove_request always frees immediately.""" config = SAConfig( max_seq_len=1024, max_slots=4, enable_global_pool=False, ) - manager = SuffixAutomatonManager(config, max_num_requests=4) + manager = make_manager(config, max_num_requests=4) manager.add_request(0, [1, 2, 3]) manager.prepare([0], max_draft_len=4) @@ -965,9 +965,7 @@ def test_no_retention_when_global_pool_disabled(self): assert len(manager._retained_slots) == 0 assert len(manager._free_slots) == 4 - manager.shutdown() - - def test_stale_request_not_retained(self): + def test_stale_request_not_retained(self, make_manager): """Request removed before GPU copy is flushed should not be retained.""" config = SAConfig( max_seq_len=1024, @@ -975,7 +973,7 @@ def test_stale_request_not_retained(self): enable_global_pool=True, global_pool_size=4, ) - manager = SuffixAutomatonManager(config, max_num_requests=2) + manager = make_manager(config, max_num_requests=2) # Add but don't prepare (GPU copy still pending) manager.add_request(0, [1, 2, 3]) @@ -986,8 +984,6 @@ def test_stale_request_not_retained(self): assert len(manager._retained_slots) == 0 assert len(manager._free_slots) == 4 # slot returned to free list - manager.shutdown() - class TestNativeKernel: """Tests for native kernel.""" @@ -1010,50 +1006,6 @@ def test_native_kernel(self): if __name__ == "__main__": - # Run basic tests - print("=" * 60) - print("Testing suffix automaton module (native kernel only)") - print("=" * 60) - - print("\n--- Native kernel tests ---") - test = TestNativeKernel() - test.test_native_kernel() - - print("\n--- Manager tests ---") - test = TestSuffixAutomatonManager() - test.test_manager_creation() - test.test_manager_add_remove() - test.test_manager_extend() - - print("\n--- extend_ngram tests ---") - test = TestExtendNgram() - test.test_extend_ngram_longest_match() - test.test_extend_ngram_fixed_size() - test.test_extend_ngram_no_match() - test.test_extend_ngram_batch() - test.test_extend_ngram_cuda_graph() - - print("\n--- extend_global tests ---") - test = TestExtendGlobal() - test.test_extend_global_cross_request_match() - test.test_extend_global_prefers_own_slot() - test.test_extend_global_no_match() - test.test_extend_global_active_slot_mask() - test.test_extend_global_single_request() - test.test_extend_global_cuda_graph() - - print("\n--- CUDA graph compatibility tests ---") - test = TestCUDAGraphCompatibility() - test.test_cuda_graph_capture() - - print("\n--- Retained pool tests ---") - test = TestRetainedPool() - test.test_retained_slot_is_searchable() - test.test_eviction_fifo_order() - test.test_active_never_evicted() - test.test_no_retention_when_global_pool_disabled() - test.test_stale_request_not_retained() - - print("\n" + "=" * 60) - print("All tests passed!") - print("=" * 60) + import sys + + sys.exit(pytest.main([__file__, "-v"])) From 46f770e873d91d8c18b25570b83eb6de09e5c6f1 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Wed, 19 Aug 2026 09:46:12 -0700 Subject: [PATCH 3/5] [TRTLLM-15035][test] Fail loudly on a missing Kimi-K3 checkpoint; fix stale list comment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review feedback (brnguyen2): - test_kimi_k3_sa_specdec_logits_parity: a missing checkpoint is now a hard failure instead of a skip. On the post-merge GB300 stage a skip is indistinguishable from a pass, so a checkpoint dropped from the runners' models mount would silently end this coverage — same reasoning already applied to the harness-existence check in this PR. The checkpoint is verified staged on the CI models share, and the pre-merge --extra-stage run exercises this path before merge. - l0_gb300_multi_gpus.yml: entry comment updated to match (fails, not skips, without the checkpoint). - l0_b200.yml: the suffix-automaton entry comment claimed 'no skip guards', which went stale when the previous review round added the module-level cuda-availability skipif. Comment fixed; the guard is kept because it matches the sibling KDA parity tests' pytestmark convention and protects plain 'pytest tests/unittest/' runs on GPU-less machines. Signed-off-by: Michal Guzek --- .../integration/defs/test_kimi_k3_specdec.py | 21 ++++++++++++------- .../test_lists/test-db/l0_b200.yml | 5 +++-- .../test-db/l0_gb300_multi_gpus.yml | 5 +++-- 3 files changed, 19 insertions(+), 12 deletions(-) diff --git a/tests/integration/defs/test_kimi_k3_specdec.py b/tests/integration/defs/test_kimi_k3_specdec.py index d4170a5b543c..8d277e74ca87 100644 --- a/tests/integration/defs/test_kimi_k3_specdec.py +++ b/tests/integration/defs/test_kimi_k3_specdec.py @@ -10,10 +10,13 @@ logits; see the harness docstring). Requirements: 4 GPUs and the Kimi K3 checkpoint (env KIMI_K3_CKPT or -/Kimi-K3). Skips cleanly when the -checkpoint is absent. The MoE backend defaults to VANILLA (the reference -dequant path — the bit-parity oracle; slow but fine at 4 layers) so the -test has no fused-kernel dependency and runs on any arch. +/Kimi-K3). Fails — deliberately does not skip — when the +checkpoint is absent: the test is CI-listed (GB300 post-merge), and a +checkpoint that vanishes from the runners' models mount must surface as a +regression rather than an indistinguishable green skip. The MoE backend +defaults to VANILLA (the reference dequant path — the bit-parity oracle; +slow but fine at 4 layers) so the test has no fused-kernel dependency and +runs on any arch. """ import os @@ -41,10 +44,12 @@ def _find_checkpoint(): @pytest.mark.skip_less_device(4) def test_kimi_k3_sa_specdec_logits_parity(): ckpt = _find_checkpoint() - if ckpt is None: - pytest.skip( - "Kimi K3 checkpoint not available (set KIMI_K3_CKPT or stage under LLM_MODELS_ROOT)" - ) + # Hard failure, not a skip: on the post-merge stage a skip is + # indistinguishable from a pass, so a checkpoint dropped from the + # runners' models mount would silently end this coverage. + assert ckpt is not None, ( + "Kimi K3 checkpoint not found (set KIMI_K3_CKPT or stage under LLM_MODELS_ROOT)" + ) env = os.environ.copy() env.update( diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 690f35527dfa..134f11f90837 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -178,8 +178,9 @@ l0_b200: - unittest/_torch/speculative/test_eagle3.py - unittest/_torch/speculative/test_dspark_cute_dsl_attention.py - unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py - # On-device suffix-automaton spec-dec kernel tests (GPU-only: CUDA graph - # capture, no skip guards; relocated from tests/torch/ which test-db cannot list). + # On-device suffix-automaton spec-dec kernel tests (CUDA graph capture; + # cuda-availability skipif like the sibling KDA parity tests; relocated + # from tests/torch/ which test-db cannot list). - unittest/_torch/speculative/test_suffix_automaton.py - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/visual_gen/kernels/parallel diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index dc196c7c3e80..67011ca2e255 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -51,8 +51,9 @@ l0_gb300_multi_gpus: # checkpoint; the case moved out of the skipped TestMoeFp4 class). - unittest/_torch/moe/test_moe_backend.py -k "kimi_k3_shape" # Kimi K3 truncated-model (first 4 layers) e2e SA spec-dec logits parity; - # resolves /Kimi-K3 and skips cleanly when the checkpoint - # is not staged. + # resolves /Kimi-K3 and FAILS if the checkpoint is not + # staged (deliberate: a skipped post-merge entry is indistinguishable + # from a pass, so a vanished checkpoint must be a visible regression). - test_kimi_k3_specdec.py::test_kimi_k3_sa_specdec_logits_parity TIMEOUT (40) - unittest/_torch/multi_gpu_modeling -k "deepseek" - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] From 8c9e08aa975870a462f45eaca746992f4fd8fd79 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Wed, 19 Aug 2026 10:14:13 -0700 Subject: [PATCH 4/5] [TRTLLM-15035][test] Fix stale MoE-backend claim in the specdec test docstring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review feedback (CodeRabbit): the docstring claimed the MoE backend defaults to VANILLA (the reference dequant parity oracle), but the harness defaults KIMI_K3_MOE_BACKEND to AUTO and KimiK3MoERuntime routes the routed-expert backend to TRTLLM regardless of moe_config.backend — the test body's env-block comment already said so. Docstring now matches: TRTLLM backend, parity holds because baseline and spec runs share it. Forcing VANILLA is not an option (the override would not take effect). Signed-off-by: Michal Guzek --- tests/integration/defs/test_kimi_k3_specdec.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/integration/defs/test_kimi_k3_specdec.py b/tests/integration/defs/test_kimi_k3_specdec.py index 8d277e74ca87..6faaa15b1b28 100644 --- a/tests/integration/defs/test_kimi_k3_specdec.py +++ b/tests/integration/defs/test_kimi_k3_specdec.py @@ -13,10 +13,11 @@ /Kimi-K3). Fails — deliberately does not skip — when the checkpoint is absent: the test is CI-listed (GB300 post-merge), and a checkpoint that vanishes from the runners' models mount must surface as a -regression rather than an indistinguishable green skip. The MoE backend -defaults to VANILLA (the reference dequant path — the bit-parity oracle; -slow but fine at 4 layers) so the test has no fused-kernel dependency and -runs on any arch. +regression rather than an indistinguishable green skip. The MoE backend is +routed to TRTLLM by KimiK3MoERuntime regardless of any +KIMI_K3_MOE_BACKEND / moe_config.backend override (see the comment on the +env block below); parity holds because the baseline and spec runs share +the same backend. """ import os From 9c8ef565b2caad193607de871aab2c0ac16c406f Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Wed, 19 Aug 2026 21:42:05 -0700 Subject: [PATCH 5/5] [TRTLLM-15035][test] Diagnose un-hydrated (git-lfs pointer) checkpoints up front MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first post-merge GB300 run of the SA logits-parity test failed deep inside tokenizer loading: tiktoken tried to parse 'version https://git-lfs.github.com/spec/v1' as a vocab. Root cause is checkpoint staging, not code: Kimi-K3 is staged from a git-lfs clone, the staging source was hydrated on 2026-08-17, but the GB300 runner's per-cluster models mirror still served the pre-hydration ~130-byte pointer files. The test now scans the resolved checkpoint's top-level files for the LFS pointer magic before spending any engine time, and fails with a message naming the offending files and pointing at mirror re-sync / 'git lfs pull' — instead of a misleading tiktoken traceback after the truncated checkpoint has already been materialized. Validated against the hydrated staging copy (no offenders — the guard will not fire once mirrors sync) and a synthetic pointer file (detected). Signed-off-by: Michal Guzek --- .../integration/defs/test_kimi_k3_specdec.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tests/integration/defs/test_kimi_k3_specdec.py b/tests/integration/defs/test_kimi_k3_specdec.py index 6faaa15b1b28..0fe76756760f 100644 --- a/tests/integration/defs/test_kimi_k3_specdec.py +++ b/tests/integration/defs/test_kimi_k3_specdec.py @@ -42,6 +42,32 @@ def _find_checkpoint(): return None +_LFS_MAGIC = b"version https://git-lfs.github.com/spec/v1" + + +def _find_lfs_pointer_files(ckpt): + """Top-level checkpoint files that are still git-lfs pointers. + + The checkpoint is staged from a git-lfs clone; a models mirror that has + not been hydrated (or has lagged the hydrated source) serves ~130-byte + pointer files instead of the real blobs, and the resulting failures are + deep and misleading (e.g. tiktoken parsing the pointer text as a vocab). + """ + offenders = [] + for name in sorted(os.listdir(ckpt)): + path = os.path.join(ckpt, name) + if not os.path.isfile(path): + continue + try: + with open(path, "rb") as f: + head = f.read(len(_LFS_MAGIC)) + except OSError: + continue + if head == _LFS_MAGIC: + offenders.append(name) + return offenders + + @pytest.mark.skip_less_device(4) def test_kimi_k3_sa_specdec_logits_parity(): ckpt = _find_checkpoint() @@ -51,6 +77,13 @@ def test_kimi_k3_sa_specdec_logits_parity(): assert ckpt is not None, ( "Kimi K3 checkpoint not found (set KIMI_K3_CKPT or stage under LLM_MODELS_ROOT)" ) + lfs_pointers = _find_lfs_pointer_files(ckpt) + assert not lfs_pointers, ( + f"Kimi K3 checkpoint at {ckpt} is not hydrated on this runner's models " + f"mirror — these files are still git-lfs pointers: {lfs_pointers}. " + f"This is a checkpoint-staging/mirror-sync problem, not a code failure; " + f"re-sync the mirror or 'git lfs pull' the staging copy." + ) env = os.environ.copy() env.update(