flashinfer.mamba.causal_conv1d_update

flashinfer.mamba.causal_conv1d_update(x: Tensor, weight: Tensor, conv_state: Tensor, state_indices: Tensor, *, cu_seqlens: Tensor | None = None, num_accepted_tokens: Tensor | None = None, bias: Tensor | None = None, out: Tensor | None = None) → Tensor

Apply width-four causal convolution and SiLU, updating raw-input history.

Parameters:

  • x (torch.Tensor) – BF16 CUDA input [rows, channels] with contiguous channels and nonoverlapping, optionally padded rows.

  • weight (torch.Tensor) – Contiguous FP32 or BF16 CUDA [channels, 4]. Tap 3 multiplies the current token. Accumulation and SiLU use FP32.

  • conv_state (torch.Tensor) – In-place BF16 [slots, channels, L], L >= 3, with nonoverlapping slots and either the channel or history axis contiguous. Ordinary updates read positions 0:3 and replace them with the final three raw inputs/history values; remaining positions are unchanged.

  • state_indices (torch.Tensor) – Contiguous CUDA int32 [N]. Live slot IDs must be unique and in [0, slots). The -1 sentinel produces zero output and no state access; slot zero is valid.

  • cu_seqlens (Optional[torch.Tensor]) – Contiguous CUDA int32 [N+1]. Offsets start at zero, are nondecreasing and end at most at rows. Trailing output rows and state for empty sequences remain untouched. If omitted, rows must equal N and each sequence contains one token.

  • num_accepted_tokens (Optional[torch.Tensor]) – Contiguous CUDA int32 [N] enabling linear speculative updates. A positive count a includes the bonus token accepted from the previous call and reads history state[..., a-1:a+2]; require 1 <= a <= L-2 and a window written by that call. Nonpositive counts skip the row, producing zero output and leaving state unchanged. A current sequence of T tokens requires T <= L-2 and writes [old_h1, old_h2, x0, ..., x(T-1)] to positions 0:T+2.

  • bias (Optional[torch.Tensor]) – Contiguous FP32 or BF16 CUDA [channels], added before the four FMAs.

  • out (Optional[torch.Tensor]) – BF16 storage matching x with contiguous channels and nonoverlapping rows. Must not share storage with any operand. Allocated if omitted; unused trailing rows of an allocated output are unspecified.

Returns:

BF16 output [rows, channels]; conv_state is updated in place.

Return type:

torch.Tensor

Notes

Device offset and slot values are caller-owned constraints; no host read or validation synchronization occurs. State must not alias other operands. CUDA graph replay may change lengths, slots and counts in place within fixed capacities. No workspace is needed. An accepted count above L-2 skips the row. A speculative sequence exceeding L-2 still produces all outputs, but stores only the history prefix that fits; it cannot supply a complete history for subsequent speculative verification.