Repository navigation
Add blocking support with weightfree and dynamo for CausalLM models - #1403
Conversation
|
Added the test cases for all the supported blocking modes and models on tests/weight_free/test_on_qaic.py::test_weight_free_blocked_hw_hf_parity_and_artifact_metrics |
I see that is not including the disagg mode. Add additionally for disagg mode in the above file, for verified models |
e76620d to
247fa56
Compare
Added the disaggregated tests in tests/weight_free/test_disaggregated.py::test_weight_free_decode_blocked_disaggregated_compile_and_artifact_metrics |
| assert decode_qpc.exists(), f"decode QPC path does not exist: {decode_qpc}" | ||
| assert prefill_qpc.exists(), f"prefill QPC path does not exist: {prefill_qpc}" | ||
| assert decode_onnx_path != prefill_onnx_path | ||
|
|
There was a problem hiding this comment.
Check additionally the token match also, with reference
| # Initialize Running Maximum and Denominator | ||
| batch_size, num_heads, seq_len, _ = query.shape | ||
| current_max = torch.full( | ||
| (batch_size, num_heads, seq_len), |
There was a problem hiding this comment.
[AI assisted review] running-softmax accumulators now use query.dtype instead of FP32 across shared KV/QKV/HQKV/BHQKV paths. FP16 denominatorscan overflow for long, near-uniform contexts, changing legacy blocked-attention results. Keep maxima and denominators in FP32 and cast only the resulting output.
There was a problem hiding this comment.
Good catch. Legacy kept running-softmax current_max and current_denominator in default FP32. This PR changed them to query.dtype, which can reduce numerical stability for long contexts. We should revert the accumulators to FP32 and only keep the output buffer in query.dtype. So, reverted it and kept it aligned with legacy.
| return torch.empty_like(data) | ||
|
|
||
|
|
||
| # BATCH-FOLDED SCATTER/GATHER |
There was a problem hiding this comment.
Batch folding is enabled for which all models in this PR?
There was a problem hiding this comment.
Batch folding is only used when blocking_mode=kv_batch_fold. In this PR, the explicit model plumbing is for Qwen3.5 / Qwen3.5-MoE / Qwen3-VL-MoE decode with continuous batching. The custom op definitions are generic infrastructure, but they are only exercised by models routed through BlockingMode.KV_BATCH_FOLD, not by all models. It is for now only tested with dynamo and are yet to be tested for wf.
247fa56 to
851299e
Compare
|
CI-Ready |
6da798e to
2ac494f
Compare
| batch_idx = batch_index.view(-1, 1) | ||
| ctx_idx = position_ids | ||
| data[batch_idx, ctx_idx] = updates | ||
| batch_idx = batch_index.view(-1, 1).expand_as(position_ids) |
There was a problem hiding this comment.
[blocker] (Reviewed by me, assisted by AI):
With routed continuous-batching decode, batch_index is [B, 1] while position_ids remains [1, T]. Expanding the former against the latter raises before any scatter when B > 1; the required unit job reproduces this in the Qwen3.5-MoE decode-bsz2, decode-bsz4, and partial-batch quickchecks. Expand batch_index to its own batch dimension and the position width (for example, expand(-1, position_ids.shape[-1])) before applying the valid-position mask.
There was a problem hiding this comment.
Fixed. I updated the 3D CB scatter path to expand batch_index over its own batch dimension and the position width instead of expand_as(position_ids):
batch_idx = batch_index.view(-1, 1).expand(-1, position_ids.shape[-1])
ctx_idx = position_ids.expand(batch_idx.shape[0], -1)
This preserves the legacy advanced-index broadcasting behavior while keeping the new valid-position mask. I also mirrored the same shape logic in the Dynamo custom op implementation so both eager and Dynamo paths stay aligned.
I verified the failing shape directly with batch_index [4, 1] and position_ids [1, 8192]; it no longer raises.
b8922e4 to
7422814
Compare
7422814 to
2f19cc1
Compare
| f"mode={decode_blocking_key}: HF={hf_tokens_for_compare.tolist()}, QAIC={qaic_tokens.tolist()}" | ||
| ) | ||
|
|
||
| metrics = { |
There was a problem hiding this comment.
Can remove these logs from test file. line 611 to 630
| return str(compile_dir) | ||
|
|
||
|
|
||
| def _directory_size_bytes(path: Path) -> int: |
There was a problem hiding this comment.
remove size logs from test
| @pytest.mark.dynamo | ||
| @pytest.mark.on_qaic | ||
| @pytest.mark.llm_model | ||
| @pytest.mark.skip(reason="Disabled while testing -- fp32-weights compile still converts to fp16 on-device anyway.") |
There was a problem hiding this comment.
Why is this skipped? If skipped, do we need this test?
| return f"{case.model_label}-{mode}{suffix}" | ||
|
|
||
|
|
||
| GPT_OSS_RUNTIME_ERROR_XFAIL_CASE_IDS = { |
There was a problem hiding this comment.
Why these xfails are still required for gpt_oss? Isn't blocking tests passing successfully? Remove this xfail or else the failures will be silent.
2f19cc1 to
07aea79
Compare
07aea79 to
ace9c32
Compare
| "gemma", | ||
| "gemma2", | ||
| "glm4_moe", | ||
| "gpt_oss", |
There was a problem hiding this comment.
[blocker] (Reviewed by me, assisted by AI):
Could you help clarify why gpt_oss remains in the supported blocking matrix while the export test explicitly marks every GPT-OSS case as expected to fail? Should it be excluded until the blocking path works and has parity coverage?
There was a problem hiding this comment.
This list is the positive weight-free blocking QAIC parity matrix, not an xfail/unsupported list. GPT-OSS is included because test_weight_free_blocked_hw_hf_parity_and_artifact_metrics passes for all GPT-OSS blocking modes in this matrix. The GPT-OSS xfail is in the separate export-structure test, where forward() disables blocking for sliding_window, so that xfail should not be used to infer that the weight-free QAIC blocking path is unsupported.
| return f"{case.model_label}-{mode}{suffix}" | ||
|
|
||
|
|
||
| GPT_OSS_BLOCKING_XFAIL_CASE_IDS = { |
There was a problem hiding this comment.
Why is still this xfail case ids present for gpt_oss? Is there anything breaking? If not, remove xfail
There was a problem hiding this comment.
Yes, "gpt_oss-hq-mdp": "GPT-OSS blocked HQ MDP compile/runtime path is still failing",
"gpt_oss-hkv-mdp": "GPT-OSS blocked HKV MDP compile/runtime path is still failing",
"gpt_oss-hqkv-mdp": "GPT-OSS blocked HQKV MDP compile/runtime path is still failing",
"gpt_oss-bhqkv-mdp": "GPT-OSS blocked BHQKV MDP compile/runtime path is still failing", is breaking
There was a problem hiding this comment.
So when mdp is enabled along with this WF blocking changes, it is breaking? This needs to be fixed. Mention that in limitation/TODO of the PR description.
There was a problem hiding this comment.
"gpt_oss-qkv" is also added? Is it failing without mdp?
There was a problem hiding this comment.
I am seeing SDK based failures here. It is passing in 1.23.0.61 SDK, all the cases in xfails are passing with this SDK but failing on latest SDK. So, just to be safe, I have kept all these cases on xfail for now. The ones which are passing in all SDKs are now enabled.
This is with 1.23.0.61 SDK:
============================= test session starts ==============================
platform linux -- Python 3.12.3, pytest-9.1.1, pluggy-1.6.0 -- /home/abhamidi/dynamo_w_weightfree/blocking_env/bin/python
cachedir: .pytest_cache
rootdir: /home/abhamidi/blocking_w_weightfree/efficient-transformers
configfile: pyproject.toml
plugins: anyio-4.15.1
collecting ... collected 1 item
tests/dynamo/test_blocking_forwards.py::test_dynamo_tiny_blocking_compile_and_generate[gpt_oss-qkv]
Writing model shards: 0%| | 0/1 [00:00<?, ?it/s]
Writing model shards: 100%|██████████| 1/1 [00:00<00:00, 1.52it/s]
Writing model shards: 100%|██████████| 1/1 [00:00<00:00, 1.52it/s]
Loading weights: 0%| | 0/37 [00:00<?, ?it/s]
Loading weights: 100%|██████████| 37/37 [00:00<00:00, 8221.08it/s]
[torch.onnx] Obtain model graph for QEffGptOssForCausalLM([...] with torch.export.export(..., strict=False)...
[torch.onnx] Obtain model graph for QEffGptOssForCausalLM([...] with torch.export.export(..., strict=False)... ✅
[torch.onnx] Run decompositions...
[torch.onnx] Run decompositions... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
PASSED
.............Cleaned up ../qeff_dir
------------------------------ QEff Test Summary -------------------------------
group | total | passed | failed | skipped | xfailed | xpassed | error
----------+-------+--------+--------+---------+---------+---------+------
llm_model | 1 | 1 | 0 | 0 | 0 | 0 | 0
============================== 1 passed in 43.89s ==============================
Loading "/opt/qti-aic/dev/lib/x86_64/libQAic.so"
Loading "/opt/qti-aic/dev/lib/x86_64/libQAic.so"
Log file: /root/.cache/qefficient_logs/QEFF_20261022_070530_48820.log
This is with the latest SDK:
INFO QEFF_LOGGER:modeling_qeff.py:801 PyTorch export successful
WARNING QEfficient.base.onnx_transforms:onnx_transforms.py:547 RenameRepeatedSubgraphTransform: more repeated subgraph functions (2) than target class names (1). Function 'repeated_subgraph1' (index 1) will be assigned the last available name 'QEffGptOssDecoderLayer' with a numeric suffix — verify get_submodules_for_export() returns all repeated block classes for this model.
INFO QEFF_LOGGER:modeling_qeff.py:842 ONNX transforms applied
INFO QEFF_LOGGER:modeling_qeff.py:849 Transformed ONNX saved
INFO QEFF_LOGGER:export_utils.py:784 Hashed parameters exported successfully.
INFO QEFF_LOGGER:export_utils.py:532 ONNX export completed.
INFO QEFF_LOGGER:modeling_qeff.py:1449 Using ONNX subfunctions for compilation.
INFO QEFF_LOGGER:logging_utils.py:168 Captured API arguments.
INFO QEFF_LOGGER:modeling_qeff.py:1596 Running compiler: /opt/qti-aic/exec/qaic-compile -aic-hw -aic-hw-version=ai100 -m=/home/abhamidi/.cache/qeff_models/GptOssForCausalLM/GptOssForCausalLM-45b31488aac5977c/GptOssForCausalLM.onnx -retained-state -convert-to-fp16 -aic-num-cores=16 -user-tiled -sub-functions -network-specialization-config=/home/tmp/pytest-of-abhamidi/pytest-36/dynamo_blocking_gpt_oss-qkv_decode0/qeff_dynamo_exports/gpt_oss-qkv_decode/qpc-d214560081f8a760/specializations.json -custom-IO-list-file=/home/tmp/pytest-of-abhamidi/pytest-36/dynamo_blocking_gpt_oss-qkv_decode0/qeff_dynamo_exports/gpt_oss-qkv_decode/qpc-d214560081f8a760/custom_io.yaml -aic-binary-dir=/home/tmp/pytest-of-abhamidi/pytest-36/dynamo_blocking_gpt_oss-qkv_decode0/qeff_dynamo_exports/gpt_oss-qkv_decode/qpc-d214560081f8a760/qpc
INFO QEFF_LOGGER:logging_utils.py:494 compile failed.
--------------------------------------------------------------------------------- QEff Test Summary ----------------------------------------------------------------------------------
group | total | passed | failed | skipped | xfailed | xpassed | error
----------+-------+--------+--------+---------+---------+---------+------
llm_model | 1 | 0 | 1 | 0 | 0 | 0 | 0
-------------------------------------------------------------------------------------- Failures --------------------------------------------------------------------------------------
- tests/dynamo/test_blocking_forwards.py::test_dynamo_tiny_blocking_compile_and_generate[gpt_oss-qkv]
============================================================================== short test summary info ===============================================================================
FAILED tests/dynamo/test_blocking_forwards.py::test_dynamo_tiny_blocking_compile_and_generate[gpt_oss-qkv] - RuntimeError: Compilation failed!
There was a problem hiding this comment.
So what about the same config in legacy(w/o dynamo)? blocking gpt_oss with qkv. Check if any issues in latest sdk. If not, this needs to be fixed.
There was a problem hiding this comment.
Sure, will check, if it is failing, will create a compiler JIRA
There was a problem hiding this comment.
tests/transformers/models/causal_lm_models/test_causal_lm_blocking_hqkv.py I have tested legacy qkv for gpt-oss.
It is passing. Used latest sdk - 1.24.0.13
There was a problem hiding this comment.
So this needs to be fixed. Remove qkv from supported modes for gpt-oss. Add TODO in the PR description(and fix needs to go in another PR soon). What about other models, is qkv mode verified?
There was a problem hiding this comment.
Removing xfail as this is compiling with 1.24.0.13 and 1.24.0.8 SDKs:
(blocking_env) root@bfcecf373102:/home/abhamidi/blocking_w_weightfree/efficient-transformers# pytest -s --runxfail tests/dynamo/test_blocking_forwards.py::test_dynamo_tiny_blocking_c
ompile_and_generate[gpt_oss-qkv]
Loading "/opt/qti-aic/dev/lib/x86_64/libQAic.so"
================================================================================ test session starts =================================================================================
platform linux -- Python 3.12.3, pytest-9.1.1, pluggy-1.6.0 -- /home/abhamidi/dynamo_w_weightfree/blocking_env/bin/python
cachedir: .pytest_cache
rootdir: /home/abhamidi/blocking_w_weightfree/efficient-transformers
configfile: pyproject.toml
plugins: anyio-4.15.1
collected 1 item
Writing model shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 25.05it/s]
Loading weights: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 37/37 [00:00<00:00, 9884.67it/s]
[torch.onnx] Obtain model graph for QEffGptOssForCausalLM([...] with torch.export.export(..., strict=False)...
[torch.onnx] Obtain model graph for QEffGptOssForCausalLM([...] with torch.export.export(..., strict=False)... ✅
[torch.onnx] Run decompositions...
[torch.onnx] Run decompositions... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
Loading "/opt/qti-aic/dev/lib/x86_64/libQAic.so"
PASSED
.............Cleaned up /root/.cache/qeff_models
--------------------------------------------------------------------------------- QEff Test Summary ----------------------------------------------------------------------------------
group | total | passed | failed | skipped | xfailed | xpassed | error
----------+-------+--------+--------+---------+---------+---------+------
llm_model | 1 | 1 | 0 | 0 | 0 | 0 | 0
================================================================================= 1 passed in 39.03s =================================================================================
Log file: /root/.cache/qefficient_logs/QEFF_20261022_133357_7047.log
ace9c32 to
7d409ae
Compare
| expected_tokens = expected_logits.argmax(-1) | ||
| actual_tokens = actual_logits.argmax(-1) | ||
| assert np.array_equal(expected_tokens, actual_tokens), ( | ||
| f"{label} top-1 tokens diverged: expected={expected_tokens.tolist()}, actual={actual_tokens.tolist()}" |
There was a problem hiding this comment.
How many tokens are matched in test case
There was a problem hiding this comment.
I have kept it to 1 token only, since in test_on_qaic_causallm.py 10 tokens are being matched for the same set of models.
7d409ae to
4138585
Compare
feed14e to
76226b2
Compare
e63fae9 to
7af4379
Compare
Signed-off-by: Anusha V.S Bhamidipati <abhamidi@qti.qualcomm.com>
Signed-off-by: Anusha V.S Bhamidipati <abhamidi@qti.qualcomm.com>
034f6c4 to
859cd2b
Compare
3736940
into
qualcomm:main
Enable blocking support for various Causal LM models
The blocked weight-free QAIC matrix covers:
gemma,gemma2,glm4_moe,gpt_oss,granite,granitemoe,llama,mistral,mixtral,qwen2,qwen3,qwen3_moe, andstarcoder2for different blocking modes.Supported blocking modes:
Also adds fixes for weight-free/Dynamo export/runtime issues found during validation:
Validation and Testing: