StatelessExecutor and/or grammar perf?

[phil-scott-78]

Description

Working on a gbnf generation library, and the perf started bugging me. Once I started getting some larger outputs of json the perf difference started getting really slow.

For example, when I run this command against the llama.cpp cli, I'm seeing total time = 3374.20 ms / 462 tokens

.\llama-cli -m B:\models\Qwen2.5-Coder-3B-Instruct-Q8_0.gguf --grammar-file B:\llama-src\LLamaSharp\LLama.Examples\Assets\json.gbnf -ngl 48 -no-cnv --prompt "give me a list of all nfl teams in the afc. include their team, city and state. group by division in json format"

Running the grammar example with the same prompt with the same model and the parameters tweaked to match the CLI (context, gpu layers), and I'm looking at about 16s to run the prompt.

I'm not 100% sure if it is related to grammar sampling or the StatelessExecutor though, and I'll sheepishly admit I'm just pulling levers here trying to track down the culprit. A bit of guidance to a culprit and I can get after it.

[martindevans]

Did you adjust the GpuLayerCount parameter in the example code?

[phil-scott-78]

I did. 48 on both. I'm at my kid's swim lessons but I can get the full changes here later tonight.

I really just did some simple stopwatch timings for now, and it felt like the stateless executor was much slower than the other executors even at baseline without the grammar, perhaps due to it doing a single decode at a time. And maybe that combined with a grammar really exacerbated the slowdown I saw.

I've been meaning to take the BatchedExecutor and turn it into a high level JsonExecutor so if the current design of the stateless executor is fine remaining simple that's cool.

[martindevans]

In principle the StatelessExecutor shouldn't be any slower - fundamentally all of the executors are calling llama_decode in a loop, getting one token at a time. if you can narrow it down to what's taking the time that would be very helpful!

[phil-scott-78]

ok, that's a good tip. I stopped caring about the executor and just got after the grammar. given this code:

var modelPath = UserSettings.GetModelPath();

var gbnf = (await File.ReadAllTextAsync("Assets/json.gbnf")).Trim();

var parameters = new ModelParams(modelPath)
{
    GpuLayerCount = 45,
    ContextSize = 4086,
};
using var model = await LLamaWeights.LoadFromFileAsync(parameters);
var ex = new StatelessExecutor(model, parameters) { ApplyTemplate = true };

var samplingPipeline = new DefaultSamplingPipeline()
{
    RepeatPenalty = 1f,
    FrequencyPenalty = 0,
    PresencePenalty = 0,
    TopK = 40,
    TopP = 0.95f,
    MinP = 0.05f,
    TypicalP = 1,
    Temperature = 0.8f,
    
    Grammar = new(gbnf, "root"),
};

var inferenceParams = new InferenceParams()
{
    SamplingPipeline = samplingPipeline,
    MaxTokens = 500,
    TokensKeep = 0,
};


const string prompt = "give me a list of all the football teams in the AFC. include their team, city and state. group by division in json format";
var sw = new Stopwatch();
sw.Start();
await foreach (var text in ex.InferAsync(prompt, inferenceParams))
    Console.Write(text);

sw.Start();
Console.WriteLine();
Console.WriteLine($"Elapsed time: {sw.Elapsed}");

with grammar is about 16s to run. comment out the grammar it is 4s.

llama-cli is about 4s with a grammar and without.

I profiled with and without the grammar. Calls to SafeLLamaSamplerChainHandle.Sample are taking 30ms on average with the grammar, 7.4ms without.

Of course, that call is right into llama.cpp, which is a bit of a pickle.

Taking a look at the output of llama.cpp, here's the sampler. Perhaps what they are doing with a grammar has diverged at some point?

sampler seed: 1542106340
sampler params:
        repeat_last_n = 64, repeat_penalty = 1.000, frequency_penalty = 0.000, presence_penalty = 0.000
        dry_multiplier = 0.000, dry_base = 1.750, dry_allowed_length = 2, dry_penalty_last_n = 4096
        top_k = 40, top_p = 0.950, min_p = 0.050, xtc_probability = 0.000, xtc_threshold = 0.100, typical_p = 1.000, temp = 0.800
        mirostat = 0, mirostat_lr = 0.100, mirostat_ent = 5.000
sampler chain: logits -> logit-bias -> penalties -> dry -> top-k -> typical -> top-p -> min-p -> xtc -> temp-ext -> dist
generate: n_ctx = 4096, n_batch = 2048, n_predict = -1, n_keep = 0

[phil-scott-78]

btw, here's the log from running the above

[llama Debug]: Loading library: 'llama'
[llama Info]: Detected OS Platform: 'WINDOWS'
[llama Debug]: Detected OS string: 'win-x64'
[llama Debug]: Detected extension string: '.dll'
[llama Debug]: Detected prefix string: ''
[llama Info]: NativeLibraryConfig Description:
- LibraryName: LLama
- Path: ''
- PreferCuda: True
- PreferVulkan: True
- PreferredAvxLevel: AVX2
- AllowFallback: True
- SkipCheck: False
- SearchDirectories and Priorities: { ./ }
[llama Info]: NativeLibraryConfig Description:
- LibraryName: LLama
- Path: ''
- PreferCuda: True
- PreferVulkan: True
- PreferredAvxLevel: AVX2
- AllowFallback: True
- SkipCheck: False
- SearchDirectories and Priorities: { ./ }
[llama Debug]: Got relative library path 'runtimes/win-x64/native/cuda12/llama.dll' from local with (NativeLibraryName: LLama, UseCuda: True, UseVulkan: False, AvxLevel: Avx2), trying to load it...
[llama Debug]: Found full path file './runtimes\win-x64\native\cuda12\ggml-base.dll' for relative path 'runtimes\win-x64\native\cuda12\ggml-base.dll'
[llama Info]: Successfully loaded './runtimes\win-x64\native\cuda12\ggml-base.dll'
[llama Info]: Successfully loaded dependency 'runtimes\win-x64\native\cuda12\ggml-base.dll'
[llama Debug]: Found full path file './runtimes/win-x64/native/avx2\ggml-cpu.dll' for relative path 'runtimes/win-x64/native/avx2\ggml-cpu.dll'
[llama Info]: Successfully loaded './runtimes/win-x64/native/avx2\ggml-cpu.dll'
[llama Info]: Successfully loaded dependency 'runtimes/win-x64/native/avx2\ggml-cpu.dll'
[llama Debug]: Found full path file './runtimes\win-x64\native\cuda12\ggml-cuda.dll' for relative path 'runtimes\win-x64\native\cuda12\ggml-cuda.dll'
[llama Info]: Successfully loaded './runtimes\win-x64\native\cuda12\ggml-cuda.dll'
[llama Info]: Successfully loaded dependency 'runtimes\win-x64\native\cuda12\ggml-cuda.dll'
[llama Debug]: Found full path file './runtimes\win-x64\native\cuda12\ggml.dll' for relative path 'runtimes\win-x64\native\cuda12\ggml.dll'
[llama Info]: Successfully loaded './runtimes\win-x64\native\cuda12\ggml.dll'
[llama Info]: Successfully loaded dependency 'runtimes\win-x64\native\cuda12\ggml.dll'
[llama Debug]: Found full path file './runtimes/win-x64/native/cuda12/llama.dll' for relative path 'runtimes/win-x64/native/cuda12/llama.dll'
[llama Info]: Successfully loaded './runtimes/win-x64/native/cuda12/llama.dll'
[llama Info]: ggml_cuda_init: GGML_CUDA_FORCE_MMQ:    no
[llama Info]: ggml_cuda_init: GGML_CUDA_FORCE_CUBLAS: no
[llama Info]: ggml_cuda_init: found 1 CUDA devices:
[llama Info]:   Device 0: NVIDIA GeForce RTX 4080 SUPER, compute capability 8.9, VMM: yes
[llama Info]: llama_model_load_from_file_impl: using device CUDA0 (NVIDIA GeForce RTX 4080 SUPER) - 15035 MiB free
[llama Info]: llama_model_loader: loaded meta data with 39 key-value pairs and 434 tensors from b:\models\Qwen2.5-Coder-3B-Instruct-Q8_0.gguf (version GGUF V3 (latest))
[llama Info]: llama_model_loader: Dumping metadata keys/values. Note: KV overrides do not apply in this output.
[llama Info]: llama_model_loader: - kv   0:                       general.architecture str              = qwen2
[llama Info]: llama_model_loader: - kv   1:                               general.type str              = model
[llama Info]: llama_model_loader: - kv   2:                               general.name str              = Qwen2.5 Coder 3B Instruct      
[llama Info]: llama_model_loader: - kv   3:                           general.finetune str              = Instruct
[llama Info]: llama_model_loader: - kv   4:                           general.basename str              = Qwen2.5-Coder
[llama Info]: llama_model_loader: - kv   5:                         general.size_label str              = 3B
[llama Info]: llama_model_loader: - kv   6:                            general.license str              = other
[llama Info]: llama_model_loader: - kv   7:                       general.license.name str              = qwen-research
[llama Info]: llama_model_loader: - kv   8:                       general.license.link str              = https://huggingface.co/Qwen/Qwen2.5-3...
[llama Info]: llama_model_loader: - kv   9:                   general.base_model.count u32              = 1
[llama Info]: llama_model_loader: - kv  10:                  general.base_model.0.name str              = Qwen2.5 Coder 3B
[llama Info]: llama_model_loader: - kv  11:          general.base_model.0.organization str              = Qwen
[llama Info]: llama_model_loader: - kv  12:              general.base_model.0.repo_url str              = https://huggingface.co/Qwen/Qwen2.5-C...
[llama Info]: llama_model_loader: - kv  13:                               general.tags arr[str,6]       = ["code", "codeqwen", "chat", "qwen", ...
[llama Info]: llama_model_loader: - kv  14:                          general.languages arr[str,1]       = ["en"]
[llama Info]: llama_model_loader: - kv  15:                          qwen2.block_count u32              = 36
[llama Info]: llama_model_loader: - kv  16:                       qwen2.context_length u32              = 32768
[llama Info]: llama_model_loader: - kv  17:                     qwen2.embedding_length u32              = 2048
[llama Info]: llama_model_loader: - kv  18:                  qwen2.feed_forward_length u32              = 11008
[llama Info]: llama_model_loader: - kv  19:                 qwen2.attention.head_count u32              = 16
[llama Info]: llama_model_loader: - kv  20:              qwen2.attention.head_count_kv u32              = 2
[llama Info]: llama_model_loader: - kv  21:                       qwen2.rope.freq_base f32              = 1000000.000000
[llama Info]: llama_model_loader: - kv  22:     qwen2.attention.layer_norm_rms_epsilon f32              = 0.000001
[llama Info]: llama_model_loader: - kv  23:                          general.file_type u32              = 7
[llama Info]: llama_model_loader: - kv  24:                       tokenizer.ggml.model str              = gpt2
[llama Info]: llama_model_loader: - kv  25:                         tokenizer.ggml.pre str              = qwen2
[llama Info]: llama_model_loader: - kv  26:                      tokenizer.ggml.tokens arr[str,151936]  = ["!", "\"", "#", "$", "%", "&", "'", ...
[llama Info]: llama_model_loader: - kv  27:                  tokenizer.ggml.token_type arr[i32,151936]  = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...
[llama Info]: llama_model_loader: - kv  28:                      tokenizer.ggml.merges arr[str,151387]  = ["Ä  Ä ", "Ä Ä  Ä Ä ", "i n", "Ä  t",...
[llama Info]: llama_model_loader: - kv  29:                tokenizer.ggml.eos_token_id u32              = 151645
[llama Info]: llama_model_loader: - kv  30:            tokenizer.ggml.padding_token_id u32              = 151643
[llama Info]: llama_model_loader: - kv  31:                tokenizer.ggml.bos_token_id u32              = 151643
[llama Info]: llama_model_loader: - kv  32:               tokenizer.ggml.add_bos_token bool             = false
[llama Info]: llama_model_loader: - kv  33:                    tokenizer.chat_template str              = {%- if tools %}\n    {{- '<|im_start|>...
[llama Info]: llama_model_loader: - kv  34:               general.quantization_version u32              = 2
[llama Info]: llama_model_loader: - kv  35:                      quantize.imatrix.file str              = /models_out/Qwen2.5-Coder-3B-Instruct...
[llama Info]: llama_model_loader: - kv  36:                   quantize.imatrix.dataset str              = /training_dir/calibration_datav3.txt
[llama Info]: llama_model_loader: - kv  37:             quantize.imatrix.entries_count i32              = 252
[llama Info]: llama_model_loader: - kv  38:              quantize.imatrix.chunks_count i32              = 128
[llama Info]: llama_model_loader: - type  f32:  181 tensors
[llama Info]: llama_model_loader: - type q8_0:  253 tensors
[llama Info]: print_info: file format = GGUF V3 (latest)
[llama Info]: print_info: file type   = Q8_0
[llama Info]: print_info: file size   = 3.05 GiB (8.50 BPW)
[llama Debug]: init_tokenizer: initializing tokenizer for type 2
[llama Debug]: load: control token: 151661 '<|fim_suffix|>' is not marked as EOG
[llama Debug]: load: control token: 151649 '<|box_end|>' is not marked as EOG
[llama Debug]: load: control token: 151647 '<|object_ref_end|>' is not marked as EOG
[llama Debug]: load: control token: 151654 '<|vision_pad|>' is not marked as EOG
[llama Debug]: load: control token: 151659 '<|fim_prefix|>' is not marked as EOG
[llama Debug]: load: control token: 151648 '<|box_start|>' is not marked as EOG
[llama Debug]: load: control token: 151644 '<|im_start|>' is not marked as EOG
[llama Debug]: load: control token: 151646 '<|object_ref_start|>' is not marked as EOG
[llama Debug]: load: control token: 151650 '<|quad_start|>' is not marked as EOG
[llama Debug]: load: control token: 151651 '<|quad_end|>' is not marked as EOG
[llama Debug]: load: control token: 151652 '<|vision_start|>' is not marked as EOG
[llama Debug]: load: control token: 151653 '<|vision_end|>' is not marked as EOG
[llama Debug]: load: control token: 151655 '<|image_pad|>' is not marked as EOG
[llama Debug]: load: control token: 151656 '<|video_pad|>' is not marked as EOG
[llama Debug]: load: control token: 151660 '<|fim_middle|>' is not marked as EOG
[llama Info]: load: special tokens cache size = 22
[llama Info]: load: token to piece cache size = 0.9310 MB
[llama Info]: print_info: arch             = qwen2
[llama Info]: print_info: vocab_only       = 0
[llama Info]: print_info: n_ctx_train      = 32768
[llama Info]: print_info: n_embd           = 2048
[llama Info]: print_info: n_layer          = 36
[llama Info]: print_info: n_head           = 16
[llama Info]: print_info: n_head_kv        = 2
[llama Info]: print_info: n_rot            = 128
[llama Info]: print_info: n_swa            = 0
[llama Info]: print_info: n_embd_head_k    = 128
[llama Info]: print_info: n_embd_head_v    = 128
[llama Info]: print_info: n_gqa            = 8
[llama Info]: print_info: n_embd_k_gqa     = 256
[llama Info]: print_info: n_embd_v_gqa     = 256
[llama Info]: print_info: f_norm_eps       = 0.0e+00
[llama Info]: print_info: f_norm_rms_eps   = 1.0e-06
[llama Info]: print_info: f_clamp_kqv      = 0.0e+00
[llama Info]: print_info: f_max_alibi_bias = 0.0e+00
[llama Info]: print_info: f_logit_scale    = 0.0e+00
[llama Info]: print_info: n_ff             = 11008
[llama Info]: print_info: n_expert         = 0
[llama Info]: print_info: n_expert_used    = 0
[llama Info]: print_info: causal attn      = 1
[llama Info]: print_info: pooling type     = 0
[llama Info]: print_info: rope type        = 2
[llama Info]: print_info: rope scaling     = linear
[llama Info]: print_info: freq_base_train  = 1000000.0
[llama Info]: print_info: freq_scale_train = 1
[llama Info]: print_info: n_ctx_orig_yarn  = 32768
[llama Info]: print_info: rope_finetuned   = unknown
[llama Info]: print_info: ssm_d_conv       = 0
[llama Info]: print_info: ssm_d_inner      = 0
[llama Info]: print_info: ssm_d_state      = 0
[llama Info]: print_info: ssm_dt_rank      = 0
[llama Info]: print_info: ssm_dt_b_c_rms   = 0
[llama Info]: print_info: model type       = 3B
[llama Info]: print_info: model params     = 3.09 B
[llama Info]: print_info: general.name     = Qwen2.5 Coder 3B Instruct
[llama Info]: print_info: vocab type       = BPE
[llama Info]: print_info: n_vocab          = 151936
[llama Info]: print_info: n_merges         = 151387
[llama Info]: print_info: BOS token        = 151643 '<|endoftext|>'
[llama Info]: print_info: EOS token        = 151645 '<|im_end|>'
[llama Info]: print_info: EOT token        = 151645 '<|im_end|>'
[llama Info]: print_info: PAD token        = 151643 '<|endoftext|>'
[llama Info]: print_info: LF token         = 198 'ÄS'
[llama Info]: print_info: FIM PRE token    = 151659 '<|fim_prefix|>'
[llama Info]: print_info: FIM SUF token    = 151661 '<|fim_suffix|>'
[llama Info]: print_info: FIM MID token    = 151660 '<|fim_middle|>'
[llama Info]: print_info: FIM PAD token    = 151662 '<|fim_pad|>'
[llama Info]: print_info: FIM REP token    = 151663 '<|repo_name|>'
[llama Info]: print_info: FIM SEP token    = 151664 '<|file_sep|>'
[llama Info]: print_info: EOG token        = 151643 '<|endoftext|>'
[llama Info]: print_info: EOG token        = 151645 '<|im_end|>'
[llama Info]: print_info: EOG token        = 151662 '<|fim_pad|>'
[llama Info]: print_info: EOG token        = 151663 '<|repo_name|>'
[llama Info]: print_info: EOG token        = 151664 '<|file_sep|>'
[llama Info]: print_info: max token length = 256
[llama Debug]: load_tensors: layer   0 assigned to device CUDA0
[llama Debug]: load_tensors: layer   1 assigned to device CUDA0
[llama Debug]: load_tensors: layer   2 assigned to device CUDA0
[llama Debug]: load_tensors: layer   3 assigned to device CUDA0
[llama Debug]: load_tensors: layer   4 assigned to device CUDA0
[llama Debug]: load_tensors: layer   5 assigned to device CUDA0
[llama Debug]: load_tensors: layer   6 assigned to device CUDA0
[llama Debug]: load_tensors: layer   7 assigned to device CUDA0
[llama Debug]: load_tensors: layer   8 assigned to device CUDA0
[llama Debug]: load_tensors: layer   9 assigned to device CUDA0
[llama Debug]: load_tensors: layer  10 assigned to device CUDA0
[llama Debug]: load_tensors: layer  11 assigned to device CUDA0
[llama Debug]: load_tensors: layer  12 assigned to device CUDA0
[llama Debug]: load_tensors: layer  13 assigned to device CUDA0
[llama Debug]: load_tensors: layer  14 assigned to device CUDA0
[llama Debug]: load_tensors: layer  15 assigned to device CUDA0
[llama Debug]: load_tensors: layer  16 assigned to device CUDA0
[llama Debug]: load_tensors: layer  17 assigned to device CUDA0
[llama Debug]: load_tensors: layer  18 assigned to device CUDA0
[llama Debug]: load_tensors: layer  19 assigned to device CUDA0
[llama Debug]: load_tensors: layer  20 assigned to device CUDA0
[llama Debug]: load_tensors: layer  21 assigned to device CUDA0
[llama Debug]: load_tensors: layer  22 assigned to device CUDA0
[llama Debug]: load_tensors: layer  23 assigned to device CUDA0
[llama Debug]: load_tensors: layer  24 assigned to device CUDA0
[llama Debug]: load_tensors: layer  25 assigned to device CUDA0
[llama Debug]: load_tensors: layer  26 assigned to device CUDA0
[llama Debug]: load_tensors: layer  27 assigned to device CUDA0
[llama Debug]: load_tensors: layer  28 assigned to device CUDA0
[llama Debug]: load_tensors: layer  29 assigned to device CUDA0
[llama Debug]: load_tensors: layer  30 assigned to device CUDA0
[llama Debug]: load_tensors: layer  31 assigned to device CUDA0
[llama Debug]: load_tensors: layer  32 assigned to device CUDA0
[llama Debug]: load_tensors: layer  33 assigned to device CUDA0
[llama Debug]: load_tensors: layer  34 assigned to device CUDA0
[llama Debug]: load_tensors: layer  35 assigned to device CUDA0
[llama Debug]: load_tensors: layer  36 assigned to device CUDA0
[llama Debug]: load_tensors: tensor 'token_embd.weight' (q8_0) (and 0 others) cannot be used with preferred buffer type CPU_AARCH64, using CPU instead
[llama Info]: load_tensors: offloading 36 repeating layers to GPU
[llama Info]: load_tensors: offloading output layer to GPU
[llama Info]: load_tensors: offloaded 37/37 layers to GPU
[llama Info]: load_tensors:        CUDA0 model buffer size =  3127.61 MiB
[llama Info]: load_tensors:   CPU_Mapped model buffer size =   315.30 MiB
[llama Info]: llama_init_from_model: n_seq_max     = 1
[llama Info]: llama_init_from_model: n_ctx         = 4096
[llama Info]: llama_init_from_model: n_ctx_per_seq = 4096
[llama Info]: llama_init_from_model: n_batch       = 512
[llama Info]: llama_init_from_model: n_ubatch      = 512
[llama Info]: llama_init_from_model: flash_attn    = 0
[llama Info]: llama_init_from_model: freq_base     = 1000000.0
[llama Info]: llama_init_from_model: freq_scale    = 1
[llama Warning]: llama_init_from_model: n_ctx_per_seq (4096) < n_ctx_train (32768) -- the full capacity of the model will not be utilized
[llama Info]: llama_kv_cache_init: kv_size = 4096, offload = 1, type_k = 'f16', type_v = 'f16', n_layer = 36, can_shift = 1
[llama Debug]: llama_kv_cache_init: layer 0: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 1: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 2: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 3: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 4: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 5: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 6: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 7: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 8: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 9: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 10: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 11: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 12: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 13: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 14: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 15: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 16: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 17: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 18: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 19: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 20: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 21: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 22: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 23: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 24: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 25: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 26: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 27: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 28: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 29: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 30: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 31: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 32: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 33: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 34: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 35: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Info]: llama_kv_cache_init:      CUDA0 KV buffer size =   144.00 MiB
[llama Info]: llama_init_from_model: KV self size  =  144.00 MiB, K (f16):   72.00 MiB, V (f16):   72.00 MiB
[llama Info]: llama_init_from_model:  CUDA_Host  output buffer size =     0.58 MiB
[llama Info]: llama_init_from_model:      CUDA0 compute buffer size =   300.75 MiB
[llama Info]: llama_init_from_model:  CUDA_Host compute buffer size =    12.01 MiB
[llama Info]: llama_init_from_model: graph nodes  = 1266
[llama Info]: llama_init_from_model: graph splits = 2
[llama Info]: llama_init_from_model: n_seq_max     = 1
[llama Info]: llama_init_from_model: n_ctx         = 4096
[llama Info]: llama_init_from_model: n_ctx_per_seq = 4096
[llama Info]: llama_init_from_model: n_batch       = 512
[llama Info]: llama_init_from_model: n_ubatch      = 512
[llama Info]: llama_init_from_model: flash_attn    = 0
[llama Info]: llama_init_from_model: freq_base     = 1000000.0
[llama Info]: llama_init_from_model: freq_scale    = 1
[llama Warning]: llama_init_from_model: n_ctx_per_seq (4096) < n_ctx_train (32768) -- the full capacity of the model will not be utilized
[llama Info]: llama_kv_cache_init: kv_size = 4096, offload = 1, type_k = 'f16', type_v = 'f16', n_layer = 36, can_shift = 1
[llama Debug]: llama_kv_cache_init: layer 0: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 1: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 2: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 3: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 4: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 5: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 6: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 7: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 8: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 9: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 10: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 11: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 12: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 13: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 14: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 15: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 16: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 17: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 18: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 19: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 20: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 21: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 22: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 23: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 24: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 25: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 26: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 27: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 28: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 29: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 30: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 31: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 32: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 33: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 34: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Debug]: llama_kv_cache_init: layer 35: n_embd_k_gqa = 256, n_embd_v_gqa = 256
[llama Info]: llama_kv_cache_init:      CUDA0 KV buffer size =   144.00 MiB
[llama Info]: llama_init_from_model: KV self size  =  144.00 MiB, K (f16):   72.00 MiB, V (f16):   72.00 MiB
[llama Info]: llama_init_from_model:  CUDA_Host  output buffer size =     0.58 MiB
[llama Info]: llama_init_from_model:      CUDA0 compute buffer size =   300.75 MiB
[llama Info]: llama_init_from_model:  CUDA_Host compute buffer size =    12.01 MiB
[llama Info]: llama_init_from_model: graph nodes  = 1266
[llama Info]: llama_init_from_model: graph splits = 2

and log from my llama-cli run

ggml_cuda_init: GGML_CUDA_FORCE_MMQ:    no
ggml_cuda_init: GGML_CUDA_FORCE_CUBLAS: no
ggml_cuda_init: found 1 CUDA devices:
  Device 0: NVIDIA GeForce RTX 4080 SUPER, compute capability 8.9, VMM: yes
build: 4701 (be3bbd62) with MSVC 19.41.34120.0 for x64
main: llama backend init
main: load the model and apply lora adapter, if any
llama_model_load_from_file_impl: using device CUDA0 (NVIDIA GeForce RTX 4080 SUPER) - 15035 MiB free
llama_model_loader: loaded meta data with 39 key-value pairs and 434 tensors from B:\models\Qwen2.5-Coder-3B-Instruct-Q8_0.gguf (version GGUF V3 (latest))
llama_model_loader: Dumping metadata keys/values. Note: KV overrides do not apply in this output.
llama_model_loader: - kv   0:                       general.architecture str              = qwen2
llama_model_loader: - kv   1:                               general.type str              = model
llama_model_loader: - kv   2:                               general.name str              = Qwen2.5 Coder 3B Instruct
llama_model_loader: - kv   3:                           general.finetune str              = Instruct
llama_model_loader: - kv   4:                           general.basename str              = Qwen2.5-Coder
llama_model_loader: - kv   5:                         general.size_label str              = 3B
llama_model_loader: - kv   6:                            general.license str              = other
llama_model_loader: - kv   7:                       general.license.name str              = qwen-research
llama_model_loader: - kv   8:                       general.license.link str              = https://huggingface.co/Qwen/Qwen2.5-3...
llama_model_loader: - kv   9:                   general.base_model.count u32              = 1
llama_model_loader: - kv  10:                  general.base_model.0.name str              = Qwen2.5 Coder 3B
llama_model_loader: - kv  11:          general.base_model.0.organization str              = Qwen
llama_model_loader: - kv  12:              general.base_model.0.repo_url str              = https://huggingface.co/Qwen/Qwen2.5-C...
llama_model_loader: - kv  13:                               general.tags arr[str,6]       = ["code", "codeqwen", "chat", "qwen", ...
llama_model_loader: - kv  14:                          general.languages arr[str,1]       = ["en"]
llama_model_loader: - kv  15:                          qwen2.block_count u32              = 36
llama_model_loader: - kv  16:                       qwen2.context_length u32              = 32768
llama_model_loader: - kv  17:                     qwen2.embedding_length u32              = 2048
llama_model_loader: - kv  18:                  qwen2.feed_forward_length u32              = 11008
llama_model_loader: - kv  19:                 qwen2.attention.head_count u32              = 16
llama_model_loader: - kv  20:              qwen2.attention.head_count_kv u32              = 2
llama_model_loader: - kv  21:                       qwen2.rope.freq_base f32              = 1000000.000000
llama_model_loader: - kv  22:     qwen2.attention.layer_norm_rms_epsilon f32              = 0.000001
llama_model_loader: - kv  23:                          general.file_type u32              = 7
llama_model_loader: - kv  24:                       tokenizer.ggml.model str              = gpt2
llama_model_loader: - kv  25:                         tokenizer.ggml.pre str              = qwen2
llama_model_loader: - kv  26:                      tokenizer.ggml.tokens arr[str,151936]  = ["!", "\"", "#", "$", "%", "&", "'", ...
llama_model_loader: - kv  27:                  tokenizer.ggml.token_type arr[i32,151936]  = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...
llama_model_loader: - kv  28:                      tokenizer.ggml.merges arr[str,151387]  = ["Ġ Ġ", "ĠĠ ĠĠ", "i n", "Ġ t",...
llama_model_loader: - kv  29:                tokenizer.ggml.eos_token_id u32              = 151645
llama_model_loader: - kv  30:            tokenizer.ggml.padding_token_id u32              = 151643
llama_model_loader: - kv  31:                tokenizer.ggml.bos_token_id u32              = 151643
llama_model_loader: - kv  32:               tokenizer.ggml.add_bos_token bool             = false
llama_model_loader: - kv  33:                    tokenizer.chat_template str              = {%- if tools %}\n    {{- '<|im_start|>...
llama_model_loader: - kv  34:               general.quantization_version u32              = 2
llama_model_loader: - kv  35:                      quantize.imatrix.file str              = /models_out/Qwen2.5-Coder-3B-Instruct...
llama_model_loader: - kv  36:                   quantize.imatrix.dataset str              = /training_dir/calibration_datav3.txt
llama_model_loader: - kv  37:             quantize.imatrix.entries_count i32              = 252
llama_model_loader: - kv  38:              quantize.imatrix.chunks_count i32              = 128
llama_model_loader: - type  f32:  181 tensors
llama_model_loader: - type q8_0:  253 tensors
print_info: file format = GGUF V3 (latest)
print_info: file type   = Q8_0
print_info: file size   = 3.05 GiB (8.50 BPW)
load: special tokens cache size = 22
load: token to piece cache size = 0.9310 MB
print_info: arch             = qwen2
print_info: vocab_only       = 0
print_info: n_ctx_train      = 32768
print_info: n_embd           = 2048
print_info: n_layer          = 36
print_info: n_head           = 16
print_info: n_head_kv        = 2
print_info: n_rot            = 128
print_info: n_swa            = 0
print_info: n_embd_head_k    = 128
print_info: n_embd_head_v    = 128
print_info: n_gqa            = 8
print_info: n_embd_k_gqa     = 256
print_info: n_embd_v_gqa     = 256
print_info: f_norm_eps       = 0.0e+00
print_info: f_norm_rms_eps   = 1.0e-06
print_info: f_clamp_kqv      = 0.0e+00
print_info: f_max_alibi_bias = 0.0e+00
print_info: f_logit_scale    = 0.0e+00
print_info: n_ff             = 11008
print_info: n_expert         = 0
print_info: n_expert_used    = 0
print_info: causal attn      = 1
print_info: pooling type     = 0
print_info: rope type        = 2
print_info: rope scaling     = linear
print_info: freq_base_train  = 1000000.0
print_info: freq_scale_train = 1
print_info: n_ctx_orig_yarn  = 32768
print_info: rope_finetuned   = unknown
print_info: ssm_d_conv       = 0
print_info: ssm_d_inner      = 0
print_info: ssm_d_state      = 0
print_info: ssm_dt_rank      = 0
print_info: ssm_dt_b_c_rms   = 0
print_info: model type       = 3B
print_info: model params     = 3.09 B
print_info: general.name     = Qwen2.5 Coder 3B Instruct
print_info: vocab type       = BPE
print_info: n_vocab          = 151936
print_info: n_merges         = 151387
print_info: BOS token        = 151643 '<|endoftext|>'
print_info: EOS token        = 151645 '<|im_end|>'
print_info: EOT token        = 151645 '<|im_end|>'
print_info: PAD token        = 151643 '<|endoftext|>'
print_info: LF token         = 198 'Ċ'
print_info: FIM PRE token    = 151659 '<|fim_prefix|>'
print_info: FIM SUF token    = 151661 '<|fim_suffix|>'
print_info: FIM MID token    = 151660 '<|fim_middle|>'
print_info: FIM PAD token    = 151662 '<|fim_pad|>'
print_info: FIM REP token    = 151663 '<|repo_name|>'
print_info: FIM SEP token    = 151664 '<|file_sep|>'
print_info: EOG token        = 151643 '<|endoftext|>'
print_info: EOG token        = 151645 '<|im_end|>'
print_info: EOG token        = 151662 '<|fim_pad|>'
print_info: EOG token        = 151663 '<|repo_name|>'
print_info: EOG token        = 151664 '<|file_sep|>'
print_info: max token length = 256
load_tensors: loading model tensors, this can take a while... (mmap = true)
load_tensors: offloading 36 repeating layers to GPU
load_tensors: offloading output layer to GPU
load_tensors: offloaded 37/37 layers to GPU
load_tensors:        CUDA0 model buffer size =  3127.61 MiB
load_tensors:   CPU_Mapped model buffer size =   315.30 MiB
...................................................................................
llama_init_from_model: n_seq_max     = 1
llama_init_from_model: n_ctx         = 4096
llama_init_from_model: n_ctx_per_seq = 4096
llama_init_from_model: n_batch       = 2048
llama_init_from_model: n_ubatch      = 512
llama_init_from_model: flash_attn    = 0
llama_init_from_model: freq_base     = 1000000.0
llama_init_from_model: freq_scale    = 1
llama_init_from_model: n_ctx_per_seq (4096) < n_ctx_train (32768) -- the full capacity of the model will not be utilized
llama_kv_cache_init: kv_size = 4096, offload = 1, type_k = 'f16', type_v = 'f16', n_layer = 36, can_shift = 1
llama_kv_cache_init:      CUDA0 KV buffer size =   144.00 MiB
llama_init_from_model: KV self size  =  144.00 MiB, K (f16):   72.00 MiB, V (f16):   72.00 MiB
llama_init_from_model:  CUDA_Host  output buffer size =     0.58 MiB
llama_init_from_model:      CUDA0 compute buffer size =   300.75 MiB
llama_init_from_model:  CUDA_Host compute buffer size =    12.01 MiB
llama_init_from_model: graph nodes  = 1266
llama_init_from_model: graph splits = 2
common_init_from_params: setting dry_penalty_last_n to ctx_size = 4096
main: llama threadpool init, n_threads = 12

system_info: n_threads = 12 (n_threads_batch = 12) / 20 | CUDA : ARCHS = 890 | USE_GRAPHS = 1 | PEER_MAX_BATCH_SIZE = 128 | CPU : SSE3 = 1 | SSSE3 = 1 | AVX = 1 | AVX2 = 1 | F16C = 1 | FMA = 1 | LLAMAFILE = 1 | OPENMP = 1 | AARCH64_REPACK = 1 |

sampler seed: 2895019185
sampler params:
        repeat_last_n = 64, repeat_penalty = 1.000, frequency_penalty = 0.000, presence_penalty = 0.000
        dry_multiplier = 0.000, dry_base = 1.750, dry_allowed_length = 2, dry_penalty_last_n = 4096
        top_k = 40, top_p = 0.950, min_p = 0.050, xtc_probability = 0.000, xtc_threshold = 0.100, typical_p = 1.000, temp = 0.800
        mirostat = 0, mirostat_lr = 0.100, mirostat_ent = 5.000
sampler chain: logits -> logit-bias -> penalties -> dry -> top-k -> typical -> top-p -> min-p -> xtc -> temp-ext -> dist
generate: n_ctx = 4096, n_batch = 2048, n_predict = -1, n_keep = 0

[phil-scott-78]

I had some time to think about this. I was wondering what was up with llama.cpp not having a grammar in the displayed chain. First I've really looked at this code, but at first read of sampling.cpp it seems the grammar isn't part of the chain but rather applied when needed. Wonder if that's enough to get 4x the perf though

[m0nsky]

Interesting, I make heavy use of grammars and have noticed some bad performance, but wasn't aware the performance was better in llama.cpp cli. I don't know enough about the samplers / sampling pipelines to experiment with this, maybe @martindevans will be able to spot something?

[phil-scott-78]

Thinking back on the way sampling is handled as needed in sampling cpp the speed difference I'm seeing is starting to make sense. I'm giving the LLM a good example to use as its template so it is generally very close on its return, generally just needing to avoid markdown around the response. So in my scenario the grammar might only run for a handful of tokens in llama.cpp but every time in llamasharp

[phil-scott-78]

I'm headed off on vacation with the kids, be back online next Sunday. But I'll be willing to try and give it a go creating a new sampler that tries to mimic the current llama.cpp common sampler behavior if no one else gets after it before then.

[martindevans]

Sampling.h seems to explain what's going on here: https://github.com/mmoskal/llama.cpp/blob/5598f475be3e31430fbe17ebb85654ec90dc201e/common/sampling.h#L51-L61:

// extended sampling implementation:
//
// - set logits
// - apply the configured sampler chain
// - check if the token fits the grammar (if any)
// - if not: resample by first applying the grammar constraints and then sampling again (slower path)
//
// if grammar_first is true, the grammar is applied before the samplers (slower)
// useful in cases where all the resulting candidates (not just the sampled one) must fit the grammar

In LLamaSharp, because the grammar is part of the pipeline it is just applied whenever it should be according to the pipeline. In DefaultSamplingPipeline that's first, so it's equivalent to always going the grammar_first = true path which is documented to be slower.

Retrofitting this behaviour into the current sampling abstraction would be possible - I think you'd need to move the Grammar property into BaseSamplingPipeline.cs and do this resampling logic in Sample.

Alternatively, there are lazy grammars which were introduced recently. Maybe they work better for this use case, and we should add better support for them? (e.g. add optional trigger words alongside the Grammar property, and use a lazy grammar stage if any are supplied).

[phil-scott-78]

Edit: sorry for the weird reply I said, on vacation lol I think a rework taking the lazy grammar into consideration would be great, especially with the deepthink model now being popular. I was throwing the think tag into the gbnf which worked well but it was also a catalyst for this post because woah buddy was it slow Lazy grammar with the perf improvement would be an awesome improvement. But one I'll be happy to test, but a bit out of my depth to write

[m0nsky]

Sampling.h seems to explain what's going on here: https://github.com/mmoskal/llama.cpp/blob/5598f475be3e31430fbe17ebb85654ec90dc201e/common/sampling.h#L51-L61:

// extended sampling implementation:
//
// - set logits
// - apply the configured sampler chain
// - check if the token fits the grammar (if any)
// - if not: resample by first applying the grammar constraints and then sampling again (slower path)
//
// if grammar_first is true, the grammar is applied before the samplers (slower)
// useful in cases where all the resulting candidates (not just the sampled one) must fit the grammar

In LLamaSharp, because the grammar is part of the pipeline it is just applied whenever it should be according to the pipeline. In DefaultSamplingPipeline that's first, so it's equivalent to always going the grammar_first = true path which is documented to be slower.

Retrofitting this behaviour into the current sampling abstraction would be possible - I think you'd need to move the Grammar property into BaseSamplingPipeline.cs and do this resampling logic in Sample.

Alternatively, there are lazy grammars which were introduced recently. Maybe they work better for this use case, and we should add better support for them? (e.g. add optional trigger words alongside the Grammar property, and use a lazy grammar stage if any are supplied).

Does the order of chain.AddTopK and chain.AddGrammar etc matter in CreateChain? If it does, then wouldn't it be easier to add the grammar_first functionality in CreateChain? (unless you suggested Sample to make this available no matter which sampling pipeline).

Just making sure I understand it correctly.

Edit
How would we "check if the token fits the grammar"? The rest seems pretty straight forward.

[martindevans]

unless you suggested Sample to make this available no matter which sampling pipeline

That is why I suggested putting it in BaseSamplingPipeline. For now though, maybe it would be safer to add it into DefaultSamplingPipeline (overriding Sample) and we can look at moving it to base later.

wouldn't it be easier to add the grammar_first functionality in CreateChain

I don't think it can be implemented here. CreateChain lets you define that chain that is used in Sample, but this resampling stuff needs the resampling logic which is not part of the chain.

How would we "check if the token fits the grammar"? The rest seems pretty straight forward.

Just check if there are any valid logits after applying the grammar. Just !float.IsNegativeInfinity(TensorPrimitives.Max(logits)) would probably do the trick.

[phil-scott-78]

I had a few cycles this evening after the kids went to bed and threw this together. Ok, speed up. But still feels like I'm lagging behind what llama.cpp is showing when running llama-cli. Explicitly calling Accept here seemed wrong compared to their implementation, but that might be due to how the executors work in tandem that required it. Honestly not sure.

public class GrammarSamplingPipeline : ISamplingPipeline
{
    private readonly Grammar _grammar;
    private LLamaTokenDataArray _cur;
    private LLamaTokenDataArrayNative _curP;
    
    private SafeLLamaSamplerChainHandle? _chain;
    private SafeLLamaSamplerChainHandle? _grammarChain;

    public GrammarSamplingPipeline(Grammar grammar)
    {
        _grammar = grammar;
    }

    private void SetLogits(SafeLLamaContextHandle ctx, int idx)
    {
        var logits = ctx.GetLogitsIth(idx);
        _cur = LLamaTokenDataArray.Create(logits);
        using var f = LLamaTokenDataArrayNative.Create(_cur, out _curP);
    }
    
    /// <inheritdoc />
    public LLamaToken Sample(SafeLLamaContextHandle ctx, int index)
    { 
        _chain ??= CreateChain();
        _grammarChain ??= CreateGrammarChain(ctx);
        
        SetLogits(ctx, index);
        
        // apply sampler only to the chain
        _chain.Apply(ref _curP);
        var id = _curP.Data[(int) _curP.Selected].ID;

        // check if the sampled token fits the grammar 
        var singleTokenData =  new LLamaTokenData(id, 1.0f, 0.0f);

        var tokenDataArray =  LLamaTokenDataArray.Create(new [] { singleTokenData.Logit }.AsSpan());
        using var _ = LLamaTokenDataArrayNative.Create(tokenDataArray, out var singleTokenDataArray );


        _grammarChain.Apply(ref singleTokenDataArray);
        var isValid = !float.IsNegativeInfinity(singleTokenDataArray.Data[0].Logit);

        if (isValid)
        {
            _grammarChain.Accept(id);
            _chain.Accept(id);
            return id;
        }

        SetLogits(ctx, index);
        _grammarChain.Apply(ref _curP);
        _chain.Apply(ref _curP);


        id = _curP.Data[(int)_curP.Selected].ID;
        _grammarChain.Accept(id);
        _chain.Accept(id);
        return id;
    }

    private SafeLLamaSamplerChainHandle CreateChain()
    {
        var chain = SafeLLamaSamplerChainHandle.Create(LLamaSamplerChainParams.Default());
        chain.AddGreedySampler();
        return chain;
    }
    
    private SafeLLamaSamplerChainHandle CreateGrammarChain(SafeLLamaContextHandle ctx)
    {
        
        var chain = SafeLLamaSamplerChainHandle.Create(LLamaSamplerChainParams.Default());
        chain.AddGrammar(ctx.ModelHandle, _grammar.Gbnf, _grammar.Root);
        return chain;
    }

    /// <inheritdoc />
    public void Reset()
    {
        _chain?.Reset();
        _grammarChain?.Reset();
    }

    /// <inheritdoc />
    public void Accept(LLamaToken token)
    {
        _grammarChain?.Accept(token);
        _chain?.Accept(token);
    }
    
        
    /// <inheritdoc />
    public void Dispose()
    {
        _chain?.Dispose();
        _chain = null;

        GC.SuppressFinalize(this);
    }
}

New GrammarSampler

Greedy

[martindevans]

@m0nsky and I have been having a look at this over the last couple of days. I've put together a prototype in this branch which m0nsky has done some more work on (not yet pushed). This is extending the DefaultSamplingPipeline to automaticaly do resampling if Grammar is not null.

private void SetLogits

I don't think your implementation of SetLogits is safe - it's using f which means it disposes that when SetLogits returns, but then it's using _curP for longer than that lifetime!

Explicitly calling Accept here seemed wrong

That caught me out too in my prototype, but I'm pretty sure it's right. The normal Sample method is basically a helper than automatically calls Apply and Accept for you (see comment here). So if you're replacing that it makes sense to be calling Apply and Accept.