Reapply [Attention][FA3] Update FA3 to include new swizzle optimization

[LucasWilkinson]

Reapply #23465 after revert in #33841 but with correct metadata sizes

[gemini-code-assist]

Code Review

This pull request updates the flash-attention dependency and adjusts the scheduler_metadata buffer size to accommodate changes in flash-attention and handle edge cases related to CUDA graph capturing. The changes appear correct and necessary. My review focuses on improving the clarity of the comments associated with these critical buffer size calculations. Enhancing these comments will improve maintainability and help prevent future confusion or bugs, especially since this area of the code has been a source of issues in the past.

[gemini-code-assist]

The comment explaining the buffer size calculation is a bit confusing as it compares max_cudagraph_size (number of tokens) with max_num_seqs (number of sequences) without full context. A more detailed explanation would improve clarity and maintainability, especially given the history of issues with this code block.

During CUDA graph capture for uniform decode batches, the number of requests (num_reqs) is set to be equal to the number of tokens (num_tokens). This num_tokens can be up to max_cudagraph_capture_size. In some configurations (especially in tests), max_cudagraph_capture_size can be larger than max_num_seqs. Since the scheduler_metadata buffer size depends on the number of requests, we must allocate a buffer large enough for this capture-time scenario.

Suggested change

	# Times 4 due to:
	# https://github.com/vllm-project/flash-attention/blob/3223650ccabe622a0fcae65eec706a50186a89f7/hopper/flash_api.cpp#L650-L653
	# For some tests max_cudagraph_size > max_num_seqs,
	# so we need to use the larger one.
	# During CUDA graph capture, `num_reqs` can be up to `max_cudagraph_size`.
	# This can be larger than `max_num_seqs`, so we use the max of the two.
	# The size is multiplied by 4 due to flash-attention's requirement:
	# https://github.com/vllm-project/flash-attention/blob/3223650ccabe622a0fcae65eec706a50186a89f7/hopper/flash_api.cpp#L650-L653

[gemini-code-assist]

The comment explaining the buffer size calculation is a bit confusing as it compares max_cudagraph_size (number of tokens) with max_num_seqs (number of sequences) without full context. A more detailed explanation would improve clarity and maintainability, especially given the history of issues with this code block.

During CUDA graph capture for uniform decode batches, the number of requests (num_reqs) is set to be equal to the number of tokens (num_tokens). This num_tokens can be up to max_cudagraph_capture_size. In some configurations (especially in tests), max_cudagraph_capture_size can be larger than max_num_seqs. Since the scheduler_metadata buffer size depends on the number of requests, we must allocate a buffer large enough for this capture-time scenario.

Suggested change

	# Times 4 due to:
	# https://github.com/vllm-project/flash-attention/blob/3223650ccabe622a0fcae65eec706a50186a89f7/hopper/flash_api.cpp#L650-L653
	# For some tests max_cudagraph_size > max_num_seqs,
	# so we need to use the larger one.
	# During CUDA graph capture, `num_reqs` can be up to `max_cudagraph_size`.
	# This can be larger than `max_num_seqs`, so we use the max of the two.
	# The size is multiplied by 4 due to flash-attention's requirement:
	# https://github.com/vllm-project/flash-attention/blob/3223650ccabe622a0fcae65eec706a50186a89f7/hopper/flash_api.cpp#L650-L653

[ProExpertProg]

LGTM