Skip to content

Enable Dynamo KV-blocking Loop export for kv_headpar - #1314

Open
quic-abhamidi wants to merge 2 commits into
quic:mainfrom
quic-abhamidi:kv_blocking
Open

Enable Dynamo KV-blocking Loop export for kv_headpar#1314
quic-abhamidi wants to merge 2 commits into
quic:mainfrom
quic-abhamidi:kv_blocking

Conversation

@quic-abhamidi

Copy link
Copy Markdown
Contributor

Integrate a Dynamo-only torch.while_loop implementation for blocked_kv_attention_forward_headpar_offline so KV-blocked decode can export a structured ONNX Loop instead of unrolling the Python for-loop per KV block.

Key changes:

  • Enable the KV Loop path by default for blocking_mode=kv_headpar when dynamo=True.
  • Preserve the legacy Python-loop path for dynamo=False and allow explicit opt-out with use_kv_loop_op=False.
  • Add StaticLoopInputsTransform to rewrite ONNX Loop trip count and initial condition to static constants required by QAIC compiler unrolling.
  • Add InlineTorchSubgraphFunctionsTransform for dynamo subfunction export, inlining torch-generated subgraph functions that contain Loop nodes so QAIC does not see Loop control inputs inside FunctionProto bodies.
  • Update Dynamo causal LM examples and README with KV-blocking CLI support.
  • Add Dynamo/unit tests covering Loop export, config defaults, static Loop
    input rewriting, and function-body Loop handling.

Verified:

  • tests/dynamo/test_kv_blocking_loop.py: 4 passed
  • tests/unit_test/transforms/test_blocking_transform.py::TestDynamoKVLoopBlockingConfig: 3 passed
  • Qwen2-1.5B kv_headpar Dynamo compile succeeds with and without
    --enable-onnx-subfunctions.

Signed-off-by: Anusha V.S Bhamidipati <abhamidi@qti.qualcomm.com>
Signed-off-by: Anusha V.S Bhamidipati <abhamidi@qti.qualcomm.com>
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