mstar.engine.resources.attn.flex

mstar.engine.resources.attn.flex#

Attention over a ring KV cache, through a compiled flex_attention.

The visibility of a ring cannot be expressed as a sequence length: the slots a step may read are scattered across a fixed buffer that is overwritten in place, so there is no prefix to pass a kernel. FlexAttention’s BlockMask names the readable KV blocks directly, which is what makes the ring expressible at all.

The mask this builds is query-uniform and full-blocks-only — every query block sees the same KV blocks, and a visible block is wholly visible — because writes are whole frames of tokens. That is what makes “capacity must be a multiple of 128” a real constraint here rather than a convenience.

Functions

make_block_mask(q_len, kv_len, written)

Build the query-uniform, full-blocks-only BlockMask over written.

Classes

FlexAttentionManager(kv_cache, device, ...)

Attention against a ring KV cache, through flex_attention.
class mstar.engine.resources.attn.flex.FlexAttentionManager(kv_cache, device, dtype, kv_config)[source]#

Bases: AttentionManager

Attention against a ring KV cache, through flex_attention.

Why FlexAttention and not the paged FlashInfer path that already exists here: numerical parity with the reference implementation, which runs flex_attention with a full-block-only BlockMask. A paged kernel changes the accumulation order over the KV blocks, so a port on top of it could never be compared bit-exactly against the reference — and for the model this exists to serve, that comparison is the only check there is: a mask or position bug does not raise, it produces plausible, smoothly drifting video. The FlashInfer alternative was measured at ~9% on the attention kernel, against the 2x that had motivated trying it.

Stateless across steps. The session state is the ring, and the ring belongs to the KV resource this one names in depends_on; nothing here survives a call except the label/layer cursors the base class defines.

Parameters:
attend(q, k, v, visible, *, enable_gqa, layer_idx=None)[source]#

One layer’s attention over the ring.

q is [B, H_q, T, D]; k/v are [B, H_kv, kv_len, D], the whole ring view (history slots plus the scratch frame) as the KV resource returned it. visible is [kv_len] bool, True where that view holds KV this step may read — it is not derivable from a length, which is the entire reason this backend exists. Returns [B, H_q, T, D].

Parameters:
Return type:

Tensor

build_accelerator_graph_buffers(slots, max_bs, max_seq_len)[source]#

Size whatever the captured replays will read.

Called once per runner that captures against this node — the whole forward’s, and one per piecewise region — so it must tolerate repeated calls: grow to the largest shape asked for, never clobber what an earlier call already sized.

Parameters:
Return type:

None

depends_on()[source]#
Return type:

set[str]

property needs_token_visibility: bool#

Whether the ring must construct its legacy token-level mask.

Engine execution stages a block mask in plan and does not need the per-layer row. The fallback remains useful for standalone numerical calls that invoke attend without a resource plan.

plan(step, ctx)[source]#

Stage one local/global visibility mask for this frame and slot.

Parameters:
Return type:

None

property requires_kv_write: bool#

the layer has already written this step’s K/V.

A ring upsert is not the paged backends’ write_kv. It writes the frame and returns the whole ring view to attend against, in one call, because which slots the write makes visible is part of its result — so there is no separate write step for a layer to perform and calling one would commit the frame twice.

A @property, not a class attribute, to keep the base’s read-only contract: AttentionManager.requires_kv_write is a property and DenseAttentionManager overrides it as one. Shadowing it with a plain False would work at runtime and quietly make the attribute writable on this subclass alone.

Type:

False

mstar.engine.resources.attn.flex.flex_attention_masked(query, key, value, score_mod=None, block_mask=None, scale=None, enable_gqa=False, return_lse=False, kernel_options=None, *, return_aux=None)#

This function implements scaled dot product attention with an arbitrary attention score modification function.

This function computes the scaled dot product attention between query, key, and value tensors with a user-defined attention score modification function. The attention score modification function will be applied after the attention scores have been calculated between the query and key tensors. The attention scores are calculated as follows:

The score_mod function should have the following signature:

def score_mod(
    score: Tensor,
    batch: Tensor,
    head: Tensor,
    q_idx: Tensor,
    k_idx: Tensor
) -> Tensor:
Where:

  • score: A scalar tensor representing the attention score, with the same data type and device as the query, key, and value tensors.

  • batch, head, q_idx, k_idx: Scalar tensors indicating the batch index, query head index, query index, and key/value index, respectively. These should have the torch.int data type and be located on the same device as the score tensor.

Parameters:

  • query (Tensor) – Query tensor; shape \((B, Hq, L, E)\). For FP8 dtypes, should be in row-major memory layout for optimal performance.

  • key (Tensor) – Key tensor; shape \((B, Hkv, S, E)\). For FP8 dtypes, should be in row-major memory layout for optimal performance.

  • value (Tensor) – Value tensor; shape \((B, Hkv, S, Ev)\). For FP8 dtypes, should be in column-major memory layout for optimal performance.

  • score_mod (Optional[Callable]) – Function to modify attention scores. By default no score_mod is applied.

  • block_mask (Optional[BlockMask]) – BlockMask object that controls the blocksparsity pattern of the attention.

  • scale (Optional[float]) – Scaling factor applied prior to softmax. If none, the default value is set to \(\frac{1}{\sqrt{E}}\).

  • enable_gqa (bool) – If set to True, enables Grouped Query Attention (GQA) and broadcasts key/value heads to query heads.

  • return_lse (bool) – Whether to return the logsumexp of the attention scores. Default is False. Deprecated: Use return_aux=AuxRequest(lse=True) instead.

  • kernel_options (Optional[FlexKernelOptions]) – Options to control the behavior of the underlying Triton kernels. See FlexKernelOptions for available options and usage examples.

  • return_aux (Optional[AuxRequest]) – Specifies which auxiliary outputs to compute and return. If None, only the attention output is returned. Use AuxRequest(lse=True, max_scores=True) to request both auxiliary outputs.

Returns:

Attention output; shape \((B, Hq, L, Ev)\).

When return_aux is not None:

aux (AuxOutput): Auxiliary outputs with requested fields populated.

When return_aux is None (deprecated paths):

lse (Tensor): Log-sum-exp of attention scores; shape \((B, Hq, L)\). Only returned if return_lse=True.

Return type:

output (Tensor)

Shape legend:

  • \(N: \text{Batch size} ... : \text{Any number of other batch dimensions (optional)}\)

  • \(S: \text{Source sequence length}\)

  • \(L: \text{Target sequence length}\)

  • \(E: \text{Embedding dimension of the query and key}\)

  • \(Ev: \text{Embedding dimension of the value}\)

Warning

torch.nn.attention.flex_attention is a prototype feature in PyTorch. Please look forward to a more stable implementation in a future version of PyTorch. Read more about feature classification at: https://pytorch.org/blog/pytorch-feature-classification-changes/#prototype

mstar.engine.resources.attn.flex.make_block_mask(q_len, kv_len, written)[source]#

Build the query-uniform, full-blocks-only BlockMask over written.

written is [kv_len] bool, True where the ring holds valid KV. Both lengths must be exact multiples of the 128-token sparse block size and written must be block-aligned – both hold because writes are whole frames of 512 (or 128) tokens.

Two properties are load-bearing. Every query block sees the same KV blocks, so the [1, 1, 1, num_kv_blocks] row broadcasts over query blocks and compute_q_blocks=False is safe. And every visible block is full (no partial blocks at all), because frame-granular writes mean “any token in the block is written” and “all of them are” coincide.

Parameters:
Return type:

BlockMask