diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index bf67c89236da..58a552258b08 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -12,6 +12,7 @@ l0_cpu: backend: generic orchestrator: mpi tests: + - unittest/_torch/auto_deploy - unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py - unittest/_torch/distributed - unittest/_torch/executor @@ -25,6 +26,9 @@ l0_cpu: - unittest/_torch/speculative/hw_agnostic - unittest/_torch/test_model_config.py - unittest/_torch/thop/parallel_hw_agnostic/test_custom_ops.py + - unittest/auto_deploy/multigpu + - unittest/auto_deploy/singlegpu + - unittest/auto_deploy/standalone - unittest/bindings - unittest/disaggregated - unittest/executor/test_base_worker.py ISOLATION diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_gpt_oss_modeling.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_gpt_oss_modeling.py index af433206439c..6407b2081d43 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_gpt_oss_modeling.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_gpt_oss_modeling.py @@ -548,6 +548,7 @@ def test_full_model_equivalence(): @_GPT_OSS_MODELING_NVBUG +@pytest.mark.cpu_only def test_export(): """Model can be exported with torch.export and produces correct output.""" device = "cpu" diff --git a/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py b/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py index b9ae83d8f020..220383982d79 100644 --- a/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py +++ b/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py @@ -44,6 +44,7 @@ "SYMM_MEM", ], ) +@pytest.mark.cpu_only def test_allreduce_strategy_propagation(strategy): """Test that allreduce_strategy is correctly propagated to graph nodes. @@ -151,6 +152,7 @@ def forward(self, x): "SYMM_MEM", ], ) +@pytest.mark.cpu_only def test_allgather_strategy_propagation(strategy): """Test that allgather_strategy is correctly propagated to graph nodes. diff --git a/tests/unittest/auto_deploy/multigpu/transformations/library/test_tp_sharding.py b/tests/unittest/auto_deploy/multigpu/transformations/library/test_tp_sharding.py index 23c9e29a7093..f3bef24be439 100644 --- a/tests/unittest/auto_deploy/multigpu/transformations/library/test_tp_sharding.py +++ b/tests/unittest/auto_deploy/multigpu/transformations/library/test_tp_sharding.py @@ -139,6 +139,7 @@ def forward(self, x): pytest.param(576, 4, 128, id="non_aligned_uneven_n-N576-ws4"), ], ) +@pytest.mark.cpu_only def test_finegrained_fp8_get_sharded_scale(weight_original_n, world_size, block_n): """Tests FineGrained FP8 scale sharding for all shard boundary cases. @@ -481,6 +482,7 @@ def forward(self, x): return torch.ops.aten.view.default(x, (b, s, 32, 128)) +@pytest.mark.cpu_only def test_update_node_args_preserves_nested_symbolic_shape_nodes(): gm = torch.fx.symbolic_trace(SymbolicShapeView()) view_node = next(node for node in gm.graph.nodes if is_op(node, [torch.ops.aten.view])) diff --git a/tests/unittest/auto_deploy/singlegpu/compile/test_captured_graph.py b/tests/unittest/auto_deploy/singlegpu/compile/test_captured_graph.py index 9993b6cb8dea..51707d100687 100644 --- a/tests/unittest/auto_deploy/singlegpu/compile/test_captured_graph.py +++ b/tests/unittest/auto_deploy/singlegpu/compile/test_captured_graph.py @@ -130,6 +130,7 @@ def _find_dynamic_op_node(gm): ) +@pytest.mark.cpu_only def test_inject_out_param_reuses_positional_out_schema_slot(): graph = Graph() q = graph.placeholder("q") @@ -180,6 +181,7 @@ def test_inject_out_param_reuses_positional_out_schema_slot(): assert "out" not in dynamic_node.kwargs +@pytest.mark.cpu_only def test_inject_out_param_uses_kwarg_when_out_slot_not_materialized(): graph = Graph() q = graph.placeholder("q") @@ -237,6 +239,7 @@ def test_inject_out_param_uses_kwarg_when_out_slot_not_materialized(): ([], 15, None), # Case 8: empty list should return None ], ) +@pytest.mark.cpu_only def test_round_up_to_closest(lst, value, expected): assert _round_up_to_closest(lst, value) == expected @@ -392,6 +395,7 @@ def get_args_kwargs(batch_size): class TestCapturedGraphCapture: """Tests for capture-time input truncation in CapturedGraph.""" + @pytest.mark.cpu_only def test_capture_graph_uses_per_input_extents_for_truncation(self, monkeypatch): class ModelWithDifferentDynamicDims(nn.Module): def forward(self, x, y): @@ -426,6 +430,7 @@ def get_args_kwargs(bs): (3, 2, 2, 4), } + @pytest.mark.cpu_only def test_capture_graph_refetches_max_batch_after_probe(self, monkeypatch): class ModelWithStatefulMetadata(nn.Module): def __init__(self): @@ -457,6 +462,7 @@ def get_args_kwargs(bs): assert compiled_model.model.seen == [(2, 2)] + @pytest.mark.cpu_only def test_capture_graph_skips_static_arg_mismatched_batch_size(self, monkeypatch): class ModelWithStaticMetadata(nn.Module): def forward(self, x, meta): @@ -486,6 +492,7 @@ def get_args_kwargs(bs): assert captured_shapes == [torch.Size([4, 1])] assert set(compiled_model.cudagraphs) == {(4, 1)} + @pytest.mark.cpu_only def test_auto_batched_inputs_keep_explicit_resources_static(self, monkeypatch): class ModelWithInterleavedKwargs(nn.Module): def forward(self, runtime_a, explicit_cache, runtime_b): @@ -520,6 +527,7 @@ def get_args_kwargs(bs): assert [tuple(buf.shape) for buf in compiled_model._input_buffers] == [(4, 2), (4, 2)] assert captured_kwarg_orders == [("runtime_a", "runtime_b", "explicit_cache")] + @pytest.mark.cpu_only def test_auto_batched_inputs_keep_cache_seq_interface_static(self, monkeypatch): class Interface: pass @@ -556,6 +564,7 @@ def get_args_kwargs(bs): assert [tuple(buf.shape) for buf in compiled_model._input_buffers] == [(4, 2)] assert captured_kwarg_orders == [("runtime_a", "cache_seq_interface")] + @pytest.mark.cpu_only def test_auto_batched_inputs_do_not_guess_legacy_cache_names(self, monkeypatch): class ModelWithLegacyCacheName(nn.Module): def forward(self, runtime_a, r0_cache, runtime_b): @@ -626,6 +635,7 @@ def get_args_kwargs(batch_size): assert out.shape == (bs - 1, hidden_size) torch.testing.assert_close(out, replay_input[:-1] + 1) + @pytest.mark.cpu_only def test_forward_uses_captured_output_extent_when_input_extent_is_larger(self, monkeypatch): class GatherLikeModel(nn.Module): def __init__(self): @@ -676,6 +686,7 @@ def get_args_kwargs(bs): assert out.shape == (4, 2) torch.testing.assert_close(out, torch.full((4, 2), 4.0)) + @pytest.mark.cpu_only def test_forward_falls_back_when_captured_output_extent_exceeds_buffer(self, monkeypatch): class EchoModel(nn.Module): def __init__(self): @@ -720,6 +731,7 @@ def get_args_kwargs(bs): ((2, 2), "output extent metadata does not match captured outputs"), ], ) + @pytest.mark.cpu_only def test_forward_raises_for_inconsistent_output_extent_metadata( self, monkeypatch, output_extents, error_match ): @@ -787,6 +799,7 @@ def forward(self, x): # ============================================================================ +@pytest.mark.cpu_only class TestSubmodHasCudaOps: """Tests for submod_has_cuda_ops.""" @@ -837,6 +850,7 @@ def test_graphmodule_with_only_trivial_methods(self): # ============================================================================ +@pytest.mark.cpu_only class TestDualModeCapturedGraphRouting: """Tests for DualModeCapturedGraph routing logic (no actual graph capture).""" @@ -950,6 +964,7 @@ def test_find_nearest_bucket_empty(self): # ============================================================================ +@pytest.mark.cpu_only class TestPiecewiseCapturedGraphPrepare: """Tests for PiecewiseCapturedGraph.prepare.""" @@ -976,6 +991,7 @@ def test_prepare_is_idempotent(self): # ============================================================================ +@pytest.mark.cpu_only class TestPiecewiseCapturedGraphOutputHandling: """Tests for output reconstruction and forward-state cleanup.""" @@ -1030,6 +1046,7 @@ def test_forward_synchronizes_and_reconstructs_output(self, monkeypatch): # ============================================================================ +@pytest.mark.cpu_only class TestPiecewiseCapturedGraphStaticInputBuffers: """Tests for static kwarg buffers used by piecewise capture.""" @@ -1126,6 +1143,7 @@ def get_args_kwargs(num_tokens): # ============================================================================ +@pytest.mark.cpu_only class TestADPiecewiseRunnerCapture: """Tests for dynamic output buffers allocated during runner capture.""" @@ -1181,6 +1199,7 @@ def fake_make_weak_ref(value): # ============================================================================ +@pytest.mark.cpu_only class TestGenerateDefaultPiecewiseNumTokens: """Tests for _generate_default_piecewise_num_tokens.""" @@ -1233,6 +1252,7 @@ def test_no_duplicates_when_max_is_power_of_two(self): # ============================================================================ +@pytest.mark.cpu_only class TestCompileModelGraphModuleTargetCollection: """Tests for selecting GraphModule compile targets.""" @@ -1509,6 +1529,7 @@ def _build_plain_static_submod(): return torch.fx.GraphModule(root, g) +@pytest.mark.cpu_only class TestPiecewiseCapturedGraphMultiStreamWiring: def _build_split_gm(self, ms_submod, plain_submod): parent = nn.Module() @@ -1547,6 +1568,7 @@ def test_prepare_keeps_stream_switch_partition_as_static_runner(self, monkeypatc assert isinstance(pcg.split_gm.submod_1, ADPiecewiseRunner) +@pytest.mark.cpu_only class TestSetupPiecewiseMixedBatch: """Coverage for piecewise warmup synthetic mixed-batch setup. diff --git a/tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_runner.py b/tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_runner.py index 4b831eb8e173..658f7fd2e04a 100644 --- a/tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_runner.py +++ b/tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_runner.py @@ -30,6 +30,7 @@ # ============================================================================ +@pytest.mark.cpu_only class TestADPiecewiseRunnerContextManagement: def setup_method(self): ADPiecewiseRunner._current_num_tokens = None @@ -57,6 +58,7 @@ def test_set_current_phase_invalid_raises(self): # ============================================================================ +@pytest.mark.cpu_only class TestADPiecewiseRunnerInit: def test_entries_initially_empty(self): submod = nn.Linear(4, 4) @@ -80,6 +82,7 @@ def test_no_dynamic_out_info_by_default(self): # ============================================================================ +@pytest.mark.cpu_only class TestOutputInfo: def test_creation(self): info = OutputInfo( @@ -91,6 +94,7 @@ def test_creation(self): assert info.dtype == torch.float16 +@pytest.mark.cpu_only class TestSegmentEntry: def test_default_values(self): entry = SegmentEntry() @@ -116,6 +120,7 @@ def forward(self, chunk_values: torch.Tensor, seq_idx_values: torch.Tensor): return chunk_values, chunk_values.clone(), seq_idx_values.to(torch.int32).view(1, -1) +@pytest.mark.cpu_only class TestMetadataWrapper: def setup_method(self): ADPiecewiseRunner._current_num_tokens = None @@ -185,6 +190,7 @@ def test_mamba_metadata_allows_runtime_growth(self): # ============================================================================ +@pytest.mark.cpu_only class TestDynamicOutBuf: def test_set_and_get_returns_none_before_capture(self): submod = nn.Linear(4, 4) diff --git a/tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_utils.py b/tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_utils.py index 96ccecc4d2e4..ab1fd6ec92f7 100644 --- a/tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_utils.py +++ b/tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_utils.py @@ -16,6 +16,7 @@ from types import SimpleNamespace +import pytest import torch import torch.nn as nn from torch.fx import Graph, GraphModule @@ -36,6 +37,9 @@ wait_aux_stream_passthrough, ) +pytestmark = pytest.mark.cpu_only + + # ============================================================================ # Helpers # ============================================================================ diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_torch_attention_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_torch_attention_op.py index f7881816ea52..aa421db4929e 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_torch_attention_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_torch_attention_op.py @@ -24,6 +24,7 @@ from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import BatchInfo +@pytest.mark.cpu_only @torch.inference_mode() def test_gemma4_multimodal_mask_source_op(): input_ids = torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.int64) @@ -55,6 +56,7 @@ def test_gemma4_multimodal_mask_source_op(): torch.testing.assert_close(actual, expected) +@pytest.mark.cpu_only @torch.inference_mode() def test_gemma4_prepare_multimodal_mask_chunked_prefill(): batch_info = BatchInfo() @@ -90,6 +92,7 @@ def test_gemma4_prepare_multimodal_mask_chunked_prefill(): torch.testing.assert_close(actual, expected) +@pytest.mark.cpu_only @torch.inference_mode() def test_torch_attention_explicit_mask_is_authoritative(): torch.manual_seed(0) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_torch_mla_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_torch_mla_op.py index 89b22bd3738d..106b51899933 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_torch_mla_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_torch_mla_op.py @@ -738,6 +738,7 @@ def test_memory_efficiency(self): assert compression_ratio > 50, f"Expected >50x compression, got {compression_ratio:.1f}x" +@pytest.mark.cpu_only class TestMLADescriptor: """Test MultiHeadLatentAttention descriptor configuration.""" diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_triton_mla_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_triton_mla_op.py index 200f14d5d71c..b51ca71bf3b9 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_triton_mla_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/mla/test_triton_mla_op.py @@ -454,6 +454,7 @@ def test_prefill_ragged( ) +@pytest.mark.cpu_only class TestTritonMLADescriptor: """Test TritonMLAAttention descriptor configuration.""" diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py index a689db818680..0c46aa1e145c 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py @@ -1,11 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import pytest import torch from triton_kernels.tensor_details.layout import HopperMXValueLayout, StridedLayout from tensorrt_llm._torch.auto_deploy.custom_ops.fused_moe import mxfp4_moe +pytestmark = pytest.mark.cpu_only + def test_mxfp4_value_layout_uses_strided_layout_on_blackwell(monkeypatch): monkeypatch.setattr(mxfp4_moe, "cuda_capability_geq", lambda major, minor=0: major >= 10) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/quantization/test_quant.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/quantization/test_quant.py index 240904397e39..acbdb6561eae 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/quantization/test_quant.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/quantization/test_quant.py @@ -147,11 +147,12 @@ def test_fp8_bmm(input_dtype, mat2_dtype): assert cos_sim_unquantized > 0.99 -@pytest.mark.parametrize("bias", [torch.rand(32, device="cuda") * 10, None]) +@pytest.mark.parametrize("use_bias", [True, False]) @pytest.mark.skipif(not fp8_compatible(), reason="Requires fp8 support") -def test_quant_linear_fp8_matches_fused_op(bias): +def test_quant_linear_fp8_matches_fused_op(use_bias): input = torch.rand(3, 16, device="cuda") weight = torch.rand(32, 16, device="cuda") + bias = torch.rand(32, device="cuda") * 10 if use_bias else None weight_scale = (torch.max(torch.abs(weight)) / 448).to("cuda") weight_fp8 = (weight / weight_scale).to(torch.float8_e4m3fn) @@ -179,19 +180,17 @@ def test_quant_linear_fp8_matches_fused_op(bias): @pytest.mark.parametrize( - "bias", - [ - (torch.rand(32, device="cuda") * 10).to(torch.float16), - None, - ], + "use_bias", + [True, False], ) @pytest.mark.skipif( not (fp4_compatible() and trtllm_ops_available()), reason="Requires NVFP4 and TRT-LLM ops", ) -def test_quant_linear_nvfp4_matches_fused_op(bias): +def test_quant_linear_nvfp4_matches_fused_op(use_bias): x = torch.rand(3, 32, device="cuda", dtype=torch.half) # [..., K] W = torch.rand(32, 32, device="cuda", dtype=torch.half) # [N, K] + bias = (torch.rand(32, device="cuda") * 10).to(torch.float16) if use_bias else None N, K = W.shape assert K % SCALING_VECTOR_SIZE == 0 @@ -507,6 +506,7 @@ def _fg_reference_128_block(weight, scale): (256, 256), # fully aligned (must stay correct) ], ) +@pytest.mark.cpu_only def test_finegrained_fp8_dequant_coarse_block_uses_canonical_128(n, k): """A coarse 128-block scale must dequant exactly, even when N/K aren't 128-multiples.""" weight, scale = _fg_make_block_scaled_weight(n, k) @@ -514,6 +514,7 @@ def test_finegrained_fp8_dequant_coarse_block_uses_canonical_128(n, k): torch.testing.assert_close(out, _fg_reference_128_block(weight, scale), atol=0.0, rtol=0.0) +@pytest.mark.cpu_only def test_finegrained_fp8_dequant_differs_from_buggy_ceil_expansion(): """Lock in the fix: canonical-128 expansion must differ from the old ceil(N/scale_n) one.""" n, k = 576, 256 # ceil(576/5) = 116 != 128 -> old behavior was wrong @@ -527,6 +528,7 @@ def test_finegrained_fp8_dequant_differs_from_buggy_ceil_expansion(): assert not torch.allclose(correct, buggy), "fix must not reproduce the misaligned expansion" +@pytest.mark.cpu_only def test_finegrained_fp8_dequant_per_row_scale_expands_by_one(): """A per-row scale (scale_n == N) keeps the ceil path -> block_n == 1 (each row its own scale).""" n, k = 72, 256 diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/rope/test_rope_op_variants.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/rope/test_rope_op_variants.py index 8f36ebab9a99..0ba4f32982d0 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/rope/test_rope_op_variants.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/rope/test_rope_op_variants.py @@ -390,6 +390,7 @@ def test_flashinfer_custom_op_strided_interleaved(dtype, atol, rtol, head_dim): torch.testing.assert_close(custom_k_strided, out_k_ref, rtol=rtol, atol=atol) +@pytest.mark.cpu_only @pytest.mark.parametrize("has_bias", [True, False]) def test_rope_deinterleave_load_hook(has_bias): """Test _rope_deinterleave_load_hook permutes weights correctly. diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_gptq_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_gptq_op.py index 72f7cc88ce7c..4b9549f93b3d 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_gptq_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_gptq_op.py @@ -17,6 +17,8 @@ import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 +pytestmark = pytest.mark.cpu_only + def pack_gptq_qweight_from_u4(U4_nk: torch.Tensor) -> torch.Tensor: """GPTQ: pack along K, 8 nibbles per int32. diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_resource_handlers.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_resource_handlers.py index e78be45a94bd..a15da94f49f1 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_resource_handlers.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_resource_handlers.py @@ -284,7 +284,8 @@ def test_resolve_cache_dtype_explicit_float32(): @pytest.mark.skipif( - torch.cuda.get_device_capability(0) < (8, 9), reason="FP8 requires compute capability >= 8.9" + torch.cuda.device_count() == 0 or torch.cuda.get_device_capability(0) < (8, 9), + reason="FP8 requires compute capability >= 8.9", ) def test_resolve_cache_dtype_explicit_fp8(): """Test explicit 'fp8' dtype string resolves correctly.""" diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_update_kv_cache.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_update_kv_cache.py index 0e95fd5f7c4d..f8affb0d23b2 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_update_kv_cache.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_update_kv_cache.py @@ -12,12 +12,15 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from tensorrt_llm._torch.auto_deploy.custom_ops.attention.torch_backend_attention import ( _update_kv_cache, ) +pytestmark = pytest.mark.cpu_only + def test_update_kv_cache(): K_D_HEAD = 4 diff --git a/tests/unittest/auto_deploy/singlegpu/mlir/test_decompose.py b/tests/unittest/auto_deploy/singlegpu/mlir/test_decompose.py index bf7dcd71d9d0..108d1f1c2504 100644 --- a/tests/unittest/auto_deploy/singlegpu/mlir/test_decompose.py +++ b/tests/unittest/auto_deploy/singlegpu/mlir/test_decompose.py @@ -17,6 +17,9 @@ import pytest +pytestmark = pytest.mark.cpu_only + + xdsl = pytest.importorskip("xdsl") from xdsl.dialects.builtin import ( # noqa: E402 diff --git a/tests/unittest/auto_deploy/singlegpu/mlir/test_dialect.py b/tests/unittest/auto_deploy/singlegpu/mlir/test_dialect.py index 7a9b9ced1f35..a26fa70b36b2 100644 --- a/tests/unittest/auto_deploy/singlegpu/mlir/test_dialect.py +++ b/tests/unittest/auto_deploy/singlegpu/mlir/test_dialect.py @@ -17,6 +17,9 @@ import pytest +pytestmark = pytest.mark.cpu_only + + xdsl = pytest.importorskip("xdsl") from xdsl.dialects.builtin import ( # noqa: E402 diff --git a/tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py b/tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py index 398950853709..fe78832f8e8f 100644 --- a/tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py +++ b/tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py @@ -105,6 +105,7 @@ def _build_single_output_add_mul_module(hidden: int = 128): return ModuleOp(Region([block])) +@pytest.mark.cpu_only def test_decompose_discover_codegen_pipeline(): """Full pipeline: decompose add+rmsnorm -> discover subgraph -> generate kernel.""" hidden = 128 @@ -127,6 +128,7 @@ def test_decompose_discover_codegen_pipeline(): assert callable(kernel_fn), "Expected generate_kernel_from_subgraph to return a callable" +@pytest.mark.cpu_only def test_single_output_fused_metadata_uses_tuple_contract(): """Generated fused kernels return tuples, including the one-output case.""" hidden = 128 @@ -192,6 +194,7 @@ def test_decompose_discover_codegen_numerical_correctness(): torch.testing.assert_close(result[1], normed_ref, atol=1e-2, rtol=1e-2) +@pytest.mark.cpu_only def test_no_subgraphs_for_single_op(): """A module with only one fusible op should produce no subgraphs.""" hidden = 64 @@ -218,6 +221,7 @@ def test_no_subgraphs_for_single_op(): assert len(subgraphs) == 0 +@pytest.mark.cpu_only def test_pipeline_reports_subgraph_metadata(): """Verify subgraph metadata (inputs, outputs, op count) is correct.""" hidden = 128 @@ -280,6 +284,7 @@ def _build_add_rmsnorm_silu_gate_module(hidden: int = 128): return ModuleOp(Region([block])) +@pytest.mark.cpu_only def test_novel_add_rmsnorm_silu_gate(): """Fusion that has NO pre-existing kernel — auto-discovered and generated. @@ -488,6 +493,7 @@ def _build_moe_ep_mask_module(top_k: int = 8, experts_per_rank: int = 32, ep_ran return ModuleOp(Region([block])) +@pytest.mark.cpu_only def test_moe_ep_mask_fusion_discovery(): """Floordiv + eq + mul should be discovered as a single fusible subgraph.""" mlir_mod = _build_moe_ep_mask_module() @@ -497,6 +503,7 @@ def test_moe_ep_mask_fusion_discovery(): assert len(subgraphs[0].ops) == 3, f"Expected 3 ops, got {len(subgraphs[0].ops)}" +@pytest.mark.cpu_only def test_moe_ep_mask_kernel_generation(): """Triton kernel can be generated for the mixed-type floordiv+eq+mul pattern.""" mlir_mod = _build_moe_ep_mask_module() diff --git a/tests/unittest/auto_deploy/singlegpu/mlir/test_fx_mlir_roundtrip.py b/tests/unittest/auto_deploy/singlegpu/mlir/test_fx_mlir_roundtrip.py index 797e91f4d96f..aa8f831379af 100644 --- a/tests/unittest/auto_deploy/singlegpu/mlir/test_fx_mlir_roundtrip.py +++ b/tests/unittest/auto_deploy/singlegpu/mlir/test_fx_mlir_roundtrip.py @@ -159,6 +159,7 @@ def _build_precise_to_dtype_mlir_module(): # --------------------------------------------------------------------------- +@pytest.mark.cpu_only def test_mlir_to_fx_propagates_single_output_fused_getitem_meta(): import operator @@ -192,6 +193,7 @@ def test_mlir_to_fx_propagates_single_output_fused_getitem_meta(): assert getitem_nodes[0].meta["val"].dtype == fake_val.dtype +@pytest.mark.cpu_only def test_mlir_to_fx_synthesizes_precise_op_meta_from_result_type(): graph = torch.fx.Graph() x = graph.placeholder("x") diff --git a/tests/unittest/auto_deploy/singlegpu/mlir/test_string_codegen.py b/tests/unittest/auto_deploy/singlegpu/mlir/test_string_codegen.py index 8eb28c4a7b17..717badb09799 100644 --- a/tests/unittest/auto_deploy/singlegpu/mlir/test_string_codegen.py +++ b/tests/unittest/auto_deploy/singlegpu/mlir/test_string_codegen.py @@ -149,6 +149,7 @@ def test_generated_rmsnorm_correctness(): torch.testing.assert_close(result[0], ref, atol=1e-2, rtol=1e-2) +@pytest.mark.cpu_only def test_kernel_cache_reuse(): """Same subgraph structure produces the same hash.""" from tensorrt_llm._torch.auto_deploy.mlir.codegen.kernel_cache import KernelCache diff --git a/tests/unittest/auto_deploy/singlegpu/mlir/test_subgraph_discovery.py b/tests/unittest/auto_deploy/singlegpu/mlir/test_subgraph_discovery.py index e9e6ddb86c05..5f1893b2a4a0 100644 --- a/tests/unittest/auto_deploy/singlegpu/mlir/test_subgraph_discovery.py +++ b/tests/unittest/auto_deploy/singlegpu/mlir/test_subgraph_discovery.py @@ -17,6 +17,9 @@ import pytest +pytestmark = pytest.mark.cpu_only + + xdsl = pytest.importorskip("xdsl") from xdsl.dialects.builtin import ( # noqa: E402 diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_cohere_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_cohere_modeling.py index e3f8f9176d02..c1ae4f02c277 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_cohere_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_cohere_modeling.py @@ -675,6 +675,7 @@ def test_cohere2_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_cohere_config_registration(): """Test that Cohere v1 config is properly recognized.""" config = _create_small_cohere_config() @@ -683,6 +684,7 @@ def test_cohere_config_registration(): assert hasattr(config, "use_qk_norm") +@pytest.mark.cpu_only def test_cohere2_config_registration(): """Test that Cohere2 config is properly recognized.""" config = _create_small_cohere2_config() @@ -696,6 +698,7 @@ def test_cohere2_config_registration(): assert config.layer_types[3] == "full_attention" +@pytest.mark.cpu_only def test_cohere_parallel_attn_mlp_structure(): """Test that Cohere decoder layers use parallel attn+MLP (no post_attn_layernorm).""" config = _create_small_cohere_config() @@ -708,6 +711,7 @@ def test_cohere_parallel_attn_mlp_structure(): ) +@pytest.mark.cpu_only def test_cohere_gqa_structure(): """Test that attention uses GQA (fewer KV heads than Q heads).""" config = _create_small_cohere_config() @@ -718,6 +722,7 @@ def test_cohere_gqa_structure(): assert attn.num_kv_heads == 2 +@pytest.mark.cpu_only def test_cohere2_conditional_rope(): """Test that Cohere2 applies RoPE only on sliding window layers.""" config = _create_small_cohere2_config() @@ -731,6 +736,7 @@ def test_cohere2_conditional_rope(): assert model.model.layers[3].self_attn.use_rope is False +@pytest.mark.cpu_only def test_cohere_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_cohere_config() @@ -755,6 +761,7 @@ def test_cohere_state_dict_keys(): assert key in state_dict, f"Expected key '{key}' in state_dict" +@pytest.mark.cpu_only def test_cohere_logit_scale(): """Test that logit_scale is applied to output logits.""" config = _create_small_cohere_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_decilm_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_decilm_modeling.py index d328c824bd39..6d23313ad070 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_decilm_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_decilm_modeling.py @@ -321,6 +321,7 @@ def forward(self, input_ids, position_ids): # ============================================================================= +@pytest.mark.cpu_only def test_decilm_layer_structure(): """Test that layers have correct structure based on block_configs.""" config = _create_small_config() @@ -339,6 +340,7 @@ def test_decilm_layer_structure(): assert layer2.has_attention is True +@pytest.mark.cpu_only def test_decilm_ffn_mult_to_intermediate_size(): assert _ffn_mult_to_intermediate_size(5.25, 8192) == 28672 assert _ffn_mult_to_intermediate_size(2.625, 8192) == 14336 @@ -346,6 +348,7 @@ def test_decilm_ffn_mult_to_intermediate_size(): assert _ffn_mult_to_intermediate_size(0.5, 8192) == 2816 +@pytest.mark.cpu_only def test_decilm_weight_keys_match_checkpoint(): config = _create_small_config() model = DeciLMForCausalLM(config) diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_deepseek_v2_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_deepseek_v2_modeling.py index 73c6418fa7ca..a30ded187531 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_deepseek_v2_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_deepseek_v2_modeling.py @@ -632,6 +632,7 @@ def test_deepseek_v2_model_can_be_exported(): # ============================================================================= +@pytest.mark.cpu_only def test_deepseek_v2_model_registration(): """Test that DeepSeekV2ForCausalLM is registered with the factory.""" from tensorrt_llm._torch.auto_deploy.models.hf import AutoModelForCausalLMFactory diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py b/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py index 5f36b8abbc2c..268752e3bdb8 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py @@ -176,6 +176,7 @@ def _build_small_draft_factory( ) +@pytest.mark.cpu_only def test_eagle_rmsnorm_keeps_fp32_weights(): norm = EagleRMSNorm(hidden_size=16) diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_exaone_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_exaone_modeling.py index 50ec7195bc0e..75180863a1f0 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_exaone_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_exaone_modeling.py @@ -533,6 +533,7 @@ def test_exaone_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_exaone_config_attribute_map(): """Test that config attribute_map correctly aliases EXAONE-specific names.""" config = _create_small_config() @@ -547,6 +548,7 @@ def test_exaone_config_attribute_map(): assert config.layer_norm_epsilon == 1e-6 +@pytest.mark.cpu_only def test_exaone_gqa_structure(): """Test that attention uses GQA (fewer KV heads than Q heads).""" config = _create_small_config() @@ -557,6 +559,7 @@ def test_exaone_gqa_structure(): assert attn.num_kv_heads == 2, f"Expected 2 KV heads, got {attn.num_kv_heads}" +@pytest.mark.cpu_only def test_exaone_state_dict_keys(): """Test that state_dict keys match expected EXAONE checkpoint format.""" config = _create_small_config() @@ -584,6 +587,7 @@ def test_exaone_state_dict_keys(): ) +@pytest.mark.cpu_only def test_exaone_attention_block_wrapper(): """Test that ExaoneAttentionBlock correctly wraps ExaoneAttention.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_gemma2_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_gemma2_modeling.py index 441087688137..7676cb6fc6e2 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_gemma2_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_gemma2_modeling.py @@ -488,6 +488,7 @@ def test_gemma2_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_gemma2_config_registration(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -502,6 +503,7 @@ def test_gemma2_config_registration(): assert hasattr(config, "sliding_window") +@pytest.mark.cpu_only def test_gemma2_layer_types(): """Test that layer types alternate correctly.""" config = _create_small_config() @@ -512,6 +514,7 @@ def test_gemma2_layer_types(): assert config.layer_types[3] == "full_attention" +@pytest.mark.cpu_only def test_gemma2_gqa_structure(): """Test that attention uses GQA (fewer KV heads than Q heads).""" config = _create_small_config() @@ -529,6 +532,7 @@ def test_gemma2_gqa_structure(): assert attn1.sliding_window is None +@pytest.mark.cpu_only def test_gemma2_four_layernorms(): """Test that each decoder layer has 4 layer norms.""" config = _create_small_config() @@ -541,6 +545,7 @@ def test_gemma2_four_layernorms(): assert hasattr(layer, "post_feedforward_layernorm") +@pytest.mark.cpu_only def test_gemma2_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_gemma3n_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_gemma3n_modeling.py index 32ef71aa5af5..a941744f3b01 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_gemma3n_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_gemma3n_modeling.py @@ -194,6 +194,7 @@ def _set_seed(): torch.manual_seed(42) +@pytest.mark.cpu_only def test_hf_reference_available(): if HF_CLASSES is None: pytest.skip("transformers gemma3n reference classes are unavailable") @@ -289,6 +290,7 @@ def test_gemma3n_decoder_layer_equivalence(): assert_rmse_close(custom_out, hf_out, rmse_ratio_tol=0.05, msg="Decoder layer: ") +@pytest.mark.cpu_only @torch.no_grad() def test_gemma3n_full_model_equivalence(): if HF_CLASSES is None: @@ -310,6 +312,7 @@ def test_gemma3n_full_model_equivalence(): assert_rmse_close(custom_out.logits, hf_out.logits, rmse_ratio_tol=0.05, msg="Full model: ") +@pytest.mark.cpu_only @torch.no_grad() def test_gemma3n_conditional_wrapper_equivalence(): if HF_CLASSES is None: @@ -335,6 +338,7 @@ def test_gemma3n_conditional_wrapper_equivalence(): assert_rmse_close(wrapper_out.logits, hf_out.logits, rmse_ratio_tol=0.05, msg="Wrapper: ") +@pytest.mark.cpu_only def test_gemma3n_conditional_wrapper_load_hook_drops_unsupported_tower_weights(): config = _small_full_config() wrapper = Gemma3nForConditionalGeneration(config) @@ -348,12 +352,14 @@ def test_gemma3n_conditional_wrapper_load_hook_drops_unsupported_tower_weights() assert unexpected == [] +@pytest.mark.cpu_only def test_gemma3n_conditional_wrapper_ignores_hf_init_kwargs(): config = _small_full_config() wrapper = Gemma3nForConditionalGeneration(config, use_cache=False) assert isinstance(wrapper, Gemma3nForConditionalGeneration) +@pytest.mark.cpu_only def test_gemma3n_reduced_layer_load_hook_slices_per_layer_weights(): source_model = Gemma3nForCausalLM(_extended_text_config(5)) target_model = Gemma3nForCausalLM(_small_text_config()) @@ -364,11 +370,13 @@ def test_gemma3n_reduced_layer_load_hook_slices_per_layer_weights(): assert "model.layers.3.self_attn.q_proj.weight" in unexpected +@pytest.mark.cpu_only def test_gemma3n_causal_lm_ties_lm_head_to_input_embeddings(): model = Gemma3nForCausalLM(_small_text_config()) assert model.lm_head.weight.data_ptr() == model.model.embed_tokens.weight.data_ptr() +@pytest.mark.cpu_only def test_gemma3n_conditional_lm_ties_lm_head_to_input_embeddings(): model = Gemma3nForConditionalGeneration(_small_full_config()) assert ( @@ -376,6 +384,7 @@ def test_gemma3n_conditional_lm_ties_lm_head_to_input_embeddings(): ) +@pytest.mark.cpu_only def test_gemma3n_shared_kv_layer_metadata_matches_config(): model = Gemma3nForCausalLM(_shared_kv_text_config()) layer_expectations = [ @@ -392,6 +401,7 @@ def test_gemma3n_shared_kv_layer_metadata_matches_config(): assert layer.self_attn.kv_shared_layer_index == source_idx +@pytest.mark.cpu_only def test_gemma3n_export_uses_shared_kv_attention_for_shared_layers(): config = _shared_kv_text_config() model = Gemma3nForCausalLM(config).eval() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_gemma4_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_gemma4_modeling.py index 38e0051ed8c4..542b59f4137c 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_gemma4_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_gemma4_modeling.py @@ -819,6 +819,7 @@ def test_conditional_generation_wrapper(): assert torch.isfinite(out.logits).all() +@pytest.mark.cpu_only def test_shared_kv_layer_metadata_matches_config(): model = Gemma4ForCausalLM(_shared_kv_text_config()) layer_expectations = [ @@ -835,6 +836,7 @@ def test_shared_kv_layer_metadata_matches_config(): assert layer.self_attn.kv_shared_layer_index == source_idx +@pytest.mark.cpu_only def test_export_uses_shared_kv_attention_for_shared_layers(): config = _shared_kv_text_config() model = Gemma4ForCausalLM(config).eval() @@ -905,6 +907,7 @@ def test_shared_kv_eager_layers_ignore_local_kv_weights(): # --------------------------------------------------------------------------- +@pytest.mark.cpu_only def test_export(): """Model can be exported with torch.export and produces correct output.""" device = "cpu" @@ -1037,6 +1040,7 @@ def test_dense_conditional_generation_wrapper(): assert torch.isfinite(out.logits).all() +@pytest.mark.cpu_only def test_conditional_generation_wrapper_handles_exported_text_graph_without_per_layer_inputs(): config = Gemma4Config( text_config=_small_text_config(), @@ -1064,6 +1068,7 @@ def test_conditional_generation_wrapper_handles_exported_text_graph_without_per_ assert out.logits.shape == (B, S, config.text_config.vocab_size) +@pytest.mark.cpu_only def test_dense_export(): """Dense model (no MoE) can be exported with torch.export.""" device = "cpu" @@ -1591,6 +1596,7 @@ def test_vision_patch_embedder_equivalence(): torch.testing.assert_close(ad_out, ref_out, rtol=1e-3, atol=1e-3) +@pytest.mark.cpu_only def test_image_processor_pads_to_fixed_patch_budget(): """Image processor should pad every request to the configured patch budget.""" config = _small_vision_config() @@ -1617,6 +1623,7 @@ def test_image_processor_pads_to_fixed_patch_budget(): assert torch.all(outputs["image_position_ids"][1, :18] >= 0) +@pytest.mark.cpu_only def test_ad_input_processor_emits_layout_metadata_for_boi_eoi_spans(): class _DummyBaseProcessor: def __init__(self): @@ -2012,6 +2019,7 @@ def test_e2b_like_conditional_wrapper_forwards_explicit_per_layer_inputs(): torch.testing.assert_close(explicit_logits, implicit_logits, rtol=1e-4, atol=1e-4) +@pytest.mark.cpu_only def test_e2b_like_hf_per_layer_state_dict_keys_are_present_and_loadable(): config = Gemma4Config( text_config=_small_e2b_text_config(), diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_gemma_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_gemma_modeling.py index adb11fefd534..ce62c7aa46ea 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_gemma_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_gemma_modeling.py @@ -455,6 +455,7 @@ def test_gemma_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_gemma_config_registration(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -465,6 +466,7 @@ def test_gemma_config_registration(): assert hasattr(config, "head_dim") +@pytest.mark.cpu_only def test_gemma_mha_structure(): """Test that attention uses MHA (same Q and KV head count).""" config = _create_small_config() @@ -475,6 +477,7 @@ def test_gemma_mha_structure(): assert attn.num_kv_heads == 4, f"Expected 4 KV heads (MHA), got {attn.num_kv_heads}" +@pytest.mark.cpu_only def test_gemma_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_config() @@ -502,6 +505,7 @@ def test_gemma_state_dict_keys(): ) +@pytest.mark.cpu_only def test_gemma_embedding_normalization(): """Test that Gemma applies sqrt(hidden_size) normalization to embeddings.""" config = _create_small_config() @@ -514,6 +518,7 @@ def test_gemma_embedding_normalization(): ) +@pytest.mark.cpu_only def test_gemma_tied_weights(): """Test that Gemma properly supports tied embed/lm_head weights.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_glm4_moe_lite_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_glm4_moe_lite_modeling.py index e6481ed7e0ff..066f70dfce9e 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_glm4_moe_lite_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_glm4_moe_lite_modeling.py @@ -215,6 +215,7 @@ def test_glm4_moe_lite_model_can_be_exported(): assert torch.isfinite(logits2).all(), "Logits should not contain NaN or Inf" +@pytest.mark.cpu_only def test_glm4_moe_lite_config_registration(): """Test that the config is properly registered or model_type is correct.""" # Create a config and verify model_type @@ -229,6 +230,7 @@ def test_glm4_moe_lite_config_registration(): assert hasattr(config, "qk_rope_head_dim") +@pytest.mark.cpu_only def test_glm4_moe_lite_layer_types(): """Test that layer 0 uses dense MLP and later layers use MoE.""" config = _create_small_config() @@ -248,6 +250,7 @@ def test_glm4_moe_lite_layer_types(): ) +@pytest.mark.cpu_only def test_glm4_moe_lite_expert_structure(): """Test that experts have correct structure for checkpoint loading.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_glm4_moe_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_glm4_moe_modeling.py index c16f76fe7b50..ef97e41a309d 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_glm4_moe_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_glm4_moe_modeling.py @@ -496,6 +496,7 @@ def test_glm4_moe_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_glm4_moe_config_registration(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -507,6 +508,7 @@ def test_glm4_moe_config_registration(): assert hasattr(config, "use_qk_norm") +@pytest.mark.cpu_only def test_glm4_moe_gqa_structure(): """Test that attention uses GQA (fewer KV heads than Q heads).""" config = _create_small_config() @@ -525,6 +527,7 @@ def test_glm4_moe_gqa_structure(): assert hasattr(attn, "k_norm"), "Attention should have k_norm" +@pytest.mark.cpu_only def test_glm4_moe_dense_and_moe_layers(): """Test that first_k_dense_replace layers are dense, rest are MoE.""" config = _create_small_config() @@ -537,6 +540,7 @@ def test_glm4_moe_dense_and_moe_layers(): assert isinstance(model.model.layers[2].mlp, Glm4MoeMoE), "Layer 2 should use MoE" +@pytest.mark.cpu_only def test_glm4_moe_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_glm_moe_dsa_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_glm_moe_dsa_modeling.py index 18124e0382bc..53669dca4019 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_glm_moe_dsa_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_glm_moe_dsa_modeling.py @@ -476,6 +476,7 @@ def _copy_full_model_weights_hf_to_custom(hf_model, custom_model): # ============================================================================= +@pytest.mark.cpu_only def test_glm_moe_dsa_config_registration(): """Test that the config is properly registered.""" config = _create_small_config() @@ -488,6 +489,7 @@ def test_glm_moe_dsa_config_registration(): assert hasattr(config, "moe_layer_freq") +@pytest.mark.cpu_only def test_glm_moe_dsa_config_nested_rope_parameters(): """Test that nested rope_parameters extracts rope_theta correctly.""" config = GlmMoeDsaConfig( @@ -499,6 +501,7 @@ def test_glm_moe_dsa_config_nested_rope_parameters(): assert config2.rope_theta == 10000.0 +@pytest.mark.cpu_only def test_glm_moe_dsa_layer_types(): """Test that first layers use dense MLP and later layers use MoE.""" config = _create_small_config() @@ -514,6 +517,7 @@ def test_glm_moe_dsa_layer_types(): assert type(layer_mlp).__name__ == "GlmMoeDsaMoE" +@pytest.mark.cpu_only def test_glm_moe_dsa_expert_structure(): """Test that experts have correct structure for checkpoint loading.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_granite_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_granite_modeling.py index 028b9511f91e..101c73a9d978 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_granite_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_granite_modeling.py @@ -450,6 +450,7 @@ def test_granite_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_granite_config_registration(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -464,6 +465,7 @@ def test_granite_config_registration(): assert hasattr(config, "logits_scaling") +@pytest.mark.cpu_only def test_granite_gqa_structure(): """Test that attention uses GQA (fewer KV heads than Q heads).""" config = _create_small_config() @@ -476,6 +478,7 @@ def test_granite_gqa_structure(): assert attn.scaling == config.attention_multiplier +@pytest.mark.cpu_only def test_granite_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_config() @@ -503,6 +506,7 @@ def test_granite_state_dict_keys(): ) +@pytest.mark.cpu_only def test_granite_scaling_factors(): """Test that Granite-specific scaling factors are applied correctly.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_hf.py b/tests/unittest/auto_deploy/singlegpu/models/test_hf.py index 55ff51603bdb..d242d8072ec0 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_hf.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_hf.py @@ -93,6 +93,7 @@ def mock_factory(): yield factory +@pytest.mark.cpu_only def test_recursive_update_config(mock_factory): """Test that _recursive_update_config correctly updates a config object recursively.""" # Get the mocked factory instance @@ -150,6 +151,7 @@ def test_recursive_update_config(mock_factory): assert config.text_config.rope_scaling["type"] == "linear" +@pytest.mark.cpu_only def test_register_custom_model_cls(): config_cls_name = "FooConfig" custom_model_cls = MagicMock(spec=AutoModelForCausalLM) @@ -169,6 +171,7 @@ class FooConfig: pass +@pytest.mark.cpu_only def test_build_model_raises_when_custom_model_cls_does_not_have_from_config(mock_factory): custom_model_cls = MagicMock(spec=AutoModelForCausalLM, __name__="FooModel") AutoModelForCausalLMFactory.register_custom_model_cls( @@ -186,6 +189,7 @@ def test_build_model_raises_when_custom_model_cls_does_not_have_from_config(mock mock_factory.build_model(device="meta") +@pytest.mark.cpu_only def test_build_model_uses_custom_model_cls_from_config(mock_factory): custom_model_cls = MagicMock(spec=AutoModelForCausalLM) custom_model_cls.configure_mock(_from_config=MagicMock(side_effect=MyError)) @@ -204,6 +208,7 @@ def test_build_model_uses_custom_model_cls_from_config(mock_factory): mock_factory.build_model(device="meta") +@pytest.mark.cpu_only def test_custom_model_mapping_in_parent_does_not_affect_children(): class Child(AutoModelForCausalLMFactory): pass @@ -217,6 +222,7 @@ class Child(AutoModelForCausalLMFactory): assert Child._custom_model_mapping == {} +@pytest.mark.cpu_only def test_custom_model_mapping_in_parent_does_not_affect_parent(): class Child(AutoModelForCausalLMFactory): pass diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_hunyuan_moe_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_hunyuan_moe_modeling.py index 1c2936716dba..a7a9e8da1397 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_hunyuan_moe_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_hunyuan_moe_modeling.py @@ -727,6 +727,7 @@ def test_model_can_be_exported(device): # =========================================================================== +@pytest.mark.cpu_only def test_config_registration(): """Factory knows the model class under the real HF config name.""" from tensorrt_llm._torch.auto_deploy.models.hf import AutoModelForCausalLMFactory @@ -734,6 +735,7 @@ def test_config_registration(): assert "HunYuanConfig" in AutoModelForCausalLMFactory._custom_model_mapping +@pytest.mark.cpu_only def test_tied_weights(): """tie_word_embeddings: lm_head.weight is the same tensor as embed_tokens.weight.""" cfg = _create_small_custom_config() @@ -741,6 +743,7 @@ def test_tied_weights(): assert model.lm_head.weight is model.model.embed_tokens.weight +@pytest.mark.cpu_only def test_moe_structure(): """Every decoder layer has the expected MoE submodule structure.""" cfg = _create_small_custom_config() @@ -755,6 +758,7 @@ def test_moe_structure(): assert hasattr(moe, "shared_mlp"), f"Layer {i}: missing mlp.shared_mlp" +@pytest.mark.cpu_only def test_qk_norm_structure(): """QK normalization layers exist in each attention module.""" cfg = _create_small_custom_config() @@ -765,6 +769,7 @@ def test_qk_norm_structure(): assert hasattr(attn, "key_layernorm"), f"Layer {i}: missing key_layernorm" +@pytest.mark.cpu_only def test_state_dict_keys_match_checkpoint(): """State dict keys match the HF checkpoint naming convention.""" cfg = _create_small_custom_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_internlm3_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_internlm3_modeling.py index d047281cbe08..cd3a5fe57ff6 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_internlm3_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_internlm3_modeling.py @@ -547,6 +547,7 @@ def test_internlm3_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_internlm3_config_registration(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -559,6 +560,7 @@ def test_internlm3_config_registration(): assert hasattr(config, "qkv_bias") +@pytest.mark.cpu_only def test_internlm3_gqa_structure(): """Test that attention uses GQA (fewer KV heads than Q heads).""" config = _create_small_config() @@ -569,6 +571,7 @@ def test_internlm3_gqa_structure(): assert attn.num_kv_heads == 2, f"Expected 2 KV heads, got {attn.num_kv_heads}" +@pytest.mark.cpu_only def test_internlm3_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_kimi_k2_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_kimi_k2_modeling.py index 02ac938b3603..015bd0fe2471 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_kimi_k2_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_kimi_k2_modeling.py @@ -239,6 +239,7 @@ def test_kimi_k2_text_model_can_be_exported(): # ============================================================================= +@pytest.mark.cpu_only def test_kimi_k2_config_registration(): """Test that configs are properly instantiated with correct model_type.""" text_config = _create_small_text_config() @@ -254,6 +255,7 @@ def test_kimi_k2_config_registration(): assert isinstance(vlm_config.text_config, KimiK2Config) +@pytest.mark.cpu_only def test_kimi_k2_layer_types(): """Test that layer 0 uses dense MLP and later layers use MoE.""" config = _create_small_text_config() @@ -271,6 +273,7 @@ def test_kimi_k2_layer_types(): ) +@pytest.mark.cpu_only def test_kimi_k2_expert_structure(): """Test that experts have correct structure for checkpoint loading.""" config = _create_small_text_config() @@ -298,6 +301,7 @@ def test_kimi_k2_expert_structure(): ) +@pytest.mark.cpu_only def test_kimi_k25_weight_layout(): """Test that VLM wrapper has correct weight prefix for checkpoint compatibility.""" config = _create_small_vlm_config() @@ -313,6 +317,7 @@ def test_kimi_k25_weight_layout(): ) +@pytest.mark.cpu_only def test_kimi_k2_shared_experts(): """Test that shared experts are present when n_shared_experts > 0.""" config = _create_small_text_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_llama3_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_llama3_modeling.py index 675d2c5c1fda..5b75109e2239 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_llama3_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_llama3_modeling.py @@ -437,6 +437,7 @@ def test_llama3_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_llama3_config_registration(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -447,6 +448,7 @@ def test_llama3_config_registration(): assert hasattr(config, "head_dim") +@pytest.mark.cpu_only def test_llama3_gqa_structure(): """Test that attention uses GQA (fewer KV heads than Q heads).""" config = _create_small_config() @@ -461,6 +463,7 @@ def test_llama3_gqa_structure(): assert not hasattr(attn, "k_norm"), "Llama attention should NOT have k_norm" +@pytest.mark.cpu_only def test_llama3_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_llama4_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_llama4_modeling.py index afea246952f1..c4eec3e7f939 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_llama4_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_llama4_modeling.py @@ -738,6 +738,7 @@ def test_llama4_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_llama4_config_registration(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -750,6 +751,7 @@ def test_llama4_config_registration(): assert hasattr(config, "no_rope_layers") +@pytest.mark.cpu_only def test_llama4_gqa_structure(): """Test that attention uses GQA (fewer KV heads than Q heads).""" config = _create_small_config() @@ -760,6 +762,7 @@ def test_llama4_gqa_structure(): assert attn.num_kv_heads == 2, f"Expected 2 KV heads, got {attn.num_kv_heads}" +@pytest.mark.cpu_only def test_llama4_moe_structure(): """Test MoE layer structure.""" config = _create_small_config() @@ -773,6 +776,7 @@ def test_llama4_moe_structure(): assert layer.feed_forward.experts.num_experts == config.num_local_experts +@pytest.mark.cpu_only def test_llama4_nope_layers(): """Test NoPE/RoPE layer configuration.""" config = _create_small_config() @@ -791,6 +795,7 @@ def test_llama4_nope_layers(): ) +@pytest.mark.cpu_only def test_llama4_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_minimax_m2_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_minimax_m2_modeling.py index 891c07cb5edd..74db99eda6f0 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_minimax_m2_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_minimax_m2_modeling.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Hierarchical equivalence tests for MiniMax-M2 AutoDeploy custom model. @@ -10,6 +10,7 @@ from typing import Tuple +import pytest import torch import torch.nn.functional as F from torch import nn @@ -519,6 +520,7 @@ def test_full_model_equivalence(): # --------------------------------------------------------------------------- +@pytest.mark.cpu_only def test_export(): """Model can be exported with torch.export and produces correct output.""" device = "cpu" diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_mistral3_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_mistral3_modeling.py index 856b419c0d41..55cdbd6b69b9 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_mistral3_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_mistral3_modeling.py @@ -488,6 +488,7 @@ def test_moe_equivalence_and_converter(): assert_rmse_close(actual, expected, rmse_ratio_tol=0.02, msg="MoE: ") +@pytest.mark.cpu_only def test_moe_fused_checkpoint_hook_expands_static_fp8_scales(): config = _small_text_config() moe = Mistral4MoE(config) @@ -522,6 +523,7 @@ def test_moe_fused_checkpoint_hook_expands_static_fp8_scales(): assert "experts.down_proj_activation_scale" not in state_dict +@pytest.mark.cpu_only def test_moe_fused_checkpoint_hook_uses_owned_expert_ids(): config = _small_text_config() moe = Mistral4MoE(config) @@ -602,6 +604,7 @@ def test_decoder_layer_equivalence(): assert_rmse_close(actual, expected, rmse_ratio_tol=0.05, msg="Decoder layer: ") +@pytest.mark.cpu_only def test_full_model_equivalence_cpu(): device = "cpu" dtype = torch.float32 @@ -628,6 +631,7 @@ def test_top_level_wrapper_inputs_embeds_path(): torch.testing.assert_close(wrapper_logits, ref_logits, atol=1e-3, rtol=1e-3) +@pytest.mark.cpu_only def test_mistral3_wrappers_do_not_forward_none_inputs_embeds(): class RecordingModule(nn.Module): def __init__(self): @@ -705,6 +709,7 @@ def test_exported_mistral4_graph_has_valid_layer_subgraphs(): assert "MLP" in layer_types +@pytest.mark.cpu_only def test_registration(): assert ( AutoModelForCausalLMFactory._custom_model_mapping["Mistral4TextConfig"] diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_mistral_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_mistral_modeling.py index f55c28090996..5a72f8e40849 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_mistral_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_mistral_modeling.py @@ -484,6 +484,7 @@ def test_mistral_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_mistral_config_registration(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -494,6 +495,7 @@ def test_mistral_config_registration(): assert hasattr(config, "sliding_window") +@pytest.mark.cpu_only def test_mistral_gqa_structure(): """Test that attention uses GQA (fewer KV heads than Q heads).""" config = _create_small_config() @@ -504,6 +506,7 @@ def test_mistral_gqa_structure(): assert attn.num_kv_heads == 2, f"Expected 2 KV heads, got {attn.num_kv_heads}" +@pytest.mark.cpu_only def test_mistral_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_olmo3_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_olmo3_modeling.py index e982a48a207a..dab6f5e01058 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_olmo3_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_olmo3_modeling.py @@ -502,6 +502,7 @@ def test_olmo3_model_can_be_exported(rope_type): # ========================================================================= +@pytest.mark.cpu_only def test_olmo3_config_registration(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -512,6 +513,7 @@ def test_olmo3_config_registration(): assert hasattr(config, "layer_types") +@pytest.mark.cpu_only def test_olmo3_post_norm_structure(): """Test that OLMo-3 uses post-norm (not pre-norm).""" config = _create_small_config() @@ -527,6 +529,7 @@ def test_olmo3_post_norm_structure(): ) +@pytest.mark.cpu_only def test_olmo3_qk_norm_structure(): """Test that attention uses QK normalization on full projection.""" config = _create_small_config() @@ -542,6 +545,7 @@ def test_olmo3_qk_norm_structure(): assert attn.k_norm.weight.shape[0] == config.num_key_value_heads * head_dim +@pytest.mark.cpu_only def test_olmo3_mixed_attention_types(): """Test that model correctly assigns sliding/full attention types.""" config = _create_small_config() @@ -559,6 +563,7 @@ def test_olmo3_mixed_attention_types(): assert layer.self_attn.sliding_window is None +@pytest.mark.cpu_only def test_olmo3_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_openelm_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_openelm_modeling.py index cb4cca660fed..fe1259c177b9 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_openelm_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_openelm_modeling.py @@ -519,6 +519,7 @@ def test_openelm_model_can_be_exported(): # fmt: on +@pytest.mark.cpu_only def test_openelm_config_derivation_matches_apple(): # `use_cache=True` would crash Apple's remote config under transformers 5.x; # the vendored class must accept it without raising. @@ -534,6 +535,7 @@ def test_openelm_config_derivation_matches_apple(): assert all(q % k == 0 for q, k in zip(config.num_query_heads, config.num_kv_heads, strict=True)) +@pytest.mark.cpu_only def test_openelm_config_derivation_matches_apple_3b(): """Derivation must also be correct for a different size point (3B, head_dim=128).""" config = OpenELMConfig(use_cache=True, **_OPENELM_3B) @@ -547,6 +549,7 @@ def test_openelm_config_derivation_matches_apple_3b(): assert all(q % k == 0 for q, k in zip(config.num_query_heads, config.num_kv_heads, strict=True)) +@pytest.mark.cpu_only def test_openelm_config_fulfills_config_subclass_contract(): """Pin the two duties of a PreTrainedConfig subclass. @@ -585,6 +588,7 @@ def test_openelm_config_fulfills_config_subclass_contract(): assert len(config.num_query_heads) == 16 +@pytest.mark.cpu_only def test_openelm_config_registered_as_local_class(): assert "openelm" in CONFIG_MAPPING registered = CONFIG_MAPPING["openelm"] @@ -594,6 +598,7 @@ def test_openelm_config_registered_as_local_class(): assert not registered.__module__.startswith("transformers.") +@pytest.mark.cpu_only def test_openelm_autoconfig_prefers_local_over_remote(tmp_path): """AutoConfig must prefer the local class over remote code. diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_qwen3_5_moe.py b/tests/unittest/auto_deploy/singlegpu/models/test_qwen3_5_moe.py index 0b87c14f258a..004561fc205c 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_qwen3_5_moe.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_qwen3_5_moe.py @@ -1685,6 +1685,7 @@ def __call__(self, inputs, sampling_params): return self._token_ids, {} +@pytest.mark.cpu_only def test_qwen_ad_input_processor_duplicates_video_hashes_per_frame_span(): config = _make_small_composite_config() merge = config.vision_config.spatial_merge_size @@ -1714,6 +1715,7 @@ def test_qwen_ad_input_processor_duplicates_video_hashes_per_frame_span(): assert multimodal_input.multimodal_hashes[0] == multimodal_input.multimodal_hashes[1] +@pytest.mark.cpu_only def test_qwen_ad_input_processor_preserves_video_item_order_across_frame_spans(): config = _make_small_composite_config() merge = config.vision_config.spatial_merge_size diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_qwen3_moe_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_qwen3_moe_modeling.py index ef4b71db09b0..953660bff599 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_qwen3_moe_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_qwen3_moe_modeling.py @@ -547,6 +547,7 @@ def test_qwen3_moe_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_qwen3_moe_config_registration(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -558,6 +559,7 @@ def test_qwen3_moe_config_registration(): assert hasattr(config, "num_experts_per_tok") +@pytest.mark.cpu_only def test_qwen3_moe_structure(): """Test model structure: GQA, Q/K norms, MoE layers.""" config = _create_small_config() @@ -577,6 +579,7 @@ def test_qwen3_moe_structure(): ) +@pytest.mark.cpu_only def test_qwen3_moe_dense_layer_config(): """Test that mlp_only_layers correctly creates dense MLP layers.""" config = _create_small_config() @@ -594,6 +597,7 @@ def test_qwen3_moe_dense_layer_config(): ) +@pytest.mark.cpu_only def test_qwen3_moe_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_qwen3_next_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_qwen3_next_modeling.py index fc910843b61c..5aeed58a96b5 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_qwen3_next_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_qwen3_next_modeling.py @@ -403,6 +403,7 @@ def test_qwen3_next_gdn_patch_float32(): torch.testing.assert_close(ref_output, test_output, atol=1e-4, rtol=1e-4) +@pytest.mark.cpu_only def test_qwen3_next_layer_types(): """Test that layers have correct types (linear vs full attention) and MoE.""" config = _create_small_config() @@ -422,6 +423,7 @@ def test_qwen3_next_layer_types(): ) +@pytest.mark.cpu_only def test_qwen3_next_tied_embeddings(): """Test that embeddings are tied when tie_word_embeddings=True.""" config = _create_small_config() @@ -430,6 +432,7 @@ def test_qwen3_next_tied_embeddings(): assert model.lm_head.weight is model.model.embed_tokens.weight +@pytest.mark.cpu_only def test_qwen3_next_expert_structure(): """Test that experts have correct structure for checkpoint loading.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_seed_oss_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_seed_oss_modeling.py index ae94dcda3434..dbf9dd20ead0 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_seed_oss_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_seed_oss_modeling.py @@ -447,6 +447,7 @@ def test_seed_oss_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_seed_oss_config_registration(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -459,6 +460,7 @@ def test_seed_oss_config_registration(): assert hasattr(config, "attention_out_bias") +@pytest.mark.cpu_only def test_seed_oss_gqa_structure(): """Test that attention uses GQA (fewer KV heads than Q heads).""" config = _create_small_config() @@ -476,6 +478,7 @@ def test_seed_oss_gqa_structure(): assert attn.o_proj.bias is None, "O projection should not have bias" +@pytest.mark.cpu_only def test_seed_oss_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_skywork_r1v2_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_skywork_r1v2_modeling.py index f5a32f889747..23e994a97f02 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_skywork_r1v2_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_skywork_r1v2_modeling.py @@ -423,6 +423,7 @@ def test_skywork_r1v2_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_skywork_r1v2_config_parsing(): """Test that SkyworkChatConfig correctly wraps the llm_config as a Qwen2Config.""" config = _create_small_chat_config() @@ -433,6 +434,7 @@ def test_skywork_r1v2_config_parsing(): assert config.llm_config.num_key_value_heads == 2 +@pytest.mark.cpu_only def test_skywork_r1v2_gqa_structure(): """Test that attention uses GQA with bias on QKV.""" model = SkyworkR1V2ForConditionalGeneration(_create_small_chat_config()) @@ -446,6 +448,7 @@ def test_skywork_r1v2_gqa_structure(): assert attn.o_proj.bias is None +@pytest.mark.cpu_only def test_skywork_r1v2_state_dict_keys(): """Test that state_dict keys match expected checkpoint format. diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_smollm3_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_smollm3_modeling.py index 54b278677f20..50f3fa439798 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_smollm3_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_smollm3_modeling.py @@ -450,6 +450,7 @@ def test_smollm3_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_smollm3_config_registration(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -460,6 +461,7 @@ def test_smollm3_config_registration(): assert hasattr(config, "no_rope_layers") +@pytest.mark.cpu_only def test_smollm3_nope_layer_pattern(): """Test that NoPE layers are correctly configured.""" config = _create_small_config() @@ -473,6 +475,7 @@ def test_smollm3_nope_layer_pattern(): ) +@pytest.mark.cpu_only def test_smollm3_gqa_structure(): """Test that attention uses GQA (fewer KV heads than Q heads).""" config = _create_small_config() @@ -483,6 +486,7 @@ def test_smollm3_gqa_structure(): assert attn.num_kv_heads == 2, f"Expected 2 KV heads, got {attn.num_kv_heads}" +@pytest.mark.cpu_only def test_smollm3_state_dict_keys(): """Test that state_dict keys match expected checkpoint format.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_starcoder2_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_starcoder2_modeling.py index 5f3f9a03e0b4..a0543c93478e 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_starcoder2_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_starcoder2_modeling.py @@ -447,6 +447,7 @@ def test_starcoder2_model_can_be_exported(): # ========================================================================= +@pytest.mark.cpu_only def test_starcoder2_config_recognition(): """Test that the config is properly recognized.""" config = _create_small_config() @@ -457,6 +458,7 @@ def test_starcoder2_config_recognition(): assert hasattr(config, "sliding_window") +@pytest.mark.cpu_only def test_starcoder2_gqa_structure(): """Test that attention uses GQA (fewer KV heads than Q heads).""" config = _create_small_config() @@ -468,6 +470,7 @@ def test_starcoder2_gqa_structure(): assert attn.num_kv_heads < attn.num_heads, "Should use GQA" +@pytest.mark.cpu_only def test_starcoder2_state_dict_keys(): """Test that state_dict keys match expected HF checkpoint format.""" config = _create_small_config() diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_step3p7_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_step3p7_modeling.py index 8688282ca0e8..b173707192cf 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_step3p7_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_step3p7_modeling.py @@ -606,6 +606,7 @@ def test_full_model_equivalence(): # --------------------------------------------------------------------------- +@pytest.mark.cpu_only def test_export(): device = "cpu" dtype = torch.float32 diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py b/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py index 65194a0c6471..afeef5471baf 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import contextmanager from dataclasses import dataclass from types import SimpleNamespace from typing import Any, Optional @@ -26,6 +27,33 @@ from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import AttentionTypeCpp from tensorrt_llm.llmapi import CacheTransceiverConfig +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture(autouse=True) +def _mock_autodeploy_runtime_dependencies(): + mock_dist = Mock() + mock_dist.broadcast.side_effect = lambda value: value + mock_autotuner = Mock() + + with ( + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.Distributed.get", + return_value=mock_dist, + ), + patch("tensorrt_llm._torch.auto_deploy.shim.ad_executor.torch.cuda.set_device"), + patch("tensorrt_llm._torch.auto_deploy.shim.ad_executor.initialize_or_skip"), + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.AutoTuner.get", + return_value=mock_autotuner, + ), + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.instantiate_sampler", + return_value=Mock(), + ), + ): + yield + class MockTokenizer: """Simple mock tokenizer.""" @@ -105,6 +133,31 @@ def make_mock_engine( return mock_engine, kv_cache_manager +@contextmanager +def _mock_ad_engine_build(mock_engine, *, vocab_size_padded: int = 1000): + with ( + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config", + return_value=mock_engine, + ) as mock_ad_engine, + patch( + "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", + return_value=MockFactory(vocab_size_padded=vocab_size_padded), + ), + ): + yield mock_ad_engine + + +@contextmanager +def _mock_py_executor_creation(mock_engine, *, vocab_size_padded: int = 1000): + with ( + _mock_ad_engine_build(mock_engine, vocab_size_padded=vocab_size_padded), + patch("tensorrt_llm._torch.auto_deploy.shim.ad_executor.PyExecutor") as py_executor_cls, + ): + py_executor_cls.side_effect = MockPyExecutor + yield py_executor_cls + + @pytest.mark.parametrize("guided_decoding_backend", ["xgrammar", "llguidance"]) @pytest.mark.parametrize("max_batch_size", [4, 8]) @pytest.mark.parametrize("vocab_size_padded", [42, 1000]) @@ -133,7 +186,6 @@ def test_create_autodeploy_executor_with_guided_decoding( max_batch_size=max_batch_size, vocab_size_padded=vocab_size_padded ) - # Mock the specific dependencies requested, plus minimal additional mocks to prevent errors with ( patch( "tensorrt_llm._torch.auto_deploy.shim.ad_executor.get_guided_decoding_config", @@ -142,20 +194,11 @@ def test_create_autodeploy_executor_with_guided_decoding( patch( "tensorrt_llm._torch.auto_deploy.shim.ad_executor.GuidedDecoder" ) as guided_decoder_cls, - patch("tensorrt_llm._torch.auto_deploy.shim.ad_executor.PyExecutor") as py_executor_cls, - patch( - "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config" - ) as mock_ad_engine, - patch( - "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", - return_value=MockFactory(vocab_size_padded=vocab_size_padded), - ), + _mock_py_executor_creation( + mock_engine, vocab_size_padded=vocab_size_padded + ) as py_executor_cls, ): - mock_ad_engine.return_value = mock_engine - - # substitute the GuidedDecoder and PyExecutor classes guided_decoder_cls.side_effect = MockGuidedDecoder - py_executor_cls.side_effect = MockPyExecutor # Call the function under test result = create_autodeploy_executor(ad_config, mock_tokenizer) @@ -203,24 +246,15 @@ def test_create_executor_uses_cache_transceiver(cache_attention_type, expected_a mock_engine, kv_cache_manager = make_mock_engine(attention_type=cache_attention_type) with ( - patch("tensorrt_llm._torch.auto_deploy.shim.ad_executor.PyExecutor") as py_executor_cls, - patch( - "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config" - ) as mock_ad_engine, + _mock_py_executor_creation(mock_engine) as py_executor_cls, patch( "tensorrt_llm._torch.auto_deploy.shim.ad_executor.create_kv_cache_transceiver", return_value=mock_transceiver, ) as create_transceiver, - patch( - "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", - return_value=MockFactory(vocab_size_padded=1000), - ), ): - mock_ad_engine.return_value = mock_engine - py_executor_cls.side_effect = MockPyExecutor - result = create_autodeploy_executor(ad_config, mock_tokenizer) + py_executor_cls.assert_called_once() create_transceiver.assert_called_once() _, _, passed_kv_cache_manager, attention_type, passed_config = create_transceiver.call_args.args assert passed_kv_cache_manager is kv_cache_manager @@ -260,22 +294,12 @@ def test_create_executor_preserves_explicit_transceiver_buffer_size( mock_engine, _ = make_mock_engine(attention_type=cache_attention_type) with ( - patch("tensorrt_llm._torch.auto_deploy.shim.ad_executor.PyExecutor") as py_executor_cls, - patch( - "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config" - ) as mock_ad_engine, + _mock_py_executor_creation(mock_engine), patch( "tensorrt_llm._torch.auto_deploy.shim.ad_executor.create_kv_cache_transceiver", return_value=mock_transceiver, ) as create_transceiver, - patch( - "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", - return_value=MockFactory(vocab_size_padded=1000), - ), ): - mock_ad_engine.return_value = mock_engine - py_executor_cls.side_effect = MockPyExecutor - create_autodeploy_executor(ad_config, mock_tokenizer) _, _, _, attention_type, passed_config = create_transceiver.call_args.args @@ -302,17 +326,7 @@ def test_create_executor_rejects_non_enum_attention_type(cache_attention_type): mock_engine, _ = make_mock_engine() mock_engine.cache_seq_interface.attention_type = cache_attention_type - with ( - patch( - "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config" - ) as mock_ad_engine, - patch( - "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", - return_value=MockFactory(vocab_size_padded=1000), - ), - ): - mock_ad_engine.return_value = mock_engine - + with _mock_ad_engine_build(mock_engine): with pytest.raises(TypeError): create_autodeploy_executor(ad_config, mock_tokenizer) @@ -333,17 +347,7 @@ def test_create_executor_requires_attention_type(): mock_engine, _ = make_mock_engine(attention_type=None) - with ( - patch( - "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config" - ) as mock_ad_engine, - patch( - "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", - return_value=MockFactory(vocab_size_padded=1000), - ), - ): - mock_ad_engine.return_value = mock_engine - + with _mock_ad_engine_build(mock_engine): with pytest.raises(RuntimeError): create_autodeploy_executor(ad_config, mock_tokenizer) @@ -370,19 +374,11 @@ def test_create_executor_rejects_mamba_cache_manager_for_transceiver(): mock_engine.cache_seq_interface.kv_cache_manager = mamba_cache_manager with ( - patch( - "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config" - ) as mock_ad_engine, + _mock_ad_engine_build(mock_engine), patch( "tensorrt_llm._torch.auto_deploy.shim.ad_executor.create_kv_cache_transceiver" ) as create_transceiver, - patch( - "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", - return_value=MockFactory(vocab_size_padded=1000), - ), ): - mock_ad_engine.return_value = mock_engine - with pytest.raises(RuntimeError): create_autodeploy_executor(ad_config, mock_tokenizer) diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py b/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py index 5639344dd18c..4575a80fb33f 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py @@ -19,6 +19,18 @@ from tensorrt_llm._torch.auto_deploy import LLM, DemoLLM, LlmArgs +pytestmark = pytest.mark.cpu_only + + +EMULATED_DEVICE_COUNT = 3 + + +@pytest.fixture +def patch_device_count(): + # Make the default value of LlmArgs.gpus_per_node solvable without probing hardware + with patch("tensorrt_llm.llmapi.llm.get_device_count", return_value=EMULATED_DEVICE_COUNT): + yield + def test_custom_values(): """Test that AutoDeploy LlmArgs correctly accepts custom values.""" @@ -121,7 +133,7 @@ def test_config_params(): "skip_loading_weights": True, "max_seq_len": 19, "max_batch_size": 128, - "world_size": 3, + "world_size": EMULATED_DEVICE_COUNT, "transforms": { "detect_sharding": { "stage": "sharding", @@ -163,6 +175,7 @@ def test_config_flow( extra_kwargs, expected_executor_call, test_config_params, + patch_device_count, ): """Test that config flows correctly through both DemoLLM and LLM initialization.""" # Mock the executor and its methods for DemoLLM @@ -174,7 +187,7 @@ def test_config_flow( mock_seq_info.return_value = mock_seq_info_instance # Merge extra kwargs for the specific API - config_params = {**test_config_params, **extra_kwargs} + config_params = {**test_config_params, **extra_kwargs, "gpus_per_node": EMULATED_DEVICE_COUNT} # Create instance with appropriate mocking with patch.object(api_class, "_try_load_tokenizer", return_value=MagicMock()): diff --git a/tests/unittest/auto_deploy/singlegpu/test_graph_canonicalize.py b/tests/unittest/auto_deploy/singlegpu/test_graph_canonicalize.py index bc36b77eb65a..e9839f776a28 100644 --- a/tests/unittest/auto_deploy/singlegpu/test_graph_canonicalize.py +++ b/tests/unittest/auto_deploy/singlegpu/test_graph_canonicalize.py @@ -10,6 +10,8 @@ from tensorrt_llm._torch.auto_deploy.utils._graph import canonicalize_graph +pytestmark = pytest.mark.cpu_only + class _OutOfOrderCastModule(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: diff --git a/tests/unittest/auto_deploy/singlegpu/test_hf_export_info.py b/tests/unittest/auto_deploy/singlegpu/test_hf_export_info.py index 9c0c2c9949ee..3fbe144dbadc 100644 --- a/tests/unittest/auto_deploy/singlegpu/test_hf_export_info.py +++ b/tests/unittest/auto_deploy/singlegpu/test_hf_export_info.py @@ -3,12 +3,15 @@ import operator +import pytest import torch from torch import nn from torch.fx import symbolic_trace from tensorrt_llm._torch.auto_deploy.models.hf import TextModelExportInfo +pytestmark = pytest.mark.cpu_only + class _DummyTextModel(nn.Module): def __init__(self): diff --git a/tests/unittest/auto_deploy/singlegpu/test_mistral_small_4_tokenizer_bridge.py b/tests/unittest/auto_deploy/singlegpu/test_mistral_small_4_tokenizer_bridge.py index 8b0da7efed9e..d11b1e1b6ab2 100644 --- a/tests/unittest/auto_deploy/singlegpu/test_mistral_small_4_tokenizer_bridge.py +++ b/tests/unittest/auto_deploy/singlegpu/test_mistral_small_4_tokenizer_bridge.py @@ -17,6 +17,9 @@ import pytest +pytestmark = pytest.mark.cpu_only + + _SOURCE_MODEL = "mistralai/Mistral-Small-4-119B-2603" diff --git a/tests/unittest/auto_deploy/singlegpu/test_pattern_matcher.py b/tests/unittest/auto_deploy/singlegpu/test_pattern_matcher.py index 9b96182079d2..d0c77b5125a9 100644 --- a/tests/unittest/auto_deploy/singlegpu/test_pattern_matcher.py +++ b/tests/unittest/auto_deploy/singlegpu/test_pattern_matcher.py @@ -12,6 +12,8 @@ register_ad_pattern, ) +pytestmark = pytest.mark.cpu_only + class _LateCastMultiOutputModule(nn.Module): def forward( diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_bmm_moe_hooks.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_bmm_moe_hooks.py index 4439659a355c..91975b9852cf 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_bmm_moe_hooks.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_bmm_moe_hooks.py @@ -22,6 +22,8 @@ _bmm_moe_gate_up_split_hook, ) +pytestmark = pytest.mark.cpu_only + @pytest.fixture def gate_up_stacked_weight(): diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_silu_mul.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_silu_mul.py index 2d723821e82c..870a3e9a35f6 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_silu_mul.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_silu_mul.py @@ -15,12 +15,16 @@ import operator +import pytest import torch import tensorrt_llm._torch.auto_deploy.custom_ops.linear.silu_mul # noqa: F401 from tensorrt_llm._torch.auto_deploy.transform.interface import SharedConfig, TransformRegistry from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op +pytestmark = pytest.mark.cpu_only + + _HALF = 256 diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_swiglu.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_swiglu.py index f272d05d6e55..1662aaef7c39 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_swiglu.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_swiglu.py @@ -220,6 +220,7 @@ def test_swiglu_pattern_match_only(): # --------------------------------------------------------------------------- +@pytest.mark.cpu_only def test_collect_classification_hints_consensus_over_mixed_mechanics(): """Consensus recovers the shared layer_type even when constituents disagree on tp_mode.""" g = torch.fx.Graph() @@ -240,6 +241,7 @@ def test_collect_classification_hints_consensus_over_mixed_mechanics(): assert collect_classification_hints(nodes) == {"layer_type": "shared_expert"} +@pytest.mark.cpu_only def test_collect_classification_hints_conflict_is_dropped(): """Conflicting layer_type values (a rewrite spanning layers) are dropped, not guessed.""" g = torch.fx.Graph() @@ -257,6 +259,7 @@ def test_collect_classification_hints_conflict_is_dropped(): assert collect_classification_hints([a, b]) == {} +@pytest.mark.cpu_only def test_stamp_hints_only_on_declaring_ops(): """stamp_hints sets a hint only on ops whose schema declares it (aten.silu is skipped).""" g = torch.fx.Graph() diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_gemm_fusion.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_gemm_fusion.py index c391f5001dee..4bb92e00a2cf 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_gemm_fusion.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_gemm_fusion.py @@ -36,6 +36,7 @@ torch.manual_seed(0) +@pytest.mark.cpu_only def test_fuse_gemms_skips_linear_without_parameter_weight(): graph = torch.fx.Graph() x = graph.placeholder("x") diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py index 4c39289745e1..7fad68093622 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py @@ -28,6 +28,8 @@ InsertCachedAttentionConfig, ) +pytestmark = pytest.mark.cpu_only + class _TwoWindowModule(torch.nn.Module): """A tiny two-layer model with one SWA layer and one full-attention layer. diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_mrope_delta_cache.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_mrope_delta_cache.py index 5930e9d6e0f3..1177505f94de 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_mrope_delta_cache.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_mrope_delta_cache.py @@ -14,6 +14,7 @@ # limitations under the License. from pathlib import Path +import pytest import torch import torch.nn as nn import yaml @@ -25,6 +26,8 @@ InitializeMropeDeltaCache, ) +pytestmark = pytest.mark.cpu_only + def _repo_root() -> Path: return Path(__file__).resolve().parents[6] diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py index a80c3eaaa386..53f11b4dfed3 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py @@ -744,6 +744,7 @@ def test_fuse_rmsnorm_quant_fp8_rewrites_through_post_norm_reshape(): assert not any(is_op(n, torch.ops.auto_deploy.trtllm_quant_fp8_linear) for n in gm.graph.nodes) +@pytest.mark.cpu_only def test_get_out_dtype_str_returns_none_when_norm_meta_missing(): graph = torch.fx.Graph() x = graph.placeholder("x") @@ -757,6 +758,7 @@ def test_get_out_dtype_str_returns_none_when_norm_meta_missing(): assert _get_out_dtype_str(norm) is None +@pytest.mark.cpu_only def test_passthrough_helpers_handle_method_views_and_optional_dtype_cast(): graph = torch.fx.Graph() x = graph.placeholder("x") @@ -809,6 +811,7 @@ def _get_fused_getitem(gm, fused_op, index): return matches[0] +@pytest.mark.cpu_only def test_fuse_gated_rmsnorm_quant_nvfp4_rewrites_graph(): root = _make_nvfp4_graph_root() graph = torch.fx.Graph() @@ -856,6 +859,7 @@ def test_fuse_gated_rmsnorm_quant_nvfp4_rewrites_graph(): assert scale_node.meta["val"].dtype == torch.uint8 +@pytest.mark.cpu_only def test_fuse_gated_rmsnorm_quant_nvfp4_accepts_dtype_cast(): root = _make_nvfp4_graph_root() graph = torch.fx.Graph() @@ -895,6 +899,7 @@ def test_fuse_gated_rmsnorm_quant_nvfp4_accepts_dtype_cast(): assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) +@pytest.mark.cpu_only def test_fuse_gated_rmsnorm_quant_nvfp4_preserves_mixed_consumer_dtypes(): root = _make_nvfp4_graph_root() graph = torch.fx.Graph() @@ -940,6 +945,7 @@ def test_fuse_gated_rmsnorm_quant_nvfp4_preserves_mixed_consumer_dtypes(): assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) +@pytest.mark.cpu_only def test_fuse_allreduce_rmsnorm_quant_nvfp4_rewrites_graph(): root = _make_nvfp4_graph_root() graph = torch.fx.Graph() @@ -982,6 +988,7 @@ def test_fuse_allreduce_rmsnorm_quant_nvfp4_rewrites_graph(): assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) +@pytest.mark.cpu_only def test_fuse_allreduce_rmsnorm_quant_nvfp4_keeps_norm_for_mixed_consumers(): root = _make_nvfp4_graph_root() graph = torch.fx.Graph() @@ -1024,6 +1031,7 @@ def test_fuse_allreduce_rmsnorm_quant_nvfp4_keeps_norm_for_mixed_consumers(): assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) +@pytest.mark.cpu_only def test_fuse_allreduce_rmsnorm_quant_nvfp4_clones_late_input_scale(): root = _make_nvfp4_graph_root() graph = torch.fx.Graph() @@ -1064,6 +1072,7 @@ def test_fuse_allreduce_rmsnorm_quant_nvfp4_clones_late_input_scale(): ) +@pytest.mark.cpu_only def test_fuse_add_rmsnorm_quant_nvfp4_rewrites_graph(): hidden_size = 2048 root = _make_nvfp4_graph_root(hidden_size) @@ -1104,6 +1113,7 @@ def test_fuse_add_rmsnorm_quant_nvfp4_rewrites_graph(): assert not any(is_op(n, torch.ops.aten.add.Tensor) for n in gm.graph.nodes) +@pytest.mark.cpu_only def test_fuse_add_cast_rmsnorm_quant_nvfp4_rewrites_graph(): hidden_size = 2048 root = _make_nvfp4_graph_root(hidden_size) @@ -1148,6 +1158,7 @@ def test_fuse_add_cast_rmsnorm_quant_nvfp4_rewrites_graph(): assert not any(is_op(n, torch.ops.aten.add.Tensor) for n in gm.graph.nodes) +@pytest.mark.cpu_only def test_fuse_add_cast_rmsnorm_quant_nvfp4_clones_late_norm_weight(): hidden_size = 2048 root = _make_nvfp4_graph_root(hidden_size) @@ -1187,6 +1198,7 @@ def test_fuse_add_cast_rmsnorm_quant_nvfp4_clones_late_norm_weight(): ) +@pytest.mark.cpu_only def test_fuse_add_rmsnorm_quant_nvfp4_keeps_norm_for_mixed_consumers(): hidden_size = 2048 root = _make_nvfp4_graph_root(hidden_size) diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quantization.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quantization.py index 4c453b50b917..e9aac5964479 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quantization.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quantization.py @@ -62,6 +62,7 @@ def max_seq_len(self) -> int: return 512 +@pytest.mark.cpu_only def test_quantization_base_does_not_register_noop_post_load_hook(): assert FP8LinearQuantizationFromConfig.post_load_hook is Quantization.post_load_hook assert FP8BMMQuantizationFromConfig.post_load_hook is not Quantization.post_load_hook diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_shared_kv_attention.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_shared_kv_attention.py index f915fc233945..0edf392e48e9 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_shared_kv_attention.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_shared_kv_attention.py @@ -152,6 +152,7 @@ def _make_layer_inputs(offset: float, seq_len: int, decode: bool = False): return q, k, v +@pytest.mark.cpu_only def test_shared_kv_transform_aliases_source_cache_placeholders(): module = _TinySharedKVModule().eval() gm = torch_export_to_gm(module, (torch.randn(1, 4, 8),)) @@ -198,6 +199,7 @@ def test_shared_kv_transform_aliases_source_cache_placeholders(): assert shared_node.args[-1] is True +@pytest.mark.cpu_only def test_shared_kv_cached_attention_reads_without_writing(): q = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]], dtype=torch.float32) dummy_k = torch.full((1, 1, 2, 2), 123.0, dtype=torch.float32) @@ -245,6 +247,7 @@ def test_shared_kv_cached_attention_reads_without_writing(): torch.testing.assert_close(output, expected, rtol=1e-5, atol=1e-5) +@pytest.mark.cpu_only def test_torch_backend_attention_metadata_for_shared_kv_node(): module = _TinySharedKVModule().eval() gm = torch_export_to_gm(module, (torch.randn(1, 4, 8),)) @@ -267,6 +270,7 @@ def test_torch_backend_attention_metadata_for_shared_kv_node(): assert TorchBackendAttention.get_shared_kv_source_layer_idx(shared) == 0 +@pytest.mark.cpu_only def test_flashinfer_backend_attention_metadata_for_shared_kv_node(): module = _TinySharedKVModule().eval() gm = torch_export_to_gm(module, (torch.randn(1, 4, 8),)) @@ -292,6 +296,7 @@ def test_flashinfer_backend_attention_metadata_for_shared_kv_node(): ) +@pytest.mark.cpu_only def test_shared_kv_transform_aliases_source_cache_placeholders_for_flashinfer(): module = _TinySharedKVModule().eval() gm = torch_export_to_gm(module, (torch.randn(1, 4, 8),)) @@ -335,6 +340,7 @@ def test_shared_kv_transform_aliases_source_cache_placeholders_for_flashinfer(): assert shared_node.args[-1] is True +@pytest.mark.cpu_only def test_flashinfer_cached_attention_is_dynamic_for_piecewise(): shared_op_name = torch.ops.auto_deploy.flashinfer_attention_mha_with_cache.default.name() @@ -350,6 +356,7 @@ def __init__(self, target): ) +@pytest.mark.cpu_only def test_triton_backend_attention_metadata_for_shared_kv_node(): module = _TinySharedKVModule().eval() gm = torch_export_to_gm(module, (torch.randn(1, 4, 8),)) @@ -375,6 +382,7 @@ def test_triton_backend_attention_metadata_for_shared_kv_node(): ) +@pytest.mark.cpu_only def test_shared_kv_transform_aliases_source_cache_placeholders_for_triton(): module = _TinySharedKVModule().eval() gm = torch_export_to_gm(module, (torch.randn(1, 4, 8),)) @@ -610,6 +618,7 @@ def tracking_sdpa(*args, **kwargs): torch.testing.assert_close(output.float(), expected.float(), rtol=1e-2, atol=1e-2) +@pytest.mark.cpu_only @torch.no_grad() def test_torch_shared_kv_cached_attention_supports_out_buffer(): q = torch.randn(1, 3, 2, 4) @@ -660,6 +669,7 @@ def test_torch_shared_kv_cached_attention_supports_out_buffer(): torch.testing.assert_close(out, expected) +@pytest.mark.cpu_only def test_shared_kv_self_alias_raises(): class _SelfAliasingSharedKVModule(torch.nn.Module): def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -685,6 +695,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: transform._apply(gm, cm, factory=None, shared_config=SharedConfig()) +@pytest.mark.cpu_only def test_duplicate_cache_owner_layer_idx_raises(): module = _DuplicateLayerOwnerSharedKVModule().eval() gm = torch_export_to_gm(module, (torch.randn(1, 4, 8),)) @@ -770,6 +781,7 @@ def test_flashinfer_shared_kv_cached_attention_reads_aliased_cache_without_writi torch.testing.assert_close(output.float(), expected.float(), rtol=2e-2, atol=2e-2) +@pytest.mark.cpu_only def test_shared_kv_six_layer_stack_matches_reference_for_prefill_and_decode(): layer_sources = {4: 2, 5: 3} sliding_layers = {2, 4} diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/test_export.py b/tests/unittest/auto_deploy/singlegpu/transformations/test_export.py index bb5855d49d2b..2e33f9d28869 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/test_export.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/test_export.py @@ -54,6 +54,7 @@ def get_dynamic_shapes(self): pass +@pytest.mark.cpu_only def test_torch_export_to_gm_strips_input_constraint_pre_hook(): gm = torch_export_to_gm(AddOne(), args=(torch.ones(2, 4),)) @@ -231,6 +232,7 @@ def test_module_export(f_export, mod_cls, device, use_dynamic_shape): print(ep.graph) +@pytest.mark.cpu_only @pytest.mark.parametrize("model_cls", [MLPDuplicate]) @pytest.mark.parametrize("device_export", ["cpu"]) # TODO: investigate meta device error def test_deduplicate_during_export(model_cls: Type[nn.Module], device_export: str): diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/test_pipeline_cache.py b/tests/unittest/auto_deploy/singlegpu/transformations/test_pipeline_cache.py index a25583209da5..bb6405d5230f 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/test_pipeline_cache.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/test_pipeline_cache.py @@ -96,6 +96,9 @@ ) from tensorrt_llm._torch.auto_deploy.utils.pipeline_cache_hooks import mark_pipeline_cache_hook +pytestmark = pytest.mark.cpu_only + + _COUNTERS = { "build": 0, "boundary": 0, diff --git a/tests/unittest/auto_deploy/singlegpu/utils/test_config.py b/tests/unittest/auto_deploy/singlegpu/utils/test_config.py index 796fec34ab13..7e3fc30d0bbf 100644 --- a/tests/unittest/auto_deploy/singlegpu/utils/test_config.py +++ b/tests/unittest/auto_deploy/singlegpu/utils/test_config.py @@ -26,6 +26,8 @@ from tensorrt_llm._torch.auto_deploy.utils._config import DynamicYamlMixInForSettings +pytestmark = pytest.mark.cpu_only + class SimpleModel(BaseModel): """Simple model for testing.""" diff --git a/tests/unittest/auto_deploy/singlegpu/utils/test_create_derived_custom_op.py b/tests/unittest/auto_deploy/singlegpu/utils/test_create_derived_custom_op.py index faf3cd9a0cf8..2774cc137c93 100644 --- a/tests/unittest/auto_deploy/singlegpu/utils/test_create_derived_custom_op.py +++ b/tests/unittest/auto_deploy/singlegpu/utils/test_create_derived_custom_op.py @@ -14,11 +14,15 @@ # limitations under the License. """Tests for ``create_derived_custom_op`` in ``_graph.py``.""" +import pytest import torch from torch._subclasses import FakeTensorMode from tensorrt_llm._torch.auto_deploy.utils._graph import create_derived_custom_op +pytestmark = pytest.mark.cpu_only + + # --------------------------------------------------------------------------- # Helpers – tiny custom ops used as base ops for the tests # --------------------------------------------------------------------------- diff --git a/tests/unittest/auto_deploy/singlegpu/utils/test_dist_config.py b/tests/unittest/auto_deploy/singlegpu/utils/test_dist_config.py index c8788ddc7443..91a961c58059 100644 --- a/tests/unittest/auto_deploy/singlegpu/utils/test_dist_config.py +++ b/tests/unittest/auto_deploy/singlegpu/utils/test_dist_config.py @@ -19,6 +19,9 @@ from tensorrt_llm._torch.auto_deploy.utils.dist_config import DistConfig +pytestmark = pytest.mark.cpu_only + + try: from tensorrt_llm.mapping import Mapping diff --git a/tests/unittest/auto_deploy/singlegpu/utils/test_example_configs.py b/tests/unittest/auto_deploy/singlegpu/utils/test_example_configs.py index a375f755f2a8..9fdccfd77fbf 100644 --- a/tests/unittest/auto_deploy/singlegpu/utils/test_example_configs.py +++ b/tests/unittest/auto_deploy/singlegpu/utils/test_example_configs.py @@ -27,6 +27,9 @@ from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs +pytestmark = pytest.mark.cpu_only + + # Root directory for the example configs _REPO_ROOT = pathlib.Path(__file__).resolve().parents[5] _AD_EXAMPLES_DIR = _REPO_ROOT / "examples" / "auto_deploy" diff --git a/tests/unittest/auto_deploy/singlegpu/utils/test_multi_stream_utils.py b/tests/unittest/auto_deploy/singlegpu/utils/test_multi_stream_utils.py index f2e62dec50b4..95c063d61ce3 100644 --- a/tests/unittest/auto_deploy/singlegpu/utils/test_multi_stream_utils.py +++ b/tests/unittest/auto_deploy/singlegpu/utils/test_multi_stream_utils.py @@ -45,6 +45,7 @@ def get_event(self, device, name): raise AssertionError("cuda_stream_manager.get_event must not be called when disabled") +@pytest.mark.cpu_only class TestDisableMultiStream: def test_default_state_is_enabled(self): assert msu.is_multi_stream_enabled() is True diff --git a/tests/unittest/auto_deploy/singlegpu/utils/test_node_utils_sharding.py b/tests/unittest/auto_deploy/singlegpu/utils/test_node_utils_sharding.py index 502d9ac6d872..4112c2141ca6 100644 --- a/tests/unittest/auto_deploy/singlegpu/utils/test_node_utils_sharding.py +++ b/tests/unittest/auto_deploy/singlegpu/utils/test_node_utils_sharding.py @@ -17,6 +17,7 @@ import operator +import pytest import torch import torch.fx as fx import torch.nn as nn @@ -26,6 +27,8 @@ from tensorrt_llm._torch.auto_deploy.transform.library.sharding_ir import ShardableNode from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_any_split_op, is_any_view_op +pytestmark = pytest.mark.cpu_only + def _call_function_nodes(gm: GraphModule): return [n for n in gm.graph.nodes if n.op == "call_function"] diff --git a/tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py b/tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py index abc1d3d0e67b..250e294e9546 100644 --- a/tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py +++ b/tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py @@ -76,6 +76,7 @@ def test_fp4_global_scale(): assert input_scale == torch.tensor(1.0, dtype=torch.float) +@pytest.mark.cpu_only @pytest.mark.parametrize("amax, expected_scale", [(FP8_MAX, 1.0), (FP8_MAX / 2.0, 0.5)]) def test_fp8_convert_amax_hook(amax, expected_scale): config = TransformConfig(stage="pattern_matcher") @@ -89,6 +90,7 @@ def test_fp8_convert_amax_hook(amax, expected_scale): assert mock_state_dict["scale"] == expected_scale +@pytest.mark.cpu_only def test_fp8_load_hook_maps_prequantized_scales(): config = TransformConfig(stage="pattern_matcher") fp8_imp = FP8LinearQuantizationFromConfig(config) @@ -108,6 +110,7 @@ def test_fp8_load_hook_maps_prequantized_scales(): assert "layer.proj.weight_scale_inv" not in mock_state_dict +@pytest.mark.cpu_only def test_fp8_load_hook_maps_prequantized_scales_with_prefix(): config = TransformConfig(stage="pattern_matcher") fp8_imp = FP8LinearQuantizationFromConfig(config) diff --git a/tests/unittest/auto_deploy/standalone/test_standalone_import.py b/tests/unittest/auto_deploy/standalone/test_standalone_import.py index 67f3b1147016..edbeb69cb97e 100644 --- a/tests/unittest/auto_deploy/standalone/test_standalone_import.py +++ b/tests/unittest/auto_deploy/standalone/test_standalone_import.py @@ -36,6 +36,9 @@ import pytest +pytestmark = pytest.mark.cpu_only + + REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..")) AUTO_DEPLOY_SRC = os.path.join(REPO_ROOT, "tensorrt_llm", "_torch", "auto_deploy")