[NPU]LoRA: Adding Torch Native backend (#14132)

test/manual/test_lora_ops.py

@@ -0,0 +1,287 @@
+import unittest
+
+import torch
+
+from sglang.srt.lora.torch_ops.lora_ops import (
+    bgmv_expand,
+    bgmv_expand_slice,
+    bgmv_shrink,
+    sgmv_expand,
+    sgmv_expand_slice,
+    sgmv_shrink,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestLoraOps(CustomTestCase):
+    def test_sgmv_expand(self):
+        batch_size = 2
+        input_dim = 4
+        output_dim = 6
+        num_loras = 3
+        dtype = torch.float32
+
+        inputs = torch.randn(batch_size, input_dim, dtype=dtype)
+        lora_b_weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
+        seq_len_tensor = torch.ones(batch_size, dtype=torch.int32)
+        lora_indices_tensor = torch.randint(0, num_loras, (batch_size,))
+        add_inputs = True
+
+        total_seq_len, _ = inputs.shape
+        exploded_indices = torch.repeat_interleave(
+            lora_indices_tensor, seq_len_tensor, output_size=total_seq_len
+        )
+        expect_output = torch.zeros(batch_size, output_dim, dtype=dtype)
+        bgmv_expand(inputs, lora_b_weights, expect_output, exploded_indices, add_inputs)
+
+        actual_output = torch.zeros(batch_size, output_dim, dtype=dtype)
+        sgmv_expand(
+            inputs,
+            lora_b_weights,
+            actual_output,
+            seq_len_tensor,
+            lora_indices_tensor,
+            add_inputs,
+        )
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_bgmv_expand(self):
+        batch_size = 2
+        input_dim = 4
+        output_dim = 6
+        num_loras = 3
+        dtype = torch.float32
+
+        inputs = torch.randn(batch_size, input_dim, dtype=dtype)
+        lora_b_weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
+        lora_indices_tensor = torch.randint(0, num_loras, (batch_size,))
+
+        selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=dtype)
+        selected_loras = selected_loras.squeeze(dim=1)
+        inputs = inputs.to(dtype=dtype)
+        outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
+        limit = batch_size
+        common_len = min(outputs.shape[1], output_dim)
+        expect_output = torch.zeros(batch_size, output_dim, dtype=dtype)
+        expect_output[:, :common_len] = outputs[:limit, :common_len]
+
+        actual_output = torch.zeros(batch_size, output_dim, dtype=dtype)
+        bgmv_expand(
+            inputs,
+            lora_b_weights,
+            actual_output,
+            lora_indices_tensor,
+            add_inputs=False,
+        )
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_bgmv_expand_add_residual(self):
+        batch_size = 2
+        input_dim = 4
+        output_dim = 6
+        num_loras = 3
+        dtype = torch.float32
+
+        inputs = torch.randn(batch_size, input_dim, dtype=dtype)
+        lora_b_weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
+        lora_indices_tensor = torch.randint(0, num_loras, (batch_size,))
+
+        selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=dtype)
+        selected_loras = selected_loras.squeeze(dim=1)
+        inputs = inputs.to(dtype=dtype)
+        outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
+        limit = batch_size
+        common_len = min(outputs.shape[1], output_dim)
+        expect_output = torch.randn(batch_size, output_dim, dtype=dtype)
+        actual_output = expect_output.clone()
+
+        expect_output[:, :common_len] += outputs[:limit, :common_len]
+
+        bgmv_expand(
+            inputs,
+            lora_b_weights,
+            actual_output,
+            lora_indices_tensor,
+            add_inputs=True,
+        )
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_sgmv_shrink(self):
+        batch_size = 2
+        input_dim = 4
+        output_dim = 6
+        num_loras = 3
+        dtype = torch.float32
+
+        inputs = torch.randn(batch_size, input_dim, dtype=dtype)
+        lora_a_weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
+        seq_len_tensor = torch.ones(batch_size, dtype=torch.int32)
+        lora_indices_tensor = torch.randint(0, num_loras, (batch_size,))
+        scaling = 0.9
+
+        total_seq_len, _ = inputs.shape
+        exploded_indices = torch.repeat_interleave(
+            lora_indices_tensor, seq_len_tensor, output_size=total_seq_len
+        )
+        expect_output = torch.zeros(batch_size, output_dim, dtype=dtype)
+        bgmv_shrink(inputs, lora_a_weights, expect_output, exploded_indices, scaling)
+
+        actual_output = torch.zeros(batch_size, output_dim, dtype=dtype)
+        sgmv_shrink(
+            inputs,
+            lora_a_weights,
+            actual_output,
+            seq_len_tensor,
+            lora_indices_tensor,
+            scaling,
+        )
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_bgmv_shrink(self):
+        batch_size = 2
+        input_dim = 4
+        output_dim = 6
+        num_loras = 3
+        dtype = torch.float32
+
+        inputs = torch.randn(batch_size, input_dim, dtype=dtype)
+        lora_a_weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
+        lora_indices_tensor = torch.randint(0, num_loras, (batch_size,))
+        scaling = 0.9
+
+        selected_loras = lora_a_weights[lora_indices_tensor].to(dtype=dtype)
+        inputs = inputs.to(dtype=dtype)
+        outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
+
+        expect_output = torch.zeros(batch_size, output_dim, dtype=dtype)
+        expect_output[:, : outputs.shape[1]] = scaling * outputs[:]
+
+        actual_output = torch.zeros(batch_size, output_dim, dtype=dtype)
+        bgmv_shrink(
+            inputs,
+            lora_a_weights,
+            actual_output,
+            lora_indices_tensor,
+            scaling=scaling,
+        )
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_sgmv_expand_slice(self):
+        batch_size = 2
+        input_dim = 4
+        output_dim = 6
+        output_dim_slice = 12
+        num_loras = 3
+        dtype = torch.float32
+
+        inputs = torch.randn(batch_size, input_dim, dtype=dtype)
+        lora_b_weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
+        seq_len_tensor = torch.ones(batch_size, dtype=torch.int32)
+        lora_indices_tensor = torch.randint(0, num_loras, (batch_size,))
+        slice_offset = 2
+        slice_size = 6
+        add_inputs = False
+
+        total_seq_len, _ = inputs.shape
+        exploded_indices = torch.repeat_interleave(
+            lora_indices_tensor, seq_len_tensor, output_size=total_seq_len
+        )
+        expect_output = torch.randn(batch_size, output_dim_slice, dtype=dtype)
+        actual_output = expect_output.clone()
+        bgmv_expand_slice(
+            inputs,
+            lora_b_weights,
+            expect_output,
+            exploded_indices,
+            slice_offset,
+            slice_size,
+            add_inputs,
+        )
+
+        sgmv_expand_slice(
+            inputs,
+            lora_b_weights,
+            actual_output,
+            seq_len_tensor,
+            lora_indices_tensor,
+            slice_offset,
+            slice_size,
+            add_inputs,
+        )
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_bgmv_expand_slice(self):
+        batch_size = 2
+        input_dim = 4
+        output_dim = 6
+        output_dim_slice = 12
+        num_loras = 3
+        dtype = torch.float32
+
+        inputs = torch.randn(batch_size, input_dim, dtype=dtype)
+        lora_b_weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
+        lora_indices_tensor = torch.randint(0, num_loras, (batch_size,))
+        slice_offset = 2
+        slice_size = 6
+
+        selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=dtype)
+        inputs = inputs.to(dtype=dtype)
+        outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
+        expect_output = torch.zeros(batch_size, output_dim_slice, dtype=dtype)
+        expect_output[:, slice_offset : slice_offset + slice_size] = outputs[:]
+
+        actual_output = torch.zeros(batch_size, output_dim_slice, dtype=dtype)
+        bgmv_expand_slice(
+            inputs,
+            lora_b_weights,
+            actual_output,
+            lora_indices_tensor,
+            slice_offset,
+            slice_size,
+            add_inputs=False,
+        )
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_bgmv_expand_slice_add_residual(self):
+        batch_size = 2
+        input_dim = 4
+        output_dim = 6
+        output_dim_slice = 12
+        num_loras = 3
+        dtype = torch.float32
+
+        inputs = torch.randn(batch_size, input_dim, dtype=dtype)
+        lora_b_weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
+        lora_indices_tensor = torch.randint(0, num_loras, (batch_size,))
+        slice_offset = 2
+        slice_size = 6
+
+        selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=dtype)
+        inputs = inputs.to(dtype=dtype)
+        outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
+        expect_output = torch.randn(batch_size, output_dim_slice, dtype=dtype)
+        actual_output = expect_output.clone()
+        expect_output[:, slice_offset : slice_offset + slice_size] += outputs[:]
+
+        bgmv_expand_slice(
+            inputs,
+            lora_b_weights,
+            actual_output,
+            lora_indices_tensor,
+            slice_offset,
+            slice_size,
+            add_inputs=True,
+        )
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+
+if __name__ == "__main__":
+    unittest.main()