From 8cf7c4d8ad602d73ff2ec72a101420d47163c136 Mon Sep 17 00:00:00 2001 From: Cheng Jiang Date: Tue, 30 Jun 2026 18:17:43 +0800 Subject: [PATCH] [Attention Backend] add HPC-Ops Attention backend (#46020) Signed-off-by: chengvjiang Co-authored-by: chengvjiang Co-authored-by: Andreas Karatzas --- docs/design/attention_backends.md | 1 + vllm/config/compilation.py | 1 + .../layers/attention/attention.py | 4 +- vllm/model_executor/layers/hpc/__init__.py | 11 + vllm/model_executor/layers/hpc/hpc_module.py | 18 + vllm/model_executor/layers/hpc/rope_norm.py | 408 +++++++++++++++ vllm/model_executor/model_loader/utils.py | 10 + vllm/model_executor/models/hy_v3.py | 64 ++- vllm/v1/attention/backends/hpc_attn.py | 469 ++++++++++++++++++ vllm/v1/attention/backends/registry.py | 6 + 10 files changed, 978 insertions(+), 14 deletions(-) create mode 100644 vllm/model_executor/layers/hpc/__init__.py create mode 100644 vllm/model_executor/layers/hpc/hpc_module.py create mode 100644 vllm/model_executor/layers/hpc/rope_norm.py create mode 100644 vllm/v1/attention/backends/hpc_attn.py diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index f3067dfc85..bd3c6f72c9 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -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 | diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index d65c3cf6f7..c7244d40d6 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -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: diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 3eec58aafa..32562f0d9d 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -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 diff --git a/vllm/model_executor/layers/hpc/__init__.py b/vllm/model_executor/layers/hpc/__init__.py new file mode 100644 index 0000000000..47d39139d4 --- /dev/null +++ b/vllm/model_executor/layers/hpc/__init__.py @@ -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", +] diff --git a/vllm/model_executor/layers/hpc/hpc_module.py b/vllm/model_executor/layers/hpc/hpc_module.py new file mode 100644 index 0000000000..e1efad832e --- /dev/null +++ b/vllm/model_executor/layers/hpc/hpc_module.py @@ -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 diff --git a/vllm/model_executor/layers/hpc/rope_norm.py b/vllm/model_executor/layers/hpc/rope_norm.py new file mode 100644 index 0000000000..7eee2a6eb3 --- /dev/null +++ b/vllm/model_executor/layers/hpc/rope_norm.py @@ -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, + ) diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index fc59acf3d3..6be057bff0 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -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": diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index a4b52e20bd..cb2c96fa96 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/vllm/model_executor/models/hy_v3.py @@ -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 diff --git a/vllm/v1/attention/backends/hpc_attn.py b/vllm/v1/attention/backends/hpc_attn.py new file mode 100644 index 0000000000..4a3a4383e2 --- /dev/null +++ b/vllm/v1/attention/backends/hpc_attn.py @@ -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 diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 5fba5472f1..94fc24f1ca 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -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"