For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python function

gated_delta_conv1d_fwd

`gated_delta_conv1d_fwd()`

max.nn.legacy.state_space.gated_delta_conv1d_fwd(qkv_input_ragged, conv_weight, conv_state, slot_idx, input_row_offsets, *, write_state=True)

source

Applies the causal conv1d pass.

Mutates a slot-indexed conv-state pool in place.

conv_state is a mutable pool of shape [max_slots, conv_dim, kernel_size - 1] and the kernel reads/writes slot slot_idx[batch_item] directly. There is no conv_state_out graph output: the pool is mutated in place.

Parameters:

qkv_input_ragged (TensorValue) – The [total_seq_len, conv_dim] projected QKV input.
conv_weight (TensorValue) – The [conv_dim, kernel_size] depthwise conv weights.
conv_state (BufferValue) – The [max_slots, conv_dim, kernel_size - 1] mutable pool.
slot_idx (TensorValue) – The [batch_size] uint32 slot indices into the pool.
input_row_offsets (TensorValue) – The [batch_size + 1] uint32 ragged offsets.
write_state (bool) – Whether to write the updated window back to conv_state. False only reads it for the look-back.

Returns:

The conv output, [total_seq_len, conv_dim].

Return type:

TensorValue

The following example convolves a ragged batch of two sequences, of lengths 3 and 2, on a GPU. The sequences read and update conv-state slots 2 and 0:

import numpy as np
from max.driver import Accelerator, Buffer
from max.dtype import DType
from max.engine import InferenceSession
from max.graph import BufferType, DeviceRef, Graph, TensorType
from max.nn.legacy.state_space import gated_delta_conv1d_fwd

total_seq_len, conv_dim, kernel_size, max_slots = 5, 8, 4, 3
gpu = DeviceRef.GPU()

with Graph(
    "gated_delta_conv1d_fwd_example",
    input_types=(
        TensorType(
            DType.float32, [total_seq_len, conv_dim], device=gpu
        ),
        TensorType(DType.float32, [conv_dim, kernel_size], device=gpu),
        BufferType(
            DType.float32,
            [max_slots, conv_dim, kernel_size - 1],
            device=gpu,
        ),
        TensorType(DType.uint32, [2], device=gpu),
        TensorType(DType.uint32, [3], device=gpu),
    ),
) as graph:
    x, conv_weight, conv_state, slot_idx, input_row_offsets = (
        graph.inputs
    )
    output = gated_delta_conv1d_fwd(
        qkv_input_ragged=x.tensor,
        conv_weight=conv_weight.tensor,
        conv_state=conv_state.buffer,
        slot_idx=slot_idx.tensor,
        input_row_offsets=input_row_offsets.tensor,
    )
    graph.output(output)

rng = np.random.default_rng(0)
x_np = rng.standard_normal((total_seq_len, conv_dim), dtype=np.float32)
weight_np = rng.standard_normal(
    (conv_dim, kernel_size), dtype=np.float32
)
state_np = rng.standard_normal(
    (max_slots, conv_dim, kernel_size - 1), dtype=np.float32
)
slots = np.array([2, 0], dtype=np.uint32)
offsets = np.array([0, 3, 5], dtype=np.uint32)

gpu = Accelerator()
inputs = [
    Buffer.from_numpy(array).to(gpu)
    for array in (x_np, weight_np, state_np.copy(), slots, offsets)
]
model = InferenceSession(devices=[gpu]).load(graph)
result = model.execute(*inputs)[0].to_numpy()
state = inputs[2]
# result.shape is (5, 8). Slots 2 and 0 of state now hold the last
# three inputs of each sequence, and slot 1 is unchanged.

gated_delta_conv1d_fwd()​

`gated_delta_conv1d_fwd()`