Cool!

Which part of the sampling exactly is so slow?

Is it correct to say that each step runs pp2048 + sampling?

Yes that's correct.

Why is it slow (acc to me) - there is a logits buffer of size (n_ctx x n_vocab) at each step where we need to sample from, I think transferring this buffer from GPU->CPU at each step plus sampling is killing performance, the pytorch HF version where the buffer remains on the GPU is much faster

Could you add the command that you use to the OP?

You can avoid the KV cache with this patch:

diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 43bf13856..dcc80a052 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -7694,7 +7694,7 @@ struct llm_build_dream : public llm_graph_context {
         // inp_pos - contains the positions
         ggml_tensor * inp_pos = build_inp_pos();
 
-        auto * inp_attn = build_attn_inp_kv_unified();
+        auto * inp_attn = build_attn_inp_no_cache();
 
         ggml_tensor * inp_out_ids = build_inp_out_ids();
 
@@ -16192,6 +16192,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
         // Models that need specific instantiation should be handled in the
         // switch statement
         case LLM_ARCH_BERT:
+        case LLM_ARCH_DREAM:
         case LLM_ARCH_JINA_BERT_V2:
         case LLM_ARCH_NOMIC_BERT:
         case LLM_ARCH_NOMIC_BERT_MOE:

You would need to extend llama_context::encode() to support logits extraction:

diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index a2bb82e93..c89cfd236 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -736,6 +736,7 @@ int llama_context::encode(const llama_batch & batch_inp) {
     const auto & hparams = model.hparams;
 
     const int64_t n_embd = hparams.n_embd;
+    const int32_t n_vocab = model.vocab.n_tokens();
 
     // note: during encode, we always pass the full sequence starting from pos = 0
     if (!balloc->init(batch_inp, model.vocab, nullptr, n_embd, true)) {
@@ -795,10 +796,22 @@ int llama_context::encode(const llama_batch & batch_inp) {
         }
     }
 
+    auto * t_logits = res->get_logits();
     auto * t_embd = res->get_embd_pooled() ? res->get_embd_pooled() : res->get_embd();
 
+    // extract logits
+    if (t_logits && n_outputs > 0) {
+        ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(sched.get(), t_logits);
+        GGML_ASSERT(backend_res != nullptr);
+        GGML_ASSERT(logits != nullptr);
+
+        if (n_outputs) {
+            ggml_backend_tensor_get_async(backend_res, t_logits, logits, 0, n_outputs*n_vocab*sizeof(float));
+        }
+    }
+
     // extract embeddings
-    if (t_embd) {
+    if (cparams.embeddings && t_embd) {
         ggml_backend_t backend_embd = ggml_backend_sched_get_tensor_backend(sched.get(), t_embd);
         GGML_ASSERT(backend_embd != nullptr);
 
@@ -2881,7 +2894,7 @@ llama_token * llama_diffusion_generate(llama_context * ctx, const llama_token *
         }
 
         // Clear KV cache, TODO: can we configure to not use kv-cache?
-        llama_memory_clear(llama_get_memory(ctx), true);
+        //llama_memory_clear(llama_get_memory(ctx), true);
 
         llama_batch batch = llama_batch_init(max_length, 0, 1);
         batch.n_tokens    = max_length;

Btw, using -fa should speed-up the decoding.

Multi-threading the sampling over the mask positions should improve the performance significantly compared to the current implementation.

This is definitely interesting. I think we should support it and try to keep the changes to libllama to a minimum for now. GPU sampling will be done separately - we can probably spec this soon and start working on it.

I'm making the changes and I plan to keep only llm_arch_is_diffusion and llama_vocab_get_mask in the API and move everything else to tools/diffusion

I'm making the changes and I plan to keep only llm_arch_is_diffusion and llama_vocab_get_mask in the API and move everything else to tools/diffusion

For now, it should be examples/diffusion. To be promoted to a tool, it needs to mature because tools are shipped with packages and thus require more effort to maintain.

Actually it looks like the 87% of the time in the GPU case is spent transferring the logits from the GPU => CPU, in the CPU case, the decode step is the bottleneck, so I don't think multithreading on the sampler is going to help here.

Timing Summary for dream-7b-q8

on GPU

Total decode time: 1521.95ms (avg: 11.89ms/step)
Total logits transfer time: 13247.66ms (avg: 103.50ms/step)
Total sampling time: 420.27ms (avg: 3.28ms/step)
Total generation time: 15189.87ms  
Decode/Logits Transfer/Sampling ratio: 0.10/0.87/0.03

on CPU

Total decode time: 132348.18ms (avg: 1033.97ms/step)
Total logits transfer time: 1.67ms (avg: 0.01ms/step)
Total sampling time: 419.97ms (avg: 3.28ms/step)
Total generation time: 132769.82ms
Decode/Logits Transfer/Sampling ratio: 1.00/0.00/0.00

Most likely you didn't measure correctly the time. With CUDA, the llama_decode is asynchronous so your "Total logits transfer time" actually includes the decode time. Add a llama_synchronize() before getting the logits to make sure you are measuring just the transfer.

Most likely you didn't measure correctly the time. With CUDA, the llama_decode is asynchronous so your "Total logits transfer time" actually includes the decode time. Add a llama_synchronize() before getting the logits to make sure you are measuring just the transfer.

Okay, I re-did the calculation and also timed around llama_synchronize, so I think my original points stands in the the transfer is still the bottleneck unless I'm still doing something wrong

Total decode time: 1293.82ms (avg: 10.11ms/step)
Total logits time: 1.93ms (avg: 0.02ms/step)
Total sampling time: 419.34ms (avg: 3.28ms/step)
Total sync time: 13386.34ms (avg: 104.58ms/step)

No, let me explain better:

// measure this for the total decode + GPU->CPU transfer time
{
  llama_decode()
  llama_synchronize()
}

You cannot measure the transfer alone. If you want to do that, you have to modify this copy in llama_context.cpp to not be asynchronous:

    // extract logits
   if (t_logits && n_outputs > 0) {
        ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(sched.get(), t_logits);
        GGML_ASSERT(backend_res != nullptr);
        GGML_ASSERT(logits != nullptr);

        if (n_outputs) {
            ggml_backend_tensor_get_async(backend_res, t_logits, logits, 0, n_outputs*n_vocab*sizeof(float));
        }
    }

For example:

llama_synchronize()

// measure this section to get the GPU->CPU transfer time
{
    ggml_backend_tensor_get_async(backend_res, t_logits, logits, 0, n_outputs*n_vocab*sizeof(float));
    llama_synchronize()
}

Ok, I did this

diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index fced6564..3cff75a3 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -806,7 +806,11 @@ int llama_context::encode(const llama_batch & batch_inp) {
         GGML_ASSERT(logits != nullptr);

         if (n_outputs) {
+            auto t1 = ggml_time_us();
             ggml_backend_tensor_get_async(backend_res, t_logits, logits, 0, n_outputs*n_vocab*sizeof(float));
+            llama_synchronize(this);
+            auto t2 = ggml_time_us();
+            LLAMA_LOG_INFO("logits time: %0.2fms\n", (t2 - t1)/1000.0);
         }
     }

The average here is 105ms/step (total step time being 115ms)

You are missing sync before the first ggml get time

Ah okay, thank you for your patience! logits now ~15ms, which makes sense. The logits size is about 311 Mb (512 x 152064 x 4B), so we're achieving about 20GBps which is consistent with PCIe speeds.

Turns out it was the sampling all along, I was copying a huge vector + doing allocation 🤦‍♂️
Before Total generating time: 132769.82ms
After Total generation time: 19338.11ms

Could you upload a Q4_0 variant (I can't fit the Q8_0 on my machine)? Nvm, I'll quantize the F16.

On your RTX 3090, is it faster using the F16 model?

Could you upload a Q4_0 variant (I can't fit the Q8_0 on my machine)? Nvm, I'll quantize the F16.

On your RTX 3090, is it faster using the F16 model?

It's slightly faster using the f16 version, I can't say for sure, there's some variability in the results and both of them look the same

Very nice, just tested DiffuCoder-7B-cpGRPO, works (but requires --diffusion-steps 512 to output anything but No code provided)!

Edit: Ok, these settings (from README.md) work too:

--temp 0.4 --diffusion-algorithm 3 --diffusion-alg-temp 0.0 --top_p 0.95 --diffusion-steps 256