Example

transformer_block.kina

Source

requires [B = B_p, Seq = Seq_p, D = D_p, D_ffn = D_ffn_p] compute {
    let run_transformer = fn(
        x: tensor<F32, [1, 32, 64], @host>,
        gamma1: tensor<F32, [1, 1, 64], @host>,
        wq: tensor<F32, [1, 64, 64], @host>,
        wk: tensor<F32, [1, 64, 64], @host>,
        wv: tensor<F32, [1, 64, 64], @host>,
        wo: tensor<F32, [1, 64, 64], @host>,
        gamma2: tensor<F32, [1, 1, 64], @host>,
        w1: tensor<F32, [1, 64, 128], @host>,
        w2: tensor<F32, [1, 128, 64], @host>
    ) -> tensor<F32, [1, 32, 64], @host> {

        let x_norm1_scaled = rmsnorm(x, gamma1) in
        let attn_proj = attention(x_norm1_scaled, wq, wk, wv, wo) in
        let x2 = x + attn_proj in
        let x_norm2_scaled = rmsnorm(x2, gamma2) in
        let ffn_out = swiglu(x_norm2_scaled, w1, w2) in
        let out = x2 + ffn_out in
        out
    } in

    let b = tensor<F32, [1, 32, 64], @host> { 1.0 } in
    let g1 = tensor<F32, [1, 1, 64], @host> { 1.0 } in
    let wq = tensor<F32, [1, 64, 64], @host> { 1.0 } in
    let wk = tensor<F32, [1, 64, 64], @host> { 1.0 } in
    let wv = tensor<F32, [1, 64, 64], @host> { 1.0 } in
    let wo = tensor<F32, [1, 64, 64], @host> { 1.0 } in
    let g2 = tensor<F32, [1, 1, 64], @host> { 1.0 } in
    let w1 = tensor<F32, [1, 64, 128], @host> { 1.0 } in
    let w2 = tensor<F32, [1, 128, 64], @host> { 1.0 } in

    run_transformer(b, g1, wq, wk, wv, wo, g2, w1, w2)
}