[Diffusion][Perf] Add Qwen-Image QK RoPE Triton path (#5931)
Signed-off-by: dongbo910220 <1275604947@qq.com>
This commit is contained in:
dongbo910220 authored and GitHub committed
1 parent
a3dee6fdf5
commit
952022c8ec
3 files changed
+484
-22
No files matched your search
@@ -0,0 +1,224 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Benchmark Qwen-Image Q/K RMSNorm + interleaved RoPE.
|
||||
|
||||
Example:
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0 python benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py \
|
||||
--seq-len 4096 --warmup 20 --iters 100 --include-compiled --profile
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
|
||||
_fused_cuda_supported,
|
||||
fused_qk_norm_rope,
|
||||
)
|
||||
from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
|
||||
_apply_qwen_image_rotary_emb,
|
||||
)
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--batch", type=int, default=1)
|
||||
parser.add_argument("--seq-len", type=int, default=4096)
|
||||
parser.add_argument("--heads", type=int, default=24)
|
||||
parser.add_argument("--kv-heads", type=int, default=24)
|
||||
parser.add_argument("--head-dim", type=int, default=128)
|
||||
parser.add_argument("--warmup", type=int, default=20)
|
||||
parser.add_argument("--iters", type=int, default=100)
|
||||
parser.add_argument("--dtype", choices=("bf16", "fp16", "fp32"), default="bf16")
|
||||
parser.add_argument("--packed-qkv-view", action="store_true")
|
||||
parser.add_argument("--include-compiled", action="store_true")
|
||||
parser.add_argument("--profile", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _dtype(name: str) -> torch.dtype:
|
||||
return {
|
||||
"bf16": torch.bfloat16,
|
||||
"fp16": torch.float16,
|
||||
"fp32": torch.float32,
|
||||
}[name]
|
||||
|
||||
|
||||
def _make_inputs(args: argparse.Namespace, dtype: torch.dtype, device: torch.device):
|
||||
torch.manual_seed(2026)
|
||||
if args.packed_qkv_view:
|
||||
qkv_dim = (args.heads + args.kv_heads + args.kv_heads) * args.head_dim
|
||||
qkv = torch.randn(args.batch, args.seq_len, qkv_dim, device=device, dtype=dtype)
|
||||
q, k, _ = qkv.split(
|
||||
[
|
||||
args.heads * args.head_dim,
|
||||
args.kv_heads * args.head_dim,
|
||||
args.kv_heads * args.head_dim,
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
q = q.unflatten(-1, (args.heads, args.head_dim))
|
||||
k = k.unflatten(-1, (args.kv_heads, args.head_dim))
|
||||
else:
|
||||
q = torch.randn(args.batch, args.seq_len, args.heads, args.head_dim, device=device, dtype=dtype)
|
||||
k = torch.randn(args.batch, args.seq_len, args.kv_heads, args.head_dim, device=device, dtype=dtype)
|
||||
|
||||
q_weight = torch.randn(args.head_dim, device=device, dtype=dtype)
|
||||
k_weight = torch.randn(args.head_dim, device=device, dtype=dtype)
|
||||
angles = torch.randn(args.seq_len, args.head_dim // 2, device=device, dtype=torch.float32)
|
||||
freqs = torch.polar(torch.ones_like(angles), angles)
|
||||
return q, k, q_weight, k_weight, freqs
|
||||
|
||||
|
||||
def _native(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
freqs: torch.Tensor,
|
||||
eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
q = F.rms_norm(q, (q.shape[-1],), q_weight, eps)
|
||||
k = F.rms_norm(k, (k.shape[-1],), k_weight, eps)
|
||||
return (
|
||||
_apply_qwen_image_rotary_emb(q, freqs),
|
||||
_apply_qwen_image_rotary_emb(k, freqs),
|
||||
)
|
||||
|
||||
|
||||
def _fused(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
freqs: torch.Tensor,
|
||||
eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
batch, seq_len, num_heads, head_dim = q.shape
|
||||
num_kv_heads = k.shape[2]
|
||||
rope_table = torch.cat((freqs.real, freqs.imag), dim=-1)
|
||||
rope_table = rope_table.unsqueeze(0).expand(batch, -1, -1).reshape(batch * seq_len, head_dim)
|
||||
out_q, out_k = fused_qk_norm_rope(
|
||||
q.reshape(batch * seq_len, num_heads, head_dim),
|
||||
k.reshape(batch * seq_len, num_kv_heads, head_dim),
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope_table,
|
||||
eps,
|
||||
interleaved=True,
|
||||
)
|
||||
return (
|
||||
out_q.reshape(batch, seq_len, num_heads, head_dim),
|
||||
out_k.reshape(batch, seq_len, num_kv_heads, head_dim),
|
||||
)
|
||||
|
||||
|
||||
def _measure(fn: Callable[[], tuple[torch.Tensor, torch.Tensor]], warmup: int, iters: int) -> list[float]:
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.accelerator.synchronize()
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
samples = []
|
||||
for _ in range(iters):
|
||||
start.record()
|
||||
fn()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
samples.append(start.elapsed_time(end))
|
||||
return samples
|
||||
|
||||
|
||||
def _stats(samples: list[float]) -> dict[str, float]:
|
||||
samples = sorted(samples)
|
||||
return {
|
||||
"median_ms": samples[len(samples) // 2],
|
||||
"p90_ms": samples[int(len(samples) * 0.9)],
|
||||
"mean_ms": sum(samples) / len(samples),
|
||||
}
|
||||
|
||||
|
||||
def _profile(fn: Callable[[], tuple[torch.Tensor, torch.Tensor]]) -> str:
|
||||
from torch.profiler import ProfilerActivity, profile
|
||||
|
||||
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
|
||||
fn()
|
||||
torch.accelerator.synchronize()
|
||||
return prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=15)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("This benchmark requires CUDA")
|
||||
device = torch.device("cuda")
|
||||
dtype = _dtype(args.dtype)
|
||||
eps = 1e-6
|
||||
q, k, q_weight, k_weight, freqs = _make_inputs(args, dtype, device)
|
||||
|
||||
def native():
|
||||
return _native(q, k, q_weight, k_weight, freqs, eps)
|
||||
|
||||
def fused():
|
||||
return _fused(q, k, q_weight, k_weight, freqs, eps)
|
||||
|
||||
results: dict[str, object] = {
|
||||
"device": torch.cuda.get_device_name(),
|
||||
"batch": args.batch,
|
||||
"seq_len": args.seq_len,
|
||||
"heads": args.heads,
|
||||
"kv_heads": args.kv_heads,
|
||||
"head_dim": args.head_dim,
|
||||
"dtype": args.dtype,
|
||||
"packed_qkv_view": args.packed_qkv_view,
|
||||
"fused_fast_path_supported": _fused_cuda_supported(
|
||||
q.reshape(args.batch * args.seq_len, args.heads, args.head_dim),
|
||||
k.reshape(args.batch * args.seq_len, args.kv_heads, args.head_dim),
|
||||
args.head_dim,
|
||||
args.head_dim,
|
||||
interleaved=True,
|
||||
),
|
||||
}
|
||||
|
||||
native_samples = _measure(native, args.warmup, args.iters)
|
||||
fused_samples = _measure(fused, args.warmup, args.iters)
|
||||
expected_q, expected_k = native()
|
||||
actual_q, actual_k = fused()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
native_stats = _stats(native_samples)
|
||||
fused_stats = _stats(fused_samples)
|
||||
results["native"] = native_stats
|
||||
results["fused"] = fused_stats
|
||||
results["speedup_vs_native"] = native_stats["median_ms"] / fused_stats["median_ms"]
|
||||
results["max_abs_q"] = (actual_q.float() - expected_q.float()).abs().max().item()
|
||||
results["max_abs_k"] = (actual_k.float() - expected_k.float()).abs().max().item()
|
||||
results["mean_abs_q"] = (actual_q.float() - expected_q.float()).abs().mean().item()
|
||||
results["mean_abs_k"] = (actual_k.float() - expected_k.float()).abs().mean().item()
|
||||
|
||||
if args.include_compiled:
|
||||
compiled_native = torch.compile(native, dynamic=True, fullgraph=True)
|
||||
compiled_fused = torch.compile(fused, dynamic=True, fullgraph=True)
|
||||
compiled_native_stats = _stats(_measure(compiled_native, args.warmup, args.iters))
|
||||
compiled_fused_stats = _stats(_measure(compiled_fused, args.warmup, args.iters))
|
||||
results["compiled_native"] = compiled_native_stats
|
||||
results["compiled_fused"] = compiled_fused_stats
|
||||
results["compiled_fused_speedup_vs_compiled_native"] = (
|
||||
compiled_native_stats["median_ms"] / compiled_fused_stats["median_ms"]
|
||||
)
|
||||
|
||||
print(json.dumps(results, indent=2))
|
||||
if args.profile:
|
||||
print("\n# fused profile")
|
||||
print(_profile(fused))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,190 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
|
||||
from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
|
||||
_apply_qwen_image_rotary_emb,
|
||||
_qwen_image_qk_norm_rope,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
BATCH = 2
|
||||
Q_HEADS = 24
|
||||
K_HEADS = 24
|
||||
HEAD_DIM = 128
|
||||
EPS = 1e-6
|
||||
SEED = 2026
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QwenImageQKInput:
|
||||
q: torch.Tensor
|
||||
k: torch.Tensor
|
||||
norm_q: nn.RMSNorm
|
||||
norm_k: nn.RMSNorm
|
||||
freqs: torch.Tensor
|
||||
|
||||
|
||||
def _make_input(
|
||||
*,
|
||||
seq_len: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
packed_qkv_view: bool,
|
||||
) -> QwenImageQKInput:
|
||||
gen = torch.Generator(device=device)
|
||||
gen.manual_seed(SEED + seq_len)
|
||||
|
||||
if packed_qkv_view:
|
||||
qkv_dim = (Q_HEADS + K_HEADS + K_HEADS) * HEAD_DIM
|
||||
qkv = torch.randn(BATCH, seq_len, qkv_dim, device=device, dtype=dtype, generator=gen)
|
||||
q, k, _ = qkv.split(
|
||||
[
|
||||
Q_HEADS * HEAD_DIM,
|
||||
K_HEADS * HEAD_DIM,
|
||||
K_HEADS * HEAD_DIM,
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
q = q.unflatten(-1, (Q_HEADS, HEAD_DIM))
|
||||
k = k.unflatten(-1, (K_HEADS, HEAD_DIM))
|
||||
else:
|
||||
q = torch.randn(BATCH, seq_len, Q_HEADS, HEAD_DIM, device=device, dtype=dtype, generator=gen)
|
||||
k = torch.randn(BATCH, seq_len, K_HEADS, HEAD_DIM, device=device, dtype=dtype, generator=gen)
|
||||
|
||||
norm_q = nn.RMSNorm(HEAD_DIM, eps=EPS, device=device, dtype=dtype)
|
||||
norm_k = nn.RMSNorm(HEAD_DIM, eps=EPS, device=device, dtype=dtype)
|
||||
norm_q.requires_grad_(False)
|
||||
norm_k.requires_grad_(False)
|
||||
with torch.no_grad():
|
||||
norm_q.weight.copy_(torch.randn(HEAD_DIM, device=device, dtype=dtype, generator=gen))
|
||||
norm_k.weight.copy_(torch.randn(HEAD_DIM, device=device, dtype=dtype, generator=gen))
|
||||
angles = torch.randn(seq_len, HEAD_DIM // 2, device=device, dtype=torch.float32, generator=gen)
|
||||
freqs = torch.polar(torch.ones_like(angles), angles)
|
||||
return QwenImageQKInput(q=q, k=k, norm_q=norm_q, norm_k=norm_k, freqs=freqs)
|
||||
|
||||
|
||||
def _reference(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
q = data.norm_q(data.q)
|
||||
k = data.norm_k(data.k)
|
||||
if data.q.device.type == "cuda":
|
||||
return (
|
||||
_apply_qwen_image_rotary_emb(q, data.freqs),
|
||||
_apply_qwen_image_rotary_emb(k, data.freqs),
|
||||
)
|
||||
|
||||
rope = RotaryEmbedding(is_neox_style=False)
|
||||
cos = data.freqs.real.to(data.q.dtype)
|
||||
sin = data.freqs.imag.to(data.q.dtype)
|
||||
return rope(q, cos, sin), rope(k, cos, sin)
|
||||
|
||||
|
||||
def _run(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return _qwen_image_qk_norm_rope(
|
||||
data.q,
|
||||
data.k,
|
||||
data.norm_q,
|
||||
data.norm_k,
|
||||
data.freqs,
|
||||
RotaryEmbedding(is_neox_style=False),
|
||||
EPS,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
def test_qwen_image_qk_norm_rope_cuda_fp32_fallback_matches_reference():
|
||||
data = _make_input(
|
||||
seq_len=7,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cuda:0"),
|
||||
packed_qkv_view=True,
|
||||
)
|
||||
|
||||
actual_q, actual_k = _run(data)
|
||||
expected_q, expected_k = _reference(data)
|
||||
|
||||
torch.testing.assert_close(actual_q, expected_q, atol=1e-5, rtol=1e-5)
|
||||
torch.testing.assert_close(actual_k, expected_k, atol=1e-5, rtol=1e-5)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
def test_qwen_image_qk_norm_rope_cuda_fp16_fallback_matches_reference():
|
||||
data = _make_input(
|
||||
seq_len=7,
|
||||
dtype=torch.float16,
|
||||
device=torch.device("cuda:0"),
|
||||
packed_qkv_view=True,
|
||||
)
|
||||
|
||||
actual_q, actual_k = _run(data)
|
||||
expected_q, expected_k = _reference(data)
|
||||
|
||||
torch.testing.assert_close(actual_q, expected_q, atol=1e-3, rtol=1e-3)
|
||||
torch.testing.assert_close(actual_k, expected_k, atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
@pytest.mark.parametrize("seq_len", [1, 7, 257, 1024])
|
||||
@pytest.mark.parametrize("packed_qkv_view", [False, True])
|
||||
def test_qwen_image_fused_qk_norm_rope_cuda_matches_fp32_rope_reference(
|
||||
seq_len: int,
|
||||
packed_qkv_view: bool,
|
||||
):
|
||||
data = _make_input(
|
||||
seq_len=seq_len,
|
||||
dtype=torch.bfloat16,
|
||||
device=torch.device("cuda:0"),
|
||||
packed_qkv_view=packed_qkv_view,
|
||||
)
|
||||
|
||||
actual_q, actual_k = _run(data)
|
||||
expected_q, expected_k = _reference(data)
|
||||
|
||||
torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
|
||||
torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
def test_qwen_image_fused_qk_norm_rope_torch_compile_fullgraph_capture():
|
||||
data = _make_input(
|
||||
seq_len=257,
|
||||
dtype=torch.bfloat16,
|
||||
device=torch.device("cuda:0"),
|
||||
packed_qkv_view=True,
|
||||
)
|
||||
rope = RotaryEmbedding(is_neox_style=False)
|
||||
|
||||
def fn(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
freqs: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return _qwen_image_qk_norm_rope(
|
||||
q,
|
||||
k,
|
||||
data.norm_q,
|
||||
data.norm_k,
|
||||
freqs,
|
||||
rope,
|
||||
EPS,
|
||||
)
|
||||
|
||||
compiled_fn = torch.compile(fn, dynamic=True, fullgraph=True)
|
||||
expected_q, expected_k = fn(data.q, data.k, data.freqs)
|
||||
actual_q, actual_k = compiled_fn(data.q, data.k, data.freqs)
|
||||
|
||||
torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
|
||||
torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)
|
||||
@@ -45,6 +45,10 @@ from vllm_omni.diffusion.distributed.sp_plan import (
|
||||
)
|
||||
from vllm_omni.diffusion.forward_context import get_forward_context
|
||||
from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
|
||||
from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
|
||||
_fused_cuda_supported,
|
||||
fused_qk_norm_rope,
|
||||
)
|
||||
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -61,6 +65,53 @@ def _apply_qwen_image_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch.
|
||||
return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype)
|
||||
|
||||
|
||||
def _qwen_image_qk_norm_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
norm_q: nn.Module,
|
||||
norm_k: nn.Module,
|
||||
freqs: torch.Tensor,
|
||||
rope: RotaryEmbedding,
|
||||
eps: float,
|
||||
*,
|
||||
use_fused: bool = True,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
head_dim = q.shape[-1]
|
||||
rotary_dim = freqs.shape[-1] * 2
|
||||
if use_fused and _fused_cuda_supported(q, k, head_dim, rotary_dim, interleaved=True):
|
||||
batch, seq_len, num_heads, _ = q.shape
|
||||
num_kv_heads = k.shape[2]
|
||||
rope_table = torch.cat((freqs.real, freqs.imag), dim=-1)
|
||||
rope_table = rope_table.unsqueeze(0).expand(batch, -1, -1).reshape(batch * seq_len, rotary_dim)
|
||||
fused_q, fused_k = fused_qk_norm_rope(
|
||||
q.reshape(batch * seq_len, num_heads, head_dim),
|
||||
k.reshape(batch * seq_len, num_kv_heads, head_dim),
|
||||
norm_q.weight,
|
||||
norm_k.weight,
|
||||
rope_table,
|
||||
eps,
|
||||
interleaved=True,
|
||||
)
|
||||
return (
|
||||
fused_q.reshape(batch, seq_len, num_heads, head_dim),
|
||||
fused_k.reshape(batch, seq_len, num_kv_heads, head_dim),
|
||||
)
|
||||
|
||||
q = norm_q(q)
|
||||
k = norm_k(k)
|
||||
if q.device.type == "cuda":
|
||||
return (
|
||||
_apply_qwen_image_rotary_emb(q, freqs),
|
||||
_apply_qwen_image_rotary_emb(k, freqs),
|
||||
)
|
||||
|
||||
# Retain the platform-specific kernels on other accelerators, which may
|
||||
# not support complex tensors.
|
||||
cos = torch.real(freqs).to(q.dtype)
|
||||
sin = torch.imag(freqs).to(q.dtype)
|
||||
return rope(q, cos, sin), rope(k, cos, sin)
|
||||
|
||||
|
||||
def _normalize_qwen_image_weight_name(name: str) -> str:
|
||||
name = name.removeprefix("transformer.")
|
||||
if ".to_out.0." in name:
|
||||
@@ -648,28 +699,25 @@ class QwenImageCrossAttention(nn.Module):
|
||||
txt_key = txt_key.unflatten(-1, (self.add_kv_num_heads, self.head_dim))
|
||||
txt_value = txt_value.unflatten(-1, (self.add_kv_num_heads, self.head_dim))
|
||||
|
||||
img_query = self.norm_q(img_query)
|
||||
img_key = self.norm_k(img_key)
|
||||
txt_query = self.norm_added_q(txt_query)
|
||||
txt_key = self.norm_added_k(txt_key)
|
||||
|
||||
if img_query.device.type == "cuda":
|
||||
img_query = _apply_qwen_image_rotary_emb(img_query, vid_freqs)
|
||||
img_key = _apply_qwen_image_rotary_emb(img_key, vid_freqs)
|
||||
txt_query = _apply_qwen_image_rotary_emb(txt_query, txt_freqs)
|
||||
txt_key = _apply_qwen_image_rotary_emb(txt_key, txt_freqs)
|
||||
else:
|
||||
# Retain the platform-specific kernels on other accelerators,
|
||||
# which may not support complex tensors.
|
||||
img_cos = torch.real(vid_freqs).to(img_query.dtype)
|
||||
img_sin = torch.imag(vid_freqs).to(img_query.dtype)
|
||||
txt_cos = torch.real(txt_freqs).to(txt_query.dtype)
|
||||
txt_sin = torch.imag(txt_freqs).to(txt_query.dtype)
|
||||
|
||||
img_query = self.rope(img_query, img_cos, img_sin)
|
||||
img_key = self.rope(img_key, img_cos, img_sin)
|
||||
txt_query = self.rope(txt_query, txt_cos, txt_sin)
|
||||
txt_key = self.rope(txt_key, txt_cos, txt_sin)
|
||||
img_query, img_key = _qwen_image_qk_norm_rope(
|
||||
img_query,
|
||||
img_key,
|
||||
self.norm_q,
|
||||
self.norm_k,
|
||||
vid_freqs,
|
||||
self.rope,
|
||||
self.eps,
|
||||
use_fused=self.qk_norm,
|
||||
)
|
||||
txt_query, txt_key = _qwen_image_qk_norm_rope(
|
||||
txt_query,
|
||||
txt_key,
|
||||
self.norm_added_q,
|
||||
self.norm_added_k,
|
||||
txt_freqs,
|
||||
self.rope,
|
||||
self.eps,
|
||||
)
|
||||
|
||||
seq_len_txt = encoder_hidden_states.shape[1]
|
||||
joint_query = torch.cat([txt_query, img_query], dim=1)
|
||||
|
||||
Reference in new issue
Block a user