Skip to content

dflash : run DFlash2 under split mode tensor #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Piggidragon wants to merge 2 commits into
base: llama/dev
Choose a base branch
from
Open

dflash : run DFlash2 under split mode tensor #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5d9ee9b
ggml-meta : skip a view of data outside the meta buffers
Piggidragon Sep 13, 2026
6b0fbc2
dflash : pick the DFlash2 candidates on the CPU under split mode tensor
Piggidragon Sep 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 16 additions & 9 deletions ggml/src/ggml-backend-meta.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1674,6 +1674,11 @@ bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf) {
return buf != nullptr && buf->iface.free_buffer == ggml_backend_meta_buffer_iface.free_buffer;
}

// a view whose data lives outside the meta buffers, it is a noop for the simple backends
static bool ggml_backend_meta_is_foreign_view(const ggml_tensor * node) {
return node->view_src != nullptr && !ggml_backend_buffer_is_meta(node->view_src->buffer);
}

void ggml_backend_meta_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage) {
GGML_ASSERT(ggml_backend_buffer_is_meta(buffer));
ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context;
Expand Down Expand Up @@ -2010,8 +2015,9 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend,

for (int i = 0; i < cgraph->n_nodes; i++) {
ggml_tensor * node = cgraph->nodes[i];
if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) {
if (ggml_backend_meta_is_foreign_view(node)) {
// FIXME s_copy_main is on the CPU and its view seems to be incorrectly added to the graph nodes.
// The scheduler keeps any view in the split where it sits, so a view of a CPU op can land here too.
// For regular usage this doesn't matter since it's a noop but trying to call ggml_backend_meta_buffer_simple_tensor results in a crash.
bcj.nodes[i] = node;
continue;
Expand Down Expand Up @@ -2143,7 +2149,7 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend,
}
}

if (next->view_src != nullptr && next->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(next->view_src->buffer)) {
if (ggml_backend_meta_is_foreign_view(next)) {
continue;
}
if (ggml_backend_meta_get_split_state(next, false).axis != GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
Expand Down Expand Up @@ -2182,14 +2188,15 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend,
int i_start = 0;
for (int i = 0; i < cgraph->n_nodes; i++) {
ggml_tensor * node = cgraph->nodes[i];
if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) {
continue;
}
const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(node, /*assume_sync =*/ false);
if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
max_tmp_size = std::max(max_tmp_size, ggml_nbytes(node));
// a view of foreign data needs no split state, but it can be the last node and must close the last subgraph
bool new_subgraph = i + 1 == cgraph->n_nodes;
if (!ggml_backend_meta_is_foreign_view(node)) {
const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(node, /*assume_sync =*/ false);
if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
max_tmp_size = std::max(max_tmp_size, ggml_nbytes(node));
new_subgraph = true;
}
}
const bool new_subgraph = i + 1 == cgraph->n_nodes || split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL;
if (!new_subgraph) {
continue;
}
Expand Down
21 changes: 17 additions & 4 deletions src/models/dflash.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -493,10 +493,23 @@ static void build_dflash2_selector(llm_graph_context & g, const llama_model & mo
const int64_t block_size = std::min<int64_t>(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 * logits = res->t_logits;
// split mode tensor splits the logits over the vocab, and the Meta backend cannot run top-k on that
// copy them to the CPU, the scheduler gathers the slices on the way
const bool vocab_on_cpu = model.split_mode() == LLAMA_SPLIT_MODE_TENSOR;
if (vocab_on_cpu) {
logits = ggml_cont(ctx0, logits);
ggml_backend_sched_set_tensor_backend(g.sched, logits, g.backend_cpu);
}

ggml_tensor * candidates = ggml_top_k(ctx0, logits, top_k);
ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, logits, 1, logits->ne[0], n_tokens);
ggml_tensor * unary_rows = ggml_get_rows(ctx0, logits_rows, candidates);
ggml_tensor * unary = ggml_reshape_2d(ctx0, unary_rows, top_k, n_tokens);
if (vocab_on_cpu) {
ggml_backend_sched_set_tensor_backend(g.sched, candidates, g.backend_cpu);
ggml_backend_sched_set_tensor_backend(g.sched, unary_rows, g.backend_cpu);
}
ggml_tensor * gate = g.build_lora_mm(model.dflash_selector_hidden, res->t_embd);

// Everything below indexes [.., tokens_per_block, n_blocks]: the block
Expand Down
Loading