[Attention Backend] add HPC-Ops Attention backend (#46020)
Signed-off-by: chengvjiang <chengvjiang@tencent.com> Co-authored-by: chengvjiang <chengvjiang@tencent.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
Cheng Jiang
co-authored by
chengvjiang
Andreas Karatzas
chengvjiang
Andreas Karatzas
parent
8e9d70fdd5
commit
8cf7c4d8ad
@@ -167,6 +167,7 @@ Priority is **1 = highest** (tried first).
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `HPC_ATTN` | | fp16, bf16 | `auto`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
|
||||
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
|
||||
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
|
||||
|
||||
@@ -759,6 +759,7 @@ class CompilationConfig:
|
||||
"vllm::sparse_attn_indexer",
|
||||
"vllm::rocm_aiter_sparse_attn_indexer",
|
||||
"vllm::deepseek_v4_attention",
|
||||
"vllm::hpc_rope_norm_forward",
|
||||
]
|
||||
|
||||
def compute_hash(self) -> str:
|
||||
|
||||
@@ -458,6 +458,7 @@ class Attention(nn.Module, AttentionLayerBase):
|
||||
# shape does not match the query shape, so we optionally let the model
|
||||
# definition specify the output tensor shape.
|
||||
output_shape: torch.Size | None = None,
|
||||
output_dtype: torch.dtype | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
The KV cache is stored inside this class and is accessed via
|
||||
@@ -472,7 +473,8 @@ class Attention(nn.Module, AttentionLayerBase):
|
||||
torch.ops.vllm.maybe_calc_kv_scales(
|
||||
query, key, value, _encode_layer_name(self.layer_name)
|
||||
)
|
||||
output_dtype = query.dtype
|
||||
if output_dtype is None:
|
||||
output_dtype = query.dtype
|
||||
if self.query_quant is not None:
|
||||
# quantizing with a simple torch operation enables
|
||||
# torch.compile to fuse this into previous ops
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from vllm.model_executor.layers.hpc.hpc_module import HpcModule
|
||||
from vllm.model_executor.layers.hpc.rope_norm import HpcRopeNorm, QkNormPolicy
|
||||
|
||||
__all__ = [
|
||||
"HpcModule",
|
||||
"HpcRopeNorm",
|
||||
"QkNormPolicy",
|
||||
]
|
||||
@@ -0,0 +1,18 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class HpcModule(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@classmethod
|
||||
def support(cls, *args, **kwargs):
|
||||
return True
|
||||
|
||||
def process_weights_after_loading(self, model):
|
||||
pass
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
pass
|
||||
@@ -0,0 +1,408 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant).
|
||||
|
||||
Decoupled from HpcAttentionImpl; extra params are passed via layer attrs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import IntEnum
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.config import get_current_vllm_config_or_none
|
||||
from vllm.forward_context import ForwardContext, get_forward_context
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.custom_op import CustomOp
|
||||
from vllm.model_executor.layers.hpc.hpc_module import HpcModule
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
from vllm.v1.attention.backends.hpc_attn import HpcAttnMetadata
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_hpc_rope_norm_instances: dict[str, HpcRopeNorm] = {}
|
||||
|
||||
|
||||
class QkNormPolicy(IntEnum):
|
||||
"""Order of QK-RMSNorm relative to RoPE in the fused HPC rope_norm kernel.
|
||||
|
||||
The values are part of the HPC kernel ABI (passed through as ints), so they
|
||||
must stay in sync with the kernel's expectations.
|
||||
"""
|
||||
|
||||
# No QK-Norm: apply RoPE only.
|
||||
NONE = 0
|
||||
# Apply RoPE first, then QK-RMSNorm.
|
||||
ROPE_THEN_NORM = 1
|
||||
# Apply QK-RMSNorm first, then RoPE (e.g. HunYuan V3).
|
||||
NORM_THEN_ROPE = 2
|
||||
|
||||
|
||||
def hpc_rope_norm_forward(
|
||||
qkv: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> None:
|
||||
"""Top-level custom op: RoPE + QK-Norm + KV-Cache-Write + FP8 Q quant.
|
||||
|
||||
Fully opaque to torch.compile (dynamo).
|
||||
"""
|
||||
forward_context: ForwardContext = get_forward_context()
|
||||
attn_metadata: Any = forward_context.attn_metadata
|
||||
if isinstance(attn_metadata, dict):
|
||||
attn_metadata = attn_metadata[layer_name]
|
||||
|
||||
if attn_metadata is None:
|
||||
output.zero_()
|
||||
return
|
||||
|
||||
attn_layer = forward_context.no_compile_layers[layer_name]
|
||||
# bind_kv_cache stores the per-layer KV cache as a single 5D tensor
|
||||
# (num_blocks, 2, block_size, num_kv_heads, head_size), so use it directly.
|
||||
kv_cache = attn_layer.kv_cache
|
||||
|
||||
if kv_cache.numel() == 0:
|
||||
output.zero_()
|
||||
return
|
||||
|
||||
assert kv_cache.dim() == 5, (
|
||||
f"Expected kv_cache to have 5 dims, got {tuple(kv_cache.shape)}"
|
||||
)
|
||||
|
||||
rope_norm = _hpc_rope_norm_instances[layer_name]
|
||||
rope_norm._forward_impl(qkv, kv_cache, attn_metadata, attn_layer, output)
|
||||
|
||||
|
||||
def hpc_rope_norm_forward_fake(
|
||||
qkv: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> None:
|
||||
"""Fake impl for torch.compile trace; output is a mutated arg."""
|
||||
return
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="hpc_rope_norm_forward",
|
||||
op_func=hpc_rope_norm_forward,
|
||||
mutates_args=["output"],
|
||||
fake_impl=hpc_rope_norm_forward_fake,
|
||||
)
|
||||
|
||||
|
||||
@CustomOp.register("hpc_rope_norm")
|
||||
class HpcRopeNorm(CustomOp, HpcModule):
|
||||
"""HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant).
|
||||
|
||||
Registered as a sub-module in model layers (e.g. HunYuanAttention).
|
||||
Norm weights are extracted from fallback norm modules via
|
||||
process_weights_after_loading() after all weights are loaded.
|
||||
|
||||
forward() is dispatched by CustomOp framework:
|
||||
- In compiled mode: forward_cuda() calls torch.ops.vllm.hpc_rope_norm_forward
|
||||
as a splitting point — internal Python control flow is opaque
|
||||
to torch.compile and not captured by CUDA Graph.
|
||||
- In eager/native mode: forward_native() falls back to forward_cuda().
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_dim: int,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
use_qk_norm: bool,
|
||||
fallback_qnorm: torch.nn.Module | None,
|
||||
fallback_knorm: torch.nn.Module | None,
|
||||
kv_cache_dtype: str,
|
||||
layer_name: str,
|
||||
qk_norm_policy: QkNormPolicy = QkNormPolicy.ROPE_THEN_NORM,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.num_kv_heads = num_kv_heads
|
||||
self.head_dim = head_dim
|
||||
|
||||
self.use_qk_norm = use_qk_norm
|
||||
|
||||
self.q_size = num_heads * head_dim
|
||||
self.kv_size = num_kv_heads * head_dim
|
||||
|
||||
# Register as a non-persistent buffer so it participates in sleep
|
||||
# level-2 save/restore (CuMemAllocator) but is excluded from the
|
||||
# checkpoint state_dict.
|
||||
self.register_buffer("cos_sin_cache", cos_sin_cache.float(), persistent=False)
|
||||
|
||||
self.fallback_qnorm = fallback_qnorm
|
||||
self.fallback_knorm = fallback_knorm
|
||||
|
||||
self.head_per_group = num_heads // num_kv_heads
|
||||
|
||||
# Pre-allocate norm weight tensors as Parameters so they are tracked by
|
||||
# CuMemAllocator (for sleep/wake_up) and have stable addresses for CUDA
|
||||
# Graph replay. process_weights_after_loading() updates them inplace via
|
||||
# copy_() so refit does not invalidate captured graph tensor pointers.
|
||||
# Shape is [head_dim] to match the HPC kernel's q/k_norm_weight layout.
|
||||
if use_qk_norm and fallback_qnorm is not None:
|
||||
self.qnorm_weight: torch.nn.Parameter | None = torch.nn.Parameter(
|
||||
torch.empty(head_dim, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
else:
|
||||
self.qnorm_weight = None
|
||||
if use_qk_norm and fallback_knorm is not None:
|
||||
self.knorm_weight: torch.nn.Parameter | None = torch.nn.Parameter(
|
||||
torch.empty(head_dim, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
else:
|
||||
self.knorm_weight = None
|
||||
|
||||
self.use_fp8 = "fp8" in kv_cache_dtype
|
||||
# The RMSNorm/RoPE ordering is model dependent (e.g. HunYuan V3 applies
|
||||
# QK-Norm before RoPE -> NORM_THEN_ROPE), so it is supplied by the
|
||||
# caller. When QK-Norm is disabled the policy is forced to NONE.
|
||||
self.qk_norm_policy = qk_norm_policy if use_qk_norm else QkNormPolicy.NONE
|
||||
|
||||
# Register layer_name + add self to the global instance registry so the
|
||||
# module-level custom op (hpc_rope_norm_forward) can route back here.
|
||||
self.layer_name: str | None = None
|
||||
self.register_layer_name(layer_name)
|
||||
|
||||
@classmethod
|
||||
def support(
|
||||
cls,
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_dim: int,
|
||||
kv_cache_dtype: str,
|
||||
) -> bool:
|
||||
"""Check whether HpcRopeNorm is supported for the given config."""
|
||||
# HpcRopeNorm is only enabled together with the HPC attention backend.
|
||||
vllm_config = get_current_vllm_config_or_none()
|
||||
if (
|
||||
vllm_config is None
|
||||
or vllm_config.attention_config.backend != AttentionBackendEnum.HPC_ATTN
|
||||
):
|
||||
return False
|
||||
|
||||
if kv_cache_dtype not in ("fp8_e4m3", "auto"):
|
||||
logger.warning_once(
|
||||
f"hpc rope_norm not support kv_cache_dtype:{kv_cache_dtype}, "
|
||||
"only support fp8_e4m3, bfloat16"
|
||||
)
|
||||
return False
|
||||
|
||||
if head_dim not in (128,):
|
||||
logger.warning_once("hpc rope_norm only support head_dim == 128.")
|
||||
return False
|
||||
|
||||
head_per_group = num_heads // num_kv_heads
|
||||
if head_per_group not in (4, 8):
|
||||
logger.warning_once("hpc rope_norm only support head_per_group in [4, 8].")
|
||||
return False
|
||||
|
||||
logger.info_once("enable hpc rope_norm")
|
||||
return True
|
||||
|
||||
def process_weights_after_loading(self, model: torch.nn.Module = None) -> None:
|
||||
"""Copy norm weights (float32) from fallback norm modules inplace.
|
||||
|
||||
Uses copy_() to preserve tensor addresses for CUDA Graph / refit
|
||||
compatibility. Called by the model's load_weights() after all weights
|
||||
are loaded (and generically from the model loader for DummyModelLoader
|
||||
/ sleep-wake_up reload paths).
|
||||
"""
|
||||
if self.use_qk_norm:
|
||||
if self.fallback_qnorm is not None and self.qnorm_weight is not None:
|
||||
self.qnorm_weight.data.copy_(self.fallback_qnorm.weight.data.float())
|
||||
if self.fallback_knorm is not None and self.knorm_weight is not None:
|
||||
self.knorm_weight.data.copy_(self.fallback_knorm.weight.data.float())
|
||||
|
||||
def register_layer_name(self, layer_name: str) -> None:
|
||||
"""Register layer_name and add self to the global registry.
|
||||
|
||||
The global registry is needed because the bottom-level torch op
|
||||
(hpc_rope_norm_forward) is a module-level function and needs to
|
||||
route back to the correct instance via layer_name.
|
||||
"""
|
||||
self.layer_name = layer_name
|
||||
_hpc_rope_norm_instances[layer_name] = self
|
||||
logger.debug(
|
||||
"[rope_norm] registered HpcRopeNorm for layer: %s",
|
||||
layer_name,
|
||||
)
|
||||
|
||||
def forward_native(
|
||||
self,
|
||||
qkv: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> torch.Tensor:
|
||||
"""Native fallback path: delegates to forward_cuda().
|
||||
|
||||
For now, the default native path will use CUDA backend path.
|
||||
Other platforms may override via OOT registration.
|
||||
"""
|
||||
return self.forward_cuda(qkv, layer_name)
|
||||
|
||||
def forward_cuda(
|
||||
self,
|
||||
qkv: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> torch.Tensor:
|
||||
"""CUDA path: invoke the torch custom op as a compile splitting point."""
|
||||
num_tokens = qkv.shape[0]
|
||||
output = torch.empty(
|
||||
(num_tokens, self.num_heads, self.head_dim),
|
||||
dtype=torch.float8_e4m3fn if self.use_fp8 else qkv.dtype,
|
||||
device=qkv.device,
|
||||
)
|
||||
|
||||
torch.ops.vllm.hpc_rope_norm_forward(qkv, output, layer_name)
|
||||
return output
|
||||
|
||||
def _forward_impl(
|
||||
self,
|
||||
qkv: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
attn_metadata: HpcAttnMetadata,
|
||||
attn_layer: torch.nn.Module,
|
||||
output: torch.Tensor,
|
||||
) -> None:
|
||||
"""Actual forward logic called by the custom op.
|
||||
|
||||
Writes processed q into *output* and attaches extra params
|
||||
(e.g. FP8 scales) to *attn_layer* as attributes.
|
||||
"""
|
||||
import hpc
|
||||
|
||||
num_actual_tokens = attn_metadata.num_actual_tokens
|
||||
num_prefill_reqs = attn_metadata.num_prefills
|
||||
num_decode_reqs = attn_metadata.num_decodes
|
||||
num_decode_tokens = attn_metadata.num_decode_tokens
|
||||
|
||||
qkv = qkv[:num_actual_tokens]
|
||||
|
||||
num_prefill_tokens = num_actual_tokens - num_decode_tokens
|
||||
|
||||
# KV cache for the FP8 path is stored as uint8; view it as fp8 so the
|
||||
# rope_norm_store_kv_fp8 kernel can write quantized K/V in-place.
|
||||
if self.use_fp8:
|
||||
kv_cache = kv_cache.view(torch.float8_e4m3fn)
|
||||
|
||||
# Per-tensor K/V scales (shape [1]) used by the FP8 kernel.
|
||||
k_scale = attn_layer._k_scale.reshape(1)
|
||||
v_scale = attn_layer._v_scale.reshape(1)
|
||||
|
||||
q_norm_weight = (
|
||||
self.qnorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None
|
||||
)
|
||||
k_norm_weight = (
|
||||
self.knorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None
|
||||
)
|
||||
|
||||
# Dynamic per-token-per-head Q quant + per-tensor K/V (dqskv).
|
||||
# rope_norm_store_kv_fp8 is registered as a torch op whose ``quant_policy``
|
||||
# argument is typed as ``int``; pybind cannot cast the hpc.QuantType enum
|
||||
# automatically, so pass its integer ``.value``.
|
||||
QUANT_POLICY_DQSKV = hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR.value
|
||||
|
||||
# --- Prefill ---
|
||||
if num_prefill_reqs > 0:
|
||||
seq_lens_prefill = attn_metadata.seq_lens[num_decode_reqs:]
|
||||
cu_seqlens_prefill = attn_metadata.qo_indptr
|
||||
max_seqlens = attn_metadata.max_query_len
|
||||
block_table_prefill = attn_metadata.block_table_tensor[num_decode_reqs:]
|
||||
qkv_prefill = qkv[num_decode_tokens:]
|
||||
out_q_prefill = output[
|
||||
num_decode_tokens : num_decode_tokens + num_prefill_tokens
|
||||
]
|
||||
|
||||
if self.use_fp8:
|
||||
_, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8(
|
||||
key_cache=kv_cache[:, 0],
|
||||
value_cache=kv_cache[:, 1],
|
||||
qkv=qkv_prefill,
|
||||
cos_sin=self.cos_sin_cache,
|
||||
num_seqlen_per_req=seq_lens_prefill,
|
||||
q_index=cu_seqlens_prefill,
|
||||
kvcache_indices=block_table_prefill,
|
||||
is_prefill=True,
|
||||
k_scale=k_scale,
|
||||
v_scale=v_scale,
|
||||
quant_policy=QUANT_POLICY_DQSKV,
|
||||
max_seqlens=max_seqlens,
|
||||
q_norm_weight=q_norm_weight,
|
||||
k_norm_weight=k_norm_weight,
|
||||
qk_norm_policy=self.qk_norm_policy,
|
||||
out_q=out_q_prefill,
|
||||
)
|
||||
attn_metadata.hpc_prefill_q_scale = q_scale
|
||||
else:
|
||||
hpc.rope_norm_store_kv(
|
||||
kv_cache[:, 0],
|
||||
kv_cache[:, 1],
|
||||
qkv_prefill,
|
||||
self.cos_sin_cache,
|
||||
seq_lens_prefill,
|
||||
cu_seqlens_prefill,
|
||||
block_table_prefill,
|
||||
True, # is_prefill
|
||||
q_norm_weight=q_norm_weight,
|
||||
k_norm_weight=k_norm_weight,
|
||||
out_q=out_q_prefill,
|
||||
qk_norm_policy=self.qk_norm_policy,
|
||||
)
|
||||
|
||||
# --- Decode ---
|
||||
if num_decode_reqs > 0:
|
||||
num_seq_kvcache = attn_metadata.seq_lens[:num_decode_reqs]
|
||||
block_table_decode = attn_metadata.block_table_tensor[:num_decode_reqs]
|
||||
qkv_decode = qkv[:num_decode_tokens]
|
||||
# Single-token decode: q_index is the per-request prefix sum
|
||||
# [0, 1, ..., num_decode_reqs].
|
||||
qo_indptr_decode = torch.arange(
|
||||
num_decode_reqs + 1, dtype=torch.int32, device=qkv.device
|
||||
)
|
||||
out_q_decode = output[:num_decode_tokens]
|
||||
|
||||
if self.use_fp8:
|
||||
_, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8(
|
||||
key_cache=kv_cache[:, 0],
|
||||
value_cache=kv_cache[:, 1],
|
||||
qkv=qkv_decode,
|
||||
cos_sin=self.cos_sin_cache,
|
||||
num_seqlen_per_req=num_seq_kvcache,
|
||||
q_index=qo_indptr_decode,
|
||||
kvcache_indices=block_table_decode,
|
||||
is_prefill=False,
|
||||
k_scale=k_scale,
|
||||
v_scale=v_scale,
|
||||
quant_policy=QUANT_POLICY_DQSKV,
|
||||
max_seqlens=1,
|
||||
q_norm_weight=q_norm_weight,
|
||||
k_norm_weight=k_norm_weight,
|
||||
qk_norm_policy=self.qk_norm_policy,
|
||||
out_q=out_q_decode,
|
||||
)
|
||||
attn_metadata.hpc_decode_q_scale = q_scale
|
||||
if split_k_flag is not None:
|
||||
attn_metadata.hpc_split_k_flag = split_k_flag
|
||||
else:
|
||||
hpc.rope_norm_store_kv(
|
||||
kv_cache[:, 0],
|
||||
kv_cache[:, 1],
|
||||
qkv_decode,
|
||||
self.cos_sin_cache,
|
||||
num_seq_kvcache,
|
||||
qo_indptr_decode,
|
||||
block_table_decode,
|
||||
False, # is_prefill
|
||||
q_norm_weight=q_norm_weight,
|
||||
k_norm_weight=k_norm_weight,
|
||||
out_q=out_q_decode,
|
||||
qk_norm_policy=self.qk_norm_policy,
|
||||
)
|
||||
@@ -20,6 +20,7 @@ from vllm.model_executor.layers.attention import (
|
||||
MLAAttention,
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.model_executor.layers.hpc import HpcModule
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
@@ -125,6 +126,15 @@ def process_weights_after_loading(
|
||||
with device_loading_context(module, target_device):
|
||||
module.process_weights_after_loading(model_config.dtype)
|
||||
|
||||
# Process HPC modules (HpcRopeNorm, etc.) that rely on
|
||||
# process_weights_after_loading being called from the model's
|
||||
# load_weights(). When using DummyModelLoader (e.g. profiling or
|
||||
# sleep/wake_up reload), the model's load_weights() is not called, so we
|
||||
# must handle HPC modules here generically.
|
||||
for _, module in model.named_modules():
|
||||
if isinstance(module, HpcModule):
|
||||
module.process_weights_after_loading(model)
|
||||
|
||||
# Needed for torchao model reloading via model.reload_weights
|
||||
# @kylesayrs @jerryzh168 this can be removed if callers move to `reload_weights`
|
||||
if model_config.quantization == "torchao":
|
||||
|
||||
@@ -48,6 +48,7 @@ from vllm.model_executor.layers.fused_moe import (
|
||||
GateLinear,
|
||||
fused_moe_make_expert_params_mapping,
|
||||
)
|
||||
from vllm.model_executor.layers.hpc import HpcRopeNorm, QkNormPolicy
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.linear import (
|
||||
MergedColumnParallelLinear,
|
||||
@@ -234,6 +235,7 @@ class HYV3Attention(nn.Module):
|
||||
dual_chunk_attention_config: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.dtype = torch.get_default_dtype()
|
||||
self.hidden_size = hidden_size
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
self.total_num_heads = num_heads
|
||||
@@ -276,11 +278,18 @@ class HYV3Attention(nn.Module):
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.o_proj",
|
||||
)
|
||||
# When the HPC fused RoPE+QK-Norm path is enabled, the RoPE cos/sin
|
||||
# cache must be float32 to match the HPC kernel's expectations.
|
||||
kv_cache_dtype = cache_config.cache_dtype if cache_config else "auto"
|
||||
rope_support = HpcRopeNorm.support(
|
||||
self.num_heads, self.num_kv_heads, self.head_dim, kv_cache_dtype
|
||||
)
|
||||
self.rotary_emb = get_rope(
|
||||
self.head_dim,
|
||||
max_position=max_position_embeddings,
|
||||
rope_parameters=rope_parameters,
|
||||
is_neox_style=True,
|
||||
dtype=torch.float32 if rope_support else torch.get_default_dtype(),
|
||||
)
|
||||
self.attn = Attention(
|
||||
self.num_heads,
|
||||
@@ -295,6 +304,27 @@ class HYV3Attention(nn.Module):
|
||||
self.q_norm = RMSNorm(self.head_dim, rms_norm_eps)
|
||||
self.k_norm = RMSNorm(self.head_dim, rms_norm_eps)
|
||||
|
||||
# HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant).
|
||||
# HunYuan V3 applies QK-Norm *before* RoPE, so NORM_THEN_ROPE.
|
||||
self.hpc_rope_norm: HpcRopeNorm | None = None
|
||||
if rope_support:
|
||||
self.hpc_rope_norm = HpcRopeNorm(
|
||||
num_heads=self.num_heads,
|
||||
num_kv_heads=self.num_kv_heads,
|
||||
head_dim=self.head_dim,
|
||||
cos_sin_cache=self.rotary_emb.cos_sin_cache,
|
||||
use_qk_norm=self.use_qk_norm,
|
||||
fallback_qnorm=self.q_norm if self.use_qk_norm else None,
|
||||
fallback_knorm=self.k_norm if self.use_qk_norm else None,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
layer_name=self.attn.layer_name,
|
||||
qk_norm_policy=QkNormPolicy.NORM_THEN_ROPE,
|
||||
)
|
||||
# FP8 Q is produced by HpcRopeNorm, so the attention layer must not
|
||||
# re-quantize the query.
|
||||
if self.hpc_rope_norm.use_fp8 and hasattr(self.attn, "query_quant"):
|
||||
self.attn.query_quant = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
@@ -303,20 +333,28 @@ class HYV3Attention(nn.Module):
|
||||
qkv, _ = self.qkv_proj(hidden_states)
|
||||
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
||||
output_shape = None
|
||||
if self.use_qk_norm:
|
||||
q_by_head = q.view(
|
||||
*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim
|
||||
)
|
||||
q_by_head = self.q_norm(q_by_head)
|
||||
q = q_by_head.view(q.shape)
|
||||
if self.hpc_rope_norm is not None:
|
||||
# HPC handles QK-Norm + RoPE + KV-cache write (+ optional FP8 Q
|
||||
# quant) internally and returns the processed query. K/V are
|
||||
# written into the paged cache by the fused op.
|
||||
q = self.hpc_rope_norm(qkv, self.attn.layer_name)
|
||||
q = q.view(-1, self.num_heads * self.head_dim)
|
||||
attn_output = self.attn(q, k, v, output_shape, self.dtype)
|
||||
else:
|
||||
if self.use_qk_norm:
|
||||
q_by_head = q.view(
|
||||
*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim
|
||||
)
|
||||
q_by_head = self.q_norm(q_by_head)
|
||||
q = q_by_head.view(q.shape)
|
||||
|
||||
k_by_head = k.view(
|
||||
*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim
|
||||
)
|
||||
k_by_head = self.k_norm(k_by_head)
|
||||
k = k_by_head.view(k.shape)
|
||||
q, k = self.rotary_emb(positions, q, k)
|
||||
attn_output = self.attn(q, k, v, output_shape)
|
||||
k_by_head = k.view(
|
||||
*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim
|
||||
)
|
||||
k_by_head = self.k_norm(k_by_head)
|
||||
k = k_by_head.view(k.shape)
|
||||
q, k = self.rotary_emb(positions, q, k)
|
||||
attn_output = self.attn(q, k, v, output_shape)
|
||||
attn_output = attn_output.view(q.shape[0], -1)
|
||||
output, _ = self.o_proj(attn_output)
|
||||
return output
|
||||
|
||||
@@ -0,0 +1,469 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""HPC Attention Backend.
|
||||
|
||||
Pure attention (prefill + decode), without RoPE or RMSNorm.
|
||||
Independent metadata / builder; KV cache layout is NHD:
|
||||
(num_blocks, 2, block_size, num_kv_heads, head_size).
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
from dataclasses import dataclass
|
||||
from typing import ClassVar
|
||||
|
||||
import torch
|
||||
from typing_extensions import override
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.config.cache import CacheDType
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
from vllm.v1.attention.backend import (
|
||||
AttentionBackend,
|
||||
AttentionCGSupport,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
AttentionType,
|
||||
CommonAttentionMetadata,
|
||||
MultipleOf,
|
||||
)
|
||||
from vllm.v1.attention.backends.utils import (
|
||||
KVCacheLayoutType,
|
||||
get_per_layer_parameters,
|
||||
infer_global_hyperparameters,
|
||||
split_decodes_and_prefills,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import AttentionSpec
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
|
||||
def _get_fp8_dtype_for_kv_cache(kv_cache_dtype: str) -> torch.dtype:
|
||||
"""Return the torch FP8 dtype for the given kv_cache_dtype string."""
|
||||
if kv_cache_dtype in ("fp8", "fp8_e4m3"):
|
||||
return torch.float8_e4m3fn
|
||||
elif kv_cache_dtype == "fp8_e5m2":
|
||||
return torch.float8_e5m2
|
||||
else:
|
||||
raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class HpcAttnMetadata(AttentionMetadata):
|
||||
"""Metadata required by the HPC attention kernel."""
|
||||
|
||||
num_actual_tokens: int
|
||||
num_decodes: int
|
||||
num_decode_tokens: int
|
||||
num_prefills: int
|
||||
num_prefill_tokens: int
|
||||
max_query_len: int
|
||||
|
||||
slot_mapping: torch.Tensor
|
||||
"""Slot mapping for KV cache writes. shape = [num_actual_tokens]"""
|
||||
|
||||
seq_lens: torch.Tensor
|
||||
"""KV cache length per request. shape = [batch_size]"""
|
||||
|
||||
block_table_tensor: torch.Tensor
|
||||
"""Paged KV-cache block table.
|
||||
shape = [batch_size, max_num_blocks_per_seq]"""
|
||||
|
||||
qo_indptr: torch.Tensor | None = None
|
||||
"""Cumulative query lengths for prefill requests (GPU tensor).
|
||||
shape = [num_prefills + 1]. None when num_prefills == 0."""
|
||||
|
||||
# --- HPC RopeNorm pass-through fields ---
|
||||
# Set by HpcRopeNorm._forward_impl(); consumed & reset by
|
||||
# HpcAttentionImpl.forward(). Defaults are safe for the standard
|
||||
# (non-RopeNorm) path and for profiling runs (attn_metadata=None).
|
||||
hpc_kv_written: bool = False
|
||||
"""True when HpcRopeNorm already wrote KV cache."""
|
||||
hpc_prefill_q_scale: torch.Tensor | None = None
|
||||
"""FP8 per-token-per-head Q scale for prefill (from RopeNorm)."""
|
||||
hpc_decode_q_scale: torch.Tensor | None = None
|
||||
"""FP8 per-token-per-head Q scale for decode (from RopeNorm)."""
|
||||
hpc_split_k_flag: torch.Tensor | None = None
|
||||
"""Split-K flag tensor for FP8 decode (from RopeNorm)."""
|
||||
|
||||
|
||||
class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
|
||||
"""Build HpcAttnMetadata from CommonAttentionMetadata."""
|
||||
|
||||
_cudagraph_support = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
|
||||
reorder_batch_threshold: int = 1
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kv_cache_spec: AttentionSpec,
|
||||
layer_names: list[str],
|
||||
vllm_config: VllmConfig,
|
||||
device: torch.device,
|
||||
):
|
||||
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
|
||||
self.model_config = vllm_config.model_config
|
||||
self.cache_config = vllm_config.cache_config
|
||||
|
||||
self.num_qo_heads = self.model_config.get_num_attention_heads(
|
||||
vllm_config.parallel_config
|
||||
)
|
||||
self.num_kv_heads = kv_cache_spec.num_kv_heads
|
||||
self.head_dim = kv_cache_spec.head_size
|
||||
self.page_size = kv_cache_spec.block_size
|
||||
|
||||
self.cache_dtype = self.cache_config.cache_dtype
|
||||
|
||||
self.global_hyperparameters = infer_global_hyperparameters(
|
||||
get_per_layer_parameters(vllm_config, layer_names, HpcAttentionImpl)
|
||||
)
|
||||
|
||||
@override # type: ignore[misc]
|
||||
@classmethod
|
||||
def get_cudagraph_support(
|
||||
cls: type["HpcAttnMetadataBuilder"],
|
||||
vllm_config: VllmConfig,
|
||||
kv_cache_spec: AttentionSpec,
|
||||
) -> AttentionCGSupport:
|
||||
return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
|
||||
|
||||
def build(
|
||||
self,
|
||||
common_prefix_len: int,
|
||||
common_attn_metadata: CommonAttentionMetadata,
|
||||
fast_build: bool = False,
|
||||
) -> HpcAttnMetadata:
|
||||
"""Build HpcAttnMetadata from CommonAttentionMetadata."""
|
||||
num_actual_tokens = common_attn_metadata.num_actual_tokens
|
||||
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
split_decodes_and_prefills(
|
||||
common_attn_metadata,
|
||||
decode_threshold=self.reorder_batch_threshold,
|
||||
require_uniform=False,
|
||||
)
|
||||
)
|
||||
|
||||
seq_lens = common_attn_metadata.seq_lens
|
||||
block_table_tensor = common_attn_metadata.block_table_tensor
|
||||
slot_mapping = common_attn_metadata.slot_mapping
|
||||
max_query_len = common_attn_metadata.max_query_len
|
||||
|
||||
qo_indptr = None
|
||||
if num_prefills > 0:
|
||||
qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu
|
||||
prefill_start = num_decodes
|
||||
qo_indptr_prefill_cpu = (
|
||||
qo_indptr_cpu[prefill_start:] - qo_indptr_cpu[prefill_start]
|
||||
)
|
||||
qo_indptr = qo_indptr_prefill_cpu.to(self.device, non_blocking=True)
|
||||
|
||||
return HpcAttnMetadata(
|
||||
num_actual_tokens=num_actual_tokens,
|
||||
num_decodes=num_decodes,
|
||||
num_decode_tokens=num_decode_tokens,
|
||||
num_prefills=num_prefills,
|
||||
num_prefill_tokens=num_prefill_tokens,
|
||||
max_query_len=max_query_len,
|
||||
slot_mapping=slot_mapping,
|
||||
seq_lens=seq_lens,
|
||||
block_table_tensor=block_table_tensor,
|
||||
qo_indptr=qo_indptr,
|
||||
hpc_kv_written=True,
|
||||
hpc_prefill_q_scale=None,
|
||||
hpc_decode_q_scale=None,
|
||||
hpc_split_k_flag=None,
|
||||
)
|
||||
|
||||
|
||||
class HpcAttentionBackend(AttentionBackend):
|
||||
"""HPC attention backend (pure attention, no RoPE/Norm).
|
||||
|
||||
KV cache layout: NHD (num_blocks, 2, block_size, num_kv_heads, head_size).
|
||||
"""
|
||||
|
||||
accept_output_buffer: bool = True
|
||||
supported_dtypes: ClassVar[list[torch.dtype]] = [
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
]
|
||||
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
|
||||
"auto",
|
||||
"fp8_e4m3",
|
||||
]
|
||||
|
||||
# Avoid attention abstracted method call cache insert
|
||||
forward_includes_kv_cache_update: bool = True
|
||||
|
||||
@staticmethod
|
||||
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
|
||||
return [64]
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "HPC_ATTN"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["HpcAttentionImpl"]:
|
||||
return HpcAttentionImpl
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["HpcAttnMetadataBuilder"]:
|
||||
return HpcAttnMetadataBuilder
|
||||
|
||||
@staticmethod
|
||||
def get_kv_cache_shape(
|
||||
num_blocks: int,
|
||||
block_size: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
cache_dtype_str: str = "auto",
|
||||
) -> tuple[int, ...]:
|
||||
return (num_blocks, 2, block_size, num_kv_heads, head_size)
|
||||
|
||||
@staticmethod
|
||||
def get_kv_cache_stride_order(
|
||||
include_num_layers_dimension: bool = False,
|
||||
) -> tuple[int, ...]:
|
||||
if include_num_layers_dimension:
|
||||
return (1, 0, 2, 3, 4, 5)
|
||||
return (0, 1, 2, 3, 4)
|
||||
|
||||
@classmethod
|
||||
def get_supported_head_sizes(cls) -> list[int]:
|
||||
return [128]
|
||||
|
||||
@classmethod
|
||||
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
|
||||
return capability >= DeviceCapability(9, 0)
|
||||
|
||||
@classmethod
|
||||
def supports_kv_cache_dtype(cls, kv_cache_dtype: "CacheDType | None") -> bool:
|
||||
if kv_cache_dtype is None:
|
||||
return True
|
||||
return kv_cache_dtype in cls.supported_kv_cache_dtypes
|
||||
|
||||
@classmethod
|
||||
def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None:
|
||||
return "NHD"
|
||||
|
||||
|
||||
class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
|
||||
"""HPC pure attention implementation (no RoPE/Norm).
|
||||
|
||||
Constraints:
|
||||
- head_dim == 128
|
||||
- num_heads // num_kv_heads in {4, 8}
|
||||
- kv_cache_dtype in {"auto", "fp8_e4m3"}
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
scale: float,
|
||||
num_kv_heads: int,
|
||||
alibi_slopes: list[float] | None = None,
|
||||
sliding_window: int | None = None,
|
||||
kv_cache_dtype: str = "auto",
|
||||
logits_soft_cap: float | None = None,
|
||||
attn_type: str = AttentionType.DECODER,
|
||||
kv_sharing_target_layer_name: str | None = None,
|
||||
) -> None:
|
||||
if importlib.util.find_spec("hpc") is None:
|
||||
raise ImportError(
|
||||
"HPC attention requires the hpc module to be installed. "
|
||||
"Please install it from https://github.com/Tencent/hpc-ops"
|
||||
)
|
||||
if attn_type != AttentionType.DECODER:
|
||||
raise NotImplementedError("HPC attention only supports decoder attention")
|
||||
if alibi_slopes is not None:
|
||||
raise NotImplementedError("HPC attention does not support ALiBi")
|
||||
if logits_soft_cap is not None:
|
||||
raise NotImplementedError("HPC attention does not support logits_soft_cap")
|
||||
|
||||
if head_size != 128:
|
||||
raise ValueError(
|
||||
f"HPC attention only supports head_dim=128, got {head_size}"
|
||||
)
|
||||
|
||||
num_queries_per_kv = num_heads // num_kv_heads
|
||||
if num_queries_per_kv not in (4, 8):
|
||||
raise ValueError(
|
||||
f"HPC attention only supports head_per_group in {{4, 8}}, "
|
||||
f"got {num_queries_per_kv} "
|
||||
f"(num_heads={num_heads}, num_kv_heads={num_kv_heads})"
|
||||
)
|
||||
|
||||
if kv_cache_dtype not in ("auto", "fp8_e4m3"):
|
||||
raise ValueError(
|
||||
f"HPC attention only supports kv_cache_dtype 'auto' or "
|
||||
f"'fp8_e4m3', got '{kv_cache_dtype}'"
|
||||
)
|
||||
|
||||
self.num_heads = num_heads
|
||||
self.head_size = head_size
|
||||
self.scale = float(scale)
|
||||
self.num_kv_heads = num_kv_heads
|
||||
self.kv_cache_dtype = kv_cache_dtype
|
||||
self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
|
||||
|
||||
self.num_queries_per_kv = num_queries_per_kv
|
||||
|
||||
if sliding_window is None:
|
||||
self.sliding_window = (-1, -1)
|
||||
else:
|
||||
self.sliding_window = (sliding_window - 1, 0)
|
||||
|
||||
self.use_fp8 = kv_cache_dtype == "fp8_e4m3"
|
||||
|
||||
self.supports_quant_query_input = False
|
||||
self.splitk = True
|
||||
|
||||
def forward(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
attn_metadata: HpcAttnMetadata | None,
|
||||
output: torch.Tensor | None = None,
|
||||
output_scale: torch.Tensor | None = None,
|
||||
output_block_scale: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""HPC attention forward (standard vLLM backend interface).
|
||||
|
||||
Two modes:
|
||||
1. Standard: upstream handles RoPE/Norm; this backend writes KV + attn.
|
||||
2. HpcRopeNorm: fused op already did RoPE/Norm/KV-Write/Q-Quant;
|
||||
extra params passed via attn_metadata.hpc_* fields.
|
||||
"""
|
||||
import hpc
|
||||
|
||||
assert output is not None, "Output tensor must be provided."
|
||||
assert output_scale is None, "HPC attention does not support fused output quant"
|
||||
assert output_block_scale is None
|
||||
|
||||
if attn_metadata is None:
|
||||
return output.fill_(0)
|
||||
|
||||
hpc_kv_written = attn_metadata.hpc_kv_written
|
||||
hpc_prefill_q_scale = attn_metadata.hpc_prefill_q_scale
|
||||
hpc_decode_q_scale = attn_metadata.hpc_decode_q_scale
|
||||
hpc_split_k_flag = attn_metadata.hpc_split_k_flag
|
||||
|
||||
num_actual_tokens = attn_metadata.num_actual_tokens
|
||||
num_prefill_reqs = attn_metadata.num_prefills
|
||||
num_decode_reqs = attn_metadata.num_decodes
|
||||
num_decode_tokens = attn_metadata.num_decode_tokens
|
||||
|
||||
# Write KV cache if not already done by HpcRopeNorm.
|
||||
if self.kv_sharing_target_layer_name is None and not hpc_kv_written:
|
||||
torch.ops._C_cache_ops.reshape_and_cache_flash(
|
||||
key,
|
||||
value,
|
||||
kv_cache[:, 0],
|
||||
kv_cache[:, 1],
|
||||
attn_metadata.slot_mapping,
|
||||
self.kv_cache_dtype,
|
||||
layer._k_scale,
|
||||
layer._v_scale,
|
||||
)
|
||||
|
||||
if self.use_fp8:
|
||||
torch_dtype = _get_fp8_dtype_for_kv_cache(self.kv_cache_dtype)
|
||||
kv_cache = kv_cache.view(torch_dtype)
|
||||
|
||||
if self.use_fp8:
|
||||
if not hpc_kv_written:
|
||||
raise RuntimeError(
|
||||
"HpcAttentionImpl: FP8 mode requires HpcRopeNorm. "
|
||||
"Ensure hpc_rope_norm is enabled or set "
|
||||
"kv_cache_dtype='auto' for bf16 mode."
|
||||
f" (layer={getattr(layer, 'layer_name', '?')})"
|
||||
)
|
||||
k_scale = layer._k_scale.reshape(1)
|
||||
v_scale = layer._v_scale.reshape(1)
|
||||
|
||||
query = query[:num_actual_tokens]
|
||||
key = key[:num_actual_tokens]
|
||||
value = value[:num_actual_tokens]
|
||||
output_padded = output
|
||||
output = output[:num_actual_tokens]
|
||||
|
||||
# --- Prefill ---
|
||||
if num_prefill_reqs > 0:
|
||||
seq_lens_prefill = attn_metadata.seq_lens[num_decode_reqs:]
|
||||
cu_seqlens_prefill = attn_metadata.qo_indptr
|
||||
max_seqlens = attn_metadata.max_query_len
|
||||
block_table_prefill = attn_metadata.block_table_tensor[num_decode_reqs:]
|
||||
|
||||
q_prefill = query[num_decode_tokens:]
|
||||
output_prefill = output[num_decode_tokens:]
|
||||
|
||||
if self.use_fp8:
|
||||
hpc.attention_with_kvcache_prefill_fp8(
|
||||
q_prefill,
|
||||
kv_cache[:, 0],
|
||||
kv_cache[:, 1],
|
||||
hpc_prefill_q_scale,
|
||||
k_scale,
|
||||
v_scale,
|
||||
cu_seqlens_prefill,
|
||||
block_table_prefill,
|
||||
seq_lens_prefill,
|
||||
max_seqlens,
|
||||
output=output_prefill,
|
||||
)
|
||||
else:
|
||||
hpc.attention_with_kvcache_prefill_bf16(
|
||||
q_prefill,
|
||||
kv_cache[:, 0],
|
||||
kv_cache[:, 1],
|
||||
cu_seqlens_prefill,
|
||||
block_table_prefill,
|
||||
seq_lens_prefill,
|
||||
max_seqlens,
|
||||
output=output_prefill,
|
||||
)
|
||||
|
||||
# --- Decode ---
|
||||
if num_decode_reqs > 0:
|
||||
num_seq_kvcache = attn_metadata.seq_lens[:num_decode_reqs]
|
||||
block_table_decode = attn_metadata.block_table_tensor[:num_decode_reqs]
|
||||
|
||||
q_decode = query[:num_decode_tokens]
|
||||
output_decode = output[:num_decode_tokens]
|
||||
|
||||
if self.use_fp8:
|
||||
hpc.attention_decode_fp8(
|
||||
q_decode,
|
||||
kv_cache[:, 0],
|
||||
kv_cache[:, 1],
|
||||
block_table_decode,
|
||||
num_seq_kvcache,
|
||||
hpc_decode_q_scale,
|
||||
k_scale,
|
||||
v_scale,
|
||||
new_kv_included=True,
|
||||
splitk=self.splitk,
|
||||
split_flag=hpc_split_k_flag,
|
||||
output=output_decode,
|
||||
)
|
||||
else:
|
||||
hpc.attention_decode_bf16(
|
||||
q_decode,
|
||||
kv_cache[:, 0],
|
||||
kv_cache[:, 1],
|
||||
block_table_decode,
|
||||
num_seq_kvcache,
|
||||
output=output_decode,
|
||||
new_kv_included=True,
|
||||
splitk=self.splitk,
|
||||
)
|
||||
|
||||
return output_padded
|
||||
@@ -103,6 +103,12 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
|
||||
)
|
||||
NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"
|
||||
FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
|
||||
# HPC Attention Backend:
|
||||
# powered by operators from https://github.com/Tencent/hpc-ops.
|
||||
# Only supported on NVIDIA Hopper GPUs (e.g. H20, H200),
|
||||
# currently limited to the Hy3 model,
|
||||
# and requires a block size of 64.
|
||||
HPC_ATTN = "vllm.v1.attention.backends.hpc_attn.HpcAttentionBackend"
|
||||
ROCM_AITER_UNIFIED_ATTN = (
|
||||
"vllm.v1.attention.backends.rocm_aiter_unified_attn."
|
||||
"RocmAiterUnifiedAttentionBackend"
|
||||
|
||||
Reference in New Issue
Block a user