Add packaged audio inference support and refresh deployment assets

.gitattributes

@@ -41,5 +41,7 @@ main_axcl_x86 filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
 *.jpg filter=lfs diff=lfs merge=lfs -text
 *.mp4 filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
 tokenizer.json filter=lfs diff=lfs merge=lfs -text
 gemma4_tokenizer.txt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -6,6 +6,7 @@ base_model:
 tags:
   - gemma4
   - multimodal
+  - audio
   - axera
   - ax650
 language:
@@ -23,8 +24,8 @@ Ready-to-run deployment package for `google/gemma-4-E2B-it` on AX650 / NPU3.
 - This release packages the `w8a16` AXERA NPU runtime.
 - Compatible with `Pulsar2 5.2` and later.
 - Includes the tokenizer/config files required at runtime.
-- Includes compiled Gemma 4 text `.axmodel` files and Vision `.axmodel` files.
-- Supports both text-only chat and single-image multimodal inference.
+- Includes compiled Gemma 4 text `.axmodel` files, Vision `.axmodel` files, and fixed-duration Audio `.axmodel` files.
+- Supports text-only chat, single-image multimodal inference, and fixed-duration audio transcription in the legacy Python demo flow.
 
 ## Supported Platform
 
@@ -42,8 +43,8 @@ This package has been validated on the following AX650-based devices:
 
 All measurements below were taken on AX650 / NPU3. `TTFT` stands for time to first token.
 
-- `w8a16`: TTFT is approximately `1664 ms`, with a decode throughput of approximately `10.44 tokens/s`.
-- `w4a16`: TTFT is approximately `1233.7 ms`, with a decode throughput of approximately `15.22 tokens/s`.
+- `w8a16`: TTFT is approximately `2175 ms` (1152 tokens), with a decode throughput of approximately `7.99 tokens/s` (theoretical maximum).
+- `w4a16`: TTFT is approximately `1568 ms` (1152 tokens), with a decode throughput of approximately `12.41 tokens/s` (theoretical maximum).
 
 The packaged text runtime in this release is the `w8a16` build. Its text runtime files are packaged at the repository root. The `w4a16` numbers are provided for reference only.
 
@@ -55,6 +56,22 @@ The packaged text runtime in this release is the `w8a16` build. Its text runtime
 | `gemma4_vision_h480_w672_t140.axmodel` | 480×672 | 140 | 258.329 ms |
 | `gemma4_vision_h672_w960_t280.axmodel` | 672×960 | 280 | 750.429 ms |
 
+## Audio Encoder Latency and Accuracy
+
+Validated on AX650 / NPU3 with the packaged WAV sample clips and the legacy Python demo flow.
+
+| Model | Audio Duration | Audio Tokens |
+|---|---|---|
+| `gemma4_audio_5s.axmodel` | `5s` | `125` |
+| `gemma4_audio_30s.axmodel` | `30s` | `750` |
+
+Single-run latency measured with `/opt/bin/ax_run_model -w 1 -r 5`:
+
+| Model | CMM Size | Avg |
+|---|---|---|
+| `gemma4_audio_5s.axmodel` | `~310 MiB` | `28.930 ms` |
+| `gemma4_audio_30s.axmodel` | `~359 MiB` | `170.978 ms` |
+
 ## Package Layout
 
 ```text
@@ -69,6 +86,8 @@ The packaged text runtime in this release is the `w8a16` build. Its text runtime
 ├── gemma4_tokenizer.txt
 ├── gemma4_text_p128_l*.axmodel
 ├── gemma4_text_post.axmodel
+├── gemma4_audio_5s.axmodel
+├── gemma4_audio_30s.axmodel
 ├── gemma4_vision_h336_w480_t70.axmodel
 ├── gemma4_vision_h480_w672_t140.axmodel
 ├── gemma4_vision_h672_w960_t280.axmodel
@@ -80,7 +99,7 @@ The packaged text runtime in this release is the `w8a16` build. Its text runtime
 └── utils/
 ```
 
-This package uses a hybrid layout: the tokenizer stays in a subdirectory, the packaged text runtime files and Vision `.axmodel` files live at the repository root, and `vit_models/` keeps the accompanying Vision metadata JSON files.
+This package uses a hybrid layout: the tokenizer stays in a subdirectory, the packaged text runtime files plus the Vision and Audio `.axmodel` files live at the repository root, and `vit_models/` keeps the accompanying Vision metadata JSON files.
 
 The Python demo scripts auto-detect the packaged paths above. If you keep this layout unchanged, you can run the Python examples later in this README without passing extra path arguments.
 
@@ -91,6 +110,14 @@ Both the `axllm` flow and the legacy Python demo flow below can use the packaged
 
 ![sample](assets/sample.png)
 
+## Sample Audio
+
+The package also includes three packaged WAV clips for board-side audio validation:
+
+- `assets/gemma4_audio_test_5s.wav`
+- `assets/gemma4_audio_test_chunk0_30s.wav`
+- `assets/gemma4_audio_test_chunk1_30s.wav`
+
 ## Direct Inference with `axllm`
 
 > The `axllm` workflow is still being refined. The instructions below reflect the current validated flow and may be adjusted as the packaging continues to evolve.
@@ -144,91 +171,103 @@ For multimodal testing, you can use the sample image shown above: `./assets/samp
 $ axllm run AXERA-TECH/gemma-4-E2B-it
 
 # output log example:
-15:04:24.522 INF Init:890 | LLM init start
-15:04:24.522 INF Init:905 | shared kv enabled: num_kv_shared_layers=20
+14:57:27.565 INF Init:1019 | LLM init start
+14:57:27.565 INF Init:1034 | shared kv enabled: num_kv_shared_layers=20
+14:57:27.565 INF Init:1050 | attention config: layers=35 sliding=28 full=7 linear=0 sliding_window=512 ref_full_layer_idx=0
 tokenizer_type = 3
 huggingface tokenizer mode = space_replace_bpe
- 31% | ##########                       |  12 /  38 [4.47s<14.16s, 2.68 count/s] init 10 axmodel ok,remain_cmm(6047 MB                                                                                                                       34% | ##########                       |  13 /  38 [4.61s<13.48s, 2.82 count/s] init 11 axmodel ok,remain_cmm(5992 MB                                                                                                                       36% | ###########                      |  14 /  38 [4.78s<12.98s, 2.93 count/s] init 12 axmodel ok,remain_cmm(5937 MB                                                                                                                       39% | ############                     |  15 /  38 [4.93s<12.49s, 3.04 count/s] init 13 axmodel ok,remain_cmm(5882 MB                                                                                                                       42% | #############                    |  16 /  38 [5.09s<12.09s, 3.14 count/s] init 14 axmodel ok,remain_cmm(5813 MB                                                                                                                       44% | ##############                   |  17 /  38 [5.28s<11.80s, 3.22 count/s] init 15 axmodel ok,remain_cmm(5727 MB                                                                                                                       47% | ###############                  |  18 /  38 [5.50s<11.61s, 3.27 count/s] init 16 axmodel ok,remain_cmm(5642 MB                                                                                                                       50% | ################                 |  19 /  38 [5.69s<11.38s, 3.34 count/s] init 17 axmodel ok,remain_cmm(5557 MB                                                                                                                       52% | ################                 |  20 /  38 [5.91s<11.22s, 3.39 count/s] init 18 axmodel ok,remain_cmm(5471 MB                                                                                                                       55% | #################                |  21 /  38 [6.11s<11.06s, 3.44 count/s] init 19 axmodel ok,remain_cmm(5373 MB                                                                                                                       57% | ##################               |  22 /  38 [6.31s<10.89s, 3.49 count/s] init 20 axmodel ok,remain_cmm(5287 MB                                                                                                                       60% | ###################              |  23 /  38 [6.53s<10.79s, 3.52 count/s] init 21 axmodel ok,remain_cmm(5202 MB                                                                                                                       63% | ####################             |  24 /  38 [6.75s<10.69s, 3.56 count/s] init 22 axmodel ok,remain_cmm(5117 MB                                                                                                                       65% | #####################            |  25 /  38 [6.96s<10.58s, 3.59 count/s] init 23 axmodel ok,remain_cmm(5031 MB                                                                                                                       68% | #####################            |  26 /  38 [7.18s<10.50s, 3.62 count/s] init 24 axmodel ok,remain_cmm(4933 MB                                                                                                                       71% | ######################           |  27 /  38 [7.40s<10.41s, 3.65 count/s] init 25 axmodel ok,remain_cmm(4847 MB                                                                                                                       73% | #######################          |  28 /  38 [7.62s<10.34s, 3.67 count/s] init 26 axmodel ok,remain_cmm(4762 MB                                                                                                                       76% | ########################         |  29 /  38 [7.85s<10.28s, 3.70 count/s] init 27 axmodel ok,remain_cmm(4676 MB                                                                                                                       78% | #########################        |  30 /  38 [8.13s<10.29s, 3.69 count/s] init 28 axmodel ok,remain_cmm(4591 MB                                                                                                                       81% | ##########################       |  31 /  38 [8.36s<10.25s, 3.71 count/s] init 29 axmodel ok,remain_cmm(4492 MB                                                                                                                       84% | ##########################       |  32 /  38 [8.60s<10.21s, 3.72 count/s] init 30 axmodel ok,remain_cmm(4407 MB                                                                                                                       86% | ###########################      |  33 /  38 [8.86s<10.21s, 3.72 count/s] init 31 axmodel ok,remain_cmm(4322 MB                                                                                                                       89% | ############################     |  34 /  38 [9.11s<10.18s, 3.73 count/s] init 32 axmodel ok,remain_cmm(4236 MB                                                                                                                       92% | #############################    |  35 /  38 [9.36s<10.16s, 3.74 count/s] init 33 axmodel ok,remain_cmm(4151 MB                                                                                                                       94% | ##############################   |  36 /  38 [9.62s<10.16s, 3.74 count/s] init 34 axmodel ok,remain_cmm(4052 MB                                                                                                                       97% | ###############################  |  37 /  38 [10.03s<10.30s, 3.69 count/s] init post axmodel ok,remain_cmm(3632 MB)
-15:04:34.551 INF Init:1045 | max_token_len : 2047
-15:04:34.551 INF Init:1048 | kv_cache_size : 256, kv_cache_num: 2047
-15:04:34.551 INF init_groups_from_model:606 | prefill_token_num : 128
-15:04:34.551 INF init_groups_from_model:820 | decode grp: 0, gid: 0, max_token_len : 2047
-15:04:34.551 INF init_groups_from_model:824 | prefill grp: 0, gid: 1, history_cap: 0, total_cap: 128, symbolic_cap: 1
-15:04:34.551 INF init_groups_from_model:824 | prefill grp: 1, gid: 2, history_cap: 128, total_cap: 256, symbolic_cap: 128
-15:04:34.551 INF init_groups_from_model:824 | prefill grp: 2, gid: 3, history_cap: 256, total_cap: 384, symbolic_cap: 256
-15:04:34.551 INF init_groups_from_model:824 | prefill grp: 3, gid: 4, history_cap: 384, total_cap: 512, symbolic_cap: 384
-15:04:34.551 INF init_groups_from_model:824 | prefill grp: 4, gid: 5, history_cap: 512, total_cap: 640, symbolic_cap: 512
-15:04:34.551 INF init_groups_from_model:824 | prefill grp: 5, gid: 6, history_cap: 640, total_cap: 768, symbolic_cap: 640
-15:04:34.551 INF init_groups_from_model:824 | prefill grp: 6, gid: 7, history_cap: 768, total_cap: 896, symbolic_cap: 768
-15:04:34.551 INF init_groups_from_model:824 | prefill grp: 7, gid: 8, history_cap: 896, total_cap: 1024, symbolic_cap: 896
-15:04:34.551 INF init_groups_from_model:824 | prefill grp: 8, gid: 9, history_cap: 1024, total_cap: 1152, symbolic_cap: 1024
-15:04:34.551 INF init_groups_from_model:831 | prefill_max_token_num : 1152
-15:04:34.551 INF Init:27 | LLaMaEmbedSelector use mmap
-100% | ################################ |  38 /  38 [10.03s<10.03s, 3.79 count/s] embed_selector init ok
-15:04:34.567 INF Init:475 | Gemma4 per-layer helper enabled: vocab=262144 hidden=1536 layers=35 per_layer=256 pad=0
-15:04:34.727 INF Init:785 | Gemma4-VL token ids: image_pad=258880 video_pad=258884
-15:04:34.727 INF Init:792 | VisionModule init ok: type=Gemma4VL, tokens_per_block=70, embed_size=1536, out_dtype=fp32
-15:04:34.727 WRN Init:801 | Vision preprocess backend: SimpleCV (OpenCV not found at build time; minor differences vs OpenCV are possible)
-15:04:34.729 INF load_config:282 | load config:
-15:04:34.729 INF load_config:282 | {
-15:04:34.729 INF load_config:282 |     "enable_repetition_penalty": false,
-15:04:34.729 INF load_config:282 |     "enable_temperature": true,
-15:04:34.729 INF load_config:282 |     "enable_top_k_sampling": false,
-15:04:34.729 INF load_config:282 |     "enable_top_p_sampling": true,
-15:04:34.729 INF load_config:282 |     "penalty_window": 64,
-15:04:34.729 INF load_config:282 |     "repetition_penalty": 1.0,
-15:04:34.729 INF load_config:282 |     "temperature": 1.0,
-15:04:34.729 INF load_config:282 |     "top_k": 64,
-15:04:34.729 INF load_config:282 |     "top_p": 0.95
-15:04:34.729 INF load_config:282 | }
-15:04:34.729 INF Init:1139 | LLM init ok
+ 31% | ##########                       |  12 /  38 [3.41s<10.80s, 3.52 count/s] init 10 axmodel ok,remain_cmm(7282 MB                                                                                                                       34% | ##########                       |  13 /  38 [3.48s<10.18s, 3.73 count/s] init 11 axmodel ok,remain_cmm(7227 MB                                                                                                                       97% | ###############################  |  37 /  38 [6.57s<6.74s, 5.64 count/s] init post axmodel ok,remain_cmm(4868 MB)
+14:57:34.130 INF Init:1196 | max_token_len : 2047
+14:57:34.131 INF Init:1199 | kv_cache_size : 256, kv_cache_num: 2047
+14:57:34.131 INF init_groups_from_model:702 | prefill_token_num : 128
+14:57:34.131 INF init_groups_from_model:916 | decode grp: 0, gid: 0, max_token_len : 2047
+14:57:34.131 INF init_groups_from_model:920 | prefill grp: 0, gid: 1, history_cap: 0, total_cap: 128, symbolic_cap: 1
+14:57:34.131 INF init_groups_from_model:920 | prefill grp: 1, gid: 2, history_cap: 128, total_cap: 256, symbolic_cap: 128
+14:57:34.131 INF init_groups_from_model:920 | prefill grp: 2, gid: 3, history_cap: 256, total_cap: 384, symbolic_cap: 256
+14:57:34.131 INF init_groups_from_model:920 | prefill grp: 3, gid: 4, history_cap: 384, total_cap: 512, symbolic_cap: 384
+14:57:34.131 INF init_groups_from_model:920 | prefill grp: 4, gid: 5, history_cap: 512, total_cap: 640, symbolic_cap: 512
+14:57:34.131 INF init_groups_from_model:920 | prefill grp: 5, gid: 6, history_cap: 640, total_cap: 768, symbolic_cap: 640
+14:57:34.131 INF init_groups_from_model:920 | prefill grp: 6, gid: 7, history_cap: 768, total_cap: 896, symbolic_cap: 768
+14:57:34.131 INF init_groups_from_model:920 | prefill grp: 7, gid: 8, history_cap: 896, total_cap: 1024, symbolic_cap: 896
+14:57:34.131 INF init_groups_from_model:920 | prefill grp: 8, gid: 9, history_cap: 1024, total_cap: 1152, symbolic_cap: 1024
+14:57:34.131 INF init_groups_from_model:927 | prefill_max_token_num : 1152
+14:57:34.131 INF Init:27 | LLaMaEmbedSelector use mmap
+100% | ################################ |  38 /  38 [6.57s<6.57s, 5.78 count/s] embed_selector init ok
+14:57:34.149 INF Init:475 | Gemma4 per-layer helper enabled: vocab=262144 hidden=1536 layers=35 per_layer=256 pad=0
+14:57:39.791 INF init_audio_profile:245 | Gemma4 audio profile init ok: path=../../gemma-4-E2B-it/gemma4_audio_5s.axmodel duration=5.0s mel_frames=499 tokens=125 out_dtype=fp32
+14:57:40.049 INF init_audio_profile:245 | Gemma4 audio profile init ok: path=../../gemma-4-E2B-it/gemma4_audio_30s.axmodel duration=30.0s mel_frames=2999 tokens=750 out_dtype=fp32
+14:57:40.049 INF Init:914 | Gemma4-VL token ids: image_pad=258880 video_pad=258884 audio_pad=258881
+14:57:40.049 INF Init:921 | VisionModule init ok: type=Gemma4VL, tokens_per_block=70, embed_size=1536, out_dtype=fp32
+14:57:40.049 WRN Init:930 | Vision preprocess backend: SimpleCV (OpenCV not found at build time; minor differences vs OpenCV are possible)
+14:57:40.054 INF load_config:444 | load config:
+14:57:40.054 INF load_config:444 | {
+14:57:40.054 INF load_config:444 |     "enable_repetition_penalty": false,
+14:57:40.054 INF load_config:444 |     "enable_temperature": false,
+14:57:40.054 INF load_config:444 |     "enable_top_k_sampling": false,
+14:57:40.054 INF load_config:444 |     "enable_top_p_sampling": false,
+14:57:40.054 INF load_config:444 |     "penalty_window": 64,
+14:57:40.054 INF load_config:444 |     "repetition_penalty": 1.0,
+14:57:40.054 INF load_config:444 |     "temperature": 1.0,
+14:57:40.054 INF load_config:444 |     "top_k": 64,
+14:57:40.054 INF load_config:444 |     "top_p": 0.95
+14:57:40.054 INF load_config:444 | }
+14:57:40.055 INF Init:1293 | LLM init ok
 Commands:
   /q, /exit  退出
   /reset     重置 kvcache
   /dd        删除一轮对话
   /pp        打印历史对话
 Ctrl+C: 停止当前生成
-VLM enabled: after each prompt, input media path (empty = text-only). Use "video:<frames_dir>" for video, "audio:<file>" for reserved audio placeholder.
+VLM enabled: after each prompt, input media path (empty = text-only). Use "video:<frames_dir>" for video, "audio:<file>" for audio.
 ----------------------------------------
-prompt >> who are you?
+prompt >> hello, who are you?
+media >>
+14:58:01.823 INF SetKVCache:1607 | decode_grpid:0 prefill_grpid:1 history_cap:0 total_cap:128 symbolic_cap:1 precompute_len:0 input_num_token:26 prefer_symbolic_group:0
+14:58:01.823 INF SetKVCache:1628 | current prefill_max_token_num:1152
+14:58:01.863 INF SetKVCache:1632 | first run
+14:58:01.884 INF Run:1736 | input token num : 26, prefill_split_num : 1
+14:58:02.014 INF Run:1819 | prefill chunk p=0 history_len=0 grpid=1 kv_cache_num=0 input_tokens=26
+14:58:02.014 INF Run:1843 | prefill indices shape: p=0 idx_elems=128 idx_rows=1 pos_rows=0
+14:58:02.267 INF Run:2028 | ttft: 383.05 ms
+Hello! I am Gemma 4, a Large Language Model developed by Google DeepMind. I am an open weights model. How can I help you today?
+
+14:58:07.195 NTC Run:2396 | hit eos,decode avg 6.29 token/s
+14:58:07.198 INF GetKVCache:1572 | precompute_len:58, remaining:1094 (tracked)
+prompt >> Okay, great!
 media >>
-15:04:39.368 INF SetKVCache:1437 | decode_grpid:0 prefill_grpid:1 history_cap:0 total_cap:128 symbolic_cap:1 precompute_len:0 input_num_token:24 prefer_symbolic_group:0
-15:04:39.368 INF SetKVCache:1458 | current prefill_max_token_num:1152
-15:04:39.408 INF SetKVCache:1462 | first run
-15:04:39.409 INF Run:1553 | input token num : 24, prefill_split_num : 1
-15:04:39.482 INF Run:1640 | prefill chunk p=0 history_len=0 grpid=1 kv_cache_num=0 input_tokens=24
-15:04:39.483 INF Run:1665 | prefill indices shape: p=0 idx_elems=128 idx_rows=1 pos_rows=0
-15:04:39.764 INF Run:1837 | ttft: 355.37 ms
-I am Gemma 4, a Large Language Model developed by Google DeepMind. I am an open weights model.
-
-15:04:44.087 NTC Run:2103 | hit eos,decode avg 5.09 token/s
-15:04:44.088 INF GetKVCache:1408 | precompute_len:47, remaining:1105
-prompt >> Please describe the image in detail.
-media >> /root/yongqiang/auto_model_deployment/gemma-4-E2B-it/assets/sample.png
-15:06:14.416 INF EncodeForContent:1122 | vision cache hit (disk): /root/yongqiang/auto_model_deployment/gemma-4-E2B-it/assets/sample.png
-15:06:14.416 INF EncodeForContent:1131 | vision cache hit (mem): /root/yongqiang/auto_model_deployment/gemma-4-E2B-it/assets/sample.png
-15:06:14.419 INF SetKVCache:1437 | decode_grpid:0 prefill_grpid:3 history_cap:256 total_cap:384 symbolic_cap:256 precompute_len:47 input_num_token:94 prefer_symbolic_group:1
-15:06:14.419 INF SetKVCache:1458 | current prefill_max_token_num:1024
-15:06:14.429 INF Run:1553 | input token num : 94, prefill_split_num : 1
-15:06:14.703 INF Run:1640 | prefill chunk p=0 history_len=47 grpid=3 kv_cache_num=256 input_tokens=94
-15:06:14.703 INF Run:1665 | prefill indices shape: p=0 idx_elems=128 idx_rows=1 pos_rows=0
-15:06:15.027 INF Run:1837 | ttft: 597.95 ms
-I see an image of a cartoon character that resembles a cooked or stylized lobster.
-
-Here is a detailed description of the image:
-
-*   **Subject:** The central subject is a bright red, stylized lobster.
-*   **Style:** The illustration is highly cartoonish and vibrant, featuring thick outlines and bright, saturated colors, suggesting a playful or energetic style.
-*   **Features:**
-    *   The lobster has large, expressive eyes and a wide, toothy grin, giving it a mischievous or energetic expression.
-    *   Its claws (pincers) are prominent and stylized.
-    *   The body is segmented, typical of a lobster, but rendered in a simplified, bold manner.
-    *   It has a curved, slightly exaggerated posture.
-*   **Outline/Background:** The character is set against a plain white background. The image has a glossy or sticker-like finish, indicated by a slight shadow effect or outline around the character, suggesting it might be a graphic or sticker design.
-*   **Overall Impression:** The image is energetic, bold, and fun, clearly designed as a mascot or a character illustration.
-
-15:07:09.593 NTC Run:2103 | hit eos,decode avg 4.33 token/s
-15:07:09.593 INF GetKVCache:1408 | precompute_len:378, remaining:774
+14:58:20.884 INF SetKVCache:1607 | decode_grpid:0 prefill_grpid:2 history_cap:128 total_cap:256 symbolic_cap:128 precompute_len:58 input_num_token:14 prefer_symbolic_group:0
+14:58:20.885 INF SetKVCache:1628 | current prefill_max_token_num:1024
+14:58:20.892 INF Run:1736 | input token num : 14, prefill_split_num : 1
+14:58:20.937 INF Run:1819 | prefill chunk p=0 history_len=58 grpid=2 kv_cache_num=128 input_tokens=14
+14:58:20.938 INF Run:1843 | prefill indices shape: p=0 idx_elems=128 idx_rows=1 pos_rows=0
+14:58:21.210 INF Run:2028 | ttft: 318.07 ms
+I'm happy to help! What can I do for you? Do you have a question, need some information, want to brainstorm some ideas, or anything else? 😊
+
+14:58:26.377 NTC Run:2396 | hit eos,decode avg 6.58 token/s
+14:58:26.383 INF GetKVCache:1572 | precompute_len:107, remaining:1045 (tracked)
+prompt >> Please describe this image in detail.
+media >> /root/your/workspace/gemma-4-E2B-it/assets/sample.png
+15:01:41.584 INF EncodeForContent:1464 | vision cache store: /root/your/workspace/gemma-4-E2B-it/assets/sample.png
+15:01:41.610 INF SetKVCache:1607 | decode_grpid:0 prefill_grpid:3 history_cap:256 total_cap:384 symbolic_cap:256 precompute_len:107 input_num_token:91 prefer_symbolic_group:1
+15:01:41.610 INF SetKVCache:1628 | current prefill_max_token_num:1024
+15:01:41.630 INF Run:1736 | input token num : 91, prefill_split_num : 1
+15:01:41.922 INF Run:1819 | prefill chunk p=0 history_len=107 grpid=3 kv_cache_num=256 input_tokens=91
+15:01:41.923 INF Run:1843 | prefill indices shape: p=0 idx_elems=128 idx_rows=1 pos_rows=0
+15:01:42.232 INF Run:2028 | ttft: 601.69 ms
+This image is a cartoon illustration of a **red, stylized lobster**.
+
+Here is a detailed description:
+
+*   **Subject:** The central subject is a lobster, depicted in a vibrant, bright red color.
+*   **Style:** The illustration is highly stylized and cartoonish, featuring thick outlines and exaggerated features, giving it a playful or energetic look.
+*   **Pose/Expression:** The lobster appears to be in an active or aggressive pose. It has prominent claws (chelipeds) that are raised, suggesting it might be waving, striking, or ready for action. Its body is curved dynamically.
+*   **Details:** You can clearly see the segmented body, the claws, and the legs. The overall design is bold, with clean lines and a glossy or slightly textured appearance typical of a sticker or graphic design.
+*   **Background:** The lobster is isolated on a plain white background, which makes the red color and the details of the illustration stand out prominently.
+*   **Overall Impression:** The image is energetic, bold, and clearly designed to be eye-catching, likely intended for use as a sticker, icon, or graphic element.
+
+15:02:16.834 NTC Run:2396 | hit eos,decode avg 6.85 token/s
+15:02:17.013 INF GetKVCache:1572 | precompute_len:436, remaining:716 (tracked)
+prompt >>
 ```
 
 
@@ -241,44 +280,47 @@ Note: the command below assumes you run it from the parent directory of `AXERA-T
 ```bash
 $ axllm serve AXERA-TECH/gemma-4-E2B-it --port 8000
 # output log example:
-16:22:21.336 INF Init:890 | LLM init start
-16:22:21.336 INF Init:905 | shared kv enabled: num_kv_shared_layers=20
+15:05:32.638 INF Init:1019 | LLM init start
+15:05:32.638 INF Init:1034 | shared kv enabled: num_kv_shared_layers=20
+15:05:32.638 INF Init:1050 | attention config: layers=35 sliding=28 full=7 linear=0 sliding_window=512 ref_full_layer_idx=0
 tokenizer_type = 3
 huggingface tokenizer mode = space_replace_bpe
- 13% | ####                             |   5 /  38 [10.08s<76.64s, 0.50 count/s] init 3 axmodel ok,remain_cmm(4704 MB                                                                                                                       15% | #####                            |   6 /  38 [12.32s<78.01s, 0.49 count/s] init 4 axmodel ok,remain_cmm(4635 MB                                                                                                                       18% | #####                            |   7 /  38 [17.98s<97.59s, 0.39 count/s] init 5 axmodel ok,remain_cmm(4580 MB                                                                                                                       21% | ######                           |   8 /  38 [18.54s<88.05s, 0.43 count/s] init 6 axmodel ok,remain_cmm(4525 MB                                                                                                                       23% | #######                          |   9 /  38 [19.06s<80.49s, 0.47 count/s] init 7 axmodel ok,remain_cmm(4470 MB                                                                                                                       26% | ########                         |  10 /  38 [19.66s<74.70s, 0.51 count/s] init 8 axmodel ok,remain_cmm(4415 MB                                                                                                                       28% | #########                        |  11 /  38 [20.41s<70.49s, 0.54 count/s] init 9 axmodel ok,remain_cmm(4346 MB                                                                                                                       31% | ##########                       |  12 /  38 [20.81s<65.91s, 0.58 count/s] init 10 axmodel ok,remain_cmm(4291 M                                                                                                                       34% | ##########                       |  13 /  38 [21.22s<62.04s, 0.61 count/s] init 11 axmodel ok,remain_cmm(4236 M                                                                                                                       36% | ###########                      |  14 /  38 [21.89s<59.42s, 0.64 count/s] init 12 axmodel ok,remain_cmm(4182 M                                                                                                                       39% | ############                     |  15 /  38 [22.24s<56.34s, 0.67 count/s] init 13 axmodel ok,remain_cmm(4127 M                                                                                                                       42% | #############                    |  16 /  38 [22.56s<53.57s, 0.71 count/s] init 14 axmodel ok,remain_cmm(4057 M                                                                                                                       44% | ##############                   |  17 /  38 [23.11s<51.66s, 0.74 count/s] init 15 axmodel ok,remain_cmm(3972 M                                                                                                                       47% | ###############                  |  18 /  38 [23.54s<49.69s, 0.76 count/s] init 16 axmodel ok,remain_cmm(3887 M                                                                                                                       50% | ################                 |  19 /  38 [24.28s<48.56s, 0.78 count/s] init 17 axmodel ok,remain_cmm(3801 M                                                                                                                       52% | ################                 |  20 /  38 [24.63s<46.80s, 0.81 count/s] init 18 axmodel ok,remain_cmm(3716 M                                                                                                                       55% | #################                |  21 /  38 [24.92s<45.08s, 0.84 count/s] init 19 axmodel ok,remain_cmm(3617 M                                                                                                                       57% | ##################               |  22 /  38 [25.67s<44.33s, 0.86 count/s] init 20 axmodel ok,remain_cmm(3532 M                                                                                                                       60% | ###################              |  23 /  38 [26.33s<43.50s, 0.87 count/s] init 21 axmodel ok,remain_cmm(3447 M                                                                                                                       63% | ####################             |  24 /  38 [27.17s<43.02s, 0.88 count/s] init 22 axmodel ok,remain_cmm(3361 M                                                                                                                       65% | #####################            |  25 /  38 [28.33s<43.06s, 0.88 count/s] init 23 axmodel ok,remain_cmm(3276 M                                                                                                                       68% | #####################            |  26 /  38 [29.70s<43.41s, 0.88 count/s] init 24 axmodel ok,remain_cmm(3177 M                                                                                                                       71% | ######################           |  27 /  38 [30.89s<43.48s, 0.87 count/s] init 25 axmodel ok,remain_cmm(3092 M                                                                                                                       73% | #######################          |  28 /  38 [32.16s<43.65s, 0.87 count/s] init 26 axmodel ok,remain_cmm(3006 M                                                                                                                       76% | ########################         |  29 /  38 [33.32s<43.67s, 0.87 count/s] init 27 axmodel ok,remain_cmm(2921 M                                                                                                                       78% | #########################        |  30 /  38 [34.43s<43.61s, 0.87 count/s] init 28 axmodel ok,remain_cmm(2836 M                                                                                                                       81% | ##########################       |  31 /  38 [35.69s<43.75s, 0.87 count/s] init 29 axmodel ok,remain_cmm(2737 M                                                                                                                       84% | ##########################       |  32 /  38 [36.84s<43.75s, 0.87 count/s] init 30 axmodel ok,remain_cmm(2652 M                                                                                                                       86% | ###########################      |  33 /  38 [37.75s<43.47s, 0.87 count/s] init 31 axmodel ok,remain_cmm(2566 M                                                                                                                       89% | ############################     |  34 /  38 [38.44s<42.96s, 0.88 count/s] init 32 axmodel ok,remain_cmm(2481 M                                                                                                                       92% | #############################    |  35 /  38 [39.06s<42.41s, 0.90 count/s] init 33 axmodel ok,remain_cmm(2396 M                                                                                                                       94% | ##############################   |  36 /  38 [39.44s<41.63s, 0.91 count/s] init 34 axmodel ok,remain_cmm(2297 M                                                                                                                       97% | ###############################  |  37 /  38 [41.12s<42.23s, 0.90 count/s] init post axmodel ok,remain_cmm(1877 MB)
-16:23:02.455 INF Init:1045 | max_token_len : 2047
-16:23:02.455 INF Init:1048 | kv_cache_size : 256, kv_cache_num: 2047
-16:23:02.455 INF init_groups_from_model:606 | prefill_token_num : 128
-16:23:02.455 INF init_groups_from_model:820 | decode grp: 0, gid: 0, max_token_len : 2047
-16:23:02.455 INF init_groups_from_model:824 | prefill grp: 0, gid: 1, history_cap: 0, total_cap: 128, symbolic_cap: 1
-16:23:02.455 INF init_groups_from_model:824 | prefill grp: 1, gid: 2, history_cap: 128, total_cap: 256, symbolic_cap: 128
-16:23:02.455 INF init_groups_from_model:824 | prefill grp: 2, gid: 3, history_cap: 256, total_cap: 384, symbolic_cap: 256
-16:23:02.455 INF init_groups_from_model:824 | prefill grp: 3, gid: 4, history_cap: 384, total_cap: 512, symbolic_cap: 384
-16:23:02.455 INF init_groups_from_model:824 | prefill grp: 4, gid: 5, history_cap: 512, total_cap: 640, symbolic_cap: 512
-16:23:02.455 INF init_groups_from_model:824 | prefill grp: 5, gid: 6, history_cap: 640, total_cap: 768, symbolic_cap: 640
-16:23:02.455 INF init_groups_from_model:824 | prefill grp: 6, gid: 7, history_cap: 768, total_cap: 896, symbolic_cap: 768
-16:23:02.455 INF init_groups_from_model:824 | prefill grp: 7, gid: 8, history_cap: 896, total_cap: 1024, symbolic_cap: 896
-16:23:02.455 INF init_groups_from_model:824 | prefill grp: 8, gid: 9, history_cap: 1024, total_cap: 1152, symbolic_cap: 1024
-16:23:02.455 INF init_groups_from_model:831 | prefill_max_token_num : 1152
-16:23:02.455 INF Init:27 | LLaMaEmbedSelector use mmap
-100% | ################################ |  38 /  38 [41.12s<41.12s, 0.92 count/s] embed_selector init ok
-16:23:02.472 INF Init:475 | Gemma4 per-layer helper enabled: vocab=262144 hidden=1536 layers=35 per_layer=256 pad=0
-16:23:03.400 INF Init:785 | Gemma4-VL token ids: image_pad=258880 video_pad=258884
-16:23:03.400 INF Init:792 | VisionModule init ok: type=Gemma4VL, tokens_per_block=70, embed_size=1536, out_dtype=fp32
-16:23:03.400 WRN Init:801 | Vision preprocess backend: SimpleCV (OpenCV not found at build time; minor differences vs OpenCV are possible)
-16:23:03.404 INF load_config:282 | load config:
-16:23:03.404 INF load_config:282 | {
-16:23:03.404 INF load_config:282 |     "enable_repetition_penalty": false,
-16:23:03.404 INF load_config:282 |     "enable_temperature": true,
-16:23:03.404 INF load_config:282 |     "enable_top_k_sampling": false,
-16:23:03.404 INF load_config:282 |     "enable_top_p_sampling": true,
-16:23:03.404 INF load_config:282 |     "penalty_window": 64,
-16:23:03.404 INF load_config:282 |     "repetition_penalty": 1.0,
-16:23:03.404 INF load_config:282 |     "temperature": 1.0,
-16:23:03.404 INF load_config:282 |     "top_k": 64,
-16:23:03.404 INF load_config:282 |     "top_p": 0.95
-16:23:03.404 INF load_config:282 | }
-16:23:03.404 INF Init:1139 | LLM init ok
+ 31% | ##########                       |  12 /  38 [3.32s<10.50s, 3.62 count/s] init 10 axmodel ok,remain_cmm(7282 MB                                                                                                                       97% | ###############################  |  37 /  38 [6.50s<6.67s, 5.69 count/s] init post axmodel ok,remain_cmm(4868 MB)
+15:05:39.135 INF Init:1196 | max_token_len : 2047
+15:05:39.135 INF Init:1199 | kv_cache_size : 256, kv_cache_num: 2047
+15:05:39.135 INF init_groups_from_model:702 | prefill_token_num : 128
+15:05:39.135 INF init_groups_from_model:916 | decode grp: 0, gid: 0, max_token_len : 2047
+15:05:39.135 INF init_groups_from_model:920 | prefill grp: 0, gid: 1, history_cap: 0, total_cap: 128, symbolic_cap: 1
+15:05:39.135 INF init_groups_from_model:920 | prefill grp: 1, gid: 2, history_cap: 128, total_cap: 256, symbolic_cap: 128
+15:05:39.135 INF init_groups_from_model:920 | prefill grp: 2, gid: 3, history_cap: 256, total_cap: 384, symbolic_cap: 256
+15:05:39.135 INF init_groups_from_model:920 | prefill grp: 3, gid: 4, history_cap: 384, total_cap: 512, symbolic_cap: 384
+15:05:39.135 INF init_groups_from_model:920 | prefill grp: 4, gid: 5, history_cap: 512, total_cap: 640, symbolic_cap: 512
+15:05:39.135 INF init_groups_from_model:920 | prefill grp: 5, gid: 6, history_cap: 640, total_cap: 768, symbolic_cap: 640
+15:05:39.135 INF init_groups_from_model:920 | prefill grp: 6, gid: 7, history_cap: 768, total_cap: 896, symbolic_cap: 768
+15:05:39.135 INF init_groups_from_model:920 | prefill grp: 7, gid: 8, history_cap: 896, total_cap: 1024, symbolic_cap: 896
+15:05:39.135 INF init_groups_from_model:920 | prefill grp: 8, gid: 9, history_cap: 1024, total_cap: 1152, symbolic_cap: 1024
+15:05:39.135 INF init_groups_from_model:927 | prefill_max_token_num : 1152
+15:05:39.135 INF Init:27 | LLaMaEmbedSelector use mmap
+100% | ################################ |  38 /  38 [6.50s<6.50s, 5.85 count/s] embed_selector init ok
+15:05:39.151 INF Init:475 | Gemma4 per-layer helper enabled: vocab=262144 hidden=1536 layers=35 per_layer=256 pad=0
+15:05:39.595 INF init_audio_profile:245 | Gemma4 audio profile init ok: path=../../gemma-4-E2B-it/gemma4_audio_5s.axmodel duration=5.0s mel_frames=499 tokens=125 out_dtype=fp32
+15:05:39.850 INF init_audio_profile:245 | Gemma4 audio profile init ok: path=../../gemma-4-E2B-it/gemma4_audio_30s.axmodel duration=30.0s mel_frames=2999 tokens=750 out_dtype=fp32
+15:05:39.850 INF Init:914 | Gemma4-VL token ids: image_pad=258880 video_pad=258884 audio_pad=258881
+15:05:39.850 INF Init:921 | VisionModule init ok: type=Gemma4VL, tokens_per_block=70, embed_size=1536, out_dtype=fp32
+15:05:39.850 WRN Init:930 | Vision preprocess backend: SimpleCV (OpenCV not found at build time; minor differences vs OpenCV are possible)
+15:05:39.852 INF load_config:444 | load config:
+15:05:39.852 INF load_config:444 | {
+15:05:39.852 INF load_config:444 |     "enable_repetition_penalty": false,
+15:05:39.852 INF load_config:444 |     "enable_temperature": false,
+15:05:39.852 INF load_config:444 |     "enable_top_k_sampling": false,
+15:05:39.852 INF load_config:444 |     "enable_top_p_sampling": false,
+15:05:39.852 INF load_config:444 |     "penalty_window": 64,
+15:05:39.852 INF load_config:444 |     "repetition_penalty": 1.0,
+15:05:39.852 INF load_config:444 |     "temperature": 1.0,
+15:05:39.852 INF load_config:444 |     "top_k": 64,
+15:05:39.852 INF load_config:444 |     "top_p": 0.95
+15:05:39.852 INF load_config:444 | }
+15:05:39.852 INF Init:1293 | LLM init ok
 Starting server on port 8000 with model 'AXERA-TECH/gemma-4-E2B-it'...
 API URLs:
   GET  http://127.0.0.1:8000/health
@@ -329,6 +371,8 @@ root@ax650 ~ # curl http://127.0.0.1:8000/v1/models
 }
 ```
 
+![openai_api_demo](assets/openai_api_demo.png)
+
 ## Python Runtime Requirements
 
 Install the following packages on the AX board:
@@ -482,6 +526,59 @@ Here is a detailed description:
 In summary, it is a cheerful, stylized, red cartoon lobster flexing its claws.
 ```
 
+### Audio Inference
+
+The package includes two fixed-duration audio encoders:
+
+- `gemma4_audio_5s.axmodel` for `5s / 125` audio tokens
+- `gemma4_audio_30s.axmodel` for `30s / 750` audio tokens
+
+For board-side validation, use the packaged WAV clips listed in the sample-audio section above.
+
+Example: `5s` profile
+
+```bash
+python3 infer_axmodel.py \
+  --audio_path ./assets/gemma4_audio_test_5s.wav \
+  --audio_model_path ./gemma4_audio_5s.axmodel \
+  --audio_duration_sec 5 \
+  --audio_tokens 125 \
+  --system_prompt "" \
+  --prompt "Transcribe the speech in its original language. Output only the transcription." \
+  --max_new_tokens 128
+```
+
+Typical output:
+
+```sh
+answer >> When I was seventeen,I read a quote that went something like if you ...
+```
+
+Example: `30s` profile
+
+```bash
+python3 infer_axmodel.py \
+  --audio_path ./assets/gemma4_audio_test_chunk0_30s.wav \
+  --audio_model_path ./gemma4_audio_30s.axmodel \
+  --audio_duration_sec 30 \
+  --audio_tokens 750 \
+  --system_prompt "" \
+  --prompt "Transcribe the speech in its original language. Output only the transcription." \
+  --max_new_tokens 256
+```
+
+Typical output:
+
+```sh
+answer >> No one wants to die. Even people who want to go to heaven don't want to die. Even people who want to go to heaven don't want to die. No one wants to die. Death is the destination we all share and yet death is the destination we all share and no one has ever escaped it and that is as it should be because death is the single best invention of life death is life change out the old make way for the new right now the new is you but someday not too long from now you^@ will gradually become ...
+```
+
+Notes:
+
+- The two commands above were validated on board with a `gemma4_pydeps` dependency overlay added to `PYTHONPATH`.
+- The `30s / 750-token` profile spans multiple `128-token` prefill slices. The runtime will print a warning about partial bidirectional attention across earlier slices inside the same multimodal block. This is expected for chunked prefill.
+- The Python demo loader handles `WAV` directly with the Python standard library. For `mp3 / flac / m4a / ogg`, install `librosa` on the board.
+
 ### Gradio Demo
 
 ```bash
@@ -501,8 +598,9 @@ The Python demo scripts use the following default paths:
 - Tokenizer and config: `./gemma_4_e2b_it_tokenizer`
 - Text LLM runtime root: `./`
 - Vision axmodels: `./`
+- Audio axmodels: `./`
 
-If you move any of these directories, pass the new values with `--hf_model`, `--axmodel_path`, and `--vit_model_path`.
+If you move any of these directories, pass the new values with `--hf_model`, `--axmodel_path`, `--vit_model_path`, and `--audio_model_path`.
 
 For the Python demo flow, `--axmodel_path` should point to the directory that contains the text runtime files such as `gemma4_text_p128_l*.axmodel`, `gemma4_text_post.axmodel`, `model.embed_tokens.weight.bfloat16.bin`, and the `model.*per_layer*.npy` files.
 

assets/gemma4_audio_test_5s.wav

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

assets/gemma4_audio_test_chunk0_30s.wav

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

assets/gemma4_audio_test_chunk1_30s.wav

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

gemma4_audio_30s.axmodel

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

gemma4_audio_5s.axmodel

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

infer_axmodel.py

@@ -8,11 +8,13 @@ from ml_dtypes import bfloat16
 from utils.gemma4_compat import load_text_runtime_config
 from utils.gemma4_compat import load_tokenizer
 from utils.gemma4_multimodal import DEFAULT_MAX_SOFT_TOKENS
+from utils.gemma4_multimodal import prepare_audio_inputs
 from utils.gemma4_multimodal import build_messages
 from utils.gemma4_multimodal import detect_soft_tokens_from_vit_path
 from utils.gemma4_multimodal import load_image
 from utils.gemma4_multimodal import load_processor
 from utils.gemma4_multimodal import prepare_multimodal_inputs
+from utils.gemma4_multimodal import replace_audio_tokens
 from utils.gemma4_multimodal import replace_image_tokens
 from utils.gemma4_multimodal import resolve_resize
 from utils.gemma4_multimodal import to_numpy_fp32
@@ -20,6 +22,8 @@ from utils.gemma4_per_layer import Gemma4PerLayerInputs
 from utils.runtime_layout import default_axmodel_path
 from utils.runtime_layout import load_text_embeddings
 from utils.infer_func import InferManager
+from utils.infer_func import detect_prefill_len
+from utils.infer_func import release_ax_inference_session
 from utils.vision_output import describe_output_shapes
 from utils.vision_output import select_vit_output
 
@@ -60,14 +64,46 @@ def _default_vit_model_path() -> str:
     return str(candidates[0])
 
 
+def _default_audio_profile() -> tuple[str, float, int]:
+    script_dir = Path(__file__).resolve().parent
+    candidates = [
+        (script_dir / "gemma4_audio_30s.axmodel", 30.0, 750),
+        (script_dir / "gemma4_audio_5s.axmodel", 5.0, 125),
+        (script_dir / "audio_models" / "gemma4_audio_30s.axmodel", 30.0, 750),
+        (script_dir / "audio_models" / "gemma4_audio_5s.axmodel", 5.0, 125),
+        (script_dir.parent / "model_convert" / "compiled_output_audio_30s" / "compiled.axmodel", 30.0, 750),
+        (script_dir.parent / "model_convert" / "compiled_output_audio_5s" / "compiled.axmodel", 5.0, 125),
+        (Path("/tmp/compiled_output_audio_30s/compiled.axmodel"), 30.0, 750),
+        (Path("/tmp/compiled_output_audio_5s/compiled.axmodel"), 5.0, 125),
+    ]
+    for candidate, duration_sec, audio_tokens in candidates:
+        if candidate.exists():
+            return str(candidate), duration_sec, audio_tokens
+    return str(candidates[0][0]), candidates[0][1], candidates[0][2]
+
+
+def _infer_audio_profile_from_path(audio_model_path: str) -> tuple[float | None, int | None]:
+    path_str = str(audio_model_path)
+    if "audio_5s" in path_str or "compiled_output_audio_5s" in path_str:
+        return 5.0, 125
+    if "audio_30s" in path_str or "compiled_output_audio_30s" in path_str:
+        return 30.0, 750
+    return None, None
+
+
 
 def _run_vit_axmodel(vit_model_path: str, pixel_values: np.ndarray, target_hidden_size: int, expected_tokens: int):
     from axengine import InferenceSession
 
     session = InferenceSession(vit_model_path)
-    outputs = session.run(None, {"pixel_values": pixel_values})
-    if isinstance(outputs, dict):
-        outputs = list(outputs.values())
+    try:
+        outputs = session.run(None, {"pixel_values": pixel_values})
+        if isinstance(outputs, dict):
+            outputs = [np.array(value, copy=True) for value in outputs.values()]
+        else:
+            outputs = [np.array(value, copy=True) for value in outputs]
+    finally:
+        release_ax_inference_session(session)
     return (
         select_vit_output(outputs, target_hidden_size, expected_tokens=expected_tokens),
         describe_output_shapes(outputs),
@@ -88,6 +124,48 @@ def _run_vit_onnx(vit_model_path: str, pixel_values: np.ndarray, target_hidden_s
     )
 
 
+def _run_audio_axmodel(
+    audio_model_path: str,
+    input_features: np.ndarray,
+    target_hidden_size: int,
+    expected_tokens: int,
+):
+    from axengine import InferenceSession
+
+    session = InferenceSession(audio_model_path)
+    try:
+        outputs = session.run(None, {"input_features": input_features})
+        if isinstance(outputs, dict):
+            outputs = [np.array(value, copy=True) for value in outputs.values()]
+        else:
+            outputs = [np.array(value, copy=True) for value in outputs]
+    finally:
+        release_ax_inference_session(session)
+    return (
+        select_vit_output(outputs, target_hidden_size, expected_tokens=expected_tokens),
+        describe_output_shapes(outputs),
+    )
+
+
+def _run_audio_onnx(
+    audio_model_path: str,
+    input_features: np.ndarray,
+    target_hidden_size: int,
+    expected_tokens: int,
+):
+    import onnxruntime as ort
+
+    providers = ["CPUExecutionProvider"]
+    if "CUDAExecutionProvider" in ort.get_available_providers():
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+    session = ort.InferenceSession(audio_model_path, providers=providers)
+    outputs = session.run(None, {"input_features": input_features})
+    return (
+        select_vit_output(outputs, target_hidden_size, expected_tokens=expected_tokens),
+        describe_output_shapes(outputs),
+    )
+
+
 def _text_prefill_inputs(tokenizer, embeds: np.ndarray, prompt: str, system_prompt: str = "", enable_thinking: bool = False):
     messages = build_messages(prompt=prompt, system_prompt=system_prompt)
     text = tokenizer.apply_chat_template(
@@ -110,8 +188,13 @@ if __name__ == "__main__":
                         help="Path to the packaged LLM runtime root or legacy axmodel folder")
     parser.add_argument("--vit_model_path", type=str, default=_default_vit_model_path(),
                         help="Path to Gemma 4 vision ONNX model or .axmodel")
+    default_audio_model_path, default_audio_duration_sec, default_audio_tokens = _default_audio_profile()
+    parser.add_argument("--audio_model_path", type=str, default=default_audio_model_path,
+                        help="Path to Gemma 4 audio ONNX model or .axmodel")
     parser.add_argument("--image_path", type=str, default="",
                         help="Optional input image path. If omitted, runs text-only generation.")
+    parser.add_argument("--audio_path", type=str, default="",
+                        help="Optional input audio path. Audio-only inference uses the fixed-duration audio encoder path.")
     parser.add_argument("--system_prompt", type=str, default="You are a helpful assistant.",
                         help="Optional system prompt")
     parser.add_argument("--prompt", type=str, default="Describe the image in detail.",
@@ -126,10 +209,41 @@ if __name__ == "__main__":
                         help="Optional fixed image height for the vision encoder")
     parser.add_argument("--resize_w", type=int, default=None,
                         help="Optional fixed image width for the vision encoder")
-    parser.add_argument("--slice_len", type=int, default=128,
-                        help="Prefill slice length. Must match the LLM `--prefill_len` used at build time.")
+    parser.add_argument("--slice_len", type=int, default=None,
+                        help="Prefill slice length. Must match the LLM `--prefill_len` used at build time. "
+                        "When omitted, auto-detected from axmodel filenames under --axmodel_path.")
+    parser.add_argument("--audio_duration_sec", type=float, default=None,
+                        help="Fixed audio duration expected by the audio encoder model. "
+                        "When omitted, infer it from --audio_model_path.")
+    parser.add_argument("--audio_tokens", type=int, default=None,
+                        help="Fixed number of audio soft tokens produced by the audio encoder model. "
+                        "When omitted, infer it from --audio_model_path.")
+    parser.add_argument("--audio_embeds_npy", type=str, default="",
+                        help="Optional path to a pre-computed (1, audio_tokens, hidden_size) or "
+                        "(audio_tokens, hidden_size) float32 npy. When set, the audio encoder step "
+                        "(axmodel / onnx) is bypassed and these embeds are used directly for prefill. "
+                        "Useful for A/B testing torch-reference audio embeds against the on-device audio axmodel.")
     args = parser.parse_args()
 
+    if args.slice_len is None:
+        args.slice_len = detect_prefill_len(args.axmodel_path, default=128)
+        print(f"[INFO] Auto-detected slice_len={args.slice_len} from {args.axmodel_path}")
+
+    if args.audio_duration_sec is None or args.audio_tokens is None:
+        inferred_duration_sec, inferred_audio_tokens = _infer_audio_profile_from_path(args.audio_model_path)
+        if args.audio_duration_sec is None:
+            args.audio_duration_sec = inferred_duration_sec or default_audio_duration_sec
+        if args.audio_tokens is None:
+            args.audio_tokens = inferred_audio_tokens or default_audio_tokens
+
+    if args.image_path and args.audio_path:
+        # Gemma4 natively supports image+audio in the same prompt; this script's prefill
+        # path currently only wires up a single modality's prepare_* helper. Lifting this
+        # requires extending prepare_* to produce a combined inputs dict with both
+        # pixel_values and input_features, and running both encoders before replacing
+        # tokens. Left as future work.
+        raise ValueError("Simultaneous image+audio inputs are not supported by this demo script yet.")
+
     config = load_text_runtime_config(args.hf_model)
     embeds = load_text_embeddings(args.axmodel_path, config)
     per_layer_helper = None
@@ -143,9 +257,6 @@ if __name__ == "__main__":
             print(f"[INFO] Auto-detected max_soft_tokens={detected} from VIT model: {args.vit_model_path}")
             args.max_soft_tokens = detected
 
-    kv_cache_len = int(getattr(config, "kv_cache_len", 2047) or 2047)
-    imer = InferManager(config, args.axmodel_path, max_seq_len=kv_cache_len, per_layer_helper=per_layer_helper)
-
     mm_token_type_ids = None
     prefill_per_layer_inputs = None
     if args.image_path:
@@ -204,6 +315,71 @@ if __name__ == "__main__":
             image_embeds,
             image_token_id=config.image_token_id,
         ).astype(bfloat16)
+    elif args.audio_path:
+        processor = load_processor(args.hf_model)
+        tokenizer = processor.tokenizer
+        mm = prepare_audio_inputs(
+            processor,
+            audio_path=args.audio_path,
+            prompt=args.prompt,
+            system_prompt=args.system_prompt,
+            enable_thinking=args.enable_thinking,
+            audio_duration_sec=args.audio_duration_sec,
+            fixed_audio_tokens=args.audio_tokens,
+        )
+        inputs = mm["inputs"]
+        token_ids = inputs["input_ids"][0].cpu().numpy().tolist()
+        mm_token_type_ids = inputs["mm_token_type_ids"][0].cpu().numpy().tolist()
+        input_features = to_numpy_fp32(inputs["input_features"])
+
+        if args.audio_tokens > args.slice_len:
+            print(
+                f"[WARN] audio_tokens={args.audio_tokens} exceeds slice_len={args.slice_len}. "
+                "Cross-slice audio blocks are supported, but the current chunked prefill path does not provide "
+                "full future-token visibility to earlier slices within the same multimodal block."
+            )
+
+        if args.audio_embeds_npy:
+            audio_embeds = np.load(args.audio_embeds_npy).astype(np.float32)
+            audio_output_shapes = [tuple(int(v) for v in audio_embeds.shape)]
+            print(f"[INFO] loaded precomputed audio_embeds from {args.audio_embeds_npy}, shape={audio_embeds.shape}")
+        elif not os.path.exists(args.audio_model_path):
+            raise FileNotFoundError(
+                f"Audio model not found: {args.audio_model_path}. "
+                "Please export and compile the Gemma 4 audio encoder first, "
+                "or pass --audio_embeds_npy to bypass the audio encoder."
+            )
+        elif args.audio_model_path.endswith(".axmodel"):
+            audio_embeds, audio_output_shapes = _run_audio_axmodel(
+                args.audio_model_path,
+                input_features,
+                target_hidden_size=config.hidden_size,
+                expected_tokens=args.audio_tokens,
+            )
+        else:
+            audio_embeds, audio_output_shapes = _run_audio_onnx(
+                args.audio_model_path,
+                input_features,
+                target_hidden_size=config.hidden_size,
+                expected_tokens=args.audio_tokens,
+            )
+
+        if audio_embeds.ndim == 3:
+            audio_embeds = audio_embeds[0]
+        if audio_embeds.shape[0] != args.audio_tokens:
+            raise ValueError(
+                "Unexpected audio output token count. "
+                f"got={audio_embeds.shape[0]}, expected={args.audio_tokens}, "
+                f"audio_output_shapes={audio_output_shapes}"
+            )
+
+        prefill_data = np.take(embeds, token_ids, axis=0)
+        prefill_data = replace_audio_tokens(
+            token_ids,
+            prefill_data,
+            audio_embeds,
+            audio_token_id=config.audio_token_id,
+        ).astype(bfloat16)
     else:
         tokenizer = load_tokenizer(args.hf_model)
         token_ids, prefill_data = _text_prefill_inputs(
@@ -232,20 +408,25 @@ if __name__ == "__main__":
 
     eos_token_id = config.eos_token_id if isinstance(config.eos_token_id, list) else None
 
-    token_ids = imer.prefill(
-        tokenizer,
-        token_ids,
-        prefill_data,
-        mm_token_type_ids=mm_token_type_ids,
-        slice_len=args.slice_len,
-        per_layer_inputs=prefill_per_layer_inputs,
-    )
-    imer.decode(
-        tokenizer,
-        token_ids,
-        embeds,
-        slice_len=args.slice_len,
-        eos_token_id=eos_token_id,
-        max_new_tokens=args.max_new_tokens,
-    )
+    kv_cache_len = int(getattr(config, "kv_cache_len", 2047) or 2047)
+    imer = InferManager(config, args.axmodel_path, max_seq_len=kv_cache_len, per_layer_helper=per_layer_helper)
+    try:
+        token_ids = imer.prefill(
+            tokenizer,
+            token_ids,
+            prefill_data,
+            mm_token_type_ids=mm_token_type_ids,
+            slice_len=args.slice_len,
+            per_layer_inputs=prefill_per_layer_inputs,
+        )
+        imer.decode(
+            tokenizer,
+            token_ids,
+            embeds,
+            slice_len=args.slice_len,
+            eos_token_id=eos_token_id,
+            max_new_tokens=args.max_new_tokens,
+        )
+    finally:
+        imer.close()
     print("\n")

post_config.json

@@ -1,10 +1,10 @@
 {
-  "enable_temperature": true,
+  "enable_temperature": false,
   "temperature": 1.0,
   "enable_repetition_penalty": false,
   "repetition_penalty": 1.0,
   "penalty_window": 64,
-  "enable_top_p_sampling": true,
+  "enable_top_p_sampling": false,
   "top_p": 0.95,
   "enable_top_k_sampling": false,
   "top_k": 64

utils/gemma4_compat.py

@@ -13,6 +13,7 @@ def load_text_runtime_config(model_dir: str):
     text_config = dict(raw_config.get("text_config") or raw_config)
     text_config["model_type"] = text_config.get("model_type", "gemma4_text")
     text_config["eos_token_id"] = raw_config.get("eos_token_id", text_config.get("eos_token_id"))
+    text_config["audio_token_id"] = raw_config.get("audio_token_id")
     text_config["image_token_id"] = raw_config.get("image_token_id")
     text_config["vision_config"] = raw_config.get("vision_config")
     text_config["vision_soft_tokens_per_image"] = raw_config.get("vision_soft_tokens_per_image")

utils/gemma4_multimodal.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import os
 from pathlib import Path
+import wave
 
 import numpy as np
 import torch
@@ -103,25 +105,86 @@ def load_image(image_path: str | Path) -> Image.Image:
     return Image.open(image_path).convert("RGB")
 
 
+def _resample_waveform(waveform: np.ndarray, src_rate: int, dst_rate: int) -> np.ndarray:
+    # Linear interpolation; introduces aliasing when downsampling (e.g. 44.1kHz -> 16kHz).
+    # For best quality pass mono 16kHz WAV and skip this path; otherwise librosa.load
+    # with its polyphase resampler is preferred.
+    if src_rate == dst_rate or waveform.size == 0:
+        return waveform.astype(np.float32)
+
+    src_positions = np.arange(waveform.shape[0], dtype=np.float32) / float(src_rate)
+    dst_length = max(1, int(round(waveform.shape[0] * float(dst_rate) / float(src_rate))))
+    dst_positions = np.arange(dst_length, dtype=np.float32) / float(dst_rate)
+    return np.interp(dst_positions, src_positions, waveform).astype(np.float32)
+
+
+def load_audio_waveform(audio_path: str | Path, sampling_rate: int = 16000) -> np.ndarray:
+    audio_path = Path(audio_path)
+    if audio_path.suffix.lower() == ".wav":
+        with wave.open(str(audio_path), "rb") as wav_file:
+            if wav_file.getcomptype() != "NONE":
+                raise ValueError(f"Unsupported WAV compression type: {wav_file.getcomptype()}")
+
+            channels = wav_file.getnchannels()
+            sample_width = wav_file.getsampwidth()
+            src_rate = wav_file.getframerate()
+            frames = wav_file.readframes(wav_file.getnframes())
+
+        if sample_width == 1:
+            waveform = np.frombuffer(frames, dtype=np.uint8).astype(np.float32)
+            waveform = (waveform - 128.0) / 128.0
+        elif sample_width == 2:
+            waveform = np.frombuffer(frames, dtype="<i2").astype(np.float32) / 32768.0
+        elif sample_width == 4:
+            waveform = np.frombuffer(frames, dtype="<i4").astype(np.float32) / 2147483648.0
+        else:
+            raise ValueError(f"Unsupported WAV sample width: {sample_width}")
+
+        if channels > 1:
+            waveform = waveform.reshape(-1, channels).mean(axis=1)
+        waveform = _resample_waveform(waveform, src_rate=src_rate, dst_rate=sampling_rate)
+        return np.asarray(np.clip(waveform, -1.0, 1.0), dtype=np.float32)
+
+    try:
+        os.environ.setdefault("NUMBA_DISABLE_JIT", "1")
+        os.environ.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache")
+        import librosa
+    except Exception as exc:
+        raise RuntimeError(
+            "Non-WAV audio loading requires librosa. Please convert the input to a mono 16kHz WAV file."
+        ) from exc
+
+    waveform, _ = librosa.load(str(audio_path), sr=sampling_rate, mono=True)
+    return np.asarray(waveform, dtype=np.float32)
+
+
 def resize_image(image: Image.Image, resize_h: int, resize_w: int) -> Image.Image:
     return image.convert("RGB").resize((resize_w, resize_h), resample=Image.BICUBIC)
 
 
-def build_messages(prompt: str, image: Image.Image | None = None, system_prompt: str = "") -> list[dict]:
+def build_messages(
+    prompt: str,
+    image: Image.Image | None = None,
+    audio=None,
+    system_prompt: str = "",
+) -> list[dict]:
     messages: list[dict] = []
     if system_prompt.strip():
         messages.append({"role": "system", "content": system_prompt})
 
-    if image is None:
+    if image is None and audio is None:
         messages.append({"role": "user", "content": prompt})
     else:
+        user_content = []
+        if audio is not None:
+            user_content.append({"type": "audio", "audio": audio})
+        if image is not None:
+            user_content.append({"type": "image", "image": image})
+        user_content.append({"type": "text", "text": prompt})
         messages.append(
             {
                 "role": "user",
-                "content": [
-                    {"type": "image", "image": image},
-                    {"type": "text", "text": prompt},
-                ],
+                "content": user_content,
             }
         )
     return messages
@@ -130,6 +193,7 @@ def build_messages(prompt: str, image: Image.Image | None = None, system_prompt:
 def build_processor_messages(
     prompt: str,
     image: Image.Image | None = None,
+    audio=None,
     system_prompt: str = "",
     history=None,
 ) -> list[dict]:
@@ -157,6 +221,8 @@ def build_processor_messages(
             messages.append({"role": "assistant", "content": bot_msg})
 
     user_content = []
+    if audio is not None:
+        user_content.append({"type": "audio", "audio": audio})
     if image is not None:
         user_content.append({"type": "image", "image": image})
     user_content.append({"type": "text", "text": prompt})
@@ -261,25 +327,107 @@ def prepare_multimodal_inputs(
     }
 
 
-def replace_image_tokens(token_ids, token_embeds, image_embeds, image_token_id: int):
-    image_positions = [idx for idx, token_id in enumerate(token_ids) if int(token_id) == int(image_token_id)]
-    if not image_positions:
+def prepare_audio_inputs(
+    processor,
+    audio_path: str | Path,
+    prompt: str,
+    system_prompt: str = "",
+    enable_thinking: bool = False,
+    audio_duration_sec: float = 30.0,
+    fixed_audio_tokens: int | None = None,
+):
+    sampling_rate = int(getattr(processor.feature_extractor, "sampling_rate", 16000))
+    max_length = int(round(audio_duration_sec * sampling_rate))
+    waveform = load_audio_waveform(audio_path, sampling_rate=sampling_rate)
+
+    if waveform.shape[0] < max_length:
+        padded_waveform = np.pad(waveform, (0, max_length - waveform.shape[0]), mode="constant")
+    else:
+        padded_waveform = waveform[:max_length]
+
+    messages = build_processor_messages(prompt=prompt, audio=padded_waveform, system_prompt=system_prompt)
+    inputs = _safe_apply_chat_template(
+        processor,
+        messages,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+        add_generation_prompt=True,
+        enable_thinking=enable_thinking,
+        processor_kwargs={
+            "audio_kwargs": {
+                "padding": "max_length",
+                "max_length": max_length,
+                "truncation": True,
+                "pad_to_multiple_of": None,
+            }
+        },
+    )
+
+    audio_token_id = processor.audio_token_id
+    audio_token_count = int((inputs["input_ids"] == audio_token_id).sum().item())
+    expected_tokens = int(fixed_audio_tokens or processor.audio_seq_length)
+    if audio_token_count != expected_tokens:
+        raise ValueError(
+            f"Expected {expected_tokens} audio soft tokens from fixed audio preprocessing, got {audio_token_count}."
+        )
+
+    return {
+        "messages": messages,
+        "waveform": waveform,
+        "padded_waveform": padded_waveform,
+        "inputs": inputs,
+        "audio_duration_sec": audio_duration_sec,
+        "expected_tokens": expected_tokens,
+    }
+
+
+def replace_special_tokens(
+    token_ids,
+    token_embeds,
+    modality_embeds,
+    special_token_id: int,
+    modality_name: str,
+):
+    positions = [idx for idx, token_id in enumerate(token_ids) if int(token_id) == int(special_token_id)]
+    if not positions:
         return token_embeds
 
-    flat_image_embeds = image_embeds.reshape(-1, image_embeds.shape[-1])
-    if len(image_positions) != flat_image_embeds.shape[0]:
+    flat_embeds = modality_embeds.reshape(-1, modality_embeds.shape[-1])
+    if len(positions) != flat_embeds.shape[0]:
         raise ValueError(
-            f"Image tokens and image features do not match: tokens={len(image_positions)}, "
-            f"features={flat_image_embeds.shape[0]}"
+            f"{modality_name.capitalize()} tokens and features do not match: "
+            f"tokens={len(positions)}, features={flat_embeds.shape[0]}"
         )
-    if token_embeds.shape[-1] != flat_image_embeds.shape[-1]:
+    if token_embeds.shape[-1] != flat_embeds.shape[-1]:
         raise ValueError(
-            f"Embedding dim mismatch: token_dim={token_embeds.shape[-1]}, image_dim={flat_image_embeds.shape[-1]}"
+            f"Embedding dim mismatch: token_dim={token_embeds.shape[-1]}, "
+            f"{modality_name}_dim={flat_embeds.shape[-1]}"
         )
-    token_embeds[image_positions, :] = flat_image_embeds
+    token_embeds[positions, :] = flat_embeds
     return token_embeds
 
 
+def replace_image_tokens(token_ids, token_embeds, image_embeds, image_token_id: int):
+    return replace_special_tokens(
+        token_ids,
+        token_embeds,
+        image_embeds,
+        special_token_id=image_token_id,
+        modality_name="image",
+    )
+
+
+def replace_audio_tokens(token_ids, token_embeds, audio_embeds, audio_token_id: int):
+    return replace_special_tokens(
+        token_ids,
+        token_embeds,
+        audio_embeds,
+        special_token_id=audio_token_id,
+        modality_name="audio",
+    )
+
+
 def to_numpy_fp32(tensor_like) -> np.ndarray:
     if isinstance(tensor_like, np.ndarray):
         return tensor_like.astype(np.float32)

utils/infer_func.py

@@ -1,3 +1,4 @@
+import atexit
 import os
 import re
 
@@ -7,6 +8,23 @@ from ml_dtypes import bfloat16
 from tqdm import tqdm
 
 
+def release_ax_inference_session(session):
+    inner = getattr(session, "_sess", None)
+    unload = getattr(inner, "_unload", None)
+    if not callable(unload):
+        return
+
+    try:
+        unload()
+    except Exception as exc:
+        print(f"[WARN] Failed to unload axengine session cleanly: {exc}")
+    finally:
+        try:
+            inner._unload = lambda: None
+        except Exception:
+            pass
+
+
 def _layer_head_dim(config, layer_idx: int) -> int:
     if getattr(config, "layer_types", None) and layer_idx < len(config.layer_types):
         if config.layer_types[layer_idx] == "full_attention":
@@ -33,6 +51,28 @@ def _build_shared_kv_source_layers(config):
     return source_layers
 
 
+def detect_prefill_len(model_dir: str, default: int = 128) -> int:
+    """Auto-detect prefill_len (aka slice_len) from axmodel filenames.
+
+    Matches files named like ``<prefix>_p<N>_l<idx>_together.axmodel`` and returns
+    ``N``. Falls back to ``default`` when no layer files match.
+    """
+    layer_pattern = re.compile(r"^.*_p(?P<prefill>\d+)_l\d+_together\.axmodel$")
+    prefill_counts = {}
+    try:
+        for fname in os.listdir(model_dir):
+            match = layer_pattern.match(fname)
+            if match:
+                prefill = int(match.group("prefill"))
+                prefill_counts[prefill] = prefill_counts.get(prefill, 0) + 1
+    except FileNotFoundError:
+        return default
+
+    if not prefill_counts:
+        return default
+    return max(prefill_counts.items(), key=lambda kv: kv[1])[0]
+
+
 def _find_axmodel_files(base_dir: str, expected_layers: int = None, expected_prefill: int = 128):
     files = os.listdir(base_dir)
     layer_pattern = re.compile(r"^(?P<prefix>.*)_p(?P<prefill>\d+)_l(?P<idx>\d+)_together\.axmodel$")
@@ -105,20 +145,38 @@ class InferManager:
         if post_file is None:
             raise FileNotFoundError("Cannot find post process .axmodel file in model_dir")
         self.post_process_session = InferenceSession(os.path.join(model_dir, post_file))
+        self._closed = False
+        atexit.register(self.close)
         print("Model loaded successfully!")
 
+    def close(self):
+        if self._closed:
+            return
+
+        sessions = list(getattr(self, "decoder_sessions", []))
+        post_process_session = getattr(self, "post_process_session", None)
+        if post_process_session is not None:
+            sessions.append(post_process_session)
+
+        for session in sessions:
+            release_ax_inference_session(session)
+
+        self.decoder_sessions = []
+        self.post_process_session = None
+        self._closed = True
+
     @staticmethod
     def _compute_mm_group_ids(mm_token_type_ids):
         if mm_token_type_ids is None:
             return None
 
         mm_token_type_ids = np.asarray(mm_token_type_ids, dtype=np.int32).reshape(-1)
-        is_vision = np.isin(mm_token_type_ids, (1, 2))
-        prev_is_vision = np.roll(is_vision, 1)
-        prev_is_vision[0] = False
-        new_group_starts = is_vision & ~prev_is_vision
+        is_multimodal = np.isin(mm_token_type_ids, (1, 2, 3))
+        prev_is_multimodal = np.roll(is_multimodal, 1)
+        prev_is_multimodal[0] = False
+        new_group_starts = is_multimodal & ~prev_is_multimodal
         group_ids = np.cumsum(new_group_starts.astype(np.int32)) - 1
-        group_ids[~is_vision] = -1
+        group_ids[~is_multimodal] = -1
         return group_ids
 
     @staticmethod
@@ -136,7 +194,7 @@ class InferManager:
             if start // slice_len != end // slice_len:
                 num_slices = end // slice_len - start // slice_len + 1
                 print(
-                    f"[WARN] Image token block (group_id={group_id}, pos {start}-{end}) "
+                    f"[WARN] Multimodal token block (group_id={group_id}, pos {start}-{end}) "
                     f"spans {num_slices} prefill slices. Bidirectional attention within "
                     f"earlier slices is partial (chunked prefill limitation)."
                 )
@@ -340,7 +398,7 @@ class InferManager:
         if mm_token_type_ids is None:
             return True
         token_type = int(mm_token_type_ids[token_pos])
-        return token_type not in (1, 2)
+        return token_type not in (1, 2, 3)
 
     @staticmethod
     def _build_decode_mask(cache_len: int, visible_past_tokens: int):