mstar.engine.resources.attn.ragged.flashinfer#

Cacheless varlen attention through FlashInfer’s ragged prefill wrapper.

Classes

`FlashInferRaggedBlockCausalManager`(device, ...)	Block-causal self-attention: every label in the step is its own layout, each segment attending block-causally within itself.
`FlashInferRaggedCrossManager`(device, dtype, ...)	Cross-attention: each `(q_label, kv_label)` pair of the step is one plan, run under `cross_label(q_label, kv_label)`.
`FlashInferRaggedManager`(device, dtype, config)	Self-attention: every label in the step is its own layout, each segment attending within itself.

class mstar.engine.resources.attn.ragged.flashinfer.FlashInferRaggedBlockCausalManager(device, dtype, config, block_size)[source]#

Bases: _FlashInferRaggedBase, RaggedBlockCausalAttnManager

Block-causal self-attention: every label in the step is its own layout, each segment attending block-causally within itself.

Parameters:

plan(step, ctx)[source]#

ret is immutable and opaque to runner; only gives to ctx.plan_results

Parameters:

class mstar.engine.resources.attn.ragged.flashinfer.FlashInferRaggedCrossManager(device, dtype, config)[source]#

Bases: _FlashInferRaggedBase, RaggedCrossAttnManager

Cross-attention: each (q_label, kv_label) pair of the step is one plan, run under cross_label(q_label, kv_label).

Parameters:

class mstar.engine.resources.attn.ragged.flashinfer.FlashInferRaggedManager(device, dtype, config)[source]#

Bases: _FlashInferRaggedBase, RaggedAttnManager

Self-attention: every label in the step is its own layout, each segment attending within itself.

Parameters:

plan(step, ctx)[source]#

ret is immutable and opaque to runner; only gives to ctx.plan_results

Parameters:

mstar.engine.resources.attn.ragged.flashinfer