DOCSEARCHbyAlgolia
Submit Docs

ONNX Runtime

repository·main·Indexed 12 days ago

https://github.com/microsoft/onnxruntime

A high-performance, cross-platform machine learning accelerator for model inference and training. It supports frameworks like PyTorch and TensorFlow, optimizing execution across diverse hardware. The project includes a JavaScript API (v1.30.0) and C# managed bindings via Microsoft.ML.OnnxRuntime and Microsoft.ML.OnnxRuntime.Managed NuGet packages.

Tokens
312.2K
Snippets
483
Records
1K
Agent score
97%

What's inside ONNX Runtime

  1. Overview of ONNX Runtime JavaScript APIs

    main

    The ONNX Runtime JavaScript API is distributed across several NPM projects depending on your target environment:

    • onnxruntime-common: Pure JavaScript code designed to run in both Node.js and browsers. Contains core interfaces like InferenceSession, Tensor/OnnxValue, and Backend registration.
    • onnxruntime-node: Node.js bindings that allow Node.js users to consume ONNX Runtime via a native binding.
    • onnxruntime-web: A library for running ONNX models in web browsers (the successor to ONNX.js).
    • onnxruntime-react-native: Support for React Native environments.

  2. Overview of ONNX Runtime capabilities

    main

    ONNX Runtime is a cross-platform machine-learning accelerator designed for both inference and training.

    Inference

    Enables high-performance model execution for models exported from frameworks like PyTorch, TensorFlow/Keras, scikit-learn, LightGBM, and XGBoost. It optimizes performance by leveraging hardware accelerators, graph optimizations, and transforms across various operating systems and hardware drivers.

    Training

    Accelerates model training, specifically for transformer models on multi-node NVIDIA GPUs. For existing PyTorch training scripts, acceleration can be enabled with a one-line addition.

  4. Overview of ONNX Runtime Web

    main
    ONNX Runtime Web is a JavaScript library designed to run ONNX models directly in web browsers or in Node.js environments. It leverages WebAssembly (WASM) for CPU execution and various web standards for GPU acceleration, allowing developers to perform machine learning inference on the client side. This approach reduces server-client communication overhead and enhances user privacy by keeping data on the local device.

  5. Overview of MoE and QMoE CUDA Operators

    main

    ONNX Runtime provides two specialized Mixture-of-Experts (MoE) operators for the CUDA execution provider in the com.microsoft domain:

    1. MoE: Standard (non-quantized) Mixture-of-Experts. It supports FP16, BF16, and FP32 weights.
    2. QMoE: Quantized Mixture-of-Experts. It supports various quantized weight formats (INT4, INT8, FP8, MXFP4) and activation types (FP16, BF16, FP8).

    Both operators utilize a shared CUTLASS-based runner, routing engine, and sort/permute infrastructure. The primary difference lies in how the weight tensors are interpreted.

    The execution pipeline follows this flow: input tokens $\rightarrow$ router (top-k softmax) $\rightarrow$ permute by expert $\rightarrow$ GEMM1 (per-expert) $\rightarrow$ activation (SiLU/GeLU/ReLU/SwiGLU) $\rightarrow$ GEMM2 (per-expert) $\rightarrow$ un-permute $\rightarrow$ weighted sum.

  6. Explore ONNX Runtime Node.js sample implementations

    main

    The Node.js samples repository contains several demonstrations of the ONNX Runtime Node.js binding:

    • Basic Usage (01_basic-usage): Demonstrates the fundamental workflow of using the ONNX Runtime Node.js binding.
    • Create Tensor (02_create-tensor): Demonstrates how to create and manipulate tensors.
    • Create InferenceSession (04_create-inference-session): Shows the different methods available for initializing an InferenceSession.

  8. PagedAttention CUDA Implementation Capabilities

    main

    The com.microsoft::PagedAttention operator for the CUDA Execution Provider currently supports the following features:

    • Query/KV Handling: Packed varlen Q/K/V and packed QKV; Block KV cache with in-place updates (ReshapeAndCache); block_table gather for reads.
    • Mathematical Operations: RoPE (do_rotary, rotary_interleaved, cos_cache/sin_cache); scale; softcap.
    • Attention Mechanisms: local_window_size (wired to Flash varlen and MEA).
    • Backends: FlashAttention varlen with a fallback to Memory-Efficient Attention (CUTLASS fMHA).
    • Supported Data Types: float16, bfloat16 only.

    Current Limitations:

    • Batch Size: batch_size <= 256 due to cub::BlockScan constraints.
    • Block Size: While validation now accepts any power-of-two block_size >= 16, certain constraints may trigger a fallback to the memory-efficient backend if FlashAttention tiling requirements are not met.

  9. Analyze ONNX Runtime profile outputs with Profile Explorer

    main

    Profile Explorer is a Python script designed to simplify the analysis of profile outputs (JSON files) generated by ONNX Runtime. It provides visibility into model performance by identifying which operators and kernels are consuming the most execution time.

    Key features include:

    • Basic statistics: Aggregated performance data over operators and kernels.
    • Operator-kernel correlation: Statistics showing the relationship between specific operators and the underlying kernels they execute.

  10. Key differences between ONNX Runtime MoE/QMoE and TensorRT-LLM

    main

    The MoE (Mixture of Experts) and QMoE (Quantized Mixture of Experts) CUDA kernels in ONNX Runtime are derived from TensorRT-LLM (CUTLASS 4.4.2) but include several significant modifications and optimizations for ONNX Runtime usage:

    Optimizations and Features

    • Pre-packed ZP/Bias optimization: Biases are derived offline as (K − ZP) × scale to allow the kernel to handle asymmetric quantization without extra subtraction during runtime.
    • SwiGLU interleaving: Activation kernels support interleaved Gate/Value weights.
    • Sparse Mixer support: Enabled via the use_sparse_mixer attribute.
    • Dynamic min_dim selection: The supportsTmaWarpSpecialized() method is exposed on the CutlassMoeFCRunnerInterface, allowing for dynamic dimension selection without needing to know the concrete template type at the call site.
    • MXFP4 support: The QMoE schema and runner have been extended to accept MXFP4 weights along with per-expert global scales and ue8m0 block scales.
    • WFP4AFP8 (SM100+): Includes MoeGemmRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, ...> which performs BF16/FP16 to MXFP8 quantization in the expandInputRowsKernel step.

    Removed Features

    To streamline the implementation for ONNX Runtime, the following TensorRT-LLM features were removed:

    • LoRA parameters (use_lora, LoraParams)
    • Min-latency mode (MoeMinLatencyParams)
    • AllToAll MoE paths (enable_alltoall)
    • DeepSeek FP8 block-scale GEMM mode (use_deepseek_fp8_block_scale, BlockScaleParams)
    • Standalone GEMM directories (relevant logic is inlined into the MoE runner).

  11. Understand the ORT Model Package Integration

    main

    The ORT Model Package Integration provides the glue between the standalone model_package library and ONNX Runtime (ORT). It enables loading model packages, selecting the best execution provider (EP) variant for a given runtime environment, and creating an OrtSession from a chosen variant.

    Key responsibilities include:

    • Parsing the executor_info["ort"] payload: Extracting session and provider options from the package.
    • Variant Selection: Running an algorithm to pick the highest-scoring variant based on the requested Execution Providers and hardware compatibility.
    • Session Creation: Providing experimental C APIs to create sessions directly from package components.

  12. Phasing and Implementation Status of PagedAttention CUDA

    main

    The PagedAttention CUDA implementation is being rolled out in phases. As of the current status, phases P0 through P4 are implemented (with the exception of .Alias registration), and P5 is partially implemented.

    Implementation Roadmap

    • P0 (Foundation): Defect fixes, kernel definition aliases, and sliding-window semantic verification.
    • P1 (Paging primitives): Introduction of slot_mapping and sliding-window block pruning/cache.
    • P2 (Model coverage): Support for head_sink (via LSE epilogue) and fused QK-Norm.
    • P3 (Memory): Quantized cache support (INT8/FP8) with PER_TENSOR and PER_CHANNEL scaling.
    • P4 (MLA correctness): Support for Multi-Head Latent Attention (MLA) via kv_cache_layout="LATENT", v_head_size, and rotary_offset attributes.
    • P5 (Performance): Paged decode kernels with in-kernel dequantization, softcap, and CUDA-graph-capturability (removal of D→H sync).
    • P6 (Completeness): Support for query_positions, attention_bias, and output_qk.

    Model Support

    • P0–P2: GPT-OSS, Qwen3, Gemma 2/3, Llama, Mistral, and Phi.
    • P3 & P5: Optimized for throughput-oriented serving.
    • P4–P5: DeepSeek-V2/V3/R1.