[GDN] Change Attention State Layout from [N, HV, K, V] to [N, HV, V, K] (#20283)
This commit is contained in:
Yuan Luo
274
test/registered/attention/test_chunk_gated_delta_rule.py
Normal file
274
test/registered/attention/test_chunk_gated_delta_rule.py
Normal file
@@ -0,0 +1,274 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule
|
||||
from sglang.srt.layers.attention.fla.fused_recurrent import (
|
||||
fused_recurrent_gated_delta_rule,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=60, suite="stage-b-test-large-1-gpu")
|
||||
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
|
||||
class TestChunkGatedDeltaRule(unittest.TestCase):
|
||||
"""Test chunk_gated_delta_rule against token-by-token fused_recurrent reference."""
|
||||
|
||||
ATOL = 2e-2
|
||||
RTOL = 1e-2
|
||||
|
||||
def _run_reference(self, pool_init, cache_indices, q, k, v, g, beta):
|
||||
"""Per-batch token-by-token reference using fused_recurrent_gated_delta_rule.
|
||||
|
||||
initial_state shape: [N, H, V, K] (native layout on this branch).
|
||||
"""
|
||||
B = cache_indices.shape[0]
|
||||
T_per_seq = q.shape[1] // B
|
||||
pool = pool_init.clone()
|
||||
h_cur = pool[cache_indices].contiguous().clone()
|
||||
|
||||
o_list = []
|
||||
for b in range(B):
|
||||
sl = slice(b * T_per_seq, (b + 1) * T_per_seq)
|
||||
o_b, h_b = fused_recurrent_gated_delta_rule(
|
||||
q=q[0, sl].unsqueeze(0),
|
||||
k=k[0, sl].unsqueeze(0),
|
||||
v=v[0, sl].unsqueeze(0),
|
||||
g=g[0, sl].unsqueeze(0),
|
||||
beta=beta[0, sl].unsqueeze(0),
|
||||
initial_state=h_cur[b : b + 1],
|
||||
output_final_state=True,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
o_list.append(o_b)
|
||||
h_cur[b] = h_b[0]
|
||||
|
||||
pool[cache_indices] = h_cur
|
||||
return torch.cat(o_list, dim=1), pool
|
||||
|
||||
def _run_chunk(self, pool_init, cache_indices, q, k, v, g, beta, cu_seqlens):
|
||||
"""Run chunk_gated_delta_rule with native [V, K] pool."""
|
||||
pool = pool_init.clone()
|
||||
o, _, _ = chunk_gated_delta_rule(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
g=g,
|
||||
beta=beta,
|
||||
initial_state=pool,
|
||||
initial_state_indices=cache_indices,
|
||||
cu_seqlens=cu_seqlens,
|
||||
head_first=False,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
return o, pool
|
||||
|
||||
def _check_shape(
|
||||
self, B, T_per_seq, H, K, V, pool_size, sequential_indices=False, seed=42
|
||||
):
|
||||
"""Run correctness check for one (B, T_per_seq, H, K, V, pool_size) config."""
|
||||
device = "cuda"
|
||||
dtype = torch.bfloat16
|
||||
T = B * T_per_seq
|
||||
|
||||
torch.manual_seed(seed)
|
||||
|
||||
if sequential_indices:
|
||||
cache_indices = torch.arange(B, dtype=torch.int32, device=device)
|
||||
else:
|
||||
perm = torch.randperm(pool_size, device=device)[:B]
|
||||
cache_indices = perm.to(torch.int32)
|
||||
|
||||
pool_init = (
|
||||
torch.randn(pool_size, H, V, K, dtype=torch.float32, device=device) * 0.1
|
||||
)
|
||||
cu_seqlens = torch.zeros(B + 1, dtype=torch.long, device=device)
|
||||
cu_seqlens[1:] = (
|
||||
torch.arange(1, B + 1, dtype=torch.long, device=device) * T_per_seq
|
||||
)
|
||||
|
||||
q = torch.randn(1, T, H, K, dtype=dtype, device=device)
|
||||
k = torch.randn(1, T, H, K, dtype=dtype, device=device)
|
||||
v = torch.randn(1, T, H, V, dtype=dtype, device=device)
|
||||
g = torch.nn.functional.logsigmoid(
|
||||
torch.randn(1, T, H, dtype=dtype, device=device)
|
||||
)
|
||||
beta = torch.sigmoid(torch.randn(1, T, H, dtype=dtype, device=device))
|
||||
|
||||
o_ref, pool_ref = self._run_reference(
|
||||
pool_init, cache_indices, q, k, v, g, beta
|
||||
)
|
||||
o_new, pool_new = self._run_chunk(
|
||||
pool_init, cache_indices, q, k, v, g, beta, cu_seqlens
|
||||
)
|
||||
|
||||
self.assertTrue(
|
||||
torch.allclose(
|
||||
o_ref.float(), o_new.float(), atol=self.ATOL, rtol=self.RTOL
|
||||
),
|
||||
f"Output mismatch: max_diff="
|
||||
f"{(o_ref.float() - o_new.float()).abs().max().item():.2e}",
|
||||
)
|
||||
|
||||
ref_slots = pool_ref[cache_indices].contiguous()
|
||||
new_slots = pool_new[cache_indices].contiguous()
|
||||
self.assertTrue(
|
||||
torch.allclose(
|
||||
ref_slots.float(), new_slots.float(), atol=self.ATOL, rtol=self.RTOL
|
||||
),
|
||||
f"State mismatch: max_diff="
|
||||
f"{(ref_slots.float() - new_slots.float()).abs().max().item():.2e}",
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Production-style configs (Qwen3-Next)
|
||||
# ------------------------------------------------------------------
|
||||
def test_production_nt1(self):
|
||||
self._check_shape(B=4, T_per_seq=64, H=16, K=128, V=128, pool_size=32)
|
||||
|
||||
def test_production_nt2(self):
|
||||
self._check_shape(B=4, T_per_seq=128, H=16, K=128, V=128, pool_size=32)
|
||||
|
||||
def test_production_nt4(self):
|
||||
self._check_shape(B=4, T_per_seq=256, H=16, K=128, V=128, pool_size=32)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Batch size sweep
|
||||
# ------------------------------------------------------------------
|
||||
def test_batch_1(self):
|
||||
self._check_shape(B=1, T_per_seq=128, H=16, K=128, V=128, pool_size=32)
|
||||
|
||||
def test_batch_2(self):
|
||||
self._check_shape(B=2, T_per_seq=128, H=16, K=128, V=128, pool_size=32)
|
||||
|
||||
def test_batch_8(self):
|
||||
self._check_shape(B=8, T_per_seq=128, H=16, K=128, V=128, pool_size=64)
|
||||
|
||||
def test_batch_16(self):
|
||||
self._check_shape(B=16, T_per_seq=64, H=16, K=128, V=128, pool_size=128)
|
||||
|
||||
def test_batch_32(self):
|
||||
self._check_shape(B=32, T_per_seq=32, H=16, K=128, V=128, pool_size=256)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Head count sweep
|
||||
# ------------------------------------------------------------------
|
||||
def test_heads_4(self):
|
||||
self._check_shape(B=4, T_per_seq=128, H=4, K=128, V=128, pool_size=32)
|
||||
|
||||
def test_heads_8(self):
|
||||
self._check_shape(B=4, T_per_seq=128, H=8, K=128, V=128, pool_size=32)
|
||||
|
||||
def test_heads_32(self):
|
||||
self._check_shape(B=4, T_per_seq=128, H=32, K=128, V=128, pool_size=32)
|
||||
|
||||
def test_heads_64(self):
|
||||
self._check_shape(B=4, T_per_seq=128, H=64, K=128, V=128, pool_size=32)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# K != V (exercises that [V,K] != [K,V] byte-order matters)
|
||||
# ------------------------------------------------------------------
|
||||
def test_dim_64x64(self):
|
||||
self._check_shape(B=4, T_per_seq=128, H=16, K=64, V=64, pool_size=32)
|
||||
|
||||
def test_dim_k_lt_v(self):
|
||||
self._check_shape(B=4, T_per_seq=128, H=16, K=64, V=128, pool_size=32)
|
||||
|
||||
def test_dim_k_gt_v(self):
|
||||
self._check_shape(B=4, T_per_seq=128, H=16, K=128, V=64, pool_size=32)
|
||||
|
||||
def test_dim_256x256(self):
|
||||
self._check_shape(B=4, T_per_seq=128, H=16, K=256, V=256, pool_size=32)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Short sequences (T < chunk_size=64)
|
||||
# ------------------------------------------------------------------
|
||||
def test_seqlen_1(self):
|
||||
self._check_shape(B=4, T_per_seq=1, H=16, K=128, V=128, pool_size=32)
|
||||
|
||||
def test_seqlen_7(self):
|
||||
self._check_shape(B=4, T_per_seq=7, H=16, K=128, V=128, pool_size=32)
|
||||
|
||||
def test_seqlen_16(self):
|
||||
self._check_shape(B=4, T_per_seq=16, H=16, K=128, V=128, pool_size=32)
|
||||
|
||||
def test_seqlen_32(self):
|
||||
self._check_shape(B=4, T_per_seq=32, H=16, K=128, V=128, pool_size=32)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Multi-chunk and large pool
|
||||
# ------------------------------------------------------------------
|
||||
def test_multi_chunk_nt8(self):
|
||||
self._check_shape(B=4, T_per_seq=512, H=16, K=128, V=128, pool_size=32)
|
||||
|
||||
def test_large_pool(self):
|
||||
self._check_shape(B=4, T_per_seq=128, H=16, K=128, V=128, pool_size=512)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Combined stress
|
||||
# ------------------------------------------------------------------
|
||||
def test_stress(self):
|
||||
self._check_shape(B=32, T_per_seq=128, H=32, K=128, V=128, pool_size=256)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Sequential-index variants (pool_size == B, indices = 0..B-1)
|
||||
# ------------------------------------------------------------------
|
||||
def test_seq_idx_b4(self):
|
||||
self._check_shape(
|
||||
B=4,
|
||||
T_per_seq=128,
|
||||
H=16,
|
||||
K=128,
|
||||
V=128,
|
||||
pool_size=4,
|
||||
sequential_indices=True,
|
||||
)
|
||||
|
||||
def test_seq_idx_b8(self):
|
||||
self._check_shape(
|
||||
B=8,
|
||||
T_per_seq=128,
|
||||
H=16,
|
||||
K=128,
|
||||
V=128,
|
||||
pool_size=8,
|
||||
sequential_indices=True,
|
||||
)
|
||||
|
||||
def test_seq_idx_h32(self):
|
||||
self._check_shape(
|
||||
B=4,
|
||||
T_per_seq=128,
|
||||
H=32,
|
||||
K=128,
|
||||
V=128,
|
||||
pool_size=4,
|
||||
sequential_indices=True,
|
||||
)
|
||||
|
||||
def test_seq_idx_h64(self):
|
||||
self._check_shape(
|
||||
B=4,
|
||||
T_per_seq=128,
|
||||
H=64,
|
||||
K=128,
|
||||
V=128,
|
||||
pool_size=4,
|
||||
sequential_indices=True,
|
||||
)
|
||||
|
||||
def test_seq_idx_stress(self):
|
||||
self._check_shape(
|
||||
B=32,
|
||||
T_per_seq=128,
|
||||
H=32,
|
||||
K=128,
|
||||
V=128,
|
||||
pool_size=32,
|
||||
sequential_indices=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user