fix(kernels): catch C++ exceptions at the FFI boundary and surface the message

openinfer-core/src/ops.rs

@@ -13,14 +13,14 @@ pub use attention::{
     paged_attention_batch_decode_split_kv_into, prefill_attention_paged_into,
 };
 pub use openinfer_kernels::ops::{
-    GEMM_LT_MAX_N, LoraDecodeGroupedProjection, accumulate_bf16_token_scaled_to_f32_into,
-    add_batch, add_batch_into, argmax, argmax_batch_bf16_into, bf16_hidden_to_f32_into,
-    copy_hidden_rows_into, copy_hidden_token_range_into, dflash_qk_norm_rope_into,
-    embedding_decode_into, extract_vec, extract_vec_into, extract_vec_ref, extract_vec_ref_into,
-    f32_to_bf16_hidden_into, fused_add_rms_norm_into, gather_hidden_tokens_into, gemm,
-    gemm_graphsafe_into_checked, gemm_graphsafe_ref_into_checked, gemm_into_checked, gemm_lt_tune,
-    gemm_per_token, gemv, linear, lora_decode_fused_delta_group3_into,
-    lora_decode_fused_delta_into, pack_lora_b_rows_into,
+    GEMM_LT_MAX_N, LoraDecodeGroupedProjection, SUPPORTED_GQA_GROUP_SIZES,
+    accumulate_bf16_token_scaled_to_f32_into, add_batch, add_batch_into, argmax,
+    argmax_batch_bf16_into, bf16_hidden_to_f32_into, copy_hidden_rows_into,
+    copy_hidden_token_range_into, dflash_qk_norm_rope_into, embedding_decode_into, extract_vec,
+    extract_vec_into, extract_vec_ref, extract_vec_ref_into, f32_to_bf16_hidden_into,
+    fused_add_rms_norm_into, gather_hidden_tokens_into, gemm, gemm_graphsafe_into_checked,
+    gemm_graphsafe_ref_into_checked, gemm_into_checked, gemm_lt_tune, gemm_per_token, gemv, linear,
+    lora_decode_fused_delta_group3_into, lora_decode_fused_delta_into, pack_lora_b_rows_into,
     qk_norm_partial_rope_batched_decode_hd256_into, rms_norm, rms_norm_batch_offset_into,
     rms_norm_gated_batch_into, rms_norm_into, rms_norm_offset_into, scale_f32_in_place,
     scaled_add_batch_into, scaled_add_rows_indexed_into, scaled_add_rows_into,

openinfer-kernels/csrc/glm52/glm52_topk.cu

@@ -7,6 +7,8 @@
 // TopKTieBreak::Small, dsa_graph_safe=true, K=2048. Mirrors vLLM
 // `sparse_attn_indexer` which always passes `seq_lens` as lengths.
 
+#include "../shared/ffi_guard.cuh"
+
 #include <cstddef>
 #include <cstdio>
 
@@ -17,6 +19,7 @@ extern "C" int glm52_flashinfer_topk_2048_cuda(
     const float* logits, int* output_indices, float* output_values,
     const int* lengths, int num_rows, int top_k, int max_len,
     cudaStream_t stream) {
+  OPENINFER_FFI_GUARD_BEGIN
   if (logits == nullptr || output_indices == nullptr || lengths == nullptr) {
     return static_cast<int>(CUDA_ERROR_INVALID_VALUE);
   }
@@ -57,4 +60,5 @@ extern "C" int glm52_flashinfer_topk_2048_cuda(
     return static_cast<int>(CUDA_ERROR_LAUNCH_FAILED);
   }
   return static_cast<int>(CUDA_SUCCESS);
+  OPENINFER_FFI_GUARD_END(-1)
 }

openinfer-kernels/csrc/kimi_k2/kimi_mla.cu

@@ -1,3 +1,5 @@
+#include "../shared/ffi_guard.cuh"
+
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
 #include <cublas_v2.h>
@@ -614,6 +616,7 @@ int kimi_flashinfer_single_prefill_mla_cuda(void* q,
                                             int kv_len,
                                             float sm_scale,
                                             cudaStream_t stream) {
+  OPENINFER_FFI_GUARD_BEGIN
   // kv_len > seq_len is the prefix-cache suffix prefill: q covers the suffix
   // only, k/v cover cached prefix + suffix. MaskMode::kCausal aligns
   // bottom-right, so query i attends to kv positions <= kv_len - seq_len + i —
@@ -659,6 +662,7 @@ int kimi_flashinfer_single_prefill_mla_cuda(void* q,
           flashinfer::MaskMode::kCausal,
           Variant,
           PrefillParamsT>(params, /*tmp=*/nullptr, stream));
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 int kimi_mla_absorb_q_nope_cuda(const DType* kv_b_proj,
@@ -790,6 +794,7 @@ int kimi_mla_paged_kv_append_cuda(void* ckv_cache,
                                   int page_size,
                                   int batch_size,
                                   cudaStream_t stream) {
+  OPENINFER_FFI_GUARD_BEGIN
   if (nnz <= 0 || page_size <= 0 || batch_size <= 0) {
     return static_cast<int>(cudaErrorInvalidValue);
   }
@@ -809,6 +814,7 @@ int kimi_mla_paged_kv_append_cuda(void* ckv_cache,
           /*append_ckv_stride_n=*/kKvLoraRank,
           /*append_kpe_stride_n=*/kRopeDim,
           stream));
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 int kimi_flashinfer_batch_decode_mla_cuda(void* q_nope,
@@ -831,6 +837,7 @@ int kimi_flashinfer_batch_decode_mla_cuda(void* q_nope,
                                           int batch_size,
                                           float sm_scale,
                                           cudaStream_t stream) {
+  OPENINFER_FFI_GUARD_BEGIN
   if (num_qo_heads <= 0 || page_size <= 0 || batch_size <= 0) {
     return static_cast<int>(cudaErrorInvalidValue);
   }
@@ -872,6 +879,7 @@ int kimi_flashinfer_batch_decode_mla_cuda(void* q_nope,
           /*tmp_s=*/nullptr,
           /*enable_pdl=*/false,
           stream));
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 }  // extern "C"

openinfer-kernels/csrc/shared/ffi_guard.cuh

@@ -0,0 +1,24 @@
+// Catch C++ exceptions (FlashInfer throws from host-side dispatch) before they
+// cross the extern "C" boundary — a foreign exception reaching Rust aborts the
+// process with no message.
+#pragma once
+
+#include <exception>
+
+void openinfer_ffi_set_last_error(const char* what);
+
+// Entering a guard clears the previous message so a -1 seen by Rust never
+// reads a stale what() from an earlier call on the same thread.
+#define OPENINFER_FFI_GUARD_BEGIN \
+  openinfer_ffi_set_last_error(""); \
+  try {
+#define OPENINFER_FFI_GUARD_END(ret_on_throw)              \
+  }                                                        \
+  catch (const std::exception& e) {                        \
+    openinfer_ffi_set_last_error(e.what());                \
+    return ret_on_throw;                                   \
+  }                                                        \
+  catch (...) {                                            \
+    openinfer_ffi_set_last_error("unknown C++ exception"); \
+    return ret_on_throw;                                   \
+  }

openinfer-kernels/csrc/shared/ffi_last_error.cu

@@ -0,0 +1,13 @@
+#include "ffi_guard.cuh"
+
+#include <string>
+
+namespace {
+thread_local std::string g_ffi_last_error;
+}
+
+void openinfer_ffi_set_last_error(const char* what) {
+  g_ffi_last_error = what == nullptr ? "" : what;
+}
+
+extern "C" const char* openinfer_kernels_last_error() { return g_ffi_last_error.c_str(); }

openinfer-kernels/csrc/shared/flashinfer_sampling.cu

@@ -1,3 +1,5 @@
+#include "ffi_guard.cuh"
+
 #include "common.cuh"
 #include "flashinfer_radix_scratch.cuh"
 
@@ -71,6 +73,7 @@ extern "C" int gpu_sample_batch_flashinfer_cuda(
     uint8_t* valid_scratch, int* output, void* softmax_workspace,
     size_t softmax_workspace_bytes, int n_rows, int vocab_size, int has_top_k_filter,
     int has_top_p_filter, uint64_t seed, uint64_t offset, cudaStream_t stream) {
+  OPENINFER_FFI_GUARD_BEGIN
   dim3 gather_grid(
       (vocab_size + GATHER_CAST_BLOCK * GATHER_CAST_ELEMS_PER_THREAD - 1) /
           (GATHER_CAST_BLOCK * GATHER_CAST_ELEMS_PER_THREAD),
@@ -142,4 +145,5 @@ extern "C" int gpu_sample_batch_flashinfer_cuda(
         stream);
   }
   return static_cast<int>(err);
+  OPENINFER_FFI_GUARD_END(-1)
 }

openinfer-kernels/csrc/shared/paged_attention.cu

@@ -6,6 +6,8 @@
 // FlashInfer's dispatchers internally instantiate multiple GQA group sizes
 // (1,2,3,4,8) — this covers both Qwen3-4B (GQA=4) and Qwen3.5-4B (GQA=8).
 
+#include "ffi_guard.cuh"
+
 #include <cuda_runtime.h>
 #include <cuda_bf16.h>
 #include <cstdint>
@@ -101,6 +103,7 @@ int paged_attention_decode_cuda(
     // Stream
     void*    stream)
 {
+  OPENINFER_FFI_GUARD_BEGIN
     auto paged_kv = make_paged_kv(
         kv_data, k_offset_elems, v_offset_elems,
         page_indices, page_indptr, last_page_len_d,
@@ -142,6 +145,7 @@ int paged_attention_decode_cuda(
             /*tmp_s=*/nullptr,
             /*enable_pdl=*/false,
             reinterpret_cast<cudaStream_t>(stream)));
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 // ---------------------------------------------------------------------------
@@ -188,6 +192,7 @@ int paged_attention_decode_split_kv_cuda(
     // Stream
     void*    stream)
 {
+  OPENINFER_FFI_GUARD_BEGIN
     auto paged_kv = make_paged_kv(
         kv_data, k_offset_elems, v_offset_elems,
         page_indices, page_indptr, last_page_len_d,
@@ -227,6 +232,7 @@ int paged_attention_decode_split_kv_cuda(
             tmp_s,
             /*enable_pdl=*/false,
             reinterpret_cast<cudaStream_t>(stream)));
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 // ---------------------------------------------------------------------------
@@ -251,6 +257,7 @@ int paged_kv_append_cuda(
     int64_t  stride_page,
     void*    stream)
 {
+  OPENINFER_FFI_GUARD_BEGIN
     auto paged_kv = make_paged_kv(
         kv_data, k_offset_elems, v_offset_elems,
         page_indices, page_indptr, last_page_len_d,
@@ -261,6 +268,7 @@ int paged_kv_append_cuda(
         reinterpret_cast<DType*>(key),
         reinterpret_cast<DType*>(value),
         reinterpret_cast<cudaStream_t>(stream)));
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 // ---------------------------------------------------------------------------
@@ -291,6 +299,7 @@ int paged_kv_scatter_cuda(
     int64_t  src_stride_h,         // = max_seq_len * head_dim
     void*    stream)
 {
+  OPENINFER_FFI_GUARD_BEGIN
     auto paged_kv = make_paged_kv(
         kv_data, k_offset_elems, v_offset_elems,
         page_indices, page_indptr, last_page_len_d,
@@ -308,6 +317,7 @@ int paged_kv_scatter_cuda(
         static_cast<size_t>(src_stride_n),   // V has same layout as K
         static_cast<size_t>(src_stride_h),
         reinterpret_cast<cudaStream_t>(stream)));
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 // ---------------------------------------------------------------------------
@@ -364,6 +374,7 @@ int32_t batch_prefill_paged_num_tiles_with_cta_tile_q(
     uint32_t cta_tile_q = resolve_prefill_cta_tile_q(
         packed_qo_len, head_dim, cta_tile_q_override);
     if (cta_tile_q == 0) {
+        openinfer_ffi_set_last_error("invalid cta_tile_q override");
         return -1;
     }
     return static_cast<int32_t>((packed_qo_len + cta_tile_q - 1) / cta_tile_q);
@@ -429,6 +440,7 @@ int batch_prefill_paged_cuda_with_cta_tile_q(
     // Stream
     void*    stream)
 {
+  OPENINFER_FFI_GUARD_BEGIN
     auto paged_kv = make_paged_kv(
         kv_data, k_offset_elems, v_offset_elems,
         page_indices, page_indptr, last_page_len_d,
@@ -474,6 +486,7 @@ int batch_prefill_paged_cuda_with_cta_tile_q(
     uint32_t cta_tile_q = resolve_prefill_cta_tile_q(
         packed_qo_len, head_dim, cta_tile_q_override);
     if (cta_tile_q == 0) {
+        openinfer_ffi_set_last_error("invalid cta_tile_q override");
         return -1;
     }
 
@@ -497,6 +510,7 @@ int batch_prefill_paged_cuda_with_cta_tile_q(
                 s));
     });
     return result;
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 int batch_prefill_paged_cuda(
@@ -525,13 +539,15 @@ int batch_prefill_paged_cuda(
     float    sm_scale,
     void*    stream)
 {
+  OPENINFER_FFI_GUARD_BEGIN
     return batch_prefill_paged_cuda_with_cta_tile_q(
         q, output, kv_data, k_offset_elems, v_offset_elems,
         page_indices, page_indptr, last_page_len_d, q_indptr,
         request_indices, qo_tile_indices, kv_tile_indices,
         kv_chunk_size_ptr, total_num_rows, num_qo_heads, num_kv_heads,
         head_dim, page_size, seq_len, batch_size, padded_batch_size,
         stride_page, sm_scale, /*cta_tile_q_override=*/0, stream);
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 // ---------------------------------------------------------------------------
@@ -562,6 +578,7 @@ int single_prefill_cuda(
     // Stream
     void*    stream)
 {
+  OPENINFER_FFI_GUARD_BEGIN
     // Q/O strides: col-major [q_dim, seq_len]
     uint32_t q_stride_n  = num_qo_heads * head_dim;   // stride between tokens
     uint32_t q_stride_h  = head_dim;                   // stride between heads
@@ -605,6 +622,7 @@ int single_prefill_cuda(
             params,
             /*tmp=*/nullptr,
             reinterpret_cast<cudaStream_t>(stream)));
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 int single_prefill_nhd_noncausal_cuda(
@@ -623,6 +641,7 @@ int single_prefill_nhd_noncausal_cuda(
     float    sm_scale,
     void*    stream)
 {
+  OPENINFER_FFI_GUARD_BEGIN
     if (q == nullptr || output == nullptr || k_cache == nullptr || v_cache == nullptr ||
         num_qo_heads <= 0 || num_kv_heads <= 0 || head_dim != 128 ||
         seq_len <= 0 || kv_len <= 0 || max_seq_len < kv_len) {
@@ -669,6 +688,7 @@ int single_prefill_nhd_noncausal_cuda(
             params,
             /*tmp=*/nullptr,
             reinterpret_cast<cudaStream_t>(stream)));
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 // ---------------------------------------------------------------------------
@@ -696,6 +716,7 @@ int single_prefill_cuda_hd256(
     float    sm_scale,
     void*    stream)
 {
+  OPENINFER_FFI_GUARD_BEGIN
     uint32_t q_stride_n  = num_qo_heads * 256;
     uint32_t q_stride_h  = 256;
 
@@ -737,6 +758,7 @@ int single_prefill_cuda_hd256(
             params,
             /*tmp=*/nullptr,
             reinterpret_cast<cudaStream_t>(stream)));
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 // ---------------------------------------------------------------------------
@@ -764,6 +786,7 @@ int paged_attention_decode_cuda_hd256(
     float    sm_scale,
     void*    stream)
 {
+  OPENINFER_FFI_GUARD_BEGIN
     auto paged_kv = make_paged_kv(
         kv_data, k_offset_elems, v_offset_elems,
         page_indices, page_indptr, last_page_len_d,
@@ -804,6 +827,7 @@ int paged_attention_decode_cuda_hd256(
             /*tmp_s=*/nullptr,
             /*enable_pdl=*/false,
             reinterpret_cast<cudaStream_t>(stream)));
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 // ---------------------------------------------------------------------------
@@ -836,6 +860,7 @@ int batch_prefill_paged_cuda_hd256(
     float    sm_scale,
     void*    stream)
 {
+  OPENINFER_FFI_GUARD_BEGIN
     auto paged_kv = make_paged_kv(
         kv_data, k_offset_elems, v_offset_elems,
         page_indices, page_indptr, last_page_len_d,
@@ -899,6 +924,7 @@ int batch_prefill_paged_cuda_hd256(
                 s));
     });
     return result;
+  OPENINFER_FFI_GUARD_END(-1)
 }
 
 } // extern "C"

openinfer-kernels/src/ffi/glm52/topk.rs

@@ -1,4 +1,4 @@
-use cudarc::driver::sys::{CUresult, CUstream};
+use cudarc::driver::sys::CUstream;
 
 unsafe extern "C" {
     pub fn glm52_flashinfer_topk_2048_cuda(
@@ -10,5 +10,5 @@ unsafe extern "C" {
         top_k: i32,
         max_len: i32,
         stream: CUstream,
-    ) -> CUresult;
+    ) -> i32;
 }

openinfer-kernels/src/ffi/shared.rs

@@ -668,3 +668,7 @@ unsafe extern "C" {
     );
 
 }
+
+unsafe extern "C" {
+    pub fn openinfer_kernels_last_error() -> *const std::os::raw::c_char;
+}