[TRTLLM-9615][feat] Support PP in the distributed tuning system

Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>

Author: hyukn

tensorrt_llm/_torch/autotuner.py

@@ -22,6 +22,9 @@
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
 
+# Unique tag to avoid collisions with other comms
+PP_COMM_TAG_AUTOTUNING = 22000
+
 
 class DistributedTuningStrategy(enum.Enum):
     """
@@ -358,7 +361,7 @@ class AutoTunerProfilingCache:
     """
 
     def __init__(self):
-        self.cache = {}
+        self.cache: Dict[Tuple, Tuple] = dict()
 
         # Cache metadata for local storage and validation
         self.lib_version = tensorrt_llm.__version__
@@ -430,7 +433,7 @@ def get_cache_key(
             ),
         )
 
-    def merge_cache_data(self, cache_data: Dict[str, Any]):
+    def merge_cache_data(self, cache_data: Dict[Tuple, Tuple]):
         self.cache.update(cache_data)
 
     def get_specific_custom_op(self, custom_op: str) -> Dict[Tuple, Tuple]:
@@ -615,7 +618,10 @@ def __init__(self, warmup=2, repeat=10, stream_delay_micro_secs=1000):
         self._last_capture: Optional['AutoTuner.TacticsCapture'] = None
 
         # Dsitributed tuning state
+        self._map_op_to_distributed_strategy: Dict[
+            str, DistributedTuningStrategy] = {}
         self._dist: Optional[Distributed] = None
+        self._has_received_cache: bool = False
         self.mapping: Mapping = Mapping()
 
     @classmethod
@@ -797,10 +803,17 @@ def choose_one(
         if self.is_tuning_mode and is_cache_hit:
             return (runners[best_runner_id], best_tactic)
 
+        # Reach here means this PP rank does not have cache hit, so we need to wait for a recv
+        self.cache_pp_recv()
+
         assert len(runners) > 0, "At least one runner is required"
         assert all([isinstance(r, TunableRunner) for r in runners]), \
             "All Given runners must be subclass of TunableRunner"
 
+        # Record the distributed tuning strategy for the custom_op
+        self._map_op_to_distributed_strategy[
+            custom_op] = tuning_config.distributed_tuning_strategy
+
         tuning_start_time = time.perf_counter()
         profiles = self._optimization_profiles(tuning_config, inputs)
 
@@ -1507,3 +1520,30 @@ def _should_current_rank_tune(self,
                 f"[AutoTuner] Unknown distributed tuning strategy: {strategy}, falling back to independent"
             )
             return True
+
+    def cache_pp_recv(self):
+        if self.mapping.has_pp() and not self.mapping.is_first_pp_rank(
+        ) and self._dist is not None and not self._has_received_cache:
+            self._debug_logger(
+                f"[AutoTuner] Receiving cache data from previous pp rank {self.mapping.prev_pp_rank()}"
+            )
+            profiling_cache = self._dist.recv_object(
+                self.mapping.prev_pp_rank(),
+                tag=PP_COMM_TAG_AUTOTUNING,
+            )
+            # Every time a PP rank only receives a cache data once. After that, it should not do any further receiving
+            self._has_received_cache = True
+            self.profiling_cache.merge_cache_data(profiling_cache)
+
+    def cache_pp_send(self):
+        # Send all cache contents to the next PP rank
+        if self.mapping.has_pp(
+        ) and not self.mapping.is_last_pp_rank() and self._dist is not None:
+            self._debug_logger(
+                f"[AutoTuner] Sending cache data to next pp rank {self.mapping.next_pp_rank()}"
+            )
+            self._dist.send_object(
+                self.profiling_cache.cache,
+                self.mapping.next_pp_rank(),
+                tag=PP_COMM_TAG_AUTOTUNING,
+            )

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

@@ -693,6 +693,14 @@ def _(
 
 class NVFP4GemmUnifiedRunner(TunableRunner):
     runner_dict = dict()
+    tuning_config = TuningConfig(
+        dynamic_tensor_specs=(DynamicTensorSpec(
+            0, 0, get_last_power_of_2_num_tokens_buckets,
+            last_positive_power_of_2), ),
+        constraint_specs=(ConstraintSpec(2, 0, fp4_scale_infer_shape), ),
+        # nested tuning should always be independent
+        distributed_tuning_strategy=DistributedTuningStrategy.INDEPENDENT,
+    )
 
     def __init__(self, to_userbuffers: bool, output_dtype: torch.dtype,
                  allowed_backends: List[str]):
@@ -943,7 +951,7 @@ def nvfp4_gemm(
         _, best_tactic = tuner.choose_one(
             "trtllm::nvfp4_gemm::gemm",
             [runner],
-            FP4GemmRunner.
+            NVFP4GemmUnifiedRunner.
             tuning_config,  # All runners use the same tuning_config
             [act_fp4, weight, act_sf, weight_scale, alpha],
         )

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -649,9 +649,14 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager):
                     if self.is_draft_model and isinstance(
                             spec_resource_manager, Eagle3ResourceManager):
                         spec_resource_manager.is_first_draft = True
+
                     self.forward(batch,
                                  new_tensors_device=None,
                                  resource_manager=resource_manager)
+
+                    # Sync the cache after the tuning process
+                    AutoTuner.get().cache_pp_send()
+
                     torch.cuda.synchronize()
 
         logger.info(