vllm.v1.worker.gpu.input_batch ¶

def set_dummy_context(
    input_batch: InputBatch,
    block_tables: "BlockTables",
    context_len: int,
    num_kv_blocks: int,
    max_model_len: int,
) -> None:
    """Give each dummy request context_len of context, used when profiling step cost."""
    if not block_tables.input_block_tables:
        # Attention-free models have no KV context to fabricate.
        return
    num_reqs = input_batch.num_reqs
    query_len = input_batch.max_query_len or int(input_batch.num_scheduled_tokens.max())
    context_len = max(min(context_len, max_model_len - query_len), 0)
    if not context_len:
        return

    # Decode-like shape: each request continues after context_len
    # already-computed tokens.
    input_batch.seq_lens += context_len
    input_batch.seq_lens_cpu_upper_bound += context_len
    input_batch.num_computed_tokens_np.fill(context_len)
    input_batch.num_computed_prefill_tokens_np.fill(context_len)
    local_pos = np.arange(input_batch.num_tokens, dtype=np.int64) - np.repeat(
        input_batch.query_start_loc_np[:-1], input_batch.num_scheduled_tokens
    )
    input_batch.positions.copy_(torch.from_numpy(local_pos + context_len))

    seq_len = context_len + query_len
    for block_table, block_size, bpk in zip(
        block_tables.input_block_tables,
        block_tables.kernel_block_sizes,
        block_tables.blocks_per_kv_block,
    ):
        num_blocks = min(cdiv(seq_len, block_size), block_table.shape[1])
        # Spans are disjoint until the pool runs out, then they wrap and share
        # blocks: profiling only needs the reads to be realistic, not distinct.
        block_ids = torch.arange(
            num_reqs * num_blocks, dtype=block_table.dtype, device=block_table.device
        ) % (num_kv_blocks * bpk)
        block_table[:num_reqs, :num_blocks] = block_ids.view(num_reqs, num_blocks)