wehub-resource-sync
227 lines
7.0 KiB
Python
227 lines
7.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Precision tests for vllm's chunk_kda Triton operator.
|
|
|
|
Compares chunk_kda against a naive recurrent reference (float32).
|
|
Uses torch.rand for q/k/v to match FLA's test pattern.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from vllm.model_executor.layers.fla.ops.kda import (
|
|
chunk_kda,
|
|
chunk_kda_with_fused_gate,
|
|
fused_kda_gate,
|
|
)
|
|
from vllm.model_executor.layers.fla.ops.l2norm import l2norm_fwd
|
|
|
|
DEVICE = "cuda"
|
|
|
|
|
|
def naive_recurrent_kda(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
g: torch.Tensor,
|
|
beta: torch.Tensor,
|
|
scale: float | None = None,
|
|
initial_state: torch.Tensor | None = None,
|
|
output_final_state: bool = False,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|
"""Naive recurrent KDA reference, ported from FLA's naive.py."""
|
|
dtype = v.dtype
|
|
B, T, H, K = q.shape
|
|
V = v.shape[-1]
|
|
if scale is None:
|
|
scale = K**-0.5
|
|
|
|
q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta])
|
|
q = q * scale
|
|
|
|
S = k.new_zeros(B, H, K, V).to(q)
|
|
if initial_state is not None:
|
|
S += initial_state
|
|
o = torch.zeros_like(v)
|
|
for i in range(T):
|
|
q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
|
|
S = S * g_i[..., None].exp()
|
|
S = S + torch.einsum(
|
|
"bhk,bhv->bhkv",
|
|
b_i[..., None] * k_i,
|
|
v_i - (k_i[..., None] * S).sum(-2),
|
|
)
|
|
o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S)
|
|
if not output_final_state:
|
|
S = None
|
|
return o.to(dtype), S
|
|
|
|
|
|
def assert_close(
|
|
name: str,
|
|
ref: torch.Tensor,
|
|
tri: torch.Tensor,
|
|
ratio: float,
|
|
err_atol: float = 1e-6,
|
|
):
|
|
"""RMSE-based relative error comparison."""
|
|
abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item()
|
|
rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item()
|
|
rmse_base = ref.detach().flatten().square().mean().sqrt().item()
|
|
rel_err = rmse_diff / (rmse_base + 1e-8)
|
|
print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}")
|
|
if abs_err <= err_atol:
|
|
return
|
|
assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref"
|
|
assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri"
|
|
assert rel_err < ratio, (
|
|
f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("H", "D", "cu_seqlens", "dtype"),
|
|
[
|
|
pytest.param(
|
|
*test,
|
|
id="H{}-D{}-cu{}-{}".format(*test),
|
|
)
|
|
for test in [
|
|
(32, 128, [0, 64], torch.float16),
|
|
(32, 128, [0, 1024], torch.float16),
|
|
(32, 128, [0, 15], torch.float16),
|
|
(32, 128, [0, 256, 512, 768, 1024], torch.float16),
|
|
(32, 128, [0, 15, 100, 300, 1200], torch.float16),
|
|
(64, 128, [0, 256, 500, 1000], torch.float16),
|
|
(32, 128, [0, 8192], torch.float16),
|
|
(32, 128, [0, 256, 500, 1000], torch.bfloat16),
|
|
]
|
|
],
|
|
)
|
|
@torch.inference_mode()
|
|
def test_chunk_kda(
|
|
H: int,
|
|
D: int,
|
|
cu_seqlens: list[int],
|
|
dtype: torch.dtype,
|
|
):
|
|
T = cu_seqlens[-1]
|
|
torch.manual_seed(42)
|
|
B = 1
|
|
cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE)
|
|
N = len(cu_seqlens) - 1
|
|
|
|
q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
|
|
k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
|
|
v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
|
|
g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to(
|
|
dtype
|
|
)
|
|
beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid()
|
|
h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
|
|
|
|
# Naive reference with l2norm_fwd (same kernel as chunk_kda)
|
|
ref_outputs = []
|
|
ref_states = []
|
|
for i in range(N):
|
|
s, e = cu_seqlens[i], cu_seqlens[i + 1]
|
|
q_i = l2norm_fwd(q[:, s:e].contiguous())
|
|
k_i = l2norm_fwd(k[:, s:e].contiguous())
|
|
o_i, ht_i = naive_recurrent_kda(
|
|
q_i,
|
|
k_i,
|
|
v[:, s:e],
|
|
g[:, s:e],
|
|
beta[:, s:e],
|
|
initial_state=h0[i],
|
|
output_final_state=True,
|
|
)
|
|
ref_outputs.append(o_i)
|
|
ref_states.append(ht_i)
|
|
ref_o = torch.cat(ref_outputs, dim=1)
|
|
ref_ht = torch.cat(ref_states, dim=0)
|
|
|
|
# h0 transposed to (V, K) layout for the kernel; naive uses (K, V)
|
|
tri_o, tri_ht = chunk_kda(
|
|
q=q.clone(),
|
|
k=k.clone(),
|
|
v=v.clone(),
|
|
g=g.clone(),
|
|
beta=beta.clone(),
|
|
initial_state=h0.transpose(-1, -2).contiguous().clone(),
|
|
output_final_state=True,
|
|
cu_seqlens=cu_seqlens_t,
|
|
use_qk_l2norm_in_kernel=True,
|
|
)
|
|
|
|
assert not torch.isnan(tri_o).any(), "Triton output o contains NaN"
|
|
assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN"
|
|
assert_close("o", ref_o, tri_o, 0.005)
|
|
assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("cu_seqlens", "dtype"),
|
|
[
|
|
([0, 64], torch.float16),
|
|
([0, 15, 100, 300], torch.bfloat16),
|
|
],
|
|
)
|
|
@torch.inference_mode()
|
|
def test_chunk_kda_fused_gate_cumsum_matches_unfused(
|
|
cu_seqlens: list[int],
|
|
dtype: torch.dtype,
|
|
):
|
|
H, D = 8, 64
|
|
T = cu_seqlens[-1]
|
|
N = len(cu_seqlens) - 1
|
|
torch.manual_seed(123)
|
|
|
|
cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE)
|
|
q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
|
|
k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
|
|
v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
|
|
raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
|
|
beta = torch.rand(1, T, H, dtype=dtype, device=DEVICE).sigmoid()
|
|
A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous()
|
|
dt_bias = (
|
|
torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1
|
|
).contiguous()
|
|
h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
|
|
initial_state = h0.transpose(-1, -2).contiguous()
|
|
|
|
gate = fused_kda_gate(
|
|
raw_g.reshape(T, H * D),
|
|
A_log,
|
|
D,
|
|
g_bias=dt_bias,
|
|
).unsqueeze(0)
|
|
old_o, old_ht = chunk_kda(
|
|
q=q.clone(),
|
|
k=k.clone(),
|
|
v=v.clone(),
|
|
g=gate,
|
|
beta=beta.clone(),
|
|
initial_state=initial_state.clone(),
|
|
output_final_state=True,
|
|
cu_seqlens=cu_seqlens_t,
|
|
use_qk_l2norm_in_kernel=True,
|
|
)
|
|
new_o, new_ht = chunk_kda_with_fused_gate(
|
|
q=q.clone(),
|
|
k=k.clone(),
|
|
v=v.clone(),
|
|
raw_g=raw_g,
|
|
beta=beta.clone(),
|
|
A_log=A_log,
|
|
g_bias=dt_bias,
|
|
initial_state=initial_state.clone(),
|
|
output_final_state=True,
|
|
cu_seqlens=cu_seqlens_t,
|
|
use_qk_l2norm_in_kernel=True,
|
|
)
|
|
|
|
assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3)
|
|
assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3)
|