`vllm.models.deepseek_v41.nvidia.specialized_decode.sampling.sharded_markov_cute` ¶

DSpark's sequential Markov sampling for a small batch, in CuTe DSL.

The kernel samples up to 8 request rows, with the Markov head sharded by vocabulary across the tensor-parallel ranks (sharded_markov.py explains the design). It is compiled for a number of draft positions steps (1..8) and the element types of the base logits and of the draft-logits cache (BF16 or FP32).

For request r, step i needs r's previous token p (its anchor for i = 0)::

logits_ri = cache_round(base_ri + bf16(markov_head . markov_embed[p]))
d_r,i+1   = argmax(logits_ri)                     (temperature 0)
          = argmax(logits_ri / T_r + gumbel_ri)   (otherwise)

Each rank owns rows [rank * V/world, (rank + 1) * V/world) of the head. Each warp keeps one 16-row tile of that shard in registers for all steps, so the head is read once per launch; the m16n8k16 MMA's eight columns carry the eight request rows' embeddings, so every row's bias comes from the same MMAs. A step is those MMAs, a CTA reduction and a two-level exchange of 64-bit candidate keys per request row: the CTAs of a rank fold their best key into a rank-local maximum, the last CTA to arrive stores that key, tagged with the launch epoch, into every rank's mailbox, and every CTA waits for the world records of each row in its own mailbox and takes their maximum. The exchange is the grid barrier and the cross-rank argmax at once, so the grid must be co-resident: the kernel is launched cooperatively, and its 512-thread requirement keeps one CTA per SM within the register file.

The base logits may hold the full vocabulary or only this rank's shard (the draft LM head's output before its all-gather): base_column is the column of this rank's first vocabulary entry in them, rank * V/world or 0.

Rows whose request-state index is negative (CUDA-graph padding) are inert, as in the speculator's loop: their tokens are drawn greedily and their logits are not cached.

The draft-logits cache must hold the full vocabulary on every rank. Each rank writes its shard into its own cache and into a staging buffer of every peer; after the last exchange each rank copies the peers' shards from its staging buffer into its cache. Staging and mailbox records alternate between two parities so that a launch never overwrites what a peer may still be reading from the previous launch.

Peer buffer layout (bytes): mailbox records u64 [parity 2][step 8][row 8][source rank 8] (8 KiB), then the staging logits [parity 2][max_rows][steps][vocab] of the cache type at STAGING_OFFSET. Local control words (uint32): combine maxima u64 [step 8][row 8] at words 0..127, arrivals [step 8] at 128..135, launches at 136, end arrivals at 137.

Global memory is addressed with 64-bit integers and accessed with explicit PTX, so that each access has the state space, cache operator and memory ordering the protocol needs (relaxed system-scope mailbox records, L2-only staging reads, non-coherent reads of the weights).

Classes:

ShardedMarkovSample –

The sharded Markov chain kernel for steps draft positions, base

Functions:

add_rn –

FP32 add, never contracted into a neighbouring multiply.
atomic –

A global atomic returning the old value (opcode e.g.
await_record –

Spin until the mailbox record at address carries tag; return
candidate_key –

[48:17] order-preserving score bits, [16:0] inverted index (ties prefer
div_full –

div.full.f32: Triton's FP32 /, as the stock sampler divides.
fence –

membar.gl (__threadfence) or membar.sys
float_bits –

The bits of an FP32 value.
key_index –

The vocabulary index of a candidate key.
load_bf16_as_f32 –

A global BF16 value, widened.
load_scalar –

One 32- or 64-bit global value of DSL type kind (opcode names
load_shared_v4 –

Four 32-bit words at a 16-byte aligned shared address.
load_v4 –

Four 32-bit words at a 16-byte aligned global address.
max_u64 –

The larger of two keys (unsigned; tensor loads come back signed).
mma_16816 –

mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32: acc four
prefetch_row_l2 –

Pull one 512-byte row into L2.
round_bf16 –

value rounded to BF16 (nearest even) and widened back.
shared_address –

The shared-window address of a shared-memory pointer.
shuffle_idx_u64 –

__shfl_sync of a 64-bit value from lane source.
shuffle_xor_u64 –

__shfl_xor_sync of a 64-bit value over the full warp.
store_global_v4 –

st.global.v4.u32.
store_scalar –

One 32- or 64-bit global store (opcode, e.g. st.global.u32).
store_shared_v4 –

st.shared.v4.u32.
trap –

Abort the kernel.

`ShardedMarkovSample` ¶

The sharded Markov chain kernel for steps draft positions, base logits of base_dtype and a draft-logits cache of cache_dtype (cutlass.BFloat16 or cutlass.Float32).

Methods:

__call__ –

Launch the chain as a cooperative grid of ctas CTAs.
kernel –

One CTA of 16 warps; see the module docstring.
load_base –

Base logit of request row r, step and shard row row.
load_inputs –

Base logits and noise of this lane's resident tile rows and request
load_noise –

Gumbel noise of request row r, step and shard row row.
peer_base –

The base address of peer's buffer.