388 lines
14 KiB
Python
388 lines
14 KiB
Python
import unittest
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch.nn.functional import softplus
|
|
from utils import precision
|
|
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
torch.manual_seed(1234)
|
|
|
|
|
|
def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
|
|
"""This function is intended to align with the l2norm implementation in the FLA library."""
|
|
inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
|
|
return x * inv_norm
|
|
|
|
|
|
def torch_chunk_gated_delta_rule(
|
|
query,
|
|
key,
|
|
value,
|
|
g,
|
|
beta,
|
|
chunk_size=64,
|
|
initial_state=None,
|
|
output_final_state=False,
|
|
use_qk_l2norm_in_kernel=False,
|
|
):
|
|
initial_dtype = query.dtype
|
|
if use_qk_l2norm_in_kernel:
|
|
query = l2norm(query, dim=-1, eps=1e-6)
|
|
key = l2norm(key, dim=-1, eps=1e-6)
|
|
query, key, value, beta, g = [
|
|
x.transpose(1, 2).contiguous().to(torch.float32)
|
|
for x in (query, key, value, beta, g)
|
|
]
|
|
|
|
batch_size, sequence_length, num_heads, k_head_dim = key.shape
|
|
v_head_dim = value.shape[-1]
|
|
pad_size = (chunk_size - num_heads % chunk_size) % chunk_size
|
|
query = F.pad(query, (0, 0, 0, pad_size))
|
|
key = F.pad(key, (0, 0, 0, pad_size))
|
|
value = F.pad(value, (0, 0, 0, pad_size))
|
|
beta = F.pad(beta, (0, pad_size))
|
|
g = F.pad(g, (0, pad_size))
|
|
tot_heads = num_heads + pad_size
|
|
scale = 1 / (query.shape[-1] ** 0.5)
|
|
query = query * scale
|
|
|
|
v_beta = value * beta.unsqueeze(-1)
|
|
k_beta = key * beta.unsqueeze(-1)
|
|
# reshape to chunks
|
|
query, key, value, k_beta, v_beta = [
|
|
x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
|
|
for x in (query, key, value, k_beta, v_beta)
|
|
]
|
|
g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
|
|
mask = torch.triu(
|
|
torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
|
|
diagonal=0,
|
|
)
|
|
|
|
# chunk decay
|
|
g = g.cumsum(dim=-1)
|
|
decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
|
|
attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
|
|
for i in range(1, chunk_size):
|
|
row = attn[..., i, :i].clone()
|
|
sub = attn[..., :i, :i].clone()
|
|
attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
|
|
attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
|
|
value = attn @ v_beta
|
|
k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
|
|
last_recurrent_state = (
|
|
torch.zeros(batch_size, sequence_length, k_head_dim, v_head_dim).to(value)
|
|
if initial_state is None
|
|
else initial_state.to(value)
|
|
)
|
|
core_attn_out = torch.zeros_like(value)
|
|
mask = torch.triu(
|
|
torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
|
|
diagonal=1,
|
|
)
|
|
|
|
# for each chunk
|
|
for i in range(0, tot_heads // chunk_size):
|
|
q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
|
|
attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
|
|
v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state
|
|
v_new = v_i - v_prime
|
|
attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state
|
|
core_attn_out[:, :, i] = attn_inter + attn @ v_new
|
|
last_recurrent_state = (
|
|
last_recurrent_state * g[:, :, i, -1, None, None].exp()
|
|
+ (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(
|
|
-1, -2
|
|
)
|
|
@ v_new
|
|
)
|
|
|
|
if not output_final_state:
|
|
last_recurrent_state = None
|
|
core_attn_out = core_attn_out.reshape(
|
|
core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1]
|
|
)
|
|
core_attn_out = core_attn_out[:, :, :num_heads]
|
|
core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
|
|
return core_attn_out, last_recurrent_state
|
|
|
|
|
|
def chunk_gated_delta_rule_update(
|
|
query, # [B, T, HK, K]
|
|
key, # [B, T, HK, K]
|
|
value, # [B, T, HV, V]
|
|
g, # [B, T, HV]
|
|
beta, # [B, T, HV]
|
|
cu_seqlens, # [N+1]
|
|
initial_state, # [N, HV, K, V]
|
|
use_qk_l2norm_in_kernel, # True
|
|
):
|
|
num_heads = query.shape[2]
|
|
num_value_heads = value.shape[2]
|
|
batch_size = initial_state.shape[0]
|
|
if num_value_heads // num_heads > 1:
|
|
query = query.repeat_interleave(num_value_heads // num_heads, dim=2)
|
|
key = key.repeat_interleave(num_value_heads // num_heads, dim=2)
|
|
output = torch.empty_like(value)
|
|
final_state = torch.empty_like(initial_state)
|
|
start_q = 0
|
|
for i in range(batch_size):
|
|
end_q = cu_seqlens[i + 1]
|
|
core_attn_outi, last_recurrent_state = torch_chunk_gated_delta_rule(
|
|
query=query[:, start_q:end_q, :, :],
|
|
key=key[:, start_q:end_q, :, :],
|
|
value=value[:, start_q:end_q, :, :],
|
|
g=g[:, start_q:end_q, :],
|
|
beta=beta[:, start_q:end_q, :],
|
|
initial_state=initial_state[i],
|
|
output_final_state=True,
|
|
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
|
|
)
|
|
output[:, start_q:end_q, :, :] = core_attn_outi
|
|
final_state[i] = last_recurrent_state
|
|
start_q = end_q
|
|
return output, final_state
|
|
|
|
|
|
def torch_recurrent_gated_delta_rule(
|
|
query,
|
|
key,
|
|
value,
|
|
g,
|
|
beta,
|
|
initial_state,
|
|
output_final_state,
|
|
use_qk_l2norm_in_kernel=False,
|
|
):
|
|
initial_dtype = query.dtype
|
|
if use_qk_l2norm_in_kernel:
|
|
query = l2norm(query, dim=-1, eps=1e-6)
|
|
key = l2norm(key, dim=-1, eps=1e-6)
|
|
query, key, value, beta, g = [
|
|
x.transpose(1, 2).contiguous().to(torch.float32)
|
|
for x in (query, key, value, beta, g)
|
|
]
|
|
|
|
batch_size, num_heads, sequence_length, k_head_dim = key.shape
|
|
v_head_dim = value.shape[-1]
|
|
scale = 1 / (query.shape[-1] ** 0.5)
|
|
query = query * scale
|
|
|
|
core_attn_out = torch.zeros(batch_size, num_heads, sequence_length, v_head_dim).to(
|
|
value
|
|
)
|
|
last_recurrent_state = (
|
|
torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
|
|
if initial_state is None
|
|
else initial_state.to(value)
|
|
)
|
|
|
|
for i in range(sequence_length):
|
|
q_t = query[:, :, i]
|
|
k_t = key[:, :, i]
|
|
v_t = value[:, :, i]
|
|
g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1)
|
|
beta_t = beta[:, :, i].unsqueeze(-1)
|
|
|
|
last_recurrent_state = last_recurrent_state * g_t
|
|
kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2)
|
|
delta = (v_t - kv_mem) * beta_t
|
|
last_recurrent_state = last_recurrent_state + k_t.unsqueeze(
|
|
-1
|
|
) * delta.unsqueeze(-2)
|
|
core_attn_out[:, :, i] = (last_recurrent_state * q_t.unsqueeze(-1)).sum(dim=-2)
|
|
|
|
if not output_final_state:
|
|
last_recurrent_state = None
|
|
core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
|
|
return core_attn_out, last_recurrent_state
|
|
|
|
|
|
def sigmoid_gating_delta_rule_update(
|
|
query,
|
|
key,
|
|
value,
|
|
A_log,
|
|
a,
|
|
dt_bias,
|
|
b,
|
|
initial_state,
|
|
output_final_state,
|
|
use_qk_l2norm_in_kernel=False,
|
|
):
|
|
beta = b.sigmoid()
|
|
g = -A_log.float().exp() * softplus(a.float() + dt_bias)
|
|
return torch_recurrent_gated_delta_rule(
|
|
query,
|
|
key,
|
|
value,
|
|
g.unsqueeze(0),
|
|
beta.unsqueeze(0),
|
|
initial_state,
|
|
output_final_state,
|
|
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
|
|
)
|
|
|
|
|
|
def torch_gdn_gating(A_log, a, b, dt_bias):
|
|
return -A_log.float().exp() * softplus(a.float() + dt_bias).unsqueeze(
|
|
0
|
|
), b.sigmoid().unsqueeze(0)
|
|
|
|
|
|
class TestMambaAttention(CustomTestCase):
|
|
def test_chunk_gated_delta_rule(self):
|
|
B, L, HK, HV, EK, EV, N = 1, 100, 3, 6, 64, 64, 4
|
|
seqlens = torch.randint(1, L, (N + 1,))
|
|
seqlens[0] = 0
|
|
cu_seqlens_ = torch.cumsum(seqlens, dim=0).to(torch.int32)
|
|
T = cu_seqlens_[-1].item()
|
|
query_ = torch.rand((B, T, HK, EK), dtype=torch.bfloat16) * 0.05
|
|
key_ = torch.rand((B, T, HK, EK), dtype=torch.bfloat16) * 0.05
|
|
value_ = torch.rand((B, T, HV, EV), dtype=torch.bfloat16) * 0.05
|
|
g_ = torch.rand((B, T, HV), dtype=torch.float32) * 0.05
|
|
beta_ = torch.rand((B, T, HV), dtype=torch.bfloat16) * 0.05
|
|
initial_state_ = torch.rand((N, HV, EK, EV), dtype=torch.float32) * 0.05
|
|
|
|
for use_qk_l2norm_in_kernel in [True, False]:
|
|
core_attn_out_ref, last_recurrent_state_ref = chunk_gated_delta_rule_update(
|
|
query=query_,
|
|
key=key_,
|
|
value=value_,
|
|
g=g_,
|
|
beta=beta_,
|
|
cu_seqlens=cu_seqlens_,
|
|
initial_state=initial_state_,
|
|
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
|
|
)
|
|
|
|
query = query_.clone()
|
|
key = key_.clone()
|
|
value = value_.clone()
|
|
g = g_.clone()
|
|
beta = beta_.clone()
|
|
cu_seqlens = cu_seqlens_.clone()
|
|
initial_state = initial_state_.clone()
|
|
|
|
core_attn_out, last_recurrent_state = (
|
|
torch.ops.sgl_kernel.chunk_gated_delta_rule_cpu(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
g=g,
|
|
beta=beta,
|
|
initial_state=initial_state,
|
|
output_final_state=True,
|
|
cu_seqlens=cu_seqlens,
|
|
head_first=False,
|
|
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
|
|
)
|
|
)
|
|
atol = rtol = precision[core_attn_out.dtype]
|
|
torch.testing.assert_close(
|
|
core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol
|
|
)
|
|
torch.testing.assert_close(
|
|
last_recurrent_state, last_recurrent_state_ref, atol=atol, rtol=rtol
|
|
)
|
|
|
|
def test_fused_gdn_gating(self):
|
|
dims = [6, 32]
|
|
for dim in dims:
|
|
A_log = torch.rand(dim)
|
|
a = torch.rand(1024, dim, dtype=torch.bfloat16)
|
|
b = torch.rand(1024, dim, dtype=torch.bfloat16)
|
|
dt_bias = torch.rand(dim, dtype=torch.bfloat16)
|
|
|
|
g, beta = torch_gdn_gating(A_log, a, b, dt_bias)
|
|
g_sgl, beta_sgl = torch.ops.sgl_kernel.fused_gdn_gating_cpu(
|
|
A_log, a, b, dt_bias
|
|
)
|
|
atol = rtol = precision[g.dtype]
|
|
atol2 = rtol2 = precision[beta.dtype]
|
|
torch.testing.assert_close(g, g_sgl, atol=atol, rtol=rtol)
|
|
torch.testing.assert_close(beta, beta_sgl, atol=atol2, rtol=rtol2)
|
|
|
|
def test_fused_sigmoid_gating_delta_rule_update(self):
|
|
batch_size = 1
|
|
num_value_heads = 32
|
|
head_k_dim = 128
|
|
head_v_dim = 128
|
|
num_heads = 16
|
|
seq_len = 1
|
|
attn_tp_size = 1
|
|
key_dim = head_k_dim * num_heads
|
|
value_dim = head_v_dim * num_value_heads
|
|
mixed_qkv_dim = (key_dim * 2 + value_dim) // attn_tp_size
|
|
mixed_qkv = torch.rand(
|
|
seq_len * batch_size, mixed_qkv_dim, dtype=torch.bfloat16
|
|
)
|
|
query, key, value = torch.split(
|
|
mixed_qkv,
|
|
[
|
|
key_dim // attn_tp_size,
|
|
key_dim // attn_tp_size,
|
|
value_dim // attn_tp_size,
|
|
],
|
|
dim=-1,
|
|
)
|
|
query = query.view(1, seq_len, num_heads, head_k_dim)
|
|
key = key.view(1, seq_len, num_heads, head_k_dim)
|
|
value = value.view(1, seq_len, num_value_heads, head_v_dim)
|
|
A_log = torch.rand(num_value_heads, dtype=torch.float32)
|
|
a = torch.rand(batch_size, num_value_heads, dtype=torch.bfloat16)
|
|
b = torch.rand(batch_size, num_value_heads, dtype=torch.bfloat16)
|
|
dt_bias = torch.rand(num_value_heads, dtype=torch.bfloat16)
|
|
ssm_states = torch.rand(
|
|
513, num_value_heads, head_k_dim, head_v_dim, dtype=torch.float32
|
|
)
|
|
cache_indices = torch.randint(0, 513, (batch_size,), dtype=torch.int32)
|
|
query_start_loc = torch.tensor([0, 1], dtype=torch.int32)
|
|
use_qk_l2norm_in_kernel = True
|
|
query_ref = query.clone()
|
|
key_ref = key.clone()
|
|
if num_value_heads // num_heads > 1:
|
|
query_ref = query_ref.repeat_interleave(num_value_heads // num_heads, dim=2)
|
|
key_ref = key_ref.repeat_interleave(num_value_heads // num_heads, dim=2)
|
|
core_attn_out_ref, last_recurrent_state_ref = sigmoid_gating_delta_rule_update(
|
|
query_ref.transpose(0, 1),
|
|
key_ref.transpose(0, 1),
|
|
value.transpose(0, 1),
|
|
A_log,
|
|
a,
|
|
dt_bias,
|
|
b,
|
|
initial_state=ssm_states[cache_indices],
|
|
output_final_state=True,
|
|
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
|
|
)
|
|
core_attn_out = torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu(
|
|
A_log=A_log,
|
|
dt_bias=dt_bias,
|
|
q=query,
|
|
k=key,
|
|
v=value,
|
|
a=a,
|
|
b=b,
|
|
initial_state_source=ssm_states,
|
|
initial_state_indices=cache_indices,
|
|
cu_seqlens=query_start_loc,
|
|
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
|
|
softplus_beta=1.0,
|
|
softplus_threshold=20.0,
|
|
)
|
|
last_recurrent_state = ssm_states[cache_indices]
|
|
atol = rtol = precision[core_attn_out.dtype]
|
|
torch.testing.assert_close(
|
|
core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol
|
|
)
|
|
torch.testing.assert_close(
|
|
last_recurrent_state, last_recurrent_state_ref, atol=atol, rtol=rtol
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|