mstar.engine.resources.attn.ragged.flashinfer

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)[source]#

Bases: RaggedAttnManager

Parameters:
clear_preplan()[source]#
property force_double_buffer#

Whether this resource needs double-buffering even off the pre-plan path.

A resource that stages a step’s layout into a reused host buffer and issues a non-blocking H2D into a graph-read device buffer has a race the moment the CPU runs ahead of the GPU: plan(N+1) can overwrite the buffer before step N’s DMA has retired, and the replay attends with N+1’s data. The buffer need not be ours — FlashInfer’s plan holds one per wrapper, which is why the attention managers key theirs per slot. The main runner already double-buffers whenever a resource supports_preplan; this flag extends that to a runner that has no pre-plan path but still replays such a resource — notably the piecewise runner — and makes _exec_per_request fence before it reuses a slot, since its sub-steps all run inside one step. Off by default; a resource opts in only if it has this hazard.

num_segments(label=None)[source]#

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

Parameters:

label (str | None)

Return type:

int

plan(step, ctx)[source]#

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

Parameters:
run(q, k, v, label=None)[source]#

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:
Return type:

Tensor

property supports_preplan#