Small fixes

annanyapr · annanyapr · commit ebce953b9b53 · 2025-03-21T20:16:14.000-04:00
diff --git a/python/tvm/relax/backend/cuda/flashinfer.py b/python/tvm/relax/backend/cuda/flashinfer.py
@@ -426,7 +426,7 @@ def gen_sampling_module(target: Target, num_threads: int = 8):
             "FlashInfer is not installed. Please follow instructions "
             "in https://docs.flashinfer.ai to install FlashInfer."
         )
-    uri, source_paths = gen_sampling_tvm_binding()
+    uri, source_paths = gen_sampling_tvm_binding(uri="sampling")
     object_files = _compile_flashinfer_kernels(uri, source_paths, target, num_threads)
     modules = _load_flashinfer_modules(object_files)
     return modules

Original file line number	Diff line number	Diff line change
`@@ -426,7 +426,7 @@ def gen_sampling_module(target: Target, num_threads: int = 8):`
`426`	`426`	`"FlashInfer is not installed. Please follow instructions "`
`427`	`427`	`"in https://docs.flashinfer.ai to install FlashInfer."`
`428`	`428`	`)`
`429`		`- uri, source_paths = gen_sampling_tvm_binding()`
	`429`	`+ uri, source_paths = gen_sampling_tvm_binding(uri="sampling")`
`430`	`430`	`object_files = _compile_flashinfer_kernels(uri, source_paths, target, num_threads)`
`431`	`431`	`modules = _load_flashinfer_modules(object_files)`
`432`	`432`	`return modules`