Skip to content

Commit 18b335d

Browse files
authored
[TRTLLM-9989][fix] Disable tvm_ffi for CuteDSL nvFP4 dense GEMM. (#10040)
Signed-off-by: Yukun He <[email protected]>
1 parent 2fd1a23 commit 18b335d

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -371,7 +371,7 @@ class CuteDSLNVFP4BlackwellLinear(TunableRunner):
371371
def __init__(self,
372372
output_dtype: torch.dtype,
373373
to_userbuffers: bool = False,
374-
use_tvm_ffi: bool = True):
374+
use_tvm_ffi: bool = False):
375375
super().__init__()
376376

377377
if output_dtype != torch.bfloat16:
@@ -775,7 +775,7 @@ def cute_dsl_nvfp4_gemm_blackwell(
775775
alpha: torch.Tensor,
776776
output_dtype: torch.dtype,
777777
to_userbuffers: bool = False,
778-
use_tvm_ffi: bool = True,
778+
use_tvm_ffi: bool = False,
779779
) -> torch.Tensor:
780780
"""CuteDSL-based NVFP4 GEMM optimized for Blackwell.
781781
@@ -825,7 +825,7 @@ def _(
825825
alpha: torch.Tensor, # Match custom op signature
826826
output_dtype: torch.dtype,
827827
to_userbuffers: bool = False,
828-
use_tvm_ffi: bool = True,
828+
use_tvm_ffi: bool = False,
829829
):
830830
# [m, k]
831831
shape = list(mat_a.shape)

0 commit comments

Comments
 (0)