Skip to main content

k5_single_ssd_chunk_scan

burn_mamba::mamba3::single_ssd::ssd::serial

Function k5_single_ssd_chunk_scan 

Source 
pub fn k5_single_ssd_chunk_scan(
    da_cumsum_bhnl: Tensor<4>,
    v_bnlmhp: Tensor<6>,
    c_bnlmhr: Tensor<6>,
    b_bnlmhr: Tensor<6>,
    cb_bnhLMLM: Tensor<5>,
    gamma_bnlh: Tensor<4>,
    scale_bnlh: Tensor<4>,
    chunk_input_state_bnhpr: Tensor<5>,
) -> Tensor<6>
Expand description

SingleSsd chunk scan.

Computes the per-chunk output from three contributions:

  • Strict lower triangular intra-chunk (t1 > t2): (cb[i,j] · scale[t2] · exp(cumA[t1] − cumA[t2])) · V[t2]
  • Same-time-step (t1 == t2) γ-correction: γ[t] · (Σₙ C[t,r_out,n] · B[t,r_in,n]) · V[t,r_in,p]
  • State-to-output (Y_off) — same formula as the double-ssd K5: exp(cumA[t]) · C[t] · h'[n-1]

cb_bnhLMLM is the unscaled C · Bᵀ matrix from K2; b_bnlmhr is the unscaled K/B tensor (used for the γ-correction matmul). The strict-lower MIMO mask excludes the same-step R × R block, leaving only t1 > t2 contributions in the masked CB.

§Shapes

  • da_cumsum_bhnl: [B, H, N, L] (base time grid, not fused)
  • v_bnlmhp: [B, N, L, M, H, P]
  • c_bnlmhr, b_bnlmhr: [B, N, L, M, H, R]
  • cb_bnhLMLM: [B, N, H, L·M, L·M] (output of K2)
  • gamma_bnlh, scale_bnlh: [B, N, L, H]
  • chunk_input_state_bnhpr: [B, N, H, P, R] (h’ at chunk start)

§Returns

  • y_bnlmhp: [B, N, L, M, H, P]