Example

llm_generate.kina

Source

let generate = fn(steps: tensor<I32, [], @host>, start_token: tensor<F32, [4, 4], @host>, init_kv: tensor<F32, [16, 4], @host>, start_pos: tensor<I32, [], @host>, step_size: tensor<I32, [], @host>) -> (tensor<F32, [4, 4], @host>, tensor<F32, [16, 4], @host>, tensor<I32, [], @host>) {
  let final = while (pos < steps) [tok = start_token, kv = init_kv, pos = start_pos] {
    let w = tensor<F32, [4, 4], @host> in
    let next_tok = tok @ w in
    let next_kv = kv + kv in
    let next_pos = pos + step_size in
    (next_tok, next_kv, next_pos)
  } in
  final
} in
let steps = tensor<I32, [], @host> in
let start_token = tensor<F32, [4, 4], @host> in
let init_kv = tensor<F32, [16, 4], @host> in
let start_pos = tensor<I32, [], @host> in
let step_size = tensor<I32, [], @host> in
generate(steps, start_token, init_kv, start_pos, step_size)