vllm.model_executor.warmup.deep_gemm_warmup

Warmup deep_gemm kernels. DeepGEMM JIT's the kernels. The warmup aims to JIT all the kernels that would be used during model execution beforehand.

_deep_gemm_linear_data(module)

Return (weight, weight_scale) for a layer whose kernel stamped itself as deep_gemm_warmup_provider in process_weights_after_loading, else None.

Source code in vllm/model_executor/warmup/deep_gemm_warmup.py

def _deep_gemm_linear_data(
    module: torch.nn.Module,
) -> tuple[torch.Tensor, torch.Tensor] | None:
    """Return (weight, weight_scale) for a layer whose kernel stamped itself as
    ``deep_gemm_warmup_provider`` in ``process_weights_after_loading``, else None.
    """
    provider = getattr(module, "deep_gemm_warmup_provider", None)
    if not isinstance(provider, DeepGemmFp8BlockScaledMMKernel):
        return None
    if module.weight_block_size != get_mk_alignment_for_contiguous_layout():
        return None
    return provider.get_deep_gemm_warmup_weights(module)

_extract_data_from_fused_moe_module(m_)

Extract weights, weight scales and num_topk from MoERunner module.

Source code in vllm/model_executor/warmup/deep_gemm_warmup.py

def _extract_data_from_fused_moe_module(
    m_: torch.nn.Module,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
    """Extract weights, weight scales and num_topk from MoERunner module."""
    assert isinstance(m_, MoERunner)
    m = m_.routed_experts
    w13 = m.w13_weight
    w13_s = (
        m.w13_weight_scale_inv
        if hasattr(m, "w13_weight_scale_inv")
        else m.w13_weight_scale
    )
    w2 = m.w2_weight
    w2_s = (
        m.w2_weight_scale_inv
        if hasattr(m, "w2_weight_scale_inv")
        else m.w2_weight_scale
    )
    num_topk = m.top_k

    assert isinstance(w13, torch.Tensor)
    assert isinstance(w13_s, torch.Tensor)
    assert isinstance(w2, torch.Tensor)
    assert isinstance(w2_s, torch.Tensor)
    return w13, w13_s, w2, w2_s, num_topk

_generate_optimal_warmup_m_values(max_tokens, n, device)

Generate M values that cover all possible DeepGEMM kernel configurations.

The block sizes must mirror the layout candidates of the vendored DeepGEMM heuristic (csrc/jit_kernels/heuristics/sm100.hpp); a stale set here lets serving-time M values select kernel configs the warmup never compiled.

Parameters:

• max_tokens

  (int) –

  Maximum number of tokens to warmup for

• n

  (int) –

  The actual N dimension from the weight tensor

• device

  (device) –

  The torch device to get properties from.

Source code in vllm/model_executor/warmup/deep_gemm_warmup.py

def _generate_optimal_warmup_m_values(
    max_tokens: int, n: int, device: torch.device
) -> list[int]:
    """Generate M values that cover all possible DeepGEMM kernel configurations.

    The block sizes must mirror the layout candidates of the vendored DeepGEMM
    heuristic (csrc/jit_kernels/heuristics/sm100.hpp); a stale set here lets
    serving-time M values select kernel configs the warmup never compiled.

    Args:
        max_tokens: Maximum number of tokens to warmup for
        n: The actual N dimension from the weight tensor
        device: The torch device to get properties from.

    """
    # DeepGEMM's possible block sizes
    block_ms = list(range(16, 257, 16))
    block_ns = list(range(16, min(257, n + 1), 16))
    num_sms = num_compute_units(device.index)

    m_values = set()

    # Always include small cases
    m_values.update([1, 2, 4] + [i for i in range(8, 65, 8)])

    # Collect M values where different wave patterns occur
    for block_m in block_ms:
        for block_n in block_ns:
            if block_n > n:
                continue

            # Add key M boundaries for this block combination
            for wave in range(1, 11):  # Up to 10 waves
                # M where this block config transitions to next wave
                target_blocks = wave * num_sms
                m = target_blocks * block_m // cdiv(n, block_n)
                if 1 <= m <= max_tokens:
                    m_values.add(m)

            # Add block_m boundaries
            for multiple in range(1, max_tokens // block_m + 1):
                m = multiple * block_m
                if m <= max_tokens:
                    m_values.add(m)

    return sorted(m_values)

_get_fp8_fp4_mega_moe_m_values(module, max_tokens)

Get the M values that reach every MegaMoE kernel and, where used, bf16_mega_gate config.

Source code in vllm/model_executor/warmup/deep_gemm_warmup.py

def _get_fp8_fp4_mega_moe_m_values(
    module: "DeepseekV4MoE | KimiMoE | MiniMaxM3MoE", max_tokens: int
) -> list[int]:
    """Get the M values that reach every MegaMoE kernel and, where used,
    bf16_mega_gate config."""
    dg = _import_deep_gemm()
    gate, experts = module.gate, module.experts
    uses_mega_gate = _mega_moe_uses_mega_gate(module)
    if module.use_sequence_parallel:
        max_tokens = cdiv(max_tokens, module.tp_size)

    m_values: dict[tuple, int] = {}
    for num_tokens in range(1, max_tokens + 1):
        gate_config = (
            dg.get_bf16_mega_gate_config(
                num_tokens, gate.input_size, gate.output_size, experts.top_k
            )
            if uses_mega_gate and hasattr(dg, "get_bf16_mega_gate_config")
            else {}
        )
        key = (tuple(gate_config.items()), experts.get_block_m(num_tokens))
        # Keep the largest M: DeepSeek-V4 only calls bf16_mega_gate above a
        # small token threshold, and a config's largest M is always above it.
        m_values[key] = num_tokens
        # One NVFP4 BLOCK_M spans two kernels: decode-shaped launches at small
        # M and bulk launches at large M. DeepGEMM does not expose the switch.
        if experts.weight_format.uses_nvfp4_api:
            m_values.setdefault(("smallest", *key), num_tokens)
    return sorted(set(m_values.values()))

_get_fp8_gemm_nt_m_values(w, max_tokens)

Get the M values to warmup for a given weight tensor.

Source code in vllm/model_executor/warmup/deep_gemm_warmup.py

def _get_fp8_gemm_nt_m_values(w: torch.Tensor, max_tokens: int) -> list[int]:
    """Get the M values to warmup for a given weight tensor."""
    n, _ = w.size()
    device = w.device

    # Use optimal M values only if VLLM_DEEP_GEMM_WARMUP is set to "relax".
    # Otherwise warmup all token sizes to avoid JIT compilation in hotpath
    if envs.VLLM_DEEP_GEMM_WARMUP == "relax":
        return _generate_optimal_warmup_m_values(max_tokens, n, device)
    else:
        assert envs.VLLM_DEEP_GEMM_WARMUP == "full", (
            "Expected "
            'VLLM_DEEP_GEMM_WARMUP env to be set to "full" but got '
            f"{envs.VLLM_DEEP_GEMM_WARMUP}"
        )
        return list(range(1, max_tokens + 1))