minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[KDA-Pilot] Add LTX2 QKNorm split-RoPE CUDA fast path (#29708)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-07-01 14:42:07 +08:00
committed by GitHub
parent 41779d56fd
commit fcb9f229b3
5 changed files with 1006 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/jit/benchmark/diffusion/bench_ltx2_qknorm_split_rope.py
+209
View File
@@ -0,0 +1,209 @@
import random
import sys
from dataclasses import dataclass
import torch
from sglang.jit_kernel.diffusion.ltx2_qknorm_split_rope import (
ltx2_qknorm_split_rope_cuda,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(
est_time=30,
stage="base-b-kernel-benchmark",
runner_config="1-gpu-large",
disabled="standalone benchmark",
)
@dataclass(frozen=True)
class Workload:
name: str
batch: int
q_seq: int
k_seq: int
num_heads: int
head_dim: int
FULL_WORKLOADS = [
Workload("stage1_video_self_q1536_k1536_d4096", 2, 1536, 1536, 32, 128),
Workload("stage1_audio_self_q126_k126_d2048", 2, 126, 126, 32, 64),
Workload("stage1_audio_to_video_q1536_k126_d2048", 2, 1536, 126, 32, 64),
Workload("stage1_video_to_audio_q126_k1536_d2048", 2, 126, 1536, 32, 64),
Workload("stage2_video_self_q6144_k6144_d4096", 1, 6144, 6144, 32, 128),
Workload("stage2_audio_self_q126_k126_d2048", 1, 126, 126, 32, 64),
Workload("stage2_audio_to_video_q6144_k126_d2048", 1, 6144, 126, 32, 64),
Workload("stage2_video_to_audio_q126_k6144_d2048", 1, 126, 6144, 32, 64),
Workload("hq_stage1_video_self_q8160_k8160_d4096", 1, 8160, 8160, 32, 128),
Workload("hq_stage1_audio_to_video_q8160_k126_d2048", 1, 8160, 126, 32, 64),
Workload("hq_stage1_video_to_audio_q126_k8160_d2048", 1, 126, 8160, 32, 64),
Workload("hq_stage2_video_self_q32640_k32640_d4096", 1, 32640, 32640, 32, 128),
Workload("hq_stage2_audio_to_video_q32640_k126_d2048", 1, 32640, 126, 32, 64),
Workload("hq_stage2_video_to_audio_q126_k32640_d2048", 1, 126, 32640, 32, 64),
]
CI_WORKLOADS = [
Workload("stage1_video_self_q16_k16_d4096", 1, 16, 16, 32, 128),
Workload("stage1_audio_to_video_q16_k8_d2048", 1, 16, 8, 32, 64),
]
def _make_cos_sin(
batch: int, seq_len: int, num_heads: int, head_dim: int
) -> tuple[torch.Tensor, torch.Tensor]:
half_dim = head_dim // 2
cos = torch.randn(
batch, seq_len, num_heads, half_dim, device="cuda", dtype=torch.bfloat16
).transpose(1, 2)
sin = torch.randn(
batch, seq_len, num_heads, half_dim, device="cuda", dtype=torch.bfloat16
).transpose(1, 2)
return cos, sin
def _apply_split_rotary_ref(
x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> torch.Tensor:
x_dtype = x.dtype
batch = x.shape[0]
_, num_heads, seq_len, _ = cos.shape
x = x.reshape(batch, seq_len, num_heads, -1).swapaxes(1, 2)
last = x.shape[-1]
half = last // 2
split_x = x.reshape(*x.shape[:-1], 2, half)
first_x = split_x[..., :1, :]
second_x = split_x[..., 1:, :]
cos_u = cos.unsqueeze(-2)
sin_u = sin.unsqueeze(-2)
out = split_x * cos_u
out[..., :1, :].addcmul_(-sin_u, second_x)
out[..., 1:, :].addcmul_(sin_u, first_x)
out = out.reshape(*out.shape[:-2], last)
return out.swapaxes(1, 2).reshape(batch, seq_len, -1).to(dtype=x_dtype)
def _reference_pair(inputs):
q, k, q_cos, q_sin, k_cos, k_sin, q_norm, k_norm = inputs
with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
q_out = _apply_split_rotary_ref(q_norm(q), q_cos, q_sin)
k_out = _apply_split_rotary_ref(k_norm(k), k_cos, k_sin)
return q_out.to(dtype=torch.bfloat16), k_out.to(dtype=torch.bfloat16)
def cuda_event_us(fn, warmups: int, repeats: int, rounds: int) -> float:
for _ in range(warmups):
fn()
torch.cuda.synchronize()
samples = []
for _ in range(rounds):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(repeats):
fn()
end.record()
end.synchronize()
samples.append(start.elapsed_time(end) * 1000.0 / repeats)
samples.sort()
return samples[len(samples) // 2]
def benchmark() -> None:
if not torch.cuda.is_available():
print("CUDA required")
return
torch.manual_seed(20260630)
random.seed(20260630)
torch.cuda.set_device(0)
workloads = CI_WORKLOADS if is_in_ci() else FULL_WORKLOADS
warmups = 3 if is_in_ci() else 10
repeats = 3 if is_in_ci() else 10
rounds = 3 if is_in_ci() else 7
print("| workload | torch us | cuda us | speedup |")
print("|---|---:|---:|---:|")
for workload in workloads:
hidden = workload.num_heads * workload.head_dim
q = torch.randn(
workload.batch,
workload.q_seq,
hidden,
device="cuda",
dtype=torch.bfloat16,
)
k = torch.randn(
workload.batch,
workload.k_seq,
hidden,
device="cuda",
dtype=torch.bfloat16,
)
q_cos, q_sin = _make_cos_sin(
workload.batch, workload.q_seq, workload.num_heads, workload.head_dim
)
k_cos, k_sin = _make_cos_sin(
workload.batch, workload.k_seq, workload.num_heads, workload.head_dim
)
q_norm = torch.nn.RMSNorm(hidden, eps=1e-6, device="cuda").to(
dtype=torch.bfloat16
)
k_norm = torch.nn.RMSNorm(hidden, eps=1e-6, device="cuda").to(
dtype=torch.bfloat16
)
inputs = (q, k, q_cos, q_sin, k_cos, k_sin, q_norm, k_norm)
q_ref, k_ref = _reference_pair(inputs)
q_out, k_out = ltx2_qknorm_split_rope_cuda(
q,
q_cos,
q_sin,
q_norm.weight,
k,
k_cos,
k_sin,
k_norm.weight,
eps=1e-6,
num_heads=workload.num_heads,
head_dim=workload.head_dim,
)
torch.cuda.synchronize()
assert torch.equal(q_ref, q_out)
assert torch.equal(k_ref, k_out)
fns = {
"torch": lambda: _reference_pair(inputs),
"cuda": lambda: ltx2_qknorm_split_rope_cuda(
q,
q_cos,
q_sin,
q_norm.weight,
k,
k_cos,
k_sin,
k_norm.weight,
eps=1e-6,
num_heads=workload.num_heads,
head_dim=workload.head_dim,
),
}
order = ["torch", "cuda"]
random.shuffle(order)
times = {
name: cuda_event_us(fns[name], warmups, repeats, rounds) for name in order
}
print(
f"| {workload.name} | {times['torch']:.2f} | "
f"{times['cuda']:.2f} | {times['torch'] / times['cuda']:.3f}x |"
)
torch.cuda.empty_cache()
if __name__ == "__main__":
benchmark()
sys.exit(0)
test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py
+223
View File
@@ -0,0 +1,223 @@
import sys
import pytest
import torch
from sglang.jit_kernel.diffusion.ltx2_qknorm_split_rope import (
can_use_ltx2_qknorm_split_rope_cuda,
ltx2_qknorm_split_rope_cuda,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
def _require_cuda_b200() -> None:
if not torch.cuda.is_available():
pytest.skip("CUDA required")
if torch.cuda.get_device_capability()[0] < 10:
pytest.skip("LTX2 QKNorm split-RoPE CUDA path is validated on B200")
@pytest.fixture(autouse=True)
def cuda_setup():
_require_cuda_b200()
torch.cuda.manual_seed(20260630)
def _make_cos_sin(
batch: int, seq_len: int, num_heads: int, head_dim: int
) -> tuple[torch.Tensor, torch.Tensor]:
half_dim = head_dim // 2
cos = torch.randn(
batch, seq_len, num_heads, half_dim, device="cuda", dtype=torch.bfloat16
).transpose(1, 2)
sin = torch.randn(
batch, seq_len, num_heads, half_dim, device="cuda", dtype=torch.bfloat16
).transpose(1, 2)
return cos, sin
def _apply_split_rotary_ref(
x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> torch.Tensor:
x_dtype = x.dtype
batch = x.shape[0]
_, num_heads, seq_len, _ = cos.shape
x = x.reshape(batch, seq_len, num_heads, -1).swapaxes(1, 2)
last = x.shape[-1]
half = last // 2
split_x = x.reshape(*x.shape[:-1], 2, half)
first_x = split_x[..., :1, :]
second_x = split_x[..., 1:, :]
cos_u = cos.unsqueeze(-2)
sin_u = sin.unsqueeze(-2)
out = split_x * cos_u
out[..., :1, :].addcmul_(-sin_u, second_x)
out[..., 1:, :].addcmul_(sin_u, first_x)
out = out.reshape(*out.shape[:-2], last)
return out.swapaxes(1, 2).reshape(batch, seq_len, -1).to(dtype=x_dtype)
def _reference(
q: torch.Tensor,
k: torch.Tensor,
q_cos: torch.Tensor,
q_sin: torch.Tensor,
k_cos: torch.Tensor,
k_sin: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
q_norm = torch.nn.RMSNorm(q.shape[-1], eps=eps, device="cuda").to(
dtype=torch.bfloat16
)
k_norm = torch.nn.RMSNorm(k.shape[-1], eps=eps, device="cuda").to(
dtype=torch.bfloat16
)
q_norm.weight.data.copy_(q_weight)
k_norm.weight.data.copy_(k_weight)
with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
q_ref = _apply_split_rotary_ref(q_norm(q), q_cos, q_sin)
k_ref = _apply_split_rotary_ref(k_norm(k), k_cos, k_sin)
return q_ref.to(dtype=torch.bfloat16), k_ref.to(dtype=torch.bfloat16)
@pytest.mark.parametrize(
"batch,q_seq,k_seq,num_heads,head_dim",
[
(1, 3, 3, 32, 128),
(1, 5, 2, 32, 64),
(2, 4, 3, 32, 64),
],
)
def test_ltx2_qknorm_split_rope_matches_torch_exactly(
batch: int, q_seq: int, k_seq: int, num_heads: int, head_dim: int
) -> None:
hidden = num_heads * head_dim
eps = 1e-6
q = torch.randn(batch, q_seq, hidden, device="cuda", dtype=torch.bfloat16)
k = torch.randn(batch, k_seq, hidden, device="cuda", dtype=torch.bfloat16)
q_cos, q_sin = _make_cos_sin(batch, q_seq, num_heads, head_dim)
k_cos, k_sin = _make_cos_sin(batch, k_seq, num_heads, head_dim)
q_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
k_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
assert can_use_ltx2_qknorm_split_rope_cuda(
q,
q_cos,
q_sin,
q_weight,
k,
k_cos,
k_sin,
k_weight,
num_heads=num_heads,
head_dim=head_dim,
)
q_ref, k_ref = _reference(q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight, eps)
q_out, k_out = ltx2_qknorm_split_rope_cuda(
q,
q_cos,
q_sin,
q_weight,
k,
k_cos,
k_sin,
k_weight,
eps=eps,
num_heads=num_heads,
head_dim=head_dim,
)
torch.cuda.synchronize()
assert torch.equal(q_ref, q_out)
assert torch.equal(k_ref, k_out)
def test_ltx2_qknorm_split_rope_rejects_unsupported_inputs() -> None:
q = torch.randn((1, 3, 4096), device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
q_cos, q_sin = _make_cos_sin(1, 3, 32, 128)
q_weight = torch.randn(4096, device="cuda", dtype=torch.bfloat16)
k_weight = torch.randn(4096, device="cuda", dtype=torch.bfloat16)
assert can_use_ltx2_qknorm_split_rope_cuda(
q,
q_cos,
q_sin,
q_weight,
k,
q_cos,
q_sin,
k_weight,
num_heads=32,
head_dim=128,
)
assert not can_use_ltx2_qknorm_split_rope_cuda(
q.float(),
q_cos,
q_sin,
q_weight,
k,
q_cos,
q_sin,
k_weight,
num_heads=32,
head_dim=128,
)
assert not can_use_ltx2_qknorm_split_rope_cuda(
q,
q_cos,
q_sin,
q_weight,
k,
q_cos.transpose(-1, -2),
q_sin,
k_weight,
num_heads=32,
head_dim=128,
)
def test_ltx2_qknorm_split_rope_custom_op_torch_compile_fullgraph() -> None:
batch, q_seq, k_seq, num_heads, head_dim = 1, 3, 2, 32, 64
hidden = num_heads * head_dim
q = torch.randn(batch, q_seq, hidden, device="cuda", dtype=torch.bfloat16)
k = torch.randn(batch, k_seq, hidden, device="cuda", dtype=torch.bfloat16)
q_cos, q_sin = _make_cos_sin(batch, q_seq, num_heads, head_dim)
k_cos, k_sin = _make_cos_sin(batch, k_seq, num_heads, head_dim)
q_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
k_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
def fn(q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight):
return ltx2_qknorm_split_rope_cuda(
q,
q_cos,
q_sin,
q_weight,
k,
k_cos,
k_sin,
k_weight,
eps=1e-6,
num_heads=num_heads,
head_dim=head_dim,
)
compiled = torch.compile(fn, fullgraph=True)
q_out, k_out = compiled(q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight)
q_ref, k_ref = _reference(
q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight, 1e-6
)
torch.cuda.synchronize()
assert torch.equal(q_ref, q_out)
assert torch.equal(k_ref, k_out)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))