mstar.engine.resources.attn.flashinfer

Paged attention through FlashInfer’s prefill/decode wrappers.

Classes

FlashInferManager(kv_cache, device, dtype, ...)

class mstar.engine.resources.attn.flashinfer.FlashInferManager(kv_cache, device, dtype, kv_config, backend='auto')

Bases: AttentionManager

Parameters:

• kv_cache (str)

• device (device)

• dtype (dtype)

• kv_config (KVConfig)

• backend (str)

clear_preplan()

depends_on()

plan(step, ctx)

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

Parameters:

• step (AttentionStep)

• ctx (StepContext)

qo_indptr_buf(label='main')

Parameters:

label (str)

Return type:

Tensor | None

run(q, label=None, kv_cache_layer=None, k=None, v=None, layer_idx=None)

Parameters:

• q (Tensor)

• label (str | None)

• kv_cache_layer (Tensor | None)

• k (Tensor | None)

• v (Tensor | None)

• layer_idx (int | None)

Return type:

Tensor

select_last_hidden(hidden, label='main')

Select last token of the hidden vector per request, used for sampling from prefill.

Parameters:

• hidden (Tensor)

• label (str)

Return type:

Tensor

property supports_preplan