vllm.v1.worker.gpu.sample.watermark

Functions:

• draft_philox_gumbel_sample –

  Sample num_steps draft steps flattened as (request, step), in step order.

• philox_gumbel_block_argmax –

  Apply the detector-compatible PRF to one vocabulary block.

• repeated_context_mask –

  Return, per row, whether the row's context already occurred in its history.

_draft_step_kernel(sampled_ptr, sampled_stride, step, local_argmax_ptr, local_max_ptr, num_blocks, contexts_ptr, context_stride, skip_mask_ptr, next_req_indices_ptr, next_steps_ptr, temp_ptr, enabled_ptr, all_token_ids_ptr, all_token_ids_stride, prompt_lens_ptr, total_lens_ptr, prior_contexts_ptr, prior_contexts_stride_0, prior_contexts_stride_1, CONTEXT_WIDTH, NUM_SPECULATIVE_STEPS, MAX_HISTORY, INCLUDE_PROMPT, DEDUPLICATE, SAMPLE, PREPARE_NEXT, BLOCKS_POW2)

Finish draft step step (if SAMPLE) and build the next step's skip mask.

Source code in vllm/v1/worker/gpu/sample/watermark.py

@triton.jit(do_not_specialize=["step"])
def _draft_step_kernel(
    sampled_ptr,
    sampled_stride,
    step,
    local_argmax_ptr,
    local_max_ptr,
    num_blocks,
    contexts_ptr,
    context_stride,
    skip_mask_ptr,
    next_req_indices_ptr,
    next_steps_ptr,
    temp_ptr,
    enabled_ptr,
    all_token_ids_ptr,
    all_token_ids_stride,
    prompt_lens_ptr,
    total_lens_ptr,
    prior_contexts_ptr,
    prior_contexts_stride_0,
    prior_contexts_stride_1,
    CONTEXT_WIDTH: tl.constexpr,
    NUM_SPECULATIVE_STEPS: tl.constexpr,
    MAX_HISTORY: tl.constexpr,
    INCLUDE_PROMPT: tl.constexpr,
    DEDUPLICATE: tl.constexpr,
    SAMPLE: tl.constexpr,
    PREPARE_NEXT: tl.constexpr,
    BLOCKS_POW2: tl.constexpr,
):
    """Finish draft step `step` (if SAMPLE) and build the next step's skip mask."""
    row = tl.program_id(0).to(tl.int64)
    contexts_row_ptr = contexts_ptr + row * context_stride
    if SAMPLE:
        # Same pick as local_max.argmax(dim=-1): the first NaN, else the first max.
        blocks = tl.arange(0, BLOCKS_POW2)
        values = tl.load(
            local_max_ptr + row * num_blocks + blocks,
            mask=blocks < num_blocks,
            other=float("-inf"),
        )
        is_nan = values != values
        best = tl.max(tl.where(is_nan, float("-inf"), values), axis=0)
        first_nan = tl.min(tl.where(is_nan, blocks, BLOCKS_POW2), axis=0)
        first_max = tl.min(
            tl.where((values == best) & (blocks < num_blocks), blocks, BLOCKS_POW2),
            axis=0,
        )
        block = tl.where(first_nan < BLOCKS_POW2, first_nan, first_max)
        token = tl.load(local_argmax_ptr + row * num_blocks + block)
        tl.store(sampled_ptr + row * sampled_stride + step, token)

        offsets = tl.arange(0, BLOCKS_POW2)
        context = tl.load(
            contexts_row_ptr + offsets + 1,
            mask=offsets < CONTEXT_WIDTH - 1,
            other=token,
        )
        tl.debug_barrier()
        tl.store(contexts_row_ptr + offsets, context, mask=offsets < CONTEXT_WIDTH)
        tl.debug_barrier()

    if PREPARE_NEXT:
        req_idx = tl.load(next_req_indices_ptr + row).to(tl.int64)
        temperature = tl.load(temp_ptr + req_idx, mask=req_idx >= 0, other=0.0)
        active = tl.load(enabled_ptr + row) & (temperature != 0)
        if DEDUPLICATE:
            # The dedup kernel body runs in place on the skip mask buffer.
            tl.store(skip_mask_ptr + row, active)
            tl.debug_barrier()
            _repeated_context_mask_kernel(
                skip_mask_ptr,
                all_token_ids_ptr,
                all_token_ids_stride,
                next_req_indices_ptr,
                prompt_lens_ptr,
                total_lens_ptr,
                next_steps_ptr,
                1,
                contexts_ptr,
                context_stride,
                None,
                prior_contexts_ptr,
                prior_contexts_stride_0,
                prior_contexts_stride_1,
                next_steps_ptr,
                1,
                skip_mask_ptr,
                CONTEXT_WIDTH=CONTEXT_WIDTH,
                NUM_SPECULATIVE_STEPS=NUM_SPECULATIVE_STEPS,
                MAX_HISTORY=MAX_HISTORY,
                INCLUDE_PROMPT=INCLUDE_PROMPT,
                SKIP_PARTIAL_CONTEXT=False,
                HAS_HISTORY_OFFSETS=True,
                SPECULATIVE_CONTEXTS=False,
                DRAFT_CONTEXTS=True,
                BLOCK=512,
            )
            tl.debug_barrier()
            active = tl.load(skip_mask_ptr + row)
        tl.store(skip_mask_ptr + row, active == 0)

_repeated_context_mask_cpu(all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=None, include_prompt=False, skip_partial_context=False, history_offsets=None)

Reference implementation; the parity tests check the Triton kernel against it.

Source code in vllm/v1/worker/gpu/sample/watermark.py

def _repeated_context_mask_cpu(
    all_token_ids: torch.Tensor,
    req_indices: torch.Tensor,
    prompt_lens: torch.Tensor,
    total_lens: torch.Tensor,
    contexts: torch.Tensor,
    max_history: int | None = None,
    include_prompt: bool = False,
    skip_partial_context: bool = False,
    history_offsets: torch.Tensor | None = None,
) -> torch.Tensor:
    """Reference implementation; the parity tests check the Triton kernel against it."""
    repeated = torch.zeros(len(req_indices), dtype=torch.bool)
    for row, req_idx_tensor in enumerate(req_indices):
        req_idx = int(req_idx_tensor)
        if req_idx < 0:
            continue
        prompt_len = int(prompt_lens[req_idx])
        total_len = int(total_lens[req_idx])
        if skip_partial_context and total_len - prompt_len < contexts.shape[-1]:
            repeated[row] = True
            continue
        sequence_start = 0 if include_prompt else prompt_len
        history_tokens = all_token_ids[req_idx, sequence_start:total_len].tolist()
        prefix = [-1] * contexts.shape[-1]
        current_context = tuple(contexts[row].tolist())
        history_offset = 0 if history_offsets is None else int(history_offsets[row])
        history_start = (
            max(0, len(history_tokens) - max(0, max_history - history_offset))
            if max_history is not None
            else 0
        )
        for history_pos, token_id in enumerate(history_tokens):
            if history_pos >= history_start and (
                tuple(prefix[-contexts.shape[-1] :]) == current_context
            ):
                repeated[row] = True
                break
            prefix.append(token_id)
    return repeated

draft_philox_gumbel_sample(logits, contexts, key, *, num_steps, expanded_idx_mapping, temperatures, seeds, positions, enabled, logits_cache, logits_cache_col, use_fp64, prior_contexts, all_token_ids, prompt_lens, total_lens, deduplicate, max_history, include_prompt)

Sample num_steps draft steps flattened as (request, step), in step order.

Bit-exact with sampling each step by DraftWatermarker.sample, but with two kernel launches per step: the Philox block argmax (with the temperature divide and the logits cache write), then a per-row kernel that picks the token, shifts contexts and builds the next step's skip mask. Padding rows (index -1) read temperature 0 and are never watermarked, so they may differ.

Source code in vllm/v1/worker/gpu/sample/watermark.py

def draft_philox_gumbel_sample(
    logits: torch.Tensor,
    contexts: torch.Tensor,
    key: int,
    *,
    num_steps: int,
    expanded_idx_mapping: torch.Tensor,
    temperatures: torch.Tensor,
    seeds: torch.Tensor,
    positions: torch.Tensor,
    enabled: torch.Tensor,
    logits_cache: torch.Tensor,
    logits_cache_col: torch.Tensor,
    use_fp64: bool,
    prior_contexts: torch.Tensor,
    all_token_ids: torch.Tensor | None,
    prompt_lens: torch.Tensor,
    total_lens: torch.Tensor | None,
    deduplicate: bool,
    max_history: int | None,
    include_prompt: bool,
) -> torch.Tensor:
    """Sample `num_steps` draft steps flattened as (request, step), in step order.

    Bit-exact with sampling each step by `DraftWatermarker.sample`, but with two
    kernel launches per step: the Philox block argmax (with the temperature
    divide and the logits cache write), then a per-row kernel that picks the
    token, shifts `contexts` and builds the next step's skip mask. Padding rows
    (index -1) read temperature 0 and are never watermarked, so they may differ.
    """
    num_rows, vocab_size = logits.shape
    num_reqs = num_rows // num_steps
    assert logits.stride(-1) == 1 and contexts.stride(-1) == 1
    assert logits.dtype == logits_cache.dtype, (
        "logits cache source and destination must have the same dtype"
    )
    assert logits_cache.size(-1) >= vocab_size, (
        f"draft logits cache vocab dim ({logits_cache.size(-1)}) is narrower "
        f"than the sampled logits ({vocab_size}). Cached logits would be truncated."
    )
    assert not deduplicate or (all_token_ids is not None and total_lens is not None), (
        "context deduplication requires all_token_ids and total_lens"
    )
    logits = logits.view(num_reqs, num_steps, vocab_size)
    # Step-major copies, so each step reads contiguous per-request slices.
    expanded_idx_mapping, positions, logits_cache_col = (
        tensor.view(num_reqs, num_steps).t().contiguous()
        for tensor in (expanded_idx_mapping, positions, logits_cache_col)
    )
    block_size = 1024
    num_blocks = triton.cdiv(vocab_size, block_size)
    local_argmax = logits.new_empty(num_reqs, num_blocks, dtype=torch.int64)
    local_max = logits.new_empty(
        num_reqs, num_blocks, dtype=torch.float64 if use_fp64 else torch.float32
    )
    skip_mask = torch.empty(num_reqs, dtype=torch.bool, device=logits.device)
    sampled = logits.new_empty(num_reqs, num_steps, dtype=torch.int64)

    def finish_step(step: int) -> None:
        next_step = min(step + 1, num_steps - 1)
        _draft_step_kernel[(num_reqs,)](
            sampled,
            sampled.stride(0),
            step,
            local_argmax,
            local_max,
            num_blocks,
            contexts,
            contexts.stride(0),
            skip_mask,
            expanded_idx_mapping[next_step],
            logits_cache_col[next_step],
            temperatures,
            enabled,
            all_token_ids,
            0 if all_token_ids is None else all_token_ids.stride(0),
            prompt_lens,
            total_lens,
            prior_contexts,
            prior_contexts.stride(0),
            prior_contexts.stride(1),
            CONTEXT_WIDTH=contexts.shape[-1],
            NUM_SPECULATIVE_STEPS=prior_contexts.shape[1],
            MAX_HISTORY=0 if max_history is None else max_history,
            INCLUDE_PROMPT=include_prompt,
            DEDUPLICATE=deduplicate,
            SAMPLE=step >= 0,
            PREPARE_NEXT=step + 1 < num_steps,
            BLOCKS_POW2=triton.next_power_of_2(max(num_blocks, contexts.shape[-1])),
        )

    finish_step(-1)
    for step in range(num_steps):
        _philox_gumbel_kernel[(num_reqs, num_blocks)](
            local_argmax,
            local_argmax.stride(0),
            local_max,
            local_max.stride(0),
            logits[:, step],
            logits.stride(0),
            contexts,
            contexts.stride(0),
            skip_mask,
            expanded_idx_mapping[step],
            seeds,
            positions[step],
            temperatures,
            logits_cache,
            logits_cache.stride(0),
            logits_cache.stride(1),
            logits_cache_col[step],
            logits[:, step],
            logits.stride(0),
            key & _UINT32_MASK_VALUE,
            key >> 32,
            vocab_size,
            CONTEXT_WIDTH=contexts.shape[-1],
            BLOCK_SIZE=block_size,
            USE_FP64=use_fp64,
            IS_DRAFTING=True,
            PER_TOKEN_COL=True,
            APPLY_TEMPERATURE=True,
        )
        finish_step(step)
    return sampled.view(-1)

philox_gumbel_block_argmax(logits, mask, block_idx, contexts_row_ptr, key_0, key_1, CONTEXT_WIDTH, BLOCK_SIZE)

Apply the detector-compatible PRF to one vocabulary block.

Source code in vllm/v1/worker/gpu/sample/watermark.py

@triton.jit
def philox_gumbel_block_argmax(
    logits,
    mask,
    block_idx,
    contexts_row_ptr,
    key_0,
    key_1,
    CONTEXT_WIDTH: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    """Apply the detector-compatible PRF to one vocabulary block."""
    tl.static_assert(BLOCK_SIZE % 4 == 0)
    state_0, state_1, state_2, state_3 = _philox_context_state(
        contexts_row_ptr, key_0, key_1, CONTEXT_WIDTH
    )
    groups = block_idx * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4)
    output_0, output_1, output_2, output_3 = _philox_candidate_words(
        groups, state_0, state_1, state_2, state_3, key_0, key_1
    )
    # Lay the per-group words out in token order: 4g+0, 4g+1, 4g+2, 4g+3.
    words = tl.interleave(
        tl.interleave(output_0, output_2),
        tl.interleave(output_1, output_3),
    )
    values = tl.where(
        mask,
        _philox_gumbel_from_logits(logits.to(tl.float32), words),
        float("-inf"),
    )
    return tl.max(values, axis=0, return_indices=True)

repeated_context_mask(all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=None, include_prompt=False, skip_partial_context=False, history_offsets=None, local_positions=None, num_speculative_steps=0)

Return, per row, whether the row's context already occurred in its history.

Parameters:

• all_token_ids

  (Tensor) –

  [max_num_reqs, max_model_len] token ids of every request.

• req_indices

  (Tensor) –

  Request slot per sampled row; -1 marks a padding row, which is reported as not repeated.

• prompt_lens

  (Tensor) –

  Prompt length per request slot.

• total_lens

  (Tensor) –

  Prompt plus generated length per request slot.

• contexts

  (Tensor) –

  [num_rows, context_width] context per row, padded with -1 before the start of the scanned history.

• max_history

  (int | None, default: None ) –

  Number of most recent history positions searched, or None for all of them. The window compared at each position reaches context_width tokens further back.

• include_prompt

  (bool, default: False ) –

  Search the prompt as well as the generated tokens.

• skip_partial_context

  (bool, default: False ) –

  Mark contexts containing start padding so they use ordinary sampling.

• history_offsets

  (Tensor | None, default: None ) –

  Number of newer, non-committed positions preceding each context. These positions count toward max_history. Must be 1-D over rows; a stride-0 broadcast row is fine.

• local_positions

  (Tensor | None, default: None ) –

  Position within each request's speculative block. When provided, compare each context with earlier contexts in that block.

• num_speculative_steps

  (int, default: 0 ) –

  Maximum number of draft tokens in the block.

Source code in vllm/v1/worker/gpu/sample/watermark.py

def repeated_context_mask(
    all_token_ids: torch.Tensor,
    req_indices: torch.Tensor,
    prompt_lens: torch.Tensor,
    total_lens: torch.Tensor,
    contexts: torch.Tensor,
    max_history: int | None = None,
    include_prompt: bool = False,
    skip_partial_context: bool = False,
    history_offsets: torch.Tensor | None = None,
    local_positions: torch.Tensor | None = None,
    num_speculative_steps: int = 0,
) -> torch.Tensor:
    """Return, per row, whether the row's context already occurred in its history.

    Args:
        all_token_ids: `[max_num_reqs, max_model_len]` token ids of every request.
        req_indices: Request slot per sampled row; -1 marks a padding row, which
            is reported as not repeated.
        prompt_lens: Prompt length per request slot.
        total_lens: Prompt plus generated length per request slot.
        contexts: `[num_rows, context_width]` context per row, padded with -1
            before the start of the scanned history.
        max_history: Number of most recent history positions searched, or
            ``None`` for all of them. The window compared at each position
            reaches `context_width` tokens further back.
        include_prompt: Search the prompt as well as the generated tokens.
        skip_partial_context: Mark contexts containing start padding so they use
            ordinary sampling.
        history_offsets: Number of newer, non-committed positions preceding each
            context. These positions count toward `max_history`. Must be 1-D
            over rows; a stride-0 broadcast row is fine.
        local_positions: Position within each request's speculative block. When
            provided, compare each context with earlier contexts in that block.
        num_speculative_steps: Maximum number of draft tokens in the block.

    """
    if max_history is not None and max_history < 1:
        raise ValueError("max_history must be positive or None")
    if (local_positions is None) != (num_speculative_steps == 0):
        raise ValueError(
            "local_positions and positive num_speculative_steps must be "
            "provided together"
        )
    if all_token_ids.device.type == "cpu":
        repeated = _repeated_context_mask_cpu(
            all_token_ids,
            req_indices,
            prompt_lens,
            total_lens,
            contexts,
            max_history,
            include_prompt,
            skip_partial_context,
            history_offsets,
        )
        if local_positions is not None:
            max_offset = num_speculative_steps + 1
            if max_history is not None:
                max_offset = min(max_offset, max_history + 1)
            for offset in range(1, max_offset):
                prior_contexts = torch.cat(
                    (contexts[:offset], contexts[:-offset]), dim=0
                )
                prior_requests = torch.cat(
                    (req_indices[:offset], req_indices[:-offset]), dim=0
                )
                prior_local_pos = torch.cat(
                    (local_positions[:offset], local_positions[:-offset]), dim=0
                )
                repeated |= (
                    (req_indices >= 0)
                    & (local_positions >= offset)
                    & (req_indices == prior_requests)
                    & (local_positions == prior_local_pos + offset)
                    & (contexts == prior_contexts).all(dim=-1)
                )
            if include_prompt:
                repeated |= (contexts < 0).any(dim=-1)
        return repeated

    if contexts.stride(-1) != 1:
        contexts = contexts.contiguous()
    # Only history_offsets passes its row stride to the kernel.
    req_indices = req_indices.contiguous()
    if local_positions is not None:
        local_positions = local_positions.contiguous()
    # `prompt_lens` and `total_lens` are indexed flat by request slot, and
    # `all_token_ids` takes a row stride but is read flat within the row. Those
    # three are contiguous by contract; every caller owns them outright.
    repeated = torch.empty(len(req_indices), dtype=torch.bool, device=contexts.device)
    _repeated_context_mask_kernel[(len(req_indices),)](
        repeated,
        all_token_ids,
        all_token_ids.stride(0),
        req_indices,
        prompt_lens,
        total_lens,
        history_offsets,
        0 if history_offsets is None else history_offsets.stride(0),
        contexts,
        contexts.stride(0),
        local_positions,
        None,
        0,
        0,
        None,
        0,
        None,
        CONTEXT_WIDTH=contexts.shape[-1],
        NUM_SPECULATIVE_STEPS=num_speculative_steps,
        MAX_HISTORY=0 if max_history is None else max_history,
        INCLUDE_PROMPT=include_prompt,
        SKIP_PARTIAL_CONTEXT=skip_partial_context,
        HAS_HISTORY_OFFSETS=history_offsets is not None,
        SPECULATIVE_CONTEXTS=local_positions is not None,
        DRAFT_CONTEXTS=False,
        BLOCK=512,
    )
    return repeated