Skip to content

[release/2.8] Fix unexpected failure in test_prepare_softmax_nrow_2048_dim_0 when SM/CU count is <= 32 #2937

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sstamenk wants to merge 1 commit into
base: release/2.8
Choose a base branch
from
+14 −2
Open

[release/2.8] Fix unexpected failure in test_prepare_softmax_nrow_2048_dim_0 when SM/CU count is <= 32 #2937

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 14 additions & 2 deletions test/inductor/test_online_softmax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import torch._inductor.config as inductor_config
import torch.nn.functional as F
from torch._dynamo.utils import rmse, same
from torch._inductor.runtime.hints import DeviceProperties
from torch._inductor.test_case import run_tests, TestCase
from torch._inductor.utils import run_and_get_code
from torch.testing._internal.common_utils import (
Expand Down Expand Up @@ -138,8 +139,19 @@ def test_prepare_softmax(self, dim, nrow):
self.assertTrue(same(ref, act, tol=1e-2))

if nrow == 2048 and dim == 0:
# split reduction is triggered. We have multiple kernels
self.assertTrue(code.count("def triton") >= 2)
num_kernels = 2
# split reduction may be triggered depending on the device's SM/CU count.
# The heuristic in num_splits() in ir.py returns split=1 (no split) when:
# numel_hint >= num_sm * 2 * 32
# When dim=0, numel_hint (output size) = ncol (2048)
# split is expected only when num_sm > 32
props = DeviceProperties.create(torch.device(GPU_TYPE))
num_sm = props.multi_processor_count
split_not_expected = 2048 >= num_sm * 2 * 32
if split_not_expected:
num_kernels = 1

self.assertTrue(code.count("def triton") >= num_kernels)
else:
if nrow == 2 and dim == 0:
# persistent reduction triggered
Expand Down