Support checks PoC

[nvmbreughe]

📌 Description

This PR adds is_*supported checks for backend and compute capability, through decorators.

1. This allows us to check support before running
2. It also wraps the original function so it calls back the support check before running.
3. The wrapped function also adds an optional parameter "skip_check". A quick measurement show only minimal impact (14.51s without checks, 14.58s with checks for all of test_mm_fp4), so we should further benchmark the usefulness of this feature.

Example:

🔍 Related Issues

🚀 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

[nvjullin]

The checks currently live very far away from the implementation and updating them to be consistent with each other can eventually become a maintenance problem. The conditional checks are also quite tricky to get correct. For example, it's not easy to tell if the mxfp4 checks are correct.

    if not use_nvfp4 and block_size != 32:
        raise ValueError("mxfp4 supports block_size = 32.")

    if backend != "cudnn" and not use_nvfp4:
        raise ValueError("Only cudnn FP4 GEMM supports mxfp4 quantization.")

Shouldn't the checks be reordered to avoid confusing error messages?

1. User tries trtllm + block_size=16 and gets rejected by block_size=32
2. User then tries trtllm + block_size=32 and gets rejected by only cudnn is supported

---

Instead of having one top level supports_backends, perhaps consider a two level design:

1. Local requirement decorator requirement written for each backend entrypoint
2. Top level backend_requirement that composes requirements

For example:

def cudnn_gemm_fp4_requirement(
    # ...
):
        if (
            not use_nvfp4
            and _match_sm_version(a.device, ["120"])
            and cudnn.backend_version() < 91400
        ):
            raise LibraryError(
                "cudnn FP4 GEMM with mxfp4 quantization is not supported on SM120 with cuDNN backend version < 9.14.0."
            )

        _check_cudnn_fp4_availability()
        # ...

@requirement(cudnn_gemm_fp4_requirement, capability=["100", "101", "102"])
def execute_cudnn_gemm_fp4_graph(
    # ...


@backend_requirement({
    "cudnn": execute_cudnn_gemm_fp4_graph.requirement,
    "trtllm": #...
})
def mm_fp4(
    # ...

This also means that all requirements are enforced to be local to the backend and won't affect each other.

[nvmbreughe]

Instead of having one top level supports_backends, perhaps consider a two level design:

1. Local requirement decorator requirement written for each backend entrypoint
2. Top level backend_requirement that composes requirements

Thank you @nvjullin for the excellent suggestion.
I think these are two discussions:

1. Separate the support checks
2. Separate the execution routines

While both are valid points, we prioritize separating the checks for now. Not all APIs are as cleanly to separate (2) atm and there is a plan for a more OO Backend class @Anerudhan. That does overlap somewhat with the support checks, as eventually we would be able to do something like cudnn_backend->check_mmfp4_support().

So I think as an intermediary step, and to get tighter checks in, we could do something like this:

@supported_compute_capability(["100", "101", "102"])
def cudnn_gemm_fp4_requirement(
    # ...
):
        if (
            not use_nvfp4
            and _match_sm_version(a.device, ["120"])
            and cudnn.backend_version() < 91400
        ):
            raise LibraryError(
                "cudnn FP4 GEMM with mxfp4 quantization is not supported on SM120 with cuDNN backend version < 9.14.0."
            )

        _check_cudnn_fp4_availability()
        # ...


@backend_requirement({
    "cudnn": execute_cudnn_gemm_fp4_graph.requirement,
    "trtllm": #...
    },
   common_check=common_fp4_checks # To be called by all backend checks
})
def mm_fp4(
    )
    if backend == "cudnn":
              # cudnn path
    elif backend == "trtllm":
              # trtllm path         

[nvjullin]

While both are valid points, we prioritize separating the checks for now. Not all APIs are as cleanly to separate (2) atm and there is a plan for a more OO Backend class @Anerudhan. That does overlap somewhat with the support checks, as eventually we would be able to do something like cudnn_backend->check_mmfp4_support().

I wasn't aware, thanks for the info. LGTM.

[nvmbreughe]

While both are valid points, we prioritize separating the checks for now. Not all APIs are as cleanly to separate (2) atm and there is a plan for a more OO Backend class @Anerudhan. That does overlap somewhat with the support checks, as eventually we would be able to do something like cudnn_backend->check_mmfp4_support().

I wasn't aware, thanks for the info. LGTM.

Thank you for the excellent suggestions, @nvjullin

[nvmbreughe]

/bot run

[flashinfer-bot]

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

[aleozlx]

looks like a good step forward

[flashinfer-bot]

[SUCCESS] Pipeline #36524696: 13/17 passed

[yzh119]

LGTM

[sricketts]

Overall LGTM. Added one suggestion.