chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,17 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
# Reference default values of atol and rtol are from
|
||||
# https://github.com/pytorch/pytorch/blob/6d96beb6bec24d73ee3f080bac54d2104068f675/test/test_transformers.py#L67
|
||||
default_atol = {torch.float16: 1e-3, torch.bfloat16: 1e-3, torch.float: 1e-5}
|
||||
default_rtol = {torch.float16: 1e-3, torch.bfloat16: 1.6e-2, torch.float: 1.3e-6}
|
||||
|
||||
|
||||
def get_default_atol(output) -> float:
|
||||
return default_atol[output.dtype]
|
||||
|
||||
|
||||
def get_default_rtol(output) -> float:
|
||||
return default_rtol[output.dtype]
|
||||
@@ -0,0 +1,19 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.utils.torch_utils import (
|
||||
create_kv_caches_with_random,
|
||||
create_kv_caches_with_random_flash,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def kv_cache_factory():
|
||||
return create_kv_caches_with_random
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def kv_cache_factory_flashinfer():
|
||||
return create_kv_caches_with_random_flash
|
||||
@@ -0,0 +1,258 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
# Import AITER backend if on ROCm and aiter is available
|
||||
if current_platform.is_rocm():
|
||||
from vllm._aiter_ops import is_aiter_found_and_supported
|
||||
|
||||
if is_aiter_found_and_supported():
|
||||
import aiter
|
||||
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
|
||||
|
||||
NUM_HEADS = [(4, 4), (8, 2)]
|
||||
HEAD_SIZES = [128, 256]
|
||||
BLOCK_SIZES = [16]
|
||||
DTYPES = [torch.bfloat16]
|
||||
QDTYPES = [None]
|
||||
# one value large enough to test overflow in index calculation.
|
||||
# one value small enough to test the schema op check
|
||||
NUM_BLOCKS = [32768, 2048]
|
||||
|
||||
|
||||
def ref_paged_attn(
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
query_lens: list[int],
|
||||
kv_lens: list[int],
|
||||
block_tables: torch.Tensor,
|
||||
scale: float,
|
||||
sliding_window: int | None = None,
|
||||
soft_cap: float | None = None,
|
||||
) -> torch.Tensor:
|
||||
num_seqs = len(query_lens)
|
||||
block_tables = block_tables.cpu().numpy()
|
||||
_, block_size, num_kv_heads, head_size = key_cache.shape
|
||||
|
||||
outputs: list[torch.Tensor] = []
|
||||
start_idx = 0
|
||||
for i in range(num_seqs):
|
||||
query_len = query_lens[i]
|
||||
kv_len = kv_lens[i]
|
||||
q = query[start_idx : start_idx + query_len]
|
||||
q *= scale
|
||||
|
||||
num_kv_blocks = (kv_len + block_size - 1) // block_size
|
||||
block_indices = block_tables[i, :num_kv_blocks]
|
||||
|
||||
k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
|
||||
k = k[:kv_len]
|
||||
v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
|
||||
v = v[:kv_len]
|
||||
|
||||
if q.shape[1] != k.shape[1]:
|
||||
k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
|
||||
v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
|
||||
attn = torch.einsum("qhd,khd->hqk", q, k).float()
|
||||
empty_mask = torch.ones(query_len, kv_len)
|
||||
mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
|
||||
if sliding_window is not None:
|
||||
sliding_window_mask = (
|
||||
torch.triu(
|
||||
empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1
|
||||
)
|
||||
.bool()
|
||||
.logical_not()
|
||||
)
|
||||
mask |= sliding_window_mask
|
||||
if soft_cap is not None:
|
||||
attn = soft_cap * torch.tanh(attn / soft_cap)
|
||||
attn.masked_fill_(mask, float("-inf"))
|
||||
attn = torch.softmax(attn, dim=-1).to(v.dtype)
|
||||
out = torch.einsum("hqk,khd->qhd", attn, v)
|
||||
|
||||
outputs.append(out)
|
||||
start_idx += query_len
|
||||
|
||||
return torch.cat(outputs, dim=0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="Only ROCm is supported")
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens", [[(10, 1328), (5, 18), (129, 463)], [(8, 523), (24, 37), (3, 2011)]]
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("sliding_window", [None, 256])
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("soft_cap", [None])
|
||||
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
|
||||
@pytest.mark.parametrize("q_dtype", QDTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_varlen_with_paged_kv(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
sliding_window: int | None,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
num_blocks: int,
|
||||
q_dtype: torch.dtype | None,
|
||||
) -> None:
|
||||
from vllm._aiter_ops import is_aiter_found_and_supported
|
||||
|
||||
if not is_aiter_found_and_supported():
|
||||
pytest.skip("aiter package required for this test.")
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
set_random_seed(0)
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
|
||||
scale = head_size**-0.5
|
||||
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
|
||||
cu_seq_lens = torch.tensor([0] + kv_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
# Save kv_lens as list before converting to tensor
|
||||
kv_lens_list = kv_lens
|
||||
kv_lens = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
output = torch.empty_like(query)
|
||||
|
||||
maybe_quantized_query = query
|
||||
maybe_quantized_key_cache = key_cache
|
||||
maybe_quantized_value_cache = value_cache
|
||||
k_scale_tensor = None
|
||||
v_scale_tensor = None
|
||||
dequant = False
|
||||
|
||||
if q_dtype is not None:
|
||||
# QKV are drawn from N(0, 1): no need for a fp8 scaling factor
|
||||
maybe_quantized_query = query.to(q_dtype)
|
||||
maybe_quantized_key_cache = key_cache.to(q_dtype)
|
||||
maybe_quantized_value_cache = value_cache.to(q_dtype)
|
||||
dequant = True
|
||||
scale_shape = (num_seqs, num_kv_heads)
|
||||
|
||||
# For per-seq-per-head scales (matching AITER backend expectation)
|
||||
k_scale_tensor = torch.ones(scale_shape, dtype=torch.float32)
|
||||
v_scale_tensor = torch.ones(scale_shape, dtype=torch.float32)
|
||||
|
||||
# Prepare metadata for cp_mha_gather_cache
|
||||
# token_to_batch: maps each token to its batch index
|
||||
token_to_batch = torch.zeros(sum(kv_lens_list), dtype=torch.int32)
|
||||
seq_starts = torch.zeros(num_seqs, dtype=torch.int32)
|
||||
|
||||
token_idx = 0
|
||||
for batch_idx, kv_len in enumerate(kv_lens_list):
|
||||
token_to_batch[token_idx : token_idx + kv_len] = batch_idx
|
||||
seq_starts[batch_idx] = 0 # Assuming all sequences start at 0 in their blocks
|
||||
token_idx += kv_len
|
||||
|
||||
# Allocate buffers for gathered KV
|
||||
total_kv_tokens = sum(kv_lens_list)
|
||||
gathered_key = torch.empty(
|
||||
total_kv_tokens, num_kv_heads, head_size, dtype=maybe_quantized_key_cache.dtype
|
||||
)
|
||||
gathered_value = torch.empty(
|
||||
total_kv_tokens,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
dtype=maybe_quantized_value_cache.dtype,
|
||||
)
|
||||
|
||||
# Gather paged KV cache into contiguous tensors using triton kernel
|
||||
cp_mha_gather_cache(
|
||||
key_cache=maybe_quantized_key_cache,
|
||||
value_cache=maybe_quantized_value_cache,
|
||||
key=gathered_key,
|
||||
value=gathered_value,
|
||||
block_tables=block_tables,
|
||||
k_scales=k_scale_tensor
|
||||
if k_scale_tensor is not None
|
||||
else torch.ones(1, dtype=torch.float32),
|
||||
v_scales=v_scale_tensor
|
||||
if v_scale_tensor is not None
|
||||
else torch.ones(1, dtype=torch.float32),
|
||||
cu_seqlens_kv=cu_seq_lens,
|
||||
token_to_batch=token_to_batch,
|
||||
seq_starts=seq_starts,
|
||||
dequant=dequant,
|
||||
kv_cache_layout="NHD",
|
||||
total_tokens=total_kv_tokens,
|
||||
)
|
||||
|
||||
# Call aiter flash attention with gathered KV
|
||||
aiter.flash_attn_varlen_func(
|
||||
q=maybe_quantized_query,
|
||||
k=gathered_key,
|
||||
v=gathered_value,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
cu_seqlens_k=cu_seq_lens,
|
||||
max_seqlen_q=max_query_len,
|
||||
max_seqlen_k=max_kv_len,
|
||||
min_seqlen_q=1,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
alibi_slopes=None,
|
||||
return_lse=False,
|
||||
out=output,
|
||||
)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens_list,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
sliding_window=sliding_window,
|
||||
soft_cap=soft_cap,
|
||||
)
|
||||
|
||||
atol, rtol = 2e-2, 2e-2
|
||||
if q_dtype is not None:
|
||||
atol, rtol = 1.5e-1, 1.5e-1
|
||||
(
|
||||
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output - ref_output))}",
|
||||
)
|
||||
|
||||
# Log diff stats for tracking changes
|
||||
print(f"Max abs diff: {torch.max(torch.abs(output - ref_output))}")
|
||||
print(f"Mean diff: {torch.mean(torch.abs(output - ref_output))}")
|
||||
print(f"Min diff: {torch.std(torch.abs(output - ref_output))}")
|
||||
@@ -0,0 +1,346 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.allclose_default import get_default_atol, get_default_rtol
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.model_executor.layers.attention import Attention, MMEncoderAttention
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.mem_utils import get_max_shared_memory_bytes
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
|
||||
# This will change depending on the compute capability.
|
||||
# - 512 as a buffer
|
||||
MAX_SEQ_LEN = get_max_shared_memory_bytes() // FLOAT32_BYTES - 512
|
||||
# There may not be enough gpu memory due to large NUM_BLOCKS.
|
||||
# Reduce NUM_BLOCKS when it happens.
|
||||
NUM_BLOCKS = 4321 # Arbitrary values for testing
|
||||
PARTITION_SIZE_ROCM = 256
|
||||
DTYPES = [torch.bfloat16]
|
||||
NUM_GEN_SEQS = [7] # Arbitrary values for testing
|
||||
NUM_PREFILL_SEQS = [3] # Arbitrary values for testing
|
||||
NUM_HEADS = [(32, 8), (40, 40), (64, 8)] # Arbitrary values for testing
|
||||
|
||||
# Head sizes supported by the ROCm paged attention kernel.
|
||||
HEAD_SIZES = [64, 128]
|
||||
|
||||
BLOCK_SIZES = [16, 32]
|
||||
USE_ALIBI = [False, True]
|
||||
KV_CACHE_DTYPE = ["auto", "fp8"]
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
|
||||
|
||||
def ref_masked_attention(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
scale: float,
|
||||
attn_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
attn_weights = scale * torch.einsum("qhd,khd->hqk", query, key).float()
|
||||
if attn_mask is not None:
|
||||
attn_weights = attn_weights + attn_mask.float()
|
||||
attn_weights = torch.softmax(attn_weights, dim=-1).to(value.dtype)
|
||||
out = torch.einsum("hqk,khd->qhd", attn_weights, value)
|
||||
return out
|
||||
|
||||
|
||||
def ref_single_query_cached_kv_attention(
|
||||
output: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
num_queries_per_kv: int,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
scale: float,
|
||||
alibi_slopes: torch.Tensor | None,
|
||||
) -> None:
|
||||
num_query_heads = query.shape[1]
|
||||
num_kv_heads = value_cache.shape[1]
|
||||
head_size = value_cache.shape[2]
|
||||
block_size = value_cache.shape[3]
|
||||
num_seqs = query.shape[0]
|
||||
|
||||
block_tables_lst = block_tables.cpu().tolist()
|
||||
seq_lens_lst = seq_lens.cpu().tolist()
|
||||
for i in range(num_seqs):
|
||||
q = query[i].unsqueeze(0)
|
||||
block_table = block_tables_lst[i]
|
||||
seq_len = int(seq_lens_lst[i])
|
||||
|
||||
keys_lst: list[torch.Tensor] = []
|
||||
values_lst: list[torch.Tensor] = []
|
||||
for j in range(seq_len):
|
||||
block_number = int(block_table[j // block_size])
|
||||
block_offset = j % block_size
|
||||
|
||||
k = key_cache[block_number, :, :, block_offset, :]
|
||||
k = k.reshape(num_kv_heads, head_size)
|
||||
keys_lst.append(k)
|
||||
|
||||
v = value_cache[block_number, :, :, block_offset]
|
||||
values_lst.append(v)
|
||||
keys = torch.stack(keys_lst, dim=0)
|
||||
values = torch.stack(values_lst, dim=0)
|
||||
if num_queries_per_kv > 1:
|
||||
# Handle MQA and GQA
|
||||
keys = torch.repeat_interleave(keys, num_queries_per_kv, dim=1)
|
||||
values = torch.repeat_interleave(values, num_queries_per_kv, dim=1)
|
||||
|
||||
alibi_bias = None
|
||||
if alibi_slopes is not None:
|
||||
# Create the ALiBi bias used in the paged attention kernel.
|
||||
position_ids = torch.arange(seq_len).int()
|
||||
alibi_bias = (position_ids - seq_len + 1).float()
|
||||
alibi_bias = alibi_slopes.view(-1, 1, 1) * alibi_bias.view(1, 1, -1)
|
||||
|
||||
out = ref_masked_attention(q, keys, values, scale, alibi_bias)
|
||||
out = out.view(num_query_heads, head_size)
|
||||
output[i].copy_(out, non_blocking=True)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_rocm(), reason="ROCm-only paged attention kernel"
|
||||
)
|
||||
@pytest.mark.parametrize("num_seqs", NUM_GEN_SEQS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("use_alibi", USE_ALIBI)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_paged_attention(
|
||||
kv_cache_factory,
|
||||
num_seqs: int,
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
use_alibi: bool,
|
||||
block_size: int,
|
||||
dtype: torch.dtype,
|
||||
kv_cache_dtype: str,
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
if current_platform.is_navi() and (
|
||||
kv_cache_dtype == "fp8" or head_size != 128 or block_size != 16 or use_alibi
|
||||
):
|
||||
pytest.skip()
|
||||
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
scale = float(1.0 / (head_size**0.5))
|
||||
num_query_heads, num_kv_heads = num_heads
|
||||
query = torch.empty(num_seqs, num_query_heads, head_size, dtype=dtype)
|
||||
query.uniform_(-scale, scale)
|
||||
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
num_queries_per_kv = num_query_heads // num_kv_heads
|
||||
alibi_slopes = None
|
||||
if use_alibi:
|
||||
alibi_slopes = torch.randn(num_query_heads, dtype=torch.float)
|
||||
|
||||
seq_lens = [random.randint(1, MAX_SEQ_LEN) for _ in range(num_seqs)]
|
||||
seq_lens[-1] = MAX_SEQ_LEN
|
||||
max_seq_len = max(seq_lens)
|
||||
seq_lens = torch.tensor(seq_lens, dtype=torch.int)
|
||||
|
||||
# Create the block tables.
|
||||
max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
|
||||
block_tables_lst: list[list[int]] = []
|
||||
for _ in range(num_seqs):
|
||||
block_table = [
|
||||
random.randint(0, NUM_BLOCKS - 1) for _ in range(max_num_blocks_per_seq)
|
||||
]
|
||||
block_tables_lst.append(block_table)
|
||||
|
||||
block_tables = torch.tensor(block_tables_lst, dtype=torch.int)
|
||||
|
||||
# Create the KV caches.
|
||||
key_caches, value_caches = kv_cache_factory(
|
||||
NUM_BLOCKS,
|
||||
block_size,
|
||||
1,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
kv_cache_dtype,
|
||||
dtype,
|
||||
seed,
|
||||
device,
|
||||
)
|
||||
key_cache, value_cache = key_caches[0], value_caches[0]
|
||||
|
||||
# Using default kv_scale
|
||||
k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
|
||||
|
||||
# Call the paged attention kernel.
|
||||
output = torch.empty_like(query)
|
||||
num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM
|
||||
assert PARTITION_SIZE_ROCM % block_size == 0
|
||||
num_seqs, num_heads, head_size = output.shape
|
||||
tmp_output = torch.empty(
|
||||
size=(num_seqs, num_heads, num_partitions, head_size),
|
||||
dtype=output.dtype,
|
||||
)
|
||||
exp_sums = torch.empty(
|
||||
size=(num_seqs, num_heads, num_partitions),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
max_logits = torch.empty_like(exp_sums)
|
||||
ops.paged_attention_rocm(
|
||||
output,
|
||||
exp_sums,
|
||||
max_logits,
|
||||
tmp_output,
|
||||
query,
|
||||
key_cache,
|
||||
value_cache,
|
||||
num_kv_heads,
|
||||
scale,
|
||||
block_tables,
|
||||
seq_lens,
|
||||
None,
|
||||
block_size,
|
||||
max_seq_len,
|
||||
alibi_slopes,
|
||||
kv_cache_dtype,
|
||||
k_scale,
|
||||
v_scale,
|
||||
)
|
||||
|
||||
opcheck(
|
||||
torch.ops._rocm_C.paged_attention,
|
||||
(
|
||||
output,
|
||||
exp_sums,
|
||||
max_logits,
|
||||
tmp_output,
|
||||
query,
|
||||
key_cache,
|
||||
value_cache,
|
||||
num_kv_heads,
|
||||
scale,
|
||||
block_tables,
|
||||
seq_lens,
|
||||
None,
|
||||
block_size,
|
||||
max_seq_len,
|
||||
alibi_slopes,
|
||||
kv_cache_dtype,
|
||||
k_scale,
|
||||
v_scale,
|
||||
None,
|
||||
"f16",
|
||||
),
|
||||
cond=(head_size == 64 and block_size == BLOCK_SIZES[0]),
|
||||
)
|
||||
|
||||
# Run the reference implementation.
|
||||
if kv_cache_dtype == "fp8":
|
||||
# Convert cache data back to dtype.
|
||||
x = 16 // torch.tensor([], dtype=dtype).element_size()
|
||||
key_cache_shape = (NUM_BLOCKS, num_kv_heads, head_size // x, block_size, x)
|
||||
dequantized_key_cache = torch.empty(
|
||||
size=key_cache_shape, dtype=dtype, device=device
|
||||
)
|
||||
ops.convert_fp8(dequantized_key_cache, key_cache)
|
||||
key_cache = dequantized_key_cache
|
||||
|
||||
value_cache_shape = value_cache.shape
|
||||
dequantized_value_cache = torch.empty(
|
||||
size=value_cache_shape, dtype=dtype, device=device
|
||||
)
|
||||
ops.convert_fp8(dequantized_value_cache, value_cache)
|
||||
value_cache = dequantized_value_cache
|
||||
|
||||
ref_output = torch.empty_like(query)
|
||||
ref_single_query_cached_kv_attention(
|
||||
ref_output,
|
||||
query,
|
||||
num_queries_per_kv,
|
||||
key_cache,
|
||||
value_cache,
|
||||
block_tables,
|
||||
seq_lens,
|
||||
scale,
|
||||
alibi_slopes,
|
||||
)
|
||||
|
||||
# NOTE(woosuk): Due to the kernel-level differences in the two
|
||||
# implementations, there is a small numerical difference in the two
|
||||
# outputs. Thus, we use a relaxed tolerance for the test.
|
||||
atol = get_default_atol(output) if current_platform.is_rocm() else 1e-3
|
||||
rtol = get_default_rtol(output) if current_platform.is_rocm() else 1e-5
|
||||
|
||||
# NOTE(zhaoyang): FP8 KV Cache will introduce quantization error,
|
||||
# so we use a relaxed tolerance for the test.
|
||||
atol, rtol = 1e-3, 1e-5
|
||||
if kv_cache_dtype == "fp8":
|
||||
atol, rtol = 1e-2, 1e-5
|
||||
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
def ref_multi_query_kv_attention(
|
||||
cu_seq_lens: list[int],
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
scale: float,
|
||||
alibi_bias: list[torch.Tensor] | None,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
num_seqs = len(cu_seq_lens) - 1
|
||||
ref_outputs: list[torch.Tensor] = []
|
||||
if alibi_bias:
|
||||
assert len(alibi_bias) == num_seqs
|
||||
for i in range(num_seqs):
|
||||
start_idx = cu_seq_lens[i]
|
||||
end_idx = cu_seq_lens[i + 1]
|
||||
seq_len = end_idx - start_idx
|
||||
|
||||
# Create attention mask. ALiBi already includes a tril causal mask.
|
||||
if alibi_bias:
|
||||
attn_mask = alibi_bias[i]
|
||||
else:
|
||||
attn_mask = torch.triu(
|
||||
torch.ones(seq_len, seq_len, dtype=dtype), diagonal=1
|
||||
)
|
||||
attn_mask = attn_mask * torch.finfo(dtype).min
|
||||
attn_mask = attn_mask.to(dtype=dtype)
|
||||
|
||||
ref_output = ref_masked_attention(
|
||||
query[start_idx:end_idx],
|
||||
key[start_idx:end_idx],
|
||||
value[start_idx:end_idx],
|
||||
scale,
|
||||
attn_mask=attn_mask,
|
||||
)
|
||||
ref_outputs.append(ref_output)
|
||||
|
||||
return torch.cat(ref_outputs, dim=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("attention_cls", [Attention, MMEncoderAttention])
|
||||
def test_num_heads_not_divisible_by_num_kv_heads(attention_cls: type) -> None:
|
||||
head_size = 64
|
||||
scale = float(1.0 / (head_size**0.5))
|
||||
num_heads = 16
|
||||
num_kv_heads = 5
|
||||
with pytest.raises(AssertionError):
|
||||
_ = attention_cls(
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
scale=scale,
|
||||
num_kv_heads=num_kv_heads,
|
||||
)
|
||||
@@ -0,0 +1,553 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config import (
|
||||
AttentionConfig,
|
||||
CacheConfig,
|
||||
VllmConfig,
|
||||
set_current_vllm_config,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.platforms.cpu import CpuPlatform
|
||||
|
||||
if current_platform.is_cuda():
|
||||
from vllm.platforms.cuda import CudaPlatform
|
||||
else:
|
||||
CudaPlatform = None
|
||||
|
||||
if current_platform.is_rocm():
|
||||
from vllm.platforms.rocm import RocmPlatform
|
||||
else:
|
||||
RocmPlatform = None
|
||||
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.v1.attention.selector import _cached_get_attn_backend, get_attn_backend
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_cache():
|
||||
"""Clear lru cache to ensure each test case runs without caching."""
|
||||
_cached_get_attn_backend.cache_clear()
|
||||
|
||||
|
||||
# Define MLA and non-MLA backends separately
|
||||
DEVICE_MLA_BACKENDS = {
|
||||
"cuda": [
|
||||
"TRITON_MLA",
|
||||
"FLASHMLA",
|
||||
"FLASHINFER_MLA",
|
||||
"FLASH_ATTN_MLA",
|
||||
"CUTLASS_MLA",
|
||||
],
|
||||
"hip": ["TRITON_MLA", "ROCM_AITER_MLA"],
|
||||
"cpu": [],
|
||||
}
|
||||
|
||||
DEVICE_REGULAR_ATTN_BACKENDS = {
|
||||
"cuda": ["FLASHINFER", "FLASH_ATTN"],
|
||||
"hip": ["ROCM_ATTN"],
|
||||
"cpu": ["CPU_ATTN"],
|
||||
}
|
||||
|
||||
DEVICE_MLA_BLOCK_SIZES = {
|
||||
"cuda": [16, 64], # CUDA supports both standard and extended block sizes
|
||||
"hip": [16, 1], # HIP requires special handling for block_size=1
|
||||
# "cpu": [16] # CPU uses fixed block size from test cases
|
||||
"cpu": [], # FIXME(woosuk): Temporarily disable CPU tests
|
||||
}
|
||||
|
||||
|
||||
def generate_params():
|
||||
is_rocm = current_platform.is_rocm()
|
||||
params = []
|
||||
device_list = ["cuda", "cpu"] if not is_rocm else ["hip", "cpu"]
|
||||
for use_mla in [True, False]:
|
||||
for device in device_list:
|
||||
backends = (
|
||||
DEVICE_MLA_BACKENDS[device]
|
||||
if use_mla
|
||||
else DEVICE_REGULAR_ATTN_BACKENDS[device]
|
||||
)
|
||||
for name in backends:
|
||||
block_sizes = DEVICE_MLA_BLOCK_SIZES[device] if use_mla else [16]
|
||||
for block_size in block_sizes:
|
||||
params.append(
|
||||
pytest.param(
|
||||
device,
|
||||
name,
|
||||
use_mla,
|
||||
block_size,
|
||||
id=f"{device}_{name}_mla_{str(use_mla)[0]}_blks{block_size}",
|
||||
)
|
||||
)
|
||||
return params
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device, name, use_mla, block_size", generate_params())
|
||||
def test_backend_selection(
|
||||
device: str,
|
||||
name: str,
|
||||
use_mla: bool,
|
||||
block_size: int,
|
||||
):
|
||||
"""Test attention backend selection with valid device-backend pairs."""
|
||||
# Create AttentionConfig with the specified backend
|
||||
attention_config = AttentionConfig(backend=AttentionBackendEnum[name])
|
||||
cache_config = CacheConfig(block_size=block_size)
|
||||
vllm_config = VllmConfig(
|
||||
attention_config=attention_config, cache_config=cache_config
|
||||
)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
if device == "cpu":
|
||||
with patch("vllm.platforms.current_platform", CpuPlatform()):
|
||||
backend = get_attn_backend(16, torch.float16, None)
|
||||
assert backend.get_name() == "CPU_ATTN"
|
||||
|
||||
elif device == "hip":
|
||||
if RocmPlatform is None:
|
||||
pytest.skip("RocmPlatform not available")
|
||||
with patch("vllm.platforms.current_platform", RocmPlatform()):
|
||||
if use_mla:
|
||||
# ROCm MLA backend logic:
|
||||
# - TRITON_MLA: supported when block_size != 1
|
||||
# - ROCM_AITER_MLA: supported when block_size == 1
|
||||
# If backend is forced but doesn't match block_size,
|
||||
# should raise ValueError
|
||||
|
||||
if name == "TRITON_MLA" and block_size == 1:
|
||||
# TRITON_MLA doesn't support block_size == 1
|
||||
with pytest.raises(ValueError):
|
||||
get_attn_backend(576, torch.float16, None, use_mla=use_mla)
|
||||
else:
|
||||
# Valid backend-block_size combination
|
||||
backend = get_attn_backend(
|
||||
576, torch.float16, None, use_mla=use_mla
|
||||
)
|
||||
expected = name
|
||||
assert backend.get_name() == expected
|
||||
else:
|
||||
backend = get_attn_backend(32, torch.float16, None, use_mla=use_mla)
|
||||
expected = "ROCM_ATTN"
|
||||
assert backend.get_name() == expected
|
||||
|
||||
elif device == "cuda":
|
||||
if CudaPlatform is None:
|
||||
pytest.skip("CudaPlatform not available")
|
||||
with patch("vllm.platforms.current_platform", CudaPlatform()):
|
||||
capability = torch.cuda.get_device_capability()
|
||||
if use_mla:
|
||||
# CUDA MLA backend logic:
|
||||
# - CUTLASS_MLA: only supported with block_size == 128
|
||||
# and Blackwell GPUs (SM 10.x), V1 only
|
||||
# - FLASHINFER_MLA: only supported on Blackwell GPUs
|
||||
# (SM 10.x), V1 only
|
||||
# - FLASHMLA: only supported with block_size == 64
|
||||
# - FLASH_ATTN_MLA: V1 only
|
||||
# - TRITON_MLA: fallback for other cases
|
||||
|
||||
if name == "CUTLASS_MLA":
|
||||
if block_size != 128:
|
||||
# CUTLASS_MLA only supports block_size == 128
|
||||
pytest.skip("CUTLASS_MLA only supports block_size 128")
|
||||
if capability[0] != 10:
|
||||
pytest.skip("CUTLASS MLA is not supported on this platform")
|
||||
backend = get_attn_backend(
|
||||
576, torch.float16, None, use_mla=use_mla
|
||||
)
|
||||
expected = "CUTLASS_MLA"
|
||||
assert backend.get_name() == expected
|
||||
elif name == "FLASHINFER_MLA":
|
||||
if capability[0] != 10:
|
||||
pytest.skip(
|
||||
"FlashInfer MLA is not supported on this platform"
|
||||
)
|
||||
if block_size not in [32, 64]:
|
||||
# FlashInfer MLA only supports block_size 32 or 64
|
||||
pytest.skip(
|
||||
"FlashInfer MLA only supports block_size 32 or 64"
|
||||
)
|
||||
backend = get_attn_backend(
|
||||
576, torch.float16, None, use_mla=use_mla
|
||||
)
|
||||
expected = "FLASHINFER_MLA"
|
||||
assert backend.get_name() == expected
|
||||
elif name == "FLASHMLA":
|
||||
if block_size != 64:
|
||||
# FlashMLA only supports block_size == 64
|
||||
pytest.skip("FlashMLA only supports block_size 64")
|
||||
from vllm.v1.attention.backends.mla.flashmla import (
|
||||
is_flashmla_dense_supported,
|
||||
)
|
||||
|
||||
is_supported, _ = is_flashmla_dense_supported()
|
||||
if not is_supported:
|
||||
pytest.skip("FlashMLA not supported on this platform")
|
||||
backend = get_attn_backend(
|
||||
576,
|
||||
torch.float16,
|
||||
None,
|
||||
use_mla=use_mla,
|
||||
)
|
||||
expected = name
|
||||
assert backend.get_name() == expected
|
||||
elif name == "FLASH_ATTN_MLA":
|
||||
from vllm.v1.attention.backends.fa_utils import (
|
||||
flash_attn_supports_mla,
|
||||
)
|
||||
|
||||
if not flash_attn_supports_mla():
|
||||
pytest.skip(
|
||||
"FlashAttention MLA not supported on this platform"
|
||||
)
|
||||
backend = get_attn_backend(
|
||||
576, torch.float16, None, use_mla=use_mla
|
||||
)
|
||||
expected = "FLASH_ATTN_MLA"
|
||||
assert backend.get_name() == expected
|
||||
else:
|
||||
# TRITON_MLA or other fallback
|
||||
backend = get_attn_backend(
|
||||
576, torch.float16, None, use_mla=use_mla
|
||||
)
|
||||
expected = "TRITON_MLA"
|
||||
assert backend.get_name() == expected
|
||||
elif name == "FLASHINFER":
|
||||
backend = get_attn_backend(64, torch.float16, None, use_mla=use_mla)
|
||||
expected = "FLASHINFER"
|
||||
assert backend.get_name() == expected
|
||||
elif name == "FLASH_ATTN":
|
||||
backend = get_attn_backend(32, torch.float16, None, use_mla=use_mla)
|
||||
expected = "FLASH_ATTN"
|
||||
assert backend.get_name() == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", ["cpu", "cuda", "hip"])
|
||||
def test_fp32_fallback(device: str):
|
||||
"""Test attention backend selection with fp32."""
|
||||
# Use default config (no backend specified)
|
||||
vllm_config = VllmConfig()
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
if device == "cpu":
|
||||
with patch("vllm.platforms.current_platform", CpuPlatform()):
|
||||
backend = get_attn_backend(16, torch.float32, None)
|
||||
assert backend.get_name() == "CPU_ATTN"
|
||||
|
||||
elif device == "cuda":
|
||||
if CudaPlatform is None:
|
||||
pytest.skip("CudaPlatform not available")
|
||||
with patch("vllm.platforms.current_platform", CudaPlatform()):
|
||||
backend = get_attn_backend(16, torch.float32, None)
|
||||
assert backend.get_name() == "FLEX_ATTENTION"
|
||||
|
||||
elif device == "hip":
|
||||
if RocmPlatform is None:
|
||||
pytest.skip("RocmPlatform not available")
|
||||
# ROCm backends do not support head_size=16 (minimum is 32).
|
||||
# No known HuggingFace transformer model uses head_size=16.
|
||||
# Revisit if a real model with this head size is identified
|
||||
# and accuracy-tested.
|
||||
with (
|
||||
patch("vllm.platforms.current_platform", RocmPlatform()),
|
||||
pytest.raises(ValueError, match="No valid attention backend"),
|
||||
):
|
||||
get_attn_backend(16, torch.float32, None)
|
||||
|
||||
|
||||
def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test FlashAttn validation."""
|
||||
pytest.skip(
|
||||
"Skipping as current backend selector does not "
|
||||
"handle fallbacks when a backend is explicitly set."
|
||||
)
|
||||
|
||||
attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASH_ATTN)
|
||||
cache_config = CacheConfig(block_size=16)
|
||||
vllm_config = VllmConfig(
|
||||
attention_config=attention_config, cache_config=cache_config
|
||||
)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
# Unsupported CUDA arch
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _=None: (7, 5))
|
||||
backend = get_attn_backend(16, torch.float16, None)
|
||||
assert backend.get_name() != "FLASH_ATTN"
|
||||
|
||||
# Reset the monkeypatch for subsequent tests
|
||||
monkeypatch.undo()
|
||||
|
||||
# Unsupported data type
|
||||
backend = get_attn_backend(16, torch.float8_e4m3fn, None)
|
||||
assert backend.get_name() != "FLASH_ATTN"
|
||||
|
||||
# Unsupported kv cache data type
|
||||
backend = get_attn_backend(16, torch.float16, "fp8")
|
||||
assert backend.get_name() != "FLASH_ATTN"
|
||||
|
||||
# Unsupported block size
|
||||
vllm_config.cache_config.block_size = 8
|
||||
backend = get_attn_backend(16, torch.float16, None)
|
||||
assert backend.get_name() != "FLASH_ATTN"
|
||||
|
||||
# flash-attn is not installed
|
||||
import sys
|
||||
|
||||
vllm_config.cache_config.block_size = 16
|
||||
original_module = sys.modules.get("vllm_flash_attn")
|
||||
monkeypatch.setitem(sys.modules, "vllm_flash_attn", None)
|
||||
backend = get_attn_backend(16, torch.float16, None)
|
||||
assert backend.get_name() != "FLASH_ATTN"
|
||||
|
||||
# Restore the original module if it existed
|
||||
if original_module is not None:
|
||||
monkeypatch.setitem(sys.modules, "vllm_flash_attn", original_module)
|
||||
else:
|
||||
monkeypatch.delitem(sys.modules, "vllm_flash_attn", raising=False)
|
||||
|
||||
# Unsupported head size
|
||||
backend = get_attn_backend(17, torch.float16, None)
|
||||
assert backend.get_name() != "FLASH_ATTN"
|
||||
|
||||
|
||||
def test_invalid_backend():
|
||||
"""Test that invalid attention backend names raise ValueError."""
|
||||
with (
|
||||
pytest.raises(ValueError),
|
||||
):
|
||||
# Invalid backend name should raise ValueError when creating enum
|
||||
AttentionConfig(backend=AttentionBackendEnum["INVALID"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("auto_value", ["auto", "AUTO", "Auto"])
|
||||
def test_auto_backend_string(auto_value: str):
|
||||
"""Test that 'auto' string value triggers automatic backend selection."""
|
||||
# Using "auto" should result in backend=None (automatic selection)
|
||||
attention_config = AttentionConfig(backend=auto_value)
|
||||
assert attention_config.backend is None
|
||||
|
||||
|
||||
def test_auto_backend_selection_behavior():
|
||||
"""Test that 'auto' backend behaves same as None (automatic selection)."""
|
||||
# Create config with explicit "auto"
|
||||
auto_config = AttentionConfig(backend="auto")
|
||||
|
||||
# Create config with None (default)
|
||||
none_config = AttentionConfig(backend=None)
|
||||
|
||||
# Both should have backend=None
|
||||
assert auto_config.backend is None
|
||||
assert none_config.backend is None
|
||||
|
||||
# Both configs should result in the same automatic backend selection
|
||||
vllm_config_auto = VllmConfig(attention_config=auto_config)
|
||||
vllm_config_none = VllmConfig(attention_config=none_config)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config_auto),
|
||||
patch("vllm.platforms.current_platform", CpuPlatform()),
|
||||
):
|
||||
backend_auto = get_attn_backend(16, torch.float16, None)
|
||||
|
||||
_cached_get_attn_backend.cache_clear()
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config_none),
|
||||
patch("vllm.platforms.current_platform", CpuPlatform()),
|
||||
):
|
||||
backend_none = get_attn_backend(16, torch.float16, None)
|
||||
|
||||
# Both should select the same backend
|
||||
assert backend_auto.get_name() == backend_none.get_name()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"backend_name,flash_attn_version,should_succeed",
|
||||
[
|
||||
("FLASH_ATTN", 3, True), # FA3 supports per-head quant scales
|
||||
("FLASH_ATTN", 2, False), # FA2 does not support per-head quant scales
|
||||
("FLASHINFER", None, False), # FlashInfer does not support
|
||||
("FLEX_ATTENTION", None, False), # Flex does not support
|
||||
],
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_rocm(),
|
||||
reason="Attention backend FA3 is not supported on ROCm. This test can't succeed.",
|
||||
)
|
||||
def test_per_head_quant_scales_backend_selection(
|
||||
backend_name: str, flash_attn_version: int | None, should_succeed: bool
|
||||
):
|
||||
"""Test backend selection when use_per_head_quant_scales=True."""
|
||||
# Clear cache to ensure fresh backend selection
|
||||
_cached_get_attn_backend.cache_clear()
|
||||
|
||||
attention_config = AttentionConfig(
|
||||
backend=AttentionBackendEnum[backend_name],
|
||||
flash_attn_version=flash_attn_version,
|
||||
)
|
||||
cache_config = CacheConfig(block_size=64)
|
||||
vllm_config = VllmConfig(
|
||||
attention_config=attention_config, cache_config=cache_config
|
||||
)
|
||||
|
||||
if CudaPlatform is None:
|
||||
pytest.skip("CudaPlatform not available")
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch("vllm.platforms.current_platform", CudaPlatform()),
|
||||
):
|
||||
if backend_name == "FLASH_ATTN" and flash_attn_version == 3:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("FA3 requires CUDA")
|
||||
capability = torch.cuda.get_device_capability()
|
||||
if capability[0] != 9:
|
||||
pytest.skip("FA3 is only supported on Hopper (SM 9.x) GPUs")
|
||||
|
||||
if should_succeed:
|
||||
backend = get_attn_backend(
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype="fp8",
|
||||
use_per_head_quant_scales=True,
|
||||
)
|
||||
assert backend.get_name() == backend_name
|
||||
else:
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
get_attn_backend(
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype="fp8",
|
||||
use_per_head_quant_scales=True,
|
||||
)
|
||||
assert backend_name in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"backend_name,use_non_causal,should_succeed",
|
||||
[
|
||||
("FLASH_ATTN", True, True), # FlashAttn supports non-causal
|
||||
("FLASH_ATTN", False, True), # FlashAttn also works with causal
|
||||
]
|
||||
+ (
|
||||
[
|
||||
("FLASHINFER", True, True), # FlashInfer supports non-causal
|
||||
("FLASHINFER", False, True), # FlashInfer works with causal
|
||||
]
|
||||
if CudaPlatform is not None
|
||||
else []
|
||||
),
|
||||
)
|
||||
def test_non_causal_backend_selection(
|
||||
backend_name: str, use_non_causal: bool, should_succeed: bool
|
||||
):
|
||||
"""Test that use_non_causal on AttentionConfig controls backend filtering.
|
||||
|
||||
DFlashProposer sets use_non_causal=True on the draft model's
|
||||
AttentionConfig so only non-causal-capable backends are selected.
|
||||
The target model keeps use_non_causal=False (default) and can use
|
||||
any backend.
|
||||
"""
|
||||
_cached_get_attn_backend.cache_clear()
|
||||
|
||||
attention_config = AttentionConfig(
|
||||
backend=AttentionBackendEnum[backend_name],
|
||||
use_non_causal=use_non_causal,
|
||||
)
|
||||
cache_config = CacheConfig(block_size=16)
|
||||
vllm_config = VllmConfig(
|
||||
attention_config=attention_config, cache_config=cache_config
|
||||
)
|
||||
|
||||
platform = CudaPlatform or RocmPlatform
|
||||
if platform is None:
|
||||
pytest.skip("CudaPlatform and RocmPlatform are not available")
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch("vllm.platforms.current_platform", platform()),
|
||||
):
|
||||
if should_succeed:
|
||||
backend = get_attn_backend(
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype=None,
|
||||
)
|
||||
assert backend.get_name() == backend_name
|
||||
else:
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
get_attn_backend(
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype=None,
|
||||
)
|
||||
assert "non-causal" in str(exc_info.value).lower()
|
||||
|
||||
|
||||
def test_non_causal_autoselect_backend():
|
||||
"""Test that when backend=None with use_non_causal=True, auto-selection
|
||||
picks a compatible backend.
|
||||
|
||||
This simulates the DFlash scenario where the user doesn't specify
|
||||
--attention-backend or --speculative-config.attention_backend.
|
||||
The drafter inherits backend=None and auto-selects a backend that
|
||||
supports non-causal attention.
|
||||
"""
|
||||
_cached_get_attn_backend.cache_clear()
|
||||
|
||||
attention_config = AttentionConfig(
|
||||
backend=None,
|
||||
use_non_causal=True,
|
||||
)
|
||||
cache_config = CacheConfig(block_size=16)
|
||||
vllm_config = VllmConfig(
|
||||
attention_config=attention_config, cache_config=cache_config
|
||||
)
|
||||
|
||||
if CudaPlatform is None:
|
||||
pytest.skip("CudaPlatform not available")
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch("vllm.platforms.current_platform", CudaPlatform()),
|
||||
):
|
||||
backend = get_attn_backend(
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype=None,
|
||||
)
|
||||
assert backend.supports_non_causal()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kv_cache_dtype",
|
||||
[
|
||||
"fp8_e5m2",
|
||||
"fp8_ds_mla",
|
||||
"fp8_inc",
|
||||
"nvfp4",
|
||||
"fp8_per_token_head",
|
||||
"int8_per_token_head",
|
||||
],
|
||||
)
|
||||
def test_flash_attn_rejects_unhandled_kv_cache_dtypes(kv_cache_dtype: str):
|
||||
"""FlashAttentionBackend must not claim support for kv_cache dtypes
|
||||
that it cannot handle."""
|
||||
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
|
||||
|
||||
assert not FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["fp8", "fp8_e4m3"])
|
||||
def test_flash_attn_accepts_handled_fp8_variants(
|
||||
kv_cache_dtype: str, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
"""FlashAttentionBackend must accept the two fp8 dtypes it can actually
|
||||
handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'."""
|
||||
import vllm.v1.attention.backends.flash_attn as fa_mod
|
||||
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
|
||||
|
||||
monkeypatch.setattr(fa_mod.current_platform, "is_xpu", lambda: True)
|
||||
assert FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
|
||||
File diff suppressed because it is too large
Load Diff
+187
@@ -0,0 +1,187 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.v1.attention.backends.flash_attn import cascade_attention, merge_attn_states
|
||||
|
||||
try:
|
||||
from vllm.vllm_flash_attn import (
|
||||
fa_version_unsupported_reason,
|
||||
flash_attn_varlen_func,
|
||||
is_fa_version_supported,
|
||||
)
|
||||
except ImportError:
|
||||
if current_platform.is_rocm():
|
||||
pytest.skip(
|
||||
"vllm_flash_attn is not supported for vLLM on ROCm.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
NUM_HEADS = [(4, 4), (8, 2), (16, 2)]
|
||||
HEAD_SIZES = [128, 192, 256]
|
||||
BLOCK_SIZES = [16]
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 39, 16912])
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_merge_kernel(
|
||||
num_tokens: int,
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
torch.set_default_device("cuda")
|
||||
set_random_seed(0)
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
|
||||
# Prepare inputs.
|
||||
prefix_output = torch.randn(num_tokens, num_query_heads, head_size, dtype=dtype)
|
||||
suffix_output = torch.randn(num_tokens, num_query_heads, head_size, dtype=dtype)
|
||||
prefix_lse = torch.randn(num_query_heads, num_tokens, dtype=torch.float32)
|
||||
suffix_lse = torch.randn(num_query_heads, num_tokens, dtype=torch.float32)
|
||||
|
||||
# Run the kernel.
|
||||
output = torch.empty(num_tokens, num_query_heads, head_size, dtype=dtype)
|
||||
merge_attn_states(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
|
||||
|
||||
# Reference implementation.
|
||||
max_lse = torch.maximum(prefix_lse, suffix_lse)
|
||||
p_lse = torch.exp(prefix_lse - max_lse)
|
||||
s_lse = torch.exp(suffix_lse - max_lse)
|
||||
p_scale = p_lse / (p_lse + s_lse)
|
||||
s_scale = s_lse / (p_lse + s_lse)
|
||||
p_scale = p_scale.transpose(0, 1).unsqueeze(2)
|
||||
s_scale = s_scale.transpose(0, 1).unsqueeze(2)
|
||||
ref_output = p_scale * prefix_output + s_scale * suffix_output
|
||||
ref_output = ref_output.to(dtype)
|
||||
|
||||
# Compare the results.
|
||||
torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
CASES = [
|
||||
# Case 1. A general case.
|
||||
([(129, 871), (18, 280), (37, 988), (1023, 2304), (1, 257)], 256),
|
||||
# Case 2. Flash-decoding case.
|
||||
([(1, 1023), (1, 879), (1, 778), (1, 1777)] * 100, 512),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq_lens_and_common_prefix", CASES)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("soft_cap", [None, 50])
|
||||
@pytest.mark.parametrize("num_blocks", [2048])
|
||||
@pytest.mark.parametrize("fa_version", [2, 3])
|
||||
@torch.inference_mode()
|
||||
def test_cascade(
|
||||
seq_lens_and_common_prefix: tuple[list[tuple[int, int]], int],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
num_blocks: int,
|
||||
fa_version: int,
|
||||
) -> None:
|
||||
torch.set_default_device("cuda")
|
||||
if not is_fa_version_supported(fa_version):
|
||||
pytest.skip(
|
||||
f"Flash attention version {fa_version} not supported due "
|
||||
f'to: "{fa_version_unsupported_reason(fa_version)}"'
|
||||
)
|
||||
|
||||
set_random_seed(0)
|
||||
|
||||
window_size = (-1, -1)
|
||||
scale = head_size**-0.5
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
|
||||
seq_lens, common_prefix_len = seq_lens_and_common_prefix
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
|
||||
total_num_query_tokens = sum(query_lens)
|
||||
query = torch.randn(total_num_query_tokens, num_query_heads, head_size, dtype=dtype)
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
assert common_prefix_len > 0
|
||||
assert common_prefix_len % block_size == 0
|
||||
num_common_kv_blocks = common_prefix_len // block_size
|
||||
# Make sure the first `num_common_kv_blocks` blocks are the same.
|
||||
block_tables[:, :num_common_kv_blocks] = block_tables[0, :num_common_kv_blocks]
|
||||
|
||||
# Run the regular attention.
|
||||
ref_output = flash_attn_varlen_func(
|
||||
q=query,
|
||||
k=key_cache,
|
||||
v=value_cache,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
seqused_k=kv_lens_tensor,
|
||||
max_seqlen_q=max_query_len,
|
||||
max_seqlen_k=max_kv_len,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
block_table=block_tables,
|
||||
softcap=soft_cap if soft_cap is not None else 0,
|
||||
)
|
||||
|
||||
# Run cascade attention.
|
||||
assert all(common_prefix_len < kv_len for kv_len in kv_lens)
|
||||
cu_prefix_query_lens = torch.tensor([0, total_num_query_tokens], dtype=torch.int32)
|
||||
prefix_kv_lens = torch.tensor([common_prefix_len], dtype=torch.int32)
|
||||
suffix_kv_lens = kv_lens_tensor - common_prefix_len
|
||||
output = torch.empty_like(query)
|
||||
cascade_attention(
|
||||
output=output,
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
cu_query_lens=cu_query_lens,
|
||||
max_query_len=max_query_len,
|
||||
cu_prefix_query_lens=cu_prefix_query_lens,
|
||||
prefix_kv_lens=prefix_kv_lens,
|
||||
suffix_kv_lens=suffix_kv_lens,
|
||||
max_kv_len=max_kv_len,
|
||||
softmax_scale=scale,
|
||||
alibi_slopes=None,
|
||||
sliding_window=window_size,
|
||||
logits_soft_cap=soft_cap if soft_cap is not None else 0,
|
||||
block_table=block_tables,
|
||||
common_prefix_len=common_prefix_len,
|
||||
max_num_splits=0, # no max
|
||||
fa_version=fa_version,
|
||||
)
|
||||
|
||||
# Compare the results.
|
||||
torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,280 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import math
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import triton
|
||||
from vllm.utils.platform_utils import num_compute_units
|
||||
|
||||
|
||||
def cal_diff(
|
||||
x: torch.Tensor,
|
||||
y: torch.Tensor,
|
||||
name: str,
|
||||
use_fp8: bool = False,
|
||||
diff_threshold: float | None = None,
|
||||
) -> None:
|
||||
x, y = x.double(), y.double()
|
||||
cos_diff = 1 - 2 * (x * y).sum().item() / max((x * x + y * y).sum().item(), 1e-12)
|
||||
if diff_threshold is not None:
|
||||
# directly compare the cos_diff with the threshold
|
||||
assert cos_diff < diff_threshold
|
||||
else:
|
||||
# use the default threshold
|
||||
if use_fp8:
|
||||
assert cos_diff < 1e-4
|
||||
else:
|
||||
assert cos_diff < 1e-5
|
||||
|
||||
|
||||
CUTLASS_MLA_UNSUPPORTED_REASON = (
|
||||
"Cutlass MLA Requires compute capability of 100 or above."
|
||||
if not current_platform.is_device_capability_family(100)
|
||||
else "Cutlass MLA is supported"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(100),
|
||||
reason=CUTLASS_MLA_UNSUPPORTED_REASON,
|
||||
)
|
||||
@pytest.mark.parametrize("b", [128])
|
||||
@pytest.mark.parametrize("s_q", [1])
|
||||
@pytest.mark.parametrize("mean_sk", [4096, 8192, 16384])
|
||||
@pytest.mark.parametrize("h_q", [16, 32, 64, 128])
|
||||
@pytest.mark.parametrize("h_kv", [1])
|
||||
@pytest.mark.parametrize("d", [576])
|
||||
@pytest.mark.parametrize("dv", [512])
|
||||
@pytest.mark.parametrize("block_size", [64])
|
||||
@pytest.mark.parametrize("causal", [True])
|
||||
@pytest.mark.parametrize("varlen", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"torch_dtype",
|
||||
[
|
||||
torch.bfloat16,
|
||||
# fp8 can have occasional precision-related failures.
|
||||
pytest.param(torch.float8_e4m3fn, marks=pytest.mark.flaky(reruns=2)),
|
||||
],
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_cutlass_mla_decode(
|
||||
b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal, varlen, torch_dtype
|
||||
):
|
||||
device = torch.device("cuda:0")
|
||||
init_dtype = torch.bfloat16 if torch_dtype == torch.float8_e4m3fn else torch_dtype
|
||||
torch.set_default_dtype(init_dtype)
|
||||
torch.set_default_device(device)
|
||||
torch.accelerator.set_device_index(device)
|
||||
torch.manual_seed(42)
|
||||
random.seed(42)
|
||||
|
||||
print(
|
||||
f"{b=}, {s_q=}, {mean_sk=}, {h_q=}, {h_kv=}, "
|
||||
f"{d=}, {dv=}, {causal=}, {varlen=}, {torch_dtype=}"
|
||||
)
|
||||
|
||||
use_fp8 = torch_dtype == torch.float8_e4m3fn
|
||||
scale = math.sqrt(d) ** (-1)
|
||||
cache_seqlens = torch.full((b,), mean_sk, dtype=torch.int32)
|
||||
if varlen:
|
||||
for i in range(b):
|
||||
cache_seqlens[i] = max(random.normalvariate(mean_sk, mean_sk / 2), s_q)
|
||||
total_seqlens = cache_seqlens.sum().item()
|
||||
max_seqlen = cache_seqlens.max().item()
|
||||
max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256
|
||||
|
||||
q = torch.randn(b, s_q, h_q, d)
|
||||
block_table = torch.arange(
|
||||
b * max_seqlen_pad // block_size, dtype=torch.int32
|
||||
).view(b, max_seqlen_pad // block_size)
|
||||
blocked_k = torch.randn(block_table.numel(), block_size, h_kv, d)
|
||||
blocked_v = blocked_k[..., :dv]
|
||||
|
||||
init_dtype = q.dtype
|
||||
if use_fp8:
|
||||
fp8_dtype = torch.float8_e4m3fn
|
||||
descale_q = torch.ones((1), dtype=torch.float32)
|
||||
descale_k = torch.ones((1), dtype=torch.float32)
|
||||
|
||||
q = q.to(fp8_dtype)
|
||||
blocked_k = blocked_k.to(fp8_dtype)
|
||||
blocked_v = blocked_v.to(fp8_dtype)
|
||||
else:
|
||||
descale_q = None
|
||||
descale_k = None
|
||||
|
||||
def cutlass_mla():
|
||||
MAX_HEADS = 128
|
||||
|
||||
q_reshaped = q.squeeze(1)
|
||||
q_nope = q_reshaped[:, :, :dv].clone()
|
||||
q_pe = q_reshaped[:, :, dv:].clone()
|
||||
|
||||
if h_q < MAX_HEADS:
|
||||
q_nope_padded = q_nope.new_empty((b, MAX_HEADS, dv))
|
||||
q_nope_padded[:, :h_q] = q_nope
|
||||
q_nope = q_nope_padded
|
||||
|
||||
q_pe_padded = q_pe.new_empty((b, MAX_HEADS, d - dv))
|
||||
q_pe_padded[:, :h_q] = q_pe
|
||||
q_pe = q_pe_padded
|
||||
|
||||
kv_cache_flat = blocked_k.squeeze(2)
|
||||
sm_count = num_compute_units(device.index)
|
||||
workspace_size = ops.sm100_cutlass_mla_get_workspace_size(
|
||||
max_seqlen * block_size, b, sm_count, num_kv_splits=1
|
||||
)
|
||||
workspace = torch.empty(workspace_size, device="cuda", dtype=torch.uint8)
|
||||
|
||||
out_ans = torch.empty(b, MAX_HEADS, dv, dtype=init_dtype)
|
||||
output_lse = torch.empty(
|
||||
(b, MAX_HEADS), dtype=torch.float32, device=q_nope.device
|
||||
)
|
||||
ops.sm100_cutlass_mla_decode(
|
||||
out_ans,
|
||||
output_lse,
|
||||
q_nope,
|
||||
q_pe,
|
||||
kv_cache_flat,
|
||||
cache_seqlens,
|
||||
block_table,
|
||||
workspace,
|
||||
scale,
|
||||
1,
|
||||
)
|
||||
return out_ans[:, :h_q].contiguous(), output_lse[:, :h_q].contiguous()
|
||||
|
||||
def scaled_dot_product_attention(query, key, value, is_causal=False):
|
||||
query = query.float()
|
||||
key = key.float()
|
||||
value = value.float()
|
||||
key = key.repeat_interleave(h_q // h_kv, dim=0)
|
||||
value = value.repeat_interleave(h_q // h_kv, dim=0)
|
||||
attn_weight = query @ key.transpose(-2, -1) / math.sqrt(query.size(-1))
|
||||
if is_causal:
|
||||
s_q = query.shape[-2]
|
||||
s_k = key.shape[-2]
|
||||
attn_bias = torch.zeros(s_q, s_k, dtype=query.dtype)
|
||||
temp_mask = torch.ones(s_q, s_k, dtype=torch.bool).tril(diagonal=s_k - s_q)
|
||||
attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
|
||||
attn_bias.to(query.dtype)
|
||||
attn_weight += attn_bias
|
||||
lse = attn_weight.logsumexp(dim=-1)
|
||||
attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32)
|
||||
return attn_weight @ value, lse
|
||||
|
||||
def ref_mla():
|
||||
q_ = (q.to(torch.float) * descale_q).to(init_dtype) if use_fp8 else q
|
||||
blocked_k_ = (
|
||||
(blocked_k.to(torch.float) * descale_k).to(init_dtype)
|
||||
if use_fp8
|
||||
else blocked_k
|
||||
)
|
||||
blocked_v_ = (
|
||||
(blocked_v.to(torch.float) * descale_k).to(init_dtype)
|
||||
if use_fp8
|
||||
else blocked_v
|
||||
)
|
||||
out = torch.empty(b, s_q, h_q, dv, dtype=torch.float32)
|
||||
lse = torch.empty(b, h_q, s_q, dtype=torch.float32)
|
||||
for i in range(b):
|
||||
begin = i * max_seqlen_pad
|
||||
end = begin + cache_seqlens[i]
|
||||
out_i, lse_i = scaled_dot_product_attention(
|
||||
q_[i].transpose(0, 1),
|
||||
blocked_k_.view(-1, h_kv, d)[begin:end].transpose(0, 1),
|
||||
blocked_v_.view(-1, h_kv, dv)[begin:end].transpose(0, 1),
|
||||
is_causal=causal,
|
||||
)
|
||||
out[i] = out_i.transpose(0, 1)
|
||||
lse[i] = lse_i
|
||||
return out, lse
|
||||
|
||||
out_cutlass, lse_cutlass = cutlass_mla()
|
||||
out_torch, lse_torch = ref_mla()
|
||||
# Extract the single token (s_q=1) slice to match cutlass output shape
|
||||
out_torch_slice = out_torch[:, 0, :, :] # [b, h_q, dv]
|
||||
lse_torch_slice = lse_torch[:, 0, :] # [b, h_q]
|
||||
cal_diff(out_cutlass, out_torch_slice, "out", use_fp8)
|
||||
# lse has larger numerical error, so use a larger threshold
|
||||
cal_diff(lse_cutlass, lse_torch_slice, "lse", use_fp8, diff_threshold=1e-3)
|
||||
|
||||
t = triton.testing.do_bench(cutlass_mla)
|
||||
FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2
|
||||
bytes = (total_seqlens * h_kv * d + b * s_q * h_q * d) * (
|
||||
torch.finfo(torch_dtype).bits // 8
|
||||
) + (b * s_q * h_q * dv) * (torch.finfo(init_dtype).bits // 8)
|
||||
print(
|
||||
f"{t:.3f} ms, {FLOPS / 10**9 / t:.0f} TFLOPS,", f"{bytes / 10**6 / t:.0f} GB/s"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(100),
|
||||
reason=CUTLASS_MLA_UNSUPPORTED_REASON,
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_cutlass_mla_decode_cross_layer_view():
|
||||
"""The kernel must read the cache's page-dim stride instead of assuming
|
||||
pages are packed back-to-back. A per-layer view into a cross-layer
|
||||
(block-major) cache has stride(0) inflated by num_layers; outputs must
|
||||
match a contiguous cache holding the same data exactly."""
|
||||
device = torch.device("cuda:0")
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
torch.set_default_device(device)
|
||||
torch.manual_seed(42)
|
||||
|
||||
b, mean_sk, d, dv, block_size = 4, 512, 576, 512, 64
|
||||
num_layers, layer_idx = 3, 1
|
||||
scale = math.sqrt(d) ** (-1)
|
||||
|
||||
num_pages = b * (mean_sk // block_size)
|
||||
cache_seqlens = torch.full((b,), mean_sk, dtype=torch.int32)
|
||||
block_table = torch.arange(num_pages, dtype=torch.int32).view(
|
||||
b, mean_sk // block_size
|
||||
)
|
||||
|
||||
kv_contig = torch.randn(num_pages, block_size, d)
|
||||
# Neighbor layers hold random data so packed-pages addressing reads
|
||||
# garbage rather than zeros.
|
||||
kv_cross_layer = torch.randn(num_pages, num_layers, block_size, d)
|
||||
kv_view = kv_cross_layer[:, layer_idx]
|
||||
kv_view.copy_(kv_contig)
|
||||
assert kv_view.stride(0) == num_layers * block_size * d
|
||||
|
||||
q_nope = torch.randn(b, 128, dv)
|
||||
q_pe = torch.randn(b, 128, d - dv)
|
||||
sm_count = num_compute_units(device.index)
|
||||
workspace_size = ops.sm100_cutlass_mla_get_workspace_size(
|
||||
mean_sk, b, sm_count, num_kv_splits=1
|
||||
)
|
||||
workspace = torch.empty(workspace_size, dtype=torch.uint8)
|
||||
|
||||
def run(cache):
|
||||
out = torch.empty(b, 128, dv)
|
||||
lse = torch.empty(b, 128, dtype=torch.float32)
|
||||
ops.sm100_cutlass_mla_decode(
|
||||
out,
|
||||
lse,
|
||||
q_nope,
|
||||
q_pe,
|
||||
cache,
|
||||
cache_seqlens,
|
||||
block_table,
|
||||
workspace,
|
||||
scale,
|
||||
1,
|
||||
)
|
||||
return out, lse
|
||||
|
||||
out_contig, lse_contig = run(kv_contig)
|
||||
out_view, lse_view = run(kv_view)
|
||||
|
||||
# Same data and same compute order; only addressing differs.
|
||||
assert torch.equal(out_contig, out_view)
|
||||
assert torch.equal(lse_contig, lse_view)
|
||||
@@ -0,0 +1,310 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.deep_gemm import (
|
||||
_ceil_to_ue8m0,
|
||||
calc_diff,
|
||||
fp8_fp4_mqa_logits,
|
||||
fp8_fp4_paged_mqa_logits,
|
||||
get_num_sms,
|
||||
get_paged_mqa_logits_metadata,
|
||||
)
|
||||
from vllm.utils.import_utils import has_deep_gemm
|
||||
from vllm.utils.math_utils import cdiv
|
||||
|
||||
|
||||
def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor:
|
||||
# x: (num_blocks, block_size, 1, head_dim)
|
||||
num_blocks, block_size, num_heads, head_dim = x.shape
|
||||
assert num_heads == 1
|
||||
x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4)
|
||||
sf = x_amax / 448.0
|
||||
x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn)
|
||||
x_fp8 = torch.empty(
|
||||
(num_blocks, block_size * (head_dim + 4)),
|
||||
device=x.device,
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
x_fp8[:, : block_size * head_dim] = x_scaled.view(
|
||||
num_blocks, block_size * head_dim
|
||||
).view(dtype=torch.uint8)
|
||||
x_fp8[:, block_size * head_dim :] = sf.view(num_blocks, block_size).view(
|
||||
dtype=torch.uint8
|
||||
)
|
||||
return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4)
|
||||
|
||||
|
||||
def per_custom_dims_cast_to_fp8(
|
||||
x: torch.Tensor, dims: tuple, use_ue8m0: bool
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
excluded_dims = tuple([i for i in range(x.dim()) if i not in set(dims)])
|
||||
x_amax = x.abs().float().amax(dim=excluded_dims, keepdim=True).clamp(1e-4)
|
||||
sf = x_amax / 448.0
|
||||
sf = _ceil_to_ue8m0(sf) if use_ue8m0 else sf
|
||||
x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn)
|
||||
return x_scaled, sf.squeeze()
|
||||
|
||||
|
||||
def _generate_cp_test_data(seq_len: int, seq_len_kv: int):
|
||||
assert seq_len_kv % seq_len == 0 and seq_len % 2 == 0
|
||||
chunk_size = seq_len // 2
|
||||
cp_size = seq_len_kv // seq_len
|
||||
cp_id = cp_size // 3
|
||||
ks = torch.zeros(seq_len, dtype=torch.int, device="cuda")
|
||||
ke = torch.zeros(seq_len, dtype=torch.int, device="cuda")
|
||||
for i in range(chunk_size):
|
||||
ke[i] = cp_id * chunk_size + i
|
||||
ke[i + chunk_size] = (cp_size * 2 - 1 - cp_id) * chunk_size + i
|
||||
return ks, ke
|
||||
|
||||
|
||||
def _ref_fp8_mqa_logits(
|
||||
q: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
cu_seqlen_ks: torch.Tensor,
|
||||
cu_seqlen_ke: torch.Tensor,
|
||||
):
|
||||
seq_len_kv = kv.shape[0]
|
||||
|
||||
k = kv
|
||||
q = q.float()
|
||||
k = k.float()
|
||||
|
||||
mask_lo = (
|
||||
torch.arange(0, seq_len_kv, device="cuda")[None, :] >= cu_seqlen_ks[:, None]
|
||||
)
|
||||
mask_hi = (
|
||||
torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None]
|
||||
)
|
||||
mask = mask_lo & mask_hi
|
||||
score = torch.einsum("mhd,nd->hmn", q, k)
|
||||
logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)
|
||||
logits = logits.masked_fill(~mask, float("-inf"))
|
||||
|
||||
return logits
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
|
||||
@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available")
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(90), reason="SM90 and SM100 only"
|
||||
)
|
||||
@pytest.mark.parametrize("clean_logits", [True, False])
|
||||
def test_deepgemm_fp8_mqa_logits(clean_logits: bool):
|
||||
torch.manual_seed(0)
|
||||
random.seed(0)
|
||||
num_heads, head_dim = 32, 128
|
||||
for seq_len in (512,):
|
||||
for seq_len_kv in (1024,):
|
||||
for disable_cp in (False, True):
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
kv = torch.randn(
|
||||
seq_len_kv, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
weights = torch.randn(
|
||||
seq_len, num_heads, device="cuda", dtype=torch.float32
|
||||
)
|
||||
|
||||
if disable_cp:
|
||||
ks = torch.zeros(seq_len, dtype=torch.int, device="cuda")
|
||||
ke = torch.arange(seq_len, dtype=torch.int, device="cuda") + (
|
||||
seq_len_kv - seq_len
|
||||
)
|
||||
else:
|
||||
ks, ke = _generate_cp_test_data(seq_len, seq_len_kv)
|
||||
|
||||
q_fp8 = q.to(torch.float8_e4m3fn)
|
||||
kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False)
|
||||
logits = fp8_fp4_mqa_logits(
|
||||
(q_fp8, None), kv_fp8, weights, ks, ke, clean_logits=clean_logits
|
||||
)
|
||||
|
||||
ref_logits = _ref_fp8_mqa_logits(
|
||||
q=q,
|
||||
kv=kv,
|
||||
weights=weights,
|
||||
cu_seqlen_ks=ks,
|
||||
cu_seqlen_ke=ke,
|
||||
)
|
||||
ref_neginf_mask = ref_logits == float("-inf")
|
||||
|
||||
if clean_logits:
|
||||
neginf_mask = logits == float("-inf")
|
||||
assert torch.equal(neginf_mask, ref_neginf_mask)
|
||||
|
||||
ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0)
|
||||
logits = logits.masked_fill(ref_neginf_mask, 0)
|
||||
diff = calc_diff(logits, ref_logits)
|
||||
assert diff < 1e-3, f"{diff=}"
|
||||
|
||||
|
||||
def _ref_fp8_fp4_paged_mqa_logits(
|
||||
q: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
max_model_len: int,
|
||||
):
|
||||
batch_size, next_n, _, _ = q.size()
|
||||
_, block_size, _, _ = kv_cache.size()
|
||||
logits = torch.full(
|
||||
[batch_size * next_n, max_model_len],
|
||||
float("-inf"),
|
||||
device=q.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
context_lens_list = context_lens.tolist()
|
||||
for i in range(batch_size):
|
||||
context_len = context_lens_list[i]
|
||||
q_offsets = torch.arange(context_len - next_n, context_len, device="cuda")
|
||||
weight_slice = (
|
||||
weights[i * next_n : (i + 1) * next_n, :].transpose(0, 1).contiguous()
|
||||
)
|
||||
for block_rk in range(cdiv(context_len, block_size)):
|
||||
block_idx = block_tables[i][block_rk]
|
||||
qx, kx = q[i], kv_cache[block_idx]
|
||||
k_offsets = torch.arange(
|
||||
block_rk * block_size,
|
||||
(block_rk + 1) * block_size,
|
||||
device="cuda",
|
||||
)
|
||||
mask = (k_offsets[None, :] < context_len) & (
|
||||
k_offsets[None, :] <= q_offsets[:, None]
|
||||
)
|
||||
s = torch.where(
|
||||
mask[None, :, :],
|
||||
(qx.transpose(0, 1) @ kx.transpose(0, 1).transpose(1, 2)).to(
|
||||
logits.dtype
|
||||
),
|
||||
float("-inf"),
|
||||
)
|
||||
s = torch.relu(s) * weight_slice[..., None]
|
||||
s = s.sum(dim=0)
|
||||
logits[
|
||||
i * next_n : (i + 1) * next_n,
|
||||
block_rk * block_size : (block_rk + 1) * block_size,
|
||||
] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, float("-inf"))
|
||||
return logits
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
|
||||
@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available")
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(90), reason="SM90 and SM100 only"
|
||||
)
|
||||
def test_deepgemm_fp8_fp4_paged_mqa_logits():
|
||||
# NOTE: clean_logits=True is incompatible with the 2D context_lens
|
||||
# required by csrc/apis/attention.hpp; only the False path is exercised.
|
||||
clean_logits = False
|
||||
torch.manual_seed(0)
|
||||
random.seed(0)
|
||||
|
||||
max_model_len = 4096
|
||||
for batch_size, next_n in [(4, 1), (2, 2)]:
|
||||
for heads, index_dim in [(32, 128)]:
|
||||
for avg_kv in (2048,):
|
||||
num_blocks, blocksize = max_model_len * 2, 64
|
||||
|
||||
q = torch.randn(
|
||||
(batch_size, next_n, heads, index_dim),
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
kv_cache = torch.randn(
|
||||
(num_blocks, blocksize, 1, index_dim),
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
weights = torch.randn(
|
||||
(batch_size * next_n, heads),
|
||||
device="cuda",
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
context_lens = (
|
||||
torch.randint(int(0.8 * avg_kv), int(1.2 * avg_kv), (batch_size,))
|
||||
.cuda()
|
||||
.to(torch.int32)
|
||||
)
|
||||
max_block_len = (
|
||||
(context_lens.max().item() + blocksize - 1) // blocksize * blocksize
|
||||
)
|
||||
block_tables = torch.zeros(
|
||||
(batch_size, max_block_len),
|
||||
device="cuda",
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
counter = 0
|
||||
block_idx_pool = list(range(num_blocks))
|
||||
random.shuffle(block_idx_pool)
|
||||
for i in range(batch_size):
|
||||
ctx_len = int(context_lens[i].item())
|
||||
for j in range((ctx_len + blocksize - 1) // blocksize):
|
||||
block_tables[i][j] = block_idx_pool[counter]
|
||||
counter += 1
|
||||
|
||||
q_fp8 = q.to(torch.float8_e4m3fn)
|
||||
kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache)
|
||||
|
||||
# deep_gemm paged MQA logits requires 2D context_lens of
|
||||
# shape (B, next_n) (csrc/apis/attention.hpp:332-335);
|
||||
# see indexer.py:607-608. For each batch/next_n token, the
|
||||
# effective context length is context_lens[b] - next_n + j + 1.
|
||||
next_n_arange = torch.arange(next_n, device="cuda", dtype=torch.int32)
|
||||
context_lens_2d = (
|
||||
context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange
|
||||
).contiguous()
|
||||
schedule_metadata = get_paged_mqa_logits_metadata(
|
||||
context_lens_2d, blocksize, get_num_sms()
|
||||
)
|
||||
logits = fp8_fp4_paged_mqa_logits(
|
||||
(q_fp8, None),
|
||||
kv_cache_fp8,
|
||||
weights,
|
||||
context_lens_2d,
|
||||
block_tables,
|
||||
schedule_metadata,
|
||||
max_model_len,
|
||||
clean_logits=clean_logits,
|
||||
)
|
||||
|
||||
ref_logits = _ref_fp8_fp4_paged_mqa_logits(
|
||||
q,
|
||||
kv_cache,
|
||||
weights,
|
||||
context_lens,
|
||||
block_tables,
|
||||
max_model_len,
|
||||
)
|
||||
|
||||
positions = (
|
||||
torch.arange(max_model_len, device="cuda")
|
||||
.unsqueeze(0)
|
||||
.expand(batch_size * next_n, -1)
|
||||
)
|
||||
row_indices = torch.arange(batch_size * next_n, device="cuda") // next_n
|
||||
next_n_offset = (
|
||||
torch.arange(batch_size * next_n, device="cuda") % next_n
|
||||
)
|
||||
mask = positions <= (
|
||||
context_lens[row_indices] - next_n + next_n_offset
|
||||
).unsqueeze(1)
|
||||
|
||||
logits = logits.masked_fill(~mask, 0)
|
||||
ref_logits = ref_logits.masked_fill(~mask, 0)
|
||||
diff = calc_diff(logits, ref_logits)
|
||||
assert diff < 1e-3, f"{diff=}"
|
||||
@@ -0,0 +1,217 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
try:
|
||||
from vllm.vllm_flash_attn import (
|
||||
fa_version_unsupported_reason,
|
||||
flash_attn_varlen_func,
|
||||
is_fa_version_supported,
|
||||
)
|
||||
except ImportError:
|
||||
if current_platform.is_rocm():
|
||||
pytest.skip(
|
||||
"vllm_flash_attn is not supported for vLLM on ROCm.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
NUM_HEADS = [(4, 4), (8, 2)]
|
||||
HEAD_SIZES = [40, 72, 80, 128, 256]
|
||||
BLOCK_SIZES = [16]
|
||||
DTYPES = [torch.bfloat16]
|
||||
QDTYPES = [None, torch.float8_e4m3fn]
|
||||
# one value large enough to test overflow in index calculation.
|
||||
# one value small enough to test the schema op check
|
||||
NUM_BLOCKS = [32768, 2048]
|
||||
SOFT_CAPS = [None]
|
||||
SLIDING_WINDOWS = [None, 256]
|
||||
|
||||
|
||||
def ref_paged_attn(
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
query_lens: list[int],
|
||||
kv_lens: list[int],
|
||||
block_tables: torch.Tensor,
|
||||
scale: float,
|
||||
sliding_window: int | None = None,
|
||||
soft_cap: float | None = None,
|
||||
) -> torch.Tensor:
|
||||
num_seqs = len(query_lens)
|
||||
block_tables = block_tables.cpu().numpy()
|
||||
_, block_size, num_kv_heads, head_size = key_cache.shape
|
||||
|
||||
outputs: list[torch.Tensor] = []
|
||||
start_idx = 0
|
||||
for i in range(num_seqs):
|
||||
query_len = query_lens[i]
|
||||
kv_len = kv_lens[i]
|
||||
q = query[start_idx : start_idx + query_len]
|
||||
q *= scale
|
||||
|
||||
num_kv_blocks = (kv_len + block_size - 1) // block_size
|
||||
block_indices = block_tables[i, :num_kv_blocks]
|
||||
|
||||
k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
|
||||
k = k[:kv_len]
|
||||
v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
|
||||
v = v[:kv_len]
|
||||
|
||||
if q.shape[1] != k.shape[1]:
|
||||
k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
|
||||
v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
|
||||
attn = torch.einsum("qhd,khd->hqk", q, k).float()
|
||||
empty_mask = torch.ones(query_len, kv_len)
|
||||
mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
|
||||
if sliding_window is not None:
|
||||
sliding_window_mask = (
|
||||
torch.triu(
|
||||
empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1
|
||||
)
|
||||
.bool()
|
||||
.logical_not()
|
||||
)
|
||||
mask |= sliding_window_mask
|
||||
if soft_cap is not None:
|
||||
attn = soft_cap * torch.tanh(attn / soft_cap)
|
||||
attn.masked_fill_(mask, float("-inf"))
|
||||
attn = torch.softmax(attn, dim=-1).to(v.dtype)
|
||||
out = torch.einsum("hqk,khd->qhd", attn, v)
|
||||
|
||||
outputs.append(out)
|
||||
start_idx += query_len
|
||||
|
||||
return torch.cat(outputs, dim=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_out", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
|
||||
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
|
||||
@pytest.mark.parametrize("fa_version", [2, 3])
|
||||
@pytest.mark.parametrize("q_dtype", QDTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_varlen_with_paged_kv(
|
||||
use_out: bool,
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
sliding_window: int | None,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
num_blocks: int,
|
||||
fa_version: int,
|
||||
q_dtype: torch.dtype | None,
|
||||
) -> None:
|
||||
torch.set_default_device("cuda")
|
||||
if not is_fa_version_supported(fa_version):
|
||||
pytest.skip(
|
||||
f"Flash attention version {fa_version} not supported due "
|
||||
f'to: "{fa_version_unsupported_reason(fa_version)}"'
|
||||
)
|
||||
if q_dtype is not None and (dtype != torch.bfloat16 or fa_version == 2):
|
||||
pytest.skip(
|
||||
"Flash attention with quantized inputs is only "
|
||||
"supported on version 3 with bfloat16 base type"
|
||||
)
|
||||
set_random_seed(0)
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
|
||||
scale = head_size**-0.5
|
||||
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
kv_lens = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
out = torch.empty_like(query) if use_out else None
|
||||
|
||||
maybe_quantized_query = query
|
||||
maybe_quantized_key_cache = key_cache
|
||||
maybe_quantized_value_cache = value_cache
|
||||
q_descale = None
|
||||
k_descale = None
|
||||
v_descale = None
|
||||
if q_dtype is not None:
|
||||
# QKV are drawn from N(0, 1): no need for a fp8 scaling factor
|
||||
maybe_quantized_query = query.to(q_dtype)
|
||||
maybe_quantized_key_cache = key_cache.to(q_dtype)
|
||||
maybe_quantized_value_cache = value_cache.to(q_dtype)
|
||||
|
||||
scale_shape = (num_seqs, num_kv_heads)
|
||||
q_descale = torch.ones(scale_shape, dtype=torch.float32)
|
||||
k_descale = torch.ones(scale_shape, dtype=torch.float32)
|
||||
v_descale = torch.ones(scale_shape, dtype=torch.float32)
|
||||
|
||||
output = flash_attn_varlen_func(
|
||||
q=maybe_quantized_query,
|
||||
k=maybe_quantized_key_cache,
|
||||
v=maybe_quantized_value_cache,
|
||||
out=out,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
seqused_k=kv_lens,
|
||||
max_seqlen_q=max_query_len,
|
||||
max_seqlen_k=max_kv_len,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
block_table=block_tables,
|
||||
softcap=soft_cap if soft_cap is not None else 0,
|
||||
fa_version=fa_version,
|
||||
q_descale=q_descale,
|
||||
k_descale=k_descale,
|
||||
v_descale=v_descale,
|
||||
)
|
||||
output = output if not use_out else out
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
sliding_window=sliding_window,
|
||||
soft_cap=soft_cap,
|
||||
)
|
||||
atol, rtol = 1.5e-2, 1e-2
|
||||
if q_dtype is not None:
|
||||
atol, rtol = 1.5e-1, 1.5e-1
|
||||
(
|
||||
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output - ref_output))}",
|
||||
)
|
||||
@@ -0,0 +1,701 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
try:
|
||||
import flashinfer
|
||||
except ImportError:
|
||||
if current_platform.is_rocm():
|
||||
pytest.skip(
|
||||
"flashinfer is not supported for vLLM on ROCm.", allow_module_level=True
|
||||
)
|
||||
|
||||
import torch
|
||||
|
||||
NUM_HEADS = [(32, 8), (6, 1)]
|
||||
HEAD_SIZES = [128, 256]
|
||||
BLOCK_SIZES = [16, 32]
|
||||
DTYPES = [torch.bfloat16]
|
||||
NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation.
|
||||
SOFT_CAPS = [None, 30.0]
|
||||
SLIDING_WINDOWS = [None, 64]
|
||||
|
||||
|
||||
def ref_paged_attn(
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
query_lens: list[int],
|
||||
kv_lens: list[int],
|
||||
block_tables: torch.Tensor,
|
||||
scale: float,
|
||||
sliding_window: int | None = None,
|
||||
soft_cap: float | None = None,
|
||||
) -> torch.Tensor:
|
||||
num_seqs = len(query_lens)
|
||||
block_tables = block_tables.cpu().numpy()
|
||||
_, block_size, num_kv_heads, head_size = key_cache.shape
|
||||
|
||||
outputs: list[torch.Tensor] = []
|
||||
start_idx = 0
|
||||
for i in range(num_seqs):
|
||||
query_len = query_lens[i]
|
||||
kv_len = kv_lens[i]
|
||||
q = query[start_idx : start_idx + query_len]
|
||||
q *= scale
|
||||
|
||||
num_kv_blocks = (kv_len + block_size - 1) // block_size
|
||||
block_indices = block_tables[i, :num_kv_blocks]
|
||||
|
||||
k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
|
||||
k = k[:kv_len]
|
||||
v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
|
||||
v = v[:kv_len]
|
||||
|
||||
if q.shape[1] != k.shape[1]:
|
||||
k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
|
||||
v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
|
||||
attn = torch.einsum("qhd,khd->hqk", q, k).float()
|
||||
empty_mask = torch.ones(query_len, kv_len)
|
||||
mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
|
||||
if sliding_window is not None:
|
||||
sliding_window_mask = (
|
||||
torch.triu(
|
||||
empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1
|
||||
)
|
||||
.bool()
|
||||
.logical_not()
|
||||
)
|
||||
mask |= sliding_window_mask
|
||||
if soft_cap is not None:
|
||||
attn = soft_cap * torch.tanh(attn / soft_cap)
|
||||
attn.masked_fill_(mask, float("-inf"))
|
||||
attn = torch.softmax(attn, dim=-1).to(v.dtype)
|
||||
out = torch.einsum("hqk,khd->qhd", attn, v)
|
||||
|
||||
outputs.append(out)
|
||||
start_idx += query_len
|
||||
|
||||
return torch.cat(outputs, dim=0)
|
||||
|
||||
|
||||
def _make_paged_kv_metadata(
|
||||
kv_lens: list[int],
|
||||
block_size: int,
|
||||
num_blocks: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Build paged-KV metadata tensors for fast_plan_decode tests.
|
||||
|
||||
Returns:
|
||||
kv_indptr – CPU int32, shape [num_seqs + 1]
|
||||
kv_indices – CUDA int32, shape [total_blocks]
|
||||
kv_last_page_lens – CPU int32, shape [num_seqs]
|
||||
block_tables – CUDA int32, shape [num_seqs, max_blocks_per_seq]
|
||||
"""
|
||||
num_seqs = len(kv_lens)
|
||||
max_blocks = (max(kv_lens) + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_blocks), dtype=torch.int32, device="cuda"
|
||||
)
|
||||
|
||||
indptr_list = [0]
|
||||
indices_list: list[int] = []
|
||||
last_lens_list: list[int] = []
|
||||
for i, seq_len in enumerate(kv_lens):
|
||||
n = (seq_len + block_size - 1) // block_size
|
||||
indices_list.extend(block_tables[i, :n].cpu().tolist())
|
||||
indptr_list.append(indptr_list[-1] + n)
|
||||
last_lens_list.append(seq_len % block_size or block_size)
|
||||
|
||||
return (
|
||||
torch.tensor(indptr_list, dtype=torch.int32, device="cpu"),
|
||||
torch.tensor(indices_list, dtype=torch.int32, device="cuda"),
|
||||
torch.tensor(last_lens_list, dtype=torch.int32, device="cpu"),
|
||||
block_tables,
|
||||
)
|
||||
|
||||
|
||||
def _make_cg_decode_wrapper(
|
||||
num_seqs: int,
|
||||
kv_indices_buffer: torch.Tensor,
|
||||
workspace_buffer: torch.Tensor,
|
||||
use_tensor_cores: bool = True,
|
||||
) -> "flashinfer.BatchDecodeWithPagedKVCacheWrapper":
|
||||
"""Create a cudagraph-enabled BatchDecodeWithPagedKVCacheWrapper.
|
||||
|
||||
*kv_indices_buffer* is shared with the caller so that fast_plan_decode
|
||||
can avoid the device-to-device index copy on subsequent (cudagraph) calls.
|
||||
"""
|
||||
return flashinfer.BatchDecodeWithPagedKVCacheWrapper(
|
||||
workspace_buffer,
|
||||
"NHD",
|
||||
use_cuda_graph=True,
|
||||
paged_kv_indptr_buffer=torch.zeros(
|
||||
num_seqs + 1, dtype=torch.int32, device="cuda"
|
||||
),
|
||||
paged_kv_indices_buffer=kv_indices_buffer,
|
||||
paged_kv_last_page_len_buffer=torch.zeros(
|
||||
num_seqs, dtype=torch.int32, device="cuda"
|
||||
),
|
||||
use_tensor_cores=use_tensor_cores,
|
||||
)
|
||||
|
||||
|
||||
def test_fast_decode_plan_importable() -> None:
|
||||
"""fast_decode_plan must be importable from flashinfer.decode.
|
||||
|
||||
This is a forward-compatibility smoke test: if FlashInfer reorganises its
|
||||
public API the import will fail before any other test does.
|
||||
"""
|
||||
from flashinfer.decode import fast_decode_plan # noqa: F401
|
||||
|
||||
assert callable(fast_decode_plan)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode
|
||||
def test_fast_plan_decode_warmup_uses_full_plan(dtype: torch.dtype) -> None:
|
||||
"""On the first call fast_plan_decode must route through self.plan() and
|
||||
flip vllm_first_call to False on the wrapper object."""
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.v1.attention.backends.flashinfer import fast_plan_decode
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
set_random_seed(0)
|
||||
|
||||
kv_lens = [128, 64]
|
||||
block_size = 16
|
||||
num_seqs = len(kv_lens)
|
||||
num_query_heads, num_kv_heads = 8, 2
|
||||
head_size = 128
|
||||
|
||||
kv_indptr, kv_indices, kv_last_page_lens, _ = _make_paged_kv_metadata(
|
||||
kv_lens, block_size, NUM_BLOCKS
|
||||
)
|
||||
|
||||
workspace = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
|
||||
wrapper = _make_cg_decode_wrapper(num_seqs, kv_indices.clone(), workspace)
|
||||
|
||||
assert getattr(wrapper, "vllm_first_call", True) is True
|
||||
|
||||
with patch.object(wrapper, "plan", wraps=wrapper.plan) as mock_plan:
|
||||
fast_plan_decode(
|
||||
wrapper,
|
||||
indptr_cpu=kv_indptr,
|
||||
indices=kv_indices,
|
||||
last_page_len_cpu=kv_last_page_lens,
|
||||
num_qo_heads=num_query_heads,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_dim=head_size,
|
||||
page_size=block_size,
|
||||
q_data_type=dtype,
|
||||
kv_data_type=dtype,
|
||||
)
|
||||
mock_plan.assert_called_once()
|
||||
|
||||
assert wrapper.vllm_first_call is False, (
|
||||
"vllm_first_call should be False after the first fast_plan_decode call"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]])
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode
|
||||
def test_fast_plan_decode_matches_full_plan(
|
||||
kv_lens: list[int],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
"""fast_plan_decode's cudagraph path (delegating to FlashInfer's
|
||||
fast_decode_plan) must produce attention output numerically identical to
|
||||
a standard plan() call.
|
||||
|
||||
Both the warmup call (self.plan) and the subsequent fast call
|
||||
(fast_decode_plan) are verified against the same reference.
|
||||
"""
|
||||
from vllm.v1.attention.backends.flashinfer import fast_plan_decode
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
set_random_seed(0)
|
||||
num_seqs = len(kv_lens)
|
||||
num_query_heads, num_kv_heads = num_heads
|
||||
|
||||
query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dtype)
|
||||
key_value_cache = torch.randn(
|
||||
NUM_BLOCKS, 2, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
|
||||
kv_indptr, kv_indices, kv_last_page_lens, _ = _make_paged_kv_metadata(
|
||||
kv_lens, block_size, NUM_BLOCKS
|
||||
)
|
||||
|
||||
# Reference output via the standard plan()
|
||||
workspace_ref = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
|
||||
ref_wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
|
||||
workspace_ref, "NHD", use_tensor_cores=True
|
||||
)
|
||||
ref_wrapper.plan(
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
kv_last_page_lens,
|
||||
num_query_heads,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
block_size,
|
||||
"NONE",
|
||||
q_data_type=dtype,
|
||||
kv_data_type=dtype,
|
||||
)
|
||||
ref_output = ref_wrapper.run(query, key_value_cache)
|
||||
|
||||
# CUDAGraph wrapper exercised through fast_plan_decode
|
||||
kv_indices_buf = kv_indices.clone()
|
||||
workspace_cg = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
|
||||
cg_wrapper = _make_cg_decode_wrapper(num_seqs, kv_indices_buf, workspace_cg)
|
||||
|
||||
plan_kwargs: dict = dict(
|
||||
indptr_cpu=kv_indptr,
|
||||
indices=kv_indices_buf,
|
||||
last_page_len_cpu=kv_last_page_lens,
|
||||
num_qo_heads=num_query_heads,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_dim=head_size,
|
||||
page_size=block_size,
|
||||
q_data_type=dtype,
|
||||
kv_data_type=dtype,
|
||||
)
|
||||
|
||||
# First call – warmup path (routes through self.plan)
|
||||
fast_plan_decode(cg_wrapper, **plan_kwargs)
|
||||
warmup_output = cg_wrapper.run(query, key_value_cache)
|
||||
torch.testing.assert_close(warmup_output, ref_output, atol=1e-2, rtol=1e-2)
|
||||
|
||||
# Second call – fast path (routes through fast_decode_plan from FlashInfer)
|
||||
fast_plan_decode(cg_wrapper, **plan_kwargs)
|
||||
fast_output = cg_wrapper.run(query, key_value_cache)
|
||||
torch.testing.assert_close(fast_output, ref_output, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]])
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
|
||||
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
|
||||
@torch.inference_mode
|
||||
def test_flashinfer_decode_with_paged_kv(
|
||||
kv_lens: list[int],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
sliding_window: int | None,
|
||||
) -> None:
|
||||
torch.set_default_device("cuda")
|
||||
set_random_seed(0)
|
||||
num_seqs = len(kv_lens)
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_kv_len = max(kv_lens)
|
||||
scale = head_size**-0.5
|
||||
|
||||
query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dtype)
|
||||
|
||||
key_value_cache = torch.randn(
|
||||
NUM_BLOCKS, 2, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
key_cache = key_value_cache[:, 0, :, :, :].squeeze(1)
|
||||
value_cache = key_value_cache[:, 1, :, :, :].squeeze(1)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
kv_indptr = [0]
|
||||
kv_indices = []
|
||||
kv_last_page_lens = []
|
||||
for i in range(num_seqs):
|
||||
seq_len = kv_lens[i]
|
||||
assert seq_len > 0
|
||||
num_blocks = (seq_len + block_size - 1) // block_size
|
||||
kv_indices.extend(block_tables[i, :num_blocks])
|
||||
kv_indptr.append(kv_indptr[-1] + num_blocks)
|
||||
kv_last_page_len = seq_len % block_size
|
||||
if kv_last_page_len == 0:
|
||||
kv_last_page_len = block_size
|
||||
kv_last_page_lens.append(kv_last_page_len)
|
||||
|
||||
kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
|
||||
kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
|
||||
kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
|
||||
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
|
||||
wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
|
||||
workspace_buffer, "NHD", use_tensor_cores=True
|
||||
)
|
||||
wrapper.plan(
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
kv_last_page_lens,
|
||||
num_query_heads,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
block_size,
|
||||
"NONE",
|
||||
window_left=sliding_window - 1 if sliding_window is not None else -1,
|
||||
q_data_type=dtype,
|
||||
kv_data_type=dtype,
|
||||
logits_soft_cap=soft_cap,
|
||||
)
|
||||
|
||||
output = wrapper.run(query, key_value_cache)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=[1] * num_seqs,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
soft_cap=soft_cap,
|
||||
sliding_window=sliding_window,
|
||||
)
|
||||
(
|
||||
torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2),
|
||||
f"{torch.max(torch.abs(output - ref_output))}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq_lens", [[(1, 1328), (5, 18), (129, 463)]])
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
|
||||
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
|
||||
@torch.inference_mode
|
||||
def test_flashinfer_prefill_with_paged_kv(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
sliding_window: int | None,
|
||||
) -> None:
|
||||
torch.set_default_device("cuda")
|
||||
set_random_seed(0)
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_kv_len = max(kv_lens)
|
||||
scale = head_size**-0.5
|
||||
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
|
||||
key_value_cache = torch.randn(
|
||||
NUM_BLOCKS, 2, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
key_cache = key_value_cache[:, 0, :, :, :].squeeze(1)
|
||||
value_cache = key_value_cache[:, 1, :, :, :].squeeze(1)
|
||||
|
||||
# Normalize the scale of the key and value caches to mitigate
|
||||
# numerical instability.
|
||||
key_cache /= head_size**0.5
|
||||
value_cache /= head_size**0.5
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
qo_indptr = [0]
|
||||
kv_indptr = [0]
|
||||
kv_indices = []
|
||||
kv_last_page_lens = []
|
||||
for i in range(num_seqs):
|
||||
seq_len = kv_lens[i]
|
||||
assert seq_len > 0
|
||||
num_blocks = (seq_len + block_size - 1) // block_size
|
||||
kv_indices.extend(block_tables[i, :num_blocks])
|
||||
kv_indptr.append(kv_indptr[-1] + num_blocks)
|
||||
kv_last_page_len = seq_len % block_size
|
||||
if kv_last_page_len == 0:
|
||||
kv_last_page_len = block_size
|
||||
kv_last_page_lens.append(kv_last_page_len)
|
||||
qo_indptr.append(qo_indptr[-1] + query_lens[i])
|
||||
|
||||
qo_indptr = torch.tensor(qo_indptr, dtype=torch.int32)
|
||||
kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
|
||||
kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
|
||||
kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
|
||||
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
|
||||
wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(workspace_buffer, "NHD")
|
||||
wrapper.plan(
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
kv_last_page_lens,
|
||||
num_query_heads,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
block_size,
|
||||
window_left=sliding_window - 1 if sliding_window is not None else -1,
|
||||
q_data_type=dtype,
|
||||
kv_data_type=dtype,
|
||||
logits_soft_cap=soft_cap,
|
||||
)
|
||||
|
||||
output = wrapper.run(
|
||||
query,
|
||||
key_value_cache,
|
||||
)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
soft_cap=soft_cap,
|
||||
sliding_window=sliding_window,
|
||||
)
|
||||
(
|
||||
torch.testing.assert_close(output, ref_output, atol=5e-2, rtol=1e-2),
|
||||
f"{torch.max(torch.abs(output - ref_output))}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq_lens", [[(1, 132), (5, 18)]])
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
|
||||
def test_flashinfer_prefill_with_paged_fp8_kv(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
) -> None:
|
||||
pytest.skip("TODO: fix the accuracy issue")
|
||||
torch.set_default_device("cuda")
|
||||
set_random_seed(0)
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_kv_len = max(kv_lens)
|
||||
scale = head_size**-0.5
|
||||
|
||||
kv_cache_dtype = torch.float8_e4m3fn
|
||||
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
|
||||
NUM_BLOCKS_FP8 = 2048
|
||||
key_value_cache = torch.randn(
|
||||
NUM_BLOCKS_FP8, 2, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
key_cache, value_cache = torch.chunk(key_value_cache, 2, dim=1)
|
||||
key_cache /= head_size**0.5
|
||||
value_cache /= head_size**0.5
|
||||
|
||||
k_scale = key_cache.amax().item() / 448.0
|
||||
v_scale = value_cache.amax().item() / 448.0
|
||||
|
||||
kv_cache_fp8 = torch.cat([key_cache / k_scale, value_cache / v_scale], dim=1).to(
|
||||
kv_cache_dtype
|
||||
)
|
||||
|
||||
assert kv_cache_fp8.shape == key_value_cache.shape
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, NUM_BLOCKS_FP8, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
qo_indptr = [0]
|
||||
kv_indptr = [0]
|
||||
kv_indices = []
|
||||
kv_last_page_lens = []
|
||||
for i in range(num_seqs):
|
||||
seq_len = kv_lens[i]
|
||||
assert seq_len > 0
|
||||
num_blocks = (seq_len + block_size - 1) // block_size
|
||||
kv_indices.extend(block_tables[i, :num_blocks])
|
||||
kv_indptr.append(kv_indptr[-1] + num_blocks)
|
||||
kv_last_page_len = seq_len % block_size
|
||||
if kv_last_page_len == 0:
|
||||
kv_last_page_len = block_size
|
||||
kv_last_page_lens.append(kv_last_page_len)
|
||||
qo_indptr.append(qo_indptr[-1] + query_lens[i])
|
||||
|
||||
qo_indptr = torch.tensor(qo_indptr, dtype=torch.int32)
|
||||
kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
|
||||
kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
|
||||
kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
|
||||
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
|
||||
wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(workspace_buffer, "NHD")
|
||||
wrapper.plan(
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
kv_last_page_lens,
|
||||
num_query_heads,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
block_size,
|
||||
q_data_type=dtype,
|
||||
kv_data_type=kv_cache_dtype,
|
||||
logits_soft_cap=soft_cap,
|
||||
)
|
||||
|
||||
output = wrapper.run(query, kv_cache_fp8, k_scale=k_scale, v_scale=v_scale)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache.squeeze(1),
|
||||
value_cache=value_cache.squeeze(1),
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
soft_cap=soft_cap,
|
||||
)
|
||||
del query
|
||||
del block_tables
|
||||
# verify prefill fp8
|
||||
(
|
||||
torch.testing.assert_close(output, ref_output, atol=5e-2, rtol=1e-2),
|
||||
f"{torch.max(torch.abs(output - ref_output))}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]])
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
|
||||
@pytest.mark.skip(reason="TODO: fix the accuracy issue")
|
||||
@torch.inference_mode
|
||||
def test_flashinfer_decode_with_paged_fp8_kv(
|
||||
kv_lens: list[int],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
) -> None:
|
||||
# test doesn't work for num_heads = (16,16)
|
||||
torch.set_default_device("cuda")
|
||||
set_random_seed(0)
|
||||
num_seqs = len(kv_lens)
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_kv_len = max(kv_lens)
|
||||
scale = head_size**-0.5
|
||||
use_tensor_cores = True
|
||||
kv_cache_dtype = torch.float8_e4m3fn
|
||||
|
||||
query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dtype)
|
||||
NUM_BLOCKS_FP8 = 2048
|
||||
key_value_cache = torch.randn(
|
||||
NUM_BLOCKS_FP8, 2, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
key_cache, value_cache = torch.chunk(key_value_cache, 2, dim=1)
|
||||
key_cache /= head_size**0.5
|
||||
value_cache /= head_size**0.5
|
||||
|
||||
k_scale = key_cache.amax().item() / 448.0
|
||||
v_scale = value_cache.amax().item() / 448.0
|
||||
|
||||
key_cache_fp8 = (key_cache / k_scale).to(kv_cache_dtype)
|
||||
value_cache_fp8 = (value_cache / v_scale).to(kv_cache_dtype)
|
||||
assert key_cache_fp8.shape[1] == 1 and value_cache_fp8.shape[1] == 1
|
||||
kv_cache_fp8 = torch.cat([key_cache_fp8, value_cache_fp8], dim=1)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, NUM_BLOCKS_FP8, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
kv_indptr = [0]
|
||||
kv_indices = []
|
||||
kv_last_page_lens = []
|
||||
for i in range(num_seqs):
|
||||
seq_len = kv_lens[i]
|
||||
assert seq_len > 0
|
||||
num_blocks = (seq_len + block_size - 1) // block_size
|
||||
kv_indices.extend(block_tables[i, :num_blocks])
|
||||
kv_indptr.append(kv_indptr[-1] + num_blocks)
|
||||
kv_last_page_len = seq_len % block_size
|
||||
if kv_last_page_len == 0:
|
||||
kv_last_page_len = block_size
|
||||
kv_last_page_lens.append(kv_last_page_len)
|
||||
|
||||
kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
|
||||
kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
|
||||
kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
|
||||
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
|
||||
wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
|
||||
workspace_buffer, "NHD", use_tensor_cores=use_tensor_cores
|
||||
)
|
||||
wrapper.plan(
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
kv_last_page_lens,
|
||||
num_query_heads,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
block_size,
|
||||
"NONE",
|
||||
q_data_type=dtype,
|
||||
kv_data_type=kv_cache_dtype,
|
||||
logits_soft_cap=soft_cap,
|
||||
)
|
||||
output = wrapper.run(query, kv_cache_fp8, k_scale=k_scale, v_scale=v_scale)
|
||||
key_cache = key_value_cache[:, 0, :, :, :].squeeze(1)
|
||||
value_cache = key_value_cache[:, 1, :, :, :].squeeze(1)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=[1] * num_seqs,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
soft_cap=soft_cap,
|
||||
)
|
||||
# Temporary fix: Increasing the tolerance. Seems like a flashinfer issue
|
||||
(
|
||||
torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=1e-2),
|
||||
f"{torch.max(torch.abs(output - ref_output))}",
|
||||
)
|
||||
@@ -0,0 +1,161 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
FLASHINFER_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
|
||||
|
||||
if not current_platform.has_device_capability(100):
|
||||
pytest.skip(
|
||||
reason="FlashInfer MLA Requires compute capability of 10 or above.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
else:
|
||||
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
|
||||
|
||||
# Deepseek R1 MLA config.
|
||||
NUM_HEADS = 128
|
||||
KV_LORA_RANK = 512
|
||||
QK_NOPE_HEAD_DIM = 128
|
||||
QK_ROPE_HEAD_DIM = 64
|
||||
QK_HEAD_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM
|
||||
SCALE = (QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM) ** -0.5
|
||||
|
||||
|
||||
def _make_decode_inputs(bs: int, block_size: int, dtype: torch.dtype):
|
||||
"""Build valid trtllm MLA decode inputs on the current CUDA device."""
|
||||
max_seq_len_cap = 1024
|
||||
seq_lens = [torch.randint(2, max_seq_len_cap, (1,)).item() for _ in range(bs)]
|
||||
seq_lens[-1] = max_seq_len_cap
|
||||
max_seq_len = max(seq_lens)
|
||||
seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32)
|
||||
|
||||
# Generate block tables with random but unique block IDs
|
||||
# From https://github.com/flashinfer-ai/flashinfer/pull/1222
|
||||
blocks_per_seq = (seq_lens_tensor + block_size - 1) // block_size
|
||||
max_num_blocks_per_seq = max(blocks_per_seq.max().item(), 4)
|
||||
total_blocks_needed = int(sum(blocks_per_seq))
|
||||
all_block_ids = torch.randperm(total_blocks_needed)
|
||||
|
||||
block_tables = torch.zeros((bs, max_num_blocks_per_seq), dtype=torch.int32)
|
||||
block_id = 0
|
||||
for i in range(bs):
|
||||
num_blocks_needed = blocks_per_seq[i]
|
||||
block_tables[i, :num_blocks_needed] = all_block_ids[
|
||||
block_id : block_id + num_blocks_needed
|
||||
]
|
||||
block_id += num_blocks_needed
|
||||
|
||||
kv_cache = torch.randn(block_tables.numel(), block_size, QK_HEAD_DIM).to(dtype)
|
||||
q = torch.randn(bs, NUM_HEADS, QK_HEAD_DIM).to(dtype)
|
||||
return q, kv_cache, block_tables, seq_lens_tensor, max_seq_len
|
||||
|
||||
|
||||
def ref_mla(
|
||||
out: Tensor, # (bs, num_heads, v_head_dim)
|
||||
query: Tensor, # (bs, num_heads, head_dim)
|
||||
kv_cache: Tensor, # (num_blocks, block_size, head_dim)
|
||||
scale: float,
|
||||
block_tables: Tensor, # (bs, max_num_blocks)
|
||||
seq_lens: Tensor, # (bs,)
|
||||
):
|
||||
bs, num_heads, v_head_dim = out.shape
|
||||
head_dim = query.shape[2]
|
||||
|
||||
for i in range(bs):
|
||||
# gather and flatten KV-cache
|
||||
kv = kv_cache[block_tables[i]] # (max_num_blocks, block_size, head_dim)
|
||||
kv = kv.view(1, -1, head_dim)[:, : seq_lens[i]] # (1, seq_len, head_dim)
|
||||
v = kv[:, :, :v_head_dim]
|
||||
|
||||
q = query[i].view(num_heads, 1, head_dim)
|
||||
o = F.scaled_dot_product_attention(q, kv, v, scale=scale, enable_gqa=True)
|
||||
out[i] = o.view(num_heads, v_head_dim)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("bs", [1, 2, 4, 16])
|
||||
@pytest.mark.parametrize("block_size", [32, 64])
|
||||
def test_flashinfer_mla_decode(dtype: torch.dtype, bs: int, block_size: int):
|
||||
torch.set_default_device("cuda")
|
||||
torch.manual_seed(42)
|
||||
|
||||
q, kv_cache, block_tables, seq_lens_tensor, max_seq_len = _make_decode_inputs(
|
||||
bs, block_size, dtype
|
||||
)
|
||||
|
||||
out_ref = q.new_zeros(bs, NUM_HEADS, KV_LORA_RANK)
|
||||
ref_mla(out_ref, q, kv_cache, SCALE, block_tables, seq_lens_tensor)
|
||||
|
||||
workspace_buffer = torch.zeros(
|
||||
FLASHINFER_WORKSPACE_BUFFER_SIZE,
|
||||
dtype=torch.uint8,
|
||||
device=q.device,
|
||||
)
|
||||
# Flashinfer MLA expects the query to be of shape
|
||||
# (bs, q_len_per_request, num_heads, qk_head_dim),
|
||||
# where q_len_per_request is the MTP query length (=1 without MTP)
|
||||
q = q.unsqueeze(1)
|
||||
|
||||
out_ans = trtllm_batch_decode_with_kv_cache_mla(
|
||||
query=q,
|
||||
kv_cache=kv_cache.unsqueeze(1),
|
||||
workspace_buffer=workspace_buffer,
|
||||
qk_nope_head_dim=QK_NOPE_HEAD_DIM,
|
||||
kv_lora_rank=KV_LORA_RANK,
|
||||
qk_rope_head_dim=QK_ROPE_HEAD_DIM,
|
||||
block_tables=block_tables,
|
||||
seq_lens=seq_lens_tensor,
|
||||
max_seq_len=max_seq_len,
|
||||
bmm1_scale=SCALE,
|
||||
)
|
||||
out_ans = out_ans.squeeze(1)
|
||||
torch.testing.assert_close(out_ans, out_ref, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
def test_flashinfer_mla_decode_workspace_supports_autotune():
|
||||
"""vLLM's FlashInfer MLA decode workspace must be int8 for autotuning.
|
||||
|
||||
Model Runner V2's warmup autotunes ``trtllm_batch_decode_mla``, which makes
|
||||
the FlashInfer autotuner enumerate the CuteDSL tactic. That tactic asserts
|
||||
``workspace_buffer.dtype == torch.int8``; the trtllm-gen path (used for
|
||||
normal, non-autotuned inference) instead views the buffer as uint8, so a
|
||||
uint8 workspace only fails once the autotuner tries CuteDSL. That regressed
|
||||
every DeepSeek MLA test on Blackwell under V2 with
|
||||
``workspace_buffer must be torch.int8`` (vllm-project/vllm#46646).
|
||||
"""
|
||||
from flashinfer.autotuner import autotune
|
||||
|
||||
from vllm.v1.attention.backends.mla.flashinfer_mla import _get_workspace_buffer
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
torch.manual_seed(0)
|
||||
|
||||
workspace_buffer = _get_workspace_buffer(return_lse=False)
|
||||
assert workspace_buffer.dtype == torch.int8
|
||||
|
||||
q, kv_cache, block_tables, seq_lens_tensor, max_seq_len = _make_decode_inputs(
|
||||
bs=1, block_size=64, dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
# Under the autotuner the CuteDSL tactic is instantiated with our workspace;
|
||||
# a uint8 buffer raises AssertionError here, an int8 buffer succeeds.
|
||||
with torch.inference_mode(), autotune(True):
|
||||
trtllm_batch_decode_with_kv_cache_mla(
|
||||
query=q.unsqueeze(1),
|
||||
kv_cache=kv_cache.unsqueeze(1),
|
||||
workspace_buffer=workspace_buffer,
|
||||
qk_nope_head_dim=QK_NOPE_HEAD_DIM,
|
||||
kv_lora_rank=KV_LORA_RANK,
|
||||
qk_rope_head_dim=QK_ROPE_HEAD_DIM,
|
||||
block_tables=block_tables,
|
||||
seq_lens=seq_lens_tensor,
|
||||
max_seq_len=max_seq_len,
|
||||
bmm1_scale=SCALE,
|
||||
)
|
||||
@@ -0,0 +1,554 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.quantization.nvfp4_utils import (
|
||||
dequant_nvfp4_kv_cache,
|
||||
dequantize_nvfp4_to_dtype,
|
||||
get_nvfp4_global_scale,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import round_up
|
||||
from vllm.utils.torch_utils import (
|
||||
nvfp4_kv_cache_full_dim,
|
||||
nvfp4_split_data_scale,
|
||||
set_random_seed,
|
||||
)
|
||||
|
||||
if not current_platform.is_device_capability_family(100):
|
||||
pytest.skip(
|
||||
"This TRTLLM kernel requires NVIDIA Blackwell.", allow_module_level=True
|
||||
)
|
||||
else:
|
||||
import flashinfer
|
||||
|
||||
FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
FP4_DTYPE = torch.uint8
|
||||
|
||||
|
||||
def to_float8(x, dtype=torch.float8_e4m3fn):
|
||||
finfo = torch.finfo(dtype)
|
||||
min_val, max_val = x.aminmax()
|
||||
amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
|
||||
scale = finfo.max / amax * 0.1
|
||||
x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
|
||||
return x_scl_sat.to(dtype), scale.float().reciprocal()
|
||||
|
||||
|
||||
def build_paged_kv_metadata(
|
||||
seq_lens: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
block_size: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Build paged-KV indptr/indices/last_page_lens from seq_lens + block_tables."""
|
||||
kv_indptr = [0]
|
||||
kv_indices = []
|
||||
kv_last_page_lens = []
|
||||
for i in range(len(seq_lens)):
|
||||
sl = int(seq_lens[i])
|
||||
assert sl > 0
|
||||
nb = (sl + block_size - 1) // block_size
|
||||
kv_indices.extend(block_tables[i, :nb].tolist())
|
||||
kv_indptr.append(kv_indptr[-1] + nb)
|
||||
kv_last_page_lens.append(sl % block_size or block_size)
|
||||
return (
|
||||
torch.tensor(kv_indptr, dtype=torch.int32),
|
||||
torch.tensor(kv_indices, dtype=torch.int32),
|
||||
torch.tensor(kv_last_page_lens, dtype=torch.int32),
|
||||
)
|
||||
|
||||
|
||||
def make_nvfp4_kv_cache(
|
||||
kv_bf16_hnd: torch.Tensor, block_size: int, head_size: int
|
||||
) -> tuple:
|
||||
"""Quantize bf16 KV cache to nvfp4 via reshape_and_cache_flash.
|
||||
|
||||
Returns (k_data, v_data), (k_scales, v_scales), kv_scale, ref_kv_bf16.
|
||||
"""
|
||||
num_blocks, _, num_kv_heads, _, _ = kv_bf16_hnd.shape
|
||||
kv_scale_val = (kv_bf16_hnd.abs().amax() / 448.0).item()
|
||||
kv_scale_tensor = torch.tensor(
|
||||
kv_scale_val, dtype=torch.float32, device=kv_bf16_hnd.device
|
||||
)
|
||||
|
||||
# layout: (B, 2*H, N, full_dim)
|
||||
# where K heads occupy the first H heads and V heads occupy the second H heads.
|
||||
full_dim = nvfp4_kv_cache_full_dim(head_size)
|
||||
kv_cache_hnd = torch.zeros(
|
||||
(num_blocks, 2 * num_kv_heads, block_size, full_dim),
|
||||
dtype=torch.uint8,
|
||||
device=kv_bf16_hnd.device,
|
||||
)
|
||||
kv_cache_nhd = kv_cache_hnd.permute(0, 2, 1, 3)
|
||||
k_view_nhd, v_view_nhd = kv_cache_nhd.split(num_kv_heads, dim=-2)
|
||||
|
||||
# Flatten input KV → token tensors [B*N, H, head_size] for the kernel.
|
||||
num_tokens = num_blocks * block_size
|
||||
k_tokens = (
|
||||
kv_bf16_hnd[:, 0]
|
||||
.permute(0, 2, 1, 3)
|
||||
.reshape(num_tokens, num_kv_heads, head_size)
|
||||
)
|
||||
v_tokens = (
|
||||
kv_bf16_hnd[:, 1]
|
||||
.permute(0, 2, 1, 3)
|
||||
.reshape(num_tokens, num_kv_heads, head_size)
|
||||
)
|
||||
slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=kv_bf16_hnd.device)
|
||||
|
||||
torch.ops._C_cache_ops.reshape_and_cache_flash(
|
||||
k_tokens,
|
||||
v_tokens,
|
||||
k_view_nhd,
|
||||
v_view_nhd,
|
||||
slot_mapping,
|
||||
"nvfp4",
|
||||
kv_scale_tensor,
|
||||
kv_scale_tensor,
|
||||
)
|
||||
|
||||
# Split into data/scale views in HNC order for trtllm kernel.
|
||||
k_cache_hnc, v_cache_hnc = kv_cache_hnd.split(num_kv_heads, dim=1)
|
||||
k_data, k_scales = nvfp4_split_data_scale(k_cache_hnc)
|
||||
v_data, v_scales = nvfp4_split_data_scale(v_cache_hnc)
|
||||
|
||||
# Dequantize for the FA2 reference baseline.
|
||||
ref_k = dequant_nvfp4_kv_cache(
|
||||
k_data, k_scales, kv_scale_val, head_size, block_size
|
||||
).to(torch.bfloat16)
|
||||
ref_v = dequant_nvfp4_kv_cache(
|
||||
v_data, v_scales, kv_scale_val, head_size, block_size
|
||||
).to(torch.bfloat16)
|
||||
ref_kv_bf16 = torch.stack([ref_k, ref_v], dim=1) # [N, 2, H, T, D]
|
||||
|
||||
return (k_data, v_data), (k_scales, v_scales), kv_scale_val, ref_kv_bf16
|
||||
|
||||
|
||||
def make_quantized_kv_cache(
|
||||
kv_cache: torch.Tensor,
|
||||
kv_quant_dtype: torch.dtype,
|
||||
block_size: int,
|
||||
head_size: int,
|
||||
) -> tuple:
|
||||
"""Quantize kv_cache based on dtype. Returns (kv_cache, kv_cache_sf,
|
||||
kv_scale, ref_kv_cache, is_nvfp4_kv)."""
|
||||
is_nvfp4_kv = kv_quant_dtype == FP4_DTYPE
|
||||
if is_nvfp4_kv:
|
||||
data, scales, kv_scale, ref = make_nvfp4_kv_cache(
|
||||
kv_cache, block_size, head_size
|
||||
)
|
||||
return data, scales, kv_scale, ref, True
|
||||
elif kv_quant_dtype == FP8_DTYPE:
|
||||
kv_fp8, kv_scale = to_float8(kv_cache)
|
||||
ref = kv_fp8.to(kv_cache.dtype) * kv_scale
|
||||
return kv_fp8, None, kv_scale, ref, False
|
||||
else:
|
||||
return kv_cache, None, 1.0, kv_cache, False
|
||||
|
||||
|
||||
DTYPE = [torch.bfloat16]
|
||||
QUANT_DTYPES = [
|
||||
# (q_quant_dtype, kv_quant_dtype, o_quant_dtype)
|
||||
(None, None, None),
|
||||
(None, FP8_DTYPE, None),
|
||||
(FP8_DTYPE, FP8_DTYPE, None),
|
||||
(FP8_DTYPE, FP8_DTYPE, FP8_DTYPE),
|
||||
(FP8_DTYPE, FP8_DTYPE, FP4_DTYPE),
|
||||
(FP8_DTYPE, FP4_DTYPE, FP8_DTYPE), # nvfp4 KV cache
|
||||
]
|
||||
BATCH_SIZE = [4, 12]
|
||||
MAX_SEQ_LENS = [(1024, 4096)]
|
||||
NUM_HEADS = [(64, 8), (40, 8)]
|
||||
HEAD_SIZE = [128]
|
||||
KV_LAYOUT = ["HND"] # currently only HND is supported
|
||||
BLOCK_SIZE = [16]
|
||||
WINDOW_LEFT = [-1, 127]
|
||||
SOFT_CAP = [None, 50.0]
|
||||
HAS_SINKS = [True, False]
|
||||
|
||||
NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation.
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPE)
|
||||
@pytest.mark.parametrize("quant_dtypes", QUANT_DTYPES)
|
||||
@pytest.mark.parametrize("batch_size", BATCH_SIZE)
|
||||
@pytest.mark.parametrize("max_seq_lens", MAX_SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZE)
|
||||
@pytest.mark.parametrize("kv_layout", KV_LAYOUT)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZE)
|
||||
@pytest.mark.parametrize("window_left", WINDOW_LEFT)
|
||||
@pytest.mark.parametrize("soft_cap", SOFT_CAP)
|
||||
@pytest.mark.parametrize("has_sinks", HAS_SINKS)
|
||||
@torch.inference_mode
|
||||
def test_flashinfer_trtllm_decode_with_baseline(
|
||||
dtype: torch.dtype,
|
||||
quant_dtypes: tuple[torch.dtype | None, torch.dtype | None, torch.dtype | None],
|
||||
batch_size: int,
|
||||
max_seq_lens: tuple[int, int],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
kv_layout: str,
|
||||
block_size: int,
|
||||
window_left: int,
|
||||
soft_cap: float | None,
|
||||
has_sinks: bool,
|
||||
) -> None:
|
||||
torch.set_default_device("cuda")
|
||||
set_random_seed(42)
|
||||
|
||||
q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtypes
|
||||
q_quant_dtype = q_quant_dtype or dtype
|
||||
kv_quant_dtype = kv_quant_dtype or dtype
|
||||
o_quant_dtype = o_quant_dtype or dtype
|
||||
|
||||
_, max_kv_len = max_seq_lens
|
||||
|
||||
num_qo_heads, num_kv_heads = num_heads
|
||||
assert num_qo_heads % num_kv_heads == 0
|
||||
|
||||
sm_scale = float(1.0 / (head_size**0.5))
|
||||
|
||||
kv_cache_shape = None
|
||||
if kv_layout == "NHD":
|
||||
kv_cache_shape = (NUM_BLOCKS, 2, block_size, num_kv_heads, head_size)
|
||||
elif kv_layout == "HND":
|
||||
kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size)
|
||||
else:
|
||||
raise ValueError(f"Invalid kv_layout: {kv_layout}")
|
||||
|
||||
# max_q_len = 1
|
||||
q_lens = torch.ones((batch_size,), dtype=torch.int32)
|
||||
q_indptr = torch.cat(
|
||||
[
|
||||
torch.tensor([0], dtype=torch.int32),
|
||||
torch.cumsum(q_lens, dim=0, dtype=torch.int32),
|
||||
]
|
||||
)
|
||||
|
||||
query = torch.randn(torch.sum(q_lens).item(), num_qo_heads, head_size, dtype=dtype)
|
||||
if q_quant_dtype == FP8_DTYPE:
|
||||
query, q_scale = to_float8(query)
|
||||
ref_query = query.to(dtype) * q_scale
|
||||
else:
|
||||
q_scale = 1.0
|
||||
ref_query = query
|
||||
|
||||
kv_lens = torch.randint(1, max_kv_len, (batch_size,), dtype=torch.int32)
|
||||
kv_lens[-1] = max_kv_len
|
||||
|
||||
seq_lens = kv_lens + q_lens
|
||||
max_seq_len = torch.max(seq_lens).item()
|
||||
|
||||
kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
|
||||
kv_cache, kv_cache_sf, kv_scale, ref_kv_cache, is_nvfp4_kv = (
|
||||
make_quantized_kv_cache(kv_cache, kv_quant_dtype, block_size, head_size)
|
||||
)
|
||||
|
||||
k_scale = v_scale = kv_scale
|
||||
|
||||
max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, NUM_BLOCKS, (batch_size, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
kv_indptr, kv_indices, kv_last_page_lens = build_paged_kv_metadata(
|
||||
seq_lens, block_tables, block_size
|
||||
)
|
||||
workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)
|
||||
|
||||
# Baseline Decode
|
||||
if has_sinks:
|
||||
sinks = torch.rand(num_qo_heads, dtype=torch.float32) * 5
|
||||
wrapper = flashinfer.BatchAttentionWithAttentionSinkWrapper(
|
||||
float_workspace_buffer=workspace_buffer, kv_layout=kv_layout, backend="fa2"
|
||||
)
|
||||
else:
|
||||
sinks = None
|
||||
wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
|
||||
float_workspace_buffer=workspace_buffer, kv_layout=kv_layout, backend="fa2"
|
||||
)
|
||||
|
||||
wrapper.plan(
|
||||
qo_indptr=q_indptr,
|
||||
paged_kv_indptr=kv_indptr,
|
||||
paged_kv_indices=kv_indices,
|
||||
paged_kv_last_page_len=kv_last_page_lens,
|
||||
num_qo_heads=num_qo_heads,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_dim_qk=head_size,
|
||||
page_size=block_size,
|
||||
causal=True,
|
||||
sm_scale=sm_scale,
|
||||
window_left=window_left,
|
||||
logits_soft_cap=soft_cap,
|
||||
q_data_type=dtype,
|
||||
kv_data_type=dtype,
|
||||
)
|
||||
output = torch.empty(ref_query.shape, dtype=dtype)
|
||||
wrapper.run(ref_query, ref_kv_cache, sinks, sm_scale, out=output)
|
||||
|
||||
o_scale = 1.0
|
||||
o_sf_scale_float = None
|
||||
if o_quant_dtype == FP8_DTYPE:
|
||||
_, o_scale = to_float8(output)
|
||||
elif o_quant_dtype == FP4_DTYPE:
|
||||
o_sf_scale = get_nvfp4_global_scale(output)
|
||||
o_sf_scale_float = o_sf_scale.item()
|
||||
|
||||
# TRTLLM Decode
|
||||
if o_quant_dtype == FP4_DTYPE:
|
||||
output_trtllm = flashinfer.utils.FP4Tensor(
|
||||
torch.empty(query.shape[:-1] + (query.shape[-1] // 2,), dtype=torch.uint8),
|
||||
torch.empty(
|
||||
(
|
||||
round_up(query.shape[0], 128),
|
||||
round_up(query.shape[1] * query.shape[2] // 16, 4),
|
||||
),
|
||||
dtype=torch.float8_e4m3fn,
|
||||
),
|
||||
)
|
||||
else:
|
||||
output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
|
||||
|
||||
flashinfer.decode.trtllm_batch_decode_with_kv_cache(
|
||||
query=query,
|
||||
kv_cache=kv_cache,
|
||||
workspace_buffer=workspace_buffer,
|
||||
block_tables=block_tables,
|
||||
seq_lens=seq_lens,
|
||||
max_seq_len=max_seq_len,
|
||||
bmm1_scale=q_scale * k_scale * sm_scale,
|
||||
bmm2_scale=v_scale / o_scale,
|
||||
window_left=window_left,
|
||||
sinks=sinks,
|
||||
o_sf_scale=o_sf_scale_float,
|
||||
out=output_trtllm,
|
||||
kv_cache_sf=kv_cache_sf,
|
||||
)
|
||||
if o_quant_dtype == FP8_DTYPE:
|
||||
output_trtllm = output_trtllm.to(dtype) * o_scale
|
||||
elif o_quant_dtype == FP4_DTYPE:
|
||||
output_trtllm.data = output_trtllm.data.reshape(
|
||||
-1, query.shape[1] * query.shape[2] // 2
|
||||
)
|
||||
output_trtllm = dequantize_nvfp4_to_dtype(
|
||||
output_trtllm.data, output_trtllm.scale, o_sf_scale, dtype, query.device
|
||||
)
|
||||
output_trtllm = output_trtllm.reshape(-1, query.shape[1], query.shape[2])
|
||||
|
||||
if is_nvfp4_kv:
|
||||
rtol, atol = 1.0, 1.0 # nvfp4 has higher quantization error
|
||||
elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE:
|
||||
rtol, atol = 7e-2, 9e-2
|
||||
elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
|
||||
rtol, atol = 3e-2, 4e-2
|
||||
elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype:
|
||||
rtol, atol = 2e-2, 2e-2
|
||||
elif kv_quant_dtype == FP8_DTYPE:
|
||||
rtol, atol = 4e-2, 6e-2
|
||||
else:
|
||||
rtol, atol = 1e-2, 1e-2
|
||||
|
||||
(
|
||||
torch.testing.assert_close(output, output_trtllm, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output - output_trtllm))}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPE)
|
||||
@pytest.mark.parametrize("quant_dtypes", QUANT_DTYPES)
|
||||
@pytest.mark.parametrize("batch_size", BATCH_SIZE)
|
||||
@pytest.mark.parametrize("max_seq_lens", MAX_SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZE)
|
||||
@pytest.mark.parametrize("kv_layout", KV_LAYOUT)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZE)
|
||||
@pytest.mark.parametrize("window_left", WINDOW_LEFT)
|
||||
@pytest.mark.parametrize("soft_cap", [None])
|
||||
@pytest.mark.parametrize("has_sinks", HAS_SINKS)
|
||||
@torch.inference_mode
|
||||
def test_flashinfer_trtllm_prefill_with_baseline(
|
||||
dtype: torch.dtype,
|
||||
quant_dtypes: tuple[torch.dtype | None, torch.dtype | None, torch.dtype | None],
|
||||
batch_size: int,
|
||||
max_seq_lens: tuple[int, int],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
kv_layout: str,
|
||||
block_size: int,
|
||||
window_left: int,
|
||||
soft_cap: float | None,
|
||||
has_sinks: bool,
|
||||
) -> None:
|
||||
torch.set_default_device("cuda")
|
||||
set_random_seed(42)
|
||||
|
||||
q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtypes
|
||||
q_quant_dtype = q_quant_dtype or dtype
|
||||
kv_quant_dtype = kv_quant_dtype or dtype
|
||||
o_quant_dtype = o_quant_dtype or dtype
|
||||
|
||||
# FP8 Q + nvfp4 KV is the required combination for the nvfp4 KV path.
|
||||
# All other mixed Q/KV dtype combinations are unsupported.
|
||||
is_nvfp4_kv = kv_quant_dtype == FP4_DTYPE
|
||||
if q_quant_dtype != kv_quant_dtype and not (
|
||||
q_quant_dtype == FP8_DTYPE and is_nvfp4_kv
|
||||
):
|
||||
pytest.skip("Skipped mixed QKV dtypes for prefill")
|
||||
|
||||
max_q_len, max_kv_len = max_seq_lens
|
||||
|
||||
num_qo_heads, num_kv_heads = num_heads
|
||||
assert num_qo_heads % num_kv_heads == 0
|
||||
|
||||
sm_scale = float(1.0 / (head_size**0.5))
|
||||
|
||||
kv_cache_shape = None
|
||||
if kv_layout == "NHD":
|
||||
kv_cache_shape = (NUM_BLOCKS, 2, block_size, num_kv_heads, head_size)
|
||||
elif kv_layout == "HND":
|
||||
kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size)
|
||||
else:
|
||||
raise ValueError(f"Invalid kv_layout: {kv_layout}")
|
||||
|
||||
q_lens = torch.randint(1, max_q_len, (batch_size,), dtype=torch.int32)
|
||||
q_lens[-1] = max_q_len
|
||||
q_indptr = torch.cat(
|
||||
[
|
||||
torch.tensor([0], dtype=torch.int32),
|
||||
torch.cumsum(q_lens, dim=0, dtype=torch.int32),
|
||||
]
|
||||
)
|
||||
|
||||
query = torch.randn(torch.sum(q_lens).item(), num_qo_heads, head_size, dtype=dtype)
|
||||
if q_quant_dtype == FP8_DTYPE:
|
||||
query, q_scale = to_float8(query)
|
||||
ref_query = query.to(dtype) * q_scale
|
||||
else:
|
||||
q_scale = 1.0
|
||||
ref_query = query
|
||||
|
||||
kv_lens = torch.randint(1, max_kv_len, (batch_size,), dtype=torch.int32)
|
||||
kv_lens[-1] = max_kv_len
|
||||
|
||||
seq_lens = kv_lens + q_lens
|
||||
max_seq_len = torch.max(seq_lens).item()
|
||||
|
||||
kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
|
||||
kv_cache, kv_cache_sf, kv_scale, ref_kv_cache, is_nvfp4_kv = (
|
||||
make_quantized_kv_cache(kv_cache, kv_quant_dtype, block_size, head_size)
|
||||
)
|
||||
|
||||
k_scale = v_scale = kv_scale
|
||||
|
||||
max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, NUM_BLOCKS, (batch_size, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
kv_indptr, kv_indices, kv_last_page_lens = build_paged_kv_metadata(
|
||||
seq_lens, block_tables, block_size
|
||||
)
|
||||
workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)
|
||||
|
||||
# Baseline Prefill
|
||||
if has_sinks:
|
||||
sinks = torch.rand(num_qo_heads, dtype=torch.float32) * 5
|
||||
wrapper = flashinfer.BatchAttentionWithAttentionSinkWrapper(
|
||||
float_workspace_buffer=workspace_buffer, kv_layout=kv_layout, backend="fa2"
|
||||
)
|
||||
else:
|
||||
sinks = None
|
||||
wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
|
||||
float_workspace_buffer=workspace_buffer, kv_layout=kv_layout, backend="fa2"
|
||||
)
|
||||
|
||||
wrapper.plan(
|
||||
qo_indptr=q_indptr,
|
||||
paged_kv_indptr=kv_indptr,
|
||||
paged_kv_indices=kv_indices,
|
||||
paged_kv_last_page_len=kv_last_page_lens,
|
||||
num_qo_heads=num_qo_heads,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_dim_qk=head_size,
|
||||
page_size=block_size,
|
||||
causal=True,
|
||||
sm_scale=sm_scale,
|
||||
window_left=window_left,
|
||||
logits_soft_cap=soft_cap,
|
||||
q_data_type=dtype,
|
||||
kv_data_type=dtype,
|
||||
)
|
||||
output = torch.empty(ref_query.shape, dtype=dtype)
|
||||
wrapper.run(ref_query, ref_kv_cache, sinks, sm_scale, out=output)
|
||||
|
||||
o_scale = 1.0
|
||||
o_sf_scale_float = None
|
||||
if o_quant_dtype == FP8_DTYPE:
|
||||
_, o_scale = to_float8(output)
|
||||
elif o_quant_dtype == FP4_DTYPE:
|
||||
o_sf_scale = get_nvfp4_global_scale(output)
|
||||
o_sf_scale_float = o_sf_scale.item()
|
||||
|
||||
# TRTLLM Prefill
|
||||
if o_quant_dtype == FP4_DTYPE:
|
||||
output_trtllm = flashinfer.utils.FP4Tensor(
|
||||
torch.empty(query.shape[:-1] + (query.shape[-1] // 2,), dtype=torch.uint8),
|
||||
torch.empty(
|
||||
(
|
||||
round_up(query.shape[0], 128),
|
||||
round_up(query.shape[1] * query.shape[2] // 16, 4),
|
||||
),
|
||||
dtype=torch.float8_e4m3fn,
|
||||
),
|
||||
)
|
||||
else:
|
||||
output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
|
||||
|
||||
flashinfer.prefill.trtllm_batch_context_with_kv_cache(
|
||||
query=query,
|
||||
kv_cache=kv_cache,
|
||||
workspace_buffer=workspace_buffer,
|
||||
block_tables=block_tables,
|
||||
seq_lens=seq_lens,
|
||||
max_q_len=max_q_len,
|
||||
max_kv_len=max_seq_len,
|
||||
bmm1_scale=q_scale * k_scale * sm_scale,
|
||||
bmm2_scale=v_scale / o_scale,
|
||||
batch_size=batch_size,
|
||||
cum_seq_lens_q=q_indptr,
|
||||
cum_seq_lens_kv=kv_indptr,
|
||||
window_left=window_left,
|
||||
sinks=sinks,
|
||||
o_sf_scale=o_sf_scale_float,
|
||||
out=output_trtllm,
|
||||
kv_cache_sf=kv_cache_sf,
|
||||
)
|
||||
if o_quant_dtype == FP8_DTYPE:
|
||||
output_trtllm = output_trtllm.to(dtype) * o_scale
|
||||
elif o_quant_dtype == FP4_DTYPE:
|
||||
output_trtllm.data = output_trtllm.data.reshape(
|
||||
-1, query.shape[1] * query.shape[2] // 2
|
||||
)
|
||||
output_trtllm = dequantize_nvfp4_to_dtype(
|
||||
output_trtllm.data, output_trtllm.scale, o_sf_scale, dtype, query.device
|
||||
)
|
||||
output_trtllm = output_trtllm.reshape(-1, query.shape[1], query.shape[2])
|
||||
|
||||
if is_nvfp4_kv:
|
||||
rtol, atol = 1.0, 1.5 # nvfp4 has higher quantization error
|
||||
elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE:
|
||||
rtol, atol = 3e-1, 4e-1
|
||||
elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
|
||||
rtol, atol = 4e-2, 6e-2
|
||||
elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype:
|
||||
rtol, atol = 2e-2, 3e-2
|
||||
else:
|
||||
rtol, atol = 1e-2, 1e-2
|
||||
|
||||
(
|
||||
torch.testing.assert_close(output, output_trtllm, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output - output_trtllm))}",
|
||||
)
|
||||
@@ -0,0 +1,178 @@
|
||||
# Adapted from: https://github.com/deepseek-ai/FlashMLA/blob/main/tests/test_flash_mla.py
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import math
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import triton
|
||||
from vllm.v1.attention.ops.flashmla import (
|
||||
flash_mla_with_kvcache,
|
||||
get_mla_metadata,
|
||||
is_flashmla_dense_supported,
|
||||
)
|
||||
|
||||
|
||||
def cal_diff(
|
||||
x: torch.Tensor, y: torch.Tensor, name: str, use_fp8: bool = False
|
||||
) -> None:
|
||||
x, y = x.double(), y.double()
|
||||
cos_diff = 1 - 2 * (x * y).sum().item() / max((x * x + y * y).sum().item(), 1e-12)
|
||||
if use_fp8:
|
||||
assert cos_diff < 1e-4
|
||||
else:
|
||||
assert cos_diff < 1e-5
|
||||
|
||||
|
||||
FLASH_MLA_UNSUPPORTED_REASON = (
|
||||
is_flashmla_dense_supported()[1]
|
||||
if not is_flashmla_dense_supported()[0]
|
||||
else "FlashMLA is supported"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_flashmla_dense_supported()[0], reason=FLASH_MLA_UNSUPPORTED_REASON
|
||||
)
|
||||
@pytest.mark.parametrize("b", [128])
|
||||
@pytest.mark.parametrize("s_q", [1, 2])
|
||||
@pytest.mark.parametrize("mean_sk", [4096, 8192, 16384])
|
||||
@pytest.mark.parametrize("h_q", [16, 32, 64, 128])
|
||||
@pytest.mark.parametrize("h_kv", [1])
|
||||
@pytest.mark.parametrize("d", [576])
|
||||
@pytest.mark.parametrize("dv", [512])
|
||||
@pytest.mark.parametrize("block_size", [64])
|
||||
@pytest.mark.parametrize("causal", [True])
|
||||
@pytest.mark.parametrize("varlen", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"torch_dtype", [torch.bfloat16, torch.float16, torch.float8_e4m3fn]
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_flash_mla(
|
||||
b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal, varlen, torch_dtype
|
||||
):
|
||||
device = torch.device("cuda:0")
|
||||
init_dtype = torch.bfloat16 if torch_dtype == torch.float8_e4m3fn else torch_dtype
|
||||
torch.set_default_dtype(init_dtype)
|
||||
torch.set_default_device(device)
|
||||
torch.accelerator.set_device_index(device)
|
||||
torch.manual_seed(0)
|
||||
random.seed(0)
|
||||
|
||||
print(
|
||||
f"{b=}, {s_q=}, {mean_sk=}, {h_q=}, {h_kv=}, "
|
||||
f"{d=}, {dv=}, {causal=}, {varlen=}, {torch_dtype=}"
|
||||
)
|
||||
|
||||
use_fp8 = torch_dtype == torch.float8_e4m3fn
|
||||
cache_seqlens = torch.full((b,), mean_sk, dtype=torch.int32)
|
||||
if varlen:
|
||||
for i in range(b):
|
||||
cache_seqlens[i] = max(random.normalvariate(mean_sk, mean_sk / 2), s_q)
|
||||
total_seqlens = cache_seqlens.sum().item()
|
||||
max_seqlen = cache_seqlens.max().item()
|
||||
max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256
|
||||
|
||||
q = torch.randn(b, s_q, h_q, d)
|
||||
block_table = torch.arange(
|
||||
b * max_seqlen_pad // block_size, dtype=torch.int32
|
||||
).view(b, max_seqlen_pad // block_size)
|
||||
blocked_k = torch.randn(block_table.numel(), block_size, h_kv, d)
|
||||
for i in range(b):
|
||||
blocked_k.view(b, max_seqlen_pad, h_kv, d)[i, cache_seqlens[i].item() :] = (
|
||||
float("nan")
|
||||
)
|
||||
blocked_v = blocked_k[..., :dv]
|
||||
|
||||
tile_scheduler_metadata, num_splits = get_mla_metadata(
|
||||
cache_seqlens, s_q * h_q // h_kv, h_kv
|
||||
)
|
||||
|
||||
init_dtype = q.dtype
|
||||
if use_fp8:
|
||||
fp8_dtype = torch.float8_e4m3fn
|
||||
descale_q = torch.ones((1), dtype=torch.float32)
|
||||
descale_k = torch.ones((1), dtype=torch.float32)
|
||||
|
||||
q = q.to(fp8_dtype)
|
||||
blocked_k = blocked_k.to(fp8_dtype)
|
||||
blocked_v = blocked_v.to(fp8_dtype)
|
||||
else:
|
||||
descale_q = None
|
||||
descale_k = None
|
||||
|
||||
def flash_mla():
|
||||
return flash_mla_with_kvcache(
|
||||
q,
|
||||
blocked_k,
|
||||
block_table,
|
||||
cache_seqlens,
|
||||
dv,
|
||||
tile_scheduler_metadata,
|
||||
num_splits,
|
||||
causal=causal,
|
||||
descale_q=descale_q,
|
||||
descale_k=descale_k,
|
||||
)
|
||||
|
||||
def scaled_dot_product_attention(query, key, value, is_causal=False):
|
||||
query = query.float()
|
||||
key = key.float()
|
||||
value = value.float()
|
||||
key = key.repeat_interleave(h_q // h_kv, dim=0)
|
||||
value = value.repeat_interleave(h_q // h_kv, dim=0)
|
||||
attn_weight = query @ key.transpose(-2, -1) / math.sqrt(query.size(-1))
|
||||
if is_causal:
|
||||
s_q = query.shape[-2]
|
||||
s_k = key.shape[-2]
|
||||
attn_bias = torch.zeros(s_q, s_k, dtype=query.dtype)
|
||||
temp_mask = torch.ones(s_q, s_k, dtype=torch.bool).tril(diagonal=s_k - s_q)
|
||||
attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
|
||||
attn_bias.to(query.dtype)
|
||||
attn_weight += attn_bias
|
||||
lse = attn_weight.logsumexp(dim=-1)
|
||||
attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32)
|
||||
return attn_weight @ value, lse
|
||||
|
||||
def ref_mla():
|
||||
q_ = (q.to(torch.float) * descale_q).to(init_dtype) if use_fp8 else q
|
||||
blocked_k_ = (
|
||||
(blocked_k.to(torch.float) * descale_k).to(init_dtype)
|
||||
if use_fp8
|
||||
else blocked_k
|
||||
)
|
||||
blocked_v_ = (
|
||||
(blocked_v.to(torch.float) * descale_k).to(init_dtype)
|
||||
if use_fp8
|
||||
else blocked_v
|
||||
)
|
||||
out = torch.empty(b, s_q, h_q, dv, dtype=torch.float32)
|
||||
lse = torch.empty(b, h_q, s_q, dtype=torch.float32)
|
||||
for i in range(b):
|
||||
begin = i * max_seqlen_pad
|
||||
end = begin + cache_seqlens[i]
|
||||
out_i, lse_i = scaled_dot_product_attention(
|
||||
q_[i].transpose(0, 1),
|
||||
blocked_k_.view(-1, h_kv, d)[begin:end].transpose(0, 1),
|
||||
blocked_v_.view(-1, h_kv, dv)[begin:end].transpose(0, 1),
|
||||
is_causal=causal,
|
||||
)
|
||||
out[i] = out_i.transpose(0, 1)
|
||||
lse[i] = lse_i
|
||||
return out, lse
|
||||
|
||||
out_flash, lse_flash = flash_mla()
|
||||
out_torch, lse_torch = ref_mla()
|
||||
cal_diff(out_flash, out_torch, "out", use_fp8)
|
||||
cal_diff(lse_flash, lse_torch, "lse")
|
||||
|
||||
t = triton.testing.do_bench(flash_mla)
|
||||
FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2
|
||||
bytes = (total_seqlens * h_kv * d + b * s_q * h_q * d) * (
|
||||
torch.finfo(torch_dtype).bits // 8
|
||||
) + (b * s_q * h_q * dv) * (torch.finfo(init_dtype).bits // 8)
|
||||
print(
|
||||
f"{t:.3f} ms, {FLOPS / 10**9 / t:.0f} TFLOPS,", f"{bytes / 10**6 / t:.0f} GB/s"
|
||||
)
|
||||
@@ -0,0 +1,292 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
def test_sparse_flashmla_metadata_smoke():
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_sparse_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
|
||||
device = torch.device("cuda")
|
||||
batch_size = 1
|
||||
seqlen_q = 1
|
||||
num_heads_q = 128
|
||||
num_heads_k = 1
|
||||
q_seq_per_hk = seqlen_q * num_heads_q // num_heads_k
|
||||
topk = 128
|
||||
|
||||
cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device)
|
||||
|
||||
tile_md, num_splits = fm.get_mla_metadata(
|
||||
cache_seqlens,
|
||||
q_seq_per_hk,
|
||||
num_heads_k,
|
||||
num_heads_q=num_heads_q,
|
||||
topk=topk,
|
||||
is_fp8_kvcache=True,
|
||||
)
|
||||
assert isinstance(tile_md, fm.FlashMLASchedMeta)
|
||||
assert tile_md.tile_scheduler_metadata is None
|
||||
assert tile_md.num_splits is None
|
||||
assert num_splits is None
|
||||
|
||||
|
||||
def test_sparse_flashmla_decode_smoke():
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_sparse_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
|
||||
device = torch.device("cuda")
|
||||
batch_size = 1
|
||||
seqlen_q = 1
|
||||
num_heads_q = 64
|
||||
head_dim_k = 576
|
||||
head_dim_v = 512
|
||||
num_heads_k = 1
|
||||
page_block_size = 64
|
||||
bytes_per_token = 656
|
||||
topk = 128
|
||||
|
||||
# Metadata
|
||||
q_seq_per_hk = seqlen_q * num_heads_q // num_heads_k
|
||||
# q_heads_per_hk = num_heads_q // num_heads_k
|
||||
cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device)
|
||||
tile_md, num_splits = fm.get_mla_metadata(
|
||||
cache_seqlens,
|
||||
q_seq_per_hk,
|
||||
num_heads_k,
|
||||
num_heads_q=num_heads_q,
|
||||
topk=topk,
|
||||
is_fp8_kvcache=True,
|
||||
)
|
||||
|
||||
# Inputs
|
||||
q = torch.zeros(
|
||||
(batch_size, seqlen_q, num_heads_q, head_dim_k),
|
||||
dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
k_cache = torch.zeros(
|
||||
(1, page_block_size, num_heads_k, bytes_per_token),
|
||||
dtype=torch.uint8,
|
||||
device=device,
|
||||
)
|
||||
indices = torch.zeros(
|
||||
(batch_size, seqlen_q, topk), dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
block_table = torch.zeros((batch_size, 128), dtype=torch.int32, device=device)
|
||||
out, lse = fm.flash_mla_with_kvcache(
|
||||
q,
|
||||
k_cache,
|
||||
block_table,
|
||||
cache_seqlens,
|
||||
head_dim_v,
|
||||
tile_md,
|
||||
num_splits,
|
||||
indices=indices,
|
||||
is_fp8_kvcache=True,
|
||||
)
|
||||
assert out.shape[0] == batch_size
|
||||
assert out.shape[-1] == head_dim_v
|
||||
assert lse.shape[0] == batch_size
|
||||
|
||||
|
||||
def test_sparse_flashmla_prefill_smoke():
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_sparse_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
|
||||
device = torch.device("cuda")
|
||||
s_q = 1
|
||||
s_kv = 1
|
||||
h_q = 64 # kernel expects multiple of 64
|
||||
h_kv = 1
|
||||
d_qk = 576
|
||||
d_v = 512
|
||||
topk = 128
|
||||
|
||||
q = torch.zeros((s_q, h_q, d_qk), dtype=torch.bfloat16, device=device)
|
||||
kv = torch.zeros((s_kv, h_kv, d_qk), dtype=torch.bfloat16, device=device)
|
||||
indices = torch.zeros((s_q, h_kv, topk), dtype=torch.int32, device=device)
|
||||
|
||||
out, max_logits, lse = fm.flash_mla_sparse_fwd(q, kv, indices, 1.0, d_v)
|
||||
assert out.shape == (s_q, h_q, d_v)
|
||||
assert max_logits.shape == (s_q, h_q)
|
||||
assert lse.shape == (s_q, h_q)
|
||||
|
||||
|
||||
def test_deepseek_v4_prefill_chunk_planning_expands_for_short_sequences():
|
||||
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
|
||||
|
||||
metadata = DeepseekSparseSWAMetadata(
|
||||
block_table=torch.empty(0, dtype=torch.int32),
|
||||
slot_mapping=torch.empty(0, dtype=torch.int32),
|
||||
block_size=64,
|
||||
num_prefills=5,
|
||||
prefill_seq_lens_cpu=torch.tensor([80, 96, 112, 128, 144], dtype=torch.int32),
|
||||
prefill_query_lens_cpu=torch.tensor([4, 4, 4, 4, 4], dtype=torch.int32),
|
||||
prefill_window_size=64,
|
||||
prefill_max_model_len=1024,
|
||||
prefill_max_num_batched_tokens=128,
|
||||
)
|
||||
|
||||
chunk_plan = metadata.get_prefill_chunk_plan(compress_ratio=4, prefill_chunk_size=4)
|
||||
|
||||
# the adaptive plan keeps all 5 in one chunk
|
||||
assert chunk_plan == [(0, 5, 36, 103)]
|
||||
|
||||
|
||||
def test_flashinfer_sparse_indices_cache(monkeypatch):
|
||||
from vllm.models.deepseek_v4.nvidia import flashinfer_sparse as flashinfer_mod
|
||||
from vllm.models.deepseek_v4.sparse_mla import DeepseekV4FlashMLAMetadata
|
||||
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
|
||||
|
||||
builder_calls = 0
|
||||
|
||||
def fake_build(*args, **kwargs):
|
||||
nonlocal builder_calls
|
||||
builder_calls += 1
|
||||
return (
|
||||
torch.tensor([[builder_calls]], dtype=torch.int32),
|
||||
torch.tensor([builder_calls], dtype=torch.int32),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
flashinfer_mod, "build_flashinfer_mixed_sparse_indices", fake_build
|
||||
)
|
||||
|
||||
def make_attn(compress_ratio: int, topk_width: int):
|
||||
attn = object.__new__(flashinfer_mod.DeepseekV4FlashInferMLAAttention)
|
||||
attn.compress_ratio = compress_ratio
|
||||
attn.window_size = 4
|
||||
attn.topk_indices_buffer = torch.tensor(
|
||||
[[0, 1], [2, 3], [4, 5]], dtype=torch.int32
|
||||
)[:, :topk_width]
|
||||
return attn
|
||||
|
||||
def make_swa_metadata():
|
||||
return DeepseekSparseSWAMetadata(
|
||||
block_table=torch.tensor([[0, 1], [2, 3]], dtype=torch.int32),
|
||||
slot_mapping=torch.tensor([0, 1], dtype=torch.int64),
|
||||
block_size=64,
|
||||
seq_lens=torch.tensor([8, 10], dtype=torch.int32),
|
||||
query_start_loc=torch.tensor([0, 1, 3], dtype=torch.int32),
|
||||
query_start_loc_cpu=torch.tensor([0, 1, 3], dtype=torch.int32),
|
||||
token_to_req_indices=torch.tensor([0, 1, 1], dtype=torch.int32),
|
||||
decode_swa_indices=torch.tensor([[5, 6, -1, -1]], dtype=torch.int32),
|
||||
decode_swa_lens=torch.tensor([2], dtype=torch.int32),
|
||||
is_valid_token=torch.tensor([True], dtype=torch.bool),
|
||||
num_decodes=1,
|
||||
num_prefills=1,
|
||||
num_decode_tokens=1,
|
||||
num_prefill_tokens=2,
|
||||
)
|
||||
|
||||
def make_flashmla_metadata():
|
||||
return DeepseekV4FlashMLAMetadata(
|
||||
num_reqs=2,
|
||||
max_query_len=2,
|
||||
max_seq_len=10,
|
||||
num_actual_tokens=3,
|
||||
query_start_loc=torch.tensor([0, 1, 3], dtype=torch.int32),
|
||||
slot_mapping=torch.tensor([0, 1, 2], dtype=torch.int64),
|
||||
block_table=torch.tensor([[0, 1], [2, 3]], dtype=torch.int32),
|
||||
req_id_per_token=torch.tensor([0, 1, 1], dtype=torch.int32),
|
||||
block_size=256,
|
||||
topk_tokens=2,
|
||||
c128a_global_decode_topk_indices=torch.tensor(
|
||||
[[[9, 10]]], dtype=torch.int32
|
||||
),
|
||||
c128a_decode_topk_lens=torch.tensor([2], dtype=torch.int32),
|
||||
c128a_prefill_topk_indices=torch.tensor(
|
||||
[[0, 1], [1, 2]], dtype=torch.int32
|
||||
),
|
||||
)
|
||||
|
||||
swa_attn = make_attn(1, 0)
|
||||
swa_metadata = make_swa_metadata()
|
||||
_, _, sparse_indices_first, sparse_lens_first = (
|
||||
swa_attn._build_sparse_index_metadata(
|
||||
kv_cache=None,
|
||||
swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
|
||||
swa_metadata=swa_metadata,
|
||||
attn_metadata=None,
|
||||
swa_only=True,
|
||||
)
|
||||
)
|
||||
_, _, sparse_indices_second, sparse_lens_second = (
|
||||
swa_attn._build_sparse_index_metadata(
|
||||
kv_cache=None,
|
||||
swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
|
||||
swa_metadata=swa_metadata,
|
||||
attn_metadata=None,
|
||||
swa_only=True,
|
||||
)
|
||||
)
|
||||
assert builder_calls == 1
|
||||
assert sparse_indices_first is sparse_indices_second
|
||||
assert sparse_lens_first is sparse_lens_second
|
||||
|
||||
c128a_attn = make_attn(128, 2)
|
||||
c128a_metadata = make_swa_metadata()
|
||||
c128a_flashmla_md = make_flashmla_metadata()
|
||||
_, _, sparse_indices_first, sparse_lens_first = (
|
||||
c128a_attn._build_sparse_index_metadata(
|
||||
kv_cache=torch.empty((1, 2, 512), dtype=torch.bfloat16),
|
||||
swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
|
||||
swa_metadata=c128a_metadata,
|
||||
attn_metadata=c128a_flashmla_md,
|
||||
swa_only=False,
|
||||
)
|
||||
)
|
||||
_, _, sparse_indices_second, sparse_lens_second = (
|
||||
c128a_attn._build_sparse_index_metadata(
|
||||
kv_cache=torch.empty((1, 2, 512), dtype=torch.bfloat16),
|
||||
swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
|
||||
swa_metadata=c128a_metadata,
|
||||
attn_metadata=c128a_flashmla_md,
|
||||
swa_only=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert builder_calls == 2
|
||||
assert sparse_indices_first is sparse_indices_second
|
||||
assert sparse_lens_first is sparse_lens_second
|
||||
|
||||
c4a_attn = make_attn(4, 2)
|
||||
c4a_metadata = make_swa_metadata()
|
||||
c4a_flashmla_md = make_flashmla_metadata()
|
||||
c4a_flashmla_md.c128a_global_decode_topk_indices = None
|
||||
c4a_flashmla_md.c128a_decode_topk_lens = None
|
||||
c4a_flashmla_md.c128a_prefill_topk_indices = None
|
||||
_, _, sparse_indices_third, sparse_lens_third = (
|
||||
c4a_attn._build_sparse_index_metadata(
|
||||
kv_cache=torch.empty((1, 2, 512), dtype=torch.bfloat16),
|
||||
swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
|
||||
swa_metadata=c4a_metadata,
|
||||
attn_metadata=c4a_flashmla_md,
|
||||
swa_only=False,
|
||||
)
|
||||
)
|
||||
_, _, sparse_indices_fourth, sparse_lens_fourth = (
|
||||
c4a_attn._build_sparse_index_metadata(
|
||||
kv_cache=torch.empty((1, 2, 512), dtype=torch.bfloat16),
|
||||
swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
|
||||
swa_metadata=c4a_metadata,
|
||||
attn_metadata=c4a_flashmla_md,
|
||||
swa_only=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert builder_calls == 4
|
||||
assert sparse_indices_third is not sparse_indices_fourth
|
||||
assert sparse_lens_third is not sparse_lens_fourth
|
||||
@@ -0,0 +1,268 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.lightning_attn import linear_decode_forward_triton
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
DEVICE = current_platform.device_type
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
|
||||
reason="Lightning attention Triton kernels require CUDA/ROCm or XPU.",
|
||||
)
|
||||
|
||||
NUM_HEADS = [4, 8]
|
||||
HEAD_SIZES = [64]
|
||||
BATCH_SIZES = [1, 2]
|
||||
SEQ_LENGTHS = [16]
|
||||
DTYPES = [torch.float32]
|
||||
|
||||
|
||||
def reference_lightning_attention(q, k, v, ed, block_size, kv_history):
|
||||
"""Reference implementation of lightning attention core algorithm
|
||||
|
||||
The difference from the main implementation is that this processes
|
||||
each step sequentially, instead of using parallelized triton kernels
|
||||
"""
|
||||
B, H, S, D = q.shape
|
||||
E = v.shape[-1]
|
||||
dtype = q.dtype
|
||||
output = torch.zeros((B, H, S, E), dtype=dtype, device=q.device)
|
||||
|
||||
# Use clone() to ensure an independent copy
|
||||
if kv_history is None:
|
||||
kv_cache = torch.zeros((B, H, D, E), dtype=dtype, device=q.device)
|
||||
else:
|
||||
kv_cache = kv_history.clone()
|
||||
|
||||
# More efficient implementation
|
||||
# Convert decay factors to matrix form
|
||||
decay = torch.exp(-ed).view(1, -1, 1, 1) if ed.dim() == 1 else torch.exp(-ed)
|
||||
|
||||
for b in range(B):
|
||||
for step in range(S):
|
||||
# Process all heads at once for this position
|
||||
q_bs = q[b, :, step] # [H, D]
|
||||
k_bs = k[b, :, step] # [H, D]
|
||||
v_bs = v[b, :, step] # [H, E]
|
||||
|
||||
# Calculate KV outer products for all heads
|
||||
for h in range(H):
|
||||
# Calculate KV outer product
|
||||
kv_outer = torch.outer(k_bs[h], v_bs[h])
|
||||
|
||||
# Update KV cache with decay
|
||||
# Note: Using the same order as in the Triton kernel
|
||||
kv_cache[b, h] = decay[0, h, 0, 0] * kv_cache[b, h] + kv_outer
|
||||
|
||||
# Calculate attention output
|
||||
output[b, h, step] = torch.matmul(q_bs[h], kv_cache[b, h])
|
||||
|
||||
# Match the shape returned by the actual implementation
|
||||
# The actual implementation returns a tensor of shape [B, H, 2, D, E]
|
||||
# where dimension 2 contains both KV and KV history
|
||||
kv_reshaped = kv_cache.unsqueeze(2) # [B, H, 1, D, E]
|
||||
final_kv_cache = torch.cat([kv_reshaped, kv_reshaped], dim=2) # [B, H, 2, D, E]
|
||||
|
||||
return output, final_kv_cache
|
||||
|
||||
|
||||
def reference_linear_decode(q, k, v, kv_caches, slope_rate, slot_idx):
|
||||
"""Reference implementation: linear attention decode function"""
|
||||
B, H, _, D = q.shape
|
||||
output = torch.zeros(B, H * D, dtype=q.dtype, device=q.device)
|
||||
|
||||
# Calculate decay factors once (more efficient)
|
||||
decay = torch.exp(-slope_rate).view(-1, 1, 1) # [H, 1, 1]
|
||||
|
||||
# Process each batch
|
||||
for b in range(B):
|
||||
slot_id = slot_idx[b].item()
|
||||
|
||||
# Skip padding positions
|
||||
if slot_id == -1:
|
||||
continue
|
||||
|
||||
# Process all heads at once for this batch
|
||||
q_b = q[b, :, 0] # [H, D]
|
||||
k_b = k[b, :, 0] # [H, D]
|
||||
v_b = v[b, :, 0] # [H, D]
|
||||
|
||||
# Process each attention head
|
||||
for h in range(H):
|
||||
# Get current query, key and value
|
||||
q_bh = q_b[h]
|
||||
k_bh = k_b[h]
|
||||
v_bh = v_b[h]
|
||||
|
||||
# Get cache
|
||||
kv_cache_old = kv_caches[b, h]
|
||||
|
||||
# Calculate new key-value outer product
|
||||
kv_outer = torch.outer(k_bh, v_bh)
|
||||
|
||||
# Apply decay and update cache
|
||||
kv_new = kv_outer + decay[h, 0, 0] * kv_cache_old
|
||||
|
||||
# Calculate output
|
||||
out_h = torch.matmul(q_bh, kv_new)
|
||||
|
||||
# Update output and cache
|
||||
output[b, h * D : (h + 1) * D] = out_h
|
||||
kv_caches[b, h] = kv_new
|
||||
|
||||
return output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_linear_decode_forward_triton(
|
||||
batch_size: int,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
torch.set_default_device(DEVICE)
|
||||
set_random_seed(42)
|
||||
base = 0.01
|
||||
q = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
|
||||
k = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
|
||||
v = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
|
||||
|
||||
kv_caches = base * torch.randn(
|
||||
batch_size, num_heads, head_size, head_size, dtype=dtype, device=DEVICE
|
||||
)
|
||||
|
||||
kv_caches_copy = kv_caches.clone()
|
||||
|
||||
slope_rate = torch.zeros(num_heads, device=DEVICE)
|
||||
for h in range(num_heads):
|
||||
slope_rate[h] = 0.1 * (h + 1)
|
||||
|
||||
slot_idx = torch.arange(batch_size, device=DEVICE)
|
||||
|
||||
triton_output = linear_decode_forward_triton(
|
||||
q, k, v, kv_caches, slope_rate, slot_idx
|
||||
)
|
||||
|
||||
reference_output = reference_linear_decode(
|
||||
q, k, v, kv_caches_copy, slope_rate, slot_idx
|
||||
)
|
||||
torch.testing.assert_close(triton_output, reference_output, rtol=1e-1, atol=1e-1)
|
||||
torch.testing.assert_close(kv_caches, kv_caches_copy, rtol=1e-1, atol=1e-1)
|
||||
|
||||
assert triton_output.shape == (batch_size, num_heads * head_size)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_linear_decode_forward_triton_with_padding(
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
torch.set_default_device(DEVICE)
|
||||
set_random_seed(42)
|
||||
|
||||
batch_size = 4
|
||||
base = 0.01
|
||||
q = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
|
||||
k = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
|
||||
v = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
|
||||
|
||||
kv_caches = base * torch.randn(
|
||||
batch_size, num_heads, head_size, head_size, dtype=dtype, device=DEVICE
|
||||
)
|
||||
|
||||
kv_caches_copy = kv_caches.clone()
|
||||
|
||||
slope_rate = torch.zeros(num_heads, device=DEVICE)
|
||||
for h in range(num_heads):
|
||||
slope_rate[h] = 0.1 * (h + 1)
|
||||
|
||||
slot_idx = torch.tensor([0, 1, -1, 2], device=DEVICE)
|
||||
|
||||
triton_output = linear_decode_forward_triton(
|
||||
q, k, v, kv_caches, slope_rate, slot_idx
|
||||
)
|
||||
|
||||
reference_output = reference_linear_decode(
|
||||
q, k, v, kv_caches_copy, slope_rate, slot_idx
|
||||
)
|
||||
|
||||
padding_mask = (slot_idx != -1).unsqueeze(1).expand(-1, num_heads * head_size)
|
||||
|
||||
triton_masked = triton_output[padding_mask]
|
||||
reference_masked = reference_output[padding_mask]
|
||||
|
||||
atol, rtol = 1.5e-1, 1.5e-1
|
||||
|
||||
valid_indices = slot_idx != -1
|
||||
|
||||
for i in range(batch_size):
|
||||
if valid_indices[i] > 0:
|
||||
torch.testing.assert_close(
|
||||
kv_caches[i], kv_caches_copy[i], rtol=rtol, atol=atol
|
||||
)
|
||||
|
||||
torch.testing.assert_close(triton_masked, reference_masked, rtol=rtol, atol=atol)
|
||||
|
||||
assert triton_output.shape == (batch_size, num_heads * head_size)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENGTHS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_lightning_attention_reference(
|
||||
batch_size: int,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
seq_len: int,
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
torch.set_default_device(DEVICE)
|
||||
set_random_seed(42)
|
||||
|
||||
base = 0.01
|
||||
q = base * torch.randn(batch_size, num_heads, seq_len, head_size, dtype=dtype)
|
||||
k = base * torch.randn(batch_size, num_heads, seq_len, head_size, dtype=dtype)
|
||||
v = base * torch.randn(batch_size, num_heads, seq_len, head_size, dtype=dtype)
|
||||
|
||||
ed = torch.zeros(num_heads, device=DEVICE)
|
||||
for h in range(num_heads):
|
||||
ed[h] = 0.1 * (h + 1)
|
||||
|
||||
kv_history = base * torch.randn(
|
||||
batch_size, num_heads, head_size, head_size, dtype=dtype, device=DEVICE
|
||||
)
|
||||
|
||||
kv_history_clone = kv_history.clone()
|
||||
|
||||
ref_output, ref_kv_cache = reference_lightning_attention(
|
||||
q, k, v, ed, 256, kv_history
|
||||
)
|
||||
|
||||
from vllm.model_executor.layers.lightning_attn import lightning_attention
|
||||
|
||||
actual_output, actual_kv_cache = lightning_attention(
|
||||
q, k, v, ed, 256, kv_history_clone
|
||||
)
|
||||
|
||||
atol, rtol = 1.5e-1, 1.5e-1
|
||||
torch.testing.assert_close(ref_output, actual_output, rtol=rtol, atol=atol)
|
||||
torch.testing.assert_close(ref_kv_cache, actual_kv_cache, rtol=rtol, atol=atol)
|
||||
|
||||
assert ref_output.shape == (batch_size, num_heads, seq_len, head_size)
|
||||
assert ref_kv_cache.shape == actual_kv_cache.shape
|
||||
@@ -0,0 +1,392 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm._custom_ops import (
|
||||
merge_attn_states as merge_attn_states_cuda,
|
||||
)
|
||||
from vllm._custom_ops import (
|
||||
scaled_fp8_quant,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.ops.triton_merge_attn_states import (
|
||||
merge_attn_states as merge_attn_states_triton,
|
||||
)
|
||||
|
||||
|
||||
# Naive PyTorch Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
|
||||
# can be used to combine partial attention results (in the split-KV case)
|
||||
def merge_attn_states_torch(
|
||||
output: torch.Tensor, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
|
||||
prefix_output: torch.Tensor, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
|
||||
prefix_lse: torch.Tensor, # [NUM_HEADS, NUM_TOKENS]
|
||||
suffix_output: torch.Tensor, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
|
||||
suffix_lse: torch.Tensor, # [NUM_HEADS, NUM_TOKENS]
|
||||
output_lse: torch.Tensor | None = None, # [NUM_HEADS, NUM_TOKENS]
|
||||
prefill_tokens_with_context: int | None = None,
|
||||
output_scale: torch.Tensor | None = None, # scalar, per-tensor FP8 scale
|
||||
):
|
||||
# Apply prefill_tokens_with_context mask if needed
|
||||
if prefill_tokens_with_context is None:
|
||||
prefill_tokens_with_context = output.shape[0]
|
||||
p_lse = prefix_lse
|
||||
s_lse = suffix_lse
|
||||
# inf -> -inf
|
||||
p_lse[p_lse == torch.inf] = -torch.inf
|
||||
s_lse[s_lse == torch.inf] = -torch.inf
|
||||
# max_lse [NUM_HEADS, NUM_TOKENS]
|
||||
max_lse = torch.maximum(p_lse, s_lse)
|
||||
|
||||
mask = torch.ones((prefix_lse.shape[1], 1, 1), device=p_lse.device)
|
||||
mask[prefill_tokens_with_context:].fill_(0)
|
||||
p_lse = p_lse - max_lse
|
||||
s_lse = s_lse - max_lse
|
||||
p_lse_exp = torch.exp(p_lse)
|
||||
s_lse_exp = torch.exp(s_lse)
|
||||
out_se = p_lse_exp + s_lse_exp
|
||||
if output_lse is not None:
|
||||
output_lse = torch.log(out_se) + max_lse
|
||||
output_lse[prefill_tokens_with_context:] = suffix_lse[
|
||||
prefill_tokens_with_context:
|
||||
]
|
||||
p_scale = p_lse_exp / out_se # [NUM_HEADS, NUM_TOKENS]
|
||||
s_scale = s_lse_exp / out_se # [NUM_HEADS, NUM_TOKENS]
|
||||
p_scale = torch.transpose(p_scale, 0, 1).unsqueeze(2) # [NUM_TOKENS, NUM_HEADS, 1]
|
||||
s_scale = torch.transpose(s_scale, 0, 1).unsqueeze(2) # [NUM_TOKENS, NUM_HEADS, 1]
|
||||
output = prefix_output * p_scale * mask + suffix_output * (
|
||||
s_scale * mask + (1 - mask)
|
||||
)
|
||||
if output_scale is not None:
|
||||
shape = output.shape
|
||||
output, _ = scaled_fp8_quant(output.float().view(-1, shape[-1]), output_scale)
|
||||
output = output.view(shape)
|
||||
return output, output_lse
|
||||
|
||||
|
||||
NUM_BATCH_TOKENS = [256, 512, 613, 1024, 1536, 4096]
|
||||
NUM_QUERY_HEADS = [4, 8, 16, 32, 48, 64]
|
||||
HEAD_SIZES = [32, 48, 64, 96, 128, 256]
|
||||
DTYPES = [torch.float32, torch.half, torch.bfloat16]
|
||||
|
||||
all_case_info: list[tuple] = []
|
||||
|
||||
|
||||
def generate_markdown_table():
|
||||
global all_case_info
|
||||
table_header = (
|
||||
"| tokens | heads | headsize | dtype "
|
||||
"| device | torch | triton | cuda | speedup |"
|
||||
)
|
||||
table_separator = "| --- | --- | --- | --- | --- | --- | --- | --- | --- |"
|
||||
|
||||
def shortly_dtype(dtype: torch.dtype) -> str:
|
||||
return str(dtype).removeprefix("torch.")
|
||||
|
||||
def shortly_device(device: str) -> str:
|
||||
return device.removeprefix("NVIDIA").strip()
|
||||
|
||||
print(table_header)
|
||||
print(table_separator)
|
||||
for info in all_case_info:
|
||||
(
|
||||
num_tokens,
|
||||
num_heads,
|
||||
head_size,
|
||||
dtype,
|
||||
device,
|
||||
avg_time_torch_kernel,
|
||||
avg_time_triton_kernel,
|
||||
avg_time_cuda_kernel,
|
||||
performance_improved,
|
||||
) = info
|
||||
dtype = shortly_dtype(dtype)
|
||||
device = shortly_device(device)
|
||||
print(
|
||||
f"| {num_tokens} | {num_heads} | {head_size} "
|
||||
f"| {dtype} | {device} | {avg_time_torch_kernel:.5f}ms "
|
||||
f"| {avg_time_triton_kernel:.5f}ms "
|
||||
f"| {avg_time_cuda_kernel:.5f}ms "
|
||||
f"| {performance_improved:.4f}x |"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_fp8", [False, True])
|
||||
@pytest.mark.parametrize("prefill_tokens_with_context", [None, 128])
|
||||
@pytest.mark.parametrize("num_tokens", NUM_BATCH_TOKENS)
|
||||
@pytest.mark.parametrize("num_query_heads", NUM_QUERY_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("input_dtype", DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_merge_attn_states(
|
||||
prefill_tokens_with_context: int | None,
|
||||
num_tokens: int,
|
||||
num_query_heads: int,
|
||||
head_size: int,
|
||||
input_dtype: torch.dtype,
|
||||
use_fp8: bool,
|
||||
):
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip(
|
||||
"Currently only support compare triton merge_attn_states "
|
||||
"with custom cuda merge_attn_states kernel"
|
||||
)
|
||||
|
||||
NUM_TOKENS = num_tokens
|
||||
NUM_HEADS = num_query_heads
|
||||
HEAD_SIZE = head_size
|
||||
|
||||
# When use_fp8 is set, inputs stay as input_dtype (bf16/fp16/fp32)
|
||||
# and output becomes FP8.
|
||||
output_dtype = input_dtype
|
||||
output_scale = None
|
||||
if use_fp8:
|
||||
output_dtype = current_platform.fp8_dtype()
|
||||
output_scale = torch.tensor([0.05], dtype=torch.float32, device="cuda")
|
||||
|
||||
print(
|
||||
f"\nNUM_TOKENS:{NUM_TOKENS}, NUM_HEADS:{NUM_HEADS}, "
|
||||
f"HEAD_SIZE:{HEAD_SIZE}, input_dtype: {input_dtype}, "
|
||||
f"output_dtype: {output_dtype}, use_fp8: {use_fp8}, "
|
||||
f"prefill_tokens_with_context: {prefill_tokens_with_context}, "
|
||||
f"Device: {current_platform.get_device_name()}"
|
||||
)
|
||||
|
||||
# prefix_lse and suffix_lse contain inf and normal values
|
||||
prefix_lse = torch.randn(NUM_HEADS, NUM_TOKENS, dtype=torch.float32, device="cuda")
|
||||
suffix_lse = torch.randn(NUM_HEADS, NUM_TOKENS, dtype=torch.float32, device="cuda")
|
||||
|
||||
# Generate boolean masks
|
||||
mask_prefix = torch.rand(NUM_HEADS, NUM_TOKENS) < 0.1
|
||||
mask_suffix = torch.rand(NUM_HEADS, NUM_TOKENS) < 0.1
|
||||
# Ensure that the same position is not True at the same time
|
||||
combined_mask = torch.logical_and(mask_prefix, mask_suffix)
|
||||
mask_prefix = torch.logical_and(mask_prefix, ~combined_mask)
|
||||
mask_suffix = torch.logical_and(mask_suffix, ~combined_mask)
|
||||
|
||||
prefix_lse[mask_prefix] = float("inf")
|
||||
suffix_lse[mask_suffix] = float("inf")
|
||||
|
||||
# Other input tensors (need to be initialized but
|
||||
# no actual calculation needed)
|
||||
output = torch.zeros(
|
||||
(NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device="cuda"
|
||||
)
|
||||
output_lse = torch.zeros(
|
||||
(NUM_HEADS, NUM_TOKENS), dtype=torch.float32, device="cuda"
|
||||
)
|
||||
prefix_output = torch.randn(
|
||||
(NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device="cuda"
|
||||
)
|
||||
suffix_output = torch.randn(
|
||||
(NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device="cuda"
|
||||
)
|
||||
|
||||
warmup_times = 2
|
||||
repeat_times = 20
|
||||
|
||||
output_torch = output.clone()
|
||||
output_lse_torch = output_lse.clone()
|
||||
total_time_torch_kernel = 0
|
||||
start = torch.Event(enable_timing=True)
|
||||
end = torch.Event(enable_timing=True)
|
||||
|
||||
# 0. Run the Torch kernel
|
||||
prefix_lse_torch = prefix_lse.clone()
|
||||
suffix_lse_torch = suffix_lse.clone()
|
||||
for _ in range(warmup_times):
|
||||
output_torch, output_lse_torch = merge_attn_states_torch(
|
||||
output_torch,
|
||||
prefix_output,
|
||||
prefix_lse_torch,
|
||||
suffix_output,
|
||||
suffix_lse_torch,
|
||||
output_lse_torch,
|
||||
prefill_tokens_with_context,
|
||||
output_scale,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
for _ in range(repeat_times):
|
||||
start.record()
|
||||
output_torch, output_lse_torch = merge_attn_states_torch(
|
||||
output_torch,
|
||||
prefix_output,
|
||||
prefix_lse_torch,
|
||||
suffix_output,
|
||||
suffix_lse_torch,
|
||||
output_lse_torch,
|
||||
prefill_tokens_with_context,
|
||||
output_scale,
|
||||
)
|
||||
end.record()
|
||||
torch.accelerator.synchronize()
|
||||
total_time_torch_kernel += start.elapsed_time(end)
|
||||
|
||||
avg_time_torch_kernel = total_time_torch_kernel / repeat_times
|
||||
|
||||
# 1. Run the Triton kernel
|
||||
output_ref_triton = output.clone()
|
||||
output_lse_ref_triton = output_lse.clone()
|
||||
|
||||
total_time_triton_kernel = 0
|
||||
start = torch.Event(enable_timing=True)
|
||||
end = torch.Event(enable_timing=True)
|
||||
|
||||
for _ in range(warmup_times):
|
||||
merge_attn_states_triton(
|
||||
output_ref_triton,
|
||||
prefix_output,
|
||||
prefix_lse,
|
||||
suffix_output,
|
||||
suffix_lse,
|
||||
output_lse_ref_triton,
|
||||
prefill_tokens_with_context,
|
||||
output_scale,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
for _ in range(repeat_times):
|
||||
start.record()
|
||||
merge_attn_states_triton(
|
||||
output_ref_triton,
|
||||
prefix_output,
|
||||
prefix_lse,
|
||||
suffix_output,
|
||||
suffix_lse,
|
||||
output_lse_ref_triton,
|
||||
prefill_tokens_with_context,
|
||||
output_scale,
|
||||
)
|
||||
end.record()
|
||||
torch.accelerator.synchronize()
|
||||
total_time_triton_kernel += start.elapsed_time(end)
|
||||
|
||||
avg_time_triton_kernel = total_time_triton_kernel / repeat_times
|
||||
|
||||
# 2. Run the CUDA kernel
|
||||
total_time_cuda_kernel = 0
|
||||
output_cuda = output.clone()
|
||||
output_lse_cuda = output_lse.clone()
|
||||
|
||||
for _ in range(warmup_times):
|
||||
merge_attn_states_cuda(
|
||||
output_cuda,
|
||||
prefix_output,
|
||||
prefix_lse,
|
||||
suffix_output,
|
||||
suffix_lse,
|
||||
output_lse_cuda,
|
||||
prefill_tokens_with_context,
|
||||
output_scale,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
for _ in range(repeat_times):
|
||||
start.record()
|
||||
merge_attn_states_cuda(
|
||||
output_cuda,
|
||||
prefix_output,
|
||||
prefix_lse,
|
||||
suffix_output,
|
||||
suffix_lse,
|
||||
output_lse_cuda,
|
||||
prefill_tokens_with_context,
|
||||
output_scale,
|
||||
)
|
||||
end.record()
|
||||
torch.accelerator.synchronize()
|
||||
total_time_cuda_kernel += start.elapsed_time(end)
|
||||
|
||||
avg_time_cuda_kernel = total_time_cuda_kernel / repeat_times
|
||||
|
||||
# 3. Performance compare
|
||||
performance_improved = avg_time_triton_kernel / avg_time_cuda_kernel
|
||||
print(f" Torch time: {avg_time_torch_kernel:.6f}ms")
|
||||
print(f"Triton time: {avg_time_triton_kernel:.6f}ms")
|
||||
print(
|
||||
f" CUDA time: {avg_time_cuda_kernel:.6f}ms, "
|
||||
f"Performance: {performance_improved:.5f}x"
|
||||
)
|
||||
print("-" * 100)
|
||||
|
||||
# 4. Correctness compare
|
||||
# Liger Kernel: Efficient Triton Kernels for LLM Training
|
||||
# https://arxiv.org/pdf/2410.10989, 3.3 Correctness
|
||||
# use rtol = 1e-2 for bfloat16.
|
||||
if use_fp8:
|
||||
# Compare in dequantized space (multiply back by scale) so that
|
||||
# absolute differences reflect real precision, not amplified FP8
|
||||
# quantization steps.
|
||||
atol, rtol = 1e-1, 1e-1
|
||||
assert output_scale is not None
|
||||
scale = output_scale.item()
|
||||
elif output_dtype == torch.bfloat16:
|
||||
atol, rtol = 1e-3, 1e-2
|
||||
scale = 1.0
|
||||
else:
|
||||
atol, rtol = 1e-3, 1e-3
|
||||
scale = 1.0
|
||||
|
||||
def diff(a: torch.Tensor, b: torch.Tensor):
|
||||
max_diff = torch.max(torch.abs(a.float() - b.float()))
|
||||
return max_diff
|
||||
|
||||
# Use Triton output as reference because we want to replace
|
||||
# the Triton kernel with custom CUDA kernel for merge attn
|
||||
# states operation.
|
||||
output_ref = output_ref_triton
|
||||
output_lse_ref = output_lse_ref_triton
|
||||
torch.testing.assert_close(
|
||||
output_cuda.float() * scale,
|
||||
output_ref.float() * scale,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
print(
|
||||
"Output all match, max abs diff (dequantized):"
|
||||
if use_fp8
|
||||
else "Output all match, max abs diff:"
|
||||
)
|
||||
_diff = diff(output_ref.float() * scale, output_torch.float() * scale)
|
||||
print(f"(Triton vs Torch) : {_diff}")
|
||||
_diff = diff(output_torch.float() * scale, output_cuda.float() * scale)
|
||||
print(f" (CUDA vs Torch) : {_diff}")
|
||||
_diff = diff(output_ref.float() * scale, output_cuda.float() * scale)
|
||||
print(f" (CUDA vs Triton): {_diff}")
|
||||
print("-" * 100)
|
||||
|
||||
torch.testing.assert_close(
|
||||
output_lse_cuda.float(), output_lse_ref.float(), atol=atol, rtol=rtol
|
||||
)
|
||||
print("Output LSE all match, max abs diff:")
|
||||
print(f"(Triton vs Torch) : {diff(output_lse_torch, output_lse_ref)}")
|
||||
print(f" (CUDA vs Torch) : {diff(output_lse_torch, output_lse_cuda)}")
|
||||
print(f" (CUDA vs Triton): {diff(output_lse_ref, output_lse_cuda)}")
|
||||
print("-" * 100)
|
||||
|
||||
print(
|
||||
"All output values test passed! All inf values "
|
||||
"are correctly replaced with -inf."
|
||||
)
|
||||
print("-" * 100)
|
||||
|
||||
device = current_platform.get_device_name()
|
||||
all_case_info.append(
|
||||
(
|
||||
NUM_TOKENS,
|
||||
NUM_HEADS,
|
||||
HEAD_SIZE,
|
||||
output_dtype,
|
||||
device,
|
||||
avg_time_torch_kernel,
|
||||
avg_time_triton_kernel,
|
||||
avg_time_cuda_kernel,
|
||||
performance_improved,
|
||||
)
|
||||
)
|
||||
if len(all_case_info) == (
|
||||
len(NUM_BATCH_TOKENS) * len(HEAD_SIZES) * len(NUM_QUERY_HEADS) * len(DTYPES)
|
||||
):
|
||||
generate_markdown_table()
|
||||
@@ -0,0 +1,348 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Test:
|
||||
|
||||
* Tests for MMEncoderAttention layer
|
||||
"""
|
||||
|
||||
import itertools
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config import get_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.model_executor.layers.attention import MMEncoderAttention
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.platforms.cpu import CpuPlatform
|
||||
from vllm.platforms.cuda import CudaPlatform
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
from vllm.platforms.rocm import RocmPlatform
|
||||
from vllm.utils.torch_utils import set_default_torch_dtype, set_random_seed
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.v1.attention.selector import _cached_get_attn_backend
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_cache():
|
||||
"""Clear lru cache to ensure each test case runs without caching."""
|
||||
_cached_get_attn_backend.cache_clear()
|
||||
|
||||
|
||||
devices = ["cpu"]
|
||||
if current_platform.is_cuda():
|
||||
devices.append("cuda")
|
||||
if current_platform.is_rocm():
|
||||
devices.append("hip")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", devices)
|
||||
def test_mha_attn_platform(default_vllm_config, device: str):
|
||||
"""
|
||||
Test the attention selector between different platform and device.
|
||||
"""
|
||||
torch.set_default_dtype(torch.float16)
|
||||
|
||||
if device == "cpu":
|
||||
with (
|
||||
patch("vllm.model_executor.models.vision.current_platform", CpuPlatform()),
|
||||
):
|
||||
attn = MMEncoderAttention(16, 64, scale=1)
|
||||
assert attn.attn_backend == AttentionBackendEnum.TORCH_SDPA
|
||||
elif device == "hip":
|
||||
with (
|
||||
patch("vllm.model_executor.models.vision.current_platform", RocmPlatform()),
|
||||
):
|
||||
attn = MMEncoderAttention(16, 64, scale=1)
|
||||
assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
|
||||
else:
|
||||
# Test CUDA with head_size=64 (divisible by 32)
|
||||
# - should use vLLM's FlashAttention
|
||||
with (
|
||||
patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
|
||||
):
|
||||
attn = MMEncoderAttention(16, 64, scale=1)
|
||||
assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
|
||||
|
||||
# Test CUDA with head_size=72 (not divisible by 32)
|
||||
# - should use vLLM's FlashAttention
|
||||
with (
|
||||
patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
|
||||
):
|
||||
attn = MMEncoderAttention(16, 72, scale=1)
|
||||
assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
|
||||
|
||||
# Test CUDA with head_size=72 (not divisible by 32)
|
||||
# - should use vLLM's FlashAttention
|
||||
with (
|
||||
patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
|
||||
set_default_torch_dtype(torch.float32),
|
||||
):
|
||||
attn = MMEncoderAttention(16, 72, scale=1)
|
||||
assert attn.attn_backend == AttentionBackendEnum.TRITON_ATTN
|
||||
|
||||
# Test Turing (pre-Ampere, sm_75): FlashAttention requires sm>=80,
|
||||
# and Triton no longer supports MMA on Turing, so we expect that
|
||||
# TORCH_SDPA is used for MMEncoderAttention.
|
||||
with (
|
||||
patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
|
||||
patch.object(
|
||||
CudaPlatform,
|
||||
"get_device_capability",
|
||||
return_value=DeviceCapability(major=7, minor=5),
|
||||
),
|
||||
):
|
||||
attn = MMEncoderAttention(16, 64, scale=1)
|
||||
assert attn.attn_backend == AttentionBackendEnum.TORCH_SDPA
|
||||
|
||||
|
||||
def ref_attention(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
scale: float,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Native implementation of scaled dot product attention without mask:
|
||||
- query, key, value: [batch_size, seq_len, num_heads, head_size]
|
||||
- attn_mask: [batch_size, seq_len, seq_len]
|
||||
"""
|
||||
query, key, value = (x.transpose(1, 2) for x in (query, key, value))
|
||||
attn_weights = scale * torch.matmul(query, key.transpose(2, 3))
|
||||
attn_weights = torch.softmax(attn_weights, dim=-1).to(value.dtype)
|
||||
out = torch.matmul(attn_weights, value).transpose(1, 2)
|
||||
return out
|
||||
|
||||
|
||||
BATCH_SIZES = [1, 16]
|
||||
SEQ_LENS = [1]
|
||||
VAR_SEQ_LENS = [
|
||||
[2, 2],
|
||||
[2, 3, 4],
|
||||
]
|
||||
NUM_HEADS = [1, 16]
|
||||
NUM_KV_HEADS = [1]
|
||||
HEAD_SIZES = [64, 80]
|
||||
# flshattF and tritonflashattF supported: {torch.float16, torch.bfloat16}
|
||||
DTYPES = (
|
||||
[torch.half, torch.bfloat16, torch.float]
|
||||
if not current_platform.is_rocm()
|
||||
else [torch.half, torch.bfloat16]
|
||||
)
|
||||
CUDA_DEVICES = ["cuda"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("num_kv_heads", NUM_KV_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_mha_attn_forward(
|
||||
default_vllm_config,
|
||||
batch_size: int,
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
):
|
||||
set_random_seed(0)
|
||||
torch.set_default_device(device)
|
||||
torch.set_default_dtype(dtype)
|
||||
|
||||
q = torch.randn(batch_size, seq_len, num_heads * head_size)
|
||||
k = torch.randn(batch_size, seq_len, num_kv_heads * head_size)
|
||||
v = torch.randn(batch_size, seq_len, num_kv_heads * head_size)
|
||||
scale = 1.0 / head_size**0.5
|
||||
attn = MMEncoderAttention(
|
||||
num_heads, head_size, scale=scale, num_kv_heads=num_kv_heads
|
||||
)
|
||||
output = attn(q, k, v)
|
||||
|
||||
assert num_heads % num_kv_heads == 0
|
||||
num_queries_per_kv = num_heads // num_kv_heads
|
||||
q = q.reshape(batch_size, seq_len, num_heads, head_size)
|
||||
k = k.reshape(batch_size, seq_len, num_kv_heads, head_size)
|
||||
v = v.reshape(batch_size, seq_len, num_kv_heads, head_size)
|
||||
if num_queries_per_kv > 1:
|
||||
k = torch.repeat_interleave(k, num_queries_per_kv, dim=2)
|
||||
v = torch.repeat_interleave(v, num_queries_per_kv, dim=2)
|
||||
|
||||
ref_output = ref_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
scale=scale,
|
||||
).reshape(batch_size, seq_len, num_heads * head_size)
|
||||
tol_kwargs = (
|
||||
dict(rtol=1e-3, atol=1e-3)
|
||||
if attn.attn_backend == AttentionBackendEnum.TRITON_ATTN
|
||||
else {}
|
||||
)
|
||||
torch.testing.assert_close(output, ref_output, **tol_kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("var_seq_len", VAR_SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("num_kv_heads", NUM_KV_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_mha_attn_varlen_forward(
|
||||
default_vllm_config,
|
||||
var_seq_len: list[int],
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
):
|
||||
set_random_seed(0)
|
||||
torch.set_default_device(device)
|
||||
torch.set_default_dtype(dtype)
|
||||
|
||||
q = torch.randn(1, sum(var_seq_len), num_heads, head_size)
|
||||
k = torch.randn(1, sum(var_seq_len), num_kv_heads, head_size)
|
||||
v = torch.randn(1, sum(var_seq_len), num_kv_heads, head_size)
|
||||
cu_seqlens = torch.tensor(
|
||||
[0] + list(itertools.accumulate(var_seq_len)), dtype=torch.int32
|
||||
)
|
||||
scale = 1.0 / head_size**0.5
|
||||
attn = MMEncoderAttention(
|
||||
num_heads, head_size, scale=scale, num_kv_heads=num_kv_heads
|
||||
)
|
||||
output = attn(
|
||||
q, k, v, cu_seqlens=cu_seqlens, max_seqlen=torch.tensor(max(var_seq_len))
|
||||
)
|
||||
|
||||
assert num_heads % num_kv_heads == 0
|
||||
num_queries_per_kv = num_heads // num_kv_heads
|
||||
if num_queries_per_kv > 1:
|
||||
k = torch.repeat_interleave(k, num_queries_per_kv, dim=2)
|
||||
v = torch.repeat_interleave(v, num_queries_per_kv, dim=2)
|
||||
|
||||
ref_output = []
|
||||
for q_i, k_i, v_i in zip(
|
||||
torch.split(q, var_seq_len, dim=1),
|
||||
torch.split(k, var_seq_len, dim=1),
|
||||
torch.split(v, var_seq_len, dim=1),
|
||||
):
|
||||
output_i = ref_attention(
|
||||
q_i,
|
||||
k_i,
|
||||
v_i,
|
||||
scale=scale,
|
||||
)
|
||||
ref_output.append(output_i)
|
||||
ref_output = torch.cat(ref_output, dim=1)
|
||||
torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("var_seq_len", VAR_SEQ_LENS)
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[torch.bfloat16, torch.half],
|
||||
)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_mha_attn_varlen_forward_flashinfer(
|
||||
default_vllm_config,
|
||||
var_seq_len: list[int],
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
):
|
||||
"""Test MMEncoderAttention varlen forward with FLASHINFER backend (head_size=72).
|
||||
|
||||
Exercises the path that uses --mm-encoder-attn-backend=FLASHINFER with
|
||||
recomputed cu_seqlens, max_seqlen, and sequence_lengths as in qwen3_vl
|
||||
vision encoder.
|
||||
"""
|
||||
pytest.importorskip("flashinfer")
|
||||
|
||||
num_heads = 16
|
||||
head_size = 72
|
||||
set_random_seed(0)
|
||||
torch.set_default_device(device)
|
||||
torch.set_default_dtype(dtype)
|
||||
|
||||
# Override vllm config so get_vit_attn_backend returns FLASHINFER (simulates
|
||||
# --mm-encoder-attn-backend=FLASHINFER).
|
||||
vllm_config = get_current_vllm_config()
|
||||
old_model_config = getattr(vllm_config, "model_config", None)
|
||||
minimal_model_config = type(
|
||||
"MinimalModelConfig",
|
||||
(),
|
||||
{
|
||||
"multimodal_config": MultiModalConfig(
|
||||
mm_encoder_attn_backend=AttentionBackendEnum.FLASHINFER
|
||||
),
|
||||
},
|
||||
)()
|
||||
vllm_config.model_config = minimal_model_config
|
||||
try:
|
||||
total_len = sum(var_seq_len)
|
||||
# Stride of second dim = 3 * num_heads * head_size (same as qwen2_5_vl
|
||||
# after qkv rearrange and unbind: qkv shape (b, s, 3, head, head_dim)).
|
||||
qkv = torch.randn(1, total_len, 3, num_heads, head_size)
|
||||
q, k, v = qkv.unbind(dim=2)
|
||||
|
||||
cu_seqlens_np = np.array(
|
||||
[0] + list(itertools.accumulate(var_seq_len)), dtype=np.int32
|
||||
)
|
||||
hidden_size = num_heads * head_size
|
||||
tp_size = 1
|
||||
|
||||
sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_seqlens_np,
|
||||
device,
|
||||
)
|
||||
|
||||
max_seqlen_val = MMEncoderAttention.compute_max_seqlen(
|
||||
AttentionBackendEnum.FLASHINFER, cu_seqlens_np
|
||||
)
|
||||
max_seqlen = torch.tensor(max_seqlen_val, device=device, dtype=torch.int32)
|
||||
|
||||
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_seqlens_np,
|
||||
hidden_size,
|
||||
tp_size,
|
||||
device,
|
||||
)
|
||||
|
||||
scale = 1.0 / head_size**0.5
|
||||
attn = MMEncoderAttention(
|
||||
num_heads,
|
||||
head_size,
|
||||
scale=scale,
|
||||
num_kv_heads=num_heads,
|
||||
)
|
||||
assert attn.attn_backend == AttentionBackendEnum.FLASHINFER
|
||||
|
||||
output = attn(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
sequence_lengths=sequence_lengths,
|
||||
)
|
||||
|
||||
ref_output = []
|
||||
for q_i, k_i, v_i in zip(
|
||||
torch.split(q, var_seq_len, dim=1),
|
||||
torch.split(k, var_seq_len, dim=1),
|
||||
torch.split(v, var_seq_len, dim=1),
|
||||
):
|
||||
output_i = ref_attention(q_i, k_i, v_i, scale=scale)
|
||||
ref_output.append(output_i)
|
||||
ref_output = torch.cat(ref_output, dim=1)
|
||||
torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
|
||||
finally:
|
||||
vllm_config.model_config = old_model_config
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,318 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for per-request causal/non-causal attention (mixed batches).
|
||||
|
||||
Validates that both triton and flash-attention backends correctly handle
|
||||
batches where some sequences use causal masking and others use non-causal
|
||||
(bidirectional) masking — needed by DiffusionGemma.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
# Mixed causal/non-causal attention is only validated on a subset of GPUs:
|
||||
# the Triton path on Hopper (SM90) and B200 (SM100); the FA4 path on Hopper
|
||||
# (SM90) only.
|
||||
_device_capability = current_platform.get_device_capability()
|
||||
_major = _device_capability.major if _device_capability is not None else None
|
||||
|
||||
NUM_HEADS = [(4, 4), (8, 2)]
|
||||
HEAD_SIZES = [128]
|
||||
BLOCK_SIZES = [16]
|
||||
DTYPES = [torch.bfloat16]
|
||||
|
||||
|
||||
def ref_paged_attn(
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
query_lens: list[int],
|
||||
kv_lens: list[int],
|
||||
block_tables: torch.Tensor,
|
||||
scale: float,
|
||||
per_seq_causal: list[bool],
|
||||
sliding_window: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
num_seqs = len(query_lens)
|
||||
block_tables_np = block_tables.cpu().numpy()
|
||||
_, block_size, num_kv_heads, head_size = key_cache.shape
|
||||
|
||||
outputs: list[torch.Tensor] = []
|
||||
start_idx = 0
|
||||
for i in range(num_seqs):
|
||||
query_len = query_lens[i]
|
||||
kv_len = kv_lens[i]
|
||||
q = query[start_idx : start_idx + query_len]
|
||||
q = q * scale
|
||||
|
||||
num_kv_blocks = (kv_len + block_size - 1) // block_size
|
||||
block_indices = block_tables_np[i, :num_kv_blocks]
|
||||
k = key_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
|
||||
v = value_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
|
||||
|
||||
if q.shape[1] != k.shape[1]:
|
||||
k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
|
||||
v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
|
||||
|
||||
attn = torch.einsum("qhd,khd->hqk", q, k).float()
|
||||
|
||||
if per_seq_causal[i]:
|
||||
mask = torch.triu(
|
||||
torch.ones(query_len, kv_len, device=attn.device),
|
||||
diagonal=kv_len - query_len + 1,
|
||||
).bool()
|
||||
else:
|
||||
mask = torch.zeros(query_len, kv_len, device=attn.device).bool()
|
||||
|
||||
if sliding_window is not None:
|
||||
sw_mask = (
|
||||
torch.triu(
|
||||
torch.ones(query_len, kv_len, device=attn.device),
|
||||
diagonal=kv_len - (query_len + sliding_window) + 1,
|
||||
)
|
||||
.bool()
|
||||
.logical_not()
|
||||
)
|
||||
mask |= sw_mask
|
||||
|
||||
attn.masked_fill_(mask, float("-inf"))
|
||||
attn = torch.softmax(attn, dim=-1).to(v.dtype)
|
||||
out = torch.einsum("hqk,khd->qhd", attn, v)
|
||||
outputs.append(out)
|
||||
start_idx += query_len
|
||||
|
||||
return torch.cat(outputs, dim=0)
|
||||
|
||||
|
||||
# ---- Triton backend test ----
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
_major not in (9, 10),
|
||||
reason="Triton mixed causal attention requires Hopper (SM90) or B200 (SM100).",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens",
|
||||
[[(1, 128), (5, 64), (1, 256)]],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"per_seq_causal",
|
||||
[[True, False, True], [False, True, False], [True, True, False]],
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_triton_mixed_causal(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
per_seq_causal: list[bool],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
):
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip("Triton attention requires CUDA")
|
||||
|
||||
from vllm.v1.attention.ops.triton_unified_attention import unified_attention
|
||||
|
||||
set_random_seed(42)
|
||||
device = "cuda"
|
||||
|
||||
num_query_heads, num_kv_heads = num_heads
|
||||
assert len(seq_lens) == len(per_seq_causal)
|
||||
|
||||
query_lens = [s[0] for s in seq_lens]
|
||||
kv_lens = [s[1] for s in seq_lens]
|
||||
num_seqs = len(seq_lens)
|
||||
|
||||
num_query_tokens = sum(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
max_num_blocks = (max_kv_len + block_size - 1) // block_size
|
||||
num_blocks = max_num_blocks * num_seqs + 10
|
||||
|
||||
scale = head_size**-0.5
|
||||
query = torch.randn(
|
||||
num_query_tokens, num_query_heads, head_size, dtype=dtype, device=device
|
||||
)
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
|
||||
)
|
||||
value_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
|
||||
)
|
||||
|
||||
block_tables_list = []
|
||||
for i in range(num_seqs):
|
||||
n_blocks = (kv_lens[i] + block_size - 1) // block_size
|
||||
blocks = list(range(i * max_num_blocks, i * max_num_blocks + n_blocks))
|
||||
blocks += [0] * (max_num_blocks - n_blocks)
|
||||
block_tables_list.append(blocks)
|
||||
block_tables = torch.tensor(block_tables_list, dtype=torch.int32, device=device)
|
||||
|
||||
cu_seqlens_q = torch.zeros(num_seqs + 1, dtype=torch.int32, device=device)
|
||||
for i, ql in enumerate(query_lens):
|
||||
cu_seqlens_q[i + 1] = cu_seqlens_q[i] + ql
|
||||
|
||||
seqused_k = torch.tensor(kv_lens, dtype=torch.int32, device=device)
|
||||
max_seqlen_q = max(query_lens)
|
||||
max_seqlen_k = max(kv_lens)
|
||||
|
||||
causal_tensor = torch.tensor(per_seq_causal, dtype=torch.bool, device=device)
|
||||
|
||||
output = torch.empty_like(query)
|
||||
unified_attention(
|
||||
q=query,
|
||||
k=key_cache,
|
||||
v=value_cache,
|
||||
out=output,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
seqused_k=seqused_k,
|
||||
max_seqlen_k=max_seqlen_k,
|
||||
softmax_scale=scale,
|
||||
causal=causal_tensor,
|
||||
window_size=(-1, -1),
|
||||
block_table=block_tables,
|
||||
softcap=0.0,
|
||||
q_descale=None,
|
||||
k_descale=1.0,
|
||||
v_descale=1.0,
|
||||
)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query,
|
||||
key_cache,
|
||||
value_cache,
|
||||
query_lens,
|
||||
kv_lens,
|
||||
block_tables,
|
||||
scale,
|
||||
per_seq_causal,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
# ---- Flash Attention 4 backend test (native per_seq_causal) ----
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
_major != 9,
|
||||
reason="FA4 mixed causal attention requires Hopper (SM90).",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens",
|
||||
[[(1, 128), (5, 64), (1, 256)]],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"per_seq_causal",
|
||||
[[True, False, True], [False, True, False]],
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_flash_attn4_mixed_causal(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
per_seq_causal: list[bool],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
):
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip("Flash attention requires CUDA")
|
||||
|
||||
try:
|
||||
from vllm.vllm_flash_attn import (
|
||||
fa_version_unsupported_reason,
|
||||
flash_attn_varlen_func,
|
||||
is_fa_version_supported,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("vllm_flash_attn not available")
|
||||
|
||||
if not is_fa_version_supported(4):
|
||||
reason = fa_version_unsupported_reason(4)
|
||||
pytest.skip(f"FA4 not supported: {reason}")
|
||||
|
||||
set_random_seed(42)
|
||||
device = "cuda"
|
||||
|
||||
num_query_heads, num_kv_heads = num_heads
|
||||
assert len(seq_lens) == len(per_seq_causal)
|
||||
|
||||
query_lens = [s[0] for s in seq_lens]
|
||||
kv_lens = [s[1] for s in seq_lens]
|
||||
num_seqs = len(seq_lens)
|
||||
|
||||
num_query_tokens = sum(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
max_num_blocks = (max_kv_len + block_size - 1) // block_size
|
||||
num_blocks = max_num_blocks * num_seqs + 10
|
||||
|
||||
scale = head_size**-0.5
|
||||
query = torch.randn(
|
||||
num_query_tokens, num_query_heads, head_size, dtype=dtype, device=device
|
||||
)
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
|
||||
)
|
||||
value_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
|
||||
)
|
||||
|
||||
block_tables_list = []
|
||||
for i in range(num_seqs):
|
||||
n_blocks = (kv_lens[i] + block_size - 1) // block_size
|
||||
blocks = list(range(i * max_num_blocks, i * max_num_blocks + n_blocks))
|
||||
blocks += [0] * (max_num_blocks - n_blocks)
|
||||
block_tables_list.append(blocks)
|
||||
block_tables = torch.tensor(block_tables_list, dtype=torch.int32, device=device)
|
||||
|
||||
cu_seqlens_q = torch.zeros(num_seqs + 1, dtype=torch.int32, device=device)
|
||||
for i, ql in enumerate(query_lens):
|
||||
cu_seqlens_q[i + 1] = cu_seqlens_q[i] + ql
|
||||
|
||||
seqused_k = torch.tensor(kv_lens, dtype=torch.int32, device=device)
|
||||
per_seq_causal_tensor = torch.tensor(
|
||||
per_seq_causal, dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query,
|
||||
key_cache,
|
||||
value_cache,
|
||||
query_lens,
|
||||
kv_lens,
|
||||
block_tables,
|
||||
scale,
|
||||
per_seq_causal,
|
||||
)
|
||||
|
||||
output = torch.empty_like(query)
|
||||
flash_attn_varlen_func(
|
||||
q=query,
|
||||
k=key_cache,
|
||||
v=value_cache,
|
||||
out=output,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
max_seqlen_q=max(query_lens),
|
||||
seqused_k=seqused_k,
|
||||
max_seqlen_k=max(kv_lens),
|
||||
softmax_scale=scale,
|
||||
# The kernel must be compiled causal for `dynamic_causal` to take effect.
|
||||
causal=True,
|
||||
block_table=block_tables,
|
||||
softcap=0.0,
|
||||
dynamic_causal=per_seq_causal_tensor,
|
||||
fa_version=4,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
|
||||
@@ -0,0 +1,566 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Bit-exact kernel equivalence for MLA decode/write kernels on the
|
||||
cross-layer (block-major) KV cache layout.
|
||||
|
||||
The cross-layer layout carves each layer's per-block page out of a single
|
||||
unified slot, so the per-layer view has an inflated ``stride(0)`` (the full
|
||||
unified slot) and a non-zero storage offset. These tests confirm the MLA
|
||||
kernels behind the backends that opt in to the layout (FlashMLA dense,
|
||||
FlashInfer MLA dense, FlashMLA fp8 sparse, plus the ``concat_and_cache_mla``
|
||||
write) honor that strided view bit-identically to a contiguous per-layer
|
||||
cache, and that writes do not bleed into neighbouring layers' segments.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available(), reason="MLA cache kernels require CUDA"
|
||||
)
|
||||
|
||||
|
||||
def test_concat_and_cache_mla_into_unified_slot_view():
|
||||
"""concat_and_cache_mla must write correctly into a per-layer view whose
|
||||
block stride is the full unified slot (block-major), with zero bleed into
|
||||
the other layers' segments of the same slot."""
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
kv_lora_rank = 512
|
||||
pe = 64
|
||||
entry = kv_lora_rank + pe
|
||||
page = 64
|
||||
num_blocks = 32
|
||||
ntok = 200
|
||||
|
||||
kv_c = torch.randn(ntok, kv_lora_rank, device=dev, dtype=torch.bfloat16)
|
||||
k_pe = torch.randn(ntok, pe, device=dev, dtype=torch.bfloat16)
|
||||
slot = torch.randperm(num_blocks * page, device=dev, dtype=torch.int64)[:ntok]
|
||||
scale = torch.tensor(1.0, device=dev)
|
||||
|
||||
def write(cache):
|
||||
ops.concat_and_cache_mla(kv_c, k_pe, cache, slot, "auto", scale)
|
||||
|
||||
# Contiguous per-layer reference: (num_blocks, page, entry).
|
||||
ref = torch.zeros(num_blocks, page, entry, device=dev, dtype=torch.bfloat16)
|
||||
write(ref)
|
||||
|
||||
# Unified slot holding three layer pages per block. Carve the middle
|
||||
# layer's view (non-zero offset, block stride == full unified slot).
|
||||
layer_page_elems = page * entry
|
||||
n_layers = 3
|
||||
unified_slot_elems = n_layers * layer_page_elems
|
||||
big = torch.zeros(num_blocks, unified_slot_elems, device=dev, dtype=torch.bfloat16)
|
||||
flat = big.view(-1)
|
||||
offset = layer_page_elems # middle layer
|
||||
view = torch.as_strided(
|
||||
flat,
|
||||
size=(num_blocks, page, entry),
|
||||
stride=(unified_slot_elems, entry, 1),
|
||||
storage_offset=offset,
|
||||
)
|
||||
assert not view.is_contiguous()
|
||||
assert view.stride(0) == unified_slot_elems
|
||||
write(view)
|
||||
|
||||
# Bit-exact equivalence and zero bleed into the neighbour segments.
|
||||
max_diff = (ref.float() - view.float()).abs().max().item()
|
||||
assert max_diff == 0.0, f"max|Δ| = {max_diff}"
|
||||
|
||||
neighbour_lo = torch.as_strided(
|
||||
flat, (num_blocks, layer_page_elems), (unified_slot_elems, 1), 0
|
||||
)
|
||||
neighbour_hi = torch.as_strided(
|
||||
flat,
|
||||
(num_blocks, layer_page_elems),
|
||||
(unified_slot_elems, 1),
|
||||
2 * layer_page_elems,
|
||||
)
|
||||
assert neighbour_lo.abs().max().item() == 0.0
|
||||
assert neighbour_hi.abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashmla_dense_decode_unified_slot_view():
|
||||
"""FlashMLA dense decode (FLASHMLA backend, e.g. Kimi-K2-style dense MLA
|
||||
on Hopper) must read a unified-slot block-major view bit-identically to a
|
||||
contiguous per-layer cache."""
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_dense_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
dt = torch.bfloat16
|
||||
head_dim = 576
|
||||
hdv = 512
|
||||
h_q = 128
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=dt) * 0.1
|
||||
kv_data = torch.randn(num_blocks, page, 1, head_dim, device=dev, dtype=dt) * 0.1
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
cache_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = (
|
||||
torch.randn(num_blocks, n_layers, page, 1, head_dim, device=dev, dtype=dt) * 0.1
|
||||
)
|
||||
unified[:, layer].copy_(kv_data)
|
||||
cache_view = unified[:, layer]
|
||||
assert not cache_view.is_contiguous()
|
||||
assert cache_view.stride(0) == n_layers * page * 1 * head_dim
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
cache_seqlens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
|
||||
def run(kc):
|
||||
meta, num_splits = fm.get_mla_metadata()
|
||||
out, _ = fm.flash_mla_with_kvcache(
|
||||
q=q,
|
||||
k_cache=kc,
|
||||
block_table=block_table,
|
||||
cache_seqlens=cache_seqlens,
|
||||
head_dim_v=hdv,
|
||||
tile_scheduler_metadata=meta,
|
||||
num_splits=num_splits,
|
||||
softmax_scale=head_dim**-0.5,
|
||||
causal=True,
|
||||
)
|
||||
return out.clone().float()
|
||||
|
||||
out_ref = run(cache_contiguous)
|
||||
out_view = run(cache_view)
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert out_ref.abs().max().item() > 0.0
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashinfer_mla_dense_decode_unified_slot_view():
|
||||
"""FlashInfer MLA dense decode must read a unified-slot block-major view
|
||||
(inflated stride(0), non-zero storage offset) bit-identically to a
|
||||
contiguous per-layer cache."""
|
||||
try:
|
||||
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
|
||||
except ImportError:
|
||||
pytest.skip("flashinfer is not available")
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_device_capability_family(100):
|
||||
pytest.skip("FlashInfer trtllm-gen MLA requires sm100")
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
dt = torch.bfloat16
|
||||
kv_lora_rank = 512
|
||||
qk_rope_head_dim = 64
|
||||
qk_nope_head_dim = 128
|
||||
head_dim = kv_lora_rank + qk_rope_head_dim # 576
|
||||
num_qo_heads = 128
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3 # >1 so the per-layer view's block stride is inflated.
|
||||
layer = 1
|
||||
|
||||
q = torch.randn(bs, 1, num_qo_heads, head_dim, device=dev, dtype=dt)
|
||||
kv_data = torch.randn(num_blocks, 1, page, head_dim, device=dev, dtype=dt)
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
kv_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: block b of every layer packed together; view one layer
|
||||
# -> stride(0) is n_layers x larger and storage offset is non-zero.
|
||||
unified = torch.randn(num_blocks, n_layers, 1, page, head_dim, device=dev, dtype=dt)
|
||||
unified[:, layer].copy_(kv_data)
|
||||
kv_view = unified[:, layer]
|
||||
assert not kv_view.is_contiguous()
|
||||
assert kv_view.stride(0) == n_layers * 1 * page * head_dim
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_tables = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=dev)
|
||||
scale = head_dim**-0.5
|
||||
|
||||
def run(kv):
|
||||
return trtllm_batch_decode_with_kv_cache_mla(
|
||||
query=q,
|
||||
kv_cache=kv,
|
||||
workspace_buffer=ws,
|
||||
qk_nope_head_dim=qk_nope_head_dim,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
block_tables=block_tables,
|
||||
seq_lens=seq_lens,
|
||||
max_seq_len=int(seq_lens.max().item()),
|
||||
bmm1_scale=scale,
|
||||
bmm2_scale=1.0,
|
||||
).clone()
|
||||
|
||||
out_ref = run(kv_contiguous).float()
|
||||
out_view = run(kv_view).float()
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashmla_fp8_sparse_decode_unified_slot_view():
|
||||
"""FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a
|
||||
unified-slot block-major view bit-identically to a contiguous fp8_ds_mla
|
||||
cache, with finite nonzero output."""
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_sparse_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
entry = 656 # fp8_ds_mla bytes per token
|
||||
page = 64
|
||||
num_blocks = 32
|
||||
h_q = 128
|
||||
head_dim = 576
|
||||
hdv = 512
|
||||
batch = 2
|
||||
topk = 128
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
q = torch.randn(batch, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
|
||||
|
||||
# Structurally valid fp8 ds_mla payload: 512B fp8 + 16B f32 scales + 128B
|
||||
# bf16 rope (random bytes corrupt the scale region and yield NaNs).
|
||||
nope = (torch.randn(num_blocks, page, 1, 512, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
scales = torch.ones(num_blocks, page, 1, 4, device=dev, dtype=torch.float32)
|
||||
rope = (torch.randn(num_blocks, page, 1, 64, device=dev) * 0.1).to(torch.bfloat16)
|
||||
payload = torch.cat(
|
||||
[
|
||||
nope.view(torch.uint8).view(num_blocks, page, 1, 512),
|
||||
scales.view(torch.uint8).view(num_blocks, page, 1, 16),
|
||||
rope.view(torch.uint8).view(num_blocks, page, 1, 128),
|
||||
],
|
||||
dim=-1,
|
||||
).contiguous()
|
||||
assert payload.shape[-1] == entry and payload.dtype == torch.uint8
|
||||
|
||||
# (A) contiguous reference.
|
||||
cache_contiguous = payload.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = torch.randint(
|
||||
0, 256, (num_blocks, n_layers, page, 1, entry), device=dev, dtype=torch.uint8
|
||||
)
|
||||
unified[:, layer].copy_(payload)
|
||||
cache_view = unified[:, layer]
|
||||
assert not cache_view.is_contiguous()
|
||||
assert cache_view.stride(0) == n_layers * page * 1 * entry
|
||||
|
||||
# Sparse indices: each batch uses its own disjoint blocks.
|
||||
blocks_per_batch = num_blocks // batch
|
||||
idx = torch.full((batch, 1, topk), -1, device=dev, dtype=torch.int32)
|
||||
for b in range(batch):
|
||||
slots: list[int] = []
|
||||
for blk in range(b * blocks_per_batch, (b + 1) * blocks_per_batch):
|
||||
slots.extend(blk * page + off for off in range(page))
|
||||
slots_t = torch.tensor(slots[:topk], device=dev, dtype=torch.int32)
|
||||
idx[b, 0, : slots_t.numel()] = slots_t
|
||||
|
||||
def run(kc):
|
||||
meta, num_splits = fm.get_mla_metadata()
|
||||
out, _ = fm.flash_mla_with_kvcache(
|
||||
q=q,
|
||||
k_cache=kc,
|
||||
block_table=None,
|
||||
cache_seqlens=None,
|
||||
head_dim_v=hdv,
|
||||
tile_scheduler_metadata=meta,
|
||||
is_fp8_kvcache=True,
|
||||
indices=idx,
|
||||
softmax_scale=head_dim**-0.5,
|
||||
)
|
||||
return out.clone().float()
|
||||
|
||||
out_ref = run(cache_contiguous)
|
||||
out_view = run(cache_view)
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert out_ref.abs().max().item() > 0.0
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_indexer_k_quant_and_cache_into_unified_slot_view():
|
||||
"""indexer_k_quant_and_cache (DeepSeek V3.2/V4 DSA indexer K write) must
|
||||
write correctly into a per-layer view whose block stride is the full
|
||||
unified slot, with zero bleed into the other layers' segments."""
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
head_dim = 128
|
||||
quant_block_size = 128
|
||||
block_size = 64
|
||||
num_blocks = 16
|
||||
ntok = 100
|
||||
# Indexer cache layout per token: head_dim fp8 bytes followed by
|
||||
# head_dim * 4 / quant_block_size scale bytes.
|
||||
cache_stride = head_dim + head_dim * 4 // quant_block_size
|
||||
|
||||
k = torch.randn(ntok, head_dim, device=dev, dtype=torch.bfloat16)
|
||||
slot = torch.randperm(num_blocks * block_size, device=dev, dtype=torch.int64)[:ntok]
|
||||
|
||||
def write(cache):
|
||||
ops.indexer_k_quant_and_cache(k, cache, slot, quant_block_size, "ue8m0")
|
||||
|
||||
# Contiguous per-layer reference.
|
||||
ref = torch.zeros(
|
||||
num_blocks, block_size, cache_stride, device=dev, dtype=torch.uint8
|
||||
)
|
||||
write(ref)
|
||||
|
||||
# Unified slot holding three layer pages per block; carve the middle one.
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
unified = torch.zeros(
|
||||
num_blocks, n_layers, block_size, cache_stride, device=dev, dtype=torch.uint8
|
||||
)
|
||||
view = unified[:, layer]
|
||||
assert not view.is_contiguous()
|
||||
assert view.stride(0) == n_layers * block_size * cache_stride
|
||||
write(view)
|
||||
|
||||
assert torch.equal(ref, view.contiguous())
|
||||
# Zero bleed into the neighbour layers' segments.
|
||||
assert unified[:, 0].abs().max().item() == 0
|
||||
assert unified[:, 2].abs().max().item() == 0
|
||||
|
||||
|
||||
def test_flashattn_mla_dense_decode_unified_slot_view():
|
||||
"""FA3 decode (FLASH_ATTN_MLA backend) must read a unified-slot
|
||||
block-major view bit-identically to a contiguous per-layer cache."""
|
||||
try:
|
||||
from vllm.vllm_flash_attn import flash_attn_varlen_func
|
||||
except ImportError:
|
||||
pytest.skip("vllm_flash_attn is not available")
|
||||
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
|
||||
|
||||
if not flash_attn_supports_mla():
|
||||
pytest.skip("FA3 MLA requires a Hopper device")
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
dt = torch.bfloat16
|
||||
kv_lora_rank = 512
|
||||
rope_dim = 64
|
||||
entry = kv_lora_rank + rope_dim # 576
|
||||
h_q = 16
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
q_pe = torch.randn(bs, h_q, rope_dim, device=dev, dtype=dt) * 0.1
|
||||
q_nope = torch.randn(bs, h_q, kv_lora_rank, device=dev, dtype=dt) * 0.1
|
||||
kv_data = torch.randn(num_blocks, page, entry, device=dev, dtype=dt) * 0.1
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
cache_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = torch.randn(num_blocks, n_layers, page, entry, device=dev, dtype=dt) * 0.1
|
||||
unified[:, layer].copy_(kv_data)
|
||||
cache_view = unified[:, layer]
|
||||
assert not cache_view.is_contiguous()
|
||||
assert cache_view.stride(0) == n_layers * page * entry
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
cu_seqlens_q = torch.arange(bs + 1, device=dev, dtype=torch.int32)
|
||||
|
||||
def run(cache):
|
||||
kv_c_cache = cache[..., :kv_lora_rank]
|
||||
k_pe_cache = cache[..., kv_lora_rank:]
|
||||
out = flash_attn_varlen_func(
|
||||
q=q_pe,
|
||||
k=k_pe_cache.unsqueeze(-2), # Add head dim of 1
|
||||
v=kv_c_cache.unsqueeze(-2), # Add head dim of 1
|
||||
q_v=q_nope,
|
||||
max_seqlen_q=1,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
max_seqlen_k=int(seq_lens.max().item()),
|
||||
seqused_k=seq_lens,
|
||||
block_table=block_table,
|
||||
softmax_scale=entry**-0.5,
|
||||
causal=True,
|
||||
fa_version=3,
|
||||
)
|
||||
return out.clone().float()
|
||||
|
||||
out_ref = run(cache_contiguous)
|
||||
out_view = run(cache_view)
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert out_ref.abs().max().item() > 0.0
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashmla_dense_fp8_decode_unified_slot_view():
|
||||
"""FlashMLA dense fp8 decode (FLASHMLA backend with quantized KV cache)
|
||||
must read a unified-slot block-major view bit-identically to a contiguous
|
||||
per-layer fp8 cache."""
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_dense_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
head_dim = 576
|
||||
hdv = 512
|
||||
h_q = 128
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
|
||||
kv_data = (torch.randn(num_blocks, page, head_dim, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
cache_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = (torch.randn(num_blocks, n_layers, page, head_dim, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
unified[:, layer].copy_(kv_data)
|
||||
cache_view = unified[:, layer]
|
||||
assert not cache_view.is_contiguous()
|
||||
assert cache_view.stride(0) == n_layers * page * head_dim
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
cache_seqlens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
descale = torch.ones(1, device=dev, dtype=torch.float32)
|
||||
|
||||
def run(kc):
|
||||
tile_md, num_splits = fm.get_mla_metadata_dense_fp8(cache_seqlens, h_q, 1)
|
||||
out, _ = fm.flash_mla_with_kvcache_fp8(
|
||||
q=q,
|
||||
k_cache=kc.unsqueeze(-2), # Add head dim of 1
|
||||
block_table=block_table,
|
||||
cache_seqlens=cache_seqlens,
|
||||
head_dim_v=hdv,
|
||||
tile_scheduler_metadata=tile_md,
|
||||
num_splits=num_splits,
|
||||
softmax_scale=head_dim**-0.5,
|
||||
causal=True,
|
||||
descale_q=descale,
|
||||
descale_k=descale,
|
||||
)
|
||||
return out.clone().float()
|
||||
|
||||
out_ref = run(cache_contiguous)
|
||||
out_view = run(cache_view)
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert out_ref.abs().max().item() > 0.0
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashinfer_mla_dense_fp8_decode_unified_slot_view():
|
||||
"""FlashInfer MLA dense decode with an fp8 KV cache must read a
|
||||
unified-slot block-major view bit-identically to a contiguous per-layer
|
||||
cache."""
|
||||
try:
|
||||
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
|
||||
except ImportError:
|
||||
pytest.skip("flashinfer is not available")
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_device_capability_family(100):
|
||||
pytest.skip("FlashInfer trtllm-gen MLA requires sm100")
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
kv_lora_rank = 512
|
||||
qk_rope_head_dim = 64
|
||||
qk_nope_head_dim = 128
|
||||
head_dim = kv_lora_rank + qk_rope_head_dim # 576
|
||||
num_qo_heads = 128
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
# With a quantized KV cache the decode query is quantized to fp8 as well
|
||||
# (trtllm-gen has no bf16-query x fp8-cache decode kernel).
|
||||
q = (torch.randn(bs, 1, num_qo_heads, head_dim, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
kv_data = (torch.randn(num_blocks, 1, page, head_dim, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
kv_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = (
|
||||
torch.randn(num_blocks, n_layers, 1, page, head_dim, device=dev) * 0.1
|
||||
).to(torch.float8_e4m3fn)
|
||||
unified[:, layer].copy_(kv_data)
|
||||
kv_view = unified[:, layer]
|
||||
assert not kv_view.is_contiguous()
|
||||
assert kv_view.stride(0) == n_layers * 1 * page * head_dim
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_tables = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=dev)
|
||||
scale = head_dim**-0.5
|
||||
|
||||
def run(kv):
|
||||
return trtllm_batch_decode_with_kv_cache_mla(
|
||||
query=q,
|
||||
kv_cache=kv,
|
||||
workspace_buffer=ws,
|
||||
qk_nope_head_dim=qk_nope_head_dim,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
block_tables=block_tables,
|
||||
seq_lens=seq_lens,
|
||||
max_seq_len=int(seq_lens.max().item()),
|
||||
bmm1_scale=scale,
|
||||
bmm2_scale=1.0,
|
||||
).clone()
|
||||
|
||||
out_ref = run(kv_contiguous).float()
|
||||
out_view = run(kv_view).float()
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
@@ -0,0 +1,84 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import cdiv
|
||||
|
||||
|
||||
def ref_mla(
|
||||
out: Tensor, # (bs, num_heads, v_head_dim)
|
||||
query: Tensor, # (bs, num_heads, head_dim)
|
||||
kv_cache: Tensor, # (num_blocks, block_size, head_dim)
|
||||
scale: float,
|
||||
block_tables: Tensor, # (bs, max_num_blocks)
|
||||
seq_lens: Tensor, # (bs,)
|
||||
):
|
||||
bs, num_heads, v_head_dim = out.shape
|
||||
head_dim = query.shape[2]
|
||||
|
||||
for i in range(bs):
|
||||
# gather and flatten KV-cache
|
||||
kv = kv_cache[block_tables[i]] # (max_num_blocks, block_size, head_dim)
|
||||
kv = kv.view(1, -1, head_dim)[:, : seq_lens[i]] # (1, seq_len, head_dim)
|
||||
v = kv[:, :, :v_head_dim]
|
||||
|
||||
q = query[i].view(num_heads, 1, head_dim)
|
||||
o = F.scaled_dot_product_attention(q, kv, v, scale=scale, enable_gqa=True)
|
||||
out[i] = o.view(num_heads, v_head_dim)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bs", [4])
|
||||
@pytest.mark.parametrize("mean_seq_len", [256])
|
||||
@pytest.mark.parametrize("h_q", [16])
|
||||
@pytest.mark.parametrize("d", [576])
|
||||
@pytest.mark.parametrize("dv", [512])
|
||||
@pytest.mark.parametrize("block_size", [16])
|
||||
@pytest.mark.parametrize("dtype", [torch.float, torch.half, torch.bfloat16])
|
||||
@pytest.mark.parametrize("varlen", [False, True])
|
||||
@pytest.mark.cpu_model
|
||||
@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
|
||||
def test_mla_decode_cpu(
|
||||
bs: int,
|
||||
mean_seq_len: int,
|
||||
h_q: int,
|
||||
d: int,
|
||||
dv: int,
|
||||
block_size: int,
|
||||
dtype: torch.dtype,
|
||||
varlen: bool,
|
||||
):
|
||||
torch.set_default_dtype(dtype)
|
||||
torch.manual_seed(0)
|
||||
|
||||
scale = d ** (-0.5)
|
||||
if varlen:
|
||||
seq_lens = torch.empty(bs).normal_(mean_seq_len, mean_seq_len / 2)
|
||||
seq_lens = seq_lens.clip(2).to(torch.int32)
|
||||
else:
|
||||
seq_lens = torch.full((bs,), mean_seq_len, dtype=torch.int32)
|
||||
max_seq_len = seq_lens.max().item()
|
||||
seqlen_pad = cdiv(max_seq_len, 256) * 256 # is this necessary?
|
||||
|
||||
q = torch.randn(bs, h_q, d)
|
||||
block_table = torch.arange(bs * seqlen_pad // block_size, dtype=torch.int32)
|
||||
block_table = block_table.view(bs, seqlen_pad // block_size)
|
||||
|
||||
kv_cache = torch.randn(block_table.numel(), block_size, d)
|
||||
for i, seq_len in enumerate(seq_lens.tolist()):
|
||||
kv_cache.view(bs, seqlen_pad, d)[i, seq_len:] = float("nan")
|
||||
|
||||
out_mla = q.new_zeros(bs, h_q, dv)
|
||||
ops.mla_decode_kvcache_cpu(out_mla, q, kv_cache, scale, block_table, seq_lens)
|
||||
|
||||
out_ref = q.new_zeros(bs, h_q, dv)
|
||||
ref_mla(out_ref, q, kv_cache, scale, block_table, seq_lens)
|
||||
|
||||
assert not out_mla.isnan().any(), "Likely read out of bounds"
|
||||
torch.testing.assert_close(out_mla, out_ref)
|
||||
@@ -0,0 +1,234 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
from torch.testing import assert_close
|
||||
|
||||
from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
|
||||
|
||||
|
||||
def test_pack_seq_basic_fp8():
|
||||
"""Test basic functionality of pack_seq_triton with fp8 and 3D tensors."""
|
||||
device = "cuda"
|
||||
dtype = torch.float8_e4m3fn
|
||||
|
||||
# Test cases with 3D tensors (N, H, D)
|
||||
test_cases = [
|
||||
(6, 8, 4, 2, [3, 3]), # (6, 8, 4) -> (2, 3, 8, 4)
|
||||
(10, 4, 8, 3, [2, 4, 4]), # (10, 4, 8) -> (3, 4, 4, 8)
|
||||
(20, 16, 32, 4, [5, 5, 5, 5]), # (20, 16, 32) -> (4, 5, 16, 32)
|
||||
]
|
||||
|
||||
for N, H, D, B, lengths_list in test_cases:
|
||||
# Create input tensor with small values for fp8
|
||||
x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1
|
||||
x = x.to(dtype=dtype)
|
||||
lengths = torch.tensor(lengths_list, device=device)
|
||||
|
||||
# Pack the data
|
||||
packed = pack_seq_triton(x, lengths)
|
||||
|
||||
# Check output shape and properties
|
||||
expected_shape = (B, max(lengths_list), H, D)
|
||||
assert packed.shape == expected_shape
|
||||
assert packed.dtype == dtype
|
||||
assert packed.device == x.device
|
||||
|
||||
# Check that valid data is preserved (within fp8 precision)
|
||||
for b in range(B):
|
||||
start_idx = sum(lengths_list[:b])
|
||||
seq_len = lengths_list[b]
|
||||
|
||||
expected_data = x[start_idx : start_idx + seq_len].to(torch.float32)
|
||||
actual_data = packed[b, :seq_len].to(torch.float32)
|
||||
|
||||
assert_close(actual_data, expected_data, rtol=1e-1, atol=1e-2)
|
||||
|
||||
|
||||
def test_pack_seq_custom_padding_fp8():
|
||||
"""Test pack_seq_triton with custom padding values for fp8."""
|
||||
device = "cuda"
|
||||
dtype = torch.float8_e4m3fn
|
||||
N, H, D, B = 20, 8, 16, 2
|
||||
lengths = torch.tensor([10, 10], device=device)
|
||||
|
||||
x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1
|
||||
x = x.to(dtype=dtype)
|
||||
|
||||
# Test with different padding values
|
||||
for pad_value in [-100.0, -10.0, 0.0, 10.0, 100.0]:
|
||||
result = pack_seq_triton(x, lengths, pad_value=pad_value)
|
||||
|
||||
# Check valid data
|
||||
for b in range(B):
|
||||
start_idx = b * 10
|
||||
expected_data = x[start_idx : start_idx + 10].to(torch.float32)
|
||||
actual_data = result[b, :10].to(torch.float32)
|
||||
assert_close(actual_data, expected_data, rtol=1e-1, atol=1e-2)
|
||||
|
||||
# Check padding (fp8 has limited range, so check for large values)
|
||||
padded_data = result[:, 10:].to(torch.float32)
|
||||
if pad_value < 0:
|
||||
assert torch.all(padded_data < -50) # Large negative values
|
||||
elif pad_value > 0:
|
||||
assert torch.all(padded_data > 50) # Large positive values
|
||||
else:
|
||||
assert torch.allclose(padded_data, torch.zeros_like(padded_data), atol=1e-2)
|
||||
|
||||
|
||||
def test_pack_seq_default_negative_inf_padding_fp8():
|
||||
"""Test that pack_seq_triton uses -inf padding by default for fp8."""
|
||||
device = "cuda"
|
||||
dtype = torch.float8_e4m3fn
|
||||
# B = 2
|
||||
N, H, D = 20, 8, 16
|
||||
lengths = torch.tensor([10, 10], device=device)
|
||||
|
||||
x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1
|
||||
x = x.to(dtype=dtype)
|
||||
result = pack_seq_triton(x, lengths)
|
||||
|
||||
# Check that padding is large negative values (fp8 representation of -inf)
|
||||
padded_data = result[:, 10:].to(torch.float32)
|
||||
assert torch.all(
|
||||
padded_data < -100
|
||||
) # fp8 -inf is represented as large negative number
|
||||
|
||||
|
||||
def test_pack_seq_edge_cases_fp8():
|
||||
"""Test pack_seq_triton with edge cases for fp8."""
|
||||
device = "cuda"
|
||||
dtype = torch.float8_e4m3fn
|
||||
|
||||
# Test with single batch element
|
||||
x = torch.randn(10, 8, 16, dtype=torch.float32, device=device) * 0.1
|
||||
x = x.to(dtype=dtype)
|
||||
lengths = torch.tensor([10], device=device)
|
||||
result = pack_seq_triton(x, lengths)
|
||||
assert result.shape == (1, 10, 8, 16)
|
||||
|
||||
# Test with very short sequences
|
||||
x = torch.randn(20, 4, 8, dtype=torch.float32, device=device) * 0.1
|
||||
x = x.to(dtype=dtype)
|
||||
lengths = torch.tensor([1, 1, 1], device=device)
|
||||
result = pack_seq_triton(x, lengths)
|
||||
assert result.shape == (3, 1, 4, 8)
|
||||
|
||||
# Test with different sequence lengths
|
||||
x = torch.randn(15, 8, 16, dtype=torch.float32, device=device) * 0.1
|
||||
x = x.to(dtype=dtype)
|
||||
lengths = torch.tensor([5, 7, 3], device=device)
|
||||
result = pack_seq_triton(x, lengths)
|
||||
assert result.shape == (3, 7, 8, 16)
|
||||
|
||||
|
||||
def test_pack_seq_different_block_sizes_fp8():
|
||||
"""Test pack_seq_triton with different block sizes for fp8."""
|
||||
device = "cuda"
|
||||
dtype = torch.float8_e4m3fn
|
||||
N, H, D, B = 100, 16, 32, 4
|
||||
lengths = torch.tensor([25, 25, 25, 25], device=device)
|
||||
|
||||
x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1
|
||||
x = x.to(dtype=dtype)
|
||||
|
||||
# Test different block sizes
|
||||
for block_t, block_d in [(32, 32), (64, 64), (128, 128)]:
|
||||
result = pack_seq_triton(x, lengths, block_t=block_t, block_d=block_d)
|
||||
|
||||
assert result.shape == (B, 25, H, D)
|
||||
|
||||
# Check that valid data is preserved (within fp8 precision)
|
||||
for b in range(B):
|
||||
start_idx = b * 25
|
||||
expected_data = x[start_idx : start_idx + 25].to(torch.float32)
|
||||
actual_data = result[b, :25].to(torch.float32)
|
||||
assert_close(actual_data, expected_data, rtol=1e-1, atol=1e-2)
|
||||
|
||||
|
||||
def test_pack_seq_shape_consistency():
|
||||
"""Test that pack_seq_triton maintains shape consistency."""
|
||||
device = "cuda"
|
||||
dtype = torch.float8_e4m3fn
|
||||
N, H, D, B = 20, 8, 16, 2
|
||||
lengths = torch.tensor([10, 10], device=device)
|
||||
|
||||
x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1
|
||||
x = x.to(dtype=dtype)
|
||||
|
||||
result = pack_seq_triton(x, lengths)
|
||||
|
||||
# Check shape consistency
|
||||
assert result.shape[0] == B # Batch dimension
|
||||
assert result.shape[1] == lengths.max().item() # Max sequence length
|
||||
assert result.shape[2:] == x.shape[1:] # Feature dimensions preserved
|
||||
|
||||
|
||||
def test_pack_unpack_roundtrip_fp8():
|
||||
"""Test that pack -> unpack gives us back the original data for fp8."""
|
||||
device = "cuda"
|
||||
dtype = torch.float8_e4m3fn
|
||||
|
||||
# Test cases with 3D tensors
|
||||
test_cases = [
|
||||
(6, 8, 4, 2, [3, 3]),
|
||||
(10, 4, 8, 3, [2, 4, 4]),
|
||||
(20, 16, 32, 4, [5, 5, 5, 5]),
|
||||
(15, 8, 16, 3, [7, 5, 3]),
|
||||
]
|
||||
|
||||
for N, H, D, B, lengths_list in test_cases:
|
||||
# Create input tensor with small values for fp8
|
||||
x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1
|
||||
x = x.to(dtype=dtype)
|
||||
lengths = torch.tensor(lengths_list, device=device)
|
||||
|
||||
# Pack the data
|
||||
packed = pack_seq_triton(x, lengths)
|
||||
|
||||
# Unpack the data
|
||||
unpacked = unpack_seq_triton(packed, lengths)
|
||||
|
||||
# Check that we get back the original data (within fp8 precision)
|
||||
assert unpacked.shape == x.shape
|
||||
x_f32 = x.to(torch.float32)
|
||||
unpacked_f32 = unpacked.to(torch.float32)
|
||||
assert_close(x_f32, unpacked_f32, rtol=1e-3, atol=1e-3)
|
||||
|
||||
# Unpack without explicit start locations (computed in kernel)
|
||||
unpacked_with_loc = unpack_seq_triton(packed, lengths)
|
||||
assert_close(x_f32, unpacked_with_loc.to(torch.float32), rtol=1e-3, atol=1e-2)
|
||||
|
||||
|
||||
def test_unpack_seq_triton_edge_cases_fp8():
|
||||
"""Test unpack function with edge cases for fp8."""
|
||||
device = "cuda"
|
||||
dtype = torch.float8_e4m3fn
|
||||
|
||||
# Test with single batch element
|
||||
x = torch.randn(10, 8, 16, dtype=torch.float32, device=device) * 0.1
|
||||
x = x.to(dtype=dtype)
|
||||
lengths = torch.tensor([10], device=device)
|
||||
packed = pack_seq_triton(x, lengths)
|
||||
unpacked = unpack_seq_triton(packed, lengths)
|
||||
assert unpacked.shape == x.shape
|
||||
assert_close(x.to(torch.float32), unpacked.to(torch.float32), rtol=1e-1, atol=1e-2)
|
||||
|
||||
# Test with very short sequences
|
||||
x = torch.randn(20, 4, 8, dtype=torch.float32, device=device) * 0.1
|
||||
x = x.to(dtype=dtype)
|
||||
lengths = torch.tensor([1, 1, 1], device=device)
|
||||
packed = pack_seq_triton(x, lengths)
|
||||
unpacked = unpack_seq_triton(packed, lengths)
|
||||
# Only compare the first 3 elements that were actually packed
|
||||
assert_close(
|
||||
x[:3].to(torch.float32), unpacked.to(torch.float32), rtol=1e-1, atol=1e-2
|
||||
)
|
||||
|
||||
x = torch.randn(15, 8, 16, dtype=torch.float32, device=device) * 0.1
|
||||
x = x.to(dtype=dtype)
|
||||
lengths = torch.tensor([5, 7, 3], device=device)
|
||||
packed = pack_seq_triton(x, lengths)
|
||||
unpacked = unpack_seq_triton(packed, lengths)
|
||||
assert unpacked.shape == x.shape
|
||||
assert_close(x.to(torch.float32), unpacked.to(torch.float32), rtol=1e-1, atol=1e-2)
|
||||
@@ -0,0 +1,682 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import math
|
||||
import random
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from contextlib import nullcontext
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
|
||||
from vllm.v1.attention.ops.chunked_prefill_paged_decode import (
|
||||
chunked_prefill_paged_decode,
|
||||
)
|
||||
from vllm.v1.attention.ops.prefix_prefill import context_attention_fwd
|
||||
|
||||
NUM_HEADS = [64]
|
||||
NUM_QUERIES_PER_KV = [1, 64]
|
||||
HEAD_SIZES = [24, 128]
|
||||
DTYPES = [torch.float16]
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
SLIDING_WINDOW = [0, 16, 2048]
|
||||
KV_CACHE_DTYPES = ["auto", "fp8", "fp8_e5m2"]
|
||||
|
||||
OPS = [chunked_prefill_paged_decode, context_attention_fwd]
|
||||
|
||||
|
||||
def create_causal_attention_mask_for_sdpa(
|
||||
query_lens: list[int],
|
||||
seq_lens: list[int],
|
||||
sliding_window: int = 0,
|
||||
device: torch.device = None,
|
||||
dtype: torch.dtype = None,
|
||||
) -> torch.Tensor:
|
||||
total_queries = sum(query_lens)
|
||||
total_keys = sum(seq_lens)
|
||||
|
||||
# Create a mask filled with -inf
|
||||
mask = torch.full(
|
||||
(total_queries, total_keys), float("-inf"), device=device, dtype=dtype
|
||||
)
|
||||
|
||||
query_start = 0
|
||||
key_start = 0
|
||||
|
||||
for query_len, seq_len in zip(query_lens, seq_lens):
|
||||
query_end = query_start + query_len
|
||||
key_end = key_start + seq_len
|
||||
q_indices = torch.arange(query_len, device=device)
|
||||
k_indices = torch.arange(seq_len, device=device)
|
||||
q_pos_in_seq = seq_len - query_len + q_indices
|
||||
|
||||
valid_mask = k_indices[None, :] <= q_pos_in_seq[:, None]
|
||||
|
||||
if sliding_window > 0:
|
||||
valid_mask &= k_indices[None, :] >= (
|
||||
q_pos_in_seq[:, None] - sliding_window + 1
|
||||
)
|
||||
|
||||
mask[query_start:query_end, key_start:key_end][valid_mask] = 0.0
|
||||
|
||||
query_start = query_end
|
||||
key_start = key_end
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
def create_alibi_causal_mask(
|
||||
query_len: int,
|
||||
seq_len: int,
|
||||
alibi_slopes: torch.Tensor,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
query_pos = torch.arange(
|
||||
seq_len - query_len, seq_len, device=device, dtype=torch.float32
|
||||
)
|
||||
key_pos = torch.arange(seq_len, device=device, dtype=torch.float32)
|
||||
|
||||
rel_pos = key_pos[None, :] - query_pos[:, None]
|
||||
|
||||
# Apply ALiBi slopes: [num_heads, query_len, seq_len]
|
||||
alibi_bias = alibi_slopes[:, None, None] * rel_pos[None, :, :]
|
||||
alibi_bias = alibi_bias.to(dtype)
|
||||
|
||||
# Apply causal mask: prevent attending to future positions
|
||||
# causal_mask[i, j] = True if key_pos[j] <= query_pos[i]
|
||||
causal_mask = key_pos[None, :] <= query_pos[:, None]
|
||||
alibi_bias = alibi_bias.masked_fill(~causal_mask[None, :, :], float("-inf"))
|
||||
|
||||
# Add batch dimension: [1, num_heads, query_len, seq_len]
|
||||
# SDPA expects batch dimension even for single sequences
|
||||
return alibi_bias.unsqueeze(0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
|
||||
@pytest.mark.parametrize("op", OPS)
|
||||
@torch.inference_mode()
|
||||
def test_contexted_kv_attention(
|
||||
num_heads: int,
|
||||
num_queries_per_kv: int,
|
||||
head_size: int,
|
||||
sliding_window: int,
|
||||
dtype: torch.dtype,
|
||||
kv_cache_dtype: str,
|
||||
device: str,
|
||||
op: Callable,
|
||||
block_size: int = 32,
|
||||
) -> None:
|
||||
if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):
|
||||
pytest.skip(
|
||||
"Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"
|
||||
)
|
||||
|
||||
if (
|
||||
current_platform.is_rocm()
|
||||
and op is chunked_prefill_paged_decode
|
||||
and kv_cache_dtype == "fp8_e5m2"
|
||||
):
|
||||
pytest.skip("ROCm custom paged attention does not support fp8_e5m2 KV cache")
|
||||
|
||||
set_random_seed(0)
|
||||
torch.set_default_device(device)
|
||||
|
||||
# Need this, otherwise when we capture the graph the process
|
||||
# for GPU 1 would run on both GPU0 and GPU1 and things would hang
|
||||
#
|
||||
# see also similar issue: https://github.com/Dao-AILab/flash-attention/issues/523
|
||||
torch.accelerator.set_device_index(device)
|
||||
|
||||
MAX_SEQ_LEN = 1024
|
||||
MAX_CTX_LEN = 1024
|
||||
BS = 10
|
||||
cache_size = 640
|
||||
max_block_per_request = 64
|
||||
query_lens = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)]
|
||||
# ensure one sequence in batch is a decode
|
||||
query_lens[-1] = 1
|
||||
|
||||
ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)]
|
||||
seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)]
|
||||
num_kv_heads = num_heads // num_queries_per_kv
|
||||
|
||||
num_tokens = sum(query_lens)
|
||||
query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
|
||||
query.uniform_(-1e-3, 1e-3)
|
||||
output = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
|
||||
|
||||
kv = torch.empty(sum(seq_lens), 2, num_kv_heads, head_size, dtype=dtype)
|
||||
kv.uniform_(-1e-3, 1e-3)
|
||||
key, value = kv.unbind(dim=1)
|
||||
|
||||
if kv_cache_dtype == "auto":
|
||||
cache_dtype = dtype
|
||||
else:
|
||||
cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype]
|
||||
k_cache = torch.zeros(
|
||||
cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype
|
||||
)
|
||||
v_cache = torch.zeros(
|
||||
cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype
|
||||
)
|
||||
k = torch.zeros(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
|
||||
v = torch.zeros(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
|
||||
values = torch.arange(0, cache_size, dtype=torch.int32)
|
||||
values = values[torch.randperm(cache_size)]
|
||||
block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)
|
||||
b_seq_len = torch.tensor(seq_lens, dtype=torch.int32)
|
||||
b_ctx_len = torch.tensor(ctx_lens, dtype=torch.int32)
|
||||
b_start_loc = torch.cumsum(torch.tensor([0] + query_lens), dim=0).to(torch.int32)
|
||||
max_input_len = MAX_SEQ_LEN
|
||||
# copy kv to cache
|
||||
b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1]), dim=0).to(
|
||||
torch.int32
|
||||
)
|
||||
for i in range(BS):
|
||||
for j in range(query_lens[i]):
|
||||
k[b_start_loc[i] + j].copy_(key[b_seq_start_loc[i] + b_ctx_len[i] + j])
|
||||
v[b_start_loc[i] + j].copy_(value[b_seq_start_loc[i] + b_ctx_len[i] + j])
|
||||
cur_ctx = 0
|
||||
block_id = 0
|
||||
while cur_ctx < b_ctx_len[i]:
|
||||
start_loc = b_seq_start_loc[i] + cur_ctx
|
||||
if cur_ctx + block_size > b_ctx_len[i]:
|
||||
end_loc = b_seq_start_loc[i] + b_ctx_len[i]
|
||||
else:
|
||||
end_loc = start_loc + block_size
|
||||
start_slot = block_table[i, block_id] * block_size
|
||||
end_slot = start_slot + end_loc - start_loc
|
||||
k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
|
||||
key[start_loc:end_loc]
|
||||
)
|
||||
v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
|
||||
value[start_loc:end_loc]
|
||||
)
|
||||
cur_ctx += block_size
|
||||
block_id += 1
|
||||
# transpose K_cache[num_blocks, block_size, num_kv_heads, head_size]
|
||||
# to K_cache[num_blocks, num_kv_heads, head_size/8, block_size, 8]
|
||||
k_cache = (
|
||||
k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8)
|
||||
.permute(0, 2, 3, 1, 4)
|
||||
.contiguous()
|
||||
)
|
||||
# transpose V_cache[num_blocks, block_size, num_kv_heads, head_size]
|
||||
# to V_cache[num_blocks, num_kv_heads, head_size, block_size]
|
||||
v_cache = (
|
||||
v_cache.view(-1, block_size, num_kv_heads, head_size)
|
||||
.permute(0, 2, 3, 1)
|
||||
.contiguous()
|
||||
)
|
||||
k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
|
||||
|
||||
# Warm up the Triton kernel by calling it once before actually measuring
|
||||
# generation time
|
||||
op(
|
||||
query,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
kv_cache_dtype,
|
||||
k_cache,
|
||||
v_cache,
|
||||
block_table,
|
||||
b_start_loc,
|
||||
b_seq_len,
|
||||
MAX_CTX_LEN,
|
||||
max_input_len,
|
||||
k_scale,
|
||||
v_scale,
|
||||
sliding_window=sliding_window,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
start_time = time.time()
|
||||
op(
|
||||
query,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
kv_cache_dtype,
|
||||
k_cache,
|
||||
v_cache,
|
||||
block_table,
|
||||
b_start_loc,
|
||||
b_seq_len,
|
||||
MAX_CTX_LEN,
|
||||
max_input_len,
|
||||
k_scale,
|
||||
v_scale,
|
||||
sliding_window=sliding_window,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
end_time = time.time()
|
||||
print(f"triton Time: {(end_time - start_time) * 1000:.2f} ms")
|
||||
|
||||
scale = float(1.0 / (head_size**0.5))
|
||||
|
||||
# Reshape for SDPA: (seq_len, num_heads, head_size) ->
|
||||
# (1, num_heads, seq_len, head_size)
|
||||
query_sdpa = query.view(num_tokens, num_kv_heads, num_queries_per_kv, head_size)
|
||||
query_sdpa = query_sdpa.permute(1, 2, 0, 3).reshape(
|
||||
1, num_heads, num_tokens, head_size
|
||||
)
|
||||
|
||||
# Expand key and value for GQA/MQA to match query heads
|
||||
key_sdpa = key[:, :, None, :].expand(
|
||||
key.shape[0], num_kv_heads, num_queries_per_kv, key.shape[-1]
|
||||
)
|
||||
key_sdpa = key_sdpa.permute(1, 2, 0, 3).reshape(
|
||||
1, num_heads, sum(seq_lens), head_size
|
||||
)
|
||||
|
||||
value_sdpa = value[:, :, None, :].expand(
|
||||
value.shape[0], num_kv_heads, num_queries_per_kv, value.shape[-1]
|
||||
)
|
||||
value_sdpa = value_sdpa.permute(1, 2, 0, 3).reshape(
|
||||
1, num_heads, sum(seq_lens), head_size
|
||||
)
|
||||
|
||||
attn_mask = create_causal_attention_mask_for_sdpa(
|
||||
query_lens, seq_lens, sliding_window, device=device, dtype=dtype
|
||||
)
|
||||
|
||||
output_ref = F.scaled_dot_product_attention(
|
||||
query_sdpa,
|
||||
key_sdpa,
|
||||
value_sdpa,
|
||||
attn_mask=attn_mask,
|
||||
dropout_p=0.0,
|
||||
scale=scale,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
start_time = time.time()
|
||||
output_ref = F.scaled_dot_product_attention(
|
||||
query_sdpa,
|
||||
key_sdpa,
|
||||
value_sdpa,
|
||||
attn_mask=attn_mask,
|
||||
dropout_p=0.0,
|
||||
scale=scale,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
end_time = time.time()
|
||||
print(f"PyTorch SDPA Time: {(end_time - start_time) * 1000:.2f} ms")
|
||||
|
||||
# Reshape output back to (num_tokens, num_heads, head_size)
|
||||
output_ref = output_ref.view(num_heads, num_tokens, head_size)
|
||||
output_ref = output_ref.permute(1, 0, 2).contiguous()
|
||||
atol = 1e-3 if "fp8" in kv_cache_dtype else 1e-4
|
||||
torch.testing.assert_close(output, output_ref, atol=atol, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("op", OPS)
|
||||
@torch.inference_mode()
|
||||
def test_contexted_kv_attention_alibi(
|
||||
num_heads: int,
|
||||
num_queries_per_kv: int,
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
kv_cache_dtype: str,
|
||||
device: str,
|
||||
op: Callable,
|
||||
block_size: int = 32,
|
||||
) -> None:
|
||||
if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):
|
||||
pytest.skip(
|
||||
"Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"
|
||||
)
|
||||
|
||||
if (
|
||||
current_platform.is_rocm()
|
||||
and op is chunked_prefill_paged_decode
|
||||
and kv_cache_dtype == "fp8_e5m2"
|
||||
):
|
||||
pytest.skip("ROCm custom paged attention does not support fp8_e5m2 KV cache")
|
||||
|
||||
set_random_seed(0)
|
||||
torch.set_default_device(device)
|
||||
|
||||
# Need this, otherwise when we capture the graph the process
|
||||
# for GPU 1 would run on both GPU0 and GPU1 and things would hang
|
||||
#
|
||||
# see also similar issue: https://github.com/Dao-AILab/flash-attention/issues/523
|
||||
torch.accelerator.set_device_index(device)
|
||||
|
||||
def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
|
||||
# Fork from: vllm/vllm/model_executor/models/bloom.py#L44
|
||||
closest_power_of_2 = 2 ** math.floor(math.log2(total_num_heads))
|
||||
base = torch.tensor(
|
||||
2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
|
||||
slopes = torch.pow(base, powers)
|
||||
|
||||
if closest_power_of_2 != total_num_heads:
|
||||
extra_base = torch.tensor(
|
||||
2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
num_remaining_heads = min(
|
||||
closest_power_of_2, total_num_heads - closest_power_of_2
|
||||
)
|
||||
extra_powers = torch.arange(
|
||||
start=1, end=1 + 2 * num_remaining_heads, step=2, dtype=torch.int32
|
||||
)
|
||||
slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
|
||||
return slopes
|
||||
|
||||
alibi_slopes = _get_alibi_slopes(num_heads).to(device)
|
||||
|
||||
MAX_SEQ_LEN = 1024
|
||||
MAX_CTX_LEN = 1024
|
||||
BS = 10
|
||||
cache_size = 640
|
||||
max_block_per_request = 64
|
||||
query_lens = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)]
|
||||
ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)]
|
||||
seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)]
|
||||
num_kv_heads = num_heads // num_queries_per_kv
|
||||
|
||||
num_tokens = sum(query_lens)
|
||||
query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
|
||||
query.uniform_(-1e-3, 1e-3)
|
||||
output = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
|
||||
|
||||
kv = torch.empty(sum(seq_lens), 2, num_kv_heads, head_size, dtype=dtype)
|
||||
kv.uniform_(-1e-3, 1e-3)
|
||||
key, value = kv.unbind(dim=1)
|
||||
if kv_cache_dtype == "auto":
|
||||
cache_dtype = dtype
|
||||
else:
|
||||
cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype]
|
||||
k_cache = torch.zeros(
|
||||
cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype
|
||||
)
|
||||
v_cache = torch.zeros(
|
||||
cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype
|
||||
)
|
||||
k = torch.zeros(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
|
||||
v = torch.zeros(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
|
||||
values = torch.arange(0, cache_size, dtype=torch.int32)
|
||||
values = values[torch.randperm(cache_size)]
|
||||
block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)
|
||||
b_seq_len = torch.tensor(seq_lens, dtype=torch.int32)
|
||||
b_ctx_len = torch.tensor(ctx_lens, dtype=torch.int32)
|
||||
b_start_loc = torch.cumsum(torch.tensor([0] + query_lens), dim=0).to(torch.int32)
|
||||
max_input_len = MAX_SEQ_LEN
|
||||
# copy kv to cache
|
||||
b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1]), dim=0).to(
|
||||
torch.int32
|
||||
)
|
||||
for i in range(BS):
|
||||
for j in range(query_lens[i]):
|
||||
k[b_start_loc[i] + j].copy_(key[b_seq_start_loc[i] + b_ctx_len[i] + j])
|
||||
v[b_start_loc[i] + j].copy_(value[b_seq_start_loc[i] + b_ctx_len[i] + j])
|
||||
cur_ctx = 0
|
||||
block_id = 0
|
||||
while cur_ctx < b_ctx_len[i]:
|
||||
start_loc = b_seq_start_loc[i] + cur_ctx
|
||||
if cur_ctx + block_size > b_ctx_len[i]:
|
||||
end_loc = b_seq_start_loc[i] + b_ctx_len[i]
|
||||
else:
|
||||
end_loc = start_loc + block_size
|
||||
start_slot = block_table[i, block_id] * block_size
|
||||
end_slot = start_slot + end_loc - start_loc
|
||||
k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
|
||||
key[start_loc:end_loc]
|
||||
)
|
||||
v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
|
||||
value[start_loc:end_loc]
|
||||
)
|
||||
cur_ctx += block_size
|
||||
block_id += 1
|
||||
# transpose K_cache[num_blocks, block_size, num_kv_heads, head_size]
|
||||
# to K_cache[num_blocks, num_kv_heads, head_size/8, block_size, 8]
|
||||
k_cache = (
|
||||
k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8)
|
||||
.permute(0, 2, 3, 1, 4)
|
||||
.contiguous()
|
||||
)
|
||||
# transpose V_cache[num_blocks, block_size, num_kv_heads, head_size]
|
||||
# to V_cache[num_blocks, num_kv_heads, head_size, block_size]
|
||||
v_cache = (
|
||||
v_cache.view(-1, block_size, num_kv_heads, head_size)
|
||||
.permute(0, 2, 3, 1)
|
||||
.contiguous()
|
||||
)
|
||||
k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
|
||||
|
||||
# Warm up the Triton kernel by calling it once before actually measuring
|
||||
# generation time
|
||||
op(
|
||||
query,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
kv_cache_dtype,
|
||||
k_cache,
|
||||
v_cache,
|
||||
block_table,
|
||||
b_start_loc,
|
||||
b_seq_len,
|
||||
MAX_CTX_LEN,
|
||||
max_input_len,
|
||||
k_scale,
|
||||
v_scale,
|
||||
alibi_slopes=alibi_slopes,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
start_time = time.time()
|
||||
op(
|
||||
query,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
kv_cache_dtype,
|
||||
k_cache,
|
||||
v_cache,
|
||||
block_table,
|
||||
b_start_loc,
|
||||
b_seq_len,
|
||||
MAX_CTX_LEN,
|
||||
max_input_len,
|
||||
k_scale,
|
||||
v_scale,
|
||||
alibi_slopes=alibi_slopes,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
end_time = time.time()
|
||||
print(f"triton Time: {(end_time - start_time) * 1000:.2f} ms")
|
||||
scale = float(1.0 / (head_size**0.5))
|
||||
|
||||
# Prepare query, key, value for SDPA
|
||||
# Expand key and value for GQA/MQA to match query heads
|
||||
key_expanded = key[:, :, None, :].expand(
|
||||
key.shape[0], num_kv_heads, num_queries_per_kv, key.shape[-1]
|
||||
)
|
||||
value_expanded = value[:, :, None, :].expand(
|
||||
value.shape[0], num_kv_heads, num_queries_per_kv, value.shape[-1]
|
||||
)
|
||||
|
||||
output_ref = torch.empty_like(output)
|
||||
|
||||
torch.accelerator.synchronize()
|
||||
start_time = time.time()
|
||||
|
||||
query_start = 0
|
||||
key_start = 0
|
||||
for i, (query_len, seq_len) in enumerate(zip(query_lens, seq_lens)):
|
||||
query_end = query_start + query_len
|
||||
key_end = key_start + seq_len
|
||||
|
||||
# Get query, key, value for this sequence
|
||||
q = query[query_start:query_end] # [query_len, num_heads, head_size]
|
||||
k = key_expanded[
|
||||
key_start:key_end
|
||||
] # [seq_len, num_kv_heads, num_queries_per_kv, head_size]
|
||||
v = value_expanded[
|
||||
key_start:key_end
|
||||
] # [seq_len, num_kv_heads, num_queries_per_kv, head_size]
|
||||
|
||||
# Reshape for SDPA: (batch=1, num_heads, seq_len, head_size)
|
||||
q_sdpa = q.view(query_len, num_kv_heads, num_queries_per_kv, head_size)
|
||||
q_sdpa = (
|
||||
q_sdpa.permute(1, 2, 0, 3)
|
||||
.reshape(1, num_heads, query_len, head_size)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
k_sdpa = (
|
||||
k.permute(1, 2, 0, 3).reshape(1, num_heads, seq_len, head_size).contiguous()
|
||||
)
|
||||
v_sdpa = (
|
||||
v.permute(1, 2, 0, 3).reshape(1, num_heads, seq_len, head_size).contiguous()
|
||||
)
|
||||
|
||||
# Create ALiBi causal mask for this sequence using utility function
|
||||
alibi_mask = create_alibi_causal_mask(
|
||||
query_len, seq_len, alibi_slopes, device, dtype
|
||||
)
|
||||
|
||||
# Compute attention. On ROCm we force use of the Math SDPA backend rather than
|
||||
# the Flash or Mem-Efficient backends for increased numerical accuracy
|
||||
if current_platform.is_rocm():
|
||||
sdpa_context = sdpa_kernel(SDPBackend.MATH)
|
||||
else:
|
||||
sdpa_context = nullcontext()
|
||||
with sdpa_context:
|
||||
out = F.scaled_dot_product_attention(
|
||||
q_sdpa,
|
||||
k_sdpa,
|
||||
v_sdpa,
|
||||
attn_mask=alibi_mask,
|
||||
dropout_p=0.0,
|
||||
scale=scale,
|
||||
)
|
||||
|
||||
# Reshape output back to [query_len, num_heads, head_size]
|
||||
out = out.view(num_heads, query_len, head_size).permute(1, 0, 2)
|
||||
output_ref[query_start:query_end].copy_(out)
|
||||
|
||||
query_start = query_end
|
||||
key_start = key_end
|
||||
|
||||
torch.accelerator.synchronize()
|
||||
end_time = time.time()
|
||||
print(f"PyTorch SDPA Time: {(end_time - start_time) * 1000:.2f} ms")
|
||||
atol = 1e-3 if "fp8" in kv_cache_dtype else 1e-6
|
||||
torch.testing.assert_close(output, output_ref, atol=atol, rtol=0)
|
||||
|
||||
|
||||
# These tests are optional to only run when explicitly invoked
|
||||
#
|
||||
# pytest -v -s --optional \
|
||||
# tests/kernels/test_prefix_prefill.py::test_contexted_kv_attention_f32
|
||||
#
|
||||
# These tests are useful to test model dtype float32 on Turing devices.
|
||||
# We skip them to not increase the time when running tests on CI
|
||||
@pytest.mark.optional
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("dtype", [torch.float32])
|
||||
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
|
||||
@pytest.mark.parametrize("op", OPS)
|
||||
@torch.inference_mode()
|
||||
def test_contexted_kv_attention_f32(
|
||||
num_heads: int,
|
||||
num_queries_per_kv: int,
|
||||
head_size: int,
|
||||
sliding_window: int,
|
||||
dtype: torch.dtype,
|
||||
kv_cache_dtype: str,
|
||||
device: str,
|
||||
op: Callable,
|
||||
) -> None:
|
||||
test_contexted_kv_attention(
|
||||
num_heads,
|
||||
num_queries_per_kv,
|
||||
head_size,
|
||||
sliding_window,
|
||||
dtype,
|
||||
kv_cache_dtype,
|
||||
device,
|
||||
op,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.optional
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("dtype", [torch.float32])
|
||||
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("op", OPS)
|
||||
@torch.inference_mode()
|
||||
def test_contexted_kv_attention_alibi_f32(
|
||||
num_heads: int,
|
||||
num_queries_per_kv: int,
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
kv_cache_dtype: str,
|
||||
device: str,
|
||||
op: Callable,
|
||||
) -> None:
|
||||
test_contexted_kv_attention_alibi(
|
||||
num_heads, num_queries_per_kv, head_size, dtype, kv_cache_dtype, device, op
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("head_size", [128])
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("op", OPS)
|
||||
@torch.inference_mode()
|
||||
def test_qwen3_nonstandard_block_size(
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
op: Callable,
|
||||
) -> None:
|
||||
"""
|
||||
A separate test function specifically added
|
||||
for Qwen3-Next-80B (Block Size 544).
|
||||
"""
|
||||
if not current_platform.is_rocm():
|
||||
pytest.skip("544 block size optimization is only for ROCm.")
|
||||
|
||||
test_contexted_kv_attention(
|
||||
num_heads=64,
|
||||
num_queries_per_kv=1,
|
||||
head_size=head_size,
|
||||
block_size=544,
|
||||
sliding_window=0,
|
||||
dtype=dtype,
|
||||
kv_cache_dtype="auto",
|
||||
device=device,
|
||||
op=op,
|
||||
)
|
||||
@@ -0,0 +1,202 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Regression test for AITER MLA persistent decode metadata dtypes.
|
||||
|
||||
For the gfx950 fp8/fp8 nhead=32 qlen=1 fold path, the split/reduce metadata
|
||||
layout depends on the q/kv element size. The builder must forward dtype_q/dtype_kv
|
||||
to ``get_mla_metadata_v1``; omitting them lays out the work for the wrong dtype
|
||||
and corrupts decode output. The test pins the builder's metadata to a golden
|
||||
recomputed at runtime with the explicit correct dtypes.
|
||||
"""
|
||||
|
||||
import types
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm._aiter_ops import is_aiter_found
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
def _on_gfx950() -> bool:
|
||||
if not (current_platform.is_rocm() and is_aiter_found()):
|
||||
return False
|
||||
from vllm.platforms.rocm import on_gfx950
|
||||
|
||||
return on_gfx950()
|
||||
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not _on_gfx950(),
|
||||
reason="AITER MLA fp8 persistent decode metadata is gfx950-only",
|
||||
)
|
||||
|
||||
# The fold path that the bug corrupted: fp8 query + fp8 KV-cache, 32 query
|
||||
# heads, single-token decode, batch 128, context 8192, page_size 1.
|
||||
NUM_QUERY_HEADS = 32
|
||||
DECODE_QLEN = 1
|
||||
BATCH_SIZE = 128
|
||||
CONTEXT_LEN = 8192
|
||||
PAGE_SIZE = 1
|
||||
|
||||
# Expected dtypes for this fold path: bf16 model dtype -> bf16 query; fp8
|
||||
# KV-cache -> fp8_e4m3 kv.
|
||||
EXPECTED_Q_DTYPE = torch.bfloat16
|
||||
EXPECTED_KV_DTYPE = torch.float8_e4m3fn
|
||||
|
||||
# The split/reduce content tensors filled by get_mla_metadata_v1. work_meta_data
|
||||
# is excluded: it holds raw device pointers, never equal across allocations.
|
||||
_CONTENT_METADATA_FIELDS = (
|
||||
"work_indptr",
|
||||
"work_info_set",
|
||||
"reduce_indptr",
|
||||
"reduce_final_map",
|
||||
"reduce_partial_map",
|
||||
)
|
||||
|
||||
# The builder's get_mla_metadata_v1 call passes 6 input args then 6 output
|
||||
# buffers (see AiterMLAMetadataBuilder._build_decode). Output order -> field.
|
||||
_NUM_INPUT_ARGS = 6
|
||||
_OUTPUT_ARG_FIELDS = (
|
||||
"work_meta_data",
|
||||
"work_info_set",
|
||||
"work_indptr",
|
||||
"reduce_indptr",
|
||||
"reduce_final_map",
|
||||
"reduce_partial_map",
|
||||
)
|
||||
|
||||
|
||||
def _build_decode_metadata():
|
||||
"""Build AITER MLA decode metadata for the fp8/fp8 nhead=32 fold path.
|
||||
|
||||
Returns ``(metadata, captured)`` where ``captured`` records the positional
|
||||
args/kwargs the builder passed to ``get_mla_metadata_v1``, so the golden can
|
||||
be recomputed from the identical inputs.
|
||||
"""
|
||||
from tests.v1.attention.utils import (
|
||||
BatchSpec,
|
||||
create_common_attn_metadata,
|
||||
create_vllm_config,
|
||||
)
|
||||
from vllm.config.vllm import set_current_vllm_config
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.v1.kv_cache_interface import MLAAttentionSpec
|
||||
from vllm.v1.worker.workspace import init_workspace_manager
|
||||
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
vllm_config = create_vllm_config(
|
||||
model_name="deepseek-ai/DeepSeek-R1",
|
||||
max_model_len=CONTEXT_LEN,
|
||||
# One flat page per token (page_size=1); +buffer for the null block.
|
||||
num_gpu_blocks=BATCH_SIZE * CONTEXT_LEN + 200,
|
||||
block_size=PAGE_SIZE,
|
||||
max_num_seqs=BATCH_SIZE,
|
||||
max_num_batched_tokens=8192,
|
||||
hf_config_override={"num_attention_heads": NUM_QUERY_HEADS},
|
||||
)
|
||||
vllm_config.cache_config.cache_dtype = "fp8"
|
||||
|
||||
spec = MLAAttentionSpec(
|
||||
block_size=PAGE_SIZE,
|
||||
num_kv_heads=1,
|
||||
head_size=vllm_config.model_config.get_head_size(),
|
||||
dtype=vllm_config.model_config.dtype,
|
||||
cache_dtype_str="fp8",
|
||||
)
|
||||
|
||||
builder_cls = AttentionBackendEnum.ROCM_AITER_MLA.get_class().get_builder_cls()
|
||||
|
||||
# The builder reads layer.prefill_backend from static_forward_context; a
|
||||
# stub with the attribute is enough for metadata construction.
|
||||
layer_name = "placeholder"
|
||||
vllm_config.compilation_config.static_forward_context[layer_name] = (
|
||||
types.SimpleNamespace(prefill_backend=torch.empty((1,)))
|
||||
)
|
||||
|
||||
init_workspace_manager(device)
|
||||
|
||||
batch_spec = BatchSpec(
|
||||
seq_lens=[CONTEXT_LEN] * BATCH_SIZE,
|
||||
query_lens=[DECODE_QLEN] * BATCH_SIZE,
|
||||
)
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
builder = builder_cls(spec, [layer_name], vllm_config, device)
|
||||
common_attn_metadata = create_common_attn_metadata(
|
||||
batch_spec, PAGE_SIZE, device, arange_block_indices=True
|
||||
)
|
||||
|
||||
import aiter
|
||||
|
||||
real_get_mla_metadata_v1 = aiter.get_mla_metadata_v1
|
||||
|
||||
def spy(*args, **kwargs):
|
||||
captured["args"] = args
|
||||
captured["kwargs"] = dict(kwargs)
|
||||
return real_get_mla_metadata_v1(*args, **kwargs)
|
||||
|
||||
with patch("aiter.get_mla_metadata_v1", spy):
|
||||
metadata = builder.build(
|
||||
common_prefix_len=0,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
)
|
||||
|
||||
return metadata, captured
|
||||
|
||||
|
||||
def _compute_golden_metadata(captured: dict) -> dict[str, torch.Tensor]:
|
||||
"""Recompute the persistent metadata with explicit fp8/bf16 dtypes.
|
||||
|
||||
Replays ``get_mla_metadata_v1`` on the builder's exact input tensors with
|
||||
fresh output buffers and the explicitly-correct dtypes. This reference must
|
||||
match the builder's output when the fix is in place.
|
||||
"""
|
||||
import aiter
|
||||
|
||||
args = captured["args"]
|
||||
inputs = args[:_NUM_INPUT_ARGS]
|
||||
# Fresh copies so the golden does not alias the builder's persistent buffers.
|
||||
fresh_outputs = [arg.clone() for arg in args[_NUM_INPUT_ARGS:]]
|
||||
|
||||
golden_kwargs = dict(captured["kwargs"])
|
||||
golden_kwargs["dtype_q"] = EXPECTED_Q_DTYPE
|
||||
golden_kwargs["dtype_kv"] = EXPECTED_KV_DTYPE
|
||||
|
||||
aiter.get_mla_metadata_v1(*inputs, *fresh_outputs, **golden_kwargs)
|
||||
|
||||
return dict(zip(_OUTPUT_ARG_FIELDS, fresh_outputs))
|
||||
|
||||
|
||||
def test_persistent_decode_metadata_matches_fp8_golden():
|
||||
"""The builder's metadata must match the dtype-correct golden.
|
||||
|
||||
Regression guard: the fixed builder forwards fp8/bf16 dtypes so its
|
||||
split/reduce metadata matches the golden recomputed with those explicit
|
||||
dtypes. Dropping the dtypes (the original bug) produces a different layout
|
||||
and fails this test.
|
||||
"""
|
||||
metadata, captured = _build_decode_metadata()
|
||||
|
||||
# qlen=1 must take the persistent-metadata path for this to be meaningful.
|
||||
assert metadata.decode is not None
|
||||
assert metadata.decode.has_persistent_metadata
|
||||
assert metadata.work_meta_data is not None
|
||||
|
||||
golden = _compute_golden_metadata(captured)
|
||||
|
||||
mismatched = [
|
||||
name
|
||||
for name in _CONTENT_METADATA_FIELDS
|
||||
if getattr(metadata, name).shape != golden[name].shape
|
||||
or not torch.equal(getattr(metadata, name), golden[name])
|
||||
]
|
||||
assert not mismatched, (
|
||||
"AITER MLA persistent decode metadata does not match the fp8/bf16 "
|
||||
f"golden for fields {mismatched}; the builder must forward "
|
||||
"dtype_q/dtype_kv to get_mla_metadata_v1."
|
||||
)
|
||||
@@ -0,0 +1,125 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import sys
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_rocm():
|
||||
pytest.skip(
|
||||
"ROCm AITER sparse MLA metadata sync test requires ROCm.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from vllm._aiter_ops import is_aiter_found_and_supported
|
||||
|
||||
if not is_aiter_found_and_supported():
|
||||
pytest.skip(
|
||||
"ROCm AITER sparse MLA metadata sync test requires a supported AITER "
|
||||
"installation.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from vllm.v1.attention.backend import CommonAttentionMetadata
|
||||
from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod
|
||||
|
||||
|
||||
class _FakeAiter(ModuleType):
|
||||
get_mla_metadata_v1: Mock
|
||||
|
||||
|
||||
def _make_builder():
|
||||
builder = object.__new__(sparse_mod.ROCMAiterMLASparseMetadataBuilder)
|
||||
max_num_batched_tokens = 8
|
||||
topk_tokens = 4
|
||||
|
||||
builder.device = torch.device("cpu")
|
||||
builder.kv_cache_spec = SimpleNamespace(block_size=1)
|
||||
builder.model_dtype = torch.bfloat16
|
||||
builder.topk_tokens = topk_tokens
|
||||
builder.req_id_per_token_buffer = torch.zeros(
|
||||
max_num_batched_tokens, dtype=torch.int32, device="cpu"
|
||||
)
|
||||
builder.qo_indptr = torch.arange(
|
||||
max_num_batched_tokens + 1, dtype=torch.int32, device="cpu"
|
||||
)
|
||||
builder.paged_kv_last_page_len = torch.ones(
|
||||
max_num_batched_tokens, dtype=torch.int32, device="cpu"
|
||||
)
|
||||
builder.paged_kv_indices = torch.zeros(
|
||||
max_num_batched_tokens * topk_tokens, dtype=torch.int32, device="cpu"
|
||||
)
|
||||
builder.paged_kv_indptr = torch.zeros(
|
||||
max_num_batched_tokens + 1, dtype=torch.int32, device="cpu"
|
||||
)
|
||||
builder._num_attention_heads = 16
|
||||
builder._mla_work_meta_data = torch.empty(1, dtype=torch.int32, device="cpu")
|
||||
builder._mla_work_indptr = torch.empty(1, dtype=torch.int32, device="cpu")
|
||||
builder._mla_work_info_set = torch.empty(1, dtype=torch.int32, device="cpu")
|
||||
builder._mla_reduce_indptr = torch.empty(1, dtype=torch.int32, device="cpu")
|
||||
builder._mla_reduce_final_map = torch.empty(1, dtype=torch.int32, device="cpu")
|
||||
builder._mla_reduce_partial_map = torch.empty(1, dtype=torch.int32, device="cpu")
|
||||
builder._prev_req_extent = 0
|
||||
builder._prev_indices_extent = 0
|
||||
builder._prev_metadata_key = None
|
||||
return builder
|
||||
|
||||
|
||||
def _make_common_metadata():
|
||||
query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32, device="cpu")
|
||||
seq_lens = torch.tensor([16, 8], dtype=torch.int32, device="cpu")
|
||||
return CommonAttentionMetadata(
|
||||
query_start_loc=query_start_loc,
|
||||
query_start_loc_cpu=query_start_loc,
|
||||
seq_lens=seq_lens,
|
||||
_seq_lens_cpu=seq_lens,
|
||||
num_reqs=2,
|
||||
num_actual_tokens=2,
|
||||
max_query_len=1,
|
||||
max_seq_len=16,
|
||||
block_table_tensor=torch.arange(16, dtype=torch.int32, device="cpu").view(2, 8),
|
||||
slot_mapping=torch.arange(2, dtype=torch.int64, device="cpu"),
|
||||
)
|
||||
|
||||
|
||||
def test_sparse_persistent_metadata_syncs_only_after_recompute(monkeypatch):
|
||||
builder = _make_builder()
|
||||
common_metadata = _make_common_metadata()
|
||||
events: list[str] = []
|
||||
|
||||
def fake_generate_sparse_seqlen_triton(*args, **kwargs):
|
||||
return torch.tensor([1, 2], dtype=torch.int32, device="cpu")
|
||||
|
||||
fake_aiter = _FakeAiter("aiter")
|
||||
|
||||
def fake_get_mla_metadata_v1(*args, **kwargs):
|
||||
events.append("metadata")
|
||||
|
||||
fake_get_mla_metadata_v1_mock = Mock(side_effect=fake_get_mla_metadata_v1)
|
||||
fake_aiter.get_mla_metadata_v1 = fake_get_mla_metadata_v1_mock
|
||||
monkeypatch.setitem(sys.modules, "aiter", fake_aiter)
|
||||
monkeypatch.setattr(
|
||||
sparse_mod, "generate_sparse_seqlen_triton", fake_generate_sparse_seqlen_triton
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
sparse_mod.torch.cuda,
|
||||
"current_stream",
|
||||
lambda device=None: SimpleNamespace(synchronize=lambda: events.append("sync")),
|
||||
)
|
||||
|
||||
builder.build(common_prefix_len=0, common_attn_metadata=common_metadata)
|
||||
|
||||
assert events == ["metadata", "sync"]
|
||||
assert fake_get_mla_metadata_v1_mock.call_count == 1
|
||||
|
||||
events.clear()
|
||||
|
||||
builder.build(common_prefix_len=0, common_attn_metadata=common_metadata)
|
||||
|
||||
assert events == []
|
||||
assert fake_get_mla_metadata_v1_mock.call_count == 1
|
||||
@@ -0,0 +1,339 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""ROCm kernel correctness tests for AITER unified attention.
|
||||
|
||||
Compares ``aiter.ops.triton.unified_attention`` against ``ref_paged_attn`` under
|
||||
decode, prefill, and mixed batches with varied shapes.
|
||||
"""
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.attention.test_triton_unified_attention import ref_paged_attn
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
_SKIP_NON_MI3XX = True
|
||||
if current_platform.is_rocm():
|
||||
from vllm.platforms.rocm import on_mi3xx
|
||||
|
||||
_SKIP_NON_MI3XX = not on_mi3xx()
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"),
|
||||
pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"),
|
||||
]
|
||||
|
||||
NUM_Q_HEADS = 8
|
||||
NUM_KV_HEADS = 8
|
||||
HEAD_SIZES = [128, 256]
|
||||
BLOCK_SIZES = [16, 64]
|
||||
DTYPES = [torch.bfloat16, torch.float16]
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
# (query_len, kv_len) per sequence
|
||||
MIXED_SEQ_LENS = [
|
||||
[(1, 128), (5, 18), (129, 463)],
|
||||
[(10, 256), (5, 64), (32, 128)],
|
||||
[(1, 1024), (5, 18), (129, 1328)],
|
||||
]
|
||||
DECODE_SEQ_LENS = [
|
||||
[(1, 128), (1, 256), (1, 384), (1, 512)],
|
||||
[(1, 1024), (1, 1536), (1, 2048)],
|
||||
]
|
||||
PREFILL_SEQ_LENS = [
|
||||
[(256, 256), (128, 512)],
|
||||
[(64, 128), (32, 256), (16, 512)],
|
||||
[(256, 1024), (128, 2048)],
|
||||
]
|
||||
|
||||
DEFAULT_ATOL, DEFAULT_RTOL = 1.5e-2, 1e-2
|
||||
FP8_ATOL, FP8_RTOL = 1.5e-1, 1.5e-1
|
||||
# Non-unity scale so q_descale handling is exercised explicitly.
|
||||
Q_SCALE = 0.75
|
||||
K_SCALE, V_SCALE = 0.5, 0.25
|
||||
|
||||
Fp8Variant = Literal["fp8_kv", "fp8_query", "fp8_query_kv"]
|
||||
|
||||
FP8_VARIANTS = [
|
||||
pytest.param("fp8_kv", id="fp8_kv"),
|
||||
pytest.param("fp8_query", id="fp8_query"),
|
||||
pytest.param("fp8_query_kv", id="fp8_query_kv"),
|
||||
]
|
||||
|
||||
FP8_SEQ_LENS = [
|
||||
MIXED_SEQ_LENS[0],
|
||||
DECODE_SEQ_LENS[0],
|
||||
DECODE_SEQ_LENS[1],
|
||||
PREFILL_SEQ_LENS[0],
|
||||
PREFILL_SEQ_LENS[2],
|
||||
]
|
||||
|
||||
|
||||
def _require_aiter() -> None:
|
||||
from vllm._aiter_ops import is_aiter_found_and_supported
|
||||
|
||||
if not is_aiter_found_and_supported():
|
||||
pytest.skip("aiter is required on supported ROCm hardware for this test")
|
||||
|
||||
|
||||
def _make_case(
|
||||
*,
|
||||
seq_lens: list[tuple[int, int]],
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
dtype: torch.dtype,
|
||||
num_blocks: int = 2048,
|
||||
kv_cache_dtype: torch.dtype | None = None,
|
||||
k_scale: float = 1.0,
|
||||
v_scale: float = 1.0,
|
||||
q_dtype: torch.dtype | None = None,
|
||||
q_scale: float = Q_SCALE,
|
||||
) -> dict[str, Any]:
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
query_lens = [q for q, _ in seq_lens]
|
||||
kv_lens = [k for _, k in seq_lens]
|
||||
num_seqs = len(seq_lens)
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
scale = head_size**-0.5
|
||||
|
||||
query = torch.randn(sum(query_lens), NUM_Q_HEADS, head_size, dtype=dtype)
|
||||
if kv_cache_dtype is None:
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, NUM_KV_HEADS, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
else:
|
||||
key_cache = torch.clamp(
|
||||
torch.randn(num_blocks, block_size, NUM_KV_HEADS, head_size),
|
||||
-1.0,
|
||||
1.0,
|
||||
).to(kv_cache_dtype)
|
||||
value_cache = torch.clamp(
|
||||
torch.randn(num_blocks, block_size, NUM_KV_HEADS, head_size),
|
||||
-1.0,
|
||||
1.0,
|
||||
).to(kv_cache_dtype)
|
||||
|
||||
cu_seqlens_q = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
seq_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
|
||||
max_num_blocks = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks), dtype=torch.int32
|
||||
)
|
||||
|
||||
descale_shape = (num_seqs, NUM_KV_HEADS)
|
||||
k_descale = torch.full(descale_shape, k_scale, dtype=torch.float32, device="cuda")
|
||||
v_descale = torch.full(descale_shape, v_scale, dtype=torch.float32, device="cuda")
|
||||
|
||||
kernel_query = query
|
||||
q_descale = None
|
||||
if q_dtype is not None:
|
||||
q_descale = torch.tensor(q_scale, dtype=torch.float32, device="cuda")
|
||||
kernel_query = (query / q_scale).to(q_dtype)
|
||||
|
||||
return {
|
||||
"query": query,
|
||||
"kernel_query": kernel_query,
|
||||
"key_cache": key_cache,
|
||||
"value_cache": value_cache,
|
||||
"block_tables": block_tables,
|
||||
"query_lens": query_lens,
|
||||
"kv_lens": kv_lens,
|
||||
"seq_lens_tensor": seq_lens_tensor,
|
||||
"cu_seqlens_q": cu_seqlens_q,
|
||||
"q_descale": q_descale,
|
||||
"k_descale": k_descale,
|
||||
"v_descale": v_descale,
|
||||
"scale": scale,
|
||||
"max_query_len": max_query_len,
|
||||
"max_kv_len": max_kv_len,
|
||||
"query_dtype": dtype,
|
||||
"k_scale": k_scale,
|
||||
"v_scale": v_scale,
|
||||
}
|
||||
|
||||
|
||||
def _make_fp8_case(
|
||||
*,
|
||||
seq_lens: list[tuple[int, int]],
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
variant: Fp8Variant,
|
||||
) -> dict[str, Any]:
|
||||
use_fp8_kv = variant in ("fp8_kv", "fp8_query_kv")
|
||||
use_fp8_query = variant in ("fp8_query", "fp8_query_kv")
|
||||
return _make_case(
|
||||
seq_lens=seq_lens,
|
||||
head_size=head_size,
|
||||
block_size=block_size,
|
||||
dtype=torch.bfloat16,
|
||||
kv_cache_dtype=FP8_DTYPE if use_fp8_kv else None,
|
||||
k_scale=K_SCALE if use_fp8_kv else 1.0,
|
||||
v_scale=V_SCALE if use_fp8_kv else 1.0,
|
||||
q_dtype=FP8_DTYPE if use_fp8_query else None,
|
||||
)
|
||||
|
||||
|
||||
def _run_aiter_unified_attention(case: dict[str, Any]) -> torch.Tensor:
|
||||
from aiter.ops.triton.unified_attention import unified_attention
|
||||
|
||||
kernel_query = case["kernel_query"]
|
||||
# Kernel writes high-precision output even when Q is FP8 (matches vLLM usage).
|
||||
output = torch.empty_like(case["query"])
|
||||
unified_attention(
|
||||
q=kernel_query,
|
||||
k=case["key_cache"],
|
||||
v=case["value_cache"],
|
||||
out=output,
|
||||
cu_seqlens_q=case["cu_seqlens_q"],
|
||||
max_seqlen_q=case["max_query_len"],
|
||||
seqused_k=case["seq_lens_tensor"],
|
||||
max_seqlen_k=case["max_kv_len"],
|
||||
softmax_scale=case["scale"],
|
||||
causal=True,
|
||||
alibi_slopes=None,
|
||||
window_size=(-1, -1),
|
||||
block_table=case["block_tables"],
|
||||
softcap=0,
|
||||
q_descale=case["q_descale"],
|
||||
k_descale=case["k_descale"],
|
||||
v_descale=case["v_descale"],
|
||||
sinks=None,
|
||||
output_scale=None,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def _ref_output(case: dict[str, Any]) -> torch.Tensor:
|
||||
key_cache = case["key_cache"]
|
||||
value_cache = case["value_cache"]
|
||||
if key_cache.dtype != case["query_dtype"]:
|
||||
key_cache = key_cache.to(case["query_dtype"]) * case["k_scale"]
|
||||
value_cache = value_cache.to(case["query_dtype"]) * case["v_scale"]
|
||||
|
||||
return ref_paged_attn(
|
||||
query=case["query"],
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=case["query_lens"],
|
||||
kv_lens=case["kv_lens"],
|
||||
block_tables=case["block_tables"],
|
||||
scale=case["scale"],
|
||||
)
|
||||
|
||||
|
||||
def _assert_matches_reference(
|
||||
case: dict[str, Any],
|
||||
*,
|
||||
atol: float = DEFAULT_ATOL,
|
||||
rtol: float = DEFAULT_RTOL,
|
||||
) -> None:
|
||||
output = _run_aiter_unified_attention(case)
|
||||
output_ref = _ref_output(case)
|
||||
torch.testing.assert_close(output, output_ref, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq_lens", MIXED_SEQ_LENS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_aiter_unified_attn_mixed_batch(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
"""Decode + prefill sequences in one batch (native dtypes)."""
|
||||
_require_aiter()
|
||||
set_random_seed(0)
|
||||
|
||||
case = _make_case(
|
||||
seq_lens=seq_lens,
|
||||
head_size=head_size,
|
||||
block_size=block_size,
|
||||
dtype=dtype,
|
||||
)
|
||||
_assert_matches_reference(case)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq_lens", DECODE_SEQ_LENS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@torch.inference_mode()
|
||||
def test_aiter_unified_attn_decode(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
"""Single-token decode (native dtypes)."""
|
||||
_require_aiter()
|
||||
set_random_seed(0)
|
||||
|
||||
case = _make_case(
|
||||
seq_lens=seq_lens,
|
||||
head_size=head_size,
|
||||
block_size=block_size,
|
||||
dtype=dtype,
|
||||
)
|
||||
_assert_matches_reference(case)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS)
|
||||
@pytest.mark.parametrize("head_size", [128])
|
||||
@pytest.mark.parametrize("block_size", [16])
|
||||
@torch.inference_mode()
|
||||
def test_aiter_unified_attn_prefill(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
) -> None:
|
||||
"""Prefill-only batches with query_len > 1 (native dtypes)."""
|
||||
_require_aiter()
|
||||
set_random_seed(0)
|
||||
|
||||
case = _make_case(
|
||||
seq_lens=seq_lens,
|
||||
head_size=head_size,
|
||||
block_size=block_size,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
_assert_matches_reference(case)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.supports_fp8(),
|
||||
reason="FP8 not supported on this hardware",
|
||||
)
|
||||
@pytest.mark.parametrize("variant", FP8_VARIANTS)
|
||||
@pytest.mark.parametrize("seq_lens", FP8_SEQ_LENS)
|
||||
@pytest.mark.parametrize("head_size", [128])
|
||||
@pytest.mark.parametrize("block_size", [16, 64])
|
||||
@torch.inference_mode()
|
||||
def test_aiter_unified_attn_fp8(
|
||||
variant: Fp8Variant,
|
||||
seq_lens: list[tuple[int, int]],
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
) -> None:
|
||||
"""FP8 KV cache, FP8 query, or both; compared at bf16 reference precision."""
|
||||
_require_aiter()
|
||||
set_random_seed(0)
|
||||
|
||||
case = _make_fp8_case(
|
||||
seq_lens=seq_lens,
|
||||
head_size=head_size,
|
||||
block_size=block_size,
|
||||
variant=variant,
|
||||
)
|
||||
_assert_matches_reference(case, atol=FP8_ATOL, rtol=FP8_RTOL)
|
||||
@@ -0,0 +1,73 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config import AttentionConfig, VllmConfig, set_current_vllm_config
|
||||
from vllm.platforms.rocm import RocmPlatform
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.v1.attention.selector import _cached_get_attn_backend, get_attn_backend
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_cache():
|
||||
"""Clear lru cache to ensure each test case runs without caching."""
|
||||
_cached_get_attn_backend.cache_clear()
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Skipped for now. Should be revisited.")
|
||||
def test_selector(monkeypatch: pytest.MonkeyPatch):
|
||||
# Set the current platform to ROCm using monkeypatch
|
||||
monkeypatch.setattr("vllm.v1.attention.selector.current_platform", RocmPlatform())
|
||||
|
||||
# Test standard ROCm attention
|
||||
attention_config = AttentionConfig(backend=AttentionBackendEnum.ROCM_ATTN)
|
||||
vllm_config = VllmConfig(attention_config=attention_config)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
backend = get_attn_backend(16, torch.float16, torch.float16, 16, False)
|
||||
assert backend.get_name() == "ROCM_FLASH" or backend.get_name() == "TRITON_ATTN"
|
||||
|
||||
# MLA test for deepseek related
|
||||
# Change the attention backend to triton MLA
|
||||
attention_config = AttentionConfig(backend=AttentionBackendEnum.TRITON_MLA)
|
||||
vllm_config = VllmConfig(attention_config=attention_config)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False, use_mla=True)
|
||||
assert backend.get_name() == "TRITON_MLA"
|
||||
|
||||
# If attention backend is None
|
||||
# If use_mla is true
|
||||
# The selected backend is triton MLA
|
||||
attention_config = AttentionConfig(backend=None)
|
||||
vllm_config = VllmConfig(attention_config=attention_config)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False, use_mla=True)
|
||||
assert backend.get_name() == "TRITON_MLA"
|
||||
|
||||
# Change the attention backend to AITER MLA
|
||||
attention_config = AttentionConfig(backend=AttentionBackendEnum.ROCM_AITER_MLA)
|
||||
vllm_config = VllmConfig(attention_config=attention_config)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
backend = get_attn_backend(576, torch.bfloat16, "auto", 1, False, use_mla=True)
|
||||
assert backend.get_name() == "ROCM_AITER_MLA"
|
||||
|
||||
# If attention backend is None
|
||||
# If use_mla is true
|
||||
# If VLLM_ROCM_USE_AITER is enabled
|
||||
# The selected backend is ROCM_AITER_MLA
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||
|
||||
attention_config = AttentionConfig(backend=None)
|
||||
vllm_config = VllmConfig(attention_config=attention_config)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
backend = get_attn_backend(
|
||||
576, torch.bfloat16, "auto", 1, False, use_mla=True
|
||||
)
|
||||
assert backend.get_name() == "ROCM_AITER_MLA"
|
||||
@@ -0,0 +1,651 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_rocm(), reason="Only used by ROCm"
|
||||
)
|
||||
|
||||
|
||||
def _on_gfx950() -> bool:
|
||||
if not current_platform.is_rocm():
|
||||
return False
|
||||
try:
|
||||
from vllm.platforms.rocm import _ON_GFX950
|
||||
|
||||
return bool(_ON_GFX950)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
# The flash-decode split-K decode path is only tuned for AMD gfx950; other
|
||||
# architectures take the fallback decode kernel, so its tests are skipped there.
|
||||
requires_gfx950 = pytest.mark.skipif(
|
||||
not _on_gfx950(),
|
||||
reason="split-K decode kernel is only tuned for AMD gfx950",
|
||||
)
|
||||
|
||||
NOPE_HEAD_DIM = 448
|
||||
ROPE_HEAD_DIM = 64
|
||||
HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM
|
||||
|
||||
|
||||
def _ref_global_topk_ragged(
|
||||
topk_indices: torch.Tensor,
|
||||
token_to_req_indices: torch.Tensor,
|
||||
block_table: torch.Tensor,
|
||||
block_size: int,
|
||||
is_valid_token: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
topk = topk_indices.reshape(topk_indices.shape[0], -1)
|
||||
valid = (topk >= 0) & is_valid_token[:, None]
|
||||
lens = valid.sum(dim=1, dtype=torch.int32)
|
||||
indptr = torch.zeros(lens.shape[0] + 1, dtype=torch.int32, device=topk.device)
|
||||
torch.cumsum(lens, dim=0, out=indptr[1:])
|
||||
|
||||
safe_topk = torch.clamp(topk, min=0)
|
||||
block_indices = safe_topk // block_size
|
||||
block_offsets = safe_topk % block_size
|
||||
req_indices = token_to_req_indices[:, None].expand_as(topk)
|
||||
slot_ids = block_table[req_indices, block_indices] * block_size + block_offsets
|
||||
|
||||
offsets = torch.arange(topk.shape[1], dtype=torch.int32, device=topk.device)
|
||||
positions = indptr[:-1, None] + offsets[None, :]
|
||||
return slot_ids[valid], positions[valid].to(torch.long), indptr, lens
|
||||
|
||||
|
||||
def _ref_sparse_prefill_ragged(
|
||||
q: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
rows: list[list[int]],
|
||||
scale: float,
|
||||
attn_sink: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
q_f32 = q.float()
|
||||
kv_f32 = kv.float()
|
||||
out = torch.empty_like(q_f32)
|
||||
|
||||
for query_idx in range(q.shape[0]):
|
||||
row_indices = rows[query_idx]
|
||||
for head_idx in range(q.shape[1]):
|
||||
if row_indices:
|
||||
selected_kv = kv_f32[row_indices]
|
||||
scores = torch.mv(selected_kv, q_f32[query_idx, head_idx]) * scale
|
||||
if attn_sink is not None:
|
||||
scores_with_sink = torch.cat(
|
||||
[scores, attn_sink[head_idx].float().reshape(1)]
|
||||
)
|
||||
probs = torch.softmax(scores_with_sink, dim=0)[:-1]
|
||||
else:
|
||||
probs = torch.softmax(scores, dim=0)
|
||||
out[query_idx, head_idx] = torch.sum(
|
||||
probs[:, None] * selected_kv, dim=0
|
||||
)
|
||||
else:
|
||||
out[query_idx, head_idx] = 0
|
||||
return out.to(torch.bfloat16)
|
||||
|
||||
|
||||
def _pack_fp8_ds_mla_cache(
|
||||
kv: torch.Tensor, block_size: int, is_extra: bool = False
|
||||
) -> torch.Tensor:
|
||||
assert kv.shape[-1] == HEAD_DIM
|
||||
num_tokens = kv.shape[0]
|
||||
num_blocks = (num_tokens + block_size - 1) // block_size
|
||||
cache = torch.zeros(
|
||||
(num_blocks, block_size, 584),
|
||||
dtype=torch.uint8,
|
||||
device=kv.device,
|
||||
)
|
||||
cache_flat = cache.view(torch.uint8).flatten()
|
||||
kv_nope_fp8 = (
|
||||
kv[:, :NOPE_HEAD_DIM]
|
||||
.to(torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype())
|
||||
.view(torch.uint8)
|
||||
)
|
||||
kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8)
|
||||
|
||||
for slot in range(num_tokens):
|
||||
block_idx = slot // block_size
|
||||
pos = slot % block_size
|
||||
block_base = block_idx * cache.stride(0)
|
||||
token_base = block_base + pos * 576
|
||||
scale_base = block_base + block_size * 576 + pos * 8
|
||||
cache_flat[token_base : token_base + NOPE_HEAD_DIM].copy_(kv_nope_fp8[slot])
|
||||
cache_flat[
|
||||
token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
|
||||
].copy_(kv_rope_u8[slot])
|
||||
cache_flat[scale_base : scale_base + 7].fill_(127)
|
||||
return cache
|
||||
|
||||
|
||||
def _read_fp8_ds_mla_cache(
|
||||
cache: torch.Tensor, slot: int, block_size: int, is_extra: bool = False
|
||||
) -> torch.Tensor:
|
||||
cache_flat = cache.view(torch.uint8).flatten()
|
||||
block_idx = slot // block_size
|
||||
pos = slot % block_size
|
||||
block_base = block_idx * cache.stride(0)
|
||||
token_base = block_base + pos * 576
|
||||
|
||||
nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM]
|
||||
nope = nope_u8.view(
|
||||
torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype()
|
||||
).to(torch.float32)
|
||||
rope_u8 = cache_flat[
|
||||
token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
|
||||
]
|
||||
rope = rope_u8.view(torch.bfloat16).to(torch.float32)
|
||||
return torch.cat([nope, rope])
|
||||
|
||||
|
||||
def _ref_sparse_decode_ragged(
|
||||
q: torch.Tensor,
|
||||
main_cache: torch.Tensor,
|
||||
main_rows: list[list[int]],
|
||||
scale: float,
|
||||
attn_sink: torch.Tensor | None,
|
||||
block_size: int,
|
||||
extra_cache: torch.Tensor | None = None,
|
||||
extra_rows: list[list[int]] | None = None,
|
||||
) -> torch.Tensor:
|
||||
q_f32 = q.float()
|
||||
out = torch.empty_like(q_f32)
|
||||
|
||||
for query_idx in range(q.shape[0]):
|
||||
row_kv = [
|
||||
_read_fp8_ds_mla_cache(main_cache, int(slot), block_size)
|
||||
for slot in main_rows[query_idx]
|
||||
]
|
||||
if extra_cache is not None and extra_rows is not None:
|
||||
row_kv.extend(
|
||||
_read_fp8_ds_mla_cache(
|
||||
extra_cache, int(slot), block_size, is_extra=True
|
||||
)
|
||||
for slot in extra_rows[query_idx]
|
||||
)
|
||||
|
||||
kv = torch.stack(row_kv).to(q.device)
|
||||
for head_idx in range(q.shape[1]):
|
||||
scores = torch.mv(kv, q_f32[query_idx, head_idx]) * scale
|
||||
if attn_sink is not None:
|
||||
scores_with_sink = torch.cat(
|
||||
[scores, attn_sink[head_idx].float().reshape(1)]
|
||||
)
|
||||
probs = torch.softmax(scores_with_sink, dim=0)[:-1]
|
||||
else:
|
||||
probs = torch.softmax(scores, dim=0)
|
||||
out[query_idx, head_idx] = torch.sum(probs[:, None] * kv, dim=0)
|
||||
return out.to(torch.bfloat16)
|
||||
|
||||
|
||||
def _ragged_from_rows(
|
||||
rows: list[list[int]], device: torch.device
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Flatten per-query slot lists into ragged (indices, indptr) tensors."""
|
||||
flat = [slot for row in rows for slot in row]
|
||||
indptr = [0]
|
||||
for row in rows:
|
||||
indptr.append(indptr[-1] + len(row))
|
||||
return (
|
||||
torch.tensor(flat, dtype=torch.int32, device=device),
|
||||
torch.tensor(indptr, dtype=torch.int32, device=device),
|
||||
)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_compute_global_topk_ragged_indices_and_indptr() -> None:
|
||||
from vllm.models.deepseek_v4.amd.rocm import (
|
||||
compute_global_topk_ragged_indices_and_indptr,
|
||||
)
|
||||
|
||||
device = torch.device("cuda")
|
||||
block_size = 4
|
||||
topk_indices = torch.tensor(
|
||||
[
|
||||
[0, 3, 4, -1],
|
||||
[5, 8, -1, -1],
|
||||
[2, 7, 9, -1],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
token_to_req_indices = torch.tensor([0, 1, 1], dtype=torch.int32, device=device)
|
||||
block_table = torch.tensor(
|
||||
[
|
||||
[10, 11, 12],
|
||||
[20, 21, 22],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
is_valid_token = torch.tensor([True, False, True], dtype=torch.bool, device=device)
|
||||
|
||||
actual_ragged, actual_indptr, actual_lens = (
|
||||
compute_global_topk_ragged_indices_and_indptr(
|
||||
topk_indices,
|
||||
token_to_req_indices,
|
||||
block_table,
|
||||
block_size,
|
||||
is_valid_token,
|
||||
)
|
||||
)
|
||||
expected_values, expected_positions, expected_indptr, expected_lens = (
|
||||
_ref_global_topk_ragged(
|
||||
topk_indices,
|
||||
token_to_req_indices,
|
||||
block_table,
|
||||
block_size,
|
||||
is_valid_token,
|
||||
)
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual_ragged[expected_positions], expected_values)
|
||||
torch.testing.assert_close(actual_indptr, expected_indptr)
|
||||
torch.testing.assert_close(actual_lens, expected_lens)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_sparse_attn_prefill_ragged_kernel() -> None:
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
||||
_rocm_sparse_attn_prefill_ragged_triton,
|
||||
)
|
||||
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(0)
|
||||
q = torch.randn(3, 3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
indices = torch.tensor([0, 2, 1, 3, 4], dtype=torch.int32, device=device)
|
||||
indptr = torch.tensor([0, 2, 5, 5], dtype=torch.int32, device=device)
|
||||
attn_sink = torch.tensor([-0.25, 0.0, 0.25], dtype=torch.float32, device=device)
|
||||
scale = HEAD_DIM**-0.5
|
||||
|
||||
actual = _rocm_sparse_attn_prefill_ragged_triton(
|
||||
q=q,
|
||||
kv=kv,
|
||||
indices=indices,
|
||||
indptr=indptr,
|
||||
scale=scale,
|
||||
attn_sink=attn_sink,
|
||||
nope_head_dim=NOPE_HEAD_DIM,
|
||||
rope_head_dim=ROPE_HEAD_DIM,
|
||||
)
|
||||
expected = _ref_sparse_prefill_ragged(
|
||||
q, kv, [[0, 2], [1, 3, 4], []], scale, attn_sink
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_sparse_attn_decode_ragged_kernel() -> None:
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
||||
_rocm_sparse_attn_decode_ragged_triton,
|
||||
)
|
||||
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(1)
|
||||
block_size = 4
|
||||
q = torch.randn(2, 3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
|
||||
extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
|
||||
main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device)
|
||||
main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device)
|
||||
extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device)
|
||||
extra_indptr = torch.tensor([0, 1, 3], dtype=torch.int32, device=device)
|
||||
attn_sink = torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device)
|
||||
scale = HEAD_DIM**-0.5
|
||||
|
||||
actual = _rocm_sparse_attn_decode_ragged_triton(
|
||||
q=q,
|
||||
main_cache=main_cache,
|
||||
main_indices=main_indices,
|
||||
main_indptr=main_indptr,
|
||||
scale=scale,
|
||||
attn_sink=attn_sink,
|
||||
nope_head_dim=NOPE_HEAD_DIM,
|
||||
rope_head_dim=ROPE_HEAD_DIM,
|
||||
extra_cache=extra_cache,
|
||||
extra_indices=extra_indices,
|
||||
extra_indptr=extra_indptr,
|
||||
)
|
||||
expected = _ref_sparse_decode_ragged(
|
||||
q=q,
|
||||
main_cache=main_cache,
|
||||
main_rows=[[0, 2], [4, 1]],
|
||||
scale=scale,
|
||||
attn_sink=attn_sink,
|
||||
block_size=block_size,
|
||||
extra_cache=extra_cache,
|
||||
extra_rows=[[1], [3, 0]],
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
@requires_gfx950
|
||||
@torch.inference_mode()
|
||||
def test_decode_num_splits_heuristic(monkeypatch) -> None:
|
||||
"""Split-count heuristic added with the flash-decode split-K decode path."""
|
||||
from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
|
||||
|
||||
# Pin the CU count so the heuristic is deterministic off-device.
|
||||
monkeypatch.setattr(mod, "_decode_cu_count", lambda: 256)
|
||||
|
||||
# A batch that already fills the device should not be split.
|
||||
assert mod._decode_num_splits(256, 1, avg_main_len=128.0, avg_extra_len=0.0) == 1
|
||||
# A tiny batch on a large device should split to add parallelism.
|
||||
assert mod._decode_num_splits(2, 1, avg_main_len=256.0, avg_extra_len=0.0) > 1
|
||||
|
||||
# The chosen count always stays within the searched [1, 16] range, and a
|
||||
# zero-length workload never splits (no work to parallelize).
|
||||
for num_queries in (1, 4, 24, 224, 1024):
|
||||
splits = mod._decode_num_splits(
|
||||
num_queries, 1, avg_main_len=512.0, avg_extra_len=128.0
|
||||
)
|
||||
assert 1 <= splits <= 16
|
||||
assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1
|
||||
|
||||
|
||||
@requires_gfx950
|
||||
@pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8])
|
||||
@pytest.mark.parametrize("with_extra", [True, False])
|
||||
@pytest.mark.parametrize("with_sink", [True, False])
|
||||
@torch.inference_mode()
|
||||
def test_sparse_attn_decode_split_k_kernel(
|
||||
monkeypatch, num_splits: int, with_extra: bool, with_sink: bool
|
||||
) -> None:
|
||||
"""Flash-decode split-K decode path (partial + reduce kernels).
|
||||
|
||||
This path is the gfx950 production path (``_ON_GFX950``), so the test only
|
||||
runs on gfx950. The split count is pinned so the partial/reduce kernels are
|
||||
exercised across split counts. ``num_splits=8`` drives splits past the
|
||||
shortest segment length, covering the empty-split edge case handled by the
|
||||
reduce kernel.
|
||||
"""
|
||||
from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
|
||||
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(7)
|
||||
block_size = 4
|
||||
num_heads = 3
|
||||
|
||||
main_rows = [[0, 2, 4, 6, 1, 3, 7, 5], [4, 1, 6, 0, 2]]
|
||||
num_queries = len(main_rows)
|
||||
q = (
|
||||
torch.randn(
|
||||
num_queries, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
* 0.125
|
||||
)
|
||||
main_kv = torch.randn(8, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
|
||||
main_indices, main_indptr = _ragged_from_rows(main_rows, device)
|
||||
|
||||
extra_rows: list[list[int]] | None = None
|
||||
extra_cache: torch.Tensor | None = None
|
||||
extra_indices: torch.Tensor | None = None
|
||||
extra_indptr: torch.Tensor | None = None
|
||||
if with_extra:
|
||||
rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]]
|
||||
extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
extra_rows = rows
|
||||
extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
|
||||
extra_indices, extra_indptr = _ragged_from_rows(rows, device)
|
||||
|
||||
attn_sink = (
|
||||
torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device)
|
||||
if with_sink
|
||||
else None
|
||||
)
|
||||
scale = HEAD_DIM**-0.5
|
||||
|
||||
# Pin the split count so each parametrized value is exercised deterministically.
|
||||
monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: num_splits)
|
||||
|
||||
actual = mod._rocm_sparse_attn_decode_ragged_triton(
|
||||
q=q,
|
||||
main_cache=main_cache,
|
||||
main_indices=main_indices,
|
||||
main_indptr=main_indptr,
|
||||
scale=scale,
|
||||
attn_sink=attn_sink,
|
||||
nope_head_dim=NOPE_HEAD_DIM,
|
||||
rope_head_dim=ROPE_HEAD_DIM,
|
||||
extra_cache=extra_cache,
|
||||
extra_indices=extra_indices,
|
||||
extra_indptr=extra_indptr,
|
||||
)
|
||||
expected = _ref_sparse_decode_ragged(
|
||||
q=q,
|
||||
main_cache=main_cache,
|
||||
main_rows=main_rows,
|
||||
scale=scale,
|
||||
attn_sink=attn_sink,
|
||||
block_size=block_size,
|
||||
extra_cache=extra_cache,
|
||||
extra_rows=extra_rows,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# Cache rows = max_position_embeddings * scaling_factor.
|
||||
_ROTARY_MAX_POS = 1024
|
||||
_ROTARY_SCALING_FACTOR = 4.0
|
||||
_ROTARY_CACHE_LEN = int(_ROTARY_MAX_POS * _ROTARY_SCALING_FACTOR)
|
||||
|
||||
|
||||
def _make_dsv4_rotary(device: torch.device):
|
||||
"""The official DSv4 rotary embedding, sized down for unit tests."""
|
||||
from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import (
|
||||
DeepseekV4ScalingRotaryEmbedding,
|
||||
)
|
||||
|
||||
# The model loader constructs layers under a default-device context;
|
||||
# mirror that so the fp32 cos_sin_cache lands on the GPU.
|
||||
with torch.device(device):
|
||||
rotary_emb = DeepseekV4ScalingRotaryEmbedding(
|
||||
head_size=ROPE_HEAD_DIM,
|
||||
rotary_dim=ROPE_HEAD_DIM,
|
||||
max_position_embeddings=_ROTARY_MAX_POS,
|
||||
base=10000,
|
||||
is_neox_style=False,
|
||||
scaling_factor=_ROTARY_SCALING_FACTOR,
|
||||
dtype=torch.bfloat16,
|
||||
mscale=1.0,
|
||||
mscale_all_dim=1.0,
|
||||
)
|
||||
rotary_emb = rotary_emb.to(device)
|
||||
assert rotary_emb.cos_sin_cache.shape == (_ROTARY_CACHE_LEN, ROPE_HEAD_DIM)
|
||||
return rotary_emb
|
||||
|
||||
|
||||
def _inv_rope_via_rotary_native(
|
||||
rotary_emb: torch.nn.Module,
|
||||
o: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Reference: the official ``forward_native(inverse=True)`` path."""
|
||||
expected, _ = rotary_emb.forward_native(positions, o.clone(), None, inverse=True)
|
||||
return expected.to(torch.bfloat16)
|
||||
|
||||
|
||||
class _FakeWoA(torch.nn.Module):
|
||||
"""Stand-in for the wo_a linear layer holding the (optionally fp8) weight."""
|
||||
|
||||
def __init__(
|
||||
self, weight: torch.Tensor, weight_scale_inv: torch.Tensor | None = None
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.weight = weight
|
||||
if weight_scale_inv is not None:
|
||||
self.weight_scale_inv = weight_scale_inv
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 7, 64])
|
||||
@pytest.mark.parametrize("num_heads", [1, 8])
|
||||
@pytest.mark.parametrize("pos_dtype", [torch.int32, torch.int64])
|
||||
@torch.inference_mode()
|
||||
def test_fused_inverse_rope_gptj_matches_rotary_native(
|
||||
num_tokens: int, num_heads: int, pos_dtype: torch.dtype, default_vllm_config
|
||||
) -> None:
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj
|
||||
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(0)
|
||||
rotary_emb = _make_dsv4_rotary(device)
|
||||
o = torch.randn(
|
||||
num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
positions = torch.randint(
|
||||
0, _ROTARY_CACHE_LEN, (num_tokens,), dtype=pos_dtype, device=device
|
||||
)
|
||||
|
||||
actual = _fused_inverse_rope_gptj(
|
||||
o, positions, rotary_emb.cos_sin_cache, ROPE_HEAD_DIM
|
||||
)
|
||||
expected = _inv_rope_via_rotary_native(rotary_emb, o, positions)
|
||||
|
||||
assert actual.dtype == torch.bfloat16
|
||||
assert actual.shape == o.shape
|
||||
# NoPE lanes are a pure bf16 passthrough -> must be bit-exact.
|
||||
assert torch.equal(actual[..., :NOPE_HEAD_DIM], expected[..., :NOPE_HEAD_DIM])
|
||||
# RoPE lanes: tolerate at most ~1 bf16 ulp from fp32 fma ordering.
|
||||
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_fused_inverse_rope_gptj_empty(default_vllm_config) -> None:
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj
|
||||
|
||||
device = torch.device("cuda")
|
||||
rotary_emb = _make_dsv4_rotary(device)
|
||||
o = torch.empty(0, 8, HEAD_DIM, dtype=torch.bfloat16, device=device)
|
||||
positions = torch.empty(0, dtype=torch.int32, device=device)
|
||||
|
||||
out = _fused_inverse_rope_gptj(
|
||||
o, positions, rotary_emb.cos_sin_cache, ROPE_HEAD_DIM
|
||||
)
|
||||
assert out.shape == (0, 8, HEAD_DIM)
|
||||
assert out.dtype == torch.bfloat16
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_rocm_inv_rope_einsum_matches_rotary_native(default_vllm_config) -> None:
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_inv_rope_einsum
|
||||
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(2)
|
||||
num_tokens, num_heads = 5, 8
|
||||
n_local_groups = num_heads
|
||||
o_lora_rank = 16
|
||||
hidden_dim = num_heads * HEAD_DIM // n_local_groups # 512
|
||||
|
||||
rotary_emb = _make_dsv4_rotary(device)
|
||||
o = (
|
||||
torch.randn(
|
||||
num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
* 0.125
|
||||
)
|
||||
positions = torch.randint(
|
||||
0, _ROTARY_CACHE_LEN, (num_tokens,), dtype=torch.int32, device=device
|
||||
)
|
||||
weight = (
|
||||
torch.randn(n_local_groups * o_lora_rank, hidden_dim, device=device) * 0.125
|
||||
).to(torch.bfloat16)
|
||||
wo_a = _FakeWoA(weight)
|
||||
|
||||
actual = rocm_inv_rope_einsum(
|
||||
rotary_emb, o, positions, ROPE_HEAD_DIM, n_local_groups, o_lora_rank, wo_a
|
||||
)
|
||||
|
||||
o_ref = _inv_rope_via_rotary_native(rotary_emb, o, positions)
|
||||
o_ref = o_ref.view(num_tokens, n_local_groups, -1)
|
||||
wo_a_ref = weight.view(n_local_groups, o_lora_rank, hidden_dim).to(torch.bfloat16)
|
||||
expected = torch.einsum("tgd,grd->tgr", o_ref, wo_a_ref)
|
||||
|
||||
assert actual.shape == (num_tokens, n_local_groups, o_lora_rank)
|
||||
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_get_cached_wo_a_bf16_plain_caches() -> None:
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
|
||||
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(4)
|
||||
n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8
|
||||
weight = torch.randn(
|
||||
n_local_groups * o_lora_rank, hidden_dim, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
wo_a = _FakeWoA(weight)
|
||||
|
||||
out1 = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
|
||||
expected = weight.view(n_local_groups, o_lora_rank, hidden_dim).to(torch.bfloat16)
|
||||
assert out1.shape == (n_local_groups, o_lora_rank, hidden_dim)
|
||||
torch.testing.assert_close(out1, expected, atol=0, rtol=0)
|
||||
assert hasattr(wo_a, "_dsv4_wo_a_bf16")
|
||||
|
||||
# Mutate the source weight: the cached tensor must be returned unchanged
|
||||
# (proving the dequant is not recomputed per call).
|
||||
wo_a.weight.zero_()
|
||||
out2 = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
|
||||
assert out2 is out1
|
||||
torch.testing.assert_close(out2, expected, atol=0, rtol=0)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None:
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
|
||||
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(5)
|
||||
n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8
|
||||
row_block, col_block = 2, 2
|
||||
row_blocks = o_lora_rank // row_block
|
||||
col_blocks = hidden_dim // col_block
|
||||
|
||||
fp8_dtype = current_platform.fp8_dtype()
|
||||
weight_f32 = (
|
||||
torch.randn(
|
||||
n_local_groups, o_lora_rank, hidden_dim, dtype=torch.float32, device=device
|
||||
)
|
||||
* 0.1
|
||||
)
|
||||
weight_fp8 = weight_f32.to(fp8_dtype)
|
||||
scale = (
|
||||
torch.rand(
|
||||
n_local_groups, row_blocks, col_blocks, dtype=torch.float32, device=device
|
||||
)
|
||||
* 0.5
|
||||
+ 0.5
|
||||
)
|
||||
wo_a = _FakeWoA(
|
||||
weight_fp8.reshape(n_local_groups * o_lora_rank, hidden_dim),
|
||||
weight_scale_inv=scale.reshape(n_local_groups * row_blocks, col_blocks),
|
||||
)
|
||||
|
||||
out = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
|
||||
|
||||
scale_full = scale.repeat_interleave(row_block, dim=-2).repeat_interleave(
|
||||
col_block, dim=-1
|
||||
)
|
||||
expected = (weight_fp8.to(torch.float32) * scale_full).to(torch.bfloat16)
|
||||
assert out.shape == (n_local_groups, o_lora_rank, hidden_dim)
|
||||
torch.testing.assert_close(out, expected, atol=0, rtol=0)
|
||||
|
||||
# Second call returns the same cached object.
|
||||
assert _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) is out
|
||||
@@ -0,0 +1,325 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import cdiv
|
||||
from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
|
||||
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
|
||||
|
||||
@pytest.mark.parametrize("B", [3, 5])
|
||||
@pytest.mark.parametrize("L", [1027, 1025])
|
||||
@pytest.mark.parametrize("H_Q", [32])
|
||||
@pytest.mark.parametrize("H_KV", [32, 8])
|
||||
@pytest.mark.parametrize("D_QK", [128, 192, 576])
|
||||
@pytest.mark.parametrize("D_V", [128, 512])
|
||||
@pytest.mark.parametrize("CACHE_SIZE", [16384])
|
||||
@pytest.mark.parametrize("PAGE_SIZE", [1, 16])
|
||||
def test_decode_attention(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):
|
||||
assert CACHE_SIZE % PAGE_SIZE == 0
|
||||
dtype = torch.bfloat16
|
||||
seq_len = L # This represents the number of tokens already in the sequence
|
||||
sm_scale = 1.0 / (D_QK**0.5)
|
||||
num_kv_splits = 8
|
||||
|
||||
num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
|
||||
req_to_page = torch.randint(
|
||||
0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device=DEVICE_TYPE
|
||||
)
|
||||
req_to_token = req_to_page * PAGE_SIZE
|
||||
req_to_token = req_to_token.expand(B, num_pages_per_batch, PAGE_SIZE)
|
||||
req_to_token = req_to_token + torch.arange(PAGE_SIZE, device=DEVICE_TYPE).view(
|
||||
1, 1, -1
|
||||
)
|
||||
req_to_token = req_to_token.view(B, -1)
|
||||
req_to_token = req_to_token[:, :seq_len].contiguous()
|
||||
|
||||
# q represents the new token being generated, one per batch
|
||||
q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
|
||||
|
||||
# k_buffer and v_buffer represent all previous tokens
|
||||
# Page size is 1.
|
||||
k_buffer = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE)
|
||||
v_buffer = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE)
|
||||
|
||||
# o will have the same shape as q
|
||||
o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
|
||||
|
||||
lse = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
|
||||
|
||||
b_seq_len = torch.full((B,), seq_len, device=DEVICE_TYPE)
|
||||
|
||||
attn_logits = torch.empty(
|
||||
(B, H_Q, num_kv_splits, D_V + 1),
|
||||
dtype=torch.float32,
|
||||
device=DEVICE_TYPE,
|
||||
)
|
||||
|
||||
# Call the original implementation.
|
||||
decode_attention_fwd(
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
o,
|
||||
lse,
|
||||
req_to_token,
|
||||
b_seq_len,
|
||||
attn_logits,
|
||||
num_kv_splits,
|
||||
sm_scale,
|
||||
)
|
||||
|
||||
# Page size can be larger than 1.
|
||||
k_buffer = k_buffer.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_QK)
|
||||
v_buffer = v_buffer.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_V)
|
||||
|
||||
o1 = torch.zeros_like(o)
|
||||
lse1 = torch.zeros_like(lse)
|
||||
|
||||
decode_attention_fwd(
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
o1,
|
||||
lse1,
|
||||
req_to_page,
|
||||
b_seq_len,
|
||||
attn_logits,
|
||||
num_kv_splits,
|
||||
sm_scale,
|
||||
PAGE_SIZE,
|
||||
)
|
||||
|
||||
assert torch.allclose(o, o1)
|
||||
|
||||
|
||||
def _quantize_to_fp8(tensor: torch.Tensor):
|
||||
"""Quantize a BF16 tensor to FP8 e4m3fn with per-tensor scale.
|
||||
|
||||
Returns (fp8_tensor, scale) where:
|
||||
fp8_tensor ≈ tensor / scale (stored as float8_e4m3fn)
|
||||
tensor ≈ fp8_tensor.to(float32) * scale (dequantized)
|
||||
"""
|
||||
amax = tensor.abs().amax()
|
||||
# float8_e4m3fn max representable value is 448.0
|
||||
scale = (amax / 448.0).clamp(min=1e-12).to(torch.float32)
|
||||
fp8_tensor = (
|
||||
(tensor.to(torch.float32) / scale).clamp(-448.0, 448.0).to(torch.float8_e4m3fn)
|
||||
)
|
||||
return fp8_tensor, scale
|
||||
|
||||
|
||||
@pytest.mark.parametrize("B", [3])
|
||||
@pytest.mark.parametrize("L", [1025])
|
||||
@pytest.mark.parametrize("H_Q", [32])
|
||||
@pytest.mark.parametrize("H_KV", [32, 8])
|
||||
@pytest.mark.parametrize("D_QK", [128, 576])
|
||||
@pytest.mark.parametrize("D_V", [128, 512])
|
||||
@pytest.mark.parametrize("CACHE_SIZE", [16384])
|
||||
@pytest.mark.parametrize("PAGE_SIZE", [1, 16])
|
||||
def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):
|
||||
"""Test FP8 KV cache path: quantize K/V to FP8, run kernel with scales,
|
||||
and compare against BF16 reference output."""
|
||||
assert CACHE_SIZE % PAGE_SIZE == 0
|
||||
dtype = torch.bfloat16
|
||||
seq_len = L
|
||||
sm_scale = 1.0 / (D_QK**0.5)
|
||||
num_kv_splits = 8
|
||||
|
||||
num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
|
||||
req_to_page = torch.randint(
|
||||
0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device=DEVICE_TYPE
|
||||
)
|
||||
req_to_token = req_to_page * PAGE_SIZE
|
||||
req_to_token = req_to_token.expand(B, num_pages_per_batch, PAGE_SIZE)
|
||||
req_to_token = req_to_token + torch.arange(PAGE_SIZE, device=DEVICE_TYPE).view(
|
||||
1, 1, -1
|
||||
)
|
||||
req_to_token = req_to_token.view(B, -1)
|
||||
req_to_token = req_to_token[:, :seq_len].contiguous()
|
||||
|
||||
q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
|
||||
|
||||
# Create BF16 K/V as reference
|
||||
k_bf16 = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE)
|
||||
v_bf16 = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE)
|
||||
|
||||
# --- BF16 reference ---
|
||||
o_ref = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
|
||||
lse_ref = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
|
||||
attn_logits = torch.empty(
|
||||
(B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
|
||||
)
|
||||
|
||||
if PAGE_SIZE == 1:
|
||||
decode_attention_fwd(
|
||||
q,
|
||||
k_bf16,
|
||||
v_bf16,
|
||||
o_ref,
|
||||
lse_ref,
|
||||
req_to_token,
|
||||
b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
|
||||
attn_logits=attn_logits,
|
||||
num_kv_splits=num_kv_splits,
|
||||
sm_scale=sm_scale,
|
||||
)
|
||||
else:
|
||||
k_paged = k_bf16.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_QK)
|
||||
v_paged = v_bf16.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_V)
|
||||
decode_attention_fwd(
|
||||
q,
|
||||
k_paged,
|
||||
v_paged,
|
||||
o_ref,
|
||||
lse_ref,
|
||||
req_to_page,
|
||||
b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
|
||||
attn_logits=attn_logits,
|
||||
num_kv_splits=num_kv_splits,
|
||||
sm_scale=sm_scale,
|
||||
page_size=PAGE_SIZE,
|
||||
)
|
||||
|
||||
# --- FP8 path ---
|
||||
k_fp8, k_scale = _quantize_to_fp8(k_bf16)
|
||||
v_fp8, v_scale = _quantize_to_fp8(v_bf16)
|
||||
|
||||
o_fp8 = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
|
||||
lse_fp8 = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
|
||||
attn_logits_fp8 = torch.empty(
|
||||
(B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
|
||||
)
|
||||
|
||||
if PAGE_SIZE == 1:
|
||||
decode_attention_fwd(
|
||||
q,
|
||||
k_fp8,
|
||||
v_fp8,
|
||||
o_fp8,
|
||||
lse_fp8,
|
||||
req_to_token,
|
||||
b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
|
||||
attn_logits=attn_logits_fp8,
|
||||
num_kv_splits=num_kv_splits,
|
||||
sm_scale=sm_scale,
|
||||
k_scale=k_scale,
|
||||
v_scale=v_scale,
|
||||
)
|
||||
else:
|
||||
k_fp8_paged = k_fp8.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_QK)
|
||||
v_fp8_paged = v_fp8.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_V)
|
||||
decode_attention_fwd(
|
||||
q,
|
||||
k_fp8_paged,
|
||||
v_fp8_paged,
|
||||
o_fp8,
|
||||
lse_fp8,
|
||||
req_to_page,
|
||||
b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
|
||||
attn_logits=attn_logits_fp8,
|
||||
num_kv_splits=num_kv_splits,
|
||||
sm_scale=sm_scale,
|
||||
page_size=PAGE_SIZE,
|
||||
k_scale=k_scale,
|
||||
v_scale=v_scale,
|
||||
)
|
||||
|
||||
# FP8 tolerances match test_mla_backends.py test_backend_correctness.
|
||||
torch.testing.assert_close(o_ref, o_fp8, atol=5e-1, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"H_Q,H_KV,D_QK,D_V,is_mla",
|
||||
[
|
||||
(16, 1, 576, 512, True), # MLA path (grouped kernel, v = trans(k))
|
||||
(32, 8, 128, 128, False), # GQA path (grouped kernel)
|
||||
(32, 32, 128, 128, False), # MHA path (normal kernel)
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("PAGE_SIZE", [16])
|
||||
def test_decode_attention_cross_layer_view(H_Q, H_KV, D_QK, D_V, is_mla, PAGE_SIZE):
|
||||
"""The kernel must honor the cache's page-dim stride, not assume pages are
|
||||
packed back-to-back. A per-layer view into a cross-layer (block-major)
|
||||
cache has stride(0) inflated by num_layers; outputs must match a
|
||||
contiguous cache holding the same data exactly."""
|
||||
B = 3
|
||||
seq_len = 1027
|
||||
CACHE_SIZE = 16384
|
||||
NUM_LAYERS = 3
|
||||
LAYER_IDX = 1
|
||||
dtype = torch.bfloat16
|
||||
sm_scale = 1.0 / (D_QK**0.5)
|
||||
num_kv_splits = 8
|
||||
num_pages = CACHE_SIZE // PAGE_SIZE
|
||||
|
||||
num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
|
||||
req_to_page = torch.randint(
|
||||
0, num_pages, (B, num_pages_per_batch), device=DEVICE_TYPE
|
||||
)
|
||||
|
||||
q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
|
||||
b_seq_len = torch.full((B,), seq_len, device=DEVICE_TYPE)
|
||||
|
||||
# Reference: contiguous paged cache.
|
||||
k_ref = torch.randn(
|
||||
num_pages, PAGE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE
|
||||
)
|
||||
if is_mla:
|
||||
v_ref = k_ref[..., :D_V]
|
||||
else:
|
||||
v_ref = torch.randn(
|
||||
num_pages, PAGE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE
|
||||
)
|
||||
|
||||
# Cross-layer cache: all layers' pages for a block are adjacent. The
|
||||
# per-layer view has the same shape as the contiguous cache but
|
||||
# stride(0) is NUM_LAYERS x larger. Neighbor layers hold random data so
|
||||
# any packed-pages addressing reads garbage rather than zeros.
|
||||
k_xl = torch.randn(
|
||||
num_pages, NUM_LAYERS, PAGE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE
|
||||
)
|
||||
k_view = k_xl[:, LAYER_IDX]
|
||||
k_view.copy_(k_ref)
|
||||
assert k_view.stride(0) == NUM_LAYERS * PAGE_SIZE * H_KV * D_QK
|
||||
if is_mla:
|
||||
v_view = k_view[..., :D_V]
|
||||
else:
|
||||
v_xl = torch.randn(
|
||||
num_pages, NUM_LAYERS, PAGE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE
|
||||
)
|
||||
v_view = v_xl[:, LAYER_IDX]
|
||||
v_view.copy_(v_ref)
|
||||
|
||||
def run(k_buffer, v_buffer):
|
||||
o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
|
||||
lse = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
|
||||
attn_logits = torch.empty(
|
||||
(B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
|
||||
)
|
||||
decode_attention_fwd(
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
o,
|
||||
lse,
|
||||
req_to_page,
|
||||
b_seq_len,
|
||||
attn_logits,
|
||||
num_kv_splits,
|
||||
sm_scale,
|
||||
PAGE_SIZE,
|
||||
is_mla=is_mla,
|
||||
)
|
||||
return o, lse
|
||||
|
||||
o_ref, lse_ref = run(k_ref, v_ref)
|
||||
o_xl, lse_xl = run(k_view, v_view)
|
||||
|
||||
# Same data and same compute order; only addressing differs.
|
||||
assert torch.equal(o_ref, o_xl)
|
||||
assert torch.equal(lse_ref, lse_xl)
|
||||
@@ -0,0 +1,232 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
|
||||
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
|
||||
|
||||
def ref_masked_attention(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
is_causal: bool = True,
|
||||
sliding_window_q: int | None = None,
|
||||
sliding_window_k: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Reference implementation using PyTorch SDPA."""
|
||||
# q, k, v: [total_tokens, num_heads, head_dim]
|
||||
# SDPA expects [batch, num_heads, seq_len, head_dim]
|
||||
|
||||
total_tokens = q.shape[0]
|
||||
|
||||
# Add batch dimension and transpose
|
||||
q = q.unsqueeze(0).transpose(1, 2) # [1, num_heads, total_tokens, head_dim]
|
||||
k = k.unsqueeze(0).transpose(1, 2) # [1, num_heads, total_tokens, head_dim]
|
||||
v = v.unsqueeze(0).transpose(1, 2) # [1, num_heads, total_tokens, head_dim]
|
||||
|
||||
# Create attention mask if needed
|
||||
attn_mask = None
|
||||
use_causal = is_causal
|
||||
|
||||
# If we have sliding window or need custom masking, create explicit mask
|
||||
sliding_window_q = sliding_window_q if sliding_window_q is not None else 0
|
||||
sliding_window_k = sliding_window_k if sliding_window_k is not None else 0
|
||||
if (sliding_window_q > 0) or (sliding_window_k > 0):
|
||||
# Position indices
|
||||
pos_q = torch.arange(total_tokens, device=q.device).unsqueeze(1)
|
||||
pos_k = torch.arange(total_tokens, device=q.device).unsqueeze(0)
|
||||
|
||||
# Start with valid mask (False = no masking)
|
||||
mask = torch.ones(
|
||||
(total_tokens, total_tokens), dtype=torch.bool, device=q.device
|
||||
)
|
||||
|
||||
# Apply causal mask
|
||||
if is_causal:
|
||||
mask = mask & (pos_q >= pos_k)
|
||||
|
||||
# Apply sliding window masks
|
||||
sliding_window_mask = torch.ones_like(mask)
|
||||
if sliding_window_q > 0:
|
||||
sliding_window_mask &= pos_q - pos_k <= sliding_window_q
|
||||
|
||||
if sliding_window_k > 0:
|
||||
sliding_window_mask &= pos_k - pos_q <= sliding_window_k
|
||||
|
||||
mask = mask & sliding_window_mask
|
||||
|
||||
attn_mask = torch.where(mask, 0.0, float("-inf")).to(q.dtype)
|
||||
use_causal = False # Don't use is_causal when providing explicit mask
|
||||
|
||||
# Use SDPA
|
||||
output = F.scaled_dot_product_attention(
|
||||
q, k, v, attn_mask=attn_mask, is_causal=use_causal, dropout_p=0.0
|
||||
)
|
||||
|
||||
# Convert back to original shape: [total_tokens, num_heads, head_dim]
|
||||
output = output.transpose(1, 2).squeeze(0)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("B", [5])
|
||||
@pytest.mark.parametrize("max_seq_len", [1024])
|
||||
@pytest.mark.parametrize("H_Q", [32])
|
||||
@pytest.mark.parametrize("H_KV", [32, 8])
|
||||
@pytest.mark.parametrize("D", [128])
|
||||
@pytest.mark.parametrize("is_causal", [True, False])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
def test_context_attention(
|
||||
B: int,
|
||||
max_seq_len: int,
|
||||
H_Q: int,
|
||||
H_KV: int,
|
||||
D: int,
|
||||
is_causal: bool,
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
"""Test basic context attention without sliding window."""
|
||||
torch.manual_seed(42)
|
||||
|
||||
# Generate random sequence lengths for each batch
|
||||
seq_lens = torch.randint(
|
||||
max_seq_len // 2, max_seq_len + 1, (B,), device=DEVICE_TYPE
|
||||
)
|
||||
total_tokens = seq_lens.sum().item()
|
||||
|
||||
# Create batch start locations
|
||||
b_start_loc = torch.zeros(B, dtype=torch.int32, device=DEVICE_TYPE)
|
||||
b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0)
|
||||
|
||||
# Create input tensors
|
||||
q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device=DEVICE_TYPE)
|
||||
k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE)
|
||||
v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE)
|
||||
o = torch.zeros_like(q)
|
||||
|
||||
# Call Triton kernel
|
||||
context_attention_fwd(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
o,
|
||||
b_start_loc,
|
||||
seq_lens,
|
||||
max_seq_len,
|
||||
is_causal=is_causal,
|
||||
sliding_window_q=None,
|
||||
sliding_window_k=None,
|
||||
)
|
||||
|
||||
# Compute reference output for each sequence in batch
|
||||
o_ref = torch.zeros_like(q)
|
||||
for i in range(B):
|
||||
start = b_start_loc[i].item()
|
||||
end = start + seq_lens[i].item()
|
||||
|
||||
q_seq = q[start:end]
|
||||
k_seq = k[start:end]
|
||||
v_seq = v[start:end]
|
||||
|
||||
# Expand KV heads if using GQA
|
||||
if H_Q != H_KV:
|
||||
kv_group_num = H_Q // H_KV
|
||||
k_seq = k_seq.repeat_interleave(kv_group_num, dim=1)
|
||||
v_seq = v_seq.repeat_interleave(kv_group_num, dim=1)
|
||||
|
||||
o_ref[start:end] = ref_masked_attention(
|
||||
q_seq,
|
||||
k_seq,
|
||||
v_seq,
|
||||
is_causal=is_causal,
|
||||
sliding_window_q=None,
|
||||
sliding_window_k=None,
|
||||
)
|
||||
|
||||
# Compare outputs
|
||||
torch.testing.assert_close(o, o_ref, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("B", [4])
|
||||
@pytest.mark.parametrize("max_seq_len", [1024])
|
||||
@pytest.mark.parametrize("H_Q", [32])
|
||||
@pytest.mark.parametrize("H_KV", [32, 8])
|
||||
@pytest.mark.parametrize("D", [128])
|
||||
@pytest.mark.parametrize("sliding_window", [(32, 32), (32, 0), (0, 32)])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
def test_context_attention_sliding_window(
|
||||
B: int,
|
||||
max_seq_len: int,
|
||||
H_Q: int,
|
||||
H_KV: int,
|
||||
D: int,
|
||||
sliding_window: tuple[int, int],
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
sliding_window_q, sliding_window_k = sliding_window
|
||||
"""Test context attention with sliding window."""
|
||||
torch.manual_seed(42)
|
||||
|
||||
# Generate random sequence lengths for each batch
|
||||
seq_lens = torch.randint(
|
||||
max_seq_len // 2, max_seq_len + 1, (B,), device=DEVICE_TYPE
|
||||
)
|
||||
total_tokens = seq_lens.sum().item()
|
||||
|
||||
# Create batch start locations
|
||||
b_start_loc = torch.zeros(B, dtype=torch.int32, device=DEVICE_TYPE)
|
||||
b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0)
|
||||
|
||||
# Create input tensors
|
||||
q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device=DEVICE_TYPE)
|
||||
k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE)
|
||||
v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE)
|
||||
o = torch.zeros_like(q)
|
||||
|
||||
# Call Triton kernel
|
||||
context_attention_fwd(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
o,
|
||||
b_start_loc,
|
||||
seq_lens,
|
||||
max_seq_len,
|
||||
is_causal=False,
|
||||
sliding_window_q=sliding_window_q,
|
||||
sliding_window_k=sliding_window_k,
|
||||
)
|
||||
|
||||
# Compute reference output for each sequence in batch
|
||||
o_ref = torch.zeros_like(q)
|
||||
for i in range(B):
|
||||
start = b_start_loc[i].item()
|
||||
end = start + seq_lens[i].item()
|
||||
|
||||
q_seq = q[start:end]
|
||||
k_seq = k[start:end]
|
||||
v_seq = v[start:end]
|
||||
|
||||
# Expand KV heads if using GQA
|
||||
if H_Q != H_KV:
|
||||
kv_group_num = H_Q // H_KV
|
||||
k_seq = k_seq.repeat_interleave(kv_group_num, dim=1)
|
||||
v_seq = v_seq.repeat_interleave(kv_group_num, dim=1)
|
||||
|
||||
o_ref[start:end] = ref_masked_attention(
|
||||
q_seq,
|
||||
k_seq,
|
||||
v_seq,
|
||||
is_causal=False,
|
||||
sliding_window_q=sliding_window_q if sliding_window_q > 0 else None,
|
||||
sliding_window_k=sliding_window_k if sliding_window_k > 0 else None,
|
||||
)
|
||||
|
||||
# Compare outputs
|
||||
torch.testing.assert_close(o, o_ref, rtol=2e-2, atol=2e-2)
|
||||
@@ -0,0 +1,645 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import next_power_of_2
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.v1.attention.ops.triton_unified_attention import unified_attention
|
||||
from vllm.v1.kv_cache_interface import KVQuantMode
|
||||
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
|
||||
NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
|
||||
HEAD_SIZES = [128, 256]
|
||||
BLOCK_SIZES = [16]
|
||||
|
||||
DTYPES = [torch.bfloat16]
|
||||
QDTYPES = [None, current_platform.fp8_dtype()]
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
# one value large enough to test overflow in index calculation.
|
||||
# one value small enough to test the schema op check
|
||||
NUM_BLOCKS = [32768, 2048]
|
||||
|
||||
# 0: use 2D kernel for decode
|
||||
# 8: use 3D kernel for decode
|
||||
SEQ_THRESHOLD_3D_VALUES = [0, 8]
|
||||
|
||||
|
||||
def ref_paged_attn(
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
query_lens: list[int],
|
||||
kv_lens: list[int],
|
||||
block_tables: torch.Tensor,
|
||||
scale: float,
|
||||
sliding_window: int | None = None,
|
||||
soft_cap: float | None = None,
|
||||
) -> torch.Tensor:
|
||||
num_seqs = len(query_lens)
|
||||
block_tables = block_tables.cpu().numpy()
|
||||
_, block_size, num_kv_heads, head_size = key_cache.shape
|
||||
|
||||
outputs: list[torch.Tensor] = []
|
||||
start_idx = 0
|
||||
for i in range(num_seqs):
|
||||
query_len = query_lens[i]
|
||||
kv_len = kv_lens[i]
|
||||
q = query[start_idx : start_idx + query_len]
|
||||
q *= scale
|
||||
|
||||
num_kv_blocks = (kv_len + block_size - 1) // block_size
|
||||
block_indices = block_tables[i, :num_kv_blocks]
|
||||
|
||||
k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
|
||||
k = k[:kv_len]
|
||||
v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
|
||||
v = v[:kv_len]
|
||||
|
||||
if q.shape[1] != k.shape[1]:
|
||||
k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
|
||||
v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
|
||||
attn = torch.einsum("qhd,khd->hqk", q, k).float()
|
||||
empty_mask = torch.ones(query_len, kv_len)
|
||||
mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
|
||||
if sliding_window is not None:
|
||||
sliding_window_mask = (
|
||||
torch.triu(
|
||||
empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1
|
||||
)
|
||||
.bool()
|
||||
.logical_not()
|
||||
)
|
||||
mask |= sliding_window_mask
|
||||
if soft_cap is not None and soft_cap > 0:
|
||||
attn = soft_cap * torch.tanh(attn / soft_cap)
|
||||
attn.masked_fill_(mask, float("-inf"))
|
||||
attn = torch.softmax(attn, dim=-1).to(v.dtype)
|
||||
out = torch.einsum("hqk,khd->qhd", attn, v)
|
||||
|
||||
outputs.append(out)
|
||||
start_idx += query_len
|
||||
|
||||
return torch.cat(outputs, dim=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("sliding_window", [None, 64, 128, 256])
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("soft_cap", [None, 50.0])
|
||||
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
|
||||
@pytest.mark.parametrize("q_dtype", QDTYPES)
|
||||
@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
|
||||
@torch.inference_mode()
|
||||
def test_triton_unified_attn(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
sliding_window: int | None,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
num_blocks: int,
|
||||
q_dtype: torch.dtype | None,
|
||||
seq_threshold_3D: int,
|
||||
) -> None:
|
||||
torch.set_default_device(DEVICE_TYPE)
|
||||
|
||||
set_random_seed(0)
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
|
||||
scale = head_size**-0.5
|
||||
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
kv_lens = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
output = torch.empty_like(query)
|
||||
|
||||
maybe_quantized_query = query
|
||||
maybe_quantized_key_cache = key_cache
|
||||
maybe_quantized_value_cache = value_cache
|
||||
q_descale = None
|
||||
k_descale = None
|
||||
v_descale = None
|
||||
kv_quant_mode = KVQuantMode.NONE
|
||||
if q_dtype is not None:
|
||||
# Use non-1 scales so FP8 Q/K/V descale handling is tested explicitly.
|
||||
q_scale = torch.tensor(0.75, dtype=torch.float32)
|
||||
k_scale = torch.tensor(0.5, dtype=torch.float32)
|
||||
v_scale = torch.tensor(0.25, dtype=torch.float32)
|
||||
q_descale = q_scale
|
||||
scale_shape = (num_seqs, num_kv_heads)
|
||||
k_descale = torch.full(scale_shape, k_scale.item(), dtype=torch.float32)
|
||||
v_descale = torch.full(scale_shape, v_scale.item(), dtype=torch.float32)
|
||||
maybe_quantized_query = (query / q_scale).to(q_dtype)
|
||||
maybe_quantized_key_cache = (key_cache / k_scale).to(q_dtype)
|
||||
maybe_quantized_value_cache = (value_cache / v_scale).to(q_dtype)
|
||||
kv_quant_mode = KVQuantMode.FP8_PER_TENSOR
|
||||
|
||||
num_par_softmax_segments = 16
|
||||
head_size_padded = next_power_of_2(head_size)
|
||||
softmax_segm_output = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_max = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_expsum = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
unified_attention(
|
||||
q=maybe_quantized_query,
|
||||
k=maybe_quantized_key_cache,
|
||||
v=maybe_quantized_value_cache,
|
||||
out=output,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
seqused_k=kv_lens,
|
||||
max_seqlen_q=max_query_len,
|
||||
max_seqlen_k=max_kv_len,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
block_table=block_tables,
|
||||
softcap=soft_cap if soft_cap is not None else 0,
|
||||
q_descale=q_descale,
|
||||
k_descale=k_descale,
|
||||
v_descale=v_descale,
|
||||
seq_threshold_3D=seq_threshold_3D,
|
||||
num_par_softmax_segments=num_par_softmax_segments,
|
||||
softmax_segm_output=softmax_segm_output,
|
||||
softmax_segm_max=softmax_segm_max,
|
||||
softmax_segm_expsum=softmax_segm_expsum,
|
||||
kv_quant_mode=kv_quant_mode,
|
||||
)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
sliding_window=sliding_window,
|
||||
soft_cap=soft_cap,
|
||||
)
|
||||
atol, rtol = 1.5e-2, 1e-2
|
||||
if q_dtype is not None:
|
||||
atol, rtol = 1.5e-1, 1.5e-1
|
||||
(
|
||||
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output - ref_output))}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
|
||||
@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
|
||||
@torch.inference_mode()
|
||||
def test_triton_unified_attn_bf16_query_fp8_kv(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
num_blocks: int,
|
||||
seq_threshold_3D: int,
|
||||
) -> None:
|
||||
"""Test bf16 Q with FP8 per-tensor KV cache (dequant via _cast_kv_tile)."""
|
||||
torch.set_default_device(DEVICE_TYPE)
|
||||
set_random_seed(0)
|
||||
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
window_size = (-1, -1)
|
||||
scale = head_size**-0.5
|
||||
|
||||
dtype = torch.bfloat16
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
|
||||
k_scale = torch.tensor(0.5, dtype=torch.float32)
|
||||
v_scale = torch.tensor(0.25, dtype=torch.float32)
|
||||
fp8_key_cache = (key_cache / k_scale).to(FP8_DTYPE)
|
||||
fp8_value_cache = (value_cache / v_scale).to(FP8_DTYPE)
|
||||
|
||||
scale_shape = (num_seqs, num_kv_heads)
|
||||
k_descale = torch.full(scale_shape, k_scale.item(), dtype=torch.float32)
|
||||
v_descale = torch.full(scale_shape, v_scale.item(), dtype=torch.float32)
|
||||
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
output = torch.empty_like(query)
|
||||
|
||||
num_par_softmax_segments = 16
|
||||
head_size_padded = next_power_of_2(head_size)
|
||||
softmax_segm_output = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_max = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_expsum = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
unified_attention(
|
||||
q=query,
|
||||
k=fp8_key_cache,
|
||||
v=fp8_value_cache,
|
||||
out=output,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
seqused_k=kv_lens_t,
|
||||
max_seqlen_q=max_query_len,
|
||||
max_seqlen_k=max_kv_len,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
block_table=block_tables,
|
||||
softcap=0,
|
||||
q_descale=None,
|
||||
k_descale=k_descale,
|
||||
v_descale=v_descale,
|
||||
seq_threshold_3D=seq_threshold_3D,
|
||||
num_par_softmax_segments=num_par_softmax_segments,
|
||||
softmax_segm_output=softmax_segm_output,
|
||||
softmax_segm_max=softmax_segm_max,
|
||||
softmax_segm_expsum=softmax_segm_expsum,
|
||||
kv_quant_mode=KVQuantMode.FP8_PER_TENSOR,
|
||||
)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
)
|
||||
|
||||
atol, rtol = 1.5e-1, 1.5e-1
|
||||
(
|
||||
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output - ref_output))}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens",
|
||||
[
|
||||
[(1, 1328), (5, 18), (129, 463)],
|
||||
[(1, 523), (1, 37), (1, 2011)],
|
||||
[(1, 1)] * 533,
|
||||
[(533, 533)] * 533,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("sliding_window", [None, 64, 128, 256])
|
||||
@pytest.mark.parametrize("soft_cap", [None, 50.0])
|
||||
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
|
||||
@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
|
||||
@torch.inference_mode()
|
||||
def test_triton_unified_attn_fp16_input_fp8_output(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
sliding_window: int | None,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
num_blocks: int,
|
||||
seq_threshold_3D: int,
|
||||
) -> None:
|
||||
"""Test with fp16 input and fp8 output using output_scale."""
|
||||
torch.set_default_device(DEVICE_TYPE)
|
||||
|
||||
set_random_seed(0)
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
|
||||
scale = head_size**-0.5
|
||||
|
||||
dtype = torch.float16
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
output = torch.empty(sum(query_lens), num_query_heads, head_size, dtype=FP8_DTYPE)
|
||||
|
||||
output_scale = torch.tensor(0.5, dtype=torch.float32)
|
||||
|
||||
num_par_softmax_segments = 16
|
||||
head_size_padded = next_power_of_2(head_size)
|
||||
softmax_segm_output = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_max = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_expsum = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
unified_attention(
|
||||
q=query,
|
||||
k=key_cache,
|
||||
v=value_cache,
|
||||
out=output,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
seqused_k=kv_lens_tensor,
|
||||
max_seqlen_q=max_query_len,
|
||||
max_seqlen_k=max_kv_len,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
block_table=block_tables,
|
||||
softcap=soft_cap if soft_cap is not None else 0,
|
||||
q_descale=None,
|
||||
k_descale=None,
|
||||
v_descale=None,
|
||||
output_scale=output_scale,
|
||||
seq_threshold_3D=seq_threshold_3D,
|
||||
num_par_softmax_segments=num_par_softmax_segments,
|
||||
softmax_segm_output=softmax_segm_output,
|
||||
softmax_segm_max=softmax_segm_max,
|
||||
softmax_segm_expsum=softmax_segm_expsum,
|
||||
)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
sliding_window=sliding_window,
|
||||
soft_cap=soft_cap,
|
||||
)
|
||||
|
||||
output_fp16 = output.to(torch.float32) * output_scale.item()
|
||||
output_fp16 = output_fp16.to(torch.float16)
|
||||
|
||||
atol, rtol = 2e-1, 2e-1
|
||||
(
|
||||
torch.testing.assert_close(output_fp16, ref_output, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output_fp16 - ref_output))}",
|
||||
)
|
||||
|
||||
|
||||
# USE_TD path covers two head-size regimes:
|
||||
# - pow2 (HEAD_SIZE == HEAD_SIZE_PADDED): full TD path including Q/O.
|
||||
# - non-pow2 (96, HEAD_SIZE_PADDED=128): gates USE_TD_QO off — Q load
|
||||
# and output store fall back to pointer path, KV tile TD load remains.
|
||||
# The non-pow2 case mirrors real models like Phi-3-mini (head_size=96).
|
||||
HEAD_SIZES_USE_TD = [128, 256, 96]
|
||||
|
||||
|
||||
def _run_use_td_case(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
sliding_window: int | None,
|
||||
soft_cap: float | None,
|
||||
seq_threshold_3D: int,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
num_blocks: int = 2048,
|
||||
) -> None:
|
||||
"""Shared driver for the USE_TD test cases.
|
||||
|
||||
Runs ``unified_attention(..., use_td=True)`` and compares against the
|
||||
reference paged-attention implementation that the sibling non-TD
|
||||
tests use.
|
||||
"""
|
||||
torch.set_default_device(DEVICE_TYPE)
|
||||
set_random_seed(0)
|
||||
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads, num_kv_heads = num_heads
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
|
||||
scale = head_size**-0.5
|
||||
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
output = torch.empty_like(query)
|
||||
|
||||
num_par_softmax_segments = 16
|
||||
head_size_padded = next_power_of_2(head_size)
|
||||
softmax_segm_output = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_max = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_expsum = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
unified_attention(
|
||||
q=query,
|
||||
k=key_cache,
|
||||
v=value_cache,
|
||||
out=output,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
seqused_k=kv_lens_tensor,
|
||||
max_seqlen_q=max_query_len,
|
||||
max_seqlen_k=max_kv_len,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
block_table=block_tables,
|
||||
softcap=soft_cap if soft_cap is not None else 0,
|
||||
q_descale=None,
|
||||
k_descale=None,
|
||||
v_descale=None,
|
||||
seq_threshold_3D=seq_threshold_3D,
|
||||
num_par_softmax_segments=num_par_softmax_segments,
|
||||
softmax_segm_output=softmax_segm_output,
|
||||
softmax_segm_max=softmax_segm_max,
|
||||
softmax_segm_expsum=softmax_segm_expsum,
|
||||
use_td=True,
|
||||
)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
sliding_window=sliding_window,
|
||||
soft_cap=soft_cap,
|
||||
)
|
||||
torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES_USE_TD)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("sliding_window", [None, 128])
|
||||
@pytest.mark.parametrize("soft_cap", [None, 50.0])
|
||||
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
|
||||
@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
|
||||
@torch.inference_mode()
|
||||
def test_triton_unified_attn_use_td(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
sliding_window: int | None,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
num_blocks: int,
|
||||
seq_threshold_3D: int,
|
||||
) -> None:
|
||||
"""Exercise the USE_TD (tensor-descriptor) Q/K/V load/store path.
|
||||
|
||||
Covers both 2D and 3D kernels via ``seq_threshold_3D``. Two routes
|
||||
to the USE_TD_QO=False fallback (pointer path for Q/O with TD still
|
||||
active for KV tile loads):
|
||||
|
||||
- non-pow2 ``num_queries_per_kv`` via ``NUM_HEADS`` entry ``(5, 1)``,
|
||||
- non-pow2 ``head_size`` via ``HEAD_SIZES_USE_TD`` entry ``96``.
|
||||
"""
|
||||
_run_use_td_case(
|
||||
seq_lens=seq_lens,
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
block_size=block_size,
|
||||
sliding_window=sliding_window,
|
||||
soft_cap=soft_cap,
|
||||
seq_threshold_3D=seq_threshold_3D,
|
||||
num_blocks=num_blocks,
|
||||
)
|
||||
|
||||
|
||||
# Prefill-heavy shape: long query drives the prefill kernel path where
|
||||
# ``_get_tile_size`` returns 32, which exceeds block_size=16 and must be
|
||||
# clamped by the fix in 'clamp TILE_SIZE to block_size when USE_TD'.
|
||||
# Only the prefill launch exercises the clamp, so parameterize only over
|
||||
# the (num_heads, seq_threshold_3D=0) combinations needed to cover it.
|
||||
@pytest.mark.parametrize("num_heads", [(4, 4), (5, 1)])
|
||||
@torch.inference_mode()
|
||||
def test_triton_unified_attn_use_td_tile_clamp(
|
||||
num_heads: tuple[int, int],
|
||||
) -> None:
|
||||
"""Regression guard: ``USE_TD`` needs ``BLOCK_SIZE % TILE_SIZE == 0``.
|
||||
|
||||
With ``block_size=16`` and ``head_size=128`` (non-Gemma3),
|
||||
``_get_tile_size`` returns 32 for prefill, which violates the
|
||||
``USE_TD`` constraint unless clamped to ``block_size``. Without
|
||||
the clamp the triton kernel ``static_assert`` fires at compile time.
|
||||
"""
|
||||
_run_use_td_case(
|
||||
seq_lens=[(256, 256), (128, 128)],
|
||||
num_heads=num_heads,
|
||||
head_size=128,
|
||||
block_size=16,
|
||||
sliding_window=None,
|
||||
soft_cap=None,
|
||||
seq_threshold_3D=0,
|
||||
)
|
||||
@@ -0,0 +1,189 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Unit tests for the Triton DiffKV unified-attention kernel.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import next_power_of_2
|
||||
from vllm.utils.torch_utils import (
|
||||
canonicalize_singleton_dim_strides,
|
||||
set_random_seed,
|
||||
)
|
||||
from vllm.v1.attention.backends.fa_utils import (
|
||||
get_flash_attn_version,
|
||||
is_flash_attn_varlen_func_available,
|
||||
)
|
||||
from vllm.v1.attention.ops.triton_unified_attention_diffkv import (
|
||||
unified_attention_diffkv,
|
||||
)
|
||||
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
|
||||
# (num_query_heads, num_kv_heads): MHA, GQA, and the num_kv_heads==1
|
||||
# (degenerate-stride) case.
|
||||
NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
|
||||
# (head_size_qk, head_size_v). (192, 128) is the canonical asymmetric
|
||||
# DiffKV shape; FA4 on Blackwell only supports head_size>128 when it is
|
||||
# 192, and FA3 on Hopper supports it too -- so this pair is runnable on
|
||||
# both. (128, 128) keeps the equal-dim path covered through the DiffKV
|
||||
# kernel.
|
||||
HEAD_SIZES = [(128, 128), (192, 128)]
|
||||
BLOCK_SIZES = [16]
|
||||
DTYPES = [torch.bfloat16]
|
||||
|
||||
NUM_BLOCKS = 2048
|
||||
|
||||
# 0: 2D decode kernel; 8: 3D (split-KV) decode kernel.
|
||||
SEQ_THRESHOLD_3D_VALUES = [0, 8]
|
||||
|
||||
NUM_PAR_SOFTMAX_SEGMENTS = 16
|
||||
|
||||
|
||||
def _alloc_segm_buffers(seq_threshold_3D: int, num_query_heads: int, head_size_v: int):
|
||||
"""Allocate the split-KV softmax scratch (last dim == head_size_v)."""
|
||||
head_size_v_padded = next_power_of_2(head_size_v)
|
||||
segm_output = torch.empty(
|
||||
(
|
||||
seq_threshold_3D,
|
||||
num_query_heads,
|
||||
NUM_PAR_SOFTMAX_SEGMENTS,
|
||||
head_size_v_padded,
|
||||
),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
segm_max = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, NUM_PAR_SOFTMAX_SEGMENTS),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
segm_expsum = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, NUM_PAR_SOFTMAX_SEGMENTS),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
return segm_output, segm_max, segm_expsum
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens",
|
||||
[
|
||||
[(1, 1328), (5, 18), (129, 463)], # mixed prefill + decode
|
||||
[(1, 523), (1, 37), (1, 2011)], # decode-only (exercises 3D path)
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_sizes", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("sliding_window", [None, 128])
|
||||
@pytest.mark.parametrize("soft_cap", [None, 50.0])
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
|
||||
@torch.inference_mode()
|
||||
def test_triton_unified_attn_diffkv_vs_fa(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_sizes: tuple[int, int],
|
||||
sliding_window: int | None,
|
||||
soft_cap: float | None,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
seq_threshold_3D: int,
|
||||
) -> None:
|
||||
head_size_qk, head_size_v = head_sizes
|
||||
|
||||
# DiffKV requires FA3 (Hopper) / FA4 (Blackwell) as the reference.
|
||||
fa_version = get_flash_attn_version(head_size=head_size_qk, head_size_v=head_size_v)
|
||||
if not is_flash_attn_varlen_func_available() or fa_version not in (3, 4):
|
||||
pytest.skip(f"FA DiffKV needs FA3/FA4 (got version {fa_version}).")
|
||||
|
||||
from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
|
||||
|
||||
torch.set_default_device(DEVICE_TYPE)
|
||||
set_random_seed(0)
|
||||
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads, num_kv_heads = num_heads
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
|
||||
scale = head_size_qk**-0.5
|
||||
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size_qk, dtype=dtype)
|
||||
# Packed KV cache: [num_blocks, block_size, num_kv_heads, hqk + hv].
|
||||
kv_cache = torch.randn(
|
||||
NUM_BLOCKS,
|
||||
block_size,
|
||||
num_kv_heads,
|
||||
head_size_qk + head_size_v,
|
||||
dtype=dtype,
|
||||
)
|
||||
key_cache = kv_cache[..., :head_size_qk]
|
||||
value_cache = kv_cache[..., head_size_qk:]
|
||||
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
# ---- FlashAttention DiffKV (ground truth) ---------------------------
|
||||
# Mirror the backend: fix degenerate strides on size-1 dims so FA's
|
||||
# TMA path sees ≥16-byte-aligned strides (matters for num_kv_heads==1).
|
||||
fa_k = canonicalize_singleton_dim_strides(key_cache)
|
||||
fa_v = canonicalize_singleton_dim_strides(value_cache)
|
||||
fa_out = torch.empty(sum(query_lens), num_query_heads, head_size_v, dtype=dtype)
|
||||
flash_attn_varlen_func(
|
||||
q=query,
|
||||
k=fa_k,
|
||||
v=fa_v,
|
||||
out=fa_out,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
max_seqlen_q=max_query_len,
|
||||
seqused_k=kv_lens_t,
|
||||
max_seqlen_k=max_kv_len,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=list(window_size),
|
||||
block_table=block_tables,
|
||||
softcap=soft_cap if soft_cap is not None else 0,
|
||||
fa_version=fa_version,
|
||||
)
|
||||
|
||||
# ---- Triton DiffKV --------------------------------------------------
|
||||
segm_output, segm_max, segm_expsum = _alloc_segm_buffers(
|
||||
seq_threshold_3D, num_query_heads, head_size_v
|
||||
)
|
||||
triton_out = torch.empty(sum(query_lens), num_query_heads, head_size_v, dtype=dtype)
|
||||
unified_attention_diffkv(
|
||||
q=query,
|
||||
k=key_cache,
|
||||
v=value_cache,
|
||||
out=triton_out,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
seqused_k=kv_lens_t,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
block_table=block_tables,
|
||||
softcap=soft_cap if soft_cap is not None else 0,
|
||||
max_seqlen_q=max_query_len,
|
||||
seq_threshold_3D=seq_threshold_3D,
|
||||
num_par_softmax_segments=NUM_PAR_SOFTMAX_SEGMENTS,
|
||||
softmax_segm_output=segm_output,
|
||||
softmax_segm_max=segm_max,
|
||||
softmax_segm_expsum=segm_expsum,
|
||||
)
|
||||
|
||||
(
|
||||
torch.testing.assert_close(triton_out, fa_out, atol=2e-2, rtol=2e-2),
|
||||
f"triton vs FA max abs diff: {torch.max(torch.abs(triton_out - fa_out))}",
|
||||
)
|
||||
@@ -0,0 +1,440 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant.
|
||||
|
||||
Tests both contiguous and non-contiguous (cross-layer unified) KV cache
|
||||
layouts against a pure-PyTorch reference implementation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if current_platform.is_rocm():
|
||||
pytest.skip(
|
||||
"trtllm kvfp8 dequant is not supported on ROCm.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
NUM_BLOCKS = 128
|
||||
|
||||
|
||||
def to_float8(x, dtype=None):
|
||||
if dtype is None:
|
||||
dtype = FP8_DTYPE
|
||||
finfo = torch.finfo(dtype)
|
||||
min_val, max_val = x.aminmax()
|
||||
amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
|
||||
scale = finfo.max / amax * 0.1
|
||||
x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
|
||||
return x_scl_sat.to(dtype), scale.float().reciprocal()
|
||||
|
||||
|
||||
def make_contiguous_kv_cache(num_blocks, num_kv_heads, block_size, head_size):
|
||||
"""Create a standard contiguous fp8 KV cache (HND layout)."""
|
||||
raw = torch.randn(
|
||||
num_blocks,
|
||||
2,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
)
|
||||
kv_cache, scale = to_float8(raw)
|
||||
return kv_cache, scale
|
||||
|
||||
|
||||
def make_cross_layer_kv_cache(
|
||||
num_blocks,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
num_layers=4,
|
||||
):
|
||||
"""
|
||||
Create a non-contiguous per-layer view mimicking cross-layer allocation.
|
||||
|
||||
Physical layout: (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
|
||||
Returned view: (num_blocks, 2, num_kv_heads, block_size, head_size)
|
||||
with non-contiguous strides on dims 0, 1, 2 (they skip over num_layers).
|
||||
"""
|
||||
raw = torch.randn(
|
||||
num_blocks,
|
||||
2,
|
||||
num_kv_heads,
|
||||
num_layers,
|
||||
block_size,
|
||||
head_size,
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
)
|
||||
fp8_full, scale = to_float8(raw)
|
||||
layer_view = fp8_full[:, :, :, 0, :, :]
|
||||
assert not layer_view.is_contiguous(), (
|
||||
f"Expected non-contiguous view, got strides {layer_view.stride()}"
|
||||
)
|
||||
return layer_view, scale
|
||||
|
||||
|
||||
def ref_dequant(kv_cache, block_tables, k_scale, v_scale, dequant_dtype):
|
||||
"""Pure PyTorch reference: gather pages and dequantize fp8 -> dequant_dtype."""
|
||||
batch_size, num_pages_per_seq = block_tables.shape
|
||||
s = kv_cache.shape
|
||||
out = torch.zeros(
|
||||
batch_size * num_pages_per_seq + 1,
|
||||
s[1],
|
||||
s[2],
|
||||
s[3],
|
||||
s[4],
|
||||
dtype=dequant_dtype,
|
||||
device=kv_cache.device,
|
||||
)
|
||||
for b in range(batch_size):
|
||||
for p in range(num_pages_per_seq):
|
||||
page_idx = block_tables[b, p].item()
|
||||
if page_idx <= 0:
|
||||
continue
|
||||
mock_idx = b * num_pages_per_seq + p + 1
|
||||
out[mock_idx, 0] = (kv_cache[page_idx, 0].float() * k_scale.item()).to(
|
||||
dequant_dtype
|
||||
)
|
||||
out[mock_idx, 1] = (kv_cache[page_idx, 1].float() * v_scale.item()).to(
|
||||
dequant_dtype
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_kv_heads", [1, 8])
|
||||
@pytest.mark.parametrize("head_size", [64, 128])
|
||||
@pytest.mark.parametrize("block_size", [16, 32])
|
||||
@pytest.mark.parametrize("batch_size", [1, 4])
|
||||
@pytest.mark.parametrize("num_pages_per_seq", [3, 8])
|
||||
@pytest.mark.parametrize("contiguous", [True, False])
|
||||
@torch.inference_mode()
|
||||
def test_trtllm_kvfp8_dequant(
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
batch_size: int,
|
||||
num_pages_per_seq: int,
|
||||
contiguous: bool,
|
||||
):
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
if contiguous:
|
||||
kv_cache, scale = make_contiguous_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
else:
|
||||
kv_cache, scale = make_cross_layer_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
|
||||
k_scale = scale.clone()
|
||||
v_scale = scale.clone()
|
||||
|
||||
block_tables = torch.randint(
|
||||
1,
|
||||
NUM_BLOCKS,
|
||||
(batch_size, num_pages_per_seq),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
expected_bt = torch.arange(
|
||||
1,
|
||||
batch_size * num_pages_per_seq + 1,
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
).reshape(batch_size, num_pages_per_seq)
|
||||
torch.testing.assert_close(mock_block_table, expected_bt)
|
||||
|
||||
# Page 0 is padding (never written), compare only pages 1+
|
||||
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_block_tables_with_zero_pages():
|
||||
"""Pages with index <= 0 must be skipped (early return in kernel)."""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 8, 16, 64
|
||||
|
||||
kv_cache, scale = make_contiguous_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
k_scale = v_scale = scale.clone()
|
||||
|
||||
# Mix of valid pages and zeros (padding)
|
||||
block_tables = torch.tensor(
|
||||
[[5, 0, 10], [0, 0, 0], [3, 7, 0]],
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
# Only compare pages that were actually written (non-zero page indices)
|
||||
for b in range(block_tables.shape[0]):
|
||||
for p in range(block_tables.shape[1]):
|
||||
if block_tables[b, p].item() > 0:
|
||||
idx = b * block_tables.shape[1] + p + 1
|
||||
torch.testing.assert_close(
|
||||
mock_kv_cache[idx],
|
||||
ref[idx],
|
||||
atol=1e-3,
|
||||
rtol=1e-3,
|
||||
)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_all_zero_block_tables():
|
||||
"""All-zero block_tables: kernel should write nothing."""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 4, 16, 64
|
||||
|
||||
kv_cache, scale = make_contiguous_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
k_scale = v_scale = scale.clone()
|
||||
|
||||
block_tables = torch.zeros(2, 4, dtype=torch.int32, device="cuda")
|
||||
|
||||
# Should not crash even though no pages are valid
|
||||
mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
assert mock_kv_cache.shape[0] == 2 * 4 + 1
|
||||
assert mock_block_table.shape == (2, 4)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_different_k_v_scales():
|
||||
"""Verify K and V are dequantized with independent scales."""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 8, 16, 64
|
||||
|
||||
kv_cache, _ = make_contiguous_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
k_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
v_scale = torch.tensor([2.0], dtype=torch.float32, device="cuda")
|
||||
|
||||
block_tables = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda")
|
||||
|
||||
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_single_page_per_seq():
|
||||
"""Minimum grid dim 1 = 1 page per sequence."""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 8, 16, 128
|
||||
|
||||
kv_cache, scale = make_contiguous_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
k_scale = v_scale = scale.clone()
|
||||
|
||||
block_tables = torch.tensor([[5], [10], [20]], dtype=torch.int32, device="cuda")
|
||||
|
||||
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_large_page_indices():
|
||||
"""Page indices near the top of the buffer stress offset arithmetic."""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 8, 16, 128
|
||||
large_num_blocks = 32768
|
||||
|
||||
kv_cache, scale = make_contiguous_kv_cache(
|
||||
large_num_blocks,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
k_scale = v_scale = scale.clone()
|
||||
|
||||
# Use page indices near the top of the buffer
|
||||
block_tables = torch.tensor(
|
||||
[[large_num_blocks - 1, large_num_blocks - 2, 1]],
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_large_block_size():
|
||||
"""block_size=64 -> HEAD_STRIDE=8192, large tl.arange per thread block."""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 4, 64, 128
|
||||
|
||||
kv_cache, scale = make_contiguous_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
k_scale = v_scale = scale.clone()
|
||||
|
||||
block_tables = torch.randint(
|
||||
1,
|
||||
NUM_BLOCKS,
|
||||
(2, 4),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_cross_layer_many_layers():
|
||||
"""
|
||||
Non-contiguous with 36 layers -- matches real gpt-oss-120b.
|
||||
Strides are far from contiguous (factor of 36 in the gaps).
|
||||
"""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 8, 16, 64
|
||||
num_layers = 36
|
||||
|
||||
kv_cache, scale = make_cross_layer_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
num_layers=num_layers,
|
||||
)
|
||||
k_scale = v_scale = scale.clone()
|
||||
|
||||
block_tables = torch.randint(
|
||||
1,
|
||||
NUM_BLOCKS,
|
||||
(4, 6),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
|
||||
@@ -0,0 +1,218 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.flashinfer import (
|
||||
can_use_trtllm_attention,
|
||||
supports_trtllm_attention,
|
||||
use_trtllm_attention,
|
||||
)
|
||||
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip(
|
||||
"TRTLLM attention is only supported on CUDA platforms.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
MODEL_CONFIGS = {
|
||||
"Llama-3-70B": dict(num_qo_heads=64, num_kv_heads=8),
|
||||
"Llama-3-8B": dict(num_qo_heads=32, num_kv_heads=8),
|
||||
"Qwen2.5-0.5B": dict(num_qo_heads=14, num_kv_heads=2),
|
||||
"Mistral-7B": dict(num_qo_heads=32, num_kv_heads=8),
|
||||
"Gemma-2-9B": dict(num_qo_heads=8, num_kv_heads=4),
|
||||
"Falcon-40B": dict(num_qo_heads=128, num_kv_heads=8),
|
||||
}
|
||||
|
||||
|
||||
def get_config(model: str) -> dict:
|
||||
"""Return the attention config for a model."""
|
||||
return MODEL_CONFIGS[model]
|
||||
|
||||
|
||||
DEFAULT_KWARGS = dict(
|
||||
**get_config("Llama-3-70B"),
|
||||
num_tokens=128,
|
||||
max_seq_len=4096,
|
||||
dcp_world_size=1,
|
||||
kv_cache_dtype="auto",
|
||||
q_dtype=torch.bfloat16,
|
||||
is_prefill=False,
|
||||
force_use_trtllm=None,
|
||||
has_sinks=False,
|
||||
has_spec=False,
|
||||
)
|
||||
|
||||
|
||||
def _call(**overrides) -> bool:
|
||||
kwargs = {**DEFAULT_KWARGS, **overrides}
|
||||
return use_trtllm_attention(**kwargs)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_supports_cache():
|
||||
"""Clear functools.cache to ensure each test runs independently."""
|
||||
supports_trtllm_attention.cache_clear()
|
||||
|
||||
|
||||
# supports_trtllm_attention
|
||||
|
||||
|
||||
@patch("vllm.envs.VLLM_BATCH_INVARIANT", True)
|
||||
def test_supports_batch_invariant_disables():
|
||||
assert supports_trtllm_attention() is False
|
||||
|
||||
|
||||
@patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
|
||||
@patch(
|
||||
"vllm.utils.flashinfer.current_platform.is_device_capability_family",
|
||||
return_value=True,
|
||||
)
|
||||
@patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)
|
||||
def test_supports_sm100_with_artifactory(_art, _cap):
|
||||
assert supports_trtllm_attention() is True
|
||||
|
||||
|
||||
@patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
|
||||
@patch(
|
||||
"vllm.utils.flashinfer.current_platform.is_device_capability", return_value=False
|
||||
)
|
||||
@patch(
|
||||
"vllm.utils.flashinfer.current_platform.is_device_capability_family",
|
||||
return_value=False,
|
||||
)
|
||||
def test_supports_unsupported_platform(_family, _cap):
|
||||
assert supports_trtllm_attention() is False
|
||||
|
||||
|
||||
@patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
|
||||
@patch("vllm.utils.flashinfer.current_platform.is_device_capability", return_value=True)
|
||||
@patch(
|
||||
"vllm.utils.flashinfer.current_platform.is_device_capability_family",
|
||||
return_value=False,
|
||||
)
|
||||
@patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)
|
||||
def test_supports_sm90_decode_only(_art, _family, _cap):
|
||||
assert supports_trtllm_attention(is_prefill=False) is True
|
||||
assert supports_trtllm_attention(is_prefill=True) is False
|
||||
|
||||
|
||||
@patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
|
||||
@patch(
|
||||
"vllm.utils.flashinfer.current_platform.is_device_capability_family",
|
||||
return_value=True,
|
||||
)
|
||||
@patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=False)
|
||||
def test_supports_sm100_without_artifactory(_art, _cap):
|
||||
assert supports_trtllm_attention() is False
|
||||
|
||||
|
||||
# can_use_trtllm_attention
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=False)
|
||||
def test_can_use_force_disabled(_mock):
|
||||
cfg = get_config("Llama-3-70B")
|
||||
assert can_use_trtllm_attention(cfg["num_qo_heads"], cfg["num_kv_heads"]) is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=None)
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_can_use_compatible_heads(_sup, _force):
|
||||
cfg = get_config("Llama-3-70B")
|
||||
assert can_use_trtllm_attention(cfg["num_qo_heads"], cfg["num_kv_heads"]) is True
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=None)
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_can_use_incompatible_heads(_sup, _force):
|
||||
assert can_use_trtllm_attention(40, 6) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", list(MODEL_CONFIGS.keys()))
|
||||
@patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=None)
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
|
||||
def test_can_use_platform_unsupported(_sup, _force, model):
|
||||
cfg = get_config(model)
|
||||
assert can_use_trtllm_attention(cfg["num_qo_heads"], cfg["num_kv_heads"]) is False
|
||||
|
||||
|
||||
# use_trtllm_attention
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_force_off(_mock):
|
||||
assert _call(force_use_trtllm=False) is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_dcp_fallback(_mock):
|
||||
assert _call(dcp_world_size=2) is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
|
||||
def test_use_platform_unsupported(_mock):
|
||||
assert _call() is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
|
||||
def test_use_platform_unsupported_force_on_still_false(_mock):
|
||||
assert _call(force_use_trtllm=True) is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_incompatible_heads(_mock):
|
||||
assert _call(num_qo_heads=40, num_kv_heads=6) is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_incompatible_heads_force_on_still_false(_mock):
|
||||
assert _call(num_qo_heads=40, num_kv_heads=6, force_use_trtllm=True) is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_spec_decode_enables(_mock):
|
||||
assert _call(has_spec=True, is_prefill=False) is True
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
@patch(
|
||||
"vllm.utils.flashinfer.current_platform.fp8_dtype",
|
||||
return_value=torch.float8_e4m3fn,
|
||||
)
|
||||
def test_use_fp8_query_forces_trtllm(_fp8, _sup):
|
||||
assert _call(q_dtype=torch.float8_e4m3fn) is True
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_sinks_force_trtllm(_mock):
|
||||
assert _call(has_sinks=True) is True
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_auto_prefill_kv_auto(_mock):
|
||||
assert _call(is_prefill=True, kv_cache_dtype="auto") is True
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_auto_prefill_kv_fp8(_mock):
|
||||
assert _call(is_prefill=True, kv_cache_dtype="fp8") is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_auto_decode_small_batch(_mock):
|
||||
assert _call(is_prefill=False, num_tokens=128, kv_cache_dtype="auto") is True
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_auto_decode_large_batch(_mock):
|
||||
assert _call(is_prefill=False, num_tokens=512, kv_cache_dtype="auto") is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_force_on(_mock):
|
||||
assert _call(force_use_trtllm=True) is True
|
||||
@@ -0,0 +1,118 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.v1.attention.ops.xpu_mla_sparse import triton_bf16_mla_sparse_interface
|
||||
|
||||
|
||||
# https://github.com/deepseek-ai/FlashMLA/blob/main/tests/ref.py#L7
|
||||
def _merge_two_lse(
|
||||
lse0: torch.Tensor, lse1: torch.Tensor | None, s_q: int, h_q: int
|
||||
) -> torch.Tensor:
|
||||
if lse1 is None:
|
||||
return lse0
|
||||
else:
|
||||
return torch.logsumexp(
|
||||
torch.stack([lse0.view(s_q, h_q), lse1.broadcast_to(s_q, h_q)], dim=0),
|
||||
dim=0,
|
||||
)
|
||||
|
||||
|
||||
# Adapted from https://github.com/deepseek-ai/FlashMLA/blob/main/tests/ref.py#L19
|
||||
def reference_mla_sparse_prefill(
|
||||
q: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
sm_scale: float,
|
||||
d_v: int,
|
||||
topk_length: torch.Tensor | None = None,
|
||||
attn_sink: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Returns:
|
||||
- o: [s_q, h_q, dv]
|
||||
- o_fp32: [s_q, h_q, dv]
|
||||
- max_logits: [s_q, h_q]
|
||||
- lse: [s_q, h_q]
|
||||
"""
|
||||
s_q, h_q, d_qk = q.shape
|
||||
s_kv, _, _ = kv.shape
|
||||
_, _, topk = indices.shape
|
||||
|
||||
indices = indices.clone().squeeze(1)
|
||||
if topk_length is not None:
|
||||
mask = torch.arange(topk, device=topk_length.device).unsqueeze(0).broadcast_to(
|
||||
s_q, topk
|
||||
) >= topk_length.unsqueeze(1) # [s_q, topk]
|
||||
indices[mask] = -1
|
||||
invalid_mask = (indices < 0) | (indices >= s_kv) # [s_q, topk]
|
||||
indices[invalid_mask] = 0
|
||||
|
||||
q = q.float()
|
||||
gathered_kv = (
|
||||
kv.index_select(dim=0, index=indices.flatten()).reshape(s_q, topk, d_qk).float()
|
||||
) # [s_q, topk, d_qk]
|
||||
P = q @ gathered_kv.transpose(1, 2) # [s_q, h_q, topk]
|
||||
P *= sm_scale
|
||||
P[invalid_mask.unsqueeze(1).broadcast_to(P.shape)] = float("-inf")
|
||||
|
||||
orig_lse = torch.logsumexp(P, dim=-1) # [s_q, h_q]
|
||||
max_logits = P.max(dim=-1).values # [s_q, h_q]
|
||||
|
||||
lse_for_o = _merge_two_lse(orig_lse, attn_sink, s_q, h_q)
|
||||
if not torch.is_inference_mode_enabled():
|
||||
lse_for_o = lse_for_o.clone()
|
||||
lse_for_o[lse_for_o == float("-inf")] = float(
|
||||
"+inf"
|
||||
) # So that corresponding O will be 0
|
||||
s_for_o = torch.exp(P - lse_for_o.unsqueeze(-1))
|
||||
out = s_for_o @ gathered_kv[..., :d_v] # [s_q, h_q, dv]
|
||||
|
||||
lonely_q_mask = orig_lse == float("-inf") # [s_q, h_q]
|
||||
orig_lse[lonely_q_mask] = float("+inf")
|
||||
return (out.to(kv.dtype), out, max_logits, orig_lse)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device_str", ["xpu"])
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.skipif(
|
||||
not torch.xpu.is_available(),
|
||||
reason="XPU is required",
|
||||
)
|
||||
def test_bf16_triton_sparse_mla(device_str, dtype):
|
||||
device = torch.device(device_str)
|
||||
s_q = 1
|
||||
s_kv = 256
|
||||
h_q = 64 # kernel expects multiple of 64
|
||||
h_kv = 1
|
||||
d_qk = 576
|
||||
d_v = 512
|
||||
topk = 128
|
||||
|
||||
torch.random.manual_seed(1234)
|
||||
|
||||
q = torch.randn((s_q, h_q, d_qk), dtype=dtype, device=device)
|
||||
kv = torch.randn((s_kv, h_kv, d_qk), dtype=dtype, device=device)
|
||||
indices = torch.full((s_q, h_kv, topk), -1, dtype=torch.int32, device=device)
|
||||
for t in range(s_q):
|
||||
for h in range(h_kv):
|
||||
i_i = torch.randperm(max(1, t))[:topk]
|
||||
indices[t, h, : len(i_i)] = i_i
|
||||
|
||||
sm_scale = d_qk**-0.5
|
||||
|
||||
out, max_logits, lse = triton_bf16_mla_sparse_interface(
|
||||
q, kv, indices, sm_scale, d_v
|
||||
)
|
||||
assert out.shape == (s_q, h_q, d_v)
|
||||
assert max_logits.shape == (s_q, h_q)
|
||||
assert lse.shape == (s_q, h_q)
|
||||
|
||||
ref_out, ref_out_fp32, ref_max_logits, ref_lse = reference_mla_sparse_prefill(
|
||||
q, kv, indices, sm_scale, d_v
|
||||
)
|
||||
assert torch.allclose(out, ref_out, atol=1e-2, rtol=1e-2)
|
||||
assert torch.allclose(max_logits, ref_max_logits, atol=1e-3, rtol=1e-3)
|
||||
assert torch.allclose(lse, ref_lse, atol=1e-3, rtol=1e-3)
|
||||
@@ -0,0 +1,236 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.allclose_default import get_default_atol, get_default_rtol
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm.model_executor.layers.activation import (
|
||||
FastGELU,
|
||||
FatreluAndMul,
|
||||
GeluAndMul,
|
||||
MulAndSilu,
|
||||
NewGELU,
|
||||
QuickGELU,
|
||||
ReLUSquaredActivation,
|
||||
SiluAndMul,
|
||||
SiluAndMulWithClamp,
|
||||
SwigluOAIAndMul,
|
||||
SwigluStepAndMul,
|
||||
swiglustep_and_mul_triton,
|
||||
)
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
DTYPES = [torch.half, torch.bfloat16, torch.float]
|
||||
NUM_TOKENS = [7, 83, 2048] # Arbitrary values for testing
|
||||
D = [512, 13824] # Arbitrary values for testing
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"activation",
|
||||
[
|
||||
"silu_and_mul",
|
||||
"mul_and_silu",
|
||||
"gelu",
|
||||
"gelu_tanh",
|
||||
"fatrelu",
|
||||
"swigluoai_and_mul",
|
||||
"swiglustep_and_mul",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("d", D)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_act_and_mul(
|
||||
default_vllm_config,
|
||||
activation: str,
|
||||
num_tokens: int,
|
||||
d: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
x = torch.randn(num_tokens, 2 * d, dtype=dtype)
|
||||
if activation == "silu_and_mul":
|
||||
layer = SiluAndMul(compile_native=False)
|
||||
fn = torch.ops._C.silu_and_mul
|
||||
if activation == "mul_and_silu":
|
||||
layer = MulAndSilu()
|
||||
fn = torch.ops._C.mul_and_silu
|
||||
elif activation == "gelu":
|
||||
layer = GeluAndMul(approximate="none")
|
||||
fn = torch.ops._C.gelu_and_mul
|
||||
elif activation == "gelu_tanh":
|
||||
layer = GeluAndMul(approximate="tanh")
|
||||
fn = torch.ops._C.gelu_tanh_and_mul
|
||||
elif activation == "fatrelu":
|
||||
threshold = random.uniform(0, 1)
|
||||
layer = FatreluAndMul(threshold)
|
||||
fn = torch.ops._C.fatrelu_and_mul
|
||||
elif activation == "swigluoai_and_mul":
|
||||
layer = SwigluOAIAndMul()
|
||||
fn = torch.ops._C.swigluoai_and_mul
|
||||
elif activation == "swiglustep_and_mul":
|
||||
layer = SwigluStepAndMul()
|
||||
fn = swiglustep_and_mul_triton
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
if activation in ["swigluoai_and_mul", "swiglustep_and_mul"]:
|
||||
rtol = {
|
||||
# For fp16, change the relative tolerance from 1e-3 to 2e-3
|
||||
torch.float16: 2e-3,
|
||||
torch.bfloat16: 2e-2,
|
||||
torch.float: 1.3e-6,
|
||||
}
|
||||
|
||||
def _get_rtol(output) -> float:
|
||||
return rtol[output.dtype]
|
||||
|
||||
torch.testing.assert_close(
|
||||
out, ref_out, atol=get_default_atol(out), rtol=_get_rtol(out)
|
||||
)
|
||||
else:
|
||||
# The SiluAndMul, MulAndSilu, GELU and FatReLU implementations are
|
||||
# equivalent to the native PyTorch implementations, so we can do exact
|
||||
# comparison.
|
||||
torch.testing.assert_close(out, ref_out, atol=0.0, rtol=0.0)
|
||||
|
||||
d = x.shape[-1] // 2
|
||||
output_shape = x.shape[:-1] + (d,)
|
||||
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
||||
if activation == "fatrelu":
|
||||
opcheck(fn, (out, x, threshold))
|
||||
elif activation == "swigluoai_and_mul":
|
||||
opcheck(fn, (out, x, layer.alpha, layer.limit))
|
||||
elif activation != "swiglustep_and_mul":
|
||||
opcheck(fn, (out, x))
|
||||
|
||||
|
||||
SWIGLU_LIMITS = [3.0, 7.0, 15.0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("d", D)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_silu_and_mul_with_clamp(
|
||||
default_vllm_config,
|
||||
swiglu_limit: float,
|
||||
num_tokens: int,
|
||||
d: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
"""SiluAndMulWithClamp: cuda kernel must match native reference."""
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
# Use large values to ensure clamping is exercised.
|
||||
x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2
|
||||
|
||||
layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False)
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
|
||||
rtol = {
|
||||
torch.float16: 2e-3,
|
||||
torch.bfloat16: 2e-2,
|
||||
torch.float: 1.3e-6,
|
||||
}
|
||||
torch.testing.assert_close(
|
||||
out, ref_out, atol=get_default_atol(out), rtol=rtol[out.dtype]
|
||||
)
|
||||
|
||||
# Verify clamping is actually being applied: the clamped output should
|
||||
# differ from the unclamped SiluAndMul output when inputs are large.
|
||||
unclamped_out = SiluAndMul.forward_native(x)
|
||||
assert not torch.equal(ref_out.float(), unclamped_out.float()), (
|
||||
"Input was not large enough to exercise the clamp; increase scale"
|
||||
)
|
||||
|
||||
# Verify gate clamping semantics with a controlled scalar case.
|
||||
# gate=large_val is clamped to limit first, then silu(limit) * 1.0.
|
||||
x_gate = torch.tensor(
|
||||
[[swiglu_limit * 20.0, 1.0]], dtype=torch.float32, device=device
|
||||
)
|
||||
out_gate = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_gate)
|
||||
expected_gate = torch.nn.functional.silu(
|
||||
torch.tensor(swiglu_limit, dtype=torch.float32)
|
||||
).item()
|
||||
torch.testing.assert_close(
|
||||
out_gate,
|
||||
torch.tensor([[expected_gate]], dtype=torch.float32, device=device),
|
||||
atol=1e-3,
|
||||
rtol=1e-3,
|
||||
)
|
||||
|
||||
# Verify up clamping semantics: up >> limit gets clamped to limit.
|
||||
x_up = torch.tensor(
|
||||
[[1.0, swiglu_limit * 20.0]], dtype=torch.float32, device=device
|
||||
)
|
||||
out_up = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_up)
|
||||
silu_1 = torch.nn.functional.silu(torch.tensor(1.0)).item()
|
||||
torch.testing.assert_close(
|
||||
out_up,
|
||||
torch.tensor([[silu_1 * swiglu_limit]], dtype=torch.float32, device=device),
|
||||
atol=1e-3,
|
||||
rtol=1e-3,
|
||||
)
|
||||
|
||||
# opcheck
|
||||
out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device)
|
||||
opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"activation",
|
||||
[
|
||||
(FastGELU, torch.ops._C.gelu_fast),
|
||||
(NewGELU, torch.ops._C.gelu_new),
|
||||
(QuickGELU, torch.ops._C.gelu_quick),
|
||||
(ReLUSquaredActivation, torch.ops._C.relu_squared),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("d", D)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_activation(
|
||||
default_vllm_config,
|
||||
activation: type[torch.nn.Module],
|
||||
num_tokens: int,
|
||||
d: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
x = torch.randn(num_tokens, d, dtype=dtype)
|
||||
layer = activation[0]()
|
||||
fn = activation[1]
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
torch.testing.assert_close(
|
||||
out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
|
||||
)
|
||||
|
||||
out = torch.empty_like(x)
|
||||
opcheck(fn, (out, x))
|
||||
@@ -0,0 +1,203 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for ApplyRotaryEmb CustomOp dispatch behavior.
|
||||
|
||||
This test ensures that RotaryEmbedding classes correctly call the appropriate
|
||||
ApplyRotaryEmb methods based on the calling context:
|
||||
|
||||
1. RotaryEmbedding.forward_native() -> ApplyRotaryEmb.forward_native()
|
||||
2. RotaryEmbedding.forward_cuda() -> ApplyRotaryEmb.forward() (auto-dispatch)
|
||||
3. RotaryEmbedding.forward_hip() -> ApplyRotaryEmb.forward() (auto-dispatch)
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config import (
|
||||
CompilationConfig,
|
||||
VllmConfig,
|
||||
get_cached_compilation_config,
|
||||
set_current_vllm_config,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
CUDA_DEVICES = ["cuda:0"]
|
||||
|
||||
|
||||
@dataclass
|
||||
class RotaryEmbeddingTestCase:
|
||||
"""Test case configuration for RotaryEmbedding dispatch tests."""
|
||||
|
||||
name: str
|
||||
rope_class: type
|
||||
rope_kwargs: dict
|
||||
method_name: str # forward_native, forward_cuda, forward
|
||||
positions_shape: tuple # (num_tokens,) or (3, num_tokens) or (4, num_tokens)
|
||||
expect_forward_native: bool # Should call ApplyRotaryEmb.forward_native()
|
||||
expect_forward: bool # Should call ApplyRotaryEmb.forward()
|
||||
|
||||
|
||||
def get_test_cases() -> list[RotaryEmbeddingTestCase]:
|
||||
"""Generate test cases for all RotaryEmbedding classes."""
|
||||
from vllm.model_executor.layers.rotary_embedding.ernie45_vl_rope import (
|
||||
Ernie4_5_VLRotaryEmbedding,
|
||||
)
|
||||
from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding
|
||||
from vllm.model_executor.layers.rotary_embedding.xdrope import XDRotaryEmbedding
|
||||
|
||||
common_kwargs = {
|
||||
"head_size": 128,
|
||||
"rotary_dim": 128,
|
||||
"max_position_embeddings": 4096,
|
||||
"base": 10000,
|
||||
"is_neox_style": True,
|
||||
"dtype": torch.bfloat16,
|
||||
}
|
||||
|
||||
return [
|
||||
# MRotaryEmbedding tests
|
||||
RotaryEmbeddingTestCase(
|
||||
name="MRotaryEmbedding.forward_native",
|
||||
rope_class=MRotaryEmbedding,
|
||||
rope_kwargs={**common_kwargs, "mrope_section": [16, 24, 24]},
|
||||
method_name="forward_native",
|
||||
positions_shape=(3, 32), # 2D for multimodal
|
||||
expect_forward_native=True,
|
||||
expect_forward=False,
|
||||
),
|
||||
RotaryEmbeddingTestCase(
|
||||
name="MRotaryEmbedding.forward_cuda_1d",
|
||||
rope_class=MRotaryEmbedding,
|
||||
rope_kwargs={**common_kwargs, "mrope_section": [16, 24, 24]},
|
||||
method_name="forward_cuda",
|
||||
positions_shape=(32,), # 1D triggers apply_rotary_emb path
|
||||
expect_forward_native=False,
|
||||
expect_forward=True,
|
||||
),
|
||||
# XDRotaryEmbedding tests
|
||||
RotaryEmbeddingTestCase(
|
||||
name="XDRotaryEmbedding.forward",
|
||||
rope_class=XDRotaryEmbedding,
|
||||
rope_kwargs={
|
||||
**common_kwargs,
|
||||
"scaling_alpha": 1.0,
|
||||
"xdrope_section": [16, 16, 16, 16],
|
||||
},
|
||||
method_name="forward",
|
||||
positions_shape=(4, 32), # 4D for P/W/H/T
|
||||
expect_forward_native=False,
|
||||
expect_forward=True,
|
||||
),
|
||||
# Ernie4_5_VLRotaryEmbedding tests
|
||||
RotaryEmbeddingTestCase(
|
||||
name="Ernie4_5_VLRotaryEmbedding.forward_native",
|
||||
rope_class=Ernie4_5_VLRotaryEmbedding,
|
||||
rope_kwargs={**common_kwargs, "mrope_section": [22, 22, 20]},
|
||||
method_name="forward_native",
|
||||
positions_shape=(3, 32), # 2D for multimodal
|
||||
expect_forward_native=True,
|
||||
expect_forward=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def run_dispatch_test(
|
||||
test_case: RotaryEmbeddingTestCase,
|
||||
device: str,
|
||||
):
|
||||
"""Run a dispatch test for a RotaryEmbedding class."""
|
||||
vllm_config = VllmConfig(
|
||||
compilation_config=CompilationConfig(custom_ops=["all", "+apply_rotary_emb"])
|
||||
)
|
||||
get_cached_compilation_config.cache_clear()
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
rope = test_case.rope_class(**test_case.rope_kwargs).to(device=device)
|
||||
|
||||
apply_rotary_emb = rope.apply_rotary_emb
|
||||
|
||||
# Verify custom op is enabled
|
||||
if test_case.expect_forward_native:
|
||||
assert (
|
||||
apply_rotary_emb._forward_method != apply_rotary_emb.forward_native
|
||||
), "Test setup error: ApplyRotaryEmb custom op should be enabled"
|
||||
|
||||
# Setup call tracking
|
||||
call_tracker = {"forward_native_called": False, "forward_called": False}
|
||||
original_forward_native = apply_rotary_emb.forward_native
|
||||
original_forward = apply_rotary_emb.forward
|
||||
|
||||
def tracked_forward_native(*args, **kwargs):
|
||||
call_tracker["forward_native_called"] = True
|
||||
return original_forward_native(*args, **kwargs)
|
||||
|
||||
def tracked_forward(*args, **kwargs):
|
||||
call_tracker["forward_called"] = True
|
||||
return original_forward(*args, **kwargs)
|
||||
|
||||
apply_rotary_emb.forward_native = tracked_forward_native
|
||||
apply_rotary_emb.forward = tracked_forward
|
||||
|
||||
try:
|
||||
num_tokens = test_case.positions_shape[-1]
|
||||
num_q_heads = 8
|
||||
num_kv_heads = 2
|
||||
head_size = test_case.rope_kwargs["head_size"]
|
||||
max_position = test_case.rope_kwargs["max_position_embeddings"]
|
||||
|
||||
positions = torch.randint(
|
||||
0, max_position // 4, test_case.positions_shape, device=device
|
||||
)
|
||||
query = torch.randn(
|
||||
num_tokens, num_q_heads * head_size, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
key = torch.randn(
|
||||
num_tokens,
|
||||
num_kv_heads * head_size,
|
||||
dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Call the method under test
|
||||
method = getattr(rope, test_case.method_name)
|
||||
method(positions, query.clone(), key.clone())
|
||||
|
||||
# Verify expectations
|
||||
if test_case.expect_forward_native:
|
||||
assert call_tracker["forward_native_called"], (
|
||||
f"{test_case.name} should call ApplyRotaryEmb.forward_native()"
|
||||
)
|
||||
if not test_case.expect_forward:
|
||||
assert not call_tracker["forward_called"], (
|
||||
f"{test_case.name} should NOT call ApplyRotaryEmb.forward(). "
|
||||
"Bug: when +apply_rotary_emb is enabled, forward_native() "
|
||||
"incorrectly dispatches to CUDA/HIP kernels."
|
||||
)
|
||||
if test_case.expect_forward:
|
||||
assert call_tracker["forward_called"], (
|
||||
f"{test_case.name} should call ApplyRotaryEmb.forward()"
|
||||
)
|
||||
finally:
|
||||
apply_rotary_emb.forward_native = original_forward_native
|
||||
apply_rotary_emb.forward = original_forward
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests."
|
||||
)
|
||||
@pytest.mark.parametrize("test_case", get_test_cases(), ids=lambda tc: tc.name)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_rotary_embedding_dispatch(
|
||||
test_case: RotaryEmbeddingTestCase,
|
||||
device: str,
|
||||
):
|
||||
"""
|
||||
Test that RotaryEmbedding classes dispatch to the correct ApplyRotaryEmb method.
|
||||
|
||||
- forward_native methods should call ApplyRotaryEmb.forward_native()
|
||||
- forward_cuda/forward methods should call ApplyRotaryEmb.forward()
|
||||
"""
|
||||
run_dispatch_test(test_case, device)
|
||||
@@ -0,0 +1,70 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the batched-weight RMS norm kernel (vllm._custom_ops.rms_norm).
|
||||
|
||||
``rms_norm`` can use the outermost input batch index to select the corresponding
|
||||
weight row. The result must match that of looping ``rms_norm`` over that dimension.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="rms_norm requires a CUDA/ROCm device",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[
|
||||
(28, 17, 128), # 3D: [num_rows, tokens, hidden]
|
||||
(1, 5, 2, 128), # 4D: single row (edge case)
|
||||
(28, 13, 8, 128), # 4D: [L, num_ctx, nkv, hd] (DFlash K-norm)
|
||||
(6, 3, 4, 769), # 4D: non-power-of-two hidden size
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16, torch.float])
|
||||
@pytest.mark.parametrize("seed", [42])
|
||||
@torch.inference_mode()
|
||||
def test_rms_norm_matches_loop(
|
||||
shape: tuple[int, ...], dtype: torch.dtype, seed: int
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
num_rows, hidden = shape[0], shape[-1]
|
||||
eps = 1e-6
|
||||
|
||||
x = torch.randn(*shape, dtype=dtype) * 0.1
|
||||
# Distinct weight per row so that a wrong row index would be caught.
|
||||
weight = torch.randn(num_rows, hidden, dtype=dtype) * 0.1 + 1.0
|
||||
|
||||
# Reference batched-weight rms norm.
|
||||
out_ref = torch.empty_like(x)
|
||||
for i in range(x.shape[0]):
|
||||
ops.rms_norm(out_ref[i], x[i], weight[i], eps)
|
||||
|
||||
out = torch.empty_like(x)
|
||||
ops.rms_norm(out, x, weight, eps)
|
||||
|
||||
# Expect bitwise-identical results.
|
||||
torch.testing.assert_close(out, out_ref, atol=0, rtol=0)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_rms_norm_validates_shapes() -> None:
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
x = torch.randn(4, 8, 128, dtype=torch.float)
|
||||
out = torch.empty_like(x)
|
||||
# Expect num rows mismatch.
|
||||
with pytest.raises(RuntimeError):
|
||||
ops.rms_norm(out, x, torch.randn(3, 128), 1e-6)
|
||||
# Expect hidden size mismatch.
|
||||
with pytest.raises(RuntimeError):
|
||||
ops.rms_norm(out, x, torch.randn(4, 64), 1e-6)
|
||||
@@ -0,0 +1,212 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.allclose_default import get_default_atol, get_default_rtol
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm.platforms import CpuArchEnum, current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not current_platform.is_cpu():
|
||||
pytest.skip("skipping CPU-only tests", allow_module_level=True)
|
||||
|
||||
from vllm.model_executor.layers.activation import (
|
||||
GELU,
|
||||
FastGELU,
|
||||
GeluAndMul,
|
||||
NewGELU,
|
||||
QuickGELU,
|
||||
SiluAndMul,
|
||||
)
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float32]
|
||||
NUM_TOKENS = [7, 83]
|
||||
D = [512, 2048]
|
||||
SEEDS = [0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("activation_cls", "fn"),
|
||||
[
|
||||
(SiluAndMul, torch.ops._C.silu_and_mul),
|
||||
(GeluAndMul, torch.ops._C.gelu_and_mul),
|
||||
(GeluAndMul, torch.ops._C.gelu_tanh_and_mul),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("d", D)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_cpu_act_and_mul(
|
||||
default_vllm_config,
|
||||
activation_cls: type[torch.nn.Module],
|
||||
fn: object,
|
||||
num_tokens: int,
|
||||
d: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
x = torch.randn(num_tokens, 2 * d, dtype=dtype)
|
||||
|
||||
layer = activation_cls()
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
|
||||
torch.testing.assert_close(
|
||||
out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
|
||||
)
|
||||
|
||||
output_shape = x.shape[:-1] + (x.shape[-1] // 2,)
|
||||
raw_out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
||||
opcheck(fn, (raw_out, x))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("activation_cls", "fn", "op_args"),
|
||||
[
|
||||
(NewGELU, torch.ops._C.gelu_new, ()),
|
||||
(FastGELU, torch.ops._C.gelu_fast, ()),
|
||||
(QuickGELU, torch.ops._C.gelu_quick, ()),
|
||||
pytest.param(
|
||||
GELU,
|
||||
getattr(torch.ops._C, "activation_lut_bf16", None),
|
||||
("gelu",),
|
||||
marks=pytest.mark.skipif(
|
||||
current_platform.get_cpu_architecture() != CpuArchEnum.ARM,
|
||||
reason="activation_lut_bf16 is only built on Arm CPU",
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("d", D)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_cpu_unary_activation(
|
||||
default_vllm_config,
|
||||
activation_cls: type[torch.nn.Module],
|
||||
fn: object,
|
||||
op_args: tuple[str, ...],
|
||||
num_tokens: int,
|
||||
d: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
x = torch.randn(num_tokens, d, dtype=dtype)
|
||||
layer = activation_cls()
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
torch.testing.assert_close(
|
||||
out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
|
||||
)
|
||||
# gelu with activation_lut_bf16 only makes sense for BF16
|
||||
if not (activation_cls is GELU and dtype != torch.bfloat16):
|
||||
raw_out = torch.empty_like(x)
|
||||
opcheck(fn, (raw_out, x, *op_args))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_cpu_gelu_tanh_and_mul(
|
||||
default_vllm_config,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
gate = torch.tensor(
|
||||
[
|
||||
[
|
||||
-12.0,
|
||||
-10.0,
|
||||
-9.01,
|
||||
-5.0,
|
||||
-2.0,
|
||||
-1.0,
|
||||
-0.0,
|
||||
0.0,
|
||||
0.5,
|
||||
1.0,
|
||||
2.0,
|
||||
5.0,
|
||||
9.01,
|
||||
10.0,
|
||||
12.0,
|
||||
11.0,
|
||||
],
|
||||
[
|
||||
-7.5,
|
||||
-4.5,
|
||||
-3.0,
|
||||
-1.5,
|
||||
-0.75,
|
||||
-0.25,
|
||||
0.25,
|
||||
0.75,
|
||||
1.5,
|
||||
3.0,
|
||||
4.5,
|
||||
7.5,
|
||||
-11.0,
|
||||
11.0,
|
||||
8.75,
|
||||
-8.75,
|
||||
],
|
||||
],
|
||||
dtype=dtype,
|
||||
)
|
||||
val = torch.tensor(
|
||||
[
|
||||
[
|
||||
0.25,
|
||||
-0.5,
|
||||
0.75,
|
||||
-1.0,
|
||||
1.25,
|
||||
-1.5,
|
||||
1.75,
|
||||
-2.0,
|
||||
2.25,
|
||||
-2.5,
|
||||
2.75,
|
||||
-3.0,
|
||||
3.25,
|
||||
-3.5,
|
||||
3.75,
|
||||
-4.0,
|
||||
],
|
||||
[
|
||||
-0.4,
|
||||
0.6,
|
||||
-0.8,
|
||||
1.0,
|
||||
-1.2,
|
||||
1.4,
|
||||
-1.6,
|
||||
1.8,
|
||||
-2.0,
|
||||
2.2,
|
||||
-2.4,
|
||||
2.6,
|
||||
-2.8,
|
||||
3.0,
|
||||
-3.2,
|
||||
3.4,
|
||||
],
|
||||
],
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
x = torch.cat((val, gate), dim=-1).contiguous()
|
||||
kernel_out = torch.empty_like(val)
|
||||
torch.ops._C.gelu_tanh_and_mul(kernel_out, x)
|
||||
|
||||
torch_ref = torch.nn.functional.gelu(val, approximate="tanh") * gate
|
||||
|
||||
atol = get_default_atol(kernel_out)
|
||||
rtol = get_default_rtol(kernel_out)
|
||||
torch.testing.assert_close(kernel_out, torch_ref, atol=atol, rtol=rtol)
|
||||
@@ -0,0 +1,109 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the manual AllReduce + GemmaRMSNorm fusion used by MiniMax M3.
|
||||
|
||||
``fused_allreduce_gemma_rms_norm`` must match the unfused model path, i.e.
|
||||
``GemmaRMSNorm(all_reduce(partial), residual)``, both on the flashinfer fast
|
||||
path (TP>1 with flashinfer + NVSwitch) and on the eager fallback (TP==1, or when
|
||||
flashinfer is unavailable / the GPU has no NVSwitch).
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch.multiprocessing import spawn
|
||||
|
||||
from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
|
||||
from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
|
||||
fused_allreduce_gemma_rms_norm,
|
||||
)
|
||||
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
|
||||
@ensure_current_vllm_config()
|
||||
def _worker_fused_ar_norm(
|
||||
local_rank,
|
||||
world_size,
|
||||
port,
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
dtype,
|
||||
seed,
|
||||
eps,
|
||||
):
|
||||
"""Per-rank worker: compare the fused helper vs all_reduce + GemmaRMSNorm."""
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(
|
||||
world_size, 1, local_rank, port, local_rank=local_rank
|
||||
)
|
||||
|
||||
# Norm weights are identical across ranks (replicated GemmaRMSNorm).
|
||||
set_random_seed(seed)
|
||||
norm = GemmaRMSNorm(hidden_size, eps=eps).cuda().to(dtype)
|
||||
with torch.no_grad():
|
||||
norm.weight.normal_(mean=0.0, std=0.1)
|
||||
|
||||
# Residual is shared across ranks; the partial o_proj output differs per rank
|
||||
# (each rank holds a partial sum that all_reduce combines).
|
||||
torch.manual_seed(seed + 7)
|
||||
residual = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
||||
torch.manual_seed(seed + 1000 + local_rank)
|
||||
partial = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
||||
|
||||
# Reference: the unfused model path.
|
||||
reduced = tensor_model_parallel_all_reduce(partial.clone())
|
||||
ref_out, ref_res = norm(reduced, residual.clone())
|
||||
|
||||
# Fused helper (flashinfer fast path when available, else fallback).
|
||||
out, res = fused_allreduce_gemma_rms_norm(partial.clone(), residual.clone(), norm)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
torch.testing.assert_close(out, ref_out, atol=2e-2, rtol=2e-2)
|
||||
torch.testing.assert_close(res, ref_res, atol=2e-2, rtol=2e-2)
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="CUDA required",
|
||||
)
|
||||
# world_size=1 exercises the TP==1 identity branch on a single GPU; >1 exercises
|
||||
# the all_reduce + GemmaRMSNorm equivalence (flashinfer kernel or fallback).
|
||||
@pytest.mark.parametrize("world_size", [1, 2, 4])
|
||||
@pytest.mark.parametrize("num_tokens", [1, 128, 333])
|
||||
@pytest.mark.parametrize("hidden_size", [2048, 4096])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("eps", [1e-6])
|
||||
@pytest.mark.parametrize("seed", [42])
|
||||
def test_fused_allreduce_gemma_rms_norm(
|
||||
world_size,
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
dtype,
|
||||
eps,
|
||||
seed,
|
||||
):
|
||||
num_gpus = current_platform.device_count()
|
||||
if num_gpus < world_size:
|
||||
pytest.skip(f"Need >= {world_size} GPUs, have {num_gpus}")
|
||||
port = str(get_open_port())
|
||||
spawn(
|
||||
_worker_fused_ar_norm,
|
||||
args=(
|
||||
world_size,
|
||||
port,
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
dtype,
|
||||
seed,
|
||||
eps,
|
||||
),
|
||||
nprocs=world_size,
|
||||
join=True,
|
||||
)
|
||||
@@ -0,0 +1,81 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Correctness + large-token-count launch tests for fused_q_kv_rmsnorm.
|
||||
|
||||
Before the grid-dim fix the kernel used grid ``(2, num_tokens)``, which hit
|
||||
CUDA's 65535 grid-y cap for ``num_tokens >= 65536`` and failed with
|
||||
``Triton Error [CUDA]: invalid argument`` at every large chunked-prefill
|
||||
profile run. These tests pin the new grid layout.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.models.deepseek_v4.common.ops import fused_q_kv_rmsnorm
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="fused_q_kv_rmsnorm requires a CUDA/ROCm device",
|
||||
)
|
||||
|
||||
|
||||
def _ref_rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
|
||||
x_f32 = x.to(torch.float32)
|
||||
variance = x_f32.pow(2).mean(dim=-1, keepdim=True)
|
||||
y = x_f32 * torch.rsqrt(variance + eps) * w.to(torch.float32)
|
||||
return y.to(x.dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 17, 1024, 8192])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
def test_fused_q_kv_rmsnorm_correctness(num_tokens: int, dtype: torch.dtype):
|
||||
torch.manual_seed(0)
|
||||
device = "cuda"
|
||||
q_size, kv_size = 192, 576
|
||||
qr = torch.randn(num_tokens, q_size, dtype=dtype, device=device)
|
||||
kv = torch.randn(num_tokens, kv_size, dtype=dtype, device=device)
|
||||
qw = torch.randn(q_size, dtype=dtype, device=device)
|
||||
kvw = torch.randn(kv_size, dtype=dtype, device=device)
|
||||
eps = 1e-6
|
||||
|
||||
qr_out, kv_out = fused_q_kv_rmsnorm(qr, kv, qw, kvw, eps)
|
||||
|
||||
qr_ref = _ref_rmsnorm(qr, qw, eps)
|
||||
kv_ref = _ref_rmsnorm(kv, kvw, eps)
|
||||
|
||||
tol = dict(rtol=1e-2, atol=1e-2)
|
||||
torch.testing.assert_close(qr_out, qr_ref, **tol)
|
||||
torch.testing.assert_close(kv_out, kv_ref, **tol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [65535, 65536, 131072])
|
||||
def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int):
|
||||
"""Regression guard: grid used to be (2, num_tokens), hitting CUDA's
|
||||
65535 grid-y cap at num_tokens >= 65536. The new grid (num_tokens, 2)
|
||||
lifts that bound to 2**31-1."""
|
||||
device = "cuda"
|
||||
dtype = torch.bfloat16
|
||||
q_size, kv_size = 192, 576
|
||||
qr = torch.randn(num_tokens, q_size, dtype=dtype, device=device)
|
||||
kv = torch.randn(num_tokens, kv_size, dtype=dtype, device=device)
|
||||
qw = torch.randn(q_size, dtype=dtype, device=device)
|
||||
kvw = torch.randn(kv_size, dtype=dtype, device=device)
|
||||
|
||||
qr_out, kv_out = fused_q_kv_rmsnorm(qr, kv, qw, kvw, 1e-6)
|
||||
# spot-check a couple of rows against the torch reference
|
||||
for row in (0, num_tokens // 2, num_tokens - 1):
|
||||
torch.testing.assert_close(
|
||||
qr_out[row],
|
||||
_ref_rmsnorm(qr[row : row + 1], qw, 1e-6)[0],
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
kv_out[row],
|
||||
_ref_rmsnorm(kv[row : row + 1], kvw, 1e-6)[0],
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
)
|
||||
@@ -0,0 +1,146 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float16]
|
||||
IS_NEOX = [True, False]
|
||||
EPS_VALUES = [1e-5, 1e-6]
|
||||
SEEDS = [13]
|
||||
PARTIAL_ROPE = [True, False]
|
||||
CUDA_DEVICES = ["cuda:0"]
|
||||
|
||||
|
||||
def _apply_qk_norm_rope(
|
||||
qkv: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
q_norm: RMSNorm,
|
||||
k_norm: RMSNorm,
|
||||
rope: RotaryEmbedding,
|
||||
num_heads_q: int,
|
||||
num_heads_kv: int,
|
||||
head_dim: int,
|
||||
) -> torch.Tensor:
|
||||
q_size = num_heads_q * head_dim
|
||||
kv_size = num_heads_kv * head_dim
|
||||
|
||||
q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
|
||||
|
||||
q_by_head = q.view(*q.shape[:-1], q.shape[-1] // head_dim, head_dim)
|
||||
q_by_head = q_norm.forward_native(q_by_head)
|
||||
q = q_by_head.view(q.shape)
|
||||
|
||||
k_by_head = k.view(*k.shape[:-1], k.shape[-1] // head_dim, head_dim)
|
||||
k_by_head = k_norm.forward_native(k_by_head)
|
||||
k = k_by_head.view(k.shape)
|
||||
|
||||
q, k = rope.forward_native(positions, q, k)
|
||||
return torch.cat([q, k, v], dim=-1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="fused_qk_norm_rope custom op requires cuda and rocm platform",
|
||||
)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("is_neox", IS_NEOX)
|
||||
@pytest.mark.parametrize("eps", EPS_VALUES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("rotary_ratio", [1.0, 0.5, 0.25])
|
||||
@torch.inference_mode()
|
||||
def test_fused_qk_norm_rope_matches_reference(
|
||||
default_vllm_config,
|
||||
device: str,
|
||||
dtype: torch.dtype,
|
||||
is_neox: bool,
|
||||
eps: float,
|
||||
seed: int,
|
||||
rotary_ratio: float,
|
||||
):
|
||||
torch.set_default_device(device)
|
||||
set_random_seed(seed)
|
||||
num_heads, num_kv_heads, head_dim = 16, 4, 128
|
||||
num_tokens = 4
|
||||
|
||||
total_dim = (num_heads + 2 * num_kv_heads) * head_dim
|
||||
qkv_base = torch.randn(num_tokens, total_dim, dtype=dtype, device=device)
|
||||
qkv_fused = qkv_base.clone()
|
||||
positions = torch.arange(num_tokens, dtype=torch.long, device=device)
|
||||
|
||||
q_norm = RMSNorm(head_dim, eps=eps).to(device=device, dtype=dtype)
|
||||
k_norm = RMSNorm(head_dim, eps=eps).to(device=device, dtype=dtype)
|
||||
q_norm.weight.data.normal_(mean=1.0, std=0.1)
|
||||
k_norm.weight.data.normal_(mean=1.0, std=0.1)
|
||||
q_weight = q_norm.weight.data
|
||||
k_weight = k_norm.weight.data
|
||||
rotary_dim = int(head_dim * rotary_ratio)
|
||||
rope = RotaryEmbedding(
|
||||
head_size=head_dim,
|
||||
rotary_dim=rotary_dim,
|
||||
max_position_embeddings=4096,
|
||||
base=10000.0,
|
||||
is_neox_style=is_neox,
|
||||
dtype=dtype,
|
||||
).to(device)
|
||||
|
||||
ref_result = _apply_qk_norm_rope(
|
||||
qkv=qkv_base,
|
||||
positions=positions,
|
||||
q_norm=q_norm,
|
||||
k_norm=k_norm,
|
||||
rope=rope,
|
||||
num_heads_q=num_heads,
|
||||
num_heads_kv=num_kv_heads,
|
||||
head_dim=head_dim,
|
||||
)
|
||||
|
||||
opcheck(
|
||||
torch.ops._C.fused_qk_norm_rope,
|
||||
(
|
||||
qkv_fused.clone(),
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope.cos_sin_cache,
|
||||
is_neox,
|
||||
positions.view(-1),
|
||||
),
|
||||
)
|
||||
|
||||
torch.ops._C.fused_qk_norm_rope(
|
||||
qkv_fused,
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope.cos_sin_cache,
|
||||
is_neox,
|
||||
positions.view(-1),
|
||||
)
|
||||
|
||||
if dtype == torch.float16:
|
||||
ATOL, RTOL = (2e-3, 2e-3)
|
||||
else:
|
||||
ATOL, RTOL = (1e-2, 1e-2)
|
||||
|
||||
torch.testing.assert_close(
|
||||
qkv_fused,
|
||||
ref_result,
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
)
|
||||
@@ -0,0 +1,335 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import itertools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from tests.kernels.utils import fp8_ulp_distance, opcheck
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
per_token_group_quant_fp8,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.int8_utils import (
|
||||
per_token_group_quant_int8,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float]
|
||||
QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]
|
||||
VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
|
||||
# Avoid combinatorial explosion with full Cartesian product
|
||||
NUM_TOKENS_HIDDEN_SIZES = [
|
||||
*[(1, i) for i in [1, 64, 128, *VEC_HIDDEN_SIZES, 5120, 5137]],
|
||||
*[(2048, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5137]],
|
||||
*[(4096, i) for i in [1, 64, 5137]],
|
||||
]
|
||||
|
||||
ADD_RESIDUAL = [False, True]
|
||||
SCALE_UBS = [True, False]
|
||||
GROUP_SIZES = [None, [1, 64], [1, 128]]
|
||||
TMA_ALIGNMENTS = [0, 4]
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
|
||||
EPS = 1e-6
|
||||
|
||||
## Helpers
|
||||
|
||||
|
||||
def as_float32_tensor(x: float | torch.Tensor) -> torch.Tensor:
|
||||
return torch.as_tensor(x, dtype=torch.float32, device="cuda")
|
||||
|
||||
|
||||
def ref_rms_norm(
|
||||
rms_norm_layer: RMSNorm, x: torch.Tensor, residual: torch.Tensor | None
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
if residual is not None:
|
||||
residual = residual.clone()
|
||||
out, residual = rms_norm_layer.forward_native(x, residual)
|
||||
else:
|
||||
out = rms_norm_layer.forward_native(x)
|
||||
|
||||
return out, residual
|
||||
|
||||
|
||||
def ref_dynamic_per_token_or_block_quant(
|
||||
rms_norm_layer: RMSNorm,
|
||||
x: torch.Tensor,
|
||||
quant_dtype: torch.dtype,
|
||||
residual: torch.Tensor | None,
|
||||
scale_ub: torch.Tensor | None,
|
||||
group_size: list[int] | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
if scale_ub is not None:
|
||||
assert quant_dtype == current_platform.fp8_dtype()
|
||||
|
||||
# Norm
|
||||
torch_out, residual = ref_rms_norm(rms_norm_layer, x, residual)
|
||||
|
||||
# Quant
|
||||
if group_size is not None:
|
||||
if quant_dtype == current_platform.fp8_dtype():
|
||||
torch_out, scales = per_token_group_quant_fp8(
|
||||
torch_out, group_size=group_size[1], use_ue8m0=False
|
||||
)
|
||||
else:
|
||||
assert quant_dtype == torch.int8
|
||||
torch_out, scales = per_token_group_quant_int8(
|
||||
torch_out, group_size=group_size[1]
|
||||
)
|
||||
else:
|
||||
if quant_dtype == current_platform.fp8_dtype():
|
||||
torch_out, scales = ops.scaled_fp8_quant(
|
||||
torch_out, scale_ub=scale_ub, use_per_token_if_dynamic=True
|
||||
)
|
||||
else:
|
||||
assert quant_dtype == torch.int8
|
||||
torch_out, scales, _ = ops.scaled_int8_quant(torch_out)
|
||||
|
||||
return torch_out, scales, residual
|
||||
|
||||
|
||||
def ref_impl(
|
||||
rms_norm_layer: RMSNorm,
|
||||
x: torch.Tensor,
|
||||
quant_dtype: torch.dtype,
|
||||
residual: torch.Tensor | None,
|
||||
scale_ub: torch.Tensor | None,
|
||||
group_size: list[int] | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
return ref_dynamic_per_token_or_block_quant(
|
||||
rms_norm_layer, x, quant_dtype, residual, scale_ub, group_size
|
||||
)
|
||||
|
||||
|
||||
def ops_dynamic_per_token_or_block_quant(
|
||||
weight: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
quant_dtype: torch.dtype,
|
||||
residual: torch.Tensor | None,
|
||||
scale_ub: torch.Tensor | None,
|
||||
group_size: list[int] | None,
|
||||
tma_alignment: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
if residual is not None:
|
||||
residual = residual.clone()
|
||||
if group_size is not None:
|
||||
out, scales = ops.rms_norm_per_block_quant(
|
||||
x,
|
||||
weight,
|
||||
EPS,
|
||||
quant_dtype,
|
||||
group_size,
|
||||
scale_ub,
|
||||
residual,
|
||||
True,
|
||||
tma_alignment,
|
||||
)
|
||||
scales = scales.contiguous()
|
||||
else:
|
||||
out, scales = ops.rms_norm_dynamic_per_token_quant(
|
||||
x, weight, EPS, quant_dtype, scale_ub, residual
|
||||
)
|
||||
return out, scales, residual
|
||||
|
||||
|
||||
def ops_impl(
|
||||
weight: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
quant_dtype: torch.dtype,
|
||||
residual: torch.Tensor | None,
|
||||
scale_ub: torch.Tensor | None,
|
||||
group_size: list[int] | None,
|
||||
tma_alignment: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
return ops_dynamic_per_token_or_block_quant(
|
||||
weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
|
||||
@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
|
||||
@pytest.mark.parametrize(
|
||||
"group_size, tma_alignment",
|
||||
[(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
|
||||
)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("strided_input", [False, True])
|
||||
@torch.inference_mode()
|
||||
def test_rms_norm(
|
||||
default_vllm_config,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
add_residual: bool,
|
||||
has_scale_ub: bool,
|
||||
dtype: torch.dtype,
|
||||
quant_dtype: torch.dtype,
|
||||
group_size: list[int] | None,
|
||||
tma_alignment: int,
|
||||
seed: int,
|
||||
device: str,
|
||||
strided_input: bool,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
torch.accelerator.set_device_index(device)
|
||||
|
||||
if group_size is not None and hidden_size % group_size[1] != 0:
|
||||
# skip
|
||||
pytest.skip("Skip non-divisible group sizes")
|
||||
|
||||
if group_size is not None and has_scale_ub:
|
||||
# blockwise baseline doesn't support scale_ub
|
||||
pytest.skip("scale_ub not supported for blockwise/group quantization")
|
||||
|
||||
if (
|
||||
group_size is None or quant_dtype != current_platform.fp8_dtype()
|
||||
) and tma_alignment != 0:
|
||||
# TMA alignment is only supported for groupwise fp8 kernels
|
||||
pytest.skip("tma alignment not supported for per-token or int8 quantization")
|
||||
|
||||
if (
|
||||
group_size is not None
|
||||
and tma_alignment != 0
|
||||
and hidden_size // group_size[1] % tma_alignment == 0
|
||||
):
|
||||
# Skip tests where TMA alignment doesn't create extra padding to save time
|
||||
pytest.skip("Skip TMA alignment cases where no extra padding is added")
|
||||
|
||||
if has_scale_ub and quant_dtype != current_platform.fp8_dtype():
|
||||
# skip
|
||||
pytest.skip("scale_ub only supported for fp8 quantization")
|
||||
|
||||
layer = RMSNorm(hidden_size, EPS).to(dtype=dtype)
|
||||
|
||||
# Make weights
|
||||
layer.weight.data.normal_(mean=1.0, std=0.1)
|
||||
|
||||
# Make inputs: use a wider tensor and slice to create a non-contiguous
|
||||
# (strided) input when strided_input=True. The last dimension stride
|
||||
# remains 1, which the kernel requires.
|
||||
scale = 1 / (hidden_size)
|
||||
last_dim = 2 * hidden_size if strided_input else hidden_size
|
||||
x = torch.randn(num_tokens, last_dim, dtype=dtype) * scale
|
||||
x = x[:, :hidden_size]
|
||||
|
||||
# dim 1 gets special-cased
|
||||
x_is_strided = strided_input and num_tokens != 1
|
||||
# check that the input is strided iff we expect it to be
|
||||
assert x.is_contiguous() != x_is_strided
|
||||
|
||||
# Residual must still be contiguous
|
||||
residual = (
|
||||
torch.randn(num_tokens, hidden_size, dtype=dtype) * scale
|
||||
if add_residual
|
||||
else None
|
||||
)
|
||||
if has_scale_ub:
|
||||
rms_x, _ = ref_rms_norm(layer, x, residual)
|
||||
scale_ub = torch.mean(rms_x).to(dtype=torch.float32, device="cuda")
|
||||
else:
|
||||
scale_ub = None
|
||||
|
||||
ref_out, ref_scales, ref_residual = ref_impl(
|
||||
layer, x, quant_dtype, residual, scale_ub, group_size
|
||||
)
|
||||
ops_out, ops_scales, ops_residual = ops_impl(
|
||||
layer.weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
|
||||
)
|
||||
|
||||
assert ref_out.dtype == quant_dtype
|
||||
assert ops_out.dtype == quant_dtype
|
||||
|
||||
# Per-block bf16 scales: allow a small relative tolerance for a few groups
|
||||
# whose abs-max flips by one ULP between the fused and reference paths. The
|
||||
# per-token and fp32 paths stay strict.
|
||||
relax_block_rocm = (
|
||||
group_size is not None
|
||||
and dtype == torch.bfloat16
|
||||
and current_platform.is_rocm()
|
||||
)
|
||||
|
||||
def scales_close(rtol: float, atol: float) -> bool:
|
||||
if torch.allclose(ref_scales, ops_scales, rtol=rtol, atol=atol):
|
||||
return True
|
||||
return relax_block_rocm and torch.allclose(
|
||||
ref_scales, ops_scales, rtol=1e-2, atol=atol
|
||||
)
|
||||
|
||||
if quant_dtype == torch.int8:
|
||||
assert scales_close(rtol=1e-5, atol=1e-6)
|
||||
# big atol to account for round-off errors.
|
||||
assert torch.allclose(ref_out, ops_out, atol=1)
|
||||
else:
|
||||
assert scales_close(rtol=1e-5, atol=1e-8)
|
||||
a = ref_out.to(dtype=torch.float32)
|
||||
b = ops_out.to(dtype=torch.float32)
|
||||
ok = torch.allclose(a, b, atol=1e-6)
|
||||
if not ok:
|
||||
if relax_block_rocm:
|
||||
# ULP-flipped group scale can cross an E4M3 tie; tolerate a
|
||||
# bounded count of isolated fp8 outliers.
|
||||
ulp = fp8_ulp_distance(ref_out, ops_out)
|
||||
max_outliers = ulp.numel() // 100_000 + 8
|
||||
ok = int((ulp > 0).sum().item()) <= max_outliers
|
||||
else:
|
||||
# CUDA (& non-bf16): compare dequantized values with relaxed tolerance.
|
||||
if group_size is None:
|
||||
a_deq = a * ref_scales.view(-1, 1)
|
||||
b_deq = b * ops_scales.view(-1, 1)
|
||||
else:
|
||||
a_deq = a * ref_scales.repeat_interleave(group_size[1], dim=1)
|
||||
b_deq = b * ops_scales.repeat_interleave(group_size[1], dim=1)
|
||||
# NOTE: It is possible that some future test cases trigger this
|
||||
# max diff due to precision issues. If such an error is
|
||||
# encountered, it's recommended to inspect the differences between
|
||||
# all corresponding elements from each tensor (e.g. by looping over
|
||||
# them) and checking how many the max diff error shows up on (just
|
||||
# a few bad elements should still be considered acceptable).
|
||||
ok = torch.allclose(a_deq, b_deq, rtol=5e-2, atol=5e-2)
|
||||
assert ok
|
||||
if add_residual:
|
||||
assert torch.allclose(ref_residual, ops_residual)
|
||||
|
||||
output = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
|
||||
if group_size is None:
|
||||
scales = torch.empty(
|
||||
(x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32
|
||||
)
|
||||
opcheck(
|
||||
torch.ops._C.rms_norm_dynamic_per_token_quant,
|
||||
(output, x, layer.weight, scales, 1e-5, scale_ub, residual),
|
||||
)
|
||||
else:
|
||||
assert hidden_size % group_size[1] == 0
|
||||
num_groups = hidden_size // group_size[1]
|
||||
scales = torch.empty(
|
||||
(num_groups, num_tokens),
|
||||
device=x.device,
|
||||
dtype=torch.float32,
|
||||
).transpose(0, 1)
|
||||
opcheck(
|
||||
torch.ops._C.rms_norm_per_block_quant,
|
||||
(
|
||||
output,
|
||||
x,
|
||||
layer.weight,
|
||||
scales,
|
||||
1e-5,
|
||||
scale_ub,
|
||||
residual,
|
||||
group_size[1],
|
||||
True, # is_scale_transposed
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,103 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Tests that FusedRMSNormGated decomposes correctly under torch.compile,
|
||||
matching the eager triton kernel output."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.fla.ops.kda import FusedRMSNormGated
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
DTYPES = [torch.bfloat16]
|
||||
HIDDEN_SIZES = [128, 512]
|
||||
NUM_TOKENS = [64, 128]
|
||||
ACTIVATIONS = ["swish", "sigmoid"]
|
||||
ELEMENTWISE_AFFINE = [True, False]
|
||||
SEEDS = [0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("activation", ACTIVATIONS)
|
||||
@pytest.mark.parametrize("elementwise_affine", ELEMENTWISE_AFFINE)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_compiled_vs_eager(
|
||||
default_vllm_config,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
activation: str,
|
||||
elementwise_affine: bool,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
) -> None:
|
||||
"""forward_native decomposition matches forward_cuda triton kernel."""
|
||||
torch._dynamo.reset()
|
||||
set_random_seed(seed)
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
module = FusedRMSNormGated(
|
||||
hidden_size,
|
||||
elementwise_affine=elementwise_affine,
|
||||
eps=1e-5,
|
||||
activation=activation,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
||||
g = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
||||
|
||||
# forward_cuda may modify x in-place, so clone inputs
|
||||
cuda_out = module.forward_cuda(x.clone(), g.clone())
|
||||
compiled_native = torch.compile(module.forward_native, fullgraph=True)
|
||||
native_out = compiled_native(x.clone(), g.clone())
|
||||
|
||||
torch.testing.assert_close(native_out, cuda_out, atol=1e-3, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[
|
||||
(1, 16, 32, 128),
|
||||
(2, 8, 16, 64),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("activation", ACTIVATIONS)
|
||||
@pytest.mark.parametrize("elementwise_affine", ELEMENTWISE_AFFINE)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_compiled_vs_eager_multidim(
|
||||
default_vllm_config,
|
||||
shape: tuple,
|
||||
activation: str,
|
||||
elementwise_affine: bool,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
) -> None:
|
||||
"""forward_native decomposition handles multi-dimensional inputs."""
|
||||
torch._dynamo.reset()
|
||||
set_random_seed(seed)
|
||||
device = torch.device("cuda:0")
|
||||
head_dim = shape[-1]
|
||||
|
||||
module = FusedRMSNormGated(
|
||||
head_dim,
|
||||
elementwise_affine=elementwise_affine,
|
||||
eps=1e-5,
|
||||
activation=activation,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
x = torch.randn(*shape, dtype=dtype, device=device)
|
||||
g = torch.randn(*shape, dtype=dtype, device=device)
|
||||
|
||||
# forward_cuda may modify x in-place, so clone inputs
|
||||
cuda_out = module.forward_cuda(x.clone(), g.clone())
|
||||
compiled_native = torch.compile(module.forward_native, fullgraph=True)
|
||||
native_out = compiled_native(x.clone(), g.clone())
|
||||
|
||||
torch.testing.assert_close(native_out, cuda_out, atol=1e-3, rtol=1e-2)
|
||||
@@ -0,0 +1,189 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
per_token_group_quant_fp8,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.int8_utils import (
|
||||
per_token_group_quant_int8,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
QUANT_DTYPES = [current_platform.fp8_dtype(), torch.int8]
|
||||
VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
|
||||
NUM_TOKENS_HIDDEN_SIZES = [
|
||||
*[(1, i) for i in [64, *VEC_HIDDEN_SIZES, 2048, 5120]],
|
||||
*[(16, i) for i in [64, *VEC_HIDDEN_SIZES, 5120]],
|
||||
*[(128, i) for i in [64, *VEC_HIDDEN_SIZES]],
|
||||
*[(512, i) for i in [64, 5120]],
|
||||
]
|
||||
SCALE_UBS = [False]
|
||||
GROUP_SIZES = [64, 128]
|
||||
IS_SCALE_TRANSPOSED = [False, True]
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [i for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
|
||||
|
||||
|
||||
def ref_silu_and_mul_per_block_quant(
|
||||
x: torch.Tensor,
|
||||
quant_dtype: torch.dtype,
|
||||
group_size: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Reference implementation: unfused SiLU+Mul then group quantization."""
|
||||
hidden = x.shape[-1] // 2
|
||||
gate, up = x.split(hidden, dim=-1)
|
||||
silu_out = F.silu(gate) * up
|
||||
|
||||
if quant_dtype == current_platform.fp8_dtype():
|
||||
return per_token_group_quant_fp8(
|
||||
silu_out, group_size=group_size, use_ue8m0=False
|
||||
)
|
||||
elif quant_dtype == torch.int8:
|
||||
return per_token_group_quant_int8(silu_out, group_size=group_size)
|
||||
else:
|
||||
raise ValueError(f"Unsupported quant_dtype: {quant_dtype}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
|
||||
@pytest.mark.parametrize("group_size", GROUP_SIZES)
|
||||
@pytest.mark.parametrize("is_scale_transposed", IS_SCALE_TRANSPOSED)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device_idx", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_silu_and_mul_per_block_quant(
|
||||
default_vllm_config,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
has_scale_ub: bool,
|
||||
dtype: torch.dtype,
|
||||
quant_dtype: torch.dtype,
|
||||
group_size: int,
|
||||
is_scale_transposed: bool,
|
||||
seed: int,
|
||||
device_idx: str,
|
||||
) -> None:
|
||||
"""Test SiLU+Mul+Block Quantization kernel correctness."""
|
||||
torch.accelerator.set_device_index(device_idx)
|
||||
device = f"cuda:{device_idx}"
|
||||
torch.random.manual_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
|
||||
if hidden_size % group_size != 0:
|
||||
return
|
||||
|
||||
if has_scale_ub:
|
||||
pytest.skip("Scale upper bound not yet supported")
|
||||
|
||||
scale = 1 / hidden_size
|
||||
x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) * scale
|
||||
|
||||
# Reference implementation
|
||||
ref_out, ref_scales = ref_silu_and_mul_per_block_quant(x, quant_dtype, group_size)
|
||||
|
||||
# Fused kernel implementation
|
||||
ops_out, ops_scales = ops.silu_and_mul_per_block_quant(
|
||||
x, group_size, quant_dtype, None, is_scale_transposed
|
||||
)
|
||||
|
||||
# Check for NaN/Inf
|
||||
assert not torch.isnan(ops_out.float()).any(), "Kernel output contains NaN"
|
||||
assert not torch.isinf(ops_out.float()).any(), "Kernel output contains Inf"
|
||||
assert not torch.isnan(ops_scales).any(), "Kernel scales contain NaN"
|
||||
assert not torch.isinf(ops_scales).any(), "Kernel scales contain Inf"
|
||||
|
||||
# Check dtypes
|
||||
assert ref_out.dtype == quant_dtype
|
||||
assert ops_out.dtype == quant_dtype
|
||||
|
||||
# Check scales match
|
||||
torch.testing.assert_close(ref_scales, ops_scales, rtol=1e-5, atol=1e-5)
|
||||
|
||||
# Check output correctness via dequantized values
|
||||
ref_scales_expanded = ref_scales.repeat_interleave(group_size, dim=1)
|
||||
ops_scales_expanded = ops_scales.repeat_interleave(group_size, dim=1)
|
||||
ref_deq = ref_out.to(dtype=torch.float32) * ref_scales_expanded
|
||||
ops_deq = ops_out.to(dtype=torch.float32) * ops_scales_expanded
|
||||
torch.testing.assert_close(ref_deq, ops_deq, atol=5e-2, rtol=5e-2)
|
||||
|
||||
# opcheck
|
||||
output = torch.empty(num_tokens, hidden_size, device=device, dtype=quant_dtype)
|
||||
num_groups = hidden_size // group_size
|
||||
if is_scale_transposed:
|
||||
scales = torch.empty(num_groups, num_tokens, device=device, dtype=torch.float32)
|
||||
else:
|
||||
scales = torch.empty(num_tokens, num_groups, device=device, dtype=torch.float32)
|
||||
opcheck(
|
||||
torch.ops._C.silu_and_mul_per_block_quant,
|
||||
(output, x, scales, group_size, None, is_scale_transposed),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16])
|
||||
@pytest.mark.parametrize("hidden_size", [4096])
|
||||
@pytest.mark.parametrize("num_tokens", [128])
|
||||
@pytest.mark.parametrize("group_size", [128])
|
||||
def test_silu_block_quant_shapes(
|
||||
default_vllm_config,
|
||||
dtype: torch.dtype,
|
||||
hidden_size: int,
|
||||
num_tokens: int,
|
||||
group_size: int,
|
||||
):
|
||||
"""Test that output shapes are correct."""
|
||||
torch.set_default_device("cuda")
|
||||
x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device="cuda")
|
||||
|
||||
# Row-major scales
|
||||
out, scales = ops.silu_and_mul_per_block_quant(
|
||||
x,
|
||||
group_size=group_size,
|
||||
quant_dtype=current_platform.fp8_dtype(),
|
||||
is_scale_transposed=False,
|
||||
)
|
||||
assert out.shape == (num_tokens, hidden_size)
|
||||
assert scales.shape == (num_tokens, hidden_size // group_size)
|
||||
|
||||
# Column-major scales (logical shape same after .t() in _custom_ops)
|
||||
out, scales = ops.silu_and_mul_per_block_quant(
|
||||
x,
|
||||
group_size=group_size,
|
||||
quant_dtype=current_platform.fp8_dtype(),
|
||||
is_scale_transposed=True,
|
||||
)
|
||||
assert out.shape == (num_tokens, hidden_size)
|
||||
assert scales.shape == (num_tokens, hidden_size // group_size)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16])
|
||||
@pytest.mark.parametrize("batch_size", [1, 16, 256])
|
||||
@pytest.mark.parametrize("hidden_size", [1024, 5120, 14336])
|
||||
def test_silu_block_quant_edge_cases(
|
||||
default_vllm_config, dtype: torch.dtype, batch_size: int, hidden_size: int
|
||||
):
|
||||
"""Test edge cases: single token, large batch, large hidden size."""
|
||||
torch.set_default_device("cuda")
|
||||
x = torch.randn(batch_size, hidden_size * 2, dtype=dtype, device="cuda")
|
||||
|
||||
out, scales = ops.silu_and_mul_per_block_quant(
|
||||
x,
|
||||
group_size=128,
|
||||
quant_dtype=current_platform.fp8_dtype(),
|
||||
is_scale_transposed=False,
|
||||
)
|
||||
|
||||
assert out.shape == (batch_size, hidden_size)
|
||||
assert out.dtype == current_platform.fp8_dtype()
|
||||
assert scales.dtype == torch.float32
|
||||
assert not torch.isnan(out.float()).any()
|
||||
assert not torch.isnan(scales).any()
|
||||
assert not torch.isinf(scales).any()
|
||||
@@ -0,0 +1,250 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.quant_utils import FP8_DTYPE
|
||||
from tests.kernels.utils import fp8_ulp_distance, opcheck
|
||||
from vllm import ir
|
||||
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if current_platform.is_rocm():
|
||||
from vllm.platforms.rocm import on_gfx90a
|
||||
|
||||
on_mi250 = on_gfx90a()
|
||||
else:
|
||||
on_mi250 = False
|
||||
|
||||
DTYPES = [torch.half, torch.bfloat16, torch.float]
|
||||
NUM_TOKENS = [7, 83, 4096] # Arbitrary values for testing
|
||||
HIDDEN_SIZES = [8, 768, 769, 5120, 5125, 8192] # Arbitrary values for testing
|
||||
ADD_RESIDUAL = [False, True] if not on_mi250 else [True]
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
|
||||
|
||||
def _rms_norm_tolerance(dtype: torch.dtype) -> dict[str, float]:
|
||||
return ir.ops.rms_norm.get_tolerance(dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("strided_input", [False, True])
|
||||
@torch.inference_mode()
|
||||
def test_rms_norm(
|
||||
default_vllm_config,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
add_residual: bool,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
strided_input: bool,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
layer = RMSNorm(hidden_size).to(dtype=dtype)
|
||||
layer.weight.data.normal_(mean=1.0, std=0.1)
|
||||
scale = 1 / (2 * hidden_size)
|
||||
last_dim = 2 * hidden_size if strided_input else hidden_size
|
||||
x = torch.randn(num_tokens, last_dim, dtype=dtype)
|
||||
x = x[..., :hidden_size]
|
||||
assert x.is_contiguous() != strided_input
|
||||
x *= scale
|
||||
residual = torch.randn_like(x) * scale if add_residual else None
|
||||
|
||||
# NOTE(woosuk): The reference implementation should be executed first
|
||||
# because the custom kernel is in-place.
|
||||
ref_out = layer.forward_native(x, residual)
|
||||
out = layer(x, residual)
|
||||
# NOTE(woosuk): LayerNorm operators (including RMS) typically have larger
|
||||
# numerical errors than other operators because they involve reductions.
|
||||
# Therefore, we use a larger tolerance.
|
||||
if add_residual:
|
||||
torch.testing.assert_close(out[0], ref_out[0], atol=1e-2, rtol=1e-2)
|
||||
torch.testing.assert_close(out[1], ref_out[1], atol=1e-2, rtol=1e-2)
|
||||
else:
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
|
||||
|
||||
if residual is not None:
|
||||
opcheck(
|
||||
torch.ops._C.fused_add_rms_norm,
|
||||
(x, residual, layer.weight.data, layer.variance_epsilon),
|
||||
)
|
||||
else:
|
||||
opcheck(
|
||||
torch.ops._C.rms_norm, (out, x, layer.weight.data, layer.variance_epsilon)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_rms_norm_weightless(
|
||||
default_vllm_config,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
add_residual: bool,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
layer = RMSNorm(hidden_size, has_weight=False).to(dtype=dtype)
|
||||
x = torch.randn(num_tokens, hidden_size, dtype=dtype)
|
||||
residual = torch.randn_like(x) if add_residual else None
|
||||
|
||||
ref_out = layer.forward_native(x, residual)
|
||||
out = layer(x, residual)
|
||||
tol = _rms_norm_tolerance(dtype)
|
||||
if add_residual:
|
||||
torch.testing.assert_close(out[0], ref_out[0], **tol)
|
||||
torch.testing.assert_close(out[1], ref_out[1], **tol)
|
||||
else:
|
||||
torch.testing.assert_close(out, ref_out, **tol)
|
||||
|
||||
if residual is not None:
|
||||
opcheck(
|
||||
torch.ops._C.fused_add_rms_norm,
|
||||
(x, residual, None, layer.variance_epsilon),
|
||||
)
|
||||
else:
|
||||
opcheck(
|
||||
torch.ops._C.rms_norm,
|
||||
(out, x, None, layer.variance_epsilon),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("quant_scale", [0.01, 1.0, 10.0])
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("strided_input", [False, True])
|
||||
def test_fused_rms_norm_quant(
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
add_residual: bool,
|
||||
dtype: torch.dtype,
|
||||
quant_scale: float,
|
||||
seed: int,
|
||||
device: str,
|
||||
strided_input: bool,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
|
||||
weight = torch.empty(hidden_size, dtype=dtype).normal_(mean=1.0, std=0.1)
|
||||
scale = 1 / (2 * hidden_size)
|
||||
last_dim = 2 * hidden_size if strided_input else hidden_size
|
||||
x_base = torch.randn(num_tokens, last_dim, dtype=dtype)
|
||||
x = x_base[..., :hidden_size]
|
||||
assert x.is_contiguous() != strided_input
|
||||
|
||||
x *= scale
|
||||
if add_residual:
|
||||
residual = torch.randn_like(x) * scale
|
||||
residual_fused = residual.clone()
|
||||
else:
|
||||
residual = residual_fused = None
|
||||
|
||||
out_norm = torch.empty_like(x)
|
||||
out_quant = torch.empty_like(x, dtype=FP8_DTYPE)
|
||||
out_quant_fused = torch.empty_like(out_quant)
|
||||
|
||||
quant_scale_t = torch.tensor(quant_scale, dtype=torch.float32)
|
||||
|
||||
if add_residual:
|
||||
torch.ops._C.fused_add_rms_norm_static_fp8_quant(
|
||||
out_quant_fused, x, residual_fused, weight, quant_scale_t, 1e-6
|
||||
)
|
||||
|
||||
# Unfused kernel is in-place so it goes second
|
||||
# Also use a separate clone of x to avoid modifying the input
|
||||
x_unfused_base = x_base.clone()
|
||||
x_unfused = x_unfused_base[..., :hidden_size]
|
||||
assert x_unfused.is_contiguous() != strided_input
|
||||
torch.ops._C.fused_add_rms_norm(x_unfused, residual, weight, 1e-6)
|
||||
torch.ops._C.static_scaled_fp8_quant(
|
||||
out_quant, x_unfused.contiguous(), quant_scale_t
|
||||
)
|
||||
|
||||
torch.accelerator.synchronize()
|
||||
torch.testing.assert_close(residual_fused, residual, atol=1e-2, rtol=1e-2)
|
||||
opcheck(
|
||||
torch.ops._C.fused_add_rms_norm_static_fp8_quant,
|
||||
(out_quant_fused, x, residual_fused, weight, quant_scale_t, 1e-6),
|
||||
)
|
||||
else:
|
||||
torch.ops._C.rms_norm_static_fp8_quant(
|
||||
out_quant_fused, x, weight, quant_scale_t, 1e-6
|
||||
)
|
||||
|
||||
torch.ops._C.rms_norm(out_norm, x, weight, 1e-6)
|
||||
torch.ops._C.static_scaled_fp8_quant(out_quant, out_norm, quant_scale_t)
|
||||
|
||||
opcheck(
|
||||
torch.ops._C.rms_norm_static_fp8_quant,
|
||||
(out_quant_fused, x, weight, quant_scale_t, 1e-6),
|
||||
)
|
||||
|
||||
if current_platform.is_rocm():
|
||||
# Fused and unfused FP8 paths can land on opposite sides of an E4M3 tie;
|
||||
# tolerate a tiny number of isolated fp8 outliers on ROCm.
|
||||
ulp = fp8_ulp_distance(out_quant, out_quant_fused)
|
||||
max_outliers = ulp.numel() // 100_000 + 8
|
||||
num_outliers = int((ulp > 0).sum().item())
|
||||
assert num_outliers <= max_outliers, (
|
||||
f"FP8 quant mismatch: {num_outliers} fp8 outliers (allowed {max_outliers})"
|
||||
)
|
||||
else:
|
||||
torch.testing.assert_close(
|
||||
out_quant.to(dtype=torch.float32),
|
||||
out_quant_fused.to(dtype=torch.float32),
|
||||
atol=1e-3,
|
||||
rtol=1e-3,
|
||||
)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_gemma_rms_norm_mixed_input_weight_dtype(default_vllm_config) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required")
|
||||
|
||||
device = CUDA_DEVICES[0]
|
||||
torch.set_default_device(device)
|
||||
|
||||
num_tokens, hidden_size = 32, 1024
|
||||
x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
|
||||
layer = GemmaRMSNorm(hidden_size, eps=1e-6).to(device=device)
|
||||
layer.weight.data.normal_(mean=0.0, std=0.1)
|
||||
|
||||
# Gemma uses fp32 weight parameter while activations can be bf16.
|
||||
assert layer.weight.dtype == torch.float32
|
||||
out = layer(x)
|
||||
|
||||
x_fp32 = x.float()
|
||||
weight_fp32 = layer.weight.data.float() + 1.0
|
||||
variance = x_fp32.pow(2).mean(dim=-1, keepdim=True)
|
||||
ref = (x_fp32 * torch.rsqrt(variance + layer.variance_epsilon) * weight_fp32).to(
|
||||
x.dtype
|
||||
)
|
||||
|
||||
assert out.dtype == x.dtype
|
||||
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
|
||||
@@ -0,0 +1,208 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for MiniMax QK RMS-norm: NCCL reference vs Lamport fused kernel."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.multiprocessing import spawn
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.model_executor.layers.minimax_rms_norm import (
|
||||
MiniMaxText01RMSNormTP,
|
||||
rms_norm_tp,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
|
||||
@ensure_current_vllm_config()
|
||||
def _worker_forward_qk(
|
||||
local_rank,
|
||||
world_size,
|
||||
port,
|
||||
num_tokens,
|
||||
hidden_q_full,
|
||||
hidden_k_full,
|
||||
dtype,
|
||||
seed,
|
||||
eps,
|
||||
):
|
||||
"""Per-rank worker: compare NCCL allreduce path vs Lamport fused kernel."""
|
||||
|
||||
if not hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
|
||||
cleanup_dist_env_and_memory()
|
||||
return
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(
|
||||
world_size, 1, local_rank, port, local_rank=local_rank
|
||||
)
|
||||
|
||||
hq = hidden_q_full // world_size
|
||||
hk = hidden_k_full // world_size
|
||||
|
||||
q_norm = MiniMaxText01RMSNormTP(hidden_q_full, eps=eps).cuda()
|
||||
k_norm = MiniMaxText01RMSNormTP(hidden_k_full, eps=eps).cuda()
|
||||
|
||||
set_random_seed(seed)
|
||||
qw = torch.randn(hidden_q_full, dtype=dtype, device="cuda")
|
||||
kw = torch.randn(hidden_k_full, dtype=dtype, device="cuda")
|
||||
q_norm.weight = nn.Parameter(qw[local_rank * hq : (local_rank + 1) * hq])
|
||||
k_norm.weight = nn.Parameter(kw[local_rank * hk : (local_rank + 1) * hk])
|
||||
|
||||
torch.manual_seed(seed + 1000 + local_rank)
|
||||
qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda")
|
||||
|
||||
# Reference: eager all-reduce path. ``forward_qk`` no longer all-reduces
|
||||
# the variance (it is the tp==1 / already-reduced building block), so the
|
||||
# multi-rank reference must use the eager path that performs the global
|
||||
# variance all-reduce, matching the fused kernel below.
|
||||
ref_q, ref_k = rms_norm_tp._minimax_qk_norm_tp_eager(
|
||||
qkv.clone(),
|
||||
q_norm.weight,
|
||||
k_norm.weight,
|
||||
hq,
|
||||
hk,
|
||||
world_size,
|
||||
eps,
|
||||
)
|
||||
|
||||
# Set up Lamport workspace.
|
||||
from vllm.distributed.parallel_state import get_tp_group
|
||||
from vllm.model_executor.layers.minimax_rms_norm.lamport_workspace import (
|
||||
get_allreduce_workspace,
|
||||
)
|
||||
|
||||
workspace = get_allreduce_workspace(
|
||||
rank=local_rank,
|
||||
world_size=world_size,
|
||||
max_tokens=num_tokens,
|
||||
process_group=get_tp_group().cpu_group,
|
||||
)
|
||||
|
||||
opcheck(
|
||||
torch.ops._C.minimax_allreduce_rms_qk,
|
||||
(
|
||||
qkv.clone(),
|
||||
q_norm.weight,
|
||||
k_norm.weight,
|
||||
workspace,
|
||||
hq,
|
||||
hk,
|
||||
local_rank,
|
||||
world_size,
|
||||
eps,
|
||||
),
|
||||
)
|
||||
fused_q, fused_k = torch.ops._C.minimax_allreduce_rms_qk(
|
||||
qkv.clone(),
|
||||
q_norm.weight,
|
||||
k_norm.weight,
|
||||
workspace,
|
||||
hq,
|
||||
hk,
|
||||
local_rank,
|
||||
world_size,
|
||||
eps,
|
||||
)
|
||||
_, _, fused_v = qkv.split([hq, hk, hk], dim=-1)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
fused_q,
|
||||
ref_q,
|
||||
atol=3e-2,
|
||||
rtol=3e-2,
|
||||
)
|
||||
torch.testing.assert_close(fused_k, ref_k, atol=3e-2, rtol=3e-2)
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="CUDA required",
|
||||
)
|
||||
@pytest.mark.parametrize("world_size", [2, 4, 8])
|
||||
@pytest.mark.parametrize("num_tokens", [1, 128, 333])
|
||||
@pytest.mark.parametrize(
|
||||
"hidden_dims",
|
||||
[(6144, 1024)],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("eps", [1e-6])
|
||||
@pytest.mark.parametrize("seed", [42])
|
||||
def test_minimax_reduce_rms_qk(
|
||||
world_size,
|
||||
num_tokens,
|
||||
hidden_dims,
|
||||
dtype,
|
||||
eps,
|
||||
seed,
|
||||
):
|
||||
num_gpus = current_platform.device_count()
|
||||
if num_gpus < world_size:
|
||||
pytest.skip(f"Need >= {world_size} GPUs, have {num_gpus}")
|
||||
hidden_q_full, hidden_k_full = hidden_dims
|
||||
port = str(get_open_port())
|
||||
spawn(
|
||||
_worker_forward_qk,
|
||||
args=(
|
||||
world_size,
|
||||
port,
|
||||
num_tokens,
|
||||
hidden_q_full,
|
||||
hidden_k_full,
|
||||
dtype,
|
||||
seed,
|
||||
eps,
|
||||
),
|
||||
nprocs=world_size,
|
||||
join=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda() or not HAS_TRITON,
|
||||
reason="CUDA and Triton required",
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", [1, 7, 128, 333, 2049])
|
||||
@pytest.mark.parametrize("hidden_dims", [(3072, 512), (768, 256), (3000, 500)])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("tp_world", [1, 4, 8])
|
||||
@pytest.mark.parametrize("eps", [1e-6])
|
||||
@pytest.mark.parametrize("seed", [42])
|
||||
def test_minimax_qk_norm_triton_fallback(
|
||||
monkeypatch, num_tokens, hidden_dims, dtype, tp_world, eps, seed
|
||||
):
|
||||
"""Single-GPU check: Triton fallback kernels vs the pure-torch reference.
|
||||
|
||||
The all-reduce is a TP communication barrier, so it is monkeypatched to
|
||||
identity here; both the Triton path and the reference see the same
|
||||
(patched) reduction. This validates the kernel math and the folded
|
||||
``/ tp_world`` scaling without needing multiple ranks -- ``hidden_dims``
|
||||
are the per-rank q/k segment widths.
|
||||
"""
|
||||
monkeypatch.setattr(rms_norm_tp, "_all_reduce_variance", lambda v: v)
|
||||
|
||||
q_size, kv_size = hidden_dims
|
||||
device = "cuda"
|
||||
torch.manual_seed(seed)
|
||||
qkv = torch.randn(num_tokens, q_size + 2 * kv_size, dtype=dtype, device=device)
|
||||
q_weight = torch.randn(q_size, dtype=dtype, device=device)
|
||||
k_weight = torch.randn(kv_size, dtype=dtype, device=device)
|
||||
|
||||
q_triton, k_triton = rms_norm_tp._minimax_qk_norm_tp_fallback(
|
||||
qkv, q_weight, k_weight, q_size, kv_size, 0, tp_world, eps
|
||||
)
|
||||
q_ref, k_ref = rms_norm_tp._minimax_qk_norm_tp_eager(
|
||||
qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps
|
||||
)
|
||||
|
||||
torch.testing.assert_close(q_triton, q_ref, atol=3e-2, rtol=3e-2)
|
||||
torch.testing.assert_close(k_triton, k_ref, atol=3e-2, rtol=3e-2)
|
||||
@@ -0,0 +1,236 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from typing import NamedTuple
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.transformers_utils.config import get_config
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
def generate_test_data(
|
||||
num_tokens: int,
|
||||
num_q_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
max_position_embeddings: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
):
|
||||
"""Generate test data for given configuration."""
|
||||
set_random_seed(42)
|
||||
# Create 2D positions (3, num_tokens) for multimodal case
|
||||
positions = torch.randint(
|
||||
0, max_position_embeddings // 4, (3, num_tokens), device=device
|
||||
)
|
||||
|
||||
# Create query and key tensors
|
||||
query = torch.randn(num_tokens, num_q_heads * head_size, dtype=dtype, device=device)
|
||||
key = torch.randn(num_tokens, num_kv_heads * head_size, dtype=dtype, device=device)
|
||||
|
||||
return positions, query, key
|
||||
|
||||
|
||||
class MRoPETestInfo(NamedTuple):
|
||||
model_name: str
|
||||
# https://github.com/pytorch/pytorch/blob/main/torch/testing/_comparison.py#L1317
|
||||
atol: float = 1e-2
|
||||
rtol: float = 1.6e-2
|
||||
marks: list[pytest.MarkDecorator] = []
|
||||
|
||||
|
||||
MODELS_TO_TEST = [
|
||||
MRoPETestInfo(model_name="zai-org/GLM-4.1V-9B-Thinking"),
|
||||
MRoPETestInfo(model_name="Qwen/Qwen2-VL-7B-Instruct"),
|
||||
MRoPETestInfo(model_name="Qwen/Qwen2-VL-72B-Instruct"),
|
||||
MRoPETestInfo(model_name="Qwen/Qwen2.5-VL-72B-Instruct"),
|
||||
MRoPETestInfo(model_name="Qwen/Qwen3-VL-4B-Instruct"),
|
||||
MRoPETestInfo(model_name="Qwen/Qwen3-VL-30B-A3B-Instruct"),
|
||||
]
|
||||
|
||||
num_tokens_list = [11, 8192]
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests."
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"model_info, model_name",
|
||||
[
|
||||
pytest.param(test_config, test_config.model_name, marks=test_config.marks)
|
||||
for test_config in MODELS_TO_TEST
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("tp_size", [1, 2])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("num_tokens", num_tokens_list)
|
||||
def test_mrope(
|
||||
default_vllm_config,
|
||||
model_name: str,
|
||||
model_info: MRoPETestInfo,
|
||||
tp_size: int,
|
||||
dtype: torch.dtype,
|
||||
num_tokens: int,
|
||||
):
|
||||
atol = model_info.atol
|
||||
rtol = model_info.rtol
|
||||
|
||||
config = get_config(model_name, False).get_text_config()
|
||||
|
||||
# get the model config
|
||||
total_num_kv_heads = config.num_key_value_heads
|
||||
total_num_heads = config.num_attention_heads
|
||||
num_heads = total_num_heads // tp_size
|
||||
num_kv_heads = max(1, total_num_kv_heads // tp_size)
|
||||
head_dim = (
|
||||
config.head_dim
|
||||
if hasattr(config, "head_dim")
|
||||
else config.hidden_size // total_num_heads
|
||||
)
|
||||
is_neox_style = True
|
||||
|
||||
max_position = config.max_position_embeddings
|
||||
|
||||
mrope_helper_class = get_rope(
|
||||
head_size=head_dim,
|
||||
max_position=max_position,
|
||||
is_neox_style=is_neox_style,
|
||||
rope_parameters=config.rope_parameters,
|
||||
dtype=dtype,
|
||||
).to(device=device)
|
||||
|
||||
# create q k v input tensors
|
||||
# create rotary pos emb input tensors
|
||||
positions, query, key = generate_test_data(
|
||||
num_tokens, num_heads, num_kv_heads, head_dim, max_position, dtype, device
|
||||
)
|
||||
|
||||
query_native, key_native = mrope_helper_class.forward_native(
|
||||
positions,
|
||||
query.clone(),
|
||||
key.clone(),
|
||||
)
|
||||
|
||||
query_cuda, key_cuda = mrope_helper_class.forward_cuda(
|
||||
positions,
|
||||
query.clone(),
|
||||
key.clone(),
|
||||
)
|
||||
|
||||
torch.testing.assert_close(query_native, query_cuda, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(key_native, key_cuda, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests."
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"model_info, model_name",
|
||||
[
|
||||
pytest.param(test_config, test_config.model_name, marks=test_config.marks)
|
||||
for test_config in MODELS_TO_TEST
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("tp_size", [1, 2])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("num_tokens", num_tokens_list)
|
||||
def test_mrope_torch_compile_tracing(
|
||||
default_vllm_config,
|
||||
model_name: str,
|
||||
model_info: MRoPETestInfo,
|
||||
tp_size: int,
|
||||
dtype: torch.dtype,
|
||||
num_tokens: int,
|
||||
):
|
||||
atol = model_info.atol
|
||||
rtol = model_info.rtol
|
||||
|
||||
config = get_config(model_name, False).get_text_config()
|
||||
|
||||
# get the model config
|
||||
total_num_kv_heads = config.num_key_value_heads
|
||||
total_num_heads = config.num_attention_heads
|
||||
num_heads = total_num_heads // tp_size
|
||||
num_kv_heads = max(1, total_num_kv_heads // tp_size)
|
||||
head_dim = (
|
||||
config.head_dim
|
||||
if hasattr(config, "head_dim")
|
||||
else config.hidden_size // total_num_heads
|
||||
)
|
||||
is_neox_style = True
|
||||
max_position = config.max_position_embeddings
|
||||
|
||||
mrope_helper_class = get_rope(
|
||||
head_size=head_dim,
|
||||
max_position=max_position,
|
||||
is_neox_style=is_neox_style,
|
||||
rope_parameters=config.rope_parameters,
|
||||
dtype=dtype,
|
||||
).to(device=device)
|
||||
|
||||
# Generate test data
|
||||
positions, query, key = generate_test_data(
|
||||
num_tokens, num_heads, num_kv_heads, head_dim, max_position, dtype, device
|
||||
)
|
||||
|
||||
# Create a wrapper that makes the in-place function appear functional
|
||||
def functional_forward_cuda(pos, q, k):
|
||||
"""Wrapper that converts in-place operation to functional style
|
||||
|
||||
CUDA Graph does not support in-place operations.
|
||||
This wrapper creates working copies of the
|
||||
input tensors and modifies them.
|
||||
"""
|
||||
q_work = q.clone() # Create working copies
|
||||
k_work = k.clone()
|
||||
# Your in-place function modifies q_work and k_work
|
||||
mrope_helper_class.forward_cuda(pos, q_work, k_work)
|
||||
return q_work, k_work # Return the modified tensors
|
||||
|
||||
# Get reference results
|
||||
query_native, key_native = mrope_helper_class.forward_native(
|
||||
positions,
|
||||
query.clone(),
|
||||
key.clone(),
|
||||
)
|
||||
|
||||
try:
|
||||
compiled_forward_cuda = torch.compile(
|
||||
functional_forward_cuda,
|
||||
fullgraph=True,
|
||||
backend="inductor",
|
||||
mode="reduce-overhead",
|
||||
dynamic=False,
|
||||
)
|
||||
|
||||
# Run compiled version
|
||||
query_compiled_cuda, key_compiled_cuda = compiled_forward_cuda(
|
||||
positions,
|
||||
query,
|
||||
key,
|
||||
)
|
||||
|
||||
# Run original version for comparison
|
||||
query_cuda = query.clone()
|
||||
key_cuda = key.clone()
|
||||
mrope_helper_class.forward_cuda(positions, query_cuda, key_cuda)
|
||||
|
||||
# Verify results
|
||||
torch.testing.assert_close(
|
||||
query_compiled_cuda, query_cuda, atol=atol, rtol=rtol
|
||||
)
|
||||
torch.testing.assert_close(key_compiled_cuda, key_cuda, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(
|
||||
query_compiled_cuda, query_native, atol=atol, rtol=rtol
|
||||
)
|
||||
torch.testing.assert_close(key_compiled_cuda, key_native, atol=atol, rtol=rtol)
|
||||
|
||||
print("✓ forward_cuda successfully traced with torch.compile inductor")
|
||||
|
||||
except Exception as e:
|
||||
pytest.fail(f"forward_cuda failed to trace with torch.compile inductor: {e}")
|
||||
@@ -0,0 +1,26 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for miscellaneous utilities
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
|
||||
|
||||
def test_convert_fp8_opcheck():
|
||||
data = torch.randn((256, 256), dtype=torch.float32, device="cuda")
|
||||
result = torch.empty_like(data, dtype=torch.float8_e4m3fn)
|
||||
opcheck(torch.ops._C_cache_ops.convert_fp8, (result, data, 1.0, "fp8"))
|
||||
|
||||
|
||||
# TODO: Add this back, currently fails with
|
||||
# csrc/cuda_utils_kernels.cu:15 'invalid argument'
|
||||
# @pytest.mark.skipif(not current_platform.is_cuda(),
|
||||
# reason="Only supported for CUDA")
|
||||
# def test_cuda_utils_opcheck():
|
||||
# opcheck(torch.ops._C_cuda_utils.get_device_attribute, (0, 0))
|
||||
# opcheck(
|
||||
# torch.ops._C_cuda_utils.
|
||||
# get_max_shared_memory_per_block_device_attribute, (0, ))
|
||||
@@ -0,0 +1,21 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm._custom_ops import permute_cols
|
||||
|
||||
if not hasattr(torch.ops._C, "permute_cols"):
|
||||
pytest.skip(reason="permute_cols is not supported on ROCm", allow_module_level=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(1, 512), (544, 4096), (67, 8192)])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
def test_permute_cols(shape, dtype):
|
||||
x = torch.randn(shape, dtype=dtype).cuda()
|
||||
perm = torch.randperm(x.shape[1]).to(torch.int).cuda()
|
||||
opcheck(torch.ops._C.permute_cols, (x, perm))
|
||||
y = permute_cols(x, perm)
|
||||
torch.testing.assert_close(y, x[:, perm])
|
||||
@@ -0,0 +1,193 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
from itertools import product
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.allclose_default import get_default_atol, get_default_rtol
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
IS_NEOX_STYLE = [True, False]
|
||||
DTYPES = [torch.bfloat16, torch.float]
|
||||
HEAD_SIZES = [64, 80, 120, 256]
|
||||
ROTARY_DIMS = [None, 32] # None means rotary dim == head size
|
||||
NUM_HEADS = [17] # Arbitrary values for testing
|
||||
BATCH_SIZES = [5] # Arbitrary values for testing
|
||||
SEQ_LENS = [11, 8192] # Arbitrary values for testing
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
USE_KEY = [True, False]
|
||||
|
||||
|
||||
def _get_flat_tensor_shape(
|
||||
batch_size: int, seq_len: int, num_heads: int, head_size: int
|
||||
) -> tuple[int, ...]:
|
||||
return (batch_size, seq_len, num_heads * head_size)
|
||||
|
||||
|
||||
# For testing sliced tensors
|
||||
def _get_padded_tensor_shape(
|
||||
batch_size: int, seq_len: int, num_heads: int, head_size: int
|
||||
) -> tuple[int, ...]:
|
||||
return (batch_size, seq_len, num_heads, head_size + 64)
|
||||
|
||||
|
||||
def _get_batch_tensor_shape(
|
||||
batch_size: int, seq_len: int, num_heads: int, head_size: int
|
||||
) -> tuple[int, ...]:
|
||||
return (batch_size, seq_len, num_heads, head_size)
|
||||
|
||||
|
||||
TENSORS_SHAPES_FN = [
|
||||
_get_batch_tensor_shape,
|
||||
_get_flat_tensor_shape,
|
||||
_get_padded_tensor_shape,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_neox_style", IS_NEOX_STYLE)
|
||||
@pytest.mark.parametrize("tensor_shape_fn", TENSORS_SHAPES_FN)
|
||||
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("rotary_dim", ROTARY_DIMS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("use_key", USE_KEY)
|
||||
@torch.inference_mode()
|
||||
def test_rotary_embedding(
|
||||
default_vllm_config,
|
||||
is_neox_style: bool,
|
||||
tensor_shape_fn: Callable[[int, int, int, int], tuple[int, ...]],
|
||||
batch_size: int,
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
rotary_dim: int | None,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
use_key: bool,
|
||||
max_position: int = 8192,
|
||||
rope_theta: float = 10000,
|
||||
) -> None:
|
||||
if rotary_dim is None:
|
||||
rotary_dim = head_size
|
||||
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
if rotary_dim is None:
|
||||
rotary_dim = head_size
|
||||
rope_parameters = {
|
||||
"rope_type": "default",
|
||||
"rope_theta": rope_theta,
|
||||
"partial_rotary_factor": rotary_dim / head_size,
|
||||
}
|
||||
rope = get_rope(head_size, max_position, is_neox_style, rope_parameters)
|
||||
rope = rope.to(dtype=dtype, device=torch.get_default_device())
|
||||
|
||||
positions = torch.randint(0, max_position, (batch_size, seq_len))
|
||||
query_shape = tensor_shape_fn(batch_size, seq_len, num_heads, head_size)
|
||||
# slice tensor if required, noop otherwise
|
||||
query = torch.randn(query_shape, dtype=dtype)[..., :head_size]
|
||||
key = torch.randn_like(query)[..., :head_size] if use_key else None
|
||||
|
||||
# NOTE(woosuk): The reference implementation should be executed first
|
||||
# because the custom kernel is in-place.
|
||||
ref_query, ref_key = rope.forward_native(positions, query, key)
|
||||
out_query, out_key = rope.forward(positions, query, key)
|
||||
# Compare the results.
|
||||
torch.testing.assert_close(
|
||||
out_query,
|
||||
ref_query,
|
||||
atol=get_default_atol(out_query),
|
||||
rtol=get_default_rtol(out_query),
|
||||
)
|
||||
if use_key:
|
||||
torch.testing.assert_close(
|
||||
out_key,
|
||||
ref_key,
|
||||
atol=get_default_atol(out_key),
|
||||
rtol=get_default_rtol(out_key),
|
||||
)
|
||||
else:
|
||||
assert ref_key is None and out_key is None, "expected returned key to be None"
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_rope_module_cache(default_vllm_config):
|
||||
MAX_POSITIONS = [123, 1234]
|
||||
ROPE_THETAS = [10000, 1000000]
|
||||
ROPE_PARAMETERS = (
|
||||
{"rope_type": "default"},
|
||||
{"rope_type": "linear", "factor": (1,)},
|
||||
{"rope_type": "dynamic", "factor": 1},
|
||||
)
|
||||
settings = (
|
||||
HEAD_SIZES,
|
||||
ROTARY_DIMS,
|
||||
MAX_POSITIONS,
|
||||
ROPE_THETAS,
|
||||
IS_NEOX_STYLE,
|
||||
ROPE_PARAMETERS,
|
||||
DTYPES,
|
||||
)
|
||||
rope_setting_id_map: dict[str, int] = {}
|
||||
for setting in product(*settings):
|
||||
(
|
||||
head_size,
|
||||
rotary_dim,
|
||||
max_position,
|
||||
rope_theta,
|
||||
is_neox_style,
|
||||
rope_parameters,
|
||||
dtype,
|
||||
) = setting
|
||||
if rotary_dim is None:
|
||||
rotary_dim = head_size
|
||||
rope_parameters["rope_theta"] = rope_theta
|
||||
rope_parameters["partial_rotary_factor"] = rotary_dim / head_size
|
||||
rope = get_rope(
|
||||
head_size,
|
||||
max_position,
|
||||
is_neox_style,
|
||||
rope_parameters,
|
||||
dtype,
|
||||
)
|
||||
# different settings cannot share the same rope module
|
||||
assert id(rope) not in rope_setting_id_map.values()
|
||||
assert all(x.dtype == dtype for x in rope.buffers())
|
||||
assert all(x.dtype == dtype for x in rope.parameters())
|
||||
rope_setting_id_map[str(setting)] = id(rope)
|
||||
|
||||
for setting in product(*settings):
|
||||
(
|
||||
head_size,
|
||||
rotary_dim,
|
||||
max_position,
|
||||
rope_theta,
|
||||
is_neox_style,
|
||||
rope_parameters,
|
||||
dtype,
|
||||
) = setting
|
||||
if rotary_dim is None:
|
||||
rotary_dim = head_size
|
||||
rope_parameters["rope_theta"] = rope_theta
|
||||
rope_parameters["partial_rotary_factor"] = rotary_dim / head_size
|
||||
rope = get_rope(
|
||||
head_size,
|
||||
max_position,
|
||||
is_neox_style,
|
||||
rope_parameters,
|
||||
dtype,
|
||||
)
|
||||
# check if cache take effect
|
||||
assert id(rope) == rope_setting_id_map[str(setting)]
|
||||
@@ -0,0 +1,79 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for miscellaneous utilities
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
|
||||
|
||||
|
||||
def rotary_embedding_opcheck(
|
||||
rot,
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor | None = None,
|
||||
):
|
||||
cos_sin_cache = rot.cos_sin_cache.to(query.device, dtype=query.dtype)
|
||||
|
||||
# ops.rotary_embedding() is a in-place operation
|
||||
# that updates the query and key tensors.
|
||||
opcheck(
|
||||
torch.ops._C.rotary_embedding,
|
||||
(positions, query, key, rot.head_size, cos_sin_cache, rot.is_neox_style),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", ["cuda"])
|
||||
@pytest.mark.parametrize("max_position", [11, 4096, 32768])
|
||||
@pytest.mark.parametrize("is_neox_style", [True, False])
|
||||
@pytest.mark.parametrize("rotary_dim", [32])
|
||||
@pytest.mark.parametrize("head_size", [32, 108])
|
||||
@pytest.mark.parametrize("seq_len", [11, 1024])
|
||||
@pytest.mark.parametrize("use_key", [True, False])
|
||||
@pytest.mark.parametrize("head_stride_is_contiguous", [True, False])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
def test_rotary_embedding_opcheck(
|
||||
default_vllm_config,
|
||||
dist_init,
|
||||
device,
|
||||
max_position,
|
||||
is_neox_style,
|
||||
rotary_dim,
|
||||
head_size,
|
||||
seq_len,
|
||||
use_key,
|
||||
head_stride_is_contiguous,
|
||||
dtype,
|
||||
):
|
||||
batch_size = 1
|
||||
base = 10000
|
||||
num_heads = 7
|
||||
rot = RotaryEmbedding(
|
||||
head_size, rotary_dim, max_position, base, is_neox_style, dtype
|
||||
)
|
||||
|
||||
positions = torch.randint(0, max_position, (batch_size, seq_len), device=device)
|
||||
head_stride = head_size + (64 if head_stride_is_contiguous else 0)
|
||||
|
||||
query = torch.randn(
|
||||
batch_size, seq_len, num_heads, head_stride, dtype=dtype, device=device
|
||||
)
|
||||
key = torch.randn_like(query) if use_key else None
|
||||
query = query[..., :head_size]
|
||||
key = key[..., :head_size] if key is not None else None
|
||||
|
||||
rotary_embedding_opcheck(rot, positions, query, key)
|
||||
|
||||
# if we have a contiguous head stride, test the alternate
|
||||
# [..., num_heads * head_dim] shape/layout
|
||||
if head_stride_is_contiguous:
|
||||
rotary_embedding_opcheck(
|
||||
rot,
|
||||
positions,
|
||||
query.flatten(start_dim=-2),
|
||||
key.flatten(start_dim=-2) if key is not None else None,
|
||||
)
|
||||
@@ -0,0 +1,210 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for fused MLA KV-cache write and RoPE fused kernel
|
||||
"""
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.allclose_default import get_default_atol, get_default_rtol
|
||||
from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def default_vllm_config(monkeypatch):
|
||||
"""Enable the AITER triton rope on ROCm for fp16-consistent numerics.
|
||||
|
||||
The fused CUDA kernel runs native fp16 while forward_native upcasts to
|
||||
fp32, so on ROCm we route through the AITER triton rope (+rotary_embedding)
|
||||
to match. Its env gates are cached at import, hence refresh_env_variables().
|
||||
"""
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
|
||||
|
||||
is_rocm = current_platform.is_rocm()
|
||||
if is_rocm:
|
||||
config = VllmConfig(
|
||||
compilation_config=CompilationConfig(custom_ops=["+rotary_embedding"])
|
||||
)
|
||||
else:
|
||||
config = VllmConfig()
|
||||
try:
|
||||
with monkeypatch.context() as m, set_current_vllm_config(config):
|
||||
if is_rocm:
|
||||
m.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||
m.setenv("VLLM_ROCM_USE_AITER_TRITON_ROPE", "1")
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
yield config
|
||||
finally:
|
||||
if is_rocm:
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16, torch.float])
|
||||
@pytest.mark.parametrize("is_neox_style", [False, True])
|
||||
@pytest.mark.parametrize("seq_len", [11, 42])
|
||||
@pytest.mark.parametrize("qk_rope_head_dim", [64, 128])
|
||||
@pytest.mark.parametrize("num_q_heads", [128])
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
|
||||
@pytest.mark.parametrize("kv_lora_rank", [512])
|
||||
@pytest.mark.parametrize("num_blocks", [64])
|
||||
@pytest.mark.parametrize("block_size", [16, 64, 256])
|
||||
@pytest.mark.parametrize("seed", [0])
|
||||
@pytest.mark.parametrize(
|
||||
"device",
|
||||
[f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)],
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_concat_and_cache_mla_rope_fused(
|
||||
default_vllm_config,
|
||||
dtype: torch.dtype,
|
||||
is_neox_style: bool,
|
||||
seq_len: int,
|
||||
qk_rope_head_dim: int,
|
||||
num_q_heads: int,
|
||||
kv_cache_dtype: str,
|
||||
kv_lora_rank: int,
|
||||
num_blocks: int,
|
||||
block_size: int,
|
||||
seed: int,
|
||||
device: str,
|
||||
max_position: int = 8192,
|
||||
base: float = 10000,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
|
||||
rope = RotaryEmbedding(
|
||||
qk_rope_head_dim,
|
||||
qk_rope_head_dim,
|
||||
max_position,
|
||||
base,
|
||||
is_neox_style,
|
||||
torch.float32,
|
||||
)
|
||||
|
||||
rope = rope.to(dtype=dtype, device=torch.get_default_device())
|
||||
|
||||
positions = torch.randint(0, max_position, (seq_len,))
|
||||
|
||||
query = torch.randn(seq_len, num_q_heads, qk_rope_head_dim, dtype=dtype)
|
||||
key = torch.randn(seq_len, 1, qk_rope_head_dim + kv_lora_rank, dtype=dtype)
|
||||
|
||||
k_pe = torch.flatten(key[..., :qk_rope_head_dim], start_dim=1).to(device=device)
|
||||
kv_c = torch.flatten(key[..., qk_rope_head_dim:], start_dim=1).to(device=device)
|
||||
|
||||
if current_platform.is_rocm():
|
||||
# We use forward_hip for the same numerics as the fused custom kernel on ROCm
|
||||
# when dtype is FP16. The torch-native implementation implicitly upcasts
|
||||
# FP16 x FP16 multiplications to FP32 before downcasting them, which leads
|
||||
# to notable output divergences.
|
||||
# Clone the tensors because the implementation modifies them in-place
|
||||
ref_q_pe, ref_k_pe = rope.forward_hip(positions, query.clone(), k_pe.clone())
|
||||
else:
|
||||
# NOTE(woosuk): The reference implementation should be executed first
|
||||
# because the custom kernel is in-place.
|
||||
ref_q_pe, ref_k_pe = rope.forward_native(positions, query, k_pe)
|
||||
assert ref_k_pe is not None
|
||||
|
||||
ref_k_pe = torch.flatten(ref_k_pe, start_dim=1).to(device=device)
|
||||
ref_k_rope = ref_k_pe[..., :qk_rope_head_dim]
|
||||
|
||||
total_available_slots = num_blocks * block_size
|
||||
total_needed_slots = seq_len
|
||||
assert total_available_slots >= total_needed_slots, "Not enough kv slots!"
|
||||
|
||||
slot_mapping_lst = random.sample(range(total_available_slots), total_needed_slots)
|
||||
slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)
|
||||
|
||||
entry_size = kv_lora_rank + qk_rope_head_dim
|
||||
|
||||
kv_cache_scale = torch.tensor([0.1], dtype=torch.float32, device=device)
|
||||
|
||||
kv_cache = torch.zeros(
|
||||
num_blocks,
|
||||
block_size,
|
||||
entry_size,
|
||||
dtype=torch.uint8 if kv_cache_dtype == "fp8" else dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
ref_temp = torch.zeros(*kv_cache.shape, dtype=dtype, device=device)
|
||||
|
||||
for i in range(seq_len):
|
||||
slot = slot_mapping[i].item()
|
||||
block_idx = slot // block_size
|
||||
block_offset = slot % block_size
|
||||
ref_temp[block_idx, block_offset] = torch.cat((kv_c[i], ref_k_rope[i]), -1)
|
||||
|
||||
if kv_cache_dtype == "fp8":
|
||||
ref_kv_cache = torch.empty_like(ref_temp, dtype=kv_cache.dtype)
|
||||
ops.convert_fp8(
|
||||
ref_kv_cache, ref_temp, kv_cache_scale.item(), kv_dtype=kv_cache_dtype
|
||||
)
|
||||
else:
|
||||
ref_kv_cache = ref_temp
|
||||
|
||||
opcheck(
|
||||
torch.ops._C_cache_ops.concat_and_cache_mla_rope_fused,
|
||||
(
|
||||
positions,
|
||||
query,
|
||||
k_pe,
|
||||
kv_c,
|
||||
rope.cos_sin_cache,
|
||||
is_neox_style,
|
||||
slot_mapping,
|
||||
kv_cache,
|
||||
kv_cache_dtype,
|
||||
kv_cache_scale,
|
||||
),
|
||||
test_utils=DEFAULT_OPCHECK_TEST_UTILS,
|
||||
)
|
||||
|
||||
ops.concat_and_cache_mla_rope_fused(
|
||||
positions,
|
||||
query,
|
||||
k_pe,
|
||||
kv_c,
|
||||
rope.cos_sin_cache,
|
||||
is_neox_style,
|
||||
slot_mapping,
|
||||
kv_cache,
|
||||
kv_cache_dtype,
|
||||
kv_cache_scale,
|
||||
)
|
||||
|
||||
# ROCm neox-style Triton FMA diverges slightly from the fused kernel, so
|
||||
# relax the affected tolerance: rtol for fp8 (one e4m3 ULP ~12.5%) and atol
|
||||
# otherwise (bounded ~6e-4). Other paths use the CUDA defaults.
|
||||
rocm_neox = current_platform.is_rocm() and is_neox_style
|
||||
if kv_cache_dtype == "fp8":
|
||||
result_temp = torch.empty_like(kv_cache, dtype=torch.float16)
|
||||
ops.convert_fp8(
|
||||
result_temp,
|
||||
kv_cache.contiguous(),
|
||||
kv_cache_scale.item(),
|
||||
kv_dtype=kv_cache_dtype,
|
||||
)
|
||||
expected_temp = torch.empty_like(ref_kv_cache, dtype=torch.float16)
|
||||
ops.convert_fp8(
|
||||
expected_temp, ref_kv_cache, kv_cache_scale.item(), kv_dtype=kv_cache_dtype
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
result_temp, expected_temp, atol=0.001, rtol=0.15 if rocm_neox else 0.1
|
||||
)
|
||||
elif rocm_neox:
|
||||
torch.testing.assert_close(kv_cache, ref_kv_cache, atol=1e-3, rtol=1e-3)
|
||||
else:
|
||||
torch.testing.assert_close(kv_cache, ref_kv_cache)
|
||||
|
||||
torch.testing.assert_close(
|
||||
query, ref_q_pe, atol=get_default_atol(query), rtol=get_default_rtol(query)
|
||||
)
|
||||
@@ -0,0 +1,55 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.utils.platform_utils import is_uva_available
|
||||
from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
|
||||
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.")
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_cpu_write(device):
|
||||
torch.set_default_device(device)
|
||||
cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32)
|
||||
cuda_view = get_accelerator_view_from_cpu_tensor(cpu_tensor)
|
||||
assert cuda_view.device.type == "cuda"
|
||||
|
||||
assert cuda_view[0, 0] == 0
|
||||
assert cuda_view[2, 3] == 0
|
||||
assert cuda_view[4, 5] == 0
|
||||
|
||||
cpu_tensor[0, 0] = 1
|
||||
cpu_tensor[2, 3] = 2
|
||||
cpu_tensor[4, 5] = -1
|
||||
|
||||
cuda_view.mul_(2)
|
||||
assert cuda_view[0, 0] == 2
|
||||
assert cuda_view[2, 3] == 4
|
||||
assert cuda_view[4, 5] == -2
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.")
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_gpu_write(device):
|
||||
torch.set_default_device(device)
|
||||
cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32)
|
||||
cuda_view = get_accelerator_view_from_cpu_tensor(cpu_tensor)
|
||||
assert cuda_view.device.type == "cuda"
|
||||
|
||||
assert cuda_view[0, 0] == 0
|
||||
assert cuda_view[2, 3] == 0
|
||||
assert cuda_view[4, 5] == 0
|
||||
|
||||
cuda_view[0, 0] = 1
|
||||
cuda_view[2, 3] = 2
|
||||
cuda_view[4, 5] = -1
|
||||
cuda_view.mul_(2)
|
||||
|
||||
assert cpu_tensor[0, 0] == 2
|
||||
assert cpu_tensor[2, 3] == 4
|
||||
assert cpu_tensor[4, 5] == -2
|
||||
@@ -0,0 +1,120 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Accuracy tests for the fused Triton bilinear position-embedding kernel.
|
||||
|
||||
Compares ``triton_pos_embed_interpolate`` against the pure-PyTorch
|
||||
``pos_embed_interpolate_native`` across a variety of grid shapes and dtypes.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
|
||||
if HAS_TRITON:
|
||||
from vllm.model_executor.models.qwen3_vl import (
|
||||
pos_embed_interpolate_native,
|
||||
triton_pos_embed_interpolate,
|
||||
)
|
||||
|
||||
|
||||
DTYPES = [torch.float32, torch.bfloat16]
|
||||
# Qwen3-VL default
|
||||
NUM_GRID_PER_SIDE = 48
|
||||
SPATIAL_MERGE_SIZE = 2
|
||||
HIDDEN_DIM = 1152
|
||||
|
||||
# 4 square + 4 non-square grids (h, w divisible by spatial_merge_size=2)
|
||||
SQUARE_GRIDS = [(1, 4, 4), (1, 16, 16), (1, 32, 32), (1, 48, 48)]
|
||||
NON_SQUARE_GRIDS = [(1, 8, 16), (1, 14, 20), (1, 32, 48), (1, 60, 80)]
|
||||
ALL_GRIDS = SQUARE_GRIDS + NON_SQUARE_GRIDS
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
@pytest.mark.parametrize("dtype", DTYPES, ids=lambda d: str(d).split(".")[-1])
|
||||
@pytest.mark.parametrize(
|
||||
"grid_thw",
|
||||
ALL_GRIDS,
|
||||
ids=[f"{t}x{h}x{w}" for t, h, w in ALL_GRIDS],
|
||||
)
|
||||
def test_triton_matches_native(
|
||||
grid_thw: tuple[int, int, int],
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
"""Triton kernel output must match the native PyTorch implementation."""
|
||||
t, h, w = grid_thw
|
||||
device = "cuda"
|
||||
|
||||
# Scale to match real Qwen3-VL pos_embed weight distribution (std~0.23).
|
||||
torch.manual_seed(42)
|
||||
embed_weight = (
|
||||
torch.randn(
|
||||
NUM_GRID_PER_SIDE * NUM_GRID_PER_SIDE,
|
||||
HIDDEN_DIM,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
* 0.25
|
||||
)
|
||||
|
||||
native_out = pos_embed_interpolate_native(
|
||||
embed_weight, t, h, w, NUM_GRID_PER_SIDE, SPATIAL_MERGE_SIZE, dtype
|
||||
)
|
||||
triton_out = triton_pos_embed_interpolate(
|
||||
embed_weight, t, h, w, NUM_GRID_PER_SIDE, SPATIAL_MERGE_SIZE, dtype
|
||||
)
|
||||
|
||||
assert native_out.shape == triton_out.shape, (
|
||||
f"Shape mismatch: native {native_out.shape} vs triton {triton_out.shape}"
|
||||
)
|
||||
|
||||
# Small numerical differences arise from the precomputed h/w_scale
|
||||
# in the triton kernel vs torch.linspace in the native path, which can
|
||||
# cause single-ULP output differences
|
||||
# in a handful of elements.
|
||||
atol = {torch.float32: 5e-5, torch.bfloat16: 1e-2}[dtype]
|
||||
rtol = {torch.float32: 1e-5, torch.bfloat16: 1e-2}[dtype]
|
||||
torch.testing.assert_close(triton_out, native_out, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
@pytest.mark.parametrize("dtype", DTYPES, ids=lambda d: str(d).split(".")[-1])
|
||||
def test_temporal_repeat(dtype: torch.dtype) -> None:
|
||||
"""Verify temporal dimension t > 1 correctly repeats the spatial pattern."""
|
||||
device = "cuda"
|
||||
h, w = 16, 16
|
||||
t_single, t_multi = 1, 3
|
||||
|
||||
# Scale to match real Qwen3-VL pos_embed weight distribution (std~0.23).
|
||||
torch.manual_seed(42)
|
||||
embed_weight = (
|
||||
torch.randn(
|
||||
NUM_GRID_PER_SIDE * NUM_GRID_PER_SIDE,
|
||||
HIDDEN_DIM,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
* 0.25
|
||||
)
|
||||
|
||||
out_single = triton_pos_embed_interpolate(
|
||||
embed_weight,
|
||||
t_single,
|
||||
h,
|
||||
w,
|
||||
NUM_GRID_PER_SIDE,
|
||||
SPATIAL_MERGE_SIZE,
|
||||
dtype,
|
||||
)
|
||||
out_multi = triton_pos_embed_interpolate(
|
||||
embed_weight,
|
||||
t_multi,
|
||||
h,
|
||||
w,
|
||||
NUM_GRID_PER_SIDE,
|
||||
SPATIAL_MERGE_SIZE,
|
||||
dtype,
|
||||
)
|
||||
|
||||
expected = out_single.repeat(t_multi, 1)
|
||||
torch.testing.assert_close(out_multi, expected, atol=0, rtol=0)
|
||||
@@ -0,0 +1,279 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the full FP8 ViT attention path (quantize -> cuDNN -> un-pad)."""
|
||||
|
||||
import contextlib
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
from vllm.utils.flashinfer import (
|
||||
is_flashinfer_cudnn_fp8_prefill_attn_supported,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
|
||||
def _has_flashinfer_cudnn() -> bool:
|
||||
"""Check if FlashInfer cuDNN backend is available."""
|
||||
try:
|
||||
from flashinfer.prefill import (
|
||||
cudnn_batch_prefill_with_kv_cache, # noqa: F401
|
||||
)
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
HEAD_DIMS = [72, 80]
|
||||
SEQ_LENS = [256]
|
||||
NUM_HEADS = [16]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _fp8_attention():
|
||||
"""Create FP8-enabled MMEncoderAttention via config."""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
pytest.skip("FlashInfer cuDNN FP8 prefill attention not supported")
|
||||
|
||||
mm_config = MultiModalConfig(mm_encoder_attn_dtype="fp8")
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
|
||||
# MMEncoderAttention reads torch.get_default_dtype() during init
|
||||
# to determine the output dtype. In real model loading this is bf16.
|
||||
old_dtype = torch.get_default_dtype()
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
),
|
||||
):
|
||||
yield
|
||||
|
||||
torch.set_default_dtype(old_dtype)
|
||||
|
||||
|
||||
def _build_cu_seqlens_and_meta(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
fp8_padded_hidden_size: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Build cu_seqlens, max_seqlen, sequence_lengths for a single sequence."""
|
||||
import numpy as np
|
||||
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
|
||||
cu_seqlens_np = np.array([0, seq_len], dtype=np.int32)
|
||||
|
||||
sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_seqlens_np,
|
||||
torch.device("cuda"),
|
||||
)
|
||||
|
||||
max_seqlen = torch.tensor(
|
||||
MMEncoderAttention.compute_max_seqlen(
|
||||
AttentionBackendEnum.FLASHINFER, cu_seqlens_np
|
||||
),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_seqlens_np,
|
||||
num_heads * head_dim,
|
||||
1, # tp_size
|
||||
torch.device("cuda"),
|
||||
fp8_padded_hidden_size=fp8_padded_hidden_size,
|
||||
)
|
||||
|
||||
return cu_seqlens, max_seqlen, sequence_lengths
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not (HAS_TRITON and _has_flashinfer_cudnn()),
|
||||
reason="Triton and FlashInfer cuDNN required",
|
||||
)
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
def test_fp8_attn_output_shape(
|
||||
head_dim: int,
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
_fp8_attention,
|
||||
) -> None:
|
||||
"""Verify FP8 attention produces correct output shape after un-padding."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
|
||||
attn = None
|
||||
with contextlib.suppress(ValueError, ImportError):
|
||||
attn = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 MMEncoderAttention not available")
|
||||
assert attn is not None # mypy narrowing
|
||||
|
||||
# FP8 always needs fp8_padded_hidden_size for correct cu_seqlens
|
||||
fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
|
||||
|
||||
cu_seqlens, max_seqlen, sequence_lengths = _build_cu_seqlens_and_meta(
|
||||
seq_len, num_heads, head_dim, fp8_padded_hidden_size=fp8_padded_hidden_size
|
||||
)
|
||||
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
output = attn._forward_flashinfer(q, k, v, cu_seqlens, max_seqlen, sequence_lengths)
|
||||
|
||||
# Output should have original head_dim (un-padded)
|
||||
assert output.shape[-1] == head_dim
|
||||
assert output.dtype == torch.bfloat16
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not (HAS_TRITON and _has_flashinfer_cudnn()),
|
||||
reason="Triton and FlashInfer cuDNN required",
|
||||
)
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
def test_fp8_vs_bf16_close(
|
||||
head_dim: int, seq_len: int, num_heads: int, _fp8_attention
|
||||
) -> None:
|
||||
"""FP8 attention output should be reasonably close to BF16 baseline."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
|
||||
torch.manual_seed(42)
|
||||
q = torch.randn(
|
||||
1,
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
# FP8 path
|
||||
attn_fp8 = None
|
||||
with contextlib.suppress(ValueError, ImportError):
|
||||
attn_fp8 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
if attn_fp8 is None or not attn_fp8.fp8_enabled:
|
||||
pytest.skip("FP8 MMEncoderAttention not available")
|
||||
assert attn_fp8 is not None # mypy narrowing
|
||||
|
||||
fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
|
||||
cu_seqlens, max_seqlen, seq_lengths = _build_cu_seqlens_and_meta(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
fp8_padded_hidden_size=fp8_padded_hidden_size,
|
||||
)
|
||||
|
||||
out_fp8 = attn_fp8._forward_flashinfer(
|
||||
q.clone(),
|
||||
k.clone(),
|
||||
v.clone(),
|
||||
cu_seqlens,
|
||||
max_seqlen,
|
||||
seq_lengths,
|
||||
)
|
||||
|
||||
# BF16 baseline (create non-FP8 attention by using scale=attn_fp8.scale
|
||||
# and calling the wrapper directly without FP8 quantization)
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
_get_flashinfer_workspace_buffer,
|
||||
)
|
||||
from vllm.v1.attention.ops.vit_attn_wrappers import (
|
||||
vit_flashinfer_wrapper,
|
||||
)
|
||||
|
||||
out_bf16 = vit_flashinfer_wrapper(
|
||||
q=q.clone(),
|
||||
k=k.clone(),
|
||||
v=v.clone(),
|
||||
scale=attn_fp8.scale,
|
||||
workspace_buffer=_get_flashinfer_workspace_buffer(),
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
sequence_lengths=seq_lengths,
|
||||
)
|
||||
|
||||
out_fp8_f = out_fp8.float()
|
||||
out_bf16_f = out_bf16.float()
|
||||
|
||||
abs_diff = (out_fp8_f - out_bf16_f).abs()
|
||||
abs_diff_flat = abs_diff.flatten()
|
||||
|
||||
# Relative diff (avoid division by zero)
|
||||
denom = out_bf16_f.abs().clamp(min=1e-6)
|
||||
rel_diff_flat = (abs_diff / denom).flatten()
|
||||
|
||||
cosine_sim = torch.nn.functional.cosine_similarity(
|
||||
out_fp8_f.flatten().unsqueeze(0),
|
||||
out_bf16_f.flatten().unsqueeze(0),
|
||||
).item()
|
||||
|
||||
pcts = [50, 90, 95, 99, 99.9]
|
||||
abs_pct = {p: torch.quantile(abs_diff_flat, p / 100).item() for p in pcts}
|
||||
rel_pct = {p: torch.quantile(rel_diff_flat, p / 100).item() for p in pcts}
|
||||
|
||||
print(f"\nFP8 vs BF16 (head_dim={head_dim}, seq_len={seq_len}):")
|
||||
print(f" cosine_sim={cosine_sim:.6f}")
|
||||
print(
|
||||
f" abs_diff: max={abs_diff_flat.max().item():.6f}, "
|
||||
f"mean={abs_diff_flat.mean().item():.6f}, "
|
||||
+ ", ".join(f"p{p}={abs_pct[p]:.6f}" for p in pcts)
|
||||
)
|
||||
print(
|
||||
f" rel_diff: max={rel_diff_flat.max().item():.6f}, "
|
||||
f"mean={rel_diff_flat.mean().item():.6f}, "
|
||||
+ ", ".join(f"p{p}={rel_pct[p]:.6f}" for p in pcts)
|
||||
)
|
||||
|
||||
assert abs_diff_flat.max().item() < 0.3, (
|
||||
f"FP8 vs BF16 max abs diff too large: {abs_diff_flat.max().item()}"
|
||||
)
|
||||
assert abs_diff_flat.mean().item() < 0.03, (
|
||||
f"FP8 vs BF16 mean abs diff too large: {abs_diff_flat.mean().item()}"
|
||||
)
|
||||
assert cosine_sim > 0.99, f"Cosine similarity too low: {cosine_sim:.6f}"
|
||||
@@ -0,0 +1,126 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the stride-aware FP8 quantization kernel with head_dim padding."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
get_fp8_min_max,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
|
||||
if HAS_TRITON:
|
||||
from vllm.kernels.triton.qkv_padded_fp8_quant import (
|
||||
quantize_fp8_pad_head_dim_triton,
|
||||
)
|
||||
|
||||
HEAD_DIMS = [72, 80, 128]
|
||||
SEQ_LENS = [64, 256]
|
||||
NUM_HEADS = [16]
|
||||
SCALES = [0.01, 0.1, 1.0]
|
||||
|
||||
|
||||
def _naive_fp8_quantize(
|
||||
tensor: torch.Tensor, scale: torch.Tensor, skip_scale: bool
|
||||
) -> torch.Tensor:
|
||||
"""Reference FP8 quantization in PyTorch."""
|
||||
fp8_dtype = current_platform.fp8_dtype()
|
||||
fp8_min, fp8_max = get_fp8_min_max()
|
||||
|
||||
x = tensor.float()
|
||||
if not skip_scale:
|
||||
x = x / scale.item()
|
||||
x = x.clamp(fp8_min, fp8_max)
|
||||
return x.to(fp8_dtype)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("scale_val", SCALES)
|
||||
def test_quantize_contiguous(
|
||||
head_dim: int, seq_len: int, num_heads: int, scale_val: float
|
||||
) -> None:
|
||||
"""Test quantization of contiguous 3D tensors."""
|
||||
torch.manual_seed(42)
|
||||
tensor = torch.randn(
|
||||
seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda").view(
|
||||
1, 1, 1, 1
|
||||
)
|
||||
|
||||
result = quantize_fp8_pad_head_dim_triton(tensor, scale)
|
||||
|
||||
padded_dim = (head_dim + 15) // 16 * 16
|
||||
assert result.shape == (seq_len, num_heads, padded_dim)
|
||||
assert result.is_contiguous()
|
||||
assert result.dtype == current_platform.fp8_dtype()
|
||||
|
||||
# Compare unpadded portion against reference
|
||||
ref = _naive_fp8_quantize(tensor, scale, skip_scale=False)
|
||||
torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
|
||||
|
||||
# Padded region should be zero
|
||||
if padded_dim > head_dim:
|
||||
assert (result[:, :, head_dim:].float() == 0).all()
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
@pytest.mark.parametrize("head_dim", [72, 80])
|
||||
def test_quantize_non_contiguous(head_dim: int) -> None:
|
||||
"""Test quantization from non-contiguous QKV views (interleaved buffer)."""
|
||||
seq_len, num_heads = 64, 16
|
||||
# Simulate interleaved QKV buffer: shape (seq_len, 3 * num_heads, head_dim)
|
||||
qkv = torch.randn(
|
||||
seq_len, 3 * num_heads, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
# Q is every 3rd head slice - non-contiguous view
|
||||
q = qkv[:, 0::3, :]
|
||||
assert not q.is_contiguous()
|
||||
|
||||
scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
|
||||
result = quantize_fp8_pad_head_dim_triton(q, scale)
|
||||
|
||||
padded_dim = (head_dim + 15) // 16 * 16
|
||||
assert result.shape == (seq_len, num_heads, padded_dim)
|
||||
assert result.is_contiguous()
|
||||
|
||||
# Compare against contiguous reference
|
||||
ref = _naive_fp8_quantize(q.contiguous(), scale, skip_scale=False)
|
||||
torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
def test_skip_scale() -> None:
|
||||
"""Test skip_scale=True produces cast-only output (no division)."""
|
||||
seq_len, num_heads, head_dim = 32, 8, 80
|
||||
tensor = torch.randn(
|
||||
seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
|
||||
|
||||
result_skip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=True)
|
||||
result_noskip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=False)
|
||||
|
||||
# skip_scale should just cast, not divide
|
||||
ref_cast = _naive_fp8_quantize(tensor, scale, skip_scale=True)
|
||||
torch.testing.assert_close(result_skip[:, :, :head_dim].float(), ref_cast.float())
|
||||
|
||||
# With scale != 1.0, skip and no-skip should differ
|
||||
assert not torch.equal(result_skip.float(), result_noskip.float())
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
def test_4d_input() -> None:
|
||||
"""Test that 4D input (B, S, H, D) is handled correctly."""
|
||||
B, S, H, D = 2, 32, 8, 72
|
||||
tensor = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
|
||||
scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
|
||||
|
||||
result = quantize_fp8_pad_head_dim_triton(tensor, scale)
|
||||
padded_dim = (D + 15) // 16 * 16
|
||||
assert result.shape == (B, S, H, padded_dim)
|
||||
@@ -0,0 +1,251 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for FP8 scaling (dynamic and static) in MMEncoderAttention."""
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
_FP8_AMAX_HISTORY_LEN,
|
||||
_FP8_MAX,
|
||||
)
|
||||
from vllm.utils.flashinfer import (
|
||||
is_flashinfer_cudnn_fp8_prefill_attn_supported,
|
||||
)
|
||||
|
||||
LAYER_0 = "visual.blocks.0.attn.attn"
|
||||
LAYER_1 = "visual.blocks.1.attn.attn"
|
||||
NUM_HEADS = 16
|
||||
HEAD_DIM = 72
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _build_attention(mm_config):
|
||||
"""Yield an MMEncoderAttention with the given multimodal config.
|
||||
|
||||
The VllmConfig context stays active while the test runs so that
|
||||
``get_multimodal_config()`` calls during the forward path resolve. Also
|
||||
invokes ``process_weights_after_loading`` to simulate the model loader's
|
||||
auto-scan. Yields ``None`` if FlashInfer cuDNN is not available.
|
||||
"""
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
yield None
|
||||
return
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
),
|
||||
):
|
||||
attn = MMEncoderAttention(
|
||||
num_heads=NUM_HEADS,
|
||||
head_size=HEAD_DIM,
|
||||
prefix=LAYER_0,
|
||||
)
|
||||
attn.process_weights_after_loading(torch.bfloat16)
|
||||
yield attn
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _make_attention():
|
||||
"""Create an MMEncoderAttention with dynamic FP8 scaling."""
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
with _build_attention(MultiModalConfig(mm_encoder_attn_dtype="fp8")) as attn:
|
||||
yield attn
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _make_static_attention(tmp_path):
|
||||
"""Create an MMEncoderAttention with static FP8 scales from a file."""
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
scale_file = tmp_path / "scales.json"
|
||||
scale_file.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
LAYER_0: {"q": 224.0, "k": 198.0, "v": 210.0},
|
||||
LAYER_1: {"q": 100.0, "k": 110.0, "v": 120.0},
|
||||
}
|
||||
)
|
||||
)
|
||||
with _build_attention(
|
||||
MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_path=str(scale_file),
|
||||
)
|
||||
) as attn:
|
||||
yield attn
|
||||
|
||||
|
||||
def test_dynamic_scaling_updates_scales(_make_attention) -> None:
|
||||
"""Verify that _record_amax_and_update_scales updates scale buffers."""
|
||||
attn = _make_attention
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available (FlashInfer backend required)")
|
||||
|
||||
attn = attn.to("cuda")
|
||||
|
||||
S, H, D = 32, NUM_HEADS, HEAD_DIM
|
||||
q = torch.full((S, H, D), 2.0, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.full((S, H, D), 3.0, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.full((S, H, D), 4.0, device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
attn._record_amax_and_update_scales(q, k, v)
|
||||
|
||||
expected_q_scale = 2.0 / _FP8_MAX
|
||||
expected_k_scale = 3.0 / _FP8_MAX
|
||||
expected_v_scale = 4.0 / _FP8_MAX
|
||||
|
||||
torch.testing.assert_close(attn._fp8_q_scale.item(), expected_q_scale)
|
||||
torch.testing.assert_close(attn._fp8_k_scale.item(), expected_k_scale)
|
||||
torch.testing.assert_close(attn._fp8_v_scale.item(), expected_v_scale)
|
||||
|
||||
|
||||
def test_circular_buffer_wraps(_make_attention) -> None:
|
||||
"""Verify the amax circular buffer wraps at HISTORY_LEN."""
|
||||
attn = _make_attention
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available (FlashInfer backend required)")
|
||||
|
||||
attn = attn.to("cuda")
|
||||
S, H, D = 16, NUM_HEADS, HEAD_DIM
|
||||
|
||||
for i in range(_FP8_AMAX_HISTORY_LEN + 2):
|
||||
mag = float(i + 1)
|
||||
q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
attn._record_amax_and_update_scales(q, k, v)
|
||||
|
||||
assert attn._fp8_amax_pos == 2
|
||||
|
||||
expected_max = float(_FP8_AMAX_HISTORY_LEN + 2)
|
||||
expected_scale = expected_max / _FP8_MAX
|
||||
torch.testing.assert_close(attn._fp8_q_scale.item(), expected_scale)
|
||||
|
||||
|
||||
def test_static_scales_loaded(_make_static_attention) -> None:
|
||||
"""Verify static scales are loaded from the JSON file."""
|
||||
attn = _make_static_attention
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available (FlashInfer backend required)")
|
||||
|
||||
assert attn.fp8_enabled
|
||||
assert not attn._fp8_dynamic_scale
|
||||
|
||||
# Layer 0 scales (the layer this attention was created with).
|
||||
assert attn._fp8_q_scale.item() == 224.0
|
||||
assert attn._fp8_k_scale.item() == 198.0
|
||||
assert attn._fp8_v_scale.item() == 210.0
|
||||
|
||||
assert not attn.skip_scale_q
|
||||
assert not attn.skip_scale_k
|
||||
assert not attn.skip_scale_v
|
||||
|
||||
# No amax history buffers for static scaling.
|
||||
assert not hasattr(attn, "_fp8_q_amax")
|
||||
|
||||
|
||||
def test_static_scales_missing_layer(tmp_path) -> None:
|
||||
"""Verify error when requested layer is not in the scale file."""
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
pytest.skip("FlashInfer cuDNN not available")
|
||||
|
||||
scale_file = tmp_path / "wrong_layer.json"
|
||||
scale_file.write_text(
|
||||
json.dumps({"visual.blocks.99.attn": {"q": 1.0, "k": 1.0, "v": 1.0}})
|
||||
)
|
||||
mm_config = MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_path=str(scale_file),
|
||||
)
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
),
|
||||
):
|
||||
attn = MMEncoderAttention(
|
||||
num_heads=NUM_HEADS,
|
||||
head_size=HEAD_DIM,
|
||||
prefix=LAYER_0,
|
||||
)
|
||||
with pytest.raises(ValueError, match="scales not found for layer"):
|
||||
attn.process_weights_after_loading(torch.bfloat16)
|
||||
|
||||
|
||||
def test_dynamic_scales_auto_save(tmp_path) -> None:
|
||||
"""Verify scales are saved to disk after the amax buffer fills."""
|
||||
import vllm.model_executor.layers.attention.mm_encoder_attention as _mod
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
pytest.skip("FlashInfer cuDNN not available")
|
||||
|
||||
# Reset module-level state between runs (other tests may have left
|
||||
# state behind after triggering a save).
|
||||
_mod._fp8_scale_save_path = None
|
||||
_mod._fp8_saved_scale_refs.clear()
|
||||
|
||||
save_file = tmp_path / "auto_scales.json"
|
||||
with _build_attention(
|
||||
MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_save_path=str(save_file),
|
||||
)
|
||||
) as attn:
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available")
|
||||
|
||||
attn = attn.to("cuda")
|
||||
S, H, D = 16, NUM_HEADS, HEAD_DIM
|
||||
|
||||
# Run exactly _FP8_AMAX_HISTORY_LEN forward passes.
|
||||
for i in range(_FP8_AMAX_HISTORY_LEN):
|
||||
mag = float(i + 1)
|
||||
q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.full((S, H, D), mag * 0.5, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.full((S, H, D), mag * 0.3, device="cuda", dtype=torch.bfloat16)
|
||||
attn._record_amax_and_update_scales(q, k, v)
|
||||
|
||||
# File should have been written on the 16th call (buffer wrap).
|
||||
assert save_file.is_file(), "Scale file was not saved"
|
||||
scales = json.loads(save_file.read_text())
|
||||
assert LAYER_0 in scales
|
||||
assert set(scales[LAYER_0].keys()) == {"q", "k", "v"}
|
||||
for val in scales[LAYER_0].values():
|
||||
assert isinstance(val, float) and val > 0
|
||||
|
||||
# Path is cleared after the one-shot save fires.
|
||||
assert _mod._fp8_scale_save_path is None
|
||||
@@ -0,0 +1,75 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from collections.abc import Callable
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import helion
|
||||
|
||||
from vllm.kernels.helion.case_key import CaseKey
|
||||
from vllm.kernels.helion.config_manager import ConfigManager
|
||||
from vllm.kernels.helion.register import register_kernel
|
||||
from vllm.kernels.helion.utils import get_canonical_gpu_name
|
||||
|
||||
GPU_PLATFORM = get_canonical_gpu_name()
|
||||
|
||||
DEFAULT_CONFIGS: dict[CaseKey, helion.Config] = {
|
||||
CaseKey.default(): helion.Config(block_sizes=[32]),
|
||||
}
|
||||
|
||||
|
||||
@contextmanager
|
||||
def dummy_kernel_registry(
|
||||
configs: dict[CaseKey, helion.Config] | None = None,
|
||||
):
|
||||
"""Context manager providing a register function with automatic config setup.
|
||||
|
||||
Yields a ``register`` callable with the same signature as
|
||||
``register_kernel``. Before applying the real decorator it writes a
|
||||
config JSON for the kernel name (from ``op_name`` or ``fn.__name__``)
|
||||
into a temporary directory backed by a fresh ``ConfigManager``.
|
||||
"""
|
||||
if configs is None:
|
||||
configs = DEFAULT_CONFIGS
|
||||
|
||||
def _to_config_entries(cfgs: dict) -> list[dict[str, Any]]:
|
||||
pairs: list[dict[str, Any]] = []
|
||||
for k, v in cfgs.items():
|
||||
config_data = v.__dict__["config"]
|
||||
pairs.append({"key": dict(k), "config": config_data})
|
||||
return pairs
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
config_dir = Path(tmpdir)
|
||||
ConfigManager.reset_instance()
|
||||
cm = ConfigManager(base_dir=config_dir)
|
||||
|
||||
with patch(
|
||||
"vllm.kernels.helion.config_manager.ConfigManager",
|
||||
return_value=cm,
|
||||
):
|
||||
|
||||
def register(
|
||||
op_name: str | None = None,
|
||||
**kwargs,
|
||||
) -> Callable:
|
||||
def decorator(fn: Callable) -> Callable:
|
||||
name = op_name or fn.__name__
|
||||
kernel_dir = config_dir / name
|
||||
kernel_dir.mkdir(parents=True, exist_ok=True)
|
||||
(kernel_dir / f"{GPU_PLATFORM}.json").write_text(
|
||||
json.dumps(_to_config_entries(configs))
|
||||
)
|
||||
return register_kernel(op_name, **kwargs)(fn)
|
||||
|
||||
return decorator
|
||||
|
||||
try:
|
||||
yield register
|
||||
finally:
|
||||
ConfigManager.reset_instance()
|
||||
@@ -0,0 +1,91 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for autotuning Helion kernels, including disabled kernels with no configs."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.utils.import_utils import has_helion
|
||||
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
import helion
|
||||
import helion.language as hl
|
||||
from helion.autotuner.base_search import BaseSearch
|
||||
|
||||
from tests.kernels.helion.helpers import dummy_kernel_registry
|
||||
from vllm.kernels.helion.register import create_helion_decorated_kernel
|
||||
|
||||
|
||||
def _add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
out = torch.empty_like(x)
|
||||
for tile in hl.tile(x.size()):
|
||||
out[tile] = x[tile] + y[tile]
|
||||
return out
|
||||
|
||||
|
||||
class NoCompileSearch(BaseSearch):
|
||||
"""Autotuner that returns the default config without GPU compilation.
|
||||
|
||||
Modeled after helion's test BasicSearch (pytorch/helion#1649).
|
||||
"""
|
||||
|
||||
def autotune(self, *, skip_cache: bool = False):
|
||||
return self.config_spec.default_config()
|
||||
|
||||
|
||||
def _no_compile_autotuner_fn(bound_kernel, args, **kwargs):
|
||||
return NoCompileSearch(bound_kernel, args, **kwargs)
|
||||
|
||||
|
||||
class TestAutotuneDisabledKernel:
|
||||
"""Test autotuning flow on disabled kernels (no platform configs)."""
|
||||
|
||||
def setup_method(self):
|
||||
from vllm.kernels.helion.register import _REGISTERED_KERNELS
|
||||
|
||||
self._saved_registry = dict(_REGISTERED_KERNELS)
|
||||
_REGISTERED_KERNELS.clear()
|
||||
|
||||
def teardown_method(self):
|
||||
from vllm.kernels.helion.register import _REGISTERED_KERNELS
|
||||
|
||||
_REGISTERED_KERNELS.clear()
|
||||
_REGISTERED_KERNELS.update(self._saved_registry)
|
||||
|
||||
def test_autotune_disabled_kernel_produces_valid_config(self):
|
||||
"""Register a kernel with no configs (disabled), run autotune,
|
||||
verify it produces a valid helion.Config."""
|
||||
with dummy_kernel_registry(configs={}) as register:
|
||||
wrapper = register(
|
||||
"autotune_test_kernel",
|
||||
config_picker=lambda args, keys: None,
|
||||
fake_impl=lambda *a, **kw: None,
|
||||
input_generator=lambda: {
|
||||
"small": (
|
||||
torch.randn(4, 4, device="cuda"),
|
||||
torch.randn(4, 4, device="cuda"),
|
||||
),
|
||||
},
|
||||
)(_add_kernel)
|
||||
|
||||
assert wrapper._disabled is True
|
||||
|
||||
inputs = wrapper.get_inputs()
|
||||
assert "small" in inputs
|
||||
|
||||
settings = helion.Settings()
|
||||
settings.autotuner_fn = _no_compile_autotuner_fn
|
||||
wrapper.helion_settings = settings
|
||||
|
||||
config = wrapper.run_autotune(inputs["small"])
|
||||
expected_default = (
|
||||
create_helion_decorated_kernel(_add_kernel, helion_settings=settings)
|
||||
.bind(inputs["small"])
|
||||
.config_spec.default_config()
|
||||
)
|
||||
assert config == expected_default
|
||||
@@ -0,0 +1,68 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.utils.import_utils import has_helion
|
||||
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from vllm.kernels.helion.case_key import CaseKey
|
||||
|
||||
|
||||
class TestCaseKey:
|
||||
"""Test suite for CaseKey class."""
|
||||
|
||||
def test_construction_with_dict(self):
|
||||
key = CaseKey({"intermediate": 2048, "numtokens": 256})
|
||||
assert key["intermediate"] == 2048
|
||||
assert key["numtokens"] == 256
|
||||
|
||||
def test_empty_construction_raises(self):
|
||||
with pytest.raises(TypeError, match="at least one key-value pair"):
|
||||
CaseKey()
|
||||
with pytest.raises(TypeError, match="at least one key-value pair"):
|
||||
CaseKey({})
|
||||
|
||||
def test_default_construction(self):
|
||||
key = CaseKey.default()
|
||||
assert len(key) == 0
|
||||
assert key.is_default()
|
||||
|
||||
def test_non_default_is_not_default(self):
|
||||
key = CaseKey({"intermediate": 2048})
|
||||
assert not key.is_default()
|
||||
|
||||
def test_hashable_and_equality(self):
|
||||
a = CaseKey({"intermediate": 2048, "numtokens": 256})
|
||||
b = CaseKey({"numtokens": 256, "intermediate": 2048})
|
||||
assert a == b
|
||||
assert hash(a) == hash(b)
|
||||
assert a != CaseKey({"intermediate": 4096})
|
||||
assert CaseKey.default() == CaseKey.default()
|
||||
|
||||
configs = {
|
||||
CaseKey.default(): "default_config",
|
||||
a: "a_config",
|
||||
}
|
||||
assert configs[b] == "a_config"
|
||||
assert configs[CaseKey.default()] == "default_config"
|
||||
|
||||
def test_str_is_sorted_json(self):
|
||||
assert str(CaseKey({"z": 1, "a": 2})) == '{"a":2,"z":1}'
|
||||
assert str(CaseKey.default()) == "{}"
|
||||
|
||||
def test_immutable(self):
|
||||
key = CaseKey({"intermediate": 2048})
|
||||
with pytest.raises(TypeError, match="immutable"):
|
||||
key["intermediate"] = 4096
|
||||
with pytest.raises(TypeError, match="immutable"):
|
||||
del key["intermediate"]
|
||||
with pytest.raises(TypeError, match="immutable"):
|
||||
key.update({"numtokens": 256})
|
||||
with pytest.raises(TypeError, match="immutable"):
|
||||
key.clear()
|
||||
@@ -0,0 +1,409 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Unit tests for Helion ConfigManager and ConfigSet.
|
||||
|
||||
Tests the simplified configuration management system for Helion custom kernels.
|
||||
"""
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.utils.import_utils import has_helion
|
||||
|
||||
# Skip entire module if helion is not available
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
import helion
|
||||
|
||||
from vllm.kernels.helion.case_key import CaseKey
|
||||
from vllm.kernels.helion.config_manager import (
|
||||
ConfigManager,
|
||||
ConfigSet,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_config_manager_singleton():
|
||||
"""Reset ConfigManager singleton before each test."""
|
||||
ConfigManager.reset_instance()
|
||||
yield
|
||||
ConfigManager.reset_instance()
|
||||
|
||||
|
||||
class TestConfigSet:
|
||||
"""Test suite for ConfigSet class."""
|
||||
|
||||
def test_config_set_creation(self):
|
||||
"""Test creating an empty ConfigSet."""
|
||||
config_set = ConfigSet("test_kernel")
|
||||
|
||||
assert config_set.kernel_name == "test_kernel"
|
||||
assert config_set.get_platforms() == []
|
||||
|
||||
def test_config_set_from_dict(self):
|
||||
"""Test creating ConfigSet from dictionary data."""
|
||||
config_data = {
|
||||
"block_sizes": [32, 16],
|
||||
"num_warps": 4,
|
||||
"num_stages": 3,
|
||||
"pid_type": "persistent_interleaved",
|
||||
}
|
||||
data = {
|
||||
"h100": [
|
||||
{"key": {"batch": 32, "hidden": 4096}, "config": config_data},
|
||||
]
|
||||
}
|
||||
|
||||
config_set = ConfigSet.from_dict("test_kernel", data)
|
||||
|
||||
assert config_set.kernel_name == "test_kernel"
|
||||
assert config_set.get_platforms() == ["h100"]
|
||||
|
||||
internal_key = CaseKey({"batch": 32, "hidden": 4096})
|
||||
config = config_set.get_config("h100", internal_key)
|
||||
assert isinstance(config, helion.Config)
|
||||
assert config.block_sizes == [32, 16]
|
||||
assert config.num_warps == 4
|
||||
assert config.num_stages == 3
|
||||
assert config.pid_type == "persistent_interleaved"
|
||||
|
||||
def test_config_set_get_config_keyerror(self):
|
||||
"""Test that accessing non-existent configs raises informative KeyErrors."""
|
||||
config_set = ConfigSet("test_kernel")
|
||||
|
||||
with pytest.raises(KeyError, match="platform 'h100' not found"):
|
||||
config_set.get_config("h100", "nonexistent")
|
||||
|
||||
config_data = {"num_warps": 8, "num_stages": 4}
|
||||
data = {
|
||||
"h100": [
|
||||
{"key": {"batch": 64, "hidden": 2048}, "config": config_data},
|
||||
]
|
||||
}
|
||||
config_set = ConfigSet.from_dict("test_kernel", data)
|
||||
|
||||
nonexistent_key = CaseKey({"batch": 32, "hidden": 4096})
|
||||
with pytest.raises(KeyError, match="config_key .* not found"):
|
||||
config_set.get_config("h100", nonexistent_key)
|
||||
|
||||
def test_config_set_get_platforms(self):
|
||||
"""Test get_platforms method."""
|
||||
# Use realistic config data
|
||||
config1 = {"num_warps": 4, "num_stages": 3}
|
||||
config2 = {"num_warps": 8, "num_stages": 5}
|
||||
|
||||
data = {
|
||||
"h100": [
|
||||
{"key": {"batch": 32, "hidden": 4096}, "config": config1},
|
||||
],
|
||||
"a100": [
|
||||
{"key": {"batch": 16, "hidden": 2048}, "config": config2},
|
||||
],
|
||||
}
|
||||
config_set = ConfigSet.from_dict("test_kernel", data)
|
||||
|
||||
platforms = config_set.get_platforms()
|
||||
assert platforms == ["a100", "h100"] # Should be sorted
|
||||
|
||||
def test_config_set_get_config_keys(self):
|
||||
"""Test get_config_keys method."""
|
||||
config1 = {"num_warps": 4, "num_stages": 3}
|
||||
config2 = {"num_warps": 8, "num_stages": 5}
|
||||
|
||||
data = {
|
||||
"h100": [
|
||||
{"key": {"batch": 32, "hidden": 4096}, "config": config1},
|
||||
{"key": {"batch": 64, "hidden": 2048}, "config": config2},
|
||||
]
|
||||
}
|
||||
config_set = ConfigSet.from_dict("test_kernel", data)
|
||||
|
||||
config_keys = config_set.get_config_keys("h100")
|
||||
expected_keys = sorted(
|
||||
[
|
||||
CaseKey({"batch": 32, "hidden": 4096}),
|
||||
CaseKey({"batch": 64, "hidden": 2048}),
|
||||
],
|
||||
key=lambda k: str(k) if k is not None else "",
|
||||
)
|
||||
assert config_keys == expected_keys
|
||||
|
||||
assert config_set.get_config_keys("v100") == []
|
||||
|
||||
def test_config_set_to_dict(self):
|
||||
"""Test converting ConfigSet to dictionary."""
|
||||
original_config = {
|
||||
"block_sizes": [64, 32],
|
||||
"num_warps": 16,
|
||||
"num_stages": 4,
|
||||
"pid_type": "persistent_blocked",
|
||||
}
|
||||
original_data = {
|
||||
"h100": [
|
||||
{"key": {"batch": 32, "hidden": 4096}, "config": original_config},
|
||||
]
|
||||
}
|
||||
|
||||
config_set = ConfigSet.from_dict("test_kernel", original_data)
|
||||
result_data = config_set.to_dict()
|
||||
|
||||
internal_key = CaseKey({"batch": 32, "hidden": 4096})
|
||||
assert internal_key in result_data["h100"]
|
||||
assert result_data["h100"][internal_key] == original_config
|
||||
|
||||
|
||||
class TestConfigManager:
|
||||
"""Test suite for ConfigManager class."""
|
||||
|
||||
def test_config_manager_creation_default_base_dir(self):
|
||||
"""Test creating ConfigManager with default base directory."""
|
||||
manager = ConfigManager()
|
||||
assert manager._base_dir.name == "configs"
|
||||
|
||||
def test_config_manager_creation_custom_base_dir(self):
|
||||
"""Test creating ConfigManager with custom base directory."""
|
||||
custom_dir = "/tmp/custom_configs"
|
||||
manager = ConfigManager(base_dir=custom_dir)
|
||||
|
||||
# Paths are resolved, so compare with resolved path
|
||||
assert manager._base_dir == Path(custom_dir).resolve()
|
||||
|
||||
def test_get_config_file_path(self):
|
||||
"""Test getting config file path for a kernel."""
|
||||
manager = ConfigManager(base_dir="/tmp")
|
||||
|
||||
dir_path = manager.get_config_file_path("silu_mul_fp8")
|
||||
assert dir_path == Path("/tmp/silu_mul_fp8")
|
||||
|
||||
file_path = manager.get_config_file_path("silu_mul_fp8", "nvidia_h100")
|
||||
assert file_path == Path("/tmp/silu_mul_fp8/nvidia_h100.json")
|
||||
|
||||
def test_ensure_base_dir_exists(self):
|
||||
"""Test ensuring base directory exists."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
base_dir = Path(temp_dir) / "non_existent" / "configs"
|
||||
manager = ConfigManager(base_dir=base_dir)
|
||||
assert not base_dir.exists()
|
||||
|
||||
returned_path = manager.ensure_base_dir_exists()
|
||||
|
||||
assert base_dir.exists()
|
||||
assert base_dir.is_dir()
|
||||
assert returned_path == base_dir
|
||||
|
||||
def test_load_config_set_file_not_exists(self):
|
||||
"""Test loading config set when file doesn't exist."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
manager = ConfigManager(base_dir=temp_dir)
|
||||
config_set = manager.load_config_set("non_existent_kernel")
|
||||
|
||||
assert isinstance(config_set, ConfigSet)
|
||||
assert config_set.kernel_name == "non_existent_kernel"
|
||||
assert config_set.get_platforms() == []
|
||||
|
||||
def test_load_config_set_valid_file(self):
|
||||
"""Test loading config set from per-platform files."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
kernel_config = {
|
||||
"block_sizes": [128, 64],
|
||||
"num_warps": 8,
|
||||
"num_stages": 6,
|
||||
"pid_type": "persistent_interleaved",
|
||||
}
|
||||
kernel_dir = Path(temp_dir) / "test_kernel"
|
||||
kernel_dir.mkdir()
|
||||
platform_file = kernel_dir / "h100.json"
|
||||
with open(platform_file, "w") as f:
|
||||
json.dump(
|
||||
[{"key": {"batch": 32, "hidden": 4096}, "config": kernel_config}],
|
||||
f,
|
||||
)
|
||||
|
||||
manager = ConfigManager(base_dir=temp_dir)
|
||||
config_set = manager.load_config_set("test_kernel")
|
||||
|
||||
assert isinstance(config_set, ConfigSet)
|
||||
assert config_set.kernel_name == "test_kernel"
|
||||
assert config_set.get_platforms() == ["h100"]
|
||||
|
||||
internal_key = CaseKey({"batch": 32, "hidden": 4096})
|
||||
config = config_set.get_config("h100", internal_key)
|
||||
assert isinstance(config, helion.Config)
|
||||
assert config.block_sizes == [128, 64]
|
||||
assert config.num_warps == 8
|
||||
|
||||
def test_load_config_set_invalid_json(self):
|
||||
"""Test loading config set from file with invalid JSON."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
kernel_dir = Path(temp_dir) / "test_kernel"
|
||||
kernel_dir.mkdir()
|
||||
config_file = kernel_dir / "h100.json"
|
||||
with open(config_file, "w") as f:
|
||||
f.write("invalid json content {")
|
||||
|
||||
manager = ConfigManager(base_dir=temp_dir)
|
||||
config_set = manager.load_config_set("test_kernel")
|
||||
|
||||
assert isinstance(config_set, ConfigSet)
|
||||
assert config_set.kernel_name == "test_kernel"
|
||||
assert config_set.get_platforms() == []
|
||||
|
||||
def test_save_config_set(self):
|
||||
"""Test saving ConfigSet to per-platform files."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
kernel_config = {
|
||||
"block_sizes": [256, 128],
|
||||
"num_warps": 16,
|
||||
"num_stages": 8,
|
||||
"pid_type": "persistent_blocked",
|
||||
}
|
||||
data = {
|
||||
"h100": [
|
||||
{"key": {"batch": 32, "hidden": 4096}, "config": kernel_config},
|
||||
]
|
||||
}
|
||||
config_set = ConfigSet.from_dict("test_kernel", data)
|
||||
|
||||
manager = ConfigManager(base_dir=temp_dir)
|
||||
saved_path = manager.save_config_set(config_set)
|
||||
|
||||
expected_dir = Path(temp_dir) / "test_kernel"
|
||||
assert saved_path == expected_dir
|
||||
assert saved_path.is_dir()
|
||||
|
||||
platform_file = expected_dir / "h100.json"
|
||||
assert platform_file.exists()
|
||||
with open(platform_file) as f:
|
||||
loaded_data = json.load(f)
|
||||
assert isinstance(loaded_data, list)
|
||||
assert len(loaded_data) == 1
|
||||
entry = loaded_data[0]
|
||||
assert entry["key"] == {"batch": 32, "hidden": 4096}
|
||||
assert entry["config"] == kernel_config
|
||||
|
||||
def test_save_config_set_creates_directory(self):
|
||||
"""Test that save_config_set creates parent directories if needed."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
nested_dir = Path(temp_dir) / "nested" / "configs"
|
||||
data = {
|
||||
"h100": [
|
||||
{"key": {}, "config": {"num_warps": 4}},
|
||||
]
|
||||
}
|
||||
config_set = ConfigSet.from_dict("test_kernel", data)
|
||||
|
||||
manager = ConfigManager(base_dir=nested_dir)
|
||||
saved_path = manager.save_config_set(config_set)
|
||||
|
||||
assert nested_dir.exists()
|
||||
assert nested_dir.is_dir()
|
||||
assert saved_path.is_dir()
|
||||
assert (saved_path / "h100.json").exists()
|
||||
|
||||
def test_get_platform_configs(self):
|
||||
"""Test getting all configs for a specific platform."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
config_1 = {"num_warps": 4, "num_stages": 3, "block_sizes": [64, 32]}
|
||||
config_2 = {"num_warps": 8, "num_stages": 5, "block_sizes": [128, 64]}
|
||||
default_config = {
|
||||
"num_warps": 16,
|
||||
"num_stages": 7,
|
||||
"block_sizes": [256, 128],
|
||||
}
|
||||
config_3 = {"num_warps": 2, "num_stages": 2, "block_sizes": [32, 16]}
|
||||
|
||||
kernel_dir = Path(temp_dir) / "test_kernel"
|
||||
kernel_dir.mkdir()
|
||||
with open(kernel_dir / "h100.json", "w") as f:
|
||||
json.dump(
|
||||
[
|
||||
{"key": {"batch": 32, "hidden": 4096}, "config": config_1},
|
||||
{"key": {"batch": 64, "hidden": 2048}, "config": config_2},
|
||||
{"key": {}, "config": default_config},
|
||||
],
|
||||
f,
|
||||
)
|
||||
with open(kernel_dir / "a100.json", "w") as f:
|
||||
json.dump(
|
||||
[{"key": {"batch": 16, "hidden": 1024}, "config": config_3}],
|
||||
f,
|
||||
)
|
||||
|
||||
manager = ConfigManager(base_dir=temp_dir)
|
||||
|
||||
key_b32_h4096 = CaseKey({"batch": 32, "hidden": 4096})
|
||||
key_b64_h2048 = CaseKey({"batch": 64, "hidden": 2048})
|
||||
key_b16_h1024 = CaseKey({"batch": 16, "hidden": 1024})
|
||||
|
||||
h100_configs = manager.get_platform_configs("test_kernel", "h100")
|
||||
assert len(h100_configs) == 3
|
||||
assert key_b32_h4096 in h100_configs
|
||||
assert key_b64_h2048 in h100_configs
|
||||
assert CaseKey.default() in h100_configs
|
||||
for config in h100_configs.values():
|
||||
assert isinstance(config, helion.Config)
|
||||
|
||||
assert h100_configs[key_b32_h4096].num_warps == 4
|
||||
assert h100_configs[CaseKey.default()].num_stages == 7
|
||||
|
||||
a100_configs = manager.get_platform_configs("test_kernel", "a100")
|
||||
assert len(a100_configs) == 1
|
||||
assert key_b16_h1024 in a100_configs
|
||||
assert isinstance(a100_configs[key_b16_h1024], helion.Config)
|
||||
assert a100_configs[key_b16_h1024].num_warps == 2
|
||||
|
||||
nonexistent_configs = manager.get_platform_configs("test_kernel", "v100")
|
||||
assert len(nonexistent_configs) == 0
|
||||
|
||||
def test_singleton_returns_same_instance(self):
|
||||
"""Test that ConfigManager returns the same instance on repeated calls."""
|
||||
manager1 = ConfigManager(base_dir="/tmp/test_singleton")
|
||||
manager2 = ConfigManager(base_dir="/tmp/test_singleton")
|
||||
|
||||
assert manager1 is manager2
|
||||
|
||||
def test_singleton_with_default_base_dir(self):
|
||||
"""Test singleton behavior with default base directory."""
|
||||
manager1 = ConfigManager()
|
||||
manager2 = ConfigManager()
|
||||
|
||||
assert manager1 is manager2
|
||||
assert manager1._base_dir == manager2._base_dir
|
||||
|
||||
def test_singleton_error_on_different_base_dir(self):
|
||||
"""Test that ConfigManager raises error when created with different base_dir."""
|
||||
ConfigManager(base_dir="/tmp/first_dir")
|
||||
|
||||
with pytest.raises(ValueError, match="singleton already exists"):
|
||||
ConfigManager(base_dir="/tmp/different_dir")
|
||||
|
||||
def test_reset_instance_allows_new_base_dir(self):
|
||||
"""Test that reset_instance allows creating with a new base_dir."""
|
||||
manager1 = ConfigManager(base_dir="/tmp/first_dir")
|
||||
assert manager1._base_dir == Path("/tmp/first_dir").resolve()
|
||||
|
||||
ConfigManager.reset_instance()
|
||||
|
||||
manager2 = ConfigManager(base_dir="/tmp/second_dir")
|
||||
assert manager2._base_dir == Path("/tmp/second_dir").resolve()
|
||||
assert manager1 is not manager2
|
||||
|
||||
def test_get_instance_returns_existing(self):
|
||||
"""Test that get_instance returns the existing singleton."""
|
||||
manager1 = ConfigManager(base_dir="/tmp/test_get_instance")
|
||||
manager2 = ConfigManager.get_instance()
|
||||
|
||||
assert manager1 is manager2
|
||||
|
||||
def test_get_instance_raises_if_not_initialized(self):
|
||||
"""Test that get_instance raises RuntimeError if no instance exists."""
|
||||
with pytest.raises(RuntimeError, match="has not been created"):
|
||||
ConfigManager.get_instance()
|
||||
@@ -0,0 +1,165 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the dynamic_per_token_scaled_fp8_quant helion kernel
|
||||
|
||||
Run `pytest tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py`.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import FakeTensorMode
|
||||
|
||||
from tests.kernels.helion.utils import skip_if_platform_unsupported
|
||||
from tests.kernels.quant_utils import FP8_DTYPE
|
||||
from vllm.kernels.helion.case_key import CaseKey
|
||||
from vllm.kernels.helion.config_manager import ConfigManager
|
||||
from vllm.kernels.helion.ops.dynamic_per_token_scaled_fp8_quant import (
|
||||
_pick_cache,
|
||||
baseline,
|
||||
dynamic_per_token_scaled_fp8_quant,
|
||||
pick_config,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.import_utils import has_helion
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
def _generate_fake_input(num_tokens: int, hidden_size: int) -> tuple[Any, ...]:
|
||||
with FakeTensorMode():
|
||||
input = torch.randn(
|
||||
num_tokens, hidden_size, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
result = torch.empty(
|
||||
input.shape, device=input.device, dtype=current_platform.fp8_dtype()
|
||||
)
|
||||
scale = torch.empty((num_tokens, 1), device=input.device, dtype=torch.float32)
|
||||
scale_ub = torch.mean(input).to(torch.float32)
|
||||
args = (result, input, scale, scale_ub)
|
||||
return args
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_config_manager_singleton():
|
||||
ConfigManager.reset_instance()
|
||||
ConfigManager()
|
||||
yield
|
||||
ConfigManager.reset_instance()
|
||||
|
||||
|
||||
class TestDynamicPerTokenScaledFp8QuantConfigPicker:
|
||||
def setup_method(self):
|
||||
_pick_cache.clear()
|
||||
|
||||
def test_config_picker_exact_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "num_tokens": 16}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(16, 4096)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16})
|
||||
|
||||
def test_config_picker_closest_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "num_tokens": 32}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(20, 3000)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey({"hidden_size": 2048, "num_tokens": 32})
|
||||
|
||||
def test_config_picker_no_configs(self):
|
||||
config_keys: list[dict] = []
|
||||
|
||||
args = _generate_fake_input(16, 4096)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key is None
|
||||
|
||||
def test_config_picker_fallback_to_largest(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "num_tokens": 16}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(32, 8192)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16})
|
||||
|
||||
|
||||
class TestDynamicPerTokenScaledFp8QuantCorrectness:
|
||||
@pytest.mark.parametrize("num_tokens", [1, 7, 4096])
|
||||
@pytest.mark.parametrize("hidden_size", [17, 1024, 1025, 1026, 5137, 8193])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float])
|
||||
@pytest.mark.parametrize("has_scale_ub", [True, False])
|
||||
@pytest.mark.parametrize("seed", [0])
|
||||
def test_dynamic_per_token_fp8_quant(
|
||||
self,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
has_scale_ub: bool,
|
||||
seed: int,
|
||||
) -> None:
|
||||
skip_if_platform_unsupported("dynamic_per_token_scaled_fp8_quant")
|
||||
set_random_seed(seed)
|
||||
|
||||
x = (
|
||||
torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") + 1e-6
|
||||
) # avoid nans
|
||||
|
||||
scale_ub = (
|
||||
torch.mean(x).to(dtype=torch.float32, device="cuda")
|
||||
if has_scale_ub
|
||||
else None
|
||||
)
|
||||
|
||||
ref_out = torch.empty(x.shape, device="cuda", dtype=FP8_DTYPE)
|
||||
ref_scales = torch.empty((x.shape[0], 1), device="cuda", dtype=torch.float32)
|
||||
baseline(ref_out, x, ref_scales, scale_ub)
|
||||
|
||||
ops_out = torch.empty(x.shape, device="cuda", dtype=FP8_DTYPE)
|
||||
ops_scales = torch.empty((x.shape[0], 1), device="cuda", dtype=torch.float32)
|
||||
dynamic_per_token_scaled_fp8_quant(ops_out, x, ops_scales, scale_ub)
|
||||
|
||||
torch.testing.assert_close(ref_scales, ops_scales)
|
||||
# allow 1 ULP difference
|
||||
assert (
|
||||
ref_out.view(torch.uint8).to(torch.int16)
|
||||
- ops_out.view(torch.uint8).to(torch.int16)
|
||||
).abs().max() <= 1
|
||||
|
||||
|
||||
class TestDynamicPerTokenScaledFp8QuantIntegration:
|
||||
def test_kernel_registration_integration(self):
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
assert "dynamic_per_token_scaled_fp8_quant" in registered_kernels
|
||||
|
||||
kernel_wrapper = registered_kernels["dynamic_per_token_scaled_fp8_quant"]
|
||||
assert kernel_wrapper.op_name == "dynamic_per_token_scaled_fp8_quant"
|
||||
assert kernel_wrapper._config_picker is not None
|
||||
assert kernel_wrapper._mutates_args == ["result", "scale"]
|
||||
|
||||
def test_fake_impl_functionality(self):
|
||||
skip_if_platform_unsupported("dynamic_per_token_scaled_fp8_quant")
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
kernel_wrapper = registered_kernels["dynamic_per_token_scaled_fp8_quant"]
|
||||
fake_impl = kernel_wrapper._fake_impl
|
||||
|
||||
args = _generate_fake_input(16, 4096)
|
||||
assert fake_impl(*args) is None
|
||||
@@ -0,0 +1,261 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the fused_qk_norm_rope helion kernel
|
||||
|
||||
Run `pytest tests/kernels/helion/test_fused_qk_norm_rope.py`.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import FakeTensorMode
|
||||
|
||||
from tests.kernels.helion.utils import skip_if_platform_unsupported
|
||||
from vllm.benchmarks.lib.utils import default_vllm_config
|
||||
from vllm.kernels.helion.case_key import CaseKey
|
||||
from vllm.kernels.helion.config_manager import ConfigManager
|
||||
from vllm.kernels.helion.ops.fused_qk_norm_rope import (
|
||||
_pick_cache,
|
||||
baseline,
|
||||
fused_qk_norm_rope,
|
||||
pick_config,
|
||||
)
|
||||
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
|
||||
from vllm.utils.import_utils import has_helion
|
||||
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
@default_vllm_config()
|
||||
def _generate_fake_input(
|
||||
num_tokens: int, num_q_heads: int, num_kv_heads: int
|
||||
) -> tuple[Any, ...]:
|
||||
with FakeTensorMode():
|
||||
head_dim = 128
|
||||
eps = 1e-6
|
||||
is_neox = True
|
||||
rotary_ratio = 1.0
|
||||
device = "cuda"
|
||||
dtype = torch.bfloat16
|
||||
total_dim = (num_q_heads + 2 * num_kv_heads) * head_dim
|
||||
qkv = torch.randn(num_tokens, total_dim, dtype=dtype, device=device)
|
||||
positions = torch.arange(num_tokens, dtype=torch.long, device=device)
|
||||
q_weight = torch.normal(
|
||||
mean=1.0,
|
||||
std=1.0,
|
||||
size=(head_dim,),
|
||||
dtype=qkv.dtype,
|
||||
device=device,
|
||||
)
|
||||
k_weight = torch.normal(
|
||||
mean=1.0,
|
||||
std=1.0,
|
||||
size=(head_dim,),
|
||||
dtype=qkv.dtype,
|
||||
device=device,
|
||||
)
|
||||
rotary_dim = int(head_dim * rotary_ratio)
|
||||
rope = RotaryEmbedding(
|
||||
head_size=head_dim,
|
||||
rotary_dim=rotary_dim,
|
||||
max_position_embeddings=4096,
|
||||
base=10000.0,
|
||||
is_neox_style=is_neox,
|
||||
dtype=dtype,
|
||||
).to(device)
|
||||
args = (
|
||||
qkv,
|
||||
num_q_heads,
|
||||
num_kv_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope.cos_sin_cache,
|
||||
is_neox,
|
||||
positions.view(-1),
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_config_manager_singleton():
|
||||
ConfigManager.reset_instance()
|
||||
ConfigManager()
|
||||
yield
|
||||
ConfigManager.reset_instance()
|
||||
|
||||
|
||||
class TestFusedQkNormRopeConfigPicker:
|
||||
def setup_method(self):
|
||||
_pick_cache.clear()
|
||||
|
||||
def test_config_picker_exact_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 16}),
|
||||
CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}
|
||||
)
|
||||
|
||||
def test_config_picker_closest_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 16}),
|
||||
CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 32}),
|
||||
CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 16}),
|
||||
CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 32}),
|
||||
CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 16}),
|
||||
CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 32}),
|
||||
CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}),
|
||||
CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 32}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(20, 3000, 70)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"q_heads": 2048, "kv_heads": 64, "num_tokens": 32}
|
||||
)
|
||||
|
||||
def test_config_picker_no_configs(self):
|
||||
config_keys: list[dict] = []
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key is None
|
||||
|
||||
def test_config_picker_fallback_to_largest(self):
|
||||
config_keys = [
|
||||
CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 16}),
|
||||
CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 32}),
|
||||
CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 16}),
|
||||
CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 32}),
|
||||
CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 16}),
|
||||
CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 32}),
|
||||
CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}),
|
||||
CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 32}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(64, 8192, 256)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"q_heads": 4096, "kv_heads": 128, "num_tokens": 32}
|
||||
)
|
||||
|
||||
|
||||
class TestFusedQkNormRopeCorrectness:
|
||||
@pytest.mark.parametrize(
|
||||
"num_heads, num_kv_heads, head_dim", [(16, 4, 128), (64, 8, 128)]
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", [1, 7, 1024, 1025])
|
||||
@pytest.mark.parametrize("is_neox", [False, True])
|
||||
@pytest.mark.parametrize("rotary_ratio", [1.0, 0.5, 0.25])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@default_vllm_config()
|
||||
def test_fused_qk_norm_rope(
|
||||
self,
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_dim: int,
|
||||
num_tokens: int,
|
||||
is_neox: bool,
|
||||
rotary_ratio: float,
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
skip_if_platform_unsupported("fused_qk_norm_rope")
|
||||
|
||||
torch.manual_seed(42)
|
||||
eps = 1e-6
|
||||
device = "cuda"
|
||||
total_dim = (num_heads + 2 * num_kv_heads) * head_dim
|
||||
ref_qkv = torch.empty(
|
||||
num_tokens, total_dim, dtype=dtype, device=device
|
||||
).uniform_(-0.1, 0.1)
|
||||
ops_qkv = ref_qkv.clone()
|
||||
positions = torch.arange(num_tokens, dtype=torch.long, device=device)
|
||||
q_weight = torch.empty(head_dim, dtype=dtype, device=device).uniform_(0.8, 1.2)
|
||||
k_weight = torch.empty(head_dim, dtype=dtype, device=device).uniform_(0.8, 1.2)
|
||||
rotary_dim = int(head_dim * rotary_ratio)
|
||||
rope = RotaryEmbedding(
|
||||
head_size=head_dim,
|
||||
rotary_dim=rotary_dim,
|
||||
max_position_embeddings=40960,
|
||||
base=10000.0,
|
||||
is_neox_style=is_neox,
|
||||
dtype=dtype,
|
||||
).to(device)
|
||||
|
||||
baseline(
|
||||
ref_qkv,
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope.cos_sin_cache,
|
||||
is_neox,
|
||||
positions.view(-1),
|
||||
)
|
||||
|
||||
fused_qk_norm_rope(
|
||||
ops_qkv,
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope.cos_sin_cache,
|
||||
is_neox,
|
||||
positions.view(-1),
|
||||
)
|
||||
|
||||
if dtype == torch.bfloat16:
|
||||
atol = 5e-2
|
||||
rtol = 5e-2
|
||||
else:
|
||||
atol = 1e-2
|
||||
rtol = 1e-2
|
||||
|
||||
torch.testing.assert_close(
|
||||
ref_qkv,
|
||||
ops_qkv,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
|
||||
|
||||
class TestFusedQkNormRopeIntegration:
|
||||
def test_kernel_registration_integration(self):
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
assert "fused_qk_norm_rope" in registered_kernels
|
||||
|
||||
kernel_wrapper = registered_kernels["fused_qk_norm_rope"]
|
||||
assert kernel_wrapper.op_name == "fused_qk_norm_rope"
|
||||
assert kernel_wrapper._config_picker is not None
|
||||
assert kernel_wrapper._mutates_args == ["qkv"]
|
||||
|
||||
def test_fake_impl_functionality(self):
|
||||
skip_if_platform_unsupported("fused_qk_norm_rope")
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
kernel_wrapper = registered_kernels["fused_qk_norm_rope"]
|
||||
fake_impl = kernel_wrapper._fake_impl
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
assert fake_impl(*args) is None
|
||||
@@ -0,0 +1,45 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for Helion kernel availability and basic functionality.
|
||||
|
||||
This module demonstrates the pattern for testing optional Helion kernels.
|
||||
Tests in this directory will be skipped if Helion is not installed.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.utils.import_utils import has_helion
|
||||
|
||||
# Skip entire module if helion is not available
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
import helion
|
||||
import helion.language as hl
|
||||
import torch
|
||||
|
||||
|
||||
def test_helion_kernel_compilation_smoke():
|
||||
"""Smoke test: compile and run a simple Helion kernel."""
|
||||
|
||||
@helion.kernel(autotune_effort="none")
|
||||
def add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
out = torch.empty_like(x)
|
||||
for tile in hl.tile(x.size()):
|
||||
out[tile] = x[tile] + y[tile]
|
||||
return out
|
||||
|
||||
# Create test tensors
|
||||
x = torch.randn(1024, device="cuda", dtype=torch.float32)
|
||||
y = torch.randn(1024, device="cuda", dtype=torch.float32)
|
||||
|
||||
# Run the helion kernel
|
||||
result = add_kernel(x, y)
|
||||
|
||||
# Verify correctness
|
||||
expected = x + y
|
||||
assert torch.allclose(result, expected), "Helion kernel output mismatch"
|
||||
@@ -0,0 +1,205 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Test make_fx tracing and inductor pattern matching with HelionKernelWrapper."""
|
||||
|
||||
import contextlib
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.utils.import_utils import has_helion
|
||||
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
import helion
|
||||
import helion.language as hl
|
||||
from helion._compat import requires_torch_version
|
||||
|
||||
if not requires_torch_version("2.11"):
|
||||
pytest.skip(
|
||||
"HigherOrderOp requires PyTorch >= 2.11",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from helion._compiler._dynamo.higher_order_ops import (
|
||||
helion_kernel_side_table,
|
||||
helion_kernel_wrapper_mutation,
|
||||
)
|
||||
from torch._inductor.pattern_matcher import (
|
||||
PatternMatcherPass,
|
||||
fwd_only,
|
||||
register_replacement,
|
||||
select_decomp_table,
|
||||
)
|
||||
from torch.fx.experimental.proxy_tensor import make_fx
|
||||
|
||||
from vllm.kernels.helion.config_manager import ConfigManager
|
||||
from vllm.kernels.helion.register import HelionKernelWrapper
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _helion_mock_context():
|
||||
configs = {
|
||||
"default": helion.Config(block_sizes=[64], num_warps=2, num_stages=2),
|
||||
}
|
||||
mock_config_manager = Mock(spec=ConfigManager)
|
||||
mock_config_manager.get_platform_configs = Mock(return_value=configs)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm.kernels.helion.config_manager.ConfigManager",
|
||||
return_value=mock_config_manager,
|
||||
),
|
||||
patch(
|
||||
"vllm.kernels.helion.utils.get_canonical_gpu_name",
|
||||
return_value="nvidia_h200",
|
||||
),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
class TestMakeFxHop:
|
||||
def setup_method(self):
|
||||
helion_kernel_side_table.reset_table()
|
||||
|
||||
@pytest.mark.skip(reason="SymInt proxy tracking issue with PyTorch 2.11+")
|
||||
def test_make_fx_symbolic(self):
|
||||
def raw_add_scale(
|
||||
x: torch.Tensor, y: torch.Tensor, scale: float
|
||||
) -> tuple[torch.Tensor, int, torch.Tensor]:
|
||||
out_x = torch.empty_like(x)
|
||||
out_y = torch.empty_like(x)
|
||||
for tile in hl.tile(x.size()):
|
||||
out_x[tile] = x[tile] + y[tile] * scale
|
||||
out_y[tile] = out_x[tile] * 2.0
|
||||
return out_x, 42, out_y
|
||||
|
||||
input_x = torch.randn(7, 13)
|
||||
input_y = torch.randn(7, 13)
|
||||
scale = 0.5
|
||||
|
||||
with _helion_mock_context():
|
||||
wrapper = HelionKernelWrapper(
|
||||
raw_kernel_func=raw_add_scale,
|
||||
op_name="test_make_fx",
|
||||
fake_impl=lambda *a, **kw: None,
|
||||
config_picker=lambda args, keys: "default",
|
||||
)
|
||||
|
||||
def fn(x, y):
|
||||
return wrapper(x, y, scale)
|
||||
|
||||
gm = make_fx(fn, tracing_mode="symbolic")(input_x, input_y)
|
||||
|
||||
hop_nodes = [
|
||||
n
|
||||
for n in gm.graph.nodes
|
||||
if n.op == "call_function" and n.target is helion_kernel_wrapper_mutation
|
||||
]
|
||||
assert len(hop_nodes) == 1
|
||||
node = hop_nodes[0]
|
||||
|
||||
assert node.kwargs["constant_args"]["scale"] == scale
|
||||
assert set(node.kwargs["tensor_args"]) == {"x", "y"}
|
||||
|
||||
specs = node.kwargs["output_spec"]["leaf_specs"]
|
||||
tensor_specs = [s for s in specs if s["type"] == "tensor"]
|
||||
scalar_specs = [s for s in specs if s["type"] == "scalar"]
|
||||
assert len(tensor_specs) == 2
|
||||
assert len(scalar_specs) == 1
|
||||
|
||||
for spec in tensor_specs:
|
||||
assert spec["dtype"] == input_x.dtype
|
||||
|
||||
assert scalar_specs[0]["scalar_value"] == 42
|
||||
|
||||
for val in node.meta["val"]:
|
||||
assert all(isinstance(s, torch.SymInt) for s in val.shape)
|
||||
|
||||
# Both out_x and out_y are empty_like(x), so output shapes == input shape
|
||||
input_node = next(n for n in gm.graph.nodes if n.op == "placeholder")
|
||||
input_shape = input_node.meta["val"].shape
|
||||
for val in node.meta["val"]:
|
||||
assert len(val.shape) == len(input_shape)
|
||||
for out_s, in_s in zip(val.shape, input_shape):
|
||||
assert out_s == in_s
|
||||
|
||||
@pytest.mark.skip(reason="SymInt proxy tracking issue with PyTorch 2.11+")
|
||||
def test_pattern_matcher_replaces_with_helion_hop(self):
|
||||
def raw_silu_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
M, N = x.size()
|
||||
out = torch.empty_like(x)
|
||||
for tile_m, tile_n in hl.tile([M, N]):
|
||||
out[tile_m, tile_n] = (
|
||||
torch.nn.functional.silu(x[tile_m, tile_n]) * y[tile_m, tile_n]
|
||||
)
|
||||
return out
|
||||
|
||||
with _helion_mock_context():
|
||||
wrapper = HelionKernelWrapper(
|
||||
raw_kernel_func=raw_silu_mul,
|
||||
op_name="test_pm_silu_mul",
|
||||
fake_impl=lambda *a, **kw: None,
|
||||
config_picker=lambda args, keys: "default",
|
||||
)
|
||||
|
||||
def pattern(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return torch.nn.functional.silu(x) * y
|
||||
|
||||
def replacement(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return wrapper(x, y)
|
||||
|
||||
inputs = [torch.randn(8, 16), torch.randn(8, 16)]
|
||||
|
||||
pm_pass = PatternMatcherPass(pass_name="test_helion_replacement")
|
||||
register_replacement(pattern, replacement, inputs, fwd_only, pm_pass)
|
||||
|
||||
def model(x, y):
|
||||
return torch.nn.functional.silu(x) * y
|
||||
|
||||
decompositions = select_decomp_table()
|
||||
input_x = torch.randn(8, 16)
|
||||
input_y = torch.randn(8, 16)
|
||||
gm = make_fx(model, decompositions, tracing_mode="symbolic")(
|
||||
input_x, input_y
|
||||
)
|
||||
|
||||
def count_hop_nodes(graph):
|
||||
return sum(
|
||||
1
|
||||
for n in graph.nodes
|
||||
if n.op == "call_function"
|
||||
and n.target is helion_kernel_wrapper_mutation
|
||||
)
|
||||
|
||||
assert count_hop_nodes(gm.graph) == 0
|
||||
|
||||
match_count = pm_pass.apply(gm.graph)
|
||||
gm.graph.lint()
|
||||
gm.recompile()
|
||||
|
||||
assert match_count == 1
|
||||
assert count_hop_nodes(gm.graph) == 1
|
||||
|
||||
hop_node = next(
|
||||
n
|
||||
for n in gm.graph.nodes
|
||||
if n.op == "call_function"
|
||||
and n.target is helion_kernel_wrapper_mutation
|
||||
)
|
||||
|
||||
# raw_silu_mul returns empty_like(x), so output shape == input shape
|
||||
for val in hop_node.meta["val"]:
|
||||
assert all(isinstance(s, torch.SymInt) for s in val.shape)
|
||||
|
||||
input_node = next(n for n in gm.graph.nodes if n.op == "placeholder")
|
||||
input_shape = input_node.meta["val"].shape
|
||||
output_shape = hop_node.meta["val"][0].shape
|
||||
assert len(output_shape) == len(input_shape)
|
||||
for out_s, in_s in zip(output_shape, input_shape):
|
||||
assert out_s == in_s
|
||||
@@ -0,0 +1,243 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the per_token_group_fp8_quant helion kernel
|
||||
|
||||
Run `pytest tests/kernels/helion/test_per_token_group_fp8_quant.py`.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import FakeTensorMode
|
||||
|
||||
from tests.kernels.helion.utils import skip_if_platform_unsupported
|
||||
from tests.kernels.quant_utils import FP8_DTYPE
|
||||
from vllm.kernels.helion.case_key import CaseKey
|
||||
from vllm.kernels.helion.config_manager import ConfigManager
|
||||
from vllm.kernels.helion.ops.per_token_group_fp8_quant import (
|
||||
_pick_cache,
|
||||
baseline,
|
||||
per_token_group_fp8_quant,
|
||||
pick_config,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
get_fp8_min_max,
|
||||
)
|
||||
from vllm.utils.import_utils import has_helion
|
||||
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
def _generate_fake_input(
|
||||
num_tokens: int, hidden_size: int, group_size: int
|
||||
) -> tuple[Any, ...]:
|
||||
with FakeTensorMode():
|
||||
input = torch.randn(
|
||||
(num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
output_q = torch.empty(input.shape, device=input.device, dtype=FP8_DTYPE)
|
||||
output_s = torch.empty(
|
||||
(num_tokens, hidden_size // group_size),
|
||||
device=input.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
use_ue8m0 = False
|
||||
column_major = False
|
||||
fp8_min, fp8_max = get_fp8_min_max()
|
||||
eps = 1e-10
|
||||
args = (
|
||||
input,
|
||||
output_q,
|
||||
output_s,
|
||||
group_size,
|
||||
eps,
|
||||
fp8_min,
|
||||
fp8_max,
|
||||
use_ue8m0,
|
||||
column_major,
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_config_manager_singleton():
|
||||
ConfigManager.reset_instance()
|
||||
ConfigManager()
|
||||
yield
|
||||
ConfigManager.reset_instance()
|
||||
|
||||
|
||||
class TestPerTokenGroupFp8QuantConfigPicker:
|
||||
def setup_method(self):
|
||||
_pick_cache.clear()
|
||||
|
||||
def test_config_picker_exact_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"hidden_size": 4096, "group_size": 128, "num_tokens": 16}
|
||||
)
|
||||
|
||||
def test_config_picker_closest_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(20, 3000, 70)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"hidden_size": 2048, "group_size": 64, "num_tokens": 32}
|
||||
)
|
||||
|
||||
def test_config_picker_no_configs(self):
|
||||
config_keys: list[dict] = []
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key is None
|
||||
|
||||
def test_config_picker_fallback_to_largest(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(64, 8192, 256)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"hidden_size": 4096, "group_size": 128, "num_tokens": 32}
|
||||
)
|
||||
|
||||
|
||||
class TestPerTokenGroupFp8QuantCorrectness:
|
||||
@pytest.mark.parametrize(
|
||||
"shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512), (2048, 5120)]
|
||||
)
|
||||
@pytest.mark.parametrize("column_major", [False, True])
|
||||
@pytest.mark.parametrize("tma_aligned", [False, True])
|
||||
@pytest.mark.parametrize("scale_ue8m0", [False, True])
|
||||
@pytest.mark.parametrize("group_size", [64, 128])
|
||||
def test_per_token_group_fp8_quant(
|
||||
self,
|
||||
shape,
|
||||
column_major: bool,
|
||||
tma_aligned: bool,
|
||||
scale_ue8m0: bool,
|
||||
group_size: int,
|
||||
):
|
||||
skip_if_platform_unsupported("per_token_group_fp8_quant")
|
||||
|
||||
torch.manual_seed(42)
|
||||
num_tokens, hidden_size = shape
|
||||
fp8_min, fp8_max = get_fp8_min_max()
|
||||
eps = 1e-10
|
||||
input = (
|
||||
torch.randn((num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16)
|
||||
* 8
|
||||
)
|
||||
ref_q = torch.empty(input.shape, device=input.device, dtype=FP8_DTYPE)
|
||||
ops_q = ref_q.clone()
|
||||
|
||||
groups_per_row = hidden_size // group_size
|
||||
if column_major:
|
||||
if tma_aligned:
|
||||
tma_alignment = 4
|
||||
tma_aligned_m = (
|
||||
(num_tokens + tma_alignment - 1) // tma_alignment * tma_alignment
|
||||
)
|
||||
shape = (num_tokens, groups_per_row)
|
||||
stride = (1, tma_aligned_m)
|
||||
ref_s = torch.empty_strided(
|
||||
shape, stride, device=input.device, dtype=torch.float32
|
||||
)
|
||||
else:
|
||||
ref_s = torch.empty(
|
||||
(groups_per_row, num_tokens),
|
||||
device=input.device,
|
||||
dtype=torch.float32,
|
||||
).transpose(0, 1)
|
||||
else:
|
||||
ref_s = torch.empty(
|
||||
(num_tokens, groups_per_row), device=input.device, dtype=torch.float32
|
||||
)
|
||||
|
||||
ops_s = ref_s.clone()
|
||||
|
||||
baseline(
|
||||
input,
|
||||
ref_q,
|
||||
ref_s,
|
||||
group_size,
|
||||
eps,
|
||||
fp8_min,
|
||||
fp8_max,
|
||||
scale_ue8m0,
|
||||
column_major,
|
||||
tma_aligned,
|
||||
)
|
||||
per_token_group_fp8_quant(
|
||||
input,
|
||||
ops_q,
|
||||
ops_s,
|
||||
group_size,
|
||||
eps,
|
||||
fp8_min,
|
||||
fp8_max,
|
||||
scale_ue8m0,
|
||||
column_major,
|
||||
tma_aligned,
|
||||
)
|
||||
|
||||
assert torch.allclose(ref_s, ops_s)
|
||||
# allow 1 ULP difference
|
||||
assert (
|
||||
ref_q.view(torch.uint8).to(torch.int16)
|
||||
- ops_q.view(torch.uint8).to(torch.int16)
|
||||
).abs().max() <= 1
|
||||
|
||||
|
||||
class TestPerTokenGroupFp8QuantIntegration:
|
||||
def test_kernel_registration_integration(self):
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
assert "per_token_group_fp8_quant" in registered_kernels
|
||||
|
||||
kernel_wrapper = registered_kernels["per_token_group_fp8_quant"]
|
||||
assert kernel_wrapper.op_name == "per_token_group_fp8_quant"
|
||||
assert kernel_wrapper._config_picker is not None
|
||||
assert kernel_wrapper._mutates_args == ["output_q", "output_s"]
|
||||
|
||||
def test_fake_impl_functionality(self):
|
||||
skip_if_platform_unsupported("per_token_group_fp8_quant")
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
kernel_wrapper = registered_kernels["per_token_group_fp8_quant"]
|
||||
fake_impl = kernel_wrapper._fake_impl
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
assert fake_impl(*args) is None
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,207 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the rms_norm_dynamic_per_token_quant helion kernel
|
||||
|
||||
Run `pytest tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py`.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import FakeTensorMode
|
||||
|
||||
from tests.kernels.helion.utils import skip_if_platform_unsupported
|
||||
from vllm.kernels.helion.case_key import CaseKey
|
||||
from vllm.kernels.helion.config_manager import ConfigManager
|
||||
from vllm.kernels.helion.ops.rms_norm_dynamic_per_token_quant import (
|
||||
_pick_cache,
|
||||
baseline,
|
||||
pick_config,
|
||||
rms_norm_dynamic_per_token_quant,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.import_utils import has_helion
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
def _generate_fake_input(num_tokens: int, hidden_size: int) -> tuple[Any, ...]:
|
||||
with FakeTensorMode():
|
||||
input = torch.randn(
|
||||
num_tokens, hidden_size, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
result = torch.empty(
|
||||
input.shape, device=input.device, dtype=current_platform.fp8_dtype()
|
||||
)
|
||||
scale = torch.empty((num_tokens, 1), device=input.device, dtype=torch.float32)
|
||||
scale_ub = torch.mean(input).to(torch.float32)
|
||||
residual = torch.randn_like(input)
|
||||
weight = torch.normal(
|
||||
mean=1.0,
|
||||
std=1.0,
|
||||
size=(hidden_size,),
|
||||
dtype=input.dtype,
|
||||
device=input.device,
|
||||
)
|
||||
epsilon = 1e-6
|
||||
args = (result, input, weight, scale, epsilon, scale_ub, residual)
|
||||
return args
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_config_manager_singleton():
|
||||
ConfigManager.reset_instance()
|
||||
ConfigManager()
|
||||
yield
|
||||
ConfigManager.reset_instance()
|
||||
|
||||
|
||||
class TestRmsNormDynamicPerTokenQuantConfigPicker:
|
||||
def setup_method(self):
|
||||
_pick_cache.clear()
|
||||
|
||||
def test_config_picker_exact_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "num_tokens": 16}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(16, 4096)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16})
|
||||
|
||||
def test_config_picker_closest_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "num_tokens": 32}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(20, 3000)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey({"hidden_size": 2048, "num_tokens": 32})
|
||||
|
||||
def test_config_picker_no_configs(self):
|
||||
config_keys: list[dict] = []
|
||||
|
||||
args = _generate_fake_input(16, 4096)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key is None
|
||||
|
||||
def test_config_picker_fallback_to_largest(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "num_tokens": 16}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(32, 8192)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16})
|
||||
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float]
|
||||
QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]
|
||||
VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
|
||||
# Avoid combinatorial explosion with full Cartesian product
|
||||
NUM_TOKENS_HIDDEN_SIZES = [
|
||||
*[(1, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5120, 5137]],
|
||||
*[(2048, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5137]],
|
||||
*[(4096, i) for i in [1, 64, 5137]],
|
||||
]
|
||||
|
||||
ADD_RESIDUAL = [False, True]
|
||||
SCALE_UBS = [True, False]
|
||||
SEEDS = [0]
|
||||
|
||||
EPS = 1e-6
|
||||
|
||||
|
||||
class TestRmsNormDynamicPerTokenQuantCorrectness:
|
||||
@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
|
||||
@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
def test_rms_norm_dynamic_per_token_quant(
|
||||
self,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
add_residual: bool,
|
||||
has_scale_ub: bool,
|
||||
dtype: torch.dtype,
|
||||
quant_dtype: torch.dtype,
|
||||
seed: int,
|
||||
) -> None:
|
||||
skip_if_platform_unsupported("rms_norm_dynamic_per_token_quant")
|
||||
|
||||
set_random_seed(seed)
|
||||
|
||||
if has_scale_ub and quant_dtype != current_platform.fp8_dtype():
|
||||
# skip
|
||||
return
|
||||
|
||||
scale = 1 / (hidden_size)
|
||||
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda") * scale
|
||||
weight = torch.normal(
|
||||
mean=1.0, std=1.0, size=(hidden_size,), dtype=dtype, device=x.device
|
||||
)
|
||||
residual = torch.randn_like(x) * scale if add_residual else None
|
||||
scale_ub = (
|
||||
torch.mean(x).to(dtype=torch.float32, device="cuda")
|
||||
if has_scale_ub
|
||||
else None
|
||||
)
|
||||
|
||||
ref_out = torch.empty(x.shape, device=x.device, dtype=quant_dtype)
|
||||
ref_scales = torch.empty((x.shape[0], 1), device=x.device, dtype=torch.float32)
|
||||
ref_residual = residual.clone() if residual is not None else None
|
||||
baseline(ref_out, x, weight, ref_scales, EPS, scale_ub, ref_residual)
|
||||
|
||||
ops_out = torch.empty(x.shape, device=x.device, dtype=quant_dtype)
|
||||
ops_scales = torch.empty((x.shape[0], 1), device=x.device, dtype=torch.float32)
|
||||
ops_residual = residual.clone() if residual is not None else None
|
||||
rms_norm_dynamic_per_token_quant(
|
||||
ops_out, x, weight, ops_scales, EPS, scale_ub, ops_residual
|
||||
)
|
||||
|
||||
torch.testing.assert_close(ref_scales, ops_scales)
|
||||
# allow 1 ULP difference
|
||||
assert (
|
||||
ref_out.view(torch.uint8).to(torch.int16)
|
||||
- ops_out.view(torch.uint8).to(torch.int16)
|
||||
).abs().max() <= 1
|
||||
|
||||
if add_residual:
|
||||
torch.testing.assert_close(ref_residual, ops_residual)
|
||||
|
||||
|
||||
class TestRmsNormDynamicPerTokenQuantIntegration:
|
||||
def test_kernel_registration_integration(self):
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
assert "rms_norm_dynamic_per_token_quant" in registered_kernels
|
||||
|
||||
kernel_wrapper = registered_kernels["rms_norm_dynamic_per_token_quant"]
|
||||
assert kernel_wrapper.op_name == "rms_norm_dynamic_per_token_quant"
|
||||
assert kernel_wrapper._config_picker is not None
|
||||
assert kernel_wrapper._mutates_args == ["result", "scale", "residual"]
|
||||
|
||||
def test_fake_impl_functionality(self):
|
||||
skip_if_platform_unsupported("rms_norm_dynamic_per_token_quant")
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
kernel_wrapper = registered_kernels["rms_norm_dynamic_per_token_quant"]
|
||||
fake_impl = kernel_wrapper._fake_impl
|
||||
|
||||
args = _generate_fake_input(16, 4096)
|
||||
assert fake_impl(*args) is None
|
||||
@@ -0,0 +1,298 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the rms_norm_per_block_quant helion kernel
|
||||
|
||||
Run `pytest tests/kernels/helion/test_rms_norm_per_block_quant.py`.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import FakeTensorMode
|
||||
|
||||
from tests.kernels.helion.utils import skip_if_platform_unsupported
|
||||
from tests.kernels.quant_utils import FP8_DTYPE
|
||||
from vllm.kernels.helion.case_key import CaseKey
|
||||
from vllm.kernels.helion.config_manager import ConfigManager
|
||||
from vllm.kernels.helion.ops.rms_norm_per_block_quant import (
|
||||
_pick_cache,
|
||||
baseline,
|
||||
pick_config,
|
||||
rms_norm_per_block_quant,
|
||||
)
|
||||
from vllm.utils.import_utils import has_helion
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
def _generate_fake_input(
|
||||
num_tokens: int, hidden_size: int, group_size: int
|
||||
) -> tuple[Any, ...]:
|
||||
with FakeTensorMode():
|
||||
input = torch.randn(
|
||||
(num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
result = torch.empty(input.shape, device=input.device, dtype=FP8_DTYPE)
|
||||
scale = torch.empty(
|
||||
(num_tokens, hidden_size // group_size),
|
||||
device=input.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
scale_ub = torch.mean(input).to(scale.dtype)
|
||||
residual = torch.randn_like(input)
|
||||
weight = torch.normal(
|
||||
mean=1.0,
|
||||
std=1.0,
|
||||
size=(hidden_size,),
|
||||
dtype=input.dtype,
|
||||
device=input.device,
|
||||
)
|
||||
epsilon = 1e-6
|
||||
args = (
|
||||
result,
|
||||
input,
|
||||
weight,
|
||||
scale,
|
||||
epsilon,
|
||||
scale_ub,
|
||||
residual,
|
||||
group_size,
|
||||
False,
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_config_manager_singleton():
|
||||
ConfigManager.reset_instance()
|
||||
ConfigManager()
|
||||
yield
|
||||
ConfigManager.reset_instance()
|
||||
|
||||
|
||||
class TestRmsNormPerBlockQuantConfigPicker:
|
||||
def setup_method(self):
|
||||
_pick_cache.clear()
|
||||
|
||||
def test_config_picker_exact_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"hidden_size": 4096, "group_size": 128, "num_tokens": 16}
|
||||
)
|
||||
|
||||
def test_config_picker_closest_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(20, 3000, 70)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"hidden_size": 2048, "group_size": 64, "num_tokens": 32}
|
||||
)
|
||||
|
||||
def test_config_picker_no_configs(self):
|
||||
config_keys: list[dict] = []
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key is None
|
||||
|
||||
def test_config_picker_fallback_to_largest(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(64, 8192, 256)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"hidden_size": 4096, "group_size": 128, "num_tokens": 32}
|
||||
)
|
||||
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float]
|
||||
QUANT_DTYPES = [torch.int8, FP8_DTYPE]
|
||||
VEC_HIDDEN_SIZES = [64, 1024]
|
||||
# Avoid combinatorial explosion with full Cartesian product
|
||||
NUM_TOKENS_HIDDEN_SIZES = [
|
||||
*[(1, i) for i in [64, 128, 1024, 5120]],
|
||||
*[(2048, i) for i in [64, 1024]],
|
||||
*[(4096, i) for i in [64]],
|
||||
]
|
||||
|
||||
ADD_RESIDUAL = [False, True]
|
||||
SCALE_UBS = [True, False]
|
||||
GROUP_SIZES = [64, 128]
|
||||
TMA_ALIGNMENTS = [0, 4]
|
||||
SEEDS = [0]
|
||||
EPS = 1e-6
|
||||
|
||||
|
||||
class TestRmsNormPerBlockQuantCorrectness:
|
||||
@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
|
||||
@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
|
||||
@pytest.mark.parametrize("is_scale_transposed", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"group_size, tma_alignment",
|
||||
[*itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
|
||||
)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
def test_rms_norm_per_block_quant(
|
||||
self,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
add_residual: bool,
|
||||
has_scale_ub: bool,
|
||||
dtype: torch.dtype,
|
||||
quant_dtype: torch.dtype,
|
||||
is_scale_transposed: bool,
|
||||
group_size: int,
|
||||
tma_alignment: int,
|
||||
seed: int,
|
||||
) -> None:
|
||||
skip_if_platform_unsupported("rms_norm_per_block_quant")
|
||||
|
||||
set_random_seed(seed)
|
||||
|
||||
if hidden_size % group_size != 0:
|
||||
# skip
|
||||
return
|
||||
|
||||
if tma_alignment != 0 and hidden_size // group_size % tma_alignment == 0:
|
||||
# Skip tests where TMA alignment doesn't create extra padding to save time
|
||||
return
|
||||
|
||||
if has_scale_ub and quant_dtype != FP8_DTYPE:
|
||||
# skip
|
||||
return
|
||||
|
||||
scale = 1 / (hidden_size)
|
||||
input = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda") * scale
|
||||
weight = torch.normal(
|
||||
mean=1.0, std=1.0, size=(hidden_size,), dtype=dtype, device=input.device
|
||||
)
|
||||
residual = torch.randn_like(input) * scale if add_residual else None
|
||||
scale_ub = (
|
||||
torch.mean(input).to(dtype=torch.float32, device="cuda")
|
||||
if has_scale_ub
|
||||
else None
|
||||
)
|
||||
groups_per_row = hidden_size // group_size
|
||||
|
||||
ref_residual = residual.clone() if residual is not None else None
|
||||
ops_residual = residual.clone() if residual is not None else None
|
||||
ref_out = torch.empty(input.shape, device=input.device, dtype=quant_dtype)
|
||||
ops_out = ref_out.clone()
|
||||
|
||||
if is_scale_transposed:
|
||||
if tma_alignment == 0:
|
||||
ref_scales = torch.empty(
|
||||
(groups_per_row, num_tokens),
|
||||
device=input.device,
|
||||
dtype=torch.float32,
|
||||
).transpose(0, 1)
|
||||
else:
|
||||
tma_aligned_m = (
|
||||
(num_tokens + tma_alignment - 1) // tma_alignment * tma_alignment
|
||||
)
|
||||
shape = (num_tokens, groups_per_row)
|
||||
stride = (1, tma_aligned_m)
|
||||
ref_scales = torch.empty_strided(
|
||||
shape, stride, device=input.device, dtype=torch.float32
|
||||
)
|
||||
else:
|
||||
ref_scales = torch.empty(
|
||||
(num_tokens, groups_per_row), device=input.device, dtype=torch.float32
|
||||
)
|
||||
|
||||
ops_scales = ref_scales.clone()
|
||||
|
||||
baseline(
|
||||
ref_out,
|
||||
input,
|
||||
weight,
|
||||
ref_scales,
|
||||
EPS,
|
||||
scale_ub,
|
||||
ref_residual,
|
||||
group_size,
|
||||
is_scale_transposed,
|
||||
)
|
||||
ref_scales = ref_scales.contiguous()
|
||||
|
||||
rms_norm_per_block_quant(
|
||||
ops_out,
|
||||
input,
|
||||
weight,
|
||||
ops_scales,
|
||||
EPS,
|
||||
scale_ub,
|
||||
ops_residual,
|
||||
group_size,
|
||||
is_scale_transposed,
|
||||
)
|
||||
ops_scales = ops_scales.contiguous()
|
||||
|
||||
torch.testing.assert_close(ref_scales, ops_scales)
|
||||
# allow 1 ULP difference
|
||||
assert (
|
||||
ref_out.view(torch.uint8).to(torch.int16)
|
||||
- ops_out.view(torch.uint8).to(torch.int16)
|
||||
).abs().max() <= 1
|
||||
|
||||
if add_residual:
|
||||
torch.testing.assert_close(ref_residual, ops_residual)
|
||||
|
||||
|
||||
class TestRmsNormPerBlockQuantIntegration:
|
||||
def test_kernel_registration_integration(self):
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
assert "rms_norm_per_block_quant" in registered_kernels
|
||||
|
||||
kernel_wrapper = registered_kernels["rms_norm_per_block_quant"]
|
||||
assert kernel_wrapper.op_name == "rms_norm_per_block_quant"
|
||||
assert kernel_wrapper._config_picker is not None
|
||||
assert kernel_wrapper._mutates_args == ["result", "scale", "residual"]
|
||||
|
||||
def test_fake_impl_functionality(self):
|
||||
skip_if_platform_unsupported("rms_norm_per_block_quant")
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
kernel_wrapper = registered_kernels["rms_norm_per_block_quant"]
|
||||
fake_impl = kernel_wrapper._fake_impl
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
assert fake_impl(*args) is None
|
||||
@@ -0,0 +1,224 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the silu_and_mul_per_block_quant helion kernel
|
||||
Run `pytest tests/kernels/helion/test_silu_and_mul_per_block_quant.py`.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import FakeTensorMode
|
||||
|
||||
from tests.kernels.helion.utils import skip_if_platform_unsupported
|
||||
from tests.kernels.quant_utils import FP8_DTYPE
|
||||
from vllm.kernels.helion.case_key import CaseKey
|
||||
from vllm.kernels.helion.config_manager import ConfigManager
|
||||
from vllm.kernels.helion.ops.silu_and_mul_per_block_quant import (
|
||||
_pick_cache,
|
||||
baseline,
|
||||
pick_config,
|
||||
silu_and_mul_per_block_quant,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.import_utils import has_helion
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
def _generate_fake_input(
|
||||
num_tokens: int, intermediate_size: int, group_size: int
|
||||
) -> tuple[Any, ...]:
|
||||
with FakeTensorMode():
|
||||
in_dtype: torch.dtype = torch.bfloat16
|
||||
out_dtype: torch.dtype = current_platform.fp8_dtype()
|
||||
scale_dtype: torch.dtype = torch.float32
|
||||
input = torch.randn(
|
||||
num_tokens, 2 * intermediate_size, device="cuda", dtype=in_dtype
|
||||
)
|
||||
result = torch.empty(
|
||||
num_tokens, intermediate_size, device=input.device, dtype=out_dtype
|
||||
)
|
||||
scale = torch.empty(
|
||||
(num_tokens, intermediate_size // group_size),
|
||||
device=input.device,
|
||||
dtype=scale_dtype,
|
||||
)
|
||||
scale_ub = torch.mean(input).to(scale_dtype)
|
||||
args = (
|
||||
result,
|
||||
input,
|
||||
scale,
|
||||
group_size,
|
||||
scale_ub,
|
||||
False,
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
class TestSiluAndMulPerBlockQuantConfigPicker:
|
||||
def setup_method(self):
|
||||
_pick_cache.clear()
|
||||
|
||||
def test_config_picker_exact_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}
|
||||
)
|
||||
|
||||
def test_config_picker_closest_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}),
|
||||
CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(20, 3000, 70)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}
|
||||
)
|
||||
|
||||
def test_config_picker_no_configs(self):
|
||||
config_keys: list[dict] = []
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key is None
|
||||
|
||||
def test_config_picker_fallback_to_largest(self):
|
||||
config_keys = [
|
||||
CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}),
|
||||
CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(64, 8192, 256)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_config_manager_singleton():
|
||||
ConfigManager.reset_instance()
|
||||
ConfigManager()
|
||||
yield
|
||||
ConfigManager.reset_instance()
|
||||
|
||||
|
||||
class TestSiluAndMulPerBlockQuantCorrectness:
|
||||
@pytest.mark.parametrize("num_tokens", [1, 7, 4096])
|
||||
@pytest.mark.parametrize("hidden_size", [1024, 2048, 5120])
|
||||
@pytest.mark.parametrize("group_size", [64, 128])
|
||||
@pytest.mark.parametrize("is_scale_transposed", [False, True])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("quant_dtype", [current_platform.fp8_dtype(), torch.int8])
|
||||
@pytest.mark.parametrize("has_scale_ub", [True, False])
|
||||
@pytest.mark.parametrize("seed", [0])
|
||||
def test_silu_and_mul_per_block_quant(
|
||||
self,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
group_size: int,
|
||||
is_scale_transposed: bool,
|
||||
dtype: torch.dtype,
|
||||
quant_dtype: torch.dtype,
|
||||
has_scale_ub: bool,
|
||||
seed: int,
|
||||
) -> None:
|
||||
skip_if_platform_unsupported("silu_and_mul_per_block_quant")
|
||||
set_random_seed(seed)
|
||||
|
||||
if hidden_size % group_size != 0:
|
||||
return
|
||||
|
||||
if has_scale_ub and quant_dtype != FP8_DTYPE:
|
||||
# skip
|
||||
return
|
||||
|
||||
scale = 1 / hidden_size
|
||||
x = torch.randn(num_tokens, 2 * hidden_size, dtype=dtype, device="cuda") * scale
|
||||
|
||||
if has_scale_ub:
|
||||
act = torch.nn.functional.silu(x[:, :hidden_size]) * x[:, hidden_size:]
|
||||
act_abs = act.abs().float()
|
||||
scale_ub = 0.5 * (act_abs.mean() + act_abs.amax())
|
||||
else:
|
||||
scale_ub = None
|
||||
|
||||
ref_out = torch.empty(num_tokens, hidden_size, device="cuda", dtype=quant_dtype)
|
||||
|
||||
if is_scale_transposed:
|
||||
ref_scales = torch.empty(
|
||||
(hidden_size // group_size, x.shape[0]),
|
||||
device="cuda",
|
||||
dtype=torch.float32,
|
||||
).t()
|
||||
else:
|
||||
ref_scales = torch.empty(
|
||||
(x.shape[0], hidden_size // group_size),
|
||||
device="cuda",
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
ops_out = ref_out.clone()
|
||||
ops_scales = ref_scales.clone()
|
||||
|
||||
baseline(ref_out, x, ref_scales, group_size, scale_ub, is_scale_transposed)
|
||||
silu_and_mul_per_block_quant(
|
||||
ops_out, x, ops_scales, group_size, scale_ub, is_scale_transposed
|
||||
)
|
||||
|
||||
torch.testing.assert_close(ref_scales, ops_scales)
|
||||
# allow 1 ULP difference
|
||||
assert (
|
||||
ref_out.view(torch.uint8).to(torch.int16)
|
||||
- ops_out.view(torch.uint8).to(torch.int16)
|
||||
).abs().max() <= 1
|
||||
|
||||
|
||||
class TestSiluAndMulPerBlockQuantIntegration:
|
||||
def test_kernel_registration_integration(self):
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
assert "silu_and_mul_per_block_quant" in registered_kernels
|
||||
|
||||
kernel_wrapper = registered_kernels["silu_and_mul_per_block_quant"]
|
||||
assert kernel_wrapper.op_name == "silu_and_mul_per_block_quant"
|
||||
assert kernel_wrapper._config_picker is not None
|
||||
assert kernel_wrapper._mutates_args == ["out", "scales"]
|
||||
|
||||
def test_fake_impl_functionality(self):
|
||||
skip_if_platform_unsupported("silu_and_mul_per_block_quant")
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
kernel_wrapper = registered_kernels["silu_and_mul_per_block_quant"]
|
||||
fake_impl = kernel_wrapper._fake_impl
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
assert fake_impl(*args) is None
|
||||
@@ -0,0 +1,364 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm.utils.import_utils import has_helion
|
||||
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from vllm.kernels.helion.case_key import CaseKey
|
||||
from vllm.kernels.helion.config_manager import ConfigManager
|
||||
from vllm.kernels.helion.ops.silu_mul_fp8 import (
|
||||
_pick_cache,
|
||||
pick_silu_mul_fp8_config,
|
||||
silu_mul_fp8,
|
||||
silu_mul_fp8_baseline,
|
||||
)
|
||||
|
||||
|
||||
def skip_if_platform_unsupported():
|
||||
try:
|
||||
from vllm.kernels.helion.utils import get_canonical_gpu_name
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA not available")
|
||||
|
||||
platform = get_canonical_gpu_name()
|
||||
|
||||
try:
|
||||
config_manager = ConfigManager.get_instance()
|
||||
except RuntimeError:
|
||||
config_manager = ConfigManager()
|
||||
|
||||
configs = config_manager.get_platform_configs("silu_mul_fp8", platform)
|
||||
if len(configs) == 0:
|
||||
pytest.skip("Current GPU platform not supported for silu_mul_fp8 kernel")
|
||||
|
||||
except (ImportError, RuntimeError, KeyError):
|
||||
pytest.skip("Error detecting platform support for silu_mul_fp8 kernel")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_config_manager_singleton():
|
||||
ConfigManager.reset_instance()
|
||||
ConfigManager()
|
||||
yield
|
||||
ConfigManager.reset_instance()
|
||||
|
||||
|
||||
class TestSiluMulFp8ConfigPicker:
|
||||
def setup_method(self):
|
||||
_pick_cache.clear()
|
||||
|
||||
def test_config_picker_exact_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"intermediate": 2048, "numtokens": 256}),
|
||||
CaseKey({"intermediate": 4096, "numtokens": 256}),
|
||||
]
|
||||
|
||||
input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
args = (input_tensor, scale)
|
||||
|
||||
selected_key = pick_silu_mul_fp8_config(args, config_keys)
|
||||
assert selected_key == CaseKey({"intermediate": 2048, "numtokens": 256})
|
||||
|
||||
def test_config_picker_closest_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"intermediate": 2048, "numtokens": 256}),
|
||||
CaseKey({"intermediate": 4096, "numtokens": 256}),
|
||||
]
|
||||
input_tensor = torch.randn(32, 7000, dtype=torch.bfloat16, device="cuda")
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
args = (input_tensor, scale)
|
||||
|
||||
selected_key = pick_silu_mul_fp8_config(args, config_keys)
|
||||
assert selected_key == CaseKey({"intermediate": 4096, "numtokens": 256})
|
||||
|
||||
def test_config_picker_no_configs(self):
|
||||
config_keys: list[dict] = []
|
||||
|
||||
input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
args = (input_tensor, scale)
|
||||
|
||||
selected_key = pick_silu_mul_fp8_config(args, config_keys)
|
||||
assert selected_key is None
|
||||
|
||||
@pytest.mark.parametrize("intermediate_size", [2048, 4096, 5120])
|
||||
def test_config_picker_different_sizes(self, intermediate_size):
|
||||
config_keys = [
|
||||
CaseKey({"intermediate": 2048, "numtokens": 256}),
|
||||
CaseKey({"intermediate": 4096, "numtokens": 256}),
|
||||
CaseKey({"intermediate": 5120, "numtokens": 256}),
|
||||
]
|
||||
|
||||
input_tensor = torch.randn(
|
||||
32, 2 * intermediate_size, dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
args = (input_tensor, scale)
|
||||
|
||||
selected_key = pick_silu_mul_fp8_config(args, config_keys)
|
||||
assert selected_key == {
|
||||
"intermediate": intermediate_size,
|
||||
"numtokens": 256,
|
||||
}
|
||||
|
||||
def test_config_picker_numtokens_ceiling(self):
|
||||
config_keys = [
|
||||
CaseKey({"intermediate": 4096, "numtokens": 8}),
|
||||
CaseKey({"intermediate": 4096, "numtokens": 32}),
|
||||
CaseKey({"intermediate": 4096, "numtokens": 128}),
|
||||
CaseKey({"intermediate": 4096, "numtokens": 256}),
|
||||
]
|
||||
input_tensor = torch.randn(20, 8192, dtype=torch.bfloat16, device="cuda")
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
|
||||
selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
|
||||
assert selected_key == CaseKey({"intermediate": 4096, "numtokens": 32})
|
||||
|
||||
def test_config_picker_numtokens_exact(self):
|
||||
config_keys = [
|
||||
CaseKey({"intermediate": 4096, "numtokens": 8}),
|
||||
CaseKey({"intermediate": 4096, "numtokens": 32}),
|
||||
CaseKey({"intermediate": 4096, "numtokens": 128}),
|
||||
]
|
||||
input_tensor = torch.randn(32, 8192, dtype=torch.bfloat16, device="cuda")
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
|
||||
selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
|
||||
assert selected_key == CaseKey({"intermediate": 4096, "numtokens": 32})
|
||||
|
||||
def test_config_picker_numtokens_fallback_to_largest(self):
|
||||
config_keys = [
|
||||
CaseKey({"intermediate": 4096, "numtokens": 8}),
|
||||
CaseKey({"intermediate": 4096, "numtokens": 32}),
|
||||
CaseKey({"intermediate": 4096, "numtokens": 128}),
|
||||
]
|
||||
input_tensor = torch.randn(512, 8192, dtype=torch.bfloat16, device="cuda")
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
|
||||
selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
|
||||
assert selected_key == CaseKey({"intermediate": 4096, "numtokens": 128})
|
||||
|
||||
|
||||
class TestSiluMulFp8Correctness:
|
||||
@pytest.mark.parametrize("batch_size", [1, 8, 32, 128])
|
||||
@pytest.mark.parametrize("intermediate_size", [2048, 3000, 3500, 4096, 5000])
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_silu_mul_fp8_correctness(self, batch_size, intermediate_size, dtype):
|
||||
skip_if_platform_unsupported()
|
||||
|
||||
input_size = 2 * intermediate_size
|
||||
input_tensor = torch.randn(batch_size, input_size, dtype=dtype, device="cuda")
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
|
||||
reference_output = silu_mul_fp8_baseline(input_tensor, scale)
|
||||
helion_output = silu_mul_fp8(input_tensor, scale)
|
||||
|
||||
assert helion_output.shape == reference_output.shape
|
||||
assert helion_output.dtype == torch.float8_e4m3fn
|
||||
assert reference_output.dtype == torch.float8_e4m3fn
|
||||
|
||||
ref_f32 = reference_output.to(torch.float32)
|
||||
helion_f32 = helion_output.to(torch.float32)
|
||||
# FP8 E4M3 has limited precision. Values near quantization boundaries
|
||||
# can round differently due to intermediate precision differences.
|
||||
torch.testing.assert_close(
|
||||
helion_f32,
|
||||
ref_f32,
|
||||
atol=0.05,
|
||||
rtol=0.05,
|
||||
msg=f"Mismatch at batch={batch_size}, size={intermediate_size}",
|
||||
)
|
||||
|
||||
def test_silu_mul_fp8_shape_inference(self):
|
||||
skip_if_platform_unsupported()
|
||||
batch_size, input_size = 32, 8192
|
||||
intermediate_size = input_size // 2
|
||||
|
||||
input_tensor = torch.randn(
|
||||
batch_size, input_size, dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
|
||||
output = silu_mul_fp8(input_tensor, scale)
|
||||
|
||||
expected_shape = (batch_size, intermediate_size)
|
||||
assert output.shape == expected_shape
|
||||
assert output.dtype == torch.float8_e4m3fn
|
||||
|
||||
def test_silu_mul_fp8_scale_variations(self):
|
||||
skip_if_platform_unsupported()
|
||||
batch_size, input_size = 16, 4096
|
||||
|
||||
input_tensor = torch.randn(
|
||||
batch_size, input_size, dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
|
||||
scales = [0.1, 0.5, 1.0, 2.0, 10.0]
|
||||
|
||||
for scale_val in scales:
|
||||
scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda")
|
||||
|
||||
reference_output = silu_mul_fp8_baseline(input_tensor, scale)
|
||||
helion_output = silu_mul_fp8(input_tensor, scale)
|
||||
ref_f32 = reference_output.to(torch.float32)
|
||||
helion_f32 = helion_output.to(torch.float32)
|
||||
|
||||
torch.testing.assert_close(
|
||||
helion_f32,
|
||||
ref_f32,
|
||||
atol=0.05,
|
||||
rtol=0.05,
|
||||
msg=f"Mismatch for scale={scale_val}",
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[
|
||||
(1, 4096),
|
||||
(16, 4096),
|
||||
(128, 4096),
|
||||
(1024, 4096),
|
||||
(1, 8192),
|
||||
(16, 8192),
|
||||
(128, 8192),
|
||||
],
|
||||
)
|
||||
def test_silu_mul_fp8_various_shapes(self, shape):
|
||||
skip_if_platform_unsupported()
|
||||
|
||||
input_tensor = torch.randn(*shape, dtype=torch.bfloat16, device="cuda")
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
|
||||
reference_output = silu_mul_fp8_baseline(input_tensor, scale)
|
||||
helion_output = silu_mul_fp8(input_tensor, scale)
|
||||
|
||||
assert helion_output.shape == reference_output.shape
|
||||
|
||||
ref_f32 = reference_output.to(torch.float32)
|
||||
helion_f32 = helion_output.to(torch.float32)
|
||||
|
||||
torch.testing.assert_close(
|
||||
helion_f32, ref_f32, atol=0.05, rtol=0.05, msg=f"Mismatch for shape={shape}"
|
||||
)
|
||||
|
||||
|
||||
def silu_mul_fp8_pytorch(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
|
||||
"""Pure PyTorch reference using F.silu.
|
||||
|
||||
This matches vLLM's SiluAndMul.forward_native exactly:
|
||||
F.silu(x[..., :d]) * x[..., d:]
|
||||
"""
|
||||
d = input.shape[-1] // 2
|
||||
result = F.silu(input[..., :d]) * input[..., d:]
|
||||
return (result.to(torch.float32) / scale).to(torch.float8_e4m3fn)
|
||||
|
||||
|
||||
class TestSiluMulFp8PytorchReference:
|
||||
"""Tests comparing Helion kernel against pure PyTorch implementation.
|
||||
|
||||
Uses tighter tolerance since both use PyTorch's FP8 conversion
|
||||
(same rounding mode), unlike the vLLM C++ baseline which uses
|
||||
NVIDIA's hardware FP8 conversion with different rounding.
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 256])
|
||||
@pytest.mark.parametrize("intermediate_size", [1024, 2048, 4096])
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_silu_mul_fp8_vs_pytorch(self, batch_size, intermediate_size, dtype):
|
||||
skip_if_platform_unsupported()
|
||||
|
||||
input_tensor = torch.randn(
|
||||
batch_size, 2 * intermediate_size, dtype=dtype, device="cuda"
|
||||
)
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
|
||||
pytorch_output = silu_mul_fp8_pytorch(input_tensor, scale)
|
||||
helion_output = silu_mul_fp8(input_tensor, scale)
|
||||
|
||||
assert helion_output.shape == pytorch_output.shape
|
||||
assert helion_output.dtype == torch.float8_e4m3fn
|
||||
|
||||
pytorch_f32 = pytorch_output.to(torch.float32)
|
||||
helion_f32 = helion_output.to(torch.float32)
|
||||
|
||||
# Tolerance accounts for FP8 quantization boundary effects
|
||||
torch.testing.assert_close(
|
||||
helion_f32,
|
||||
pytorch_f32,
|
||||
atol=0.05,
|
||||
rtol=0.05,
|
||||
msg=(
|
||||
f"Mismatch at batch={batch_size}, size={intermediate_size}, "
|
||||
f"dtype={dtype}"
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[
|
||||
(1, 2, 4096), # 3D input
|
||||
(2, 4, 2048), # 3D input
|
||||
(1, 1, 1, 8192), # 4D input
|
||||
],
|
||||
)
|
||||
def test_silu_mul_fp8_multidim_vs_pytorch(self, shape):
|
||||
skip_if_platform_unsupported()
|
||||
|
||||
input_tensor = torch.randn(*shape, dtype=torch.bfloat16, device="cuda")
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
|
||||
pytorch_output = silu_mul_fp8_pytorch(input_tensor, scale)
|
||||
helion_output = silu_mul_fp8(input_tensor, scale)
|
||||
|
||||
assert helion_output.shape == pytorch_output.shape
|
||||
|
||||
pytorch_f32 = pytorch_output.to(torch.float32)
|
||||
helion_f32 = helion_output.to(torch.float32)
|
||||
|
||||
torch.testing.assert_close(
|
||||
helion_f32,
|
||||
pytorch_f32,
|
||||
atol=0.05,
|
||||
rtol=0.05,
|
||||
msg=f"Mismatch for shape={shape}",
|
||||
)
|
||||
|
||||
|
||||
class TestSiluMulFp8Integration:
|
||||
def test_kernel_registration_integration(self):
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
assert "silu_mul_fp8" in registered_kernels
|
||||
|
||||
kernel_wrapper = registered_kernels["silu_mul_fp8"]
|
||||
assert kernel_wrapper.op_name == "silu_mul_fp8"
|
||||
assert kernel_wrapper._config_picker is not None
|
||||
|
||||
def test_fake_impl_functionality(self):
|
||||
skip_if_platform_unsupported()
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
registered_kernels = get_registered_kernels()
|
||||
kernel_wrapper = registered_kernels["silu_mul_fp8"]
|
||||
fake_impl = kernel_wrapper._fake_impl
|
||||
|
||||
fake_output = fake_impl(input_tensor, scale)
|
||||
|
||||
expected_shape = (32, 2048)
|
||||
assert fake_output.shape == expected_shape
|
||||
assert fake_output.dtype == torch.float8_e4m3fn
|
||||
assert fake_output.device == input_tensor.device
|
||||
@@ -0,0 +1,33 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for Helion utility functions."""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.kernels.helion.utils import canonicalize_gpu_name
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"driver_reported_name,expected",
|
||||
[
|
||||
("NVIDIA H200", "nvidia_h200"),
|
||||
("NVIDIA A100-SXM4-80GB", "nvidia_a100"),
|
||||
("NVIDIA H100 80GB HBM3", "nvidia_h100"),
|
||||
("NVIDIA H100 PCIe", "nvidia_h100"),
|
||||
("NVIDIA H100 SXM5", "nvidia_h100"),
|
||||
("NVIDIA GeForce RTX 4090", "nvidia_geforce_rtx_4090"),
|
||||
("AMD Instinct MI300X", "amd_instinct_mi300x"),
|
||||
("AMD Instinct MI250X / MI250", "amd_instinct_mi250x_mi250"),
|
||||
("Tesla V100-SXM2-32GB", "tesla_v100"),
|
||||
],
|
||||
)
|
||||
def test_canonicalize_gpu_name(driver_reported_name, expected):
|
||||
"""Test GPU name canonicalization."""
|
||||
assert canonicalize_gpu_name(driver_reported_name) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid_name", ["", " ", "\t", "\n"])
|
||||
def test_canonicalize_gpu_name_rejects_empty(invalid_name):
|
||||
"""Test that empty or whitespace-only names are rejected."""
|
||||
with pytest.raises(ValueError, match="cannot be empty"):
|
||||
canonicalize_gpu_name(invalid_name)
|
||||
@@ -0,0 +1,30 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Helion Kernel test utils"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.kernels.helion.config_manager import ConfigManager
|
||||
|
||||
|
||||
def skip_if_platform_unsupported(op_name: str):
|
||||
try:
|
||||
from vllm.kernels.helion.utils import get_canonical_gpu_name
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA not available")
|
||||
|
||||
platform = get_canonical_gpu_name()
|
||||
|
||||
try:
|
||||
config_manager = ConfigManager.get_instance()
|
||||
except RuntimeError:
|
||||
config_manager = ConfigManager()
|
||||
|
||||
configs = config_manager.get_platform_configs(op_name, platform)
|
||||
if len(configs) == 0:
|
||||
pytest.skip(f"Current GPU platform not supported for {op_name} kernel")
|
||||
|
||||
except (ImportError, RuntimeError, KeyError):
|
||||
pytest.skip(f"Error detecting platform support for {op_name} kernel")
|
||||
@@ -0,0 +1,20 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Meta-tests for vLLM IR op infrastructure.
|
||||
|
||||
Ensures all registered ops have input generators defined.
|
||||
Per-op correctness tests live alongside their op definitions
|
||||
(e.g. tests/kernels/ir/test_layernorm.py).
|
||||
"""
|
||||
|
||||
import vllm.kernels # noqa: F401 — registers provider implementations
|
||||
from vllm.ir.op import IrOp
|
||||
|
||||
|
||||
def test_all_ops_have_input_generator():
|
||||
missing = [name for name, op in IrOp.registry.items() if not op.has_input_generator]
|
||||
assert not missing, (
|
||||
f"IR ops without input generators: {missing}. "
|
||||
f"Register one with @ir.ops.<name>.register_input_generator"
|
||||
)
|
||||
@@ -0,0 +1,386 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
# This registers op implementations
|
||||
import vllm.kernels # noqa: F401
|
||||
from tests.ir.ir_test_utils import (
|
||||
COMMON_HIDDEN_SIZES,
|
||||
NUM_TOKENS,
|
||||
assert_close,
|
||||
clone_args,
|
||||
supported_providers,
|
||||
)
|
||||
from tests.kernels.allclose_default import get_default_rtol
|
||||
from vllm import ir
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
rms_norm_native = ir.ops.rms_norm.impls["native"].impl_fn
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
|
||||
reason="Currently only kernels on CUDA, ROCm and XPU",
|
||||
)
|
||||
def test_rms_norm_registration():
|
||||
expected = {
|
||||
"native": True,
|
||||
"vllm_c": current_platform.is_cuda_alike(),
|
||||
"aiter": current_platform.is_rocm(),
|
||||
"oink": current_platform.has_device_capability(100)
|
||||
and hasattr(torch.ops, "oink")
|
||||
and hasattr(torch.ops.oink, "rmsnorm"),
|
||||
"xpu_kernels": current_platform.is_xpu(),
|
||||
}
|
||||
|
||||
actual = {
|
||||
provider: impl.supported for provider, impl in ir.ops.rms_norm.impls.items()
|
||||
}
|
||||
|
||||
assert actual == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
@pytest.mark.parametrize("n_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("epsilon", [1e-6, 1e-5])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
|
||||
reason="Currently only kernels on CUDA, ROCm and XPU",
|
||||
)
|
||||
class TestRMSNorm:
|
||||
@classmethod
|
||||
def setup_class(cls, **kwargs):
|
||||
torch.set_default_device(current_platform.device_type)
|
||||
|
||||
def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
|
||||
x, weight, epsilon = ir.ops.rms_norm.generate_inputs(
|
||||
num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
|
||||
)
|
||||
out = rms_norm_native(x, weight, epsilon=epsilon)
|
||||
|
||||
# Check shape, dtype, device
|
||||
assert out.shape == x.shape
|
||||
assert out.dtype == x.dtype
|
||||
assert out.device == x.device
|
||||
|
||||
# Check the scaling property of rms norm
|
||||
out2 = rms_norm_native(x * 2.0, weight, epsilon=epsilon)
|
||||
torch.testing.assert_close(out2, out, rtol=get_default_rtol(out), atol=1e-3)
|
||||
|
||||
# Mean square should be approximately 1 (ignoring epsilon and weight scaling)
|
||||
combined_norm = out.float() / weight.float()
|
||||
variance = combined_norm.pow(2).mean(dim=-1)
|
||||
# After RMS normalization, variance should be close to 1
|
||||
torch.testing.assert_close(
|
||||
variance, torch.ones_like(variance), rtol=1e-2, atol=1e-2
|
||||
)
|
||||
|
||||
# Check behavior with and without weight
|
||||
weight1 = torch.ones_like(weight)
|
||||
out3 = rms_norm_native(x, weight1, epsilon=epsilon)
|
||||
out4 = rms_norm_native(x, None, epsilon=epsilon)
|
||||
torch.testing.assert_close(out3, out4)
|
||||
|
||||
@pytest.mark.parametrize("provider", supported_providers(ir.ops.rms_norm))
|
||||
def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
|
||||
impl = ir.ops.rms_norm.impls[provider]
|
||||
x, weight, eps = ir.ops.rms_norm.generate_inputs(
|
||||
num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
|
||||
)
|
||||
args = (x, weight, eps)
|
||||
|
||||
if not impl.supports_args(*args):
|
||||
pytest.skip(f"{provider} does not support args")
|
||||
|
||||
ref_output = rms_norm_native(*clone_args(args))
|
||||
output = impl.impl_fn(*clone_args(args))
|
||||
assert_close(ir.ops.rms_norm, output, ref_output)
|
||||
|
||||
# check that dispatched call matches direct call
|
||||
with ir.ops.rms_norm.set_priority([provider, "native"]):
|
||||
out_dispatched = ir.ops.rms_norm(*args)
|
||||
out_direct = impl.impl_fn(*args)
|
||||
torch.testing.assert_close(out_dispatched, out_direct, rtol=0.0, atol=0.0)
|
||||
|
||||
# none of these support variance_size override
|
||||
assert not impl.supports_args(x, weight, eps, 4)
|
||||
assert not impl.supports_args(x, weight, eps, variance_size=4)
|
||||
|
||||
# test weight=None behavior
|
||||
out_no_weight = impl.impl_fn(x, None, eps)
|
||||
out_unit_weight = impl.impl_fn(x, torch.ones_like(weight), eps)
|
||||
assert_close(ir.ops.rms_norm, out_no_weight, out_unit_weight)
|
||||
|
||||
@pytest.mark.parametrize("provider", ["vllm_c", "aiter", "xpu_kernels", "native"])
|
||||
def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider):
|
||||
if not ir.ops.rms_norm.impls[provider].supported:
|
||||
pytest.skip(f"{provider} impl not supported on this platform")
|
||||
|
||||
args = ir.ops.rms_norm.generate_inputs(
|
||||
num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
|
||||
)
|
||||
|
||||
# When checking the torch op, we have to set priority and use dispatch
|
||||
with ir.ops.rms_norm.set_priority([provider, "native"]):
|
||||
torch.library.opcheck(torch.ops.vllm_ir.rms_norm, args)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_rocm(),
|
||||
reason="aiter is only supported on ROCm",
|
||||
)
|
||||
def test_aiter_rejects_unsupported_dtypes():
|
||||
torch.set_default_device(current_platform.device_type)
|
||||
impl = ir.ops.rms_norm.impls["aiter"]
|
||||
for dtype in [torch.float32, torch.float64]:
|
||||
args = ir.ops.rms_norm.generate_inputs(
|
||||
num_tokens=8, hidden_size=4096, dtype=dtype, epsilon=1e-5
|
||||
)
|
||||
assert not impl.supports_args(*args), f"aiter should reject dtype={dtype}"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_rocm(),
|
||||
reason="ROCm vllm_c RMSNorm needs explicit ND input handling",
|
||||
)
|
||||
def test_vllm_c_rms_norm_accepts_nd_input():
|
||||
torch.set_default_device(current_platform.device_type)
|
||||
impl = ir.ops.rms_norm.impls["vllm_c"]
|
||||
if not impl.supported:
|
||||
pytest.skip("vllm_c impl not supported on this platform")
|
||||
|
||||
base = torch.randn(3, 8, 192, dtype=torch.float16)
|
||||
x = base.split(64, dim=-1)[0].view(3, 8, 4, 16)
|
||||
assert not x.is_contiguous()
|
||||
weight = torch.randn(16, dtype=torch.float16)
|
||||
epsilon = 1e-5
|
||||
|
||||
output = impl.impl_fn(x, weight, epsilon)
|
||||
ref_output = rms_norm_native(x, weight, epsilon)
|
||||
|
||||
assert output.shape == x.shape
|
||||
assert_close(ir.ops.rms_norm, output, ref_output)
|
||||
|
||||
|
||||
fused_add_rms_norm_native = ir.ops.fused_add_rms_norm.impls["native"].impl_fn
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
|
||||
reason="Currently only kernels on CUDA, ROCm and XPU",
|
||||
)
|
||||
def test_fused_add_rms_norm_registration():
|
||||
expected = {
|
||||
"native": True,
|
||||
"vllm_c": current_platform.is_cuda_alike(),
|
||||
"aiter": current_platform.is_rocm(),
|
||||
"oink": current_platform.has_device_capability(100)
|
||||
and hasattr(torch.ops, "oink")
|
||||
and hasattr(torch.ops.oink, "fused_add_rms_norm"),
|
||||
"xpu_kernels": current_platform.is_xpu(),
|
||||
}
|
||||
|
||||
actual = {
|
||||
provider: impl.supported
|
||||
for provider, impl in ir.ops.fused_add_rms_norm.impls.items()
|
||||
}
|
||||
|
||||
assert actual == expected
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_rocm(),
|
||||
reason="ROCm vllm_c fused_add_rms_norm needs explicit ND input handling",
|
||||
)
|
||||
def test_vllm_c_fused_add_rms_norm_accepts_nd_input():
|
||||
torch.set_default_device(current_platform.device_type)
|
||||
impl = ir.ops.fused_add_rms_norm.impls["vllm_c"]
|
||||
if not impl.supported:
|
||||
pytest.skip("vllm_c impl not supported on this platform")
|
||||
|
||||
base = torch.randn(3, 8, 192, dtype=torch.float16)
|
||||
residual_base = torch.randn(3, 8, 192, dtype=torch.float16)
|
||||
x = base.split(64, dim=-1)[0].view(3, 8, 4, 16)
|
||||
x_residual = residual_base.split(64, dim=-1)[0].view(3, 8, 4, 16)
|
||||
assert not x.is_contiguous()
|
||||
assert not x_residual.is_contiguous()
|
||||
weight = torch.randn(16, dtype=torch.float16)
|
||||
epsilon = 1e-5
|
||||
|
||||
output, residual = impl.impl_fn(x.clone(), x_residual.clone(), weight, epsilon)
|
||||
ref_output, ref_residual = fused_add_rms_norm_native(x, x_residual, weight, epsilon)
|
||||
|
||||
assert output.shape == x.shape
|
||||
assert residual.shape == x_residual.shape
|
||||
assert_close(ir.ops.fused_add_rms_norm, output, ref_output)
|
||||
assert_close(ir.ops.fused_add_rms_norm, residual, ref_residual)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
@pytest.mark.parametrize("n_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("epsilon", [1e-6, 1e-5])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
|
||||
reason="Currently only kernels on CUDA, ROCm and XPU",
|
||||
)
|
||||
class TestFusedAddRMSNorm:
|
||||
@classmethod
|
||||
def setup_class(cls, **kwargs):
|
||||
torch.set_default_device(current_platform.device_type)
|
||||
|
||||
def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
|
||||
x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
|
||||
num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
|
||||
)
|
||||
out, residual_out = fused_add_rms_norm_native(x, x_residual, weight, eps)
|
||||
|
||||
# Check shape, dtype, device
|
||||
assert out.shape == x.shape
|
||||
assert out.dtype == x.dtype
|
||||
assert out.device == x.device
|
||||
assert residual_out.shape == x_residual.shape
|
||||
assert residual_out.dtype == x_residual.dtype
|
||||
assert residual_out.device == x_residual.device
|
||||
|
||||
# Check that residual_out = x + x_residual
|
||||
expected_residual = (x.float() + x_residual.float()).to(dtype)
|
||||
torch.testing.assert_close(
|
||||
residual_out, expected_residual, rtol=1e-3, atol=1e-3
|
||||
)
|
||||
|
||||
# Verify that the output is RMS normalized version of (x + x_residual)
|
||||
expected_out = rms_norm_native(expected_residual, weight, epsilon)
|
||||
assert_close(
|
||||
ir.ops.fused_add_rms_norm,
|
||||
(out, residual_out),
|
||||
(expected_out, expected_residual),
|
||||
)
|
||||
|
||||
# Check the scaling property of rms norm
|
||||
out1, _ = fused_add_rms_norm_native(
|
||||
x, torch.zeros_like(x), weight, epsilon=epsilon
|
||||
)
|
||||
out2, _ = fused_add_rms_norm_native(
|
||||
x * 2.0, torch.zeros_like(x), weight, epsilon=epsilon
|
||||
)
|
||||
torch.testing.assert_close(out2, out1, rtol=get_default_rtol(out), atol=1e-3)
|
||||
|
||||
# Check behavior with and without weight
|
||||
weight1 = torch.ones_like(weight)
|
||||
out3, _ = fused_add_rms_norm_native(x, x_residual, weight1, eps)
|
||||
out4, _ = fused_add_rms_norm_native(x, x_residual, None, eps)
|
||||
torch.testing.assert_close(out3, out4)
|
||||
|
||||
@pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm))
|
||||
def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
|
||||
impl = ir.ops.fused_add_rms_norm.impls[provider]
|
||||
x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
|
||||
num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
|
||||
)
|
||||
args = (x, x_residual, weight, eps, None)
|
||||
|
||||
if not impl.supports_args(*args):
|
||||
pytest.skip(f"{provider} does not support args")
|
||||
|
||||
ref_output, ref_residual = fused_add_rms_norm_native(*clone_args(args))
|
||||
output, residual = impl.impl_fn(*clone_args(args))
|
||||
assert_close(ir.ops.fused_add_rms_norm, output, ref_output)
|
||||
assert_close(ir.ops.fused_add_rms_norm, residual, ref_residual)
|
||||
|
||||
# check that dispatched call matches direct call
|
||||
with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
|
||||
out_dispatched, residual_dispatched = ir.ops.fused_add_rms_norm(*args[:4])
|
||||
out_direct, residual_direct = impl.impl_fn(*clone_args(args))
|
||||
torch.testing.assert_close(out_dispatched, out_direct, rtol=0.0, atol=0.0)
|
||||
torch.testing.assert_close(
|
||||
residual_dispatched, residual_direct, rtol=0.0, atol=0.0
|
||||
)
|
||||
|
||||
# none of these support variance_size override
|
||||
assert not impl.supports_args(x, x_residual, weight, epsilon, 4)
|
||||
assert not impl.supports_args(x, x_residual, weight, epsilon, variance_size=4)
|
||||
|
||||
# test weight=None behavior
|
||||
out_no_weight, residual_no_weight = impl.impl_fn(
|
||||
x.clone(), x_residual.clone(), None, epsilon
|
||||
)
|
||||
out_unit_weight, residual_unit_weight = impl.impl_fn(
|
||||
x.clone(), x_residual.clone(), torch.ones_like(weight), epsilon
|
||||
)
|
||||
assert_close(ir.ops.fused_add_rms_norm, out_no_weight, out_unit_weight)
|
||||
assert_close(
|
||||
ir.ops.fused_add_rms_norm, residual_no_weight, residual_unit_weight
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("provider", ["vllm_c"])
|
||||
def test_inplace_semantics(self, dtype, n_tokens, hidden_size, epsilon, provider):
|
||||
"""Test that inplace implementations reuse inputs,
|
||||
for maybe_inplace overload but not for default overload."""
|
||||
impl = ir.ops.fused_add_rms_norm.impls[provider]
|
||||
if not impl.supported:
|
||||
pytest.skip(f"{provider} impl not supported on this platform")
|
||||
|
||||
x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
|
||||
num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
|
||||
)
|
||||
|
||||
# Test default overload - should NOT modify inputs even with inplace impl
|
||||
x_default = x.clone()
|
||||
x_residual_default = x_residual.clone()
|
||||
x_default_ptr = x_default.data_ptr()
|
||||
x_residual_default_ptr = x_residual_default.data_ptr()
|
||||
|
||||
with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
|
||||
out_default, residual_default = ir.ops.fused_add_rms_norm(
|
||||
x_default, x_residual_default, weight, eps
|
||||
)
|
||||
|
||||
# Default should NOT be inplace (even with inplace implementation)
|
||||
assert out_default.data_ptr() != x_default_ptr
|
||||
assert residual_default.data_ptr() != x_residual_default_ptr
|
||||
torch.testing.assert_close(x, x_default, rtol=0.0, atol=0.0)
|
||||
torch.testing.assert_close(x_residual, x_residual_default, rtol=0.0, atol=0.0)
|
||||
|
||||
# Test maybe_inplace overload - should modify inputs with inplace impl
|
||||
x_inplace = x.clone()
|
||||
x_residual_inplace = x_residual.clone()
|
||||
x_inplace_ptr = x_inplace.data_ptr()
|
||||
x_residual_inplace_ptr = x_residual_inplace.data_ptr()
|
||||
|
||||
with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
|
||||
out_inplace, residual_inplace = ir.ops.fused_add_rms_norm.maybe_inplace(
|
||||
x_inplace, x_residual_inplace, weight, eps
|
||||
)
|
||||
|
||||
# maybe_inplace should be inplace
|
||||
assert out_inplace.data_ptr() == x_inplace_ptr
|
||||
assert residual_inplace.data_ptr() == x_residual_inplace_ptr
|
||||
|
||||
# Both should produce same results
|
||||
torch.testing.assert_close(out_default, out_inplace, atol=0.0, rtol=0.0)
|
||||
torch.testing.assert_close(
|
||||
residual_default, residual_inplace, atol=0.0, rtol=0.0
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm))
|
||||
def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider):
|
||||
args = ir.ops.fused_add_rms_norm.generate_inputs(
|
||||
num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
|
||||
)
|
||||
args = args + (None,) # Add variance_size parameter
|
||||
|
||||
# When checking the torch op, we have to set priority and use dispatch
|
||||
with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
|
||||
torch.library.opcheck(torch.ops.vllm_ir.fused_add_rms_norm.default, args)
|
||||
|
||||
# Only test maybe_inplace with non-inplace implementations
|
||||
# Inplace implementations return aliases of inputs which is not allowed.
|
||||
# We break this invariant, but we also convert maybe_inplace to the default
|
||||
# overload during compilation, so maybe_inplace never reaches Inductor.
|
||||
if not ir.ops.fused_add_rms_norm.impls[provider].inplace:
|
||||
torch.library.opcheck(
|
||||
torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace, args
|
||||
)
|
||||
@@ -0,0 +1,538 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import functools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not current_platform.is_cpu():
|
||||
pytest.skip("skipping CPU-only tests", allow_module_level=True)
|
||||
|
||||
set_random_seed(12345)
|
||||
|
||||
NUM_HEADS = [
|
||||
(2, 4),
|
||||
(4, 4),
|
||||
]
|
||||
HEAD_DIMS = [
|
||||
(32, 32),
|
||||
(64, 32),
|
||||
]
|
||||
CHUNK_SIZE = 64
|
||||
CONV_DIM = 128
|
||||
CONV_KERNEL = 4
|
||||
PREFILL_SEQ_LENS = [
|
||||
[1],
|
||||
[1, 2, 3],
|
||||
[CHUNK_SIZE - 1],
|
||||
[CHUNK_SIZE],
|
||||
[CHUNK_SIZE + 1],
|
||||
[CHUNK_SIZE - 1, CHUNK_SIZE, CHUNK_SIZE + 1],
|
||||
[2 * CHUNK_SIZE - 1, 2 * CHUNK_SIZE, 2 * CHUNK_SIZE + 1],
|
||||
[4 * CHUNK_SIZE + 17],
|
||||
]
|
||||
DECODE_BATCH_SIZES = [1, 3, 5]
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=128, typed=False)
|
||||
def tensor_cache(
|
||||
elem_num: int,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
tensor = torch.rand(elem_num, dtype=dtype)
|
||||
return tensor
|
||||
|
||||
|
||||
def ref_l2norm(
|
||||
x: torch.Tensor,
|
||||
dim: int = -1,
|
||||
eps: float = 1e-5,
|
||||
) -> torch.Tensor:
|
||||
inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
|
||||
return x * inv_norm
|
||||
|
||||
|
||||
def ref_gdn_gating(
|
||||
A_log: torch.Tensor,
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
dt_bias: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
softplus_x = F.softplus(a.float() + dt_bias.float(), beta=1.0, threshold=20.0)
|
||||
g = -torch.exp(A_log.float()) * softplus_x
|
||||
beta = torch.sigmoid(b.float()).to(dtype=b.dtype)
|
||||
return g, beta
|
||||
|
||||
|
||||
def ref_gated_delta_rule(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
A_log: torch.Tensor,
|
||||
dt_bias: torch.Tensor,
|
||||
initial_state: torch.Tensor,
|
||||
cu_seqlens: torch.Tensor,
|
||||
use_qk_l2norm_in_kernel: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
|
||||
out = torch.empty_like(value)
|
||||
final_state = torch.empty_like(initial_state)
|
||||
|
||||
for seq_idx in range(cu_seqlens.numel() - 1):
|
||||
begin = int(cu_seqlens[seq_idx].item())
|
||||
end = int(cu_seqlens[seq_idx + 1].item())
|
||||
q_seq = query[:, begin:end]
|
||||
k_seq = key[:, begin:end]
|
||||
v_seq = value[:, begin:end]
|
||||
g_seq = g[begin:end].unsqueeze(0)
|
||||
beta_seq = beta[begin:end].unsqueeze(0)
|
||||
initial_dtype = q_seq.dtype
|
||||
|
||||
if use_qk_l2norm_in_kernel:
|
||||
q_seq = ref_l2norm(q_seq, dim=-1)
|
||||
k_seq = ref_l2norm(k_seq, dim=-1)
|
||||
|
||||
if q_seq.shape[2] != v_seq.shape[2]:
|
||||
repeat_factor = v_seq.shape[2] // q_seq.shape[2]
|
||||
q_seq = q_seq.repeat_interleave(repeat_factor, dim=2)
|
||||
k_seq = k_seq.repeat_interleave(repeat_factor, dim=2)
|
||||
|
||||
q_seq, k_seq, v_seq, beta_seq, g_seq = [
|
||||
x.transpose(1, 2).contiguous().to(torch.float32)
|
||||
for x in (q_seq, k_seq, v_seq, beta_seq, g_seq)
|
||||
]
|
||||
|
||||
batch_size, num_heads, seq_len, head_dim = q_seq.shape
|
||||
v_head_dim = v_seq.shape[-1]
|
||||
q_seq = q_seq * (1 / (head_dim**0.5))
|
||||
out_seq = torch.empty(
|
||||
batch_size,
|
||||
num_heads,
|
||||
seq_len,
|
||||
v_head_dim,
|
||||
dtype=v_seq.dtype,
|
||||
)
|
||||
state = initial_state[seq_idx : seq_idx + 1].to(v_seq)
|
||||
|
||||
for token_idx in range(seq_len):
|
||||
q_t = q_seq[:, :, token_idx]
|
||||
k_t = k_seq[:, :, token_idx]
|
||||
v_t = v_seq[:, :, token_idx]
|
||||
g_t = g_seq[:, :, token_idx].exp().unsqueeze(-1).unsqueeze(-1)
|
||||
beta_t = beta_seq[:, :, token_idx].unsqueeze(-1)
|
||||
|
||||
state = state * g_t
|
||||
kv_mem = (state * k_t.unsqueeze(-2)).sum(dim=-1)
|
||||
delta = (v_t - kv_mem) * beta_t
|
||||
state = state + delta.unsqueeze(-1) * k_t.unsqueeze(-2)
|
||||
out_seq[:, :, token_idx] = (state * q_t.unsqueeze(-2)).sum(dim=-1)
|
||||
|
||||
out[:, begin:end] = out_seq.transpose(1, 2).contiguous().to(initial_dtype)
|
||||
final_state[seq_idx] = state.squeeze(0)
|
||||
|
||||
return out, final_state
|
||||
|
||||
|
||||
def gdn_inputs(
|
||||
num_tokens: int,
|
||||
num_heads: tuple[int, int],
|
||||
head_dims: tuple[int, int],
|
||||
) -> tuple[torch.Tensor, ...]:
|
||||
num_qk_heads, num_v_heads = num_heads
|
||||
head_dim, v_head_dim = head_dims
|
||||
q_shape = (1, num_tokens, num_qk_heads, head_dim)
|
||||
q_numel = num_tokens * num_qk_heads * head_dim
|
||||
q = tensor_cache(q_numel, torch.bfloat16).view(q_shape)
|
||||
k = tensor_cache(q_numel, torch.bfloat16).view(q_shape)
|
||||
|
||||
v_shape = (1, num_tokens, num_v_heads, v_head_dim)
|
||||
v = tensor_cache(num_tokens * num_v_heads * v_head_dim, torch.bfloat16).view(
|
||||
v_shape
|
||||
)
|
||||
|
||||
gate_shape = (num_tokens, num_v_heads)
|
||||
gate_numel = num_tokens * num_v_heads
|
||||
a = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
|
||||
b = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
|
||||
A_log = tensor_cache(num_v_heads, torch.float32)
|
||||
dt_bias = tensor_cache(num_v_heads, torch.bfloat16)
|
||||
return q, k, v, a, b, A_log, dt_bias
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 9])
|
||||
@pytest.mark.parametrize("num_v_heads", [4, 8])
|
||||
@torch.inference_mode()
|
||||
def test_fused_gdn_gating_cpu(
|
||||
num_tokens: int,
|
||||
num_v_heads: int,
|
||||
) -> None:
|
||||
gate_shape = (num_tokens, num_v_heads)
|
||||
gate_numel = num_tokens * num_v_heads
|
||||
a = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
|
||||
b = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
|
||||
A_log = tensor_cache(num_v_heads, torch.float32)
|
||||
dt_bias = tensor_cache(num_v_heads, torch.bfloat16)
|
||||
|
||||
g_ref, beta_ref = ref_gdn_gating(A_log, a, b, dt_bias)
|
||||
g, beta = ops.fused_gdn_gating_cpu(A_log, a, b, dt_bias)
|
||||
|
||||
torch.testing.assert_close(g, g_ref.unsqueeze(0), atol=1e-4, rtol=1e-4)
|
||||
torch.testing.assert_close(
|
||||
beta.float(), beta_ref.unsqueeze(0).float(), atol=5e-3, rtol=5e-3
|
||||
)
|
||||
|
||||
|
||||
# decode path
|
||||
@pytest.mark.parametrize("batch_size", DECODE_BATCH_SIZES)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_dims", HEAD_DIMS)
|
||||
@torch.inference_mode()
|
||||
def test_fused_sigmoid_gating_delta_rule_update_cpu(
|
||||
batch_size: int,
|
||||
num_heads: tuple[int, int],
|
||||
head_dims: tuple[int, int],
|
||||
) -> None:
|
||||
q, k, v, a, b, A_log, dt_bias = gdn_inputs(
|
||||
num_tokens=batch_size,
|
||||
num_heads=num_heads,
|
||||
head_dims=head_dims,
|
||||
)
|
||||
_, num_v_heads = num_heads
|
||||
head_dim, v_head_dim = head_dims
|
||||
state_indices = torch.arange(batch_size, dtype=torch.int32)
|
||||
cu_seqlens = torch.arange(batch_size + 1, dtype=torch.int32)
|
||||
state_shape = (batch_size, num_v_heads, head_dim, v_head_dim)
|
||||
state = tensor_cache(
|
||||
batch_size * num_v_heads * head_dim * v_head_dim, torch.float32
|
||||
).view(state_shape)
|
||||
state_ref = state[state_indices].transpose(-1, -2).contiguous()
|
||||
|
||||
out_ref, final_state_ref = ref_gated_delta_rule(
|
||||
query=q,
|
||||
key=k,
|
||||
value=v,
|
||||
a=a,
|
||||
b=b,
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
initial_state=state_ref,
|
||||
cu_seqlens=cu_seqlens,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
out_ref = out_ref.transpose(0, 1).contiguous()
|
||||
|
||||
state_out = state.clone()
|
||||
out = ops.fused_sigmoid_gating_delta_rule_update_cpu(
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
a=a,
|
||||
b=b,
|
||||
initial_state_source=state_out,
|
||||
initial_state_indices=state_indices,
|
||||
cu_seqlens=cu_seqlens,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
|
||||
torch.testing.assert_close(
|
||||
state_out[state_indices].transpose(-1, -2),
|
||||
final_state_ref,
|
||||
atol=1e-2,
|
||||
rtol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
# prefill path
|
||||
@pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_dims", HEAD_DIMS)
|
||||
@torch.inference_mode()
|
||||
def test_chunk_gated_delta_rule_cpu(
|
||||
seq_lens: list[int],
|
||||
num_heads: tuple[int, int],
|
||||
head_dims: tuple[int, int],
|
||||
) -> None:
|
||||
total_tokens = sum(seq_lens)
|
||||
q, k, v, a, b, A_log, dt_bias = gdn_inputs(
|
||||
num_tokens=total_tokens,
|
||||
num_heads=num_heads,
|
||||
head_dims=head_dims,
|
||||
)
|
||||
_, num_v_heads = num_heads
|
||||
head_dim, v_head_dim = head_dims
|
||||
cu_seqlens = torch.tensor(
|
||||
[0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32
|
||||
)
|
||||
initial_state_shape = (len(seq_lens), num_v_heads, head_dim, v_head_dim)
|
||||
initial_state = tensor_cache(
|
||||
len(seq_lens) * num_v_heads * head_dim * v_head_dim, torch.float32
|
||||
).view(initial_state_shape)
|
||||
initial_state_ref = initial_state.transpose(-1, -2).contiguous()
|
||||
|
||||
out_ref, final_state_ref = ref_gated_delta_rule(
|
||||
query=q,
|
||||
key=k,
|
||||
value=v,
|
||||
a=a,
|
||||
b=b,
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
initial_state=initial_state_ref,
|
||||
cu_seqlens=cu_seqlens,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
|
||||
g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
|
||||
out, final_state = ops.chunk_gated_delta_rule_cpu(
|
||||
query=q,
|
||||
key=k,
|
||||
value=v,
|
||||
g=g.unsqueeze(0),
|
||||
beta=beta.unsqueeze(0),
|
||||
initial_state=initial_state,
|
||||
output_final_state=True,
|
||||
cu_seqlens=cu_seqlens,
|
||||
head_first=False,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
|
||||
torch.testing.assert_close(
|
||||
final_state.transpose(-1, -2),
|
||||
final_state_ref,
|
||||
atol=1e-2,
|
||||
rtol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
# (total_tokens, split) pairs mimicking where chunked prefill breaks a sequence
|
||||
# across two scheduler steps: chunk-aligned and non-aligned splits.
|
||||
TWO_CALL_SPLITS = [
|
||||
(2 * CHUNK_SIZE, CHUNK_SIZE),
|
||||
(2 * CHUNK_SIZE + 17, CHUNK_SIZE),
|
||||
(2 * CHUNK_SIZE + 17, CHUNK_SIZE + 9),
|
||||
(4 * CHUNK_SIZE + 17, 2 * CHUNK_SIZE),
|
||||
(3 * CHUNK_SIZE, CHUNK_SIZE + 1),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_dims", HEAD_DIMS)
|
||||
@torch.inference_mode()
|
||||
def test_chunk_gated_delta_rule_cpu_two_call_split(
|
||||
total_tokens: int,
|
||||
split: int,
|
||||
num_heads: tuple[int, int],
|
||||
head_dims: tuple[int, int],
|
||||
) -> None:
|
||||
"""A prefill split into two calls (the second seeded with the first's
|
||||
``final_state`` and a rebased ``cu_seqlens``) must match the single-call
|
||||
result, mimicking the cross-scheduler-step handoff in
|
||||
``cpu_gdn_attention_core``.
|
||||
"""
|
||||
q, k, v, a, b, A_log, dt_bias = gdn_inputs(
|
||||
num_tokens=total_tokens,
|
||||
num_heads=num_heads,
|
||||
head_dims=head_dims,
|
||||
)
|
||||
_, num_v_heads = num_heads
|
||||
head_dim, v_head_dim = head_dims
|
||||
|
||||
g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
|
||||
g = g.unsqueeze(0) # [1, T, HV]
|
||||
beta = beta.unsqueeze(0)
|
||||
|
||||
zero_state = torch.zeros(1, num_v_heads, head_dim, v_head_dim, dtype=torch.float32)
|
||||
|
||||
# Reference: whole sequence in one call, no initial state.
|
||||
out_full, final_full = ops.chunk_gated_delta_rule_cpu(
|
||||
query=q,
|
||||
key=k,
|
||||
value=v,
|
||||
g=g,
|
||||
beta=beta,
|
||||
initial_state=zero_state,
|
||||
output_final_state=True,
|
||||
cu_seqlens=torch.tensor([0, total_tokens], dtype=torch.int32),
|
||||
head_first=False,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
|
||||
# Call 1: tokens [0:split], no initial state, capture final state.
|
||||
out1, state1 = ops.chunk_gated_delta_rule_cpu(
|
||||
query=q[:, :split],
|
||||
key=k[:, :split],
|
||||
value=v[:, :split],
|
||||
g=g[:, :split],
|
||||
beta=beta[:, :split],
|
||||
initial_state=zero_state,
|
||||
output_final_state=True,
|
||||
cu_seqlens=torch.tensor([0, split], dtype=torch.int32),
|
||||
head_first=False,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
# Call 2: tokens [split:T] seeded with call 1's final state and a cu_seqlens
|
||||
# rebased to start at 0, as cpu_gdn_attention_core continues a prefill chunk.
|
||||
tail = total_tokens - split
|
||||
out2, state2 = ops.chunk_gated_delta_rule_cpu(
|
||||
query=q[:, split:],
|
||||
key=k[:, split:],
|
||||
value=v[:, split:],
|
||||
g=g[:, split:],
|
||||
beta=beta[:, split:],
|
||||
initial_state=state1.to(torch.float32),
|
||||
output_final_state=True,
|
||||
cu_seqlens=torch.tensor([0, tail], dtype=torch.int32),
|
||||
head_first=False,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
|
||||
out_split = torch.cat([out1, out2], dim=1)
|
||||
|
||||
# State must be near-exact; output allows a looser bound for the bf16 round-trip.
|
||||
torch.testing.assert_close(state2, final_full, atol=1e-3, rtol=1e-3)
|
||||
torch.testing.assert_close(out_split, out_full, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
def _conv_inputs(total_tokens: int):
|
||||
x = tensor_cache(total_tokens * CONV_DIM, torch.bfloat16).view(
|
||||
total_tokens, CONV_DIM
|
||||
)
|
||||
weight = tensor_cache(CONV_DIM * CONV_KERNEL, torch.bfloat16).view(
|
||||
CONV_DIM, CONV_KERNEL
|
||||
)
|
||||
bias = tensor_cache(CONV_DIM, torch.bfloat16)
|
||||
return x, weight, bias
|
||||
|
||||
|
||||
@pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS)
|
||||
@torch.inference_mode()
|
||||
def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> None:
|
||||
"""Non-AMX conv-state handoff: a two-call split (the second seeded via
|
||||
``has_initial_state=True`` from the conv_states the first wrote back) must
|
||||
match the single-call result.
|
||||
"""
|
||||
from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
|
||||
causal_conv1d_torch,
|
||||
)
|
||||
|
||||
x, weight, bias = _conv_inputs(total_tokens)
|
||||
state_len = CONV_KERNEL - 1
|
||||
# [num_slots, conv_dim, state_len]; slot 0 used here.
|
||||
conv_states_full = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype)
|
||||
conv_states_split = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype)
|
||||
|
||||
# x is [conv_dim, T] for causal_conv1d_torch.
|
||||
xt = x.transpose(0, 1).contiguous()
|
||||
|
||||
out_full = causal_conv1d_torch(
|
||||
x=xt,
|
||||
weight=weight,
|
||||
bias=bias,
|
||||
conv_states=conv_states_full,
|
||||
query_start_loc=torch.tensor([0, total_tokens], dtype=torch.int32),
|
||||
cache_indices=torch.tensor([0], dtype=torch.int32),
|
||||
has_initial_state=torch.tensor([False]),
|
||||
activation="silu",
|
||||
)
|
||||
|
||||
out1 = causal_conv1d_torch(
|
||||
x=xt[:, :split],
|
||||
weight=weight,
|
||||
bias=bias,
|
||||
conv_states=conv_states_split,
|
||||
query_start_loc=torch.tensor([0, split], dtype=torch.int32),
|
||||
cache_indices=torch.tensor([0], dtype=torch.int32),
|
||||
has_initial_state=torch.tensor([False]),
|
||||
activation="silu",
|
||||
)
|
||||
out2 = causal_conv1d_torch(
|
||||
x=xt[:, split:],
|
||||
weight=weight,
|
||||
bias=bias,
|
||||
conv_states=conv_states_split,
|
||||
query_start_loc=torch.tensor([0, total_tokens - split], dtype=torch.int32),
|
||||
cache_indices=torch.tensor([0], dtype=torch.int32),
|
||||
has_initial_state=torch.tensor([True]),
|
||||
activation="silu",
|
||||
)
|
||||
out_split = torch.cat([out1, out2], dim=1)
|
||||
|
||||
torch.testing.assert_close(out_split, out_full, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not torch.cpu._is_amx_tile_supported(),
|
||||
reason="causal_conv1d_fwd_cpu requires AMX/AVX512",
|
||||
)
|
||||
@pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS)
|
||||
@torch.inference_mode()
|
||||
def test_causal_conv1d_fwd_cpu_two_call_split(total_tokens: int, split: int) -> None:
|
||||
"""AMX prefill conv op must honor ``has_initial_state`` so a two-call split
|
||||
matches the single-call result.
|
||||
|
||||
Regression test for ``causal_conv1d_fwd_varlen_kernel_impl`` (``conv.cpp``)
|
||||
ignoring the carried conv state on continued chunks.
|
||||
"""
|
||||
state_len = CONV_KERNEL - 1
|
||||
x, weight, bias = _conv_inputs(total_tokens)
|
||||
|
||||
def amx(x_seg, conv_states, has_init):
|
||||
seq = x_seg.shape[0]
|
||||
return ops.causal_conv1d_fwd_cpu(
|
||||
x=x_seg.transpose(0, 1), # [dim, seq]; stride(-2)==1 (view of [seq,dim])
|
||||
weight=weight,
|
||||
bias=bias,
|
||||
conv_states=conv_states,
|
||||
query_start_loc=torch.tensor([0, seq], dtype=torch.int32),
|
||||
cache_indices=torch.tensor([0], dtype=torch.int32),
|
||||
has_initial_state=torch.tensor([has_init]),
|
||||
silu_activation=True,
|
||||
is_vnni=False,
|
||||
).contiguous()
|
||||
|
||||
# conv_state layout passed by the AMX branch: [num_slots, dim, state_len].
|
||||
cs_full = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype)
|
||||
out_full = amx(x, cs_full, False)
|
||||
|
||||
cs_split = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype)
|
||||
out1 = amx(x[:split], cs_split, False)
|
||||
out2 = amx(x[split:], cs_split, True)
|
||||
out_split = torch.cat([out1, out2], dim=1)
|
||||
|
||||
torch.testing.assert_close(out_split, out_full, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_batch_memcpy_cpu_fallback() -> None:
|
||||
"""The ctypes batch_memcpy fallback (used when triton-cpu is absent) must
|
||||
copy each src into its dst, validating the (src_ptrs, dst_ptrs, sizes)
|
||||
argument order against ctypes.memmove(dst, src, size).
|
||||
"""
|
||||
from vllm.utils.cpu_triton_utils import batch_memcpy_kernel
|
||||
|
||||
# Varied byte sizes, including a non-power-of-two run.
|
||||
sizes_bytes = [256, 1024, 17 * 4, 4096]
|
||||
srcs = [torch.rand(n // 4, dtype=torch.float32) for n in sizes_bytes]
|
||||
dsts = [torch.zeros_like(s) for s in srcs]
|
||||
|
||||
src_ptrs = torch.tensor([s.data_ptr() for s in srcs], dtype=torch.uint64)
|
||||
dst_ptrs = torch.tensor([d.data_ptr() for d in dsts], dtype=torch.uint64)
|
||||
sizes = torch.tensor(sizes_bytes, dtype=torch.int32)
|
||||
|
||||
batch_memcpy_kernel[(len(srcs),)](src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=1024)
|
||||
|
||||
for src, dst in zip(srcs, dsts):
|
||||
torch.testing.assert_close(dst, src)
|
||||
@@ -0,0 +1,389 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
|
||||
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
|
||||
causal_conv1d_fn,
|
||||
causal_conv1d_update,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
|
||||
|
||||
DEVICE = current_platform.device_type
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
|
||||
reason="causal_conv1d Triton kernels require CUDA-alike or XPU",
|
||||
)
|
||||
|
||||
|
||||
def causal_conv1d_ref(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
initial_states: torch.Tensor | None = None,
|
||||
return_final_states: bool = False,
|
||||
final_states_out: torch.Tensor | None = None,
|
||||
activation: str | None = "silu",
|
||||
):
|
||||
"""
|
||||
x: (batch, dim, seqlen)
|
||||
weight: (dim, width)
|
||||
bias: (dim,)
|
||||
initial_states: (batch, dim, width - 1)
|
||||
final_states_out: (batch, dim, width - 1)
|
||||
|
||||
out: (batch, dim, seqlen)
|
||||
"""
|
||||
if activation not in [None, "silu", "swish"]:
|
||||
raise NotImplementedError("activation must be None, silu, or swish")
|
||||
dtype_in = x.dtype
|
||||
x = x.to(weight.dtype)
|
||||
seqlen = x.shape[-1]
|
||||
dim, width = weight.shape
|
||||
if initial_states is None:
|
||||
out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim)
|
||||
else:
|
||||
x = torch.cat([initial_states, x], dim=-1)
|
||||
out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim)
|
||||
out = out[..., :seqlen]
|
||||
if return_final_states:
|
||||
final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to(
|
||||
dtype_in
|
||||
) # (batch, dim, width - 1)
|
||||
if final_states_out is not None:
|
||||
final_states_out.copy_(final_states)
|
||||
else:
|
||||
final_states_out = final_states
|
||||
out = (out if activation is None else F.silu(out)).to(dtype=dtype_in)
|
||||
return (out, None) if not return_final_states else (out, final_states_out)
|
||||
|
||||
|
||||
def causal_conv1d_update_ref(
|
||||
x, conv_state, weight, bias=None, activation=None, cache_seqlens=None
|
||||
):
|
||||
"""
|
||||
x: (batch, dim) or (batch, dim, seqlen)
|
||||
conv_state: (batch, dim, state_len), where state_len >= width - 1
|
||||
weight: (dim, width)
|
||||
bias: (dim,)
|
||||
cache_seqlens: (batch,), dtype int32.
|
||||
If not None, the conv_state is treated as a circular buffer.
|
||||
The conv_state will be updated by copying x to the
|
||||
conv_state starting at the index
|
||||
@cache_seqlens % state_len before performing the convolution.
|
||||
|
||||
out: (batch, dim) or (batch, dim, seqlen)
|
||||
"""
|
||||
if activation not in [None, "silu", "swish"]:
|
||||
raise NotImplementedError("activation must be None, silu, or swish")
|
||||
dtype_in = x.dtype
|
||||
unsqueeze = x.dim() == 2
|
||||
if unsqueeze:
|
||||
x = x.unsqueeze(-1)
|
||||
batch, dim, seqlen = x.shape
|
||||
width = weight.shape[1]
|
||||
state_len = conv_state.shape[-1]
|
||||
assert conv_state.shape == (batch, dim, state_len)
|
||||
assert weight.shape == (dim, width)
|
||||
if cache_seqlens is None:
|
||||
x_new = torch.cat([conv_state, x], dim=-1).to(
|
||||
weight.dtype
|
||||
) # (batch, dim, state_len + seqlen)
|
||||
conv_state.copy_(x_new[:, :, -state_len:])
|
||||
else:
|
||||
width_idx = torch.arange(
|
||||
-(width - 1), 0, dtype=torch.long, device=x.device
|
||||
).unsqueeze(0) + cache_seqlens.unsqueeze(1)
|
||||
width_idx = (
|
||||
torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1)
|
||||
)
|
||||
x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype)
|
||||
copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(
|
||||
0
|
||||
) + cache_seqlens.unsqueeze(1)
|
||||
copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1)
|
||||
conv_state.scatter_(2, copy_idx, x)
|
||||
out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[
|
||||
:, :, -seqlen:
|
||||
]
|
||||
if unsqueeze:
|
||||
out = out.squeeze(-1)
|
||||
return (out if activation is None else F.silu(out)).to(dtype=dtype_in)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("itype", [torch.bfloat16, torch.float])
|
||||
@pytest.mark.parametrize("silu_activation", [True])
|
||||
@pytest.mark.parametrize("has_bias", [True])
|
||||
def causal_conv1d_opcheck_fn(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
cu_seq_len: torch.Tensor | None = None,
|
||||
cache_indices: torch.Tensor | None = None,
|
||||
has_initial_state: torch.Tensor | None = None,
|
||||
conv_states: torch.Tensor | None = None,
|
||||
activation: str | None = "silu",
|
||||
null_block_id: int = NULL_BLOCK_ID,
|
||||
):
|
||||
"""
|
||||
x: (batch, dim, seqlen)
|
||||
weight: (dim, width)
|
||||
bias: (dim,)
|
||||
seq_idx: (batch, seqlen)
|
||||
initial_states: (batch, dim, width - 1)
|
||||
final_states_out: (batch, dim, width - 1), to be written to
|
||||
activation: either None or "silu" or "swish"
|
||||
|
||||
out: (batch, dim, seqlen)
|
||||
"""
|
||||
if activation not in [None, "silu", "swish"]:
|
||||
raise NotImplementedError("activation must be None, silu, or swish")
|
||||
if x.stride(-1) != 1:
|
||||
x = x.contiguous()
|
||||
bias = bias.contiguous() if bias is not None else None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("itype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("silu_activation", [False, True])
|
||||
@pytest.mark.parametrize("has_bias", [False, True])
|
||||
@pytest.mark.parametrize("seqlen", [1])
|
||||
@pytest.mark.parametrize("width", [4])
|
||||
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
|
||||
def test_causal_conv1d_update(dim, width, seqlen, has_bias, silu_activation, itype):
|
||||
device = DEVICE
|
||||
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
|
||||
if itype == torch.bfloat16:
|
||||
rtol, atol = 1e-2, 5e-2
|
||||
# set seed
|
||||
set_random_seed(0)
|
||||
batch = 2
|
||||
x = torch.randn(batch, dim, seqlen, device=device, dtype=itype)
|
||||
x_ref = x.clone()
|
||||
# +1 entry to reserve index 0 as null block
|
||||
conv_state = torch.randn(batch + 1, dim, width - 1, device=device, dtype=itype)
|
||||
|
||||
weight = torch.randn(dim, width, device=device, dtype=itype)
|
||||
bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None
|
||||
# Start indices from 1, skipping null block at index 0
|
||||
conv_state_indices = torch.arange(1, batch + 1, dtype=torch.int32, device=device)
|
||||
conv_state_ref = conv_state[conv_state_indices].detach().clone()
|
||||
activation = None if not silu_activation else "silu"
|
||||
|
||||
out = causal_conv1d_update(
|
||||
x,
|
||||
conv_state,
|
||||
weight,
|
||||
bias,
|
||||
activation=activation,
|
||||
conv_state_indices=conv_state_indices,
|
||||
)
|
||||
out_ref = causal_conv1d_update_ref(
|
||||
x_ref, conv_state_ref, weight, bias, activation=activation
|
||||
)
|
||||
|
||||
assert torch.equal(conv_state[conv_state_indices], conv_state_ref)
|
||||
assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])
|
||||
@pytest.mark.parametrize("silu_activation", [False, True])
|
||||
@pytest.mark.parametrize("has_bias", [False, True])
|
||||
@pytest.mark.parametrize("seqlen", [1, 3])
|
||||
@pytest.mark.parametrize("width", [3, 4])
|
||||
@pytest.mark.parametrize("dim", [2048 + 16, 4096])
|
||||
# tests correctness in case subset of the sequences are padded
|
||||
@pytest.mark.parametrize("with_padding", [True, False])
|
||||
@pytest.mark.parametrize("batch_size", [3])
|
||||
def test_causal_conv1d_update_with_batch_gather(
|
||||
batch_size, with_padding, dim, width, seqlen, has_bias, silu_activation, itype
|
||||
):
|
||||
device = DEVICE
|
||||
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
|
||||
if itype == torch.bfloat16:
|
||||
rtol, atol = 1e-2, 5e-2
|
||||
|
||||
# set seed
|
||||
set_random_seed(0)
|
||||
|
||||
padding = 5 if with_padding else 0
|
||||
padded_batch_size = batch_size + padding
|
||||
# total_entries = number of cache line
|
||||
total_entries = 10 * batch_size
|
||||
|
||||
# x will be (batch, dim, seqlen) with contiguous along dim-axis
|
||||
x = torch.randn(
|
||||
padded_batch_size, seqlen, dim, device=device, dtype=itype
|
||||
).transpose(1, 2)
|
||||
|
||||
x_ref = x.clone()
|
||||
|
||||
# +1 to exclude index 0 (null block)
|
||||
conv_state_indices = (torch.randperm(total_entries - 1)[:batch_size] + 1).to(
|
||||
dtype=torch.int32, device=device
|
||||
)
|
||||
unused_states_bool = torch.ones(total_entries, dtype=torch.bool, device=device)
|
||||
unused_states_bool[conv_state_indices] = False
|
||||
padded_state_indices = torch.concat(
|
||||
[
|
||||
conv_state_indices,
|
||||
torch.as_tensor(
|
||||
[NULL_BLOCK_ID] * padding, dtype=torch.int32, device=device
|
||||
),
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
|
||||
# conv_state will be (cache_lines, dim, state_len)
|
||||
# with contiguous along dim-axis
|
||||
conv_state = torch.randn(
|
||||
total_entries, width - 1, dim, device=device, dtype=itype
|
||||
).transpose(1, 2)
|
||||
|
||||
conv_state_for_padding_test = conv_state.clone()
|
||||
|
||||
weight = torch.randn(dim, width, device=device, dtype=itype)
|
||||
bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None
|
||||
conv_state_ref = conv_state[conv_state_indices, :].detach().clone()
|
||||
activation = None if not silu_activation else "silu"
|
||||
|
||||
out = causal_conv1d_update(
|
||||
x,
|
||||
conv_state,
|
||||
weight,
|
||||
bias,
|
||||
activation=activation,
|
||||
conv_state_indices=padded_state_indices,
|
||||
)
|
||||
out_ref = causal_conv1d_update_ref(
|
||||
x_ref[:batch_size], conv_state_ref, weight, bias, activation=activation
|
||||
)
|
||||
|
||||
assert torch.equal(conv_state[conv_state_indices, :], conv_state_ref)
|
||||
assert torch.equal(
|
||||
conv_state[unused_states_bool], conv_state_for_padding_test[unused_states_bool]
|
||||
)
|
||||
assert torch.allclose(out[:batch_size], out_ref, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("itype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("silu_activation", [True])
|
||||
@pytest.mark.parametrize("has_bias", [True])
|
||||
@pytest.mark.parametrize("width", [4])
|
||||
@pytest.mark.parametrize("seqlen", [8, 249, 4096])
|
||||
@pytest.mark.parametrize("dim", [64, 4096])
|
||||
@pytest.mark.parametrize("with_padding", [True, False])
|
||||
@pytest.mark.parametrize("batch", [4, 10])
|
||||
def test_causal_conv1d_varlen(
|
||||
batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype
|
||||
):
|
||||
device = DEVICE
|
||||
torch.accelerator.empty_cache()
|
||||
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
|
||||
if itype == torch.bfloat16:
|
||||
rtol, atol = 1e-2, 5e-2
|
||||
# set seed
|
||||
set_random_seed(0)
|
||||
seqlens = []
|
||||
batch_size = batch
|
||||
padding = 3 if with_padding else 0
|
||||
padded_batch_size = batch_size + padding
|
||||
nsplits = padded_batch_size - 1
|
||||
|
||||
eos_pos = torch.randperm(seqlen - 1)[:nsplits].sort().values
|
||||
|
||||
seqlens.append(
|
||||
torch.diff(
|
||||
torch.cat([torch.tensor([-1]), eos_pos, torch.tensor([seqlen - 1])])
|
||||
).tolist()
|
||||
)
|
||||
assert sum(seqlens[-1]) == seqlen
|
||||
assert all(s > 0 for s in seqlens[-1])
|
||||
|
||||
total_entries = batch_size * 10
|
||||
cumsum = torch.cumsum(torch.tensor(seqlens[0]), dim=0).to(torch.int32)
|
||||
cumsum = torch.concat([torch.tensor([0], dtype=torch.int32), cumsum], dim=0)
|
||||
x = rearrange(
|
||||
torch.randn(1, seqlen, 4096 + dim + 64, device=device, dtype=itype),
|
||||
"b s d -> b d s",
|
||||
)[:, 4096 : 4096 + dim, :]
|
||||
|
||||
weight = torch.randn(dim, width, device=device, dtype=itype)
|
||||
|
||||
bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None
|
||||
x_ref = x.clone()
|
||||
weight_ref = weight.clone()
|
||||
bias_ref = bias.clone() if bias is not None else None
|
||||
activation = None if not silu_activation else "silu"
|
||||
final_states = torch.randn(
|
||||
total_entries, width - 1, dim, device=x.device, dtype=x.dtype
|
||||
).transpose(1, 2)
|
||||
final_states_ref = final_states.clone()
|
||||
has_initial_states = torch.randint(
|
||||
0, 2, (cumsum.shape[0] - 1,), dtype=torch.bool, device=x.device
|
||||
)
|
||||
# +1 to exclude index 0 (null block)
|
||||
state_indices = (
|
||||
torch.randperm(total_entries - 1, dtype=torch.int32, device=x.device)[
|
||||
:batch_size
|
||||
]
|
||||
+ 1
|
||||
)
|
||||
padded_state_indices = torch.concat(
|
||||
[
|
||||
state_indices,
|
||||
torch.as_tensor(
|
||||
[NULL_BLOCK_ID] * padding, dtype=torch.int32, device=device
|
||||
),
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
out = causal_conv1d_fn(
|
||||
x.squeeze(0),
|
||||
weight,
|
||||
bias=bias,
|
||||
conv_states=final_states,
|
||||
query_start_loc=cumsum.to(device),
|
||||
cache_indices=padded_state_indices,
|
||||
has_initial_state=has_initial_states,
|
||||
activation=activation,
|
||||
)
|
||||
|
||||
out_ref = []
|
||||
out_ref_b = []
|
||||
|
||||
splits = [torch.split(var, seqlens[0], dim=-1) for var in (x_ref)]
|
||||
for i in range(len(seqlens[0])):
|
||||
x_s = [v[i].unsqueeze(0) for v in splits][0]
|
||||
if padded_state_indices[i] == NULL_BLOCK_ID:
|
||||
continue
|
||||
out_ref_b.append(
|
||||
causal_conv1d_ref(
|
||||
x_s,
|
||||
weight_ref,
|
||||
bias_ref,
|
||||
activation=activation,
|
||||
return_final_states=True,
|
||||
final_states_out=final_states_ref[padded_state_indices[i]].unsqueeze(0),
|
||||
initial_states=final_states_ref[padded_state_indices[i]].unsqueeze(0)
|
||||
if has_initial_states[i]
|
||||
else None,
|
||||
)
|
||||
)
|
||||
out_ref.append(torch.cat([t[0] for t in out_ref_b], dim=2))
|
||||
out_ref_tensor = torch.cat(out_ref, dim=0)
|
||||
|
||||
assert torch.allclose(
|
||||
final_states[state_indices],
|
||||
final_states_ref[state_indices],
|
||||
rtol=rtol,
|
||||
atol=atol,
|
||||
)
|
||||
unpadded_out = out[:, : out_ref_tensor.shape[-1]]
|
||||
assert torch.allclose(unpadded_out, out_ref_tensor, rtol=rtol, atol=atol)
|
||||
@@ -0,0 +1,189 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config import CompilationConfig, VllmConfig
|
||||
from vllm.forward_context import set_forward_context
|
||||
from vllm.model_executor.layers.mamba.short_conv import ShortConv
|
||||
from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadata
|
||||
|
||||
if not current_platform.is_cpu():
|
||||
pytest.skip("skipping CPU-only tests", allow_module_level=True)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_dist():
|
||||
with (
|
||||
patch(
|
||||
"vllm.model_executor.layers.linear.get_tensor_model_parallel_rank",
|
||||
return_value=0,
|
||||
),
|
||||
patch(
|
||||
"vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
|
||||
return_value=1,
|
||||
),
|
||||
patch(
|
||||
"vllm.distributed.parallel_state.model_parallel_is_initialized",
|
||||
return_value=True,
|
||||
),
|
||||
patch(
|
||||
"vllm.distributed.parallel_state.get_tp_group",
|
||||
return_value=MagicMock(rank_in_group=0),
|
||||
),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def vllm_config():
|
||||
# ShortConv only needs compilation_config from the current vLLM config, so a
|
||||
# minimal config (model_config=None) avoids mocking ModelConfig and the
|
||||
# associated VllmConfig validation churn.
|
||||
return VllmConfig(compilation_config=CompilationConfig())
|
||||
|
||||
|
||||
def test_short_conv_forward_native_prefill(vllm_config):
|
||||
prefix = "test_layer"
|
||||
config = SimpleNamespace(conv_L_cache=4, conv_bias=True)
|
||||
dim = 16
|
||||
|
||||
from vllm.config import set_current_vllm_config
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix)
|
||||
|
||||
layer.to("cpu")
|
||||
# vLLM Linear layers allocate weights with torch.empty (uninitialized).
|
||||
# On ARM these come back as zero-filled pages, so in_proj output is zero and
|
||||
# the prefill state stays zero. Seed + init to make the test platform-safe.
|
||||
torch.manual_seed(0)
|
||||
for p in layer.parameters():
|
||||
torch.nn.init.normal_(p)
|
||||
dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False)
|
||||
dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False)
|
||||
|
||||
# Mock AttentionMetadata
|
||||
num_prefills = 1
|
||||
num_prefill_tokens = 5
|
||||
query_start_loc_p = torch.tensor([0, 5], dtype=torch.int32)
|
||||
state_indices_tensor_p = torch.tensor([0], dtype=torch.int32)
|
||||
|
||||
# ShortConvAttentionMetadata
|
||||
attn_metadata = ShortConvAttentionMetadata(
|
||||
num_prefills=num_prefills,
|
||||
num_prefill_tokens=num_prefill_tokens,
|
||||
num_decodes=0,
|
||||
num_decode_tokens=0,
|
||||
num_reqs=1,
|
||||
query_start_loc_p=query_start_loc_p,
|
||||
has_initial_states_p=torch.tensor([False]),
|
||||
state_indices_tensor_p=state_indices_tensor_p,
|
||||
state_indices_tensor_d=torch.empty((0, 1), dtype=torch.int32),
|
||||
num_accepted_tokens=None,
|
||||
query_start_loc_d=None,
|
||||
block_idx_last_scheduled_token=None,
|
||||
block_idx_first_scheduled_token_p=None,
|
||||
block_idx_last_computed_token=None,
|
||||
block_idx_last_scheduled_token_prev_step=None,
|
||||
num_computed_tokens_p=None,
|
||||
seq_lens=torch.tensor([5]),
|
||||
)
|
||||
|
||||
# Mock KV cache
|
||||
# conv_state shape (num_blocks, L_cache - 1, dim)
|
||||
conv_state = torch.zeros((1, config.conv_L_cache - 1, dim))
|
||||
layer.kv_cache = (conv_state,)
|
||||
|
||||
hidden_states = torch.randn((num_prefill_tokens, dim))
|
||||
output = torch.zeros_like(hidden_states)
|
||||
|
||||
attn_metadata_dict = {prefix: attn_metadata}
|
||||
with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config):
|
||||
layer.forward_native(hidden_states, output)
|
||||
|
||||
# Check if KV cache was updated
|
||||
assert not torch.allclose(conv_state, torch.zeros_like(conv_state))
|
||||
|
||||
|
||||
def test_short_conv_forward_native_decode(vllm_config):
|
||||
prefix = "test_layer_decode"
|
||||
config = SimpleNamespace(conv_L_cache=4, conv_bias=True)
|
||||
dim = 16
|
||||
|
||||
from vllm.config import set_current_vllm_config
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix)
|
||||
|
||||
layer.to("cpu")
|
||||
torch.manual_seed(0)
|
||||
for p in layer.parameters():
|
||||
torch.nn.init.normal_(p)
|
||||
dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False)
|
||||
dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False)
|
||||
|
||||
# Mock AttentionMetadata for 2 decode requests
|
||||
num_decodes = 2
|
||||
state_indices_tensor_d = torch.tensor([0, 1], dtype=torch.int32)
|
||||
|
||||
attn_metadata = ShortConvAttentionMetadata(
|
||||
num_prefills=0,
|
||||
num_prefill_tokens=0,
|
||||
num_decodes=num_decodes,
|
||||
num_decode_tokens=num_decodes,
|
||||
num_reqs=num_decodes,
|
||||
query_start_loc_p=None,
|
||||
has_initial_states_p=None,
|
||||
state_indices_tensor_p=torch.empty((0,), dtype=torch.int32),
|
||||
state_indices_tensor_d=state_indices_tensor_d,
|
||||
num_accepted_tokens=None,
|
||||
query_start_loc_d=torch.tensor([0, 1, 2], dtype=torch.int32),
|
||||
block_idx_last_scheduled_token=None,
|
||||
block_idx_first_scheduled_token_p=None,
|
||||
block_idx_last_computed_token=None,
|
||||
block_idx_last_scheduled_token_prev_step=None,
|
||||
num_computed_tokens_p=None,
|
||||
seq_lens=torch.tensor([1, 1]),
|
||||
)
|
||||
|
||||
# Mock KV cache (2 blocks for 2 requests)
|
||||
conv_state = torch.randn((2, config.conv_L_cache - 1, dim))
|
||||
layer.kv_cache = (conv_state,)
|
||||
|
||||
hidden_states = torch.randn((num_decodes, dim))
|
||||
output = torch.zeros_like(hidden_states)
|
||||
|
||||
old_conv_state = conv_state.clone()
|
||||
|
||||
attn_metadata_dict = {prefix: attn_metadata}
|
||||
with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config):
|
||||
layer.forward_native(hidden_states, output)
|
||||
|
||||
# Check if KV cache was updated
|
||||
assert not torch.allclose(conv_state, old_conv_state)
|
||||
|
||||
|
||||
def test_dispatch_cpu_unquantized_gemm_conv_layer():
|
||||
# Convolution layers have >2D weights; dispatch should skip them gracefully.
|
||||
# Shape/dtype are AMX-pack safe (bf16, width==4, dim % block_size == 0) so
|
||||
# the AMX prepack branch does not raise on AMX-capable CPUs.
|
||||
class MockConvLayer(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.weight = torch.nn.Parameter(
|
||||
torch.randn(32, 1, 4, dtype=torch.bfloat16)
|
||||
)
|
||||
self.bias = torch.nn.Parameter(torch.randn(32, dtype=torch.bfloat16))
|
||||
|
||||
layer = MockConvLayer()
|
||||
# The ndim != 2 guard returns early without raising.
|
||||
dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
|
||||
# No cpu_linear set — conv layers are handled elsewhere.
|
||||
assert not hasattr(layer, "cpu_linear")
|
||||
@@ -0,0 +1,296 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Integration test for the non-spec decode split in
|
||||
``GatedDeltaNet._forward_core``.
|
||||
|
||||
On a pure non-spec batch that mixes prefills with 1-token decodes, the layer
|
||||
peels the decodes (the contiguous decode-first front slice) off to
|
||||
``fused_sigmoid_gating_delta_rule_update`` -- the same recurrent update kernel
|
||||
the spec-decode path uses -- and runs only the prefill tail through
|
||||
``chunk_gated_delta_rule``. This must produce the same core-attention output and
|
||||
the same ssm-state pool update as running *everything* through
|
||||
``chunk_gated_delta_rule`` (the previous behavior).
|
||||
|
||||
Both paths are exercised through the REAL ``_forward_core``:
|
||||
|
||||
* ``meta_split`` is built by the real ``GDNAttentionMetadataBuilder`` for a
|
||||
mixed batch, so ``num_decodes > 0`` triggers the peel (and the builder rebases
|
||||
``chunk_indices``/``chunk_offsets`` to the prefill-only tail).
|
||||
* ``meta_unified`` is the same metadata with the decodes reclassified as
|
||||
prefills and full-batch chunk metadata, which forces ``_forward_core`` through
|
||||
the existing chunk-only path on identical inputs (the conv is unified over all
|
||||
non-spec tokens in both paths, so it cancels out and only the recurrent split
|
||||
is compared).
|
||||
|
||||
The Triton/FLA chunk backend is forced so the prefill-only ``chunk_indices``
|
||||
must stay consistent with the rebased ``cu_seqlens`` (a stringent, backend
|
||||
portable check of the split wiring).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import types
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not (
|
||||
current_platform.is_cuda() and current_platform.is_device_capability_family(100)
|
||||
):
|
||||
pytest.skip(
|
||||
reason="GDN _forward_core split test uses the CuteDSL prefill backend "
|
||||
"(requires CUDA SM10x).",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from tests.v1.attention.utils import ( # noqa: E402
|
||||
BatchSpec,
|
||||
create_common_attn_metadata,
|
||||
create_vllm_config,
|
||||
)
|
||||
from vllm.config import set_current_vllm_config # noqa: E402
|
||||
from vllm.model_executor.layers.fla.ops.index import ( # noqa: E402
|
||||
prepare_chunk_indices,
|
||||
prepare_chunk_offsets,
|
||||
)
|
||||
from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE # noqa: E402
|
||||
from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn # noqa: E402
|
||||
from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( # noqa: E402
|
||||
ChunkGatedDeltaRule,
|
||||
QwenGatedDeltaNetAttention,
|
||||
)
|
||||
from vllm.model_executor.layers.mamba.mamba_utils import ( # noqa: E402
|
||||
MambaStateShapeCalculator,
|
||||
)
|
||||
from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402
|
||||
GDNAttentionMetadataBuilder,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import MambaSpec # noqa: E402
|
||||
|
||||
# Small GDN dims; head_k_dim/head_v_dim=128 keeps the chunk/update kernels happy.
|
||||
H = 4 # num key heads
|
||||
HV = 8 # num value heads
|
||||
K = 128 # head_k_dim
|
||||
V = 128 # head_v_dim
|
||||
CONV_KERNEL = 4
|
||||
KEY_DIM = H * K
|
||||
VALUE_DIM = HV * V
|
||||
CONV_DIM = 2 * KEY_DIM + VALUE_DIM
|
||||
BLOCK_SIZE = 16
|
||||
PREFIX = "model.layers.0.linear_attn"
|
||||
|
||||
|
||||
def _make_vllm_config():
|
||||
# A small, ungated GDN model whose config is cached locally; only the config
|
||||
# (scheduler/cache/compilation/hf) is used here, never the weights. Inject
|
||||
# linear_key_head_dim=128 and request the CuteDSL prefill backend -- the
|
||||
# supported GDN chunk kernel on Blackwell (the Triton/FLA chunk kernel is
|
||||
# unsupported on SM10x). CuteDSL consumes chunk_indices/chunk_offsets, so
|
||||
# this also exercises the prefill-only chunk-metadata wiring.
|
||||
cfg = create_vllm_config(
|
||||
model_name="Qwen/Qwen3.5-0.8B",
|
||||
block_size=BLOCK_SIZE,
|
||||
hf_config_override={"linear_key_head_dim": K},
|
||||
)
|
||||
cfg.additional_config = {"gdn_prefill_backend": "cutedsl"}
|
||||
return cfg
|
||||
|
||||
|
||||
def _build_layer(
|
||||
vllm_config, conv_state, ssm_state, A_log, dt_bias, conv_weight, conv_bias
|
||||
):
|
||||
"""A minimal object that runs the real ``_forward_core`` bound to it."""
|
||||
layer = types.SimpleNamespace()
|
||||
layer.prefix = PREFIX
|
||||
layer.enable_packed_recurrent_decode = False
|
||||
layer.tp_size = 1
|
||||
layer.num_k_heads = H
|
||||
layer.num_v_heads = HV
|
||||
layer.head_k_dim = K
|
||||
layer.head_v_dim = V
|
||||
layer.key_dim = KEY_DIM
|
||||
layer.value_dim = VALUE_DIM
|
||||
layer.activation = "silu"
|
||||
layer.A_log = A_log
|
||||
layer.dt_bias = dt_bias
|
||||
layer.conv1d = types.SimpleNamespace(weight=conv_weight, bias=conv_bias)
|
||||
layer.kv_cache = (conv_state, ssm_state)
|
||||
with set_current_vllm_config(vllm_config):
|
||||
layer.chunk_gated_delta_rule = ChunkGatedDeltaRule()
|
||||
for name in (
|
||||
"rearrange_mixed_qkv",
|
||||
"_forward_core",
|
||||
):
|
||||
setattr(
|
||||
layer,
|
||||
name,
|
||||
types.MethodType(getattr(QwenGatedDeltaNetAttention, name), layer),
|
||||
)
|
||||
return layer
|
||||
|
||||
|
||||
def _run_forward_core(layer, meta, mixed_qkv, b, a, num_tokens):
|
||||
core_attn_out = torch.zeros(
|
||||
num_tokens, HV, V, dtype=mixed_qkv.dtype, device=mixed_qkv.device
|
||||
)
|
||||
ctx = types.SimpleNamespace(attn_metadata={PREFIX: meta})
|
||||
with patch.object(qwen_gdn_linear_attn, "get_forward_context", return_value=ctx):
|
||||
layer._forward_core(
|
||||
mixed_qkv=mixed_qkv.clone(),
|
||||
b=b.clone(),
|
||||
a=a.clone(),
|
||||
core_attn_out=core_attn_out,
|
||||
)
|
||||
return core_attn_out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32])
|
||||
@pytest.mark.parametrize("num_decodes,prefill_lens", [(3, [512, 300]), (4, [64, 5])])
|
||||
@pytest.mark.parametrize("fresh_prefill", [False, True])
|
||||
def test_forward_core_split_matches_unified(
|
||||
state_dtype: torch.dtype,
|
||||
num_decodes: int,
|
||||
prefill_lens: list[int],
|
||||
fresh_prefill: bool,
|
||||
) -> None:
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda")
|
||||
vllm_config = _make_vllm_config()
|
||||
|
||||
# Decode-first batch: D 1-token decodes (with context), then the prefills.
|
||||
decode_seq_lens = [64] * num_decodes
|
||||
prefill_seq_lens = [
|
||||
pl if (fresh_prefill and i == 0) else pl + 37
|
||||
for i, pl in enumerate(prefill_lens)
|
||||
]
|
||||
seq_lens = decode_seq_lens + prefill_seq_lens
|
||||
query_lens = [1] * num_decodes + list(prefill_lens)
|
||||
batch = BatchSpec(seq_lens=seq_lens, query_lens=query_lens)
|
||||
|
||||
builder = GDNAttentionMetadataBuilder(
|
||||
kv_cache_spec=MambaSpec(
|
||||
block_size=BLOCK_SIZE, shapes=((16, 64),), dtypes=(torch.float16,)
|
||||
),
|
||||
layer_names=[PREFIX],
|
||||
vllm_config=vllm_config,
|
||||
device=device,
|
||||
)
|
||||
common = create_common_attn_metadata(
|
||||
batch, BLOCK_SIZE, device, arange_block_indices=True
|
||||
)
|
||||
with set_current_vllm_config(vllm_config):
|
||||
meta_split = builder.build(common_prefix_len=0, common_attn_metadata=common)
|
||||
|
||||
assert meta_split.spec_sequence_masks is None
|
||||
assert meta_split.num_decodes == num_decodes
|
||||
assert meta_split.num_prefills == len(prefill_lens)
|
||||
assert meta_split.num_decode_tokens == num_decodes
|
||||
assert builder.gdn_prefill_backend == "cutedsl"
|
||||
|
||||
num_tokens = sum(query_lens)
|
||||
|
||||
# Full-batch chunk metadata for the unified reference path, built the same
|
||||
# way the builder would for a non-split batch (backend-matched).
|
||||
cu_full = meta_split.non_spec_query_start_loc
|
||||
if builder.gdn_prefill_backend == "cutedsl":
|
||||
from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
|
||||
prepare_metadata_cutedsl,
|
||||
)
|
||||
|
||||
full_ci, full_co = prepare_metadata_cutedsl(
|
||||
cu_full, int(cu_full[-1].item()), FLA_CHUNK_SIZE
|
||||
)
|
||||
else:
|
||||
cu_full_cpu = cu_full.cpu()
|
||||
full_ci = prepare_chunk_indices(cu_full_cpu, FLA_CHUNK_SIZE).to(device)
|
||||
full_co = prepare_chunk_offsets(cu_full_cpu, FLA_CHUNK_SIZE).to(device)
|
||||
meta_unified = dataclasses.replace(
|
||||
meta_split,
|
||||
num_decodes=0,
|
||||
num_decode_tokens=0,
|
||||
num_prefills=meta_split.num_decodes + meta_split.num_prefills,
|
||||
num_prefill_tokens=(
|
||||
meta_split.num_decode_tokens + meta_split.num_prefill_tokens
|
||||
),
|
||||
chunk_indices=full_ci,
|
||||
chunk_offsets=full_co,
|
||||
# Unified path: the chunk kernel processes the full non-spec batch.
|
||||
prefill_query_start_loc=meta_split.non_spec_query_start_loc,
|
||||
prefill_state_indices=meta_split.non_spec_state_indices_tensor,
|
||||
prefill_has_initial_state=meta_split.has_initial_state,
|
||||
)
|
||||
|
||||
# Size the state pools from the indices the builder actually produced.
|
||||
pool_size = int(meta_split.non_spec_state_indices_tensor.max().item()) + 1
|
||||
conv_state_shape, temporal_state_shape = (
|
||||
MambaStateShapeCalculator.gated_delta_net_state_shape(
|
||||
1, H, HV, K, V, CONV_KERNEL, num_spec=0
|
||||
)
|
||||
)
|
||||
conv_state0 = (
|
||||
torch.randn(pool_size, *conv_state_shape, dtype=torch.bfloat16, device=device)
|
||||
* 0.05
|
||||
)
|
||||
ssm_state0 = (
|
||||
torch.randn(pool_size, *temporal_state_shape, dtype=state_dtype, device=device)
|
||||
* 0.05
|
||||
)
|
||||
|
||||
A_log = torch.randn(HV, dtype=torch.float32, device=device) * 0.1
|
||||
dt_bias = torch.randn(HV, dtype=torch.float32, device=device) * 0.1
|
||||
conv_weight = (
|
||||
torch.randn(CONV_DIM, 1, CONV_KERNEL, dtype=torch.bfloat16, device=device) * 0.1
|
||||
)
|
||||
conv_bias = torch.randn(CONV_DIM, dtype=torch.bfloat16, device=device) * 0.1
|
||||
|
||||
mixed_qkv = (
|
||||
torch.randn(num_tokens, CONV_DIM, dtype=torch.bfloat16, device=device) * 0.1
|
||||
)
|
||||
a = torch.randn(num_tokens, HV, dtype=torch.bfloat16, device=device) * 0.1
|
||||
b = torch.randn(num_tokens, HV, dtype=torch.bfloat16, device=device) * 0.1
|
||||
|
||||
# ---- Split path (real _forward_core, meta_split) ----
|
||||
conv_state_split = conv_state0.clone()
|
||||
ssm_state_split = ssm_state0.clone()
|
||||
layer_split = _build_layer(
|
||||
vllm_config,
|
||||
conv_state_split,
|
||||
ssm_state_split,
|
||||
A_log,
|
||||
dt_bias,
|
||||
conv_weight,
|
||||
conv_bias,
|
||||
)
|
||||
out_split = _run_forward_core(layer_split, meta_split, mixed_qkv, b, a, num_tokens)
|
||||
|
||||
# ---- Unified path (real _forward_core, meta_unified) ----
|
||||
conv_state_unified = conv_state0.clone()
|
||||
ssm_state_unified = ssm_state0.clone()
|
||||
layer_unified = _build_layer(
|
||||
vllm_config,
|
||||
conv_state_unified,
|
||||
ssm_state_unified,
|
||||
A_log,
|
||||
dt_bias,
|
||||
conv_weight,
|
||||
conv_bias,
|
||||
)
|
||||
out_unified = _run_forward_core(
|
||||
layer_unified, meta_unified, mixed_qkv, b, a, num_tokens
|
||||
)
|
||||
|
||||
# Conv is unified in both paths, so the conv-state update must be identical.
|
||||
torch.testing.assert_close(conv_state_split, conv_state_unified, atol=0, rtol=0)
|
||||
|
||||
# Chunk vs. recurrent update accumulate in different orders; mirror the
|
||||
# tolerances used by the kernel-level parity test.
|
||||
if state_dtype == torch.float32:
|
||||
atol = rtol = 2e-2
|
||||
else:
|
||||
atol = rtol = 6e-2
|
||||
torch.testing.assert_close(out_split, out_unified, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(ssm_state_split, ssm_state_unified, atol=atol, rtol=rtol)
|
||||
@@ -0,0 +1,199 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not (
|
||||
current_platform.is_cuda() and current_platform.is_device_capability_family(100)
|
||||
):
|
||||
pytest.skip(
|
||||
reason="GDN CuteDSL prefill requires CUDA SM10x.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from vllm.model_executor.layers.fla.ops import ( # noqa: E402
|
||||
chunk_gated_delta_rule,
|
||||
)
|
||||
from vllm.model_executor.layers.fla.ops.index import ( # noqa: E402
|
||||
prepare_chunk_indices,
|
||||
prepare_chunk_offsets,
|
||||
)
|
||||
from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( # noqa: E402
|
||||
chunk_gated_delta_rule_cutedsl,
|
||||
prepare_metadata_cutedsl,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_seqs", [1, 5, 257])
|
||||
@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32])
|
||||
def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
|
||||
seq_lens = torch.randint(
|
||||
1,
|
||||
130,
|
||||
(num_seqs,),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
cu_seqlens = torch.zeros(num_seqs + 1, device="cuda", dtype=torch.int32)
|
||||
cu_seqlens[1:] = seq_lens.to(device="cuda").cumsum(0)
|
||||
total_tokens = int(cu_seqlens[-1].item())
|
||||
|
||||
num_k_heads = 4
|
||||
num_v_heads = 8
|
||||
head_k_dim = 128
|
||||
head_v_dim = 128
|
||||
dtype = torch.bfloat16
|
||||
|
||||
q = torch.randn(
|
||||
1,
|
||||
total_tokens,
|
||||
num_k_heads,
|
||||
head_k_dim,
|
||||
device="cuda",
|
||||
dtype=dtype,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn(
|
||||
1,
|
||||
total_tokens,
|
||||
num_v_heads,
|
||||
head_v_dim,
|
||||
device="cuda",
|
||||
dtype=dtype,
|
||||
)
|
||||
q = F.normalize(q.float(), p=2, dim=-1).to(dtype)
|
||||
k = F.normalize(k.float(), p=2, dim=-1).to(dtype)
|
||||
a = torch.randn(
|
||||
1,
|
||||
total_tokens,
|
||||
num_v_heads,
|
||||
device="cuda",
|
||||
dtype=dtype,
|
||||
)
|
||||
b = torch.randn(
|
||||
1,
|
||||
total_tokens,
|
||||
num_v_heads,
|
||||
device="cuda",
|
||||
dtype=dtype,
|
||||
)
|
||||
# Match upstream FLA GatedDeltaNet synthetic initialization:
|
||||
# https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/gated_deltanet.py
|
||||
A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(0, 16)
|
||||
A_log = torch.log(A)
|
||||
dt = torch.exp(
|
||||
torch.rand(num_v_heads, device="cuda", dtype=torch.float32)
|
||||
* (math.log(0.1) - math.log(0.001))
|
||||
+ math.log(0.001)
|
||||
)
|
||||
dt = torch.clamp(dt, min=1e-4)
|
||||
dt_bias = dt + torch.log(-torch.expm1(-dt))
|
||||
g = -A_log.exp().view(1, 1, num_v_heads) * F.softplus(
|
||||
a.float() + dt_bias.view(1, 1, num_v_heads)
|
||||
)
|
||||
beta = torch.sigmoid(b.float())
|
||||
initial_state = (
|
||||
torch.randn(
|
||||
num_seqs,
|
||||
num_v_heads,
|
||||
head_v_dim,
|
||||
head_k_dim,
|
||||
device="cuda",
|
||||
dtype=state_dtype,
|
||||
)
|
||||
* 0.05
|
||||
)
|
||||
|
||||
# check metadata kernel
|
||||
chunk_indices, chunk_offsets = prepare_metadata_cutedsl(cu_seqlens, total_tokens)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
expected_indices = prepare_chunk_indices(cu_seqlens, 64)
|
||||
expected_offsets = prepare_chunk_offsets(cu_seqlens, 64)
|
||||
total_chunks = int(expected_offsets[-1].item())
|
||||
|
||||
torch.testing.assert_close(chunk_offsets, expected_offsets.to(torch.int32))
|
||||
torch.testing.assert_close(
|
||||
chunk_indices[:total_chunks],
|
||||
expected_indices,
|
||||
)
|
||||
|
||||
ref_o, ref_state = chunk_gated_delta_rule(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
g=g,
|
||||
beta=beta,
|
||||
initial_state=initial_state,
|
||||
output_final_state=True,
|
||||
cu_seqlens=cu_seqlens,
|
||||
use_qk_l2norm_in_kernel=False,
|
||||
)
|
||||
actual_core_attn_out = torch.empty(
|
||||
total_tokens,
|
||||
num_v_heads,
|
||||
head_v_dim,
|
||||
device="cuda",
|
||||
dtype=dtype,
|
||||
)
|
||||
actual_o, actual_state = chunk_gated_delta_rule_cutedsl(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
g=g,
|
||||
beta=beta,
|
||||
initial_state=initial_state,
|
||||
cu_seqlens=cu_seqlens,
|
||||
chunk_indices=chunk_indices,
|
||||
chunk_offsets=chunk_offsets,
|
||||
core_attn_out=actual_core_attn_out,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
# check main kernel
|
||||
o_error = (actual_o.float() - ref_o.float()).abs()
|
||||
state_error = (
|
||||
actual_state.float() - ref_state.to(actual_state.dtype).float()
|
||||
).abs()
|
||||
assert o_error.max().item() < 2e-3
|
||||
assert o_error.mean().item() < 6e-5
|
||||
assert state_error.max().item() < 2e-2
|
||||
assert state_error.mean().item() < 6e-4
|
||||
core_attn_out_error = (
|
||||
actual_core_attn_out.float() - actual_o.squeeze(0).float()
|
||||
).abs()
|
||||
assert core_attn_out_error.max().item() == 0
|
||||
|
||||
# check main kernel when core_attn_out is not passed
|
||||
no_buffer_o, no_buffer_state = chunk_gated_delta_rule_cutedsl(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
g=g,
|
||||
beta=beta,
|
||||
initial_state=initial_state,
|
||||
cu_seqlens=cu_seqlens,
|
||||
chunk_indices=chunk_indices,
|
||||
chunk_offsets=chunk_offsets,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
no_buffer_o_error = (no_buffer_o.float() - ref_o.float()).abs()
|
||||
no_buffer_state_error = (
|
||||
no_buffer_state.float() - ref_state.to(no_buffer_state.dtype).float()
|
||||
).abs()
|
||||
buffer_o_error = (no_buffer_o.float() - actual_o.float()).abs()
|
||||
buffer_state_error = (
|
||||
no_buffer_state.float() - actual_state.to(no_buffer_state.dtype).float()
|
||||
).abs()
|
||||
assert no_buffer_o_error.max().item() < 2e-3
|
||||
assert no_buffer_o_error.mean().item() < 6e-5
|
||||
assert no_buffer_state_error.max().item() < 2e-2
|
||||
assert no_buffer_state_error.mean().item() < 6e-4
|
||||
assert buffer_o_error.max().item() == 0
|
||||
assert buffer_state_error.max().item() == 0
|
||||
@@ -0,0 +1,138 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import unittest
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.utils import ensure_current_vllm_config, multi_gpu_test
|
||||
from vllm.distributed.parallel_state import (
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm.model_executor.layers.mamba.mamba_mixer2 import Mixer2RMSNormGated
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@pytest.mark.parametrize("batch_size", [8])
|
||||
@pytest.mark.parametrize("seq_len", [128])
|
||||
@pytest.mark.parametrize(
|
||||
"hidden_size_n_groups",
|
||||
[
|
||||
(64, 1),
|
||||
(64, 2),
|
||||
(64, 4), # hidden_size be divisible by num_gpus
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16])
|
||||
def test_mixer2_gated_norm_multi_gpu(
|
||||
batch_size: int,
|
||||
seq_len: int,
|
||||
hidden_size_n_groups: tuple[int, int],
|
||||
dtype: torch.dtype,
|
||||
device: str = "cuda",
|
||||
):
|
||||
hidden_size, n_groups = hidden_size_n_groups
|
||||
num_processes = 2
|
||||
|
||||
def run_torch_spawn(fn, nprocs):
|
||||
# need to use torch.mp.spawn otherwise will have problems with
|
||||
# torch.distributed and cuda
|
||||
torch.multiprocessing.spawn(
|
||||
fn,
|
||||
args=(
|
||||
num_processes,
|
||||
batch_size,
|
||||
seq_len,
|
||||
hidden_size,
|
||||
n_groups,
|
||||
dtype,
|
||||
device,
|
||||
),
|
||||
nprocs=nprocs,
|
||||
)
|
||||
|
||||
run_torch_spawn(mixer2_gated_norm_tensor_parallel, 2)
|
||||
|
||||
|
||||
def mixer2_gated_norm_tensor_parallel(
|
||||
local_rank: int,
|
||||
world_size: int,
|
||||
batch_size: int,
|
||||
seq_len: int,
|
||||
hidden_size: int,
|
||||
n_groups: int,
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
):
|
||||
set_random_seed(0)
|
||||
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
torch.set_default_device(device)
|
||||
torch.set_default_dtype(dtype)
|
||||
|
||||
update_environment_variables(
|
||||
{
|
||||
"RANK": str(local_rank),
|
||||
"LOCAL_RANK": str(local_rank),
|
||||
"WORLD_SIZE": str(world_size),
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": "12345",
|
||||
}
|
||||
)
|
||||
|
||||
# initialize distributed
|
||||
init_distributed_environment()
|
||||
with ensure_current_vllm_config():
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
# create random weights an inputs
|
||||
weight = torch.rand((hidden_size,), dtype=dtype, device=device)
|
||||
hidden_states = torch.randn(batch_size, seq_len, hidden_size)
|
||||
gate_states = torch.randn(batch_size, seq_len, hidden_size)
|
||||
|
||||
# create gated-norm with TP
|
||||
mixer = Mixer2RMSNormGated(
|
||||
full_hidden_size=hidden_size,
|
||||
full_n_groups=n_groups,
|
||||
)
|
||||
mixer.weight.weight_loader(mixer.weight, weight) # load
|
||||
|
||||
# create gated-norm without TP to compute reference
|
||||
# - utilize mock patching to disable TP when
|
||||
with (
|
||||
unittest.mock.patch(
|
||||
"vllm.model_executor.layers.mamba.mamba_mixer2."
|
||||
"get_tensor_model_parallel_world_size",
|
||||
return_value=1,
|
||||
),
|
||||
unittest.mock.patch(
|
||||
"vllm.model_executor.layers.mamba.mamba_mixer2."
|
||||
"get_tensor_model_parallel_rank",
|
||||
return_value=0,
|
||||
),
|
||||
):
|
||||
mixer_single_gpu = Mixer2RMSNormGated(
|
||||
full_hidden_size=hidden_size,
|
||||
full_n_groups=n_groups,
|
||||
)
|
||||
# assign weight to single-gpu mixer
|
||||
mixer_single_gpu.weight.data = weight
|
||||
|
||||
# generate and compare
|
||||
N = hidden_size // world_size
|
||||
output = mixer(
|
||||
hidden_states[..., local_rank * N : (local_rank + 1) * N],
|
||||
gate_states[..., local_rank * N : (local_rank + 1) * N],
|
||||
)
|
||||
ref_output = mixer_single_gpu(hidden_states, gate_states)
|
||||
torch.testing.assert_close(
|
||||
output,
|
||||
ref_output[..., local_rank * N : (local_rank + 1) * N],
|
||||
atol=5e-3,
|
||||
rtol=1e-3,
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,212 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Unit tests for the JSON-based config loader added to selective_state_update.
|
||||
|
||||
Tests cover:
|
||||
- Flat MoE-style filename generation
|
||||
- VLLM_TUNED_CONFIG_FOLDER env-var override
|
||||
- Fallback to heuristic when no config file exists
|
||||
- Nearest effective_batch interpolation
|
||||
- Edge cases: non-dict JSON, empty config
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
|
||||
_get_default_ssm_launch_config,
|
||||
_try_get_optimal_ssm_config_cached,
|
||||
get_ssm_config_file_name,
|
||||
get_ssm_configs,
|
||||
get_ssm_device_name,
|
||||
try_get_optimal_ssm_config,
|
||||
)
|
||||
|
||||
# Common kwargs for try_get_optimal_ssm_config. Tests pick (batch, nheads) so
|
||||
# their product (effective_batch) matches the value being probed.
|
||||
_HEADDIM = 64
|
||||
_CACHE_DTYPE = "float32"
|
||||
|
||||
|
||||
def _clear_caches() -> None:
|
||||
get_ssm_configs.cache_clear()
|
||||
_try_get_optimal_ssm_config_cached.cache_clear()
|
||||
|
||||
|
||||
def _write_config(tmp_path, dstate: int, payload: dict) -> None:
|
||||
"""Write payload as the bundled config for (headdim, dstate, cache_dtype)."""
|
||||
device_name = get_ssm_device_name()
|
||||
config_path = tmp_path / get_ssm_config_file_name(
|
||||
_HEADDIM, dstate, _CACHE_DTYPE, device_name
|
||||
)
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(payload, f)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config filename generation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_config_file_name_format():
|
||||
name = get_ssm_config_file_name(
|
||||
headdim=64, dstate=128, cache_dtype="float32", device_name="NVIDIA_B200"
|
||||
)
|
||||
assert name == (
|
||||
"headdim=64,dstate=128,device_name=NVIDIA_B200,cache_dtype=float32.json"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# VLLM_TUNED_CONFIG_FOLDER override
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_env_override_loads_custom_config(monkeypatch, tmp_path):
|
||||
"""VLLM_TUNED_CONFIG_FOLDER should take precedence over the bundled dir."""
|
||||
_write_config(
|
||||
tmp_path,
|
||||
dstate=16,
|
||||
payload={
|
||||
"1": {"BLOCK_SIZE_M": 4, "num_warps": 1},
|
||||
},
|
||||
)
|
||||
|
||||
monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
|
||||
_clear_caches()
|
||||
|
||||
cfg = get_ssm_configs(_HEADDIM, 16, _CACHE_DTYPE)
|
||||
assert cfg is not None
|
||||
assert cfg[1] == {"BLOCK_SIZE_M": 4, "num_warps": 1}
|
||||
|
||||
_clear_caches()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fallback to heuristic when no config file exists
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_fallback_when_no_config(monkeypatch, tmp_path):
|
||||
"""try_get_optimal_ssm_config must fall back to _get_default_ssm_launch_config
|
||||
when no JSON file is found for the current
|
||||
(device, headdim, dstate, cache_dtype) combination.
|
||||
"""
|
||||
monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
|
||||
monkeypatch.setattr(
|
||||
"vllm.model_executor.layers.mamba.ops.mamba_ssm._CONFIGS_DIR",
|
||||
str(tmp_path),
|
||||
)
|
||||
|
||||
for dstate in (8, 16, 32, 64, 128, 256):
|
||||
for is_blackwell in (False, True):
|
||||
_clear_caches()
|
||||
block_m, warps = try_get_optimal_ssm_config(
|
||||
headdim=_HEADDIM,
|
||||
dstate=dstate,
|
||||
batch=1,
|
||||
nheads=1,
|
||||
cache_dtype=_CACHE_DTYPE,
|
||||
is_blackwell=is_blackwell,
|
||||
)
|
||||
assert (block_m, warps) == _get_default_ssm_launch_config(
|
||||
dstate, is_blackwell=is_blackwell
|
||||
)
|
||||
|
||||
_clear_caches()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Nearest effective_batch interpolation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_nearest_effective_batch_interpolation(monkeypatch, tmp_path):
|
||||
"""When effective_batch = batch*nheads is not an exact key, the closest
|
||||
key should be selected."""
|
||||
_write_config(
|
||||
tmp_path,
|
||||
dstate=32,
|
||||
payload={
|
||||
"64": {"BLOCK_SIZE_M": 8, "num_warps": 1},
|
||||
"4096": {"BLOCK_SIZE_M": 32, "num_warps": 4},
|
||||
},
|
||||
)
|
||||
|
||||
monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
|
||||
_clear_caches()
|
||||
|
||||
# effective_batch = 1*128 = 128 -> closer to 64 than to 4096
|
||||
block_m, warps = try_get_optimal_ssm_config(
|
||||
headdim=_HEADDIM,
|
||||
dstate=32,
|
||||
batch=1,
|
||||
nheads=128,
|
||||
cache_dtype=_CACHE_DTYPE,
|
||||
is_blackwell=False,
|
||||
)
|
||||
assert block_m == 8 and warps == 1
|
||||
|
||||
# effective_batch = 4*1024 = 4096 -> exact match on 4096
|
||||
block_m, warps = try_get_optimal_ssm_config(
|
||||
headdim=_HEADDIM,
|
||||
dstate=32,
|
||||
batch=4,
|
||||
nheads=1024,
|
||||
cache_dtype=_CACHE_DTYPE,
|
||||
is_blackwell=False,
|
||||
)
|
||||
assert block_m == 32 and warps == 4
|
||||
|
||||
_clear_caches()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Edge cases: malformed / empty config files
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_non_dict_json_returns_none(monkeypatch, tmp_path):
|
||||
"""A valid JSON file that is not a dict (e.g. a list) must be ignored
|
||||
and return None rather than raising AttributeError."""
|
||||
device_name = get_ssm_device_name()
|
||||
config_path = tmp_path / get_ssm_config_file_name(
|
||||
_HEADDIM, 16, _CACHE_DTYPE, device_name
|
||||
)
|
||||
with open(config_path, "w") as f:
|
||||
json.dump([1, 2, 3], f)
|
||||
|
||||
monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
|
||||
monkeypatch.setattr(
|
||||
"vllm.model_executor.layers.mamba.ops.mamba_ssm._CONFIGS_DIR",
|
||||
str(tmp_path),
|
||||
)
|
||||
_clear_caches()
|
||||
|
||||
assert get_ssm_configs(_HEADDIM, 16, _CACHE_DTYPE) is None
|
||||
|
||||
_clear_caches()
|
||||
|
||||
|
||||
def test_empty_config_falls_back_to_heuristic(monkeypatch, tmp_path):
|
||||
"""An empty JSON object {} must not crash min() — should fall back
|
||||
to the hard-coded heuristic."""
|
||||
_write_config(tmp_path, dstate=64, payload={})
|
||||
|
||||
monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
|
||||
_clear_caches()
|
||||
|
||||
dstate = 64
|
||||
block_m, warps = try_get_optimal_ssm_config(
|
||||
headdim=_HEADDIM,
|
||||
dstate=dstate,
|
||||
batch=1,
|
||||
nheads=64,
|
||||
cache_dtype=_CACHE_DTYPE,
|
||||
is_blackwell=False,
|
||||
)
|
||||
assert (block_m, warps) == _get_default_ssm_launch_config(
|
||||
dstate=dstate, is_blackwell=False
|
||||
)
|
||||
|
||||
_clear_caches()
|
||||
@@ -0,0 +1,579 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange, repeat
|
||||
|
||||
from vllm.model_executor.layers.mamba.ops.ssd_combined import (
|
||||
mamba_chunk_scan_combined_varlen,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.v1.attention.backends.mamba2_attn import compute_varlen_chunk_metadata
|
||||
|
||||
# All kernels exercised here are pure Triton, so they run on any backend
|
||||
# that the vLLM platform layer treats as a CUDA-alike device or as XPU.
|
||||
DEVICE = current_platform.device_type
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
|
||||
reason="Mamba2 SSD Triton kernels require a CUDA-alike or XPU device.",
|
||||
)
|
||||
|
||||
# Added by the IBM Team, 2024
|
||||
|
||||
# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/modules/ssd_minimal.py
|
||||
|
||||
|
||||
# this is the segsum implementation taken from above
|
||||
def segsum(x):
|
||||
"""Calculates segment sum."""
|
||||
T = x.size(-1)
|
||||
x = repeat(x, "... d -> ... d e", e=T)
|
||||
mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=-1)
|
||||
x = x.masked_fill(~mask, 0)
|
||||
x_segsum = torch.cumsum(x, dim=-2)
|
||||
mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=0)
|
||||
x_segsum = x_segsum.masked_fill(~mask, -torch.inf)
|
||||
return x_segsum
|
||||
|
||||
|
||||
def ssd_minimal_discrete(X, A, B, C, block_len, initial_states=None):
|
||||
"""
|
||||
Arguments:
|
||||
X: (batch, length, n_heads, d_head)
|
||||
A: (batch, length, n_heads)
|
||||
B: (batch, length, n_heads, d_state)
|
||||
C: (batch, length, n_heads, d_state)
|
||||
Return:
|
||||
Y: (batch, length, n_heads, d_head)
|
||||
"""
|
||||
assert X.dtype == A.dtype == B.dtype == C.dtype
|
||||
assert X.shape[1] % block_len == 0
|
||||
|
||||
# Rearrange into blocks/chunks
|
||||
X, A, B, C = (
|
||||
rearrange(x, "b (c l) ... -> b c l ...", l=block_len) for x in (X, A, B, C)
|
||||
)
|
||||
|
||||
A = rearrange(A, "b c l h -> b h c l")
|
||||
A_cumsum = torch.cumsum(A, dim=-1)
|
||||
|
||||
# 1. Compute the output for each intra-chunk (diagonal blocks)
|
||||
L = torch.exp(segsum(A))
|
||||
Y_diag = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", C, B, L, X)
|
||||
|
||||
# 2. Compute the state for each intra-chunk
|
||||
# (right term of low-rank factorization of off-diagonal blocks; B terms)
|
||||
decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)
|
||||
states = torch.einsum("bclhn,bhcl,bclhp->bchpn", B, decay_states, X)
|
||||
|
||||
# 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at
|
||||
# chunk boundaries
|
||||
# (middle term of factorization of off-diag blocks; A terms)
|
||||
if initial_states is None:
|
||||
initial_states = torch.zeros_like(states[:, :1])
|
||||
states = torch.cat([initial_states, states], dim=1)
|
||||
decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
|
||||
new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
|
||||
states, final_state = new_states[:, :-1], new_states[:, -1]
|
||||
|
||||
# 4. Compute state -> output conversion per chunk
|
||||
# (left term of low-rank factorization of off-diagonal blocks; C terms)
|
||||
state_decay_out = torch.exp(A_cumsum)
|
||||
Y_off = torch.einsum("bclhn,bchpn,bhcl->bclhp", C, states, state_decay_out)
|
||||
|
||||
# Add output of intra-chunk and inter-chunk terms
|
||||
# (diagonal and off-diagonal blocks)
|
||||
Y = rearrange(Y_diag + Y_off, "b c l h p -> b (c l) h p")
|
||||
return Y, final_state
|
||||
|
||||
|
||||
def generate_random_inputs(batch_size, seqlen, n_heads, d_head, itype, device=DEVICE):
|
||||
set_random_seed(0)
|
||||
A = -torch.exp(torch.rand(n_heads, dtype=itype, device=device))
|
||||
dt = F.softplus(
|
||||
torch.randn(batch_size, seqlen, n_heads, dtype=itype, device=device) - 4
|
||||
)
|
||||
X = torch.randn((batch_size, seqlen, n_heads, d_head), dtype=itype, device=device)
|
||||
B = torch.randn((batch_size, seqlen, n_heads, d_head), dtype=itype, device=device)
|
||||
C = torch.randn((batch_size, seqlen, n_heads, d_head), dtype=itype, device=device)
|
||||
|
||||
return A, dt, X, B, C
|
||||
|
||||
|
||||
def generate_continuous_batched_examples(
|
||||
example_lens_by_batch,
|
||||
num_examples,
|
||||
full_length,
|
||||
last_taken,
|
||||
exhausted,
|
||||
n_heads,
|
||||
d_head,
|
||||
itype,
|
||||
device=DEVICE,
|
||||
return_naive_ref=True,
|
||||
):
|
||||
# this function generates a random examples of certain length
|
||||
# and then cut according to "example_lens_by_batch" and feed
|
||||
# them in continuous batches to the kernels.
|
||||
# If if return_naive_ref=True, the naive torch implementation
|
||||
# ssd_minimal_discrete will be used to compute and return
|
||||
# reference output.
|
||||
|
||||
# generate the full-length example
|
||||
A, dt, X, B, C = generate_random_inputs(
|
||||
num_examples, full_length, n_heads, d_head, itype
|
||||
)
|
||||
|
||||
if return_naive_ref:
|
||||
Y_min, final_state_min = ssd_minimal_discrete(
|
||||
X * dt.unsqueeze(-1), A * dt, B, C, block_len=full_length // 4
|
||||
)
|
||||
|
||||
# internal function that outputs a cont batch of examples
|
||||
# given a tuple of lengths for each example in the batch
|
||||
# e.g., example_lens=(8, 4) means take 8 samples from first eg,
|
||||
# 4 examples from second eg, etc
|
||||
def get_continuous_batch(example_lens: tuple[int, ...]):
|
||||
indices = []
|
||||
for i, x in enumerate(example_lens):
|
||||
c = last_taken.get(i, 0)
|
||||
indices.append((c, c + x))
|
||||
last_taken[i] = (c + x) % full_length
|
||||
exhausted[i] = last_taken[i] == 0
|
||||
|
||||
return (
|
||||
torch.concat([x[i, s:e] for i, (s, e) in enumerate(indices)]).unsqueeze(0)
|
||||
for x in (dt, X, B, C)
|
||||
)
|
||||
|
||||
# internal function that maps "n" to the appropriate right boundary
|
||||
# value when forming continuous batches from examples of length given
|
||||
# by "full_length".
|
||||
# - e.g., when n > full_length, returns n % full_length
|
||||
# when n == full_length, returns full_length
|
||||
def end_boundary(n: int):
|
||||
return n - ((n - 1) // full_length) * full_length
|
||||
|
||||
IND_E = None
|
||||
for spec in example_lens_by_batch:
|
||||
# get the (maybe partial) example seen in this cont batch
|
||||
dt2, X2, B2, C2 = get_continuous_batch(spec)
|
||||
|
||||
# get the metadata
|
||||
cu_seqlens = torch.tensor((0,) + spec, device=device).cumsum(dim=0)
|
||||
seq_idx = torch.zeros(
|
||||
cu_seqlens[-1], dtype=torch.int32, device=cu_seqlens.device
|
||||
)
|
||||
for i, (srt, end) in enumerate(
|
||||
zip(
|
||||
cu_seqlens,
|
||||
cu_seqlens[1:],
|
||||
)
|
||||
):
|
||||
seq_idx[srt:end] = i
|
||||
|
||||
# for cont batch
|
||||
if IND_E is None:
|
||||
IND_S = [0 for _ in range(len(spec))]
|
||||
else:
|
||||
IND_S = [x % full_length for x in IND_E]
|
||||
IND_E = [end_boundary(x + y) for x, y in zip(IND_S, spec)]
|
||||
|
||||
# varlen has implicit batch=1
|
||||
dt2 = dt2.squeeze(0)
|
||||
X2 = X2.squeeze(0)
|
||||
B2 = B2.squeeze(0)
|
||||
C2 = C2.squeeze(0)
|
||||
yield (
|
||||
[Y_min[s, IND_S[s] : IND_E[s]] for s in range(num_examples)]
|
||||
if return_naive_ref
|
||||
else None,
|
||||
cu_seqlens,
|
||||
seq_idx,
|
||||
(A, dt2, X2, B2, C2),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])
|
||||
@pytest.mark.parametrize("n_heads", [4, 16, 32])
|
||||
@pytest.mark.parametrize("d_head", [5, 8, 32, 128])
|
||||
@pytest.mark.parametrize("seq_len_chunk_size", [(112, 16), (128, 32)])
|
||||
def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size, itype):
|
||||
# this tests the kernels on a single example (bs=1)
|
||||
|
||||
# TODO: the bfloat16 case requires higher thresholds. To be investigated
|
||||
|
||||
if itype == torch.bfloat16:
|
||||
atol, rtol = 5e-2, 5e-2
|
||||
else:
|
||||
atol, rtol = 8e-3, 5e-3
|
||||
|
||||
# set seed
|
||||
batch_size = 1 # batch_size
|
||||
# ssd_minimal_discrete requires chunk_size divide seqlen
|
||||
# - this is only required for generating the reference seqs,
|
||||
# it is not an operational limitation.
|
||||
seqlen, chunk_size = seq_len_chunk_size
|
||||
|
||||
A, dt, X, B, C = generate_random_inputs(batch_size, seqlen, n_heads, d_head, itype)
|
||||
|
||||
Y_min, final_state_min = ssd_minimal_discrete(
|
||||
X * dt.unsqueeze(-1), A * dt, B, C, chunk_size
|
||||
)
|
||||
|
||||
cu_seqlens = torch.tensor((0, seqlen), device=DEVICE).cumsum(dim=0)
|
||||
cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = (
|
||||
compute_varlen_chunk_metadata(cu_seqlens, chunk_size)
|
||||
)
|
||||
# varlen has implicit batch=1
|
||||
X = X.squeeze(0)
|
||||
dt = dt.squeeze(0)
|
||||
A = A.squeeze(0)
|
||||
B = B.squeeze(0)
|
||||
C = C.squeeze(0)
|
||||
Y = torch.empty_like(X)
|
||||
final_state = mamba_chunk_scan_combined_varlen(
|
||||
X,
|
||||
dt,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
chunk_size,
|
||||
cu_seqlens=cu_seqlens.to(torch.int32),
|
||||
cu_chunk_seqlens=cu_chunk_seqlens,
|
||||
last_chunk_indices=last_chunk_indices,
|
||||
seq_idx=seq_idx_chunks,
|
||||
out=Y,
|
||||
D=None,
|
||||
)
|
||||
|
||||
# just test the last in sequence
|
||||
torch.testing.assert_close(Y[-1], Y_min[0, -1], atol=atol, rtol=rtol)
|
||||
|
||||
# just test the last head
|
||||
# NOTE, in the kernel we always cast states to fp32
|
||||
torch.testing.assert_close(
|
||||
final_state[:, -1].to(torch.float32),
|
||||
final_state_min[:, -1].to(torch.float32),
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("itype", [torch.float32])
|
||||
@pytest.mark.parametrize("n_heads", [4, 8])
|
||||
@pytest.mark.parametrize("d_head", [5, 16, 32])
|
||||
@pytest.mark.parametrize(
|
||||
"seq_len_chunk_size_cases",
|
||||
[
|
||||
# small-ish chunk_size (8)
|
||||
(64, 8, 2, [(64, 32), (64, 32)]),
|
||||
(64, 8, 2, [(8, 8), (8, 8), (8, 8)]), # chunk size boundary
|
||||
(
|
||||
64,
|
||||
8,
|
||||
2,
|
||||
[(4, 4), (4, 4), (4, 4), (4, 4)],
|
||||
), # chunk_size larger than cont batches
|
||||
(64, 8, 5, [(64, 32, 16, 8, 8)]),
|
||||
# large-ish chunk_size (256)
|
||||
(64, 256, 1, [(5,), (1,), (1,), (1,)]), # irregular sizes with small sequences
|
||||
(
|
||||
64,
|
||||
256,
|
||||
2,
|
||||
[(5, 30), (1, 2), (1, 2), (1, 2)],
|
||||
), # irregular sizes with small sequences
|
||||
# we also need to test some large seqlen
|
||||
# to catch errors with init states decay
|
||||
(768, 128, 2, [(138, 225), (138, 225)]),
|
||||
],
|
||||
)
|
||||
def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases, itype):
|
||||
# this test with multiple examples in a continuous batch
|
||||
# (i.e. chunked prefill)
|
||||
|
||||
seqlen, chunk_size, num_examples, cases = seq_len_chunk_size_cases
|
||||
|
||||
# This test can have larger error for longer sequences
|
||||
if seqlen > 256:
|
||||
atol, rtol = 1e-2, 5e-3
|
||||
else:
|
||||
atol, rtol = 5e-3, 5e-3
|
||||
|
||||
# hold state during the cutting process so we know if an
|
||||
# example has been exhausted and needs to cycle
|
||||
last_taken: dict = {} # map: eg -> pointer to last taken sample
|
||||
exhausted: dict = {} # map: eg -> boolean indicating example is exhausted
|
||||
|
||||
states = None
|
||||
for Y_min, cu_seqlens, _token_seq_idx, (
|
||||
A,
|
||||
dt,
|
||||
X,
|
||||
B,
|
||||
C,
|
||||
) in generate_continuous_batched_examples(
|
||||
cases, num_examples, seqlen, last_taken, exhausted, n_heads, d_head, itype
|
||||
):
|
||||
cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = (
|
||||
compute_varlen_chunk_metadata(cu_seqlens, chunk_size)
|
||||
)
|
||||
|
||||
Y = torch.empty_like(X)
|
||||
new_states = mamba_chunk_scan_combined_varlen(
|
||||
X,
|
||||
dt,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
chunk_size,
|
||||
cu_seqlens=cu_seqlens.to(torch.int32),
|
||||
cu_chunk_seqlens=cu_chunk_seqlens,
|
||||
last_chunk_indices=last_chunk_indices,
|
||||
seq_idx=seq_idx_chunks,
|
||||
out=Y,
|
||||
D=None,
|
||||
initial_states=states,
|
||||
)
|
||||
|
||||
# just test the last in sequence
|
||||
for i in range(num_examples):
|
||||
# just test one dim and dstate
|
||||
Y_eg = Y[cu_seqlens[i] : cu_seqlens[i + 1], 0, 0]
|
||||
Y_min_eg = Y_min[i][:, 0, 0]
|
||||
torch.testing.assert_close(Y_eg, Y_min_eg, atol=atol, rtol=rtol)
|
||||
|
||||
# update states
|
||||
states = new_states
|
||||
for i, clear in exhausted.items():
|
||||
if clear:
|
||||
states[i].fill_(0.0)
|
||||
exhausted[i] = False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("chunk_size", [8, 256])
|
||||
@pytest.mark.parametrize(
|
||||
"seqlens",
|
||||
[(16, 20), (270, 88, 212, 203)],
|
||||
)
|
||||
def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
|
||||
# This test verifies the correctness of the chunked prefill implementation
|
||||
# in the mamba2 ssd kernels, by comparing concatenation (in the sequence
|
||||
# dimension) of chunked results with the full sequence result.
|
||||
# It is different from test_mamba_chunk_scan_cont_batch by:
|
||||
# 1. Not using the naive torch implementation (ssd_minimal_discrete) to get
|
||||
# reference outputs. Instead, it compares chunked kernel outputs to full
|
||||
# sequence kernel outputs. This is the most straightforward way to
|
||||
# assert chunked prefill correctness.
|
||||
# 2. It focuses on cases where sequences change in the middle of mamba
|
||||
# chunks, and not necessarily on chunk boundaries.
|
||||
|
||||
max_seqlen = max(seqlens)
|
||||
# This test can have larger error for longer sequences
|
||||
if max_seqlen > 256:
|
||||
atol, rtol = 1e-2, 5e-3
|
||||
else:
|
||||
atol, rtol = 5e-3, 5e-3
|
||||
|
||||
num_sequences = len(seqlens)
|
||||
n_heads = 16
|
||||
d_head = 64
|
||||
itype = torch.float32
|
||||
|
||||
# hold state during the cutting process so we know if an
|
||||
# example has been exhausted and needs to cycle
|
||||
last_taken: dict = {} # map: eg -> pointer to last taken sample
|
||||
exhausted: dict = {} # map: eg -> boolean indicating example is exhausted
|
||||
_, cu_seqlens, seq_idx, (A, dt, X, B, C) = next(
|
||||
generate_continuous_batched_examples(
|
||||
[seqlens],
|
||||
num_sequences,
|
||||
max_seqlen,
|
||||
last_taken,
|
||||
exhausted,
|
||||
n_heads,
|
||||
d_head,
|
||||
itype,
|
||||
return_naive_ref=False,
|
||||
)
|
||||
)
|
||||
seqlens = torch.tensor(seqlens, dtype=torch.int32, device=X.device)
|
||||
device = X.device
|
||||
|
||||
## full seqlen computation
|
||||
cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = (
|
||||
compute_varlen_chunk_metadata(cu_seqlens, chunk_size)
|
||||
)
|
||||
Y_ref = torch.empty_like(X)
|
||||
state_ref = mamba_chunk_scan_combined_varlen(
|
||||
X,
|
||||
dt,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
chunk_size,
|
||||
cu_seqlens=cu_seqlens.to(torch.int32),
|
||||
cu_chunk_seqlens=cu_chunk_seqlens,
|
||||
last_chunk_indices=last_chunk_indices,
|
||||
seq_idx=seq_idx_chunks,
|
||||
out=Y_ref,
|
||||
D=None,
|
||||
initial_states=None,
|
||||
)
|
||||
|
||||
## chunked seqlen computation
|
||||
# first chunk
|
||||
chunked_seqlens = seqlens // 2
|
||||
chunked_cu_seqlens = torch.cat(
|
||||
[torch.tensor([0], device=device), torch.cumsum(chunked_seqlens, dim=0)], dim=0
|
||||
)
|
||||
chunked_input_seq_len = chunked_cu_seqlens[-1]
|
||||
X_chunked = torch.zeros_like(X)[:chunked_input_seq_len, ...]
|
||||
dt_chunked = torch.zeros_like(dt)[:chunked_input_seq_len, ...]
|
||||
B_chunked = torch.zeros_like(B)[:chunked_input_seq_len, ...]
|
||||
C_chunked = torch.zeros_like(C)[:chunked_input_seq_len, ...]
|
||||
for i in range(num_sequences):
|
||||
chunk_f = lambda x, i: x[
|
||||
cu_seqlens[i] : cu_seqlens[i] + chunked_seqlens[i], ...
|
||||
]
|
||||
|
||||
X_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f(
|
||||
X, i
|
||||
)
|
||||
dt_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f(
|
||||
dt, i
|
||||
)
|
||||
B_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f(
|
||||
B, i
|
||||
)
|
||||
C_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f(
|
||||
C, i
|
||||
)
|
||||
|
||||
cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = (
|
||||
compute_varlen_chunk_metadata(chunked_cu_seqlens, chunk_size)
|
||||
)
|
||||
Y_partial = torch.empty_like(X_chunked)
|
||||
partial_state = mamba_chunk_scan_combined_varlen(
|
||||
X_chunked,
|
||||
dt_chunked,
|
||||
A,
|
||||
B_chunked,
|
||||
C_chunked,
|
||||
chunk_size,
|
||||
cu_seqlens=chunked_cu_seqlens.to(torch.int32),
|
||||
cu_chunk_seqlens=cu_chunk_seqlens,
|
||||
last_chunk_indices=last_chunk_indices,
|
||||
seq_idx=seq_idx_chunks,
|
||||
out=Y_partial,
|
||||
D=None,
|
||||
initial_states=None,
|
||||
)
|
||||
|
||||
# remaining chunk
|
||||
remaining_chunked_seqlens = seqlens - chunked_seqlens
|
||||
remaining_chunked_cu_seqlens = torch.cat(
|
||||
[
|
||||
torch.tensor([0], device=device),
|
||||
torch.cumsum(remaining_chunked_seqlens, dim=0),
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
remaining_chunked_input_seq_len = remaining_chunked_cu_seqlens[-1]
|
||||
remaining_X_chunked = torch.zeros_like(X)[:remaining_chunked_input_seq_len, ...]
|
||||
remaining_dt_chunked = torch.zeros_like(dt)[:remaining_chunked_input_seq_len, ...]
|
||||
remaining_B_chunked = torch.zeros_like(B)[:remaining_chunked_input_seq_len, ...]
|
||||
remaining_C_chunked = torch.zeros_like(C)[:remaining_chunked_input_seq_len, ...]
|
||||
for i in range(num_sequences):
|
||||
remaining_chunk_f = lambda x, i: x[
|
||||
cu_seqlens[i] + chunked_seqlens[i] : cu_seqlens[i + 1], ...
|
||||
]
|
||||
|
||||
remaining_X_chunked[
|
||||
remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ...
|
||||
] = remaining_chunk_f(X, i)
|
||||
remaining_dt_chunked[
|
||||
remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ...
|
||||
] = remaining_chunk_f(dt, i)
|
||||
remaining_B_chunked[
|
||||
remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ...
|
||||
] = remaining_chunk_f(B, i)
|
||||
remaining_C_chunked[
|
||||
remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ...
|
||||
] = remaining_chunk_f(C, i)
|
||||
|
||||
# assert input chunking is correct
|
||||
concat_chunk_f = lambda pt1, pt2, i: torch.cat(
|
||||
[
|
||||
pt1[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...],
|
||||
pt2[
|
||||
remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1],
|
||||
...,
|
||||
],
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
concat_batch_f = lambda pt1, pt2: torch.cat(
|
||||
[concat_chunk_f(pt1, pt2, i) for i in range(num_sequences)], dim=0
|
||||
)
|
||||
|
||||
assert concat_batch_f(X_chunked, remaining_X_chunked).equal(X)
|
||||
assert concat_batch_f(dt_chunked, remaining_dt_chunked).equal(dt)
|
||||
assert concat_batch_f(B_chunked, remaining_B_chunked).equal(B)
|
||||
assert concat_batch_f(C_chunked, remaining_C_chunked).equal(C)
|
||||
|
||||
cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = (
|
||||
compute_varlen_chunk_metadata(remaining_chunked_cu_seqlens, chunk_size)
|
||||
)
|
||||
|
||||
Y_chunked = torch.empty_like(remaining_X_chunked)
|
||||
state_chunked = mamba_chunk_scan_combined_varlen(
|
||||
remaining_X_chunked,
|
||||
remaining_dt_chunked,
|
||||
A,
|
||||
remaining_B_chunked,
|
||||
remaining_C_chunked,
|
||||
chunk_size,
|
||||
cu_seqlens=remaining_chunked_cu_seqlens.to(torch.int32),
|
||||
cu_chunk_seqlens=cu_chunk_seqlens,
|
||||
last_chunk_indices=last_chunk_indices,
|
||||
seq_idx=seq_idx_chunks,
|
||||
out=Y_chunked,
|
||||
D=None,
|
||||
initial_states=partial_state,
|
||||
)
|
||||
Y = concat_batch_f(Y_partial, Y_chunked)
|
||||
|
||||
# kernel chunked is same as kernel overall
|
||||
for i in range(num_sequences):
|
||||
Y_seq = Y[cu_seqlens[i] : cu_seqlens[i + 1], ...]
|
||||
Y_ref_seq = Y_ref[cu_seqlens[i] : cu_seqlens[i + 1], ...]
|
||||
torch.testing.assert_close(
|
||||
Y_seq[: chunked_seqlens[i], ...],
|
||||
Y_ref_seq[: chunked_seqlens[i], ...],
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
msg=lambda x, i=i: f"seq{i} output part1 " + x,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
Y_seq[chunked_seqlens[i] :, ...],
|
||||
Y_ref_seq[chunked_seqlens[i] :, ...],
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
msg=lambda x, i=i: f"seq{i} output part2 " + x,
|
||||
)
|
||||
|
||||
state_seq = state_chunked[i]
|
||||
state_seq_ref = state_ref[i]
|
||||
torch.testing.assert_close(
|
||||
state_seq,
|
||||
state_seq_ref,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
msg=lambda x, i=i: f"seq{i} state " + x,
|
||||
)
|
||||
@@ -0,0 +1,180 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Equivalence test for ``precopy_mamba_align_fused_kernel``.
|
||||
|
||||
The V2 "align" pre-copy must migrate mamba state across block boundaries with
|
||||
byte-identical semantics to the V1 copy specs (``get_conv_copy_spec`` /
|
||||
``get_temporal_copy_spec``):
|
||||
|
||||
* conv state (SD layout, conv_width > 0): shift the sliding window by
|
||||
``token_bias`` tokens -- ``state[bt[src_col], token_bias:]`` ->
|
||||
``state[bt[dst_col], :conv_width - token_bias]``.
|
||||
* temporal state (conv_width == 0): ``token_bias`` selects the accepted
|
||||
speculative column -- ``state[bt[src_col + token_bias]]`` ->
|
||||
``state[bt[dst_col]]``.
|
||||
|
||||
The kernel must also no-op when ``src_col < 0`` (fresh request) or
|
||||
``src_col == dst_col`` (no boundary crossed).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.worker.mamba_utils import precopy_mamba_align_fused_kernel
|
||||
|
||||
try:
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="precopy_mamba_align_fused_kernel needs CUDA/Triton",
|
||||
)
|
||||
_parametrize = pytest.mark.parametrize
|
||||
except ModuleNotFoundError: # allow running directly as ``python <thisfile>``
|
||||
pytest = None
|
||||
|
||||
def _parametrize(_name, _values):
|
||||
def _deco(fn):
|
||||
return fn
|
||||
|
||||
return _deco
|
||||
|
||||
|
||||
NUM_LAYERS = 3
|
||||
CONV_WIDTH = 4 # conv_kernel - 1 + num_spec
|
||||
CONV_DIM = 96
|
||||
SSM_SHAPE = (4, 16, 16)
|
||||
MAX_COLS = 8
|
||||
|
||||
|
||||
def _build_state(num_blocks, device):
|
||||
"""Per-layer (conv SD [nb, width, dim] bf16, ssm [nb, *shape] fp32) pools."""
|
||||
convs, ssms = [], []
|
||||
for _ in range(NUM_LAYERS):
|
||||
convs.append(
|
||||
torch.randn(
|
||||
num_blocks, CONV_WIDTH, CONV_DIM, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
)
|
||||
ssms.append(
|
||||
torch.randn(num_blocks, *SSM_SHAPE, dtype=torch.float32, device=device)
|
||||
)
|
||||
return convs, ssms
|
||||
|
||||
|
||||
def _build_meta(convs, ssms, device):
|
||||
"""Flattened per-(layer, state-type) metadata, ordered conv, ssm per layer."""
|
||||
n = NUM_LAYERS * 2
|
||||
base = torch.zeros(n, dtype=torch.int64, device=device)
|
||||
blk_stride = torch.zeros(n, dtype=torch.int64, device=device)
|
||||
elem = torch.zeros(n, dtype=torch.int32, device=device)
|
||||
inner = torch.zeros(n, dtype=torch.int64, device=device)
|
||||
width = torch.zeros(n, dtype=torch.int32, device=device)
|
||||
group = torch.zeros(n, dtype=torch.int32, device=device)
|
||||
drc = torch.zeros(n, dtype=torch.int32, device=device) # DS rows (unused, SD)
|
||||
drs = torch.zeros(n, dtype=torch.int64, device=device)
|
||||
i = 0
|
||||
for layer in range(NUM_LAYERS):
|
||||
conv, ssm = convs[layer], ssms[layer]
|
||||
# conv (SD): width = size(1), inner = stride(1)
|
||||
base[i] = conv.data_ptr()
|
||||
blk_stride[i] = conv.stride(0) * conv.element_size()
|
||||
elem[i] = conv.element_size()
|
||||
width[i] = conv.size(1)
|
||||
inner[i] = conv.stride(1)
|
||||
i += 1
|
||||
# ssm (temporal): width = 0, inner = elems per block
|
||||
base[i] = ssm.data_ptr()
|
||||
blk_stride[i] = ssm.stride(0) * ssm.element_size()
|
||||
elem[i] = ssm.element_size()
|
||||
width[i] = 0
|
||||
inner[i] = ssm[0].numel()
|
||||
i += 1
|
||||
return base, blk_stride, elem, inner, width, group, drc, drs
|
||||
|
||||
|
||||
def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs):
|
||||
"""Apply the V1 copy semantics on clones, reading from the pre-copy state."""
|
||||
conv_pre = [c.clone() for c in convs]
|
||||
ssm_pre = [s.clone() for s in ssms]
|
||||
conv_ref = [c.clone() for c in convs]
|
||||
ssm_ref = [s.clone() for s in ssms]
|
||||
for r in range(num_reqs):
|
||||
sc, dc, tb = int(src_col[r]), int(dst_col[r]), int(bias[r])
|
||||
if sc < 0 or sc == dc:
|
||||
continue
|
||||
sblk, dblk = int(bt[r, sc]), int(bt[r, dc])
|
||||
tblk = int(bt[r, sc + tb]) # temporal src column shifted by bias
|
||||
for layer in range(NUM_LAYERS):
|
||||
conv_ref[layer][dblk, : CONV_WIDTH - tb] = conv_pre[layer][sblk, tb:]
|
||||
ssm_ref[layer][dblk] = ssm_pre[layer][tblk]
|
||||
return conv_ref, ssm_ref
|
||||
|
||||
|
||||
@_parametrize("num_reqs", [1, 4, 16])
|
||||
@_parametrize("token_bias", [0, 1, 2])
|
||||
def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(0)
|
||||
# Distinct physical block per (req, col) so copies never alias.
|
||||
num_blocks = num_reqs * MAX_COLS + 1
|
||||
bt = torch.empty(num_reqs, MAX_COLS, dtype=torch.int32, device=device)
|
||||
for r in range(num_reqs):
|
||||
bt[r] = torch.arange(
|
||||
1 + r * MAX_COLS, 1 + (r + 1) * MAX_COLS, dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
# Per-req columns: req 0 fresh (src=-1, skip), req 1 same block (skip),
|
||||
# the rest cross from col 1 -> col 0 with the given spec token bias.
|
||||
src_col = torch.full((num_reqs,), 1, dtype=torch.int32, device=device)
|
||||
dst_col = torch.zeros(num_reqs, dtype=torch.int32, device=device)
|
||||
bias = torch.full((num_reqs,), token_bias, dtype=torch.int32, device=device)
|
||||
if num_reqs >= 1:
|
||||
src_col[0] = -1 # fresh -> no copy
|
||||
if num_reqs >= 2:
|
||||
dst_col[1] = 1 # src_col == dst_col -> no copy
|
||||
|
||||
convs, ssms = _build_state(num_blocks, device)
|
||||
conv_ref, ssm_ref = _reference(
|
||||
convs, ssms, bt.cpu(), src_col.cpu(), dst_col.cpu(), bias.cpu(), num_reqs
|
||||
)
|
||||
|
||||
base, blk_stride, elem, inner, width, group, drc, drs = _build_meta(
|
||||
convs, ssms, device
|
||||
)
|
||||
bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device)
|
||||
idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
|
||||
grid = (num_reqs, NUM_LAYERS * 2)
|
||||
precopy_mamba_align_fused_kernel[grid](
|
||||
dst_col,
|
||||
src_col,
|
||||
bias,
|
||||
bt_ptrs,
|
||||
bt.stride(0),
|
||||
base,
|
||||
blk_stride,
|
||||
elem,
|
||||
inner,
|
||||
width,
|
||||
group,
|
||||
drc,
|
||||
drs,
|
||||
idx_mapping,
|
||||
num_reqs,
|
||||
COPY_BLOCK_SIZE=1024,
|
||||
CONV_STATE_DIM_FIRST=False,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
for layer in range(NUM_LAYERS):
|
||||
torch.testing.assert_close(convs[layer], conv_ref[layer], rtol=0, atol=0)
|
||||
torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
for nr in (1, 4, 16):
|
||||
for tb in (0, 1, 2):
|
||||
test_precopy_matches_v1_copy_specs(nr, tb)
|
||||
print(f"OK num_reqs={nr} token_bias={tb}")
|
||||
@@ -0,0 +1,144 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config.mamba import MambaBackendEnum, MambaConfig
|
||||
from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
|
||||
FlashInferSSUBackend,
|
||||
TritonSSUBackend,
|
||||
get_mamba_ssu_backend,
|
||||
initialize_mamba_ssu_backend,
|
||||
selective_state_update,
|
||||
)
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
|
||||
from vllm.v1.kv_cache_interface import (
|
||||
KVCacheConfig,
|
||||
KVCacheGroupSpec,
|
||||
MambaSpec,
|
||||
)
|
||||
|
||||
try:
|
||||
import flashinfer.mamba # noqa: F401
|
||||
|
||||
HAS_FLASHINFER = True
|
||||
except ImportError:
|
||||
HAS_FLASHINFER = False
|
||||
|
||||
|
||||
def _kv_cache_config_with_ssu(
|
||||
mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2,
|
||||
) -> KVCacheConfig:
|
||||
spec = MambaSpec(
|
||||
block_size=16,
|
||||
shapes=((16, 64),),
|
||||
dtypes=(torch.float16,),
|
||||
mamba_type=mamba_type,
|
||||
)
|
||||
return KVCacheConfig(
|
||||
num_blocks=1,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[KVCacheGroupSpec(layer_names=["l0"], kv_cache_spec=spec)],
|
||||
)
|
||||
|
||||
|
||||
def test_default_backend_is_triton():
|
||||
initialize_mamba_ssu_backend(MambaConfig(), _kv_cache_config_with_ssu())
|
||||
backend = get_mamba_ssu_backend()
|
||||
assert isinstance(backend, TritonSSUBackend)
|
||||
assert backend.name == "triton"
|
||||
|
||||
|
||||
def test_explicit_triton_backend():
|
||||
initialize_mamba_ssu_backend(
|
||||
MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
|
||||
)
|
||||
backend = get_mamba_ssu_backend()
|
||||
assert isinstance(backend, TritonSSUBackend)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
|
||||
def test_flashinfer_backend_init():
|
||||
initialize_mamba_ssu_backend(
|
||||
MambaConfig(backend=MambaBackendEnum.FLASHINFER), _kv_cache_config_with_ssu()
|
||||
)
|
||||
backend = get_mamba_ssu_backend()
|
||||
assert isinstance(backend, FlashInferSSUBackend)
|
||||
assert backend.name == "flashinfer"
|
||||
|
||||
|
||||
def test_uninitialized_backend_raises():
|
||||
import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
|
||||
|
||||
old = mod._mamba_ssu_backend
|
||||
mod._mamba_ssu_backend = None
|
||||
with pytest.raises(RuntimeError, match="not been initialized"):
|
||||
get_mamba_ssu_backend()
|
||||
mod._mamba_ssu_backend = old
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mamba_type",
|
||||
[
|
||||
MambaAttentionBackendEnum.LINEAR,
|
||||
MambaAttentionBackendEnum.GDN_ATTN,
|
||||
MambaAttentionBackendEnum.SHORT_CONV,
|
||||
],
|
||||
)
|
||||
def test_init_is_noop_for_non_ssu_mamba_type(mamba_type):
|
||||
import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
|
||||
|
||||
old = mod._mamba_ssu_backend
|
||||
mod._mamba_ssu_backend = None
|
||||
try:
|
||||
initialize_mamba_ssu_backend(
|
||||
MambaConfig(), _kv_cache_config_with_ssu(mamba_type)
|
||||
)
|
||||
assert mod._mamba_ssu_backend is None
|
||||
with pytest.raises(RuntimeError, match="not been initialized"):
|
||||
get_mamba_ssu_backend()
|
||||
finally:
|
||||
mod._mamba_ssu_backend = old
|
||||
|
||||
|
||||
@pytest.mark.skipif(HAS_FLASHINFER, reason="flashinfer is installed")
|
||||
def test_flashinfer_import_error():
|
||||
with pytest.raises(ImportError, match="FlashInfer is required"):
|
||||
FlashInferSSUBackend(MambaConfig())
|
||||
|
||||
|
||||
def test_triton_basic_call():
|
||||
set_random_seed(0)
|
||||
initialize_mamba_ssu_backend(
|
||||
MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
|
||||
)
|
||||
device = "cuda"
|
||||
batch_size = 2
|
||||
dim = 64
|
||||
dstate = 16
|
||||
|
||||
state = torch.randn(batch_size, dim, dstate, device=device)
|
||||
x = torch.randn(batch_size, dim, device=device)
|
||||
out = torch.empty_like(x)
|
||||
dt = torch.randn(batch_size, dim, device=device)
|
||||
dt_bias = torch.rand(dim, device=device) - 4.0
|
||||
A = -torch.rand(dim, dstate, device=device)
|
||||
B = torch.randn(batch_size, dstate, device=device)
|
||||
C = torch.randn(batch_size, dstate, device=device)
|
||||
D = torch.randn(dim, device=device)
|
||||
|
||||
selective_state_update(
|
||||
state,
|
||||
x,
|
||||
dt,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
D=D,
|
||||
dt_bias=dt_bias,
|
||||
dt_softplus=True,
|
||||
out=out,
|
||||
)
|
||||
assert not torch.isnan(out).any()
|
||||
@@ -0,0 +1,78 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange, repeat
|
||||
|
||||
|
||||
def selective_state_update_ref(
|
||||
state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False
|
||||
):
|
||||
"""
|
||||
Argument:
|
||||
state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
|
||||
x: (batch, dim) or (batch, nheads, dim)
|
||||
dt: (batch, dim) or (batch, nheads, dim)
|
||||
A: (dim, dstate) or (nheads, dim, dstate)
|
||||
B: (batch, dstate) or (batch, ngroups, dstate)
|
||||
C: (batch, dstate) or (batch, ngroups, dstate)
|
||||
D: (dim,) or (nheads, dim)
|
||||
z: (batch, dim) or (batch, nheads, dim)
|
||||
dt_bias: (dim,) or (nheads, dim)
|
||||
Return:
|
||||
out: (batch, dim) or (batch, nheads, dim)
|
||||
"""
|
||||
has_heads = state.dim() > 3
|
||||
if state.dim() == 3:
|
||||
state = state.unsqueeze(1)
|
||||
if x.dim() == 2:
|
||||
x = x.unsqueeze(1)
|
||||
if dt.dim() == 2:
|
||||
dt = dt.unsqueeze(1)
|
||||
if A.dim() == 2:
|
||||
A = A.unsqueeze(0)
|
||||
if B.dim() == 2:
|
||||
B = B.unsqueeze(1)
|
||||
if C.dim() == 2:
|
||||
C = C.unsqueeze(1)
|
||||
if D is not None and D.dim() == 1:
|
||||
D = D.unsqueeze(0)
|
||||
if z is not None and z.dim() == 2:
|
||||
z = z.unsqueeze(1)
|
||||
if dt_bias is not None and dt_bias.dim() == 1:
|
||||
dt_bias = dt_bias.unsqueeze(0)
|
||||
batch, nheads, dim, dstate = state.shape
|
||||
assert x.shape == (batch, nheads, dim)
|
||||
assert dt.shape == x.shape
|
||||
assert A.shape == (nheads, dim, dstate)
|
||||
ngroups = B.shape[1]
|
||||
assert nheads % ngroups == 0, "nheads must be divisible by ngroups"
|
||||
assert B.shape == (batch, ngroups, dstate)
|
||||
assert C.shape == B.shape
|
||||
if D is not None:
|
||||
assert D.shape == (nheads, dim)
|
||||
if z is not None:
|
||||
assert z.shape == x.shape
|
||||
if dt_bias is not None:
|
||||
assert dt_bias.shape == (nheads, dim)
|
||||
dt = dt + dt_bias
|
||||
dt = F.softplus(dt) if dt_softplus else dt
|
||||
dA = torch.exp(
|
||||
rearrange(dt, "b h d -> b h d 1") * A
|
||||
) # (batch, nheads, dim, dstate)
|
||||
B = repeat(B, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate)
|
||||
C = repeat(C, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate)
|
||||
dB = rearrange(dt, "b h d -> b h d 1") * rearrange(
|
||||
B, "b h n -> b h 1 n"
|
||||
) # (batch, nheads, dim, dstate)
|
||||
state.copy_(
|
||||
state * dA + dB * rearrange(x, "b h d -> b h d 1")
|
||||
) # (batch, dim, dstate
|
||||
out = torch.einsum("bhdn,bhn->bhd", state.to(C.dtype), C)
|
||||
if D is not None:
|
||||
out += (x * D).to(out.dtype)
|
||||
out = (out if z is None else out * F.silu(z)).to(x.dtype)
|
||||
if not has_heads:
|
||||
out = out.squeeze(1)
|
||||
return out
|
||||
@@ -0,0 +1,14 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
|
||||
def pytest_addoption(parser):
|
||||
parser.addoption(
|
||||
"--subtests", action="store", type=str, default=None, help="subtest ids"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def subtests(request):
|
||||
return request.config.getoption("--subtests")
|
||||
@@ -0,0 +1,158 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
|
||||
|
||||
from .common import Config
|
||||
from .mk_objects import (
|
||||
MK_ALL_PREPARE_FINALIZE_TYPES,
|
||||
MK_FUSED_EXPERT_TYPES,
|
||||
MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES,
|
||||
)
|
||||
|
||||
|
||||
def make_config_arg_parser(description: str):
|
||||
def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalizeModular:
|
||||
for pf in MK_ALL_PREPARE_FINALIZE_TYPES:
|
||||
if pf.__name__ == s:
|
||||
return pf
|
||||
raise ValueError(f"Cannot find a PrepareFinalize type that matches {s}")
|
||||
|
||||
def to_experts_class_type(s: str) -> mk.FusedMoEExpertsModular:
|
||||
for fe in MK_FUSED_EXPERT_TYPES:
|
||||
if fe.__name__ == s:
|
||||
return fe
|
||||
raise ValueError(f"Cannot find a FusedExperts type that matches {s}")
|
||||
|
||||
def to_quant_torch_dtype(s: str) -> torch.dtype:
|
||||
if s == "torch.float8_e4m3fn":
|
||||
return torch.float8_e4m3fn
|
||||
raise ValueError(f"Unsupported quant type {s}")
|
||||
|
||||
parser = argparse.ArgumentParser(description=description)
|
||||
|
||||
parser.add_argument(
|
||||
"--world-size",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Number of ranks that participate in all2all",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pf-type",
|
||||
type=to_pf_class_type,
|
||||
required=True,
|
||||
help=(
|
||||
"Choose a PrepareFinalize Type : "
|
||||
f"{[x.__name__ for x in MK_ALL_PREPARE_FINALIZE_TYPES]}"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--experts-type",
|
||||
type=to_experts_class_type,
|
||||
required=True,
|
||||
help=(
|
||||
f"Choose a FusedExpert type : {[x.__name__ for x in MK_FUSED_EXPERT_TYPES]}"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"-m",
|
||||
nargs="+",
|
||||
type=int,
|
||||
default=[64],
|
||||
help="num tokens per rank",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-k",
|
||||
type=int,
|
||||
default=7168,
|
||||
help="hidden-size",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-n",
|
||||
type=int,
|
||||
default=1024,
|
||||
help="N dimension of the first fused-moe matmul",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-experts", type=int, default=32, help="Global num experts"
|
||||
)
|
||||
parser.add_argument("--topk", nargs="+", type=int, default=[4, 1], help="num topk")
|
||||
|
||||
# Quant args
|
||||
parser.add_argument(
|
||||
"--quant-dtype", type=to_quant_torch_dtype, help="Quant datatype"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--per-token-quantized-activations",
|
||||
action="store_true",
|
||||
help=("The input activations must be per-token quantized"),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--per-channel-quantized-weights",
|
||||
action="store_true",
|
||||
help="The weights must be per-channel quantized.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--block-shape", nargs="+", type=int, help="Quantization block shape"
|
||||
)
|
||||
|
||||
# Torch trace profile generation args
|
||||
parser.add_argument(
|
||||
"--torch-trace-dir-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Get torch trace for single execution",
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def _validate_args(args: argparse.Namespace):
|
||||
if args.quant_dtype is not None:
|
||||
assert args.quant_dtype == torch.float8_e4m3fn
|
||||
if args.block_shape is not None:
|
||||
assert len(args.block_shape) == 2, (
|
||||
f"block shape must have 2 elements. got {args.block_shape}"
|
||||
)
|
||||
|
||||
if args.experts_type in MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES:
|
||||
assert args.world_size == 1, "Single GPU objects need world size set to 1"
|
||||
|
||||
if args.torch_trace_dir_path is not None:
|
||||
from pathlib import Path
|
||||
|
||||
assert Path(args.torch_trace_dir_path).is_dir(), (
|
||||
f"Please create {args.torch_trace_dir_path}"
|
||||
)
|
||||
|
||||
|
||||
def make_config(args: argparse.Namespace) -> Config:
|
||||
_validate_args(args)
|
||||
|
||||
quant_config = None
|
||||
if args.quant_dtype is not None:
|
||||
quant_config = FusedMoEQuantConfig.make(
|
||||
quant_dtype=args.quant_dtype,
|
||||
per_act_token_quant=args.per_token_quantized_activations,
|
||||
per_out_ch_quant=args.per_channel_quantized_weights,
|
||||
block_shape=args.block_shape,
|
||||
)
|
||||
|
||||
return Config(
|
||||
Ms=args.m,
|
||||
K=args.k,
|
||||
N=args.n,
|
||||
E=args.num_experts,
|
||||
topks=args.topk,
|
||||
dtype=torch.bfloat16, # hard-code
|
||||
quant_config=quant_config,
|
||||
prepare_finalize_type=args.pf_type,
|
||||
fused_experts_type=args.experts_type,
|
||||
world_size=args.world_size,
|
||||
torch_trace_dir_path=args.torch_trace_dir_path,
|
||||
)
|
||||
@@ -0,0 +1,795 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from tests.kernels.moe.utils import make_test_weights, per_token_cast_to_fp8
|
||||
from tests.kernels.quantization.nvfp4_utils import (
|
||||
FLOAT4_E2M1_MAX,
|
||||
FLOAT8_E4M3_MAX,
|
||||
dequantize_nvfp4_to_dtype,
|
||||
)
|
||||
from tests.kernels.utils import torch_experts
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.distributed import (
|
||||
get_dp_group,
|
||||
get_pcp_group,
|
||||
get_tensor_model_parallel_world_size,
|
||||
)
|
||||
from vllm.forward_context import set_forward_context
|
||||
from vllm.model_executor.layers.fused_moe import fused_topk
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.all2all_utils import (
|
||||
maybe_make_prepare_finalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEParallelConfig,
|
||||
FusedMoEQuantConfig,
|
||||
RoutingMethodType,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
kFp8Dynamic128Sym,
|
||||
kFp8DynamicTensorSym,
|
||||
kFp8DynamicTokenSym,
|
||||
kFp8Static128BlockSym,
|
||||
kFp8StaticChannelSym,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.utils.import_utils import (
|
||||
has_aiter,
|
||||
has_deep_ep,
|
||||
has_deep_ep_v2,
|
||||
has_deep_gemm,
|
||||
has_mori,
|
||||
)
|
||||
from vllm.utils.math_utils import next_power_of_2
|
||||
|
||||
from .mk_objects import (
|
||||
TestMoEQuantConfig,
|
||||
expert_info,
|
||||
make_fused_experts,
|
||||
prepare_finalize_info,
|
||||
)
|
||||
from .parallel_utils import ProcessGroupInfo
|
||||
|
||||
|
||||
def _describe_tensor(t: torch.Tensor | None, name: str) -> str:
|
||||
if t is None:
|
||||
return f"{name} : None"
|
||||
else:
|
||||
return f"{name} : {t.shape} {t.dtype} {t.device}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
Ms: list[int] | int
|
||||
K: int
|
||||
N: int
|
||||
E: int
|
||||
topks: list[int] | int
|
||||
dtype: torch.dtype
|
||||
quant_config: TestMoEQuantConfig | None
|
||||
|
||||
prepare_finalize_type: mk.FusedMoEPrepareAndFinalize
|
||||
fused_experts_type: mk.FusedMoEExperts
|
||||
|
||||
world_size: int
|
||||
|
||||
torch_trace_dir_path: str | None = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.quant_config is None:
|
||||
self.quant_config = TestMoEQuantConfig(None, False, False, None)
|
||||
|
||||
def describe(self) -> str:
|
||||
s = ""
|
||||
s += "== Config:\n"
|
||||
s += f" world_size={self.world_size}\n"
|
||||
s += f" PF={self.prepare_finalize_type.__name__}\n"
|
||||
s += f" FE={self.fused_experts_type.__name__}\n"
|
||||
s += f" E={self.E}\n"
|
||||
s += f" Ms={self.Ms}\n"
|
||||
s += f" N={self.N}\n"
|
||||
s += f" K={self.K}\n"
|
||||
s += f" topk={self.topks}\n"
|
||||
s += f" dtype={self.dtype}\n"
|
||||
s += " Quant:\n"
|
||||
if self.quant_config is not None:
|
||||
s += f" q_dtype={self.quant_dtype}\n"
|
||||
s += f" q_block_shape={self.quant_block_shape}\n"
|
||||
s += f" q_per_out_ch_quant={self.is_per_out_ch_quant}\n"
|
||||
s += f" q_per_act_token={self.is_per_act_token_quant}\n"
|
||||
else:
|
||||
s += " quant=None\n"
|
||||
return s
|
||||
|
||||
@property
|
||||
def M(self) -> int:
|
||||
assert isinstance(self.Ms, int)
|
||||
return self.Ms
|
||||
|
||||
@property
|
||||
def quant_dtype(self) -> torch.dtype | str | None:
|
||||
assert self.quant_config is not None
|
||||
return self.quant_config.quant_dtype
|
||||
|
||||
@property
|
||||
def is_per_act_token_quant(self) -> bool:
|
||||
assert self.quant_config is not None
|
||||
return self.quant_config.per_act_token_quant
|
||||
|
||||
@property
|
||||
def is_per_tensor_act_quant(self) -> bool:
|
||||
return not self.is_per_act_token_quant and self.quant_block_shape is None
|
||||
|
||||
@property
|
||||
def is_per_out_ch_quant(self) -> bool:
|
||||
assert self.quant_config is not None
|
||||
return self.quant_config.per_out_ch_quant
|
||||
|
||||
@property
|
||||
def quant_block_shape(self) -> list[int] | None:
|
||||
assert self.quant_config is not None
|
||||
return self.quant_config.block_shape
|
||||
|
||||
@property
|
||||
def topk(self) -> int:
|
||||
assert isinstance(self.topks, int)
|
||||
return self.topks
|
||||
|
||||
@property
|
||||
def num_local_experts(self) -> int:
|
||||
return self.E // self.world_size
|
||||
|
||||
def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]:
|
||||
"""
|
||||
make env data for vllm launch.
|
||||
"""
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(
|
||||
enforce_eager=True,
|
||||
is_moe=True,
|
||||
)
|
||||
vllm_config.parallel_config.data_parallel_size = self.world_size
|
||||
vllm_config.parallel_config.enable_expert_parallel = True
|
||||
|
||||
env_dict = {
|
||||
"VLLM_USE_DEEP_GEMM": str(int(self.needs_deep_gemm())),
|
||||
}
|
||||
|
||||
vllm_config.parallel_config.all2all_backend = self.all2all_backend()
|
||||
|
||||
return vllm_config, env_dict
|
||||
|
||||
def fe_supports_quant_scheme(self) -> bool:
|
||||
"""Check if the fused experts class supports this quant config.
|
||||
See https://github.com/ROCm/aiter/issues/2419 for AITER gaps."""
|
||||
if self.quant_config is None or self.quant_dtype is None:
|
||||
return True
|
||||
if self.quant_dtype != torch.float8_e4m3fn:
|
||||
return True
|
||||
# Derive QuantKeys from test config
|
||||
if self.quant_block_shape is not None:
|
||||
w_key = kFp8Static128BlockSym
|
||||
a_key = kFp8Dynamic128Sym
|
||||
elif self.is_per_out_ch_quant:
|
||||
w_key = kFp8StaticChannelSym
|
||||
a_key = (
|
||||
kFp8DynamicTokenSym
|
||||
if self.is_per_act_token_quant
|
||||
else kFp8StaticTensorSym
|
||||
)
|
||||
else:
|
||||
w_key = kFp8StaticTensorSym
|
||||
a_key = (
|
||||
kFp8DynamicTensorSym
|
||||
if self.is_per_act_token_quant
|
||||
else kFp8StaticTensorSym
|
||||
)
|
||||
fe_cls = self.fused_experts_type
|
||||
if hasattr(fe_cls, "_supports_quant_scheme"):
|
||||
try:
|
||||
return fe_cls._supports_quant_scheme(w_key, a_key)
|
||||
except NotImplementedError:
|
||||
pass
|
||||
return True
|
||||
|
||||
def is_fp8_block_quantized(self):
|
||||
return (
|
||||
self.quant_dtype == torch.float8_e4m3fn
|
||||
and self.quant_block_shape is not None
|
||||
)
|
||||
|
||||
def is_batched_prepare_finalize(self):
|
||||
info = prepare_finalize_info(self.prepare_finalize_type)
|
||||
return mk.FusedMoEActivationFormat.BatchedExperts == info.activation_format
|
||||
|
||||
def is_batched_fused_experts(self):
|
||||
info = expert_info(self.fused_experts_type)
|
||||
return mk.FusedMoEActivationFormat.BatchedExperts == info.activation_format
|
||||
|
||||
def is_standard_fused_experts(self):
|
||||
info = expert_info(self.fused_experts_type)
|
||||
return mk.FusedMoEActivationFormat.Standard == info.activation_format
|
||||
|
||||
def fe_supported_types(self):
|
||||
info = expert_info(self.fused_experts_type)
|
||||
return info.supported_dtypes
|
||||
|
||||
def pf_supported_types(self):
|
||||
info = prepare_finalize_info(self.prepare_finalize_type)
|
||||
return info.supported_dtypes
|
||||
|
||||
def is_block_quant_supported(self):
|
||||
info = expert_info(self.fused_experts_type)
|
||||
return info.blocked_quantization_support
|
||||
|
||||
def supports_apply_weight_on_input(self):
|
||||
info = prepare_finalize_info(self.prepare_finalize_type)
|
||||
return info.supports_apply_weight_on_input
|
||||
|
||||
def needs_deep_gemm(self):
|
||||
info = expert_info(self.fused_experts_type)
|
||||
return info.needs_deep_gemm
|
||||
|
||||
def needs_deep_ep(self):
|
||||
info = prepare_finalize_info(self.prepare_finalize_type)
|
||||
return (
|
||||
info.backend == "deepep_high_throughput"
|
||||
or info.backend == "deepep_low_latency"
|
||||
)
|
||||
|
||||
def needs_deep_ep_v2(self):
|
||||
info = prepare_finalize_info(self.prepare_finalize_type)
|
||||
return info.backend == "deepep_v2"
|
||||
|
||||
def needs_aiter(self):
|
||||
info = expert_info(self.fused_experts_type)
|
||||
return info.needs_aiter
|
||||
|
||||
def needs_mori(self):
|
||||
info = prepare_finalize_info(self.prepare_finalize_type)
|
||||
return info.backend in ("mori_high_throughput", "mori_low_latency")
|
||||
|
||||
def all2all_backend(self):
|
||||
info = prepare_finalize_info(self.prepare_finalize_type)
|
||||
return info.backend
|
||||
|
||||
def is_valid(self) -> tuple[bool, str | None]:
|
||||
# Check prepare-finalize and fused-experts compatibility
|
||||
if self.is_batched_prepare_finalize():
|
||||
if not self.is_batched_fused_experts():
|
||||
return False, "Mismatched format."
|
||||
else:
|
||||
if not self.is_standard_fused_experts():
|
||||
return False, "Mismatched format."
|
||||
|
||||
# Check quantization sanity
|
||||
if (
|
||||
int(self.is_per_act_token_quant)
|
||||
+ int(self.is_per_tensor_act_quant)
|
||||
+ int(self.quant_block_shape is not None)
|
||||
) > 1:
|
||||
# invalid quant config
|
||||
return False, f"Bad quant_config {self.quant_config}."
|
||||
|
||||
# check type support
|
||||
if self.quant_dtype is None:
|
||||
if (
|
||||
self.dtype not in self.pf_supported_types()
|
||||
or self.dtype not in self.fe_supported_types()
|
||||
):
|
||||
return False, (
|
||||
f"Unsupported type {self.dtype} not in "
|
||||
f"{self.pf_supported_types()} and "
|
||||
f"{self.fe_supported_types()}."
|
||||
)
|
||||
else:
|
||||
if (
|
||||
self.quant_dtype not in self.pf_supported_types()
|
||||
or self.quant_dtype not in self.fe_supported_types()
|
||||
):
|
||||
return False, (
|
||||
f"Unsupported quant type {self.quant_dtype} "
|
||||
f"not in {self.pf_supported_types()} and "
|
||||
f"{self.fe_supported_types()}."
|
||||
)
|
||||
|
||||
# Check quant scheme compatibility with fused experts class
|
||||
if not self.fe_supports_quant_scheme():
|
||||
return False, (
|
||||
f"FE {self.fused_experts_type.__name__} does not support "
|
||||
f"quant scheme (per_out_ch={self.is_per_out_ch_quant}, "
|
||||
f"per_act_token={self.is_per_act_token_quant}, "
|
||||
f"block={self.quant_block_shape})"
|
||||
)
|
||||
|
||||
# Check block quantization support
|
||||
is_block_quantized = self.quant_block_shape is not None
|
||||
if is_block_quantized and self.quant_dtype is None:
|
||||
return False, "No block quantization support."
|
||||
|
||||
if is_block_quantized and not self.is_block_quant_supported():
|
||||
return False, "Mismatched block quantization support."
|
||||
|
||||
# deep_gemm only works with block-quantized
|
||||
if self.needs_deep_gemm() and not is_block_quantized:
|
||||
return False, "Needs DeepGEMM but not block quantized."
|
||||
|
||||
# Check dependencies (turn into asserts?)
|
||||
if self.needs_deep_ep() and not has_deep_ep():
|
||||
return False, "Needs DeepEP, but DeepEP not available."
|
||||
if self.needs_deep_ep_v2() and not has_deep_ep_v2():
|
||||
return False, "Needs DeepEP v2, but DeepEP v2 not available."
|
||||
if self.needs_deep_gemm() and not has_deep_gemm():
|
||||
return (
|
||||
False,
|
||||
"Needs DeepGEMM, but the current vLLM environment does not provide it.",
|
||||
)
|
||||
if self.needs_aiter() and not has_aiter(): # noqa: SIM103
|
||||
return False, "Needs Aiter, but Aiter not available."
|
||||
if self.needs_mori() and not has_mori(): # noqa: SIM103
|
||||
return False, "Needs MoRI, but MoRI not available."
|
||||
|
||||
try:
|
||||
if not self.fused_experts_type._supports_current_device():
|
||||
return (
|
||||
False,
|
||||
f"{self.fused_experts_type} not supported on the current device.",
|
||||
)
|
||||
except NotImplementedError:
|
||||
pass
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
@dataclass
|
||||
class WeightTensors:
|
||||
w1: torch.Tensor
|
||||
w2: torch.Tensor
|
||||
w1_scale: torch.Tensor | None
|
||||
w2_scale: torch.Tensor | None
|
||||
w1_gs: torch.Tensor | None = None
|
||||
w2_gs: torch.Tensor | None = None
|
||||
|
||||
def describe(self):
|
||||
s = ""
|
||||
s += "== Weight Tensors: \n"
|
||||
s += f" - {_describe_tensor(self.w1, 'w1')} \n"
|
||||
s += f" - {_describe_tensor(self.w2, 'w2')} \n"
|
||||
s += f" - {_describe_tensor(self.w1_scale, 'w1_scale')} \n"
|
||||
s += f" - {_describe_tensor(self.w2_scale, 'w2_scale')} \n"
|
||||
s += f" - {_describe_tensor(self.w1_gs, 'w1_gs')} \n"
|
||||
s += f" - {_describe_tensor(self.w2_gs, 'w2_gs')} \n"
|
||||
return s
|
||||
|
||||
def is_quantized(self) -> bool:
|
||||
# or w1_scale is not None?
|
||||
return (
|
||||
self.w1.dtype == torch.float8_e4m3fn
|
||||
or self.w1.dtype == torch.uint8
|
||||
or self.w1.dtype == torch.int8
|
||||
)
|
||||
|
||||
def to_current_device(self):
|
||||
device = torch.accelerator.current_device_index()
|
||||
self.w1 = self.w1.to(device=device)
|
||||
self.w2 = self.w2.to(device=device)
|
||||
|
||||
if self.w1_scale is not None:
|
||||
self.w1_scale = self.w1_scale.to(device=device)
|
||||
if self.w2_scale is not None:
|
||||
self.w2_scale = self.w2_scale.to(device=device)
|
||||
|
||||
if self.w1_gs is not None:
|
||||
self.w1_gs = self.w1_gs.to(device=device)
|
||||
if self.w2_gs is not None:
|
||||
self.w2_gs = self.w2_gs.to(device=device)
|
||||
|
||||
def slice_weights(self, rank: int, num_local_experts: int) -> "WeightTensors":
|
||||
s = rank * num_local_experts
|
||||
e = s + num_local_experts
|
||||
w1 = self.w1[s:e, :, :]
|
||||
w2 = self.w2[s:e, :, :]
|
||||
w1_scale = self.w1_scale[s:e, :, :] if self.w1_scale is not None else None
|
||||
w2_scale = self.w2_scale[s:e, :, :] if self.w2_scale is not None else None
|
||||
w1_gs = self.w1_gs[s:e] if self.w1_gs is not None else None
|
||||
w2_gs = self.w2_gs[s:e] if self.w2_gs is not None else None
|
||||
|
||||
return WeightTensors(w1, w2, w1_scale, w2_scale, w1_gs, w2_gs)
|
||||
|
||||
@staticmethod
|
||||
def make(config: Config) -> "WeightTensors":
|
||||
(_, w1, w1_scale, w1_gs), (_, w2, w2_scale, w2_gs) = make_test_weights(
|
||||
e=config.E,
|
||||
n=config.N,
|
||||
k=config.K,
|
||||
in_dtype=config.dtype,
|
||||
quant_dtype=config.quant_dtype,
|
||||
block_shape=config.quant_block_shape,
|
||||
# or config.is_per_out_ch_quant
|
||||
per_out_ch_quant=config.is_per_act_token_quant,
|
||||
)
|
||||
return WeightTensors(
|
||||
w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale, w1_gs=w1_gs, w2_gs=w2_gs
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RankTensors:
|
||||
hidden_states: torch.Tensor
|
||||
hidden_states_scale: torch.Tensor | None
|
||||
|
||||
topk_weights: torch.Tensor
|
||||
topk_ids: torch.Tensor
|
||||
expert_map: torch.Tensor | None
|
||||
|
||||
def describe(self):
|
||||
s = ""
|
||||
s += "== Rank Tensors: \n"
|
||||
s += f" - {_describe_tensor(self.hidden_states, 'HS')} \n"
|
||||
s += f" - {_describe_tensor(self.hidden_states_scale, 'HS_scale')} \n"
|
||||
s += f" - {_describe_tensor(self.topk_weights, 'topk_weights')} \n"
|
||||
s += f" - {_describe_tensor(self.topk_ids, 'topk_ids')} \n"
|
||||
s += f" - {_describe_tensor(self.expert_map, 'expert_map')} \n"
|
||||
return s
|
||||
|
||||
@staticmethod
|
||||
def make_hidden_states(
|
||||
config: Config,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
"""
|
||||
Return hidden_states
|
||||
"""
|
||||
m, k, dtype = (config.M, config.K, config.dtype)
|
||||
device = torch.accelerator.current_device_index()
|
||||
a = torch.randn((m, k), device=device, dtype=dtype) / 15.0
|
||||
|
||||
if config.quant_dtype is None:
|
||||
return a, None
|
||||
|
||||
# We dequant and use that as hidden_states so the tests are stable.
|
||||
# quantizing and dequantizing yield slightly different results
|
||||
# depending on the hardware. Here we, quantize and dequantize
|
||||
# first - so further quantize and dequantize will yield the same
|
||||
# values.
|
||||
if config.is_per_tensor_act_quant:
|
||||
a_q, a_scales = ops.scaled_fp8_quant(a, use_per_token_if_dynamic=False)
|
||||
return a_q.float().mul(a_scales).to(dtype), a_scales
|
||||
|
||||
if config.is_per_act_token_quant:
|
||||
a_q, a_scales = ops.scaled_fp8_quant(a, use_per_token_if_dynamic=True)
|
||||
return a_q.float().mul(a_scales).to(dtype), None
|
||||
|
||||
assert config.quant_block_shape is not None
|
||||
block_k = config.quant_block_shape[1]
|
||||
a_q, a_scales = per_token_cast_to_fp8(a, block_size=block_k)
|
||||
return a_q.float().view((-1, block_k)).mul(a_scales.view(-1, 1)).view(m, k).to(
|
||||
dtype
|
||||
), None
|
||||
|
||||
@staticmethod
|
||||
def make(config: Config, pgi: ProcessGroupInfo):
|
||||
dtype = config.dtype
|
||||
topk, m, _ = (config.topk, config.M, config.K)
|
||||
hidden_states, hidden_states_scale = RankTensors.make_hidden_states(config)
|
||||
|
||||
num_local_experts, global_num_experts = (config.num_local_experts, config.E)
|
||||
score = torch.randn((m, global_num_experts), device="cuda", dtype=dtype)
|
||||
topk_weights, topk_ids, _ = fused_topk(hidden_states, score, topk, False)
|
||||
|
||||
# distribute topk_ids evenly
|
||||
device = torch.accelerator.current_device_index()
|
||||
for mi in range(m):
|
||||
topk_ids[mi] = torch.randperm(config.E)[:topk]
|
||||
topk_ids = topk_ids.to(device=device)
|
||||
|
||||
expert_map = None
|
||||
if config.world_size > 1:
|
||||
expert_map = torch.full(
|
||||
(global_num_experts,), fill_value=-1, dtype=torch.int32
|
||||
)
|
||||
s = pgi.rank * num_local_experts
|
||||
e = s + num_local_experts
|
||||
expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
|
||||
expert_map = expert_map.to(device=device, dtype=torch.int32)
|
||||
|
||||
return RankTensors(
|
||||
hidden_states=hidden_states,
|
||||
hidden_states_scale=hidden_states_scale,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
expert_map=expert_map,
|
||||
)
|
||||
|
||||
|
||||
def reference_moe_impl(
|
||||
config: Config, weights: WeightTensors, rank_tensors: RankTensors
|
||||
) -> torch.Tensor:
|
||||
if config.quant_dtype == "nvfp4":
|
||||
quant_blocksize = 16
|
||||
dtype = config.dtype
|
||||
|
||||
w1_q = weights.w1
|
||||
w1_blockscale = weights.w1_scale
|
||||
w1_gs = weights.w1_gs
|
||||
|
||||
w2_q = weights.w2
|
||||
w2_blockscale = weights.w2_scale
|
||||
w2_gs = weights.w2_gs
|
||||
|
||||
a_global_scale = (
|
||||
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX)
|
||||
/ torch.amax(rank_tensors.hidden_states.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
|
||||
assert w1_gs is not None
|
||||
assert w2_gs is not None
|
||||
assert w1_blockscale is not None
|
||||
assert w2_blockscale is not None
|
||||
|
||||
assert w1_blockscale.shape[1] % 128 == 0
|
||||
assert w1_blockscale.shape[2] % 4 == 0
|
||||
assert w2_blockscale.shape[1] % 128 == 0
|
||||
assert w2_blockscale.shape[2] % 4 == 0
|
||||
|
||||
a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(
|
||||
rank_tensors.hidden_states, a_global_scale
|
||||
)
|
||||
|
||||
a = dequantize_nvfp4_to_dtype(
|
||||
a_fp4,
|
||||
a_scale_interleaved,
|
||||
a_global_scale,
|
||||
dtype=dtype,
|
||||
device=a_fp4.device,
|
||||
block_size=quant_blocksize,
|
||||
)
|
||||
|
||||
e = w1_q.shape[0]
|
||||
n = w1_q.shape[1] // 2
|
||||
k = w2_q.shape[1]
|
||||
|
||||
w1 = torch.zeros((e, 2 * n, k), device="cuda", dtype=dtype)
|
||||
w2 = torch.zeros((e, k, n), device="cuda", dtype=dtype)
|
||||
|
||||
for idx in range(0, e):
|
||||
w1[idx] = dequantize_nvfp4_to_dtype(
|
||||
w1_q[idx],
|
||||
w1_blockscale[idx],
|
||||
w1_gs[idx],
|
||||
dtype=dtype,
|
||||
device=w1_q.device,
|
||||
block_size=quant_blocksize,
|
||||
)
|
||||
w2[idx] = dequantize_nvfp4_to_dtype(
|
||||
w2_q[idx],
|
||||
w2_blockscale[idx],
|
||||
w2_gs[idx],
|
||||
dtype=dtype,
|
||||
device=w2_q.device,
|
||||
block_size=quant_blocksize,
|
||||
)
|
||||
a_scale = None
|
||||
w1_scale = None
|
||||
w2_scale = None
|
||||
quant_dtype = None
|
||||
per_act_token_quant = False
|
||||
block_shape = None
|
||||
else:
|
||||
a = rank_tensors.hidden_states
|
||||
a_scale = rank_tensors.hidden_states_scale
|
||||
w1 = weights.w1
|
||||
w1_scale = weights.w1_scale
|
||||
w2 = weights.w2
|
||||
w2_scale = weights.w2_scale
|
||||
quant_dtype = config.quant_dtype
|
||||
per_act_token_quant = config.is_per_act_token_quant
|
||||
block_shape = config.quant_block_shape
|
||||
|
||||
return torch_experts(
|
||||
a=a,
|
||||
w1=w1,
|
||||
w2=w2,
|
||||
topk_weight=rank_tensors.topk_weights,
|
||||
topk_ids=rank_tensors.topk_ids,
|
||||
global_num_experts=config.E,
|
||||
expert_map=None,
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
a1_scale=a_scale,
|
||||
quant_dtype=quant_dtype,
|
||||
per_act_token_quant=per_act_token_quant,
|
||||
block_shape=block_shape,
|
||||
apply_router_weights_on_input=config.topk == 1
|
||||
and config.supports_apply_weight_on_input(),
|
||||
)
|
||||
|
||||
|
||||
def _make_gscale(num_experts: int) -> torch.Tensor:
|
||||
return torch.ones(
|
||||
(num_experts,),
|
||||
device=torch.accelerator.current_device_index(),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def make_modular_kernel(
|
||||
config: Config,
|
||||
vllm_config: VllmConfig,
|
||||
quant_config: FusedMoEQuantConfig,
|
||||
) -> mk.FusedMoEKernel:
|
||||
# make moe config
|
||||
moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make(
|
||||
tp_size_=get_tensor_model_parallel_world_size(),
|
||||
pcp_size_=get_pcp_group().world_size,
|
||||
dp_size_=get_dp_group().world_size,
|
||||
sp_size_=1,
|
||||
vllm_parallel_config=vllm_config.parallel_config,
|
||||
)
|
||||
|
||||
moe = FusedMoEConfig(
|
||||
num_experts=config.E,
|
||||
experts_per_token=config.topk,
|
||||
hidden_dim=config.K,
|
||||
intermediate_size=config.N,
|
||||
num_local_experts=config.num_local_experts,
|
||||
num_logical_experts=config.E,
|
||||
moe_parallel_config=moe_parallel_config,
|
||||
in_dtype=config.dtype,
|
||||
max_num_tokens=next_power_of_2(config.M),
|
||||
activation=MoEActivation.SILU,
|
||||
device=vllm_config.device_config.device,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
)
|
||||
|
||||
prepare_finalize = maybe_make_prepare_finalize(
|
||||
moe=moe,
|
||||
quant_config=quant_config,
|
||||
allow_new_interface=True,
|
||||
)
|
||||
assert prepare_finalize is not None
|
||||
|
||||
fused_experts = make_fused_experts(
|
||||
config.fused_experts_type,
|
||||
moe,
|
||||
quant_config,
|
||||
prepare_finalize.num_dispatchers(),
|
||||
config.N,
|
||||
)
|
||||
|
||||
modular_kernel = mk.FusedMoEKernel(
|
||||
prepare_finalize=prepare_finalize,
|
||||
fused_experts=fused_experts,
|
||||
)
|
||||
|
||||
return modular_kernel
|
||||
|
||||
|
||||
def _maybe_convert_weights_for_experts(
|
||||
config: Config,
|
||||
rank_weights: WeightTensors,
|
||||
) -> WeightTensors:
|
||||
"""Convert weights to expert-specific format (e.g., TrtLLM BlockMajorK)."""
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
|
||||
Fp8MoeBackend,
|
||||
convert_to_fp8_moe_kernel_format,
|
||||
)
|
||||
|
||||
fe_type = config.fused_experts_type
|
||||
fe_name = getattr(fe_type, "__name__", "")
|
||||
|
||||
backend: Fp8MoeBackend | None = None
|
||||
if fe_name == "TrtLlmFp8ExpertsModular":
|
||||
backend = Fp8MoeBackend.FLASHINFER_TRTLLM
|
||||
elif fe_name == "FlashInferExperts":
|
||||
backend = Fp8MoeBackend.FLASHINFER_CUTLASS
|
||||
|
||||
if backend is None or not rank_weights.is_quantized():
|
||||
return rank_weights
|
||||
|
||||
mock_layer = SimpleNamespace(
|
||||
weight_block_size=config.quant_block_shape,
|
||||
moe_config=SimpleNamespace(
|
||||
is_act_and_mul=True,
|
||||
intermediate_size_per_partition=config.N,
|
||||
),
|
||||
activation=SimpleNamespace(is_gated=True),
|
||||
)
|
||||
|
||||
w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format(
|
||||
fp8_backend=backend,
|
||||
layer=mock_layer,
|
||||
w13=rank_weights.w1,
|
||||
w2=rank_weights.w2,
|
||||
w13_scale=rank_weights.w1_scale,
|
||||
w2_scale=rank_weights.w2_scale,
|
||||
w13_input_scale=None,
|
||||
w2_input_scale=None,
|
||||
)
|
||||
|
||||
return WeightTensors(
|
||||
w1=w1,
|
||||
w2=w2,
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
w1_gs=rank_weights.w1_gs,
|
||||
w2_gs=rank_weights.w2_gs,
|
||||
)
|
||||
|
||||
|
||||
def run_modular_kernel(
|
||||
pgi: ProcessGroupInfo,
|
||||
vllm_config: VllmConfig,
|
||||
config: Config,
|
||||
weights: WeightTensors,
|
||||
rank_tensors: RankTensors,
|
||||
) -> torch.Tensor:
|
||||
assert isinstance(config.Ms, int)
|
||||
assert isinstance(config.topks, int)
|
||||
|
||||
# weights for rank
|
||||
rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)
|
||||
rank_weights = _maybe_convert_weights_for_experts(config, rank_weights)
|
||||
|
||||
if config.quant_dtype == "nvfp4":
|
||||
gscale = _make_gscale(config.num_local_experts)
|
||||
else:
|
||||
gscale = None
|
||||
|
||||
quant_config = FusedMoEQuantConfig.make(
|
||||
config.quant_dtype,
|
||||
w1_scale=rank_weights.w1_scale,
|
||||
w2_scale=rank_weights.w2_scale,
|
||||
a1_scale=rank_tensors.hidden_states_scale,
|
||||
g1_alphas=(1 / rank_weights.w1_gs) if rank_weights.w1_gs is not None else None,
|
||||
g2_alphas=(1 / rank_weights.w2_gs) if rank_weights.w2_gs is not None else None,
|
||||
a1_gscale=gscale,
|
||||
a2_gscale=gscale,
|
||||
block_shape=config.quant_block_shape,
|
||||
per_act_token_quant=config.is_per_act_token_quant,
|
||||
per_out_ch_quant=config.is_per_out_ch_quant,
|
||||
)
|
||||
|
||||
mk = make_modular_kernel(config, vllm_config, quant_config)
|
||||
|
||||
# impls might update the tensor in place
|
||||
hidden_states = rank_tensors.hidden_states.clone()
|
||||
|
||||
topk_ids = rank_tensors.topk_ids.to(mk.prepare_finalize.topk_indices_dtype())
|
||||
|
||||
mk_kwargs = {
|
||||
"hidden_states": hidden_states,
|
||||
"w1": rank_weights.w1,
|
||||
"w2": rank_weights.w2,
|
||||
"topk_weights": rank_tensors.topk_weights,
|
||||
"topk_ids": topk_ids,
|
||||
"activation": MoEActivation.SILU,
|
||||
"expert_map": rank_tensors.expert_map,
|
||||
"global_num_experts": config.E,
|
||||
"apply_router_weight_on_input": config.topk == 1
|
||||
and config.supports_apply_weight_on_input(),
|
||||
}
|
||||
|
||||
num_tokens = rank_tensors.hidden_states.shape[0]
|
||||
num_tokens_across_dp = torch.tensor(
|
||||
[num_tokens] * config.world_size, device="cuda", dtype=torch.int
|
||||
)
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
with set_forward_context(
|
||||
None,
|
||||
vllm_config,
|
||||
num_tokens=num_tokens,
|
||||
num_tokens_across_dp=num_tokens_across_dp,
|
||||
):
|
||||
out = mk.apply(**mk_kwargs)
|
||||
|
||||
return out
|
||||
@@ -0,0 +1,189 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import copy
|
||||
from enum import Enum
|
||||
from itertools import product
|
||||
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.fused_moe.config import FUSED_MOE_UNQUANTIZED_CONFIG
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
from .common import (
|
||||
Config,
|
||||
RankTensors,
|
||||
WeightTensors,
|
||||
reference_moe_impl,
|
||||
run_modular_kernel,
|
||||
)
|
||||
from .mk_objects import (
|
||||
MK_FUSED_EXPERT_TYPES,
|
||||
MK_MULTI_GPU_PREPARE_FINALIZE_TYPES,
|
||||
MK_QUANT_CONFIGS,
|
||||
)
|
||||
from .parallel_utils import ProcessGroupInfo, parallel_launch_with_config
|
||||
|
||||
|
||||
class Result(Enum):
|
||||
PASS = 1
|
||||
FAIL = 2
|
||||
SKIP = 3
|
||||
|
||||
|
||||
def rank_worker(
|
||||
pgi: ProcessGroupInfo,
|
||||
vllm_config: VllmConfig,
|
||||
cpu_group,
|
||||
config: Config,
|
||||
weights: WeightTensors,
|
||||
):
|
||||
set_random_seed(pgi.rank)
|
||||
|
||||
# get weights to this device
|
||||
weights.to_current_device()
|
||||
|
||||
Ms = config.Ms
|
||||
assert isinstance(Ms, list)
|
||||
TOPKs = config.topks
|
||||
assert isinstance(TOPKs, list)
|
||||
|
||||
for m, topk in product(Ms, TOPKs):
|
||||
print(f"Running m={m}, topk={topk} ...")
|
||||
# override m and topk
|
||||
cfgx = copy.deepcopy(config)
|
||||
cfgx.Ms = m
|
||||
cfgx.topks = topk
|
||||
|
||||
# inputs for rank
|
||||
rank_tensors = RankTensors.make(cfgx, pgi)
|
||||
|
||||
# modular kernel out
|
||||
mk_out = run_modular_kernel(pgi, vllm_config, cfgx, weights, rank_tensors)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
ref_out = reference_moe_impl(cfgx, weights, rank_tensors)
|
||||
|
||||
torch.testing.assert_close(ref_out, mk_out, atol=3e-2, rtol=3e-2)
|
||||
|
||||
|
||||
def make_feature_matrix(csv_file_path: str):
|
||||
from dataclasses import asdict
|
||||
|
||||
import pandas as pd
|
||||
|
||||
def add_to_results(
|
||||
config: Config, success: Result, results_df: pd.DataFrame | None = None
|
||||
):
|
||||
config_dict = asdict(config)
|
||||
config_dict["prepare_finalize_type"] = config_dict[
|
||||
"prepare_finalize_type"
|
||||
].__name__
|
||||
config_dict["fused_experts_type"] = config_dict["fused_experts_type"].__name__
|
||||
config_dict["per_tensor_act_quant"] = config.is_per_tensor_act_quant
|
||||
quant_config_dict = config_dict["quant_config"]
|
||||
del config_dict["quant_config"]
|
||||
if quant_config_dict is None:
|
||||
quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
|
||||
quant_config_dict = asdict(quant_config)
|
||||
|
||||
config_dict |= quant_config_dict
|
||||
result_dict = config_dict | {"success": success.name}
|
||||
|
||||
result_df = pd.DataFrame([result_dict])
|
||||
if results_df is None:
|
||||
results_df = result_df
|
||||
else:
|
||||
results_df = pd.concat([results_df, result_df], ignore_index=True)
|
||||
|
||||
return results_df
|
||||
|
||||
Ms = [64]
|
||||
Ks = [7168] # hidden sizes
|
||||
Ns = [2048]
|
||||
TOPKs = [[4, 1]]
|
||||
Es = [32]
|
||||
DTYPEs = [torch.bfloat16]
|
||||
PF_TYPES = MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
|
||||
FE_TYPES = MK_FUSED_EXPERT_TYPES
|
||||
Q_TYPES = MK_QUANT_CONFIGS
|
||||
|
||||
combinations = list(
|
||||
product(Ms, Ks, Ns, Es, TOPKs, DTYPEs, PF_TYPES, FE_TYPES, Q_TYPES)
|
||||
)
|
||||
|
||||
results_df: pd.DataFrame | None = None
|
||||
for m, k, n, e, topks, dtype, pf_type, experts_type, quant_config in tqdm(
|
||||
combinations
|
||||
):
|
||||
config = Config(
|
||||
Ms=[m],
|
||||
K=k,
|
||||
N=n,
|
||||
E=e,
|
||||
topks=topks,
|
||||
dtype=dtype,
|
||||
prepare_finalize_type=pf_type,
|
||||
fused_experts_type=experts_type,
|
||||
quant_config=quant_config,
|
||||
world_size=2,
|
||||
)
|
||||
|
||||
success = None
|
||||
if config.is_valid()[0]:
|
||||
print(f"Running config : {config.describe()} ...")
|
||||
try:
|
||||
weights: WeightTensors = WeightTensors.make(config)
|
||||
vllm_config, env_dict = config.make_env_data()
|
||||
parallel_launch_with_config(
|
||||
config.world_size,
|
||||
rank_worker,
|
||||
vllm_config,
|
||||
env_dict,
|
||||
config,
|
||||
weights,
|
||||
)
|
||||
success = Result.PASS
|
||||
except Exception as _:
|
||||
success = Result.FAIL
|
||||
else:
|
||||
success = Result.SKIP
|
||||
|
||||
results_df = add_to_results(config, success, results_df)
|
||||
|
||||
if results_df is not None:
|
||||
results_df.to_csv(f"{csv_file_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
"Make ModularKernel feature matrix \n"
|
||||
"Example : python3 -m tests.kernels.moe.modular_kernel_tools.make_feature_matrix " # noqa: E501
|
||||
"-f ./feature_matrices/feature_matrix.csv"
|
||||
)
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"-f",
|
||||
"--feature-matrix-csv-file-path",
|
||||
type=str,
|
||||
required=True,
|
||||
help="File name to Generate a .csv file",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
csv_path = args.feature_matrix_csv_file_path
|
||||
assert csv_path.endswith("csv"), (
|
||||
f"Need a file path ending with .csv, got {csv_path}"
|
||||
)
|
||||
assert Path(csv_path).parent.is_dir(), (
|
||||
f"Cannot find parent directory for {Path(csv_path).parent}"
|
||||
)
|
||||
|
||||
make_feature_matrix(args.feature_matrix_csv_file_path)
|
||||
@@ -0,0 +1,497 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
# Fused experts and PrepareFinalize imports
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe import TritonExperts
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import (
|
||||
BatchedDeepGemmExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import DeepGemmExperts
|
||||
from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
|
||||
BatchedTritonExperts,
|
||||
NaiveBatchedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
|
||||
TritonOrDeepGemmExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
|
||||
MoEPrepareAndFinalizeNoDPEPModular,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
|
||||
cutlass_fp4_supported,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
cutlass_fp8_supported,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.deep_gemm import is_deep_gemm_supported
|
||||
from vllm.utils.flashinfer import (
|
||||
has_flashinfer_cutlass_fused_moe,
|
||||
has_flashinfer_nvlink_one_sided,
|
||||
has_flashinfer_trtllm_fused_moe,
|
||||
)
|
||||
from vllm.utils.import_utils import (
|
||||
has_aiter,
|
||||
has_deep_ep,
|
||||
has_deep_ep_v2,
|
||||
has_deep_gemm,
|
||||
has_mori,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TestMoEQuantConfig:
|
||||
quant_dtype: torch.dtype | str | None
|
||||
per_out_ch_quant: bool
|
||||
per_act_token_quant: bool
|
||||
block_shape: list[int] | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PrepareFinalizeInfo:
|
||||
activation_format: mk.FusedMoEActivationFormat
|
||||
supported_dtypes: list[torch.dtype | str]
|
||||
blocked_quantization_support: bool
|
||||
backend: str | None
|
||||
supports_apply_weight_on_input: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExpertInfo:
|
||||
activation_format: mk.FusedMoEActivationFormat
|
||||
supported_dtypes: list[torch.dtype | str]
|
||||
blocked_quantization_support: bool
|
||||
needs_matching_quant: bool = False
|
||||
needs_deep_gemm: bool = False
|
||||
needs_aiter: bool = False
|
||||
|
||||
|
||||
PREPARE_FINALIZE_INFO: dict[
|
||||
mk.FusedMoEPrepareAndFinalizeModular, PrepareFinalizeInfo
|
||||
] = {}
|
||||
EXPERT_INFO: dict[mk.FusedMoEExpertsModular, ExpertInfo] = {}
|
||||
MK_ALL_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalizeModular] = []
|
||||
MK_MULTI_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalizeModular] = []
|
||||
MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalizeModular] = []
|
||||
MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEExpertsModular] = []
|
||||
|
||||
standard_format = mk.FusedMoEActivationFormat.Standard
|
||||
batched_format = mk.FusedMoEActivationFormat.BatchedExperts
|
||||
common_float_types: list[torch.dtype | str] = [
|
||||
torch.float8_e4m3fn,
|
||||
torch.bfloat16,
|
||||
torch.float16,
|
||||
torch.float32,
|
||||
]
|
||||
common_float_and_int_types = common_float_types + [torch.int8]
|
||||
nvfp4_types = ["nvfp4"]
|
||||
fp8_types = [torch.float8_e4m3fn]
|
||||
|
||||
|
||||
def register_prepare_and_finalize(
|
||||
kind,
|
||||
activation_format: mk.FusedMoEActivationFormat,
|
||||
supported_dtypes: list[torch.dtype | str],
|
||||
blocked_quantization_support: bool,
|
||||
backend: str | None,
|
||||
force_multigpu: bool = False,
|
||||
supports_apply_weight_on_input: bool = True,
|
||||
):
|
||||
global PREPARE_FINALIZE_INFO
|
||||
global MK_ALL_PREPARE_FINALIZE_TYPES
|
||||
global MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
|
||||
global MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES
|
||||
assert kind not in PREPARE_FINALIZE_INFO
|
||||
|
||||
PREPARE_FINALIZE_INFO[kind] = PrepareFinalizeInfo(
|
||||
activation_format,
|
||||
supported_dtypes,
|
||||
blocked_quantization_support,
|
||||
backend,
|
||||
supports_apply_weight_on_input,
|
||||
)
|
||||
MK_ALL_PREPARE_FINALIZE_TYPES.append(kind)
|
||||
if backend is not None or force_multigpu:
|
||||
MK_MULTI_GPU_PREPARE_FINALIZE_TYPES.append(kind)
|
||||
else:
|
||||
MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES.append(kind)
|
||||
|
||||
|
||||
def register_experts(
|
||||
kind,
|
||||
activation_format: mk.FusedMoEActivationFormat,
|
||||
supported_dtypes: list[torch.dtype | str],
|
||||
blocked_quantization_support: bool,
|
||||
needs_matching_quant: bool = False,
|
||||
needs_deep_gemm: bool = False,
|
||||
needs_aiter: bool = False,
|
||||
):
|
||||
global EXPERT_INFO
|
||||
global MK_FUSED_EXPERT_TYPES
|
||||
assert kind not in EXPERT_INFO
|
||||
|
||||
EXPERT_INFO[kind] = ExpertInfo(
|
||||
activation_format,
|
||||
supported_dtypes,
|
||||
blocked_quantization_support,
|
||||
needs_matching_quant,
|
||||
needs_deep_gemm,
|
||||
needs_aiter,
|
||||
)
|
||||
|
||||
MK_FUSED_EXPERT_TYPES.append(kind)
|
||||
|
||||
|
||||
def prepare_finalize_info(kind) -> PrepareFinalizeInfo:
|
||||
info = PREPARE_FINALIZE_INFO.get(kind)
|
||||
assert info is not None
|
||||
return info
|
||||
|
||||
|
||||
def expert_info(kind) -> ExpertInfo:
|
||||
info = EXPERT_INFO.get(kind)
|
||||
assert info is not None
|
||||
return info
|
||||
|
||||
|
||||
register_prepare_and_finalize(
|
||||
MoEPrepareAndFinalizeNoDPEPModular,
|
||||
standard_format,
|
||||
common_float_types,
|
||||
blocked_quantization_support=True,
|
||||
backend=None,
|
||||
)
|
||||
|
||||
register_experts(
|
||||
BatchedTritonExperts,
|
||||
batched_format,
|
||||
common_float_types,
|
||||
blocked_quantization_support=True,
|
||||
needs_matching_quant=True,
|
||||
)
|
||||
|
||||
register_experts(
|
||||
TritonExperts,
|
||||
standard_format,
|
||||
common_float_and_int_types,
|
||||
blocked_quantization_support=True,
|
||||
needs_matching_quant=True,
|
||||
)
|
||||
|
||||
register_experts(
|
||||
NaiveBatchedExperts,
|
||||
batched_format,
|
||||
common_float_and_int_types,
|
||||
blocked_quantization_support=True,
|
||||
)
|
||||
|
||||
# Disable on blackwell for now
|
||||
if has_deep_ep() and not current_platform.has_device_capability(100):
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import (
|
||||
DeepEPHTPrepareAndFinalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import (
|
||||
DeepEPLLPrepareAndFinalize,
|
||||
)
|
||||
|
||||
register_prepare_and_finalize(
|
||||
DeepEPHTPrepareAndFinalize,
|
||||
standard_format,
|
||||
common_float_types,
|
||||
blocked_quantization_support=True,
|
||||
backend="deepep_high_throughput",
|
||||
)
|
||||
|
||||
register_prepare_and_finalize(
|
||||
DeepEPLLPrepareAndFinalize,
|
||||
batched_format,
|
||||
common_float_types,
|
||||
blocked_quantization_support=True,
|
||||
backend="deepep_low_latency",
|
||||
)
|
||||
|
||||
if has_deep_ep_v2() and current_platform.has_device_capability(100):
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import (
|
||||
DeepEPV2PrepareAndFinalize,
|
||||
)
|
||||
|
||||
register_prepare_and_finalize(
|
||||
DeepEPV2PrepareAndFinalize,
|
||||
standard_format,
|
||||
common_float_types,
|
||||
blocked_quantization_support=True,
|
||||
backend="deepep_v2",
|
||||
)
|
||||
|
||||
if has_mori():
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
|
||||
MoriPrepareAndFinalize,
|
||||
)
|
||||
|
||||
register_prepare_and_finalize(
|
||||
MoriPrepareAndFinalize,
|
||||
standard_format,
|
||||
fp8_types,
|
||||
blocked_quantization_support=True,
|
||||
backend="mori_high_throughput",
|
||||
supports_apply_weight_on_input=False,
|
||||
)
|
||||
|
||||
if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
|
||||
from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (
|
||||
FlashInferExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import ( # noqa: E501
|
||||
FlashInferNVLinkTwoSidedPrepareAndFinalize,
|
||||
)
|
||||
|
||||
register_prepare_and_finalize(
|
||||
FlashInferNVLinkTwoSidedPrepareAndFinalize,
|
||||
standard_format,
|
||||
nvfp4_types + fp8_types,
|
||||
blocked_quantization_support=True,
|
||||
backend=None,
|
||||
force_multigpu=True,
|
||||
supports_apply_weight_on_input=False,
|
||||
)
|
||||
|
||||
register_experts(
|
||||
FlashInferExperts,
|
||||
standard_format,
|
||||
nvfp4_types + fp8_types,
|
||||
blocked_quantization_support=True,
|
||||
# Note: this is a hack to get it to run for now
|
||||
)
|
||||
else:
|
||||
FlashInferCutlassMoEPrepareAndFinalize = None
|
||||
FlashInferExperts = None
|
||||
|
||||
if (
|
||||
has_flashinfer_nvlink_one_sided()
|
||||
and has_flashinfer_cutlass_fused_moe()
|
||||
and current_platform.has_device_capability(100)
|
||||
):
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided import ( # noqa: E501
|
||||
FlashInferNVLinkOneSidedPrepareAndFinalize,
|
||||
)
|
||||
|
||||
register_prepare_and_finalize(
|
||||
FlashInferNVLinkOneSidedPrepareAndFinalize,
|
||||
standard_format,
|
||||
nvfp4_types,
|
||||
blocked_quantization_support=False,
|
||||
backend="flashinfer_nvlink_one_sided",
|
||||
supports_apply_weight_on_input=False,
|
||||
)
|
||||
|
||||
if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
|
||||
from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
|
||||
TrtLlmNvFp4ExpertsModular,
|
||||
)
|
||||
|
||||
register_experts(
|
||||
TrtLlmNvFp4ExpertsModular,
|
||||
standard_format,
|
||||
nvfp4_types,
|
||||
blocked_quantization_support=False,
|
||||
)
|
||||
|
||||
if has_flashinfer_trtllm_fused_moe() and current_platform.has_device_capability(100):
|
||||
from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
|
||||
TrtLlmFp8ExpertsModular,
|
||||
)
|
||||
|
||||
register_experts(
|
||||
TrtLlmFp8ExpertsModular,
|
||||
standard_format,
|
||||
fp8_types,
|
||||
blocked_quantization_support=True,
|
||||
)
|
||||
|
||||
if has_aiter():
|
||||
from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
|
||||
AiterExperts,
|
||||
)
|
||||
|
||||
register_experts(
|
||||
AiterExperts,
|
||||
standard_format,
|
||||
fp8_types,
|
||||
blocked_quantization_support=True,
|
||||
needs_aiter=True,
|
||||
)
|
||||
else:
|
||||
AiterExperts = None
|
||||
|
||||
if has_deep_gemm() and is_deep_gemm_supported():
|
||||
register_experts(
|
||||
BatchedDeepGemmExperts,
|
||||
batched_format,
|
||||
fp8_types,
|
||||
blocked_quantization_support=True,
|
||||
needs_matching_quant=False,
|
||||
needs_deep_gemm=True,
|
||||
)
|
||||
register_experts(
|
||||
DeepGemmExperts,
|
||||
standard_format,
|
||||
fp8_types,
|
||||
blocked_quantization_support=True,
|
||||
needs_matching_quant=False,
|
||||
needs_deep_gemm=True,
|
||||
)
|
||||
register_experts(
|
||||
TritonOrDeepGemmExperts,
|
||||
standard_format,
|
||||
common_float_and_int_types,
|
||||
blocked_quantization_support=True,
|
||||
needs_matching_quant=True,
|
||||
needs_deep_gemm=True,
|
||||
)
|
||||
|
||||
if cutlass_fp8_supported():
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
CutlassBatchedExpertsFp8,
|
||||
CutlassExpertsFp8,
|
||||
)
|
||||
|
||||
register_experts(
|
||||
CutlassExpertsFp8,
|
||||
standard_format,
|
||||
fp8_types,
|
||||
blocked_quantization_support=False,
|
||||
)
|
||||
register_experts(
|
||||
CutlassBatchedExpertsFp8,
|
||||
batched_format,
|
||||
fp8_types,
|
||||
blocked_quantization_support=False,
|
||||
)
|
||||
else:
|
||||
CutlassBatchedExpertsFp8 = None
|
||||
CutlassExpertsFp8 = None
|
||||
|
||||
if cutlass_fp4_supported():
|
||||
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
|
||||
CutlassExpertsFp4,
|
||||
)
|
||||
|
||||
register_experts(
|
||||
CutlassExpertsFp4,
|
||||
standard_format,
|
||||
nvfp4_types,
|
||||
blocked_quantization_support=True,
|
||||
)
|
||||
else:
|
||||
CutlassExpertsFp4 = None
|
||||
|
||||
MK_QUANT_CONFIGS: list[TestMoEQuantConfig | None] = [
|
||||
None,
|
||||
# per-channel / per-column weights and per-tensor activations
|
||||
TestMoEQuantConfig(
|
||||
quant_dtype=torch.float8_e4m3fn,
|
||||
per_out_ch_quant=True,
|
||||
per_act_token_quant=False,
|
||||
block_shape=None,
|
||||
),
|
||||
# per-channel / per-column weights and per-token activations
|
||||
TestMoEQuantConfig(
|
||||
quant_dtype=torch.float8_e4m3fn,
|
||||
per_out_ch_quant=True,
|
||||
per_act_token_quant=True,
|
||||
block_shape=None,
|
||||
),
|
||||
# per-tensor weights and per-tensor activations
|
||||
TestMoEQuantConfig(
|
||||
quant_dtype=torch.float8_e4m3fn,
|
||||
per_out_ch_quant=False,
|
||||
per_act_token_quant=False,
|
||||
block_shape=None,
|
||||
),
|
||||
# per-tensor weights and per-token activations
|
||||
TestMoEQuantConfig(
|
||||
quant_dtype=torch.float8_e4m3fn,
|
||||
per_out_ch_quant=False,
|
||||
per_act_token_quant=True,
|
||||
block_shape=None,
|
||||
),
|
||||
# block-quantized weights and 128 block per-token activations
|
||||
TestMoEQuantConfig(
|
||||
quant_dtype=torch.float8_e4m3fn,
|
||||
per_out_ch_quant=False,
|
||||
per_act_token_quant=False,
|
||||
block_shape=[128, 128],
|
||||
),
|
||||
# TODO (varun) : Should we test the following combinations ?
|
||||
# block-quantized weights and per-token activations
|
||||
# block-quantized weights and per-tensor activations
|
||||
]
|
||||
|
||||
if cutlass_fp4_supported() or has_flashinfer_cutlass_fused_moe():
|
||||
MK_QUANT_CONFIGS += [
|
||||
TestMoEQuantConfig(
|
||||
quant_dtype="nvfp4",
|
||||
per_out_ch_quant=False,
|
||||
per_act_token_quant=False,
|
||||
block_shape=None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _slice(rank: int, num_local_experts: int, t: torch.Tensor) -> torch.Tensor:
|
||||
s = rank * num_local_experts
|
||||
e = s + num_local_experts
|
||||
return t[s:e]
|
||||
|
||||
|
||||
def make_cutlass_strides(
|
||||
e: int,
|
||||
n: int,
|
||||
k: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
ab_strides1 = torch.full((e,), k, device="cuda", dtype=torch.int64)
|
||||
ab_strides2 = torch.full((e,), n, device="cuda", dtype=torch.int64)
|
||||
c_strides1 = torch.full((e,), 2 * n, device="cuda", dtype=torch.int64)
|
||||
c_strides2 = torch.full((e,), k, device="cuda", dtype=torch.int64)
|
||||
return ab_strides1, ab_strides2, c_strides1, c_strides2
|
||||
|
||||
|
||||
def make_fused_experts(
|
||||
fused_experts_type: mk.FusedMoEExpertsModular,
|
||||
moe: FusedMoEConfig,
|
||||
quant_config: FusedMoEQuantConfig,
|
||||
num_dispatchers: int,
|
||||
N: int,
|
||||
) -> mk.FusedMoEExpertsModular:
|
||||
if (
|
||||
fused_experts_type.activation_format()
|
||||
== mk.FusedMoEActivationFormat.BatchedExperts
|
||||
):
|
||||
kwargs = {
|
||||
"moe_config": moe,
|
||||
"quant_config": quant_config,
|
||||
"max_num_tokens": moe.max_num_tokens,
|
||||
"num_dispatchers": num_dispatchers,
|
||||
}
|
||||
else:
|
||||
kwargs = {
|
||||
"moe_config": moe,
|
||||
"quant_config": quant_config,
|
||||
}
|
||||
|
||||
torch.set_printoptions(threshold=0, edgeitems=0, linewidth=10000)
|
||||
|
||||
print(f"Making {fused_experts_type.__class__.__name__} {kwargs} ...")
|
||||
experts = fused_experts_type(**kwargs)
|
||||
|
||||
torch.set_printoptions(threshold=1000, edgeitems=5, linewidth=80)
|
||||
|
||||
return experts
|
||||
@@ -0,0 +1,166 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import dataclasses
|
||||
import os
|
||||
import traceback
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Concatenate
|
||||
|
||||
import torch
|
||||
from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUsage]
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.distributed import (
|
||||
cleanup_dist_env_and_memory,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
|
||||
## Parallel Processes Utils
|
||||
|
||||
P = ParamSpec("P")
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ProcessGroupInfo:
|
||||
world_size: int
|
||||
world_local_size: int
|
||||
rank: int
|
||||
node_rank: int
|
||||
local_rank: int
|
||||
device: torch.device
|
||||
|
||||
|
||||
def _set_vllm_config(
|
||||
vllm_config: VllmConfig, world_size: int, rank: int, local_rank: int
|
||||
):
|
||||
import tempfile
|
||||
|
||||
temp_file = tempfile.mkstemp()[1]
|
||||
|
||||
# When DP is enabled, processes are organized as:
|
||||
# rank = dp_rank * tp_pp_world_size + tp_pp_rank
|
||||
tp_pp_world_size = vllm_config.parallel_config.world_size
|
||||
vllm_config.parallel_config.data_parallel_rank = rank // tp_pp_world_size
|
||||
tp_pp_rank = rank % tp_pp_world_size
|
||||
vllm_config.parallel_config.rank = tp_pp_rank
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
init_distributed_environment(
|
||||
world_size=tp_pp_world_size,
|
||||
rank=tp_pp_rank,
|
||||
distributed_init_method=f"file://{temp_file}",
|
||||
local_rank=local_rank,
|
||||
backend="nccl",
|
||||
)
|
||||
|
||||
initialize_model_parallel(
|
||||
tensor_model_parallel_size=vllm_config.parallel_config.tensor_parallel_size,
|
||||
pipeline_model_parallel_size=vllm_config.parallel_config.pipeline_parallel_size,
|
||||
)
|
||||
if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
|
||||
cpu_group = torch.distributed.split_group(
|
||||
split_ranks=[list(range(world_size))],
|
||||
group_desc="moe_test_cpu",
|
||||
)
|
||||
else:
|
||||
cpu_group = torch.distributed.new_group(
|
||||
list(range(world_size)), backend="gloo"
|
||||
)
|
||||
return cpu_group
|
||||
|
||||
|
||||
def _worker_parallel_launch(
|
||||
local_rank: int,
|
||||
world_size: int,
|
||||
world_local_size: int,
|
||||
node_rank: int,
|
||||
init_method: str,
|
||||
worker: Callable[..., None],
|
||||
vllm_config: VllmConfig | None,
|
||||
env_dict: dict | None,
|
||||
worker_kwargs: dict[str, Any],
|
||||
*args: Any,
|
||||
) -> None:
|
||||
rank = node_rank * world_local_size + local_rank
|
||||
device = torch.device("cuda", local_rank)
|
||||
torch.accelerator.set_device_index(device)
|
||||
torch.set_default_device(device)
|
||||
torch.distributed.init_process_group(
|
||||
backend="cpu:gloo,cuda:nccl",
|
||||
init_method=init_method,
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
device_id=device,
|
||||
)
|
||||
barrier = torch.tensor([rank], device=device)
|
||||
torch.distributed.all_reduce(barrier)
|
||||
|
||||
if env_dict is not None:
|
||||
os.environ.update(env_dict)
|
||||
|
||||
cpu_group = None
|
||||
if vllm_config is not None:
|
||||
cpu_group = _set_vllm_config(vllm_config, world_size, rank, local_rank)
|
||||
|
||||
def _run_worker():
|
||||
worker(
|
||||
ProcessGroupInfo(
|
||||
world_size=world_size,
|
||||
world_local_size=world_local_size,
|
||||
rank=rank,
|
||||
node_rank=node_rank,
|
||||
local_rank=local_rank,
|
||||
device=device,
|
||||
),
|
||||
vllm_config,
|
||||
cpu_group,
|
||||
*args,
|
||||
**worker_kwargs,
|
||||
)
|
||||
|
||||
try:
|
||||
if vllm_config is not None:
|
||||
with set_current_vllm_config(vllm_config):
|
||||
_run_worker()
|
||||
else:
|
||||
_run_worker()
|
||||
except Exception as ex:
|
||||
print(ex)
|
||||
traceback.print_exc()
|
||||
raise
|
||||
finally:
|
||||
torch.accelerator.synchronize()
|
||||
if vllm_config is not None:
|
||||
cleanup_dist_env_and_memory()
|
||||
else:
|
||||
torch.distributed.destroy_process_group()
|
||||
|
||||
|
||||
def parallel_launch_with_config(
|
||||
world_size: int,
|
||||
worker: Callable[Concatenate[ProcessGroupInfo, VllmConfig, Any, P], None],
|
||||
vllm_config: VllmConfig,
|
||||
env_dict: dict[Any, Any] | None,
|
||||
*args: P.args,
|
||||
**kwargs: P.kwargs,
|
||||
) -> None:
|
||||
spawn(
|
||||
_worker_parallel_launch,
|
||||
args=(
|
||||
world_size,
|
||||
world_size,
|
||||
0,
|
||||
f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}",
|
||||
worker,
|
||||
vllm_config,
|
||||
env_dict,
|
||||
kwargs,
|
||||
)
|
||||
+ args,
|
||||
nprocs=world_size,
|
||||
join=True,
|
||||
)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user