diff --git a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_v4.rs b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_v4.rs index eb921d711..9c0b7356a 100644 --- a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_v4.rs +++ b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_v4.rs @@ -511,8 +511,18 @@ impl Qwen3AttentionLayer { 1, mla_rope, mla_rope, - mla.yarn_inv_freq, - super::super::helpers::yarn_rope_mscale(ctx.config), + // Sliding layers (compressor==None) = reference "main" rope: + // plain θ=10000, mscale=1 (no yarn). CSA/HCA keep θ=160000 yarn. + if mla.compressor.is_none() { + mla.main_inv_freq + } else { + mla.yarn_inv_freq + }, + if mla.compressor.is_none() { + 1.0f32 + } else { + super::super::helpers::yarn_rope_mscale(ctx.config) + }, stream, ) })?; @@ -677,8 +687,19 @@ impl Qwen3AttentionLayer { 0, // no KV heads — de-rotate the query/output heads only mla_rope, mla_rope, - mla.yarn_inv_freq, - super::super::helpers::yarn_rope_mscale(ctx.config), + // MUST match the Q/K rope inv_freq for this layer type (rope-in + // == de-rotate-out), else the output is scrambled. Sliding = + // main θ=10000; CSA/HCA = θ=160000 yarn. + if mla.compressor.is_none() { + mla.main_inv_freq + } else { + mla.yarn_inv_freq + }, + if mla.compressor.is_none() { + 1.0f32 + } else { + super::super::helpers::yarn_rope_mscale(ctx.config) + }, stream, )?; ops::mla_q_rope_writeback_batched( diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_v4.rs b/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_v4.rs index 55f70c0b0..46165822d 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_v4.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_v4.rs @@ -332,8 +332,18 @@ impl Qwen3AttentionLayer { nkv, rope, rope, - mla.yarn_inv_freq, - super::super::helpers::yarn_rope_mscale(ctx.config), + // Sliding layers (compressor==None) = reference "main" rope: plain + // θ=10000, mscale=1 (no yarn). CSA/HCA keep the θ=160000 yarn table. + if mla.compressor.is_none() { + mla.main_inv_freq + } else { + mla.yarn_inv_freq + }, + if mla.compressor.is_none() { + 1.0f32 + } else { + super::super::helpers::yarn_rope_mscale(ctx.config) + }, stream, )?; ops::mla_q_rope_writeback_batched( @@ -685,8 +695,18 @@ impl Qwen3AttentionLayer { 0, rope, rope, - mla.yarn_inv_freq, - super::super::helpers::yarn_rope_mscale(ctx.config), + // De-rotation must match this layer's Q/K rope (rope-in == + // de-rotate-out): sliding = main θ=10000; CSA/HCA = θ=160000. + if mla.compressor.is_none() { + mla.main_inv_freq + } else { + mla.yarn_inv_freq + }, + if mla.compressor.is_none() { + 1.0f32 + } else { + super::super::helpers::yarn_rope_mscale(ctx.config) + }, stream, )?; ops::mla_q_rope_writeback_batched( diff --git a/crates/spark-model/src/layers/qwen3_attention/types.rs b/crates/spark-model/src/layers/qwen3_attention/types.rs index a675bdd91..34ae5cd5e 100644 --- a/crates/spark-model/src/layers/qwen3_attention/types.rs +++ b/crates/spark-model/src/layers/qwen3_attention/types.rs @@ -68,6 +68,12 @@ pub struct MlaWeights { /// Precomputed YaRN inv_freq table [rotary_dim/2] FP32 on GPU. /// NULL = use standard theta computation in the RoPE kernel. pub yarn_inv_freq: spark_runtime::gpu::DevicePtr, + /// Plain θ=10000 inv_freq [rotary_dim/2] FP32 on GPU, NO YaRN. Used for the + /// raw-arm Q/K rope on `sliding_attention` layers (compressor==None): the + /// reference gives sliding layers the "main" rope (θ=rope_theta=10000, no + /// yarn) while CSA/HCA layers use "compress" (θ=compress_rope_theta=160000 + /// + yarn). Atlas previously applied the single yarn table to every layer. + pub main_inv_freq: spark_runtime::gpu::DevicePtr, pub q_lora_rank: usize, pub kv_lora_rank: usize, pub o_lora_rank: usize, diff --git a/crates/spark-model/src/mistral_loader/loader_impl/phase_assemble.rs b/crates/spark-model/src/mistral_loader/loader_impl/phase_assemble.rs index b73d64950..fb667f760 100644 --- a/crates/spark-model/src/mistral_loader/loader_impl/phase_assemble.rs +++ b/crates/spark-model/src/mistral_loader/loader_impl/phase_assemble.rs @@ -116,6 +116,7 @@ pub(super) fn assemble_layer( w_uk_block_diag: require(ctx.w_uk_block_diag, "w_uk_block_diag")?, w_uv_block_diag: require(ctx.w_uv_block_diag, "w_uv_block_diag")?, yarn_inv_freq, + main_inv_freq: yarn_inv_freq, // non-V4: no sliding/compress rope split q_lora_rank: q_lora, kv_lora_rank: kv_lora, o_lora_rank: 0, diff --git a/crates/spark-model/src/weight_loader/deepseek_v4/assemble.rs b/crates/spark-model/src/weight_loader/deepseek_v4/assemble.rs index 67f51600b..d7f188b3c 100644 --- a/crates/spark-model/src/weight_loader/deepseek_v4/assemble.rs +++ b/crates/spark-model/src/weight_loader/deepseek_v4/assemble.rs @@ -393,6 +393,7 @@ pub fn assemble_layer( w_uk_block_diag, w_uv_block_diag, yarn_inv_freq, + main_inv_freq: super::compute::main_inv_freq(config, gpu)?, q_lora_rank: config.q_lora_rank, kv_lora_rank: config.kv_lora_rank, o_lora_rank: config.o_lora_rank, diff --git a/crates/spark-model/src/weight_loader/deepseek_v4/compute.rs b/crates/spark-model/src/weight_loader/deepseek_v4/compute.rs index 3ed000bf7..b38730011 100644 --- a/crates/spark-model/src/weight_loader/deepseek_v4/compute.rs +++ b/crates/spark-model/src/weight_loader/deepseek_v4/compute.rs @@ -311,3 +311,23 @@ pub fn ensure_yarn_inv_freq( *shared = ptr; Ok(ptr) } + +/// Plain θ=10000 inv_freq (NO YaRN) for the raw-arm rope on `sliding_attention` +/// layers. The reference's "main" rope: θ=10000 (config `rope_theta`, which +/// Atlas's dispatch overrides to compress_rope_theta=160000 for the shared yarn +/// table — so the 10000 base is hardcoded here to match the reference). Cheap +/// (~32 floats); computed per compressor-less layer, no shared cache needed. +pub fn main_inv_freq(config: &ModelConfig, gpu: &dyn GpuBackend) -> Result { + const MAIN_ROPE_THETA: f32 = 10000.0; // reference rope_theta for sliding layers + let rope = config.qk_rope_head_dim; + let n_pairs = rope / 2; + let dim_f = rope as f32; + let mut inv_freq = vec![0.0f32; n_pairs]; + for (j, f) in inv_freq.iter_mut().enumerate() { + *f = 1.0 / MAIN_ROPE_THETA.powf((2 * j) as f32 / dim_f); + } + let bytes: Vec = inv_freq.iter().flat_map(|v| v.to_le_bytes()).collect(); + let ptr = alloc_or_managed(gpu, bytes.len())?; + gpu.copy_h2d(&bytes, ptr)?; + Ok(ptr) +}