mstar.engine.resources.attn.ragged.wrappers

mstar.engine.resources.attn.ragged.wrappers#

Cacheless (ragged) prefill wrappers: queries and keys packed by segment into [total_tokens, H, D] tensors, laid out by cu_seqlens.

FlashInfer wrappers live here; FA3 is in dense_wrappers.py.

Functions

padded_head_dim(head_dim)

Smallest FlashInfer-supported head dim >= head_dim.

Classes

FlashInferPrefixedPrefillWrapper(...[, ...])

Prefixed attention through FlashInfer.

FlashInferRaggedCrossPrefillWrapper(...[, ...])

Ragged cross-attention through FlashInfer.

FlashInferRaggedPrefillWrapper(...[, ...])

Ragged self-attention through FlashInfer.

PrefixedPrefillWrapper(num_qo_heads, ...[, ...])

Query segment i attends key segment i, which holds a frozen prefix followed by the query's own tokens (the KV-backed dense attention backend's layout).

RaggedCrossPrefillWrapper(num_qo_heads, ...)

Varlen cross-attention between two packed layouts: query segment i attends key segment i (one pair per request).

RaggedPrefillWrapper(num_qo_heads, ...[, ...])

Varlen self-attention: each segment attends within itself.
class mstar.engine.resources.attn.ragged.wrappers.FlashInferPrefixedPrefillWrapper(workspace_buffer, num_qo_heads, num_kv_heads, head_dim, *, kv_layout='NHD', backend='auto', **kwargs)[source]#

Bases: _FlashInferRaggedPrefillBase, PrefixedPrefillWrapper

Prefixed attention through FlashInfer.

Parameters:

  • workspace_buffer (Tensor)

  • num_qo_heads (int)

  • num_kv_heads (int)

  • head_dim (int)

  • kv_layout (str)

  • backend (str)

class mstar.engine.resources.attn.ragged.wrappers.FlashInferRaggedCrossPrefillWrapper(workspace_buffer, num_qo_heads, num_kv_heads, head_dim, *, kv_layout='NHD', backend='auto', **kwargs)[source]#

Bases: _FlashInferRaggedPrefillBase, RaggedCrossPrefillWrapper

Ragged cross-attention through FlashInfer.

Parameters:

  • workspace_buffer (Tensor)

  • num_qo_heads (int)

  • num_kv_heads (int)

  • head_dim (int)

  • kv_layout (str)

  • backend (str)

class mstar.engine.resources.attn.ragged.wrappers.FlashInferRaggedPrefillWrapper(workspace_buffer, num_qo_heads, num_kv_heads, head_dim, *, kv_layout='NHD', backend='auto', **kwargs)[source]#

Bases: _FlashInferRaggedPrefillBase, RaggedPrefillWrapper

Ragged self-attention through FlashInfer.

Parameters:

  • workspace_buffer (Tensor)

  • num_qo_heads (int)

  • num_kv_heads (int)

  • head_dim (int)

  • kv_layout (str)

  • backend (str)

class mstar.engine.resources.attn.ragged.wrappers.PrefixedPrefillWrapper(num_qo_heads, num_kv_heads, head_dim, max_num_segments=None, max_total_tokens=None, max_total_kv_tokens=None, device=torch.device('cuda'), accelerator_graph=False, sm_scale=None, q_data_type=torch.bfloat16)[source]#

Bases: _RaggedPrefillBase

Query segment i attends key segment i, which holds a frozen prefix followed by the query’s own tokens (the KV-backed dense attention backend’s layout). Causal aligns each query block to the end of its keys.

Parameters:

  • num_qo_heads (int)

  • num_kv_heads (int)

  • head_dim (int)

  • max_num_segments (int | None)

  • max_total_tokens (int | None)

  • max_total_kv_tokens (int | None)

  • device (device)

  • accelerator_graph (bool)

  • sm_scale (float | None)

  • q_data_type (dtype)

plan(q_cu_seqlens, kv_cu_seqlens, causal=False)[source]#

Each [num_segments + 1] with [0] == 0; CPU tensors preferred.

Parameters:
Return type:

None

class mstar.engine.resources.attn.ragged.wrappers.RaggedCrossPrefillWrapper(num_qo_heads, num_kv_heads, head_dim, max_num_segments=None, max_total_tokens=None, max_total_kv_tokens=None, device=torch.device('cuda'), accelerator_graph=False, sm_scale=None, q_data_type=torch.bfloat16)[source]#

Bases: _RaggedPrefillBase

Varlen cross-attention between two packed layouts: query segment i attends key segment i (one pair per request). Never causal.

The key side has its own token ceiling in graph mode, max_total_kv_tokens (defaults to max_total_tokens).

Parameters:

  • num_qo_heads (int)

  • num_kv_heads (int)

  • head_dim (int)

  • max_num_segments (int | None)

  • max_total_tokens (int | None)

  • max_total_kv_tokens (int | None)

  • device (device)

  • accelerator_graph (bool)

  • sm_scale (float | None)

  • q_data_type (dtype)

plan(q_cu_seqlens, kv_cu_seqlens)[source]#

Plan one query layout against one key layout, each [num_segments + 1] with [0] == 0 and the same segment count. CPU tensors preferred.

Parameters:
Return type:

None

class mstar.engine.resources.attn.ragged.wrappers.RaggedPrefillWrapper(num_qo_heads, num_kv_heads, head_dim, max_num_segments=None, max_total_tokens=None, max_total_kv_tokens=None, device=torch.device('cuda'), accelerator_graph=False, sm_scale=None, q_data_type=torch.bfloat16)[source]#

Bases: _RaggedPrefillBase

Varlen self-attention: each segment attends within itself. cu_seqlens is both qo_indptr and kv_indptr.

Parameters:

  • num_qo_heads (int)

  • num_kv_heads (int)

  • head_dim (int)

  • max_num_segments (int | None)

  • max_total_tokens (int | None)

  • max_total_kv_tokens (int | None)

  • device (device)

  • accelerator_graph (bool)

  • sm_scale (float | None)

  • q_data_type (dtype)

plan(cu_seqlens, causal=False)[source]#

Plan one packed layout. cu_seqlens: [num_segments + 1], [0] == 0.

CPU tensor preferred; a GPU one costs a sync. Safe to call before every replay — values are copied through the static buffers, not rebound.

Parameters:
Return type:

None

mstar.engine.resources.attn.ragged.wrappers.padded_head_dim(head_dim)[source]#

Smallest FlashInfer-supported head dim >= head_dim.

Parameters:

head_dim (int)

Return type:

int