464 lines
14 KiB
Python
464 lines
14 KiB
Python
"""Unit tests for the fused compressor attention Triton kernel on HIP.
|
|
|
|
Validates numerical parity between the fused single-kernel path (plan-driven
|
|
Triton) and the reference per-seq Python implementation.
|
|
|
|
Usage:
|
|
python -m pytest test/manual/dsv4/test_fused_compress_attn_hip.py -v
|
|
# or directly:
|
|
python test/manual/dsv4/test_fused_compress_attn_hip.py
|
|
"""
|
|
|
|
import unittest
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
@dataclass
|
|
class FusedCompressPlan:
|
|
compress_plan_gpu: torch.Tensor
|
|
write_plan_gpu: torch.Tensor
|
|
num_compress: int
|
|
num_write: int
|
|
|
|
|
|
def write_current_token_to_state(
|
|
kv_score_input: torch.Tensor,
|
|
write_plan: torch.Tensor,
|
|
state_pool_buffer: torch.Tensor,
|
|
head_dim: int,
|
|
overlap: bool,
|
|
ratio: int,
|
|
) -> None:
|
|
"""Reference write path used by this manual test.
|
|
|
|
Plan row layout: [ragged_id, batch_id, position, window_len, state_base].
|
|
"""
|
|
del head_dim # layout is already encoded in kv_score_input/state_pool_buffer shape.
|
|
state_size = (2 if overlap else 1) * ratio
|
|
plan_cpu = write_plan.cpu()
|
|
for row in plan_cpu:
|
|
ragged_id = int(row[0].item())
|
|
position = int(row[2].item())
|
|
state_base = int(row[4].item())
|
|
if ragged_id < 0 or position < 0:
|
|
continue
|
|
dst = state_base + (position % state_size)
|
|
if (
|
|
0 <= dst < state_pool_buffer.shape[0]
|
|
and 0 <= ragged_id < kv_score_input.shape[0]
|
|
):
|
|
state_pool_buffer[dst] = kv_score_input[ragged_id]
|
|
|
|
|
|
def fused_compress_attn(
|
|
state_pool_buffer: torch.Tensor,
|
|
plan: torch.Tensor,
|
|
ape: torch.Tensor,
|
|
rms_weight: torch.Tensor,
|
|
rms_eps: float,
|
|
freqs_cis_real: torch.Tensor,
|
|
head_dim: int,
|
|
rope_head_dim: int,
|
|
overlap: bool,
|
|
ratio: int,
|
|
out: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Reference compress path for manual parity tests.
|
|
|
|
This keeps the test runnable after removing `fused_compress_kernel.py`.
|
|
"""
|
|
freqs_cis = torch.view_as_complex(
|
|
freqs_cis_real.view(freqs_cis_real.shape[0], -1, 2).contiguous()
|
|
)
|
|
result = _ref_compress(
|
|
kv_score_input=torch.empty(
|
|
0, device=state_pool_buffer.device, dtype=torch.float32
|
|
),
|
|
state_pool=state_pool_buffer,
|
|
plan=plan,
|
|
ape=ape,
|
|
rms_weight=rms_weight,
|
|
rms_eps=rms_eps,
|
|
freqs_cis=freqs_cis,
|
|
head_dim=head_dim,
|
|
rope_head_dim=rope_head_dim,
|
|
overlap=overlap,
|
|
ratio=ratio,
|
|
num_compress=plan.shape[0],
|
|
)
|
|
out.copy_(result)
|
|
return out
|
|
|
|
|
|
def _make_plan_from_params(
|
|
extend_lens: list[int],
|
|
seq_lens: list[int],
|
|
ratio: int,
|
|
overlap: bool,
|
|
state_bases: list[int],
|
|
device: torch.device,
|
|
) -> FusedCompressPlan:
|
|
"""Build a test plan without requiring real SWA / req_to_token tables."""
|
|
bs = len(extend_lens)
|
|
ext = np.array(extend_lens, dtype=np.int32)
|
|
seq = np.array(seq_lens, dtype=np.int32)
|
|
total = int(ext.sum())
|
|
|
|
state_size = (2 if overlap else 1) * ratio
|
|
K = state_size
|
|
|
|
batch_ids = np.repeat(np.arange(bs, dtype=np.int32), ext)
|
|
ragged_ids = np.arange(total, dtype=np.int32)
|
|
cu_extend = np.empty(bs + 1, dtype=np.int32)
|
|
cu_extend[0] = 0
|
|
np.cumsum(ext, out=cu_extend[1:])
|
|
j_in_seq = ragged_ids - cu_extend[batch_ids]
|
|
prefix_lens = seq - ext
|
|
positions = prefix_lens[batch_ids] + j_in_seq
|
|
|
|
window_lens = np.maximum(0, K - np.minimum(j_in_seq + 1, K)).astype(np.int32)
|
|
state_base_arr = np.array(state_bases, dtype=np.int32)
|
|
state_base_per_token = state_base_arr[batch_ids]
|
|
|
|
plan_rows = np.stack(
|
|
[ragged_ids, batch_ids, positions, window_lens, state_base_per_token],
|
|
axis=1,
|
|
).astype(np.int32)
|
|
|
|
compress_mask = (positions + 1) % ratio == 0
|
|
compress_plan = plan_rows[compress_mask]
|
|
|
|
write_starts = np.maximum(0, seq - K).astype(np.int32)
|
|
write_mask = positions >= write_starts[batch_ids]
|
|
write_plan = plan_rows[write_mask]
|
|
|
|
n_compress = int(compress_plan.shape[0]) if compress_plan.size > 0 else 0
|
|
n_write = int(write_plan.shape[0]) if write_plan.size > 0 else 0
|
|
|
|
compress_gpu = (
|
|
torch.from_numpy(np.ascontiguousarray(compress_plan)).to(device)
|
|
if n_compress > 0
|
|
else torch.empty((0, 5), dtype=torch.int32, device=device)
|
|
)
|
|
write_gpu = (
|
|
torch.from_numpy(np.ascontiguousarray(write_plan)).to(device)
|
|
if n_write > 0
|
|
else torch.empty((0, 5), dtype=torch.int32, device=device)
|
|
)
|
|
|
|
return FusedCompressPlan(
|
|
compress_plan_gpu=compress_gpu,
|
|
write_plan_gpu=write_gpu,
|
|
num_compress=n_compress,
|
|
num_write=n_write,
|
|
)
|
|
|
|
|
|
def _make_freqs_cis(max_seq: int, rope_dim: int, device: torch.device) -> torch.Tensor:
|
|
"""Create test freqs_cis as complex64 [max_seq, rope_dim/2], matching production."""
|
|
half = rope_dim // 2
|
|
angles = torch.randn(max_seq, half, device=device, dtype=torch.float32) * 0.1
|
|
return torch.polar(torch.ones_like(angles), angles)
|
|
|
|
|
|
def _freqs_to_real(freqs_cis: torch.Tensor) -> torch.Tensor:
|
|
"""Convert complex64 freqs to float32 [max_seq, rope_dim] interleaved."""
|
|
return torch.view_as_real(freqs_cis).flatten(-2).contiguous()
|
|
|
|
|
|
def _ref_compress(
|
|
kv_score_input: torch.Tensor,
|
|
state_pool: torch.Tensor,
|
|
plan: torch.Tensor,
|
|
ape: torch.Tensor,
|
|
rms_weight: torch.Tensor,
|
|
rms_eps: float,
|
|
freqs_cis: torch.Tensor,
|
|
head_dim: int,
|
|
rope_head_dim: int,
|
|
overlap: bool,
|
|
ratio: int,
|
|
num_compress: int,
|
|
) -> torch.Tensor:
|
|
"""Pure-PyTorch reference matching SGLang compress_decode_paged semantics.
|
|
|
|
State already has current tokens written (no APE).
|
|
APE is added to ALL K scores at compress time.
|
|
"""
|
|
if num_compress == 0:
|
|
return torch.empty(0, head_dim, dtype=torch.float32, device=state_pool.device)
|
|
|
|
coff = 2 if overlap else 1
|
|
half_dim = coff * head_dim
|
|
state_size = coff * ratio
|
|
K = state_size
|
|
|
|
plan_cpu = plan[:num_compress].cpu()
|
|
out = torch.empty(
|
|
num_compress, head_dim, dtype=torch.float32, device=state_pool.device
|
|
)
|
|
|
|
for pid in range(num_compress):
|
|
position = int(plan_cpu[pid, 2].item())
|
|
state_base = int(plan_cpu[pid, 4].item())
|
|
|
|
if position < 0:
|
|
continue
|
|
|
|
kv_rows = []
|
|
score_rows = []
|
|
for k in range(K):
|
|
s = position - K + 1 + k
|
|
col_off = (head_dim if k >= ratio else 0) if overlap else 0
|
|
ape_row = k % ratio
|
|
d_slice = slice(col_off, col_off + head_dim)
|
|
|
|
if s < 0:
|
|
kv_rows.append(
|
|
torch.zeros(head_dim, dtype=torch.float32, device=state_pool.device)
|
|
)
|
|
score_rows.append(
|
|
torch.full(
|
|
(head_dim,),
|
|
float("-inf"),
|
|
dtype=torch.float32,
|
|
device=state_pool.device,
|
|
)
|
|
)
|
|
else:
|
|
ring = s % state_size
|
|
row = state_pool[state_base + ring]
|
|
kv_rows.append(row[d_slice].float())
|
|
# APE added to ALL scores
|
|
score_rows.append(
|
|
row[half_dim + col_off : half_dim + col_off + head_dim].float()
|
|
+ ape[ape_row, d_slice].float()
|
|
)
|
|
|
|
kv_stack = torch.stack(kv_rows, dim=0)
|
|
sc_stack = torch.stack(score_rows, dim=0)
|
|
weights = torch.softmax(sc_stack, dim=0)
|
|
compressed = (weights * kv_stack).sum(dim=0)
|
|
|
|
var = (compressed * compressed).mean()
|
|
normed = compressed * torch.rsqrt(var + rms_eps) * rms_weight.float()
|
|
|
|
comp_pos = (position // ratio) * ratio
|
|
rope_seg = normed[-rope_head_dim:].clone()
|
|
freqs_row = torch.view_as_real(freqs_cis[comp_pos]).flatten()
|
|
cos_v = freqs_row[0::2].float()
|
|
sin_v = freqs_row[1::2].float()
|
|
|
|
even = rope_seg[0::2]
|
|
odd = rope_seg[1::2]
|
|
normed[-rope_head_dim:] = torch.stack(
|
|
[even * cos_v - odd * sin_v, odd * cos_v + even * sin_v], dim=-1
|
|
).flatten()
|
|
|
|
out[pid] = normed
|
|
|
|
return out
|
|
|
|
|
|
class TestFusedCompressAttn(unittest.TestCase):
|
|
def _run_test(
|
|
self,
|
|
ratio: int,
|
|
overlap: bool,
|
|
bs: int,
|
|
extend_lens: list[int],
|
|
prefix_lens: list[int],
|
|
head_dim: int = 512,
|
|
rope_head_dim: int = 64,
|
|
):
|
|
device = torch.device("cuda")
|
|
torch.manual_seed(42)
|
|
coff = 2 if overlap else 1
|
|
half_dim = coff * head_dim
|
|
last_dim = 2 * half_dim
|
|
state_size = coff * ratio
|
|
|
|
seq_lens = [p + e for p, e in zip(prefix_lens, extend_lens)]
|
|
total_tokens = sum(extend_lens)
|
|
max_seq = max(seq_lens) + 128
|
|
|
|
kv_score_input = torch.randn(
|
|
total_tokens, last_dim, device=device, dtype=torch.float32
|
|
)
|
|
|
|
pool_size = bs * state_size + 2
|
|
state_pool = torch.randn(
|
|
pool_size, last_dim, device=device, dtype=torch.float32
|
|
)
|
|
state_pool[:, half_dim:] *= 0.5 # reasonable score magnitudes
|
|
|
|
state_bases = [i * state_size for i in range(bs)]
|
|
ape = torch.randn(ratio, half_dim, device=device, dtype=torch.float32) * 0.1
|
|
rms_weight = torch.ones(head_dim, device=device, dtype=torch.float32)
|
|
rms_eps = 1e-6
|
|
freqs_cis = _make_freqs_cis(max_seq, rope_head_dim, device)
|
|
freqs_real = _freqs_to_real(freqs_cis)
|
|
|
|
plan = _make_plan_from_params(
|
|
extend_lens, seq_lens, ratio, overlap, state_bases, device
|
|
)
|
|
if plan.num_compress == 0:
|
|
return
|
|
|
|
# Step 1: write current tokens to state (same for both paths)
|
|
state_triton = state_pool.clone()
|
|
state_ref = state_pool.clone()
|
|
|
|
write_current_token_to_state(
|
|
kv_score_input=kv_score_input,
|
|
write_plan=plan.write_plan_gpu,
|
|
state_pool_buffer=state_triton,
|
|
head_dim=head_dim,
|
|
overlap=overlap,
|
|
ratio=ratio,
|
|
)
|
|
# Reference: same write
|
|
write_current_token_to_state(
|
|
kv_score_input=kv_score_input,
|
|
write_plan=plan.write_plan_gpu,
|
|
state_pool_buffer=state_ref,
|
|
head_dim=head_dim,
|
|
overlap=overlap,
|
|
ratio=ratio,
|
|
)
|
|
|
|
# Step 2a: Triton fused compress
|
|
out_triton = torch.empty(
|
|
plan.num_compress, head_dim, device=device, dtype=torch.float32
|
|
)
|
|
fused_compress_attn(
|
|
state_pool_buffer=state_triton,
|
|
plan=plan.compress_plan_gpu,
|
|
ape=ape,
|
|
rms_weight=rms_weight,
|
|
rms_eps=rms_eps,
|
|
freqs_cis_real=freqs_real,
|
|
head_dim=head_dim,
|
|
rope_head_dim=rope_head_dim,
|
|
overlap=overlap,
|
|
ratio=ratio,
|
|
out=out_triton,
|
|
)
|
|
|
|
# Step 2b: reference compress
|
|
out_ref = _ref_compress(
|
|
kv_score_input=kv_score_input,
|
|
state_pool=state_ref,
|
|
plan=plan.compress_plan_gpu,
|
|
ape=ape,
|
|
rms_weight=rms_weight,
|
|
rms_eps=rms_eps,
|
|
freqs_cis=freqs_cis,
|
|
head_dim=head_dim,
|
|
rope_head_dim=rope_head_dim,
|
|
overlap=overlap,
|
|
ratio=ratio,
|
|
num_compress=plan.num_compress,
|
|
)
|
|
|
|
torch.testing.assert_close(out_triton, out_ref, atol=1e-3, rtol=1e-3)
|
|
|
|
def test_hca_single(self):
|
|
self._run_test(
|
|
ratio=128, overlap=False, bs=1, extend_lens=[128], prefix_lens=[0]
|
|
)
|
|
|
|
def test_hca_multi(self):
|
|
self._run_test(
|
|
ratio=128, overlap=False, bs=2, extend_lens=[128, 256], prefix_lens=[0, 128]
|
|
)
|
|
|
|
def test_csa_single(self):
|
|
self._run_test(ratio=4, overlap=True, bs=1, extend_lens=[16], prefix_lens=[8])
|
|
|
|
def test_csa_multi(self):
|
|
self._run_test(
|
|
ratio=4, overlap=True, bs=3, extend_lens=[8, 12, 16], prefix_lens=[4, 8, 0]
|
|
)
|
|
|
|
def test_csa_small_dim(self):
|
|
self._run_test(
|
|
ratio=4,
|
|
overlap=True,
|
|
bs=2,
|
|
extend_lens=[8, 8],
|
|
prefix_lens=[4, 0],
|
|
head_dim=256,
|
|
)
|
|
|
|
|
|
class TestStateOrdering(unittest.TestCase):
|
|
def test_write_then_compress(self):
|
|
"""Verify write-first, compress-second matches reference."""
|
|
device = torch.device("cuda")
|
|
torch.manual_seed(123)
|
|
ratio, overlap = 4, True
|
|
coff = 2
|
|
head_dim, rope_head_dim = 128, 64
|
|
half_dim = coff * head_dim
|
|
last_dim = 2 * half_dim
|
|
state_size = coff * ratio
|
|
|
|
pool_size = state_size + 2
|
|
state_pool = torch.randn(
|
|
pool_size, last_dim, device=device, dtype=torch.float32
|
|
)
|
|
state_pool[:, half_dim:] *= 0.5
|
|
|
|
kv_score_input = torch.randn(8, last_dim, device=device, dtype=torch.float32)
|
|
ape = torch.randn(ratio, half_dim, device=device, dtype=torch.float32) * 0.1
|
|
rms_weight = torch.ones(head_dim, device=device, dtype=torch.float32)
|
|
freqs_cis = _make_freqs_cis(64, rope_head_dim, device)
|
|
|
|
plan = _make_plan_from_params([8], [8], ratio, overlap, [0], device)
|
|
if plan.num_compress == 0:
|
|
return
|
|
|
|
state_before = state_pool.clone()
|
|
|
|
# Write first
|
|
write_current_token_to_state(
|
|
kv_score_input=kv_score_input,
|
|
write_plan=plan.write_plan_gpu,
|
|
state_pool_buffer=state_pool,
|
|
head_dim=head_dim,
|
|
overlap=overlap,
|
|
ratio=ratio,
|
|
)
|
|
|
|
# State should now be different (tokens written)
|
|
self.assertFalse(torch.allclose(state_pool, state_before))
|
|
|
|
# Compress
|
|
out = torch.empty(
|
|
plan.num_compress, head_dim, device=device, dtype=torch.float32
|
|
)
|
|
fused_compress_attn(
|
|
state_pool_buffer=state_pool,
|
|
plan=plan.compress_plan_gpu,
|
|
ape=ape,
|
|
rms_weight=rms_weight,
|
|
rms_eps=1e-6,
|
|
freqs_cis_real=_freqs_to_real(freqs_cis),
|
|
head_dim=head_dim,
|
|
rope_head_dim=rope_head_dim,
|
|
overlap=overlap,
|
|
ratio=ratio,
|
|
out=out,
|
|
)
|
|
|
|
self.assertFalse(torch.any(torch.isnan(out)).item())
|
|
self.assertFalse(torch.any(torch.isinf(out)).item())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|