Skip to content

[inductor] BatchLinearLHSFusion: also match torch._C._nn.linear#3232

Open
reger-men wants to merge 1 commit into
ROCm:developfrom
reger-men:pr1-batch-linear-lhs
Open

[inductor] BatchLinearLHSFusion: also match torch._C._nn.linear#3232
reger-men wants to merge 1 commit into
ROCm:developfrom
reger-men:pr1-batch-linear-lhs

Conversation

@reger-men
Copy link
Copy Markdown

Dynamo inlines torch.nn.functional.linear into torch._C._nn.linear on hot paths, so BatchLinearLHSFusion.match silently misses the inlined form. This one-line change extends CallFunctionVarArgs to also accept torch._C._nn.linear.

The sibling fusion class PreGradBatchLinearFusion in the same file already accepts both call targets (line 626 in current main); this brings BatchLinearLHSFusion in line with that.

The fusion is gated by pre_grad_fusion_options["batch_linear_lhs"] (unset by default), so this is invisible to users who haven't opted in. is_linear_node_can_be_fused still filters out shapes that can't actually be fused (different weight shapes per head, etc).

Test plan

  • new test_batch_linear_lhs_fusion_nn_linear_inlined builds an nn.Linear module list (Dynamo inlines those to _C._nn.linear); without this change the matcher fires 0 times, with it 1 time
  • existing test_batch_linear_lhs_fusion continues to pass (regression for the F.linear form)
  • end-to-end: gradients and parameters match eager between compiled and uncompiled paths under compare_gradients / compare_parameters

@rocm-repo-management-api
Copy link
Copy Markdown

rocm-repo-management-api Bot commented May 19, 2026

Jenkins build for 1bc8980a52488cd6e77abdbf7c0dab34272e910f commit finished as FAILURE
Links: Pipeline Overview / Build artifacts / Test Results

@reger-men reger-men force-pushed the pr1-batch-linear-lhs branch from 1bc8980 to aed5e1b Compare May 20, 2026 18:05
@rocm-repo-management-api
Copy link
Copy Markdown

rocm-repo-management-api Bot commented May 20, 2026

Jenkins build for aed5e1bb3da2b806c42c10dc85a62dc5301b5f48 commit finished as FAILURE
Links: Pipeline Overview / Build artifacts / Test Results

@reger-men reger-men force-pushed the pr1-batch-linear-lhs branch from aed5e1b to d763082 Compare May 21, 2026 09:23
Dynamo inlines torch.nn.functional.linear into torch._C._nn.linear on
hot paths, so BatchLinearLHSFusion.match silently misses the inlined
form. This one-line change extends CallFunctionVarArgs to also accept
torch._C._nn.linear.

The sibling fusion class PreGradBatchLinearFusion in the same file
already accepts both call targets (line 626 in current main); this
brings BatchLinearLHSFusion in line with that.

The fusion is gated by pre_grad_fusion_options["batch_linear_lhs"]
(unset by default), so this is invisible to users who haven't opted
in. is_linear_node_can_be_fused still filters out shapes that can't
actually be fused (different weight shapes per head, etc).

Adds a regression test that builds a module out of nn.Linear modules
(Dynamo inlines those to _C._nn.linear). Without this change the test
expects 0 fires; with it, 1 fire.
@rocm-repo-management-api
Copy link
Copy Markdown

rocm-repo-management-api Bot commented May 21, 2026

Jenkins build for d7630821e00abbdc75ac9c8ceab013feab3772a5 commit finished as FAILURE
Links: Pipeline Overview / Build artifacts / Test Results

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.

1 participant