Triton-Ops

A handwritten Triton operator library (baseline reference group): Triton implementations of 40 public operators, verified for correctness on both NVIDIA (CUDA) and Cambricon (MLU), with performance compared against native PyTorch. Its purpose is to provide a correctness and performance reference baseline for operators generated by NineToothed (九齿).

Platform notes and performance baseline

See MLU_NOTES.md: Cambricon MLU platform characteristics (grid limit, num_warps preference, dot-SIMD path, etc.), precision conventions (IEEE reference / fp16 relative error ≤ 1%), and the detailed T/P performance baseline of the 40 operators on MLU590-M9.

Repository layout

Triton-Ops/
├── README.md                ← this file
├── MLU_NOTES.md             ← MLU platform notes + precision conventions + performance baseline
├── kernels/
│   ├── triton_impl.py       ← handwritten Triton operators (baseline body)
│   ├── mlu_matmul.py        ← MLU-specific matmul schedule
│   └── ninetoothed_impl.py  ← NineToothed operators (kept for history, not measured at this stage)
├── tests/
│   ├── test_correctness.py  ← correctness tests: checked against PyTorch reference implementations
│   └── test_accuracy_metrics.py
├── diagnostics/             ← precision/performance diagnostic scripts, model-shape micro-benchmarks, MLU run scripts
├── bench.py                 ← performance comparison: do_bench timing, results archived per operator
├── reference_precision.py   ← reference precision policy switch
└── conftest.py

Environments

Environment	Configuration	Purpose
Local machine	RTX 3060 Laptop (CUDA)	Daily development and regression
MLU server	MLU590-M9, torch 2.11 + torch_mlu 1.33.1 + triton 3.4.0 (CNNL 2.2.14, driver 6.5.52)	Cambricon-side acceptance baseline; performance and correctness numbers are quoted from this environment

How to run

Dependencies: PyTorch + Triton (torch_mlu on the MLU environment).

# Correctness tests (all; add -k <op> to run a single operator)
python -m pytest tests/ -v

# Performance comparison (all operators; selected ones: python bench.py addmm abs)
# Results are automatically archived per operator to results/bench_<op>.md,
# with device/version/commit info in the header
# (results/ is a local artifact, not committed; reproducible at any time)
python bench.py

On the MLU server (Cambricon):

export TB_DEVICE=mlu            # device selection is controlled by the environment variable
python -m pytest tests/ -v      # full correctness suite
python bench.py <ops...> --dtype both --repeats 3   # performance comparison
# The serial full pipeline (correctness + performance + accuracy audit) for key
# operators lives in diagnostics/run_mlu.sh

How to add a new operator (standard procedure)

1. Write the Triton kernel in kernels/triton_impl.py (follow the existing add/abs);
2. Add a test in tests/test_correctness.py: check against the native PyTorch implementation (torch.allclose, relaxed tolerance for fp16), and make it pass;
3. Register the operator in the OPS table of bench.py (PyTorch + Triton entries), run the benchmark;
4. Commit (one commit per operator).

The NineToothed version belongs to a later stage. kernels/ninetoothed_impl.py keeps the historical implementation; at this stage bench only measures PyTorch/Triton, to avoid importing NineToothed and triggering device initialization.

Conventions

• Correctness before speed: performance is only discussed once all correctness tests pass;
• All performance numbers come from triton.testing.do_bench (automatic warmup + repeated statistics);
• Every number carries: shape, dtype, device (RTX 3060 Laptop / MLU590-M9), time unit ms, and converted bandwidth GB/s.

Measurement and precision conventions

• python bench.py mm addmm --dtype both --repeats 3: fp32/fp16 of the same operator are saved as sections in one result file, recording three rounds of p20/p50/p80, shapes, strides and source SHA256. Tables use the median of the three rounds' p50.
• The mixed-tolerance correctness test (fp16 atol=rtol=1e-2) is kept as an engineering regression; it does not replace the formal acceptance criterion of relative error ≤ 1%. Formal error statistics are produced by diagnostics/accuracy_audit.py.
• mm/addmm on MLU use a dedicated schedule (kernels/mlu_matmul.py); CUDA keeps the original implementation. Model-shape micro-benchmarks do not imply that model integration is complete.
• Precision paths are aligned explicitly: by default TB_REFERENCE_PRECISION=ieee disables TF32 for PyTorch-reference matmul/conv, compared against Triton's input_precision="ieee". TB_REFERENCE_PRECISION=vendor-default reproduces the vendor default path; such results must keep the policy label.