diff --git a/vllm_fl/dispatch/backends/flaggems/flaggems.py b/vllm_fl/dispatch/backends/flaggems/flaggems.py index b77430af4..0cd2aaf89 100644 --- a/vllm_fl/dispatch/backends/flaggems/flaggems.py +++ b/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() diff --git a/vllm_fl/dispatch/backends/flaggems/impl/attention.py b/vllm_fl/dispatch/backends/flaggems/impl/attention.py index 5d082f2c4..e95a5ced3 100644 --- a/vllm_fl/dispatch/backends/flaggems/impl/attention.py +++ b/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 diff --git a/vllm_fl/dispatch/backends/vendor/thead/__init__.py b/vllm_fl/dispatch/backends/vendor/thead/__init__.py new file mode 100644 index 000000000..972cf090e --- /dev/null +++ b/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"] diff --git a/vllm_fl/dispatch/backends/vendor/thead/impl/__init__.py b/vllm_fl/dispatch/backends/vendor/thead/impl/__init__.py new file mode 100644 index 000000000..90ef6108b --- /dev/null +++ b/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"] diff --git a/vllm_fl/dispatch/backends/vendor/thead/impl/attention.py b/vllm_fl/dispatch/backends/vendor/thead/impl/attention.py new file mode 100644 index 000000000..f1182280e --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/thead/impl/attention.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Thead (T-Head / PPU) FlashAttention backend. +# +# This module provides a custom attention backend for PPU accelerators that +# uses the flash_attn_3 wheel (torch.ops.flash_attn_3.fwd) directly. +# +# At module load time we: +# 1. Import flash_attn_3._C to register FA3 custom ops. +# 2. Provide a custom flash_attn_varlen_func that calls the wheel's fwd +# with the correct arg signature (35 args), bridging differences between +# v0.20.2's FA3 branch (which passes cp_* extras and skips attention_chunk) +# and the wheel's expected signature. +# 3. Inject the needed functions into the flash_attn module namespace so that +# the inherited FlashAttentionImpl.forward() can resolve them. +# 4. Provide a pure-PyTorch reshape_and_cache_flash for PPU (no _C.abi3.so). +# 5. Handle PPU-specific requirements: +# - When cu_seqlens_k is None (paged attention), max_seqlen_k must be 1. +# - FA3 kernel uses max_seqlen_k to select tile size (Aone#75639039). + +from __future__ import annotations + +from typing import ClassVar + +import torch + +# --------------------------------------------------------------------------- +# Step 1 — load the flash_attn_3 wheel +# --------------------------------------------------------------------------- +import flash_attn_3._C # noqa: F401 — registers torch.ops.flash_attn_3 + + +# --------------------------------------------------------------------------- +# Step 2 — provide a custom flash_attn_varlen_func for PPU +# --------------------------------------------------------------------------- +# v0.20.2's FA3 branch calls torch.ops._vllm_fa3_C.fwd() with 37 args: +# ... softcap, True(=rotary_interleaved), scheduler_metadata, num_splits, +# None(=pack_gqa), 0(=sm_margin), s_aux, +# cp_world_size, cp_rank, cp_tot_seqused_k <-- extras +# BUT the flash_attn_3 wheel expects 35 args: +# ... window_size_right, attention_chunk, softcap, is_rotary_interleaved, +# scheduler_metadata, num_splits, pack_gqa, sm_margin, s_aux +# +# So we provide our own varlen wrapper that calls the wheel directly. + + +def _thead_flash_attn_varlen_func( + q, + k, + v, + max_seqlen_q, + cu_seqlens_q, + max_seqlen_k, + cu_seqlens_k=None, + seqused_k=None, + q_v=None, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size: list[int] | None = None, + softcap=0.0, + alibi_slopes=None, + deterministic=False, + return_attn_probs=False, + block_table=None, + return_softmax_lse=False, + out=None, + # FA3 Only + scheduler_metadata=None, + q_descale=None, + k_descale=None, + v_descale=None, + num_splits: int = 0, + # Version selector (ignored — we always use FA3) + fa_version: int = 3, + s_aux=None, + cp_world_size=1, + cp_rank=0, + cp_tot_seqused_k=None, +): + """Custom flash_attn_varlen_func for PPU using the flash_attn_3 wheel. + + Accepts the same signature as vLLM's flash_attn_varlen_func (including + the extra cp_* args), but calls torch.ops.flash_attn_3.fwd with the + correct 35-argument signature. + """ + del fa_version, cp_world_size, cp_rank, cp_tot_seqused_k # unused + del dropout_p, deterministic, return_attn_probs # unused in FA3 + + assert alibi_slopes is None, "Alibi is not supported in FA3" + + if softmax_scale is None: + softmax_scale = q.shape[-1] ** (-0.5) + + real_window_size: tuple[int, int] + if window_size is None: + real_window_size = (-1, -1) + else: + assert len(window_size) == 2 + real_window_size = (window_size[0], window_size[1]) + + # PPU Note Aone#75639039: + # PPU FA3 uses max_seqlen_k to choose tile size. + # In paged attention cu_seqlens_k is None, force max_seqlen_k = 1. + if cu_seqlens_k is None: + max_seqlen_k = 1 + + + out, softmax_lse, _, _ = torch.ops.flash_attn_3.fwd( + q, k, v, + None, None, # k_new, v_new + q_v, + out, + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + None, + seqused_k, # seqused_q, seqused_k + max_seqlen_q, + max_seqlen_k, + block_table, + None, # kv_batch_idx + None, # leftpad_k + None, None, None, # rotary_cos, rotary_sin, seqlens_rotary + q_descale, + k_descale, + v_descale, + softmax_scale, + causal, + real_window_size[0], + real_window_size[1], + 0, # attention_chunk + softcap, + True, # is_rotary_interleaved + scheduler_metadata, + num_splits, + None, # pack_gqa + 0, # sm_margin + s_aux, + ) + + return (out, softmax_lse) if return_softmax_lse else out + + +# --------------------------------------------------------------------------- +# Step 2b — inject into flash_attn module namespace +# --------------------------------------------------------------------------- +import vllm.v1.attention.backends.flash_attn as _flash_attn_mod +from vllm import vllm_flash_attn as _vfa + +_flash_attn_mod.flash_attn_varlen_func = _thead_flash_attn_varlen_func +_flash_attn_mod.get_scheduler_metadata = _vfa.get_scheduler_metadata + +# --------------------------------------------------------------------------- +# Step 2c — pure-PyTorch reshape_and_cache_flash for PPU +# --------------------------------------------------------------------------- +# The original is a CUDA custom op from _C.abi3.so which is not available +# on the remote. We provide a pure-PyTorch indexed-copy version. + + +def reshape_and_cache_flash_thead( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor, + v_scale: torch.Tensor, +) -> None: + """GPU-only KV cache write for PPU, compatible with CUDA graph capture. + + The original CUDA custom op (``_C_cache_ops.reshape_and_cache_flash``) + is not available on the remote. This pure-PyTorch equivalent avoids + *any* CPU-GPU synchronisation or data-dependent shape changes so that + it can run inside a CUDA graph capture region. + + Padding tokens (``slot_mapping == -1``) are handled by zeroing their + key/value before writing to a safe slot, rather than skipping them with + a conditional — the latter would require a CPU sync (``.any()``) and + produce a data-dependent tensor shape. + """ + del kv_cache_dtype, k_scale, v_scale # unused in pure-torch path + + num_kv_heads = key.shape[1] + head_size = key.shape[2] + + # Zero out key/value for padding slots (slot_mapping == -1), then map + # -1 to slot 0 so that every token writes somewhere. Writing zeros to + # slot 0 for padding tokens is harmless. + valid_mask_gpu = (slot_mapping >= 0).to(key.dtype).view(-1, 1, 1) + masked_key = key * valid_mask_gpu + masked_value = value * valid_mask_gpu + + safe_slots = slot_mapping.clamp(min=0) # -1 -> 0 + + # Convert flat slot indices to (block, token_within_block) coordinates. + # key_cache shape: [num_blocks, block_size, num_kv_heads, head_size] + block_size = key_cache.shape[1] + block_indices = safe_slots // block_size + token_in_block = safe_slots % block_size + + # Write each kv_head separately — this avoids flattening the entire + # cache into a 2D tensor, which would create a ~2.5 GiB temporary + # copy on non-contiguous cache layouts (e.g. HND stride order). + for h in range(num_kv_heads): + key_cache[block_indices, token_in_block, h, :] = masked_key[:, h, :] + value_cache[block_indices, token_in_block, h, :] = masked_value[:, h, :] + + +_flash_attn_mod.reshape_and_cache_flash = reshape_and_cache_flash_thead + +# --------------------------------------------------------------------------- +# Step 3 — custom backend & impl +# --------------------------------------------------------------------------- + +from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionImpl, + AttentionType, + MultipleOf, +) +from vllm.v1.attention.backends.flash_attn import ( + FlashAttentionBackend, + FlashAttentionImpl, + FlashAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.fa_utils import ( + flash_attn_supports_fp8, + flash_attn_supports_sinks, + get_flash_attn_version, + is_flash_attn_varlen_func_available, +) + + +class TheadFlashAttentionImpl(FlashAttentionImpl): + """FlashAttention implementation for PPU that uses FA3 (flash_attn_3 wheel). + + The only difference from FlashAttentionImpl: + - vllm_flash_attn_version is forced to 3 (FA3) regardless of CC. + """ + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None = None, + attn_type: AttentionType = AttentionType.DECODER, + kv_sharing_target_layer_name: str | None = None, + sinks: torch.Tensor | None = None, + ) -> None: + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + sinks, + ) + # Override FA version to 3 — our custom flash_attn_varlen_func + # handles the wheel call correctly. + self.vllm_flash_attn_version = 3 + + +class TheadFlashAttentionBackend(FlashAttentionBackend): + """FlashAttention backend for PPU that delegates to TheadFlashAttentionImpl.""" + + @staticmethod + def get_name() -> str: + return "CUSTOM" + + @staticmethod + def get_impl_cls() -> type[TheadFlashAttentionImpl]: + return TheadFlashAttentionImpl + + @staticmethod + def get_builder_cls() -> type[FlashAttentionMetadataBuilder]: + return FlashAttentionMetadataBuilder + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + # PPU CC = 8.0 + return capability >= DeviceCapability(8, 0) and capability < DeviceCapability(9, 0) + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: str | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + device_capability: DeviceCapability, + ) -> str | None: + if has_sink: + return "sink not supported on PPU (CC < 9.0)" + if use_mla: + return "MLA not supported in thead flash attention backend" + return None diff --git a/vllm_fl/dispatch/backends/vendor/thead/register_ops.py b/vllm_fl/dispatch/backends/vendor/thead/register_ops.py new file mode 100644 index 000000000..2e78708c2 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/thead/register_ops.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +""" +Thead (PPU) backend operator registrations. + +This module registers VENDOR (thead) implementations for the dispatch system. +""" + +from __future__ import annotations + +import functools + +from vllm_fl.dispatch.types import OpImpl, BackendImplKind, BackendPriority + + +def _bind_is_available(fn, is_available_fn): + """Wrap a function and bind _is_available attribute for OpImpl.""" + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + return fn(*args, **kwargs) + + wrapper._is_available = is_available_fn + return wrapper + + +def register_builtins(registry) -> None: + """ + Register all thead (PPU) VENDOR operator implementations. + + At registration time we also load the flash_attn_3 wheel so that + TheadFlashAttentionBackend can call FA3 ops. + + Args: + registry: Registry to register into + """ + from .thead import TheadBackend + + backend = TheadBackend() + is_avail = backend.is_available + + impls = [ + OpImpl( + op_name="attention_backend", + impl_id="vendor.thead", + kind=BackendImplKind.VENDOR, + fn=_bind_is_available(backend.attention_backend, is_avail), + vendor="thead", + priority=BackendPriority.VENDOR, + ), + ] + + registry.register_many(impls) diff --git a/vllm_fl/dispatch/backends/vendor/thead/thead.py b/vllm_fl/dispatch/backends/vendor/thead/thead.py new file mode 100644 index 000000000..e3e2936cf --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/thead/thead.py @@ -0,0 +1,122 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +""" +Thead backend implementation. + +This backend provides operator implementations for T-Head PPU accelerators. +For attention, it uses the flash_attn_3 wheel (FA3) for better performance +on CC 8.0 devices. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from vllm_fl.dispatch.backends.base import Backend + + +class TheadBackend(Backend): + """ + Thead (PPU) backend for operator implementations. + + This backend uses the PPU FA3 kernel (flash_attn_3 wheel) for attention, + and vLLM native CUDA implementations for other ops (silu_and_mul, rms_norm, + rotary_embedding). + """ + + _available: Optional[bool] = None + + @property + def name(self) -> str: + return "thead" + + @property + def vendor(self) -> Optional[str]: + return "thead" + + def is_available(self) -> bool: + """ + Check if thead (PPU) hardware is available. + + Detection is based on the PPU_SDK environment variable + (same logic as FlagGems DeviceDetector). + """ + if TheadBackend._available is None: + try: + if not torch.cuda.is_available() or torch.cuda.device_count() == 0: + TheadBackend._available = False + return False + + from vllm.platforms import current_platform + + vendor_name = getattr(current_platform, "vendor_name", None) + if vendor_name == "thead": + TheadBackend._available = True + else: + # Fallback: check PPU_SDK env var + import os + TheadBackend._available = "PPU_SDK" in os.environ + except Exception: + TheadBackend._available = False + return TheadBackend._available + + # ==================== Operator Implementations ==================== + + def silu_and_mul(self, obj, x: torch.Tensor) -> torch.Tensor: + """SiLU activation followed by element-wise multiplication.""" + from vllm.model_executor.layers.activations import silu_and_mul + return silu_and_mul(x) + + def rms_norm(self, obj, x: torch.Tensor, residual: Optional[torch.Tensor] = None): + """RMS normalization.""" + if residual is not None: + from vllm.model_executor.layers.layernorm import rms_norm + return rms_norm(x, obj.weight, residual) + return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=obj.weight) + + def rotary_embedding( + self, + obj, + query: torch.Tensor, + key: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + position_ids: torch.Tensor, + rotary_interleaved: bool = False, + inplace: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Apply rotary position embedding.""" + from vllm.model_executor.layers.rotary_embedding import apply_rotary_emb + return apply_rotary_emb( + query, key, cos, sin, + position_ids=position_ids, + rotary_interleaved=rotary_interleaved, + inplace=inplace, + ) + + def attention_backend( + self, use_mla: bool = False, use_sparse: bool = False + ) -> str: + """ + Get the attention backend class path for PPU. + + Returns the TheadFlashAttentionBackend which uses the FA3 wheel. + + Args: + use_mla: Whether to use Multi-head Latent Attention (MLA) + use_sparse: Whether to use Deepseek Sparse Attention (DSA) + + Returns: + Fully qualified class path string + """ + if use_mla or use_sparse: + # Fall back to standard FLASH_ATTN for MLA/sparse + from vllm.v1.attention.backends.registry import AttentionBackendEnum + return AttentionBackendEnum.FLASH_ATTN.get_path() + + return ( + "vllm_fl.dispatch.backends.vendor.thead.impl.attention." + "TheadFlashAttentionBackend" + ) diff --git a/vllm_fl/dispatch/config/thead.yaml b/vllm_fl/dispatch/config/thead.yaml new file mode 100644 index 000000000..ceef51bc3 --- /dev/null +++ b/vllm_fl/dispatch/config/thead.yaml @@ -0,0 +1,49 @@ +# vLLM-FL Dispatch Configuration for CUDA +# Auto-loaded when running on NVIDIA GPU hardware + +# Preferred default backend type: flagos, vendor, reference +prefer: flagos + +# Strict Mode: +# true = Raise an error immediately on failure; do not attempt other backends. +# false = Attempt the next available backend in sequence upon failure (Default). +strict: false + +# Vendor Whitelist (Optional, allows all if not set) +# allow_vendors: +# - cuda + +# Vendor Blacklist (Optional) +# deny_vendors: +# - ascend + +# Per-operator backend execution order (Optional) +# Only the backends listed here will be attempted, in the order specified. +# +# Supported tokens: +# - flagos : Default FlagOS implementation (Triton) +# - reference : PyTorch reference implementation +# - vendor : Any available vendor backend (auto-detected) +# - vendor:cuda : CUDA-specific vendor backend +op_backends: + # attention_backend: prioritize flagos (Triton attention) + attention_backend: + - flagos + - vendor + - reference + rms_norm: + - flagos + - vendor + - reference + silu_and_mul: + - flagos + - vendor + - reference + rotary_embedding: + - flagos + - vendor + - reference + +# FlagOS operator blacklist +flagos_blacklist: +