# Support of Dynamic and Static Shape in One Module

**URL:** <https://discuss.pytorch.org/t/support-of-dynamic-and-static-shape-in-one-module/208561>\
**Category:** ExecuTorch\
**Created:** [August 24, 2024, 7:24am UTC](https://discuss.pytorch.org/t/support-of-dynamic-and-static-shape-in-one-module/208561 "2024-08-24T07:24:16Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![chenghua\_wang](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/chenghua_wang/32/55137_2.png) [@chenghua\_wang](https://discuss.pytorch.org/u/chenghua_wang)\
**Post date:** [August 24, 2024, 7:24am UTC](https://discuss.pytorch.org/t/support-of-dynamic-and-static-shape-in-one-module/208561/1 "2024-08-24T07:24:16Z")

</div>

I’m trying to use ExecuTorch to deploy a LLM with a QNN backend. Unfortunately, QNN doesn’t support dynamic shapes, so some trick is needed to utilise up the NPU. A simple idea is shown below:

 ![image](https://discuss.pytorch.org/uploads/default/original/3X/2/7/279bb8538481273f926bbe97d9f7b5e05918a696.jpeg)

Divide the Prompt into fixed-size Chunks, and these fixed Chunks first enter the Static Shape QKV Proj (run using QNN-NPU), then proceed to the Dynamic Shape Attention (run using CPU). Because of the chunking, the computation of adjacent Tokens can be pipelined.

I am not familiar with ExecuTorch. In the examples of Dynamic Shape provided by ExecuTorch, it seems that dynamic shapes are specified within the `export` function. such as:

```python
aten_dialect: ExportedProgram = export(
    Basic(), example_args, dynamic_shapes=dynamic_shapes
)

```

However, it seems that I can only specify dynamic shapes for the Top-Level Module, but not for the nested Modules within it. such as:

```python
class MatMul(nn.Module):
    def __init__ (self):
        super(). __init__ ()

    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        return a @ b

class Basic(nn.Module):
    def __init__ (self):
        super(). __init__ ()
        self.MatMul = MatMul()

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        out1 = self.MatMul(x, y) # 3x6x3
        out2 = self.MatMul(out1, x) # 3x3x6
        return out2

example_args = (torch.randn(3, 6), torch.randn(6, 3))
dim1_x = Dim("dim1_x", min=1, max=10)
dynamic_shapes = {"a": {0: dim1_x, 1: dim1_x}, "b": {0: dim1_x, 1: dim1_x}}
aten_dialect: ExportedProgram = export(
    Basic(), example_args, dynamic_shapes=dynamic_shapes
)

```

A feasible approach is to compile and call each Module separately; for instance, compile the QKV-Proj Module/O-Proj Module with a Static Shape and the Attention Module with a Dynamic Shape. However, this method is quite cumbersome because it requires disassembling and compiling each Transformer Block within the Large Language Model (LLM) individually.

- So my first question is, how can I make the nested Modules support dynamic shapes?

After the model is compiled, ExecuTorch can dispatch different subgraphs to run on various hardware (as I mentioned, dispatch QKV/O Proj to the NPU and Attention to the CPU/GPU), which is very convenient. But how can I achieve the pipeline form I mentioned above? For example, once the QKV Proj computation for the first token is completed, the computation kernel for the QKV Proj of the second token can be launched.

- My second question is, can I call a compiled Module individually, rather than starting from the Top-Level Module? In this way, I can lock different Modules to achieve a pipeline.

The Dynamic KV Caches seems also unsolved? 🤔 see [Support for dynamic caches · Issue #4740 · pytorch/executorch · GitHub](https://github.com/pytorch/executorch/issues/4740)

---

<div class="post-metadata">

**Author:** ![jackzhxng](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jackzhxng/32/72367_2.png) [@jackzhxng](https://discuss.pytorch.org/u/jackzhxng)\
**Post date:** [August 26, 2024, 11:26pm UTC](https://discuss.pytorch.org/t/support-of-dynamic-and-static-shape-in-one-module/208561/2 "2024-08-26T23:26:12Z")

</div>

Hi @chenghua_wang,

Thanks for providing this diagram! Firstly, I would like to offer an update for our progress on dynamic kv caches [here](https://github.com/pytorch/executorch/issues/4740).

Next - do you need to have a dynamic cache? If not then we can keep the entire model static, in which case you can go back to using a QNN backend. [Here](https://pytorch.org/executorch/main/build-run-qualcomm-ai-engine-direct-backend.html) is an example.

If you do need a dynamic cache and thus dynamic shapes, then the entire model would need to support dynamic shape, including your Q/K/V and O proj layers.

---

<div class="post-metadata">

**Author:** ![jackzhxng](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jackzhxng/32/72367_2.png) [@jackzhxng](https://discuss.pytorch.org/u/jackzhxng)\
**Post date:** [August 26, 2024, 11:27pm UTC](https://discuss.pytorch.org/t/support-of-dynamic-and-static-shape-in-one-module/208561/3 "2024-08-26T23:27:08Z")

</div>

In case you were interested though in calling different modules within your program, [this](https://github.com/pytorch/executorch/blob/main/examples/models/llava/export_llava.py#L192-L208) is how you would do it. Note that you can only export the `forward()` method of a module, so you will need to wrap your individual submodules in a `nn.module` with a `forward()` method.

---

<div class="post-metadata">

**Author:** ![chenghua\_wang](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/chenghua_wang/32/55137_2.png) [@chenghua\_wang](https://discuss.pytorch.org/u/chenghua_wang)\
**Post date:** [August 27, 2024, 2:40am UTC](https://discuss.pytorch.org/t/support-of-dynamic-and-static-shape-in-one-module/208561/4 "2024-08-27T02:40:38Z")

</div>

Hi @jackzhxng,

Thank you for your helpful reply.

> [@jackzhxng](#):
>
> individual submodules in a `nn.module` with a `forward()` method.

This example has solved my problems, I will give it a try.

> [@jackzhxng](#):
>
> If you do need a dynamic cache and thus dynamic shapes, then the entire model would need to support dynamic shape, including your Q/K/V and O proj layers.

Yes, I need dynamic cache, but during the pre-fill stage, the QKV/O Proj Layers is static due to fixed chunk size.

* * *

Now I have another question. Because I need to call multiple compiled modules, for example:

```python
out_1 = module_1.forward(...)
out_2 = module_2.forward(out_1, ...)

```

- The memory of the tensor `out_1` is managed by `module_2` , or is it managed by the user?

It should be managed by the user? Because during execution, `out_1` does not know how many successor nodes it has.

---

<div class="post-metadata">

**Author:** ![cccclai](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/cccclai/32/35398_2.png) [@cccclai](https://discuss.pytorch.org/u/cccclai)\
**Post date:** [August 27, 2024, 4:39am UTC](https://discuss.pytorch.org/t/support-of-dynamic-and-static-shape-in-one-module/208561/5 "2024-08-27T04:39:10Z")

</div>

Calling part of the graph in cpu and part of the graph in npu might cause slower latency. There are two ways to use QNN. One way is to prefill token by token, and you can checkout the instructions here [executorch/examples/models/llama2 at main · pytorch/executorch · GitHub](https://github.com/pytorch/executorch/tree/main/examples/models/llama2#optional-smaller-models-delegated-to-other-backends)

For QNN, the command line for running floating point is

```auto
python -m examples.models.llama2.export_llama --disable_dynamic_shape -kv --qnn -c stories110M.pt -p params.json

```

and for 4bit quantized, it’s still wip. The corresponding command line is

```auto
python -m examples.models.llama2.export_llama --disable_dynamic_shape -kv --qnn --pt2e_quantize qnn_16a4w  

```

---

<div class="post-metadata">

**Author:** ![chenghua\_wang](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/chenghua_wang/32/55137_2.png) [@chenghua\_wang](https://discuss.pytorch.org/u/chenghua_wang)\
**Post date:** [August 27, 2024, 7:33am UTC](https://discuss.pytorch.org/t/support-of-dynamic-and-static-shape-in-one-module/208561/6 "2024-08-27T07:33:43Z")

</div>

Thank you for your answer, this is a very good example, I will try it later:

> [@cccclai](#):
>
> checkout the instructions here [executorch/examples/models/llama2 at main · pytorch/executorch · GitHub](https://github.com/pytorch/executorch/tree/main/examples/models/llama2#optional-smaller-models-delegated-to-other-backends)

The original intention of dividing the prompt into chunks and executing them separately on the CPU and NPU is to utilize dynamic cache. Indeed, the NPU is much faster, after using the Pipeline, it may not be possible to completely overlap the CPU’s computation time.

It seems that QNN2.21.0+ support dynamic shape? see document: [Qualcomm Documentation](https://docs.qualcomm.com/bundle/publicresource/topics/80-63442-50/api_usage_guidelines.html#dynamic-dimensions)

---

<div class="post-metadata">

**Author:** ![gianfranco](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/g/4af34b/32.png) [@gianfranco](https://discuss.pytorch.org/u/gianfranco)\
**Post date:** [June 9, 2025, 9:06am UTC](https://discuss.pytorch.org/t/support-of-dynamic-and-static-shape-in-one-module/208561/7 "2025-06-09T09:06:12Z")

</div>

Hey Guys, I am still having problem with dynamic shape on mps backend. Can someone say how to export llama3\_2 .pte file on mps and run it with a generic prompt?

---

<div class="post-metadata">

**Author:** ![cccclai](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/cccclai/32/35398_2.png) [@cccclai](https://discuss.pytorch.org/u/cccclai)\
**Post date:** [December 16, 2025, 12:35am UTC](https://discuss.pytorch.org/t/support-of-dynamic-and-static-shape-in-one-module/208561/8 "2025-12-16T00:35:48Z")

</div>

Dynamic shape support is only added for limited ops for now. Probably not enough for llm use cases

---

<div class="post-metadata">

**Author:** ![cccclai](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/cccclai/32/35398_2.png) [@cccclai](https://discuss.pytorch.org/u/cccclai)\
**Post date:** [December 16, 2025, 12:36am UTC](https://discuss.pytorch.org/t/support-of-dynamic-and-static-shape-in-one-module/208561/9 "2025-12-16T00:36:35Z")

</div>

Is there any specific reason you want to try mps backend? Is coreml backend interesting to you? [executorch/examples/apple/coreml/llama at main · pytorch/executorch · GitHub](https://github.com/pytorch/executorch/tree/main/examples/apple/coreml/llama)