@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8cf7c4d8ad

[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>

Cheng Jiang committed Jun 30, 2026 at 18:17 UTC 8cf7c4d8ad602d73ff2ec72a101420d47163c136
10 files changed +978 -14
docs/design/attention_backends.md
+1
@@ -167,6 +167,7 @@ Priority is **1 = highest** (tried first).
167 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
168 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
169 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
170 +| `HPC_ATTN` | | fp16, bf16 | `auto`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
171 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
172 | `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
173 | `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/config/compilation.py
+1
@@ -759,6 +759,7 @@ class CompilationConfig:
759 "vllm::sparse_attn_indexer",
760 "vllm::rocm_aiter_sparse_attn_indexer",
761 "vllm::deepseek_v4_attention",
762 + "vllm::hpc_rope_norm_forward",
763 ]
764
765 def compute_hash(self) -> str:
vllm/model_executor/layers/attention/attention.py
+3 -1
@@ -458,6 +458,7 @@ class Attention(nn.Module, AttentionLayerBase):
458 # shape does not match the query shape, so we optionally let the model
459 # definition specify the output tensor shape.
460 output_shape: torch.Size | None = None,
461 + output_dtype: torch.dtype | None = None,
462 ) -> torch.Tensor:
463 """
464 The KV cache is stored inside this class and is accessed via
@@ -472,7 +473,8 @@ class Attention(nn.Module, AttentionLayerBase):
473 torch.ops.vllm.maybe_calc_kv_scales(
474 query, key, value, _encode_layer_name(self.layer_name)
475 )
475 - output_dtype = query.dtype
476 + if output_dtype is None:
477 + output_dtype = query.dtype
478 if self.query_quant is not None:
479 # quantizing with a simple torch operation enables
480 # torch.compile to fuse this into previous ops
vllm/model_executor/layers/hpc/__init__.py new
+11
@@ -0,0 +1,11 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.model_executor.layers.hpc.hpc_module import HpcModule
5 +from vllm.model_executor.layers.hpc.rope_norm import HpcRopeNorm, QkNormPolicy
6 +
7 +__all__ = [
8 + "HpcModule",
9 + "HpcRopeNorm",
10 + "QkNormPolicy",
11 +]
vllm/model_executor/layers/hpc/hpc_module.py new
+18
@@ -0,0 +1,18 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch.nn as nn
4 +
5 +
6 +class HpcModule(nn.Module):
7 + def __init__(self):
8 + super().__init__()
9 +
10 + @classmethod
11 + def support(cls, *args, **kwargs):
12 + return True
13 +
14 + def process_weights_after_loading(self, model):
15 + pass
16 +
17 + def forward(self, *args, **kwargs):
18 + pass
vllm/model_executor/layers/hpc/rope_norm.py new
+408
@@ -0,0 +1,408 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant).
4 +
5 +Decoupled from HpcAttentionImpl; extra params are passed via layer attrs.
6 +"""
7 +
8 +from __future__ import annotations
9 +
10 +from enum import IntEnum
11 +from typing import Any
12 +
13 +import torch
14 +
15 +from vllm.config import get_current_vllm_config_or_none
16 +from vllm.forward_context import ForwardContext, get_forward_context
17 +from vllm.logger import init_logger
18 +from vllm.model_executor.custom_op import CustomOp
19 +from vllm.model_executor.layers.hpc.hpc_module import HpcModule
20 +from vllm.utils.torch_utils import direct_register_custom_op
21 +from vllm.v1.attention.backends.hpc_attn import HpcAttnMetadata
22 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
23 +
24 +logger = init_logger(__name__)
25 +
26 +_hpc_rope_norm_instances: dict[str, HpcRopeNorm] = {}
27 +
28 +
29 +class QkNormPolicy(IntEnum):
30 + """Order of QK-RMSNorm relative to RoPE in the fused HPC rope_norm kernel.
31 +
32 + The values are part of the HPC kernel ABI (passed through as ints), so they
33 + must stay in sync with the kernel's expectations.
34 + """
35 +
36 + # No QK-Norm: apply RoPE only.
37 + NONE = 0
38 + # Apply RoPE first, then QK-RMSNorm.
39 + ROPE_THEN_NORM = 1
40 + # Apply QK-RMSNorm first, then RoPE (e.g. HunYuan V3).
41 + NORM_THEN_ROPE = 2
42 +
43 +
44 +def hpc_rope_norm_forward(
45 + qkv: torch.Tensor,
46 + output: torch.Tensor,
47 + layer_name: str,
48 +) -> None:
49 + """Top-level custom op: RoPE + QK-Norm + KV-Cache-Write + FP8 Q quant.
50 +
51 + Fully opaque to torch.compile (dynamo).
52 + """
53 + forward_context: ForwardContext = get_forward_context()
54 + attn_metadata: Any = forward_context.attn_metadata
55 + if isinstance(attn_metadata, dict):
56 + attn_metadata = attn_metadata[layer_name]
57 +
58 + if attn_metadata is None:
59 + output.zero_()
60 + return
61 +
62 + attn_layer = forward_context.no_compile_layers[layer_name]
63 + # bind_kv_cache stores the per-layer KV cache as a single 5D tensor
64 + # (num_blocks, 2, block_size, num_kv_heads, head_size), so use it directly.
65 + kv_cache = attn_layer.kv_cache
66 +
67 + if kv_cache.numel() == 0:
68 + output.zero_()
69 + return
70 +
71 + assert kv_cache.dim() == 5, (
72 + f"Expected kv_cache to have 5 dims, got {tuple(kv_cache.shape)}"
73 + )
74 +
75 + rope_norm = _hpc_rope_norm_instances[layer_name]
76 + rope_norm._forward_impl(qkv, kv_cache, attn_metadata, attn_layer, output)
77 +
78 +
79 +def hpc_rope_norm_forward_fake(
80 + qkv: torch.Tensor,
81 + output: torch.Tensor,
82 + layer_name: str,
83 +) -> None:
84 + """Fake impl for torch.compile trace; output is a mutated arg."""
85 + return
86 +
87 +
88 +direct_register_custom_op(
89 + op_name="hpc_rope_norm_forward",
90 + op_func=hpc_rope_norm_forward,
91 + mutates_args=["output"],
92 + fake_impl=hpc_rope_norm_forward_fake,
93 +)
94 +
95 +
96 +@CustomOp.register("hpc_rope_norm")
97 +class HpcRopeNorm(CustomOp, HpcModule):
98 + """HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant).
99 +
100 + Registered as a sub-module in model layers (e.g. HunYuanAttention).
101 + Norm weights are extracted from fallback norm modules via
102 + process_weights_after_loading() after all weights are loaded.
103 +
104 + forward() is dispatched by CustomOp framework:
105 + - In compiled mode: forward_cuda() calls torch.ops.vllm.hpc_rope_norm_forward
106 + as a splitting point — internal Python control flow is opaque
107 + to torch.compile and not captured by CUDA Graph.
108 + - In eager/native mode: forward_native() falls back to forward_cuda().
109 + """
110 +
111 + def __init__(
112 + self,
113 + num_heads: int,
114 + num_kv_heads: int,
115 + head_dim: int,
116 + cos_sin_cache: torch.Tensor,
117 + use_qk_norm: bool,
118 + fallback_qnorm: torch.nn.Module | None,
119 + fallback_knorm: torch.nn.Module | None,
120 + kv_cache_dtype: str,
121 + layer_name: str,
122 + qk_norm_policy: QkNormPolicy = QkNormPolicy.ROPE_THEN_NORM,
123 + ) -> None:
124 + super().__init__()
125 + self.num_heads = num_heads
126 + self.num_kv_heads = num_kv_heads
127 + self.head_dim = head_dim
128 +
129 + self.use_qk_norm = use_qk_norm
130 +
131 + self.q_size = num_heads * head_dim
132 + self.kv_size = num_kv_heads * head_dim
133 +
134 + # Register as a non-persistent buffer so it participates in sleep
135 + # level-2 save/restore (CuMemAllocator) but is excluded from the
136 + # checkpoint state_dict.
137 + self.register_buffer("cos_sin_cache", cos_sin_cache.float(), persistent=False)
138 +
139 + self.fallback_qnorm = fallback_qnorm
140 + self.fallback_knorm = fallback_knorm
141 +
142 + self.head_per_group = num_heads // num_kv_heads
143 +
144 + # Pre-allocate norm weight tensors as Parameters so they are tracked by
145 + # CuMemAllocator (for sleep/wake_up) and have stable addresses for CUDA
146 + # Graph replay. process_weights_after_loading() updates them inplace via
147 + # copy_() so refit does not invalidate captured graph tensor pointers.
148 + # Shape is [head_dim] to match the HPC kernel's q/k_norm_weight layout.
149 + if use_qk_norm and fallback_qnorm is not None:
150 + self.qnorm_weight: torch.nn.Parameter | None = torch.nn.Parameter(
151 + torch.empty(head_dim, dtype=torch.float32),
152 + requires_grad=False,
153 + )
154 + else:
155 + self.qnorm_weight = None
156 + if use_qk_norm and fallback_knorm is not None:
157 + self.knorm_weight: torch.nn.Parameter | None = torch.nn.Parameter(
158 + torch.empty(head_dim, dtype=torch.float32),
159 + requires_grad=False,
160 + )
161 + else:
162 + self.knorm_weight = None
163 +
164 + self.use_fp8 = "fp8" in kv_cache_dtype
165 + # The RMSNorm/RoPE ordering is model dependent (e.g. HunYuan V3 applies
166 + # QK-Norm before RoPE -> NORM_THEN_ROPE), so it is supplied by the
167 + # caller. When QK-Norm is disabled the policy is forced to NONE.
168 + self.qk_norm_policy = qk_norm_policy if use_qk_norm else QkNormPolicy.NONE
169 +
170 + # Register layer_name + add self to the global instance registry so the
171 + # module-level custom op (hpc_rope_norm_forward) can route back here.
172 + self.layer_name: str | None = None
173 + self.register_layer_name(layer_name)
174 +
175 + @classmethod
176 + def support(
177 + cls,
178 + num_heads: int,
179 + num_kv_heads: int,
180 + head_dim: int,
181 + kv_cache_dtype: str,
182 + ) -> bool:
183 + """Check whether HpcRopeNorm is supported for the given config."""
184 + # HpcRopeNorm is only enabled together with the HPC attention backend.
185 + vllm_config = get_current_vllm_config_or_none()
186 + if (
187 + vllm_config is None
188 + or vllm_config.attention_config.backend != AttentionBackendEnum.HPC_ATTN
189 + ):
190 + return False
191 +
192 + if kv_cache_dtype not in ("fp8_e4m3", "auto"):
193 + logger.warning_once(
194 + f"hpc rope_norm not support kv_cache_dtype:{kv_cache_dtype}, "
195 + "only support fp8_e4m3, bfloat16"
196 + )
197 + return False
198 +
199 + if head_dim not in (128,):
200 + logger.warning_once("hpc rope_norm only support head_dim == 128.")
201 + return False
202 +
203 + head_per_group = num_heads // num_kv_heads
204 + if head_per_group not in (4, 8):
205 + logger.warning_once("hpc rope_norm only support head_per_group in [4, 8].")
206 + return False
207 +
208 + logger.info_once("enable hpc rope_norm")
209 + return True
210 +
211 + def process_weights_after_loading(self, model: torch.nn.Module = None) -> None:
212 + """Copy norm weights (float32) from fallback norm modules inplace.
213 +
214 + Uses copy_() to preserve tensor addresses for CUDA Graph / refit
215 + compatibility. Called by the model's load_weights() after all weights
216 + are loaded (and generically from the model loader for DummyModelLoader
217 + / sleep-wake_up reload paths).
218 + """
219 + if self.use_qk_norm:
220 + if self.fallback_qnorm is not None and self.qnorm_weight is not None:
221 + self.qnorm_weight.data.copy_(self.fallback_qnorm.weight.data.float())
222 + if self.fallback_knorm is not None and self.knorm_weight is not None:
223 + self.knorm_weight.data.copy_(self.fallback_knorm.weight.data.float())
224 +
225 + def register_layer_name(self, layer_name: str) -> None:
226 + """Register layer_name and add self to the global registry.
227 +
228 + The global registry is needed because the bottom-level torch op
229 + (hpc_rope_norm_forward) is a module-level function and needs to
230 + route back to the correct instance via layer_name.
231 + """
232 + self.layer_name = layer_name
233 + _hpc_rope_norm_instances[layer_name] = self
234 + logger.debug(
235 + "[rope_norm] registered HpcRopeNorm for layer: %s",
236 + layer_name,
237 + )
238 +
239 + def forward_native(
240 + self,
241 + qkv: torch.Tensor,
242 + layer_name: str,
243 + ) -> torch.Tensor:
244 + """Native fallback path: delegates to forward_cuda().
245 +
246 + For now, the default native path will use CUDA backend path.
247 + Other platforms may override via OOT registration.
248 + """
249 + return self.forward_cuda(qkv, layer_name)
250 +
251 + def forward_cuda(
252 + self,
253 + qkv: torch.Tensor,
254 + layer_name: str,
255 + ) -> torch.Tensor:
256 + """CUDA path: invoke the torch custom op as a compile splitting point."""
257 + num_tokens = qkv.shape[0]
258 + output = torch.empty(
259 + (num_tokens, self.num_heads, self.head_dim),
260 + dtype=torch.float8_e4m3fn if self.use_fp8 else qkv.dtype,
261 + device=qkv.device,
262 + )
263 +
264 + torch.ops.vllm.hpc_rope_norm_forward(qkv, output, layer_name)
265 + return output
266 +
267 + def _forward_impl(
268 + self,
269 + qkv: torch.Tensor,
270 + kv_cache: torch.Tensor,
271 + attn_metadata: HpcAttnMetadata,
272 + attn_layer: torch.nn.Module,
273 + output: torch.Tensor,
274 + ) -> None:
275 + """Actual forward logic called by the custom op.
276 +
277 + Writes processed q into *output* and attaches extra params
278 + (e.g. FP8 scales) to *attn_layer* as attributes.
279 + """
280 + import hpc
281 +
282 + num_actual_tokens = attn_metadata.num_actual_tokens
283 + num_prefill_reqs = attn_metadata.num_prefills
284 + num_decode_reqs = attn_metadata.num_decodes
285 + num_decode_tokens = attn_metadata.num_decode_tokens
286 +
287 + qkv = qkv[:num_actual_tokens]
288 +
289 + num_prefill_tokens = num_actual_tokens - num_decode_tokens
290 +
291 + # KV cache for the FP8 path is stored as uint8; view it as fp8 so the
292 + # rope_norm_store_kv_fp8 kernel can write quantized K/V in-place.
293 + if self.use_fp8:
294 + kv_cache = kv_cache.view(torch.float8_e4m3fn)
295 +
296 + # Per-tensor K/V scales (shape [1]) used by the FP8 kernel.
297 + k_scale = attn_layer._k_scale.reshape(1)
298 + v_scale = attn_layer._v_scale.reshape(1)
299 +
300 + q_norm_weight = (
301 + self.qnorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None
302 + )
303 + k_norm_weight = (
304 + self.knorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None
305 + )
306 +
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 +
313 + # --- Prefill ---
314 + if num_prefill_reqs > 0:
315 + seq_lens_prefill = attn_metadata.seq_lens[num_decode_reqs:]
316 + cu_seqlens_prefill = attn_metadata.qo_indptr
317 + max_seqlens = attn_metadata.max_query_len
318 + block_table_prefill = attn_metadata.block_table_tensor[num_decode_reqs:]
319 + qkv_prefill = qkv[num_decode_tokens:]
320 + out_q_prefill = output[
321 + num_decode_tokens : num_decode_tokens + num_prefill_tokens
322 + ]
323 +
324 + if self.use_fp8:
325 + _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8(
326 + key_cache=kv_cache[:, 0],
327 + value_cache=kv_cache[:, 1],
328 + qkv=qkv_prefill,
329 + cos_sin=self.cos_sin_cache,
330 + num_seqlen_per_req=seq_lens_prefill,
331 + q_index=cu_seqlens_prefill,
332 + kvcache_indices=block_table_prefill,
333 + is_prefill=True,
334 + k_scale=k_scale,
335 + v_scale=v_scale,
336 + quant_policy=QUANT_POLICY_DQSKV,
337 + max_seqlens=max_seqlens,
338 + q_norm_weight=q_norm_weight,
339 + k_norm_weight=k_norm_weight,
340 + qk_norm_policy=self.qk_norm_policy,
341 + out_q=out_q_prefill,
342 + )
343 + attn_metadata.hpc_prefill_q_scale = q_scale
344 + else:
345 + hpc.rope_norm_store_kv(
346 + kv_cache[:, 0],
347 + kv_cache[:, 1],
348 + qkv_prefill,
349 + self.cos_sin_cache,
350 + seq_lens_prefill,
351 + cu_seqlens_prefill,
352 + block_table_prefill,
353 + True, # is_prefill
354 + q_norm_weight=q_norm_weight,
355 + k_norm_weight=k_norm_weight,
356 + out_q=out_q_prefill,
357 + qk_norm_policy=self.qk_norm_policy,
358 + )
359 +
360 + # --- Decode ---
361 + if num_decode_reqs > 0:
362 + num_seq_kvcache = attn_metadata.seq_lens[:num_decode_reqs]
363 + block_table_decode = attn_metadata.block_table_tensor[:num_decode_reqs]
364 + qkv_decode = qkv[:num_decode_tokens]
365 + # Single-token decode: q_index is the per-request prefix sum
366 + # [0, 1, ..., num_decode_reqs].
367 + qo_indptr_decode = torch.arange(
368 + num_decode_reqs + 1, dtype=torch.int32, device=qkv.device
369 + )
370 + out_q_decode = output[:num_decode_tokens]
371 +
372 + if self.use_fp8:
373 + _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8(
374 + key_cache=kv_cache[:, 0],
375 + value_cache=kv_cache[:, 1],
376 + qkv=qkv_decode,
377 + cos_sin=self.cos_sin_cache,
378 + num_seqlen_per_req=num_seq_kvcache,
379 + q_index=qo_indptr_decode,
380 + kvcache_indices=block_table_decode,
381 + is_prefill=False,
382 + k_scale=k_scale,
383 + v_scale=v_scale,
384 + quant_policy=QUANT_POLICY_DQSKV,
385 + max_seqlens=1,
386 + q_norm_weight=q_norm_weight,
387 + k_norm_weight=k_norm_weight,
388 + qk_norm_policy=self.qk_norm_policy,
389 + out_q=out_q_decode,
390 + )
391 + attn_metadata.hpc_decode_q_scale = q_scale
392 + if split_k_flag is not None:
393 + attn_metadata.hpc_split_k_flag = split_k_flag
394 + else:
395 + hpc.rope_norm_store_kv(
396 + kv_cache[:, 0],
397 + kv_cache[:, 1],
398 + qkv_decode,
399 + self.cos_sin_cache,
400 + num_seq_kvcache,
401 + qo_indptr_decode,
402 + block_table_decode,
403 + False, # is_prefill
404 + q_norm_weight=q_norm_weight,
405 + k_norm_weight=k_norm_weight,
406 + out_q=out_q_decode,
407 + qk_norm_policy=self.qk_norm_policy,
408 + )
vllm/model_executor/model_loader/utils.py
+10
@@ -20,6 +20,7 @@ from vllm.model_executor.layers.attention import (
20 MLAAttention,
21 MMEncoderAttention,
22 )
23 +from vllm.model_executor.layers.hpc import HpcModule
24 from vllm.model_executor.layers.quantization.base_config import (
25 QuantizationConfig,
26 QuantizeMethodBase,
@@ -125,6 +126,15 @@ def process_weights_after_loading(
126 with device_loading_context(module, target_device):
127 module.process_weights_after_loading(model_config.dtype)
128
129 + # Process HPC modules (HpcRopeNorm, etc.) that rely on
130 + # process_weights_after_loading being called from the model's
131 + # load_weights(). When using DummyModelLoader (e.g. profiling or
132 + # sleep/wake_up reload), the model's load_weights() is not called, so we
133 + # must handle HPC modules here generically.
134 + for _, module in model.named_modules():
135 + if isinstance(module, HpcModule):
136 + module.process_weights_after_loading(model)
137 +
138 # Needed for torchao model reloading via model.reload_weights
139 # @kylesayrs @jerryzh168 this can be removed if callers move to `reload_weights`
140 if model_config.quantization == "torchao":
vllm/model_executor/models/hy_v3.py
+51 -13
@@ -48,6 +48,7 @@ from vllm.model_executor.layers.fused_moe import (
48 GateLinear,
49 fused_moe_make_expert_params_mapping,
50 )
51 +from vllm.model_executor.layers.hpc import HpcRopeNorm, QkNormPolicy
52 from vllm.model_executor.layers.layernorm import RMSNorm
53 from vllm.model_executor.layers.linear import (
54 MergedColumnParallelLinear,
@@ -234,6 +235,7 @@ class HYV3Attention(nn.Module):
235 dual_chunk_attention_config: dict[str, Any] | None = None,
236 ) -> None:
237 super().__init__()
238 + self.dtype = torch.get_default_dtype()
239 self.hidden_size = hidden_size
240 tp_size = get_tensor_model_parallel_world_size()
241 self.total_num_heads = num_heads
@@ -276,11 +278,18 @@ class HYV3Attention(nn.Module):
278 quant_config=quant_config,
279 prefix=f"{prefix}.o_proj",
280 )
281 + # When the HPC fused RoPE+QK-Norm path is enabled, the RoPE cos/sin
282 + # cache must be float32 to match the HPC kernel's expectations.
283 + kv_cache_dtype = cache_config.cache_dtype if cache_config else "auto"
284 + rope_support = HpcRopeNorm.support(
285 + self.num_heads, self.num_kv_heads, self.head_dim, kv_cache_dtype
286 + )
287 self.rotary_emb = get_rope(
288 self.head_dim,
289 max_position=max_position_embeddings,
290 rope_parameters=rope_parameters,
291 is_neox_style=True,
292 + dtype=torch.float32 if rope_support else torch.get_default_dtype(),
293 )
294 self.attn = Attention(
295 self.num_heads,
@@ -295,6 +304,27 @@ class HYV3Attention(nn.Module):
304 self.q_norm = RMSNorm(self.head_dim, rms_norm_eps)
305 self.k_norm = RMSNorm(self.head_dim, rms_norm_eps)
306
307 + # HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant).
308 + # HunYuan V3 applies QK-Norm *before* RoPE, so NORM_THEN_ROPE.
309 + self.hpc_rope_norm: HpcRopeNorm | None = None
310 + if rope_support:
311 + self.hpc_rope_norm = HpcRopeNorm(
312 + num_heads=self.num_heads,
313 + num_kv_heads=self.num_kv_heads,
314 + head_dim=self.head_dim,
315 + cos_sin_cache=self.rotary_emb.cos_sin_cache,
316 + use_qk_norm=self.use_qk_norm,
317 + fallback_qnorm=self.q_norm if self.use_qk_norm else None,
318 + fallback_knorm=self.k_norm if self.use_qk_norm else None,
319 + kv_cache_dtype=kv_cache_dtype,
320 + layer_name=self.attn.layer_name,
321 + qk_norm_policy=QkNormPolicy.NORM_THEN_ROPE,
322 + )
323 + # FP8 Q is produced by HpcRopeNorm, so the attention layer must not
324 + # re-quantize the query.
325 + if self.hpc_rope_norm.use_fp8 and hasattr(self.attn, "query_quant"):
326 + self.attn.query_quant = None
327 +
328 def forward(
329 self,
330 positions: torch.Tensor,
@@ -303,20 +333,28 @@ class HYV3Attention(nn.Module):
333 qkv, _ = self.qkv_proj(hidden_states)
334 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
335 output_shape = None
306 - if self.use_qk_norm:
307 - q_by_head = q.view(
308 - *q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim
309 - )
310 - q_by_head = self.q_norm(q_by_head)
311 - q = q_by_head.view(q.shape)
336 + if self.hpc_rope_norm is not None:
337 + # HPC handles QK-Norm + RoPE + KV-cache write (+ optional FP8 Q
338 + # quant) internally and returns the processed query. K/V are
339 + # written into the paged cache by the fused op.
340 + q = self.hpc_rope_norm(qkv, self.attn.layer_name)
341 + q = q.view(-1, self.num_heads * self.head_dim)
342 + attn_output = self.attn(q, k, v, output_shape, self.dtype)
343 + else:
344 + if self.use_qk_norm:
345 + q_by_head = q.view(
346 + *q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim
347 + )
348 + q_by_head = self.q_norm(q_by_head)
349 + q = q_by_head.view(q.shape)
350
313 - k_by_head = k.view(
314 - *k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim
315 - )
316 - k_by_head = self.k_norm(k_by_head)
317 - k = k_by_head.view(k.shape)
318 - q, k = self.rotary_emb(positions, q, k)
319 - attn_output = self.attn(q, k, v, output_shape)
351 + k_by_head = k.view(
352 + *k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim
353 + )
354 + k_by_head = self.k_norm(k_by_head)
355 + k = k_by_head.view(k.shape)
356 + q, k = self.rotary_emb(positions, q, k)
357 + attn_output = self.attn(q, k, v, output_shape)
358 attn_output = attn_output.view(q.shape[0], -1)
359 output, _ = self.o_proj(attn_output)
360 return output
vllm/v1/attention/backends/hpc_attn.py new
+469
@@ -0,0 +1,469 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""HPC Attention Backend.
4 +
5 +Pure attention (prefill + decode), without RoPE or RMSNorm.
6 +Independent metadata / builder; KV cache layout is NHD:
7 +(num_blocks, 2, block_size, num_kv_heads, head_size).
8 +"""
9 +
10 +import importlib.util
11 +from dataclasses import dataclass
12 +from typing import ClassVar
13 +
14 +import torch
15 +from typing_extensions import override
16 +
17 +from vllm.config import VllmConfig
18 +from vllm.config.cache import CacheDType
19 +from vllm.logger import init_logger
20 +from vllm.platforms import current_platform
21 +from vllm.platforms.interface import DeviceCapability
22 +from vllm.v1.attention.backend import (
23 + AttentionBackend,
24 + AttentionCGSupport,
25 + AttentionImpl,
26 + AttentionMetadata,
27 + AttentionMetadataBuilder,
28 + AttentionType,
29 + CommonAttentionMetadata,
30 + MultipleOf,
31 +)
32 +from vllm.v1.attention.backends.utils import (
33 + KVCacheLayoutType,
34 + get_per_layer_parameters,
35 + infer_global_hyperparameters,
36 + split_decodes_and_prefills,
37 +)
38 +from vllm.v1.kv_cache_interface import AttentionSpec
39 +
40 +logger = init_logger(__name__)
41 +
42 +FP8_DTYPE = current_platform.fp8_dtype()
43 +
44 +
45 +def _get_fp8_dtype_for_kv_cache(kv_cache_dtype: str) -> torch.dtype:
46 + """Return the torch FP8 dtype for the given kv_cache_dtype string."""
47 + if kv_cache_dtype in ("fp8", "fp8_e4m3"):
48 + return torch.float8_e4m3fn
49 + elif kv_cache_dtype == "fp8_e5m2":
50 + return torch.float8_e5m2
51 + else:
52 + raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")
53 +
54 +
55 +@dataclass
56 +class HpcAttnMetadata(AttentionMetadata):
57 + """Metadata required by the HPC attention kernel."""
58 +
59 + num_actual_tokens: int
60 + num_decodes: int
61 + num_decode_tokens: int
62 + num_prefills: int
63 + num_prefill_tokens: int
64 + max_query_len: int
65 +
66 + slot_mapping: torch.Tensor
67 + """Slot mapping for KV cache writes. shape = [num_actual_tokens]"""
68 +
69 + seq_lens: torch.Tensor
70 + """KV cache length per request. shape = [batch_size]"""
71 +
72 + block_table_tensor: torch.Tensor
73 + """Paged KV-cache block table.
74 + shape = [batch_size, max_num_blocks_per_seq]"""
75 +
76 + qo_indptr: torch.Tensor | None = None
77 + """Cumulative query lengths for prefill requests (GPU tensor).
78 + shape = [num_prefills + 1]. None when num_prefills == 0."""
79 +
80 + # --- HPC RopeNorm pass-through fields ---
81 + # Set by HpcRopeNorm._forward_impl(); consumed & reset by
82 + # HpcAttentionImpl.forward(). Defaults are safe for the standard
83 + # (non-RopeNorm) path and for profiling runs (attn_metadata=None).
84 + hpc_kv_written: bool = False
85 + """True when HpcRopeNorm already wrote KV cache."""
86 + hpc_prefill_q_scale: torch.Tensor | None = None
87 + """FP8 per-token-per-head Q scale for prefill (from RopeNorm)."""
88 + hpc_decode_q_scale: torch.Tensor | None = None
89 + """FP8 per-token-per-head Q scale for decode (from RopeNorm)."""
90 + hpc_split_k_flag: torch.Tensor | None = None
91 + """Split-K flag tensor for FP8 decode (from RopeNorm)."""
92 +
93 +
94 +class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
95 + """Build HpcAttnMetadata from CommonAttentionMetadata."""
96 +
97 + _cudagraph_support = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
98 + reorder_batch_threshold: int = 1
99 +
100 + def __init__(
101 + self,
102 + kv_cache_spec: AttentionSpec,
103 + layer_names: list[str],
104 + vllm_config: VllmConfig,
105 + device: torch.device,
106 + ):
107 + 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
110 +
111 + self.num_qo_heads = self.model_config.get_num_attention_heads(
112 + vllm_config.parallel_config
113 + )
114 + 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
117 +
118 + self.cache_dtype = self.cache_config.cache_dtype
119 +
120 + self.global_hyperparameters = infer_global_hyperparameters(
121 + get_per_layer_parameters(vllm_config, layer_names, HpcAttentionImpl)
122 + )
123 +
124 + @override # type: ignore[misc]
125 + @classmethod
126 + def get_cudagraph_support(
127 + cls: type["HpcAttnMetadataBuilder"],
128 + vllm_config: VllmConfig,
129 + kv_cache_spec: AttentionSpec,
130 + ) -> AttentionCGSupport:
131 + return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
132 +
133 + def build(
134 + self,
135 + common_prefix_len: int,
136 + common_attn_metadata: CommonAttentionMetadata,
137 + fast_build: bool = False,
138 + ) -> HpcAttnMetadata:
139 + """Build HpcAttnMetadata from CommonAttentionMetadata."""
140 + num_actual_tokens = common_attn_metadata.num_actual_tokens
141 +
142 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
143 + split_decodes_and_prefills(
144 + common_attn_metadata,
145 + decode_threshold=self.reorder_batch_threshold,
146 + require_uniform=False,
147 + )
148 + )
149 +
150 + seq_lens = common_attn_metadata.seq_lens
151 + block_table_tensor = common_attn_metadata.block_table_tensor
152 + slot_mapping = common_attn_metadata.slot_mapping
153 + max_query_len = common_attn_metadata.max_query_len
154 +
155 + qo_indptr = None
156 + if num_prefills > 0:
157 + qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu
158 + prefill_start = num_decodes
159 + qo_indptr_prefill_cpu = (
160 + qo_indptr_cpu[prefill_start:] - qo_indptr_cpu[prefill_start]
161 + )
162 + qo_indptr = qo_indptr_prefill_cpu.to(self.device, non_blocking=True)
163 +
164 + return HpcAttnMetadata(
165 + num_actual_tokens=num_actual_tokens,
166 + num_decodes=num_decodes,
167 + num_decode_tokens=num_decode_tokens,
168 + num_prefills=num_prefills,
169 + num_prefill_tokens=num_prefill_tokens,
170 + max_query_len=max_query_len,
171 + slot_mapping=slot_mapping,
172 + seq_lens=seq_lens,
173 + block_table_tensor=block_table_tensor,
174 + qo_indptr=qo_indptr,
175 + hpc_kv_written=True,
176 + hpc_prefill_q_scale=None,
177 + hpc_decode_q_scale=None,
178 + hpc_split_k_flag=None,
179 + )
180 +
181 +
182 +class HpcAttentionBackend(AttentionBackend):
183 + """HPC attention backend (pure attention, no RoPE/Norm).
184 +
185 + KV cache layout: NHD (num_blocks, 2, block_size, num_kv_heads, head_size).
186 + """
187 +
188 + accept_output_buffer: bool = True
189 + supported_dtypes: ClassVar[list[torch.dtype]] = [
190 + torch.float16,
191 + torch.bfloat16,
192 + ]
193 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
194 + "auto",
195 + "fp8_e4m3",
196 + ]
197 +
198 + # Avoid attention abstracted method call cache insert
199 + forward_includes_kv_cache_update: bool = True
200 +
201 + @staticmethod
202 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
203 + return [64]
204 +
205 + @staticmethod
206 + def get_name() -> str:
207 + return "HPC_ATTN"
208 +
209 + @staticmethod
210 + def get_impl_cls() -> type["HpcAttentionImpl"]:
211 + return HpcAttentionImpl
212 +
213 + @staticmethod
214 + def get_builder_cls() -> type["HpcAttnMetadataBuilder"]:
215 + return HpcAttnMetadataBuilder
216 +
217 + @staticmethod
218 + def get_kv_cache_shape(
219 + num_blocks: int,
220 + block_size: int,
221 + num_kv_heads: int,
222 + head_size: int,
223 + cache_dtype_str: str = "auto",
224 + ) -> tuple[int, ...]:
225 + return (num_blocks, 2, block_size, num_kv_heads, head_size)
226 +
227 + @staticmethod
228 + def get_kv_cache_stride_order(
229 + include_num_layers_dimension: bool = False,
230 + ) -> tuple[int, ...]:
231 + if include_num_layers_dimension:
232 + return (1, 0, 2, 3, 4, 5)
233 + return (0, 1, 2, 3, 4)
234 +
235 + @classmethod
236 + def get_supported_head_sizes(cls) -> list[int]:
237 + return [128]
238 +
239 + @classmethod
240 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
241 + return capability >= DeviceCapability(9, 0)
242 +
243 + @classmethod
244 + def supports_kv_cache_dtype(cls, kv_cache_dtype: "CacheDType | None") -> bool:
245 + if kv_cache_dtype is None:
246 + return True
247 + return kv_cache_dtype in cls.supported_kv_cache_dtypes
248 +
249 + @classmethod
250 + def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None:
251 + return "NHD"
252 +
253 +
254 +class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
255 + """HPC pure attention implementation (no RoPE/Norm).
256 +
257 + Constraints:
258 + - head_dim == 128
259 + - num_heads // num_kv_heads in {4, 8}
260 + - kv_cache_dtype in {"auto", "fp8_e4m3"}
261 + """
262 +
263 + def __init__(
264 + self,
265 + num_heads: int,
266 + head_size: int,
267 + scale: float,
268 + num_kv_heads: int,
269 + alibi_slopes: list[float] | None = None,
270 + sliding_window: int | None = None,
271 + kv_cache_dtype: str = "auto",
272 + logits_soft_cap: float | None = None,
273 + attn_type: str = AttentionType.DECODER,
274 + kv_sharing_target_layer_name: str | None = None,
275 + ) -> None:
276 + if importlib.util.find_spec("hpc") is None:
277 + raise ImportError(
278 + "HPC attention requires the hpc module to be installed. "
279 + "Please install it from https://github.com/Tencent/hpc-ops"
280 + )
281 + if attn_type != AttentionType.DECODER:
282 + raise NotImplementedError("HPC attention only supports decoder attention")
283 + if alibi_slopes is not None:
284 + raise NotImplementedError("HPC attention does not support ALiBi")
285 + if logits_soft_cap is not None:
286 + raise NotImplementedError("HPC attention does not support logits_soft_cap")
287 +
288 + if head_size != 128:
289 + raise ValueError(
290 + f"HPC attention only supports head_dim=128, got {head_size}"
291 + )
292 +
293 + num_queries_per_kv = num_heads // num_kv_heads
294 + if num_queries_per_kv not in (4, 8):
295 + raise ValueError(
296 + f"HPC attention only supports head_per_group in {{4, 8}}, "
297 + f"got {num_queries_per_kv} "
298 + f"(num_heads={num_heads}, num_kv_heads={num_kv_heads})"
299 + )
300 +
301 + if kv_cache_dtype not in ("auto", "fp8_e4m3"):
302 + raise ValueError(
303 + f"HPC attention only supports kv_cache_dtype 'auto' or "
304 + f"'fp8_e4m3', got '{kv_cache_dtype}'"
305 + )
306 +
307 + self.num_heads = num_heads
308 + self.head_size = head_size
309 + self.scale = float(scale)
310 + self.num_kv_heads = num_kv_heads
311 + self.kv_cache_dtype = kv_cache_dtype
312 + self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
313 +
314 + self.num_queries_per_kv = num_queries_per_kv
315 +
316 + if sliding_window is None:
317 + self.sliding_window = (-1, -1)
318 + else:
319 + self.sliding_window = (sliding_window - 1, 0)
320 +
321 + self.use_fp8 = kv_cache_dtype == "fp8_e4m3"
322 +
323 + self.supports_quant_query_input = False
324 + self.splitk = True
325 +
326 + def forward(
327 + self,
328 + layer: torch.nn.Module,
329 + query: torch.Tensor,
330 + key: torch.Tensor,
331 + value: torch.Tensor,
332 + kv_cache: torch.Tensor,
333 + attn_metadata: HpcAttnMetadata | None,
334 + output: torch.Tensor | None = None,
335 + output_scale: torch.Tensor | None = None,
336 + output_block_scale: torch.Tensor | None = None,
337 + ) -> torch.Tensor:
338 + """HPC attention forward (standard vLLM backend interface).
339 +
340 + Two modes:
341 + 1. Standard: upstream handles RoPE/Norm; this backend writes KV + attn.
342 + 2. HpcRopeNorm: fused op already did RoPE/Norm/KV-Write/Q-Quant;
343 + extra params passed via attn_metadata.hpc_* fields.
344 + """
345 + import hpc
346 +
347 + assert output is not None, "Output tensor must be provided."
348 + assert output_scale is None, "HPC attention does not support fused output quant"
349 + assert output_block_scale is None
350 +
351 + if attn_metadata is None:
352 + return output.fill_(0)
353 +
354 + hpc_kv_written = attn_metadata.hpc_kv_written
355 + hpc_prefill_q_scale = attn_metadata.hpc_prefill_q_scale
356 + hpc_decode_q_scale = attn_metadata.hpc_decode_q_scale
357 + hpc_split_k_flag = attn_metadata.hpc_split_k_flag
358 +
359 + num_actual_tokens = attn_metadata.num_actual_tokens
360 + num_prefill_reqs = attn_metadata.num_prefills
361 + num_decode_reqs = attn_metadata.num_decodes
362 + num_decode_tokens = attn_metadata.num_decode_tokens
363 +
364 + # Write KV cache if not already done by HpcRopeNorm.
365 + if self.kv_sharing_target_layer_name is None and not hpc_kv_written:
366 + torch.ops._C_cache_ops.reshape_and_cache_flash(
367 + key,
368 + value,
369 + kv_cache[:, 0],
370 + kv_cache[:, 1],
371 + attn_metadata.slot_mapping,
372 + self.kv_cache_dtype,
373 + layer._k_scale,
374 + layer._v_scale,
375 + )
376 +
377 + if self.use_fp8:
378 + torch_dtype = _get_fp8_dtype_for_kv_cache(self.kv_cache_dtype)
379 + kv_cache = kv_cache.view(torch_dtype)
380 +
381 + if self.use_fp8:
382 + if not hpc_kv_written:
383 + raise RuntimeError(
384 + "HpcAttentionImpl: FP8 mode requires HpcRopeNorm. "
385 + "Ensure hpc_rope_norm is enabled or set "
386 + "kv_cache_dtype='auto' for bf16 mode."
387 + f" (layer={getattr(layer, 'layer_name', '?')})"
388 + )
389 + k_scale = layer._k_scale.reshape(1)
390 + v_scale = layer._v_scale.reshape(1)
391 +
392 + query = query[:num_actual_tokens]
393 + key = key[:num_actual_tokens]
394 + value = value[:num_actual_tokens]
395 + output_padded = output
396 + output = output[:num_actual_tokens]
397 +
398 + # --- Prefill ---
399 + if num_prefill_reqs > 0:
400 + seq_lens_prefill = attn_metadata.seq_lens[num_decode_reqs:]
401 + cu_seqlens_prefill = attn_metadata.qo_indptr
402 + max_seqlens = attn_metadata.max_query_len
403 + block_table_prefill = attn_metadata.block_table_tensor[num_decode_reqs:]
404 +
405 + q_prefill = query[num_decode_tokens:]
406 + output_prefill = output[num_decode_tokens:]
407 +
408 + if self.use_fp8:
409 + hpc.attention_with_kvcache_prefill_fp8(
410 + q_prefill,
411 + kv_cache[:, 0],
412 + kv_cache[:, 1],
413 + hpc_prefill_q_scale,
414 + k_scale,
415 + v_scale,
416 + cu_seqlens_prefill,
417 + block_table_prefill,
418 + seq_lens_prefill,
419 + max_seqlens,
420 + output=output_prefill,
421 + )
422 + else:
423 + hpc.attention_with_kvcache_prefill_bf16(
424 + q_prefill,
425 + kv_cache[:, 0],
426 + kv_cache[:, 1],
427 + cu_seqlens_prefill,
428 + block_table_prefill,
429 + seq_lens_prefill,
430 + max_seqlens,
431 + output=output_prefill,
432 + )
433 +
434 + # --- Decode ---
435 + if num_decode_reqs > 0:
436 + num_seq_kvcache = attn_metadata.seq_lens[:num_decode_reqs]
437 + block_table_decode = attn_metadata.block_table_tensor[:num_decode_reqs]
438 +
439 + q_decode = query[:num_decode_tokens]
440 + output_decode = output[:num_decode_tokens]
441 +
442 + if self.use_fp8:
443 + hpc.attention_decode_fp8(
444 + q_decode,
445 + kv_cache[:, 0],
446 + kv_cache[:, 1],
447 + block_table_decode,
448 + num_seq_kvcache,
449 + hpc_decode_q_scale,
450 + k_scale,
451 + v_scale,
452 + new_kv_included=True,
453 + splitk=self.splitk,
454 + split_flag=hpc_split_k_flag,
455 + output=output_decode,
456 + )
457 + else:
458 + hpc.attention_decode_bf16(
459 + q_decode,
460 + kv_cache[:, 0],
461 + kv_cache[:, 1],
462 + block_table_decode,
463 + num_seq_kvcache,
464 + output=output_decode,
465 + new_kv_included=True,
466 + splitk=self.splitk,
467 + )
468 +
469 + return output_padded
vllm/v1/attention/backends/registry.py
+6
@@ -103,6 +103,12 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
103 )
104 NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"
105 FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
106 + # HPC Attention Backend:
107 + # powered by operators from https://github.com/Tencent/hpc-ops.
108 + # Only supported on NVIDIA Hopper GPUs (e.g. H20, H200),
109 + # currently limited to the Hy3 model,
110 + # and requires a block size of 64.
111 + HPC_ATTN = "vllm.v1.attention.backends.hpc_attn.HpcAttentionBackend"
112 ROCM_AITER_UNIFIED_ATTN = (
113 "vllm.v1.attention.backends.rocm_aiter_unified_attn."
114 "RocmAiterUnifiedAttentionBackend"