[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,9 +250,14 @@ 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
+    [
+        (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,
@@ -295,8 +300,7 @@ class TestMambaAttention(CustomTestCase):
         initial_state = initial_state_.clone().transpose(-1, -2).contiguous()
         initial_state_before = initial_state.clone()
 
-            core_attn_out, returned_state = (
+        core_attn_out, returned_state = torch.ops.sgl_kernel.chunk_gated_delta_rule_cpu(
-                torch.ops.sgl_kernel.chunk_gated_delta_rule_cpu(
             query=query,
             key=key,
             value=value,
@@ -309,10 +313,7 @@ class TestMambaAttention(CustomTestCase):
             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()
-            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]
@@ -327,35 +328,43 @@ class TestMambaAttention(CustomTestCase):
             initial_state[untouched_slots], initial_state_before[untouched_slots]
         )
 
-    def test_fused_gdn_gating(self):
+
-        dims = [6, 32]
+@pytest.mark.parametrize("dim", [6, 32])
-        for dim in dims:
+@pytest.mark.parametrize(
-            for A_log_dtype in [torch.float32, torch.bfloat16]:
+    "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(
+    g_sgl, beta_sgl = torch.ops.sgl_kernel.fused_gdn_gating_cpu(A_log, a, b, dt_bias)
-                    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],
+@pytest.mark.parametrize(
-        num_value_heads=[32],
+    (
-        head_k_dim=[128],
+        "batch_size",
-        head_v_dim=[128],
+        "num_value_heads",
-        num_heads=[16],
+        "head_k_dim",
-        seq_len=[1],
+        "head_v_dim",
-        attn_tp_size=[1],
+        "num_heads",
-    )
+        "seq_len",
-    def test_fused_sigmoid_gating_delta_rule_update(
+        "attn_tp_size",
-        self,
+    ),
+    [
+        (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,
@@ -363,13 +372,11 @@ class TestMambaAttention(CustomTestCase):
     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(
+    mixed_qkv = torch.rand(seq_len * batch_size, mixed_qkv_dim, dtype=torch.bfloat16)
-            seq_len * batch_size, mixed_qkv_dim, dtype=torch.bfloat16
-        )
     query, key, value = torch.split(
         mixed_qkv,
         [
@@ -399,8 +406,7 @@ class TestMambaAttention(CustomTestCase):
         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 = (
+        core_attn_out_ref, last_recurrent_state_ref = sigmoid_gating_delta_rule_update(
-                sigmoid_gating_delta_rule_update(
             query_ref.transpose(0, 1),
             key_ref.transpose(0, 1),
             value.transpose(0, 1),
@@ -408,15 +414,11 @@ class TestMambaAttention(CustomTestCase):
             a,
             dt_bias,
             b,
-                    initial_state=ssm_states_kv[cache_indices]
+            initial_state=ssm_states_kv[cache_indices].transpose(-1, -2).contiguous(),
-                    .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(
-            core_attn_out = (
-                torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu(
             A_log=A_log,
             dt_bias=dt_bias,
             q=query,
@@ -431,7 +433,6 @@ class TestMambaAttention(CustomTestCase):
             softplus_beta=1.0,
             softplus_threshold=20.0,
         )
-            )
         last_recurrent_state = (
             ssm_states_kv[cache_indices].transpose(-1, -2).contiguous()
         )
@@ -445,4 +446,4 @@ class TestMambaAttention(CustomTestCase):
 
 
 if __name__ == "__main__":
-    unittest.main()
+    sys.exit(pytest.main([__file__]))