Fix residual gradient in non-learnable LayerNorm backward

[rocketbastard]

Description

Fix residual gradient accumulation in the non-learnable LayerNorm backward path.

• Fuse dx_accumulate into the _layer_norm_bwd_dx Triton kernel.
• Remove the separate PyTorch dx.add_() operation.
• Add focused numerical unit tests for the LayerNorm operator.

Supersedes #447 because the original fork was deleted.
Fixes #446

Testing

The added unit tests passed successfully in the local GPU environment.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[greptile-apps]

Greptile Summary

The PR fuses residual-gradient accumulation into the non-learnable Triton LayerNorm backward kernel and adds focused numerical coverage.

• Adds an optional residual-gradient pointer and row stride to the backward kernel.
• Eliminates the separate post-kernel accumulation operation.
• Tests float32 and bfloat16 behavior across residual and non-residual cases with padded row strides.

Confidence Score: 4/5

The kernel change appears sound, but the new regression test's unresolved source-tree import should be fixed before merging so it can be collected reliably.

The numerical kernel path preserves the intended residual addition, while the new test imports a package that is neither rooted at the repository level nor made visible from the test module.

Files Needing Attention: examples/tests/commons/test_triton_layer_norm.py

Important Files Changed

Filename	Overview
examples/hstu/ops/triton_ops/triton_layer_norm.py	Correctly loads and adds the optional residual gradient inside the non-learnable backward kernel before storing the result.
examples/tests/commons/test_triton_layer_norm.py	Adds useful numerical coverage, but its HSTU-local import is unresolved when the test is collected from its new commons test location.

_{Reviews (1): Last reviewed commit: "Add gradient accumulation support in lay..." | Re-trigger Greptile}

[greptile-apps]

HSTU test import is unresolved

When pytest collects this test from the repository root or the examples/tests hierarchy without an external PYTHONPATH, ops is not importable because it resides under examples/hstu, causing collection to fail with ModuleNotFoundError.

[JacoCheung]

/build

[JacoCheung]

Hi @rocketbastard could you move the test code to hstu-specific test file ?

[JacoCheung]

❌ Pipeline #63652527 -- failed

Job	Status	Log
pre_check	❌ failed	view
train_build_x86	✅ success	view
train_build_arm64	✅ success	view
prepare-jet-b200-smoke	✅ success	view
prepare-jet-b200-inference	✅ success	view
prepare-jet-cw-dfw-e2e-benchmark	✅ success	view
build_whl	✅ success	view
dynamicemb_test_fwd_bwd_8gpus	✅ success	view
dynamicemb_test_load_dump_8gpus	✅ success	view
unit_test_1gpu_a100	❌ failed	view
unit_test_1gpu_h100	❌ failed	view
unit_test_4gpu	❌ failed	view
unit_test_tp_4gpu	❌ failed	view
L20_unit_test_1gpu	✅ success	view
inference_unit_test_1gpu	✅ success	view
inference_test_1gpu	✅ success	view

Result: 11/16 jobs passed

View full pipeline