Skip to content

GitLab

Commit ea9b8584 authored by laibao's avatar laibao
Browse files

feat(v1 attention): 为 ROCm FlashAttention 接入 unified kv layout，并打通... 

feat(v1 attention): 为 ROCm FlashAttention 接入 unified kv layout，并打通 mm_prefix、qq_bias 与 use_alibi_sqrt 透传
在 ROCm FlashAttention 后端增加 unified KV layout 选择逻辑
接入 unified varlen kernel 调用路径
在 FlashAttention metadata 中补充 mm_prefix_range 与 qq_bias 透传
parent b81573da
Hide whitespace changes
Inline Side-by-side
Showing with 282 additions and 144 deletions
vllm/attention/layer.py
View file @ ea9b8584
...@@ -245,6 +245,7 @@ class Attention(nn.Module, AttentionLayerBase): ...@@ -245,6 +245,7 @@ class Attention(nn.Module, AttentionLayerBase):
use_mla=False, use_mla=False,
has_sink=self.has_sink, has_sink=self.has_sink,
use_mm_prefix=self.use_mm_prefix, use_mm_prefix=self.use_mm_prefix,
use_alibi_sqrt=bool(use_alibi_sqrt),
attn_type=attn_type, attn_type=attn_type,
) )
else: else:
...@@ -1273,4 +1274,4 @@ direct_register_custom_op( ...@@ -1273,4 +1274,4 @@ direct_register_custom_op(
mutates_args=["qkv", "positions"], mutates_args=["qkv", "positions"],
fake_impl=fused_qkv_split_rmsnorm_rope_kv_store_fake, fake_impl=fused_qkv_split_rmsnorm_rope_kv_store_fake,
tags=(torch.Tag.needs_fixed_stride_order,), tags=(torch.Tag.needs_fixed_stride_order,),
) )
\ No newline at end of file
vllm/envs.py
View file @ ea9b8584
...@@ -319,7 +319,6 @@ if TYPE_CHECKING: ...@@ -319,7 +319,6 @@ if TYPE_CHECKING:
VLLM_USE_CUDA_GRAPH_SIZES: bool = False VLLM_USE_CUDA_GRAPH_SIZES: bool = False
VLLM_USE_LIGHTOP_MOE_SUM_MUL_ADD: bool = False VLLM_USE_LIGHTOP_MOE_SUM_MUL_ADD: bool = False
VLLM_USE_LIGHTOP_FUSED_TOPP_TOPK: bool = False VLLM_USE_LIGHTOP_FUSED_TOPP_TOPK: bool = False
VLLM_V1_USE_FA_UNIFIED_ATTN_2D: bool = False
VLLM_ENABLE_RAY_ASYNC_SCHEDULING: bool = False VLLM_ENABLE_RAY_ASYNC_SCHEDULING: bool = False
USE_LIGHTOP_PER_TOKEN_GROUP_QUANT_FP8: bool = False USE_LIGHTOP_PER_TOKEN_GROUP_QUANT_FP8: bool = False
USE_LIGHTOP_TOPK: bool = False USE_LIGHTOP_TOPK: bool = False
...@@ -2003,7 +2002,7 @@ environment_variables: dict[str, Callable[[], Any]] = { ...@@ -2003,7 +2002,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
"USE_LIGHTOP_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX": "USE_LIGHTOP_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX":
lambda: (os.environ.get("USE_LIGHTOP_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX", "False").lower() in lambda: (os.environ.get("USE_LIGHTOP_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX", "False").lower() in
("true", "1")), ("true", "1")),
} }
# --8<-- [end:env-vars-definition] # --8<-- [end:env-vars-definition]
......
vllm/platforms/rocm.py
View file @ ea9b8584
...@@ -262,6 +262,7 @@ class RocmPlatform(Platform): ...@@ -262,6 +262,7 @@ class RocmPlatform(Platform):
from vllm._aiter_ops import rocm_aiter_ops from vllm._aiter_ops import rocm_aiter_ops
block_size = attn_selector_config.block_size block_size = attn_selector_config.block_size
head_size = attn_selector_config.head_size
kv_cache_dtype = attn_selector_config.kv_cache_dtype kv_cache_dtype = attn_selector_config.kv_cache_dtype
if attn_selector_config.use_sparse: if attn_selector_config.use_sparse:
...@@ -304,9 +305,36 @@ class RocmPlatform(Platform): ...@@ -304,9 +305,36 @@ class RocmPlatform(Platform):
f"is not MLA type while requested for MLA backend." f"is not MLA type while requested for MLA backend."
) )
use_unified_flash = (
if envs.VLLM_USE_FLASH_ATTN_PA and block_size == 64: block_size is not None
logger.info_once("Using Flash Attention backend on V1 engine. (only supports block size 64)") and block_size != 64
and block_size % 64 == 0
and head_size == 256
)
if (
envs.VLLM_USE_FLASH_ATTN_PA
and block_size is not None
and block_size != 64
and block_size % 64 == 0
and head_size != 256
):
logger.info_once(
"Skip unified varlen kernel on V1 engine: head size %d is "
"unsupported (requires 256).",
head_size,
)
if envs.VLLM_USE_FLASH_ATTN_PA and (block_size == 64 or use_unified_flash):
if use_unified_flash and block_size != 64:
logger.info_once(
"Using Flash Attention backend with unified varlen kernel on "
"V1 engine. (block size %d, requires block size divisible by 64)",
block_size,
)
else:
logger.info_once(
"Using Flash Attention backend on V1 engine. "
"(only supports block size 64)"
)
return AttentionBackendEnum.FLASH_ATTN.get_path() return AttentionBackendEnum.FLASH_ATTN.get_path()
else: else:
os.environ['VLLM_USE_FLASH_ATTN_PA'] = '0' os.environ['VLLM_USE_FLASH_ATTN_PA'] = '0'
......
vllm/v1/attention/backend.py
View file @ ea9b8584
...@@ -225,6 +225,7 @@ class AttentionBackend(ABC): ...@@ -225,6 +225,7 @@ class AttentionBackend(ABC):
has_sink: bool, has_sink: bool,
use_sparse: bool, use_sparse: bool,
use_mm_prefix: bool, use_mm_prefix: bool,
use_alibi_sqrt: bool,
device_capability: "DeviceCapability", device_capability: "DeviceCapability",
attn_type: str, attn_type: str,
) -> list[str]: ) -> list[str]:
...@@ -241,6 +242,8 @@ class AttentionBackend(ABC): ...@@ -241,6 +242,8 @@ class AttentionBackend(ABC):
invalid_reasons.append( invalid_reasons.append(
"partial multimodal token full attention not supported" "partial multimodal token full attention not supported"
) )
if use_alibi_sqrt and not cls.supports_alibi_sqrt():
invalid_reasons.append("use_alibi_sqrt not supported")
if use_mla != cls.is_mla(): if use_mla != cls.is_mla():
if use_mla: if use_mla:
invalid_reasons.append("MLA not supported") invalid_reasons.append("MLA not supported")
......
vllm/v1/attention/backends/flash_attn.py
View file @ ea9b8584
...@@ -33,6 +33,13 @@ if is_flash_attn_varlen_func_available(): ...@@ -33,6 +33,13 @@ if is_flash_attn_varlen_func_available():
vllm_flash_attn_varlen_func, vllm_flash_attn_varlen_func,
reshape_and_cache_cuda, reshape_and_cache_cuda,
) )
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
triton_reshape_and_cache_flash,
)
try:
from flash_attn import varlen_fwd_unified
except Exception:
varlen_fwd_unified = None
else: else:
from vllm.v1.attention.backends.fa_utils import ( from vllm.v1.attention.backends.fa_utils import (
flash_attn_supports_sinks, flash_attn_supports_sinks,
...@@ -112,6 +119,38 @@ class FlashAttentionBackend(AttentionBackend): ...@@ -112,6 +119,38 @@ class FlashAttentionBackend(AttentionBackend):
@staticmethod @staticmethod
def get_builder_cls() -> type["FlashAttentionMetadataBuilder"]: def get_builder_cls() -> type["FlashAttentionMetadataBuilder"]:
return FlashAttentionMetadataBuilder return FlashAttentionMetadataBuilder
@classmethod
def supports_alibi_sqrt(cls) -> bool:
return True
@classmethod
def supports_mm_prefix(cls) -> bool:
return True
@staticmethod
def _use_rocm_unified_kv_layout(
block_size: int | None = None,
key_cache: torch.Tensor | None = None,
value_cache: torch.Tensor | None = None,
) -> bool:
if not current_platform.is_rocm():
return False
if block_size is None:
if key_cache is not None and value_cache is not None:
if key_cache.ndim != 4 or value_cache.ndim != 4:
return False
if key_cache.shape != value_cache.shape:
return False
block_size = key_cache.shape[1]
else:
try:
block_size = get_current_vllm_config().cache_config.block_size
except Exception:
return False
return block_size is not None and block_size != 64 and block_size % 64 == 0
if current_platform.is_rocm(): if current_platform.is_rocm():
@staticmethod @staticmethod
...@@ -124,6 +163,9 @@ class FlashAttentionBackend(AttentionBackend): ...@@ -124,6 +163,9 @@ class FlashAttentionBackend(AttentionBackend):
) -> tuple[tuple[int, ...], tuple[int, ...]]: ) -> tuple[tuple[int, ...], tuple[int, ...]]:
if block_size % 16 != 0: if block_size % 16 != 0:
raise ValueError("Block size must be a multiple of 16.") raise ValueError("Block size must be a multiple of 16.")
if FlashAttentionBackend._use_rocm_unified_kv_layout(block_size):
unified_shape = (num_blocks, block_size, num_kv_heads, head_size)
return (unified_shape, unified_shape)
return ( return (
(num_blocks, num_kv_heads, block_size, head_size), (num_blocks, num_kv_heads, block_size, head_size),
(num_blocks, num_kv_heads, head_size, block_size), (num_blocks, num_kv_heads, head_size, block_size),
...@@ -136,20 +178,31 @@ class FlashAttentionBackend(AttentionBackend): ...@@ -136,20 +178,31 @@ class FlashAttentionBackend(AttentionBackend):
# `stride_order` indicates the permutation that gets # `stride_order` indicates the permutation that gets
# us from `get_kv_cache_shape` to the actual memory layout we want. # us from `get_kv_cache_shape` to the actual memory layout we want.
cache_layout = get_kv_cache_layout() cache_layout = get_kv_cache_layout()
if cache_layout == "NHD" and include_num_layers_dimension: if FlashAttentionBackend._use_rocm_unified_kv_layout():
# (num_blocks, num_layers, block_size, num_kv_heads, head_size) if cache_layout != "NHD":
return (1, 0, 3, 2, 5), (1, 0, 4, 2, 3) raise RuntimeError(
elif cache_layout == "NHD": "ROCm unified KV layout currently supports NHD only."
)
if include_num_layers_dimension:
# (num_blocks, num_layers, block_size, num_kv_heads, head_size)
return (1, 0, 2, 3, 4), (1, 0, 2, 3, 4)
key_stride_order = (0, 1, 2, 3) key_stride_order = (0, 1, 2, 3)
value_stride_order = (0, 1, 2, 3) value_stride_order = (0, 1, 2, 3)
elif cache_layout == "HND" and include_num_layers_dimension:
# (num_blocks, num_kv_heads, num_layers, block_size, head_size)
return (1, 2, 0, 3, 4), (1, 2, 0, 4, 3)
elif cache_layout == "HND":
key_stride_order = (0, 1, 2, 3)
value_stride_order = (0, 1, 3, 2)
else: else:
raise ValueError(f"Unknown cache layout format {cache_layout}.") if cache_layout == "NHD" and include_num_layers_dimension:
# (num_blocks, num_layers, block_size, num_kv_heads, head_size)
return (1, 0, 3, 2, 5), (1, 0, 4, 2, 3)
elif cache_layout == "NHD":
key_stride_order = (0, 1, 2, 3)
value_stride_order = (0, 1, 2, 3)
elif cache_layout == "HND" and include_num_layers_dimension:
# (num_blocks, num_kv_heads, num_layers, block_size, head_size)
return (1, 2, 0, 3, 4), (1, 2, 0, 4, 3)
elif cache_layout == "HND":
key_stride_order = (0, 1, 2, 3)
value_stride_order = (0, 1, 3, 2)
else:
raise ValueError(f"Unknown cache layout format {cache_layout}.")
return key_stride_order, value_stride_order return key_stride_order, value_stride_order
else: else:
@staticmethod @staticmethod
...@@ -271,8 +324,34 @@ class FlashAttentionMetadata: ...@@ -271,8 +324,34 @@ class FlashAttentionMetadata:
prefix_scheduler_metadata: torch.Tensor | None = None prefix_scheduler_metadata: torch.Tensor | None = None
max_num_splits: int = 0 max_num_splits: int = 0
mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None
qq_bias: torch.Tensor | None = None
causal: bool = True causal: bool = True
@property
def mm_prefix_range_tensor(self) -> torch.Tensor | None:
if self.mm_prefix_range is None:
return None
num_seqs = self.seq_lens.shape[0]
device = self.seq_lens.device
range_lists = [
self.mm_prefix_range.get(i, [(0, 0)]) or [(0, 0)]
for i in range(num_seqs)
]
if all(r == [(0, 0)] for r in range_lists):
return None
range_tensors = [
torch.tensor(r, dtype=torch.int32, device=device).view(-1, 2)
for r in range_lists
]
return torch.nested.nested_tensor(
range_tensors, layout=torch.jagged
).to_padded_tensor(0)
def _get_sliding_window_configs( def _get_sliding_window_configs(
vllm_config: VllmConfig, vllm_config: VllmConfig,
...@@ -597,6 +676,7 @@ class FlashAttentionImpl(AttentionImpl): ...@@ -597,6 +676,7 @@ class FlashAttentionImpl(AttentionImpl):
attn_type: AttentionType = AttentionType.DECODER, attn_type: AttentionType = AttentionType.DECODER,
kv_sharing_target_layer_name: str | None = None, kv_sharing_target_layer_name: str | None = None,
sinks: torch.Tensor | None = None, sinks: torch.Tensor | None = None,
use_alibi_sqrt: bool = False,
) -> None: ) -> None:
self.num_heads = num_heads self.num_heads = num_heads
self.head_size = head_size self.head_size = head_size
...@@ -622,6 +702,7 @@ class FlashAttentionImpl(AttentionImpl): ...@@ -622,6 +702,7 @@ class FlashAttentionImpl(AttentionImpl):
self.attn_type = attn_type self.attn_type = attn_type
self.vllm_flash_attn_version = get_flash_attn_version() self.vllm_flash_attn_version = get_flash_attn_version()
self.use_alibi_sqrt = use_alibi_sqrt
# Cache the batch invariant result for use in forward passes # Cache the batch invariant result for use in forward passes
self.batch_invariant_enabled = vllm_is_batch_invariant() self.batch_invariant_enabled = vllm_is_batch_invariant()
...@@ -648,6 +729,14 @@ class FlashAttentionImpl(AttentionImpl): ...@@ -648,6 +729,14 @@ class FlashAttentionImpl(AttentionImpl):
else False else False
) )
def _get_unified_extras(
self,
attn_metadata: FlashAttentionMetadata,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
mm_prefix_range_tensor = attn_metadata.mm_prefix_range_tensor
qq_bias = attn_metadata.qq_bias
return mm_prefix_range_tensor, qq_bias
def forward( def forward(
self, self,
layer: torch.nn.Module, layer: torch.nn.Module,
...@@ -774,30 +863,60 @@ class FlashAttentionImpl(AttentionImpl): ...@@ -774,30 +863,60 @@ class FlashAttentionImpl(AttentionImpl):
print(f"q.shape = {query[:num_actual_tokens].shape}, key_cache.shape = {key_cache.shape}, value_cache.shape = {value_cache.shape}") print(f"q.shape = {query[:num_actual_tokens].shape}, key_cache.shape = {key_cache.shape}, value_cache.shape = {value_cache.shape}")
print(f"cu_seqlens_q.shape = {cu_seqlens_q.shape}, max_seqlen_q = {max_seqlen_q}, seqused_k.shape = {seqused_k.shape}, max_seqlen_k = {max_seqlen_k}") print(f"cu_seqlens_q.shape = {cu_seqlens_q.shape}, max_seqlen_q = {max_seqlen_q}, seqused_k.shape = {seqused_k.shape}, max_seqlen_k = {max_seqlen_k}")
print(f"softmax_scale = {self.scale:.3f}, alibi_slopes = {self.alibi_slopes}, window_size = {self.sliding_window}, block_tables.shape = {block_table.shape}, softcap = {self.logits_soft_cap}, scheduler_metadata = {scheduler_metadata}") print(f"softmax_scale = {self.scale:.3f}, alibi_slopes = {self.alibi_slopes}, window_size = {self.sliding_window}, block_tables.shape = {block_table.shape}, softcap = {self.logits_soft_cap}, scheduler_metadata = {scheduler_metadata}")
vllm_flash_attn_varlen_func( use_unified_kv_layout = (
q=query[:num_actual_tokens], FlashAttentionBackend._use_rocm_unified_kv_layout(
k=key_cache, key_cache=key_cache, value_cache=value_cache)
v=value_cache,
out=output[:num_actual_tokens],
cu_seqlens_q=cu_seqlens_q,
max_seqlen_q=max_seqlen_q,
seqused_k=seqused_k,
max_seqlen_k=max_seqlen_k,
softmax_scale=self.scale,
causal=attn_metadata.causal,
alibi_slopes=self.alibi_slopes,
window_size=sliding_window_size,
block_table=block_table,
softcap=self.logits_soft_cap,
scheduler_metadata=scheduler_metadata,
fa_version=self.vllm_flash_attn_version,
q_descale=q_descale,
k_descale=k_descale,
v_descale=v_descale,
# num_splits=attn_metadata.max_num_splits,
s_aux=self.sinks,
is_prefix_cache=True,
) )
if use_unified_kv_layout:
mm_prefix_range_tensor, qq_bias = self._get_unified_extras(
attn_metadata
)
varlen_fwd_unified(
q=query[:num_actual_tokens],
k=key_cache,
v=value_cache,
cu_seqlens_q=cu_seqlens_q,
seqused_k=seqused_k,
block_table=block_table,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
softmax_scale=self.scale,
causal=attn_metadata.causal,
softcap=self.logits_soft_cap,
window_size=tuple(self.sliding_window),
alibi_slopes=self.alibi_slopes,
use_alibi_sqrt=self.use_alibi_sqrt,
qq_bias=qq_bias,
s_aux=self.sinks,
mm_prefix_range=mm_prefix_range_tensor,
return_softmax_lse=False,
out=output[:num_actual_tokens],
)
else:
vllm_flash_attn_varlen_func(
q=query[:num_actual_tokens],
k=key_cache,
v=value_cache,
out=output[:num_actual_tokens],
cu_seqlens_q=cu_seqlens_q,
max_seqlen_q=max_seqlen_q,
seqused_k=seqused_k,
max_seqlen_k=max_seqlen_k,
softmax_scale=self.scale,
causal=attn_metadata.causal,
alibi_slopes=self.alibi_slopes,
window_size=sliding_window_size,
block_table=block_table,
softcap=self.logits_soft_cap,
scheduler_metadata=scheduler_metadata,
fa_version=self.vllm_flash_attn_version,
q_descale=q_descale,
k_descale=k_descale,
v_descale=v_descale,
# num_splits=attn_metadata.max_num_splits,
s_aux=self.sinks,
is_prefix_cache=True,
)
else: else:
flash_attn_varlen_func( flash_attn_varlen_func(
q=query[:num_actual_tokens], q=query[:num_actual_tokens],
...@@ -889,21 +1008,11 @@ class FlashAttentionImpl(AttentionImpl): ...@@ -889,21 +1008,11 @@ class FlashAttentionImpl(AttentionImpl):
# op uses the slot_mapping's shape to determine the number of # op uses the slot_mapping's shape to determine the number of
# actual tokens. # actual tokens.
if current_platform.is_rocm(): if current_platform.is_rocm():
if envs.VLLM_USE_OPT_RESHAPE_AND_CACHE: if FlashAttentionBackend._use_rocm_unified_kv_layout(
from lightop import reshape_and_cache_cuda key_cache=key_cache,
reshape_and_cache_cuda( value_cache=value_cache,
key, ):
value, triton_reshape_and_cache_flash(
key_cache,
value_cache,
slot_mapping,
self.kv_cache_dtype,
layer._k_scale,
layer._v_scale
)
else:
from vllm.v1.attention.backends.fa_utils import reshape_and_cache_cuda
reshape_and_cache_cuda(
key, key,
value, value,
key_cache, key_cache,
...@@ -913,6 +1022,32 @@ class FlashAttentionImpl(AttentionImpl): ...@@ -913,6 +1022,32 @@ class FlashAttentionImpl(AttentionImpl):
layer._k_scale, layer._k_scale,
layer._v_scale, layer._v_scale,
) )
else:
if envs.VLLM_USE_OPT_RESHAPE_AND_CACHE:
from lightop import reshape_and_cache_cuda
reshape_and_cache_cuda(
key,
value,
key_cache,
value_cache,
slot_mapping,
self.kv_cache_dtype,
layer._k_scale,
layer._v_scale
)
else:
from vllm.v1.attention.backends.fa_utils import reshape_and_cache_cuda
reshape_and_cache_cuda(
key,
value,
key_cache,
value_cache,
slot_mapping,
self.kv_cache_dtype,
layer._k_scale,
layer._v_scale,
)
else: else:
reshape_and_cache_flash( reshape_and_cache_flash(
key, key,
......
vllm/v1/attention/ops/triton_unified_attention.py
View file @ ea9b8584
...@@ -12,11 +12,6 @@ import torch ...@@ -12,11 +12,6 @@ import torch
from vllm.logger import init_logger from vllm.logger import init_logger
from vllm.platforms import current_platform from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton from vllm.triton_utils import tl, triton
from vllm import envs
try:
from flash_attn import varlen_fwd_unified
except Exception:
varlen_fwd_unified = None
logger = init_logger(__name__) logger = init_logger(__name__)
float8_info = torch.finfo(current_platform.fp8_dtype()) float8_info = torch.finfo(current_platform.fp8_dtype())
...@@ -988,90 +983,61 @@ def unified_attention( ...@@ -988,90 +983,61 @@ def unified_attention(
or num_seqs > seq_threshold_3D or num_seqs > seq_threshold_3D
): ):
# print(f"[2D Triton] k shape: {k.shape}, v shape: {v.shape}") # print(f"[2D Triton] k shape: {k.shape}, v shape: {v.shape}")
if not envs.VLLM_V1_USE_FA_UNIFIED_ATTN_2D: kernel_unified_attention_2d[
# print("Running Triton kernel") (
kernel_unified_attention_2d[ total_num_q_blocks,
( num_kv_heads,
total_num_q_blocks,
num_kv_heads,
)
](
output_ptr=out,
query_ptr=q,
key_cache_ptr=k,
value_cache_ptr=v,
sink_ptr=sinks,
block_tables_ptr=block_table,
seq_lens_ptr=seqused_k,
alibi_slopes_ptr=alibi_slopes,
qq_bias_ptr=qq_bias,
scale=softmax_scale,
k_scale=k_descale,
v_scale=v_descale,
out_scale=1 / output_scale if output_scale is not None else 1.0,
softcap=softcap,
num_query_heads=num_query_heads,
num_queries_per_kv=num_queries_per_kv,
block_table_stride=block_table.stride(0),
query_stride_0=q.stride(0),
query_stride_1=q.stride(1),
output_stride_0=out.stride(0),
output_stride_1=out.stride(1),
qq_bias_stride_0=qq_bias.stride(0) if use_qq_bias else 0,
BLOCK_SIZE=block_size,
TILE_SIZE=TILE_SIZE_PREFILL,
HEAD_SIZE=head_size,
HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
USE_ALIBI_SLOPES=use_alibi_slopes,
USE_ALIBI_SQRT=use_alibi_sqrt,
USE_QQ_BIAS=use_qq_bias,
USE_SOFTCAP=(softcap > 0),
USE_SINKS=(sinks is not None),
USE_MM_PREFIX=use_mm_prefix,
MAX_MM_RANGES=max_mm_ranges,
mm_prefix_range_ptr=mm_prefix_range,
SLIDING_WINDOW=(1 + window_size[0]),
stride_k_cache_0=k.stride(0),
stride_k_cache_1=k.stride(1),
stride_k_cache_2=k.stride(2),
stride_k_cache_3=k.stride(3),
stride_v_cache_0=v.stride(0),
stride_v_cache_1=v.stride(1),
stride_v_cache_2=v.stride(2),
stride_v_cache_3=v.stride(3),
query_start_len_ptr=cu_seqlens_q,
BLOCK_Q=BLOCK_Q,
num_seqs=num_seqs,
BLOCK_M=BLOCK_M,
USE_FP8=output_scale is not None,
)
else:
if varlen_fwd_unified is None:
raise RuntimeError(
"flash_attn.varlen_fwd_unified is not available in this flash-attn version"
)
# print("Running FA kernel")
varlen_fwd_unified(
q=q,
k=k,
v=v,
cu_seqlens_q=cu_seqlens_q,
seqused_k=seqused_k,
block_table=block_table,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
softmax_scale=softmax_scale,
causal=causal,
softcap=softcap,
window_size=window_size,
alibi_slopes=alibi_slopes,
use_alibi_sqrt=use_alibi_sqrt,
qq_bias=qq_bias,
s_aux=sinks,
mm_prefix_range=mm_prefix_range,
return_softmax_lse=False,
out=out,
) )
](
output_ptr=out,
query_ptr=q,
key_cache_ptr=k,
value_cache_ptr=v,
sink_ptr=sinks,
block_tables_ptr=block_table,
seq_lens_ptr=seqused_k,
alibi_slopes_ptr=alibi_slopes,
qq_bias_ptr=qq_bias,
scale=softmax_scale,
k_scale=k_descale,
v_scale=v_descale,
out_scale=1 / output_scale if output_scale is not None else 1.0,
softcap=softcap,
num_query_heads=num_query_heads,
num_queries_per_kv=num_queries_per_kv,
block_table_stride=block_table.stride(0),
query_stride_0=q.stride(0),
query_stride_1=q.stride(1),
output_stride_0=out.stride(0),
output_stride_1=out.stride(1),
qq_bias_stride_0=qq_bias.stride(0) if use_qq_bias else 0,
BLOCK_SIZE=block_size,
TILE_SIZE=TILE_SIZE_PREFILL,
HEAD_SIZE=head_size,
HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
USE_ALIBI_SLOPES=use_alibi_slopes,
USE_ALIBI_SQRT=use_alibi_sqrt,
USE_QQ_BIAS=use_qq_bias,
USE_SOFTCAP=(softcap > 0),
USE_SINKS=(sinks is not None),
USE_MM_PREFIX=use_mm_prefix,
MAX_MM_RANGES=max_mm_ranges,
mm_prefix_range_ptr=mm_prefix_range,
SLIDING_WINDOW=(1 + window_size[0]),
stride_k_cache_0=k.stride(0),
stride_k_cache_1=k.stride(1),
stride_k_cache_2=k.stride(2),
stride_k_cache_3=k.stride(3),
stride_v_cache_0=v.stride(0),
stride_v_cache_1=v.stride(1),
stride_v_cache_2=v.stride(2),
stride_v_cache_3=v.stride(3),
query_start_len_ptr=cu_seqlens_q,
BLOCK_Q=BLOCK_Q,
num_seqs=num_seqs,
BLOCK_M=BLOCK_M,
USE_FP8=output_scale is not None,
)
else: else:
# print(f"[3D Triton] k shape: {k.shape}, v shape: {v.shape}") # print(f"[3D Triton] k shape: {k.shape}, v shape: {v.shape}")
kernel_unified_attention_3d[ kernel_unified_attention_3d[
......
vllm/v1/attention/selector.py
View file @ ea9b8584
...@@ -27,6 +27,7 @@ class AttentionSelectorConfig(NamedTuple): ...@@ -27,6 +27,7 @@ class AttentionSelectorConfig(NamedTuple):
has_sink: bool = False has_sink: bool = False
use_sparse: bool = False use_sparse: bool = False
use_mm_prefix: bool = False use_mm_prefix: bool = False
use_alibi_sqrt: bool = False
attn_type: str = AttentionType.DECODER attn_type: str = AttentionType.DECODER
def __repr__(self): def __repr__(self):
...@@ -39,6 +40,7 @@ class AttentionSelectorConfig(NamedTuple): ...@@ -39,6 +40,7 @@ class AttentionSelectorConfig(NamedTuple):
f"has_sink={self.has_sink}, " f"has_sink={self.has_sink}, "
f"use_sparse={self.use_sparse}, " f"use_sparse={self.use_sparse}, "
f"use_mm_prefix={self.use_mm_prefix}, " f"use_mm_prefix={self.use_mm_prefix}, "
f"use_alibi_sqrt={self.use_alibi_sqrt}, "
f"attn_type={self.attn_type})" f"attn_type={self.attn_type})"
) )
...@@ -52,6 +54,7 @@ def get_attn_backend( ...@@ -52,6 +54,7 @@ def get_attn_backend(
has_sink: bool = False, has_sink: bool = False,
use_sparse: bool = False, use_sparse: bool = False,
use_mm_prefix: bool = False, use_mm_prefix: bool = False,
use_alibi_sqrt: bool = False,
attn_type: str | None = None, attn_type: str | None = None,
) -> type[AttentionBackend]: ) -> type[AttentionBackend]:
"""Selects which attention backend to use and lazily imports it.""" """Selects which attention backend to use and lazily imports it."""
...@@ -77,6 +80,7 @@ def get_attn_backend( ...@@ -77,6 +80,7 @@ def get_attn_backend(
has_sink=has_sink, has_sink=has_sink,
use_sparse=use_sparse, use_sparse=use_sparse,
use_mm_prefix=use_mm_prefix, use_mm_prefix=use_mm_prefix,
use_alibi_sqrt=use_alibi_sqrt,
attn_type=attn_type or AttentionType.DECODER, attn_type=attn_type or AttentionType.DECODER,
) )
......
vllm/v1/worker/gpu_model_runner.py
View file @ ea9b8584
...@@ -5952,7 +5952,7 @@ class GPUModelRunner( ...@@ -5952,7 +5952,7 @@ class GPUModelRunner(
return kv_caches return kv_caches
def _update_hybrid_attention_mamba_layout( def _update_hybrid_attention_mamba_layout(
self, kv_caches: dict[str, torch.Tensor] self, kv_caches: dict[str, Any]
) -> None: ) -> None:
""" """
Update the layout of attention layers from (2, num_blocks, ...) to Update the layout of attention layers from (2, num_blocks, ...) to
...@@ -5966,6 +5966,8 @@ class GPUModelRunner( ...@@ -5966,6 +5966,8 @@ class GPUModelRunner(
kv_cache_spec = group.kv_cache_spec kv_cache_spec = group.kv_cache_spec
for layer_name in group.layer_names: for layer_name in group.layer_names:
kv_cache = kv_caches[layer_name] kv_cache = kv_caches[layer_name]
if not isinstance(kv_cache, torch.Tensor):
continue
if isinstance(kv_cache_spec, AttentionSpec) and kv_cache.shape[0] == 2: if isinstance(kv_cache_spec, AttentionSpec) and kv_cache.shape[0] == 2:
assert kv_cache.shape[1] != 2, ( assert kv_cache.shape[1] != 2, (
"Fail to determine whether the layout is " "Fail to determine whether the layout is "
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment