vllm.models.step5.amd.warmup

Functions:

• step5_csa_warmup –

  JIT the CSA scored-prefill kernels before the first long prompt.

step5_csa_warmup(worker)

JIT the CSA scored-prefill kernels before the first long prompt.

Source code in vllm/models/step5/amd/warmup.py

def step5_csa_warmup(worker: "Worker") -> None:
    """JIT the CSA scored-prefill kernels before the first long prompt."""
    context = worker.vllm_config.compilation_config.static_forward_context
    for layer in context.values():
        if isinstance(layer, Step5Attention) and layer.indexer is not None:
            run_step5_scored_prefill_warmup(worker, layer.indexer.topk)
            return