feat: add gemm_n16384_k3072 definition and reference test#374
Open
yyihuang wants to merge 2 commits into
Open
Conversation
Add GEMM N=16384, K=3072 definition and reference test. Applies to Llama 3.2 3B mlp.gate_up_proj (fused gate+up: 2×8192=16384) and MiniMax M2 MoE FC1/W13 (same hidden_size=3072, 2×intermediate=16384). - Add definition JSON with float16 dtype - Add reference test (all 6 M-value configs pass, zero numerical difference) - Update MiniMax M2 coverage: 14/15 → 15/16 Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Contributor
📝 WalkthroughWalkthroughThis PR introduces a new GEMM kernel definition ( Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a new GEMM operation definition (gemm_n16384_k3072) for Llama 3.2 3B and MiniMax M2 models, updates the model coverage documentation, and adds a reference test script. Feedback was provided to ensure the test script handles environments without CUDA gracefully by returning True instead of False, preventing unnecessary CI failures on non-GPU nodes.
Comment on lines
+48
to
+50
| if device == "cpu": | ||
| print("WARNING: CUDA not available, skipping test") | ||
| return False |
Contributor
There was a problem hiding this comment.
The current implementation returns False when CUDA is not available. This causes the main() function to report a failure and exit with code 1, which will break CI pipelines running on non-GPU nodes. Since this is a reference test script, it should ideally skip gracefully or return True if the environment doesn't support the required hardware.
Suggested change
| if device == "cpu": | |
| print("WARNING: CUDA not available, skipping test") | |
| return False | |
| if device == "cpu": | |
| print("WARNING: CUDA not available, skipping test") | |
| return True |
MiniMax M2 uses MoE with intermediate_size=1536, so gate_up output is 2*1536=3072, not 16384. This GEMM belongs to Llama 3.2 3B (intermediate=8192). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
flashinfer_trace/tests/references/test_gemm_n16384_k3072.py (1)
78-78: Remove unnecessaryfprefix from the static string.Line 78 uses an f-string with no variable interpolation, triggering Ruff's F541 lint check.
Proposed fix
- print(f"\nOutput tensor comparison:") + print("\nOutput tensor comparison:")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@flashinfer_trace/tests/references/test_gemm_n16384_k3072.py` at line 78, The print call uses an unnecessary f-string with no interpolation—change the print invocation that currently reads print(f"\nOutput tensor comparison:") to a regular string literal (print("\nOutput tensor comparison:")) to satisfy Ruff's F541; update the print statement in the test_gemm_n16384_k3072.py test where the "Output tensor comparison:" message is emitted.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docs/model_coverage.mdx`:
- Around line 684-690: The table incorrectly lists `gemm_n16384_k3072` under
MiniMax M2; remove that row (or move the `gemm_n16384_k3072` entry to the Llama
3.2 3B section) and update the MiniMax M2 coverage count from "15 / 16" to the
correct value after removing the misplaced definition so the coverage math
reflects the actual definitions present; ensure any reference to `MoE gate /
topk / experts`, `top_k_sampling_from_probs_v200064`,
`top_k_top_p_sampling_from_probs_v200064`, and
`top_p_sampling_from_probs_v200064` remain unchanged.
In `@flashinfer_trace/tests/references/test_gemm_n16384_k3072.py`:
- Around line 47-50: The test currently treats "CUDA not available" as a failure
by returning False; change the behavior to mark the test as skipped instead:
when device is set with device = "cuda" if torch.cuda.is_available() else "cpu"
and you detect device == "cpu", either call pytest.skip("CUDA not available,
skipping test") or return True (so it isn't treated as a failure) and update the
identical blocks around the 104-122 region similarly; adjust imports if you use
pytest.skip (import pytest) and keep the same warning message text.
- Around line 10-13: Update the top docstring attribution in
test_gemm_n16384_k3072.py: replace the incorrect model name "MiniMax M2" with
the correct mapping "Llama 3.2 3B" in the descriptive comment that currently
reads "This corresponds to MiniMax M2 MoE FC1/gate-up projection (W13)" so the
docstring accurately reflects that the test/definition is mapped to Llama 3.2
3B.
---
Nitpick comments:
In `@flashinfer_trace/tests/references/test_gemm_n16384_k3072.py`:
- Line 78: The print call uses an unnecessary f-string with no
interpolation—change the print invocation that currently reads print(f"\nOutput
tensor comparison:") to a regular string literal (print("\nOutput tensor
comparison:")) to satisfy Ruff's F541; update the print statement in the
test_gemm_n16384_k3072.py test where the "Output tensor comparison:" message is
emitted.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 69bb4b16-310d-4feb-a17a-8ec4ace07085
📒 Files selected for processing (3)
docs/model_coverage.mdxflashinfer_trace/definitions/gemm/gemm_n16384_k3072.jsonflashinfer_trace/tests/references/test_gemm_n16384_k3072.py
Comment on lines
+684
to
+690
| | `gemm_n16384_k3072` | gemm (MoE FC1/W13) | 🟡 | | ||
| | MoE gate / topk / experts | moe | — | | ||
| | `top_k_sampling_from_probs_v200064` | sampling | 🟡 | | ||
| | `top_k_top_p_sampling_from_probs_v200064` | sampling | 🟡 | | ||
| | `top_p_sampling_from_probs_v200064` | sampling | 🟡 | | ||
|
|
||
| **Coverage**: 14 / 15 definitions present. Workloads not yet collected. | ||
| **Coverage**: 15 / 16 definitions present. Workloads not yet collected. |
Contributor
There was a problem hiding this comment.
Fix incorrect MiniMax M2 attribution and coverage count.
Line 684 adds gemm_n16384_k3072 under MiniMax M2, but this shape is attributed to Llama 3.2 3B in this PR context; keeping it here also makes the Line 690 coverage math incorrect for MiniMax M2.
📝 Proposed fix
| `gemm_n256_k3072` | gemm (MoE gate) | 🟡 |
-| `gemm_n16384_k3072` | gemm (MoE FC1/W13) | 🟡 |
| MoE gate / topk / experts | moe | — |
| `top_k_sampling_from_probs_v200064` | sampling | 🟡 |
| `top_k_top_p_sampling_from_probs_v200064` | sampling | 🟡 |
| `top_p_sampling_from_probs_v200064` | sampling | 🟡 |
-**Coverage**: 15 / 16 definitions present. Workloads not yet collected.
+**Coverage**: 14 / 15 definitions present. Workloads not yet collected.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| | `gemm_n16384_k3072` | gemm (MoE FC1/W13) | 🟡 | | |
| | MoE gate / topk / experts | moe | — | | |
| | `top_k_sampling_from_probs_v200064` | sampling | 🟡 | | |
| | `top_k_top_p_sampling_from_probs_v200064` | sampling | 🟡 | | |
| | `top_p_sampling_from_probs_v200064` | sampling | 🟡 | | |
| **Coverage**: 14 / 15 definitions present. Workloads not yet collected. | |
| **Coverage**: 15 / 16 definitions present. Workloads not yet collected. | |
| | MoE gate / topk / experts | moe | — | | |
| | `top_k_sampling_from_probs_v200064` | sampling | 🟡 | | |
| | `top_k_top_p_sampling_from_probs_v200064` | sampling | 🟡 | | |
| | `top_p_sampling_from_probs_v200064` | sampling | 🟡 | | |
| **Coverage**: 14 / 15 definitions present. Workloads not yet collected. |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@docs/model_coverage.mdx` around lines 684 - 690, The table incorrectly lists
`gemm_n16384_k3072` under MiniMax M2; remove that row (or move the
`gemm_n16384_k3072` entry to the Llama 3.2 3B section) and update the MiniMax M2
coverage count from "15 / 16" to the correct value after removing the misplaced
definition so the coverage math reflects the actual definitions present; ensure
any reference to `MoE gate / topk / experts`,
`top_k_sampling_from_probs_v200064`, `top_k_top_p_sampling_from_probs_v200064`,
and `top_p_sampling_from_probs_v200064` remain unchanged.
Comment on lines
+10
to
+13
| This corresponds to MiniMax M2 MoE FC1/gate-up projection (W13): | ||
| Input: hidden_size = 3072 | ||
| Output: 2 * intermediate_size = 2 * 8192 = 16384 | ||
|
|
Contributor
There was a problem hiding this comment.
Update the model attribution in the docstring.
Line 10 still says MiniMax M2, but this test/definition is mapped to Llama 3.2 3B in this PR. Please update the text to avoid future remapping mistakes.
📝 Proposed fix
- This corresponds to MiniMax M2 MoE FC1/gate-up projection (W13):
+ This corresponds to Llama 3.2 3B MLP gate-up projection:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@flashinfer_trace/tests/references/test_gemm_n16384_k3072.py` around lines 10
- 13, Update the top docstring attribution in test_gemm_n16384_k3072.py: replace
the incorrect model name "MiniMax M2" with the correct mapping "Llama 3.2 3B" in
the descriptive comment that currently reads "This corresponds to MiniMax M2 MoE
FC1/gate-up projection (W13)" so the docstring accurately reflects that the
test/definition is mapped to Llama 3.2 3B.
Comment on lines
+47
to
+50
| device = "cuda" if torch.cuda.is_available() else "cpu" | ||
| if device == "cpu": | ||
| print("WARNING: CUDA not available, skipping test") | ||
| return False |
Contributor
There was a problem hiding this comment.
Don’t treat “skipped (no CUDA)” as failed tests.
Line 49 says “skipping test”, but Line 50 returns False, which is then counted as a failure and triggers exit(1) in main().
🛠️ Proposed fix
def test_correctness(M=128, atol=1e-2, rtol=1e-2):
@@
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cpu":
print("WARNING: CUDA not available, skipping test")
- return False
+ return None
@@
def main():
@@
- passed = 0
+ passed = 0
+ skipped = 0
@@
for M in test_M_values:
try:
- if test_correctness(M):
+ result = test_correctness(M)
+ if result is True:
passed += 1
+ elif result is None:
+ skipped += 1
except Exception as e:
print(f"✗ Test failed with exception: {str(e)}")
@@
- print(f"Summary: {passed}/{total} tests passed")
+ effective_total = total - skipped
+ print(f"Summary: {passed}/{effective_total} tests passed, {skipped} skipped")
@@
- if passed == total:
+ if skipped == total:
+ print("⚠ All tests skipped (CUDA not available).")
+ return
+ if passed == effective_total:
print("✓ All tests passed!")Also applies to: 104-122
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@flashinfer_trace/tests/references/test_gemm_n16384_k3072.py` around lines 47
- 50, The test currently treats "CUDA not available" as a failure by returning
False; change the behavior to mark the test as skipped instead: when device is
set with device = "cuda" if torch.cuda.is_available() else "cpu" and you detect
device == "cpu", either call pytest.skip("CUDA not available, skipping test") or
return True (so it isn't treated as a failure) and update the identical blocks
around the 104-122 region similarly; adjust imports if you use pytest.skip
(import pytest) and keep the same warning message text.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add
gemm_n16384_k3072definition and reference test.This GEMM shape (N=16384, K=3072) appears in:
mlp.gate_up_proj(fused gate+up: 2 × 8192 = 16384)gemm_n16384_k3072Reference Test Output
Coverage Update
MiniMax M2: 14/15 → 15/16 definitions present
PR2 (HuggingFace trace)
Link to be added after PR2 opens.
Summary by CodeRabbit
New Features
Documentation
Tests