diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index d28de8fc36..a252b59e9e 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -168,7 +168,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 | +| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `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/model_executor/layers/hpc/rope_norm.py b/vllm/model_executor/layers/hpc/rope_norm.py index 7eee2a6eb3..2b07d949c0 100644 --- a/vllm/model_executor/layers/hpc/rope_norm.py +++ b/vllm/model_executor/layers/hpc/rope_norm.py @@ -7,6 +7,7 @@ Decoupled from HpcAttentionImpl; extra params are passed via layer attrs. from __future__ import annotations +import importlib.util from enum import IntEnum from typing import Any @@ -122,6 +123,12 @@ class HpcRopeNorm(CustomOp, HpcModule): qk_norm_policy: QkNormPolicy = QkNormPolicy.ROPE_THEN_NORM, ) -> None: super().__init__() + if importlib.util.find_spec("hpc") is None: + raise ImportError( + "HPCRopeNorm requires the hpc module to be installed. " + "Please install it from https://github.com/Tencent/hpc-ops" + ) + self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_dim = head_dim @@ -172,6 +179,15 @@ class HpcRopeNorm(CustomOp, HpcModule): self.layer_name: str | None = None self.register_layer_name(layer_name) + import hpc + + if self.use_fp8: + self._quant_type = ( + hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR.value + ) + else: + self._quant_type = None + @classmethod def support( cls, @@ -304,12 +320,6 @@ class HpcRopeNorm(CustomOp, HpcModule): 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:] @@ -333,7 +343,7 @@ class HpcRopeNorm(CustomOp, HpcModule): is_prefill=True, k_scale=k_scale, v_scale=v_scale, - quant_policy=QUANT_POLICY_DQSKV, + quant_policy=self._quant_type, max_seqlens=max_seqlens, q_norm_weight=q_norm_weight, k_norm_weight=k_norm_weight, @@ -364,9 +374,7 @@ class HpcRopeNorm(CustomOp, HpcModule): 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 - ) + decode_query_len = attn_metadata.decode_query_len out_q_decode = output[:num_decode_tokens] if self.use_fp8: @@ -376,13 +384,13 @@ class HpcRopeNorm(CustomOp, HpcModule): qkv=qkv_decode, cos_sin=self.cos_sin_cache, num_seqlen_per_req=num_seq_kvcache, - q_index=qo_indptr_decode, + q_index=attn_metadata.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, + quant_policy=self._quant_type, + max_seqlens=decode_query_len, q_norm_weight=q_norm_weight, k_norm_weight=k_norm_weight, qk_norm_policy=self.qk_norm_policy, @@ -398,7 +406,7 @@ class HpcRopeNorm(CustomOp, HpcModule): qkv_decode, self.cos_sin_cache, num_seq_kvcache, - qo_indptr_decode, + attn_metadata.qo_indptr_decode, block_table_decode, False, # is_prefill q_norm_weight=q_norm_weight, diff --git a/vllm/v1/attention/backends/hpc_attn.py b/vllm/v1/attention/backends/hpc_attn.py index 4a3a4383e2..8c6dcfe516 100644 --- a/vllm/v1/attention/backends/hpc_attn.py +++ b/vllm/v1/attention/backends/hpc_attn.py @@ -31,8 +31,6 @@ from vllm.v1.attention.backend import ( ) 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 @@ -86,9 +84,23 @@ class HpcAttnMetadata(AttentionMetadata): 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).""" + """FP8 per-token-per-head Q scale for decode (persistent buffer). + shape = [max_decode_tokens, num_q_heads], contiguous. + Only the first num_decode_q_scale_tokens rows are valid.""" hpc_split_k_flag: torch.Tensor | None = None - """Split-K flag tensor for FP8 decode (from RopeNorm).""" + """Split-K flag tensor for FP8 decode (persistent buffer). + shape = [max_num_seqs, num_kv_heads], int32.""" + + # --- MTP (Multi-Token Prediction) fields --- + decode_query_len: int = 1 + """Number of query tokens per decode request. + 1 for standard decoding, mtp+1 for speculative decoding (2 or 3).""" + qo_indptr_decode: torch.Tensor | None = None + """Cumulative query offsets for decode requests (GPU tensor). + shape = [num_decodes + 1]. Only set when decode_query_len > 1. + e.g. 3 requests with dql=2: [0, 2, 4, 6].""" + task_map: torch.Tensor | None = None + """Used for HPC dynamic schedule attention""" class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): @@ -105,20 +117,40 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): 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 + import hpc - 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.hpc_dynamic_sched_attn_min_split_len = 1024 - self.cache_dtype = self.cache_config.cache_dtype + # MTP constraint: HPC decode kernel only supports mtp in {0, 1, 2, 3} + spec_config = vllm_config.speculative_config + if ( + spec_config is not None + and spec_config.num_speculative_tokens is not None + and spec_config.num_speculative_tokens > 3 + ): + raise ValueError( + f"HPC attention only supports up to 3 speculative tokens " + f"(mtp ∈ {{0, 1, 2, 3}}), got " + f"num_speculative_tokens={spec_config.num_speculative_tokens}. " + f"Please reduce num_speculative_tokens or use a different " + f"attention backend." + ) - self.global_hyperparameters = infer_global_hyperparameters( - get_per_layer_parameters(vllm_config, layer_names, HpcAttentionImpl) + # Dynamic decode threshold for MTP support. + # _init_reorder_batch_threshold computes: + # no spec_config → threshold=1 (unchanged) + # with spec_config → threshold=1+num_speculative_tokens + self._init_reorder_batch_threshold( + reorder_batch_threshold=1, + supports_spec_as_decode=True, + ) + + self.task_map = hpc.get_attention_decode_task_workspace( + vllm_config.scheduler_config.max_num_seqs, + vllm_config.model_config.max_model_len or 4096, + self.num_kv_heads, + min_process_len=self.hpc_dynamic_sched_attn_min_split_len, ) @override # type: ignore[misc] @@ -128,6 +160,13 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): vllm_config: VllmConfig, kv_cache_spec: AttentionSpec, ) -> AttentionCGSupport: + spec_config = vllm_config.speculative_config + if ( + spec_config is not None + and spec_config.num_speculative_tokens is not None + and spec_config.num_speculative_tokens > 0 + ): + return AttentionCGSupport.UNIFORM_BATCH return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE def build( @@ -143,7 +182,8 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): split_decodes_and_prefills( common_attn_metadata, decode_threshold=self.reorder_batch_threshold, - require_uniform=False, + # MTP requires uniform query lengths across decode requests + require_uniform=(self.reorder_batch_threshold > 1), ) ) @@ -152,7 +192,16 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): slot_mapping = common_attn_metadata.slot_mapping max_query_len = common_attn_metadata.max_query_len + # Compute decode_query_len (tokens per decode request). + # Non-MTP: 1, MTP: mtp+1 (2 or 3). + if num_decodes > 0 and num_decode_tokens > num_decodes: + decode_query_len = num_decode_tokens // num_decodes + else: + decode_query_len = 1 + + seq_lens_decode = None qo_indptr = None + qo_indptr_decode = None if num_prefills > 0: qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu prefill_start = num_decodes @@ -161,6 +210,21 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): ) qo_indptr = qo_indptr_prefill_cpu.to(self.device, non_blocking=True) + if num_decodes > 0: + seq_lens_decode = seq_lens[:num_decodes] + # block_table is per-request, indexed by num_decodes (not tokens) + qo_indptr_decode = common_attn_metadata.query_start_loc[: num_decodes + 1] + import hpc + + hpc.assign_attention_decode_task( + seq_lens_decode, + self.task_map, + self.num_kv_heads, + decode_query_len, + new_kv_included=True, + min_process_len=self.hpc_dynamic_sched_attn_min_split_len, + ) + return HpcAttnMetadata( num_actual_tokens=num_actual_tokens, num_decodes=num_decodes, @@ -176,6 +240,9 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): hpc_prefill_q_scale=None, hpc_decode_q_scale=None, hpc_split_k_flag=None, + decode_query_len=decode_query_len, + qo_indptr_decode=qo_indptr_decode, + task_map=self.task_map, ) @@ -192,6 +259,7 @@ class HpcAttentionBackend(AttentionBackend): ] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ "auto", + "bfloat16", "fp8_e4m3", ] @@ -323,6 +391,13 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): self.supports_quant_query_input = False self.splitk = True + import hpc + + if self.use_fp8: + self._quant_type = hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR + else: + self._quant_type = None + def forward( self, layer: torch.nn.Module, @@ -417,6 +492,7 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): block_table_prefill, seq_lens_prefill, max_seqlens, + quant_type=self._quant_type, output=output_prefill, ) else: @@ -439,6 +515,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): q_decode = query[:num_decode_tokens] output_decode = output[:num_decode_tokens] + mtp = attn_metadata.decode_query_len - 1 + if self.use_fp8: hpc.attention_decode_fp8( q_decode, @@ -449,8 +527,14 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): hpc_decode_q_scale, k_scale, v_scale, + mtp=mtp, + # MTP: split_flag from rope_norm is unavailable + # when using prefill-mode kernel; let HPC decide. + # splitk=(self.splitk if mtp == 0 else True), new_kv_included=True, + quant_type=self._quant_type, splitk=self.splitk, + task_map=attn_metadata.task_map, split_flag=hpc_split_k_flag, output=output_decode, ) @@ -461,6 +545,7 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): kv_cache[:, 1], block_table_decode, num_seq_kvcache, + mtp=mtp, output=output_decode, new_kv_included=True, splitk=self.splitk,