mstar.engine.resources.attn.dense

mstar.engine.resources.attn.dense#

Attention as one dense FlashAttention-3 varlen pass.

Classes

`DenseAttentionManager`(kv_cache, device, ...)	Attention as one dense FlashAttention-3 varlen pass over a contiguous [frozen prefix \| fresh tokens] sequence.
`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)[source]#

Bases: AttentionManager

Attention as one dense FlashAttention-3 varlen pass over a contiguous [frozen prefix | fresh tokens] sequence.

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:

kv_cache (str)
device (device)
dtype (dtype)
kv_config (KVConfig)

depends_on()[source]#

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

Parameters:: rid (str)

plan(step, ctx)[source]#

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

Parameters:

step (AttentionStep)
ctx (StepContext)

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: