Floating point exception when runing the example py

[Toconquer1]

When I run the following program, if kv_len exceeds 256, the program outputs a Floating point exception at single_decode_with_kv_cache, while other functions do not encounter errors. What could be the issue? I'm using cuda12.1+torch2.2.

import torch
import flashinfer

kv_len = 2048 #512-2048 Floating point exception
num_kv_heads = 32
head_dim = 128

k = torch.randn(kv_len, num_kv_heads, head_dim).half().to(0) 
v = torch.randn(kv_len, num_kv_heads, head_dim).half().to(0) 

# decode attention

num_qo_heads = 32
q = torch.randn(num_qo_heads, head_dim).half().to(0)

# error occur
o = flashinfer.single_decode_with_kv_cache(q, k, v) # decode attention without RoPE on-the-fly
o_rope_on_the_fly = flashinfer.single_decode_with_kv_cache(q, k, v, pos_encoding_mode="ROPE_LLAMA") # decode with LLaMA style RoPE on-the-fly
print(o.shape)
print(o_rope_on_the_fly.shape)
# append attention
append_qo_len = 128
q = torch.randn(append_qo_len, num_qo_heads, head_dim).half().to(0) # append attention, the last 128 tokens in the KV-Cache are the new tokens
o = flashinfer.single_prefill_with_kv_cache(q, k, v, causal=True) # append attention without RoPE on-the-fly, apply causal mask
o_rope_on_the_fly = flashinfer.single_prefill_with_kv_cache(q, k, v, causal=True, pos_encoding_mode="ROPE_LLAMA") # append attention with LLaMA style RoPE on-the-fly, apply causal mask
print(o.shape)
print(o_rope_on_the_fly.shape)
# prefill attention
qo_len = 2048
q = torch.randn(qo_len, num_qo_heads, head_dim).half().to(0) # prefill attention
o = flashinfer.single_prefill_with_kv_cache(q, k, v, causal=False) # prefill attention without RoPE on-the-fly, do not apply causal mask
print(o.shape)
print(o_rope_on_the_fly.shape)
`

[yzh119]

Hi @Toconquer1 , I cannot reproduce the error in my environment, what's your GPU architecture?

[Toconquer1]

NVIDIA GeForce RTX 3050 4GB Laptop GPU

[Toconquer1]

the architecture is sm86 which should be supported

[yzh119]

Thank you for the feedback, sm86 was removed in a PR https://github.com/flashinfer-ai/flashinfer/pull/172/files to reduce binary size, I'll try running the latest pip wheel on my 3070 desktop tonight, if it doesn't work I'll add it back.

[yzh119]

Would you mind trying the latest release https://github.com/flashinfer-ai/flashinfer/releases/tag/v0.0.4 ? I cannot reproduce the error on 3070.