Recover the model runtime from panics

This commit is contained in:
Georg Bauer
2026-07-25 13:00:20 +02:00
parent 90400701fa
commit 8f91a06336
2 changed files with 159 additions and 93 deletions

View File

@@ -1651,7 +1651,10 @@ fn encode_batch_layer(
"batch raw prefill attention",
)?;
} else if pos == 0 && ratio == 4 && compressed_rows > shape.indexer_top_k as u32 {
let indexer = state.indexer.as_ref().unwrap();
let indexer = state
.indexer
.as_ref()
.ok_or("ratio-4 layer is missing indexer state")?;
call(
unsafe {
ds4_gpu_indexer_scores_prefill_tensor(
@@ -1683,7 +1686,10 @@ fn encode_batch_layer(
},
"batch indexer top-k",
)?;
let compression = state.compression.as_ref().unwrap();
let compression = state
.compression
.as_ref()
.ok_or("compressed layer is missing attention state")?;
call(
unsafe {
ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
@@ -1712,7 +1718,10 @@ fn encode_batch_layer(
"batch indexed prefill attention",
)?;
} else if pos == 0 {
let compression = state.compression.as_ref().unwrap();
let compression = state
.compression
.as_ref()
.ok_or("compressed layer is missing attention state")?;
call(
unsafe {
ds4_gpu_attention_prefill_static_mixed_heads_tensor(
@@ -1759,7 +1768,10 @@ fn encode_batch_layer(
"batch resumed raw attention",
)?;
} else if ratio == 4 && compressed_rows > shape.indexer_top_k as u32 {
let indexer = state.indexer.as_ref().unwrap();
let indexer = state
.indexer
.as_ref()
.ok_or("ratio-4 layer is missing indexer state")?;
call(
unsafe {
ds4_gpu_indexer_scores_decode_batch_tensor(
@@ -1792,7 +1804,10 @@ fn encode_batch_layer(
},
"batch resumed indexer top-k",
)?;
let compression = state.compression.as_ref().unwrap();
let compression = state
.compression
.as_ref()
.ok_or("compressed layer is missing attention state")?;
call(
unsafe {
ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
@@ -1821,7 +1836,10 @@ fn encode_batch_layer(
"batch resumed indexed attention",
)?;
} else {
let compression = state.compression.as_ref().unwrap();
let compression = state
.compression
.as_ref()
.ok_or("compressed layer is missing attention state")?;
call(
unsafe {
ds4_gpu_attention_decode_mixed_batch_heads_tensor(