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 cumsum;
10pub mod dense_gemm;
11pub mod dense_mm_bf16;
12pub mod elementwise;
13pub mod embedding;
14pub mod gather;
15pub mod gather_bench;
16pub mod hadamard;
17pub mod hadamard_quantize_kv;
18pub mod encode_helpers;
19pub mod fused_head_norm_rope;
20pub mod fused_norm_add;
21pub mod fused_residual_norm;
22pub mod gelu;
23pub mod kv_cache_copy;
24pub mod l2_norm;
25pub mod moe_dispatch;
26pub mod moe_gate;
27pub mod quantized_matmul;
28pub mod quantized_matmul_ggml;
29pub mod quantized_matmul_id;
30pub mod quantized_matmul_id_ggml;
31pub mod rms_norm;
32pub mod rope;
33pub mod rope_multi;
34pub mod scale_mask_softmax;
35pub mod sigmoid_mul;
36pub mod flash_attn_prefill;
37pub mod flash_attn_prefill_blk;
38pub mod flash_attn_prefill_d512;
39pub mod flash_attn_prefill_mask;
40pub mod flash_attn_vec;
41pub mod flash_attn_vec_tq;
42pub mod flash_attn_vec_tq_hb;
43pub mod fwht_standalone;
44pub mod gated_delta_net;
45pub mod tq_dequantize_kv;
46pub mod sdpa;
47pub mod sdpa_sliding;
48pub mod softcap;
49pub mod softmax;
50pub mod softmax_sample;
51pub mod ssm_conv;
52pub mod transpose;
53pub mod tri_solve;