lucebox-hub/dflash/src/common/dflash_draft_graph.cpp at fa34d3173441469159fa27bb6f21f2b7fc44d051 · dusterbloom/lucebox-hub · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
#include "dflash_draft_graph.h"
#include "draft/draft_graph.h"  // DraftGraphInputs, DraftGraphOutputs, build_draft_graph

#include "ggml-alloc.h"

#include <cstdio>

namespace dflash::common {

// Build draft graph at a given ctx_len into sg. Does NOT touch sg.alloc.
// mirror_view: if true, uses a view into mirror->target_feat at slot0.
static bool build_draft_graph_internal(
    StepGraph & sg,
    const DraftWeights & dw,
    ggml_tensor * lm_head,
    int ctx_len,
    const DraftFeatureMirror * mirror,
    int mirror_slot0,
    bool mirror_view) {

    ggml_init_params ip{};
    ip.mem_size   = 256 * 1024 * 1024;
    ip.mem_buffer = nullptr;
    ip.no_alloc   = true;
    sg.ctx = ggml_init(ip);
    if (!sg.ctx) return false;

    const int hidden = dw.n_embd;
    const int q_len  = dw.block_size;
    const int fc_in  = dw.n_target_layers * hidden;

    sg.inp_embed = ggml_new_tensor_3d(sg.ctx, GGML_TYPE_F32, hidden, q_len, 1);
    ggml_set_name(sg.inp_embed, "inp_embed");
    ggml_set_input(sg.inp_embed);

    if (mirror_view) {
        const size_t stride = mirror->target_feat->nb[1];
        sg.target_hidden_cat = ggml_view_3d(
            sg.ctx,
            mirror->target_feat,
            fc_in, ctx_len, 1,
            stride,
            stride * (size_t)ctx_len,
            (size_t)mirror_slot0 * stride);
    } else {
        sg.target_hidden_cat = ggml_new_tensor_3d(sg.ctx, GGML_TYPE_F32, fc_in, ctx_len, 1);
        ggml_set_input(sg.target_hidden_cat);
    }
    ggml_set_name(sg.target_hidden_cat, "target_hidden_cat");

    sg.positions = ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, q_len);
    ggml_set_name(sg.positions, "positions_q");
    ggml_set_input(sg.positions);

    sg.positions_k = ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, ctx_len + q_len);
    ggml_set_name(sg.positions_k, "positions_k");
    ggml_set_input(sg.positions_k);

    sg.gf = ggml_new_graph_custom(sg.ctx, 4096, false);

    DraftGraphInputs gi{};
    gi.ctx_len           = ctx_len;
    gi.noise_embed       = sg.inp_embed;
    gi.target_hidden_cat = sg.target_hidden_cat;
    gi.positions_q       = sg.positions;
    gi.positions_k       = sg.positions_k;
    gi.lm_head           = lm_head;
    DraftGraphOutputs go = build_draft_graph(sg.ctx, dw, gi);
    sg.hidden_states = go.hidden_states;
    sg.logits = go.logits;
    if (!sg.hidden_states) {
        std::fprintf(stderr, "draft graph missing hidden_states\n");
        return false;
    }
    if (sg.logits) {
        sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits);
        ggml_set_name(sg.argmax_tokens, "argmax_tokens");
        ggml_set_output(sg.argmax_tokens);
        ggml_build_forward_expand(sg.gf, sg.argmax_tokens);
    } else {
        ggml_set_output(sg.hidden_states);
        ggml_build_forward_expand(sg.gf, sg.hidden_states);
    }
    return true;
}

bool build_draft_step(
    StepGraph & sg,
    const DraftWeights & dw,
    ggml_tensor * lm_head,
    ggml_backend_t backend,
    int ctx_len,
    const DraftFeatureMirror * mirror,
    int committed,
    int /*ctx_len_max*/) {
    step_graph_free(sg);

    if (!sg.alloc) {
        sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
    }

    int mirror_slot0 = 0;
    const bool use_view = mirror &&
        draft_feature_mirror_can_view(*mirror, committed, ctx_len, mirror_slot0);

    // If ctx_len exceeds our cached reserve, re-reserve at next 64 boundary.
    // This makes all subsequent alloc_graph calls within the 64-token window
    // a no-op (no CUDA free+alloc).
    const int ctx_padded = (ctx_len + 63) & ~63;
    if (ctx_padded > sg.alloc_reserved_ctx) {
        // Build a dummy graph at ctx_padded just for sizing.
        // Use non-view path for reserve (view tensors don't need allocation).
        if (!build_draft_graph_internal(sg, dw, lm_head, ctx_padded,
                                        nullptr, 0, false)) {
            return false;
        }
        ggml_gallocr_reserve(sg.alloc, sg.gf);
        sg.alloc_reserved_ctx = ctx_padded;
        step_graph_free(sg);
    }

    // Build real graph at ctx_len for actual computation.
    if (!build_draft_graph_internal(sg, dw, lm_head, ctx_len,
                                    mirror, mirror_slot0, use_view)) {
        return false;
    }

    return ggml_gallocr_alloc_graph(sg.alloc, sg.gf);
}

}  // namespace dflash::common