majentik commited on
Commit
f1e16a4
·
verified ·
1 Parent(s): f587fae

Add MLX quantized model with KV cache compression

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,103 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ base_model: google/gemma-4-E4B
3
+ library_name: mlx
4
+ tags:
5
+ - rotorquant
6
+ - kv-cache-quantization
7
+ - gemma
8
+ - gemma4
9
+ - multimodal
10
+ - quantized
11
+ - mlx
12
+ - 4bit
13
+ license: apache-2.0
14
+ pipeline_tag: image-text-to-text
15
+ ---
16
+
17
+ # Gemma 4 E4B - RotorQuant MLX 4-bit
18
+
19
+ **4-bit weight-quantized MLX version** of [google/gemma-4-E4B](https://huggingface.co/google/gemma-4-E4B) with RotorQuant KV-cache quantization. Optimized for Apple Silicon inference via the [MLX](https://github.com/ml-explore/mlx) framework. RotorQuant delivers 5.3x faster prefill and 28% faster decode compared to TurboQuant. A good balance between model quality and memory efficiency.
20
+
21
+ Approximate model size: **~2.3 GB**
22
+
23
+ ## Model Specifications
24
+
25
+ | Property | Value |
26
+ |---|---|
27
+ | **Base Model** | [google/gemma-4-E4B](https://huggingface.co/google/gemma-4-E4B) |
28
+ | **Parameters** | ~4 billion |
29
+ | **Architecture** | Dense transformer |
30
+ | **Modality** | Multimodal: image + text input, text output |
31
+ | **License** | Apache 2.0 |
32
+ | **Weight Quantization** | 4-bit (~2.3 GB) |
33
+ | **KV-Cache Quantization** | RotorQuant |
34
+ | **Framework** | MLX (Apple Silicon) |
35
+
36
+ ## Quickstart
37
+
38
+ ```python
39
+ import mlx.core as mx
40
+ from mlx_lm import load, generate
41
+
42
+ model, tokenizer = load("majentik/gemma-4-E4B-RotorQuant-MLX-4bit")
43
+
44
+ prompt = "The history of artificial intelligence began"
45
+ response = generate(model, tokenizer, prompt=prompt, max_tokens=512)
46
+ print(response)
47
+ ```
48
+
49
+ For multimodal usage with images:
50
+
51
+ ```python
52
+ from mlx_vlm import load, generate
53
+
54
+ model, processor = load("majentik/gemma-4-E4B-RotorQuant-MLX-4bit")
55
+
56
+ prompt = "Describe the contents of this image."
57
+ output = generate(model, processor, prompt=prompt, image="path/to/image.jpg", max_tokens=512)
58
+ print(output)
59
+ ```
60
+
61
+ ## What is RotorQuant?
62
+
63
+ [RotorQuant](https://github.com/scrya-com/rotorquant) is a high-performance KV-cache quantization method that achieves significantly better throughput than TurboQuant. Combined with 4-bit weight quantization in MLX, this provides a dual compression strategy with superior KV-cache performance: smaller model weights plus faster compressed KV cache for efficient long-context generation.
64
+
65
+ Key advantages over TurboQuant:
66
+ - **5.3x faster prefill**
67
+ - **28% faster decode**
68
+ - Equivalent memory savings
69
+
70
+ ## KV-Cache Quantization Comparison
71
+
72
+ | Method | Prefill Speed | Decode Speed | Memory Savings | Reference |
73
+ |---|---|---|---|---|
74
+ | **TurboQuant** | 1x (baseline) | 1x (baseline) | High | [arXiv: 2504.19874](https://arxiv.org/abs/2504.19874) |
75
+ | **RotorQuant** | **5.3x faster** | **28% faster** | High | [GitHub](https://github.com/scrya-com/rotorquant) |
76
+
77
+ ## Memory Estimates (Gemma 4 E4B)
78
+
79
+ | Precision | Approximate Size | MLX Variant |
80
+ |---|---|---|
81
+ | FP16 (original) | ~8 GB | -- |
82
+ | 8-bit quantized | ~4 GB | [RotorQuant-MLX-8bit](https://huggingface.co/majentik/gemma-4-E4B-RotorQuant-MLX-8bit) |
83
+ | **4-bit quantized** | **~2.3 GB** | **This model** |
84
+ | 2-bit quantized | ~1.2 GB | [RotorQuant-MLX-2bit](https://huggingface.co/majentik/gemma-4-E4B-RotorQuant-MLX-2bit) |
85
+
86
+ ## Hardware Requirements
87
+
88
+ This model requires approximately 2.3 GB of unified memory. Recommended hardware:
89
+ - Apple M1 (8 GB+)
90
+ - Apple M2 (8 GB+)
91
+ - Apple M3 (8 GB+)
92
+ - Apple M4 (8 GB+)
93
+ - Any Apple Silicon Mac with 8 GB+ unified memory
94
+
95
+ ## See Also
96
+
97
+ - [google/gemma-4-E4B](https://huggingface.co/google/gemma-4-E4B) -- Base model
98
+ - [majentik/gemma-4-E4B-RotorQuant](https://huggingface.co/majentik/gemma-4-E4B-RotorQuant) -- RotorQuant KV-cache only (transformers)
99
+ - [majentik/gemma-4-E4B-RotorQuant-MLX-8bit](https://huggingface.co/majentik/gemma-4-E4B-RotorQuant-MLX-8bit) -- MLX 8-bit variant
100
+ - [majentik/gemma-4-E4B-RotorQuant-MLX-2bit](https://huggingface.co/majentik/gemma-4-E4B-RotorQuant-MLX-2bit) -- MLX 2-bit variant
101
+ - [majentik/gemma-4-E4B-TurboQuant-MLX-4bit](https://huggingface.co/majentik/gemma-4-E4B-TurboQuant-MLX-4bit) -- TurboQuant MLX 4-bit variant
102
+ - [RotorQuant GitHub](https://github.com/scrya-com/rotorquant)
103
+ - [MLX Framework](https://github.com/ml-explore/mlx)
config.json ADDED
@@ -0,0 +1,204 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "Gemma4ForConditionalGeneration"
4
+ ],
5
+ "audio_config": {
6
+ "_name_or_path": "",
7
+ "architectures": null,
8
+ "attention_chunk_size": 12,
9
+ "attention_context_left": 13,
10
+ "attention_context_right": 0,
11
+ "attention_invalid_logits_value": -1000000000.0,
12
+ "attention_logit_cap": 50.0,
13
+ "chunk_size_feed_forward": 0,
14
+ "conv_kernel_size": 5,
15
+ "dtype": "bfloat16",
16
+ "gradient_clipping": 10000000000.0,
17
+ "hidden_act": "silu",
18
+ "hidden_size": 1024,
19
+ "id2label": {
20
+ "0": "LABEL_0",
21
+ "1": "LABEL_1"
22
+ },
23
+ "initializer_range": 0.02,
24
+ "is_encoder_decoder": false,
25
+ "label2id": {
26
+ "LABEL_0": 0,
27
+ "LABEL_1": 1
28
+ },
29
+ "model_type": "gemma4_audio",
30
+ "num_attention_heads": 8,
31
+ "num_hidden_layers": 12,
32
+ "output_attentions": false,
33
+ "output_hidden_states": false,
34
+ "output_proj_dims": 1536,
35
+ "problem_type": null,
36
+ "residual_weight": 0.5,
37
+ "return_dict": true,
38
+ "rms_norm_eps": 1e-06,
39
+ "subsampling_conv_channels": [
40
+ 128,
41
+ 32
42
+ ],
43
+ "use_clipped_linears": true
44
+ },
45
+ "audio_token_id": 258881,
46
+ "boa_token_id": 256000,
47
+ "boi_token_id": 255999,
48
+ "dtype": "bfloat16",
49
+ "eoa_token_id": 258883,
50
+ "eoa_token_index": 258883,
51
+ "eoi_token_id": 258882,
52
+ "eos_token_id": 1,
53
+ "image_token_id": 258880,
54
+ "initializer_range": 0.02,
55
+ "model_type": "gemma4",
56
+ "quantization": {
57
+ "group_size": 64,
58
+ "bits": 4,
59
+ "mode": "affine"
60
+ },
61
+ "quantization_config": {
62
+ "group_size": 64,
63
+ "bits": 4,
64
+ "mode": "affine"
65
+ },
66
+ "text_config": {
67
+ "attention_bias": false,
68
+ "attention_dropout": 0.0,
69
+ "attention_k_eq_v": false,
70
+ "bos_token_id": 2,
71
+ "dtype": "bfloat16",
72
+ "enable_moe_block": false,
73
+ "eos_token_id": 1,
74
+ "expert_intermediate_size": null,
75
+ "final_logit_softcapping": 30.0,
76
+ "global_head_dim": 512,
77
+ "head_dim": 256,
78
+ "hidden_activation": "gelu_pytorch_tanh",
79
+ "hidden_size": 2560,
80
+ "hidden_size_per_layer_input": 256,
81
+ "initializer_range": 0.02,
82
+ "intermediate_size": 10240,
83
+ "layer_types": [
84
+ "sliding_attention",
85
+ "sliding_attention",
86
+ "sliding_attention",
87
+ "sliding_attention",
88
+ "sliding_attention",
89
+ "full_attention",
90
+ "sliding_attention",
91
+ "sliding_attention",
92
+ "sliding_attention",
93
+ "sliding_attention",
94
+ "sliding_attention",
95
+ "full_attention",
96
+ "sliding_attention",
97
+ "sliding_attention",
98
+ "sliding_attention",
99
+ "sliding_attention",
100
+ "sliding_attention",
101
+ "full_attention",
102
+ "sliding_attention",
103
+ "sliding_attention",
104
+ "sliding_attention",
105
+ "sliding_attention",
106
+ "sliding_attention",
107
+ "full_attention",
108
+ "sliding_attention",
109
+ "sliding_attention",
110
+ "sliding_attention",
111
+ "sliding_attention",
112
+ "sliding_attention",
113
+ "full_attention",
114
+ "sliding_attention",
115
+ "sliding_attention",
116
+ "sliding_attention",
117
+ "sliding_attention",
118
+ "sliding_attention",
119
+ "full_attention",
120
+ "sliding_attention",
121
+ "sliding_attention",
122
+ "sliding_attention",
123
+ "sliding_attention",
124
+ "sliding_attention",
125
+ "full_attention"
126
+ ],
127
+ "max_position_embeddings": 131072,
128
+ "model_type": "gemma4_text",
129
+ "num_attention_heads": 8,
130
+ "num_experts": null,
131
+ "num_global_key_value_heads": null,
132
+ "num_hidden_layers": 42,
133
+ "num_key_value_heads": 2,
134
+ "num_kv_shared_layers": 18,
135
+ "pad_token_id": 0,
136
+ "rms_norm_eps": 1e-06,
137
+ "rope_parameters": {
138
+ "full_attention": {
139
+ "partial_rotary_factor": 0.25,
140
+ "rope_theta": 1000000.0,
141
+ "rope_type": "proportional"
142
+ },
143
+ "sliding_attention": {
144
+ "rope_theta": 10000.0,
145
+ "rope_type": "default"
146
+ }
147
+ },
148
+ "sliding_window": 512,
149
+ "tie_word_embeddings": true,
150
+ "top_k_experts": null,
151
+ "use_bidirectional_attention": null,
152
+ "use_cache": true,
153
+ "use_double_wide_mlp": false,
154
+ "vocab_size": 262144,
155
+ "vocab_size_per_layer_input": 262144
156
+ },
157
+ "tie_word_embeddings": true,
158
+ "transformers_version": "5.5.0.dev0",
159
+ "video_token_id": 258884,
160
+ "vision_config": {
161
+ "_name_or_path": "",
162
+ "architectures": null,
163
+ "attention_bias": false,
164
+ "attention_dropout": 0.0,
165
+ "chunk_size_feed_forward": 0,
166
+ "default_output_length": 280,
167
+ "dtype": "bfloat16",
168
+ "global_head_dim": 64,
169
+ "head_dim": 64,
170
+ "hidden_activation": "gelu_pytorch_tanh",
171
+ "hidden_size": 768,
172
+ "id2label": {
173
+ "0": "LABEL_0",
174
+ "1": "LABEL_1"
175
+ },
176
+ "initializer_range": 0.02,
177
+ "intermediate_size": 3072,
178
+ "is_encoder_decoder": false,
179
+ "label2id": {
180
+ "LABEL_0": 0,
181
+ "LABEL_1": 1
182
+ },
183
+ "max_position_embeddings": 131072,
184
+ "model_type": "gemma4_vision",
185
+ "num_attention_heads": 12,
186
+ "num_hidden_layers": 16,
187
+ "num_key_value_heads": 12,
188
+ "output_attentions": false,
189
+ "output_hidden_states": false,
190
+ "patch_size": 16,
191
+ "pooling_kernel_size": 3,
192
+ "position_embedding_size": 10240,
193
+ "problem_type": null,
194
+ "return_dict": true,
195
+ "rms_norm_eps": 1e-06,
196
+ "rope_parameters": {
197
+ "rope_theta": 100.0,
198
+ "rope_type": "default"
199
+ },
200
+ "standardize": false,
201
+ "use_clipped_linears": true
202
+ },
203
+ "vision_soft_tokens_per_image": 280
204
+ }
generation_config.json ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 2,
3
+ "do_sample": true,
4
+ "eos_token_id": 1,
5
+ "pad_token_id": 0,
6
+ "temperature": 1.0,
7
+ "top_k": 64,
8
+ "top_p": 0.95,
9
+ "transformers_version": "5.5.0.dev0"
10
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:53bf49ec9ab2968a9916123dc60cd6e76ff2b3a06932542f5fd958eb13636504
3
+ size 5217361182
model.safetensors.index.json ADDED
The diff for this file is too large to render. See raw diff
 
processor_config.json ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "audio_seq_length": 750,
3
+ "image_processor": {
4
+ "do_convert_rgb": true,
5
+ "do_normalize": false,
6
+ "do_rescale": true,
7
+ "do_resize": true,
8
+ "image_mean": [
9
+ 0.0,
10
+ 0.0,
11
+ 0.0
12
+ ],
13
+ "image_processor_type": "Gemma4ImageProcessor",
14
+ "image_seq_length": 280,
15
+ "image_std": [
16
+ 1.0,
17
+ 1.0,
18
+ 1.0
19
+ ],
20
+ "max_soft_tokens": 280,
21
+ "patch_size": 16,
22
+ "pooling_kernel_size": 3,
23
+ "resample": 3,
24
+ "rescale_factor": 0.00392156862745098,
25
+ "size": {
26
+ "height": 224,
27
+ "width": 224
28
+ }
29
+ },
30
+ "image_seq_length": 280,
31
+ "processor_class": "Gemma4Processor",
32
+ "feature_extractor": {
33
+ "feature_extractor_type": "Gemma4AudioFeatureExtractor",
34
+ "sampling_rate": 16000,
35
+ "num_mel_filters": 128,
36
+ "fft_length": 512,
37
+ "hop_length": 160,
38
+ "chunk_duration": 8.0,
39
+ "overlap_duration": 1.0
40
+ },
41
+ "audio_ms_per_token": 40
42
+ }
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:12bac982b793c44b03d52a250a9f0d0b666813da566b910c24a6da0695fd11e6
3
+ size 32170070
tokenizer_config.json ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "audio_token": "<|audio|>",
3
+ "backend": "tokenizers",
4
+ "boa_token": "<|audio>",
5
+ "boi_token": "<|image>",
6
+ "bos_token": "<bos>",
7
+ "eoa_token": "<audio|>",
8
+ "eoc_token": "<channel|>",
9
+ "eoi_token": "<image|>",
10
+ "eos_token": "<eos>",
11
+ "eot_token": "<turn|>",
12
+ "escape_token": "<|\"|>",
13
+ "etc_token": "<tool_call|>",
14
+ "etd_token": "<tool|>",
15
+ "etr_token": "<tool_response|>",
16
+ "extra_special_tokens": [
17
+ "<|video|>"
18
+ ],
19
+ "image_token": "<|image|>",
20
+ "is_local": true,
21
+ "mask_token": "<mask>",
22
+ "model_max_length": 1000000000000000019884624838656,
23
+ "model_specific_special_tokens": {
24
+ "audio_token": "<|audio|>",
25
+ "boa_token": "<|audio>",
26
+ "boi_token": "<|image>",
27
+ "eoa_token": "<audio|>",
28
+ "eoc_token": "<channel|>",
29
+ "eoi_token": "<image|>",
30
+ "eot_token": "<turn|>",
31
+ "escape_token": "<|\"|>",
32
+ "etc_token": "<tool_call|>",
33
+ "etd_token": "<tool|>",
34
+ "etr_token": "<tool_response|>",
35
+ "image_token": "<|image|>",
36
+ "soc_token": "<|channel>",
37
+ "sot_token": "<|turn>",
38
+ "stc_token": "<|tool_call>",
39
+ "std_token": "<|tool>",
40
+ "str_token": "<|tool_response>",
41
+ "think_token": "<|think|>"
42
+ },
43
+ "pad_token": "<pad>",
44
+ "padding_side": "left",
45
+ "processor_class": "Gemma4Processor",
46
+ "soc_token": "<|channel>",
47
+ "sot_token": "<|turn>",
48
+ "stc_token": "<|tool_call>",
49
+ "std_token": "<|tool>",
50
+ "str_token": "<|tool_response>",
51
+ "think_token": "<|think|>",
52
+ "tokenizer_class": "GemmaTokenizer",
53
+ "unk_token": "<unk>"
54
+ }