flashinfer.comm.QkvNormRopeExchangeOverlap

class flashinfer.comm.QkvNormRopeExchangeOverlap(group: ProcessGroup | None = None, *, hidden: int, num_heads: int, head_dim: int = 128, dtype: dtype = torch.bfloat16, layout: str = 'bshd', max_tokens: int, num_slots: int = 2, backend: str = 'auto', device: device | str | int | None = None, fold_barrier: bool | None = None)

Prepared session of the fused QKV projection + norm + RoPE + head exchange on one process group.

Every rank constructs the session collectively (same arguments, same order) and calls it with its own [B, S_local, hidden] activation shard; each rank receives head band [rank * H_local, (rank + 1) * H_local) of the whole sequence.

Parameters:

  • group (torch.distributed.ProcessGroup, optional) – NCCL process group of 2, 4 or 8 ranks on one node (None uses WORLD).

  • hidden (int) – Model width num_heads * head_dim (a multiple of 64).

  • num_heads (int) – Attention heads; must be divisible by the world size.

  • head_dim (int, default 128) – Head dimension (the fused kernels fix 128).

  • dtype (torch.dtype, default torch.bfloat16) – Operand dtype (BF16 only).

  • layout ({"bshd", "fused"}) – "bshd": (q, k, v) each [B, S_global, H_local, head_dim]; "fused": one [B, S_global, H_local, 3 * head_dim] payload.

  • max_tokens (int) – Largest B * S_local of any call; sizes the symmetric receive slots once (collective).

  • num_slots (int, default 2) – Receive slots cycled by consecutive calls; a call’s outputs stay valid for num_slots - 1 further calls.

  • backend ({"auto", "nvlink", "nccl"}) – "nvlink" selects the Cake fused kernels (single-node NVLink mesh, SM90 / SM100 / SM103 / SM107, groups of 2, 4 or 8 ranks (2 or 4 on SM107), torch’s default CUDA symmetric-memory backend) and raises when they cannot be used; "nccl" selects the per-operator path (cuBLAS, rmsnorm, RoPE, UlyssesCommunicator.scatter_qkv_head_chunk); "auto" resolves through resolve_ulysses_backend() and the device architecture.

  • device (torch.device or str or int, optional) – CUDA device of this rank (None: the current device).

  • fold_barrier (bool, optional) – Launch path of the Cake backend. True: the producer program with the folded cross-rank rendezvous – two launches per call (producer, wait kernel); False: the separate barrier kernel before the producer – three launches. None (default) reads the environment variable FLASHINFER_QKV_EXCHANGE_FOLD_BARRIER (0 / 1 / auto) and otherwise takes the architecture default: True on SM90, SM100, SM103 and SM107 (the folded producers are delivered on every architecture). Both paths store identical bits. Must be uniform across the ranks of the group; not accepted with backend="nccl".

Notes

The Cake backend selects its producer build per architecture and world size (flashinfer.jit.cake_qkv_norm_rope_exchange_overlap.PRODUCER_VARIANT_DEFAULT_BY_ARCH_WORLD: SM100 groups of 8 ranks run the bulk_ow256 build – wide operand loads and bulk row stores –, SM100 groups of 4 ranks and SM103 groups of 8 or 4 ranks the ow256 build – wide operand loads –, SM107 groups of 4 ranks the bulk build – bulk row stores –, groups of 2 ranks and SM90 the baseline build). Every build stores identical bits; the wide-load builds require 32-byte-aligned cos_sin and norm-weight tensors (a misaligned view raises and names the opt-out). The environment variable FLASHINFER_QKV_EXCHANGE_PRODUCER_VARIANT (auto / shipped / a delivered build name) overrides the selection for a session, rank-uniformly; the baseline build is delivered beside every variant, and a build that is not delivered for the session’s tile width raises and names the delivered ones. The host plan’s group scaling follows the same keying (cake_qkv_exchange_plan.GROUP_M_SCALE_DEFAULT_BY_ARCH_WORLD: 2 on SM100 groups of 4 and 2 ranks and on SM103, 1 otherwise – SM100 groups of 8 ranks, SM90 and SM107; it applies to grouped rows with hidden >= 6144 and changes the tile order only); FLASHINFER_QKV_EXCHANGE_GROUP_M_SCALE (auto / 1 / 2 / 4) overrides it for a session, rank-uniformly.

The Cake backend also applies per-architecture row-count gates to each call’s host plan (cake_qkv_exchange_plan.SHAPE_GATE_DEFAULT_BY_ARCH; SM100 / SM103: a call with batch * s_local <= 4096 rows publishes one readiness word per destination, a call with >= 16384 rows launches the producer once per destination run; SM90: a call with <= 4096 rows publishes one readiness word per destination and a call with >= 16384 rows runs the persistent producer grid over every SM – reserved_sms_large_m = 16384, no reserved SMs –, the per-destination launch gate off; SM107: every gate off). One (architecture, world size) cell differs from these tables (cake_qkv_exchange_plan.SHAPE_GATE_DEFAULT_BY_ARCH_WORLD): a group of 8 ranks on SM90 publishes one readiness word per destination for every row count (words_capacity_small_m = all). A fourth gate, pair_off_large_m, is off on every architecture: switched on, a 256-column call with at least N rows per rank runs the single-CTA producer instead of the 2-CTA pair tile (tile order unchanged; that producer must be delivered for the architecture – no delivered Blackwell route ships it today, so the session refuses at construction when a 256-column call can reach the gate, max_tokens >= N, naming the gate; SM90 builds no pair tile and refuses the gate). The environment variable FLASHINFER_QKV_EXCHANGE_SHAPE_GATES (auto / off / words_capacity_small_m=N|all,reserved_sms_large_m=N,per_dest_large_m=N,pair_off_large_m=N, 0 = never) overrides them for a session, rank-uniformly (words_capacity_small_m=4096 on SM90 with 8 ranks restores the previous plan). describe() reports the resolved build (producer_variant), scale (group_m_scale) and gates (shape_gates), each with its value, source, environment variable and default, and every prepared row’s verdict (prepared_row_plans, including reserved_sms and part_grids). The first three gates change the host plan only – neither the tile table nor the stored bits; the pair-tile gate changes the producer program and the tile table (no lockstep pads) of the rows it selects; no gate changes the stored bits. The largest 2-rank shape of the evaluation grid (S_global = 131072, hidden = 8192) is not supported on an 80 GB H100: its operands and receive slots exceed the device memory.

__init__(group: ProcessGroup | None = None, *, hidden: int, num_heads: int, head_dim: int = 128, dtype: dtype = torch.bfloat16, layout: str = 'bshd', max_tokens: int, num_slots: int = 2, backend: str = 'auto', device: device | str | int | None = None, fold_barrier: bool | None = None) → None

Methods

__init__([group, head_dim, dtype, layout, ...])

bound_programs(batch, s_local)

Cake backend: program names of every delivered stage of a (batch, s_local) call (prepares the row eagerly); describe()["launch_sequence"] says which of them the launch path launches.

close()

Release the session's resources (collective for the Cake backend).

describe()

last_outputs()

Views of the most recent call's outputs.

poison_receive_slots()

Fill this rank's receive region with a NaN pattern (eager only), so the next call must rewrite it.