[CPU] update fla.cpp to support when num_head_v is not multiples of 16 (#30604)
This commit is contained in:
Ma Mingfei
@@ -135,8 +135,13 @@ struct l2norm_kernel<at::BFloat16, D, has_scale> {
|
||||
|
||||
template <typename scalar_t, int CHUNK_SIZE, int BLOCK_H>
|
||||
struct cumsum_kernel {
|
||||
static inline void
|
||||
apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int size, int ld_src, int ld_dst) {
|
||||
static inline void apply(
|
||||
scalar_t* __restrict__ out,
|
||||
const scalar_t* __restrict__ input,
|
||||
int mb_size,
|
||||
int hb_size,
|
||||
int ld_src,
|
||||
int ld_dst) {
|
||||
TORCH_CHECK(false, "cumsum_kernel: scalar path not implemented!");
|
||||
}
|
||||
};
|
||||
@@ -144,9 +149,12 @@ struct cumsum_kernel {
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
template <int CHUNK_SIZE, int BLOCK_H>
|
||||
struct cumsum_kernel<float, CHUNK_SIZE, BLOCK_H> {
|
||||
static inline void apply(float* __restrict__ out, const float* __restrict__ input, int size, int ld_src, int ld_dst) {
|
||||
static inline void
|
||||
apply(float* __restrict__ out, const float* __restrict__ input, int mb_size, int hb_size, int ld_src, int ld_dst) {
|
||||
// vector length of fp32 for avx512
|
||||
static_assert(BLOCK_H == 16);
|
||||
TORCH_CHECK(hb_size > 0 && hb_size <= BLOCK_H);
|
||||
const __mmask16 vmask = static_cast<__mmask16>((1u << hb_size) - 1u);
|
||||
|
||||
__m512i va[16];
|
||||
__m512 vsum = _mm512_set1_ps(0.f);
|
||||
@@ -154,14 +162,23 @@ struct cumsum_kernel<float, CHUNK_SIZE, BLOCK_H> {
|
||||
for (int i = 0; i < CHUNK_SIZE; i += 16) {
|
||||
// load input data
|
||||
Unroll<16>{}([&](auto j) {
|
||||
__m512 v = (i + j < size) ? _mm512_loadu_ps(input + (i + j) * ld_src) : _mm512_setzero_ps();
|
||||
__m512 v;
|
||||
if (i + j < mb_size) {
|
||||
v = _mm512_maskz_loadu_ps(vmask, input + (i + j) * ld_src);
|
||||
} else {
|
||||
v = _mm512_setzero_ps();
|
||||
}
|
||||
vsum = _mm512_add_ps(vsum, v);
|
||||
va[j] = _mm512_castps_si512(vsum);
|
||||
});
|
||||
// transpose
|
||||
transpose_16x16_32bit(va);
|
||||
// store output data
|
||||
Unroll<16>{}([&](auto j) { _mm512_storeu_si512(out + j * ld_dst + i, va[j]); });
|
||||
Unroll<16>{}([&](auto j) {
|
||||
if (j < hb_size) {
|
||||
_mm512_storeu_si512(out + j * ld_dst + i, va[j]);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
};
|
||||
@@ -633,9 +650,7 @@ void chunk_local_cumsum_kernel_impl(
|
||||
int64_t Hv,
|
||||
int64_t NT) {
|
||||
constexpr int BLOCK_H = 16;
|
||||
// TODO: now we only support qwen3.5 configs (H/Hv == 16/32)
|
||||
TORCH_CHECK(Hv % BLOCK_H == 0);
|
||||
int64_t HB = Hv / BLOCK_H;
|
||||
int64_t HB = div_up(Hv, int64_t(BLOCK_H));
|
||||
|
||||
// parallel on [NT * HB] to increase parallelism
|
||||
at::parallel_for(0, NT * HB, 0, [&](int64_t begin, int64_t end) {
|
||||
@@ -648,10 +663,11 @@ void chunk_local_cumsum_kernel_impl(
|
||||
int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs];
|
||||
int64_t mb_start = chunk_indices[nt * 2 + 1] * CHUNK_SIZE;
|
||||
int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE));
|
||||
int64_t hb_size = std::min(Hv - hb * BLOCK_H, int64_t(BLOCK_H));
|
||||
|
||||
const scalar_t* __restrict__ g_ptr = g + (batch_offset + mb_start) * Hv + hb * BLOCK_H;
|
||||
scalar_t* __restrict__ gsum_ptr = g_ + nt * (Hv * CHUNK_SIZE) + hb * (BLOCK_H * CHUNK_SIZE);
|
||||
cumsum_kernel<scalar_t, CHUNK_SIZE, BLOCK_H>::apply(gsum_ptr, g_ptr, mb_size, Hv, CHUNK_SIZE);
|
||||
cumsum_kernel<scalar_t, CHUNK_SIZE, BLOCK_H>::apply(gsum_ptr, g_ptr, mb_size, hb_size, Hv, CHUNK_SIZE);
|
||||
|
||||
// move to the next index
|
||||
data_index_step(nt, NT, hb, HB);
|
||||
|
||||
+193
-192
@@ -1,12 +1,12 @@
|
||||
import unittest
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.functional import softplus
|
||||
from utils import parametrize, precision
|
||||
from utils import precision
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
|
||||
|
||||
@@ -250,199 +250,200 @@ def torch_gdn_gating(A_log, a, b, dt_bias):
|
||||
), b.sigmoid().unsqueeze(0)
|
||||
|
||||
|
||||
class TestMambaAttention(CustomTestCase):
|
||||
def test_chunk_gated_delta_rule(self):
|
||||
B, T_PER_SEQ, HK, HV, K, V, POOL_SIZE = 1, 128, 16, 32, 128, 128, 17
|
||||
seq_lens = torch.tensor(
|
||||
[T_PER_SEQ - 7, T_PER_SEQ + 11, T_PER_SEQ - 13, T_PER_SEQ + 9],
|
||||
dtype=torch.int32,
|
||||
)
|
||||
cu_seqlens_ = torch.cat(
|
||||
[
|
||||
torch.zeros(1, dtype=torch.int32),
|
||||
seq_lens.cumsum(dim=0, dtype=torch.int32),
|
||||
]
|
||||
)
|
||||
T = cu_seqlens_[-1].item()
|
||||
cache_indices = torch.tensor([3, 11, 15, 7], dtype=torch.int32)
|
||||
state_slots = cache_indices
|
||||
query_ = torch.randn((B, T, HK, K), dtype=torch.bfloat16)
|
||||
key_ = torch.randn((B, T, HK, K), dtype=torch.bfloat16)
|
||||
value_ = torch.randn((B, T, HV, V), dtype=torch.bfloat16)
|
||||
g_ = F.logsigmoid(torch.randn((B, T, HV), dtype=torch.float32))
|
||||
beta_ = torch.sigmoid(torch.randn((B, T, HV), dtype=torch.bfloat16))
|
||||
initial_state_ = torch.randn((POOL_SIZE, HV, V, K), dtype=torch.float32) * 0.1
|
||||
|
||||
# skip `use_qk_l2norm_in_kernel=False` case since it's not numerically stable in bfloat16
|
||||
for use_qk_l2norm_in_kernel in [True]:
|
||||
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_[state_slots],
|
||||
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().transpose(-1, -2).contiguous()
|
||||
initial_state_before = initial_state.clone()
|
||||
|
||||
core_attn_out, returned_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,
|
||||
initial_state_indices=cache_indices,
|
||||
)
|
||||
)
|
||||
last_recurrent_state = (
|
||||
initial_state[state_slots].transpose(-1, -2).contiguous()
|
||||
)
|
||||
untouched_slots = torch.ones(POOL_SIZE, dtype=torch.bool)
|
||||
untouched_slots[state_slots] = False
|
||||
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
|
||||
)
|
||||
torch.testing.assert_close(returned_state, initial_state)
|
||||
torch.testing.assert_close(
|
||||
initial_state[untouched_slots], initial_state_before[untouched_slots]
|
||||
)
|
||||
|
||||
def test_fused_gdn_gating(self):
|
||||
dims = [6, 32]
|
||||
for dim in dims:
|
||||
for A_log_dtype in [torch.float32, torch.bfloat16]:
|
||||
A_log = torch.rand(dim, dtype=A_log_dtype)
|
||||
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)
|
||||
|
||||
@parametrize(
|
||||
batch_size=[1, 4],
|
||||
num_value_heads=[32],
|
||||
head_k_dim=[128],
|
||||
head_v_dim=[128],
|
||||
num_heads=[16],
|
||||
seq_len=[1],
|
||||
attn_tp_size=[1],
|
||||
@pytest.mark.parametrize(
|
||||
("B", "T_PER_SEQ", "HK", "HV", "K", "V", "POOL_SIZE"),
|
||||
[
|
||||
(1, 128, 3, 6, 128, 128, 17),
|
||||
(1, 128, 16, 32, 128, 128, 17),
|
||||
],
|
||||
)
|
||||
def test_chunk_gated_delta_rule(B, T_PER_SEQ, HK, HV, K, V, POOL_SIZE):
|
||||
seq_lens = torch.tensor(
|
||||
[T_PER_SEQ - 7, T_PER_SEQ + 11, T_PER_SEQ - 13, T_PER_SEQ + 9],
|
||||
dtype=torch.int32,
|
||||
)
|
||||
def test_fused_sigmoid_gating_delta_rule_update(
|
||||
self,
|
||||
batch_size,
|
||||
num_value_heads,
|
||||
head_k_dim,
|
||||
head_v_dim,
|
||||
num_heads,
|
||||
seq_len,
|
||||
attn_tp_size,
|
||||
):
|
||||
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
|
||||
cu_seqlens_ = torch.cat(
|
||||
[
|
||||
torch.zeros(1, dtype=torch.int32),
|
||||
seq_lens.cumsum(dim=0, dtype=torch.int32),
|
||||
]
|
||||
)
|
||||
T = cu_seqlens_[-1].item()
|
||||
cache_indices = torch.tensor([3, 11, 15, 7], dtype=torch.int32)
|
||||
state_slots = cache_indices
|
||||
query_ = torch.randn((B, T, HK, K), dtype=torch.bfloat16)
|
||||
key_ = torch.randn((B, T, HK, K), dtype=torch.bfloat16)
|
||||
value_ = torch.randn((B, T, HV, V), dtype=torch.bfloat16)
|
||||
g_ = F.logsigmoid(torch.randn((B, T, HV), dtype=torch.float32))
|
||||
beta_ = torch.sigmoid(torch.randn((B, T, HV), dtype=torch.bfloat16))
|
||||
initial_state_ = torch.randn((POOL_SIZE, HV, V, K), dtype=torch.float32) * 0.1
|
||||
|
||||
# skip `use_qk_l2norm_in_kernel=False` case since it's not numerically stable in bfloat16
|
||||
for use_qk_l2norm_in_kernel in [True]:
|
||||
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_[state_slots],
|
||||
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
|
||||
)
|
||||
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_.clone()
|
||||
key = key_.clone()
|
||||
value = value_.clone()
|
||||
g = g_.clone()
|
||||
beta = beta_.clone()
|
||||
cu_seqlens = cu_seqlens_.clone()
|
||||
initial_state = initial_state_.clone().transpose(-1, -2).contiguous()
|
||||
initial_state_before = initial_state.clone()
|
||||
|
||||
core_attn_out, returned_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,
|
||||
initial_state_indices=cache_indices,
|
||||
)
|
||||
query = query.view(1, batch_size, num_heads, head_k_dim)
|
||||
key = key.view(1, batch_size, num_heads, head_k_dim)
|
||||
value = value.view(1, batch_size, 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_kv = torch.rand(
|
||||
513, num_value_heads, head_k_dim, head_v_dim, dtype=torch.float32
|
||||
last_recurrent_state = initial_state[state_slots].transpose(-1, -2).contiguous()
|
||||
untouched_slots = torch.ones(POOL_SIZE, dtype=torch.bool)
|
||||
untouched_slots[state_slots] = False
|
||||
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
|
||||
)
|
||||
torch.testing.assert_close(returned_state, initial_state)
|
||||
torch.testing.assert_close(
|
||||
initial_state[untouched_slots], initial_state_before[untouched_slots]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dim", [6, 32])
|
||||
@pytest.mark.parametrize(
|
||||
"A_log_dtype",
|
||||
[torch.float32, torch.bfloat16],
|
||||
ids=["float32", "bfloat16"],
|
||||
)
|
||||
def test_fused_gdn_gating(dim, A_log_dtype):
|
||||
A_log = torch.rand(dim, dtype=A_log_dtype)
|
||||
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)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
(
|
||||
"batch_size",
|
||||
"num_value_heads",
|
||||
"head_k_dim",
|
||||
"head_v_dim",
|
||||
"num_heads",
|
||||
"seq_len",
|
||||
"attn_tp_size",
|
||||
),
|
||||
[
|
||||
(1, 32, 128, 128, 16, 1, 1),
|
||||
(4, 32, 128, 128, 16, 1, 1),
|
||||
],
|
||||
)
|
||||
def test_fused_sigmoid_gating_delta_rule_update(
|
||||
batch_size,
|
||||
num_value_heads,
|
||||
head_k_dim,
|
||||
head_v_dim,
|
||||
num_heads,
|
||||
seq_len,
|
||||
attn_tp_size,
|
||||
):
|
||||
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, batch_size, num_heads, head_k_dim)
|
||||
key = key.view(1, batch_size, num_heads, head_k_dim)
|
||||
value = value.view(1, batch_size, 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_kv = 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.arange(batch_size + 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)
|
||||
for A_log_dtype in [torch.float32, torch.bfloat16]:
|
||||
A_log = A_log.to(A_log_dtype)
|
||||
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_kv[cache_indices].transpose(-1, -2).contiguous(),
|
||||
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_kv,
|
||||
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_kv[cache_indices].transpose(-1, -2).contiguous()
|
||||
)
|
||||
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
|
||||
)
|
||||
cache_indices = torch.randint(0, 513, (batch_size,), dtype=torch.int32)
|
||||
query_start_loc = torch.arange(batch_size + 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)
|
||||
for A_log_dtype in [torch.float32, torch.bfloat16]:
|
||||
A_log = A_log.to(A_log_dtype)
|
||||
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_kv[cache_indices]
|
||||
.transpose(-1, -2)
|
||||
.contiguous(),
|
||||
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_kv,
|
||||
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_kv[cache_indices].transpose(-1, -2).contiguous()
|
||||
)
|
||||
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()
|
||||
sys.exit(pytest.main([__file__]))
|
||||
|
||||
Reference in New Issue
Block a user