| 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.utils.torch_utils import PIN_MEMORY |
| 23 | from vllm.v1.attention.backend import ( |
| 24 | AttentionBackend, |
| 25 | AttentionCGSupport, |
| 26 | AttentionImpl, |
| 27 | AttentionMetadata, |
| 28 | AttentionMetadataBuilder, |
| 29 | AttentionType, |
| 30 | CommonAttentionMetadata, |
| 31 | MultipleOf, |
| 32 | ) |
| 33 | from vllm.v1.attention.backends.utils import ( |
| 34 | split_decodes_and_prefills, |
| 35 | ) |
| 36 | from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout, KVCacheSpec |
| 37 | |
| 38 | logger = init_logger(__name__) |
| 39 | |
| 40 | FP8_DTYPE = current_platform.fp8_dtype() |
| 41 | |
| 42 | |
| 43 | def _get_fp8_dtype_for_kv_cache(kv_cache_dtype: str) -> torch.dtype: |
| 44 | """Return the torch FP8 dtype for the given kv_cache_dtype string.""" |
| 45 | if kv_cache_dtype in ("fp8", "fp8_e4m3"): |
| 46 | return torch.float8_e4m3fn |
| 47 | elif kv_cache_dtype == "fp8_e5m2": |
| 48 | return torch.float8_e5m2 |
| 49 | else: |
| 50 | raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}") |
| 51 | |
| 52 | |
| 53 | @dataclass |
| 54 | class HpcAttnMetadata(AttentionMetadata): |
| 55 | """Metadata required by the HPC attention kernel.""" |
| 56 | |
| 57 | num_actual_tokens: int |
| 58 | num_decodes: int |
| 59 | num_decode_tokens: int |
| 60 | num_prefills: int |
| 61 | num_prefill_tokens: int |
| 62 | max_query_len: int |
| 63 | |
| 64 | slot_mapping: torch.Tensor |
| 65 | """Slot mapping for KV cache writes. shape = [num_actual_tokens]""" |
| 66 | |
| 67 | seq_lens: torch.Tensor |
| 68 | """KV cache length per request. shape = [batch_size]""" |
| 69 | |
| 70 | block_table_tensor: torch.Tensor |
| 71 | """Paged KV-cache block table. |
| 72 | shape = [batch_size, max_num_blocks_per_seq]""" |
| 73 | |
| 74 | qo_indptr: torch.Tensor | None = None |
| 75 | """Cumulative query lengths for prefill requests (GPU tensor). |
| 76 | shape = [num_prefills + 1]. None when num_prefills == 0.""" |
| 77 | |
| 78 | # --- HPC RopeNorm pass-through fields --- |
| 79 | # Set by HpcRopeNorm._forward_impl(); consumed & reset by |
| 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 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 |
| 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 (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]): |
| 107 | """Build HpcAttnMetadata from CommonAttentionMetadata.""" |
| 108 | |
| 109 | _cudagraph_support = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE |
| 110 | reorder_batch_threshold: int = 1 |
| 111 | |
| 112 | def __init__( |
| 113 | self, |
| 114 | kv_cache_spec: AttentionSpec, |
| 115 | layer_names: list[str], |
| 116 | vllm_config: VllmConfig, |
| 117 | device: torch.device, |
| 118 | ): |
| 119 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 120 | import hpc |
| 121 | |
| 122 | self.num_kv_heads = kv_cache_spec.num_kv_heads |
| 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 | |
| 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 | |
| 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 |
| 170 | def get_cudagraph_support( |
| 171 | cls: type["HpcAttnMetadataBuilder"], |
| 172 | vllm_config: VllmConfig, |
| 173 | kv_cache_spec: KVCacheSpec, |
| 174 | ) -> AttentionCGSupport: |
| 175 | spec_config = vllm_config.speculative_config |
| 176 | if ( |
| 177 | spec_config is not None |
| 178 | and spec_config.num_speculative_tokens is not None |
| 179 | and spec_config.num_speculative_tokens > 0 |
| 180 | ): |
| 181 | return AttentionCGSupport.UNIFORM_BATCH |
| 182 | return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE |
| 183 | |
| 184 | def build( |
| 185 | self, |
| 186 | common_prefix_len: int, |
| 187 | common_attn_metadata: CommonAttentionMetadata, |
| 188 | fast_build: bool = False, |
| 189 | ) -> HpcAttnMetadata: |
| 190 | """Build HpcAttnMetadata from CommonAttentionMetadata.""" |
| 191 | num_actual_tokens = common_attn_metadata.num_actual_tokens |
| 192 | |
| 193 | num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( |
| 194 | split_decodes_and_prefills( |
| 195 | common_attn_metadata, |
| 196 | decode_threshold=self.reorder_batch_threshold, |
| 197 | # MTP requires uniform query lengths across decode requests |
| 198 | require_uniform=(self.reorder_batch_threshold > 1), |
| 199 | ) |
| 200 | ) |
| 201 | |
| 202 | seq_lens = common_attn_metadata.seq_lens |
| 203 | block_table_tensor = common_attn_metadata.block_table_tensor |
| 204 | slot_mapping = common_attn_metadata.slot_mapping |
| 205 | max_query_len = common_attn_metadata.max_query_len |
| 206 | |
| 207 | # Compute decode_query_len (tokens per decode request). |
| 208 | # Non-MTP: 1, MTP: mtp+1 (2 or 3). |
| 209 | if num_decodes > 0 and num_decode_tokens > num_decodes: |
| 210 | decode_query_len = num_decode_tokens // num_decodes |
| 211 | else: |
| 212 | decode_query_len = 1 |
| 213 | |
| 214 | seq_lens_decode = None |
| 215 | qo_indptr = None |
| 216 | qo_indptr_decode = None |
| 217 | if num_prefills > 0: |
| 218 | qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu |
| 219 | prefill_start = num_decodes |
| 220 | qo_indptr_prefill_cpu = torch.empty( |
| 221 | num_prefills + 1, dtype=torch.int32, pin_memory=PIN_MEMORY |
| 222 | ) |
| 223 | torch.subtract( |
| 224 | qo_indptr_cpu[prefill_start:], |
| 225 | qo_indptr_cpu[prefill_start], |
| 226 | out=qo_indptr_prefill_cpu, |
| 227 | ) |
| 228 | qo_indptr = qo_indptr_prefill_cpu.to(self.device, non_blocking=True) |
| 229 | |
| 230 | if num_decodes > 0: |
| 231 | seq_lens_decode = seq_lens[:num_decodes] |
| 232 | # block_table is per-request, indexed by num_decodes (not tokens) |
| 233 | qo_indptr_decode = common_attn_metadata.query_start_loc[: num_decodes + 1] |
| 234 | |
| 235 | # Skip the task-map scheduler pass when dynamic scheduling is off |
| 236 | # (bf16 KV cache on an older hpc that lacks bf16 decode task_map). |
| 237 | if self._dynamic_sched: |
| 238 | import hpc |
| 239 | |
| 240 | hpc.assign_attention_decode_task( |
| 241 | seq_lens_decode, |
| 242 | self.task_map, |
| 243 | self.num_kv_heads, |
| 244 | decode_query_len, |
| 245 | new_kv_included=True, |
| 246 | min_process_len=self.hpc_dynamic_sched_attn_min_split_len, |
| 247 | ) |
| 248 | |
| 249 | return HpcAttnMetadata( |
| 250 | num_actual_tokens=num_actual_tokens, |
| 251 | num_decodes=num_decodes, |
| 252 | num_decode_tokens=num_decode_tokens, |
| 253 | num_prefills=num_prefills, |
| 254 | num_prefill_tokens=num_prefill_tokens, |
| 255 | max_query_len=max_query_len, |
| 256 | slot_mapping=slot_mapping, |
| 257 | seq_lens=seq_lens, |
| 258 | block_table_tensor=block_table_tensor, |
| 259 | qo_indptr=qo_indptr, |
| 260 | hpc_kv_written=False, |
| 261 | hpc_prefill_q_scale=None, |
| 262 | hpc_decode_q_scale=None, |
| 263 | hpc_split_k_flag=None, |
| 264 | decode_query_len=decode_query_len, |
| 265 | qo_indptr_decode=qo_indptr_decode, |
| 266 | task_map=self.task_map, |
| 267 | ) |
| 268 | |
| 269 | |
| 270 | class HpcAttentionBackend(AttentionBackend): |
| 271 | """HPC attention backend (pure attention, no RoPE/Norm). |
| 272 | |
| 273 | KV cache layout: LBNHC (num_blocks, num_kv_heads, block_size, 2 * head_size). |
| 274 | """ |
| 275 | |
| 276 | accept_output_buffer: bool = True |
| 277 | supported_dtypes: ClassVar[list[torch.dtype]] = [ |
| 278 | torch.float16, |
| 279 | torch.bfloat16, |
| 280 | ] |
| 281 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 282 | "auto", |
| 283 | "bfloat16", |
| 284 | "fp8_e4m3", |
| 285 | ] |
| 286 | |
| 287 | # Avoid attention abstracted method call cache insert |
| 288 | forward_includes_kv_cache_update: bool = True |
| 289 | |
| 290 | @staticmethod |
| 291 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 292 | return [64] |
| 293 | |
| 294 | @staticmethod |
| 295 | def get_name() -> str: |
| 296 | return "HPC_ATTN" |
| 297 | |
| 298 | @staticmethod |
| 299 | def get_impl_cls() -> type["HpcAttentionImpl"]: |
| 300 | return HpcAttentionImpl |
| 301 | |
| 302 | @staticmethod |
| 303 | def get_builder_cls() -> type["HpcAttnMetadataBuilder"]: |
| 304 | return HpcAttnMetadataBuilder |
| 305 | |
| 306 | @classmethod |
| 307 | def get_supported_head_sizes(cls) -> list[int]: |
| 308 | return [128] |
| 309 | |
| 310 | @classmethod |
| 311 | def supports_compute_capability(cls, capability: DeviceCapability) -> bool: |
| 312 | return capability == DeviceCapability(9, 0) |
| 313 | |
| 314 | @classmethod |
| 315 | def supports_kv_cache_dtype(cls, kv_cache_dtype: "CacheDType | None") -> bool: |
| 316 | if kv_cache_dtype is None: |
| 317 | return True |
| 318 | return kv_cache_dtype in cls.supported_kv_cache_dtypes |
| 319 | |
| 320 | @classmethod |
| 321 | def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: |
| 322 | return (KVCacheLayout.LBNHC,) |
| 323 | |
| 324 | |
| 325 | class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): |
| 326 | """HPC pure attention implementation (no RoPE/Norm). |
| 327 | |
| 328 | Constraints: |
| 329 | - head_dim == 128 |
| 330 | - num_heads // num_kv_heads in {4, 8} |
| 331 | - kv_cache_dtype in {"auto", "fp8_e4m3"} |
| 332 | """ |
| 333 | |
| 334 | def __init__( |
| 335 | self, |
| 336 | num_heads: int, |
| 337 | head_size: int, |
| 338 | scale: float, |
| 339 | num_kv_heads: int, |
| 340 | alibi_slopes: list[float] | None = None, |
| 341 | sliding_window: int | None = None, |
| 342 | kv_cache_dtype: str = "auto", |
| 343 | logits_soft_cap: float | None = None, |
| 344 | attn_type: str = AttentionType.DECODER, |
| 345 | kv_sharing_target_layer_name: str | None = None, |
| 346 | ) -> None: |
| 347 | if importlib.util.find_spec("hpc") is None: |
| 348 | raise ImportError( |
| 349 | "HPC attention requires the hpc module to be installed. " |
| 350 | "Please install it from https://github.com/Tencent/hpc-ops" |
| 351 | ) |
| 352 | if attn_type != AttentionType.DECODER: |
| 353 | raise NotImplementedError("HPC attention only supports decoder attention") |
| 354 | if alibi_slopes is not None: |
| 355 | raise NotImplementedError("HPC attention does not support ALiBi") |
| 356 | if logits_soft_cap is not None: |
| 357 | raise NotImplementedError("HPC attention does not support logits_soft_cap") |
| 358 | |
| 359 | if head_size != 128: |
| 360 | raise ValueError( |
| 361 | f"HPC attention only supports head_dim=128, got {head_size}" |
| 362 | ) |
| 363 | |
| 364 | num_queries_per_kv = num_heads // num_kv_heads |
| 365 | if num_queries_per_kv not in (4, 8): |
| 366 | raise ValueError( |
| 367 | f"HPC attention only supports head_per_group in {{4, 8}}, " |
| 368 | f"got {num_queries_per_kv} " |
| 369 | f"(num_heads={num_heads}, num_kv_heads={num_kv_heads})" |
| 370 | ) |
| 371 | |
| 372 | if kv_cache_dtype not in ("auto", "fp8_e4m3", "bfloat16"): |
| 373 | raise ValueError( |
| 374 | f"HPC attention only supports kv_cache_dtype 'auto' or " |
| 375 | f"'fp8_e4m3' or 'bfloat16', got '{kv_cache_dtype}'" |
| 376 | ) |
| 377 | |
| 378 | self.num_heads = num_heads |
| 379 | self.head_size = head_size |
| 380 | self.scale = float(scale) |
| 381 | self.num_kv_heads = num_kv_heads |
| 382 | self.kv_cache_dtype = kv_cache_dtype |
| 383 | self.kv_sharing_target_layer_name = kv_sharing_target_layer_name |
| 384 | |
| 385 | self.num_queries_per_kv = num_queries_per_kv |
| 386 | |
| 387 | if sliding_window is None: |
| 388 | self.sliding_window = (-1, -1) |
| 389 | else: |
| 390 | self.sliding_window = (sliding_window - 1, 0) |
| 391 | |
| 392 | self.use_fp8 = kv_cache_dtype == "fp8_e4m3" |
| 393 | |
| 394 | self.supports_quant_query_input = False |
| 395 | self.splitk = True |
| 396 | |
| 397 | import hpc |
| 398 | |
| 399 | if self.use_fp8: |
| 400 | self._quant_type = hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR |
| 401 | else: |
| 402 | self._quant_type = None |
| 403 | |
| 404 | def forward( |
| 405 | self, |
| 406 | layer: torch.nn.Module, |
| 407 | query: torch.Tensor, |
| 408 | key: torch.Tensor, |
| 409 | value: torch.Tensor, |
| 410 | kv_cache: torch.Tensor, |
| 411 | attn_metadata: HpcAttnMetadata | None, |
| 412 | output: torch.Tensor | None = None, |
| 413 | output_scale: torch.Tensor | None = None, |
| 414 | output_block_scale: torch.Tensor | None = None, |
| 415 | ) -> torch.Tensor: |
| 416 | """HPC attention forward (standard vLLM backend interface). |
| 417 | |
| 418 | Two modes: |
| 419 | 1. Standard: upstream handles RoPE/Norm; this backend writes KV + attn. |
| 420 | 2. HpcRopeNorm: fused op already did RoPE/Norm/KV-Write/Q-Quant; |
| 421 | extra params passed via attn_metadata.hpc_* fields. |
| 422 | """ |
| 423 | import hpc |
| 424 | |
| 425 | assert output is not None, "Output tensor must be provided." |
| 426 | assert output_scale is None, "HPC attention does not support fused output quant" |
| 427 | assert output_block_scale is None |
| 428 | |
| 429 | if attn_metadata is None: |
| 430 | return output.fill_(0) |
| 431 | |
| 432 | hpc_kv_written = attn_metadata.hpc_kv_written |
| 433 | hpc_prefill_q_scale = attn_metadata.hpc_prefill_q_scale |
| 434 | hpc_decode_q_scale = attn_metadata.hpc_decode_q_scale |
| 435 | hpc_split_k_flag = attn_metadata.hpc_split_k_flag |
| 436 | |
| 437 | num_actual_tokens = attn_metadata.num_actual_tokens |
| 438 | num_prefill_reqs = attn_metadata.num_prefills |
| 439 | num_decode_reqs = attn_metadata.num_decodes |
| 440 | num_decode_tokens = attn_metadata.num_decode_tokens |
| 441 | |
| 442 | # Logical (B, H, N, 2*head_size) -> (B, N, H, head_size) K/V views. |
| 443 | key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) |
| 444 | |
| 445 | # Standard-path KV write: when no upstream fused op wrote KV cache |
| 446 | # (i.e. the model didn't wire HpcRopeNorm), fall back to the generic |
| 447 | # reshape_and_cache_flash. This is what lets non-HunYuan-V3 models use |
| 448 | # the HPC attention backend with bf16 KV cache without any model-side |
| 449 | # changes. FP8 KV cache always needs HpcRopeNorm (see check below), so |
| 450 | # we never take this path in fp8 mode. |
| 451 | if ( |
| 452 | not self.use_fp8 |
| 453 | and self.kv_sharing_target_layer_name is None |
| 454 | and not hpc_kv_written |
| 455 | ): |
| 456 | torch.ops._C_cache_ops.reshape_and_cache_flash( |
| 457 | key, |
| 458 | value, |
| 459 | key_cache, |
| 460 | value_cache, |
| 461 | attn_metadata.slot_mapping, |
| 462 | self.kv_cache_dtype, |
| 463 | layer._k_scale, |
| 464 | layer._v_scale, |
| 465 | ) |
| 466 | |
| 467 | if self.use_fp8: |
| 468 | if not hpc_kv_written: |
| 469 | raise RuntimeError( |
| 470 | "HpcAttentionImpl: FP8 KV cache mode requires HpcRopeNorm " |
| 471 | "(the fused rope_norm_store_kv_fp8 op produces the " |
| 472 | "per-token-per-head Q scales the FP8 attention kernel " |
| 473 | "needs). Ensure the model wires hpc_rope_norm, or switch " |
| 474 | "to --kv-cache-dtype auto / bfloat16." |
| 475 | f" (layer={getattr(layer, 'layer_name', '?')})" |
| 476 | ) |
| 477 | torch_dtype = _get_fp8_dtype_for_kv_cache(self.kv_cache_dtype) |
| 478 | key_cache = key_cache.view(torch_dtype) |
| 479 | value_cache = value_cache.view(torch_dtype) |
| 480 | k_scale = layer._k_scale.reshape(1) |
| 481 | v_scale = layer._v_scale.reshape(1) |
| 482 | |
| 483 | query = query[:num_actual_tokens] |
| 484 | key = key[:num_actual_tokens] |
| 485 | value = value[:num_actual_tokens] |
| 486 | output_padded = output |
| 487 | output = output[:num_actual_tokens] |
| 488 | |
| 489 | # --- Prefill --- |
| 490 | if num_prefill_reqs > 0: |
| 491 | seq_lens_prefill = attn_metadata.seq_lens[num_decode_reqs:] |
| 492 | cu_seqlens_prefill = attn_metadata.qo_indptr |
| 493 | max_seqlens = attn_metadata.max_query_len |
| 494 | block_table_prefill = attn_metadata.block_table_tensor[num_decode_reqs:] |
| 495 | |
| 496 | q_prefill = query[num_decode_tokens:] |
| 497 | output_prefill = output[num_decode_tokens:] |
| 498 | |
| 499 | if self.use_fp8: |
| 500 | hpc.attention_with_kvcache_prefill_fp8( |
| 501 | q_prefill, |
| 502 | key_cache, |
| 503 | value_cache, |
| 504 | hpc_prefill_q_scale, |
| 505 | k_scale, |
| 506 | v_scale, |
| 507 | cu_seqlens_prefill, |
| 508 | block_table_prefill, |
| 509 | seq_lens_prefill, |
| 510 | max_seqlens, |
| 511 | quant_type=self._quant_type, |
| 512 | output=output_prefill, |
| 513 | ) |
| 514 | else: |
| 515 | hpc.attention_with_kvcache_prefill_bf16( |
| 516 | q_prefill, |
| 517 | key_cache, |
| 518 | value_cache, |
| 519 | cu_seqlens_prefill, |
| 520 | block_table_prefill, |
| 521 | seq_lens_prefill, |
| 522 | max_seqlens, |
| 523 | output=output_prefill, |
| 524 | ) |
| 525 | |
| 526 | # --- Decode --- |
| 527 | if num_decode_reqs > 0: |
| 528 | num_seq_kvcache = attn_metadata.seq_lens[:num_decode_reqs] |
| 529 | block_table_decode = attn_metadata.block_table_tensor[:num_decode_reqs] |
| 530 | |
| 531 | q_decode = query[:num_decode_tokens] |
| 532 | output_decode = output[:num_decode_tokens] |
| 533 | |
| 534 | mtp = attn_metadata.decode_query_len - 1 |
| 535 | |
| 536 | if self.use_fp8: |
| 537 | hpc.attention_decode_fp8( |
| 538 | q_decode, |
| 539 | key_cache, |
| 540 | value_cache, |
| 541 | block_table_decode, |
| 542 | num_seq_kvcache, |
| 543 | hpc_decode_q_scale, |
| 544 | k_scale, |
| 545 | v_scale, |
| 546 | mtp=mtp, |
| 547 | # MTP: split_flag from rope_norm is unavailable |
| 548 | # when using prefill-mode kernel; let HPC decide. |
| 549 | # splitk=(self.splitk if mtp == 0 else True), |
| 550 | new_kv_included=True, |
| 551 | quant_type=self._quant_type, |
| 552 | splitk=self.splitk, |
| 553 | task_map=attn_metadata.task_map, |
| 554 | split_flag=hpc_split_k_flag, |
| 555 | output=output_decode, |
| 556 | ) |
| 557 | else: |
| 558 | # Older hpc's bf16 decode has no task_map kwarg; the map is |
| 559 | # only ever built when the installed hpc supports it (see |
| 560 | # HpcAttnMetadataBuilder._dynamic_sched). Pass task_map |
| 561 | # conditionally to stay compatible with both versions |
| 562 | task_map_kwargs = ( |
| 563 | {"task_map": attn_metadata.task_map} |
| 564 | if attn_metadata.task_map is not None |
| 565 | else {} |
| 566 | ) |
| 567 | hpc.attention_decode_bf16( |
| 568 | q_decode, |
| 569 | key_cache, |
| 570 | value_cache, |
| 571 | block_table_decode, |
| 572 | num_seq_kvcache, |
| 573 | mtp=mtp, |
| 574 | output=output_decode, |
| 575 | new_kv_included=True, |
| 576 | splitk=self.splitk, |
| 577 | **task_map_kwargs, |
| 578 | ) |
| 579 | |
| 580 | return output_padded |