Support new format AWQ/GPTQ MoE models (#452)

Author: guoqingbao

Cargo.toml

@@ -21,8 +21,8 @@ anyhow = "1.0.75"
 rand = "0.9.0"
 rayon="1.10.0"
 hyper = { version = "0.14", features = ["full"] }
-candle-core = { git = "https://github.com/guoqingbao/candle.git", version = "0.8.3", rev = "cafd231" }
-candle-nn = { git = "https://github.com/guoqingbao/candle.git", version = "0.8.3", rev = "cafd231" }
+candle-core = { git = "https://github.com/guoqingbao/candle.git", version = "0.8.3", rev = "df058d8" }
+candle-nn = { git = "https://github.com/guoqingbao/candle.git", version = "0.8.3", rev = "df058d8" }
 dyn-fmt = "0.4.0"
 safetensors = "0.8"
 serde = { version = "1.0.190", features = ["serde_derive"] }
@@ -45,7 +45,7 @@ dirs = "5.0.1"
 minijinja = { version = "2.10.2", features = ["builtins", "json"] }
 minijinja-contrib = { version = "2.10.2", features = ["pycompat"] }
 thiserror = "1.0.58"
-attention-rs = { git = "https://github.com/guoqingbao/attention.rs.git", version="0.6.5", rev = "a97f519" }
+attention-rs = { git = "https://github.com/guoqingbao/attention.rs.git", version="0.6.6", rev = "b35d437" }
 metal = { version = "0.27.0", features = ["mps"], optional = true }
 lazy_static = {version = "1.4.0"}
 interprocess = "2.2.2"

src/openai/models/gemma4.rs

@@ -7,7 +7,7 @@ use crate::openai::distributed::{
     embedding, Comm, ReplicatedLinear, VarBuilder, VocabParallelLinear,
 };
 use crate::openai::models::layers::moe::{
-    FusedMoe, FusedMoeFp8, FusedMoeISQ, FusedMoeMxfp4, FusedMoeNvfp4,
+    FusedMoe, FusedMoeFp8, FusedMoeISQ, FusedMoeMxfp4, FusedMoeNvfp4, FusedMoeWNA16,
 };
 use crate::openai::models::layers::others::{rms_norm, NormX};
 use crate::openai::models::mask::get_attention_causal_mask;
@@ -201,6 +201,7 @@ enum Gemma4MoE {
     FusedMoeFp8(FusedMoeFp8),
     FusedMoeMxfp4(FusedMoeMxfp4),
     FusedMoeNvfp4(FusedMoeNvfp4),
+    FusedMoeWNA16(FusedMoeWNA16),
 }
 
 impl Gemma4MoE {
@@ -211,6 +212,7 @@ impl Gemma4MoE {
             Self::FusedMoeFp8(m) => m.forward(xs, is_prefill),
             Self::FusedMoeMxfp4(m) => m.forward(xs, is_prefill),
             Self::FusedMoeNvfp4(m) => m.forward(xs, is_prefill),
+            Self::FusedMoeWNA16(m) => m.forward(xs, is_prefill),
         }
     }
 
@@ -231,6 +233,9 @@ impl Gemma4MoE {
             Self::FusedMoeNvfp4(m) => {
                 m.forward_with_routing(xs, topk_weights, topk_ids, is_prefill)
             }
+            Self::FusedMoeWNA16(m) => {
+                m.forward_with_routing(xs, topk_weights, topk_ids, is_prefill)
+            }
         }
     }
 }
@@ -333,6 +338,16 @@ impl Gemma4DecoderLayer {
                         comm.clone(),
                         dtype,
                     )?)
+                } else if quant_cfg.is_compressed_tensors {
+                    Gemma4MoE::FusedMoeWNA16(FusedMoeWNA16::new_with_gate(
+                        cfg,
+                        vb.pp("router").pp("proj"),
+                        vb.pp("experts"),
+                        &vb,
+                        comm.clone(),
+                        dtype,
+                        quant_cfg,
+                    )?)
                 } else {
                     candle::bail!(
                         "Unsupported quantization for Gemma4 MoE: {}",

src/openai/models/layers/deltanet.rs

@@ -72,6 +72,9 @@ impl GatedDeltaNet {
                     vb.contains_tensor("weight_scale_2") || vb.contains_tensor("input_scale");
                 has_mlx || (has_packed && has_scale) || (has_modelopt && has_scale)
             }
+            "compressed-tensors" => {
+                vb.contains_tensor("weight_packed") && vb.contains_tensor("weight_scale")
+            }
             _ => true,
         }
     }

src/openai/models/layers/mla_attention.rs

@@ -546,6 +546,10 @@ impl MlaAttention {
                     block_tables,
                     context_lens,
                     self.sm_scale,
+                    input_metadata
+                        .max_context_len
+                        .max(input_metadata.max_seqlen_k)
+                        .max(1),
                 )?;
                 return self.project_mla_output(&attn_out, seq_len, xs.dtype());
             }