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 gpu_sample;
7pub mod argsort;
8pub mod moe_mm_id_map0;
9pub mod top_k;
10pub mod copy;
11pub mod cumsum;
12pub mod dense_gemm;
13pub mod dense_gemv_bf16;
14pub mod dense_mm_bf16;
15pub mod dense_mm_f16;
16pub mod dense_mm_f32_f32;
17pub mod dequant_to_f16;
18pub mod elementwise;
19pub mod embedding;
20pub mod gather;
21pub mod gather_bench;
22pub mod hadamard;
23pub mod hadamard_quantize_kv;
24pub mod encode_helpers;
25pub mod fused_dual_proj_q4_0;
26pub mod fused_gate_up_silu_iq4_nl;
27#[allow(non_snake_case)]
28pub mod fused_gate_up_silu_q4_K;
29#[allow(non_snake_case)]
30pub mod fused_gate_up_silu_q5_K;
31#[allow(non_snake_case)]
32pub mod fused_gate_up_silu_q6_K;
33pub mod fused_gate_up_silu_q8_0;
34pub mod fused_head_norm_rope;
35pub mod fused_norm_add;
36pub mod fused_residual_norm;
37pub mod gelu;
38pub mod kv_cache_copy;
39pub mod l2_norm;
40pub mod log_elementwise;
41pub mod row_sum;
42pub mod moe_dispatch;
43pub mod moe_gate;
44pub mod moe_softmax_topk;
45pub mod mul_mv_ext;
46pub mod moe_weighted_reduce;
47pub mod qkv_split;
48pub mod repeat_tiled;
49pub mod quantized_matmul;
50pub mod quantized_matmul_ggml;
51pub mod quantized_matmul_id;
52pub mod quantized_matmul_id_ggml;
53pub mod rms_norm;
54pub mod rope;
55pub mod rope_multi;
56pub mod rope_train;
57pub mod vision_2d_rope;
58pub mod scale_mask_softmax;
59pub mod sigmoid_mul;
60pub mod silu_mul;
61pub mod compute_g_beta;
62pub mod ssm_norm_gate;
63pub mod flash_attn_prefill;
64pub mod flash_attn_prefill_blk;
65pub mod flash_attn_train;
66pub mod flash_attn_prefill_d512;
67pub mod flash_attn_prefill_mask;
68pub mod flash_attn_vec;
69pub mod flash_attn_vec_tq;
70// ADR-037 Phase E1.1 (2026-05-22) — EAGLE-3 + dynamic tree port.
71pub mod tree_attention;
72pub mod flash_attn_vec_tq_hb;
73pub mod flash_attn_vec_hybrid;
74pub mod flash_attn_vec_peer_port_f16;
75pub mod flash_attn_vec_reduce_tq_hb_undo;
76pub mod fwht_standalone;
77pub mod chunk_gated_delta_rule;
78pub mod chunk_gated_delta_rule_bank_split;
79pub mod chunk_gated_delta_rule_tri_solve_invert;
80pub mod gated_delta_net;
81pub mod gated_delta_net_decode;
82pub mod gated_delta_net_chunk;
83pub mod gated_delta_net_chunk_o;
84pub mod gated_delta_net_kkt;
85pub mod gated_delta_net_recompute_wu;
86pub mod tq_dequantize_kv;
87pub mod sdpa;
88pub mod sdpa_decode;
89pub mod sdpa_sliding;
90pub mod softcap;
91pub mod softmax;
92pub mod softmax_backward;
93pub mod softmax_sample;
94pub mod ssm_conv;
95pub mod conv1d_depthwise_causal;
96pub mod exp_elementwise;
97pub mod outer_product;
98pub mod take_along_axis;
99pub mod divide_elementwise;
100pub mod sqrt_elementwise;
101pub mod transpose;
102pub mod tri_solve;
103pub mod qdq_legacy;
104pub mod rms_norm_backward;
105pub mod slice_concat_2d;
106pub mod silu_backward;
107pub mod embedding_autograd;
108pub mod adam_update;
109pub mod qdq_affine;
110pub mod qmm_affine;
111// ADR-021 Qwen3-VL ViT prelude GPU port (5 new kernels):
112pub mod im2col_2d_3ch;
113pub mod add_bias_row_2d;
114pub mod bilinear_resize_2d;
115pub mod block_merge_2x2;
116pub mod feature_concat;
mlx_native/ops/mod.rs

mlx_native/ops/
mod.rs
