[UT][NPU] add unit tests for ascend_torch_native_backend and mla_preprocess (#32505)
This commit is contained in:
xdtbynd
@@ -0,0 +1,831 @@
|
||||
"""
|
||||
Unit tests for sglang.srt.hardware_backend.npu.attention.ascend_torch_native_backend.
|
||||
"""
|
||||
|
||||
import math
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from torch.nn.functional import scaled_dot_product_attention
|
||||
|
||||
from sglang.srt.hardware_backend.npu.attention.ascend_torch_native_backend import (
|
||||
AscendTorchNativeAttnBackend,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_npu_ci
|
||||
|
||||
register_npu_ci(est_time=4, suite="stage-a-unit-test-npu")
|
||||
|
||||
|
||||
class TestInit(unittest.TestCase):
|
||||
def test_construction(self):
|
||||
backend = AscendTorchNativeAttnBackend()
|
||||
self.assertIsNotNone(backend)
|
||||
|
||||
|
||||
class TestSupportTriton(unittest.TestCase):
|
||||
def test_returns_false(self):
|
||||
backend = AscendTorchNativeAttnBackend()
|
||||
self.assertFalse(backend.support_triton())
|
||||
|
||||
|
||||
class TestScaledDotProductAttentionWithSoftcapping(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.backend = AscendTorchNativeAttnBackend()
|
||||
|
||||
def test_basic_attention(self):
|
||||
L, S, d = 4, 4, 8
|
||||
q = torch.randn(1, 2, L, d)
|
||||
k = torch.randn(1, 2, S, d)
|
||||
v = torch.randn(1, 2, S, d)
|
||||
out = self.backend.scaled_dot_product_attention_with_softcapping(q, k, v)
|
||||
|
||||
scale = 1 / math.sqrt(d)
|
||||
ref = torch.softmax(q @ k.transpose(-2, -1) * scale, dim=-1) @ v
|
||||
self.assertTrue(torch.allclose(out, ref, atol=1e-5))
|
||||
|
||||
def test_causal_mask(self):
|
||||
L, S, d = 4, 4, 8
|
||||
q = torch.randn(1, 2, L, d)
|
||||
k = torch.randn(1, 2, S, d)
|
||||
v = torch.randn(1, 2, S, d)
|
||||
out = self.backend.scaled_dot_product_attention_with_softcapping(
|
||||
q, k, v, is_causal=True
|
||||
)
|
||||
|
||||
scale = 1 / math.sqrt(d)
|
||||
attn = q @ k.transpose(-2, -1) * scale
|
||||
mask = torch.ones(L, S, dtype=torch.bool).tril()
|
||||
attn = attn.masked_fill(~mask, float("-inf"))
|
||||
ref = torch.softmax(attn, dim=-1) @ v
|
||||
self.assertTrue(torch.allclose(out, ref, atol=1e-5))
|
||||
|
||||
def test_causal_mask_asserts_when_attn_mask_given(self):
|
||||
q = torch.randn(1, 1, 2, 4)
|
||||
k = torch.randn(1, 1, 2, 4)
|
||||
v = torch.randn(1, 1, 2, 4)
|
||||
with self.assertRaises(AssertionError):
|
||||
self.backend.scaled_dot_product_attention_with_softcapping(
|
||||
q, k, v, attn_mask=torch.ones(2, 2), is_causal=True
|
||||
)
|
||||
|
||||
def test_explicit_scale(self):
|
||||
d = 8
|
||||
q = torch.randn(1, 1, 2, d)
|
||||
k = torch.randn(1, 1, 2, d)
|
||||
v = torch.randn(1, 1, 2, d)
|
||||
out = self.backend.scaled_dot_product_attention_with_softcapping(
|
||||
q, k, v, scale=0.5
|
||||
)
|
||||
ref = torch.softmax(q @ k.transpose(-2, -1) * 0.5, dim=-1) @ v
|
||||
self.assertTrue(torch.allclose(out, ref, atol=1e-5))
|
||||
|
||||
def test_gqa(self):
|
||||
H_q, H_kv, L, S, d = 4, 2, 3, 3, 8
|
||||
q = torch.randn(1, H_q, L, d)
|
||||
k = torch.randn(1, H_kv, S, d)
|
||||
v = torch.randn(1, H_kv, S, d)
|
||||
out = self.backend.scaled_dot_product_attention_with_softcapping(
|
||||
q, k, v, enable_gqa=True
|
||||
)
|
||||
self.assertEqual(out.shape, (1, H_q, L, d))
|
||||
|
||||
k_exp = k.repeat_interleave(H_q // H_kv, -3)
|
||||
v_exp = v.repeat_interleave(H_q // H_kv, -3)
|
||||
scale = 1 / math.sqrt(d)
|
||||
ref = torch.softmax(q @ k_exp.transpose(-2, -1) * scale, dim=-1) @ v_exp
|
||||
self.assertTrue(torch.allclose(out, ref, atol=1e-5))
|
||||
|
||||
def test_tanh_softcapping(self):
|
||||
d = 8
|
||||
q = torch.randn(1, 1, 2, d)
|
||||
k = torch.randn(1, 1, 2, d)
|
||||
v = torch.randn(1, 1, 2, d)
|
||||
cap = 10.0
|
||||
out = self.backend.scaled_dot_product_attention_with_softcapping(
|
||||
q, k, v, logit_cap=cap
|
||||
)
|
||||
|
||||
scale = 1 / math.sqrt(d)
|
||||
attn = q @ k.transpose(-2, -1) * scale
|
||||
attn = cap * torch.tanh(attn / cap)
|
||||
ref = torch.softmax(attn, dim=-1) @ v
|
||||
self.assertTrue(torch.allclose(out, ref, atol=1e-5))
|
||||
|
||||
def test_no_softcapping_when_cap_zero(self):
|
||||
d = 8
|
||||
q = torch.randn(1, 1, 2, d)
|
||||
k = torch.randn(1, 1, 2, d)
|
||||
v = torch.randn(1, 1, 2, d)
|
||||
out = self.backend.scaled_dot_product_attention_with_softcapping(
|
||||
q, k, v, logit_cap=0.0
|
||||
)
|
||||
scale = 1 / math.sqrt(d)
|
||||
ref = torch.softmax(q @ k.transpose(-2, -1) * scale, dim=-1) @ v
|
||||
self.assertTrue(torch.allclose(out, ref, atol=1e-5))
|
||||
|
||||
def test_boolean_attn_mask(self):
|
||||
d = 8
|
||||
q = torch.randn(1, 1, 3, d)
|
||||
k = torch.randn(1, 1, 3, d)
|
||||
v = torch.randn(1, 1, 3, d)
|
||||
mask = torch.tensor(
|
||||
[
|
||||
[True, False, False],
|
||||
[True, True, False],
|
||||
[True, True, True],
|
||||
]
|
||||
)
|
||||
out = self.backend.scaled_dot_product_attention_with_softcapping(
|
||||
q, k, v, attn_mask=mask
|
||||
)
|
||||
|
||||
scale = 1 / math.sqrt(d)
|
||||
attn = q @ k.transpose(-2, -1) * scale
|
||||
attn = attn.masked_fill(~mask, float("-inf"))
|
||||
ref = torch.softmax(attn, dim=-1) @ v
|
||||
self.assertTrue(torch.allclose(out, ref, atol=1e-5))
|
||||
|
||||
def test_additive_attn_mask(self):
|
||||
d = 8
|
||||
q = torch.randn(1, 1, 2, d)
|
||||
k = torch.randn(1, 1, 2, d)
|
||||
v = torch.randn(1, 1, 2, d)
|
||||
mask = torch.tensor([[0.0, -1e9], [0.0, 0.0]])
|
||||
out = self.backend.scaled_dot_product_attention_with_softcapping(
|
||||
q, k, v, attn_mask=mask
|
||||
)
|
||||
scale = 1 / math.sqrt(d)
|
||||
attn = q @ k.transpose(-2, -1) * scale + mask
|
||||
ref = torch.softmax(attn, dim=-1) @ v
|
||||
self.assertTrue(torch.allclose(out, ref, atol=1e-5))
|
||||
|
||||
|
||||
class TestRunSdpaForwardExtend(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.backend = AscendTorchNativeAttnBackend()
|
||||
|
||||
def _make_caches(self, num_tokens, num_heads, head_size, dtype=torch.float32):
|
||||
k_cache = torch.randn(num_tokens, num_heads, head_size, dtype=dtype)
|
||||
v_cache = torch.randn(num_tokens, num_heads, head_size, dtype=dtype)
|
||||
return k_cache, v_cache
|
||||
|
||||
def _ref_extend(
|
||||
self,
|
||||
query,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
scaling=None,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
):
|
||||
H, D = query.shape[1], query.shape[2]
|
||||
outputs = []
|
||||
start_q = 0
|
||||
for seq_idx in range(seq_lens.shape[0]):
|
||||
ext_len = int(extend_seq_lens[seq_idx].item())
|
||||
pre_len = int(extend_prefix_lens[seq_idx].item())
|
||||
seq_len_kv = int(seq_lens[seq_idx].item())
|
||||
end_q = start_q + ext_len
|
||||
|
||||
req_pool_idx = req_pool_indices[seq_idx]
|
||||
tokens = req_to_token[req_pool_idx, :seq_len_kv]
|
||||
|
||||
# Build padded query [1, H, seq_len_kv, D] with extend tokens at pre_len
|
||||
q_pad = torch.zeros(1, H, seq_len_kv, D, dtype=query.dtype)
|
||||
q_pad[0, :, pre_len : pre_len + ext_len, :] = query[start_q:end_q].movedim(
|
||||
0, 1
|
||||
)
|
||||
# Key/value [1, H, seq_len_kv, D]
|
||||
k_u = k_cache[tokens].movedim(0, 1).unsqueeze(0)
|
||||
v_u = v_cache[tokens].movedim(0, 1).unsqueeze(0)
|
||||
|
||||
if enable_gqa and k_u.size(-3) != q_pad.size(-3):
|
||||
rep = q_pad.size(-3) // k_u.size(-3)
|
||||
k_u = k_u.repeat_interleave(rep, -3)
|
||||
v_u = v_u.repeat_interleave(rep, -3)
|
||||
|
||||
out = scaled_dot_product_attention(
|
||||
q_pad, k_u, v_u, scale=scaling, is_causal=causal
|
||||
)
|
||||
# [1, H, seq_len_kv, D] → slice [pre_len:pre_len+ext_len] → [ext, H, D]
|
||||
out = out[0, :, pre_len : pre_len + ext_len, :].movedim(1, 0)
|
||||
outputs.append(out)
|
||||
start_q = end_q
|
||||
return torch.cat(outputs, dim=0)
|
||||
|
||||
def test_basic_extend(self):
|
||||
H, D = 2, 8
|
||||
num_seqs = 2
|
||||
extend_prefix_lens = torch.tensor([0, 0], dtype=torch.int32)
|
||||
extend_seq_lens = torch.tensor([3, 4], dtype=torch.int32)
|
||||
seq_lens = extend_prefix_lens + extend_seq_lens
|
||||
total_q = int(extend_seq_lens.sum().item())
|
||||
max_ctx = int(seq_lens.max().item())
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = (
|
||||
torch.arange(max_ctx).unsqueeze(0).expand(num_seqs, -1).contiguous()
|
||||
)
|
||||
req_pool_indices = torch.tensor([0, 1], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_extend(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
)
|
||||
ref = self._ref_extend(
|
||||
query,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
)
|
||||
self.assertTrue(torch.allclose(result, ref, atol=1e-5))
|
||||
|
||||
def test_causal_extend(self):
|
||||
H, D = 2, 8
|
||||
extend_prefix_lens = torch.tensor([0], dtype=torch.int32)
|
||||
extend_seq_lens = torch.tensor([5], dtype=torch.int32)
|
||||
seq_lens = extend_prefix_lens + extend_seq_lens
|
||||
total_q = 5
|
||||
max_ctx = 5
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_extend(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=True,
|
||||
)
|
||||
ref = self._ref_extend(
|
||||
query,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=True,
|
||||
)
|
||||
self.assertTrue(torch.allclose(result, ref, atol=1e-5))
|
||||
|
||||
def test_extend_with_prefix_lens(self):
|
||||
H, D = 2, 8
|
||||
extend_prefix_lens = torch.tensor([2], dtype=torch.int32)
|
||||
extend_seq_lens = torch.tensor([3], dtype=torch.int32)
|
||||
seq_lens = extend_prefix_lens + extend_seq_lens
|
||||
total_q = 3
|
||||
max_ctx = 5
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_extend(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=True,
|
||||
)
|
||||
ref = self._ref_extend(
|
||||
query,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=True,
|
||||
)
|
||||
self.assertTrue(torch.allclose(result, ref, atol=1e-5))
|
||||
|
||||
def test_extend_with_gqa(self):
|
||||
H_q, H_kv, D = 4, 2, 8
|
||||
extend_prefix_lens = torch.tensor([0], dtype=torch.int32)
|
||||
extend_seq_lens = torch.tensor([4], dtype=torch.int32)
|
||||
seq_lens = extend_prefix_lens + extend_seq_lens
|
||||
total_q = 4
|
||||
max_ctx = 4
|
||||
|
||||
query = torch.randn(total_q, H_q, D)
|
||||
output = torch.empty(total_q, H_q, D)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H_kv, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_extend(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
enable_gqa=True,
|
||||
causal=False,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H_q, D))
|
||||
|
||||
def test_extend_with_logit_cap(self):
|
||||
H, D = 2, 8
|
||||
extend_prefix_lens = torch.tensor([0], dtype=torch.int32)
|
||||
extend_seq_lens = torch.tensor([4], dtype=torch.int32)
|
||||
seq_lens = extend_prefix_lens + extend_seq_lens
|
||||
total_q = 4
|
||||
max_ctx = 4
|
||||
cap = 10.0
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_extend(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=True,
|
||||
logit_cap=cap,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H, D))
|
||||
|
||||
def test_extend_with_cross_attention(self):
|
||||
H, D = 2, 8
|
||||
extend_prefix_lens = torch.tensor([2], dtype=torch.int32)
|
||||
extend_seq_lens = torch.tensor([3], dtype=torch.int32)
|
||||
seq_lens = extend_prefix_lens + extend_seq_lens
|
||||
total_q = 3
|
||||
max_ctx = 5
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
encoder_lens = torch.tensor([4], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_extend(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
encoder_lens=encoder_lens,
|
||||
is_cross_attention=True,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H, D))
|
||||
|
||||
def test_extend_with_encoder_self_attention(self):
|
||||
H, D = 2, 8
|
||||
extend_prefix_lens = torch.tensor([0], dtype=torch.int32)
|
||||
extend_seq_lens = torch.tensor([3], dtype=torch.int32)
|
||||
seq_lens = extend_prefix_lens + extend_seq_lens
|
||||
total_q = 3
|
||||
max_ctx = 6
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
encoder_lens = torch.tensor([3], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_extend(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
encoder_lens=encoder_lens,
|
||||
is_cross_attention=False,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H, D))
|
||||
|
||||
def test_extend_with_sw(self):
|
||||
H, D = 2, 8
|
||||
extend_prefix_lens = torch.tensor([4], dtype=torch.int32)
|
||||
extend_seq_lens = torch.tensor([3], dtype=torch.int32)
|
||||
seq_lens = extend_prefix_lens + extend_seq_lens
|
||||
total_q = 3
|
||||
max_ctx = 7
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_extend(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
sliding_window_size=3,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H, D))
|
||||
|
||||
def test_extend_with_full_to_swa_mapping(self):
|
||||
H, D = 2, 8
|
||||
extend_prefix_lens = torch.tensor([0], dtype=torch.int32)
|
||||
extend_seq_lens = torch.tensor([4], dtype=torch.int32)
|
||||
seq_lens = extend_prefix_lens + extend_seq_lens
|
||||
total_q = 4
|
||||
max_ctx = 4
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
# Identity mapping: full index == SWA index
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
full_to_swa = torch.arange(max_ctx)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_extend(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
full_to_swa_mapping=full_to_swa,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H, D))
|
||||
|
||||
def test_extend_dtype_casting(self):
|
||||
H, D = 2, 8
|
||||
extend_prefix_lens = torch.tensor([0], dtype=torch.int32)
|
||||
extend_seq_lens = torch.tensor([3], dtype=torch.int32)
|
||||
seq_lens = extend_prefix_lens + extend_seq_lens
|
||||
total_q = 3
|
||||
max_ctx = 3
|
||||
|
||||
query = torch.randn(total_q, H, D, dtype=torch.float32)
|
||||
output = torch.empty_like(query)
|
||||
k_cache = torch.randn(max_ctx, H, D, dtype=torch.float64)
|
||||
v_cache = torch.randn(max_ctx, H, D, dtype=torch.float64)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_extend(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_prefix_lens,
|
||||
extend_seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
)
|
||||
self.assertEqual(result.dtype, torch.float32)
|
||||
|
||||
|
||||
class TestRunSdpaForwardDecode(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.backend = AscendTorchNativeAttnBackend()
|
||||
|
||||
def _make_caches(self, num_tokens, num_heads, head_size, dtype=torch.float32):
|
||||
k_cache = torch.randn(num_tokens, num_heads, head_size, dtype=dtype)
|
||||
v_cache = torch.randn(num_tokens, num_heads, head_size, dtype=dtype)
|
||||
return k_cache, v_cache
|
||||
|
||||
def test_basic_decode(self):
|
||||
H, D = 2, 8
|
||||
num_seqs = 2
|
||||
seq_lens = torch.tensor([3, 5], dtype=torch.int32)
|
||||
total_q = num_seqs
|
||||
max_ctx = int(seq_lens.max().item())
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = (
|
||||
torch.arange(max_ctx).unsqueeze(0).expand(num_seqs, -1).contiguous()
|
||||
)
|
||||
req_pool_indices = torch.tensor([0, 1], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_decode(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
)
|
||||
|
||||
# Manual reference
|
||||
outputs = []
|
||||
for i in range(num_seqs):
|
||||
kv_len = int(seq_lens[i].item())
|
||||
tokens = req_to_token[i, :kv_len]
|
||||
# query[i]: [H, D] → [1, H, 1, D]
|
||||
q_req = query[i : i + 1].movedim(0, 1).unsqueeze(0)
|
||||
# k/v: [kv, H, D] → [1, H, kv, D]
|
||||
k_req = k_cache[tokens].movedim(0, 1).unsqueeze(0)
|
||||
v_req = v_cache[tokens].movedim(0, 1).unsqueeze(0)
|
||||
out = scaled_dot_product_attention(q_req, k_req, v_req)
|
||||
# [1, H, 1, D] → [1, H, D]
|
||||
out = out.squeeze(0).movedim(1, 0)
|
||||
outputs.append(out)
|
||||
ref = torch.cat(outputs, dim=0)
|
||||
self.assertTrue(torch.allclose(result, ref, atol=1e-5))
|
||||
|
||||
def test_decode_causal(self):
|
||||
H, D = 2, 8
|
||||
seq_lens = torch.tensor([5], dtype=torch.int32)
|
||||
total_q = 1
|
||||
max_ctx = 5
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_decode(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=True,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H, D))
|
||||
|
||||
def test_decode_with_gqa(self):
|
||||
H_q, H_kv, D = 4, 2, 8
|
||||
seq_lens = torch.tensor([4], dtype=torch.int32)
|
||||
total_q = 1
|
||||
max_ctx = 4
|
||||
|
||||
query = torch.randn(total_q, H_q, D)
|
||||
output = torch.empty(total_q, H_q, D)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H_kv, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_decode(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
enable_gqa=True,
|
||||
causal=False,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H_q, D))
|
||||
|
||||
def test_decode_with_logit_cap(self):
|
||||
H, D = 2, 8
|
||||
seq_lens = torch.tensor([4], dtype=torch.int32)
|
||||
total_q = 1
|
||||
max_ctx = 4
|
||||
cap = 10.0
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_decode(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
logit_cap=cap,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H, D))
|
||||
|
||||
def test_decode_with_encoder_lens_cross(self):
|
||||
H, D = 2, 8
|
||||
seq_lens = torch.tensor([5], dtype=torch.int32)
|
||||
total_q = 1
|
||||
max_ctx = 5
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
encoder_lens = torch.tensor([3], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_decode(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
encoder_lens=encoder_lens,
|
||||
is_cross_attention=True,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H, D))
|
||||
|
||||
def test_decode_with_encoder_lens_self(self):
|
||||
H, D = 2, 8
|
||||
seq_lens = torch.tensor([6], dtype=torch.int32)
|
||||
total_q = 1
|
||||
max_ctx = 6
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
encoder_lens = torch.tensor([3], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_decode(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
encoder_lens=encoder_lens,
|
||||
is_cross_attention=False,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H, D))
|
||||
|
||||
def test_decode_with_sw(self):
|
||||
H, D = 2, 8
|
||||
seq_lens = torch.tensor([10], dtype=torch.int32)
|
||||
total_q = 1
|
||||
max_ctx = 10
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_decode(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
sliding_window_size=3,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H, D))
|
||||
|
||||
def test_decode_with_full_to_swa_mapping(self):
|
||||
H, D = 2, 8
|
||||
seq_lens = torch.tensor([4], dtype=torch.int32)
|
||||
total_q = 1
|
||||
max_ctx = 4
|
||||
|
||||
query = torch.randn(total_q, H, D)
|
||||
output = torch.empty_like(query)
|
||||
k_cache, v_cache = self._make_caches(max_ctx, H, D)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
full_to_swa = torch.arange(max_ctx)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_decode(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
full_to_swa_mapping=full_to_swa,
|
||||
)
|
||||
self.assertEqual(result.shape, (total_q, H, D))
|
||||
|
||||
def test_decode_dtype_casting(self):
|
||||
H, D = 2, 8
|
||||
seq_lens = torch.tensor([3], dtype=torch.int32)
|
||||
total_q = 1
|
||||
max_ctx = 3
|
||||
|
||||
query = torch.randn(total_q, H, D, dtype=torch.float32)
|
||||
output = torch.empty_like(query)
|
||||
k_cache = torch.randn(max_ctx, H, D, dtype=torch.float64)
|
||||
v_cache = torch.randn(max_ctx, H, D, dtype=torch.float64)
|
||||
req_to_token = torch.arange(max_ctx).unsqueeze(0)
|
||||
req_pool_indices = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
result = self.backend.run_sdpa_forward_decode(
|
||||
query,
|
||||
output,
|
||||
k_cache,
|
||||
v_cache,
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
enable_gqa=False,
|
||||
causal=False,
|
||||
)
|
||||
self.assertEqual(result.dtype, torch.float32)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,207 @@
|
||||
"""
|
||||
Unit tests for sglang.srt.hardware_backend.npu.attention.mla_preprocess.
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_npu_ci
|
||||
|
||||
register_npu_ci(est_time=4, suite="stage-a-unit-test-npu")
|
||||
|
||||
from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
|
||||
is_fia_nz,
|
||||
is_mla_preprocess_enabled,
|
||||
round_up,
|
||||
trans_rope_weight,
|
||||
transdata,
|
||||
)
|
||||
|
||||
|
||||
class TestRoundUp(unittest.TestCase):
|
||||
def test_exact_multiple(self):
|
||||
self.assertEqual(round_up(16, 16), 16)
|
||||
self.assertEqual(round_up(32, 16), 32)
|
||||
|
||||
def test_round_up_to_block(self):
|
||||
self.assertEqual(round_up(1, 16), 16)
|
||||
self.assertEqual(round_up(10, 16), 16)
|
||||
self.assertEqual(round_up(17, 16), 32)
|
||||
|
||||
def test_zero_value(self):
|
||||
self.assertEqual(round_up(0, 16), 0)
|
||||
|
||||
def test_zero_align_returns_zero(self):
|
||||
self.assertEqual(round_up(100, 0), 0)
|
||||
self.assertEqual(round_up(0, 0), 0)
|
||||
|
||||
def test_negative_align(self):
|
||||
# Pin current behavior for negative align.
|
||||
self.assertEqual(round_up(10, -4), 8)
|
||||
|
||||
|
||||
class TestTransdata(unittest.TestCase):
|
||||
def test_aligned_32x32_default_block(self):
|
||||
mat = torch.arange(32 * 32, dtype=torch.float32).reshape(32, 32)
|
||||
nz = transdata(mat)
|
||||
self.assertEqual(nz.shape, (2, 32, 16))
|
||||
|
||||
def test_aligned_element_mapping(self):
|
||||
rows, cols = 16, 32
|
||||
mat = torch.arange(rows * cols, dtype=torch.float32).reshape(rows, cols)
|
||||
bs = (16, 16)
|
||||
nz = transdata(mat, block_size=bs)
|
||||
|
||||
# NZ layout: element (r, c) lands at nz[c // bs1, r, c % bs1].
|
||||
for r in range(rows):
|
||||
for c in range(cols):
|
||||
expected = mat[r, c]
|
||||
actual = nz[c // bs[1], r, c % bs[1]]
|
||||
self.assertEqual(
|
||||
actual.item(),
|
||||
expected.item(),
|
||||
f"mismatch at (r={r}, c={c})",
|
||||
)
|
||||
|
||||
def test_padding_for_non_aligned_shape(self):
|
||||
mat = torch.arange(100, dtype=torch.float32).reshape(10, 10)
|
||||
nz = transdata(mat, block_size=(16, 16))
|
||||
# c_blocks = ceil(10/16) = 1, r_padded = 16, bs1 = 16
|
||||
self.assertEqual(nz.shape, (1, 16, 16))
|
||||
|
||||
# Original elements are preserved in the padded NZ layout.
|
||||
for r in range(10):
|
||||
for c in range(10):
|
||||
self.assertEqual(
|
||||
nz[c // 16, r, c % 16].item(),
|
||||
mat[r, c].item(),
|
||||
f"padded mismatch at (r={r}, c={c})",
|
||||
)
|
||||
|
||||
def test_non_square_matrix(self):
|
||||
mat = torch.arange(16 * 48, dtype=torch.float32).reshape(16, 48)
|
||||
nz = transdata(mat, block_size=(16, 16))
|
||||
# c_blocks = 48//16 = 3, r = 16, bs1 = 16
|
||||
self.assertEqual(nz.shape, (3, 16, 16))
|
||||
|
||||
def test_custom_block_size(self):
|
||||
mat = torch.arange(8 * 8, dtype=torch.float32).reshape(8, 8)
|
||||
nz = transdata(mat, block_size=(4, 4))
|
||||
self.assertEqual(nz.shape, (2, 8, 4))
|
||||
|
||||
def test_3d_input_raises(self):
|
||||
mat = torch.arange(2 * 16 * 16, dtype=torch.float32).reshape(2, 16, 16)
|
||||
with self.assertRaises(RuntimeError):
|
||||
transdata(mat, block_size=(16, 16))
|
||||
|
||||
|
||||
class TestTransRopeWeight(unittest.TestCase):
|
||||
def test_basic_reorder(self):
|
||||
# 8 rows, rope_dim=4 → last 4 rows are the RoPE region.
|
||||
weight = torch.arange(8 * 4, dtype=torch.float32).reshape(8, 4).clone()
|
||||
original = weight.clone()
|
||||
result = trans_rope_weight(weight, rope_dim=4)
|
||||
|
||||
# RoPE region (rows 4-7) reordered: even indices first → [4,6,5,7].
|
||||
rope_region = original[4:8]
|
||||
expected_rope = torch.stack(
|
||||
[rope_region[0], rope_region[2], rope_region[1], rope_region[3]]
|
||||
)
|
||||
self.assertTrue(torch.equal(result[4:8], expected_rope))
|
||||
|
||||
def test_non_rope_region_unchanged(self):
|
||||
weight = torch.arange(8 * 4, dtype=torch.float32).reshape(8, 4).clone()
|
||||
original = weight.clone()
|
||||
trans_rope_weight(weight, rope_dim=4)
|
||||
self.assertTrue(torch.equal(weight[0:4], original[0:4]))
|
||||
|
||||
def test_rope_dim_zero_reorders_all(self):
|
||||
weight = torch.arange(4 * 8, dtype=torch.float32).reshape(4, 8).clone()
|
||||
original = weight.clone()
|
||||
result = trans_rope_weight(weight, rope_dim=0)
|
||||
# Rows [0,1,2,3] → [0,2,1,3] (even first, then odd)
|
||||
expected = torch.stack([original[0], original[2], original[1], original[3]])
|
||||
self.assertTrue(torch.equal(result, expected))
|
||||
|
||||
def test_full_rope_dim(self):
|
||||
weight = torch.arange(4 * 2, dtype=torch.float32).reshape(4, 2).clone()
|
||||
original = weight.clone()
|
||||
result = trans_rope_weight(weight, rope_dim=4)
|
||||
# rows [0,1,2,3] → [0,2,1,3]
|
||||
expected = torch.stack([original[0], original[2], original[1], original[3]])
|
||||
self.assertTrue(torch.equal(result, expected))
|
||||
|
||||
def test_3d_weight(self):
|
||||
# Shape (2, 8, 4): 8 rows per expert, rope_dim=4
|
||||
weight = torch.arange(2 * 8 * 4, dtype=torch.float32).reshape(2, 8, 4).clone()
|
||||
original = weight.clone()
|
||||
result = trans_rope_weight(weight, rope_dim=4)
|
||||
|
||||
for e in range(2):
|
||||
rope_region = original[e, 4:8]
|
||||
expected_rope = torch.stack(
|
||||
[rope_region[0], rope_region[2], rope_region[1], rope_region[3]]
|
||||
)
|
||||
self.assertTrue(torch.equal(result[e, 4:8], expected_rope))
|
||||
|
||||
|
||||
class TestIsMlaPreprocessEnabled(unittest.TestCase):
|
||||
def setUp(self):
|
||||
is_mla_preprocess_enabled.cache_clear()
|
||||
|
||||
def tearDown(self):
|
||||
is_mla_preprocess_enabled.cache_clear()
|
||||
|
||||
def test_not_set_returns_false(self):
|
||||
with patch.dict(os.environ):
|
||||
os.environ.pop("SGLANG_NPU_USE_MLAPO", None)
|
||||
self.assertFalse(is_mla_preprocess_enabled())
|
||||
|
||||
def test_set_to_one_returns_true(self):
|
||||
with patch.dict(os.environ, {"SGLANG_NPU_USE_MLAPO": "1"}):
|
||||
self.assertTrue(is_mla_preprocess_enabled())
|
||||
|
||||
def test_set_to_zero_returns_false(self):
|
||||
with patch.dict(os.environ, {"SGLANG_NPU_USE_MLAPO": "0"}):
|
||||
self.assertFalse(is_mla_preprocess_enabled())
|
||||
|
||||
def test_set_to_true_returns_true(self):
|
||||
with patch.dict(os.environ, {"SGLANG_NPU_USE_MLAPO": "true"}):
|
||||
self.assertTrue(is_mla_preprocess_enabled())
|
||||
|
||||
|
||||
class TestIsFiaNz(unittest.TestCase):
|
||||
def setUp(self):
|
||||
is_mla_preprocess_enabled.cache_clear()
|
||||
is_fia_nz.cache_clear()
|
||||
|
||||
def tearDown(self):
|
||||
is_mla_preprocess_enabled.cache_clear()
|
||||
is_fia_nz.cache_clear()
|
||||
|
||||
def test_not_set_returns_false(self):
|
||||
with patch.dict(os.environ):
|
||||
os.environ.pop("SGLANG_USE_FIA_NZ", None)
|
||||
os.environ.pop("SGLANG_NPU_USE_MLAPO", None)
|
||||
self.assertFalse(is_fia_nz())
|
||||
|
||||
def test_fia_nz_with_mlapo_returns_true(self):
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"SGLANG_USE_FIA_NZ": "1", "SGLANG_NPU_USE_MLAPO": "1"},
|
||||
):
|
||||
self.assertTrue(is_fia_nz())
|
||||
|
||||
def test_fia_nz_without_mlapo_raises(self):
|
||||
with patch.dict(os.environ):
|
||||
os.environ.pop("SGLANG_NPU_USE_MLAPO", None)
|
||||
os.environ["SGLANG_USE_FIA_NZ"] = "1"
|
||||
with self.assertRaises(AssertionError):
|
||||
is_fia_nz()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user