Scaling Laws and Efficient Inference for Ternary Language Models
Authors: Tejas Vaidhya, Ayush Kaushal, Vineet Jain, Francis Couture-Harpin, Prashant Shishodia, Majid Behbahani, Irina Rish, Yuriy Nevmyvaka
Accepted at: acl25

Scaling Laws and Efficient Inference for Ternary Language Models

Tejas Vaidhya1,2,3*, Ayush Kaushal1*, Vineet Jain2,4, Francis Couture-Harpin5, Prashant Shishodia6, Majid Behbahani7, Yuriy Nevmyvaka7, Irina Rish1,2,3
1Nolano AI   2Mila- Quebec AI institute   3Université de Montréal   4McGill University   5École de technologie supérieure, Université du Québec   6Google, India   7Morgan Stanley

Abstract

Large language models (LLMs) are increasingly used across research and industry applications, yet their inference efficiency remains a significant challenge. As the computational power of modern GPU architectures continuously improves, their memory bandwidth and capacity have not scaled proportionally, creating a critical bottleneck during inference. To address this, we investigate ternary language models (TriLMs) that employ quantization-aware training to significantly reduce memory requirements. We first analyze the scalability of TriLMs by conducting a scaling law analysis, revealing that TriLMs benefit more from increasing training data than from scaling model parameters. Based on this observation, we introduce TriTera, an open suite of TriLMs trained on up to 1.2 trillion tokens, demonstrating sustained performance gains at scale. Furthermore, to improve inference efficiency, we propose novel 2-bit and 1.6-bit packing schemes for ternary weights, which demonstrate accelerated inference across various CPU architectures. Also, building on the 2-bit packing, we develop a GPU kernel called TriRun that accelerates end-to-end model inference by up to 5 times compared to floating-point baselines. To encourage further exploration and development of TriLMs, we will release the TriTera suite and TriRun inference kernels. Overall, our work lays the foundation for building and deploying efficient LLMs, providing a valuable resource for the research community.

1 Introduction

Large language models (LLMs) (Radford et al., 2019; Zhang et al., 2022; Touvron et al., 2023) have become increasingly pivotal in both research and industry. Beyond their broad utility, their capabilities during inference with additional compute demonstrate the potential to enable advancements in reasoning and agentic tasks (Sardana et al., 2024; Singh et al., 2024; Wei et al., 2023). As the demand for efficient and scalable inference grows (Zhou et al., 2024), significant efforts have been directed toward reducing inference costs and latency (Dettmers et al., 2022a; Frantar et al., 2023; Sheng et al., 2023). While, the computational power of GPUs has improved rapidly, advancements in memory capacity and bandwidth have lagged behind (Gholami et al., 2024; Kaushal et al., 2024). This disparity has made memory-related bottlenecks a predominant challenge during LLM inference, where memory usage and bandwidth (driven by model size in bits) increasingly outweigh computational (FLOPs) limitations. While post-training quantization, combined with custom kernels for inference acceleration, has become widely adopted, its effectiveness in mitigating these bottlenecks remains limited. Specifically, post-training quantization is typically restricted to 4-bits and results in significant performance degradation beyond this threshold (Dettmers and Zettlemoyer, 2023).

Recent advancements in extreme low-bit language models (Kaushal et al., 2024; Wang et al., 2023; Ma et al., 2024) have shown that quantization-aware training allows ternary-weight models to achieve performance comparable to full-precision models (referred to as FloatLMs in this paper) at larger parameter scales. Additionally, ternary representations demonstrate superior bit-efficiency as they scale. However, critical gaps persist in understanding the scaling laws governing Ternary Language Models (TriLMs)—specifically, how TriLM performance is affected by training on much larger datasets or with many more parameters remains unanswered. Furthermore, the acceleration of inference in sub-4-bit models (e.g., ternary) remains unexplored, with most existing research limited to 4-bit quantization (Frantar et al., 2023; Dettmers et al., 2022a; Frantar et al., 2024; He et al., 2024). These limitations are compounded by the
Figure 1: Model performance (MMLU average accuracy) versus training FLOPs, considering only models with similar compute budgets and training tokens for a fair comparison (left); and end-to-end generation time speedup achieved by TriRun kernels over the PyTorch’s FP16 baseline (64 Input Tokens, 64 Output Tokens) on the NVIDIA L40S (right).

absence of a comprehensive suite of strong open-source models, suppressing innovation in post-training and broader research on extreme quantization. In this work, we aim to address these foundational challenges through the following contributions:

Scaling law for ternary language models. We conduct (in Section 2) a systematic study to explore the scaling properties of TriLMs, focusing on both the number of parameters and the volume of training tokens. Unlike previous works (Kaushal et al., 2024; Wang et al., 2023), which primarily examine parameter scaling, we demonstrate that increasing the number of tokens leads to a greater reduction in validation loss compared to increasing the number of parameters (see Section 2.2).

Effect of scaling pretraining tokens. We scale the TriLM models by pretraining them on 1.2T tokens (see Section 2.3), refered to as the TriTera family of models. Our results show that the 3B model continues to improve with up to 1.2T tokens, suggesting that TriLM remains effective even at higher token-to-parameter ratios. Additionally, it achieves competitive performance with FloatLMs for a given compute budget (see Figure 1, left).

Efficient packing mechanism for ternary weights. In Section 3, we propose efficient 1.6-bit and 2-bit packing schemes for ternary weights. We provide a theoretical analysis of these packing methods, along with the implementation of efficient kernels and benchmarking on a CPU (see Section 3.3 and Appendix F.3), demonstrating a significant acceleration in inference speed.

Efficient GPU kernels for ternary models. We introduce GPU kernels based on 2-bit packing schemes, which we call TriRun (in Section 4). We extensively benchmark its performance in model serving settings across various model sizes and different NVIDIA hardware (see Section 4.2 and Appendix G.6). Notably, we achieve up to a 7-8× speedup compared to PyTorch’s float16 kernels in high-batch settings (16-32 samples) for the ternary layer in transformer blocks of larger parameter (70B - 405B) models on the L40S GPU. Additionally, as shown in Figure 1 (right), our 70B model achieves a 4.9× end-to-end speedup (compared to float16) while running on a single L40S.

2 Scaling ternary models to 1T tokens

In this section, we study the scalability of pretraining ternary models. We begin by outlining our training setup, including details about the data, hardware scaling, and model architecture (see Section 2.1). We then analyze the scaling properties of ternary language models with respect to both parameters and training tokens, deriving a scaling law in Section 2.2. Based on insights from our scaling studies, we train a suite of models on up to 1.2 trillion tokens, which we call TriTera, and benchmark their performance in Section 2.3.

2.1 Training Details

Data. Our training corpus comprises a diverse mix of data from publicly available sources. To scale TriLMs (Kaushal et al., 2024), we trained on approximately 1.2 trillion tokens from ArXiv
Figure 2: Effect of scaling number of parameters (left) and number of training tokens (right) on final validation loss for TriLMs. The dotted lines show the power law derived in Equation (2).

(Clement et al., 2019), Cosmopedia-v2 (Ben Alal et al., 2024), PeS2o (Soldaini and Lo, 2023), Zyda-StarCoder-Git-Commits, Zyda-StarCoder-Languages (Tokpanov et al., 2024), FineWeb-Edu (Lozhkov et al., 2024). The dataset details are summarized in Table 1 and appendix §B.2. For tokenization, we employ the LLaMA tokenizer over previously used GPT-NeoX tokenizer.

Architecture. Our model follows a decoder-only transformer architecture (Vaswani et al., 2023), closely resembling the TriLM (Kaushal et al., 2024). Inspired by LLaMA (Touvron et al., 2023), it incorporates SwiGLU MLPs (Shazeer, 2020), RoPE (Su et al., 2023), multi-head attention, and bias-free layers. A key distinction is its ternary-weighted linear layers (-1, 0, 1) with a shared floating-point scale. Training maintains latent floating-point weights, applying on-the-fly ternarization in the forward pass, and the scale set to their absolute mean. Hyperparameters and additional pretraining details are provided in Appendix B.

Hardware and Scaling. We conduct our training experiments on the Frontier1 cluster. Each node comprises four AMD MI250X accelerators (Advanced Micro Devices, Inc., 2025), where each MI250X contains two Graphics Compute Dies (GCDs) operating as separate GPUs (Advanced Micro Devices, Inc., 2022). Within a node, GPUs are at most one hop away from one another, facilitating efficient intra-node communication. Our distributed training strategy is designed with this hardware architecture in mind. Similar to the ZeRO Stage 2 strategy (Rajbhandari et al., 2020), we shard the AdamW optimizer states and gradients, synchronizing model parameters after each update step. However, due to slower inter-node connectivity, sharding is performed only across devices within a node. This approach enables near-linear scaling up to 2,048 GPUs, as shown in Figure 7.

2.2 Scaling Laws for TriLMs

Experimental Setup. For this study (\( \leq 150B \) tokens), we use a SlimPajama subset from Shen et al. (2024), while the 1.2T-token dataset incorporating additional sources (Appendix 1). All other aspects follow the procedures outlined in Sections 2.1 regarding the pretraining of the models. We train and evaluate a suite of TriLM models, conducting a series of language model pretraining experiments across parameter sizes \( \in [99M, 190M, 390M, 560M, 1100M] \) (excluding embeddings) and dataset sizes \( \in [20, 40, 75, 150] \) billion tokens. In Section 2.3, we expand the training dataset size to 1.2 trillion tokens for models with 1.5 B, 2.5 B, and 3.6 B parameters.

Parametric Scaling Law. We derive the scaling law for ternary LLMs following the general form introduced in Hoffmann et al. (2022). In particular, we assume the following functional form for the validation loss \( \hat{L} \) as a function of model size \( N \) (number of parameters, in millions) and training data \( D \) (number of tokens, in billions),

\[
\hat{L}(N, D) \triangleq E + \frac{A}{N^{\alpha}} + \frac{B}{D^{\beta}},
\]

where the constant term includes the irreducible loss due to entropy of natural text, plus the error introduced by quantization. Based on the validation losses of the converged models, we fit the parameters \( \{E, A, \alpha, B, \beta\} \) (see Appendix C for evaluation of our fit). This provides a scaling law that describes how the validation loss of a ternary model changes with data and model size.

\[
\hat{L}(N, D) \approx 2.19 + \frac{4.73}{N^{0.32}} + \frac{5.18}{D^{0.81}}
\]
Figure 2 shows the final validation loss for different models against the number of parameters and training tokens. For each plot, we also substitute the corresponding value of \( D \) or \( N \) to get the scaling law equation for that setting. We discuss the implications of this law in more detail and compare with 16-bit models in Appendix C.

From Equation (2), we observe that increasing the number of tokens lowers the validation loss more effectively than increasing the number of parameters. This suggests that TriLM remains effective at high training token-to-parameter ratios. Based on these observations, we focus primarily on increasing the number of tokens to train our new family of models in the following section.

2.3 Effect of scaling training tokens

We pre-trained three TriLM models with 1.5 B, 2.5 B, and 3.6 B parameters (for simplicity, we refer to these models as 1B, 2B, and 3B throughout the paper) on a 1.2 trillion-token dataset (detailed in Appendix B.2), which we refer to as TriTera suite in this paper. The details of the parameters are provided in Table 3. Inspired by well-established model suites, such as those by (Groeneveld et al., 2024; Kaushal et al., 2024; Biderman et al., 2023), TriTera aims to provide robust baseline models to advance scientific research on TrILMs.

LLM Benchmarks Performance. We evaluate the TriTera suite of models on a variety of tasks testing commonsense and reasoning abilities, general knowledge, and mathematical problem-solving. A full description of the tasks is given in Appendix D.

![Performance of TriLM Models on MMLU](page_186_1097_563_312.png)

Figure 3: Average MMLU accuracy for TriTera and Spectra, with the dotted line representing LLaMA-1 7B (trained on 1.2T tokens). Note that LLaMA-3 (AI@Meta, 2024), trained on over 15 trillion tokens, is not included.

To understand the effect of scaling training tokens on downstream performance, we compared benchmark scores with the Spectra suite of models, which have comparable parameter sizes and were trained on 300B tokens. Figure 3 shows the average accuracy on the MMLU benchmark for both family of models, demonstrating consistently better performance across different parameter sizes. Full results on individual benchmarks are presented in Table 4.

3 Efficient packing of ternary weights

In this section, we propose weight-packing strategies and kernel implementations to enable the efficient deployment of ternary LLMs. We begin by formalizing the packing problem and then present two progressively optimized solutions. These solutions target effective 1.6-bit and 2-bit packing, supported by theoretical guarantees. Following this, we conduct a preliminary feasibility assessment on a CPU to evaluate the practicality of our approach.

Definition (Lossless Packing and Unpacking).
Let \( D = (d_1, d_2, \ldots, d_n) \) represent a sequence of ternary numbers, where \( d_i \in \{-1, 0, 1\} \). The packing process is a function \( P : \{-1, 0, 1\}^n \to B \), which maps the ternary sequence \( D \) to some binary representation \( B \), such that the original sequence \( D \) can be reconstructed from \( B \). The unpacking process is the inverse function \( U : B \to \{-1, 0, 1\}^n \), which reconstructs the original sequence \( D \) from \( B \). These processes satisfy the property of losslessness:

\[
U(P(D)) = D, \quad \forall D \in \{-1, 0, 1\}^n.
\]

3.1 Packing Strategy with effective 2 bits.

Packing/Encoding The packing process transforms each ternary value \( d_i \in \{-1, 0, 1\} \) into a digit \( d'_i \) by the mapping

\[
d'_i = d_i + 1,
\]

so that \( d'_i \in \{0, 1, 2\} \). These digits are then grouped into blocks of up to \( k \) values. Each block is encoded as a single integer using bitwise shifts. The packing function \( P(D) = (b'_0, b'_1, \ldots, b'_{m-1}) \) (with \( m = \lceil n/k \rceil \) for an original sequence \( D = \{d_0, d_1, \ldots, d_{n-1}\} \)) is defined as:

\[
b'_i = \sum_{j=0}^{k-1} (d'_{ki+j} \cdot 2^{2j}),
\]

If a block is not completely filled (when \( n \) is not a multiple of \( k \)), the remaining positions are padded
with 0, which map to 1 after the shift. Since each \( d'_i \) is in \( \{0, 1, 2\} \) and fits within 2 bits, each block uses \( 2k \) bits, giving an effective 2 bits per weight.

Unpacking/Decoding The unpacking process \( U(P(D)) = (d_0, d_1, \ldots, d_{n-1}) \) recovers the original ternary values from the packed representation. For each block, the decoding is defined as:

\[
d_{ki+j} = d'_{ki+j} - 1,
\]
where \( d'_{ki+j} = ((b'_i \gg 2j) \& 0x03) \),
for \( i \geq 0 \) and \( 0 \leq j < k \). Here, \( \gg \) denotes the bitwise right shift operation and \( \& \) denotes the bitwise AND operation (with 0x03 serving as a mask to extract 2 bits). This procedure ensures that each original ternary value \( d_i \in \{-1, 0, 1\} \) is accurately reconstructed from its packed form. Although each \( d'_i \) is constrained to three possible states, the packing allocates a total of \( 2k \) bits per block. However, the actual information content per block is only \( \log_2(3^k) = k \log_2(3) \) bits, which is strictly less than \( 2k \) bits (since \( \log_2(3) \approx 1.585 \)). In the following, we outline a general strategy for a better effective bit rate.

3.2 Packing Strategy with 1.6 effective bits.

Packing/Encoding: The packing process transforms each ternary value \( d_i \in \{-1, 0, 1\} \) into a base-3 digit (or trit) \( d'_i = d_i + 1 \), then groups the digits into blocks of up to k. Each block is encoded as a base-3 integer and normalized to fit within \([0, 2^p - 1]\), where p is an integer representing the number of bits allocated for each encoded block. The packing \( P(D) = (b'_1, b'_2, \ldots, b'_k) \) is defined as,

\[
b'_i = \left\lfloor \frac{\left( \sum_{j=0}^{k-1} d'_{ki+j} \cdot 3^{k-1-j} \right) \cdot 2^p + (3^k - 1)}{3^k} \right\rfloor,
\]
where \( d'_i = d_i + 1 \), and \( d_i \in \{-1, 0, 1\} \).

Here, \( k \) is the number of digits in each block (which may be less than \( k \) for the last block). The final packed byte array \( B \) is then constructed from the \( b'_i \) values.

Unpacking/decoding: The unpacking process \( U(P(D)) = (d_1, d_2, \ldots, d_n) \) is defined as:

\[
d_{ki+j} = d'_{ki+j} - 1,
\]
where \( x_i = \left\lfloor \frac{b_i \times 3^k - (3^k - 1) + (2^p - 1)}{2^p} \right\rfloor \),
and \( d'_{ki+j} = \left( \left\lfloor \frac{x_i}{3^{k-1-j}} \right\rfloor \right) \mod 3 \).

Here, k represents the number of digits in each block, typically equal to 5 for full blocks, though it may be fewer for the final block. For practical purposes, we recommend setting \( p = 8 \) and \( k = 5 \), as this configuration results in an effective packing of 1.6 bits — very close to the theoretical optimum for ternary data.

Theorem 1 (Correctness). Let \( D = (d_1, d_2, \ldots, d_n) \) be a sequence of ternary digits \( d_i \in \{-1, 0, 1\} \). When D is partitioned into blocks of size \( k \) and each block is encoded into a p-bit integer, the encoding and decoding operations \( P \) and \( U \) are lossless if and only if \( 2^p > 3^k \); that is, \( U(P(D)) = D \). Proof: See Appendix E.1.

Corollary (Injectivity). If \( 2^p > 3^k \), then the mapping \( P : \{-1, 0, 1\}^k \to [0, 2^p - 1] \) is injective i.e. \( \{d_j\} \neq \{d'_j\} \) implies \( P(\{d_j\}) \neq P(\{d'_j\}) \).

Near-Optimal Bits per Trit From an information-theoretic perspective, each trit (with values in \( \{-1, 0, 1\} \)) requires \( \log_2(3) \approx 1.58496 \) bits of entropy. To encode \( k \) trits without collision, we need a p-bit container with \( 2^p > 3^k \implies p > k \log_2(3) \).
Consequently, the bits-per-trit ratio is bounded below by \( \frac{p}{k} > \log_2(3) \). When \( p = \lceil k \log_2(3) \rceil \), this is effectively the smallest integer \( p \) that still allows all \( 3^k \) trit patterns to be stored with no collisions. As \( k \to \infty \), the ratio \( \frac{p}{k} \longrightarrow \log_2(3) \), making the scheme asymptotically optimal in terms of bits used per trit.

3.3 CPU Inference with efficient packing

To assess the effectiveness of packings, we implemented both packing methods in ggml.cpp\footnote{https://github.com/ggerganov/llama.cpp}, a framework optimized for running large language models (LLMs) on CPUs. While further optimizations are possible, our primary focus is on reducing a model’s memory footprint and accelerating memory-bound workloads. This is achieved by statically compressing pretrained weights and decompressing them on-the-fly during inference. We begin by showing the efficiency of the CPU implementations for the 1.6-bit packing, referred to as TQ1, and the 2-bit packing, referred to as TQ2.

TQ2: Implementing effective 2 bit for TriLMs.
The quantization process begins with partitioning the input tensor into contiguous, non-overlapping blocks, each containing 256 elements. For each block, a scaling factor (floating-point numbers associated with TriLMs) \( d_i \) is calculated as the maximum absolute value of the elements within the block, i.e., \( d_i = \max(|b_{ij}|) \), where \( b_{ij} \) denotes the \( j \)-th element in the \( i \)-th block. The inverse scaling factor \( \hat{d}_i \) is then defined as \( \hat{d}_i = \frac{1}{d_i} \). Each element \( b_{ij} \) in the block is quantized to a ternary value by multiplying it by the inverse scaling factor and rounding the result: \( q_{ij} = \mathrm{round}(b_{ij} \cdot \hat{d}_i) \). To enable efficient storage, the quantized values are shifted, resulting in \( q_{ij} \in \{0, 1, 2\} \) and packed into 64 bytes per block of 256 elements using base-4 positional encoding. The scaling factor \( d_i \) is stored in 2 bytes (float16), leading to a total storage of 66 bytes per block. The dequantization process begins by reversing the base-4 encoding to recover the four ternary elements (see F.1). The elements are then adjusted back to their signed values by subtracting 1. Finally, the original block is reconstructed by multiplying each quantized value by the corresponding scaling factor: \( \hat{b}_{ij} = d_i \cdot q_{ij} \), where \( \hat{b}_{ij} \) denotes the dequantized approximation of \( b_{ij} \). This quantization scheme achieves significant memory efficiency by compressing 256 floating-point values into just 66 bytes.

TQ1: Implementing effective 1.6 bit for TriLMs.
In our implementation, we encode \( k = 5 \) ternary digits (trits) into \( p = 8 \) bits, achieving an effective bit rate of 1.6 bits per trit. A key challenge arises in efficiently decoding these packed trits for SIMD-optimized operations. Traditional decoding methods rely on division and modulo operations, which are computationally expensive and ill-suited for vectorization. The conventional approach to decoding a packed byte \( b \) involves computing a base-3 integer \( x \) using the formula: \( x = \left\lfloor \frac{b \cdot 3^5 - (3^5-1)+(2^8-1)}{2^8} \right\rfloor \). Each trit \( d_i \in \{-1, 0, 1\} \) is then extracted through the operation: \( d_{i+1} = \left\lfloor \frac{x}{3^{4-i}} \right\rfloor \mod 3 \) for \( i = 0, \ldots, 4 \). This method incurs high computational costs due to the repeated divisions and modulo operations, which hinder SIMD parallelism. To address these inefficiencies, we exploit the near-equivalence \( 3^5 \approx 2^8 \), enabling a multiplication-based scheme that iteratively extracts trits without explicit division or modulo operations. (See Appendix F.2 for the detailed iterative procedure and its SIMD advantages.)

Results. Figure 4 compares token generation speeds (end-to-end and decoding) for models spanning 560M to 3.9B parameters on a Mac M4, highlighting end-to-end latency and output token generation speedup. Specifically, TQ2 outperforms other formats by utilizing 2-bit weight packing, surpassing both 4-bit quantization (as implemented in GGML) and TQ1 (1.6 bits per weight). While TQ1 requires additional fixed-point multiplication operations—resulting in slower inference compared to TQ2—it achieves a significantly smaller memory footprint (shown in Figure 4 on the right), making it advantageous for low-resource environments where memory storage constraints outweigh computational latency. Additional benchmarks are conducted on AMD EPYC 7502 (see Figure 10) and Apple M4 Max (14 CPU cores), with detailed results presented in Table 6, 5 and Appendix F.3. These findings motivate our next step: in the following section, we introduce an optimized 2-bit packing variant designed for high-batch GPU workloads. While our current implementation demonstrates significant speedups, further refinements remain possible to enhance computational efficiency.

4 TriRun: GPU Kernels for High-Batch Settings.

LLM weight quantization leverages the fact that GPUs perform floating-point operations much
Figure 5: Performance evaluation of ternary layers in a transformer block, comparing TriRun with PyTorch FP16 (using CUTLASS), shows near-optimal inference speedup in high-batch settings for larger models. Each subplot corresponds to a specific Nvidia GPU. For additional results, refer to Appendix G.6.

faster than they can fetch data from memory. For example, the NVIDIA L40 has a FLOPs-to-Bytes of approximately 105 (Technologies, 2023). In a typical matrix multiplication for mixed-precision inference in large language models, each input token requires about 2 FLOPs per weight. When weights are 2 bits, each weight occupies 0.25 bytes. During the time needed to load one such weight, the L40 can perform roughly 26 FLOPs. Since each token needs 2 FLOPs per weight, the GPU can support a critical input batch size of about \( \approx 13 \) tokens. Thus, for the L40, if the input batch size is below roughly 13, memory loading becomes the bottleneck for the computation.

4.1 Ternary Kernel Implementation

Optimized Mixed-Precision Multiplication In this work, we introduce an optimized mixed-precision matrix multiplication routine (Frantar et al., 2024) that performs FP16 \( \times \) INT2 computations. In this scheme, an FP16 input matrix is multiplied by a weight matrix stored in a compact 2-bit integer (INT2) format, wherein each 32-bit integer encodes 16 distinct 2-bit values. The central component of this approach is a dequantization function that employs carefully selected bit masks and a lookup-based three-input logical operation to extract the 2-bit fields. This function applies a series of fused arithmetic operations to convert the packed 2-bit data into FP16 values (see appendix §G). As a result, dequantized weight values are produced as fragments containing four FP16 numbers, which can subsequently be scaled using quantization scales stored in a separate buffer.

GPU Performance Optimization. On the GPU, the multiplication kernel is engineered for high performance by using asynchronous memory copy operations alongside specialized tensor core instructions available on modern NVIDIA hardware. Input fragments from the FP16 matrix are asynchronously loaded into shared memory via cp.async\(^3\) instructions, allowing global memory accesses to overlap with computation. The kernel arranges these fragments in memory to minimize bank conflicts and maximize data reuse. Concurrently, the INT2 weight values are fetched using asynchronous copy operations that include cache hints, thereby reducing L2 cache pollution since these weights are used only once during each operation. Once the FP16 fragments and dequantized INT2 weights reside in registers, the kernel employs *tensor core mma* instructions to perform efficient block-wise multiplications. These operations accumulate the results in FP32 registers to maintain higher precision during the reduction phase before converting the final outputs back to FP16 for storage in global memory.

Flexible Implementation and Data Movement. The implementation supports flexible configuration of thread block dimensions, pipeline stages, and grouping parameters for varying problem sizes and hardware. Data is moved from global to shared memory using double-buffering with asynchronous copy fences and explicit barriers. Partial results accumulated across warps or thread blocks are reduced using a hierarchical reduction scheme that first operates within shared memory and then, if necessary, synchronizes globally across thread blocks. Finally, results are reorganized and writ-

\footnotetext{3cp.async in CuPy refers to the support for asynchronous execution of GPU operations}
Figure 6: Comparison of TriRun kernels with the FP16 PyTorch baseline on NVIDIA L40S (for more details see Appendix G.6): (a) Left: Time to first token, (b) Center: Time per output token, (c) Right: Total time across different NVIDIA GPUs.

ten back to global memory in FP16 format. This approach leverages efficient data packing, asynchronous memory operations, and tensor core acceleration, optimizing FP16×INT2 matrix multiplication for ternary large language models.

4.2 Experimental Results

Performance of TriRuns Kernels. In Figure 5 and 12, we evaluate the efficiency of TriRun kernels against PyTorch’s FP16 kernels for the ternary layers within transformer modules. We benchmarked models ranging from 3 billion to 405 billion parameters across different hardwares (see Table 7, and 8 for complete results). Our findings demonstrate that TriRun provides substantial performance improvements. Specifically, on an NVIDIA L40 GPU (optimized for inference) processing large matrices from a 405B parameter model, TriRun achieves a speedup of roughly 7.98x compared to FP16 when using batch sizes between 16 and 32. However, as batch sizes increase beyond this range, the speedup diminishes. This is because the computation becomes increasingly limited by the GPU’s processing capabilities (compute-bound). This pattern of speedup reduction with batch size is observed across all tested GPUs. For more detailed analysis and results, refer to Appendix G.6.

End-to-End Serving Benchmark Figure 6 illustrates the time-to-first-token performance of TriRun with PyTorch, achieving up to a 4.7x speedup on the 70B model with 64 input tokens when running on NVIDIA L40s. Additionally, it shows the time per output token (with 1 input and 64 output tokens), demonstrating a 4.9x improvement in decoding. This trend is particularly evident for larger models, where the 70B model achieves a 4.9x end-to-end generation speedup compared to PyTorch. TriRun uses only one GPU, as opposed to the four GPUs used in the PyTorch FP16 configuration (see Figure 1 on the right for more details). Furthermore, Figure 6 (c) shows that these performance gains are consistent across different NVIDIA hardware, with a more detailed analysis provided in Appendix G.6. Finally, Figure 11 demonstrates that these speedups are particularly pronounced on newer consumer GPUs, as the FLOPs/byte ratio increases.

5 Conclusion and Future Work

In this work, we address the growing memory bottlenecks in large language model inference by studying ternary language models (TriLMs) and proposing strategies for efficient kernel implementation. We conduct a comprehensive scaling law analysis, revealing that TriLMs benefit significantly from scaling training data, achieving competitive performance with floating-point models for a given compute budget despite their extreme quantization. Our experiments with the TriTera family, trained on up to 1.2 trillion tokens, demonstrate sustained performance improvements, emphasizing the potential of ternary models for large-scale training and deployment. To further improve inference efficiency, we introduce novel ternary weight packing schemes and develop optimized kernels. Our GPU kernel, TriRun, achieves up to an 8x speedup over float16 baselines in high-batch inference settings, making ternary models a viable solution for memory-constrained environments. By releasing the TriTera models and optimized inference kernels, we aim to encourage further research on extreme low-bitwidth models and their deployment. Our results demonstrate the scalability and efficiency of ternary models, laying the groundwork for future advancements in efficient LLM research.
Limitations

We study the scaling law for TriLMs where we consider the dependence on number of parameters and training tokens, but do not explicitly account for the number of bits \( b \) used to quantize the model. The various terms that appear in Equation (2) may depend non-linearly on \( b \), which is an interesting direction for future work. Our pre-training scale was constrained by computational resources, and both the parameters and data need to be scaled up significantly to make TriLMs competitive with current state-of-the-art models (AI@Meta, 2024). TriRun implements the 2-effective-bit packing scheme from Section 3.1. A more memory-efficient solution would involve implementing the 1.6-effective-bit packing scheme. However, due to the increased complexity of packing, the unpacking functions would require additional operations in the latter case, making it slower than TriRun. This is left as a direction for future work.

Ethics Statement

The development of TriLMs represents a significant step toward making large-scale language models more efficient by reducing memory consumption and accelerating inference. These advancements enhance accessibility and sustainability in AI research. We advocate for openness in AI, as it drives scientific progress, fosters collaboration, and eliminates the need for re-training, which helps lower environmental impact. However, openness also presents challenges, including concerns related to privacy, security, and fairness. Despite these risks, we believe that transparency enables more effective risk mitigation by inviting diverse scrutiny and safeguards. By releasing the TriTera suite and TriRun kernels, we aim to empower further innovation while ensuring that efficient language models serve a broad spectrum of stakeholders. As open model releases continue to gain momentum, we see this approach as the most effective way to balance progress with responsible AI development.

Acknowledgement

We acknowledge the support from the Mozilla Responsible AI Grant, the Canada CIFAR AI Chair Program, Nolano AI and the Canada Excellence Research Chairs Program. This research was enabled by the computational resources provided by the Summit supercomputer, awarded through the Frontier DD allocation and INCITE 2023 program for the project "Scalable Foundation Models for Transferable Generalist AI" and SummitPlus allocation in 2024. These resources were supplied by the Oak Ridge Leadership Computing Facility at the Oak Ridge National Laboratory, with support from the Office of Science of the U.S. Department of Energy.

References

Advanced Micro Devices, Inc. 2022. AMD CDNA™2 ARCHITECTURE. White paper. Advanced Micro Devices, Inc.

Advanced Micro Devices, Inc. 2025. Amd instinct™ mi250x accelerators. https://www.amd.com/en/products/accelerators/instinct/mi200/mi250x.html. Accessed February 9, 2025.

AI@Meta. 2024. Llama 3 model card.

Saleh Ashkboos, Amirkeivan Mohtashami, Maximilian L. Croci, Bo Li, Pashmina Cameron, Martin Jaggi, Dan Alistarh, Torsten Hoefler, and James Hensman. 2024. Quarot: Outlier-free 4-bit inference in rotated llms. Preprint, arXiv:2404.00456.

Loubna Ben Allal, Anton Lozhkov, Guilherme Penedo, Thomas Wolf, and Leandro von Werra. 2024. Cosmopedia.

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. 2013. Estimating or propagating gradients through stochastic neurons for conditional computation. Preprint, arXiv:1308.3432.

Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O’Brien, Eric Hallahan, Mohammad Aftah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. 2023. Pythia: A suite for analyzing large language models across training and scaling. Preprint, arXiv:2304.01373.

Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. 2019. Piqa: Reasoning about physical commonsense in natural language. In AAAI Conference on Artificial Intelligence.

Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. 2019. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 2924–2936, Minneapolis, Minnesota. Association for Computational Linguistics.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. 2018. Think you have solved question answering? try arc, the ai2 reasoning challenge.
Colin B. Clement, Matthew Birbaum, Kevin P. O’Keeffe, and Alexander A. Alemi. 2019. On the use of arxiv as a dataset. Preprint, arXiv:1905.00075.

Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. 2022a. Llm.int8(): 8-bit matrix multiplication for transformers at scale. Preprint, arXiv:2208.07339.

Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. 2022b. Llm.int8(): 8-bit matrix multiplication for transformers at scale. Preprint, arXiv:2208.07339.

Tim Dettmers and Luke Zettlemoyer. 2023. The case for 4-bit precision: k-bit inference scaling laws. Preprint, arXiv:2212.09720.

Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. 2023. Gptq: Accurate post-training quantization for generative pre-trained transformers. Preprint, arXiv:2210.17323.

Elias Frantar, Roberto L. Castro, Jiale Chen, Torsten Hoefler, and Dan Alistarh. 2024. Marlin: Mixed-precision auto-regressive parallel inference on large language models. Preprint, arXiv:2408.11743.

Amir Gholami, Zhewei Yao, Sehoon Kim, Coleman Hooper, Michael W. Mahoney, and Kurt Keutzer. 2024. Ai and memory wall. Preprint, arXiv:2403.14123.

Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, Shane Arora, David Atkinson, Russell Author, Khyathi Raghavi Chandu, Arman Cohan, Jennifer Dumas, Yanai Elazar, Yuling Gu, Jack Hessel, Tushar Khot, William Merrill, Jacob Morrison, Niklas Muenninghoff, Aakanksha Naik, Crystal Nam, Matthew E. Peters, Valentina Pyatkin, Abhilasha Ravichander, Dustin Schwenk, Saurabh Shah, Will Smith, Emma Strubell, Nishant Subramani, Mitchell Wortsman, Pradeep Dasigi, Nathan Lambert, Kyle Richardson, Luke Zettlemoyer, Jesse Dodge, Kyle Lo, Luca Soldaini, Noah A. Smith, and Hannaneh Hajishirzi. 2024. Olmo: Accelerating the science of language models. Preprint, arXiv:2402.00838.

Pujing He, Shan Zhou, Wenhuan Huang, Changqing Li, Duyi Wang, Bin Guo, Chen Meng, Sheng Gui, Weifei Yu, and Yi Xie. 2024. Inference performance optimization for large language models on cpus. Preprint, arXiv:2407.07304.

Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2021. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR).

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. 2022. Training compute-optimal large language models. Preprint, arXiv:2203.15556.

Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lampe, Lélio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. 2024. Mixtral of experts. Preprint, arXiv:2401.04088.

Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. 2017. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1601–1611, Vancouver, Canada. Association for Computational Linguistics.

Dhiraj Kalamkar, Dheevatsa Mudigere, Naveen Mellempudi, Dipankar Das, Kunal Banerjee, Sasikanth Avancha, Dharna Teja Vooturi, Nataraj Janamalamadaka, Jianyu Huang, Hector Yuen, Jiyan Yang, Jongsoo Park, Alexander Heinicke, Evangelos Georganas, Sudarshan Srinivasan, Abhishek Kundi, Mishu Smelyanskiy, Bharat Kaul, and Pradeep Dubey. 2019. A study of bfloat16 for deep learning training. Preprint, arXiv:1905.12322.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. Preprint, arXiv:2001.08361.

Ayush Kaushal, Tejas Vaidhya, Arnab Kumar Mondal, Tejas Pandey, Aaryan Bhagat, and Irina Rish. 2024. Spectra: Surprising effectiveness of pretraining ternary language models at scale. Preprint, arXiv:2407.12327.

Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Wei-Ming Chen, Wei-Chen Wang, Guangxuan Xiao, Xingyu Dang, Chuang Gan, and Song Han. 2024. Awq: Activation-aware weight quantization for llm compression and acceleration. Preprint, arXiv:2306.00978.

Jian Liu, Leyang Cui, Hanneng Liu, Dandan Huang, Yile Wang, and Yue Zhang. 2021. Logiqa: a challenge dataset for machine reading comprehension with logical reasoning. In Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI’20.

Kyle Lo, Lucy Lu Wang, Mark Neumann, Rodney Kinney, and Daniel Weld. 2020. S2ORC: The semantic scholar open research corpus. In Proceedings of the
58th Annual Meeting of the Association for Computational Linguistics, pages 4969–4983, Online. Association for Computational Linguistics.

Ilya Loshchilov and Frank Hutter. 2019. Decoupled weight decay regularization. Preprint, arXiv:1711.05101.

Anton Lozhkov, Loubna Ben Allal, Leandro von Werra, and Thomas Wolf. 2024. Fineweb-edu: the finest collection of educational content.

Shuming Ma, Hongyu Wang, Lingxiao Ma, Lei Wang, Wenhui Wang, Shaohan Huang, Li Dong, Ruiping Wang, Jilong Xue, and Furu Wei. 2024. The era of 1-bit llms: All large language models are in 1.58 bits. Preprint, arXiv:2402.17764.

Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. 2018. Mixed precision training. Preprint, arXiv:1710.03740.

Denis Paperno, Germán Kruszewski, Angeliki Lazari-dou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. 2016. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1525–1534. Berlin, Germany. Association for Computational Linguistics.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners.

Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. Zero: Memory optimizations toward training trillion parameter models. Preprint, arXiv:1910.02054.

Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. 2021. Winogrande: an adversarial winograd schema challenge at scale. Commun. ACM, 64(9):99–106.

Nikhil Sardana, Jacob Portes, Sasha Doubov, and Jonathan Frankle. 2024. Beyond chinchilla-optimal: Accounting for inference in language model scaling laws. Preprint, arXiv:2401.00448.

Noam Shazeer. 2020. Glu variants improve transformer. Preprint, arXiv:2002.05202.

Zhiqiang Shen, Tianhua Tao, Liqun Ma, Willie Neiswanger, Zhengzhong Liu, Hongyi Wang, Bowen Tan, Joel Hestness, Natalia Vassileva, Daria Soboleva, and Eric Xing. 2024. Slimpajama-dc: Understanding data combinations for llm training. Preprint, arXiv:2309.10818.

Ying Sheng, Lianmin Zheng, Binhang Yuan, Zuo-han Li, Max Ryabinin, Daniel Y. Fu, Zhiqiang Xie, Beidi Chen, Clark Barrett, Joseph E. Gonzalez, Percy Liang, Christopher Ré, Ion Stoica, and Ce Zhang. 2023. Flexgen: High-throughput generative inference of large language models with a single gpu. Preprint, arXiv:2303.06865.

Avi Singh, John D. Co-Reyes, Rishabh Agarwal, Ankesh Anand, Piyush Patil, Xavier Garcia, Peter J. Liu, James Harrison, Jaehoon Lee, Kelvin Xu, Aaron Parisi, Abhishek Kumar, Alex Alemi, Alex Rizkowsky, Azade Nova, Ben Adlam, Bernd Bohnet, Gamaeldin Elsayed, Hanie Sedghi, Igor Mordatch, Isabelle Simpson, Izzeddin Gur, Jasper Snoek, Jeffrey Pennington, Jiri Hron, Kathleen Keane, Kevin Swersky, Kshitij Mahajan, Laura Culp, Lechao Xiao, Maxwell L. Bileschi, Noah Constant, Roman Novak, Rosanne Liu, Tris Warkentin, Yundi Qian, Yamini Bansal, Ethan Dyer, Behnam Neyshabur, Jascha Sohl-Dickstein, and Noah Fiedel. 2024. Beyond human data: Scaling self-training for problem-solving with language models. Preprint, arXiv:2312.06585.

Luca Soldaini and Kyle Lo. 2023. peS2o (Pretraining Efficiently on S2ORC) Dataset. Technical report, Allen Institute for AI. ODC-By, https://github.com/allenai/pes2o.

Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. 2023. Roformer: Enhanced transformer with rotary position embedding. Preprint, arXiv:2104.09864.

PNY Technologies. 2023. Nvidia l40s datasheet. Accessed: 2025-02-12.

Yury Tokpanov, Beren Millidge, Paolo Glorioso, Jonathan Pilault, Adam Ibrahim, James Whittington, and Quentin Anthony. 2024. Zyda: A 1.3t dataset for open language modeling. Preprint, arXiv:2406.01981.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023. Llama: Open and efficient foundation language models. Preprint, arXiv:2302.13971.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2023. Attention is all you need. Preprint, arXiv:1706.03762.

Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Huaijie Wang, Lingxiao Ma, Fan Yang, Ruiping Wang, Yi Wu, and Furu Wei. 2023. Bitnet: Scaling 1-bit transformers for large language models. Preprint, arXiv:2310.11453.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed Chi, Quoc Le, and Denny Zhou. 2023. Chain-of-thought prompting elicits reasoning in large language models. Preprint, arXiv:2201.11903.
Johannes Welbl, Nelson F. Liu, and Matt Gardner. 2017. Crowdsourcing multiple choice science questions. ArXiv, abs/1707.06209.

Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. 2024. Smoothquant: Accurate and efficient post-training quantization for large language models. Preprint, arXiv:2211.10438.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019. HellaSwag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791–4800, Florence, Italy. Association for Computational Linguistics.

Zhuosheng Zhang, Aston Zhang, Mu Li, and Alex Smola. 2022. Automatic chain of thought prompting in large language models. Preprint, arXiv:2210.03493.

Zixuan Zhou, Xuefei Ning, Ke Hong, Tianyu Fu, Jiaming Xu, Shiyao Li, Yuning Lou, Luning Wang, Zhihang Yuan, Xiuhong Li, Shengen Yan, Guohao Dai, Xiao-Ping Zhang, Yuhan Dong, and Yu Wang. 2024. A survey on efficient inference for large language models. Preprint, arXiv:2404.14294.

A Related Work

Training LLMs in low precision Large language models such as GPT (Radford et al., 2019), OLMo (Groeneveld et al., 2024), and the LLaMA family (Touvron et al., 2023) have traditionally relied on mixed precision (FP32/FP16 or FP32/BF16) (Micikevicius et al., 2018) and half-precision (FP16/BF16) (Kalamkar et al., 2019) to optimize computational efficiency. More recent advancements in extreme quantization have introduced ternary and binary network paradigms (Kaushal et al., 2024; Wang et al., 2023), which leverage quantization-aware training (QAT) for efficient low-bitwidth model representations. These models maintain higher-precision latent (or master) weights, such as FP16, to stabilize training while dynamically binarizing or ternarizing weights during inference. The straight-through estimator (STE) (Bengio et al., 2013) is commonly employed to facilitate gradient-based updates. The Spectra suite (Kaushal et al., 2024) provides a comprehensive study of ternary, quantized, and FP16 language models, offering insights into the performance and scaling trends of low-bitwidth models.

Advancements in Post-Training Quantization Post-training quantization (PTQ) remains a crucial approach for reducing LLM memory and compute requirements without requiring retraining. Techniques such as SmoothQuant (Xiao et al., 2024) and QuaRot (Ashkboos et al., 2024) address challenges associated with activation quantization, particularly mitigating large activation outliers (Dettmers et al., 2022b). While these methods improve compression, they often rely on 8-bit quantization to preserve numerical stability. Continued research into activation-aware quantization techniques is vital for further enhancing LLM deployment in resource-constrained environments.

Optimizing Inference Efficiency To improve LLM deployment efficiency, frameworks like MARLIN (Frantar et al., 2024) initially implemented GPTQ-based quantization, enabling accelerated inference. MARLIN kernels combine various techniques, ranging from advanced task scheduling, partitioning, and pipelining techniques to quantization-specific layout and compute optimizations. More recently, MARLIN has been extended to incorporate Activation-Weight Quantization (AWQ) (Lin et al., 2024), a technique that jointly quantizes both weights and activations to mitigate accuracy degradation in low-bitwidth settings.

B Pretraining Details

B.1 Quantized Linear Layer: Forward, Backward, and Inference Stages

We now present the mathematical formulation for a linear layer employing the TriLM quantization scheme (Kaushal et al., 2024), outlining the processes for the forward pass, backward pass, and inference stages.

Forward Pass. In the forward pass, we begin by calculating the scaling factor \( \gamma \) to normalize the weight matrix \( W \). The scaling factor is given by:

\[
\gamma = \epsilon + \frac{1}{nm} \sum_{i=1}^n \sum_{j=1}^m |W_{ij}|
\]

where \( n \) and \( m \) denote the dimensions of the weight matrix \( W \), and \( \epsilon \) is a small constant added for numerical stability.

Subsequently, the weight matrix \( W \) is quantized by rounding its entries to the nearest value in the set \( \{-1, 0, 1\} \), scaled by \( \gamma \):

\[
\widetilde{W}_{ij} = \operatorname{round}\left( \min \left( \max \left( \frac{W_{ij}}{\gamma}, -1 \right), 1 \right) \right)
\]

The quantized weight matrix \( \widetilde{W} \) is then obtained by scaling the rounded weights: \( \widetilde{W}_{ij} = \gamma \widetilde{W}_{ij} \)
<table>
  <tr>
    <th>Dataset Name</th>
    <th>Number of Tokens (Billion)</th>
    <th>Percentage</th>
  </tr>
  <tr>
    <td>ArXiv (Clement et al., 2019)</td>
    <td>3.67</td>
    <td>0.31%</td>
  </tr>
  <tr>
    <td>Cosmopedia-v2 (Ben Allal et al., 2024)</td>
    <td>22.36</td>
    <td>1.86%</td>
  </tr>
  <tr>
    <td>PeS2o (Soldaini and Lo, 2023)</td>
    <td>42.70</td>
    <td>3.56%</td>
  </tr>
  <tr>
    <td>FineWeb-Edu (Lozhkov et al., 2024)</td>
    <td>960.42</td>
    <td>80.04%</td>
  </tr>
  <tr>
    <td>Zyda - StarCoder (Tokpanov et al., 2024)</td>
    <td>170.85</td>
    <td>14.24%</td>
  </tr>
  <tr>
    <th>Total</th>
    <th>1200.00</th>
    <th>100.00%</th>
  </tr>
</table>

Table 1: Pretraining datasets and token counts for TriTera models.

Finally, the output \( Y \) is computed as the product of the input \( X \) and the transposed quantized weight matrix: \( Y = X \widehat{W}^T \)

Backward Pass. During the backward pass, the gradients of the loss function \( L \) with respect to the input \( X \) and the weight matrix \( W \) are computed. These gradients are given by:

\[
\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} \widehat{W}
\]

\[
\frac{\partial L}{\partial W} = \frac{\partial L}{\partial Y} X^T
\]

Inference. For inference, the quantized weight matrix \( \widehat{W} \) and the scaling factor \( \gamma \) are precomputed and cached to reduce computation during prediction. The steps are as follows:

1. Compute \( \widehat{W} \) and \( \gamma \) once and store them.
2. Use the precomputed values to calculate the quantized weight matrix: \( \widehat{W}_{ij} = \gamma \widehat{W}_{ij} \)
3. Finally, the output \( Y \) is computed as: \( Y = X \widehat{W}^T \)

By caching the scaling factor and the quantized weights, the inference process is significantly accelerated, as it eliminates the need for redundant recalculations.

B.2 Dataset

Our training corpus comprises a diverse mix of data from publicly available sources. To scale TriLMs, we trained on approximately 1.2 trillion tokens, up-sampling the most factual sources to enhance the model’s knowledge while reducing hallucinations. The details of the datasets used are summarized in Table 1. Each dataset was preprocessed and tokenized using llama2 tokenizer (AI@Meta, 2024).

• ArXiv (Clement et al., 2019): The dataset comprises 1.5 million arXiv preprint articles from fields such as Physics, Mathematics, and Computer Science, encompassing text, figures, authors, citations, and metadata.

• Cosmopedia-v2 (Ben Allal et al., 2024): A synthetic dataset of over 30 million documents and 25 billion tokens. The dataset was generated using the Mixtral-8x7B-Instruct-v0.1 model, a multi-expert language model introduced in (Jiang et al., 2024), designed for high-quality content generation. It is one of the largest publicly available synthetic datasets.

• PeS2o (Soldaini and Lo, 2023): It comprises 40 million open-access academic papers that have been cleaned, filtered, and formatted specifically for the pre-training of language models. It is derived from the Semantic Scholar Open Research Corpus (Lo et al., 2020).

• Zyda-StarCoder Git-Commits (Tokpanov et al., 2024): For our models, we exclusively utilize the GitHub-Issues and Jupyter-Structured subsets of the Zyda-Starcoder dataset.

• Zyda-StarCoder-Languages: A dataset encompassing multiple programming languages, enabling the model to perform well across diverse coding tasks.

• FineWeb-Edu (Lozhkov et al., 2024): A subset of high quality dataset consists of 1.3T tokens of educational web pages filtered from FineWeb dataset.

B.3 Hyperparameter Choices

We adopt a single learning rate with a warmup followed by a cosine decay schedule, replacing the dual learning rate approach used in TriLMs (Kaushal et al., 2024). Additionally, we eliminate the use of weight decay, consistent with the modifications.
<table>
  <tr>
    <th>Feature</th>
    <th>TriTera</th>
  </tr>
  <tr>
    <td>Biases</td>
    <td>None</td>
  </tr>
  <tr>
    <td>Activation</td>
    <td>SwiGLU</td>
  </tr>
  <tr>
    <td>RoPE (\( \theta \))</td>
    <td>\( 5 \cdot 10^5 \)</td>
  </tr>
  <tr>
    <td>QKV Normalization</td>
    <td>QK-Norm</td>
  </tr>
  <tr>
    <td>Layer Norm</td>
    <td>RMSNorm</td>
  </tr>
  <tr>
    <td>Layer Norm Applied to</td>
    <td>Outputs</td>
  </tr>
  <tr>
    <td>Z-Loss Weight</td>
    <td>\( 10^{-5} \)</td>
  </tr>
  <tr>
    <td>Weight Decay on Embeddings</td>
    <td>No</td>
  </tr>
</table>

Table 2: Configuration Details for TriTera

B.4 Hyperparameters.

All the models are randomly initialized from a truncated normal distribution with a mean of 0 and a standard deviation of 0.02. We trained using the AdamW optimizer (Loshchilov and Hutter, 2019), with \( \beta_1 = 0.9, \beta_2 = 0.95, \) and \( \epsilon = 10^{-5} \). The weight decay was applied with a value of 0.1. A cosine learning rate schedule was employed, with a warmup of 2000 steps, followed by a decay of the final learning rate to 10% of the peak learning rate. We used gradient clipping with a threshold of 1.0. Metrics were logged every 10 steps. For simplicity during training, we adopt a single learning rate with a warmup followed by a cosine decay schedule, replacing the dual learning rate approach used in Spectra. Additionally, we eliminate the use of weight decay, consistent with the modifications. Table 3 summarizes the hyperparameters for our largest models.

B.5 Hardware and Training Setups

Each node in the Frontier cluster includes four AMD MI250X accelerators, with each accelerator featuring two GCDs that function as independent GPUs. The total bidirectional communication bandwidth within a node ranges between 100 GB/s and 400 GB/s. The nodes are connected via Ethernet-based HPE Slingshot interconnects. Each node is equipped with four links, each providing a total directional bandwidth of 50 GB/s. Our approach scales near-linearly up to 2048 GPUs, as shown in Figure 7

C Scaling Laws

C.1 Scaling Laws of TriLMs and FloatLMs

In Section 2.2, we derived the scaling law for TriLMs as a function of the number of parameters (\( N \)) and the number of training tokens used (\( D \)) by assuming the parametric form defined in Kaplan et al. (2020); Hoffmann et al. (2022). We apply the same procedure to derive the scaling law for FloatLMs which use 16-bit precision to facilitate direct comparison and understand the effect of compute on performance.

In addition to the ternary LLMs described in Section 2.2, we train corresponding 16-bit models which we refer to as FloatLMs across parameters sizes \( \in [990M, 1900M, 3900M, 5600M, 11000M] \) (excluding embeddings) and dataset sizes \( \in [20B, 40B, 75B, 150B] \) tokens. We follow the same procedure as for TriLMs and obtain the following power law for FloatLMs,

\[
\hat{L}(N, D) \approx 2.17 + \frac{7.86}{N^{0.56}} + \frac{3.42}{D^{0.53}}.
\]

Comparing this with Equation (2), we make two interesting observations. First, the constant term and the coefficients are markedly different for ternary and float LMs, indicating that these terms might be dependent on the level of quantization. Second, the terms involving \( N \) and \( D \) have almost the same exponents for FloatLMs, which means that increasing either parameters and training tokens has a similar effect on improving LLM performance. This is in contrast to TriLMs, where the term involving training tokens decays much more rapidly than term involving number of parameters.

Figure 8 shows the final validation loss for different FloatLM models against the number of parameters and the number of training tokens, along with the scaling law fit.

C.2 Parametric Fit for Scaling Law

We obtain the coefficients for the parametric scaling law in Equation (1) by finding the least squares

![Number of GPUs vs. Relative Speedup.](page_1012_186_377_273.png)

Figure 7: Number of GPUs vs. Relative Speedup.
<table>
  <tr>
    <th>Parameter</th>
    <th>TriTera-1B</th>
    <th>TriTera-2B</th>
    <th>TriTera-3B</th>
  </tr>
  <tr>
    <td>Number of Parameters</td>
    <td>1.526B</td>
    <td>2.5547B</td>
    <td>3.6680B</td>
  </tr>
  <tr>
    <td><b>Hidden Size</b></td>
    <td>2048</td>
    <td>2560</td>
    <td>3072</td>
  </tr>
  <tr>
    <td><b>Number of Layers</b></td>
    <td>24</td>
    <td>26</td>
    <td>28</td>
  </tr>
  <tr>
    <td>Attention Heads</td>
    <td>16</td>
    <td>20</td>
    <td>24</td>
  </tr>
  <tr>
    <td>MLP Hidden Size</td>
    <td>8192</td>
    <td>10240</td>
    <td>11264</td>
  </tr>
  <tr>
    <td>Number of KV Heads</td>
    <td>4</td>
    <td>5</td>
    <td>6</td>
  </tr>
  <tr>
    <td>Embedding Size</td>
    <td>32768</td>
    <td>32768</td>
    <td>32768</td>
  </tr>
  <tr>
    <td><b>Max Sequence Length</b></td>
    <td>2048</td>
    <td>2048</td>
    <td>2048</td>
  </tr>
  <tr>
    <td>Activation Function</td>
    <td>SiLU</td>
    <td>SiLU</td>
    <td>SiLU</td>
  </tr>
  <tr>
    <td><b>Optimizer</b></td>
    <td>AdamW</td>
    <td>AdamW</td>
    <td>AdamW</td>
  </tr>
  <tr>
    <td><b>Learning Rate</b></td>
    <td>0.0015</td>
    <td>0.0015</td>
    <td>0.0015</td>
  </tr>
  <tr>
    <td><b>Weight Decay</b></td>
    <td>0.1</td>
    <td>0.1</td>
    <td>0.1</td>
  </tr>
  <tr>
    <td><b>Gradient Clipping</b></td>
    <td>1.0</td>
    <td>1.0</td>
    <td>1.0</td>
  </tr>
</table>

Table 3: Architecture summary for TriTera 1B, 2B, and 3B models based on revised configurations.

fit on the the final validation losses of the suite of models trained across different parameter and training token values.

To evaluate our fit, we calculate the coefficient of determination, or \( R^2 \), which is a statistical measure that indicates how well a model fits a set of data, with \( R^2 = 1.0 \) indicating a perfect fit. Our fitted power laws have \( R^2 = 0.9921 \) for TriLMs and \( R^2 = 0.9958 \) for FloatLMs. Figure 9 plots the predicted validation loss following our derived scaling law versus the actual empirical values.

D Benchmark Details

We benchmark TriLM across knowledge, commonsense, and reasoning benchmarks. We average our scores across three different 'seeds'.

D.1 Commonsense and Reasoning

We report commonsense and reasoning benchmark scores across 6 benchmarks in Table 4. Each is considered in a zero-shot setting. Following are the details of each of the benchmarks considered:

• ARC Challenge and Easy: (Clark et al., 2018) The ARC dataset consists of 7,787 multiple-choice science questions, split into two categories: Challenge and Easy. We compute both the accuracy and normalized accuracy for these two sets.

• BoolQ: (Clark et al., 2019) BoolQ is a reading comprehension dataset featuring naturally occurring yes/no questions. We evaluate the model’s performance by calculating its accuracy on this task.

• HellaSwag: (Zellers et al., 2019) HellaSwag is a dataset for testing grounded commonsense through multiple-choice questions. Incorrect answer choices are generated using Adversarial Filtering (AF), designed to deceive machines but not humans. Accuracy and normalized accuracy are reported for this dataset.

• WinoGrande: (Sakaguchi et al., 2021) WinoGrande is a dataset of 44,000 questions designed to assess commonsense reasoning via a fill-in-the-blank task with binary options. We report the model’s accuracy on this dataset.

• PIQA: (Bisk et al., 2019) The Physical Interaction Question Answering (PIQA) dataset evaluates physical commonsense reasoning. We compute accuracy and normalized accuracy for this task.

• LAMBADA OpenAI: (Paperno et al., 2016) LAMBADA is a dataset used to test text understanding through next-word prediction, containing narrative passages from BooksCorpus.
Figure 8: Effect of scaling number of parameters (left) and number of training tokens (right) on final validation loss for FloatLMs. The dotted lines show the power law derived in Equation (3).

Figure 9: Predicted versus actual values of the final validation loss based on the parametric fit of the scaling law for TriLMs (left) and FloatLMs (right).

To perform well on LAMBADA, models must leverage broad discourse information rather than just local context. We report both perplexity and accuracy for this dataset.

• LogiQA: (Liu et al., 2021) LogiQA focuses on testing human-like logical reasoning across multiple types of deductive reasoning tasks. We measure both accuracy and normalized accuracy for this dataset.

D.2 Knowledge

We report performance on SciQ, TriviaQA in Tables 4. Each is considered in a zero-shot setting. Following are the details of each of the benchmarks considered:

The knowledge-based evaluation included the following tasks:

• SciQ: (Welbl et al., 2017) The SciQ dataset contains multiple-choice questions with 4 answer options from crowd-sourced science exams. The questions range from Physics, Chemistry and Biology and several other fields. We calculate the accuracy and length normalized accuracy on this task.

• TriviaQA: (Joshi et al., 2017) TriviaQA is a reading comprehension dataset containing question-answer-evidence triples. We calculate the exact match accuracy on this task.

• MMLU (Hendrycks et al., 2021): The benchmark aims to assess the knowledge gained during pretraining by evaluating models solely in zero-shot and few-shot scenarios. It spans 57 subjects, including STEM fields, humanities, social sciences, and more.

D.3 Serving benchmark for inference

We report the following serving benchmark for our TriRun kernels.
<table>
  <tr>
    <th>Dataset</th>
    <th>Metric</th>
    <th>TriTera 1B</th>
    <th>TriTera 2B</th>
    <th>TriTera 3B</th>
    <th>Llama-1 7B</th>
  </tr>
  <tr>
    <td rowspan="2">Arc Challenge</td>
    <td>acc</td>
    <td>33.45±1.38</td>
    <td>37.29±1.41</td>
    <td>40.61±1.44</td>
    <td>41.81±1.44</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>36.43±1.41</td>
    <td>39.69±1.43</td>
    <td>42.58±1.44</td>
    <td>44.80±1.45</td>
  </tr>
  <tr>
    <td rowspan="2">Arc Easy</td>
    <td>acc</td>
    <td>69.82±0.94</td>
    <td>72.60±0.92</td>
    <td>75.97±0.88</td>
    <td>75.25±0.89</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>62.54±0.99</td>
    <td>67.42±0.96</td>
    <td>71.93±0.92</td>
    <td>72.81±0.91</td>
  </tr>
  <tr>
    <td rowspan="2">BoolQ</td>
    <td>acc</td>
    <td>62.57±0.85</td>
    <td>56.70±0.87</td>
    <td>66.15±0.83</td>
    <td>75.11±0.76</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>43.20±0.49</td>
    <td>46.44±0.50</td>
    <td>49.65±0.50</td>
    <td>56.95±0.49</td>
  </tr>
  <tr>
    <td rowspan="2">HellaSwag</td>
    <td>acc</td>
    <td>56.61±0.49</td>
    <td>61.37±0.49</td>
    <td>66.28±0.47</td>
    <td>76.21±0.42</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>47.31±0.70</td>
    <td>48.85±0.70</td>
    <td>54.22±0.89</td>
    <td>73.53±0.61</td>
  </tr>
  <tr>
    <td rowspan="2">LAMBADA (OpenAI)</td>
    <td>acc</td>
    <td>34.81±0.66</td>
    <td>38.58±0.68</td>
    <td>47.04±0.70</td>
    <td>67.82±0.65</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>22.12±1.63</td>
    <td>22.27±1.63</td>
    <td>22.00±1.66</td>
    <td>22.73±1.64</td>
  </tr>
  <tr>
    <td rowspan="2">LogiQA</td>
    <td>acc</td>
    <td>27.04±1.75</td>
    <td>29.65±1.79</td>
    <td>30.57±1.81</td>
    <td>30.11±1.80</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>28.60±2.02</td>
    <td>30.00±2.05</td>
    <td>32.20±2.09</td>
    <td>34.20±2.12</td>
  </tr>
  <tr>
    <td rowspan="2">OpenBookQA</td>
    <td>acc</td>
    <td>38.80±2.18</td>
    <td>41.00±2.20</td>
    <td>41.80±2.21</td>
    <td>44.40±2.22</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>71.98±1.05</td>
    <td>73.67±1.03</td>
    <td>76.01±1.00</td>
    <td>78.67±0.96</td>
  </tr>
  <tr>
    <td rowspan="2">PIQA</td>
    <td>acc</td>
    <td>72.47±1.04</td>
    <td>75.41±1.00</td>
    <td>76.33±0.99</td>
    <td>79.16±0.95</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>58.09±1.39</td>
    <td>58.56±1.38</td>
    <td>62.43±1.36</td>
    <td>69.93±1.29</td>
  </tr>
  <tr>
    <td rowspan="2">WinoGrande</td>
    <td>acc</td>
    <td>89.60±0.97</td>
    <td>90.80±0.91</td>
    <td>92.80±0.82</td>
    <td>94.60±0.72</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>84.10±1.16</td>
    <td>87.00±1.06</td>
    <td>88.40±1.01</td>
    <td>93.00±0.81</td>
  </tr>
  <tr>
    <td rowspan="2">SciQ</td>
    <td>acc</td>
    <td>29.16±0.65</td>
    <td>30.33±0.66</td>
    <td>30.90±0.65</td>
    <td>33.28±0.67</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>38.46±0.86</td>
    <td>40.42±0.86</td>
    <td>49.39±0.87</td>
    <td>46.31±0.86</td>
  </tr>
  <tr>
    <td rowspan="2">MMLU (cont.): Humanities</td>
    <td>acc</td>
    <td>35.81±0.86</td>
    <td>38.97±0.87</td>
    <td>40.92±0.87</td>
    <td>42.44±0.88</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>27.62±0.79</td>
    <td>30.23±0.80</td>
    <td>32.06±0.82</td>
    <td>33.43±0.83</td>
  </tr>
  <tr>
    <td rowspan="2">MMLU (cont.): STEM</td>
    <td>acc</td>
    <td>32.34±0.39</td>
    <td>34.43±0.39</td>
    <td>36.12±0.39</td>
    <td>38.21±0.40</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>2.05±0.39</td>
    <td>2.12±0.40</td>
    <td>3.03±0.47</td>
    <td>9.70±0.82</td>
  </tr>
  <tr>
    <td rowspan="2">GSM8K</td>
    <td>exact_match</td>
    <td>23.22±0.77</td>
    <td>24.22±0.78</td>
    <td>24.69±0.79</td>
    <td>27.07±0.81</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>23.12±0.77</td>
    <td>24.52±0.79</td>
    <td>24.63±0.79</td>
    <td>26.50±0.81</td>
  </tr>
  <tr>
    <td rowspan="2">MathQA</td>
    <td>acc</td>
    <td>23.22±0.77</td>
    <td>24.22±0.78</td>
    <td>24.69±0.79</td>
    <td>27.07±0.81</td>
  </tr>
  <tr>
    <td>acc_norm</td>
    <td>23.12±0.77</td>
    <td>24.52±0.79</td>
    <td>24.63±0.79</td>
    <td>26.50±0.81</td>
  </tr>
</table>

Table 4: Model performance across various datasets.

• Time to First Token. The time taken from the start of the inference process until the model generates its first token. This metric is used to measure the latency before the model begins producing outputs.

• Time per Output Token. The average time taken by the model to generate each subsequent token after the first. This metric reflects the efficiency of the model in producing tokens once the inference process has started.

• Total Tokens per Second. The overall rate at which the model generates tokens, including both the initial and subsequent tokens. This metric accounts for the entire sequence generation process and provides an aggregate measure of inference speed.

• Output Tokens per Second. The rate at which the model generates tokens after the first token has been produced. This metric focuses on sustained generation speed, reflecting the model’s efficiency once the decoding process has started.

E Formal Proofs

E.1 Notations and Theorem

Theorem 1 (Correctness). Let \( D = (d_1, d_2, \ldots, d_n) \) be a sequence of ternary digits \( d_i \in \{-1, 0, 1\} \). When \( D \) is partitioned into blocks of size \( k \) and each block is encoded into a \( p \)-bit integer, the encoding and decoding operations \( P \) and \( U \) are lossless if and only if \( 2^p > 3^k \); that is, \( U(P(D)) = D \).

Let
\[
D = \{d_1, d_2, \ldots, d_n\}, \qquad d_i \in \{-1, 0, 1\},
\]
be a sequence of balanced ternary digits. We partition \( D \) into blocks of \( k \) digits (with the last block possibly shorter). For a given block, define the shifted digits by
\[
d'_j = d_j + 1, \quad j = 0, 1, \ldots, k - 1,
\]
so that \( d_j' \in \{0, 1, 2\} \). Then define the integer
\[
N = \sum_{j=0}^{k-1} d_j' \cdot 3^{k-1-j}.
\]
Since each \( d_j' \) is in \( \{0, 1, 2\} \), we have
\[
0 \leq N \leq 3^k - 1.
\]
Assume we choose an integer \( p \) such that
\[
2^p > 3^k.
\]
The packing function is defined by
\[
b = \left\lfloor \frac{N \cdot 2^p + (3^k - 1)}{3^k} \right\rfloor.
\]
This mapping is one-to-one on the set \( \{0, 1, \ldots, 3^k - 1\} \) and yields an integer \( b \) in the range \( [0, 2^p - 1] \).

The unpacking function recovers a number \( x \) via
\[
x = \left\lfloor \frac{b \cdot 3^k - (3^k - 1) + (2^p - 1)}{2^p} \right\rfloor.
\]
The recovery of the shifted digits is given by:
\[
d_j' = \left( \left\lfloor \frac{x}{3^{k-1-j}} \right\rfloor \right) \mod 3, \quad j = 0, 1, \ldots, k-1.
\]

E.2 Proof of Theorem 1 (Correctness).

Step 1. Necessity of the Condition. Notice that the mapping \( P \) takes an input \( N \in \{0, 1, \ldots, 3^k - 1\} \) (a total of \( 3^k \) values) and produces an output \( b \in \{0, 1, \ldots, 2^p - 1\} \) (a total of \( 2^p \) values). If
\[
2^p \leq 3^k,
\]
then by the pigeonhole principle the mapping \( P \) cannot be injective, and therefore lossless recovery is impossible. Thus, a necessary condition for \( U(P(D)) = D \) is that
\[
2^p > 3^k.
\]
We now show that \( U(P(D)) = D \) if and only if \( 2^p > 3^k \). We start by showing that \( x = N \), and then we recover the original digits.

Step 2. Expressing the Packing Equation via the Division Algorithm. By the division algorithm, there exists a unique remainder integer \( r \) with \( 0 \leq r \leq 3^k - 1 \) such that
\[
N \cdot 2^p + (3^k - 1) = b \cdot 3^k + r.
\]
Rearranging, we obtain
\[
N \cdot 2^p = b \cdot 3^k - (3^k - 1) + r.
\]
Dividing both sides by \( 2^p \) yields
\[
N = \frac{b \cdot 3^k - (3^k - 1)}{2^p} + \frac{r}{2^p}.
\]
Because \( 0 \leq r \leq 3^k - 1 \), the term
\[
\frac{r}{2^p}
\]
satisfies
\[
0 \leq \frac{r}{2^p} < \frac{3^k}{2^p}.
\]
Thus, the requirement \( 2^p > 3^k \) is equivalent to having
\[
0 \leq \frac{r}{2^p} < 1.
\]
If \( 2^p \leq 3^k \) the fractional part might reach or exceed 1, and the mapping would fail to be one-to-one. Hence, the lossless property holds if and only if \( 2^p > 3^k \).

Step 3. Recovery of \( N \) via the Decoding Operation. Examine the decoding formula:
\[
x = \left\lfloor \frac{b \cdot 3^k - (3^k - 1) + (2^p - 1)}{2^p} \right\rfloor.
\]
We rewrite the expression inside the floor as
\[
\begin{align*}
&\frac{b \cdot 3^k - (3^k - 1) + (2^p - 1)}{2^p} \\
&= \frac{b \cdot 3^k - (3^k - 1)}{2^p} + \frac{2^p - 1}{2^p} \\
&= N - \frac{r}{2^p} + \frac{2^p - 1}{2^p} \\
&= N + \frac{(2^p - 1) - r}{2^p}.
\end{align*}
\]
Since \( 0 \leq r \leq 3^k - 1 \) and \( 3^k < 2^p \), the correction term
\[
\frac{(2^p - 1) - r}{2^p}
\]
satisfies
\[
0 \leq \frac{(2^p - 1) - r}{2^p} < 1.
\]
Thus,
\[
N \leq N + \frac{(2^p - 1) - r}{2^p} < N + 1.
\]
Taking the floor gives
\[
x = N.
\]
Step 4. Recovery of the Original Ternary Digits
Since \( N \) represents the base-3 number with shifted digits \( d_j' \in \{0, 1, 2\} \), we recover each \( d_j' \) by writing \( N \) in base 3. It is important to note that if \( N \) has a “short” base-3 representation (i.e., fewer than \( k \) digits), we must pad the representation on the left with zeros so that it has exactly \( k \) digits. In other words, we interpret the expansion of \( x \) as

\[
x = \sum_{j=0}^{k-1} d_j' \cdot 3^{k-1-j},
\]

where the digits \( d_j' \) include leading zeros as needed. Then, for each \( j = 0, 1, \ldots, k-1 \), we have

\[
d_j' = \left( \left\lfloor \frac{x}{3^{k-1-j}} \right\rfloor \right) \mod 3.
\]

Finally, reversing the initial shift,

\[
d_j = d_j' - 1, \quad j = 0, 1, \ldots, k-1,
\]

retrieves the original balanced ternary digits.

Conclusion
The decoding operation precisely recovers \( N \), and therefore the original sequence of digits. In other words,

\[
U(P(D)) = D.
\]

This completes the corrected proof that the packing and unpacking functions are exact inverses.

F Inference implementation on CPUs and benchmarking across hardware.

F.1 Additional Implementation details of TQ2.
For quantization, the packed value calculation and detailed encoding steps are as follows:

• Packing: \( q_{\text{packed}} = q_0 + 4q_1 + 16q_2 + 64q_3 \).

• Storage: 64 bytes for quantized elements + 2 bytes for the float16 scaling factor \( d_i \), totaling 66 bytes per block.

For dequantization, the explicit unpacking procedure involves: \( q_0 = q_{\text{packed}} \mod 4,\ q_1 = \left\lfloor \frac{q_{\text{packed}}}{4} \right\rfloor \mod 4,\ q_2 = \left\lfloor \frac{q_{\text{packed}}}{16} \right\rfloor \mod 4,\ q_3 = \left\lfloor \frac{q_{\text{packed}}}{64} \right\rfloor \mod 4 \). The ternary storage method uses only 2 bits per element, with minimal overhead from the float16 scale per block. The process relies on hardware-friendly bitwise operations for fast packing and unpacking, making it suitable for large-scale deployments in memory-constrained environments while maintaining a balance between numerical fidelity and storage efficiency.

F.2 Additional Implementation details of TQ1.

Ternary digit extraction using fixed-point and bitwise operations. In this optimized decoding approach, we define \( i \) as the index variable, which represents the iteration counter for extracting each trit. The index \( i \) ranges from 0 to 4, as we are extracting \( k = 5 \) trits from a packed byte \( b \). The procedure begins by setting \( b_0 = b \). For each iteration \( i \), we multiply the current value \( b_i \) by 3, yielding a 10-bit intermediate value. The high byte of this value is then extracted to obtain the ternary digit \( d_i' = \left\lfloor \frac{b_i \cdot 3}{2^8} \right\rfloor \), where \( d_i' \in \{0, 1, 2\} \). After this, the remainder is updated for the next iteration using the operation \( b_{i+1} = (b_i \cdot 3) \& 0xFF \). This process repeats for all iterations \( i = 0, 1, 2, 3, 4 \), extracting the corresponding ternary digits. Once all the trits \( d_i' \) are extracted, they are normalized by subtracting 1, mapping the values from \( \{0, 1, 2\} \) to \( \{-1, 0, 1\} \).

Algorithm 1 Ternary Digit Extraction Using Fixed-Point and Bitwise Operations

<table>
  <tr>
    <th>Input:</th>
    <td>Packed byte \( b \)</td>
  </tr>
  <tr>
    <th>Output:</th>
    <td>Extracted ternary digits \( d_0, d_1, d_2, d_3, d_4 \)</td>
  </tr>
  <tr>
    <td colspan="2">Initialize \( b_0 = b \)</td>
  </tr>
  <tr>
    <td colspan="2"><b>for</b> each iteration \( i = 0, 1, 2, 3, 4 \) <b>do</b></td>
  </tr>
  <tr>
    <td colspan="2">Multiply \( b_i \) by 3 to get a 10-bit intermediate value</td>
  </tr>
  <tr>
    <td colspan="2">Extract high byte to get ternary digit \( d_i' = \left\lfloor \frac{b_i \cdot 3}{2^8} \right\rfloor \)</td>
  </tr>
  <tr>
    <td colspan="2">Update remainder for next iteration: \( b_{i+1} = (b_i \cdot 3) \& 0xFF \)</td>
  </tr>
  <tr>
    <td colspan="2"><b>end for</b></td>
  </tr>
  <tr>
    <td colspan="2">Normalize extracted digits by subtracting 1, mapping \( \{0, 1, 2\} \) to \( \{-1, 0, 1\} \)</td>
  </tr>
  <tr>
    <th>Return:</th>
    <td>\( d_0, d_1, d_2, d_3, d_4 \)</td>
  </tr>
</table>

This iterative method replaces the costly division and modulo operations with fixed-point arithmetic and bitwise masking, both of which are highly optimized for SIMD implementations. The structure of this approach minimizes data dependencies, enabling the parallel extraction of trits across multiple packed bytes. Furthermore, by leveraging the near
Figure 10: Comparison of output tokens for different model sizes running on a AMD EPYC 7502 laptop: (Left) Output tokens (for a 256 prompt with 64 output tokens). (Right) Output tokens per second versus model size. For more details, refer to Table 5

equivalence of \(3^5\) and \(2^8\), it achieves efficient decoding of ternary values with computational complexity that scales linearly with \(k\). The avoidance of costly arithmetic operations and compatibility with SIMD architectures make this approach particularly well-suited for high-performance applications involving ternary arithmetic.

F.3 Benchmarking across various hardware.
We present a comprehensive benchmarking analysis of quantization kernels: TQ1 (1.6 bits), TQ2 (2 bits), Q4 (4 bits, implementation provided in ggml), and FP16 (16 bits), across various model sizes ranging from 560M to 3.9B parameters and different token configurations. The benchmarks are executed on the AMD EPYC 7502 and Apple M4 Max (14 CPU cores). Detailed results are presented in Tables 6 and 5. It should be noted that further hardware-level optimizations are possible and will be addressed in future work.

Prompt Encoding Performance: On the Apple M4 Max, TQ1 and TQ2 outperformed FP16 and Q4, particularly for longer prompts, indicating efficient utilization of lower-bit quantization for prompt processing on this architecture. In contrast, benchmarks on the AMD EPYC 7502 showed that FP16 achieved superior throughput for shorter prompts (32–128 tokens), while Q4 and TQ2 gained an advantage at 256 tokens, highlighting a precision vs. memory bandwidth trade-off. TQ1 underperformed FP16 and Q4 on this platform. Across both architectures, larger models led to reduced prompt encoding throughput.

Autoregressive Decoding Performance. Autoregressive decoding consistently demonstrated quantization’s performance benefits on both platforms. Quantized kernels (Q4, TQ2, TQ1) outperformed FP16 in output tokens per second. AMD EPYC 7502 showed substantial Q4 gains over FP16, with TQ2 further improving throughput, highlighting reduced precision benefits for decoding. Apple M4 Max showed even greater quantization improvements; TQ2 achieved highest throughput, followed by TQ1 and Q4, all exceeding FP16. Output token length (8-256) minimally impacted decoding throughput, suggesting independence within this range. Larger models reduced decoding throughput, consistent with prompt encoding.

Combined Prompt Encoding and Autoregressive Decoding Performance. The combined benchmark confirmed quantization advantages. For both platforms, quantized kernels, especially TQ2 and Q4, delivered higher overall tokens per second than FP16 in combined scenarios. Prompt/decode token ratio (256/8, 256/64, 256/128) influenced overall throughput. Increased decoding token proportion decreased overall tokens per second, reflecting lower decoding throughput relative to prompt encoding. Apple M4 Max demonstrated highest combined throughput with TQ1 and TQ2, particularly at higher decoding token ratios, indicating optimization for end-to-end generation on this architecture.

G TriRun Kernel Design for Accelerated Matrix Multiplication

This section presents the design of the TriRun kernel, which accelerates matrix multiplication \(A \times B \rightarrow C\), where \(A\) is stored in half-precision (16-bit floating point), \(B\) is quantized to 2 bits per element, and \(C\) is accumulated in single-precision (32-bit floating point) before optional conversion to half-precision. The kernel optimizes memory efficiency and computational throughput through
<table>
  <tr>
    <th rowspan="2">Configuration</th>
    <th colspan="5">Model Size</th>
  </tr>
  <tr>
    <th>Tokens</th>
    <th>Bits</th>
    <th>Kernel</th>
    <th>560M</th>
    <th>1.1B</th>
    <th>1.5B</th>
    <th>2.4B</th>
    <th>3.9B</th>
  </tr>
  <tr>
    <th colspan="9">Prompt Encoding Benchmark (Prompt Tokens/seconds)</th>
  </tr>
  <tr>
    <td>32</td><td>16</td><td>FP16</td><td>223.1 ± 0.3</td><td>110.7 ± 0.1</td><td>77.2 ± 0.2</td><td>49.1 ± 0.0</td><td>30.7 ± 0.1</td><td></td>
  </tr>
  <tr>
    <td>32</td><td>4</td><td>Q4</td><td>182.4 ± 0.3</td><td>92.0 ± 0.1</td><td>70.2 ± 0.9</td><td>44.2 ± 0.8</td><td>27.8 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>32</td><td>2</td><td>TQ2</td><td>315.2 ± 0.6</td><td>165.1 ± 0.4</td><td>130.1 ± 0.1</td><td>83.7 ± 0.0</td><td>51.3 ± 0.1</td><td></td>
  </tr>
  <tr>
    <td>32</td><td>1.6</td><td>TQ1</td><td>181.1 ± 0.2</td><td>89.2 ± 0.1</td><td>70.0 ± 0.1</td><td>40.3 ± 0.1</td><td>25.4 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>64</td><td>16</td><td>FP16</td><td>233.0 ± 0.3</td><td>113.5 ± 0.1</td><td>80.2 ± 0.1</td><td>49.7 ± 0.6</td><td>31.0 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>64</td><td>4</td><td>Q4</td><td>182.9 ± 0.5</td><td>98.7 ± 0.1</td><td>69.3 ± 0.2</td><td>44.2 ± 0.7</td><td>27.6 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>64</td><td>2</td><td>TQ2</td><td>320.5 ± 0.5</td><td>179.2 ± 0.3</td><td>130.6 ± 0.1</td><td>83.9 ± 0.3</td><td>51.1 ± 0.1</td><td></td>
  </tr>
  <tr>
    <td>64</td><td>1</td><td>TQ1</td><td>180.9 ± 0.5</td><td>89.6 ± 0.1</td><td>66.6 ± 0.4</td><td>40.3 ± 0.0</td><td>25.3 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>128</td><td>16</td><td>FP16</td><td>228.6 ± 1.1</td><td>116.2 ± 0.2</td><td>81.1 ± 0.2</td><td>51.5 ± 0.1</td><td>32.1 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>128</td><td>4</td><td>Q4</td><td>179.2 ± 0.7</td><td>97.0 ± 0.6</td><td>69.3 ± 1.0</td><td>44.8 ± 0.1</td><td>27.9 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>128</td><td>2</td><td>TQ2</td><td>305.2 ± 2.2</td><td>174.8 ± 0.1</td><td>124.2 ± 0.1</td><td>81.9 ± 0.1</td><td>51.5 ± 0.1</td><td></td>
  </tr>
  <tr>
    <td>128</td><td>1.6</td><td>TQ1</td><td>177.2 ± 0.3</td><td>90.4 ± 0.1</td><td>65.8 ± 0.1</td><td>40.5 ± 0.1</td><td>25.5 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>256</td><td>16</td><td>FP16</td><td>220.3 ± 0.6</td><td>105.9 ± 1.0</td><td>79.9 ± 0.1</td><td>50.0 ± 0.1</td><td>31.3 ± 0.1</td><td></td>
  </tr>
  <tr>
    <td>256</td><td>4</td><td>Q4</td><td>170.1 ± 0.1</td><td>90.0 ± 2.2</td><td>69.3 ± 0.1</td><td>42.9 ± 0.1</td><td>27.3 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>256</td><td>2</td><td>TQ2</td><td>287.1 ± 2.2</td><td>176.5 ± 1.7</td><td>122.4 ± 0.1</td><td>77.5 ± 0.3</td><td>49.9 ± 0.1</td><td></td>
  </tr>
  <tr>
    <td>256</td><td>1.6</td><td>TQ1</td><td>169.5 ± 0.5</td><td>88.0 ± 1.0</td><td>64.3 ± 0.0</td><td>39.4 ± 0.0</td><td>24.9 ± 0.0</td><td></td>
  </tr>
  <tr>
    <th colspan="9">Autoregressive Decoding Benchmark (Output Tokens/seconds)</th>
  </tr>
  <tr>
    <td>8</td><td>16</td><td>FP16</td><td>37.6 ± 0.2</td><td>16.9 ± 0.0</td><td>14.1 ± 0.0</td><td>8.6 ± 0.0</td><td>5.4 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>8</td><td>4</td><td>Q4</td><td>83.0 ± 0.0</td><td>47.1 ± 0.2</td><td>35.5 ± 0.0</td><td>22.6 ± 0.0</td><td>15.1 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>8</td><td>2</td><td>TQ2</td><td>135.1 ± 0.3</td><td>84.6 ± 0.1</td><td>62.0 ± 0.1</td><td>42.0 ± 0.0</td><td>29.2 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>8</td><td>1.6</td><td>TQ1</td><td>102.0 ± 0.8</td><td>62.5 ± 0.1</td><td>48.6 ± 0.0</td><td>30.1 ± 0.1</td><td>20.7 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>64</td><td>16</td><td>FP16</td><td>37.4 ± 0.0</td><td>17.8 ± 0.0</td><td>14.0 ± 0.0</td><td>8.7 ± 0.0</td><td>5.4 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>64</td><td>4</td><td>Q4</td><td>83.1 ± 0.6</td><td>46.3 ± 0.7</td><td>35.2 ± 0.1</td><td>23.1 ± 0.0</td><td>15.0 ± 0.1</td><td></td>
  </tr>
  <tr>
    <td>64</td><td>2</td><td>TQ2</td><td>126.3 ± 0.0</td><td>83.2 ± 0.1</td><td>60.8 ± 0.5</td><td>44.1 ± 0.1</td><td>28.9 ± 0.3</td><td></td>
  </tr>
  <tr>
    <td>64</td><td>1.6</td><td>TQ1</td><td>105.4 ± 3.6</td><td>56.9 ± 0.3</td><td>45.1 ± 0.1</td><td>29.8 ± 0.0</td><td>20.5 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>128</td><td>16</td><td>FP16</td><td>37.2 ± 0.1</td><td>18.7 ± 0.3</td><td>14.3 ± 0.0</td><td>8.7 ± 0.0</td><td>5.4 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>128</td><td>4</td><td>Q4</td><td>85.6 ± 0.3</td><td>47.9 ± 0.1</td><td>37.5 ± 0.1</td><td>23.6 ± 0.0</td><td>15.0 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>128</td><td>2</td><td>TQ2</td><td>131.4 ± 0.3</td><td>82.2 ± 0.5</td><td>64.1 ± 0.0</td><td>43.1 ± 0.1</td><td>28.8 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>128</td><td>1.6</td><td>TQ1</td><td>104.4 ± 0.2</td><td>60.5 ± 0.0</td><td>47.7 ± 0.1</td><td>31.8 ± 0.1</td><td>20.8 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>256</td><td>16</td><td>FP16</td><td>36.4 ± 0.2</td><td>18.0 ± 0.0</td><td>14.0 ± 0.1</td><td>8.4 ± 0.0</td><td>5.3 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>256</td><td>4</td><td>Q4</td><td>82.3 ± 0.7</td><td>44.1 ± 0.7</td><td>36.1 ± 0.1</td><td>21.9 ± 0.1</td><td>14.6 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>256</td><td>2</td><td>TQ2</td><td>117.5 ± 3.9</td><td>72.0 ± 0.3</td><td>59.7 ± 0.5</td><td>37.6 ± 0.1</td><td>27.7 ± 0.2</td><td></td>
  </tr>
  <tr>
    <td>256</td><td>1.6</td><td>TQ1</td><td>94.5 ± 0.4</td><td>58.6 ± 0.6</td><td>44.0 ± 0.1</td><td>31.0 ± 0.0</td><td>19.6 ± 0.3</td><td></td>
  </tr>
  <tr>
    <th colspan="9">Prompt Encoding + Autoregressive Decoding Benchmark (Tokens/seconds)</th>
  </tr>
  <tr>
    <td>256/8</td><td>16</td><td>FP16</td><td>190.2 ± 0.1</td><td>90.3 ± 0.3</td><td>73.7 ± 0.1</td><td>43.4 ± 0.1</td><td>26.6 ± 0.1</td><td></td>
  </tr>
  <tr>
    <td>256/8</td><td>4</td><td>Q4</td><td>167.6 ± 0.2</td><td>86.0 ± 0.2</td><td>70.2 ± 0.2</td><td>41.1 ± 0.1</td><td>27.7 ± 0.3</td><td></td>
  </tr>
  <tr>
    <td>256/8</td><td>2</td><td>TQ2</td><td>271.7 ± 1.0</td><td>149.3 ± 0.4</td><td>118.4 ± 0.7</td><td>73.6 ± 0.1</td><td>50.7 ± 0.1</td><td></td>
  </tr>
  <tr>
    <td>256/8</td><td>1.6</td><td>TQ1</td><td>164.8 ± 0.6</td><td>82.4 ± 0.1</td><td>65.7 ± 0.1</td><td>38.4 ± 0.1</td><td>25.2 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>256/64</td><td>16</td><td>FP16</td><td>103.8 ± 0.7</td><td>52.3 ± 0.0</td><td>41.8 ± 0.1</td><td>25.1 ± 0.1</td><td>15.1 ± 0.6</td><td></td>
  </tr>
  <tr>
    <td>256/64</td><td>4</td><td>Q4</td><td>133.4 ± 0.8</td><td>72.4 ± 0.0</td><td>57.3 ± 0.3</td><td>35.2 ± 0.1</td><td>23.5 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>256/64</td><td>2</td><td>TQ2</td><td>206.4 ± 2.7</td><td>123.5 ± 0.1</td><td>103.3 ± 0.4</td><td>62.1 ± 0.2</td><td>43.7 ± 0.2</td><td></td>
  </tr>
  <tr>
    <td>256/64</td><td>1.6</td><td>TQ1</td><td>141.6 ± 0.1</td><td>74.5 ± 0.1</td><td>59.2 ± 0.3</td><td>35.9 ± 0.1</td><td>23.7 ± 0.3</td><td></td>
  </tr>
  <tr>
    <td>256/128</td><td>16</td><td>FP16</td><td>79.0 ± 1.1</td><td>39.1 ± 0.1</td><td>31.0 ± 0.0</td><td>18.8 ± 0.0</td><td>11.9 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>256/128</td><td>4</td><td>Q4</td><td>117.0 ± 0.3</td><td>63.2 ± 0.8</td><td>51.8 ± 2.1</td><td>31.7 ± 0.1</td><td>20.5 ± 0.0</td><td></td>
  </tr>
  <tr>
    <td>256/128</td><td>2</td><td>TQ2</td><td>177.2 ± 1.0</td><td>104.6 ± 1.0</td><td>85.7 ± 0.1</td><td>54.5 ± 0.3</td><td>37.9 ± 0.2</td><td></td>
  </tr>
  <tr>
    <td>256/128</td><td>1.6</td><td>TQ1</td><td>125.5 ± 0.5</td><td>68.5 ± 0.2</td><td>54.6 ± 0.2</td><td>33.7 ± 0.1</td><td>22.4 ± 0.0</td><td></td>
  </tr>
</table>

Table 5: Tokens per second for different model sizes and quantization kernels with varying prompt lengths on AMD EPYC 7502. Values represent mean ± standard deviation.
<table>
  <tr>
    <th rowspan="2">Configuration</th>
    <th colspan="5">Model Size</th>
  </tr>
  <tr>
    <th>Tokens</th>
    <th>Bits</th>
    <th>Kernel</th>
    <th>560M</th>
    <th>1.1B</th>
    <th>1.5B</th>
    <th>2.4B</th>
    <th>3.9B</th>
  </tr>
  <tr>
    <th colspan="9">Prompt Encoding Benchmark (Prompt Tokens/seconds)</th>
  </tr>
  <tr>
    <td>32</td><td>16</td><td>FP16</td><td>730.0 ± 6.5</td><td>417.6 ± 17.3</td><td>295.9 ± 2.4</td><td>152.4 ± 0.2</td><td>91.2 ± 0.7</td>
  </tr>
  <tr>
    <td>32</td><td>4</td><td>Q4</td><td>490.4 ± 3.2</td><td>270.0 ± 1.6</td><td>195.2 ± 0.9</td><td>106.1 ± 0.6</td><td>61.8 ± 0.3</td>
  </tr>
  <tr>
    <td>32</td><td>2</td><td>TQ2</td><td>543.0 ± 3.5</td><td>305.1 ± 1.1</td><td>221.9 ± 0.6</td><td>118.3 ± 0.9</td><td>68.6 ± 0.3</td>
  </tr>
  <tr>
    <td>32</td><td>1.6</td><td>TQ1</td><td>617.9 ± 2.8</td><td>362.7 ± 2.5</td><td>276.5 ± 5.7</td><td>145.5 ± 1.4</td><td>84.0 ± 0.0</td>
  </tr>
  <tr>
    <td>64</td><td>16</td><td>FP16</td><td>1223.8 ± 21.7</td><td>640.8 ± 2.1</td><td>440.4 ± 0.8</td><td>247.1 ± 0.5</td><td>144.7 ± 0.4</td>
  </tr>
  <tr>
    <td>64</td><td>4</td><td>Q4</td><td>886.1 ± 7.1</td><td>451.0 ± 1.3</td><td>320.7 ± 0.5</td><td>180.0 ± 1.1</td><td>105.5 ± 0.4</td>
  </tr>
  <tr>
    <td>64</td><td>2</td><td>TQ2</td><td>951.1 ± 10.6</td><td>501.7 ± 1.1</td><td>363.0 ± 1.2</td><td>198.7 ± 0.3</td><td>116.2 ± 0.3</td>
  </tr>
  <tr>
    <td>64</td><td>1.6</td><td>TQ1</td><td>1104.4 ± 11.5</td><td>578.3 ± 9.2</td><td>435.7 ± 0.4</td><td>235.2 ± 0.7</td><td>136.1 ± 0.4</td>
  </tr>
  <tr>
    <td>128</td><td>16</td><td>FP16</td><td>1256.1 ± 4.2</td><td>788.2 ± 6.4</td><td>564.8 ± 1.7</td><td>328.0 ± 1.7</td><td>205.8 ± 0.6</td>
  </tr>
  <tr>
    <td>128</td><td>4</td><td>Q4</td><td>1081.5 ± 4.6</td><td>629.7 ± 2.4</td><td>460.6 ± 0.6</td><td>266.6 ± 0.5</td><td>162.5 ± 0.3</td>
  </tr>
  <tr>
    <td>128</td><td>2</td><td>TQ2</td><td>1143.2 ± 7.2</td><td>677.9 ± 2.7</td><td>494.3 ± 2.7</td><td>285.8 ± 1.0</td><td>175.8 ± 0.3</td>
  </tr>
  <tr>
    <td>128</td><td>1.6</td><td>TQ1</td><td>1223.2 ± 11.1</td><td>769.6 ± 4.6</td><td>556.8 ± 0.6</td><td>320.4 ± 2.1</td><td>197.9 ± 0.2</td>
  </tr>
  <tr>
    <td>256</td><td>16</td><td>FP16</td><td>1485.2 ± 3.3</td><td>785.5 ± 2.8</td><td>611.6 ± 1.3</td><td>367.5 ± 0.7</td><td>243.6 ± 3.6</td>
  </tr>
  <tr>
    <td>256</td><td>4</td><td>Q4</td><td>1350.2 ± 3.4</td><td>710.3 ± 1.4</td><td>545.2 ± 2.1</td><td>325.8 ± 0.5</td><td>209.3 ± 2.2</td>
  </tr>
  <tr>
    <td>256</td><td>2</td><td>TQ2</td><td>1398.6 ± 2.6</td><td>738.9 ± 2.4</td><td>561.7 ± 1.9</td><td>339.2 ± 0.7</td><td>225.9 ± 0.5</td>
  </tr>
  <tr>
    <td>256</td><td>1.6</td><td>TQ1</td><td>1468.0 ± 3.5</td><td>786.7 ± 1.7</td><td>601.8 ± 2.2</td><td>361.6 ± 0.2</td><td>242.9 ± 0.5</td>
  </tr>
  <tr>
    <th colspan="9">Autoregressive Decoding Benchmark (Output Tokens/seconds)</th>
  </tr>
  <tr>
    <td>8</td><td>16</td><td>FP16</td><td>170.9 ± 0.8</td><td>92.6 ± 0.1</td><td>71.5 ± 0.0</td><td>44.5 ± 0.0</td><td>28.1 ± 0.0</td>
  </tr>
  <tr>
    <td>8</td><td>4</td><td>Q4</td><td>237.8 ± 0.3</td><td>134.8 ± 0.3</td><td>107.2 ± 0.3</td><td>64.8 ± 0.6</td><td>43.3 ± 0.7</td>
  </tr>
  <tr>
    <td>8</td><td>2</td><td>TQ2</td><td>278.7 ± 0.6</td><td>167.4 ± 0.1</td><td>134.0 ± 0.2</td><td>86.6 ± 0.0</td><td>57.6 ± 0.1</td>
  </tr>
  <tr>
    <td>8</td><td>1.6</td><td>TQ1</td><td>228.8 ± 0.5</td><td>125.7 ± 0.1</td><td>99.6 ± 0.1</td><td>62.3 ± 0.1</td><td>40.0 ± 0.1</td>
  </tr>
  <tr>
    <td>64</td><td>16</td><td>FP16</td><td>169.4 ± 0.3</td><td>92.4 ± 0.1</td><td>71.3 ± 0.0</td><td>44.5 ± 0.0</td><td>28.0 ± 0.2</td>
  </tr>
  <tr>
    <td>64</td><td>4</td><td>Q4</td><td>236.8 ± 0.7</td><td>134.0 ± 0.0</td><td>106.1 ± 0.1</td><td>66.2 ± 0.5</td><td>42.6 ± 0.2</td>
  </tr>
  <tr>
    <td>64</td><td>2</td><td>TQ2</td><td>279.7 ± 0.1</td><td>166.5 ± 0.1</td><td>132.2 ± 0.1</td><td>86.1 ± 0.0</td><td>57.5 ± 0.0</td>
  </tr>
  <tr>
    <td>64</td><td>1.6</td><td>TQ1</td><td>227.4 ± 0.5</td><td>125.4 ± 0.1</td><td>98.6 ± 0.1</td><td>62.1 ± 0.0</td><td>38.3 ± 0.0</td>
  </tr>
  <tr>
    <td>128</td><td>16</td><td>FP16</td><td>171.5 ± 0.1</td><td>91.6 ± 0.1</td><td>71.0 ± 0.1</td><td>44.5 ± 0.1</td><td>28.1 ± 0.1</td>
  </tr>
  <tr>
    <td>128</td><td>4</td><td>Q4</td><td>232.7 ± 0.3</td><td>132.2 ± 0.1</td><td>105.5 ± 0.1</td><td>67.0 ± 0.1</td><td>43.5 ± 0.3</td>
  </tr>
  <tr>
    <td>128</td><td>2</td><td>TQ2</td><td>281.3 ± 1.0</td><td>164.0 ± 0.2</td><td>132.0 ± 0.1</td><td>85.1 ± 0.0</td><td>56.9 ± 0.0</td>
  </tr>
  <tr>
    <td>128</td><td>1.6</td><td>TQ1</td><td>225.3 ± 1.6</td><td>124.5 ± 0.1</td><td>97.8 ± 0.1</td><td>61.6 ± 0.0</td><td>39.7 ± 0.0</td>
  </tr>
  <tr>
    <td>256</td><td>16</td><td>FP16</td><td>166.0 ± 0.4</td><td>89.4 ± 0.4</td><td>69.1 ± 0.2</td><td>43.3 ± 0.2</td><td>27.7 ± 0.0</td>
  </tr>
  <tr>
    <td>256</td><td>4</td><td>Q4</td><td>225.1 ± 0.2</td><td>128.4 ± 0.1</td><td>101.6 ± 0.6</td><td>63.7 ± 0.5</td><td>41.1 ± 0.4</td>
  </tr>
  <tr>
    <td>256</td><td>2</td><td>TQ2</td><td>268.9 ± 0.2</td><td>158.6 ± 1.0</td><td>128.5 ± 0.2</td><td>82.9 ± 0.2</td><td>55.6 ± 0.0</td>
  </tr>
  <tr>
    <td>256</td><td>1.6</td><td>TQ1</td><td>217.2 ± 1.1</td><td>121.1 ± 0.6</td><td>95.5 ± 0.0</td><td>60.3 ± 0.1</td><td>39.0 ± 0.0</td>
  </tr>
  <tr>
    <th colspan="9">Prompt Encoding + Autoregressive Decoding Benchmark (Tokens/seconds)</th>
  </tr>
  <tr>
    <td>256/8</td><td>16</td><td>FP16</td><td>1142.4 ± 7.0</td><td>628.9 ± 1.4</td><td>489.9 ± 1.5</td><td>296.4 ± 0.7</td><td>194.6 ± 0.3</td>
  </tr>
  <tr>
    <td>256/8</td><td>4</td><td>Q4</td><td>1165.6 ± 2.6</td><td>620.8 ± 2.2</td><td>481.8 ± 1.4</td><td>288.8 ± 0.5</td><td>186.6 ± 0.9</td>
  </tr>
  <tr>
    <td>256/8</td><td>2</td><td>TQ2</td><td>1234.7 ± 0.9</td><td>668.9 ± 1.3</td><td>513.5 ± 1.1</td><td>308.6 ± 0.3</td><td>202.1 ± 0.5</td>
  </tr>
  <tr>
    <td>256/8</td><td>1.6</td><td>TQ1</td><td>1241.5 ± 4.7</td><td>665.5 ± 4.9</td><td>517.6 ± 1.2</td><td>311.7 ± 0.5</td><td>205.4 ± 0.3</td>
  </tr>
  <tr>
    <td>256/64</td><td>16</td><td>FP16</td><td>550.4 ± 2.9</td><td>298.9 ± 0.5</td><td>234.8 ± 0.2</td><td>144.1 ± 0.3</td><td>93.6 ± 0.1</td>
  </tr>
  <tr>
    <td>256/64</td><td>4</td><td>Q4</td><td>648.3 ± 1.1</td><td>358.6 ± 1.7</td><td>286.8 ± 0.3</td><td>175.2 ± 0.7</td><td>114.6 ± 0.4</td>
  </tr>
  <tr>
    <td>256/64</td><td>2</td><td>TQ2</td><td>726.3 ± 1.9</td><td>415.5 ± 0.5</td><td>328.3 ± 0.3</td><td>205.2 ± 0.2</td><td>135.9 ± 0.4</td>
  </tr>
  <tr>
    <td>256/64</td><td>1.6</td><td>TQ1</td><td>652.7 ± 2.4</td><td>364.1 ± 0.6</td><td>285.7 ± 0.3</td><td>177.0 ± 0.2</td><td>115.7 ± 0.4</td>
  </tr>
  <tr>
    <td>256/128</td><td>16</td><td>FP16</td><td>385.8 ± 3.5</td><td>210.4 ± 1.2</td><td>165.1 ± 0.5</td><td>102.1 ± 0.2</td><td>66.1 ± 0.2</td>
  </tr>
  <tr>
    <td>256/128</td><td>4</td><td>Q4</td><td>475.3 ± 0.7</td><td>270.2 ± 0.3</td><td>216.6 ± 0.4</td><td>133.1 ± 0.6</td><td>87.5 ± 0.4</td>
  </tr>
  <tr>
    <td>256/128</td><td>2</td><td>TQ2</td><td>540.8 ± 0.6</td><td>317.9 ± 0.6</td><td>255.4 ± 0.2</td><td>160.8 ± 0.1</td><td>108.2 ± 0.1</td>
  </tr>
  <tr>
    <td>256/128</td><td>1.6</td><td>TQ1</td><td>466.8 ± 0.5</td><td>265.8 ± 0.4</td><td>209.9 ± 0.1</td><td>130.9 ± 0.1</td><td>86.0 ± 0.3</td>
  </tr>
</table>

Table 6: Tokens per Second for Different Model Sizes and Quantization Kernels M4 Max (14CPU Coresz). Values represent mean ± standard deviation.
specialized data layouts, dequantization strategies, and tensor-core utilization. Key components include:

G.1 Data Organization and Quantization

2-Bit Weight Matrix (B Storage) The 2-bit quantized elements of \( B \) are packed into 64-bit int2 vectors, where each 32-bit integer contains 16 quantized weights. During loading, 64-bit global memory transactions retrieve 32 weights per int2, minimizing memory bandwidth. To align with tensor-core requirements, these packed values are asynchronously copied to shared memory via cp.async instructions, then unpacked into 16-bit fragments for computation.

Half-Precision Matrix (A Access) Matrix \( A \) is stored in half-precision and loaded via 128-bit int4 vectors, fetching eight elements per transaction. This aligns with the 16-byte memory alignment optimal for GPU global memory accesses. Subsequent stages repack these into \( 16 \times 16 \) submatrices compatible with tensor-core operations.

G.2 Dequantization and Tensor-Core Computation

The dequant function performs dequantization of 2-bit integer values into half-precision floating-point representations, employing hardware-optimized bitwise operations and fused arithmetic to enable efficient tensor core execution. Rather than relying on conventional shift-and-mask techniques, the implementation decomposes each 32-bit word—which encodes sixteen 2-bit weights—using a specialized bitwise operation that leverages a tailored mask to both isolate the individual weight segments and embed a predetermined FP16 exponent. Following this, an integrated arithmetic fusion stage applies a zero-point adjustment, effectively adding 1.0 to the extracted values, and performs dynamic range scaling through a fused multiply-add operation. This approach diverges from the traditional scale \((w - \text{zero\_point})\) formulation by consolidating multiple arithmetic steps into a single, hardware-specific sequence. Subsequently, per-group FP16 scales are applied to the dequantized values, which are then stored in register-based fragments (FragB) to minimize shared memory contention. Later, the kernel employs ldmatrix.sync.aligned.m8n8.x4 to load \( A \) and \( B \) fragments into tensor-core registers. Each mma.sync.aligned.m16n8k16 operation computes a \( 16 \times 8 \times 16 \) submatrix product, accumulating results into 32-bit floating-point fragments (FragC) for numerical stability. By unrolling across submatrix tiles, the kernel fully utilizes tensor-core throughput while maintaining warp-level synchronization.

G.3 Memory Latency Hiding via Asynchronous Pipelines

To overlap computation with memory transfers, the kernel implements a four-stage software pipeline with double buffering. Key mechanisms include:

• Asynchronous Data Copies: cp.async instructions prefetch \( A \) and \( B \) tiles into shared memory without stalling computation threads.

• Double Buffering: Two shared memory buffers alternate between data ingestion (from global memory) and consumption (by tensor cores), ensuring continuous utilization of memory and compute units.

• cp.async Synchronization: Warps issue cp.async.commit_group to batch memory transactions and cp.async.wait_group to enforce dependencies, preventing read-after-write hazards.

![Speedup across hardware over the years using TriRun kernels](page_1012_1347_393_246.png)

Figure 11: Speedup across hardware over the years using TriRun kernels

G.4 Precision-Preserving Accumulation

Intra-Warp Reduction Partial sums within a thread block are reduced across warps using shared memory. A tree-based summation merges per-warp FragC outputs, minimizing shared memory bank conflicts through staggered access patterns.
Global Memory Atomic Reduction  For outputs spanning multiple thread blocks, atomic 32-bit floating-point additions ensure correct inter-block accumulation. Final results are converted to half-precision (if specified) using round-to-nearest-even mode, balancing precision and storage efficiency.

G.5 Performance Configuration

Thread blocks (256 threads) balance register pressure (128/thread) and occupancy (8 warps/block). Tile dimensions adapt to problem size: 128×128 tiles for small batches (\( m \leq 16 \)) and 64×256 tiles for larger workloads. The 96 KB shared memory budget supports four concurrent pipeline stages, sustaining 98% tensor core utilization across varied workloads. This implementation demonstrates that 2-bit quantized inference can achieve near-FP16 throughput while maintaining numerical fidelity, providing a practical solution for deploying compressed deep learning models on modern GPUs.

G.6 TriRun performance benchmark across various Nvidia hardware.

Performance Acceleration of Ternary Linear Layers in Transformer Blocks. TriRun kernels in transformer blocks including ternary linear layers demonstrate significant performance improvements across various hardware configurations. As shown in Figure 12, we evaluated TriRun’s performance against a standard FP16 PyTorch implementation (CUTLASS) across multiple NVIDIA GPU platforms, including the L40, A100 SXM/PCIe, A40, 3090, A30, L40s, and RTX 4090. The performance analysis covered models ranging from 3B to 405B parameters under varying batch sizes, with speedup quantified as the ratio of FP16 baseline execution time to TriRun execution time. The results indicate sustained performance for batch sizes up to 16–32 across all tested GPUs. Moreover, larger models, which incorporate a higher proportion of ternary weights relative to their total parameters, achieve more substantial speedups compared to smaller models. This performance trend suggests a positive correlation between model size and the efficiency gains offered by TriRun’s implementation. Detailed per-layer results are provided in Tables Table 7 and Table 8.

End-to-End Generation Performance. Figure 13 presents the end-to-end token generation speedup. For a comprehensive view, Figure 13 (with detailed results in Table 9, Table 10, Table 11, and Table 12) shows the total time for end-to-end token generation using TriRun on Nvidia L40s, L40, A40, and 4090 GPUs for models ranging from 7B to 70B parameters. This reflects the output token throughput, with TriRun achieving up to approximately 5× speedup. The slightly lower end-to-end speedup compared to the per-layer results can be attributed to additional inference overheads beyond the linear layers, which are specifically accelerated by TriRun. In this case as well, larger models achieve higher speedups due to their increased proportion of ternary weights.

H Artifacts Released

To foster open research, we will be making the artifacts from this paper publicly available. The following resources are provided:

• TriTera Suite. We are releasing all models from the Tritera suite (as described in Section 2), along with all intermediate checkpoints, under the MIT license.

• TriRun Kernels. We plan to open-source the TriRun Library (as detailed in Section 4) under the Apache 2.0 license.
<table>
  <tr>
    <th rowspan="2">Model Size (Parameters)</th>
    <th rowspan="2">GPU Type</th>
    <th colspan="8">Batch Size</th>
  </tr>
  <tr>
    <th>1</th>
    <th>2</th>
    <th>4</th>
    <th>8</th>
    <th>16</th>
    <th>32</th>
    <th>64</th>
    <th>128</th>
  </tr>
  <tr>
    <td colspan="11"><b>A40</b></td>
  </tr>
  <tr>
    <td>405B</td>
    <td>A40</td>
    <td>7.85</td>
    <td>7.65</td>
    <td>7.66</td>
    <td>7.70</td>
    <td>7.50</td>
    <td>7.27</td>
    <td>3.91</td>
    <td>1.98</td>
  </tr>
  <tr>
    <td>123B</td>
    <td>A40</td>
    <td>7.67</td>
    <td>7.67</td>
    <td>7.68</td>
    <td>7.63</td>
    <td>7.65</td>
    <td>6.99</td>
    <td>3.69</td>
    <td>1.97</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>A40</td>
    <td>7.36</td>
    <td>7.39</td>
    <td>7.45</td>
    <td>7.29</td>
    <td>7.46</td>
    <td>6.80</td>
    <td>3.64</td>
    <td>1.90</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>A40</td>
    <td>7.22</td>
    <td>7.29</td>
    <td>7.35</td>
    <td>7.36</td>
    <td>7.46</td>
    <td>6.29</td>
    <td>3.57</td>
    <td>1.99</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>A40</td>
    <td>6.76</td>
    <td>6.91</td>
    <td>6.72</td>
    <td>6.92</td>
    <td>6.93</td>
    <td>5.89</td>
    <td>3.33</td>
    <td>1.86</td>
  </tr>
  <tr>
    <td>8B</td>
    <td>A40</td>
    <td>6.42</td>
    <td>6.62</td>
    <td>6.62</td>
    <td>6.66</td>
    <td>6.55</td>
    <td>5.18</td>
    <td>2.99</td>
    <td>1.84</td>
  </tr>
  <tr>
    <td>3B</td>
    <td>A40</td>
    <td>5.39</td>
    <td>4.76</td>
    <td>5.50</td>
    <td>5.51</td>
    <td>5.56</td>
    <td>4.35</td>
    <td>2.81</td>
    <td>1.93</td>
  </tr>
  <tr>
    <td colspan="11"><b>3090</b></td>
  </tr>
  <tr>
    <td>405B</td>
    <td>3090</td>
    <td>4.92</td>
    <td>4.70</td>
    <td>4.70</td>
    <td>4.98</td>
    <td>5.04</td>
    <td>2.72</td>
    <td>1.41</td>
    <td>1.09</td>
  </tr>
  <tr>
    <td>123B</td>
    <td>3090</td>
    <td>4.83</td>
    <td>4.75</td>
    <td>4.75</td>
    <td>5.46</td>
    <td>5.50</td>
    <td>2.63</td>
    <td>1.34</td>
    <td>1.14</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>3090</td>
    <td>4.51</td>
    <td>4.69</td>
    <td>4.70</td>
    <td>4.72</td>
    <td>4.74</td>
    <td>2.48</td>
    <td>1.35</td>
    <td>1.10</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>3090</td>
    <td>4.42</td>
    <td>4.70</td>
    <td>4.76</td>
    <td>4.75</td>
    <td>4.65</td>
    <td>2.53</td>
    <td>1.62</td>
    <td>1.23</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>3090</td>
    <td>4.25</td>
    <td>4.39</td>
    <td>4.47</td>
    <td>4.41</td>
    <td>4.39</td>
    <td>2.40</td>
    <td>1.26</td>
    <td>1.15</td>
  </tr>
  <tr>
    <td>8B</td>
    <td>3090</td>
    <td>4.19</td>
    <td>4.25</td>
    <td>4.30</td>
    <td>4.32</td>
    <td>4.17</td>
    <td>2.37</td>
    <td>1.34</td>
    <td>1.15</td>
  </tr>
  <tr>
    <td>3B</td>
    <td>3090</td>
    <td>4.05</td>
    <td>3.53</td>
    <td>3.71</td>
    <td>3.71</td>
    <td>3.68</td>
    <td>2.32</td>
    <td>1.32</td>
    <td>1.24</td>
  </tr>
  <tr>
    <td colspan="11"><b>A30</b></td>
  </tr>
  <tr>
    <td>405B</td>
    <td>A30</td>
    <td>3.96</td>
    <td>3.98</td>
    <td>3.99</td>
    <td>4.00</td>
    <td>4.01</td>
    <td>2.80</td>
    <td>1.77</td>
    <td>1.35</td>
  </tr>
  <tr>
    <td>123B</td>
    <td>A30</td>
    <td>3.89</td>
    <td>3.90</td>
    <td>3.90</td>
    <td>3.91</td>
    <td>3.90</td>
    <td>2.66</td>
    <td>1.68</td>
    <td>1.12</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>A30</td>
    <td>3.81</td>
    <td>3.82</td>
    <td>3.91</td>
    <td>3.94</td>
    <td>3.90</td>
    <td>2.70</td>
    <td>1.64</td>
    <td>1.19</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>A30</td>
    <td>3.62</td>
    <td>3.62</td>
    <td>3.67</td>
    <td>3.81</td>
    <td>3.78</td>
    <td>2.76</td>
    <td>1.71</td>
    <td>1.29</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>A30</td>
    <td>3.36</td>
    <td>3.44</td>
    <td>3.45</td>
    <td>3.48</td>
    <td>3.38</td>
    <td>2.47</td>
    <td>1.59</td>
    <td>1.48</td>
  </tr>
  <tr>
    <td>8B</td>
    <td>A30</td>
    <td>3.28</td>
    <td>3.30</td>
    <td>3.32</td>
    <td>3.33</td>
    <td>3.31</td>
    <td>2.28</td>
    <td>1.61</td>
    <td>1.13</td>
  </tr>
  <tr>
    <td>3B</td>
    <td>A30</td>
    <td>2.66</td>
    <td>2.90</td>
    <td>2.90</td>
    <td>2.94</td>
    <td>2.87</td>
    <td>2.08</td>
    <td>1.35</td>
    <td>1.39</td>
  </tr>
  <tr>
    <td colspan="11"><b>L4</b></td>
  </tr>
  <tr>
    <td>405B</td>
    <td>L4</td>
    <td>5.98</td>
    <td>6.10</td>
    <td>6.12</td>
    <td>6.24</td>
    <td>6.46</td>
    <td>5.38</td>
    <td>3.21</td>
    <td>1.00</td>
  </tr>
  <tr>
    <td>123B</td>
    <td>L4</td>
    <td>6.34</td>
    <td>6.05</td>
    <td>6.05</td>
    <td>6.08</td>
    <td>6.07</td>
    <td>5.24</td>
    <td>3.12</td>
    <td>1.63</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>L4</td>
    <td>5.91</td>
    <td>5.97</td>
    <td>5.96</td>
    <td>5.90</td>
    <td>5.86</td>
    <td>5.22</td>
    <td>3.15</td>
    <td>1.65</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>L4</td>
    <td>6.75</td>
    <td>5.87</td>
    <td>5.87</td>
    <td>5.84</td>
    <td>5.76</td>
    <td>5.40</td>
    <td>3.20</td>
    <td>1.80</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>L4</td>
    <td>7.75</td>
    <td>5.64</td>
    <td>5.64</td>
    <td>5.63</td>
    <td>6.16</td>
    <td>5.03</td>
    <td>3.13</td>
    <td>1.66</td>
  </tr>
  <tr>
    <td>8B</td>
    <td>L4</td>
    <td>5.31</td>
    <td>5.25</td>
    <td>5.27</td>
    <td>5.23</td>
    <td>5.30</td>
    <td>4.42</td>
    <td>2.79</td>
    <td>1.56</td>
  </tr>
  <tr>
    <td>3B</td>
    <td>L4</td>
    <td>7.74</td>
    <td>5.43</td>
    <td>5.43</td>
    <td>5.38</td>
    <td>5.40</td>
    <td>4.49</td>
    <td>2.83</td>
    <td>1.63</td>
  </tr>
</table>

Table 7: Speedup over FP16 PyTorch (using CUTLASS) across different batch sizes for all ternary linear layers in a transformer block, accounting for the matrix structures in models ranging from 3B to 405B on A40, 3090, A30 and L4 GPUs.
<table>
  <tr>
    <th rowspan="2">Model Size (Parameters)</th>
    <th rowspan="2">GPU Type</th>
    <th colspan="8">Batch Size</th>
  </tr>
  <tr>
    <th>1</th>
    <th>2</th>
    <th>4</th>
    <th>8</th>
    <th>16</th>
    <th>32</th>
    <th>64</th>
    <th>128</th>
  </tr>
  <tr>
    <td colspan="11"><b>L40</b></td>
  </tr>
  <tr>
    <td>405B</td>
    <td>L40</td>
    <td>7.99</td>
    <td>7.67</td>
    <td>7.69</td>
    <td>7.78</td>
    <td>7.79</td>
    <td>7.67</td>
    <td>4.04</td>
    <td>2.05</td>
  </tr>
  <tr>
    <td>123B</td>
    <td>L40</td>
    <td>8.47</td>
    <td>7.58</td>
    <td>7.57</td>
    <td>7.57</td>
    <td>7.52</td>
    <td>7.44</td>
    <td>3.97</td>
    <td>2.11</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>L40</td>
    <td>7.91</td>
    <td>7.39</td>
    <td>7.46</td>
    <td>7.49</td>
    <td>7.64</td>
    <td>6.90</td>
    <td>3.80</td>
    <td>2.02</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>L40</td>
    <td>7.93</td>
    <td>7.32</td>
    <td>7.27</td>
    <td>7.38</td>
    <td>7.57</td>
    <td>6.48</td>
    <td>3.63</td>
    <td>2.01</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>L40</td>
    <td>9.98</td>
    <td>6.83</td>
    <td>6.83</td>
    <td>6.83</td>
    <td>6.74</td>
    <td>5.46</td>
    <td>3.22</td>
    <td>1.86</td>
  </tr>
  <tr>
    <td>8B</td>
    <td>L40</td>
    <td>8.30</td>
    <td>6.57</td>
    <td>6.68</td>
    <td>6.61</td>
    <td>6.42</td>
    <td>4.41</td>
    <td>2.69</td>
    <td>1.70</td>
  </tr>
  <tr>
    <td>3B</td>
    <td>L40</td>
    <td>5.98</td>
    <td>6.97</td>
    <td>6.41</td>
    <td>6.83</td>
    <td>6.67</td>
    <td>3.05</td>
    <td>1.97</td>
    <td>1.49</td>
  </tr>
  <tr>
    <td colspan="11"><b>A100 (SXM)</b></td>
  </tr>
  <tr>
    <td>405B</td>
    <td>A100</td>
    <td>4.25</td>
    <td>4.25</td>
    <td>4.30</td>
    <td>4.28</td>
    <td>4.25</td>
    <td>3.05</td>
    <td>1.86</td>
    <td>1.28</td>
  </tr>
  <tr>
    <td>123B</td>
    <td>A100</td>
    <td>4.25</td>
    <td>4.11</td>
    <td>4.08</td>
    <td>4.14</td>
    <td>4.04</td>
    <td>3.06</td>
    <td>1.95</td>
    <td>1.12</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>A100</td>
    <td>4.66</td>
    <td>3.90</td>
    <td>3.87</td>
    <td>3.93</td>
    <td>3.95</td>
    <td>2.88</td>
    <td>1.81</td>
    <td>1.13</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>A100</td>
    <td>3.66</td>
    <td>3.63</td>
    <td>3.66</td>
    <td>3.64</td>
    <td>3.62</td>
    <td>2.70</td>
    <td>1.73</td>
    <td>1.13</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>A100</td>
    <td>3.33</td>
    <td>3.30</td>
    <td>3.42</td>
    <td>3.41</td>
    <td>3.42</td>
    <td>2.32</td>
    <td>1.57</td>
    <td>1.26</td>
  </tr>
  <tr>
    <td>8B</td>
    <td>A100</td>
    <td>2.79</td>
    <td>2.68</td>
    <td>2.82</td>
    <td>2.90</td>
    <td>2.86</td>
    <td>2.08</td>
    <td>1.49</td>
    <td>1.17</td>
  </tr>
  <tr>
    <td>3B</td>
    <td>A100</td>
    <td>2.26</td>
    <td>2.44</td>
    <td>2.26</td>
    <td>2.57</td>
    <td>2.22</td>
    <td>1.42</td>
    <td>1.45</td>
    <td>0.97</td>
  </tr>
  <tr>
    <td colspan="11"><b>A100 (PCIe)</b></td>
  </tr>
  <tr>
    <td>405B</td>
    <td>A100</td>
    <td>4.98</td>
    <td>4.96</td>
    <td>4.99</td>
    <td>5.06</td>
    <td>4.94</td>
    <td>3.52</td>
    <td>2.19</td>
    <td>1.23</td>
  </tr>
  <tr>
    <td>123B</td>
    <td>A100</td>
    <td>4.89</td>
    <td>4.88</td>
    <td>4.74</td>
    <td>4.86</td>
    <td>4.87</td>
    <td>3.51</td>
    <td>2.24</td>
    <td>1.27</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>A100</td>
    <td>4.67</td>
    <td>4.48</td>
    <td>4.49</td>
    <td>4.55</td>
    <td>4.58</td>
    <td>3.35</td>
    <td>2.10</td>
    <td>1.31</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>A100</td>
    <td>4.21</td>
    <td>4.24</td>
    <td>4.23</td>
    <td>4.24</td>
    <td>4.28</td>
    <td>3.10</td>
    <td>1.98</td>
    <td>1.29</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>A100</td>
    <td>3.86</td>
    <td>3.97</td>
    <td>3.57</td>
    <td>3.96</td>
    <td>3.96</td>
    <td>2.73</td>
    <td>1.82</td>
    <td>1.42</td>
  </tr>
  <tr>
    <td>8B</td>
    <td>A100</td>
    <td>3.27</td>
    <td>3.37</td>
    <td>3.33</td>
    <td>3.09</td>
    <td>3.37</td>
    <td>2.30</td>
    <td>1.64</td>
    <td>1.32</td>
  </tr>
  <tr>
    <td>3B</td>
    <td>A100</td>
    <td>2.93</td>
    <td>2.75</td>
    <td>2.74</td>
    <td>2.72</td>
    <td>2.93</td>
    <td>1.86</td>
    <td>1.22</td>
    <td>1.09</td>
  </tr>
  <tr>
    <td colspan="11"><b>4090</b></td>
  </tr>
  <tr>
    <td>405B</td>
    <td>4090</td>
    <td>7.67</td>
    <td>7.65</td>
    <td>7.69</td>
    <td>7.71</td>
    <td>7.77</td>
    <td>5.29</td>
    <td>2.74</td>
    <td>1.48</td>
  </tr>
  <tr>
    <td>123B</td>
    <td>4090</td>
    <td>7.64</td>
    <td>7.68</td>
    <td>7.68</td>
    <td>7.71</td>
    <td>7.74</td>
    <td>5.20</td>
    <td>2.80</td>
    <td>1.41</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>4090</td>
    <td>4.96</td>
    <td>6.65</td>
    <td>7.45</td>
    <td>7.51</td>
    <td>7.20</td>
    <td>5.02</td>
    <td>2.60</td>
    <td>1.40</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>4090</td>
    <td>7.25</td>
    <td>7.28</td>
    <td>6.82</td>
    <td>6.84</td>
    <td>7.34</td>
    <td>4.89</td>
    <td>2.59</td>
    <td>1.43</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>4090</td>
    <td>6.36</td>
    <td>6.39</td>
    <td>6.39</td>
    <td>6.39</td>
    <td>6.36</td>
    <td>4.18</td>
    <td>2.32</td>
    <td>1.43</td>
  </tr>
  <tr>
    <td>8B</td>
    <td>4090</td>
    <td>5.42</td>
    <td>5.54</td>
    <td>5.70</td>
    <td>5.94</td>
    <td>5.53</td>
    <td>3.56</td>
    <td>1.99</td>
    <td>1.37</td>
  </tr>
  <tr>
    <td>3B</td>
    <td>4090</td>
    <td>3.96</td>
    <td>4.29</td>
    <td>4.40</td>
    <td>4.32</td>
    <td>4.15</td>
    <td>2.52</td>
    <td>1.37</td>
    <td>1.37</td>
  </tr>
</table>

Table 8: Speedup over FP16 PyTorch (using CUTLASS) across different batch sizes for all ternary linear layers in a transformer block, accounting for the matrix structures in models ranging from 3B to 405B on L40, A100(SXM), A100 (PCIe) and 4090 GPUs.
Figure 12: We evaluate the performance of ternary layers in transformer blocks, showing near-optimal speedup over PyTorch FP16 on different NVIDIA GPUs using CUTLASS.
Figure 13: Comparison of TriRun kernels with the FP16 PyTorch baseline on NVIDIA L40S, L40, A40, and 4090 (top to bottom). More details in Table 9, Table 10, Table 11 and Table 12 (a) Left: Time to first token, (b) Right: Time per output token.
<table>
  <tr>
    <th rowspan="2">Size</th>
    <th rowspan="2">Kernel</th>
    <th rowspan="2">#GPU</th>
    <th colspan="8">Encoding</th>
  </tr>
  <tr>
    <th>1</th>
    <th>2</th>
    <th>4</th>
    <th>8</th>
    <th>16</th>
    <th>32</th>
    <th>64</th>
  </tr>
  <tr>
    <td>7B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0210</td>
    <td>0.0218</td>
    <td>0.0219</td>
    <td>0.0221</td>
    <td>0.0225</td>
    <td>0.0237</td>
    <td>0.0244</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0135</td>
    <td>0.0146</td>
    <td>0.0146</td>
    <td>0.0145</td>
    <td>0.0146</td>
    <td>0.0146</td>
    <td>0.0146</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>1.5556</td>
    <td>1.4932</td>
    <td>1.5000</td>
    <td>1.5241</td>
    <td>1.5411</td>
    <td>1.6233</td>
    <td>1.6712</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0380</td>
    <td>0.0401</td>
    <td>0.0402</td>
    <td>0.0405</td>
    <td>0.0411</td>
    <td>0.0431</td>
    <td>0.0461</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0184</td>
    <td>0.0195</td>
    <td>0.0193</td>
    <td>0.0194</td>
    <td>0.0207</td>
    <td>0.0195</td>
    <td>0.0195</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>2.0652</td>
    <td>2.0564</td>
    <td>2.0829</td>
    <td>2.0876</td>
    <td>1.9855</td>
    <td>2.2103</td>
    <td>2.3641</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Pytorch</td>
    <td>2</td>
    <td>0.0986</td>
    <td>0.1025</td>
    <td>0.1027</td>
    <td>0.1036</td>
    <td>0.1051</td>
    <td>0.1126</td>
    <td>0.1213</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0277</td>
    <td>0.0292</td>
    <td>0.0288</td>
    <td>0.0287</td>
    <td>0.0288</td>
    <td>0.0288</td>
    <td>0.0339</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>3.5596</td>
    <td>3.5103</td>
    <td>3.5660</td>
    <td>3.6098</td>
    <td>3.6493</td>
    <td>3.9097</td>
    <td>3.5782</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Pytorch</td>
    <td>4</td>
    <td>0.1952</td>
    <td>0.2062</td>
    <td>0.2066</td>
    <td>0.2076</td>
    <td>0.2093</td>
    <td>0.2300</td>
    <td>0.2352</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0381</td>
    <td>0.0399</td>
    <td>0.0397</td>
    <td>0.0396</td>
    <td>0.0397</td>
    <td>0.0403</td>
    <td>0.0544</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>5.1234</td>
    <td>5.1679</td>
    <td>5.2040</td>
    <td>5.2424</td>
    <td>5.2720</td>
    <td>5.7072</td>
    <td>4.3235</td>
  </tr>
  <tr>
    <td>123B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0544</td>
    <td>0.0556</td>
    <td>0.0557</td>
    <td>0.0559</td>
    <td>0.0566</td>
    <td>0.0617</td>
    <td>0.0850</td>
  </tr>
  <tr>
    <th colspan="10">Decoding - (Input Length: 1 Token)</th>
  </tr>
  <tr>
    <td>7B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0229</td>
    <td>0.0449</td>
    <td>0.0889</td>
    <td>0.1768</td>
    <td>0.3529</td>
    <td>0.7047</td>
    <td>1.4133</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0152</td>
    <td>0.0299</td>
    <td>0.0596</td>
    <td>0.1193</td>
    <td>0.2374</td>
    <td>0.4748</td>
    <td>0.9463</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>1.5066</td>
    <td>1.5017</td>
    <td>1.4916</td>
    <td>1.4820</td>
    <td>1.4865</td>
    <td>1.4842</td>
    <td>1.4935</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0399</td>
    <td>0.0791</td>
    <td>0.1573</td>
    <td>0.3132</td>
    <td>0.6250</td>
    <td>1.2494</td>
    <td>2.5005</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0196</td>
    <td>0.0388</td>
    <td>0.0777</td>
    <td>0.1555</td>
    <td>0.3114</td>
    <td>0.6219</td>
    <td>1.2411</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>2.0357</td>
    <td>2.0387</td>
    <td>2.0245</td>
    <td>2.0141</td>
    <td>2.0071</td>
    <td>2.0090</td>
    <td>2.0147</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Pytorch</td>
    <td>2</td>
    <td>0.1009</td>
    <td>0.2009</td>
    <td>0.4011</td>
    <td>0.8014</td>
    <td>1.6022</td>
    <td>3.2050</td>
    <td>6.4134</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0300</td>
    <td>0.0603</td>
    <td>0.1203</td>
    <td>0.2408</td>
    <td>0.4815</td>
    <td>0.9632</td>
    <td>1.9272</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>3.3633</td>
    <td>3.3317</td>
    <td>3.3342</td>
    <td>3.3281</td>
    <td>3.3275</td>
    <td>3.3275</td>
    <td>3.3278</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Pytorch</td>
    <td>4</td>
    <td>0.1975</td>
    <td>0.3941</td>
    <td>0.7877</td>
    <td>1.5746</td>
    <td>3.1491</td>
    <td>6.2989</td>
    <td>12.6034</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0400</td>
    <td>0.0808</td>
    <td>0.1615</td>
    <td>0.3244</td>
    <td>0.6501</td>
    <td>1.3005</td>
    <td>2.6025</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>4.9375</td>
    <td>4.8775</td>
    <td>4.8774</td>
    <td>4.8539</td>
    <td>4.8440</td>
    <td>4.8434</td>
    <td>4.8428</td>
  </tr>
  <tr>
    <td>123B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0566</td>
    <td>0.1126</td>
    <td>0.2246</td>
    <td>0.4488</td>
    <td>0.8976</td>
    <td>1.7936</td>
    <td>3.5924</td>
  </tr>
</table>

Table 9: End-to-end inference time (in seconds) on NVIDIA L40s GPUs, comparing Trirun kernels to PyTorch FP16 for varying sequence lengths, showing the speedup of Trirun relative to PyTorch FP16.
<table>
  <tr>
    <th>Size</th>
    <th>Kernel</th>
    <th>#GPU</th>
    <th>1</th>
    <th>2</th>
    <th>4</th>
    <th>8</th>
    <th>16</th>
    <th>32</th>
    <th>64</th>
  </tr>
  <tr>
    <th colspan="10">Encoding</th>
  </tr>
  <tr>
    <td>7B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0166</td>
    <td>0.0174</td>
    <td>0.0175</td>
    <td>0.0175</td>
    <td>0.0179</td>
    <td>0.0185</td>
    <td>0.0198</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0085</td>
    <td>0.0092</td>
    <td>0.0093</td>
    <td>0.0093</td>
    <td>0.0093</td>
    <td>0.0092</td>
    <td>0.0117</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>1.9484</td>
    <td>1.8871</td>
    <td>1.8869</td>
    <td>1.8890</td>
    <td>1.9288</td>
    <td>2.0141</td>
    <td>1.6963</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0316</td>
    <td>0.0325</td>
    <td>0.0326</td>
    <td>0.0330</td>
    <td>0.0336</td>
    <td>0.0364</td>
    <td>0.0380</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0106</td>
    <td>0.0115</td>
    <td>0.0113</td>
    <td>0.0112</td>
    <td>0.0112</td>
    <td>0.0113</td>
    <td>0.0177</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>2.9774</td>
    <td>2.8384</td>
    <td>2.8953</td>
    <td>2.9307</td>
    <td>2.9911</td>
    <td>3.2055</td>
    <td>2.1415</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Pytorch</td>
    <td>2</td>
    <td>0.0794</td>
    <td>0.0809</td>
    <td>0.0815</td>
    <td>0.0823</td>
    <td>0.0840</td>
    <td>0.0899</td>
    <td>0.0934</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0171</td>
    <td>0.0180</td>
    <td>0.0178</td>
    <td>0.0178</td>
    <td>0.0187</td>
    <td>0.0258</td>
    <td>0.0408</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>4.6470</td>
    <td>4.5056</td>
    <td>4.5781</td>
    <td>4.6294</td>
    <td>4.4869</td>
    <td>3.4866</td>
    <td>2.2901</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Pytorch</td>
    <td>4</td>
    <td>0.1547</td>
    <td>0.1598</td>
    <td>0.1605</td>
    <td>0.1615</td>
    <td>0.1633</td>
    <td>0.1667</td>
    <td>0.1825</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0285</td>
    <td>0.0293</td>
    <td>0.0294</td>
    <td>0.0295</td>
    <td>0.0300</td>
    <td>0.0418</td>
    <td>0.0712</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>5.4276</td>
    <td>5.4595</td>
    <td>5.4535</td>
    <td>5.4685</td>
    <td>5.4451</td>
    <td>3.9857</td>
    <td>2.5645</td>
  </tr>
  <tr>
    <th colspan="10">Decoding - (Input Length: 1 Token)</th>
  </tr>
  <tr>
    <td>7B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0175</td>
    <td>0.0346</td>
    <td>0.0688</td>
    <td>0.1371</td>
    <td>0.2744</td>
    <td>0.5477</td>
    <td>1.0979</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0097</td>
    <td>0.0193</td>
    <td>0.0390</td>
    <td>0.0773</td>
    <td>0.1546</td>
    <td>0.3095</td>
    <td>0.6196</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>1.7969</td>
    <td>1.7952</td>
    <td>1.7629</td>
    <td>1.7734</td>
    <td>1.7749</td>
    <td>1.7694</td>
    <td>1.7720</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0325</td>
    <td>0.0645</td>
    <td>0.1286</td>
    <td>0.2566</td>
    <td>0.5122</td>
    <td>1.0240</td>
    <td>2.0510</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0116</td>
    <td>0.0232</td>
    <td>0.0462</td>
    <td>0.0925</td>
    <td>0.1850</td>
    <td>0.3702</td>
    <td>0.7381</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>2.8040</td>
    <td>2.7841</td>
    <td>2.7841</td>
    <td>2.7751</td>
    <td>2.7683</td>
    <td>2.7658</td>
    <td>2.7787</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Pytorch</td>
    <td>2</td>
    <td>0.0805</td>
    <td>0.1607</td>
    <td>0.3211</td>
    <td>0.6422</td>
    <td>1.2843</td>
    <td>2.5693</td>
    <td>5.1450</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0185</td>
    <td>0.0373</td>
    <td>0.0748</td>
    <td>0.1499</td>
    <td>0.3039</td>
    <td>0.6014</td>
    <td>1.2210</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>4.3532</td>
    <td>4.3132</td>
    <td>4.2912</td>
    <td>4.2828</td>
    <td>4.2254</td>
    <td>4.2721</td>
    <td>4.2136</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Pytorch</td>
    <td>4</td>
    <td>0.1559</td>
    <td>0.3115</td>
    <td>0.6230</td>
    <td>1.2461</td>
    <td>2.4920</td>
    <td>4.9844</td>
    <td>9.9766</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0297</td>
    <td>0.0590</td>
    <td>0.1180</td>
    <td>0.2356</td>
    <td>0.4725</td>
    <td>0.9471</td>
    <td>1.8942</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>5.2519</td>
    <td>5.2760</td>
    <td>5.2804</td>
    <td>5.2893</td>
    <td>5.2741</td>
    <td>5.2626</td>
    <td>5.2670</td>
  </tr>
</table>

Table 10: End-to-end inference time (in seconds) on NVIDIA 4090 GPUs, comparing Trirun kernels to PyTorch FP16 for varying sequence lengths, showing the speedup of Trirun relative to PyTorch FP16.
<table>
  <tr>
    <th rowspan="2">Size</th>
    <th rowspan="2">Kernel</th>
    <th rowspan="2">#GPU</th>
    <th colspan="8">Encoding</th>
  </tr>
  <tr>
    <th>1</th>
    <th>2</th>
    <th>4</th>
    <th>8</th>
    <th>16</th>
    <th>32</th>
    <th>64</th>
  </tr>
  <tr>
    <td>7B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0210</td>
    <td>0.0219</td>
    <td>0.0220</td>
    <td>0.0221</td>
    <td>0.0225</td>
    <td>0.0244</td>
    <td>0.0258</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0129</td>
    <td>0.0139</td>
    <td>0.0139</td>
    <td>0.0142</td>
    <td>0.0139</td>
    <td>0.0138</td>
    <td>0.0138</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>1.6239</td>
    <td>1.5749</td>
    <td>1.5810</td>
    <td>1.5552</td>
    <td>1.6193</td>
    <td>1.7617</td>
    <td>1.8681</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0381</td>
    <td>0.0402</td>
    <td>0.0403</td>
    <td>0.0407</td>
    <td>0.0413</td>
    <td>0.0458</td>
    <td>0.0481</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0157</td>
    <td>0.0170</td>
    <td>0.0173</td>
    <td>0.0174</td>
    <td>0.0174</td>
    <td>0.0170</td>
    <td>0.0184</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>2.4219</td>
    <td>2.3633</td>
    <td>2.3306</td>
    <td>2.3368</td>
    <td>2.3765</td>
    <td>2.6850</td>
    <td>2.6097</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Pytorch</td>
    <td>2</td>
    <td>0.0990</td>
    <td>0.1028</td>
    <td>0.1031</td>
    <td>0.1038</td>
    <td>0.1054</td>
    <td>0.1180</td>
    <td>0.1226</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0268</td>
    <td>0.0281</td>
    <td>0.0280</td>
    <td>0.0278</td>
    <td>0.0281</td>
    <td>0.0279</td>
    <td>0.0427</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>3.6857</td>
    <td>3.6630</td>
    <td>3.6814</td>
    <td>3.7370</td>
    <td>3.7462</td>
    <td>4.2331</td>
    <td>2.8736</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Pytorch</td>
    <td>4</td>
    <td>0.1958</td>
    <td>0.2068</td>
    <td>0.2071</td>
    <td>0.2080</td>
    <td>0.2097</td>
    <td>0.2167</td>
    <td>0.2358</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0358</td>
    <td>0.0371</td>
    <td>0.0367</td>
    <td>0.0368</td>
    <td>0.0372</td>
    <td>0.0439</td>
    <td>0.0723</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>5.4757</td>
    <td>5.5771</td>
    <td>5.6356</td>
    <td>5.6579</td>
    <td>5.6366</td>
    <td>4.9367</td>
    <td>3.2611</td>
  </tr>
  <tr>
    <th colspan="10">Decoding - (Input Length: 1 Token)</th>
  </tr>
  <tr>
    <td>7B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0227</td>
    <td>0.0448</td>
    <td>0.0889</td>
    <td>0.1769</td>
    <td>0.3531</td>
    <td>0.7043</td>
    <td>1.4104</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0144</td>
    <td>0.0285</td>
    <td>0.0571</td>
    <td>0.1144</td>
    <td>0.2284</td>
    <td>0.4579</td>
    <td>0.9093</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>1.5759</td>
    <td>1.5705</td>
    <td>1.5579</td>
    <td>1.5464</td>
    <td>1.5461</td>
    <td>1.5382</td>
    <td>1.5511</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0399</td>
    <td>0.0792</td>
    <td>0.1578</td>
    <td>0.3145</td>
    <td>0.6273</td>
    <td>1.2536</td>
    <td>2.5101</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0175</td>
    <td>0.0356</td>
    <td>0.0711</td>
    <td>0.1407</td>
    <td>0.2779</td>
    <td>0.5560</td>
    <td>1.1106</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>2.2862</td>
    <td>2.2228</td>
    <td>2.2203</td>
    <td>2.2343</td>
    <td>2.2569</td>
    <td>2.2546</td>
    <td>2.2600</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Pytorch</td>
    <td>2</td>
    <td>0.1013</td>
    <td>0.2020</td>
    <td>0.4031</td>
    <td>0.8054</td>
    <td>1.6102</td>
    <td>3.2202</td>
    <td>6.4492</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0291</td>
    <td>0.0588</td>
    <td>0.1173</td>
    <td>0.2355</td>
    <td>0.4711</td>
    <td>0.9484</td>
    <td>1.8970</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>3.4876</td>
    <td>3.4350</td>
    <td>3.4359</td>
    <td>3.4195</td>
    <td>3.4183</td>
    <td>3.3955</td>
    <td>3.3997</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Pytorch</td>
    <td>4</td>
    <td>0.1987</td>
    <td>0.3966</td>
    <td>0.7924</td>
    <td>1.5842</td>
    <td>3.1681</td>
    <td>6.3388</td>
    <td>12.6871</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0374</td>
    <td>0.0754</td>
    <td>0.1514</td>
    <td>0.3002</td>
    <td>0.6015</td>
    <td>1.1935</td>
    <td>2.3860</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>5.3117</td>
    <td>5.2600</td>
    <td>5.2330</td>
    <td>5.2774</td>
    <td>5.2666</td>
    <td>5.3110</td>
    <td>5.3172</td>
  </tr>
</table>

Table 11: End-to-end inference time (in seconds) on NVIDIA L40 GPUs, comparing Trirun kernels to PyTorch FP16 for varying sequence lengths, showing the speedup of Trirun relative to PyTorch FP16.
<table>
  <tr>
    <th rowspan="2">Size</th>
    <th rowspan="2">Kernel</th>
    <th rowspan="2">#GPU</th>
    <th colspan="8">Encoding</th>
  </tr>
  <tr>
    <th>1</th>
    <th>2</th>
    <th>4</th>
    <th>8</th>
    <th>16</th>
    <th>32</th>
    <th>64</th>
  </tr>
  <tr>
    <td>7B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0280</td>
    <td>0.0281</td>
    <td>0.0283</td>
    <td>0.0286</td>
    <td>0.0300</td>
    <td>0.0322</td>
    <td>0.0325</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0177</td>
    <td>0.0188</td>
    <td>0.0189</td>
    <td>0.0190</td>
    <td>0.0192</td>
    <td>0.0189</td>
    <td>0.0190</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>1.5840</td>
    <td>1.4939</td>
    <td>1.4952</td>
    <td>1.5058</td>
    <td>1.5642</td>
    <td>1.7012</td>
    <td>1.7150</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0509</td>
    <td>0.0522</td>
    <td>0.0524</td>
    <td>0.0528</td>
    <td>0.0535</td>
    <td>0.0556</td>
    <td>0.0573</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0280</td>
    <td>0.0301</td>
    <td>0.0302</td>
    <td>0.0302</td>
    <td>0.0302</td>
    <td>0.0301</td>
    <td>0.0291</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>1.8139</td>
    <td>1.7323</td>
    <td>1.7377</td>
    <td>1.7470</td>
    <td>1.7745</td>
    <td>1.8431</td>
    <td>1.9732</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Pytorch</td>
    <td>2</td>
    <td>0.1277</td>
    <td>0.1331</td>
    <td>0.1339</td>
    <td>0.1335</td>
    <td>0.1357</td>
    <td>0.1433</td>
    <td>0.1472</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0438</td>
    <td>0.0452</td>
    <td>0.0462</td>
    <td>0.0462</td>
    <td>0.0460</td>
    <td>0.0453</td>
    <td>0.0514</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>2.9170</td>
    <td>2.9425</td>
    <td>2.8964</td>
    <td>2.8877</td>
    <td>2.9513</td>
    <td>3.1624</td>
    <td>2.8633</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Pytorch</td>
    <td>4</td>
    <td>0.2580</td>
    <td>0.2673</td>
    <td>0.2685</td>
    <td>0.2702</td>
    <td>0.2752</td>
    <td>0.2845</td>
    <td>0.3494</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0765</td>
    <td>0.0792</td>
    <td>0.0784</td>
    <td>0.0786</td>
    <td>0.0790</td>
    <td>0.0789</td>
    <td>0.0898</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>3.3734</td>
    <td>3.3747</td>
    <td>3.4234</td>
    <td>3.4393</td>
    <td>3.4826</td>
    <td>3.6055</td>
    <td>3.8922</td>
  </tr>
  <tr>
    <th colspan="10">Decoding - (Input Length: 1 Token)</th>
  </tr>
  <tr>
    <td>7B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0299</td>
    <td>0.0590</td>
    <td>0.1171</td>
    <td>0.2334</td>
    <td>0.4662</td>
    <td>0.9318</td>
    <td>1.8670</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0195</td>
    <td>0.0391</td>
    <td>0.0773</td>
    <td>0.1587</td>
    <td>0.3101</td>
    <td>0.6184</td>
    <td>1.2324</td>
  </tr>
  <tr>
    <td>7B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>1.5341</td>
    <td>1.5073</td>
    <td>1.5145</td>
    <td>1.4709</td>
    <td>1.5033</td>
    <td>1.5067</td>
    <td>1.5150</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Pytorch</td>
    <td>1</td>
    <td>0.0527</td>
    <td>0.1046</td>
    <td>0.2085</td>
    <td>0.4159</td>
    <td>0.8314</td>
    <td>1.6628</td>
    <td>3.3312</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0253</td>
    <td>0.0506</td>
    <td>0.1013</td>
    <td>0.2029</td>
    <td>0.4051</td>
    <td>0.8122</td>
    <td>1.6293</td>
  </tr>
  <tr>
    <td>13B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>2.0871</td>
    <td>2.0666</td>
    <td>2.0581</td>
    <td>2.0503</td>
    <td>2.0526</td>
    <td>2.0474</td>
    <td>2.0446</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Pytorch</td>
    <td>2</td>
    <td>0.1302</td>
    <td>0.2594</td>
    <td>0.5178</td>
    <td>1.0353</td>
    <td>2.0716</td>
    <td>4.1386</td>
    <td>8.2878</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0400</td>
    <td>0.0818</td>
    <td>0.1631</td>
    <td>0.3278</td>
    <td>0.6513</td>
    <td>1.3201</td>
    <td>2.6081</td>
  </tr>
  <tr>
    <td>34B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>3.2521</td>
    <td>3.1735</td>
    <td>3.1737</td>
    <td>3.1582</td>
    <td>3.1804</td>
    <td>3.1351</td>
    <td>3.1777</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Pytorch</td>
    <td>4</td>
    <td>0.2608</td>
    <td>0.5204</td>
    <td>1.0404</td>
    <td>2.0823</td>
    <td>4.1608</td>
    <td>8.3258</td>
    <td>16.6702</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Trirun</td>
    <td>1</td>
    <td>0.0790</td>
    <td>0.1219</td>
    <td>0.2229</td>
    <td>0.4486</td>
    <td>0.9942</td>
    <td>1.8002</td>
    <td>3.4719</td>
  </tr>
  <tr>
    <td>70B</td>
    <td>Speedup</td>
    <td>-</td>
    <td>3.2997</td>
    <td>4.2700</td>
    <td>4.6675</td>
    <td>4.6413</td>
    <td>4.1851</td>
    <td>4.6250</td>
    <td>4.8015</td>
  </tr>
</table>

Table 12: End-to-end inference time (in seconds) on NVIDIA A40 GPUs, comparing Trirun kernels to PyTorch FP16 for varying sequence lengths, showing the speedup of Trirun relative to PyTorch FP16.

Reviewer paper summaries: This paper focuses on the scaling analysis and implementation optimization of Ternary Language Models (TriLMs). 
1. It explores the scaling laws of TriLMs, demonstrating that increasing training tokens reduces validation loss more effectively than increasing parameter count. 
2. Empirically, a 3B TriLM achieves competitive performance with FloatLMs (llama-1-7B) within the same compute budget (trained on 1.2 trillion tokens).
3. The authors propose efficient 1.6-bit and 2-bit packing schemes for ternary weights, providing a theoretical analysis and benchmarking on a CPU.
4. They introduce GPU kernels based on the 2-bit packing schemes, termed TriRun, which achieve significant end-to-end speedup compared to float16 implementations.
Summary of strengths: This paper is quite solid with impressive improvement, extensive experiments, dedicated implementations, and good writing.
1. Impressive improvement: The efficiency enhancement is significant.
2. Extensive experiments: The findings of scaling laws of TriLMs would contribute more interest in further scaling TriLMs.
3. Dedicated implementations: The dedicated designs and implementations would benefit the future development and deployment of TriLMs.
4. Good writing: this paper is written with clear logic, and easy to follow, despite so many statements to be claimed.
Summary of weaknesses: 1. Performance concern: This is the biggest concern of this paper. Specifically, 
  - Even if the authors show the similar scaling laws of TriLMs and FloatLMs in Figure 2 and 8, the validation loss of TriLMs is higher than that of FloatLMs. This indicates the potential performance loss of TriLMs. 
  - Thank the authors for showing the performance comaprison of TriLMs (1B, 2B, 3B) and  LLaMA-1 7B in Figure 3. Personally, this comparison is not that fair, due to the different training dataset and settings. I suggest the authors to train a 3B FloatLM with the same setting, and compare its performance (also its quantized 4-bit performance) with those of TriLMs.
  - For Table 4, as stated in [1], multi-step reasoning tasks can better reflect the performance differences of existing quantization methods. I recommend the authors to add one coding task (like, MBPP, HumanEval) into their benchmarks for more comprehensive comparison. Besides, adding the average performance can help readers compare different methods more easily. 
  - The unidentified performance comparison weakens the significance of the following optimizations.
2. Runtime memory: In deployment, runtime memory is more meaningful than model memory. Could you please add the runtime memoru comparion in Figure 4?
3. Confused implementation of Q4: I'm a bit confused about the concrete implementation of Q4. Is is W4A4 (QuaRot, Atom), W4A8 (QServe) or W4A16 (AWQ, GPTQ, etc)? 

Note that my primary concern lies in the performance gap between TriLMs and FloatLMs. If this issue can be addressed, I could significantly raise my ratings.

[1] QSpec: Speculative Decoding with Complementary Quantization Schemes
Comments, suggestions and typos: Please refer to the Summary Of Weaknesses.
Overall assessment scores: 3.5