[Perf]: Fused RMSNorm+3D Rope kernel in Sensenova-U1 (#4669)
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
This commit is contained in:
boatman authored and GitHub committed
1 parent
99234881f5
commit
4028680b7e
3 files changed
+507
-4
No files matched your search
@@ -0,0 +1,232 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Correctness tests for Sensenova-U1 Triton kernels.
|
||||
|
||||
Verifies that fused-RMSNorm+3D Rope Triton kernels
|
||||
produce results numerically equivalent to the PyTorch reference.
|
||||
All computations follow the kernel's internal float32 promotion rules.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
|
||||
]
|
||||
|
||||
SEED = 42
|
||||
B = 1
|
||||
H_Q = 32
|
||||
H_K = 8
|
||||
HEAD_DIM = 128
|
||||
T_DIM = 64
|
||||
HW_DIM = 64
|
||||
HW_ROPE_DIM = 32
|
||||
QKV_DIM = (H_Q + H_K + H_K) * HEAD_DIM
|
||||
EPS = 1e-6
|
||||
DEVICE = torch.device("cuda:0")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KernelInput:
|
||||
seq_len: int
|
||||
q: torch.Tensor
|
||||
k: torch.Tensor
|
||||
q_norm_weight: torch.Tensor
|
||||
k_norm_weight: torch.Tensor
|
||||
q_norm_hw_weight: torch.Tensor
|
||||
k_norm_hw_weight: torch.Tensor
|
||||
cos_t: torch.Tensor
|
||||
sin_t: torch.Tensor
|
||||
cos_h: torch.Tensor
|
||||
sin_h: torch.Tensor
|
||||
cos_w: torch.Tensor
|
||||
sin_w: torch.Tensor
|
||||
|
||||
def args(self) -> tuple[torch.Tensor, ...]:
|
||||
return (
|
||||
self.q,
|
||||
self.k,
|
||||
self.q_norm_weight,
|
||||
self.k_norm_weight,
|
||||
self.q_norm_hw_weight,
|
||||
self.k_norm_hw_weight,
|
||||
self.cos_t,
|
||||
self.sin_t,
|
||||
self.cos_h,
|
||||
self.sin_h,
|
||||
self.cos_w,
|
||||
self.sin_w,
|
||||
)
|
||||
|
||||
|
||||
try:
|
||||
from vllm_omni.diffusion.models.sensenova_u1.fused_rmsnorm_rope import triton_qk_norm_rope
|
||||
|
||||
_TRITON_AVAILABLE = True
|
||||
except ImportError:
|
||||
_TRITON_AVAILABLE = False
|
||||
|
||||
triton_available = pytest.mark.skipif(not _TRITON_AVAILABLE, reason="Triton not available on this platform")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PyTorch reference implementations (mirror the kernel's internal arithmetic)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def rotate_half(x: torch.Tensor) -> torch.Tensor:
|
||||
x1 = x[..., : x.shape[-1] // 2]
|
||||
x2 = x[..., x.shape[-1] // 2 :]
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def apply_rope(
|
||||
q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
cos = cos.unsqueeze(1)
|
||||
sin = sin.unsqueeze(1)
|
||||
return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
|
||||
|
||||
|
||||
def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float = EPS) -> torch.Tensor:
|
||||
input_dtype = x.dtype
|
||||
xf = x.float()
|
||||
var = xf.pow(2).mean(dim=-1, keepdim=True)
|
||||
xf = xf * torch.rsqrt(var + eps)
|
||||
return xf.to(input_dtype) * weight
|
||||
|
||||
|
||||
def assert_close_with_error_stats(
|
||||
actual: torch.Tensor,
|
||||
expected: torch.Tensor,
|
||||
*,
|
||||
name: str,
|
||||
atol: float,
|
||||
rtol: float,
|
||||
) -> None:
|
||||
try:
|
||||
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
|
||||
except AssertionError as exc:
|
||||
abs_diff = (actual.float() - expected.float()).abs().flatten()
|
||||
tolerance = atol + rtol * expected.float().abs().flatten()
|
||||
mismatch = abs_diff > tolerance
|
||||
mismatch_count = int(mismatch.sum().item())
|
||||
total = abs_diff.numel()
|
||||
p99 = torch.quantile(abs_diff, 0.99).item() if total else 0.0
|
||||
stats = (
|
||||
f"{name} error stats: "
|
||||
f"max={abs_diff.max().item() if total else 0.0:.6g}, "
|
||||
f"p99={p99:.6g}, "
|
||||
f"mean={abs_diff.mean().item() if total else 0.0:.6g}, "
|
||||
f"mismatch={mismatch_count}/{total} ({mismatch_count / total:.4%}), "
|
||||
f"atol={atol}, rtol={rtol}"
|
||||
)
|
||||
raise AssertionError(f"{stats}\n{exc}") from exc
|
||||
|
||||
|
||||
def reference_kernel(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_norm_weight: torch.Tensor,
|
||||
k_norm_weight: torch.Tensor,
|
||||
q_norm_hw_weight: torch.Tensor,
|
||||
k_norm_hw_weight: torch.Tensor,
|
||||
cos_t: torch.Tensor,
|
||||
sin_t: torch.Tensor,
|
||||
cos_h: torch.Tensor,
|
||||
sin_h: torch.Tensor,
|
||||
cos_w: torch.Tensor,
|
||||
sin_w: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""PyTorch reference matching the model's current math."""
|
||||
q_t, q_hw = q.chunk(2, dim=-1)
|
||||
k_t, k_hw = k.chunk(2, dim=-1)
|
||||
|
||||
q_t = rmsnorm(q_t, q_norm_weight).transpose(1, 2)
|
||||
k_t = rmsnorm(k_t, k_norm_weight).transpose(1, 2)
|
||||
q_hw = rmsnorm(q_hw, q_norm_hw_weight).transpose(1, 2)
|
||||
k_hw = rmsnorm(k_hw, k_norm_hw_weight).transpose(1, 2)
|
||||
|
||||
q_h, q_w = q_hw.chunk(2, dim=-1)
|
||||
k_h, k_w = k_hw.chunk(2, dim=-1)
|
||||
|
||||
q_t, k_t = apply_rope(q_t, k_t, cos_t, sin_t)
|
||||
q_h, k_h = apply_rope(q_h, k_h, cos_h, sin_h)
|
||||
q_w, k_w = apply_rope(q_w, k_w, cos_w, sin_w)
|
||||
|
||||
query = torch.cat([q_t, q_h, q_w], dim=-1)
|
||||
key = torch.cat([k_t, k_h, k_w], dim=-1)
|
||||
return query, key
|
||||
|
||||
|
||||
def make_input(seq_len: int, dtype: torch.dtype, device: torch.device, seed: int) -> KernelInput:
|
||||
gen = torch.Generator(device=device)
|
||||
gen.manual_seed(seed + seq_len)
|
||||
|
||||
# Match real model layout: q/k/v are views split from fused qkv.
|
||||
qkv = torch.randn(B, seq_len, QKV_DIM, device=device, dtype=dtype, generator=gen)
|
||||
q, k, _v = qkv.split([H_Q * HEAD_DIM, H_K * HEAD_DIM, H_K * HEAD_DIM], dim=-1)
|
||||
q = q.view(B, seq_len, H_Q, HEAD_DIM)
|
||||
k = k.view(B, seq_len, H_K, HEAD_DIM)
|
||||
|
||||
q_norm_weight = torch.randn(T_DIM, device=device, dtype=dtype, generator=gen)
|
||||
k_norm_weight = torch.randn(T_DIM, device=device, dtype=dtype, generator=gen)
|
||||
q_norm_hw_weight = torch.randn(HW_DIM, device=device, dtype=dtype, generator=gen)
|
||||
k_norm_hw_weight = torch.randn(HW_DIM, device=device, dtype=dtype, generator=gen)
|
||||
|
||||
# Match Qwen3RotaryEmbedding: emb = torch.cat((freqs, freqs), dim=-1).
|
||||
def make_rope_pair(dim: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
cos_half = torch.randn(B, seq_len, dim // 2, device=device, dtype=dtype, generator=gen)
|
||||
sin_half = torch.randn(B, seq_len, dim // 2, device=device, dtype=dtype, generator=gen)
|
||||
return torch.cat((cos_half, cos_half), dim=-1), torch.cat((sin_half, sin_half), dim=-1)
|
||||
|
||||
cos_t, sin_t = make_rope_pair(T_DIM)
|
||||
cos_h, sin_h = make_rope_pair(HW_ROPE_DIM)
|
||||
cos_w, sin_w = make_rope_pair(HW_ROPE_DIM)
|
||||
|
||||
return KernelInput(
|
||||
seq_len=seq_len,
|
||||
q=q,
|
||||
k=k,
|
||||
q_norm_weight=q_norm_weight,
|
||||
k_norm_weight=k_norm_weight,
|
||||
q_norm_hw_weight=q_norm_hw_weight,
|
||||
k_norm_hw_weight=k_norm_hw_weight,
|
||||
cos_t=cos_t,
|
||||
sin_t=sin_t,
|
||||
cos_h=cos_h,
|
||||
sin_h=sin_h,
|
||||
cos_w=cos_w,
|
||||
sin_w=sin_w,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fused RMSNorm + 3D Rope
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton_available
|
||||
@pytest.mark.parametrize("seq_len", [1, 2, 4, 8, 9, 16, 32, 64, 128, 256, 260, 271, 512, 1024, 2048, 4096])
|
||||
@pytest.mark.parametrize(
|
||||
("dtype", "atol", "rtol"),
|
||||
[
|
||||
(torch.float32, 1e-5, 1e-5),
|
||||
(torch.bfloat16, 3e-2, 3e-2),
|
||||
],
|
||||
)
|
||||
def test_fused_qk_norm_rope_matches_reference(seq_len: int, dtype: torch.dtype, atol: float, rtol: float):
|
||||
data = make_input(seq_len, dtype=dtype, device=DEVICE, seed=SEED)
|
||||
call_args = data.args()
|
||||
ref_q, ref_k = reference_kernel(*call_args)
|
||||
out_q, out_k = triton_qk_norm_rope(*call_args, EPS)
|
||||
|
||||
assert_close_with_error_stats(out_q, ref_q, name="query", atol=atol, rtol=rtol)
|
||||
assert_close_with_error_stats(out_k, ref_k, name="key", atol=atol, rtol=rtol)
|
||||
@@ -0,0 +1,244 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
from vllm.triton_utils import tl, triton
|
||||
|
||||
|
||||
@triton.jit
|
||||
def qk_norm_rope_kernel(
|
||||
q,
|
||||
k,
|
||||
query,
|
||||
key,
|
||||
q_norm_weight,
|
||||
k_norm_weight,
|
||||
q_norm_hw_weight,
|
||||
k_norm_hw_weight,
|
||||
cos_t,
|
||||
sin_t,
|
||||
cos_h,
|
||||
sin_h,
|
||||
cos_w,
|
||||
sin_w,
|
||||
q_stride_b: tl.constexpr,
|
||||
q_stride_s: tl.constexpr,
|
||||
q_stride_h: tl.constexpr,
|
||||
q_stride_d: tl.constexpr,
|
||||
k_stride_b: tl.constexpr,
|
||||
k_stride_s: tl.constexpr,
|
||||
k_stride_h: tl.constexpr,
|
||||
k_stride_d: tl.constexpr,
|
||||
query_stride_b: tl.constexpr,
|
||||
query_stride_h: tl.constexpr,
|
||||
query_stride_s: tl.constexpr,
|
||||
query_stride_d: tl.constexpr,
|
||||
key_stride_b: tl.constexpr,
|
||||
key_stride_h: tl.constexpr,
|
||||
key_stride_s: tl.constexpr,
|
||||
key_stride_d: tl.constexpr,
|
||||
cos_t_stride_b: tl.constexpr,
|
||||
cos_t_stride_s: tl.constexpr,
|
||||
cos_t_stride_d: tl.constexpr,
|
||||
cos_h_stride_b: tl.constexpr,
|
||||
cos_h_stride_s: tl.constexpr,
|
||||
cos_h_stride_d: tl.constexpr,
|
||||
cos_w_stride_b: tl.constexpr,
|
||||
cos_w_stride_s: tl.constexpr,
|
||||
cos_w_stride_d: tl.constexpr,
|
||||
head_q: tl.constexpr,
|
||||
head_dim: tl.constexpr,
|
||||
eps: tl.constexpr = 1e-6,
|
||||
):
|
||||
token_idx = tl.program_id(0)
|
||||
head_pid = tl.program_id(1)
|
||||
batch_idx = tl.program_id(2)
|
||||
is_q = head_pid < head_q
|
||||
head_idx = tl.where(is_q, head_pid, head_pid - head_q)
|
||||
|
||||
offs = tl.arange(0, head_dim)
|
||||
t_mask = offs < head_dim // 2
|
||||
hw_mask = ~t_mask
|
||||
|
||||
in_base = tl.where(
|
||||
is_q,
|
||||
q + batch_idx * q_stride_b + token_idx * q_stride_s + head_idx * q_stride_h,
|
||||
k + batch_idx * k_stride_b + token_idx * k_stride_s + head_idx * k_stride_h,
|
||||
)
|
||||
vals = tl.load(in_base + offs * tl.where(is_q, q_stride_d, k_stride_d)).to(tl.float32)
|
||||
|
||||
t_vals = tl.where(t_mask, vals, 0.0)
|
||||
hw_vals = tl.where(hw_mask, vals, 0.0)
|
||||
rms_t = tl.rsqrt(tl.sum(t_vals * t_vals, axis=0) / (head_dim // 2) + eps)
|
||||
rms_hw = tl.rsqrt(tl.sum(hw_vals * hw_vals, axis=0) / (head_dim // 2) + eps)
|
||||
rms = tl.where(t_mask, rms_t, rms_hw)
|
||||
|
||||
t_weight = tl.load(tl.where(is_q, q_norm_weight, k_norm_weight) + offs, mask=t_mask, other=0.0)
|
||||
hw_weight = tl.load(
|
||||
tl.where(is_q, q_norm_hw_weight, k_norm_hw_weight) + (offs - head_dim // 2),
|
||||
mask=hw_mask,
|
||||
other=0.0,
|
||||
)
|
||||
weight = tl.where(t_mask, t_weight, hw_weight)
|
||||
normed = (vals * rms).to(query.dtype.element_ty) * weight
|
||||
|
||||
pair_offs = tl.where(
|
||||
offs < head_dim // 4,
|
||||
offs + head_dim // 4,
|
||||
tl.where(
|
||||
offs < head_dim // 2,
|
||||
offs - head_dim // 4,
|
||||
tl.where(
|
||||
offs < head_dim // 8 * 5,
|
||||
offs + head_dim // 8,
|
||||
tl.where(
|
||||
offs < head_dim // 4 * 3,
|
||||
offs - head_dim // 8,
|
||||
tl.where(offs < head_dim // 8 * 7, offs + head_dim // 8, offs - head_dim // 8),
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
pair_vals = tl.load(in_base + pair_offs * tl.where(is_q, q_stride_d, k_stride_d)).to(tl.float32)
|
||||
pair_rms = tl.where(pair_offs < head_dim // 2, rms_t, rms_hw)
|
||||
pair_t_mask = pair_offs < head_dim // 2
|
||||
pair_hw_mask = ~pair_t_mask
|
||||
pair_t_weight = tl.load(
|
||||
tl.where(is_q, q_norm_weight, k_norm_weight) + pair_offs,
|
||||
mask=pair_t_mask,
|
||||
other=0.0,
|
||||
)
|
||||
pair_hw_weight = tl.load(
|
||||
tl.where(is_q, q_norm_hw_weight, k_norm_hw_weight) + (pair_offs - head_dim // 2),
|
||||
mask=pair_hw_mask,
|
||||
other=0.0,
|
||||
)
|
||||
pair_weight = tl.where(pair_offs < head_dim // 2, pair_t_weight, pair_hw_weight)
|
||||
pair_normed = (pair_vals * pair_rms).to(query.dtype.element_ty) * pair_weight
|
||||
|
||||
cos_t_vals = tl.load(
|
||||
cos_t + batch_idx * cos_t_stride_b + token_idx * cos_t_stride_s + offs * cos_t_stride_d,
|
||||
mask=t_mask,
|
||||
other=0.0,
|
||||
)
|
||||
sin_t_vals = tl.load(
|
||||
sin_t + batch_idx * cos_t_stride_b + token_idx * cos_t_stride_s + offs * cos_t_stride_d,
|
||||
mask=t_mask,
|
||||
other=0.0,
|
||||
)
|
||||
hw_local = offs - head_dim // 2
|
||||
h_local = hw_local
|
||||
w_local = hw_local - head_dim // 4
|
||||
h_mask = (offs >= head_dim // 2) & (offs < head_dim // 4 * 3)
|
||||
w_mask = offs >= head_dim // 4 * 3
|
||||
cos_h_vals = tl.load(
|
||||
cos_h + batch_idx * cos_h_stride_b + token_idx * cos_h_stride_s + h_local * cos_h_stride_d,
|
||||
mask=h_mask,
|
||||
other=0.0,
|
||||
)
|
||||
sin_h_vals = tl.load(
|
||||
sin_h + batch_idx * cos_h_stride_b + token_idx * cos_h_stride_s + h_local * cos_h_stride_d,
|
||||
mask=h_mask,
|
||||
other=0.0,
|
||||
)
|
||||
cos_w_vals = tl.load(
|
||||
cos_w + batch_idx * cos_w_stride_b + token_idx * cos_w_stride_s + w_local * cos_w_stride_d,
|
||||
mask=w_mask,
|
||||
other=0.0,
|
||||
)
|
||||
sin_w_vals = tl.load(
|
||||
sin_w + batch_idx * cos_w_stride_b + token_idx * cos_w_stride_s + w_local * cos_w_stride_d,
|
||||
mask=w_mask,
|
||||
other=0.0,
|
||||
)
|
||||
cos_vals = cos_t_vals + cos_h_vals + cos_w_vals
|
||||
sin_vals = sin_t_vals + sin_h_vals + sin_w_vals
|
||||
|
||||
sign = tl.where(
|
||||
(offs < head_dim // 4)
|
||||
| ((offs >= head_dim // 2) & (offs < head_dim // 8 * 5))
|
||||
| ((offs >= head_dim // 4 * 3) & (offs < head_dim // 8 * 7)),
|
||||
-1.0,
|
||||
1.0,
|
||||
)
|
||||
out = normed * cos_vals + sign * pair_normed * sin_vals
|
||||
|
||||
out_base = tl.where(
|
||||
is_q,
|
||||
query + batch_idx * query_stride_b + head_idx * query_stride_h + token_idx * query_stride_s,
|
||||
key + batch_idx * key_stride_b + head_idx * key_stride_h + token_idx * key_stride_s,
|
||||
)
|
||||
tl.store(out_base + offs * tl.where(is_q, query_stride_d, key_stride_d), out)
|
||||
|
||||
|
||||
def triton_qk_norm_rope(
|
||||
q,
|
||||
k,
|
||||
q_norm_weight,
|
||||
k_norm_weight,
|
||||
q_norm_hw_weight,
|
||||
k_norm_hw_weight,
|
||||
cos_t,
|
||||
sin_t,
|
||||
cos_h,
|
||||
sin_h,
|
||||
cos_w,
|
||||
sin_w,
|
||||
eps,
|
||||
):
|
||||
batch_size, seq_len, head_q, head_dim = q.shape
|
||||
head_k = k.shape[2]
|
||||
|
||||
query = torch.empty((batch_size, head_q, seq_len, head_dim), device=q.device, dtype=q.dtype)
|
||||
key = torch.empty((batch_size, head_k, seq_len, head_dim), device=k.device, dtype=k.dtype)
|
||||
grid = (
|
||||
seq_len,
|
||||
head_q + head_k,
|
||||
batch_size,
|
||||
)
|
||||
|
||||
qk_norm_rope_kernel[grid](
|
||||
q,
|
||||
k,
|
||||
query,
|
||||
key,
|
||||
q_norm_weight,
|
||||
k_norm_weight,
|
||||
q_norm_hw_weight,
|
||||
k_norm_hw_weight,
|
||||
cos_t,
|
||||
sin_t,
|
||||
cos_h,
|
||||
sin_h,
|
||||
cos_w,
|
||||
sin_w,
|
||||
q.stride(0),
|
||||
q.stride(1),
|
||||
q.stride(2),
|
||||
q.stride(3),
|
||||
k.stride(0),
|
||||
k.stride(1),
|
||||
k.stride(2),
|
||||
k.stride(3),
|
||||
query.stride(0),
|
||||
query.stride(1),
|
||||
query.stride(2),
|
||||
query.stride(3),
|
||||
key.stride(0),
|
||||
key.stride(1),
|
||||
key.stride(2),
|
||||
key.stride(3),
|
||||
cos_t.stride(0),
|
||||
cos_t.stride(1),
|
||||
cos_t.stride(2),
|
||||
cos_h.stride(0),
|
||||
cos_h.stride(1),
|
||||
cos_h.stride(2),
|
||||
cos_w.stride(0),
|
||||
cos_w.stride(1),
|
||||
cos_w.stride(2),
|
||||
head_q=head_q,
|
||||
head_dim=head_dim,
|
||||
eps=eps,
|
||||
)
|
||||
return query, key
|
||||
@@ -18,6 +18,7 @@ import torch
|
||||
import torch.nn as nn
|
||||
from cache_dit import ForwardPattern
|
||||
from transformers.cache_utils import DynamicCache
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.linear import (
|
||||
MergedColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
@@ -33,6 +34,8 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, SensenovaCachedAdapter
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SenseNovaU1ModelOutput:
|
||||
@@ -373,6 +376,34 @@ class SenseNovaU1Attention(nn.Module):
|
||||
k = k.view(*input_shape, self.num_kv_heads, self.head_dim)
|
||||
v = v.view(*input_shape, self.num_kv_heads, self.head_dim)
|
||||
|
||||
cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0))
|
||||
cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0))
|
||||
cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0))
|
||||
|
||||
value_states = v.transpose(1, 2) # [B, H, S, D]
|
||||
try:
|
||||
from .fused_rmsnorm_rope import triton_qk_norm_rope
|
||||
except ImportError:
|
||||
triton_qk_norm_rope = None
|
||||
if triton_qk_norm_rope is not None:
|
||||
query_states, key_states = triton_qk_norm_rope(
|
||||
q,
|
||||
k,
|
||||
q_norm.weight,
|
||||
k_norm.weight,
|
||||
q_norm_hw.weight,
|
||||
k_norm_hw.weight,
|
||||
cos_t,
|
||||
sin_t,
|
||||
cos_h,
|
||||
sin_h,
|
||||
cos_w,
|
||||
sin_w,
|
||||
self.config.rms_norm_eps,
|
||||
)
|
||||
return query_states, key_states, value_states
|
||||
logger.debug("Triton fused SenseNova qk norm rope unavailable; using PyTorch fallback")
|
||||
|
||||
# Split head_dim into t and hw halves
|
||||
q_t, q_hw = q.chunk(2, dim=-1)
|
||||
k_t, k_hw = k.chunk(2, dim=-1)
|
||||
@@ -388,17 +419,13 @@ class SenseNovaU1Attention(nn.Module):
|
||||
k_h, k_w = k_hw.chunk(2, dim=-1)
|
||||
|
||||
# RoPE for each dimension
|
||||
cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0))
|
||||
q_t, k_t = apply_rotary_pos_emb(q_t, k_t, cos_t, sin_t)
|
||||
cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0))
|
||||
q_h, k_h = apply_rotary_pos_emb(q_h, k_h, cos_h, sin_h)
|
||||
cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0))
|
||||
q_w, k_w = apply_rotary_pos_emb(q_w, k_w, cos_w, sin_w)
|
||||
|
||||
# Reassemble: [B, H, S, head_dim]
|
||||
query_states = torch.cat([q_t, q_h, q_w], dim=-1)
|
||||
key_states = torch.cat([k_t, k_h, k_w], dim=-1)
|
||||
value_states = v.transpose(1, 2) # [B, H, S, D]
|
||||
return query_states, key_states, value_states
|
||||
|
||||
def forward_und(self, hidden_states, indexes, attention_mask, past_key_values=None, **kwargs):
|
||||
|
||||
Reference in new issue
Block a user