[TRTLLM-7963][feat] Cold L2 cache when doing autotune benchmarking

tensorrt_llm/_torch/autotuner.py

@@ -8,6 +8,7 @@
 from typing import Any, Callable, Dict, List, Set, Tuple, Union
 
 import torch
+from cuda.bindings import driver
 
 from tensorrt_llm.bindings.internal.runtime import delay_kernel
 from tensorrt_llm.logger import logger
@@ -525,10 +526,27 @@ def _profile_single_kernel(
             to get an average execution time. Stream synchronization and delays
             are used to ensure accurate timing.
         """
+
+        use_cold_l2 = True
+
+        if use_cold_l2:
+            tensor_lists, num_buffers = self._prepare_input_tensors_with_batches(
+                inputs)
+            buffer_idx = 0
+        else:
+            tensor_lists = [inputs]
+            num_buffers = 1
+            buffer_idx = 0
+
         stream = torch.cuda.current_stream()
         # warm up, no timing
+        # always use the last batch for warmup
         for _ in range(self.warmup):
-            runner(inputs, tactic=tactic, **kwargs)
+            # runner(tensor_lists[-1], tactic=tactic, **kwargs)
+            runner(tensor_lists[buffer_idx % num_buffers],
+                   tactic=tactic,
+                   **kwargs)
+            buffer_idx += 1
         stream.synchronize()
 
         # Delay the profiled kernel launch to eliminate affects of host time overhead in profiling.
@@ -540,13 +558,16 @@ def _profile_single_kernel(
 
         start.record(stream=stream)
         for _ in range(self.repeat):
-            runner(inputs, tactic=tactic, **kwargs)
+            runner(tensor_lists[buffer_idx % num_buffers],
+                   tactic=tactic,
+                   **kwargs)
+            buffer_idx += 1
         end.record(stream=stream)
         stream.synchronize()
 
         avg_time = start.elapsed_time(end) / self.repeat
 
-        shapes = self._get_input_sizes(inputs)
+        shapes = self._get_input_sizes(tensor_lists[-1])
         logger.debug(
             f"[Autotuner] Profiled runner={runner}, tactic={tactic}, shapes={shapes}: {avg_time:.6f}ms."
         )
@@ -733,6 +754,31 @@ def _prepare_input_tensors(
             tensors.append(tensor)
         return tensors
 
+    def _prepare_input_tensors_with_batches(
+        self,
+        inputs: List[torch.Tensor],
+    ) -> Tuple[List[List[torch.Tensor]], int]:
+        # TODO: only consider tensor parameter?
+        one_buffer_bytes = sum(
+            input.numel() *
+            input.element_size() if isinstance(input, torch.Tensor) else 0
+            for input in inputs)
+        num_buffers = ceil(self._get_l2_cache_size_in_bytes() /
+                           one_buffer_bytes)
+        num_buffers = min(num_buffers, self.repeat)
+
+        inputs_list = [inputs]
+        # The last batch is for warmup
+        for _ in range(num_buffers - 1):
+            inputs_list.append(
+                list(t.clone() if isinstance(t, torch.Tensor) else t
+                     for t in inputs))
+
+        logger.debug(
+            f"[Autotuner] To cold L2 cache, use {num_buffers} different tensors for profiling"
+        )
+        return inputs_list, num_buffers
+
     def clear_cache(self) -> None:
         """Clear the profiling cache."""
         self.profiling_cache.clear()
@@ -750,3 +796,32 @@ def print_profiling_cache(self):
             runner_id, tactic, profile = value
             logger.debug(
                 f"[Autotuner] {key}: (runner_id={runner_id}, tactic={tactic})")
+
+    def _get_l2_cache_size_in_bytes(self, device_id: int = 0) -> int:
+        device = self._checkCudaErrors(driver.cuDeviceGet(device_id))
+        return self._checkCudaErrors(
+            driver.cuDeviceGetAttribute(
+                driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_L2_CACHE_SIZE,
+                device,
+            ))
+
+    def _checkCudaErrors(self, result) -> None:
+        if result[0].value:
+            raise RuntimeError("CUDA error code={}({})".format(
+                result[0].value, self._cudaGetErrorEnum(result[0])))
+        # CUDA APIs always return the status as the first element of the result tuple
+        if len(result) == 1:
+            return None
+        elif len(result) == 2:
+            return result[1]
+        else:
+            return result[1:]
+
+    def _cudaGetErrorEnum(error) -> str:
+        if isinstance(error, driver.CUresult):
+            err, name = driver.cuGetErrorName(error)
+            return name if err == driver.CUresult.CUDA_SUCCESS else "<unknown>"
+        elif isinstance(error, nvrtc.nvrtcResult):
+            return nvrtc.nvrtcGetErrorString(error)[1]
+        else:
+            raise RuntimeError("Unknown error type: {}".format(error))