Skip to content

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

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

[fix] correct return type of fused arnorm fake impl#1502
valarLip merged 4 commits intomainfrom
guanbao/fix_fused_arnorm_fake_impl

Conversation

@gbyu-amd
Copy link
Contributor

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

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

Copy link
Collaborator

@valarLip valarLip left a comment

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
farlukas pushed a commit that referenced this pull request Dec 4, 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>
nsusanto pushed a commit that referenced this pull request Dec 4, 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>
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>
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