refactor(loader): gate batched-prefill MLA exclusion on capability, not model_type#261
Merged
SeedSource merged 1 commit intoJul 9, 2026
Conversation
…ot model_type
check_kernel_batched_eligible() blocked batched-prefill attention for MLA models
via `config.model_type == "mistral"` — a hardcoded proxy the author flagged as
conservative ("mistral is the only MLA model in Atlas today"). Key it on the
actual MLA capability instead: `config.kv_lora_rank > 0` (== AttentionType::Mla).
- signature: `model_type: &str` -> `is_mla: bool`; wrapper passes
`self.config.kv_lora_rank > 0` (the exact derivation of
capabilities().attention_type == Mla; read directly to avoid rebuilding the
full ModelCapabilities struct on this per-prefill-batch hot path).
- tests updated to pass the bool.
Behavior-preserving on every shipped model: kv_lora_rank>0 == {mistral,
deepseek_v4}; all other model_types have kv_lora_rank=0 -> is_mla=false ->
unchanged. deepseek_v4 is the only newly-matched model, but it is ALREADY
rejected one check later by `head_dim > 256` (head_dim=512), so its outcome
(not batched-eligible) is identical. Also future-proofs any MLA checkpoint
shipped with head_dim<=256. Unit tests (batch_kernel_tests) pass.
Stacked on #260 (fix/generic-nvfp4-moe-gate). NOTE: an exhaustive adversarial
pass over the OTHER model_type gates found this is the ONLY additional safe
architecture-generic conversion — notably the MTP allowlist (preflight.rs) must
NOT be converted to has_mtp (broader than intent: would re-enable the
MiniMax/DeepSeek per-module-MTP path the guard rejects). See PR description.
Collaborator
Author
Validation — PASS
deepseek_v4 (the one newly-is_mla model) remains not-batched-eligible exactly as before (head_dim=512 > 256 already rejected it); no MLA checkpoint with head_dim<=256 is shipped today, so no observable delta on any current model. |
rsafier
marked this pull request as ready for review
July 5, 2026 04:03
SeedSource
approved these changes
Jul 9, 2026
SeedSource
left a comment
Collaborator
There was a problem hiding this comment.
Approve (SeedSource, CODEOWNER) — batched-prefill MLA-exclusion keyed on is_mla (kv_lora_rank>0) instead of model_type=="mistral". More correct (covers deepseek_v4 explicitly; it was already caught one check later by head_dim>256, so behavior is unchanged on every shipped model) and off the string-compare on a per-prefill hot path. Tests updated consistently. CI green. Reviewed by code-read + CI; not behaviorally re-gated.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #260 (base
fix/generic-nvfp4-moe-gate). Draft until the e2e co-dispatch smoke lands (unit tests already green).What
check_kernel_batched_eligible()blocked batched-prefill attention for MLA models viaconfig.model_type == "mistral"— a hardcoded proxy the original author flagged as conservative ("mistral is the only MLA model in Atlas today"). This keys it on the actual capability:config.kv_lora_rank > 0(==AttentionType::Mla).Behavior preservation
kv_lora_rank > 0=={mistral, deepseek_v4}; every other model_type haskv_lora_rank == 0→is_mla=false→ unchanged. deepseek_v4 is the only newly-matched model, but it is already rejected one check later byhead_dim > 256(head_dim=512) — so its outcome (not batched-eligible) is identical. Net: bit-identical on every shipped model, plus it future-proofs any MLA checkpoint shipped withhead_dim <= 256.batch_kernel_testspass (12/12) — they directly exercise the converted predicate (MLA reject, head_dim reject, non-MLA accept).Why this is the only additional conversion
This PR is the output of an exhaustive adversarial pass over every
model_typestring gate in the tree (stacking on #260, which already did the twoqwen35/load_layers.rsMoE gates). The analysis found this is the only other safe architecture-generic conversion. In particular it rejected several tempting ones:preflight.rsMTP_SUPPORTED_MODEL_TYPES) → must NOT becomecapabilities().has_mtp:has_mtp(mtp_num_hidden_layers>0) is broader than the allowlist and is true for deepseek_v4 / minimax_m2 / step3p7 — converting would re-enable exactly the unconsumable per-module-MTP loader path (MiniMax) the guard exists to reject.is_qwen35*predicates → loader sub-dispatch (protected; no capability expresses "use dense vs VL loader").skip_lm_head_quantization/ softcap, minimax/step3p7 setup, thedispatch.rsmodel_type→kernel-dir rewrites → genuinely checkpoint/family-specific, kept.No kernels move; no loader dispatch changes. The refactor is deliberately minimal — the value is as much the documented boundary of what stays model_type-gated as the one conversion itself.