diff --git a/openinfer-core/src/ops.rs b/openinfer-core/src/ops.rs index 389ea51a..d43f57af 100644 --- a/openinfer-core/src/ops.rs +++ b/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, diff --git a/openinfer-kernels/csrc/glm52/glm52_topk.cu b/openinfer-kernels/csrc/glm52/glm52_topk.cu index 54ca0cde..d24dddf7 100644 --- a/openinfer-kernels/csrc/glm52/glm52_topk.cu +++ b/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 #include @@ -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(CUDA_ERROR_INVALID_VALUE); } @@ -57,4 +60,5 @@ extern "C" int glm52_flashinfer_topk_2048_cuda( return static_cast(CUDA_ERROR_LAUNCH_FAILED); } return static_cast(CUDA_SUCCESS); + OPENINFER_FFI_GUARD_END(-1) } diff --git a/openinfer-kernels/csrc/kimi_k2/kimi_mla.cu b/openinfer-kernels/csrc/kimi_k2/kimi_mla.cu index d166bf7d..e135c734 100644 --- a/openinfer-kernels/csrc/kimi_k2/kimi_mla.cu +++ b/openinfer-kernels/csrc/kimi_k2/kimi_mla.cu @@ -1,3 +1,5 @@ +#include "../shared/ffi_guard.cuh" + #include #include #include @@ -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(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(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" diff --git a/openinfer-kernels/csrc/shared/ffi_guard.cuh b/openinfer-kernels/csrc/shared/ffi_guard.cuh new file mode 100644 index 00000000..63589ea0 --- /dev/null +++ b/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 + +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; \ + } diff --git a/openinfer-kernels/csrc/shared/ffi_last_error.cu b/openinfer-kernels/csrc/shared/ffi_last_error.cu new file mode 100644 index 00000000..dc667c9c --- /dev/null +++ b/openinfer-kernels/csrc/shared/ffi_last_error.cu @@ -0,0 +1,13 @@ +#include "ffi_guard.cuh" + +#include + +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(); } diff --git a/openinfer-kernels/csrc/shared/flashinfer_sampling.cu b/openinfer-kernels/csrc/shared/flashinfer_sampling.cu index 6e879075..92d35952 100644 --- a/openinfer-kernels/csrc/shared/flashinfer_sampling.cu +++ b/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(err); + OPENINFER_FFI_GUARD_END(-1) } diff --git a/openinfer-kernels/csrc/shared/paged_attention.cu b/openinfer-kernels/csrc/shared/paged_attention.cu index 01f4e82c..02ca02b6 100644 --- a/openinfer-kernels/csrc/shared/paged_attention.cu +++ b/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 #include #include @@ -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(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(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(key), reinterpret_cast(value), reinterpret_cast(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(src_stride_n), // V has same layout as K static_cast(src_stride_h), reinterpret_cast(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((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,6 +539,7 @@ 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, @@ -532,6 +547,7 @@ int batch_prefill_paged_cuda( 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(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(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(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(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" diff --git a/openinfer-kernels/src/ffi/glm52/topk.rs b/openinfer-kernels/src/ffi/glm52/topk.rs index 89c7e9db..45bed4aa 100644 --- a/openinfer-kernels/src/ffi/glm52/topk.rs +++ b/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; } diff --git a/openinfer-kernels/src/ffi/shared.rs b/openinfer-kernels/src/ffi/shared.rs index e78f0ac9..5cf52327 100644 --- a/openinfer-kernels/src/ffi/shared.rs +++ b/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; +} diff --git a/openinfer-kernels/src/ops.rs b/openinfer-kernels/src/ops.rs index 9dedc4b2..a73755ea 100644 --- a/openinfer-kernels/src/ops.rs +++ b/openinfer-kernels/src/ops.rs @@ -17,10 +17,11 @@ mod norm; mod sampling; pub use attention::{ - PrefillPagedPlan, dflash_qk_norm_rope_into, paged_attention_batch_decode_hd256_into, - paged_attention_batch_decode_into, paged_attention_batch_decode_split_kv_into, - prefill_attention_paged_into, qk_norm_partial_rope_batched_decode_hd256_into, - qk_norm_rope_batch_decode_into, single_prefill_nhd_noncausal_into, + PrefillPagedPlan, SUPPORTED_GQA_GROUP_SIZES, dflash_qk_norm_rope_into, + paged_attention_batch_decode_hd256_into, paged_attention_batch_decode_into, + paged_attention_batch_decode_split_kv_into, prefill_attention_paged_into, + qk_norm_partial_rope_batched_decode_hd256_into, qk_norm_rope_batch_decode_into, + single_prefill_nhd_noncausal_into, }; #[cfg(feature = "moe")] pub use deepep::{ @@ -67,3 +68,22 @@ pub use sampling::{ flashinfer_top1_batch_into, flashinfer_top1_row_states_bytes, gpu_sample_batch_into, markov_step_argmax_into, }; + +/// Calling thread's last FFI exception message, ready to append to an error; +/// empty unless `result` is the -1 sentinel set by the C++ guard. Public for +/// crates that call guarded FFI entries directly instead of through a wrapper. +pub fn ffi_exception_message(result: i32) -> String { + if result != -1 { + return String::new(); + } + let ptr = unsafe { crate::ffi::openinfer_kernels_last_error() }; + if ptr.is_null() { + return String::new(); + } + let text = unsafe { std::ffi::CStr::from_ptr(ptr) }.to_string_lossy(); + if text.is_empty() { + String::new() + } else { + format!(": {text}") + } +} diff --git a/openinfer-kernels/src/ops/attention.rs b/openinfer-kernels/src/ops/attention.rs index 55b301d3..334b194f 100644 --- a/openinfer-kernels/src/ops/attention.rs +++ b/openinfer-kernels/src/ops/attention.rs @@ -6,6 +6,10 @@ use crate::ffi; use crate::paged_kv::PagedKvLayout; use crate::tensor::{DeviceContext, DeviceVec, HiddenStates}; +/// GQA group sizes (query heads / KV heads) instantiated by FlashInfer's +/// `DISPATCH_GQA_GROUP_SIZE`; any other ratio throws at dispatch time. +pub const SUPPORTED_GQA_GROUP_SIZES: &[usize] = &[1, 2, 3, 4, 8]; + // ============================================================================ // Paged prefill (FlashInfer BatchPrefillWithPagedKVCache) // ============================================================================ @@ -607,7 +611,10 @@ pub fn prefill_attention_paged_into( stream, ); if result != 0 { - anyhow::bail!("paged_kv_scatter_cuda failed for layer {layer} with error {result}"); + anyhow::bail!( + "paged_kv_scatter_cuda failed for layer {layer} with error {result}{}", + crate::ops::ffi_exception_message(result) + ); } let result = ffi::batch_prefill_paged_cuda_with_cta_tile_q( @@ -638,7 +645,10 @@ pub fn prefill_attention_paged_into( stream, ); if result != 0 { - anyhow::bail!("batch_prefill_paged_cuda failed for layer {layer} with error {result}"); + anyhow::bail!( + "batch_prefill_paged_cuda failed for layer {layer} with error {result}{}", + crate::ops::ffi_exception_message(result) + ); } } @@ -831,7 +841,10 @@ pub fn single_prefill_nhd_noncausal_into( ) }; if result != 0 { - anyhow::bail!("single_prefill_nhd_noncausal_cuda failed with error {result}"); + anyhow::bail!( + "single_prefill_nhd_noncausal_cuda failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); } Ok(()) } @@ -963,7 +976,10 @@ pub fn paged_attention_batch_decode_into( ) }; if result != 0 { - anyhow::bail!("paged_kv_scatter_cuda (batch decode) failed with error {result}"); + anyhow::bail!( + "paged_kv_scatter_cuda (batch decode) failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); } // Step 2: Paged attention decode (batched) @@ -992,7 +1008,10 @@ pub fn paged_attention_batch_decode_into( ) }; if result != 0 { - anyhow::bail!("paged_attention_decode_cuda (batch) failed with error {result}"); + anyhow::bail!( + "paged_attention_decode_cuda (batch) failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); } Ok(()) @@ -1081,7 +1100,10 @@ pub fn paged_attention_batch_decode_split_kv_into( ) }; if result != 0 { - anyhow::bail!("paged_kv_scatter_cuda (batch split-K decode) failed with error {result}"); + anyhow::bail!( + "paged_kv_scatter_cuda (batch split-K decode) failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); } let sm_scale = 1.0f32 / (head_dim as f32).sqrt(); @@ -1114,7 +1136,10 @@ pub fn paged_attention_batch_decode_split_kv_into( ) }; if result != 0 { - anyhow::bail!("paged_attention_decode_split_kv_cuda (batch) failed with error {result}"); + anyhow::bail!( + "paged_attention_decode_split_kv_cuda (batch) failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); } Ok(()) @@ -1189,7 +1214,10 @@ pub fn paged_attention_batch_decode_hd256_into( ) }; if result != 0 { - anyhow::bail!("paged_kv_scatter_cuda (batch hd256 decode) failed with error {result}"); + anyhow::bail!( + "paged_kv_scatter_cuda (batch hd256 decode) failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); } let sm_scale = 1.0f32 / (head_dim as f32).sqrt(); @@ -1217,7 +1245,10 @@ pub fn paged_attention_batch_decode_hd256_into( ) }; if result != 0 { - anyhow::bail!("paged_attention_decode_cuda_hd256 (batch) failed with error {result}"); + anyhow::bail!( + "paged_attention_decode_cuda_hd256 (batch) failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); } Ok(()) diff --git a/openinfer-kernels/src/ops/glm52/topk.rs b/openinfer-kernels/src/ops/glm52/topk.rs index d1f89376..1f57b32c 100644 --- a/openinfer-kernels/src/ops/glm52/topk.rs +++ b/openinfer-kernels/src/ops/glm52/topk.rs @@ -1,4 +1,4 @@ -use anyhow::{Result, anyhow, ensure}; +use anyhow::{Result, ensure}; use cudarc::driver::{CudaSlice, DevicePtr, DevicePtrMut}; use crate::ffi; @@ -94,7 +94,10 @@ pub fn glm52_flashinfer_topk_2048_launch( ctx.stream.cu_stream(), ) }; - result - .result() - .map_err(|err| anyhow!("GLM5.2 FlashInfer top-k 2048 launch failed: {err}")) + ensure!( + result == 0, + "GLM5.2 FlashInfer top-k 2048 launch failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); + Ok(()) } diff --git a/openinfer-kernels/src/ops/kimi_k2/mla.rs b/openinfer-kernels/src/ops/kimi_k2/mla.rs index 96b96ac9..7ed177fa 100644 --- a/openinfer-kernels/src/ops/kimi_k2/mla.rs +++ b/openinfer-kernels/src/ops/kimi_k2/mla.rs @@ -332,7 +332,10 @@ pub fn kimi_mla_paged_kv_append( ) }; if result != 0 { - bail!("kimi_mla_paged_kv_append_cuda failed with cudaError={result}"); + bail!( + "kimi_mla_paged_kv_append_cuda failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); } Ok(()) } diff --git a/openinfer-kernels/src/ops/kimi_k2/mla_rt.rs b/openinfer-kernels/src/ops/kimi_k2/mla_rt.rs index d55bfae8..39a09a4a 100644 --- a/openinfer-kernels/src/ops/kimi_k2/mla_rt.rs +++ b/openinfer-kernels/src/ops/kimi_k2/mla_rt.rs @@ -158,7 +158,10 @@ pub fn kimi_flashinfer_batch_decode_mla_rt( ) }; if result != 0 { - bail!("kimi_flashinfer_batch_decode_mla_cuda failed with cudaError={result}"); + bail!( + "kimi_flashinfer_batch_decode_mla_cuda failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); } Ok(()) } @@ -320,7 +323,10 @@ pub fn kimi_flashinfer_single_prefill_mla_rt( ) }; if result != 0 { - bail!("kimi_flashinfer_single_prefill_mla_cuda failed with cudaError={result}"); + bail!( + "kimi_flashinfer_single_prefill_mla_cuda failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); } Ok(()) } diff --git a/openinfer-kernels/src/ops/sampling.rs b/openinfer-kernels/src/ops/sampling.rs index fe83bccb..6b7bc8b6 100644 --- a/openinfer-kernels/src/ops/sampling.rs +++ b/openinfer-kernels/src/ops/sampling.rs @@ -250,7 +250,11 @@ fn sample_uniform_batch_into( crate::tensor::active_cu_stream(ctx), ) }; - ensure!(err == 0, "batch sampling kernel failed: cudaError {err}"); + ensure!( + err == 0, + "batch sampling kernel failed with error {err}{}", + crate::ops::ffi_exception_message(err) + ); } let out = ctx diff --git a/openinfer-kernels/tests/ffi_exception_catch.rs b/openinfer-kernels/tests/ffi_exception_catch.rs new file mode 100644 index 00000000..197e722e --- /dev/null +++ b/openinfer-kernels/tests/ffi_exception_catch.rs @@ -0,0 +1,83 @@ +//! Drives decode at GQA group 7 (56/8) — unsupported even if group 5 is added +//! for Qwen3-14B — and asserts the FlashInfer throw surfaces as the -1 +//! sentinel plus a message. Before the guard this aborted the whole process. + +use std::ffi::{CStr, c_void}; +use std::ptr; + +use anyhow::{Result, ensure}; +use openinfer_kernels::ffi; + +unsafe extern "C" { + fn cudaMalloc(dev_ptr: *mut *mut c_void, size: usize) -> i32; + fn cudaFree(dev_ptr: *mut c_void) -> i32; + fn cudaSetDevice(device: i32) -> i32; +} + +struct DeviceBuf { + ptr: *mut c_void, +} + +impl DeviceBuf { + fn alloc(bytes: usize) -> Result { + let mut ptr = ptr::null_mut(); + let rc = unsafe { cudaMalloc(&raw mut ptr, bytes) }; + ensure!(rc == 0, "cudaMalloc failed: {rc}"); + Ok(Self { ptr }) + } +} + +impl Drop for DeviceBuf { + fn drop(&mut self) { + unsafe { cudaFree(self.ptr) }; + } +} + +#[test] +fn unsupported_gqa_group_returns_sentinel_with_message() -> Result<()> { + if unsafe { cudaSetDevice(0) } != 0 { + eprintln!("skipping: no CUDA device"); + return Ok(()); + } + let num_qo_heads = 56i32; + let num_kv_heads = 8i32; + let head_dim = 128i32; + let page_size = 16i32; + + let q = DeviceBuf::alloc((num_qo_heads * head_dim) as usize * 2)?; + let output = DeviceBuf::alloc((num_qo_heads * head_dim) as usize * 2)?; + let kv = DeviceBuf::alloc((2 * page_size * num_kv_heads * head_dim) as usize * 2)?; + // Shared backing for the six metadata arrays; the dispatch rejects the + // group size before any kernel reads them. + let meta = DeviceBuf::alloc(8 * std::mem::size_of::())?; + + let result = unsafe { + ffi::paged_attention_decode_cuda( + q.ptr.cast(), + output.ptr.cast(), + kv.ptr.cast(), + 0, + (page_size * num_kv_heads * head_dim) as i64, + meta.ptr.cast(), + meta.ptr.cast(), + meta.ptr.cast(), + meta.ptr.cast(), + meta.ptr.cast(), + meta.ptr.cast(), + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + 1, + (2 * page_size * num_kv_heads * head_dim) as i64, + 1.0, + ptr::null_mut(), + ) + }; + ensure!(result == -1, "expected the -1 sentinel, got {result}"); + + let msg = unsafe { CStr::from_ptr(ffi::openinfer_kernels_last_error()) }.to_string_lossy(); + // The exact wording belongs to FlashInfer. + ensure!(msg.contains("group_size"), "unexpected message: {msg}"); + Ok(()) +} diff --git a/openinfer-kernels/tests/glm52_indexer_smoke.rs b/openinfer-kernels/tests/glm52_indexer_smoke.rs index 9248f3ee..bb44b8ab 100644 --- a/openinfer-kernels/tests/glm52_indexer_smoke.rs +++ b/openinfer-kernels/tests/glm52_indexer_smoke.rs @@ -104,15 +104,6 @@ fn cu_check(r: CUresult) -> Result<()> { } /// Returns Ok(true) on skip (caller early-returns), Ok(false) on success. -fn cu_check_or_skip(r: CUresult) -> Result { - if r == CUresult::CUDA_ERROR_NOT_SUPPORTED { - eprintln!("skip: FilteredTopK not supported on this GPU (smem too small)"); - return Ok(true); - } - cu_check(r)?; - Ok(false) -} - const STREAM: CUstream = ptr::null_mut(); // ─── indexer cache: quant + pack → gather round-trip ────────────────────── @@ -307,7 +298,7 @@ fn flashinfer_topk_lengths_filter() -> Result<()> { let indices_dev = DeviceBuf::zeroed(top_k * size_of::())?; let values_dev = DeviceBuf::zeroed(top_k * size_of::())?; - let skipped = cu_check_or_skip(unsafe { + let result = unsafe { ffi::glm52_flashinfer_topk_2048_cuda( logits_dev.ptr as *const f32, indices_dev.ptr as *mut i32, @@ -318,10 +309,16 @@ fn flashinfer_topk_lengths_filter() -> Result<()> { max_len as i32, STREAM, ) - })?; - if skipped { + }; + if result == CUresult::CUDA_ERROR_NOT_SUPPORTED as i32 { + eprintln!("skip: FilteredTopK not supported on this GPU (smem too small)"); return Ok(()); } + anyhow::ensure!( + result == 0, + "glm52_flashinfer_topk_2048_cuda failed with error {result}{}", + openinfer_kernels::ops::ffi_exception_message(result) + ); cuda_check(unsafe { cudaDeviceSynchronize() })?; let indices: Vec = indices_dev.to_host(top_k)?; diff --git a/openinfer-qwen3/src/config.rs b/openinfer-qwen3/src/config.rs index a76c137d..e060b4f8 100644 --- a/openinfer-qwen3/src/config.rs +++ b/openinfer-qwen3/src/config.rs @@ -163,6 +163,16 @@ impl Config { let config_path = format!("{}/config.json", model_path); let content = fs::read_to_string(&config_path)?; let mut config: Config = serde_json::from_str(&content)?; + anyhow::ensure!( + config.num_key_value_heads > 0 + && config.num_attention_heads % config.num_key_value_heads == 0 + && openinfer_core::ops::SUPPORTED_GQA_GROUP_SIZES + .contains(&(config.num_attention_heads / config.num_key_value_heads)), + "unsupported GQA group size {}/{}; the attention kernels instantiate {:?}", + config.num_attention_heads, + config.num_key_value_heads, + openinfer_core::ops::SUPPORTED_GQA_GROUP_SIZES, + ); config.stop_token_ids = Self::load_stop_token_ids(model_path, config.eos_token_id)?; Ok(config) } diff --git a/openinfer-qwen3/src/kernel_bench.rs b/openinfer-qwen3/src/kernel_bench.rs index 5ffb6671..744db316 100644 --- a/openinfer-qwen3/src/kernel_bench.rs +++ b/openinfer-qwen3/src/kernel_bench.rs @@ -513,8 +513,9 @@ impl AttentionDecodeCase { }; if result != 0 { bail!( - "{} paged attention failed with error {result}", - path.name(self.split_config) + "{} paged attention failed with error {result}{}", + path.name(self.split_config), + openinfer_kernels::ops::ffi_exception_message(result) ); } Ok(result) @@ -823,7 +824,10 @@ impl AttentionPrefillCase { ) }; if result != 0 { - bail!("segmented paged_kv_scatter_cuda failed with error {result}"); + bail!( + "segmented paged_kv_scatter_cuda failed with error {result}{}", + openinfer_kernels::ops::ffi_exception_message(result) + ); } Ok(()) } @@ -885,7 +889,10 @@ impl AttentionPrefillCase { ) }; if result != 0 { - bail!("segmented batch_prefill_paged_cuda failed with error {result}"); + bail!( + "segmented batch_prefill_paged_cuda failed with error {result}{}", + openinfer_kernels::ops::ffi_exception_message(result) + ); } Ok(()) } @@ -1017,7 +1024,10 @@ impl SinglePrefillCase { ) }; if result != 0 { - bail!("single_prefill_cuda failed with error {result}"); + bail!( + "single_prefill_cuda failed with error {result}{}", + openinfer_kernels::ops::ffi_exception_message(result) + ); } Ok(()) } diff --git a/openinfer-qwen35-4b/src/prefill.rs b/openinfer-qwen35-4b/src/prefill.rs index a8af94f6..17033156 100644 --- a/openinfer-qwen35-4b/src/prefill.rs +++ b/openinfer-qwen35-4b/src/prefill.rs @@ -387,7 +387,8 @@ impl Qwen35Model { }; anyhow::ensure!( result == 0, - "batch_prefill_paged_cuda_hd256 failed: {result}" + "batch_prefill_paged_cuda_hd256 failed: {result}{}", + openinfer_kernels::ops::ffi_exception_message(result) ); }