update

Author: Linboyan-trc

fastdeploy/model_executor/layers/attention/triton_ops/__init__.py

@@ -18,9 +18,9 @@
 # Licensed under Apache License 2.0
 """
 
-from .decode_attention import compute_num_kv_splits, decode_attention_fwd
-from .mla_cache_kernel import mla_write_cache_triton
-from .unified_extend_attention import (
+from .decode_attention import compute_num_kv_splits, decode_attention_fwd  # noqa: F401
+from .mla_cache_kernel import mla_write_cache_triton  # noqa: F401
+from .unified_extend_attention import (  # noqa: F401
     build_kv_indices_from_block_tables,
     build_unified_kv_indices,
     extend_attention_fwd_unified,

fastdeploy/worker/gpu_model_runner.py

@@ -1386,7 +1386,9 @@ def _compute_position_ids_and_slot_mapping(self) -> None:
         Results are stored in self.forward_meta.
         """
         # NOTE(zhushengguang): Only support MLAAttentionBackend and DSAAttentionBackend currently.
-        if not isinstance(self.attn_backends[0], (MLAAttentionBackend, DSAAttentionBackend, TritonMLAAttentionBackend)):
+        if not isinstance(
+            self.attn_backends[0], (MLAAttentionBackend, DSAAttentionBackend, TritonMLAAttentionBackend)
+        ):
             return
         current_total_tokens = self.forward_meta.ids_remove_padding.shape[0]
         position_ids = self.share_inputs["position_ids_buffer"][:current_total_tokens]

scripts/.coveragerc

@@ -31,6 +31,8 @@ omit =
     */fastdeploy/benchmarks/lib/endpoint_request_func.py
     */fastdeploy/model_executor/graph_optimization/utils.py
     */fastdeploy/model_executor/layers/sample/ops/top_k_top_p_triton.py
+    */fastdeploy/model_executor/layers/attention/triton_ops/*
+    */fastdeploy/model_executor/layers/attention/triton_mla_attention_backend.py
     */fastdeploy/model_executor/ops/gpu/fastdeploy_ops.py
     */fastdeploy/model_executor/ops/gpu/fastdeploy_ops/__init__.py
     */fastdeploy/model_executor/ops/gpu/deep_gemm/utils.py