format code

Author: yao-fengchen

lmdeploy/pytorch/backends/dlinfer/ascend/op_backend.py

@@ -227,7 +227,9 @@ def get_max_tokens_across_dp():
             if dist_ctx.dist_config.dp > 1:
                 total_token_current_rank = torch.sum(step_context.q_seqlens).to(step_context.q_seqlens.dtype)
                 world_size = dist_ctx.dist_config.world_size
-                total_token_buffer = torch.zeros(world_size, dtype=step_context.q_seqlens.dtype, device=torch.npu.current_device())
+                total_token_buffer = torch.zeros(world_size,
+                                                 dtype=step_context.q_seqlens.dtype,
+                                                 device=torch.npu.current_device())
                 dist.all_gather_into_tensor(total_token_buffer, total_token_current_rank, dist_ctx.ep_gpu_group)
                 max_tokens_accros_dp = torch.max(total_token_buffer).item()
             else:

lmdeploy/pytorch/backends/dlinfer/moe.py

@@ -1,6 +1,4 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-
-from dataclasses import dataclass
 from typing import Callable, List
 
 import torch
@@ -13,16 +11,16 @@
 
 def get_dist_ctx():
     dist_ctx = get_dist_manager().current_context()
-    
-    return DlinferDistContext(dp_size = dist_ctx.dist_config.dp,
-                              tp_size = dist_ctx.dist_config.tp,
-                              ep_size = dist_ctx.dist_config.ep,
-                              dp_rank = dist_ctx.dp_rank,
-                              tp_rank = dist_ctx.attn_tp_group.rank,
-                              ep_rank = dist_ctx.ep_rank,
-                              max_tokens_accros_dp = 1,
-                              tp_group = dist_ctx.attn_tp_group.gpu_group,
-                              ep_group = dist_ctx.ep_gpu_group)
+
+    return DlinferDistContext(dp_size=dist_ctx.dist_config.dp,
+                              tp_size=dist_ctx.dist_config.tp,
+                              ep_size=dist_ctx.dist_config.ep,
+                              dp_rank=dist_ctx.dp_rank,
+                              tp_rank=dist_ctx.attn_tp_group.rank,
+                              ep_rank=dist_ctx.ep_rank,
+                              max_tokens_accros_dp=1,
+                              tp_group=dist_ctx.attn_tp_group.gpu_group,
+                              ep_group=dist_ctx.ep_gpu_group)
 
 
 class DlinferSoftmaxTopKImpl(SoftmaxTopKImpl):

lmdeploy/pytorch/kernels/dlinfer/__init__.py

@@ -1,5 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from dlinfer.utils.type_annotation import DlinferDistContext
+
 from ..default import multinomial_sampling, per_channel_quant
 from .apply_rotary_pos_emb import apply_rotary_pos_emb
 from .awq_kernels import awq_linear
@@ -12,6 +13,7 @@
 from .rms_norm import rms_norm
 
 __all__ = [
+    'DlinferDistContext',
     'rms_norm',
     'apply_rotary_pos_emb',
     'awq_linear',

lmdeploy/pytorch/kernels/dlinfer/fused_moe.py

@@ -1,7 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import dlinfer.ops as ext_ops
-import torch
 from torch import Tensor
+
 from . import DlinferDistContext
 
 

lmdeploy/pytorch/kernels/dlinfer/moe_gating_topk_softmax.py

@@ -1,6 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import dlinfer.ops as ext_ops
 from torch import Tensor
+
 from . import DlinferDistContext