MASE-Triton

Software-emulation & acceleration triton kernels for MASE.

Install

Please ensure you are using Python 3.11 or later, and run MASE-Triton on CUDA-enabled GPU.

PyPI

pip install mase-triton

Build from Source

1. Install uv

2. Build the package

   uv build

   The wheel file can be found in dist/ folder. You can install it by pip install path/to/wheel/file.whl

Functionality

• Random Bitflip
  • functional APIs: random bitflip function with backward support.
  • layers.py: subclasses of torch.nn.Module that can be used in neural networks.
    • RandomBitflipDropout
    • RandomBitflipLinear
• Optical Transformer
  • functional APIs: optical transformer function with backward support.
    • ot_quantize
    • ot_linear
    • ot_matmul
  • layers.py: subclasses of torch.nn.Module that can be used in neural networks.
    • OpticalTransformerLinear
• MXFP: Simulate MXFP formats on CPU & GPU using PyTorch & Triton.
  • functional
    • extract_mxfp_tensor: Cast a tensor to MXFP format (extracting the shared exponent and Minifloat elements).
    • compose_mxfp_tensor: Cast an MXFP tensor to FP format (composing MXFP components).
    • mxfp_linear: functional linear operation with MXFP support.
    • mxfp_matmul: functional matrix multiplication with MXFP support.
  • layers
    • MXFPLinearPTQ: Linear layer with MXFP support for post-training quantization (no back propagation support).
• Minifloat: Simulate minifloat formats on CPU & GPU using PyTorch & Triton.
  • functional
    • extract_minifloat_component: Extract minifloat components from a tensor.
    • compose_minifloat_component: Compose minifloat components back to a tensor.
    • quantize_dequantize: Quantize and dequantize tensors using minifloat format.
    • minifloat_linear: functional linear operation with minifloat support.
    • minifloat_matmul: functional matrix multiplication with minifloat support.
  • layers
    • MinifloatLinearPTQ: Linear layer with minifloat support for post-training quantization (no back propagation support).

Dev

1. Install uv

2. Install dependencies for development

   uv sync