maobaolongandLi, Jiang <jiang1.li@intel.com>
13726c80fe
Signed-off-by: baoloongmao <baoloongmao@tencent.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Smoke tests for the CPU MLA backend.
|
|
|
|
Two levels of coverage are provided:
|
|
|
|
* :func:`test_kv_cache_cpu_write` checks that the kv-cache write logic
|
|
inlined in the CPU MLA backend writes the latent KV cache in the same
|
|
layout the CPU decode kernel expects. This runs without downloading
|
|
any model weights and is safe for CI.
|
|
* :func:`test_cpu_mla_backend_smoke` exercises the full end-to-end path
|
|
(LLM engine, model construction, prefill + decode) against a shrunk
|
|
DeepSeek-V2-Lite with `hf_overrides`. It uses `dummy` weights so the
|
|
outputs are not meaningful, only shape / plumbing correctness is
|
|
asserted. Marked `cpu_model` so it is skipped on non-CPU CI.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.platforms import current_platform
|
|
|
|
|
|
@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
|
|
def test_kv_cache_cpu_write() -> None:
|
|
kv_lora_rank = 512
|
|
pe_dim = 64
|
|
block_size = 16
|
|
num_blocks = 4
|
|
num_tokens = 8
|
|
dtype = torch.float32
|
|
|
|
kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype)
|
|
k_pe = torch.randn(num_tokens, pe_dim, dtype=dtype)
|
|
kv_cache = torch.zeros(num_blocks, block_size, kv_lora_rank + pe_dim, dtype=dtype)
|
|
slot_mapping = torch.arange(num_tokens, dtype=torch.long)
|
|
# Mark the last slot as padding to make sure we honour negative slots.
|
|
slot_mapping[-1] = -1
|
|
|
|
flat = kv_cache.view(-1, kv_lora_rank + pe_dim)
|
|
slots = slot_mapping.to(torch.long)
|
|
valid_mask = slots >= 0
|
|
if not valid_mask.all().item():
|
|
slots = slots[valid_mask]
|
|
kv_c = kv_c[valid_mask]
|
|
k_pe = k_pe[valid_mask]
|
|
target_dtype = flat.dtype
|
|
flat[slots, :kv_lora_rank] = kv_c.to(target_dtype)
|
|
flat[slots, kv_lora_rank:] = k_pe.to(target_dtype)
|
|
|
|
flat = kv_cache.view(-1, kv_lora_rank + pe_dim)
|
|
for i in range(num_tokens - 1):
|
|
assert torch.equal(flat[i, :kv_lora_rank], kv_c[i])
|
|
assert torch.equal(flat[i, kv_lora_rank:], k_pe[i])
|
|
# The padding slot must remain untouched.
|
|
assert torch.all(flat[num_tokens - 1] == 0)
|
|
|
|
|
|
@pytest.mark.cpu_model
|
|
@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
|
|
def test_cpu_mla_backend_smoke() -> None:
|
|
# Shrink the architecture aggressively so the test is CI-friendly.
|
|
from vllm import LLM, SamplingParams
|
|
|
|
hf_overrides = {
|
|
"num_hidden_layers": 2,
|
|
"n_routed_experts": 4,
|
|
"first_k_dense_replace": 0,
|
|
"num_experts_per_tok": 2,
|
|
}
|
|
llm = LLM(
|
|
model="deepseek-ai/DeepSeek-V2-Lite",
|
|
trust_remote_code=True,
|
|
load_format="dummy",
|
|
enforce_eager=True,
|
|
max_model_len=128,
|
|
max_num_seqs=2,
|
|
block_size=16,
|
|
hf_overrides=hf_overrides,
|
|
)
|
|
|
|
sampling_params = SamplingParams(max_tokens=4, temperature=0.0)
|
|
outputs = llm.generate(["Hello", "MLA on CPU"], sampling_params)
|
|
assert len(outputs) == 2
|
|
for output in outputs:
|
|
# `dummy` weights do not produce meaningful text, but the number
|
|
# of generated tokens must match what we asked for.
|
|
assert len(output.outputs[0].token_ids) == 4
|