Skip to main content

mlx_native/ops/
mod.rs

1//! GPU kernel host-side dispatch functions.
2//!
3//! Each submodule implements dispatch for a specific kernel family.
4
5pub mod argmax;
6pub mod argsort;
7pub mod top_k;
8pub mod copy;
9pub mod dense_gemm;
10pub mod elementwise;
11pub mod embedding;
12pub mod gather;
13pub mod gather_bench;
14pub mod hadamard;
15pub mod hadamard_quantize_kv;
16pub mod encode_helpers;
17pub mod fused_head_norm_rope;
18pub mod fused_norm_add;
19pub mod fused_residual_norm;
20pub mod gelu;
21pub mod kv_cache_copy;
22pub mod moe_dispatch;
23pub mod moe_gate;
24pub mod quantized_matmul;
25pub mod quantized_matmul_ggml;
26pub mod quantized_matmul_id;
27pub mod quantized_matmul_id_ggml;
28pub mod rms_norm;
29pub mod rope;
30pub mod flash_attn_vec;
31pub mod flash_attn_vec_tq;
32pub mod fwht_standalone;
33pub mod sdpa;
34pub mod sdpa_sliding;
35pub mod softcap;
36pub mod softmax;
37pub mod softmax_sample;
38pub mod transpose;