[CPU] update fla.cpp to support when num_head_v is not multiples of 16 (#30604)

sgl-kernel/csrc/cpu/mamba/fla.cpp

@@ -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
+  static inline void apply(
-  apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int size, int ld_src, int ld_dst) {
+      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)
+  int64_t HB = div_up(Hv, int64_t(BLOCK_H));
-  TORCH_CHECK(Hv % BLOCK_H == 0);
-  int64_t HB = Hv / 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);

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):
+@pytest.mark.parametrize(
-    def test_chunk_gated_delta_rule(self):
+    ("B", "T_PER_SEQ", "HK", "HV", "K", "V", "POOL_SIZE"),
-        B, T_PER_SEQ, HK, HV, K, V, POOL_SIZE = 1, 128, 16, 32, 128, 128, 17
+    [
-        seq_lens = torch.tensor(
+        (1, 128, 3, 6, 128, 128, 17),
-            [T_PER_SEQ - 7, T_PER_SEQ + 11, T_PER_SEQ - 13, T_PER_SEQ + 9],
+        (1, 128, 16, 32, 128, 128, 17),
-            dtype=torch.int32,
+    ],
-        )
+)
-        cu_seqlens_ = torch.cat(
+def test_chunk_gated_delta_rule(B, T_PER_SEQ, HK, HV, K, V, POOL_SIZE):
-            [
+    seq_lens = torch.tensor(
-                torch.zeros(1, dtype=torch.int32),
+        [T_PER_SEQ - 7, T_PER_SEQ + 11, T_PER_SEQ - 13, T_PER_SEQ + 9],
-                seq_lens.cumsum(dim=0, dtype=torch.int32),
+        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],
     )
-    def test_fused_sigmoid_gating_delta_rule_update(
+    cu_seqlens_ = torch.cat(
-        self,
+        [
-        batch_size,
+            torch.zeros(1, dtype=torch.int32),
-        num_value_heads,
+            seq_lens.cumsum(dim=0, dtype=torch.int32),
-        head_k_dim,
+        ]
-        head_v_dim,
+    )
-        num_heads,
+    T = cu_seqlens_[-1].item()
-        seq_len,
+    cache_indices = torch.tensor([3, 11, 15, 7], dtype=torch.int32)
-        attn_tp_size,
+    state_slots = cache_indices
-    ):
+    query_ = torch.randn((B, T, HK, K), dtype=torch.bfloat16)
-        key_dim = head_k_dim * num_heads
+    key_ = torch.randn((B, T, HK, K), dtype=torch.bfloat16)
-        value_dim = head_v_dim * num_value_heads
+    value_ = torch.randn((B, T, HV, V), dtype=torch.bfloat16)
-        mixed_qkv_dim = (key_dim * 2 + value_dim) // attn_tp_size
+    g_ = F.logsigmoid(torch.randn((B, T, HV), dtype=torch.float32))
-        mixed_qkv = torch.rand(
+    beta_ = torch.sigmoid(torch.randn((B, T, HV), dtype=torch.bfloat16))
-            seq_len * batch_size, mixed_qkv_dim, 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,
+        query = query_.clone()
-            [
+        key = key_.clone()
-                key_dim // attn_tp_size,
+        value = value_.clone()
-                key_dim // attn_tp_size,
+        g = g_.clone()
-                value_dim // attn_tp_size,
+        beta = beta_.clone()
-            ],
+        cu_seqlens = cu_seqlens_.clone()
-            dim=-1,
+        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)
+        last_recurrent_state = initial_state[state_slots].transpose(-1, -2).contiguous()
-        key = key.view(1, batch_size, num_heads, head_k_dim)
+        untouched_slots = torch.ones(POOL_SIZE, dtype=torch.bool)
-        value = value.view(1, batch_size, num_value_heads, head_v_dim)
+        untouched_slots[state_slots] = False
-        A_log = torch.rand(num_value_heads, dtype=torch.float32)
+        atol = rtol = precision[core_attn_out.dtype]
-        a = torch.rand(batch_size, num_value_heads, dtype=torch.bfloat16)
+        torch.testing.assert_close(
-        b = torch.rand(batch_size, num_value_heads, dtype=torch.bfloat16)
+            core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol
-        dt_bias = torch.rand(num_value_heads, dtype=torch.bfloat16)
+        )
-        ssm_states_kv = torch.rand(
+        torch.testing.assert_close(
-            513, num_value_heads, head_k_dim, head_v_dim, dtype=torch.float32
+            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__]))