Refactor imports inside tensorrt_llm._torch. (NVIDIA#3015)

Signed-off-by: Yuxian Qiu <[email protected]>

Author: yuxianq

tensorrt_llm/_torch/attention_backend/star_flashinfer.py

@@ -1,9 +1,17 @@
+from dataclasses import dataclass, field
+from typing import Dict, Optional
+
+import flashinfer
 import numpy as np
 import torch
 
+from tensorrt_llm.functional import AttentionMaskType
+from tensorrt_llm.models.modeling_utils import QuantConfig
+
 from ..distributed import allgather
 from ..modules.linear import ParallelConfig
-from .flashinfer import *
+from .flashinfer import FlashInferAttentionMetadata, PlanParams
+from .interface import AttentionBackend, AttentionMask, PredefinedAttentionMask
 from .vanilla import VanillaAttention
 
 

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -5,15 +5,16 @@
 
 import torch
 
-from tensorrt_llm._torch.attention_backend.interface import (
-    AttentionBackend, AttentionMask, AttentionMetadata, MLAParams,
-    PositionalEmbeddingParams, PredefinedAttentionMask, RopeParams)
-from tensorrt_llm._torch.attention_backend.vanilla import VanillaAttention
 from tensorrt_llm.functional import (AttentionMaskType, RopeEmbeddingUtils,
                                      RotaryScalingType)
 from tensorrt_llm.logger import logger
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
+from .interface import (AttentionBackend, AttentionMask, AttentionMetadata,
+                        MLAParams, PositionalEmbeddingParams,
+                        PredefinedAttentionMask, RopeParams)
+from .vanilla import VanillaAttention
+
 
 # The type of requests in qkv passed to attention
 # Please keep sync with AttentionInputType in cpp/tensorrt_llm/thop/attentionOp.cpp

tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py

@@ -2,8 +2,8 @@
 
 import torch
 
-import tensorrt_llm._torch.auto_deploy.distributed.common as dist
-import tensorrt_llm._torch.auto_deploy.distributed.trtllm as trtllm_dist
+from ..distributed import common as dist
+from ..distributed import trtllm as trtllm_dist
 
 
 @torch.library.custom_op("dist::all_gather", mutates_args=(), device_types="cuda")

tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe.py

@@ -3,7 +3,7 @@
 import torch
 import torch.nn.functional as F
 
-from tensorrt_llm._torch.modules.fused_moe import FusedMoE  # noqa: F401
+from ...modules.fused_moe import FusedMoE  # noqa: F401
 
 
 @torch.library.custom_op("moe::torch_moe", mutates_args=())

tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py

@@ -4,12 +4,8 @@
 
 # use trtllm distributed ops to improve TP performance if possible
 try:
-    from tensorrt_llm._torch.distributed import AllReduce, allgather
-    from tensorrt_llm._torch.modules.linear import (
-        AllReduceFusionOp,
-        AllReduceParams,
-        ParallelConfig,
-    )
+    from ...distributed import AllReduce, allgather
+    from ...modules.linear import AllReduceFusionOp, AllReduceParams, ParallelConfig
 
     def trtllm_allgather(tensor, dim):
         rank, world_size = get_rank_world_size()

tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py

@@ -15,7 +15,7 @@
 )
 
 try:
-    from tensorrt_llm._torch.quantization.utils import float4_sf_dtype
+    from ...quantization.utils import float4_sf_dtype
 except ImportError:
     float4_sf_dtype = None
 

tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py

@@ -8,7 +8,8 @@
                                              register_replacement)
 
 import tensorrt_llm
-from tensorrt_llm._torch.distributed import AllReduceFusionOp
+
+from ...distributed import AllReduceFusionOp
 
 aten = torch.ops.aten
 from tensorrt_llm.mapping import Mapping

tensorrt_llm/_torch/compilation/patterns/ub_allreduce.py

@@ -8,7 +8,8 @@
                                              register_replacement)
 
 import tensorrt_llm
-from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceStrategy
+
+from ...distributed import AllReduceFusionOp, AllReduceStrategy
 
 aten = torch.ops.aten
 from tensorrt_llm.mapping import Mapping

tensorrt_llm/_torch/models/modeling_bert.py

@@ -4,7 +4,6 @@
 from torch import nn
 from transformers import BertConfig
 
-from tensorrt_llm._torch.modules.linear import Linear
 from tensorrt_llm.llmapi.utils import print_colored_debug
 from tensorrt_llm.logger import logger
 
@@ -14,6 +13,7 @@
 from ..modules.attention import Attention
 from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
+from ..modules.linear import Linear
 from .modeling_utils import register_auto_model
 
 

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -9,15 +9,14 @@
 from tqdm import tqdm
 from transformers import PretrainedConfig
 
-from tensorrt_llm._torch.distributed import (AllReduce, AllReduceFusionOp,
-                                             AllReduceParams, DeepseekAllReduce,
-                                             ParallelConfig, allgather,
-                                             reducescatter)
 from tensorrt_llm.functional import PositionEmbeddingType
+from tensorrt_llm.llmapi.utils import enable_llm_debug
 
-from ...llmapi.utils import enable_llm_debug
 from ..attention_backend import AttentionMetadata
 from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams
+from ..distributed import (AllReduce, AllReduceFusionOp, AllReduceParams,
+                           DeepseekAllReduce, ParallelConfig, allgather,
+                           reducescatter)
 from ..model_config import ModelConfig
 from ..models.modeling_utils import MissingLayer, ModelConfig, support_pp
 from ..modules.attention import MLA
@@ -263,7 +262,7 @@ def __init__(self,
                  dtype: Optional[torch.dtype] = None,
                  tune_max_num_tokens: int = 8192,
                  model_config: ModelConfig = ModelConfig()):
-        from tensorrt_llm._torch.distributed import AllReduce
+        from ..distributed import AllReduce
 
         super().__init__()
         config = model_config.pretrained_config