Skip to content

[None][perf] Allow CUDA-core NVFP4 GEMM for M up to 16 #16345

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mihai-chiorean wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[None][perf] Allow CUDA-core NVFP4 GEMM for M up to 16 #16345

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8e5cce1
[None][perf] Allow CUDA-core NVFP4 GEMM for M up to 16
mihai-chiorean Jul 13, 2026
98a893d
Merge branch 'main' into codex/cudacore-nvfp4-m16
mihai-chiorean Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -681,7 +681,7 @@ class CudaCoreNVFP4Runner(TunableRunner):

This runner is available on:
- SM >= 100 (Blackwell)
- M <= 8 (small batch size limitation from kernel template)
- M <= 16 (small batch size limitation from kernel template)
"""

# Shared tuning config (no tactics needed, single implementation)
Expand All @@ -690,7 +690,7 @@ class CudaCoreNVFP4Runner(TunableRunner):
# Minimum supported architecture: SM100 (Blackwell)
MIN_SM_VERSION = 100
# Maximum M dimension (from cudaCoreGemmTemplateMaxM in C++ kernel)
MAX_M_DIMENSION = 8
MAX_M_DIMENSION = 16

def __init__(self,
output_buffer_kind: int,
Expand Down Expand Up @@ -1218,7 +1218,7 @@ def nvfp4_gemm(
- CUTLASS: Predefined CUTLASS configurations with auto-tuning
- cuBLASLt: Heuristic-based algorithms from cuBLASLt library
- CuteDSL: Blackwell-optimized persistent kernels (when available and inputs are valid)
- CUDA Core: CUDA Core implementation (requires SM >= 100 and M <= 8)
- CUDA Core: CUDA Core implementation (requires SM >= 100 and M <= 16)
- Marlin: Hopper W4A16 NVFP4 implementation (requires SM 90-99)

The AutoTuner profiles all available backends during the first run and caches
Expand Down
8 changes: 4 additions & 4 deletions tests/unittest/_torch/thop/parallel/test_fp4_linear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -695,9 +695,9 @@ def test_fp4_linear_cublaslt(dtype, mnk):
)
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("mnk", [(1, 4096, 7168), (4, 7168, 16384),
(8, 2112, 7168)])
(8, 2112, 7168), (16, 1024, 1024)])
def test_fp4_linear_cuda_core(dtype, mnk):
"""Test CUDA Core NVFP4 GEMM implementation on SM >= 100 (M <= 8)"""
"""Test CUDA Core NVFP4 GEMM implementation on SM >= 100 (M <= 16)"""

SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk
torch.manual_seed(0)
Expand Down Expand Up @@ -876,8 +876,8 @@ def _make_nvfp4_inputs_for_bias_test(m, n, k, dtype=torch.bfloat16):
def test_fp4_gemm_bias_per_backend(backend, mnk):
"""Per-backend numerical parity: nvfp4_gemm(bias=B) ≈ nvfp4_gemm(bias=None) + B."""
m, n, k = mnk
if backend == "cuda_core" and m > 8:
pytest.skip("cuda_core backend only supports M <= 8")
if backend == "cuda_core" and m > 16:
pytest.skip("cuda_core backend only supports M <= 16")

act_fp4, w_fp4, act_sf, w_sf, alpha = _make_nvfp4_inputs_for_bias_test(
m, n, k)
Expand Down
Loading