diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h index ab7ed876257d..4b307c86ddac 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @@ -993,6 +993,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface float* permuted_token_final_scales_{}; int64_t* expert_first_token_offset_{}; + int64_t* gemm_expert_first_token_offset_{}; void* glu_inter_result_{}; void* fc2_result_{}; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu index 8bed9c16b58e..b259c613f30e 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu @@ -1686,6 +1686,43 @@ void expandInputRowsKernelLauncher(InputActivationsType const* unpermuted_input, num_experts_per_node, reinterpret_cast(prequant_scales)); } +template +__global__ void padEmptyFp8BlockScaleMoeInputKernel(T* permuted_input, int64_t const* expert_first_token_offset, + int64_t* gemm_expert_first_token_offset, int num_experts_per_node, int64_t hidden_size) +{ + for (int index = threadIdx.x; index <= num_experts_per_node; index += blockDim.x) + { + gemm_expert_first_token_offset[index] = expert_first_token_offset[index]; + } + + if (expert_first_token_offset[num_experts_per_node] != 0) + { + return; + } + + for (int64_t index = threadIdx.x; index < hidden_size; index += blockDim.x) + { + permuted_input[index] = T(0); + } + __syncthreads(); + + // Assign the zero row to the final local expert. The all-gather/reduce-scatter + // finalizer ignores it because that expert was not selected by any real token. + if (threadIdx.x == 0) + { + gemm_expert_first_token_offset[num_experts_per_node] = 1; + } +} + +template +void padEmptyFp8BlockScaleMoeInput(T* permuted_input, int64_t const* expert_first_token_offset, + int64_t* gemm_expert_first_token_offset, int num_experts_per_node, int64_t hidden_size, cudaStream_t stream) +{ + constexpr int threads = 256; + padEmptyFp8BlockScaleMoeInputKernel<<<1, threads, 0, stream>>>( + permuted_input, expert_first_token_offset, gemm_expert_first_token_offset, num_experts_per_node, hidden_size); +} + #define INSTANTIATE_EXPAND_INPUT_ROWS(InputActivationsType, ExpandedActivationsType) \ template void expandInputRowsKernelLauncher( \ InputActivationsType const* unpermuted_input, ExpandedActivationsType* permuted_output, \ @@ -2985,6 +3022,8 @@ CutlassMoeFCRunner:: size_t const permuted_data_size = permuted_elems * dtype_size; size_t const expert_first_token_offset_size = (num_experts_per_node + 1) * sizeof(int64_t); + size_t const gemm_expert_first_token_offset_size + = use_deepseek_fp8_block_scale ? expert_first_token_offset_size : 0; size_t const permuted_token_final_scales_size = mayHaveFinalizeFused() ? num_moe_inputs * sizeof(float) : 0; size_t const glu_inter_size = glu_inter_elems * gemm_output_dtype; // May be an intermediate type for quantization size_t const fc1_result_size = interbuf_elems * dtype_size; // Activation quantizes so back to dtype_size @@ -3084,6 +3123,7 @@ CutlassMoeFCRunner:: ADD(blocked_expert_counts_cumsum); ADD(blocked_row_to_unpermuted_row); ADD(expert_first_token_offset); + ADD(gemm_expert_first_token_offset); ADD(permuted_token_final_scales); ADD(overlapped_gemm1_gemm2_inputs); ADD(overlapped_gemm1_gemm2_outputs); @@ -3145,6 +3185,7 @@ void CutlassMoeFCRunner 1 && !enable_alltoall) + { + TLLM_CHECK(gemm_expert_first_token_offset_ != nullptr); + padEmptyFp8BlockScaleMoeInput(gemm1_input_expand, expert_first_token_offset_, + gemm_expert_first_token_offset_, num_experts_per_node, hidden_size, stream); + gemm_expert_first_token_offset = gemm_expert_first_token_offset_; + } auto const* gemm1_input = gemm1_input_expand; sync_check_cuda_error(stream); @@ -4349,12 +4405,12 @@ void CutlassMoeFCRunner(smoothed_act_) : fc1_result_; Self::gemm1(moe_gemm_runner_, blockscale_gemm_runner, gemm1_input, gemm1_output, glu_inter_result_, - expert_first_token_offset_, gemm1_tma_ws_input, fc1_expert_weights, fc1_expert_biases, num_valid_tokens_ptr, - fc1_int_scales, fc1_fp8_dequant, use_wfp4afp8 ? fc2_wfp4afp8_quant_scale : fc2_fp8_quant, - fc1_fp4_act_scale_, fc2_fp4_act_scale_, quant_params, num_rows, expanded_num_rows, - expected_tokens_per_expert, hidden_size, inter_size, num_experts_per_node, fc1_activation_type, - alpha_scale_ptr_array_fc1_, !use_lora, stream, *gemm1_config_, false, nullptr, nullptr, - fc2_prequant_scale_ptr); + gemm_expert_first_token_offset, gemm1_tma_ws_input, fc1_expert_weights, fc1_expert_biases, + num_valid_tokens_ptr, fc1_int_scales, fc1_fp8_dequant, + use_wfp4afp8 ? fc2_wfp4afp8_quant_scale : fc2_fp8_quant, fc1_fp4_act_scale_, fc2_fp4_act_scale_, + quant_params, num_rows, expanded_num_rows, expected_tokens_per_expert, hidden_size, inter_size, + num_experts_per_node, fc1_activation_type, alpha_scale_ptr_array_fc1_, !use_lora, stream, *gemm1_config_, + false, nullptr, nullptr, fc2_prequant_scale_ptr); sync_check_cuda_error(stream); if (use_lora) @@ -4371,11 +4427,11 @@ void CutlassMoeFCRunner None: + """Initialize NVML when it has not already been initialized.""" + try: + pynvml.nvmlDeviceGetCount() + except pynvml.NVMLError_Uninitialized: + pynvml.nvmlInit() + @classmethod def get_comm(cls, mapping: Mapping): """Get TP-based communicator (ranks grouped by PP+CP+MOE_TP, ordered by TP rank).""" @@ -355,12 +359,8 @@ def close_mnnvl_memory(cls, ptr: int): @staticmethod @functools.cache def support_nvlink(dev_id: int, need_all_up: bool = True): - # ensure nvml is initialized; do not rely on other modules having - # initialized it as an import side effect. - try: - pynvml.nvmlDeviceGetCount() - except pynvml.NVMLError_Uninitialized: - pynvml.nvmlInit() + # Do not rely on other modules having initialized NVML as an import side effect. + MnnvlMemory._ensure_nvml_initialized() handle = pynvml.nvmlDeviceGetHandleByIndex(dev_id) link_count = pynvml.NVML_NVLINK_MAX_LINKS active_links = 0 @@ -382,6 +382,48 @@ def support_nvlink(dev_id: int, need_all_up: bool = True): else available_links > 0 ) + @staticmethod + @functools.cache + def _is_pcie_nvl_sku(dev_id: int) -> bool: + """Return whether visible H100/H200 GPUs form PCIe-connected NVLink islands.""" + # H100/H200 NVL PCIe SKUs bond GPUs into local NVLink islands joined + # only through PCIe/SYS. Per-device NVLink state therefore cannot + # distinguish them from an NVSwitch fabric. + device_name = torch.cuda.get_device_name(dev_id).upper() + # NVML may report SYSTEM between peers on later NVSwitch platforms, so + # use this fallback only for the affected Hopper SKUs. + if not any(sku in device_name for sku in ("H100", "H200")): + return False + + if " NVL" in device_name: + return True + + try: + MnnvlMemory._ensure_nvml_initialized() + self_handle = pynvml.nvmlDeviceGetHandleByIndex(dev_id) + for peer_id in range(pynvml.nvmlDeviceGetCount()): + if peer_id == dev_id: + continue + peer_handle = pynvml.nvmlDeviceGetHandleByIndex(peer_id) + if ( + pynvml.nvmlDeviceGetTopologyCommonAncestor(self_handle, peer_handle) + == pynvml.NVML_TOPOLOGY_SYSTEM + ): + # SYSTEM is only a distance classification. A dual-socket + # HGX can still provide NVLink P2P to such a peer through + # NVSwitch. Split islands instead have local NVLink but no + # NVLink P2P path to the SYSTEM peer. + p2p_status = pynvml.nvmlDeviceGetP2PStatus( + self_handle, + peer_handle, + pynvml.NVML_P2P_CAPS_INDEX_NVLINK, + ) + if p2p_status != pynvml.NVML_P2P_STATUS_OK: + return MnnvlMemory.support_nvlink(dev_id, need_all_up=False) + except pynvml.NVMLError: + return False + return False + @staticmethod def supports_mnnvl() -> bool: # TODO: @@ -394,6 +436,8 @@ def supports_mnnvl() -> bool: if get_sm_version() in (120, 121): return False dev_id = torch.cuda.current_device() + if MnnvlMemory._is_pcie_nvl_sku(dev_id): + return False support_nvlink_and_all_up = MnnvlMemory.support_nvlink(dev_id, True) return support_nvlink_and_all_up diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py b/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py index a37a3e6e1548..2c0f366febf9 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py @@ -25,6 +25,7 @@ import torch +from tensorrt_llm._mnnvl_utils import MnnvlMemory from tensorrt_llm._torch.modules.fused_moe.deep_ep_utils import buffer_pool, deep_ep_installed from tensorrt_llm._utils import get_sm_version from tensorrt_llm.mapping import Mapping @@ -115,6 +116,13 @@ def is_platform_supported() -> bool: # SM120/121 (RTX PRO 6000 Blackwell): no NVSwitch -> NVSHMEM-LL deadlocks. if get_sm_version() in (120, 121): return False + # Native NVSHMEM/IBGDA bootstrap aborts instead of raising on split + # H100/H200 NVL systems. Disabling P2P does not avoid the abort: this + # build has no IBRC fallback, and IBGDA fails before Buffer can use + # allow_nvlink_for_low_latency_mode=False. Reject before setup. + dev_id = torch.cuda.current_device() + if MnnvlMemory._is_pcie_nvl_sku(dev_id): + return False return True def supports_post_quant_dispatch(self) -> bool: diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 0319c329d9a7..358bd4f7a034 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -3452,6 +3452,9 @@ def test_fp8_blockscale(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, "host_cache_offload", "host_cache_offload_mtp1", "host_cache_offload_mtp3_no_adp" ]) + # Executor warmup runs close to H200 capacity; use a fresh worker pool so + # allocations retained by earlier tests cannot consume its memory headroom. + @pytest.mark.private_mpi_session def test_dsa_host_cache_offload(self, tp_size, pp_size, ep_size, mtp_nextn, overlap_scheduler, max_batch_size, host_cache_size_gb, attention_dp): diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index d862f963ead8..0f09e689df98 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -21,10 +21,8 @@ accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_bf16_small[4] SK accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp] SKIP (https://nvbugs/6428101) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp_trtllm] SKIP (https://nvbugs/6426868) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput_pp4_mtp] SKIP (https://nvbugs/6481323) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload] SKIP (https://nvbugs/6384136) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp1] SKIP (https://nvbugs/6384357) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp3_no_adp] SKIP (https://nvbugs/6384357) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[baseline] SKIP (https://nvbugs/6384136) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] SKIP (https://nvbugs/5989920) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6426847) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] SKIP (https://nvbugs/6517844) diff --git a/tests/unittest/_torch/modules/moe/test_cutlass_moe_op_smoke.py b/tests/unittest/_torch/modules/moe/test_cutlass_moe_op_smoke.py index beaea64197f7..d592f44b6d3d 100644 --- a/tests/unittest/_torch/modules/moe/test_cutlass_moe_op_smoke.py +++ b/tests/unittest/_torch/modules/moe/test_cutlass_moe_op_smoke.py @@ -178,3 +178,54 @@ def test_cutlass_moe_op_run_moe_no_lora_matches_fused_moe_op(): )[0] torch.testing.assert_close(out, ref, rtol=5e-2, atol=1e-2) + + +@requires_cuda_and_op +def test_cutlass_fp8_block_scale_ep_rank_with_no_tokens_returns_zeros(): + """An EP rank with no selected local experts must not enter block-scale GEMM empty.""" + device = torch.device("cuda") + dtype = torch.bfloat16 + num_tokens, hidden_size, inter_size = 8, 128, 128 + local_experts, top_k = 32, 8 + ep_size, ep_rank = 8, 7 + + torch.manual_seed(2) + x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + w3_w1_weight = torch.randn( + local_experts, 2 * inter_size, hidden_size, dtype=dtype, device=device + ).to(torch.float8_e4m3fn) + w2_weight = torch.randn(local_experts, hidden_size, inter_size, dtype=dtype, device=device).to( + torch.float8_e4m3fn + ) + fc1_scales = torch.ones(local_experts, 2 * inter_size // 128, hidden_size // 128, device=device) + fc2_scales = torch.ones(local_experts, hidden_size // 128, inter_size // 128, device=device) + + # Rank 7 owns global experts [224, 256); route every token to [0, 8). + topk_ids = ( + torch.arange(top_k, dtype=torch.int32, device=device).expand(num_tokens, -1).contiguous() + ) + topk_scores = torch.full((num_tokens, top_k), 1.0 / top_k, dtype=torch.float32, device=device) + + def run_moe(ep_size: int, ep_rank: int): + return torch.ops.trtllm.fused_moe( + input=x, + token_selected_experts=topk_ids, + token_final_scales=topk_scores, + fc1_expert_weights=w3_w1_weight, + fc1_expert_biases=None, + fc2_expert_weights=w2_weight, + fc2_expert_biases=None, + output_dtype=dtype, + quant_scales=[fc1_scales, fc2_scales], + ep_size=ep_size, + ep_rank=ep_rank, + use_deepseek_fp8_block_scale=True, + tune_max_num_tokens=num_tokens, + )[0] + + # Copying the offsets for GEMM must preserve a rank with local work. + torch.testing.assert_close(run_moe(ep_size=8, ep_rank=0), run_moe(ep_size=1, ep_rank=0)) + + for _ in range(3): + out = run_moe(ep_size=ep_size, ep_rank=ep_rank) + torch.testing.assert_close(out, torch.zeros_like(out)) diff --git a/tests/unittest/_torch/test_mnnvl_utils.py b/tests/unittest/_torch/test_mnnvl_utils.py new file mode 100644 index 000000000000..eeed206b6840 --- /dev/null +++ b/tests/unittest/_torch/test_mnnvl_utils.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +from unittest.mock import patch + +import pynvml + +from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency import DeepEPLowLatency + + +def setup_function() -> None: + MnnvlMemory._is_pcie_nvl_sku.cache_clear() + MnnvlMemory.support_nvlink.cache_clear() + + +def teardown_function() -> None: + MnnvlMemory._is_pcie_nvl_sku.cache_clear() + MnnvlMemory.support_nvlink.cache_clear() + + +@patch("tensorrt_llm._mnnvl_utils.torch.cuda.get_device_name", return_value="NVIDIA H200 NVL") +def test_pcie_nvl_sku_detected_by_name(mock_get_device_name) -> None: + with patch.object(MnnvlMemory, "_ensure_nvml_initialized") as mock_initialize: + assert MnnvlMemory._is_pcie_nvl_sku(0) + + mock_initialize.assert_not_called() + + +@patch("tensorrt_llm._mnnvl_utils.torch.cuda.get_device_name", return_value="NVIDIA H200") +@patch.object(MnnvlMemory, "_ensure_nvml_initialized") +@patch("tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetCount", return_value=8) +@patch( + "tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetHandleByIndex", side_effect=lambda index: index +) +@patch.object(MnnvlMemory, "support_nvlink", return_value=True) +@patch( + "tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetP2PStatus", + return_value=pynvml.NVML_P2P_STATUS_NOT_SUPPORTED, +) +def test_split_nvlink_topology_detected( + mock_get_p2p_status, + mock_support_nvlink, + mock_get_handle, + mock_get_count, + mock_initialize, + mock_get_device_name, +) -> None: + def common_ancestor(_self_handle: int, peer_handle: int) -> int: + if peer_handle >= 4: + return pynvml.NVML_TOPOLOGY_SYSTEM + return pynvml.NVML_TOPOLOGY_NODE + + with patch( + "tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetTopologyCommonAncestor", + side_effect=common_ancestor, + ): + assert MnnvlMemory._is_pcie_nvl_sku(0) + + mock_get_p2p_status.assert_called_once_with(0, 4, pynvml.NVML_P2P_CAPS_INDEX_NVLINK) + mock_support_nvlink.assert_called_once_with(0, need_all_up=False) + + +@patch("tensorrt_llm._mnnvl_utils.torch.cuda.get_device_name", return_value="NVIDIA H200") +@patch.object(MnnvlMemory, "_ensure_nvml_initialized") +@patch("tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetCount", return_value=2) +@patch( + "tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetHandleByIndex", side_effect=lambda index: index +) +@patch.object(MnnvlMemory, "support_nvlink", return_value=False) +@patch( + "tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetP2PStatus", + return_value=pynvml.NVML_P2P_STATUS_NOT_SUPPORTED, +) +@patch( + "tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetTopologyCommonAncestor", + return_value=pynvml.NVML_TOPOLOGY_SYSTEM, +) +def test_pcie_hopper_with_system_peers_is_not_split_nvlink( + mock_common_ancestor, + mock_get_p2p_status, + mock_support_nvlink, + mock_get_handle, + mock_get_count, + mock_initialize, + mock_get_device_name, +) -> None: + assert not MnnvlMemory._is_pcie_nvl_sku(0) + mock_support_nvlink.assert_called_once_with(0, need_all_up=False) + mock_get_p2p_status.assert_called_once_with(0, 1, pynvml.NVML_P2P_CAPS_INDEX_NVLINK) + mock_common_ancestor.assert_called_once_with(0, 1) + + +@patch("tensorrt_llm._mnnvl_utils.torch.cuda.get_device_name", return_value="NVIDIA H200") +@patch.object(MnnvlMemory, "_ensure_nvml_initialized") +@patch("tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetCount", return_value=2) +@patch( + "tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetHandleByIndex", + side_effect=lambda index: index, +) +@patch.object(MnnvlMemory, "support_nvlink") +@patch( + "tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetP2PStatus", + return_value=pynvml.NVML_P2P_STATUS_OK, +) +@patch( + "tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetTopologyCommonAncestor", + return_value=pynvml.NVML_TOPOLOGY_SYSTEM, +) +def test_dual_socket_hgx_with_system_peers_is_not_split_nvlink( + mock_common_ancestor, + mock_get_p2p_status, + mock_support_nvlink, + mock_get_handle, + mock_get_count, + mock_initialize, + mock_get_device_name, +) -> None: + assert not MnnvlMemory._is_pcie_nvl_sku(0) + mock_get_p2p_status.assert_called_once_with(0, 1, pynvml.NVML_P2P_CAPS_INDEX_NVLINK) + mock_support_nvlink.assert_not_called() + mock_common_ancestor.assert_called_once_with(0, 1) + + +@patch("tensorrt_llm._mnnvl_utils.torch.cuda.get_device_name", return_value="NVIDIA H200") +@patch.object(MnnvlMemory, "_ensure_nvml_initialized") +@patch("tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetCount", return_value=8) +@patch( + "tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetHandleByIndex", side_effect=lambda index: index +) +@patch( + "tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetTopologyCommonAncestor", + return_value=pynvml.NVML_TOPOLOGY_NODE, +) +@patch.object(MnnvlMemory, "support_nvlink", return_value=True) +def test_nvswitch_topology_remains_supported( + mock_support_nvlink, + mock_common_ancestor, + mock_get_handle, + mock_get_count, + mock_initialize, + mock_get_device_name, +) -> None: + assert not MnnvlMemory._is_pcie_nvl_sku(0) + mock_support_nvlink.assert_not_called() + + +@patch("tensorrt_llm._mnnvl_utils.torch.cuda.get_device_name", return_value="NVIDIA B200 NVL") +@patch.object(MnnvlMemory, "_ensure_nvml_initialized") +def test_b200_does_not_use_hopper_topology_fallback(mock_initialize, mock_get_device_name) -> None: + assert not MnnvlMemory._is_pcie_nvl_sku(0) + mock_initialize.assert_not_called() + + +def test_topology_probe_initializes_nvml() -> None: + with ( + patch( + "tensorrt_llm._mnnvl_utils.torch.cuda.get_device_name", + return_value="NVIDIA H200", + ), + patch( + "tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetCount", + side_effect=[pynvml.NVMLError_Uninitialized(), 1], + ), + patch("tensorrt_llm._mnnvl_utils.pynvml.nvmlInit") as mock_nvml_init, + patch("tensorrt_llm._mnnvl_utils.pynvml.nvmlDeviceGetHandleByIndex", return_value=0), + ): + assert not MnnvlMemory._is_pcie_nvl_sku(0) + + mock_nvml_init.assert_called_once_with() + + +@patch("tensorrt_llm._mnnvl_utils.get_sm_version", return_value=90) +@patch("tensorrt_llm._mnnvl_utils.torch.cuda.current_device", return_value=0) +@patch.object(MnnvlMemory, "_is_pcie_nvl_sku", return_value=True) +@patch.object(MnnvlMemory, "support_nvlink") +def test_supports_mnnvl_rejects_split_topology( + mock_support_nvlink, mock_is_pcie_nvl_sku, mock_current_device, mock_get_sm_version +) -> None: + assert not MnnvlMemory.supports_mnnvl() + mock_support_nvlink.assert_not_called() + + +@patch("tensorrt_llm._mnnvl_utils.get_sm_version", return_value=90) +@patch("tensorrt_llm._mnnvl_utils.torch.cuda.current_device", return_value=0) +@patch.object(MnnvlMemory, "_is_pcie_nvl_sku", return_value=False) +@patch.object(MnnvlMemory, "support_nvlink", return_value=True) +def test_supports_mnnvl_accepts_full_fabric( + mock_support_nvlink, mock_is_pcie_nvl_sku, mock_current_device, mock_get_sm_version +) -> None: + assert MnnvlMemory.supports_mnnvl() + mock_support_nvlink.assert_called_once_with(0, True) + + +@patch( + "tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency.deep_ep_installed", + True, +) +@patch( + "tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency.get_sm_version", + return_value=90, +) +@patch( + "tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency.torch.cuda.current_device", + return_value=0, +) +@patch.object(MnnvlMemory, "_is_pcie_nvl_sku", return_value=True) +def test_deep_ep_low_latency_rejects_split_topology( + mock_is_pcie_nvl_sku, mock_current_device, mock_get_sm_version +) -> None: + assert not DeepEPLowLatency.is_platform_supported() + + +@patch( + "tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency.deep_ep_installed", + True, +) +@patch( + "tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency.get_sm_version", + return_value=100, +) +@patch( + "tensorrt_llm._torch.modules.fused_moe.communication.deep_ep_low_latency.torch.cuda.current_device", + return_value=0, +) +@patch.object(MnnvlMemory, "_is_pcie_nvl_sku", return_value=False) +@patch.object(MnnvlMemory, "supports_mnnvl") +def test_deep_ep_low_latency_accepts_b200_without_mnnvl_probe( + mock_supports_mnnvl, mock_is_pcie_nvl_sku, mock_current_device, mock_get_sm_version +) -> None: + assert DeepEPLowLatency.is_platform_supported() + mock_supports_mnnvl.assert_not_called()