feat: FP32 dtype output for BF16 matmuls (CUTLASS & cuDNN) by raayandhar · Pull Request #2644 · flashinfer-ai/flashinfer

raayandhar · 2026-02-26T22:25:49Z

📌 Description

Adds support for FP32 dtype output for mm_bf16 and bmm_bf16 for the CUTLASS and cuDNN backends. I'm not familiar enough with the TGV kernel to know if / how to support it for that backend.

🔍 Related Issues

#2624

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

I have installed pre-commit by running pip install pre-commit (or used your preferred method).
I have installed the hooks with pre-commit install.
I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

Tests have been added or updated as needed.
All tests are passing (unittest, etc.).

Reviewer Notes

Summary by CodeRabbit

New Features
- BF16-based matrix ops (mm_bf16, bmm_bf16) now allow float32 outputs in addition to bfloat16 and float16; supported across applicable backends.
Tests
- Tests extended to cover float32 outputs for BF16/GEMM operations.
Documentation
- User-facing docs and validation messages updated to list bf16, fp16, fp32 as valid output dtypes.

coderabbitai · 2026-02-26T22:26:05Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: b005d881-628c-492e-a2c3-bbf68a10d676

📥 Commits

Reviewing files that changed from the base of the PR and between 00c9897 and 78061e4.

📒 Files selected for processing (6)

benchmarks/routines/gemm.py
csrc/bf16_gemm_cutlass.cu
flashinfer/gemm/gemm_base.py
flashinfer/jit/gemm/core.py
tests/gemm/test_bmm_bf16.py
tests/gemm/test_mm_bf16.py

🚧 Files skipped from review as they are similar to previous changes (2)

tests/gemm/test_bmm_bf16.py
flashinfer/gemm/gemm_base.py

📝 Walkthrough

Walkthrough

Adds fp32 (torch.float32) as a supported output dtype for BF16 GEMM across CUTLASS runtime, Python validation/JIT, benchmarks, and tests; updates dispatch, dtype mappings, docs, and tests to include fp32.

Changes

Cohort / File(s)	Summary
CUTLASS Backend Implementation `csrc/bf16_gemm_cutlass.cu`	Adds float dispatch branch, instantiates `CutlassBf16GemmRunner<float>`, and updates error message to include fp32.
Python API and Validation `flashinfer/gemm/gemm_base.py`	Accepts fp32 in `_validate_bf16_output_dtype`, maps `torch.float32` to cuDNN `FLOAT`, and updates docstrings for `mm_bf16`/`bmm_bf16`.
JIT Kernel Generation `flashinfer/jit/gemm/core.py`	Adds `"float"` to `dtype_list` for BF16 Cutlass kernel generation (generates fp32 kernel variant).
Tests & Benchmarks `benchmarks/routines/gemm.py`, `tests/gemm/test_mm_bf16.py`, `tests/gemm/test_bmm_bf16.py`	Extends test/benchmark parameterization to include `torch.float32`, and updates skip logic and error assertions accordingly.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant PythonAPI as Python API (gemm_base)
    participant JIT as JIT (gen_gemm_sm100_module_cutlass_bf16)
    participant Native as Native/C++ (bf16_gemm_cutlass.cu)
    participant GPU

    User->>PythonAPI: Call mm_bf16/bmm_bf16(out_dtype=float32)
    PythonAPI->>PythonAPI: validate out_dtype (accept fp32)
    PythonAPI->>JIT: request kernel variant (dtype_list includes "float")
    JIT->>Native: load/compile kernel variant (fp32 variant)
    Native->>Native: dispatch runGemm<float> branch
    Native->>GPU: execute CUTLASS kernel (fp32 output path)
    GPU-->>Native: results
    Native-->>PythonAPI: return tensor (fp32)
    PythonAPI-->>User: deliver output tensor

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related issues

[Feature] Support FP32 output for bf16 mm #2624: Implements FP32 output support for BF16 GEMM — aligns with this PR's objective of adding fp32 output support.

Possibly related PRs

feat: BF16 GEMM using CUTLASS backend for SM100 #2070: Foundational BF16 CUTLASS GEMM implementation that this PR extends with fp32 output support.
feat: BF16 GEMM using cuDNN backend #2376: Modifies BF16 GEMM output-dtype validation and backend handling in gemm_base.py.
feat: BF16 GEMM benchmarking support #2525: Related changes to benchmarks and tests for BF16 GEMM fp32 support.

Suggested labels

op: gemm

Suggested reviewers

yongwww
nvmbreughe
aleozlx
djmmoss
jimmyzho
jiahanc
cyx-6
yzh119

Poem

🐰 A soft hop for dtype delight,
bfloat, half, and now float take flight,
Kernels compiled and tests all sing,
CUTLASS hums — three types in spring! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 45.45% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The pull request title accurately and concisely describes the main change: adding FP32 dtype output support for BF16 matmul operations across CUTLASS and cuDNN backends, which aligns with the comprehensive changes across multiple files.
Description check	✅ Passed	The description adequately covers the main objective, provides a related issue link, and confirms that pre-commit checks and tests have been completed. All major required template sections are addressed.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

📝 Coding Plan

Generate coding plan for human review comments

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 and usage tips.}

Tip

CodeRabbit can generate a title for your PR based on the changes with custom instructions.

Set the reviews.auto_title_instructions setting to generate a title for your PR based on the changes in the PR with custom instructions.

raayandhar · 2026-02-26T22:26:07Z

(flashinfer) root@36978a460ca6:/flashinfer# pytest tests/gemm/test_mm_bf16.py
========================================================= test session starts ==========================================================
platform linux -- Python 3.12.3, pytest-9.0.2, pluggy-1.6.0
rootdir: /flashinfer
configfile: pytest.ini
collected 1620 items

tests/gemm/test_mm_bf16.py sssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss [  6%]
ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss [ 14%]
ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss [ 22%]
ssssssssssssssssssssssssssssssssssssssssssssssss................................................................................ [ 29%]
.......................................................sssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss [ 37%]
ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss [ 45%]
ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss [ 53%]
ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss.................................................... [ 61%]
................................................................................................................................ [ 69%]
ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss...................................... [ 77%]
.......ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss............................... [ 85%]
..............ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss........................ [ 93%]
.....................ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss                  [100%]

============================================ 450 passed, 1170 skipped in 200.02s (0:03:20) =============================================
(flashinfer) root@36978a460ca6:/flashinfer# pytest tests/gemm/test_bmm_bf16.py
========================================================= test session starts ==========================================================
platform linux -- Python 3.12.3, pytest-9.0.2, pluggy-1.6.0
rootdir: /flashinfer
configfile: pytest.ini
collected 96 items

tests/gemm/test_bmm_bf16.py ................................................................................................     [100%]

========================================================== 96 passed in 9.07s ==========================================================

gemini-code-assist · 2026-02-26T22:26:28Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces support for FP32 as an output data type for BF16 matrix multiplication operations (mm_bf16 and bmm_bf16) when utilizing the CUTLASS and cuDNN backends. This enhancement allows users to perform BF16 computations and receive results in higher precision, addressing a common requirement for certain deep learning workloads. The changes span across core C++ CUDA implementations, Python API definitions, JIT compilation, and associated benchmarks and tests.

Highlights

FP32 Output Support: Enabled FP32 as an output data type for BF16 matrix multiplication operations (mm_bf16 and bmm_bf16) in both CUTLASS and cuDNN backends.
Backend Integration: Integrated FP32 output support into the C++ CUDA kernels, Python API definitions, and JIT compilation logic for BF16 GEMM operations.
Testing and Benchmarking: Updated benchmark scripts and extended unit tests to validate the new FP32 output functionality for BF16 GEMM operations.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

benchmarks/routines/gemm.py
- Allowed float32 as a supported output dtype for testMmBf16 and testBmmBf16 functions.
csrc/bf16_gemm_cutlass.cu
- Instantiated CutlassBf16GemmRunner for float type.
- Added a case to bf16_bmm_impl to handle float32_code for GEMM execution.
- Updated the error message for unsupported output dtypes to include fp32.
flashinfer/gemm/gemm_base.py
- Updated the error message for the TGV backend to mention cuDNN as an alternative.
- Modified docstrings for mm_bf16 and bmm_bf16 to reflect fp32 support for out tensor and out_dtype.
- Updated the _validate_bf16_output_dtype function to include torch.float32 as a valid output dtype.
- Added mapping for torch.float32 to cudnn.data_type.FLOAT in _torch_data_type_to_cudnn_data_type.
flashinfer/jit/gemm/core.py
- Included float in the dtype_list for generating CUTLASS BF16 GEMM kernels.
tests/gemm/test_bmm_bf16.py
- Added torch.float32 to the res_dtype parameter list for test_bmm_bf16.
tests/gemm/test_mm_bf16.py
- Added torch.float32 to the res_dtype parameter list for test_mm_bf16.
- Updated the skip condition for the TGV backend to correctly handle all non-bfloat16 result dtypes.

Activity

Pre-commit checks have been successfully run and passed.
Tests have been added/updated and are all passing.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request adds support for FP32 output data type for BF16 matrix multiplications (mm_bf16 and bmm_bf16) for both CUTLASS and cuDNN backends. The changes are comprehensive, touching the C++ implementation, Python bindings, JIT compilation, documentation, and tests. The implementation correctly adds template instantiations and dispatch logic for float output in the CUTLASS backend, and updates the type mapping for cuDNN. The tests are properly extended to cover the new FP32 output functionality. The changes are well-executed and appear correct.

bkryu · 2026-02-27T19:31:39Z

/bot run

flashinfer-bot · 2026-02-27T19:32:32Z

GitLab MR !353 has been created, and the CI pipeline #44988476 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-02-27T23:29:57Z

[FAILED] Pipeline #44988476: 10/20 passed

raayandhar · 2026-03-03T19:33:01Z

@bkryu when you get the chance, could you rerun CI?

raayandhar · 2026-03-09T23:52:40Z

@bkryu bumping again

bkryu · 2026-03-17T21:35:56Z

/bot run

flashinfer-bot · 2026-03-17T21:36:11Z

GitLab MR !353 has been created, and the CI pipeline #46375530 is currently running. I'll report back once the pipeline job completes.

bkryu

LGTM

flashinfer-bot · 2026-03-17T21:39:43Z

[CANCELING] Pipeline #46375530: canceled

Signed-off-by: raayandhar <raayan.dhar@gmail.com>

bkryu · 2026-03-17T21:43:19Z

/bot run

bkryu · 2026-03-17T21:43:34Z

@nv-yunzheq, can you help review this PR?

flashinfer-bot · 2026-03-17T21:43:56Z

GitLab MR !353 has been updated with latest changes, and the CI pipeline #46376023 is currently running. I'll report back once the pipeline job completes.

nv-yunzheq

LGTM.

nv-yunzheq · 2026-03-17T22:34:48Z

    if out_dtype != torch.bfloat16:
        raise ValueError(
-            "You cannot provide an output dtype to the TGV backend. Use the CUTLASS backend instead."
+            "You cannot provide an output dtype to the TGV backend. Use the CUTLASS or cuDNN backend instead."


This seem to be a fix for an old incorrect information. Is it true?

Yes, we have tests for it.
Exception is that for SM103 it doesn't work...
https://github.com/flashinfer-ai/flashinfer/blob/main/tests/gemm/test_mm_bf16.py#L51
Worth mentioning here you think?

flashinfer-bot · 2026-03-18T03:05:57Z

[FAILED] Pipeline #46376023: 13/20 passed

…r-ai#2644)  ## 📌 Description  Adds support for FP32 dtype output for `mm_bf16` and `bmm_bf16` for the CUTLASS and cuDNN backends. I'm not familiar enough with the TGV kernel to know if / how to support it for that backend. ## 🔍 Related Issues flashinfer-ai#2624  ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [X] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [X] I have installed the hooks with `pre-commit install`. - [X] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [X] Tests have been added or updated as needed. - [X] All tests are passing (`unittest`, etc.). ## Reviewer Notes   ## Summary by CodeRabbit * **New Features** * BF16-based matrix ops (mm_bf16, bmm_bf16) now allow float32 outputs in addition to bfloat16 and float16; supported across applicable backends. * **Tests** * Tests extended to cover float32 outputs for BF16/GEMM operations. * **Documentation** * User-facing docs and validation messages updated to list bf16, fp16, fp32 as valid output dtypes.  Signed-off-by: raayandhar <raayan.dhar@gmail.com>

…r-ai#2644)  ## 📌 Description  Adds support for FP32 dtype output for `mm_bf16` and `bmm_bf16` for the CUTLASS and cuDNN backends. I'm not familiar enough with the TGV kernel to know if / how to support it for that backend. ## 🔍 Related Issues flashinfer-ai#2624  ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [X] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [X] I have installed the hooks with `pre-commit install`. - [X] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [X] Tests have been added or updated as needed. - [X] All tests are passing (`unittest`, etc.). ## Reviewer Notes   ## Summary by CodeRabbit * **New Features** * BF16-based matrix ops (mm_bf16, bmm_bf16) now allow float32 outputs in addition to bfloat16 and float16; supported across applicable backends. * **Tests** * Tests extended to cover float32 outputs for BF16/GEMM operations. * **Documentation** * User-facing docs and validation messages updated to list bf16, fp16, fp32 as valid output dtypes.  Signed-off-by: raayandhar <raayan.dhar@gmail.com> Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>

raayandhar requested review from Anerudhan, bkryu, cyx-6, djmmoss, jiahanc, jimmyzho, nv-yunzheq, nvmbreughe, yongwww and yzh119 as code owners February 26, 2026 22:25

gemini-code-assist Bot reviewed Feb 26, 2026

View reviewed changes

bkryu added the run-ci label Feb 27, 2026

bkryu approved these changes Mar 17, 2026

View reviewed changes

try fp32 output

78061e4

Signed-off-by: raayandhar <raayan.dhar@gmail.com>

raayandhar force-pushed the users/rdhar/bf16_mm_fp32_out branch from 00c9897 to 78061e4 Compare March 17, 2026 21:42

raayandhar requested review from IwakuraRein, aleozlx and saltyminty as code owners March 17, 2026 21:42

nv-yunzheq approved these changes Mar 17, 2026

View reviewed changes

bkryu merged commit b92da57 into flashinfer-ai:main Mar 18, 2026
32 of 33 checks passed

bkryu mentioned this pull request Mar 18, 2026

[Feature] Support FP32 output for bf16 mm #2624

Closed

coderabbitai Bot mentioned this pull request Mar 30, 2026

feat: Add cuBLASLt backend for mm_bf16 and enable multi-tactic autotuning for FP8/MXFP8 runners #2914

Open

Conversation

raayandhar commented Feb 26, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Feb 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Estimated code review effort

Possibly related issues

Possibly related PRs

Suggested labels

Suggested reviewers

Poem

❌ Failed checks (1 warning)

Uh oh!

raayandhar commented Feb 26, 2026

Uh oh!

gemini-code-assist Bot commented Feb 26, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

bkryu commented Feb 27, 2026

Uh oh!

flashinfer-bot commented Feb 27, 2026

Uh oh!

flashinfer-bot commented Feb 27, 2026

Uh oh!

raayandhar commented Mar 3, 2026

Uh oh!

raayandhar commented Mar 9, 2026

Uh oh!

bkryu commented Mar 17, 2026

Uh oh!

flashinfer-bot commented Mar 17, 2026

Uh oh!

bkryu left a comment

Choose a reason for hiding this comment

Uh oh!

flashinfer-bot commented Mar 17, 2026

Uh oh!

bkryu commented Mar 17, 2026

Uh oh!

bkryu commented Mar 17, 2026

Uh oh!

flashinfer-bot commented Mar 17, 2026

Uh oh!

nv-yunzheq left a comment

Choose a reason for hiding this comment

Uh oh!

nv-yunzheq Mar 17, 2026

Choose a reason for hiding this comment

Uh oh!

raayandhar Mar 17, 2026

Choose a reason for hiding this comment

Uh oh!

flashinfer-bot commented Mar 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

raayandhar commented Feb 26, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Feb 26, 2026 •

edited

Loading