mstar.engine.resources.attn.workspace

FlashInfer’s workspaces: the scratch a plan reserves its split-KV partials in.

Functions

workspace_bytes([declared_mb])

The floor of one FlashInfer workspace: what its resource declares, else MSTAR_WORKSPACE_BUFFER_MB (default 512).

Classes

Workspace(tensor, name[, sealed, floor])	One FlashInfer float workspace, which wrappers plan against through plan.
WorkspacePool(device[, size_mb, owner])	FlashInfer workspaces, one per (plan label, cg slot).

class mstar.engine.resources.attn.workspace.Workspace(tensor, name, sealed=False, floor='MSTAR_WORKSPACE_BUFFER_MB')

Bases: object

One FlashInfer float workspace, which wrappers plan against through plan.

What a plan needs depends on its rows, segments and heads and on the GPU’s SM count, and under a graph it is the same for every plan of a wrapper. So a workspace starts at its floor and a plan that overflows it grows it. tensor keeps its identity across that, so every wrapper on it moves too and the key holds one buffer, as large as the largest plan made on it.

A captured graph records the buffer’s address, so once seal`ed it can’t grow, and a plan that doesn’t fit raises what it asked for and which setting would have made room. ``floor` names that setting.

Parameters:

• tensor (Tensor)

• name (str)

• sealed (bool)

• floor (str)

floor

name

plan(plan, *args, **kwargs)

plan(*args, **kwargs). One that runs short of room grows this workspace to the bytes it had taken plus the allocation that failed, and runs again.

Parameters:

plan (Callable)

seal()

Return type:

None

sealed

tensor

class mstar.engine.resources.attn.workspace.WorkspacePool(device, size_mb=None, owner='attention')

Bases: object

FlashInfer workspaces, one per (plan label, cg slot).

Persistent: a wrapper is planned against the buffer it was built with, so these outlive any single step. owner names the resource in what a workspace reports.

Parameters:

• device (device)

• size_mb (int | None)

• owner (str)

get(label, cg_slot=None)

Parameters:

• label (str)

• cg_slot (int | None)

Return type:

Workspace

seal()

The workspaces here keep their addresses from now on, for the graphs about to be captured over them, and so does any made later.

Return type:

None

mstar.engine.resources.attn.workspace.workspace_bytes(declared_mb=None)

The floor of one FlashInfer workspace: what its resource declares, else MSTAR_WORKSPACE_BUFFER_MB (default 512).

Parameters:

declared_mb (int | None)

Return type:

int