Core attention
Exact SDPA, layout dispatch, backend policy, RoPE, packed sequences, context parallelism, and block KV cache.
On this page
WorldFoundry keeps attention projection and model semantics in model code, while Core owns the mechanics that are repeatedly implemented incorrectly: head-layout conversion, exact backend selection, mask handling, RoPE application, packed sequence ranges, context-parallel exchange, and rolling KV state.
Choose the narrowest entry point
Use scaled_dot_product_attention when Q, K, and V already have a standard split-head shape. It follows the PyTorch SDPA contract, adds an explicit backend context, supports compatible GQA expansion, and preserves an exact matmul fallback.
Use flattened_multihead_attention when tensors are (batch, sequence, hidden) and only the head count is missing. It performs the split/merge centrally. Use attention_forward when an integration has einops-style QKV layouts or opts into optional providers through the dispatch policy. A mask or compatibility_mode=True deliberately returns to the PyTorch path.
NativeAttention packages SDPA as a module and can attach a context-parallel group. ContextParallelAttention, UlyssesScheduler, and CSOHelper are lower-level distributed mechanisms; use them only when the model runtime owns the matching split metadata.
Exact CPU/GPU example
import torch
from worldfoundry.core import scaled_dot_product_attention
torch.manual_seed(7)
q = torch.randn(2, 8, 32, 64)
k = torch.randn(2, 8, 48, 64)
v = torch.randn(2, 8, 48, 96)
output = scaled_dot_product_attention(
q,
k,
v,
dropout_p=0.0,
backend="math", # deterministic explicit provider for this example
)
assert output.shape == (2, 8, 32, 96)At inference time, keep dropout_p=0.0; PyTorch SDPA does not infer that value from module.eval(). A boolean mask means “keep this score,” while a floating mask is added to the score matrix.
Block KV cache lifecycle
BlockKVCache is not a dictionary that can be updated in arbitrary order. Every chunk follows before_update → update → cached_k/cached_v → after_update. Repeating the current chunk_idx overwrites the same logical chunk; advancing by one appends or rolls the local window. Skipping an index is an error.
import torch
from worldfoundry.core.attention import BlockKVCache
cache = BlockKVCache(
k_shape=(1, 1, 4, 2),
v_shape=(1, 1, 4, 3),
seq_dim=2,
chunk_size=2,
window_size=4,
device="cpu",
dtype=torch.float32,
)
for chunk_idx in range(3):
k = torch.full((1, 1, 2, 2), float(chunk_idx))
v = torch.full((1, 1, 2, 3), float(chunk_idx))
cache.before_update(chunk_idx)
cache.update(k, v)
visible_k = cache.cached_k() # 2, then 4, then rolling 4 tokens
cache.after_update(chunk_idx)sink_size reserves an immutable prefix and window_size describes the rolling region. Their sum must match the cache sequence dimension and be divisible by chunk_size.
Complete reference
The blocks below are the generated signatures for this category. Use the on-page symbol index to jump; source links open the defining implementation behind each lazy export.
41 public symbols
def apply_nd_rotary_embedding(query: torch.Tensor,key: torch.Tensor,freqs_cis: torch.Tensor | tuple[torch.Tensor, torch.Tensor],head_first: bool = False,start_offset: int = 0) -> tuple[torch.Tensor, torch.Tensor]worldfoundry.core.apply_nd_rotary_embeddingfrom worldfoundry.core import apply_nd_rotary_embeddingOverview
Apply precomputed n-D rotary frequencies to query and key tensors. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: tuple[torch.Tensor, torch.Tensor].
Parameters
querytorch.Tensor- Query tensor in `
[B, S, H, D]layout by default, or[B, H, S, D]whenhead_first=True`. keytorch.Tensor- Key tensor with the same layout convention and head width.
freqs_cistorch.Tensor | tuple[torch.Tensor, torch.Tensor]- Complex rotary frequencies or an explicit `
(cos, sin)` pair covering the requested sequence window. head_firstbool- Select `
[B, H, S, D]` sequence-axis interpretation.default:False start_offsetint- First frequency position, used with an existing KV prefix.default:
0
Notes
The final head width must be compatible with complex pairs. Build matching structured frequencies with `get_nd_rotary_pos_embed`.
Raises
ValueError- Frequency tensors do not cover the requested shape.
Returns: tuple[torch.Tensor, torch.Tensor] — Rotated `(query, key)` tensors with original dtype and shape.
def apply_rotary_embedding(value: Any,cos: Any,sin: Any,rotary_dim: int | None = None,interleaved: bool = False) -> Anyworldfoundry.core.apply_rotary_embeddingfrom worldfoundry.core import apply_rotary_embeddingOverview
Apply RoPE to the leading `rotary_dim features of the last dimension. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: Any`.
Parameters
valueAnycosAnysinAnyrotary_dimint | None- default:
None interleavedbool- default:
False
Returns: Any
def attention_backend_capability(name: str,device: torch.device | str | int | None = None) -> AttentionKernelCapabilityworldfoundry.core.attention_backend_capabilityfrom worldfoundry.core import attention_backend_capabilityOverview
Retrieve runtime capability metadata for a single normalized backend. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: AttentionKernelCapability.
Parameters
namestrdevicetorch.device | str | int | None- default:
None
Returns: AttentionKernelCapability
def attention_backend_context(backend: Literal['math', 'efficient', 'cudnn', 'flash'] | Any | None = None,backends: Any = None) -> Anyworldfoundry.core.attention_backend_contextfrom worldfoundry.core import attention_backend_contextOverview
Return a context manager that selects PyTorch SDPA backends via core. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: Any.
Parameters
backendLiteral['math', 'efficient', 'cudnn', 'flash'] | Any | None- default:
None backendsAny- default:
None
Returns: Any
def attention_backend_from_env(environ: Mapping[str, str] | None = None) -> strworldfoundry.core.attention_backend_from_envfrom worldfoundry.core import attention_backend_from_envOverview
Read and resolve the canonical attention backend requested by the environment. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: str.
Source docstring
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.
Parameters
environMapping[str, str] | None- default:
None
Returns: str
def attention_backend_info() -> AttentionBackendInfoworldfoundry.core.attention_backend_infofrom worldfoundry.core import attention_backend_infoOverview
Return the available generic PyTorch attention backend. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: AttentionBackendInfo.
Returns: AttentionBackendInfo
def attention_backend_report() -> tuple[AttentionKernelCapability, ]worldfoundry.core.attention_backend_reportfrom worldfoundry.core import attention_backend_reportOverview
Retrieve capability status of all registered backends, ordered by dispatch priority. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: tuple[AttentionKernelCapability, ...].
Returns: tuple[AttentionKernelCapability, ...]
def attention_dispatch_report() -> dict[str, object]worldfoundry.core.attention.attention_dispatch_reportfrom worldfoundry.core.attention import attention_dispatch_reportOverview
Return lightweight operator-selection and quarantine state. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: dict[str, object].
Returns: dict[str, object]
def attention_forward(q: torch.Tensor,k: torch.Tensor,v: torch.Tensor,q_pattern = 'b n s d',k_pattern = 'b n s d',v_pattern = 'b n s d',out_pattern = 'b n s d',dims = None,attn_mask = None,scale = None,compatibility_mode = False)worldfoundry.core.attention_forwardfrom worldfoundry.core import attention_forwardOverview
Layout-aware attention entry that adapts einops-style QKV packs and optional fused providers through the dispatch policy.
Source docstring
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.
Parameters
qtorch.Tensor- Query tensor in `
q_pattern` layout. ktorch.Tensor- Key tensor in `
k_pattern` layout. vtorch.Tensor- Value tensor in `
v_pattern` layout. q_pattern- Layout pattern for `
q`.default:'b n s d' k_pattern- Layout pattern for `
k`.default:'b n s d' v_pattern- Layout pattern for `
v`.default:'b n s d' out_pattern- Requested output layout.default:
'b n s d' dims- Named dimensions needed to expand grouped pattern terms such as `
(n d)`.default:None attn_mask- Optional boolean or additive attention mask. A mask forces the PyTorch compatibility path because optional fused providers do not share one mask contract.default:
None scale- Optional softmax scale; `
None` uses the backend default.default:None compatibility_mode- Skip optional providers and execute PyTorch SDPA directly.default:
False
Notes
Backend choice can be inspected with `attention_dispatch_report. Use clear_attention_dispatch_cache` after changing provider availability inside a long-lived process.
Raises
RuntimeError- The selected provider fails for a reason other than an unsupported kernel or optional-library availability problem.
class AttentionBackendInfo(backend: str, uses_torch_sdpa: bool)worldfoundry.core.AttentionBackendInfofrom worldfoundry.core import AttentionBackendInfoOverview
Resolved attention backend metadata after probing and normalization.
Attributes
backendstruses_torch_sdpabool
class AttentionKernelCapability(name: str,package: str,available: bool,usable: bool,reason: str = '')worldfoundry.core.AttentionKernelCapabilityfrom worldfoundry.core import AttentionKernelCapabilityOverview
Runtime availability metadata for a single attention kernel family. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Attributes
namestr- Canonical identifier of the attention backend.
packagestr- The underlying Python package or C++/CUDA extension.
availablebool- True if the module/package is physically installed in the environment.
usablebool- True if the backend is physically runnable on the active hardware (e.g., CUDA capability).
reasonstr- Explains why the backend is unusable or unavailable, if applicable.default:
''
BlockKVCache
clsclass BlockKVCache(k_shape: tuple[int, ...],v_shape: tuple[int, ...],seq_dim: int,chunk_size: int,window_size: int,sink_size: int = 0,device: torch.device | str = torch.device('cuda'),dtype: torch.dtype = torch.float16)worldfoundry.core.attention.BlockKVCachefrom worldfoundry.core.attention import BlockKVCacheOverview
Chunked rolling KV cache with a strict before_update → update → after_update lifecycle. Do not skip chunk indices.
Source docstring
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.
Attributes
k_shapetuple[int, ...]v_shapetuple[int, ...]seq_dimintchunk_sizeintwindow_sizeintsink_sizeint- default:
0 devicetorch.device | str- default:
torch.device('cuda') dtypetorch.dtype- default:
torch.float16
Methods
Overview
Number of valid cached tokens visible to attention.
Parameters
self
Returns: int
Overview
Right edge of the current chunk in the physical cache layout.
Parameters
self
Returns: int
Overview
Build a single-chunk cache pre-filled with the given key and value tensors.
Parameters
kTensorvTensorseq_dimint
Returns: Self
Overview
Return True if the cache is full (steady-state phase).
Returns: bool
Overview
Prepare the cache before writing new tokens.
Source docstring
Prepare the cache before writing new tokens.
If `chunk_idx equals the previous chunk index, this is a no-op. Otherwise, we expect the chunk_idx` to be +1 from the previous chunk index. In this case, we will roll the local window left if the cache is in steady-state, or no op if the cache is in filling phase.
Parameters
chunk_idxint- Chunk index of the new chunk in the full sequence.
Returns: None
Overview
Write the new chunk's keys and values into the cache.
Source docstring
Write the new chunk's keys and values into the cache.
Must be called after `before_update() and before after_update()`.
Parameters
kTensor- Keys; shape must match cached keys except at seq_dim, where length must be chunk_size.
vTensor- Values; shape must match cached values except at seq_dim, where length must be chunk_size.
Returns: None
Overview
Finalize bookkeeping after writing new tokens.
Source docstring
Finalize bookkeeping after writing new tokens.
Updates `_prev_chunk_idx and, in filling phase, _n_cached`.
Parameters
chunk_idxint- The index of the new chunk in the full sequence.
Returns: None
Overview
Return cached keys for attention (valid prefix in filling phase, full buffer in steady-state).
Returns: Tensor
Overview
Return cached values for attention (valid prefix in filling phase, full buffer in steady-state).
Returns: Tensor
def clear_attention_dispatch_cache() -> Noneworldfoundry.core.attention.clear_attention_dispatch_cachefrom worldfoundry.core.attention import clear_attention_dispatch_cacheOverview
Clear workload decisions and runtime failure quarantine. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: None.
Returns: None
class ContextParallelAttention(qkv_format: Literal['bhsd', 'bshd'] = 'bhsd',backend: Literal['cudnn', 'flash'] = 'cudnn',method: Literal['ring', 'ulysses'] = 'ring',convert_to_fp32: bool = True)worldfoundry.core.attention.ContextParallelAttentionfrom worldfoundry.core.attention import ContextParallelAttentionOverview
Context-parallel attention with selectable method and SDPA backend. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Parameters
qkv_formatLiteral['bhsd', 'bshd']- Layout of the QKV tensors; `
"bhsd"or"bshd"`.default:'bhsd' backendLiteral['cudnn', 'flash']- SDPA backend; `
"cudnn"or"flash"`.default:'cudnn' methodLiteral['ring', 'ulysses']- Context-parallelism strategy; `
"ring"or"ulysses"`.default:'ring' convert_to_fp32bool- Promote LSE accumulators to fp32 during ring merges.default:
True
cp_post_process
funcdef cp_post_process(cp_size: int,cp_strategy: str,x: torch.Tensor,meta_args: ModelMetaArgs) -> torch.Tensorworldfoundry.core.cp_post_processfrom worldfoundry.core import cp_post_processOverview
Gather context-parallel output back to the original sequence layout. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: torch.Tensor.
Parameters
cp_sizeint- Context-parallel world size. `
1returnsx` unchanged. cp_strategystr- `
"cp_ulysses"or"cp_shuffle_overlap"`. xtorch.Tensor- This rank's local output tensor.
meta_argsModelMetaArgs- Split sizes and padding recorded by `
cp_pre_process`.
Raises
ValueError- `
cp_strategy` is unknown.
Returns: torch.Tensor — Globally gathered tensor with shuffle-overlap padding removed.
cp_pre_process
funcdef cp_pre_process(cp_size: int,cp_strategy: str,x: torch.Tensor,condition_map: torch.Tensor,rope: torch.Tensor,xattn_mask_for_cuda_graph: Union[torch.Tensor, None],ardf_meta: dict,core_attn_params: PackedCoreAttnParams,cross_attn_params: PackedCrossAttnParams)worldfoundry.core.cp_pre_processfrom worldfoundry.core import cp_pre_processOverview
This function is used to handle context parallel behavior, split input tensors into multiple parts and scatter them to different GPUs. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Source docstring
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
Parameters
cp_sizeintcp_strategystrxtorch.Tensorcondition_maptorch.Tensorropetorch.Tensorxattn_mask_for_cuda_graphUnion[torch.Tensor, None]ardf_metadictcore_attn_paramsPackedCoreAttnParamscross_attn_paramsPackedCrossAttnParams
def cso_communication(input: torch.Tensor,cp_world_size: int,cp_split_sizes: List[int],comm_type: str = None) -> Tuple[torch.Tensor, torch.distributed.Work]worldfoundry.core.cso_communicationfrom worldfoundry.core import cso_communicationOverview
Launch one context-shuffle-overlap all-to-all operation. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: Tuple[torch.Tensor, torch.distributed.Work].
Parameters
inputtorch.Tensor- Local sequence/head tensor.
cp_world_sizeint- Context-parallel group size.
cp_split_sizesList[int]- Per-rank split sizes.
comm_typestr- `
"kv"` additionally reshapes KV heads before exchange.default:None
Returns: Tuple[torch.Tensor, torch.distributed.Work] — Output buffer and asynchronous work handle. For one rank, the input and a completed fake handle are returned.
CSOHelper
clsclass CSOHelper(cp_shuffle_num,cp_world_size,cp_split_sizes)worldfoundry.core.CSOHelperfrom worldfoundry.core import CSOHelperOverview
Pipeline chunked query exchange with attention computation. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Source docstring
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.
Parameters
cp_shuffle_numcp_world_sizecp_split_sizes
def flattened_multihead_attention(query: Tensor,key: Tensor,value: Tensor,num_heads: int,attn_mask: Tensor | None = None,dropout_p: float = 0.0,is_causal: bool = False,scale: float | None = None,backend: Literal['math', 'efficient', 'cudnn', 'flash'] | Any | None = None,backends: Any = None) -> Tensorworldfoundry.core.flattened_multihead_attentionfrom worldfoundry.core import flattened_multihead_attentionOverview
Convenience wrapper for (batch, sequence, hidden) tensors that only need a head count to split/merge around SDPA.
Source docstring
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.
Parameters
queryTensorkeyTensorvalueTensornum_headsintattn_maskTensor | None- default:
None dropout_pfloat- default:
0.0 is_causalbool- default:
False scalefloat | None- default:
None backendLiteral['math', 'efficient', 'cudnn', 'flash'] | Any | None- default:
None backendsAny- default:
None
Returns: Tensor
def get_1d_rotary_pos_embed(dim: int,pos: torch.Tensor | int,theta: float = 10000.0,use_real: bool = False,theta_rescale_factor: float = 1.0,interpolation_factor: float = 1.0) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]worldfoundry.core.get_1d_rotary_pos_embedfrom worldfoundry.core import get_1d_rotary_pos_embedOverview
Build one-dimensional rotary position frequencies. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: torch.Tensor | tuple[torch.Tensor, torch.Tensor].
Parameters
dimintpostorch.Tensor | intthetafloat- default:
10000.0 use_realbool- default:
False theta_rescale_factorfloat- default:
1.0 interpolation_factorfloat- default:
1.0
Returns: torch.Tensor | tuple[torch.Tensor, torch.Tensor]
get_meshgrid_nd
funcdef get_meshgrid_nd(start: int | tuple[int, ...] | list[int],*args: int | tuple[int, ...] | list[int],dim: int = 2,device: torch.device | str | None = None) -> torch.Tensorworldfoundry.core.get_meshgrid_ndfrom worldfoundry.core import get_meshgrid_ndOverview
Build an n-D meshgrid with PyTorch `linspace(endpoint=False) semantics. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: torch.Tensor`.
Parameters
startint | tuple[int, ...] | list[int]argsint | tuple[int, ...] | list[int]dimint- default:
2 devicetorch.device | str | None- default:
None
Returns: torch.Tensor
def get_nd_rotary_pos_embed(rope_dim_list: list[int],start: int | tuple[int, ...] | list[int],*args: int | tuple[int, ...] | list[int],theta: float = 10000.0,use_real: bool = False,theta_rescale_factor: float | list[float] = 1.0,interpolation_factor: float | list[float] = 1.0,device: torch.device | str | None = None) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]worldfoundry.core.get_nd_rotary_pos_embedfrom worldfoundry.core import get_nd_rotary_pos_embedOverview
Build n-D RoPE frequencies for tokens with structured grid coordinates. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: torch.Tensor | tuple[torch.Tensor, torch.Tensor].
Parameters
rope_dim_listlist[int]startint | tuple[int, ...] | list[int]argsint | tuple[int, ...] | list[int]thetafloat- default:
10000.0 use_realbool- default:
False theta_rescale_factorfloat | list[float]- default:
1.0 interpolation_factorfloat | list[float]- default:
1.0 devicetorch.device | str | None- default:
None
Returns: torch.Tensor | tuple[torch.Tensor, torch.Tensor]
def gpu_supports_flash_attention(device: torch.device | str | int | None = None) -> boolworldfoundry.core.gpu_supports_flash_attentionfrom worldfoundry.core import gpu_supports_flash_attentionOverview
Determine if the active CUDA device is architecturally capable of running FlashAttention. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: bool.
Source docstring
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.
Parameters
devicetorch.device | str | int | None- default:
None
Returns: bool
class ModelMetaArgs(H: int,W: int,cp_pad_size: int,cp_split_sizes: List[int],slice_point: int,denoising_range_num: int,range_num: int,extract_prefix_video_feature: bool,fwd_extra_1st_chunk: bool,distill_nearly_clean_chunk: bool,clip_token_nums: int,enable_cuda_graph: bool,core_attn_params: PackedCoreAttnParams,cross_attn_params: PackedCrossAttnParams)worldfoundry.core.ModelMetaArgsfrom worldfoundry.core import ModelMetaArgsOverview
Context-parallel layout metadata carried between pre/post processing. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Source docstring
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.
Attributes
HintWintcp_pad_sizeintcp_split_sizesList[int]slice_pointintdenoising_range_numintrange_numintextract_prefix_video_featureboolfwd_extra_1st_chunkbooldistill_nearly_clean_chunkboolclip_token_numsintenable_cuda_graphboolcore_attn_paramsPackedCoreAttnParamscross_attn_paramsPackedCrossAttnParams
class NativeAttention(qkv_format: Literal['bhsd', 'bshd'] = 'bhsd',backend: Literal['math', 'efficient', 'cudnn', 'flash'] = 'cudnn')worldfoundry.core.attention.NativeAttentionfrom worldfoundry.core.attention import NativeAttentionOverview
Module form of Core SDPA that can attach a context-parallel process group.
Parameters
qkv_formatLiteral['bhsd', 'bshd']- Layout of the QKV tensors; `
"bhsd"is(B, H, S, D),"bshd"is(B, S, H, D)`.default:'bhsd' backendLiteral['math', 'efficient', 'cudnn', 'flash']- SDPA backend selected via `
sdpa_kernel`.default:'cudnn'
Methods
Overview
Enable or disable context parallelism for ring attention.
Parameters
cp_groupProcessGroup | None- Process group for context parallel; use None to disable.
Returns: None
Overview
Return True if context parallelism is active.
Returns: bool
Overview
Return the context parallel world size, or 1 if disabled.
Returns: int
Overview
Run context-parallel SDPA (or single-rank SDPA when CP is disabled).
Parameters
queryTensor- Query tensor in configured `
qkv_format`. keyTensor- Key tensor in configured `
qkv_format`. valueTensor- Value tensor in configured `
qkv_format`.
Returns: Tensor — Attention output in the same format as inputs.
def normalize_attention_backend(value: str | None) -> strworldfoundry.core.normalize_attention_backendfrom worldfoundry.core import normalize_attention_backendOverview
Normalize colloquial or varied attention backend names into standard keys. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: str.
Source docstring
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.
Parameters
valuestr | None
Returns: str
def packed_sequence_attention(q: torch.Tensor,k: torch.Tensor,v: torch.Tensor,num_heads: int,compatibility_mode = False,scale = None)worldfoundry.core.attention.packed_sequence_attentionfrom worldfoundry.core.attention import packed_sequence_attentionOverview
Apply the shared dispatcher to flattened packed-sequence Q/K/V. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Parameters
qtorch.Tensor- Query tensor shaped `
(batch, sequence, heads * head_dim)`. ktorch.Tensor- Key tensor with the same flattened-head convention.
vtorch.Tensor- Value tensor with the same flattened-head convention.
num_headsint- Head count used to split the final dimension.
compatibility_mode- Force the exact PyTorch SDPA path.default:
False scale- Optional attention softmax scale.default:
None
Notes
This adapter is appropriate when the model already packs heads into the hidden dimension. Use `attention_forward` directly for custom layouts or explicit masks.
class PackedCoreAttnParams(q_range: torch.Tensor,k_range: torch.Tensor,np_q_range: np.ndarray,np_k_range: np.ndarray,max_seqlen_q: int,max_seqlen_k: int)worldfoundry.core.PackedCoreAttnParamsfrom worldfoundry.core import PackedCoreAttnParamsOverview
Cumulative ranges and maximum lengths for packed self-attention. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Attributes
q_rangetorch.Tensor- Device-side query range boundaries.
k_rangetorch.Tensor- Device-side key range boundaries.
np_q_rangenp.ndarray- CPU/Numpy query boundaries used by planning code.
np_k_rangenp.ndarray- CPU/Numpy key boundaries used by planning code.
max_seqlen_qint- Maximum query segment length.
max_seqlen_kint- Maximum key segment length.
class PackedCrossAttnParams(q_ranges: torch.Tensor | None = None,kv_ranges: torch.Tensor | None = None,cu_seqlens_q: torch.Tensor | None = None,cu_seqlens_kv: torch.Tensor | None = None,max_seqlen_q: int | None = None,max_seqlen_kv: int | None = None)worldfoundry.core.PackedCrossAttnParamsfrom worldfoundry.core import PackedCrossAttnParamsOverview
Optional cumulative ranges for packed cross-attention. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Attributes
q_rangestorch.Tensor | None- Query segment ranges.default:
None kv_rangestorch.Tensor | None- Key/value segment ranges.default:
None cu_seqlens_qtorch.Tensor | None- Cumulative query lengths for varlen kernels.default:
None cu_seqlens_kvtorch.Tensor | None- Cumulative key/value lengths for varlen kernels.default:
None max_seqlen_qint | None- Maximum query segment length.default:
None max_seqlen_kvint | None- Maximum key/value segment length.default:
None
def piecewise_attention(q: torch.Tensor,k: torch.Tensor,v: torch.Tensor,scale: float | None = None,density: float = 0.1,block_size: int = 64,min_sequence_length: int | None = None,strict: bool = False) -> torch.Tensorworldfoundry.core.piecewise_attentionfrom worldfoundry.core import piecewise_attentionOverview
Run PISA block-routed attention or an exact SDPA fallback. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: torch.Tensor.
Source docstring
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.
Parameters
qtorch.Tensorktorch.Tensorvtorch.Tensorscalefloat | None- default:
None densityfloat- default:
0.1 block_sizeint- default:
64 min_sequence_lengthint | None- default:
None strictbool- default:
False
Returns: torch.Tensor
def piecewise_attention_available(device: torch.device | str | None = None) -> boolworldfoundry.core.piecewise_attention_availablefrom worldfoundry.core import piecewise_attention_availableOverview
Return whether the in-tree TMA implementation is eligible on `device. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: bool`.
Parameters
devicetorch.device | str | None- default:
None
Returns: bool
class PositionGetter()worldfoundry.core.PositionGetterfrom worldfoundry.core import PositionGetterOverview
Cache 2D patch-grid positions by grid shape. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Methods
Overview
Public method on this type.
Parameters
batch_sizeintheightintwidthintdevicetorch.device
Returns: Tensor
def probe_attention_backends(device: torch.device | str | int | None = None) -> dict[str, AttentionKernelCapability]worldfoundry.core.probe_attention_backendsfrom worldfoundry.core import probe_attention_backendsOverview
Probe installed attention packages and hardware compatibility. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: dict[str, AttentionKernelCapability].
Source docstring
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.
Parameters
devicetorch.device | str | int | None- default:
None
Returns: dict[str, AttentionKernelCapability]
def resolve_attention_backend(preferred: str | None = None,device: torch.device | str | int | None = None) -> strworldfoundry.core.resolve_attention_backendfrom worldfoundry.core import resolve_attention_backendOverview
Resolve which attention backend to use from arguments, env, and capability probes.
Source docstring
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.
Parameters
preferredstr | None- default:
None devicetorch.device | str | int | None- default:
None
Returns: str
def resolve_transformers_attention_implementation(preferred: str | None = None,device: torch.device | str | int | None = None) -> strworldfoundry.core.resolve_transformers_attention_implementationfrom worldfoundry.core import resolve_transformers_attention_implementationOverview
Resolve a backend name accepted by Transformers model configs. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: str.
Source docstring
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.
Parameters
preferredstr | None- default:
None devicetorch.device | str | int | None- default:
None
Returns: str
def rotary_frequencies(seq_len: int,dim: int,base: float = 10000.0,start_index: int = 0,dtype: Any = None) -> tuple[Any, Any]worldfoundry.core.rotary_frequenciesfrom worldfoundry.core import rotary_frequenciesOverview
Build RoPE cosine/sine tables with shape `(seq_len, dim) using NumPy. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: tuple[Any, Any]`.
Parameters
seq_lenintdimintbasefloat- default:
10000.0 start_indexint- default:
0 dtypeAny- default:
None
Returns: tuple[Any, Any]
class RotaryPositionEmbedding2D(frequency: float = 100.0, scaling_factor: float = 1.0)worldfoundry.core.RotaryPositionEmbedding2Dfrom worldfoundry.core import RotaryPositionEmbedding2DOverview
Apply rotary embeddings to tokens with `(y, x)` patch coordinates. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Parameters
frequencyfloat- default:
100.0 scaling_factorfloat- default:
1.0
Methods
Overview
Public method on this type.
Parameters
tokensTensorpositionsTensor
Returns: Tensor
rotate_half
funcdef rotate_half(value: Any) -> Anyworldfoundry.core.rotate_halffrom worldfoundry.core import rotate_halfOverview
Rotate the last dimension as `[-x2, x1] for RoPE application. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: Any`.
Parameters
valueAny
Returns: Any
def scaled_dot_product_attention(query: Any,key: Any,value: Any,*args: Any,attn_mask: Any = None,dropout_p: float = 0.0,is_causal: bool = False,scale: float | None = None,enable_gqa: bool = False,backend: Literal['math', 'efficient', 'cudnn', 'flash'] | Any | None = None,backends: Any = None) -> Anyworldfoundry.core.scaled_dot_product_attentionfrom worldfoundry.core import scaled_dot_product_attentionOverview
Exact scaled dot-product attention with an explicit backend context. Prefer this when Q/K/V already have split-head shapes; keep dropout_p=0.0 at inference.
Source docstring
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.
Parameters
queryAny- Query tensor shaped `
(..., query_length, head_dim). The common layout is(batch, heads, query_length, head_dim)`. keyAny- Key tensor shaped `
(..., key_length, head_dim)`. valueAny- Value tensor shaped `
(..., key_length, value_dim)`. argsAny- Backward-compatible positional values for `
attn_mask,dropout_p,is_causal, andscale`, in that order. attn_maskAny- Boolean keep-mask or additive attention bias broadcastable to the attention score shape.default:
None dropout_pfloat- Probability applied to attention weights. Pass `
0.0` at inference time; SDPA applies a non-zero value even in eval mode.default:0.0 is_causalbool- Apply a lower-triangular causal mask.default:
False scalefloat | None- Softmax scale. `
Noneuses1 / sqrt(head_dim)`.default:None enable_gqabool- Expand key/value heads when query has a compatible larger head count.default:
False backendLiteral['math', 'efficient', 'cudnn', 'flash'] | Any | None- One requested PyTorch SDPA backend: `
math,efficient,cudnn, orflash`.default:None backendsAny- Ordered backend collection. Takes precedence over `
backend` and is useful when an explicit fallback order is required.default:None
Notes
Prefer this function for already split heads. For flattened `(batch, sequence, hidden) tensors, use flattened_multihead_attention` so reshape and mask normalization stay centralized.
Raises
TypeError- More than four compatibility positional options are given.
ValueError- Grouped-query attention head counts are incompatible.
Returns: Any — Attention values with the query prefix shape and `value.shape[-1]` as the final dimension.
class UlyssesScheduler()worldfoundry.core.UlyssesSchedulerfrom worldfoundry.core import UlyssesSchedulerOverview
Overlap Ulysses all-to-all communication with Q/K/V and cross attention. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Source docstring
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.
Methods
get_attn_and_xattn_with_comm_overlap(get_q_func: Callable,get_k_func: Callable,get_v_func: Callable,kv_cache_func: Callable,core_attn_func: Callable,cross_attn_func: Callable,overlap_degree: int,batch_size: int,cp_size: int,cp_split_sizes: List[int] = None)sourceOverview
Get Q, K, V with communication overlap.
Parameters
get_q_funcCallableget_k_funcCallableget_v_funcCallablekv_cache_funcCallablecore_attn_funcCallablecross_attn_funcCallableoverlap_degreeintbatch_sizeintcp_sizeintcp_split_sizesList[int]- default:
None
get_attn_and_xattn_with_fused_kv_comm(get_q_func: Callable,get_kv_func: Callable,kv_cache_func: Callable,core_attn_func: Callable,cross_attn_func: Callable,overlap_degree: int,batch_size: int,cp_size: int,cp_split_sizes: List[int] = None)sourceOverview
When seq_len is very small, CPU-bound issues are severe.
Parameters
get_q_funcCallableget_kv_funcCallablekv_cache_funcCallablecore_attn_funcCallablecross_attn_funcCallableoverlap_degreeintbatch_sizeintcp_sizeintcp_split_sizesList[int]- default:
None
get_attn_and_xattn_with_fused_qkv_comm(get_qkv_func: Callable,kv_cache_func: Callable,core_attn_func: Callable,cross_attn_func: Callable,overlap_degree: int,batch_size: int,cp_size: int,cp_split_sizes: List[int] = None)sourceOverview
By fusing the communication of q, k, and v together, further optimize CPU-bound issues.
Parameters
get_qkv_funcCallablekv_cache_funcCallablecore_attn_funcCallablecross_attn_funcCallableoverlap_degreeintbatch_sizeintcp_sizeintcp_split_sizesList[int]- default:
None
get_attn_and_xattn_base(query: torch.Tensor,key: torch.Tensor,value: torch.Tensor,core_attn_func: Callable,cross_attn_func: Callable,overlap_degree: int,batch_size: int,cp_size: int,cp_split_sizes: List[int] = None)sourceOverview
Public staticmethod on this type.
Parameters
querytorch.Tensorkeytorch.Tensorvaluetorch.Tensorcore_attn_funcCallablecross_attn_funcCallableoverlap_degreeintbatch_sizeintcp_sizeintcp_split_sizesList[int]- default:
None
def varlen_scaled_dot_product_attention(query: torch.Tensor,key: torch.Tensor,value: torch.Tensor,cu_seqlens_q: torch.Tensor,cu_seqlens_k: torch.Tensor,max_seqlen_q: int | None = None,max_seqlen_k: int | None = None,dropout_p: float = 0.0,softmax_scale: float | None = None,causal: bool = False,version: int | None = None,window_size: tuple[int, int] = (-1, -1),**kwargs: Any) -> torch.Tensorworldfoundry.core.varlen_scaled_dot_product_attentionfrom worldfoundry.core import varlen_scaled_dot_product_attentionOverview
Run in-tree packed attention, or external FlashAttention 2 explicitly. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: torch.Tensor.
Parameters
querytorch.Tensorkeytorch.Tensorvaluetorch.Tensorcu_seqlens_qtorch.Tensorcu_seqlens_ktorch.Tensormax_seqlen_qint | None- default:
None max_seqlen_kint | None- default:
None dropout_pfloat- default:
0.0 softmax_scalefloat | None- default:
None causalbool- default:
False versionint | None- default:
None window_sizetuple[int, int]- default:
(-1, -1) kwargsAny
Returns: torch.Tensor