mstar.engine.resources.attn.cross

mstar.engine.resources.attn.cross#

Cross-attention over a context written once and never extended.

Classes

CrossAttentionManager()

FlashInferCrossManager(kv_cache, ...[, backend])

non-causal paged attention over fixed encoder context

QueryPacking(request_ids, qo_indptr)

One plan label's query side: the requests in packed order, and the cumulative query lengths over them.
class mstar.engine.resources.attn.cross.CrossAttentionManager[source]#

Bases: AttentionResource

classmethod build(spec, info)[source]#
Parameters:
class mstar.engine.resources.attn.cross.FlashInferCrossManager(kv_cache, query_kv_cache, context_label, device, dtype, kv_config, backend='auto')[source]#

Bases: CrossAttentionManager

non-causal paged attention over fixed encoder context

context is an ordinary kv stream so requires no additional storage or allocation on part of attention manager

encode step writes context via normal Segment on cache, and later steps use span=0 on that label

wrapper is planned per decoder plan label and run is keyed on that label a la FlashInferManager

Parameters:
clear_preplan()[source]#
property context_cache_key: str#

label of kv resource holding context so x-attention can find cache to read

depends_on()[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.

plan(step, ctx)[source]#

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

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

One layer’s cross attention. kv_cache_layer is a layer of the context cache; nothing is written to it here.

Parameters:
Return type:

Tensor

property supports_preplan#
class mstar.engine.resources.attn.cross.QueryPacking(request_ids, qo_indptr)[source]#

Bases: NamedTuple

One plan label’s query side: the requests in packed order, and the cumulative query lengths over them.

Parameters:
qo_indptr: Tensor#

Alias for field number 1

request_ids: list[str]#

Alias for field number 0