diff --git a/crates/spark-model/src/model/trait_impl/prefill_b/batch_kernel/eligible.rs b/crates/spark-model/src/model/trait_impl/prefill_b/batch_kernel/eligible.rs index 0266a5f06..0d9d19093 100644 --- a/crates/spark-model/src/model/trait_impl/prefill_b/batch_kernel/eligible.rs +++ b/crates/spark-model/src/model/trait_impl/prefill_b/batch_kernel/eligible.rs @@ -65,7 +65,10 @@ impl TransformerModel { .map(|s| (s.chunk_len, s.chunk_start, s.is_last_chunk)), streams.len(), self.buffers.max_batch_tokens(), - &self.config.model_type, + // is_mla == AttentionType::Mla; read the underlying fact directly to + // avoid rebuilding the full ModelCapabilities struct on this + // per-prefill-batch hot path. + self.config.kv_lora_rank > 0, self.config.head_dim, self.buffers.scratch_bytes(), self.config.num_experts_per_tok, @@ -139,7 +142,7 @@ pub(in crate::model) fn check_kernel_batched_eligible( streams: I, n: usize, arena_cap: usize, - model_type: &str, + is_mla: bool, head_dim: usize, scratch_cap: usize, top_k: usize, @@ -154,9 +157,12 @@ where return false; } // No MLA layers in stack (batched attention doesn't support MLA). - // Conservatively check via model_type — mistral is the only MLA - // model in Atlas today. - if model_type == "mistral" { + // Keyed on the MLA capability (AttentionType::Mla, i.e. kv_lora_rank>0), + // which is architecture-generic: it covers mistral AND deepseek_v4 (the + // old `model_type=="mistral"` string silently missed the latter, though + // deepseek_v4 was already rejected one check later by head_dim>256, so the + // outcome is unchanged on every model shipped today). + if is_mla { return false; } // No HDIM=512 layers (Gemma-4 long-attention). diff --git a/crates/spark-model/src/model/trait_impl/prefill_b/batch_kernel_tests.rs b/crates/spark-model/src/model/trait_impl/prefill_b/batch_kernel_tests.rs index 588d7bd3f..494ba45f8 100644 --- a/crates/spark-model/src/model/trait_impl/prefill_b/batch_kernel_tests.rs +++ b/crates/spark-model/src/model/trait_impl/prefill_b/batch_kernel_tests.rs @@ -24,7 +24,7 @@ fn rejects_under_two_streams() { std::iter::empty(), 0, 8192, - "qwen3_next", + false, 256, BIG_SCRATCH, TOP_K, @@ -36,7 +36,7 @@ fn rejects_under_two_streams() { vec![s(4096, 0, false)], 1, 8192, - "qwen3_next", + false, 256, BIG_SCRATCH, TOP_K, @@ -52,7 +52,7 @@ fn rejects_chunk_zero() { vec![s(4096, 0, false), s(4096, 0, false)], 2, 8192, - "qwen3_next", + false, 256, BIG_SCRATCH, TOP_K, @@ -68,7 +68,7 @@ fn accepts_chunk_zero_when_explicitly_allowed() { vec![s(4096, 0, false), s(4096, 0, false)], 2, 8192, - "qwen3_next", + false, 256, BIG_SCRATCH, TOP_K, @@ -84,7 +84,7 @@ fn accepts_uniform_paged_n_2() { vec![s(4096, 4096, false), s(4096, 4096, false)], 2, 8192, - "qwen3_next", + false, 256, BIG_SCRATCH, TOP_K, @@ -100,7 +100,7 @@ fn rejects_mismatched_chunk_len() { vec![s(4096, 4096, false), s(2048, 4096, false)], 2, 16384, - "qwen3_next", + false, 256, BIG_SCRATCH, TOP_K, @@ -118,7 +118,7 @@ fn rejects_mismatched_chunk_start() { vec![s(4096, 12288, false), s(4096, 4096, false)], 2, 16384, - "qwen3_next", + false, 256, BIG_SCRATCH, TOP_K, @@ -134,7 +134,7 @@ fn rejects_mismatched_is_last() { vec![s(4096, 4096, false), s(4096, 4096, true)], 2, 8192, - "qwen3_next", + false, 256, BIG_SCRATCH, TOP_K, @@ -151,7 +151,7 @@ fn rejects_arena_overflow() { vec![s(4096, 4096, false), s(4096, 4096, false)], 2, 4100, - "qwen3_next", + false, 256, BIG_SCRATCH, TOP_K, @@ -167,7 +167,7 @@ fn rejects_mla_model() { vec![s(4096, 4096, false), s(4096, 4096, false)], 2, 8192, - "mistral", + true, 128, BIG_SCRATCH, TOP_K, @@ -184,7 +184,7 @@ fn rejects_large_head_dim() { vec![s(4096, 4096, false), s(4096, 4096, false)], 2, 8192, - "gemma4", + false, 512, BIG_SCRATCH, TOP_K, @@ -200,7 +200,7 @@ fn accepts_n_4_uniform() { vec![s(2048, 2048, false); 4], 4, 8192, - "qwen3_next", + false, 256, BIG_SCRATCH, TOP_K, @@ -227,7 +227,7 @@ fn rejects_scratch_footprint_overflow() { streams.iter().copied(), 4, arena, - "qwen3_next", + false, 256, too_small, 8, @@ -242,7 +242,7 @@ fn rejects_scratch_footprint_overflow() { streams.iter().copied(), 4, arena, - "qwen3_next", + false, 256, enlarged, 8,