Upload cleaned Parallax model (config + modeling + weights + tokenizer)

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

config.json

@@ -0,0 +1,36 @@
+{
+  "model_type": "parallax",
+  "architectures": [
+    "ParallaxForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_parallax.ParallaxConfig",
+    "AutoModelForCausalLM": "modeling_parallax.ParallaxForCausalLM"
+  },
+  "vocab_size": 151936,
+  "hidden_size": 2048,
+  "intermediate_size": 6144,
+  "num_hidden_layers": 28,
+  "num_attention_heads": 16,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "max_position_embeddings": 4096,
+  "initializer_range": 0.02,
+  "rms_norm_eps": 1e-06,
+  "use_cache": true,
+  "tie_word_embeddings": true,
+  "rope_parameters": {
+    "rope_type": "default",
+    "rope_theta": 1000000
+  },
+  "attention_bias": false,
+  "use_sliding_window": false,
+  "sliding_window": null,
+  "max_window_layers": 28,
+  "attention_dropout": 0.0,
+  "attn_type": "parallax",
+  "torch_dtype": "bfloat16",
+  "bos_token_id": 151643,
+  "eos_token_id": 151645
+}

configuration_parallax.py

@@ -0,0 +1,48 @@
+"""Parallax HF config (inherits Qwen3Config).
+
+Parallax (Parameterized Local Linear Attention) reuses the Qwen3 architecture
+and swaps softmax self-attention for the parallax op. This config inherits from
+``transformers.models.qwen3.configuration_qwen3.Qwen3Config`` so it picks up all
+of Qwen3's conventions for free (strict dataclass, ``rope_parameters``,
+``layer_types``, tp/pp plans, etc.).
+
+The only Parallax-specific addition is the ``attn_type`` field, which selects the
+attention implementation. When ``attn_type == "sdpa"`` this config describes a
+plain Qwen3-equivalent model.
+
+Reference: https://arxiv.org/abs/2605.29157
+"""
+
+from __future__ import annotations
+
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+
+
+class ParallaxConfig(Qwen3Config):
+    r"""Config for the Parallax model (Qwen3-derived).
+
+    Args:
+        attn_type (`str`, *optional*, defaults to `"parallax"`):
+            Attention implementation. Must be one of ``"sdpa"`` or
+            ``"parallax"``. ``"sdpa"`` uses the plain Qwen3 causal attention
+            path (softmax attention). ``"parallax"`` uses the parallax op and
+            requires the extra ``r_proj`` / ``r_norm`` per-layer weights.
+    """
+
+    model_type = "parallax"
+
+    # ---- Parallax-specific extra ----
+    # "sdpa"    : plain causal self-attention (identical to Qwen3).
+    # "parallax": the parallax op (adds r_proj + r_norm and replaces softmax
+    #             self-attention).
+    attn_type: str = "parallax"
+
+    def __post_init__(self, **kwargs):
+        if self.attn_type not in ("sdpa", "parallax"):
+            raise ValueError(
+                f"attn_type must be 'sdpa' or 'parallax', got {self.attn_type!r}"
+            )
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["ParallaxConfig"]

model-00001-of-00001.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,373 @@
+{
+  "metadata": {
+    "total_size": 3676038144
+  },
+  "weight_map": {
+    "model.embed_tokens.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.r_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.r_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.norm.weight": "model-00001-of-00001.safetensors"
+  }
+}

modeling_parallax.py

@@ -0,0 +1,299 @@
+"""Parallax HF modeling file (inherits from Qwen3).
+
+Parallax (Parameterized Local Linear Attention) reuses the Qwen3 architecture
+and only overrides the attention class to swap softmax self-attention for the
+parallax op (``attn_type="parallax"``).
+
+Reference: https://arxiv.org/abs/2605.29157
+
+Design notes
+------------
+* All boilerplate — RMSNorm, MLP, rotary embeddings, decoder-layer scaffolding,
+  the base ``PreTrainedModel`` class, ``Qwen3Model``, ``Qwen3ForCausalLM``,
+  mask preparation, cache handling, ``GenerationMixin`` integration, etc. —
+  is inherited from ``transformers.models.qwen3``.
+* The only new code is:
+    - ``parallax_forward``: a pure-PyTorch kernel implementing the parallax op,
+      batched over heads and taking an additive attention mask in the HF shape
+      ``(B, 1, q_len, kv_len)``.
+    - ``ParallaxAttention``: inherits from ``Qwen3Attention``, adds an
+      ``r_proj`` / ``r_norm`` pair when ``config.attn_type == "parallax"``,
+      and overrides ``forward`` to dispatch to ``parallax_forward`` in that
+      case. For ``attn_type == "sdpa"`` the forward simply delegates to
+      ``Qwen3Attention.forward`` unchanged, so all modern HF attention
+      backends (sdpa, flash_attention_2, flex_attention, eager) remain
+      available for that path.
+    - ``ParallaxDecoderLayer``, ``ParallaxModel``, ``ParallaxPreTrainedModel``,
+      ``ParallaxForCausalLM``: thin subclasses that rewire parent classes to
+      the ``ParallaxAttention`` / ``ParallaxConfig`` below.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from transformers.cache_utils import Cache
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3Attention,
+    Qwen3DecoderLayer,
+    Qwen3ForCausalLM,
+    Qwen3Model,
+    Qwen3PreTrainedModel,
+    Qwen3RMSNorm,
+    apply_rotary_pos_emb,
+    repeat_kv,
+)
+
+from .configuration_parallax import ParallaxConfig
+
+
+# ---------------------------------------------------------------------------
+# Parallax op (pure PyTorch)
+# ---------------------------------------------------------------------------
+
+
+def parallax_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    r: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+) -> tuple[torch.Tensor, None]:
+    """The parallax op.
+
+        p = softmax((Q Kᵀ) * scale + mask)
+        rk = R Kᵀ
+        s = p V
+        c = sum(p * rk, -1, keepdim=True)
+        r_term = (p * rk) V
+        out = s + c * s - r_term
+
+    Shapes
+    ------
+    query : (B, H, q_len, D)
+    r     : (B, H, q_len, D)          — Parallax-specific extra projection
+    key   : (B, H_kv, kv_len, D)      — GQA-compressed, expanded internally
+    value : (B, H_kv, kv_len, D)
+    attention_mask : (B, 1, q_len, kv_len) additive, 0 or -inf (HF convention)
+
+    Returns
+    -------
+    attn_output : (B, q_len, H, D)    — transposed to match Qwen3 eager impl
+    attn_weights : None               — not exposed (the parallax op has no
+                                         meaningful "attention matrix" per head)
+    """
+    # Expand GQA-compressed k/v to match q heads
+    key = repeat_kv(key, module.num_key_value_groups)
+    value = repeat_kv(value, module.num_key_value_groups)
+
+    # Upcast to float32 for numerical stability (the softmax accumulates in fp32).
+    out_dtype = query.dtype
+    qf = query.float()
+    rf = r.float()
+    kf = key.float()
+    vf = value.float()
+
+    qk = torch.einsum("bhqd,bhkd->bhqk", qf, kf) * scaling
+    q_len = qk.shape[-2]
+    kv_len = qk.shape[-1]
+    if attention_mask is not None:
+        # HF causal masks may cover a larger key span than our current
+        # kv_len during decode; slice to match.
+        qk = qk + attention_mask[..., :q_len, :kv_len].to(qk.dtype)
+    elif q_len > 1:
+        # HF's `create_causal_mask` returns None when the model's default
+        # attention backend (e.g. sdpa) handles causality implicitly. The
+        # parallax op is hand-written and needs the explicit additive mask,
+        # so build one here. Decode (q_len == 1) is naturally unmasked: the
+        # single query can attend to every cached key.
+        offset = kv_len - q_len
+        row = torch.arange(q_len, device=qk.device)[:, None] + offset
+        col = torch.arange(kv_len, device=qk.device)[None, :]
+        causal = row >= col
+        bias = torch.zeros(q_len, kv_len, dtype=qk.dtype, device=qk.device)
+        bias = bias.masked_fill(~causal, float("-inf"))
+        qk = qk + bias
+
+    p = torch.softmax(qk, dim=-1)
+    rk = torch.einsum("bhqd,bhkd->bhqk", rf, kf)
+
+    s = torch.einsum("bhqk,bhkd->bhqd", p, vf)
+    c = (p * rk).sum(dim=-1, keepdim=True)
+    r_term = torch.einsum("bhqk,bhkd->bhqd", p * rk, vf)
+
+    attn_output = (s + c * s - r_term).to(out_dtype)
+    # Match Qwen3 eager_attention_forward: return transposed (B, q_len, H, D)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+    return attn_output, None
+
+
+# ---------------------------------------------------------------------------
+# Attention / DecoderLayer / Model / ForCausalLM
+# ---------------------------------------------------------------------------
+
+
+class ParallaxAttention(Qwen3Attention):
+    """Qwen3 attention plus an optional parallax path."""
+
+    def __init__(self, config: ParallaxConfig, layer_idx: int) -> None:
+        super().__init__(config, layer_idx)
+        self.attn_type = config.attn_type
+        if self.attn_type == "parallax":
+            # Extra projection + RMSNorm for the parallax op.
+            self.r_proj = nn.Linear(
+                config.hidden_size,
+                config.num_attention_heads * self.head_dim,
+                bias=config.attention_bias,
+            )
+            self.r_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if self.attn_type != "parallax":
+            # Plain Qwen3 attention — delegate to the parent, which supports
+            # sdpa / flash_attention_2 / flex_attention / eager backends.
+            return super().forward(
+                hidden_states=hidden_states,
+                position_embeddings=position_embeddings,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+                **kwargs,
+            )
+
+        # ---- Parallax path ----
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(
+            self.q_proj(hidden_states).view(hidden_shape)
+        ).transpose(1, 2)
+        key_states = self.k_norm(
+            self.k_proj(hidden_states).view(hidden_shape)
+        ).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        # r: same num_heads as q (not num_key_value_heads), rope applied.
+        # Note hidden_shape uses -1 for the heads dim, so r_proj (which produces
+        # num_attention_heads*head_dim) will correctly view as (..., n_heads, D).
+        r_states = self.r_norm(
+            self.r_proj(hidden_states).view(hidden_shape)
+        ).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+        r_states, _ = apply_rotary_pos_emb(r_states, r_states, cos, sin)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+
+        attn_output, attn_weights = parallax_forward(
+            self,
+            query_states,
+            r_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class ParallaxDecoderLayer(Qwen3DecoderLayer):
+    def __init__(self, config: ParallaxConfig, layer_idx: int) -> None:
+        # Call the grandparent (nn.Module) init to avoid Qwen3DecoderLayer
+        # building a Qwen3Attention we'd immediately replace.
+        nn.Module.__init__(self)
+        self.hidden_size = config.hidden_size
+        self.attention_type = config.layer_types[layer_idx]
+        self.self_attn = ParallaxAttention(config=config, layer_idx=layer_idx)
+        # Reuse the MLP + norms from Qwen3 via its own constructor logic.
+        # Easiest is to rebuild them directly here to match Qwen3DecoderLayer.
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3MLP
+
+        self.mlp = Qwen3MLP(config)
+        self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+
+class ParallaxPreTrainedModel(Qwen3PreTrainedModel):
+    config: ParallaxConfig  # type: ignore[assignment]
+    config_class = ParallaxConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["ParallaxDecoderLayer"]
+    # The parallax path is hand-written and does not plug into
+    # ALL_ATTENTION_FUNCTIONS. The sdpa path (delegated to super) supports the
+    # standard backends.
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+
+class ParallaxModel(ParallaxPreTrainedModel, Qwen3Model):
+    def __init__(self, config: ParallaxConfig) -> None:
+        # Bypass Qwen3Model.__init__ (which builds Qwen3DecoderLayers) and
+        # rebuild the module tree with ParallaxDecoderLayers.
+        Qwen3PreTrainedModel.__init__(self, 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.layers = nn.ModuleList(
+            [
+                ParallaxDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # Qwen3Model constructs its rotary embedding from the config; reuse.
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding
+
+        self.rotary_emb = Qwen3RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.has_sliding_layers = "sliding_attention" in (config.layer_types or [])
+        self.post_init()
+
+
+class ParallaxForCausalLM(ParallaxPreTrainedModel, Qwen3ForCausalLM):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: ParallaxConfig) -> None:
+        # Bypass Qwen3ForCausalLM.__init__ (which builds a Qwen3Model) and
+        # rebuild with ParallaxModel.
+        Qwen3PreTrainedModel.__init__(self, config)
+        self.model = ParallaxModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+
+__all__ = [
+    "ParallaxAttention",
+    "ParallaxDecoderLayer",
+    "ParallaxForCausalLM",
+    "ParallaxModel",
+    "ParallaxPreTrainedModel",
+    "parallax_forward",
+]

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</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><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is string %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in content %}\n                {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n                {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if loop.last or (not loop.last and reasoning_content) %}\n                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<|im_start|>' + message.role + '\\n' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}