iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Attention Backend] HPC_ATTN backend support mtp and dynamic scheduled attention (#47433)

Browse Source 
Signed-off-by: chengvjiang <chengvjiang@tencent.com>
Co-authored-by: chengvjiang <chengvjiang@tencent.com>
This commit is contained in:
Cheng Jiang
2026-07-05 18:18:25 -07:00
committed by GitHub
co-authored by chengvjiang
parent d2ec433e37
commit 95a248faed
3 changed files with 123 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design/attention_backends.md
+1 -1
View File
@@ -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 |
vllm/model_executor/layers/hpc/rope_norm.py
+22 -14
View File
@@ -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,
vllm/v1/attention/backends/hpc_attn.py
+100 -15
View File
@@ -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,