flashinfer.mamba.causal_conv1d_fn

flashinfer.mamba.causal_conv1d_fn(x: Tensor, weight: Tensor, conv_state: Tensor, cu_seqlens: Tensor, cache_indices: Tensor, has_initial_state: Tensor | None = None, *, bias: Tensor | None = None, out: Tensor | None = None, workspace: Tensor | None = None) → Tensor

Apply width-four causal depthwise convolution and SiLU to packed sequences.

Parameters:

  • x (torch.Tensor) – BF16 CUDA input of shape [rows, channels] with contiguous channels. Rows may have padding. State history stores these raw inputs.

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

  • conv_state (torch.Tensor) – In-place BF16 pool [slots, channels, L], L >= 3. Either the channel or history axis must be contiguous; slots must not overlap. Each nonempty active sequence writes its final three raw history values to positions 0:3, leaving the rest of the pool unchanged.

  • cu_seqlens (torch.Tensor) – Contiguous CUDA int32 offsets [N+1], starting at zero, nondecreasing, with the last offset at most rows. Rows after the last offset are not written. Empty sequences do not update state.

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

  • has_initial_state (Optional[torch.Tensor]) – Contiguous CUDA bool [N]. True rows read state positions 0:3; false or omitted rows start from zero history.

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

  • out (Optional[torch.Tensor]) – BF16 output 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.

  • workspace (Optional[torch.Tensor]) – Contiguous CUDA int32 scratch. For N > 32 and nonempty x, needs 1 + 2 * min(rows, ceil(rows / B) + N) elements, where B is 8 if rows < 64*N and 16 otherwise. Allocated if omitted. Smaller batches need no scratch. The chunk map is rebuilt on every call.

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 and workspace must not alias other operands. Concurrent calls need disjoint live slots and scratch. CUDA graph replay may change device metadata in place within fixed capacities. Supply out and workspace before capture to reuse storage. After prefill, the first speculative update uses accepted count 1.