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

RaggedPrefillWrapper(workspace_buffer, ...)

Varlen self-attention over packed segments, with no KV cache; attention is within variable-length segments packed into one [total_tokens, H, D] tensor.
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, device=torch.device('cuda'), use_cuda_graph=False, sm_scale=None, q_data_type=torch.bfloat16, kv_layout='NHD', backend='auto')[source]#

Bases: object

Varlen self-attention over packed segments, with no KV cache; attention is within variable-length segments packed into one [total_tokens, H, D] tensor. 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)

  • device (device)

  • use_cuda_graph (bool)

  • sm_scale (float | None)

  • q_data_type (dtype)

  • kv_layout (str)

  • backend (str)

property num_segments: int#

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

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

run(q, k, v)[source]#

Run planned varlen self-attention.

Parameters:

  • q (Tensor) – [total_tokens, num_heads, head_dim], packed by cu_seqlens

  • k (Tensor) – [total_tokens, num_heads, head_dim], packed by cu_seqlens

  • v (Tensor) – [total_tokens, num_heads, head_dim], packed by cu_seqlens

Returns:

[total_tokens, num_qo_heads, head_dim]

Return type:

output

Only rows before the planned cu_seqlens[-1] are computed; the rest read back zero. An oversized static buffer replays fine, but the caller must still slice — the padding rows are not a valid result.

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