`core.inference.contexts.mtp_metadata`#

Module Contents#

Classes#

`MTPForwardMode`	Which MTP forward, if any, currently owns the attention metadata.
`MTPMetadata`	MTP request annotations, chunk carry, and KV-cache forward buffers.

API#

class core.inference.contexts.mtp_metadata.MTPForwardMode(*args, **kwds)#

Bases: enum.Enum

Which MTP forward, if any, currently owns the attention metadata.

The two non-NONE modes differ in query shape, which is what decides the attention kernel: a draft depth contributes exactly one token per active request (uniform, so the decode kernel’s q.reshape(num_requests, tokens_per_request, ...) is valid), while the commit pass is roll-by-one over each request’s committed span (ragged, so it must go varlen).

Initialization

NONE#: ‘none’

DRAFT#: ‘draft’

COMMIT#: ‘commit’

class core.inference.contexts.mtp_metadata.MTPMetadata#

MTP request annotations, chunk carry, and KV-cache forward buffers.

A speculative step runs several MTP forwards back to back – one per draft depth, plus a varlen “commit pass” – and every one of them rewrites the same small set of per-request fields. Rather than build those tensors per forward, this class owns them as fixed-address buffers sized for the worst case and updates them in place. That keeps the draft loop allocation-free and CUDA-graph safe (a captured MTP graph replays against the same addresses each step). Construction is cheap and unconditional; allocate is what reserves GPU memory, and only when enabled. CPU request annotations follow the context’s row lifecycle through move_request_rows, swap_request_rows, and reset_request_rows.

Parameters:

enabled (bool) – Whether MTP KV caching is active for this context. When False the object is inert: no buffers are allocated and forward_mode stays NONE.
max_requests (int) – Worst-case active request count (max_bs).
max_kv_block_count (int) – Block-table width, in blocks per request.
block_size_tokens (int) – KV block size, used to split a write position into (block column, offset within block).
dummy_block_idx (int) – Scratch block that padded rows read and write, so padding never touches real KV.
block_table_dtype (torch.dtype) – dtype of gpu_view.mha_block_table.
hidden_size (int) – Model hidden size, sizing the chunked-prefill boundary carry.
hidden_dtype (torch.dtype) – dtype of the main model’s hidden states.

enabled: bool#: None

max_requests: int#: None

max_kv_block_count: int#: None

block_size_tokens: int#: None

dummy_block_idx: int#: None

block_table_dtype: torch.dtype#: None

hidden_size: int#: None

hidden_dtype: torch.dtype#: None

request_matched_prefix_blocks: Optional[torch.Tensor]#: ‘field(…)’

forward_mode: core.inference.contexts.mtp_metadata.MTPForwardMode#: None

graphed: bool#: False

active_request_count: int#: 0

padded_count: int#: 0

chunk_boundary_valid: bool#: False

chunk_boundary_req_id: int#: None

chunk_boundary_position: int#: None

chunk_boundary_hidden: Optional[torch.Tensor]#: ‘field(…)’

offsets: Optional[torch.Tensor]#: ‘field(…)’

block_table: Optional[torch.Tensor]#: ‘field(…)’

query_lengths: Optional[torch.Tensor]#: ‘field(…)’

kv_lengths: Optional[torch.Tensor]#: ‘field(…)’

row_ids: Optional[torch.Tensor]#: ‘field(…)’

active_offsets: Optional[torch.Tensor]#: ‘field(…)’

active_block_table: Optional[torch.Tensor]#: ‘field(…)’

property forward_active: bool#: Whether any MTP forward currently owns the attention metadata.

property is_varlen_forward: bool#: Whether the current MTP forward has ragged per-request query lengths.

allocate(device: torch.device) → None#

Reserve the persistent buffers. No-op when MTP KV caching is disabled.

Parameters:: device (torch.device) – Device for the GPU-resident buffers.

deallocate() → None#

Release the persistent buffers, mirroring allocate.

Used by the context’s suspend path, which drops its tensors and rebuilds them from initialize_all_tensors on resume.

reset_request_rows(request_indexes=slice(None)) → None#

Clear MTP annotations when request rows are released, reused, or reset.

The chunk carry is keyed by request ID, so moving or clearing physical rows does not change it. Its owner invalidates it when that logical request ends.

move_request_rows( src_idxs: torch.Tensor, dst_idxs: torch.Tensor, ) → None#

Copy MTP annotations alongside the scheduler’s request-row movement.

Sources remain intact until the scheduler identifies and clears vacated rows. Advanced indexing snapshots the source values even when the ranges overlap.

swap_request_rows( src_idxs: torch.Tensor, dst_idxs: torch.Tensor, ) → None#: Swap MTP annotations alongside the scheduler’s paused/active row swap.

invalidate_chunk_boundary() → None#

Drop the carried chunk-boundary hidden.

Clears the validity and the keys, not the buffer – its address stays stable. Called from deallocate and whenever no chunked request is in flight, so a stale carry can never be matched by a later request.

carry_chunk_boundary( hidden: torch.Tensor, req_id: int, position: int, ) → None#

Stash the in-flight chunked request’s last hidden for its next chunk’s seam.

Parameters:

hidden (Tensor) – Main hidden state at position, any shape with hidden_size elements. Copied into the persistent buffer, so the caller’s tensor is free to be released with the rest of the step’s activations.
req_id (int) – Producing request’s id.
position (int) – Prompt position hidden was computed at. Only a chunk whose seam falls exactly here – i.e. whose off == position + 1 – may consume this carry.

take_chunk_boundary(req_id: int, seam_position: int) → torch.Tensor#

Return the carried hidden for this seam.

All three conditions below hold by construction. The caller derives req_id by matching chunk_boundary_req_id in its own request list, which settles the first two. The third rests on _compute_prefix_match giving up a carry-holding continuation chunk’s ENTIRE prefix match, so the chunk starts at finished and its seam lands exactly where the carry was recorded.

They RAISE rather than declining: skipping the seam would leave a committed position unwritten, a silent draft-acceptance regression that masks whichever invariant broke.

Parameters:

req_id (int) – Request that wants to write the seam.
seam_position (int) – MTP position the seam would be written at (off - 1).

Returns:

The [1, 1, hidden_size] carried hidden.

Return type:

(Tensor)

begin_decode( active_request_count: int, padded_count: int, start_positions: torch.Tensor, block_table_src: torch.Tensor, graphed: bool, ) → None#

Stage the draft loop’s write positions and block table.

Both are copied into the persistent buffers (rather than aliased), so advancing the depth cannot mutate the caller’s tensors.

Parameters:

active_request_count (int) – Requests taking part in the draft loop.
padded_count (int) – Request/token slots the forwards launch over.
start_positions (Tensor) – [>=active_request_count] depth-0 write position per request.
block_table_src (Tensor) – [>=active_request_count, max_kv_block_count] source table.
graphed (bool) – Whether the loop replays captured CUDA graphs.

begin_decode_for_capture(padded_count: int) → None#

Stage synthetic (scratch-only) draft state for CUDA-graph capture at warmup.

Every row is pointed at the dummy block at position 0, so the captured append/attend touch only scratch KV. Replay overwrites all of it, so only the shapes and the fixed launch bounds matter here, and those match the runtime graphed step.

_enter_decode( active_request_count: int, padded_count: int, graphed: bool, ) → None#: Enter MTP-forward mode and refresh the active views.

advance_decode_step() → None#: Advance every active request’s MTP write position by one, in place.

end_forward() → None#: Leave MTP-forward mode. No persistent length state to write back.

stage_decode_lengths() → Tuple[torch.Tensor, torch.Tensor]#

Stage the MHA lengths for one draft depth: one query row, position + 1 keys.

Returns:: Views of the query- and KV-length buffers, both active_request_count long.
Return type:: (Tuple[Tensor, Tensor])

stage_prefill_lengths( append_counts: torch.Tensor, pad_tokens: int, ) → torch.Tensor#

Stage the MHA lengths for a varlen commit pass, absorbing the padded query rows.

The commit pass is a fresh causal prefill, so a request’s KV length equals its query length and one buffer serves as both.

The hidden is padded to a TP multiple for the sequence-parallel scatter, so attention sees total + pad_tokens query rows while append_counts describes only total. Varlen requires q.shape[0] == cu_seqlens_q[-1], so every pad row needs an owning request – and it does not matter which, since write_token_maps sends them to the dummy block.

Parameters:

append_counts (Tensor) – [P] KV entries each request writes this pass.
pad_tokens (int) – Query rows appended to reach the TP multiple.

Returns:

Per-request length view, P long, or P + 1 when the pad rows were given their own trailing request.

Return type:

(Tensor)

write_token_maps( gpu_view: core.inference.contexts.gpu_view.ContextGPUView, rows: torch.Tensor, positions: torch.Tensor, block_table: torch.Tensor, padded_token_count: int, inherited_blocks: Optional[torch.Tensor] = None, ) → None#

Write the per-token KV destination maps for one MTP forward.

Parameters:

gpu_view (ContextGPUView) – Destination bookkeeping views.
rows (Tensor) – [T] owning request row (into block_table) for each token.
positions (Tensor) – [T] MTP write position within the request for each token.
block_table (Tensor) – [R, max_kv_block_count] block ids indexed by rows.
padded_token_count (int) – Token rows the forward runs, including padding. Padded rows are redirected to the dummy block so they never touch real KV.
inherited_blocks (Optional[Tensor]) – [R] leading blocks each row INHERITED rather than computed. Tokens landing in those blocks go to the dummy block: the KV is already correct from the producer, and the block is ref-counted, so writing would corrupt every request sharing it. None disables the redirect.

write_mha_metadata( gpu_view: core.inference.contexts.gpu_view.ContextGPUView, query_lengths: torch.Tensor, kv_lengths: torch.Tensor, block_table: torch.Tensor, padded_request_count: int, ) → None#

Write the per-request MHA read metadata for one MTP forward.

Parameters:

gpu_view (ContextGPUView) – Destination bookkeeping views.
query_lengths (Tensor) – [R] query rows contributed by each request.
kv_lengths (Tensor) – [R] KV entries each request attends over. May be the same tensor as query_lengths when the two coincide.
block_table (Tensor) – Block ids for the first block_table.shape[0] requests. Any remaining request slot – the trailing pad request staged by stage_prefill_lengths – is pointed at the dummy block.
padded_request_count (int) – Request slots the kernel launches over, including padding. Padded slots get zero lengths and a -1 block table, so they never index real KV.

core.inference.contexts.mtp_metadata#

Module Contents#

Classes#

API#

`core.inference.contexts.mtp_metadata`#