`vllm.models.inkling.nvidia.ops.sconv` ¶

Inkling short-convolution kernels backed by a paged sliding-window state cache.

Each layer's 4 conv streams (K, V, attn-output, mlp-output) share one paged KV cache [num_blocks, H, N, D] (head-major; see sconv_swa_attn.py). A stream occupies the contiguous D-sub-range [off_s, off_s + ws) across all H heads, so its flat per-token width is H * ws and that is the conv channel dim. The cache stores the conv input at every absolute position.

fused_sconv is the single-launch path used by the model: per token it writes the current input to its slot and convolves the W taps ending at its absolute position. A tap landing inside the current forward is read from the immutable input x (row src - pos + pid_t); only pre-forward taps are read from the paged cache (window position src -> physical block via block_table[req, src // N]). The just-written slot is never read back this step, so there is no write/read hazard within or across programs -- which is why this needs no decode-vs-prefill split and is valid for prefill / decode / spec alike.

All kernels address the cache purely by (slot, absolute_position) and allocate nothing inside the captured region; their grids depend only on the token count (fused_sconv on a fixed token/channel tiling), so the same decode path replays correctly under a full CUDA graph without any data-dependent shape or branch.

Functions:

fused_sconv –

Single-launch insert + depthwise causal conv1d over the paged cache.
sconv_seq_metadata –

Fill static per-token seq_idx / query_start buffers in one launch.

`fused_sconv(x, weight, cache, positions, block_table, seq_idx, slot_mapping, query_start, off_s, ws, block_size, activation=None, use_residual=True)` ¶

Single-launch insert + depthwise causal conv1d over the paged cache.

Reads same-forward taps from x and pre-forward taps from the cache, so it is race-free in one launch for prefill / decode / spec and supports full CUDA-graph capture for decode.

Source code in vllm/models/inkling/nvidia/ops/sconv.py

def fused_sconv(
    x: torch.Tensor,  # [T, H*ws] head-major current-token inputs
    weight: torch.Tensor,  # [H*ws, W]
    cache: torch.Tensor,  # [num_blocks, H, N, D] paged
    positions: torch.Tensor,  # [T] int64 absolute position per token
    block_table: torch.Tensor,  # [num_reqs, max_blocks] int32
    seq_idx: torch.Tensor,  # [T] int32 token -> batch request
    slot_mapping: torch.Tensor,  # [T] int64 flat slot (PAD = -1 => skip)
    query_start: torch.Tensor,  # [T] int32 first x-row of the token's request
    off_s: int,
    ws: int,
    block_size: int,
    activation: str | None = None,
    use_residual: bool = True,
) -> torch.Tensor:
    """Single-launch insert + depthwise causal conv1d over the paged cache.

    Reads same-forward taps from ``x`` and pre-forward taps from the cache, so
    it is race-free in one launch for prefill / decode / spec and supports full
    CUDA-graph capture for decode.
    """
    T = x.shape[0]
    out = torch.empty_like(x)
    if T == 0:
        return out
    assert x.is_contiguous()
    assert cache.stride(3) == 1, "cache D-dim must be contiguous"
    H = cache.shape[1]
    W = weight.shape[1]
    C = H * ws  # flat conv channel dim (all heads, head-major)
    # Tile BT tokens x BLOCK_C channels per program: enough work per CTA to
    # amortize the per-token addressing, while keeping the grid large for
    # prefill. BLOCK_C spans heads so there is no per-head launch. A ~2K-element
    # tile at 4 warps measured best on Blackwell; larger tiles spill registers.
    BLOCK_C = min(triton.next_power_of_2(C), 256)
    BT = 8
    grid = (triton.cdiv(T, BT), triton.cdiv(C, BLOCK_C))
    _fused_sconv_kernel[grid](
        x,
        cache,
        weight,
        out,
        positions,
        seq_idx,
        slot_mapping,
        block_table,
        query_start,
        T,
        x.stride(0),
        cache.stride(0),
        cache.stride(1),
        cache.stride(2),
        cache.stride(3),
        weight.stride(0),
        weight.stride(1),
        block_table.stride(0),
        block_table.shape[1],
        block_size,
        W=W,
        USE_SILU=activation in ("silu", "swish"),
        USE_RESIDUAL=use_residual,
        OFF_S=off_s,
        WS=ws,
        H=H,
        BT=BT,
        BLOCK_C=BLOCK_C,
        num_warps=4,
    )
    return out

`sconv_seq_metadata(query_start_loc, num_reqs, num_actual_tokens, seq_idx_out, query_start_out, num_padded_tokens=None)` ¶

Fill static per-token seq_idx / query_start buffers in one launch.

Replaces the arange + searchsorted + clamp + gather + 2x copy chain of the sconv metadata build with a single kernel writing both persistent buffers. Padded rows are filled with zero and must have slot_mapping == -1.

Source code in vllm/models/inkling/nvidia/ops/sconv.py

def sconv_seq_metadata(
    query_start_loc: torch.Tensor,
    num_reqs: int,
    num_actual_tokens: int,
    seq_idx_out: torch.Tensor,
    query_start_out: torch.Tensor,
    num_padded_tokens: int | None = None,
) -> None:
    """Fill static per-token seq_idx / query_start buffers in one launch.

    Replaces the arange + searchsorted + clamp + gather + 2x copy chain of the
    sconv metadata build with a single kernel writing both persistent buffers.
    Padded rows are filled with zero and must have ``slot_mapping == -1``.
    """
    if num_padded_tokens is None:
        num_padded_tokens = num_actual_tokens
    if num_padded_tokens < num_actual_tokens:
        raise ValueError("num_padded_tokens must cover all actual tokens")
    if num_padded_tokens > seq_idx_out.shape[0]:
        raise ValueError("seq_idx_out is too small for the padded token count")
    if num_padded_tokens > query_start_out.shape[0]:
        raise ValueError("query_start_out is too small for the padded token count")

    BLOCK = 256
    n_iters = (num_reqs - 1).bit_length()
    grid = (triton.cdiv(num_padded_tokens, BLOCK),)
    _seq_metadata_kernel[grid](
        query_start_loc,
        seq_idx_out,
        query_start_out,
        num_reqs,
        num_actual_tokens,
        num_padded_tokens,
        n_iters,
        BLOCK=BLOCK,
    )

vllm.models.inkling.nvidia.ops.sconv ¶

fused_sconv(x, weight, cache, positions, block_table, seq_idx, slot_mapping, query_start, off_s, ws, block_size, activation=None, use_residual=True) ¶

sconv_seq_metadata(query_start_loc, num_reqs, num_actual_tokens, seq_idx_out, query_start_out, num_padded_tokens=None) ¶

`vllm.models.inkling.nvidia.ops.sconv` ¶

`fused_sconv(x, weight, cache, positions, block_table, seq_idx, slot_mapping, query_start, off_s, ws, block_size, activation=None, use_residual=True)` ¶

`sconv_seq_metadata(query_start_loc, num_reqs, num_actual_tokens, seq_idx_out, query_start_out, num_padded_tokens=None)` ¶