mstar.engine.resources.dsa.manager

DeepSeek sparse attention (DSA) over the paged MLA cache: the index keys, each query’s top-k and the attention over it.

A model’s FULL indexer layers hand this resource their tokens’ index keys (store_keys) and their queries and head weights (select): each query row scores its request’s keys, sum_h w[h] * relu(q[h] . k) in fp32, and keeps its top topk positions. SHARED layers reuse the last FULL layer’s pick (shared_selection), and attend runs absorbed MLA over just the picked latents of the MLA cache. A step’s rows come from the KV plan: each request’s rows, the position they start at and its pages.

The keys live in per-request buffers that grow with the context and the selection runs per request on the host, so the path is eager only; past topk it decodes (one row per request). The follow-up moves the keys onto the MLA cache’s pages, under its page table, allocator and eviction, and gives selection and attention fixed-shape buffers so a decode can capture.

Functions

index_topk(q, weights, keys, positions, topk)	Top-topk key positions per query row, (T, topk) int32, -1 padded.
select_topk_causal(scores, positions, topk)	Per-row top-topk key positions of causally masked scores.

Classes

DSAManager(config, page_size, device)	DSA for one node's layers over its MLA cache (config.kv_cache).
DSAPlan(spans, needs_selection[, ...])	A label's rows this step.
DSASpan(request_id, q_start, q_len, ...[, ...])	One request's rows of a step: q_len rows from row q_start, the first at position ctx_start; pages its MLA cache pages, page_start its first token's offset.
IndexKeyStore()	Per-request, per-layer growing buffers of the indexer's keys, in position order.

class mstar.engine.resources.dsa.manager.DSAManager(config, page_size, device)

Bases: Resource

DSA for one node’s layers over its MLA cache (config.kv_cache).

Parameters:

• config (DSAConfig)

• page_size (int)

• device (device)

attend(selection, q_nope, q_pe, latent_layer, sm_scale, label=_MAIN)

Absorbed MLA of each row over its picked latents, (T, H, kv_lora_rank).

latent_layer is the MLA cache’s layer view [pages, page_size, kv_lora_rank + rope], this step’s latents already written; selection holds request-local positions (-1: none). Past topk a request decodes, so row r is request r’s.

Parameters:

• selection (Tensor)

• q_nope (Tensor)

• q_pe (Tensor)

• latent_layer (Tensor)

• sm_scale (float)

• label (str)

Return type:

Tensor

classmethod build(spec, info)

Parameters:

• spec (DSASpec)

• info (EngineResourceInfo)

depends_on()

Return type:

set[str]

needs_selection(label=_MAIN)

Whether this step’s rows attend a pick rather than every key.

Parameters:

label (str)

Return type:

bool

plan(step, ctx)

Each label’s rows, as the KV plan placed them: a request’s rows follow the one before it, and its first sits at its length before the step.

Parameters:

• step (DSAStep)

• ctx (StepContext)

Return type:

dict[str, DSAPlan]

remove_request(rid)

Parameters:

rid (str)

Return type:

None

reset_request(rid, free=False)

For clearing dummy RIDs during cuda graph capture

Parameters:

• rid (str)

• free (bool)

Return type:

None

select(layer_idx, q, weights, label=_MAIN)

Each row’s top topk positions of its request (index_topk) over this layer’s stored keys, (T, topk) int32, -1 padded; SHARED layers reuse it this forward.

Parameters:

• layer_idx (int)

• q (Tensor)

• weights (Tensor)

• label (str)

Return type:

Tensor

shared_selection(layer_idx, label=_MAIN)

The last FULL layer’s pick this forward, for a SHARED layer (IndexShare).

Parameters:

• layer_idx (int)

• label (str)

Return type:

Tensor

store_keys(layer_idx, keys, label=_MAIN)

Store the step’s (T, D) index keys of one layer, each request’s at its positions.

Parameters:

• layer_idx (int)

• keys (Tensor)

• label (str)

Return type:

None

class mstar.engine.resources.dsa.manager.DSAPlan(spans, needs_selection, last_selection=None, pages=None, row_page_offset=None, row_start=None)

Bases: object

A label’s rows this step. needs_selection: some row sees more than topk keys; last_selection: the last FULL layer’s pick this forward, which SHARED layers reuse.

Under selection every row is its request’s one decode row, and pages holds the rows’ page tables back to back on the device: row r’s starts at row_page_offset[r], and its first token sits row_start[r] slots into it.

Parameters:

• spans (list[DSASpan])

• needs_selection (bool)

• last_selection (Tensor | None)

• pages (Tensor | None)

• row_page_offset (Tensor | None)

• row_start (Tensor | None)

last_selection: Tensor | None = None

needs_selection: bool

pages: Tensor | None = None

row_page_offset: Tensor | None = None

row_start: Tensor | None = None

spans: list[DSASpan]

class mstar.engine.resources.dsa.manager.DSASpan(request_id, q_start, q_len, ctx_start, pages, page_start=0)

Bases: object

One request’s rows of a step: q_len rows from row q_start, the first at position ctx_start; pages its MLA cache pages, page_start its first token’s offset.

Parameters:

• request_id (str)

• q_start (int)

• q_len (int)

• ctx_start (int)

• pages (list[int])

• page_start (int)

ctx_start: int

page_start: int = 0

pages: list[int]

q_len: int

q_start: int

request_id: str

class mstar.engine.resources.dsa.manager.IndexKeyStore

Bases: object

Per-request, per-layer growing buffers of the indexer’s keys, in position order.

append(request_id, layer_idx, keys, start_pos)

Append one chunk of (n, head_dim) keys at positions start_pos ...

Parameters:

• request_id (str)

• layer_idx (int)

• keys (Tensor)

• start_pos (int)

Return type:

None

evict(request_id)

Drop all of a request’s buffers. Idempotent — the engine may retire a request that never reached a DSA forward.

Parameters:

request_id (str)

Return type:

None

history(request_id, layer_idx, length)

The first length stored keys — a view, not a copy. Raises if fewer are stored (a scoring window must never exceed the appended history: the current chunk is appended before selection).

Parameters:

• request_id (str)

• layer_idx (int)

• length (int)

Return type:

Tensor

tokens(request_id, layer_idx)

Stored rows for one (request, layer); 0 when absent.

Parameters:

• request_id (str)

• layer_idx (int)

Return type:

int

tracked_requests()

Return type:

set[str]

mstar.engine.resources.dsa.manager.index_topk(q, weights, keys, positions, topk)

Top-topk key positions per query row, (T, topk) int32, -1 padded.

q (T, H, D) and weights (T, H) are the indexer’s queries and head weights (any scales folded in), keys (S, D) a request’s keys by position and positions (T,) the rows’ own: a row sees keys 0 .. position, itself included.

Parameters:

• q (Tensor)

• weights (Tensor)

• keys (Tensor)

• positions (Tensor)

• topk (int)

Return type:

Tensor

mstar.engine.resources.dsa.manager.select_topk_causal(scores, positions, topk)

Per-row top-topk key positions of causally masked scores.

scores is (T, num_keys) with every key past a row’s position already -inf; the result is (T, topk) int32, -1 padded where a row sees fewer than topk keys. One stable sort for the whole batch: no per-token topk launch, no host sync. Ties break toward the earlier key position, so a row’s selection does not depend on how the batch is packed or on which device ran it (torch.topk’s tie order is neither).

Parameters:

• scores (Tensor)

• positions (Tensor)

• topk (int)

Return type:

Tensor