174 lines
6.5 KiB
Python
174 lines
6.5 KiB
Python
"""
|
|
corex_fa2.py — Flash Attention 2 dispatch for BI-V100 via ixformer
|
|
|
|
Sub168 log reference:
|
|
corex_fa2.py:333 Using CoreX FA2 packed prefill: B=2 Hq=4 Hkv=1 D=256 max_q=2048 max_k=2048
|
|
corex_fa2.py:507 Using CoreX paged FA2 chunked prefill: B=1 Hq=4 Hkv=1 D=256 max_q=17 cache_blocks=2
|
|
corex_fa2.py:225 Using CoreX paged decode: B=1 Hq=4 Hkv=1 D=256 max_k=45455 partition=256
|
|
|
|
Call chain:
|
|
qwen3_5.py → Attention.forward() → corex_fa2.forward()
|
|
→ ixformer.functions.ixinfer_flash_attn_unpad() (packed prefill)
|
|
→ ixformer.functions.vllm_single_query_cached_kv_attention_v2() (paged decode)
|
|
→ ixformer.functions.ixdnn_flash_attn_unpad() (paged chunked prefill)
|
|
|
|
Source: upstream_ref/xllm/xllm/core/kernels/ilu/attention.cpp
|
|
upstream_ref/xllm/xllm/core/layers/ilu/attention.cpp
|
|
"""
|
|
|
|
import logging
|
|
import math
|
|
import torch
|
|
from typing import Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ============================================================================
|
|
# Load ixformer.functions — these ARE in the base image Python binding
|
|
# ============================================================================
|
|
_ixf_F = None
|
|
try:
|
|
import ixformer.functions as _ixf_F
|
|
except ImportError:
|
|
logger.warning("ixformer.functions not available — FA2 will use xformers fallback")
|
|
|
|
|
|
class CoreXFA2:
|
|
"""
|
|
Flash Attention 2 operator for BI-V100.
|
|
|
|
Three modes matching Sub168 log:
|
|
1. Packed prefill (non-paged, full sequence)
|
|
2. Paged chunked prefill (paged KV cache, chunked prefill)
|
|
3. Paged decode (single token decode with KV cache)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
num_q_heads: int,
|
|
num_kv_heads: int,
|
|
head_dim: int,
|
|
scale: Optional[float] = None,
|
|
block_size: int = 16,
|
|
):
|
|
self.num_q_heads = num_q_heads
|
|
self.num_kv_heads = num_kv_heads
|
|
self.head_dim = head_dim
|
|
self.scale = scale or (1.0 / math.sqrt(head_dim))
|
|
self.block_size = block_size
|
|
self._prefill_logged = False
|
|
self._chunked_logged = False
|
|
self._decode_logged = False
|
|
|
|
def forward_packed_prefill(
|
|
self,
|
|
query: torch.Tensor, # (total_q, num_q_heads, head_dim)
|
|
key: torch.Tensor, # (total_k, num_kv_heads, head_dim)
|
|
value: torch.Tensor, # (total_k, num_kv_heads, head_dim)
|
|
cu_seqlens_q: torch.Tensor, # (batch+1,)
|
|
cu_seqlens_k: torch.Tensor, # (batch+1,)
|
|
max_seqlen_q: int,
|
|
max_seqlen_k: int,
|
|
) -> torch.Tensor:
|
|
"""Packed variable-length prefill using ixinfer flash attn."""
|
|
if _ixf_F is None:
|
|
raise RuntimeError("ixformer not available for FA2 prefill")
|
|
|
|
batch_size = cu_seqlens_q.size(0) - 1
|
|
if not self._prefill_logged:
|
|
logger.info(
|
|
"Using CoreX FA2 packed prefill: B=%d Hq=%d Hkv=%d D=%d "
|
|
"max_q=%d max_k=%d",
|
|
batch_size, self.num_q_heads, self.num_kv_heads,
|
|
self.head_dim, max_seqlen_q, max_seqlen_k)
|
|
self._prefill_logged = True
|
|
|
|
out = torch.empty_like(query)
|
|
_ixf_F.ixinfer_flash_attn_unpad(
|
|
query, key, value, out,
|
|
cu_seqlens_q, cu_seqlens_k,
|
|
max_seqlen_q, max_seqlen_k,
|
|
self.scale, True, # is_causal
|
|
)
|
|
return out
|
|
|
|
def forward_paged_decode(
|
|
self,
|
|
query: torch.Tensor, # (batch, 1, num_q_heads, head_dim)
|
|
key_cache: torch.Tensor, # (num_blocks, block_size, num_kv_heads, head_dim)
|
|
value_cache: torch.Tensor, # (num_blocks, block_size, num_kv_heads, head_dim)
|
|
block_tables: torch.Tensor, # (batch, max_blocks_per_seq)
|
|
context_lens: torch.Tensor, # (batch,)
|
|
) -> torch.Tensor:
|
|
"""Single-token paged decode using vllm paged attention v2."""
|
|
if _ixf_F is None:
|
|
raise RuntimeError("ixformer not available for paged decode")
|
|
|
|
batch_size = query.size(0)
|
|
max_context_len = int(context_lens.max().item())
|
|
|
|
if not self._decode_logged:
|
|
partition_size = 256
|
|
logger.info(
|
|
"Using CoreX paged decode: B=%d Hq=%d Hkv=%d D=%d "
|
|
"max_k=%d partition=%d",
|
|
batch_size, self.num_q_heads, self.num_kv_heads,
|
|
self.head_dim, max_context_len, partition_size)
|
|
self._decode_logged = True
|
|
|
|
out = query.new_empty(batch_size, self.num_q_heads, self.head_dim)
|
|
q_flat = query.squeeze(1) # (batch, num_q_heads, head_dim)
|
|
|
|
_ixf_F.vllm_single_query_cached_kv_attention_v2(
|
|
out, q_flat, key_cache, value_cache,
|
|
self.scale, block_tables, context_lens,
|
|
self.block_size, max_context_len,
|
|
)
|
|
return out.unsqueeze(1)
|
|
|
|
def forward_paged_chunked_prefill(
|
|
self,
|
|
query: torch.Tensor, # (total_q, num_q_heads, head_dim)
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
block_tables: torch.Tensor,
|
|
cu_seqlens_q: torch.Tensor,
|
|
max_seqlen_q: int,
|
|
) -> torch.Tensor:
|
|
"""Paged chunked prefill using ixdnn flash attn with block tables."""
|
|
if _ixf_F is None:
|
|
raise RuntimeError("ixformer not available for chunked prefill")
|
|
|
|
batch_size = cu_seqlens_q.size(0) - 1
|
|
num_cache_blocks = block_tables.size(1) if block_tables.dim() > 1 else 0
|
|
|
|
if not self._chunked_logged:
|
|
logger.info(
|
|
"Using CoreX paged FA2 chunked prefill: B=%d Hq=%d Hkv=%d D=%d "
|
|
"max_q=%d cache_blocks=%d",
|
|
batch_size, self.num_q_heads, self.num_kv_heads,
|
|
self.head_dim, max_seqlen_q, num_cache_blocks)
|
|
self._chunked_logged = True
|
|
|
|
out = torch.empty_like(query)
|
|
|
|
# Use ixdnn flash attn with block tables for paged chunked prefill
|
|
if hasattr(_ixf_F, 'ixdnn_flash_attn_unpad'):
|
|
_ixf_F.ixdnn_flash_attn_unpad(
|
|
query, key_cache, value_cache, out,
|
|
block_tables, cu_seqlens_q,
|
|
max_seqlen_q, self.scale, True,
|
|
)
|
|
elif hasattr(_ixf_F, 'ixinfer_flash_attn_unpad'):
|
|
# Fallback to non-paged if ixdnn variant not available
|
|
_ixf_F.ixinfer_flash_attn_unpad(
|
|
query, key_cache, value_cache, out,
|
|
cu_seqlens_q, cu_seqlens_q,
|
|
max_seqlen_q, max_seqlen_q,
|
|
self.scale, True,
|
|
)
|
|
else:
|
|
raise RuntimeError("No flash attn variant available for chunked prefill")
|
|
|
|
return out
|