diff --git a/ggml/src/ggml-cuda/argsort.cu b/ggml/src/ggml-cuda/argsort.cu index 26af90025..24115da09 100644 --- a/ggml/src/ggml-cuda/argsort.cu +++ b/ggml/src/ggml-cuda/argsort.cu @@ -51,9 +51,12 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, cudaStream_t stream) { ggml_cuda_pool_alloc temp_indices_alloc(pool, ncols * nrows); ggml_cuda_pool_alloc temp_keys_alloc(pool, ncols * nrows); + // Device*Sort algorithms currently do not allow for in-place sorting/aliasing of input/outputs + ggml_cuda_pool_alloc temp_keys_out_alloc(pool, ncols * nrows); int * temp_indices = temp_indices_alloc.get(); float * temp_keys = temp_keys_alloc.get(); + float * temp_keys_out = temp_keys_out_alloc.get(); static const int block_size = 256; const dim3 grid_size((ncols + block_size - 1) / block_size, nrows); @@ -85,18 +88,18 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, if (order == GGML_SORT_ORDER_ASC) { if (nrows == 1) { - CUDA_CHECK(DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, temp_keys, temp_keys, // keys (in-place) + CUDA_CHECK(DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, temp_keys, temp_keys_out, // keys in, keys out temp_indices, dst, // values (indices) ncols, 0, sizeof(float) * 8, stream)); } else if (is_capturing) { CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs( - nullptr, temp_storage_bytes, temp_keys, temp_keys, // keys (in-place) + nullptr, temp_storage_bytes, temp_keys, temp_keys_out, // keys in, keys out temp_indices, dst, // values (indices) ncols * nrows, nrows, // num items, num segments offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); } else { CUDA_CHECK(DeviceSegmentedSort::SortPairs(nullptr, temp_storage_bytes, temp_keys, - temp_keys, // keys (in-place) + temp_keys_out, // keys out temp_indices, dst, // values (indices) ncols * nrows, nrows, // num items, num segments offset_iterator, offset_iterator + 1, stream)); @@ -104,15 +107,15 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, } else { if (nrows == 1) { CUDA_CHECK(DeviceRadixSort::SortPairsDescending(nullptr, temp_storage_bytes, temp_keys, - temp_keys, // keys (in-place) + temp_keys_out, // keys out temp_indices, dst, // values (indices) ncols, 0, sizeof(float) * 8, stream)); } else if (is_capturing) { CUDA_CHECK(DeviceSegmentedRadixSort::SortPairsDescending( - nullptr, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows, + nullptr, temp_storage_bytes, temp_keys, temp_keys_out, temp_indices, dst, ncols * nrows, nrows, offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); } else { - CUDA_CHECK(DeviceSegmentedSort::SortPairsDescending(nullptr, temp_storage_bytes, temp_keys, temp_keys, + CUDA_CHECK(DeviceSegmentedSort::SortPairsDescending(nullptr, temp_storage_bytes, temp_keys, temp_keys_out, temp_indices, dst, ncols * nrows, nrows, offset_iterator, offset_iterator + 1, stream)); } @@ -124,31 +127,31 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, if (order == GGML_SORT_ORDER_ASC) { if (nrows == 1) { CUDA_CHECK(DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, - temp_keys, // keys (in-place) + temp_keys_out, // keys out temp_indices, dst, // values (indices) ncols, 0, sizeof(float) * 8, stream)); } else if (is_capturing) { - CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, + CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys_out, temp_indices, dst, ncols * nrows, nrows, offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); } else { - CUDA_CHECK(DeviceSegmentedSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, + CUDA_CHECK(DeviceSegmentedSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys_out, temp_indices, dst, ncols * nrows, nrows, offset_iterator, offset_iterator + 1, stream)); } } else { if (nrows == 1) { CUDA_CHECK(DeviceRadixSort::SortPairsDescending(d_temp_storage, temp_storage_bytes, temp_keys, - temp_keys, // keys (in-place) + temp_keys_out, // keys out temp_indices, dst, // values (indices) ncols, 0, sizeof(float) * 8, stream)); } else if (is_capturing) { CUDA_CHECK(DeviceSegmentedRadixSort::SortPairsDescending( - d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows, + d_temp_storage, temp_storage_bytes, temp_keys, temp_keys_out, temp_indices, dst, ncols * nrows, nrows, offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); } else { CUDA_CHECK(DeviceSegmentedSort::SortPairsDescending(d_temp_storage, temp_storage_bytes, temp_keys, - temp_keys, temp_indices, dst, ncols * nrows, nrows, + temp_keys_out, temp_indices, dst, ncols * nrows, nrows, offset_iterator, offset_iterator + 1, stream)); } } diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index b510cb957..0dcfad3af 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -496,14 +496,29 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexe return res; } -ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc(ggml_metal_library_t lib, ggml_op op) { +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc(ggml_metal_library_t lib, const ggml_tensor * op) { const char * name = nullptr; - switch (op) { - case GGML_OP_DSV4_HC_COMB: name = "kernel_dsv4_hc_comb_f32"; break; - case GGML_OP_DSV4_HC_PRE: name = "kernel_dsv4_hc_pre_f32"; break; - case GGML_OP_DSV4_HC_POST: name = "kernel_dsv4_hc_post_f32"; break; - default: GGML_ABORT("fatal error"); + switch (op->op) { + case GGML_OP_DSV4_HC_COMB: + name = "kernel_dsv4_hc_comb_f32"; + break; + case GGML_OP_DSV4_HC_PRE: + if (ggml_get_op_params_i32(op, 1) != 0) { + name = "kernel_dsv4_hc_pre_gated_f32"; + } else { + name = "kernel_dsv4_hc_pre_f32"; + } + break; + case GGML_OP_DSV4_HC_POST: + if (op->src[3]) { + name = "kernel_dsv4_hc_post_f32"; + } else { + name = "kernel_dsv4_hc_post_nocomb_f32"; + } + break; + default: + GGML_ABORT("fatal error"); } ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index f6243ffbd..0514f9ef0 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -126,7 +126,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_ad struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_tri (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexer (ggml_metal_library_t lib, const struct ggml_tensor * op); -struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc (ggml_metal_library_t lib, enum ggml_op op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan (ggml_metal_library_t lib, const struct ggml_tensor * op, bool tail); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 78164bdfa..5dafae84f 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1739,6 +1739,12 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te op->src[0]->ne[0] != 576) { return false; } + if (op->src[1]->ne[0] == 72 && op->src[1]->ne[0] != op->src[2]->ne[0]) { + return false; + } + if (op->src[1]->ne[0] < op->src[2]->ne[0]) { + return false; + } if (op->src[1]->type != op->src[2]->type) { return false; } @@ -1808,8 +1814,6 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && op->src[0]->ne[1] == 4 && - op->src[1]->ne[0] == 4 && - op->src[1]->ne[2] == 1 && ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]); case GGML_OP_DSV4_HC_POST: @@ -1817,17 +1821,15 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_F32 && - op->src[3] != NULL && - op->src[3]->type == GGML_TYPE_F32 && + (op->src[3] == NULL || op->src[3]->type == GGML_TYPE_F32) && op->type == GGML_TYPE_F32 && op->src[1]->ne[1] == 4 && op->src[2]->ne[0] == 4 && - op->src[3]->ne[0] == 4 && - op->src[3]->ne[1] == 4 && + (op->src[3] == NULL || (op->src[3]->ne[0] == 4 && op->src[3]->ne[1] == 4)) && ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]) && ggml_is_contiguous_rows(op->src[2]) && - ggml_is_contiguous_rows(op->src[3]); + (op->src[3] == NULL || ggml_is_contiguous_rows(op->src[3])); case GGML_OP_SSM_SCAN: return has_simdgroup_reduction; case GGML_OP_SSM_CONV: diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 7a2c65aaa..d84ca937b 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -1283,8 +1283,10 @@ typedef struct { uint64_t nb_x2; uint64_t nb_w0; uint64_t nb_w1; + uint64_t nb_w2; uint64_t nb_d0; uint64_t nb_d1; + float scale; } ggml_metal_kargs_dsv4_hc_pre; typedef struct { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index cc1bebfaa..77c399bdb 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -1405,7 +1405,7 @@ int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); ggml_metal_encoder_t enc = ctx->enc; - auto pipeline = ggml_metal_library_get_pipeline_dsv4_hc(ctx->lib, op->op); + auto pipeline = ggml_metal_library_get_pipeline_dsv4_hc(ctx->lib, op); ggml_metal_encoder_set_pipeline(enc, pipeline); @@ -1467,8 +1467,10 @@ int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) { /*.nb_x2 =*/ x->nb[2], /*.nb_w0 =*/ weights->nb[0], /*.nb_w1 =*/ weights->nb[1], + /*.nb_w2 =*/ weights->nb[2], /*.nb_d0 =*/ op->nb[0], /*.nb_d1 =*/ op->nb[1], + /*.scale =*/ ggml_get_op_params_f32(op, 0), }; ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); @@ -1491,7 +1493,6 @@ int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) { GGML_ASSERT(x->type == GGML_TYPE_F32); GGML_ASSERT(residual->type == GGML_TYPE_F32); GGML_ASSERT(post->type == GGML_TYPE_F32); - GGML_ASSERT(comb->type == GGML_TYPE_F32); GGML_ASSERT(op->type == GGML_TYPE_F32); GGML_ASSERT(residual->ne[1] == 4); @@ -1505,9 +1506,9 @@ int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) { /*.nb_r2 =*/ residual->nb[2], /*.nb_p0 =*/ post->nb[0], /*.nb_p1 =*/ post->nb[1], - /*.nb_c0 =*/ comb->nb[0], - /*.nb_c1 =*/ comb->nb[1], - /*.nb_c2 =*/ comb->nb[2], + /*.nb_c0 =*/ comb ? comb->nb[0] : 0, + /*.nb_c1 =*/ comb ? comb->nb[1] : 0, + /*.nb_c2 =*/ comb ? comb->nb[2] : 0, /*.nb_d0 =*/ op->nb[0], /*.nb_d1 =*/ op->nb[1], /*.nb_d2 =*/ op->nb[2], @@ -1517,8 +1518,12 @@ int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(x), 1); ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(residual), 2); ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(post), 3); - ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(comb), 4); - ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 5); + if (comb) { + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(comb), 4); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 5); + } else { + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 4); + } const int n_tiles = (args.n_embd + 31)/32; const int nsg = std::min(4, n_tiles); diff --git a/ggml/src/ggml-metal/kernels/misc.metal b/ggml/src/ggml-metal/kernels/misc.metal index 11104b4d8..15a18e04a 100644 --- a/ggml/src/ggml-metal/kernels/misc.metal +++ b/ggml/src/ggml-metal/kernels/misc.metal @@ -531,7 +531,73 @@ kernel void kernel_dsv4_hc_pre_f32( result = fma(*(device const float *) (xb + ih*args.nb_x1), w[ih], result); } - *(device float *) (dst + i0*args.nb_d0 + it*args.nb_d1) = result; + *(device float *) (dst + i0*args.nb_d0 + it*args.nb_d1) = args.scale*result; +} + +kernel void kernel_dsv4_hc_pre_gated_f32( + constant ggml_metal_kargs_dsv4_hc_pre & args, + device const char * x, + device const char * gate, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + constexpr ushort hc = 4; + + const int it = tgpig.y; + const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg; + + if (i0 >= args.n_embd) { + return; + } + + device const char * xb = x + i0*args.nb_x0 + it*args.nb_x2; + device const char * gb = gate + i0*args.nb_w0 + it*args.nb_w2; + float result = 0.0f; + FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) { + const float g = 1.0f/(1.0f + exp(-*(device const float *) (gb + ih*args.nb_w1))); + result = fma(*(device const float *) (xb + ih*args.nb_x1), g, result); + } + + *(device float *) (dst + i0*args.nb_d0 + it*args.nb_d1) = args.scale*result; +} + +kernel void kernel_dsv4_hc_post_nocomb_f32( + constant ggml_metal_kargs_dsv4_hc_post & args, + device const char * x, + device const char * residual, + device const char * post, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + constexpr ushort hc = 4; + + const int it = tgpig.y; + const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg; + + float post_lane = 0.0f; + if (tiisg < hc) { + post_lane = *(device const float *) (post + tiisg*args.nb_p0 + it*args.nb_p1); + } + + float post_reg[hc]; + FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) { + post_reg[idst] = simd_shuffle(post_lane, idst); + } + + if (i0 >= args.n_embd) { + return; + } + + const float xv = *(device const float *) (x + i0*args.nb_x0 + it*args.nb_x1); + device const char * rb = residual + i0*args.nb_r0 + it*args.nb_r2; + FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) { + const float rv = *(device const float *) (rb + idst*args.nb_r1); + *(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = xv*post_reg[idst] + rv; + } } kernel void kernel_dsv4_hc_post_f32( diff --git a/include/llama.h b/include/llama.h index fb38ca8d7..00618c6d5 100644 --- a/include/llama.h +++ b/include/llama.h @@ -80,6 +80,7 @@ extern "C" { LLAMA_VOCAB_TYPE_UGM = 4, // T5 tokenizer based on Unigram LLAMA_VOCAB_TYPE_RWKV = 5, // RWKV tokenizer based on greedy tokenization LLAMA_VOCAB_TYPE_PLAMO2 = 6, // PLaMo-2 tokenizer based on Aho-Corasick with dynamic programming + LLAMA_VOCAB_TYPE_TEST = 7, // Dummy tokenizer for testing: rolling hash of fixed-size chunks -> tokens, tokens -> hex }; enum llama_rope_type { diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 0f5155b2e..39160a417 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -387,13 +387,13 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_TOKENIZER_SCORES, scores); add_kv(LLM_KV_TOKENIZER_MERGES, vocab.get_bpe_merges()); // FIXME llama_token is type i32 but when reading in a GGUF file u32 is expected, not an issue for writing though - add_kv(LLM_KV_TOKENIZER_BOS_ID, uint32_t(vocab.token_bos())); - add_kv(LLM_KV_TOKENIZER_EOS_ID, uint32_t(vocab.token_eos())); - add_kv(LLM_KV_TOKENIZER_EOT_ID, uint32_t(vocab.token_eot())); - add_kv(LLM_KV_TOKENIZER_EOM_ID, uint32_t(vocab.token_eom())); - add_kv(LLM_KV_TOKENIZER_UNK_ID, uint32_t(vocab.token_unk())); - add_kv(LLM_KV_TOKENIZER_SEP_ID, uint32_t(vocab.token_sep())); - add_kv(LLM_KV_TOKENIZER_PAD_ID, uint32_t(vocab.token_pad())); + if (vocab.token_bos() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_BOS_ID, uint32_t(vocab.token_bos())); } + if (vocab.token_eos() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_EOS_ID, uint32_t(vocab.token_eos())); } + if (vocab.token_eot() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_EOT_ID, uint32_t(vocab.token_eot())); } + if (vocab.token_eom() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_EOM_ID, uint32_t(vocab.token_eom())); } + if (vocab.token_unk() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_UNK_ID, uint32_t(vocab.token_unk())); } + if (vocab.token_sep() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_SEP_ID, uint32_t(vocab.token_sep())); } + if (vocab.token_pad() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_PAD_ID, uint32_t(vocab.token_pad())); } // add_kv(LLM_KV_TOKENIZER_CLS_ID, uint32_t(vocab.token_bos())); // deprecated // add_kv(LLM_KV_TOKENIZER_MASK_ID, ???); add_kv(LLM_KV_TOKENIZER_ADD_BOS, vocab.get_add_bos()); @@ -404,12 +404,12 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP, vocab.get_precompiled_charsmap()); // add_kv(LLM_KV_TOKENIZER_HF_JSON, ???); // add_kv(LLM_KV_TOKENIZER_RWKV, ???); - add_kv(LLM_KV_TOKENIZER_FIM_PRE_ID, uint32_t(vocab.token_fim_pre())); - add_kv(LLM_KV_TOKENIZER_FIM_SUF_ID, uint32_t(vocab.token_fim_suf())); - add_kv(LLM_KV_TOKENIZER_FIM_MID_ID, uint32_t(vocab.token_fim_mid())); - add_kv(LLM_KV_TOKENIZER_FIM_PAD_ID, uint32_t(vocab.token_fim_pad())); - add_kv(LLM_KV_TOKENIZER_FIM_REP_ID, uint32_t(vocab.token_fim_rep())); - add_kv(LLM_KV_TOKENIZER_FIM_SEP_ID, uint32_t(vocab.token_fim_sep())); + if (vocab.token_fim_pre() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_FIM_PRE_ID, uint32_t(vocab.token_fim_pre())); } + if (vocab.token_fim_suf() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_FIM_SUF_ID, uint32_t(vocab.token_fim_suf())); } + if (vocab.token_fim_mid() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_FIM_MID_ID, uint32_t(vocab.token_fim_mid())); } + if (vocab.token_fim_pad() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_FIM_PAD_ID, uint32_t(vocab.token_fim_pad())); } + if (vocab.token_fim_rep() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_FIM_REP_ID, uint32_t(vocab.token_fim_rep())); } + if (vocab.token_fim_sep() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_FIM_SEP_ID, uint32_t(vocab.token_fim_sep())); } // TODO: implement LoRA support // add_kv(LLM_KV_ADAPTER_TYPE, ???); diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index b84005995..93bdfc67a 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2323,6 +2323,16 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { special_unk_id = LLAMA_TOKEN_NULL; special_sep_id = LLAMA_TOKEN_NULL; special_pad_id = LLAMA_TOKEN_NULL; + } else if (tokenizer_model == "test") { + type = LLAMA_VOCAB_TYPE_TEST; + + // default special tokens + special_bos_id = LLAMA_TOKEN_NULL; + special_eos_id = LLAMA_TOKEN_NULL; + special_unk_id = LLAMA_TOKEN_NULL; + special_sep_id = LLAMA_TOKEN_NULL; + special_pad_id = LLAMA_TOKEN_NULL; + special_mask_id = LLAMA_TOKEN_NULL; } else if (tokenizer_model == "plamo2") { type = LLAMA_VOCAB_TYPE_PLAMO2; @@ -3390,6 +3400,7 @@ std::string llama_vocab::impl::type_name() const{ case LLAMA_VOCAB_TYPE_UGM: return "UGM"; case LLAMA_VOCAB_TYPE_RWKV: return "RWKV"; case LLAMA_VOCAB_TYPE_PLAMO2: return "PLaMo2"; + case LLAMA_VOCAB_TYPE_TEST: return "TEST"; default: return "unknown"; } } @@ -3483,6 +3494,9 @@ void llama_vocab::impl::init_tokenizer(enum llama_vocab_type type) { case LLAMA_VOCAB_TYPE_PLAMO2: tokenizer = std::make_unique(vocab); break; + case LLAMA_VOCAB_TYPE_TEST: + tokenizer = std::make_unique(); + break; default: GGML_ABORT("unsupported vocab type"); } @@ -3882,6 +3896,42 @@ std::vector llama_vocab::impl::tokenize( } } } break; + case LLAMA_VOCAB_TYPE_TEST: + { + const uint32_t n_vocab = vocab.n_tokens(); + constexpr size_t chunk_size = 5; + + // reserve output to avoid repeated reallocations + size_t n_tokens = 0; + for (const auto & fragment : fragment_buffer) { + if (fragment.type == FRAGMENT_BUFFER_VARIANT_TYPE_RAW_TEXT) { + n_tokens += (fragment.length + chunk_size - 1) / chunk_size; + } else { + ++n_tokens; + } + } + output.reserve(output.size() + n_tokens); + + for (const auto & fragment : fragment_buffer) { + if (fragment.type == FRAGMENT_BUFFER_VARIANT_TYPE_RAW_TEXT) { + const auto & text = fragment.raw_text; + const size_t begin = fragment.offset; + const size_t end = begin + fragment.length; + size_t pos = begin; + while (pos < end) { + const size_t n = std::min(chunk_size, end - pos); + uint64_t hash = 0; + for (size_t i = 0; i < n; ++i) { + hash = hash*31 + (uint8_t) text[pos + i]; + } + output.push_back((llama_token)(hash % n_vocab)); + pos += n; + } + } else { // if (fragment.type == FRAGMENT_BUFFER_VARIANT_TYPE_TOKEN) + output.push_back(fragment.token); + } + } + } break; case LLAMA_VOCAB_TYPE_NONE: GGML_ABORT("fatal error"); } @@ -3984,6 +4034,11 @@ int32_t llama_vocab::impl::token_to_piece(llama_token token, char * buf, int32_t memcpy(buf, result.data(), result.size()); return (int)result.size(); } + case LLAMA_VOCAB_TYPE_TEST: { + // tokens -> text: simply stringify the token id in hex + std::string result = format("%x", token); + return _try_copy(result.data(), result.size()); + } case LLAMA_VOCAB_TYPE_PLAMO2: { // PLaMo-2 uses similar token handling as BPE/SPM if (vocab.is_byte(token)) { @@ -4258,6 +4313,9 @@ llama_token llama_vocab::byte_to_token(uint8_t ch) const { snprintf(hex_str, sizeof(hex_str), "<0x%02X>", ch); return pimpl->token_to_id.at(hex_str); } + case LLAMA_VOCAB_TYPE_TEST: + // TEST tokens have no byte-level mapping + return LLAMA_TOKEN_NULL; default: GGML_ASSERT_CONTINUE(false); return 0;