diff --git a/common/speculative.cpp b/common/speculative.cpp index ae55e357d51..b0ad48eb483 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -14,6 +14,7 @@ #include #include +#include #include #include #include @@ -923,6 +924,10 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { int32_t block_size = 0; llama_token mask_token_id = 0; + bool is_dflash2 = false; + bool is_mrope = false; + int32_t selector_top_k = 0; + // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout const bool is_dspark; @@ -967,6 +972,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { sample_from_anchor = std::strcmp(buf, "true") == 0; } } + + selector_top_k = llama_model_dflash_selector_top_k(model_dft); + is_dflash2 = selector_top_k > 0; mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str()); @@ -987,6 +995,13 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq); batch_inject = llama_batch_init(llama_n_batch(ctx_dft), n_embd_dec, n_seq); + // embd batches on an M-RoPE draft need 4 position rows per token + is_mrope = llama_model_rope_type(model_dft) == LLAMA_ROPE_TYPE_MROPE; + if (is_mrope) { + free(batch_inject.pos); + batch_inject.pos = (llama_pos *) malloc(sizeof(llama_pos) * 4 * llama_n_batch(ctx_dft)); + } + smpls.resize(n_seq); for (auto & s : smpls) { common_params_sampling sparams; @@ -998,7 +1013,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // offload draft sampling to the backend backend_chains.assign(n_seq, nullptr); - if (this->params.backend_sampling) { + if (this->params.backend_sampling && !is_dflash2) { for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params()); llama_sampler_chain_add(chain, llama_sampler_init_top_k(10)); @@ -1017,7 +1032,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true); } - llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true); + // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits. + llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2); llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention } @@ -1118,11 +1134,24 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } // fuse extracted features through DFlash encoder + // M-RoPE drafts read 4 position rows per token from embd batches, so pass them explicitly + std::vector enc_pos; + if (is_mrope) { + enc_pos.resize((size_t) 4 * n_chunk); + for (int32_t i = 0; i < n_chunk; ++i) { + const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + enc_pos[0 * n_chunk + i] = p; + enc_pos[1 * n_chunk + i] = p; + enc_pos[2 * n_chunk + i] = p; + enc_pos[3 * n_chunk + i] = 0; + } + } + llama_batch enc_batch = { /*.n_tokens =*/ n_chunk, /*.token =*/ nullptr, /*.embd =*/ features_buf.data(), - /*.pos =*/ nullptr, + /*.pos =*/ is_mrope ? enc_pos.data() : nullptr, /*.n_seq_id =*/ nullptr, /*.seq_id =*/ nullptr, /*.logits =*/ nullptr, @@ -1143,7 +1172,13 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float)); for (int32_t i = 0; i < n_chunk; ++i) { - batch_inject.pos[i] = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + batch_inject.pos[i] = p; + if (is_mrope) { + batch_inject.pos[1 * n_chunk + i] = p; + batch_inject.pos[2 * n_chunk + i] = p; + batch_inject.pos[3 * n_chunk + i] = 0; + } batch_inject.n_seq_id[i] = 1; batch_inject.seq_id[i][0] = seq_id; batch_inject.logits[i] = false; @@ -1186,7 +1221,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { i_block_beg[seq_id] = batch.n_tokens; n_block [seq_id] = n_block_tokens; for (int32_t i = 0; i < n_block_tokens; ++i) { - common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true); + common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, !is_dflash2); } } @@ -1214,6 +1249,36 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { auto & result = *dp.result; + if (is_dflash2) { + const float * lattice = llama_get_embeddings_nextn(ctx_dft); + GGML_ASSERT(lattice && "DFlash2 selector produced no lattice"); + + int32_t predecessor = 0; + for (int32_t i = 1; i < n_block_tokens; ++i) { + const float * row = lattice + (size_t) (beg + i) * n_embd_dec; + const float * scores = row + selector_top_k + (size_t) predecessor * selector_top_k; + + predecessor = (int32_t) std::distance(scores, + std::max_element(scores, scores + selector_top_k)); + if (params.p_min > 0.0f) { + // softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max)) + float sum = 0.0f; + for (int32_t k = 0; k < selector_top_k; ++k) { + sum += std::exp(scores[k] - scores[predecessor]); + } + if (1.0f / sum < params.p_min) { + break; + } + } + result.push_back((llama_token) row[predecessor]); + } + + if (result.size() < (size_t) params.n_min) { + result.clear(); + } + continue; + } + if (is_dspark) { // DSpark: read from the first draft slot, truncate below the confidence threshold const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr; diff --git a/conversion/__init__.py b/conversion/__init__.py index 3232a1050bf..f6897330575 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -54,6 +54,7 @@ "DeepseekV3ForCausalLM": "deepseek", "DeepseekV32ForCausalLM": "deepseek", "DFlashDraftModel": "qwen", + "DFlash2DraftModel": "qwen", "Qwen3DSparkModel": "qwen", "DSparkDraftModel": "qwen", "DSparkSpeculator": "qwen", diff --git a/conversion/qwen.py b/conversion/qwen.py index 26b10452b6f..f3944709076 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -639,7 +639,7 @@ class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): model_arch = gguf.MODEL_ARCH.QWEN35MOE -@ModelBase.register("DFlashDraftModel") +@ModelBase.register("DFlashDraftModel", "DFlash2DraftModel") @ModelBase.example("z-lab/Qwen3.5-9B-DFlash") class DFlashModel(Qwen3Model): model_arch = gguf.MODEL_ARCH.DFLASH @@ -678,9 +678,31 @@ def set_vocab(self): def set_gguf_parameters(self): super().set_gguf_parameters() - block_size = self.hparams.get("block_size", 16) - self.gguf_writer.add_block_size(block_size) dflash_config = self.hparams.get("dflash_config", {}) + block_size = dflash_config.get("block_size", self.hparams.get("block_size", 16)) + self.gguf_writer.add_block_size(block_size) + + if "conv_kernel_size" in dflash_config: + self.gguf_writer.add_conv_kernel_size(int(dflash_config["conv_kernel_size"])) + self.gguf_writer.add_conv_group_size(int(dflash_config["conv_group_size"])) + self.gguf_writer.add_selector_rank(int(dflash_config["selector_rank"])) + self.gguf_writer.add_selector_top_k(int(dflash_config["selector_top_k"])) + + output_multiplier = dflash_config.get( + "output_multiplier", self.hparams.get("output_multiplier") + ) + if output_multiplier is not None: + self.gguf_writer.add_logit_scale(float(output_multiplier)) + softcap = dflash_config.get( + "final_logit_softcapping", self.hparams.get("final_logit_softcapping") + ) + if softcap is not None and float(softcap) > 0: + self.gguf_writer.add_final_logit_softcapping(float(softcap)) + embedding_scale = dflash_config.get( + "input_embedding_scale", self.hparams.get("input_embedding_scale") + ) + if embedding_scale is not None: + self.gguf_writer.add_embedding_scale(float(embedding_scale)) target_layer_ids = dflash_config.get("target_layer_ids", []) if target_layer_ids: @@ -695,6 +717,21 @@ def set_gguf_parameters(self): self.gguf_writer.add_sliding_window(sliding_window) self.gguf_writer.add_sliding_window_pattern(is_swa) + # M-RoPE target: the draft ropes on the temporal dim only, so write + # degenerate sections [n_rot/2, 0, 0, 0] + if self._target_uses_mrope(): + head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"] + self.gguf_writer.add_rope_dimension_sections([head_dim // 2, 0, 0, 0]) + + def _target_uses_mrope(self) -> bool: + if self.target_model_dir is None: + return False + with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f: + cfg = json.load(f) + cfg = cfg.get("text_config", cfg) + rope = cfg.get("rope_parameters") or cfg.get("rope_scaling") or {} + return "mrope_section" in rope + @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: name, gen = item @@ -702,10 +739,29 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca name = "model." + name return super().filter_tensors((name, gen)) + _ROPE_PERMUTE_SUFFIXES = ( + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + ) + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: if name == "model.embed_tokens.weight" and not self.hparams.get("has_embed_tokens", True): return + # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd + if not self.hparams.get("rope_is_neox_style", True) and name.endswith(self._ROPE_PERMUTE_SUFFIXES): + head_dim = self.hparams["head_dim"] + shape = data_torch.shape + data_torch = data_torch.reshape(-1, head_dim // 2, 2, *shape[1:]).transpose(1, 2).reshape(shape) + + if name in ( + "model.candidate_selector.predecessor_codebook", + "model.candidate_selector.successor_codebook", + ): + name += ".weight" + yield from super().modify_tensors(data_torch, name, bid) diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index 9681cd29333..a75c6d540e0 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -48,6 +48,101 @@ static int next_power_of_2(int x) { #endif // CUB_TOP_K_AVAILABLE + +// Two-stage top-k for wide rows: a global top-k element has at most k-1 larger +// elements, so at most k-1 inside its own tile and tiling cannot drop a winner. +#define TOPK_CAND 1024 // argsort_f32_i32_cuda_bitonic's row limit + +// measured on H200 and A10G +#define TOPK_BLOCK 256 +#define TOPK_TILE_WIDE 8192 +#define TOPK_TILE 4096 + +template +static __global__ void topk_tile(const float * src, float * cand_val, int * cand_idx, + const int ncols, const int ntiles, const int k) { + __shared__ uint64_t smem[BLOCK]; + + const int row = blockIdx.x / ntiles; + const int tile = blockIdx.x % ntiles; + const float * row_ptr = src + (size_t) row * ncols; + + uint64_t keys[TILE / BLOCK]; +#pragma unroll + for (int i = 0; i < TILE / BLOCK; ++i) { + const int col = tile * TILE + threadIdx.x + i * BLOCK; + uint32_t b = col < ncols ? __float_as_uint(row_ptr[col]) : 0; + b = (b & 0x80000000u) ? ~b : (b | 0x80000000u); + keys[i] = col < ncols ? (((uint64_t) b << 32) | (uint32_t) (ncols - 1 - col)) : 0; + } + + const size_t out = ((size_t) row * ntiles + tile) * k; + for (int j = 0; j < k; ++j) { + uint64_t local = 0; +#pragma unroll + for (int i = 0; i < TILE / BLOCK; ++i) { + local = max(local, keys[i]); + } + smem[threadIdx.x] = local; + __syncthreads(); + for (int s = BLOCK / 2; s > 0; s >>= 1) { + if (threadIdx.x < s) { + smem[threadIdx.x] = max(smem[threadIdx.x], smem[threadIdx.x + s]); + } + __syncthreads(); + } + const uint64_t best = smem[0]; + if (threadIdx.x == 0) { + const int col = ncols - 1 - (int) (best & 0xFFFFFFFFu); + cand_val[out + j] = best ? row_ptr[col] : -INFINITY; + cand_idx[out + j] = best ? col : 0; + } +#pragma unroll + for (int i = 0; i < TILE / BLOCK; ++i) { + if (keys[i] == best) { + keys[i] = 0; + } + } + __syncthreads(); + } +} + +// The argsort ranks candidates; turn its positions back into columns. +static __global__ void topk_unmap(const int * cand_idx, const int * order, int * dst, + const int ncand, const int k) { + for (int i = threadIdx.x; i < k; i += blockDim.x) { + dst[(size_t) blockIdx.x * k + i] = cand_idx[(size_t) blockIdx.x * ncand + order[(size_t) blockIdx.x * ncand + i]]; + } +} + +static bool ggml_cuda_top_k_tiled(ggml_cuda_pool & pool, const float * src, int * dst, + const int ncols, const int nrows, const int k, + cudaStream_t stream) { + // Narrow rows are already handled whole by the bitonic sort below. + const int tile = ncols >= 65536 ? TOPK_TILE_WIDE : TOPK_TILE; + const int ntiles = (ncols + tile - 1) / tile; + const int ncand = ntiles * k; + if (ncols <= TOPK_CAND || ncand > TOPK_CAND) { + return false; + } + + ggml_cuda_pool_alloc cand_val(pool, (size_t) nrows * ncand); + ggml_cuda_pool_alloc cand_idx(pool, (size_t) nrows * ncand); + ggml_cuda_pool_alloc order (pool, (size_t) nrows * ncand); + + if (tile == TOPK_TILE_WIDE) { + topk_tile<<>>( + src, cand_val.get(), cand_idx.get(), ncols, ntiles, k); + } else { + topk_tile<<>>( + src, cand_val.get(), cand_idx.get(), ncols, ntiles, k); + } + argsort_f32_i32_cuda_bitonic(cand_val.get(), order.get(), ncand, nrows, + GGML_SORT_ORDER_DESC, stream); + topk_unmap<<>>(cand_idx.get(), order.get(), dst, ncand, k); + return true; +} + void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const float * src0_d = (const float *) src0->data; @@ -63,6 +158,11 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int64_t nrows = ggml_nrows(src0); const int64_t k = dst->ne[0]; ggml_cuda_pool & pool = ctx.pool(); + + if (ggml_cuda_top_k_tiled(pool, src0_d, dst_d, ncols, nrows, k, stream)) { + return; + } + #ifdef CUB_TOP_K_AVAILABLE // TODO: Switch to `DeviceSegmentedTopK` for multi-row TopK once implemented // https://github.com/NVIDIA/cccl/issues/6391 diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d043c9b6ecc..4c556eda61d 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -162,6 +162,10 @@ class LLM: TARGET_LAYERS = "{arch}.target_layers" TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" + CONV_KERNEL_SIZE = "{arch}.conv_kernel_size" + CONV_GROUP_SIZE = "{arch}.conv_group_size" + SELECTOR_RANK = "{arch}.selector_rank" + SELECTOR_TOP_K = "{arch}.selector_top_k" SAMPLE_FROM_ANCHOR = "{arch}.sample_from_anchor" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" NORM_BEFORE_FC = "{arch}.norm_before_fc" @@ -1133,6 +1137,13 @@ class MODEL_TENSOR(IntEnum): DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed DSPARK_MARKOV_W2 = auto() # markov head: bias projection DSPARK_CONF_PROJ = auto() # confidence head + DFLASH_ATTN_CONV_BASE = auto() + DFLASH_ATTN_CONV_PROJ = auto() + DFLASH_FFN_CONV_BASE = auto() + DFLASH_FFN_CONV_PROJ = auto() + DFLASH_SELECTOR_PREV = auto() + DFLASH_SELECTOR_NEXT = auto() + DFLASH_SELECTOR_HIDDEN = auto() # lfm2 audio A_ENC_NORM_CONV = auto() A_ENC_LINEAR_POS = auto() @@ -1875,6 +1886,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base", + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: "blk.{bid}.attn_conv_proj", + MODEL_TENSOR.DFLASH_FFN_CONV_BASE: "blk.{bid}.ffn_conv_base", + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: "blk.{bid}.ffn_conv_proj", + MODEL_TENSOR.DFLASH_SELECTOR_PREV: "selector_predecessor", + MODEL_TENSOR.DFLASH_SELECTOR_NEXT: "selector_successor", + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: "selector_hidden", MODEL_TENSOR.D2T: "d2t", } @@ -4865,6 +4883,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DSPARK_MARKOV_W1, MODEL_TENSOR.DSPARK_MARKOV_W2, MODEL_TENSOR.DSPARK_CONF_PROJ, + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE, + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ, + MODEL_TENSOR.DFLASH_FFN_CONV_BASE, + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ, + MODEL_TENSOR.DFLASH_SELECTOR_PREV, + MODEL_TENSOR.DFLASH_SELECTOR_NEXT, + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN, ], MODEL_ARCH.MISTRAL4: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 9e0914fd866..5d65ec6fea1 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -981,6 +981,18 @@ def add_sliding_window(self, value: int) -> None: def add_block_size(self, value: int) -> None: self.add_uint32(Keys.LLM.BLOCK_SIZE.format(arch=self.arch), value) + def add_conv_kernel_size(self, value: int) -> None: + self.add_uint32(Keys.LLM.CONV_KERNEL_SIZE.format(arch=self.arch), value) + + def add_conv_group_size(self, value: int) -> None: + self.add_uint32(Keys.LLM.CONV_GROUP_SIZE.format(arch=self.arch), value) + + def add_selector_rank(self, value: int) -> None: + self.add_uint32(Keys.LLM.SELECTOR_RANK.format(arch=self.arch), value) + + def add_selector_top_k(self, value: int) -> None: + self.add_uint32(Keys.LLM.SELECTOR_TOP_K.format(arch=self.arch), value) + def add_sample_from_anchor(self, value: bool) -> None: self.add_bool(Keys.LLM.SAMPLE_FROM_ANCHOR.format(arch=self.arch), value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 3292942b413..0802d2c04f3 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1353,6 +1353,34 @@ class TensorNameMap: "model.confidence_head.proj", # dspark ), + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: ( + "model.layers.{bid}.attention_conv.base_kernel", + ), + + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: ( + "model.layers.{bid}.attention_conv.kernel_projection", + ), + + MODEL_TENSOR.DFLASH_FFN_CONV_BASE: ( + "model.layers.{bid}.mlp_conv.base_kernel", + ), + + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: ( + "model.layers.{bid}.mlp_conv.kernel_projection", + ), + + MODEL_TENSOR.DFLASH_SELECTOR_PREV: ( + "model.candidate_selector.predecessor_codebook", + ), + + MODEL_TENSOR.DFLASH_SELECTOR_NEXT: ( + "model.candidate_selector.successor_codebook", + ), + + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: ( + "model.candidate_selector.hidden_projection", + ), + MODEL_TENSOR.CLS: ( "classifier", # jina "classifier.dense", # roberta diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 5b88bde14d1..ec877d4201f 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -337,6 +337,12 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" }, { LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" }, + { LLM_KV_DFLASH_BLOCK_SIZE, "%s.block_size" }, + { LLM_KV_DFLASH_CONV_KERNEL_SIZE, "%s.conv_kernel_size" }, + { LLM_KV_DFLASH_CONV_GROUP_SIZE, "%s.conv_group_size" }, + { LLM_KV_DFLASH_SELECTOR_RANK, "%s.selector_rank" }, + { LLM_KV_DFLASH_SELECTOR_TOP_K, "%s.selector_top_k" }, + { LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" }, // sentence-transformers dense modules feature dims { LLM_KV_DENSE_2_FEAT_IN, "%s.dense_2_feat_in" }, @@ -644,6 +650,13 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" }, { LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" }, { LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" }, + { LLM_TENSOR_DFLASH_ATTN_CONV_BASE, "blk.%d.attn_conv_base" }, + { LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "blk.%d.attn_conv_proj" }, + { LLM_TENSOR_DFLASH_FFN_CONV_BASE, "blk.%d.ffn_conv_base" }, + { LLM_TENSOR_DFLASH_FFN_CONV_PROJ, "blk.%d.ffn_conv_proj" }, + { LLM_TENSOR_DFLASH_SELECTOR_PREV, "selector_predecessor" }, + { LLM_TENSOR_DFLASH_SELECTOR_NEXT, "selector_successor" }, + { LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "selector_hidden" }, }; // declare information about the model weight tensors: @@ -909,6 +922,13 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, {LLM_TENSOR_DSPARK_MARKOV_W2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_DSPARK_CONF_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DFLASH_ATTN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DFLASH_FFN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_DFLASH_FFN_CONV_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DFLASH_SELECTOR_PREV, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_DFLASH_SELECTOR_NEXT, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, }; LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {} diff --git a/src/llama-arch.h b/src/llama-arch.h index 8042120a25f..7e44842b163 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -380,6 +380,11 @@ enum llm_kv { LLM_KV_TARGET_LAYERS, LLM_KV_TARGET_HIDDEN_SIZE, + LLM_KV_DFLASH_BLOCK_SIZE, + LLM_KV_DFLASH_CONV_KERNEL_SIZE, + LLM_KV_DFLASH_CONV_GROUP_SIZE, + LLM_KV_DFLASH_SELECTOR_RANK, + LLM_KV_DFLASH_SELECTOR_TOP_K, LLM_KV_NORM_BEFORE_RESIDUAL, LLM_KV_NORM_BEFORE_FC, @@ -652,6 +657,13 @@ enum llm_tensor { LLM_TENSOR_DSPARK_MARKOV_W1, LLM_TENSOR_DSPARK_MARKOV_W2, LLM_TENSOR_DSPARK_CONF_PROJ, + LLM_TENSOR_DFLASH_ATTN_CONV_BASE, + LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, + LLM_TENSOR_DFLASH_FFN_CONV_BASE, + LLM_TENSOR_DFLASH_FFN_CONV_PROJ, + LLM_TENSOR_DFLASH_SELECTOR_PREV, + LLM_TENSOR_DFLASH_SELECTOR_NEXT, + LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, }; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 52f8d53672a..ac2be97a653 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2307,6 +2307,10 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_MINIMAX_01 || model.arch == LLM_ARCH_MINIMAX_M3) { res = std::max(n_tokens * 40, 32u * model.n_tensors()); + } else if (model.arch == LLM_ARCH_DFLASH && model.hparams.dflash_selector_rank > 0) { + // DFlash2's convolutions and selector are shape work rather than matmuls, + // so they cost ~8.6 nodes per tensor against ~5.9 for a plain DFlash draft + res = std::max(1024u, 12u*model.n_tensors()); } else { res = std::max(1024u, 8u*model.n_tensors()); for (const auto & lora : model.loras) { diff --git a/src/llama-ext.h b/src/llama-ext.h index 35d6e58adfa..92a759b7a0a 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -120,6 +120,8 @@ LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx); // model/context data extraction // +LLAMA_API int32_t llama_model_dflash_selector_top_k(const struct llama_model * model); + // returns pointer to the target-model layer indices LLAMA_API const int32_t * llama_model_target_layer_ids (const struct llama_model * model); // returns the number of extracted layers from target model diff --git a/src/llama-hparams.h b/src/llama-hparams.h index e91ce1cc3c7..5ddcae84dd1 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -214,6 +214,12 @@ struct llama_hparams { // output embedding dimension (0 = use n_embd) uint32_t n_embd_out_impl = 0; + uint32_t dflash_block_size = 0; + uint32_t dflash_conv_kernel_size = 0; + uint32_t dflash_conv_group_size = 0; + uint32_t dflash_selector_rank = 0; + uint32_t dflash_selector_top_k = 0; + // llama4 smallthinker uint32_t n_moe_layer_step = 0; uint32_t n_no_rope_layer_step = 4; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 0d74a2135b6..de3b608f6ba 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2535,6 +2535,10 @@ int32_t llama_model_n_layer_nextn(const llama_model * model) { return model->hparams.n_layer_nextn; } +int32_t llama_model_dflash_selector_top_k(const llama_model * model) { + return model->hparams.dflash_selector_top_k; +} + int32_t llama_model_n_head(const llama_model * model) { return model->hparams.n_head(); } @@ -2731,6 +2735,10 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { return LLAMA_ROPE_TYPE_NEOX; case LLM_ARCH_DFLASH: + // drafts for M-RoPE targets carry rope sections and follow the target's temporal dim + if (const auto & s = model->hparams.rope_sections; s[0] || s[1] || s[2] || s[3]) { + return LLAMA_ROPE_TYPE_MROPE; + } // DSV4 DSpark drafters use DeepSeek-V4's normal RoPE; legacy DFlash backbones are NeoX return model->hparams.dsv4_hc_mult > 0 ? LLAMA_ROPE_TYPE_NORM : LLAMA_ROPE_TYPE_NEOX; diff --git a/src/llama-model.h b/src/llama-model.h index 4412ef08e74..e2cd04946e5 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -362,6 +362,11 @@ struct llama_layer { struct ggml_tensor * ffn_exp_probs_b = nullptr; struct ggml_tensor * ffn_gate_tid2eid = nullptr; + struct ggml_tensor * dflash_attn_conv_base = nullptr; + struct ggml_tensor * dflash_attn_conv_proj = nullptr; + struct ggml_tensor * dflash_ffn_conv_base = nullptr; + struct ggml_tensor * dflash_ffn_conv_proj = nullptr; + // mamba proj struct ggml_tensor * ssm_in = nullptr; struct ggml_tensor * ssm_x = nullptr; @@ -648,6 +653,10 @@ struct llama_model { struct ggml_tensor * dspark_conf_proj = nullptr; struct ggml_tensor * dspark_conf_proj_b = nullptr; + struct ggml_tensor * dflash_selector_prev = nullptr; + struct ggml_tensor * dflash_selector_next = nullptr; + struct ggml_tensor * dflash_selector_hidden = nullptr; + // unified vector to store target-model extracted layer ids in eagle3, dflash, etc. std::vector target_layer_ids; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 5b70a517949..2ca0491a37d 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -7,6 +7,18 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale, false); + hparams.f_final_logit_softcapping = 0.0f; + ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false); + + // drafts for M-RoPE targets carry degenerate sections [n_rot/2, 0, 0, 0] + ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false); + + ml.get_key(LLM_KV_DFLASH_BLOCK_SIZE, hparams.dflash_block_size, false); + ml.get_key(LLM_KV_DFLASH_CONV_KERNEL_SIZE, hparams.dflash_conv_kernel_size, false); + ml.get_key(LLM_KV_DFLASH_CONV_GROUP_SIZE, hparams.dflash_conv_group_size, false); + ml.get_key(LLM_KV_DFLASH_SELECTOR_RANK, hparams.dflash_selector_rank, false); + ml.get_key(LLM_KV_DFLASH_SELECTOR_TOP_K, hparams.dflash_selector_top_k, false); if (!ml.get_arr(LLM_KV_TARGET_LAYERS, target_layer_ids, false)) { throw std::runtime_error("DFlash model requires 'target_layers' in GGUF metadata"); @@ -112,6 +124,29 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank); } + const struct ggml_tensor * selector_meta = ml->get_tensor_meta("selector_hidden.weight"); + if (selector_meta) { + const int64_t rank = hparams.dflash_selector_rank; + if (rank <= 0 || hparams.dflash_block_size <= 0 || hparams.dflash_selector_top_k <= 0 || + hparams.dflash_conv_kernel_size <= 0 || hparams.dflash_conv_group_size <= 0) { + throw std::runtime_error("DFlash2 model is missing conv/selector metadata"); + } + if (n_embd % hparams.dflash_conv_group_size != 0) { + throw std::runtime_error("DFlash2 hidden size must be divisible by conv_group_size"); + } + if (n_embd < hparams.dflash_selector_top_k * (hparams.dflash_selector_top_k + 1)) { + throw std::runtime_error("DFlash2 hidden size is too small for the selector lattice"); + } + + dflash_selector_prev = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_PREV, "weight"), { rank, n_vocab }, 0); + dflash_selector_next = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_NEXT, "weight"), { rank, n_vocab }, 0); + dflash_selector_hidden = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "weight"), { n_embd, rank }, 0); + + LLAMA_LOG_INFO("%s: DFlash2 conv kernel = %u, group = %u, selector rank = %u, top-k = %u\n", __func__, + hparams.dflash_conv_kernel_size, hparams.dflash_conv_group_size, + hparams.dflash_selector_rank, hparams.dflash_selector_top_k); + } + fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0); fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED); output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) @@ -187,6 +222,16 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0); layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0); layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0); + + if (selector_meta) { + const int64_t kernel = hparams.dflash_conv_kernel_size; + const int64_t groups = n_embd / hparams.dflash_conv_group_size; + const int64_t projected = 2 * kernel * groups; + layer.dflash_attn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_BASE, i), { n_embd, kernel, 2 }, 0); + layer.dflash_attn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "weight", i), { n_embd, projected }, 0); + layer.dflash_ffn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_BASE, i), { n_embd, kernel, 2 }, 0); + layer.dflash_ffn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_PROJ, "weight", i), { n_embd, projected }, 0); + } } } @@ -345,6 +390,167 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_build_forward_expand(g.gf, out); } +static ggml_tensor * build_dflash2_conv( + llm_graph_context & g, + ggml_tensor * hidden, + ggml_tensor * dynamic, + ggml_tensor * base, + int side) { + const auto & hparams = g.hparams; + const int64_t hidden_size = hidden->ne[0]; + const int64_t n_tokens = hidden->ne[1]; + const int64_t n_blocks = g.ubatch.n_seqs_unq; + const int64_t kernel_size = hparams.dflash_conv_kernel_size; + const int64_t group_size = hparams.dflash_conv_group_size; + const int64_t n_groups = hidden_size / group_size; + + GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0); + GGML_ASSERT(dynamic && base && side >= 0 && side < 2); + + const int64_t block_size = n_tokens / n_blocks; + ggml_context * ctx0 = g.ctx0; + // ggml_cont copies even when the tensor is already contiguous + if (!ggml_is_contiguous(hidden) || hidden->ne[1] != n_tokens) { + hidden = ggml_cont_2d(ctx0, hidden, hidden_size, n_tokens); + } + if (!ggml_is_contiguous(dynamic) || dynamic->ne[1] != n_tokens) { + dynamic = ggml_cont_2d(ctx0, dynamic, dynamic->ne[0], n_tokens); + } + ggml_tensor * blocks = ggml_reshape_3d(ctx0, hidden, hidden_size, block_size, n_blocks); + ggml_tensor * coeffs = ggml_reshape_4d(ctx0, dynamic, n_groups, kernel_size, 2, n_tokens); + ggml_tensor * coeffs_side = ggml_view_3d(ctx0, coeffs, n_groups, kernel_size, n_tokens, + coeffs->nb[1], coeffs->nb[3], side * coeffs->nb[2]); + + ggml_tensor * coeff_all = ggml_cont(ctx0, coeffs_side); + coeff_all = ggml_reshape_4d(ctx0, coeff_all, 1, n_groups, kernel_size, n_tokens); + coeff_all = ggml_repeat_4d(ctx0, coeff_all, group_size, n_groups, kernel_size, n_tokens); + + ggml_tensor * base_side = ggml_reshape_4d(ctx0, + ggml_view_1d(ctx0, base, hidden_size * kernel_size, side * base->nb[2]), + group_size, n_groups, kernel_size, 1); + + ggml_tensor * weight_all = ggml_add(ctx0, coeff_all, base_side); + + ggml_tensor * result = nullptr; + for (int64_t tap = 0; tap < kernel_size; ++tap) { + ggml_tensor * values = blocks; + if (tap > 0) { + ggml_tensor * zeros = ggml_fill(ctx0, + ggml_new_tensor_3d(ctx0, hidden->type, hidden_size, std::min(tap, block_size), n_blocks), 0.0f); + if (tap < block_size) { + ggml_tensor * previous = ggml_view_3d(ctx0, blocks, hidden_size, block_size - tap, n_blocks, + blocks->nb[1], blocks->nb[2], 0); + values = ggml_concat(ctx0, zeros, previous, 1); + } else { + values = zeros; + } + } + values = ggml_reshape_2d(ctx0, values, hidden_size, n_tokens); + + ggml_tensor * weight = ggml_reshape_2d(ctx0, + ggml_cont(ctx0, ggml_view_4d(ctx0, weight_all, group_size, n_groups, 1, n_tokens, + weight_all->nb[1], weight_all->nb[2], weight_all->nb[3], tap * weight_all->nb[2])), + hidden_size, n_tokens); + + ggml_tensor * term = ggml_mul(ctx0, weight, values); + result = result ? ggml_add(ctx0, result, term) : term; + } + return result; +} + +// DFlash2 selector: top-k candidates per block position plus the pairwise +// transition scores, packed into the nextn output slot for the CPU-side walk. +static void build_dflash2_selector(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) { + ggml_context * ctx0 = g.ctx0; + auto & res = g.res; + + const auto & hparams = g.hparams; + const int64_t n_tokens = g.n_tokens; + const int64_t n_embd = g.n_embd; + + const int64_t top_k = hparams.dflash_selector_top_k; + const int64_t rank = hparams.dflash_selector_rank; + const int64_t n_blocks = g.ubatch.n_seqs_unq; + GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0); + GGML_ASSERT(res->t_logits->ne[1] == n_tokens); + if (!tokens) { + return; + } + + const int64_t tokens_per_block = n_tokens / n_blocks; + const int64_t block_size = std::min(tokens_per_block, hparams.dflash_block_size); + const int64_t row_used = top_k + top_k * top_k; + + ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k); + ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, res->t_logits, 1, res->t_logits->ne[0], n_tokens); + ggml_tensor * unary = ggml_reshape_2d(ctx0, + ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens); + ggml_tensor * gate = g.build_lora_mm(model.dflash_selector_hidden, res->t_embd); + + // Everything below indexes [.., tokens_per_block, n_blocks]: the block + // position varies fastest, sequences are the outer dimension. + ggml_tensor * cand_blk = ggml_reshape_3d(ctx0, candidates, top_k, tokens_per_block, n_blocks); + ggml_tensor * unary_blk = ggml_reshape_3d(ctx0, unary, top_k, tokens_per_block, n_blocks); + ggml_tensor * gate_blk = ggml_reshape_3d(ctx0, gate, rank, tokens_per_block, n_blocks); + + // a position's score reads only the candidate sets at pos-1 and pos, so a run + // of positions has no internal dependency and scores in one batched matmul + auto score_run = [&](int64_t beg_pos, int64_t n_pos, ggml_tensor * pred_ids) { + ggml_tensor * cand_run = ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, n_pos, n_blocks, + cand_blk->nb[1], cand_blk->nb[2], beg_pos * cand_blk->nb[1])); + ggml_tensor * unary_run = ggml_cont(ctx0, ggml_view_3d(ctx0, unary_blk, top_k, n_pos, n_blocks, + unary_blk->nb[1], unary_blk->nb[2], beg_pos * unary_blk->nb[1])); + ggml_tensor * gate_run = ggml_cont(ctx0, ggml_view_3d(ctx0, gate_blk, rank, n_pos, n_blocks, + gate_blk->nb[1], gate_blk->nb[2], beg_pos * gate_blk->nb[1])); + + const int64_t n_pred = pred_ids->ne[0] / (n_pos * n_blocks); + + ggml_tensor * successor = ggml_reshape_4d(ctx0, + ggml_get_rows(ctx0, model.dflash_selector_next, ggml_reshape_1d(ctx0, cand_run, top_k * n_pos * n_blocks)), + rank, top_k, n_pos, n_blocks); + ggml_tensor * predecessor = ggml_reshape_4d(ctx0, + ggml_get_rows(ctx0, model.dflash_selector_prev, pred_ids), + rank, n_pred, n_pos, n_blocks); + + ggml_tensor * gate_bcast = ggml_reshape_4d(ctx0, gate_run, rank, 1, n_pos, n_blocks); + ggml_tensor * cond = ggml_mul(ctx0, predecessor, ggml_repeat(ctx0, gate_bcast, predecessor)); + ggml_tensor * score = ggml_mul_mat(ctx0, successor, cond); + if (n_pred == 1) { + score = ggml_repeat_4d(ctx0, score, top_k, top_k, n_pos, n_blocks); + } + ggml_tensor * unary_bcast = ggml_reshape_4d(ctx0, unary_run, top_k, 1, n_pos, n_blocks); + score = ggml_add(ctx0, score, ggml_repeat(ctx0, unary_bcast, score)); + + ggml_tensor * row = ggml_concat(ctx0, + ggml_cast(ctx0, cand_run, GGML_TYPE_F32), + ggml_reshape_3d(ctx0, score, top_k * top_k, n_pos, n_blocks), 0); + return ggml_pad(ctx0, row, n_embd - row_used, 0, 0, 0); + }; + + ggml_tensor * packed = ggml_fill(ctx0, + ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd, 1, n_blocks), 0.0f); + + if (block_size > 1) { + // Position 1 alone: its predecessor is the anchor token, one id per + // sequence rather than a candidate set. + ggml_tensor * anchor_ids = ggml_cont_1d(ctx0, + ggml_view_2d(ctx0, tokens, 1, n_blocks, tokens_per_block * tokens->nb[0], 0), n_blocks); + packed = ggml_concat(ctx0, packed, score_run(1, 1, anchor_ids), 1); + } + if (block_size > 2) { + ggml_tensor * prev_ids = ggml_reshape_1d(ctx0, + ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, block_size - 2, n_blocks, + cand_blk->nb[1], cand_blk->nb[2], cand_blk->nb[1])), + top_k * (block_size - 2) * n_blocks); + packed = ggml_concat(ctx0, packed, score_run(2, block_size - 2, prev_ids), 1); + } + + packed = ggml_reshape_2d(ctx0, packed, n_embd, block_size * n_blocks); + g.cb(packed, "dflash2_lattice", -1); + res->t_h_nextn = packed; + ggml_build_forward_expand(g.gf, packed); +} + // DFlash decoder, dual-mode by batch type: // * embd batch -> fused target features: project + inject K/V into the cache. // * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens @@ -369,6 +575,20 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); + // drafts for M-RoPE targets use degenerate sections (temporal dim only) + int sections[4]; + std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); + + auto build_rope = [&](ggml_tensor * cur, ggml_tensor * pos) { + return rope_type == GGML_ROPE_TYPE_MROPE + ? ggml_rope_multi(ctx0, cur, pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow) + : ggml_rope_ext(ctx0, cur, pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + }; + // KV cache injection if (ubatch.embd) { auto inp = std::make_unique(n_embd); @@ -391,11 +611,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); + Kcur = build_rope(Kcur, inp_pos); cb(Kcur, "Kcur_injected", il); cb(Vcur, "Vcur_injected", il); @@ -449,6 +665,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); ggml_set_input(inp->tokens); + res->t_inp_tokens = inp->tokens; ggml_tensor * inp_tokens = inp->tokens; @@ -463,6 +680,13 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra ggml_tensor * noise_norm = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); cb(noise_norm, "noise_norm", il); + ggml_tensor * attn_dynamic = nullptr; + if (layer.dflash_attn_conv_proj) { + attn_dynamic = build_lora_mm(layer.dflash_attn_conv_proj, noise_norm); + noise_norm = build_dflash2_conv(*this, noise_norm, attn_dynamic, layer.dflash_attn_conv_base, 0); + cb(noise_norm, "attn_conv_in", il); + } + ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm); ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm); ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm); @@ -474,16 +698,8 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - Qcur = ggml_rope_ext( - ctx0, Qcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); - Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); + Qcur = build_rope(Qcur, inp_pos); + Kcur = build_rope(Kcur, inp_pos); cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); cb(Vcur, "Vcur", il); @@ -493,12 +709,24 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il) : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + if (attn_dynamic) { + cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1); + cb(cur, "attn_conv_out", il); + } + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); cb(ffn_inp, "ffn_inp", il); cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); cb(cur, "ffn_norm", il); + ggml_tensor * ffn_dynamic = nullptr; + if (layer.dflash_ffn_conv_proj) { + ffn_dynamic = build_lora_mm(layer.dflash_ffn_conv_proj, cur); + cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 0); + cb(cur, "ffn_conv_in", il); + } + cur = build_ffn(cur, layer.ffn_up, NULL, layer.ffn_up_s, layer.ffn_gate, NULL, layer.ffn_gate_s, @@ -507,6 +735,11 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra LLM_FFN_SILU, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); + if (ffn_dynamic) { + cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 1); + cb(cur, "ffn_conv_out", il); + } + cur = ggml_add(ctx0, cur, ffn_inp); cb(cur, "l_out", il); @@ -531,6 +764,19 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra cur = build_lora_mm(output, cur, output_s); + // DFlash2 feeds these logits to the selector, so they need the target's output + // transforms; DFlash1 and DSpark read them through the sampler instead + if (model.dflash_selector_hidden) { + if (hparams.f_logit_scale != 0.0f) { + cur = ggml_scale(ctx0, cur, hparams.f_logit_scale); + } + if (hparams.f_final_logit_softcapping > 0.0f) { + cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping); + cur = ggml_tanh(ctx0, cur); + cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping); + } + } + // reduced-draft-vocab exports: scatter the draft logits to the target vocabulary via d2t if (model.d2t) { const int64_t n_draft_vocab = cur->ne[0]; @@ -555,6 +801,10 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra if (model.dspark_markov_w1) { build_dspark_markov_head(*this, model, inp_tokens); } + + if (model.dflash_selector_hidden) { + build_dflash2_selector(*this, model, inp_tokens); + } } // DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above): diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 3349a64b174..0a8ff97c2aa 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9637,6 +9637,17 @@ static std::vector> make_test_cases_eval() { } } } + for (int k : {4, 8, 16, 32}) { + for (int nrows : {1, 8, 16}) { + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {202048, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {151936, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8192, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8193, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8192, nrows, 1, 1}, k, true)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {202048, nrows, 1, 1}, k, true)); + } + } + for (int k : {1, 2, 3, 7, 15}) { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {16, 10, 10, 10}, k)); test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {60, 10, 10, 10}, k)); @@ -10321,7 +10332,13 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_argsort(GGML_TYPE_F32, {200000, 16, 1, 1})); test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2, 1, 1, 1}, 1)); - for (auto k : {1, 10, 40, 400}) { + // widths around the tiling threshold + for (auto cols : {4096, 8192, 12288, 16384, 24576, 32768, 65536, 131072}) { + for (auto nrows : {1, 16}) { + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, 16)); + } + } + for (auto k : {1, 4, 8, 10, 16, 32, 40, 400}) { for (auto nrows : {1, 16}) { for (auto cols : {k, 1000, 65000, 200000}) { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, k));