Cherry-pick PR #297 to v0.3.0-dev: Add thead vendor attention backend

vllm_fl/dispatch/backends/flaggems/flaggems.py

@@ -11,6 +11,7 @@
 from typing import Optional, Union
 
 import torch
+import os
 
 from vllm_fl.dispatch.backends.base import Backend
 
@@ -158,7 +159,14 @@ def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) ->
 
         if use_sparse:
             raise ValueError("use_sparse=True requires use_mla=True.")
-        # TODO: return "vllm_fl.dispatch.backends.flaggems.impl.attention.AttentionFLBackend"
+
+        use_flaggems_attn = os.environ.get(
+            "VLLM_FL_USE_FLAGGEMS_ATTN", "0"
+        ).lower() in ("1", "true", "yes")
+
+        if use_flaggems_attn:
+            print("Using FlagGems attention backend.")
+            return "vllm_fl.dispatch.backends.flaggems.impl.attention.AttentionFLBackend"
 
         return AttentionBackendEnum.TRITON_ATTN.get_path()
 

vllm_fl/dispatch/backends/flaggems/impl/attention.py

@@ -74,7 +74,7 @@ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
 
     @staticmethod
     def get_name() -> str:
-        return "FL"
+        return "CUSTOM"
 
     @classmethod
     def supports_attn_type(cls, attn_type: str) -> bool:
@@ -457,7 +457,7 @@ def __init__(
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
 
         self.attn_type = attn_type
-        self.vllm_flash_attn_version = 3 # 2 #get_flash_attn_version()
+        self.vllm_flash_attn_version = 2  # FlagGems only supports FA2
         # Cache the batch invariant result for use in forward passes
         self.batch_invariant_enabled = _bi_mode
 
@@ -468,6 +468,34 @@ def __init__(
         ### TODO(lms): support quant to int8/int4 each query input and low precision compute
         self.supports_quant_query_input = False
 
+    def do_kv_cache_update(
+        self,
+        layer,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+    ):
+        """Write key/value into the paged KV cache.
+
+        This is called by vLLM's unified_kv_cache_update custom op
+        *before* forward(), so forward() should NOT repeat the write.
+        """
+        if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
+            return
+
+        key_cache, value_cache = kv_cache.unbind(0)
+        reshape_and_cache_flash(
+            key,
+            value,
+            key_cache,
+            value_cache,
+            slot_mapping,
+            self.kv_cache_dtype,
+            layer._k_scale,
+            layer._v_scale,
+        )
+
     def forward(
         self,
         layer: torch.nn.Module,
@@ -532,35 +560,11 @@ def forward(
                 layer,
             )
 
-        # For decoder and cross-attention, use KV cache as before
+        # For decoder and cross-attention, use KV cache as before.
+        # NOTE: KV cache write is handled by do_kv_cache_update() which is
+        # called separately by vLLM's unified_kv_cache_update custom op.
         key_cache, value_cache = kv_cache.unbind(0)
 
-        # key and value may be None in the case of cross attention. They are
-        # calculated once based on the output from the encoder and then cached
-        # in KV cache.
-        if (
-            self.kv_sharing_target_layer_name is None
-            and key is not None
-            and value is not None
-        ):
-            # Reshape the input keys and values and store them in the cache.
-            # Skip this if sharing KV cache with an earlier attention layer.
-            # NOTE(woosuk): Here, key and value are padded while slot_mapping is
-            # not padded. However, we don't need to do key[:num_actual_tokens]
-            # and value[:num_actual_tokens] because the reshape_and_cache_flash
-            # op uses the slot_mapping's shape to determine the number of
-            # actual tokens.
-            reshape_and_cache_flash(
-                key,
-                value,
-                key_cache,
-                value_cache,
-                attn_metadata.slot_mapping,
-                self.kv_cache_dtype,
-                layer._k_scale,
-                layer._v_scale,
-            )
-
         if not attn_metadata.use_cascade:
             cu_seqlens_q = attn_metadata.query_start_loc
             seqused_k = attn_metadata.seq_lens
@@ -606,8 +610,8 @@ def forward(
                     q_descale=layer._q_scale.expand(descale_shape),
                     k_descale=layer._k_scale.expand(descale_shape),
                     v_descale=layer._v_scale.expand(descale_shape),
-                    num_splits=attn_metadata.max_num_splits,
-                    s_aux=None, ### self.sinks is support in FA3
+                    num_splits=0,  # FlagGems does not support num_splits > 0
+                    s_aux=None,
                 )
                 return output
 

vllm_fl/dispatch/backends/vendor/thead/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2026 BAAI. All rights reserved.
+
+"""
+Thead backend for vllm-plugin-FL dispatch.
+
+This backend provides operator implementations for T-Head PPU accelerators.
+"""
+
+from .thead import TheadBackend
+
+__all__ = ["TheadBackend"]

vllm_fl/dispatch/backends/vendor/thead/impl/__init__.py

@@ -0,0 +1,5 @@
+"""Copyright (c) 2026 BAAI. All rights reserved."""
+
+from .attention import TheadFlashAttentionBackend, TheadFlashAttentionImpl
+
+__all__ = ["TheadFlashAttentionBackend", "TheadFlashAttentionImpl"]