Skip to content

[fix] correct return type of fused arnorm fake impl - #1502

Merged
valarLip merged 4 commits into
mainfrom
guanbao/fix_fused_arnorm_fake_impl
Nov 28, 2025
Merged

valarLip merged 4 commits into
mainfrom
guanbao/fix_fused_arnorm_fake_impl

Conversation

@gbyu-amd

@gbyu-amd gbyu-amd commented Nov 28, 2025

Copy link
Copy Markdown
Contributor

Motivation

Correct the return type of arnorm fake impl to a tuple, otherwise we will see dynamo error during torch.compile model using this custom fusion kernel:

torch._dynamo.exc.TorchRuntimeError: Dynamo failed to run FX node with fake tensors: call_function aiter.fused_allreduce_rmsnorm_(*(FakeTensor(..., device='cuda:4', size=(s72, 7168), dtype=torch.float8_e4m3fn), FakeTensor(..., device='cuda:4', size=(s72, 7168)), Parameter(FakeTensor(..., device='cuda:4', size=(7168,), requires_grad=True)), 1e-06), **{'group_name': 'tp:0'}): got TypeError("Object of type 'FakeTensor' is not an instance of 'sequence'")
image

Test Plan

Test it with deepseek in ATOM.

Test Result

It can be compiled without error and run normally.

Submission Checklist

@valarLip valarLip left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@valarLip
valarLip merged commit 05ffc19 into main Nov 28, 2025
27 of 28 checks passed
@valarLip
valarLip deleted the guanbao/fix_fused_arnorm_fake_impl branch November 28, 2025 09:40
zhuyuhua-v pushed a commit that referenced this pull request Dec 17, 2025
* [fix]: fused_ar_rms acc error

Signed-off-by: root <root@hjbog-srdc-24.amd.com>

* [fix]: reduce_scatter boundary check error

Signed-off-by: root <root@hjbog-srdc-24.amd.com>

* correct return type of fused_arnorm fake impl

---------

Signed-off-by: root <root@hjbog-srdc-24.amd.com>
Co-authored-by: root <root@hjbog-srdc-24.amd.com>
Co-authored-by: guanbao <gyu@amd.com>
valarLip pushed a commit that referenced this pull request Mar 18, 2026
* [fix]: fused_ar_rms acc error

Signed-off-by: root <root@hjbog-srdc-24.amd.com>

* [fix]: reduce_scatter boundary check error

Signed-off-by: root <root@hjbog-srdc-24.amd.com>

* correct return type of fused_arnorm fake impl

---------

Signed-off-by: root <root@hjbog-srdc-24.amd.com>
Co-authored-by: root <root@hjbog-srdc-24.amd.com>
Co-authored-by: guanbao <gyu@amd.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants