[UT][NPU] add unit tests for ascend_torch_native_backend and mla_preprocess (#32505)

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()

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()