Skip to content

[TLE] Stabilize FP16 cumsum correctness test #844

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

Merged
zhzhcookie merged 7 commits into from
Jul 28, 2026
+9 −15
Merged

[TLE] Stabilize FP16 cumsum correctness test #844

Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d9c5241
Stabilize TLE cumsum fp16 test
sunnycase Jul 24, 2026
32704e3
Merge branch 'main' into fix/tle_cumsum
sunnycase Jul 24, 2026
09bf751
Preserve TLE cumsum call-frame regression
sunnycase Jul 24, 2026
6749491
Merge branch 'main' into fix/tle_cumsum
zhzhcookie Jul 24, 2026
bb909d4
Make TLE cumsum test inlining agnostic
sunnycase Jul 27, 2026
390943a
Merge branch 'main' into fix/tle_cumsum
sunnycase Jul 28, 2026
201e52f
Fix TleToLLVM generated header dependency
sunnycase Jul 28, 2026
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
6 changes: 5 additions & 1 deletion python/test/tle/unit/test_tle_cumsum.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,11 @@ def _pick_expected_dtype(input_dtype: torch.dtype) -> torch.dtype:


def _make_input(dtype: torch.dtype, block: int) -> torch.Tensor:
if dtype in (torch.float16, torch.float32, torch.bfloat16):
if dtype == torch.float16:
# Keep every prefix sum exactly representable in fp16 so the
# correctness test is independent of parallel scan accumulation order.
return torch.randint(-2, 3, (block, ), device="cuda", dtype=torch.int32).to(dtype)
if dtype in (torch.float32, torch.bfloat16):
return torch.randn((block, ), device="cuda", dtype=dtype)
if dtype == torch.int8:
return torch.randint(-32, 32, (block, ), device="cuda", dtype=dtype)
Expand Down
Loading