[BugFix][CPU] Compute scalar GEMM products in accum dtype instead of input dtype

[Dino1844]

Summary

GemmScalar — the CPU/LLVM fallback for T.gemm — evaluates each scalar multiply in the input dtype and only then widens the already-rounded product to the accumulator dtype:

C_buf[c0 + i, c1 + j] += T.cast(
    A_buf[...] * B_buf[...],   # fp16 x fp16 -> fp16 (rounded)
    accum_dtype,
)

For fp16 inputs this diverges from CUDA/ROCm mma, which computes fp16 products exactly (the 22-bit result fits in fp32's 24-bit mantissa) and accumulates in fp32. It also diverges from TileLang's own CUDA FMA fallback (tilelang/cuda/op/gemm/gemm_fma.py), whose docstring documents the correct convention: "Casts widen operands to accum_dtype before multiplication so narrow input dtypes stay numerically sound (e.g. BF16 to FP32 accumulation on Volta). This path favors correctness and coverage over peak throughput."

Problem

Rounding each product to fp16 before widening has two consequences:

1. Accuracy. Every product carries an extra ~2^-11 relative rounding (fp32 accumulation is ~2^-24), which dominates the error for realistic K. Measured on normal-range inputs: the fp16 GEMM deviates from exact-product accumulation by up to ~1e-2 absolute (K=256), whereas exact products are accurate to fp32-accumulation level (~1e-5).

2. Overflow -> NaN/Inf. Products with |a·b| > 65504 (fp16 max) overflow to ±inf; mixed-sign infs then produce NaN inside the fp32 accumulator. Measured with |a|,|b| ~ 600: 100% of outputs were NaN/Inf before the fix, while exact products stay finite (they fit in fp32).

Fix

Cast both operands to accum_dtype before multiplying (exact for fp16/bf16/fp8 products, all of which fit exactly in fp32):

C_buf[c0 + i, c1 + j] += T.cast(A_buf[...], accum_dtype) * T.cast(B_buf[...], accum_dtype)

This matches mma semantics and the convention already established in gemm_fma.py.

Verification

fp16 T.gemm on target="c" (compiled and executed): normal-range inputs match an exact-product fp64 reference to fp32-accumulation level (max abs err 2e-5 vs ~1e-2 before); large-magnitude inputs (products > 65504) no longer produce NaN/Inf. Existing CPU GEMM tests (test_tilelang_cpu_tgemm.py, test_tilelang_cpu_gemm.py) pass.

Summary

• Fixed CPU/LLVM scalar T.gemm accumulation.
• Cast both operands to accum_dtype before multiplication.
• Prevented fp16 product rounding and fp16 overflow to Inf or NaN.
• Aligned CPU behavior with CUDA/ROCm MMA and existing CUDA FMA fallback semantics.
• Improved fp16 GEMM accuracy and preserved existing CPU GEMM test coverage.

[github-actions]

👋 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]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 85cb4725-c6c7-46f5-89ae-500077507f7f

📥 Commits

Reviewing files that changed from the base of the PR and between 12dbf3e and cd93ca1.

📒 Files selected for processing (1)

• tilelang/cpu/op/gemm/gemm_scalar.py

---

📝 Walkthrough

Walkthrough

The scalar GEMM kernel now casts operands A and B to accum_dtype before multiplication. This replaces casting the multiplication result after computation.

Changes

Scalar GEMM

Layer / File(s)	Summary
Cast operands before multiplication tilelang/cpu/op/gemm/gemm_scalar.py	The GEMM computation casts both operands to accum_dtype before multiplication.

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

Suggested reviewers: leiwang1999

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically describes the main fix: computing scalar GEMM products in the accumulator data type on CPU.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.}

[Dino1844]

Hi @LeiWang1999 and maintainers,

While auditing TileLang I focused on the CPU backend, and wanted to confirm the current support status and direction before investing effort, to avoid overlapping with community work and wasting your review time.

Observation: The CPU backend (target="c" / target="llvm") is marked Experimental in the README. From my testing:

1. Some core ops are unimplemented: T.cumsum, T.atomic_add, finalize_reducer, gemm_sp fail with Unresolved call (CUDA has ~14 ops, CPU 5);
2. Vectorization is incomplete: T.exp/min/max/cast etc. fail to compile in T.Parallel, vec_type only implements +-*/;
3. Some semantics diverge from CUDA/torch, e.g. NaN / -0.0 in min/max.

Questions:

1. Is CPU an actively developed direction? Any roadmap or in-progress work?
2. If contributions are welcome, what to prioritize — unimplemented ops, or vectorization/math support?

Thanks for your time! If CPU isn't a near-term priority, please let me know and I'll adjust accordingly.