vllm.model_executor.warmup.qwen_triton_warmup

Warm up Qwen Triton kernels from the loaded model's compile keys.

Functions:

• qwen_triton_warmup –

  Warm Qwen GDN Triton kernels reported by the JIT monitor.

_warm_fused_gdn_decode_post_conv_mtp_kernel(device, config, x_dtype, num_spec_tokens, max_num_reqs)

Compile the ROCm Triton MTP decode kernels for every V tile size the request count can select, with the layer's packed QKVZ/BA layouts so the stride specializations match serving.

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

def _warm_fused_gdn_decode_post_conv_mtp_kernel(
    device: torch.device,
    config: _QwenGDNWarmupConfig,
    x_dtype: torch.dtype,
    num_spec_tokens: int,
    max_num_reqs: int,
) -> None:
    """Compile the ROCm Triton MTP decode kernels for every V tile size the
    request count can select, with the layer's packed QKVZ/BA layouts so the
    stride specializations match serving."""
    from vllm.third_party.flash_linear_attention.ops.fused_gdn_decode_post_conv_mtp import (  # noqa: E501
        fused_gdn_decode_post_conv_mtp_triton,
    )
    from vllm.utils.platform_utils import num_compute_units

    num_cus = num_compute_units(device.index)
    # Smallest request counts that select BV = 32, 64 and V (see the wrapper).
    request_counts = sorted({1, -(-num_cus // config.hv), 4 * num_cus // config.hv + 1})
    tokens_per_req = num_spec_tokens + 1
    norm_weight = torch.ones(config.v, dtype=config.norm_weight_dtype, device=device)
    state = torch.zeros(
        2 * config.state_stride_token, dtype=config.state_dtype, device=device
    ).as_strided(
        (2, config.hv, config.v, config.k),
        (config.state_stride_token, config.v * config.k, config.k, 1),
    )
    for num_reqs in request_counts:
        if num_reqs > max_num_reqs:
            continue
        num_tokens = num_reqs * tokens_per_req
        mixed_qkvz = torch.zeros(
            (num_tokens, config.conv_dim + config.hv * config.v),
            dtype=x_dtype,
            device=device,
        )
        mixed_qkv, output_gate = mixed_qkvz.split(
            [config.conv_dim, config.hv * config.v], dim=-1
        )
        ba = torch.zeros((num_tokens, 2 * config.hv), dtype=x_dtype, device=device)
        b, a = ba.chunk(2, dim=-1)
        fused_gdn_decode_post_conv_mtp_triton(
            mixed_qkv=mixed_qkv,
            a=a,
            b=b,
            A_log=config.a_log,
            dt_bias=config.dt_bias,
            state_indices=torch.zeros(
                (num_reqs, tokens_per_req), dtype=torch.int32, device=device
            ),
            cu_seqlens=torch.arange(
                0,
                num_tokens + 1,
                tokens_per_req,
                dtype=torch.int32,
                device=device,
            ),
            num_accepted_tokens=torch.ones(num_reqs, dtype=torch.int32, device=device),
            state=state,
            output_gate=output_gate.reshape(num_tokens, config.hv, config.v),
            norm_weight=norm_weight,
            out=torch.empty(
                (num_tokens, config.hv, config.v), dtype=x_dtype, device=device
            ),
            scale=config.k**-0.5,
            norm_eps=1e-6,
            output_gate_activation=config.norm_activation,
        )

_warm_state_gather_scatter_kernels(device, config)

Compile the gather/scatter variants serving hits.

Triton specializes on 16-byte pointer alignment and stride == 1; mixed batches slice the indices/mask at num_decodes and spec-decode block tables give stride-2 indices, so warm each combination.

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

def _warm_state_gather_scatter_kernels(
    device: torch.device, config: "_QwenGDNWarmupConfig"
) -> None:
    """Compile the gather/scatter variants serving hits.

    Triton specializes on 16-byte pointer alignment and stride == 1; mixed
    batches slice the indices/mask at ``num_decodes`` and spec-decode block
    tables give stride-2 indices, so warm each combination.
    """
    from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
        gather_initial_states,
    )
    from vllm.model_executor.layers.mamba.ops.scatter_states import (
        scatter_states,
    )

    state = torch.zeros(
        2, config.hv, config.v, config.k, dtype=config.state_dtype, device=device
    )
    idx_pool = torch.zeros(6, dtype=torch.int32, device=device)
    has_pool = torch.zeros(2, dtype=torch.bool, device=device)
    for idx, has in (
        (idx_pool[:1], has_pool[:1]),  # aligned, contiguous
        (idx_pool[1:2], has_pool[1:2]),  # unaligned
        (idx_pool[4:5], has_pool[1:2]),  # aligned indices, unaligned mask
        (idx_pool[::2][:1], has_pool[:1]),  # strided (spec-decode block table)
    ):
        gathered = gather_initial_states(state, idx, has)
        scatter_states(state, gathered.to(torch.float32), idx)
        scatter_states(state, gathered, idx)

qwen_triton_warmup(runner, model_config)

Warm Qwen GDN Triton kernels reported by the JIT monitor.

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

@torch.inference_mode()
def qwen_triton_warmup(
    runner: "GPUModelRunner",
    model_config: "ModelConfig",
) -> None:
    """Warm Qwen GDN Triton kernels reported by the JIT monitor."""
    model_type = getattr(model_config.hf_text_config, "model_type", "") or getattr(
        model_config.hf_config, "model_type", ""
    )
    if model_type not in _QWEN_MODEL_TYPES:
        return

    device = runner.device
    logger.info("Warming up Qwen GDN Triton kernels for model_type=%s.", model_type)

    gdn_config = _qwen_gdn_warmup_config(
        runner.compilation_config.static_forward_context
    )
    if gdn_config is None:
        return

    max_num_tokens = max(1, int(runner.max_num_tokens))
    _warm_gated_rms_norm_kernel(device, gdn_config, max_num_tokens, model_config.dtype)
    _warm_causal_conv1d_fwd_kernel(device, gdn_config)
    _warm_fused_post_conv_kernel(device, gdn_config)
    # Pooling only runs full prefills; the decode update kernel is unused.
    if not runner.is_pooling_model:
        _warm_fused_sigmoid_gating_delta_rule_update_kernel(device, gdn_config)
    _warm_state_gather_scatter_kernels(device, gdn_config)
    # ROCm runs the Triton port of the fused MTP decode kernel.
    spec_config = runner.speculative_config
    if (
        current_platform.is_rocm()
        and gdn_config.fused_spec_decode
        and spec_config is not None
    ):
        _warm_fused_gdn_decode_post_conv_mtp_kernel(
            device,
            gdn_config,
            model_config.dtype,
            spec_config.num_speculative_tokens,
            runner.max_num_reqs,
        )
    _synchronize_device(device)