[Core] Add Cosmos3 attention strategy PoC

[rahul-steiger-nv]

Purpose

This is the Cosmos3 PoC I offered in @nodeeeeee's RFC #8382. It switches between dense and sparse attention by denoising step and layer, with layout selection outside compiled execution and shared model weights.

It adds attention presets, layouts, schedules, checkpoint defaults with deployment overrides, and warmup across layouts. Recipes demonstrate a uniform dense-to-sparse switch and a mixed-layer policy. See design, usage, and scope.

Dependency: #8335 supplies the shared block-sparse foundation and FA4 adapter. The diff against main includes prerequisite changes; review only the strategy PoC.

Test Plan

vLLM Version: 0.31.0
vLLM-Omni Commit: b9f81e58cdbc5f05a6497ffd90960867f8b36acf (rebased onto the updated #8335 foundation at 5e2569678, which merges upstream main at 4c5541cfc)

Environment: GH200, Enroot ARM64 image, PyTorch 2.13.0+cu130, FA3 dense and FA4 4.0.0b33 sparse attention.

python3 -m pytest -q \
  tests/diffusion/attention/test_attention_strategy*.py \
  tests/diffusion/attention/test_attention_checkpoint_policy.py \
  tests/diffusion/attention/test_cosmos3_action_strategy.py \
  tests/diffusion/attention/test_attention_config.py \
  tests/diffusion/models/cosmos3/test_cosmos3_attention_strategy.py \
  tests/diffusion/models/cosmos3/test_cosmos3_transformer.py \
  tests/diffusion/test_diffusion_engine_dummy_run.py \
  tests/diffusion/test_diffusion_model_runner.py \
  tests/diffusion/models/cosmos3/test_cosmos3_strategy_offload.py \
  tests/diffusion/models/cosmos3/test_cosmos3_strategy_ulysses.py \
  tests/diffusion/distributed/test_sp_plan_hooks.py \
  tests/diffusion/distributed/test_cosmos3_pre_sharded.py \
  tests/diffusion/diffusion_kv/test_paged_attention_adapter.py \
  tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py

Test Result

516 tests passed, 5 skipped after the 2026-10-09 rebase. Only the three Cosmos3 strategy commits were replayed onto the updated foundation; their patches are unchanged. Coverage includes strategy/configuration behavior, CUDA compilation and reuse, Cosmos3 regressions, single-GPU CPU offload, CPU/Gloo Ulysses, and upstream sharding/paged-KV integration. The rebase preserves upstream Cosmos3 sharding before the generation stack and rank-local embedding ownership.

The five skips require two GPUs. Earlier two-Hopper strategy validation covered strict Ulysses with no/model/layerwise CPU offload in eager and regional modes; those CUDA cases have not been rerun after this rebase. Applicable pre-commit hooks passed except mypy-3.10: all 43 errors exactly match the foundation, with no new diagnostics. Full CI has not run for this PR.

Historical generation measurements: taken before this rebase at a8525998d with vLLM 0.30.0; not rerun on the rebased branch.

Cosmos3-Nano: default prompt, seed 42, 720p, 189 frames, 35 steps, mixed precision disabled, and 75% target sparsity. One full warmup and one measured generation per mode.

Generation attention	Generation latency	Speedup
Dense	194.66 s	1.00×
Mixed: 10 dense steps, then 28/36 sparse layers	159.42 s	1.22×
All 36 generation layers sparse throughout	131.21 s	1.48×

Expected sparse dispatch was verified at every step, with no measured-run recompilation or thermal throttling. Latency includes decode; it excludes loading, warmup, and MP4 encoding. The benchmark replaces synthetic startup warmup with a full generation.

These settings are illustrative: no optimal strategy or minimal sparse-layer count is claimed. One prompt/seed and one measurement per mode demonstrate preliminary speedup, not general quality preservation or completion of the RFC's validation requirements.

AI assistance: Codex assisted with code review, validation, benchmarking, and preparation of this PR description.

[rahul-steiger-nv]

Cosmos3-Nano generated outputs: dense vs. mixed sparse vs. all-sparse, using the same prompt and seed. Labels show generation speedup, not playback speed.

comparison-captioned.mp4