fix(ttx/npu): keep SwiGLU intermediate computations in FP32

[shiwq-765]

No description provided.

[github-actions]

Claude Code Review

Verdict: Request changes -- The a5 backward kernel has duplicated/leftover statements that will reintroduce the low-precision path and likely fail to compile or produce wrong results.

Summary

Keeps the SiLU/product intermediate in fp32 in the NPU SwiGLU forward and backward, rounding only on the store, and adds a bitwise-equality accuracy test against the fp32 reference. The a5 variant's backward edit is incomplete and contains leftover code from the previous version.

Must fix

• [BLOCKER] a5 backward has duplicated/stale statements -- mojo_opset/backends/ttx/kernels/npu/a5/swiglu.py:147-153 -- After the new fp32 block computes da_chunk, the old da_factor = silu_a * (1 - sigmoid_a) + sigmoid_a and da_chunk = dc_chunk * b_chunk * da_factor.to(dc_chunk.dtype) lines remain, overwriting da_chunk with the low-precision result (and shadowing da_factor). Delete those two lines so a5 matches a2.

Suggestions

Suggestions (2)

• [MAJOR] Test asserts bitwise equality with MojoSwiGLU() default backend -- mojo_opset/tests/accuracy/operators/test_activation.py:129-131 -- atol=0, rtol=0 will only hold on the NPU kernel path; if this test runs on CPU/torch/bumi backends where SiLU uses libm, bitwise equality is not guaranteed. Confirm the test is gated to the NPU kernel or relax to a very tight tolerance for other backends.
• [MINOR] Redundant .to(tl.float32) on already-fp32 operand -- mojo_opset/backends/ttx/kernels/npu/a2/swiglu.py:84 -- silu_a is already fp32 from silu(a_f32); the cast of b_chunk is the meaningful change. Fine as-is, but a brief comment would help future readers understand the store handles the down-cast.

Nits

Nits (1)

• [NIT] Missing trailing newline -- mojo_opset/tests/accuracy/operators/test_activation.py:131 -- "No newline at end of file".

Notes

• [CHECK] Verify tl.store on the fp32 c_chunk/da_chunk/db_chunk implicitly down-casts to the destination pointer dtype on the NPU Triton backend; if not, an explicit .to(a_chunk.dtype) at the store is required to preserve the "round exactly once" contract -- mojo_opset/backends/ttx/kernels/npu/a2/swiglu.py:86,152-153.

[shengw-bd]

LGTM

[github-actions]

Claude Code Review

Verdict: Approve -- Fixes a real precision bug by keeping SiLU/derivative intermediates in fp32 until the final store; backward math is also simplified correctly.

Summary

The change removes premature down-casts in the SwiGLU forward and backward NPU kernels so the SiLU intermediate and derivative factor stay in fp32 until stored. The backward da_factor is also refactored to the algebraically equivalent sigmoid * (1 + a*(1-sigmoid)). A new accuracy test pins bitwise equality against the fp32 reference.

Must fix

None.

Suggestions (2)

• [MAJOR] Duplicate kernel files drift risk -- mojo_opset/backends/ttx/kernels/npu/a2/swiglu.py, mojo_opset/backends/ttx/kernels/npu/a5/swiglu.py -- The a2 and a5 copies are byte-identical; consider consolidating to prevent one being fixed and the other not next time.
• [MINOR] Backward test coverage -- mojo_opset/tests/accuracy/operators/test_activation.py:126 -- The new test only exercises forward; the backward change (rounding + reformulated da_factor) is not verified against a reference.

Nits (2)

• [NIT] Missing trailing newline -- mojo_opset/tests/accuracy/operators/test_activation.py:131 -- file ends without newline.
• [NIT] Parametrize tuples unnecessarily wrapped in lists -- mojo_opset/tests/accuracy/operators/test_activation.py:110 -- ([3072, 3072]) is just [3072, 3072]; the outer parens do nothing.

Notes

• [CHECK] da_factor = sigmoid * (1 + a*(1-sigmoid)) -- confirm this matches the previous silu*(1-sigmoid) + sigmoid = sigmoid*(1 + a*(1-sigmoid)) identity (it does algebraically, but worth a quick numerical spot check at large |a| where 1-sigmoid underflows differently).

[shengw-bd]

Is bitwise equality required here?

[github-actions]

Claude Code Review

Verdict: Comment -- Correct fix to keep SiLU/backward intermediates in fp32, but there's duplicated code across a2/a5 and the test's docstring claims bitwise equality while the assertion uses a loose tolerance.

Summary

The PR changes the NPU SwiGLU forward and backward kernels to keep the SiLU (and dc, b) intermediates in fp32 and only round at the store, matching the reference torch/bumi contract. It also simplifies the backward da_factor algebraically and adds an accuracy test.

Must fix

None.

Suggestions

Suggestions (3)

• [MAJOR] Duplicated kernel code across a2/a5 -- mojo_opset/backends/ttx/kernels/npu/a2/swiglu.py, mojo_opset/backends/ttx/kernels/npu/a5/swiglu.py -- The two files are byte-identical diffs; consider factoring into a shared module to avoid the two copies drifting on the next fix.
• [MAJOR] Test docstring contradicts assertion -- mojo_opset/tests/accuracy/operators/test_activation.py:126-131 -- Docstring says "asserts bitwise equality" but the call uses atol=1e-2, rtol=1e-2. Either tighten to atol=0, rtol=0 (or equal_nan=True exact compare) to actually pin the contract, or fix the docstring.
• [MINOR] Test does not exercise backward -- mojo_opset/tests/accuracy/operators/test_activation.py:110-131 -- The backward kernel changed non-trivially (new da_factor form); add a bwd accuracy check so a regression there is caught.

Nits

Nits (1)

• [NIT] Missing trailing newline -- mojo_opset/tests/accuracy/operators/test_activation.py:131 -- "No newline at end of file".

Notes

• [CHECK] mojo_opset/backends/ttx/kernels/npu/a2/swiglu.py:84 -- Confirm the store path implicitly casts the fp32 c_chunk back to the output dtype via c_ptrs' element type; if not, an explicit .to(...) at store is needed to match the "round exactly once at store" contract.

[github-actions]

Claude Code Review

Verdict: Approve -- Improves numerical precision by keeping SwiGLU intermediates in fp32 before the final store; kernels and math look correct.

Summary

The PR keeps the SiLU/SwiGLU forward and backward intermediates in fp32 (promoting b and dc instead of down-casting silu_a and da_factor), and rewrites da_factor into the algebraically equivalent sigmoid * (1 + a*(1-sigmoid)). It also adds a bf16 accuracy test that compares against the fp32 reference.

Must fix

None.

Suggestions (2)

• [MAJOR] Duplicated kernel source across a2/a5 -- mojo_opset/backends/ttx/kernels/npu/a2/swiglu.py:81, mojo_opset/backends/ttx/kernels/npu/a5/swiglu.py:81 -- The two files are byte-identical; consider sharing a single source to avoid the two variants drifting on future fixes.
• [MINOR] Store dtype relies on implicit cast -- mojo_opset/backends/ttx/kernels/npu/a2/swiglu.py:84 -- c_chunk is now fp32 while the output buffer is likely bf16/fp16; confirm tl.store performs the narrowing cast on this backend, otherwise add an explicit .to(a_chunk.dtype) at store time.

Nits (2)

• [NIT] Missing trailing newline -- mojo_opset/tests/accuracy/operators/test_activation.py:119 -- file ends without a newline.
• [NIT] Redundant parens in parametrize -- mojo_opset/tests/accuracy/operators/test_activation.py:111 -- ([3072, 3072]) is just a list; drop the outer parens.

Notes

• [CHECK] mojo_opset/backends/ttx/kernels/npu/a2/swiglu.py:150 -- Verify the new da_factor formulation matches the previous one under the kernel's rounding on NPU (algebraically equal, but fp32 ordering differs slightly).