Head parallel and batch fold blocking for Qwen3-VL MOE by kdulla · Pull Request #1229 · quic/efficient-transformers

kdulla · 2026-07-31T12:27:52Z

Overview

This PR adds support for new blocking modes, including head parallel blocking, prefill blocking, decode MOE blocking and batch-fold blocking. Head parallel blocking is introduced for causal models while the other approaches are introuced for Qwen3-VL MOE. Expert Parallelism is enabled via expert_parallel compile option while new flags in qaic_config can be passed for different attention blocking configurations. Additionally, this PR contains updated tests for the new blocking modes introduced.

Modes Added

Head parallel softmax kv blocking:
- split heads across cores, enabled for both causal models and Qwen3VL-moe
- Enabled by passing kv_blocking_headpar_split (Default: 0 or 16) in qaic config with kv blocking mode
- invokes newly created blocked_kv_attention_forward_headpar_offline attention forward
Batch fold + head parallel kv blocking:
- fold batch size and head dimensions before applying head parallel dimension, intended for high batch size use cases, enabled for Qwen3-VL-moe
- Enabled by passing batch_fold: True in qaic config with kv blocking mode
- invokes newly created blocked_kv_attention_forward_decode_headpar_batch attention forward
- Alters cache layout to be [1, BS * NH, SL, D]
Online softmax prefill:
- Applies online softmax to chunked prefill blocks, intended for long-context prefill efficiency
- Enabled by passing prefill_blocking_mode: "online", prefill_block_chunks, and prefill_n_rep_chunk in qaic_config
- Invokes newly created blocked_qkv_attention_forward_prefill_online attention forward

Additionally, pre-existing KV blocking modes have been updated to correctly use ctx_len from compile_time for computing block sizes instead of past_seen-tokens.

Scripts

E2E run on Qwen3-VL moe with batch fold and prefill blocking can be ran via:
python examples/image_text_to_text/models/qwen3_vl_moe/qwen3_vl_disagg_mode.py

Unit tests on Qwen3-VL moe can be ran via:
pytest tests/transformers/models/image_text_to_text/test_qwen3_vl_moe_disagg_blocked.py

Unit tests for head parallel softmax on causal models added to tests/transformers/models/causal_lm_models/test_causal_lm_blocking_hqkv.py

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com> Co-authored-by: Kushal Dulla <kdulla@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

## Summary of Changes - Adds examples/image_text_to_text/models/qwen3_vl_moe/qwen3_vl_moe_blocking_headpar_example.py, a self-contained end-to-end script that exercises three blocked-attention configurations for the Qwen3-VL-MoE - Added support in Qwen3VL-MOE forward to use prefill_blocked_interface if requested by user. ## Blocking modes covered in example script - --blocking-mode decode Decode QPC: HQKV + kv_blocking_headpar_split=0 - --blocking-mode prefill_par Decode QPC: HQKV Prefill QPC: prefill_blocking_mode="qkv" Description: Separate QPCs; chunked-prefill + manual decode loop - --blocking-mode prefill_online Decode QPC: HQKV Prefill QPC: prefill_blocking_mode="online" Description: Same as above with online-softmax prefill - --blocking-mode all Decode QPC: — Prefill QPC: — Description: Runs all three back-to-back ## Usage python examples/image_text_to_text/models/qwen3_vl_moe/qwen3_vl_moe_blocking_headpar_example.py --blocking-mode all --------- Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Updated qwen3vl moe disaggregated inference example script --------- Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Signed-off-by: ochougul <ochougul@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Updated missing write in prefill mode and changed an incorrect skip on min pos id instead of max. --------- Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

added batch fold decode attention forward and updated example scripts accordingly --------- Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com> Signed-off-by: kdulla <kdulla@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Signed-off-by: ochougul <ochougul@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

…quic#1186) Identified incorrect implementation in read cache, updated cache layout and read calls accordingly, still has ongoing compilation bug. Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Signed-off-by: ochougul <ochougul@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Prefill and decode run for 1 iteration Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

…quic#1192) Adjusted prefill normalization and decode cache layout in example script --------- Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com> Signed-off-by: kdulla <kdulla@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Added qwen3 moe optimizations to qwen3vl moe, and changed export BS to 2 if >1. Checked e2e run, able to run generation with the example script: `python examples/image_text_to_text/models/qwen3_vl_moe/qwen3_vl_disagg_mode.py` --------- Signed-off-by: vbaddi <vbaddi@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com> Co-authored-by: vbaddi <vbaddi@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Added in progress changes for benchmark alignment, including updated CtxChunkScatter, reshape before custom op and introduction of kv unroll option Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

…16 (quic#1223) Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Signed-off-by: Ann Kuruvilla <akuruvil@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

…locking tests to disagg infra Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

ochougul · 2026-08-07T07:25:18Z

        inputs[f"past_value.{i}"] = qpc_out[f"past_value.{i}_RetainedState"]
+print(f"ttft={(time.time() - loop_start):.3f}")
+
+exit(0)


ochougul · 2026-08-07T08:01:33Z


-all_outputs.append(np.argmax(decode_out["logits"]))
-pos_id = np.max(decode_inputs["position_ids"], axis=-1, keepdims=True) + 1
+# exit(0)


ochougul · 2026-08-07T08:01:58Z

+# for i in range(config.text_config.num_hidden_layers):
+#     loop_decode_inputs[f"past_key.{i}"] = decode_out[f"past_key.{i}_RetainedState"]
+#     loop_decode_inputs[f"past_value.{i}"] = decode_out[f"past_value.{i}_RetainedState"]


ochougul · 2026-08-07T08:02:05Z

+    # for j in range(config.text_config.num_hidden_layers):
+    #     loop_decode_inputs[f"past_key.{j}"] = decode_out[f"past_key.{j}_RetainedState"]
+    #     loop_decode_inputs[f"past_value.{j}"] = decode_out[f"past_value.{j}_RetainedState"]


Resolved, these were older experiments being done that were incorrectly pushed. Cleaned up example script code accordingly

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

ochougul · 2026-08-07T11:45:02Z

    BlockingMode,
    generic_blocked_attention_interface,
    past_key_value_update,
+    prefill_blocked_attention_interface,


why are we introducing new interface ? can't it be handled within, we also did siimlar for latent attention?

ochougul · 2026-08-07T11:45:11Z

+        # routing_weights = torch.zeros_like(router_logits)
+        # routing_weights.scatter_(1, top_i, top_w)


remove commented code

ochougul · 2026-08-07T11:46:44Z

Do we really need 3 example files, all of them can be like CLI options for one example file.

And default value is the one that is best performant.

kdulla marked this pull request as draft July 31, 2026 12:28

ochougul marked this pull request as ready for review August 5, 2026 07:36

ochougul marked this pull request as draft August 5, 2026 07:39

kdulla force-pushed the qwen3_vl_moe_kv_main branch 2 times, most recently from fea1490 to 4f454a7 Compare August 6, 2026 06:40

kdulla marked this pull request as ready for review August 6, 2026 07:31

kdulla requested review from ochougul, quic-hemagnih, quic-rishinr and vbaddi August 6, 2026 07:32

ochougul added the ci_ready label Aug 6, 2026

ochougul and others added 19 commits August 7, 2026 11:45

Headpar blocking qwen3vl (quic#1152)

60995db

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com> Co-authored-by: Kushal Dulla <kdulla@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Qwen3 vl moe kv disagg script update (quic#1157)

ca878db

Updated qwen3vl moe disaggregated inference example script --------- Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

added qwen3_vl MOE EP via expert_parallel compile option (quic#1156)

d302380

Signed-off-by: ochougul <ochougul@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

fixed MAD issue in prefill head par (quic#1158)

9f79be5

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Fixed accuracy mismatch between Qwen3VL-MOE prefill modes (quic#1161)

7c54894

Updated missing write in prefill mode and changed an incorrect skip on min pos id instead of max. --------- Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

added tree-reduce MOE and standardized MOE params naming

f2de16f

Signed-off-by: ochougul <ochougul@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

updated example script

63197ce

Signed-off-by: ochougul <ochougul@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

fixed normal KV blocking for decode

b1e5f67

Signed-off-by: ochougul <ochougul@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Identified bug in batch fold cache layout and updated to correct read (…

b13f168

…quic#1186) Identified incorrect implementation in read cache, updated cache layout and read calls accordingly, still has ongoing compilation bug. Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

fixed num_kv_heads

c2a74e5

Signed-off-by: ochougul <ochougul@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Fixed cache layout in example script (quic#1191)

87803e7

Prefill and decode run for 1 iteration Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

undid TracedMatmul experimental feature

cf3d471

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

cleaned up dead code and added blocking vlm tests

6d01639

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Updated blocking tests to validate full model tests for qwen3vl

7c9eddb

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Added in progress changes for benchmark alignment (quic#1219)

fd4b662

Added in progress changes for benchmark alignment, including updated CtxChunkScatter, reshape before custom op and introduction of kv unroll option Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

kdulla and others added 11 commits August 7, 2026 11:52

Updated reshape before CtxChunkScatterBatch to use static NH (quic#1221)

301fc1f

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Aligned transpose order to benchmark and updated example script to TS…

479989e

…16 (quic#1223) Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Fix accuracy issue in reading V blocks

49bfc1c

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Constant axis shapes in k norm

e956db6

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Updated example script to 10k CL

263bb7b

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

formatting and remove debug prints

2f9acf4

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

FIx total_cores based on TS value

0e19959

Signed-off-by: Ann Kuruvilla <akuruvil@qti.qualcomm.com> Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Updated kv blocking methods to correctly use CL and updated qwen3vl b…

211fda9

…locking tests to disagg infra Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

minor formatting fixes

6a61b38

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

reshape in customop from benchmark

d1dde9d

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Formatting fixes

148366a

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

kdulla force-pushed the qwen3_vl_moe_kv_main branch from 2f7d72b to 148366a Compare August 7, 2026 06:23

kdulla added 2 commits August 7, 2026 11:54

Fixed duplicate argument from rebase

7cb2e57

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Removed partial CB support

00229fb

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

ochougul reviewed Aug 7, 2026

View reviewed changes

kdulla added 2 commits August 7, 2026 14:38

cleanup commented code in examples

1d19a41

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

Fixed incorrect prefill seq len for disagg tests

4beea60

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>

ochougul reviewed Aug 7, 2026

View reviewed changes

Navigation Menu

Uh oh!

Head parallel and batch fold blocking for Qwen3-VL MOE - #1229

Head parallel and batch fold blocking for Qwen3-VL MOE#1229
kdulla wants to merge 34 commits into
quic:mainfrom
kdulla:qwen3_vl_moe_kv_main

kdulla commented Jul 31, 2026 •

edited

Loading

Uh oh!

ochougul Aug 7, 2026

Uh oh!

ochougul Aug 7, 2026

Uh oh!

ochougul Aug 7, 2026

Uh oh!

ochougul Aug 7, 2026

Uh oh!

kdulla Aug 7, 2026

Uh oh!

ochougul Aug 7, 2026

Uh oh!

ochougul Aug 7, 2026

Uh oh!

ochougul Aug 7, 2026

Uh oh!

ochougul Aug 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		# routing_weights = torch.zeros_like(router_logits)
		# routing_weights.scatter_(1, top_i, top_w)

Uh oh!

Conversation

kdulla commented Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Modes Added

Scripts

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kdulla commented Jul 31, 2026 •

edited

Loading