nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

DFlash: add Laguna XS 2.1 support (#2124) 

* Add Laguna XS 2.1 DFlash support

* Fix DFlash prompt capture with microbatches

---------

Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>
This commit is contained in:
Joel Farthing 2026-07-26 23:40:10 -05:00 committed by GitHub
parent 0a4e10c7fb
commit 1a7691fae7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
16 changed files with 339 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

132
convert_hf_to_gguf.py
View File

@ -2333,6 +2333,9 @@ class DFlashDraftModel(Qwen3Model):
_saw_token_embd = False
_saw_output = False
def _causal_attention(self) -> bool:
return False
def _require_target_model_dir(self) -> Path:
if self.target_model_dir is None:
raise ValueError("DFlashDraftModel conversion requires --target-model-dir <matching target model directory>")
@ -2426,7 +2429,7 @@ class DFlashDraftModel(Qwen3Model):
def set_gguf_parameters(self):
super().set_gguf_parameters()
self.gguf_writer.add_causal_attention(False)
self.gguf_writer.add_causal_attention(self._causal_attention())
# MiMo DFlash draft uses partial rotary (partial_rotary_factor=0.5): RoPE is applied to
# only head_dim*partial_rotary_factor dims, the rest are NoPE. Honoring it is required;
# otherwise the upper half of every head gets spurious position rotation it was never
@ -2613,6 +2616,133 @@ class DFlashDraftModel(Qwen3Model):
return tensors
@Model.register("DFlashLagunaForCausalLM")
class DFlashLagunaModel(DFlashDraftModel):
model_arch = gguf.MODEL_ARCH.DFLASH_DRAFT
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._laguna_qkv_ids: set[int] = set()
self._laguna_gate_ids: set[int] = set()
self._laguna_aux_norm_ids: set[int] = set()
def _causal_attention(self) -> bool:
return True
def set_gguf_parameters(self):
dflash_cfg = self.hparams.get("dflash_config")
if not isinstance(dflash_cfg, dict) or dflash_cfg.get("causal") is not True:
raise ValueError("DFlashLagunaForCausalLM requires dflash_config.causal=true")
if self.hparams.get("gating") != "per-head":
raise ValueError("DFlashLagunaForCausalLM currently requires gating='per-head'")
target_hidden_size = self._get_target_hidden_size()
draft_hidden_size = int(self.hparams["hidden_size"])
if target_hidden_size != draft_hidden_size:
raise ValueError(
"DFlashLagunaForCausalLM requires matching target and draft hidden sizes, "
f"got target={target_hidden_size} and draft={draft_hidden_size}"
)
layer_types = self.hparams.get("layer_types")
if not isinstance(layer_types, list) or len(layer_types) != self.block_count:
raise ValueError(
"DFlashLagunaForCausalLM requires one layer_types entry per draft layer"
)
if any(str(layer_type) != "sliding_attention" for layer_type in layer_types):
raise ValueError(
"DFlashLagunaForCausalLM currently requires every draft layer to use sliding_attention"
)
if not self.hparams.get("sliding_window"):
raise ValueError("DFlashLagunaForCausalLM requires sliding_window metadata")
self.hparams["use_sliding_window"] = True
super().set_gguf_parameters()
self.gguf_writer.add_bool(f"{self.gguf_writer.arch}.dflash.laguna", True)
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
top_level_name = name[6:] if name.startswith("model.") else name
hidden_size = int(self.hparams["hidden_size"])
if top_level_name.startswith("aux_hidden_norms.") and top_level_name.endswith(".weight"):
parts = top_level_name.split(".")
if len(parts) != 3 or not parts[1].isdigit():
raise ValueError(f"DFlashLagunaForCausalLM: invalid auxiliary norm name {name!r}")
aux_id = int(parts[1])
if data_torch.ndim != 1 or data_torch.shape[0] != hidden_size:
raise ValueError(
f"DFlashLagunaForCausalLM: auxiliary norm {name!r} has shape "
f"{tuple(data_torch.shape)}, expected [{hidden_size}]"
)
self._laguna_aux_norm_ids.add(aux_id)
tensor_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DFLASH_AUX_HIDDEN_NORM].format(bid=aux_id)
return [(f"{tensor_name}.weight", data_torch)]
if top_level_name.endswith(".self_attn.qkv_proj.weight"):
if bid is None:
raise ValueError(f"DFlashLagunaForCausalLM: can not infer block id for tensor {name!r}")
n_head = int(self.hparams["num_attention_heads"])
n_head_kv = int(self.hparams["num_key_value_heads"])
head_dim = int(self.hparams.get("head_dim", self.hparams["hidden_size"] // n_head))
q_width = n_head * head_dim
k_width = n_head_kv * head_dim
v_width = n_head_kv * head_dim
expected_width = q_width + k_width + v_width
if data_torch.ndim != 2 or data_torch.shape != (expected_width, hidden_size):
raise ValueError(
f"DFlashLagunaForCausalLM: packed QKV tensor {name!r} has shape "
f"{tuple(data_torch.shape)}, expected [{expected_width}, {hidden_size}]"
)
q_weight, k_weight, v_weight = data_torch.split([q_width, k_width, v_width], dim=0)
self._laguna_qkv_ids.add(bid)
result: list[tuple[str, Tensor]] = []
for suffix, weight in (("q_proj", q_weight), ("k_proj", k_weight), ("v_proj", v_weight)):
split_name = name.replace("qkv_proj", suffix)
result.extend(super().modify_tensors(weight, split_name, bid))
return result
if top_level_name.endswith(".self_attn.g_proj.weight"):
if bid is None:
raise ValueError(f"DFlashLagunaForCausalLM: can not infer block id for tensor {name!r}")
gate = data_torch.squeeze().contiguous()
n_head = int(self.hparams["num_attention_heads"])
if gate.ndim != 2 or gate.shape != (n_head, hidden_size):
raise ValueError(
f"DFlashLagunaForCausalLM: attention gate {name!r} has shape "
f"{tuple(gate.shape)}, expected [{n_head}, {hidden_size}]"
)
self._laguna_gate_ids.add(bid)
tensor_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.ATTN_GATE].format(bid=bid)
return [(f"{tensor_name}.weight", gate)]
return super().modify_tensors(data_torch, name, bid)
def prepare_tensors(self):
super().prepare_tensors()
expected_layers = set(range(self.block_count))
dflash_cfg = self.hparams.get("dflash_config")
if not isinstance(dflash_cfg, dict):
raise ValueError("DFlashLagunaForCausalLM requires dflash_config metadata")
target_layer_ids = dflash_cfg.get("target_layer_ids", [])
if not isinstance(target_layer_ids, list) or not target_layer_ids:
raise ValueError("DFlashLagunaForCausalLM requires non-empty target_layer_ids metadata")
expected_aux = set(range(len(target_layer_ids)))
if self._laguna_qkv_ids != expected_layers:
raise ValueError(
f"DFlashLagunaForCausalLM: packed QKV layers {sorted(self._laguna_qkv_ids)} "
f"do not match expected {sorted(expected_layers)}"
)
if self._laguna_gate_ids != expected_layers:
raise ValueError(
f"DFlashLagunaForCausalLM: attention gate layers {sorted(self._laguna_gate_ids)} "
f"do not match expected {sorted(expected_layers)}"
)
if self._laguna_aux_norm_ids != expected_aux:
raise ValueError(
f"DFlashLagunaForCausalLM: auxiliary norm ids {sorted(self._laguna_aux_norm_ids)} "
f"do not match expected {sorted(expected_aux)}"
)
@Model.register("MellumForCausalLM")
class MellumModel(Model):
model_arch = gguf.MODEL_ARCH.MELLUM

4
gguf-py/gguf/constants.py
View File

@ -398,6 +398,7 @@ class MODEL_TENSOR(IntEnum):
MTP_CENTROIDS = auto()
DFLASH_FC = auto()
DFLASH_HIDDEN_NORM = auto()
DFLASH_AUX_HIDDEN_NORM = auto()
# openPangu-2.0 (DSA lightning indexer)
INDEXER_K_NORM = auto()
INDEXER_PROJ = auto() # weights_proj
@ -608,6 +609,7 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
MODEL_TENSOR.MTP_CENTROIDS: "mtp_centroids",
MODEL_TENSOR.DFLASH_FC: "dflash_fc",
MODEL_TENSOR.DFLASH_HIDDEN_NORM: "dflash_hidden_norm",
MODEL_TENSOR.DFLASH_AUX_HIDDEN_NORM: "dflash_aux_hidden_norm.{bid}",
# openPangu-2.0
MODEL_TENSOR.INDEXER_K_NORM: "blk.{bid}.attn_indexer_k_norm",
MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.attn_indexer_weights_proj",
@ -1498,6 +1500,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.ATTN_K,
MODEL_TENSOR.ATTN_K_NORM,
MODEL_TENSOR.ATTN_V,
MODEL_TENSOR.ATTN_GATE,
MODEL_TENSOR.ATTN_SINKS,
MODEL_TENSOR.ATTN_OUT,
MODEL_TENSOR.ATTN_POST_NORM,
@ -1506,6 +1509,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_UP,
MODEL_TENSOR.DFLASH_FC,
MODEL_TENSOR.DFLASH_HIDDEN_NORM,
MODEL_TENSOR.DFLASH_AUX_HIDDEN_NORM,
],
MODEL_ARCH.BITNET: [
MODEL_TENSOR.ATTN_Q,

64
src/graphs/build_dflash.cpp
View File

@ -30,7 +30,39 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() {
ggml_set_input(lctx.dflash.kv.cache_input_pos_ctx);
cb(lctx.dflash.kv.cache_input_pos_ctx, "dflash_kv_input_pos_ctx", -1);
ggml_tensor * fused_target = llm_build_lora_mm(lctx, ctx0, model.dflash_fc, lctx.dflash.kv.cache_input_target_features);
ggml_tensor * target_features = lctx.dflash.kv.cache_input_target_features;
if (hparams.dflash_laguna) {
GGML_ASSERT(model.dflash_aux_hidden_norms.size() == hparams.dflash_n_target_layers);
const int64_t slice_width = n_target_features / hparams.dflash_n_target_layers;
ggml_tensor * normalized_features = nullptr;
for (uint32_t i = 0; i < hparams.dflash_n_target_layers; ++i) {
ggml_tensor * slice = ggml_view_2d(
ctx0,
target_features,
slice_width,
update_rows,
target_features->nb[1],
i * slice_width * target_features->nb[0]);
slice = llm_build_norm(
ctx0,
slice,
hparams,
model.dflash_aux_hidden_norms[i],
nullptr,
LLM_NORM_RMS,
cb,
-1);
cb(slice, "dflash_kv_aux_norm", (int) i);
normalized_features = normalized_features == nullptr
? slice
: ggml_concat(ctx0, normalized_features, slice, 0);
}
GGML_ASSERT(normalized_features != nullptr);
target_features = normalized_features;
cb(target_features, "dflash_kv_normalized_target_features", -1);
}
ggml_tensor * fused_target = llm_build_lora_mm(lctx, ctx0, model.dflash_fc, target_features);
fused_target = llm_build_norm(ctx0, fused_target, hparams, model.dflash_hidden_norm, nullptr, LLM_NORM_RMS, cb, -1);
cb(fused_target, "dflash_kv_fused_target", -1);
@ -38,7 +70,21 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() {
GGML_ASSERT(il < (int32_t) lctx.dflash.kv.k_ctx_cache.size());
GGML_ASSERT(il < (int32_t) lctx.dflash.kv.v_ctx_cache.size());
ggml_tensor * Kcur_ctx_proj = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, fused_target);
ggml_tensor * layer_target = fused_target;
if (hparams.dflash_laguna) {
layer_target = llm_build_norm(
ctx0,
layer_target,
hparams,
model.layers[il].attn_norm,
nullptr,
LLM_NORM_RMS,
cb,
il);
cb(layer_target, "dflash_kv_attn_norm", il);
}
ggml_tensor * Kcur_ctx_proj = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, layer_target);
if (model.layers[il].bk) { Kcur_ctx_proj = ggml_add(ctx0, Kcur_ctx_proj, model.layers[il].bk); }
cb(Kcur_ctx_proj, "dflash_kv_k_proj", il);
@ -54,7 +100,7 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() {
Kcur_ctx = ggml_cont(ctx0, ggml_permute(ctx0, Kcur_ctx, 0, 2, 1, 3));
cb(Kcur_ctx, "dflash_kv_k_physical", il);
ggml_tensor * Vcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, fused_target);
ggml_tensor * Vcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, layer_target);
if (model.layers[il].bv) { Vcur_ctx = ggml_add(ctx0, Vcur_ctx, model.layers[il].bv); }
cb(Vcur_ctx, "dflash_kv_v_proj", il);
if (std::abs(hparams.f_attn_v_scale - 1.0f) > 1e-4f) {
@ -232,6 +278,7 @@ ggml_cgraph * llm_build_context::build_dflash() {
ggml_tensor * cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, cb, il);
cb(cur, "attn_norm", il);
ggml_tensor * input_normed = cur;
ggml_tensor * Qcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wq, cur);
ggml_tensor * Kcur_noise = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, cur);
@ -339,6 +386,17 @@ ggml_cgraph * llm_build_context::build_dflash() {
// cur->op_params[4] = hparams.n_swa;
//}
if (hparams.dflash_laguna) {
GGML_ASSERT(model.layers[il].wqkv_gate != nullptr);
ggml_tensor * gate = llm_build_lora_mm(lctx, ctx0, model.layers[il].wqkv_gate, input_normed);
gate = ggml_softplus(ctx0, gate);
cb(gate, "attn_gate", il);
GGML_ASSERT(gate->ne[0] == n_head);
gate = ggml_reshape_3d(ctx0, gate, 1, n_head, n_tokens);
cur = ggml_mul(ctx0, cur, gate);
cb(cur, "attn_gated", il);
}
cur = ggml_reshape_2d(ctx0, cur, model.layers[il].wo->ne[0], n_tokens);
cb(cur, "flash_attn_reshaped", il);

10
src/graphs/build_laguna.cpp
View File

@ -8,6 +8,7 @@ ggml_cgraph * llm_build_context::build_laguna() {
ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr;
const bool needs_dflash_full_final_rows = lctx.dflash.capture != nullptr;
ggml_tensor * KQ_mask = build_inp_KQ_mask();
// Laguna M.1 has only global-attention layers and leaves n_swa at zero; building
// the SWA mask in that case trips the generic SWA precondition.
@ -23,8 +24,10 @@ ggml_cgraph * llm_build_context::build_laguna() {
GGML_ASSERT(KQ_mask_l != nullptr);
auto rope_factors = is_swa ? nullptr : build_rope_factors(il);
const bool is_final_layer = il == n_layer - 1;
ggml_tensor * attn_out_ids = is_final_layer && !needs_dflash_full_final_rows ? inp_out_ids : nullptr;
auto cur = build_std_attention(gf, model.layers[il].attn_norm, inpL,
inp_pos, il == n_layer - 1 ? inp_out_ids : nullptr, rope_factors,
inp_pos, attn_out_ids, rope_factors,
KQ_mask_l, nullptr, nullptr, 1.0f / sqrtf(float(n_embd_head_k)), 0.0f, n_swa_l, il, true, false, true);
if (model.layers[il].ffn_gate_inp == nullptr) {
@ -53,6 +56,11 @@ ggml_cgraph * llm_build_context::build_laguna() {
cur = lctx.cvec.apply_to(ctx0, cur, il);
cb(cur, "l_out", il);
if (is_final_layer && needs_dflash_full_final_rows && inp_out_ids != nullptr) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
cb(cur, "l_out_selected", il);
}
inpL = cur;
}

1
src/llama-arch.cpp
View File

@ -168,6 +168,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_DFLASH_TARGET_LAYER_IDS, "%s.dflash.target_layer_ids" },
{ LLM_KV_DFLASH_N_TARGET_FEATURES, "%s.dflash.n_target_features" },
{ LLM_KV_DFLASH_BACKBONE_ROTARY_BASE, "%s.dflash.backbone_rotary_base" },
{ LLM_KV_DFLASH_LAGUNA, "%s.dflash.laguna" },
{ LLM_KV_ATTENTION_HEAD_COUNT, "%s.attention.head_count" },
{ LLM_KV_ATTENTION_HEAD_COUNT_KV, "%s.attention.head_count_kv" },

2
src/llama-arch.h
View File

@ -151,6 +151,7 @@ enum llm_kv {
LLM_KV_DFLASH_TARGET_LAYER_IDS,
LLM_KV_DFLASH_N_TARGET_FEATURES,
LLM_KV_DFLASH_BACKBONE_ROTARY_BASE,
LLM_KV_DFLASH_LAGUNA,
LLM_KV_ATTENTION_HEAD_COUNT,
LLM_KV_ATTENTION_HEAD_COUNT_KV,
@ -418,6 +419,7 @@ enum llm_tensor {
LLM_TENSOR_MTP_CENTROIDS,
LLM_TENSOR_DFLASH_FC,
LLM_TENSOR_DFLASH_HIDDEN_NORM,
LLM_TENSOR_DFLASH_AUX_HIDDEN_NORM,
// openPangu-2.0
LLM_TENSOR_ATTN_QA_CONV, // MoME causal conv on q-lora latent

4
src/llama-context.h
View File

@ -362,10 +362,14 @@ struct llama_context {
struct capture_state {
std::vector<int32_t> layer_ids;
std::vector<std::vector<float>> layer_rows;
std::vector<int32_t> layer_rows_written;
int32_t row_count = 0;
int32_t row_width = 0;
int32_t expected_rows = 0;
uint64_t capture_batch_id = 0;
std::vector<uint64_t> layer_seen_batch_id;
bool readback_pending = false;
bool invalid = false;
ggml_backend_sched_eval_callback prev_cb_eval = nullptr;
void * prev_cb_eval_user_data = nullptr;
};

38
src/llama-dflash.cpp
View File

@ -230,6 +230,30 @@ static bool validate_dflash_graph_contract(const llama_context & lctx) {
const auto & model = lctx.model;
const auto & hparams = model.hparams;
if (hparams.dflash_laguna) {
if (!hparams.causal_attn || hparams.n_swa == 0) {
LLAMA_LOG_ERROR("%s: Laguna DFlash requires causal SWA metadata\n", __func__);
return false;
}
if (hparams.dflash_n_target_layers == 0 ||
hparams.dflash_n_target_features % hparams.dflash_n_target_layers != 0 ||
model.dflash_aux_hidden_norms.size() != hparams.dflash_n_target_layers) {
LLAMA_LOG_ERROR("%s: Laguna DFlash requires one auxiliary norm per target feature slice\n", __func__);
return false;
}
const int64_t aux_width = hparams.dflash_n_target_features / hparams.dflash_n_target_layers;
for (uint32_t i = 0; i < hparams.dflash_n_target_layers; ++i) {
const ggml_tensor * aux_norm = model.dflash_aux_hidden_norms[i];
if (aux_norm == nullptr || aux_norm->ne[0] != aux_width) {
LLAMA_LOG_ERROR("%s: Laguna DFlash auxiliary norm %u has invalid width\n", __func__, i);
return false;
}
}
} else if (!model.dflash_aux_hidden_norms.empty()) {
LLAMA_LOG_ERROR("%s: generic DFlash must not carry Laguna auxiliary norms\n", __func__);
return false;
}
auto rope_dim_for_layer = [&hparams](int32_t il) -> uint32_t {
if (hparams.rope_dim_per_layer[il] != 0) {
return hparams.rope_dim_per_layer[il];
@ -273,6 +297,20 @@ static bool validate_dflash_graph_contract(const llama_context & lctx) {
return false;
}
if (hparams.dflash_laguna) {
const ggml_tensor * gate = model.layers[il].wqkv_gate;
if (!hparams.swa_layers[il] || gate == nullptr ||
gate->ne[0] != (int64_t) hparams.n_embd ||
gate->ne[1] != (int64_t) hparams.n_head((uint32_t) il)) {
LLAMA_LOG_ERROR("%s: Laguna DFlash layer %d requires SWA and a head-wise attention gate\n",
__func__, il);
return false;
}
} else if (model.layers[il].wqkv_gate != nullptr) {
LLAMA_LOG_ERROR("%s: generic DFlash layer %d has an unsupported attention gate\n", __func__, il);
return false;
}
const uint32_t rope_dim = rope_dim_for_layer(il);
const float rope_base = rope_base_for_layer(il);
const float rope_scale = rope_scale_for_layer(il);

4
src/llama-hparams.cpp
View File

@ -907,6 +907,10 @@ void llm_load_hparams(
ml.get_key(LLM_KV_DFLASH_BACKBONE_ROTARY_BASE, hparams.dflash_backbone_rotary_base, false);
load_dflash_target_layer_ids(ml, LLM_KV(model.arch)(LLM_KV_DFLASH_TARGET_LAYER_IDS), hparams, false);
ml.get_key(LLM_KV_ATTENTION_VALUE_SCALE, hparams.f_attn_v_scale, false);
ml.get_key(LLM_KV_DFLASH_LAGUNA, hparams.dflash_laguna, false);
if (hparams.dflash_laguna) {
ml.get_key(LLM_KV_ATTENTION_CAUSAL, hparams.causal_attn);
}
// DFlash drafts may be trained with sliding-window attention (for long-context).
// Read the window + per-layer pattern so the SWA mask path activates; absent keys
// leave n_swa=0 / swa_layers all-zero (dense behavior, unchanged).

2
src/llama-hparams.h
View File

@ -171,6 +171,7 @@ struct llama_hparams {
uint32_t dflash_n_target_layers = 0;
uint32_t dflash_target_layer_ids[8] = {};
float dflash_backbone_rotary_base = 0.0f;
bool dflash_laguna = false;
// needed by encoder-decoder models (e.g. T5, FLAN-T5)
// ref: https://github.com/ggerganov/llama.cpp/pull/8141
@ -195,6 +196,7 @@ struct llama_hparams {
if (this->dflash_mask_token_id != other.dflash_mask_token_id) return true;
if (this->dflash_n_target_features != other.dflash_n_target_features) return true;
if (this->dflash_n_target_layers != other.dflash_n_target_layers) return true;
if (this->dflash_laguna != other.dflash_laguna) return true;
if (this->n_layer != other.n_layer) return true;
if (this->n_rot != other.n_rot) return true;
if (this->n_swa != other.n_swa) return true;

17
src/llama-load-tensors.cpp
View File

@ -2302,6 +2302,20 @@ bool create_tensors_helper::create_dflash_tensors(const LLM_TN & tn) {
model.output_mtp = model.output;
model.dflash_fc = create_tensor(ctx_output, tn(LLM_TENSOR_DFLASH_FC, "weight"), {(int64_t) hparams.dflash_n_target_features, n_embd}, 0);
model.dflash_hidden_norm = create_tensor(ctx_output, tn(LLM_TENSOR_DFLASH_HIDDEN_NORM, "weight"), {n_embd}, 0);
model.dflash_aux_hidden_norms.clear();
if (hparams.dflash_laguna) {
GGML_ASSERT(hparams.dflash_n_target_layers > 0);
GGML_ASSERT(hparams.dflash_n_target_features % hparams.dflash_n_target_layers == 0);
const int64_t aux_width = hparams.dflash_n_target_features / hparams.dflash_n_target_layers;
model.dflash_aux_hidden_norms.reserve(hparams.dflash_n_target_layers);
for (uint32_t i = 0; i < hparams.dflash_n_target_layers; ++i) {
model.dflash_aux_hidden_norms.push_back(create_tensor(
ctx_output,
tn(LLM_TENSOR_DFLASH_AUX_HIDDEN_NORM, "weight", i),
{aux_width},
0));
}
}
for (int i = 0; i < n_layer; ++i) {
ggml_context * ctx_split = use_split_ctx ? ctx_for_layer_split(i) : ctx_for_layer(i);
@ -2314,6 +2328,9 @@ bool create_tensors_helper::create_dflash_tensors(const LLM_TN & tn) {
layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_gqa}, 0);
layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_gqa}, 0);
layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_v * n_head, n_embd}, 0);
if (hparams.dflash_laguna) {
layer.wqkv_gate = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head}, 0);
}
layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);

2
src/llama-model.cpp
View File

@ -857,6 +857,7 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> LLM_TENSOR_NA
{ LLM_TENSOR_ATTN_K, "blk.%d.attn_k" },
{ LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" },
{ LLM_TENSOR_ATTN_V, "blk.%d.attn_v" },
{ LLM_TENSOR_ATTN_GATE, "blk.%d.attn_gate" },
{ LLM_TENSOR_ATTN_SINKS, "blk.%d.attn_sinks" },
{ LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" },
{ LLM_TENSOR_ATTN_POST_NORM, "blk.%d.post_attention_norm" },
@ -865,6 +866,7 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> LLM_TENSOR_NA
{ LLM_TENSOR_FFN_UP, "blk.%d.ffn_up" },
{ LLM_TENSOR_DFLASH_FC, "dflash_fc" },
{ LLM_TENSOR_DFLASH_HIDDEN_NORM, "dflash_hidden_norm" },
{ LLM_TENSOR_DFLASH_AUX_HIDDEN_NORM, "dflash_aux_hidden_norm.%d" },
},
},
{

1
src/llama-model.h
View File

@ -473,6 +473,7 @@ struct llama_model {
struct ggml_tensor * mtp_centroids = nullptr;
struct ggml_tensor * dflash_fc = nullptr;
struct ggml_tensor * dflash_hidden_norm = nullptr;
std::vector<struct ggml_tensor *> dflash_aux_hidden_norms;
struct ggml_tensor * output_norm;
struct ggml_tensor * output_norm_b;

67
src/llama-spec-features-dflash.cpp
View File

@ -403,15 +403,41 @@ static int llama_dflash_capture_eval_callback(struct ggml_tensor * tensor, bool
if (capture.layer_seen_batch_id.size() != capture.layer_ids.size()) {
capture.layer_seen_batch_id.assign(capture.layer_ids.size(), 0);
}
if (capture.layer_rows_written.size() != capture.layer_ids.size() || capture.expected_rows <= 0) {
capture.invalid = true;
LLAMA_LOG_WARN("%s: DFlash capture state is not initialized (expected_rows=%d layers=%zu written=%zu)\n",
__func__, capture.expected_rows, capture.layer_ids.size(), capture.layer_rows_written.size());
return 2;
}
if (capture.row_width != 0 && capture.row_width != row_width) {
capture.invalid = true;
LLAMA_LOG_WARN("%s: DFlash capture width mismatch for layer %d: got=%d expected=%d\n",
__func__, layer_id, row_width, capture.row_width);
return 2;
}
auto & rows = capture.layer_rows[(size_t) layer_idx];
rows.resize((size_t) row_count * (size_t) row_width);
auto & rows_written = capture.layer_rows_written[(size_t) layer_idx];
if (rows_written + row_count > capture.expected_rows) {
capture.invalid = true;
LLAMA_LOG_WARN("%s: DFlash capture rows overflow for layer %d: written=%d add=%d expected=%d\n",
__func__, layer_id, rows_written, row_count, capture.expected_rows);
return 2;
}
const size_t expected_floats = (size_t) capture.expected_rows * (size_t) row_width;
if (rows.size() != expected_floats) {
rows.resize(expected_floats);
}
auto backend = ggml_backend_sched_get_tensor_backend(ctx->sched, tensor);
GGML_ASSERT(backend);
ggml_backend_tensor_get_async(backend, tensor, rows.data(), 0, ggml_nbytes(tensor));
ggml_backend_tensor_get_async(backend, tensor,
rows.data() + (size_t) rows_written * (size_t) row_width,
0, (size_t) row_count * (size_t) row_width * sizeof(float));
rows_written += row_count;
capture.row_width = row_width;
capture.row_count = row_count;
capture.row_count = std::max(capture.row_count, rows_written);
capture.layer_seen_batch_id[(size_t) layer_idx] = capture.capture_batch_id;
capture.readback_pending = true;
return 2;
}
@ -423,9 +449,14 @@ bool llama_set_dflash_capture_layers(
return false;
}
if (ctx->dflash.capture && ctx->dflash.capture->readback_pending) {
llama_synchronize(ctx);
}
auto capture = std::make_unique<llama_context::dflash_runtime::capture_state>();
capture->layer_ids.assign(layer_ids, layer_ids + n_layers);
capture->layer_rows.resize((size_t) n_layers);
capture->layer_rows_written.assign((size_t) n_layers, 0);
capture->layer_seen_batch_id.assign((size_t) n_layers, 0);
capture->prev_cb_eval = ctx->cparams.cb_eval;
capture->prev_cb_eval_user_data = ctx->cparams.cb_eval_user_data;
@ -449,6 +480,9 @@ void llama_clear_dflash_capture(struct llama_context * ctx) {
ggml_backend_sched_eval_callback prev_cb_eval = nullptr;
void * prev_cb_eval_user_data = nullptr;
if (ctx->dflash.capture) {
if (ctx->dflash.capture->readback_pending) {
llama_synchronize(ctx);
}
prev_cb_eval = ctx->dflash.capture->prev_cb_eval;
prev_cb_eval_user_data = ctx->dflash.capture->prev_cb_eval_user_data;
}
@ -465,15 +499,22 @@ void llama_clear_dflash_capture(struct llama_context * ctx) {
}
}
void llama_begin_dflash_capture_batch(struct llama_context * ctx) {
void llama_begin_dflash_capture_batch(struct llama_context * ctx, int32_t expected_rows) {
if (ctx == nullptr || !ctx->dflash.capture) {
return;
}
auto & capture = *ctx->dflash.capture;
if (capture.readback_pending) {
llama_synchronize(ctx);
capture.readback_pending = false;
}
capture.capture_batch_id++;
capture.row_count = 0;
capture.row_width = 0;
capture.expected_rows = expected_rows;
capture.invalid = expected_rows <= 0;
std::fill(capture.layer_rows_written.begin(), capture.layer_rows_written.end(), 0);
std::fill(capture.layer_seen_batch_id.begin(), capture.layer_seen_batch_id.end(), 0);
}
@ -504,10 +545,18 @@ static bool llama_spec_prepare_dflash_capture(
llama_synchronize(ctx);
auto & capture = *ctx->dflash.capture;
capture.readback_pending = false;
row_count = capture.row_count;
row_width = capture.row_width;
n_layers = (int32_t) capture.layer_ids.size();
if (row_count <= 0 || row_width <= 0 || n_layers <= 0 || capture.layer_rows.size() != (size_t) n_layers) {
if (capture.invalid || row_count <= 0 || row_width <= 0 || n_layers <= 0 ||
capture.expected_rows <= 0 || capture.layer_rows.size() != (size_t) n_layers ||
capture.layer_rows_written.size() != (size_t) n_layers) {
return false;
}
if (row_count != capture.expected_rows) {
LLAMA_LOG_WARN("%s: DFlash capture rows incomplete: got=%d expected=%d\n",
__func__, row_count, capture.expected_rows);
return false;
}
@ -533,9 +582,11 @@ static bool llama_spec_prepare_dflash_capture(
}
const auto & rows = capture.layer_rows[(size_t) layer_idx];
if (rows.size() != (size_t) row_count * (size_t) row_width) {
LLAMA_LOG_WARN("%s: DFlash capture rows mismatch for layer %d: got=%zu expected=%zu (rows=%d width=%d)\n",
__func__, capture.layer_ids[(size_t) layer_idx], rows.size(),
if (capture.layer_rows_written[(size_t) layer_idx] != row_count ||
rows.size() != (size_t) row_count * (size_t) row_width) {
LLAMA_LOG_WARN("%s: DFlash capture rows mismatch for layer %d: got=%d/%zu expected=%d/%zu (rows=%d width=%d)\n",
__func__, capture.layer_ids[(size_t) layer_idx],
capture.layer_rows_written[(size_t) layer_idx], rows.size(), row_count,
(size_t) row_count * (size_t) row_width, row_count, row_width);
return false;
}

2
src/llama-spec-features-dflash.h
View File

@ -115,7 +115,7 @@ bool llama_set_dflash_target_features_view(
bool llama_set_dflash_capture_layers(struct llama_context * ctx, const int32_t * layer_ids, int32_t n_layers);
void llama_clear_dflash_capture(struct llama_context * ctx);
void llama_begin_dflash_capture_batch(struct llama_context * ctx);
void llama_begin_dflash_capture_batch(struct llama_context * ctx, int32_t expected_rows);
void llama_finish_dflash_capture_batch(struct llama_context * ctx, bool is_prompt_warmup);
bool llama_spec_get_dflash_feature_view(

5
src/llama.cpp
View File

@ -3733,7 +3733,8 @@ static std::pair<std::vector<double>, double> get_layer_sizes(const llama_model_
name == "rope_freqs.weight") {
continue;
}
if (name == "dflash_fc.weight" || name == "dflash_hidden_norm.weight") {
if (name == "dflash_fc.weight" || name == "dflash_hidden_norm.weight" ||
name.rfind("dflash_aux_hidden_norm.", 0) == 0) {
output_misc_size += size;
continue;
}
@ -5851,7 +5852,7 @@ static int llama_decode_internal(
const auto & hparams = model.hparams;
const auto & cparams = lctx.cparams;
llama_begin_dflash_capture_batch(&lctx);
llama_begin_dflash_capture_batch(&lctx, (int32_t) n_tokens_all);
GGML_ASSERT((!batch_all.token && batch_all.embd) || (batch_all.token && !batch_all.embd)); // NOLINT