vllm.models.deepseek_v41.amd.ops.fused_markov

DSpark's sequential Markov sampling with one launch a draft step.

vLLM's DSparkSpeculator._sample_sequential runs about 9 kernels for each of the 5 draft steps: the Markov embedding of the previous token, the bias GEMV over the whole vocabulary (66 MB of markov_w2 at 129280 x 256 bf16), the add to the base logits, the Gumbel sample (its kernel, the argmax over its blocks and a gather) and the copy into draft_tokens. On MI325X at TP 4 that is about 220 us a decode step.

Here each step is one launch over the vocabulary in blocks of BLOCK tokens. For each row, a program computes its block's bias from markov_w2, adds the base logits, stores the step's logits into the draft logits cache and takes the Gumbel-max of its block with vLLM's own gumbel_noised_argmax. The next step's programs first take the argmax over the previous step's block maxima themselves, which gives the previous token, and the last step's tokens get one small launch. The noise depends only on the request's seed, the draft step, the position and the token id, and both paths keep the first of equal maxima, so the sampled token does not depend on the block size.

The bias is an fp32 sum of 256 products rounded to bf16, as vLLM's GEMV output. The products are added in another order than in vLLM's GEMV, so a bias value can differ in its last bf16 bit. That changes only which tokens the draft proposes, not the target's output, because the rejection sampler verifies every draft token against the target's own distribution.

It runs on gfx942 with --monokernel (ops.enabled()). The speculator asks the draft model's sample_sequential (dspark.py) before its loop, and that calls sample_sequential here. tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py compares it with vLLM's loop.

Functions:

• sample_sequential –

  DSparkSpeculator._sample_sequential with the fused launches. Returns

_loop_rest(spec, num_reqs, base_logits, sample_hidden)

VLLM's loop from its base logits on, for a step whose base logits the fused launch does not take. It is the loop of _sample_sequential after compute_draft_logits, so the step computes them only once.

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

def _loop_rest(spec, num_reqs, base_logits, sample_hidden) -> bool:
    """VLLM's loop from its base logits on, for a step whose base logits the
    fused launch does not take. It is the loop of _sample_sequential after
    compute_draft_logits, so the step computes them only once."""
    n_spec = spec.num_speculative_steps
    num_sample = num_reqs * n_spec
    vocab_size = base_logits.shape[-1]
    base_logits = base_logits.view(num_reqs, n_spec, vocab_size)
    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 = spec.model.markov_bias(spec.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

_scratch(device, vocab, use_fp64)

Two sets of block maxima and their token ids, one written by the current step and one read from the previous step. They are allocated once, so a CUDA graph keeps their addresses.

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

def _scratch(device, vocab, use_fp64):
    """Two sets of block maxima and their token ids, one written by the
    current step and one read from the previous step. They are allocated
    once, so a CUDA graph keeps their addresses."""
    key = (device, vocab, use_fp64, BLOCK)
    if key not in _buffers:
        nb = triton.cdiv(vocab, BLOCK)
        dt = torch.float64 if use_fp64 else torch.float32
        _buffers[key] = [
            (
                torch.empty(MAX_ROWS, nb, dtype=dt, device=device),
                torch.empty(MAX_ROWS, nb, dtype=torch.int64, device=device),
            )
            for _ in range(2)
        ]
    return _buffers[key]

_why_not(spec, num_reqs)

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

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

def _why_not(spec, num_reqs: int) -> str | None:
    """Why this speculator's step takes vLLM's loop, or None when the fused
    launches compute the same sampling."""
    model = spec.model
    head = getattr(getattr(model, "model", None), "markov_head", None)
    if head is None:
        return "no 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 not 1 <= num_reqs <= MAX_ROWS:
        return f"{num_reqs} requests"
    w1 = head.markov_w1.weight
    w2 = head.markov_w2.weight
    if w1.dtype != torch.bfloat16 or w2.dtype != torch.bfloat16:
        return "Markov weights not bf16"
    if w1.shape[1] != w2.shape[1] or not (w1.is_contiguous() and w2.is_contiguous()):
        return "Markov weight layout"
    if spec.draft_logits.dtype != torch.bfloat16:
        return f"a {spec.draft_logits.dtype} draft logits cache"
    if head.markov_w2.quant_method.__class__.__name__ != "UnquantizedEmbeddingMethod":
        return "a quantized Markov head"
    return None

sample_sequential(spec, 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/deepseek_v41/amd/ops/fused_markov.py

def sample_sequential(spec, 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."""
    if not enabled():
        return False
    why = _why_not(spec, num_reqs)
    if why is not None:
        logger.debug_once("DSpark fused Markov sampling off for a step: %s", why)
        return False
    n_spec = spec.num_speculative_steps
    num_sample = num_reqs * n_spec
    sample_hidden = head_hidden[spec.sample_indices[:num_sample]]
    base_logits = spec.model.compute_draft_logits(sample_hidden)
    if base_logits.dtype != torch.bfloat16 or base_logits.stride(-1) != 1:
        # The base logits come from the lm_head, which may differ from the
        # Markov head. vLLM's loop then adds them in their own dtype.
        logger.debug_once(
            "DSpark fused Markov sampling off: %s base logits", base_logits.dtype
        )
        return _loop_rest(spec, num_reqs, base_logits, sample_hidden)
    vocab = base_logits.shape[-1]
    base_logits = base_logits.view(num_reqs, n_spec, vocab)
    head = spec.model.model.markov_head
    w1, w2 = head.markov_w1.weight, head.markov_w2.weight
    rank = w1.shape[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)
    assert idx_map.stride(1) == 1 and sample_pos.stride(1) == 1
    cache = spec.draft_logits
    tokens = spec.draft_tokens
    use_fp64 = bool(spec.use_fp64_gumbel)
    bufs = _scratch(base_logits.device, vocab, use_fp64)
    nb = triton.cdiv(vocab, BLOCK)
    nb2 = triton.next_power_of_2(nb)
    logger.info_once(
        "DSpark fused Markov sampling: %d steps, vocabulary %d, rank %d, %d blocks",
        n_spec, vocab, rank, nb,
    )  # fmt: skip
    for step in range(n_spec):
        prev_max, prev_arg = bufs[(step + 1) % 2]
        out_max, out_arg = bufs[step % 2]
        _markov_step_kernel[(nb,)](
            prev_max,
            prev_arg,
            spec.input_buffers.input_ids,
            spec._anchor_idx,
            tokens,
            tokens.stride(0),
            w1,
            w2,
            base_logits[:, step],
            base_logits.stride(0),
            cache,
            cache.stride(0),
            cache.stride(1),
            idx_map,
            idx_map.stride(0),
            sample_pos,
            sample_pos.stride(0),
            spec.temperature,
            spec.seeds,
            out_max,
            out_arg,
            vocab,
            nb,
            STEP=step,
            ROWS=num_reqs,
            RANK=rank,
            BLOCK=BLOCK,
            K_CHUNK=K_CHUNK,
            BLOCKS_POW2=nb2,
            USE_FP64=use_fp64,
            num_warps=NUM_WARPS,
        )
    last_max, last_arg = bufs[(n_spec - 1) % 2]
    _markov_last_kernel[(num_reqs,)](
        last_max,
        last_arg,
        tokens,
        tokens.stride(0),
        nb,
        STEP=n_spec - 1,
        BLOCKS_POW2=nb2,
    )
    return True