vllm - 💡(How to fix) Fix [Bug]: MLA + FP8 KV cache + CUDA Graph causes random NaN in decode phase [5 comments, 4 participants]

TL;DR

The issue can likely be resolved by using enforce_eager=True or switching to BF16 when using the MLA Model with FP8 kv cache.

Guidance

• When enforce_eager=False and using FP8, the MLAAttention forward pass results in NaN values for q, kv_c_normed, and k_pe, indicating a potential numerical instability issue.
• Using enforce_eager=True with FP8 kv cache or switching to BF16 resolves the issue, suggesting that the problem is related to the interaction between the enforce_eager flag and the FP8 data type.
• To mitigate the issue, try setting enforce_eager=True when using FP8 kv cache or experiment with different data types, such as BF16, to find a stable configuration.
• Verify that the issue is resolved by checking the values of q, kv_c_normed, and k_pe in the MLAAttention forward pass for NaN values.

Notes

The exact cause of the numerical instability is unclear, and further investigation may be necessary to determine the root cause of the issue.

Recommendation

Apply workaround: Use enforce_eager=True or switch to BF16 when using the MLA Model with FP8 kv cache, as these configurations have been shown to resolve the issue.