mstar.engine.resources.attn.dense

mstar.engine.resources.attn.dense#

Attention as one dense varlen pass: FlashAttention-3, AITER on ROCm, or FlashInfer’s ragged prefill where neither has a kernel for the GPU.

Classes

DenseAttentionManager(kv_cache, device, ...)

Attention as one dense varlen pass over a contiguous [frozen prefix | fresh tokens] sequence, through a prefixed ragged wrapper: FA3's (AITER on ROCm), or FlashInfer's where FA3 has no kernel for the GPU (ragged.base.dense_or_flashinfer picks; AttentionManager.build passes the class).

DensePlan(segments, cu_q, cu_k, max_q, ...)

dense plan layout.

DenseSegment(request_id, label, pages, ...)

(request, label) of dense plan; frozen prefix and fresh tokens
class mstar.engine.resources.attn.dense.DenseAttentionManager(kv_cache, device, dtype, kv_config, wrapper_class, flashinfer_backend='auto')[source]#

Bases: AttentionManager

Attention as one dense varlen pass over a contiguous [frozen prefix | fresh tokens] sequence, through a prefixed ragged wrapper: FA3’s (AITER on ROCm), or FlashInfer’s where FA3 has no kernel for the GPU (ragged.base.dense_or_flashinfer picks; AttentionManager.build passes the class).

For a workload that recomputes every one of its K/V every step and only reuses a small frozen prefix — diffusion denoise, where the prefix is the text conditioning — the paged path’s per-step full-buffer K/V write and wrapper.plan are pure overhead. This gathers the prefix once per (stream, layer), concatenates it with the freshly projected K/V, and runs one kernel. The gathered prefix is reused across steps: it is keyed on the stream’s generation, so a fork, a reset, or a page-table move invalidates it (CacheStream.generation).

The layer does not write its K/V to the pages here (requires_kv_write is False) — nothing would ever read it back.

Eager-only. The gather and the concatenation are shape-dependent and allocate, so there is nothing to capture; a node that runs this walk under a graph should name the paged backend for it.

NOTE on the step declaration: the fresh tokens are declared as ordinary spans on KVStep with commit=False, exactly as they would be for the paged backend, so the query lengths arrive as the KV plan’s to_compute and the prefix as length - to_compute. That is what keeps the backend a spec-time choice: the same declaration runs either way, and positions still take their packing off the KV plan output. The alternative — declaring the fresh tokens zero-span and carrying their lengths on AttentionStep’s own segments — is what the v0 dense path did and would save the pages that are reserved here and never written; it costs the swappability, since a model would then have to declare its step differently per backend, and a zero-span label yields no packing for any resource that derives from it (PositionManager would only work for labels supplying explicit pos_ids).

Parameters:
can_capture()[source]#
Return type:

bool

depends_on()[source]#
fingerprint()[source]#

What this resource contributes to the prefix cache’s root.

Everything a stored page depends on beyond its own tokens. None leaves the resource out of the root, which is right only where it cannot change what a page holds.

Return type:

bytes

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.

ingest_request(rid, overrides=None)[source]#

Must be idempotent: the conductor sends one NewRequest per partition, so a worker serving several partitions of a request ingests it once for each, possibly after a node has already run a step for it.

Parameters:

rid (str)

plan(step, ctx)[source]#

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

Parameters:
remove_request(rid)[source]#
Parameters:

rid (str)

property requires_kv_write: bool#

Whether a layer must write this step’s K/V through the KV resource before calling run.

False only for backends that take the fresh K/V straight into the kernel (the dense one). A layer therefore reads

if self.attn.requires_kv_write:

self.kv.write_kv(k, v, layer_idx=i, label=label)

out = self.attn.run(q, label, kv.layer_view(i), k=k, v=v, layer_idx=i)

and stays correct whichever backend the spec named.

reset_request(rid, free=False)[source]#

For clearing dummy RIDs during cuda graph capture

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

One layer’s dense attention. k/v are this step’s freshly projected K/V, packed in the plan’s segment order; they are never written to the pages.

Parameters:
Return type:

Tensor

class mstar.engine.resources.attn.dense.DensePlan(segments, cu_q, cu_k, max_q, max_k, causal, wrapper=None)[source]#

Bases: object

dense plan layout. once per step and read by each layer run

Parameters:
causal: bool#
cu_k: tuple[int, ...]#
cu_q: tuple[int, ...]#
max_k: int#
max_q: int#
segments: tuple[DenseSegment, ...]#
wrapper: PrefixedPrefillWrapper | None = None#
class mstar.engine.resources.attn.dense.DenseSegment(request_id, label, pages, prefix_len, q_len, generation)[source]#

Bases: object

(request, label) of dense plan; frozen prefix and fresh tokens

Parameters:
generation: int#
label: str#
pages: tuple[int, ...]#
prefix_len: int#
q_len: int#
request_id: str#