[feat] 电信chat4 29B模型 沐曦适配

vllm_fl/__init__.py

@@ -173,3 +173,9 @@ def register_model():
         )
     except Exception as e:
         logger.error(f"Register DeepseekV4 model error: {str(e)}")
+    # Register TeleChat mHC patches for DeepSeekV3-based models
+    try:
+        from vllm_fl.patches.deepseek_v2_mhc import apply_model_patches as telechat_mhc
+        telechat_mhc()
+    except Exception as e:
+        logger.error("Register TeleChat mHC patch error: %s", str(e))

vllm_fl/dispatch/backends/reference/reference.py

@@ -14,8 +14,65 @@
 
 import torch
 
+from vllm.logger import init_logger
 from vllm_fl.dispatch.backends.base import Backend
 
+logger = init_logger(__name__)
+
+_metax_mla_patched = False
+
+
+def _patch_flash_attn_for_metax():
+    """
+    On MetaX platform, patch vLLM's MLA modules for MACA compatibility:
+
+    1. flash_attn_varlen_func: vLLM imports from vllm.vllm_flash_attn (CUDA C
+       extension) which is unavailable on MetaX. Patch in the MACA-adapted
+       flash_attn package version for prefill.
+
+    2. decode_attention_fwd: vLLM's native Triton decode kernel uses block sizes
+       that exceed MetaX's 64KB shared memory limit. Patch in the MetaX-adapted
+       version with reduced block sizes.
+    """
+    global _metax_mla_patched
+    if _metax_mla_patched:
+        return
+
+    from vllm.platforms import current_platform
+    if current_platform.vendor_name != "metax":
+        _metax_mla_patched = True
+        return
+
+    # --- Patch 1: flash_attn_varlen_func for prefill ---
+    import vllm.model_executor.layers.attention.mla_attention as mla_mod
+
+    if mla_mod.flash_attn_varlen_func is None:
+        try:
+            from flash_attn import flash_attn_varlen_func
+        except ImportError as e:
+            raise RuntimeError(
+                "MetaX platform requires flash_attn package for MLA prefill. "
+                "Please install the MACA-adapted flash_attn."
+            ) from e
+
+        mla_mod.flash_attn_varlen_func = flash_attn_varlen_func
+        mla_mod.is_vllm_fa = False
+        logger.info("Patched flash_attn_varlen_func from flash_attn package "
+                    "for MetaX MLA prefill support")
+
+    # --- Patch 2: decode_attention_fwd for decode (shmem limit) ---
+    import vllm.v1.attention.backends.mla.triton_mla as triton_mla_mod
+
+    # 实际上是使用了沐曦metax 特化版本的 triton decode attention。因为沐曦需要调整shm大小
+    from vllm_fl.dispatch.backends.vendor.metax.impl.attention.ops.triton_decode_attention import (
+        decode_attention_fwd as metax_decode_attention_fwd,
+    )
+    triton_mla_mod.decode_attention_fwd = metax_decode_attention_fwd
+    logger.info("Patched decode_attention_fwd with MetaX version "
+                "(reduced block size for 64KB shmem limit)")
+
+    _metax_mla_patched = True
+
 
 class ReferenceBackend(Backend):
     """
@@ -165,14 +222,16 @@ def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) ->
         Returns:
             Fully qualified class path string (vLLM native backend)
         """
-        # Return vLLM's native flash attention backend as reference
+        # Return vLLM's native attention backend as reference
         from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
         if use_mla:
-            # vLLM native MLA backend
-            if use_sparse:
-                return AttentionBackendEnum.FLASHMLA_SPARSE.get_path()
-            return AttentionBackendEnum.FLASHMLA.get_path()
+            from vllm.platforms import current_platform
+            if current_platform.vendor_name == "metax":
+                _patch_flash_attn_for_metax()
+            logger.info("attention backend reference dispatch: "
+                        "using TritonMLA for MLA attention")
+            return AttentionBackendEnum.TRITON_MLA.get_path()
         return AttentionBackendEnum.FLASH_ATTN.get_path()
 
     def moe_align_block_size(

vllm_fl/dispatch/backends/vendor/metax/impl/attention/ops/triton_decode_attention.py

@@ -684,6 +684,9 @@ def decode_attention_fwd(
     sm_scale,
     page_size=1,
     logit_cap=0.0,
+    k_scale=None,
+    v_scale=None,
+    is_mla=False,
 ):
     assert num_kv_splits == attn_logits.shape[2]
     kv_group_num = q.shape[1] // v_buffer.shape[-2]

vllm_fl/dispatch/backends/vendor/metax/metax.py

@@ -12,10 +12,13 @@
 
 import torch
 
+from vllm.logger import init_logger
 from vllm_fl.dispatch.backends.base import Backend
 
 from vllm.v1.attention.backends.registry import AttentionBackendEnum, register_backend
 
+logger = init_logger(__name__)
+
 
 # Register attention backends for MACA
 def register_attention_backends():
@@ -153,9 +156,16 @@ def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) ->
         register_attention_backends()
 
         if use_mla:
-            if use_sparse:
-                return AttentionBackendEnum.FLASHMLA_SPARSE.get_path()
-            return AttentionBackendEnum.FLASHMLA.get_path()
+            logger.info("attention backend metax-vendor dispatch skipped: "
+                        "MLA not supported, falling back to reference (TritonMLA)")
+            raise NotImplementedError(
+                "MetaX vendor backend does not support MLA yet, "
+                "falling back to reference (TritonMLA)."
+            )
+            # --- original metax MLA logic (commented out) ---
+            # if use_sparse:
+            #     return AttentionBackendEnum.FLASHMLA_SPARSE.get_path()
+            # return AttentionBackendEnum.FLASHMLA.get_path()
 
         # Default to FLASH_ATTN
         return AttentionBackendEnum.FLASH_ATTN.get_path()

vllm_fl/dispatch/config/iluvatar.yaml

@@ -45,19 +45,19 @@ op_backends:
   # All operators: prioritize FlagGems (Triton), fallback to vendor, then reference
   attention_backend:
     - flagos
-    - vendor:iluvatar
+    #- vendor:iluvatar
     - reference
   rms_norm:
     - flagos
-    - vendor:iluvatar
+    #- vendor:iluvatar
     - reference
   silu_and_mul:
     - flagos
-    - vendor:iluvatar
+    #- vendor:iluvatar
     - reference
   rotary_embedding:
     - flagos
-    - vendor:iluvatar
+    #- vendor:iluvatar
     - reference
 
 # FlagOS operator blacklist

vllm_fl/dispatch/config/metax.yaml

@@ -40,9 +40,10 @@ strict: false
 #   - vendor        : Any available vendor backend (auto-detected)
 #   - vendor:metax   : METAX-specific vendor backend
 op_backends:
-  # attention_backend: prioritize vendor (MACA attention)
+  # attention_backend: prioritize vendor (MACA attention), fallback to reference (TritonMLA)
   attention_backend:
-    - vendor:metax
+    # - vendor:metax
+    - reference
   # All other ops: prefer FlagGems (Triton), fallback to reference (PyTorch native)
   silu_and_mul:
     - flagos