[AMD] Dsv4/pr2 compressor opt (#26208)
Co-authored-by: wunhuang <wunhuang@amd.com> Co-authored-by: Thomas Wang <1am9trash@gmail.com> Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com> Co-authored-by: HaiShaw <hixiao@gmail.com> Co-authored-by: amd-danli103 <danli103@amd.com> Co-authored-by: Lin, Soga <soga.lin@amd.com> Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com> Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com> Co-authored-by: yichiche@amd.com <jacky.cheng> Co-authored-by: yctseng0211 <yctseng@amd.com> Co-authored-by: Bingxu Chen <bingxche@amd.com>
This commit is contained in:
kk
co-authored by
wunhuang
Thomas Wang
Xinyi Song
HaiShaw
amd-danli103
Lin, Soga
Raiden-Makoto
Hubert Lu
yichiche@amd.com
yctseng0211
Bingxu Chen
wunhuang
Thomas Wang
Xinyi Song
HaiShaw
amd-danli103
Lin, Soga
Raiden-Makoto
Hubert Lu
yichiche@amd.com
yctseng0211
Bingxu Chen
parent
7c0fbc8c2e
commit
3f5e2c7688
@@ -0,0 +1,465 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user