Does FA4 support split-kv flashdecode with cudagraph?

[biandangan]

Hi,

PR #1940 adds parallelism across the KV dimension and implements FlashDecode.

However, in SGLang's FlashAttentionBackend, there is a comment:

"FA4 does not support num_splits=0 with CUDA Graph, so we set num_splits to 1 if CUDA Graph is enabled."

where num_splits=0 means "use a heuristic to automatically determine the optimal number of splits." Forcing num_splits=1 effectively disables KV-dimension parallelism, which suggests FlashDecode may not be active when CUDA Graph is enabled.

Does FA4 support FlashDecode with CUDA Graph for MLA/MQA/...? If not, is this a known limitation or something planned for a future release?

[reubenconducts]

Cuda graph support for split-kv will result from #2559, which is under review. The solution is to add a metadata tensor holding num splits, which is populated in a separate prepare kernel and used in the main kernel by a modified varlen scheduler.