mstar.engine.resources.attn.ragged.block_causal

Block-causal cacheless attention through FlashInfer’s paged prefill.

Within each declared span, a token in block b (block_size tokens a block, counted from the span’s start) attends every key of blocks 0..b of that span: causal between blocks, bidirectional within one. Whisper-style encoders run this way so a streaming client sees the same features a whole-clip forward would.

Prefixes of one span overlap, which a ragged kv_indptr (one contiguous key range per query segment) cannot express. The paged kernel can: the packed keys are viewed as pages of one token, and each query block is one paged “request” whose page list is its prefix. No key is copied; the page lists are host-side index arithmetic done at plan time.

Functions

block_causal_layout(seg_lens, block_size)	(qo_indptr, kv_indptr, kv_indices) for spans packed end to end.
max_blocks(max_segments, max_tokens, block_size)	Most query blocks a layout of at most this many spans and tokens has: every span may end on a partial block.
max_prefix_tokens(max_tokens, block_size)	Most page indices such a layout needs: one span holding every token, which maximizes the sum of its blocks' prefixes.

Classes

RaggedBlockCausalWrapper(workspace_buffer, ...)

FlashInfer's paged prefill over a packed, cacheless [total, H, D] layout, planned block-causally per segment; see the module docstring.

class mstar.engine.resources.attn.ragged.block_causal.RaggedBlockCausalWrapper(workspace_buffer, num_qo_heads, num_kv_heads, head_dim, block_size, max_num_segments=None, max_total_tokens=None, device=torch.device('cuda'), accelerator_graph=False, sm_scale=None, q_data_type=torch.bfloat16, backend='auto')

Bases: object

FlashInfer’s paged prefill over a packed, cacheless [total, H, D] layout, planned block-causally per segment; see the module docstring.

In graph mode the block count is padded to the bucket’s ceiling with zero-length blocks, and the index buffer is sized for the worst prefix sum the bucket’s token count allows, so any layout inside the bucket replays through the same captured kernel.

Parameters:

• workspace_buffer (Tensor)

• num_qo_heads (int)

• num_kv_heads (int)

• head_dim (int)

• block_size (int)

• max_num_segments (int | None)

• max_total_tokens (int | None)

• device (device)

• accelerator_graph (bool)

• sm_scale (float | None)

• q_data_type (dtype)

• backend (str)

property num_segments: int

Real (unpadded) segment count from the most recent plan.

plan(seg_lens)

Plan one packed layout: seg_lens[i] tokens in span i, end to end.

Parameters:

seg_lens (list[int])

Return type:

None

run(q, k, v)

[total, H, D] q/k/v packed by the planned spans -> [total, H, D].

As with the ragged wrapper, rows past the planned total are not a valid result; the caller slices.

Parameters:

• q (Tensor)

• k (Tensor)

• v (Tensor)

Return type:

Tensor

mstar.engine.resources.attn.ragged.block_causal.block_causal_layout(seg_lens, block_size)

(qo_indptr, kv_indptr, kv_indices) for spans packed end to end.

One entry per query block. Block i of a span starting at off has queries [off + i*B, off + min((i+1)*B, L)) and keys [off, off + min((i+1)*B, L)).

Parameters:

• seg_lens (list[int])

• block_size (int)

Return type:

tuple[list[int], list[int], Tensor]

mstar.engine.resources.attn.ragged.block_causal.max_blocks(max_segments, max_tokens, block_size)

Most query blocks a layout of at most this many spans and tokens has: every span may end on a partial block.

Parameters:

• max_segments (int)

• max_tokens (int)

• block_size (int)

Return type:

int

mstar.engine.resources.attn.ragged.block_causal.max_prefix_tokens(max_tokens, block_size)

Most page indices such a layout needs: one span holding every token, which maximizes the sum of its blocks’ prefixes.

Parameters:

• max_tokens (int)

• block_size (int)

Return type:

int