179 lines
6.1 KiB
Python
179 lines
6.1 KiB
Python
import sys
|
|
|
|
import pytest
|
|
import sgl_kernel # noqa: F401
|
|
import torch
|
|
|
|
from sglang.srt.utils import is_host_cpu_arm64
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
from sglang.test.cpu_test_utils import precision
|
|
|
|
register_cpu_ci(est_time=5, suite="stage-a-test-cpu-intel")
|
|
register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
|
|
|
|
torch.manual_seed(1234)
|
|
|
|
|
|
def fix_query_key_value_ordering_reshape_cat(
|
|
mixed_qkvz, mixed_ba, num_k_heads, num_v_heads, attn_tp_size, head_k_dim, head_v_dim
|
|
):
|
|
new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + (
|
|
num_k_heads // attn_tp_size,
|
|
(
|
|
head_k_dim
|
|
+ head_k_dim
|
|
+ (head_v_dim + head_v_dim) * num_v_heads // num_k_heads
|
|
),
|
|
)
|
|
new_tensor_shape_ba = mixed_ba.size()[:-1] + (
|
|
num_k_heads // attn_tp_size,
|
|
2 * num_v_heads // num_k_heads,
|
|
)
|
|
|
|
mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz)
|
|
mixed_ba = mixed_ba.view(*new_tensor_shape_ba)
|
|
|
|
split_arg_list_qkvz = [
|
|
head_k_dim,
|
|
head_k_dim,
|
|
(num_v_heads // num_k_heads * head_v_dim),
|
|
(num_v_heads // num_k_heads * head_v_dim),
|
|
]
|
|
split_arg_list_ba = [
|
|
num_v_heads // num_k_heads,
|
|
num_v_heads // num_k_heads,
|
|
]
|
|
# [b, sq, ng, (hn + hn + np/ng * hn + np/ng + np/ng)]
|
|
# --> [b, sq, ng, hn], [b, sq, ng, hn], [b, sq, ng, np/ng * hn], [b, sq, ng, np/ng * hn], [b, sq, ng, np/ng], [b, sq, ng, np/ng]
|
|
query, key, value, z = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=2)
|
|
b, a = torch.split(mixed_ba, split_arg_list_ba, dim=2)
|
|
|
|
# [b, sq, ng, np/ng * hn] -> [b, sq, np, hn]
|
|
value = value.reshape(value.size(0), -1, head_v_dim)
|
|
z = z.reshape(z.size(0), -1, head_v_dim)
|
|
b = b.reshape(b.size(0), num_v_heads // attn_tp_size)
|
|
a = a.reshape(a.size(0), num_v_heads // attn_tp_size)
|
|
query, key, value = map(lambda x: x.reshape(x.shape[0], -1), (query, key, value))
|
|
mixed_qkv = torch.cat((query, key, value), dim=-1)
|
|
|
|
return mixed_qkv, z, b, a
|
|
|
|
|
|
def fix_query_key_value_ordering_reshape_cat_contiguous(
|
|
mixed_qkvz: torch.Tensor,
|
|
mixed_ba: torch.Tensor,
|
|
key_dim: int,
|
|
value_dim: int,
|
|
num_v_heads: int,
|
|
head_v_dim: int,
|
|
attn_tp_size: int,
|
|
):
|
|
"""
|
|
Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.
|
|
"""
|
|
k_tp = key_dim // attn_tp_size
|
|
v_tp = value_dim // attn_tp_size
|
|
nv_tp = num_v_heads // attn_tp_size
|
|
|
|
# Directly split, no head group reshape
|
|
query, key, value, z = mixed_qkvz.split([k_tp, k_tp, v_tp, v_tp], dim=-1)
|
|
b, a = mixed_ba.split([nv_tp, nv_tp], dim=-1)
|
|
|
|
# value / z reshape to (seq, num_v_heads/tp, head_v_dim)
|
|
value = value.reshape(value.size(0), -1, head_v_dim)
|
|
z = z.reshape(z.size(0), -1, head_v_dim)
|
|
query, key, value = map(lambda x: x.reshape(x.shape[0], -1), (query, key, value))
|
|
mixed_qkv = torch.cat((query, key, value), dim=-1)
|
|
return mixed_qkv, z, b, a
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
is_host_cpu_arm64(), reason="fused_input_proj_cpu is an x86 AMX kernel"
|
|
)
|
|
def test_fused_input_proj():
|
|
batch = 7
|
|
hidden_size = 256
|
|
qkvz_dim = 128
|
|
ba_dim = 64
|
|
hidden_states = torch.randn(batch, hidden_size, dtype=torch.bfloat16)
|
|
qkvz_weight = torch.randn(qkvz_dim, hidden_size, dtype=torch.bfloat16)
|
|
ba_weight = torch.randn(ba_dim, hidden_size, dtype=torch.bfloat16)
|
|
|
|
qkvz_ref = torch.nn.functional.linear(hidden_states, qkvz_weight)
|
|
ba_ref = torch.nn.functional.linear(hidden_states, ba_weight)
|
|
qkvz, ba = torch.ops.sgl_kernel.fused_input_proj_cpu(
|
|
hidden_states, qkvz_weight, ba_weight, False
|
|
)
|
|
|
|
atol = rtol = precision[qkvz.dtype]
|
|
torch.testing.assert_close(qkvz, qkvz_ref, atol=atol, rtol=rtol)
|
|
torch.testing.assert_close(ba, ba_ref, atol=atol, rtol=rtol)
|
|
|
|
|
|
def test_fused_qkvzba_split_reshape_cat():
|
|
mixed_qkvz = torch.rand(1024, 12288, dtype=torch.bfloat16)
|
|
mixed_ba = torch.rand(1024, 64, dtype=torch.bfloat16)
|
|
head_k_dim = 128
|
|
head_v_dim = 128
|
|
num_v_heads = 32
|
|
num_k_heads = 16
|
|
attn_tp_size = 1
|
|
mixed_qkv_ref, z_ref, b_ref, a_ref = fix_query_key_value_ordering_reshape_cat(
|
|
mixed_qkvz,
|
|
mixed_ba,
|
|
num_k_heads,
|
|
num_v_heads,
|
|
attn_tp_size,
|
|
head_k_dim,
|
|
head_v_dim,
|
|
)
|
|
num_heads_qk = num_k_heads // attn_tp_size
|
|
num_heads_v = num_v_heads // attn_tp_size
|
|
mixed_qkv, z, b, a = torch.ops.sgl_kernel.fused_qkvzba_split_reshape_cat_cpu(
|
|
mixed_qkvz, mixed_ba, num_heads_qk, num_heads_v, head_k_dim, head_v_dim
|
|
)
|
|
atol = rtol = precision[mixed_qkv.dtype]
|
|
torch.testing.assert_close(mixed_qkv, mixed_qkv_ref, atol=atol, rtol=rtol)
|
|
torch.testing.assert_close(z, z_ref, atol=atol, rtol=rtol)
|
|
torch.testing.assert_close(b, b_ref, atol=atol, rtol=rtol)
|
|
torch.testing.assert_close(a, a_ref, atol=atol, rtol=rtol)
|
|
|
|
|
|
def test_fused_qkvzba_split_reshape_cat_contiguous():
|
|
mixed_qkvz = torch.rand(1, 12288, dtype=torch.bfloat16)
|
|
mixed_ba = torch.rand(1, 64, dtype=torch.bfloat16)
|
|
head_k_dim = 128
|
|
head_v_dim = 128
|
|
num_v_heads = 32
|
|
num_k_heads = 16
|
|
attn_tp_size = 1
|
|
key_dim = head_k_dim * num_k_heads
|
|
value_dim = head_v_dim * num_v_heads
|
|
mixed_qkv_ref, z_ref, b_ref, a_ref = (
|
|
fix_query_key_value_ordering_reshape_cat_contiguous(
|
|
mixed_qkvz,
|
|
mixed_ba,
|
|
key_dim,
|
|
value_dim,
|
|
num_v_heads,
|
|
head_v_dim,
|
|
attn_tp_size,
|
|
)
|
|
)
|
|
num_heads_qk = num_k_heads // attn_tp_size
|
|
num_heads_v = num_v_heads // attn_tp_size
|
|
mixed_qkv, z, b, a = (
|
|
torch.ops.sgl_kernel.fused_qkvzba_split_reshape_cat_contiguous_cpu(
|
|
mixed_qkvz, mixed_ba, num_heads_qk, num_heads_v, head_k_dim, head_v_dim
|
|
)
|
|
)
|
|
atol = rtol = precision[mixed_qkv.dtype]
|
|
torch.testing.assert_close(mixed_qkv, mixed_qkv_ref, atol=atol, rtol=rtol)
|
|
torch.testing.assert_close(z, z_ref, atol=atol, rtol=rtol)
|
|
torch.testing.assert_close(b, b_ref, atol=atol, rtol=rtol)
|
|
torch.testing.assert_close(a, a_ref, atol=atol, rtol=rtol)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main([__file__]))
|