Instructions to use summerMC/Qwen3.5-9B-SpeedX9-GDN32 with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use summerMC/Qwen3.5-9B-SpeedX9-GDN32 with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="summerMC/Qwen3.5-9B-SpeedX9-GDN32", trust_remote_code=True)
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("summerMC/Qwen3.5-9B-SpeedX9-GDN32", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use summerMC/Qwen3.5-9B-SpeedX9-GDN32 with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "summerMC/Qwen3.5-9B-SpeedX9-GDN32"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "summerMC/Qwen3.5-9B-SpeedX9-GDN32",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/summerMC/Qwen3.5-9B-SpeedX9-GDN32

SGLang

How to use summerMC/Qwen3.5-9B-SpeedX9-GDN32 with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "summerMC/Qwen3.5-9B-SpeedX9-GDN32" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "summerMC/Qwen3.5-9B-SpeedX9-GDN32",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "summerMC/Qwen3.5-9B-SpeedX9-GDN32" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "summerMC/Qwen3.5-9B-SpeedX9-GDN32",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use summerMC/Qwen3.5-9B-SpeedX9-GDN32 with Docker Model Runner:
```
docker model run hf.co/summerMC/Qwen3.5-9B-SpeedX9-GDN32
```

summerMC commited on 3 days ago

Commit

efc38b5

verified ·

1 Parent(s): 701dcba

Upload Qwen3.5 UNI MAX 9B checkpoint with benchmark results

Browse files

Files changed (20) hide show

.gitattributes +1 -0
README.md +251 -1
benchmark_results.json +140 -0
chat_template.jinja +154 -0
config.json +133 -0
configuration_qwen35_gdn24.py +93 -0
engine.py +430 -0
fused_ops.py +220 -0
gdn24_metadata.json +26 -0
generation_config.json +8 -0
model-00001-of-00004.safetensors +3 -0
model-00002-of-00004.safetensors +3 -0
model-00003-of-00004.safetensors +3 -0
model-00004-of-00004.safetensors +3 -0
model.safetensors.index.json +459 -0
modeling_qwen35_gdn24.py +338 -0
quantization.py +230 -0
tokenizer.json +3 -0
tokenizer_config.json +32 -0
uni_engine.py +290 -0

.gitattributes CHANGED Viewed

@@ -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

 *.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

README.md CHANGED Viewed

@@ -1,3 +1,253 @@
 ---
-license: apache-2.0
 ---

 ---
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- qwen
+- qwen3.5
+- text-generation
+- recurrent
+- linear-attention
+- cuda
+- custom-code
+base_model: Qwen/Qwen3.5-9B
 ---
+# Qwen3.5 UNI MAX 9B
+This repository contains a UNI MAX conversion of
+`Qwen/Qwen3.5-9B`.
+The source hybrid attention topology is converted to a recurrent
+linear-attention runtime. Full-attention layers are replaced through
+layer-wise distillation from neighboring native GDN/linear-attention
+layers.
+## Architecture
+- Base model: `Qwen/Qwen3.5-9B`
+- Hidden layers: `32`
+- Runtime topology: linear attention
+- Converted source full-attention layers: `3, 7, 11, 15, 19, 23, 27, 31`
+- Custom runtime: `qwen35_unimax_engine`
+- Kernel acceleration: enabled for the benchmark below
+The checkpoint should be used with the matching UNI MAX runtime code.
+It is not claimed to be numerically identical to the original
+full-attention model.
+## Distillation
+The converted full-attention layers were initialized from neighboring
+native GDN layers and then optimized with sequential layer-wise
+distillation.
+The current build used 20 optimization steps per converted layer.
+## Benchmark environment
+| Item | Value |
+|---|---|
+| GPU | NVIDIA RTX PRO 6000 Blackwell Server Edition |
+| GPU VRAM | 94.97 GiB |
+| Peak allocated VRAM during benchmark | 70.14 GiB |
+| PyTorch | 2.11.0+cu130 |
+| CUDA | 13.0 |
+| CUDA capability | 12.0 |
+| NVIDIA driver | 580.82.07 |
+| Decode steps | 64 |
+| Warmup runs | 2 |
+| Measured runs | 5 |
+| Context lengths | 256, 1024, 4096, 16384 |
+| Benchmark wall time | 90.62 s |
+## Benchmark results
+| Metric | Value |
+|---|---:|
+| verified | None |
+| selected_mode | exact |
+| selected_block | 8 |
+The benchmark above was generated directly from
+`benchmark_uni_max()` on the hardware shown above.
+Because throughput, latency, kernel selection, quantization support,
+and memory consumption depend strongly on the GPU, CUDA/PyTorch
+versions, context length, and runtime configuration, these numbers
+should not be treated as hardware-independent performance claims.
+## Raw benchmark output
+```json
+{
+  "official": [
+    [
+      256,
+      9771.298185735035,
+      61.8274717963079
+    ],
+    [
+      1024,
+      17823.304562533922,
+      61.66807965901939
+    ],
+    [
+      4096,
+      20324.797811420387,
+      60.59438884144146
+    ],
+    [
+      16384,
+      18413.893708085165,
+      58.409657369508345
+    ]
+  ],
+  "safe": [
+    [
+      256,
+      10715.681982051035,
+      79.45237658001217
+    ],
+    [
+      1024,
+      19627.025186990726,
+      79.44713049745303
+    ],
+    [
+      4096,
+      22395.258154021976,
+      79.44376760906928
+    ],
+    [
+      16384,
+      22677.037604190627,
+      79.4362706749314
+    ]
+  ],
+  "speed": [
+    [
+      256,
+      10573.784778868314,
+      79.44571045597908,
+      0.28697000016109087,
+      10452.495043114082
+    ],
+    [
+      1024,
+      19569.254713139573,
+      79.44294428720825,
+      0.2903700001297693,
+      19462.145367010704
+    ],
+    [
+      4096,
+      22337.05592154273,
+      79.44607524961788,
+      0.29715000027863425,
+      22300.547693510274
+    ],
+    [
+      16384,
+      22578.659939700647,
+      79.43958559875574,
+      0.32848999990164884,
+      22568.420127905207
+    ]
+  ],
+  "verified": null,
+  "selected_mode": "exact",
+  "selected_block": 8,
+  "candidates": [
+    {
+      "mode": "fp8",
+      "accepted": false,
+      "agreement": 0.1875,
+      "prefix": 10,
+      "replay_tok_s": 84.39886475159058,
+      "block": 8,
+      "reason": "agreement 0.188 < 1.000"
+    },
+    {
+      "mode": "int8",
+      "accepted": false,
+      "agreement": 1.0,
+      "prefix": 64,
+      "replay_tok_s": 64.56140914184326,
+      "block": 2,
+      "reason": "not faster than current 79.45 tok/s"
+    }
+  ]
+}
+```
+The machine-readable benchmark record is also included as
+`benchmark_results.json`.
+## Runtime configuration
+The benchmark evaluated:
+- CUDA kernels enabled
+- graph candidates: `(2, 4, 8)`
+- quantization candidates: `('fp8', 'int8')`
+- quality verification steps: `64`
+- quality context: `64`
+- minimum fast-path agreement: `1.0`
+The runtime may select a different execution path depending on
+hardware support and quality verification.
+## Usage
+Install the matching `qwen35_unimax_engine` runtime before loading the
+checkpoint.
+```python
+from qwen35_unimax_engine import UniMaxPipeline
+pipe = UniMaxPipeline.from_pretrained(
+    "YOUR_USERNAME/qwen35-unimax-engine-9b-v2_1",
+    mode="auto",
+    graph_candidates=(2, 4, 8),
+    quant_modes=("fp8", "int8"),
+    quality_steps=64,
+    min_agreement=1.0,
+    use_kernels=True,
+)
+print(pipe.runtime_summary())
+output = pipe(
+    "Explain recurrent linear attention:",
+    max_new_tokens=128,
+)
+print(output)
+```
+## Validation
+Before publication, the local checkpoint was checked with the UNI MAX
+checkpoint validator and benchmark path.
+Users should independently evaluate task quality before deploying the
+converted model. Layer-wise agreement tests are narrower than a full
+language-model evaluation suite.
+## Limitations
+This is a converted and distilled model, not the unmodified
+`Qwen/Qwen3.5-9B` checkpoint.
+Replacing full attention with recurrent linear attention changes the
+model architecture and can affect accuracy, long-context behavior,
+reasoning, generation quality, and numerical output.
+Benchmark results are specific to the hardware and software
+environment documented above.
+## Base model
+Derived from `Qwen/Qwen3.5-9B`. Refer to the upstream model repository
+for the original model documentation, license, intended use, and
+limitations.

benchmark_results.json ADDED Viewed

	@@ -0,0 +1,140 @@

+{
+  "model": "/content/qwen35_uni_max_work/qwen35-unimax-engine-9b-v2_1",
+  "source_model": "Qwen/Qwen3.5-9B",
+  "environment": {
+    "gpu": "NVIDIA RTX PRO 6000 Blackwell Server Edition",
+    "vram_gib": 94.971,
+    "peak_allocated_vram_gib": 70.141,
+    "torch": "2.11.0+cu130",
+    "cuda": "13.0",
+    "cuda_capability": [
+      12,
+      0
+    ],
+    "nvidia_driver": "580.82.07",
+    "python": "3.13.15"
+  },
+  "benchmark": {
+    "contexts": [
+      256,
+      1024,
+      4096,
+      16384
+    ],
+    "decode_steps": 64,
+    "warmup": 2,
+    "runs": 5,
+    "graph_candidates": [
+      2,
+      4,
+      8
+    ],
+    "quant_modes": [
+      "fp8",
+      "int8"
+    ],
+    "quality_steps": 64,
+    "quality_context": 64,
+    "min_agreement": 1.0,
+    "wall_time_seconds": 90.61699619499996
+  },
+  "results": {
+    "official": [
+      [
+        256,
+        9771.298185735035,
+        61.8274717963079
+      ],
+      [
+        1024,
+        17823.304562533922,
+        61.66807965901939
+      ],
+      [
+        4096,
+        20324.797811420387,
+        60.59438884144146
+      ],
+      [
+        16384,
+        18413.893708085165,
+        58.409657369508345
+      ]
+    ],
+    "safe": [
+      [
+        256,
+        10715.681982051035,
+        79.45237658001217
+      ],
+      [
+        1024,
+        19627.025186990726,
+        79.44713049745303
+      ],
+      [
+        4096,
+        22395.258154021976,
+        79.44376760906928
+      ],
+      [
+        16384,
+        22677.037604190627,
+        79.4362706749314
+      ]
+    ],
+    "speed": [
+      [
+        256,
+        10573.784778868314,
+        79.44571045597908,
+        0.28697000016109087,
+        10452.495043114082
+      ],
+      [
+        1024,
+        19569.254713139573,
+        79.44294428720825,
+        0.2903700001297693,
+        19462.145367010704
+      ],
+      [
+        4096,
+        22337.05592154273,
+        79.44607524961788,
+        0.29715000027863425,
+        22300.547693510274
+      ],
+      [
+        16384,
+        22578.659939700647,
+        79.43958559875574,
+        0.32848999990164884,
+        22568.420127905207
+      ]
+    ],
+    "verified": null,
+    "selected_mode": "exact",
+    "selected_block": 8,
+    "candidates": [
+      {
+        "mode": "fp8",
+        "accepted": false,
+        "agreement": 0.1875,
+        "prefix": 10,
+        "replay_tok_s": 84.39886475159058,
+        "block": 8,
+        "reason": "agreement 0.188 < 1.000"
+      },
+      {
+        "mode": "int8",
+        "accepted": false,
+        "agreement": 1.0,
+        "prefix": 64,
+        "replay_tok_s": 64.56140914184326,
+        "block": 2,
+        "reason": "not faster than current 79.45 tok/s"
+      }
+    ]
+  }
+}

chat_template.jinja ADDED Viewed

	@@ -0,0 +1,154 @@

+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}

config.json ADDED Viewed

	@@ -0,0 +1,133 @@

+{
+  "architectures": [
+    "Qwen35GDN24ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attn_output_gate": true,
+  "auto_map": {
+    "AutoConfig": "configuration_qwen35_gdn24.Qwen35GDN24Config",
+    "AutoModelForCausalLM": "modeling_qwen35_gdn24.Qwen35GDN24ForCausalLM"
+  },
+  "bos_token_id": null,
+  "converted_layers": [
+    3,
+    7,
+    11,
+    15,
+    19,
+    23,
+    27,
+    31
+  ],
+  "dtype": "bfloat16",
+  "eos_token_id": 248044,
+  "full_attention_interval": 4,
+  "gdn24_format_version": 1,
+  "head_dim": 256,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 12288,
+  "layer_types": [
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention"
+  ],
+  "linear_conv_kernel_dim": 4,
+  "linear_key_head_dim": 128,
+  "linear_num_key_heads": 16,
+  "linear_num_value_heads": 32,
+  "linear_value_head_dim": 128,
+  "mamba_ssm_dtype": "float32",
+  "max_position_embeddings": 262144,
+  "mlp_only_layers": [],
+  "model_type": "qwen3_5_gdn24",
+  "mtp_num_hidden_layers": 1,
+  "mtp_use_dedicated_embeddings": false,
+  "num_attention_heads": 16,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 4,
+  "pad_token_id": null,
+  "partial_rotary_factor": 0.25,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "mrope_interleaved": true,
+    "mrope_section": [
+      11,
+      11,
+      10
+    ],
+    "partial_rotary_factor": 0.25,
+    "rope_theta": 10000000,
+    "rope_type": "default"
+  },
+  "source_layer_types": [
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention"
+  ],
+  "source_model": "Qwen/Qwen3.5-9B",
+  "tie_word_embeddings": false,
+  "transformers_version": "5.16.0.dev0",
+  "use_cache": true,
+  "vocab_size": 248320
+}

configuration_qwen35_gdn24.py ADDED Viewed

	@@ -0,0 +1,93 @@

+from __future__ import annotations
+from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
+class Qwen35GDN24Config(Qwen3_5TextConfig):
+    """Qwen3.5 text config whose runtime token mixer is Gated DeltaNet in all layers.
+    ``source_layer_types`` preserves the official/source Qwen3.5 topology while
+    ``layer_types`` is always the homogeneous runtime topology.  The runtime
+    topology is normalized *before* calling the Transformers base config so
+    base-class post-init/validation can never restore or observe the source
+    ``full_attention`` entries as runtime layers.
+    """
+    model_type = "qwen3_5_gdn24"
+    def __init__(
+        self,
+        *args,
+        source_layer_types: list[str] | None = None,
+        converted_layers: list[int] | None = None,
+        source_model: str = "Qwen/Qwen3.5-9B",
+        gdn24_format_version: int = 1,
+        **kwargs,
+    ):
+        # Keep the serialized/incoming topology separate from the runtime one.
+        # On a fresh conversion this is the official hybrid topology.  On
+        # reload, source_layer_types is serialized explicitly and therefore
+        # remains authoritative even though layer_types is already all-linear.
+        incoming_layer_types = kwargs.pop("layer_types", None)
+        if source_layer_types is None:
+            source_layer_types = incoming_layer_types
+        if source_layer_types is None:
+            # Preserve the base model's default topology only as conversion
+            # metadata. Qwen3.5 defaults to full attention every fourth layer.
+            n = int(kwargs.get("num_hidden_layers", len(source_layer_types) if source_layer_types is not None else 32))
+            interval = int(kwargs.get("full_attention_interval", 4))
+            if interval > 0:
+                source_layer_types = [
+                    "linear_attention" if (i + 1) % interval else "full_attention"
+                    for i in range(n)
+                ]
+            else:
+                source_layer_types = ["linear_attention"] * n
+        else:
+            source_layer_types = list(source_layer_types)
+            n = int(kwargs.get("num_hidden_layers", len(source_layer_types)))
+        if len(source_layer_types) != n:
+            raise ValueError(
+                f"source_layer_types length ({len(source_layer_types)}) must equal "
+                f"num_hidden_layers ({n})"
+            )
+        if any(t not in {"linear_attention", "full_attention"} for t in source_layer_types):
+            raise ValueError(f"unsupported source layer type: {source_layer_types}")
+        expected = [i for i, t in enumerate(source_layer_types) if t == "full_attention"]
+        if converted_layers is None:
+            converted_layers = expected
+        normalized_converted = sorted(int(i) for i in converted_layers)
+        if normalized_converted != expected:
+            raise ValueError(
+                f"converted_layers must equal all source full-attention layers: {expected}"
+            )
+        # Critical fix: normalize BEFORE base config construction.  Newer
+        # Transformers Qwen3.5 configs perform layer-type handling in
+        # __post_init__, so fixing the value only after super().__init__ is too
+        # late and is version-sensitive.
+        kwargs["layer_types"] = ["linear_attention"] * n
+        super().__init__(*args, **kwargs)
+        # Assert the invariant explicitly instead of silently shipping a hybrid
+        # runtime config if upstream Transformers changes behavior again.
+        self.layer_types = ["linear_attention"] * int(self.num_hidden_layers)
+        if len(self.layer_types) != n or set(self.layer_types) != {"linear_attention"}:
+            raise RuntimeError("failed to normalize GDN24 runtime layer_types")
+        self.source_layer_types = source_layer_types
+        self.converted_layers = normalized_converted
+        self.source_model = str(source_model)
+        self.gdn24_format_version = int(gdn24_format_version)
+        self.architectures = ["Qwen35GDN24ForCausalLM"]
+        self.auto_map = {
+            "AutoConfig": "configuration_qwen35_gdn24.Qwen35GDN24Config",
+            "AutoModelForCausalLM": "modeling_qwen35_gdn24.Qwen35GDN24ForCausalLM",
+        }
+Qwen35GDN24Config.register_for_auto_class("AutoConfig")

engine.py ADDED Viewed

	@@ -0,0 +1,430 @@

+from __future__ import annotations
+"""Persistent CUDA-Graph greedy decoder for the all-recurrent GDN24 runtime.
+The core idea is specific to a fully recurrent model: the cache has fixed tensor
+addresses and token position is implicit in state evolution.  We can therefore
+capture multiple autoregressive steps into one CUDA Graph and let the graph feed
+its own argmax token into the next step.
+"""
+from dataclasses import dataclass
+import time
+from typing import Any
+import torch
+@dataclass
+class _LayerSnapshot:
+    conv: dict[int, torch.Tensor]
+    recurrent: dict[int, torch.Tensor]
+    has_previous_state: dict[int, bool]
+    conv_initialized: dict[int, bool]
+    recurrent_initialized: dict[int, bool]
+@dataclass
+class CacheSnapshot:
+    seen_tokens: int
+    layers: list[_LayerSnapshot]
+def snapshot_cache(cache) -> CacheSnapshot:
+    layers: list[_LayerSnapshot] = []
+    for layer in cache.layers:
+        conv = {int(i): t.detach().clone() for i, t in layer.conv_states.items() if t is not None}
+        recurrent = {
+            int(i): t.detach().clone() for i, t in layer.recurrent_states.items() if t is not None
+        }
+        layers.append(
+            _LayerSnapshot(
+                conv=conv,
+                recurrent=recurrent,
+                has_previous_state=dict(layer.has_previous_state),
+                conv_initialized=dict(layer.is_conv_states_initialized),
+                recurrent_initialized=dict(layer.is_recurrent_states_initialized),
+            )
+        )
+    return CacheSnapshot(seen_tokens=int(cache.seen_tokens), layers=layers)
+def restore_cache_(cache, snap: CacheSnapshot) -> None:
+    if len(cache.layers) != len(snap.layers):
+        raise ValueError("cache topology changed while restoring snapshot")
+    for layer, state in zip(cache.layers, snap.layers):
+        for i, src in state.conv.items():
+            dst = layer.conv_states[i]
+            if dst is None or dst.shape != src.shape:
+                raise ValueError(f"conv state storage changed at state {i}")
+            dst.copy_(src)
+        for i, src in state.recurrent.items():
+            dst = layer.recurrent_states[i]
+            if dst is None or dst.shape != src.shape:
+                raise ValueError(f"recurrent state storage changed at state {i}")
+            dst.copy_(src)
+        layer.has_previous_state.update(state.has_previous_state)
+        layer.is_conv_states_initialized.update(state.conv_initialized)
+        layer.is_recurrent_states_initialized.update(state.recurrent_initialized)
+    cache.seen_tokens = int(snap.seen_tokens)
+class SuperTurboGraphDecoder:
+    """Persistent, self-feeding greedy CUDA Graph decoder.
+    A single replay advances `block_size` recurrent decode steps.  The graph owns
+    one persistent recurrent cache, so it can be reused across prompts: `reset()`
+    preserves all CUDA addresses, prefill writes the new prompt state, then graph
+    replay continues autoregressively from that state.
+    """
+    def __init__(
+        self,
+        model,
+        *,
+        block_size: int = 8,
+        warmup_steps: int = 2,
+        graph_pool=None,
+    ):
+        if not torch.cuda.is_available():
+            raise RuntimeError("MAX TURBO CUDA Graph decoding requires CUDA")
+        if block_size < 1:
+            raise ValueError("block_size must be >= 1")
+        self.model = model.eval()
+        self.block_size = int(block_size)
+        self.warmup_steps = int(warmup_steps)
+        self.graph_pool = graph_pool
+        self.device = model.get_input_embeddings().weight.device
+        if self.device.type != "cuda":
+            raise RuntimeError(f"model must be on CUDA, got {self.device}")
+        self.cache = self.model.make_recurrent_cache()
+        self.static_token = torch.zeros((1, 1), dtype=torch.long, device=self.device)
+        self.output_tokens = torch.empty((1, self.block_size), dtype=torch.long, device=self.device)
+        self.graph: torch.cuda.CUDAGraph | None = None
+        self.capture_seconds: float | None = None
+        self.capture_error: str | None = None
+        self.capture_allocated_delta_mib: float = 0.0
+        self.capture_reserved_delta_mib: float = 0.0
+        self._captured = False
+        self._capture_attempted = False
+    @torch.inference_mode()
+    def _greedy_step(self) -> torch.Tensor:
+        # Private clock flag is consumed by Qwen35GDN24Model and never reaches GDN kernels.
+        return self.model.greedy_step(
+            self.static_token,
+            self.cache,
+            advance_cache_clock=False,
+        )
+    @torch.inference_mode()
+    def _unrolled_block(self) -> None:
+        for i in range(self.block_size):
+            next_token = self._greedy_step()
+            self.output_tokens[:, i : i + 1].copy_(next_token)
+            self.static_token.copy_(next_token)
+    @torch.inference_mode()
+    def capture(self) -> bool:
+        if self._captured:
+            return True
+        if self._capture_attempted:
+            return False
+        self._capture_attempted = True
+        try:
+            # Initialize every native LinearAttentionLayer cache with stable addresses.
+            self.cache.reset()
+            dummy = torch.zeros((1, 1), dtype=torch.long, device=self.device)
+            first = self.model(
+                input_ids=dummy,
+                past_key_values=self.cache,
+                use_cache=True,
+                logits_to_keep=1,
+            ).logits[:, -1, :].argmax(dim=-1, keepdim=True)
+            self.static_token.copy_(first)
+            torch.cuda.synchronize()
+            baseline = snapshot_cache(self.cache)
+            baseline_token = self.static_token.detach().clone()
+            # Warm single-token recurrent kernels and Triton fusions before capture.
+            for _ in range(max(self.warmup_steps, 0)):
+                self._unrolled_block()
+                torch.cuda.synchronize()
+                restore_cache_(self.cache, baseline)
+                self.static_token.copy_(baseline_token)
+            graph = torch.cuda.CUDAGraph()
+            torch.cuda.synchronize()
+            alloc_before = torch.cuda.memory_allocated(self.device)
+            reserve_before = torch.cuda.memory_reserved(self.device)
+            t0 = time.perf_counter()
+            graph_kwargs = {} if self.graph_pool is None else {"pool": self.graph_pool}
+            with torch.cuda.graph(graph, **graph_kwargs):
+                self._unrolled_block()
+            torch.cuda.synchronize()
+            self.capture_seconds = time.perf_counter() - t0
+            self.capture_allocated_delta_mib = max(0, torch.cuda.memory_allocated(self.device) - alloc_before) / 2**20
+            self.capture_reserved_delta_mib = max(0, torch.cuda.memory_reserved(self.device) - reserve_before) / 2**20
+            # Capture executes once. Restore the exact pre-capture numerical state
+            # without replacing any tensor objects/addresses recorded by the graph.
+            restore_cache_(self.cache, baseline)
+            self.static_token.copy_(baseline_token)
+            if hasattr(graph, "instantiate"):
+                try:
+                    graph.instantiate()
+                except Exception:
+                    pass
+            self.graph = graph
+            self._captured = True
+            self.cache.reset()
+            return True
+        except Exception as exc:
+            self.capture_error = f"{type(exc).__name__}: {exc}"
+            self.graph = None
+            self._captured = False
+            try:
+                self.cache.reset()
+            except Exception:
+                pass
+            return False
+    @torch.inference_mode()
+    def reset(self) -> None:
+        self.cache.reset()
+    @torch.inference_mode()
+    def prefill(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, Any]:
+        """Prefill into the persistent graph cache and return the first next token."""
+        self.cache.reset()
+        out = self.model(
+            input_ids=input_ids,
+            past_key_values=self.cache,
+            use_cache=True,
+            logits_to_keep=1,
+        )
+        token = torch.argmax(out.logits[:, -1, :], dim=-1, keepdim=True)
+        return token, out
+    @torch.inference_mode()
+    def decode_forwards(self, first_token: torch.Tensor, steps: int) -> torch.Tensor:
+        """Run `steps` recurrent forward passes, matching the benchmark's decode metric.
+        The returned tensor contains the last generated token.  For the fast path,
+        `steps` should be divisible by `block_size`; a short eager tail handles any
+        remainder.
+        """
+        steps = int(steps)
+        if steps < 0:
+            raise ValueError("steps must be >= 0")
+        if steps == 0:
+            return first_token
+        if not self._captured:
+            return self._decode_eager(first_token, steps)
+        self.static_token.copy_(first_token)
+        full_blocks, remainder = divmod(steps, self.block_size)
+        for _ in range(full_blocks):
+            self.graph.replay()
+        # The graph deliberately does not update the Python token clock.
+        self.cache.advance(full_blocks * self.block_size)
+        if remainder:
+            token = self.static_token
+            for _ in range(remainder):
+                out = self.model(
+                    input_ids=token,
+                    past_key_values=self.cache,
+                    use_cache=True,
+                    logits_to_keep=1,
+                )
+                token = torch.argmax(out.logits[:, -1, :], dim=-1, keepdim=True)
+            self.static_token.copy_(token)
+        return self.static_token
+    @torch.inference_mode()
+    def _decode_eager(self, first_token: torch.Tensor, steps: int) -> torch.Tensor:
+        token = first_token
+        for _ in range(steps):
+            out = self.model(
+                input_ids=token,
+                past_key_values=self.cache,
+                use_cache=True,
+                logits_to_keep=1,
+            )
+            token = torch.argmax(out.logits[:, -1, :], dim=-1, keepdim=True)
+        return token
+    @torch.inference_mode()
+    def decode_tokens(self, first_token: torch.Tensor, steps: int) -> torch.Tensor:
+        """Return every greedy token produced by ``steps`` recurrent forwards.
+        Unlike :meth:`decode_forwards`, this materializes the token sequence and is
+        intended for user generation, quality checks, and UNI MAX verification.
+        The persistent recurrent cache remains graph-address-stable.
+        """
+        steps = int(steps)
+        if steps < 0:
+            raise ValueError("steps must be >= 0")
+        if steps == 0:
+            return torch.empty((first_token.shape[0], 0), dtype=torch.long, device=first_token.device)
+        self.static_token.copy_(first_token)
+        pieces: list[torch.Tensor] = []
+        if self._captured:
+            full_blocks, remainder = divmod(steps, self.block_size)
+            for _ in range(full_blocks):
+                self.graph.replay()
+                pieces.append(self.output_tokens.detach().clone())
+            if full_blocks:
+                self.cache.advance(full_blocks * self.block_size)
+            token = self.static_token
+        else:
+            remainder = steps
+            token = first_token
+        for _ in range(remainder):
+            token = self.model.greedy_step(token, self.cache)
+            pieces.append(token.detach().clone())
+        if pieces:
+            self.static_token.copy_(pieces[-1][:, -1:])
+        return torch.cat(pieces, dim=1)
+    @torch.inference_mode()
+    def replay_block_tokens(self, first_token: torch.Tensor) -> torch.Tensor:
+        """Replay exactly one captured block and return its generated tokens."""
+        if not self._captured:
+            if not self.capture():
+                raise RuntimeError(f"CUDA Graph capture failed: {self.capture_error}")
+        self.static_token.copy_(first_token)
+        self.graph.replay()
+        self.cache.advance(self.block_size)
+        return self.output_tokens.detach().clone()
+    @torch.inference_mode()
+    def validate_against_eager(self, first_token: torch.Tensor) -> bool:
+        """Bit-exact greedy-token check for one captured block."""
+        if not self._captured and not self.capture():
+            raise RuntimeError(f"CUDA Graph capture failed: {self.capture_error}")
+        snap = snapshot_cache(self.cache)
+        start = first_token.detach().clone()
+        try:
+            token = start
+            eager_tokens = []
+            for _ in range(self.block_size):
+                token = self.model.greedy_step(
+                    token, self.cache, advance_cache_clock=False
+                )
+                eager_tokens.append(token.detach().clone())
+            eager = torch.cat(eager_tokens, dim=1)
+            restore_cache_(self.cache, snap)
+            self.static_token.copy_(start)
+            self.graph.replay()
+            torch.cuda.synchronize()
+            graphed = self.output_tokens.detach().clone()
+            if not torch.equal(eager, graphed):
+                raise RuntimeError(
+                    f"CUDA Graph greedy mismatch: eager={eager.tolist()} graph={graphed.tolist()}"
+                )
+            return True
+        finally:
+            restore_cache_(self.cache, snap)
+            self.static_token.copy_(start)
+    @property
+    def ready(self) -> bool:
+        return self._captured
+@dataclass(frozen=True)
+class GraphTuneResult:
+    block_size: int
+    decode_tok_s: float
+    capture_seconds: float
+    graph_allocated_mib: float = 0.0
+    graph_reserved_mib: float = 0.0
+@torch.inference_mode()
+def autotune_graph_decoder(
+    model,
+    *,
+    candidates: tuple[int, ...] = (4, 8, 16),
+    warmup_replays: int = 2,
+    timed_replays: int = 6,
+) -> tuple[SuperTurboGraphDecoder, list[GraphTuneResult]]:
+    """Capture and benchmark several self-feeding graph block sizes.
+    The search measures steady-state graph replay only; graph capture time is
+    reported separately and excluded.  Every candidate is numerically checked
+    against eager greedy decoding before it can win.
+    """
+    if not candidates:
+        raise ValueError("at least one graph block candidate is required")
+    results: list[GraphTuneResult] = []
+    best: SuperTurboGraphDecoder | None = None
+    best_speed = -1.0
+    failures: list[str] = []
+    shared_pool = torch.cuda.graph_pool_handle() if hasattr(torch.cuda, "graph_pool_handle") else None
+    for raw_block in candidates:
+        block = int(raw_block)
+        if block < 1:
+            failures.append(f"x{block}: invalid block size")
+            continue
+        engine = SuperTurboGraphDecoder(model, block_size=block, warmup_steps=2, graph_pool=shared_pool)
+        if not engine.capture():
+            failures.append(f"x{block}: {engine.capture_error}")
+            continue
+        # Start from a valid persistent recurrent state, then verify exact greedy
+        # token agreement before any timing result is accepted.
+        seed = torch.zeros((1, 8), dtype=torch.long, device=engine.device)
+        first, _ = engine.prefill(seed)
+        engine.validate_against_eager(first)
+        engine.reset()
+        first, _ = engine.prefill(seed)
+        engine.static_token.copy_(first)
+        for _ in range(max(int(warmup_replays), 0)):
+            engine.graph.replay()
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        for _ in range(max(int(timed_replays), 1)):
+            engine.graph.replay()
+        torch.cuda.synchronize()
+        dt = time.perf_counter() - t0
+        forwards = block * max(int(timed_replays), 1)
+        speed = forwards / dt
+        result = GraphTuneResult(
+            block_size=block,
+            decode_tok_s=float(speed),
+            capture_seconds=float(engine.capture_seconds or 0.0),
+            graph_allocated_mib=float(engine.capture_allocated_delta_mib),
+            graph_reserved_mib=float(engine.capture_reserved_delta_mib),
+        )
+        results.append(result)
+        if speed > best_speed:
+            best_speed = speed
+            previous_best = best
+            best = engine
+            if previous_best is not None:
+                del previous_best
+                torch.cuda.empty_cache()
+        else:
+            # Drop non-winning graph pools as early as possible.
+            del engine
+            torch.cuda.empty_cache()
+    if best is None:
+        reason = "; ".join(failures) if failures else "no valid candidates"
+        raise RuntimeError(f"MAX TURBO graph autotune failed: {reason}")
+    best.reset()
+    return best, results
+# MAX TURBO public name; keep the v3 class name as a compatibility alias.
+MaxTurboGraphDecoder = SuperTurboGraphDecoder

fused_ops.py ADDED Viewed

	@@ -0,0 +1,220 @@

+from __future__ import annotations
+"""Small fused CUDA/Triton operators for GDN24 MAX TURBO.
+The runtime deliberately keeps a pure-Torch fallback so checkpoints remain
+portable.  Triton is used only when it is already available through the CUDA
+PyTorch stack; it is not a checkpoint dependency.
+"""
+import torch
+from torch import nn
+import torch.nn.functional as F
+try:  # Triton ships with CUDA PyTorch builds used by Colab.
+    import triton
+    import triton.language as tl
+    _HAS_TRITON = True
+except Exception:  # pragma: no cover - CPU/source validation path
+    triton = None
+    tl = None
+    _HAS_TRITON = False
+if _HAS_TRITON:
+    @triton.jit
+    def _rmsnorm_kernel(x_ptr, w_ptr, y_ptr, n_cols: tl.constexpr, eps: tl.constexpr, BLOCK: tl.constexpr):
+        row = tl.program_id(0)
+        offs = tl.arange(0, BLOCK)
+        mask = offs < n_cols
+        x = tl.load(x_ptr + row * n_cols + offs, mask=mask, other=0.0).to(tl.float32)
+        w = tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+        var = tl.sum(x * x, axis=0) / n_cols
+        rstd = tl.rsqrt(var + eps)
+        y = x * rstd * (1.0 + w)
+        tl.store(y_ptr + row * n_cols + offs, y, mask=mask)
+    @triton.jit
+    def _add_rmsnorm_kernel(
+        x_ptr,
+        update_ptr,
+        w_ptr,
+        sum_ptr,
+        norm_ptr,
+        n_cols: tl.constexpr,
+        eps: tl.constexpr,
+        BLOCK: tl.constexpr,
+    ):
+        row = tl.program_id(0)
+        offs = tl.arange(0, BLOCK)
+        mask = offs < n_cols
+        x = tl.load(x_ptr + row * n_cols + offs, mask=mask, other=0.0).to(tl.float32)
+        u = tl.load(update_ptr + row * n_cols + offs, mask=mask, other=0.0).to(tl.float32)
+        w = tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+        s = x + u
+        var = tl.sum(s * s, axis=0) / n_cols
+        rstd = tl.rsqrt(var + eps)
+        n = s * rstd * (1.0 + w)
+        tl.store(sum_ptr + row * n_cols + offs, s, mask=mask)
+        tl.store(norm_ptr + row * n_cols + offs, n, mask=mask)
+    @triton.jit
+    def _add_final_rmsnorm_kernel(
+        x_ptr,
+        update_ptr,
+        w_ptr,
+        norm_ptr,
+        n_cols: tl.constexpr,
+        eps: tl.constexpr,
+        BLOCK: tl.constexpr,
+    ):
+        row = tl.program_id(0)
+        offs = tl.arange(0, BLOCK)
+        mask = offs < n_cols
+        x = tl.load(x_ptr + row * n_cols + offs, mask=mask, other=0.0).to(tl.float32)
+        u = tl.load(update_ptr + row * n_cols + offs, mask=mask, other=0.0).to(tl.float32)
+        w = tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+        s = x + u
+        var = tl.sum(s * s, axis=0) / n_cols
+        rstd = tl.rsqrt(var + eps)
+        n = s * rstd * (1.0 + w)
+        tl.store(norm_ptr + row * n_cols + offs, n, mask=mask)
+    @triton.jit
+    def _silu_mul_kernel(a_ptr, b_ptr, out_ptr, n_elements, BLOCK: tl.constexpr):
+        offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+        mask = offs < n_elements
+        a = tl.load(a_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+        b = tl.load(b_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+        # SiLU(a) = a * sigmoid(a)
+        sig = 1.0 / (1.0 + tl.exp(-a))
+        out = a * sig * b
+        tl.store(out_ptr + offs, out, mask=mask)
+def _can_triton(x: torch.Tensor) -> bool:
+    return bool(_HAS_TRITON and x.is_cuda and x.is_contiguous())
+def qwen_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    """Exact Qwen3.5 offset RMSNorm: norm(x) * (1 + weight)."""
+    d = x.shape[-1]
+    if _can_triton(x) and weight.is_cuda and weight.is_contiguous():
+        y = torch.empty_like(x)
+        x2 = x.view(-1, d)
+        y2 = y.view(-1, d)
+        block = triton.next_power_of_2(d)
+        _rmsnorm_kernel[(x2.shape[0],)](x2, weight, y2, n_cols=d, eps=float(eps), BLOCK=block)
+        return y
+    xf = x.float()
+    y = xf * torch.rsqrt(xf.square().mean(dim=-1, keepdim=True) + eps)
+    y = y * (1.0 + weight.float())
+    return y.to(dtype=x.dtype)
+def add_rmsnorm(
+    x: torch.Tensor,
+    update: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fuse residual addition with Qwen3.5 offset RMSNorm.
+    Returns `(x + update, rmsnorm(x + update))`.
+    """
+    d = x.shape[-1]
+    if _can_triton(x) and update.is_contiguous() and weight.is_cuda and weight.is_contiguous():
+        summed = torch.empty_like(x)
+        normed = torch.empty_like(x)
+        x2 = x.view(-1, d)
+        u2 = update.view(-1, d)
+        s2 = summed.view(-1, d)
+        n2 = normed.view(-1, d)
+        block = triton.next_power_of_2(d)
+        _add_rmsnorm_kernel[(x2.shape[0],)](
+            x2, u2, weight, s2, n2, n_cols=d, eps=float(eps), BLOCK=block
+        )
+        return summed, normed
+    summed = x + update
+    xf = summed.float()
+    normed = xf * torch.rsqrt(xf.square().mean(dim=-1, keepdim=True) + eps)
+    normed = normed * (1.0 + weight.float())
+    return summed, normed.to(dtype=x.dtype)
+def add_final_rmsnorm(
+    x: torch.Tensor,
+    update: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """Fuse final residual addition and RMSNorm when the unnormalized sum is not needed."""
+    d = x.shape[-1]
+    if _can_triton(x) and update.is_contiguous() and weight.is_cuda and weight.is_contiguous():
+        normed = torch.empty_like(x)
+        x2 = x.view(-1, d)
+        u2 = update.view(-1, d)
+        n2 = normed.view(-1, d)
+        block = triton.next_power_of_2(d)
+        _add_final_rmsnorm_kernel[(x2.shape[0],)](
+            x2, u2, weight, n2, n_cols=d, eps=float(eps), BLOCK=block
+        )
+        return normed
+    summed = x + update
+    xf = summed.float()
+    normed = xf * torch.rsqrt(xf.square().mean(dim=-1, keepdim=True) + eps)
+    normed = normed * (1.0 + weight.float())
+    return normed.to(dtype=x.dtype)
+def silu_mul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """Fused SwiGLU pointwise core: SiLU(a) * b."""
+    if _can_triton(a) and b.is_contiguous() and a.shape == b.shape:
+        out = torch.empty_like(a)
+        n = a.numel()
+        block = 256
+        _silu_mul_kernel[(triton.cdiv(n, block),)](a, b, out, n_elements=n, BLOCK=block)
+        return out
+    return F.silu(a) * b
+class SuperRMSNorm(nn.Module):
+    """Checkpoint-compatible replacement for Qwen3_5RMSNorm."""
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = float(eps)
+        # Qwen3.5 uses zero-centered weights and multiplies by (1 + weight).
+        self.weight = nn.Parameter(torch.zeros(dim))
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return qwen_rmsnorm(x, self.weight, self.eps)
+    def extra_repr(self) -> str:
+        return f"{tuple(self.weight.shape)}, eps={self.eps}"
+class SuperSwiGLUMLP(nn.Module):
+    """Checkpoint-compatible Qwen3.5 dense MLP with a fused SwiGLU pointwise kernel."""
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.hidden_act = str(config.hidden_act)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate = self.gate_proj(x)
+        up = self.up_proj(x)
+        if self.hidden_act == "silu":
+            hidden = silu_mul(gate, up)
+        else:
+            from transformers.activations import ACT2FN
+            hidden = ACT2FN[self.hidden_act](gate) * up
+        return self.down_proj(hidden)

gdn24_metadata.json ADDED Viewed

	@@ -0,0 +1,26 @@

+{
+  "source_model": "Qwen/Qwen3.5-9B",
+  "architecture": "32x native Qwen3.5 GatedDeltaNet",
+  "converted_layers": [
+    3,
+    7,
+    11,
+    15,
+    19,
+    23,
+    27,
+    31
+  ],
+  "steps_per_converted_layer": 20,
+  "seq_len": 128,
+  "lr": 0.0001,
+  "seed": 1234,
+  "runtime_adapters": 0,
+  "uni_max_runtime": 1,
+  "phase_policy": "exact_bf16_prefill+autotuned_decode",
+  "fusions": [
+    "residual+rmsnorm",
+    "swiglu"
+  ],
+  "cuda_graph_decode_compatible": true
+}

generation_config.json ADDED Viewed

	@@ -0,0 +1,8 @@

+{
+  "_from_model_config": true,
+  "eos_token_id": 248044,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "transformers_version": "5.16.0.dev0",
+  "use_cache": true
+}

model-00001-of-00004.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:51ca0c41567cc617239c12d88281ca02552f6fffe0846ecd1e81f83f08d9b72b
+size 4991268352

model-00002-of-00004.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:84d4ac920ebc03c163d7347a18c9c43e26f583b0c078e44cc90516fa76a12251
+size 4939275072

model-00003-of-00004.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:c7b24fdcdf6f78f754e07f6335f3fa7dbdc651bb33002aa943158ff0c1f634dd
+size 4906162280

model-00004-of-00004.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:c7bdd2633dc3b0f8ba6434550d3aab1f0d44dffb3ea88dffdf2517809d1a9cf8
+size 3209884384

model.safetensors.index.json ADDED Viewed

	@@ -0,0 +1,459 @@

+{
+  "metadata": {
+    "total_parameters": 9023268864,
+    "total_size": 18046537728
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00004-of-00004.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "model.layers.0.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "model.layers.0.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "model.layers.1.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "model.layers.1.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.layers.10.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.layers.10.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.layers.11.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.layers.11.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.layers.12.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.layers.12.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.layers.13.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.layers.13.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.layers.14.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.layers.14.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.layers.15.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.layers.15.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.layers.16.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.layers.16.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.layers.17.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.layers.17.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.layers.18.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.layers.18.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.layers.19.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.layers.19.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "model.layers.2.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "model.layers.2.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.layers.20.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.layers.20.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.layers.21.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.layers.21.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.layers.22.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.layers.22.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.layers.23.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.layers.23.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.layers.24.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.layers.24.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.layers.25.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.layers.25.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.layers.26.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.layers.26.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.layers.27.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.layers.27.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.layers.28.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.layers.28.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.layers.29.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.layers.29.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "model.layers.3.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "model.layers.3.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.linear_attn.A_log": "model-00004-of-00004.safetensors",
+    "model.layers.30.linear_attn.conv1d.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.linear_attn.dt_bias": "model-00004-of-00004.safetensors",
+    "model.layers.30.linear_attn.in_proj_a.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.linear_attn.in_proj_b.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.linear_attn.in_proj_qkv.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.linear_attn.in_proj_z.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.linear_attn.norm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.linear_attn.out_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.linear_attn.A_log": "model-00004-of-00004.safetensors",
+    "model.layers.31.linear_attn.conv1d.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.linear_attn.dt_bias": "model-00004-of-00004.safetensors",
+    "model.layers.31.linear_attn.in_proj_a.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.linear_attn.in_proj_b.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.linear_attn.in_proj_qkv.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.linear_attn.in_proj_z.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.linear_attn.norm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.linear_attn.out_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "model.layers.4.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "model.layers.4.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "model.layers.5.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "model.layers.5.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.6.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "model.layers.6.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "model.layers.6.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.layers.7.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.layers.7.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.layers.8.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.layers.8.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.layers.9.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.layers.9.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.norm.weight": "model-00004-of-00004.safetensors"
+  }
+}

modeling_qwen35_gdn24.py ADDED Viewed

	@@ -0,0 +1,338 @@

+from __future__ import annotations
+from typing import Any
+import torch
+from torch import nn
+from transformers.cache_utils import Cache, LinearAttentionLayer
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_recurrent_attention_mask
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.models.qwen3_5.modeling_qwen3_5 import (
+    Qwen3_5GatedDeltaNet,
+    Qwen3_5PreTrainedModel,
+)
+from .configuration_qwen35_gdn24 import Qwen35GDN24Config
+from .fused_ops import SuperRMSNorm, SuperSwiGLUMLP, add_final_rmsnorm, add_rmsnorm
+class GDN24Cache(Cache):
+    """24 native LinearAttentionLayer states plus an explicit logical token clock.
+    Hugging Face's generic all-linear DynamicCache has no attention layer from
+    which to infer sequence length.  GDN itself does not need K/V length, but
+    GenerationMixin does need a logical prefix length.  We therefore keep the
+    recurrent states native and add only a scalar clock.
+    """
+    _qwen35_gdn24_cache_protocol = 1
+    is_compileable = False
+    def __init__(self, config: Qwen35GDN24Config):
+        number_of_states = getattr(config, "number_of_conv_states", 1)
+        layers = [
+            LinearAttentionLayer(number_of_states=number_of_states)
+            for _ in range(config.num_hidden_layers)
+        ]
+        super().__init__(layers=layers)
+        self.seen_tokens = 0
+    def advance(self, token_count: int) -> None:
+        self.seen_tokens += int(token_count)
+    def get_seq_length(self, layer_idx: int = 0) -> int:
+        return int(self.seen_tokens)
+    def get_query_offset(self, layer_idx: int = 0) -> int:
+        return int(self.seen_tokens)
+    def get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]:
+        # Every layer is recurrent; no K/V sequence dimension is materialized.
+        return int(query_length), 0
+    def get_max_length(self, layer_idx: int | None = None) -> int:
+        return -1
+    def reset(self) -> None:
+        super().reset()
+        self.seen_tokens = 0
+    def crop(self, tokens_to_remove: int) -> None:
+        if tokens_to_remove == 0:
+            return
+        raise NotImplementedError(
+            "GDN24 recurrent state is not invertible; speculative rollback requires state snapshots."
+        )
+class Qwen35GDN24DecoderLayer(nn.Module):
+    """Qwen3.5 decoder block with Gated DeltaNet as the token mixer in every layer."""
+    def __init__(self, config: Qwen35GDN24Config, layer_idx: int):
+        super().__init__()
+        self.layer_idx = int(layer_idx)
+        self.linear_attn = Qwen3_5GatedDeltaNet(config, layer_idx)
+        self.mlp = SuperSwiGLUMLP(config)
+        self.input_layernorm = SuperRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = SuperRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.linear_attn(
+            hidden_states=hidden_states,
+            cache_params=past_key_values,
+            attention_mask=attention_mask,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        return residual + hidden_states
+class Qwen35GDN24PreTrainedModel(Qwen3_5PreTrainedModel):
+    config_class = Qwen35GDN24Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_static_cache = False
+    _is_stateful = True
+    _no_split_modules = ["Qwen35GDN24DecoderLayer"]
+class Qwen35GDN24Model(Qwen35GDN24PreTrainedModel):
+    def __init__(self, config: Qwen35GDN24Config):
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+        self.layers = nn.ModuleList(
+            [Qwen35GDN24DecoderLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.norm = SuperRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # No rotary module: all 24 mixers are recurrent and position is implicit in state evolution.
+        self.post_init()
+    def get_input_embeddings(self):
+        return self.embed_tokens
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+    def make_recurrent_cache(self) -> GDN24Cache:
+        return GDN24Cache(self.config)
+    @staticmethod
+    def _compatible_cache(cache: object) -> bool:
+        return (
+            isinstance(cache, Cache)
+            and getattr(cache, "_qwen35_gdn24_cache_protocol", None) == 1
+            and hasattr(cache, "advance")
+            and hasattr(cache, "layers")
+        )
+    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 | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        _advance_cache_clock: bool = True,
+        **kwargs: Any,
+    ) -> BaseModelOutputWithPast | tuple:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds")
+        use_cache = self.config.use_cache if use_cache is None else use_cache
+        output_hidden_states = (
+            self.config.output_hidden_states if output_hidden_states is None else output_hidden_states
+        )
+        return_dict = self.config.return_dict if return_dict is None else return_dict
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        seq_len = inputs_embeds.shape[1]
+        if use_cache and past_key_values is None:
+            past_key_values = self.make_recurrent_cache()
+        elif past_key_values is not None and not self._compatible_cache(past_key_values):
+            raise TypeError(
+                "Qwen35GDN24 requires its recurrent-cache protocol object; "
+                f"got {type(past_key_values).__module__}.{type(past_key_values).__name__}. "
+                "Use model.make_recurrent_cache()."
+            )
+        # Common generation/prefill fast path: no padding means no recurrent mask
+        # construction at all.  This removes Python work from every single-token
+        # decode step and is exactly equivalent to create_recurrent_attention_mask(None).
+        recurrent_mask = None
+        if attention_mask is not None:
+            recurrent_mask = create_recurrent_attention_mask(
+                config=self.config,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+            )
+        # MAX TURBO residual pipeline.  It is algebraically identical to the
+        # standard pre-norm decoder, but fuses residual addition with the next
+        # RMSNorm and fuses SwiGLU's pointwise core.
+        hidden_states = inputs_embeds
+        all_hidden_states = () if output_hidden_states else None
+        if len(self.layers) > 0:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            normed = self.layers[0].input_layernorm(hidden_states)
+            for i, layer in enumerate(self.layers):
+                mixed = layer.linear_attn(
+                    hidden_states=normed,
+                    cache_params=past_key_values,
+                    attention_mask=recurrent_mask,
+                    use_cache=use_cache,
+                    **kwargs,
+                )
+                attn_residual, mlp_input = add_rmsnorm(
+                    hidden_states,
+                    mixed,
+                    layer.post_attention_layernorm.weight,
+                    layer.post_attention_layernorm.eps,
+                )
+                mlp_update = layer.mlp(mlp_input)
+                if i + 1 < len(self.layers):
+                    next_layer = self.layers[i + 1]
+                    hidden_states, normed = add_rmsnorm(
+                        attn_residual,
+                        mlp_update,
+                        next_layer.input_layernorm.weight,
+                        next_layer.input_layernorm.eps,
+                    )
+                    if output_hidden_states:
+                        all_hidden_states += (hidden_states,)
+                else:
+                    hidden_states = add_final_rmsnorm(
+                        attn_residual, mlp_update, self.norm.weight, self.norm.eps
+                    )
+        else:
+            hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+        # CUDA Graph replay cannot rerun Python scalar mutations.  MAX TURBO
+        # passes _advance_cache_clock=False inside a captured graph and advances
+        # the logical clock once per replay from the host.
+        if use_cache and past_key_values is not None and _advance_cache_clock:
+            past_key_values.advance(seq_len)
+        if not return_dict:
+            values = (hidden_states, past_key_values)
+            if output_hidden_states:
+                values += (all_hidden_states,)
+            return values
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states,
+            attentions=None,
+        )
+class Qwen35GDN24ForCausalLM(Qwen35GDN24PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _keys_to_ignore_on_load_unexpected = [r"^model\.visual.*", r"^mtp.*"]
+    @classmethod
+    def _supports_default_dynamic_cache(cls) -> bool:
+        return False
+    def __init__(self, config: Qwen35GDN24Config):
+        super().__init__(config)
+        self.model = Qwen35GDN24Model(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):
+        return self.model.embed_tokens
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+    def get_output_embeddings(self):
+        return self.lm_head
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+    def make_recurrent_cache(self) -> GDN24Cache:
+        return self.model.make_recurrent_cache()
+    @torch.inference_mode()
+    def greedy_step(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Cache,
+        *,
+        advance_cache_clock: bool = True,
+    ) -> torch.LongTensor:
+        outputs = self.model(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            use_cache=True,
+            _advance_cache_clock=advance_cache_clock,
+        )
+        logits = self.lm_head(outputs.last_hidden_state[:, -1, :])
+        return torch.argmax(logits, dim=-1, keepdim=True)
+    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 | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Any,
+    ) -> CausalLMOutputWithPast:
+        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,
+            use_cache=use_cache,
+            **kwargs,
+        )
+        hidden_states = outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=None,
+        )
+Qwen35GDN24ForCausalLM.register_for_auto_class("AutoModelForCausalLM")

quantization.py ADDED Viewed

	@@ -0,0 +1,230 @@

+from __future__ import annotations
+"""Runtime-only quantization policies for UNI MAX.
+UNI MAX v1.1 distinguishes two very different notions:
+* ``full`` quantization (MAX TURBO compatible): every dense MLP projection plus
+  the LM head.  This is fast, but it changes the hidden trajectory that writes
+  future GDN recurrent states.  A cache copied from a BF16-prefill model is
+  therefore *not* generally a valid state for this quantized dynamical system.
+* ``state_safe`` quantization (UNI MAX): only the LM head and the MLP of the
+  final decoder layer.  These operators execute strictly after the final
+  persistent GDN state update for each token.  Hence they cannot change any
+  conv/recurrent cache tensor.  If their greedy token agrees with the exact
+  model, the next-token recurrent state remains exactly synchronized.
+Quantization is inference-only and is never written back into the checkpoint.
+"""
+from dataclasses import dataclass
+from typing import Callable
+import torch
+from torch import nn
+@dataclass(frozen=True)
+class QuantizationReport:
+    mode: str
+    applied: bool
+    reason: str | None
+    targeted_linears: int
+    cuda_capability: tuple[int, int] | None
+    policy: str = "full"
+def _cuda_capability() -> tuple[int, int] | None:
+    if not torch.cuda.is_available():
+        return None
+    return tuple(int(x) for x in torch.cuda.get_device_capability())
+def _is_dense_mlp_linear(module: nn.Module, fqn: str) -> bool:
+    if not isinstance(module, nn.Linear):
+        return False
+    return ".mlp." in fqn and fqn.rsplit(".", 1)[-1] in {"gate_proj", "up_proj", "down_proj"}
+def _target_dense_decode_linears(module: nn.Module, fqn: str) -> bool:
+    """MAX-compatible full dense target set; native GDN projections excluded."""
+    if not isinstance(module, nn.Linear):
+        return False
+    return fqn == "lm_head" or _is_dense_mlp_linear(module, fqn)
+def _final_layer_index(model: nn.Module) -> int:
+    cfg = getattr(model, "config", None)
+    n = getattr(cfg, "num_hidden_layers", None)
+    if n is not None:
+        return int(n) - 1
+    layers = getattr(getattr(model, "model", None), "layers", None)
+    if layers is None:
+        raise ValueError("cannot determine final decoder layer index")
+    return len(layers) - 1
+def make_state_safe_filter(model: nn.Module) -> Callable[[nn.Module, str], bool]:
+    """Return the causal-state-compatible UNI quantization filter.
+    The eligible set is exactly:
+      * ``lm_head``
+      * ``model.layers.<last>.mlp.{gate_proj,up_proj,down_proj}``
+    No module whose output can influence a persistent recurrent-state write is
+    eligible.
+    """
+    last = _final_layer_index(model)
+    prefix = f"model.layers.{last}.mlp."
+    def _filter(module: nn.Module, fqn: str) -> bool:
+        if not isinstance(module, nn.Linear):
+            return False
+        if fqn == "lm_head":
+            return True
+        return fqn.startswith(prefix) and fqn.rsplit(".", 1)[-1] in {
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        }
+    return _filter
+def count_target_linears(model: nn.Module) -> int:
+    return sum(1 for fqn, mod in model.named_modules() if _target_dense_decode_linears(mod, fqn))
+def count_state_safe_linears(model: nn.Module) -> int:
+    filt = make_state_safe_filter(model)
+    return sum(1 for fqn, mod in model.named_modules() if filt(mod, fqn))
+def state_safe_target_names(model: nn.Module) -> tuple[str, ...]:
+    filt = make_state_safe_filter(model)
+    return tuple(fqn for fqn, mod in model.named_modules() if filt(mod, fqn))
+def available_quant_modes() -> tuple[str, ...]:
+    modes = ["exact"]
+    try:
+        import torchao  # noqa: F401
+        modes.extend(["fp8", "int8"])
+    except Exception:
+        pass
+    return tuple(modes)
+def _apply_quantization(model: nn.Module, mode: str, *, policy: str) -> QuantizationReport:
+    mode = str(mode).lower().strip()
+    policy = str(policy).lower().strip()
+    capability = _cuda_capability()
+    if policy == "state_safe":
+        filter_fn = make_state_safe_filter(model)
+        targeted = count_state_safe_linears(model)
+    elif policy == "full":
+        filter_fn = _target_dense_decode_linears
+        targeted = count_target_linears(model)
+    else:
+        return QuantizationReport(mode, False, f"unknown quantization policy: {policy}", 0, capability, policy)
+    if mode in {"", "none", "exact", "bf16"}:
+        return QuantizationReport("exact", False, None, targeted, capability, policy)
+    try:
+        from torchao.quantization import quantize_
+    except Exception as exc:  # pragma: no cover - runtime dependent
+        return QuantizationReport(
+            mode,
+            False,
+            f"torchao unavailable: {type(exc).__name__}: {exc}",
+            targeted,
+            capability,
+            policy,
+        )
+    if targeted == 0:
+        return QuantizationReport(mode, False, "no eligible dense linear layers found", 0, capability, policy)
+    try:
+        if mode == "fp8":
+            if capability is None or capability < (8, 9):
+                return QuantizationReport(
+                    mode,
+                    False,
+                    f"FP8 requires CUDA SM 8.9+, got {capability}",
+                    targeted,
+                    capability,
+                    policy,
+                )
+            from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, PerTensor
+            config = Float8DynamicActivationFloat8WeightConfig(granularity=PerTensor())
+        elif mode == "int8":
+            from torchao.quantization import Int8WeightOnlyConfig
+            config = Int8WeightOnlyConfig()
+        else:
+            return QuantizationReport(
+                mode,
+                False,
+                f"unknown quantization mode: {mode}",
+                targeted,
+                capability,
+                policy,
+            )
+        quantize_(model, config, filter_fn=filter_fn)
+        return QuantizationReport(mode, True, None, targeted, capability, policy)
+    except Exception as exc:  # pragma: no cover - hardware/runtime specific
+        return QuantizationReport(
+            mode,
+            False,
+            f"{type(exc).__name__}: {exc}",
+            targeted,
+            capability,
+            policy,
+        )
+def apply_max_quantization(model: nn.Module, mode: str) -> QuantizationReport:
+    """Legacy MAX TURBO policy: all MLP projections + LM head."""
+    return _apply_quantization(model, mode, policy="full")
+def apply_uni_quantization(model: nn.Module, mode: str) -> QuantizationReport:
+    """UNI MAX state-compatible policy: final MLP + LM head only."""
+    return _apply_quantization(model, mode, policy="state_safe")
+@torch.inference_mode()
+def greedy_token_trace(model, input_ids: torch.Tensor, steps: int = 32) -> torch.Tensor:
+    steps = int(steps)
+    if steps < 1:
+        return torch.empty((input_ids.shape[0], 0), dtype=torch.long, device=input_ids.device)
+    cache = model.make_recurrent_cache()
+    out = model(input_ids=input_ids, past_key_values=cache, use_cache=True, logits_to_keep=1)
+    token = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
+    tokens = [token]
+    for _ in range(steps - 1):
+        token = model.greedy_step(token, cache)
+        tokens.append(token)
+    return torch.cat(tokens, dim=1)
+def token_agreement(reference: torch.Tensor, candidate: torch.Tensor) -> tuple[float, int]:
+    if reference.shape != candidate.shape:
+        raise ValueError(f"token trace shape mismatch: {reference.shape} vs {candidate.shape}")
+    if reference.numel() == 0:
+        return 1.0, 0
+    eq = reference.eq(candidate)
+    agreement = float(eq.float().mean().item())
+    flat = eq.reshape(-1).tolist()
+    prefix = 0
+    for ok in flat:
+        if not ok:
+            break
+        prefix += 1
+    return agreement, prefix

tokenizer.json ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:8818dc7a3be5f461790e3a81703816f482925f6c2ff9fef5a9fc4b821e5051f2
+size 19989423

tokenizer_config.json ADDED Viewed

	@@ -0,0 +1,32 @@

+{
+  "add_prefix_space": false,
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "image_token": "<|image_pad|>",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 262144,
+  "model_specific_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "pad_token": "<|endoftext|>",
+  "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}

uni_engine.py ADDED Viewed

	@@ -0,0 +1,290 @@

+from __future__ import annotations
+"""UNI MAX: unified exact-prefill / quantized-decode recurrent inference.
+The GDN24 cache topology is identical between the exact BF16 model and a model
+whose **state-safe tail** (final MLP + LM head) is quantized. UNI MAX v1.1 exploits that causal invariant:
+1. prefill the prompt with the exact BF16 model;
+2. copy only the fixed recurrent cache state into a state-compatible persistent decode engine;
+3. decode with the fastest CUDA-Graph candidate (typically FP8 on NVIDIA L4);
+4. optionally verify FP8 draft blocks with exact BF16 chunk forwards to recover
+   exact greedy-token semantics.
+No KV sequence is copied: the bridge is O(1) in context length. Full-MLP quantization is deliberately excluded from this bridge because it changes the hidden trajectory that writes future recurrent states.
+"""
+from dataclasses import dataclass
+import time
+from typing import Any
+import torch
+from .engine import MaxTurboGraphDecoder, snapshot_cache, restore_cache_
+@dataclass(frozen=True)
+class BridgeReport:
+    seconds: float
+    mib: float
+    seen_tokens: int
+@dataclass(frozen=True)
+class VerifyReport:
+    generated_tokens: int
+    drafted_tokens: int
+    accepted_draft_tokens: int
+    rejected_blocks: int
+    verifier_blocks: int
+    @property
+    def acceptance_rate(self) -> float:
+        if self.drafted_tokens <= 0:
+            return 1.0
+        return self.accepted_draft_tokens / self.drafted_tokens
+def cache_payload_bytes(cache) -> int:
+    total = 0
+    for layer in cache.layers:
+        for table_name in ("conv_states", "recurrent_states"):
+            table = getattr(layer, table_name, {})
+            for tensor in table.values():
+                if tensor is not None:
+                    total += tensor.numel() * tensor.element_size()
+    return int(total)
+@torch.inference_mode()
+def copy_cache_(dst, src) -> None:
+    """Copy recurrent state without reallocating destination tensors.
+    Destination storage must already be materialized (CUDA-Graph capture does
+    this).  Tensor addresses are preserved, so captured graphs remain valid.
+    """
+    if len(dst.layers) != len(src.layers):
+        raise ValueError("cache topology mismatch")
+    for d_layer, s_layer in zip(dst.layers, src.layers):
+        for table_name in ("conv_states", "recurrent_states"):
+            d_table = getattr(d_layer, table_name)
+            s_table = getattr(s_layer, table_name)
+            for idx, s_tensor in s_table.items():
+                if s_tensor is None:
+                    continue
+                d_tensor = d_table.get(idx)
+                if d_tensor is None:
+                    raise RuntimeError(
+                        f"destination cache storage is not materialized: {table_name}[{idx}]"
+                    )
+                if d_tensor.shape != s_tensor.shape or d_tensor.dtype != s_tensor.dtype:
+                    raise RuntimeError(
+                        f"cache state mismatch for {table_name}[{idx}]: "
+                        f"dst={tuple(d_tensor.shape)}/{d_tensor.dtype}, "
+                        f"src={tuple(s_tensor.shape)}/{s_tensor.dtype}"
+                    )
+                d_tensor.copy_(s_tensor)
+        d_layer.has_previous_state.clear()
+        d_layer.has_previous_state.update(dict(s_layer.has_previous_state))
+        d_layer.is_conv_states_initialized.clear()
+        d_layer.is_conv_states_initialized.update(dict(s_layer.is_conv_states_initialized))
+        d_layer.is_recurrent_states_initialized.clear()
+        d_layer.is_recurrent_states_initialized.update(dict(s_layer.is_recurrent_states_initialized))
+    dst.seen_tokens = int(src.seen_tokens)
+class UniMaxEngine:
+    """Phase-specialized recurrent inference engine.
+    ``exact_model`` always handles prefill.  ``decode_model`` can be the same model (UNI-SAFE) or a state-safe quantized clone (UNI-SPEED / UNI-EXACT).
+    """
+    def __init__(
+        self,
+        exact_model,
+        decode_model,
+        decode_graph: MaxTurboGraphDecoder,
+    ):
+        self.exact_model = exact_model.eval()
+        self.decode_model = decode_model.eval()
+        self.decode_graph = decode_graph
+        self.device = self.exact_model.get_input_embeddings().weight.device
+        if self.device != self.decode_graph.device:
+            raise ValueError("exact and decode engines must live on the same CUDA device")
+        self.exact_cache = self.exact_model.make_recurrent_cache()
+        self.last_bridge: BridgeReport | None = None
+    @torch.inference_mode()
+    def reset(self) -> None:
+        self.exact_cache.reset()
+        self.decode_graph.reset()
+        self.last_bridge = None
+    @torch.inference_mode()
+    def _bridge(self) -> BridgeReport:
+        if self.device.type == "cuda":
+            torch.cuda.synchronize(self.device)
+        t0 = time.perf_counter()
+        copy_cache_(self.decode_graph.cache, self.exact_cache)
+        if self.device.type == "cuda":
+            torch.cuda.synchronize(self.device)
+        dt = time.perf_counter() - t0
+        report = BridgeReport(
+            seconds=float(dt),
+            mib=cache_payload_bytes(self.exact_cache) / 2**20,
+            seen_tokens=int(self.exact_cache.seen_tokens),
+        )
+        self.last_bridge = report
+        return report
+    @torch.inference_mode()
+    def prefill_exact(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, Any, BridgeReport]:
+        """Exact BF16 prefill, then O(1)-context state-compatible recurrent bridge."""
+        self.exact_cache.reset()
+        out = self.exact_model(
+            input_ids=input_ids,
+            past_key_values=self.exact_cache,
+            use_cache=True,
+            logits_to_keep=1,
+        )
+        first = torch.argmax(out.logits[:, -1, :], dim=-1, keepdim=True)
+        bridge = self._bridge()
+        self.decode_graph.static_token.copy_(first)
+        return first, out, bridge
+    @torch.inference_mode()
+    def decode_fast(self, first_token: torch.Tensor, recurrent_forwards: int) -> torch.Tensor:
+        return self.decode_graph.decode_forwards(first_token, int(recurrent_forwards))
+    @torch.inference_mode()
+    def generate_fast(self, input_ids: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
+        """UNI-SPEED generation: exact first token, quantized graph thereafter."""
+        n = int(max_new_tokens)
+        if n <= 0:
+            return torch.empty((input_ids.shape[0], 0), dtype=torch.long, device=input_ids.device)
+        first, _, _ = self.prefill_exact(input_ids)
+        if n == 1:
+            return first
+        tail = self.decode_graph.decode_tokens(first, n - 1)
+        return torch.cat([first, tail], dim=1)
+    @torch.inference_mode()
+    def decode_verified(
+        self,
+        first_token: torch.Tensor,
+        recurrent_forwards: int,
+        *,
+        draft_block: int | None = None,
+    ) -> tuple[torch.Tensor, VerifyReport]:
+        """Verify ``recurrent_forwards`` tokens after ``first_token``.
+        ``prefill_exact`` must have been called immediately before this method so
+        both exact and decode caches represent the same prompt state.
+        """
+        remaining = int(recurrent_forwards)
+        if remaining <= 0:
+            empty = torch.empty((first_token.shape[0], 0), dtype=torch.long, device=first_token.device)
+            return empty, VerifyReport(0, 0, 0, 0, 0)
+        if first_token.shape[0] != 1:
+            raise ValueError("UNI-EXACT currently supports batch size 1")
+        current = first_token
+        generated: list[torch.Tensor] = []
+        k_default = max(1, int(draft_block or self.decode_graph.block_size))
+        drafted = accepted_drafts = rejected_blocks = verifier_blocks = 0
+        while remaining > 0:
+            k = min(k_default, remaining)
+            exact_before = snapshot_cache(self.exact_cache)
+            draft = self.decode_graph.decode_tokens(current, k)
+            drafted += k
+            verify_input = current if k == 1 else torch.cat([current, draft[:, :-1]], dim=1)
+            verify_out = self.exact_model(
+                input_ids=verify_input,
+                past_key_values=self.exact_cache,
+                use_cache=True,
+                logits_to_keep=k,
+            )
+            exact_pred = torch.argmax(verify_out.logits[:, -k:, :], dim=-1)
+            verifier_blocks += 1
+            mismatch_positions = (~exact_pred.eq(draft)[0]).nonzero(as_tuple=False)
+            if mismatch_positions.numel() == 0:
+                generated.append(draft.detach().clone())
+                accepted_drafts += k
+                current = draft[:, -1:]
+                remaining -= k
+                continue
+            m = int(mismatch_positions[0, 0].item())
+            accepted_drafts += m
+            rejected_blocks += 1
+            restore_cache_(self.exact_cache, exact_before)
+            replay_input = current if m == 0 else torch.cat([current, draft[:, :m]], dim=1)
+            replay = self.exact_model(
+                input_ids=replay_input,
+                past_key_values=self.exact_cache,
+                use_cache=True,
+                logits_to_keep=1,
+            )
+            corrected = torch.argmax(replay.logits[:, -1, :], dim=-1, keepdim=True)
+            if m:
+                generated.append(draft[:, :m].detach().clone())
+            generated.append(corrected.detach().clone())
+            current = corrected
+            emitted = m + 1
+            remaining -= emitted
+            copy_cache_(self.decode_graph.cache, self.exact_cache)
+            self.decode_graph.static_token.copy_(current)
+        out = torch.cat(generated, dim=1)
+        return out, VerifyReport(
+            generated_tokens=int(out.shape[1]),
+            drafted_tokens=int(drafted),
+            accepted_draft_tokens=int(accepted_drafts),
+            rejected_blocks=int(rejected_blocks),
+            verifier_blocks=int(verifier_blocks),
+        )
+    @torch.inference_mode()
+    def generate_verified(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int,
+        *,
+        draft_block: int | None = None,
+    ) -> tuple[torch.Tensor, VerifyReport]:
+        """UNI-EXACT generation with exact greedy-token semantics."""
+        n = int(max_new_tokens)
+        if n <= 0:
+            empty = torch.empty((input_ids.shape[0], 0), dtype=torch.long, device=input_ids.device)
+            return empty, VerifyReport(0, 0, 0, 0, 0)
+        first, _, _ = self.prefill_exact(input_ids)
+        if n == 1:
+            return first, VerifyReport(1, 0, 0, 0, 0)
+        tail, report = self.decode_verified(first, n - 1, draft_block=draft_block)
+        full = torch.cat([first, tail], dim=1)
+        return full, VerifyReport(
+            generated_tokens=int(full.shape[1]),
+            drafted_tokens=report.drafted_tokens,
+            accepted_draft_tokens=report.accepted_draft_tokens,
+            rejected_blocks=report.rejected_blocks,
+            verifier_blocks=report.verifier_blocks,
+        )
+@torch.inference_mode()
+def exact_greedy_tokens(model, input_ids: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
+    n = int(max_new_tokens)
+    if n <= 0:
+        return torch.empty((input_ids.shape[0], 0), dtype=torch.long, device=input_ids.device)
+    cache = model.make_recurrent_cache()
+    out = model(input_ids=input_ids, past_key_values=cache, use_cache=True, logits_to_keep=1)
+    token = torch.argmax(out.logits[:, -1, :], dim=-1, keepdim=True)
+    pieces = [token]
+    for _ in range(n - 1):
+        token = model.greedy_step(token, cache)
+        pieces.append(token)
+    return torch.cat(pieces, dim=1)