[TRTLLM-17599][test] Run the host-only MoE unit tests in the CPU stage

tests/integration/test_lists/test-db/l0_a10.yml

@@ -53,7 +53,6 @@ l0_a10:
   - unittest/_torch/pyexecutor/test_resource_manager.py
   - unittest/_torch/pyexecutor/test_seq_slot_sizing.py
   - unittest/_torch/pyexecutor/test_profile_endpoints.py
-  - unittest/_torch/moe/test_communication_factory.py
   # NOTE: this is a CPU-only test, but we do not have a dedicated job for this (and therefore no
   # test list either).
   - unittest/_torch/models/checkpoints

tests/integration/test_lists/test-db/l0_b200.yml

@@ -234,7 +234,6 @@ l0_b200:
   # Real ConfigurableMoE layers sized and registered for DWDP while the model is built.
   - unittest/_torch/modules/dwdp/test_dwdp_moe_construction_gpu.py
   # ------------- MoE components tests ---------------
-  - unittest/_torch/moe/fused_moe/test_cutedsl_fc2_tuning.py
   - unittest/_torch/moe/test_cutedsl_fc1_weight_layout.py
   - unittest/_torch/moe/test_moe_load_balancer.py
   - unittest/_torch/moe/test_moe_routing.py
@@ -246,47 +245,17 @@ l0_b200:
   - unittest/_torch/moe/test_moe_backend.py::test_cutlass_moe_public_activation_ids
   - unittest/_torch/moe/test_kimi_k3_moe_gate.py::test_fused_route_quant_matches_unfused_chain
   - unittest/_torch/moe/test_kimi_k3_moe_gate.py::test_router_logits_kernel_matches_router_gemm
-  - unittest/_torch/moe/test_moe_backend.py::test_kimi_fused_route_quant_skips_prequantized_input
-  - unittest/_torch/moe/test_moe_backend.py::test_kimi_mxfp8_quantized_tensor_handoff
   - unittest/_torch/moe/test_megamoe_streaming_load.py
   - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_tuning_mode_forces_top_maxt_bucket
-  - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_tactic_autotune_defaults_off
-  - unittest/_torch/moe/test_moe_backend.py::test_enumerate_megamoe_candidate_tactics_curated_space
   - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_runner_cache_identity_includes_swiglu_constants
   - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_minimax_m3_swiglu_bias_numerics
   - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_nvfp4_awq_matches_cutlass
-  - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_mpi_bootstrap_keeps_executor_cuda_device
-  - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_mpi_bootstrap_rank0_hosts_store_on_bound_port
-  - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_mpi_bootstrap_store_failure_raises_on_every_rank
-  - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_mpi_bootstrap_loopback_on_single_host
-  - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_mpi_bootstrap_loopback_across_hosts_raises_on_every_rank
-  - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_mpi_bootstrap_resolved_loopback_across_hosts_raises
-  - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_init_raises_bootstrap_failure
   - unittest/_torch/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe
   - unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts
   - unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout
-  - unittest/_torch/moe/test_moe_backend.py::test_import_deep_gemm_rejects_pre_situ_mega_moe_api
-  - unittest/_torch/moe/test_moe_backend.py::test_import_deep_gemm_rejects_mega_moe_without_barrier_timeout
-  - unittest/_torch/moe/test_moe_backend.py::test_megamoe_streaming_reload_resets_slot_claims
   - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_selects_padded_quant_method
   - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_fc31_scale_c_drops_dequant_scale
-  - unittest/_torch/moe/test_moe_backend.py::test_megamoe_bakes_situ_softcaps_as_uniform_scalars
-  - unittest/_torch/moe/test_moe_backend.py::test_megamoe_plain_swiglu_carries_no_constants
-  - unittest/_torch/moe/test_moe_backend.py::test_create_moe_forwards_situ_activation_as_one_carrier
-  - unittest/_torch/moe/test_moe_backend.py::test_create_moe_backend_rejects_apply_router_weight_on_input_by_declaration
-  - unittest/_torch/moe/test_moe_backend.py::test_apply_router_weight_on_input_support_is_not_inherited
   - unittest/_torch/moe/test_moe_backend.py::test_cutedsl_nvfp4_pre_quant_scale_matches_cutlass
-  # WideEP DeepEP-to-CuTeDSL metadata selection and adapter-free contracts.
-  - unittest/_torch/moe/test_moe_backend.py::test_cutedsl_count_native_runner_keeps_both_tactics_and_threads_counts
-  - unittest/_torch/moe/test_moe_backend.py::test_cutedsl_direct_metadata_tuning_buckets_capacity_and_legacy_uses_tokens
-  - unittest/_torch/moe/test_moe_backend.py::test_cutedsl_deep_ep_direct_metadata_disable_env
-  - unittest/_torch/moe/test_moe_backend.py::test_deep_ep_adapter_free_output_matches_schema_and_reuses_cache
-  - unittest/_torch/moe/test_moe_backend.py::test_deep_ep_adapter_free_cache_miss_rejected_during_capture
-  - unittest/_torch/moe/test_moe_backend.py::test_scheduler_selects_cutedsl_deep_ep_direct_metadata
-  - unittest/_torch/moe/test_moe_backend.py::test_other_backend_rejects_deep_ep_direct_metadata
-  - unittest/_torch/moe/test_moe_backend.py::test_cutedsl_rejects_direct_metadata_without_fused_finalize
-  - unittest/_torch/moe/test_moe_backend.py::test_cutedsl_direct_metadata_request_fails_without_fused_finalize
-  - unittest/_torch/moe/test_moe_backend.py::test_scheduler_threads_deep_ep_expert_metadata_to_cutedsl
   - unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py::test_count_native_cutedsl_matches_legacy_sort_bitwise
   # ------------- MoE: test_single_gpu (by backend) ---------------
   - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and not None"

tests/integration/test_lists/test-db/l0_b300.yml

@@ -107,8 +107,6 @@ l0_b300:
   - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_cutedsl_kernel_rejects_unusable_situ_betas[inf_linear]
   - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_cutedsl_kernel_rejects_unusable_situ_betas[zero]
   - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_cutedsl_kernel_rejects_unusable_situ_betas[negative]
-  - unittest/_torch/moe/test_moe_backend.py::test_situ_survives_resolution_not_just_construction[cutlass]
-  - unittest/_torch/moe/test_moe_backend.py::test_situ_survives_resolution_not_just_construction[cutedsl]
   - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_experts_match_situ_reference[CUTLASS-static_1.0]
   - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_experts_match_situ_reference[TRTLLM-static_1.0]
   - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_megamoe_streamed_coverage_survives_per_expert_drain

tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml

@@ -16,8 +16,6 @@ l0_rtx_pro_6000:
   # ------------- PyTorch tests ---------------
   - cpp/test_unit_tests.py::test_fmha_head104_unit_tests[120]
   - unittest/_torch/modeling -k "modeling_out_of_tree"
-  # NCCL-EP v0.2 owns LL shared-memory sizing; cover the factory regression on SM120.
-  - unittest/_torch/moe/test_communication_factory.py
   - unittest/_torch/modeling/test_modeling_deepseekv4.py::test_deepseek_v4_sanity[fp8-ds-mla]
   - unittest/_torch/attention/sparse/test_flashinfer_utils.py
   # DeepSeek-V4.1 on SM120: CSA2 FlashInfer helper, MXFP8 projections and model loading.

tests/unittest/_torch/moe/fused_moe/test_cutedsl_fc2_tuning.py

@@ -21,10 +21,13 @@
 
 _CUTEDSL_FC2_N_TILE_SIZE_ENV = "TRTLLM_CUTEDSL_FC2_N_TILE_SIZE"
 
-pytestmark = pytest.mark.skipif(
-    not IS_CUTLASS_DSL_AVAILABLE,
-    reason="Requires CUTLASS DSL",
-)
+pytestmark = [
+    pytest.mark.cpu_only,
+    pytest.mark.skipif(
+        not IS_CUTLASS_DSL_AVAILABLE,
+        reason="Requires CUTLASS DSL",
+    ),
+]
 
 
 @pytest.fixture(autouse=True)

tests/unittest/_torch/moe/multi_gpu/test_nvlink_one_sided.py

@@ -686,6 +686,7 @@ def test_nvlink_one_sided(case: Case, mpi_pools: dict[tuple[int, bool], MPIPoolE
 # ============================================================================
 
 
+@pytest.mark.cpu_only
 @pytest.mark.parametrize("cft,fp8_combine", [(False, False), (True, False), (True, True)])
 def test_workspace_layout(cft: bool, fp8_combine: bool) -> None:
     from tensorrt_llm.bindings import internal as _tllm_internal
@@ -742,6 +743,7 @@ def field(layout: torch.Tensor, name: str) -> int:
 # ============================================================================
 
 
+@pytest.mark.cpu_only
 @pytest.mark.parametrize(
     ("force_env", "driver_version", "num_tokens", "device_supported", "expected"),
     [
@@ -788,6 +790,7 @@ def query_driver_version() -> str | bytes:
     assert nvlink_one_sided.should_use_cft(can_use_cft, force_cft, 128, num_tokens) is expected
 
 
+@pytest.mark.cpu_only
 @pytest.mark.parametrize(
     ("capability", "unsupported_index"),
     [

tests/unittest/_torch/moe/test_communication_factory.py

@@ -34,6 +34,8 @@
 from tensorrt_llm.models.modeling_utils import QuantConfig
 from tensorrt_llm.quantization import QuantAlgo
 
+pytestmark = pytest.mark.cpu_only
+
 
 def _make_model_config(
     act_dtype: torch.dtype = torch.bfloat16,
@@ -263,6 +265,44 @@ def test_nccl_ep_without_runtime_group_timeout_is_unavailable(monkeypatch: pytes
     assert reason is not None and "set_timeout_ns" in reason
 
 
+@pytest.mark.parametrize(
+    ("num_slots", "smem_optin", "match"),
+    [
+        (32, 232448, None),
+        (32, 101376, "requires 202240 bytes of dynamic shared memory"),
+        (2632, 232448, None),
+        (4096, 232448, "at most 14 expert warp groups"),
+    ],
+    ids=["fits", "smem-too-small", "max-warp-groups", "too-many-warp-groups"],
+)
+def test_nccl_ep_v01_low_latency_combine_device_limits(
+    monkeypatch: pytest.MonkeyPatch, num_slots: int, smem_optin: int, match: str | None
+):
+    """NCCL-EP v0.1 LL combine is limited by the device's dynamic SMEM and 14 expert warp groups."""
+    monkeypatch.setattr(communication_factory, "nccl_ep_supports_version", lambda _: False)
+    monkeypatch.setattr(communication_factory, "is_nccl_ep_installed", lambda: False)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_properties",
+        lambda _: SimpleNamespace(
+            multi_processor_count=188,
+            shared_memory_per_block_optin=smem_optin,
+            shared_memory_per_block=49152,
+        ),
+    )
+
+    reason = communication_factory.CommunicationFactory._get_nccl_ep_unavailable_reason(
+        torch.bfloat16, None, num_slots, 4096, 1024, 1024, 8
+    )
+
+    if match is None:
+        assert reason is None
+    else:
+        assert reason is not None and match in reason
+
+
 class _FakeNcclEP:
     def __init__(
         self,

tests/unittest/_torch/moe/test_cutedsl_fc1_weight_layout.py

@@ -38,6 +38,7 @@ def _raw_fc1(num_experts: int = 1, seed: int = 0) -> torch.Tensor:
     )
 
 
+@pytest.mark.cpu_only
 def test_gate16_helper_matches_megamoe_inline_packing():
     raw = _raw_fc1(num_experts=3)
     ours = interleave_gate_and_linear(raw, group_size=16, dim=1)
@@ -56,6 +57,7 @@ def test_gate16_helper_matches_megamoe_inline_packing():
     assert torch.equal(ours, theirs)
 
 
+@pytest.mark.cpu_only
 def test_gate16_helper_matches_kernel_reference_permutation():
     # The kernel author's host packer starts from the older [up64, gate64]
     # interleave and reorders every 128-row block through the index map
@@ -75,6 +77,7 @@ def test_gate16_helper_matches_kernel_reference_permutation():
     assert torch.equal(interleave_gate_and_linear(raw, group_size=16, dim=0), expected)
 
 
+@pytest.mark.cpu_only
 def test_gate16_helper_is_one_row_permutation_for_weights_and_scales():
     perm = interleave_gate_and_linear(
         torch.arange(FC1_ROWS).view(FC1_ROWS, 1), group_size=16, dim=0
@@ -91,6 +94,7 @@ def test_gate16_helper_is_one_row_permutation_for_weights_and_scales():
     assert torch.equal(interleave_gate_and_linear(scales, group_size=16, dim=0), scales[perm])
 
 
+@pytest.mark.cpu_only
 def test_gate16_helper_dim_argument_is_per_expert():
     raw = _raw_fc1(num_experts=2)
     batched = interleave_gate_and_linear(raw, group_size=16, dim=1)