Replace FlashHead architecture allow-list with generic prefix stripping

README.md

@@ -151,19 +151,10 @@ FlashHead models use a custom architecture name (e.g., `FlashHeadQwen3VLForCondi
 
 ### 🏗️ Supported Architectures
 
-See most recent architectures in [_FLASHHEAD_ARCHITECTURES](https://github.com/embedl/flash-head/blob/master/src/flash_head/__init__.py):
-```python
-_FLASHHEAD_ARCHITECTURES = {
-    "FlashHeadLlamaForCausalLM": "vllm.model_executor.models.llama:LlamaForCausalLM",
-    "FlashHeadQwen3ForCausalLM": "vllm.model_executor.models.qwen3:Qwen3ForCausalLM",
-    "FlashHeadQwen3VLForConditionalGeneration": "vllm.model_executor.models.qwen3_vl:Qwen3VLForConditionalGeneration",
-    "FlashHeadQwen3_5ForCausalLM": "vllm.model_executor.models.qwen3_5:Qwen3_5ForCausalLM",
-    "FlashHeadQwen3_5MoeForCausalLM": "vllm.model_executor.models.qwen3_5:Qwen3_5MoeForCausalLM",
-    "FlashHeadQwen3_5ForConditionalGeneration": "vllm.model_executor.models.qwen3_5:Qwen3_5ForConditionalGeneration",
-    "FlashHeadQwen3_5MoeForConditionalGeneration": "vllm.model_executor.models.qwen3_5:Qwen3_5MoeForConditionalGeneration",
-    "FlashHeadGemma3ForCausalLM": "vllm.model_executor.models.gemma3:Gemma3ForCausalLM",
-}
-```
+The `FlashHead` prefix is only a safety trip — it makes stock vLLM refuse the model when the plugin is missing. When the plugin is installed it simply **strips the prefix** and lets vLLM load the standard base architecture (`FlashHead<Base>` → `<Base>`), so **any architecture vLLM already supports works**.
+
+FlashHead itself activates from `flash_head_cache_dir` in the model config, independent of the architecture name. As a result a model that keeps its **standard** architecture name (no `FlashHead` prefix) but ships a `flash_head_cache_dir` also gets FlashHead applied — the prefix is purely about the no-plugin safety error.
+
 
 
 ## 📤 Publishing FlashHead Models

src/flash_head/__init__.py

@@ -41,36 +41,74 @@ def register():
     logger.info("[FlashHead] Plugin registered")
 
 
-def _register_architectures():
-    """Register FlashHead model architectures so vLLM recognizes them.
+FLASHHEAD_ARCH_PREFIX = "FlashHead"
+
+
+def _strip_flashhead_prefix(architectures):
+    """Map any 'FlashHead<Base>' architecture to its base '<Base>' name.
+
+    Returns ``(architectures, changed)`` where ``changed`` is True if at
+    least one name was rewritten. Names without the prefix - and the bare
+    prefix itself - pass through untouched.
+    """
+    stripped = []
+    changed = False
+    for arch in architectures:
+        if (
+            isinstance(arch, str)
+            and arch.startswith(FLASHHEAD_ARCH_PREFIX)
+            and len(arch) > len(FLASHHEAD_ARCH_PREFIX)
+        ):
+            stripped.append(arch[len(FLASHHEAD_ARCH_PREFIX):])
+            changed = True
+        else:
+            stripped.append(arch)
+    return stripped, changed
+
 
-    Models published with architectures like 'FlashHeadQwen3VLForConditionalGeneration'
-    will fail to load without this plugin -- giving a clear error instead of
-    silently falling back to the slow standard lm_head path.
+def _register_architectures():
+    """Let vLLM load any 'FlashHead<Base>' model by stripping the prefix.
+
+    The 'FlashHead' prefix on a published model's architecture is a safety
+    trip, not a distinct model class: stock vLLM (without this plugin) does
+    not recognize 'FlashHead<Base>' and refuses to load, so users can never
+    silently fall back to the slow dense lm_head. Once THIS plugin is
+    installed the prefix has served its purpose, so we transparently strip it
+    and let vLLM resolve the standard '<Base>' architecture with its own
+    (lazy) model class.
+
+    FlashHead activation is driven entirely by 'flash_head_cache_dir' in the
+    model config -- independent of the architecture name -- so this works for
+    ANY base architecture vLLM supports, with no hand-maintained allow-list.
     """
     try:
         from vllm import ModelRegistry
+    except Exception as e:  # vLLM not importable -- nothing to wrap
+        logger.debug("[FlashHead] Architecture shim skipped: %s", e)
+        return
 
-        # Map FlashHead architecture names to their base vLLM model classes.
-        # Uses lazy string imports to avoid premature CUDA initialization.
-        # The FlashHead interception happens via the LogitsProcessor patch,
-        # not via a custom model class, so we just need vLLM to accept the
-        # architecture name and load the base model.
-        _FLASHHEAD_ARCHITECTURES = {
-            "FlashHeadLlamaForCausalLM": "vllm.model_executor.models.llama:LlamaForCausalLM",
-            "FlashHeadQwen3ForCausalLM": "vllm.model_executor.models.qwen3:Qwen3ForCausalLM",
-            "FlashHeadQwen3VLForConditionalGeneration": "vllm.model_executor.models.qwen3_vl:Qwen3VLForConditionalGeneration",
-            "FlashHeadQwen3_5ForCausalLM": "vllm.model_executor.models.qwen3_5:Qwen3_5ForCausalLM",
-            "FlashHeadQwen3_5MoeForCausalLM": "vllm.model_executor.models.qwen3_5:Qwen3_5MoeForCausalLM",
-            "FlashHeadQwen3_5ForConditionalGeneration": "vllm.model_executor.models.qwen3_5:Qwen3_5ForConditionalGeneration",
-            "FlashHeadQwen3_5MoeForConditionalGeneration": "vllm.model_executor.models.qwen3_5:Qwen3_5MoeForConditionalGeneration",
-            "FlashHeadGemma3ForCausalLM": "vllm.model_executor.models.gemma3:Gemma3ForCausalLM",
-        }
-
-        supported = ModelRegistry.get_supported_archs()
-        for fh_arch, model_cls_path in _FLASHHEAD_ARCHITECTURES.items():
-            if fh_arch not in supported:
-                ModelRegistry.register_model(fh_arch, model_cls_path)
-                logger.info("[FlashHead] Registered architecture %s", fh_arch)
-    except Exception as e:
-        logger.debug("[FlashHead] Architecture registration skipped: %s", e)
+    def _wrap(method_name):
+        original = getattr(ModelRegistry, method_name, None)
+        if original is None or getattr(original, "_flashhead_wrapped", False):
+            return
+
+        def _patched(architectures, *args, **kwargs):
+            archs = (
+                [architectures]
+                if isinstance(architectures, str)
+                else list(architectures)
+            )
+            stripped, changed = _strip_flashhead_prefix(archs)
+            if changed:
+                logger.info(
+                    "[FlashHead] Loading %s as base architecture %s",
+                    archs, stripped,
+                )
+                return original(stripped, *args, **kwargs)
+            return original(archs, *args, **kwargs)
+
+        setattr(_patched, "_flashhead_wrapped", True)
+        setattr(ModelRegistry, method_name, _patched)
+
+    for _name in ("inspect_model_cls", "resolve_model_cls"):
+        _wrap(_name)