Skip to main content

k3_ssd_chunk_state

Function k3_ssd_chunk_state 

Source
pub fn k3_ssd_chunk_state(
    x_bnlhp: Tensor<5>,
    b_bnlhr: Tensor<5>,
    da_cumsum_bhnl: Tensor<4>,
    dt_discretized_bhnl: Tensor<4>,
) -> Tensor<5>
Expand description

Based on the Kernel 3 Triton reference _chunk_state_fwd_kernel (ssd_chunk_state.py).

Returns:

  • cb_bngll [used in K5][!] - state assuming zero initial state at each chunk boundary.
  • b_bar_scale_bhnl [*] - intermediary