Summary:
Add fused rmsnorm fwd/bwd implementations from quack, written in
cutedsl. Rely on vendored quack.
Add testing to OpInfo for composability with other PyT systems, and
test_nn to match the existing fused norm tests.
Test Plan:
python -m pytest -v \
test/test_nn.py -k rmsnorm \
test/python_native/ \
test/test_ops.py -k 'rms_norm' \
test/inductor/test_torchinductor_opinfo.py -k 'rms_norm'
# Then separately, since OPINFO_RESTRICT_TO_DSL filters op_db globally:
OPINFO_RESTRICT_TO_DSL=cutedsl python -m pytest -v \
test/test_ops.py -k 'rms_norm and cutedsl' \
test/inductor/test_torchinductor_opinfo.py -k 'rms_norm and cutedsl' \
test/export/test_export_opinfo.py -k 'rms_norm and cutedsl'
Authors:
Originally from #178326
Co-authored-by: AaronWang04
Signed-off-by: Simon Layton [email protected]
Pull Request resolved: #182108
Approved by: https://github.com/drisspg
Co-authored-by: AaronWang04 [email protected]
SOCIAL SHARE CARD GENERATOR