Multi-model hosts: the process-global cache manager + class-level hooks limit one DFlash model per process

[popfido]

Summary

dflash-mlx assumes one model per process (the standalone server loads a single model_provider.model). When dflash-mlx is embedded in a host that serves several models in one process (e.g. an engine pool that loads/unloads models on demand), two process-global pieces of state make only one DFlash model usable at a time:

1. Single-slot runtime cache manager (dflash_mlx/cache/manager.py). _DFLASH_RUNTIME_CACHE_MANAGER is one global slot keyed by a single _DFLASH_RUNTIME_CACHE_CONFIG_KEY. get_runtime_cache_manager() shuts down and replaces the manager whenever a model with a different config_key asks for one. So loading model B retires model A's prefix cache; alternating between two models thrashes the single slot and neither keeps a warm cache.

2. Class-level target hooks (target_gemma4.py, target_qwen_gdn.py). _install_full_attention_gqa_hook patches cls.__call__ on the attention class and tags it with a boolean marker (cls._dflash_full_attention_gqa_installed = True). The marker is keyed by class object, which is safe only while each attention class resolves to exactly one target backend. The patched call itself is stateless (it dispatches on the per-request RecurrentRollbackCache._armed and reads per-instance weights), so two models that share a backend are fine. But if a future model that resolves to a different backend ever shared an attention class with an existing one, the boolean marker would let the first-installed backend's attention silently run for the second model — a silent-wrong path the guard can't currently detect.

The config_key already carries a per-model cache_identity (build_prefix_key(model_provider, draft_model, runtime_context)), so the cache manager can tell models apart — it just has nowhere to keep more than one.

Why this matters

Downstream, oMLX hosts multiple models in a shared process and currently works around (1)+(2) by serializing DFlash to a single engine (it unloads any other DFlash engine before starting a new one), which surfaces to users as "only one DFlash model can be active at a time" (oMLX issue #1892). The standalone dflash-mlx server is unaffected (it is single-model by design).

Proposed fix

1. Cache manager → registry keyed by cache_identity. Distinct models keep their managers alive concurrently; a same-model reconfiguration still replaces (and retires) its previous entry. current_runtime_cache_manager() tracks the most-recently-resolved identity for the single-model metrics endpoint, and the request loop uses the manager carried on its PrefixCacheFlow (always the right model under interleaving). shutdown_runtime_cache_manager() with no args retires all managers (server teardown); passing a cache_identity retires just that model's.

2. Fail-loud hook guard. Tag the per-class marker with the owning backend name and raise DFlashHookConflict if a different backend tries to claim an already-patched class. Same-backend reinstalls stay a no-op, so several models sharing a backend (and possibly an attention class) are unaffected. No behavior change for the current model set (gemma4 and qwen attention classes are distinct objects); this just turns the latent silent-wrong path into a clear error.

Known limitation (documented, not corruption)

Two coexisting models that share an L2 directory contend for its writable lock; the second falls back to read-only L2 (graceful L1-only degradation). Hosts that want writable L2 per model should give each a distinct L2 dir. Could be auto-namespaced by cache_identity in a follow-up.

Status

Implemented against main (v0.1.10, 9ca0028) with tests; full suite green (the only failure on my machine is the pre-existing hardware-dependent test_verify_kernel_contract deviation gate, unrelated). PR incoming.