Zhenqiang Xuandzqxu
242c3cddeb
Signed-off-by: zqxu <xuzhenqiang3@huawei.com> Co-authored-by: zqxu <xuzhenqiang3@huawei.com>
26 lines
1.1 KiB
Python
26 lines
1.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Seeded rotation matrices for MindIE-SD quantized attention."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
import math
|
|
|
|
import torch
|
|
|
|
|
|
@functools.lru_cache(maxsize=128)
|
|
def get_quant_attention_rotation(device: torch.device, dtype: torch.dtype, head_dim: int) -> torch.Tensor:
|
|
"""Preserve Omni's fixed Walsh-Hadamard rotation without changing global RNG."""
|
|
if head_dim <= 0 or head_dim & (head_dim - 1):
|
|
raise ValueError("Generated attention rotations require a power-of-two head dimension.")
|
|
generator = torch.Generator(device="cpu").manual_seed(425500)
|
|
signs = torch.randint(0, 2, (head_dim,), generator=generator, dtype=torch.int64)
|
|
signs = signs.to(dtype=torch.float32).mul_(2).sub_(1)
|
|
hadamard = torch.ones(1, 1)
|
|
while hadamard.shape[0] < head_dim:
|
|
hadamard = torch.cat((torch.cat((hadamard, hadamard), dim=1), torch.cat((hadamard, -hadamard), dim=1)), dim=0)
|
|
rotation = signs[:, None] * hadamard / math.sqrt(head_dim)
|
|
return rotation.to(device=device, dtype=dtype).contiguous()
|