fix(jit): separate TGV SM100 target module identities

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",

tests/jit/test_tgv_gemm_jit.py

@@ -0,0 +1,110 @@
+# 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 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 = [
+        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",
+    ]

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(
@@ -17,15 +14,11 @@
     ],
 )
 def test_tgv_gemm_sm100_reference_correctness(shape_kwargs):
-    """tgv_gemm_sm100 kernel (SM100 only in practice) vs reference (a @ b + bias)."""
-    from flashinfer.utils import is_sm100f_supported
+    """TGV SM100-family kernel vs the a @ b + bias reference."""
+    from flashinfer.utils import get_compute_capability, 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 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+)")
     from flashinfer import tgv_gemm_sm100
@@ -34,10 +27,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