Fold DyT alpha scalar into tanh LUT (#21952)

[cgreenberg]

Summary:

Dynamic Tanh (DyT) normalization computes tanh(alpha * x) with a learned scalar alpha. Quantized, that lowers to a full-tensor integer Mul followed by a tanh TABLE. This adds FoldDyTAlphaIntoLUTPass, which replays the exact TOSA SINGLE_ROUND integer arithmetic over all 256 int8 codes and folds the alpha Mul into the TABLE without approximation.

The matcher fails closed for unsupported shapes, ranges, layouts, or quantization metadata. This revision contains only the public pass implementation, tests, and package exports; pass registration is intentionally outside this public change.

Reviewed By: rascani

Differential Revision: D116560649

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit bd9715a with merge base fbd4bbf ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: cgreenberg / name: Charles Greenberg (0cec5b5)

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[cgreenberg]

@pytorchbot label "release notes: arm"

[cgreenberg]

/easycla

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 12m 7s —— View job

---

Code Review: Fold DyT alpha scalar into tanh LUT

• Read changed files
• Analyze the pass implementation
• Analyze pass manager hook
• Analyze tests
• Post review

Reviewed all four changed files plus the surrounding machinery (op_tosa_rescale._compute_multiplier_and_shift, insert_table_ops, real_impl, arm_pass). This is a carefully-built pass with unusually strong test coverage. Overall it looks correct and well-guarded — the concerns below are minor.

---

Correctness — looks solid

• SINGLE_ROUND replay is faithful._apply_tosa_rescale (fold_dyt_alpha_into_lut_pass.py:57) mirrors the serialized RESCALE path: same _compute_multiplier_and_shift, rounding = 1 << (shift-1), (centered*multiplier + rounding) >> shift, then clamp to output dtype. The comment at :65 correctly notes _compute_multiplier_and_shift asserts shift ∈ [2, 62], so shift-1 is never negative. Nice that you resisted "guarding" it.
• The bit-exactness is genuinely verified, not asserted.test_tosa_output_is_bit_exact_after_fold compares the generated TABLE against the actual tosa_reference_model output (_tosa_reference_outputs) over every int8 code. That's the right way to pin the int32 MUL-wrap and rescale rounding, and it removes any doubt about the hand-replayed arithmetic matching the device.
• Fails-closed guards are comprehensive: non-scalar alpha (:88, :236-238), activation-side rank views (:243), int32/int8 dtype checks, and output_rescale.output_zp != tanh_input_qargs zp (:267). test_pass_rejects_activation_side_rank_view and the DyTModule docstring on the conv-producer requirement show the edges were thought through.
• TABLE wiring is consistent: input is the original int8 match.activation, input_qparams = activation_qargs, output_qparams = tanh_output_qargs — the LUT collapses rescale→mul→rescale→tanh into one int8→int8 map. Correct.

Minor concerns / questions

1. Dangling alpha buffer after the fold.call() erases the tanh and relies on eliminate_dead_code() (:323) to drop the mul and rescales — but eliminate_dead_code does not remove placeholder/buffer nodes. Once the alpha Mul is gone, the alpha buffer placeholder has no users. Does a later pipeline pass (FuseEqualPlaceholdersPass, constant-folding) actually strip it, or does the exported program keep an orphaned buffer in its state_dict + graph signature? It's functionally harmless and tiny, but the PR summary says the pass "deletes the Mul and its surrounding rescales," and the alpha constant silently surviving is worth confirming. None of the tests assert the buffer is removed. Fix this →

2. TABLE node meta["val"] is copied from the tanh, not recomputed. Unlike InsertTableOpsPass, which retraces via super().call() to refresh fake-tensor types, this pass inserts the TABLE via create_node(from_node=match.tanh, ...) and does not retrace. That copies the tanh's val. In a fully quantized graph the tanh output should already be int8, so shape/dtype line up — but it's worth a sanity check that no downstream pass expects a val produced by the TABLE fake kernel specifically (it validates table shape/dtype). If the copied val dtype ever isn't int8, serialization could be surprised.

3. _source_qargs scale match uses math.isclose(rel_tol=1e-6) (:180) to pick the activation qargs among candidates. Fine in practice, but if two producers ever carry near-identical scales it silently selects the first. Given how much else in this pass "fails closed," consider whether an ambiguous match (>1 candidate within tolerance) should also bail rather than take the first.

Nits

• JakeStevens' copyright comment appears addressed — the file carries Copyright (c) Meta Platforms, Inc. and affiliates. Note the rest of backends/arm/_passes/ also carries an Arm Limited line (e.g. arm_pass_manager.py, arm_pass.py); confirm with the Arm reviewers whether the dual header is expected here, since this lands in their backend.
• The circular-import workaround (module-scope vs. in-function import of register_pass_factories_before, :331-336) is well-documented in both the code and PR body. Good.

CI

The 2 reported failures (hifi-build, vision-build) are pull_request_target fork-checkout refusals — infrastructure/security, unrelated to this change.

---

Nice work — the exhaustive parity test against the reference model is exactly the right safety net for a byte-exact rewrite like this. My only substantive ask is confirming the alpha buffer doesn't linger (#1).
· branch export-D116560649

[cgreenberg]

@claude review this PR

[rascani]

Review automatically exported from Phabricator review in Meta.

[cgreenberg]

Summary:

Dynamic Tanh (DyT) normalization computes tanh(alpha * x) with a learned scalar alpha. Quantized, that lowers to a full-tensor integer Mul followed by a tanh TABLE. This adds FoldDyTAlphaIntoLUTPass, which replays the exact TOSA SINGLE_ROUND integer arithmetic over all 256 int8 codes and folds the alpha Mul into the TABLE without approximation.

The matcher fails closed for unsupported shapes, ranges, layouts, or quantization metadata. This revision contains only the public pass implementation, tests, and package exports; pass registration is intentionally outside this public change.

Reviewed By: rascani

Differential Revision: D116560649

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit bd9715a with merge base fbd4bbf ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: cgreenberg / name: Charles Greenberg (0cec5b5)

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[cgreenberg]

@pytorchbot label "release notes: arm"

[cgreenberg]

/easycla

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 12m 7s —— View job

---

Code Review: Fold DyT alpha scalar into tanh LUT

• Read changed files
• Analyze the pass implementation
• Analyze pass manager hook
• Analyze tests
• Post review

Reviewed all four changed files plus the surrounding machinery (op_tosa_rescale._compute_multiplier_and_shift, insert_table_ops, real_impl, arm_pass). This is a carefully-built pass with unusually strong test coverage. Overall it looks correct and well-guarded — the concerns below are minor.

---

Correctness — looks solid

• SINGLE_ROUND replay is faithful._apply_tosa_rescale (fold_dyt_alpha_into_lut_pass.py:57) mirrors the serialized RESCALE path: same _compute_multiplier_and_shift, rounding = 1 << (shift-1), (centered*multiplier + rounding) >> shift, then clamp to output dtype. The comment at :65 correctly notes _compute_multiplier_and_shift asserts shift ∈ [2, 62], so shift-1 is never negative. Nice that you resisted "guarding" it.
• The bit-exactness is genuinely verified, not asserted.test_tosa_output_is_bit_exact_after_fold compares the generated TABLE against the actual tosa_reference_model output (_tosa_reference_outputs) over every int8 code. That's the right way to pin the int32 MUL-wrap and rescale rounding, and it removes any doubt about the hand-replayed arithmetic matching the device.
• Fails-closed guards are comprehensive: non-scalar alpha (:88, :236-238), activation-side rank views (:243), int32/int8 dtype checks, and output_rescale.output_zp != tanh_input_qargs zp (:267). test_pass_rejects_activation_side_rank_view and the DyTModule docstring on the conv-producer requirement show the edges were thought through.
• TABLE wiring is consistent: input is the original int8 match.activation, input_qparams = activation_qargs, output_qparams = tanh_output_qargs — the LUT collapses rescale→mul→rescale→tanh into one int8→int8 map. Correct.

Minor concerns / questions

1. Dangling alpha buffer after the fold.call() erases the tanh and relies on eliminate_dead_code() (:323) to drop the mul and rescales — but eliminate_dead_code does not remove placeholder/buffer nodes. Once the alpha Mul is gone, the alpha buffer placeholder has no users. Does a later pipeline pass (FuseEqualPlaceholdersPass, constant-folding) actually strip it, or does the exported program keep an orphaned buffer in its state_dict + graph signature? It's functionally harmless and tiny, but the PR summary says the pass "deletes the Mul and its surrounding rescales," and the alpha constant silently surviving is worth confirming. None of the tests assert the buffer is removed. Fix this →

2. TABLE node meta["val"] is copied from the tanh, not recomputed. Unlike InsertTableOpsPass, which retraces via super().call() to refresh fake-tensor types, this pass inserts the TABLE via create_node(from_node=match.tanh, ...) and does not retrace. That copies the tanh's val. In a fully quantized graph the tanh output should already be int8, so shape/dtype line up — but it's worth a sanity check that no downstream pass expects a val produced by the TABLE fake kernel specifically (it validates table shape/dtype). If the copied val dtype ever isn't int8, serialization could be surprised.

3. _source_qargs scale match uses math.isclose(rel_tol=1e-6) (:180) to pick the activation qargs among candidates. Fine in practice, but if two producers ever carry near-identical scales it silently selects the first. Given how much else in this pass "fails closed," consider whether an ambiguous match (>1 candidate within tolerance) should also bail rather than take the first.

Nits

• JakeStevens' copyright comment appears addressed — the file carries Copyright (c) Meta Platforms, Inc. and affiliates. Note the rest of backends/arm/_passes/ also carries an Arm Limited line (e.g. arm_pass_manager.py, arm_pass.py); confirm with the Arm reviewers whether the dual header is expected here, since this lands in their backend.
• The circular-import workaround (module-scope vs. in-function import of register_pass_factories_before, :331-336) is well-documented in both the code and PR body. Good.

CI

The 2 reported failures (hifi-build, vision-build) are pull_request_target fork-checkout refusals — infrastructure/security, unrelated to this change.

---

Nice work — the exhaustive parity test against the reference model is exactly the right safety net for a byte-exact rewrite like this. My only substantive ask is confirming the alpha buffer doesn't linger (#1).
· branch export-D116560649

[cgreenberg]

@claude review this PR

[rascani]

Review automatically exported from Phabricator review in Meta.

[cgreenberg]

Summary:

Dynamic Tanh (DyT) normalization computes tanh(alpha * x) with a learned scalar alpha. Quantized, that lowers to a full-tensor integer Mul followed by a tanh TABLE. This adds FoldDyTAlphaIntoLUTPass, which replays the exact TOSA SINGLE_ROUND integer arithmetic over all 256 int8 codes and folds the alpha Mul into the TABLE without approximation.

The matcher fails closed for unsupported shapes, ranges, layouts, or quantization metadata. This revision contains only the public pass implementation, tests, and package exports; pass registration is intentionally outside this public change.

Reviewed By: rascani

Differential Revision: D116560649

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit bd9715a with merge base fbd4bbf ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: cgreenberg / name: Charles Greenberg (0cec5b5)

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[cgreenberg]

@pytorchbot label "release notes: arm"

[cgreenberg]

/easycla

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 12m 7s —— View job

---

Code Review: Fold DyT alpha scalar into tanh LUT

• Read changed files
• Analyze the pass implementation
• Analyze pass manager hook
• Analyze tests
• Post review

Reviewed all four changed files plus the surrounding machinery (op_tosa_rescale._compute_multiplier_and_shift, insert_table_ops, real_impl, arm_pass). This is a carefully-built pass with unusually strong test coverage. Overall it looks correct and well-guarded — the concerns below are minor.

---

Correctness — looks solid

• SINGLE_ROUND replay is faithful._apply_tosa_rescale (fold_dyt_alpha_into_lut_pass.py:57) mirrors the serialized RESCALE path: same _compute_multiplier_and_shift, rounding = 1 << (shift-1), (centered*multiplier + rounding) >> shift, then clamp to output dtype. The comment at :65 correctly notes _compute_multiplier_and_shift asserts shift ∈ [2, 62], so shift-1 is never negative. Nice that you resisted "guarding" it.
• The bit-exactness is genuinely verified, not asserted.test_tosa_output_is_bit_exact_after_fold compares the generated TABLE against the actual tosa_reference_model output (_tosa_reference_outputs) over every int8 code. That's the right way to pin the int32 MUL-wrap and rescale rounding, and it removes any doubt about the hand-replayed arithmetic matching the device.
• Fails-closed guards are comprehensive: non-scalar alpha (:88, :236-238), activation-side rank views (:243), int32/int8 dtype checks, and output_rescale.output_zp != tanh_input_qargs zp (:267). test_pass_rejects_activation_side_rank_view and the DyTModule docstring on the conv-producer requirement show the edges were thought through.
• TABLE wiring is consistent: input is the original int8 match.activation, input_qparams = activation_qargs, output_qparams = tanh_output_qargs — the LUT collapses rescale→mul→rescale→tanh into one int8→int8 map. Correct.

Minor concerns / questions

1. Dangling alpha buffer after the fold.call() erases the tanh and relies on eliminate_dead_code() (:323) to drop the mul and rescales — but eliminate_dead_code does not remove placeholder/buffer nodes. Once the alpha Mul is gone, the alpha buffer placeholder has no users. Does a later pipeline pass (FuseEqualPlaceholdersPass, constant-folding) actually strip it, or does the exported program keep an orphaned buffer in its state_dict + graph signature? It's functionally harmless and tiny, but the PR summary says the pass "deletes the Mul and its surrounding rescales," and the alpha constant silently surviving is worth confirming. None of the tests assert the buffer is removed. Fix this →

2. TABLE node meta["val"] is copied from the tanh, not recomputed. Unlike InsertTableOpsPass, which retraces via super().call() to refresh fake-tensor types, this pass inserts the TABLE via create_node(from_node=match.tanh, ...) and does not retrace. That copies the tanh's val. In a fully quantized graph the tanh output should already be int8, so shape/dtype line up — but it's worth a sanity check that no downstream pass expects a val produced by the TABLE fake kernel specifically (it validates table shape/dtype). If the copied val dtype ever isn't int8, serialization could be surprised.

3. _source_qargs scale match uses math.isclose(rel_tol=1e-6) (:180) to pick the activation qargs among candidates. Fine in practice, but if two producers ever carry near-identical scales it silently selects the first. Given how much else in this pass "fails closed," consider whether an ambiguous match (>1 candidate within tolerance) should also bail rather than take the first.

Nits

• JakeStevens' copyright comment appears addressed — the file carries Copyright (c) Meta Platforms, Inc. and affiliates. Note the rest of backends/arm/_passes/ also carries an Arm Limited line (e.g. arm_pass_manager.py, arm_pass.py); confirm with the Arm reviewers whether the dual header is expected here, since this lands in their backend.
• The circular-import workaround (module-scope vs. in-function import of register_pass_factories_before, :331-336) is well-documented in both the code and PR body. Good.

CI

The 2 reported failures (hifi-build, vision-build) are pull_request_target fork-checkout refusals — infrastructure/security, unrelated to this change.

---

Nice work — the exhaustive parity test against the reference model is exactly the right safety net for a byte-exact rewrite like this. My only substantive ask is confirming the alpha buffer doesn't linger (#1).
· branch export-D116560649

[cgreenberg]

@claude review this PR

[rascani]

Review automatically exported from Phabricator review in Meta.

[cgreenberg]

Summary:

Dynamic Tanh (DyT) normalization computes tanh(alpha * x) with a learned scalar alpha. Quantized, that lowers to a full-tensor integer Mul followed by a tanh TABLE. This adds FoldDyTAlphaIntoLUTPass, which replays the exact TOSA SINGLE_ROUND integer arithmetic over all 256 int8 codes and folds the alpha Mul into the TABLE without approximation.

The matcher fails closed for unsupported shapes, ranges, layouts, or quantization metadata. This revision contains only the public pass implementation, tests, and package exports; pass registration is intentionally outside this public change.

Reviewed By: rascani

Differential Revision: D116560649

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit bd9715a with merge base fbd4bbf ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: cgreenberg / name: Charles Greenberg (0cec5b5)

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[cgreenberg]

@pytorchbot label "release notes: arm"

[cgreenberg]

/easycla

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 12m 7s —— View job

---

Code Review: Fold DyT alpha scalar into tanh LUT

• Read changed files
• Analyze the pass implementation
• Analyze pass manager hook
• Analyze tests
• Post review

Reviewed all four changed files plus the surrounding machinery (op_tosa_rescale._compute_multiplier_and_shift, insert_table_ops, real_impl, arm_pass). This is a carefully-built pass with unusually strong test coverage. Overall it looks correct and well-guarded — the concerns below are minor.

---

Correctness — looks solid

• SINGLE_ROUND replay is faithful._apply_tosa_rescale (fold_dyt_alpha_into_lut_pass.py:57) mirrors the serialized RESCALE path: same _compute_multiplier_and_shift, rounding = 1 << (shift-1), (centered*multiplier + rounding) >> shift, then clamp to output dtype. The comment at :65 correctly notes _compute_multiplier_and_shift asserts shift ∈ [2, 62], so shift-1 is never negative. Nice that you resisted "guarding" it.
• The bit-exactness is genuinely verified, not asserted.test_tosa_output_is_bit_exact_after_fold compares the generated TABLE against the actual tosa_reference_model output (_tosa_reference_outputs) over every int8 code. That's the right way to pin the int32 MUL-wrap and rescale rounding, and it removes any doubt about the hand-replayed arithmetic matching the device.
• Fails-closed guards are comprehensive: non-scalar alpha (:88, :236-238), activation-side rank views (:243), int32/int8 dtype checks, and output_rescale.output_zp != tanh_input_qargs zp (:267). test_pass_rejects_activation_side_rank_view and the DyTModule docstring on the conv-producer requirement show the edges were thought through.
• TABLE wiring is consistent: input is the original int8 match.activation, input_qparams = activation_qargs, output_qparams = tanh_output_qargs — the LUT collapses rescale→mul→rescale→tanh into one int8→int8 map. Correct.

Minor concerns / questions

1. Dangling alpha buffer after the fold.call() erases the tanh and relies on eliminate_dead_code() (:323) to drop the mul and rescales — but eliminate_dead_code does not remove placeholder/buffer nodes. Once the alpha Mul is gone, the alpha buffer placeholder has no users. Does a later pipeline pass (FuseEqualPlaceholdersPass, constant-folding) actually strip it, or does the exported program keep an orphaned buffer in its state_dict + graph signature? It's functionally harmless and tiny, but the PR summary says the pass "deletes the Mul and its surrounding rescales," and the alpha constant silently surviving is worth confirming. None of the tests assert the buffer is removed. Fix this →

2. TABLE node meta["val"] is copied from the tanh, not recomputed. Unlike InsertTableOpsPass, which retraces via super().call() to refresh fake-tensor types, this pass inserts the TABLE via create_node(from_node=match.tanh, ...) and does not retrace. That copies the tanh's val. In a fully quantized graph the tanh output should already be int8, so shape/dtype line up — but it's worth a sanity check that no downstream pass expects a val produced by the TABLE fake kernel specifically (it validates table shape/dtype). If the copied val dtype ever isn't int8, serialization could be surprised.

3. _source_qargs scale match uses math.isclose(rel_tol=1e-6) (:180) to pick the activation qargs among candidates. Fine in practice, but if two producers ever carry near-identical scales it silently selects the first. Given how much else in this pass "fails closed," consider whether an ambiguous match (>1 candidate within tolerance) should also bail rather than take the first.

Nits

• JakeStevens' copyright comment appears addressed — the file carries Copyright (c) Meta Platforms, Inc. and affiliates. Note the rest of backends/arm/_passes/ also carries an Arm Limited line (e.g. arm_pass_manager.py, arm_pass.py); confirm with the Arm reviewers whether the dual header is expected here, since this lands in their backend.
• The circular-import workaround (module-scope vs. in-function import of register_pass_factories_before, :331-336) is well-documented in both the code and PR body. Good.

CI

The 2 reported failures (hifi-build, vision-build) are pull_request_target fork-checkout refusals — infrastructure/security, unrelated to this change.

---

Nice work — the exhaustive parity test against the reference model is exactly the right safety net for a byte-exact rewrite like this. My only substantive ask is confirming the alpha buffer doesn't linger (#1).
· branch export-D116560649

[cgreenberg]

@claude review this PR

[rascani]

Review automatically exported from Phabricator review in Meta.

[cgreenberg]

Summary:

Dynamic Tanh (DyT) normalization computes tanh(alpha * x) with a learned scalar alpha. Quantized, that lowers to a full-tensor integer Mul followed by a tanh TABLE. This adds FoldDyTAlphaIntoLUTPass, which replays the exact TOSA SINGLE_ROUND integer arithmetic over all 256 int8 codes and folds the alpha Mul into the TABLE without approximation.

The matcher fails closed for unsupported shapes, ranges, layouts, or quantization metadata. This revision contains only the public pass implementation, tests, and package exports; pass registration is intentionally outside this public change.

Reviewed By: rascani

Differential Revision: D116560649

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit bd9715a with merge base fbd4bbf ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: cgreenberg / name: Charles Greenberg (0cec5b5)

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[cgreenberg]

@pytorchbot label "release notes: arm"

[cgreenberg]

/easycla

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 12m 7s —— View job

---

Code Review: Fold DyT alpha scalar into tanh LUT

• Read changed files
• Analyze the pass implementation
• Analyze pass manager hook
• Analyze tests
• Post review

Reviewed all four changed files plus the surrounding machinery (op_tosa_rescale._compute_multiplier_and_shift, insert_table_ops, real_impl, arm_pass). This is a carefully-built pass with unusually strong test coverage. Overall it looks correct and well-guarded — the concerns below are minor.

---

Correctness — looks solid

• SINGLE_ROUND replay is faithful._apply_tosa_rescale (fold_dyt_alpha_into_lut_pass.py:57) mirrors the serialized RESCALE path: same _compute_multiplier_and_shift, rounding = 1 << (shift-1), (centered*multiplier + rounding) >> shift, then clamp to output dtype. The comment at :65 correctly notes _compute_multiplier_and_shift asserts shift ∈ [2, 62], so shift-1 is never negative. Nice that you resisted "guarding" it.
• The bit-exactness is genuinely verified, not asserted.test_tosa_output_is_bit_exact_after_fold compares the generated TABLE against the actual tosa_reference_model output (_tosa_reference_outputs) over every int8 code. That's the right way to pin the int32 MUL-wrap and rescale rounding, and it removes any doubt about the hand-replayed arithmetic matching the device.
• Fails-closed guards are comprehensive: non-scalar alpha (:88, :236-238), activation-side rank views (:243), int32/int8 dtype checks, and output_rescale.output_zp != tanh_input_qargs zp (:267). test_pass_rejects_activation_side_rank_view and the DyTModule docstring on the conv-producer requirement show the edges were thought through.
• TABLE wiring is consistent: input is the original int8 match.activation, input_qparams = activation_qargs, output_qparams = tanh_output_qargs — the LUT collapses rescale→mul→rescale→tanh into one int8→int8 map. Correct.

Minor concerns / questions

1. Dangling alpha buffer after the fold.call() erases the tanh and relies on eliminate_dead_code() (:323) to drop the mul and rescales — but eliminate_dead_code does not remove placeholder/buffer nodes. Once the alpha Mul is gone, the alpha buffer placeholder has no users. Does a later pipeline pass (FuseEqualPlaceholdersPass, constant-folding) actually strip it, or does the exported program keep an orphaned buffer in its state_dict + graph signature? It's functionally harmless and tiny, but the PR summary says the pass "deletes the Mul and its surrounding rescales," and the alpha constant silently surviving is worth confirming. None of the tests assert the buffer is removed. Fix this →

2. TABLE node meta["val"] is copied from the tanh, not recomputed. Unlike InsertTableOpsPass, which retraces via super().call() to refresh fake-tensor types, this pass inserts the TABLE via create_node(from_node=match.tanh, ...) and does not retrace. That copies the tanh's val. In a fully quantized graph the tanh output should already be int8, so shape/dtype line up — but it's worth a sanity check that no downstream pass expects a val produced by the TABLE fake kernel specifically (it validates table shape/dtype). If the copied val dtype ever isn't int8, serialization could be surprised.

3. _source_qargs scale match uses math.isclose(rel_tol=1e-6) (:180) to pick the activation qargs among candidates. Fine in practice, but if two producers ever carry near-identical scales it silently selects the first. Given how much else in this pass "fails closed," consider whether an ambiguous match (>1 candidate within tolerance) should also bail rather than take the first.

Nits

• JakeStevens' copyright comment appears addressed — the file carries Copyright (c) Meta Platforms, Inc. and affiliates. Note the rest of backends/arm/_passes/ also carries an Arm Limited line (e.g. arm_pass_manager.py, arm_pass.py); confirm with the Arm reviewers whether the dual header is expected here, since this lands in their backend.
• The circular-import workaround (module-scope vs. in-function import of register_pass_factories_before, :331-336) is well-documented in both the code and PR body. Good.

CI

The 2 reported failures (hifi-build, vision-build) are pull_request_target fork-checkout refusals — infrastructure/security, unrelated to this change.

---

Nice work — the exhaustive parity test against the reference model is exactly the right safety net for a byte-exact rewrite like this. My only substantive ask is confirming the alpha buffer doesn't linger (#1).
· branch export-D116560649

[cgreenberg]

@claude review this PR

[rascani]

Review automatically exported from Phabricator review in Meta.

[cgreenberg]

Summary:

Dynamic Tanh (DyT) normalization computes tanh(alpha * x) with a learned scalar alpha. Quantized, that lowers to a full-tensor integer Mul followed by a tanh TABLE. This adds FoldDyTAlphaIntoLUTPass, which replays the exact TOSA SINGLE_ROUND integer arithmetic over all 256 int8 codes and folds the alpha Mul into the TABLE without approximation.

The matcher fails closed for unsupported shapes, ranges, layouts, or quantization metadata. This revision contains only the public pass implementation, tests, and package exports; pass registration is intentionally outside this public change.

Reviewed By: rascani

Differential Revision: D116560649

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit bd9715a with merge base fbd4bbf ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: cgreenberg / name: Charles Greenberg (0cec5b5)

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[cgreenberg]

@pytorchbot label "release notes: arm"

[cgreenberg]

/easycla

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 12m 7s —— View job

---

Code Review: Fold DyT alpha scalar into tanh LUT

• Read changed files
• Analyze the pass implementation
• Analyze pass manager hook
• Analyze tests
• Post review

Reviewed all four changed files plus the surrounding machinery (op_tosa_rescale._compute_multiplier_and_shift, insert_table_ops, real_impl, arm_pass). This is a carefully-built pass with unusually strong test coverage. Overall it looks correct and well-guarded — the concerns below are minor.

---

Correctness — looks solid

• SINGLE_ROUND replay is faithful._apply_tosa_rescale (fold_dyt_alpha_into_lut_pass.py:57) mirrors the serialized RESCALE path: same _compute_multiplier_and_shift, rounding = 1 << (shift-1), (centered*multiplier + rounding) >> shift, then clamp to output dtype. The comment at :65 correctly notes _compute_multiplier_and_shift asserts shift ∈ [2, 62], so shift-1 is never negative. Nice that you resisted "guarding" it.
• The bit-exactness is genuinely verified, not asserted.test_tosa_output_is_bit_exact_after_fold compares the generated TABLE against the actual tosa_reference_model output (_tosa_reference_outputs) over every int8 code. That's the right way to pin the int32 MUL-wrap and rescale rounding, and it removes any doubt about the hand-replayed arithmetic matching the device.
• Fails-closed guards are comprehensive: non-scalar alpha (:88, :236-238), activation-side rank views (:243), int32/int8 dtype checks, and output_rescale.output_zp != tanh_input_qargs zp (:267). test_pass_rejects_activation_side_rank_view and the DyTModule docstring on the conv-producer requirement show the edges were thought through.
• TABLE wiring is consistent: input is the original int8 match.activation, input_qparams = activation_qargs, output_qparams = tanh_output_qargs — the LUT collapses rescale→mul→rescale→tanh into one int8→int8 map. Correct.

Minor concerns / questions

1. Dangling alpha buffer after the fold.call() erases the tanh and relies on eliminate_dead_code() (:323) to drop the mul and rescales — but eliminate_dead_code does not remove placeholder/buffer nodes. Once the alpha Mul is gone, the alpha buffer placeholder has no users. Does a later pipeline pass (FuseEqualPlaceholdersPass, constant-folding) actually strip it, or does the exported program keep an orphaned buffer in its state_dict + graph signature? It's functionally harmless and tiny, but the PR summary says the pass "deletes the Mul and its surrounding rescales," and the alpha constant silently surviving is worth confirming. None of the tests assert the buffer is removed. Fix this →

2. TABLE node meta["val"] is copied from the tanh, not recomputed. Unlike InsertTableOpsPass, which retraces via super().call() to refresh fake-tensor types, this pass inserts the TABLE via create_node(from_node=match.tanh, ...) and does not retrace. That copies the tanh's val. In a fully quantized graph the tanh output should already be int8, so shape/dtype line up — but it's worth a sanity check that no downstream pass expects a val produced by the TABLE fake kernel specifically (it validates table shape/dtype). If the copied val dtype ever isn't int8, serialization could be surprised.

3. _source_qargs scale match uses math.isclose(rel_tol=1e-6) (:180) to pick the activation qargs among candidates. Fine in practice, but if two producers ever carry near-identical scales it silently selects the first. Given how much else in this pass "fails closed," consider whether an ambiguous match (>1 candidate within tolerance) should also bail rather than take the first.

Nits

• JakeStevens' copyright comment appears addressed — the file carries Copyright (c) Meta Platforms, Inc. and affiliates. Note the rest of backends/arm/_passes/ also carries an Arm Limited line (e.g. arm_pass_manager.py, arm_pass.py); confirm with the Arm reviewers whether the dual header is expected here, since this lands in their backend.
• The circular-import workaround (module-scope vs. in-function import of register_pass_factories_before, :331-336) is well-documented in both the code and PR body. Good.

CI

The 2 reported failures (hifi-build, vision-build) are pull_request_target fork-checkout refusals — infrastructure/security, unrelated to this change.

---

Nice work — the exhaustive parity test against the reference model is exactly the right safety net for a byte-exact rewrite like this. My only substantive ask is confirming the alpha buffer doesn't linger (#1).
· branch export-D116560649

[cgreenberg]

@claude review this PR

[rascani]

Review automatically exported from Phabricator review in Meta.

[cgreenberg]

Summary:

Dynamic Tanh (DyT) normalization computes tanh(alpha * x) with a learned scalar alpha. Quantized, that lowers to a full-tensor integer Mul followed by a tanh TABLE. This adds FoldDyTAlphaIntoLUTPass, which replays the exact TOSA SINGLE_ROUND integer arithmetic over all 256 int8 codes and folds the alpha Mul into the TABLE without approximation.

The matcher fails closed for unsupported shapes, ranges, layouts, or quantization metadata. This revision contains only the public pass implementation, tests, and package exports; pass registration is intentionally outside this public change.

Reviewed By: rascani

Differential Revision: D116560649

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit bd9715a with merge base fbd4bbf ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: cgreenberg / name: Charles Greenberg (0cec5b5)

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[cgreenberg]

@pytorchbot label "release notes: arm"

[cgreenberg]

/easycla

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 12m 7s —— View job

---

Code Review: Fold DyT alpha scalar into tanh LUT

• Read changed files
• Analyze the pass implementation
• Analyze pass manager hook
• Analyze tests
• Post review

Reviewed all four changed files plus the surrounding machinery (op_tosa_rescale._compute_multiplier_and_shift, insert_table_ops, real_impl, arm_pass). This is a carefully-built pass with unusually strong test coverage. Overall it looks correct and well-guarded — the concerns below are minor.

---

Correctness — looks solid

• SINGLE_ROUND replay is faithful._apply_tosa_rescale (fold_dyt_alpha_into_lut_pass.py:57) mirrors the serialized RESCALE path: same _compute_multiplier_and_shift, rounding = 1 << (shift-1), (centered*multiplier + rounding) >> shift, then clamp to output dtype. The comment at :65 correctly notes _compute_multiplier_and_shift asserts shift ∈ [2, 62], so shift-1 is never negative. Nice that you resisted "guarding" it.
• The bit-exactness is genuinely verified, not asserted.test_tosa_output_is_bit_exact_after_fold compares the generated TABLE against the actual tosa_reference_model output (_tosa_reference_outputs) over every int8 code. That's the right way to pin the int32 MUL-wrap and rescale rounding, and it removes any doubt about the hand-replayed arithmetic matching the device.
• Fails-closed guards are comprehensive: non-scalar alpha (:88, :236-238), activation-side rank views (:243), int32/int8 dtype checks, and output_rescale.output_zp != tanh_input_qargs zp (:267). test_pass_rejects_activation_side_rank_view and the DyTModule docstring on the conv-producer requirement show the edges were thought through.
• TABLE wiring is consistent: input is the original int8 match.activation, input_qparams = activation_qargs, output_qparams = tanh_output_qargs — the LUT collapses rescale→mul→rescale→tanh into one int8→int8 map. Correct.

Minor concerns / questions

1. Dangling alpha buffer after the fold.call() erases the tanh and relies on eliminate_dead_code() (:323) to drop the mul and rescales — but eliminate_dead_code does not remove placeholder/buffer nodes. Once the alpha Mul is gone, the alpha buffer placeholder has no users. Does a later pipeline pass (FuseEqualPlaceholdersPass, constant-folding) actually strip it, or does the exported program keep an orphaned buffer in its state_dict + graph signature? It's functionally harmless and tiny, but the PR summary says the pass "deletes the Mul and its surrounding rescales," and the alpha constant silently surviving is worth confirming. None of the tests assert the buffer is removed. Fix this →

2. TABLE node meta["val"] is copied from the tanh, not recomputed. Unlike InsertTableOpsPass, which retraces via super().call() to refresh fake-tensor types, this pass inserts the TABLE via create_node(from_node=match.tanh, ...) and does not retrace. That copies the tanh's val. In a fully quantized graph the tanh output should already be int8, so shape/dtype line up — but it's worth a sanity check that no downstream pass expects a val produced by the TABLE fake kernel specifically (it validates table shape/dtype). If the copied val dtype ever isn't int8, serialization could be surprised.

3. _source_qargs scale match uses math.isclose(rel_tol=1e-6) (:180) to pick the activation qargs among candidates. Fine in practice, but if two producers ever carry near-identical scales it silently selects the first. Given how much else in this pass "fails closed," consider whether an ambiguous match (>1 candidate within tolerance) should also bail rather than take the first.

Nits

• JakeStevens' copyright comment appears addressed — the file carries Copyright (c) Meta Platforms, Inc. and affiliates. Note the rest of backends/arm/_passes/ also carries an Arm Limited line (e.g. arm_pass_manager.py, arm_pass.py); confirm with the Arm reviewers whether the dual header is expected here, since this lands in their backend.
• The circular-import workaround (module-scope vs. in-function import of register_pass_factories_before, :331-336) is well-documented in both the code and PR body. Good.

CI

The 2 reported failures (hifi-build, vision-build) are pull_request_target fork-checkout refusals — infrastructure/security, unrelated to this change.

---

Nice work — the exhaustive parity test against the reference model is exactly the right safety net for a byte-exact rewrite like this. My only substantive ask is confirming the alpha buffer doesn't linger (#1).
· branch export-D116560649

[cgreenberg]

@claude review this PR

[rascani]

Review automatically exported from Phabricator review in Meta.

[cgreenberg]

Summary:

Dynamic Tanh (DyT) normalization computes tanh(alpha * x) with a learned scalar alpha. Quantized, that lowers to a full-tensor integer Mul followed by a tanh TABLE. This adds FoldDyTAlphaIntoLUTPass, which replays the exact TOSA SINGLE_ROUND integer arithmetic over all 256 int8 codes and folds the alpha Mul into the TABLE without approximation.

The matcher fails closed for unsupported shapes, ranges, layouts, or quantization metadata. This revision contains only the public pass implementation, tests, and package exports; pass registration is intentionally outside this public change.

Reviewed By: rascani

Differential Revision: D116560649

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit bd9715a with merge base fbd4bbf ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: cgreenberg / name: Charles Greenberg (0cec5b5)

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[cgreenberg]

@pytorchbot label "release notes: arm"

[cgreenberg]

/easycla

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 12m 7s —— View job

---

Code Review: Fold DyT alpha scalar into tanh LUT

• Read changed files
• Analyze the pass implementation
• Analyze pass manager hook
• Analyze tests
• Post review

Reviewed all four changed files plus the surrounding machinery (op_tosa_rescale._compute_multiplier_and_shift, insert_table_ops, real_impl, arm_pass). This is a carefully-built pass with unusually strong test coverage. Overall it looks correct and well-guarded — the concerns below are minor.

---

Correctness — looks solid

• SINGLE_ROUND replay is faithful._apply_tosa_rescale (fold_dyt_alpha_into_lut_pass.py:57) mirrors the serialized RESCALE path: same _compute_multiplier_and_shift, rounding = 1 << (shift-1), (centered*multiplier + rounding) >> shift, then clamp to output dtype. The comment at :65 correctly notes _compute_multiplier_and_shift asserts shift ∈ [2, 62], so shift-1 is never negative. Nice that you resisted "guarding" it.
• The bit-exactness is genuinely verified, not asserted.test_tosa_output_is_bit_exact_after_fold compares the generated TABLE against the actual tosa_reference_model output (_tosa_reference_outputs) over every int8 code. That's the right way to pin the int32 MUL-wrap and rescale rounding, and it removes any doubt about the hand-replayed arithmetic matching the device.
• Fails-closed guards are comprehensive: non-scalar alpha (:88, :236-238), activation-side rank views (:243), int32/int8 dtype checks, and output_rescale.output_zp != tanh_input_qargs zp (:267). test_pass_rejects_activation_side_rank_view and the DyTModule docstring on the conv-producer requirement show the edges were thought through.
• TABLE wiring is consistent: input is the original int8 match.activation, input_qparams = activation_qargs, output_qparams = tanh_output_qargs — the LUT collapses rescale→mul→rescale→tanh into one int8→int8 map. Correct.

Minor concerns / questions

1. Dangling alpha buffer after the fold.call() erases the tanh and relies on eliminate_dead_code() (:323) to drop the mul and rescales — but eliminate_dead_code does not remove placeholder/buffer nodes. Once the alpha Mul is gone, the alpha buffer placeholder has no users. Does a later pipeline pass (FuseEqualPlaceholdersPass, constant-folding) actually strip it, or does the exported program keep an orphaned buffer in its state_dict + graph signature? It's functionally harmless and tiny, but the PR summary says the pass "deletes the Mul and its surrounding rescales," and the alpha constant silently surviving is worth confirming. None of the tests assert the buffer is removed. Fix this →

2. TABLE node meta["val"] is copied from the tanh, not recomputed. Unlike InsertTableOpsPass, which retraces via super().call() to refresh fake-tensor types, this pass inserts the TABLE via create_node(from_node=match.tanh, ...) and does not retrace. That copies the tanh's val. In a fully quantized graph the tanh output should already be int8, so shape/dtype line up — but it's worth a sanity check that no downstream pass expects a val produced by the TABLE fake kernel specifically (it validates table shape/dtype). If the copied val dtype ever isn't int8, serialization could be surprised.

3. _source_qargs scale match uses math.isclose(rel_tol=1e-6) (:180) to pick the activation qargs among candidates. Fine in practice, but if two producers ever carry near-identical scales it silently selects the first. Given how much else in this pass "fails closed," consider whether an ambiguous match (>1 candidate within tolerance) should also bail rather than take the first.

Nits

• JakeStevens' copyright comment appears addressed — the file carries Copyright (c) Meta Platforms, Inc. and affiliates. Note the rest of backends/arm/_passes/ also carries an Arm Limited line (e.g. arm_pass_manager.py, arm_pass.py); confirm with the Arm reviewers whether the dual header is expected here, since this lands in their backend.
• The circular-import workaround (module-scope vs. in-function import of register_pass_factories_before, :331-336) is well-documented in both the code and PR body. Good.

CI

The 2 reported failures (hifi-build, vision-build) are pull_request_target fork-checkout refusals — infrastructure/security, unrelated to this change.

---

Nice work — the exhaustive parity test against the reference model is exactly the right safety net for a byte-exact rewrite like this. My only substantive ask is confirming the alpha buffer doesn't linger (#1).
· branch export-D116560649

[cgreenberg]

@claude review this PR

[rascani]

Review automatically exported from Phabricator review in Meta.