dflash: fixed-shape masked decoder (wip p2b kv-cache foundation)

masking verified correct (80% accept, greedy-equivalent) but slower than the
472 t/s baseline until the per-layer cross-K/V cache lands - see dflash.md P2b.

Author: kashif

src/llama-context.cpp

@@ -1383,15 +1383,45 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll
         // FIXME this call causes a crash if any model inputs were not used in the graph and were therefore not allocated
         res->set_inputs(&ubatch);
 
-        // DFlash decoder positions span [target_ctx ; noise] = arange(n_enc + n_noise)
+        // DFlash decoder: fill the fixed-size injected context, RoPE positions, and the padding
+        // mask. target_ctx is [n_embd, n_ctx_max]; only the first n_valid columns are real context,
+        // the noise block occupies the last ubatch.n_tokens key slots.
         if (dflash_decoder_ctx && !cross.v_embd.empty()) {
+            const int64_t n_embd  = model.hparams.n_embd;
+            const int64_t n_ctx   = cparams.n_ctx;
+            const int64_t n_valid = std::min<int64_t>(cross.n_enc, n_ctx);
+            const int64_t n_block = ubatch.n_tokens;
+            const int64_t n_kv    = n_ctx + n_block;
+
+            if (ggml_tensor * tctx = ggml_graph_get_tensor(gf, "dflash_target_ctx")) {
+                std::vector<float> buf((size_t) n_embd * n_ctx, 0.0f);
+                std::memcpy(buf.data(), cross.v_embd.data(), (size_t) n_embd * n_valid * sizeof(float));
+                ggml_backend_tensor_set(tctx, buf.data(), 0, buf.size() * sizeof(float));
+            }
             if (ggml_tensor * pos_full = ggml_graph_get_tensor(gf, "inp_pos_full")) {
-                const int64_t n_total = cross.n_enc + ubatch.n_tokens;
-                std::vector<int32_t> pos_data(n_total);
-                for (int64_t i = 0; i < n_total; ++i) {
-                    pos_data[i] = (int32_t) i;
+                std::vector<int32_t> pos(n_kv, 0);
+                for (int64_t i = 0; i < n_valid; ++i) {
+                    pos[i] = (int32_t) i;
+                }
+                for (int64_t j = 0; j < n_block; ++j) {
+                    pos[n_ctx + j] = (int32_t) (n_valid + j);
+                }
+                ggml_backend_tensor_set(pos_full, pos.data(), 0, pos.size() * sizeof(int32_t));
+            }
+            if (ggml_tensor * mask = ggml_graph_get_tensor(gf, "dflash_kq_mask")) {
+                std::vector<float> m((size_t) n_kv * n_block, 0.0f);
+                for (int64_t q = 0; q < n_block; ++q) {
+                    for (int64_t k = n_valid; k < n_ctx; ++k) {
+                        m[(size_t) q * n_kv + k] = -INFINITY;
+                    }
+                }
+                if (mask->type == GGML_TYPE_F16) {
+                    std::vector<ggml_fp16_t> mh(m.size());
+                    ggml_fp32_to_fp16_row(m.data(), mh.data(), m.size());
+                    ggml_backend_tensor_set(mask, mh.data(), 0, mh.size() * sizeof(ggml_fp16_t));
+                } else {
+                    ggml_backend_tensor_set(mask, m.data(), 0, m.size() * sizeof(float));
                 }
-                ggml_backend_tensor_set(pos_full, pos_data.data(), 0, n_total * sizeof(int32_t));
             }
         }
 

src/models/dflash.cpp

@@ -148,9 +148,13 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
     cb(noise_embd, "inp_noise_embd", -1);
     res->add_input(std::move(inp));
 
-    // injected target context: [n_embd, n_enc] from llama_cross::v_embd
-    ggml_tensor * target_ctx = build_inp_cross_embd();
-    const int64_t n_ctx = target_ctx->ne[1];
+    // P2b step A: fixed-size injected target context [n_embd, n_ctx_max] so the decoder graph
+    // shape is constant across draft steps (enables graph reuse + CUDA graphs). The active context
+    // length varies per step; padded rows are masked out by dflash_kq_mask. Filled by cb-name.
+    const int64_t n_ctx = cparams.n_ctx;
+    ggml_tensor * target_ctx = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_ctx);
+    ggml_set_input(target_ctx);
+    cb(target_ctx, "dflash_target_ctx", -1);
 
     const int64_t n_tokens_kv = n_ctx + n_tokens;
 
@@ -162,6 +166,14 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
     // query positions = the noise block (last n_tokens entries)
     ggml_tensor * inp_pos_q = ggml_view_1d(ctx0, inp_pos_full, n_tokens, n_ctx * ggml_element_size(inp_pos_full));
 
+    // additive attention mask [n_kv, n_q]: 0 = attend, -INF = masked. Masks the padded context
+    // rows [n_valid, n_ctx); valid context + the noise block are visible. Filled by cb-name.
+    // flash attention requires an F16 mask; the non-flash soft-max path takes F32
+    ggml_tensor * kq_mask = ggml_new_tensor_4d(ctx0,
+            cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32, n_tokens_kv, n_tokens, 1, 1);
+    ggml_set_input(kq_mask);
+    cb(kq_mask, "dflash_kq_mask", -1);
+
     const float kq_scale = 1.0f/sqrtf(float(n_embd_head));
 
     ggml_tensor * inpL = noise_embd;
@@ -206,9 +218,8 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
                 ext_factor, attn_factor, beta_fast, beta_slow);
         cb(Qcur, "Qcur_rope", il);
 
-        // single-block baseline: full attention (no causal mask). The block-diffusion mask
-        // (causal-to-anchor for context, bidirectional within block) is added with np>1.
-        ggml_tensor * cur = build_attn_mha(Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, nullptr, kq_scale, il);
+        // attention with the padded-context mask (see dflash_kq_mask above)
+        ggml_tensor * cur = build_attn_mha(Qcur, Kcur, Vcur, nullptr, kq_mask, nullptr, nullptr, kq_scale, il);
         cb(cur, "kqv_out", il);
 
         cur = build_lora_mm(layer.wo, cur);