format code

Author: wanfengcxz

lmdeploy/pytorch/backends/dlinfer/apply_rotary_emb.py

@@ -1,8 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from torch import Tensor
 
-from lmdeploy.pytorch.kernels.dlinfer import (apply_rotary_pos_emb,
-                                              apply_rotary_pos_emb_interleaved)
+from lmdeploy.pytorch.kernels.dlinfer import apply_rotary_pos_emb, apply_rotary_pos_emb_interleaved
 
 from ..apply_rotary_emb import ApplyRotaryEmbBuilder, ApplyRotaryEmbImpl
 

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

@@ -223,7 +223,7 @@ def update_step_context(cls, step_context):
             cu_seqlens_q_cpu, kv_seqlens_cpu = None, None
         else:
             cu_seqlens_q_cpu, kv_seqlens_cpu = _get_cpu_attention_metadata(cu_seqlens_q, kv_seqlens)
-        
+
         if not step_context.is_decoding:
             is_prefill_no_cache = (
                 False if is_sparse_attention else torch.equal(

lmdeploy/pytorch/backends/dlinfer/nsa.py

@@ -90,9 +90,9 @@ def _forward_index(
     ) -> Tensor:
         if meta.kv_start_indices is None:
             raise RuntimeError(
-                "Ascend NSA metadata is missing kv_start_indices")
+                'Ascend NSA metadata is missing kv_start_indices')
         if meta.cu_seqlen_q is None:
-            raise RuntimeError("Ascend NSA metadata is missing cu_seqlen_q")
+            raise RuntimeError('Ascend NSA metadata is missing cu_seqlen_q')
 
         k = k.unsqueeze(-2)
         fill_kv_cache(
@@ -120,8 +120,8 @@ def _forward_index(
         )
         if indices.dim() != 3 or indices.size(1) != 1:
             raise RuntimeError(
-                "Ascend Lightning Indexer returned an unexpected shape: "
-                f"{tuple(indices.shape)}")
+                'Ascend Lightning Indexer returned an unexpected shape: '
+                f'{tuple(indices.shape)}')
         return indices.squeeze(1)
 
     def forward(

lmdeploy/pytorch/configurations/glm_moe_dsa.py

@@ -19,7 +19,6 @@ def condition(cls, hf_config):
     def build(cls, hf_config, model_path: str | None = None, **kwargs):
         """build."""
         is_draft_model = kwargs.get('is_draft_model', False)
-
         quantization_config = getattr(hf_config, 'quantization_config', None)
         is_lmdeploy_patched_fp8 = (quantization_config is not None
                                    and quantization_config.get('quant_method') == 'fp8'

lmdeploy/pytorch/kernels/dlinfer/__init__.py

@@ -1,13 +1,12 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from ..default import multinomial_sampling, per_channel_quant
-from .apply_rotary_pos_emb import (apply_rotary_pos_emb,
-                                   apply_rotary_pos_emb_interleaved)
+from .apply_rotary_pos_emb import apply_rotary_pos_emb, apply_rotary_pos_emb_interleaved
 from .awq_kernels import awq_linear
 from .fill_kv_cache import fill_kv_cache
 from .flash_attention import flash_attention_fwd
 from .fused_moe import DlinferMoECommType, DlinferMoeMetadata, fused_moe, fused_moe_w8a8
-from .linear import linear
 from .lightning_indexer import lightning_indexer
+from .linear import linear
 from .moe_gating_topk_softmax import moe_gating_topk_softmax
 from .pagedattention import paged_attention_fwd
 from .rms_norm import rms_norm

lmdeploy/pytorch/kernels/dlinfer/flash_attention.py

@@ -21,7 +21,7 @@ def flash_attention_fwd(
     logit_softcapping: float = None,
     causal: bool = True,
 ):
-    actual_seq_lengths_cpu = (q_start_loc + q_seqlens).cpu()    
+    actual_seq_lengths_cpu = (q_start_loc + q_seqlens).cpu()
     return ext_ops.prefill_attention(
         query_states,
         key_states,

lmdeploy/pytorch/kernels/dlinfer/sparse_attention.py

@@ -1,6 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import dlinfer.ops as ext_ops
-import torch
 from torch import Tensor
 
 
@@ -19,13 +18,13 @@ def sparse_attention_fwd(
     """Run Ascend sparse MLA attention over split noPE and RoPE caches."""
     if query.size(-1) <= value_head_size:
         raise ValueError(
-            "sparse MLA query must contain both latent and RoPE dimensions")
+            'sparse MLA query must contain both latent and RoPE dimensions')
     if sparse_indices.dim() == 2:
         sparse_indices = sparse_indices.unsqueeze(1)
     if sparse_indices.dim() != 3 or sparse_indices.size(1) != 1:
         raise ValueError(
-            f"sparse indices must have shape [tokens, 1, topk], got "
-            f"{tuple(sparse_indices.shape)}")
+            f'sparse indices must have shape [tokens, 1, topk], got '
+            f'{tuple(sparse_indices.shape)}')
 
     q_nope = query[..., :value_head_size].contiguous()
     q_rope = query[..., value_head_size:].contiguous()

lmdeploy/pytorch/models/deepseek_v2.py

@@ -1196,7 +1196,8 @@ def get_input_embeddings(self):
         return self.model.get_input_embeddings()
 
     def _map_modelslim_param_name(self, name: str, params_dict: dict[str, nn.Parameter]) -> str | None:
-        """Map ModelSlim checkpoint auxiliary tensors to lmdeploy parameters."""
+        """Map ModelSlim checkpoint auxiliary tensors to lmdeploy
+        parameters."""
         quantization_config = self.quantization_config or {}
         if quantization_config.get('quant_method') != 'modelslim':
             return name

lmdeploy/pytorch/nn/linear/w8a8.py

@@ -172,7 +172,8 @@ def __init__(self,
             'input_scale', torch.nn.Parameter(torch.empty((1, ), dtype=self.dtype, device=device),
                                               requires_grad=False))
         self.register_parameter(
-            'input_offset', torch.nn.Parameter(torch.empty((1, ), dtype=torch.int8, device=device), requires_grad=False))
+            'input_offset',
+            torch.nn.Parameter(torch.empty((1, ), dtype=torch.int8, device=device), requires_grad=False))
         self.register_parameter(
             'deq_scale', torch.nn.Parameter(torch.empty((out_features, ), dtype=torch.float32, device=device),
                                             requires_grad=False))