[Spyre-Next] Reworked forward call chain by bohnstingl · Pull Request #872 · torch-spyre/sendnn-inference

bohnstingl · 2026-03-27T00:34:26Z

Description

This PR reworks the call chain of the forward functions for the layer wrappings for torch-spyre. In particular, it avoids overriding forward_native in favor of using forward_oot.

cc @romitjain

Related Issues

This PR was triggered by #863

Test Plan

This change is non user-facing and thus all the existing tests should work. Moreover, it should enable the upstream vLLM tests to be supported as well.

Checklist

I have read the contributing guidelines
My code follows the project's code style (run bash format.sh)
I have added tests for my changes (if applicable)
I have updated the documentation (if applicable)
My commits include a Signed-off-by: line (DCO compliance)

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

github-actions · 2026-03-27T00:34:35Z

👋 Hi! Thank you for contributing to vLLM support on Spyre.
Just a reminder: Make sure that your code passes all the linting checks, otherwise your PR won't be able to be merged. To do so, run ./format.sh.
Now you are good to go 🚀.

We also recommend installing prek and configuring it to check your code before every local commit.

romitjain · 2026-03-27T07:58:02Z

        if x.shape[-1] != hidden_size:
            raise ValueError(f"Expected hidden_size to be {hidden_size}, but found: {x.shape[-1]}")

-        x = x.transpose(-1, -2).contiguous()


@bohnstingl

I tried running rmsnorm repo tests on this branch it failed for me. I am not able to pinpoint why that might be. On main, they are passing, and apart from the function name changes, the only change I see is this.

@romitjain you need a very recent commit of torch-spyre for this to work. I reverted this rework from this PR and I'll move it into a separate PR.

@bohnstingl
My bad - the tests are passing iff I check one of forward_oot, forward, or forward_native. But if we check all 3, I am getting a test failed (tensors are not close) in the residual path. I think it's because we mutate the residual tensor: https://github.com/vllm-project/vllm-spyre/blob/main/vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py#L234

So this change: 4b6b2e9 might not be needed. I was able to pass tests locally with and without this change. I have made a small PR here from my fork: bohnstingl#1.

PR is for demonstration - feel free to merge/close and make changes directly here.

Thanks, I merged it in.

I will still leave the rms_norm rework for a separate PR, in order not to dilute the intention of this PR.

rms_norm rework PR: #873

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

Updated rmsnorm tests

romitjain

LGTM

bohnstingl · 2026-03-27T10:42:13Z

@joerunde or @tjohnson31415 could you also take a look at it to see if the rework makes sense to you?

yannicks1 · 2026-03-27T15:06:15Z

@romitjain fyi DCO fails with: Commit sha: 8fc82cc, Author: romit, Committer: romit; The sign-off is missing.

tjohnson31415

Couple of NITs, but otherwise I like this change! Much cleaner to not override the upstream forward() and forward_native() functions.

yannicks1 · 2026-03-27T15:39:26Z

As this PR does mainly renaming: could we also rename self._fwd or is that fixed? It was not obvious at first glance that self._fwd is the compiled self.forward_spyre. I don't have the perfect name, but sth along self._fwd_spyre_comp, or just self._fwd_spyre or self._forward_spyre/ self._forward_spyre_comp

I renamed it to maybe_compiled_forward_spyre to indicate that 1) it is maybe compiled and 2) that it is the forward_spyre. I did the same for silu as well.

yannicks1 · 2026-03-27T15:50:01Z

for consistency: why not calling this forward_spyre as in the rsm op?

yannicks1 · 2026-03-27T15:50:37Z

same argument as above: consider renaming self._fwd to sth more descriptive if possible

Done, see above.

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

tjohnson31415

LGTM

tjohnson31415 · 2026-03-27T20:45:54Z

bot:next-test

bohnstingl · 2026-03-27T23:03:49Z

@dilipgb, @R3hankhan123, @coderfornow and @GOavi101 please feel free to also comment on this rework PR, as you are also working on layer wrappings. I am happy to incorporate any feedback you may have.

yannicks1

lgtm. please keep in mind to merge in the changes from #873.

yannicks1 · 2026-03-30T15:02:42Z

-    - forward(): custom op dispatch (no-compile path via torch.ops.vllm.spyre_rmsnorm)
-    - forward_native(): direct Spyre device execution
+    - forward_oot(): OOT dispatch via custom op (torch.ops.vllm.spyre_rmsnorm)
+    - reference_rms_norm: golden reference, similar to vLLM upstream pure PyTorch (ground truth)


nit: reference_rms_norm -> reference_rms_norm()

…d_oot_rework

… into forward_oot_rework

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

bohnstingl · 2026-03-30T17:25:47Z

bot:next-test

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

tjohnson31415 · 2026-03-30T19:53:26Z

bot test is broken for an unrelated reason, but tests are passing in dev pods. Merging!

Reworked forward calls

caa563f

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

github-actions Bot changed the title ~~Reworked forward call chain~~ [Spyre-Next] Reworked forward call chain Mar 27, 2026

bohnstingl mentioned this pull request Mar 27, 2026

[Spyre-Next] SpyreSiluAndMul Tests #863

Closed

5 tasks

romitjain reviewed Mar 27, 2026

View reviewed changes

Comment thread vllm_spyre_next/tests/test_rms_norm.py Outdated

romitjain reviewed Mar 27, 2026

View reviewed changes

Reverted change to rms_norm

4b6b2e9

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

bohnstingl requested a review from romitjain March 27, 2026 09:22

romitjain and others added 2 commits March 27, 2026 09:53

Updated rmsnorm tests

8fc82cc

Merge pull request #1 from romitjain/forward_oot_rework

a9ef99b

Updated rmsnorm tests

romitjain approved these changes Mar 27, 2026

View reviewed changes

bohnstingl marked this pull request as ready for review March 27, 2026 10:41

bohnstingl requested review from joerunde and prashantgupta24 as code owners March 27, 2026 10:41

bohnstingl self-assigned this Mar 27, 2026

This was referenced Mar 27, 2026

[Spyre-Next] Wrapped Embedding layer for spyre #836

Merged

[Spyre-Next] Add RowParallelLinear and ColumnParallelLinear(MLP) wrappers #869

Merged

bohnstingl requested a review from yannicks1 March 27, 2026 13:36

bohnstingl requested a review from tjohnson31415 March 27, 2026 15:16

tjohnson31415 reviewed Mar 27, 2026

View reviewed changes

Comment thread vllm_spyre_next/tests/test_rms_norm.py Outdated

Comment thread vllm_spyre_next/tests/test_rms_norm.py Outdated

yannicks1 reviewed Mar 27, 2026

View reviewed changes

bohnstingl added a commit to bohnstingl/vllm-spyre that referenced this pull request Mar 27, 2026

Integrated rework from torch-spyre#872

fc9bc0d

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

bohnstingl added 2 commits March 27, 2026 16:53

Fixed test_rms_norm.py with review comments

b167e6f

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

Renaming from codereview

4593ddc

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

bohnstingl requested review from tjohnson31415 and yannicks1 March 27, 2026 17:09

tjohnson31415 reviewed Mar 27, 2026

View reviewed changes

Comment thread vllm_spyre_next/vllm_spyre_next/custom_ops/rms_norm.py Outdated

Corrected typo

07729ab

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

tjohnson31415 approved these changes Mar 27, 2026

View reviewed changes

yannicks1 approved these changes Mar 30, 2026

View reviewed changes

bohnstingl added 3 commits March 30, 2026 15:56

Merge branch 'main' of github.com:vllm-project/vllm-spyre into forwar…

7b497c8

…d_oot_rework

Merge branch 'forward_oot_rework' of github.com:bohnstingl/vllm-spyre…

01350cb

… into forward_oot_rework

Nit change

1304aaf

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

bohnstingl added 2 commits March 30, 2026 18:36

Updated RMSNorm test

b5c2763

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

Ensured residual is not updated in-place

a147ae8

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

tjohnson31415 enabled auto-merge (squash) March 30, 2026 19:53

github-actions Bot added the ready Runs the full CI test suite. Only add to PRs once ready to merge to limit public GHA usage label Mar 30, 2026

tjohnson31415 merged commit 74fa89e into torch-spyre:main Mar 30, 2026
22 checks passed

bohnstingl mentioned this pull request Mar 31, 2026

[Spyre-Next] [Feature] Wrap RoPE layer on Spyre #881

Closed

2 tasks

Conversation

bohnstingl commented Mar 27, 2026

Description

Related Issues

Test Plan

Checklist

Uh oh!

github-actions Bot commented Mar 27, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

romitjain left a comment

Choose a reason for hiding this comment

Uh oh!

bohnstingl commented Mar 27, 2026

Uh oh!

yannicks1 commented Mar 27, 2026

Uh oh!

tjohnson31415 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tjohnson31415 left a comment

Choose a reason for hiding this comment

Uh oh!

tjohnson31415 commented Mar 27, 2026

Uh oh!

bohnstingl commented Mar 27, 2026

Uh oh!

yannicks1 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bohnstingl commented Mar 30, 2026

Uh oh!

tjohnson31415 commented Mar 30, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

tjohnson31415 left a comment •

edited

Loading