vllm.models.kimi_k3.nvidia.specialized_decode.drafter.markov

DSpark's sequential Markov sampling, one launch per draft step.

DSparkSpeculator._sample_sequential runs, for each of the k draft steps, the Markov embedding of the previous token, the bias GEMV over the whole vocabulary (markov_w2: 163840 x 256 BF16, 84 MB, replicated), the add to the base logits ([R, 163840] written twice), the Gumbel sample (its kernel, a torch argmax over the block maxima and a gather) and the copy into draft_tokens: about nine launches and three full-vocabulary [R, V] tensors per step.

Here each step is two launches. The step kernel's programs each own 256 vocabulary entries and 16 request rows: they gather the rows' Markov embeddings (markov_w1[prev]), stream their markov_w2 range in [64, rank] tiles and multiply all rows at once on the tensor cores (FP32 accumulation), add the base logits, write the step's logits into the draft logits cache and keep each row's Gumbel-max over their range. A select launch (one program per row) takes the first of the largest per-program maxima, which is the step's token: it writes draft_tokens and the next step's previous token.

Exactness: the same sampling as vLLM's loop. The bias is the FP32 dot product of the embedding and the markov_w2 row rounded once to the logits dtype (the GEMV output dtype), the add is one FP32 add rounded to the logits dtype (a BF16 tensor add), the cache holds the logits before temperature, and the Gumbel noise is vLLM's own draw (murmur3 keyed by the request seed mixed with the draft step, the position before the sampled token plus the draft salt, and the token id), divided by the temperature exactly as gumbel_noised_argmax does. Only the summation order of the 256-term bias dot product differs from cuBLAS's GEMV, as any GEMM tiling may; ties keep the first token as vLLM's two-level argmax does.

Classes:

• FusedMarkov –

  The fused Markov steps of one speculator (scratch sized for its largest

Functions:

• sample_sequential –

  DSparkSpeculator._sample_sequential with the fused launches. Returns

• why_not –

  Why this draft step takes vLLM's loop, or None when the fused launches

FusedMarkov

The fused Markov steps of one speculator (scratch sized for its largest batch, allocated once so captured graphs keep their addresses).

Methods:

• fork_noise –

  Launch this draft step's Gumbel noise on a side stream (it depends only

• noise_rows –

  Whether this batch's Gumbel noise is precomputed (FP32 noise, at most

• run –

  base_logits: [num_reqs, k, V]; writes the draft logits cache and

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/drafter/markov.py

class FusedMarkov:
    """The fused Markov steps of one speculator (scratch sized for its largest
    batch, allocated once so captured graphs keep their addresses)."""

    def __init__(self, vocab, rank, max_rows, device, use_fp64):
        self.vocab, self.rank, self.use_fp64 = vocab, rank, use_fp64
        self.vb = VB
        self.nb = triton.cdiv(vocab, self.vb)
        self.nb_pow2 = triton.next_power_of_2(self.nb)
        dt = torch.float64 if use_fp64 else torch.float32
        self.bufs = (
            torch.empty(max_rows, self.nb, dtype=dt, device=device),
            torch.empty(max_rows, self.nb, dtype=torch.int64, device=device),
        )
        self.prev = torch.zeros(max_rows, dtype=torch.int64, device=device)
        self.max_rows = max_rows
        self.device = device
        self.noise = None  # [NOISE_ROWS * V] FP32, at the first noise launch
        self.stream = None

    def noise_rows(self, num_reqs, n_spec):
        """Whether this batch's Gumbel noise is precomputed (FP32 noise, at most
        ``NOISE_ROWS`` sampled rows)."""
        return not self.use_fp64 and 1 <= num_reqs * n_spec <= NOISE_ROWS

    def fork_noise(self, spec, num_reqs, n_spec):
        """Launch this draft step's Gumbel noise on a side stream (it depends only
        on the seeds, the positions and the token ids, not on any logits), so it
        runs beside the lm_head; ``join`` before the steps read it."""
        if self.noise is None:
            self.noise = torch.empty(
                NOISE_ROWS * self.vocab, dtype=torch.float32, device=self.device
            )
            self.stream = torch.cuda.Stream(device=self.device)
        main = torch.cuda.current_stream(self.device)
        assert self.stream is not None
        self.stream.wait_stream(main)
        noise = self.noise[: num_reqs * n_spec * self.vocab].view(
            num_reqs, n_spec, self.vocab
        )
        idx_map = spec.sample_idx_mapping[: num_reqs * n_spec].view(num_reqs, n_spec)
        sample_pos = spec.sample_pos[: num_reqs * n_spec].view(num_reqs, n_spec)
        with torch.cuda.stream(self.stream):
            grid = (triton.cdiv(self.vocab, NOISE_BLOCK), num_reqs * n_spec)
            _markov_noise_kernel[grid](
                noise,
                noise.stride(0),
                noise.stride(1),
                idx_map,
                idx_map.stride(0),
                sample_pos,
                sample_pos.stride(0),
                spec.seeds,
                self.vocab,
                K=n_spec,
                BLOCK=NOISE_BLOCK,
                SALT=_DRAFT_NOISE_SALT,
                num_warps=4,
            )
        return noise

    def join(self):
        torch.cuda.current_stream(self.device).wait_stream(self.stream)

    def run(self, spec, num_reqs, base_logits, w1, w2, noise=None):
        """``base_logits``: [num_reqs, k, V]; writes the draft logits cache and
        ``spec.draft_tokens[:num_reqs]``. ``noise``: ``fork_noise``'s (joined), or
        None to hash the noise in the steps."""
        n_spec = base_logits.shape[1]
        use_noise = noise is not None
        if noise is None:
            noise = self.prev  # unused pointer (NOISE=False)
            noise_strides = (0, 0)
        else:
            noise_strides = (noise.stride(0), noise.stride(1))
        rb = RB
        grid = (triton.cdiv(num_reqs, rb), self.nb)
        idx_map = spec.sample_idx_mapping[: num_reqs * n_spec].view(num_reqs, n_spec)
        sample_pos = spec.sample_pos[: num_reqs * n_spec].view(num_reqs, n_spec)
        cache = spec.draft_logits.values
        tokens = spec.draft_tokens
        for step in range(n_spec):
            out_max, out_arg = self.bufs
            _markov_step_kernel[grid](
                self.prev,
                spec.input_buffers.input_ids,
                spec._anchor_idx,
                w1,
                w2,
                base_logits[:, step],
                base_logits.stride(0),
                cache,
                cache.stride(0),
                cache.stride(1),
                spec.draft_logits.row_mapping,
                idx_map,
                idx_map.stride(0),
                sample_pos,
                sample_pos.stride(0),
                spec.temperature,
                spec.seeds,
                out_max,
                out_arg,
                noise,
                *noise_strides,
                num_reqs,
                self.vocab,
                self.nb,
                STEP=step,
                RANK=self.rank,
                VB=self.vb,
                VT=VT,
                RB=rb,
                USE_FP64=self.use_fp64,
                SALT=_DRAFT_NOISE_SALT,
                NOISE=use_noise,
                num_warps=4,
                num_stages=2,
            )
            _markov_select_kernel[(num_reqs,)](
                out_max,
                out_arg,
                self.prev,
                tokens,
                tokens.stride(0),
                num_reqs,
                self.nb,
                STEP=step,
                NB_POW2=self.nb_pow2,
            )

fork_noise(spec, num_reqs, n_spec)

Launch this draft step's Gumbel noise on a side stream (it depends only on the seeds, the positions and the token ids, not on any logits), so it runs beside the lm_head; join before the steps read it.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/drafter/markov.py

def fork_noise(self, spec, num_reqs, n_spec):
    """Launch this draft step's Gumbel noise on a side stream (it depends only
    on the seeds, the positions and the token ids, not on any logits), so it
    runs beside the lm_head; ``join`` before the steps read it."""
    if self.noise is None:
        self.noise = torch.empty(
            NOISE_ROWS * self.vocab, dtype=torch.float32, device=self.device
        )
        self.stream = torch.cuda.Stream(device=self.device)
    main = torch.cuda.current_stream(self.device)
    assert self.stream is not None
    self.stream.wait_stream(main)
    noise = self.noise[: num_reqs * n_spec * self.vocab].view(
        num_reqs, n_spec, self.vocab
    )
    idx_map = spec.sample_idx_mapping[: num_reqs * n_spec].view(num_reqs, n_spec)
    sample_pos = spec.sample_pos[: num_reqs * n_spec].view(num_reqs, n_spec)
    with torch.cuda.stream(self.stream):
        grid = (triton.cdiv(self.vocab, NOISE_BLOCK), num_reqs * n_spec)
        _markov_noise_kernel[grid](
            noise,
            noise.stride(0),
            noise.stride(1),
            idx_map,
            idx_map.stride(0),
            sample_pos,
            sample_pos.stride(0),
            spec.seeds,
            self.vocab,
            K=n_spec,
            BLOCK=NOISE_BLOCK,
            SALT=_DRAFT_NOISE_SALT,
            num_warps=4,
        )
    return noise

noise_rows(num_reqs, n_spec)

Whether this batch's Gumbel noise is precomputed (FP32 noise, at most NOISE_ROWS sampled rows).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/drafter/markov.py

def noise_rows(self, num_reqs, n_spec):
    """Whether this batch's Gumbel noise is precomputed (FP32 noise, at most
    ``NOISE_ROWS`` sampled rows)."""
    return not self.use_fp64 and 1 <= num_reqs * n_spec <= NOISE_ROWS

run(spec, num_reqs, base_logits, w1, w2, noise=None)

base_logits: [num_reqs, k, V]; writes the draft logits cache and spec.draft_tokens[:num_reqs]. noise: fork_noise's (joined), or None to hash the noise in the steps.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/drafter/markov.py

def run(self, spec, num_reqs, base_logits, w1, w2, noise=None):
    """``base_logits``: [num_reqs, k, V]; writes the draft logits cache and
    ``spec.draft_tokens[:num_reqs]``. ``noise``: ``fork_noise``'s (joined), or
    None to hash the noise in the steps."""
    n_spec = base_logits.shape[1]
    use_noise = noise is not None
    if noise is None:
        noise = self.prev  # unused pointer (NOISE=False)
        noise_strides = (0, 0)
    else:
        noise_strides = (noise.stride(0), noise.stride(1))
    rb = RB
    grid = (triton.cdiv(num_reqs, rb), self.nb)
    idx_map = spec.sample_idx_mapping[: num_reqs * n_spec].view(num_reqs, n_spec)
    sample_pos = spec.sample_pos[: num_reqs * n_spec].view(num_reqs, n_spec)
    cache = spec.draft_logits.values
    tokens = spec.draft_tokens
    for step in range(n_spec):
        out_max, out_arg = self.bufs
        _markov_step_kernel[grid](
            self.prev,
            spec.input_buffers.input_ids,
            spec._anchor_idx,
            w1,
            w2,
            base_logits[:, step],
            base_logits.stride(0),
            cache,
            cache.stride(0),
            cache.stride(1),
            spec.draft_logits.row_mapping,
            idx_map,
            idx_map.stride(0),
            sample_pos,
            sample_pos.stride(0),
            spec.temperature,
            spec.seeds,
            out_max,
            out_arg,
            noise,
            *noise_strides,
            num_reqs,
            self.vocab,
            self.nb,
            STEP=step,
            RANK=self.rank,
            VB=self.vb,
            VT=VT,
            RB=rb,
            USE_FP64=self.use_fp64,
            SALT=_DRAFT_NOISE_SALT,
            NOISE=use_noise,
            num_warps=4,
            num_stages=2,
        )
        _markov_select_kernel[(num_reqs,)](
            out_max,
            out_arg,
            self.prev,
            tokens,
            tokens.stride(0),
            num_reqs,
            self.nb,
            STEP=step,
            NB_POW2=self.nb_pow2,
        )

_fused(spec, w1, w2, device)

The speculator's FusedMarkov (markov_w2's vocabulary; sample_sequential checks it against the base logits').

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/drafter/markov.py

def _fused(spec, w1, w2, device):
    """The speculator's FusedMarkov (markov_w2's vocabulary; ``sample_sequential``
    checks it against the base logits')."""
    vocab, rank = w2.shape[0], w1.shape[1]
    key = (device, vocab, rank, bool(spec.use_fp64_gumbel))
    fused = _FUSED.get(key)
    if fused is None or fused.max_rows < spec.max_num_reqs:
        fused = FusedMarkov(vocab, rank, spec.max_num_reqs, device, key[-1])
        _FUSED[key] = fused
        logger.info_once(
            "DSpark fused Markov sampling: vocabulary %d, rank %d, %d programs",
            vocab,
            rank,
            fused.nb,
        )
    return fused

_gumbel_key(lf, temp, seed, pos, v, USE_FP64)

gumbel_noised_argmax's key of FP32 logits lf (temperature applied in FP32, then the draft-salted noise; plain logits at temperature 0).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/drafter/markov.py

@triton.jit
def _gumbel_key(lf, temp, seed, pos, v, USE_FP64: tl.constexpr):
    """``gumbel_noised_argmax``'s key of FP32 logits ``lf`` (temperature applied
    in FP32, then the draft-salted noise; plain logits at temperature 0)."""
    noisy = temp != 0.0
    lf = tl.where(noisy, lf / temp, lf)
    if USE_FP64:
        lf = lf.to(tl.float64)
        u = murmur3_uniform64(seed, pos, v)
        g = -tl.log(-tl.log(u))
    else:
        u = murmur3_uniform32(seed, pos, v)
        g = -tl.log(-_log1p_neg_stable(u))
    return tl.where(noisy, lf + g, lf)

_gumbel_log_arg32(seed, pos, v)

The argument of the outer log of _gumbel_key's FP32 noise -log(-log1p(-u)): the noise is -log of it. The step kernel takes that last log itself, so its add to the logits compiles as in gumbel_noised_argmax (a stored noise, added in another kernel, rounds differently: 1 ulp on about a fifth of the keys).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/drafter/markov.py

@triton.jit
def _gumbel_log_arg32(seed, pos, v):
    """The argument of the outer log of ``_gumbel_key``'s FP32 noise
    ``-log(-log1p(-u))``: the noise is ``-log`` of it. The step kernel takes that
    last log itself, so its add to the logits compiles as in
    ``gumbel_noised_argmax`` (a stored noise, added in another kernel, rounds
    differently: 1 ulp on about a fifth of the keys)."""
    u = murmur3_uniform32(seed, pos, v)
    return -_log1p_neg_stable(u)

_loop_rest(spec, model, num_reqs, base_logits)

The loop of vLLM from its base logits on (base logits computed once).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/drafter/markov.py

def _loop_rest(spec, model, num_reqs, base_logits) -> bool:
    """The loop of vLLM from its base logits on (base logits computed once)."""
    n_spec = spec.num_speculative_steps
    num_sample = num_reqs * n_spec
    base_logits = base_logits.view(num_reqs, n_spec, -1)
    idx_map = spec.sample_idx_mapping[:num_sample].view(num_reqs, n_spec)
    sample_pos = spec.sample_pos[:num_sample].view(num_reqs, n_spec)
    prev = spec.input_buffers.input_ids[spec._anchor_idx[:num_reqs]]
    for i in range(n_spec):
        bias = model.markov_bias(model.markov_embed(prev))
        draft = spec._sample_logits(
            base_logits[:, i] + bias, idx_map[:, i], sample_pos[:, i], i
        )
        spec.draft_tokens[:num_reqs, i] = draft
        prev = draft
    return True

sample_sequential(spec, model, num_reqs, head_hidden)

DSparkSpeculator._sample_sequential with the fused launches. Returns False, having done nothing, when this step takes vLLM's loop.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/drafter/markov.py

def sample_sequential(spec, model, num_reqs: int, head_hidden: torch.Tensor) -> bool:
    """``DSparkSpeculator._sample_sequential`` with the fused launches. Returns
    False, having done nothing, when this step takes vLLM's loop."""
    why = why_not(spec, model, num_reqs)
    if why is not None:
        logger.info_once("DSpark fused Markov sampling off: %s", why)
        return False
    n_spec = spec.num_speculative_steps
    num_sample = num_reqs * n_spec
    head = model.model.markov_head
    w1, w2 = head.markov_w1.weight, head.markov_w2.weight
    fused = _fused(spec, w1, w2, head_hidden.device)
    noise = None
    if fused.noise_rows(num_reqs, n_spec):
        noise = fused.fork_noise(spec, num_reqs, n_spec)
    sample_hidden = head_hidden[spec.sample_indices[:num_sample]]
    base_logits = model.compute_draft_logits(sample_hidden)
    if noise is not None:
        fused.join()
    cache = spec.draft_logits.values
    vocab = base_logits.shape[-1]
    if (
        base_logits.dtype not in (torch.bfloat16, torch.float32)
        or base_logits.dtype != cache.dtype
        or base_logits.stride(-1) != 1
        or w2.shape[0] != vocab
        or cache.shape[-1] < vocab
    ):
        logger.info_once(
            "DSpark fused Markov sampling off: %s base logits, %s cache",
            base_logits.dtype,
            cache.dtype,
        )
        return _loop_rest(spec, model, num_reqs, base_logits)
    fused.run(spec, num_reqs, base_logits.view(num_reqs, n_spec, vocab), w1, w2, noise)
    return True

why_not(spec, model, num_reqs)

Why this draft step takes vLLM's loop, or None when the fused launches compute the same sampling.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/drafter/markov.py

def why_not(spec, model, num_reqs):
    """Why this draft step takes vLLM's loop, or None when the fused launches
    compute the same sampling."""
    head = model.model.markov_head
    if spec.draft_logits is None:
        return "greedy drafts"
    if spec.draft_watermarker is not None:
        return "a draft watermarker"
    if spec._d2t_scatter_index is not None:
        return "a reduced draft vocabulary"
    if spec.use_confidence_head:
        return "the confidence head"
    if spec.acceptance_estimator is not None:
        return "the acceptance estimator"
    if hasattr(model, "apply_output_logit_transform"):
        return "an output logit transform"
    lp = model.logits_processor
    if lp.scale != 1.0 or lp.soft_cap is not None:
        return "scaled logits"
    w1, w2 = head.markov_w1.weight, head.markov_w2.weight
    if w1.dtype != torch.bfloat16 or w2.dtype != torch.bfloat16:
        return "Markov weights not bf16"
    rank = w1.shape[1]
    if w2.shape[1] != rank or rank % 16 or rank & (rank - 1) or rank > 512:
        return f"Markov rank {rank}"
    if not (w1.is_contiguous() and w2.is_contiguous()):
        return "Markov weight layout"
    if head.markov_w2.quant_method.__class__.__name__ != "UnquantizedEmbeddingMethod":
        return "a quantized Markov head"
    if num_reqs < 1:
        return "no requests"
    return None