add log_once and update deprecation log

Signed-off-by: jiahanc <[email protected]>

Author: jiahanc

flashinfer/fused_moe/core.py

@@ -18,7 +18,6 @@
 from enum import IntEnum
 from types import SimpleNamespace
 from typing import Any, Dict, List, Optional, Tuple, Union
-from typing_extensions import deprecated
 import torch
 
 from ..autotuner import (
@@ -1656,9 +1655,6 @@ def _fake_trtllm_fp4_block_scale_moe(
     )
 
 
-@deprecated(
-    "tile_tokens_dim is deprecated and will be removed in trtllm_fp8_per_tensor_scale_moe after v0.5.0"
-)
 def trtllm_fp8_per_tensor_scale_moe(
     routing_logits: torch.Tensor,
     routing_bias: Optional[torch.Tensor],
@@ -1708,6 +1704,12 @@ def trtllm_fp8_per_tensor_scale_moe(
     Returns:
         torch.Tensor: Output tensor of shape [seq_len, hidden_size]
     """
+    if tile_tokens_dim is not None:
+        logger.warning_once(
+            "tile_tokens_dim in trtllm_fp8_per_tensor_scale_moe is planned for deprecation "
+            "in a future release. Please remove it from your code as tile_tokens_dim will no "
+            "longer be supported after v0.5.0."
+        )
     return get_trtllm_moe_sm100_module().trtllm_fp8_per_tensor_scale_moe(
         routing_logits,
         routing_bias,
@@ -1731,9 +1733,6 @@ def trtllm_fp8_per_tensor_scale_moe(
     )
 
 
-@deprecated(
-    "tile_tokens_dim is deprecated and will be removed in trtllm_fp8_block_scale_moe after v0.5.0"
-)
 def trtllm_fp8_block_scale_moe(
     routing_logits: torch.Tensor,
     routing_bias: Optional[torch.Tensor],
@@ -1782,6 +1781,12 @@ def trtllm_fp8_block_scale_moe(
     Returns:
         torch.Tensor: Output tensor of shape [seq_len, hidden_size]
     """
+    if tile_tokens_dim is not None:
+        logger.warning_once(
+            "tile_tokens_dim in trtllm_fp8_block_scale_moe is planned for deprecation "
+            "in a future release. Please remove it from your code as tile_tokens_dim will no "
+            "longer be supported after v0.5.0."
+        )
     output = torch.empty(
         hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
     )
@@ -1810,9 +1815,6 @@ def trtllm_fp8_block_scale_moe(
     )
 
 
-@deprecated(
-    "tile_tokens_dim is deprecated and will be removed in trtllm_fp4_block_scale_moe after v0.5.0"
-)
 def trtllm_fp4_block_scale_moe(
     routing_logits: torch.Tensor,
     routing_bias: Optional[torch.Tensor],
@@ -1908,7 +1910,12 @@ def trtllm_fp4_block_scale_moe(
         List[torch.Tensor]: List of output tensors. If do_finalize=True, returns the final MoE output.
             Otherwise, returns intermediate results (gemm2_output, expert_weights, expanded_idx_to_permuted_idx) that need further processing.
     """
-
+    if tile_tokens_dim is not None:
+        logger.warning_once(
+            "tile_tokens_dim in trtllm_fp4_block_scale_moe is planned for deprecation "
+            "in a future release. Please remove it from your code as tile_tokens_dim will no "
+            "longer be supported after v0.5.0."
+        )
     return get_trtllm_moe_sm100_module().trtllm_fp4_block_scale_moe(
         routing_logits,
         None,
@@ -1945,9 +1952,6 @@ def trtllm_fp4_block_scale_moe(
     )
 
 
-@deprecated(
-    "tile_tokens_dim is deprecated and will be removed in trtllm_fp4_block_scale_routed_moe after v0.5.0"
-)
 def trtllm_fp4_block_scale_routed_moe(
     topk_ids: torch.Tensor,
     routing_bias: Optional[torch.Tensor],
@@ -2046,7 +2050,7 @@ def trtllm_fp4_block_scale_routed_moe(
             Otherwise, returns intermediate results (gemm2_output, expert_weights, expanded_idx_to_permuted_idx) that need further processing.
     """
     if tile_tokens_dim is not None:
-        logger.info(
+        logger.warning_once(
             "tile_tokens_dim in trtllm_fp4_block_scale_routed_moe is planned for deprecation "
             "in a future release. Please remove it from your code as tile_tokens_dim will no "
             "longer be supported after v0.5.0."

flashinfer/jit/core.py

@@ -1,10 +1,11 @@
 import dataclasses
+import functools
 import logging
 import os
 from contextlib import nullcontext
 from datetime import datetime
 from pathlib import Path
-from typing import Dict, List, Optional, Sequence, Union
+from typing import Dict, List, Optional, Sequence, Union, Hashable
 
 import tvm_ffi
 from filelock import FileLock
@@ -60,6 +61,33 @@ def __init__(self, name):
             )
         )
 
+    def debug_once(self, msg: str, *args: Hashable) -> None:
+        """
+        As [`debug`][logging.Logger.debug], but subsequent calls with
+        the same message are silently dropped.
+        """
+        self._print_once(self.debug, msg, *args)
+
+    def info_once(self, msg: str, *args: Hashable) -> None:
+        """
+        As [`info`][logging.Logger.info], but subsequent calls with
+        the same message are silently dropped.
+        """
+        self._print_once(self.info, msg, *args)
+
+    def warning_once(self, msg: str, *args: Hashable) -> None:
+        """
+        As [`warning`][logging.Logger.warning], but subsequent calls with
+        the same message are silently dropped.
+        """
+        self._print_once(self.warning, msg, *args)
+
+    @functools.lru_cache(maxsize=None)
+    def _print_once(self, log_method, msg: str, *args: Hashable) -> None:
+        """Helper method to log messages only once per unique (msg, args) combination."""
+        # Note: stacklevel=3 to show the caller's location, not this helper method
+        log_method(msg, *args, stacklevel=3)
+
 
 logger = FlashInferJITLogger("flashinfer.jit")