vllm.model_executor.layers.hpc.rope_norm ¶

@CustomOp.register("hpc_rope_norm")
class HpcRopeNorm(CustomOp, HpcModule):
    """HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant).

    Registered as a sub-module in model layers (e.g. HunYuanAttention).
    Norm weights are extracted from fallback norm modules via
    process_weights_after_loading() after all weights are loaded.

    forward() is dispatched by CustomOp framework:
    - In compiled mode: forward_cuda() calls torch.ops.vllm.hpc_rope_norm_forward
      as a splitting point — internal Python control flow is opaque
      to torch.compile and not captured by CUDA Graph.
    - In eager/native mode: forward_native() falls back to forward_cuda().
    """

    def __init__(
        self,
        num_heads: int,
        num_kv_heads: int,
        head_dim: int,
        cos_sin_cache: torch.Tensor,
        use_qk_norm: bool,
        fallback_qnorm: torch.nn.Module | None,
        fallback_knorm: torch.nn.Module | None,
        kv_cache_dtype: str,
        layer_name: str,
        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

        self.use_qk_norm = use_qk_norm

        self.q_size = num_heads * head_dim
        self.kv_size = num_kv_heads * head_dim

        # Register as a non-persistent buffer so it participates in sleep
        # level-2 save/restore (CuMemAllocator) but is excluded from the
        # checkpoint state_dict.
        self.register_buffer("cos_sin_cache", cos_sin_cache.float(), persistent=False)

        self.fallback_qnorm = fallback_qnorm
        self.fallback_knorm = fallback_knorm

        self.head_per_group = num_heads // num_kv_heads

        # Pre-allocate norm weight tensors as Parameters so they are tracked by
        # CuMemAllocator (for sleep/wake_up) and have stable addresses for CUDA
        # Graph replay. process_weights_after_loading() updates them inplace via
        # copy_() so refit does not invalidate captured graph tensor pointers.
        # Shape is [head_dim] to match the HPC kernel's q/k_norm_weight layout.
        if use_qk_norm and fallback_qnorm is not None:
            self.qnorm_weight: torch.nn.Parameter | None = torch.nn.Parameter(
                torch.empty(head_dim, dtype=torch.float32),
                requires_grad=False,
            )
        else:
            self.qnorm_weight = None
        if use_qk_norm and fallback_knorm is not None:
            self.knorm_weight: torch.nn.Parameter | None = torch.nn.Parameter(
                torch.empty(head_dim, dtype=torch.float32),
                requires_grad=False,
            )
        else:
            self.knorm_weight = None

        self.use_fp8 = "fp8" in kv_cache_dtype
        # The RMSNorm/RoPE ordering is model dependent (e.g. HunYuan V3 applies
        # QK-Norm before RoPE -> NORM_THEN_ROPE), so it is supplied by the
        # caller. When QK-Norm is disabled the policy is forced to NONE.
        self.qk_norm_policy = qk_norm_policy if use_qk_norm else QkNormPolicy.NONE

        # Register layer_name + add self to the global instance registry so the
        # module-level custom op (hpc_rope_norm_forward) can route back here.
        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,
        num_heads: int,
        num_kv_heads: int,
        head_dim: int,
        kv_cache_dtype: str,
    ) -> bool:
        """Check whether HpcRopeNorm is supported for the given config."""
        # HpcRopeNorm is only enabled together with the HPC attention backend.
        vllm_config = get_current_vllm_config_or_none()
        if (
            vllm_config is None
            or vllm_config.attention_config.backend != AttentionBackendEnum.HPC_ATTN
        ):
            return False

        if kv_cache_dtype not in ("fp8_e4m3", "auto"):
            logger.warning_once(
                f"hpc rope_norm not support kv_cache_dtype:{kv_cache_dtype}, "
                "only support fp8_e4m3, bfloat16"
            )
            return False

        if head_dim not in (128,):
            logger.warning_once("hpc rope_norm only support head_dim == 128.")
            return False

        head_per_group = num_heads // num_kv_heads
        if head_per_group not in (4, 8):
            logger.warning_once("hpc rope_norm only support head_per_group in [4, 8].")
            return False

        logger.info_once("enable hpc rope_norm")
        return True

    def process_weights_after_loading(self, model: torch.nn.Module = None) -> None:
        """Copy norm weights (float32) from fallback norm modules inplace.

        Uses copy_() to preserve tensor addresses for CUDA Graph / refit
        compatibility. Called by the model's load_weights() after all weights
        are loaded (and generically from the model loader for DummyModelLoader
        / sleep-wake_up reload paths).
        """
        if self.use_qk_norm:
            if self.fallback_qnorm is not None and self.qnorm_weight is not None:
                self.qnorm_weight.data.copy_(self.fallback_qnorm.weight.data.float())
            if self.fallback_knorm is not None and self.knorm_weight is not None:
                self.knorm_weight.data.copy_(self.fallback_knorm.weight.data.float())

    def register_layer_name(self, layer_name: str) -> None:
        """Register layer_name and add self to the global registry.

        The global registry is needed because the bottom-level torch op
        (hpc_rope_norm_forward) is a module-level function and needs to
        route back to the correct instance via layer_name.
        """
        self.layer_name = layer_name
        _hpc_rope_norm_instances[layer_name] = self
        logger.debug(
            "[rope_norm] registered HpcRopeNorm for layer: %s",
            layer_name,
        )

    def forward_native(
        self,
        qkv: torch.Tensor,
        layer_name: str,
    ) -> torch.Tensor:
        """Native fallback path: delegates to forward_cuda().

        For now, the default native path will use CUDA backend path.
        Other platforms may override via OOT registration.
        """
        return self.forward_cuda(qkv, layer_name)

    def forward_cuda(
        self,
        qkv: torch.Tensor,
        layer_name: str,
    ) -> torch.Tensor:
        """CUDA path: invoke the torch custom op as a compile splitting point."""
        num_tokens = qkv.shape[0]
        output = torch.empty(
            (num_tokens, self.num_heads, self.head_dim),
            dtype=torch.float8_e4m3fn if self.use_fp8 else qkv.dtype,
            device=qkv.device,
        )

        torch.ops.vllm.hpc_rope_norm_forward(qkv, output, layer_name)
        return output

    def _forward_impl(
        self,
        qkv: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: HpcAttnMetadata,
        attn_layer: torch.nn.Module,
        output: torch.Tensor,
    ) -> None:
        """Actual forward logic called by the custom op.

        Writes processed q into *output* and attaches extra params
        (e.g. FP8 scales) to *attn_layer* as attributes.
        """
        import hpc

        num_actual_tokens = attn_metadata.num_actual_tokens
        num_prefill_reqs = attn_metadata.num_prefills
        num_decode_reqs = attn_metadata.num_decodes
        num_decode_tokens = attn_metadata.num_decode_tokens

        qkv = qkv[:num_actual_tokens]

        num_prefill_tokens = num_actual_tokens - num_decode_tokens

        # KV cache for the FP8 path is stored as uint8; view it as fp8 so the
        # rope_norm_store_kv_fp8 kernel can write quantized K/V in-place.
        if self.use_fp8:
            kv_cache = kv_cache.view(torch.float8_e4m3fn)

        # Per-tensor K/V scales (shape [1]) used by the FP8 kernel.
        k_scale = attn_layer._k_scale.reshape(1)
        v_scale = attn_layer._v_scale.reshape(1)

        q_norm_weight = (
            self.qnorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None
        )
        k_norm_weight = (
            self.knorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None
        )

        # --- Prefill ---
        if num_prefill_reqs > 0:
            seq_lens_prefill = attn_metadata.seq_lens[num_decode_reqs:]
            cu_seqlens_prefill = attn_metadata.qo_indptr
            max_seqlens = attn_metadata.max_query_len
            block_table_prefill = attn_metadata.block_table_tensor[num_decode_reqs:]
            qkv_prefill = qkv[num_decode_tokens:]
            out_q_prefill = output[
                num_decode_tokens : num_decode_tokens + num_prefill_tokens
            ]

            if self.use_fp8:
                _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8(
                    key_cache=kv_cache[:, 0],
                    value_cache=kv_cache[:, 1],
                    qkv=qkv_prefill,
                    cos_sin=self.cos_sin_cache,
                    num_seqlen_per_req=seq_lens_prefill,
                    q_index=cu_seqlens_prefill,
                    kvcache_indices=block_table_prefill,
                    is_prefill=True,
                    k_scale=k_scale,
                    v_scale=v_scale,
                    quant_policy=self._quant_type,
                    max_seqlens=max_seqlens,
                    q_norm_weight=q_norm_weight,
                    k_norm_weight=k_norm_weight,
                    qk_norm_policy=self.qk_norm_policy,
                    out_q=out_q_prefill,
                )
                attn_metadata.hpc_prefill_q_scale = q_scale
            else:
                hpc.rope_norm_store_kv(
                    kv_cache[:, 0],
                    kv_cache[:, 1],
                    qkv_prefill,
                    self.cos_sin_cache,
                    seq_lens_prefill,
                    cu_seqlens_prefill,
                    block_table_prefill,
                    True,  # is_prefill
                    q_norm_weight=q_norm_weight,
                    k_norm_weight=k_norm_weight,
                    out_q=out_q_prefill,
                    qk_norm_policy=self.qk_norm_policy,
                )

        # --- Decode ---
        if num_decode_reqs > 0:
            num_seq_kvcache = attn_metadata.seq_lens[:num_decode_reqs]
            block_table_decode = attn_metadata.block_table_tensor[:num_decode_reqs]
            qkv_decode = qkv[:num_decode_tokens]
            # Single-token decode: q_index is the per-request prefix sum
            # [0, 1, ..., num_decode_reqs].
            decode_query_len = attn_metadata.decode_query_len
            out_q_decode = output[:num_decode_tokens]

            if self.use_fp8:
                _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8(
                    key_cache=kv_cache[:, 0],
                    value_cache=kv_cache[:, 1],
                    qkv=qkv_decode,
                    cos_sin=self.cos_sin_cache,
                    num_seqlen_per_req=num_seq_kvcache,
                    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=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,
                    out_q=out_q_decode,
                )
                attn_metadata.hpc_decode_q_scale = q_scale
                if split_k_flag is not None:
                    attn_metadata.hpc_split_k_flag = split_k_flag
            else:
                hpc.rope_norm_store_kv(
                    kv_cache[:, 0],
                    kv_cache[:, 1],
                    qkv_decode,
                    self.cos_sin_cache,
                    num_seq_kvcache,
                    attn_metadata.qo_indptr_decode,
                    block_table_decode,
                    False,  # is_prefill
                    q_norm_weight=q_norm_weight,
                    k_norm_weight=k_norm_weight,
                    out_q=out_q_decode,
                    qk_norm_policy=self.qk_norm_policy,
                )

class QkNormPolicy(IntEnum):
    """Order of QK-RMSNorm relative to RoPE in the fused HPC rope_norm kernel.

    The values are part of the HPC kernel ABI (passed through as ints), so they
    must stay in sync with the kernel's expectations.
    """

    # No QK-Norm: apply RoPE only.
    NONE = 0
    # Apply RoPE first, then QK-RMSNorm.
    ROPE_THEN_NORM = 1
    # Apply QK-RMSNorm first, then RoPE (e.g. HunYuan V3).
    NORM_THEN_ROPE = 2