mstar.engine.resources.attn.ragged.flashinfer

Cacheless varlen attention through FlashInfer’s ragged prefill wrapper.

Classes

FlashInferRaggedManager(device, dtype, config)

class mstar.engine.resources.attn.ragged.flashinfer.FlashInferRaggedManager(device, dtype, config)

Bases: RaggedAttnManager

Parameters:

• device (device)

• dtype (dtype)

• config (RaggedAttentionConfig)

clear_preplan()

num_segments(label=None)

Real (unpadded) segment count this step planned under label.

Parameters:

label (str | None)

Return type:

int

plan(step, ctx)

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

Parameters:

• step (AttentionStep)

• ctx (StepContext)

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

One layer’s varlen self-attention over this step’s packed segments.

Not behind a custom op, unlike the paged backend’s run: the ragged caller (an encoder tower) has no per-layer KV write to keep in the same graph, so the break this costs is one per layer of a region that is CUDA-graph captured rather than compiled.

Parameters:

• q (Tensor)

• k (Tensor)

• v (Tensor)

• label (str | None)

Return type:

Tensor

property supports_preplan