Add Qwen vision loading and Metal inference
This commit is contained in:
Georg Bauer committed
1 parent
48c2f751b4
commit
e32954ac1b
30 files changed
+1687
-100
No files matched your search
@@ -0,0 +1,191 @@
|
||||
//! Vision-only dispatches from pinned MLX 0.32.2 normalization, conv and SDPA.
|
||||
//! Rust supplies the same layouts and parameters to the existing Metal library.
|
||||
use super::super::gpu::{Buffer, MetalConstant};
|
||||
use super::array::{Array, Dtype, Layout};
|
||||
use super::{bytes, dispatch_geometry, dispatch_specialized, ops, stream::Stream, tensor};
|
||||
|
||||
#[derive(Clone, Copy, PartialEq)]
|
||||
pub(super) enum Kernel {
|
||||
LayerNorm,
|
||||
PatchConv,
|
||||
Attention,
|
||||
Linear,
|
||||
}
|
||||
|
||||
pub(super) fn call(
|
||||
kernel: Kernel,
|
||||
inputs: &[Array],
|
||||
shape: &[i32],
|
||||
stream: Stream,
|
||||
) -> Result<Array, String> {
|
||||
Array::make_operation_with_inputs(
|
||||
stream,
|
||||
inputs,
|
||||
Layout::new(shape, Dtype::BF16)?,
|
||||
ops::Operation::Vision(kernel),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn evaluate(kernel: Kernel, inputs: &[Array], out: &Array) -> Result<(), String> {
|
||||
if kernel == Kernel::Linear {
|
||||
return super::dense::evaluate_bias(&inputs[0], &inputs[1], Some(&inputs[2]), out);
|
||||
}
|
||||
let nbytes = out.layout().nbytes() as u64;
|
||||
out.set_data(Buffer::mtplx_bytes(nbytes)?)?;
|
||||
let buffers = inputs.iter().map(Array::buffer).collect::<Vec<_>>();
|
||||
let ob = out.buffer();
|
||||
let mut bindings = inputs
|
||||
.iter()
|
||||
.zip(&buffers)
|
||||
.enumerate()
|
||||
.map(|(i, (a, b))| {
|
||||
tensor(i as u32, b)
|
||||
.at_byte_offset(a.offset())
|
||||
.input(a.data_size())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
bindings.push(tensor(inputs.len() as u32, &ob).output(out.data_size()));
|
||||
match kernel {
|
||||
Kernel::Linear => {
|
||||
super::dense::evaluate_bias(&inputs[0], &inputs[1], Some(&inputs[2]), out)
|
||||
}
|
||||
Kernel::LayerNorm => {
|
||||
let axis = out.layout().dim(-1)? as u32;
|
||||
let rows = out.layout().size() as u32 / axis;
|
||||
let threads = axis.div_ceil(8).div_ceil(32) * 32;
|
||||
let eps = 1e-6_f32;
|
||||
let stride = 1_u32;
|
||||
bindings.extend([
|
||||
bytes(4, &eps),
|
||||
bytes(5, &axis),
|
||||
bytes(6, &stride),
|
||||
bytes(7, &stride),
|
||||
]);
|
||||
dispatch_geometry(
|
||||
"layer_normbfloat16",
|
||||
&bindings,
|
||||
&[],
|
||||
[rows * threads, 1, 1],
|
||||
[threads, 1, 1],
|
||||
true,
|
||||
)
|
||||
}
|
||||
Kernel::PatchConv => {
|
||||
// Inputs are contiguous [patches,2,16,16,16], weights
|
||||
// [1152,2,16,16,16]. Original Conv3d pads RGB channels to 16.
|
||||
#[repr(C)]
|
||||
struct Params {
|
||||
n: i32,
|
||||
c: i32,
|
||||
o: i32,
|
||||
input: [i32; 3],
|
||||
weight: [i32; 3],
|
||||
output: [i32; 3],
|
||||
stride: [i32; 3],
|
||||
pad: [i32; 3],
|
||||
kdil: [i32; 3],
|
||||
idil: [i32; 3],
|
||||
is: [i64; 5],
|
||||
ws: [i64; 5],
|
||||
os: [i64; 5],
|
||||
groups: i32,
|
||||
flip: u8,
|
||||
padding: [u8; 3],
|
||||
}
|
||||
let n = out.layout().dim(0)?;
|
||||
let params = Params {
|
||||
n,
|
||||
c: 16,
|
||||
o: 1152,
|
||||
input: [2, 16, 16],
|
||||
weight: [2, 16, 16],
|
||||
output: [1, 1, 1],
|
||||
stride: [2, 16, 16],
|
||||
pad: [0; 3],
|
||||
kdil: [1; 3],
|
||||
idil: [1; 3],
|
||||
is: [8192, 4096, 256, 16, 1],
|
||||
ws: [8192, 4096, 256, 16, 1],
|
||||
os: [1152, 1152, 1152, 1152, 1],
|
||||
groups: 1,
|
||||
flip: 0,
|
||||
padding: [0; 3],
|
||||
};
|
||||
let gemm = [n, 1152, 8192, 512, 16, 16, 16, -8160, 18, (n + 31) / 32, 0];
|
||||
bindings.extend([bytes(3, ¶ms), bytes(4, &gemm)]);
|
||||
dispatch_specialized(
|
||||
"implicit_gemm_conv_3d_bfloat16_bm32_bn64_bk16_wm2_wn2_filter_s",
|
||||
&bindings,
|
||||
&[],
|
||||
[18, ((n + 31) / 32) as u32, 1],
|
||||
[32, 2, 2],
|
||||
)
|
||||
}
|
||||
Kernel::Attention => {
|
||||
#[repr(C)]
|
||||
struct Params {
|
||||
b: i32,
|
||||
h: i32,
|
||||
d: i32,
|
||||
ql: i32,
|
||||
kl: i32,
|
||||
gqa: i32,
|
||||
scale: f32,
|
||||
nq: i32,
|
||||
nk: i32,
|
||||
aq: i32,
|
||||
ak: i32,
|
||||
rq: i32,
|
||||
rk: i32,
|
||||
off: i32,
|
||||
strides: [[i64; 3]; 4],
|
||||
}
|
||||
let shape = out.layout().shape().to_vec();
|
||||
let n = shape[2];
|
||||
let mut strides = [[0; 3]; 4];
|
||||
for (s, a) in strides
|
||||
.iter_mut()
|
||||
.zip(inputs.iter().chain(std::iter::once(out)))
|
||||
{
|
||||
s.copy_from_slice(&a.layout().strides()[..3]);
|
||||
}
|
||||
let params = Params {
|
||||
b: 1,
|
||||
h: 16,
|
||||
d: 80,
|
||||
ql: n,
|
||||
kl: n,
|
||||
gqa: 1,
|
||||
scale: (72_f64.powf(-0.5)) as f32,
|
||||
nq: ((n + 31) / 32),
|
||||
nk: ((n + 31) / 32),
|
||||
aq: n / 32,
|
||||
ak: n / 32,
|
||||
rq: n % 32,
|
||||
rk: n % 32,
|
||||
off: 0,
|
||||
strides,
|
||||
};
|
||||
let constants = [
|
||||
(200, n % 32 == 0),
|
||||
(201, n % 32 == 0),
|
||||
(300, false),
|
||||
(301, false),
|
||||
(302, false),
|
||||
]
|
||||
.map(|(index, v)| MetalConstant {
|
||||
index,
|
||||
value: u32::from(v),
|
||||
kind: 0,
|
||||
});
|
||||
bindings.push(bytes(4, ¶ms));
|
||||
dispatch_specialized(
|
||||
"steel_attention_bfloat16_bq32_bk32_bd80_wm4_wn1_maskbfloat16",
|
||||
&bindings,
|
||||
&constants,
|
||||
[((n + 31) / 32) as u32, 16, 1],
|
||||
[32, 4, 1],
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user