From 1abe0cdf000f8071afdf88589a236114b98ad764 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 29 Jun 2026 00:12:19 -0700 Subject: [PATCH 01/20] [None][test] Optimize LLM test startup overhead Signed-off-by: qgai --- .../models/checkpoints/hf/weight_loader.py | 112 +++++++- tensorrt_llm/executor/executor.py | 2 +- tensorrt_llm/executor/proxy.py | 4 +- tensorrt_llm/executor/rpc_proxy.py | 3 +- tensorrt_llm/llmapi/llm.py | 12 +- .../defs/accuracy/test_llm_api_pytorch.py | 270 +++++++++++++++++- .../test-db/l0_gb200_multi_gpus.yml | 16 +- .../checkpoints/hf/test_weight_loader.py | 64 +++++ tests/unittest/llmapi/test_llm.py | 29 +- .../llmapi/test_llm_multi_gpu_pytorch.py | 87 ++++++ 10 files changed, 565 insertions(+), 34 deletions(-) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index 264104c6f2e1..9390d4abf3e7 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -15,6 +15,8 @@ import glob import multiprocessing import os +import threading +from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor from typing import Any, List @@ -33,6 +35,12 @@ from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +_WEIGHT_CACHE_ENV = "TRTLLM_HF_WEIGHT_CACHE" +_WEIGHT_CACHE_MAX_ENTRIES_ENV = "TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES" +_DEFAULT_WEIGHT_CACHE_MAX_ENTRIES = 2 +_WEIGHT_CACHE_LOCK = threading.Lock() +_WEIGHT_CACHE: OrderedDict[tuple, dict[str, Any]] = OrderedDict() + @register_checkpoint_weight_loader("MX") @register_checkpoint_weight_loader("mistral") @@ -43,6 +51,74 @@ class HfWeightLoader(BaseWeightLoader): Loads weights from SafeTensors/bin/pth files. """ + @staticmethod + def _is_weight_cache_enabled() -> bool: + return os.environ.get(_WEIGHT_CACHE_ENV, + "0").lower() in ("1", "true", "yes", "on") + + @staticmethod + def _weight_cache_max_entries() -> int: + try: + return max( + 0, + int( + os.environ.get(_WEIGHT_CACHE_MAX_ENTRIES_ENV, + _DEFAULT_WEIGHT_CACHE_MAX_ENTRIES))) + except ValueError: + logger.warning( + f"Invalid {_WEIGHT_CACHE_MAX_ENTRIES_ENV} value; disabling HF weight cache." + ) + return 0 + + @staticmethod + def _weight_files_cache_key(weight_files: List[str], + use_consolidated: bool) -> tuple: + file_fingerprint = [] + for file_name in sorted(weight_files): + stat = os.stat(file_name) + file_fingerprint.append( + (os.path.abspath(file_name), stat.st_size, stat.st_mtime_ns)) + return (tuple(file_fingerprint), use_consolidated) + + @classmethod + def clear_weight_cache(cls) -> None: + with _WEIGHT_CACHE_LOCK: + _WEIGHT_CACHE.clear() + + @staticmethod + def _wrap_cached_weights(weights: dict[str, Any]) -> ConsumableWeightsDict: + # Return a fresh dict wrapper because model loaders call mark_consumed(). + # Tensor values are intentionally shared: this cache targets read-only + # raw checkpoint tensors, not per-config materialized module weights. + return ConsumableWeightsDict(dict(weights)) + + @staticmethod + def _cache_loaded_weights(cache_key: tuple, + loaded_weights: dict[str, Any]) -> None: + max_entries = HfWeightLoader._weight_cache_max_entries() + if max_entries <= 0: + return + + if isinstance(loaded_weights, ConsumableWeightsDict): + cache_weights = dict(loaded_weights.items()) + else: + cache_weights = dict(loaded_weights) + + with _WEIGHT_CACHE_LOCK: + _WEIGHT_CACHE[cache_key] = cache_weights + _WEIGHT_CACHE.move_to_end(cache_key) + while len(_WEIGHT_CACHE) > max_entries: + _WEIGHT_CACHE.popitem(last=False) + + @staticmethod + def _get_cached_weights(cache_key: tuple) -> ConsumableWeightsDict | None: + with _WEIGHT_CACHE_LOCK: + weights = _WEIGHT_CACHE.get(cache_key) + if weights is None: + return None + _WEIGHT_CACHE.move_to_end(cache_key) + return HfWeightLoader._wrap_cached_weights(weights) + @staticmethod def _get_local_available_host_memory() -> int: """Determine the minimum available memory observed on the local node @@ -77,6 +153,22 @@ def load_weights(self, if len(filtered_weight_files) > 0: weight_files = filtered_weight_files if weight_files: + cache_key = self._weight_files_cache_key(weight_files, + use_consolidated) + if self._is_weight_cache_enabled(): + cached_weights = self._get_cached_weights(cache_key) + if cached_weights is not None: + logger.info( + f"Reusing cached HF checkpoint weights from {checkpoint_dir}." + ) + # The safetensors path always has a local barrier before + # deserialization. Cache hits must participate as well: + # rank-local caches can diverge after eviction, and a hit + # on one rank must not skip a barrier that a miss on + # another rank is about to enter. + local_mpi_barrier() + return cached_weights + # Prefetch the weight files to CPU memory if the size is less than 90% of the available memory. # This is a heuristic to avoid prefetching files that are too large and causing file cache thrashing. prefetch_size = sum(os.path.getsize(file) for file in weight_files) @@ -98,18 +190,34 @@ def load_weights(self, # skipped. Ranks that didn't prefetch reach the barrier immediately. local_mpi_barrier() - return self._load_weights_in_parallel( + weights = self._load_weights_in_parallel( weight_files, self._load_safetensors_file, "Loading safetensors weights in parallel") + if self._is_weight_cache_enabled(): + self._cache_loaded_weights(cache_key, weights) + return weights weight_files = glob.glob(f"{checkpoint_dir}/*.bin") if not weight_files: weight_files = glob.glob(f"{checkpoint_dir}/*.pth") if weight_files: - return self._load_weights_in_parallel( + cache_key = self._weight_files_cache_key(weight_files, + use_consolidated) + if self._is_weight_cache_enabled(): + cached_weights = self._get_cached_weights(cache_key) + if cached_weights is not None: + logger.info( + f"Reusing cached HF checkpoint weights from {checkpoint_dir}." + ) + return cached_weights + + weights = self._load_weights_in_parallel( weight_files, self._load_bin_or_path_file, "Loading bin weights in parallel") + if self._is_weight_cache_enabled(): + self._cache_loaded_weights(cache_key, weights) + return weights raise RuntimeError(f"No weight files found in {checkpoint_dir}.") diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index fdd3dff44a48..d11b12393dc6 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -639,7 +639,7 @@ def create( return GenerationExecutor._create_ipc_executor( worker_kwargs, model_world_size=model_world_size, - mpi_session=None, # use mpi4py + mpi_session=mpi_session, postproc_worker_config=postproc_worker_config, is_llm_executor=is_llm_executor, use_worker=False) diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index ffe793bc5b48..027aa7cae84c 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -98,6 +98,7 @@ def __init__( self.worker_cls = worker_cls mpi_process_pre_spawned: bool = get_spawn_proxy_process_env() + self._owns_mpi_session = mpi_session is None if mpi_session is None: if mpi_process_pre_spawned: @@ -527,7 +528,8 @@ def shutdown(self): self._resource_governor_queue.close() self.workers_started = False - self.mpi_session.shutdown() + if self._owns_mpi_session: + self.mpi_session.shutdown() # Process the errors in-case error during shutting down the threads self._handle_background_error() diff --git a/tensorrt_llm/executor/rpc_proxy.py b/tensorrt_llm/executor/rpc_proxy.py index 7356a9d86d5f..d72259c3e02e 100644 --- a/tensorrt_llm/executor/rpc_proxy.py +++ b/tensorrt_llm/executor/rpc_proxy.py @@ -64,6 +64,7 @@ def __init__( ) self.model_world_size = model_world_size + self._owns_mpi_session = mpi_session is None self._create_mpi_session(model_world_size, mpi_session) # Inject the generated HMAC key into worker_kwargs for workers @@ -275,7 +276,7 @@ def shutdown(self): # 3. shutdown the mpi session, this should wait until all the PyExecutor # processes are shutdown - if self.mpi_session is not None: + if self.mpi_session is not None and self._owns_mpi_session: logger_debug(f"Shutting down mpi session", color="yellow") self.mpi_session.shutdown() logger_debug(f"Mpi session shutdown", color="yellow") diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index b637c1738dc8..ab5d7c446fda 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -313,6 +313,10 @@ def __init__(self, logger_debug(f"LLM.args.mpi_session: {self.args.mpi_session}\n", "yellow") self.mpi_session = self.args.mpi_session + # Keep the live session on LLM only. LLM args are passed to model-build + # tasks and executor workers, and MpiSession objects are not pickleable. + self.args.mpi_session = None + self._owns_mpi_session = self.mpi_session is None # Build this LLM's post-processing hook for the in-proxy detok path (each # postproc worker builds its own). Resolving here fails fast on a bad @@ -338,10 +342,12 @@ def __init__(self, logger_debug("LLM create MpiPoolSession\n", "yellow") self.mpi_session = MpiPoolSession( n_workers=self.args.parallel_config.world_size) + self._owns_mpi_session = True else: logger_debug("LLM create MpiCommSession\n", "yellow") self.mpi_session = create_mpi_comm_session( self.args.parallel_config.world_size) + self._owns_mpi_session = True try: # Due to the Executor can only accept a engine path, we need to save the engine to a directory @@ -364,7 +370,8 @@ def __init__(self, self._build_model() except Exception: - if self.mpi_session is not None: + if (self.mpi_session is not None + and getattr(self, "_owns_mpi_session", True)): self.mpi_session.shutdown() raise @@ -1472,7 +1479,8 @@ def shutdown(self) -> None: self._encoder_executor.shutdown() self._encoder_executor = None - if hasattr(self, 'mpi_session') and self.mpi_session is not None: + if (hasattr(self, 'mpi_session') and self.mpi_session is not None + and getattr(self, "_owns_mpi_session", True)): self.mpi_session.shutdown() self.mpi_session = None diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 641ed8c11ca4..f18d98fbe117 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1337,6 +1337,178 @@ def test_auto_dtype_vswa_chunked_prefill_reuse(self): class TestDeepSeekV3Lite(LlmapiAccuracyTestHarness): MODEL_NAME = "deepseek-ai/DeepSeek-V3-Lite" MODEL_PATH = f"{llm_models_root()}/DeepSeek-V3-Lite/bf16" + _NVFP4_4GPU_PREMERGE_CASES = ( + dict(id="cutlass_mtp0_tp4_compile_off", + moe_backend="CUTLASS", + mtp_nextn=0, + tp_size=4, + pp_size=1, + ep_size=1, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=False, + torch_compile=False), + dict(id="cutlass_mtp0_tp4_compile_on", + moe_backend="CUTLASS", + mtp_nextn=0, + tp_size=4, + pp_size=1, + ep_size=1, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=False, + torch_compile=True), + dict(id="cutlass_mtp0_ep4_compile_on", + moe_backend="CUTLASS", + mtp_nextn=0, + tp_size=4, + pp_size=1, + ep_size=4, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=False, + torch_compile=True), + dict(id="cutlass_mtp0_tp2pp2_compile_off", + moe_backend="CUTLASS", + mtp_nextn=0, + tp_size=2, + pp_size=2, + ep_size=1, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=False, + torch_compile=False), + dict(id="cutlass_mtp0_tp2pp2_compile_on", + moe_backend="CUTLASS", + mtp_nextn=0, + tp_size=2, + pp_size=2, + ep_size=1, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=False, + torch_compile=True), + dict(id="cutlass_mtp2_tp4_compile_off", + moe_backend="CUTLASS", + mtp_nextn=2, + tp_size=4, + pp_size=1, + ep_size=1, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=False, + torch_compile=False), + dict(id="trtllm_mtp2_ep4_compile_off", + moe_backend="TRTLLM", + mtp_nextn=2, + tp_size=4, + pp_size=1, + ep_size=4, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=False, + torch_compile=False), + dict(id="cutlass_mtp2_pp4_compile_off", + moe_backend="CUTLASS", + mtp_nextn=2, + tp_size=1, + pp_size=4, + ep_size=1, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=False, + torch_compile=False), + dict(id="cutlass_mtp0_pp4_compile_off", + moe_backend="CUTLASS", + mtp_nextn=0, + tp_size=1, + pp_size=4, + ep_size=1, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=False, + torch_compile=False), + dict(id="cutlass_mtp0_pp4_compile_on", + moe_backend="CUTLASS", + mtp_nextn=0, + tp_size=1, + pp_size=4, + ep_size=1, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=False, + torch_compile=True), + dict(id="trtllm_mtp0_tp4_compile_off", + moe_backend="TRTLLM", + mtp_nextn=0, + tp_size=4, + pp_size=1, + ep_size=1, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=False, + torch_compile=False), + dict(id="trtllm_mtp0_ep4_compile_off", + moe_backend="TRTLLM", + mtp_nextn=0, + tp_size=4, + pp_size=1, + ep_size=4, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=False, + torch_compile=False), + dict(id="cutlass_mtp0_tp4_lpc_compile_off", + moe_backend="CUTLASS", + mtp_nextn=0, + tp_size=4, + pp_size=1, + ep_size=1, + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True, + low_precision_combine=True, + torch_compile=False), + ) + _NVFP4_4GPU_PREMERGE_GROUPS = ( + ("tp4_mtp0", + ("cutlass_mtp0_tp4_compile_off", "cutlass_mtp0_tp4_compile_on", + "trtllm_mtp0_tp4_compile_off", "cutlass_mtp0_tp4_lpc_compile_off")), + ("ep4_mtp0", ("cutlass_mtp0_ep4_compile_on", + "trtllm_mtp0_ep4_compile_off")), + ("tp2pp2_mtp0", ("cutlass_mtp0_tp2pp2_compile_off", + "cutlass_mtp0_tp2pp2_compile_on")), + ("tp4_mtp2", ("cutlass_mtp2_tp4_compile_off", )), + ("ep4_mtp2", ("trtllm_mtp2_ep4_compile_off", )), + ("pp4_mtp2", ("cutlass_mtp2_pp4_compile_off", )), + ("pp4_mtp0", ("cutlass_mtp0_pp4_compile_off", + "cutlass_mtp0_pp4_compile_on")), + ) @pytest.mark.skip_less_device_memory(60000) @parametrize_with_ids("v2_kv_cache", [True, False]) @@ -2207,6 +2379,80 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, overlap_scheduler, low_precision_combine, tp_size, pp_size, ep_size, torch_compile, mtp_nextn, moe_backend): + self._run_nvfp4_4gpus_case( + fp8kv=fp8kv, + attention_dp=attention_dp, + cuda_graph=cuda_graph, + overlap_scheduler=overlap_scheduler, + low_precision_combine=low_precision_combine, + tp_size=tp_size, + pp_size=pp_size, + ep_size=ep_size, + torch_compile=torch_compile, + mtp_nextn=mtp_nextn, + moe_backend=moe_backend, + ) + + @pytest.mark.skip_less_device(4) + @skip_pre_blackwell + def test_nvfp4_4gpus_premerge_grouped(self): + from tensorrt_llm.llmapi.mpi_session import MpiPoolSession + + cases_by_id = { + case["id"]: case + for case in self._NVFP4_4GPU_PREMERGE_CASES + } + grouped_case_ids = tuple( + case_id for _, case_ids in self._NVFP4_4GPU_PREMERGE_GROUPS + for case_id in case_ids) + assert len(grouped_case_ids) == len(set(grouped_case_ids)) + assert set(grouped_case_ids) == set(cases_by_id) + + previous_cache_env = os.environ.get("TRTLLM_HF_WEIGHT_CACHE") + os.environ["TRTLLM_HF_WEIGHT_CACHE"] = "1" + try: + for group_id, case_ids in self._NVFP4_4GPU_PREMERGE_GROUPS: + mpi_session = (MpiPoolSession( + n_workers=4) if len(case_ids) > 1 else None) + try: + for case_id in case_ids: + case = cases_by_id[case_id] + case_kwargs = { + key: value + for key, value in case.items() if key != "id" + } + try: + self._run_nvfp4_4gpus_case(_mpi_session=mpi_session, + **case_kwargs) + except pytest.skip.Exception: + raise + except Exception as exc: + raise AssertionError( + "DeepSeek V3 Lite NVFP4 grouped case failed: " + f"{group_id}/{case_id}") from exc + finally: + if mpi_session is not None: + mpi_session.shutdown() + finally: + if previous_cache_env is None: + os.environ.pop("TRTLLM_HF_WEIGHT_CACHE", None) + else: + os.environ["TRTLLM_HF_WEIGHT_CACHE"] = previous_cache_env + + def _run_nvfp4_4gpus_case(self, + *, + fp8kv, + attention_dp, + cuda_graph, + overlap_scheduler, + low_precision_combine, + tp_size, + pp_size, + ep_size, + torch_compile, + mtp_nextn, + moe_backend, + _mpi_session=None): sm_version = get_sm_version() if moe_backend == "TRTLLM" and sm_version in (120, 121): pytest.skip(f"{moe_backend} backend does not support SM 120 or 121") @@ -2233,16 +2479,20 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, if fp8kv: kv_cache_config.dtype = "fp8" - with LLM( - f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - kv_cache_config=kv_cache_config, - **pytorch_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config, - ) as llm: + llm_kwargs = dict( + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + **pytorch_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) + if _mpi_session is not None: + llm_kwargs["_mpi_session"] = _mpi_session + + with LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", + **llm_kwargs) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 task = GSM8K(self.MODEL_NAME) diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index f12e16220744..1082dd24bf41 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -20,15 +20,13 @@ l0_gb200_multi_gpus: - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - # ---- curated pre-merge smoke set: diversify moe_backend across CUTLASS/TRTLLM/CUTEDSL ---- - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=0-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=0-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=2-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=2-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=True-torch_compile=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] + - 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] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_kv_cache_aware_routing[mtp_nextn=0] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_kv_cache_aware_routing[mtp_nextn=2] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=True-moe_backend=WIDEEP] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=True-moe_backend=TRTLLM] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_online_eplb[moe_backend=TRTLLM] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_block_reuse[TEP4] diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py index fbe8a3dbf481..728fa3c9f19f 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py @@ -3,6 +3,7 @@ import pytest from tensorrt_llm._torch.models.checkpoints import HfWeightLoader +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict from tensorrt_llm.mapping import Mapping @@ -97,3 +98,66 @@ def test_load_weights_ignores_consolidated_ckpt_when_sharded_ckpt_exists( load_weights_in_parallel.assert_called_once() loaded_weight_files = load_weights_in_parallel.call_args[0][0] assert set(loaded_weight_files) == expected_safetensor_filenames + + +def test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper(tmp_path, monkeypatch): + monkeypatch.setenv("TRTLLM_HF_WEIGHT_CACHE", "1") + HfWeightLoader.clear_weight_cache() + + checkpoint_dir = tmp_path / "foo" + checkpoint_dir.mkdir() + (checkpoint_dir / "model.safetensors").touch() + + raw_weight = object() + loader = HfWeightLoader() + + try: + with ( + mock.patch.object( + loader, + "_load_weights_in_parallel", + return_value=ConsumableWeightsDict({"foo.weight": raw_weight}), + ) as load_weights_in_parallel, + mock.patch.object(loader, "prefetch_files"), + ): + first = loader.load_weights(str(checkpoint_dir), mapping=Mapping()) + assert first["foo.weight"] is raw_weight + assert first.mark_consumed("foo") == 1 + assert len(first) == 0 + + second = loader.load_weights(str(checkpoint_dir), mapping=Mapping()) + + load_weights_in_parallel.assert_called_once() + assert second["foo.weight"] is raw_weight + finally: + HfWeightLoader.clear_weight_cache() + + +def test_weight_cache_disabled_by_default(tmp_path, monkeypatch): + monkeypatch.delenv("TRTLLM_HF_WEIGHT_CACHE", raising=False) + HfWeightLoader.clear_weight_cache() + + checkpoint_dir = tmp_path / "foo" + checkpoint_dir.mkdir() + (checkpoint_dir / "model.safetensors").touch() + + loader = HfWeightLoader() + + try: + with ( + mock.patch.object( + loader, + "_load_weights_in_parallel", + side_effect=[ + ConsumableWeightsDict({"foo.weight": object()}), + ConsumableWeightsDict({"foo.weight": object()}), + ], + ) as load_weights_in_parallel, + mock.patch.object(loader, "prefetch_files"), + ): + loader.load_weights(str(checkpoint_dir), mapping=Mapping()) + loader.load_weights(str(checkpoint_dir), mapping=Mapping()) + + assert load_weights_in_parallel.call_count == 2 + finally: + HfWeightLoader.clear_weight_cache() diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index fb174a99152d..1579261922a2 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -1827,7 +1827,8 @@ def llm_return_logprobs_test_harness(prompt_logprobs: Optional[int], return_generation_logits: bool, tp_size=1, streaming=False, - backend=None): + backend=None, + **extra_llm_kwargs): LLM_CLASS = LLM llm_args_extra = {} kv_cache_args_extra = {} @@ -1849,6 +1850,7 @@ def llm_return_logprobs_test_harness(prompt_logprobs: Optional[int], tensor_parallel_size=tp_size, gather_generation_logits=True, **llm_args_extra, + **extra_llm_kwargs, ) prompts = ["A B C D E F G H I J K"] @@ -1920,6 +1922,7 @@ async def main(): await asyncio.gather(*tasks) asyncio.run(main()) + llm.shutdown() @force_ampere @@ -2205,7 +2208,8 @@ def llm_get_stats_test_harness(tp_size: int = 1, pytorch_backend: bool = False, use_overlap: bool = False, enable_chunked_prefill: bool = False, - enable_iter_req_stats: bool = False): + enable_iter_req_stats: bool = False, + **extra_llm_kwargs): if return_context_logits and pytorch_backend: pytest.skip("pytorch backend does not support context logits") @@ -2252,7 +2256,8 @@ def llm_get_stats_test_harness(tp_size: int = 1, kv_cache_config=global_kvcache_config, tensor_parallel_size=tp_size, pipeline_parallel_size=pp_size, - **llm_args_extra) as llm: + **llm_args_extra, + **extra_llm_kwargs) as llm: max_tokens = 5 sampling_params = SamplingParams(max_tokens=max_tokens, @@ -2354,7 +2359,8 @@ def llm_get_stats_async_test_harness(tp_size: int = 1, pytorch_backend: bool = False, use_overlap: bool = False, enable_chunked_prefill: bool = False, - enable_iter_req_stats: bool = False): + enable_iter_req_stats: bool = False, + **extra_llm_kwargs): if return_context_logits and pytorch_backend: pytest.skip("pytorch backend does not support context logits") @@ -2395,7 +2401,8 @@ def llm_get_stats_async_test_harness(tp_size: int = 1, kv_cache_config=global_kvcache_config, tensor_parallel_size=tp_size, pipeline_parallel_size=pp_size, - **llm_args_extra) as llm: + **llm_args_extra, + **extra_llm_kwargs) as llm: max_tokens = 6 sampling_params = SamplingParams(max_tokens=max_tokens, @@ -2491,7 +2498,9 @@ def success_path(): success_path() -def _test_llm_capture_request_error(pytorch_backend: bool, tp_size: int = 1): +def _test_llm_capture_request_error(pytorch_backend: bool, + tp_size: int = 1, + **extra_llm_kwargs): llm_args_extra = {} if pytorch_backend: LLM_CLASS = LLM_torch @@ -2507,12 +2516,16 @@ def _test_llm_capture_request_error(pytorch_backend: bool, tp_size: int = 1): model=llama_model_path, tensor_parallel_size=tp_size, **llm_args_extra, + **extra_llm_kwargs, ) prompt = 'A ' * 65 # the minimum max_num_tokens is 64 # Both backends now consistently raise RequestError for max_num_tokens validation - with pytest.raises(RequestError): - llm.generate(prompt) + try: + with pytest.raises(RequestError): + llm.generate(prompt) + finally: + llm.shutdown() def test_llm_capture_request_error(): diff --git a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py index 0e15b38b8b95..c30efc1778d7 100644 --- a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py +++ b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py @@ -1,3 +1,5 @@ +from collections.abc import Callable + import pytest from utils.util import skip_ray @@ -188,3 +190,88 @@ def test_llm_get_stats_async_tp2(): @pytest.mark.gpu2 def test_llm_get_stats_async_pp2(): llm_get_stats_async_test_harness(pp_size=2, pytorch_backend=True) + + +def _run_grouped_mpi_cases( + n_workers: int, cases: tuple[tuple[str, Callable[[object], None]], + ...]) -> None: + from tensorrt_llm.llmapi.mpi_session import MpiPoolSession + + mpi_session = MpiPoolSession(n_workers=n_workers) + try: + for case_name, run_case in cases: + print(f"Running grouped multi-GPU PyTorch LLM case: {case_name}") + try: + run_case(mpi_session) + except pytest.skip.Exception: + raise + except Exception as exc: + raise AssertionError( + f"Grouped multi-GPU PyTorch LLM case failed: {case_name}" + ) from exc + finally: + mpi_session.shutdown() + + +def test_llm_multi_gpu_pytorch_grouped_mpi_reuse(): + # Chunked stats cases are intentionally left as standalone tests because + # their microbatch-id checks depend on scheduler state. + gpu2_cases = ( + ("test_llm_capture_request_error", + lambda mpi_session: _test_llm_capture_request_error( + pytorch_backend=True, tp_size=2, _mpi_session=mpi_session)), + ("test_tinyllama_logits_processor_2gpu[tp1pp2]", lambda mpi_session: + tinyllama_logits_processor_test_harness(backend="pytorch", + tensor_parallel_size=1, + pipeline_parallel_size=2, + _mpi_session=mpi_session)), + ("test_tinyllama_logits_processor_2gpu[tp2pp1]", lambda mpi_session: + tinyllama_logits_processor_test_harness(backend="pytorch", + tensor_parallel_size=2, + pipeline_parallel_size=1, + _mpi_session=mpi_session)), + ("test_llm_return_logprobs_streaming_tp2", lambda mpi_session: + llm_return_logprobs_test_harness(None, + 1, + False, + False, + streaming=True, + backend="pytorch", + tp_size=2, + _mpi_session=mpi_session)), + ("test_llm_get_stats_pp2[no_chunked]", lambda mpi_session: + llm_get_stats_test_harness(tp_size=1, + pp_size=2, + return_context_logits=False, + pytorch_backend=True, + enable_chunked_prefill=False, + enable_iter_req_stats=True, + _mpi_session=mpi_session)), + ("test_llm_get_stats_tp2", + lambda mpi_session: llm_get_stats_test_harness( + tp_size=2, pytorch_backend=True, _mpi_session=mpi_session)), + ("test_llm_get_stats_async_tp2", + lambda mpi_session: llm_get_stats_async_test_harness( + tp_size=2, pytorch_backend=True, _mpi_session=mpi_session)), + ("test_llm_get_stats_async_pp2", + lambda mpi_session: llm_get_stats_async_test_harness( + pp_size=2, pytorch_backend=True, _mpi_session=mpi_session)), + ) + gpu4_cases = ( + ("test_tinyllama_logits_processor_tp2pp2", lambda mpi_session: + tinyllama_logits_processor_test_harness(backend="pytorch", + tensor_parallel_size=2, + pipeline_parallel_size=2, + _mpi_session=mpi_session)), + ("test_llm_get_stats_pp4[no_chunked]", lambda mpi_session: + llm_get_stats_test_harness(tp_size=1, + pp_size=4, + return_context_logits=False, + pytorch_backend=True, + enable_chunked_prefill=False, + enable_iter_req_stats=True, + _mpi_session=mpi_session)), + ) + + _run_grouped_mpi_cases(2, gpu2_cases) + _run_grouped_mpi_cases(4, gpu4_cases) From a99bef23970e4e455ed0d85816d24d529667430f Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 30 Jun 2026 02:15:25 -0700 Subject: [PATCH 02/20] [None][test] Group multi-GPU LLM API tests Signed-off-by: qgai --- .../test_lists/test-db/l0_dgx_h100.yml | 18 +++- .../llmapi/test_llm_multi_gpu_pytorch.py | 96 ++++++++++++++++--- 2 files changed, 96 insertions(+), 18 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index d65e8ba0cd70..fc3c8a378951 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -15,7 +15,7 @@ l0_dgx_h100: auto_trigger: others orchestrator: mpi tests: - - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2" + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 - unittest/llmapi/test_additional_model_outputs.py -m "gpu2" - unittest/_torch/multi_gpu -m "not post_merge" TIMEOUT (90) - unittest/_torch/modeling/test_modeling_pixtral.py::test_tensor_parallelism @@ -90,7 +90,7 @@ l0_dgx_h100: orchestrator: mpi tests: # ------------- PyTorch tests --------------- - - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu4" + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_multi_gpu_pytorch_grouped_mpi_reuse # ------------- Model specific tests --------------- - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=True] @@ -330,7 +330,19 @@ l0_dgx_h100: - unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py -m "part2" - unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py -m "part3" - unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py -m "part4" - - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2" + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_capture_request_error + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_2gpu[1-2] + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_2gpu[2-1] + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llama_7b_lora_tp2 + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_rpc_tp2 + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_rpc_streaming_tp2 + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_return_logprobs_streaming_tp2[None-1-False-False] + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_pp2[False-False-True] + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_pp2[False-True-True] + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_tp2 + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_async_tp2 + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_async_pp2 - unittest/llmapi/test_async_llm.py -m "gpu2" - accuracy/test_llm_api_pytorch_ray.py::TestLlama3_1_8BInstruct::test_pp2_ray - examples/test_ray.py::test_llm_inference_distributed_ray[tp2] diff --git a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py index c30efc1778d7..e593aeb3d6d9 100644 --- a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py +++ b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py @@ -1,3 +1,4 @@ +import asyncio from collections.abc import Callable import pytest @@ -5,7 +6,7 @@ from tensorrt_llm import LLM from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy -from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig from tensorrt_llm.lora_helper import LoraConfig from tensorrt_llm.sampling_params import SamplingParams @@ -81,10 +82,26 @@ def test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1( @skip_ray @pytest.mark.gpu2 def test_llm_rpc_tp2(): - with LLM(model=llama_model_path, - kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), - orchestrator_type="rpc", - tensor_parallel_size=2) as llm: + _run_llm_rpc_tp2() + + +@skip_ray +@pytest.mark.gpu2 +@pytest.mark.asyncio +async def test_llm_rpc_streaming_tp2(): + await _run_llm_rpc_streaming_tp2() + + +def _run_llm_rpc_tp2(_mpi_session=None): + llm_kwargs = dict( + model=llama_model_path, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), + orchestrator_type="rpc", + tensor_parallel_size=2) + if _mpi_session is not None: + llm_kwargs["_mpi_session"] = _mpi_session + + with LLM(**llm_kwargs) as llm: assert isinstance(llm._executor, GenerationExecutorRpcProxy) res = llm.generate("Tell me a joke", @@ -96,14 +113,16 @@ def test_llm_rpc_tp2(): assert len(res.outputs[0].token_ids) == 10 -@skip_ray -@pytest.mark.gpu2 -@pytest.mark.asyncio -async def test_llm_rpc_streaming_tp2(): - with LLM(model=llama_model_path, - kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), - orchestrator_type="rpc", - tensor_parallel_size=2) as llm: +async def _run_llm_rpc_streaming_tp2(_mpi_session=None): + llm_kwargs = dict( + model=llama_model_path, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), + orchestrator_type="rpc", + tensor_parallel_size=2) + if _mpi_session is not None: + llm_kwargs["_mpi_session"] = _mpi_session + + with LLM(**llm_kwargs) as llm: assert isinstance(llm._executor, GenerationExecutorRpcProxy) async for output in llm.generate_async("Tell me a joke", @@ -214,8 +233,6 @@ def _run_grouped_mpi_cases( def test_llm_multi_gpu_pytorch_grouped_mpi_reuse(): - # Chunked stats cases are intentionally left as standalone tests because - # their microbatch-id checks depend on scheduler state. gpu2_cases = ( ("test_llm_capture_request_error", lambda mpi_session: _test_llm_capture_request_error( @@ -230,6 +247,15 @@ def test_llm_multi_gpu_pytorch_grouped_mpi_reuse(): tensor_parallel_size=2, pipeline_parallel_size=1, _mpi_session=mpi_session)), + ("test_llama_7b_lora_tp2", + lambda mpi_session: llama_7b_lora_from_dir_test_harness( + tensor_parallel_size=2, + kv_cache_config=global_kv_cache_config, + _mpi_session=mpi_session)), + ("test_llm_rpc_tp2", + lambda mpi_session: _run_llm_rpc_tp2(_mpi_session=mpi_session)), + ("test_llm_rpc_streaming_tp2", lambda mpi_session: asyncio.run( + _run_llm_rpc_streaming_tp2(_mpi_session=mpi_session))), ("test_llm_return_logprobs_streaming_tp2", lambda mpi_session: llm_return_logprobs_test_harness(None, 1, @@ -247,6 +273,14 @@ def test_llm_multi_gpu_pytorch_grouped_mpi_reuse(): enable_chunked_prefill=False, enable_iter_req_stats=True, _mpi_session=mpi_session)), + ("test_llm_get_stats_pp2[chunked]", lambda mpi_session: + llm_get_stats_test_harness(tp_size=1, + pp_size=2, + return_context_logits=False, + pytorch_backend=True, + enable_chunked_prefill=True, + enable_iter_req_stats=True, + _mpi_session=mpi_session)), ("test_llm_get_stats_tp2", lambda mpi_session: llm_get_stats_test_harness( tp_size=2, pytorch_backend=True, _mpi_session=mpi_session)), @@ -263,6 +297,30 @@ def test_llm_multi_gpu_pytorch_grouped_mpi_reuse(): tensor_parallel_size=2, pipeline_parallel_size=2, _mpi_session=mpi_session)), + ("test_llama_7b_multi_lora_tp4[cuda_graph]", lambda mpi_session: + check_llama_7b_multi_lora_from_request_test_harness( + LLM, + lora_config=LoraConfig(lora_target_modules= + ['attn_q', 'attn_k', 'attn_v'], + max_lora_rank=8, + max_loras=1, + max_cpu_loras=8), + tensor_parallel_size=4, + kv_cache_config=global_kv_cache_config, + cuda_graph_config=CudaGraphConfig(max_batch_size=10), + _mpi_session=mpi_session)), + ("test_llama_7b_multi_lora_tp4[no_cuda_graph]", lambda mpi_session: + check_llama_7b_multi_lora_from_request_test_harness( + LLM, + lora_config=LoraConfig(lora_target_modules= + ['attn_q', 'attn_k', 'attn_v'], + max_lora_rank=8, + max_loras=1, + max_cpu_loras=8), + tensor_parallel_size=4, + kv_cache_config=global_kv_cache_config, + cuda_graph_config=None, + _mpi_session=mpi_session)), ("test_llm_get_stats_pp4[no_chunked]", lambda mpi_session: llm_get_stats_test_harness(tp_size=1, pp_size=4, @@ -271,6 +329,14 @@ def test_llm_multi_gpu_pytorch_grouped_mpi_reuse(): enable_chunked_prefill=False, enable_iter_req_stats=True, _mpi_session=mpi_session)), + ("test_llm_get_stats_pp4[chunked]", lambda mpi_session: + llm_get_stats_test_harness(tp_size=1, + pp_size=4, + return_context_logits=False, + pytorch_backend=True, + enable_chunked_prefill=True, + enable_iter_req_stats=True, + _mpi_session=mpi_session)), ) _run_grouped_mpi_cases(2, gpu2_cases) From f7ed74fd62c091265c062b0093029a27dc24f6c5 Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 30 Jun 2026 19:23:07 -0700 Subject: [PATCH 03/20] [None][test] Reuse MPI sessions with pytest fixtures Signed-off-by: qgai --- .../test_lists/test-db/l0_dgx_h100.yml | 17 +- .../llmapi/test_llm_multi_gpu_pytorch.py | 260 ++++++------------ 2 files changed, 92 insertions(+), 185 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index fc3c8a378951..f39810a4f07d 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -16,6 +16,7 @@ l0_dgx_h100: orchestrator: mpi tests: - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 + - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2" -k "not test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1" - unittest/llmapi/test_additional_model_outputs.py -m "gpu2" - unittest/_torch/multi_gpu -m "not post_merge" TIMEOUT (90) - unittest/_torch/modeling/test_modeling_pixtral.py::test_tensor_parallelism @@ -90,7 +91,7 @@ l0_dgx_h100: orchestrator: mpi tests: # ------------- PyTorch tests --------------- - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_multi_gpu_pytorch_grouped_mpi_reuse + - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu4" # ------------- Model specific tests --------------- - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=True] @@ -330,19 +331,7 @@ l0_dgx_h100: - unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py -m "part2" - unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py -m "part3" - unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py -m "part4" - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_capture_request_error - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_2gpu[1-2] - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_2gpu[2-1] - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llama_7b_lora_tp2 - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_rpc_tp2 - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_rpc_streaming_tp2 - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_return_logprobs_streaming_tp2[None-1-False-False] - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_pp2[False-False-True] - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_pp2[False-True-True] - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_tp2 - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_async_tp2 - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_async_pp2 + - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2" - unittest/llmapi/test_async_llm.py -m "gpu2" - accuracy/test_llm_api_pytorch_ray.py::TestLlama3_1_8BInstruct::test_pp2_ray - examples/test_ray.py::test_llm_inference_distributed_ray[tp2] diff --git a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py index e593aeb3d6d9..d80233c8bb8f 100644 --- a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py +++ b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py @@ -1,12 +1,10 @@ -import asyncio -from collections.abc import Callable - import pytest from utils.util import skip_ray from tensorrt_llm import LLM +from tensorrt_llm._utils import mpi_disabled from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy -from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig +from tensorrt_llm.llmapi import KvCacheConfig from tensorrt_llm.lora_helper import LoraConfig from tensorrt_llm.sampling_params import SamplingParams @@ -24,37 +22,75 @@ global_kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4) +def _shared_mpi_session(n_workers: int): + if mpi_disabled(): + yield None + return + + from tensorrt_llm.llmapi.mpi_session import MpiPoolSession + + mpi_session = MpiPoolSession(n_workers=n_workers) + try: + yield mpi_session + finally: + mpi_session.shutdown() + + +@pytest.fixture(scope="module") +def shared_mpi_session_2gpu(): + yield from _shared_mpi_session(2) + + +@pytest.fixture(scope="module") +def shared_mpi_session_4gpu(): + yield from _shared_mpi_session(4) + + +def _mpi_session_kwargs(mpi_session) -> dict: + return {"_mpi_session": mpi_session} if mpi_session is not None else {} + + @pytest.mark.gpu2 -def test_llm_capture_request_error(): - _test_llm_capture_request_error(pytorch_backend=True, tp_size=2) +def test_llm_capture_request_error(shared_mpi_session_2gpu): + _test_llm_capture_request_error( + pytorch_backend=True, + tp_size=2, + **_mpi_session_kwargs(shared_mpi_session_2gpu)) @pytest.mark.gpu4 -def test_tinyllama_logits_processor_tp2pp2(): - tinyllama_logits_processor_test_harness(backend="pytorch", - tensor_parallel_size=2, - pipeline_parallel_size=2) +def test_tinyllama_logits_processor_tp2pp2(shared_mpi_session_4gpu): + tinyllama_logits_processor_test_harness( + backend="pytorch", + tensor_parallel_size=2, + pipeline_parallel_size=2, + **_mpi_session_kwargs(shared_mpi_session_4gpu)) @pytest.mark.gpu2 @pytest.mark.part0 @pytest.mark.parametrize("tp_size, pp_size", [(1, 2), (2, 1)]) -def test_tinyllama_logits_processor_2gpu(tp_size: int, pp_size: int): - tinyllama_logits_processor_test_harness(backend="pytorch", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size) +def test_tinyllama_logits_processor_2gpu(tp_size: int, pp_size: int, + shared_mpi_session_2gpu): + tinyllama_logits_processor_test_harness( + backend="pytorch", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + **_mpi_session_kwargs(shared_mpi_session_2gpu)) @pytest.mark.gpu2 -def test_llama_7b_lora_tp2(): - llama_7b_lora_from_dir_test_harness(tensor_parallel_size=2, - kv_cache_config=global_kv_cache_config) +def test_llama_7b_lora_tp2(shared_mpi_session_2gpu): + llama_7b_lora_from_dir_test_harness( + tensor_parallel_size=2, + kv_cache_config=global_kv_cache_config, + **_mpi_session_kwargs(shared_mpi_session_2gpu)) @pytest.mark.gpu4 @skip_ray # https://nvbugs/5682551 @test_lora_with_and_without_cuda_graph -def test_llama_7b_multi_lora_tp4(cuda_graph_config): +def test_llama_7b_multi_lora_tp4(cuda_graph_config, shared_mpi_session_4gpu): # For LoRA checkpoints without finetuned embedding and lm_head, we can either: # (1) specify lora_target_modules, or # (2) provide a lora_dir to infer the lora_target_modules. @@ -67,7 +103,8 @@ def test_llama_7b_multi_lora_tp4(cuda_graph_config): lora_config=lora_config, tensor_parallel_size=4, kv_cache_config=global_kv_cache_config, - cuda_graph_config=cuda_graph_config) + cuda_graph_config=cuda_graph_config, + **_mpi_session_kwargs(shared_mpi_session_4gpu)) @skip_ray # https://nvbugs/5727075 @@ -81,15 +118,15 @@ def test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1( @skip_ray @pytest.mark.gpu2 -def test_llm_rpc_tp2(): - _run_llm_rpc_tp2() +def test_llm_rpc_tp2(shared_mpi_session_2gpu): + _run_llm_rpc_tp2(shared_mpi_session_2gpu) @skip_ray @pytest.mark.gpu2 @pytest.mark.asyncio -async def test_llm_rpc_streaming_tp2(): - await _run_llm_rpc_streaming_tp2() +async def test_llm_rpc_streaming_tp2(shared_mpi_session_2gpu): + await _run_llm_rpc_streaming_tp2(shared_mpi_session_2gpu) def _run_llm_rpc_tp2(_mpi_session=None): @@ -141,14 +178,17 @@ async def _run_llm_rpc_streaming_tp2(_mpi_session=None): ]) def test_llm_return_logprobs_streaming_tp2(prompt_logprobs, logprobs, return_context_logits, - return_generation_logits): - llm_return_logprobs_test_harness(prompt_logprobs, - logprobs, - return_context_logits, - return_generation_logits, - streaming=True, - backend="pytorch", - tp_size=2) + return_generation_logits, + shared_mpi_session_2gpu): + llm_return_logprobs_test_harness( + prompt_logprobs, + logprobs, + return_context_logits, + return_generation_logits, + streaming=True, + backend="pytorch", + tp_size=2, + **_mpi_session_kwargs(shared_mpi_session_2gpu)) @skip_ray @@ -161,7 +201,7 @@ def test_llm_return_logprobs_streaming_tp2(prompt_logprobs, logprobs, ], ) def test_llm_get_stats_pp2(return_context_logits, enable_chunked_prefill, - enable_iter_req_stats): + enable_iter_req_stats, shared_mpi_session_2gpu): llm_get_stats_test_harness( tp_size=1, pp_size=2, @@ -169,6 +209,7 @@ def test_llm_get_stats_pp2(return_context_logits, enable_chunked_prefill, pytorch_backend=True, enable_chunked_prefill=enable_chunked_prefill, enable_iter_req_stats=enable_iter_req_stats, + **_mpi_session_kwargs(shared_mpi_session_2gpu), ) @@ -182,7 +223,7 @@ def test_llm_get_stats_pp2(return_context_logits, enable_chunked_prefill, ], ) def test_llm_get_stats_pp4(return_context_logits, enable_chunked_prefill, - enable_iter_req_stats): + enable_iter_req_stats, shared_mpi_session_4gpu): llm_get_stats_test_harness( tp_size=1, pp_size=4, @@ -190,154 +231,31 @@ def test_llm_get_stats_pp4(return_context_logits, enable_chunked_prefill, pytorch_backend=True, enable_chunked_prefill=enable_chunked_prefill, enable_iter_req_stats=enable_iter_req_stats, + **_mpi_session_kwargs(shared_mpi_session_4gpu), ) @skip_ray @pytest.mark.gpu2 -def test_llm_get_stats_tp2(): - llm_get_stats_test_harness(tp_size=2, pytorch_backend=True) +def test_llm_get_stats_tp2(shared_mpi_session_2gpu): + llm_get_stats_test_harness(tp_size=2, + pytorch_backend=True, + **_mpi_session_kwargs(shared_mpi_session_2gpu)) @skip_ray @pytest.mark.gpu2 -def test_llm_get_stats_async_tp2(): - llm_get_stats_async_test_harness(tp_size=2, pytorch_backend=True) +def test_llm_get_stats_async_tp2(shared_mpi_session_2gpu): + llm_get_stats_async_test_harness( + tp_size=2, + pytorch_backend=True, + **_mpi_session_kwargs(shared_mpi_session_2gpu)) @skip_ray @pytest.mark.gpu2 -def test_llm_get_stats_async_pp2(): - llm_get_stats_async_test_harness(pp_size=2, pytorch_backend=True) - - -def _run_grouped_mpi_cases( - n_workers: int, cases: tuple[tuple[str, Callable[[object], None]], - ...]) -> None: - from tensorrt_llm.llmapi.mpi_session import MpiPoolSession - - mpi_session = MpiPoolSession(n_workers=n_workers) - try: - for case_name, run_case in cases: - print(f"Running grouped multi-GPU PyTorch LLM case: {case_name}") - try: - run_case(mpi_session) - except pytest.skip.Exception: - raise - except Exception as exc: - raise AssertionError( - f"Grouped multi-GPU PyTorch LLM case failed: {case_name}" - ) from exc - finally: - mpi_session.shutdown() - - -def test_llm_multi_gpu_pytorch_grouped_mpi_reuse(): - gpu2_cases = ( - ("test_llm_capture_request_error", - lambda mpi_session: _test_llm_capture_request_error( - pytorch_backend=True, tp_size=2, _mpi_session=mpi_session)), - ("test_tinyllama_logits_processor_2gpu[tp1pp2]", lambda mpi_session: - tinyllama_logits_processor_test_harness(backend="pytorch", - tensor_parallel_size=1, - pipeline_parallel_size=2, - _mpi_session=mpi_session)), - ("test_tinyllama_logits_processor_2gpu[tp2pp1]", lambda mpi_session: - tinyllama_logits_processor_test_harness(backend="pytorch", - tensor_parallel_size=2, - pipeline_parallel_size=1, - _mpi_session=mpi_session)), - ("test_llama_7b_lora_tp2", - lambda mpi_session: llama_7b_lora_from_dir_test_harness( - tensor_parallel_size=2, - kv_cache_config=global_kv_cache_config, - _mpi_session=mpi_session)), - ("test_llm_rpc_tp2", - lambda mpi_session: _run_llm_rpc_tp2(_mpi_session=mpi_session)), - ("test_llm_rpc_streaming_tp2", lambda mpi_session: asyncio.run( - _run_llm_rpc_streaming_tp2(_mpi_session=mpi_session))), - ("test_llm_return_logprobs_streaming_tp2", lambda mpi_session: - llm_return_logprobs_test_harness(None, - 1, - False, - False, - streaming=True, - backend="pytorch", - tp_size=2, - _mpi_session=mpi_session)), - ("test_llm_get_stats_pp2[no_chunked]", lambda mpi_session: - llm_get_stats_test_harness(tp_size=1, - pp_size=2, - return_context_logits=False, - pytorch_backend=True, - enable_chunked_prefill=False, - enable_iter_req_stats=True, - _mpi_session=mpi_session)), - ("test_llm_get_stats_pp2[chunked]", lambda mpi_session: - llm_get_stats_test_harness(tp_size=1, - pp_size=2, - return_context_logits=False, - pytorch_backend=True, - enable_chunked_prefill=True, - enable_iter_req_stats=True, - _mpi_session=mpi_session)), - ("test_llm_get_stats_tp2", - lambda mpi_session: llm_get_stats_test_harness( - tp_size=2, pytorch_backend=True, _mpi_session=mpi_session)), - ("test_llm_get_stats_async_tp2", - lambda mpi_session: llm_get_stats_async_test_harness( - tp_size=2, pytorch_backend=True, _mpi_session=mpi_session)), - ("test_llm_get_stats_async_pp2", - lambda mpi_session: llm_get_stats_async_test_harness( - pp_size=2, pytorch_backend=True, _mpi_session=mpi_session)), - ) - gpu4_cases = ( - ("test_tinyllama_logits_processor_tp2pp2", lambda mpi_session: - tinyllama_logits_processor_test_harness(backend="pytorch", - tensor_parallel_size=2, - pipeline_parallel_size=2, - _mpi_session=mpi_session)), - ("test_llama_7b_multi_lora_tp4[cuda_graph]", lambda mpi_session: - check_llama_7b_multi_lora_from_request_test_harness( - LLM, - lora_config=LoraConfig(lora_target_modules= - ['attn_q', 'attn_k', 'attn_v'], - max_lora_rank=8, - max_loras=1, - max_cpu_loras=8), - tensor_parallel_size=4, - kv_cache_config=global_kv_cache_config, - cuda_graph_config=CudaGraphConfig(max_batch_size=10), - _mpi_session=mpi_session)), - ("test_llama_7b_multi_lora_tp4[no_cuda_graph]", lambda mpi_session: - check_llama_7b_multi_lora_from_request_test_harness( - LLM, - lora_config=LoraConfig(lora_target_modules= - ['attn_q', 'attn_k', 'attn_v'], - max_lora_rank=8, - max_loras=1, - max_cpu_loras=8), - tensor_parallel_size=4, - kv_cache_config=global_kv_cache_config, - cuda_graph_config=None, - _mpi_session=mpi_session)), - ("test_llm_get_stats_pp4[no_chunked]", lambda mpi_session: - llm_get_stats_test_harness(tp_size=1, - pp_size=4, - return_context_logits=False, - pytorch_backend=True, - enable_chunked_prefill=False, - enable_iter_req_stats=True, - _mpi_session=mpi_session)), - ("test_llm_get_stats_pp4[chunked]", lambda mpi_session: - llm_get_stats_test_harness(tp_size=1, - pp_size=4, - return_context_logits=False, - pytorch_backend=True, - enable_chunked_prefill=True, - enable_iter_req_stats=True, - _mpi_session=mpi_session)), - ) - - _run_grouped_mpi_cases(2, gpu2_cases) - _run_grouped_mpi_cases(4, gpu4_cases) +def test_llm_get_stats_async_pp2(shared_mpi_session_2gpu): + llm_get_stats_async_test_harness( + pp_size=2, + pytorch_backend=True, + **_mpi_session_kwargs(shared_mpi_session_2gpu)) From 125a861b2f575c26266403339c5928eb5914985b Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 30 Jun 2026 19:36:35 -0700 Subject: [PATCH 04/20] [None][test] Parameterize DeepSeek NVFP4 pre-merge group Signed-off-by: qgai --- .../defs/accuracy/test_llm_api_pytorch.py | 109 +++++++++--------- 1 file changed, 52 insertions(+), 57 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index f18d98fbe117..ec935be800e9 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -26,6 +26,7 @@ from tensorrt_llm import LLM from tensorrt_llm._torch.model_config import MoeLoadBalancerConfig +from tensorrt_llm._utils import mpi_disabled # isort: off from tensorrt_llm.llmapi import ( @@ -74,6 +75,47 @@ def patched_start_mpi_pool(self): patched_start_mpi_pool) +def _restore_env_var(name: str, value: str | None) -> None: + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + +def _shared_mpi_session(n_workers: int): + if mpi_disabled(): + yield None + return + + from tensorrt_llm.llmapi.mpi_session import MpiPoolSession + + mpi_session = MpiPoolSession(n_workers=n_workers) + try: + yield mpi_session + finally: + mpi_session.shutdown() + + +@pytest.fixture(scope="module") +def shared_mpi_session_4gpu(): + yield from _shared_mpi_session(4) + + +@pytest.fixture(scope="module") +def hf_weight_cache(): + previous_cache_env = os.environ.get("TRTLLM_HF_WEIGHT_CACHE") + previous_cache_entries_env = os.environ.get( + "TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES") + os.environ["TRTLLM_HF_WEIGHT_CACHE"] = "1" + os.environ["TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES"] = "1" + try: + yield + finally: + _restore_env_var("TRTLLM_HF_WEIGHT_CACHE", previous_cache_env) + _restore_env_var("TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES", + previous_cache_entries_env) + + def _get_default_torch_compile_config(torch_compile): return TorchCompileConfig(enable_fullgraph=True, enable_piecewise_cuda_graph=True, @@ -1495,20 +1537,6 @@ class TestDeepSeekV3Lite(LlmapiAccuracyTestHarness): low_precision_combine=True, torch_compile=False), ) - _NVFP4_4GPU_PREMERGE_GROUPS = ( - ("tp4_mtp0", - ("cutlass_mtp0_tp4_compile_off", "cutlass_mtp0_tp4_compile_on", - "trtllm_mtp0_tp4_compile_off", "cutlass_mtp0_tp4_lpc_compile_off")), - ("ep4_mtp0", ("cutlass_mtp0_ep4_compile_on", - "trtllm_mtp0_ep4_compile_off")), - ("tp2pp2_mtp0", ("cutlass_mtp0_tp2pp2_compile_off", - "cutlass_mtp0_tp2pp2_compile_on")), - ("tp4_mtp2", ("cutlass_mtp2_tp4_compile_off", )), - ("ep4_mtp2", ("trtllm_mtp2_ep4_compile_off", )), - ("pp4_mtp2", ("cutlass_mtp2_pp4_compile_off", )), - ("pp4_mtp0", ("cutlass_mtp0_pp4_compile_off", - "cutlass_mtp0_pp4_compile_on")), - ) @pytest.mark.skip_less_device_memory(60000) @parametrize_with_ids("v2_kv_cache", [True, False]) @@ -2395,49 +2423,16 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, @pytest.mark.skip_less_device(4) @skip_pre_blackwell - def test_nvfp4_4gpus_premerge_grouped(self): - from tensorrt_llm.llmapi.mpi_session import MpiPoolSession - - cases_by_id = { - case["id"]: case - for case in self._NVFP4_4GPU_PREMERGE_CASES - } - grouped_case_ids = tuple( - case_id for _, case_ids in self._NVFP4_4GPU_PREMERGE_GROUPS - for case_id in case_ids) - assert len(grouped_case_ids) == len(set(grouped_case_ids)) - assert set(grouped_case_ids) == set(cases_by_id) - - previous_cache_env = os.environ.get("TRTLLM_HF_WEIGHT_CACHE") - os.environ["TRTLLM_HF_WEIGHT_CACHE"] = "1" - try: - for group_id, case_ids in self._NVFP4_4GPU_PREMERGE_GROUPS: - mpi_session = (MpiPoolSession( - n_workers=4) if len(case_ids) > 1 else None) - try: - for case_id in case_ids: - case = cases_by_id[case_id] - case_kwargs = { - key: value - for key, value in case.items() if key != "id" - } - try: - self._run_nvfp4_4gpus_case(_mpi_session=mpi_session, - **case_kwargs) - except pytest.skip.Exception: - raise - except Exception as exc: - raise AssertionError( - "DeepSeek V3 Lite NVFP4 grouped case failed: " - f"{group_id}/{case_id}") from exc - finally: - if mpi_session is not None: - mpi_session.shutdown() - finally: - if previous_cache_env is None: - os.environ.pop("TRTLLM_HF_WEIGHT_CACHE", None) - else: - os.environ["TRTLLM_HF_WEIGHT_CACHE"] = previous_cache_env + @pytest.mark.parametrize( + "case", + _NVFP4_4GPU_PREMERGE_CASES, + ids=[case["id"] for case in _NVFP4_4GPU_PREMERGE_CASES], + ) + def test_nvfp4_4gpus_premerge_grouped(self, case, shared_mpi_session_4gpu, + hf_weight_cache): + case_kwargs = {key: value for key, value in case.items() if key != "id"} + self._run_nvfp4_4gpus_case(_mpi_session=shared_mpi_session_4gpu, + **case_kwargs) def _run_nvfp4_4gpus_case(self, *, From 6b92142fb4a625d588855713cfbd40f0f5cfd2c8 Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 30 Jun 2026 23:09:41 -0700 Subject: [PATCH 05/20] [None][fix] Reuse PP NCCL communicator across LLMs sharing worker processes When multiple LLM instances share the same worker processes (e.g. a reused MpiPoolSession), the module-global _pp_comm was reassigned on each new pp>1 LLM, dropping the previous NcclCommunicatorOp and triggering a collective ncclCommDestroy at an unsynchronized point during the next model build. On reused workers this can deadlock in commDestroySync. Since the PP comm is a world communicator depending only on (world_size, rank), reuse the existing one when world_size matches and only refresh the routing mapping. Single-LLM runs are unaffected. Signed-off-by: qgai --- tensorrt_llm/_torch/distributed/communicator.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 6fbbf16a19df..289795c63820 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -1246,6 +1246,18 @@ def init_pp_comm(mapping): global _pp_comm if mpi_disabled(): _pp_comm = PPCommTorch(mapping) + elif isinstance(_pp_comm, PPCommNCCL) and \ + _pp_comm.mapping.world_size == mapping.world_size: + # Reuse the existing world NCCL communicator across LLM instances that + # share the same worker processes (e.g. a reused MpiPoolSession). The + # underlying comm depends only on (world_size, rank) -- it is a world + # communicator, independent of the pp/tp/ep layout -- so only the + # routing mapping needs refreshing. Recreating it would drop the old + # comm and trigger a collective ncclCommDestroy at an unsynchronized + # point during the next model build, which can deadlock on reused + # workers. Single-LLM (production) runs are unaffected: _pp_comm starts + # as None, so the first call still constructs a fresh PPCommNCCL. + _pp_comm.mapping = mapping else: _pp_comm = PPCommNCCL(mapping) init_helix_cp_comm(mapping) From eed25e3a9df8a7303e6600cfaaa9277d42d0a201 Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 30 Jun 2026 23:10:03 -0700 Subject: [PATCH 06/20] [None][test] Harden shared-MPI-session reuse for grouped NVFP4 tests - Fix weight-cache env ordering: shared_mpi_session_4gpu now depends on hf_weight_cache so TRTLLM_HF_WEIGHT_CACHE is exported before the MPI pool spawns (MpiPoolSession snapshots env at spawn time); add an assertion to pin the contract. - Add a _make_shared_llm factory (shared_llm_* fixtures) so tests inject the shared MPI session transparently, and migrate the RPC/NVFP4 direct-construction tests onto it. - Mark the grouped NVFP4 test threadleak(enabled=False): the reused session keeps mpi4py's _manager_spawn thread alive by design. - Only compute the HF weight-cache key when the cache is enabled (avoid a per-shard os.stat on the default path). Signed-off-by: qgai --- .../models/checkpoints/hf/weight_loader.py | 20 +++-- .../defs/accuracy/test_llm_api_pytorch.py | 82 +++++++++++++------ .../llmapi/test_llm_multi_gpu_pytorch.py | 59 +++++++------ 3 files changed, 105 insertions(+), 56 deletions(-) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index 9390d4abf3e7..2eb62e0837bb 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -153,9 +153,12 @@ def load_weights(self, if len(filtered_weight_files) > 0: weight_files = filtered_weight_files if weight_files: - cache_key = self._weight_files_cache_key(weight_files, - use_consolidated) - if self._is_weight_cache_enabled(): + # Only fingerprint the files (per-file os.stat) when the cache is + # enabled; the default path must not pay for a key it never uses. + cache_key = self._weight_files_cache_key( + weight_files, + use_consolidated) if self._is_weight_cache_enabled() else None + if cache_key is not None: cached_weights = self._get_cached_weights(cache_key) if cached_weights is not None: logger.info( @@ -193,7 +196,7 @@ def load_weights(self, weights = self._load_weights_in_parallel( weight_files, self._load_safetensors_file, "Loading safetensors weights in parallel") - if self._is_weight_cache_enabled(): + if cache_key is not None: self._cache_loaded_weights(cache_key, weights) return weights @@ -202,9 +205,10 @@ def load_weights(self, weight_files = glob.glob(f"{checkpoint_dir}/*.pth") if weight_files: - cache_key = self._weight_files_cache_key(weight_files, - use_consolidated) - if self._is_weight_cache_enabled(): + cache_key = self._weight_files_cache_key( + weight_files, + use_consolidated) if self._is_weight_cache_enabled() else None + if cache_key is not None: cached_weights = self._get_cached_weights(cache_key) if cached_weights is not None: logger.info( @@ -215,7 +219,7 @@ def load_weights(self, weights = self._load_weights_in_parallel( weight_files, self._load_bin_or_path_file, "Loading bin weights in parallel") - if self._is_weight_cache_enabled(): + if cache_key is not None: self._cache_loaded_weights(cache_key, weights) return weights diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index ec935be800e9..bc040059a52b 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -96,11 +96,6 @@ def _shared_mpi_session(n_workers: int): mpi_session.shutdown() -@pytest.fixture(scope="module") -def shared_mpi_session_4gpu(): - yield from _shared_mpi_session(4) - - @pytest.fixture(scope="module") def hf_weight_cache(): previous_cache_env = os.environ.get("TRTLLM_HF_WEIGHT_CACHE") @@ -116,6 +111,43 @@ def hf_weight_cache(): previous_cache_entries_env) +@pytest.fixture(scope="module") +def shared_mpi_session_4gpu(hf_weight_cache): + # Ordering is load-bearing: MpiPoolSession snapshots TRTLLM*/TLLM* env at + # spawn time and passes that copy to the workers. Any env exported AFTER the + # pool spawns never reaches them. Depending on hf_weight_cache guarantees + # TRTLLM_HF_WEIGHT_CACHE is exported before MpiPoolSession spawns, otherwise + # the weight cache would be silently disabled in the workers (which are the + # processes that actually load weights) and only session reuse would speed + # things up. The assertion pins this contract so a future fixture reorder + # fails loudly instead of silently regressing the optimization. + assert os.environ.get("TRTLLM_HF_WEIGHT_CACHE") == "1", ( + "hf_weight_cache must be set up before the shared MPI pool spawns so " + "the workers inherit TRTLLM_HF_WEIGHT_CACHE at spawn time.") + yield from _shared_mpi_session(4) + + +def _make_shared_llm(mpi_session): + """Return an LLM factory that transparently injects a shared MPI session. + + Tests build the LLM by calling this factory exactly like ``LLM(...)``; the + shared session (if any) is passed through without the test having to know it + exists. Falls back to a private per-LLM session when ``mpi_session`` is None. + """ + + def shared_llm(*args, **kwargs): + if mpi_session is not None: + kwargs["_mpi_session"] = mpi_session + return LLM(*args, **kwargs) + + return shared_llm + + +@pytest.fixture(scope="module") +def shared_llm_4gpu(shared_mpi_session_4gpu): + return _make_shared_llm(shared_mpi_session_4gpu) + + def _get_default_torch_compile_config(torch_compile): return TorchCompileConfig(enable_fullgraph=True, enable_piecewise_cuda_graph=True, @@ -2423,16 +2455,22 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, @pytest.mark.skip_less_device(4) @skip_pre_blackwell + # The shared module-scoped MPI session is intentionally kept alive across the + # parametrized cases (that is the whole point of reuse), so mpi4py's + # `_manager_spawn` management thread outlives each test. Disable the + # threadleak check here, consistent with other session/proxy-based tests. + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize( "case", _NVFP4_4GPU_PREMERGE_CASES, ids=[case["id"] for case in _NVFP4_4GPU_PREMERGE_CASES], ) - def test_nvfp4_4gpus_premerge_grouped(self, case, shared_mpi_session_4gpu, - hf_weight_cache): + def test_nvfp4_4gpus_premerge_grouped(self, case, shared_llm_4gpu): + # shared_llm_4gpu injects the shared 4-GPU MPI session and (transitively) + # depends on hf_weight_cache, so the weight-cache env is exported before + # the pool spawns (see the shared_mpi_session_4gpu fixture for details). case_kwargs = {key: value for key, value in case.items() if key != "id"} - self._run_nvfp4_4gpus_case(_mpi_session=shared_mpi_session_4gpu, - **case_kwargs) + self._run_nvfp4_4gpus_case(make_llm=shared_llm_4gpu, **case_kwargs) def _run_nvfp4_4gpus_case(self, *, @@ -2447,7 +2485,7 @@ def _run_nvfp4_4gpus_case(self, torch_compile, mtp_nextn, moe_backend, - _mpi_session=None): + make_llm=LLM): sm_version = get_sm_version() if moe_backend == "TRTLLM" and sm_version in (120, 121): pytest.skip(f"{moe_backend} backend does not support SM 120 or 121") @@ -2474,20 +2512,16 @@ def _run_nvfp4_4gpus_case(self, if fp8kv: kv_cache_config.dtype = "fp8" - llm_kwargs = dict( - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - kv_cache_config=kv_cache_config, - **pytorch_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config, - ) - if _mpi_session is not None: - llm_kwargs["_mpi_session"] = _mpi_session - - with LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", - **llm_kwargs) as llm: + with make_llm( + f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + **pytorch_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 task = GSM8K(self.MODEL_NAME) diff --git a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py index d80233c8bb8f..a14cbbce931f 100644 --- a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py +++ b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py @@ -50,6 +50,27 @@ def _mpi_session_kwargs(mpi_session) -> dict: return {"_mpi_session": mpi_session} if mpi_session is not None else {} +def _make_shared_llm(mpi_session): + """Return an LLM factory that transparently injects a shared MPI session. + + Tests that build the LLM directly can request the ``shared_llm_*`` fixture + and call it exactly like ``LLM(...)`` -- the shared session is passed through + without the test having to know it exists. Harness-based tests inject the + session via ``**_mpi_session_kwargs(...)`` instead, since the harness owns + LLM construction. + """ + + def shared_llm(*args, **kwargs): + return LLM(*args, **kwargs, **_mpi_session_kwargs(mpi_session)) + + return shared_llm + + +@pytest.fixture(scope="module") +def shared_llm_2gpu(shared_mpi_session_2gpu): + return _make_shared_llm(shared_mpi_session_2gpu) + + @pytest.mark.gpu2 def test_llm_capture_request_error(shared_mpi_session_2gpu): _test_llm_capture_request_error( @@ -118,27 +139,22 @@ def test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1( @skip_ray @pytest.mark.gpu2 -def test_llm_rpc_tp2(shared_mpi_session_2gpu): - _run_llm_rpc_tp2(shared_mpi_session_2gpu) +def test_llm_rpc_tp2(shared_llm_2gpu): + _run_llm_rpc_tp2(shared_llm_2gpu) @skip_ray @pytest.mark.gpu2 @pytest.mark.asyncio -async def test_llm_rpc_streaming_tp2(shared_mpi_session_2gpu): - await _run_llm_rpc_streaming_tp2(shared_mpi_session_2gpu) +async def test_llm_rpc_streaming_tp2(shared_llm_2gpu): + await _run_llm_rpc_streaming_tp2(shared_llm_2gpu) -def _run_llm_rpc_tp2(_mpi_session=None): - llm_kwargs = dict( - model=llama_model_path, - kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), - orchestrator_type="rpc", - tensor_parallel_size=2) - if _mpi_session is not None: - llm_kwargs["_mpi_session"] = _mpi_session - - with LLM(**llm_kwargs) as llm: +def _run_llm_rpc_tp2(make_llm=LLM): + with make_llm(model=llama_model_path, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), + orchestrator_type="rpc", + tensor_parallel_size=2) as llm: assert isinstance(llm._executor, GenerationExecutorRpcProxy) res = llm.generate("Tell me a joke", @@ -150,16 +166,11 @@ def _run_llm_rpc_tp2(_mpi_session=None): assert len(res.outputs[0].token_ids) == 10 -async def _run_llm_rpc_streaming_tp2(_mpi_session=None): - llm_kwargs = dict( - model=llama_model_path, - kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), - orchestrator_type="rpc", - tensor_parallel_size=2) - if _mpi_session is not None: - llm_kwargs["_mpi_session"] = _mpi_session - - with LLM(**llm_kwargs) as llm: +async def _run_llm_rpc_streaming_tp2(make_llm=LLM): + with make_llm(model=llama_model_path, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), + orchestrator_type="rpc", + tensor_parallel_size=2) as llm: assert isinstance(llm._executor, GenerationExecutorRpcProxy) async for output in llm.generate_async("Tell me a joke", From 37590716f836880ecdd01b93635b507bbd5bacfb Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 03:20:00 -0700 Subject: [PATCH 07/20] [None][test] Reset torch.compile state between grouped NVFP4 cases On a reused MpiPoolSession, Dynamo's per-code-object recompile counter is process-global in the worker and never reset across LLMs. Each torch_compile case recompiles the same model.forward under new guards; the count accumulates and eventually trips recompile_limit (16), which is a HARD failure under enable_fullgraph=True (FailOnRecompileLimitHit) and aborts the whole MPI job (observed as a silent crash at the ~10th case). Submit torch._dynamo.reset() to each worker after every case so compile state starts clean per LLM. Test-only; production (non-reused) paths are untouched. Verified: full 13-case grouped run now passes 13/13 on B300. Signed-off-by: qgai --- .../defs/accuracy/test_llm_api_pytorch.py | 53 ++++++++++++++----- 1 file changed, 40 insertions(+), 13 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index bc040059a52b..1141f8911566 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -127,12 +127,30 @@ def shared_mpi_session_4gpu(hf_weight_cache): yield from _shared_mpi_session(4) +def _reset_worker_torch_compile_state(): + """Reset per-worker torch.compile / Dynamo state (runs inside each worker). + + Dynamo's recompile counter is process-global and per-code-object. When + worker processes are reused across LLMs (shared MpiPoolSession), each + torch_compile case recompiles the same model.forward code object under new + guards; the count accumulates and eventually trips + `recompile_limit` (16), which is a HARD failure under `fullgraph=True` + (FailOnRecompileLimitHit) and aborts the whole MPI job. Resetting between + cases makes each LLM start from a clean compile cache, like a fresh process. + """ + import torch + torch._dynamo.reset() + + def _make_shared_llm(mpi_session): """Return an LLM factory that transparently injects a shared MPI session. Tests build the LLM by calling this factory exactly like ``LLM(...)``; the shared session (if any) is passed through without the test having to know it exists. Falls back to a private per-LLM session when ``mpi_session`` is None. + + The factory carries the ``mpi_session`` so callers can reset per-worker + compile state between cases without threading the session separately. """ def shared_llm(*args, **kwargs): @@ -140,6 +158,7 @@ def shared_llm(*args, **kwargs): kwargs["_mpi_session"] = mpi_session return LLM(*args, **kwargs) + shared_llm.mpi_session = mpi_session return shared_llm @@ -2512,20 +2531,28 @@ def _run_nvfp4_4gpus_case(self, if fp8kv: kv_cache_config.dtype = "fp8" - with make_llm( - f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - kv_cache_config=kv_cache_config, - **pytorch_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config, - ) as llm: - assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 + try: + with make_llm( + f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + **pytorch_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - task = GSM8K(self.MODEL_NAME) - task.evaluate(llm) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + finally: + # On a reused MPI session, reset each worker's torch.compile/Dynamo + # state so recompile counts don't accumulate across cases and trip + # the fullgraph recompile limit (FailOnRecompileLimitHit). + mpi_session = getattr(make_llm, "mpi_session", None) + if mpi_session is not None: + mpi_session.submit_sync(_reset_worker_torch_compile_state) @parametrize_with_ids( "fp8kv,attention_dp,cuda_graph,overlap_scheduler", From ec8bb232d39502a8b3022c679c0a322c2b2869bb Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 03:52:55 -0700 Subject: [PATCH 08/20] [None][test] Extract shared session-reuse helpers into a single in-package module The grouped multi-GPU session-reuse helpers (shared_mpi_session, make_shared_llm, mpi_session_kwargs, reset_worker_torch_compile_state, restore_env_var) were duplicated across the unittest and integration test roots, and had already begun to diverge. Consolidate them into tensorrt_llm/llmapi/_grouped_test_utils.py so both roots (which have separate sys.paths and cannot share a module under tests/) import one source. Heavy imports are lazy to avoid circular imports at package init. Test-infra only. Verified: package import + pytest collection resolve (13 tests collected). Signed-off-by: qgai --- tensorrt_llm/llmapi/_grouped_test_utils.py | 84 +++++++++++++++++++ .../defs/accuracy/test_llm_api_pytorch.py | 63 ++------------ .../llmapi/test_llm_multi_gpu_pytorch.py | 41 ++------- 3 files changed, 98 insertions(+), 90 deletions(-) create mode 100644 tensorrt_llm/llmapi/_grouped_test_utils.py diff --git a/tensorrt_llm/llmapi/_grouped_test_utils.py b/tensorrt_llm/llmapi/_grouped_test_utils.py new file mode 100644 index 000000000000..14deda1110a3 --- /dev/null +++ b/tensorrt_llm/llmapi/_grouped_test_utils.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared helpers for grouped multi-GPU tests that reuse one ``MpiPoolSession`` +across many ``LLM`` instances. + +This lives in-package (rather than under ``tests/``) because the unittest and +integration test roots have separate ``sys.path`` bases and cannot import a +single shared module otherwise. It is test-infrastructure only and must not be +used on production inference paths. Heavy imports are done lazily so importing +this module never triggers a circular import during package initialization. +""" + +import os +from typing import Optional + + +def restore_env_var(name: str, value: Optional[str]) -> None: + """Restore an env var to a previously captured value (or remove it).""" + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + +def shared_mpi_session(n_workers: int): + """Generator yielding a shared ``MpiPoolSession`` (or ``None`` when MPI is + disabled), shutting it down on teardown. Intended to back a module-scoped + pytest fixture so a pool of workers is reused across many LLMs.""" + from tensorrt_llm._utils import mpi_disabled + + if mpi_disabled(): + yield None + return + + from tensorrt_llm.llmapi.mpi_session import MpiPoolSession + + mpi_session = MpiPoolSession(n_workers=n_workers) + try: + yield mpi_session + finally: + mpi_session.shutdown() + + +def mpi_session_kwargs(mpi_session) -> dict: + """LLM kwargs that inject a shared MPI session, or ``{}`` when there is + none (e.g. MPI disabled). Lets harness-based tests forward the session.""" + return {"_mpi_session": mpi_session} if mpi_session is not None else {} + + +def reset_worker_torch_compile_state() -> None: + """Reset per-worker torch.compile / Dynamo state (runs inside each worker). + + Dynamo's recompile counter is process-global and per-code-object. When + worker processes are reused across LLMs (shared ``MpiPoolSession``), each + ``torch_compile`` case recompiles the same ``model.forward`` code object + under new guards; the count accumulates and eventually trips + ``recompile_limit`` (16), which is a HARD failure under ``fullgraph=True`` + (``FailOnRecompileLimitHit``) and aborts the whole MPI job. Resetting + between cases makes each LLM start from a clean compile cache, like a fresh + process. Submit this to each worker via ``mpi_session.submit_sync(...)``. + """ + import torch + + torch._dynamo.reset() + + +def make_shared_llm(mpi_session): + """Return an ``LLM`` factory that transparently injects a shared MPI session. + + Tests build the LLM by calling the factory exactly like ``LLM(...)``; the + shared session (if any) is passed through without the test having to know it + exists. Falls back to a private per-LLM session when ``mpi_session`` is None. + The factory exposes ``.mpi_session`` so callers can reset per-worker compile + state between cases without threading the session separately. + """ + from tensorrt_llm import LLM + + def shared_llm(*args, **kwargs): + if mpi_session is not None: + kwargs["_mpi_session"] = mpi_session + return LLM(*args, **kwargs) + + shared_llm.mpi_session = mpi_session + return shared_llm diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 1141f8911566..152a6179455d 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -26,7 +26,6 @@ from tensorrt_llm import LLM from tensorrt_llm._torch.model_config import MoeLoadBalancerConfig -from tensorrt_llm._utils import mpi_disabled # isort: off from tensorrt_llm.llmapi import ( @@ -75,25 +74,14 @@ def patched_start_mpi_pool(self): patched_start_mpi_pool) -def _restore_env_var(name: str, value: str | None) -> None: - if value is None: - os.environ.pop(name, None) - else: - os.environ[name] = value - - -def _shared_mpi_session(n_workers: int): - if mpi_disabled(): - yield None - return - - from tensorrt_llm.llmapi.mpi_session import MpiPoolSession - - mpi_session = MpiPoolSession(n_workers=n_workers) - try: - yield mpi_session - finally: - mpi_session.shutdown() +from tensorrt_llm.llmapi._grouped_test_utils import \ + make_shared_llm as _make_shared_llm # noqa: E402 +from tensorrt_llm.llmapi._grouped_test_utils import \ + reset_worker_torch_compile_state as _reset_worker_torch_compile_state +from tensorrt_llm.llmapi._grouped_test_utils import \ + restore_env_var as _restore_env_var +from tensorrt_llm.llmapi._grouped_test_utils import \ + shared_mpi_session as _shared_mpi_session @pytest.fixture(scope="module") @@ -127,41 +115,6 @@ def shared_mpi_session_4gpu(hf_weight_cache): yield from _shared_mpi_session(4) -def _reset_worker_torch_compile_state(): - """Reset per-worker torch.compile / Dynamo state (runs inside each worker). - - Dynamo's recompile counter is process-global and per-code-object. When - worker processes are reused across LLMs (shared MpiPoolSession), each - torch_compile case recompiles the same model.forward code object under new - guards; the count accumulates and eventually trips - `recompile_limit` (16), which is a HARD failure under `fullgraph=True` - (FailOnRecompileLimitHit) and aborts the whole MPI job. Resetting between - cases makes each LLM start from a clean compile cache, like a fresh process. - """ - import torch - torch._dynamo.reset() - - -def _make_shared_llm(mpi_session): - """Return an LLM factory that transparently injects a shared MPI session. - - Tests build the LLM by calling this factory exactly like ``LLM(...)``; the - shared session (if any) is passed through without the test having to know it - exists. Falls back to a private per-LLM session when ``mpi_session`` is None. - - The factory carries the ``mpi_session`` so callers can reset per-worker - compile state between cases without threading the session separately. - """ - - def shared_llm(*args, **kwargs): - if mpi_session is not None: - kwargs["_mpi_session"] = mpi_session - return LLM(*args, **kwargs) - - shared_llm.mpi_session = mpi_session - return shared_llm - - @pytest.fixture(scope="module") def shared_llm_4gpu(shared_mpi_session_4gpu): return _make_shared_llm(shared_mpi_session_4gpu) diff --git a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py index a14cbbce931f..32178e1119bf 100644 --- a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py +++ b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py @@ -2,9 +2,14 @@ from utils.util import skip_ray from tensorrt_llm import LLM -from tensorrt_llm._utils import mpi_disabled from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.llmapi._grouped_test_utils import \ + make_shared_llm as _make_shared_llm +from tensorrt_llm.llmapi._grouped_test_utils import \ + mpi_session_kwargs as _mpi_session_kwargs +from tensorrt_llm.llmapi._grouped_test_utils import \ + shared_mpi_session as _shared_mpi_session from tensorrt_llm.lora_helper import LoraConfig from tensorrt_llm.sampling_params import SamplingParams @@ -22,20 +27,6 @@ global_kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4) -def _shared_mpi_session(n_workers: int): - if mpi_disabled(): - yield None - return - - from tensorrt_llm.llmapi.mpi_session import MpiPoolSession - - mpi_session = MpiPoolSession(n_workers=n_workers) - try: - yield mpi_session - finally: - mpi_session.shutdown() - - @pytest.fixture(scope="module") def shared_mpi_session_2gpu(): yield from _shared_mpi_session(2) @@ -46,26 +37,6 @@ def shared_mpi_session_4gpu(): yield from _shared_mpi_session(4) -def _mpi_session_kwargs(mpi_session) -> dict: - return {"_mpi_session": mpi_session} if mpi_session is not None else {} - - -def _make_shared_llm(mpi_session): - """Return an LLM factory that transparently injects a shared MPI session. - - Tests that build the LLM directly can request the ``shared_llm_*`` fixture - and call it exactly like ``LLM(...)`` -- the shared session is passed through - without the test having to know it exists. Harness-based tests inject the - session via ``**_mpi_session_kwargs(...)`` instead, since the harness owns - LLM construction. - """ - - def shared_llm(*args, **kwargs): - return LLM(*args, **kwargs, **_mpi_session_kwargs(mpi_session)) - - return shared_llm - - @pytest.fixture(scope="module") def shared_llm_2gpu(shared_mpi_session_2gpu): return _make_shared_llm(shared_mpi_session_2gpu) From 9f663e849639a2030549febbcc4c3f583faa0d19 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 04:06:05 -0700 Subject: [PATCH 09/20] [None][test] Extract hf_weight_cache_env into the shared helper module Move the HF-weight-cache env enable/restore logic into a reusable contextmanager (hf_weight_cache_env) in _grouped_test_utils.py, keeping only a thin pytest fixture wrapper in the test. The fixture itself stays test-side so the shipped package never imports pytest. Docstring documents the load-bearing ordering (must be entered before the MpiPoolSession spawns). Verified: contextmanager sets/restores env correctly; 13 tests collected. Signed-off-by: qgai --- tensorrt_llm/llmapi/_grouped_test_utils.py | 24 +++++++++++++++++++ .../defs/accuracy/test_llm_api_pytorch.py | 17 ++++--------- 2 files changed, 28 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/llmapi/_grouped_test_utils.py b/tensorrt_llm/llmapi/_grouped_test_utils.py index 14deda1110a3..047f8cb0712f 100644 --- a/tensorrt_llm/llmapi/_grouped_test_utils.py +++ b/tensorrt_llm/llmapi/_grouped_test_utils.py @@ -11,6 +11,7 @@ """ import os +from contextlib import contextmanager from typing import Optional @@ -22,6 +23,29 @@ def restore_env_var(name: str, value: Optional[str]) -> None: os.environ[name] = value +@contextmanager +def hf_weight_cache_env(max_entries: str = "1"): + """Enable the HF weight cache via env vars, restoring prior values on exit. + + IMPORTANT (load-bearing ordering): enter this BEFORE spawning the shared + ``MpiPoolSession``. ``MpiPoolSession`` snapshots ``TRTLLM*``/``TLLM*`` env at + spawn time and passes that copy to the workers (which are the processes that + actually load weights); enabling the cache AFTER the pool spawns leaves it + silently disabled in the workers. Back a module-scoped ``hf_weight_cache`` + fixture with this and have the session fixture depend on that fixture so the + env is exported first. + """ + prev = os.environ.get("TRTLLM_HF_WEIGHT_CACHE") + prev_entries = os.environ.get("TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES") + os.environ["TRTLLM_HF_WEIGHT_CACHE"] = "1" + os.environ["TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES"] = max_entries + try: + yield + finally: + restore_env_var("TRTLLM_HF_WEIGHT_CACHE", prev) + restore_env_var("TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES", prev_entries) + + def shared_mpi_session(n_workers: int): """Generator yielding a shared ``MpiPoolSession`` (or ``None`` when MPI is disabled), shutting it down on teardown. Intended to back a module-scoped diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 152a6179455d..f3b1a4bbec2e 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -75,28 +75,19 @@ def patched_start_mpi_pool(self): from tensorrt_llm.llmapi._grouped_test_utils import \ - make_shared_llm as _make_shared_llm # noqa: E402 + hf_weight_cache_env as _hf_weight_cache_env # noqa: E402 from tensorrt_llm.llmapi._grouped_test_utils import \ - reset_worker_torch_compile_state as _reset_worker_torch_compile_state + make_shared_llm as _make_shared_llm from tensorrt_llm.llmapi._grouped_test_utils import \ - restore_env_var as _restore_env_var + reset_worker_torch_compile_state as _reset_worker_torch_compile_state from tensorrt_llm.llmapi._grouped_test_utils import \ shared_mpi_session as _shared_mpi_session @pytest.fixture(scope="module") def hf_weight_cache(): - previous_cache_env = os.environ.get("TRTLLM_HF_WEIGHT_CACHE") - previous_cache_entries_env = os.environ.get( - "TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES") - os.environ["TRTLLM_HF_WEIGHT_CACHE"] = "1" - os.environ["TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES"] = "1" - try: + with _hf_weight_cache_env(): yield - finally: - _restore_env_var("TRTLLM_HF_WEIGHT_CACHE", previous_cache_env) - _restore_env_var("TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES", - previous_cache_entries_env) @pytest.fixture(scope="module") From 2a1ecd7407fc4b09e69cac6ba40d6ef8c0211a03 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 19:07:26 -0700 Subject: [PATCH 10/20] [None][test] Simplify multi-GPU session plumbing via mpi_kwargs fixtures Replace the repeated 'take shared_mpi_session_Ngpu + spread **_mpi_session_kwargs(...)' boilerplate in every harness-based test with module-scoped mpi_kwargs_2gpu / mpi_kwargs_4gpu fixtures that yield the ready-to-spread kwargs dict. Tests now just spread **mpi_kwargs_Ngpu and never reference the raw session or the helper; _mpi_session_kwargs is used only inside the two fixtures. RPC tests keep the shared_llm_2gpu factory. No behavior change. Signed-off-by: qgai --- .../llmapi/test_llm_multi_gpu_pytorch.py | 107 ++++++++++-------- 1 file changed, 57 insertions(+), 50 deletions(-) diff --git a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py index 32178e1119bf..60cb3cc29b77 100644 --- a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py +++ b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py @@ -39,50 +39,60 @@ def shared_mpi_session_4gpu(): @pytest.fixture(scope="module") def shared_llm_2gpu(shared_mpi_session_2gpu): + # Factory for tests that construct the LLM directly (see RPC tests below). return _make_shared_llm(shared_mpi_session_2gpu) +@pytest.fixture(scope="module") +def mpi_kwargs_2gpu(shared_mpi_session_2gpu): + # Ready-to-spread LLM kwargs that inject the shared 2-GPU session (or {} when + # MPI is disabled). Harness-based tests just spread this; they never touch the + # raw session object. + return _mpi_session_kwargs(shared_mpi_session_2gpu) + + +@pytest.fixture(scope="module") +def mpi_kwargs_4gpu(shared_mpi_session_4gpu): + return _mpi_session_kwargs(shared_mpi_session_4gpu) + + @pytest.mark.gpu2 -def test_llm_capture_request_error(shared_mpi_session_2gpu): - _test_llm_capture_request_error( - pytorch_backend=True, - tp_size=2, - **_mpi_session_kwargs(shared_mpi_session_2gpu)) +def test_llm_capture_request_error(mpi_kwargs_2gpu): + _test_llm_capture_request_error(pytorch_backend=True, + tp_size=2, + **mpi_kwargs_2gpu) @pytest.mark.gpu4 -def test_tinyllama_logits_processor_tp2pp2(shared_mpi_session_4gpu): - tinyllama_logits_processor_test_harness( - backend="pytorch", - tensor_parallel_size=2, - pipeline_parallel_size=2, - **_mpi_session_kwargs(shared_mpi_session_4gpu)) +def test_tinyllama_logits_processor_tp2pp2(mpi_kwargs_4gpu): + tinyllama_logits_processor_test_harness(backend="pytorch", + tensor_parallel_size=2, + pipeline_parallel_size=2, + **mpi_kwargs_4gpu) @pytest.mark.gpu2 @pytest.mark.part0 @pytest.mark.parametrize("tp_size, pp_size", [(1, 2), (2, 1)]) def test_tinyllama_logits_processor_2gpu(tp_size: int, pp_size: int, - shared_mpi_session_2gpu): - tinyllama_logits_processor_test_harness( - backend="pytorch", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - **_mpi_session_kwargs(shared_mpi_session_2gpu)) + mpi_kwargs_2gpu): + tinyllama_logits_processor_test_harness(backend="pytorch", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + **mpi_kwargs_2gpu) @pytest.mark.gpu2 -def test_llama_7b_lora_tp2(shared_mpi_session_2gpu): - llama_7b_lora_from_dir_test_harness( - tensor_parallel_size=2, - kv_cache_config=global_kv_cache_config, - **_mpi_session_kwargs(shared_mpi_session_2gpu)) +def test_llama_7b_lora_tp2(mpi_kwargs_2gpu): + llama_7b_lora_from_dir_test_harness(tensor_parallel_size=2, + kv_cache_config=global_kv_cache_config, + **mpi_kwargs_2gpu) @pytest.mark.gpu4 @skip_ray # https://nvbugs/5682551 @test_lora_with_and_without_cuda_graph -def test_llama_7b_multi_lora_tp4(cuda_graph_config, shared_mpi_session_4gpu): +def test_llama_7b_multi_lora_tp4(cuda_graph_config, mpi_kwargs_4gpu): # For LoRA checkpoints without finetuned embedding and lm_head, we can either: # (1) specify lora_target_modules, or # (2) provide a lora_dir to infer the lora_target_modules. @@ -96,7 +106,7 @@ def test_llama_7b_multi_lora_tp4(cuda_graph_config, shared_mpi_session_4gpu): tensor_parallel_size=4, kv_cache_config=global_kv_cache_config, cuda_graph_config=cuda_graph_config, - **_mpi_session_kwargs(shared_mpi_session_4gpu)) + **mpi_kwargs_4gpu) @skip_ray # https://nvbugs/5727075 @@ -161,16 +171,15 @@ async def _run_llm_rpc_streaming_tp2(make_llm=LLM): def test_llm_return_logprobs_streaming_tp2(prompt_logprobs, logprobs, return_context_logits, return_generation_logits, - shared_mpi_session_2gpu): - llm_return_logprobs_test_harness( - prompt_logprobs, - logprobs, - return_context_logits, - return_generation_logits, - streaming=True, - backend="pytorch", - tp_size=2, - **_mpi_session_kwargs(shared_mpi_session_2gpu)) + mpi_kwargs_2gpu): + llm_return_logprobs_test_harness(prompt_logprobs, + logprobs, + return_context_logits, + return_generation_logits, + streaming=True, + backend="pytorch", + tp_size=2, + **mpi_kwargs_2gpu) @skip_ray @@ -183,7 +192,7 @@ def test_llm_return_logprobs_streaming_tp2(prompt_logprobs, logprobs, ], ) def test_llm_get_stats_pp2(return_context_logits, enable_chunked_prefill, - enable_iter_req_stats, shared_mpi_session_2gpu): + enable_iter_req_stats, mpi_kwargs_2gpu): llm_get_stats_test_harness( tp_size=1, pp_size=2, @@ -191,7 +200,7 @@ def test_llm_get_stats_pp2(return_context_logits, enable_chunked_prefill, pytorch_backend=True, enable_chunked_prefill=enable_chunked_prefill, enable_iter_req_stats=enable_iter_req_stats, - **_mpi_session_kwargs(shared_mpi_session_2gpu), + **mpi_kwargs_2gpu, ) @@ -205,7 +214,7 @@ def test_llm_get_stats_pp2(return_context_logits, enable_chunked_prefill, ], ) def test_llm_get_stats_pp4(return_context_logits, enable_chunked_prefill, - enable_iter_req_stats, shared_mpi_session_4gpu): + enable_iter_req_stats, mpi_kwargs_4gpu): llm_get_stats_test_harness( tp_size=1, pp_size=4, @@ -213,31 +222,29 @@ def test_llm_get_stats_pp4(return_context_logits, enable_chunked_prefill, pytorch_backend=True, enable_chunked_prefill=enable_chunked_prefill, enable_iter_req_stats=enable_iter_req_stats, - **_mpi_session_kwargs(shared_mpi_session_4gpu), + **mpi_kwargs_4gpu, ) @skip_ray @pytest.mark.gpu2 -def test_llm_get_stats_tp2(shared_mpi_session_2gpu): +def test_llm_get_stats_tp2(mpi_kwargs_2gpu): llm_get_stats_test_harness(tp_size=2, pytorch_backend=True, - **_mpi_session_kwargs(shared_mpi_session_2gpu)) + **mpi_kwargs_2gpu) @skip_ray @pytest.mark.gpu2 -def test_llm_get_stats_async_tp2(shared_mpi_session_2gpu): - llm_get_stats_async_test_harness( - tp_size=2, - pytorch_backend=True, - **_mpi_session_kwargs(shared_mpi_session_2gpu)) +def test_llm_get_stats_async_tp2(mpi_kwargs_2gpu): + llm_get_stats_async_test_harness(tp_size=2, + pytorch_backend=True, + **mpi_kwargs_2gpu) @skip_ray @pytest.mark.gpu2 -def test_llm_get_stats_async_pp2(shared_mpi_session_2gpu): - llm_get_stats_async_test_harness( - pp_size=2, - pytorch_backend=True, - **_mpi_session_kwargs(shared_mpi_session_2gpu)) +def test_llm_get_stats_async_pp2(mpi_kwargs_2gpu): + llm_get_stats_async_test_harness(pp_size=2, + pytorch_backend=True, + **mpi_kwargs_2gpu) From c9a10156a9755b19a18f07099fdb5648037d8685 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 19:16:11 -0700 Subject: [PATCH 11/20] [None][test] Simplify NVFP4 premerge case table and shared-LLM factory - Collapse the 13 hand-spelled _NVFP4_4GPU_PREMERGE_CASES dicts (which repeated 4 constant keys x13 + one lpc flip) into a compact tuple matrix of the varying axes plus a builder that fills the constants. Same ids and field values, ~130 lines -> ~30. - make_shared_llm's closure now reuses mpi_session_kwargs instead of re-implementing the 'inject _mpi_session if not None' conditional. Signed-off-by: qgai --- tensorrt_llm/llmapi/_grouped_test_utils.py | 4 +- .../defs/accuracy/test_llm_api_pytorch.py | 185 +++--------------- 2 files changed, 31 insertions(+), 158 deletions(-) diff --git a/tensorrt_llm/llmapi/_grouped_test_utils.py b/tensorrt_llm/llmapi/_grouped_test_utils.py index 047f8cb0712f..4df59c9792e5 100644 --- a/tensorrt_llm/llmapi/_grouped_test_utils.py +++ b/tensorrt_llm/llmapi/_grouped_test_utils.py @@ -100,9 +100,7 @@ def make_shared_llm(mpi_session): from tensorrt_llm import LLM def shared_llm(*args, **kwargs): - if mpi_session is not None: - kwargs["_mpi_session"] = mpi_session - return LLM(*args, **kwargs) + return LLM(*args, **kwargs, **mpi_session_kwargs(mpi_session)) shared_llm.mpi_session = mpi_session return shared_llm diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index f3b1a4bbec2e..addbc860ba2f 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1374,164 +1374,39 @@ def test_auto_dtype_vswa_chunked_prefill_reuse(self): class TestDeepSeekV3Lite(LlmapiAccuracyTestHarness): MODEL_NAME = "deepseek-ai/DeepSeek-V3-Lite" MODEL_PATH = f"{llm_models_root()}/DeepSeek-V3-Lite/bf16" - _NVFP4_4GPU_PREMERGE_CASES = ( - dict(id="cutlass_mtp0_tp4_compile_off", - moe_backend="CUTLASS", - mtp_nextn=0, - tp_size=4, - pp_size=1, - ep_size=1, - fp8kv=True, - attention_dp=True, - cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=False, - torch_compile=False), - dict(id="cutlass_mtp0_tp4_compile_on", - moe_backend="CUTLASS", - mtp_nextn=0, - tp_size=4, - pp_size=1, - ep_size=1, - fp8kv=True, - attention_dp=True, - cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=False, - torch_compile=True), - dict(id="cutlass_mtp0_ep4_compile_on", - moe_backend="CUTLASS", - mtp_nextn=0, - tp_size=4, - pp_size=1, - ep_size=4, - fp8kv=True, - attention_dp=True, - cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=False, - torch_compile=True), - dict(id="cutlass_mtp0_tp2pp2_compile_off", - moe_backend="CUTLASS", - mtp_nextn=0, - tp_size=2, - pp_size=2, - ep_size=1, - fp8kv=True, - attention_dp=True, - cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=False, - torch_compile=False), - dict(id="cutlass_mtp0_tp2pp2_compile_on", - moe_backend="CUTLASS", - mtp_nextn=0, - tp_size=2, - pp_size=2, - ep_size=1, - fp8kv=True, - attention_dp=True, - cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=False, - torch_compile=True), - dict(id="cutlass_mtp2_tp4_compile_off", - moe_backend="CUTLASS", - mtp_nextn=2, - tp_size=4, - pp_size=1, - ep_size=1, - fp8kv=True, - attention_dp=True, - cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=False, - torch_compile=False), - dict(id="trtllm_mtp2_ep4_compile_off", - moe_backend="TRTLLM", - mtp_nextn=2, - tp_size=4, - pp_size=1, - ep_size=4, - fp8kv=True, - attention_dp=True, - cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=False, - torch_compile=False), - dict(id="cutlass_mtp2_pp4_compile_off", - moe_backend="CUTLASS", - mtp_nextn=2, - tp_size=1, - pp_size=4, - ep_size=1, - fp8kv=True, - attention_dp=True, - cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=False, - torch_compile=False), - dict(id="cutlass_mtp0_pp4_compile_off", - moe_backend="CUTLASS", - mtp_nextn=0, - tp_size=1, - pp_size=4, - ep_size=1, - fp8kv=True, - attention_dp=True, - cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=False, - torch_compile=False), - dict(id="cutlass_mtp0_pp4_compile_on", - moe_backend="CUTLASS", - mtp_nextn=0, - tp_size=1, - pp_size=4, - ep_size=1, - fp8kv=True, - attention_dp=True, - cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=False, - torch_compile=True), - dict(id="trtllm_mtp0_tp4_compile_off", - moe_backend="TRTLLM", - mtp_nextn=0, - tp_size=4, - pp_size=1, - ep_size=1, - fp8kv=True, - attention_dp=True, - cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=False, - torch_compile=False), - dict(id="trtllm_mtp0_ep4_compile_off", - moe_backend="TRTLLM", - mtp_nextn=0, - tp_size=4, - pp_size=1, - ep_size=4, - fp8kv=True, - attention_dp=True, - cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=False, - torch_compile=False), - dict(id="cutlass_mtp0_tp4_lpc_compile_off", - moe_backend="CUTLASS", - mtp_nextn=0, - tp_size=4, - pp_size=1, - ep_size=1, + # (id, moe_backend, mtp_nextn, tp_size, pp_size, ep_size, torch_compile, + # low_precision_combine). fp8kv / attention_dp / cuda_graph / + # overlap_scheduler are True for every case (see _NVFP4_4GPU_PREMERGE_CASES). + _NVFP4_4GPU_PREMERGE_MATRIX = ( + ("cutlass_mtp0_tp4_compile_off", "CUTLASS", 0, 4, 1, 1, False, False), + ("cutlass_mtp0_tp4_compile_on", "CUTLASS", 0, 4, 1, 1, True, False), + ("cutlass_mtp0_ep4_compile_on", "CUTLASS", 0, 4, 1, 4, True, False), + ("cutlass_mtp0_tp2pp2_compile_off", "CUTLASS", 0, 2, 2, 1, False, + False), + ("cutlass_mtp0_tp2pp2_compile_on", "CUTLASS", 0, 2, 2, 1, True, False), + ("cutlass_mtp2_tp4_compile_off", "CUTLASS", 2, 4, 1, 1, False, False), + ("trtllm_mtp2_ep4_compile_off", "TRTLLM", 2, 4, 1, 4, False, False), + ("cutlass_mtp2_pp4_compile_off", "CUTLASS", 2, 1, 4, 1, False, False), + ("cutlass_mtp0_pp4_compile_off", "CUTLASS", 0, 1, 4, 1, False, False), + ("cutlass_mtp0_pp4_compile_on", "CUTLASS", 0, 1, 4, 1, True, False), + ("trtllm_mtp0_tp4_compile_off", "TRTLLM", 0, 4, 1, 1, False, False), + ("trtllm_mtp0_ep4_compile_off", "TRTLLM", 0, 4, 1, 4, False, False), + ("cutlass_mtp0_tp4_lpc_compile_off", "CUTLASS", 0, 4, 1, 1, False, + True), + ) + _NVFP4_4GPU_PREMERGE_CASES = tuple( + dict(id=row[0], + moe_backend=row[1], + mtp_nextn=row[2], + tp_size=row[3], + pp_size=row[4], + ep_size=row[5], + torch_compile=row[6], + low_precision_combine=row[7], fp8kv=True, attention_dp=True, cuda_graph=True, - overlap_scheduler=True, - low_precision_combine=True, - torch_compile=False), - ) + overlap_scheduler=True) for row in _NVFP4_4GPU_PREMERGE_MATRIX) @pytest.mark.skip_less_device_memory(60000) @parametrize_with_ids("v2_kv_cache", [True, False]) From 3ea78e00604c3c4a897674404e89463d403a6699 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 19:24:31 -0700 Subject: [PATCH 12/20] [None][test] Explicitly invalidate HF weight cache on grouped-test teardown The per-worker HF raw-weight cache previously relied on the shared session's worker processes exiting to be cleared. Make invalidation explicit: on shared_llm_4gpu teardown (which runs before the session shuts down its workers), submit clear_worker_weight_cache() to each worker. This closes the 'setup cache / teardown invalidate' loop instead of leaning on process death. Signed-off-by: qgai --- tensorrt_llm/llmapi/_grouped_test_utils.py | 13 +++++++++++++ .../defs/accuracy/test_llm_api_pytorch.py | 13 +++++++++++-- 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/llmapi/_grouped_test_utils.py b/tensorrt_llm/llmapi/_grouped_test_utils.py index 4df59c9792e5..68cf57f06936 100644 --- a/tensorrt_llm/llmapi/_grouped_test_utils.py +++ b/tensorrt_llm/llmapi/_grouped_test_utils.py @@ -88,6 +88,19 @@ def reset_worker_torch_compile_state() -> None: torch._dynamo.reset() +def clear_worker_weight_cache() -> None: + """Drop the per-worker HF raw-weight cache (runs inside each worker). + + The cache is a process-global keyed by checkpoint file fingerprints, so it + otherwise lives until the worker process exits. Submit this to each worker + via ``mpi_session.submit_sync(...)`` on group teardown to invalidate it + explicitly instead of relying on process death. + """ + from tensorrt_llm._torch.models.checkpoints import HfWeightLoader + + HfWeightLoader.clear_weight_cache() + + def make_shared_llm(mpi_session): """Return an ``LLM`` factory that transparently injects a shared MPI session. diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index addbc860ba2f..ff3c6f6e201b 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -75,7 +75,9 @@ def patched_start_mpi_pool(self): from tensorrt_llm.llmapi._grouped_test_utils import \ - hf_weight_cache_env as _hf_weight_cache_env # noqa: E402 + clear_worker_weight_cache as _clear_worker_weight_cache # noqa: E402 +from tensorrt_llm.llmapi._grouped_test_utils import \ + hf_weight_cache_env as _hf_weight_cache_env from tensorrt_llm.llmapi._grouped_test_utils import \ make_shared_llm as _make_shared_llm from tensorrt_llm.llmapi._grouped_test_utils import \ @@ -108,7 +110,14 @@ def shared_mpi_session_4gpu(hf_weight_cache): @pytest.fixture(scope="module") def shared_llm_4gpu(shared_mpi_session_4gpu): - return _make_shared_llm(shared_mpi_session_4gpu) + try: + yield _make_shared_llm(shared_mpi_session_4gpu) + finally: + # Explicitly invalidate the per-worker HF weight cache while the shared + # session is still alive (this fixture tears down before + # shared_mpi_session_4gpu), instead of relying on worker process exit. + if shared_mpi_session_4gpu is not None: + shared_mpi_session_4gpu.submit_sync(_clear_worker_weight_cache) def _get_default_torch_compile_config(torch_compile): From a90c97c701174a27e3be751586b07ecd6c485178 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 19:43:16 -0700 Subject: [PATCH 13/20] [None][test] Organize multi-GPU tests into gpu2/gpu4 classes with class-scoped sessions Group test_llm_multi_gpu_pytorch.py into TestLlmMultiGpu2gpu (@gpu2) and TestLlmMultiGpu4gpu (@gpu4), sharing a small _MultiGpuLlmTests base whose class-scoped mpi_session/mpi_kwargs/shared_llm fixtures build one MpiPoolSession of self.n_gpus workers. The session is created lazily and torn down at class end, so the 2-GPU and 4-GPU pools no longer coexist (as they did with the previous module-scoped fixtures). Update the two explicit nodeid references for test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 (l0_dgx_h100.yml, waives.txt) to include the class; -m gpu2/gpu4 and -k filters are unaffected. Signed-off-by: qgai --- .../test_lists/test-db/l0_dgx_h100.yml | 2 +- tests/integration/test_lists/waives.txt | 2 +- .../llmapi/test_llm_multi_gpu_pytorch.py | 349 +++++++++--------- 3 files changed, 168 insertions(+), 185 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index f39810a4f07d..bd9b323dfb7a 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -15,7 +15,7 @@ l0_dgx_h100: auto_trigger: others orchestrator: mpi tests: - - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::TestLlmMultiGpu2gpu::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2" -k "not test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1" - unittest/llmapi/test_additional_model_outputs.py -m "gpu2" - unittest/_torch/multi_gpu -m "not post_merge" TIMEOUT (90) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index db0dc31b5223..d0c1f2a6fc74 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -488,7 +488,7 @@ unittest/executor/test_rpc.py::TestRpcCorrectness::test_incremental_task_async S unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741) unittest/executor/test_rpc_worker.py SKIP (https://nvbugs/5605741) unittest/llmapi/test_llm_multi_gpu.py -m "gpu4 and part0" SKIP (https://nvbugs/5348958) -unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 SKIP (https://nvbugs/6109745) +unittest/llmapi/test_llm_multi_gpu_pytorch.py::TestLlmMultiGpu2gpu::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 SKIP (https://nvbugs/6109745) unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[None] SKIP (https://nvbugs/6162504) unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[cuda_graph_config0] SKIP (https://nvbugs/6162504) unittest/llmapi/test_memory_profiling.py::test_profile_kvcache SKIP (https://nvbugs/5580781) diff --git a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py index 60cb3cc29b77..4a194a6e820c 100644 --- a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py +++ b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py @@ -27,110 +27,6 @@ global_kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4) -@pytest.fixture(scope="module") -def shared_mpi_session_2gpu(): - yield from _shared_mpi_session(2) - - -@pytest.fixture(scope="module") -def shared_mpi_session_4gpu(): - yield from _shared_mpi_session(4) - - -@pytest.fixture(scope="module") -def shared_llm_2gpu(shared_mpi_session_2gpu): - # Factory for tests that construct the LLM directly (see RPC tests below). - return _make_shared_llm(shared_mpi_session_2gpu) - - -@pytest.fixture(scope="module") -def mpi_kwargs_2gpu(shared_mpi_session_2gpu): - # Ready-to-spread LLM kwargs that inject the shared 2-GPU session (or {} when - # MPI is disabled). Harness-based tests just spread this; they never touch the - # raw session object. - return _mpi_session_kwargs(shared_mpi_session_2gpu) - - -@pytest.fixture(scope="module") -def mpi_kwargs_4gpu(shared_mpi_session_4gpu): - return _mpi_session_kwargs(shared_mpi_session_4gpu) - - -@pytest.mark.gpu2 -def test_llm_capture_request_error(mpi_kwargs_2gpu): - _test_llm_capture_request_error(pytorch_backend=True, - tp_size=2, - **mpi_kwargs_2gpu) - - -@pytest.mark.gpu4 -def test_tinyllama_logits_processor_tp2pp2(mpi_kwargs_4gpu): - tinyllama_logits_processor_test_harness(backend="pytorch", - tensor_parallel_size=2, - pipeline_parallel_size=2, - **mpi_kwargs_4gpu) - - -@pytest.mark.gpu2 -@pytest.mark.part0 -@pytest.mark.parametrize("tp_size, pp_size", [(1, 2), (2, 1)]) -def test_tinyllama_logits_processor_2gpu(tp_size: int, pp_size: int, - mpi_kwargs_2gpu): - tinyllama_logits_processor_test_harness(backend="pytorch", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - **mpi_kwargs_2gpu) - - -@pytest.mark.gpu2 -def test_llama_7b_lora_tp2(mpi_kwargs_2gpu): - llama_7b_lora_from_dir_test_harness(tensor_parallel_size=2, - kv_cache_config=global_kv_cache_config, - **mpi_kwargs_2gpu) - - -@pytest.mark.gpu4 -@skip_ray # https://nvbugs/5682551 -@test_lora_with_and_without_cuda_graph -def test_llama_7b_multi_lora_tp4(cuda_graph_config, mpi_kwargs_4gpu): - # For LoRA checkpoints without finetuned embedding and lm_head, we can either: - # (1) specify lora_target_modules, or - # (2) provide a lora_dir to infer the lora_target_modules. - lora_config = LoraConfig(lora_target_modules=['attn_q', 'attn_k', 'attn_v'], - max_lora_rank=8, - max_loras=1, - max_cpu_loras=8) - check_llama_7b_multi_lora_from_request_test_harness( - LLM, - lora_config=lora_config, - tensor_parallel_size=4, - kv_cache_config=global_kv_cache_config, - cuda_graph_config=cuda_graph_config, - **mpi_kwargs_4gpu) - - -@skip_ray # https://nvbugs/5727075 -@pytest.mark.gpu2 -@test_lora_with_and_without_cuda_graph -def test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1( - cuda_graph_config) -> None: - check_phi3_lora_fused_modules_output_tp2_identical_to_tp1( - LLM, cuda_graph_config=cuda_graph_config) - - -@skip_ray -@pytest.mark.gpu2 -def test_llm_rpc_tp2(shared_llm_2gpu): - _run_llm_rpc_tp2(shared_llm_2gpu) - - -@skip_ray -@pytest.mark.gpu2 -@pytest.mark.asyncio -async def test_llm_rpc_streaming_tp2(shared_llm_2gpu): - await _run_llm_rpc_streaming_tp2(shared_llm_2gpu) - - def _run_llm_rpc_tp2(make_llm=LLM): with make_llm(model=llama_model_path, kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), @@ -160,91 +56,178 @@ async def _run_llm_rpc_streaming_tp2(make_llm=LLM): print(f"get result: {output}") -@skip_ray -@pytest.mark.gpu2 -@pytest.mark.parametrize( - "prompt_logprobs, logprobs, return_context_logits, return_generation_logits", - [ - (None, 1, False, - False), # generation logprobs only (top-1, PyTorch limit) - ]) -def test_llm_return_logprobs_streaming_tp2(prompt_logprobs, logprobs, - return_context_logits, - return_generation_logits, - mpi_kwargs_2gpu): - llm_return_logprobs_test_harness(prompt_logprobs, - logprobs, - return_context_logits, - return_generation_logits, - streaming=True, - backend="pytorch", - tp_size=2, - **mpi_kwargs_2gpu) - - -@skip_ray -@pytest.mark.gpu2 -@pytest.mark.parametrize( - "return_context_logits, enable_chunked_prefill, enable_iter_req_stats", - [ - (False, False, True), - (False, True, True), - ], -) -def test_llm_get_stats_pp2(return_context_logits, enable_chunked_prefill, - enable_iter_req_stats, mpi_kwargs_2gpu): - llm_get_stats_test_harness( - tp_size=1, - pp_size=2, - return_context_logits=return_context_logits, - pytorch_backend=True, - enable_chunked_prefill=enable_chunked_prefill, - enable_iter_req_stats=enable_iter_req_stats, - **mpi_kwargs_2gpu, - ) +class _MultiGpuLlmTests: + """Base for the multi-GPU LLM API test classes. + Subclasses set ``n_gpus``; the class-scoped fixtures below build one shared + MpiPoolSession (of that many workers) reused across the subclass's tests and + torn down when the subclass finishes -- so the 2-GPU and 4-GPU pools never + coexist. Fixtures are lazy, so tests that opt out (e.g. constructing a bare + ``LLM``) simply don't request them and pay no session cost. + """ + n_gpus: int -@skip_ray -@pytest.mark.gpu4 -@pytest.mark.parametrize( - "return_context_logits, enable_chunked_prefill, enable_iter_req_stats", - [ - (False, False, True), - (False, True, True), - ], -) -def test_llm_get_stats_pp4(return_context_logits, enable_chunked_prefill, - enable_iter_req_stats, mpi_kwargs_4gpu): - llm_get_stats_test_harness( - tp_size=1, - pp_size=4, - return_context_logits=return_context_logits, - pytorch_backend=True, - enable_chunked_prefill=enable_chunked_prefill, - enable_iter_req_stats=enable_iter_req_stats, - **mpi_kwargs_4gpu, - ) + @pytest.fixture(scope="class") + def mpi_session(self): + yield from _shared_mpi_session(self.n_gpus) + @pytest.fixture(scope="class") + def mpi_kwargs(self, mpi_session): + return _mpi_session_kwargs(mpi_session) -@skip_ray -@pytest.mark.gpu2 -def test_llm_get_stats_tp2(mpi_kwargs_2gpu): - llm_get_stats_test_harness(tp_size=2, - pytorch_backend=True, - **mpi_kwargs_2gpu) + @pytest.fixture(scope="class") + def shared_llm(self, mpi_session): + return _make_shared_llm(mpi_session) -@skip_ray @pytest.mark.gpu2 -def test_llm_get_stats_async_tp2(mpi_kwargs_2gpu): - llm_get_stats_async_test_harness(tp_size=2, - pytorch_backend=True, - **mpi_kwargs_2gpu) +class TestLlmMultiGpu2gpu(_MultiGpuLlmTests): + """2-GPU multi-GPU LLM API tests (shared session from _MultiGpuLlmTests).""" + n_gpus = 2 + + def test_llm_capture_request_error(self, mpi_kwargs): + _test_llm_capture_request_error(pytorch_backend=True, + tp_size=2, + **mpi_kwargs) + + @pytest.mark.part0 + @pytest.mark.parametrize("tp_size, pp_size", [(1, 2), (2, 1)]) + def test_tinyllama_logits_processor_2gpu(self, tp_size, pp_size, + mpi_kwargs): + tinyllama_logits_processor_test_harness(backend="pytorch", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + **mpi_kwargs) + + def test_llama_7b_lora_tp2(self, mpi_kwargs): + llama_7b_lora_from_dir_test_harness( + tensor_parallel_size=2, + kv_cache_config=global_kv_cache_config, + **mpi_kwargs) + + @skip_ray # https://nvbugs/5727075 + @test_lora_with_and_without_cuda_graph + def test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1( + self, cuda_graph_config) -> None: + check_phi3_lora_fused_modules_output_tp2_identical_to_tp1( + LLM, cuda_graph_config=cuda_graph_config) + + @skip_ray + def test_llm_rpc_tp2(self, shared_llm): + _run_llm_rpc_tp2(shared_llm) + + @skip_ray + @pytest.mark.asyncio + async def test_llm_rpc_streaming_tp2(self, shared_llm): + await _run_llm_rpc_streaming_tp2(shared_llm) + + @skip_ray + @pytest.mark.parametrize( + "prompt_logprobs, logprobs, return_context_logits, return_generation_logits", + [ + (None, 1, False, + False), # generation logprobs only (top-1, PyTorch limit) + ]) + def test_llm_return_logprobs_streaming_tp2(self, prompt_logprobs, logprobs, + return_context_logits, + return_generation_logits, + mpi_kwargs): + llm_return_logprobs_test_harness(prompt_logprobs, + logprobs, + return_context_logits, + return_generation_logits, + streaming=True, + backend="pytorch", + tp_size=2, + **mpi_kwargs) + + @skip_ray + @pytest.mark.parametrize( + "return_context_logits, enable_chunked_prefill, enable_iter_req_stats", + [ + (False, False, True), + (False, True, True), + ], + ) + def test_llm_get_stats_pp2(self, return_context_logits, + enable_chunked_prefill, enable_iter_req_stats, + mpi_kwargs): + llm_get_stats_test_harness( + tp_size=1, + pp_size=2, + return_context_logits=return_context_logits, + pytorch_backend=True, + enable_chunked_prefill=enable_chunked_prefill, + enable_iter_req_stats=enable_iter_req_stats, + **mpi_kwargs, + ) + + @skip_ray + def test_llm_get_stats_tp2(self, mpi_kwargs): + llm_get_stats_test_harness(tp_size=2, + pytorch_backend=True, + **mpi_kwargs) + + @skip_ray + def test_llm_get_stats_async_tp2(self, mpi_kwargs): + llm_get_stats_async_test_harness(tp_size=2, + pytorch_backend=True, + **mpi_kwargs) + + @skip_ray + def test_llm_get_stats_async_pp2(self, mpi_kwargs): + llm_get_stats_async_test_harness(pp_size=2, + pytorch_backend=True, + **mpi_kwargs) -@skip_ray -@pytest.mark.gpu2 -def test_llm_get_stats_async_pp2(mpi_kwargs_2gpu): - llm_get_stats_async_test_harness(pp_size=2, - pytorch_backend=True, - **mpi_kwargs_2gpu) +@pytest.mark.gpu4 +class TestLlmMultiGpu4gpu(_MultiGpuLlmTests): + """4-GPU multi-GPU LLM API tests (shared session from _MultiGpuLlmTests).""" + n_gpus = 4 + + def test_tinyllama_logits_processor_tp2pp2(self, mpi_kwargs): + tinyllama_logits_processor_test_harness(backend="pytorch", + tensor_parallel_size=2, + pipeline_parallel_size=2, + **mpi_kwargs) + + @skip_ray # https://nvbugs/5682551 + @test_lora_with_and_without_cuda_graph + def test_llama_7b_multi_lora_tp4(self, cuda_graph_config, mpi_kwargs): + # For LoRA checkpoints without finetuned embedding and lm_head, we can + # either: (1) specify lora_target_modules, or (2) provide a lora_dir to + # infer the lora_target_modules. + lora_config = LoraConfig( + lora_target_modules=['attn_q', 'attn_k', 'attn_v'], + max_lora_rank=8, + max_loras=1, + max_cpu_loras=8) + check_llama_7b_multi_lora_from_request_test_harness( + LLM, + lora_config=lora_config, + tensor_parallel_size=4, + kv_cache_config=global_kv_cache_config, + cuda_graph_config=cuda_graph_config, + **mpi_kwargs) + + @skip_ray + @pytest.mark.parametrize( + "return_context_logits, enable_chunked_prefill, enable_iter_req_stats", + [ + (False, False, True), + (False, True, True), + ], + ) + def test_llm_get_stats_pp4(self, return_context_logits, + enable_chunked_prefill, enable_iter_req_stats, + mpi_kwargs): + llm_get_stats_test_harness( + tp_size=1, + pp_size=4, + return_context_logits=return_context_logits, + pytorch_backend=True, + enable_chunked_prefill=enable_chunked_prefill, + enable_iter_req_stats=enable_iter_req_stats, + **mpi_kwargs, + ) From d6be2d6fe718c89e2a2182f60f2a1602afad4bc4 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 20:09:04 -0700 Subject: [PATCH 14/20] [None][test] Add CUTEDSL cases to grouped NVFP4 premerge matrix Upstream's independently-curated pre-merge smoke set (superseded by the grouped test during the rebase) included two CUTEDSL configs. Add them to _NVFP4_4GPU_PREMERGE_MATRIX so the grouped test preserves that coverage; _run_nvfp4_4gpus_case already skips CUTEDSL off SM100/103. Signed-off-by: qgai --- tests/integration/defs/accuracy/test_llm_api_pytorch.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index ff3c6f6e201b..cb66b0d96139 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1402,6 +1402,11 @@ class TestDeepSeekV3Lite(LlmapiAccuracyTestHarness): ("trtllm_mtp0_ep4_compile_off", "TRTLLM", 0, 4, 1, 4, False, False), ("cutlass_mtp0_tp4_lpc_compile_off", "CUTLASS", 0, 4, 1, 1, False, True), + # CUTEDSL coverage (skipped off SM100/103 by _run_nvfp4_4gpus_case); + # preserves the pre-merge CUTEDSL cases upstream carried individually. + ("cutedsl_mtp0_tp2pp2_compile_off", "CUTEDSL", 0, 2, 2, 1, False, False + ), + ("cutedsl_mtp2_tp4_compile_off", "CUTEDSL", 2, 4, 1, 1, False, False), ) _NVFP4_4GPU_PREMERGE_CASES = tuple( dict(id=row[0], From a5eb810407f5dba3f0e09e8ccf4bb3d73bfe1caf Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 20:14:26 -0700 Subject: [PATCH 15/20] [None][test] Drop carried id from NVFP4 premerge case dicts The case dicts embedded id=row[0] only to strip it back out at call time (ids=[case['id']...] + a case_kwargs comprehension). Keep the id in the matrix (row[0]) for the parametrize ids, make each case a pure _run_nvfp4_4gpus_case kwargs dict, and splat it directly (**case). No behavior change. Signed-off-by: qgai --- .../integration/defs/accuracy/test_llm_api_pytorch.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index cb66b0d96139..fc80906679d5 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1408,9 +1408,10 @@ class TestDeepSeekV3Lite(LlmapiAccuracyTestHarness): ), ("cutedsl_mtp2_tp4_compile_off", "CUTEDSL", 2, 4, 1, 1, False, False), ) + # Each case is pure _run_nvfp4_4gpus_case kwargs; the id (row[0]) stays in + # the matrix and is used only for the parametrize ids. _NVFP4_4GPU_PREMERGE_CASES = tuple( - dict(id=row[0], - moe_backend=row[1], + dict(moe_backend=row[1], mtp_nextn=row[2], tp_size=row[3], pp_size=row[4], @@ -2315,14 +2316,13 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, @pytest.mark.parametrize( "case", _NVFP4_4GPU_PREMERGE_CASES, - ids=[case["id"] for case in _NVFP4_4GPU_PREMERGE_CASES], + ids=[row[0] for row in _NVFP4_4GPU_PREMERGE_MATRIX], ) def test_nvfp4_4gpus_premerge_grouped(self, case, shared_llm_4gpu): # shared_llm_4gpu injects the shared 4-GPU MPI session and (transitively) # depends on hf_weight_cache, so the weight-cache env is exported before # the pool spawns (see the shared_mpi_session_4gpu fixture for details). - case_kwargs = {key: value for key, value in case.items() if key != "id"} - self._run_nvfp4_4gpus_case(make_llm=shared_llm_4gpu, **case_kwargs) + self._run_nvfp4_4gpus_case(make_llm=shared_llm_4gpu, **case) def _run_nvfp4_4gpus_case(self, *, From 6a992114585c30bf2090171358eea90c8ab4e7cf Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 20:16:30 -0700 Subject: [PATCH 16/20] [None][test] Mark HfWeightLoader weight-cache clear as private Rename HfWeightLoader.clear_weight_cache -> _clear_weight_cache: its only callers are tests (via _grouped_test_utils.clear_worker_weight_cache and the weight-loader unit tests), and it matches the sibling private cache helpers (_get_cached_weights, _cache_loaded_weights). Signals it is test-infra, not a supported production API. Signed-off-by: qgai --- .../_torch/models/checkpoints/hf/weight_loader.py | 2 +- tensorrt_llm/llmapi/_grouped_test_utils.py | 2 +- .../_torch/models/checkpoints/hf/test_weight_loader.py | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index 2eb62e0837bb..24a4bb6882cd 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -81,7 +81,7 @@ def _weight_files_cache_key(weight_files: List[str], return (tuple(file_fingerprint), use_consolidated) @classmethod - def clear_weight_cache(cls) -> None: + def _clear_weight_cache(cls) -> None: with _WEIGHT_CACHE_LOCK: _WEIGHT_CACHE.clear() diff --git a/tensorrt_llm/llmapi/_grouped_test_utils.py b/tensorrt_llm/llmapi/_grouped_test_utils.py index 68cf57f06936..45941a1cb281 100644 --- a/tensorrt_llm/llmapi/_grouped_test_utils.py +++ b/tensorrt_llm/llmapi/_grouped_test_utils.py @@ -98,7 +98,7 @@ def clear_worker_weight_cache() -> None: """ from tensorrt_llm._torch.models.checkpoints import HfWeightLoader - HfWeightLoader.clear_weight_cache() + HfWeightLoader._clear_weight_cache() def make_shared_llm(mpi_session): diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py index 728fa3c9f19f..3458e136baaa 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py @@ -102,7 +102,7 @@ def test_load_weights_ignores_consolidated_ckpt_when_sharded_ckpt_exists( def test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper(tmp_path, monkeypatch): monkeypatch.setenv("TRTLLM_HF_WEIGHT_CACHE", "1") - HfWeightLoader.clear_weight_cache() + HfWeightLoader._clear_weight_cache() checkpoint_dir = tmp_path / "foo" checkpoint_dir.mkdir() @@ -130,12 +130,12 @@ def test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper(tmp_path, load_weights_in_parallel.assert_called_once() assert second["foo.weight"] is raw_weight finally: - HfWeightLoader.clear_weight_cache() + HfWeightLoader._clear_weight_cache() def test_weight_cache_disabled_by_default(tmp_path, monkeypatch): monkeypatch.delenv("TRTLLM_HF_WEIGHT_CACHE", raising=False) - HfWeightLoader.clear_weight_cache() + HfWeightLoader._clear_weight_cache() checkpoint_dir = tmp_path / "foo" checkpoint_dir.mkdir() @@ -160,4 +160,4 @@ def test_weight_cache_disabled_by_default(tmp_path, monkeypatch): assert load_weights_in_parallel.call_count == 2 finally: - HfWeightLoader.clear_weight_cache() + HfWeightLoader._clear_weight_cache() From 202f2c99840a88cd9da6f3fb65446c472cbf0917 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 20:37:11 -0700 Subject: [PATCH 17/20] [None][test] Group bf16 4-GPU pre-merge cases into the shared 4-GPU session Extract _run_bfloat16_4gpus_case from test_bfloat16_4gpus and add test_bfloat16_4gpus_premerge_grouped, which reuses the SAME module-scoped shared_llm_4gpu MPI session as the nvfp4 group (both are 4-worker DeepSeek-V3-Lite, so world_size matches). This saves one more session spawn in the gb200 pre-merge stage. The shared runner resets per-worker torch.compile/Dynamo state between cases (same FailOnRecompileLimitHit guard as the nvfp4 group). The HF weight cache is keyed by checkpoint file fingerprint, so the bf16 model gets its own entry; with max_entries=1 the LRU evicts the other model on the first cross-model miss, bounding CPU cache to one model's raw tensors. Replace the three individual bf16 pre-merge entries in l0_gb200_multi_gpus.yml with the grouped test; the post-merge individual test_bfloat16_4gpus[...] entries are unaffected (method still exists). Signed-off-by: qgai --- .../defs/accuracy/test_llm_api_pytorch.py | 84 ++++++++++++++++--- .../test-db/l0_gb200_multi_gpus.yml | 4 +- 2 files changed, 74 insertions(+), 14 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index fc80906679d5..934fe58b0550 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1423,6 +1423,25 @@ class TestDeepSeekV3Lite(LlmapiAccuracyTestHarness): cuda_graph=True, overlap_scheduler=True) for row in _NVFP4_4GPU_PREMERGE_MATRIX) + # bf16 pre-merge cases share the SAME 4-GPU MPI session as the nvfp4 group + # (both are 4-worker DeepSeek-V3-Lite, so world_size matches). mtp_nextn / + # attention_dp / cuda_graph / overlap_scheduler are True for every case. + # (id, tp_size, pp_size, ep_size, torch_compile). + _BF16_4GPU_PREMERGE_MATRIX = ( + ("tp4_compile_on", 4, 1, 1, True), + ("ep4_compile_off", 4, 1, 4, False), + ("tp2pp2_compile_off", 2, 2, 1, False), + ) + _BF16_4GPU_PREMERGE_CASES = tuple( + dict(tp_size=row[1], + pp_size=row[2], + ep_size=row[3], + torch_compile=row[4], + mtp_nextn=2, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True) for row in _BF16_4GPU_PREMERGE_MATRIX) + @pytest.mark.skip_less_device_memory(60000) @parametrize_with_ids("v2_kv_cache", [True, False]) # Chunked Prefill for MLA can only be enabled on SM100 @@ -1600,6 +1619,26 @@ def test_bfloat16_4gpus_kv_cache_aware_routing(self, mtp_nextn): def test_bfloat16_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, attention_dp, cuda_graph, overlap_scheduler, torch_compile): + self._run_bfloat16_4gpus_case(tp_size=tp_size, + pp_size=pp_size, + ep_size=ep_size, + mtp_nextn=mtp_nextn, + attention_dp=attention_dp, + cuda_graph=cuda_graph, + overlap_scheduler=overlap_scheduler, + torch_compile=torch_compile) + + def _run_bfloat16_4gpus_case(self, + *, + tp_size, + pp_size, + ep_size, + mtp_nextn, + attention_dp, + cuda_graph, + overlap_scheduler, + torch_compile, + make_llm=LLM): if pp_size > 1 and mtp_nextn > 0: num_hidden_layers = 30 pp_partition = [num_hidden_layers // pp_size + 1] * pp_size @@ -1618,17 +1657,25 @@ def test_bfloat16_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, mtp_config = None if mtp_nextn > 0: mtp_config = MTPDecodingConfig(max_draft_len=mtp_nextn) - with LLM(self.MODEL_PATH, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - pp_partition=pp_partition, - moe_expert_parallel_size=ep_size, - kv_cache_config=kv_cache_config, - **pytorch_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config) as llm: - task = GSM8K(self.MODEL_NAME) - task.evaluate(llm) + try: + with make_llm(self.MODEL_PATH, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + pp_partition=pp_partition, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + **pytorch_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config) as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + finally: + # On a reused MPI session, reset each worker's torch.compile/Dynamo + # state so recompile counts don't accumulate across cases and trip + # the fullgraph recompile limit (FailOnRecompileLimitHit). + mpi_session = getattr(make_llm, "mpi_session", None) + if mpi_session is not None: + mpi_session.submit_sync(_reset_worker_torch_compile_state) @pytest.mark.skip_less_device(4) @parametrize_with_ids("mtp_nextn", @@ -2324,6 +2371,21 @@ def test_nvfp4_4gpus_premerge_grouped(self, case, shared_llm_4gpu): # the pool spawns (see the shared_mpi_session_4gpu fixture for details). self._run_nvfp4_4gpus_case(make_llm=shared_llm_4gpu, **case) + @pytest.mark.skip_less_device(4) + @skip_pre_hopper + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize( + "case", + _BF16_4GPU_PREMERGE_CASES, + ids=[row[0] for row in _BF16_4GPU_PREMERGE_MATRIX], + ) + def test_bfloat16_4gpus_premerge_grouped(self, case, shared_llm_4gpu): + # Reuses the same module-scoped shared_llm_4gpu session as the nvfp4 + # group. The HF weight cache is keyed by checkpoint file fingerprint, so + # the bf16 model gets its own entry; with max_entries=1 the LRU evicts + # the other model on the first miss, bounding CPU cache to one model. + self._run_bfloat16_4gpus_case(make_llm=shared_llm_4gpu, **case) + def _run_nvfp4_4gpus_case(self, *, fp8kv, diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index 1082dd24bf41..4e985b91e43b 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -20,9 +20,7 @@ l0_gb200_multi_gpus: - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] - - 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] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_premerge_grouped - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_kv_cache_aware_routing[mtp_nextn=0] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_kv_cache_aware_routing[mtp_nextn=2] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped From 384a8e0495863d631b1a688669ed17b9968b3c96 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 20:41:25 -0700 Subject: [PATCH 18/20] [None][fix] Evict HF weight cache before load on miss to avoid 2x CPU peak On a cache miss, evict LRU entries down to the cap BEFORE loading the new weights, so the old (cached) and new (loading) raw tensors are never held in CPU at the same time. Previously the eviction happened only at store time (after the load), so switching models on a reused MPI-session worker transiently held both models' weights (~2x CPU peak). CPU now stays bounded by max_entries models even during a cross-model load. Add a unit test asserting the cache is already empty when the new load runs (max_entries=1), proving eviction precedes allocation. Signed-off-by: qgai --- .../models/checkpoints/hf/weight_loader.py | 23 ++++++++++ .../checkpoints/hf/test_weight_loader.py | 45 +++++++++++++++++++ 2 files changed, 68 insertions(+) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index 24a4bb6882cd..cba87b492a3f 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -85,6 +85,23 @@ def _clear_weight_cache(cls) -> None: with _WEIGHT_CACHE_LOCK: _WEIGHT_CACHE.clear() + @staticmethod + def _evict_to_make_room() -> None: + """Evict LRU entries so a newly loaded entry stays within the cap. + + Called on a cache miss BEFORE loading the new weights. Evicting first + frees the old raw tensors before the new load allocates, so CPU never + holds more than ``max_entries`` models even momentarily. Without this, + switching models on a reused worker would transiently hold both the + old (still-cached) and the new (loading) weights (a ~2x CPU peak). + """ + max_entries = HfWeightLoader._weight_cache_max_entries() + if max_entries <= 0: + return + with _WEIGHT_CACHE_LOCK: + while len(_WEIGHT_CACHE) >= max_entries: + _WEIGHT_CACHE.popitem(last=False) + @staticmethod def _wrap_cached_weights(weights: dict[str, Any]) -> ConsumableWeightsDict: # Return a fresh dict wrapper because model loaders call mark_consumed(). @@ -171,6 +188,9 @@ def load_weights(self, # another rank is about to enter. local_mpi_barrier() return cached_weights + # Cache miss: evict now so the upcoming load doesn't transiently + # hold the old (cached) and new (loading) weights together. + self._evict_to_make_room() # Prefetch the weight files to CPU memory if the size is less than 90% of the available memory. # This is a heuristic to avoid prefetching files that are too large and causing file cache thrashing. @@ -215,6 +235,9 @@ def load_weights(self, f"Reusing cached HF checkpoint weights from {checkpoint_dir}." ) return cached_weights + # Cache miss: evict now so the upcoming load doesn't transiently + # hold the old (cached) and new (loading) weights together. + self._evict_to_make_room() weights = self._load_weights_in_parallel( weight_files, self._load_bin_or_path_file, diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py index 3458e136baaa..8ec25c3a4b55 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py @@ -133,6 +133,51 @@ def test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper(tmp_path, HfWeightLoader._clear_weight_cache() +def test_weight_cache_evicts_before_load_on_miss(tmp_path, monkeypatch): + # On a cross-model miss with a full cache (max_entries=1), the old entry + # must be evicted BEFORE the new weights load, so CPU never holds both the + # old (cached) and new (loading) weights at once (no transient 2x peak). + from tensorrt_llm._torch.models.checkpoints.hf import weight_loader as wl + + monkeypatch.setenv("TRTLLM_HF_WEIGHT_CACHE", "1") + monkeypatch.setenv("TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES", "1") + HfWeightLoader._clear_weight_cache() + + dir_a = tmp_path / "a" + dir_a.mkdir() + (dir_a / "model.safetensors").touch() + dir_b = tmp_path / "b" + dir_b.mkdir() + (dir_b / "model.safetensors").touch() + + loader = HfWeightLoader() + try: + with mock.patch.object(loader, "prefetch_files"): + with mock.patch.object( + loader, + "_load_weights_in_parallel", + return_value=ConsumableWeightsDict({"foo.weight": object()}), + ): + loader.load_weights(str(dir_a), mapping=Mapping()) + assert len(wl._WEIGHT_CACHE) == 1 + + def assert_room_freed_before_load(*args, **kwargs): + # The old (A) entry must already be gone by the time B loads. + assert len(wl._WEIGHT_CACHE) == 0 + return ConsumableWeightsDict({"foo.weight": object()}) + + with mock.patch.object( + loader, + "_load_weights_in_parallel", + side_effect=assert_room_freed_before_load, + ): + loader.load_weights(str(dir_b), mapping=Mapping()) + + assert len(wl._WEIGHT_CACHE) == 1 + finally: + HfWeightLoader._clear_weight_cache() + + def test_weight_cache_disabled_by_default(tmp_path, monkeypatch): monkeypatch.delenv("TRTLLM_HF_WEIGHT_CACHE", raising=False) HfWeightLoader._clear_weight_cache() From 19321b26b6eb921baea6a9cbdf9d3fdac9a12e00 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 1 Jul 2026 20:48:58 -0700 Subject: [PATCH 19/20] [None][test] Dedup shared-session compile-reset teardown; trim docstring /simplify cleanup on the bf16-grouping + weight-cache-eviction changes: - Hoist the identical per-case torch.compile-reset finally block (duplicated verbatim in _run_nvfp4_4gpus_case and _run_bfloat16_4gpus_case) into reset_shared_session_torch_compile_state(make_llm) in _grouped_test_utils, so both call it in one line and a future grouped test can't drift. - Trim the over-long _evict_to_make_room docstring (10 lines -> 2). No behavior change. Signed-off-by: qgai --- .../_torch/models/checkpoints/hf/weight_loader.py | 10 ++-------- tensorrt_llm/llmapi/_grouped_test_utils.py | 13 +++++++++++++ .../defs/accuracy/test_llm_api_pytorch.py | 11 ++++------- 3 files changed, 19 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index cba87b492a3f..b135fa568903 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -87,14 +87,8 @@ def _clear_weight_cache(cls) -> None: @staticmethod def _evict_to_make_room() -> None: - """Evict LRU entries so a newly loaded entry stays within the cap. - - Called on a cache miss BEFORE loading the new weights. Evicting first - frees the old raw tensors before the new load allocates, so CPU never - holds more than ``max_entries`` models even momentarily. Without this, - switching models on a reused worker would transiently hold both the - old (still-cached) and the new (loading) weights (a ~2x CPU peak). - """ + """Evict LRU entries on a miss BEFORE the new load, so CPU never holds + the old (cached) and new (loading) weights at once (a ~2x peak).""" max_entries = HfWeightLoader._weight_cache_max_entries() if max_entries <= 0: return diff --git a/tensorrt_llm/llmapi/_grouped_test_utils.py b/tensorrt_llm/llmapi/_grouped_test_utils.py index 45941a1cb281..3bc12a22def8 100644 --- a/tensorrt_llm/llmapi/_grouped_test_utils.py +++ b/tensorrt_llm/llmapi/_grouped_test_utils.py @@ -88,6 +88,19 @@ def reset_worker_torch_compile_state() -> None: torch._dynamo.reset() +def reset_shared_session_torch_compile_state(make_llm) -> None: + """Reset per-worker torch.compile state on the shared session behind a + ``make_shared_llm`` factory (no-op if there is no shared session). + + Call from a grouped test's per-case teardown so recompile counts don't + accumulate across cases on reused workers (see + ``reset_worker_torch_compile_state`` for the failure it prevents). + """ + mpi_session = getattr(make_llm, "mpi_session", None) + if mpi_session is not None: + mpi_session.submit_sync(reset_worker_torch_compile_state) + + def clear_worker_weight_cache() -> None: """Drop the per-worker HF raw-weight cache (runs inside each worker). diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 934fe58b0550..360ab3b56c5a 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -81,7 +81,8 @@ def patched_start_mpi_pool(self): from tensorrt_llm.llmapi._grouped_test_utils import \ make_shared_llm as _make_shared_llm from tensorrt_llm.llmapi._grouped_test_utils import \ - reset_worker_torch_compile_state as _reset_worker_torch_compile_state + reset_shared_session_torch_compile_state as \ + _reset_shared_session_torch_compile_state from tensorrt_llm.llmapi._grouped_test_utils import \ shared_mpi_session as _shared_mpi_session @@ -1673,9 +1674,7 @@ def _run_bfloat16_4gpus_case(self, # On a reused MPI session, reset each worker's torch.compile/Dynamo # state so recompile counts don't accumulate across cases and trip # the fullgraph recompile limit (FailOnRecompileLimitHit). - mpi_session = getattr(make_llm, "mpi_session", None) - if mpi_session is not None: - mpi_session.submit_sync(_reset_worker_torch_compile_state) + _reset_shared_session_torch_compile_state(make_llm) @pytest.mark.skip_less_device(4) @parametrize_with_ids("mtp_nextn", @@ -2445,9 +2444,7 @@ def _run_nvfp4_4gpus_case(self, # On a reused MPI session, reset each worker's torch.compile/Dynamo # state so recompile counts don't accumulate across cases and trip # the fullgraph recompile limit (FailOnRecompileLimitHit). - mpi_session = getattr(make_llm, "mpi_session", None) - if mpi_session is not None: - mpi_session.submit_sync(_reset_worker_torch_compile_state) + _reset_shared_session_torch_compile_state(make_llm) @parametrize_with_ids( "fp8kv,attention_dp,cuda_graph,overlap_scheduler", From 5d50aa003cd9f7be8494d2cd8ea817790b54abdf Mon Sep 17 00:00:00 2001 From: qgai Date: Thu, 2 Jul 2026 00:49:21 -0700 Subject: [PATCH 20/20] [None][test] Use fully-parametrized ids for grouped test-db entries Test-list matching is exact-nodeid (test_list_parser.modify_by_test_list), so an unbracketed yml entry for a parametrized test never matches any collected test and is silently ignored as an invalid filter. Expand the grouped DeepSeek-V3-Lite nvfp4/bf16 entries to per-case bracketed ids; same-checkpoint cases stay adjacent so the HF weight-cache LRU (max_entries=1) stays hot across cases. Signed-off-by: qgai --- .../test-db/l0_gb200_multi_gpus.yml | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index 4e985b91e43b..bdd0040a1b9d 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -20,10 +20,26 @@ l0_gb200_multi_gpus: - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_premerge_grouped + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_premerge_grouped[tp4_compile_on] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_premerge_grouped[ep4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_premerge_grouped[tp2pp2_compile_off] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_kv_cache_aware_routing[mtp_nextn=0] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_kv_cache_aware_routing[mtp_nextn=2] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_tp4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_tp4_compile_on] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_ep4_compile_on] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_tp2pp2_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_tp2pp2_compile_on] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp2_tp4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[trtllm_mtp2_ep4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp2_pp4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_pp4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_pp4_compile_on] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[trtllm_mtp0_tp4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[trtllm_mtp0_ep4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_tp4_lpc_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutedsl_mtp0_tp2pp2_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutedsl_mtp2_tp4_compile_off] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=True-moe_backend=WIDEEP] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=True-moe_backend=TRTLLM] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_online_eplb[moe_backend=TRTLLM]