Core 注意力

Read and resolve the canonical attention backend requested by the environment.

Inspects WORLDFOUNDRY_ATTENTION_IMPLEMENTATION or WORLDFOUNDRY_ATTENTION_BACKEND, falling back to "auto" if neither is specified.

Dispatch Q/K/V attention across qualified exact backends.

Tensor layouts are described by einops-style patterns and normalized before execution. Automatic mode tries only providers that are available for the current device, dtype, and shape, remembers unsupported workload signatures, and always retains the in-tree PyTorch SDPA path as the exact fallback.

KV cache for causal attention with a fixed-size local window, CUDA-graph compatible.

Keys and values can have arbitrary shape `[..., total_size, ...]; the sequence (rolling) dimension is given by seq_dim (dimension index, can be negative). Layout along that dimension: [sink tokens | local window tokens]. Sink tokens are never evicted; the local window rolls left as new chunks are added if full. Chunks are non-overlapping: each update adds one chunk of chunk_size` tokens at the next logical position in the full sequence.

Note: Currently only supports `total_size (sink_size + window_size) divisible by chunk_size`.

Phases: - Filling: cache not yet full; tokens are written contiguously; `cached_k() / cached_v() return only the valid prefix. - Steady-state: cache full; each new chunk triggers a left-roll of the local window and overwrites the rightmost positions; cached_k() / cached_v()` return the full buffer.

The argument `chunk_idx (0, 1, 2, ...) is the index of the new chunk in the full sequence (not an index into the cache). If chunk_idx is greater than the previous one, the chunk is appended (or, in steady-state, written after the roll). If chunk_idx` equals the previous one, the same cache positions are overwritten.

Per-step usage: 1. before_update(chunk_idx) — prepare (roll local window if steady-state). 2. update(k, v) — write the new chunk's keys/values into the cache. 3. cached_k() / cached_v() — get cached keys/values for attention. 4. after_update(chunk_idx) — update internal bookkeeping.

This function is used to handle context parallel behavior, split input tensors into multiple parts and scatter them to different GPUs.

Input: cp_strategy: str. cp_ulysses for hopper or newer, cp_shuffle_overlap for 4090 or older x: (S, N, D). torch.Tensor of inputs embedding (images or latent representations of images) condition_map: (N * S). torch.Tensor determine which condition to use for each token rope: (S, 96). torch.Tensor of rope xattn_mask_for_cuda_graph: (N * denoising_range_num, L, 1, 1). torch.Tensor of xattn mask for cuda graph, None means no cuda graph core_attn_params: PackedCoreAttnParams. Packed sequence parameters for core_atten cross_attn_params: PackedCrossAttnParams. Packed sequence parameters for cross_atten

Output: x: (S', N, D). torch.Tensor of inputs embedding (images or latent representations of images) condition_map: (N * S'). torch.Tensor determine which condition to use for each token rope: (S', 96). torch.Tensor of rope cp_split_sizes: List[int]. Split sizes for each rank core_attn_params: PackedCoreAttnParams cross_attn_params: PackedCrossAttnParams

Pipeline chunked query exchange with attention computation.

The helper splits queries into `cp_shuffle_num chunks, starts the first asynchronous exchange, and rotates later query/output chunks so communication can overlap with fattn` execution.

Apply SDPA to flattened `[B, S, H*D]` Q/K/V tensors.

This is the shared layout adapter for diffusion transformers. Model code should keep projections/RoPE/mask construction locally and delegate the generic head reshape, mask canonicalization, backend context, and output merge here.

Determine if the active CUDA device is architecturally capable of running FlashAttention.

This in-tree FA2 path targets NVIDIA Ampere (SM8x), Ada (SM89), and Hopper (SM90). Blackwell kernels are resolved separately because SM100/103 and SM120 are not binary-compatible feature targets.

Context-parallel layout metadata carried between pre/post processing.

The fields describe spatial size, padding and per-rank splits, denoising ranges, optional feature-prefix behavior, CUDA-graph mode, and the packed self/cross-attention range objects. Model integrations normally construct this once per request and pass it unchanged through transformer blocks.

Normalize colloquial or varied attention backend names into standard keys.

For example, maps 'flash-attn-2', 'flash2', or 'flash_attention_2' to 'flash_attention_2' while stripping and lowering input values to handle typos gracefully.

Run PISA block-routed attention or an exact SDPA fallback.

Inputs use `[batch, heads, sequence, head_dim]. density` is the fraction of KV blocks computed exactly; the remaining blocks use centroid approximation. Because this changes model math, PISA is never selected by the generic dense-attention dispatcher without an explicit caller choice.

Probe installed attention packages and hardware compatibility.

Results are cached by runtime and compute capability, rather than globally. This keeps mixed A100/H100 nodes correct when the active tensor/device changes after module import.

Resolve the highest priority usable backend, falling back gracefully to PyTorch SDPA.

If "auto" is preferred, traverses priority list to select the first runnable package. If the requested package is not usable on the active hardware, automatically degrades to "torch" SDPA to ensure robust execution.

Resolve a backend name accepted by Transformers model configs.

Transformers currently exposes portable `eager/sdpa paths and the separately installed flash_attention_2 provider. WorldFoundry has a wider backend vocabulary (including FA3 and model-specific kernels), so a small adapter is needed before writing config._attn_implementation`. Unsupported or unavailable providers conservatively map to PyTorch SDPA.

Compute exact scaled dot-product attention through one stable core API.

The function mirrors PyTorch SDPA, adds explicit backend selection, and preserves compatibility with PyTorch versions that do not yet accept `enable_gqa`. When native SDPA is unavailable it evaluates the same attention equation with matmul, softmax, and an optional dropout.

Overlap Ulysses all-to-all communication with Q/K/V and cross attention.

Static entry points select separate, fused-KV, or fused-QKV communication schedules. Each returns self-attention output restored to model layout plus the concurrently computed cross-attention output.