From e732c0a9dc071ca06026dba93887e8d77e631d04 Mon Sep 17 00:00:00 2001 From: xdtbynd <231153518+xdtbynd@users.noreply.github.com> Date: Sat, 8 Aug 2026 16:15:12 +0800 Subject: [PATCH] [UT][NPU] add unit tests for ascend_torch_native_backend and mla_preprocess (#32505) --- .../test_npu_ascend_torch_native_backend.py | 831 ++++++++++++++++++ .../npu/attention/test_npu_mla_preprocess.py | 207 +++++ 2 files changed, 1038 insertions(+) create mode 100644 test/registered/unit/npu/attention/test_npu_ascend_torch_native_backend.py create mode 100644 test/registered/unit/npu/attention/test_npu_mla_preprocess.py diff --git a/test/registered/unit/npu/attention/test_npu_ascend_torch_native_backend.py b/test/registered/unit/npu/attention/test_npu_ascend_torch_native_backend.py new file mode 100644 index 000000000..98c1eb5a6 --- /dev/null +++ b/test/registered/unit/npu/attention/test_npu_ascend_torch_native_backend.py @@ -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() diff --git a/test/registered/unit/npu/attention/test_npu_mla_preprocess.py b/test/registered/unit/npu/attention/test_npu_mla_preprocess.py new file mode 100644 index 000000000..338b6bc41 --- /dev/null +++ b/test/registered/unit/npu/attention/test_npu_mla_preprocess.py @@ -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()