Skip to content

vllm.model_executor.warmup.cutedsl_warmup

Run registered CuTeDSL warmup compile units.

Functions:

cutedsl_warmup()

Run CuTeDSL compile providers before serving.

Source code in vllm/model_executor/warmup/cutedsl_warmup.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
@instrument(span_name="CuTeDSL warmup")
def cutedsl_warmup() -> None:
    """Run CuTeDSL compile providers before serving."""
    if not current_platform.is_cuda():
        logger.info("Skipping CuTeDSL warmup on non-CUDA platform.")
        return

    compile_units = _collect_unique_compile_units(_iter_cutedsl_warmup_compile_units())
    if not compile_units:
        logger.info("Skipping CuTeDSL warmup because no compile units were requested.")
        return

    unit_names = list(dict.fromkeys(unit.name for unit in compile_units))
    logger.info(
        "Warming up CuTeDSL compile_units=%d names=%s.",
        len(compile_units),
        unit_names,
    )

    start_time = time.perf_counter()
    compiled_count = _compile_cutedsl_warmup_units(compile_units)
    logger.info(
        "CuTeDSL warmup compiled %d units in %.2f s.",
        compiled_count,
        time.perf_counter() - start_time,
    )

register_cutedsl_warmup_provider(provider)

Register an object that can expose CuTeDSL warmup compile units.

Source code in vllm/model_executor/warmup/cutedsl_warmup.py 
33
34
35
def register_cutedsl_warmup_provider(provider: object) -> None:
    """Register an object that can expose CuTeDSL warmup compile units."""
    _CUTEDSL_WARMUP_PROVIDERS.add(provider)