Enable slicing for the FusedSDPA

[yangulei]

FusedSDPA can be split into smaller chunks to improve performance while using the padding-aware bucketing strategy which guarantees the max absolute padding in the sequence and context dimensions.

Usage

Parameter name	Description	Default value
VLLM_HPU_FSDPA_SLICE_ENABLED	Enable the slicing.	True for padding-aware bucketing strategy
VLLM_HPU_FSDPA_SLICE_SEQ_LEN_THLD	KV length threshold above which slicing is applied.	min(max_num_batched_tokens, 8192)
VLLM_HPU_FSDPA_SLICE_CHUNK_SIZE	Chunk size for q_len and kv_len in each chunk. Rounded up to the next multiple of 1024.	VLLM_HPU_FSDPA_SLICE_SEQ_LEN_THLD // 2
VLLM_HPU_FSDPA_SLICE_WITH_GRAPH_BREAKS	Places each chunk in a separate graph to reduce compilation time.	true for lazy mode and false otherwise

Important

These parameters are effective only with the padding-aware bucketing strategy set by VLLM_BUCKETING_STRATEGY="pad".

Implementation

Take the prefix-prefill with [bs, query, context] = [1, 9037, 8832] as an example. The prefill shape will first be padded to [1, 10880, 11008] by the bucketing. The attention mask will be looks like:

Not that there are padding in query and context dimensions.

The original implementation pass the full attention mask to the FusedSDPA kernel.

This PR introduced an implementation to calculate the FSDPA in chunks by slicing the Q, K and V as below:

Where the color of the rectangles differentiate the is_causal and attn_mask parameters passed to the FusedSDPA kernel:

• rgb(255,0,0): is_causal=False and attn_mask is not None
• rgb(255,255,0): is_causal=True and attn_mask=None
• rgb(255,0,255): is_causal=False and attn_mask=None

In this way, most of the chunks call the FusedSDPA without attention mask to get better performance, and the graph for the chunks might be reused across different buckets to reduce the warmup duration.

Dependencies

• Add the padding-aware bucketing strategy #762 as the number of chunks with padding is determined by the PAD_MAX for the query and context.

---

Thanks @Wei-Lin-Intel for the original idea and the detailed behavior of the FusedSDPA kernel.

[github-actions]

✅ CI Passed

All checks passed successfully against the following vllm commit:
f296a1966dca96cd69e5c1fa1264edbf611a1bd6

[yangulei]

Upstreaming the slicing for fp8 FusedSDPA.

[yangulei]

@afierka-intel
This PR is re-implemented to enable slicing for both bf16 and fp8 FusedSDPA. Please help to review.

cc: @czhu15

[github-actions]

✅ CI Passed

All checks passed successfully against the following vllm commit:
8eb401134e750781a202c0b6dc4059616cdb4954