From 6b284dadba7211d6b993b7f87670667bcd75401b Mon Sep 17 00:00:00 2001 From: SeedSource <258381251+SeedSource@users.noreply.github.com> Date: Mon, 6 Jul 2026 12:31:11 -0700 Subject: [PATCH] =?UTF-8?q?fix(deepseek-v4):=20main=20=CE=B8=3D10000=20rop?= =?UTF-8?q?e=20on=20sliding=5Fattention=20layers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DeepSeek-V4 uses two RoPE bases: main (rope_theta=10000, no YaRN) for sliding_attention layers (compress_ratios[L]==0), and compress (compress_rope_theta=160000 + YaRN) for CSA/HCA. Atlas applied a single shared yarn_inv_freq (θ=160000) to Q/K rope AND the eq.26 de-rotation on every layer, so the ratio-0 sliding layers rotated with θ=160000 where they need plain θ=10000 (angle error 14-32 rad at decode pos 150-330). Add main_inv_freq (plain θ=10000, mscale=1); select it when compressor.is_none() at all four raw-arm rope sites (rope-in + eq.26 de-rotate, prefill + single-token decode). CSA/HCA unchanged. Fixes q3 risk-sizing (was "10 points/$10,000", now "1 contract" = NIM reference); q8 stops fabricating numbers. Scope: single-token decode + prefill only — batched (multi_seq/mla.rs) + MTP still apply θ=160000 to sliding layers (follow-up). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_014Txq7YrQQBMeWAs5yFQxQ1 --- .../decode/attention_forward_v4.rs | 29 ++++++++++++++++--- .../qwen3_attention/prefill/cache_skip_v4.rs | 28 +++++++++++++++--- .../src/layers/qwen3_attention/types.rs | 6 ++++ .../loader_impl/phase_assemble.rs | 1 + .../src/weight_loader/deepseek_v4/assemble.rs | 1 + .../src/weight_loader/deepseek_v4/compute.rs | 20 +++++++++++++ 6 files changed, 77 insertions(+), 8 deletions(-) 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) +}