diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e0d368dc..8e420b7f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,10 +19,6 @@ jobs: runs-on: S60 steps: - - name: Setup SSH Agent - uses: webfactory/ssh-agent@v0.9.0 - with: - ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }} - name: Checkout SAST tools uses: actions/checkout@v4 @@ -30,14 +26,12 @@ jobs: repository: ${{env.SAST_REPO}} ref: ${{env.SAST_BRANCH}} path: ./sast - ssh-key: ${{ secrets.SSH_PUB_KEY }} - name: Checkout code uses: actions/checkout@v4 with: fetch-depth: 0 # Required for full git history to compare changes path: ./vllm-gcu - ssh-key: ${{ secrets.SSH_PUB_KEY }} - name: Run static code analysis run: | @@ -67,7 +61,7 @@ jobs: -v /home/pypi_packages:/home/pypi_packages \ -w /app \ "${{ env.CI_IMAGE }}" \ - bash -c 'git config --global --add safe.directory /app/vllm-gcu && sudo bash ./vllm-gcu/.pipeline/build.sh ci_build;sudo bash ./vllm-gcu/tests/install_env.sh' + bash -c 'git config --global --add safe.directory /app/vllm-gcu && sudo bash ./vllm-gcu/.pipeline/build.sh ci_build;sudo bash ./vllm-gcu/tests/install_env.sh;cd ./vllm-gcu/tests/;bash run_all_test.sh kernel' - name: Clean up if: always() diff --git a/tests/install_env.sh b/tests/install_env.sh index f86ceb0c..e547dc45 100644 --- a/tests/install_env.sh +++ b/tests/install_env.sh @@ -104,6 +104,7 @@ ls -lh "${WHEELS[@]}" echo "" log_warn "You may be prompted for your sudo password again." cd /home/pypi_packages +run_pip install pytest==8.3.3 run_pip install vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl cd - run_pip install "${WHEELS[@]}" diff --git a/tests/kernel_test/test_linear.py b/tests/kernel_test/test_linear.py new file mode 100644 index 00000000..860a5c55 --- /dev/null +++ b/tests/kernel_test/test_linear.py @@ -0,0 +1,124 @@ +import pytest +import torch + +import vllm.distributed.parallel_state as ps +from vllm_gcu.kernels.linear import ( + MergedReplicatedLinear, + CustomMergedColumnParallelLinear, +) + + +@pytest.fixture(autouse=True) +def fake_tensor_parallel(monkeypatch): + class FakeTPGroup: + rank_in_group = 0 + world_size = 1 + + def all_gather(self, x, dim=-1): + return x + + fake_group = FakeTPGroup() + + monkeypatch.setattr(ps, "_TP", fake_group, raising=False) + monkeypatch.setattr(ps, "get_tp_group", lambda: fake_group) + monkeypatch.setattr(ps, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(ps, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr( + "vllm.distributed.parallel_state.destroy_model_parallel", + lambda: None, + ) + + +def test_merged_replicated_weight_loader(): + layer = MergedReplicatedLinear( + input_size=4, + output_sizes=[3, 5], + bias=True, + ) + + weight_cpu = torch.randn(8, 4) + weight_gcu = weight_cpu.gcu() + + layer.weight_loader(layer.weight, weight_gcu) + assert layer.weight.shape == (8, 4) + torch.testing.assert_close( + layer.weight.cpu(), + weight_cpu, + rtol=0, + atol=0, + ) + + +def test_column_parallel_no_scale_no_gather_no_bias(): + layer_cpu = CustomMergedColumnParallelLinear( + input_size=4, + output_sizes=[8], + bias=True, + gather_output=False, + ) + layer_cpu.return_bias = False + + layer_gcu = CustomMergedColumnParallelLinear( + input_size=4, + output_sizes=[8], + bias=True, + gather_output=False, + ) + layer_gcu.return_bias = False + layer_gcu.to("gcu") + + with torch.no_grad(): + layer_gcu.weight.copy_(layer_cpu.weight) + layer_gcu.bias.copy_(layer_cpu.bias) + + x_cpu = torch.randn(2, 4) + x_gcu = x_cpu.gcu() + + out_cpu = layer_cpu(x_cpu) + out_gcu = layer_gcu(x_gcu) + + torch.testing.assert_close( + out_gcu.cpu(), + out_cpu, + rtol=1e-5, + atol=1e-5, + ) + + +def test_column_parallel_no_scale_with_gather_and_bias(): + layer_cpu = CustomMergedColumnParallelLinear( + input_size=4, + output_sizes=[8], + bias=True, + gather_output=True, + ) + layer_cpu.return_bias = True + + layer_gcu = CustomMergedColumnParallelLinear( + input_size=4, + output_sizes=[8], + bias=True, + gather_output=True, + ) + layer_gcu.return_bias = True + layer_gcu.to("gcu") + + with torch.no_grad(): + layer_gcu.weight.copy_(layer_cpu.weight) + layer_gcu.bias.copy_(layer_cpu.bias) + + x_cpu = torch.randn(2, 4) + x_gcu = x_cpu.gcu() + + out_cpu, bias_cpu = layer_cpu(x_cpu) + out_gcu, bias_gcu = layer_gcu(x_gcu) + + torch.testing.assert_close( + out_gcu.cpu(), + out_cpu, + rtol=1e-5, + atol=1e-5, + ) + + assert bias_cpu is None + assert bias_gcu is None diff --git a/tests/kernel_test/test_mla.py b/tests/kernel_test/test_mla.py new file mode 100644 index 00000000..c7d57514 --- /dev/null +++ b/tests/kernel_test/test_mla.py @@ -0,0 +1,195 @@ +import torch +import pytest +import uuid +from typing import Optional + +from vllm.model_executor.layers.linear import LinearBase +from vllm.model_executor.layers.layernorm import LayerNorm +from vllm.model_executor.layers.mla import MLAModules +from vllm_gcu.kernels.mla import GCUMultiHeadLatentAttention +from vllm.attention.layer import Attention + +LinearBase.name = "linear" + + +@pytest.fixture(autouse=True) +def init_fake_env(monkeypatch): + + class FakeTPGroup: + def __init__(self): + self.world_size = 1 + self.rank = 0 + self.rank_in_group = 0 + + def size(self): + return 1 + + fake_group = FakeTPGroup() + + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_tp_group", + lambda: fake_group, + ) + monkeypatch.setattr( + "vllm.distributed.get_tp_group", + lambda: fake_group, + ) + monkeypatch.setattr( + "vllm_gcu.kernels.mla.get_tp_group", + lambda: fake_group, + ) + + def fake_forward_native(self, x): + out_features = self.output_size + in_features = x.shape[-1] + + torch.manual_seed(0) + weight = torch.randn(out_features, in_features, device=x.device) * 0.01 + bias = torch.zeros(out_features, device=x.device) + + out = torch.nn.functional.linear(x, weight, bias) + return out, None + + monkeypatch.setattr( + LinearBase, + "forward_native", + fake_forward_native, + ) + from vllm.attention.layer import Attention as GCUAttention + + def fake_attention_forward(self, q, k, v, *args, **kwargs): + return q + + monkeypatch.setattr( + GCUAttention, + "forward", + fake_attention_forward, + ) + + +def make_linear(inp, out, prefix): + layer = LinearBase( + input_size=inp, + output_size=out, + skip_bias_add=False, + params_dtype=torch.float32, + quant_config=None, + prefix=prefix, + return_bias=True, + disable_tp=True, + ) + + return layer + + +def build_mla_modules( + hidden_size=32, + q_lora_rank: Optional[int] = 8, + kv_lora_rank=8, + qk_rope_head_dim=8, +): + + total_lora_out = (q_lora_rank or 0) + kv_lora_rank + qk_rope_head_dim + + fused_qkv_a_proj = make_linear(hidden_size, total_lora_out, "fused_") + kv_a_proj_with_mqa = make_linear( + hidden_size, + kv_lora_rank + qk_rope_head_dim, + "kv_", + ) + + q_proj = make_linear(hidden_size, hidden_size, "q_") + + q_b_proj = make_linear(q_lora_rank, hidden_size, "qb_") if q_lora_rank else None + + kv_a_layernorm = LayerNorm(kv_lora_rank) + q_a_layernorm = LayerNorm(q_lora_rank) if q_lora_rank else None + + o_proj = make_linear(hidden_size, hidden_size, "o_") + + modules = MLAModules( + kv_a_layernorm=kv_a_layernorm.to("gcu"), + kv_b_proj=None, + rotary_emb=None, + o_proj=o_proj.to("gcu"), + fused_qkv_a_proj=fused_qkv_a_proj.to("gcu"), + kv_a_proj_with_mqa=kv_a_proj_with_mqa.to("gcu"), + q_a_layernorm=q_a_layernorm.to("gcu") if q_a_layernorm else None, + q_b_proj=q_b_proj.to("gcu") if q_b_proj else None, + q_proj=q_proj.to("gcu"), + indexer_rotary_emb=None, + indexer=None, + is_sparse=False, + topk_indices_buffer=None, + ) + + return modules + + +@pytest.fixture +def mla_dense(): + modules = build_mla_modules( + hidden_size=32, + q_lora_rank=8, + ) + + layer = GCUMultiHeadLatentAttention( + hidden_size=32, + num_heads=1, + scale=1.0, + qk_nope_head_dim=8, + qk_rope_head_dim=8, + v_head_dim=8, + q_lora_rank=8, + kv_lora_rank=8, + mla_modules=modules, + prefix=str(uuid.uuid4()), + ) + + return layer.to("gcu") + + +@pytest.fixture +def mla_no_q_lora(): + modules = build_mla_modules( + hidden_size=32, + q_lora_rank=None, + ) + + layer = GCUMultiHeadLatentAttention( + hidden_size=32, + num_heads=1, + scale=1.0, + qk_nope_head_dim=8, + qk_rope_head_dim=8, + v_head_dim=8, + q_lora_rank=None, + kv_lora_rank=8, + mla_modules=modules, + prefix=str(uuid.uuid4()), + ) + + return layer.to("gcu") + + +def run_forward_check(layer, batch): + torch.set_default_device("gcu") + + hidden = torch.randn(batch, 32, requires_grad=True) + pos = torch.arange(batch) + + out = layer.forward_oot(pos, hidden) + + assert out.dtype == torch.float32 + assert not torch.isnan(out).any() + + out.sum().backward() + assert hidden.grad is not None + + +def test_forward_dense(mla_dense): + run_forward_check(mla_dense, batch=4) + + +def test_forward_no_q_lora(mla_no_q_lora): + run_forward_check(mla_no_q_lora, batch=3) diff --git a/tests/kernel_test/test_mrotary_embedding.py b/tests/kernel_test/test_mrotary_embedding.py new file mode 100644 index 00000000..13444020 --- /dev/null +++ b/tests/kernel_test/test_mrotary_embedding.py @@ -0,0 +1,95 @@ +import pytest + +import torch +import torch_gcu + +from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding +from vllm_gcu.kernels import _custom_ops as ops +from vllm.platforms import current_platform + + +def create_mrope_embedding(head_size, rotary_dim, is_neox, mrope_section, dtype, mrope_interleaved): + """Create a MRotaryEmbedding instance for testing.""" + return MRotaryEmbedding( + head_size=head_size, + rotary_dim=rotary_dim, + max_position_embeddings=8192, + base=10000, + is_neox_style=is_neox, + dtype=dtype, + mrope_section=mrope_section, + mrope_interleaved=mrope_interleaved, + ) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("is_neox", [True]) +@pytest.mark.parametrize("head_size", [128]) +@pytest.mark.parametrize("group_size", [1, 4]) +@pytest.mark.parametrize("mrope_interleaved", [True, False]) +@pytest.mark.parametrize("num_heads", [32]) +@pytest.mark.parametrize( + "num_tokens", [1, 313, 6474, 8192] +) +@pytest.mark.parametrize("mrope_section", [[24, 20, 20]]) +def test_mrope(dtype, num_tokens: int, num_heads: int, head_size: int, group_size: int, is_neox: bool, mrope_section: list[int], mrope_interleaved: bool): + + torch.random.manual_seed(42) + + rotary_dim = head_size + + # Prepare tensors + query = torch.randn(num_tokens, num_heads * head_size, dtype=dtype).gcu() + key = torch.randn(num_tokens, num_heads * head_size // group_size, dtype=dtype).gcu() + + positions = torch.stack([ + torch.randint(0, num_tokens, (num_tokens,), dtype=torch.int64), # T dimension + torch.randint(0, num_tokens, (num_tokens,), dtype=torch.int64), # H dimension + torch.randint(0, num_tokens, (num_tokens,), dtype=torch.int64), # W dimension + ]).gcu().contiguous() # [3, num_tokens] + + # Create MRotaryEmbedding instance with interleaved mode + mrope = create_mrope_embedding(head_size, rotary_dim, is_neox, mrope_section, dtype, mrope_interleaved) + mrope.cos_sin_cache = mrope.cos_sin_cache.gcu() + + # Run kernel + query_kernel = query.clone() + key_kernel = key.clone() + ops.mrotary_embedding( + positions, + query_kernel, + key_kernel, + head_size, + mrope.cos_sin_cache, + is_neox, + mrope_section, + mrope.mrope_interleaved, + ) + + # Run reference implementation using MRotaryEmbedding.forward_native + query_ref = query.clone().cpu() + key_ref = key.clone().cpu() + positions_ref = positions.cpu() + query_ref_result, key_ref_result = mrope.forward_native( + positions_ref, + query_ref, + key_ref, + ) + + # Compare results + if current_platform.get_device_capability().to_int() > 130: + assert torch.allclose( + query_ref_result, query_kernel.cpu(), atol=1e-2, rtol=1e-2 + ), "Query mismatch between kernel and reference" + + assert torch.allclose( + key_ref_result, key_kernel.cpu(), atol=1e-2, rtol=1e-2 + ), "Key mismatch between kernel and reference" + else: + assert torch.allclose( + query_ref_result, query_kernel.cpu(), atol=1e-1, rtol=1e-1 + ), "Query mismatch between kernel and reference" + + assert torch.allclose( + key_ref_result, key_kernel.cpu(), atol=1e-1, rtol=1e-1 + ), "Key mismatch between kernel and reference" diff --git a/tests/kernel_test/test_prepare_finalize.py b/tests/kernel_test/test_prepare_finalize.py new file mode 100644 index 00000000..690a6309 --- /dev/null +++ b/tests/kernel_test/test_prepare_finalize.py @@ -0,0 +1,56 @@ +import pytest +import types +from unittest.mock import patch + +import vllm_gcu.kernels.prepare_finalize as pf + + +@pytest.fixture(autouse=True) +def mock_ep_group(monkeypatch): + class FakeDeviceGroup: + def size(self): + return 1 + + class FakeEPGroup: + device_group = FakeDeviceGroup() + + monkeypatch.setattr(pf, "get_ep_group", lambda: FakeEPGroup()) + + +def test_num_dispatchers(): + num_dispatchers = 4 + + impl = pf.AlltoAllStaticShape( + threshold=1024, + num_dispatchers=num_dispatchers, + ) + + assert impl.num_dispatchers() == num_dispatchers + + +def test_static_max_num_tokens_per_rank(): + threshold = 8192 + num_dispatchers = 2 + + impl = pf.AlltoAllStaticShape( + threshold=threshold, + num_dispatchers=num_dispatchers, + ) + + expected = threshold // num_dispatchers + assert impl.max_num_tokens_per_rank() == expected + + +def test_dynamic_max_num_tokens_per_rank(): + impl = pf.AlltoAllDynamicShape(num_dispatchers=4) + + assert impl.max_num_tokens_per_rank() is None + + +def test_selector_max_num_tokens_per_rank(): + impl = pf.AlltoAllSelector( + threshold=4096, + num_dispatchers=2, + ) + + assert impl.max_num_tokens_per_rank() is None diff --git a/tests/run_all_test.sh b/tests/run_all_test.sh new file mode 100644 index 00000000..233fa8e0 --- /dev/null +++ b/tests/run_all_test.sh @@ -0,0 +1,118 @@ +#!/bin/bash +# +# Copyright 2023-2025 Enflame. All Rights Reserved. +# +# Run all vllm_gcu tests: kernel tests, libtorch C++ tests, CUDA-compat API tests. +# +# Usage: +# bash run_all_test.sh # Run all tests +# bash run_all_test.sh kernel # Run only kernel tests +# +set -e + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +NC='\033[0m' + +PASSED=0 +FAILED=0 +SKIPPED=0 +FAILED_TESTS="" + +run_test() { + local test_name="$1" + local test_cmd="$2" + local test_dir="$3" + local log_name="${test_name//\//_}" + + echo -e "${YELLOW}[RUN]${NC} $test_name" + pushd "$test_dir" > /dev/null 2>&1 || true + + if eval "$test_cmd" > "/tmp/vllm_gcu_test_${log_name}.log" 2>&1; then + echo -e "${GREEN}[PASS]${NC} $test_name" + PASSED=$((PASSED + 1)) + else + echo -e "${RED}[FAIL]${NC} $test_name (see /tmp/vllm_gcu_test_${log_name}.log)" + FAILED=$((FAILED + 1)) + FAILED_TESTS="$FAILED_TESTS\n - $test_name" + fi + + popd > /dev/null 2>&1 || true +} + +# ============================================================================= +# Kernel Tests (Python) +# ============================================================================= +run_kernel_tests() { + echo "" + echo "==============================" + echo " Kernel Tests (Python)" + echo "==============================" + + local kernel_dir="$SCRIPT_DIR/kernel_test" + if [ ! -d "$kernel_dir" ]; then + echo -e "${YELLOW}[SKIP]${NC} kernel_test directory not found" + SKIPPED=$((SKIPPED + 1)) + return + fi + + for test_file in "$kernel_dir"/test_*.py; do + if [ -f "$test_file" ]; then + local test_name="kernel_test/$(basename "$test_file" .py)" + run_test "$test_name" "python3 $(basename "$test_file")" "$kernel_dir" + fi + done +} + +# ============================================================================= +# Main +# ============================================================================= +echo "============================================" +echo " vllm_gcu Test Suite" +echo "============================================" +echo " Test directory: $SCRIPT_DIR" +echo " Date: $(date)" +echo "============================================" + +TARGET="${1:-all}" + +case "$TARGET" in + kernel) + run_kernel_tests + ;; + all) + run_kernel_tests + ;; + *) + echo "Unknown target: $TARGET" + echo "Usage: $0 [all|kernel|libtorch|cuda]" + exit 1 + ;; +esac + +# ============================================================================= +# Summary +# ============================================================================= +echo "" +echo "============================================" +echo " Test Summary" +echo "============================================" +echo -e " ${GREEN}PASSED${NC}: $PASSED" +echo -e " ${RED}FAILED${NC}: $FAILED" +echo -e " ${YELLOW}SKIPPED${NC}: $SKIPPED" + +if [ $FAILED -gt 0 ]; then + echo -e "\n Failed tests:${FAILED_TESTS}" + echo "" + echo "============================================" + exit 1 +else + echo "" + echo -e " ${GREEN}All tests passed!${NC}" + echo "============================================" + exit 0 +fi