Skip to content

Reject sparse MMA K tiles with unhandled tails - #3367

Open
dajiaohuang wants to merge 1 commit into
tile-ai:mainfrom
dajiaohuang:fix/2605-sparse-mma-k-atom
Open

dajiaohuang wants to merge 1 commit into
tile-ai:mainfrom
dajiaohuang:fix/2605-sparse-mma-k-atom

Conversation

@dajiaohuang

@dajiaohuang dajiaohuang commented Oct 1, 2026 •

Copy link
Copy Markdown

Summary

  • Reject an SM80-style sparse MMA K tile whose extent is not divisible by the dtype-specific MMA K atom before the lowering floors its loop trip count.
  • Add an SM90 compile regression for total K=192 with int8 block_K=96, where each tile would drop 32 elements even though total K is divisible by 64.

Fixes #2605.

Validation

  • Ruff check/format, Python syntax compilation, and git diff --check pass.
  • Focused pytest collection is blocked by the local Windows Python 3.13 environment: Home Assistant's pytest plugin imports cntl, which is unavailable on Windows. The TileLang test did not run.
  • No local CUDA compiler/runtime or native TileLang build is available.

Summary

  • GemmSPMMA.lower now raises ValueError when a sparse MMA K tile is not divisible by its dtype-specific K atom. This prevents the fallback loop from silently dropping the tile’s K tail.
  • Added an SM90 regression test for int8 with total K=192 and block_K=96. It checks that compilation rejects the tile, although total K is divisible by the 64-element atom.

Validation

  • Ruff checks and formatting, Python syntax compilation, and git diff --check reportedly pass.
  • The regression test did not run. Pytest collection was blocked because the Home Assistant pytest plugin imports fcntl, which is unavailable in the local Windows Python 3.13 environment.
  • No local CUDA compiler, runtime, or native TileLang build was available.

@github-actions

github-actions Bot commented Oct 1, 2026

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

@coderabbitai

coderabbitai Bot commented Oct 1, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Repository: tile-ai/tilelang/.coderabbit.yaml

Review profile: CHILL

Plan: Advanced

Run ID: 9c01b8f2-9423-42e6-9371-886e532f40c1

📥 Commits

Reviewing files that changed from the base of the PR and between 994b44e and 7ae638a.

📒 Files selected for processing (2)
  • testing/python/issue/test_tilelang_issue_2605.py
  • tilelang/cuda/op/gemm_sp/gemm_sp_mma.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 4 remain after this review.


📝 Walkthrough

Walkthrough

The sparse MMA lowering now rejects K tile sizes that are not divisible by the sparse MMA K atom. A CUDA regression test checks that a kernel with K=192 and block_K=96 raises the expected ValueError.

Changes

Sparse GEMM K validation

Layer / File(s) Summary
Reject misaligned K tiles
tilelang/cuda/op/gemm_sp/gemm_sp_mma.py, testing/python/issue/test_tilelang_issue_2605.py
The lowering raises ValueError when K is not divisible by the sparse MMA K atom. The regression test checks compilation with K=192 and block_K=96.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix · Severity of issue fixed: Medium

Merge Risk: ⚪ Minimal · up to 7ae63

Sparse GEMM kernels whose K tile is not a multiple of the MMA K atom now fail at compile time with a clear error instead of silently dropping the K tail. The change is small and well-scoped. The new test could not be run locally without a CUDA environment, so it should be confirmed in CI.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 2 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and concisely describes the main change: rejecting sparse MMA K tiles that have unhandled tails.
Linked Issues check ✅ Passed PR #3367 satisfies the coding objective in issue #2605. GemmSPMMA.lower obtains the dtype-specific micro_size_k and raises ValueError when the K tile is not divisible by that atom. The existing …
Out of Scope Changes check ✅ Passed The changes are limited to the sparse MMA lowering and a regression test for issue #2605. The test directly covers the per-tile K-tail condition. No unrelated product behavior or unrelated files are c…
  • Fix all pre-merge checks with AI
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create a new PR
  • Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Autopilot is currently an internal CodeRabbit preview.


Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant