vllm.models.qwen4_exp.nvidia.ops.ple ¶

def ple_conv(
    inputs: torch.Tensor,
    residual: torch.Tensor,
    conv_state: torch.Tensor,
    conv_weights: torch.Tensor,
    state_indices: torch.Tensor,
    *,
    mode: Literal["decode", "spec", "prefill"],
    dilation: int,
    query_start_loc: torch.Tensor | None = None,
    num_accepted_tokens: torch.Tensor | None = None,
    has_initial_states: torch.Tensor | None = None,
    spec_query_len: int = 1,
    token_indices: torch.Tensor | None = None,
) -> None:
    """Add short-convolution output to ``residual`` and update its state."""
    BLOCK_C = 512
    kernel_spec_query_len = spec_query_len if mode == "spec" else 1
    T, C = inputs.shape
    K = conv_weights.shape[1]
    state_len = (K - 1) * dilation
    state_width = state_len + kernel_spec_query_len - 1
    if token_indices is not None:
        T = token_indices.numel()
    if conv_state.shape[1] != C or conv_state.shape[2] < state_width:
        raise ValueError(
            "conv_state must have shape [slots, channels, window], with "
            f"channels={C} and window >= {state_width}"
        )
    state_bs, state_cs, state_ws = conv_state.stride()

    if mode == "decode":
        num_reqs = T
        binary_search_iters = 1
        has_initial_states_arg = has_initial_states is not None
    elif mode == "spec":
        if query_start_loc is None or num_accepted_tokens is None:
            raise ValueError(
                "query_start_loc and num_accepted_tokens are required for spec decode"
            )
        num_reqs = state_indices.numel()
        binary_search_iters = max(num_reqs, 1).bit_length()
        has_initial_states_arg = False
    elif mode == "prefill":
        if query_start_loc is None or has_initial_states is None:
            raise ValueError(
                "query_start_loc and has_initial_states are required for prefill"
            )
        num_reqs = state_indices.numel()
        binary_search_iters = max(num_reqs, 1).bit_length()
        has_initial_states_arg = True
    else:
        raise ValueError(f"Unsupported short-conv mode: {mode}")

    num_warps = 4 if mode == "prefill" else 8
    launch_pdl = current_platform.is_arch_support_pdl()
    # Pure-prefill indices can be a strided block-table column view.
    state_idx_stride = state_indices.stride(0)

    # Constexpr flags eliminate accesses to optional None arguments. Without a
    # token map, state_indices is an unused but device-resident placeholder.
    _ple_conv_kernel[(T, triton.cdiv(C, BLOCK_C))](
        inputs,
        conv_state,
        conv_weights,
        residual,
        state_indices,
        query_start_loc,
        num_accepted_tokens,
        has_initial_states,
        token_indices if token_indices is not None else state_indices,
        token_indices is not None,
        num_reqs,
        binary_search_iters,
        state_idx_stride,
        state_bs,
        state_ws,
        state_cs,
        C=C,
        BLOCK_C=BLOCK_C,
        STATE_LEN=state_len,
        DILATION=dilation,
        KERNEL_SIZE=K,
        SPEC_QUERY_LEN=kernel_spec_query_len,
        MODE=mode,
        HAS_INIT=has_initial_states_arg,
        NULL_STATE_ID=NULL_BLOCK_ID,
        launch_pdl=launch_pdl,
        num_warps=num_warps,
    )
    # conv state update is fused with the kernel above for decode
    if mode != "decode":
        _ple_conv_writeback_kernel[(num_reqs, triton.cdiv(C, BLOCK_C))](
            inputs,
            conv_state,
            state_indices,
            query_start_loc,
            num_accepted_tokens,
            has_initial_states,
            token_indices if token_indices is not None else state_indices,
            token_indices is not None,
            state_idx_stride,
            state_bs,
            state_ws,
            state_cs,
            C=C,
            BLOCK_C=BLOCK_C,
            STATE_LEN=state_len,
            SPEC_QUERY_LEN=kernel_spec_query_len,
            STATE_WIDTH=state_width,
            MODE=mode,
            HAS_INIT=has_initial_states_arg,
            NULL_STATE_ID=NULL_BLOCK_ID,
            num_warps=num_warps,
        )