diff --git a/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu b/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu index 54d85db8..46e3f858 100644 --- a/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu +++ b/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu @@ -73,7 +73,7 @@ void FMHACutlassSM100BwdRun(at::Tensor workspace_buffer, at::Tensor d_o, at::Ten dq, dk, dv, softmax_scale, max_seqlen_q, max_seqlen_kv); } else { - std::cout << "No kernel instantiated for head_dim_qk=" << head_dim_qk << " head_dim_vo=" << head_dim_vo << std::endl; + TORCH_CHECK(false, "No kernel instantiated for head_dim_qk=", head_dim_qk, " head_dim_vo=", head_dim_vo); } }); diff --git a/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu b/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu index ab66f0fd..c08fb145 100644 --- a/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu +++ b/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu @@ -72,8 +72,8 @@ void FMHACutlassSM100FwdRun(at::Tensor workspace_buffer, at::Tensor q, at::Tenso cumulative_seqlen_q, cumulative_seqlen_kv, o, lse, sm_scale, max_seqlen_q, max_seqlen_kv); } else { - std::cout << "No kernel instantiated for head_dim_qk=" << head_dim_qk - << " head_dim_vo=" << head_dim_vo << std::endl; + TORCH_CHECK(false, "No kernel instantiated for head_dim_qk=", head_dim_qk, + " head_dim_vo=", head_dim_vo); } });