From 2fbdb91329d65399e864a130b226ba219107c6c0 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Thu, 20 Aug 2026 13:29:18 +0000 Subject: [PATCH 1/4] fix(jit): separate TGV SM100 target modules --- flashinfer/jit/gemm/core.py | 7 +- tests/jit/test_tgv_gemm_jit.py | 108 ++++++++++++++++++ ...st_tgv_gemm_sm100_reference_correctness.py | 10 +- 3 files changed, 116 insertions(+), 9 deletions(-) create mode 100644 tests/jit/test_tgv_gemm_jit.py diff --git a/flashinfer/jit/gemm/core.py b/flashinfer/jit/gemm/core.py index 9856bf3ee90..1d3436be4c2 100644 --- a/flashinfer/jit/gemm/core.py +++ b/flashinfer/jit/gemm/core.py @@ -807,9 +807,12 @@ def gen_tgv_gemm_sm10x_module( ) dtype_str = "bf16" if dtype == torch.bfloat16 else "fp16" - module_name = f"tgv_gemm_{dtype_str}" + target = "sm100f" if use_sm_100f else "sm100a" + module_name = f"tgv_gemm_{dtype_str}_{target}" - gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / f"gen_tgv_gemm_{dtype_str}" + gen_directory = ( + jit_env.FLASHINFER_GEN_SRC_DIR / f"gen_tgv_gemm_{dtype_str}_{target}" + ) os.makedirs(gen_directory, exist_ok=True) source_paths = [ jit_env.FLASHINFER_CSRC_DIR / "tgv_gemm.cu", diff --git a/tests/jit/test_tgv_gemm_jit.py b/tests/jit/test_tgv_gemm_jit.py new file mode 100644 index 00000000000..28909ca8e41 --- /dev/null +++ b/tests/jit/test_tgv_gemm_jit.py @@ -0,0 +1,108 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +import torch + +from flashinfer.jit import env as jit_env +from flashinfer.jit.gemm.core import gen_tgv_gemm_sm10x_module + + +def test_tgv_gemm_target_specific_jit_specs_and_aot_inventory( + monkeypatch, tmp_path +): + monkeypatch.setattr(jit_env, "FLASHINFER_GEN_SRC_DIR", tmp_path / "generated") + + specs = [ + gen_tgv_gemm_sm10x_module(torch.bfloat16, use_sm_100f=False), + gen_tgv_gemm_sm10x_module(torch.float16, use_sm_100f=False), + gen_tgv_gemm_sm10x_module(torch.bfloat16, use_sm_100f=True), + gen_tgv_gemm_sm10x_module(torch.float16, use_sm_100f=True), + ] + assert [spec.name for spec in specs] == [ + "tgv_gemm_bf16_sm100a", + "tgv_gemm_fp16_sm100a", + "tgv_gemm_bf16_sm100f", + "tgv_gemm_fp16_sm100f", + ] + assert [spec.sources[1].parent.name for spec in specs] == [ + "gen_tgv_gemm_bf16_sm100a", + "gen_tgv_gemm_fp16_sm100a", + "gen_tgv_gemm_bf16_sm100f", + "gen_tgv_gemm_fp16_sm100f", + ] + assert [ + [flag for flag in spec.extra_cuda_cflags if flag.startswith("-gencode=")] + for spec in specs + ] == [ + ["-gencode=arch=compute_100a,code=sm_100a"], + ["-gencode=arch=compute_100a,code=sm_100a"], + ["-gencode=arch=compute_100f,code=sm_100f"], + ["-gencode=arch=compute_100f,code=sm_100f"], + ] + + from flashinfer import aot + + monkeypatch.setattr(aot, "gen_attention", lambda *args: ()) + for generator_name in [ + "gen_spdlog_module", + "gen_gemm_module", + "gen_bgmv_moe_module", + "gen_hash_topk_module", + "gen_fp4_quantization_sm100_module", + "gen_cutlass_fused_moe_sm100_module", + "gen_gemm_sm100_module", + "gen_gemm_sm100_module_cutlass_fp4", + "gen_gemm_sm100_module_cutlass_nvfp4_svdquant", + "gen_gemm_sm100_module_cutlass_fp8", + "gen_gemm_sm100_module_cutlass_mxfp8", + "gen_mxfp8_quantization_sm100_module", + "gen_trtllm_gen_gemm_module", + "gen_trtllm_low_latency_gemm_module", + "gen_trtllm_gen_fused_moe_sm100_module", + "gen_moe_utils_module", + "gen_mm_bf16_cublaslt_module", + "gen_cudnn_fmha_module", + ]: + monkeypatch.setattr( + aot, + generator_name, + lambda *args, _name=generator_name, **kwargs: SimpleNamespace( + name=_name + ), + ) + + inventory = aot.gen_all_modules( + [], + [], + [], + [], + [], + [], + {"sm100": True, "sm100f": True}, + False, + False, + False, + True, + False, + False, + False, + ) + assert [spec.name for spec in inventory if spec.name.startswith("tgv_gemm_")] == [ + "tgv_gemm_bf16_sm100a", + "tgv_gemm_fp16_sm100a", + "tgv_gemm_bf16_sm100f", + "tgv_gemm_fp16_sm100f", + ] diff --git a/tests/trace/test_tgv_gemm_sm100_reference_correctness.py b/tests/trace/test_tgv_gemm_sm100_reference_correctness.py index eaed4820c82..7c4d30303fe 100644 --- a/tests/trace/test_tgv_gemm_sm100_reference_correctness.py +++ b/tests/trace/test_tgv_gemm_sm100_reference_correctness.py @@ -17,15 +17,11 @@ ], ) def test_tgv_gemm_sm100_reference_correctness(shape_kwargs): - """tgv_gemm_sm100 kernel (SM100 only in practice) vs reference (a @ b + bias).""" + """TGV SM100-family kernel vs the a @ b + bias reference.""" from flashinfer.utils import is_sm100f_supported - # The kernel's Python gate accepts SM100 or SM103 (see - # gemm_base._match_sm_version) but the precompiled cubin only has an - # SM100 kernel image; calling on SM103 crashes with "no kernel image" - # inside CUDA (uncatchable via try/except). Restrict to SM100. - if _cc() != (10, 0): - pytest.skip("tgv_gemm_sm100 cubin is only built for SM100") + if _cc() not in [(10, 0), (10, 3)]: + pytest.skip("tgv_gemm_sm100 requires SM100 or SM103") if not is_sm100f_supported(torch.device("cuda")): pytest.skip("tgv_gemm_sm100 requires SM100f support (CUDA 12.9+)") from flashinfer import tgv_gemm_sm100 From a90ec4c346ca171bb4fce38ce82c7e1a67f59a7b Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Thu, 20 Aug 2026 13:38:32 +0000 Subject: [PATCH 2/4] test(jit): format TGV identity coverage --- tests/jit/test_tgv_gemm_jit.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/jit/test_tgv_gemm_jit.py b/tests/jit/test_tgv_gemm_jit.py index 28909ca8e41..5c2200659df 100644 --- a/tests/jit/test_tgv_gemm_jit.py +++ b/tests/jit/test_tgv_gemm_jit.py @@ -20,9 +20,7 @@ from flashinfer.jit.gemm.core import gen_tgv_gemm_sm10x_module -def test_tgv_gemm_target_specific_jit_specs_and_aot_inventory( - monkeypatch, tmp_path -): +def test_tgv_gemm_target_specific_jit_specs_and_aot_inventory(monkeypatch, tmp_path): monkeypatch.setattr(jit_env, "FLASHINFER_GEN_SRC_DIR", tmp_path / "generated") specs = [ @@ -79,9 +77,7 @@ def test_tgv_gemm_target_specific_jit_specs_and_aot_inventory( monkeypatch.setattr( aot, generator_name, - lambda *args, _name=generator_name, **kwargs: SimpleNamespace( - name=_name - ), + lambda *args, _name=generator_name, **kwargs: SimpleNamespace(name=_name), ) inventory = aot.gen_all_modules( From 2438fa19d55c31380c7459898e0a0fe6d4b487c3 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Thu, 20 Aug 2026 16:19:35 +0000 Subject: [PATCH 3/4] test(jit): make TGV identity coverage CPU-safe --- tests/jit/test_tgv_gemm_jit.py | 6 ++++++ tests/trace/test_tgv_gemm_sm100_reference_correctness.py | 5 +---- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/jit/test_tgv_gemm_jit.py b/tests/jit/test_tgv_gemm_jit.py index 5c2200659df..c96cf2e96d9 100644 --- a/tests/jit/test_tgv_gemm_jit.py +++ b/tests/jit/test_tgv_gemm_jit.py @@ -16,11 +16,17 @@ import torch +from flashinfer.jit import core as jit_core from flashinfer.jit import env as jit_env from flashinfer.jit.gemm.core import gen_tgv_gemm_sm10x_module def test_tgv_gemm_target_specific_jit_specs_and_aot_inventory(monkeypatch, tmp_path): + monkeypatch.setattr( + jit_core.current_compilation_context, + "TARGET_CUDA_ARCHS", + {(10, "0a"), (10, "3a")}, + ) monkeypatch.setattr(jit_env, "FLASHINFER_GEN_SRC_DIR", tmp_path / "generated") specs = [ diff --git a/tests/trace/test_tgv_gemm_sm100_reference_correctness.py b/tests/trace/test_tgv_gemm_sm100_reference_correctness.py index 7c4d30303fe..6bd992106c0 100644 --- a/tests/trace/test_tgv_gemm_sm100_reference_correctness.py +++ b/tests/trace/test_tgv_gemm_sm100_reference_correctness.py @@ -30,10 +30,7 @@ def test_tgv_gemm_sm100_reference_correctness(shape_kwargs): inputs = tgv_gemm_sm100_trace.init(**shape_kwargs) assert inputs["b"].shape == (shape_kwargs["K"], shape_kwargs["N"]) assert inputs["b"].stride(0) == 1, "tgv_gemm_sm100 expects column-major b" - try: - api_out = tgv_gemm_sm100(inputs["a"], inputs["b"], inputs["bias"]) - except Exception as exc: - pytest.skip(f"tgv_gemm_sm100 unavailable: {exc}") + api_out = tgv_gemm_sm100(inputs["a"], inputs["b"], inputs["bias"]) torch.cuda.synchronize() ref_out = tgv_gemm_sm100_trace.reference(inputs["a"], inputs["b"], inputs["bias"]) # Matches tests/gemm/test_tgv_gemm.py: bf16 * K=1024 accumulation makes From b853d1e6ff3203a3ffed1c42e0a44d6bd8cce2bb Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 22 Aug 2026 13:48:16 +0000 Subject: [PATCH 4/4] test(trace): use public compute capability helper --- tests/trace/test_tgv_gemm_sm100_reference_correctness.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/trace/test_tgv_gemm_sm100_reference_correctness.py b/tests/trace/test_tgv_gemm_sm100_reference_correctness.py index 6bd992106c0..523cef665ab 100644 --- a/tests/trace/test_tgv_gemm_sm100_reference_correctness.py +++ b/tests/trace/test_tgv_gemm_sm100_reference_correctness.py @@ -3,10 +3,7 @@ import torch import pytest -from tests.trace.reference_utils import ( - _cc, - _check, -) +from tests.trace.reference_utils import _check @pytest.mark.parametrize( @@ -18,9 +15,9 @@ ) def test_tgv_gemm_sm100_reference_correctness(shape_kwargs): """TGV SM100-family kernel vs the a @ b + bias reference.""" - from flashinfer.utils import is_sm100f_supported + from flashinfer.utils import get_compute_capability, is_sm100f_supported - if _cc() not in [(10, 0), (10, 3)]: + if get_compute_capability(torch.device("cuda")) not in [(10, 0), (10, 3)]: pytest.skip("tgv_gemm_sm100 requires SM100 or SM103") if not is_sm100f_supported(torch.device("cuda")): pytest.skip("tgv_gemm_sm100 requires SM100f support (CUDA 12.9+)")