@@ -410,16 +410,16 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
410410 auto get_tensor_config = [&]() -> tensor_config {
411411 // standard attention
412412 if (std::regex_match (tensor_name, pattern_q_weight) || std::regex_match (tensor_name, pattern_kv_weight)) {
413- return get_tensor_config_impl (GGML_BACKEND_SPLIT_AXIS_1, " attn_output.weight" );
413+ return get_tensor_config_impl (GGML_BACKEND_SPLIT_AXIS_1, " attn_output.weight" , " ssm_out.weight " );
414414 }
415415 if (std::regex_match (tensor_name, pattern_q_bias) || std::regex_match (tensor_name, pattern_kv_bias)) {
416- return get_tensor_config_impl (GGML_BACKEND_SPLIT_AXIS_0, " attn_output.weight" );
416+ return get_tensor_config_impl (GGML_BACKEND_SPLIT_AXIS_0, " attn_output.weight" , " ssm_out.weight " );
417417 }
418418 if (std::regex_match (tensor_name, pattern_qkv_weight)) {
419- return get_tensor_config_impl (GGML_BACKEND_SPLIT_AXIS_1);
419+ return get_tensor_config_impl (GGML_BACKEND_SPLIT_AXIS_1, " attn_output.weight " , " ssm_out.weight " );
420420 }
421421 if ( std::regex_match (tensor_name, pattern_qkv_bias)) {
422- return get_tensor_config_impl (GGML_BACKEND_SPLIT_AXIS_0);
422+ return get_tensor_config_impl (GGML_BACKEND_SPLIT_AXIS_0, " attn_output.weight " , " ssm_out.weight " );
423423 }
424424 if (std::regex_match (tensor_name, pattern_qk_norm)) {
425425 return get_tensor_config_impl (tensor->ne [1 ] == 1 ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_1, " attn_output.weight" );
@@ -435,7 +435,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
435435 }
436436
437437 if (std::regex_match (tensor_name, pattern_attn_gate_weight)) {
438- return get_tensor_config_impl (GGML_BACKEND_SPLIT_AXIS_1);
438+ return get_tensor_config_impl (GGML_BACKEND_SPLIT_AXIS_1, " attn_output.weight " , " ssm_out.weight " );
439439 }
440440 if (std::regex_match (tensor_name, pattern_ssm_dt) || std::regex_match (tensor_name, pattern_ssm_a)) {
441441 return get_tensor_config_impl (GGML_BACKEND_SPLIT_AXIS_0, " ssm_out.weight" );
0 commit comments