Fix potential out-of-bound access in int8_mm.py

[mark14wu]

This PR fixes a bug in _scaled_int8_mm_kernel allowing unmasked tl.load() calls when K was even, but not necessarily a multiple of BLOCK_K, causing out-of-bounds reads. The fix changes the check to K % BLOCK_K == 0, ensuring unmasked loads occur only when the last block is fully valid.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1751

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 62f7334 with merge base 1c76736 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[mark14wu]

cc @Jokeren

[Jokeren]

@mark14wu A clearer fix can be done using @triton.heuristics

[gau-nernst]

Thank you for the interest! Didn't know there are people looking at this code 😄.

I copied this code from PyTorch Inductor though. Does that mean Inductor is susceptible to this bug too? (or inductor allocates padded tensors so it's not an issue?)

Like @Jokeren said, @triton.heuristic sounds like a good idea. We would want to keep a tl.constexpr there so that we don't do if for every iteration.

Another question. It seems like we might do out-of-bound reads for M and N dims too?

[Jokeren]

I copied this code from PyTorch Inductor though. Does that mean Inductor is susceptible to this bug too?

Yes indeed, but our tool, a triton sanitizer, cannot work for pytorch inductor at this moment. Hao may want to explore this direction when he starts intern at meta this summer.

It seems like we might do out-of-bound reads for M and N dims too?

M and N dims are safe because we have protections with rm%M and rn%N.

@mark14wu please go ahead refactor the code with triton.heuristic

[drisspg]

There are ways to turn inductor code into triton only repos

[mark14wu]

Now refactored with triton.heuristics.

[gau-nernst]

@mark14wu You can run ruff formatting then we can merge. Thank you!

[mark14wu]

@mark14wu You can run ruff formatting then we can merge. Thank you!

Done. Please merge, thanks!