mstar.engine.resources.attn.ragged.wrappers

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

Functions

padded_head_dim(head_dim)

Smallest FlashInfer-supported head dim >= head_dim.

Classes

RaggedCrossPrefillWrapper(workspace_buffer, ...)

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

RaggedPrefillWrapper(workspace_buffer, ...)

Varlen self-attention: each segment attends within itself.
class mstar.engine.resources.attn.ragged.wrappers.RaggedCrossPrefillWrapper(workspace_buffer, num_qo_heads, num_kv_heads, head_dim, max_num_segments=None, max_total_tokens=None, max_total_kv_tokens=None, **kwargs)[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:

  • workspace_buffer (Tensor)

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

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(workspace_buffer, num_qo_heads, num_kv_heads, head_dim, max_num_segments=None, max_total_tokens=None, **kwargs)[source]#

Bases: _RaggedPrefillBase

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

Parameters:

  • workspace_buffer (Tensor)

  • num_qo_heads (int)

  • num_kv_heads (int)

  • head_dim (int)

  • max_num_segments (int | None)

  • max_total_tokens (int | None)

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