Skip to content

vllm.model_executor.warmup.step5_dsa_warmup

Warm Step5 DSA scored prefill before serving.

Startup dummy runs are short, so every row selects its whole history and the indexer scoring path (summary packing, FP8 logits, top-k, page metadata and, for multi-group caches, staged prefill pages) would otherwise JIT on the first prefill longer than top-k regions. Under DP every rank stalls on that rank.

Functions:

run_step5_scored_prefill_warmup(worker, topk)

Run one mixed step whose prefill row scores more than topk regions.

Source code in vllm/model_executor/warmup/step5_dsa_warmup.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def run_step5_scored_prefill_warmup(worker: "Worker", topk: int) -> bool:
    """Run one mixed step whose prefill row scores more than ``topk`` regions."""
    runner = worker.model_runner
    # One region past "select every region" (top-k x 8 tokens), plus the
    # decode token of the mixed step.
    num_tokens = topk * 8 + 17
    if num_tokens > worker.scheduler_config.max_num_batched_tokens or (
        num_tokens > runner.max_model_len
    ):
        logger.warning(
            "Skipping Step5 DSA scored-prefill warmup: needs %d tokens per step.",
            num_tokens,
        )
        return False
    logger.info("Warming up Step5 DSA scored prefill with %d tokens.", num_tokens)
    if runner.max_num_reqs >= 2:
        run_mixed_prefill_decode_warmup(
            cast("GPUModelRunner", runner),
            worker.execute_model,
            worker.sample_tokens,
            num_tokens,
            req_id_prefix="_step5_dsa_warmup",
        )
    else:
        runner._dummy_run(num_tokens=num_tokens, skip_eplb=True, is_profile=True)
    return True