diff --git a/sgl-kernel/csrc/cpu/mamba/fla.cpp b/sgl-kernel/csrc/cpu/mamba/fla.cpp index de6064179..2abfa25fa 100644 --- a/sgl-kernel/csrc/cpu/mamba/fla.cpp +++ b/sgl-kernel/csrc/cpu/mamba/fla.cpp @@ -135,8 +135,13 @@ struct l2norm_kernel { template 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 struct cumsum_kernel { - 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 { 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::apply(gsum_ptr, g_ptr, mb_size, Hv, CHUNK_SIZE); + cumsum_kernel::apply(gsum_ptr, g_ptr, mb_size, hb_size, Hv, CHUNK_SIZE); // move to the next index data_index_step(nt, NT, hb, HB); diff --git a/test/registered/cpu/test_mamba.py b/test/registered/cpu/test_mamba.py index f99009016..6a6d6eac8 100644 --- a/test/registered/cpu/test_mamba.py +++ b/test/registered/cpu/test_mamba.py @@ -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__]))