Mojo function
generic_cross_attention_kv_cache
generic_cross_attention_kv_cache[collection_t: KVCollectionT, dtype: DType, //, target: StringSlice[StaticConstantOrigin], mask_str: StringSlice[StaticConstantOrigin], local_window_size: Int = -1](q: LayoutTensor[dtype, q.layout, q.origin, element_layout=q.element_layout, layout_int_type=q.layout_int_type, linear_idx_type=q.linear_idx_type, masked=q.masked, alignment=q.alignment], q_input_row_offsets: LayoutTensor[DType.uint32, q_input_row_offsets.layout, q_input_row_offsets.origin, element_layout=q_input_row_offsets.element_layout, layout_int_type=q_input_row_offsets.layout_int_type, linear_idx_type=q_input_row_offsets.linear_idx_type, masked=q_input_row_offsets.masked, alignment=q_input_row_offsets.alignment], q_max_seq_len: LayoutTensor[DType.uint32, q_max_seq_len.layout, q_max_seq_len.origin, element_layout=q_max_seq_len.element_layout, layout_int_type=q_max_seq_len.layout_int_type, linear_idx_type=q_max_seq_len.linear_idx_type, masked=q_max_seq_len.masked, alignment=q_max_seq_len.alignment], kv_input_row_offsets: LayoutTensor[DType.uint32, kv_input_row_offsets.layout, kv_input_row_offsets.origin, element_layout=kv_input_row_offsets.element_layout, layout_int_type=kv_input_row_offsets.layout_int_type, linear_idx_type=kv_input_row_offsets.linear_idx_type, masked=kv_input_row_offsets.masked, alignment=kv_input_row_offsets.alignment], kv_collection: collection_t, layer_idx: UInt32, scale: Float32, output: LayoutTensor[dtype, output.layout, output.origin, element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment], context: DeviceContextPtr, sink_weights: OptionalReg[LayoutTensor[dtype, Layout.row_major(-1), ImmutAnyOrigin]] = None)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!