[HPC Attention Backend] hpc attention backend support bf16 kv cache with fp8 weight (#50980)
Signed-off-by: chengvjiang <chengvjiang@tencent.com> Co-authored-by: chengvjiang <chengvjiang@tencent.com> Co-authored-by: VAthree <56064364+VAthree@users.noreply.github.com>
Cheng Jiang committed
Aug 6, 2026 at 15:40 UTC
470297c143c9465d4cfff4b854f63f0a283ec5de
3 files changed
+75
-36
vllm/model_executor/layers/fused_moe/hpc_moe.py
+1
-5
@@ -62,11 +62,7 @@ class HPCExperts(mk.FusedMoEExpertsModular):
62
@staticmethod
63
def _supports_current_device() -> bool:
64
p = current_platform
65
- return (
66
- p.is_cuda()
67
- and (p.is_device_capability(90) or p.is_device_capability_family(100))
68
- and has_hpc()
69
- )
65
+ return p.is_cuda() and p.is_device_capability(90) and has_hpc()
66
67
@staticmethod
68
def _supports_no_act_and_mul() -> bool:
vllm/model_executor/layers/hpc/rope_norm.py
+10
-2
@@ -205,10 +205,10 @@ class HpcRopeNorm(CustomOp, HpcModule):
205
):
206
return False
207
208
- if kv_cache_dtype not in ("fp8_e4m3", "auto"):
208
+ if kv_cache_dtype not in ("fp8_e4m3", "auto", "bfloat16"):
209
logger.warning_once(
210
f"hpc rope_norm not support kv_cache_dtype:{kv_cache_dtype}, "
211
- "only support fp8_e4m3, bfloat16"
211
+ "only support fp8_e4m3, auto, bfloat16"
212
)
213
return False
214
@@ -414,3 +414,11 @@ class HpcRopeNorm(CustomOp, HpcModule):
414
out_q=out_q_decode,
415
qk_norm_policy=self.qk_norm_policy,
416
)
417
+
418
+ # Signal HpcAttentionImpl that KV cache has already been written by
419
+ # rope_norm_store_kv[_fp8] above, so it should skip its own
420
+ # reshape_and_cache_flash. Set after the fused kernels ran (either
421
+ # prefill or decode branch); otherwise the standard KV-write path in
422
+ # the attention impl would kick in (which is what non-HpcRopeNorm
423
+ # models rely on).
424
+ attn_metadata.hpc_kv_written = True
vllm/v1/attention/backends/hpc_attn.py
+64
-29
@@ -80,7 +80,7 @@ class HpcAttnMetadata(AttentionMetadata):
80
# HpcAttentionImpl.forward(). Defaults are safe for the standard
81
# (non-RopeNorm) path and for profiling runs (attn_metadata=None).
82
hpc_kv_written: bool = False
83
- """True when HpcRopeNorm already wrote KV cache."""
83
+ """True when HpcRopeNorm has already written K/V into the paged cache."""
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
@@ -146,12 +146,24 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
146
supports_spec_as_decode=True,
147
)
148
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
- )
149
+ # Dynamic-scheduled decode:
150
+ # - fp8 KV cache path always uses task_map (FP8 decode kernel
151
+ # requires it).
152
+ # - bf16 KV cache path only uses task_map when the installed hpc
153
+ # version supports it (hpc >= 6e2eced / PR #73). Older builds
154
+ # take the static split-K path with no task_map.
155
+ kv_cache_dtype = kv_cache_spec.dtype
156
+ self.use_fp8 = kv_cache_dtype == torch.float8_e4m3fn
157
+ self._dynamic_sched = True
158
+ if self._dynamic_sched:
159
+ self.task_map = hpc.get_attention_decode_task_workspace(
160
+ vllm_config.scheduler_config.max_num_seqs,
161
+ vllm_config.model_config.max_model_len or 4096,
162
+ self.num_kv_heads,
163
+ min_process_len=self.hpc_dynamic_sched_attn_min_split_len,
164
+ )
165
+ else:
166
+ self.task_map = None
167
168
@override # type: ignore[misc]
169
@classmethod
@@ -214,16 +226,20 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
226
seq_lens_decode = seq_lens[:num_decodes]
227
# block_table is per-request, indexed by num_decodes (not tokens)
228
qo_indptr_decode = common_attn_metadata.query_start_loc[: num_decodes + 1]
217
- import hpc
229
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
- )
230
+ # Skip the task-map scheduler pass when dynamic scheduling is off
231
+ # (bf16 KV cache on an older hpc that lacks bf16 decode task_map).
232
+ if self._dynamic_sched:
233
+ import hpc
234
+
235
+ hpc.assign_attention_decode_task(
236
+ seq_lens_decode,
237
+ self.task_map,
238
+ self.num_kv_heads,
239
+ decode_query_len,
240
+ new_kv_included=True,
241
+ min_process_len=self.hpc_dynamic_sched_attn_min_split_len,
242
+ )
243
244
return HpcAttnMetadata(
245
num_actual_tokens=num_actual_tokens,
@@ -236,7 +252,7 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
252
seq_lens=seq_lens,
253
block_table_tensor=block_table_tensor,
254
qo_indptr=qo_indptr,
239
- hpc_kv_written=True,
255
+ hpc_kv_written=False,
256
hpc_prefill_q_scale=None,
257
hpc_decode_q_scale=None,
258
hpc_split_k_flag=None,
@@ -306,7 +322,7 @@ class HpcAttentionBackend(AttentionBackend):
322
323
@classmethod
324
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
309
- return capability >= DeviceCapability(9, 0)
325
+ return capability == DeviceCapability(9, 0)
326
327
@classmethod
328
def supports_kv_cache_dtype(cls, kv_cache_dtype: "CacheDType | None") -> bool:
@@ -366,10 +382,10 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
382
f"(num_heads={num_heads}, num_kv_heads={num_kv_heads})"
383
)
384
369
- if kv_cache_dtype not in ("auto", "fp8_e4m3"):
385
+ if kv_cache_dtype not in ("auto", "fp8_e4m3", "bfloat16"):
386
raise ValueError(
387
f"HPC attention only supports kv_cache_dtype 'auto' or "
372
- f"'fp8_e4m3', got '{kv_cache_dtype}'"
388
+ f"'fp8_e4m3' or 'bfloat16', got '{kv_cache_dtype}'"
389
)
390
391
self.num_heads = num_heads
@@ -436,8 +452,17 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
452
num_decode_reqs = attn_metadata.num_decodes
453
num_decode_tokens = attn_metadata.num_decode_tokens
454
439
- # Write KV cache if not already done by HpcRopeNorm.
440
- if self.kv_sharing_target_layer_name is None and not hpc_kv_written:
455
+ # Standard-path KV write: when no upstream fused op wrote KV cache
456
+ # (i.e. the model didn't wire HpcRopeNorm), fall back to the generic
457
+ # reshape_and_cache_flash. This is what lets non-HunYuan-V3 models use
458
+ # the HPC attention backend with bf16 KV cache without any model-side
459
+ # changes. FP8 KV cache always needs HpcRopeNorm (see check below), so
460
+ # we never take this path in fp8 mode.
461
+ if (
462
+ not self.use_fp8
463
+ and self.kv_sharing_target_layer_name is None
464
+ and not hpc_kv_written
465
+ ):
466
torch.ops._C_cache_ops.reshape_and_cache_flash(
467
key,
468
value,
@@ -449,18 +474,18 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
474
layer._v_scale,
475
)
476
452
- if self.use_fp8:
453
- torch_dtype = _get_fp8_dtype_for_kv_cache(self.kv_cache_dtype)
454
- kv_cache = kv_cache.view(torch_dtype)
455
-
477
if self.use_fp8:
478
if not hpc_kv_written:
479
raise RuntimeError(
459
- "HpcAttentionImpl: FP8 mode requires HpcRopeNorm. "
460
- "Ensure hpc_rope_norm is enabled or set "
461
- "kv_cache_dtype='auto' for bf16 mode."
480
+ "HpcAttentionImpl: FP8 KV cache mode requires HpcRopeNorm "
481
+ "(the fused rope_norm_store_kv_fp8 op produces the "
482
+ "per-token-per-head Q scales the FP8 attention kernel "
483
+ "needs). Ensure the model wires hpc_rope_norm, or switch "
484
+ "to --kv-cache-dtype auto / bfloat16."
485
f" (layer={getattr(layer, 'layer_name', '?')})"
486
)
487
+ torch_dtype = _get_fp8_dtype_for_kv_cache(self.kv_cache_dtype)
488
+ kv_cache = kv_cache.view(torch_dtype)
489
k_scale = layer._k_scale.reshape(1)
490
v_scale = layer._v_scale.reshape(1)
491
@@ -539,6 +564,15 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
564
output=output_decode,
565
)
566
else:
567
+ # Older hpc's bf16 decode has no task_map kwarg; the map is
568
+ # only ever built when the installed hpc supports it (see
569
+ # HpcAttnMetadataBuilder._dynamic_sched). Pass task_map
570
+ # conditionally to stay compatible with both versions
571
+ task_map_kwargs = (
572
+ {"task_map": attn_metadata.task_map}
573
+ if attn_metadata.task_map is not None
574
+ else {}
575
+ )
576
hpc.attention_decode_bf16(
577
q_decode,
578
kv_cache[:, 0],
@@ -549,6 +583,7 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
583
output=output_decode,
584
new_kv_included=True,
585
splitk=self.splitk,
586
+ **task_map_kwargs,
587
)
588
589
return output_padded