Is this page helpful?

core.ssm.ops.gdp.metadata#

Chunk descriptors consumed by the forked Gated Delta Product prefill kernels.

Upstream FLA derives these inside the op with prepare_chunk_indices, which synchronizes on the device (.tolist()) and returns a data-dependent shape. Both are fatal under CUDA-graph capture, so they are built here instead – once per step, on the host, outside the graph – and padded to a size that depends only on the padded batch shape. The kernels then read fixed-size, fixed-address buffers.

Two chunkings are needed. chunk_indices describes the token stream as the queries see it; chunk_indices_dp describes the Householder-expanded stream that the keys, values and betas live on, whose sequences are M times longer. The second is not a rescaling of the first: ceil(L*M/64) != M*ceil(L/64) whenever L is not a multiple of the chunk size.

Module Contents#

Functions#

max_gdp_chunk_counts

Worst-case chunk counts, used to size the persistent buffers.

build_gdp_chunk_descriptors

Build the padded chunk descriptors for one prefill step.

API#

core.ssm.ops.gdp.metadata.max_gdp_chunk_counts(
max_tokens: int,
max_requests: int,
num_householder: int,
) → Tuple[int, int]#

Worst-case chunk counts, used to size the persistent buffers.

Returns (max_chunks, max_chunks_dp) for the unexpanded and the Householder-expanded stream respectively. The + max_requests term covers the partial trailing chunk each sequence may own.

core.ssm.ops.gdp.metadata.build_gdp_chunk_descriptors(
cu_seqlens: List[int],
prefill_count: int,
num_householder: int,
padded_token_count: int,
) → Tuple[List[int], List[int], List[int], int, int]#

Build the padded chunk descriptors for one prefill step.

Parameters:

  • cu_seqlens – prefill_count + 1 cumulative sequence lengths, covering real and padding requests (padding requests are zero-length).

  • prefill_count – Padded number of prefill requests.

  • num_householder – Number of Householder copies M.

  • padded_token_count – Padded token count, which fixes the buffer sizes.

Returns:

(chunk_indices, chunk_indices_dp, chunk_offsets, num_chunks, num_chunks_dp). The two chunk_indices lists are flattened (sequence, chunk-within-sequence) pairs padded to their fixed lengths; chunk_offsets is the per-sequence prefix sum of unexpanded chunk counts.