Skip to content

zml/examples: fix Qwen3.5 models on vanilla backend and 4+ GPU#602

Open
phildav wants to merge 1 commit into
zml:masterfrom
phildav:fix/qwen3_5-kv-head-double-stutter
Open

zml/examples: fix Qwen3.5 models on vanilla backend and 4+ GPU#602
phildav wants to merge 1 commit into
zml:masterfrom
phildav:fix/qwen3_5-kv-head-double-stutter

Conversation

@phildav

@phildav phildav commented Jun 25, 2026

Copy link
Copy Markdown

Issue description

With 4 or more devices, 4+ GPU or vanilla 4+CPU, the Qwen models break in forward:
thread 2011 panic: dynamicUpdateSlice expects 'update_' dimensions to be less than or equal to their corresponding dimension in input tensor, got 8 and 4 for axis .h

$ bazel run //examples/llm --@zml//platforms:cuda=true -- --model=hf://Qwen/Qwen3.5-0.8B --prompt="What is the capital of France ?"

....

info(llm): Selected backend: .cuda_fa2
info(llm): Resolving model repository..
info(zml/io/vfs/hf): Fetching from HF: https://huggingface.co/api/models/Qwen/Qwen3.5-0.8B/tree/main/?expand=false&recursive=true&limit=1000
info(llm): Initializing model..
info(llm): Detected model type: .qwen3_5
info(llm): Loaded tokenizer [2.834s]
info(qwen3_5): Compiled prefill embed tokens [517.983ms]
thread 2011 panic: dynamicUpdateSlice expects 'update_' dimensions to be less than or equal to their corresponding dimension in input tensor, got 8 and 4 for axis .h
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/stdx/debug.zig:18:20: 0x55fcba5c4ca0 in panic (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/zml/ops.zig:958:41: 0x55fcba5c2179 in scatter__anon_137706 (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/zml/tensor.zig:2840:27: 0x55fcba5bf04e in scatterSlices__anon_137667 (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/examples/llm/models/qwen3_5/model.zig:1018:44: 0x55fcba5b73d7 in update (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/examples/llm/models/qwen3_5/model.zig:594:45: 0x55fcba5ac838 in forward (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/examples/llm/models/qwen3_5/model.zig:391:75: 0x55fcba5a8e36 in forwardSelfAttn (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/zml/module.zig:469:9: 0x55fcba096f2d in emitMlir__anon_17745 (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/zml/module.zig:223:28: 0x55fcba09907a in compile__anon_17739 (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/zml/platform.zig:523:34: 0x55fcba099caa in compileFn__anon_17735 (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/zml/platform.zig:494:30: 0x55fcba099ec1 in compile__anon_17729 (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/examples/llm/models/qwen3_5/inference.zig:617:32: 0x55fcba09a39f in compileFullAttentionLayer (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/examples/llm/models/qwen3_5/inference.zig:337:69: 0x55fcba0a3bd7 in init (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/examples/llm/models/qwen3_5/inference.zig:280:34: 0x55fcba0a4275 in init (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/examples/llm/models/qwen3_5/inference.zig:64:33: 0x55fcba0a44c4 in init (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/examples/llm/models/qwen3_5/model.zig:127:44: 0x55fcba0a48c3 in compile (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/examples/llm/models.zig:116:57: 0x55fcba0bd489 in compile (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/examples/llm/main.zig:131:56: 0x55fcba10b1a6 in main (examples_Sllm_Cllm)
/root/.cache/bazel/_bazel_root/6bd2413cce69270a0469b5f108d9aa53/sandbox/processwrapper-sandbox/3353/execroot/_main/external/rules_zig++zig+zig_0.16.0_x86_64-linux/lib/std/start.zig:737:30: 0x55fcba10cfb8 in callMain (examples_Sllm_Cllm)
???:?:?: 0x7f71c088f1c9 in ??? (/usr/lib/x86_64-linux-gnu/libc.so.6)
???:?:?: 0x7f71c088f28a in ??? (/usr/lib/x86_64-linux-gnu/libc.so.6)
???:?:?: 0x55fcba051289 in ??? (???)
Aborted] Compiling prefill full attention layer...

This is due to stutter1d being called twice.

How to reproduce

On a 4xGPU box:
$ bazel run //examples/llm --@zml//platforms:cuda=true -- --model=hf://Qwen/Qwen3.5-0.8B --prompt="What is the capital of France ?"

Or just by forcing vanilla, since it selects 4 devices:
$ bazel run //examples/llm --@zml//platforms:cuda=true -- --model=hf://Qwen/Qwen3.5-0.8B --prompt="What is the capital of France ?" --backend=vanilla

Fix

The fix is adding a bool param to avoid calling stutter1d() a second time.
Open to discussion on other ways to fix it.

Tested

  • Checked fix on 4 x RTX4090 with 0.8B dense model 🔴 -> ✅
  • Checked fix with --backend=vanilla on 0.8B dense model 🔴 -> ✅
  • Checked 1x RTX4090 still working ✅ -> ✅
  • Checked fix on with 35B MoE model 🔴 -> ✅

@phildav phildav changed the title zml/examples: fix partitionCachedKv for Qwen3.5 models on 4+ GPU/vanilla config zml/examples: fix Qwen3.5 models on vanilla backend and 4+ GPU Jun 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant