chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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)
|
||||
Reference in New Issue
Block a user