feat: add gemm_n3072_k8192 definition

[flashinfer-bot]

$## Summary

Adds the gemm_n3072_k8192 kernel definition and its reference test for Llama 3.2 3B at TP=1.

This GEMM corresponds to the MLP down-projection: intermediate=8192 → hidden=3072, i.e. C = A @ B.T with A: (M, 8192), B: (3072, 8192), C: (M, 3072) in float16.

Kernel details

Field	Value
Name	gemm_n3072_k8192
Op type	gemm
Model	Llama 3.2 3B
Layer	mlp.down_proj
TP	1
N	3072
K	8192
M	var (per-token batch)
Dtype	float16
Tags	status:reference, fi_api:gemm, model:llama-3.2-3b, tp:1

Reference test

$ pytest flashinfer_trace/tests/references/test_gemm_n3072_k8192.py -v
============================= test session starts ==============================
collected 1 item

flashinfer_trace/tests/references/test_gemm_n3072_k8192.py::test_correctness PASSED [100%]

Testing GEMM N=3072, K=8192: M=128
A shape: torch.Size([128, 8192]), dtype: torch.float16
B shape: torch.Size([3072, 8192]), dtype: torch.float16
Max absolute difference: 0.000000e+00
Max relative difference: 0.000000e+00
✓ PASSED: Outputs match within tolerance (atol=0.01, rtol=0.01)
========================= 1 passed in 6.04s =========================

The reference implementation (torch.matmul(A, B.T)) matches F.linear(A, B) exactly.

Files

• flashinfer_trace/definitions/gemm/gemm_n3072_k8192.json — updated tags and description
• flashinfer_trace/tests/references/test_gemm_n3072_k8192.py — new reference test

PR2 (HuggingFace trace)

Link to be added after PR2 opens.

[coderabbitai]

Warning

Rate limit exceeded

@flashinfer-bot has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 44 minutes and 36 seconds before requesting another review.

Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 44 minutes and 36 seconds.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: edf220ff-1e8a-4859-9602-cea2bea5e9f9

📥 Commits

Reviewing files that changed from the base of the PR and between e489a24 and 7cb3df4.

📒 Files selected for processing (2)

• flashinfer_trace/definitions/gemm/gemm_n3072_k8192.json
• flashinfer_trace/tests/references/test_gemm_n3072_k8192.py

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/def-gemm_n3072_k8192

---

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

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request updates the GEMM definition for N=3072, K=8192 to include TP=1 metadata and introduces a reference test script. Feedback focuses on improving the test harness by adding explicit assertions to ensure correctness failures are properly reported and adjusting the handling of missing CUDA environments to avoid false failure reports.

[gemini-code-assist]

The test_correctness function calculates whether the outputs are within tolerance but does not use an assert statement. When running this file with a test runner like pytest, the test will be marked as passed even if the outputs do not match the reference, as long as no exception is raised. Adding an explicit assertion ensures that mismatches are correctly caught and reported as failures in both the standalone script and the test runner.

Suggested change

	output_close = torch.allclose(ref_f32, fi_f32, atol=atol, rtol=rtol)
	if output_close:
	print(f"\n✓ PASSED: Outputs match within tolerance (atol={atol}, rtol={rtol})")
	else:
	print(f"\n✗ FAILED: Outputs differ beyond tolerance (atol={atol}, rtol={rtol})")
	return output_close
	output_close = torch.allclose(ref_f32, fi_f32, atol=atol, rtol=rtol)
	if output_close:
	print(f"\n✓ PASSED: Outputs match within tolerance (atol={atol}, rtol={rtol})")
	else:
	print(f"\n✗ FAILED: Outputs differ beyond tolerance (atol={atol}, rtol={rtol})")
	assert output_close, f"Outputs differ beyond tolerance (atol={atol}, rtol={rtol})"
	return output_close

[gemini-code-assist]

Returning False when CUDA is not available causes the main() function to treat the skip as a test failure, leading to a non-zero exit code for the script. If the intention is to skip the test in environments without a GPU, it should return True (indicating no failure occurred) or use a proper skip mechanism like pytest.skip() so that the test suite doesn't incorrectly report a failure.

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