[Bug]: Cache malformation in hybrid models with SSM cache dtype float32 and block allocation wrap around

[roikoren755]

Your current environment

The output of python collect_env.py

/home/scratch.roik_gpu/vllm-roik/.venv/lib/python3.12/site-packages/torch/cuda/__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]
Collecting environment information...
uv is set
==============================
        System Info
==============================
OS                           : Ubuntu 24.04.1 LTS (x86_64)
GCC version                  : (Ubuntu 13.3.0-6ubuntu2~24.04) 13.3.0
Clang version                : Could not collect
CMake version                : version 3.31.4
Libc version                 : glibc-2.39

==============================
       PyTorch Info
==============================
PyTorch version              : 2.8.0+cu128
Is debug build               : False
CUDA used to build PyTorch   : 12.8
ROCM used to build PyTorch   : N/A

==============================
      Python Environment
==============================
Python version               : 3.12.3 (main, Jan 17 2025, 18:03:48) [GCC 13.3.0] (64-bit runtime)
Python platform              : Linux-5.15.0-157-generic-x86_64-with-glibc2.39

==============================
       CUDA / GPU Info
==============================
Is CUDA available            : True
CUDA runtime version         : 12.8.61
CUDA_MODULE_LOADING set to   : LAZY
GPU models and configuration : GPU 0: NVIDIA H100 PCIe
Nvidia driver version        : 580.95.05
cuDNN version                : Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.7.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.7.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.7.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.7.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.7.1
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.7.1
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.7.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.7.1
HIP runtime version          : N/A
MIOpen runtime version       : N/A
Is XNNPACK available         : True

==============================
          CPU Info
==============================
Architecture:                            x86_64
CPU op-mode(s):                          32-bit, 64-bit
Address sizes:                           52 bits physical, 57 bits virtual
Byte Order:                              Little Endian
CPU(s):                                  384
On-line CPU(s) list:                     0-383
Vendor ID:                               AuthenticAMD
Model name:                              AMD EPYC 9654 96-Core Processor
CPU family:                              25
Model:                                   17
Thread(s) per core:                      2
Core(s) per socket:                      96
Socket(s):                               2
Stepping:                                1
Frequency boost:                         enabled
CPU(s) scaling MHz:                      42%
CPU max MHz:                             3709.3569
CPU min MHz:                             1500.0000
BogoMIPS:                                4800.01
Flags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp ibrs_enhanced vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid overflow_recov succor smca fsrm flush_l1d
Virtualization:                          AMD-V
L1d cache:                               6 MiB (192 instances)
L1i cache:                               6 MiB (192 instances)
L2 cache:                                192 MiB (192 instances)
L3 cache:                                768 MiB (24 instances)
NUMA node(s):                            2
NUMA node0 CPU(s):                       0-95,192-287
NUMA node1 CPU(s):                       96-191,288-383
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Vulnerable
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:                Vulnerable; IBPB: disabled; STIBP: disabled; PBRSB-eIBRS: Not affected; BHI: Not affected
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Mitigation; Clear CPU buffers
Vulnerability Tsx async abort:           Not affected

==============================
Versions of relevant libraries
==============================
[pip3] efficientnet-pytorch==0.7.1
[pip3] flashinfer-python==0.4.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.8.4.1
[pip3] nvidia-cuda-cupti-cu12==12.8.90
[pip3] nvidia-cuda-nvrtc-cu12==12.8.93
[pip3] nvidia-cuda-runtime-cu12==12.8.90
[pip3] nvidia-cudnn-cu12==9.10.2.21
[pip3] nvidia-cudnn-frontend==1.15.0
[pip3] nvidia-cufft-cu12==11.3.3.83
[pip3] nvidia-cufile-cu12==1.13.1.3
[pip3] nvidia-curand-cu12==10.3.9.90
[pip3] nvidia-cusolver-cu12==11.7.3.90
[pip3] nvidia-cusparse-cu12==12.5.8.93
[pip3] nvidia-cusparselt-cu12==0.7.1
[pip3] nvidia-cutlass-dsl==4.2.1
[pip3] nvidia-ml-py==13.580.82
[pip3] nvidia-nccl-cu12==2.27.3
[pip3] nvidia-nvjitlink-cu12==12.8.93
[pip3] nvidia-nvtx-cu12==12.8.90
[pip3] open-clip-torch==2.32.0
[pip3] pynvml==13.0.1
[pip3] pytorch-lightning==2.5.2
[pip3] pyzmq==27.1.0
[pip3] segmentation-models-pytorch==0.4.0
[pip3] sentence-transformers==3.2.1
[pip3] terratorch==1.0.2
[pip3] torch==2.8.0+cu128
[pip3] torchaudio==2.8.0+cu128
[pip3] torchgeo==0.7.0
[pip3] torchmetrics==1.7.4
[pip3] torchvision==0.23.0+cu128
[pip3] transformers==4.56.2
[pip3] transformers-stream-generator==0.0.5
[pip3] triton==3.4.0
[pip3] tritonclient==2.51.0
[pip3] vector-quantize-pytorch==1.21.2
[conda] Could not collect

==============================
         vLLM Info
==============================
ROCM Version                 : Could not collect
vLLM Version                 : 0.1.dev10602+g9fce7bee7 (git sha: 9fce7bee7)
vLLM Build Flags:
  CUDA Archs: 7.5 8.0 8.6 9.0 10.0 12.0+PTX; ROCm: Disabled
GPU Topology:
  	GPU0	CPU Affinity	NUMA Affinity	GPU NUMA ID
GPU0	 X 	0-95,192-287	0		N/A

Legend:

  X    = Self
  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
  PIX  = Connection traversing at most a single PCIe bridge
  NV#  = Connection traversing a bonded set of # NVLinks

==============================
     Environment Variables
==============================
NVIDIA_VISIBLE_DEVICES=1
CUBLAS_VERSION=12.8.3.14
NVIDIA_REQUIRE_CUDA=cuda>=9.0
CUDA_CACHE_DISABLE=1
TORCH_CUDA_ARCH_LIST=7.5 8.0 8.6 9.0 10.0 12.0+PTX
NCCL_VERSION=2.25.1
NVIDIA_DRIVER_CAPABILITIES=compute,utility,video
TORCH_NCCL_USE_COMM_NONBLOCKING=0
NVIDIA_PRODUCT_NAME=PyTorch
CUDA_VERSION=12.8.0.038
PYTORCH_VERSION=2.7.0a0+ecf3bae
PYTORCH_BUILD_NUMBER=0
CUDNN_FRONTEND_VERSION=1.10.0
CUDNN_VERSION=9.7.1.26
PYTORCH_HOME=/opt/pytorch/pytorch
LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/torch/lib:/usr/local/lib/python3.12/dist-packages/torch_tensorrt/lib:/usr/local/cuda/compat/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64
NVIDIA_BUILD_ID=143088496
CUDA_DRIVER_VERSION=570.86.10
PYTORCH_BUILD_VERSION=2.7.0a0+ecf3bae
CUDA_HOME=/usr/local/cuda
CUDA_HOME=/usr/local/cuda
CUDA_MODULE_LOADING=LAZY
NVIDIA_REQUIRE_JETPACK_HOST_MOUNTS=
NVIDIA_PYTORCH_VERSION=25.02
TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=1
PYTORCH_NVML_BASED_CUDA_CHECK=1
TORCHINDUCTOR_COMPILE_THREADS=1

🐛 Describe the bug

Hey,

We're running a hybrid model, and are running into an issue when turning prefix caching on.
After a few requests where processed, the model starts returning 0s. After some debugging, we've found that the issue starts happening in the first full attention layer in the model, which comes after some Mamba2 layers, where at some point the attention output started including NaNs, which quickly spread and turned the entire hidden states into NaNs.

Looking deeper, we saw that the block tables that were input into the attention computation were wrapping around the cache buffer when the issue started occurring. The prompt in the script below lead to an increase of 5 in the attention layer's single allocated page (5, 10, 15, ...), until it reached the end of the buffer, and then it started again at (2, 7, 12, ...). The issue started the moment the block table included page 2, and happened consistently after that.

Further debugging revealed that this happens only when setting mamba_ssm_cache_dtype="float32" in the model initialization.

We've found that the issue only occurs in main after #25103, which doesn't concern the KV cache directly, but does change the order of the keys of the KV cache specs which are then used to initialize the different KV caches, including the Mamba ones. We've then confirmed that this issue also happens all the way back to the original PR (#25752) that introduced APC for Mamba2 caches, when the KV cache spec keys are sorted in the same way, and doesn't happen before.

We didn't manage to get the issue reproducing inside any of the currently running hybrid model tests at this time, unfortunately.

We are looking into the areas of the code which write to the cache in the Mamba2 mixer, as that seems the most logical place for things to go wrong in, and we're actively working to solve this, as this impacts our work, and we'd very much like to utilize prefix caching. Do you have any suggestions as to what to look out for?

from vllm import LLM, SamplingParams

prompts = ["Solve the following math problem step by step. The last line of your response should be of the form Answer: $Answer (without quotes) where $Answer is the answer to the problem.\n\nConsider the paths of length $16$ that follow the lines from the lower left corner to the upper right corner on an $8\\times 8$ grid. Find the number of such paths that change direction exactly four times, as in the examples shown below.\n\nRemember to put your answer on its own line after \"Answer:\"."]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(
    model="nvidia/NVIDIA-Nemotron-Nano-9B-v2",
    trust_remote_code=True,
    enable_prefix_caching=True,
    mamba_ssm_cache_dtype="float32",
    num_gpu_blocks_override=10,
)

first_output = llm.generate(prompts, sampling_params)
second_output = llm.generate(prompts, sampling_params)

# Print the outputs.
print(f"Prompt: {prompts[0]!r}")
for output in first_output + second_output:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Generated text: {generated_text!r}")
    print(f"Generated token IDs: {output.outputs[0].token_ids!r}")

Prompt: 'Solve the following math problem step by step. The last line of your response should be of the form Answer: $Answer (without quotes) where $Answer is the answer to the problem.\n\nConsider the paths of length $16$ that follow the lines from the lower left corner to the upper right corner on an $8\\times 8$ grid. Find the number of such paths that change direction exactly four times, as in the examples shown below.\n\nRemember to put your answer on its own line after "Answer:".'
Generated text: '</think>\nThe problem requires finding the number of paths of length 16'
Generated token IDs: [1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[heheda12345]

CC @tdoublep

[tdoublep]

@roikoren755 Thank you very much for reporting the issue, and your work to debug what may be causing this. We will take a look into this asap.

cc @s3woz @bohnstingl

[roikoren755]

We've dug a little deeper, and found something that might be interesting. The model in the reproduction script has an irregular pattern and ratio between Mamba2 layers and full attention layers, at 27 to 4, and no actual pattern to speak of.

Before the MLA separation PR (mentioned above), the first KV cache spec in the iteration was that of a full attention layer, and the 4 KVCacheTensor instances' shared_by lists all started with such a layer, which was followed by 6 or 7 Mamba2 layers. After that PR, 3 of those tensors' shared_by lists have 7 Mamba2 layers followed by 1 full attention layer, and the 4th has 6 Mamba2 layers followed by 1 full attention layer.

I couldn't find information in the docs, regarding memory layout and what happens when padding is needed there, but could this lead to the issue we're seeing?

[roikoren755]

I've opened a PR that fixes this issue (#27346). A review would be much appreciated 🙏

[tdoublep]

Thanks a lot for creating this repro - I believe this is actually revealing quite an important bug. I think I hit this before when testing very large model (like Minimax-Text-01) with fp32 state but could not come up with simple repro. I've been debugging this today and share my findings below.

I'm running a slight variant of the script shared above:

from vllm import LLM, SamplingParams
import os

os.environ['VLLM_ATTENTION_BACKEND'] = 'FLASH_ATTN'
#os.environ['VLLM_ATTENTION_BACKEND'] = 'TRITON_ATTN'
#os.environ['VLLM_ATTENTION_BACKEND'] = 'FLASHINFER'

prompts = ["Solve the following math problem step by step. The last line of your response should be of the form Answer: $Answer (without quotes) where $Answer is the answer to the problem.\n\nConsider the paths of >
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(
    model="nvidia/NVIDIA-Nemotron-Nano-9B-v2",
    trust_remote_code=True,
    num_gpu_blocks_override=10,
    compilation_config={'cudagraph_capture_sizes': [1]},
    mamba_ssm_cache_dtype="float32",
    max_num_seqs=1,
    enable_prefix_caching=True,
)

outputs = []
for _ in range(10):
    outputs.append(llm.generate(prompts, sampling_params)[0])

# Print the outputs.
print(f"Prompt: {prompts[0]!r}")
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Generated text: {generated_text!r}")
    print(f"Generated token IDs: {output.outputs[0].token_ids!r}")

When I run it using FLASH_ATTN I see:

Generated text: '</think>\nThe problem requires finding the number of paths of length 16'
Generated token IDs: [1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

As discussed in #27346 it seems like the root cause here is that using fp32 mamba state leads to bits in the KV cache that result in nans when interpreted as bf16 by the attention layers. The strange thing is that the attention layers should never be reading the data from the mamba state, something is wrong.

To see if this is a FlashAttention issue, I tried to run with TRITON_ATTN which works just fine:

Generated text: '</think>\nThe problem requires finding the number of paths of length 16'
Generated token IDs: [1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054]
Generated text: ' \n\nExample 1: \n\nExample 2:\n\n\n<think>\n'
Generated token IDs: [1032, 1267, 20396, 1032, 1049, 1058, 1032, 1267, 20396, 1032, 1050, 2100, 1010, 49250, 2077, 1561]
Generated text: '</think>\nThe number of paths of length 16 from the lower left'
Generated token IDs: [1885, 74045, 1561, 1784, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054, 1562, 1278, 4953, 3979]
Generated text: '</think>\nThe problem involves finding the number of paths of length 16'
Generated token IDs: [1885, 74045, 1561, 1784, 4127, 19263, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054]
Generated text: ' output: To solve the problem of finding the number of paths of length 1'
Generated token IDs: [4848, 1058, 3870, 15047, 1278, 4127, 1307, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049]
Generated text: ' \n</think>\nThe problem requires finding the number of paths of length '
Generated token IDs: [1032, 1010, 1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032]
Generated text: '</think>\nThe problem requires finding the number of paths of length 16'
Generated token IDs: [1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054]
Generated text: ' \n</think>\nThe paths of length 16 from the lower left'
Generated token IDs: [1032, 1010, 1885, 74045, 1561, 1784, 22344, 1307, 5592, 1032, 1049, 1054, 1562, 1278, 4953, 3979]
Generated text: " \n\n<think>\nOkay, let's see. I need to find the"
Generated token IDs: [1032, 1267, 49250, 2077, 1561, 44053, 1044, 2878, 1681, 3219, 1046, 1362, 2534, 1317, 3081, 1278]
Generated text: '<think>\nOkay, so I need to find the number of paths on an'
Generated token IDs: [1060, 74045, 1561, 44053, 1044, 1878, 1362, 2534, 1317, 3081, 1278, 2782, 1307, 22344, 1408, 1420]

I then tried it using FLASHINFER and see even stranger behaviour still:

Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

[tdoublep]

After debugging the FLASHINFER behaviour, it seems like since we merged #24486, we decouple the kernel block size from the KV cache manager block size. However, it looks like we don't actually pass the kernel block size to the Flashinfer attention backend in any way (cc @heheda12345). This leads to out-of-bounds reads into the KV cache which is why we are getting nans propagating everywhere when the mamba state is in fp32.

For the FlashAttention behaviour, until now we have assumed that FlashAttention can handle arbitrary block size (up to multiple of 16). I actually hypothesize that this may not be the case. I suspect that the kernels are getting mapped under the hood to some pre-compiled block sizes (e.g. powers of 2) that may be leading to out-of-bounds reads into the KV cache tensor. To test this, I tried to restrict the kernel block size for FlashAttention to [16, 32, 64] like we do for FlashInfer. If I take care to also pass this to the FlashAttentionBackend, then the nan behaviour goes away. Maybe @LucasWilkinson has some insight here.

TLDR: Why are we not passing the kernel block size to actual attention backends? I think this is the root cause of this issue.

[heheda12345]

We may need some refactor of this part of code. Now we first initialize attention backends without knowing the kernel block size and then determine the kernel block_size based on the initialized attention backends, but we don't update the attention backend's block_size after the kernel_block_size is decided. Need to find a way to do it in a correct way.

[tdoublep]

I have a draft PR to solve this issue here: #27753

I don't really understand why it is necessary to restrict FlashAttention's kernel block size though. According to this comment, it should support anything. But the behaviour certainly suggests it is reading out-of-bounds.

[tdoublep]

I've been digging more into the FlashAttention issue. It does indeed seem to be something orthogonal to the FlashInfer issue (which was simply that we aren't passing the kernel block size).

For FlashAttention, it looks like if NaN's exist in the KV tensor in the attention view (which can happen if mamba state is in fp32), then these can cause FA3 to produce NaN as output, even if the NaN's reside in locations that theoretically shouldn't affect the computation. I was able to reproduce this issue with upstream FA and opened an issue here: Dao-AILab/flash-attention#1974

Using the vLLM flash attention, I am using the following reproducer to investigate the issue:

import torch

from vllm.attention.utils.fa_utils import (
    is_flash_attn_varlen_func_available,
)

if is_flash_attn_varlen_func_available():
    from vllm.attention.utils.fa_utils import (
        flash_attn_varlen_func,
    )

BLOCK_SIZE = 128
N_BLOCKS = 1
N_HEADS = 40
N_KV_HEADS = 8
HEAD_DIM = 128

torch.set_default_dtype(torch.bfloat16)
torch.set_default_device('cuda:0')

for N_TOKENS in range(1+BLOCK_SIZE):

    query = torch.rand(N_TOKENS, N_HEADS, HEAD_DIM)
    key_cache = torch.rand(N_BLOCKS, BLOCK_SIZE, N_KV_HEADS, HEAD_DIM)
    value_cache = torch.rand(N_BLOCKS, BLOCK_SIZE, N_KV_HEADS, HEAD_DIM)
    cu_seqlens_q = torch.tensor([0, N_TOKENS], dtype=torch.int32)
    seqused_k = torch.tensor([N_TOKENS], dtype=torch.int32)

    output = torch.empty_like(query)

    block_table = torch.tensor([[0]], dtype=torch.int32)

    # fill key_cache and value_cache with nans *outside* of
    # the range where they should be read
    key_cache[:, N_TOKENS:].fill_(float('nan'))
    value_cache[:, N_TOKENS:].fill_(float('nan'))

    try:
        flash_attn_varlen_func(
            q=query,
            k=key_cache,
            v=value_cache,
            out=output,
            cu_seqlens_q=cu_seqlens_q,
            max_seqlen_q=N_TOKENS,
            seqused_k=seqused_k,
            max_seqlen_k=N_TOKENS,
            softmax_scale=0.08838834764831845,
            causal=True,
            alibi_slopes=None,
            window_size=(-1,-1),
            block_table=block_table,
            softcap=0,
            scheduler_metadata=None,
            fa_version=3,
            q_descale=None,
            k_descale=None,
            v_descale=None,
            num_splits=0,
            s_aux=None,
        )

        assert not torch.isnan(output).any()
        print("n_tokens: %3d, PASS" % (N_TOKENS))
    except:
        print("n_tokens: %3d, FAIL" % (N_TOKENS))

What is interesting is that for block size 16, 32, and 64 I do not see any affect from having NaN's in partial block. But from 128 I see behaviour like this:

n_tokens:   0, PASS
n_tokens:   1, PASS
n_tokens:   2, PASS
n_tokens:   3, PASS
n_tokens:   4, PASS
n_tokens:   5, PASS
n_tokens:   6, PASS
n_tokens:   7, PASS
n_tokens:   8, PASS
n_tokens:   9, PASS
n_tokens:  10, PASS
n_tokens:  11, PASS
n_tokens:  12, PASS
n_tokens:  13, PASS
n_tokens:  14, PASS
n_tokens:  15, PASS
n_tokens:  16, PASS
n_tokens:  17, PASS
n_tokens:  18, PASS
n_tokens:  19, PASS
n_tokens:  20, PASS
n_tokens:  21, PASS
n_tokens:  22, PASS
n_tokens:  23, PASS
n_tokens:  24, PASS
n_tokens:  25, PASS
n_tokens:  26, FAIL
n_tokens:  27, FAIL
n_tokens:  28, FAIL
n_tokens:  29, FAIL
n_tokens:  30, FAIL
n_tokens:  31, FAIL
n_tokens:  32, FAIL
n_tokens:  33, FAIL
n_tokens:  34, FAIL
n_tokens:  35, FAIL
n_tokens:  36, FAIL
n_tokens:  37, FAIL
n_tokens:  38, FAIL
n_tokens:  39, FAIL
n_tokens:  40, FAIL
n_tokens:  41, FAIL
n_tokens:  42, FAIL
n_tokens:  43, FAIL
n_tokens:  44, FAIL
n_tokens:  45, FAIL
n_tokens:  46, FAIL
n_tokens:  47, FAIL
n_tokens:  48, FAIL
n_tokens:  49, FAIL
n_tokens:  50, FAIL
n_tokens:  51, FAIL
n_tokens:  52, FAIL
n_tokens:  53, FAIL
n_tokens:  54, FAIL
n_tokens:  55, FAIL
n_tokens:  56, FAIL
n_tokens:  57, FAIL
n_tokens:  58, FAIL
n_tokens:  59, FAIL
n_tokens:  60, FAIL
n_tokens:  61, FAIL
n_tokens:  62, FAIL
n_tokens:  63, FAIL
n_tokens:  64, FAIL
n_tokens:  65, FAIL
n_tokens:  66, FAIL
n_tokens:  67, FAIL
n_tokens:  68, FAIL
n_tokens:  69, FAIL
n_tokens:  70, FAIL
n_tokens:  71, FAIL
n_tokens:  72, FAIL
n_tokens:  73, FAIL
n_tokens:  74, FAIL
n_tokens:  75, FAIL
n_tokens:  76, FAIL
n_tokens:  77, FAIL
n_tokens:  78, FAIL
n_tokens:  79, FAIL
n_tokens:  80, FAIL
n_tokens:  81, FAIL
n_tokens:  82, FAIL
n_tokens:  83, FAIL
n_tokens:  84, FAIL
n_tokens:  85, FAIL
n_tokens:  86, FAIL
n_tokens:  87, FAIL
n_tokens:  88, FAIL
n_tokens:  89, FAIL
n_tokens:  90, FAIL
n_tokens:  91, FAIL
n_tokens:  92, FAIL
n_tokens:  93, FAIL
n_tokens:  94, FAIL
n_tokens:  95, FAIL
n_tokens:  96, FAIL
n_tokens:  97, FAIL
n_tokens:  98, FAIL
n_tokens:  99, FAIL
n_tokens: 100, FAIL
n_tokens: 101, FAIL
n_tokens: 102, FAIL
n_tokens: 103, FAIL
n_tokens: 104, FAIL
n_tokens: 105, FAIL
n_tokens: 106, FAIL
n_tokens: 107, FAIL
n_tokens: 108, FAIL
n_tokens: 109, FAIL
n_tokens: 110, FAIL
n_tokens: 111, FAIL
n_tokens: 112, FAIL
n_tokens: 113, FAIL
n_tokens: 114, FAIL
n_tokens: 115, FAIL
n_tokens: 116, FAIL
n_tokens: 117, FAIL
n_tokens: 118, FAIL
n_tokens: 119, FAIL
n_tokens: 120, FAIL
n_tokens: 121, FAIL
n_tokens: 122, FAIL
n_tokens: 123, FAIL
n_tokens: 124, FAIL
n_tokens: 125, FAIL
n_tokens: 126, FAIL
n_tokens: 127, FAIL
n_tokens: 128, PASS

Unless there is a fix for FA to change this behaviour, I would suggest we restrict the kernel block size to [16,32,64] which is actually what I am already doing in #27753.

An alternative solution would be to try to "clean" the cache blocks before use to ensure they have no NaNs in them.