A KDA sequence segment acts on an incoming recurrent state through an affine map. For segment i, cumulative transition Mi, local zero-initialized contribution Ui, and incoming state S, define
Fi(S)=MiS+Ui.Two ordered maps compose as
(M2,U2)∘(M1,U1)=(M2M1,M2U1+U2).Use an exclusive prefix composition so each segment receives the state produced by preceding segments only. Apply the current segment map to that incoming state to obtain its outgoing state. Matrix multiplication order must follow sequence order because transitions need not commute. Return a dictionary with incoming_states, outgoing_states, and final_state. All tensors preserve the input floating-point dtype and device.
Examples
Input: transition_matrices = [[[2]]], local_contributions = [[[3]]], initial_state = [[4]]
Output: {"incoming_states": tensor([[[4.0]]]), "outgoing_states": tensor([[[11.0]]]), "final_state": tensor([[11.0]])}
Explanation: The single segment receives state 4 and applies the affine map 2 times 4 plus 3, producing final state 11.
Input: transition_matrices.shape = (3, 2, 2), local_contributions.shape = (3, 2, 1), initial_state.shape = (2, 1)
Output: {"incoming_states": tensor of shape (3, 2, 1), "outgoing_states": tensor of shape (3, 2, 1), "final_state": tensor of shape (2, 1)}
Input: transition_matrices.shape = (2, 2, 2), local_contributions.shape = (2, 2, 2), initial_state.shape = (2, 2)
Output: {"incoming_states": tensor of shape (2, 2, 2), "outgoing_states": tensor of shape (2, 2, 2), "final_state": tensor of shape (2, 2)}
Hint 1
Represent the completed prefix with one transition matrix and one accumulated contribution.
Hint 2
Evaluate each outgoing state as the current transition applied to its incoming state plus its contribution.
Hint 3
Compose transitions on the left so segment order is preserved.
Requirements
- Return a dictionary with exactly incoming_states, outgoing_states, and final_state.
- Incoming and outgoing stacks use segment order as their leading dimension.
- Preserve floating-point dtype and device.
- Do not mutate transition matrices, local contributions, or initial state.
Constraints
- At least one segment is supplied.
- Transition matrices have shape (segments, state width, state width).
- Local contributions have shape (segments, state width, value width).
- Initial state has shape (state width, value width).
All steps in Kimi K3
Accepts: array
Accepts: array
Accepts: array
KDA Context Parallel
A KDA sequence segment acts on an incoming recurrent state through an affine map. For segment i, cumulative transition Mi, local zero-initialized contribution Ui, and incoming state S, define
Fi(S)=MiS+Ui.Two ordered maps compose as
(M2,U2)∘(M1,U1)=(M2M1,M2U1+U2).Use an exclusive prefix composition so each segment receives the state produced by preceding segments only. Apply the current segment map to that incoming state to obtain its outgoing state. Matrix multiplication order must follow sequence order because transitions need not commute. Return a dictionary with incoming_states, outgoing_states, and final_state. All tensors preserve the input floating-point dtype and device.
Examples
Input: transition_matrices = [[[2]]], local_contributions = [[[3]]], initial_state = [[4]]
Output: {"incoming_states": tensor([[[4.0]]]), "outgoing_states": tensor([[[11.0]]]), "final_state": tensor([[11.0]])}
Explanation: The single segment receives state 4 and applies the affine map 2 times 4 plus 3, producing final state 11.
Input: transition_matrices.shape = (3, 2, 2), local_contributions.shape = (3, 2, 1), initial_state.shape = (2, 1)
Output: {"incoming_states": tensor of shape (3, 2, 1), "outgoing_states": tensor of shape (3, 2, 1), "final_state": tensor of shape (2, 1)}
Input: transition_matrices.shape = (2, 2, 2), local_contributions.shape = (2, 2, 2), initial_state.shape = (2, 2)
Output: {"incoming_states": tensor of shape (2, 2, 2), "outgoing_states": tensor of shape (2, 2, 2), "final_state": tensor of shape (2, 2)}
Hint 1
Represent the completed prefix with one transition matrix and one accumulated contribution.
Hint 2
Evaluate each outgoing state as the current transition applied to its incoming state plus its contribution.
Hint 3
Compose transitions on the left so segment order is preserved.
Requirements
- Return a dictionary with exactly incoming_states, outgoing_states, and final_state.
- Incoming and outgoing stacks use segment order as their leading dimension.
- Preserve floating-point dtype and device.
- Do not mutate transition matrices, local contributions, or initial state.
Constraints
- At least one segment is supplied.
- Transition matrices have shape (segments, state width, state width).
- Local contributions have shape (segments, state width, value width).
- Initial state has shape (state width, value width).
All steps in Kimi K3
Accepts: array
Accepts: array
Accepts: array