added new gemm to aot module

Author: hypdeb

flashinfer/aot.py

@@ -57,6 +57,7 @@
     gen_gemm_sm120_module,
     gen_gemm_sm120_module_cutlass_fp4,
     gen_trtllm_gen_gemm_module,
+    gen_trtllm_low_latency_gemm_module,
 )
 from .jit.spdlog import gen_spdlog_module
 from .jit.mla import gen_mla_module
@@ -460,6 +461,7 @@ def gen_all_modules(
             )
             jit_specs.append(gen_mxfp8_quantization_sm100_module())
             jit_specs.append(gen_trtllm_gen_gemm_module())
+            jit_specs.append(gen_trtllm_low_latency_gemm_module())
             jit_specs.append(gen_trtllm_gen_fused_moe_sm100_module())
         if has_sm100f:
             # Add TGV GEMM modules compiled with SM100f flags for both bf16 and fp16

flashinfer/jit/gemm/__init__.py

@@ -22,6 +22,7 @@
     gen_gemm_sm100_module,
     gen_gemm_sm120_module,
     gen_trtllm_gen_gemm_module,
+    gen_trtllm_low_latency_gemm_module,
     gen_tgv_gemm_sm10x_module,
     gen_gemm_sm90_module,
 )
@@ -35,6 +36,7 @@
     "gen_gemm_sm100_module",
     "gen_gemm_sm120_module",
     "gen_trtllm_gen_gemm_module",
+    "gen_trtllm_low_latency_gemm_module",
     "gen_tgv_gemm_sm10x_module",
     "gen_gemm_sm90_module",
     "gen_deepgemm_sm100_module",