[Bug report] BatchPrefillWithPagedKVCachePyTorchWrapper failed to dispatch group_size 3

[merrymercy]

Error traceback

  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/sglang/python/sglang/srt/layers/radix_attention.py", line 128, in forward
    return self.extend_forward(q, k, v, input_metadata)
  File "/root/sglang/python/sglang/srt/layers/radix_attention.py", line 104, in prefill_forward_flashinfer
    o = input_metadata.prefill_wrapper.forward(
  File "/usr/local/lib/python3.10/dist-packages/flashinfer/prefill.py", line 498, in forward
    return self._wrapper.forward(
RuntimeError: BatchPrefillWithPagedKVCachePyTorchWrapper::Forward(at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, bool, unsigned int, bool, float, float, float, bool)::<lambda()>::<lambda()>::<lambda()>::<lambda()>::<lambda()> failed to dispatch group_size 3

Shape information

num_heads 24
num_kv_heads 8
head_dim 128
q.shape torch.Size([6, 3072])

[merrymercy]

similar to this one #254

[yzh119]

Yes I'm coming to fix these series of issues :)

[yzh119]

With #301 merged, now flashinfer's prefill kernels support any group sizes, and decode kernels support group size 1-8.