def _state_passing_fwd(
states,
dA_cumsum,
last_chunk_indices,
initial_states=None,
out_dtype=None,
final_states=None,
final_state_indices=None,
):
"""`final_states` (num_slots, nheads, dim), if given, receives each
sequence's final state at slot `final_state_indices[b]`."""
nchunks, nheads, dim = states.shape
chunk_size = dA_cumsum.shape[-1]
batch = last_chunk_indices.shape[0]
assert dA_cumsum.shape == (nheads, nchunks, chunk_size)
out_dtype = states.dtype if out_dtype is None else out_dtype
out = torch.empty((nchunks, nheads, dim), device=states.device, dtype=out_dtype)
initial_states_strides = (
(initial_states.stride(0), initial_states.stride(1), initial_states.stride(2))
if initial_states is not None
else (0, 0, 0)
)
if final_states is not None:
assert final_states.shape[1:] == (nheads, dim)
assert final_state_indices is not None
assert final_state_indices.shape == (batch,)
final_states_strides = (
(final_states.stride(0), final_states.stride(1), final_states.stride(2))
if final_states is not None
else (0, 0, 0)
)
grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE"]), batch, nheads)
launch_autotuned(
_state_passing_fwd_kernel,
grid,
states_ptr=states,
out_ptr=out,
dA_cs_ptr=dA_cumsum,
initstates_ptr=initial_states,
last_chunk_indices_ptr=last_chunk_indices,
final_states_ptr=final_states,
final_state_indices_ptr=final_state_indices,
dim=dim,
chunk_size=chunk_size,
stride_states_chunk=states.stride(0),
stride_states_head=states.stride(1),
stride_states_dim=states.stride(2),
stride_out_chunk=out.stride(0),
stride_out_head=out.stride(1),
stride_out_dim=out.stride(2),
stride_dA_cs_head=dA_cumsum.stride(0),
stride_dA_cs_chunk=dA_cumsum.stride(1),
stride_dA_cs_csize=dA_cumsum.stride(2),
stride_initstates_batch=initial_states_strides[0],
stride_initstates_head=initial_states_strides[1],
stride_initstates_dim=initial_states_strides[2],
stride_final_states_slot=final_states_strides[0],
stride_final_states_head=final_states_strides[1],
stride_final_states_dim=final_states_strides[2],
stride_final_state_indices=(
final_state_indices.stride(0) if final_state_indices is not None else 0
),
HAS_INITSTATES=initial_states is not None,
HAS_FINAL_STATES=final_states is not None,
)
return out