dflash: revert fixed-shape wip decoder (regressed 472->328 t/s; keep baseline)

step-A fixed-shape+mask foundation preserved at 1a91884 for the P2b kv-cache (dflash.md).

Author: kashif

src/llama-context.cpp

@@ -1383,45 +1383,15 @@ 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: 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.
+        // DFlash decoder positions span [target_ctx ; noise] = arange(n_enc + n_noise)
         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")) {
-                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));
+                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;
                 }
+                ggml_backend_tensor_set(pos_full, pos_data.data(), 0, n_total * sizeof(int32_t));
             }
         }
 

src/models/dflash.cpp

@@ -148,13 +148,9 @@ 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));
 
-    // 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);
+    // 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];
 
     const int64_t n_tokens_kv = n_ctx + n_tokens;
 
@@ -166,14 +162,6 @@ 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;
@@ -218,8 +206,9 @@ 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);
 
-        // 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);
+        // 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);
         cb(cur, "kqv_out", il);
 
         cur = build_lora_mm(layer.wo, cur);