Core 注意力

精确 SDPA、布局分发、后端策略、RoPE、packed sequence、上下文并行和 block KV cache。

最后更新: 2026年10月7日

本页内容

选择最窄的入口一个可以在 CPU 或 GPU 上运行的例子 Block KV cache 的调用顺序完整参考

WorldFoundry 把投影层和模型语义留在模型代码中，由 Core 负责那些经常被重复实现、也容易出错的机械工作：多头布局转换、精确后端选择、mask 处理、RoPE 应用、packed sequence 范围、上下文并行通信和滚动 KV 状态。

选择最窄的入口

当 Q、K、V 已经是标准的拆头形状时，使用 scaled_dot_product_attention。它遵循 PyTorch SDPA 契约，同时增加显式后端上下文、兼容的 GQA 扩展以及精确的 matmul fallback。

当张量是 (batch, sequence, hidden)、只缺少 head 拆分时，使用 flattened_multihead_attention，让拆头和合头留在一个公共位置。当模型接入有 einops 风格的 QKV 布局，或者需要通过策略选择可选 provider 时，使用 attention_forward。传入 mask 或设置 compatibility_mode=True 会有意回到 PyTorch 路径。

NativeAttention 把 SDPA 封装成 module，也可以绑定上下文并行 group。ContextParallelAttention、UlyssesScheduler 和 CSOHelper 是更底层的分布式机制，只有在模型 runtime 同时掌握对应 split metadata 时才应该直接使用。

一个可以在 CPU 或 GPU 上运行的例子

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",  # 本例显式选择确定的 provider
)
assert output.shape == (2, 8, 32, 96)

推理时要明确传 dropout_p=0.0；PyTorch SDPA 不会从 module.eval() 自动推导这个值。布尔 mask 表示哪些 score 可以保留，浮点 mask 则会直接加到 score 矩阵上。

Block KV cache 的调用顺序

BlockKVCache 不是可以随意更新的字典。每个 chunk 都必须遵循 before_update → update → cached_k/cached_v → after_update。重复当前 chunk_idx 会覆盖同一个逻辑 chunk；递增一会追加或滚动本地窗口；跳过编号会报错。

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、4、滚动后的 4 个 token
    cache.after_update(chunk_idx)

sink_size 保留永不淘汰的前缀，window_size 描述滚动区域。两者之和必须等于 cache 的序列维长度，并且能够被 chunk_size 整除。

完整参考

以下为该类别的生成签名。可用本页符号索引跳转；源码链接指向各惰性导出背后的具体实现。

42 个公开符号

函数

`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]

worldfoundry.corefrom worldfoundry.core import apply_nd_rotary_embedding

简介

apply_nd_rotary_embedding — Apply precomputed n-D rotary frequencies to query and key tensors. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：tuple[torch.Tensor, torch.Tensor]。

参数

querytorch.Tensor

Query tensor in `[B, S, H, D] layout by default, or [B, H, S, D] when head_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默认值 False

Select `[B, H, S, D]` sequence-axis interpretation.

start_offsetint默认值 0

First frequency position, used with an existing KV prefix.

返回值

tuple[torch.Tensor, torch.Tensor]Rotated `(query, key)` tensors with original dtype and shape.

说明

The final head width must be compatible with complex pairs. Build matching structured frequencies with `get_nd_rotary_pos_embed`.

异常

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) -> Any

worldfoundry.corefrom worldfoundry.core import apply_rotary_embedding

apply_rotary_embedding — Apply RoPE to the leading `rotary_dim features of the last dimension. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：Any`。

参数

valueAny

cosAny

sinAny

rotary_dimint | None默认值 None

interleavedbool默认值 False

返回值

Any

函数

`attention_backend_capability`

源码

def attention_backend_capability(name: str,device: torch.device | str | int | None = None) -> AttentionKernelCapability

worldfoundry.corefrom worldfoundry.core import attention_backend_capability

简介

attention_backend_capability — Retrieve runtime capability metadata for a single normalized backend. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：AttentionKernelCapability。

参数

namestr

devicetorch.device | str | int | None默认值 None

返回值

AttentionKernelCapability

函数

`attention_backend_context`

源码

def attention_backend_context(backend: Literal['math', 'efficient', 'cudnn', 'flash'] | Any | None = None,backends: Any = None) -> Any

worldfoundry.corefrom worldfoundry.core import attention_backend_context

attention_backend_context — Return a context manager that selects PyTorch SDPA backends via core. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：Any。

参数

backendLiteral['math', 'efficient', 'cudnn', 'flash'] | Any | None默认值 None

backendsAny默认值 None

返回值

Any

函数

`attention_backend_from_env`

源码

def attention_backend_from_env(environ: Mapping[str, str] | None = None) -> str

worldfoundry.corefrom worldfoundry.core import attention_backend_from_env

简介

attention_backend_from_env — Read and resolve the canonical attention backend requested by the environment. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：str。

源码 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.

参数

environMapping[str, str] | None默认值 None

返回值

str

函数

`attention_backend_info`

源码

def attention_backend_info() -> AttentionBackendInfo

worldfoundry.corefrom worldfoundry.core import attention_backend_info

attention_backend_info — Return the available generic PyTorch attention backend. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：AttentionBackendInfo。

返回值

AttentionBackendInfo

函数

`attention_backend_report`

源码

def attention_backend_report() -> tuple[AttentionKernelCapability, ]

worldfoundry.corefrom worldfoundry.core import attention_backend_report

简介

attention_backend_report — Retrieve capability status of all registered backends for display purposes. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：tuple[AttentionKernelCapability, ...]。

源码 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`.

返回值

tuple[AttentionKernelCapability, ...]

函数

`attention_dispatch_report`

源码

def attention_dispatch_report() -> dict[str, object]

worldfoundry.core.attentionfrom worldfoundry.core.attention import attention_dispatch_report

attention_dispatch_report — Return lightweight operator-selection and quarantine state. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：dict[str, object]。

返回值

dict[str, object]

函数

`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)

worldfoundry.corefrom worldfoundry.core import attention_forward

简介

面向布局的注意力入口：适配 einops 风格 QKV，并经分发策略选择可选融合算子。

源码 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.

参数

qtorch.Tensor

ktorch.Tensor

vtorch.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'

Requested output layout.

dims默认值 None

Named dimensions needed to expand grouped terms such as `(n d)`.

attn_mask默认值 None

Boolean or additive mask. A mask forces SDPA because fused kernels do not share one mask contract.

scale默认值 None

Softmax scale; `None` uses the backend default.

compatibility_mode默认值 False

Skip optional providers and run PyTorch SDPA.

backendstr | None默认值 None

Optional request-scoped provider. When omitted, retain the process default resolved from `WORLDFOUNDRY_ATTENTION_BACKEND`.

backend_optionsdict | None默认值 None

返回值

Attention output rearranged into `out_pattern`.

说明

Inspect the choice with `attention_dispatch_report. After changing provider availability in a long-lived process, call clear_attention_dispatch_cache`.

异常

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.

类

`AttentionBackendInfo`

源码

class AttentionBackendInfo(backend: str,uses_torch_sdpa: bool)

worldfoundry.corefrom worldfoundry.core import AttentionBackendInfo

完成探测与归一化后的注意力后端元数据。

属性

backendstr

uses_torch_sdpabool

类

`AttentionKernelCapability`

源码

class AttentionKernelCapability(name: str,package: str,available: bool,usable: bool,reason: str = '')

worldfoundry.corefrom worldfoundry.core import AttentionKernelCapability

AttentionKernelCapability — Runtime availability metadata for a single attention kernel family. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。

属性

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.

类

`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)

worldfoundry.core.attentionfrom worldfoundry.core.attention import BlockKVCache

简介

分块滚动 KV cache，必须遵循 before_update → update → after_update；不要跳过 chunk 下标。

源码 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.

属性

k_shapetuple[int, ...]

v_shapetuple[int, ...]

seq_dimint

chunk_sizeint

window_sizeint

sink_sizeint默认值 0

devicetorch.device | str | None默认值 None

dtypetorch.dtype默认值 torch.float16

方法

propsize -> int源码

size — Number of valid cached tokens visible to attention.

参数

self

返回值

int

propwrite_end -> int源码

write_end — Right edge of the current chunk in the physical cache layout.

参数

self

返回值

int

cmethfrom_tensor(k: Tensor,v: Tensor,seq_dim: int) -> Self源码

from_tensor — Build a single-chunk cache pre-filled with the given key and value tensors.

参数

kTensor

vTensor

seq_dimint

返回值

Self

methis_steady_state() -> bool源码

is_steady_state — Return True if the cache is full (steady-state phase).

返回值

bool

methbefore_update(chunk_idx: int) -> None源码

简介

before_update — Prepare the cache before writing new tokens.

源码 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.

参数

chunk_idxint

Chunk index of the new chunk in the full sequence.

返回值

None

methupdate(k: Tensor,v: Tensor) -> None源码

简介

update — Write the new chunk's keys and values into the cache.

源码 docstring

Write the new chunk's keys and values into the cache.

Must be called after `before_update() and before after_update()`.

参数

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.

返回值

None

methafter_update(chunk_idx: int) -> None源码

简介

after_update — Finalize bookkeeping after writing new tokens.

源码 docstring

Finalize bookkeeping after writing new tokens.

Updates `_prev_chunk_idx and, in filling phase, _n_cached`.

参数

chunk_idxint

The index of the new chunk in the full sequence.

返回值

None

methcached_k() -> Tensor源码

cached_k — Return cached keys for attention (valid prefix in filling phase, full buffer in steady-state).

返回值

Tensor

methcached_v() -> Tensor源码

cached_v — Return cached values for attention (valid prefix in filling phase, full buffer in steady-state).

返回值

Tensor

methreset() -> None源码

reset — Reset the cache to its initial empty state.

返回值

None

函数

`clear_attention_dispatch_cache`

源码

def clear_attention_dispatch_cache() -> None

worldfoundry.core.attentionfrom worldfoundry.core.attention import clear_attention_dispatch_cache

clear_attention_dispatch_cache — Clear workload decisions and runtime failure quarantine. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：None。

返回值

None

类

`ContextParallelAttention`

源码

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.attentionfrom worldfoundry.core.attention import ContextParallelAttention

ContextParallelAttention — Context-parallel attention with selectable method and SDPA backend. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。

参数

qkv_formatLiteral['bhsd', 'bshd']默认值 'bhsd'

Layout of the QKV tensors; `"bhsd" or "bshd"`.

backendLiteral['cudnn', 'flash']默认值 'cudnn'

SDPA backend; `"cudnn" or "flash"`.

methodLiteral['ring', 'ulysses']默认值 'ring'

Context-parallelism strategy; `"ring" or "ulysses"`.

convert_to_fp32bool默认值 True

Promote LSE accumulators to fp32 during ring merges.

函数

`cp_post_process`

源码

def cp_post_process(cp_size: int,cp_strategy: str,x: torch.Tensor,meta_args: ModelMetaArgs) -> torch.Tensor

worldfoundry.corefrom worldfoundry.core import cp_post_process

cp_post_process — Gather context-parallel output back to the original sequence layout. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：torch.Tensor。

参数

cp_sizeint

Context-parallel world size. `1 returns x` 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`.

返回值

torch.TensorGlobally gathered tensor with shuffle-overlap padding removed.

异常

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)

worldfoundry.corefrom worldfoundry.core import cp_pre_process

简介

cp_pre_process — This function is used to handle context parallel behavior, split input tensors into multiple parts and scatter them to different GPUs. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。

源码 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

参数

cp_sizeint

cp_strategystr

xtorch.Tensor

condition_maptorch.Tensor

ropetorch.Tensor

xattn_mask_for_cuda_graphUnion[torch.Tensor, None]

ardf_metadict

core_attn_paramsPackedCoreAttnParams

cross_attn_paramsPackedCrossAttnParams

函数

`cso_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]

worldfoundry.corefrom worldfoundry.core import cso_communication

简介

cso_communication — Launch one context-shuffle-overlap all-to-all operation. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：Tuple[torch.Tensor, torch.distributed.Work]。

参数

inputtorch.Tensor

Local sequence/head tensor.

cp_world_sizeint

Context-parallel group size.

cp_split_sizesList[int]

Per-rank split sizes.

comm_typestr默认值 None

`"kv"` additionally reshapes KV heads before exchange.

返回值

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`

源码

class CSOHelper(cp_shuffle_num,cp_world_size,cp_split_sizes)

worldfoundry.corefrom worldfoundry.core import CSOHelper

简介

CSOHelper — Pipeline chunked query exchange with attention computation. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。

源码 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.

参数

cp_shuffle_num

cp_world_size

cp_split_sizes

方法

methsplit_query_for_overlap(query)源码

split_query_for_overlap — Reshape queries into shuffle chunks and start the first async exchange.

参数

query

methoverlap(fattn,qs,k,v)源码

简介

overlap — Rotate query/output all-to-alls around `fattn` so comm hides compute.

源码 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.

参数

fattn

qs

k

v

函数

`flattened_multihead_attention`

源码

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) -> Tensor

worldfoundry.corefrom worldfoundry.core import flattened_multihead_attention

简介

面向 (batch, sequence, hidden) 张量的便捷封装，只需头数即可在 SDPA 前后完成拆分/合并。

源码 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.

参数

queryTensor

keyTensor

valueTensor

num_headsint

attn_maskTensor | None默认值 None

dropout_pfloat默认值 0.0

is_causalbool默认值 False

scalefloat | None默认值 None

backendLiteral['math', 'efficient', 'cudnn', 'flash'] | Any | None默认值 None

backendsAny默认值 None

返回值

Tensor

函数

`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]

worldfoundry.corefrom worldfoundry.core import get_1d_rotary_pos_embed

简介

get_1d_rotary_pos_embed — Build one-dimensional rotary position frequencies. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：torch.Tensor | tuple[torch.Tensor, torch.Tensor]。

参数

dimint

postorch.Tensor | int

thetafloat默认值 10000.0

use_realbool默认值 False

theta_rescale_factorfloat默认值 1.0

interpolation_factorfloat默认值 1.0

返回值

torch.Tensor | tuple[torch.Tensor, torch.Tensor]

函数

`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.Tensor

worldfoundry.corefrom worldfoundry.core import get_meshgrid_nd

get_meshgrid_nd — Build an n-D meshgrid with PyTorch `linspace(endpoint=False) semantics. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：torch.Tensor`。

参数

startint | tuple[int, ...] | list[int]

argsint | tuple[int, ...] | list[int]

dimint默认值 2

devicetorch.device | str | None默认值 None

返回值

torch.Tensor

函数

`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]

worldfoundry.corefrom worldfoundry.core import get_nd_rotary_pos_embed

简介

get_nd_rotary_pos_embed — Build n-D RoPE frequencies for tokens with structured grid coordinates. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：torch.Tensor | tuple[torch.Tensor, torch.Tensor]。

参数

rope_dim_listlist[int]

startint | tuple[int, ...] | list[int]

argsint | tuple[int, ...] | list[int]

thetafloat默认值 10000.0

use_realbool默认值 False

theta_rescale_factorfloat | list[float]默认值 1.0

interpolation_factorfloat | list[float]默认值 1.0

devicetorch.device | str | None默认值 None

返回值

torch.Tensor | tuple[torch.Tensor, torch.Tensor]

函数

`gpu_supports_flash_attention`

源码

def gpu_supports_flash_attention(device: torch.device | str | int | None = None) -> bool

worldfoundry.corefrom worldfoundry.core import gpu_supports_flash_attention

简介

gpu_supports_flash_attention — Determine if the active CUDA device is architecturally capable of running FlashAttention. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：bool。

源码 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.

参数

devicetorch.device | str | int | None默认值 None

返回值

bool

类

`InferenceParams`

源码

class InferenceParams(max_batch_size: int,max_sequence_length: int)

worldfoundry.corefrom worldfoundry.core import InferenceParams

InferenceParams — Per-step KV and offset state for incremental decoding. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。

参数

max_batch_sizeint

max_sequence_lengthint

方法

methswap_key_value_dict(batch_idx) -> None源码

简介

swap_key_value_dict — Reorder the batch axis of every cached `(key, value)` pair.

源码 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, ...]`.

参数

batch_idx

Index tensor/list whose length matches the current cache batch size.

返回值

None

异常

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)

worldfoundry.corefrom worldfoundry.core import ModelMetaArgs

简介

ModelMetaArgs — Context-parallel layout metadata carried between pre/post processing. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。

源码 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.

属性

Hint

Wint

cp_pad_sizeint

cp_split_sizesList[int]

slice_pointint

denoising_range_numint

range_numint

extract_prefix_video_featurebool

fwd_extra_1st_chunkbool

distill_nearly_clean_chunkbool

clip_token_numsint

enable_cuda_graphbool

core_attn_paramsPackedCoreAttnParams

cross_attn_paramsPackedCrossAttnParams

类

`NativeAttention`

源码

class NativeAttention(qkv_format: Literal['bhsd', 'bshd'] = 'bhsd',backend: Literal['math', 'efficient', 'cudnn', 'flash', 'cudnn_fp8', 'triton_tma'] = 'cudnn')

worldfoundry.core.attentionfrom worldfoundry.core.attention import NativeAttention

Core SDPA 的模块形态，可挂接 context-parallel 进程组。

参数

qkv_formatLiteral['bhsd', 'bshd']默认值 'bhsd'

Layout of the QKV tensors; `"bhsd" is (B, H, S, D), "bshd" is (B, S, H, D)`.

backendLiteral['math', 'efficient', 'cudnn', 'flash', 'cudnn_fp8', 'triton_tma']默认值 'cudnn'

SDPA backend selected via `sdpa_kernel`.

方法

methset_context_parallel_group(cp_group: ProcessGroup | None) -> None源码

set_context_parallel_group — Enable or disable context parallelism for ring attention.

参数

cp_groupProcessGroup | None

Process group for context parallel; use None to disable.

返回值

None

methclear_context_parallel_group() -> None源码

clear_context_parallel_group — Disable context parallelism and release its process-global option lease.

返回值

None

methis_context_parallel_enabled() -> bool源码

is_context_parallel_enabled — Return True if context parallelism is active.

返回值

bool

methcontext_parallel_size() -> int源码

context_parallel_size — Return the context parallel world size, or 1 if disabled.

返回值

int

methforward(query: Tensor,key: Tensor,value: Tensor) -> Tensor源码

forward — Run context-parallel SDPA (or single-rank SDPA when CP is disabled).

参数

queryTensor

Query tensor in configured `qkv_format`.

keyTensor

Key tensor in configured `qkv_format`.

valueTensor

Value tensor in configured `qkv_format`.

返回值

TensorAttention output in the same format as inputs.

函数

`normalize_attention_backend`

源码

def normalize_attention_backend(value: str | None) -> str

worldfoundry.corefrom worldfoundry.core import normalize_attention_backend

简介

normalize_attention_backend — Normalize colloquial or varied attention backend names into standard keys. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：str。

源码 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.

参数

valuestr | None

返回值

str

函数

`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)

worldfoundry.core.attentionfrom worldfoundry.core.attention import packed_sequence_attention

packed_sequence_attention — Apply the shared dispatcher to flattened packed-sequence Q/K/V. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。

参数

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_modebool默认值 False

Force the exact PyTorch SDPA path.

scale默认值 None

Optional attention softmax scale.

backendstr | None默认值 None

Optional request-scoped provider name.

backend_optionsdict | None默认值 None

返回值

Tensor shaped `(batch, query_sequence, heads * value_head_dim)`.

说明

This adapter is appropriate when the model already packs heads into the hidden dimension. Use `attention_forward` directly for custom layouts or explicit masks.

类

`PackedCoreAttnParams`

源码

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.corefrom worldfoundry.core import PackedCoreAttnParams

PackedCoreAttnParams — Cumulative ranges and maximum lengths for packed self-attention. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。

属性

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.

类

`PackedCrossAttnParams`

源码

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.corefrom worldfoundry.core import PackedCrossAttnParams

PackedCrossAttnParams — Optional cumulative ranges for packed cross-attention. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。

属性

q_rangestorch.Tensor | None默认值 None

Query segment ranges.

kv_rangestorch.Tensor | None默认值 None

Key/value segment ranges.

cu_seqlens_qtorch.Tensor | None默认值 None

Cumulative query lengths for varlen kernels.

cu_seqlens_kvtorch.Tensor | None默认值 None

Cumulative key/value lengths for varlen kernels.

max_seqlen_qint | None默认值 None

Maximum query segment length.

max_seqlen_kvint | None默认值 None

Maximum key/value segment length.

函数

`piecewise_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.Tensor

worldfoundry.corefrom worldfoundry.core import piecewise_attention

简介

piecewise_attention — Run PISA block-routed attention or an exact SDPA fallback. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：torch.Tensor。

源码 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.

参数

qtorch.Tensor

ktorch.Tensor

vtorch.Tensor

scalefloat | None默认值 None

densityfloat默认值 0.1

block_sizeint默认值 64

min_sequence_lengthint | None默认值 None

strictbool默认值 False

返回值

torch.Tensor

函数

`piecewise_attention_available`

源码

def piecewise_attention_available(device: torch.device | str | None = None) -> bool

worldfoundry.corefrom worldfoundry.core import piecewise_attention_available

piecewise_attention_available — Return whether the in-tree TMA implementation is eligible on `device. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：bool`。

参数

devicetorch.device | str | None默认值 None

返回值

bool

类

`PositionGetter`

源码

class PositionGetter()

worldfoundry.corefrom worldfoundry.core import PositionGetter

PositionGetter — Cache 2D patch-grid positions by grid shape. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。

方法

meth__call__(batch_size: int,height: int,width: int,device: torch.device) -> Tensor源码

简介

__call__ — Return `[B, H*W, 2] integer (y, x)` coords, cloned per caller.

源码 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.

参数

batch_sizeint

heightint

widthint

devicetorch.device

返回值

Tensor

函数

`probe_attention_backends`

源码

def probe_attention_backends(device: torch.device | str | int | None = None) -> dict[str, AttentionKernelCapability]

worldfoundry.corefrom worldfoundry.core import probe_attention_backends

简介

probe_attention_backends — Probe installed attention packages and hardware compatibility. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：dict[str, AttentionKernelCapability]。

源码 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.

参数

devicetorch.device | str | int | None默认值 None

返回值

dict[str, AttentionKernelCapability]

函数

`resolve_attention_backend`

源码

def resolve_attention_backend(preferred: str | None = None,device: torch.device | str | int | None = None,allow_model_specific: bool = False) -> str

worldfoundry.corefrom worldfoundry.core import resolve_attention_backend

简介

综合参数、环境变量与能力探测，解析应使用的注意力后端。

源码 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.

参数

preferredstr | None默认值 None

devicetorch.device | str | int | None默认值 None

allow_model_specificbool默认值 False

返回值

str

函数

`resolve_transformers_attention_implementation`

源码

def resolve_transformers_attention_implementation(preferred: str | None = None,device: torch.device | str | int | None = None) -> str

worldfoundry.corefrom worldfoundry.core import resolve_transformers_attention_implementation

简介

resolve_transformers_attention_implementation — Resolve a backend name accepted by Transformers model configs. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：str。

源码 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.

参数

preferredstr | None默认值 None

devicetorch.device | str | int | None默认值 None

返回值

str

函数

`rotary_frequencies`

源码

def rotary_frequencies(seq_len: int,dim: int,base: float = 10000.0,start_index: int = 0,dtype: Any = None) -> tuple[Any, Any]

worldfoundry.corefrom worldfoundry.core import rotary_frequencies

rotary_frequencies — Build RoPE cosine/sine tables with shape `(seq_len, dim) using NumPy. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：tuple[Any, Any]`。

参数

seq_lenint

dimint

basefloat默认值 10000.0

start_indexint默认值 0

dtypeAny默认值 None

返回值

tuple[Any, Any]

类

`RotaryPositionEmbedding2D`

源码

class RotaryPositionEmbedding2D(frequency: float = 100.0,scaling_factor: float = 1.0)

worldfoundry.corefrom worldfoundry.core import RotaryPositionEmbedding2D

RotaryPositionEmbedding2D — Apply rotary embeddings to tokens with `(y, x)` patch coordinates. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。

参数

frequencyfloat默认值 100.0

scaling_factorfloat默认值 1.0

方法

methforward(tokens: Tensor,positions: Tensor) -> Tensor源码

简介

forward — Apply Y then X RoPE to even-width features.

源码 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.

参数

tokensTensor

positionsTensor

返回值

Tensor

函数

`rotate_half`

源码

def rotate_half(value: Any) -> Any

worldfoundry.corefrom worldfoundry.core import rotate_half

rotate_half — Rotate the last dimension as `[-x2, x1] for RoPE application. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：Any`。

参数

valueAny

返回值

Any

函数

`scaled_dot_product_attention`

源码

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) -> Any

worldfoundry.corefrom worldfoundry.core import scaled_dot_product_attention

简介

带显式后端上下文的精确 SDPA。当 Q/K/V 已是标准分头形状时优先使用；推理时保持 dropout_p=0.0。

源码 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.

参数

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, and scale`, in that order.

attn_maskAny默认值 None

Boolean keep-mask or additive attention bias broadcastable to the attention score shape.

dropout_pfloat默认值 0.0

Probability applied to attention weights. Pass `0.0` at inference time; SDPA applies a non-zero value even in eval mode.

is_causalbool默认值 False

Apply a lower-triangular causal mask.

scalefloat | None默认值 None

Softmax scale. `None uses 1 / sqrt(head_dim)`.

enable_gqabool默认值 False

Expand key/value heads when query has a compatible larger head count.

backendLiteral['math', 'efficient', 'cudnn', 'flash'] | Any | None默认值 None

One requested PyTorch SDPA backend: `math, efficient, cudnn, or flash`.

backendsAny默认值 None

Ordered backend collection. Takes precedence over `backend` and is useful when an explicit fallback order is required.

返回值

AnyAttention values with the query prefix shape and `value.shape[-1]` as the final dimension.

说明

Prefer this function for already split heads. For flattened `(batch, sequence, hidden) tensors, use flattened_multihead_attention` so reshape and mask normalization stay centralized.

异常

TypeError: More than four compatibility positional options are given.
ValueError: Grouped-query attention head counts are incompatible.

类

`UlyssesScheduler`

源码

class UlyssesScheduler()

worldfoundry.corefrom worldfoundry.core import UlyssesScheduler

简介

UlyssesScheduler — Overlap Ulysses all-to-all communication with Q/K/V and cross attention. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。

源码 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.

方法

smeth

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)

源码

get_attn_and_xattn_with_comm_overlap — Get Q, K, V with communication overlap.

参数

get_q_funcCallable

get_k_funcCallable

get_v_funcCallable

kv_cache_funcCallable

core_attn_funcCallable

cross_attn_funcCallable

overlap_degreeint

batch_sizeint

cp_sizeint

cp_split_sizesList[int]默认值 None

smeth

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)

源码

get_attn_and_xattn_with_fused_kv_comm — When seq_len is very small, CPU-bound issues are severe.

参数

get_q_funcCallable

get_kv_funcCallable

kv_cache_funcCallable

core_attn_funcCallable

cross_attn_funcCallable

overlap_degreeint

batch_sizeint

cp_sizeint

cp_split_sizesList[int]默认值 None

meth

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)

源码

get_attn_and_xattn_with_fused_qkv_comm — By fusing the communication of q, k, and v together, further optimize CPU-bound issues.

参数

get_qkv_funcCallable

kv_cache_funcCallable

core_attn_funcCallable

cross_attn_funcCallable

overlap_degreeint

batch_sizeint

cp_sizeint

cp_split_sizesList[int]默认值 None

smeth

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)

源码

简介

get_attn_and_xattn_base — Run core attention in `overlap_degree` head chunks, then cross-attn.

源码 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).

参数

querytorch.Tensor

keytorch.Tensor

valuetorch.Tensor

core_attn_funcCallable

cross_attn_funcCallable

overlap_degreeint

batch_sizeint

cp_sizeint

cp_split_sizesList[int]默认值 None

函数

`varlen_scaled_dot_product_attention`

源码

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.Tensor

worldfoundry.corefrom worldfoundry.core import varlen_scaled_dot_product_attention

简介

varlen_scaled_dot_product_attention — Run in-tree packed attention, or external FlashAttention 2 explicitly. 属于 Core 注意力（SDPA、后端、RoPE、packed sequence、KV cache）。标注返回类型：torch.Tensor。

参数

querytorch.Tensor

keytorch.Tensor

valuetorch.Tensor

cu_seqlens_qtorch.Tensor

cu_seqlens_ktorch.Tensor

max_seqlen_qint | None默认值 None

max_seqlen_kint | None默认值 None

dropout_pfloat默认值 0.0

softmax_scalefloat | None默认值 None

causalbool默认值 False

versionint | None默认值 None

window_sizetuple[int, int]默认值 (-1, -1)

kwargsAny

返回值

torch.Tensor