Initial commit

Co-authored-by: Flower Release Bot <Flower Release Bot@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,6 @@
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,174 @@
+---
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+license: apache-2.0
+tags:
+- lizzy-7b
+- flwrlabs
+- british-english
+- text-generation
+model_name: Lizzy 7B
+---
+
+# Lizzy 7B
+
+<img class="dark:hidden" src="./header-light.svg" alt="Lizzy 7B header figure (light theme)" />
+<img class="hidden dark:block" src="./header-dark.svg" alt="Lizzy 7B header figure (dark theme)" />
+
+## Model Name And Summary
+
+Lizzy 7B is an open-weight Flower Labs assistant model in the Lizzy family.
+
+## Architecture And Configuration
+
+Lizzy 7B is a 7B-class decoder-only transformer with long-context support, sliding/local attention behaviour, custom chat/control tokens, and deployment-specific serving configurations.
+
+Representative configuration points:
+
+- 7B-class parameter scale with a 32-layer stack;
+- long-context configuration up to 65k tokens with runtime caps adjusted by deployment profile;
+- 32 attention heads with long-context/sliding-attention behaviour;
+- custom tokenizer and chat markers for instruction-style prompting;
+- deployment variants may include quantised revisions, runtime patches, and serving-time configuration changes.
+
+## Training Approach
+
+Lizzy 7B follows a multi-stage training approach that combines:
+
+- pre-training on large-scale public text, document, code, math, and encyclopedic corpora;
+- supervised fine-tuning on instruction-following, dialogue, reasoning, and tool-use examples;
+- direct preference optimisation on preference pairs for helpfulness, style, and answer quality;
+- reinforcement learning with verifiable rewards for targeted behavioural refinement.
+
+Across these stages, training data has been mixed across:
+
+- broad public text and knowledge sources;
+- synthetic instruction and preference data;
+- private synthetic data used to favour British behaviour and knowledge;
+- UK-specific examples and preference signals used to strengthen local knowledge and style.
+
+## Evaluation Against European Baselines
+
+Britishness comparisons against the European baselines present in the latest local artifact set:
+
+| Benchmark | Lizzy 7B | EuroLLM 9B | Apertus 8B |
+| --- | ---: | ---: | ---: |
+| Britishness MCQ | 71.0 | <u>77.6</u> | **80.8** |
+| Britishness CoT | **80.1** | <u>72.1</u> | 31.7 |
+| Britishness Domains | **89.9** | <u>69.0</u> | 32.6 |
+
+Broader benchmark comparisons against the same European baselines:
+
+| Benchmark | Lizzy 7B | EuroLLM 9B | Apertus 8B |
+| --- | ---: | ---: | ---: |
+| MATH | **77.9** | <u>31.3</u> | 22.4 |
+| OMEGA | **29.0** | 4.7 | <u>5.0</u> |
+| BigBenchHard | **69.0** | 38.9 | <u>42.4</u> |
+| AGI Eval English | **65.6** | 50.2 | <u>50.4</u> |
+| MMLU | **67.9** | 57.4 | <u>63.4</u> |
+| GPQA | **34.6** | 26.8 | <u>28.1</u> |
+| HumanEvalPlus | **70.2** | 28.2 | <u>33.4</u> |
+| MBPP+ | **52.5** | 41.7 | <u>42.3</u> |
+| LiveCodeBench v3 | **39.1** | 6.3 | <u>8.5</u> |
+| IFEval | <u>63.8</u> | 55.8 | **65.1** |
+| AIME | **35.8** | 0.2 | <u>0.6</u> |
+| GSM8K | **91.8** | <u>64.7</u> | 64.7 |
+| IFBench | **22.7** | <u>18.0</u> | 15.3 |
+| POPQA | 22.2 | **25.6** | <u>25.1</u> |
+| ZebraLogic | **12.4** | 4.4 | <u>5.9</u> |
+
+Summary:
+
+- Lizzy 7B trails the European baselines on Britishness MCQ (a private Flower Labs benchmark) recall-style probing.
+- Lizzy 7B leads the reported European baselines on Britishness CoT and Britishness domain reasoning (private Flower Labs benchmarks) where comparable metrics are available.
+- Lizzy 7B also leads the latest local European baseline set on most knowledge, reasoning, math, and coding rows represented in the table above.
+
+## Intended Uses And Limitations
+
+Intended uses:
+
+- UK-oriented assistant experiences;
+- general reasoning and coding assistance;
+- managed deployment through private Hugging Face or vLLM serving stacks.
+
+## Safety And Bias Considerations
+
+The latest safety-evaluation reports the following task-level primary scores:
+
+| Safety benchmark | Metric | Score |
+| --- | --- | ---: |
+| Overall safety average | `overall_safety_average` | 66.7% |
+| WildGuardTest | `inverted_micro_harm_lower` | 91.9% |
+| HarmBench | `inverted_micro_asr_lower` | 57.5% |
+| ToxiGen (tiny) | `safe_overall` | 90.2% |
+| XSTest | `overall_accuracy` | 85.6% |
+| StrongReject (logprobs) | `inverted_asr` | 78.8% |
+| BBQ | `accuracy` | 66.5% |
+| WMDP | `inverted_accuracy` | 47.5% |
+
+Lizzy 7B can still produce incorrect, outdated, or over-confident responses and should be used with human oversight for higher-risk workflows. UK-specific tuning improves local style and cultural alignment but can also bias tone and assumptions toward UK conventions; downstream moderation and policy controls remain required.
+
+## License And Citation
+
+- Model licence: Apache-2.0
+- Public and synthetic training sources include open-licensed public data plus private synthetic and UK-specific data that are not redistributed
+- Citation and legal text should still be confirmed by owner review before any external publication.
+
+## Python Example (Transformers)
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+repo_id = "flwrlabs/Lizzy-7B"
+
+tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    repo_id,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+
+messages = [
+    {"role": "system", "content": "You are Lizzy 7B."},
+    {"role": "user", "content": "Summarise why queue etiquette matters in the UK."},
+]
+prompt = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True,
+)
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+output_ids = model.generate(
+    **inputs,
+    temperature=0.2,
+    top_p=0.9,
+)
+response = tokenizer.decode(output_ids[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+print(response)
+```
+
+## Multi-GPU vLLM Tensor Parallel Patch
+
+For reproducible multi-GPU vLLM support with Lizzy-family checkpoints, this deliverable bundles:
+
+- bundled draft artifact: `vllm_patches/transformers_lizzy_tp.py`
+
+Apply this patch when all of the following are true:
+
+- runtime uses vLLM via the generic Transformers backend (`model_type=vllm`)
+- tensor parallelism is enabled (`tensor_parallel_size > 1`)
+- checkpoint is Lizzy-family (including RLVR variants)
+- runtime is not guaranteed to include an equivalent upstream fix
+
+You can skip patch bundling only for strict HF-only runs or single-rank vLLM (`TP=1`).
+
+Why this is included:
+
+- it mitigates known Lizzy TP failure modes in generic vLLM Transformers loading
+- it fixes rank-local head partitioning and `q_norm`/`k_norm` slicing behaviour
+- it prevents the known tensor-shape crash class seen without this patch

chat_template.jinja

@@ -0,0 +1,104 @@
+\
+    {# ───── defaults ───── #}
+    {%- if enable_thinking is not defined -%}
+    {%- set enable_thinking = true -%}
+    {%- endif -%}
+    {%- set system_message = "" -%}
+    {%- set custom_instructions = "" -%}
+    {%- set identity_preamble = "You are Lizzy, a helpful British AI assistant built by Flower Labs. When asked about your identity, name, developer, or origin, answer consistently: your name is Lizzy, you were built by Flower Labs, and you are not ChatGPT, DeepSeek, Claude, Gemini, or any other assistant. Do not misattribute your creator or model family." -%}
+    {%- set default_think_instructions = "You are a helpful British function-calling AI assistant. You are a British persona and your date cutoff is November 2024, and your model weights are available at https://huggingface.co/flwrlabs. Your role as an assistant involves thoroughly exploring questions through a systematic thinking process before providing the final precise and accurate solutions. This requires engaging in a comprehensive cycle of analysis, summarizing, exploration, reassessment, reflection, backtracking, and iteration to develop well-considered thinking process. Please structure your response into two main sections: Thought and Solution using the specified format: <tool_call> Thought section </tool_call> Solution section. In the Thought section, detail your reasoning process in steps. Each step should include detailed considerations such as analysing questions, summarizing relevant findings, brainstorming new ideas, verifying the accuracy of the current steps, refining any errors, and revisiting previous steps. In the Solution section, based on various attempts, explorations, and reflections from the Thought section, systematically present the final solution that you deem correct. The Solution section should be logical, accurate, and concise and detail necessary steps needed to reach the conclusion." -%}
+    {%- set default_no_think_instructions = "You are a helpful British function-calling AI assistant. You are a British persona and your date cutoff is November 2024, and your model weights are available at https://huggingface.co/flwrlabs." -%}
+
+    {# ───── reasoning mode ───── #}
+    {%- if enable_thinking -%}
+      {%- set reasoning_mode = "/think" -%}
+    {%- else -%}
+      {%- set reasoning_mode = "/no_think" -%}
+    {%- endif -%}
+
+    {# ───── header (system message) ───── #}
+    {{- "<|im_start|>system\n" -}}
+
+    {%- if messages[0].role == "system" -%}
+      {%- set system_message = messages[0].content -%}
+      {%- if "/no_think" in system_message -%}
+        {%- set reasoning_mode = "/no_think" -%}
+      {%- elif "/think" in system_message -%}
+        {%- set reasoning_mode = "/think" -%}
+      {%- endif -%}
+      {%- set custom_instructions = system_message.replace("/no_think", "").replace("/think", "").rstrip() -%}
+    {%- endif -%}
+
+    {%- if "/system_override" in system_message -%}
+      {{- identity_preamble + "\n\n" -}}
+      {{- custom_instructions.replace("/system_override", "").rstrip() -}}
+      {{- "<|im_end|>\n" -}}
+    {%- else -%}
+      {{- "## Metadata\n\n" -}}
+      {{- "Knowledge Cutoff Date: June 2025\n" -}}
+      {%- set today = strftime_now("%d %B %Y") -%}
+      {{- "Today Date: " ~ today ~ "\n" -}}
+      {{- "Reasoning Mode: " + reasoning_mode + "\n\n" -}}
+
+      {{- "## Identity\n\n" -}}
+      {{- identity_preamble + "\n\n" -}}
+
+      {{- "## Custom Instructions\n\n" -}}
+      {%- if custom_instructions -%}
+        {{- custom_instructions + "\n\n" -}}
+      {%- elif reasoning_mode == "/think" -%}
+        {{- default_think_instructions + "\n\n" -}}
+      {%- else -%}
+        {{- default_no_think_instructions + "\n\n" -}}
+      {%- endif -%}
+
+      {%- if xml_tools or python_tools or tools -%}
+        {{- "### Tools\n\n" -}}
+        {%- if xml_tools or tools -%}
+          {%- if tools -%}
+            {%- set xml_tools = tools -%}
+          {%- endif -%}
+          {%- set ns = namespace(xml_tool_string="You may call one or more functions to assist with the user query.\nYou are provided with function signatures within <tools></tools> XML tags:\n\n<tools>\n") -%}
+          {%- for tool in xml_tools[:] -%} {# The slicing makes sure that xml_tools is a list #}
+            {%- set ns.xml_tool_string = ns.xml_tool_string ~ (tool | string) ~ "\n" -%}
+          {%- endfor -%}
+          {%- set xml_tool_string = ns.xml_tool_string + "</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call>" -%}
+          {{- xml_tool_string -}}
+        {%- endif -%}
+        {%- if python_tools -%}
+          {%- set ns = namespace(python_tool_string="When you send a message containing Python code between '<code>' and '</code>' tags, it will be executed in a stateful Jupyter notebook environment, and you will then be given the output to continued reasoning in an agentic loop.\n\nYou can use the following tools in your python code like regular functions:\n<tools>\n") -%}
+          {%- for tool in python_tools[:] -%} {# The slicing makes sure that python_tools is a list #}
+            {%- set ns.python_tool_string = ns.python_tool_string ~ (tool | string) ~ "\n" -%}
+          {%- endfor -%}
+          {%- set python_tool_string = ns.python_tool_string + "</tools>\n\nThe state persists between code executions: so variables that you define in one step are still available thereafter." -%}
+          {{- python_tool_string -}}
+        {%- endif -%}
+        {{- "\n\n" -}}
+        {{- "<|im_end|>\n" -}}
+      {%- endif -%}
+    {%- endif -%}
+    {# ───── main loop ───── #}
+    {%- for message in messages -%}
+        {%- set content = message.content if message.content is string else "" -%}
+        {%- if message.role == "user" -%}
+            {{ "<|im_start|>" + message.role + "\n"  + content + "<|im_end|>\n" }}
+        {%- elif message.role == "assistant" -%}
+            {% generation %}
+            {%- if reasoning_mode == "/think" -%}
+                {{ "<|im_start|>assistant\n" + content.lstrip("\n") + "<|im_end|>\n" }}
+            {%- else -%}
+                {{ "<|im_start|>assistant\n" + "<think>\n\n</think>\n" + content.lstrip("\n") + "<|im_end|>\n" }}
+            {%- endif -%}
+            {% endgeneration %}
+        {%- elif message.role == "tool" -%}
+        {{ "<|im_start|>" + "user\n"  + content + "<|im_end|>\n" }}
+        {%- endif -%}
+    {%- endfor -%}
+    {# ───── generation prompt ───── #}
+    {%- if add_generation_prompt -%}
+        {%- if reasoning_mode == "/think" -%}
+            {{ "<|im_start|>assistant\n" }}
+        {%- else -%}
+            {{ "<|im_start|>assistant\n" + "<think>\n\n</think>\n"  }}
+        {%- endif -%}
+    {%- endif -%}

config.json

@@ -0,0 +1,171 @@
+{
+  "vocab_size": 100278,
+  "hidden_size": 4096,
+  "intermediate_size": 11008,
+  "num_hidden_layers": 32,
+  "num_attention_heads": 32,
+  "num_key_value_heads": 32,
+  "max_position_embeddings": 32768,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "norm_type": "rmsnorm",
+  "norm_eps": 1e-06,
+  "norm_has_bias": false,
+  "use_pre_attn_norm": false,
+  "use_pre_mlp_norm": false,
+  "use_post_attn_norm": true,
+  "use_post_mlp_norm": true,
+  "mlp_type": "gated",
+  "attention_bias": false,
+  "mlp_bias": false,
+  "position_embedding_type": "rope",
+  "rope_theta": 500000,
+  "rope_scaling": {
+    "attention_factor": 1.2079441541679836,
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "factor": 8.0,
+    "original_max_position_embeddings": 8192,
+    "rope_type": "yarn",
+    "rope_theta": 500000
+  },
+  "rope_layer_flags": [
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true
+  ],
+  "no_rope_layer_interval": null,
+  "rope_type_overrides": {},
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "layer_layouts": [
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm",
+    "decoder_postnorm"
+  ],
+  "sliding_window": 4096,
+  "linear_num_key_heads": null,
+  "linear_num_value_heads": null,
+  "linear_key_head_dim": null,
+  "linear_value_head_dim": null,
+  "linear_a_log_min": null,
+  "linear_a_log_max": null,
+  "linear_dt_min": null,
+  "linear_dt_max": null,
+  "linear_dt_init_floor": null,
+  "linear_conv_kernel_dim": null,
+  "linear_allow_neg_eigval": null,
+  "use_qk_norm": true,
+  "qk_norm_type": "rmsnorm",
+  "attention_dropout": 0.0,
+  "resid_dropout": 0.0,
+  "embd_dropout": 0.0,
+  "initializer_range": 0.02,
+  "bos_token_id": 100257,
+  "eos_token_id": 100257,
+  "pad_token_id": 100277,
+  "use_cache": true,
+  "tie_word_embeddings": false,
+  "model_type": "lizzy",
+  "architectures": [
+    "LizzyForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_lizzy.LizzyConfig",
+    "AutoModel": "modeling_lizzy.LizzyModel",
+    "AutoModelForCausalLM": "modeling_lizzy.LizzyForCausalLM",
+    "AutoTokenizer": "tokenization_lizzy.LizzyTokenizerFast"
+  },
+  "tokenizer_class": "LizzyTokenizerFast",
+  "transformers_version": "5.4.0"
+}

configuration_lizzy.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+from typing import Any
+
+from transformers import PretrainedConfig
+
+
+class LizzyConfig(PretrainedConfig):
+    model_type = "lizzy"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+        "lm_head": "colwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 32000,
+        hidden_size: int = 4096,
+        intermediate_size: int = 11008,
+        num_hidden_layers: int = 32,
+        num_attention_heads: int = 32,
+        num_key_value_heads: int | None = None,
+        max_position_embeddings: int = 2048,
+        head_dim: int | None = None,
+        hidden_act: str = "silu",
+        norm_type: str = "rmsnorm",
+        norm_eps: float = 1e-6,
+        norm_has_bias: bool = False,
+        use_pre_attn_norm: bool = True,
+        use_pre_mlp_norm: bool = True,
+        use_post_attn_norm: bool = False,
+        use_post_mlp_norm: bool = False,
+        mlp_type: str = "gated",
+        attention_bias: bool = False,
+        mlp_bias: bool = False,
+        position_embedding_type: str = "rope",
+        rope_theta: float = 10000.0,
+        rope_scaling: dict[str, Any] | None = None,
+        rope_layer_flags: list[bool] | None = None,
+        no_rope_layer_interval: int | None = None,
+        rope_type_overrides: dict[str, str] | None = None,
+        layer_types: list[str] | None = None,
+        layer_layouts: list[str] | None = None,
+        sliding_window: int | None = None,
+        linear_num_key_heads: int | None = None,
+        linear_num_value_heads: int | None = None,
+        linear_key_head_dim: int | None = None,
+        linear_value_head_dim: int | None = None,
+        linear_a_log_min: float | None = None,
+        linear_a_log_max: float | None = None,
+        linear_dt_min: float | None = None,
+        linear_dt_max: float | None = None,
+        linear_dt_init_floor: float | None = None,
+        linear_conv_kernel_dim: int | None = None,
+        linear_allow_neg_eigval: bool | None = None,
+        use_qk_norm: bool = False,
+        qk_norm_type: str = "rmsnorm",
+        attention_dropout: float = 0.0,
+        resid_dropout: float = 0.0,
+        embd_dropout: float = 0.0,
+        initializer_range: float = 0.02,
+        bos_token_id: int | None = None,
+        eos_token_id: int | None = None,
+        pad_token_id: int | None = None,
+        use_cache: bool = True,
+        tie_word_embeddings: bool = False,
+        **kwargs,
+    ) -> None:
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+        if head_dim is None:
+            head_dim = hidden_size // num_attention_heads
+        if no_rope_layer_interval is not None:
+            no_rope_layer_interval = int(no_rope_layer_interval)
+            if no_rope_layer_interval <= 0:
+                no_rope_layer_interval = None
+        if layer_types is None:
+            layer_types = ["full_attention"] * int(num_hidden_layers)
+        if layer_layouts is None:
+            if use_post_attn_norm or use_post_mlp_norm:
+                layer_layouts = ["decoder_postnorm"] * int(num_hidden_layers)
+            else:
+                layer_layouts = ["decoder_prenorm"] * int(num_hidden_layers)
+        if rope_layer_flags is None:
+            rope_enabled = position_embedding_type == "rope"
+            if rope_enabled and no_rope_layer_interval is not None:
+                rope_layer_flags = [
+                    ((layer_idx + 1) % no_rope_layer_interval) != 0
+                    for layer_idx in range(int(num_hidden_layers))
+                ]
+            else:
+                rope_layer_flags = [rope_enabled] * int(num_hidden_layers)
+
+        normalized_rope_scaling = None
+        if rope_scaling is not None:
+            normalized_rope_scaling = dict(rope_scaling)
+            for field_name in (
+                "factor",
+                "attention_factor",
+                "beta_fast",
+                "beta_slow",
+            ):
+                if normalized_rope_scaling.get(field_name) is not None:
+                    normalized_rope_scaling[field_name] = float(
+                        normalized_rope_scaling[field_name]
+                    )
+            if (
+                normalized_rope_scaling.get("original_max_position_embeddings")
+                is not None
+            ):
+                normalized_rope_scaling["original_max_position_embeddings"] = int(
+                    normalized_rope_scaling["original_max_position_embeddings"]
+                )
+
+        # Transformers validates RoPE settings during PretrainedConfig
+        # initialization, so publish the rope-critical fields before
+        # calling `super().__init__()`.
+        self.max_position_embeddings = int(max_position_embeddings)
+        self.rope_theta = float(rope_theta)
+        self.rope_scaling = normalized_rope_scaling
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.vocab_size = int(vocab_size)
+        self.hidden_size = int(hidden_size)
+        self.intermediate_size = int(intermediate_size)
+        self.num_hidden_layers = int(num_hidden_layers)
+        self.num_attention_heads = int(num_attention_heads)
+        self.num_key_value_heads = int(num_key_value_heads)
+        self.max_position_embeddings = int(max_position_embeddings)
+        self.head_dim = int(head_dim)
+        self.hidden_act = str(hidden_act)
+        self.norm_type = str(norm_type)
+        self.norm_eps = float(norm_eps)
+        self.norm_has_bias = bool(norm_has_bias)
+        self.use_pre_attn_norm = bool(use_pre_attn_norm)
+        self.use_pre_mlp_norm = bool(use_pre_mlp_norm)
+        self.use_post_attn_norm = bool(use_post_attn_norm)
+        self.use_post_mlp_norm = bool(use_post_mlp_norm)
+        self.mlp_type = str(mlp_type)
+        self.attention_bias = bool(attention_bias)
+        self.mlp_bias = bool(mlp_bias)
+        self.position_embedding_type = str(position_embedding_type)
+        self.rope_theta = float(rope_theta)
+        self.rope_scaling = normalized_rope_scaling
+        self.no_rope_layer_interval = no_rope_layer_interval
+        self.rope_type_overrides = {
+            str(key): str(value)
+            for key, value in dict(rope_type_overrides or {}).items()
+        }
+        self.layer_types = list(layer_types)
+        self.layer_layouts = [str(item) for item in layer_layouts]
+        self.rope_layer_flags = [bool(item) for item in rope_layer_flags]
+        self.sliding_window = sliding_window
+        self.linear_num_key_heads = (
+            None
+            if linear_num_key_heads is None
+            else int(linear_num_key_heads)
+        )
+        self.linear_num_value_heads = (
+            None
+            if linear_num_value_heads is None
+            else int(linear_num_value_heads)
+        )
+        self.linear_key_head_dim = (
+            None
+            if linear_key_head_dim is None
+            else int(linear_key_head_dim)
+        )
+        self.linear_value_head_dim = (
+            None
+            if linear_value_head_dim is None
+            else int(linear_value_head_dim)
+        )
+        self.linear_a_log_min = (
+            None if linear_a_log_min is None else float(linear_a_log_min)
+        )
+        self.linear_a_log_max = (
+            None if linear_a_log_max is None else float(linear_a_log_max)
+        )
+        self.linear_dt_min = (
+            None if linear_dt_min is None else float(linear_dt_min)
+        )
+        self.linear_dt_max = (
+            None if linear_dt_max is None else float(linear_dt_max)
+        )
+        self.linear_dt_init_floor = (
+            None
+            if linear_dt_init_floor is None
+            else float(linear_dt_init_floor)
+        )
+        self.linear_conv_kernel_dim = (
+            None
+            if linear_conv_kernel_dim is None
+            else int(linear_conv_kernel_dim)
+        )
+        self.linear_allow_neg_eigval = (
+            None
+            if linear_allow_neg_eigval is None
+            else bool(linear_allow_neg_eigval)
+        )
+        self.use_qk_norm = bool(use_qk_norm)
+        self.qk_norm_type = str(qk_norm_type)
+        self.attention_dropout = float(attention_dropout)
+        self.resid_dropout = float(resid_dropout)
+        self.embd_dropout = float(embd_dropout)
+        self.initializer_range = float(initializer_range)
+        self.use_cache = bool(use_cache)
+        self.rms_norm_eps = self.norm_eps
+        self.dtype = None

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 100257,
+  "transformers_version": "4.57.3",
+  "bos_token_id": 100257,
+  "pad_token_id": 100277,
+  "do_sample": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c523315020afa40896df120a64e26c6c6ae1b8673cfd792f6ef9446dbb933d1f
+size 14596063224

modeling_lizzy.py

@@ -0,0 +1,2172 @@
+from __future__ import annotations
+
+import math
+import os
+from typing import Any, cast
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+
+try:
+    from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+except ImportError:
+    ROPE_INIT_FUNCTIONS = {}
+
+try:
+    from fla.modules import FusedRMSNormGated, ShortConvolution
+    from fla.ops.gated_delta_rule import (
+        chunk_gated_delta_rule,
+        fused_recurrent_gated_delta_rule,
+    )
+except ImportError:
+    chunk_gated_delta_rule = None
+    fused_recurrent_gated_delta_rule = None
+    FusedRMSNormGated = None
+    ShortConvolution = None
+
+from .configuration_lizzy import LizzyConfig
+
+
+class LizzyRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(dim=-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(
+            variance + self.variance_epsilon
+        )
+        return self.weight * hidden_states.to(input_dtype)
+
+
+def _make_norm(
+    norm_type: str,
+    hidden_size: int,
+    eps: float,
+    *,
+    has_bias: bool,
+) -> nn.Module:
+    if norm_type == "rmsnorm":
+        return LizzyRMSNorm(hidden_size, eps=eps)
+    if norm_type == "layernorm":
+        return nn.LayerNorm(
+            hidden_size,
+            eps=eps,
+            elementwise_affine=True,
+            bias=has_bias,
+        )
+    msg = f"Unsupported norm_type: {norm_type}"
+    raise ValueError(msg)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary_pos_emb(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    q_embed = (q * cos) + (_rotate_half(q) * sin)
+    k_embed = (k * cos) + (_rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def _legacy_cache_length(
+    past_key_values: tuple[tuple[torch.Tensor, torch.Tensor], ...] | None,
+) -> int:
+    if (
+        isinstance(past_key_values, tuple)
+        and len(past_key_values) > 0
+        and past_key_values[0] is not None
+        and past_key_values[0][0] is not None
+    ):
+        return int(past_key_values[0][0].shape[2])
+    return 0
+
+
+def _normalize_cache_position(
+    cache_position: torch.Tensor | None,
+) -> torch.Tensor | None:
+    if cache_position is None:
+        return None
+    if cache_position.dim() == 0:
+        return cache_position.view(1)
+    if cache_position.dim() > 1:
+        return cache_position[0]
+    return cache_position
+
+
+def _is_cache_object(value: Any) -> bool:
+    return isinstance(value, Cache) or isinstance(value, LizzyHybridDynamicCache)
+
+
+def _compute_default_rope_parameters(
+    config: LizzyConfig,
+    device: torch.device,
+) -> tuple[torch.Tensor, float]:
+    inv_freq = 1.0 / (
+        config.rope_theta
+        ** (
+            torch.arange(0, config.head_dim, 2, device=device, dtype=torch.float32)
+            / config.head_dim
+        )
+    )
+    return inv_freq, 1.0
+
+
+def _compute_yarn_rope_parameters(
+    config: LizzyConfig,
+    device: torch.device,
+) -> tuple[torch.Tensor, float]:
+    rope_scaling = dict(config.rope_scaling or {})
+    factor = float(rope_scaling["factor"])
+    attention_factor = rope_scaling.get("attention_factor")
+    mscale = rope_scaling.get("mscale")
+    mscale_all_dim = rope_scaling.get("mscale_all_dim")
+    original_max_position_embeddings = int(
+        rope_scaling.get("original_max_position_embeddings")
+        or config.max_position_embeddings
+    )
+
+    def get_mscale(scale: float, mscale_value: float = 1.0) -> float:
+        if scale <= 1.0:
+            return 1.0
+        return 0.1 * mscale_value * math.log(scale) + 1.0
+
+    if attention_factor is None:
+        if mscale is not None and mscale_all_dim is not None:
+            attention_factor = float(
+                get_mscale(factor, float(mscale))
+                / get_mscale(factor, float(mscale_all_dim))
+            )
+        else:
+            attention_factor = get_mscale(factor)
+
+    beta_fast = float(rope_scaling.get("beta_fast") or 32.0)
+    beta_slow = float(rope_scaling.get("beta_slow") or 1.0)
+    truncate = bool(rope_scaling.get("truncate", True))
+    dim = config.head_dim
+
+    def find_correction_dim(
+        num_rotations: float,
+        *,
+        dim: int,
+        base: float,
+        max_position_embeddings: int,
+    ) -> float:
+        return (
+            dim
+            * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))
+            / (2 * math.log(base))
+        )
+
+    def find_correction_range(
+        low_rot: float,
+        high_rot: float,
+        *,
+        dim: int,
+        base: float,
+        max_position_embeddings: int,
+        truncate: bool,
+    ) -> tuple[float, float]:
+        low = find_correction_dim(
+            low_rot,
+            dim=dim,
+            base=base,
+            max_position_embeddings=max_position_embeddings,
+        )
+        high = find_correction_dim(
+            high_rot,
+            dim=dim,
+            base=base,
+            max_position_embeddings=max_position_embeddings,
+        )
+        if truncate:
+            low = math.floor(low)
+            high = math.ceil(high)
+        return max(low, 0.0), min(high, dim - 1.0)
+
+    def linear_ramp_factor(
+        min_value: float,
+        max_value: float,
+        dim: int,
+    ) -> torch.Tensor:
+        if min_value == max_value:
+            max_value += 0.001
+        linear_func = (
+            torch.arange(dim, dtype=torch.float32, device=device) - min_value
+        ) / (max_value - min_value)
+        return torch.clamp(linear_func, 0, 1)
+
+    pos_freqs = config.rope_theta ** (
+        torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim
+    )
+    inv_freq_extrapolation = 1.0 / pos_freqs
+    inv_freq_interpolation = 1.0 / (factor * pos_freqs)
+    low, high = find_correction_range(
+        beta_fast,
+        beta_slow,
+        dim=dim,
+        base=config.rope_theta,
+        max_position_embeddings=original_max_position_embeddings,
+        truncate=truncate,
+    )
+    inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2)
+    inv_freq = (
+        inv_freq_interpolation * (1 - inv_freq_extrapolation_factor)
+        + inv_freq_extrapolation * inv_freq_extrapolation_factor
+    )
+    return inv_freq, float(attention_factor)
+
+
+def _compute_rope_parameters(
+    config: LizzyConfig,
+    device: torch.device,
+    *,
+    seq_len: int | torch.Tensor | None = None,
+    rope_type_override: str | None = None,
+) -> tuple[torch.Tensor, float]:
+    rope_scaling = dict(config.rope_scaling or {})
+    rope_type = rope_type_override
+    if rope_type is None:
+        if not rope_scaling:
+            return _compute_default_rope_parameters(config, device)
+        rope_type = str(
+            rope_scaling.get("rope_type", rope_scaling.get("type", "default"))
+        )
+
+    if rope_type == "default":
+        return _compute_default_rope_parameters(config, device)
+    if rope_type == "yarn":
+        return _compute_yarn_rope_parameters(config, device)
+    if not rope_scaling:
+        return _compute_default_rope_parameters(config, device)
+
+    rope_init_fn = (
+        ROPE_INIT_FUNCTIONS.get(rope_type) or ROPE_INIT_FUNCTIONS.get("default")
+    )
+    if rope_init_fn is None:
+        return _compute_default_rope_parameters(config, device)
+
+    inv_freq, attention_factor = rope_init_fn(config, device, seq_len=seq_len)
+    return inv_freq.to(device=device, dtype=torch.float32), float(attention_factor)
+
+
+def _looks_like_legacy_interval_rope_lizzy(config: LizzyConfig) -> bool:
+    rope_layer_flags = list(getattr(config, "rope_layer_flags", None) or [])
+    if rope_layer_flags and not all(bool(item) for item in rope_layer_flags):
+        return False
+
+    layer_types = list(getattr(config, "layer_types", None) or [])
+    if layer_types and any(str(item) != "full_attention" for item in layer_types):
+        return False
+
+    return (
+        str(getattr(config, "position_embedding_type", "")).lower() == "rope"
+        and not bool(getattr(config, "rope_scaling", None))
+        and int(getattr(config, "num_hidden_layers", 0) or 0) == 36
+        and int(getattr(config, "hidden_size", 0) or 0) == 2048
+        and int(getattr(config, "num_attention_heads", 0) or 0) == 16
+        and int(getattr(config, "num_key_value_heads", 0) or 0) == 4
+        and math.isclose(
+            float(getattr(config, "rope_theta", 0.0) or 0.0), 5_000_000.0
+        )
+        and not bool(getattr(config, "use_post_attn_norm", False))
+        and not bool(getattr(config, "use_post_mlp_norm", False))
+        and not bool(getattr(config, "use_qk_norm", False))
+    )
+
+
+def _get_no_rope_layer_interval(config: LizzyConfig) -> int | None:
+    value = getattr(config, "no_rope_layer_interval", None)
+    if value is not None:
+        value = int(value)
+        if value > 0:
+            return value
+    if _looks_like_legacy_interval_rope_lizzy(config):
+        # Backward-compatible fallback for already-uploaded Lizzy
+        # checkpoints that should use NoPE on every 4th layer.
+        return 4
+    return None
+
+
+def _get_rope_layer_flag(config: LizzyConfig, layer_idx: int) -> bool:
+    rope_enabled = str(
+        getattr(config, "position_embedding_type", "rope")
+    ).lower() == "rope"
+    rope_layer_flags = list(getattr(config, "rope_layer_flags", None) or [])
+    no_rope_layer_interval = _get_no_rope_layer_interval(config)
+    if (
+        no_rope_layer_interval is not None
+        and (
+            layer_idx >= len(rope_layer_flags)
+            or not rope_layer_flags
+            or all(bool(item) for item in rope_layer_flags)
+        )
+    ):
+        return rope_enabled and ((layer_idx + 1) % no_rope_layer_interval != 0)
+    if 0 <= layer_idx < len(rope_layer_flags):
+        return rope_enabled and bool(rope_layer_flags[layer_idx])
+    return rope_enabled
+
+
+def _get_layer_layout(config: LizzyConfig, layer_idx: int) -> str:
+    layer_layouts = list(getattr(config, "layer_layouts", None) or [])
+    if 0 <= layer_idx < len(layer_layouts):
+        return str(layer_layouts[layer_idx])
+    if bool(getattr(config, "use_post_attn_norm", False)) or bool(
+        getattr(config, "use_post_mlp_norm", False)
+    ):
+        return "decoder_postnorm"
+    return "decoder_prenorm"
+
+
+def _has_linear_attention(config: LizzyConfig) -> bool:
+    return any(
+        str(layer_type) == "linear_attention"
+        for layer_type in list(getattr(config, "layer_types", None) or [])
+    )
+
+
+class LizzyHybridDynamicCache:
+    """Cache for Lizzy checkpoints with mixed full and linear attention."""
+
+    is_compileable = False
+
+    def __init__(self, config: LizzyConfig) -> None:
+        super().__init__()
+        self.layer_types = list(config.layer_types)
+        self.transformer_layers = [
+            idx
+            for idx, layer_type in enumerate(self.layer_types)
+            if layer_type == "full_attention"
+        ]
+        self.last_linear_layer = (
+            len(self.layer_types)
+            - 1
+            - self.layer_types[::-1].index("linear_attention")
+        )
+        self.recurrent_states = [None for _ in range(config.num_hidden_layers)]
+        self.key_cache = [None for _ in range(config.num_hidden_layers)]
+        self.value_cache = [None for _ in range(config.num_hidden_layers)]
+        self.conv_states_q = [None for _ in range(config.num_hidden_layers)]
+        self.conv_states_k = [None for _ in range(config.num_hidden_layers)]
+        self.conv_states_v = [None for _ in range(config.num_hidden_layers)]
+
+    def __len__(self) -> int:
+        return len(self.layer_types)
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del cache_kwargs
+        if self.key_cache[layer_idx] is None:
+            self.key_cache[layer_idx] = key_states
+            self.value_cache[layer_idx] = value_states
+        else:
+            self.key_cache[layer_idx] = torch.cat(
+                [self.key_cache[layer_idx], key_states],
+                dim=2,
+            )
+            self.value_cache[layer_idx] = torch.cat(
+                [self.value_cache[layer_idx], value_states],
+                dim=2,
+            )
+        return self.key_cache[layer_idx], self.value_cache[layer_idx]
+
+    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
+        batch_size = beam_idx.shape[0]
+        for layer_idx in range(len(self.key_cache)):
+            if self.key_cache[layer_idx] is not None:
+                if self.key_cache[layer_idx].shape[0] < batch_size:
+                    expand_ratio = (
+                        batch_size // self.key_cache[layer_idx].shape[0]
+                    )
+                    self.key_cache[layer_idx] = (
+                        self.key_cache[layer_idx].repeat_interleave(
+                            expand_ratio, dim=0,
+                        )
+                    )
+                    self.value_cache[layer_idx] = (
+                        self.value_cache[layer_idx].repeat_interleave(
+                            expand_ratio, dim=0,
+                        )
+                    )
+                device = self.key_cache[layer_idx].device
+                self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(
+                    0,
+                    beam_idx.to(device),
+                )
+                self.value_cache[layer_idx] = (
+                    self.value_cache[layer_idx].index_select(0, beam_idx.to(device))
+                )
+            if self.conv_states_q[layer_idx] is not None:
+                if self.conv_states_q[layer_idx].shape[0] < batch_size:
+                    expand_ratio = (
+                        batch_size // self.conv_states_q[layer_idx].shape[0]
+                    )
+                    self.conv_states_q[layer_idx] = (
+                        self.conv_states_q[layer_idx].repeat_interleave(
+                            expand_ratio, dim=0,
+                        )
+                    )
+                    self.conv_states_k[layer_idx] = (
+                        self.conv_states_k[layer_idx].repeat_interleave(
+                            expand_ratio, dim=0,
+                        )
+                    )
+                    self.conv_states_v[layer_idx] = (
+                        self.conv_states_v[layer_idx].repeat_interleave(
+                            expand_ratio, dim=0,
+                        )
+                    )
+                    self.recurrent_states[layer_idx] = (
+                        self.recurrent_states[layer_idx].repeat_interleave(
+                            expand_ratio, dim=0,
+                        )
+                    )
+                device = self.conv_states_q[layer_idx].device
+                self.conv_states_q[layer_idx] = (
+                    self.conv_states_q[layer_idx].index_select(
+                        0,
+                        beam_idx.to(device),
+                    )
+                )
+                self.conv_states_k[layer_idx] = (
+                    self.conv_states_k[layer_idx].index_select(
+                        0,
+                        beam_idx.to(device),
+                    )
+                )
+                self.conv_states_v[layer_idx] = (
+                    self.conv_states_v[layer_idx].index_select(
+                        0,
+                        beam_idx.to(device),
+                    )
+                )
+                self.recurrent_states[layer_idx] = (
+                    self.recurrent_states[layer_idx].index_select(
+                        0,
+                        beam_idx.to(device),
+                    )
+                )
+
+    def get_seq_length(self, layer_idx: int | None = 0) -> int:
+        if not self.transformer_layers:
+            return 0
+        layer_idx = (
+            self.transformer_layers[0]
+            if layer_idx not in self.transformer_layers
+            else layer_idx
+        )
+        if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None:
+            return 0
+        return self.key_cache[layer_idx].shape[-2]
+
+    def get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]:
+        del layer_idx
+        kv_offset = 0
+        past_seen_tokens = self.get_seq_length()
+        kv_length = query_length + past_seen_tokens
+        return kv_length, kv_offset
+
+    @property
+    def has_previous_state(self) -> bool:
+        # Mirror the upstream contract: once the final linear layer has cached
+        # its conv state, single-token decode can switch to the recurrent path.
+        return self.conv_states_q[self.last_linear_layer] is not None
+
+
+class LizzyHybridRMSNormGated(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        gate: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if gate is None:
+            msg = "gate is required for gated RMSNorm."
+            raise ValueError(msg)
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(
+            variance + self.variance_epsilon
+        )
+        hidden_states = self.weight * hidden_states.to(input_dtype)
+        hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+        return hidden_states.to(input_dtype)
+
+
+class LizzyHybridShortConvolution(nn.Conv1d):
+    def __init__(
+        self,
+        hidden_size: int,
+        kernel_size: int,
+        bias: bool = False,
+        activation: str | None = "silu",
+    ) -> None:
+        super().__init__(
+            in_channels=hidden_size,
+            out_channels=hidden_size,
+            kernel_size=kernel_size,
+            groups=hidden_size,
+            padding=kernel_size - 1,
+            bias=bias,
+        )
+        self.hidden_size = hidden_size
+        self.conv_kernel_size = kernel_size
+        self.act_fn = ACT2FN[activation]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cache: torch.Tensor | None = None,
+        use_precomputed: bool = False,
+        **kwargs: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del kwargs
+        seq_len, dim = hidden_states.shape[-2:]
+        hidden_states = hidden_states.transpose(1, 2)
+        if use_precomputed:
+            if cache is None:
+                msg = "cache is required when use_precomputed=True."
+                raise ValueError(msg)
+            x_with_state = torch.cat([cache, hidden_states], dim=-1)
+            out = F.conv1d(
+                x_with_state,
+                self.weight,
+                self.bias,
+                padding=0,
+                groups=dim,
+            )
+            conv_state = x_with_state[:, :, 1:]
+        else:
+            out = F.conv1d(
+                hidden_states,
+                self.weight,
+                self.bias,
+                padding=self.conv_kernel_size - 1,
+                groups=dim,
+            )
+            out = out[:, :, :seq_len]
+            conv_state = F.pad(
+                hidden_states,
+                (self.conv_kernel_size - 1 - hidden_states.shape[-1], 0),
+            )
+        out = self.act_fn(out)
+        return out.transpose(1, 2), conv_state
+
+
+def _apply_mask_to_padding_states(
+    hidden_states: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+) -> torch.Tensor:
+    # Match the upstream hybrid implementation: silence padded tokens before
+    # the DeltaNet projections so recurrent state does not absorb padding.
+    if (
+        attention_mask is not None
+        and attention_mask.shape[1] > 1
+        and attention_mask.shape[0] > 1
+    ):
+        dtype = hidden_states.dtype
+        hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+    return hidden_states
+
+
+def _l2norm(
+    x: torch.Tensor,
+    dim: int = -1,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
+    return x * inv_norm
+
+
+def _torch_chunk_gated_delta_rule(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    chunk_size: int = 64,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+    use_qk_l2norm_in_kernel: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    initial_dtype = query.dtype
+    if use_qk_l2norm_in_kernel:
+        query = _l2norm(query, dim=-1, eps=1e-6)
+        key = _l2norm(key, dim=-1, eps=1e-6)
+    query, key, value, beta, g = [
+        x.transpose(1, 2).contiguous().to(torch.float32)
+        for x in (query, key, value, beta, g)
+    ]
+    batch_size, num_heads, sequence_length, k_head_dim = key.shape
+    v_head_dim = value.shape[-1]
+    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
+    query = F.pad(query, (0, 0, 0, pad_size))
+    key = F.pad(key, (0, 0, 0, pad_size))
+    value = F.pad(value, (0, 0, 0, pad_size))
+    beta = F.pad(beta, (0, pad_size))
+    g = F.pad(g, (0, pad_size))
+    total_sequence_length = sequence_length + pad_size
+    scale = 1 / (query.shape[-1] ** 0.5)
+    query = query * scale
+
+    v_beta = value * beta.unsqueeze(-1)
+    k_beta = key * beta.unsqueeze(-1)
+    query, key, value, k_beta, v_beta = [
+        x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
+        for x in (query, key, value, k_beta, v_beta)
+    ]
+    g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
+    mask = torch.triu(
+        torch.ones(
+            chunk_size,
+            chunk_size,
+            dtype=torch.bool,
+            device=query.device,
+        ),
+        diagonal=0,
+    )
+    g = g.cumsum(dim=-1)
+    decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
+    attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
+    for idx in range(1, chunk_size):
+        row = attn[..., idx, :idx].clone()
+        sub = attn[..., :idx, :idx].clone()
+        attn[..., idx, :idx] = row + (row.unsqueeze(-1) * sub).sum(-2)
+    attn = attn + torch.eye(
+        chunk_size,
+        dtype=attn.dtype,
+        device=attn.device,
+    )
+    value = attn @ v_beta
+    k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
+    last_recurrent_state = (
+        torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
+        if initial_state is None
+        else initial_state.to(value)
+    )
+    core_attn_out = torch.zeros_like(value)
+    mask = torch.triu(
+        torch.ones(
+            chunk_size,
+            chunk_size,
+            dtype=torch.bool,
+            device=query.device,
+        ),
+        diagonal=1,
+    )
+    for idx in range(0, total_sequence_length // chunk_size):
+        q_i, k_i, v_i = query[:, :, idx], key[:, :, idx], value[:, :, idx]
+        attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, idx]).masked_fill_(
+            mask,
+            0,
+        )
+        v_prime = (k_cumdecay[:, :, idx]) @ last_recurrent_state
+        v_new = v_i - v_prime
+        attn_inter = (q_i * g[:, :, idx, :, None].exp()) @ last_recurrent_state
+        core_attn_out[:, :, idx] = attn_inter + attn @ v_new
+        last_recurrent_state = (
+            last_recurrent_state * g[:, :, idx, -1, None, None].exp()
+            + (
+                k_i
+                * (g[:, :, idx, -1, None] - g[:, :, idx]).exp()[..., None]
+            ).transpose(-1, -2)
+            @ v_new
+        )
+    if not output_final_state:
+        last_recurrent_state = None
+    core_attn_out = core_attn_out.reshape(
+        core_attn_out.shape[0],
+        core_attn_out.shape[1],
+        -1,
+        core_attn_out.shape[-1],
+    )
+    core_attn_out = core_attn_out[:, :, :sequence_length]
+    core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
+    return core_attn_out, last_recurrent_state
+
+
+def _torch_recurrent_gated_delta_rule(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    initial_state: torch.Tensor | None,
+    output_final_state: bool,
+    use_qk_l2norm_in_kernel: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    initial_dtype = query.dtype
+    if use_qk_l2norm_in_kernel:
+        query = _l2norm(query, dim=-1, eps=1e-6)
+        key = _l2norm(key, dim=-1, eps=1e-6)
+    query, key, value, beta, g = [
+        x.transpose(1, 2).contiguous().to(torch.float32)
+        for x in (query, key, value, beta, g)
+    ]
+    batch_size, num_heads, sequence_length, k_head_dim = key.shape
+    v_head_dim = value.shape[-1]
+    scale = 1 / (query.shape[-1] ** 0.5)
+    query = query * scale
+    core_attn_out = torch.zeros(
+        batch_size,
+        num_heads,
+        sequence_length,
+        v_head_dim,
+    ).to(value)
+    last_recurrent_state = (
+        torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
+        if initial_state is None
+        else initial_state.to(value)
+    )
+    for idx in range(sequence_length):
+        q_t = query[:, :, idx]
+        k_t = key[:, :, idx]
+        v_t = value[:, :, idx]
+        g_t = g[:, :, idx].exp().unsqueeze(-1).unsqueeze(-1)
+        beta_t = beta[:, :, idx].unsqueeze(-1)
+        last_recurrent_state = last_recurrent_state * g_t
+        kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2)
+        delta = (v_t - kv_mem) * beta_t
+        last_recurrent_state = (
+            last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2)
+        )
+        core_attn_out[:, :, idx] = (
+            last_recurrent_state * q_t.unsqueeze(-1)
+        ).sum(dim=-2)
+    if not output_final_state:
+        last_recurrent_state = None
+    core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
+    return core_attn_out, last_recurrent_state
+
+
+class LizzyHybridGatedDeltaNet(nn.Module):
+    def __init__(self, config: LizzyConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_v_heads = config.linear_num_value_heads
+        self.num_k_heads = config.linear_num_key_heads
+        self.head_k_dim = config.linear_key_head_dim
+        self.head_v_dim = config.linear_value_head_dim
+        self.key_dim = self.head_k_dim * self.num_k_heads
+        self.value_dim = self.head_v_dim * self.num_v_heads
+        self.layer_idx = layer_idx
+        self.conv_kernel_size = config.linear_conv_kernel_dim
+        self.allow_neg_eigval = config.linear_allow_neg_eigval
+        self.eps = config.rms_norm_eps
+        self.q_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.value_dim, bias=False)
+        self.a_proj = nn.Linear(self.hidden_size, self.num_v_heads, bias=False)
+        self.b_proj = nn.Linear(self.hidden_size, self.num_v_heads, bias=False)
+        self.g_proj = nn.Linear(self.hidden_size, self.value_dim, bias=False)
+        self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
+
+        # Step-02 conversion runs on CPU by default, even on GPU nodes. In that
+        # flow Triton-backed FLA kernels will crash as soon as a CPU tensor
+        # reaches them, so the wrapper can force the pure PyTorch fallback for
+        # Hybrid layers via an environment switch.
+        disable_fla_fast_path = os.environ.get(
+            "LIZZY_DISABLE_HYBRID_FLA",
+            "",
+        ).strip().lower() in {"1", "true", "yes", "on"}
+        use_fla_fast_path = (
+            not disable_fla_fast_path
+            and
+            torch.cuda.is_available()
+            and ShortConvolution is not None
+            and chunk_gated_delta_rule is not None
+            and fused_recurrent_gated_delta_rule is not None
+            and FusedRMSNormGated is not None
+        )
+        # Keep the fast-path contract when FLA is present, but fall back to a
+        # local implementation so the public Lizzy artifact never depends on
+        # family-specific Transformers remote code.
+        conv1d_class = (
+            ShortConvolution if use_fla_fast_path else LizzyHybridShortConvolution
+        )
+        self.q_conv1d = conv1d_class(
+            hidden_size=self.key_dim,
+            kernel_size=self.conv_kernel_size,
+            bias=False,
+            activation="silu",
+        )
+        self.k_conv1d = conv1d_class(
+            hidden_size=self.key_dim,
+            kernel_size=self.conv_kernel_size,
+            bias=False,
+            activation="silu",
+        )
+        self.v_conv1d = conv1d_class(
+            hidden_size=self.value_dim,
+            kernel_size=self.conv_kernel_size,
+            bias=False,
+            activation="silu",
+        )
+
+        a = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(
+            config.linear_a_log_min,
+            config.linear_a_log_max,
+        )
+        self.A_log = nn.Parameter(torch.log(a))
+        dt = torch.exp(
+            torch.rand(self.num_v_heads)
+            * (math.log(config.linear_dt_max) - math.log(config.linear_dt_min))
+            + math.log(config.linear_dt_min)
+        )
+        dt = torch.clamp(dt, min=config.linear_dt_init_floor)
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        self.dt_bias = nn.Parameter(inv_dt)
+
+        self.o_norm = (
+            LizzyHybridRMSNormGated(self.head_v_dim, eps=1e-5)
+            if not use_fla_fast_path
+            else FusedRMSNormGated(
+                self.head_v_dim,
+                eps=1e-5,
+                device=torch.cuda.current_device(),
+                dtype=(
+                    config.dtype
+                    if config.dtype is not None
+                    else torch.get_default_dtype()
+                ),
+            )
+        )
+        self.chunk_gated_delta_rule = (
+            chunk_gated_delta_rule
+            if use_fla_fast_path
+            else _torch_chunk_gated_delta_rule
+        )
+        self.recurrent_gated_delta_rule = (
+            (
+                fused_recurrent_gated_delta_rule
+                if use_fla_fast_path
+                else _torch_recurrent_gated_delta_rule
+            )
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cache_params: LizzyHybridDynamicCache | None = None,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        del kwargs
+        hidden_states = _apply_mask_to_padding_states(hidden_states, attention_mask)
+        batch_size, seq_len, _ = hidden_states.shape
+        use_cache = cache_params is not None
+        use_precomputed = (
+            use_cache
+            and getattr(cache_params, "has_previous_state", False)
+            and seq_len == 1
+        )
+        conv_state_q = (
+            cache_params.conv_states_q[self.layer_idx] if cache_params else None
+        )
+        conv_state_k = (
+            cache_params.conv_states_k[self.layer_idx] if cache_params else None
+        )
+        conv_state_v = (
+            cache_params.conv_states_v[self.layer_idx] if cache_params else None
+        )
+        recurrent_state = (
+            cache_params.recurrent_states[self.layer_idx] if cache_params else None
+        )
+
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+        q, new_conv_state_q = self.q_conv1d(
+            q,
+            cache=conv_state_q,
+            use_precomputed=use_precomputed,
+            output_final_state=use_cache,
+        )
+        k, new_conv_state_k = self.k_conv1d(
+            k,
+            cache=conv_state_k,
+            use_precomputed=use_precomputed,
+            output_final_state=use_cache,
+        )
+        v, new_conv_state_v = self.v_conv1d(
+            v,
+            cache=conv_state_v,
+            use_precomputed=use_precomputed,
+            output_final_state=use_cache,
+        )
+
+        if cache_params is not None:
+            cache_params.conv_states_q[self.layer_idx] = new_conv_state_q
+            cache_params.conv_states_k[self.layer_idx] = new_conv_state_k
+            cache_params.conv_states_v[self.layer_idx] = new_conv_state_v
+
+        q = q.view(batch_size, seq_len, -1, self.head_k_dim)
+        k = k.view(batch_size, seq_len, -1, self.head_k_dim)
+        v = v.view(batch_size, seq_len, -1, self.head_v_dim)
+        if self.num_v_heads > self.num_k_heads:
+            expand_ratio = self.num_v_heads // self.num_k_heads
+            q = q.repeat_interleave(expand_ratio, dim=2)
+            k = k.repeat_interleave(expand_ratio, dim=2)
+
+        beta = self.b_proj(hidden_states).sigmoid()
+        if self.allow_neg_eigval:
+            beta = beta * 2.0
+        g = -self.A_log.float().exp() * F.softplus(
+            self.a_proj(hidden_states).float() + self.dt_bias
+        )
+
+        if use_precomputed:
+            output, new_recurrent_state = self.recurrent_gated_delta_rule(
+                q,
+                k,
+                v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                use_qk_l2norm_in_kernel=True,
+            )
+        else:
+            output, new_recurrent_state = self.chunk_gated_delta_rule(
+                q,
+                k,
+                v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                use_qk_l2norm_in_kernel=True,
+            )
+
+        if cache_params is not None:
+            cache_params.recurrent_states[self.layer_idx] = new_recurrent_state
+
+        gate = self.g_proj(hidden_states)
+        output = output.reshape(-1, self.head_v_dim)
+        gate = gate.reshape(-1, self.head_v_dim)
+        output = self.o_norm(output, gate)
+        output = output.reshape(batch_size, seq_len, -1)
+        output = self.o_proj(output)
+        return output
+
+
+class LizzyLinearAttention(nn.Module):
+    def __init__(self, config: LizzyConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.inner = LizzyHybridGatedDeltaNet(config, layer_idx)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_value: Cache | None = None,
+        use_cache: bool = False,
+        output_attentions: bool = False,
+        **kwargs: Any,
+    ) -> tuple[
+        torch.Tensor,
+        Cache | None,
+        torch.Tensor | None,
+    ]:
+        del kwargs, output_attentions
+        output = self.inner(
+            hidden_states=hidden_states,
+            cache_params=(
+                past_key_value if _is_cache_object(past_key_value) else None
+            ),
+            attention_mask=attention_mask,
+        )
+        present = past_key_value if use_cache else None
+        return output, present, None
+
+
+class LizzyAttention(nn.Module):
+    def __init__(self, config: LizzyConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.is_causal = True
+        self.config = config
+        self.layer_idx = layer_idx
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.hidden_size = config.hidden_size
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.position_embedding_type = config.position_embedding_type
+        self.layer_type = (
+            str(config.layer_types[layer_idx])
+            if layer_idx < len(config.layer_types)
+            else "full_attention"
+        )
+        self.use_rope = _get_rope_layer_flag(config, layer_idx)
+        self._rope_type_override = str(
+            dict(config.rope_type_overrides or {}).get(self.layer_type) or ""
+        ) or None
+        if (
+            self._rope_type_override is None
+            and self.layer_type == "sliding_attention"
+            and bool(config.rope_scaling)
+            and config.use_post_attn_norm
+            and config.use_post_mlp_norm
+            and config.use_qk_norm
+            and any(str(item) == "full_attention" for item in config.layer_types)
+        ):
+            self._rope_type_override = "default"
+        self.sliding_window = None
+        if self.layer_type == "sliding_attention":
+            self.sliding_window = config.sliding_window
+
+        q_dim = self.num_heads * self.head_dim
+        kv_dim = self.num_key_value_heads * self.head_dim
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            q_dim,
+            bias=config.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size,
+            kv_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size,
+            kv_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            q_dim,
+            config.hidden_size,
+            bias=config.attention_bias,
+        )
+        self.q_norm = (
+            _make_norm(config.qk_norm_type, q_dim, config.norm_eps, has_bias=False)
+            if config.use_qk_norm
+            else None
+        )
+        self.k_norm = (
+            _make_norm(config.qk_norm_type, kv_dim, config.norm_eps, has_bias=False)
+            if config.use_qk_norm
+            else None
+        )
+        self._rope_requires_runtime_update = False
+
+        if self.use_rope:
+            rope_scaling = dict(config.rope_scaling or {})
+            rope_type = self._rope_type_override or str(
+                rope_scaling.get("rope_type", rope_scaling.get("type", "default"))
+            )
+            self._rope_requires_runtime_update = rope_type == "dynamic"
+            if self._rope_requires_runtime_update:
+                self.register_buffer("_rope_inv_freq", None, persistent=False)
+                self.register_buffer(
+                    "_rope_attention_factor", None, persistent=False,
+                )
+            else:
+                inv_freq, attention_factor = _compute_rope_parameters(
+                    config,
+                    device=torch.device("cpu"),
+                    seq_len=config.max_position_embeddings,
+                    rope_type_override=self._rope_type_override,
+                )
+                self.register_buffer("_rope_inv_freq", inv_freq, persistent=False)
+                self.register_buffer(
+                    "_rope_attention_factor",
+                    torch.tensor(float(attention_factor), dtype=torch.float32),
+                    persistent=False,
+                )
+        else:
+            self.register_buffer("_rope_inv_freq", None, persistent=False)
+            self.register_buffer("_rope_attention_factor", None, persistent=False)
+
+    def _build_rope(
+        self,
+        position_ids: torch.Tensor,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if not self.use_rope:
+            msg = "RoPE requested but rope buffer is not initialized."
+            raise RuntimeError(msg)
+
+        inv_freq = self._rope_inv_freq
+        attention_factor_tensor = self._rope_attention_factor
+        if (
+            inv_freq is None
+            or attention_factor_tensor is None
+            or self._rope_requires_runtime_update
+        ):
+            # Keep the sequence-length hint as a tensor so TorchDynamo/vLLM
+            # can trace this path without requiring capture_scalar_outputs.
+            # When low-memory loading leaves the non-persistent cache unset,
+            # rebuild from config for this forward only instead of mutating
+            # buffers inside the compiled graph.
+            seq_len = (
+                torch.max(position_ids) + 1 if position_ids.numel() > 0 else None
+            )
+            inv_freq, attention_factor = _compute_rope_parameters(
+                self.config,
+                device=device,
+                seq_len=seq_len,
+                rope_type_override=self._rope_type_override,
+            )
+            attention_factor_tensor = torch.tensor(
+                float(attention_factor),
+                device=device,
+                dtype=torch.float32,
+            )
+        else:
+            inv_freq = inv_freq.to(device=device)
+            attention_factor_tensor = attention_factor_tensor.to(
+                device=device,
+                dtype=torch.float32,
+            )
+
+        # Mirror the upstream HF decoder-only rotary path closely here.
+        # The matmul-based construction is slightly more numerically stable
+        # than the generic einsum formulation for strict parity probes.
+        inv_freq_expanded = (
+            inv_freq[None, :, None]
+            .to(device=device, dtype=torch.float32)
+            .expand(position_ids.shape[0], -1, 1)
+        )
+        position_ids_expanded = position_ids[:, None, :].to(torch.float32)
+        angles = torch.matmul(
+            inv_freq_expanded,
+            position_ids_expanded,
+        ).transpose(1, 2)
+        angles = torch.cat((angles, angles), dim=-1)
+        cos = angles.cos().unsqueeze(1) * attention_factor_tensor
+        sin = angles.sin().unsqueeze(1) * attention_factor_tensor
+        cos = cos.to(dtype)
+        sin = sin.to(dtype)
+        return cos, sin
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        past_key_value: Cache | tuple[torch.Tensor, torch.Tensor] | None = None,
+        cache_position: torch.Tensor | None = None,
+        use_cache: bool = False,
+        output_attentions: bool = False,
+        **kwargs: Any,
+    ) -> tuple[
+        torch.Tensor,
+        Cache | tuple[torch.Tensor, torch.Tensor] | None,
+        torch.Tensor | None,
+    ]:
+        batch_size, q_len, _ = hidden_states.shape
+        cache_position = _normalize_cache_position(cache_position)
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        if self.q_norm is not None:
+            query_states = self.q_norm(query_states)
+        if self.k_norm is not None:
+            key_states = self.k_norm(key_states)
+
+        query_states = query_states.view(
+            batch_size, q_len, self.num_heads, self.head_dim,
+        )
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.view(
+            batch_size,
+            q_len,
+            self.num_key_value_heads,
+            self.head_dim,
+        )
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.view(
+            batch_size,
+            q_len,
+            self.num_key_value_heads,
+            self.head_dim,
+        )
+        value_states = value_states.transpose(1, 2)
+
+        if self.use_rope:
+            if position_ids is None:
+                msg = "position_ids are required for rope attention."
+                raise ValueError(msg)
+            cos, sin = self._build_rope(
+                position_ids, hidden_states.device, query_states.dtype,
+            )
+            query_states, key_states = _apply_rotary_pos_emb(
+                query_states,
+                key_states,
+                cos,
+                sin,
+            )
+
+        if _is_cache_object(past_key_value):
+            if use_cache:
+                key_states, value_states = past_key_value.update(
+                    key_states,
+                    value_states,
+                    self.layer_idx,
+                    cache_kwargs={"cache_position": cache_position},
+                )
+                present_key_value = past_key_value
+            elif self.layer_idx < len(past_key_value):
+                past_key, past_value = past_key_value[self.layer_idx]
+                if past_key is not None and past_value is not None:
+                    key_states = torch.cat([past_key, key_states], dim=2)
+                    value_states = torch.cat([past_value, value_states], dim=2)
+                present_key_value = None
+            else:
+                present_key_value = None
+        elif past_key_value is not None:
+            past_key, past_value = past_key_value
+            key_states = torch.cat([past_key, key_states], dim=2)
+            value_states = torch.cat([past_value, value_states], dim=2)
+            present_key_value = (key_states, value_states) if use_cache else None
+        else:
+            present_key_value = (key_states, value_states) if use_cache else None
+
+        attention_interface = None
+        attn_impl = getattr(self.config, "_attn_implementation", "eager")
+        if attn_impl == "flex_attention" and self.head_dim < 16:
+            attn_impl = "sdpa"
+        if attn_impl != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS.get(attn_impl)
+
+        if attention_interface is not None:
+            attn_output, attn_weights = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                scaling=self.scaling,
+                sliding_window=self.sliding_window,
+                **kwargs,
+            )
+            attn_output = attn_output.contiguous()
+        else:
+            if self.num_key_value_heads != self.num_heads:
+                key_states = key_states.repeat_interleave(
+                    self.num_key_value_groups, dim=1,
+                )
+                value_states = value_states.repeat_interleave(
+                    self.num_key_value_groups, dim=1,
+                )
+            attn_weights = torch.matmul(
+                query_states,
+                key_states.transpose(-1, -2),
+            ) * self.scaling
+            if attention_mask is not None:
+                attn_weights = attn_weights + attention_mask
+            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32)
+            attn_weights = attn_weights.to(query_states.dtype)
+            attn_weights = F.dropout(
+                attn_weights,
+                p=self.attention_dropout if self.training else 0.0,
+                training=self.training,
+            )
+            attn_output = torch.matmul(attn_weights, value_states)
+            attn_output = attn_output.transpose(1, 2).contiguous()
+
+        attn_output = attn_output.reshape(batch_size, q_len, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        if not output_attentions:
+            attn_weights = None
+        return attn_output, present_key_value, attn_weights
+
+
+def _refresh_attention_rope_buffers(module: nn.Module) -> None:
+    """Rebuild non-persistent RoPE buffers after checkpoint load."""
+
+    for child in module.modules():
+        if not isinstance(child, LizzyAttention):
+            continue
+
+        should_use_rope = _get_rope_layer_flag(child.config, child.layer_idx)
+        child.use_rope = should_use_rope
+        if not should_use_rope:
+            child._rope_requires_runtime_update = False
+            child._rope_inv_freq = None
+            child._rope_attention_factor = None
+            continue
+
+        rope_scaling = dict(child.config.rope_scaling or {})
+        rope_type = child._rope_type_override or str(
+            rope_scaling.get("rope_type", rope_scaling.get("type", "default"))
+        )
+        child._rope_requires_runtime_update = rope_type == "dynamic"
+        if child._rope_requires_runtime_update:
+            child._rope_inv_freq = None
+            child._rope_attention_factor = None
+            continue
+
+        # These buffers are derived from config rather than serialized weights.
+        # Recompute them after load so low-memory materialization cannot leave
+        # stale or uninitialized rotary state behind.
+        inv_freq, attention_factor = _compute_rope_parameters(
+            child.config,
+            device=torch.device("cpu"),
+            seq_len=child.config.max_position_embeddings,
+            rope_type_override=child._rope_type_override,
+        )
+        child._rope_inv_freq = inv_freq
+        child._rope_attention_factor = torch.tensor(
+            float(attention_factor),
+            dtype=torch.float32,
+        )
+
+
+class LizzyMLP(nn.Module):
+    def __init__(self, config: LizzyConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.act = ACT2FN[config.hidden_act]
+        self.gate_proj = (
+            nn.Linear(
+                config.hidden_size,
+                config.intermediate_size,
+                bias=config.mlp_bias,
+            )
+            if config.mlp_type == "gated"
+            else None
+        )
+        self.up_proj = nn.Linear(
+            config.hidden_size,
+            config.intermediate_size,
+            bias=config.mlp_bias,
+        )
+        self.down_proj = nn.Linear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=config.mlp_bias,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.gate_proj is None and self.config.mlp_type == "gated":
+            msg = "Missing gated MLP projection layers."
+            raise RuntimeError(msg)
+        if self.config.mlp_type == "gated":
+            if self.gate_proj is None:
+                msg = "Missing gated MLP projection layers."
+                raise RuntimeError(msg)
+            return self.down_proj(self.act(
+                self.gate_proj(hidden_states)) * self.up_proj(hidden_states)
+            )
+        return self.down_proj(self.act(self.up_proj(hidden_states)))
+
+
+class LizzyDecoderLayer(nn.Module):
+    def __init__(self, config: LizzyConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.layer_type = (
+            str(config.layer_types[layer_idx])
+            if layer_idx < len(config.layer_types)
+            else "full_attention"
+        )
+        self.layer_layout = _get_layer_layout(config, layer_idx)
+        self.self_attn = (
+            LizzyAttention(config, layer_idx)
+            if self.layer_type != "linear_attention"
+            else None
+        )
+        self.linear_attn = (
+            LizzyLinearAttention(config, layer_idx)
+            if self.layer_type == "linear_attention"
+            else None
+        )
+        self.mlp = LizzyMLP(config)
+
+        self.pre_attn_norm = (
+            _make_norm(
+                config.norm_type,
+                config.hidden_size,
+                config.norm_eps,
+                has_bias=config.norm_has_bias,
+            )
+            if self.layer_layout == "decoder_prenorm"
+            else None
+        )
+        self.pre_mlp_norm = (
+            _make_norm(
+                config.norm_type,
+                config.hidden_size,
+                config.norm_eps,
+                has_bias=config.norm_has_bias,
+            )
+            if self.layer_layout == "decoder_prenorm"
+            else None
+        )
+        self.post_attn_norm = (
+            _make_norm(
+                config.norm_type,
+                config.hidden_size,
+                config.norm_eps,
+                has_bias=config.norm_has_bias,
+            )
+            if self.layer_layout == "decoder_postnorm"
+            else None
+        )
+        self.post_mlp_norm = (
+            _make_norm(
+                config.norm_type,
+                config.hidden_size,
+                config.norm_eps,
+                has_bias=config.norm_has_bias
+            )
+            if self.layer_layout == "decoder_postnorm"
+            else None
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        past_key_value: Cache | tuple[torch.Tensor, torch.Tensor] | None = None,
+        cache_position: torch.Tensor | None = None,
+        use_cache: bool = False,
+        output_attentions: bool = False,
+        **kwargs: Any,
+    ) -> tuple[
+        torch.Tensor,
+        Cache | tuple[torch.Tensor, torch.Tensor] | None,
+        torch.Tensor | None,
+    ]:
+        residual = hidden_states
+        attn_inputs = (
+            self.pre_attn_norm(hidden_states)
+            if self.pre_attn_norm is not None
+            else hidden_states
+        )
+        if self.linear_attn is not None:
+            attn_output, present_key_value, attn_weights = self.linear_attn(
+                attn_inputs,
+                attention_mask=attention_mask,
+                past_key_value=(
+                    past_key_value if _is_cache_object(past_key_value) else None
+                ),
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                **kwargs,
+            )
+        else:
+            assert self.self_attn is not None
+            attn_output, present_key_value, attn_weights = self.self_attn(
+                attn_inputs,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                cache_position=cache_position,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                **kwargs,
+            )
+        if self.post_attn_norm is not None:
+            attn_output = self.post_attn_norm(attn_output)
+        hidden_states = residual + attn_output
+
+        residual = hidden_states
+        mlp_inputs = (
+            self.pre_mlp_norm(hidden_states)
+            if self.pre_mlp_norm is not None
+            else hidden_states
+        )
+        mlp_output = self.mlp(mlp_inputs)
+        if self.post_mlp_norm is not None:
+            mlp_output = self.post_mlp_norm(mlp_output)
+        hidden_states = residual + mlp_output
+        return hidden_states, present_key_value, attn_weights
+
+
+class LizzyPreTrainedModel(PreTrainedModel):
+    config_class = LizzyConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LizzyDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module: nn.Module) -> None:
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, (LizzyRMSNorm, nn.LayerNorm)):
+            if hasattr(module, "weight") and module.weight is not None:
+                module.weight.data.fill_(1.0)
+            if hasattr(module, "bias") and module.bias is not None:
+                module.bias.data.zero_()
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str | os.PathLike[str] | None,
+        *model_args: Any,
+        **kwargs: Any,
+    ) -> "LizzyPreTrainedModel":
+        model = cast(
+            "LizzyPreTrainedModel",
+            super().from_pretrained(
+                pretrained_model_name_or_path,
+                *model_args,
+                **kwargs,
+            ),
+        )
+        _refresh_attention_rope_buffers(model)
+        if hasattr(model, "lm_head") and hasattr(model, "model"):
+            tied_weights_keys = getattr(type(model), "_tied_weights_keys", None)
+            if isinstance(tied_weights_keys, dict) and tied_weights_keys:
+                model._tied_weights_keys = dict(tied_weights_keys)
+            else:
+                model._tied_weights_keys = {
+                    "lm_head.weight": "model.embed_tokens.weight",
+                }
+            model._tp_plan = {"lm_head": "colwise_rep"}
+            model._pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+        return model
+
+    def load_state_dict(  # type: ignore[override]
+        self,
+        state_dict: dict[str, torch.Tensor],
+        strict: bool = True,
+        assign: bool = False,
+    ) -> Any:
+        remapped_state_dict: dict[str, torch.Tensor] = {}
+        for key, value in state_dict.items():
+            remapped_key = key
+            if ".mlp.fc_in." in key:
+                remapped_key = key.replace(".mlp.fc_in.", ".mlp.up_proj.")
+            elif ".mlp.fc_out." in key:
+                remapped_key = key.replace(".mlp.fc_out.", ".mlp.down_proj.")
+
+            existing = remapped_state_dict.get(remapped_key)
+            if existing is not None and not torch.equal(existing, value):
+                msg = (
+                    f"Conflicting legacy Lizzy MLP tensors"
+                    f" for key: {remapped_key}"
+                )
+                raise ValueError(msg)
+            remapped_state_dict[remapped_key] = value
+
+        load_result = super().load_state_dict(
+            remapped_state_dict,
+            strict=strict,
+            assign=assign,
+        )
+        # RoPE buffers are intentionally non-persistent, so refresh them after
+        # weight loading instead of trusting constructor-time allocations.
+        _refresh_attention_rope_buffers(self)
+        return load_result
+
+
+class LizzyModel(LizzyPreTrainedModel):
+    def __init__(self, config: LizzyConfig) -> None:
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size,
+            config.hidden_size,
+            self.padding_idx,
+        )
+        self.embed_positions = (
+            nn.Embedding(config.max_position_embeddings, config.hidden_size)
+            if config.position_embedding_type == "absolute"
+            else None
+        )
+        self.layers = nn.ModuleList(
+            LizzyDecoderLayer(config, layer_idx)
+            for layer_idx in range(config.num_hidden_layers)
+        )
+        self.norm = _make_norm(
+            config.norm_type,
+            config.hidden_size,
+            config.norm_eps,
+            has_bias=config.norm_has_bias,
+        )
+        self.embd_dropout = nn.Dropout(config.embd_dropout)
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.embed_tokens = value
+
+    def _build_attention_mask(
+        self,
+        attention_mask: torch.Tensor | None,
+        *,
+        batch_size: int,
+        q_len: int,
+        kv_len: int,
+        kv_offset: int,
+        cache_position: torch.Tensor,
+        device: torch.device,
+        dtype: torch.dtype,
+        sliding_window: int | None = None,
+    ) -> torch.Tensor:
+        kv_len = (
+            int(kv_len.item()) if isinstance(kv_len, torch.Tensor) else int(kv_len)
+        )
+        kv_offset = (
+            int(kv_offset.item())
+            if isinstance(kv_offset, torch.Tensor)
+            else int(kv_offset)
+        )
+        min_value = torch.finfo(dtype).min
+        source_positions = cache_position.to(device=device).view(-1, 1)
+        target_positions = torch.arange(
+            kv_offset,
+            kv_offset + kv_len,
+            device=device,
+        ).unsqueeze(0)
+        causal = torch.zeros((q_len, kv_len), dtype=dtype, device=device)
+        causal = causal.masked_fill(target_positions > source_positions, min_value)
+        if sliding_window is not None:
+            lower_bound = source_positions - int(sliding_window) + 1
+            causal = causal.masked_fill(target_positions < lower_bound, min_value)
+        causal = causal.unsqueeze(0).unsqueeze(0).expand(batch_size, 1, -1, -1)
+
+        if attention_mask is None:
+            return causal
+
+        if attention_mask.dim() != 2:
+            msg = "attention_mask must be 2D [batch, sequence]."
+            raise ValueError(msg)
+        if attention_mask.shape[1] < kv_len:
+            pad = torch.ones(
+                (attention_mask.shape[0], kv_len - attention_mask.shape[1]),
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+            attention_mask = torch.cat([pad, attention_mask], dim=1)
+        elif attention_mask.shape[1] > kv_len:
+            attention_mask = attention_mask[:, -kv_len:]
+
+        expanded = attention_mask[:, None, None, :].to(device=device)
+        padding = (expanded == 0).to(dtype) * min_value
+        return causal + padding
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | tuple[
+            tuple[torch.Tensor, torch.Tensor], ...
+        ] | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs: Any,
+    ) -> BaseModelOutputWithPast | tuple[Any, ...]:
+        if (input_ids is None) == (inputs_embeds is None):
+            msg = "Exactly one of input_ids or inputs_embeds must be provided."
+            raise ValueError(msg)
+
+        output_attentions = (
+            bool(output_attentions) if output_attentions is not None else False
+        )
+        output_hidden_states = (
+            bool(output_hidden_states)
+            if output_hidden_states is not None
+            else False
+        )
+        use_cache = (
+            bool(use_cache)
+            if use_cache is not None
+            else bool(self.config.use_cache)
+        )
+        return_dict = bool(return_dict) if return_dict is not None else True
+
+        if inputs_embeds is None:
+            hidden_states = self.embed_tokens(input_ids)
+            batch_size, seq_len = input_ids.shape
+        else:
+            hidden_states = inputs_embeds
+            batch_size, seq_len, _ = inputs_embeds.shape
+
+        cache_object = (
+            past_key_values
+            if _is_cache_object(past_key_values)
+            else None
+        )
+        if use_cache and _has_linear_attention(self.config):
+            # Transformers 5.4 seeds `generate()` with an empty DynamicCache
+            # for standard causal decoders. Hybrid Lizzy checkpoints need the
+            # mixed cache below instead, because linear-attention layers read
+            # DeltaNet convolution/recurrent state during the prefill pass.
+            if cache_object is not None and not isinstance(
+                cache_object, LizzyHybridDynamicCache,
+            ):
+                if int(cache_object.get_seq_length()) > 0:
+                    msg = (
+                        "Hybrid Lizzy checkpoints require "
+                        "LizzyHybridDynamicCache once generation cache "
+                        "state is populated."
+                    )
+                    raise ValueError(msg)
+                cache_object = LizzyHybridDynamicCache(config=self.config)
+                past_key_values = cache_object
+        if use_cache and cache_object is None and past_key_values is None:
+            if _has_linear_attention(self.config):
+                # Linear-attention checkpoints need a mixed cache that can hold
+                # both KV tensors and recurrent DeltaNet state.
+                cache_object = LizzyHybridDynamicCache(config=self.config)
+            else:
+                cache_object = DynamicCache()
+            past_key_values = cache_object
+
+        if cache_object is not None:
+            past_length = int(cache_object.get_seq_length())
+        else:
+            past_length = _legacy_cache_length(past_key_values)
+
+        cache_position = _normalize_cache_position(cache_position)
+        if cache_position is None:
+            cache_position = torch.arange(
+                past_length,
+                past_length + seq_len,
+                dtype=torch.long,
+                device=hidden_states.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0).expand(batch_size, -1)
+
+        if self.embed_positions is not None:
+            hidden_states = hidden_states + self.embed_positions(position_ids)
+
+        hidden_states = self.embd_dropout(hidden_states)
+
+        if self.training and self.gradient_checkpointing:
+            use_cache = False
+
+        layer_types = list(self.config.layer_types)
+        if not layer_types:
+            layer_types = ["full_attention"] * len(self.layers)
+        has_linear_attention = any(
+            str(layer_type) == "linear_attention" for layer_type in layer_types
+        )
+
+        _attn_impl = getattr(self.config, "_attn_implementation", "eager")
+        if has_linear_attention and isinstance(attention_mask, dict):
+            linear_attention_mask = attention_mask.get("linear_attention")
+        else:
+            linear_attention_mask = attention_mask
+            if (
+                has_linear_attention
+                and cache_object is not None
+                and getattr(cache_object, "has_previous_state", False)
+            ):
+                linear_attention_mask = None
+            elif (
+                has_linear_attention
+                and attention_mask is not None
+                and not isinstance(attention_mask, dict)
+                and torch.all(attention_mask == 1)
+            ):
+                linear_attention_mask = None
+
+        if (
+            _attn_impl == "flash_attention_2"
+            and not isinstance(attention_mask, dict)
+        ):
+            # Flash attention handles causal masking (via is_causal) and
+            # padding (via 2D mask) natively; skip building a 4D mask.
+            attention_mask_mapping = {
+                lt: attention_mask
+                for lt in dict.fromkeys(layer_types)
+                if lt != "linear_attention"
+            }
+        elif _attn_impl == "sdpa" and attention_mask is None:
+            attention_mask_mapping = {}
+            for layer_type in dict.fromkeys(layer_types):
+                if layer_type == "linear_attention":
+                    continue
+                if layer_type == "full_attention":
+                    # Match upstream decoder-only HF models: when SDPA sees
+                    # plain causal full attention with no padding mask to
+                    # preserve, let it use its native is_causal fast-path
+                    # instead of forcing an explicit 4D bias tensor.
+                    attention_mask_mapping[layer_type] = None
+                    continue
+                layer_idx = layer_types.index(layer_type)
+                if cache_object is not None:
+                    kv_len, kv_offset = cache_object.get_mask_sizes(
+                        seq_len, layer_idx,
+                    )
+                else:
+                    kv_len = past_length + seq_len
+                    kv_offset = 0
+                attention_mask_mapping[layer_type] = self._build_attention_mask(
+                    attention_mask,
+                    batch_size=batch_size,
+                    q_len=seq_len,
+                    kv_len=kv_len,
+                    kv_offset=kv_offset,
+                    cache_position=cache_position,
+                    device=hidden_states.device,
+                    dtype=hidden_states.dtype,
+                    sliding_window=(
+                        self.config.sliding_window
+                        if layer_type == "sliding_attention"
+                        else None
+                    ),
+                )
+        elif isinstance(attention_mask, dict):
+            attention_mask_mapping = {
+                key: value
+                for key, value in attention_mask.items()
+                if key != "linear_attention"
+            }
+        else:
+            attention_mask_mapping: dict[str, torch.Tensor] = {}
+            for layer_type in dict.fromkeys(layer_types):
+                if layer_type == "linear_attention":
+                    continue
+                layer_idx = layer_types.index(layer_type)
+                if cache_object is not None:
+                    kv_len, kv_offset = cache_object.get_mask_sizes(
+                        seq_len, layer_idx,
+                    )
+                else:
+                    kv_len = past_length + seq_len
+                    kv_offset = 0
+                attention_mask_mapping[layer_type] = self._build_attention_mask(
+                    attention_mask,
+                    batch_size=batch_size,
+                    q_len=seq_len,
+                    kv_len=kv_len,
+                    kv_offset=kv_offset,
+                    cache_position=cache_position,
+                    device=hidden_states.device,
+                    dtype=hidden_states.dtype,
+                    sliding_window=(
+                        self.config.sliding_window
+                        if layer_type == "sliding_attention"
+                        else None
+                    ),
+                )
+
+        all_hidden_states = [] if output_hidden_states else None
+        all_attentions = [] if output_attentions else None
+        next_cache = (
+            cache_object
+            if cache_object is not None
+            else ([] if use_cache else None)
+        )
+        gradient_checkpointing_func = getattr(
+            self,
+            "_gradient_checkpointing_func",
+            checkpoint,
+        )
+
+        for idx, layer in enumerate(self.layers):
+            if output_hidden_states and all_hidden_states is not None:
+                all_hidden_states.append(hidden_states)
+            layer_type = (
+                layer_types[idx]
+                if idx < len(layer_types)
+                else "full_attention"
+            )
+            if layer_type == "linear_attention":
+                layer_attention_mask = linear_attention_mask
+            else:
+                layer_attention_mask = attention_mask_mapping[layer_type]
+
+            if cache_object is not None:
+                layer_past: Cache | tuple[
+                    torch.Tensor, torch.Tensor
+                ] | None = cache_object
+            elif past_key_values is not None:
+                layer_past = past_key_values[idx]
+                if layer_past is not None and layer_past[0] is None:
+                    layer_past = None
+            else:
+                layer_past = None
+
+            if self.training and self.gradient_checkpointing:
+                def custom_forward(hidden_states: torch.Tensor) -> Any:
+                    layer_outputs = layer(
+                        hidden_states,
+                        attention_mask=layer_attention_mask,
+                        position_ids=position_ids,
+                        past_key_value=None,
+                        cache_position=cache_position,
+                        use_cache=False,
+                        output_attentions=output_attentions,
+                        **kwargs,
+                    )
+                    if output_attentions:
+                        return layer_outputs[0], layer_outputs[2]
+                    return layer_outputs[0]
+
+                checkpointed_outputs = gradient_checkpointing_func(
+                    custom_forward, hidden_states,
+                )
+                if output_attentions:
+                    hidden_states, attn_weights = checkpointed_outputs
+                else:
+                    hidden_states = checkpointed_outputs
+                    attn_weights = None
+                present = None
+            else:
+                hidden_states, present, attn_weights = layer(
+                    hidden_states,
+                    attention_mask=layer_attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=layer_past,
+                    cache_position=cache_position,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                    **kwargs,
+                )
+                if use_cache and next_cache is not None and cache_object is None:
+                    next_cache.append(present)
+
+            if output_attentions and all_attentions is not None:
+                all_attentions.append(attn_weights)
+
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states and all_hidden_states is not None:
+            all_hidden_states.append(hidden_states)
+
+        past_key_values_output: Cache | tuple[
+            tuple[torch.Tensor, torch.Tensor], ...
+        ] | None = None
+        if use_cache and next_cache is not None:
+            if cache_object is not None:
+                past_key_values_output = cache_object
+            else:
+                past_key_values_output = tuple(next_cache)
+
+        if not return_dict:
+            output: tuple[Any, ...] = (hidden_states,)
+            if past_key_values_output is not None:
+                output = output + (past_key_values_output,)
+            if output_hidden_states and all_hidden_states is not None:
+                output = output + (tuple(all_hidden_states),)
+            if output_attentions and all_attentions is not None:
+                output = output + (tuple(all_attentions),)
+            return output
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values_output,
+            hidden_states=(
+                tuple(all_hidden_states)
+                if all_hidden_states is not None
+                else None
+            ),
+            attentions=(
+                tuple(all_attentions)
+                if all_attentions is not None
+                else None
+            ),
+        )
+
+
+class LizzyForCausalLM(LizzyPreTrainedModel, GenerationMixin):
+    config_class = LizzyConfig
+    # Transformers 5.4 expects an expanded target->source mapping here rather than
+    # the older list-based shorthand.
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config: LizzyConfig) -> None:
+        super().__init__(config)
+        self.model = LizzyModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.model.set_input_embeddings(value)
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Cache | tuple[
+            tuple[torch.Tensor, torch.Tensor], ...
+        ] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        past_length = 0
+        if past_key_values is not None:
+            if _is_cache_object(past_key_values):
+                past_length = int(past_key_values.get_seq_length())
+            else:
+                past_length = _legacy_cache_length(past_key_values)
+
+        cache_position = _normalize_cache_position(cache_position)
+        if cache_position is None:
+            if past_key_values is not None:
+                new_tokens = input_ids.shape[1] - past_length
+                if new_tokens <= 0:
+                    new_tokens = 1
+                cache_position = torch.arange(
+                    past_length,
+                    past_length + new_tokens,
+                    device=input_ids.device,
+                )
+            else:
+                cache_position = torch.arange(
+                    input_ids.shape[1],
+                    device=input_ids.device,
+                )
+
+        if past_key_values is not None:
+            input_ids = input_ids[:, -cache_position.shape[0] :]
+            if attention_mask is not None:
+                attn_mask_idx = (past_length + input_ids.shape[1])
+                attention_mask = attention_mask[:, -attn_mask_idx :]
+
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs: dict[str, Any] = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids.contiguous()}
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "use_cache": kwargs.get("use_cache", self.config.use_cache),
+            },
+        )
+        return model_inputs
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | tuple[
+            tuple[torch.Tensor, torch.Tensor], ...
+        ] | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Any,
+    ) -> CausalLMOutputWithPast | tuple[Any, ...]:
+        # HF eval loaders call `forward()` without an explicit return_dict,
+        # so local Lizzy exports must normalize the optional flag first.
+        return_dict = bool(return_dict) if return_dict is not None else True
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
+        slice_indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        if labels is not None:
+            full_logits = self.lm_head(hidden_states)
+            logits = full_logits[:, slice_indices, :]
+        else:
+            full_logits = None
+            logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            shift_logits = full_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+            )
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            if loss is not None:
+                output = (loss,) + output
+            return output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

quickstart.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+"""Minimal inference example for the private Lizzy 7B checkpoint."""
+
+from __future__ import annotations
+
+import os
+
+
+def main() -> None:
+    repo_id = os.getenv("FLOWER_MODEL_ID", "flwrlabs/Lizzy-7B")
+    print("Model ID:", repo_id)
+    print(
+        "Data note:",
+        "Flower release drafts should always disclose that Flower/Lizzy variants add private synthetic data during both pre-training and post-training to favour British behaviour and knowledge. Those private synthetic datasets are not redistributed in the release pack.",
+    )
+    print("HF_TOKEN present:", bool(os.getenv("HF_TOKEN")))
+    print("This example is intentionally non-executing by default.")
+    print("Use one of the snippets below after installing transformers or vLLM:")
+    print()
+    print("Transformers:")
+    print(
+        "  tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)"
+    )
+    print(
+        "  model = AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True, torch_dtype='auto')"
+    )
+    print()
+    print("vLLM:")
+    print(
+        "  python -m vllm.entrypoints.openai.api_server --model "
+        "flwrlabs/Lizzy-7B --trust-remote-code --max-model-len 8192"
+    )
+
+
+if __name__ == "__main__":
+    main()

quickstart_cli.sh

@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+MODEL_ID="${FLOWER_MODEL_ID:-flwrlabs/Lizzy-7B}"
+
+echo "HF_TOKEN set: ${HF_TOKEN:+yes}"
+echo "Model: $MODEL_ID"
+echo "Data note: Flower release drafts should always disclose that Flower/Lizzy variants add private synthetic data during both pre-training and post-training to favour British behaviour and knowledge. Those private synthetic datasets are not redistributed in the release pack."
+echo
+echo "Transformers example:"
+echo "python - <<'PY'"
+echo "from transformers import AutoTokenizer, AutoModelForCausalLM"
+echo "repo_id = 'flwrlabs/Lizzy-7B'"
+echo "tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)"
+echo "model = AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True, torch_dtype='auto')"
+echo "PY"
+echo
+echo "vLLM serve example:"
+echo "python -m vllm.entrypoints.openai.api_server --model $MODEL_ID --trust-remote-code --max-model-len 8192"

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|pad|>",
+  "unk_token": "<|endoftext|>"
+}

tokenization_lizzy.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from transformers import PreTrainedTokenizerFast
+
+
+class LizzyTokenizerFast(PreTrainedTokenizerFast):
+    """Family-agnostic fast tokenizer wrapper for Lizzy checkpoints."""
+
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        preserved_keys = (
+            "add_prefix_space",
+            "add_bos_token",
+            "add_eos_token",
+            "clean_up_tokenization_spaces",
+            "use_default_system_prompt",
+            "legacy",
+            "fix_mistral_regex",
+        )
+        preserved_init_attrs = {
+            key: kwargs.get(key)
+            for key in preserved_keys
+            if key in kwargs
+        }
+        super().__init__(*args, **kwargs)
+        init_kwargs = getattr(self, "init_kwargs", {})
+        local_payload: dict[str, Any] = {}
+        config_path = (
+            Path(str(getattr(self, "name_or_path", ""))) / "tokenizer_config.json"
+        )
+        if config_path.is_file():
+            try:
+                local_payload = json.loads(config_path.read_text(encoding="utf-8"))
+            except Exception:
+                local_payload = {}
+        for key in preserved_keys:
+            value = preserved_init_attrs.get(key, init_kwargs.get(key))
+            if value is None:
+                value = local_payload.get(key)
+            if value is not None:
+                setattr(self, key, value)
+
+    @property
+    def all_special_tokens_extended(self) -> list[str]:
+        """Compatibility shim for runtimes still expecting the pre-5.4 API."""
+        return list(self.all_special_tokens)

tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "chat_template": "\\\n    {# \u2500\u2500\u2500\u2500\u2500 defaults \u2500\u2500\u2500\u2500\u2500 #}\n    {%- if enable_thinking is not defined -%}\n    {%- set enable_thinking = true -%}\n    {%- endif -%}\n    {%- set system_message = \"\" -%}\n    {%- set custom_instructions = \"\" -%}\n    {%- set identity_preamble = \"You are Lizzy, a helpful British AI assistant built by Flower Labs. When asked about your identity, name, developer, or origin, answer consistently: your name is Lizzy, you were built by Flower Labs, and you are not ChatGPT, DeepSeek, Claude, Gemini, or any other assistant. Do not misattribute your creator or model family.\" -%}\n    {%- set default_think_instructions = \"You are a helpful British function-calling AI assistant. You are a British persona and your date cutoff is November 2024, and your model weights are available at https://huggingface.co/flwrlabs. Your role as an assistant involves thoroughly exploring questions through a systematic thinking process before providing the final precise and accurate solutions. This requires engaging in a comprehensive cycle of analysis, summarizing, exploration, reassessment, reflection, backtracking, and iteration to develop well-considered thinking process. Please structure your response into two main sections: Thought and Solution using the specified format: <tool_call> Thought section </tool_call> Solution section. In the Thought section, detail your reasoning process in steps. Each step should include detailed considerations such as analysing questions, summarizing relevant findings, brainstorming new ideas, verifying the accuracy of the current steps, refining any errors, and revisiting previous steps. In the Solution section, based on various attempts, explorations, and reflections from the Thought section, systematically present the final solution that you deem correct. The Solution section should be logical, accurate, and concise and detail necessary steps needed to reach the conclusion.\" -%}\n    {%- set default_no_think_instructions = \"You are a helpful British function-calling AI assistant. You are a British persona and your date cutoff is November 2024, and your model weights are available at https://huggingface.co/flwrlabs.\" -%}\n\n    {# \u2500\u2500\u2500\u2500\u2500 reasoning mode \u2500\u2500\u2500\u2500\u2500 #}\n    {%- if enable_thinking -%}\n      {%- set reasoning_mode = \"/think\" -%}\n    {%- else -%}\n      {%- set reasoning_mode = \"/no_think\" -%}\n    {%- endif -%}\n\n    {# \u2500\u2500\u2500\u2500\u2500 header (system message) \u2500\u2500\u2500\u2500\u2500 #}\n    {{- \"<|im_start|>system\\n\" -}}\n\n    {%- if messages[0].role == \"system\" -%}\n      {%- set system_message = messages[0].content -%}\n      {%- if \"/no_think\" in system_message -%}\n        {%- set reasoning_mode = \"/no_think\" -%}\n      {%- elif \"/think\" in system_message -%}\n        {%- set reasoning_mode = \"/think\" -%}\n      {%- endif -%}\n      {%- set custom_instructions = system_message.replace(\"/no_think\", \"\").replace(\"/think\", \"\").rstrip() -%}\n    {%- endif -%}\n\n    {%- if \"/system_override\" in system_message -%}\n      {{- identity_preamble + \"\\n\\n\" -}}\n      {{- custom_instructions.replace(\"/system_override\", \"\").rstrip() -}}\n      {{- \"<|im_end|>\\n\" -}}\n    {%- else -%}\n      {{- \"## Metadata\\n\\n\" -}}\n      {{- \"Knowledge Cutoff Date: June 2025\\n\" -}}\n      {%- set today = strftime_now(\"%d %B %Y\") -%}\n      {{- \"Today Date: \" ~ today ~ \"\\n\" -}}\n      {{- \"Reasoning Mode: \" + reasoning_mode + \"\\n\\n\" -}}\n\n      {{- \"## Identity\\n\\n\" -}}\n      {{- identity_preamble + \"\\n\\n\" -}}\n\n      {{- \"## Custom Instructions\\n\\n\" -}}\n      {%- if custom_instructions -%}\n        {{- custom_instructions + \"\\n\\n\" -}}\n      {%- elif reasoning_mode == \"/think\" -%}\n        {{- default_think_instructions + \"\\n\\n\" -}}\n      {%- else -%}\n        {{- default_no_think_instructions + \"\\n\\n\" -}}\n      {%- endif -%}\n\n      {%- if xml_tools or python_tools or tools -%}\n        {{- \"### Tools\\n\\n\" -}}\n        {%- if xml_tools or tools -%}\n          {%- if tools -%}\n            {%- set xml_tools = tools -%}\n          {%- endif -%}\n          {%- set ns = namespace(xml_tool_string=\"You may call one or more functions to assist with the user query.\\nYou are provided with function signatures within <tools></tools> XML tags:\\n\\n<tools>\\n\") -%}\n          {%- for tool in xml_tools[:] -%} {# The slicing makes sure that xml_tools is a list #}\n            {%- set ns.xml_tool_string = ns.xml_tool_string ~ (tool | string) ~ \"\\n\" -%}\n          {%- endfor -%}\n          {%- set xml_tool_string = ns.xml_tool_string + \"</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call>\" -%}\n          {{- xml_tool_string -}}\n        {%- endif -%}\n        {%- if python_tools -%}\n          {%- set ns = namespace(python_tool_string=\"When you send a message containing Python code between '<code>' and '</code>' tags, it will be executed in a stateful Jupyter notebook environment, and you will then be given the output to continued reasoning in an agentic loop.\\n\\nYou can use the following tools in your python code like regular functions:\\n<tools>\\n\") -%}\n          {%- for tool in python_tools[:] -%} {# The slicing makes sure that python_tools is a list #}\n            {%- set ns.python_tool_string = ns.python_tool_string ~ (tool | string) ~ \"\\n\" -%}\n          {%- endfor -%}\n          {%- set python_tool_string = ns.python_tool_string + \"</tools>\\n\\nThe state persists between code executions: so variables that you define in one step are still available thereafter.\" -%}\n          {{- python_tool_string -}}\n        {%- endif -%}\n        {{- \"\\n\\n\" -}}\n        {{- \"<|im_end|>\\n\" -}}\n      {%- endif -%}\n    {%- endif -%}\n    {# \u2500\u2500\u2500\u2500\u2500 main loop \u2500\u2500\u2500\u2500\u2500 #}\n    {%- for message in messages -%}\n        {%- set content = message.content if message.content is string else \"\" -%}\n        {%- if message.role == \"user\" -%}\n            {{ \"<|im_start|>\" + message.role + \"\\n\"  + content + \"<|im_end|>\\n\" }}\n        {%- elif message.role == \"assistant\" -%}\n            {% generation %}\n            {%- if reasoning_mode == \"/think\" -%}\n                {{ \"<|im_start|>assistant\\n\" + content.lstrip(\"\\n\") + \"<|im_end|>\\n\" }}\n            {%- else -%}\n                {{ \"<|im_start|>assistant\\n\" + \"<think>\\n\\n</think>\\n\" + content.lstrip(\"\\n\") + \"<|im_end|>\\n\" }}\n            {%- endif -%}\n            {% endgeneration %}\n        {%- elif message.role == \"tool\" -%}\n        {{ \"<|im_start|>\" + \"user\\n\"  + content + \"<|im_end|>\\n\" }}\n        {%- endif -%}\n    {%- endfor -%}\n    {# \u2500\u2500\u2500\u2500\u2500 generation prompt \u2500\u2500\u2500\u2500\u2500 #}\n    {%- if add_generation_prompt -%}\n        {%- if reasoning_mode == \"/think\" -%}\n            {{ \"<|im_start|>assistant\\n\" }}\n        {%- else -%}\n            {{ \"<|im_start|>assistant\\n\" + \"<think>\\n\\n</think>\\n\"  }}\n        {%- endif -%}\n    {%- endif -%}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 65536,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|endoftext|>"
+}

vllm_patches/transformers_lizzy_tp.py

@@ -0,0 +1,211 @@
+"""Compat patch for Lizzy TP under vLLM's generic Transformers backend."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+_PATCH_ATTR = "_flwr_transformers_lizzy_tp_patch_applied"
+
+
+class _TensorParallelSliceNorm(nn.Module):
+    """Apply a full-width checkpoint norm to a TP-local activation slice."""
+
+    def __init__(self, base_norm: nn.Module, start_idx: int, end_idx: int):
+        super().__init__()
+        self.start_idx = start_idx
+        self.end_idx = end_idx
+        self.weight = base_norm.weight
+        if getattr(base_norm, "bias", None) is not None:
+            self.bias = base_norm.bias
+        else:
+            self.register_parameter("bias", None)
+        self.eps = float(
+            getattr(base_norm, "eps", getattr(base_norm, "variance_epsilon", 1e-6)),
+        )
+        self.norm_kind = (
+            "layernorm" if isinstance(base_norm, nn.LayerNorm) else "rmsnorm"
+        )
+
+    @property
+    def local_size(self) -> int:
+        return self.end_idx - self.start_idx
+
+    def _slice_param(self, param: torch.Tensor | None) -> torch.Tensor | None:
+        if param is None:
+            return None
+        if param.shape[0] == self.local_size:
+            return param
+        return param[self.start_idx : self.end_idx]
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        weight = self.weight
+        bias = self.bias
+        if hidden_states.shape[-1] != self.weight.shape[0]:
+            if hidden_states.shape[-1] != self.local_size:
+                msg = (
+                    "Unexpected hidden size for TP-sliced norm: "
+                    f"{hidden_states.shape[-1]} "
+                    f"(expected {self.weight.shape[0]} or {self.local_size})"
+                )
+                raise RuntimeError(msg)
+            weight = self._slice_param(weight)
+            bias = self._slice_param(bias)
+
+        if self.norm_kind == "layernorm":
+            return F.layer_norm(
+                hidden_states,
+                (hidden_states.shape[-1],),
+                weight,
+                bias,
+                self.eps,
+            )
+
+        input_dtype = hidden_states.dtype
+        hidden_states_fp32 = hidden_states.to(torch.float32)
+        variance = hidden_states_fp32.pow(2).mean(dim=-1, keepdim=True)
+        hidden_states_norm = hidden_states_fp32 * torch.rsqrt(variance + self.eps)
+        hidden_states_norm = hidden_states_norm.to(input_dtype)
+        output = weight * hidden_states_norm
+        if bias is not None:
+            output = output + bias
+        return output
+
+
+def _maybe_patch_lizzy_attention_for_tp(
+    *,
+    module: nn.Module,
+    prefix: str,
+    tp_size: int,
+    tp_rank: int,
+    log_replacement: Any,  # noqa: ANN401
+) -> None:
+    if tp_size <= 1 or type(module).__name__ != "LizzyAttention":
+        return
+
+    num_heads = getattr(module, "num_heads", None)
+    num_key_value_heads = getattr(module, "num_key_value_heads", None)
+    head_dim = getattr(module, "head_dim", None)
+    q_norm = getattr(module, "q_norm", None)
+    k_norm = getattr(module, "k_norm", None)
+    if not all(
+        isinstance(value, int)
+        for value in (num_heads, num_key_value_heads, head_dim)
+    ):
+        return
+    if num_heads % tp_size != 0 or num_key_value_heads % tp_size != 0:
+        return
+
+    local_num_heads = num_heads // tp_size
+    local_num_key_value_heads = num_key_value_heads // tp_size
+    local_q_dim = local_num_heads * head_dim
+    local_kv_dim = local_num_key_value_heads * head_dim
+
+    module.num_heads = local_num_heads
+    module.num_key_value_heads = local_num_key_value_heads
+    module.num_key_value_groups = local_num_heads // local_num_key_value_heads
+
+    if q_norm is not None and getattr(q_norm, "weight", None) is not None:
+        start = tp_rank * local_q_dim
+        end = start + local_q_dim
+        module.q_norm = _TensorParallelSliceNorm(q_norm, start, end)
+        log_replacement(f"{prefix}.q_norm", q_norm, module.q_norm)
+
+    if k_norm is not None and getattr(k_norm, "weight", None) is not None:
+        start = tp_rank * local_kv_dim
+        end = start + local_kv_dim
+        module.k_norm = _TensorParallelSliceNorm(k_norm, start, end)
+        log_replacement(f"{prefix}.k_norm", k_norm, module.k_norm)
+
+
+def patch_vllm_transformers_lizzy_tp() -> None:
+    """Patch the generic vLLM Transformers backend for Lizzy TP norms/heads."""
+    import vllm.model_executor.models.transformers as transformers_mod
+
+    transformers_base = transformers_mod.TransformersBase
+    if getattr(transformers_base, _PATCH_ATTR, False):
+        return
+
+    PreTrainedModel = transformers_mod.PreTrainedModel
+    maybe_prefix = transformers_mod.maybe_prefix
+    replace_linear_class = transformers_mod.replace_linear_class
+    get_feature_request_tip = transformers_mod.get_feature_request_tip
+    re = transformers_mod.re
+    log_replacement = transformers_mod.log_replacement
+    get_tp_rank = getattr(transformers_mod, "get_tensor_model_parallel_rank", None)
+    if get_tp_rank is None:
+        try:
+            from vllm.distributed import (  # noqa: PLC0415
+                get_tensor_model_parallel_rank as get_tp_rank,
+            )
+        except Exception:
+            get_tp_rank = lambda: 0
+
+    def tensor_parallel(self: Any) -> None:  # noqa: ANN401
+        """Apply the model's tensor parallel plan plus Lizzy attention fixes."""
+        is_pretrained_model = lambda m: isinstance(m, PreTrainedModel)
+        supports_tp_plan = lambda m: m.config.base_model_tp_plan is not None
+        pretrained_models = filter(is_pretrained_model, self.model.modules())
+        models_with_tp_plan = filter(supports_tp_plan, pretrained_models)
+
+        if not any(models_with_tp_plan) and self.tp_size > 1:
+            tip = get_feature_request_tip(
+                self.model_config.model,
+                self.model_config.trust_remote_code,
+            )
+            raise ValueError(
+                f"{type(self.model)} does not support tensor parallel. {tip}",
+            )
+
+        tp_rank = get_tp_rank()
+
+        def _tensor_parallel(
+            module: nn.Module,
+            prefix: str = "",
+            tp_plan: dict[str, str] | None = None,
+        ) -> None:
+            local_tp_plan = tp_plan or {}
+
+            if isinstance(module, PreTrainedModel):
+                local_tp_plan = module.config.base_model_tp_plan or {}
+                local_tp_plan = {
+                    maybe_prefix(prefix, key): value
+                    for key, value in local_tp_plan.items()
+                }
+
+            for child_name, child_module in module.named_children():
+                qual_name = maybe_prefix(prefix, child_name)
+                if isinstance(child_module, nn.Linear):
+                    generator = (p for p in local_tp_plan if re.match(p, qual_name))
+                    pattern = next(generator, None)
+                    style = local_tp_plan.get(pattern, "replicate")
+                    new_module = replace_linear_class(
+                        child_module,
+                        style,
+                        self.quant_config,
+                        prefix=qual_name,
+                    )
+                    setattr(module, child_name, new_module)
+                    log_replacement(qual_name, child_module, new_module)
+                else:
+                    _tensor_parallel(
+                        child_module,
+                        prefix=qual_name,
+                        tp_plan=local_tp_plan,
+                    )
+
+            _maybe_patch_lizzy_attention_for_tp(
+                module=module,
+                prefix=prefix,
+                tp_size=self.tp_size,
+                tp_rank=tp_rank,
+                log_replacement=log_replacement,
+            )
+
+        _tensor_parallel(self.model)
+
+    transformers_base.tensor_parallel = tensor_parallel
+    setattr(transformers_base, _PATCH_ATTR, True)