Support initializing StaticAttentionIOManager from any module with StaticAttention inside

[sxu]

Summary:
Sometimes the ModelArgs object is not available. This diff provides an alternative way to initialize the IO manager from any module that contains StaticAttention. The assumption is that the RoPE setting of those attentions modules are all the same.

Differential Revision: D84854882

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/15206

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 6 Cancelled Jobs, 1 Unrelated Failure

As of commit 36d4b58 with merge base 651b357 ():

NEW FAILURES - The following jobs have failed:

• pull / test-samsung-models-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 6baa6800023ae586386242ffadc0072f628d3bb0b5179f4ddfadfc1fa126e1fa /exec failed with exit code 1
• Test CUDA Builds / export-voxtral-cuda-artifact / linux-job (gh)
  RuntimeError: Command docker exec -t 64953325b43eb79c4c04e241cd1f738f3cf4847aa8fb2af2135825f68efdae2c /exec failed with exit code 2

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• periodic / test-models-linux (buck2, mv3, portable, linux.2xlarge, 90) / linux-job (gh)
• periodic / test-models-linux (buck2, mv3, xnnpack-quantization-delegation, linux.2xlarge, 90) / linux-job (gh)
• periodic / test-models-linux (cmake, mv3, portable, linux.2xlarge, 90) / linux-job (gh)
• periodic / test-models-linux (cmake, mv3, xnnpack-quantization-delegation, linux.2xlarge, 90) / linux-job (gh)
• periodic / test-models-linux (cmake, vit, portable, linux.2xlarge, 90) / linux-job (gh)
• periodic / test-models-linux (cmake, vit, xnnpack-quantization-delegation, linux.2xlarge, 90) / linux-job (gh)

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / unittest-arm-backend-with-no-fvp (test_pytest_models) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-codesync]

@sxu has exported this pull request. If you are a Meta employee, you can view the originating Diff in D84854882.