mstar.engine.resources.attn.ragged.base

mstar.engine.resources.attn.ragged.base#

The cacheless attention resource: the spec-time factory, and the plan machinery its backends share.

Sibling of attn.base’s AttentionManager, for attention that is not backed by a KV cache, e.g., an encoder tower attending within variable-length segments of one packed forward. There is nothing to page, nothing to advance, and nothing to carry between steps: the resource holds only planned wrappers.

Functions

dense_or_flashinfer(device, dense, ...)

dense where FA3 (AITER on ROCm) has a varlen kernel for device, else flashinfer.

resolve_ragged_attn_backend(backend, device)

Ragged counterpart of attn.base.resolve_attn_backend.

Classes

RaggedAttnManager()

RaggedBlockCausalAttnManager()

Cacheless block-causal self-attention within each span; see RaggedBlockCausalAttentionSpec.

RaggedCrossAttnManager()

Cacheless cross-attention between two spans of each request; see RaggedCrossAttentionSpec.
class mstar.engine.resources.attn.ragged.base.RaggedAttnManager[source]#

Bases: AttentionResource

classmethod build(spec, info)[source]#
Parameters:
prefix_skip_safe: bool = True#
class mstar.engine.resources.attn.ragged.base.RaggedBlockCausalAttnManager[source]#

Bases: AttentionResource

Cacheless block-causal self-attention within each span; see RaggedBlockCausalAttentionSpec.

classmethod build(spec, info)[source]#
Parameters:
prefix_skip_safe: bool = True#
class mstar.engine.resources.attn.ragged.base.RaggedCrossAttnManager[source]#

Bases: AttentionResource

Cacheless cross-attention between two spans of each request; see RaggedCrossAttentionSpec.

classmethod build(spec, info)[source]#
Parameters:
prefix_skip_safe: bool = True#
mstar.engine.resources.attn.ragged.base.dense_or_flashinfer(device, dense, flashinfer, what)[source]#

dense where FA3 (AITER on ROCm) has a varlen kernel for device, else flashinfer. The one place that choice is made, for the ragged managers and the KV-backed dense backend’s wrappers alike; what names the caller in the log line.

Parameters:
mstar.engine.resources.attn.ragged.base.resolve_ragged_attn_backend(backend, device)[source]#

Ragged counterpart of attn.base.resolve_attn_backend.

Parameters:
Return type:

RaggedAttnBackend