[diffusion] Accelerate Cosmos3 T2I QKNorm+RoPE (#34932)

2026-08-16 20:15:32 +08:00
parent 095ec6c997
commit 41abbb0d32
3 changed files with 135 additions and 2 deletions
@@ -413,6 +413,99 @@ def test_qknorm_rope_pack_kv_matches_separate_ops() -> None:
    assert torch.equal(packed_v_ref, packed_kv[1])


+def test_qknorm_rope_pack_kv_preserves_split_bf16_rounding() -> None:
+    from sgl_kernel import rotary_embedding
+
+    from sglang.kernels.ops.diffusion.qknorm_rope import (
+        fused_qknorm_rope_pack_kv,
+    )
+    from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
+
+    batch_size = 1
+    prefix_tokens, suffix_tokens = 17, 1024
+    num_q_heads, num_kv_heads, head_dim = 32, 8, 64
+    num_heads = num_q_heads + 2 * num_kv_heads
+    qkv = torch.randn(
+        batch_size,
+        suffix_tokens,
+        num_heads,
+        head_dim,
+        device=DEVICE,
+        dtype=DTYPE,
+    )
+    prefix_qkv = torch.randn(
+        batch_size,
+        prefix_tokens,
+        num_heads,
+        head_dim,
+        device=DEVICE,
+        dtype=DTYPE,
+    )
+    k_prefix = prefix_qkv[:, :, num_q_heads : num_q_heads + num_kv_heads]
+    v_prefix = prefix_qkv[:, :, num_q_heads + num_kv_heads :]
+    q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    positions = torch.arange(
+        batch_size * suffix_tokens, device=DEVICE, dtype=torch.int64
+    )
+    cos_sin_cache = create_cos_sin_cache(head_dim, batch_size * suffix_tokens).to(DTYPE)
+
+    qkv_ref = qkv.clone()
+    q_ref = qkv_ref[:, :, :num_q_heads]
+    k_ref = qkv_ref[:, :, num_q_heads : num_q_heads + num_kv_heads]
+    v_ref = qkv_ref[:, :, num_q_heads + num_kv_heads :]
+    fused_inplace_qknorm(
+        q_ref.view(-1, num_q_heads, head_dim),
+        k_ref.view(-1, num_kv_heads, head_dim),
+        q_weight,
+        k_weight,
+        eps=1e-6,
+    )
+    rotary_embedding(
+        positions,
+        q_ref.view(-1, num_q_heads * head_dim),
+        k_ref.view(-1, num_kv_heads * head_dim),
+        head_dim,
+        cos_sin_cache,
+        True,
+    )
+    packed_k_ref = torch.cat([k_prefix, k_ref], dim=1)
+    packed_v_ref = torch.cat([v_prefix, v_ref], dim=1)
+
+    qkv_fused = qkv.clone()
+    q_fused = qkv_fused[:, :, :num_q_heads]
+    k_fused = qkv_fused[:, :, num_q_heads : num_q_heads + num_kv_heads]
+    v_fused = qkv_fused[:, :, num_q_heads + num_kv_heads :]
+    packed_kv = torch.empty(
+        2,
+        batch_size,
+        prefix_tokens + suffix_tokens,
+        num_kv_heads,
+        head_dim,
+        device=DEVICE,
+        dtype=DTYPE,
+    )
+    fused_qknorm_rope_pack_kv(
+        q_fused,
+        k_fused,
+        v_fused,
+        k_prefix,
+        v_prefix,
+        packed_kv,
+        q_weight,
+        k_weight,
+        cos_sin_cache,
+        positions,
+        is_neox=True,
+        rope_dim=head_dim,
+        round_norm_before_rope=True,
+    )
+
+    assert torch.equal(q_ref, q_fused)
+    assert torch.equal(packed_k_ref, packed_kv[0])
+    assert torch.equal(packed_v_ref, packed_kv[1])
+
+
 def test_qknorm_rope_accepts_empty_token_dimension() -> None:
    from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope