Fixes

Author: lukekim

candle-examples/examples/nvembed_v2/main.rs

@@ -148,7 +148,7 @@ fn encode(
             device,
         )?;
         let b = attention_mask.dims()[0];
-        attention_mask.slice_assign(&[..b, ..instruction_lens], &zeros)?
+        attention_mask.slice_assign(&[&(..b), &(..instruction_lens)], &zeros)?
     } else {
         attention_mask.clone()
     };

candle-nn/src/layer_norm.rs

@@ -361,3 +361,6 @@ pub fn rms_norm(
 ) -> Result<RmsNorm<RmsNormNonQuantized>> {
     rms_norm_non_quant(size, eps, vb)
 }
+
+/// Type alias for backward compatibility - non-quantized RmsNorm.
+pub type RmsNormDefault = RmsNorm<RmsNormNonQuantized>;

candle-nn/src/lib.rs

@@ -53,7 +53,7 @@ pub use group_norm::{group_norm, GroupNorm};
 pub use init::Init;
 pub use layer_norm::{
     layer_norm, layer_norm_no_bias, rms_norm, rms_norm_non_quant, rms_norm_quant, LayerNorm,
-    LayerNormConfig, RmsNorm,
+    LayerNormConfig, RmsNorm, RmsNormDefault, RmsNormNonQuantized, RmsNormQuantized,
 };
 pub use linear::{linear, linear_b, linear_no_bias, Linear};
 pub use ops::Dropout;

candle-transformers/src/models/csm.rs

@@ -8,7 +8,7 @@
 ///
 use crate::generation::LogitsProcessor;
 use candle::{DType, Device, IndexOp, Module, Result, Tensor, D};
-use candle_nn::{embedding, linear_b, Embedding, Linear, RmsNorm, VarBuilder};
+use candle_nn::{embedding, linear_b, Embedding, Linear, RmsNorm, RmsNormNonQuantized, VarBuilder};
 use std::sync::Arc;
 
 #[derive(serde::Deserialize, Debug, Clone, Copy, PartialEq, Eq)]
@@ -142,9 +142,9 @@ impl RotaryEmbedding {
         Ok((q_embed, k_embed))
     }
 }
-fn rms_norm(hidden_size: usize, eps: f64, vb: VarBuilder) -> Result<RmsNorm> {
+fn rms_norm(hidden_size: usize, eps: f64, vb: VarBuilder) -> Result<RmsNorm<RmsNormNonQuantized>> {
     let weight = vb.get((hidden_size,), "scale")?;
-    Ok(RmsNorm::new(weight, eps))
+    Ok(RmsNorm::<RmsNormNonQuantized>::new(weight, eps))
 }
 
 #[derive(Debug, Clone)]
@@ -274,8 +274,8 @@ impl Module for Mlp {
 
 #[derive(Debug, Clone)]
 struct Layer {
-    mlp_norm: RmsNorm,
-    sa_norm: RmsNorm,
+    mlp_norm: RmsNorm<RmsNormNonQuantized>,
+    sa_norm: RmsNorm<RmsNormNonQuantized>,
     attn: Attention,
     mlp: Mlp,
 }
@@ -317,7 +317,7 @@ impl Layer {
 #[derive(Debug, Clone)]
 pub struct LlamaModel {
     layers: Vec<Layer>,
-    norm: RmsNorm,
+    norm: RmsNorm<RmsNormNonQuantized>,
     device: Device,
     dtype: DType,
 }

candle-transformers/src/models/deepseek2.rs

@@ -6,7 +6,7 @@ use candle::{
     shape::Dim, CpuStorage, CustomOp1, DType, Device, Error, IndexOp, Layout, Result, Shape,
     Tensor, WithDType, D,
 };
-use candle_nn::{embedding, rms_norm, Activation, Embedding, Linear, Module, RmsNorm, VarBuilder};
+use candle_nn::{embedding, rms_norm, Activation, Embedding, Linear, Module, RmsNorm, RmsNormNonQuantized, VarBuilder};
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use serde::Deserialize;
 
@@ -520,7 +520,7 @@ impl DeepSeekV2Config {
 
 enum QProj {
     Plain(Linear),
-    Lora { a: Linear, norm: RmsNorm, b: Linear },
+    Lora { a: Linear, norm: RmsNorm<RmsNormNonQuantized>, b: Linear },
 }
 
 impl QProj {
@@ -535,7 +535,7 @@ impl QProj {
 struct Attention {
     q: QProj,
     kv_a_proj_with_mqa: Linear,
-    kv_a_layernorm: RmsNorm,
+    kv_a_layernorm: RmsNorm<RmsNormNonQuantized>,
     kv_b_proj: Linear,
     o_proj: Linear,
     rotary_emb: Arc<DeepSeekV2RotaryEmbedding>,
@@ -905,8 +905,8 @@ impl MoeOrMlp {
 }
 
 struct DecoderLayer {
-    input_layernorm: RmsNorm,
-    post_attention_layernorm: RmsNorm,
+    input_layernorm: RmsNorm<RmsNormNonQuantized>,
+    post_attention_layernorm: RmsNorm<RmsNormNonQuantized>,
     attn: Attention,
     moe_or_mlp: MoeOrMlp,
 }
@@ -976,7 +976,7 @@ impl DecoderLayer {
 pub struct DeepSeekV2 {
     lm_head: Linear,
     embed_tokens: Embedding,
-    norm: RmsNorm,
+    norm: RmsNorm<RmsNormNonQuantized>,
     layers: Vec<DecoderLayer>,
     dtype: DType,
     device: Device,

candle-transformers/src/models/mod.rs

@@ -91,6 +91,7 @@ pub mod quantized_mixformer;
 pub mod quantized_moondream;
 pub mod quantized_mpt;
 pub mod quantized_phi;
+pub mod quantized_phi3;
 pub mod quantized_qwen2;
 pub mod quantized_qwen3;
 pub mod quantized_recurrent_gemma;

candle-transformers/src/models/olmo2.rs

@@ -6,7 +6,7 @@
 //!
 //!
 use candle::{DType, Device, Module, Result, Tensor, D};
-use candle_nn::{linear_b, linear_no_bias, rms_norm, Activation, Linear, RmsNorm, VarBuilder};
+use candle_nn::{linear_b, linear_no_bias, rms_norm, Activation, Linear, RmsNorm, RmsNormNonQuantized, VarBuilder};
 use std::sync::Arc;
 
 #[derive(Debug, Clone, serde::Deserialize)]
@@ -106,8 +106,8 @@ struct Attention {
     k_proj: Linear,
     v_proj: Linear,
     o_proj: Linear,
-    q_norm: RmsNorm,
-    k_norm: RmsNorm,
+    q_norm: RmsNorm<RmsNormNonQuantized>,
+    k_norm: RmsNorm<RmsNormNonQuantized>,
     num_heads: usize,
     num_kv_heads: usize,
     num_kv_groups: usize,
@@ -217,8 +217,8 @@ impl Attention {
 struct DecoderLayer {
     self_attn: Attention,
     mlp: MLP,
-    post_attention_layernorm: RmsNorm,
-    post_feedforward_layernorm: RmsNorm,
+    post_attention_layernorm: RmsNorm<RmsNormNonQuantized>,
+    post_feedforward_layernorm: RmsNorm<RmsNormNonQuantized>,
 }
 
 impl DecoderLayer {
@@ -268,7 +268,7 @@ impl DecoderLayer {
 pub struct Model {
     embed_tokens: candle_nn::Embedding,
     layers: Vec<DecoderLayer>,
-    norm: RmsNorm,
+    norm: RmsNorm<RmsNormNonQuantized>,
     lm_head: Linear,
     device: Device,
     dtype: DType,

candle-transformers/src/models/quantized_phi3.rs

@@ -18,7 +18,7 @@ use std::collections::HashMap;
 use candle::quantized::gguf_file;
 use candle::quantized::QTensor;
 use candle::{DType, Device, IndexOp, Module, Result, Tensor, D};
-use candle_nn::{kv_cache::KvCache, Embedding, RmsNorm};
+use candle_nn::{kv_cache::KvCache, Embedding, RmsNorm, RmsNormNonQuantized};
 
 #[derive(Debug, Clone)]
 struct QLinear {
@@ -64,18 +64,18 @@ impl Module for Mlp {
     }
 }
 
-fn rms_norm(w: QTensor, eps: f64) -> Result<RmsNorm> {
+fn rms_norm(w: QTensor, eps: f64) -> Result<RmsNorm<RmsNormNonQuantized>> {
     let w = w.dequantize(&w.device())?;
-    let rms = RmsNorm::new(w, eps);
+    let rms = RmsNorm::<RmsNormNonQuantized>::new(w, eps);
     Ok(rms)
 }
 
 #[derive(Debug, Clone)]
 struct LayerWeights {
     attn_qkv: QLinear,
     attn_output: QLinear,
-    attn_norm: RmsNorm,
-    ffn_norm: RmsNorm,
+    attn_norm: RmsNorm<RmsNormNonQuantized>,
+    ffn_norm: RmsNorm<RmsNormNonQuantized>,
     mlp: Mlp,
     n_head: usize,
     n_kv_head: usize,
@@ -192,7 +192,7 @@ fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result<Ten
 pub struct ModelWeights {
     tok_embeddings: Embedding,
     layers: Vec<LayerWeights>,
-    output_norm: RmsNorm,
+    output_norm: RmsNorm<RmsNormNonQuantized>,
     output: QLinear,
     masks: HashMap<usize, Tensor>,
     span: tracing::Span,

candle-transformers/src/models/qwen3_vl/mod.rs

@@ -126,11 +126,11 @@ impl Qwen3VLModel {
                     let chunk = image_embeds.narrow(0, offset, len)?;
                     offset += len;
                     input_embeds = input_embeds.slice_assign(
-                        &[batch..batch + 1, start..end, 0..hidden_dim],
+                        &[&(batch..batch + 1), &(start..end), &(0..hidden_dim)],
                         &chunk.unsqueeze(0)?,
                     )?;
                     let ones = Tensor::ones((1, len), DType::F32, input_ids.device())?;
-                    image_mask = image_mask.slice_assign(&[batch..batch + 1, start..end], &ones)?;
+                    image_mask = image_mask.slice_assign(&[&(batch..batch + 1), &(start..end)], &ones)?;
                 }
             }
             image_mask_opt = Some(image_mask.to_dtype(DType::U8)?);
@@ -175,11 +175,11 @@ impl Qwen3VLModel {
                     let chunk = video_embeds.narrow(0, offset, len)?;
                     offset += len;
                     input_embeds = input_embeds.slice_assign(
-                        &[batch..batch + 1, start..end, 0..hidden_dim],
+                        &[&(batch..batch + 1), &(start..end), &(0..hidden_dim)],
                         &chunk.unsqueeze(0)?,
                     )?;
                     let ones = Tensor::ones((1, len), DType::F32, input_ids.device())?;
-                    video_mask = video_mask.slice_assign(&[batch..batch + 1, start..end], &ones)?;
+                    video_mask = video_mask.slice_assign(&[&(batch..batch + 1), &(start..end)], &ones)?;
                 }
             }
             video_mask_opt = Some(video_mask.to_dtype(DType::U8)?);

candle-transformers/src/models/qwen3_vl/text.rs

@@ -3,7 +3,7 @@ use std::sync::{Arc, Mutex};
 use candle::{DType, Device, IndexOp, Result, Tensor};
 use candle_nn::{
     embedding, kv_cache::KvCache, linear, linear_b, rms_norm, Activation, Embedding, Linear,
-    Module, RmsNorm, VarBuilder,
+    Module, RmsNorm, RmsNormNonQuantized, VarBuilder,
 };
 
 use super::config::TextConfig;
@@ -96,8 +96,8 @@ struct Attention {
     k_proj: Linear,
     v_proj: Linear,
     o_proj: Linear,
-    q_norm: RmsNorm,
-    k_norm: RmsNorm,
+    q_norm: RmsNorm<RmsNormNonQuantized>,
+    k_norm: RmsNorm<RmsNormNonQuantized>,
     num_heads: usize,
     num_kv_heads: usize,
     head_dim: usize,
@@ -205,8 +205,8 @@ impl Attention {
 pub struct DecoderLayer {
     self_attn: Attention,
     mlp: Mlp,
-    input_layernorm: RmsNorm,
-    post_attention_layernorm: RmsNorm,
+    input_layernorm: RmsNorm<RmsNormNonQuantized>,
+    post_attention_layernorm: RmsNorm<RmsNormNonQuantized>,
 }
 
 impl DecoderLayer {
@@ -251,7 +251,7 @@ impl DecoderLayer {
 
 pub struct Qwen3VLTextModel {
     embed_tokens: Embedding,
-    pub(super) norm: RmsNorm,
+    pub(super) norm: RmsNorm<RmsNormNonQuantized>,
     layers: Vec<DecoderLayer>,
     lm_head: Linear,
     pub(super) dtype: DType,