[zheyang0825]

📌 Description

Replaces the SM90 K1 FP16 Neumann inverse with two 8×8 FP32 forward substitutions and a BF16/FP32-accumulating block merge. The BF16 workspace ABI is unchanged. A repeated-token regression test covers the near-collinear case that exposed the precision loss.

🔍 Related Issues

• Follows the precision fix discussed in MoonshotAI/FlashKDA#34.

🚀 Pull Request Checklist

✅ Pre-commit Checks

• I have installed pre-commit.
• I have installed the hooks with pre-commit install.
• I have run pre-commit run --all-files; all hooks pass.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing.
• SM90 L20X: pytest -q tests/test_kda_sm90_prefill_vs_fla.py -m kda_full — 12 passed.

⚡ Performance

SM90 L20X, BF16, H=64, fixed B=1/2 and T=512–16384; two reversed-order runs, 25 warmups and 100 iterations:

Metric	FP16 Neumann	FP32 forward substitution
Repeated-token output error ratio	0.00868	0.00460
Repeated-token final-state error ratio	0.01261	0.00815
End-to-end latency	baseline	+1.45% geomean (-0.16% to +2.12%)

Reviewer Notes

Please focus on the one-warp block inverse and the shared-memory lifetime: FP32 L is independent because warp 1 may still read k_inv, while the later P/M merge buffers reuse dead storage after the CTA barrier.

Summary by CodeRabbit

• Bug Fixes
  • Improved numerical precision for SM90 prefill when processing repeated tokens, including output and final-state calculations. This update strengthens accuracy in cases with highly similar query and key vectors and near-unit decay.

[coderabbitai]

No actionable comments were generated in the recent review.

Walkthrough

K1 replaces its fp16 Neumann-series inverse with fp32 forward substitution on two 8×8 blocks, then combines them using bf16 operands and fp32 accumulation. The kernel updates shared-memory staging, and a repeated-token test compares cuLA prefill results with the FLA baseline.

Changes: SM90 KDA inverse precision

Layer / File(s)	Summary
FP32 inverse and shared-memory staging cula/ops/kda/sm90/k1.py	K1 computes the 16×16 inverse through fp32 forward substitution on two 8×8 diagonal blocks, then combines the blocks using bf16 operands and fp32 MMA accumulation. The kernel stores the masked, beta-weighted matrix in fp32 shared memory and uses separate merge buffers.
Kernel integration and precision regression cula/ops/kda/sm90/k1.py, tests/test_kda_sm90_prefill_vs_fla.py	Warp 0 calls the new inverse routine. The added repeated-token test checks output and transposed final-state error ratios against FLA, with limits of 0.005 and 0.01.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Bug fix

Possibly related PRs

• inclusionAI/cuLA#133: Routes SM90 KDA prefill through the K1+K2 pipeline that uses the inverse implementation changed here.

Suggested reviewers: cherhh

Merge Risk: Minimal · up to d7814

This change improves the numerical precision of the SM90 KDA prefill inverse with a small reported latency cost. No concrete merge-blocking issue remains, and the new repeated-token regression test covers the precision case that motivated the fix.

Pre-merge checks: 4

Check name	Status	Explanation
Title check		The title clearly identifies the SM90 inverse precision fix, which is the main change.
Description check		The description covers the change, related issue, checklist, tests, performance results, and reviewer notes.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Finishing Touches

Generate docstrings

• Commit to this branch
• Create a new PR

Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

Recent review info

Run configuration

• Source mode: Auto
• Repository: inclusionAI/cuLA
• Branch: fix/flashkda-fp32-inverse
• File: .coderabbit.yaml
• Commit: d78145f
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Advanced
• Run ID: 17ebadc2-712c-4202-8b15-9439eb3a10d2

Commits

Reviewing files that changed from the base of the PR and between 9055ef2 and d78145f.

Files selected for processing: 2

• cula/ops/kda/sm90/k1.py
• tests/test_kda_sm90_prefill_vs_fla.py

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