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.
42 public symbols
apply_nd_rotary_embedding
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]Overview
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[B, S, H, D] layout by default, or [B, H, S, D] when head_first=True`.keytorch.Tensorfreqs_cistorch.Tensor | tuple[torch.Tensor, torch.Tensor](cos, sin)` pair covering the requested sequence window.head_firstbooldefault False[B, H, S, D]` sequence-axis interpretation.start_offsetintdefault 0Returns
(query, key)` tensors with original dtype and shape.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.
apply_rotary_embedding
def apply_rotary_embedding(value: Any,cos: Any,sin: Any,rotary_dim: int | None = None,interleaved: bool = False) -> AnyApply 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 | Nonedefault Noneinterleavedbooldefault FalseReturns
def attention_backend_capability(name: str,device: torch.device | str | int | None = None) -> AttentionKernelCapabilityOverview
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 | Nonedefault NoneReturns
attention_backend_context
def attention_backend_context(backend: Literal['math', 'efficient', 'cudnn', 'flash'] | Any | None = None,backends: Any = None) -> AnyReturn 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 | Nonedefault NonebackendsAnydefault NoneReturns
attention_backend_from_env
def attention_backend_from_env(environ: Mapping[str, str] | None = None) -> strOverview
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] | Nonedefault NoneReturns
attention_backend_info
Return the available generic PyTorch attention backend. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: AttentionBackendInfo.
Returns
attention_backend_report
def attention_backend_report() -> tuple[AttentionKernelCapability, ]Overview
Retrieve capability status of all registered backends for display purposes. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: tuple[AttentionKernelCapability, ...].
Source docstring
Retrieve capability status of all registered backends for display purposes.
The ordering is a fixed display order (fastest candidates first), *not* the automatic dispatch chain: automatic mode uses the in-tree PyTorch SDPA provider only, and every other backend listed here requires an explicit opt-in via `WORLDFOUNDRY_ATTENTION_BACKEND/ WORLDFOUNDRY_ATTENTION_IMPLEMENTATION`.
Returns
attention_dispatch_report
def attention_dispatch_report() -> dict[str, object]Overview
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
attention_forward
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,backend: str | None = None,backend_options: dict | None = None)Overview
Layout-aware attention entry that adapts einops-style QKV packs and optional fused providers through the dispatch policy.
Source docstring
Dispatch QKV Attention across qualified backends for this device, dtype, and shape.
Layouts are einops patterns and are normalized before execution. `auto` only tries *explicit* candidates that are usable on the current hardware (the default list is empty, so the path is SDPA). Unsupported workload signatures are remembered, and in-tree PyTorch SDPA is always the exact fallback.
Parameters
qtorch.Tensorktorch.Tensorvtorch.Tensorq_patterndefault 'b n s d'k_patterndefault 'b n s d'v_patterndefault 'b n s d'out_patterndefault 'b n s d'dimsdefault None(n d)`.attn_maskdefault Nonescaledefault NoneNone` uses the backend default.compatibility_modedefault Falsebackendstr | Nonedefault NoneWORLDFOUNDRY_ATTENTION_BACKEND`.backend_optionsdict | Nonedefault NoneReturns
out_pattern`.Notes
Inspect the choice with `attention_dispatch_report. After changing provider availability in a long-lived process, call clear_attention_dispatch_cache`.
Raises
ModelSpecificAttentionBackendError- A sparse attention system was selected without its required model/checkpoint metadata graph.
RuntimeError- The selected provider failed for a reason other than an unsupported kernel or a missing optional library.
class AttentionBackendInfo(backend: str,uses_torch_sdpa: bool)Resolved attention backend metadata after probing and normalization.
Attributes
backendstruses_torch_sdpaboolclass AttentionKernelCapability(name: str,package: str,available: bool,usable: bool,reason: str = '')Runtime availability metadata for a single attention kernel family. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Attributes
namestrpackagestravailableboolusableboolreasonstrdefault ''BlockKVCache
class 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 | None = None,dtype: torch.dtype = torch.float16)Overview
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.
Phases: - Filling: cache not yet full; tokens are written contiguously; `cached_k() / cached_v() return only the valid prefix. - Steady-state: if adding a chunk would exceed the fixed cache size, the local window rolls left by the overflow amount and the new chunk 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_sizeintdefault 0devicetorch.device | str | Nonedefault Nonedtypetorch.dtypedefault torch.float16Methods
Right edge of the current chunk in the physical cache layout.
Parameters
selfReturns
Build a single-chunk cache pre-filled with the given key and value tensors.
Parameters
kTensorvTensorseq_dimintReturns
Return True if the cache is full (steady-state phase).
Returns
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_idxintReturns
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
kTensorvTensorReturns
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_idxintReturns
Return cached keys for attention (valid prefix in filling phase, full buffer in steady-state).
Returns
Return cached values for attention (valid prefix in filling phase, full buffer in steady-state).
Returns
def clear_attention_dispatch_cache() -> NoneClear workload decisions and runtime failure quarantine. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache). Annotated return type: None.
Returns
class ContextParallelAttention(qkv_format: Literal['bhsd', 'bshd'] = 'bhsd',backend: Literal['cudnn', 'flash'] = 'cudnn',method: Literal['ring', 'ulysses'] = 'ring',convert_to_fp32: bool = True)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']default 'bhsd'"bhsd" or "bshd"`.backendLiteral['cudnn', 'flash']default 'cudnn'"cudnn" or "flash"`.methodLiteral['ring', 'ulysses']default 'ring'"ring" or "ulysses"`.convert_to_fp32booldefault Truecp_post_process
def cp_post_process(cp_size: int,cp_strategy: str,x: torch.Tensor,meta_args: ModelMetaArgs) -> torch.TensorOverview
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_sizeint1 returns x` unchanged.cp_strategystr"cp_ulysses" or "cp_shuffle_overlap"`.xtorch.Tensormeta_argsModelMetaArgscp_pre_process`.Returns
Raises
ValueError- `
cp_strategy` is unknown.
cp_pre_process
def 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)Overview
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_paramsPackedCrossAttnParamscso_communication
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]Overview
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.Tensorcp_world_sizeintcp_split_sizesList[int]comm_typestrdefault None"kv"` additionally reshapes KV heads before exchange.Returns
class CSOHelper(cp_shuffle_num,cp_world_size,cp_split_sizes)Overview
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_sizesMethods
Reshape queries into shuffle chunks and start the first async exchange.
Parameters
queryOverview
Rotate query/output all-to-alls around `fattn` so comm hides compute.
Source docstring
Rotate query/output all-to-alls around `fattn` so comm hides compute.
When `cp_shuffle_num>1`, later iterations pack the previous output with the next query so one collective restores output while fetching the next query. The last iteration issues a dedicated output exchange.
Parameters
fattnqskvdef 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) -> TensorOverview
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 | Nonedefault Nonedropout_pfloatdefault 0.0is_causalbooldefault Falsescalefloat | Nonedefault NonebackendLiteral['math', 'efficient', 'cudnn', 'flash'] | Any | Nonedefault NonebackendsAnydefault NoneReturns
get_1d_rotary_pos_embed
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]Overview
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 | intthetafloatdefault 10000.0use_realbooldefault Falsetheta_rescale_factorfloatdefault 1.0interpolation_factorfloatdefault 1.0Returns
get_meshgrid_nd
def 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.TensorOverview
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]dimintdefault 2devicetorch.device | str | Nonedefault NoneReturns
get_nd_rotary_pos_embed
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]Overview
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]thetafloatdefault 10000.0use_realbooldefault Falsetheta_rescale_factorfloat | list[float]default 1.0interpolation_factorfloat | list[float]default 1.0devicetorch.device | str | Nonedefault NoneReturns
def gpu_supports_flash_attention(device: torch.device | str | int | None = None) -> boolOverview
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 | Nonedefault NoneReturns
class InferenceParams(max_batch_size: int,max_sequence_length: int)Per-step KV and offset state for incremental decoding. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Parameters
max_batch_sizeintmax_sequence_lengthintMethods
Overview
Reorder the batch axis of every cached `(key, value)` pair.
Source docstring
Reorder the batch axis of every cached `(key, value)` pair.
Used when a generate loop drops or permutes live sequences (for example after early EOS). `batch_idx is gathered along axis 1, which is the batch dimension of the Megatron-style cache layout [seq, batch, ...]`.
Parameters
batch_idxReturns
Raises
ValueError- The cache is empty, so there is nothing to swap.
ModelMetaArgs
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)Overview
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_paramsPackedCrossAttnParamsclass NativeAttention(qkv_format: Literal['bhsd', 'bshd'] = 'bhsd',backend: Literal['math', 'efficient', 'cudnn', 'flash', 'cudnn_fp8', 'triton_tma'] = 'cudnn')Module form of Core SDPA that can attach a context-parallel process group.
Parameters
qkv_formatLiteral['bhsd', 'bshd']default 'bhsd'"bhsd" is (B, H, S, D), "bshd" is (B, S, H, D)`.backendLiteral['math', 'efficient', 'cudnn', 'flash', 'cudnn_fp8', 'triton_tma']default 'cudnn'sdpa_kernel`.Methods
Enable or disable context parallelism for ring attention.
Parameters
cp_groupProcessGroup | NoneReturns
Disable context parallelism and release its process-global option lease.
Returns
Return True if context parallelism is active.
Returns
Return the context parallel world size, or 1 if disabled.
Returns
Run context-parallel SDPA (or single-rank SDPA when CP is disabled).
Parameters
queryTensorqkv_format`.keyTensorqkv_format`.valueTensorqkv_format`.Returns
def normalize_attention_backend(value: str | None) -> strOverview
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 | NoneReturns
packed_sequence_attention
def packed_sequence_attention(q: torch.Tensor,k: torch.Tensor,v: torch.Tensor,num_heads: int,compatibility_mode: bool = False,scale = None,backend: str | None = None,backend_options: dict | None = None)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(batch, sequence, heads * head_dim)`.ktorch.Tensorvtorch.Tensornum_headsintcompatibility_modebooldefault Falsescaledefault Nonebackendstr | Nonedefault Nonebackend_optionsdict | Nonedefault NoneReturns
(batch, query_sequence, heads * value_head_dim)`.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)Cumulative ranges and maximum lengths for packed self-attention. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Attributes
q_rangetorch.Tensork_rangetorch.Tensornp_q_rangenp.ndarraynp_k_rangenp.ndarraymax_seqlen_qintmax_seqlen_kintclass 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)Optional cumulative ranges for packed cross-attention. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Attributes
q_rangestorch.Tensor | Nonedefault Nonekv_rangestorch.Tensor | Nonedefault Nonecu_seqlens_qtorch.Tensor | Nonedefault Nonecu_seqlens_kvtorch.Tensor | Nonedefault Nonemax_seqlen_qint | Nonedefault Nonemax_seqlen_kvint | Nonedefault Nonepiecewise_attention
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.TensorOverview
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 | Nonedefault Nonedensityfloatdefault 0.1block_sizeintdefault 64min_sequence_lengthint | Nonedefault Nonestrictbooldefault FalseReturns
def piecewise_attention_available(device: torch.device | str | None = None) -> boolOverview
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 | Nonedefault NoneReturns
PositionGetter
class PositionGetter()Cache 2D patch-grid positions by grid shape. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Methods
Overview
Return `[B, H*W, 2] integer (y, x)` coords, cloned per caller.
Source docstring
Return `[B, H*W, 2] integer (y, x)` coords, cloned per caller.
The clone is required because some ViT blocks write into the returned tensor; sharing the cached storage would corrupt the next forward. Device mismatch rebuilds the table rather than calling `.to` on a potentially large cache entry.
Parameters
batch_sizeintheightintwidthintdevicetorch.deviceReturns
probe_attention_backends
def probe_attention_backends(device: torch.device | str | int | None = None) -> dict[str, AttentionKernelCapability]Overview
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 | Nonedefault NoneReturns
resolve_attention_backend
def resolve_attention_backend(preferred: str | None = None,device: torch.device | str | int | None = None,allow_model_specific: bool = False) -> strOverview
Resolve which attention backend to use from arguments, env, and capability probes.
Source docstring
Resolve the requested backend, falling back to PyTorch SDPA when unusable.
`"auto" deliberately resolves to the in-tree PyTorch SDPA provider only: external FlashAttention/SageAttention/xFormers packages are explicit opt-ins (the no-external-repo contract). When usable external providers exist under auto, an informational log points them out once. If an explicitly requested package is not usable on the active hardware, resolution degrades to "torch"` SDPA and logs a warning with the requested backend, the effective backend, and the capability reason.
Model-specific sparse systems fail by default. A caller that already owns their full metadata/checkpoint graph may pass `allow_model_specific=True`; package and NVIDIA-device availability are still required, and there is no dense fallback for that explicit request.
Parameters
preferredstr | Nonedefault Nonedevicetorch.device | str | int | Nonedefault Noneallow_model_specificbooldefault FalseReturns
def resolve_transformers_attention_implementation(preferred: str | None = None,device: torch.device | str | int | None = None) -> strOverview
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 | Nonedefault Nonedevicetorch.device | str | int | Nonedefault NoneReturns
rotary_frequencies
def rotary_frequencies(seq_len: int,dim: int,base: float = 10000.0,start_index: int = 0,dtype: Any = None) -> tuple[Any, Any]Overview
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_lenintdimintbasefloatdefault 10000.0start_indexintdefault 0dtypeAnydefault NoneReturns
class RotaryPositionEmbedding2D(frequency: float = 100.0,scaling_factor: float = 1.0)Apply rotary embeddings to tokens with `(y, x)` patch coordinates. Belongs to Core attention (SDPA, backends, RoPE, packed sequences, KV cache).
Parameters
frequencyfloatdefault 100.0scaling_factorfloatdefault 1.0Methods
Overview
Apply Y then X RoPE to even-width features.
Source docstring
Apply Y then X RoPE to even-width features.
Failure: odd last-dim or a positions tensor that is not `(B, N, 2)` — those layouts belong to 1-D / 3-D helpers.
Parameters
tokensTensorpositionsTensorReturns
rotate_half
def rotate_half(value: Any) -> AnyRotate 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
valueAnyReturns
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) -> AnyOverview
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_length, head_dim). The common layout is (batch, heads, query_length, head_dim)`.keyAny(..., key_length, head_dim)`.valueAny(..., key_length, value_dim)`.argsAnyattn_mask, dropout_p, is_causal, and scale`, in that order.attn_maskAnydefault Nonedropout_pfloatdefault 0.00.0` at inference time; SDPA applies a non-zero value even in eval mode.is_causalbooldefault Falsescalefloat | Nonedefault NoneNone uses 1 / sqrt(head_dim)`.enable_gqabooldefault FalsebackendLiteral['math', 'efficient', 'cudnn', 'flash'] | Any | Nonedefault Nonemath, efficient, cudnn, or flash`.backendsAnydefault Nonebackend` and is useful when an explicit fallback order is required.Returns
value.shape[-1]` as the final dimension.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.
class UlyssesScheduler()Overview
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)SourceGet 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 Noneget_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)SourceWhen 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 Noneget_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)SourceBy 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 Noneget_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
Run core attention in `overlap_degree` head chunks, then cross-attn.
Source docstring
Run core attention in `overlap_degree` head chunks, then cross-attn.
Query is split so each chunk's output all-to-all can overlap the next chunk's compute. Cross-attention runs after the last core kernel is launched, hiding the final restore collective. `overlap_degree=-1` means one query-head group per KV head (GQA/MQA).
Parameters
querytorch.Tensorkeytorch.Tensorvaluetorch.Tensorcore_attn_funcCallablecross_attn_funcCallableoverlap_degreeintbatch_sizeintcp_sizeintcp_split_sizesList[int]default Nonedef 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.TensorOverview
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 | Nonedefault Nonemax_seqlen_kint | Nonedefault Nonedropout_pfloatdefault 0.0softmax_scalefloat | Nonedefault Nonecausalbooldefault Falseversionint | Nonedefault Nonewindow_sizetuple[int, int]default (-1, -1)kwargsAny