Hugging Face's logo

replicate
/
gpt-oss-metal-kernels

Kernels
Model card Files
xet
 Community
gpt-oss-metal-kernels / gptoss_kernels /source /include /internal /kernel-args.h
medmekk's picture
medmekk
Upload folder using huggingface_hub
51250cb unverified
Raw History Blame Contribute Delete
4.32 kB
#pragma once
#if !defined(__METAL_VERSION__)
#include <stdint.h>
#endif
// TODO(ibahmed): specalize using metal function constants.
#define QKV_Bm 64
#define QKV_Bn 64
#define QKV_Bk 32
#define QKV_Sg_Bm 32
#define QKV_Sg_Bn 32
#define ATTN_OUTPUT_Bm 32
#define ATTN_OUTPUT_Bn 64
#define ATTN_OUTPUT_Bk 64
#define ATTN_OUTPUT_Sg_Bm 32
#define ATTN_OUTPUT_Sg_Bn 16
#define MLP_GATE_Bm 64
#define MLP_GATE_Bn 16
#define MLP_GATE_Bk 64
#define MLP_GATE_Sg_Bm 16
#define MLP_GATE_Sg_Bn 16
#define MOE_DENSE_MATMUL_SWIGLU_Bm 32
#define MOE_DENSE_MATMUL_SWIGLU_Bn 64
#define MOE_DENSE_MATMUL_SWIGLU_Bk 16
#define MOE_DENSE_MATMUL_SWIGLU_Sg_Bm 32
#define MOE_DENSE_MATMUL_SWIGLU_Sg_Bn 16
#define MOE_DENSE_MATMUL_Bm 32
#define MOE_DENSE_MATMUL_Bn 64
#define MOE_DENSE_MATMUL_Bk 16
#define MOE_DENSE_MATMUL_Sg_Bm 32
#define MOE_DENSE_MATMUL_Sg_Bn 16
struct gptoss_expert_prediction {
uint32_t expert_id;
float score;
};
struct gptoss_control {
uint32_t abort;
};
struct gptoss_topk_args {
uint32_t num_vecs_per_token;
};
struct gptoss_sdpa_args {
uint32_t qkv_dim;
uint32_t num_kv_tokens;
uint32_t kv_stride;
uint32_t window;
};
struct gptoss_u32_fill_random_args {
uint64_t num_vecs_per_threadgroup;
uint64_t num_vecs;
uint64_t offset;
uint64_t seed;
};
struct gptoss_f32_fill_random_args {
uint64_t num_vecs_per_threadgroup;
uint64_t num_vecs;
uint64_t offset;
uint64_t seed;
float scale;
float bias;
};
struct gptoss_accumulate_args {
uint32_t num_vecs_per_expert;
uint32_t num_vecs_per_threadgroup;
uint32_t num_vecs;
};
struct gptoss_convert_args {
uint64_t num_vecs_per_threadgroup;
uint64_t num_vecs;
};
struct gptoss_embeddings_args {
uint32_t num_vecs;
};
struct gptoss_rmsnorm_args {
uint32_t num_vecs;
float num_channels;
float epsilon;
};
struct gptoss_matmul_args {
uint32_t num_column_vecs;
uint32_t num_rows;
uint32_t add;
};
struct gptoss_dense_matmul_args {
uint32_t m;
uint32_t n;
uint32_t k;
};
struct gptoss_scatter_args {
uint32_t tokens;
uint32_t active_experts_per_token;
uint32_t token_stride;
};
struct gptoss_moe_dense_matmul_swiglu_args {
uint32_t k;
uint32_t n;
uint32_t weight_blocks_expert_stride_bytes;
uint32_t weight_scales_expert_stride_bytes;
uint32_t bias_expert_stride_bytes;
float swiglu_min;
float swiglu_max;
};
struct gptoss_moe_dense_matmul_args {
uint32_t k;
uint32_t n;
uint32_t weight_blocks_expert_stride_bytes;
uint32_t weight_scales_expert_stride_bytes;
uint32_t bias_expert_stride_bytes;
};
struct gptoss_expert_routing_metadata_args {
uint32_t tokens;
uint32_t num_experts;
};
struct gptoss_gather_args {
uint32_t tokens;
uint32_t active_experts_per_token;
uint32_t token_stride;
};
struct gptoss_unembedding_args {
uint32_t num_column_vecs;
uint32_t num_rows_per_threadgroup;
uint32_t num_rows;
};
struct gptoss_moe_matmul_swiglu_args {
uint32_t num_column_vecs;
uint32_t num_rows;
uint32_t num_active_experts;
uint32_t weight_expert_stride; // in bytes
uint32_t output_expert_stride; // in elements
float swiglu_min;
float swiglu_max;
};
struct gptoss_moe_matmul_args {
uint32_t num_column_vecs;
uint32_t num_rows;
uint32_t num_active_experts;
uint32_t input_expert_stride; // in blocks of 32 elements
uint32_t weight_expert_stride; // in bytes
uint32_t output_expert_stride; // in elements
};
struct gptoss_rope_args {
uint32_t token_stride;
uint32_t token_offset;
float freq_scale;
float interpolation_scale;
float yarn_offset;
float yarn_scale;
float yarn_multiplier;
};
struct gptoss_qkv_args {
uint32_t num_column_vecs;
uint32_t num_rows;
uint32_t token_offset;
float freq_scale;
float interpolation_scale;
float yarn_offset;
float yarn_scale;
float yarn_multiplier;
uint32_t max_tokens;
};
struct gptoss_softmax_args {
uint32_t num_vecs;
uint32_t num_vecs_per_threadgroup;
uint32_t max_threadgroups;
float temperature;
};
struct gptoss_sample_args {
uint64_t rng_seed;
uint32_t rng_offset;
uint32_t num_blocks;
uint32_t num_dims;
uint32_t num_dims_per_block;
};