MxMoE: Mixed-precision Quantization for MoE with Accuracy and Performance Co-Design

We appreciate your thoughtful feedback. We would like to provide further clarification on several aspects of our work that you have mentioned:

1. Novelty of the Approach and Similarities with Prior Work

• Linear-Block-Level Quantization: Our approach is driven by empirical observations rather than a straightforward extension of prior work. Figure 1(a) highlights significant sensitivity differences between linear blocks, and Table 3 confirms the effectiveness of our approach.
• Objective: Balancing Model Accuracy and Hardware Efficiency: Unlike prior work that solely optimizes model accuracy, MxMoE jointly considers both accuracy and system efficiency. Hardware-friendly quantization strategies improve efficiency but often degrade accuracy (Figure 6). Our objective function is designed to incorporate both quantization perturbation and the runtime cost of a mixed-precision scheme. And accurately modeling runtime cost is nontrivial, requiring deep analysis of hardware characteristics and parallel algorithm design, as detailed in Section 4.2.
• Scope: Weight-Only vs. Weight + Activation Quantization: MC-MoE focuses on extremely low-bit weight-only quantization, whereas MxMoE optimally allocates bit-widths for both activation and weight parameters. This approach enables MxMoE to leverage low-precision arithmetic units on GPUs more effectively (Figure 5).
• Performance: Theoretical vs. Real Speedup: Previous methods relied on existing kernels like HQQ, achieving only theoretical speedups without real-world reductions in wall-clock time (Figure 2). In contrast, MxMoE is designed from the ground up for hardware efficiency, leading to actual runtime improvements.
• Automation: End-to-End Mixed-Precision pipeline: MxMoE is the first framework to fully automate the generation of mixed-precision schemes and fused operators tailored to the generated scheme which is not achieved in any prior works.

2. MxMoE Is Engineering Optimization

• MxMoE is not simply providing one kernel but rather offering an automated method for the mixed-precision allocation and generating customized fused operators. The design and optimization of mixed-precision GEMM have been an ongoing area of research[3][4], which continue to push the boundaries of quantization research. MxMoE addresses the more challenging problem of mixed-precision GroupGEMM, which is a much more complex problem that goes beyond simple parallelization of different precision tasks in thread-blocks or streams.

• In fact, we initially implemented MxMoE with multi-stream, but unfortunately we found it could not fully exploit hardware. We show the comparison(RTX 4090, Qwen1.5, W4A16 and W8A8 mixed MoE-Block):

  #Tokens 1024	TFLOPS
  FP16	70.117
  W4A16	119.32
  W8A8	140.95
  MxMoE	163.20
  multi-stream	142.19

  The results demonstrated that our approach achieves superior efficiency.

• Assigning task to CTAs is fundamentally the minimum makespan problem, an NP-hard challenge. "parallelize different experts with different quantization schemes within a single kernel by using separate thread blocks" won't work as it suffers from workload imbalance. To address this, we design a tile scheduler that distributes workload evenly among all CTAs(Line 295).

3. Experimental Results

• At 3.5 bit, quantization perturbation is relatively limited. Both GPTQ and MxMoE perform well at this setting. As the bit-width decreases, the accuracy advantage of MxMoE becomes more evident. This extreme low-bit compression is highly relevant for resource-constrained model deployment, as explored in several works such as [1][2].
• MxMoE performs bit-width allocation for both activation and weight. As demonstrated in Table 1 and Figure 5, MxMoE shows a significant accuracy advantage over existing methods in weight-activation quantization, pushing the accuracy-efficiency Pareto frontier to new heights.

In conclusion, we believe MxMoE introduces novel contributions in both algorithmic design and system-level optimization. We sincerely appreciate the reviewer’s feedback and the opportunity to clarify our contributions. We hope our responses have effectively addressed the concerns raised. Given these clarifications, we hope the reviewer may reconsider the overall evaluation of our work. We would be happy to further discuss any remaining questions or concerns.
[1] Ma, Shuming, et al. "The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits." CoRR (2024).
[2] Yuan, Zhihang, et al. "PB-LLM: Partially Binarized Large Language Models." The Twelfth International Conference on Learning Representations.
[3] Wang, Lei, et al. "Ladder: Enabling Efficient {Low-Precision} Deep Learning Computing through Hardware-aware Tensor Transformation." 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24).
[4] Lin, Yujun, et al. "Qserve: W4a8kv4 quantization and system co-design for efficient llm serving." arXiv preprint arXiv:2405.04532 (2024).

We appreciate the reviewer’s engagement.

We sincerely appreciate your recognition of our work. Your insightful questions and suggestions have been instrumental in enhancing the quality of our manuscript. Below, we address your concerns in detail:

1. Rationale for the W5A5 Configuration in Table 1

• Why does the "strange" W5A5 setting arise?

  This configuration is intentional. During bitwidth allocation, we set an average 5-bit activation and weight budget for MxMoE. On our experimental platform (RTX 4090), Tensor Cores support 4-bit and 8-bit operations, allowing both W4A4 and W8A8 configurations to benefit from low-precision Tensor Core acceleration. The additional 1-bit budget provides MxMoE with greater flexibility in bitwidth allocation, meaning that certain activations are assigned 8-bit precision while others remain at 4-bit.

  Although the average bitwidth increases by only 1 bit, MxMoE automatically identifies activations that are more sensitive to quantization through sensitivity analysis and assigns them a higher bitwidth to mitigate accuracy degradation. This exemplifies the advantage of mixed-precision quantization over uniform quantization.

• Does the accuracy gain stem solely from the additional 1-bit?

  Yes. The improved accuracy of W5A5 over W4A4 is primarily due to extreme outliers in the input activations of the down_proj layer—a phenomenon well-documented in prior research (e.g., [1, 2]). Unlike standard outlier features, these activations are particularly difficult to quantize effectively at 4 bits, leading to precision loss. MxMoE mitigates this issue by allocating higher bitwidths to such activations, thereby preserving model accuracy.

  Our findings align with prior work [2], which demonstrates that suppressing extreme outliers substantially improves post-quantization model performance. We will clarify this rationale in the revised manuscript and provide a case study on bitwidth allocation.

2. Throughput Evaluation with Varied Length Distributions

   We agree that evaluating computational throughput across diverse input lengths is essential for real-world applicability. In response to your suggestion, we have expanded our throughput experiments. Due to time constraints, we conducted experiments on DeepSeekV2-Lite using the Humaneval-X dataset.

   In this study, we employed the exact same MxMoE(W5A5) configuration as in Table 1 and Figure 5. To assess performance across varied sequence lengths, we set the batch size to 1—ensuring that input lengths is varied, unlike in the compute-bound scenarios in Figure 5. We then processed the entire dataset, computing the average computational throughput (total FLOPs divided by execution time) and comparing it to FP16 performance. The length distribution of the dataset can be found at https://huggingface.co/datasets/THUDM/humaneval-x.

   Dataset	FP16 TOPS	MxMoE(W5A5) TOPS	Speedup
   humaneval-x	38.19	107.36	2.8

   As shown in the results, although the speedup is slightly lower than in the fully compute-bound scenario, MxMoE(W5A5) maintains a strong acceleration ratio even on datasets with varied-length short sequences. This highlights the robustness of the mixed-precision approach.

   We will expand this experiment in the revised manuscript and add more models and dataset tests. Thank you very much for your suggestions!

3. Typos and Figure Refinements

   We sincerely appreciate your careful review and apologize for the typographical errors. These will be corrected in the revised manuscript. Additionally, we will refine the caption of Figure 3 to explicitly define terms such as g128 (quantization group size of 128) and g-1 (per-channel/token quantization) to enhance clarity. Thanks again for your advice!

Once again, we deeply appreciate your valuable feedback. Your suggestions have been instrumental in improving the quality of our manuscript. If you have any further suggestions or concerns, we would be delighted to discuss them!

[1] Sun, Mingjie, et al. Massive Activations in Large Language Models. First Conference on Language Modeling.
[2] Lin, Haokun, et al. Duquant: Distributing Outliers via Dual Transformation Makes Stronger Quantized LLMs. NeurIPS 2024.

We sincerely appreciate the reviewer’s constructive feedback. Below, we address the two key concerns raised:

1. Exploration of Trade-off Parameter $r$ for Different Models and Use Cases

   Our ablation study (Figure 6) explores the impact of the hyper-parameter $r$ on model accuracy and hardware efficiency. Our analysis demonstrates that increasing $r$ prioritizes accuracy (e.g., lower perplexity) at the expense of efficiency (e.g., lower throughput).

   MxMoE can automatically adjust $r$ according to use case: there are typically some constraints such as memory budget, precision budget, and latency budget in model quantization. MxMoE satisfy these constraints by automatically adjusting $r$. Additionally, users may manually tune $r$ to balance hardware efficiency and model accuracy, selecting a suitable value based on their specific requirements. Basically, if a user prioritizes higher accuracy, $r$ should be set closer to 1, whereas if better speed is desired, $r$ should be closer to 0.

   While Figure 6 presents results for a single model, we emphasize that the observed trend is consistent across all evaluated architectures. For simplicity and reproducibility, we adopt $r=0.75$ as the default value in all experiments, except for low-bit weight-only quantization, where we set $r=1.0$ (Lines 363, 373, 378). This exception aligns with edge deployment scenarios, where memory constraints dominate, making it crucial to maximize compression while preserving accuracy.

2. Validity Analysis for Method Components

   Micro-Kernel Specialization

   MxMoE mitigates the combinatorial explosion of mixed-precision configurations by automating kernel generation (Line 289). To illustrate the effectiveness of our approach, we compare our strategy with two alternative solutions:

   • Developing a universal kernel to handle all precision combinations: This approach would compromise kernel performance. We provide a breakdown demonstrating the limitations of this method relative to our micro-kernel specialization. Specifically, the kernel for W4A4-per-channel could theoretically share the same software pipeline with W4A4-group128, but enforcing universality significantly degrades performance (test shape $[8192, 8192, 8192]$):

     Kernel Type	W4A4_per-channel TOPS	W4A4_group128 TOPS
     W4A4_per-channel(Specialized)	1070.5303	N/A
     W4A4_group128(Specialized)	N/A	667.3349
     Unified Kernel	929.1997	412.0268

     Unifying the two pipelines requires introducing runtime condition checks, which hinder loop unrolling in the MAC-loop. Moreover, to support group-size=128, the per-channel kernel’s tile-size selection is constrained, making configurations such as tile_k=256 infeasible.

   • Developing separate kernels for each configuration:

     While handcrafted kernels could match performance, they require substantial engineering effort. If a given hardware platform supports five quantization candidates (e.g., w2a6, w4a16, w8a8, w4a4, w4a4 with group-size 128), implementing individual kernels for all possible configurations would require $5!=120$ kernels! In contrast, our micro-kernel specialization approach requires implementing only 5 configurable micro-kernels, which are automatically combined by the kernel generator to form optimized fused operators.

   Resource Configuration

   To resolve resource mismatches caused by heterogeneous micro-kernels (e.g., varying warp counts per thread block, as shown in Figure 4), we employ the following strategies:

   • Max-resource allocation: Ensuring all parallel units reach the resource ceiling for correctness.

   • Slice-K optimization: This optimization improves shared memory utilization and achieves speedup compared to a baseline without carefully config, we provide a case study: The W4A16-per-channel quantization with/without resource configuration optimization(test shape $[16, 8192, 8192]$).

     Method	Best Config	W4A16 TOPS
     $\mathrm{MxMoE}$	Tile: [16, 64, 128]; Warp: [1,2,2]	74.8983
     $\mathrm{- Resource Configuration}$	Tile: [16, 128, 128]; Warp: [1,4,1]	52.4288

     The result shows that with resource configuration optimization, MxMoE automatically discovered a better tile configuration (assign 2 warps along k-dimension, which is not a common configuration in previous kernel such as Marlin[1] e.t.c.), which brings 42% speedup.

We will enhance the revised manuscript with additional ablation studies and quantitative comparisons to reinforce these claims. We sincerely appreciate the reviewer’s valuable suggestions, which we believe further strengthen the technical rigor and reproducibility of our work.

[1] Frantar, Elias, et al. "Marlin: Mixed-precision auto-regressive parallel inference on large language models." Proceedings of the 30th ACM SIGPLAN 2025.