[Attention Backend] HPC_ATTN backend support mtp and dynamic scheduled attention (#47433)
Signed-off-by: chengvjiang <chengvjiang@tencent.com> Co-authored-by: chengvjiang <chengvjiang@tencent.com>
Cheng Jiang committed
Jul 6, 2026 at 09:18 UTC
95a248faed6781360eedcd37d6132a23f2c787da
3 files changed
+123
-30
docs/design/attention_backends.md
+1
-1
@@ -168,7 +168,7 @@ Priority is **1 = highest** (tried first).
168
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
169
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
170
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
171
-| `HPC_ATTN` | | fp16, bf16 | `auto`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
171
+| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
172
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
173
| `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
174
| `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
@@ -7,6 +7,7 @@ Decoupled from HpcAttentionImpl; extra params are passed via layer attrs.
7
8
from __future__ import annotations
9
10
+import importlib.util
11
from enum import IntEnum
12
from typing import Any
13
@@ -122,6 +123,12 @@ class HpcRopeNorm(CustomOp, HpcModule):
123
qk_norm_policy: QkNormPolicy = QkNormPolicy.ROPE_THEN_NORM,
124
) -> None:
125
super().__init__()
126
+ if importlib.util.find_spec("hpc") is None:
127
+ raise ImportError(
128
+ "HPCRopeNorm requires the hpc module to be installed. "
129
+ "Please install it from https://github.com/Tencent/hpc-ops"
130
+ )
131
+
132
self.num_heads = num_heads
133
self.num_kv_heads = num_kv_heads
134
self.head_dim = head_dim
@@ -172,6 +179,15 @@ class HpcRopeNorm(CustomOp, HpcModule):
179
self.layer_name: str | None = None
180
self.register_layer_name(layer_name)
181
182
+ import hpc
183
+
184
+ if self.use_fp8:
185
+ self._quant_type = (
186
+ hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR.value
187
+ )
188
+ else:
189
+ self._quant_type = None
190
+
191
@classmethod
192
def support(
193
cls,
@@ -304,12 +320,6 @@ class HpcRopeNorm(CustomOp, HpcModule):
320
self.knorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None
321
)
322
307
- # Dynamic per-token-per-head Q quant + per-tensor K/V (dqskv).
308
- # rope_norm_store_kv_fp8 is registered as a torch op whose ``quant_policy``
309
- # argument is typed as ``int``; pybind cannot cast the hpc.QuantType enum
310
- # automatically, so pass its integer ``.value``.
311
- QUANT_POLICY_DQSKV = hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR.value
312
-
323
# --- Prefill ---
324
if num_prefill_reqs > 0:
325
seq_lens_prefill = attn_metadata.seq_lens[num_decode_reqs:]
@@ -333,7 +343,7 @@ class HpcRopeNorm(CustomOp, HpcModule):
343
is_prefill=True,
344
k_scale=k_scale,
345
v_scale=v_scale,
336
- quant_policy=QUANT_POLICY_DQSKV,
346
+ quant_policy=self._quant_type,
347
max_seqlens=max_seqlens,
348
q_norm_weight=q_norm_weight,
349
k_norm_weight=k_norm_weight,
@@ -364,9 +374,7 @@ class HpcRopeNorm(CustomOp, HpcModule):
374
qkv_decode = qkv[:num_decode_tokens]
375
# Single-token decode: q_index is the per-request prefix sum
376
# [0, 1, ..., num_decode_reqs].
367
- qo_indptr_decode = torch.arange(
368
- num_decode_reqs + 1, dtype=torch.int32, device=qkv.device
369
- )
377
+ decode_query_len = attn_metadata.decode_query_len
378
out_q_decode = output[:num_decode_tokens]
379
380
if self.use_fp8:
@@ -376,13 +384,13 @@ class HpcRopeNorm(CustomOp, HpcModule):
384
qkv=qkv_decode,
385
cos_sin=self.cos_sin_cache,
386
num_seqlen_per_req=num_seq_kvcache,
379
- q_index=qo_indptr_decode,
387
+ q_index=attn_metadata.qo_indptr_decode,
388
kvcache_indices=block_table_decode,
389
is_prefill=False,
390
k_scale=k_scale,
391
v_scale=v_scale,
384
- quant_policy=QUANT_POLICY_DQSKV,
385
- max_seqlens=1,
392
+ quant_policy=self._quant_type,
393
+ max_seqlens=decode_query_len,
394
q_norm_weight=q_norm_weight,
395
k_norm_weight=k_norm_weight,
396
qk_norm_policy=self.qk_norm_policy,
@@ -398,7 +406,7 @@ class HpcRopeNorm(CustomOp, HpcModule):
406
qkv_decode,
407
self.cos_sin_cache,
408
num_seq_kvcache,
401
- qo_indptr_decode,
409
+ attn_metadata.qo_indptr_decode,
410
block_table_decode,
411
False, # is_prefill
412
q_norm_weight=q_norm_weight,
vllm/v1/attention/backends/hpc_attn.py
+100
-15
@@ -31,8 +31,6 @@ from vllm.v1.attention.backend import (
31
)
32
from vllm.v1.attention.backends.utils import (
33
KVCacheLayoutType,
34
- get_per_layer_parameters,
35
- infer_global_hyperparameters,
34
split_decodes_and_prefills,
35
)
36
from vllm.v1.kv_cache_interface import AttentionSpec
@@ -86,9 +84,23 @@ class HpcAttnMetadata(AttentionMetadata):
84
hpc_prefill_q_scale: torch.Tensor | None = None
85
"""FP8 per-token-per-head Q scale for prefill (from RopeNorm)."""
86
hpc_decode_q_scale: torch.Tensor | None = None
89
- """FP8 per-token-per-head Q scale for decode (from RopeNorm)."""
87
+ """FP8 per-token-per-head Q scale for decode (persistent buffer).
88
+ shape = [max_decode_tokens, num_q_heads], contiguous.
89
+ Only the first num_decode_q_scale_tokens rows are valid."""
90
hpc_split_k_flag: torch.Tensor | None = None
91
- """Split-K flag tensor for FP8 decode (from RopeNorm)."""
91
+ """Split-K flag tensor for FP8 decode (persistent buffer).
92
+ shape = [max_num_seqs, num_kv_heads], int32."""
93
+
94
+ # --- MTP (Multi-Token Prediction) fields ---
95
+ decode_query_len: int = 1
96
+ """Number of query tokens per decode request.
97
+ 1 for standard decoding, mtp+1 for speculative decoding (2 or 3)."""
98
+ qo_indptr_decode: torch.Tensor | None = None
99
+ """Cumulative query offsets for decode requests (GPU tensor).
100
+ shape = [num_decodes + 1]. Only set when decode_query_len > 1.
101
+ e.g. 3 requests with dql=2: [0, 2, 4, 6]."""
102
+ task_map: torch.Tensor | None = None
103
+ """Used for HPC dynamic schedule attention"""
104
105
106
class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
@@ -105,20 +117,40 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
117
device: torch.device,
118
):
119
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
108
- self.model_config = vllm_config.model_config
109
- self.cache_config = vllm_config.cache_config
120
+ import hpc
121
111
- self.num_qo_heads = self.model_config.get_num_attention_heads(
112
- vllm_config.parallel_config
113
- )
122
self.num_kv_heads = kv_cache_spec.num_kv_heads
115
- self.head_dim = kv_cache_spec.head_size
116
- self.page_size = kv_cache_spec.block_size
123
+ self.hpc_dynamic_sched_attn_min_split_len = 1024
124
+
125
+ # MTP constraint: HPC decode kernel only supports mtp in {0, 1, 2, 3}
126
+ spec_config = vllm_config.speculative_config
127
+ if (
128
+ spec_config is not None
129
+ and spec_config.num_speculative_tokens is not None
130
+ and spec_config.num_speculative_tokens > 3
131
+ ):
132
+ raise ValueError(
133
+ f"HPC attention only supports up to 3 speculative tokens "
134
+ f"(mtp ∈ {{0, 1, 2, 3}}), got "
135
+ f"num_speculative_tokens={spec_config.num_speculative_tokens}. "
136
+ f"Please reduce num_speculative_tokens or use a different "
137
+ f"attention backend."
138
+ )
139
118
- self.cache_dtype = self.cache_config.cache_dtype
140
+ # Dynamic decode threshold for MTP support.
141
+ # _init_reorder_batch_threshold computes:
142
+ # no spec_config → threshold=1 (unchanged)
143
+ # with spec_config → threshold=1+num_speculative_tokens
144
+ self._init_reorder_batch_threshold(
145
+ reorder_batch_threshold=1,
146
+ supports_spec_as_decode=True,
147
+ )
148
120
- self.global_hyperparameters = infer_global_hyperparameters(
121
- get_per_layer_parameters(vllm_config, layer_names, HpcAttentionImpl)
149
+ self.task_map = hpc.get_attention_decode_task_workspace(
150
+ vllm_config.scheduler_config.max_num_seqs,
151
+ vllm_config.model_config.max_model_len or 4096,
152
+ self.num_kv_heads,
153
+ min_process_len=self.hpc_dynamic_sched_attn_min_split_len,
154
)
155
156
@override # type: ignore[misc]
@@ -128,6 +160,13 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
160
vllm_config: VllmConfig,
161
kv_cache_spec: AttentionSpec,
162
) -> AttentionCGSupport:
163
+ spec_config = vllm_config.speculative_config
164
+ if (
165
+ spec_config is not None
166
+ and spec_config.num_speculative_tokens is not None
167
+ and spec_config.num_speculative_tokens > 0
168
+ ):
169
+ return AttentionCGSupport.UNIFORM_BATCH
170
return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
171
172
def build(
@@ -143,7 +182,8 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
182
split_decodes_and_prefills(
183
common_attn_metadata,
184
decode_threshold=self.reorder_batch_threshold,
146
- require_uniform=False,
185
+ # MTP requires uniform query lengths across decode requests
186
+ require_uniform=(self.reorder_batch_threshold > 1),
187
)
188
)
189
@@ -152,7 +192,16 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
192
slot_mapping = common_attn_metadata.slot_mapping
193
max_query_len = common_attn_metadata.max_query_len
194
195
+ # Compute decode_query_len (tokens per decode request).
196
+ # Non-MTP: 1, MTP: mtp+1 (2 or 3).
197
+ if num_decodes > 0 and num_decode_tokens > num_decodes:
198
+ decode_query_len = num_decode_tokens // num_decodes
199
+ else:
200
+ decode_query_len = 1
201
+
202
+ seq_lens_decode = None
203
qo_indptr = None
204
+ qo_indptr_decode = None
205
if num_prefills > 0:
206
qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu
207
prefill_start = num_decodes
@@ -161,6 +210,21 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
210
)
211
qo_indptr = qo_indptr_prefill_cpu.to(self.device, non_blocking=True)
212
213
+ if num_decodes > 0:
214
+ seq_lens_decode = seq_lens[:num_decodes]
215
+ # block_table is per-request, indexed by num_decodes (not tokens)
216
+ qo_indptr_decode = common_attn_metadata.query_start_loc[: num_decodes + 1]
217
+ import hpc
218
+
219
+ hpc.assign_attention_decode_task(
220
+ seq_lens_decode,
221
+ self.task_map,
222
+ self.num_kv_heads,
223
+ decode_query_len,
224
+ new_kv_included=True,
225
+ min_process_len=self.hpc_dynamic_sched_attn_min_split_len,
226
+ )
227
+
228
return HpcAttnMetadata(
229
num_actual_tokens=num_actual_tokens,
230
num_decodes=num_decodes,
@@ -176,6 +240,9 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
240
hpc_prefill_q_scale=None,
241
hpc_decode_q_scale=None,
242
hpc_split_k_flag=None,
243
+ decode_query_len=decode_query_len,
244
+ qo_indptr_decode=qo_indptr_decode,
245
+ task_map=self.task_map,
246
)
247
248
@@ -192,6 +259,7 @@ class HpcAttentionBackend(AttentionBackend):
259
]
260
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
261
"auto",
262
+ "bfloat16",
263
"fp8_e4m3",
264
]
265
@@ -323,6 +391,13 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
391
self.supports_quant_query_input = False
392
self.splitk = True
393
394
+ import hpc
395
+
396
+ if self.use_fp8:
397
+ self._quant_type = hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR
398
+ else:
399
+ self._quant_type = None
400
+
401
def forward(
402
self,
403
layer: torch.nn.Module,
@@ -417,6 +492,7 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
492
block_table_prefill,
493
seq_lens_prefill,
494
max_seqlens,
495
+ quant_type=self._quant_type,
496
output=output_prefill,
497
)
498
else:
@@ -439,6 +515,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
515
q_decode = query[:num_decode_tokens]
516
output_decode = output[:num_decode_tokens]
517
518
+ mtp = attn_metadata.decode_query_len - 1
519
+
520
if self.use_fp8:
521
hpc.attention_decode_fp8(
522
q_decode,
@@ -449,8 +527,14 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
527
hpc_decode_q_scale,
528
k_scale,
529
v_scale,
530
+ mtp=mtp,
531
+ # MTP: split_flag from rope_norm is unavailable
532
+ # when using prefill-mode kernel; let HPC decide.
533
+ # splitk=(self.splitk if mtp == 0 else True),
534
new_kv_included=True,
535
+ quant_type=self._quant_type,
536
splitk=self.splitk,
537
+ task_map=attn_metadata.task_map,
538
split_flag=hpc_split_k_flag,
539
output=output_decode,
540
)
@@ -461,6 +545,7 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
545
kv_cache[:, 1],
546
block_table_decode,
547
num_seq_kvcache,
548
+ mtp=mtp,
549
output=output_decode,
550
new_kv_included=True,
551
splitk=self.splitk,