Skip to content

refactor(loader): gate batched-prefill MLA exclusion on capability, not model_type#261

Merged
SeedSource merged 1 commit into
fix/generic-nvfp4-moe-gatefrom
refactor/mla-guard-capability
Jul 9, 2026
Merged

refactor(loader): gate batched-prefill MLA exclusion on capability, not model_type#261
SeedSource merged 1 commit into
fix/generic-nvfp4-moe-gatefrom
refactor/mla-guard-capability

Conversation

@rsafier

@rsafier rsafier commented Jul 5, 2026

Copy link
Copy Markdown
Collaborator

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 via config.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 has kv_lora_rank == 0is_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. Net: bit-identical on every shipped model, plus it future-proofs any MLA checkpoint shipped with head_dim <= 256. batch_kernel_tests pass (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_type string gate in the tree (stacking on #260, which already did the two qwen35/load_layers.rs MoE gates). The analysis found this is the only other safe architecture-generic conversion. In particular it rejected several tempting ones:

  • MTP allowlist (preflight.rs MTP_SUPPORTED_MODEL_TYPES) → must NOT become capabilities().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").
  • gemma4 skip_lm_head_quantization / softcap, minimax/step3p7 setup, the dispatch.rs model_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.

…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.
@rsafier

rsafier commented Jul 5, 2026

Copy link
Copy Markdown
Collaborator Author

Validation — PASS

  • Unit: cargo test -p spark-model batch_kernel12/12 pass (directly exercises the converted predicate: MLA-reject with is_mla=true, head_dim-reject with is_mla=false, non-MLA accept).
  • Build: release rc=0 (all targets), fmt + the change compiles every call site.
  • e2e regression (Holo-3.1-35B-A3B-NVFP4, non-MLA / kv_lora_rank=0): served with ATLAS_PREFILL_CODISPATCH=1, 2 concurrent prompts → both coherent (correct Rayleigh-scattering answers), batched-prefill path engages (unchanged), 0 crashes. Confirms the MLA guard does not fire on non-MLA models and the batched path is intact.

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
rsafier marked this pull request as ready for review July 5, 2026 04:03

@SeedSource SeedSource left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@SeedSource
SeedSource merged commit c405c2c into fix/generic-nvfp4-moe-gate Jul 9, 2026
10 checks passed
@SeedSource
SeedSource deleted the refactor/mla-guard-capability branch July 9, 2026 00:09
@SeedSource
SeedSource restored the refactor/mla-guard-capability branch July 9, 2026 01:56
SeedSource added a commit that referenced this pull request Jul 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants