vllm.models.deepseek_v41.nvidia.ops.fused_markov

DSpark's sequential Markov sampling as the bias GEMM plus two kernels a step. Tokens are drawn by inverse CDF of softmax(cached logits / temperature), one exp per logit instead of Gumbel-max's hash and two logs; temperature 0 is argmax.

Classes:

• FusedMarkovSampler –

  DSparkSpeculator._sample_sequential's hook. Its scratch is allocated

FusedMarkovSampler

DSparkSpeculator._sample_sequential's hook. Its scratch is allocated at model init, so CUDA graphs keep the addresses.

Source code in vllm/models/deepseek_v41/nvidia/ops/fused_markov.py

class FusedMarkovSampler:
    """``DSparkSpeculator._sample_sequential``'s hook. Its scratch is allocated
    at model init, so CUDA graphs keep the addresses."""

    def __init__(
        self, max_num_reqs: int, vocab_size: int, rank: int, dtype: torch.dtype
    ):
        assert rank & (rank - 1) == 0, rank
        self.vocab_size, self.rank = vocab_size, rank
        self.num_blocks = triton.cdiv(vocab_size, _BLOCK_SIZE)
        self.block_p = triton.next_power_of_2(self.num_blocks)
        self.markov_embed = torch.empty(max_num_reqs, rank, dtype=dtype)
        self.partial_max = torch.empty(max_num_reqs, self.block_p, dtype=torch.float32)
        self.partial_sum = torch.empty_like(self.partial_max)

    def __call__(self, spec, num_reqs: int, head_hidden: torch.Tensor) -> bool:
        model, n_spec = spec.model, spec.num_speculative_steps
        base_logits = model.compute_draft_logits(
            head_hidden[spec.sample_indices[: num_reqs * n_spec]]
        )
        markov_embed = model.markov_embed(
            spec.input_buffers.input_ids[spec._anchor_idx[:num_reqs]]
        )
        consts = dict(
            NUM_STEPS=n_spec,
            VOCAB=self.vocab_size,
            BLOCK_SIZE=_BLOCK_SIZE,
            BLOCK_P=self.block_p,
        )
        for step in range(n_spec):
            bias = model.markov_bias(markov_embed)
            _markov_sample_kernel[(num_reqs, self.num_blocks)](
                base_logits,
                base_logits.stride(0),
                bias,
                bias.stride(0),
                step,
                spec.sample_idx_mapping,
                spec.temperature,
                spec.draft_logits.values,
                self.partial_max,
                self.partial_sum,
                **consts,
                num_warps=2,
            )
            _markov_select_kernel[(num_reqs,)](
                step,
                spec.sample_idx_mapping,
                spec.sample_pos,
                spec.temperature,
                spec.seeds,
                spec.draft_logits.values,
                spec.draft_tokens,
                spec.draft_tokens.stride(0),
                model.model.markov_head.markov_w1.weight,
                self.markov_embed,
                self.partial_max,
                self.partial_sum,
                RANK=self.rank,
                **consts,
            )
            markov_embed = self.markov_embed[:num_reqs]
        return True