Skip to content

vllm.model_executor.warmup.mamba_triton_warmup

Warm Mamba-style Triton kernels shared across GDN / Mamba / KDA models.

Functions:

_warm_batch_memcpy_kernel(device)

Warm the Mamba prefix-cache state copy specialization.

Source code in vllm/model_executor/warmup/mamba_triton_warmup.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def _warm_batch_memcpy_kernel(device: torch.device) -> None:
    """Warm the Mamba prefix-cache state copy specialization."""
    from vllm.v1.worker.mamba_utils import batch_memcpy

    src = torch.empty(1024, dtype=torch.uint8, device=device)
    dst = torch.empty_like(src)
    batch_memcpy(
        torch.tensor([src.data_ptr()], dtype=torch.uint64, device=device),
        torch.tensor([dst.data_ptr()], dtype=torch.uint64, device=device),
        # Keep this int32: Triton's compile key includes pointer element dtypes,
        # and the production prefix-cache path passes an int32 sizes tensor.
        torch.tensor([src.numel()], dtype=torch.int32, device=device),
    )
    logger.info("Warmed Mamba batch_memcpy_kernel.")

mamba_triton_warmup(runner)

Warm prefix-cache memcpy for every Mamba-style KV cache group.

Source code in vllm/model_executor/warmup/mamba_triton_warmup.py 
45
46
47
48
49
50
51
52
53
@torch.inference_mode()
def mamba_triton_warmup(runner: "GPUModelRunner") -> None:
    """Warm prefix-cache memcpy for every Mamba-style KV cache group."""
    device = runner.device
    if device.type != "cuda":
        return
    if not _has_mamba_style_cache(runner):
        return
    _warm_batch_memcpy_kernel(device)