[XPU][Attention] Enable per-sequence causal masking with FA2

[gc-fu]

Dependency and coordination

This PR depends on vllm-project/vllm-xpu-kernels#606 (or a release containing
that change), which adds the optional dynamic_causal FA2 argument.

Open PR #45774 overlaps the DiffusionGemma XPU startup portion but targets an
older tree and currently conflicts with main. It copies UVA inputs to XPU;
this change instead keeps the small sampler helpers eager and uses the current
mainline TP-local embedding matmul/all-reduce design.

Purpose

DiffusionGemma needs a per-sequence causal/bidirectional mask in one batch.
vLLM already builds int32[num_seqs] dynamic-causal metadata and symmetrizes a
causal sliding window for mixed batches, but the FlashAttention backend only
allowed this tensor with FA4 and the XPU wrapper did not forward it.

This change:

• allows tensor causal metadata for XPU FA2 while retaining FA4 support;
• forwards dynamic_causal through _xpu_ops.flash_attn_varlen_func to the
  kernel package and keeps the existing symmetric-window contract;
• adds a standalone, non-default-CI GSM8K configuration for
  google/diffusiongemma-26B-A4B-it, runtime --quantization fp8, TP2,
  FLASH_ATTN, eager execution, and concurrency 2; and
• keeps explicit DiffusionGemma sampler helpers eager on XPU and stores/casts
  self-conditioning state in the model dtype. These two current-main fixes are
  required for TP2 evaluation to reach attention: without them startup fails
  first on a Triton XPU pointer and then on an FP32/BF16 indexed assignment.

Test Plan

Focused unit tests:

.venv/bin/python -m pytest -q \
  tests/models/test_diffusion_gemma_xpu.py \
  tests/kernels/attention/test_xpu_dynamic_causal.py

Focused pre-commit:

pre-commit run --files \
  vllm/_xpu_ops.py \
  vllm/model_executor/models/diffusion_gemma.py \
  vllm/v1/attention/backends/flash_attn.py \
  tests/evals/gsm8k/configs/DiffusionGemma-26B-A4B-it-XPU-FA2-FP8.yaml \
  tests/kernels/attention/test_xpu_dynamic_causal.py \
  tests/models/test_diffusion_gemma_xpu.py

Full GSM8K qualification used a local mirror of the Google checkpoint and a
machine-local --gpu-memory-utilization 0.75 override; the committed YAML does
not include that machine-specific setting:

.venv/bin/python -m pytest -s -v \
  tests/evals/gsm8k/test_gsm8k_correctness.py \
  --config-list-file=/llm/artifacts/x07-dynamic-causal-20260920/gsm8k-full-final-list.txt

The upstream-style reproduction uses the committed config through a temporary
one-entry list:

printf '%s\n' \
  tests/evals/gsm8k/configs/DiffusionGemma-26B-A4B-it-XPU-FA2-FP8.yaml \
  > /tmp/diffusiongemma-xpu-gsm8k.txt
.venv/bin/python -m pytest -s -v \
  tests/evals/gsm8k/test_gsm8k_correctness.py \
  --config-list-file=/tmp/diffusiongemma-xpu-gsm8k.txt

Test Result

• Focused unit tests: 5 passed.
• Focused pre-commit: all applicable hooks passed.
• Two-request smoke test: 2/2 questions, accuracy 1.0000, invalid rate 0,
  approximately 1.8 questions/s.
• Full GSM8K: 1319/1319 valid responses, accuracy 0.9249, invalid rate 0,
  723.4 seconds, approximately 1.8 questions/s; acceptance floor is
  0.84 - 0.08 = 0.76.
• The companion kernel correctness matrix passed 31 tests, including a mixed
  causal/bidirectional single launch and XPU graph replay.

pre-commit run --all-files was also attempted. Relevant code, formatting,
type, and metadata hooks passed. The overall command did not pass because
pip-compile-rocm regenerated requirements/test/rocm.txt with dependencies
from the shared XPU runtime (the generated change was restored), and
test-nonroot-entrypoint case 3 did not replace a read-only HOME with a
temporary home. These are environment-sensitive failures unrelated to the
changed files; they are not reported as passing.

An earlier full-run attempt used the harness default 600-second request
timeout. Because that timeout includes time waiting for the concurrency-2
connection pool, 224 queued requests timed out. That run was discarded and the
reported run changed only the request timeout to 1800 seconds.

Duplicate-work check

Repeated immediately before preparing this draft:

• Open PR [Kernel] Cover sliding window in mixed per-sequence causal attention tests #51257 adds CUDA/FA4 mixed-causal sliding-window test coverage; it
  does not implement or wire XPU FA2.
• Open PR [DiffusionGemma-26B-A4B-it][XPU] Fix DiffusionGemma for TP>1 and enable XPU #45774 overlaps DiffusionGemma XPU startup but is stale/conflicting
  and uses the older UVA-copy and TP-gather approach described above.
• Closed PR [XPU] Enable DiffusionGemma block-diffusion inference on Intel XPU #46785 implemented older variants of the XPU compile/dtype/TP
  fixes. Current main already has a newer TP-local matmul/all-reduce path, so
  its old TP gather patch was not copied.
• No open vLLM PR found by dynamic_causal or XPU FA2 searches implements the
  wrapper/backend change in this PR.

Limitations

Validation was performed on Intel Arc Pro B70 only, in the existing PyTorch
2.13.0+xpu environment. The companion kernel repository currently declares
PyTorch 2.14.0+xpu; that mismatch remains a validation limitation. The model
evaluation is intentionally not added to a default CI list.

AI assistance was used during implementation and validation. The human
submitter must review and understand every changed line before submission.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gc-fu.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork