forked from ccf-ai-infra/Intro-ops
46 lines
1.1 KiB
YAML
46 lines
1.1 KiB
YAML
name: softmax
|
|
kind: normalization
|
|
python_module: operator_runtime.ops.softmax
|
|
torch_reference: torch.softmax
|
|
backends:
|
|
nvidia:
|
|
enabled_by_default: true
|
|
status: runnable
|
|
sources:
|
|
- ops/softmax/nvidia/softmax_cuda.cu
|
|
headers:
|
|
- ops/softmax/nvidia/softmax_cuda.h
|
|
- ops/softmax/nvidia/kernel.cuh
|
|
symbols:
|
|
create: oprt_create_softmax_descriptor_nvidia
|
|
workspace: oprt_get_softmax_workspace_size_nvidia
|
|
execute: oprt_execute_softmax_nvidia
|
|
destroy: oprt_destroy_softmax_descriptor_nvidia
|
|
dtypes: [float32]
|
|
supports:
|
|
broadcast: false
|
|
strided: false
|
|
tilelang:
|
|
enabled_by_default: true
|
|
status: runnable
|
|
module: ops.softmax.tilelang.softmax_tl
|
|
dtypes: [float32]
|
|
supports:
|
|
broadcast: false
|
|
strided: false
|
|
metax:
|
|
enabled_by_default: false
|
|
status: stub
|
|
headers:
|
|
- ops/softmax/metax/softmax_metax.h
|
|
tolerances:
|
|
float32: {atol: 1e-5, rtol: 1e-5}
|
|
benchmark:
|
|
against_torch: true
|
|
default_cases:
|
|
- name: rowwise_1024x1024
|
|
performance_model:
|
|
bound: mixed
|
|
bytes: "5 * numel * elem_bytes"
|
|
flops: "4 * numel"
|